diff --git a/.github/skills/pepsy-vmc/SKILL.md b/.github/skills/pepsy-vmc/SKILL.md new file mode 100644 index 0000000..ea5af00 --- /dev/null +++ b/.github/skills/pepsy-vmc/SKILL.md @@ -0,0 +1,298 @@ +--- +name: pepsy-vmc +description: Design, implement, debug, or review Pepsy variational Monte Carlo workflows with PyTorch, Quimb PEPS/MPS amplitudes, Symmray fermionic tensors, Metropolis-Hastings sampling, local estimators, observables, gradients, and stochastic reconfiguration. Use for work in src/pepsy/vmc, Fermi-Hubbard VMC examples, PEPS amplitude batching, symmetry-constrained samplers, or integrations inspired by sjdu10/vmc_torch. +--- + +# Pepsy VMC + +Use this skill for Pepsy VMC work. Keep the public workflow small and +automatic while making the internal state layout, conserved charges, fermionic +signs, amplitude contraction mode, and numerical scaling explicit. + +## Read first + +- Read `AGENTS.md` and nested instructions. +- Read the current VMC implementation in `src/pepsy/vmc/torch.py` and its + focused tests in `tests/test_vmc_torch.py`. +- Read the native fermion guidance in + `.github/skills/pepsy-fermion-operators/SKILL.md` and + `.github/skills/pepsy-fermion-operators/references/design.md` when changing + spinful operators, basis conventions, or Symmray signs. +- For Fermi-Hubbard workflows, inspect + `/home/reza.haghshenas@quantinuum.com/pepsy_examples/fermi_hubbard/fh_peps.ipynb` + and `fh_mps.ipynb`. Treat notebook source cells as authoritative over + stale stored outputs. +- Read `references/vmc_torch.md` when using patterns from the upstream + `sjdu10/vmc_torch` implementation. + +## Core architecture + +Keep these responsibilities separate: + +1. **PEPS amplitude model**: owns packed Quimb/Symmray tensor parameters and + evaluates `psi(configs)` or stable `(phase, log_abs)` values. +2. **Fermion/layout adapter**: derives PEPS site order, local basis, geometry, + symmetry, and charge-sector metadata from the PEPS and `Fermion` object. +3. **Sampler**: proposes valid configurations, evaluates acceptance ratios, + and retains current configurations and amplitudes in sync. +4. **Connection compiler**: turns local Fermion operators into connected + configurations, coefficients, and source-sample indices. +5. **Estimator**: evaluates local observables and returns mean, variance, + standard error, samples, and diagnostics. +6. **VMC driver**: orchestrates sampling, observable evaluation, + log-derivatives, SR/minSR or optimizer updates, checkpoints, and optional + distributed reductions. + +Do not make the sampler construct Hamiltonians or make the amplitude model +decide fermionic operator signs. The public VMC entry point should normally +accept a PEPS and an explicit observable mapping; a `Fermion` is needed when +the default Hamiltonian must be constructed, but can be omitted when the PEPS +exposes symmetry metadata and the caller supplies `terms`. `Lx`, `Ly`, +`site_order`, local encoding, and sampler-rule details should be inferred +internally. + +The current native Torch entry point is `TorchFermionVMC`. Treat it as the +integration seam: it should prepare validated metadata and delegate amplitude, +connection, sampling, and estimation work to the existing lower-level kernels. + +The stateful native sampler is `TorchMetropolisSampler`. Its public sampling +contract follows NetKet-style chain semantics: `n_samples` is the requested +total across chains, `n_discard_per_chain` controls per-chain burn-in, and +`sweep_size` controls thinning. Accept `n_discard` and `n_thin` as convenience +aliases at user-facing boundaries. Preserve the chain axis in returned arrays +with shape `(n_samples_per_chain, n_chains, n_sites)`; flatten only when a +local-estimator kernel requires a batch matrix. For spinful fermions, the +proposal remains symmetry-aware and is not necessarily a literal one-site +flip. + +When the amplitude model exposes `forward_log(configs) -> (phase, log_abs)`, +Metropolis acceptance must use the clipped log ratio +`exp(min(0, 2 * (log_abs_new - log_abs_old)))`; retain the raw amplitudes for +local-estimator compatibility and fall back to raw ratios only for models +without a log interface. For chain-shaped scalar observables, use +`torch_chain_diagnostics(...)` to report `r_hat`, integrated autocorrelation +time, and effective sample size. Do not report diagnostics from flattened +samples because the chain axis is required for between-chain variance. + +For an approximate global proposal, use `PepsBpSampler` with +`TorchBPMetropolisSampler`. Quimb's public `sample_d2bp` currently samples +binary output indices, so four-state spinful PEPS legs must be represented as +two occupation-bit legs in a private dense BP copy. Preserve the original +Symmray PEPS for Torch amplitudes, gradients, local estimators, and fermionic +signs. Pass the resolved `FermionSiteEncoding` explicitly when constructing +the BP sampler, or let `TorchFermionVMC.make_bp_sampler(...)` do so. + +Treat BP as a proposal `q_BP`, not as the VMC target. For independent BP +Metropolis proposals, retain `log q_BP` for every current chain and accept +`x -> y` using +`min(1, |psi(y)|**2 q_BP(x) / (|psi(x)|**2 q_BP(y)))`. Initialize chains from +BP when `q_BP` is not available for caller-supplied initial configurations. +For direct importance sampling, use self-normalized weights +`|psi(x)|**2 / q_BP(x)` and report effective sample size. Reject or discard +non-finite/zero proposal probabilities and configurations outside the target +sector; do not silently reinterpret them as valid fermion states. + +The BP adapter and sector validator support spinful `U1`, `U1U1`, `Z2`, and +`Z2Z2`. `Z2Z2` moves must preserve each flavor parity independently; a move +that toggles only one empty/double site preserves total `Z2` parity but is not +valid for `Z2Z2`. + +## Metadata inference and validation + +Derive, in this order: + +- PEPS `sites` and physical-index order; +- `Lx`, `Ly`, and lattice edges from the PEPS geometry/tags; +- physical dimension and Symmray sector/index metadata from PEPS arrays; +- spinful/spinless local space, symmetry, operator basis, and native terms from + `Fermion`; +- the actual global charge sector from PEPS metadata or a validated initial + configuration. + +Validate PEPS and Fermion compatibility before sampling. At minimum check +site count, physical dimension, local sector map, spinful flag, and symmetry. +Never silently guess a charge sector or silently permute local basis values. +If the PEPS cannot expose its target sector, raise a concise error requesting +an explicit sector or valid initial configuration. + +Keep the native Pepsy/Symmray metadata as the source of truth, but distinguish +two local orders: native `Fermion.dense_operator(...)` matrices are built from +`empty, up, down, double`, while a Symmray PEPS physical index is ordered by +its charge map and commonly exposes the physical codes +`empty, down, up, double`. Upstream `vmc_torch` examples use a different +semantic site-code order in some paths. Derive the PEPS encoding from the +physical charge map and convert exactly once; never infer it from integer +labels or dense-operator positions alone. + +## Symmetry-constrained sampling + +Choose the proposal rule from the resolved symmetry and sector: + +| symmetry | conserved sector | valid default move | +| --- | --- | --- | +| spinless `U1` | total particle number | exchange or hop one particle | +| spinful `U1` | total `N_up + N_down` | any ergodic move preserving total number | +| spinful `U1U1` | `(N_up, N_down)` separately | same-spin hop or full-site exchange | +| spinful `Z2` | total parity | exchange/hop, spin flip, and empty/double pair toggle | +| `Z2Z2` | parity per flavor | flavor-parity-preserving custom rule | + +For a spinful `U1` state, do not impose separate spin-number conservation +merely because the local space has two flavors. Separate conservation can be +an additional physical-sector choice, but `U1U1` must enforce it. + +Initialize walkers inside the target sector. Every proposal must preserve the +sector, and the move set must be checked for ergodicity within that sector. +Track no-op, proposed, accepted, and effective configuration-changing moves. +For symmetric exchange/hop proposals, the acceptance ratio is simply based on +`abs(psi_new)**2 / abs(psi_old)**2`. If proposal probabilities differ in the +forward and reverse directions, include the Hastings correction. + +Prefer fully vectorized Torch proposals that keep configurations, random +choices, and masks on the target device. Retain current amplitudes and update +only accepted walkers. For `torch.compile`, use fixed-shape full-batch forward +passes when variable changed-subset shapes would trigger recompilation. + +## Amplitude evaluation + +Use Quimb's public packing/unpacking path or the existing Pepsy wrappers: + +- pack the PEPS tensor leaves and retain an immutable skeleton; +- register numeric leaves as Torch parameters without losing Symmray metadata; +- evaluate configurations in PEPS site order; +- use exact contraction for tiny reference systems and boundary/HOTRG/CTMRG + style contraction for production; +- batch over configurations, not parameters; +- use `torch.vmap` only when the traced contraction path is pure, non-mutating, + and compatible with the selected backend; +- keep gradients enabled only for gradient/optimization phases. + +There is no PEPS analogue of a globally canonical MPS form. Do not assume +canonicalization makes PEPS amplitude evaluation cheap. The contraction mode, +bond cutoff, boundary bond, and approximation error are part of the VMC +configuration and should be reported. + +Support a stable log-amplitude representation for wide dynamic ranges. A +convenient contract is `(phase, log_abs)` or `(mantissa, exponent)`. Compute +Metropolis ratios and connected-amplitude ratios from log magnitudes or +exponent differences; do not divide underflowed raw amplitudes. Preserve +complex phases for local estimators. + +## Observable and Hamiltonian connections + +Accept explicit term containers, including a mapping, a Fermion Hamiltonian, +or an object exposing `.terms`. Compile each supported local term into a +batched connection structure equivalent to: + +- `configs`: connected configurations; +- `coeffs`: matrix elements for each connection; +- `batch_ids`: source walker for each connection. + +Use output axes followed by input axes for dense one- and two-site operators, +and validate rank, local dimension, site support, and PEPS site order. Reuse +parent amplitudes for diagonal connections. Chunk connected amplitudes when +the number of walkers times nonzero operator transitions is large. + +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 +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 +same connection/estimator path. + +Construct native Symmray fermion terms with Pepsy's `Fermion` operators. Do +not add Jordan-Wigner strings to the native graded path. A dense/JW reference +is acceptable only as a separately named small-system oracle. Never combine +native fermionic signs with an additional manually inserted parity sign. + +## Optimization and gradients + +Separate measurement from optimization: + +- inference mode: Metropolis, amplitudes, connections, local estimators; +- gradient mode: differentiable amplitudes and log-derivative matrix; +- update mode: Adam/SGD or SR/minSR direction and parameter update. + +For a parameterized amplitude, use log derivatives +`O_k(x) = d log(psi(x)) / d theta_k` and the standard covariance form for the +energy gradient. Center sample means consistently. For SR, expose diagonal +shift, solver choice, convergence, residual, and condition diagnostics. Use +double precision for reference and ill-conditioned SR calculations; mixed +precision requires explicit validation. + +Do not detach connected amplitudes when a differentiable energy estimator is +requested. Conversely, do not retain autograd graphs during ordinary +Metropolis sampling. Reset gradients and temporary contraction/cache state at +the end of every VMC step. + +## Diagnostics and validation + +Always validate the smallest available system before performance work: + +- direct Quimb amplitude versus Torch amplitude for every local basis state; +- native Fermion local basis and sector map; +- sampler invariants for `U1`, `U1U1`, and `Z2`; +- two-site hopping signs against the native Fermion operator and a dense/JW + oracle kept as a reference path; +- connection local estimator versus direct dense expectation; +- zero, empty, full, single-particle, and doublon configurations; +- PEPS deterministic local expectation versus VMC mean within sampling error; +- acceptance rate, no-op rate, effective sample count, and chain mixing; +- finite, non-NaN amplitudes, local energies, gradients, and SR directions. + +Report contraction approximation (`exact`, boundary, HOTRG, or other), bond +cutoff, sample count, burn-in, sweeps between measurements, acceptance, and +standard error. Keep energy and observable measurement tests separate from +optimizer tests. + +## Pepsy-specific validation + +Use Python 3.12 and run focused checks first: + +```bash +source ~/envs/py312/bin/activate +pytest -q tests/test_vmc_torch.py +pytest -q tests/test_symmetric_tensors.py +``` + +For public API changes also run the public API and package-layout tests. For +notebook changes, do not modify generated outputs unless explicitly requested; +validate the relevant setup and measurement cells in a controlled run. + +## Performance baseline and run planning + +The native `TorchFermionVMC` path is functionally usable for the current +4x6, `D=4`, spinful `U1U1` workflow. A completed CPU run with target sector +`(N_up, N_down) = (12, 12)` produced finite amplitudes, valid-sector +diagnostics, stable SR linear solves, consistent boundary/CTMRG estimates, +and a final JSONL record with `status="ok"`. The deterministic PEPS energy +was `E/N = -0.425322`; the post-optimization boundary estimate was +`E/N = -0.424972 +/- 0.000691`, and the exact final diagnostic was +`E/N = -0.425612 +/- 0.000726`. These values are a functional baseline, not +an optimizer or ground-state convergence certificate. + +The same run took about 10.5 hours on CPU. The dominant cost was repeated +local-estimator amplitude contraction, not DMRG or SR: each 2,000-sample +boundary or CTMRG measurement took roughly 2.1--2.3 hours, while the eight +SR steps took only minutes. The final exact 2,000-sample diagnostic took about +1.3 hours. Boundary and CTMRG were numerically indistinguishable at the +reported precision in this baseline, but both were still evaluated. The +connected-amplitude cache also recorded tens of thousands of fallback +contractions, which explains why sampler acceptance and chain mixing alone do +not predict runtime. + +For development runs, use one measurement contraction, a smaller measurement +bond (for example `chi=32`), 256--512 diagnostic samples, and one or two VMC +steps. Omit redundant `measurement_grid` entries until the amplitude and +estimator paths are validated. Re-enable larger samples, boundary/CTMRG +cross-checks, and final diagnostics only for a production measurement. CPU +thread limits such as `OMP_NUM_THREADS=1` and `MKL_NUM_THREADS=1` should be +changed only in controlled benchmarks because extra threads can increase +memory use and oversubscription. + +Preserve unrelated user changes in the worktree. Do not recreate Gaugy, +Tensy, or an upstream `vmc_torch` package inside Pepsy. diff --git a/.github/skills/pepsy-vmc/agents/openai.yaml b/.github/skills/pepsy-vmc/agents/openai.yaml new file mode 100644 index 0000000..572c750 --- /dev/null +++ b/.github/skills/pepsy-vmc/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Pepsy VMC" + short_description: "Build and debug Pepsy VMC workflows" + default_prompt: "Use $pepsy-vmc to design or implement a Pepsy VMC sampling and measurement workflow." diff --git a/.github/skills/pepsy-vmc/references/vmc_torch.md b/.github/skills/pepsy-vmc/references/vmc_torch.md new file mode 100644 index 0000000..b522193 --- /dev/null +++ b/.github/skills/pepsy-vmc/references/vmc_torch.md @@ -0,0 +1,131 @@ +# Upstream `sjdu10/vmc_torch` reference + +Use this document to borrow implementation patterns, not to copy the package +or override Pepsy's native conventions. + +Source repository: + +## Useful separation of concerns + +The upstream repository separates: + +- `vmc_torch/VMC.py`: high-level optimization driver, statistics, MPI + reductions, checkpointing, and optional preconditioning; +- `vmc_torch/sampler.py`: CPU Metropolis state, burn-in, chain sampling, and + configuration conversion; +- `vmc_torch/GPU/sampler.py`: GPU-vectorized proposal and acceptance kernels; +- `vmc_torch/GPU/VMC.py`: GPU driver that orchestrates sampler, local energy, + gradient, SR, optimizer, and distributed synchronization; +- `vmc_torch/hamiltonian_torch.py`: Hilbert-space and operator connection + abstractions; +- `vmc_torch/GPU/vmc_utils.py`: batched energy and gradient accumulation; +- `vmc_torch/GPU/models/pureTNS.py`: packed Quimb tensor-network parameters, + boundary contractions, basis remapping, and optional compiled/vmapped PEPS + amplitude paths. + +Preserve this division in Pepsy. A sampler should only propose and accept +walkers. An estimator should only build connections and evaluate local +estimators. A VMC driver should coordinate the phases. + +## Configuration conventions + +The upstream spinful Hilbert representation commonly uses a flattened mode +configuration: + +- first all spin-up modes; +- then all spin-down modes. + +Its Quimb site encoding is commonly `0=empty, 1=down, 2=up, 3=double`, with +conversion to/from the flattened mode representation. This is not the current +native Pepsy `Fermion` basis. Native `Fermion` dense local operators use +`empty, up, down, double`, while the Symmray PEPS physical index is ordered +by its charge map and commonly appears as `empty, down, up, double`. + +Therefore: + +1. Choose the ordered Symmray PEPS physical charge map as the public PEPS + configuration contract; for the usual `U1U1` map this is + `empty, down, up, double`. +2. Keep the dense native `Fermion` operator basis (`empty, up, down, double`) + separate from the PEPS physical-index order. +3. Convert upstream-style or NetKet-style flattened configurations exactly + once at an explicit boundary. +4. Test the conversion on all four local states and on hopping transitions. +5. Never infer a permutation from the integer values alone. + +## Sampler patterns + +The upstream GPU sampler has several useful patterns: + +- keep current configurations and amplitudes synchronized; +- propose all walkers for one graph edge at once; +- evaluate only changed walkers in eager mode; +- evaluate a fixed full batch when export/compile requires static shapes; +- accept using the squared-amplitude ratio; +- provide log-amplitude mode using `2 * (log_abs_new - log_abs_old)`; +- count accepted and effective configuration-changing moves; +- cache graph edges and same-spin mode-neighbor tables on the device. + +Its spinful exchange/hopping rule can preserve both spin counts by combining: + +- full-site exchange on a lattice edge; +- a same-spin particle hop along a lattice edge; +- local four-state transitions `empty <-> up/down <-> double` that preserve + each flavor count. + +For Pepsy, choose the move set from the resolved symmetry. The upstream +spinful rule is suitable for `U1U1`; for total `U1`, add spin flips so the +chain can change spin-resolved counts. For spinful `Z2`, add an +`empty <-> double` pair toggle so the chain can change particle number by two +while preserving parity. + +## Amplitude patterns + +The upstream PEPS/TNS path follows the Quimb packing pattern: + +1. copy the tensor network so the source network is not mutated; +2. call `qtn.pack` to obtain numeric leaves and a skeleton; +3. flatten leaves into Torch parameters; +4. reconstruct with `qtn.unpack` inside the amplitude function; +5. select physical indices in the network's site order; +6. use exact or boundary contraction; +7. optionally trace/vmap/compile over configurations. + +Some upstream fPEPS models clear a Symmray physical-leg line map and store a +local basis permutation. That permutation is model-specific. In Pepsy, read +the live PEPS/Symmray index map and the `Fermion` registry instead of copying a +hard-coded permutation such as `[0, 2, 3, 1]`. + +For large dynamic ranges, the upstream code supports sign/phase plus log +magnitude or mantissa/exponent outputs. Preserve this idea in Pepsy, but keep +complex phase and backend/autograd behavior correct for native Symmray data. + +## Local-energy and gradient patterns + +The upstream VMC estimator follows the standard connected-configuration +contract: + +- each operator produces connected configurations and matrix elements; +- the model evaluates amplitudes on those configurations; +- local values are matrix-element-weighted amplitude ratios; +- sample means and variances are accumulated by chain/rank; +- gradients use the covariance between local energy and log-amplitude + derivatives; +- SR/minSR preconditions the gradient before the optimizer update. + +Use this as the conceptual model for Pepsy's `TorchConnections` and local +energy functions. Keep connection construction separate from the PEPS +contraction backend so exact, boundary, and future compiled paths can share +the estimator. + +## Distributed and diagnostic patterns + +The upstream driver supports MPI/`torch.distributed` reductions, rank-local +sample batches, parameter synchronization, checkpointing, and diagnostics for +local-energy outliers, log-amplitude spread, gradient norms, NaNs, and Infs. + +Pepsy should first provide a correct single-process/device implementation. +Add distributed execution only after sector invariants, amplitude ratios, +local estimators, and SR results are validated. Reuse the same diagnostics: +acceptance/no-op rates, log-amplitude range, local-energy variance, effective +sample count, gradient norm, SR residual, and parameter norm. diff --git a/.github/skills/stabilizer-tensor-networks/SKILL.md b/.github/skills/stabilizer-tensor-networks/SKILL.md index 2743c9e..174c18e 100644 --- a/.github/skills/stabilizer-tensor-networks/SKILL.md +++ b/.github/skills/stabilizer-tensor-networks/SKILL.md @@ -53,8 +53,11 @@ basis; magic / non-stabilizerness lives in $|\nu\rangle$. ## Diagnostic contract -- Treat `track_infidelity` as a **normalized-unitary norm-loss proxy**, not exact overlap - fidelity, per-gate discarded SVD weight, or a general non-unitary error metric. +- Treat `track_infidelity` as the canonical ``infidelity`` API. For unitary + updates it is the normalized norm-loss proxy; for compressed dense + non-unitary matrices the retained norm ratio is measured against the local + physical ``G^dagger G`` target. It is not a general ideal-circuit overlap + calculation. - Read `1 - ||p||^2` from the tracked one-site canonical centre, including Quimb's MPS `exponent`. Do not build an uncapped target, contract ``, or renormalize after a unitary update. @@ -62,26 +65,31 @@ basis; magic / non-stabilizerness lives in $|\nu\rangle$. unitaries, measurements/projectors, arbitrary non-unitary matrices, and `submpo` events emit no sample. Never sum this list; each emitted value is already cumulative within its normalized unitary segment. -- An unnormalized non-unitary matrix or coefficient-frame `submpo` invalidates the proxy, - so later unitaries remain unreported. A normalized projective collapse starts a fresh - segment at zero but does not itself append a sample. +- An unnormalized non-unitary matrix invalidates the current unitary norm + segment after recording its ``G^dagger G``-normalized compression loss. + A coefficient-frame `submpo` has no certified physical target norm and + invalidates the unitary proxy without adding a sample. A normalized + projective collapse starts a fresh segment at zero but does not itself append + a unitary sample. - Projective measurement/reset boundaries are recorded separately in `.norm_events`. Each event snapshots the pre-collapse unitary segment norm, the physical Born branch probability, the actual projected norm before renormalization, and the post-normalized norm. Compare `projected_norm_sq` to `pre_norm_sq * branch_probability` to get the projector-compression proxy `projector_infidelity`. Use `norm_diagnostics()` for product/geometric summaries that multiply unitary and projector compression-survival - factors, but never measurement probabilities. Prefer `norm_infidelity`, - `norm_survival`, and `norm`; the older `total_*_proxy` keys are compatibility - aliases. `geometric_mean_norm` is only the per-segment geometric mean, not the + factors, but never measurement probabilities. Prefer `infidelity`, `fidelity`, + `norm_survival`, and `norm`; `norm_infidelity` and + the older `total_*_proxy` keys are compatibility aliases. `geometric_mean_norm` + is only the per-segment geometric mean, not the total norm summary. - A selected `TrajectoryEvent` Kraus outcome is a normalized trajectory boundary. Before applying its non-unitary matrix, snapshot the current segment with its branch probability; normalize the selected branch at the canonical centre, reset the proxy, and commit a `"trajectory_kraus"` norm event. This lets later unitary steps track a fresh segment without treating - the Born probability as compression loss. STN progress reports only - `norm_infidelity` plus a compact stream `part` label. + the Born probability as compression loss. STN progress reports the shared + `infidelity` field plus a compact stream `part` label and retains + `norm_infidelity` as a compatibility alias. - Treat stream-local stochastic entries as the primary Pepsy noise design. `("x_error", p, q)`, `("depolarize1", p, q)`, `("depolarize2", p, q0, q1)`, `("pauli_channel1", probs, q)`, diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index 008d4ae..4529de3 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -1,6 +1,6 @@ --- name: tree-optimizer -description: 'Run, review, debug, benchmark, or extend pepsy.TreeOptimizer -- the rooted tree-tensor-network (TTN) gate-stream circuit simulator of Seitz et al. (Quantum 7, 964, 2023; arXiv:2206.01000). Use when the user asks to replay a circuit on a tree tensor network; build or score a TreeLayoutFinder / TreePlan (entanglement-adapted recursive spectral bisection); tune the exact 2q-gate threading + single canonical compression sweep; work on the sibling-leaf fast path, canonical single-/multi-site local_expectation (Steiner subtree), measure/reset, the BLAS thread cap, the canonical-centre tracker, the tid cache, convergence_sweep / bond_report diagnostics, or copy(); or asks why TTN truncation is more accurate than per-hop truncation, how the tree geodesic threading works, or how to add sampling / stream-wired control events. Not for MPS (pepsy.MpsOptimizer -> mps-optimizer skill), MERA (qmera-energy-optimizer), stabilizer TN (stabilizer-tensor-networks), or BP.' +description: 'Run, review, debug, benchmark, or extend pepsy.TreeOptimizer -- the rooted tree-tensor-network (TTN) gate-stream circuit simulator of Seitz et al. (Quantum 7, 964, 2023; arXiv:2206.01000). Use when the user asks to replay a circuit on a tree tensor network; build or score a TreeLayoutFinder / TreePlan (entanglement-adapted recursive spectral bisection); handle exact product-state handoff, TTN layout safety, or Torch/CuPy backend compatibility; tune the exact 2q-gate threading + single canonical compression sweep; work on TreeTensorNetwork canonical single-/multi-site local_expectation (Steiner subtree), measurement/reset, independent or coalesced noisy trajectories, the BLAS thread cap, the canonical-centre tracker, the tid cache, convergence_sweep / bond_report diagnostics, or copy(); or asks why TTN truncation is more accurate than per-hop truncation, how the tree geodesic threading works, or how to add sampling / stream-wired control events. Not for MPS (pepsy.MpsOptimizer -> mps-optimizer skill), MERA (qmera-energy-optimizer), stabilizer TN (stabilizer-tensor-networks), or BP.' --- # Tree Optimizer in pepsy @@ -29,6 +29,49 @@ whose leaves carry the physical qubit indices; a bundled gate stream (2q); supports with `len(where) >= 3` route through `apply_subtree_operator` (see *Multi-qubit / sub-MPO application*). +Preferred public handoff: + +```python +finder = TreeLayoutFinder(gate_stream, n=n, objective="congestion") +choice = finder.recommend_arities((2, 3, 4)) +optimizer = TreeOptimizer(gate_stream, tree=choice["plan"], chi=chi) +``` + +Alternatively pass the finder itself with `layout=finder`. An initial +non-product or entangled state must be passed explicitly as `state=` (alias +`tn=`). `tree=` and `layout=` accept only `TreePlan` / `TreeLayoutFinder`; a +`TreeTensorNetwork` passed there must raise a clear error rather than being +silently replaced by the default `|0...0>` state. + +### State/layout handoff and backend contract + +`TreeLayoutFinder` is **circuit-only**: it accepts a gate stream or explicit +supports, never a `TreeTensorNetwork`. Passing a TTN to it must raise a clear +`TypeError`; construct the plan from the circuit, then pass the state separately +to `TreeOptimizer`. + +- An entangled `TreeTensorNetwork` (`max_bond() > 1`) can be installed only on + its matching `TreePlan`. A different `tree=` / `layout=` must raise before + tensor work: implicit relayout would be lossy and hide a fidelity decision. +- A product `TreeTensorNetwork` (`max_bond() == 1`) may be remounted **exactly** + on a requested plan; warn that its geometry changed. Preserve the product + vectors and any distributed global scalar/phase. +- A bond-one Quimb `MatrixProductState` is also a geometry-neutral product input + and may be mounted exactly on the selected tree. Reject an entangled MPS; + converting it to a tree is an explicit caller-controlled operation. +- All live state tensors must use one backend, dtype, and device. + `backend_info()` reports this. Reject a mixed state at construction or + `set_tn`, rather than choosing an arbitrary execution backend. +- Callers should convert every gate/operator/sub-MPO/observable/cap vector with + the same backend converter as the state. Explicit array mismatches are + coerced for compatibility with a one-time `UserWarning` per source/target + signature; untyped Python lists/scalars are convenience inputs and are + materialized silently. Internal Pauli, projector, reset, and tree-MPO helper + tensors must be built on the state backend without warning. +- Backend-aware contractions must stay in Autoray/Quimb. Convert only scalar + readouts intentionally (`to_float`); `to_dense()` intentionally returns a + host NumPy vector for interoperability while the live TTN remains native. + ## Tree state class (`TreeTensorNetwork`) `pepsy.TreeTensorNetwork` (also `pepsy.optimizers.TreeTensorNetwork`, source @@ -37,19 +80,24 @@ whose leaves carry the physical qubit indices; a bundled gate stream `quimb.tensor.TensorNetworkGenVector` (import from `quimb.tensor`, **not** the deprecated `tensor_arbgeom`). It owns a `TreePlan` plus the node/site/index naming, so all inherited quimb methods (`canonize_around`, `canonize_between`, -`compress_between`, `gate_inds`, `local_expectation`, `to_dense`, `copy`) work -directly. + `compress_between`, `gate_inds`, `to_dense`, `copy`) work directly. The + state overrides `local_expectation` with a canonical Steiner-subtree + contraction that also supports native Symmray observables. - `_EXTRA_PROPS = ("_sites", "_site_tag_id", "_site_ind_id", "_plan", - "_node_tag_id")` -- these are copied on `.copy()`/every quimb view. The - `__init__` copy-branch guard `if isinstance(ts, TensorNetwork): super(). - __init__(ts, **o); return` lets the base copy the extra props without the - fresh-construction defaults clobbering `_plan`. + "_node_tag_id", "_canonical_region", "_symmetry", "_fermionic", + "_physical_sectors")` -- these are copied on `.copy()`/ + every quimb view. The `__init__` copy-branch guard + `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 -- `local_expectation(G, where=[q], max_bond=None, - optimize="auto")` works (emits a cosmetic "not a compressed tree" warning). + 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. - `node_tid(nid)` is a self-healing tid cache kept in `__dict__` (not `_EXTRA_PROPS`) so a copy starts with a fresh, independent cache. - Builders: `from_plan(plan)` (product `|0...0>`), `from_order(order, @@ -70,9 +118,10 @@ directly. 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}`. -- Virtual bond between adjacent nodes `u,v` = `_tb{lo}_{hi}` with `lo the two leaf factors are isometric, the parent is the -new centre). Both new bonds keep their canonical `_tb...` names via `bond_ind=`. -This is the common case in a good layout and avoids the QR hops and double-bond -fusion. Leaves are never directly bonded (both bond only to the parent), so the -correlation flows through the parent blob -- this is exact up to the truncation. +parent, so no threading is needed. Both direct-SVD and Quimb-MPO factors are +absorbed into their leaves, then the two leaves and parent are contracted into +one blob and re-split by two truncating SVDs (`absorb="right"` -> the two leaf +factors are isometric, the parent is the new centre). Both new bonds keep their +canonical `_tb...` names via `bond_ind=`. This is the common case in a good +layout and avoids QR hops and double-bond fusion. Leaves are never directly +bonded (both bond only to the parent), so the correlation flows through the +parent blob -- this is exact up to the truncation. ## Multi-qubit / sub-MPO application (`apply_subtree_operator`) `apply_subtree_operator(op, where, *, max_bond=None, cutoff=None, renormalize=False)` applies a general operator on `k >= 1` qubits in one shot -- a `k`-qubit gate, a multi-site **non-unitary / Kraus** operator, or a whole -**Trotter block**. It is the arbitrary-`k` generalisation of `_apply_2q_siblings` -to the whole spanning subtree, the tree analogue of a sub-MPO applied over a +**Trotter block**. It extends the two-factor path-thread kernel to the whole +spanning subtree: the tree analogue of a sub-MPO applied over a covering range then compressed (quimb's `gate_with_submpo` is `MatrixProductState` -only; the tree base `TensorNetworkGenVector` has no such method). 1. `snodes = _steiner_nodes(leaves)` -- minimal connected subtree spanning the target leaves. -2. Move the centre onto a target leaf if it is not already inside `snodes` - (`_move_center(leaves[0])`, incremental) so the **whole exterior is isometric - toward the subtree** -- each re-split truncation then measures true state - error. -3. Contract copies of all `snodes` tensors into one blob; apply `op` - (`op[o_0..,i_0..]`, output indices first) on the physical legs. -4. `_peel_order(snodes)` -> peel each current subtree-leaf toward a hub with a - truncating SVD (`absorb="right"` -> peeled tensor isometric, norm rides into - the shrinking blob); the last node is the hub / new centre. -5. `renormalize=True` renormalises afterwards (for Kraus/projection). - -**Critical gotcha:** identify each node's kept legs from the **live tensors**, -not from `_bond_name`. Gate application renames bonds off the deterministic -`_tb{lo}_{hi}` scheme (quimb mints `_...` uuids in every split), so -`_subtree_owned_inds` reads the real boundary-bond names via -`qtn.bonds(t_u, t_w)` for exterior neighbours `w`; only the **new** peel bonds -are (re)named with `_bond_name(u, v)` and registered on the hub-side node so a -later peel keeps them. `apply_gate` routes `len(where) >= 3` here; `k == 1`/`k == -2` still take the optimised leaf-absorb / threading paths (but `k == 1` -non-unitary and `k == 2` Kraus can be sent here explicitly). +2. Move the centre onto a target leaf (`_move_center(leaves[0])`, incremental) + so the **whole exterior is isometric toward the subtree**. +3. Factor `op` into an exact tree-MPO on the same Steiner tree by packing each + `(output,input)` physical pair into a dimension-four leg and applying + leaf-to-hub SVDs. +4. Absorb the tree-MPO into copied local state tensors. For each + `_peel_order(snodes)` edge, QR-split the child message while retaining all + 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). + +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 +bonds are private to the temporary tree-MPO. `apply_gate` routes +`len(where) >= 3` here; `k == 1`/`k == 2` still take the optimised +leaf-absorb / threading paths (but `k == 1` non-unitary and `k == 2` Kraus +can be sent here explicitly). + +### Native streamed sub-MPOs + +An explicit `("submpo", mpo, where)` stream event first attempts the native +leaf-to-hub QR-routing sweep. Quimb MPO payloads expose their active site tags, +tensor map, and operator bond indices, so their virtual bonds can be carried +through the TTN without calling `mpo.to_dense()`, then compressed once over the +affected subtree. `estimate_bonds()` uses the product of MPO bond dimensions +crossing a cut as a conservative operator-Schmidt bound. Payloads without that +interface use the dense `to_dense()` fallback and remain subject to +`max_operator_qubits`. ## Readout -- `local_expectation(op, where)`: single-site contracts only the centre tensor - with `op` (canonical). **Multi-site contracts only the minimal Steiner - subtree spanning the target leaves** (`_steiner_nodes` = union of - `node_path(leaves[0], leaf)`): move the centre inside that subtree, rename the - subtree-**internal** bonds in the bra (fresh `rand_uuid`) while keeping the - **boundary** bonds shared between bra and ket. Sharing a boundary bond name - implements the isometric-exterior identity `sum_b ket[..b] bra[..b] = - sum_{b,b'} ket delta(b,b') bra`. In a binary tree the leaf-leaf path passes - only through internal (physical-index-free) nodes, so the subtree's leaves are - exactly the targets. Cost scales with operator spread, not `n`. +- `TreeTensorNetwork.local_expectation(op, where)`: dense/nonfermionic + single-site readout contracts the centre tensor; dense multi-site readout + contracts the minimal Steiner subtree. Native fermionic readout instead + inserts the Symmray operator with `contract=False` and contracts the complete + doubled tree. This preserves graded boundary phases and deliberately avoids + the ordinary isometric-exterior shortcut. Its `max_bond` argument is + compatibility-only: the exact native contraction is not truncated. Dense + readout restores a known canonical centre/region and uses a temporary copy + 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 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)` = `measure` then conditional `X`. `seed` in `__init__` sets `self.rng`; `copy()` - gets a fresh independent RNG. -- `to_dense()` returns the statevector in `k0, k1, ..., k(n-1)` order. -- `bond_report()` / `max_bond()` / `convergence_sweep(...)` are diagnostics. + derives a deterministic child seed for a fresh independent RNG. +- Stream control events follow the MPS tuple/mapping contract for `measure`, + `cap`, `reset`, and `measure_reset`. Pauli measurements support product observables + on distinct qubits, use `+1`/`-1` eigenvalue outcomes, and append + `(pauli, where, outcome, probability)` to `measurements`; reset measurements + are internal and are not recorded. A cap contracts and removes one leaf, + compacts labels above it by default, and absorbs into the unique tree parent; + `stable_labels=True` / `compact_labels=False` preserves caller-facing labels + while storage remains compact. `measure_pauli` returns outcome and Born + probability directly; `project_pauli(..., renormalize=False)` preserves the + 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. +- `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; + the reference resets after control or explicitly non-unitary events. This is + display-only, not a substitute for truncation history. +- `bond_report()` / `estimate_bonds()` / `max_bond()` / + `convergence_sweep(...)` are diagnostics. `estimate_bonds()` is the + non-mutating Eq. (4) dry run: it multiplies operator-Schmidt ranks on each + crossed edge and can conservatively flag a `chi` that will truncate. +- `TreeTensorNetwork.validate()` checks the live tensor/bond structure against + its `TreePlan`; use `check_canonical=True` for the more expensive isometry + check. `TreeOptimizer.preflight(...)` adds `max_bond`, + `max_operator_qubits`, and `max_subtree_nodes` resource limits before replay; + the constructor defaults to finite dense/operator-subtree guards, and + `None` disables either guard. Product-Pauli measurement uses a factorized + parity projector rather than a dense `4**k` matrix. `convergence_sweep` builds the tree once and reuses it across `chi` so the comparison isolates truncation from layout; it reports `fidelity` (only when `2**n <= dense_cap`) and reference-free `max_drift`. +- `record_history=False` disables retained per-edge and per-update history for + long replays. `TreeTensorNetwork` invalidates its canonical-region metadata + and native norm cache after direct Quimb mutators; use + `invalidate_canonical_form()` after raw tensor edits. The optimizer's + state-aware wrappers restore a known centre only for operations proven to + preserve canonicality. +- `truncation_report()` returns the per-edge compression / split history with + before/after dimensions. `track_truncation=True` additionally probes the + untruncated local singular spectrum and records absolute discarded weight and + relative discarded fraction; native Symmray reports use the actually kept + charge-block spectrum. Leave it false on performance runs: the spectrum + probe adds local SVD work per truncation edge. The report's gate-level + `updates` group edge events by support and include cumulative relative loss, + analogous to the MPS infidelity trace. + +## Noisy trajectory replay + +`run_trajectory_shots` and `run_coalesced_trajectory_shots` support +`TreeOptimizer` factories as well as MPS and stabilizer-TN factories. Use them +for trajectory simulation without forming a density matrix: + +- Independent replay samples random-unitary mixtures, Pauli/depolarizing + channels, and state-dependent Kraus channels. For a Kraus event, the runner + applies each branch to a copied TTN, obtains its squared norm, samples the + conditional probability, then applies and normalizes the selected branch on + the live TTN. +- Coalesced replay shares deterministic prefixes and branches exact + mid-circuit `measure`, `reset`, and `measure_reset` events. Tree measurement + probabilities come from `TreeOptimizer.expectation_pauli`; each resulting + leaf remains normalized. +- The runner converts generated dense matrices through the live state backend. + When constructing a direct Tree stream, use matrix-valued gate payloads such + as `pepsy.h()`; textual MPS gate aliases are not normalized by the Tree gate + parser. +- Regression coverage lives in `tests/test_trajectory_noise.py`, including + Tree state-dependent Kraus sampling and coalesced measurement branching. ## Performance / stability (do not regress) @@ -247,7 +385,8 @@ non-unitary and `k == 2` Kraus can be sent here explicitly). rebuilt via `gate_inds_`); a stale entry just misses and is recomputed. Tensor ids are unique and never reused, so this is always safe. - `copy()` shares the immutable `TreePlan`, owns `self.tn.copy()`, resets the - tid cache, and gets a fresh RNG. + tid cache, and derives a deterministic child seed for a fresh independent + RNG. ## Layout (TreeLayoutFinder / TreePlan) @@ -262,7 +401,8 @@ leaves, minimising the geodesic a 2q gate must thread. - `score(plan)` uses the **pure** `pair_weights` (weighted geodesic sum, lower is better) -- keep it separate from the augmented partition similarity. - `report(plan)` compares against a naive `structure="balanced"` index-order - tree (`score_ratio_vs_balanced`). Quality must be `<=` balanced. + tree (`score_ratio_vs_balanced`). The path-objective quality layout should be + `<=` balanced; congestion mode is selected by edge-load cost instead. - `TreePlan.from_order(order, weights=, structure=, max_arity=2, community_frac=, star_frac=)`: `"quality"` = spectral reorder + split; `"balanced"` = split `order` directly (useful in tests to force sibling relationships, e.g. @@ -276,10 +416,10 @@ builder, geometry query, `ascii_tree`, the optimizer's geodesic threading + sibling fast path, and `TreeSampler` loop over `plan.children[nid]` generically. The *only* binary-specific piece was construction. Controls: -- `max_arity` (default `2`) caps children per internal node. `2` reproduces the - strictly-binary tree **exactly** (cut points use `floor(i*L/k)` so the 2-way - case matches the old `mid = L//2` bisection); larger values / `None` give - flatter `k`-ary trees. +- `max_arity` defaults to the candidate search `(2, 3, 4)`; pass scalar `2` to + force a strictly binary tree. `2` reproduces the original binary partition + exactly (cut points use `floor(i*L/k)` so the 2-way case matches the old + `mid = L//2` bisection); larger values / `None` give flatter `k`-ary trees. - `structure="adaptive"` branches per strongly coupled community (`community_frac` of the level's strongest edge) and collapses a near-clique (present-strong-edge fraction `>= star_frac`) into a flat **star** node @@ -288,10 +428,40 @@ The *only* binary-specific piece was construction. Controls: builds an arbitrary hand-specified tree (checks single parent, leaf/qubit coverage of `0..n-1`, reachability, no cycles). - `TreePlan.max_arity()` / `TreePlan.is_binary()` report the shape. +- `TreeLayoutFinder.recommend_arities((2, 3, 4))` compares candidate arities + using path cost or the rank-aware congestion objective and reports local + virtual degree, edge load, and peak bond growth. `report(plan)` also exposes + `is_binary`, `max_arity`, and `arity_histogram`. +- `objective="congestion"` compares interaction, congestion-aware, and + balanced candidates using predicted `log2` operator-Schmidt rank load on + each edge. `objective="path"` remains the backward-compatible default. +- `weight_mode` accepts `count`, `auto`, `angle`, and `operator_schmidt` for + interaction-graph event weighting. `TreeOptimizer` exposes these as + `layout_objective` and `layout_weight_mode`. +- `layered(block_size=...)` is deliberately a fixed construction: it uses the + spectral order but does not score block sizes or use `chi`. Prefer + `recommend_layered((2, 3, 4), chi=...)` for selection; when `chi` is omitted + it inherits the finder's configured value (pass `chi=None` to compare blind). + `weight_mode="operator_schmidt"` is a two-site entangling-strength proxy; + `objective="congestion"` performs the actual rank-per-tree-edge selection. +- `objective="hybrid"` combines normalized weighted path, maximum edge load, + and total edge load using `hybrid_weights=(path, max_edge_load, + total_edge_load)`. It is a fixed-tree replay/runtime-and-bond-cost proxy; + it does not permit layout rewrites during simulation. +- `recommend_layered` / `recommend_arities` accept `refine="greedy"` for a + bounded deterministic pre-simulation adjacent leaf-label swap pass. It keeps + parent/child topology fixed and returns planning metadata per candidate. + `search="nevergrad"` is a separate optional offline search, seeded and + budgeted per candidate, requiring `pepsy[layout]`; it starts from the + spectral/greedy plan and retains only objective improvements. Keep both + opt-in: default layout construction stays fast and reproducible. +- Layout hot paths matter for long streams: `edge_loads` must traverse only the + support's Steiner subtree, and a dense gate's Schmidt-rank cache key must use + wire positions rather than global labels so one repeated CNOT/CZ is factored + once per distinct partition. - All these are plumbed identically through `TreeLayoutFinder`, `TreeOptimizer`, `TreeOptimizer.find_tree_layout`, `TreeOptimizer.convergence_sweep`, and - `TreeTensorNetwork.from_order`. The `TreeLayoutFinder` default is still binary - so existing behaviour is unchanged. + `TreeTensorNetwork.from_order`. ## Gotchas / teaching notes @@ -305,12 +475,13 @@ The *only* binary-specific piece was construction. Controls: ## Roadmap / not yet implemented -- **Sampling** (batched bitstring sampling from the canonical tree) is the next - natural feature and is *not* done. It should reuse the `measure` centre-on-leaf - Born-probability read plus conditional propagation, mirroring `MpsSampler`. -- Stream-wired control events: `measure`/`reset` are explicit methods; the - `[(gate, where)]` stream does not yet carry measurement/reset event markers - (unlike `MpsOptimizer` control events). +- Chain-only MPS execution modes (`svd`, `dmrg`, `mpo`, `swap`, `perm`, `su`, + and `mix`) are not meaningful on arbitrary tree geometry. Native structured + sub-MPO payloads are routed through the Quimb MPO site interface; only + payloads that do not expose that interface use the guarded dense fallback. + Pauli and computational-basis measurement helpers are dense two-level qubit + APIs and intentionally reject native fermionic TTNs, whose local observables + must be supplied through the fermion/Symmray model layer. ## Validation @@ -331,6 +502,14 @@ sphinx-build -W -b html docs docs/_build/html ``` The safety-net tests are `test_tree_matches_statevector` (untruncated fidelity -must stay exactly 1.0) and the multi-site / sibling / measurement regressions. +must stay exactly 1.0), the multi-site / sibling / measurement regressions, and +the state-handoff/backend cases: exact product TTN/MPS mounting, rejected +entangled relayouts, native Torch controls/readout, and mixed-backend rejection. Add a regression test for every new behaviour and prefer `structure="balanced"` plans when a test needs deterministic sibling relationships. + +For noisy trajectory changes, also run: + +```bash +PYTHONDONTWRITEBYTECODE=1 pytest -q tests/test_trajectory_noise.py -k 'not benchmark' +``` diff --git a/.tmp_update_fh_vmc_api.py b/.tmp_update_fh_vmc_api.py new file mode 100644 index 0000000..1af89d3 --- /dev/null +++ b/.tmp_update_fh_vmc_api.py @@ -0,0 +1,154 @@ +import json +import os +import shutil +import stat +from pathlib import Path + + +path = Path("../pepsy_examples/fermi_hubbard/fh_peps.ipynb").resolve() +backup = Path("/tmp/fh_peps_before_vmc_api.ipynb") +tmp_path = path.with_suffix(path.suffix + ".tmp") + +with path.open("r", encoding="utf-8") as f: + notebook = json.load(f) + +new_source = [ + "# VMC energy estimate on the post-imaginary-time PEPS.\n", + "from pepsy.vmc import (\n", + " FermionSiteEncoding,\n", + " TorchPEPSBoundaryAmplitude,\n", + " TorchSquareLattice,\n", + " TorchVMCDriver,\n", + ")\n", + "\n", + "vmc_encoding = FermionSiteEncoding.vmc_torch()\n", + "vmc_device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", + "vmc_sites = tuple(native_peps.sites)\n", + "vmc_seed_config = torch.tensor(\n", + " [[\n", + " vmc_encoding.up\n", + " if half_filled_occupations[site] == (1, 0)\n", + " else vmc_encoding.down\n", + " for site in vmc_sites\n", + " ]],\n", + " dtype=torch.long,\n", + " device=vmc_device,\n", + ")\n", + "n_vmc_walkers = 100\n", + "vmc_configs = vmc_seed_config.repeat(n_vmc_walkers, 1)\n", + "\n", + "vmc_model = TorchPEPSBoundaryAmplitude(\n", + " native_peps,\n", + " contraction=\"boundary\",\n", + " chi=32,\n", + " cutoff=1.0e-10,\n", + " contraction_opts={\"compress_opts\": {\"method\": \"cholesky\", \"absorb\": \"both\"}},\n", + " dtype=torch.complex128,\n", + " device=vmc_device,\n", + ")\n", + "\n", + "# The supplied native terms define the local energy; no t/U guessing.\n", + "vmc_driver = TorchVMCDriver(\n", + " vmc_model,\n", + " TorchSquareLattice(Lx, Ly),\n", + " vmc_configs,\n", + " terms=ham.terms,\n", + " site_order=vmc_sites,\n", + " proposal=\"spinful\",\n", + " encoding=vmc_encoding,\n", + " chunk_size=16,\n", + " generator=torch.Generator(device=vmc_device).manual_seed(7),\n", + ")\n", + "\n", + "# 100 measured samples, with burn-in and extra mixing before measurement.\n", + "with torch.inference_mode():\n", + " vmc_estimate = vmc_driver.estimate_energy(\n", + " burn_in=10,\n", + " n_measurements=1,\n", + " sweeps_between=3,\n", + " progress=True,\n", + " )\n", + "\n", + "vmc_energy_total = float(torch.real(vmc_estimate.energy_mean).cpu())\n", + "vmc_energy_density = vmc_energy_total / num_sites\n", + "vmc_stderr_density = float(vmc_estimate.energy_stderr.cpu()) / num_sites\n", + "print(\n", + " f\"VMC boundary-MPS: E={vmc_energy_total:+.8f}, \"\n", + " f\"E/N={vmc_energy_density:+.8f} +/- {vmc_stderr_density:.3e} \"\n", + " f\"({vmc_estimate.n_samples} samples, \"\n", + " f\"acceptance={vmc_estimate.acceptance_rate:.3f}, \"\n", + " f\"samples/s={vmc_estimate.samples_per_second:.2f})\"\n", + ")\n", + "print(\n", + " f\"deterministic boundary-MPS check: E/N=\"\n", + " f\"{native_energy_density(native_peps, max_bond=32):+.8f}\"\n", + ")\n", +] + +found = False +for cell in notebook["cells"]: + if cell.get("id") == "vmc-energy-estimate-code": + cell["source"] = new_source + found = True + break +if not found: + raise SystemExit("VMC code cell not found; refusing to edit.") + +if not any( + cell.get("id") == "vmc-bp-importance" for cell in notebook["cells"] +): + index = next( + i for i, cell in enumerate(notebook["cells"]) + if cell.get("id") == "vmc-energy-estimate-code" + ) + 1 + notebook["cells"][index:index] = [ + { + "cell_type": "markdown", + "id": "vmc-bp-importance", + "metadata": {}, + "source": [ + "### Optional BP importance proposal\n", + "\n", + "Set `use_bp_importance=True` to draw proposals with BP and " + "measure their amplitudes and local energies with torch.\n", + ], + }, + { + "cell_type": "code", + "execution_count": None, + "id": "vmc-bp-importance-code", + "metadata": {}, + "outputs": [], + "source": [ + "from pepsy.sampling import PepsBpSampler\n", + "\n", + "use_bp_importance = False\n", + "if use_bp_importance:\n", + " bp_importance = vmc_driver.importance_energy_estimate(\n", + " PepsBpSampler(native_peps),\n", + " n_samples=100,\n", + " sample_kwargs={\n", + " \"method\": \"mps\",\n", + " \"chi\": 32,\n", + " \"cutoff\": 1.0e-10,\n", + " },\n", + " progress=True,\n", + " )\n", + " print(\n", + " f\"BP importance: E/N={float(torch.real(bp_importance.energy_mean).cpu()) / num_sites:+.8f} \"\n", + " f\"ESS={float(bp_importance.effective_sample_size.cpu()):.1f}/\"\n", + " f\"{bp_importance.n_valid}\"\n", + " )\n", + ], + }, + ] + +shutil.copy2(path, backup) +mode = stat.S_IMODE(path.stat().st_mode) +with tmp_path.open("w", encoding="utf-8") as f: + json.dump(notebook, f, indent=1, ensure_ascii=False) + f.write("\n") +os.chmod(tmp_path, mode) +os.replace(tmp_path, path) +print(f"updated {path}") +print(f"backup {backup}") diff --git a/.tmp_validate_new_fh_vmc.py b/.tmp_validate_new_fh_vmc.py new file mode 100644 index 0000000..c6272bc --- /dev/null +++ b/.tmp_validate_new_fh_vmc.py @@ -0,0 +1,33 @@ +import nbformat +from nbclient import NotebookClient + + +path = "../pepsy_examples/fermi_hubbard/fh_peps.ipynb" +nb = nbformat.read(path, as_version=4) +for cell in nb.cells: + if cell.cell_type != "code": + continue + cell.source = cell.source.replace( + "tau = [0.2, 0.1, 0.03, 0.01]", + "tau = [0.01]", + ) + cell.source = cell.source.replace( + "steps = [50, 50, 50, 50]", + "steps = [1]", + ) + if cell.get("id") == "vmc-energy-estimate-code": + cell.source = cell.source.replace("burn_in=10", "burn_in=1") + cell.source = cell.source.replace("sweeps_between=3", "sweeps_between=1") + +limit = next( + i for i, cell in enumerate(nb.cells) + if cell.get("id") == "vmc-bp-importance" +) +nb.cells = nb.cells[:limit] +NotebookClient(nb, timeout=600, kernel_name="python3").execute() +for cell in nb.cells: + for output in cell.get("outputs", ()): + if output.get("output_type") == "stream": + text = "".join(output.get("text", [])) + if "VMC boundary-MPS" in text or "deterministic boundary-MPS" in text: + print(text, end="") diff --git a/AGENTS.md b/AGENTS.md index 98f5bad..121bcc3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -103,6 +103,35 @@ At the start of a new task: with persistent layouts. Measurement/reset bookkeeping remains in logical site labels. +### Tree optimizer workflow + +- Read `.github/skills/tree-optimizer/SKILL.md` before changing + `pepsy.optimizers.tree` code, tests, or documentation. +- `TreeLayoutFinder` is circuit-only: derive a `TreePlan` from a gate stream or + supports, then pass a coefficient state separately to `TreeOptimizer`. + Reject a `TreeTensorNetwork` supplied to the finder rather than guessing a + circuit or silently relayouting it. +- An entangled TTN must retain its existing plan. A requested mismatched + `tree=`/`layout=` is an error, never an implicit compression or warning-only + conversion. Bond-one TTNs and bond-one Quimb MPSs may be remounted exactly on + a selected tree; warn only when replacing a product TTN's old geometry. +- Treat backend, dtype, and device as one state invariant: every live TTN + tensor must match. Reject mixed initial states. Callers should provide gates, + operators, sub-MPOs, observables, and cap vectors on the state backend; + preserve the optimizer's one-time warning when an explicit array must be + coerced. Internal Pauli/control tensors must follow the live backend without + producing user-facing transfer warnings. +- Layout is fixed before replay. Never rewrite a live tree during optimization; + `refine="greedy"` and Nevergrad are bounded pre-simulation leaf-label searches + that retain the parent/child topology. +- `TreeOptimizer` supports the shared noisy-trajectory runner. Independent + replay handles random-unitary mixtures, depolarizing channels, and + state-dependent Kraus channels by evaluating branch norms on copied TTNs; + selected branches are normalized before replay continues. Coalesced replay + branches `measure`, `reset`, and `measure_reset` through + `expectation_pauli`. Use matrix-valued gates such as `pepsy.h()` in direct + Tree streams; textual MPS gate aliases are not normalized by the Tree parser. + ### Stabilizer tensor-network workflow - Read `.github/skills/stabilizer-tensor-networks/SKILL.md` and its method/API @@ -201,7 +230,9 @@ At the start of a new task: - `infidelity(...)` returns a dict with `infidelity`, `norm`, `norm_target`, `overlap`, and reused/created boundary handles. - PEPS-like networks should carry lattice tags `X{i}`, `Y{j}`, `I...`; physical outer indices conventionally use `k...` for ket legs and `b...` for bra/operator-output legs. - Gate streams should use canonical bundled entries like `[(gate, where), ...]`. The ambiguous single bundled alias `(gate, where)` is intentionally rejected. -- User-provided gate tensors are not backend-coerced automatically; keep gate tensors and tensor-network arrays on compatible backends. +- Keep user-provided gate tensors and tensor-network arrays on compatible + backends. Where an optimizer intentionally performs a compatibility coercion, + retain its explicit warning/diagnostic rather than silently transferring data. - `OneDMap` supports `snake`, `snake-row-major`, `row-major`, `col-major`, `hilbert`, `hilbert-row-major`, and `diag` modes. PEPO conversion is restricted; check `tests/test_ham.py` before changing mapper behavior. ## Public API Rules @@ -269,6 +300,12 @@ NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/mplconfig PYTHONPYCACHEPREFIX - Do not modify generated notebook outputs unless explicitly requested. - When changing public examples, verify that imports use current namespaces and public API names. - If a notebook or helper still shows an old flat import path, migrate it deliberately instead of copying that pattern into new code. +- For any matplotlib figure (examples, notebooks, benchmark plots), aim for + clear, legible, decent-looking plots (labeled axes, readable fonts, a legend + when there is more than one series, a light grid, a distinct reference curve, + and vector/PNG saves for keep-worthy figures). See `plot_policy.md` at the + repo root for the principles and a recommended house-style starting point — + it is guidance, not a rigid template to copy verbatim. ## Validation diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..d2db2fc --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,61 @@ +# Changelog + +All notable Pepsy changes are documented here. + +Pepsy follows [Semantic Versioning](https://semver.org/): + +- **MAJOR** versions may contain incompatible public API changes. +- **MINOR** versions add backwards-compatible public functionality. +- **PATCH** versions contain backwards-compatible fixes and documentation updates. + +## [Unreleased] + +Changes for the next release should be added here before the version is bumped. + +## [0.3.0] - 2026-07-24 + +This release consolidates the tensor-network API refresh and the new native +TreeOptimizer and symmetric-tensor workflows. + +### Added + +- Native fermionic TreeTensorNetwork evolution, observables, measurements, and + state-versioned norm caching with explicit mutation invalidation. +- TreeOptimizer support for direct and MPO execution paths, including native + subtree and multi-site operator routing. +- Backend-aware symmetric-tensor and Symmray sweep support with regression + coverage for Torch-backed block arrays. +- Public trajectory, stabilizer tensor-network, fermionic, and VMC workflow + APIs with corresponding documentation and examples. + +### Changed + +- TreeOptimizer execution modes are now limited to `auto`, `direct`, and + `mpo`; unsupported legacy mode names fail clearly. +- Dense and native TreeOptimizer measurements use consistent gauge and norm + diagnostics semantics. +- Progress reporting uses a common norm-infidelity proxy, and Symmray + truncation diagnostics use the actual retained block spectra. +- Public imports and package documentation are organized around the current + `pepsy.backends`, `pepsy.boundary`, `pepsy.operators`, `pepsy.optimizers`, + `pepsy.sampling`, `pepsy.solvers`, and `pepsy.tensors` namespaces. + +### Fixed + +- Fermionic local expectations and norm calculations now agree with complete + graded-network reference contractions, including nonzero hopping terms. +- Norm-cache invalidation covers public optimizer mutation and normalization + paths, including constructor normalization. +- Symmray backend conversion, soft MPO bond caps, and blockwise discarded + weight reporting are now handled without dense global-spectrum assumptions. + +### Removed + +- Stale benchmark and example artifacts that no longer represent the current + public API. + +## [0.2.0] - Baseline + +`0.2.0` is the package metadata baseline that preceded this changelog. Earlier +changes were not recorded in a versioned changelog, so historical entries are +intentionally not reconstructed here. diff --git a/PLAN.md b/PLAN.md index a6ea052..17b8392 100644 --- a/PLAN.md +++ b/PLAN.md @@ -9,7 +9,7 @@ PEPS norm contraction and DMRG fitting). It is the "what/when" companion to the conceptual notes in `learning/` (the "why/how"), the session journal in `history/` (the "what happened"), and the published documentation in `docs/`. -Current package version: `0.2.0` (`pepsy.__version__` / `pyproject.toml`). +Current package version: `0.3.0` (`pepsy.__version__` / `pyproject.toml`). Four headline workstreams drive the roadmap: diff --git a/README.md b/README.md index cd6ea98..02fdf5f 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,8 @@ `pepsy` is a Python package for circuit simulation and related DMRG fitting workflows. -Current package version: `0.2.0` (from `pyproject.toml` / `pepsy.__version__`). +Current package version: `0.3.0` (from `pyproject.toml` / `pepsy.__version__`). +See the [changelog](CHANGELOG.md) for release history and versioned changes. ## Package Layout - `src/pepsy/`: installable library code diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py deleted file mode 100644 index f242678..0000000 --- a/benchmarks/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Local benchmark helpers for Pepsy development.""" diff --git a/benchmarks/bp_expansions.py b/benchmarks/bp_expansions.py deleted file mode 100644 index 0e235ab..0000000 --- a/benchmarks/bp_expansions.py +++ /dev/null @@ -1,302 +0,0 @@ -"""Compare BP, loop-series, loop-cluster, and PNE contractions. - -This benchmark helper is intentionally library-like: callers construct the -norm-1 or norm-2 tensor network, then pass it to -:func:`benchmark_bp_expansions`. The returned records are JSON-ready through -``record.as_dict()`` and can be used from notebooks or a command-line driver. -""" - -from __future__ import annotations - -from dataclasses import dataclass -import time -from typing import Any - -import numpy as np - -from pepsy.bp import ( - loop_cluster_expand, - loop_series_expand, - partitioned_expand, - recursive_partitioned_expand, - select_pne_partitions, - weight_pass, -) - - -def _numpy_value(value): - if hasattr(value, "data") and hasattr(value, "inds"): - value = value.data - return np.asarray(value) - - -def _json_value(value): - if value is None: - return None - if isinstance(value, tuple) and len(value) == 2: - return [_json_value(value[0]), _json_value(value[1])] - array = _numpy_value(value) - if array.size != 1: - return { - "shape": list(array.shape), - "norm": float(np.linalg.norm(array.reshape(-1))), - } - scalar = array.reshape(-1)[0] - if np.iscomplexobj(scalar): - scalar = complex(scalar) - return {"real": float(scalar.real), "imag": float(scalar.imag)} - return float(scalar) - - -def _relative_error(estimate, exact): - estimate = _numpy_value(estimate) - exact = _numpy_value(exact) - if estimate.shape != exact.shape: - return float("nan") - denominator = max(float(np.linalg.norm(exact.reshape(-1))), 1e-300) - return float(np.linalg.norm((estimate - exact).reshape(-1)) / denominator) - - -@dataclass(frozen=True) -class ExpansionBenchmarkRecord: - """One JSON-ready expansion benchmark result.""" - - method: str - estimate: Any - exact: Any - relative_error: float - wall_seconds: float - num_terms: int | None - residue_norm: float | None - bp_converged: bool | None - error: str | None = None - - def as_dict(self): - return { - "method": self.method, - "estimate": _json_value(self.estimate), - "exact": _json_value(self.exact), - "relative_error": self.relative_error, - "wall_seconds": self.wall_seconds, - "num_terms": self.num_terms, - "residue_norm": self.residue_norm, - "bp_converged": self.bp_converged, - "error": self.error, - } - - -def _run_record(method, function, exact): - start = time.perf_counter() - try: - result = function() - except Exception as exc: # benchmark tables should retain failed methods - return ExpansionBenchmarkRecord( - method=method, - estimate=None, - exact=exact, - relative_error=float("nan"), - wall_seconds=time.perf_counter() - start, - num_terms=None, - residue_norm=None, - bp_converged=None, - error=f"{type(exc).__name__}: {exc}", - ) - - if method == "loop_cluster": - num_terms = ( - len(result.region_counts) - if result.region_counts is not None - else None - ) - else: - num_terms = len(result.terms) - return ExpansionBenchmarkRecord( - method=method, - estimate=result.estimate, - exact=exact, - relative_error=_relative_error(result.estimate, exact), - wall_seconds=time.perf_counter() - start, - num_terms=num_terms, - residue_norm=getattr(result, "residue_norm", None), - bp_converged=result.bp_converged, - ) - - -def benchmark_bp_expansions( - tn, - *, - norm: str = "2norm", - exact=None, - loop_cutoff: int = 4, - cluster_cutoff: int = 4, - partition_inds=None, - recursive_levels=None, - partition_form: str = "linear", - include_residue: bool = False, - max_iterations: int = 1000, - tol: float = 5e-6, - optimize: str = "auto-hq", - contract_opts: dict[str, Any] | None = None, - auto_select: bool = False, - max_partitions: int = 1, - candidate_inds=None, - partition_opts: dict[str, Any] | None = None, - weight_passing_rank: int | None = None, - weight_passing_opts: dict[str, Any] | None = None, -): - """Benchmark the available BP corrections on one scalar contraction. - - Parameters - ---------- - tn - A closed scalar network for ``norm="1norm"`` or a PEPS-like network - for ``norm="2norm"``. - exact - Optional exact reference. If omitted, ``tn.contract()`` is used. - partition_inds - PNE partition indices. If omitted, the first internal index is used, - unless ``auto_select=True``. - auto_select - Rank candidate indices by the one-index PNE residue heuristic. - max_partitions, candidate_inds, partition_opts - Controls for ``auto_select``. - recursive_levels - Optional fixed recursive PNE schedule, e.g. - ``(("e0", "e1"), ("e2",))``. - weight_passing_rank - If set, add a higher-rank PNE record produced by Appendix-C weight - passing. This route currently targets closed pairwise norm-1 inputs. - weight_passing_opts - Options forwarded to :func:`pepsy.bp.weight_pass`. - """ - if exact is None: - exact = tn.contract(optimize=optimize, **(contract_opts or {})) - contract_opts = {} if contract_opts is None else dict(contract_opts) - common = { - "norm": norm, - "max_iterations": max_iterations, - "tol": tol, - "optimize": optimize, - "contract_opts": contract_opts, - } - internal = tuple(sorted(tn.inner_inds(), key=repr)) - if partition_inds is None and auto_select: - selection_options = { - "max_iterations": max_iterations, - "tol": tol, - } - selection_options.update(partition_opts or {}) - selection = select_pne_partitions( - tn, - norm=norm, - max_partitions=max_partitions, - candidate_inds=candidate_inds, - partition_opts=selection_options, - ) - partition_inds = selection.indices - if partition_inds is None: - if not internal: - raise ValueError("no internal indices available for the PNE benchmark") - partition_inds = (internal[0],) - - records = [] - records.append( - _run_record( - "bp", - lambda: loop_series_expand( - tn.copy(), - gloops=0, - multi_excitation_correct=False, - **common, - ), - exact, - ) - ) - records.append( - _run_record( - "loop_series", - lambda: loop_series_expand( - tn.copy(), gloops=loop_cutoff, **common - ), - exact, - ) - ) - records.append( - _run_record( - "loop_cluster", - lambda: loop_cluster_expand( - tn.copy(), gloops=cluster_cutoff, **common - ), - exact, - ) - ) - records.append( - _run_record( - "pne", - lambda: partitioned_expand( - tn.copy(), - partition_inds=partition_inds, - form=partition_form, - include_residue=include_residue, - **common, - ), - exact, - ) - ) - if recursive_levels is not None: - records.append( - _run_record( - "pne_recursive", - lambda: recursive_partitioned_expand( - tn.copy(), - recursive_levels, - form=partition_form, - include_residue=include_residue, - **common, - ), - exact, - ) - ) - if weight_passing_rank is not None: - weight_options = {} if weight_passing_opts is None else dict(weight_passing_opts) - records.append( - _run_record( - "pne_weight_pass", - lambda: _weight_passing_expansion( - tn, - norm=norm, - partition_inds=partition_inds, - rank=weight_passing_rank, - include_residue=include_residue, - optimize=optimize, - contract_opts=contract_opts, - weight_options=weight_options, - ), - exact, - ) - ) - return tuple(records) - - -def _weight_passing_expansion( - tn, - *, - norm, - partition_inds, - rank, - include_residue, - optimize, - contract_opts, - weight_options, -): - weight_result = weight_pass(tn.copy(), **weight_options) - return partitioned_expand( - weight_result.network, - partition_inds=partition_inds, - norm=norm, - projectors=weight_result.projectors(rank=rank), - run_bp=False, - include_residue=include_residue, - optimize=optimize, - contract_opts=contract_opts, - ) diff --git a/benchmarks/coalesced_trajectory_scaling.py b/benchmarks/coalesced_trajectory_scaling.py deleted file mode 100644 index faf4f9a..0000000 --- a/benchmarks/coalesced_trajectory_scaling.py +++ /dev/null @@ -1,355 +0,0 @@ -"""Benchmark independent versus count-coalesced MPS Pauli trajectories. - -The benchmark sweeps fault probability, circuit depth, and shot count. It -compares :func:`pepsy.run_noisy_shots`, which retains one optimizer per shot, -against :func:`pepsy.run_coalesced_noisy_shots`, which shares equal no-error -prefixes and keeps a count per final branch. - -Examples --------- - source ~/envs/py312/bin/activate - python benchmarks/coalesced_trajectory_scaling.py - python benchmarks/coalesced_trajectory_scaling.py \ - --p-list 0,1e-4,1e-3,1e-2 --depth-list 2,4,8 --shots-list 32,128,512 - python benchmarks/coalesced_trajectory_scaling.py --backend torch --device cuda - -The GPU option accelerates the ordinary MPS gate and terminal tensor work. It -does not make divergent noise branches a uniform ``vmap`` batch; the reported -``coalesced_states`` makes the structural sharing visible directly. -""" - -from __future__ import annotations - -import argparse -import json -import os -import time -from pathlib import Path - -os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") - -import numpy as np - - -def _parse_csv(value, cast): - """Parse a nonempty comma-separated command-line list.""" - values = [cast(item) for item in str(value).split(",") if item.strip()] - if not values: - raise ValueError("expected at least one comma-separated value.") - return values - - -def brickwall_gate_stream( - n: int, depth: int, *, backend: str = "numpy", device: str | None = None -): - """Return a small entangling MPS gate stream with deterministic structure.""" - import quimb as qu - - backend = str(backend).lower() - - def on_backend(matrix): - if backend == "numpy": - return matrix - if backend in {"torch", "cuda"}: - import torch - - target = "cuda" if backend == "cuda" and device is None else (device or "cpu") - return torch.as_tensor( - np.array(matrix, copy=True), dtype=torch.complex128, device=target - ) - raise ValueError("backend must be 'numpy', 'torch', or 'cuda'.") - - stream = [] - for layer in range(int(depth)): - for site in range(int(n)): - if (layer + site) % 2 == 0: - stream.append((on_backend(qu.hadamard()), site)) - for left in range(layer % 2, int(n) - 1, 2): - stream.append((on_backend(qu.CNOT()), (left, left + 1))) - return stream - - -def _state_factory(n: int, chi: int, backend: str, device: str | None): - """Build a repeatable ordinary-MPS factory, optionally with Torch tensors.""" - import quimb.tensor as qtn - from pepsy import MpsOptimizer - - state = qtn.MPS_computational_state("0" * int(n), dtype="complex128") - backend = str(backend).lower() - if backend == "numpy": - pass - elif backend in {"torch", "cuda"}: - import torch - - device = "cuda" if backend == "cuda" and device is None else (device or "cpu") - if str(device).startswith("cuda") and not torch.cuda.is_available(): - raise RuntimeError("Torch CUDA was requested but no CUDA device is available.") - state.apply_to_arrays( - lambda array: torch.as_tensor( - array, dtype=torch.complex128, device=device - ) - ) - else: - raise ValueError("backend must be 'numpy', 'torch', or 'cuda'.") - - return lambda: MpsOptimizer(state, chi=int(chi), mode="mpo") - - -def noisy_target_count(n: int, depth: int) -> int: - """Count post-gate Pauli-channel targets in :func:`brickwall_gate_stream`.""" - targets = 0 - for layer in range(int(depth)): - targets += sum((layer + site) % 2 == 0 for site in range(int(n))) - targets += 2 * len(range(layer % 2, int(n) - 1, 2)) - return targets - - -def _median_elapsed(fn, repeats: int, *, synchronize=None): - """Return the median wall time and final value of a repeated callable.""" - elapsed = [] - value = None - for _ in range(int(repeats)): - if synchronize is not None: - synchronize() - start = time.perf_counter() - value = fn() - if synchronize is not None: - synchronize() - elapsed.append(time.perf_counter() - start) - return float(np.median(elapsed)), value - - -def run_case( - *, - n: int, - depth: int, - shots: int, - probability: float, - chi: int, - seed: int, - repeats: int = 1, - backend: str = "numpy", - device: str | None = None, -): - """Time one independent/coalesced pair and return JSON-ready metrics.""" - from pepsy import ( - PauliErrorModel, - run_coalesced_noisy_shots, - run_noisy_shots, - ) - - n = int(n) - depth = int(depth) - shots = int(shots) - if n < 2 or depth < 1 or shots < 1: - raise ValueError("n >= 2, depth >= 1, and shots >= 1 are required.") - stream = brickwall_gate_stream(n, depth, backend=backend, device=device) - model = PauliErrorModel.depolarizing(float(probability)) - noisy_targets = noisy_target_count(n, depth) - factory = _state_factory(n, chi, backend, device) - run_kwargs = {"progbar": False, "fidelity_samples": 0} - synchronize = None - torch_device = "cuda" if backend == "cuda" and device is None else device - if backend in {"torch", "cuda"} and str(torch_device or "").startswith("cuda"): - import torch - - synchronize = lambda: torch.cuda.synchronize(torch_device) - - baseline_s, baseline = _median_elapsed( - lambda: run_noisy_shots( - factory, stream, model, shots, seed=seed, run_kwargs=run_kwargs - ), - repeats, - synchronize=synchronize, - ) - coalesced_s, coalesced = _median_elapsed( - lambda: run_coalesced_noisy_shots( - factory, stream, model, shots, seed=seed, run_kwargs=run_kwargs - ), - repeats, - synchronize=synchronize, - ) - branches = coalesced.branches - return { - "n": n, - "depth": depth, - "gates": len(stream), - "noisy_targets": noisy_targets, - "expected_faults": noisy_targets * float(probability), - "shots": shots, - "probability": float(probability), - "chi": int(chi), - "backend": str(backend), - "device": device, - "independent_s": baseline_s, - "coalesced_s": coalesced_s, - "speedup": baseline_s / coalesced_s if coalesced_s > 0.0 else None, - "independent_states": len(baseline.optimizers), - "coalesced_states": branches, - "state_reduction": shots / branches, - "represented_shots": coalesced.shots, - } - - -def _case_key(row): - """Return the resume key for one benchmark case.""" - return ( - row["n"], - row["depth"], - row["shots"], - row["probability"], - row["chi"], - row["backend"], - row["device"], - ) - - -def _with_derived_metrics(row): - """Fill report-only fields omitted by checkpoints from older scripts.""" - row = dict(row) - noisy_targets = noisy_target_count(row["n"], row["depth"]) - row.setdefault("noisy_targets", noisy_targets) - row.setdefault("expected_faults", noisy_targets * row["probability"]) - return row - - -def _report_config(args, probabilities, depths, shots_list): - """Build the JSON-stable benchmark configuration section.""" - return { - "n": int(args.n), - "p_list": probabilities, - "depth_list": depths, - "shots_list": shots_list, - "chi": int(args.chi), - "seed": int(args.seed), - "repeats": int(args.repeats), - "backend": args.backend, - "device": args.device, - } - - -def run(args, *, existing_results=(), checkpoint=None): - """Run the requested probability/depth/shot sweep.""" - probabilities = _parse_csv(args.p_list, float) - depths = _parse_csv(args.depth_list, int) - shots_list = _parse_csv(args.shots_list, int) - report = { - "config": _report_config(args, probabilities, depths, shots_list), - "results": [], - } - existing = {_case_key(row): row for row in existing_results} - for probability in probabilities: - for depth in depths: - for shots in shots_list: - key = ( - int(args.n), - int(depth), - int(shots), - float(probability), - int(args.chi), - str(args.backend), - args.device, - ) - if key in existing: - report["results"].append(_with_derived_metrics(existing[key])) - continue - row = run_case( - n=args.n, - depth=depth, - shots=shots, - probability=probability, - chi=args.chi, - seed=args.seed, - repeats=args.repeats, - backend=args.backend, - device=args.device, - ) - report["results"].append(row) - if getattr(args, "progress", False): - print( - "completed " - f"p={probability:.2e}, depth={depth}, shots={shots}: " - f"speedup={row['speedup']:.2f}, " - f"leaves={row['coalesced_states']}", - flush=True, - ) - if checkpoint is not None: - checkpoint(report) - return report - - -def _print_table(report): - """Print a compact comparison table alongside optional JSON output.""" - header = ( - f"{'p':>9} {'lambda':>8} {'depth':>5} {'shots':>7} {'ind[s]':>10} {'coal[s]':>10} " - f"{'speedup':>8} {'states':>12} {'reduction':>10}" - ) - print(header) - print("-" * len(header)) - for row in report["results"]: - print( - f"{row['probability']:>9.2e} {row['expected_faults']:>8.2e} " - f"{row['depth']:>5} {row['shots']:>7} " - f"{row['independent_s']:>10.4f} {row['coalesced_s']:>10.4f} " - f"{row['speedup']:>8.2f} " - f"{row['coalesced_states']:>5}/{row['independent_states']:<6} " - f"{row['state_reduction']:>10.2f}" - ) - - -def build_arg_parser(): - """Build the benchmark command-line parser.""" - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--n", type=int, default=8, help="number of qubits") - parser.add_argument("--depth-list", default="2,4,8") - parser.add_argument("--shots-list", default="32,128,512") - parser.add_argument("--p-list", default="0,1e-4,1e-3,1e-2") - parser.add_argument("--chi", type=int, default=16) - parser.add_argument("--seed", type=int, default=20260715) - parser.add_argument("--repeats", type=int, default=1) - parser.add_argument("--backend", choices=("numpy", "torch", "cuda"), default="numpy") - parser.add_argument("--device", default=None, help="Torch device, e.g. cuda:0") - parser.add_argument("--json", default=None, help="optional path for JSON report") - parser.add_argument( - "--progress", - action="store_true", - help="print each completed probability/depth/shot case immediately", - ) - parser.add_argument( - "--resume", - action="store_true", - help="reuse completed matching cases from --json and checkpoint every case", - ) - return parser - - -def _write_json(path, report): - """Atomically checkpoint a JSON-ready benchmark report.""" - path = Path(path) - temporary = path.with_suffix(path.suffix + ".tmp") - with temporary.open("w", encoding="utf-8") as handle: - json.dump(report, handle, indent=2) - temporary.replace(path) - - -def main(): - """Run the command-line benchmark.""" - args = build_arg_parser().parse_args() - if args.resume and not args.json: - raise SystemExit("--resume requires --json PATH.") - existing_results = () - if args.resume and Path(args.json).is_file(): - with open(args.json, encoding="utf-8") as handle: - existing_results = json.load(handle).get("results", ()) - checkpoint = ( - (lambda report: _write_json(args.json, report)) if args.json else None - ) - report = run(args, existing_results=existing_results, checkpoint=checkpoint) - _print_table(report) - if args.json: - _write_json(args.json, report) - print(f"\nwrote {args.json}") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/mps_symmray_sampling.py b/benchmarks/mps_symmray_sampling.py new file mode 100644 index 0000000..d8e7b06 --- /dev/null +++ b/benchmarks/mps_symmray_sampling.py @@ -0,0 +1,305 @@ +"""Benchmark sparse Symmray MPS sampling against dense alternatives. + +Examples +-------- +Run all sampler variants for a charge-restricted bosonic U1 batch: + +``python benchmarks/mps_symmray_sampling.py --symmetry U1 --samples 4096`` + +Run only the sparse path while comparing its prefix policies: + +``python benchmarks/mps_symmray_sampling.py --fermionic --variants symmray --strategy prefix`` +``python benchmarks/mps_symmray_sampling.py --fermionic --variants symmray --strategy serial`` + +Each variant runs in a fresh process. ``peak_rss_mib`` captures its complete +setup high-water mark, while ``resident_rss_mib`` is measured after temporary +conversion data is released. This avoids inheriting allocations from another +route while keeping conversion cost visible. +Fermionic Z2/Z2Z2 inputs are reported as sparse-only: directly expanding their +graded virtual legs into ordinary dense MPS tensors does not preserve the state. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import multiprocessing as mp +import os +import sys +import traceback +from time import perf_counter + +import numpy as np +import quimb.tensor as qtn + +from pepsy.sampling import MpsSampler +from pepsy.tensors import Fermion, SymMPS, site_charge_from_occupations + + +def _nonfermionic_setup(symmetry, length): + if symmetry == "U1": + return {0: 1, 1: 1, 2: 1}, (1,) * length + if symmetry == "Z2": + return {0: 1, 1: 1}, tuple(site % 2 for site in range(length)) + sectors = {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1} + charges = tuple( + (1, 0) if site % 2 == 0 else (0, 1) + for site in range(length) + ) + return sectors, charges + + +def _build_state(args): + if args.fermionic: + fermion = Fermion(spinful=True, symmetry=args.symmetry) + sectors = fermion.physical_sectors + occupations = fermion.half_filled_occupations(args.length) + else: + sectors, occupations = _nonfermionic_setup(args.symmetry, args.length) + return SymMPS.random( + args.length, + symmetry=args.symmetry, + fermionic=args.fermionic, + bond_dim=args.bond_dim, + phys_dim=sectors, + site_charge=site_charge_from_occupations(occupations), + seed=args.seed, + dtype="complex128", + ).mps + + +def _dense_mps_from_symmray(psi): + """Expand local sparse blocks into a dense quimb MPS baseline.""" + arrays = [] + for site in range(psi.L): + tensor = psi[site] + site_ind = psi.site_ind(site) + left_ind = psi.bond(site - 1, site) if site else None + right_ind = psi.bond(site, site + 1) if site < psi.L - 1 else None + if left_ind is None: + data = tensor.transpose(site_ind, right_ind).data + data = data.to_dense().reshape((1, data.shape[0], data.shape[1])) + elif right_ind is None: + data = tensor.transpose(left_ind, site_ind).data + data = data.to_dense().reshape((data.shape[0], data.shape[1], 1)) + else: + data = tensor.transpose(left_ind, site_ind, right_ind).data.to_dense() + arrays.append(np.asarray(data)) + + # quimb's ``lrp`` convention stores the physical leg last, whereas the + # sampler's dense internal convention is ``lpr``. + return qtn.MatrixProductState( + [ + arrays[0][0].T, + *(array.transpose(0, 2, 1) for array in arrays[1:-1]), + arrays[-1][:, :, 0], + ], + shape="lrp", + ) + + +def _peak_rss_mib(): + """Return the process high-water RSS in MiB on Unix-like systems.""" + try: + import resource # pylint: disable=import-outside-toplevel + except ImportError: # pragma: no cover - Windows fallback + return None + + rss = float(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) + if sys.platform == "darwin": # macOS reports bytes rather than KiB. + return rss / (1024.0**2) + return rss / 1024.0 + + +def _resident_rss_mib(): + """Return current RSS in MiB when the platform exposes it.""" + if not sys.platform.startswith("linux"): + return None + try: + with open("/proc/self/statm", encoding="utf-8") as statm: + resident_pages = int(statm.read().split()[1]) + return resident_pages * os.sysconf("SC_PAGE_SIZE") / (1024.0**2) + except (IndexError, OSError, ValueError): # pragma: no cover - procfs edge case + return None + + +def _dense_baseline_unavailable_reason(args): + """Return why an ordinary dense MPS would be an invalid comparator.""" + if args.fermionic and args.symmetry in {"Z2", "Z2Z2"}: + return ( + "Fermionic Z2 graded virtual-leg phases cannot be represented by " + "locally calling to_dense() and handing the result to an ordinary " + "quimb MPS. Only MpsSampler's Symmray route is valid." + ) + return None + + +def _benchmark_variant(config, variant): + """Build, warm, and time one sampler variant in its own process.""" + args = argparse.Namespace(**config) + unavailable_reason = _dense_baseline_unavailable_reason(args) + if variant != "symmray" and unavailable_reason is not None: + return { + "status": "unsupported", + "reason": unavailable_reason, + } + + started = perf_counter() + sparse_psi = _build_state(args) + if variant == "symmray": + psi = sparse_psi + sampler = MpsSampler( + psi, + backend="symmray", + prefix_strategy=args.strategy, + max_prefix_groups=args.max_prefix_groups, + ) + else: + psi = _dense_mps_from_symmray(sparse_psi) + del sparse_psi + sampler = MpsSampler(psi, backend=variant) + setup_seconds = perf_counter() - started + + warmup_samples = min(args.samples, 32) + configs, probabilities = sampler.sample_arrays(warmup_samples, seed=args.seed) + if not np.allclose( + sampler.probabilities(configs), + probabilities, + rtol=1e-10, + atol=1e-12, + ): + raise RuntimeError(f"{variant} sampler returned inconsistent probabilities.") + + elapsed = [] + stats = [] + for repeat in range(args.repeats): + start = perf_counter() + sampler.sample_arrays(args.samples, seed=args.seed + repeat) + elapsed.append(perf_counter() - start) + if variant == "symmray": + stats.append(sampler.symmray_sampling_stats) + + median_seconds = float(np.median(elapsed)) + gc.collect() + resident_rss_mib = _resident_rss_mib() + peak_rss_mib = _peak_rss_mib() + if resident_rss_mib is not None and peak_rss_mib is not None: + # Some container runtimes report ``ru_maxrss`` at a coarser cadence + # than procfs; a peak cannot be lower than the current residency. + peak_rss_mib = max(peak_rss_mib, resident_rss_mib) + result = { + "status": "ok", + "backend": sampler.resolved_backend, + "setup_seconds": float(setup_seconds), + "seconds_median": median_seconds, + "samples_per_second": float(args.samples / median_seconds), + "peak_rss_mib": peak_rss_mib, + "resident_rss_mib": resident_rss_mib, + } + if stats: + result.update( + { + "conditional_evaluations_median": float( + np.median([item["conditional_evaluations"] for item in stats]) + ), + "candidate_contractions_median": float( + np.median([item["candidate_contractions"] for item in stats]) + ), + "charge_pruned_branches_median": float( + np.median([item["charge_pruned_branches"] for item in stats]) + ), + "max_active_prefix_groups": max( + item["max_active_prefix_groups"] for item in stats + ), + "serial_fallback": any(item["serial_fallback"] for item in stats), + "adaptive_serial_fallback": any( + item["adaptive_serial_fallback"] for item in stats + ), + } + ) + return result + + +def _benchmark_worker(config, variant, result_queue): + """Send a worker result or traceback back to the benchmark parent.""" + try: + result_queue.put(("ok", _benchmark_variant(config, variant))) + except BaseException: # pragma: no cover - propagated to the CLI parent + result_queue.put(("error", traceback.format_exc())) + + +def _run_isolated_variant(config, variant): + """Run one variant separately so high-water RSS remains comparable.""" + context = mp.get_context("spawn") + result_queue = context.Queue() + process = context.Process( + target=_benchmark_worker, + args=(config, variant, result_queue), + ) + process.start() + process.join() + if process.exitcode: + raise RuntimeError( + f"{variant} benchmark worker exited with status {process.exitcode}." + ) + status, result = result_queue.get() + if status != "ok": + raise RuntimeError(f"{variant} benchmark worker failed:\n{result}") + return result + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--symmetry", + choices=("Z2", "U1", "U1U1", "Z2Z2"), + default="U1", + ) + parser.add_argument("--fermionic", action="store_true") + parser.add_argument("--length", type=int, default=12) + parser.add_argument("--bond-dim", type=int, default=8) + parser.add_argument("--samples", type=int, default=4096) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument( + "--strategy", + choices=("auto", "prefix", "serial"), + default="auto", + ) + parser.add_argument("--max-prefix-groups", type=int, default=256) + parser.add_argument( + "--variants", + nargs="+", + choices=("symmray", "native", "quimb"), + default=("symmray", "native", "quimb"), + help=( + "Sampler implementations to measure. Each runs in an isolated " + "process for comparable peak RSS." + ), + ) + args = parser.parse_args(argv) + if args.length < 2 or args.bond_dim < 1 or args.samples < 1 or args.repeats < 1: + parser.error( + "length >= 2, bond-dim >= 1, samples >= 1, and repeats >= 1 " + "are required" + ) + + payload = { + "symmetry": args.symmetry, + "fermionic": args.fermionic, + "length": args.length, + "bond_dim": args.bond_dim, + "samples": args.samples, + "strategy": args.strategy, + "max_prefix_groups": args.max_prefix_groups, + "variants": { + variant: _run_isolated_variant(vars(args), variant) + for variant in args.variants + }, + } + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/stabilizer_tn_magic_scaling.py b/benchmarks/stabilizer_tn_magic_scaling.py deleted file mode 100644 index 6650848..0000000 --- a/benchmarks/stabilizer_tn_magic_scaling.py +++ /dev/null @@ -1,289 +0,0 @@ -"""Benchmark the Stabilizer Tensor Network (STN) magic-vs-chi scaling. - -Compares three STN treatments of T-doped Clifford circuits: - -* ``direct`` applies every T rotation to ``|nu>``; -* ``immediate`` teleports every T through one recycled ancilla and immediately - projects it; -* ``deferred`` implements the MAST protocol: one fresh magic ancilla per T, - then end-of-circuit basis-updating projections in a chosen order. - -The report separates circuit replay and final-projection time for deferred -MAST, and records the peak coefficient-MPS bond rather than only its final -value. Clifford gates are free tableau updates; non-Clifford resource handling -is the quantity being compared. - -For each ``N`` the script uses the same deterministic random circuit in every -mode, records peak/final ``|nu>`` bond, pseudo-stabilizer rank (for small -systems), total wall time, and deferred replay/projection times, then emits JSON -and a human-readable table. Use ``--no-exact-cooling`` to isolate injection from -the constructive exact-cooling pre-check. - -Examples --------- - source ~/envs/py312/bin/activate - python benchmarks/stabilizer_tn_magic_scaling.py - python benchmarks/stabilizer_tn_magic_scaling.py --n-list 8,16,32,64 --t 4 - python benchmarks/stabilizer_tn_magic_scaling.py --t 3 --backend torch -""" - -from __future__ import annotations - -import argparse -import json -import os -import time - -os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") - -import numpy as np - - -_ONE_Q_CLIFFORD = ("h", "s", "sdg", "x", "y", "z") - - -def random_clifford_t_circuit(n, t, depth, seed): - """Return a deterministic Clifford gate stream with ``t`` interspersed T gates. - - Each of ``depth`` layers applies a random single-qubit Clifford per qubit and - a brick-wall of CNOTs; ``t`` ``("t", q)`` gates are placed at distinct - (layer, qubit) slots. The T-count is exactly ``t`` regardless of ``n``. - """ - rng = np.random.default_rng(seed) - t = min(int(t), depth * n) - slots = set() - while len(slots) < t: - slots.add((int(rng.integers(depth)), int(rng.integers(n)))) - stream = [] - for layer in range(depth): - for q in range(n): - stream.append((_ONE_Q_CLIFFORD[int(rng.integers(len(_ONE_Q_CLIFFORD)))], q)) - if (layer, q) in slots: - stream.append(("t", q)) - for a in range(layer % 2, n - 1, 2): # brick-wall CNOTs - stream.append(("cnot", a, a + 1)) - return stream - - -def _resolve_backend(name): - if name in (None, "", "numpy"): - return None - import pepsy as py - - if name == "torch": - import torch - - return py.backend_torch(dtype=torch.complex128, device="cpu") - if name == "cupy": - return py.backend_cupy() - if name == "jax": - return py.backend_jax(dtype="complex128") - raise ValueError(f"unknown backend {name!r} (use numpy/torch/cupy/jax).") - - -def run_case( - n, - t, - depth, - seed, - chi, - mode, - to_backend, - rank_max_n, - exact_cooling=True, - deferred_projection_order="middle_out", -): - """Run one (n, mode) case and return a metrics dict.""" - from pepsy.optimizers import MpsStabOptimizer - - stream = random_clifford_t_circuit(n, t, depth, seed) - n_two_qubit = sum(1 for e in stream if e[0] == "cnot") - - mode = str(mode).strip().lower() - common = { - "chi": chi, - "to_backend": to_backend, - "exact_cooling": exact_cooling, - } - start = time.perf_counter() - projection_report = None - if mode in ("immediate", "injection"): - sim = MpsStabOptimizer.with_injection( - n, stream, n_ancilla=1, **common - ) - projection_report = sim.last_immediate_injection_report - mode = "immediate" - elif mode == "deferred": - sim = MpsStabOptimizer.with_deferred_injection( - n, - stream, - projection_order=deferred_projection_order, - **common, - ) - projection_report = sim.last_deferred_injection_report - elif mode == "direct": - sim = MpsStabOptimizer(n, **common).apply(stream) - else: - raise ValueError( - f"unknown mode {mode!r}; use direct, immediate, deferred, or injection." - ) - elapsed = time.perf_counter() - start - - rank = sim.pseudo_stabilizer_rank() if sim.n <= rank_max_n else None - return { - "n": int(n), - "t": int(t), - "mode": mode, - "gates": len(stream), - "two_qubit_gates": int(n_two_qubit), - "peak_nu_bond": int(max(sim.bond_history)), - # Backward-compatible name retained for callers of the original benchmark. - "max_nu_bond": int(max(sim.bond_history)), - "final_nu_bond": int(sim.state.max_bond()), - "bond_bound_2_to_t": int(2 ** t), - "pseudo_stabilizer_rank": None if rank is None else int(rank), - "replay_elapsed_s": float( - elapsed - if mode == "direct" - else ( - projection_report["replay_elapsed_s"] - if mode == "deferred" - else elapsed - projection_report["projection_elapsed_s"] - ) - ), - "projection_elapsed_s": float( - 0.0 if mode == "direct" else projection_report["projection_elapsed_s"] - ), - "pre_projection_peak_bond": ( - None if mode != "deferred" else projection_report["pre_projection_peak_bond"] - ), - "projection_peak_bond": ( - None if mode == "direct" else projection_report["projection_peak_bond"] - ), - "elapsed_s": float(elapsed), - } - - -def run(args): - n_list = [int(x) for x in str(args.n_list).split(",") if x.strip()] - modes = [m.strip() for m in str(args.modes).split(",") if m.strip()] - to_backend = _resolve_backend(args.backend) - results = [] - for n in n_list: - for mode in modes: - res = run_case( - n=n, - t=args.t, - depth=args.depth, - seed=args.seed, - chi=args.chi, - mode=mode, - to_backend=to_backend, - rank_max_n=args.rank_max_n, - exact_cooling=args.exact_cooling, - deferred_projection_order=args.deferred_projection_order, - ) - results.append(res) - return { - "config": { - "n_list": n_list, - "t": int(args.t), - "depth": int(args.depth), - "seed": int(args.seed), - "chi": args.chi, - "modes": modes, - "backend": args.backend or "numpy", - "exact_cooling": bool(args.exact_cooling), - "deferred_projection_order": args.deferred_projection_order, - }, - "results": results, - } - - -def _print_table(report): - cfg = report["config"] - print( - f"# STN magic-vs-chi scaling (t={cfg['t']} T-gates, depth={cfg['depth']}, " - f"chi={cfg['chi']}, backend={cfg['backend']})" - ) - print( - f"# exact_cooling={cfg['exact_cooling']}; deferred order=" - f"{cfg['deferred_projection_order']}\n" - ) - header = ( - f"{'mode':>10} {'N':>5} {'gates':>7} {'2q':>5} {'peak':>6} " - f"{'final':>6} {'rank':>6} {'replay[s]':>10} {'proj-bond':>10} {'proj[s]':>9} {'total[s]':>9}" - ) - print(header) - print("-" * len(header)) - for r in report["results"]: - rank = "-" if r["pseudo_stabilizer_rank"] is None else r["pseudo_stabilizer_rank"] - replay = "-" if r["replay_elapsed_s"] is None else f"{r['replay_elapsed_s']:.3f}" - projection = ( - f"{r['projection_elapsed_s']:.3f}" - ) - projection_bond = ( - "-" if r["projection_peak_bond"] is None else r["projection_peak_bond"] - ) - print( - f"{r['mode']:>10} {r['n']:>5} {r['gates']:>7} {r['two_qubit_gates']:>5} " - f"{r['peak_nu_bond']:>6} {r['final_nu_bond']:>6} {str(rank):>6} " - f"{replay:>10} {str(projection_bond):>10} {projection:>9} {r['elapsed_s']:>9.3f}" - ) - print() - for mode in cfg["modes"]: - label = "immediate" if mode == "injection" else mode - bonds = [r["peak_nu_bond"] for r in report["results"] if r["mode"] == label] - print( - f"{label}: peak |nu> bond over N = {bonds}" - ) - print( - "\nnote: immediate injection recycles one ancilla, whereas deferred MAST " - "uses one fresh ancilla per injected rotation and reports its final " - "projection cost separately." - ) - - -def build_arg_parser(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--n-list", default="6,12,24,48", - help="comma-separated qubit counts to sweep") - parser.add_argument("--t", type=int, default=3, help="number of T gates (fixed)") - parser.add_argument("--depth", type=int, default=6, help="Clifford circuit depth") - parser.add_argument("--seed", type=int, default=20260708) - parser.add_argument("--chi", type=int, default=None, - help="max |nu> bond (None = exact)") - parser.add_argument("--modes", default="direct,immediate,deferred", - help="comma-separated: direct, immediate, deferred (or injection alias)") - parser.add_argument( - "--exact-cooling", - action=argparse.BooleanOptionalAction, - default=True, - help="enable constructive exact cooling (use --no-exact-cooling to isolate MAST)", - ) - parser.add_argument( - "--deferred-projection-order", - default="middle_out", - choices=("middle_out", "input", "min_span"), - help="end-of-circuit magic-register projection order for deferred mode", - ) - parser.add_argument("--backend", default=None, - help="numpy (default), torch, cupy, or jax") - parser.add_argument("--rank-max-n", type=int, default=12, - help="report pseudo-stabilizer rank only for N <= this") - parser.add_argument("--json", default=None, help="optional path to write JSON") - return parser - - -def main(): - args = build_arg_parser().parse_args() - report = run(args) - _print_table(report) - if args.json: - with open(args.json, "w", encoding="utf-8") as fh: - json.dump(report, fh, indent=2) - print(f"\nwrote {args.json}") - - -if __name__ == "__main__": - main() diff --git a/docs/api/optimizers/energy.md b/docs/api/optimizers/energy.md index c4f13f2..884e614 100644 --- a/docs/api/optimizers/energy.md +++ b/docs/api/optimizers/energy.md @@ -14,11 +14,17 @@ import pepsy estimate = pepsy.MpsEnergyOptimizer( fermionic_state, - hamiltonian=hamiltonian.terms, + terms=hamiltonian, energy_per_site=False, ).energy() ``` +``terms=`` accepts either a local-term mapping or the complete +``SymHamiltonian`` returned by ``fermion.hamiltonian(terms)``. For an explicit +one-site-plus-two-site Hamiltonian, passing the ``SymHamiltonian`` preserves +the fermionic ordering metadata and lets a native fermionic MPS evaluate the +terms directly without constructing the bosonic MPO. + The MPO returned by a Fermi-Hubbard ``SymHamiltonian.to_mpo(...)`` uses the bosonic/Jordan-Wigner representation required by the DMRG MPO path. Pepsy now rejects an implicit contraction of that MPO with a native fermionic MPS, since diff --git a/docs/api/optimizers/mps.md b/docs/api/optimizers/mps.md index fbbe76f..b3050cd 100644 --- a/docs/api/optimizers/mps.md +++ b/docs/api/optimizers/mps.md @@ -50,6 +50,10 @@ If an independent physical copy is needed, use: physical = opt.p_ungauged.copy() ``` +For Symmray block-sparse MPS data, `gate_simple` automatically uses Quimb's +full two-site `split` path so symmetry and fermionic fusion metadata are +preserved. Dense MPS data keeps the faster `reduce-split` path by default. + `mode="swap"` applies non-local two-site gates through a swap-and-split path and swaps the sites back after each gate. `mode="perm"` uses the same swap-and-split path but leaves the swaps in place, tracking the current @@ -86,32 +90,36 @@ For unitary streams, 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; no diagnostic -flag is needed. For dense two-site non-unitary gates, the target +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 +dense two-site non-unitary gates, the target norm is obtained from the local expectation of `G†G`, so no copied target MPS is needed. Symmray and general sub-MPO backends use a raw target-norm fallback where that local expectation is not available. DMRG still materializes its target because FIT needs it, but the diagnostic uses FIT's final local norm trace as the current retained norm. Temporary fallback targets never modify the live `info_c` cache. -The `mpo`, `swap`, and `svd` progress bars always show the same cumulative -`infidelity` field, starting at zero before the first compressed two-site gate. +When tracking is enabled, the `mpo`, `swap`, and `svd` progress bars show the +same cumulative `infidelity` field, starting at zero before the first +compressed two-site gate. `mode="exact"` and `mode="su"` deliberately skip canonical metadata; switching back to an MPS mode rebuilds and canonicalizes the contracted state. The result API is intentionally small: ```python -opt.run(non_unitary=True, normalize_every=True) +opt.run(non_unitary=True, normalize_every=True, track_infidelity=False) -opt.get_infidelities() # [0.0, cumulative infidelity, ...] -opt.get_infidelity_samples() # detailed per-compression records +opt.get_infidelities() # [0.0] when tracking is disabled +opt.get_infidelity_samples() # [] when tracking is disabled opt.get_normalizations() # scale events and accumulated exponents ``` -Infidelity is recorded automatically whenever a compressed two-site update -occurs; no sampling or tracking option is needed. `get_infidelities()` is the -cheap cumulative trace for progress and stopping criteria. Use +When enabled, infidelity is recorded automatically whenever a compressed +two-site update occurs. `get_infidelities()` is the cheap cumulative trace for +progress and stopping criteria. Use `get_infidelity_samples()` when the target norm, retained norm, local ratio, or step metadata is needed. diff --git a/docs/api/optimizers/noise.md b/docs/api/optimizers/noise.md index 03be8f4..ff5f4f1 100644 --- a/docs/api/optimizers/noise.md +++ b/docs/api/optimizers/noise.md @@ -189,8 +189,8 @@ default `0.1` threshold when profiling a different workload. `TrajectoryEvent` is the general independent noise-simulation interface. Put one directly inside an ordinary gate stream and run independently sampled shots -with either `MpsOptimizer` or `MpsStabOptimizer`. It does not require Stim or a -density matrix. +with `MpsOptimizer`, `TreeOptimizer`, or `MpsStabOptimizer`. It does not require +Stim or a density matrix. Use a `mixture` for a user-defined random-unitary channel. Its outcomes have explicit probabilities, so `sample_trajectory_stream(...)` can make a concrete @@ -235,9 +235,9 @@ print(result.records[0]) `TrajectoryChannel.kraus([("no_jump", K0), ("jump", K1)])` accepts any complete local qubit channel (`sum(K.conj().T @ K) == I`) on the corresponding -one- or multi-qubit `TrajectoryEvent` support. For ordinary MPS replay, replace -the factory above with a fresh `MpsOptimizer(initial_mps, ...)` and pass its -usual options through `run_kwargs`. +one- or multi-qubit `TrajectoryEvent` support. For ordinary MPS or TTN replay, +replace the factory above with a fresh `MpsOptimizer(initial_mps, ...)` or +`TreeOptimizer(...)` and pass its usual options through `run_kwargs`. For `MpsStabOptimizer(track_infidelity=True)`, a selected Kraus outcome is a normalized trajectory boundary, just like a measurement/reset: its Born weight diff --git a/docs/api/optimizers/peps.md b/docs/api/optimizers/peps.md index a802ea0..4d0cbc7 100644 --- a/docs/api/optimizers/peps.md +++ b/docs/api/optimizers/peps.md @@ -39,9 +39,9 @@ SWAP routing. Use `route_opts` for routing controls such as `sequence`, `measure_final_infidelity=True`. If final measurement is disabled, the fallback optimizer loss can come from the coarser `boundary_chi` environment while the pre-check used `evaluation_chi`. -- Sweep mode defaults to the optional NLopt `LD_LBFGS` solver. Install NLopt - or pass an explicit `optimizer` / `sweep_optimize_kwargs` value when NLopt is - unavailable. +- Sweep mode defaults to Torch Adam for Torch-backed Symmray states and to the + optional NLopt `LD_LBFGS` solver otherwise. Pass an explicit `optimizer` / + `sweep_optimize_kwargs` value to override this choice. - `non_unitary=True` normalizes generated targets and candidates; it does not implement the interval scheduling or norm-proxy machinery available in `MpsOptimizer`. diff --git a/docs/api/optimizers/stabilizer_tn.md b/docs/api/optimizers/stabilizer_tn.md index c8264ad..3d7ba10 100644 --- a/docs/api/optimizers/stabilizer_tn.md +++ b/docs/api/optimizers/stabilizer_tn.md @@ -89,7 +89,8 @@ print(sim.stim_sample.faults) With `progbar=True`, the STN progress bar reports the current stream `part` (`clifford`, `T`, `measurement`, `reset`, `nonclifford`, ...) and the -single diagnostic field `norm_infidelity`. +MPS-compatible diagnostic field `infidelity`. The legacy +`norm_infidelity` field is emitted with the same value. `exact_cooling=True` is the default constructive pre-check for multi-site non-Clifford Pauli rotations. If the frame image has an isolated product @@ -251,10 +252,12 @@ no reference-state copy or overlap contraction. For normalized unitary evolution it records the cumulative proxy `1 - ||nu||**2` after compressed coefficient-MPS updates, reading the norm from the tracked one-site canonical centre. Unitary updates are not renormalized, so lost norm remains visible. -Projectors, measurements, coefficient-frame sub-MPOs, and arbitrary -non-unitary matrices do not emit infidelity samples; an unnormalized -non-unitary map also suspends later samples until projection restores a -normalized baseline. The public `infidelities` trace is therefore sparse and +For dense multi-qubit non-unitary matrices, the target norm is measured from +the local physical `G†G` expectation and the retained norm ratio is reported as +`infidelity`. Coefficient-frame sub-MPOs and arbitrary physical maps without a +certified target norm do not emit such samples; an unnormalized non-unitary +map also suspends later unitary samples until projection restores a normalized +baseline. The public `infidelities` trace is therefore sparse and historical, is not index-aligned with `bond_history`, and must not be summed. Projective measurement/reset boundaries do preserve the current segment before normalization in `norm_events`: the event records the pre-collapse norm, @@ -266,8 +269,9 @@ The post-collapse state is then normalized. Use `norm_diagnostics()` to form product/geometric-mean survival summaries across completed segments plus the current open segment; these summaries multiply unitary- and projector- compression survival factors, but not measurement probabilities. The preferred -summary keys are `norm_infidelity`, `norm_survival`, and `norm`; the older -`total_*_proxy` keys remain as compatibility aliases. +summary keys are `infidelity`, `fidelity`, `norm_survival`, and `norm`; +`norm_infidelity` and the older +`total_*_proxy` keys remain compatibility aliases. The proxy is not exact overlap fidelity or a discarded-SVD-weight report; validate physical accuracy independently when that distinction matters. diff --git a/docs/api/optimizers/sweep.md b/docs/api/optimizers/sweep.md index 5fc2869..23ea653 100644 --- a/docs/api/optimizers/sweep.md +++ b/docs/api/optimizers/sweep.md @@ -4,10 +4,15 @@ - `boundary_engine="dmrg"` uses Pepsy `BdyMPS` plus `CompBdy`. - `boundary_engine="quimb-mps"` uses Quimb MPS environments and scalar - `contract_boundary(...)` contractions. + `contract_boundary(...)` contractions. During a half-sweep it builds the + opposite-side environments once, then advances the moving boundary one row + or column at a time. - `boundary_engine="auto"` keeps dense inputs on `dmrg` and routes Symmray-looking inputs to `quimb-mps`. +Torch-backed Symmray blocks use the Torch autograd local solver. NumPy-backed +Symmray blocks retain the finite-difference fallback. + ```{eval-rst} .. automodule:: pepsy.optimizers.sweep.optimizer :members: diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index f84f9c1..f8ac2b8 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -6,9 +6,11 @@ gate stream `[(gate, where), ...]` on a **rooted tree tensor network**, after Huang, Mendl; Quantum 7, 964, 2023; [arXiv:2206.01000](https://arxiv.org/abs/2206.01000)). The state is stored with one leaf tensor per qubit. Internal nodes may have -**any arity** -- the default is a strictly-binary tree, but flatter `k`-ary -trees or gate-connectivity-driven communities (see *Tree structure*) work -through the same machinery. Gates are absorbed into the tree: +**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. Gates are absorbed +into the tree: - **single-qubit gates** are contracted into the leaf tensor with no bond growth; a unitary one-qubit gate preserves the tree canonical form regardless @@ -24,8 +26,19 @@ through the same machinery. Gates are absorbed into the tree: - **operators on three or more qubits** -- a `k`-qubit gate (Toffoli, Fredkin), a multi-site non-unitary / Kraus operator, or a whole Trotter block -- are applied *in one shot* over their minimal spanning subtree by - `apply_subtree_operator`; `apply_gate` routes any support with `len(where) >= 3` - there automatically (see *Multi-qubit / sub-MPO application*). + `apply_subtree_operator`. All open operator bonds are QR-routed to a subtree + hub before a final canonical compression sweep truncates every touched edge + once; `apply_gate` routes any support with `len(where) >= 3` there + automatically (see *Multi-qubit / sub-MPO application*). +- **stream events** -- MPS-compatible `measure`, `cap`, `reset`, and + `measure_reset` entries can be mixed into the stream. Measurements use Pauli + eigenvalue outcomes (`+1`/`-1`) and are appended to `measurements`; + explicit `submpo` markers use the same native QR-routing and final subtree + 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 + labels above it, and keeps the live tree canonical. The orthogonality centre is a single node id tracked on the `TreeTensorNetwork` itself (`orthogonality_center`), so the state -- not any one @@ -45,6 +58,26 @@ 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. +`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. +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; +it also invalidates the native fermionic norm cache. The optimizer's +state-aware wrappers do both automatically and restore the centre only for +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 +`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 +the required outer-leg phase flips), while an unknown centre falls back to an +exact complete doubled-network contraction. + ## Range / subtree canonicalisation The single orthogonality centre generalises to a connected **canonical region** @@ -55,7 +88,7 @@ when the region spans more than one node `orthogonality_center` honestly reads `None`. `TreeTensorNetwork.canonize_subtree_(nodes)` gauges every tensor *outside* a connected subtree to point inward (Quimb `canonize_around` with `which="any"`), so the whole state norm is carried by the region tensors -- -contracting just the region against its conjugate reproduces the squared norm, +contracting just the region against its graded conjugate reproduces the squared norm, 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 @@ -75,12 +108,15 @@ generalisation of the two-qubit gate: a `k`-qubit gate, a multi-site **non-unitary / Kraus** operator, or a whole **Trotter block**. It is the tree 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 operator is contracted onto the **minimal connected subtree** (Steiner -subtree) spanning the target leaves and the result re-split back into that -subtree with truncating SVDs. The orthogonality centre is first moved onto a -target leaf so the whole exterior is isometric and every re-split truncation sees -the complete operator against an isometric environment; the centre is left inside -the subtree, so the state stays in canonical form. +only). The dense operator is first factorized into an exact tree-MPO on the +**minimal connected subtree** (Steiner subtree) spanning the target leaves. +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. `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 @@ -89,7 +125,92 @@ afterwards (e.g. after a Kraus/projection operator). `max_bond` / `cutoff` default to the optimizer's `chi` / `cutoff`. `apply_gate` dispatches `len(where) == 1` and `== 2` to the optimised leaf-absorb / geodesic-threading paths and any larger support to `apply_subtree_operator`; the cost scales with the operator's -spread (the boundary of its spanning subtree) rather than the whole tree. +spread and factor ranks, using recursive edge messages rather than one dense +state tensor for the whole spanning subtree. + +An explicit MPS-style sub-MPO marker, `("submpo", mpo, where)` (or the +equivalent mapping form), is accepted in a TreeOptimizer stream. Quimb MPOs are +applied natively by carrying their virtual operator bonds through the same +lossless leaf-to-hub QR sweep followed by one subtree compression sweep; bond +estimates use MPO bond dimensions as a conservative Schmidt-rank bound. +Payloads without the required site interface fall back to `mpo.to_dense()`, +which must produce an operator on the declared support. +`TreeOptimizer.submpo_event(...)` builds the tuple form. + +Both two-site implementations preserve native Symmray gates and their +block-sparse fermionic grading. Direct mode splits the rank-four gate; MPO mode +asks Quimb to make the equivalent two-tensor sub-MPO. The resulting two factors +then enter the *same* TTN kernel: attach one factor, QR-thread its shared +operator bond along the unique path, attach the other, then make one final +compression sweep. Neither path converts the gate to a dense qubit array or +splits it into base-2 legs. + +For an ordinary two-site gate, choose `mode="direct"` to use the +specialised gate-SVD/QR threading implementation, or `mode="mpo"` to +use the same Quimb gate-to-sub-MPO route. `mode="auto"` is the default: +it selects direct factorization for every backend, including native Symmray +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. + +`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. +The old `run(mode="tree")`/`"ttn"` selector is a deprecated no-op retained only +for shared frontends. + +`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 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 +operate on dense two-level qubit coefficient states and do not require tableau +metadata; they intentionally reject native fermionic TTNs. +`measure_pauli` returns `(outcome, probability)` and accepts an optional +`return_diagnostics=True` flag. `project_pauli` normalizes by default; pass +`renormalize=False` to retain the branch norm. Both APIs can report projection +diagnostics containing the norm ratio and support, spanning-tree, and bond +snapshots before and after the update. The records are also available through +`projection_diagnostics` and `get_projection_diagnostics()`. + +## Coefficient-backend feature boundary + +`TreeOptimizer` covers the state operations shared with `MpsOptimizer`: +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 +`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 +that mapping. Native fermionic TTNs support native Symmray gates and MPOs, but +the qubit Pauli/measurement/reset helpers intentionally reject them; use the +fermion model's native observable/projector with +`TreeTensorNetwork.local_expectation` instead of silently treating a graded +local space as a qubit. + +The shared trajectory runner supports dense-qubit `TreeOptimizer` instances as +well. Independent trajectories can sample Pauli mixtures, depolarizing +channels, and state-dependent Kraus channels; branch probabilities are +evaluated from copied TTNs and selected branches are normalized before replay +continues. Coalesced trajectory replay supports exact branching of mid-circuit +measurement, reset, and measure-reset events through `expectation_pauli`. Use +matrix-valued gate payloads in tree streams (for example `pepsy.h()`), since +textual MPS gate aliases are not normalized by the tree gate parser. Native +fermionic trajectories may use native gates/MPOs, but Pauli/control events +require a model-native observable or projector. + +MPS execution modes such as `svd`, `dmrg`, `mpo`, `swap`, `perm`, `su`, and +`mix` are chain algorithms and are intentionally not copied into +`TreeOptimizer`. Tree layout is part of the TTN geometry and is selected with +`tree=`/`layout=` at construction. A later `TreeStabOptimizer` can therefore +delegate numerical coefficient updates to this class while keeping tableau +state and stabilizer-specific bookkeeping above it. ## Tree state class @@ -97,17 +218,19 @@ spread (the boundary of its spanning subtree) rather than the whole tree. geometry-owning subclass of Quimb's arbitrary-geometry vector class `quimb.tensor.TensorNetworkGenVector`. It *is* a Quimb tensor network, so all of Quimb's arbitrary-geometry methods (`canonize_around`, `canonize_between`, -`compress_between`, `gate_inds`, `local_expectation`, `to_dense`, `copy`, ...) -apply directly; the class adds the naming and geometry glue on top of a +`compress_between`, `gate_inds`, `to_dense`, `copy`, ...) apply directly; the +class adds the naming and geometry glue on top of a `TreePlan`: - every node (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 `site_tag_id.format(q)` (default `"I{}"`) and physical index `site_ind_id.format(q)` (default `"k{}"`) - for qubit `q`, so the inherited site / `local_expectation` machinery treats the - leaves as the sites; -- adjacent nodes share the deterministic virtual bond `_tb{lo}_{hi}`. + for qubit `q`, so Quimb treats the leaves as the sites; +- 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 + current live index. 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. @@ -118,6 +241,41 @@ product state in one step), or `TreeTensorNetwork.rand(plan, D=..., seed=...)` builds and evolves its state on this class, delegating all node/qubit naming and 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 +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 +API compatibility but cannot truncate this exact doubled-network contraction. +Observable readout deliberately belongs to the state, not to `TreeOptimizer`; +use `optimizer.tn.local_expectation(...)`. + +Readout is gauge-preserving: a dense expectation restores the previously +tracked canonical centre/region, while an unknown dense gauge is evaluated on a +temporary state copy. Native fermionic expectations do not move the gauge. +The repeated normalized native-readout denominator is cached and invalidated by +state mutation, copying, caps, and canonical/gate updates. + +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`. + +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. +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 +`(n_up, n_down)` occupations to choose each spin explicitly. The completed +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. + `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 @@ -138,13 +296,36 @@ 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. +For circuits with gates of different operator-Schmidt ranks, use +`TreeLayoutFinder(..., objective="congestion")` or +`TreeOptimizer(..., layout_objective="congestion")`. This evaluates interaction, +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. +`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 +`weight_mode` / `layout_weight_mode` option accepts `count`, `auto`, `angle`, or +`operator_schmidt` for interaction-graph weighting. + +`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. + The structure is **not restricted to binary trees**. Internal nodes may have any arity, controlled by two knobs on `TreeLayoutFinder` / `TreePlan.from_order` / `TreeOptimizer`: -- `max_arity` caps the children per internal node. The default `2` reproduces - the strictly-binary tree exactly; larger values give flatter `k`-ary trees - with shorter geodesics; `None` leaves the arity unbounded. +- `max_arity` caps the children per internal node. It accepts a scalar (a + 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. - `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 @@ -162,6 +343,197 @@ 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. +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). +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: + +```python +finder = TreeLayoutFinder(gate_stream, n=n, objective="congestion") +choice = finder.recommend_arities((2, 3, 4)) +opt = TreeOptimizer(gate_stream, tree=choice["plan"], chi=chi) +``` + +The `path` and `congestion` objectives are `chi`-blind cost proxies: they score +geodesic length and additive edge load, so they can favour a wider block or +arity whose widest bond induces a qubit bipartition too large to fit `chi`. +Every bond splitting `k` of the `n` qubits from the rest can carry a Schmidt +rank up to `2 ** min(k, n - k)`, so `TreePlan.max_bond_cut()` is a purely +structural accuracy ceiling: the tree can hold an arbitrary state exactly only +when `chi >= 2 ** max_bond_cut`. Pass `chi=` to `recommend_layered` or +`recommend_arities` to make the search `chi`-aware -- candidates are ranked +first by `chi_overflow` (how far the widest bond exceeds `log2(chi)`), so a +structure that stays exact at `chi` is preferred and the layout objective only +breaks ties. Each candidate then also reports `max_bond_cut`, `chi_overflow`, +and `exact_at_chi`: + +```python +choice = finder.recommend_layered((2, 3, 4), chi=chi) # prefers a chi-exact block +opt = TreeOptimizer(gate_stream, tree=choice["plan"], chi=chi) +``` + +For a fixed layered family, prefer an explicit recommendation to a hard-coded +block size: + +```python +finder = TreeLayoutFinder( + gates, + n=L, + objective="congestion", + weight_mode="operator_schmidt", + chi=chi, +) +choice = finder.recommend_layered(block_sizes=(2, 3, 4)) +tree_plan = choice["plan"] +``` + +`layered(block_size=4)` remains the right API when the block size is an +intentional experimental control: it spectral-orders the qubits and builds +exactly that fixed structure, but it does not score alternatives or use `chi`. +`recommend_layered()` inherits the finder's `chi` when its own `chi` argument +is omitted; pass `chi=None` explicitly for a chi-blind comparison. Inspect its +per-candidate `max_edge_load`, `peak_bond_growth`, `max_bond_cut`, and +`chi_overflow` rather than relying only on the chosen block size. + +### Fixed-plan refinement and Nevergrad search + +The tree topology is fixed before `TreeOptimizer` begins replay. Moving the +canonical centre and threading a gate along a path are tensor operations, not +layout rewrites. For a stronger *pre-simulation* layout search, +`recommend_layered` and `recommend_arities` can refine each candidate through +adjacent leaf-label swaps. This preserves every parent/child edge in the plan: +only which qubit label occupies each leaf changes. + +```python +finder = TreeLayoutFinder( + gate_stream, + n=n, + objective="hybrid", + hybrid_weights=(1.0, 1.0, 0.25), + chi=chi, +) +choice = finder.recommend_layered( + block_sizes=(2, 3, 4), + refine="greedy", + refine_budget=64, +) +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 +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. +Install the optional dependency with `pip install pepsy[layout]`: + +```python +choice = finder.recommend_layered( + block_sizes=(2, 3, 4), + refine="greedy", + search="nevergrad", + search_budget=128, + seed=0, +) +``` + +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 +`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. +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 +included, `peak_bond_growth_log2` remains finite even when the human-readable +`peak_bond_growth` would overflow floating point. + +Both helpers are also available from the package-level API: + +```python +import pepsy as py + +finder = py.TreeLayoutFinder(gate_stream, n=n, objective="congestion") +opt = py.TreeOptimizer(gate_stream, layout=finder, chi=chi) +``` + +To evolve a non-product or entangled initial state, pass it explicitly as +`state=` (or the backward-compatible `tn=`): + +```python +opt = TreeOptimizer(gate_stream, layout=finder, state=initial_ttn, chi=chi) +``` + +`tree=` / `layout=` accept only a `TreePlan` or `TreeLayoutFinder`; passing a +`TreeTensorNetwork` there raises an error so an entangled state cannot be +silently replaced by the default `|0...0⟩` product state. + +### Initial-state layout handoff and array backends + +`TreeLayoutFinder` is deliberately circuit-only: it consumes gate supports and +weights to choose a plan, never an already-entangled coefficient state. Pass +the resulting finder/plan *and* a state separately to `TreeOptimizer`. +An entangled `TreeTensorNetwork` must already own that same plan. Supplying a +different `tree=` or `layout=` raises an error before any tensor is changed: +there is no generally exact, cheap relayout of an entangled TTN, and silently +compressing it would hide a fidelity loss. + +Product states are the safe exception. A `TreeTensorNetwork` with +`max_bond() == 1` is rebuilt exactly on the requested plan (and emits a warning +that the requested layout replaced its old geometry). A bond-one Quimb +`MatrixProductState` is likewise accepted and mounted exactly on the selected +tree, so a caller may choose the tree layout after preparing an MPS product +state. Entangled MPS inputs are rejected rather than implicitly converted. + +The live TTN has one array contract: every tensor must have the same backend, +dtype, and device. `backend_info()` reports that contract. Construct the full +initial state and every user-supplied gate/operator with the same converter. +If a gate, sub-MPO, observable, or cap vector does not match, TreeOptimizer +converts it for compatibility but emits one warning per source/target +combination; this makes an unintended CPU/GPU transfer or dtype cast visible. +Mixed-backend initial states fail immediately because there is no unambiguous +safe execution backend. Internal Pauli/projector tensors follow the state +backend automatically. + +```python +import pepsy as py +import torch + +to_backend = py.backend_torch(device="cuda", dtype=torch.complex128) +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)) + +# Convert user-provided gate arrays once, at their source. +native_gates = [(to_backend(gate), where) for gate, where in gates] +opt = py.TreeOptimizer(native_gates, layout=plan, state=state, chi=chi) +assert opt.backend_info()["backend"] == "torch" +``` + +The same rule applies to CuPy; choose `py.backend_cupy(...)` and convert every +state tensor and payload with that converter. `to_dense()` intentionally returns +a host NumPy vector for interoperability; the live state remains on its native +backend. + ## Diagnostics The dominant lever for accuracy at fixed `chi` is the tree structure, so the @@ -170,28 +542,63 @@ finder and optimizer expose diagnostics to choose it: - `TreeLayoutFinder.report(plan=None)` summarises the leaf-to-leaf 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`). + (`balanced_score`, `score_ratio_vs_balanced`). It also reports + `edge_loads`, `max_edge_load`, and `peak_bond_growth` for the rank-aware + congestion estimate. - `TreeOptimizer.bond_report()` reports the current `max_bond`, `mean_bond`, and tensor/bond counts -- bonds pinned at `chi` mean truncation is active. +- `TreeOptimizer.estimate_bonds()` performs the paper's non-mutating dry run: + it multiplies the operator-Schmidt ranks of gates crossing each tree edge, + returning the conservative Eq. (4) bound before replay. This is useful for + choosing `chi`; it deliberately ignores cancellations and can overestimate + the live dimensions. +- `TreeOptimizer.preflight(...)` turns that bound into explicit resource + protection. `max_bond`, `max_operator_qubits`, and `max_subtree_nodes` can + reject a replay with `MemoryError`; the same limits can be passed to the + constructor for automatic checking before eager replay. The constructor + defaults to `max_operator_qubits=8` and `max_subtree_nodes=128`; pass `None` + to disable either guard. Product-Pauli measurements use a factorized parity + projector and do not materialize a `4**k` dense operator. +- `TreeOptimizer.truncation_report()` exposes the per-edge compression and + SVD-split history, including before/after bond dimensions. Pass + `track_truncation=True` to also collect each local full singular spectrum's + absolute discarded weight and relative discarded fraction. Dense states use + 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 + 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`, `norm`, observable `expectations`, `fidelity` against the untruncated state (when `2**n <= dense_cap`), and observable `max_drift` between consecutive - `chi` -- a reference-free convergence signal for large systems. + `chi` -- a reference-free convergence signal for large systems. Optional + observables are evaluated by the underlying `TreeTensorNetwork`; they are + not a `TreeOptimizer` readout API. ## Readout `to_dense()` returns the dense statevector in index order `k0, k1, ..., k(n-1)`. -`local_expectation(op, where)` returns ` / `. A single-site -operator contracts only the canonical centre tensor; a multi-site operator -contracts only the **minimal subtree spanning the target leaves**, with the -centre moved inside that subtree so the isometric exterior cancels to the -identity on the shared boundary bonds (cost scales with the operator's spread, -not the whole tree). `measure(q, outcome=None)` projectively measures a qubit in -the computational basis -- reading the exact Born probabilities from the -centred leaf, sampling (or forcing) an outcome, and renormalising -- and -`reset(q)` returns a qubit to `|0>`. `normalize()` rescales the represented -state to unit norm and `max_bond()` reports the largest virtual bond. +`run(progbar=True)` shows a tqdm replay bar with one-/two-/multi-qubit event +counts, current bond usage, state norm, and a norm-based truncation proxy. +Both dense and native fermionic replay report +`1 - (norm / reference_norm)^2`; the reference is established at run start and +reset after control or explicitly non-unitary events. This is display-only and +is not a replacement for the recorded truncation history. The bar is disabled +by default. +For dense two-level qubit TTNs, `measure(q, outcome=None)` projectively +measures a qubit in the computational basis and returns a bit; `reset(q)` +returns a qubit to `|0>`. Native fermionic TTNs deliberately do not expose +these qubit readouts. +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 +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 +`get_infidelity_samples()` when spectrum tracking is enabled. ## Performance and stability @@ -213,7 +620,8 @@ state to unit norm and `max_bond()` reports the largest virtual bond. - **Lazy canonical centre.** A freshly built product state has every virtual bond at dimension 1, so it is already canonical with the root as orthogonality centre; `from_plan` records that centre on the network rather - than recomputing it on the first gate. + than recomputing it on the first gate. Native fermionic product trees are + additionally normalized by their exact graded norm readout. - **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 @@ -222,6 +630,12 @@ state to unit norm and `max_bond()` reports the largest virtual bond. - **Self-healing tid cache.** Node-to-tensor lookups are cached and validated against the live tensor map, so the hot path avoids re-scanning tags while staying correct when a gate rebuilds a tensor. +- **Resource guards.** `max_intermediate_bond`, `max_operator_qubits`, and + `max_subtree_nodes` provide preflight and direct-application limits; the + latter two default to conservative finite values and accept `None` to opt + out. A dense `k`-qubit operator still has `4**k` payload values, while + product-Pauli measurement uses a factorized parity projector and recursive + edge messages. - **`copy()`.** Returns an independent optimizer that shares the immutable `TreePlan` but owns its own tensor network (which carries the tracked orthogonality centre), for branching experiments or trial gate sequences. diff --git a/docs/api/sampling/samplers.md b/docs/api/sampling/samplers.md index 72a28ff..3ade096 100644 --- a/docs/api/sampling/samplers.md +++ b/docs/api/sampling/samplers.md @@ -22,6 +22,110 @@ require quimb to canonicalize Torch or CuPy tensors before sampling. It caches those environments for the current MPS tensors: after modifying the MPS, call `sampler.refresh()` before sampling again. +For a Symmray MPS, `MpsSampler` detects the block-sparse tensor data and +selects its native symmetry-aware route automatically. It caches a +right-canonical copy, preserves the source physical-code/charge map, then +samples by projecting one physical sector at a time and absorbing it into a +block-sparse boundary. No MPS tensor is converted with `to_dense()`. The +physical-code reconstruction is generic over Symmray's abelian charge maps, +including non-fermionic and fermionic `Z2`, `U1`, `U1U1`, and `Z2Z2` states. +In particular, a degenerate physical sector such as spinful `U1` remains a +selected sparse block rather than forcing the MPS dense. The sampler does not +need a Fermion object: it preserves the generic source basis map for every +site: + +```python +sampler = MpsSampler(psi) # Symmray MPS +print(sampler.physical_code_maps) +# ({0: (charge_0, 0), 1: (charge_1, 0), ...}, ...) +``` + +Each map is `physical_code -> (Symmray charge, offset within that charge +sector)`. It includes source sectors with zero amplitude that canonicalization +pruned, so it remains suitable for interpreting user configurations. +The symmetry-aware route currently requires an open-boundary MPS; periodic +MPS sampling needs a separate cyclic conditional-environment algorithm. + +For a batch, all shots that have the same sampled prefix share its normalized +block-sparse boundary and local conditional distribution. This makes product, +low-entanglement, and charge-restricted states substantially cheaper to sample +while retaining the source MPS and its symmetry metadata unchanged. This +supports fermionic U1U1 starts directly: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") +psi = pepsy.ps_to_mps(8, fermion=fermion) +sampler = MpsSampler(psi) # resolved_backend == "symmray" +configs, probs = sampler.sample_arrays(256, seed=0) +``` + +Use `backend="symmray"` to require this route explicitly. If Symmray blocks +are backed by Torch or CuPy, contractions and local probability vectors remain +on that device; only the final discrete code choice is synchronized with the +Python control flow. + +### Fermionic configuration codes + +`configs` are always physical-index codes, not universal occupation labels. +Bind the `Fermion` definition when constructing the sampler to attach an +explicit, symmetry-aware code map to every batch before using it in a VMC or +local-estimator workflow: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="Z2") +sampler = MpsSampler(psi, backend="symmray", fermion=fermion) +batch = sampler.sample_batch(4096, seed=0) + +physical_configs = batch.configs # shape (batch, n_sites) +occupations = batch.occupations() # shape (batch, n_sites, 2), (n_up, n_down) +encoding = batch.configuration_encoding +assert np.array_equal(encoding.encode(occupations), physical_configs) +``` + +Passing `fermion=...` directly to `sample_batch` remains supported when one +sampler is intentionally shared across compatible workflows. A bound Fermion +also lets `fermion_configuration_encoding()` and the diagonal-observable +helpers omit their repeated Fermion argument. + +This is essential for collapsed sectors: spinful Symmray `Z2` uses physical +codes in `empty, double, up, down` order, while resolved `U1`, `U1U1`, and +`Z2Z2` use their sector-derived code order. The +`FermionConfigurationEncoding` object is immutable, site-aware, and rejects +invalid codes instead of silently interpreting them with a VMC-specific +default. + +For predictable memory on high-entropy states, choose the prefix policy when +constructing the sampler: + +```python +sampler = MpsSampler( + psi, + backend="symmray", + prefix_strategy="auto", # "prefix" or "serial" are also available + max_prefix_groups=256, +) +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. +Use `python benchmarks/mps_symmray_sampling.py --help` for comparable +product/entangled, fermionic/non-fermionic sampling measurements. By default, +it runs the sparse Symmray, dense-native, and Quimb sampler variants in +separate processes and reports median throughput, setup peak RSS, and +post-setup resident RSS for each variant. +Use `--variants symmray` when comparing prefix policies only. +For fermionic `Z2`/`Z2Z2` inputs it reports the dense variants as unsupported: +naively expanding their graded virtual legs is not a state-preserving dense MPS +conversion, so only the Symmray result is a valid comparison. + Torch sampling is inference-only by default, avoiding an autograd graph for the discrete draw and its sampled probabilities. Use `track_grad=True` only when those sampled Born probabilities need gradients: @@ -44,6 +148,40 @@ probs = sampler.probabilities(configs, to_numpy=False) Both methods evaluate the whole batch in one backend-native sweep. +### Fermionic diagonal observables + +`MpsSampler` can estimate diagonal observables of a +`pepsy.tensors.Fermion` directly from Born samples. This safely handles the +block-ordering of spinful Symmray `Z2`, `U1`, and `U1U1` physical indices: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") +sampler = MpsSampler(psi, backend="symmray", fermion=fermion) +estimate = sampler.estimate_fermion_diagonal( + "doublon", # average n_up n_down per site + n_samples=16_384, + seed=0, +) +print(estimate.mean, estimate.standard_error) + +density_corr = sampler.estimate_fermion_diagonal( + "density_correlation", # average n_i n_j over the listed pairs + pairs=((0, 1), (2, 3)), + n_samples=16_384, + seed=1, +) +``` + +Supported names are `"occupation"` (mean occupation on `sites`), +`"total_charge"` (occupation sum), `"doublon"`, and +`"density_correlation"`. `fermion_diagonal_values(configs, fermion, ...)` +evaluates the same observable on an existing batch, which is useful for exact +short-chain probability sums. The returned `MpsDiagonalEstimate` contains the +sample mean and a standard error based on the unbiased sample variance. +Hopping, pairing, and spin-flip +operators are not diagonal and require a separate fermionic local-estimator +calculation. + ```{eval-rst} .. automodule:: pepsy.sampling.samplers :members: diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index 8c66391..98affbe 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -77,11 +77,39 @@ psi = py.ps_to_mps( ) ``` -With ``chi=1`` this creates a charge-fixed product MPS. A larger ``chi`` uses -charge-preserving random-unitary growth while keeping the same global sector. -``hrps_to_mps`` retains its ordinary local-Haar meaning and is not used for -fixed-charge Fermion states, since a generic local Haar vector breaks the -conserved sector. +``ps_to_mps`` always creates a bond-one, charge-fixed product MPS. For a +random charge-preserving MPS with a requested bond dimension, use +``hrs_to_mps``: + +```python +psi = py.hrs_to_mps( + 8, + fermion=fh, + occupations=fh.half_filled_occupations(8), + chi=4, + seed=7, +) +``` + +The random-unitary growth fills only symmetry-compatible sectors and keeps the +same global charge. The historical alias ``hrps_to_mps`` points to the same +constructor. For direct Symmray block filling, use ``method="direct"``: + +```python +psi = py.hrs_to_mps( + 8, + fermion=fh, + occupations=fh.half_filled_occupations(8), + chi=4, + method="direct", + subsizes="maximal", + seed=7, +) +``` + +The direct method uses Symmray's ``MPS_fermionic_rand`` constructor, then +right-canonicalizes and normalizes the MPS. It is useful when direct random +block filling is preferred; ``method="unitary"`` remains the default. The matching PEPS constructor is ``ps_to_peps``. Pass ``(Lx, Ly)`` (or one integer for a square lattice), a coordinate occupation mapping, and the same @@ -97,14 +125,31 @@ seed = py.ps_to_peps( (Lx, Ly), fermion=fh, occupations=occupations, - chi=1, seed=7, ) ``` -This returns the underlying PEPS with native fermionic Symmray tensors; use -``SymPEPS`` only when the wrapper methods or stored Hamiltonian metadata are -needed. +``ps_to_peps`` always returns a bond-one, charge-fixed product PEPS. For a +random Symmray block-filled PEPS with a requested bond dimension, use +``hrs_to_peps``: + +```python +psi = py.hrs_to_peps( + (Lx, Ly), + fermion=fh, + occupations=occupations, + chi=4, + method="direct", + subsizes="maximal", + seed=7, +) +``` + +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. ## Unified native fermion helper @@ -167,6 +212,16 @@ double_occupancy = spinful.operator_term( sites=(i,), ) +# Add the correctly ordered fermionic Hermitian conjugate automatically. +hopping_up = spinful.operator_term( + [(1.0, ((0, "create_up"), (1, "annihilate_up")))], + sites=(0, 1), + add_hc=True, +) + +# Built-in convenience for Delta_0^dagger Delta_1 + h.c. +eta_pair = spinful.eta_pair_operator() + # A complete explicit Hamiltonian can be wrapped for MPO/DMRG conversion. terms = { edge: -t * spinful.hopping_operator() @@ -298,6 +353,46 @@ object when a workflow prefers a model-family entry point. ``U1``, ``0`` for ``Z2``, and ``(1, 1)`` for ``U1U1``/``Z2Z2``. Pass those charges to ``SymMPS.measure`` when evaluating a charged pairing correlator. +For explicit rectangular PEPS workflows, ``Fermion.lattice_half_filling`` +prepares only the lattice and charge metadata. It does not build a PEPS, +Hamiltonian, or gate stream, so those steps remain visible in the calling +notebook: + +```python +setup = fermion.lattice_half_filling( + Lx, + Ly, + pattern="checkerboard", + cyclic=True, # periodic physical Hamiltonian edges +) + +sites = setup.sites +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} +ham = fermion.hamiltonian(terms) + +gate_stream = fermion.strang_gate_stream( + edges, + dt=0.01, + sites=sites, + imaginary=True, +) +``` + +The ``cyclic`` option controls only the physical edge list in ``setup``. Native +fermionic ``ps_to_peps(..., cyclic=True)`` also supports periodic Symmray PEPS +bonds. An OBC MPS or PEPS can still represent a PBC Hamiltonian by using +``cyclic=True`` for the physical edge list and ``cyclic=False`` for the state; +nonlocal wrap-around gates are then routed through the open state network. + +For ``U1U1``, ``setup.occupations`` contains ``(n_up, n_down)`` pairs. For +spinful ``U1`` and ``Z2``, it contains scalar total occupations while +``setup.spin_occupations`` retains the up/down checkerboard pattern for +initial walker configurations. + ## Random-unitary MPS starts For DMRG-style random starts, prefer ``SymMPS.random_unitary_evolution`` or the @@ -483,6 +578,28 @@ ham = py.SymHamiltonian.from_edges( mpo = ham.to_mpo(L=3, compress=True, max_bond=16, cutoff=1e-12) ``` +### MPO compression diagnostics + +Every MPO built by ``SymHamiltonian.to_mpo`` carries a construction record in +``mpo.pepsy_compression_report``. It records the raw and final maximum bond, +the requested cap, whether compression reduced the represented rank, and +whether Symmray returned a bond larger than the requested cap: + +```python +report = mpo.pepsy_compression_report +print(report["raw_max_bond"], report["final_max_bond"]) +if report["max_bond_exceeded"]: + print("The requested max_bond was a soft cap for this compression.") +``` + +Pepsy emits a ``RuntimeWarning`` in the latter case. This can occur when a +positive cutoff leaves singular values tied at the selected threshold. The +report's ``rank_reduced`` flag says that the MPO bond rank changed; it is not +an error estimate. For a hard numerical cap, use ``cutoff=0.0`` and verify the +returned ``final_max_bond``. When the cap binds, also check convergence against +a larger cap because the compressed MPO can represent an approximation to the +original Hamiltonian. + ``to_mpo`` also accepts ``to_backend=`` to map each stored Symmray block to an array backend, and ``dtype=`` to choose the dense local operator dtype used before conversion to block-sparse arrays. The implementation is validated by diff --git a/docs/api/vmc.md b/docs/api/vmc.md index 088f79b..90ef806 100644 --- a/docs/api/vmc.md +++ b/docs/api/vmc.md @@ -11,12 +11,34 @@ Use the model-specific builders for NetKet-backed VMC: - `build_ising_vmc(...)` - `build_heisenberg_vmc(...)` - `build_fermi_hubbard_vmc(...)` +- `build_fermion_vmc(...)` — general spinful-fermion models (not only + Fermi-Hubbard), built from a `pepsy.Fermion`, explicit symbolic terms, or a + ready NetKet operator, with optional measurement observables. 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: +main path compact. 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. +- `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`. + +For fermions, the supported JIT configuration intentionally combines a flat +`Z2` Symmray PEPS ansatz with NetKet's fixed `(N_up, N_down)` `U1U1` +`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. ```python import pepsy.vmc as pvmc @@ -42,6 +64,160 @@ Lower-level functions such as `pack_peps_ansatz(...)`, `make_peps_batched_amplitude_function(...)`, and `make_netket_vmc_driver(...)` remain public for custom flows and profiling. +### Portable Torch / NetKet VMC façade + +The native Torch and NetKet entry points deliberately retain their own +performance controls. The preferred portable workflow follows NetKet's shape: +an `MCState` owns the PEPS ansatz and sampling recipe, while `VMC` owns the +Hamiltonian and drives sampling, expectations, and optimization. Both +backends return `VMCSamples`, `VMCMeasurement`, and +`VMCOptimizationResult`; each result retains the engine-specific value in +`.native`. + +```python +import pepsy.vmc as pvmc + +from pepsy.vmc import ( + LocalMatrixTerm, + MCState, + OperatorFactor, + OperatorSum, + ProductTerm, + VMC, +) + +hamiltonian = OperatorSum.from_terms([ + ProductTerm( + coefficient=-1.0, + factors=( + OperatorFactor(0, "fermion", spin="up", dagger=True), + OperatorFactor(1, "fermion", spin="up", dagger=False), + ), + ), + LocalMatrixTerm(support=(0,), matrix=onsite_matrix), +]) +state = MCState( + peps, + n_samples=16_384, # total over all chains, as in NetKet + n_chains=32, + n_discard_per_chain=64, + symmetry="U1U1", +) + +torch_vmc = VMC( + hamiltonian, + state, + backend="torch", + fermion=fermion, + observables={"density": density_operator}, + contraction=pvmc.ContractionConfig(method="boundary", chi=32), +) +netket_vmc = VMC( + hamiltonian, + state, + backend="netket", + fermion=fermion, + contraction=pvmc.ContractionConfig(method="exact"), +) + +samples = torch_vmc.sample() +energy = torch_vmc.expect() +history = torch_vmc.run( + optimization=pvmc.OptimizationConfig(n_steps=100, method="sr"), +) +``` + +`VMCProblem` and `build_torch_vmc(...)` / `build_netket_vmc(...)` remain as +compatible lower-level construction APIs. `MCState.to_problem(...)` bridges +to that representation when integrating existing code. + +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. + +`compile_operator_sum_torch(...)` lowers the common terms to Torch connection +tables (using a supplied `Fermion` object for graded symbolic factors), while +`compile_operator_sum_netket(...)` lowers them to a NetKet operator. The +existing `TorchFermionVMC(..., terms=operator_sum)` and +`build_fermion_vmc(..., hamiltonian=operator_sum)` entry points use these +compilers automatically. + +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) +torch_samples = torch_vmc.sample(sampling) +netket_samples = netket_vmc.sample() # its sampler was built from MCState + +optimization = OptimizationConfig( + n_steps=100, + method="sr", # "sgd" also works on both backends + learning_rate=2.0e-2, + diag_shift=1.0e-2, +) +torch_history = torch_vmc.optimize(optimization) +netket_history = netket_vmc.optimize(optimization) +``` + +### Three sampling layers + +The portable Torch VMC interface keeps three deliberately separate workflows: + +1. Metropolis-Hastings measurement: call `vmc.measure()` or + `vmc.expect()` with no supplied batch. +2. Measurement of one or more observables from a supplied batch: pass + `samples=...`. For a weighted empirical estimate also pass `weights=...`. + For importance sampling, prefer `proposal_log_probs=log_q`, which causes + the current `|psi(x)|**2 / q(x)` weights to be calculated stably. +3. Autodiff energy optimization: call `vmc.optimize(...)` normally for fresh + Metropolis samples, or pass the same supplied batch and `proposal_log_probs` + to reuse an importance-sampling proposal over several autodiff SR or minSR + updates. + +```python +samples = torch_vmc.sample() +measurement = torch_vmc.measure(samples=samples) + +# `external_configs` was drawn from q(x), whose log-density is known. +proposal_samples = pvmc.VMCSamples( + configs=external_configs, + proposal_log_probs=external_log_q, +) +weighted_measurement = torch_vmc.measure( + {"density": density_operator}, + samples=proposal_samples, +) +history = torch_vmc.optimize( + pvmc.OptimizationConfig(n_steps=20, method="sr"), + samples=proposal_samples, +) +``` + +Fixed `weights=` are useful for a single weighted empirical estimate or update, +but they do not change as the ansatz changes. Retain `proposal_log_probs` for a +reused optimization batch so the importance weights are refreshed at every +parameter value. NetKet's portable adapter currently supports the first layer +through its jitted `MCState`; it raises `VMCBackendCapabilityError` rather than +silently accepting external weighted batches. + +For NetKet, `n_chains`, seeds, and sampler choice are fixed at construction, +so supply them on `MCState` (or `sampling=` on the compatible legacy +builders) and then call `sample()` without a new seed. Portable NetKet +currently supports the fixed spinful `U1U1` sector; +use the native builders or `build_torch_vmc` for the other symmetry families. +`thin`, a replacement proposal, and post-construction reseeding raise +`VMCBackendCapabilityError` rather than being silently ignored. `minsr` is +currently Torch-specific; use portable `method="sr"` for NetKet. + +The original `TorchFermionVMC`, `TorchVMCDriver`, and `NetKetPEPSVMC` APIs +remain available when native result objects or backend-specific controls are +required. + ## Torch sampling and local-energy kernels `pepsy.vmc.torch` provides lightweight PyTorch kernels for custom VMC loops and @@ -79,6 +255,33 @@ sample = pvmc.metropolis_exchange_sweep( hopping_rate=0.25, encoding=encoding, ) + +# NetKet-like stateful sampling. ``n_samples`` is total across chains. +sampler = pvmc.TorchMetropolisSampler( + model, + graph, + configs, + proposal="spinful", + n_chains=128, # normally inferred from configs + seed=2, +) +samples = sampler.sample( + n_samples=1024, + n_discard_per_chain=500, + sweep_size=graph.n_sites, +) +print(samples.configs.shape) # (chain_length, n_chains, n_sites) + +# ``n_discard`` and ``n_thin`` are convenience aliases. +samples = pvmc.metropolis_local_sampler( + configs, + model, + graph, + n_samples=1024, + n_discard=500, + n_thin=graph.n_sites, + seed=2, +) conn = pvmc.spinful_fermi_hubbard_connections( sample.configs, graph, @@ -94,12 +297,59 @@ eloc = pvmc.local_energy_from_connections( ) ``` +If the amplitude model exposes `forward_log(configs)` and returns +`(phase, log_abs)`, Metropolis acceptance uses +`exp(min(0, 2 * (log_abs_new - log_abs_old)))` instead of dividing raw +amplitudes. Models without that method retain the raw-amplitude fallback. +For a small float64 PEPS where raw amplitudes are safely representable, pass +`log_amplitude_fn=False` to `TorchMetropolisSampler`, `TorchVMCDriver`, or +`TorchFermionVMC` to skip the additional stable-log contractions; leave the +default enabled for larger states that can underflow. + +`TorchPEPSBoundaryAmplitude` and `TorchFermionVMC` also accept +`proposal_batching="auto"`, `"vmap"`, or `"cache"`. The automatic policy +uses full native `torch.vmap` proposal batches on CUDA once enough walkers +changed; it retains the cached boundary-environment path when the selected +Symmray representation or contraction cannot be vmapped. Use `"vmap"` and +`"cache"` to benchmark the two numerically equivalent routes explicitly. +Independent configuration batches have a separate +`amplitude_batching="auto"` / `"vmap"` / `"serial"` policy on +`TorchPEPSAmplitude`, `TorchPEPSBoundaryAmplitude`, and `TorchFermionVMC`. +`"auto"` probes the pure `torch.vmap` path and falls back permanently to +serial contraction when the backend is not vmappable; use `"vmap"` for flat +Z2 PEPS and `"serial"` for native U1/U1U1 block-sparse PEPS. The last selected +route is available as `model.last_amplitude_batching` for profiling. +`TorchFermionVMC` defaults to `contraction="boundary"` with `chi=4`; pass a +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 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. +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, +`result.energy_stderr` is the autocorrelation-corrected +`sqrt(variance / effective_sample_size)`, while +`result.energy_stderr_naive` retains the independent-sample value and +`result.effective_sample_size` is available directly on the result. For a +single retained chain or sample, both error fields use the ordinary sample +count. + `TorchPEPSAmplitude.parameters()` returns the packed PEPS tensor leaves as torch parameters, so plain torch optimizers and the lightweight SR/minSR helpers can update the tensor network directly. ```python -log_derivatives = pvmc.torch_log_derivative_matrix(model, sample.configs) +log_derivatives = pvmc.torch_log_derivative_matrix( + model, + sample.configs, + derivative_backend="auto", # batched PEPS Jacobian with scalar fallback +) sr = pvmc.solve_torch_sr( log_derivatives, eloc, @@ -109,9 +359,49 @@ sr = pvmc.solve_torch_sr( pvmc.apply_torch_sr_update(model, sr.direction, learning_rate=0.02) ``` -For boundary-MPS contractions, `TorchPEPSBoundaryAmplitude` also reuses each -parent walker's row or column environments when evaluating local Hamiltonian -connections: +`diag_shift` can instead be a callable of the SR update number. The solver +uses Cholesky for the positive-definite SR metric and falls back to a Hermitian +pseudoinverse when a small batch leaves it rank-deficient. For noisy, +small-batch updates, `sr_momentum` enables a SPRING-style update: it retains +only the component of the previous SR direction outside the current sampled +tangent span, rather than adding conventional momentum directly. + +```python +def sr_shift(step): + return max(1.0e-3, 1.0e-1 * 0.8**step) + +step = vmc.step( + sample_sweeps=1, + sr=True, + learning_rate=2.0e-3, + sr_diag_shift=sr_shift, + sr_pinv_rtol=1.0e-10, + sr_momentum=0.8, # optional; disabled by default +) +print(step.sr.info["solver"], step.sr.info["relative_residual"]) +``` + +For PEPS models, `derivative_backend="auto"` uses one batched Torch Jacobian +when the amplitude model exposes its tensor parameters, and falls back to the +scalar loop for generic models. Use `"batched"` to require the optimized path +or `"loop"`/`"scalar"` for compatibility diagnostics. + +Local-energy evaluation also coalesces repeated `(walker, configuration)` +connections and batches each unique off-diagonal target amplitude once before +scattering the values back to the Hamiltonian terms. + +For boundary-MPS contractions, `make_torch_peps_amplitude_model(..., +contraction="boundary")` now constructs `TorchPEPSBoundaryAmplitude`. It +reuses bounded row or column environment caches for local Hamiltonian +connections and for local Metropolis proposals, rather than rebuilding a full +boundary contraction for every proposal. During a local-energy measurement, +connected configurations are grouped by parent and preferred boundary strip. +A cached parent-selected strip is copied and only its changed physical +projectors are replaced, so all terms in that strip share two-sided boundary +environments and unchanged tensor data. The cache is keyed by the parent and +target configurations and automatically invalidated when torch PEPS leaves +change. The optional VMC profile exposes `num_groups`, +`num_strip_cache_hits`, and `num_strip_builds` alongside reuse/fallback counts: ```python model = pvmc.TorchPEPSBoundaryAmplitude( @@ -122,32 +412,341 @@ model = pvmc.TorchPEPSBoundaryAmplitude( ) ``` +For approximate PEPS, pass a cotengra optimizer in +`final_contract_opts`. It controls every remaining *exact scalar closure*: +the ordinary boundary result, cached local-energy/proposal closures, and the +final closure after CTMRG. It does not change the boundary/CTMRG truncation +set by `chi` and `cutoff`. + +```python +from pepsy import build_optimizer + +final_contraction_opt = build_optimizer( + progbar=False, + directory="ctg_cash", + parallel=True, +) +boundary_opts = { + "mode": "mps", + "final_contract": True, + "final_contract_opts": {"optimize": final_contraction_opt}, + "progbar": False, +} +model = pvmc.TorchPEPSBoundaryAmplitude( + peps, + chi=64, + cutoff=1e-10, + contraction_opts=boundary_opts, +) +``` + +If a reusable cotengra search cannot produce a tree for a previously unseen +closure, VMC emits one warning and retries that closure with `"auto-hq"`. + For a compact no-JIT torch loop, `TorchVMCDriver` keeps walker configurations, current amplitudes, Hamiltonian connection metadata, sampling, local-energy evaluation, and optional SR updates together. Local-energy evaluation reuses diagonal connected amplitudes, supports chunked off-diagonal amplitude calls, and uses `connected_amplitudes(...)` automatically when the model provides it. +For native Pepsy Hamiltonians, pass the explicit term mapping and let the +driver build all connected configurations directly from those one- and +two-site operators: ```python +ham = fermion.hamiltonian(terms) driver = pvmc.TorchVMCDriver( model, graph, configs, - connection_fn="fermi_hubbard", - connection_kwargs={"t": 1.0, "U": 8.0, "encoding": encoding}, + terms=ham.terms, + site_order=tuple(peps.sites), proposal="spinful", hopping_rate=0.25, chunk_size=64, ) -result = driver.step( +result = driver.estimate_observable( + burn_in=20, + n_measurements=8, + sweeps_between=2, + progress=True, +) +print( + result.energy_mean, + result.energy_stderr, + result.samples_per_second, +) +``` + +For a repeated VMC workflow, use the driver's phase-specific methods rather +than writing an outer sampling loop. `warmup_proposal_mix(...)` is optional +adaptive local-move warm-up, `burn_in(...)` is fixed-kernel equilibration, +`optimize(...)` performs repeated `step(...)` updates, and +`estimate_observable(...)` performs the final measurement. All accept +`progress=True` where useful; the internal optimization bar reports energy per +site, acceptance, optional no-op fraction, and the SR solver. + +```python +driver.warmup_proposal_mix(n_sweeps=16, progress=True) +driver.burn_in(50, progress=True, track_proposal_stats=True) +history = driver.optimize( + 10, + progress=True, + sample_sweeps=1, sr=True, - learning_rate=0.02, - sr_diag_shift=1e-3, + learning_rate=2.0e-3, + sr_method="minsr", + track_proposal_stats=True, +) +final_step = history[-1] +``` + +`run(n_steps, ...)` remains a compatibility alias for `optimize(n_steps, ...)`. + +The local exchange/hopping sampler can expose move-wise diagnostics without +using a BP proposal. Pass `track_proposal_stats=True` to a sweep, `sample`, or +`step` to obtain `selected`, `no_op`, `proposed`, and `accepted` counts for +`exchange`, `hopping`, `spin_flip`, and `pair_toggle` branches. This helps +separate rejected PEPS-amplitude proposals from invalid local moves: + +```python +step = driver.step(track_proposal_stats=True) +print(step.proposal_stats) +``` + +For a short local-sampler warm-up, `warmup_proposal_mix(...)` uses those +counters to rebalance the enabled exchange/hop, spin-flip, and pair-toggle +branches. It holds all rates fixed during each whole graph sweep, adapts only +between completed warm-up sweeps, and then leaves the final rates fixed for +normal VMC sampling. The warm-up configurations are not production samples: + +```python +tuning = driver.warmup_proposal_mix( + n_sweeps=32, + adaptation_rate=1.0, + min_probability=0.05, +) +print(tuning["rates"]) + +# The following measurement uses the frozen tuned local proposal mix. +result = driver.estimate_observable( + burn_in=20, + n_measurements=64, + sweeps_between=2, ) -print(result.energy_mean, result.acceptance_rate) ``` +An exactly disabled branch (`rate=0`) or an always-selected branch (`rate=1`) +is treated as an explicit user choice and is not re-enabled by the tuner. + +When energy and another native-term observable (for example an eta-pair +correlator) are needed from the same Markov chain, use +`estimate_observables`. It merges matching `(walker, connected configuration)` +targets before amplitude evaluation, so boundary-MPS environments are shared +as well: + +```python +measurements = driver.estimate_observables( + { + "energy": ham.terms, + "eta": eta_terms, + }, + burn_in=50, + n_measurements=50, + sweeps_between=1, + profile=True, +) +energy = measurements["energy"] +eta = measurements["eta"] +print(energy.energy_mean, eta.energy_mean) +print(energy.profile) # sampling / connection / estimator timings and cache stats +``` + +With the modern `n_samples=...` interface, `progress=True` first reports the +Metropolis sweeps and then shows a staged evaluation bar for shared connection +building, all requested observable contractions, and final statistics. + +`local_observables({"energy": ham.terms, "eta": eta_terms})` provides the +same connected-amplitude sharing for the driver's current walkers without a +sampling pass. For one observable, pass `profile=True` to +`estimate_observable` or `step` to record phase timings and the available PEPS +boundary-cache counters. + +To measure samples saved from an earlier run without starting another +Metropolis chain, pass the `TorchMCMCSamples` object directly to +`measure_samples`. A chain-shaped configuration tensor with shape +`(n_samples_per_chain, n_chains, n_sites)` is accepted as well; stored +amplitudes are reused when available, or recomputed with the active batching +backend when omitted. The returned estimate retains the chain axis and +reports autocorrelation-aware error bars. Repeated parent configurations and +connected targets are deduplicated before amplitude contraction by default; +the profile includes the number of unique parent samples: + +```python +saved = driver.sample(n_samples=4096, n_chains=16, n_thin=2) +energy = driver.measure_samples(saved) +observables = driver.measure_samples( + saved, + observables={"energy": ham.terms, "eta": eta_terms}, +) +print(energy.energy_mean, energy.chain_diagnostics) +``` + +For repeated runs with a fixed walker shape, set `compile_kernels=True` on +`TorchVMCDriver`, `TorchFermionVMC`, or `TorchMetropolisSampler`. It only +compiles cheap tensor work: proposal construction, connection-key packing, and +the local-energy `index_add` scatter. PEPS/Symmray configuration selection and +all contractions remain eager. If the active Torch installation cannot compile +those kernels, Pepsy automatically keeps the eager implementation. + +The explicit-term route is Hamiltonian-agnostic and is the preferred +measurement path when the caller already has `ham.terms`. The older +`connection_fn="fermi_hubbard"` / `connection_kwargs={...}` route remains +available for lightweight custom loops. `FermionSiteEncoding.vmc_torch()` is +the default torch spinful encoding: `0=empty, 1=down, 2=up, 3=double`. +Native graded fermion terms are compiled into sparse configuration transitions +on their first measurement and reused across subsequent VMC sweeps. The +compiled form includes the endpoint crossing phase and the parity string for +non-adjacent sites, and applies the Jordan–Wigner prefix for odd one-site +operators. Its cache is separated by term placement and `coefficient_cutoff`. +Flat Z2 Symmray physical indices are supported as well; their duplicated +sector blocks are interpreted using the canonical fermion physical ordering. +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`: + +```python +from pepsy import Fermion + +fermion = Fermion(spinful=True, symmetry="U1U1", t=t, U=U) +vmc = pvmc.TorchFermionVMC( + peps, + fermion, + n_walkers=128, + contraction="boundary", # or "exact" / "ctmrg" / "hotrg" + chi=32, # required for an approximate contraction + dtype=torch.complex128, + seed=1, +) +result = vmc.estimate_observable( + burn_in=20, + n_measurements=8, + sweeps_between=2, +) +``` + +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 +supports are added to the proposal graph, including separated and +long-range terms, while the estimator retains the exact term placements and +fermionic parity strings. + +For one VMC update, call ``vmc.step(sr=True, ...)``; for repeated updates and +internal progress reporting, call ``vmc.optimize(n_steps, sr=True, +progress=True, ...)``. Boundary-MPS, CTMRG, and HOTRG use the differentiable +truncated Torch contraction path when ``chi`` is supplied. These approximate +SR updates are more sensitive to cutoff and gauge choices than exact +contraction, so monitor finite amplitudes and the SR linear-solve residual +while tuning them. + +For the chain-preserving interface, call +`vmc.estimate_observable(n_samples=..., n_discard_per_chain=...,` +`sweep_size=...)`. The older `burn_in`/`n_measurements`/ +`sweeps_between` form remains available. For spinful fermions, the selected +proposal is symmetry-aware rather than a literal single-site flip: `U1U1` +preserves both flavor counts, `U1` preserves total particle number, and `Z2` +preserves parity. + +With the older sweep-based interface, `target_effective_sample_size` makes +`n_measurements` a maximum rather than a fixed count. The driver checks both +ESS and (by default) `R-hat <= 1.05`; with several observables it waits until +all of them meet the target. `auto_thin=True` uses the current integrated +autocorrelation-time estimate as the spacing for later measurements: + +```python +result = vmc.estimate_observable( + burn_in=50, + n_measurements=200, # safety cap + sweeps_between=1, + target_effective_sample_size=128, + min_measurements=8, + rhat_threshold=1.05, + auto_thin=True, +) +print(result.n_measurements, result.effective_sample_size, result.energy_stderr) +``` + +This adaptive stopping path intentionally does not change the explicit +chain-preserving `n_samples=...` sampler interface, whose sample count remains +fixed and reproducible. + +`U1U1` uses moves that preserve `(N_up, N_down)`. For spinful `U1`, the +default proposal also includes single-site spin flips, so only +`N_up + N_down` is fixed. Spinful `Z2` adds empty/double pair toggles, so it +preserves parity without freezing total particle number. Pass an explicit +`terms={site_or_edge: operator}` mapping to measure a custom observable; site +labels may be PEPS labels or positional integers. If a dense PEPS does not +expose its charge sector, pass `sector=...` or valid `configs=...` explicitly. +The adapter refuses to apply an unverified local basis permutation. + +For an importance proposal from Pepsy's dense 2-norm BP sampler, use BP only +to propose configurations and let the torch PEPS model measure amplitudes and +local energies: + +```python +from pepsy.sampling import PepsBpSampler + +bp_sampler = PepsBpSampler(native_peps) +importance = driver.importance_energy_estimate( + bp_sampler, + n_samples=512, + sample_kwargs={"method": "mps", "chi": 32, "cutoff": 1e-10}, + progress=True, +) +print(importance.energy_mean, importance.effective_sample_size) +``` + +The importance result uses self-normalized weights +`|psi(x)|**2 / q_BP(x)` and reports the effective sample size, so a small BP +proposal overlap is visible instead of being mistaken for a precise VMC +estimate. + +For large nonlocal BP moves with an exact Metropolis correction, use the BP +sampler as an independence proposal. `TorchFermionVMC` infers the PEPS +encoding, symmetry, and sector automatically: + +```python +bp_mcmc = vmc.make_bp_sampler( + n_chains=64, + bp_sampler_kwargs={"max_iterations": 100}, + sample_kwargs={"method": "mps", "chi": 32, "cutoff": 1e-10}, + seed=3, +) +samples = bp_mcmc.sample( + n_samples=1024, + n_discard=100, + n_thin=4, +) +``` + +Each proposal uses + +`min(1, |psi(y)|**2 q_BP(x) / (|psi(x)|**2 q_BP(y)))`. + +The BP adapter supports four-state spinful fermion PEPS with `U1`, `U1U1`, +`Z2`, and `Z2Z2` physical symmetries. Since Quimb's current D2BP interface +samples binary output legs, a four-state physical leg is represented as two +occupation bits in a private dense BP copy. The original Symmray PEPS remains +block-sparse and is still used for Torch amplitudes and local energies. BP +proposals outside a fixed Fermion sector are rejected before they can enter a +chain. + The torch PEPS wrapper is validated for dense quimb PEPS and Symmray block-sparse fermionic PEPS. Symmray tensors are packed through their own `to_pytree` / `from_pytree` metadata via `quimb.tensor.pack`, then the numeric @@ -160,11 +759,11 @@ Symmray's batch-GPU fermionic-amplitude example. Do not assume the same flat supported through sparse block contractions, so their batching/performance profile should be benchmarked separately. -The default spinful encoding matches Pepsy/Symmray physical indices -`0=empty, 1=double, 2=up, 3=down`. Use -`FermionSiteEncoding.vmc_torch()` when consuming configs from -`sjdu10/vmc_torch`, where the convention is -`0=empty, 1=down, 2=up, 3=double`. +Torch VMC uses the physical-index order +`0=empty, 1=down, 2=up, 3=double`, matching the native four-sector fermionic +PEPS and the `sjdu10/vmc_torch` convention. Use +`FermionSiteEncoding.symmray()` only when interoperating with a caller that +explicitly uses its alternate legacy labels. Available connected-config helpers include `spinful_fermi_hubbard_connections`, `heisenberg_connections`, and @@ -271,13 +870,96 @@ driver = setup.make_driver( ) ``` -The fermionic batch path mirrors Symmray's batch GPU amplitude example: pack +### 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(...)`. + +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. + +```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 + +# Clean API: infer edges + PBC from the native terms, build the NetKet model. +setup = pvmc.build_fermion_vmc( + peps, fermion=fermion, terms=ham.terms, Lx=4, Ly=4, + contraction="ctmrg", chi=16, +) +``` + +```python +fermion = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0, V=0.5) +setup = pvmc.build_fermion_vmc( + peps, + fermion=fermion, + Lx=4, + Ly=4, + contraction="ctmrg", + chi=16, + observables={ + "double_occupancy": pvmc.standard_fermion_observables( + # any SpinOrbitalFermions Hilbert space with the same orbitals + )["double_occupancy"], + }, +) + +result = setup.optimize( + 2000, + driver="vmc_sr", + learning_rate=0.01, + energy_shift=-8.0 / 4, # arXiv:2511.02125 E/N - U/4 convention + per_site=setup.n_sites, +) +result.plot(per_site=setup.n_sites) +stats = setup.measure() # {name: nk.stats.Stats} +``` + +Two lower-level helpers back this path: + +- `netket_fermion_operator(hilbert, terms, *, constant=0.0, conserving="auto")` + builds any jittable NetKet fermion operator from symbolic + `(coefficient, ops)` terms, where each op is `(site, sz, dagger)` with + `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. +- `standard_fermion_observables(hilbert)` returns common observables + (`n_up`, `n_down`, `n_total`, `double_occupancy`) as a ready + `{name: operator}` mapping for `setup.measure(...)`. 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 fermionic `U1`, `Z2Z2`, and `U1U1` PEPS remain supported by the non-jitted batch amplitude function and by `TorchPEPSAmplitude`, but should use those sparse-block paths unless a matching flat backend is available and benchmarked. +For native sparse `U1U1` PEPS, `TorchPEPSAmplitude(..., +graded_torch=True, contraction="exact")` enables a separate fixed-shape Torch +projector. It densifies only the numerical leaves, compiles the U1U1 charge +transitions and Symmray Koszul phase masks once, and then contracts the +projected leaves under `torch.vmap`; it does not use the unsupported +`U1U1FermionicArrayFlat` conversion or a naive `to_dense()` contraction. +This opt-in path is exact and differentiable, but currently supports only +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 diff --git a/docs/development/package_layout.md b/docs/development/package_layout.md index 4de607d..48315e5 100644 --- a/docs/development/package_layout.md +++ b/docs/development/package_layout.md @@ -17,7 +17,7 @@ pepsy.solvers gradient and finite-difference parameter solvers pepsy.fitting local tensor fitting routines pepsy.optimizers high-level MPS, MPO, and PEPS optimizers pepsy.sampling MPS, vector, and PEPS samplers -pepsy.vmc optional NetKet/JAX variational Monte Carlo bridges +pepsy.vmc optional Torch and NetKet/JAX VMC adapters (`MCState` + `VMC` portable façade) pepsy._internal private formatting and utility helpers ``` @@ -29,7 +29,13 @@ Use the clearer namespaces for submodule imports: from pepsy.boundary import BdyMPS, contract_boundary from pepsy.optimizers import SweepOptimizer, MpsOptimizer, PepsOptimizer, SimpleUpdateGen from pepsy.operators import rx, rzz, gate -from pepsy.tensors import haar_random_state, ps_to_peps, ps_to_3dpeps, ps_to_mps +from pepsy.tensors import ( + haar_random_state, + ps_to_peps, + ps_to_3dpeps, + ps_to_mps, + ps_to_ttn, +) ``` When a leaf module is needed, import the new implementation path directly: diff --git a/examples/FermiHubbardJW/fermi_hubbard_jw.ipynb b/examples/FermiHubbardJW/fermi_hubbard_jw.ipynb deleted file mode 100644 index c2f0a1b..0000000 --- a/examples/FermiHubbardJW/fermi_hubbard_jw.ipynb +++ /dev/null @@ -1,246 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Fermi-Hubbard in the Jordan-Wigner spin picture\n", - "\n", - "A compact, end-to-end demonstration of Pepsy's Jordan-Wigner (bosonic) workflow\n", - "for the spinful Fermi-Hubbard model, using **one** Jordan-Wigner conversion for\n", - "both the Hamiltonian MPO and the time-evolution gates:\n", - "\n", - "1. build the `U1xU1` Fermi-Hubbard Hamiltonian on a chain;\n", - "2. ground state via `SymHamiltonian.to_mpo` + `SymDMRG2`;\n", - "3. `jw_bond_layout` to see which bonds are two-site-gate-able;\n", - "4. imaginary-time evolution with `jw_trotter_gates`, read out with `jw_energy`\n", - " (it converges to the DMRG ground energy);\n", - "5. a real-time quench measuring the average double occupancy.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-15T16:10:39.468342Z", - "iopub.status.busy": "2026-07-15T16:10:39.468125Z", - "iopub.status.idle": "2026-07-15T16:10:42.986960Z", - "shell.execute_reply": "2026-07-15T16:10:42.986399Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "adjacent bonds : [(0, 1), (1, 2), (2, 3)]\n", - "long-range bonds: []\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "SymDMRG2 ground energy: -1.11717241\n" - ] - } - ], - "source": [ - "import numpy as np\n", - "import pepsy\n", - "from pepsy.tensors import SymMPS, SymHamiltonian, site_charge_from_occupations\n", - "\n", - "# Spinful Fermi-Hubbard on an L-site chain, U1xU1 (per-spin particle number).\n", - "L, t, U = 4, 1.0, 8.0\n", - "edges = [(i, i + 1) for i in range(L - 1)]\n", - "occ = [(1, 0), (0, 1)] * (L // 2) # half filling, Sz = 0\n", - "sc = site_charge_from_occupations(occ)\n", - "\n", - "ham = SymHamiltonian.from_edges(\"fermi_hubbard_u1u1\", \"U1U1\", edges, t=t, U=U)\n", - "\n", - "# One Jordan-Wigner conversion feeds BOTH the MPO and the gates.\n", - "mpo = ham.to_mpo(L=L, compress=False)\n", - "\n", - "# Every chain bond is nearest-neighbour, so each is a valid two-site JW gate.\n", - "layout = ham.jw_bond_layout()\n", - "print(\"adjacent bonds :\", layout[\"adjacent\"])\n", - "print(\"long-range bonds:\", layout[\"long_range\"])\n", - "\n", - "# Ground state via symmetric DMRG (Jordan-Wigner bosonic representation).\n", - "gs = SymMPS.for_model(\"fermi_hubbard_u1u1\", L, bond_dim=1, site_charge=sc,\n", - " seed=1, dtype=\"complex128\")\n", - "opt = pepsy.SymDMRG2(mpo, gs, bond_dims=[1, 2, 4, 8, 16], cutoffs=[1e-10],\n", - " mixer=\"density_matrix\", compute_initial_energy=False)\n", - "opt.solve(max_sweeps=30, sweep_sequence=\"RL\", tol=1e-11)\n", - "E_gs = float(opt.energy)\n", - "print(f\"SymDMRG2 ground energy: {E_gs:.8f}\")\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Imaginary-time evolution converges to the ground state\n", - "\n", - "`ham.jw_trotter_gates(dt, imaginary=True)` builds the bosonic Jordan-Wigner\n", - "Trotter step from the *same* conversion as the MPO, and `ham.jw_energy(state)`\n", - "reads the energy back by summing local term expectations.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-15T16:10:42.988953Z", - "iopub.status.busy": "2026-07-15T16:10:42.988740Z", - "iopub.status.idle": "2026-07-15T16:10:45.222055Z", - "shell.execute_reply": "2026-07-15T16:10:45.221433Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "step 50 E = -1.04246841\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "step 100 E = -1.09909621\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "step 150 E = -1.11315117\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "step 200 E = -1.11623032\n", - "\n", - "JW imaginary-time energy: -1.11623032\n", - "SymDMRG2 ground energy : -1.11717241\n", - "difference : 9.42e-04\n" - ] - } - ], - "source": [ - "# Imaginary-time evolution in the Jordan-Wigner (bosonic spin) picture.\n", - "psi = SymMPS.for_model(\"fermi_hubbard_u1u1\", L, bond_dim=1, site_charge=sc,\n", - " seed=2, dtype=\"complex128\", fermionic=False)\n", - "step = ham.jw_trotter_gates(0.05, imaginary=True, order=2)\n", - "\n", - "for k in range(200):\n", - " psi.apply_gates(step, method=\"direct\", max_bond=16, cutoff=1e-10, normalize=True)\n", - " if (k + 1) % 50 == 0:\n", - " print(f\"step {k + 1:3d} E = {ham.jw_energy(psi):.8f}\")\n", - "\n", - "print(f\"\\nJW imaginary-time energy: {ham.jw_energy(psi):.8f}\")\n", - "print(f\"SymDMRG2 ground energy : {E_gs:.8f}\")\n", - "print(f\"difference : {abs(ham.jw_energy(psi) - E_gs):.2e}\")\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Real-time quench: average double occupancy\n", - "\n", - "Starting from the Neel product state (no double occupancy), real-time hopping\n", - "builds up double occupancy. Local observables use the state's symmetry-aware\n", - "`measure`.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-15T16:10:45.223705Z", - "iopub.status.busy": "2026-07-15T16:10:45.223574Z", - "iopub.status.idle": "2026-07-15T16:10:45.532837Z", - "shell.execute_reply": "2026-07-15T16:10:45.531672Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " t \n", - "0.00 0.00000\n", - "0.50 0.04336\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "1.00 0.04244\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "1.50 0.04314\n" - ] - } - ], - "source": [ - "from pepsy.tensors import symm_operator_from_dense, default_physical_sectors\n", - "\n", - "# Double-occupancy projector |up,down>\")\n", - "print(f\"{0.0:.2f} {mean_double_occ(psi_t):.5f}\")\n", - "for k in range(30):\n", - " psi_t.apply_gates(rt, method=\"direct\", max_bond=16, cutoff=1e-10)\n", - " if (k + 1) % 10 == 0:\n", - " print(f\"{(k + 1) * 0.05:.2f} {mean_double_occ(psi_t):.5f}\")\n" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/examples/FermiHubbardQMera/fermion_qmera_energy.py b/examples/FermiHubbardQMera/fermion_qmera_energy.py deleted file mode 100644 index e253176..0000000 --- a/examples/FermiHubbardQMera/fermion_qmera_energy.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Optimize a small native Symmray Fermi-Hubbard qMERA energy.""" - -from __future__ import annotations - -import argparse -import importlib.util -import os -import pprint - -import numpy as np - -os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/pepsy_numba_cache") - -from pepsy import Fermion -from pepsy.backends import backend_torch -from pepsy.optimizers.mera import ( - QMeraBuilder, - QMeraSymmrayFermionBackend, - symmray_fermion_gate_registry, -) - - -def _as_float(value): - """Convert a scalar-like NumPy or Torch value to a real float.""" - detach = getattr(value, "detach", None) - if callable(detach): - value = detach().cpu().numpy() - return float(np.real(np.asarray(value))) - - -def build_demo(*, steps=0): - """Build and optionally optimize a two-site spinful qMERA.""" - if steps and importlib.util.find_spec("torch") is None: - raise RuntimeError("--steps requires the optional Torch dependency.") - - array_backend = None - if steps: - import torch - - array_backend = backend_torch(dtype=torch.complex128) - - backend = QMeraSymmrayFermionBackend(to_backend=array_backend) - registry = symmray_fermion_gate_registry(backend=backend) - - # The register order is (0, up), (1, up), (0, down), (1, down). This - # product state has one up fermion on site 0 and one down fermion on site 1. - occupations = {0: 1, 1: 0, 2: 0, 3: 1} - - def product_state_factory(schedule, sites, **kwargs): - return backend.product_state( - schedule, - sites, - occupations=occupations, - **kwargs, - ) - - builder = QMeraBuilder( - shape=2, - site_modes=("up", "down"), - mode_order="mode-major", - gate_registry=registry, - array_backend=array_backend, - disentangler={"block_size": 2, "circuit_depth": 0}, - isometry={ - "block_size": 2, - "circuit_depth": 1, - "gate_family": "symmray-fsim", - }, - max_layers=1, - seed=7, - 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) - optimizer = builder.fermion_parametric_optimizer( - fermion, - energy_per_site=False, - ) - initial = optimizer.loss(energy_per_site=False) - report = { - "num_terms": len(terms), - "initial_energy": _as_float(initial), - } - - if steps: - result = optimizer.run( - solver="torch-adam", - n_steps=steps, - log_every=1, - options={"lr": 0.01}, - ) - report["energy_history"] = [float(value) for value in result.history] - report["final_energy"] = float(result.history[-1]) - return report - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--steps", - type=int, - default=0, - help="Optional number of native Torch Adam steps.", - ) - args = parser.parse_args() - pprint.pp(build_demo(steps=args.steps)) - - -if __name__ == "__main__": - main() diff --git a/examples/FermiHubbardUnified/fermion_workflow.py b/examples/FermiHubbardUnified/fermion_workflow.py deleted file mode 100644 index 683ea41..0000000 --- a/examples/FermiHubbardUnified/fermion_workflow.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Use one ``Fermion`` model for site-native MPS and qMERA workflows.""" - -from __future__ import annotations - -import pepsy -from pepsy.optimizers.mera import QMeraGeometry - - -def build_workflow(length=3): - """Build matching site terms, MPS gates, and qMERA mode terms.""" - fermion = pepsy.Fermion( - spinful=True, - symmetry="U1U1", - t=1.0, - U=4.0, - mu=0.0, - ) - edges = tuple((site, site + 1) for site in range(length - 1)) - - # Four-state site-native terms and a canonical stream for MPS evolution. - site_terms = fermion.local_terms(edges) - gate_stream = fermion.gate_stream( - edges, - dt=0.01, - sites=range(length), - order=2, - ) - - # qMERA deliberately expands each physical site into two two-state modes. - geometry = QMeraGeometry( - shape=length, - site_modes=("up", "down"), - ) - qmera_terms = fermion.local_terms(geometry, layout="qmera") - - return { - "fermion": fermion, - "site_terms": site_terms, - "gate_stream": gate_stream, - "geometry": geometry, - "qmera_terms": qmera_terms, - } - - -if __name__ == "__main__": - workflow = build_workflow() - print(f"site terms: {len(workflow['site_terms'])}") - print(f"MPS gates: {len(workflow['gate_stream'])}") - print(f"qMERA mode terms: {len(workflow['qmera_terms'])}") diff --git a/examples/QMeraEnergy/qmera_energy.py b/examples/QMeraEnergy/qmera_energy.py deleted file mode 100644 index 303ef5d..0000000 --- a/examples/QMeraEnergy/qmera_energy.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Schedule-first qMERA local-energy example.""" - -from __future__ import annotations - -import argparse -import importlib.util -import os -import pprint - -import numpy as np - -os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") - -from pepsy.optimizers import QMeraBuilder, build_qmera_contraction_optimizer - - -def zz_operator(): - """Return a two-site ZZ operator in quimb gate tensor shape.""" - z_op = np.diag([1.0, -1.0]) - return np.kron(z_op, z_op).reshape(2, 2, 2, 2) - - -def zz_chain_hamiltonian(length): - """Return nearest-neighbor ZZ terms on an open chain.""" - h2 = zz_operator() - return {(site, site + 1): h2 for site in range(length - 1)} - - -def as_float(value): - """Convert NumPy/Torch/JAX scalar-like values to a display float.""" - detach = getattr(value, "detach", None) - if callable(detach): - value = detach().cpu().numpy() - return float(np.real(np.asarray(value))) - - -def run_demo(*, steps=0, cache_dir=False): - """Build a small qMERA, evaluate compiled local-cone energy, optionally train.""" - builder = QMeraBuilder( - shape=4, - gate_family="rxx", - isometry_gate_family="rzz", - seed=7, - param_scale=0.05, - ) - schedule = builder.build_schedule() - params = builder.initialize_parameters(schedule) - hamiltonian = zz_chain_hamiltonian(schedule.geometry.num_sites) - chunks = builder.parametric_lightcone_chunks(hamiltonian, schedule) - contraction_opt = build_qmera_contraction_optimizer( - directory=cache_dir, - max_repeats=4, - max_time=0.25, - progbar=False, - ) - - energy = builder.parametric_loss( - params, - schedule=schedule, - chunks=chunks, - energy_per_site=False, - contraction_opt=contraction_opt, - ) - compiled = builder.compile_parametric_lightcones( - schedule=schedule, - chunks=chunks, - contraction_opt=contraction_opt, - ) - compiled_energy = builder.compiled_parametric_loss( - params, - schedule=schedule, - compiled_chunks=compiled, - energy_per_site=False, - ) - - report = { - "num_sites": schedule.geometry.num_sites, - "num_layers": len(schedule.layers), - "num_gates": schedule.num_gates, - "num_terms": len(chunks), - "max_lightcone_gates": max(chunk.num_gates for chunk in chunks), - "energy": as_float(energy), - "compiled_energy": as_float(compiled_energy), - } - - if steps: - if importlib.util.find_spec("torch") is None: - report["torch_optimization"] = "skipped: torch is not installed" - else: - opt = builder.parametric_optimizer( - hamiltonian, - schedule=schedule, - chunks=chunks, - parameters=params, - energy_per_site=False, - contraction_opt=contraction_opt, - ) - result = opt.run( - solver="torch-adam", - n_steps=steps, - log_every=1, - options={"lr": 0.05}, - compiled=True, - ) - report["torch_initial_energy"] = float(result.history[0]) - report["torch_final_energy"] = float(result.history[-1]) - - return report - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--steps", - type=int, - default=0, - help="Optional number of Torch Adam steps to run.", - ) - parser.add_argument( - "--cache-dir", - default=False, - help="Optional cotengra path-cache directory. Omit for memory-only paths.", - ) - args = parser.parse_args() - pprint.pp(run_demo(steps=args.steps, cache_dir=args.cache_dir)) - - -if __name__ == "__main__": - main() diff --git a/examples/RelayBP/odd_cycle_stress.py b/examples/RelayBP/odd_cycle_stress.py deleted file mode 100644 index b40506d..0000000 --- a/examples/RelayBP/odd_cycle_stress.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Stress Relay D1BP on polarized, near-deterministic odd cycles. - -This is a convergence stress test, not an accuracy claim. The networks are -positive but strongly frustrated, so their exact scalar contractions remain -available while ordinary loopy BP has a poor uncontrolled approximation. -""" - -from __future__ import annotations - -import warnings - -import numpy as np -import quimb.tensor as qtn -from quimb.tensor.belief_propagation import D1BP - -from pepsy.bp import one_norm_bp, relay_bp - - -def _odd_antiferromagnetic_cycle(epsilon): - """Return a positive triangle whose edge factors nearly flip a bit.""" - factor = np.array([[epsilon, 1.0], [1.0, epsilon]]) - return qtn.TensorNetwork( - [ - qtn.Tensor(factor, inds=("ab", "ca")), - qtn.Tensor(factor, inds=("ab", "bc")), - qtn.Tensor(factor, inds=("bc", "ca")), - ] - ) - - -def _polarized_messages(tn): - """Choose a deterministic non-fixed initial message for every edge end.""" - return { - key: np.array([1.0, 0.0]) - for key in D1BP(tn, update="parallel").messages - } - - -def run_stress_cases(): - """Run two exact-reference cases where parallel D1BP stalls.""" - records = [] - for epsilon in (1e-3, 1e-2): - tn = _odd_antiferromagnetic_cycle(epsilon) - exact = float(tn.contract()) - initial = _polarized_messages(tn) - common = { - "method": "d1bp", - "init_messages": initial, - "update": "parallel", - "diis": False, - "max_iterations": 100, - "tol": 1e-10, - } - # The stall is expected and recorded below, so avoid duplicating it as - # a Quimb warning in this runnable comparison. - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message="Belief propagation did not converge.*", - category=UserWarning, - ) - plain = one_norm_bp(tn, **common) - relay = relay_bp( - tn, - **common, - num_relays=5, - memory_first_leg=True, - gamma_range=(0.2, 0.8), - seed=0, - ) - relay_estimate = float(relay.contract()) - records.append( - { - "epsilon": epsilon, - "exact": exact, - "plain_converged": plain.converged, - "plain_max_mdiff": plain.max_mdiff, - "relay_converged": relay.converged, - "relay_iterations": relay.iterations, - "relay_num_legs": relay.num_legs_run, - "relay_max_mdiff": relay.max_mdiff, - "relay_relative_error": abs(relay_estimate - exact) / abs(exact), - } - ) - return records - - -def main(): - """Print strict convergence and exact-reference accuracy for each case.""" - for record in run_stress_cases(): - print( - f"epsilon={record['epsilon']:.0e} " - f"plain_converged={record['plain_converged']} " - f"plain_max_mdiff={record['plain_max_mdiff']:.3e} " - f"relay_converged={record['relay_converged']} " - f"relay_iterations={record['relay_iterations']} " - f"relay_relative_error={record['relay_relative_error']:.3e}" - ) - - -if __name__ == "__main__": - main() diff --git a/examples/RelayBP/simple_update_relay_comparison.py b/examples/RelayBP/simple_update_relay_comparison.py deleted file mode 100644 index bdb3e1c..0000000 --- a/examples/RelayBP/simple_update_relay_comparison.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Compare raw, simple-update-initialized, and Relay D1BP to exact contraction. - -Run from the repository root after installing pepsy and quimb:: - - python examples/RelayBP/simple_update_relay_comparison.py - -The network is deliberately tiny so that its exact contraction is an -independent reference. It is still loopy, so D1BP is an approximation: the -example records accuracy rather than treating agreement with exact contraction -as a convergence criterion. -""" - -from __future__ import annotations - -import quimb.tensor as qtn - -from pepsy.bp import one_norm_bp, run_d1bp_from_simple_update_gauges - - -def run_comparison(): - """Run the three D1BP initializations and return reproducible metrics.""" - tn = qtn.TN2D_classical_ising_partition_function(3, 3, beta=0.2) - exact = float(tn.contract(optimize="auto-hq")) - - # This is Quimb's actual simple-update routine. It conditions the tensors - # while moving the diagonal bond environments into ``gauges``. - core = tn.copy() - gauges = {} - core.gauge_all_simple_(gauges=gauges, max_iterations=100, tol=1e-12) - rebuilt = core.copy() - rebuilt.gauge_simple_insert(gauges) - su_representation_error = abs(float(rebuilt.contract()) - exact) / abs(exact) - - runs = { - "plain_d1bp": one_norm_bp( - tn, - method="d1bp", - max_iterations=1000, - tol=1e-10, - ), - "su_initialized_d1bp": run_d1bp_from_simple_update_gauges( - core, - gauges, - run_opts={"max_iterations": 1000, "tol": 1e-10}, - ), - "su_initialized_relay_d1bp": run_d1bp_from_simple_update_gauges( - core, - gauges, - use_relay=True, - run_opts={"max_iterations": 1000, "tol": 1e-10}, - relay_opts={ - "num_relays": 3, - # A plain first leg preserves the ordinary fixed point on an - # easy instance; the other legs exercise per-node memory. - "memory_first_leg": False, - "gamma_range": (0.1, 0.2), - "seed": 0, - }, - ), - } - - records = {} - for name, result in runs.items(): - estimate = float(result.contract()) - records[name] = { - "converged": result.converged, - "quimb_converged": result.quimb_converged, - "iterations": result.iterations, - "num_legs": result.num_legs_run, - "max_mdiff": result.max_mdiff, - "estimate": estimate, - "relative_error": abs(estimate - exact) / abs(exact), - } - - return { - "exact": exact, - "simple_update_representation_error": su_representation_error, - "runs": records, - } - - -def main(): - """Print the convergence and exact-reference accuracy comparison.""" - comparison = run_comparison() - print(f"exact contraction: {comparison['exact']:.12g}") - print( - "simple-update representation relative error: " - f"{comparison['simple_update_representation_error']:.3e}" - ) - print( - "name converged legs iterations " - "relative error" - ) - for name, record in comparison["runs"].items(): - print( - f"{name:28} {str(record['converged']):9} " - f"{record['num_legs']:5d} {record['iterations']:11d} " - f"{record['relative_error']:.3e}" - ) - - -if __name__ == "__main__": - main() diff --git a/plot_policy.md b/plot_policy.md new file mode 100644 index 0000000..1a48a48 --- /dev/null +++ b/plot_policy.md @@ -0,0 +1,96 @@ +# Plot Policy (guidance for clear, decent figures) + +Guidance for agents producing matplotlib figures in this workspace (pepsy +examples, notebooks, benchmark figures). + +This is **guidance, not a rigid template.** The goal is simply: produce clear, +legible, good-looking plots. Use judgment and adapt to the data — do not +mechanically copy every setting below. The concrete style in the last section +is a recommended starting point (matching the owner's taste), not a mandate. + +## What "decent and clear" means (the actual requirements) + +Aim for these principles on every figure: + +- **Readable text.** Axis labels and ticks large enough to read in a paper/slide + (labels ~14–17 pt, ticks ~12–13 pt). Don't ship default-tiny fonts. +- **Labeled axes.** Every axis has a label with units/meaning; use LaTeX math + (`r'$\delta Z^2$'`) where it helps. Add a legend whenever there is more than + one series. +- **Sensible layout.** Reasonable figure size (e.g. ~10×6 single panel), a light + grid to guide the eye, and `bbox_inches='tight'` / `tight_layout()` so nothing + is clipped. +- **Good color.** For a parameter sweep, sample a perceptual colormap + (`plasma`, `viridis`, `cividis`) so the series read in order. Avoid clashing + hand-picked colors and don't rely on the raw default cycle for many curves. +- **Distinguish reference from data.** Plot an exact/analytic/reference curve + distinctly (e.g. solid black, no marker) so it stands out from approximate data. +- **Right scale.** Use a log axis for errors / quantities spanning orders of + magnitude; set limits when they improve readability. +- **Show uncertainty.** If you have error bars, draw them (`errorbar(..., capsize=...)`). +- **Save reusably.** For figures worth keeping, save a vector format (PDF/SVG) + and, if a raster is needed, PNG at `dpi=300`. Keep image files out of `src/`. + +If a simple plot communicates the point clearly, a plain `fig, ax` with labels, +a legend, and a grid is entirely fine — clarity beats decoration. + +## Recommended house style (a good default, not required) + +This reproduces the owner's preferred look; use it as a convenient starting +point and tweak freely. + +```python +import numpy as np +import matplotlib.pyplot as plt + +plt.rcParams.update({ + 'axes.titlesize': 18, + 'axes.labelsize': 14, + 'lines.linewidth': 2.1, + 'lines.marker': 'o', + 'xtick.labelsize': 13, + 'ytick.labelsize': 13, + 'legend.fontsize': 13, + 'figure.figsize': (10, 6), +}) +plt.grid(color='gray', linestyle='--', linewidth=0.5, alpha=0.7) + +# One curve per swept parameter, colored from a perceptual colormap. +bonds = [4, 8, 16, 32, 64] +colors = plt.cm.plasma(np.linspace(0.1, 0.9, len(bonds))) + +# Optional exact/reference curve: solid black, no marker. +# plt.plot(x, exact, label='exact', color='black', linestyle='-', marker='', alpha=0.8) + +for i, bnd in enumerate(bonds): + plt.plot( + x, y[i], + label=rf'$D={bnd}$', + color=colors[i], + marker='>', markersize=10, + linewidth=2, alpha=0.5, linestyle='-', + ) + # with uncertainties: plt.errorbar(x, y[i], yerr=err[i], capsize=5, ...) + +plt.xlabel(r'depth', fontsize=17, labelpad=10) +plt.ylabel(r'$Z^2$', fontsize=17, labelpad=10) +# plt.yscale('log') # for errors / wide dynamic range +plt.legend(loc='best', frameon=True, shadow=True, fontsize=13) +plt.savefig('figure.png', dpi=300, bbox_inches='tight') +plt.savefig('figure.pdf', bbox_inches='tight') +plt.show() +``` + +In notebooks, crisp inline output helps: `%config InlineBackend.figure_formats = ['svg']`. + +For stacked diagnostics, `plt.subplots(n, 1, sharex=True)` with the same +label/grid conventions per axis works well; mark regime boundaries with +`ax.axvline(..., color='0.35', linestyle='--', linewidth=1)`. + +## Minimal checklist + +- [ ] Labeled axes (units/meaning), legend if >1 series, readable font sizes. +- [ ] Light grid; nothing clipped (`tight` layout / bbox). +- [ ] Perceptual colormap for parameter sweeps; reference curve distinct. +- [ ] Log scale for errors / wide ranges; error bars drawn when available. +- [ ] Keep-worthy figures saved as vector (PDF/SVG) + PNG `dpi=300`, not under `src/`. diff --git a/pyproject.toml b/pyproject.toml index daa6d0e..cff6389 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pepsy" -version = "0.2.0" +version = "0.3.0" description = "Boundary-MPS tools for PEPS norm contraction and DMRG fitting" readme = "README.md" requires-python = ">=3.10" @@ -23,6 +23,7 @@ dependencies = [ torch = ["torch"] solvers = ["scipy", "nlopt"] viz = ["matplotlib"] +layout = ["nevergrad"] docs = ["sphinx>=7", "myst-parser>=2", "pydata-sphinx-theme>=0.15", "nbsphinx>=0.9"] dev = ["pytest", "pyflakes", "nbclient", "nbconvert"] vmc = ["jax", "flax", "optax", "netket", "symmray"] diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index 0b4aa3f..aadcb8e 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -131,6 +131,7 @@ "SimpleUpdateGen": ".optimizers", "SymDMRG2": ".optimizers", "SweepOptimizer": ".optimizers", + "TreeLayoutFinder": ".optimizers", "TreeOptimizer": ".optimizers", "TreeTensorNetwork": ".optimizers", "compile_stim_circuit": ".optimizers", @@ -147,6 +148,8 @@ "sample_stim_circuits": ".optimizers", "sample_trajectory_stream": ".optimizers", "sample_coalesced_bits": ".optimizers", + "FermionConfigurationEncoding": ".sampling", + "MpsDiagonalEstimate": ".sampling", "MpsBatchSampleResult": ".sampling", "MpsSampleResult": ".sampling", "MpsSampler": ".sampling", @@ -163,6 +166,7 @@ "optimize_packed_params": ".solvers", "OneDMap": ".tensors", "Fermion": ".tensors", + "FermionLatticeSetup": ".tensors", "SpinfulFermion": ".tensors", "SpinfulFermionHubbard": ".tensors", "SymmFermions": ".tensors", @@ -192,14 +196,19 @@ "fermion_interaction_param_gen": ".tensors", "expec_mpo": ".tensors", "haar_random_state": ".tensors", + "hrs_to_mps": ".tensors", + "hrs_to_peps": ".tensors", + "hrs_to_ttn": ".tensors", "hrps_to_mps": ".tensors", "hrps_to_peps": ".tensors", + "hrps_to_ttn": ".tensors", "id_to_mpo": ".tensors", "id_to_pepo": ".tensors", "measure_obs": ".tensors", "ps_to_3dpeps": ".tensors", "ps_to_mpo": ".tensors", "ps_to_mps": ".tensors", + "ps_to_ttn": ".tensors", "ps_to_pepo": ".tensors", "ps_to_peps": ".tensors", "random_haar_qubit": ".tensors", @@ -348,6 +357,7 @@ "SimpleUpdateGen", "SymDMRG2", "SweepOptimizer", + "TreeLayoutFinder", "TreeOptimizer", "TreeTensorNetwork", "compile_stim_circuit", @@ -364,6 +374,8 @@ "sample_stim_circuits", "sample_trajectory_stream", "sample_coalesced_bits", + "FermionConfigurationEncoding", + "MpsDiagonalEstimate", "MpsBatchSampleResult", "MpsSampleResult", "MpsSampler", @@ -376,6 +388,7 @@ "FDSolver", "OneDMap", "Fermion", + "FermionLatticeSetup", "SpinfulFermion", "SpinfulFermionHubbard", "SymmFermions", @@ -400,14 +413,19 @@ "fermion_interaction_param_gen", "expec_mpo", "haar_random_state", + "hrs_to_mps", + "hrs_to_peps", + "hrs_to_ttn", "hrps_to_mps", "hrps_to_peps", + "hrps_to_ttn", "id_to_mpo", "id_to_pepo", "measure_obs", "ps_to_3dpeps", "ps_to_mpo", "ps_to_mps", + "ps_to_ttn", "ps_to_pepo", "ps_to_peps", "random_haar_qubit", @@ -518,12 +536,13 @@ 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, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeOptimizer, run_stabilizer_mps_stream - from .sampling import MpsBatchSampleResult, MpsSampleResult, MpsSampler, PEPSSampleResult, PepsBpSampler, TreeBatchSampleResult, TreeSampleResult, TreeSampler, VecSampler + from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeLayoutFinder, TreeOptimizer, run_stabilizer_mps_stream + from .sampling import FermionConfigurationEncoding, MpsDiagonalEstimate, MpsBatchSampleResult, MpsSampleResult, MpsSampler, PEPSSampleResult, PepsBpSampler, TreeBatchSampleResult, TreeSampleResult, TreeSampler, VecSampler from .solvers import FDSolver from .tensors import ( OneDMap, Fermion, + FermionLatticeSetup, SpinfulFermion, SpinfulFermionHubbard, SymmFermions, @@ -548,14 +567,19 @@ def __getattr__(name): fermion_interaction_param_gen, expec_mpo, haar_random_state, + hrs_to_mps, + hrs_to_peps, + hrs_to_ttn, hrps_to_mps, hrps_to_peps, + hrps_to_ttn, id_to_mpo, id_to_pepo, measure_obs, ps_to_3dpeps, ps_to_mpo, ps_to_mps, + ps_to_ttn, ps_to_pepo, ps_to_peps, random_haar_qubit, diff --git a/src/pepsy/boundary/metrics.py b/src/pepsy/boundary/metrics.py index 4032f4f..051f9d0 100644 --- a/src/pepsy/boundary/metrics.py +++ b/src/pepsy/boundary/metrics.py @@ -67,6 +67,25 @@ def _warn_nonstandard_physical_outer_inds(tn, role): ) +def _uses_symmray_arrays(tn): + """Return whether a tensor network stores Symmray block-sparse arrays.""" + tensor_map = getattr(tn, "tensor_map", None) + tensors = tensor_map.values() if tensor_map else tn + try: + iterator = iter(tensors) + except TypeError: + iterator = () + for tensor in iterator: + data = getattr(tensor, "data", None) + if data is None: + continue + if type(data).__module__.split(".", maxsplit=1)[0] == "symmray": + return True + if hasattr(data, "blocks") and hasattr(data, "apply_to_arrays"): + return True + return False + + def _to_python_scalar(value): """Convert backend scalar-like objects (torch/numpy) to python scalar.""" obj = value @@ -921,7 +940,9 @@ def peps_normalize( stripped representation before mutating ``p`` and emits a ``RuntimeWarning``. balance_bonds : bool, default=True - If ``True``, call ``balance_bonds_()`` after rescaling the state. + If ``True``, call ``balance_bonds_()`` after rescaling a dense state. + Symmray block-sparse states always skip this step because Quimb's bond + balancer currently needs a contraction unsupported by Symmray. Returns ------- @@ -979,7 +1000,7 @@ def peps_normalize( _ensure_finite_norm_value(cost) old_norm = _format_scaled_output(cost, strip_exponent=strip_exponent) _normalize_by_scaled_norm(ket_tagged, cost) - if balance_bonds: + if balance_bonds and not _uses_symmray_arrays(ket_tagged): ket_tagged.balance_bonds_() return old_norm diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index ac4f0bd..6d72684 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -2,8 +2,9 @@ This module implements the ``P + Q`` expansion of - G. Park, J. Gray, and G. K.-L. Chan, *Loop Series Expansions for Tensor - Networks*, Phys. Rev. Research 7, 013123 (2025), arXiv:2409.03108. + G. Evenbly, N. Pancotti, A. Milsted, J. Gray, and G. K.-L. Chan, *Loop + series expansions for tensor networks*, Phys. Rev. Research 8, 013245 + (2026), arXiv:2409.03108. It is deliberately separate from :mod:`pepsy.bp.cluster`. A loop-cluster expansion is indexed by tensor regions and inclusion--exclusion counting diff --git a/src/pepsy/operators/gates.py b/src/pepsy/operators/gates.py index ee1b410..362e505 100644 --- a/src/pepsy/operators/gates.py +++ b/src/pepsy/operators/gates.py @@ -75,6 +75,47 @@ def _stop_gradient(x): return x +def _is_block_sparse_array(value): + """Return whether ``value`` looks like a Symmray block-sparse array.""" + return hasattr(value, "blocks") and hasattr(value, "indices") + + +def _has_block_sparse_data(tn): + """Return whether any tensor in ``tn`` stores block-sparse data.""" + return any( + _is_block_sparse_array(tensor.data) + for tensor in getattr(tn, "tensors", ()) + ) + + +def _select_simple_update_contract(tn, gate, contract): + """Select a split strategy that preserves block-sparse fusion metadata. + + The reduced two-site split path is often faster and works for many cases, + including numerous Symmray workloads. However, some block-sparse MPS/PEPS + setups have shown fusion-metadata edge cases with reduced splitting. The + full ``split`` path is slower but is used as a conservative fallback for + Symmray block-sparse data. Dense tensor data keeps the faster reduced path + by default. + """ + if contract is not None: + contract = str(contract).strip().lower() + if contract == "auto": + contract = None + elif contract not in {"split", "reduce-split"}: + raise ValueError( + "gate_simple() contract must be 'auto', 'split', or " + "'reduce-split'." + ) + + if contract is None: + if _has_block_sparse_data(tn) or _is_block_sparse_array(gate): + return "split" + return "reduce-split" + + return contract + + def rx(theta): """Return a one-qubit RX gate for angle ``theta``. @@ -1268,7 +1309,15 @@ def _symmray_dense_index_map_from_chargemap(chargemap): return index_map -def _symmray_dense_gate_from_site_maps(tn, gate, output_sites, input_sites, ind_id): +def _symmray_dense_gate_from_site_maps( + tn, + gate, + output_sites, + input_sites, + ind_id, + *, + inferred_converter=None, +): """Convert a dense site gate to a Symmray array using live site sectors.""" sample = _symmray_sample_data_from_tn(tn) if sample is None: @@ -1292,16 +1341,38 @@ def _symmray_dense_gate_from_site_maps(tn, gate, output_sites, input_sites, ind_ if symmetry is not None: kwargs["symmetry"] = symmetry + # Symmray product symmetries (for example ``U1U1`` and ``Z2Z2``) use + # tuple-valued charges. ``charge=0`` is only valid for scalar symmetries; + # use the neutral charge with the same shape as the live physical sectors. + sample_charge = next(iter(index_maps[0].values()), 0) + zero_charge = ( + tuple(0 for _ in sample_charge) + if isinstance(sample_charge, tuple) + else 0 + ) + + dense_gate = np.asarray(gate) + if inferred_converter is not None: + dense_gate = inferred_converter(dense_gate) + return array_cls.from_dense( - np.asarray(gate), + dense_gate, index_maps=tuple(index_maps), duals=(False,) * nout + (True,) * nin, - charge=0, + charge=zero_charge, **kwargs, ) -def _symmray_swap_gate_for_site_pair(tn, site_a, site_b, *, ind_id=None, dtype="complex128"): +def _symmray_swap_gate_for_site_pair( + tn, + site_a, + site_b, + *, + ind_id=None, + dtype="complex128", + inferred_converter=None, +): """Return a Symmray block-sparse SWAP for two live physical site legs.""" sample = _symmray_sample_data_from_tn(tn) if sample is None: @@ -1316,6 +1387,7 @@ def _symmray_swap_gate_for_site_pair(tn, site_a, site_b, *, ind_id=None, dtype=" output_sites=(site_b, site_a), input_sites=(site_a, site_b), ind_id=ind_id, + inferred_converter=inferred_converter, ) @@ -1345,6 +1417,7 @@ def _swap_gate_for_site_pair( site_b, ind_id=ind_id, dtype=dtype, + inferred_converter=inferred_converter, ) if symmray_swap is not None: return symmray_swap @@ -1911,6 +1984,7 @@ def gate_simple( max_bond=None, cutoff=1e-12, cutoff_mode="rsum2", + contract=None, sequence="auto", path_canonize=False, path_canonize_distance=1, @@ -1970,6 +2044,12 @@ def gate_simple( cutoff_mode : str, optional Cutoff mode passed to ``gate_simple_`` (e.g. ``'rsum2'``, ``'rel'``). Default ``'rsum2'``. + 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 + Symmray data and ``"reduce-split"`` for dense data. Explicit + ``contract="reduce-split"`` remains supported for block-sparse runs + when that path is validated for a workload. sequence : object, optional Long-range routing preference. Defaults to ``"auto"``, which chooses a shortest 2D/3D path with the smallest current bond-dimension bottleneck. @@ -1998,14 +2078,20 @@ def gate_simple( ) which_default = _normalize_gate_which(which) - gate_opts = { + gate_opts_base = { "cutoff": cutoff, "cutoff_mode": cutoff_mode, } if max_bond is not None: - gate_opts["max_bond"] = int(max_bond) + gate_opts_base["max_bond"] = int(max_bond) for gate_payload, where_payload, which_payload in entries: + gate_opts = dict(gate_opts_base) + gate_opts["contract"] = _select_simple_update_contract( + tn_work, + gate_payload, + contract, + ) if _is_explicit_index_where(where_payload): raise ValueError( "gate_simple() requires lattice site coordinates, not explicit " @@ -2363,7 +2449,11 @@ def _gate_simple_one_with_current_site_ind_id( # Non-adjacent: route through a SWAP chain. Each SWAP is built from the # live physical dimensions because routed mixed-dimensional sites exchange # their physical index sizes as they move along the path. - backend_sample = resolve_backend_sample_data_from_tn(tn_work) + symmray_sample = _symmray_sample_data_from_tn(tn_work) + if symmray_sample is not None: + backend_sample = next(iter(symmray_sample.blocks.values()), None) + else: + backend_sample = resolve_backend_sample_data_from_tn(tn_work) inferred_converter = infer_backend_converter_from_sample( backend_sample, cast_complex_to_real=True, diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index e6323c8..a779e11 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -65,7 +65,7 @@ ) from .sym_dmrg import SymDMRG2 from .sweep import SweepOptimizer -from .tree import TreeOptimizer, TreeTensorNetwork +from .tree import TreeLayoutFinder, TreeOptimizer, TreeTensorNetwork __all__ = [ "CoalescedMeasurementRecord", @@ -113,6 +113,7 @@ "SimpleUpdateGen", "SymDMRG2", "SweepOptimizer", + "TreeLayoutFinder", "TreeOptimizer", "TreeTensorNetwork", "compile_stim_circuit", diff --git a/src/pepsy/optimizers/energy/peps.py b/src/pepsy/optimizers/energy/peps.py index fcb635a..673d0e3 100644 --- a/src/pepsy/optimizers/energy/peps.py +++ b/src/pepsy/optimizers/energy/peps.py @@ -434,6 +434,75 @@ def _state_uses_symmray(state): return True return False + @staticmethod + def _is_symmray_array(value): + """Return whether ``value`` is a native Symmray block array.""" + return hasattr(value, "blocks") and hasattr(value, "indices") + + @classmethod + def _terms_use_symmray(cls, terms): + return all(cls._is_symmray_array(term) for term in terms.values()) + + @staticmethod + def _term_sites(state, where): + """Normalize a local-term key to the corresponding PEPS sites.""" + has_site = getattr(state, "has_site", None) + if callable(has_site) and has_site(where): + return (where,) + if not isinstance(where, (tuple, list)): + return (where,) + sites = tuple(where) + if len(sites) not in {1, 2}: + raise ValueError("PEPS exact energy terms must act on one or two sites.") + return sites + + @classmethod + def _symmray_exact_local_expectation( + cls, + state, + terms, + *, + optimize, + normalized, + contract_opts, + ): + """Contract native Symmray local terms without forming a dense RDM. + + Quimb's generic ``compute_local_expectation_exact`` forms a reduced + density matrix and fuses its physical legs. Individual Symmray blocks + do not carry the full physical rank, so that dense-only fusion fails. + Directly contracting each operator-inserted ket with the bra preserves + the native block structure and fermionic metadata. + """ + bra = state.H + total = 0 + for where, term in terms.items(): + sites = cls._term_sites(state, where) + inds = [state.site_ind(site) for site in sites] + gated = qtn.tensor_network_gate_inds( + state, + term, + inds, + contract="split", + tags=[], + info=None, + inplace=False, + ) + total = total + (bra | gated).contract( + all, + optimize=optimize, + **contract_opts, + ) + + if normalized: + norm = (state.H & state).contract( + all, + optimize=optimize, + **contract_opts, + ) + total = total / norm + return total + @classmethod def _stabilize_state(cls, state): if hasattr(state, "balance_bonds_") and not cls._state_uses_symmray(state): @@ -468,22 +537,31 @@ def _loss_state( kwargs = dict(compute_kwargs or {}) mode = cls._boundary_mode(boundary_mode) if mode == "exact": - exact_expectation = getattr( - state, - "compute_local_expectation_exact", - None, - ) - if not callable(exact_expectation): - raise TypeError( - "boundary_mode='exact' requires a PEPS-like state with " - "compute_local_expectation_exact()." + if cls._state_uses_symmray(state) and cls._terms_use_symmray(terms): + value = cls._symmray_exact_local_expectation( + state, + terms, + optimize=contraction_opt, + normalized=bool(normalized), + contract_opts=kwargs, + ) + else: + exact_expectation = getattr( + state, + "compute_local_expectation_exact", + None, + ) + if not callable(exact_expectation): + raise TypeError( + "boundary_mode='exact' requires a PEPS-like state with " + "compute_local_expectation_exact()." + ) + value = exact_expectation( + terms, + optimize=contraction_opt, + normalized=bool(normalized), + **kwargs, ) - value = exact_expectation( - terms, - optimize=contraction_opt, - normalized=bool(normalized), - **kwargs, - ) else: value = state.compute_local_expectation( terms, @@ -706,7 +784,10 @@ class MpsEnergyOptimizer(PepsEnergyOptimizer): ``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 - Hamiltonian, or a plain local-term mapping. + Hamiltonian, or a plain local-term mapping. Use the explicit ``terms=`` + constructor keyword when passing a local-term mapping or a Pepsy + ``SymHamiltonian``; ``hamiltonian=`` remains supported as a compatibility + alias. """ _LOSS_KEYS = frozenset({ @@ -722,8 +803,9 @@ class MpsEnergyOptimizer(PepsEnergyOptimizer): def __init__( self, state, - hamiltonian, + hamiltonian=None, *, + terms=None, normalized: bool = True, energy_per_site: bool = True, real: bool = True, @@ -733,6 +815,9 @@ def __init__( loss_kwargs: Mapping[str, Any] | None = None, allow_encoding_conversion: bool = False, ): + if hamiltonian is not None and terms is not None: + raise TypeError("pass either hamiltonian or terms, not both") + hamiltonian = terms if terms is not None else hamiltonian self.state = self._as_mps_state(state) self.hamiltonian = hamiltonian self.terms = self._terms_from_hamiltonian(hamiltonian) @@ -1207,9 +1292,9 @@ def _can_use_native_local_terms(cls, hamiltonian, state): if not cls._terms_use_fermionic_symmray(hamiltonian.terms): return False return all( - len(edge) == 2 - and all(isinstance(site, Integral) for site in edge) - for edge in hamiltonian.edges + len(where) in {1, 2} + and all(isinstance(site, Integral) for site in where) + for where in hamiltonian.edges ) @classmethod @@ -1388,7 +1473,10 @@ def make_tn_optimizer( incoming_constants = dict(loss_constants or {}) terms = incoming_constants.pop("terms", self.terms) if self._is_fermionic_sym_hamiltonian(terms): - terms = self._fermionic_hamiltonian_mpo_for_state(terms, self.state) + if self._can_use_native_local_terms(terms, self.state): + terms = terms.terms + else: + terms = self._fermionic_hamiltonian_mpo_for_state(terms, self.state) if self._is_mpo_hamiltonian(terms): constants = {"terms": terms} constants.update(incoming_constants) diff --git a/src/pepsy/optimizers/mps/layout.py b/src/pepsy/optimizers/mps/layout.py index 21ddc83..4a85b0e 100644 --- a/src/pepsy/optimizers/mps/layout.py +++ b/src/pepsy/optimizers/mps/layout.py @@ -337,7 +337,11 @@ def _gate_stream_event_weights( if weight_mode == "auto": weight = _angle_weight(payload) if weight is None and event_type == "gate": - key = (id(payload), tuple(support), int(schmidt_max_dim)) + # This proxy only handles two-site operators, and its + # singular values are independent of the global site + # labels. Reuse one small SVD when the same gate object is + # replayed across many pairs. + key = (id(payload), len(support), int(schmidt_max_dim)) if key not in schmidt_cache: schmidt_cache[key] = _operator_schmidt_weight( payload, diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index 2fd60d3..b53c6e2 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -42,14 +42,13 @@ center tensor, and accumulates the removed scale into ``p.exponent``. Quimb includes that exponent in ``p.norm()``, so ``p.norm()`` still reports the represented state norm; inspect a copy with ``exponent=0`` to see the -rescaled data norm. Compression infidelity is -measured at every compressed two-site update. The local fidelity is the -retained canonical norm ratio, and cumulative fidelity is the product of -those local fidelities. This is the estimator used in the MPS -circuit-simulation literature and does not require a copied perfect-state -target for unitary streams. Infidelity is always recorded after compressed -two-site updates and is available through ``get_infidelities()`` and -``get_infidelity_samples()``. +rescaled data norm. Compression infidelity is measured at every compressed +two-site update by default. The local fidelity is the retained canonical norm +ratio, and cumulative fidelity is the product of those local fidelities. This +is the estimator used in the MPS circuit-simulation literature and does not +require a copied perfect-state target for unitary streams. Set +``track_infidelity=False`` on the optimizer, or pass it to ``run``, to skip +these diagnostic norm calculations and samples. """ from __future__ import annotations @@ -382,7 +381,11 @@ def _parse_control_mapping(name, entry, default_axis=None): absorb = _normalize_absorb(entry.get("absorb", "left")) return ( "cap", - {"vec": np.asarray(vec, dtype=complex).ravel(), "absorb": absorb}, + { + "vec": np.asarray(vec, dtype=complex).ravel(), + "absorb": absorb, + "compact_labels": bool(entry.get("compact_labels", True)), + }, _normalize_control_where(where, single=True), ) if name == "reset": @@ -564,6 +567,9 @@ class MpsOptimizer: # pylint: disable=too-many-instance-attributes is mutated in place and is exposed as :attr:`gauges`. If omitted, the optimizer initializes it with ``p.gauge_all_simple_(...)`` before the first simple-update gate. + track_infidelity : bool, default=True + Whether to compute and store compression-infidelity diagnostics. Set + this to ``False`` to skip diagnostic norm targets and samples. Attributes ---------- @@ -581,11 +587,12 @@ class MpsOptimizer: # pylint: disable=too-many-instance-attributes applies ``p.exponent``. infidelities : list[float] Cumulative canonical norm-ratio infidelity trace, starting at ``0.0``. - A value is appended after every compressed two-site update. + A value is appended after every compressed two-site update when + :attr:`track_infidelity` is enabled. infidelity_samples : list[dict] Per-update canonical norm-ratio records. Each record contains the target and retained canonical norms, local fidelity, and cumulative - infidelity. + infidelity. This remains empty when tracking is disabled. gauges : dict Simple-update bond gauges. In ``mode="su"``, ``p`` is the gauged core and the physical state is recovered with ``p.gauge_simple_insert(gauges)``. @@ -886,6 +893,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument ind_id="k{}", inplace=False, gauges=None, + track_infidelity=True, ): if chi is None: if isinstance(gates, Integral): @@ -907,6 +915,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument if gauges is not None and not isinstance(gauges, dict): raise TypeError("gauges must be a mutable dictionary or None.") self.gauges = {} if gauges is None else gauges + self.track_infidelity = bool(track_infidelity) self.p_ungauged = None self._su_gauges_supplied = gauges is not None self._su_gauges_ready = False @@ -1046,16 +1055,10 @@ def _is_nearest_neighbor_1d(where): return abs(int(site0) - int(site1)) == 1 def _validate_symmray_mode_support(self): - """Fail early for Symmray/MPS mode combinations with known bad paths.""" + """Fail early for Symmray/MPS combinations with known bad paths.""" if not self._has_symmray_data(self.p): return - if self.mode == "su": - raise NotImplementedError( - "mode='su' currently requires dense MPS tensor data; " - "quimb's simple-update split is not Symmray-safe yet." - ) - if self.mode == "dmrg" and self.p.max_bond() < self.chi: raise ValueError( "Symmray MPS data cannot be automatically expanded for " @@ -1221,6 +1224,7 @@ def copy(self) -> "MpsOptimizer": ind_id=self.ind_id, inplace=True, gauges=deepcopy(self.gauges), + track_infidelity=self.track_infidelity, ) # Restore the source's tracked centre afterwards. It describes the # same represented state and must remain optimizer-local. @@ -1927,6 +1931,7 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments layout_report=True, measure_renormalize=True, seed=None, + track_infidelity=None, ): """Run the currently queued gates. @@ -2001,6 +2006,9 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments seed : int | None, default=None If given, reseed the internal RNG used to sample ``measure``/ ``reset`` outcomes before running, for reproducible collapses. + track_infidelity : bool | None, default=None + Override :attr:`track_infidelity` for this run. ``None`` keeps the + constructor setting. Returns ------- @@ -2013,6 +2021,9 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments if seed is not None: self._rng = np.random.default_rng(seed) + if track_infidelity is not None: + self.track_infidelity = bool(track_infidelity) + self.last_layout_plan = self._persistent_layout_plan G_seq = list(self.G) where_seq = list(self.where) @@ -2107,7 +2118,9 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments non_unitary = bool(non_unitary) self._unitary_initial_norm = None self._unitary_previous_norm = None - self._unitary_global_norm_tracking = not non_unitary and not has_control + self._unitary_global_norm_tracking = ( + self.track_infidelity and not non_unitary and not has_control + ) if not non_unitary: if normalize_every is not None and normalize_every is not False: raise ValueError("normalize_every requires non_unitary=True.") @@ -2587,7 +2600,7 @@ def _state_expectation(self, pauli, where): compatibility with older Quimb versions without that method. """ p = self.p - op = self._pauli_operator(pauli, where) + op = self._to_state_backend(self._pauli_operator(pauli, where)) local_expectation = getattr(p, "local_expectation_canonical", None) if callable(local_expectation): return self._real_float( @@ -3641,9 +3654,10 @@ def _run_mix( # pylint: disable=too-many-arguments,too-many-positional-argument "fallback": fallback_steps, "bond": f"{entry['end_bond']}/{self.chi}", } - postfix["infidelity"] = self._format_progress_infidelity( - self.infidelities[-1] - ) + if self.track_infidelity: + postfix["infidelity"] = self._format_progress_infidelity( + self.infidelities[-1] + ) pbar.set_postfix(postfix) pbar.update(1) finally: @@ -3681,7 +3695,7 @@ def _run_dmrg( last_where = self._current_orthog(p) last_normalized_step = None cumulative_infidelity = None - track_unitary_norm = not non_unitary + track_unitary_norm = self.track_infidelity and not non_unitary if track_unitary_norm and ( not self._unitary_global_norm_tracking or self._unitary_previous_norm is None @@ -3739,7 +3753,7 @@ def _run_dmrg( cutoff, cutoff_mode, ) - if not track_unitary_norm: + if self.track_infidelity and not track_unitary_norm: target_norm = self._raw_state_norm(p_g) fit = FIT( p_g, @@ -3752,24 +3766,25 @@ def _run_dmrg( ) fit.run_gate(n_iter=n_iter, verbose=False) - approx_norm = fit.local_norm_trace[-1] p = self._install_represented_norm(fit.p) self.p = p self._record_orthog_span(p, (xmin, xmax)) - if track_unitary_norm: - sample = self._append_unitary_compression_infidelity_sample( - approx_norm, - step=idx + 1, - where=(xmin, xmax), - ) - else: - sample = self._append_compression_infidelity_sample( - approx_norm, - target_norm, - step=idx + 1, - where=(xmin, xmax), - ) - cumulative_infidelity = sample["infidelity"] + if self.track_infidelity: + approx_norm = fit.local_norm_trace[-1] + if track_unitary_norm: + sample = self._append_unitary_compression_infidelity_sample( + approx_norm, + step=idx + 1, + where=(xmin, xmax), + ) + else: + sample = self._append_compression_infidelity_sample( + approx_norm, + target_norm, + step=idx + 1, + where=(xmin, xmax), + ) + cumulative_infidelity = sample["infidelity"] idx += 1 advanced = 1 last_where = (xmin, xmax) @@ -3787,7 +3802,7 @@ def _run_dmrg( self.canonize_mps(p, (xmin, xmax)) target_norm = None p_g = self._build_dmrg_batch_target(p, batch_G, batch_where, cutoff, cutoff_mode) - if not track_unitary_norm: + if self.track_infidelity and not track_unitary_norm: target_norm = self._raw_state_norm(p_g) fit = FIT( p_g, @@ -3799,24 +3814,25 @@ def _run_dmrg( ) fit.run_gate(n_iter=n_iter, verbose=False) - approx_norm = fit.local_norm_trace[-1] p = self._install_represented_norm(fit.p) self.p = p self._record_orthog_span(p, (xmin, xmax)) - if track_unitary_norm: - sample = self._append_unitary_compression_infidelity_sample( - approx_norm, - step=next_idx, - where=(xmin, xmax), - ) - else: - sample = self._append_compression_infidelity_sample( - approx_norm, - target_norm, - step=next_idx, - where=(xmin, xmax), - ) - cumulative_infidelity = sample["infidelity"] + if self.track_infidelity: + approx_norm = fit.local_norm_trace[-1] + if track_unitary_norm: + sample = self._append_unitary_compression_infidelity_sample( + approx_norm, + step=next_idx, + where=(xmin, xmax), + ) + else: + sample = self._append_compression_infidelity_sample( + approx_norm, + target_norm, + step=next_idx, + where=(xmin, xmax), + ) + cumulative_infidelity = sample["infidelity"] advanced = next_idx - idx idx = next_idx last_where = (xmin, xmax) @@ -3837,11 +3853,12 @@ def _run_dmrg( "2q": two_qubit_count, "bnd": p.max_bond(), } - postfix["infidelity"] = self._format_progress_infidelity( - self.infidelities[-1] - if cumulative_infidelity is None - else cumulative_infidelity - ) + if self.track_infidelity: + postfix["infidelity"] = self._format_progress_infidelity( + self.infidelities[-1] + if cumulative_infidelity is None + else cumulative_infidelity + ) pbar.set_postfix(postfix) pbar.update(advanced) @@ -3956,7 +3973,7 @@ def _run_mpo( # pylint: disable=too-many-locals Uses :meth:`qtn.MatrixProductState.gate_nonlocal_` for two-qubit gates. """ p = self.p - if not non_unitary: + if self.track_infidelity and not non_unitary: self._start_unitary_norm_tracking(p) two_qubit_count = 0 submpo_count = 0 @@ -3989,7 +4006,7 @@ def _run_mpo( # pylint: disable=too-many-locals method = self._normalize_submpo_method(submpo_method) submpo_compress_opts = self._submpo_compress_opts(method) self.canonize_mps(p, (xmin, xmax)) - if non_unitary: + if self.track_infidelity and non_unitary: p_target = p.copy() p_target.gate_with_submpo_( gate, @@ -4022,21 +4039,22 @@ def _run_mpo( # pylint: disable=too-many-locals advanced = 1 last_where = (xmin, xmax) compressed = True - approx_norm = self._canonical_span_norm(p, (xmin, xmax)) - if non_unitary: - sample = self._append_compression_infidelity_sample( - approx_norm, - target_norm, - step=idx, - where=(xmin, xmax), - ) - else: - sample = self._append_unitary_compression_infidelity_sample( - approx_norm, - step=idx, - where=(xmin, xmax), - ) - norm_cumulative_infidelity = sample["infidelity"] + if self.track_infidelity: + approx_norm = self._canonical_span_norm(p, (xmin, xmax)) + if non_unitary: + sample = self._append_compression_infidelity_sample( + approx_norm, + target_norm, + step=idx, + where=(xmin, xmax), + ) + else: + sample = self._append_unitary_compression_infidelity_sample( + approx_norm, + step=idx, + where=(xmin, xmax), + ) + norm_cumulative_infidelity = sample["infidelity"] elif len(where) == 1: self._apply_gate( p, @@ -4059,11 +4077,10 @@ def _run_mpo( # pylint: disable=too-many-locals xmin, xmax = sorted(where) use_symmray_auto_swap = ( self._has_symmray_data(p) - and not self._is_nearest_neighbor_1d(where) ) self.canonize_mps(p, (xmin, xmax)) p_target = None - if non_unitary: + if self.track_infidelity and non_unitary: if use_symmray_auto_swap: p_target = self._build_symmray_auto_swap_target( p, gate, where, cutoff, cutoff_mode @@ -4104,21 +4121,22 @@ def _run_mpo( # pylint: disable=too-many-locals advanced = 1 last_where = (xmin, xmax) compressed = True - approx_norm = self._canonical_span_norm(p, (xmin, xmax)) - if non_unitary: - sample = self._append_compression_infidelity_sample( - approx_norm, - target_norm, - step=idx, - where=(xmin, xmax), - ) - else: - sample = self._append_unitary_compression_infidelity_sample( - approx_norm, - step=idx, - where=(xmin, xmax), - ) - norm_cumulative_infidelity = sample["infidelity"] + if self.track_infidelity: + approx_norm = self._canonical_span_norm(p, (xmin, xmax)) + if non_unitary: + sample = self._append_compression_infidelity_sample( + approx_norm, + target_norm, + step=idx, + where=(xmin, xmax), + ) + else: + sample = self._append_unitary_compression_infidelity_sample( + approx_norm, + step=idx, + where=(xmin, xmax), + ) + norm_cumulative_infidelity = sample["infidelity"] event = self._maybe_normalize_after_step( p, @@ -4137,9 +4155,10 @@ def _run_mpo( # pylint: disable=too-many-locals } if submpo_count: postfix["mpo"] = submpo_count - postfix["infidelity"] = self._format_progress_infidelity( - norm_cumulative_infidelity - ) + if self.track_infidelity: + postfix["infidelity"] = self._format_progress_infidelity( + norm_cumulative_infidelity + ) pbar.set_postfix(postfix) pbar.update(advanced) @@ -4192,7 +4211,7 @@ def _run_swap_network( # pylint: disable=too-many-locals,too-many-arguments and the right endpoint remains at the left endpoint's neighbour. """ p = self.p - if not non_unitary: + if self.track_infidelity and not non_unitary: self._start_unitary_norm_tracking(p) two_qubit_count = 0 cumulative_infidelity = self.infidelities[-1] @@ -4244,7 +4263,7 @@ def _run_swap_network( # pylint: disable=too-many-locals,too-many-arguments xmin, xmax = sorted(where) self.canonize_mps(p, (xmin, xmax)) p_target = None - if non_unitary: + if self.track_infidelity and non_unitary: target_norm = self._gate_target_norm_from_expectation( p, gate, where ) @@ -4272,21 +4291,22 @@ def _run_swap_network( # pylint: disable=too-many-locals,too-many-arguments advanced = 1 last_where = (xmin, xmax) compressed = True - approx_norm = self._canonical_span_norm(p, (xmin, xmax)) - if non_unitary: - sample = self._append_compression_infidelity_sample( - approx_norm, - target_norm, - step=idx, - where=(xmin, xmax), - ) - else: - sample = self._append_unitary_compression_infidelity_sample( - approx_norm, - step=idx, - where=(xmin, xmax), - ) - cumulative_infidelity = sample["infidelity"] + if self.track_infidelity: + approx_norm = self._canonical_span_norm(p, (xmin, xmax)) + if non_unitary: + sample = self._append_compression_infidelity_sample( + approx_norm, + target_norm, + step=idx, + where=(xmin, xmax), + ) + else: + sample = self._append_unitary_compression_infidelity_sample( + approx_norm, + step=idx, + where=(xmin, xmax), + ) + cumulative_infidelity = sample["infidelity"] event = self._maybe_normalize_after_step( p, @@ -4303,9 +4323,10 @@ def _run_swap_network( # pylint: disable=too-many-locals,too-many-arguments "2q": two_qubit_count, "bnd": p.max_bond(), } - postfix["infidelity"] = self._format_progress_infidelity( - cumulative_infidelity - ) + if self.track_infidelity: + postfix["infidelity"] = self._format_progress_infidelity( + cumulative_infidelity + ) pbar.set_postfix(postfix) pbar.update(advanced) @@ -4343,11 +4364,11 @@ def _run_svd( # pylint: disable=too-many-locals Two-site gates are applied with ``contract="reduce-split"`` then compressed on the local span to ``max_bond=self.chi``. Symmray-backed - MPS data use quimb's block-aware auto-swap split path instead, since - ``reduce-split`` loses Symmray fusion metadata in current quimb/Symmray. + MPS data use quimb's block-aware auto-swap split path by default as a + conservative choice for block-sparse edge cases. """ p = self.p - if not non_unitary: + if self.track_infidelity and not non_unitary: self._start_unitary_norm_tracking(p) two_qubit_count = 0 cumulative_infidelity = self.infidelities[-1] @@ -4397,7 +4418,7 @@ def _run_svd( # pylint: disable=too-many-locals use_symmray_auto_swap = self._has_symmray_data(p) self.canonize_mps(p, (xmin, xmax)) if use_symmray_auto_swap: - if not non_unitary: + if not self.track_infidelity or not non_unitary: target_norm = None p_target = None else: @@ -4418,7 +4439,7 @@ def _run_svd( # pylint: disable=too-many-locals max_bond=self.chi, ) else: - if not non_unitary: + if not self.track_infidelity or not non_unitary: target_norm = None else: target_norm = None @@ -4431,7 +4452,7 @@ def _run_svd( # pylint: disable=too-many-locals cutoff_mode=cutoff_mode, inplace=True, ) - if non_unitary and target_norm is None: + if self.track_infidelity and non_unitary and target_norm is None: target_norm = self._canonical_span_norm(p, (xmin, xmax)) self.canonize_mps(p, (xmin, xmax)) @@ -4448,21 +4469,22 @@ def _run_svd( # pylint: disable=too-many-locals advanced = 1 last_where = (xmin, xmax) compressed = True - approx_norm = self._canonical_span_norm(p, (xmin, xmax)) - if non_unitary: - sample = self._append_compression_infidelity_sample( - approx_norm, - target_norm, - step=idx, - where=(xmin, xmax), - ) - else: - sample = self._append_unitary_compression_infidelity_sample( - approx_norm, - step=idx, - where=(xmin, xmax), - ) - cumulative_infidelity = sample["infidelity"] + if self.track_infidelity: + approx_norm = self._canonical_span_norm(p, (xmin, xmax)) + if non_unitary: + sample = self._append_compression_infidelity_sample( + approx_norm, + target_norm, + step=idx, + where=(xmin, xmax), + ) + else: + sample = self._append_unitary_compression_infidelity_sample( + approx_norm, + step=idx, + where=(xmin, xmax), + ) + cumulative_infidelity = sample["infidelity"] event = self._maybe_normalize_after_step( p, @@ -4479,9 +4501,10 @@ def _run_svd( # pylint: disable=too-many-locals "2q": two_qubit_count, "bnd": p.max_bond(), } - postfix["infidelity"] = self._format_progress_infidelity( - cumulative_infidelity - ) + if self.track_infidelity: + postfix["infidelity"] = self._format_progress_infidelity( + cumulative_infidelity + ) pbar.set_postfix(postfix) pbar.update(advanced) diff --git a/src/pepsy/optimizers/noise.py b/src/pepsy/optimizers/noise.py index e33dab6..f6797b0 100644 --- a/src/pepsy/optimizers/noise.py +++ b/src/pepsy/optimizers/noise.py @@ -1320,7 +1320,7 @@ def _is_stabilizer_trajectory_optimizer(optimizer) -> bool: def _trajectory_norm_squared(optimizer) -> float: - """Read the represented state norm through either public MPS optimizer API.""" + """Read the represented state norm through the optimizer's public API.""" norm = getattr(optimizer, "norm", None) if callable(norm): value = _trajectory_real_scalar(norm(), label="trajectory state norm") @@ -1361,6 +1361,30 @@ def _mps_outcome_norm_squared(optimizer, matrix, where) -> float: return value * value +def _tree_outcome_norm_squared(optimizer, matrix, where) -> float: + """Evaluate one Kraus branch on a copied ordinary TTN without mutation.""" + copy = getattr(optimizer, "copy", None) + if not callable(copy): + raise TypeError( + "State-dependent trajectory channels require an optimizer with copy()." + ) + candidate = copy() + apply_gate = getattr(candidate, "apply_gate", None) + if not callable(apply_gate): + raise TypeError( + "State-dependent trajectory channels require TreeOptimizer " + "apply_gate(...)." + ) + matrix = _to_trajectory_backend(matrix, candidate) + support = _trajectory_where(where) + target = support[0] if len(support) == 1 else support + apply_gate(matrix, target) + value = _trajectory_real_scalar(candidate.norm(), label="Kraus branch norm") + if not np.isfinite(value) or value < 0.0: + raise ValueError("Kraus branch produced an invalid TTN norm.") + return value * value + + def _stn_outcome_norm_squared(optimizer, matrix, where) -> float: """Evaluate one physical Kraus branch in an independent STN frame copy.""" candidate = optimizer.copy() @@ -1372,9 +1396,14 @@ def _stn_outcome_norm_squared(optimizer, matrix, where) -> float: def _to_trajectory_backend(matrix, optimizer): - """Convert generated NumPy matrices to the ordinary MPS state backend.""" + """Convert generated NumPy matrices to the live MPS or TTN backend.""" converter = getattr(optimizer, "_to_state_backend", None) - return converter(matrix) if callable(converter) else matrix + if callable(converter): + return converter(matrix) + converter = getattr(optimizer, "_as_state_backend", None) + if callable(converter): + return converter(matrix, warn=False) + return matrix def _kraus_probabilities(optimizer, channel: TrajectoryChannel, where) -> np.ndarray: @@ -1388,6 +1417,14 @@ def _kraus_probabilities(optimizer, channel: TrajectoryChannel, where) -> np.nda ], dtype=float, ) + elif callable(getattr(optimizer, "apply_gate", None)): + branch_norm_squared = np.asarray( + [ + _tree_outcome_norm_squared(optimizer, outcome.gate, where) + for outcome in channel.outcomes + ], + dtype=float, + ) else: branch_norm_squared = np.asarray( [ @@ -1782,7 +1819,7 @@ def _check_coalesced_optimizer(optimizer): if not callable(getattr(optimizer, "copy", None)): raise TypeError( "coalesced trajectory replay requires an optimizer with copy(); " - "use MpsOptimizer or MpsStabOptimizer." + "use MpsOptimizer, TreeOptimizer, or MpsStabOptimizer." ) @@ -1939,12 +1976,17 @@ def _coalesced_measurement_probability(optimizer, pauli, where) -> float: arg = where[0] if len(where) == 1 else where value = expectation(pauli, arg) else: + expectation_pauli = getattr(optimizer, "expectation_pauli", None) + if callable(expectation_pauli): + arg = where[0] if len(where) == 1 else where + value = expectation_pauli(pauli, arg) + return min(max(0.5 * (1.0 + float(value)), 0.0), 1.0) mapped = getattr(optimizer, "_logical_to_physical_where", None) state_expectation = getattr(optimizer, "_state_expectation", None) if not callable(mapped) or not callable(state_expectation): raise TypeError( - "coalesced measurement branching requires MpsOptimizer or " - "MpsStabOptimizer expectation support." + "coalesced measurement branching requires MpsOptimizer, " + "MpsStabOptimizer, or TreeOptimizer expectation support." ) value = state_expectation(pauli, mapped(where)) return min(max(0.5 * (1.0 + float(value)), 0.0), 1.0) @@ -2199,19 +2241,20 @@ def run_trajectory_shots( strategy: str = "independent", max_branches: int | None = _AUTO_MAX_BRANCHES, ) -> TrajectoryShotResult | CoalescedTrajectoryResult: - """Replay user-defined noisy gate-stream trajectories on either MPS optimizer. + """Replay user-defined noisy gate-stream trajectories on MPS or tree optimizers. Insert :class:`TrajectoryEvent` objects or Pepsy stochastic entries directly into an ordinary gate stream. A ``mixture`` selects a known unitary branch by its explicit probability. A ``kraus`` channel evaluates all local branch - norms on the current MPS, samples the conditional probability, applies the + norms on the current MPS or TTN, samples the conditional probability, applies the chosen branch, and normalizes before evolution continues. This includes non-Pauli channels such as ``("amplitude_damping", gamma, q)`` without forming a density matrix. - ``optimizer_factory`` must create a fresh :class:`MpsOptimizer` or - :class:`MpsStabOptimizer` per shot. Gate segments between channel events - are batched, so a trajectory does not rebuild an optimizer for every gate. + ``optimizer_factory`` must create a fresh :class:`MpsOptimizer`, + :class:`TreeOptimizer`, or :class:`MpsStabOptimizer` per shot. Gate segments + between channel events are batched, so a trajectory does not rebuild an + optimizer for every gate. Set ``strategy="coalesced"`` to share deterministic prefixes and retain one optimizer per distinct sampled branch. ``strategy="auto"`` tries coalescing and restarts independently if ``max_branches`` would be exceeded. @@ -2921,10 +2964,7 @@ def sample_stim_circuits(circuit, shots: int, *, seed=None) -> list[StimNoiseSam def _stream_on_optimizer_backend(stream, optimizer): - """Convert library-generated dense gates to an ordinary MPS backend.""" - converter = getattr(optimizer, "_to_state_backend", None) - if not callable(converter): - return stream + """Convert library-generated dense gates to the live MPS/TTN backend.""" converted = [] for entry in stream: if ( @@ -2933,7 +2973,7 @@ def _stream_on_optimizer_backend(stream, optimizer): and not isinstance(entry[0], str) and hasattr(entry[0], "shape") ): - converted.append((converter(entry[0]), entry[1])) + converted.append((_to_trajectory_backend(entry[0], optimizer), entry[1])) else: converted.append(entry) return tuple(converted) diff --git a/src/pepsy/optimizers/peps/optimizer.py b/src/pepsy/optimizers/peps/optimizer.py index 15390ce..eef7c3e 100644 --- a/src/pepsy/optimizers/peps/optimizer.py +++ b/src/pepsy/optimizers/peps/optimizer.py @@ -17,6 +17,8 @@ from ..sweep.environments import ( canonical_boundary_engine_selector, normalize_boundary_engine, + symmray_array_backends, + uses_symmray_arrays, ) __all__ = ["PepsOptimizer"] @@ -146,9 +148,9 @@ class PepsOptimizer: # pylint: disable=too-many-instance-attributes ``measure_final_infidelity=True``. If final measurement is disabled, the optimizer loss used as a fallback can come from the coarser ``boundary_chi`` environment while the pre-check used ``evaluation_chi``. - - Sweep mode defaults to the optional NLopt ``LD_LBFGS`` solver. Install - NLopt or pass an explicit ``optimizer`` / ``sweep_optimize_kwargs`` value - for environments where NLopt is unavailable. + - Symmray PEPS inputs must be Torch-backed. Their sweep cleanup defaults + to NLopt ``LD_LBFGS`` while obtaining gradients from Torch autograd; + provide both the PEPS and gates with the same Torch Symmray backend. - ``non_unitary=True`` normalizes generated targets and candidates; it does not implement the interval scheduling or norm-proxy machinery available in :class:`pepsy.optimizers.mps.MpsOptimizer`. @@ -223,7 +225,8 @@ class PepsOptimizer: # pylint: disable=too-many-instance-attributes Compact optimizer selection for the chosen variational backend. In sweep mode this maps to :class:`SweepOptimizer`'s local solver name; when left as ``None`` the sweep local solver defaults to NLopt - ``LD_LBFGS``. In global mode, ``"nlopt"`` together with + ``LD_LBFGS``. Torch-backed Symmray states still supply gradients via + Torch autograd. In global mode, ``"nlopt"`` together with ``optimizer_options={"algorithm": "LD_VAR2"}`` routes to :meth:`GlobalOptimizer.optimize_nlopt`. optimizer_options : mapping, optional @@ -319,6 +322,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self.which = which self.inplace = bool(inplace) self.state = state if self.inplace else (state.copy() if hasattr(state, "copy") else state) + self._require_torch_symmray_backend(self.state, role="state") self.gates = _normalize_gate_queue(gates) self.normalize_initial = bool(normalize_initial) @@ -377,6 +381,28 @@ def _validate_boundary_chi(cls, chi): ) return cls._validate_scalar_chi(chi, name="boundary_chi") + @staticmethod + def _require_torch_symmray_backend(*states, role="inputs"): + """Require a single Torch backend whenever Symmray data is present.""" + states = tuple(state for state in states if state is not None) + symmray_flags = tuple(uses_symmray_arrays(state) for state in states) + if not any(symmray_flags): + return + if not all(symmray_flags): + raise TypeError( + "PepsOptimizer requires matching Torch-backed Symmray arrays " + f"for {role}; do not mix Symmray and dense tensor networks." + ) + backends = symmray_array_backends(*states) + if backends == {"torch"}: + return + found = ", ".join(sorted(backends)) or "unknown" + raise TypeError( + "PepsOptimizer requires Torch-backed Symmray arrays for autograd; " + f"{role} use backend(s): {found}. Convert the PEPS and every gate " + "with pepsy.backend_torch(...) before optimization." + ) + def _boundary_chi_max(self): if isinstance(self.boundary_chi, tuple): return max(int(self.boundary_chi[0]), int(self.boundary_chi[1])) @@ -1257,6 +1283,11 @@ def _optimize_state( global_kwargs, global_optimize_kwargs, ): + self._require_torch_symmray_backend( + state, + target, + role="state and target", + ) if mode == "sweep": return self._optimize_with_sweep( state, @@ -1431,6 +1462,7 @@ def run( # pylint: disable=too-many-arguments,too-many-locals,too-many-branches normalize_kwargs=normalize_kwargs, normalize_chi=normalize_chi, ) + self._require_torch_symmray_backend(self.state, role="state") pbar = None if show_progress: @@ -1479,6 +1511,10 @@ def run( # pylint: disable=too-many-arguments,too-many-locals,too-many-branches opts=opts, inplace=True, ) + self._require_torch_symmray_backend( + self.state, + role="state after gate application", + ) if normalize_target: self._normalize_state( self.state, @@ -1514,6 +1550,11 @@ def run( # pylint: disable=too-many-arguments,too-many-locals,too-many-branches cutoff_mode=cutoff_mode, gate_kwargs=gate_kwargs, ) + self._require_torch_symmray_backend( + state_before, + target, + role="state and gate-generated target", + ) if normalize_target: self._normalize_state( target, diff --git a/src/pepsy/optimizers/stabilizer_tn/PLAN.md b/src/pepsy/optimizers/stabilizer_tn/PLAN.md index 26043a8..9fec3ee 100644 --- a/src/pepsy/optimizers/stabilizer_tn/PLAN.md +++ b/src/pepsy/optimizers/stabilizer_tn/PLAN.md @@ -114,9 +114,11 @@ new simulator features. sequential left fold. - Unitary truncation diagnostics use the cumulative coefficient-norm loss `1 - ||nu||^2`, evaluated from the tracked one-site canonical centre. This - avoids uncapped target copies and overlap contractions. Non-unitary and - projective updates emit no `.infidelities` sample; projective normalization - starts a fresh unitary segment. Measurement/reset boundaries are preserved in + avoids uncapped target copies and overlap contractions. Dense multi-qubit + non-unitary matrices use a local physical `G†G` target norm to report retained + compression `infidelity`; coefficient-frame sub-MPOs remain uncertified. + Projective normalization starts a fresh unitary segment. Measurement/reset + boundaries are preserved in `.norm_events`: pre-collapse norm, Born branch probability, actual projected norm before normalization, and `projector_infidelity` for MPS compression incurred by the projector itself. `norm_diagnostics()` combines unitary and @@ -195,7 +197,9 @@ new simulator features. - **Initial states** — `STNState.zero/from_bits/ghz/from_tableau_and_state` and the matching `MpsStabOptimizer.from_bits/ghz/from_tableau_and_state` classmethods. - **Progress bar + diagnostics** — `run(progbar=True)` (tqdm, reports the current - stream part and `norm_infidelity`); `norm()` returns the `|nu>` norm. + stream part and MPS-compatible `infidelity`, with `norm_infidelity` retained as + an alias); `norm_diagnostics()` reports the same multiplicative `infidelity` + and `fidelity` names, and `norm()` returns the `|nu>` norm. - `StabilizerMps` is kept as a backward-compatible alias for `MpsStabOptimizer`. - **Amplitude / observable API** — `amplitude(bits)`/`probability(bits)`; `expectation(pauli, where=None)` (also full-register strings like `"ZIZ"`); diff --git a/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py b/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py index b917e9e..db65fec 100644 --- a/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py +++ b/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py @@ -626,10 +626,11 @@ class MpsStabOptimizer: If ``True``, record ``1 - ||nu||**2`` after compressed unitary updates. The normalized initial coefficient state is not renormalized during unitary evolution, so this is a cheap cumulative norm-loss proxy read - from the canonical centre. Non-unitary updates do not produce samples. - Projective measurement/reset boundaries additionally snapshot the - current segment norm in :attr:`norm_events` before normalizing the - selected branch. + from the canonical centre. Compressing a dense multi-qubit non-unitary + matrix also records its retained norm relative to the local + ``G^dagger G`` target norm. Projective measurement/reset boundaries + additionally snapshot the current segment norm in :attr:`norm_events` + before normalizing the selected branch. exact_cooling : bool If ``True`` (default), recognize the constructive exact-cooling case before a multi-site Pauli rotation. A usable separable stabilizer site @@ -666,7 +667,9 @@ class MpsStabOptimizer: state : STNState The evolving stabilizer tensor-network state. infidelities : list[float] - Cumulative norm-loss samples from compressed unitary ``|nu>`` updates. + Cumulative ``infidelity`` samples from compressed updates. Dense + non-unitary entries use the local ``G^dagger G`` target norm; projective + boundaries are additionally represented in :attr:`norm_events`. norm_events : list[NormEventRecord] Segment-boundary snapshots made immediately before projective measurement/reset normalization. These preserve the pre-collapse @@ -797,6 +800,7 @@ def __init__( self._queue: List[object] = [] self.infidelities: List[float] = [] + self._nonunitary_infidelities: List[float] = [] self.norm_events: List[NormEventRecord] = [] self.bond_history: List[int] = [self.state.max_bond()] self.exact_cooling_events: List[dict] = [] @@ -2371,7 +2375,8 @@ def run(self, *, progbar: bool = False) -> "MpsStabOptimizer": ---------- progbar : bool Show a ``tqdm`` progress bar reporting the current stream part and - ``norm_infidelity`` diagnostic. + the MPS-compatible ``infidelity`` diagnostic. The legacy + ``norm_infidelity`` key is emitted alongside it. """ queue = tuple(self._queue) completed = 0 @@ -2388,14 +2393,14 @@ def run(self, *, progbar: bool = False) -> "MpsStabOptimizer": if pbar is not None: pbar.update(1) diagnostics = self.norm_diagnostics() - norm_infidelity = diagnostics["norm_infidelity"] + infidelity = diagnostics["infidelity"] + formatted_infidelity = self._format_progress_infidelity( + infidelity + ) pbar.set_postfix( part=part, - norm_infidelity=( - "n/a" - if norm_infidelity is None - else f"{norm_infidelity:.2e}" - ), + infidelity=formatted_infidelity, + norm_infidelity=formatted_infidelity, ) finally: if pbar is not None: @@ -2408,6 +2413,24 @@ def apply(self, gates, *, progbar: bool = False) -> "MpsStabOptimizer": """Convenience: queue ``gates`` and run immediately.""" return self.set_gates(gates).run(progbar=progbar) + @staticmethod + def _format_progress_infidelity(value) -> str: + """Format the shared MPS/STN progress-bar infidelity field.""" + if value is None: + return "n/a" + text = f"{float(value):#.0e}" + if "e" not in text: + return text + mantissa, exponent = text.split("e", 1) + sign = exponent[0] if exponent[:1] in "+-" else "" + digits = exponent[1:] if sign else exponent + digits = digits.lstrip("0") or "0" + return f"{mantissa}e{sign}{digits}" + + def get_infidelities(self): + """Return the cumulative ``infidelity`` trace like ``MpsOptimizer``.""" + return self.infidelities + def to_statevector(self) -> np.ndarray: """Dense statevector ``|psi> = C|nu>`` (small ``n`` only).""" if self._layout_is_identity(): @@ -2564,7 +2587,8 @@ def norm_diagnostics(self, *, include_current: bool = True) -> dict: norm is also folded into the product/geometric summaries. The returned values are compression/norm-survival proxies only; measurement branch probabilities are kept in the individual events and are not multiplied - into the truncation total. + into the truncation total. Dense non-unitary matrix updates contribute + their ``G^dagger G``-normalized compression loss. """ completed = [ event @@ -2579,6 +2603,7 @@ def norm_diagnostics(self, *, include_current: bool = True) -> dict: for event in completed if event.get("projector_infidelity") is not None ] + completed_nonunitary_losses = list(self._nonunitary_infidelities) completed_survivals = [] for event, unitary_loss in zip(completed, completed_unitary_losses): unitary_survival = min(1.0, max(0.0, 1.0 - unitary_loss)) @@ -2589,6 +2614,10 @@ def norm_diagnostics(self, *, include_current: bool = True) -> dict: unitary_survival * min(1.0, max(0.0, float(projector_survival))) ) survivals = list(completed_survivals) + survivals.extend( + min(1.0, max(0.0, 1.0 - loss)) + for loss in completed_nonunitary_losses + ) current_loss = None if ( include_current @@ -2601,15 +2630,17 @@ def norm_diagnostics(self, *, include_current: bool = True) -> dict: survivals.append(min(1.0, max(0.0, 1.0 - current_loss))) if survivals: - total_survival = float(np.prod(survivals)) if any(survival <= 0.0 for survival in survivals): + total_survival = 0.0 geometric_mean_survival = 0.0 else: + log_survival = sum(math.log(survival) for survival in survivals) + total_survival = float(math.exp(log_survival)) geometric_mean_survival = float( - math.exp(sum(math.log(survival) for survival in survivals) - / len(survivals)) + math.exp(log_survival / len(survivals)) ) event_losses = [1.0 - survival for survival in completed_survivals] + event_losses.extend(completed_nonunitary_losses) if current_loss is not None: event_losses.append(current_loss) mean_segment_infidelity = float(sum(event_losses) / len(event_losses)) @@ -2641,6 +2672,7 @@ def norm_diagnostics(self, *, include_current: bool = True) -> dict: "completed_projector_infidelities": [ event["projector_infidelity"] for event in completed ], + "completed_nonunitary_infidelities": completed_nonunitary_losses, "completed_combined_infidelities": [ float(1.0 - survival) for survival in completed_survivals ], @@ -2648,6 +2680,11 @@ def norm_diagnostics(self, *, include_current: bool = True) -> dict: "current_segment_infidelity": current_loss, "norm_survival": norm_survival, "norm_infidelity": norm_infidelity, + # MpsOptimizer-compatible public names. ``infidelity`` is the + # cumulative multiplicative compression infidelity; it never + # includes stochastic measurement branch probabilities. + "fidelity": norm_survival, + "infidelity": norm_infidelity, "norm": norm, "total_survival_proxy": norm_survival, "total_infidelity_proxy": norm_infidelity, @@ -2778,6 +2815,14 @@ def copy(self) -> "MpsStabOptimizer": copied._norm_infidelity_valid = self._norm_infidelity_valid copied._current_norm_infidelity = self._current_norm_infidelity copied._norm_segment_open = self._norm_segment_open + copied.infidelities = list(self.infidelities) + copied._nonunitary_infidelities = list(self._nonunitary_infidelities) + copied.norm_events = [ + NormEventRecord(**event.as_dict()) + if isinstance(event, NormEventRecord) + else dict(event) + for event in self.norm_events + ] copied.logical_order = list(self.logical_order) copied._refresh_layout_map() copied.layout_plan = None if self.layout_plan is None else dict(self.layout_plan) @@ -4867,6 +4912,65 @@ def with_deferred_injection( # ------------------------------------------------------------------ # # Explicit gate matrices # ------------------------------------------------------------------ # + def _dense_gate_target_norm(self, gate: np.ndarray, where) -> float: + """Evaluate ``||G|psi>||`` from the local physical ``G^dagger G``. + + The Gram operator is Pauli-decomposed and each physical Pauli is + evaluated in the coefficient frame. This is the STN analogue of the + ordinary MPS canonical local-expectation path and avoids copying an + uncompressed target MPS. + """ + where = ( + (int(where),) + if isinstance(where, Integral) + else tuple(int(site) for site in where) + ) + k = len(where) + norm_squared = self._norm_squared() + if not np.isfinite(norm_squared): + raise ValueError( + "Cannot evaluate a non-unitary gate target norm from an invalid " + f"coefficient norm squared {norm_squared!r}." + ) + if norm_squared <= 0.0: + return 0.0 + gram = np.asarray(gate).conj().T @ np.asarray(gate) + expectation = 0.0 + 0.0j + for labels, coefficient in pauli_decomposition( + gram, k, tol=self.operator_tol + ): + physical = pauli_string(labels, where, self.n) + frame_terms, sign = hermitian_pauli_terms( + self.state.frame_pauli(physical) + ) + expectation += complex(coefficient) * self._pauli_expectation( + frame_terms, sign + ) + target_squared = float(np.real(expectation)) * norm_squared + if target_squared < 0.0: + if target_squared > -1.0e-10: + target_squared = 0.0 + else: + raise ValueError( + "G^dagger G produced a negative target norm squared: " + f"{target_squared!r}." + ) + return float(target_squared ** 0.5) + + def _nonunitary_compression_infidelity(self, target_norm: float) -> float: + """Record local compression infidelity relative to ``G^dagger G``.""" + target_norm = float(target_norm) + approx_norm = float(self.norm()) + if target_norm <= 1.0e-15: + fidelity = 1.0 if approx_norm <= 1.0e-15 else 0.0 + else: + fidelity = (approx_norm / target_norm) ** 2 + fidelity = min(1.0, max(0.0, fidelity)) + infidelity = float(1.0 - fidelity) + self._nonunitary_infidelities.append(infidelity) + self._invalidate_norm_infidelity() + return infidelity + def _apply_matrix(self, gate: np.ndarray, where) -> None: where = (int(where),) if isinstance(where, Integral) else tuple(int(w) for w in where) dim = gate.shape[0] @@ -4936,6 +5040,11 @@ def _apply_dense_gate( "already expressed in the coefficient frame; or raise the " "limit explicitly." ) + target_norm = ( + None + if unitary or k < 2 + else self._dense_gate_target_norm(gate, where) + ) decomp = pauli_decomposition(gate, k, tol=self.operator_tol) branches = [] # (weight, {site: axis}) for labels, coeff in decomp: @@ -4948,9 +5057,17 @@ def _apply_dense_gate( 0 < len(branches) <= _MAX_PAULI_SUM_SUBMPO_TERMS and len(support) >= 2 ): - self._record(self._apply_pauli_sum_submpo(branches, unitary=unitary)) + self._record( + self._apply_pauli_sum_submpo( + branches, unitary=unitary, target_norm=target_norm + ) + ) else: - self._record(self._apply_operator_sum(branches, unitary=unitary)) + self._record( + self._apply_operator_sum( + branches, unitary=unitary, target_norm=target_norm + ) + ) def _coalesce_operator_sum(self, branches): """Combine equal Pauli-string branches and prune exact/tolerant zeros.""" @@ -4967,25 +5084,33 @@ def _coalesce_operator_sum(self, branches): if abs(weight) > tol ) - def _apply_pauli_sum_submpo(self, branches, *, unitary: bool) -> Optional[float]: + def _apply_pauli_sum_submpo( + self, branches, *, unitary: bool, target_norm: Optional[float] = None + ) -> Optional[float]: """Apply a sparse Pauli-product sum as one exact coefficient-frame MPO.""" mapped = tuple( (weight, self._mps_terms(sites)) for weight, sites in branches ) mpo, where = pauli_sum_submpo(mapped, self.n, dtype=self.dtype) - infidelity = self._evolve_p(self._bk_mpo(mpo), where, unitary=unitary) - if not unitary: + if unitary: + return self._evolve_p(self._bk_mpo(mpo), where, unitary=True) + self._evolve_p(self._bk_mpo(mpo), where) + if target_norm is None: self._invalidate_norm_infidelity() - return infidelity + return None + return self._nonunitary_compression_infidelity(target_norm) - def _apply_operator_sum(self, branches, *, unitary: bool) -> Optional[float]: + def _apply_operator_sum( + self, branches, *, unitary: bool, target_norm: Optional[float] = None + ) -> Optional[float]: """Apply ``M = sum_j w_j (prod_i P_i)`` to the coefficient MPS ``p``. Each branch scales a copy of ``p`` by ``w_j`` and applies its (bond-preserving) single-qubit Paulis; the branches are summed and compressed to ``chi``/``cutoff``. Unitary sums return the cumulative - norm-loss proxy; arbitrary non-unitary sums invalidate that diagnostic. + norm-loss proxy. Dense non-unitary sums return the retained norm ratio + when a physical ``G^dagger G`` target norm was supplied. """ p = self.state.p branches = tuple(branches) @@ -4993,6 +5118,8 @@ def _apply_operator_sum(self, branches, *, unitary: bool) -> Optional[float]: self._set_zero_coefficient_state() if unitary: return self._unitary_norm_infidelity() + if target_norm is not None: + return self._nonunitary_compression_infidelity(target_norm) self._invalidate_norm_infidelity() return None @@ -5039,6 +5166,8 @@ def build(max_bond): self.state.info["cur_orthog"] = (0, 0) if unitary: return self._unitary_norm_infidelity() + if target_norm is not None: + return self._nonunitary_compression_infidelity(target_norm) self._invalidate_norm_infidelity() return None diff --git a/src/pepsy/optimizers/sweep/README.md b/src/pepsy/optimizers/sweep/README.md index e94006c..dd0be3e 100644 --- a/src/pepsy/optimizers/sweep/README.md +++ b/src/pepsy/optimizers/sweep/README.md @@ -16,7 +16,9 @@ build_bra_ket(...) -> BdyMPS(...) -> CompBdy.move_bdy/move_step_bdy(...) ``` `boundary_engine="quimb-mps"` uses Quimb MPS environments for local row/column -boundaries. The adapter lives in `environments.py` and exposes the legacy +boundaries. At the start of each half-sweep it computes the opposite-side +environments once, then moves the active boundary one row or column at a time +and caches it. The adapter lives in `environments.py` and exposes the legacy surface expected by the current local objective: - `mps_b` @@ -28,6 +30,9 @@ surface expected by the current local objective: `boundary_engine="auto"` keeps dense inputs on the Pepsy path and routes Symmray-looking inputs to Quimb MPS. +Torch-backed Symmray arrays remain on Torch and use autograd local solvers; +NumPy-backed Symmray arrays use the finite-difference compatibility path. + ## Extraction map - `optimizer.py`: `SweepOptimizer` and the current sweep orchestration. diff --git a/src/pepsy/optimizers/sweep/environments.py b/src/pepsy/optimizers/sweep/environments.py index 14b86a2..1a09c79 100644 --- a/src/pepsy/optimizers/sweep/environments.py +++ b/src/pepsy/optimizers/sweep/environments.py @@ -10,6 +10,7 @@ "QuimbMpsBoundaryStore", "canonical_boundary_engine_selector", "normalize_boundary_engine", + "symmray_array_backends", "uses_symmray_arrays", ] @@ -70,6 +71,24 @@ def uses_symmray_arrays(*states): return False +def symmray_array_backends(*states): + """Return the array backends used by Symmray tensors in ``states``.""" + backends = set() + for state in states: + if state is None: + continue + tensor_map = getattr(state, "tensor_map", None) + tensors = tensor_map.values() if tensor_map else () + for tensor in tensors: + data = getattr(tensor, "data", None) + if not _is_symmray_array_data(data): + continue + backend = getattr(data, "backend", None) + if backend is not None: + backends.add(str(backend).lower()) + return frozenset(backends) + + def canonical_boundary_engine_selector(engine): """Canonicalize a boundary engine selector without resolving ``"auto"``.""" key = "auto" if engine is None else str(engine).strip().lower().replace("_", "-") @@ -206,9 +225,12 @@ def normalize(self): def _compute_kwargs(self, progress=False): opts = dict(self.contract_boundary_opts) - if "progress" in opts and "progbar" not in opts: - opts["progbar"] = opts.pop("progress") - opts.setdefault("progbar", bool(progress)) + # Quimb's compute_*_environments forwards kwargs into lower-level + # boundary contraction internals whose progress kwarg name differs + # across versions. Avoid forwarding either spelling here to keep the + # boundary refresh path version-agnostic. + opts.pop("progress", None) + opts.pop("progbar", None) opts.update( { "max_bond": self.chi, @@ -226,7 +248,12 @@ def _compute_kwargs(self, progress=False): return opts def update_axis(self, tn, axis, *, progress=False): - """Recompute Quimb environments for one sweep axis.""" + """Recompute Quimb environments for one sweep axis. + + The returned environment dictionary is passed back to Quimb so that + the store remains a real cache, rather than only retaining the result + after a full rebuild. + """ if axis == "y": compute_fn = getattr(tn, "compute_y_environments", None) elif axis == "x": @@ -241,13 +268,183 @@ def update_axis(self, tn, axis, *, progress=False): envs = _call_with_accepted_kwargs( compute_fn, + envs={}, **self._compute_kwargs(progress=progress), ) + axis_labels = {f"{axis}min", f"{axis}max"} + self.envs = { + key: value + for key, value in self.envs.items() + if not (isinstance(key, tuple) and key[0] in axis_labels) + } self.envs.update(envs) + self.mps_b = { + key: value + for key, value in self.mps_b.items() + if not str(key).startswith(axis.upper()) + } self._sync_axis_mps_b(tn, axis, envs) self.update_count += 1 return self + def start_sweep(self, tn, axis, update_side, *, progress=False): + """Prepare one half-sweep with one static and one moving boundary. + + Quimb computes the opposite-side environments once. The selected + side is then advanced one row or column at a time with + :meth:`advance_sweep`, after the local tensor update has been applied. + """ + if update_side not in {"left", "right"}: + raise ValueError("update_side must be 'left' or 'right'") + if axis not in {"x", "y"}: + raise ValueError("axis must be 'x' or 'y'") + + # The public Quimb names are ``compute_ymin_environments`` etc. + compute_fn = getattr(tn, f"compute_{axis}{'max' if update_side == 'left' else 'min'}_environments") + envs = _call_with_accepted_kwargs( + compute_fn, + envs={}, + **self._compute_kwargs(progress=progress), + ) + + axis_labels = {f"{axis}min", f"{axis}max"} + self.envs = { + key: value + for key, value in self.envs.items() + if not (isinstance(key, tuple) and key[0] in axis_labels) + } + self.envs.update(envs) + self.mps_b = { + key: value + for key, value in self.mps_b.items() + if not str(key).startswith(axis.upper()) + } + self._sync_axis_mps_b(tn, axis, envs) + self._sweep_axis = axis + self._sweep_update_side = update_side + self._sweep_moving_env = None + self.update_count += 1 + return self + + def _extend_moving_boundary(self, tn, index, *, axis, update_side): + """Add one updated row or column to the moving boundary MPS.""" + axis_tag = axis.upper() + column = tn.select(f"{axis_tag}{index}").copy() + if update_side == "left": + if self._sweep_moving_env is None: + return column + env = self._sweep_moving_env.copy() + env.retag_( + { + tag: f"{axis_tag}0" + for tag in tuple(env.tags) + if isinstance(tag, str) and tag.startswith(axis_tag) + } + ) + column.retag_({f"{axis_tag}{index}": f"{axis_tag}1"}) + first, second = env, column + direction = f"{axis}min" + else: + if self._sweep_moving_env is None: + return column + env = self._sweep_moving_env.copy() + env.retag_( + { + tag: f"{axis_tag}1" + for tag in tuple(env.tags) + if isinstance(tag, str) and tag.startswith(axis_tag) + } + ) + column.retag_({f"{axis_tag}{index}": f"{axis_tag}0"}) + first, second = column, env + direction = f"{axis}max" + + def _rebase_plane_tags(network, coordinate): + """Give a compressed boundary one logical 2D plane coordinate. + + ``contract_boundary_from_*`` identifies sites using ``I{x,y}`` + tags, not the ``X*`` / ``Y*`` tags alone. A moving environment + carries every absorbed plane's original site tags, so it must be + rebased before combining it with one new physical plane. + """ + tag_map = {} + for x in range(int(tn.Lx)): + for y in range(int(tn.Ly)): + old_tag = tn.site_tag(x, y) + if old_tag not in network.tags: + continue + new_tag = ( + tn.site_tag(x, coordinate) + if axis == "y" + else tn.site_tag(coordinate, y) + ) + tag_map[old_tag] = new_tag + if tag_map: + network.retag_(tag_map) + + # The two-plane temporary network always has coordinates 0 and 1. + # Rebase both the compressed environment and incoming physical plane + # so Quimb actually contracts the pair rather than merely returning + # the original boundary side unchanged. + _rebase_plane_tags(first, 0) + _rebase_plane_tags(second, 1) + + if axis == "x": + pair = (first | second).view_like(tn, Lx=2, Ly=tn.Ly) + xrange, yrange = (0, 1), (0, tn.Ly - 1) + else: + pair = (first | second).view_like(tn, Lx=tn.Lx, Ly=2) + xrange, yrange = (0, tn.Lx - 1), (0, 1) + + opts = self._compute_kwargs() + opts.pop("dense", None) + opts.pop("envs", None) + getattr(pair, f"contract_boundary_from_{direction}_")( + xrange=xrange, + yrange=yrange, + **opts, + ) + # Quimb keeps the boundary on the side it contracted *from*. For a + # ymin/xmin move that is the zero-labelled row/column, whereas a + # ymax/xmax return move leaves the boundary on the one-labelled side. + # Returning ``axis_tag0`` unconditionally made every backward cached + # step retain the uncontracted new slice instead of the compressed + # moving environment. + retained_index = 0 if update_side == "left" else 1 + return pair.select(f"{axis_tag}{retained_index}").copy() + + def advance_sweep(self, tn, index, *, axis=None, update_side=None): + """Advance the moving boundary by one updated row or column.""" + axis = self._sweep_axis if axis is None else axis + update_side = ( + self._sweep_update_side if update_side is None else update_side + ) + n = int(getattr(tn, "Ly" if axis == "y" else "Lx")) + if update_side == "left": + next_index = index + 1 + if next_index >= n: + return self + boundary_index = next_index + boundary_label = f"{axis}min" + else: + next_index = index - 1 + if next_index < 0: + return self + boundary_index = next_index + boundary_label = f"{axis}max" + + self._sweep_moving_env = self._extend_moving_boundary( + tn, + index, + axis=axis, + update_side=update_side, + ) + env_key = boundary_label, boundary_index + self.envs[env_key] = self._sweep_moving_env + self._sync_axis_mps_b(tn, axis, {env_key: self._sweep_moving_env}) + self.update_count += 1 + return self + def _sync_axis_mps_b(self, tn, axis, envs): if axis == "y": ly = int(getattr(tn, "Ly", 0)) diff --git a/src/pepsy/optimizers/sweep/optimizer.py b/src/pepsy/optimizers/sweep/optimizer.py index c930c4c..f837955 100644 --- a/src/pepsy/optimizers/sweep/optimizer.py +++ b/src/pepsy/optimizers/sweep/optimizer.py @@ -23,6 +23,7 @@ from .environments import ( QuimbMpsBoundaryStore, normalize_boundary_engine, + symmray_array_backends, uses_symmray_arrays, ) @@ -188,6 +189,24 @@ def _as_scaled_scalar(value, *, name="value"): return value[0], float(value[1]) return value, 0.0 + @staticmethod + def _unwrap_contracted_scalar(value, *, name): + """Extract the data of a scalar Quimb tensor without touching arrays. + + Quimb normally returns an array scalar from ``contract(all)`` but can + retain a zero-index ``Tensor`` wrapper for Symmray contractions. Its + data is the backend scalar (a Torch scalar in the differentiable + path); unwrapping only this wrapper preserves its autograd graph. + """ + if not isinstance(value, qtn.Tensor): + return value + if value.inds: + raise ValueError( + f"{name} contraction did not produce a scalar; " + f"remaining indices are {value.inds!r}." + ) + return value.data + @staticmethod def _network_exponent(tn): """Return the real base-10 exponent carried by a tensor network.""" @@ -223,12 +242,21 @@ def _scaled_overlap_fidelity(cls, overlap, norm, target_norm): overlap_m, overlap_e = cls._as_scaled_scalar(overlap, name="overlap") norm_m, norm_e = cls._as_scaled_scalar(norm, name="norm") target_m, target_e = cls._as_scaled_scalar(target_norm, name="target_norm") - - fid_m = (ar.do("abs", overlap_m) ** 2) / ( - ar.do("abs", norm_m) * ar.do("abs", target_m) + overlap_m = cls._unwrap_contracted_scalar(overlap_m, name="overlap") + norm_m = cls._unwrap_contracted_scalar(norm_m, name="norm") + target_m = cls._unwrap_contracted_scalar(target_m, name="target_norm") + + # ``contract(..., strip_exponent=True)`` can return Quimb's scalar + # wrapper even when the tensor blocks themselves are Torch. Autoray + # correctly dispatches most array operations here, but it has no + # ``quimb.abs`` registration for that wrapper. Python's ``abs`` + # calls the scalar / array ``__abs__`` implementation instead: it + # works for Quimb qarrays and retains autograd for Torch tensors. + fid_m = (abs(overlap_m) ** 2) / ( + abs(norm_m) * abs(target_m) ) fid_e = 2.0 * overlap_e - norm_e - target_e - return ar.do("abs", fid_m) * cls._safe_pow10(fid_e) + return abs(fid_m) * cls._safe_pow10(fid_e) def _local_norm_exponent(self): """Exponent for the represented local norm environment.""" @@ -350,6 +378,7 @@ def __init__( track_boundary_fidelity: bool | None = None, ): self._ensure_no_common_internal_indices(state, state_target) + self._validate_symmray_input_backends(state, state_target) bdy_obj, bdy_holder = self._resolve_boundary_arg(bdy, "bdy") bdy_overlap_obj, bdy_overlap_holder = self._resolve_boundary_arg( @@ -397,7 +426,27 @@ def __init__( # Store chi as-is (may be int or (int, int) tuple). self.chi = chi if chi is not None else getattr(bdy_obj, "chi", None) self.single_layer = single_layer - self.simplify = simplify + self.simplify = bool(simplify) + # Symmray block-sparse states need special handling in the local slice + # optimizer. NumPy-backed blocks use the finite-difference fallback, + # while Torch-backed blocks remain on Torch and use autograd normally. + self._symmray_state = bool(uses_symmray_arrays(self.state, self.state_target)) + self._symmray_backends = symmray_array_backends( + self.state, + self.state_target, + ) + self._symmray_torch = ( + self._symmray_state and self._symmray_backends == {"torch"} + ) + self._symmray_requires_fd = self._symmray_state and not self._symmray_torch + if self.simplify and self._symmray_state: + warnings.warn( + "Disabling sweep local full_simplify for Symmray-backed states " + "to avoid backend incompatibilities.", + UserWarning, + stacklevel=2, + ) + self.simplify = False self.direction = direction if direction is not None else "y" self.max_separation = max_separation if max_separation is not None else 1 @@ -409,6 +458,7 @@ def __init__( self.Lx, self.Ly = self._infer_shape(self.state) self._reset_run_traces() + self._warned_nested_param_tree = False init_renormalize_kwargs = self._collect_init_renormalize_kwargs( chi=self.chi, @@ -444,6 +494,27 @@ def _ensure_no_common_internal_indices(state, state_target): warnings.warn(msg, UserWarning, stacklevel=3) raise ValueError(msg) + @staticmethod + def _validate_symmray_input_backends(*states): + """Reject mixed dense/Symmray and mixed-backend sweep objectives.""" + states = tuple(state for state in states if state is not None) + symmray_flags = tuple(uses_symmray_arrays(state) for state in states) + if not any(symmray_flags): + return + if not all(symmray_flags): + raise TypeError( + "SweepOptimizer cannot mix Symmray and dense state/target " + "tensor networks. Convert both inputs to the same backend." + ) + backends = symmray_array_backends(*states) + if len(backends) != 1: + found = ", ".join(sorted(backends)) or "unknown" + raise TypeError( + "Symmray sweep inputs must use one common array backend; " + f"found {found}. Convert both state and state_target before " + "constructing SweepOptimizer." + ) + @staticmethod def _call_with_accepted_kwargs(fn, *args, **kwargs): """Call ``fn`` while dropping unsupported keyword arguments.""" @@ -1149,6 +1220,26 @@ def _bra_with_reindexed_inner(local_tn): ) return bra + def _overlap_bra(self, local_tn): + """Build a local overlap bra matching the cached double-layer indices. + + ``build_bra_ket(ket=state_target, bra=state)`` renames only the + bra-side internal indices that collide with ``state_target``. A target + with mangled internal bonds therefore retains the state's original + labels on the bra layer. The local objective must follow that same + convention in order to attach cached overlap boundaries correctly. + """ + bra = local_tn.conj() + shared_inner = set(self.state_target.inner_inds()) & set(self.state.inner_inds()) + bra.reindex_( + { + idx: f"{idx}_*" + for idx in bra.ind_map + if idx in shared_inner + } + ) + return bra + @staticmethod def _resolve_user_solver(solver): """Validate solver names and emit practical warnings.""" @@ -1170,6 +1261,167 @@ def _resolve_user_solver(solver): return solver + @staticmethod + def _needs_nested_param_flatten(params): + """Return whether ``params`` contains nested Mapping/list/tuple values.""" + if not isinstance(params, Mapping): + return False + return any(isinstance(value, (Mapping, list, tuple)) for value in params.values()) + + @classmethod + def _flatten_param_tree(cls, params): + """Flatten a nested parameter pytree into leaf mapping and structure spec.""" + if not isinstance(params, Mapping): + raise TypeError("params must be a mapping") + + flat = {} + spec = {} + + def _walk(path, value): + if isinstance(value, Mapping): + keys = tuple(value.keys()) + spec[path] = ("mapping", keys) + for key in keys: + _walk(path + (("k", key),), value[key]) + return + if isinstance(value, list): + spec[path] = ("list", len(value)) + for idx, item in enumerate(value): + _walk(path + (("i", idx),), item) + return + if isinstance(value, tuple): + spec[path] = ("tuple", len(value)) + for idx, item in enumerate(value): + _walk(path + (("i", idx),), item) + return + flat[path] = value + + _walk((), params) + return flat, spec + + @classmethod + def _unflatten_param_tree(cls, flat, spec): + """Rebuild a nested parameter pytree from flattened leaves and spec.""" + if not isinstance(flat, Mapping): + raise TypeError("flat must be a mapping") + + def _build(path): + if path in spec: + kind, payload = spec[path] + if kind == "mapping": + return { + key: _build(path + (("k", key),)) + for key in payload + } + if kind == "list": + return [ + _build(path + (("i", idx),)) + for idx in range(payload) + ] + if kind == "tuple": + return tuple( + _build(path + (("i", idx),)) + for idx in range(payload) + ) + raise ValueError(f"Unknown tree container kind: {kind!r}") + return flat[path] + + return _build(()) + + @classmethod + def _clone_param_tree(cls, value): + """Detach/clone tensor leaves recursively while preserving tree shape.""" + if isinstance(value, Mapping): + return {key: cls._clone_param_tree(val) for key, val in value.items()} + if isinstance(value, list): + return [cls._clone_param_tree(val) for val in value] + if isinstance(value, tuple): + return tuple(cls._clone_param_tree(val) for val in value) + if hasattr(value, "detach"): + return value.detach().clone() + copy_fn = getattr(value, "copy", None) + if callable(copy_fn): + try: + return copy_fn() + except Exception: # pragma: no cover - defensive copy fallback + pass + return value + + @staticmethod + def _match_leaf_backend(value, reference): + """Coerce ``value`` to the array backend of ``reference``. + + The local gradient solver flattens parameters to Torch tensors even + when the input state uses a different backend (for example NumPy-backed + Symmray blocks). Writing those Torch leaves back into an otherwise + NumPy state produces a mixed-backend network that breaks subsequent + boundary contractions. This restores each optimized leaf to the backend + it came from so the state's array backend is preserved. + """ + try: + ref_backend = ar.infer_backend(reference) + val_backend = ar.infer_backend(value) + except Exception: # pragma: no cover - defensive backend inference + return value + if ref_backend == val_backend: + return value + try: + numpy_value = ar.to_numpy(value) + except Exception: # pragma: no cover - fall back to raw value + return value + if ref_backend == "numpy": + return numpy_value + try: + return ar.do("array", numpy_value, like=reference) + except Exception: # pragma: no cover - last-resort passthrough + return value + + @classmethod + def _restore_leaf_backends(cls, params_opt, params_ref): + """Coerce optimized flat leaves back to their original backends.""" + if not isinstance(params_opt, Mapping) or not isinstance(params_ref, Mapping): + return params_opt + restored = {} + for key, value in params_opt.items(): + reference = params_ref.get(key) + if reference is None: + restored[key] = value + else: + restored[key] = cls._match_leaf_backend(value, reference) + return restored + + @staticmethod + def _coerce_leaf_to_numpy(value): + """Convert a single array leaf to a NumPy array when needed.""" + try: + backend = ar.infer_backend(value) + except Exception: # pragma: no cover - defensive backend inference + return value + if backend == "numpy": + return value + try: + return ar.to_numpy(value) + except Exception: # pragma: no cover - last-resort passthrough + return value + + @classmethod + def _coerce_param_tree_numpy(cls, value): + """Recursively coerce array leaves in a param tree to NumPy. + + The autograd solvers hand the local loss function Torch tensors even + for NumPy-backed Symmray states. Contracting those against the NumPy + target/boundary tensors would mix backends inside Symmray's block + ``tensordot``. Coercing every leaf to NumPy keeps the local contraction + on a single backend so gradients can be taken by finite differences. + """ + if isinstance(value, Mapping): + return {key: cls._coerce_param_tree_numpy(val) for key, val in value.items()} + if isinstance(value, list): + return [cls._coerce_param_tree_numpy(val) for val in value] + if isinstance(value, tuple): + return tuple(cls._coerce_param_tree_numpy(val) for val in value) + return cls._coerce_leaf_to_numpy(value) + def _estimate_slice_contraction_metrics( self, *, @@ -1182,7 +1434,7 @@ def _estimate_slice_contraction_metrics( """Estimate FLOP/peak complexity for local norm and overlap networks.""" local0 = qtn.unpack(params_init, skeleton) bra_norm0 = self._bra_with_reindexed_inner(local0) - bra_overlap0 = local0.conj() + bra_overlap0 = self._overlap_bra(local0) norm_net0 = self._attach_boundaries( local0 | bra_norm0, self.bdy.mps_b, @@ -1212,6 +1464,37 @@ def _estimate_slice_contraction_metrics( "peak_overlap": peak_overlap, } + @staticmethod + def _finite_difference_solver(solver, opts): + """Route a non-autograd local objective to a finite-difference solver. + + An explicitly requested ``fd-*`` solver is preserved. Autograd solver + requests map to dependency-free ``fd-adam``; explicit SciPy and NLopt + requests keep their respective finite-difference implementations so + missing optional dependencies produce their normal, actionable errors. + """ + key = str(solver).strip().lower() + if key.startswith("fd-"): + fd_solver = key + elif key.startswith(("torch", "jax")): + fd_solver = "fd-adam" + elif "scipy" in key: + fd_solver = "fd-scipy" + elif "nlopt" in key or key.startswith(("ld_", "ln_", "gd_", "gn_")): + fd_solver = "fd-nlopt" + else: + fd_solver = "fd-adam" + opts = dict(opts) + # ``lr`` is meaningful for finite-difference Adam, but not for the + # SciPy/NLopt finite-difference backends. + if fd_solver != "fd-adam": + opts.pop("lr", None) + return fd_solver, opts + + # Kept as a private compatibility alias for callers/tests introduced with + # the original Symmray-only finite-difference fallback. + _symmray_fd_solver = _finite_difference_solver + def _optimize_packed_params( self, params_init, @@ -1221,6 +1504,17 @@ def _optimize_packed_params( solver_options=None, ): opts = self._merge_solver_options(solver_options) + if getattr(self, "_active_local_requires_finite_differences", False): + solver, opts = self._finite_difference_solver(solver, opts) + elif ( + isinstance(solver, str) + and ("scipy" in solver.lower() or "nlopt" in solver.lower()) + and "lr" not in dict(solver_options or {}) + ): + # ``lr`` is an Adam-style default, not an NLopt/SciPy control. + # Do not emit an irrelevant warning unless the caller explicitly + # requested it. + opts.pop("lr", None) n_steps = int(opts.pop("n_steps", 30)) # Allow callers to request a per-step gradient progress bar by setting # progress=True inside optimizer_options. We pop it here so it never @@ -1236,6 +1530,64 @@ def _optimize_packed_params( result = runner.run(params_init=params_init, loss_fn=loss_fn) return result.params, result.history + @classmethod + def _param_tree_backends(cls, value): + """Return inferred array backends for all leaves of a param tree.""" + if isinstance(value, Mapping): + backends = set() + for child in value.values(): + backends.update(cls._param_tree_backends(child)) + return frozenset(backends) + if isinstance(value, (list, tuple)): + backends = set() + for child in value: + backends.update(cls._param_tree_backends(child)) + return frozenset(backends) + try: + return frozenset({str(ar.infer_backend(value)).lower()}) + except Exception: # pragma: no cover - non-array metadata leaf + return frozenset() + + @classmethod + def _params_require_finite_differences(cls, params): + """Return whether local params cannot use Torch autograd safely.""" + backends = cls._param_tree_backends(params) + if not backends or backends == {"torch"}: + return False + if len(backends) != 1: + found = ", ".join(sorted(backends)) + raise TypeError( + "Local sweep parameters must use one common array backend; " + f"found {found}." + ) + return True + + @classmethod + def _match_param_tree_backends(cls, value, reference): + """Recursively restore ``value`` leaves to ``reference`` backends.""" + if isinstance(reference, Mapping): + if not isinstance(value, Mapping): + raise TypeError("Parameter-tree structure changed during optimization.") + return { + key: cls._match_param_tree_backends(value[key], ref) + for key, ref in reference.items() + } + if isinstance(reference, list): + if not isinstance(value, list) or len(value) != len(reference): + raise TypeError("Parameter-tree structure changed during optimization.") + return [ + cls._match_param_tree_backends(item, ref) + for item, ref in zip(value, reference) + ] + if isinstance(reference, tuple): + if not isinstance(value, tuple) or len(value) != len(reference): + raise TypeError("Parameter-tree structure changed during optimization.") + return tuple( + cls._match_param_tree_backends(item, ref) + for item, ref in zip(value, reference) + ) + return cls._match_leaf_backend(value, reference) + def _apply_slice_update(self, index, params_opt, skeleton, axis): tn_opt = qtn.unpack(params_opt, skeleton) if not uses_symmray_arrays(tn_opt): @@ -1258,9 +1610,32 @@ def _optimize_axis_slice_with_current_env( slice_state = self.state.select([f"{axis_tag}{index}"], "any") slice_target = self.state_target.select([f"{axis_tag}{index}"], "any") - params_init, skeleton = qtn.pack(slice_state) + params_init_tree, skeleton = qtn.pack(slice_state) + requires_finite_differences = self._params_require_finite_differences( + params_init_tree + ) + if self._needs_nested_param_flatten(params_init_tree): + params_init_opt, param_tree_spec = self._flatten_param_tree(params_init_tree) + if not self._warned_nested_param_tree: + warnings.warn( + "Detected nested local slice params (for example Symmray " + "block trees); enabling tree-flattened sweep optimization path.", + UserWarning, + stacklevel=2, + ) + self._warned_nested_param_tree = True + + def _restore_params(params_in): + return self._unflatten_param_tree(params_in, param_tree_spec) + + else: + params_init_opt = params_init_tree + + def _restore_params(params_in): + return params_in + metrics = self._estimate_slice_contraction_metrics( - params_init=params_init, + params_init=params_init_tree, skeleton=skeleton, slice_target=slice_target, right_key=right_key, @@ -1268,9 +1643,20 @@ def _optimize_axis_slice_with_current_env( ) def loss_fn(params_in): - local = qtn.unpack(params_in, skeleton) + params_tree = _restore_params(params_in) + if requires_finite_differences: + # Finite-difference solvers use Torch work tensors internally. + # Rebuild the local network with the original array backend + # before it touches cached target/boundary environments. + params_tree = self._match_param_tree_backends( + params_tree, + params_init_tree, + ) + local = qtn.unpack(params_tree, skeleton) bra_norm = self._bra_with_reindexed_inner(local) - bra_overlap = local.conj() + # Keep the local overlap bra index convention identical to the + # full double layer used to build ``bdy_overlap``. + bra_overlap = self._overlap_bra(local) norm_net = self._attach_boundaries( local | bra_norm, @@ -1317,21 +1703,32 @@ def loss_fn(params_in): infid = 1. - fid return infid - initial_loss = float(loss_fn(params_init)) + initial_loss = float(loss_fn(params_init_opt)) - params_opt, history = self._optimize_packed_params( - params_init, - loss_fn, - solver=solver, - solver_options=solver_options, + previous_requires_fd = getattr( + self, + "_active_local_requires_finite_differences", + False, ) + self._active_local_requires_finite_differences = requires_finite_differences + try: + params_opt, history = self._optimize_packed_params( + params_init_opt, + loss_fn, + solver=solver, + solver_options=solver_options, + ) + finally: + self._active_local_requires_finite_differences = previous_requires_fd history_values = self._to_float_history(history) final_loss = history_values[-1] if history_values else initial_loss - params_opt = { - k: v.detach().clone() if hasattr(v, "detach") else v - for k, v in params_opt.items() - } - self._apply_slice_update(index, params_opt, skeleton, axis) + params_opt_tree = _restore_params(params_opt) + params_opt_tree = self._match_param_tree_backends( + params_opt_tree, + params_init_tree, + ) + params_opt_tree = self._clone_param_tree(params_opt_tree) + self._apply_slice_update(index, params_opt_tree, skeleton, axis) # Track the best state using the minimum non-negative loss observed # during this local gradient optimization. self._maybe_store_best_state(self._best_nonnegative_from_history(history_values)) @@ -1380,6 +1777,21 @@ def _refresh_quimb_axis_boundaries(self, norm_tn, overlap_tn, axis, *, progress= self.bdy_overlap.update_axis(overlap_tn, axis, progress=progress) return {"norm": None, "overlap": None} + def _update_quimb_double_layer_slice(self, norm_tn, overlap_tn, index, axis): + """Update the cached double layers after one local slice optimization.""" + for site_tag in self._site_tensor_tags(axis, index): + source = next(iter(self.state.select(site_tag).tensor_map.values())) + data = source.data + + for tn, layer, layer_data in ( + (norm_tn, "KET", data), + (norm_tn, "BRA", data.conj()), + (overlap_tn, "BRA", data.conj()), + ): + selected = tn.select([site_tag, layer], "all") + for tensor in selected.tensor_map.values(): + tensor.modify(data=layer_data) + def _advance_boundary_one_step( self, index, @@ -1432,29 +1844,73 @@ def _run_axis_half_sweep( # pylint: disable=too-many-arguments,too-many-locals debug_loss_kwargs=None, ): """Run a single forward or backward half-sweep over *indices*.""" + indices = tuple(indices) runs = [] comp_norm = None comp_overlap = None uses_quimb = getattr(self, "boundary_engine", "dmrg") == "quimb-mps" - for index in indices: + if uses_quimb: norm_tn, overlap_tn = self._prepare_current_double_layers() - if uses_quimb: - comp_norm = None - comp_overlap = None - elif comp_norm is None: - comp_norm, comp_overlap = self._make_comp_pair(norm_tn, overlap_tn) - else: - self._set_comp_norms(comp_norm, comp_overlap, norm_tn=norm_tn, overlap_tn=overlap_tn) - - t0 = time.perf_counter() - if uses_quimb: - boundary_fidelity = self._refresh_quimb_axis_boundaries( + self.bdy.start_sweep(norm_tn, axis, update_side, progress=False) + self.bdy_overlap.start_sweep( + overlap_tn, + axis, + update_side, + progress=False, + ) + # A return-forward pass starts at ``1`` because site ``0`` was + # just optimized by the preceding backward pass. Seed the + # moving left/bottom boundary with that fixed endpoint before + # the first local objective is assembled. + if update_side == "left" and indices and indices[0] > 0: + self.bdy.advance_sweep( norm_tn, + 0, + axis=axis, + update_side=update_side, + ) + self.bdy_overlap.advance_sweep( overlap_tn, - axis, - progress=False, + 0, + axis=axis, + update_side=update_side, + ) + # A round-trip backward pass starts at ``n - 2`` because the + # endpoint was just optimized by the preceding forward pass. Seed + # the moving right/top boundary with that fixed endpoint before + # the first local objective is assembled. + if update_side == "right" and indices and indices[0] < self._axis_n(axis) - 1: + endpoint = self._axis_n(axis) - 1 + self.bdy.advance_sweep( + norm_tn, + endpoint, + axis=axis, + update_side=update_side, ) + self.bdy_overlap.advance_sweep( + overlap_tn, + endpoint, + axis=axis, + update_side=update_side, + ) + + for index in indices: + if not uses_quimb: + norm_tn, overlap_tn = self._prepare_current_double_layers() + if comp_norm is None: + comp_norm, comp_overlap = self._make_comp_pair(norm_tn, overlap_tn) + else: + self._set_comp_norms( + comp_norm, + comp_overlap, + norm_tn=norm_tn, + overlap_tn=overlap_tn, + ) + + t0 = time.perf_counter() + if uses_quimb: + boundary_fidelity = {"norm": None, "overlap": None} else: boundary_fidelity = self._advance_boundary_one_step( index, @@ -1482,6 +1938,28 @@ def _run_axis_half_sweep( # pylint: disable=too-many-arguments,too-many-locals ) t_opt = time.perf_counter() - t0 + if uses_quimb: + t0 = time.perf_counter() + self._update_quimb_double_layer_slice( + norm_tn, + overlap_tn, + index, + axis, + ) + self.bdy.advance_sweep( + norm_tn, + index, + axis=axis, + update_side=update_side, + ) + self.bdy_overlap.advance_sweep( + overlap_tn, + index, + axis=axis, + update_side=update_side, + ) + t_bdy += time.perf_counter() - t0 + run_info["sweep"] = sweep_name run_info["time_boundary"] = t_bdy run_info["time_optimize"] = t_opt @@ -1563,7 +2041,8 @@ def optimize_axis( old_norm = self._normalize_state(env_n_iter) self.norm_trace.append({"state_norm": float(abs(complex(old_norm)))}) - self._refresh_right_boundaries_once(axis, env_n_iter=env_n_iter) + if getattr(self, "boundary_engine", "dmrg") != "quimb-mps": + self._refresh_right_boundaries_once(axis, env_n_iter=env_n_iter) sweep_kwargs = dict( axis=axis, diff --git a/src/pepsy/optimizers/sym_dmrg.py b/src/pepsy/optimizers/sym_dmrg.py index 3b06ea7..719419c 100644 --- a/src/pepsy/optimizers/sym_dmrg.py +++ b/src/pepsy/optimizers/sym_dmrg.py @@ -767,6 +767,13 @@ def __init__(self, optimizer, left, right_inds, *, layout="unfused"): self.compiled_block_plan_output_blocks = 0 self.compiled_block_plan_mode = None self.compiled_block_plan_disabled_reason = None + # Once a plan has been matched against the first result layout, the + # local projected problem is immutable until the plan is discarded. + # Avoid rebuilding the same sector/shape/dtype dictionary on every + # Krylov matvec, while retaining one validation at plan activation. + self.compiled_block_plan_layout_frozen = False + self.compiled_block_plan_layout_checks = 0 + self.compiled_block_plan_layout_fastpath_uses = 0 if not self.shared: self.left_axis = None @@ -851,10 +858,16 @@ def _fused_output_inds(self, inds): return tuple(ind for ind in fused_inds if ind != self.fused_shared_ind) def _compiled_plan_matches(self, right): - return ( - self.compiled_block_plan is not None - and self.compiled_right_layout == _block_data_layout_map(right.data) - ) + if self.compiled_block_plan is None or self.compiled_right_layout is None: + return False + if self.compiled_block_plan_layout_frozen: + self.compiled_block_plan_layout_fastpath_uses += 1 + return True + self.compiled_block_plan_layout_checks += 1 + matches = self.compiled_right_layout == _block_data_layout_map(right.data) + if matches: + self.compiled_block_plan_layout_frozen = True + return matches def _can_compile_block_plan(self, right): if _is_fermionic_symmray_array(self.left.data) or _is_fermionic_symmray_array( @@ -872,6 +885,14 @@ def _can_compile_block_plan(self, right): return True def _compile_block_plan(self, right, output): + # A previous plan may have been invalidated by a changed block layout. + # 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_right_layout = None + self.compiled_output_template = None + self.compiled_block_plan_layout_frozen = False + self.compiled_block_plan_mode = None if not self._can_compile_block_plan(right): return @@ -955,6 +976,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_layout_frozen = False 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) @@ -1139,6 +1161,15 @@ def summary(self, prefix): self.compiled_block_plan_output_blocks ), f"{prefix}_compiled_block_plan_mode": self.compiled_block_plan_mode, + f"{prefix}_compiled_block_plan_layout_frozen": bool( + self.compiled_block_plan_layout_frozen + ), + f"{prefix}_compiled_block_plan_layout_checks": int( + self.compiled_block_plan_layout_checks + ), + f"{prefix}_compiled_block_plan_layout_fastpath_uses": int( + self.compiled_block_plan_layout_fastpath_uses + ), f"{prefix}_compiled_block_plan_disabled_reason": ( self.compiled_block_plan_disabled_reason ), @@ -1156,6 +1187,13 @@ def __init__(self, optimizer, site, theta): self.block_layout = self._block_layout(theta) self.layout = optimizer.matvec_layout self.input_map = {ind: optimizer._input_ind(ind) for ind in self.inds} + # Reindexing only changes the Symmray index metadata. Cache that + # metadata once and reuse the current theta blocks on every matvec. + # The block layout is guarded by ``matches`` below, so this does not + # widen or otherwise alter the local variational space. + theta_input_template = theta.reindex(self.input_map, inplace=False) + self.theta_input_inds = tuple(theta_input_template.inds) + self.theta_input_data_template = theta_input_template.data self.output_zeros = { sector: np.zeros_like(_to_numpy(block)) for sector, block in theta.data.blocks.items() @@ -1181,8 +1219,6 @@ def __init__(self, optimizer, site, theta): if right_env is not None else w_right ) - theta_input_inds = tuple(self.input_map.get(ind, ind) for ind in self.inds) - self.theta_input_inds = theta_input_inds left_stats = self._tensor_block_stats(self.left_projector, "left_projector") right_stats = self._tensor_block_stats(self.right_projector, "right_projector") # Keep the original right-first route unless the static projector @@ -1192,7 +1228,7 @@ def __init__(self, optimizer, site, theta): self.left_contraction = _BlockPairContraction( optimizer, self.left_projector, - theta_input_inds, + self.theta_input_inds, layout=self.layout, ) self.right_contraction = _BlockPairContraction( @@ -1206,7 +1242,7 @@ def __init__(self, optimizer, site, theta): self.right_contraction = _BlockPairContraction( optimizer, self.right_projector, - theta_input_inds, + self.theta_input_inds, layout=self.layout, ) self.left_contraction = _BlockPairContraction( @@ -1260,7 +1296,17 @@ def _copy_block(block): def apply(self, theta, *, timings=None): step_start = time.perf_counter() if timings is not None else None - theta_in = theta.reindex(self.input_map, inplace=False) + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + theta_input_data = _array_with_blocks_like( + self.theta_input_data_template, + theta.data.blocks, + ) + theta_in = qtn.Tensor( + data=theta_input_data, + inds=self.theta_input_inds, + tags=theta.tags, + ) _add_elapsed(timings, "matvec_input_reindex_elapsed", step_start) if self.contraction_order == "left_first": step_start = time.perf_counter() if timings is not None else None @@ -1293,7 +1339,8 @@ def apply(self, theta, *, timings=None): ) _add_elapsed(timings, "matvec_left_contract_elapsed", step_start) step_start = time.perf_counter() if timings is not None else None - out = out.transpose(*self.inds) + if tuple(out.inds) != self.inds: + out = out.transpose(*self.inds) _add_elapsed(timings, "matvec_transpose_elapsed", step_start) step_start = time.perf_counter() if timings is not None else None @@ -2082,7 +2129,13 @@ def _tensor_block_stats(tensor): } def profile_summary(self): - """Return aggregate timing/count information for profiling events.""" + """Return aggregate timing/count information for profiling events. + + ``total_elapsed`` is retained as the cumulative sum of all phase + timers, including nested timers. ``wall_elapsed`` and + ``phase_wall_fractions`` use the top-level solve timer and are the + appropriate values for wall-time attribution. + """ phase_totals = {} phase_counts = {} matvec_timing_totals = {} @@ -2097,6 +2150,14 @@ def profile_summary(self): matvec_timing_totals.get(key, 0.0) + float(value) ) total_elapsed = sum(phase_totals.values()) + wall_elapsed = float(phase_totals.get("solve", 0.0)) + if wall_elapsed > 0.0: + phase_wall_fractions = { + phase: float(elapsed / wall_elapsed) + for phase, elapsed in phase_totals.items() + } + else: + phase_wall_fractions = {phase: None for phase in phase_totals} lanczos_matvecs = [ int(entry["num_matvecs"]) for entry in self.local_solve_diagnostics @@ -2114,7 +2175,9 @@ def profile_summary(self): "enabled": self.profile, "num_events": len(self.profile_diagnostics), "total_elapsed": float(total_elapsed), + "wall_elapsed": wall_elapsed, "phase_totals": phase_totals, + "phase_wall_fractions": phase_wall_fractions, "phase_counts": phase_counts, "matvec_timing_totals": matvec_timing_totals, "num_matvecs": phase_counts.get("matvec", 0), @@ -2475,9 +2538,19 @@ def _apply_bond_schedule_convergence_gate(self, diagnostic, active_bond_dim): active_bond_dim = int(active_bond_dim) # Energy deltas during a ramp compare different variational spaces. Wait # until at least one previous sweep has also used the final bond target. + # An explicit schedule can already contain held final-bond sweeps, so + # don't require an additional repeated sweep beyond that schedule. + previous_bond_dim = None + if sweep > 0: + previous_bond_dim = ( + self.bond_dims[sweep - 1] + if sweep - 1 < len(self.bond_dims) + else final_bond_dim + ) + previous_bond_is_final = previous_bond_dim == final_bond_dim schedule_ready = bool( active_bond_dim >= final_bond_dim - and sweep >= len(self.bond_dims) + and previous_bond_is_final ) diagnostic.update( { @@ -2486,6 +2559,8 @@ def _apply_bond_schedule_convergence_gate(self, diagnostic, active_bond_dim): "active_bond_dim": active_bond_dim, "final_bond_dim": final_bond_dim, "bond_schedule_length": len(self.bond_dims), + "previous_bond_dim": previous_bond_dim, + "previous_bond_is_final": previous_bond_is_final, } ) if raw_converged and not schedule_ready: @@ -4167,7 +4242,43 @@ def _prune_theta_to_projected_support(self, site, theta): profile_start = self._profile_start() problem, _ = self._get_projected_problem(site, theta) - live_sectors = problem.structural_output_sectors() & set(theta.data.blocks) + all_sectors = set(theta.data.blocks) + live_sectors = problem.structural_output_sectors() & all_sectors + + # If every current sector is structurally reachable, numerical block + # norms cannot remove anything. This is the common path during the + # held-bond phase, and avoiding the scan matters because it runs before + # every local eigensolve. + if live_sectors == all_sectors: + diagnostic = { + "site": int(site), + "right_site": int(site + 1), + "input_blocks": len(all_sectors), + "kept_blocks": len(all_sectors), + "removed_blocks": 0, + "input_dim": self._theta_dim(theta), + "kept_dim": self._theta_dim(theta), + "live_blocks": len(live_sectors), + "nonzero_input_blocks": None, + "nonzero_input_scan_skipped": True, + "method": "structural", + } + self._record_profile_elapsed( + "variational_sector_prune", + profile_start, + site=int(site), + right_site=int(site + 1), + input_blocks=diagnostic["input_blocks"], + kept_blocks=diagnostic["kept_blocks"], + removed_blocks=0, + input_dim=diagnostic["input_dim"], + kept_dim=diagnostic["kept_dim"], + live_blocks=diagnostic["live_blocks"], + method=diagnostic["method"], + nonzero_input_scan_skipped=True, + ) + return theta, diagnostic + nonzero_input_sectors = { sector for sector, block in theta.data.blocks.items() if self._block_norm(block) > 0.0 @@ -4184,6 +4295,7 @@ def _prune_theta_to_projected_support(self, site, theta): "kept_dim": self._theta_dim(theta), "live_blocks": len(live_sectors), "nonzero_input_blocks": len(nonzero_input_sectors), + "nonzero_input_scan_skipped": False, "method": "structural", } self._record_profile_elapsed( @@ -4220,6 +4332,7 @@ def _prune_theta_to_projected_support(self, site, theta): "kept_dim": self._theta_dim(pruned), "live_blocks": len(live_sectors), "nonzero_input_blocks": len(nonzero_input_sectors), + "nonzero_input_scan_skipped": False, "method": "structural", } self._record_profile_elapsed( diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index fb80602..8183b3f 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -17,22 +17,299 @@ 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. Binary trees remain the default and a valid special case -(``max_arity=2``). +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. """ from __future__ import annotations +from collections.abc import Mapping +from numbers import Integral + +import autoray as ar +import numpy as np + from ..mps.layout import ( _gate_stream_adjacency, + _gate_stream_event_weights, _gate_stream_pair_weights, _gate_stream_spectral_order, _normalize_layout_gate_queue, _normalize_layout_support, + _normalize_weight_mode, ) +from ..mps.optimizer import _control_event_parts as _mps_control_event_parts __all__ = ["TreePlan", "TreeLayoutFinder"] +_DEFAULT_MAX_ARITY = object() +_DEFAULT_CHI = object() +_DEFAULT_SEARCH_OPTION = object() + + +def _looks_like_tree_tensor_network(value): + """Identify a TTN input without importing ``ttn`` (avoids a cycle).""" + return ( + getattr(value, "plan", None) is not None + and getattr(value, "tensor_map", None) is not None + ) + + +def _normalize_hybrid_weights(weights): + """Validate path / peak-load / total-load hybrid objective weights.""" + if weights is None: + values = (1.0, 1.0, 0.25) + elif isinstance(weights, Mapping): + aliases = { + "path": "path", + "max_edge_load": "max_edge_load", + "peak_load": "max_edge_load", + "total_edge_load": "total_edge_load", + "total_load": "total_edge_load", + } + normalized = {} + for key, value in weights.items(): + name = aliases.get(str(key).replace("-", "_").strip().lower()) + if name is None: + raise ValueError( + "hybrid_weights keys must be 'path', 'max_edge_load', " + "or 'total_edge_load'." + ) + if name in normalized: + raise ValueError(f"duplicate hybrid weight {name!r}.") + normalized[name] = value + values = tuple( + normalized.get(name, 0.0) + for name in ("path", "max_edge_load", "total_edge_load") + ) + else: + try: + values = tuple(weights) + except TypeError as exc: + raise ValueError( + "hybrid_weights must be a three-item sequence, mapping, or None." + ) from exc + if len(values) != 3: + raise ValueError( + "hybrid_weights must contain path, max-edge-load, and " + "total-edge-load weights." + ) + try: + values = tuple(float(value) for value in values) + except (TypeError, ValueError) as exc: + raise ValueError("hybrid_weights must be finite non-negative numbers.") from exc + if any(not np.isfinite(value) or value < 0.0 for value in values): + raise ValueError("hybrid_weights must be finite non-negative numbers.") + if not any(values): + raise ValueError("at least one hybrid weight must be positive.") + return values + + +def _normalize_layout_refinement(refine): + """Normalize an optional deterministic fixed-plan refinement mode.""" + if refine is None or refine is False: + return None + name = str(refine).replace("-", "_").strip().lower() + aliases = { + "adjacent": "greedy", + "adjacent_swaps": "greedy", + "local": "greedy", + } + name = aliases.get(name, name) + if name != "greedy": + raise ValueError("refine must be None or 'greedy'.") + return name + + +def _normalize_layout_search(search): + """Normalize an optional offline fixed-plan search mode.""" + if search is None or search is False: + return None + name = str(search).replace("-", "_").strip().lower() + aliases = {"ng": "nevergrad", "never_grad": "nevergrad"} + name = aliases.get(name, name) + if name != "nevergrad": + raise ValueError("search must be None or 'nevergrad'.") + return name + + +def _validate_search_budget(value, name): + """Validate a positive bounded layout-search evaluation budget.""" + try: + value = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be a positive integer.") from exc + if value < 1: + raise ValueError(f"{name} must be a positive integer.") + return value + + +def _safe_exp2(value): + """Return ``2**value`` without emitting overflow warnings.""" + if value > np.log2(np.finfo(float).max): + return float("inf") + return float(np.exp2(value)) + + +def _normalize_layout_objective(objective): + """Normalize the tree-layout objective name.""" + name = str(objective).replace("-", "_").strip().lower() + aliases = { + "distance": "path", + "path_length": "path", + "edge": "congestion", + "edge_load": "congestion", + "bond": "congestion", + "bond_load": "congestion", + "combined": "hybrid", + } + name = aliases.get(name, name) + if name not in {"path", "congestion", "hybrid"}: + raise ValueError( + f"Unknown tree layout objective {objective!r}. " + "Expected 'path', 'congestion', or 'hybrid'." + ) + return name + + +def _operator_schmidt_rank(payload, support, left_support): + """Return an operator-Schmidt rank across a support bipartition.""" + support = tuple(support) + left_support = tuple(left_support) + left_set = set(left_support) + if not left_set or left_set == set(support): + return 1 + try: + array = ar.to_numpy(payload) + except Exception: + return 2 + if array.size != 4 ** len(support): + return 2 + try: + array = array.reshape((2,) * (2 * len(support))) + positions = {site: pos for pos, site in enumerate(support)} + left_positions = [positions[site] for site in left_support] + right_positions = [ + positions[site] for site in support if site not in left_set + ] + axes = ( + left_positions + + [len(support) + pos for pos in left_positions] + + right_positions + + [len(support) + pos for pos in right_positions] + ) + matrix = array.transpose(axes).reshape( + 4 ** len(left_positions), + 4 ** len(right_positions), + ) + return max(1, int(np.linalg.matrix_rank(matrix))) + except (TypeError, ValueError, np.linalg.LinAlgError): + return 2 + + +def _submpo_schmidt_rank_bound(payload, support, left_support): + """Return an MPO-bond upper bound without calling ``to_dense``. + + An MPO's operator Schmidt rank across a site bipartition is bounded by the + product of the virtual MPO bonds crossing that bipartition. This is a + conservative diagnostic, but unlike lowering an MPO to a dense matrix it + remains cheap for wide supports. ``None`` means that ``payload`` does not + expose the Quimb MPO site interface. + """ + gen_sites = getattr(payload, "gen_sites_present", None) + site_tag = getattr(payload, "site_tag", None) + tag_map = getattr(payload, "tag_map", None) + tensor_map = getattr(payload, "tensor_map", None) + if not all((callable(gen_sites), callable(site_tag), + tag_map is not None, tensor_map is not None)): + return None + try: + present = tuple(gen_sites()) + support = tuple(support) + if set(present) != set(support): + return None + tensors = [] + for site in present: + tids = tuple(tag_map[site_tag(site)]) + if len(tids) != 1: + return None + tensors.append(tensor_map[tids[0]]) + left = set(left_support) + if not left or left == set(support): + return 1 + rank = 1 + for left_site, right_site, left_tensor, right_tensor in zip( + present, present[1:], tensors, tensors[1:] + ): + if (left_site in left) == (right_site in left): + continue + shared = set(left_tensor.inds).intersection(right_tensor.inds) + if len(shared) != 1: + return None + rank *= int(payload.ind_size(next(iter(shared)))) + return max(1, int(rank)) + except (AttributeError, KeyError, TypeError, ValueError): + return None + + +def _validate_chi(chi): + """Coerce and validate an optional ``chi`` selection budget.""" + if chi is None: + return None + chi = int(chi) + if chi < 1: + raise ValueError("chi must be a positive integer.") + return chi + + +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; + the representative single arity is what the legacy single-plan builders use + and is the first concrete candidate. + """ + if max_arity is None: + return None, None + if isinstance(max_arity, Integral): + return int(max_arity), None + if isinstance(max_arity, (str, bytes)): + return int(max_arity), None + if hasattr(max_arity, "__iter__"): + cand = [] + for a in max_arity: + key = None if a is None else int(a) + if key is not None and key < 2: + raise ValueError("arity candidates must be >= 2 or None.") + if key not in cand: + cand.append(key) + if not cand: + raise ValueError("max_arity iterable must be non-empty.") + representative = next((a for a in cand if a is not None), None) + return representative, tuple(cand) + return int(max_arity), None + + + +def _chi_cut_fields(plan, chi): + """Return ``{max_bond_cut, chi_overflow, exact_at_chi}`` for ``plan``. + + ``max_bond_cut`` is the widest qubit bipartition any bond induces. With a + finite ``chi`` the structure can hold an arbitrary state exactly only when + ``2 ** max_bond_cut <= chi``; ``chi_overflow`` is how many qubits the widest + bond exceeds ``log2(chi)`` (0 when the structure is exact at ``chi``). + """ + mbc = plan.max_bond_cut() + fields = {"max_bond_cut": mbc} + if chi is not None: + log_chi = float(np.log2(chi)) + fields["chi_overflow"] = max(0.0, mbc - log_chi) + fields["exact_at_chi"] = mbc <= log_chi + return fields + class TreePlan: """A rooted tree over ``n`` qubit leaves (any internal-node arity). @@ -53,6 +330,7 @@ def __init__(self, root, children, parent, qubit_of_leaf): 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._path_cache = {} # -- construction --------------------------------------------------------- @@ -96,6 +374,20 @@ def from_order(cls, order, *, weights=None, structure="quality", Maximum subsystem size for dense spectral reordering. """ order = list(order) + if not order: + 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))): + raise ValueError( + "order must be a permutation of qubit labels 0..n-1." + ) + if structure not in {"quality", "balanced", "adaptive"}: + raise ValueError( + "structure must be 'quality', 'balanced', or 'adaptive'." + ) if max_arity is not None: max_arity = int(max_arity) if max_arity < 2: @@ -140,7 +432,7 @@ def kary_split(qs): the previous ``mid = len(qs) // 2`` bisection exactly. """ length = len(qs) - k = 2 if max_arity is None else max_arity + k = length if max_arity is None else max_arity k = min(k, length) if k <= 1: return [qs] @@ -324,6 +616,116 @@ def from_children(cls, children, qubit_of_leaf, *, root=None): ) return cls(root, children, parent, qubit_of_leaf) + #: 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): + """Build a fixed-structure layered tree with a ternary top tensor. + + The structure is fixed; only ``block_size`` is tunable: + + * **First layer** (leaf-parent "blocking" nodes): each node groups + ``block_size`` consecutive qubits from ``order`` into one virtual + bond. This is the only choosable layer. + * **Middle layers**: strictly binary (two bonds in, one out). + * **Top tensor (root)**: always :attr:`LAYERED_ROOT_ARITY` (three) + children when there are at least three blocks; fewer only in the + degenerate small-``n`` case where three blocks do not exist. + + 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 + entanglement-adapted ordering, or + :meth:`TreeLayoutFinder.recommend_layered` to also search + ``block_size``. + block_size : int + Number of physical qubits per leaf-parent node. Default 4. + """ + order = list(order) + if not order: + raise ValueError("order must be non-empty.") + order = [int(q) for q in order] + n = len(order) + if sorted(order) != list(range(n)): + raise ValueError("order 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) + if block_size < 1: + raise ValueError("block_size must be >= 1.") + + counter = [0] + children_map = {} + qubit_of_leaf = {} + + def new_node(): + nid = counter[0] + counter[0] += 1 + return nid + + # Leaves: one node per qubit in the given order. + leaf_ids = [] + for q in order: + nid = new_node() + children_map[nid] = () + qubit_of_leaf[nid] = q + leaf_ids.append(nid) + + # First layer: group block_size leaves into one blocking node. + # A single-leaf chunk skips the parent and uses the leaf directly. + block_nodes = [] + for start in range(0, n, block_size): + chunk = leaf_ids[start: start + block_size] + if len(chunk) == 1: + block_nodes.append(chunk[0]) + else: + nid = new_node() + children_map[nid] = tuple(chunk) + block_nodes.append(nid) + + # Middle layers: binary tree over block_nodes. + def binary_subtree(nodes): + if len(nodes) == 1: + return nodes[0] + mid = len(nodes) // 2 + left = binary_subtree(nodes[:mid]) + right = binary_subtree(nodes[mid:]) + nid = new_node() + children_map[nid] = (left, right) + return nid + + # Top tensor: fixed ternary root (or fewer only when < 3 blocks exist). + num_blocks = len(block_nodes) + if num_blocks == 1: + # 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. + return cls.from_children( + children_map, qubit_of_leaf, root=block_nodes[0] + ) + root_arity = min(cls.LAYERED_ROOT_ARITY, num_blocks) + if num_blocks <= root_arity: + # Fewer blocks than the target arity: root takes them all directly. + root_nid = new_node() + children_map[root_nid] = tuple(block_nodes) + else: + # Split blocks into root_arity contiguous groups; each group is a + # binary middle subtree whose root becomes a direct child of the + # top tensor. + root_children = [] + for i in range(root_arity): + start = num_blocks * i // root_arity + end = num_blocks * (i + 1) // root_arity + root_children.append(binary_subtree(block_nodes[start:end])) + root_nid = new_node() + children_map[root_nid] = tuple(root_children) + + return cls.from_children(children_map, qubit_of_leaf, root=root_nid) + # -- queries -------------------------------------------------------------- def nodes(self): @@ -345,10 +747,47 @@ def is_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 max_bond_cut(self): + """Return the largest qubit bipartition induced by any tree bond. + + Every parent-child bond splits the qubits into the child's subtree + (``k`` qubits) and the rest (``n - k``). The Schmidt rank that bond can + carry is bounded by ``2 ** min(k, n - k)``, so this maximum + ``min(k, n - k)`` over all bonds is a purely structural, ``chi``-free + accuracy ceiling: the tree can represent an *arbitrary* state exactly + only when ``chi >= 2 ** max_bond_cut``. A structure whose + ``max_bond_cut`` exceeds ``log2(chi)`` must truncate at its widest bond + regardless of the gate stream. + """ + # One post-order pass to size every subtree, then reduce over bonds. + visit = [] + stack = [self.root] + while stack: + x = stack.pop() + visit.append(x) + stack.extend(self.children[x]) + size = {} + for x in reversed(visit): + ch = self.children[x] + size[x] = sum(size[c] for c in ch) if ch else 1 + best = 0 + for x, s in size.items(): + if x == self.root: + continue + best = max(best, min(s, self.n - s)) + return best + def node_path(self, a, b): """Return the node id path from node ``a`` to node ``b`` (inclusive).""" + if a not in self.children or b not in self.children: + raise ValueError(f"nodes {a!r} and {b!r} must belong to the tree") + cached = self._path_cache.get((a, b)) + if cached is not None: + return list(cached) if a == b: - return [a] + result = [a] + self._path_cache[(a, b)] = tuple(result) + return result ancestors = [] x = a while x is not None: @@ -363,7 +802,33 @@ def node_path(self, a, b): if x is None: raise ValueError("nodes are not in the same tree") lca = x - return ancestors[: depth[lca] + 1] + list(reversed(tail)) + result = ancestors[: depth[lca] + 1] + list(reversed(tail)) + self._path_cache[(a, b)] = tuple(result) + return result + + def subtree_qubit_masks(self): + """Return an integer bit mask of qubits below every node. + + Integer masks make repeated layout and preflight cut tests much cheaper + than rebuilding Python ``set`` objects for every edge. Python integers + remain exact for arbitrary qubit counts. + """ + visit = [] + stack = [self.root] + while stack: + node = stack.pop() + visit.append(node) + 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 + return masks def tree_distance(self, qa, qb): """Return the leaf-to-leaf path length between qubits ``qa`` and ``qb``.""" @@ -371,6 +836,45 @@ def tree_distance(self, qa, qb): lb = self.leaf_of_qubit[qb] return len(self.node_path(la, lb)) - 1 + def remove_leaf(self, q): + """Return a plan with qubit ``q`` capped and its unary parent removed. + + The remaining logical labels are compacted in the same way as a + one-dimensional MPS cap: labels above ``q`` shift down by one. The + surviving parent node is retained when possible, which keeps tensor + identities stable for callers holding live node references. + """ + if q not in self.leaf_of_qubit: + raise ValueError(f"qubit {q!r} is not present in the tree.") + if self.n <= 1: + raise ValueError("cannot remove the only qubit from a tree.") + leaf = self.leaf_of_qubit[q] + parent = self.parent.get(leaf) + if parent is None: + raise ValueError("cannot remove the root leaf from a multi-qubit tree.") + + children = {node: tuple(ch) for node, ch in self.children.items()} + qubit_of_leaf = dict(self.qubit_of_leaf) + children[parent] = tuple(c for c in children[parent] if c != leaf) + 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: + child = children[parent][0] + children[parent] = children[child] + del children[child] + if child in qubit_of_leaf: + qubit_of_leaf[parent] = qubit_of_leaf.pop(child) + + for node, old_q in tuple(qubit_of_leaf.items()): + if old_q > q: + qubit_of_leaf[node] = old_q - 1 + return type(self).from_children( + children, qubit_of_leaf, root=self.root + ) + def __repr__(self): n_internal = sum(1 for nid in self.nodes() if not self.is_leaf(nid)) return ( @@ -386,7 +890,9 @@ class TreeLayoutFinder: ---------- gates : bundled gate stream, optional ``[(gate, where), ...]`` entries. Two-qubit ``where`` supports define - the weighted interaction graph. Ignored when ``supports`` is given. + the weighted interaction graph. This finder does not accept a tensor + network state: pass that separately as ``state=`` to + :class:`TreeOptimizer`. Ignored when ``supports`` is given. n : int, optional Number of qubits. Inferred from the stream when omitted. supports : sequence of sequences, optional @@ -397,9 +903,17 @@ class TreeLayoutFinder: and ``"balanced"`` build strictly-binary trees when ``max_arity=2``; ``"adaptive"`` lets each level branch into its strongly coupled communities so the arity follows the gate connectivity. - max_arity : int or None - Maximum children per internal node (``2`` gives the binary tree; larger - values or ``None`` give flatter / wider trees). + 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. + 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. community_frac : float Strong-edge fraction for ``structure="adaptive"`` (see :meth:`TreePlan.from_order`). @@ -408,63 +922,1150 @@ class TreeLayoutFinder: (see :meth:`TreePlan.from_order`). dense_max : int Maximum subsystem size for dense spectral reordering. + objective : {"path", "congestion", "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``. + 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)``. + refine : {None, "greedy"} + 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. + refine_budget : int, optional + Maximum greedy swap proposals per candidate plan. Defaults to at most + 64 proposals when refinement is enabled. + search : {None, "nevergrad"} + Optional offline derivative-free refinement. It is never run unless + requested and requires the optional ``nevergrad`` package. + search_budget : int + Number of Nevergrad objective evaluations per candidate plan. + seed : int + Reproducible seed used by the optional Nevergrad stage. + nevergrad_optimizer : str + Nevergrad registry optimizer name, default ``"OnePlusOne"``. + weight_mode : {"count", "auto", "angle", "operator_schmidt"} + Event weighting used for the interaction graph. `"count"` is the + backward-compatible default. """ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", - max_arity=2, community_frac=0.35, star_frac=0.75, - dense_max=512): + 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"): + if ( + _looks_like_tree_tensor_network(gates) + or _looks_like_tree_tensor_network(supports) + ): + raise TypeError( + "TreeLayoutFinder accepts a circuit gate stream or supports, " + "not a TreeTensorNetwork. Build the layout from the circuit " + "and pass the TTN separately as TreeOptimizer(state=...)." + ) if supports is None: - supports = self._supports_from_gates(gates) + payloads, wheres, event_types = self._events_from_gates(gates) + supports = wheres + else: + supports = list(supports) + payloads = [None] * len(supports) + event_types = ["support"] * len(supports) supports = [tuple(_normalize_layout_support(s)) for s in supports] - self.supports = supports - + self.payloads = tuple(payloads) + self.event_types = tuple(event_types) inferred = -1 for support in supports: for site in support: - if isinstance(site, int): + if isinstance(site, Integral): inferred = max(inferred, site) if n is None: n = inferred + 1 + try: + n = int(n) + except (TypeError, ValueError) as exc: + raise ValueError("n must be a positive integer.") from exc if n <= 0: raise ValueError( "Could not infer qubit count; pass n explicitly." ) - self.n = int(n) + normalized_supports = [] + for support in supports: + if len(set(support)) != len(support): + raise ValueError( + f"layout support contains duplicate qubits: {support!r}." + ) + for site in support: + if not isinstance(site, Integral): + raise ValueError( + "tree layout supports must contain integer qubits; " + f"got {site!r}." + ) + if not 0 <= int(site) < n: + raise ValueError( + f"layout support qubit {site!r} is outside 0..{n - 1}." + ) + normalized_supports.append(tuple(int(site) for site in support)) + self.n = n + self.supports = tuple(normalized_supports) self.structure = structure - self.max_arity = None if max_arity is None else int(max_arity) + 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) self.dense_max = int(dense_max) + self.objective = _normalize_layout_objective(objective) + self.hybrid_weights = _normalize_hybrid_weights(hybrid_weights) + self.weight_mode = _normalize_weight_mode(weight_mode) + self.refine = _normalize_layout_refinement(refine) + if refine_budget is not None: + refine_budget = _validate_search_budget(refine_budget, "refine_budget") + self.refine_budget = refine_budget + self.search = _normalize_layout_search(search) + self.search_budget = _validate_search_budget(search_budget, "search_budget") + try: + self.seed = int(seed) + except (TypeError, ValueError) as exc: + raise ValueError("seed must be an integer.") from exc + self.nevergrad_optimizer = str(nevergrad_optimizer) + if not self.nevergrad_optimizer: + raise ValueError("nevergrad_optimizer must be a non-empty string.") + 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." + ) + self.max_operator_qubits = max_operator_qubits + + # Layout search asks for the same structural quantities several times + # (once per candidate arity and once per diagnostic). Keep these + # caches local to this immutable stream description; callers still get + # fresh dictionaries from the public diagnostic methods below. + self._plan_cache = {} + self._edge_load_cache = {} + self._schmidt_rank_cache = {} + self._similarity_cache = {} + self._congestion_weights_cache = None + self._balanced_plan_cache = None sites = list(range(self.n)) - self.pair_weights = _gate_stream_pair_weights(supports, sites) + self.event_weights = tuple( + _gate_stream_event_weights( + self.payloads, + self.supports, + self.event_types, + weight_mode=self.weight_mode, + ) + ) + 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) + ) + self.pair_weights = _gate_stream_pair_weights( + supports, sites, self.event_weights + ) @staticmethod - def _supports_from_gates(gates): + def _events_from_gates(gates): + """Return payloads, supports, and event types for layout analysis.""" if gates is None: - return [] - _, wheres, event_types = _normalize_layout_gate_queue(gates) - supports = [] - for where, etype in zip(wheres, event_types): - support = _normalize_layout_support(where) - if len(support) >= 2: - supports.append(support) - return supports + return [], [], [] - def run(self): - """Return a :class:`TreePlan` for the interaction graph.""" - order = list(range(self.n)) - return TreePlan.from_order( - order, - weights=self._similarity_weights(), - structure=self.structure, - max_arity=self.max_arity, + # ``_normalize_gate_entries`` intentionally accepts concrete bundled + # sequences, not one-shot iterators. A gate stream is commonly a + # generator, and the optimizer may pass the same stream to both its + # queue normalizer and this finder, so materialize it exactly once at + # the boundary. + if hasattr(gates, "__next__"): + gates = list(gates) + + # Control events carry support but are not gate entries. Strip them + # before delegating ordinary entries to the MPS layout normalizer, so + # a mixed stream can still build an interaction-aware tree. + control = _mps_control_event_parts(gates) + if control is not None: + return [None], [control[2]], [control[0]] + if isinstance(gates, (tuple, list)): + if not any( + _mps_control_event_parts(entry) is not None + for entry in gates + ): + return _normalize_layout_gate_queue(gates) + payloads = [] + supports = [] + event_types = [] + for entry in gates: + control = _mps_control_event_parts(entry) + if control is not None: + payloads.append(None) + supports.append(control[2]) + event_types.append(control[0]) + continue + one_payload, one_where, one_type = _normalize_layout_gate_queue( + (entry,) + ) + payloads.extend(one_payload) + supports.extend(one_where) + event_types.extend(one_type) + return payloads, supports, event_types + + payloads, supports, event_types = _normalize_layout_gate_queue(gates) + return payloads, supports, event_types + + @staticmethod + def _supports_from_gates(gates): + """Return only multi-site supports for compatibility with old callers.""" + _payloads, supports, _event_types = TreeLayoutFinder._events_from_gates(gates) + return [ + tuple(_normalize_layout_support(support)) + for support in supports + if len(_normalize_layout_support(support)) >= 2 + ] + + def _resolve_search_settings( + self, + *, + 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, + ): + """Resolve method overrides against finder-owned search defaults.""" + if refine is _DEFAULT_SEARCH_OPTION: + refine = self.refine + else: + refine = _normalize_layout_refinement(refine) + if refine_budget is _DEFAULT_SEARCH_OPTION: + refine_budget = self.refine_budget + elif refine_budget is not None: + refine_budget = _validate_search_budget( + refine_budget, "refine_budget" + ) + if refine is not None and refine_budget is None: + refine_budget = max(1, min(self.n - 1, 64)) + + if search is _DEFAULT_SEARCH_OPTION: + search = self.search + else: + search = _normalize_layout_search(search) + if search_budget is _DEFAULT_SEARCH_OPTION: + search_budget = self.search_budget + else: + search_budget = _validate_search_budget(search_budget, "search_budget") + if seed is _DEFAULT_SEARCH_OPTION: + seed = self.seed + else: + try: + seed = int(seed) + except (TypeError, ValueError) as exc: + raise ValueError("seed must be an integer.") from exc + if nevergrad_optimizer is _DEFAULT_SEARCH_OPTION: + nevergrad_optimizer = self.nevergrad_optimizer + else: + nevergrad_optimizer = str(nevergrad_optimizer) + if not nevergrad_optimizer: + raise ValueError("nevergrad_optimizer must be a non-empty string.") + return { + "refine": refine, + "refine_budget": refine_budget, + "search": search, + "search_budget": search_budget, + "seed": seed, + "nevergrad_optimizer": nevergrad_optimizer, + } + + @staticmethod + def _leaf_nodes(plan): + """Return the deterministic leaf-position order of a plan.""" + return tuple(sorted(plan.qubit_of_leaf)) + + def _leaf_order(self, plan): + """Return the qubit label assigned to each deterministic leaf position.""" + return tuple(plan.qubit_of_leaf[leaf] for leaf in self._leaf_nodes(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.") + 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 + ) + + def _plan_with_leaf_swap(self, plan, left_leaf, right_leaf): + """Swap two labels while retaining the tree topology exactly.""" + qubit_of_leaf = dict(plan.qubit_of_leaf) + qubit_of_leaf[left_leaf], qubit_of_leaf[right_leaf] = ( + qubit_of_leaf[right_leaf], + qubit_of_leaf[left_leaf], + ) + return TreePlan.from_children( + plan.children, qubit_of_leaf, root=plan.root + ) + + def _path_score_and_max(self, plan): + """Return the weighted interaction path sum and longest active path.""" + score = 0.0 + max_path = 0 + for (qa, qb), weight in self.pair_weights.items(): + distance = plan.tree_distance(qa, qb) + score += float(weight) * distance + max_path = max(max_path, distance) + return float(score), int(max_path) + + def _path_score_after_leaf_swap(self, plan, left_leaf, right_leaf, score): + """Return the exact path-score update for a two-label leaf swap.""" + qa = plan.qubit_of_leaf[left_leaf] + qb = plan.qubit_of_leaf[right_leaf] + change = 0.0 + for q in range(self.n): + if q == qa or q == qb: + continue + weight_a = self.pair_weights.get(tuple(sorted((qa, q))), 0.0) + weight_b = self.pair_weights.get(tuple(sorted((qb, q))), 0.0) + if weight_a: + change += float(weight_a) * ( + plan.tree_distance(qb, q) - plan.tree_distance(qa, q) + ) + if weight_b: + change += float(weight_b) * ( + plan.tree_distance(qa, q) - plan.tree_distance(qb, q) + ) + return float(score + change) + + @staticmethod + def _normalized_cost(value, reference): + """Normalize a non-negative layout metric against a fixed baseline.""" + if reference > 0.0: + return float(value / reference) + return 0.0 if value == 0.0 else float(value) + + def _hybrid_key(self, plan): + """Return normalized distance and rank-load cost for hybrid selection.""" + score, max_path = self._path_score_and_max(plan) + loads = self.edge_loads(plan) + max_load = max(loads.values(), default=0.0) + total_load = sum(loads.values()) + + balanced = self._balanced_plan() + balanced_score, _ = self._path_score_and_max(balanced) + balanced_loads = self.edge_loads(balanced) + balanced_max_load = max(balanced_loads.values(), default=0.0) + balanced_total_load = sum(balanced_loads.values()) + path_weight, max_load_weight, total_load_weight = self.hybrid_weights + hybrid = ( + path_weight * self._normalized_cost(score, balanced_score) + + max_load_weight * self._normalized_cost(max_load, balanced_max_load) + + total_load_weight * self._normalized_cost( + total_load, balanced_total_load + ) + ) + return ( + float(hybrid), + float(max_load), + float(total_load), + float(score), + int(max_path), + ) + + 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) + return self._hybrid_key(plan) + + def _selection_key(self, plan, chi): + """Return the objective key with the optional chi feasibility prefix.""" + key = self._objective_key(plan) + if chi is not None: + return (_chi_cut_fields(plan, chi)["chi_overflow"],) + key + return key + + def _selection_loss(self, plan, chi): + """Return a scalar surrogate for derivative-free layout search.""" + key = self._objective_key(plan) + if self.objective == "path": + value = key[0] + elif self.objective == "congestion": + value = key[0] + 1.0e-6 * key[1] + 1.0e-12 * key[2] + else: + value = key[0] + if chi is not None: + value += 1.0e6 * _chi_cut_fields(plan, chi)["chi_overflow"] + return float(value) + + def _discard_plan_cache(self, plan): + """Release diagnostics retained only for a rejected temporary 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)] + + def _refine_plan_greedy(self, plan, *, chi, budget): + """Greedily improve a fixed topology through adjacent leaf swaps.""" + initial_key = self._selection_key(plan, chi) + leaf_nodes = self._leaf_nodes(plan) + if len(leaf_nodes) < 2 or budget < 1: + return plan, { + "method": "greedy", + "evaluations": 0, + "accepted_moves": 0, + "initial_key": initial_key, + "final_key": initial_key, + } + + current = plan + current_key = initial_key + current_path_score = self.score(current) + evaluations = 0 + accepted_moves = 0 + position = 0 + while position < len(leaf_nodes) - 1 and evaluations < budget: + left_leaf = leaf_nodes[position] + right_leaf = leaf_nodes[position + 1] + evaluations += 1 + if self.objective == "path": + candidate_path_score = self._path_score_after_leaf_swap( + current, left_leaf, right_leaf, current_path_score + ) + if candidate_path_score >= current_path_score - 1.0e-12: + position += 1 + continue + candidate = self._plan_with_leaf_swap( + current, left_leaf, right_leaf + ) + candidate_key = self._selection_key(candidate, chi) + else: + candidate = self._plan_with_leaf_swap( + current, left_leaf, right_leaf + ) + candidate_key = self._selection_key(candidate, chi) + candidate_path_score = None + + if candidate_key < current_key: + self._discard_plan_cache(current) + current = candidate + current_key = candidate_key + if candidate_path_score is None: + current_path_score = self.score(current) + else: + current_path_score = candidate_path_score + accepted_moves += 1 + position = max(0, position - 1) + else: + self._discard_plan_cache(candidate) + position += 1 + return current, { + "method": "greedy", + "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 + ): + """Use Nevergrad to refine a leaf assignment before simulation starts.""" + try: + import nevergrad as ng + except ImportError as exc: + raise ImportError( + "Nevergrad tree-layout search requires the optional dependency. " + "Install it with `pip install pepsy[layout]`." + ) from exc + + try: + optimizer_class = ng.optimizers.registry[optimizer_name] + except KeyError as exc: + raise ValueError( + f"Unknown Nevergrad optimizer {optimizer_name!r}." + ) from exc + + 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 + parametrization = ng.p.Array(init=priorities) + if hasattr(parametrization, "set_bounds"): + parametrization.set_bounds(-float(self.n), float(2 * self.n)) + 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 = {} + + def loss(values): + values = np.asarray(values) + order = tuple( + int(q) for q in np.argsort(values, kind="stable") + ) + cached = losses.get(order) + if cached is not None: + return cached + candidate = self._plan_with_leaf_order(initial_plan, order) + value = self._selection_loss(candidate, chi) + losses[order] = value + self._discard_plan_cache(candidate) + return value + + recommendation = optimizer.minimize(loss) + final_order = tuple( + int(q) + for q 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) + if candidate_key < initial_key: + final_plan = candidate + improved = True + else: + self._discard_plan_cache(candidate) + final_plan = initial_plan + improved = False + return final_plan, { + "method": "nevergrad", + "optimizer": optimizer_name, + "budget": budget, + "evaluations": len(losses), + "seed": seed, + "initial_key": initial_key, + "final_key": self._selection_key(final_plan, chi), + "improved": improved, + } + + def _improve_plan(self, plan, *, chi, settings): + """Run the requested pre-simulation plan refinements in sequence.""" + initial_order = self._leaf_order(plan) + initial_key = self._selection_key(plan, chi) + info = { + "initial_order": initial_order, + "initial_key": initial_key, + "refinement": None, + "search": None, + } + if settings["refine"] == "greedy": + plan, info["refinement"] = self._refine_plan_greedy( + plan, chi=chi, budget=settings["refine_budget"] + ) + if settings["search"] == "nevergrad": + plan, info["search"] = self._refine_plan_nevergrad( + plan, + chi=chi, + budget=settings["search_budget"], + seed=settings["seed"], + optimizer_name=settings["nevergrad_optimizer"], + ) + info["final_order"] = self._leaf_order(plan) + info["final_key"] = self._selection_key(plan, chi) + return plan, info + + def _build_plan(self, weights, *, structure=None, + max_arity=_DEFAULT_MAX_ARITY): + """Build one deterministic candidate tree from pair weights.""" + if max_arity is _DEFAULT_MAX_ARITY: + max_arity = self.max_arity + structure = self.structure if structure is None else structure + key = (structure, max_arity, id(weights)) + # ``weights`` is normally one of the finder-owned cached mappings. The + # identity component avoids returning a plan built from a different + # caller-supplied weighting mapping with the same arity. + cached = self._plan_cache.get(key) + if cached is not None and cached[0] is weights: + return cached[1] + plan = TreePlan.from_order( + range(self.n), + weights=weights, + structure=structure, + max_arity=max_arity, community_frac=self.community_frac, star_frac=self.star_frac, dense_max=self.dense_max, ) + self._plan_cache[key] = (weights, plan) + return plan + + def _schmidt_rank(self, payload, support, left_support): + """Return a cached operator-Schmidt rank for layout diagnostics.""" + 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 + # 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 + # pairs should therefore reuse the small SVD. A structured MPO carries + # explicit site labels, so retain its label-sensitive cache key. + support = tuple(support) + left_set = frozenset(left_support) + is_structured_mpo = callable(getattr(payload, "gen_sites_present", None)) + if is_structured_mpo: + partition_key = left_set + support_key = support + else: + positions = {site: pos for pos, site in enumerate(support)} + partition_key = tuple( + positions[site] for site in support if site in left_set + ) + support_key = len(support) + key = (id(payload), support_key, partition_key) + cached = self._schmidt_rank_cache.get(key) + 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 + + def _candidate_plans(self, max_arity): + """Build the candidate plans considered by the selected objective.""" + interaction_plan = self._build_plan(self._similarity_weights()) + if max_arity != self.max_arity: + interaction_plan = self._build_plan( + self._similarity_weights(), max_arity=max_arity + ) + if self.objective == "path": + return {"interaction": interaction_plan} + + congestion_plan = self._build_plan( + self._similarity_weights(self._congestion_pair_weights()), + max_arity=max_arity, + ) + balanced_plan = self._build_plan( + self._similarity_weights(), structure="balanced", + max_arity=max_arity, + ) + return { + "interaction": interaction_plan, + "congestion": congestion_plan, + "balanced": balanced_plan, + } + + def _select_plan(self, max_arity): + """Select one plan without changing the finder's stored diagnostics.""" + candidates = self._candidate_plans(max_arity) + selected = min( + candidates, + key=lambda name: self._objective_key(candidates[name]), + ) + return candidates[selected] + + 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 + 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. + + Returns + ------- + list of int + A permutation of ``0..n-1``. + """ + weights = self._similarity_weights() + order = _gate_stream_spectral_order( + list(range(self.n)), weights, dense_max=self.dense_max + ) + return order if order else list(range(self.n)) + + def layered(self, block_size=4, *, order=None): + """Build a fixed layered tree for a chosen ``block_size`` (no search). + + This is the direct, single-``block_size`` counterpart to + :meth:`recommend_layered`. It orders the qubits with the spectral + :meth:`qubit_order` (so strongly coupled qubits share a block) and + returns the :class:`TreePlan` from :meth:`TreePlan.build_layered` + straight away -- no candidate sweep, no wrapper dict. + + ``block_size`` is a cost/accuracy knob, not something to maximize + blindly. A blocking node fuses ``block_size`` physical qubits into one + tensor, so every intra-block correlation is represented *exactly*, but + that tensor has ``2 ** block_size`` physical dimension. Larger blocks + are therefore more accurate at fixed ``chi`` yet exponentially more + expensive: pick the largest block that fits your memory budget rather + than searching, since a congestion search cannot see the block tensor's + exponential cost and simply trends toward the widest block. + + Parameters + ---------- + block_size : int + Number of physical qubits per leaf-parent (blocking) node. + order : sequence of int, optional + Qubit order fed to :meth:`TreePlan.build_layered`. Defaults to the + spectral :meth:`qubit_order`. + + Returns + ------- + TreePlan + The layered plan (blocking layer, binary middle, ternary top). + """ + if order is None: + order = self.qubit_order() + else: + order = [int(q) for q in order] + return TreePlan.build_layered(order, block_size=block_size) + + def recommend_layered( + self, + block_sizes=(2, 3, 4), + *, + order=None, + 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, + ): + """Optimize the fixed layered structure over ``block_size``. + + The structure family is fixed by :meth:`TreePlan.build_layered` + (a ``block_size`` blocking layer, binary middle layers, and a ternary + top tensor); only the blocking width is free. This builds one layered + plan per candidate ``block_size`` on the entanglement-adapted qubit + order and returns the plan that minimizes the selected layout + objective, mirroring :meth:`recommend_arities`. + + Parameters + ---------- + block_sizes : iterable of int + Candidate blocking widths (physical qubits per leaf-parent node). + order : sequence of int, optional + Qubit order fed to :meth:`TreePlan.build_layered`. Defaults to the + spectral :meth:`qubit_order` so strongly coupled qubits share a + block. + chi : int, optional + Bond-dimension budget. The path/congestion objectives are + ``chi``-blind cost proxies that can favour a wider block whose + widest bond overflows ``chi`` (see :meth:`TreePlan.max_bond_cut`). + When ``chi`` is given the recommendation is made ``chi``-aware: + candidates are ranked first by ``chi_overflow`` (how far the widest + bond exceeds ``log2(chi)``), so a structure that is *exact* at + ``chi`` is preferred, and the layout objective only breaks ties + among equally-overflowing candidates. Each candidate additionally + reports ``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 : {None, "greedy"}, optional + Override the finder refinement setting. `"greedy"` performs a + bounded adjacent leaf-swap search on each candidate tree. + refine_budget : int, optional + Maximum greedy proposals per candidate. When omitted, an enabled + greedy search uses at most ``min(n - 1, 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. + search_budget, seed, nevergrad_optimizer + Optional Nevergrad configuration for each candidate plan. + + Returns + ------- + dict + ``{"objective", "recommended_block_size", "order", "chi", "plan", + "candidates"}``. Each candidate carries its ``block_size``, the + :class:`TreePlan`, ``max_bond_cut``, and the same structural/cost + summary fields as :meth:`recommend_arities`. + """ + 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, + ) + options = [] + for bs in block_sizes: + key = int(bs) + if key < 1: + raise ValueError("block_sizes must be >= 1.") + if key not in options: + options.append(key) + if not options: + raise ValueError("block_sizes must contain at least one option.") + + if order is None: + order = self.qubit_order() + else: + order = [int(q) for q in order] + + candidates = [] + for bs in options: + plan = TreePlan.build_layered(order, block_size=bs) + plan, planning = self._improve_plan( + plan, chi=chi, settings=settings + ) + report = self.report( + plan, include_edge_loads=self.objective != "path" + ) + arity_histogram = {} + for node, children in plan.children.items(): + if not children: + continue + arity_histogram[len(children)] = ( + arity_histogram.get(len(children), 0) + 1 + ) + candidates.append({ + "block_size": bs, + "actual_max_arity": plan.max_arity(), + "root_arity": len(plan.children[plan.root]), + "arity_histogram": arity_histogram, + "score": report["score"], + "max_path": report["max_path"], + "max_edge_load": report["max_edge_load"], + "peak_bond_growth": report["peak_bond_growth"], + **_chi_cut_fields(plan, chi), + "order": self._leaf_order(plan), + "planning": planning, + "plan": plan, + }) + + def candidate_key(candidate): + return self._selection_key(candidate["plan"], chi) + ( + candidate["block_size"], + ) + + recommended = min(candidates, key=candidate_key) + return { + "objective": self.objective, + "recommended_block_size": recommended["block_size"], + "initial_order": tuple(order), + "order": recommended["order"], + "chi": chi, + "refine": settings["refine"], + "search": settings["search"], + "plan": recommended["plan"], + "candidates": candidates, + } + + def run(self): + """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 + :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. + """ + if self.arity_candidates is not None: + rec = self.recommend_arities(self.arity_candidates, chi=self.chi) + self._last_arity_recommendation = rec + self._selected_candidate = f"arity={rec['recommended_max_arity']}" + self._last_candidate_scores = { + f"arity={cand['max_arity']}": self._selection_key( + cand["plan"], rec["chi"] + ) + for cand in rec["candidates"] + } + 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] + for name, plan in candidates.items() + } + selected = min( + candidates, + key=lambda name: self._selection_key(candidates[name], self.chi), + ) + self._last_candidates = candidates + self._last_candidate_scores = { + name: self._selection_key(plan, self.chi) + for name, plan in candidates.items() + } + self._selected_candidate = selected + return candidates[selected] + + def recommend_arities( + self, + max_arities=(2, 3, 4), + *, + 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, + ): + """Compare binary and wider trees and return the best candidate. + + The returned mapping contains the recommended :class:`TreePlan` under + ``"plan"`` and candidate plans alongside structural/cost summaries + under ``"candidates"``. Wider arities shorten paths but increase local + tensor degree, so the recommendation uses the selected layout + objective and reports both effects. + + Parameters + ---------- + max_arities : iterable of int or None + Candidate maximum arities (``2`` = binary; ``None`` = unbounded). + chi : int, optional + Bond-dimension budget. When given, the recommendation is made + ``chi``-aware exactly as in :meth:`recommend_layered`: candidates + are ranked first by ``chi_overflow`` so a structure that stays + exact at ``chi`` (widest bond ``<= log2(chi)``) is preferred, and + the layout objective only breaks ties. Each candidate reports + ``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 + 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. + """ + 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, + ) + options = [] + for arity in max_arities: + if arity is None: + key = None + else: + key = int(arity) + if key < 2: + raise ValueError("max_arities must be >= 2 or None.") + if key not in options: + options.append(key) + if not options: + raise ValueError("max_arities must contain at least one option.") + + candidates = [] + for arity in options: + plan = self._select_plan(arity) + plan, planning = self._improve_plan( + plan, chi=chi, settings=settings + ) + report = self.report( + plan, include_edge_loads=self.objective != "path" + ) + arity_histogram = {} + for node, children in plan.children.items(): + if not children: + continue + arity_histogram[len(children)] = ( + arity_histogram.get(len(children), 0) + 1 + ) + candidates.append({ + "max_arity": arity, + "actual_max_arity": plan.max_arity(), + "is_binary": plan.is_binary(), + "arity_histogram": arity_histogram, + "max_virtual_degree": max( + ( + len(children) + (1 if node in plan.parent else 0) + for node, children in plan.children.items() + if children + ), + default=0, + ), + "score": report["score"], + "max_path": report["max_path"], + "max_edge_load": report["max_edge_load"], + "peak_bond_growth": report["peak_bond_growth"], + **_chi_cut_fields(plan, chi), + "order": self._leaf_order(plan), + "planning": planning, + "plan": plan, + }) + + def candidate_key(candidate): + return self._selection_key(candidate["plan"], chi) + ( + candidate["actual_max_arity"], + ) + + recommended = min(candidates, key=candidate_key) + return { + "objective": self.objective, + "recommended_max_arity": recommended["max_arity"], + "chi": chi, + "refine": settings["refine"], + "search": settings["search"], + "plan": recommended["plan"], + "candidates": candidates, + } + + def recommend_layout(self, max_arities=(2, 3, 4), **kwargs): + """Alias for :meth:`recommend_arities` with a layout-oriented name.""" + return self.recommend_arities(max_arities=max_arities, **kwargs) + + def _congestion_pair_weights(self): + """Return pair weights proportional to gate Schmidt load.""" + cached = getattr(self, "_congestion_weights_cache", None) + if cached is not None: + return cached + event_weights = [] + for payload, support, event_type in zip( + self.payloads, self.supports, self.event_types + ): + if len(support) < 2 or str(event_type).lower() in { + "measure", "reset", "measure_reset", "cap" + }: + event_weights.append(0.0) + continue + if payload is None: + event_weights.append(1.0) + continue + support = tuple(dict.fromkeys(support)) + logs = [] + if len(support) <= 8: + for mask in range(1, (1 << len(support)) - 1): + left = tuple( + site for i, site in enumerate(support) + if mask & (1 << i) + ) + rank = self._schmidt_rank(payload, support, left) + logs.append(float(np.log2(rank))) + else: + 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)) + self._congestion_weights_cache = _gate_stream_pair_weights( + self.supports, + range(self.n), + event_weights, + ) + return self._congestion_weights_cache + + def _subtree_qubits(self, plan): + """Return the qubits below each node of plan.""" + return { + node: frozenset( + q for q in range(self.n) if mask & (1 << q) + ) + for node, mask in plan.subtree_qubit_masks().items() + } + + def edge_loads(self, plan=None): + """Return predicted log-bond growth for every tree edge.""" + if plan is None: + plan = self.run() + cache_key = id(plan) + cached = self._edge_load_cache.get(cache_key) + if cached is not None and cached[0] is plan: + return dict(cached[1]) + below = plan.subtree_qubit_masks() + loads = { + (parent, child): 0.0 + for parent, children in plan.children.items() + for child in children + } + for payload, support, event_type in zip( + self.payloads, self.supports, self.event_types + ): + support = tuple(dict.fromkeys(support)) + if len(support) < 2 or str(event_type).lower() in { + "measure", "reset", "measure_reset", "cap" + }: + continue + support_mask = 0 + for site in support: + support_mask |= 1 << site + + # An edge crosses the support iff it belongs to the minimal + # subtree spanning the support leaves. 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: + # 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]) + 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)) + crossed_edges = [ + (parent, node) + for node in span_nodes + if (parent := plan.parent.get(node)) in span_nodes + ] + + 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) + ) + rank = ( + self._schmidt_rank(payload, support, left) + if payload is not None else 2 + ) + loads[edge] += float(np.log2(rank)) + # 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)) + return dict(loads) + + def _congestion_key(self, plan): + """Return the lexicographic key used by the load-aware objective.""" + loads = self.edge_loads(plan) + values = tuple(loads.values()) + return ( + max(values, default=0.0), + sum(values), + 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, + ), + ) - def _similarity_weights(self): + def _similarity_weights(self, pair_weights=None): """Return the qubit-pair similarity of Seitz et al. (Eq. 1). ``s(qi, qj) = |G(qi) & G(qj)| + 1 / (|G(qi)| + |G(qj)|)`` where ``G(q)`` @@ -476,17 +2077,25 @@ def _similarity_weights(self): subtrees when co-occurrence counts tie. Only the bisection uses this augmented similarity; :meth:`score` keeps the pure interaction weight. """ + if pair_weights is None: + pair_weights = self.pair_weights + cache_key = id(pair_weights) + cached = self._similarity_cache.get(cache_key) + if cached is not None and cached[0] is pair_weights: + return cached[1] + degree = {q: 0.0 for q in range(self.n)} for support in self.supports: for site in set(support): if isinstance(site, int) and 0 <= site < self.n: degree[site] += 1.0 - sim = dict(self.pair_weights) + sim = dict(pair_weights) for qi in range(self.n): for qj in range(qi + 1, self.n): deg = degree[qi] + degree[qj] if deg > 0.0: sim[(qi, qj)] = sim.get((qi, qj), 0.0) + 1.0 / deg + self._similarity_cache[cache_key] = (pair_weights, sim) return sim def score(self, plan): @@ -495,12 +2104,17 @@ def score(self, plan): Lower is better: this is the quantity the tree structure minimises (short leaf-to-leaf paths for strongly coupled qubits). """ - total = 0.0 - for (qa, qb), weight in self.pair_weights.items(): - total += float(weight) * plan.tree_distance(qa, qb) - return total + return self._path_score_and_max(plan)[0] + + 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" + ) + return self._balanced_plan_cache - def report(self, plan=None): + 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 @@ -522,11 +2136,46 @@ def report(self, plan=None): weighted_sum += float(weight) * d total_weight += float(weight) n_pairs = len(dists) - balanced = TreePlan.from_order(range(self.n), structure="balanced") + balanced = self._balanced_plan() balanced_score = self.score(balanced) + if include_edge_loads: + loads = self.edge_loads(plan) + balanced_loads = self.edge_loads(balanced) + else: + loads = None + balanced_loads = None + 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 = ( + max(balanced_loads.values(), default=0.0) + if balanced_loads is not None else None + ) + balanced_total_load = ( + sum(balanced_loads.values()) + if balanced_loads is not None else None + ) + hybrid_cost = None + if self.objective == "hybrid" and loads is not None: + hybrid_cost = self._hybrid_key(plan)[0] + arity_histogram = {} + for node, children in plan.children.items(): + if children: + arity_histogram[len(children)] = ( + arity_histogram.get(len(children), 0) + 1 + ) return { "n_qubits": self.n, "n_interacting_pairs": n_pairs, + "objective": self.objective, + "weight_mode": self.weight_mode, + "hybrid_weights": ( + self.hybrid_weights if self.objective == "hybrid" else None + ), + "hybrid_cost": hybrid_cost, + "root": plan.root, + "is_binary": plan.is_binary(), + "max_arity": plan.max_arity(), + "arity_histogram": arity_histogram, "score": float(weighted_sum), "max_path": int(max(dists)) if dists else 0, "mean_path": float(sum(dists) / n_pairs) if n_pairs else 0.0, @@ -537,4 +2186,31 @@ def report(self, plan=None): "score_ratio_vs_balanced": ( float(weighted_sum / balanced_score) if balanced_score else 0.0 ), + "edge_loads": loads, + "total_edge_load": float(total_load) if total_load is not None else None, + "max_edge_load": float(max_load) if max_load is not None else None, + "peak_bond_growth": ( + _safe_exp2(max_load) if max_load is not None else None + ), + "balanced_max_edge_load": ( + float(balanced_max_load) + if balanced_max_load is not None else None + ), + "balanced_total_edge_load": ( + float(balanced_total_load) + if balanced_total_load is not None else None + ), + "balanced_peak_bond_growth": ( + _safe_exp2(balanced_max_load) + if balanced_max_load is not None else None + ), + "peak_bond_growth_log2": ( + float(max_load) if max_load is not None else None + ), + "balanced_peak_bond_growth_log2": ( + float(balanced_max_load) + if balanced_max_load is not None else None + ), + "selected_candidate": getattr(self, "_selected_candidate", "interaction"), + "candidate_scores": getattr(self, "_last_candidate_scores", {}), } diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 18e9f1f..35276d7 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -37,13 +37,27 @@ from __future__ import annotations import contextlib +from copy import deepcopy +import heapq from numbers import Integral +import warnings +import autoray as ar import numpy as np import quimb.tensor as qtn +from ...backends import infer_backend_converter_from_sample, to_float from ...operators.gates import _normalize_gate_entries -from .layout import TreeLayoutFinder, TreePlan +from ..mps.optimizer import ( + _control_event_parts as _mps_control_event_parts, + normalize_submpo_where, + submpo_event_parts, +) +from .layout import ( + TreeLayoutFinder, + TreePlan, + _submpo_schmidt_rank_bound, +) from .ttn import TreeTensorNetwork __all__ = ["TreeOptimizer"] @@ -65,26 +79,253 @@ def _normalize_where(where): return tuple(int(site) for site in where) +def _submpo_to_dense(submpo, where): + """Materialize an explicit sub-MPO on its declared support. + + TTNs do not have a 1D sub-MPO absorption primitive. Their recursive + subtree-operator path is the equivalent operation, so stream markers are + lowered to the same dense support operator used by ``apply_gate``. The + support-size guard is checked by the caller before this conversion. + """ + to_dense = getattr(submpo, "to_dense", None) + if not callable(to_dense): + raise TypeError( + "TreeOptimizer submpo events require an MPO-like payload with " + "a to_dense() method." + ) + try: + dense = to_dense() + except Exception as exc: + raise ValueError("could not materialize the submpo event payload.") from exc + expected = 4 ** len(where) + size = int(np.prod(ar.shape(dense))) + if size != expected: + raise ValueError( + f"submpo payload has {size} entries, but support {where!r} " + f"requires {expected}." + ) + return ar.do("reshape", dense, (2 ** len(where), 2 ** len(where))) + + +def _same_tree_plan(left, right): + """Return whether two plans describe the same rooted tree geometry.""" + return ( + isinstance(left, TreePlan) + and isinstance(right, TreePlan) + and left.root == right.root + and left.children == right.children + and left.qubit_of_leaf == right.qubit_of_leaf + ) + + +def _is_product_tensor_network(state): + """Return whether every virtual bond of ``state`` has dimension one.""" + max_bond = getattr(state, "max_bond", None) + if not callable(max_bond): + return False + try: + return int(max_bond()) <= 1 + except (TypeError, ValueError): + return False + + +def _is_symmray_array(array): + """Return whether ``array`` is a Symmray block-sparse / fermionic array. + + Symmray gate payloads (e.g. native Fermi-Hubbard hopping/onsite gates) are + block-sparse arrays whose physical legs already carry the correct symmetry + charges and, for fermions, the anticommutation signs. They must never be + reshaped into base-2 sub-legs the way a dense qubit gate is; the raw + symmetric array is handed straight to the Symmray-aware Quimb split / + contract primitives. + """ + try: + return ar.infer_backend(array) == "symmray" + except Exception: # pragma: no cover - defensive backend inference + return False + + +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) + + +def _operator_schmidt_rank(op, where, left_where): + """Return the operator-Schmidt rank across a support bipartition. + + ``where`` fixes the operator's qubit ordering. The returned rank is the + exact rank of the operator viewed as a matrix from the input/output legs on + ``left_where`` to those on the complementary support. For a two-qubit + gate this is the ``k`` introduced by the gate SVD in the paper; for a + higher-order operator it is the corresponding conservative generalisation. + """ + where = tuple(where) + left_where = tuple(left_where) + left = {q: i for i, q in enumerate(where)} + left_set = set(left_where) + left_positions = [left[q] for q in left_where] + right_positions = [ + i for i, q in enumerate(where) if q not in left_set + ] + k = len(where) + try: + array = ar.to_numpy(op).reshape((2,) * (2 * k)) + except Exception as exc: + # Symmray arrays intentionally cannot be flattened through the generic + # autoray ``to_numpy`` path: their blocks carry charge sectors and, for + # fermions, graded signs. Preflight only needs a safe upper bound, so + # use the explicit operator leg dimensions without materialising the + # structured array. This is exact for the usual dense qubit shape and + # conservative for a native/sparse backend. + if not _is_symmray_array(op): + raise ValueError( + f"operator on support {where!r} must contain " + f"{4 ** k} entries." + ) from exc + try: + shape = tuple(int(dim) for dim in ar.shape(op)) + except (AttributeError, TypeError, ValueError) as shape_exc: + raise ValueError( + f"operator on support {where!r} must contain " + f"{4 ** k} entries." + ) from shape_exc + if len(shape) == 2 and shape == (2 ** k, 2 ** k): + shape = (2,) * (2 * k) + if len(shape) != 2 * k or any(dim < 1 for dim in shape): + raise ValueError( + f"operator on support {where!r} must have {2 * k} " + f"operator legs; got shape {shape}." + ) from exc + left_dim = int(np.prod([ + shape[i] * shape[k + i] for i in left_positions + ], dtype=int)) + right_dim = int(np.prod([ + shape[i] * shape[k + i] for i in right_positions + ], dtype=int)) + return max(1, min(left_dim, right_dim)) + + axes = ( + left_positions + + [k + i for i in left_positions] + + right_positions + + [k + i for i in right_positions] + ) + matrix = array.transpose(axes).reshape( + 4 ** len(left_positions), 4 ** len(right_positions) + ) + # A zero operator has Schmidt rank zero mathematically, but it cannot + # reduce a pre-existing bond. Treat it as rank one for this upper bound. + return max(1, int(np.linalg.matrix_rank(matrix))) + + +_PAULI_1Q = { + "X": np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex), + "Y": np.array([[0.0, -1.0j], [1.0j, 0.0]], dtype=complex), + "Z": np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex), +} +_RESET_FLIP_AXES = {"X": "Z", "Y": "X", "Z": "X"} +_DEFAULT_MAX_OPERATOR_QUBITS = 8 +_DEFAULT_MAX_SUBTREE_NODES = 128 + + +def _normalize_control_where(where): + """Return a non-empty tuple of integer qubit labels for a control.""" + if isinstance(where, Integral): + return (int(where),) + if ( + isinstance(where, (tuple, list)) + and where + and all(isinstance(site, Integral) for site in where) + ): + return tuple(int(site) for site in where) + raise ValueError( + "Tree control event where must be an int or non-empty sequence of ints." + ) + + +def _normalize_control_axes(axes, where, *, event): + """Return one X/Y/Z axis per control-event site.""" + axes = tuple(char for char in str(axes).upper() if not char.isspace()) + if not axes or any(axis not in _PAULI_1Q for axis in axes): + raise ValueError(f"{event} basis must use only X, Y, or Z axes.") + if len(axes) == 1 and len(where) > 1: + axes = axes * len(where) + if len(axes) != len(where): + raise ValueError( + f"{event} basis has {len(axes)} axis/axes but where {where!r} " + f"has {len(where)} site(s)." + ) + return axes + + +def _normalize_measure_axes(pauli, where): + """Return one X/Y/Z axis per measured site.""" + axes = tuple(char for char in str(pauli).upper() if not char.isspace()) + if not axes or any(axis not in _PAULI_1Q for axis in axes): + raise ValueError("measure pauli must use only X, Y, or Z axes.") + if len(axes) != len(where): + raise ValueError( + f"measure pauli {pauli!r} has {len(axes)} axis/axes but where " + f"{where!r} has {len(where)} site(s)." + ) + return axes + + class TreeOptimizer: """Replay a bundled gate stream on a rooted tree tensor network. Parameters ---------- gates : bundled gate stream, optional - ``[(gate, where), ...]`` entries with ``where`` an int (one-qubit) or a - pair of ints (two-qubit). Replayed eagerly on construction when given. + ``[(gate, where), ...]`` entries with ``where`` an int or a sequence + of distinct qubits. Replayed eagerly on construction when given. n : int, optional Number of qubits. Inferred from ``gates`` / ``tree`` when omitted. chi : int Maximum virtual bond dimension enforced during two-qubit threading. 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. structure : {"quality", "balanced", "adaptive"} Tree-structure strategy used when ``tree`` is not supplied. - max_arity : int or None - Maximum children per internal node for the auto-built structure. ``2`` - (default) gives a binary tree; larger values or ``None`` give flatter / - wider trees. Ignored when an explicit ``tree`` is supplied. + 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. + layout_objective : {"path", "congestion", "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; + ``"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 : 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 + arity/cost diagnostics. tree : TreePlan, optional Explicit tree structure (any arity). When omitted a :class:`TreeLayoutFinder` builds one from the gate stream. @@ -101,6 +342,34 @@ class TreeOptimizer: seed : int or None Seed for the internal random generator used by :meth:`measure` and :meth:`reset`. + track_truncation : bool + 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. + 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 + :meth:`estimate_bonds` predicts a larger intermediate bond. + max_operator_qubits : int, optional + Maximum support size allowed for dense multi-qubit operators. The + default is conservative protection against accidental ``4**k`` + allocations; pass ``None`` to disable it. Product-Pauli measurements + use a factorized path and do not consume this dense budget. + max_subtree_nodes : int, optional + 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. + 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; + a bond-dimension-one product TTN is instead remounted exactly on the + requested plan (with a warning). A bond-dimension-one Quimb MPS is also + accepted and mounted exactly. When omitted, the optimizer starts from + the product state ``|0...0>``. + state : TreeTensorNetwork or product MatrixProductState, optional + Explicit alias for ``tn``. All initial-state tensors must share one + backend, dtype, and device. User-provided gates/operators should use + that same backend; a mismatch is converted with a one-time warning. run : bool Whether to replay ``gates`` immediately (default ``True``). @@ -114,11 +383,107 @@ class TreeOptimizer: Node id of the current orthogonality centre (``None`` if unknown). """ + @staticmethod + def _normalize_mode(mode): + """Validate and normalize the two-site gate implementation mode.""" + mode = str(mode).strip().lower() + if mode not in {"auto", "direct", "mpo"}: + raise ValueError("mode must be 'auto', 'direct', or 'mpo'.") + return mode + def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, - structure="quality", max_arity=2, community_frac=0.35, - star_frac=0.75, tree=None, dtype=complex, threads=1, - seed=None, run=True): - self.G, self.where = self._normalize_gate_queue(gates) + 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, + dtype=complex, threads=1, seed=None, run=True, tn=None, + state=None, track_truncation=False, + max_intermediate_bond=None, + max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS, + max_subtree_nodes=_DEFAULT_MAX_SUBTREE_NODES, + record_history=True): + # 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 + # iterator and silently degrade to an interaction-free layout. + if hasattr(gates, "__next__"): + gates = list(gates) + if layout is not None: + if tree is not None: + raise ValueError("pass either layout= or tree=, not both.") + if isinstance(layout, TreeLayoutFinder): + if n is not None and int(n) != layout.n: + raise ValueError("n does not match the supplied layout.") + n = layout.n if n is None else n + tree = layout.run() + elif isinstance(layout, TreePlan): + tree = layout + else: + raise TypeError( + "layout must be a TreeLayoutFinder or TreePlan; " + "pass an entangled TreeTensorNetwork as state= or tn=." + ) + if state is not None: + if tn is not None: + raise ValueError("pass either state= or tn=, not both.") + tn = state + if isinstance(tree, TreeTensorNetwork): + raise TypeError( + "tree= expects a TreePlan, not a TreeTensorNetwork; " + "pass the entangled state as state= or tn=." + ) + self.G, self.where, self.event_types = self._normalize_gate_queue(gates) + self.layout_finder = layout if isinstance(layout, TreeLayoutFinder) else None + + product_state_source = None + if tn is not None: + if isinstance(tn, TreeTensorNetwork): + tn.validate() + source_n = tn.nqubits + if tree is not None and tree is not tn.plan: + if _same_tree_plan(tree, tn.plan): + tree = tn.plan + elif _is_product_tensor_network(tn): + product_state_source = tn + warnings.warn( + "Requested tree/layout differs from a product " + "TreeTensorNetwork; rebuilding the product state " + "exactly on the selected tree plan.", + UserWarning, + stacklevel=2, + ) + else: + raise ValueError( + "tree/layout geometry differs from the supplied " + "entangled TreeTensorNetwork. Refusing a potentially " + "lossy relayout; construct the state on the selected " + "TreePlan or explicitly convert it first." + ) + else: + tree = tn.plan + elif isinstance(tn, qtn.MatrixProductState): + if not _is_product_tensor_network(tn): + raise TypeError( + "TreeOptimizer accepts an MPS initial state only when " + "max_bond() == 1. Convert an entangled MPS explicitly " + "to a TreeTensorNetwork on the requested TreePlan." + ) + source_n = int(tn.nsites) + sites = tuple(tn.sites) + if tuple(sorted(sites)) != tuple(range(source_n)): + raise ValueError( + "Product MPS sites must be the compact labels 0..n-1 " + "to initialize TreeOptimizer." + ) + product_state_source = tn + else: + raise TypeError( + "tn/state must be a TreeTensorNetwork, or a product " + "quimb MatrixProductState with max_bond() == 1." + ) + if n is not None and int(n) != source_n: + raise ValueError("n does not match the supplied initial state.") + n = source_n if n is None: if tree is not None: @@ -131,37 +496,109 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, self.n = int(n) if self.n <= 0: raise ValueError("Could not infer qubit count; pass n explicitly.") + # 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. + self._logical_qubits = list(range(self.n)) + self._logical_positions = {q: q for q in self._logical_qubits} self.chi = int(chi) + if self.chi < 1: + raise ValueError("chi must be a positive integer.") self.cutoff = float(cutoff) + if self.cutoff < 0.0: + raise ValueError("cutoff must be non-negative.") + self.mode = self._normalize_mode(mode) + if two_site_mode is not None: + legacy_mode = self._normalize_mode(two_site_mode) + if self.mode != "auto" and self.mode != legacy_mode: + raise ValueError( + "pass either mode= or two_site_mode=, or give them " + "the same value." + ) + warnings.warn( + "two_site_mode= is deprecated; use mode= instead.", + DeprecationWarning, + stacklevel=2, + ) + self.mode = legacy_mode self.structure = structure - self.max_arity = None if max_arity is None else int(max_arity) + # ``max_arity`` may be a scalar (fixed tree), ``None`` (unbounded), or an + # 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.community_frac = float(community_frac) self.star_frac = float(star_frac) + self.layout_objective = str(layout_objective) + self.layout_weight_mode = str(layout_weight_mode) self.dtype = dtype 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.rng = np.random.default_rng(seed) + self.track_truncation = bool(track_truncation) + self.max_intermediate_bond = self._positive_limit( + max_intermediate_bond, "max_intermediate_bond" + ) + self.max_operator_qubits = self._positive_limit( + max_operator_qubits, "max_operator_qubits" + ) + self.max_subtree_nodes = self._positive_limit( + max_subtree_nodes, "max_subtree_nodes" + ) + self.record_history = bool(record_history) + self.measurements = [] + self.truncation_history = [] + self.update_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. + self.infidelities = [0.0] + self.infidelity_samples = [] + self.normalizations = [] + self.projection_diagnostics = [] + self._backend_conversion_warnings = set() + self._active_update = None + self._truncation_survival = 1.0 if tree is None: - tree = TreeLayoutFinder( - gates=gates, n=self.n, structure=structure, + self.layout_finder = TreeLayoutFinder( + 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, - ).run() + objective=self.layout_objective, + weight_mode=self.layout_weight_mode, + chi=self.chi, + max_operator_qubits=self.max_operator_qubits, + ) + tree = self.layout_finder.run() if not isinstance(tree, TreePlan): raise TypeError("tree must be a TreePlan or None.") self.plan = tree - self.tn = self._build_product_state() - # A freshly built product state has every virtual bond at dimension 1, - # so each tensor is trivially isometric and the state is already in - # canonical form with the root as orthogonality centre (norm 1). - # ``from_plan`` records that centre on the network; assert it here so the - # first gate skips an O(N) canonicalisation. - self.center = self.plan.root + if product_state_source is not None: + self.tn = self._remount_product_state(product_state_source) + self.center = self.plan.root + elif tn is None: + self.tn = self._build_product_state() + # A freshly built product state has every virtual bond at dimension + # 1, so it is already canonical at the root. + self.center = self.plan.root + else: + self._install_tn(tn) self._thread_ind = None if run and self.G: + if ( + self.max_intermediate_bond is not None + or self.max_operator_qubits is not None + or self.max_subtree_nodes is not None + ): + self.preflight( + max_bond=self.max_intermediate_bond, + max_operator_qubits=self.max_operator_qubits, + max_subtree_nodes=self.max_subtree_nodes, + ) self.run() # -- stream normalization ------------------------------------------------- @@ -169,17 +606,282 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, @staticmethod def _normalize_gate_queue(gates): if gates is None: - return [], [] + return [], [], [] + + # The low-level normalizer deliberately distinguishes a bundled + # sequence from a single gate. Materialize one-shot iterators here so + # generator-based gate streams behave like the documented bundled + # sequence and remain available for automatic layout discovery. + if hasattr(gates, "__next__"): + gates = list(gates) + + submpo_parts = submpo_event_parts(gates) + if submpo_parts is not None: + submpo, where = submpo_parts + return [submpo], [normalize_submpo_where(where)], ["submpo"] + + control_parts = _mps_control_event_parts(gates) + if control_parts is not None: + name, payload, where = control_parts + return [payload], [where], [name] + + if isinstance(gates, (tuple, list)) and any( + submpo_event_parts(entry) is not None + or _mps_control_event_parts(entry) is not None + for entry in gates + ): + payloads = [] + wheres = [] + event_types = [] + for entry in gates: + submpo_parts = submpo_event_parts(entry) + if submpo_parts is not None: + submpo, where = submpo_parts + payloads.append(submpo) + wheres.append(normalize_submpo_where(where)) + event_types.append("submpo") + continue + control_parts = _mps_control_event_parts(entry) + if control_parts is not None: + name, payload, where = control_parts + payloads.append(payload) + wheres.append(where) + event_types.append(name) + continue + gate_entries = _normalize_gate_entries( + (entry,), where=None, allow_empty=False + ) + gate, where = gate_entries[0] + payloads.append(gate) + wheres.append(_normalize_where(where)) + event_types.append("gate") + return payloads, wheres, event_types + entries = _normalize_gate_entries(gates, where=None, allow_empty=True) payloads = [g for g, _ in entries] wheres = [_normalize_where(w) for _, w in entries] - return payloads, wheres + return payloads, wheres, ["gate"] * len(payloads) + + def _layout_gate_stream(self): + """Return a layout stream remapped through preceding cap events. + + The live TTN compacts physical positions after a cap, while callers may + keep either compact or stable logical labels. Automatic layout is built + once before replay, so it must translate every later support back to + the original leaf labels before constructing the initial tree. + """ + active_labels = list(range(self.n)) + original_labels = list(range(self.n)) + stream = [] + for payload, where, event_type in zip( + self.G, self.where, self.event_types + ): + logical_where = _normalize_where(where) + try: + original_where = tuple( + original_labels[active_labels.index(q)] + for q in logical_where + ) + except ValueError as exc: + raise ValueError( + "cannot build a tree layout: event support references " + f"inactive labels {logical_where!r} after a cap." + ) from exc + + if event_type == "gate": + stream.append((payload, original_where)) + elif event_type == "submpo": + stream.append(self.submpo_event(payload, original_where)) + elif event_type == "measure": + stream.append({ + "kind": "measure", + "pauli": payload["pauli"], + "where": original_where, + "outcome": payload.get("outcome"), + }) + elif event_type == "reset": + stream.append({ + "kind": "reset", + "where": original_where, + "basis": "".join(payload["axes"]), + }) + elif event_type == "measure_reset": + stream.append({ + "kind": "measure_reset", + "where": original_where, + "basis": "".join(payload["axes"]), + "outcome": payload.get("outcomes"), + }) + elif event_type == "cap": + stream.append({ + "kind": "cap", + "where": original_where[0], + "vec": payload["vec"], + "absorb": payload.get("absorb", "left"), + "compact_labels": payload.get("compact_labels", True), + }) + else: # pragma: no cover - normalized streams are exhaustive + raise ValueError(f"unknown tree layout event {event_type!r}.") + + if event_type == "cap": + capped = logical_where[0] + position = active_labels.index(capped) + active_labels.pop(position) + original_labels.pop(position) + if payload.get("compact_labels", True): + active_labels = [ + label - 1 if label > capped else label + for label in active_labels + ] + return stream + + def _validate_event_stream_for_run(self): + """Validate stream labels, including optional stable-label caps.""" + active = list(self._logical_qubits) + for step, (payload, where, event_type) in enumerate( + zip(self.G, self.where, self.event_types), start=1 + ): + support = _normalize_where(where) + if event_type == "cap": + if len(support) != 1: + raise ValueError( + f"cap event at step {step} must reference one qubit." + ) + q = support[0] + if len(active) <= 1: + raise ValueError( + f"cap event at step {step} cannot remove the only qubit." + ) + if q not in active: + raise ValueError( + f"cap event at step {step} references qubit {q}, " + f"outside the current active labels {active!r}." + ) + active.remove(q) + if payload.get("compact_labels", True): + active = [label - 1 if label > q else label for label in active] + continue + out_of_range = [q for q in support if q not in active] + if out_of_range: + raise ValueError( + f"event at step {step} references qubit(s) outside the " + f"current active labels {active!r}: {out_of_range!r}." + ) # -- construction --------------------------------------------------------- + @staticmethod + def _positive_limit(value, name): + """Validate an optional positive integer resource limit.""" + if value is None: + return None + try: + value = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be a positive integer or None.") from exc + if value < 1: + raise ValueError(f"{name} must be a positive integer or None.") + return value + + @staticmethod + def _format_progress_infidelity(value): + """Format a norm-based progress infidelity compactly.""" + text = f"{float(value):#.0e}" + if "e" not in text: + return text + mantissa, exponent = text.split("e", 1) + sign = exponent[0] if exponent[:1] in "+-" else "" + digits = exponent[1:] if sign else exponent + digits = digits.lstrip("0") or "0" + return f"{mantissa}e{sign}{digits}" + def _phys(self, q): return self.tn.site_ind(q) + def _state_like(self): + """Return a live state array for backend-coercing operator payloads.""" + return self.tn.node_tensor(self.plan.root).data + + def _invalidate_state_norm_cache(self): + """Invalidate state-owned readout caches before direct tensor updates.""" + invalidate = getattr(self.tn, "_invalidate_norm_cache", None) + if callable(invalidate): + invalidate() + + @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} + + def backend_info(self): + """Return the common backend, dtype, and device of the live TTN.""" + return self._state_backend_info(self.tn) + + def _as_state_backend(self, array, *, warn=True): + """Return an operator payload compatible with the live TTN backend. + + Matching arrays pass through unchanged. A mismatched gate is converted + for compatibility, but emits one warning per source/target signature so + callers can keep data transfer and dtype promotion under their control. + """ + like = self._state_like() + state_info = self.backend_info() + target_signature = ( + state_info["backend"], state_info["dtype"], state_info["device"] + ) + 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 state_info["backend"] == "numpy": + return ar.to_numpy(array) + return ar.do("array", array, like=like) + + def _coerce_tensor_network_backend(self, tn, *, warn=True): + """Convert every tensor of an operator TN to the live state backend.""" + for tensor in tn.tensor_map.values(): + tensor.modify(data=self._as_state_backend(tensor.data, warn=warn)) + return tn + def _tag(self, nid): return self.tn.node_tag(nid) @@ -189,9 +891,28 @@ def _tid(self, nid): def _thread_ctx(self): """Context manager capping BLAS/OpenMP threads for the small tree ops.""" - if _THREAD_CONTROLLER is not None and self.threads is not None: - return _THREAD_CONTROLLER.limit(limits=self.threads) - return contextlib.nullcontext() + @contextlib.contextmanager + def managed(): + depth = getattr(self, "_thread_ctx_depth", 0) + if depth: + self._thread_ctx_depth = depth + 1 + try: + yield + finally: + self._thread_ctx_depth -= 1 + return + + self._thread_ctx_depth = 1 + try: + if _THREAD_CONTROLLER is not None and self.threads is not None: + with _THREAD_CONTROLLER.limit(limits=self.threads): + yield + else: + yield + finally: + self._thread_ctx_depth = 0 + + return managed() def _bond_name(self, u, v): lo, hi = (u, v) if u < v else (v, u) @@ -208,8 +929,209 @@ def _steiner_nodes(self, leaves): def _build_product_state(self): return TreeTensorNetwork.from_plan(self.plan, dtype=self.dtype) + @staticmethod + 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)) + physical_index = state.site_ind(q) + else: + site_tag = state.site_tag(q) + tids = tuple(state.tag_map[site_tag]) + if len(tids) != 1: + raise ValueError( + f"Product MPS site {q} must own exactly one tensor." + ) + tensor = state.tensor_map[tids[0]] + physical_index = state.site_ind(q) + try: + physical_axis = tensor.inds.index(physical_index) + except ValueError as exc: + raise ValueError( + f"Initial product state is missing physical index {physical_index!r}." + ) from exc + moved = ar.do("moveaxis", tensor.data, physical_axis, 0) + if int(np.prod(ar.shape(moved))) != 2: + raise ValueError( + "TreeOptimizer product-state remounting requires qubit " + "physical dimension two and unit virtual bonds." + ) + return ar.do("reshape", moved, (2,)) + + def _remount_product_state(self, state): + """Rebuild a product TTN exactly on ``self.plan`` using state vectors.""" + self._state_backend_info(state) + if not _is_product_tensor_network(state): + raise ValueError("only bond-dimension-one product states can be remounted.") + target = TreeTensorNetwork.from_plan(self.plan, dtype=complex) + 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] + vector = self._product_site_vector(state, q) + tensor.modify(data=ar.do("reshape", vector, ar.shape(tensor.data))) + else: + tensor.modify( + data=ar.do("ones", ar.shape(tensor.data), like=sample) + ) + + # In a product TTN, every internal tensor is a scalar because every + # virtual bond has dimension one. Preserve any distributed global + # normalization/phase by collecting those scalars on the new root. + if isinstance(state, TreeTensorNetwork): + factor = None + for node in state.plan.nodes(): + if state.plan.is_leaf(node): + continue + scalar = ar.do("reshape", state.node_tensor(node).data, ()) + factor = scalar if factor is None else factor * scalar + if factor is not None: + root_tensor = target.node_tensor(target.plan.root) + root_tensor.modify(data=root_tensor.data * factor) + + target._with_center(self.plan.root).validate() + self._state_backend_info(target) + return target + + def _install_tn(self, tn): + """Install an independent, canonical copy of a supplied tree state.""" + self._state_backend_info(tn) + tn.validate() + self.tn = tn.copy() + self.plan = self.tn.plan + self.tn.validate() + self.n = self.tn.nqubits + self._logical_qubits = list(range(self.n)) + self._logical_positions = {q: q for q in self._logical_qubits} + if ( + self.tn.canonical_region is None + or not self.tn.is_subtree_canonical_form() + ): + self.tn.canonize_around_node_(self.plan.root) + + def set_tn(self, tn): + """Replace the live tree state with a canonical independent copy.""" + if not isinstance(tn, TreeTensorNetwork): + raise TypeError("tn must be a TreeTensorNetwork.") + self._install_tn(tn) + self.truncation_history.clear() + self.update_history.clear() + self.infidelities[:] = [0.0] + self.infidelity_samples.clear() + self.normalizations.clear() + self.projection_diagnostics.clear() + self._truncation_survival = 1.0 + return self + + def set_p(self, tn): + """MPS-compatible alias for :meth:`set_tn`. + + The name lets a shared coefficient-backend frontend install either an + MPS or a tree coefficient state without branching on assignment APIs. + """ + return self.set_tn(tn) + + def _validate_qubit(self, q): + """Return the compact TTN position for a logical qubit label.""" + if not isinstance(q, Integral): + raise ValueError(f"qubit label must be an integer; got {q!r}.") + q = int(q) + try: + return self._logical_positions[q] + except KeyError: + raise ValueError(f"qubit {q} is outside the tree state.") + + def _validate_support(self, where, *, min_size=1, resolve=True): + """Validate a logical support and optionally return TTN positions.""" + where = tuple(int(q) for q in where) + if len(where) < min_size: + raise ValueError( + f"gate support must contain at least {min_size} qubit(s); " + f"got {where}." + ) + 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: + raise ValueError(f"tree position {q} is outside the state.") + return where + + def _require_dense_qubit_state(self, operation): + """Reject qubit-only helpers on native graded physical spaces.""" + native = bool(getattr(self.tn, "fermionic", False)) + if not native: + try: + native = any( + _is_symmray_array(tensor.data) + for tensor in self.tn.tensor_map.values() + ) + except AttributeError: + native = False + if native: + raise NotImplementedError( + f"TreeOptimizer.{operation} is a dense two-level qubit API; " + "native fermionic TTNs require a model-native Symmray " + "observable or projector via TreeTensorNetwork.local_expectation." + ) + + def _check_operator_limits(self, where, *, dense=True): + """Reject dense operators and oversized operator subtrees.""" + if ( + dense + and + self.max_operator_qubits is not None + and len(where) > self.max_operator_qubits + ): + raise MemoryError( + f"operator support of {len(where)} qubits exceeds the configured " + 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) + if len(span) > self.max_subtree_nodes: + raise MemoryError( + f"operator Steiner subtree has {len(span)} nodes, exceeding " + f"max_subtree_nodes={self.max_subtree_nodes}." + ) + + 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)) + bonds = {} + for node in span: + for neighbour in self._neighbors(node): + if neighbour not in span or node > neighbour: + continue + edge = (node, neighbour) + bonds[edge] = int( + self.tn.ind_size(self.tn.bond(node, neighbour)) + ) + return { + "support": tuple(self._logical_qubits[q] for q in where), + "span": tuple(sorted(span)), + "bonds": bonds, + "max_bond": max(bonds.values(), default=1), + } + # -- canonical centre tracking ------------------------------------------- + @property + def p(self): + """The current :class:`TreeTensorNetwork` state. + + A convenience alias for :attr:`tn` that mirrors the ``MpsOptimizer.p`` + interface, so code written against either engine can read the state via + ``engine.p``. + """ + return self.tn + + @p.setter + def p(self, value): + self._install_tn(value) + @property def center(self): """Node id of the current orthogonality centre (``None`` if unknown). @@ -240,11 +1162,38 @@ def _move_center(self, target): Delegates to :meth:`TreeTensorNetwork.shift_orthogonality_center`: a no-op when already centred, an incremental per-edge QR walk along the - tree geodesic from a known centre, or a single O(N) canonicalisation - when the centre is unknown. + tree geodesic from a known centre, or regional QR recovery followed by + 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) + def _nearest_anchor(self, nodes): + """Choose the closest node to the current centre or canonical region. + + Ties preserve the supplied order, keeping gate application + deterministic. If no canonical location is known, the first node is + used and the subsequent centre move establishes one. + """ + nodes = tuple(nodes) + if not nodes: + raise ValueError("at least one anchor node is required.") + center = self.center + region = self.canonical_region + if center is not None: + return min( + nodes, + key=lambda node: len(self.plan.node_path(center, node)), + ) + if region: + return min( + nodes, + key=lambda node: min( + len(self.plan.node_path(r, node)) for r in region + ), + ) + return nodes[0] + def shift_orthogonality_center(self, node): """Move the orthogonality centre to ``node`` along the tree geodesic. @@ -281,6 +1230,16 @@ def canonical_region(self): def canonical_region(self, value): self.tn.canonical_region = value + def layout_report(self): + """Return the stored layout diagnostics for the live tree, if available.""" + if self.layout_finder is None: + return { + "n_qubits": self.n, + "is_binary": self.plan.is_binary(), + "max_arity": self.plan.max_arity(), + } + return self.layout_finder.report(self.plan) + def canonize_subtree(self, nodes, *, span=False): """Canonicalise the state around the connected subtree ``nodes``. @@ -315,47 +1274,511 @@ def is_subtree_canonical_form(self, nodes=None, *, span=False, tol=1e-9): """ return self.tn.is_subtree_canonical_form(nodes, span=span, tol=tol) - # -- gate application ----------------------------------------------------- + def canonize_mps(self, p, where, *, info=None): + """Compatibility canonicalization entry point for shared frontends. - def run(self, gates=None): - """Replay ``gates`` (or the construction stream) on the tree. Returns self.""" - if gates is not None: - self.G, self.where = self._normalize_gate_queue(gates) - for gate, where in zip(self.G, self.where): - self.apply_gate(gate, where) - return self + The name is retained because ``MpsOptimizer`` exposes it publicly. For + a tree, an integer or singleton support moves the centre to that + qubit's leaf; a two-qubit support canonicalizes the minimal spanning + subtree. ``info['cur_orthog']`` is updated when an info mapping is + supplied, matching the MPS metadata convention. + """ + if not isinstance(p, TreeTensorNetwork): + raise TypeError("p must be a TreeTensorNetwork for tree canonicalization.") + if isinstance(where, Integral): + sites = (int(where),) + else: + try: + sites = tuple(int(q) for q in where) + except (TypeError, ValueError) as exc: + raise ValueError( + "where must be an int or a one-/two-qubit support." + ) from exc + 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: + 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]]) + target = (sites[0], sites[0]) + else: + p.canonize_around_qubits_(sites) + target = (min(sites), max(sites)) + if info is not None: + info["cur_orthog"] = target + return target - def apply_gate(self, gate, where): - """Apply a gate at qubit support ``where``. + # -- gate application ----------------------------------------------------- - One- and two-qubit gates take the optimised leaf-absorb / geodesic - threading paths (:meth:`apply_1q`, :meth:`apply_2q`); gates on three or - more qubits are applied in one shot over their minimal spanning subtree - via :meth:`apply_subtree_operator`. - """ - where = _normalize_where(where) + def _begin_update(self, kind, where): + """Start aggregating edge truncations for one state update.""" + if self._active_update is not None: + return False + self._active_update = { + "kind": str(kind), + "support": tuple(int(q) for q in where), + "edge_start": len(self.truncation_history), + } + return True + + def _abort_update(self): + """Discard a partial aggregation after a failed state update.""" + if self._active_update is None: + return + start = self._active_update["edge_start"] + del self.truncation_history[start:] + self._active_update = None + + def _finish_update(self): + """Commit one gate-level truncation aggregation.""" + active = self._active_update + if active is None: + return + if not self.record_history: + self._active_update = None + return + start = active["edge_start"] + edge_events = self.truncation_history[start:] + tracked = [ + event for event in edge_events + 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) + absolute_loss = float( + sum(event["discarded_weight"] for event in tracked) + ) + max_edge_loss = float( + max(event["discarded_weight"] for event in tracked) + ) + 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) + 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) + else: + relative_loss = None + absolute_loss = None + max_edge_loss = None + max_edge_fraction = None + cumulative_loss = None + + self.update_history.append({ + "update": len(self.update_history), + "kind": active["kind"], + "support": active["support"], + "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), + "absolute_discarded_weight": absolute_loss, + "relative_discarded_weight": relative_loss, + "cumulative_relative_discarded_weight": cumulative_loss, + "max_edge_discarded_weight": max_edge_loss, + "max_edge_discarded_fraction": max_edge_fraction, + }) + if tracked: + local_infidelity = float(relative_loss) + self.infidelities.append(float(cumulative_loss)) + self.infidelity_samples.append({ + "step": len(self.infidelity_samples) + 1, + "where": active["support"], + "edge_count": len(edge_events), + "local_fidelity": float(1.0 - local_infidelity), + "local_infidelity": local_infidelity, + "infidelity": float(cumulative_loss), + "cumulative_infidelity": float(cumulative_loss), + "method": "tree_edge_spectrum", + }) + self._active_update = None + + def apply_gate(self, gate, where, *, renormalize=False): + """Apply a gate and aggregate its edge truncation diagnostics.""" + started = self._begin_update("gate", _normalize_where(where)) + try: + result = self._apply_gate_impl(gate, where, renormalize=renormalize) + except Exception: + if started: + self._abort_update() + raise + if started: + self._finish_update() + return result + + def _apply_gate_impl(self, gate, where, *, renormalize=False): + """Apply a gate without opening a nested diagnostic update.""" + logical_where = _normalize_where(where) + where = self._validate_support(logical_where) + self._check_operator_limits(where) with self._thread_ctx(): if len(where) == 1: - self.apply_1q(gate, where[0]) + self.apply_1q( + gate, logical_where[0], renormalize=renormalize + ) elif len(where) == 2: if where[0] == where[1]: raise ValueError( "A two-qubit gate needs two distinct qubits; " f"got where={where}." ) - self.apply_2q(gate, where[0], where[1]) + self.apply_2q(gate, logical_where[0], logical_where[1]) + if renormalize: + self.normalize() else: if len(set(where)) != len(where): raise ValueError( "A multi-qubit gate needs distinct qubits; " f"got where={where}." ) - self.apply_subtree_operator(gate, where) + self.apply_subtree_operator( + gate, logical_where, renormalize=renormalize + ) + return self + + def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, + normalize_every=False, normalize_final=False, + normalize_eps=1e-15, seed=None): + """Replay ``gates`` (or the construction stream) on the tree. + + Parameters + ---------- + 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 + selectors for shared coefficient frontends. + non_unitary : bool, default=False + Mark the stream as non-unitary when using automatic normalization. + normalize_every : bool, default=False + Normalize after every replay event. Requires ``non_unitary=True``. + normalize_final : bool, default=False + Normalize 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 + Reseed measurement/reset sampling before replay. + """ + 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.", + DeprecationWarning, + stacklevel=2, + ) + else: + self.mode = self._normalize_mode(requested_mode) + non_unitary = bool(non_unitary) + 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: + raise ValueError("normalize_final requires non_unitary=True.") + normalize_every = bool(normalize_every) + normalize_eps = float(normalize_eps) + if normalize_eps < 0.0: + raise ValueError("normalize_eps must be non-negative.") + if seed is not None: + 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._validate_event_stream_for_run() + pbar = None + if progbar: + from tqdm import tqdm # pylint: disable=import-outside-toplevel + + pbar = tqdm( + total=len(self.G), + desc="tree", + leave=True, + position=0, + ascii=True, + colour="green", + ) + + one_qubit_count = 0 + two_qubit_count = 0 + multi_qubit_count = 0 + control_count = 0 + progress_reference_norm = ( + self.norm() + if pbar is not None + else None + ) + + try: + for step, (payload, where, event_type) in enumerate(zip( + self.G, self.where, self.event_types + ), start=1): + support = _normalize_where(where) + if event_type == "gate": + if len(support) == 1: + one_qubit_count += 1 + elif len(support) == 2: + two_qubit_count += 1 + else: + 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() + multi_qubit_count += 1 + else: + control_count += 1 + started = self._begin_update(event_type, support) + try: + self._apply_control_event(event_type, payload, support) + except Exception: + if started: + self._abort_update() + raise + if started: + 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", + }) + + 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 multi_qubit_count: + postfix["kq"] = multi_qubit_count + if control_count: + postfix["ctrl"] = control_count + pbar.set_postfix(postfix) + pbar.update(1) + finally: + 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", + }) + return self + + 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) + return self + + def add_gates(self, gates): + """Append gates to the queued stream and return ``self``.""" + G_new, where_new, event_types_new = self._normalize_gate_queue(gates) + self.G.extend(G_new) + self.where.extend(where_new) + self.event_types.extend(event_types_new) + return self + + @staticmethod + def measure_event(pauli, where, outcome=None): + """Return a Pauli-measurement event in the MPS-compatible format.""" + where = _normalize_control_where(where) + _normalize_measure_axes(pauli, where) + if outcome is None: + return ("measure", str(pauli), where) + if not isinstance(outcome, Integral) or int(outcome) not in (-1, 1): + raise ValueError("measure event outcome must be +1 or -1.") + return ("measure", str(pauli), where, int(outcome)) + + @staticmethod + def cap_event(where, vec, absorb="left", *, compact_labels=True): + """Return an MPS-compatible physical-index cap event. + + A tree cap contracts a leaf into its parent; ``absorb`` is accepted for + stream compatibility and must be ``"left"`` or ``"right"``. Set + ``compact_labels=False`` to preserve the caller-facing logical label + gap; this form is represented as a mapping so the option survives + stream normalization. + """ + where = _normalize_control_where(where) + if len(where) != 1: + raise ValueError("cap event where must reference exactly one site.") + if absorb not in {"left", "right"}: + raise ValueError("cap absorb direction must be 'left' or 'right'.") + try: + vec = ar.do("reshape", vec, (-1,)) + except (TypeError, ValueError) as exc: + raise ValueError("cap event vector must be array-like.") from exc + if int(ar.shape(vec)[0]) != 2: + raise ValueError("cap event vector must have exactly two entries.") + if compact_labels: + return ("cap", where[0], vec, absorb) + return { + "kind": "cap", "where": where[0], "vec": vec, + "absorb": absorb, "compact_labels": False, + } + + @staticmethod + def reset_event(where, basis="Z"): + """Return a mid-circuit reset event in the MPS-compatible format.""" + where = _normalize_control_where(where) + axes = _normalize_control_axes(basis, where, event="reset") + if all(axis == "Z" for axis in axes): + return ("reset", where) + return ("reset", where, "".join(axes)) + + @staticmethod + def measure_reset_event(pauli, where, outcome=None): + """Return a measure-then-reset event in the MPS-compatible format.""" + where = _normalize_control_where(where) + axes = _normalize_control_axes(pauli, where, event="measure_reset") + if outcome is None: + return ("measure_reset", "".join(axes), where) + if isinstance(outcome, Integral): + outcome = int(outcome) + if outcome not in (-1, 1): + raise ValueError("measure_reset outcome must be +1 or -1.") + return ("measure_reset", "".join(axes), where, outcome) + outcomes = tuple(int(value) for value in outcome) + if len(outcomes) != len(where) or any(value not in (-1, 1) for value in outcomes): + raise ValueError( + "measure_reset outcomes must be +1/-1 and match where." + ) + return ("measure_reset", "".join(axes), where, outcomes) + + @staticmethod + def control_event_parts(entry): + """Return ``(name, payload, where)`` for an MPS-style control event.""" + return _mps_control_event_parts(entry) + + @staticmethod + def is_control_event(entry): + """Return whether ``entry`` is an MPS-style control event.""" + return _mps_control_event_parts(entry) is not None + + @staticmethod + def submpo_event(submpo, where): + """Return an explicit sub-MPO stream marker for TTN replay.""" + return ("submpo", submpo, normalize_submpo_where(where)) + + @staticmethod + def submpo_event_parts(entry): + """Return ``(submpo, where)`` for an explicit marker, else ``None``.""" + return submpo_event_parts(entry, normalize_where=True) + + @staticmethod + def is_submpo_event(entry): + """Return whether ``entry`` is an explicit sub-MPO stream marker.""" + return submpo_event_parts(entry) is not None - def apply_1q(self, gate, q): + def apply_1q(self, gate, q, *, renormalize=False): """Absorb a one-qubit gate into the leaf tensor of qubit ``q``.""" - gate = np.asarray(gate).reshape(2, 2) + self._invalidate_state_norm_cache() + with self._thread_ctx(): + return self._apply_1q_impl(gate, q, renormalize=renormalize) + + def _apply_1q_impl(self, gate, q, *, renormalize=False): + """Apply a one-qubit gate without opening another thread context.""" + q = self._validate_qubit(q) + gate = self._as_state_backend(gate) + d = int(self.tn.ind_size(self._phys(q))) + if _is_symmray_array(gate): + # 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). + unitary = False + else: + if tuple(ar.shape(gate)) != (d, d): + gate = ar.do("reshape", gate, (d, d)) + gate_np = ar.to_numpy(gate) + unitary = np.allclose( + gate_np.conj().T @ gate_np, np.eye(d, dtype=gate_np.dtype), + rtol=1e-10, atol=1e-12, + ) + if not unitary: + self._move_center(self.plan.leaf_of_qubit[q]) + region = self.tn.canonical_region 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. + self.tn.canonical_region = region + if not unitary: + self.center = self.plan.leaf_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``. @@ -367,67 +1790,263 @@ def apply_2q(self, gate, qa, qb): are present is a single canonical compression sweep run back along the path, so every bond truncation sees the complete gate. """ - plan = self.plan - la = plan.leaf_of_qubit[qa] - lb = plan.leaf_of_qubit[qb] - pa, pb = self._phys(qa), self._phys(qb) + self._invalidate_state_norm_cache() + with self._thread_ctx(): + return self._apply_2q_impl(gate, qa, qb) - # Fast path: sibling leaves share a parent, so the gate correlation is - # carried entirely by the parent blob -- a single two-site update - # (contract the three tensors, apply the gate, re-split by SVD) avoids - # the QR threading and double-bond fusion of the general geodesic route. - parent = plan.parent.get(la) - if parent is not None and plan.parent.get(lb) == parent: - self._apply_2q_siblings(gate, qa, qb, la, lb, parent) - return + @staticmethod + def _as_gate_tensor4(gate, da, db): + """Return a two-site gate as a rank-4 (out_a, out_b, in_a, in_b) tensor. + + Dense qubit gates arrive as a ``(da*db, da*db)`` matrix (or an already + rank-4 array) and are reshaped to the ``(da, db, da, db)`` convention. + A Symmray gate is already a rank-4 symmetric array whose legs carry the + charge/fermion-sign structure, so it is returned untouched -- reshaping + it into base-2 sub-legs would corrupt the block structure. + """ + if _is_symmray_array(gate): + if len(ar.shape(gate)) != 4: + raise ValueError( + "A Symmray two-site gate must be a rank-4 " + "(out_a, out_b, in_a, in_b) array; got shape " + f"{tuple(ar.shape(gate))}." + ) + return gate + if tuple(ar.shape(gate)) == (da, db, da, db): + return gate + return ar.do("reshape", gate, (da, db, da, db)) + + def _apply_2q_impl(self, gate, qa, qb, *, max_bond=None, cutoff=None): + """Apply a two-site gate using the selected direct or MPO route. + + ``mode='direct'`` splits the gate locally and ``mode='mpo'`` lets Quimb + 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. + """ + gate = self._as_state_backend(gate) + if self.mode in {"auto", "direct"}: + return self._apply_2q_path_thread_impl( + gate, qa, qb, max_bond=max_bond, cutoff=cutoff + ) + return self._apply_2q_mpo_impl( + gate, qa, qb, max_bond=max_bond, cutoff=cutoff + ) - # Centre on the source leaf so the exact threading pushes an isometric - # front toward leaf b (the orthogonality centre rides the virtual bond). - self._move_center(la) + def _apply_2q_mpo_impl( + self, gate, qa, qb, *, max_bond=None, cutoff=None, + ): + """Apply a two-site gate as a Quimb sub-MPO without nested updates.""" + + qa = self._validate_qubit(qa) + qb = self._validate_qubit(qb) + if qa == qb: + raise ValueError("A two-qubit gate needs two distinct qubits.") + self._check_operator_limits((qa, qb)) + + if _is_symmray_array(gate): + # ``TensorNetwork.ind_size`` currently reports the size of a + # selected Symmray block after a native one-site gate, rather than + # the complete graded physical dimension. The native rank-four + # gate's explicit axes are authoritative for ``from_dense``. + gate_shape = tuple(int(dim) for dim in ar.shape(gate)) + if len(gate_shape) != 4: + raise ValueError( + "A native two-site gate must have rank four; got shape " + f"{gate_shape}." + ) + da, db = gate_shape[:2] + else: + da = int(self.tn.ind_size(self._phys(qa))) + db = int(self.tn.ind_size(self._phys(qb))) + gate = self._as_gate_tensor4(gate, da, db) + + # ``MatrixProductOperator.from_dense`` is backend-generic: for a + # Symmray FermionicArray its block-aware SVD returns native fermionic + # MPO tensors, including virtual-sector data and graded index + # 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, + ) + factors = self._two_site_mpo_factors(submpo, qa, qb) + if factors is None: + raise TypeError( + "two-site gate could not be represented as a Quimb sub-MPO " + "with one local factor per requested site." + ) + return self._apply_2q_factors_impl( + *factors, qa, qb, max_bond=max_bond, cutoff=cutoff + ) - gate = np.asarray(gate).reshape(2, 2, 2, 2) # (out_a, out_b, in_a, in_b) - self._thread_ind = qtn.rand_uuid() - gt = qtn.Tensor(gate, inds=("_na", "_nb", pa, pb)) - left, right = gt.split( - left_inds=("_na", pa), + def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): + """Extract and normalize a two-site Quimb MPO into local factors. + + The returned factors have exactly the same interface as a direct + gate-SVD factorization: each has one output leg, one state physical + input leg, and their shared operator-Schmidt bond. Crucially, factors + are found by their site tags rather than MPO tensor order, so a + descending support ``(qb, qa)`` remains correct. + """ + if site_where is None: + site_where = (qa, qb) + else: + site_where = tuple(site_where) + if len(site_where) != 2: + raise ValueError("a two-site MPO support must contain two sites.") + site_a, site_b = site_where + gen_sites = getattr(submpo, "gen_sites_present", None) + site_tag = getattr(submpo, "site_tag", None) + upper_id = getattr(submpo, "upper_ind_id", None) + lower_id = getattr(submpo, "lower_ind_id", None) + tag_map = getattr(submpo, "tag_map", None) + tensor_map = getattr(submpo, "tensor_map", None) + if not all((gen_sites, callable(site_tag), upper_id, lower_id, + tag_map is not None, tensor_map is not None)): + return None + try: + if set(gen_sites()) != {site_a, site_b}: + return None + raw_factors = {} + for qubit, site in ((qa, site_a), (qb, site_b)): + tids = tuple(tag_map[site_tag(site)]) + if len(tids) != 1: + return None + factor = tensor_map[tids[0]].copy() + upper = upper_id.format(site) + lower = lower_id.format(site) + if upper not in factor.inds or lower not in factor.inds: + return None + raw_factors[qubit] = (factor, upper, lower) + bonds = tuple(qtn.bonds( + raw_factors[qa][0], raw_factors[qb][0] + )) + except (KeyError, TypeError, ValueError): + return None + if len(bonds) != 1: + return None + + thread_ind = qtn.rand_uuid() + factors = {} + outputs = {} + for qubit, (factor, upper, lower) in raw_factors.items(): + output = qtn.rand_uuid() + factor.reindex_({ + upper: output, + lower: self._phys(qubit), + bonds[0]: thread_ind, + }) + factors[qubit] = factor + outputs[qubit] = output + return factors, outputs, thread_ind + + def _apply_2q_path_thread_impl( + self, gate, qa, qb, *, max_bond=None, cutoff=None, + ): + """Apply a two-qubit gate without opening another thread context.""" + gate = self._as_state_backend(gate) + qa = self._validate_qubit(qa) + qb = self._validate_qubit(qb) + if qa == qb: + raise ValueError("A two-qubit gate needs two distinct qubits.") + self._check_operator_limits((qa, qb)) + 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=self._thread_ind, + bond_ind=thread_ind, + ) + return self._apply_2q_factors_impl( + {qa: left, qb: right}, + {qa: output_a, qb: output_b}, + thread_ind, + qa, + qb, + max_bond=max_bond, + cutoff=cutoff, ) - # Absorb the left factor into leaf a; its new physical index keeps name. - ta = self.tn.tensor_map[self._tid(la)] - merged_a = qtn.tensor_contract(ta, left).reindex_({"_na": pa}) - ta.modify(data=merged_a.data, inds=merged_a.inds) - - # Thread the virtual bond exactly along the tree geodesic to leaf b. - path = plan.node_path(la, lb) - for u, v in zip(path, path[1:]): - self._thread_hop(u, v) + def _apply_2q_factors_impl( + self, factors, outputs, thread_ind, qa, qb, *, max_bond=None, + cutoff=None, + ): + """Apply two local operator factors with one shared path-thread kernel. + + ``direct`` supplies factors from its gate SVD and ``mpo`` supplies the + two factors Quimb made for its sub-MPO. From here onward they have + identical handling: choose the closest source leaf, attach its factor, + QR-thread the shared operator bond, attach the other factor, then make + one canonical compression sweep. + """ + plan = self.plan + la = plan.leaf_of_qubit[qa] + lb = plan.leaf_of_qubit[qb] + parent = plan.parent.get(la) + if 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, + ) - # Absorb the right factor into leaf b, consuming the virtual bond. - tb = self.tn.tensor_map[self._tid(lb)] - merged_b = qtn.tensor_contract(right, tb).reindex_({"_nb": pb}) - tb.modify(data=merged_b.data, inds=merged_b.inds) - self.center = lb - self._thread_ind = None + source, destination = ( + (qa, qb) + 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) + self._thread_ind = thread_ind + try: + source_tensor = self.tn.tensor_map[self._tid(source_leaf)] + merged_source = qtn.tensor_contract( + source_tensor, factors[source] + ).reindex_({outputs[source]: self._phys(source)}) + source_tensor.modify( + data=merged_source.data, inds=merged_source.inds, + ) - # Canonical compression sweep back along the path, now with both gate - # factors present, truncating every touched bond to ``chi``. - self._compress_path(path) + path = plan.node_path(source_leaf, destination_leaf) + for u, v in zip(path, path[1:]): + self._thread_hop(u, v) - def _apply_2q_siblings(self, gate, qa, qb, la, lb, parent): - """Apply a two-qubit gate to sibling leaves sharing ``parent``. + destination_tensor = self.tn.tensor_map[self._tid(destination_leaf)] + 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._compress_path(path, max_bond=max_bond, cutoff=cutoff) + finally: + self._thread_ind = None + return self - The two leaves and their shared parent are contracted into one blob, the - gate is applied, and the blob is re-split by two truncating SVDs -- a - direct two-site update (cf. the MPS two-site gate) that needs no bond - threading because both leaves already meet at ``parent``. With the - centre first moved onto ``parent`` the surrounding tree is isometric, so - each SVD truncates against an isometric environment and both gate - factors are present before any truncation. + def _apply_2q_sibling_factors( + self, factors, outputs, qa, qb, la, lb, parent, *, max_bond=None, + cutoff=None, + ): + """Apply two local factors to sibling leaves through their parent. + + There is no intermediate tree edge on which to thread the shared + operator bond. Absorb both factors, contract the two leaves and parent + into one blob (which consumes that bond), then make the usual two + truncating splits. This is the sibling specialization of + :meth:`_apply_2q_factors_impl`, shared by direct and MPO factors. """ pa, pb = self._phys(qa), self._phys(qb) self._move_center(parent) @@ -437,29 +2056,31 @@ def _apply_2q_siblings(self, gate, qa, qb, la, lb, parent): e_la = self._bond_name(la, parent) e_lb = self._bond_name(lb, parent) - blob = qtn.tensor_contract(tla, tp, tlb) - g = qtn.Tensor( - np.asarray(gate).reshape(2, 2, 2, 2), - inds=(pa + "*", pb + "*", pa, pb), + merged_a = qtn.tensor_contract(tla, factors[qa]).reindex_( + {outputs[qa]: pa} ) - blob = qtn.tensor_contract(blob, g).reindex_( - {pa + "*": pa, pb + "*": pb} + merged_b = qtn.tensor_contract(tlb, factors[qb]).reindex_( + {outputs[qb]: pb} ) + blob = qtn.tensor_contract(merged_a, tp, merged_b) # Split off leaf a (isometric), then leaf b, leaving the centre at the # parent; both new bonds keep their canonical tree-edge names. - la_t, rem = blob.split( - left_inds=[pa], method="svd", max_bond=self.chi, - cutoff=self.cutoff, absorb="right", get="tensors", bond_ind=e_la, + la_t, rem = self._split_with_diagnostics( + 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, ) - lb_t, p_t = rem.split( - left_inds=[pb], method="svd", max_bond=self.chi, - cutoff=self.cutoff, absorb="right", get="tensors", bond_ind=e_lb, + 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, ) 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) self.center = parent + return self def _thread_hop(self, u, v): """Thread the virtual bond exactly from node ``u`` to adjacent ``v``. @@ -469,9 +2090,14 @@ def _thread_hop(self, u, v): upper-triangular ``R`` carries the virtual bond forward to ``v``, moving the orthogonality centre with it. QR is lossless and cheaper than the SVD used for the final truncating sweep, so a single gate grows a - crossed bond by at most the gate rank ``k <= 4`` above its pre-gate - size, and that growth is undone by :meth:`_compress_path`. + crossed bond by at most its operator-Schmidt rank (bounded by + ``min(d_a**2, d_b**2)`` for local dimensions ``d_a`` and ``d_b``) + 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) + tu = self.tn.tensor_map[self._tid(u)] tv = self.tn.tensor_map[self._tid(v)] edge = next(iter(qtn.bonds(tu, tv))) @@ -486,7 +2112,271 @@ def _thread_hop(self, u, v): tu.modify(data=keep.data, inds=keep.inds) tv.modify(data=merged_v.data, inds=merged_v.inds) - def _compress_path(self, path): + def _fermionic_thread_hop(self, u, v): + """QR-route the operator bond without leaving native graded arrays.""" + tu = self.tn.tensor_map[self._tid(u)] + tv = self.tn.tensor_map[self._tid(v)] + edge = self.tn.bond(u, v) + left_inds = [ + index + for index in tu.inds + if index not in (edge, self._thread_ind) + ] + right_inds = [ + index + for index in tu.inds + if index not in left_inds + ] + keep, carry = tu.split( + left_inds=left_inds, + right_inds=right_inds, + method="qr", + absorb="right", + cutoff=0.0, + get="tensors", + ) + merged_v = qtn.tensor_contract(carry, tv) + 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, + ) + + @staticmethod + def _tensor_ind_size(tensor, ind): + """Return the dimension of ``ind`` on ``tensor``.""" + return int(tensor.shape[tensor.inds.index(ind)]) + + @staticmethod + def _split_rank_bound(tensor, left_inds): + """Return the matrix-rank upper bound for a tensor split.""" + left_inds = set(left_inds) + left_dim = int(np.prod([ + TreeOptimizer._tensor_ind_size(tensor, ind) for ind in left_inds + ], dtype=int)) + right_dim = int(np.prod([ + TreeOptimizer._tensor_ind_size(tensor, ind) + for ind in tensor.inds if ind not in left_inds + ], dtype=int)) + return min(left_dim, right_dim) + + @staticmethod + def _spectrum_to_numpy(values): + """Flatten a dense or block-sparse singular spectrum by magnitude.""" + if hasattr(values, "blocks"): + parts = [ + np.abs(np.asarray(ar.to_numpy(block))).ravel() + for block in values.blocks.values() + ] + spectrum = ( + np.concatenate(parts) + if parts + else np.empty(0, dtype=float) + ) + else: + spectrum = np.abs(np.asarray(ar.to_numpy(values))).ravel() + # Keep a deterministic ordering for dense diagnostics and rank display. + return np.sort(spectrum.astype(float, copy=False))[::-1] + + @staticmethod + def _spectrum_payload(values, kept_values=None): + """Package full/kept spectra for exact native truncation diagnostics.""" + payload = {"values": TreeOptimizer._spectrum_to_numpy(values)} + if hasattr(values, "blocks") and kept_values is not None: + payload["kept_values"] = TreeOptimizer._spectrum_to_numpy( + kept_values + ) + return payload + + @staticmethod + def _probe_bond_spectrum(ta, tb, *, max_bond=None, cutoff=0.0): + """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, + ) + values = info.get("singular_values") + if values is None: + return {"values": np.empty(0, dtype=float)} + kept_values = None + if hasattr(values, "blocks") and ( + max_bond is not None or float(cutoff) != 0.0 + ): + kept_info = {} + qtn.tensor_compress_bond( + ta.copy(), tb.copy(), max_bond=max_bond, + cutoff=cutoff, cutoff_mode="rel", absorb=None, + info=kept_info, + ) + kept_values = kept_info.get("singular_values") + return TreeOptimizer._spectrum_payload(values, kept_values) + + @staticmethod + def _probe_split_spectrum( + tensor, left_inds, *, max_bond=None, cutoff=0.0, + ): + """Return full/kept singular spectra for a tensor split.""" + _, values, _ = tensor.split( + left_inds=left_inds, + method="svd", + cutoff=0.0, + max_bond=None, + cutoff_mode="rel", + absorb=None, + get="arrays", + ) + kept_values = None + if hasattr(values, "blocks") and ( + max_bond is not None or float(cutoff) != 0.0 + ): + _, kept_values, _ = tensor.split( + left_inds=left_inds, + method="svd", + cutoff=cutoff, + max_bond=max_bond, + cutoff_mode="rel", + absorb=None, + get="arrays", + ) + return TreeOptimizer._spectrum_payload(values, kept_values) + + def _record_truncation( + self, *, kind, edge, before_bond, after_bond, bond_ind, + full_spectrum=None, max_bond=None, cutoff=None, + ): + """Record one edge split/compression and optional discarded weight.""" + if not self.record_history: + return + discarded_weight = None + discarded_fraction = None + spectrum_norm_sq = None + spectrum_rank = None + if full_spectrum is not None: + if isinstance(full_spectrum, dict): + spectrum = np.asarray( + full_spectrum["values"], dtype=float, + ).ravel() + kept = full_spectrum.get("kept_values") + else: + spectrum = np.asarray(full_spectrum, dtype=float).ravel() + kept = None + spectrum_rank = int(spectrum.size) + spectrum_norm_sq = float(np.sum(spectrum * spectrum)) + if kept is not None: + kept = np.asarray(kept, dtype=float).ravel() + kept_norm_sq = float(np.sum(kept * kept)) + discarded_weight = max(0.0, spectrum_norm_sq - kept_norm_sq) + else: + discarded = spectrum[int(after_bond):] + discarded_weight = float(np.sum(discarded * discarded)) + if spectrum_norm_sq > 0.0: + discarded_fraction = float(discarded_weight / spectrum_norm_sq) + else: + discarded_fraction = 0.0 + + self.truncation_history.append({ + "kind": str(kind), + "edge": tuple(int(x) for x in edge), + "bond": bond_ind, + "before_bond": int(before_bond), + "after_bond": int(after_bond), + "truncated": bool(after_bond < before_bond), + "spectrum_rank": spectrum_rank, + "spectrum_norm_sq": spectrum_norm_sq, + "discarded_weight": discarded_weight, + "discarded_fraction": discarded_fraction, + # ``None`` is meaningful: native MPO routing uses an uncapped, + # 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), + }) + + def _split_with_diagnostics( + self, tensor, left_inds, *, edge, bond_ind, max_bond, cutoff, + ): + """Split ``tensor`` and record the resulting virtual edge.""" + before_bond = self._split_rank_bound(tensor, left_inds) + # 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 + full_spectrum = ( + self._probe_split_spectrum( + tensor, + left_inds, + max_bond=max_bond, + cutoff=cutoff, + ) + 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, + ) + after_bond = self._tensor_ind_size(left, bond_ind) + self._record_truncation( + kind="split", edge=edge, before_bond=before_bond, + after_bond=after_bond, bond_ind=bond_ind, + full_spectrum=full_spectrum, + max_bond=max_bond, + cutoff=cutoff, + ) + return left, right + + def _compress_edge_with_diagnostics( + self, u, v, *, max_bond=None, cutoff=None, + ): + """Compress one live tree edge and record its truncation diagnostics.""" + max_bond = self.chi if max_bond is None else int(max_bond) + 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 + and before_bond <= max_bond + ): + # No singular value can be removed in this case. A lossless QR + # still moves the centre across the edge, but avoids an SVD. + self.tn.canonize_edge_(u, v, absorb="right") + bond_after = self.tn.bond(u, v) + self._record_truncation( + kind="canonize", edge=(u, v), before_bond=before_bond, + after_bond=int(self.tn.ind_size(bond_after)), + bond_ind=bond_after, max_bond=max_bond, cutoff=cutoff, + ) + return + full_spectrum = None + if self.track_truncation: + ta = self.tn.tensor_map[self._tid(u)] + tb = self.tn.tensor_map[self._tid(v)] + full_spectrum = self._probe_bond_spectrum( + ta, + tb, + max_bond=max_bond, + cutoff=cutoff, + ) + + # 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" + ) + bond_after = self.tn.bond(u, v) + after_bond = int(self.tn.ind_size(bond_after)) + self._record_truncation( + kind="compress", edge=(u, v), before_bond=before_bond, + after_bond=after_bond, bond_ind=bond_after, + full_spectrum=full_spectrum, max_bond=max_bond, cutoff=cutoff, + ) + + def _compress_path(self, path, *, max_bond=None, cutoff=None): """Canonically compress every bond along ``path`` down to ``chi``. The orthogonality centre sits at ``path[-1]`` on entry; sweeping back to @@ -496,19 +2386,272 @@ def _compress_path(self, path): sweep of Seitz et al. (Fig. 6) applied along the gate geodesic. """ for v, u in zip(path[::-1], path[-2::-1]): - self.tn.compress_between( - self._tag(v), - self._tag(u), - max_bond=self.chi, - cutoff=self.cutoff, - absorb="right", + self._compress_edge_with_diagnostics( + v, u, max_bond=max_bond, cutoff=cutoff, ) self.center = path[0] + def _compress_subtree(self, snodes, hub, *, max_bond=None, cutoff=None): + """Canonically compress every edge of a connected updated subtree. + + 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. + """ + snodes = frozenset(snodes) + self._move_center(hub) + + def descend(node, parent): + children = sorted( + neighbor + for neighbor in self._neighbors(node) + if neighbor in snodes and neighbor != parent + ) + for child in children: + self._compress_edge_with_diagnostics( + node, child, max_bond=max_bond, cutoff=cutoff, + ) + descend(child, node) + self.tn.canonize_edge_(child, node, absorb="right") + + descend(hub, None) + self.center = hub + + @staticmethod + def _qr_route_message(tensor, left_inds, *, bond_ind): + """Split one subtree message losslessly while carrying operator legs.""" + return tensor.split( + left_inds=left_inds, + method="qr", + absorb="right", + get="tensors", + bond_ind=bond_ind, + ) + + def _route_subtree_messages( + self, local, state_inds, operator_inds, order, *, token, + ): + """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] + ] + 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] + # -- general multi-qubit / sub-MPO application ---------------------------- def apply_subtree_operator(self, op, where, *, max_bond=None, cutoff=None, renormalize=False): + """Apply a subtree operator and aggregate its edge truncations.""" + self._invalidate_state_norm_cache() + started = self._begin_update("subtree", _normalize_where(where)) + try: + result = self._apply_subtree_operator_impl( + op, where, max_bond=max_bond, cutoff=cutoff, + renormalize=renormalize, + ) + except Exception: + if started: + self._abort_update() + raise + if started: + self._finish_update() + return result + + 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. + """ + self._invalidate_state_norm_cache() + logical_where = _normalize_where(where) + where = self._validate_support(logical_where) + return self._apply_submpo_resolved( + submpo, where, logical_where=logical_where, + max_bond=max_bond, cutoff=cutoff, + ) + + 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.""" + where = tuple(where) + if logical_where is None: + logical_where = where + else: + logical_where = tuple(logical_where) + if len(logical_where) != len(where): + raise ValueError( + "logical and compact sub-MPO supports must have equal length." + ) + self._check_operator_limits(where, dense=False) + max_bond = self.chi if max_bond is None else int(max_bond) + cutoff = self.cutoff if cutoff is None else float(cutoff) + if max_bond < 1 or cutoff < 0.0: + raise ValueError("max_bond must be positive and cutoff non-negative.") + started = self._begin_update("submpo", where) + try: + 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]), + ) + if factors is not None: + self._apply_2q_factors_impl( + *factors, where[0], where[1], + max_bond=max_bond, cutoff=cutoff, + ) + applied = True + if applied is None: + applied = self._try_apply_native_submpo( + submpo, where, payload_where=logical_where, + max_bond=max_bond, cutoff=cutoff, + ) + if applied is None: + self._check_operator_limits(where) + self._apply_subtree_operator_impl( + _submpo_to_dense(submpo, logical_where), logical_where, + max_bond=max_bond, cutoff=cutoff, + ) + except Exception: + if started: + self._abort_update() + raise + if started: + self._finish_update() + return self + + def apply_pauli_rotation(self, theta, pauli, where, *, sign=1.0): + """Apply ``exp(-i theta * sign * P / 2)`` on a Pauli support. + + The operator is represented as a bond-two MPO on the true support + window, so this remains efficient when ``where`` is sparse or long. + This method is deliberately frame-neutral: callers such as a future + stabilizer wrapper may pass a tableau-conjugated Pauli here. + """ + self._require_dense_qubit_state("apply_pauli_rotation") + from ..stabilizer_tn.operators import ( + pauli_combo_submpo, + single_qubit_rotation_matrix, + ) + + logical_where = _normalize_where(where) + where = self._validate_support(logical_where) + axes = _normalize_measure_axes(pauli, logical_where) + sign = float(sign) + if sign not in (-1.0, 1.0): + raise ValueError("Pauli rotation sign must be +1 or -1.") + if len(where) == 1: + self.apply_gate( + self._as_state_backend( + single_qubit_rotation_matrix( + float(theta), axes[0], sign=sign, dtype=self.dtype + ), + warn=False, + ), + logical_where, + ) + return self + terms = dict(zip(where, axes)) + c = np.cos(float(theta) / 2.0) + coef = -1j * sign * np.sin(float(theta) / 2.0) + mpo, mpo_where = pauli_combo_submpo( + c, coef, terms, self.n, dtype=self.dtype + ) + self._coerce_tensor_network_backend(mpo, warn=False) + return self._apply_submpo_resolved(mpo, mpo_where) + + def apply_pauli_sum(self, weighted_terms, *, max_bond=None, cutoff=None): + """Apply a weighted sum of Pauli products as one tree MPO. + + ``weighted_terms`` contains ``(coefficient, mapping)`` pairs, where + each mapping is ``{qubit: 'X'|'Y'|'Z'}``. The exact MPO bond is bounded + by the number of branches and is compressed by the usual tree path + message sweep. + """ + self._require_dense_qubit_state("apply_pauli_sum") + from ..stabilizer_tn.operators import pauli_sum_submpo + + terms = tuple(weighted_terms) + if not terms: + raise ValueError("weighted_terms must contain at least one term.") + resolved_terms = [] + for weight, term in terms: + resolved_terms.append(( + weight, + { + self._validate_qubit(q): axis + for q, axis in term.items() + }, + )) + mpo, where = pauli_sum_submpo( + resolved_terms, self.n, dtype=self.dtype + ) + self._coerce_tensor_network_backend(mpo, warn=False) + return self._apply_submpo_resolved( + mpo, where, max_bond=max_bond, cutoff=cutoff + ) + + def expectation_pauli(self, pauli, where, *, sign=1.0): + """Return the normalized expectation of a product Pauli operator.""" + self._require_dense_qubit_state("expectation_pauli") + where = self._validate_support(_normalize_where(where)) + axes = _normalize_measure_axes(pauli, where) + sign = float(sign) + if sign not in (-1.0, 1.0): + raise ValueError("Pauli expectation sign must be +1 or -1.") + return sign * to_float( + self._product_pauli_expectation(axes, where), real=True + ) + + def project_pauli(self, pauli, where, outcome, *, sign=1.0, + renormalize=True, normalize=None, + return_diagnostics=False): + """Project onto a product-Pauli eigenvalue. + + By default the post-projection state is normalized. Set + ``renormalize=False`` (or the compatibility alias ``normalize=False``) + to retain the physical projection norm, which is useful for composing + Kraus branches and for survival-probability accounting. With + ``return_diagnostics=True`` a dictionary containing the norm change and + support/tree/bond snapshots is returned instead of ``self``. + """ + self._require_dense_qubit_state("project_pauli") + logical_where = _normalize_where(where) + where = self._validate_support(logical_where) + if not isinstance(outcome, Integral) or int(outcome) not in (-1, 1): + raise ValueError("Pauli projection outcome must be +1 or -1.") + sign = float(sign) + if sign not in (-1.0, 1.0): + raise ValueError("Pauli projection sign must be +1 or -1.") + if normalize is not None: + renormalize = bool(normalize) + axes = _normalize_measure_axes(pauli, logical_where) + diagnostics = self._apply_product_pauli_projector( + axes, where, int(outcome) * int(sign), + renormalize=bool(renormalize), + return_diagnostics=return_diagnostics, + logical_support=logical_where, + ) + return diagnostics if return_diagnostics else self + + def _apply_subtree_operator_impl(self, op, where, *, max_bond=None, + cutoff=None, renormalize=False): """Apply a general multi-qubit operator over its minimal subtree. The one-shot generalisation of :meth:`apply_2q` to an operator on @@ -516,15 +2659,17 @@ def apply_subtree_operator(self, op, where, *, max_bond=None, non-unitary / Kraus operator, or a whole Trotter block -- applied as a single object rather than decomposed into one- and two-qubit gates. - The operator is contracted onto the *minimal connected subtree* (Steiner - subtree) spanning the target leaves and the result is re-split back into - that subtree with truncating SVDs: the tree 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 orthogonality centre is first moved onto a target leaf so the - whole exterior is isometric and every re-split truncation sees the - complete operator against an isometric environment; the centre is left on - a node inside the subtree. + 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 + 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 + sub-MPO applied over a covering range and then compressed (cf. + ``quimb``'s ``MatrixProductState.gate_with_submpo``, which exists for + the 1D chain only). Once the complete operator has reached the hub, a + canonical subtree sweep truncates every affected edge once against 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 @@ -533,16 +2678,74 @@ def apply_subtree_operator(self, op, where, *, max_bond=None, (e.g. after a Kraus/projection operator). ``max_bond`` / ``cutoff`` default to the optimizer's ``chi`` / ``cutoff``. Returns ``self``. """ - where = _normalize_where(where) - k = len(where) + logical_where = _normalize_where(where) + k = len(logical_where) + if k < 1: + raise ValueError("apply_subtree_operator needs at least one qubit.") + where = self._validate_support(logical_where) + self._check_operator_limits(where) if len(set(where)) != k: raise ValueError( f"apply_subtree_operator needs distinct qubits; got {where}." ) max_bond = self.chi if max_bond is None else int(max_bond) + if max_bond < 1: + raise ValueError("max_bond must be a positive integer.") cutoff = self.cutoff if cutoff is None else float(cutoff) + if cutoff < 0.0: + raise ValueError("cutoff must be non-negative.") + + if _is_symmray_array(op): + op_shape = tuple(int(dim) for dim in ar.shape(op)) + if len(op_shape) != 2 * k: + raise ValueError( + "native subtree operators must have one output and one " + f"input leg per site; got shape {op_shape}." + ) + if tuple(op_shape[:k]) != tuple(op_shape[k:]): + raise ValueError( + "native subtree operator output and input dimensions " + f"must match; got shape {op_shape}." + ) + with self._thread_ctx(): + if k == 1: + self._apply_1q_impl(op, logical_where[0]) + elif k == 2: + # Keep the public logical labels here. The two-site kernel + # resolves them exactly once, which is essential after a + # stable-label cap. + self._apply_2q_impl( + op, logical_where[0], logical_where[1], + max_bond=max_bond, cutoff=cutoff, + ) + else: + # Native multi-site arrays are lowered to a native MPO; + # unlike the dense path this preserves charge sectors and + # graded signs. The generated MPO uses compact site tags, + # while the outer update retains the caller's labels. + submpo = qtn.MatrixProductOperator.from_dense( + self._as_state_backend(op), + dims=op_shape[:k], + sites=where, + L=self.n, + max_bond=None, + cutoff=0.0, + ) + self._apply_submpo_resolved( + submpo, + where, + logical_where=where, + max_bond=max_bond, + cutoff=cutoff, + ) + if renormalize: + self.normalize() + return self + phys = [self._phys(q) for q in where] - op_arr = np.asarray(op).reshape([2] * (2 * k)) + op_arr = ar.do( + "reshape", self._as_state_backend(op), [2] * (2 * k) + ) with self._thread_ctx(): if k == 1: @@ -559,99 +2762,474 @@ def apply_subtree_operator(self, op, where, *, max_bond=None, 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 - # the subtree; every re-split truncation then measures true state - # error (the isometric exterior cancels between bra and ket). - if self.center not in snodes: - self._move_center(leaves[0]) - - # Record, from the *live* tensors, the index names each subtree node - # owns and keeps after contraction -- its physical leg and its - # boundary bonds to exterior nodes (gate application renames bonds off - # the deterministic scheme, so these must be read, not reconstructed). - snode_set = set(snodes) - owned = self._subtree_owned_inds(snodes, snode_set) - - # Contract the whole subtree (copies) into one blob, apply the - # operator to the physical legs, then re-split back into the subtree. + # 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) + if self.center != anchor: + self._move_center(anchor) + + # Factor the operator into an exact tree-MPO and apply it by + # leaf-to-hub message passing. This is the tree generalisation of + # the paper's two-qubit SVD + thread + compress update. In + # particular, no state tensor in the Steiner subtree is contracted + # 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) - ket_ts = [ - self.tn.tensor_map[self._tid(nid)].copy() for nid in snodes - ] - blob = qtn.tensor_contract(*ket_ts) - gt = qtn.Tensor(op_arr, inds=[p + "*" for p in phys] + phys) - blob = qtn.tensor_contract(blob, gt).reindex_( - {p + "*": p for p in phys} - ) - - # Peel each subtree-leaf toward the hub with a truncating SVD; the - # peeled tensor is isometric (``absorb="right"``) and the norm rides - # into the shrinking blob, ending on the hub tensor. Each new bond - # is registered on its hub-side node so a later peel keeps it. - for u, v in order: - left_inds = [ix for ix in blob.inds if ix in owned[u]] - new_bond = self._bond_name(u, v) - tu, blob = blob.split( - left_inds=left_inds, method="svd", max_bond=max_bond, - cutoff=cutoff, absorb="right", get="tensors", - bond_ind=new_bond, - ) - node_t = self.tn.tensor_map[self._tid(u)] - node_t.modify(data=tu.data, inds=tu.inds) - owned[v].add(new_bond) - hub_t = self.tn.tensor_map[self._tid(hub)] - hub_t.modify(data=blob.data, inds=blob.inds) - self.center = hub + op_factors, op_bonds = self._decompose_tree_operator( + op_arr, where, snodes, order, hub, + ) + self._apply_factorized_subtree_operator_impl( + op_factors, op_bonds, where, snodes, order, hub, + max_bond=max_bond, cutoff=cutoff, + ) if renormalize: self.normalize() return self - def _subtree_owned_inds(self, snodes, snode_set): - """Return ``{node: {index names it keeps after contraction}}``. - - For each subtree node this is its physical index (if a leaf) plus the - actual boundary-bond indices it shares with nodes *outside* the subtree, - read from the live tensors (gate application renames bonds off the - deterministic scheme, so the real names must be observed). Bonds - internal to the subtree are omitted -- they are contracted into the blob. - """ - owned = {} - for u in snodes: - t_u = self.tn.tensor_map[self._tid(u)] - ixs = set() - q = self.plan.qubit_of_leaf.get(u) - if q is not None: - ixs.add(self._phys(q)) - for w in self._neighbors(u): - if w not in snode_set: - t_w = self.tn.tensor_map[self._tid(w)] - ixs.update(qtn.bonds(t_u, t_w)) - owned[u] = ixs - return owned + def _try_apply_native_submpo( + self, submpo, where, *, max_bond, cutoff, payload_where=None, + ): + """Apply an MPO payload without materialising ``to_dense()``. + + A MatrixProductOperator is itself a tensor network. Its open operator + bonds are QR-routed through the same leaf-to-hub sweep as the TTN state + bonds: each peeled state tensor retains all currently open MPO bonds, + and they contract when their partner reaches the same accumulated + subtree. Only after the operator is complete is the affected subtree + canonically SVD-compressed. This works for arbitrary MPO bond + dimensions and does not require the MPO chain to match the TTN geometry. + + ``None`` is returned when the payload does not expose the Quimb MPO + site interface; callers then use the legacy dense fallback. + """ + where = tuple(where) + if payload_where is None: + payload_where = where + else: + payload_where = tuple(payload_where) + if len(payload_where) != len(where): + raise ValueError( + "logical and compact sub-MPO supports must have equal length." + ) + payload_for_compact = dict(zip(where, payload_where)) + gen_sites = getattr(submpo, "gen_sites_present", None) + site_tag = getattr(submpo, "site_tag", None) + upper_id = getattr(submpo, "upper_ind_id", None) + lower_id = getattr(submpo, "lower_ind_id", None) + tag_map = getattr(submpo, "tag_map", None) + tensor_map = getattr(submpo, "tensor_map", None) + if not all((gen_sites, callable(site_tag), upper_id, lower_id, + tag_map is not None, tensor_map is not None)): + return None + try: + present = tuple(gen_sites()) + except Exception: + return None + 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)) + order, hub = self._peel_order(snodes) + local = {} + state_inds = {} + operator_inds = {} + 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) + if q is None: + local[nid] = state_t + operator_inds[nid] = set() + continue + try: + payload_q = payload_for_compact[q] + tids = tuple(tag_map[site_tag(payload_q)]) + if len(tids) != 1: + return None + op_t = tensor_map[tids[0]].copy() + op_t.modify(data=self._as_state_backend(op_t.data)) + upper = upper_id.format(payload_q) + lower = lower_id.format(payload_q) + if upper not in op_t.inds or lower not in op_t.inds: + return None + op_t.reindex_({upper: self._phys(q) + "*", + lower: self._phys(q)}) + operator_inds[nid] = set(op_t.inds) - { + self._phys(q) + "*", self._phys(q) + } + local[nid] = qtn.tensor_contract(state_t, op_t).reindex_( + {self._phys(q) + "*": self._phys(q)} + ) + except (KeyError, TypeError, ValueError): + return None + + self._route_subtree_messages( + local, state_inds, operator_inds, order, token=qtn.rand_uuid(), + ) + + if operator_inds[hub]: + raise ValueError( + "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) + self._compress_subtree( + snodes, hub, max_bond=max_bond, cutoff=cutoff, + ) + return True + + def _apply_factorized_subtree_operator_impl( + self, op_factors, op_bonds, where, snodes, order, hub, *, + max_bond, cutoff, + ): + """Apply a factorized tree-MPO by QR routing then one final sweep.""" + local = {} + state_inds = {} + operator_inds = {} + for nid in snodes: + 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) + 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 + # with the live state physical index. + op_t = self._expand_tree_operator_leaf( + op_t, + 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].reindex_({f"{self._phys(q)}*": self._phys(q)}) + operator_inds[nid] = set(local[nid].inds) - state_inds[nid] + + self._route_subtree_messages( + local, state_inds, operator_inds, order, token=qtn.rand_uuid(), + ) + if operator_inds[hub]: + raise ValueError( + "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._compress_subtree( + snodes, hub, max_bond=max_bond, cutoff=cutoff, + ) + + def _decompose_tree_operator(self, op_arr, where, snodes, order, hub): + """Factor a dense operator into tensors on a Steiner tree. + + The operator is first viewed as a tensor with one dimension-four leg + per target qubit, where that leg packs (output, input). Repeated + leaf-to-hub SVDs then produce the exact hierarchical/tree-MPO factors. + Unlike the old state contraction path, this decomposition never sees a + state tensor and is performed once on the operator payload itself. + + Returns + ------- + factors : dict[int, qtn.Tensor] + One operator factor for each Steiner node. + bonds : dict + (child, parent) -> operator bond plus + "physical" -> qubit -> packed physical index. + """ + where = tuple(where) + op_axes = [f"_ttn_op_phys_{qtn.rand_uuid()}_{q}" for q in where] + # op_arr has all output axes followed by all input axes. Interleave + # them before packing each pair into one dimension-four operator leg. + interleaved = ar.do( + "transpose", + op_arr, + [i for q in range(len(where)) for i in (q, len(where) + q)], + ) + 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 + } + owned = {nid: set() for nid in snodes} + for leaf, q in leaf_for_q.items(): + owned[leaf].add(op_axes[where.index(q)]) + + factors = {} + op_bonds = {"physical": dict(zip(where, op_axes))} + for u, v in order: + left_inds = [ix for ix in blob.inds if ix in owned[u]] + if not left_inds: + raise RuntimeError( + "tree-MPO decomposition encountered an empty child block " + f"at node {u} on edge {(u, v)}." + ) + bond = f"_ttn_op_bond_{qtn.rand_uuid()}_{u}_{v}" + left, blob = blob.split( + left_inds=left_inds, + method="svd", + max_bond=None, + cutoff=0.0, + absorb="right", + bond_ind=bond, + ) + factors[u] = left + op_bonds[(u, v)] = bond + owned[v].add(bond) + + factors[hub] = blob + return factors, op_bonds + + def _apply_product_pauli_projector_impl( + self, axes, where, snodes, order, hub, outcome, *, max_bond, cutoff, + ): + """Apply ``(I + outcome P) / 2`` with a dimension-two branch index.""" + target_axes = dict(zip(where, axes)) + local = {} + branch_index = {} + for nid in snodes: + state_t = self.tn.tensor_map[self._tid(nid)].copy() + q = self.plan.qubit_of_leaf.get(nid) + if q in target_axes: + p = self._phys(q) + branch = f"_ttn_pauli_branch_{qtn.rand_uuid()}" + operators = np.stack( + ( + np.eye(2, dtype=complex), + outcome * _PAULI_1Q[target_axes[q]], + ), + axis=-1, + ) + op_t = qtn.Tensor( + self._as_state_backend(operators, warn=False), + inds=(p + "*", p, branch), + ) + local[nid] = qtn.tensor_contract(state_t, op_t).reindex_( + {p + "*": p} + ) + branch_index[nid] = branch + else: + local[nid] = state_t + + # Carry one dimension-two branch through the tree. When multiple child + # messages meet, a three-leg copy tensor enforces the same branch while + # keeping the representation linear in the node degree (rather than + # forming a high-rank 2**degree copy tensor at a wide hub). + update_token = qtn.rand_uuid() + for u, v in order: + state_bond = self.tn.bond(u, v) + branch = branch_index[u] + left_inds = [ + ix for ix in local[u].inds + if ix not in {state_bond, branch} + ] + new_bond = f"_ttn_pauli_apply_{update_token}_{u}_{v}" + tu, message = self._split_with_diagnostics( + local[u], left_inds, edge=(u, v), bond_ind=new_bond, + max_bond=max_bond, cutoff=cutoff, + ) + local[u] = tu + + parent_branch = branch_index.get(v) + if parent_branch is None: + local[v] = qtn.tensor_contract(local[v], message) + branch_index[v] = branch + else: + branch_out = f"_ttn_pauli_branch_{qtn.rand_uuid()}" + copy_tensor = np.zeros((2, 2, 2), dtype=complex) + copy_tensor[0, 0, 0] = 1.0 + copy_tensor[1, 1, 1] = 1.0 + copy_tensor = qtn.Tensor( + self._as_state_backend(copy_tensor, warn=False), + inds=(parent_branch, branch, branch_out), + ) + local[v] = qtn.tensor_contract( + local[v], message, copy_tensor + ) + branch_index[v] = branch_out + + root_branch = branch_index[hub] + local[hub] = qtn.tensor_contract( + local[hub], + qtn.Tensor( + self._as_state_backend( + np.array([0.5, 0.5], dtype=complex), warn=False + ), + inds=(root_branch,), + ), + ) + 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.center = hub + + def _apply_product_pauli_projector( + self, axes, where, outcome, *, renormalize=True, + return_diagnostics=False, logical_support=None, probability=None, + ): + """Apply a product-Pauli parity projector without dense materialization.""" + # Callers of this private helper have already resolved logical labels + # to compact TTN positions. + where = self._validate_support(where, resolve=False) + self._check_operator_limits(where, dense=False) + if len(set(where)) != len(where): + raise ValueError( + f"product-Pauli measurement needs distinct qubits; got {where}." + ) + before_norm = self.norm() + before = self._projection_snapshot(where) + started = self._begin_update("measure", where) + try: + if len(where) == 1: + projector = 0.5 * ( + np.eye(2, dtype=complex) + + outcome * _PAULI_1Q[axes[0]] + ) + self.apply_subtree_operator( + self._as_state_backend(projector, warn=False), + tuple(self._logical_qubits[q] for q in where), + max_bond=self.chi, + cutoff=self.cutoff, renormalize=renormalize, + ) + else: + leaves = [self.plan.leaf_of_qubit[q] for q in where] + snodes = self._steiner_nodes(leaves) + anchor = self._nearest_anchor(leaves) + if self.center != anchor: + self._move_center(anchor) + order, hub = self._peel_order(snodes) + with self._thread_ctx(): + self._apply_product_pauli_projector_impl( + axes, where, snodes, order, hub, outcome, + max_bond=self.chi, cutoff=self.cutoff, + ) + if renormalize: + self.normalize() + except Exception: + if started: + self._abort_update() + raise + if started: + self._finish_update() + after_norm = self.norm() + after = self._projection_snapshot(where) + diagnostics = { + "support": tuple( + before["support"] if logical_support is None + else tuple(logical_support) + ), + "span_before": before["span"], + "span_after": after["span"], + "bonds_before": before["bonds"], + "bonds_after": after["bonds"], + "max_bond_before": before["max_bond"], + "max_bond_after": after["max_bond"], + "norm_before": float(before_norm), + "norm_after": float(after_norm), + "norm_ratio": float(after_norm / before_norm) + if before_norm > 0.0 else 0.0, + "renormalized": bool(renormalize), + "outcome": int(outcome), + } + if probability is not None: + diagnostics["probability"] = float(probability) + self.projection_diagnostics.append(diagnostics) + if started and self.record_history and self.update_history: + self.update_history[-1]["projection"] = deepcopy(diagnostics) + return diagnostics if return_diagnostics else None + + 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) + if self.center not in snodes: + self._move_center(leaves[0]) + + internal = set() + for nid in snodes: + for nb in self._neighbors(nid): + if nb in snodes: + internal.add(self.tn.bond(nid, nb)) + ket = qtn.TensorNetwork([ + self.tn.tensor_map[self._tid(nid)].copy() for nid in snodes + ]) + phys = [self._phys(q) for q in where] + with self._thread_ctx(): + internal_map = {ix: qtn.rand_uuid() for ix in internal} + bra_num = ket.H.reindex({ + **internal_map, + **{p: p + "*" for p in phys}, + }) + numerator = bra_num & ket + for axis, p in zip(axes, phys): + numerator = numerator & qtn.Tensor( + self._as_state_backend(_PAULI_1Q[axis], warn=False), + inds=(p + "*", p), + ) + num = numerator.contract(output_inds=[]) + den = (ket.H.reindex(internal_map) & ket).contract(output_inds=[]) + return num / den + + @staticmethod + def _expand_tree_operator_leaf(op_tensor, packed_ind, physical_ind): + """Unpack one dimension-four tree-MPO leg into output and input legs.""" + axis = op_tensor.inds.index(packed_ind) + data = ar.do("moveaxis", op_tensor.data, axis, 0) + data = ar.do("reshape", data, (2, 2) + tuple(data.shape[1:])) + rest_inds = [ind for ind in op_tensor.inds if ind != packed_ind] + return qtn.Tensor( + data, + inds=[f"{physical_ind}*", physical_ind] + rest_inds, + ) def _peel_order(self, snodes): - """Return ``(peels, hub)`` for re-splitting a contracted subtree. + """Return ``(peels, hub)`` for recursive leaf-to-hub application. ``peels`` is a list of ``(u, v)`` edges: repeatedly a current subtree-leaf ``u`` (with a single remaining subtree neighbour ``v``) is peeled off toward ``v`` until a single ``hub`` node remains -- the node - the orthogonality centre ends on. Deterministic (smallest id first). + the orthogonality centre ends on. Deterministic (smallest id first). """ remaining = set(snodes) adj = { - u: [w for w in self._neighbors(u) if w in remaining] + u: tuple(w for w in self._neighbors(u) if w in remaining) for u in remaining } + degree = { + u: sum(w in remaining for w in neighbours) + for u, neighbours in adj.items() + } + leaves = [u for u, degree_u in degree.items() if degree_u == 1] + heapq.heapify(leaves) peels = [] while len(remaining) > 1: - leaf = next( - u for u in sorted(remaining) - if sum(1 for w in adj[u] if w in remaining) == 1 - ) + while leaves and leaves[0] not in remaining: + heapq.heappop(leaves) + if not leaves: + raise ValueError("subtree peel order requires a connected tree") + leaf = heapq.heappop(leaves) v = next(w for w in adj[leaf] if w in remaining) peels.append((leaf, v)) remaining.discard(leaf) - hub = next(iter(remaining)) + degree[leaf] = 0 + degree[v] -= 1 + if degree[v] == 1: + heapq.heappush(leaves, v) + hub = min(remaining) return peels, hub # -- readout -------------------------------------------------------------- @@ -660,15 +3238,16 @@ def max_bond(self): """Return the largest virtual bond dimension in the tree.""" return self.tn.max_bond() - def show(self, *, bond_dims=True, node_ids=False): + 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``. + 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) + self.tn.show(bond_dims=bond_dims, node_ids=node_ids, color=color) def bond_report(self): """Return a summary of the current virtual bond dimensions. @@ -691,94 +3270,502 @@ def bond_report(self): "chi": self.chi, } - def to_dense(self): - """Return the dense statevector with index order ``k0, k1, ..., k(n-1)``.""" + def estimate_bonds(self, gates=None, *, max_operator_qubits=None): + """Estimate untruncated bond growth from the fixed tree and gate stream. + + This is the paper's conservative dry-run bound (Eq. 4), not a + contraction of the live state. Each gate contributes its + operator-Schmidt rank to every tree edge that separates its support; + the reported edge dimensions are the product of those ranks over the + stream. Single-qubit gates and gates wholly contained on one side of + an edge contribute one. Because cancellations and state-specific rank + deficiencies are intentionally ignored, the estimate is safe but can + overestimate the dimensions reached by an actual replay. + + Parameters + ---------- + gates : bundled gate stream, optional + Stream to estimate. When omitted, the optimizer's queued stream is + used. Control events are included in the event trace but do not + change the bound. + max_operator_qubits : int, optional + If supplied, skip dense rank calculation for larger operators and + mark those events with ``rank_skipped=True``. This lets + :meth:`preflight` reject oversized operators before allocating a + dense ``4**k`` array. + + Returns + ------- + dict + ``edge_bonds`` maps undirected tree edges to their estimated bond + dimensions; ``max_bond`` is their maximum (or ``1`` for a one-leaf + tree); ``requires_truncation`` compares that bound to ``chi``; and + ``events`` records the crossing edges and ranks contributed by each + stream item. + """ + max_operator_qubits = self._positive_limit( + max_operator_qubits, "max_operator_qubits" + ) + if gates is None: + payloads = self.G + wheres = self.where + event_types = self.event_types + else: + payloads, wheres, event_types = self._normalize_gate_queue(gates) + + def edge_key(a, b): + return (a, b) if a < b else (b, a) + + sim_plan = self.plan + active = list(self._logical_qubits) + edge_bonds = {} + + def plan_steiner(plan, support): + leaves = [plan.leaf_of_qubit[q] for q in support] + if len(leaves) == 1: + return {leaves[0]} + nodes = set() + anchor = leaves[0] + for leaf in leaves[1:]: + nodes.update(plan.node_path(anchor, leaf)) + return nodes + + def plan_edges(plan): + return { + edge_key(parent, child): child + for parent, children in plan.children.items() + for child in children + } + + events = [] + for index, (payload, where, event_type) in enumerate( + zip(payloads, wheres, event_types) + ): + logical_support = _normalize_where(where) + try: + support = tuple(active.index(q) for q in logical_support) + except ValueError as exc: + raise ValueError( + f"event at step {index + 1} references qubit(s) outside " + f"the current active labels {active!r}: {logical_support!r}." + ) from exc + if len(support) > 1 and len(set(support)) != len(support): + raise ValueError( + f"gate support must contain distinct qubits; got {support}." + ) + span_nodes = len(plan_steiner(sim_plan, support)) + subtree_masks = sim_plan.subtree_qubit_masks() + edge_sides = { + edge: subtree_masks[child] + for edge, child in plan_edges(sim_plan).items() + } + current_edges = set(edge_sides) + edge_bonds = { + edge: edge_bonds.get(edge, 1) for edge in current_edges + } + crossing = {} + rank_skipped = bool( + max_operator_qubits is not None + and event_type in {"gate", "submpo"} + and len(logical_support) > max_operator_qubits + ) + rank_payload = payload + if ( + event_type in {"gate", "submpo"} + and len(logical_support) > 1 + and not rank_skipped + ): + support_mask = 0 + for site in support: + support_mask |= 1 << site + for edge, side in edge_sides.items(): + left_mask = support_mask & side + if not left_mask or left_mask == support_mask: + continue + left_positions = tuple( + q for q in support if left_mask & (1 << q) + ) + left_where = tuple( + logical_support[support.index(q)] for q in left_positions + ) + if event_type == "submpo": + rank = _submpo_schmidt_rank_bound( + rank_payload, logical_support, left_where + ) + if rank is None: + rank = _operator_schmidt_rank( + _submpo_to_dense( + rank_payload, logical_support + ), + logical_support, + left_where, + ) + else: + rank = _operator_schmidt_rank( + rank_payload, logical_support, left_where + ) + crossing[edge] = rank + edge_bonds[edge] *= rank + events.append({ + "index": index, + "kind": event_type, + "support": support, + "span_nodes": span_nodes, + "rank_skipped": rank_skipped, + "crossing_edges": dict(crossing), + "edge_bonds": dict(edge_bonds), + }) + + if event_type == "cap": + if len(logical_support) != 1: + raise ValueError( + f"cap event at step {index + 1} must reference one site." + ) + if len(active) <= 1: + raise ValueError( + f"cap event at step {index + 1} cannot remove the only site." + ) + sim_plan = sim_plan.remove_leaf(support[0]) + capped = logical_support[0] + active.remove(capped) + if payload.get("compact_labels", True): + active = [label - 1 if label > capped else label + for label in active] + + max_bond = max(edge_bonds.values(), default=1) + return { + "edge_bonds": edge_bonds, + "max_bond": int(max_bond), + "chi": self.chi, + "requires_truncation": bool(max_bond > self.chi), + "events": events, + } + + def preflight(self, gates=None, *, max_bond=None, + max_operator_qubits=None, max_subtree_nodes=None, + raise_on_error=True): + """Check conservative resource limits before replaying a stream. + + The bond limit uses :meth:`estimate_bonds` and is deliberately an + upper bound: it ignores cancellations and state-specific rank loss. The + operator and subtree limits protect dense operator allocation and bound + the number of recursive Steiner-tree messages. + With ``raise_on_error=False`` the method returns a report containing the + violations instead of raising :class:`MemoryError`. + + The optimizer invokes this automatically before eager replay when any + of its ``max_*`` constructor limits are set. Direct calls to + :meth:`apply_gate` and :meth:`apply_subtree_operator` still enforce the + operator/subtree limits immediately. + """ + max_bond = self._positive_limit(max_bond, "max_bond") + max_operator_qubits = self._positive_limit( + max_operator_qubits, "max_operator_qubits" + ) + max_subtree_nodes = self._positive_limit( + max_subtree_nodes, "max_subtree_nodes" + ) + report = self.estimate_bonds( + gates, max_operator_qubits=max_operator_qubits + ) + violations = [] + if max_bond is not None and report["max_bond"] > max_bond: + violations.append( + f"estimated max bond {report['max_bond']} exceeds " + f"max_bond={max_bond}" + ) + for event in report["events"]: + if ( + max_operator_qubits is not None + and event["kind"] in {"gate", "submpo"} + and len(event["support"]) > max_operator_qubits + ): + violations.append( + f"event {event['index']} has {len(event['support'])} " + f"operator qubits, exceeding max_operator_qubits=" + f"{max_operator_qubits}" + ) + if ( + max_subtree_nodes is not None + and event["span_nodes"] > max_subtree_nodes + ): + violations.append( + f"event {event['index']} spans {event['span_nodes']} " + f"tree nodes, exceeding max_subtree_nodes=" + f"{max_subtree_nodes}" + ) + result = dict(report) + result["limits"] = { + "max_bond": max_bond, + "max_operator_qubits": max_operator_qubits, + "max_subtree_nodes": max_subtree_nodes, + } + result["violations"] = violations + result["ok"] = not violations + if violations and raise_on_error: + raise MemoryError( + "TreeOptimizer preflight failed: " + "; ".join(violations) + ) + return result + + def truncation_report(self): + """Return per-edge truncation diagnostics collected during replay. + + The ``events`` list contains both ordinary edge compressions and local + SVD splits from sibling or subtree updates. Each event reports the + bond dimension before and after the update. When + ``track_truncation=True`` was requested, it also reports the local + singular-spectrum norm, absolute discarded weight, and relative + discarded fraction. Spectrum-based fields are ``None`` otherwise so + callers can distinguish an untracked run from a zero-loss truncation. + """ + events = deepcopy(self.truncation_history) + tracked = [ + event for event in events + if event["discarded_weight"] is not None + ] + if tracked: + total_discarded = float( + sum(event["discarded_weight"] for event in tracked) + ) + max_discarded = float( + max(event["discarded_weight"] for event in tracked) + ) + max_fraction = float( + max(event["discarded_fraction"] for event in tracked) + ) + else: + total_discarded = None + max_discarded = None + max_fraction = None + return { + "track_truncation": self.track_truncation, + "n_events": len(events), + "n_truncated": sum(event["truncated"] for event in events), + "n_tracked": len(tracked), + "total_discarded_weight": total_discarded, + "max_discarded_weight": max_discarded, + "max_discarded_fraction": max_fraction, + "events": events, + "updates": deepcopy(self.update_history), + } + + def to_dense(self, logical_order=True): + """Return the dense statevector in logical qubit order. + + ``logical_order`` is accepted for MPS interface compatibility. Tree + storage has no separate physical permutation, so both values produce + the same ``k0, k1, ..., k(n-1)`` ordering. + """ + _ = logical_order with self._thread_ctx(): return self.tn.to_statevector(range(self.n)) + @property + def qubits(self): + """Return the active caller-facing qubit labels in logical order. + + With the default compact cap behavior this is ``range(self.n)``. A + stable-label cap can leave gaps here while the underlying TTN remains + compact. + """ + return list(self._logical_qubits) + + @property + def logical_order(self): + """Return active logical labels ordered by their compact TTN position.""" + return list(self._logical_qubits) + + def logical_site(self, position): + """Return the logical qubit at compact tree position ``position``.""" + position = int(position) + if not 0 <= position < self.n: + raise IndexError( + f"tree position {position} is outside the range [0, {self.n})." + ) + return self._logical_qubits[position] + + def position(self, site): + """Return the tree position of logical qubit ``site``.""" + return self._validate_qubit(site) + + def remap_sample(self, config): + """Return a sample in active logical order.""" + if isinstance(config, dict): + return dict(config) + config = np.asarray(config) + if config.ndim == 0 or config.shape[-1] != self.n: + raise ValueError( + "sample configuration must have tree size as its final " + f"dimension, got shape {config.shape}." + ) + return config.copy() + + def restore_qubit_order(self): + """Return the state; tree storage is always already in logical order.""" + return self.tn + def norm(self): """Return the state norm ``**0.5``. - When the orthogonality centre is known the norm is the norm of that - single centre tensor: every other tensor is isometric and telescopes to - the identity between bra and ket. This is the tree analogue of the - one-site canonical norm tracked by :class:`pepsy.MpsOptimizer`. With an - unknown centre it falls back to the full doubled-tree contraction. + Native fermionic trees use a graded one-tensor center contraction when + 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. """ + 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)] val = qtn.tensor_contract(t.H, t, output_inds=[]) - return float(np.sqrt(np.abs(val))) + return float(np.sqrt(abs(to_float(val, real=True)))) with self._thread_ctx(): val = (self.tn.H & self.tn).contract(output_inds=[]) - return float(np.sqrt(np.abs(val))) + return float(np.sqrt(abs(to_float(val, real=True)))) - def normalize(self): - """Normalise the state in place. Returns the previous norm.""" + def _leaf_canonical_norm(self): + """Return a cheap dense canonical norm or the exact fermionic norm. + + Dense/nonfermionic trees move the orthogonality centre onto a + first-layer leaf and read that tensor's Frobenius norm. Native + fermionic trees must retain the complete graded exterior, so they + dispatch to :meth:`norm` instead of using the one-tensor shortcut. + """ + 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)] + val = qtn.tensor_contract(t.H, t, output_inds=[]) + return float(np.sqrt(abs(to_float(val, real=True)))) + + def normalize(self, eps=1e-15, insert=None): + """Normalise the state in place and return the previous norm. + + ``eps`` and ``insert`` are accepted for compatibility with + :meth:`MpsOptimizer.normalize`; a tree has no chain insertion site, so + ``insert`` is intentionally ignored. + """ + eps = float(eps) + if eps < 0.0: + raise ValueError("eps must be non-negative.") nrm = self.norm() - if nrm > 0: - tid = self._tid(self.center) if self.center is not None else self._tid( - self.plan.root - ) + if nrm > eps: + 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) t = self.tn.tensor_map[tid] t.modify(data=t.data / nrm) return nrm - def local_expectation(self, op, where): - """Return `` / `` for an operator on ``where``. - - For a single-site operator this moves the tracked orthogonality centre - onto the target leaf and contracts only that centre tensor with ``op``: - every other tensor is isometric and cancels between bra and ket, the - tree analogue of ``MpsOptimizer.local_expectation_canonical``. A - multi-site operator is contracted over only the **minimal subtree - spanning the target leaves**: with the centre inside that subtree every - tensor outside it is isometric and its bra/ket copies cancel to the - identity on the shared boundary bond, so the cost scales with the - operator's spread rather than the whole tree. + def _measure_pauli(self, pauli, where, outcome=None, *, renormalize=True, + return_diagnostics=False): + """Measure a Pauli product, collapse, and return ``(outcome, prob)``.""" + self._require_dense_qubit_state("measure_pauli") + logical_where = _normalize_control_where(where) + where = self._validate_support(logical_where) + axes = _normalize_measure_axes(pauli, logical_where) + expectation = to_float( + self._product_pauli_expectation(axes, where), real=True + ) + p_plus = min(max(0.5 * (1.0 + expectation), 0.0), 1.0) + if outcome is None: + outcome = 1 if self.rng.random() < p_plus else -1 + elif not isinstance(outcome, Integral) or int(outcome) not in (-1, 1): + raise ValueError("measure event outcome must be +1 or -1.") + else: + outcome = int(outcome) + probability = p_plus if outcome > 0 else 1.0 - p_plus + if probability <= 1e-12: + raise ValueError( + f"forced measure outcome {outcome} has ~0 probability " + f"({probability:.2e})." + ) + diagnostics = self._apply_product_pauli_projector( + axes, where, outcome, renormalize=renormalize, + return_diagnostics=return_diagnostics, + logical_support=logical_where, probability=probability, + ) + if return_diagnostics: + return int(outcome), float(probability), diagnostics + return int(outcome), float(probability) + + def measure_pauli(self, pauli, where, outcome=None, *, renormalize=True, + return_diagnostics=False): + """Measure a product Pauli and return ``(outcome, probability)``. + + ``pauli`` supplies one ``X``, ``Y``, or ``Z`` axis per site in + ``where``. An outcome is sampled unless forced to ``+1`` or ``-1``. + The state is normalized after collapse by default; set + ``renormalize=False`` to retain the branch norm. Requesting + diagnostics adds a third return value containing projection norm, + support/span, and bond-dimension before/after information. """ - where = _normalize_where(where) - if len(where) == 1: - q = where[0] - leaf = self.plan.leaf_of_qubit[q] - self._move_center(leaf) - t = self.tn.tensor_map[self._tid(leaf)] - p = self._phys(q) - mat = np.asarray(op).reshape(2, 2) - gt = qtn.Tensor(mat, inds=(p + "*", p)) - bra = t.H.reindex_({p: p + "*"}) - num = qtn.tensor_contract(bra, gt, t, output_inds=[]) - den = qtn.tensor_contract(t.H, t, output_inds=[]) - return num / den + return self._measure_pauli( + pauli, where, outcome, renormalize=renormalize, + return_diagnostics=return_diagnostics, + ) - phys = [self._phys(q) for q in where] - leaves = [self.plan.leaf_of_qubit[q] for q in where] - snodes = self._steiner_nodes(leaves) - # The reduced contraction is exact only when the centre lies inside the - # spanning subtree (then the isometric exterior cancels); ensure it. - if self.center not in snodes: - self._move_center(leaves[0]) - # Bonds internal to the subtree must be renamed in the bra so they do not - # collide with the ket; boundary bonds stay shared so the isometric - # exterior contributes an identity between bra and ket. - internal = set() - for nid in snodes: - for nb in self._neighbors(nid): - if nb in snodes: - internal.add(self._bond_name(nid, nb)) - ket_ts = [self.tn.tensor_map[self._tid(nid)].copy() for nid in snodes] - ket = qtn.TensorNetwork(ket_ts) - mat = np.asarray(op).reshape([2] * (2 * len(where))) - gt = qtn.Tensor(mat, inds=[p + "*" for p in phys] + phys) + def _reset_pauli(self, q, axis): + """Reset one qubit to the ``+1`` eigenstate of ``axis``.""" + outcome, _ = self._measure_pauli(axis, (q,)) + if outcome < 0: + flip = _RESET_FLIP_AXES[axis] + self.apply_1q( + self._as_state_backend(_PAULI_1Q[flip], warn=False), q + ) + + def _apply_control_event(self, name, payload, where): + """Apply one normalized measure/reset event from the queued stream.""" with self._thread_ctx(): - internal_map = {ix: qtn.rand_uuid() for ix in internal} - bra_num = ket.H.reindex( - {**internal_map, **{p: p + "*" for p in phys}} + return self._apply_control_event_impl(name, payload, where) + + def _apply_control_event_impl(self, name, payload, where): + """Apply a control event without opening another thread context.""" + if name == "cap": + self.cap( + where[0], payload["vec"], + absorb=payload.get("absorb", "left"), + compact_labels=payload.get("compact_labels", True), ) - num = (bra_num & gt & ket).contract(output_inds=[]) - bra_den = ket.H.reindex(internal_map) - den = (bra_den & ket).contract(output_inds=[]) - return num / den + return self + if name == "measure": + outcome, probability = self._measure_pauli( + payload["pauli"], where, payload.get("outcome") + ) + self.measurements.append( + (str(payload["pauli"]), tuple(where), outcome, probability) + ) + return self + if name == "reset": + for q, axis in zip(where, payload["axes"]): + self._reset_pauli(q, axis) + return self + if name == "measure_reset": + for q, axis, forced in zip( + where, payload["axes"], payload["outcomes"] + ): + outcome, probability = self._measure_pauli(axis, (q,), forced) + self.measurements.append( + (axis, (q,), outcome, probability) + ) + if outcome < 0: + self.apply_1q( + self._as_state_backend( + _PAULI_1Q[_RESET_FLIP_AXES[axis]], warn=False + ), + q, + ) + return self + raise ValueError(f"Unknown TreeOptimizer control event {name!r}.") def measure(self, q, outcome=None): """Projectively measure qubit ``q`` in the computational basis. @@ -789,14 +3776,17 @@ def measure(self, q, outcome=None): Returns the outcome bit. Because the centre sits on the leaf 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)] p = self._phys(q) ax = t.inds.index(p) - arr = np.moveaxis(t.data, ax, 0).reshape(2, -1) - w = np.einsum("ij,ij->i", arr, arr.conj()).real + arr = ar.do("reshape", ar.do("moveaxis", t.data, ax, 0), (2, -1)) + w = ar.do("sum", arr * ar.do("conj", arr), axis=1) + w = np.real(ar.to_numpy(w)) total = float(w.sum()) if total <= 0: raise ValueError("Cannot measure a zero-norm state.") @@ -805,20 +3795,70 @@ def measure(self, q, outcome=None): if outcome is None: outcome = int(self.rng.choice(2, p=probs)) else: + if not isinstance(outcome, Integral) or int(outcome) not in (0, 1): + raise ValueError("measurement outcome must be 0 or 1.") outcome = int(outcome) - proj = np.zeros((2, 2), dtype=self.dtype) + if probs[outcome] <= 1e-12: + raise ValueError( + f"forced measure outcome {outcome} has ~0 probability." + ) + proj = np.zeros((2, 2), dtype=complex) proj[outcome, outcome] = 1.0 - self.apply_1q(proj, q) + self.apply_1q(self._as_state_backend(proj, warn=False), q) self.normalize() return outcome def reset(self, q): """Reset qubit ``q`` to ``|0>`` (measure, then flip if it was ``|1>``).""" if self.measure(q) == 1: - x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=self.dtype) - self.apply_1q(x, q) + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + self.apply_1q(self._as_state_backend(x, warn=False), q) return 0 + def cap(self, q, vec, *, absorb="left", compact_labels=True, + stable_labels=None): + """Contract and remove qubit ``q`` from the tree state. + + This is the tree analogue of an MPS physical-index cap. Trees have no + left/right chain neighbour at a leaf, so ``absorb`` is accepted for + stream compatibility but the unique parent is used. By default the + remaining caller-facing labels are compacted above ``q``. Set + ``compact_labels=False`` (or ``stable_labels=True``) to retain the + original logical labels while the internal TTN stays compact. + """ + if stable_labels is not None: + compact_labels = not bool(stable_labels) + if not isinstance(compact_labels, (bool, np.bool_)): + raise ValueError("compact_labels must be boolean.") + self._invalidate_state_norm_cache() + logical_q = int(q) + q = self._validate_qubit(logical_q) + if absorb not in {"left", "right"}: + raise ValueError("cap absorb direction must be 'left' or 'right'.") + started = self._begin_update("cap", (q,)) + try: + self.tn.cap_qubit_(q, self._as_state_backend(vec)) + self.plan = self.tn.plan + self.n = self.tn.nqubits + remaining = [label for label in self._logical_qubits if label != logical_q] + if compact_labels: + remaining = [ + label - 1 if label > logical_q else label + for label in remaining + ] + self._logical_qubits = remaining + self._logical_positions = { + label: position for position, label in enumerate(remaining) + } + self.layout_finder = None + except Exception: + if started: + self._abort_update() + raise + if started: + self._finish_update() + return self + def copy(self): """Return an independent optimizer at the current tree state. @@ -828,38 +3868,92 @@ def copy(self): sequences. The immutable :class:`TreePlan` is shared; the gate queue is retained (gate payloads are not copied) but not replayed. """ + child_seed = int(self.rng.integers(0, 2**63, dtype=np.uint64)) other = type(self)( None, n=self.n, chi=self.chi, cutoff=self.cutoff, + mode=self.mode, structure=self.structure, + max_arity=self.max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, tree=self.plan, dtype=self.dtype, threads=self.threads, + layout_objective=self.layout_objective, + layout_weight_mode=self.layout_weight_mode, + track_truncation=self.track_truncation, + 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, + seed=child_seed, run=False, + tn=self.tn, ) - # ``TreeTensorNetwork.copy`` carries ``_canonical_region`` (an - # _EXTRA_PROPS field), so the copied network already reports the right - # centre / canonical region. - other.tn = self.tn.copy() other.G = list(self.G) other.where = list(self.where) + other.event_types = list(self.event_types) + other.layout_finder = self.layout_finder + other.measurements = deepcopy(self.measurements) + other.truncation_history = deepcopy(self.truncation_history) + other.update_history = deepcopy(self.update_history) + other.infidelities = list(self.infidelities) + other.infidelity_samples = deepcopy(self.infidelity_samples) + other.normalizations = deepcopy(self.normalizations) + other.projection_diagnostics = deepcopy(self.projection_diagnostics) + other._logical_qubits = list(self._logical_qubits) + other._logical_positions = dict(self._logical_positions) + other._truncation_survival = self._truncation_survival return other + def get_infidelities(self): + """Return cumulative tracked tree-truncation infidelities. + + The first value is ``0.0``. Additional values are recorded after + updates for which ``track_truncation=True`` supplied singular spectra. + Unlike a chain MPS, a tree has several compression edges per update, + so each value aggregates the retained weight over all touched edges. + """ + return self.infidelities + + def get_infidelity_samples(self): + """Return detailed cumulative tree-truncation sample records.""" + return self.infidelity_samples + + 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. + """ + return self.normalizations + + def get_projection_diagnostics(self): + """Return projection norm/support/span/bond diagnostics in order.""" + return self.projection_diagnostics + @classmethod def find_tree_layout(cls, gates, n=None, *, structure="quality", - max_arity=2, community_frac=0.35, star_frac=0.75): + max_arity=(2, 3, 4), community_frac=0.35, + star_frac=0.75, layout_objective="path", + layout_weight_mode="count", + 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, + star_frac=star_frac, objective=layout_objective, + weight_mode=layout_weight_mode, + max_operator_qubits=max_operator_qubits, ).run() @classmethod def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, - ops=None, structure="quality", max_arity=2, + ops=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, tree=None, dense_cap=1 << 14): """Replay ``gates`` at several ``chi`` and report convergence. @@ -892,6 +3986,8 @@ def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, ``max_drift`` (max ``|Delta|`` from the previous ``chi``, or ``None`` for the first). """ + if hasattr(gates, "__next__"): + gates = list(gates) 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, @@ -918,7 +4014,13 @@ def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, expectations = {} vals = [] for i, (op, where) in enumerate(ops): - val = complex(opt.local_expectation(op, where)) + # Observable evaluation belongs to the state object. Keep the + # replay optimizer focused on state evolution rather than a + # qubit-only local-expectation implementation. + val = complex(opt.tn.local_expectation( + op, _normalize_where(where), max_bond=None, + optimize="auto", + )) expectations[f"op{i}"] = val vals.append(val) fidelity = None diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index c48e2a5..0093447 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -17,9 +17,9 @@ structural 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 the physical index - ``site_ind_id.format(q)`` (default ``"k{}"``) for qubit ``q`` -- so the - inherited ``quimb`` site/ ``local_expectation`` machinery treats the leaves as - the sites and the internal nodes as ancillary bond carriers; + ``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; * adjacent nodes ``a`` and ``b`` share the deterministic virtual bond index ``_tb{lo}_{hi}`` with ``lo, hi = sorted((a, b))``. @@ -31,27 +31,120 @@ from __future__ import annotations +import re + +import autoray as ar import numpy as np import quimb.tensor as qtn -from quimb.tensor import TensorNetworkGenVector from quimb.tensor.tensor_core import TensorNetwork +from numbers import Integral from .layout import TreePlan __all__ = ["TreeTensorNetwork"] +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 + class TensorNetworkGenVector(TensorNetwork): + """Small compatibility base for quimb versions without GenVector.""" + + @property + def nsites(self): + return len(getattr(self, "_sites", ())) + + @property + def sites(self): + return tuple(getattr(self, "_sites", ())) + + @property + def site_ind_id(self): + return self._site_ind_id + + @property + def site_tag_id(self): + return self._site_tag_id + + def site_ind(self, site): + return self._site_ind_id.format(site) + + def site_tag(self, site): + return self._site_tag_id.format(site) + + def local_expectation(self, operator, where, *, optimize="auto-hq", **kwargs): + """Evaluate a local expectation using a generic TN contraction.""" + if isinstance(where, Integral): + where = (int(where),) + else: + where = tuple(where) + operated = qtn.tensor_network_gate_inds( + self, + operator, + [self.site_ind(site) for site in where], + contract=False, + inplace=False, + tags=[], + ) + numerator = (self.H | operated).contract(all, optimize=optimize) + denominator = (self.H | self).contract(all, optimize=optimize) + return numerator / denominator + + def _bond_index(a, b): """Return the deterministic virtual-bond index name for edge ``(a, b)``.""" lo, hi = (a, b) if a < b else (b, a) return f"_tb{lo}_{hi}" +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def _is_symmray_array(value): + """Whether ``value`` is a native Symmray block-sparse array.""" + try: + return ar.infer_backend(value) == "symmray" + except (AttributeError, TypeError): + return hasattr(value, "blocks") and hasattr(value, "indices") + + +def _visible_len(s): + """Length of ``s`` ignoring any ANSI colour escape sequences.""" + return len(_ANSI_RE.sub("", s)) + + def _ascii_place(s, width, col): - """Return ``s`` padded to ``width`` with its centre aligned at column ``col``.""" - start = col - (len(s) - 1) // 2 - start = max(0, min(start, width - len(s))) - return " " * start + s + " " * (width - start - len(s)) + """Return ``s`` padded to ``width`` with its centre aligned at column ``col``. + + ANSI-colour aware: the padding is computed from the *visible* length so that + embedded colour escapes never shift the drawing. + """ + vis = _visible_len(s) + start = col - (vis - 1) // 2 + start = max(0, min(start, width - vis)) + return " " * start + s + " " * (width - start - vis) + + +# Depth-cycled palette for the internal-node markers, a distinct leaf colour, +# and a dim style for the bond numbers / connector lines so the coloured nodes +# stand out. 256-colour SGR codes (rendered by VS Code notebooks and modern +# terminals); on a mono terminal they degrade to plain text. +_LAYER_COLORS = ( + "\x1b[1;38;5;213m", # root -> bright magenta / pink + "\x1b[1;38;5;45m", # bright cyan + "\x1b[1;38;5;82m", # bright green + "\x1b[1;38;5;220m", # gold + "\x1b[1;38;5;208m", # orange + "\x1b[1;38;5;99m", # violet +) +_LEAF_COLOR = "\x1b[1;38;5;39m" # leaves (◆) -> strong blue +_DIM_STYLE = "\x1b[38;5;244m" # bond dims + connectors -> grey +_RESET = "\x1b[0m" + + +def _color(s, code, enable): + """Wrap ``s`` in the SGR ``code`` when ``enable`` is set, else return ``s``.""" + return f"{code}{s}{_RESET}" if enable and code else s @@ -61,9 +154,9 @@ class TreeTensorNetwork(TensorNetworkGenVector): Subclasses :class:`quimb.tensor.TensorNetworkGenVector`, so it *is* a ``quimb`` tensor network: all of ``quimb``'s arbitrary-geometry methods (``canonize_around``, ``canonize_between``, ``compress_between``, - ``gate_inds``, ``local_expectation``, ``to_dense``, ``copy`` ...) work - directly. The tree geometry is owned by a :class:`TreePlan`; this class adds - the node/site/index naming glue on top. + ``gate_inds``, ``to_dense``, ``copy`` ...) work directly. This class provides + the tree-specific ``local_expectation`` implementation and owns the + geometry/node/site/index naming glue on top of a :class:`TreePlan`. Prefer the builders :meth:`from_plan`, :meth:`from_order`, and :meth:`rand` over calling the constructor with raw tensors. @@ -90,15 +183,50 @@ class TreeTensorNetwork(TensorNetworkGenVector): "_plan", "_node_tag_id", "_canonical_region", + "_symmetry", + "_fermionic", + "_physical_sectors", ) def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", - site_ind_id="k{}", node_tag_id="N{}", **tn_opts): + site_ind_id="k{}", node_tag_id="N{}", symmetry=None, + fermionic=False, physical_sectors=None, **tn_opts): # Copy / cast path: quimb's base ``__init__`` copies ``_EXTRA_PROPS`` # (``_plan``, ``_node_tag_id``, ...) straight off ``ts``; returning here # avoids clobbering them with the fresh-construction defaults below. if isinstance(ts, TensorNetwork): super().__init__(ts, **tn_opts) + if isinstance(ts, TreeTensorNetwork): + if plan is not None and ( + plan.root != ts.plan.root + or plan.children != ts.plan.children + or plan.qubit_of_leaf != ts.plan.qubit_of_leaf + ): + raise ValueError( + "plan does not match the TreeTensorNetwork being copied." + ) + self._fermionic_norm_cache = None + self._fermionic_norm_cache_version = 0 + self._fermionic_norm_cache_value_version = None + return + if plan is None: + raise TypeError( + "casting a plain TensorNetwork to TreeTensorNetwork requires " + "an explicit TreePlan." + ) + self._plan = plan + self._sites = tuple(range(plan.n)) if sites is None else tuple(sites) + self._site_tag_id = site_tag_id + self._site_ind_id = site_ind_id + self._node_tag_id = node_tag_id + self._canonical_region = None + self._symmetry = symmetry + self._fermionic = bool(fermionic) + self._physical_sectors = physical_sectors + self._fermionic_norm_cache = None + self._fermionic_norm_cache_version = 0 + self._fermionic_norm_cache_value_version = None + self.validate() return super().__init__(ts, **tn_opts) if plan is None: @@ -110,12 +238,70 @@ def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", self._site_tag_id = site_tag_id self._site_ind_id = site_ind_id self._node_tag_id = node_tag_id + self._symmetry = symmetry + self._fermionic = bool(fermionic) + self._physical_sectors = physical_sectors + self._fermionic_norm_cache = None + self._fermionic_norm_cache_version = 0 + self._fermionic_norm_cache_value_version = None # Frozenset of node ids forming the canonicalised subtree (``None`` if # unknown); a one-node region is exactly an orthogonality centre. # 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 + # -- mutation / canonical metadata -------------------------------------- + + def invalidate_canonical_form(self): + """Forget the tracked canonical region after an unmanaged mutation. + + Quimb exposes mutating tensor-network methods and tensors themselves + can be modified in place. The TTN cannot intercept every such mutation, + so direct callers should use this method after changing tensor data + outside the state-aware wrappers below. + """ + self._canonical_region = None + self._invalidate_norm_cache() + return self + + def _invalidate_norm_cache(self): + """Invalidate the cached native-fermion norm denominator.""" + self._fermionic_norm_cache = None + self._fermionic_norm_cache_value_version = None + self._fermionic_norm_cache_version = ( + getattr(self, "_fermionic_norm_cache_version", 0) + 1 + ) + return self + + def _invalidate_after_mutation(self, result): + self._canonical_region = None + self._invalidate_norm_cache() + return result + + def gate_inds_(self, *args, **kwargs): + """Apply a Quimb gate and invalidate canonical metadata.""" + return self._invalidate_after_mutation( + super().gate_inds_(*args, **kwargs) + ) + + def canonize_between(self, *args, **kwargs): + """Canonicalize through Quimb and invalidate the tracked centre.""" + return self._invalidate_after_mutation( + super().canonize_between(*args, **kwargs) + ) + + def compress_between(self, *args, **kwargs): + """Compress through Quimb and invalidate the tracked centre.""" + return self._invalidate_after_mutation( + super().compress_between(*args, **kwargs) + ) + + def canonize_around_(self, *args, **kwargs): + """Canonicalize around tags through Quimb and invalidate metadata.""" + return self._invalidate_after_mutation( + super().canonize_around_(*args, **kwargs) + ) + # -- geometry / naming ---------------------------------------------------- @property @@ -128,11 +314,238 @@ def node_tag_id(self): """Format string for the structural node tag (e.g. ``"N{}"``).""" return self._node_tag_id + @property + def symmetry(self): + """Native Symmray symmetry label, or ``None`` for dense tensors.""" + return self._symmetry + + @property + def fermionic(self): + """Whether the live tensor data uses Symmray's fermionic arrays.""" + return self._fermionic + + def _fermionic_norm_squared(self): + """Return the exact native-fermion norm squared with caching. + + A known centre uses Symmray's graded one-tensor contraction. If the + gauge is unknown, contract the complete doubled network so Quimb's + fermionic contraction machinery retains all graded boundary phases. + The result is cached until a state mutation invalidates it. + """ + if not self.fermionic: + raise TypeError("fermionic norm readout requires a fermionic TTN.") + + cache = getattr(self, "_fermionic_norm_cache", None) + version = getattr(self, "_fermionic_norm_cache_version", 0) + if ( + cache is not None + and getattr(self, "_fermionic_norm_cache_value_version", None) + == version + ): + return cache + + center = self.orthogonality_center + if center is not None: + value = self._fermionic_center_norm_squared(center) + else: + value = (self.H | self).contract(all, optimize="auto") + self._fermionic_norm_cache = value + self._fermionic_norm_cache_value_version = version + return value + + def _fermionic_center_norm_squared(self, center=None): + """Read a canonical center norm with Symmray's graded conjugation. + + ``Tensor.H`` only conjugates the data. For a fermionic tensor, the + one-tensor network conjugation also applies parity phase flips on its + outer legs. Those flips are the graded identity supplied by the + canonical exterior, so this remains a one-tensor readout while + retaining the correct fermionic norm. If no single center is known, + use the complete doubled-network contraction instead. + """ + if not self.fermionic: + raise TypeError("fermionic center norms require a fermionic TTN.") + if center is None: + center = self.orthogonality_center + if center is None: + return self._fermionic_norm_squared() + tensor = self.node_tensor(center).copy() + singleton = qtn.TensorNetwork([tensor]) + return (singleton.H & singleton).contract(all, optimize="auto") + + 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 + + def _restore_readout_region(self, region): + """Restore a dense readout's tracked canonical region.""" + if region is None: + return + if len(region) == 1: + self.shift_orthogonality_center(next(iter(region))) + else: + self.canonize_subtree_(region) + + @property + def physical_sectors(self): + """Native Symmray physical-sector map when one was supplied.""" + return self._physical_sectors + @property def root(self): """The root node id of the tree.""" return self._plan.root + def local_expectation( + self, operator, where, *, max_bond=None, optimize="auto", + normalized=True, **kwargs, + ): + """Evaluate a local observable with a backend-specific exact path. + + Dense/nonfermionic trees use the canonical target leaf 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. + + ``max_bond`` and extra keyword arguments are accepted for Quimb API + compatibility. This exact tree contraction does not truncate. + """ + preserve_gauge = bool(kwargs.pop("_preserve_gauge", True)) + _ = max_bond, kwargs + 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.leaf_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: + original_region = self.canonical_region + if original_region is None: + # An unknown dense gauge cannot be reconstructed after a + # target canonicalisation, so perform the readout on an + # independent copy. Known regions use the cheaper round-trip + # restoration below. + work = self.copy() + return work.local_expectation( + operator, + where, + max_bond=max_bond, + optimize=optimize, + normalized=normalized, + _preserve_gauge=False, + ) + else: + original_region = None + + leaves = [self.plan.leaf_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): + raise TypeError( + "native Symmray TTNs require a native Symmray observable; " + "use Fermion.observable(...) or another Symmray operator." + ) + if self.fermionic: + expected_rank = 2 * len(where) + if len(ar.shape(op)) != expected_rank: + raise ValueError( + f"a {len(where)}-site Symmray observable must have " + f"rank {expected_rank}." + ) + return self._fermionic_local_expectation( + op, + where, + optimize=optimize, + normalized=normalized, + ) + if len(where) == 1: + self.shift_orthogonality_center(leaves[0]) + tensor = self.node_tensor(leaves[0]) + physical = phys[0] + if not _is_symmray_array(op): + dim = int(tensor.shape[tensor.inds.index(physical)]) + op = ar.do("reshape", op, (dim, dim)) + elif len(ar.shape(op)) != 2: + raise ValueError("a one-site Symmray observable must be rank two.") + gate = qtn.Tensor(op, inds=(physical + "*", physical)) + bra = tensor.H.reindex_({physical: physical + "*"}) + numerator = qtn.tensor_contract(bra, gate, tensor, output_inds=[]) + denominator = qtn.tensor_contract(tensor.H, tensor, output_inds=[]) + result = numerator / denominator if normalized else numerator + if preserve_gauge: + self._restore_readout_region(original_region) + return result + + span = self.steiner_nodes(leaves) + if self.orthogonality_center not in span: + self.shift_orthogonality_center(leaves[0]) + + internal = { + self.bond(node, neighbor) + for node in span + for neighbor in self.neighbors(node) + if neighbor in span + } + ket = qtn.TensorNetwork([ + self.node_tensor(node).copy() for node in span + ]) + if not _is_symmray_array(op): + dims = [ + int(self.node_tensor(leaf).shape[ + self.node_tensor(leaf).inds.index(physical) + ]) + for leaf, physical in zip(leaves, phys) + ] + op = ar.do("reshape", op, tuple(dims + dims)) + elif len(ar.shape(op)) != 2 * len(where): + raise ValueError( + "a multi-site Symmray observable must have one output and " + "one input leg per site." + ) + gate = qtn.Tensor(op, inds=[p + "*" for p in phys] + phys) + internal_map = {index: qtn.rand_uuid() for index in internal} + bra_num = ket.H.reindex({ + **internal_map, + **{physical: physical + "*" for physical in phys}, + }) + numerator = (bra_num & gate & ket).contract( + output_inds=[], optimize=optimize, + ) + if not normalized: + if preserve_gauge: + self._restore_readout_region(original_region) + return numerator + bra_den = ket.H.reindex(internal_map) + denominator = (bra_den & ket).contract( + output_inds=[], optimize=optimize, + ) + result = numerator / denominator + if preserve_gauge: + self._restore_readout_region(original_region) + return result + @property def orthogonality_center(self): """Node id of the tracked orthogonality centre (``None`` if unknown). @@ -140,7 +553,8 @@ def orthogonality_center(self): This is the tree analogue of an MPS canonical centre: when it is a node ``c`` every *other* tensor is an isometry whose legs point toward ``c`` (``absorb="right"`` convention), so the whole state norm collapses onto - the single centre tensor. It is the one-node special case of + the single centre tensor under Symmray's graded singleton contraction. + It is the one-node special case of :attr:`canonical_region`; it is updated in place by :meth:`shift_orthogonality_center` and :meth:`canonize_around_node_`, and -- being derived from a field declared in :attr:`_EXTRA_PROPS` -- it @@ -172,8 +586,9 @@ def canonical_region(self): it is a connected node set ``R`` every tensor *outside* ``R`` is an isometry whose legs point inward toward ``R`` (``absorb="right"`` convention), so the entire state norm is carried by the region tensors - -- contracting just the region against its conjugate gives the squared - norm, exactly as the single centre tensor does for a one-node region. + -- contracting just the region against its graded conjugate gives the + squared norm, exactly as the single centre tensor does for a one-node + region. It is updated in place by :meth:`canonize_subtree_` (and its qubit-level entry point :meth:`canonize_around_qubits_`) and, being declared in :attr:`_EXTRA_PROPS`, survives ``.copy()`` and every ``quimb`` view. @@ -225,10 +640,153 @@ def node_tensor(self, nid): return self.tensor_map[self.node_tid(nid)] def bond(self, a, b): - """Return the shared virtual-bond index name for adjacent nodes ``a``/``b``.""" + """Return the live shared virtual-bond index for adjacent nodes.""" if b not in self.neighbors(a): raise ValueError(f"nodes {a} and {b} are not adjacent in the tree.") - return _bond_index(a, b) + shared = qtn.bonds(self.node_tensor(a), self.node_tensor(b)) + if len(shared) != 1: + raise ValueError( + f"nodes {a} and {b} must share exactly one bond; " + f"found {sorted(shared)}." + ) + return next(iter(shared)) + + def validate(self, *, check_canonical=False, tol=1e-9): + """Validate the live network against its :class:`TreePlan`. + + 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 + ``check_canonical=True`` to additionally verify the tracked canonical + region; that part performs tensor contractions and is more expensive. + + Raises + ------ + ValueError + If the plan and live tensor network disagree. The message names + the first invariant that failed. Returns ``self`` when valid. + """ + plan_nodes = tuple(self._plan.nodes()) + plan_node_set = set(plan_nodes) + node_tids = {} + for nid in plan_nodes: + tag = self.node_tag(nid) + tids = set(self.tag_map.get(tag, ())) + if len(tids) != 1: + raise ValueError( + f"tree node {nid} must have exactly one tensor tagged " + f"{tag!r}; found {len(tids)}." + ) + node_tids[nid] = next(iter(tids)) + + live_tids = set(self.tensor_map) + if set(node_tids.values()) != live_tids: + extra = live_tids - set(node_tids.values()) + missing = set(node_tids.values()) - live_tids + raise ValueError( + "TreeTensorNetwork tensor set disagrees with TreePlan " + f"(extra={sorted(extra)!r}, missing={sorted(missing)!r})." + ) + + physical_inds = {self.site_ind(q) for q in range(self._plan.n)} + owners = {ind: [] for ind in physical_inds} + 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] + expected = self.site_ind(q) + if expected not in tensor.inds: + raise ValueError( + f"leaf 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"{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}." + ) + elif node_phys: + raise ValueError( + f"internal node {nid} unexpectedly carries physical " + f"indices {sorted(node_phys)!r}." + ) + for ind in node_phys: + owners[ind].append(nid) + + 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"found nodes {ind_owners!r}." + ) + + expected_edges = { + frozenset((parent, child)) + for parent, children in self._plan.children.items() + for child in children + } + edge_bonds = {} + for parent, children in self._plan.children.items(): + for child in children: + shared = qtn.bonds( + self.tensor_map[node_tids[parent]], + self.tensor_map[node_tids[child]], + ) + if len(shared) != 1: + raise ValueError( + f"tree edge ({parent}, {child}) must have exactly one " + f"live bond; found {sorted(shared)!r}." + ) + edge_bonds[frozenset((parent, child))] = next(iter(shared)) + + actual_inner = set(self.inner_inds()) + if actual_inner != set(edge_bonds.values()): + raise ValueError( + "live internal indices do not match the TreePlan edges " + f"(extra={sorted(actual_inner - set(edge_bonds.values()))!r}, " + f"missing={sorted(set(edge_bonds.values()) - actual_inner)!r})." + ) + tid_to_node = {tid: nid for nid, tid in node_tids.items()} + for ind in actual_inner: + tids = set(self.ind_map.get(ind, ())) + if len(tids) != 2: + raise ValueError( + f"internal index {ind!r} must occur on two tensors; " + f"found {len(tids)}." + ) + owners = frozenset(tid_to_node[tid] for tid in tids) + if owners not in expected_edges: + raise ValueError( + f"internal index {ind!r} joins non-adjacent nodes " + f"{sorted(owners)!r}." + ) + if int(self.ind_size(ind)) < 1: + raise ValueError( + f"internal index {ind!r} has invalid dimension " + f"{self.ind_size(ind)!r}." + ) + + region = self.canonical_region + if region is not None: + if not set(region).issubset(plan_node_set): + raise ValueError( + f"canonical region contains unknown nodes " + f"{sorted(set(region) - plan_node_set)!r}." + ) + 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.") + return self # -- plan delegators ------------------------------------------------------ @@ -276,6 +834,11 @@ def steiner_nodes(self, leaves): 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: @@ -335,7 +898,78 @@ def _toward_region(self, nid, region): # -- edge-level canonical / compression helpers --------------------------- - def _track_edge_center(self, a, b, absorb): + def _fermionic_canonize_edge_(self, a, b, absorb): + """Move a native graded centre across one edge by explicit QR.""" + if absorb == "right": + isometric_node, reduced_node = a, b + elif absorb == "left": + isometric_node, reduced_node = b, a + else: + raise ValueError("absorb must be 'right' or 'left'.") + + isometric = self.node_tensor(isometric_node) + 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( + left_inds=left_inds, + right_inds=(bond,), + method="qr", + absorb="right", + cutoff=0.0, + get="tensors", + ) + merged = qtn.tensor_contract(carry, 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 + + def _fermionic_compress_edge_( + self, a, b, *, max_bond, cutoff, absorb, + ): + """Compress one native graded tree cut by an explicit two-node SVD.""" + if absorb == "right": + isometric_node, reduced_node = a, b + elif absorb == "left": + isometric_node, reduced_node = b, a + else: + raise ValueError("absorb must be 'right' or 'left'.") + + isometric = self.node_tensor(isometric_node) + reduced = self.node_tensor(reduced_node) + bond = self.bond(isometric_node, reduced_node) + left_inds = [index for index in isometric.inds if index != bond] + theta = qtn.tensor_contract(isometric, reduced) + kept, remainder = theta.split( + left_inds=left_inds, + method="svd", + max_bond=max_bond, + cutoff=cutoff, + absorb="right", + get="tensors", + bond_ind=bond, + ) + isometric.modify( + data=kept.data, + inds=kept.inds, + left_inds=kept.left_inds, + ) + reduced.modify( + data=remainder.data, + inds=remainder.inds, + left_inds=None, + ) + return self + + def _track_edge_center(self, a, b, absorb, *, previous=None): """Update the tracked centre after a gauge move across edge ``a -> b``. A single ``absorb="right"`` move makes ``a`` isometric and pushes the @@ -344,7 +978,7 @@ def _track_edge_center(self, a, b, absorb): longer the global centre after a lone edge move, so it is set to ``None`` (unknown) rather than left lying about the canonical form. """ - cur = self.orthogonality_center + cur = self.orthogonality_center if previous is None else previous if absorb == "right" and cur == a: self._canonical_region = frozenset({b}) elif absorb == "left" and cur == b: @@ -355,32 +989,54 @@ def _track_edge_center(self, a, b, absorb): def canonize_edge_(self, a, b, absorb="right"): """Canonicalise across the tree edge ``a -> b`` in place. - Thin wrapper over :meth:`quimb.tensor.TensorNetwork.canonize_between` - that resolves node ids to node tags; ``absorb="right"`` leaves node ``a`` - isometric and pushes the orthogonality centre onto node ``b``. The - tracked :attr:`orthogonality_center` is advanced accordingly. + 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``. """ - self.canonize_between(self.node_tag(a), self.node_tag(b), absorb=absorb) - self._track_edge_center(a, b, absorb) + previous = self.orthogonality_center + if self.fermionic: + self._fermionic_canonize_edge_(a, b, absorb) + else: + self.canonize_between( + self.node_tag(a), + self.node_tag(b), + absorb=absorb, + method="qr", + cutoff=0.0, + ) + self._invalidate_norm_cache() + 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"): """Compress the tree edge ``a -> b`` in place. - Thin wrapper over :meth:`quimb.tensor.TensorNetwork.compress_between` - (local reduced compression, ``canonize_distance=0``) resolving node ids - to node tags. The tracked :attr:`orthogonality_center` is advanced as - for :meth:`canonize_edge_` (compression moves the gauge the same way). + 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_`. """ - self.compress_between( - self.node_tag(a), - self.node_tag(b), - max_bond=max_bond, - cutoff=cutoff, - absorb=absorb, - ) - self._track_edge_center(a, b, absorb) + previous = self.orthogonality_center + if self.fermionic: + self._fermionic_compress_edge_( + a, + b, + max_bond=max_bond, + cutoff=cutoff, + absorb=absorb, + ) + else: + self.compress_between( + self.node_tag(a), + self.node_tag(b), + max_bond=max_bond, + cutoff=cutoff, + absorb=absorb, + ) + self._invalidate_norm_cache() + self._track_edge_center(a, b, absorb, previous=previous) return self def canonize_around_node_(self, nid): @@ -415,7 +1071,13 @@ 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] - self.canonize_around_(tags, which="any", absorb=absorb) + self.canonize_around_( + tags, + which="any", + absorb=absorb, + method="qr", + cutoff=0.0, + ) self._canonical_region = region return self @@ -428,9 +1090,57 @@ def canonize_around_qubits_(self, qubits, *, absorb="right"): those qubits is captured by that subtree. Equivalent to ``canonize_subtree_(leaves_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) + def _recover_center_from_region(self, region, target, *, absorb="right"): + """Recover one centre by peeling a tracked canonical region. + + A multi-node canonical region already has every exterior branch + pointing inward. Peeling the region's leaves toward ``target`` with + lossless QR therefore establishes a single centre without touching + tensors outside the region. + """ + region = set(region) + if target not in region: + raise ValueError("target centre must lie inside the canonical region.") + if absorb not in {"right", "left"}: + raise ValueError("absorb must be 'right' or 'left'.") + + remaining = set(region) + while len(remaining) > 1: + candidates = [ + node for node in remaining + if node != target + and sum( + neighbour in remaining + for neighbour in self.neighbors(node) + ) == 1 + ] + if not candidates: + raise ValueError( + "canonical region is not a connected tree containing target." + ) + # The region is a tree, so any non-target leaf can be peeled. + # Sorting keeps the QR sequence deterministic across set order. + node = min(candidates) + neighbour = next( + neighbour + for neighbour in self.neighbors(node) + if neighbour in remaining + ) + if absorb == "right": + a, b = node, neighbour + else: + a, b = neighbour, node + self.canonize_edge_(a, b, absorb=absorb) + remaining.remove(node) + + self._canonical_region = frozenset({target}) + return self + def shift_orthogonality_center(self, new, *, absorb="right"): """Move the tracked orthogonality centre to node ``new`` in place. @@ -441,9 +1151,9 @@ def shift_orthogonality_center(self, new, *, absorb="right"): on that path are touched, so a nearby move is O(path length), not O(N). * If the centre is already ``new`` this is a no-op (idempotent). - * If the centre is currently unknown (``None``) it falls back once to the - O(N) :meth:`canonize_around_node_`, then subsequent moves are - incremental. + * If the centre is unknown but a multi-node canonical region is tracked, + only that region is QR-canonicalised first. Otherwise it falls back + once to the O(N) :meth:`canonize_around_node_`. Returns ``self`` so moves can be chained. """ @@ -453,11 +1163,33 @@ def shift_orthogonality_center(self, new, *, absorb="right"): if cur == new: return self if cur is None: - return self.canonize_around_node_(new) + region = self.canonical_region + if region: + if new in region: + return self._recover_center_from_region( + region, new, absorb=absorb + ) + entry = min( + region, + key=lambda node: len(self._plan.node_path(node, new)), + ) + self._recover_center_from_region( + region, entry, absorb=absorb + ) + cur = entry + else: + return self.canonize_around_node_(new) path = self._plan.node_path(cur, new) for u, v in zip(path, path[1:]): - self.canonize_between(self.node_tag(u), self.node_tag(v), - absorb=absorb) + if absorb == "right": + a, b = u, v + elif absorb == "left": + # ``absorb='left'`` centres the first tag. Reverse the edge + # orientation so that the target ``v`` receives the factor. + a, b = v, u + else: + raise ValueError("absorb must be 'right' or 'left'.") + self.canonize_edge_(a, b, absorb=absorb) self._canonical_region = frozenset({new}) return self @@ -501,13 +1233,129 @@ def is_subtree_canonical_form(self, nodes=None, *, span=False, tol=1e-9): toward = self._toward_region(nid, region) t = self.node_tensor(nid) bond = next(iter(qtn.bonds(t, self.node_tensor(toward)))) - tc = t.H.reindex({bond: bond + "*"}) - prod = qtn.tensor_contract(t, tc, output_inds=[bond, bond + "*"]) + if self.fermionic: + # A singleton TensorNetwork.H applies the parity phase flips + # on all outer legs. The contraction order must also follow + # the dual orientation of the open bond, as required by + # Symmray's graded matrix product. + singleton = qtn.TensorNetwork([t.copy()]) + tc = next(iter(singleton.H.tensor_map.values())) + tc = tc.reindex({bond: bond + "*"}) + bond_dual = t.data.indices[t.inds.index(bond)].dual + if bond_dual: + output_inds = [bond, bond + "*"] + prod = qtn.tensor_contract(t, tc, output_inds=output_inds) + else: + output_inds = [bond + "*", bond] + prod = qtn.tensor_contract(tc, t, output_inds=output_inds) + else: + tc = t.H.reindex({bond: bond + "*"}) + output_inds = [bond, bond + "*"] + prod = qtn.tensor_contract(t, tc, output_inds=output_inds) d = int(prod.shape[0]) - if not np.allclose(np.asarray(prod.data), np.eye(d), atol=tol): + 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. + 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): return False return True + def cap_qubit_(self, q, vec): + """Contract qubit ``q`` with ``vec`` and remove that leaf 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. + """ + q = int(q) + self._invalidate_norm_cache() + if q not in self._plan.leaf_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 + try: + compatible = ( + ar.infer_backend(vec) == ar.infer_backend(like) + and ar.get_dtype_name(vec) == ar.get_dtype_name(like) + and str(getattr(vec, "device", None)) + == str(getattr(like, "device", None)) + ) + except (AttributeError, KeyError, TypeError, ValueError): + compatible = False + if not compatible: + vec = ar.do("array", vec, like=like) + vec = ar.do("reshape", vec, (-1,)) + phys = self.site_ind(q) + if int(ar.shape(vec)[0]) != int(self.ind_size(phys)): + raise ValueError( + f"cap vector length {ar.shape(vec)[0]} does not match the " + f"physical dimension {self.ind_size(phys)} of qubit {q}." + ) + + # 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) + self._canonical_region = None + leaf_t = self.node_tensor(leaf).copy() + parent_t = self.node_tensor(parent).copy() + cap_t = qtn.Tensor(vec, inds=(phys,)) + leaf_message = qtn.tensor_contract(leaf_t, cap_t) + merged = qtn.tensor_contract(parent_t, leaf_message) + + collapse = len(self._plan.children[parent]) == 2 + child = None + tags = set(parent_t.tags) + if collapse: + child = next(c for c in self._plan.children[parent] if c != leaf) + child_t = self.node_tensor(child).copy() + merged = qtn.tensor_contract(merged, child_t) + tags.update(set(child_t.tags) - {self.node_tag(child)}) + tags.discard(self.node_tag(leaf)) + tags.add(self.node_tag(parent)) + + self.delete(self.node_tag(leaf)) + if child is not None: + self.delete(self.node_tag(child)) + live_parent = self.node_tensor(parent) + live_parent.modify(data=merged.data, inds=merged.inds, tags=tags) + + # Compact physical indices and site tags with temporary names to avoid + # collisions when k(q+1) is renamed to the removed kq. + 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_leaf(q) + self._sites = tuple(range(self._plan.n)) + self.__dict__.pop("_node_tid_cache", None) + self._canonical_region = frozenset({parent}) + self.validate() + return self + # -- dense read-out ------------------------------------------------------- def to_statevector(self, order=None): @@ -521,16 +1369,19 @@ def to_statevector(self, order=None): if order is None: order = range(self._plan.n) out_inds = [self.site_ind(q) for q in order] - return np.asarray(self.to_dense(out_inds)).reshape(-1) + return ar.to_numpy(self.to_dense(out_inds)).reshape(-1) # -- ascii drawing -------------------------------------------------------- def _bond_dim(self, a, b): """Return the virtual-bond dimension of the tree edge ``(a, b)`` (1 if absent).""" - ix = _bond_index(a, b) + try: + ix = self.bond(a, b) + except ValueError: + return 1 return int(self.ind_size(ix)) if ix in self.ind_map else 1 - def ascii_tree(self, *, bond_dims=True, node_ids=False): + 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 @@ -555,22 +1406,32 @@ def ascii_tree(self, *, bond_dims=True, node_ids=False): Annotate each branch with its virtual-bond dimension (default True). node_ids : bool Also print the structural node id next to each ``●`` (default False). + color : bool + Colour the ``●`` markers by tree depth (a per-layer palette), the + ``◆`` leaves in a distinct colour, and dim the bond numbers and + connector lines with ANSI SGR codes so each layer reads clearly + (default False). The drawing stays width-correct because padding + ignores the colour escapes. """ plan = self._plan gap = 3 # blank columns between sibling subtrees - def render(nid): + 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]}" w = max(1, len(label)) col = (w - 1) // 2 - return [_ascii_place("◆", w, col), label.center(w)], col, w + marker = _color("◆", _LEAF_COLOR, color) + lbl = _color(label, _LEAF_COLOR, color) + return [_ascii_place(marker, w, col), + _ascii_place(lbl, w, col)], col, w - marker = f"●{nid}" if node_ids else "●" + dot = f"●{nid}" if node_ids else "●" + marker = _color(dot, _LAYER_COLORS[depth % len(_LAYER_COLORS)], color) blocks = [] for child in plan.children[nid]: - lines, col, w = render(child) + lines, col, w = render(child, depth + 1) if bond_dims: d = str(self._bond_dim(child, nid)) if len(d) > w: # widen so a fat bond number still fits @@ -580,7 +1441,7 @@ def render(nid): for ln in lines ] col, w = col + lp, len(d) - lines = [_ascii_place(d, w, col)] + lines + lines = [_ascii_place(_color(d, _DIM_STYLE, color), w, col)] + lines blocks.append((lines, col, w)) # lay the child subtrees side by side, recording each root column @@ -605,6 +1466,9 @@ def render(nid): conn[pcol] = { " ": "┴", "─": "┴", "┌": "├", "┐": "┤", "┬": "┼", }[conn[pcol]] + raw = "".join(conn) + drawn = raw.rstrip() # keep trailing pad outside the colour escape + conn_line = _color(drawn, _DIM_STYLE, color) + raw[len(drawn):] # stack the child blocks under parent marker + connector rows height = max(len(lines) for lines, _, _ in blocks) @@ -619,28 +1483,29 @@ def render(nid): lines = [ _ascii_place(marker, total_w, pcol), - "".join(conn), + conn_line, ] + body return lines, pcol, total_w - lines, _, _ = render(plan.root) + lines, _, _ = render(plan.root, 0) return "\n".join(line.rstrip() for line in lines) - def show(self, *, bond_dims=True, node_ids=False): + 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. + 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)) + print(self.ascii_tree(bond_dims=bond_dims, node_ids=node_ids, color=color)) # -- builders ------------------------------------------------------------- @classmethod - def from_plan(cls, plan, *, dtype=complex, site_tag_id="I{}", + def from_plan(cls, plan, *, dtype=complex, phys_dim=2, site_tag_id="I{}", site_ind_id="k{}", node_tag_id="N{}"): """Build the product state ``|0...0>`` on the geometry ``plan``. @@ -657,7 +1522,7 @@ def from_plan(cls, plan, *, dtype=complex, site_tag_id="I{}", if leaf: q = plan.qubit_of_leaf[nid] inds.append(site_ind_id.format(q)) - shape.append(2) + shape.append(int(phys_dim)) tags.append(site_tag_id.format(q)) for child in plan.children[nid]: inds.append(_bond_index(nid, child)) @@ -678,7 +1543,140 @@ def from_plan(cls, plan, *, dtype=complex, site_tag_id="I{}", site_tag_id=site_tag_id, site_ind_id=site_ind_id, node_tag_id=node_tag_id, - )._with_center(plan.root) + )._with_center(plan.root).validate() + + @classmethod + def from_symmray_plan( + cls, + plan, + *, + symmetry, + physical_sectors, + leaf_charges, + bond_dim=1, + fermionic=False, + seed=None, + dtype="complex128", + site_tag_id="I{}", + site_ind_id="k{}", + node_tag_id="N{}", + subsizes="maximal", + ): + """Build a native Symmray TTN on ``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. + """ + try: + import symmray as sr + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "TreeTensorNetwork.from_symmray_plan requires the optional " + "dependency 'symmray'." + ) from exc + + bond_dim = int(bond_dim) + if bond_dim < 1: + raise ValueError("bond_dim must be a positive integer.") + leaf_charges = dict(leaf_charges) + expected = set(range(plan.n)) + if set(leaf_charges) != expected: + raise ValueError( + "leaf_charges must map every qubit label 0 .. n - 1." + ) + + def zero_charge(value): + if isinstance(value, tuple): + return tuple(0 for _ in value) + return 0 + + neutral = zero_charge(next(iter(leaf_charges.values()))) + edges = [ + (parent, child) + for parent, children in plan.children.items() + for child in children + ] + constructor = ( + sr.TN_fermionic_from_edges_rand + if fermionic + else sr.TN_abelian_from_edges_rand + ) + + if edges: + base = constructor( + symmetry, + edges, + bond_dim=bond_dim, + phys_dim=None, + seed=seed, + dtype=dtype, + site_tag_id=node_tag_id, + site_charge=lambda nid: ( + leaf_charges[plan.qubit_of_leaf[nid]] + if plan.is_leaf(nid) + else neutral + ), + subsizes=subsizes, + ) + else: + base = qtn.TensorNetwork() + + for nid in plan.nodes(): + if edges: + tensor = base[node_tag_id.format(nid)] + bond_indices = list(tensor.data.indices) + bond_duals = list(tensor.data.duals) + inds = list(tensor.inds) + else: + tensor = None + bond_indices = [] + bond_duals = [] + inds = [] + + if plan.is_leaf(nid): + q = plan.qubit_of_leaf[nid] + physical_index = sr.utils.rand_index( + symmetry, + physical_sectors, + dual=False, + subsizes=subsizes, + seed=None if seed is None else int(seed) + nid, + ) + data = sr.utils.get_rand( + symmetry, + shape=[*bond_indices, physical_index], + duals=[*bond_duals, False], + charge=leaf_charges[q], + fermionic=fermionic, + label=nid, + seed=None if seed is None else int(seed) + nid, + dtype=dtype, + subsizes=subsizes, + ) + tags = (node_tag_id.format(nid), site_tag_id.format(q)) + inds.append(site_ind_id.format(q)) + if tensor is None: + base |= qtn.Tensor(data=data, inds=inds, tags=tags) + else: + tensor.modify(data=data, inds=inds, tags=tags) + + ttn = cls( + base, + plan=plan, + site_tag_id=site_tag_id, + site_ind_id=site_ind_id, + node_tag_id=node_tag_id, + symmetry=symmetry, + fermionic=fermionic, + physical_sectors=physical_sectors, + ) + # The generic Symmray constructor creates a valid charge-conserving + # tree. Quimb's native block-aware QR establishes the canonical root + # and accumulates product-state normalization there as well. + ttn.canonize_around_node_(plan.root) + return ttn._with_center(plan.root).validate() @classmethod def from_order(cls, order, *, weights=None, structure="quality", @@ -752,7 +1750,7 @@ def _rand(shape): ) if canonicalize: ttn.canonize_around_node_(plan.root) - return ttn + return ttn.validate() # -- repr ----------------------------------------------------------------- diff --git a/src/pepsy/sampling/__init__.py b/src/pepsy/sampling/__init__.py index 5e662f3..6f9028d 100644 --- a/src/pepsy/sampling/__init__.py +++ b/src/pepsy/sampling/__init__.py @@ -3,6 +3,8 @@ from importlib import import_module from .samplers import ( + FermionConfigurationEncoding, + MpsDiagonalEstimate, MpsBatchSampleResult, MpsSampleResult, MpsSampler, @@ -13,6 +15,8 @@ from .tree import TreeBatchSampleResult, TreeSampleResult, TreeSampler __all__ = [ + "FermionConfigurationEncoding", + "MpsDiagonalEstimate", "MpsBatchSampleResult", "MpsSampleResult", "MpsSampler", diff --git a/src/pepsy/sampling/mps.py b/src/pepsy/sampling/mps.py index 8b17323..5cb6b87 100644 --- a/src/pepsy/sampling/mps.py +++ b/src/pepsy/sampling/mps.py @@ -1,5 +1,19 @@ """MPS sampler facade.""" -from .samplers import MpsBatchSampleResult, MpsSampleResult, MpsSampler, VecSampler +from .samplers import ( + FermionConfigurationEncoding, + MpsBatchSampleResult, + MpsDiagonalEstimate, + MpsSampleResult, + MpsSampler, + VecSampler, +) -__all__ = ["MpsBatchSampleResult", "MpsSampleResult", "MpsSampler", "VecSampler"] +__all__ = [ + "FermionConfigurationEncoding", + "MpsBatchSampleResult", + "MpsDiagonalEstimate", + "MpsSampleResult", + "MpsSampler", + "VecSampler", +] diff --git a/src/pepsy/sampling/samplers.py b/src/pepsy/sampling/samplers.py index 8130e79..bd5b2f7 100644 --- a/src/pepsy/sampling/samplers.py +++ b/src/pepsy/sampling/samplers.py @@ -13,6 +13,8 @@ build_optimizer = None __all__ = [ + "FermionConfigurationEncoding", + "MpsDiagonalEstimate", "MpsBatchSampleResult", "MpsSampleResult", "MpsSampler", @@ -72,6 +74,9 @@ def _normalize_mps_sampler_backend(backend): "pytorch": "torch", "cupy": "cupy", "cp": "cupy", + "symmray": "symmray", + "symmetric": "symmray", + "block_sparse": "symmray", } try: return aliases[key] @@ -82,6 +87,27 @@ def _normalize_mps_sampler_backend(backend): ) from exc +def _normalize_symmray_prefix_strategy(strategy): + if strategy is None: + return "auto" + key = str(strategy).strip().lower().replace("-", "_") + aliases = { + "auto": "auto", + "prefix": "prefix", + "shared_prefix": "prefix", + "serial": "serial", + "one_by_one": "serial", + } + try: + return aliases[key] + except KeyError as exc: + allowed = ", ".join(sorted(aliases)) + raise ValueError( + "Unknown Symmray prefix strategy " + f"{strategy!r}. Expected one of: {allowed}." + ) from exc + + def _mps_array_backend(array): module = type(array).__module__.split(".", 1)[0] if module == "torch": @@ -104,6 +130,234 @@ def _backend_array_to_numpy(array): return np.asarray(array) +def _fermion_symmray_occupations(charge, offset, fermion): + """Decode one Symmray physical code into on-site occupations. + + A physical code is an index into a tensor leg, not a universal fermion + label. In particular, collapsed U1/Z2 sectors retain an offset within a + charge sector. Keeping this conversion next to the sampler avoids mixing + the U1U1 and parity-code conventions at VMC boundaries. + """ + symmetry = str(fermion.symmetry).upper() + spinful = bool(fermion.spinful) + if not spinful: + if symmetry not in {"U1", "Z2"}: + raise ValueError( + "Spinless fermion sampling requires symmetry='U1' or 'Z2'." + ) + occupation = int(charge) + if occupation not in {0, 1} or int(offset) != 0: + raise ValueError( + f"Unexpected spinless {symmetry} physical sector " + f"{(charge, offset)!r}." + ) + return (occupation,) + + offset = int(offset) + if symmetry == "Z2": + charge = int(charge) + if charge == 0: + if offset == 0: + return (0, 0) + if offset == 1: + return (1, 1) + elif charge == 1: + # Symmray's parity-collapsed physical basis is empty, double, + # up, down. This differs from the resolved U1/U1U1 ordering. + if offset == 0: + return (1, 0) + if offset == 1: + return (0, 1) + raise ValueError( + "Spinful Z2 physical sectors must be empty/double or up/down " + f"pairs; got {(charge, offset)!r}." + ) + + if symmetry == "U1": + occupation = int(charge) + if occupation == 0 and offset == 0: + return (0, 0) + if occupation == 1: + if offset == 0: + return (0, 1) + if offset == 1: + return (1, 0) + if occupation == 2 and offset == 0: + return (1, 1) + raise ValueError( + "Spinful U1 physical sectors must be empty, down/up, or double; " + f"got {(charge, offset)!r}." + ) + + if symmetry in {"U1U1", "Z2Z2"}: + occupation = tuple(int(value) for value in charge) + if len(occupation) != 2 or any(value not in {0, 1} for value in occupation): + raise ValueError( + f"Unexpected spinful {symmetry} physical charge {charge!r}." + ) + if offset != 0: + raise ValueError( + f"Spinful {symmetry} physical sectors must not be degenerate." + ) + return occupation + + raise ValueError( + "Unsupported Fermion symmetry for sampled physical-code decoding: " + f"{symmetry!r}." + ) + + +def _fermion_code_order(encoding, *, default=(0, 1, 2, 3)): + """Return physical codes ordered by ``(n_up, n_down)`` bits.""" + if encoding is None: + codes = { + "empty": int(default[0]), + "down": int(default[1]), + "up": int(default[2]), + "double": int(default[3]), + } + else: + try: + codes = { + "empty": int(encoding.empty), + "double": int(encoding.double), + "up": int(encoding.up), + "down": int(encoding.down), + } + except (AttributeError, TypeError, ValueError) as exc: + raise TypeError( + "encoding must expose integer empty, double, up, and down " + "codes, for example FermionSiteEncoding." + ) from exc + values = ( + codes["empty"], + codes["down"], + codes["up"], + codes["double"], + ) + if sorted(values) != [0, 1, 2, 3]: + raise ValueError( + "A spinful fermion encoding must contain exactly the physical " + "codes 0, 1, 2, and 3." + ) + return values + + +def _infer_fermion_code_order(tn, sites): + """Infer a four-state code order from PEPS symmetry metadata when possible.""" + default = (0, 1, 2, 3) + if not sites: + return default + tensor = tn[sites[0]] + data = getattr(tensor, "data", None) + symmetry = str(getattr(data, "symmetry", "")).upper() + if symmetry == "Z2": + # The two even states precede the two odd states in Symmray's + # charge-collapsed physical index. + return (0, 3, 2, 1) + if symmetry not in {"U1U1", "Z2Z2"}: + return default + try: + site_ind = tn.site_ind(sites[0]) + axis = tensor.inds.index(site_ind) + charges = tuple(data.indices[axis].chargemap) + position = {tuple(charge): i for i, charge in enumerate(charges)} + required = ((0, 0), (0, 1), (1, 0), (1, 1)) + if all(charge in position for charge in required): + return tuple(position[charge] for charge in required) + except (AttributeError, IndexError, KeyError, TypeError, ValueError): + pass + return default + + +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) + + +def _prepare_bp_binary_network(tn, *, site_order=None, encoding=None): + """Prepare a binary-output copy for Quimb's binary ``sample_d2bp``. + + Quimb's public D2BP sampler currently samples each output index from + ``[0, 1]``. A four-state fermionic physical leg is therefore represented + as two occupation legs, while all tensor data are converted to dense + NumPy arrays in the private BP copy. The source network is never mutated. + """ + network = getattr(tn, "tn", tn) + if not hasattr(network, "sites") or not hasattr(network, "copy"): + return network, {}, tuple(), (0, 1, 2, 3) + + sites = tuple(network.sites if site_order is None else site_order) + missing = [site for site in sites if site not in network.sites] + if missing: + raise ValueError(f"site_order contains site(s) not in PEPS: {missing!r}") + + work = network.copy() + code_order = ( + _fermion_code_order(encoding) + if encoding is not None + else _infer_fermion_code_order(network, sites) + ) + split_inds = {} + existing_inds = set(work.ind_map) + + for site in sites: + tensor = work[site] + site_ind = work.site_ind(site) + try: + axis = tensor.inds.index(site_ind) + except ValueError as exc: + raise ValueError( + f"Could not locate physical index for PEPS site {site!r}." + ) from exc + + data = _to_dense_numpy(tensor.data) + physical_dim = int(data.shape[axis]) + if physical_dim == 2: + if hasattr(tensor.data, "to_dense"): + tensor.modify(data=data) + continue + if physical_dim != 4: + raise ValueError( + "Quimb BP sampling supports binary physical legs directly " + "and spinful fermion legs through a four-state adapter; got " + f"dimension {physical_dim} at site {site!r}." + ) + + # ``code_order`` maps flattened (up, down) bits back to the PEPS + # physical-index order. Reshaping after this permutation gives BP two + # binary output indices with the intended fermion convention. + data = np.take(data, code_order, axis=axis) + data = data.reshape( + data.shape[:axis] + (2, 2) + data.shape[axis + 1:] + ) + up_ind = f"{site_ind}_bp_up" + down_ind = f"{site_ind}_bp_down" + suffix = 0 + while up_ind in existing_inds or down_ind in existing_inds: + suffix += 1 + up_ind = f"{site_ind}_bp_up_{suffix}" + down_ind = f"{site_ind}_bp_down_{suffix}" + existing_inds.update((up_ind, down_ind)) + inds = list(tensor.inds) + inds[axis:axis + 1] = [up_ind, down_ind] + tensor.modify(data=data, inds=inds) + split_inds[site] = (up_ind, down_ind) + + return work, split_inds, sites, code_order + + def _configs_to_sample_result(configs, probs, *, Lx, Ly, one_d_to_two_d): configs_1d = [] configs_2d = [] @@ -184,6 +438,216 @@ def magnetizations(self) -> np.ndarray: ]) +@dataclass(frozen=True) +class MpsDiagonalEstimate: + """Monte Carlo estimate of a diagonal MPS observable. + + Attributes + ---------- + mean + Sample mean of the observable. + standard_error + Standard error estimated from the unbiased sample variance. It is + ``nan`` when only one sample was requested, because no variance + estimate is available. + n_samples + Number of Born samples used for the estimate. + observable + Canonical observable name accepted by + :meth:`MpsSampler.estimate_fermion_diagonal`. + sites + Physical sites averaged or summed by a one-site observable. + pairs + Physical pairs averaged by a density-correlation observable. + """ + + mean: float + standard_error: float + n_samples: int + observable: str + sites: tuple[int, ...] = () + pairs: tuple[tuple[int, int], ...] = () + + +@dataclass(frozen=True) +class FermionConfigurationEncoding: + """Symmetry-aware meaning of sampled fermionic physical codes. + + ``MpsSampler`` always returns *physical codes*: integer positions in the + MPS physical legs. They must be decoded before a caller interprets them as + spinful occupations. For example, a spinful parity ``Z2`` MPS uses the + physical order ``empty, double, up, down``, whereas its resolved U1 path + uses ``empty, down, up, double``. + + The encoding is site-aware, immutable, and can convert between physical + configurations and occupation configurations without relying on either + convention implicitly. Spinful occupations have shape + ``(batch, n_sites, 2)`` in ``(n_up, n_down)`` order; spinless occupations + have shape ``(batch, n_sites)``. + """ + + symmetry: str + spinful: bool + code_to_occupations: tuple[tuple[tuple[int, ...], ...], ...] + + def __post_init__(self): + symmetry = str(self.symmetry).upper() + spinful = bool(self.spinful) + width = 2 if spinful else 1 + tables = tuple( + tuple(tuple(int(value) for value in occupation) for occupation in table) + for table in self.code_to_occupations + ) + if not tables or any(not table for table in tables): + raise ValueError("A fermion configuration encoding needs every site map.") + for table in tables: + if len(set(table)) != len(table): + raise ValueError("Each physical code must represent one occupation.") + if any( + len(occupation) != width + or any(value not in {0, 1} for value in occupation) + for occupation in table + ): + raise ValueError( + "Fermion occupation entries must contain binary " + f"{'(n_up, n_down)' if spinful else 'occupation'} values." + ) + object.__setattr__(self, "symmetry", symmetry) + object.__setattr__(self, "spinful", spinful) + object.__setattr__(self, "code_to_occupations", tables) + + @property + def n_sites(self) -> int: + """Number of physical MPS sites covered by this encoding.""" + return len(self.code_to_occupations) + + @property + def physical_dims(self) -> tuple[int, ...]: + """Per-site physical-code dimensions.""" + return tuple(len(table) for table in self.code_to_occupations) + + def site_code_map(self, site: int) -> dict[int, tuple[int, ...]]: + """Return a copy of the physical-code map for one site.""" + site = int(site) + if not 0 <= site < self.n_sites: + raise ValueError(f"site must be in 0..{self.n_sites - 1}.") + return dict(enumerate(self.code_to_occupations[site])) + + def _config_rows(self, physical_configs): + rows = np.asarray(_backend_array_to_numpy(physical_configs), dtype=np.int64) + if rows.ndim != 2 or rows.shape[1] != self.n_sites: + raise ValueError( + "physical_configs must have shape " + f"(batch, n_sites={self.n_sites}); got {tuple(rows.shape)}." + ) + for site, dim in enumerate(self.physical_dims): + invalid = (rows[:, site] < 0) | (rows[:, site] >= dim) + if np.any(invalid): + values = np.unique(rows[invalid, site]).tolist() + raise ValueError( + f"physical_configs contain invalid code(s) at site {site}: " + f"{values!r}." + ) + return rows + + def decode(self, physical_configs, *, to_numpy: bool = False): + """Decode physical codes into on-site occupations. + + The result remains on Torch/CuPy when ``physical_configs`` is on that + backend, unless ``to_numpy=True`` is requested. + """ + rows = self._config_rows(physical_configs) + backend = _mps_array_backend(physical_configs) + width = 2 if self.spinful else 1 + + if to_numpy or backend not in {"torch", "cupy"}: + out = np.empty( + rows.shape + ((width,) if self.spinful else ()), + dtype=np.int64, + ) + for site, table in enumerate(self.code_to_occupations): + values = np.asarray(table, dtype=np.int64) + decoded = values[rows[:, site]] + if self.spinful: + out[:, site, :] = decoded + else: + out[:, site] = decoded[:, 0] + return out + + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + codes = physical_configs.to(dtype=torch.long) + shape = tuple(codes.shape) + ((width,) if self.spinful else ()) + out = torch.empty(shape, dtype=torch.long, device=codes.device) + for site, table in enumerate(self.code_to_occupations): + values = torch.as_tensor(table, dtype=torch.long, device=codes.device) + decoded = values[codes[:, site]] + if self.spinful: + out[:, site, :] = decoded + else: + out[:, site] = decoded[:, 0] + return out + + import cupy as cp # pylint: disable=import-outside-toplevel + + codes = cp.asarray(physical_configs, dtype=cp.int64) + shape = tuple(codes.shape) + ((width,) if self.spinful else ()) + out = cp.empty(shape, dtype=cp.int64) + for site, table in enumerate(self.code_to_occupations): + values = cp.asarray(table, dtype=cp.int64) + decoded = values[codes[:, site]] + if self.spinful: + out[:, site, :] = decoded + else: + out[:, site] = decoded[:, 0] + return out + + occupations = decode + + def encode(self, occupations, *, to_numpy: bool = False): + """Encode occupations as the physical codes of the sampled MPS.""" + values = np.asarray(_backend_array_to_numpy(occupations), dtype=np.int64) + expected_shape = ( + (values.shape[0], self.n_sites, 2) + if self.spinful and values.ndim >= 1 + else (values.shape[0], self.n_sites) + if not self.spinful and values.ndim >= 1 + else None + ) + if expected_shape is None or tuple(values.shape) != expected_shape: + suffix = ", 2" if self.spinful else "" + raise ValueError( + "occupations must have shape " + f"(batch, n_sites={self.n_sites}{suffix}); got {tuple(values.shape)}." + ) + if np.any((values < 0) | (values > 1)): + raise ValueError("occupations must contain only zero and one values.") + + rows = np.empty((values.shape[0], self.n_sites), dtype=np.int64) + for site, table in enumerate(self.code_to_occupations): + inverse = {occupation: code for code, occupation in enumerate(table)} + site_values = values[:, site, :] if self.spinful else values[:, site, None] + for row, occupation in enumerate(site_values): + try: + rows[row, site] = inverse[tuple(int(value) for value in occupation)] + except KeyError as exc: + raise ValueError( + f"occupation {tuple(occupation)!r} is unavailable at site {site}." + ) from exc + + backend = _mps_array_backend(occupations) + if to_numpy or backend not in {"torch", "cupy"}: + return rows + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + return torch.as_tensor(rows, dtype=torch.long, device=occupations.device) + import cupy as cp # pylint: disable=import-outside-toplevel + + return cp.asarray(rows, dtype=cp.int64) + + @dataclass class MpsBatchSampleResult: """Backend-native batched MPS samples. @@ -193,6 +657,7 @@ class MpsBatchSampleResult: configs Array-like object with shape ``(n_samples, L)``. With the native sampler this can be a NumPy array, Torch tensor, or CuPy array. + Symmray MPS samples use the backend of their underlying blocks. probs Born probabilities for ``configs`` with shape ``(n_samples,)``. Lx, Ly @@ -211,6 +676,7 @@ class MpsBatchSampleResult: Ly: int one_d_to_two_d: dict[int, tuple[int, int]] backend: str = "numpy" + configuration_encoding: FermionConfigurationEncoding | None = None def __len__(self): return int(self.configs.shape[0]) @@ -234,6 +700,7 @@ def to_numpy(self) -> "MpsBatchSampleResult": Ly=self.Ly, one_d_to_two_d=dict(self.one_d_to_two_d), backend="numpy", + configuration_encoding=self.configuration_encoding, ) def configs_1d(self) -> list[list[int]]: @@ -245,6 +712,15 @@ def configs_2d(self) -> list[np.ndarray]: """Return configurations as ``(Ly, Lx)`` NumPy grids.""" return self.to_sample_result().configs_2d + def occupations(self, *, to_numpy: bool = False): + """Decode fermionic physical codes with the attached configuration map.""" + if self.configuration_encoding is None: + raise ValueError( + "This batch has no fermion configuration encoding. Pass " + "fermion=... to MpsSampler.sample_batch(...)." + ) + return self.configuration_encoding.decode(self.configs, to_numpy=to_numpy) + def magnetizations(self, *, to_numpy: bool = False): """Per-sample magnetization ``(1 / L) * sum_i (1 - 2 * spin_i)``.""" backend = _mps_array_backend(self.configs) @@ -288,21 +764,45 @@ class MpsSampler: Mapping from 1D site index to (x, y) lattice coordinate. When omitted, a trivial single-row 1D layout ``{i: (i, 0)}`` inferred from the MPS length is used, so a plain 1D chain can be sampled without a 2D map. - backend : {"quimb", "native", "auto", "numpy", "torch", "cupy"} + backend : {"quimb", "native", "auto", "numpy", "torch", "cupy", "symmray"} Sampling implementation. ``"quimb"`` preserves the historical CPU - behavior. ``"native"`` accepts dense NumPy/Torch/CuPy tensors. - ``"auto"`` tries native sampling and falls back to ``"quimb"`` when - the MPS layout is unsupported. + behavior for dense MPSs. Symmray-backed MPSs are detected and use the + native block-sparse sampler rather than being densified. ``"native"`` + accepts dense NumPy/Torch/CuPy tensors and Symmray tensors, while + ``"symmray"`` requires a Symmray MPS explicitly. ``"auto"`` tries a + native sampler and falls back to ``"quimb"`` when the MPS layout is + unsupported. torch_compile : bool, default=False 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" + 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. + 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. + fermion : pepsy.tensors.Fermion, optional + Fermionic physical-space convention associated with this sampler. When + supplied, :meth:`sample_batch` attaches its symmetry-aware + configuration encoding by default, and the fermionic diagonal helpers + can omit the repeated ``fermion`` argument. A per-call ``fermion=`` + argument remains supported and takes precedence. Notes ----- - Native right environments are cached. Call :meth:`refresh` after changing - the source MPS; otherwise the sampler continues to represent its previous - tensor data. + Dense native right environments and Symmray right-canonical copies are + cached. The Symmray route retains the source physical-code map before + canonicalization, then samples by slicing one charge-aware local state and + absorbing it into a block-sparse boundary. Its batched route shares each + distinct sampled prefix, including when a physical charge sector has + degeneracy greater than one (for example spinful fermionic Z2 or U1). + Call :meth:`refresh` after changing the source MPS; otherwise the sampler + continues to represent its previous tensor data. """ def __init__( @@ -312,6 +812,9 @@ def __init__( *, backend: str | None = "quimb", torch_compile: bool = False, + prefix_strategy: str = "auto", + max_prefix_groups: int | None = 256, + fermion=None, ): if one_d_to_two_d is None: inferred_L = getattr(psi, "L", None) @@ -333,6 +836,17 @@ def __init__( if not isinstance(torch_compile, (bool, np.bool_)): raise TypeError("torch_compile must be a boolean.") self.torch_compile = bool(torch_compile) + 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)): + raise TypeError("max_prefix_groups must be a positive integer or None.") + if int(max_prefix_groups) < 1: + raise ValueError( + "max_prefix_groups must be a positive integer or None." + ) + max_prefix_groups = int(max_prefix_groups) + self.max_prefix_groups = max_prefix_groups + self.fermion = fermion self.resolved_backend = None self._source_psi = None self._native_arrays = None @@ -341,12 +855,29 @@ def __init__( self._evaluation_backend = None self._evaluation_arrays = None self._evaluation_site_ops = None + self._symmray_state = None + self._last_symmray_sampling_stats = None self._psi = None self._torch_compiled_sample_fns = {} self._torch_compile_disabled = False self.refresh(psi) + def _resolve_fermion(self, fermion): + """Use a call-specific Fermion or the sampler's bound convention.""" + fermion = self.fermion if fermion is None else fermion + if fermion is None: + raise TypeError( + "fermion is required. Pass fermion=... when constructing " + "MpsSampler or to this call." + ) + if not all(hasattr(fermion, name) for name in ("spinful", "symmetry")): + raise TypeError( + "fermion must be a pepsy.tensors.Fermion instance or expose " + "spinful and symmetry attributes." + ) + return fermion + def refresh(self, psi=None): """Refresh cached state from ``psi`` or the original source MPS. @@ -381,10 +912,34 @@ def refresh(self, psi=None): self._evaluation_backend = None self._evaluation_arrays = None self._evaluation_site_ops = None + self._symmray_state = None + self._last_symmray_sampling_stats = None self._psi = None self._torch_compiled_sample_fns.clear() self._torch_compile_disabled = False + source_backends = { + _mps_array_backend(psi[site].data) + for site in range(int(psi.L)) + } + if "symmray" in source_backends: + if source_backends != {"symmray"}: + raise ValueError( + "MPS tensors use mixed dense and Symmray array backends." + ) + if self.backend in {"quimb", "native", "auto", "symmray"}: + self._symmray_state = self._prepare_symmray_state(psi) + self.resolved_backend = "symmray" + return self + raise ValueError( + f"MpsSampler backend={self.backend!r} requested for a Symmray " + "MPS. Use backend='symmray', 'native', or 'auto'." + ) + if self.backend == "symmray": + raise ValueError( + "MpsSampler backend='symmray' requires Symmray tensor data." + ) + if self.backend != "quimb": try: native_backend, native_arrays = self._prepare_native_arrays(psi) @@ -411,6 +966,35 @@ def refresh(self, psi=None): ) return self + @property + def physical_code_maps(self): + """Per-site Symmray ``physical_code -> (charge, sector_offset)`` maps. + + The maps describe the source MPS physical basis, including charge + sectors pruned from the private canonical sampling copy. They are + ``None`` for a dense MPS, whose physical codes are already ordinary + positional indices. A fresh set of dictionaries is returned on each + access so callers cannot mutate sampler state. + """ + if self._symmray_state is None: + return None + return tuple( + dict(code_map) + for code_map in self._symmray_state["physical_code_maps"] + ) + + @property + def symmray_sampling_stats(self): + """Diagnostics from the most recent Symmray sampling call, if any. + + ``conditional_evaluations`` counts distinct local distributions built, + which is the useful work reduced by prefix sharing. ``None`` means the + sampler has not yet taken the Symmray route. + """ + if self._last_symmray_sampling_stats is None: + return None + return dict(self._last_symmray_sampling_stats) + def _get_evaluation_arrays(self): if self._native_arrays is not None: return self.resolved_backend, self._native_arrays @@ -482,6 +1066,1156 @@ def _prepare_native_arrays(self, psi): ) return backend, arrays + @staticmethod + def _prepare_symmray_state(psi): + """Cache a canonical Symmray MPS without altering its physical basis.""" + if getattr(psi, "cyclic", False): + raise ValueError("Symmray MPS sampling currently requires an open chain.") + + try: + import symmray as sr # pylint: disable=import-outside-toplevel + except ImportError as exc: # pragma: no cover - guarded by data type + raise ImportError( + "Symmray tensor data requires the optional 'symmray' package." + ) from exc + + source_data = tuple(psi[site].data for site in range(psi.L)) + block_backends = {str(data.backend) for data in source_data} + if len(block_backends) != 1: + raise ValueError( + "Symmray MPS tensors must use one common underlying block backend; " + f"got {sorted(block_backends)!r}." + ) + array_backend = next(iter(block_backends)) + if array_backend not in {"numpy", "torch", "cupy"}: + raise ValueError( + "Symmray MPS sampling currently supports NumPy, Torch, or CuPy " + f"blocks, not {array_backend!r}." + ) + + # Quimb's canonicalization runs Symmray QR/SVD blockwise. It can prune + # identically-zero physical charge sectors, so retain the source basis + # map below and never expose the canonical copy as the input state. + canonical = psi.right_canonicalize(normalize=True) + sites = [] + physical_code_maps = [] + for site in range(psi.L): + source_tensor = psi[site] + tensor = canonical[site] + source_phys_ind = psi.site_ind(site) + phys_ind = canonical.site_ind(site) + source_axis = source_tensor.inds.index(source_phys_ind) + phys_axis = tensor.inds.index(phys_ind) + source_index = source_tensor.data.indices[source_axis] + phys_index = tensor.data.indices[phys_axis] + + source_offsets = {} + source_code_metadata = [] + offset = 0 + for charge, size in source_index.chargemap.items(): + source_offsets[charge] = offset + for sector_offset in range(int(size)): + source_code_metadata.append((charge, sector_offset)) + offset += int(size) + + code_map = [] + for charge, size in phys_index.chargemap.items(): + try: + source_size = int(source_index.chargemap[charge]) + except KeyError as exc: + raise ValueError( + "Canonical Symmray MPS changed a physical charge sector " + f"at site {site}: {charge!r}." + ) from exc + size = int(size) + if size > source_size: + raise ValueError( + "Canonical Symmray MPS enlarged a physical charge sector " + f"at site {site}: {charge!r}." + ) + code_map.extend(range(source_offsets[charge], source_offsets[charge] + size)) + if len(code_map) != int(tensor.data.shape[phys_axis]): + raise ValueError( + "Could not reconstruct the physical code map for canonical " + f"Symmray site {site}." + ) + + left_axis = None + if site: + remaining_inds = tuple(ind for ind in tensor.inds if ind != phys_ind) + left_axis = remaining_inds.index(canonical.bond(site - 1, site)) + + # Physical selection on a fermionic Symmray array must happen + # before absorbing the left boundary: selecting it afterwards can + # lose the dummy-mode ordering needed for an odd physical leg. + # Cache these immutable local branch tensors once, so every + # sampled prefix shares both the slices and their block metadata. + locals_ = [] + local_left_charges = [] + nonempty_local_codes = [] + for local_code in range(int(tensor.data.shape[phys_axis])): + item = [slice(None)] * tensor.data.ndim + item[phys_axis] = local_code + local = tensor.data[tuple(item)] + locals_.append(local) + blocks = getattr(local, "blocks", None) + if isinstance(blocks, dict): + nonempty = bool(blocks) + else: + nonempty = True + if nonempty: + nonempty_local_codes.append(local_code) + + charges = None + if left_axis is not None: + if isinstance(blocks, dict): + charges = frozenset( + sector[left_axis] for sector in blocks + ) + else: + try: + charges = frozenset( + local.indices[left_axis].chargemap + ) + except (AttributeError, IndexError, TypeError): + charges = None + local_left_charges.append(charges) + sites.append( + { + "data": tensor.data, + "phys_axis": phys_axis, + "left_axis": left_axis, + "codes": tuple(code_map), + "code_metadata": tuple( + source_code_metadata[code] for code in code_map + ), + "code_to_local": {code: local for local, code in enumerate(code_map)}, + "locals": tuple(locals_), + "nonempty_local_codes": tuple(nonempty_local_codes), + "local_left_charges": tuple(local_left_charges), + } + ) + physical_code_maps.append(dict(enumerate(source_code_metadata))) + + template = source_data[0].get_any_array() + # Keep the import local so Symmray remains optional for ordinary MPSs. + return { + "sr": sr, + "mps": canonical, + "sites": tuple(sites), + "physical_code_maps": tuple(physical_code_maps), + "array_backend": array_backend, + "template": template, + } + + @staticmethod + def _symmray_weight(value, state): + """Return ``||value||**2`` without converting an array to dense. + + Symmray returns an ordinary scalar/block when a contraction has no + remaining charge structure. This is common for the final local branch + of fermionic Z2 and U1 MPSs with degenerate physical sectors. Such a + block is already a selected sector rather than a densified state. + """ + blocks = getattr(value, "blocks", None) + if isinstance(blocks, dict) and not blocks: + return 0.0 + if np.isscalar(value): + return abs(value) ** 2 + backend = _mps_array_backend(value) + if backend == "torch": + return (value.conj() * value).sum().real + if backend == "cupy": + return (value.conj() * value).sum().real + if backend == "numpy": + return np.sum(np.abs(value) ** 2).real + return state["sr"].linalg.norm(value) ** 2 + + @staticmethod + def _symmray_scalar(value): + """Extract a scalar after a fully contracted Symmray MPS branch.""" + blocks = getattr(value, "blocks", None) + if isinstance(blocks, dict) and not blocks: + return 0.0 + if hasattr(value, "get_scalar_element"): + return value.phase_sync().get_scalar_element() + return value + + @staticmethod + def _symmray_distribution(weights, state): + """Normalize scalar weights on the backend of the Symmray blocks.""" + backend = state["array_backend"] + template = state["template"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + dtype = template.real.dtype + values = torch.stack([ + torch.as_tensor(weight, dtype=dtype, device=template.device).real + for weight in weights + ]) + values = values.clamp_min(0.0) + total = values.sum() + if not bool(torch.isfinite(total).detach().cpu().item()) or not bool( + (total > 0).detach().cpu().item() + ): + raise ValueError("MPS has a zero or non-finite conditional norm.") + return values / total + + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + values = cp.stack([ + cp.asarray(weight, dtype=template.real.dtype).real + for weight in weights + ]) + values = cp.maximum(values, 0.0) + total = values.sum() + if not bool(cp.isfinite(total).item()) or not bool((total > 0).item()): + raise ValueError("MPS has a zero or non-finite conditional norm.") + return values / total + + values = np.asarray(weights, dtype=np.asarray(template).real.dtype).real + values = np.maximum(values, 0.0) + total = values.sum() + if not np.isfinite(total) or total <= 0.0: + raise ValueError("MPS has a zero or non-finite conditional norm.") + return values / total + + @staticmethod + def _symmray_draw(probs, state, rng): + """Draw one local code, leaving probabilities on their native backend.""" + backend = state["array_backend"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + choice = int(torch.multinomial(probs, 1, generator=rng).reshape(()).item()) + elif backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + cdf = cp.cumsum(probs) + choice = int(cp.sum(rng.random() > cdf).item()) + choice = min(choice, int(probs.shape[0]) - 1) + else: + choice = int(rng.choice(len(probs), p=probs)) + return choice, probs[choice] + + @staticmethod + def _symmray_draw_many(probs, n_draws, state, rng): + """Draw a prefix group and return its local choices on the host. + + The probability vector and random-number generation stay on the + Symmray block backend. Only the integer decisions cross to Python so + that one block-sparse boundary can be retained for every distinct + prefix, rather than once per requested shot. + """ + n_draws = int(n_draws) + if n_draws < 1: # pragma: no cover - internal guard + raise ValueError("n_draws must be positive.") + if n_draws == 1: + choice, _ = MpsSampler._symmray_draw(probs, state, rng) + return np.asarray((choice,), dtype=np.int64) + + backend = state["array_backend"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + choices = torch.multinomial( + probs, + n_draws, + replacement=True, + generator=rng, + ) + return choices.detach().cpu().numpy().astype(np.int64, copy=False) + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + cdf = cp.cumsum(probs) + 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( + rng.choice(len(probs), size=n_draws, p=probs), + dtype=np.int64, + ) + + @staticmethod + def _symmray_rng(state, seed): + """Create the random generator associated with the block backend.""" + backend = state["array_backend"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + if seed is None: + return None + generator = torch.Generator(device=state["template"].device) + generator.manual_seed(int(seed)) + return generator + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + return cp.random.default_rng(seed) + return np.random.default_rng(seed) + + @staticmethod + def _symmray_sqrt(value, state): + backend = state["array_backend"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + return torch.sqrt(value) + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + return cp.sqrt(value) + return np.sqrt(value) + + @staticmethod + def _symmray_positive(value, state): + """Test a scalar branch norm without moving tensor data to dense CPU.""" + backend = state["array_backend"] + if backend == "torch": + positive = value > 0 + return bool( + positive.detach().cpu().item() + if hasattr(positive, "detach") + else positive + ) + if backend == "cupy": + positive = value > 0 + return bool(positive.item() if hasattr(positive, "item") else positive) + return bool(value > 0) + + @staticmethod + def _symmray_one(state, *, complex_value=False): + backend = state["array_backend"] + template = state["template"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + dtype = template.dtype if complex_value else template.real.dtype + return torch.ones((), dtype=dtype, device=template.device) + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + dtype = template.dtype if complex_value else template.real.dtype + return cp.ones((), dtype=dtype) + dtype = np.asarray(template).dtype + if not complex_value: + dtype = np.asarray(template).real.dtype + return np.ones((), dtype=dtype) + + @classmethod + def _symmray_sample_one(cls, state, rng): + """Draw one configuration by slice-and-absorb on a canonical MPS.""" + boundary = None + config = [] + probability = cls._symmray_one(state) + for site, site_state in enumerate(state["sites"]): + site_state, local_codes, candidates, weights = cls._symmray_candidates( + state, + site, + boundary, + ) + probs = cls._symmray_distribution(weights, state) + choice_index, choice_prob = cls._symmray_draw(probs, state, rng) + local_code = local_codes[choice_index] + config.append(site_state["codes"][local_code]) + probability = probability * choice_prob + if site < len(state["sites"]) - 1: + boundary = candidates[choice_index] / cls._symmray_sqrt( + weights[choice_index], + state, + ) + return config, probability + + @staticmethod + def _symmray_boundary_charges(boundary): + """Return the possible outgoing charge labels of a prefix boundary.""" + if boundary is None: + return None + blocks = getattr(boundary, "blocks", None) + if isinstance(blocks, dict): + return frozenset(sector[0] for sector in blocks) + try: + return frozenset(boundary.indices[0].chargemap) + except (AttributeError, IndexError, TypeError): + return None + + @classmethod + def _symmray_candidate_codes(cls, site_state, boundary): + """Skip cached local branches incompatible with the prefix charge.""" + local_codes = site_state["nonempty_local_codes"] + if boundary is None or site_state["left_axis"] is None: + return local_codes + boundary_charges = cls._symmray_boundary_charges(boundary) + if not boundary_charges: + return local_codes + return tuple( + local_code + for local_code in local_codes + if ( + site_state["local_left_charges"][local_code] is None + or boundary_charges + & site_state["local_left_charges"][local_code] + ) + ) + + @classmethod + def _symmray_candidates(cls, state, site, boundary): + """Build one charge-pruned conditional from cached local branches.""" + site_state = state["sites"][site] + local_codes = cls._symmray_candidate_codes(site_state, boundary) + candidates = [] + weights = [] + for local_code in local_codes: + local = site_state["locals"][local_code] + if boundary is None: + candidate = local + else: + candidate = state["sr"].tensordot( + boundary, + local, + axes=((0,), (site_state["left_axis"],)), + ) + candidates.append(candidate) + weights.append(cls._symmray_weight(candidate, state)) + return site_state, local_codes, candidates, weights + + @staticmethod + def _symmray_note_candidates(stats, site_state, local_codes): + """Record the sparse branch work avoided by cache/pruning.""" + if stats is None: + return + stats["candidate_contractions"] += len(local_codes) + stats["static_pruned_branches"] += ( + len(site_state["codes"]) - len(site_state["nonempty_local_codes"]) + ) + stats["charge_pruned_branches"] += ( + len(site_state["nonempty_local_codes"]) - len(local_codes) + ) + + @staticmethod + def _symmray_stack(values, state, *, integer=False, complex_value=False): + """Stack scalar results using the backend of the Symmray blocks.""" + backend = state["array_backend"] + template = state["template"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + if integer: + return torch.as_tensor(values, dtype=torch.long, device=template.device) + dtype = template.dtype if complex_value else template.real.dtype + return torch.stack([ + torch.as_tensor(value, dtype=dtype, device=template.device) + for value in values + ]) + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + dtype = ( + cp.int64 + if integer + else (template.dtype if complex_value else template.real.dtype) + ) + return cp.asarray(values, dtype=dtype) + dtype = ( + np.int64 + if integer + else ( + np.asarray(template).dtype + if complex_value + else np.asarray(template).real.dtype + ) + ) + return np.asarray(values, dtype=dtype) + + @classmethod + def _symmray_sample_from_prefix( + cls, + state, + rng, + config, + probability, + boundary, + start_site, + stats, + ): + """Complete one shot from an already sampled normalized prefix.""" + for site in range(start_site, len(state["sites"])): + site_state, local_codes, candidates, weights = cls._symmray_candidates( + state, + site, + boundary, + ) + stats["conditional_evaluations"] += 1 + cls._symmray_note_candidates(stats, site_state, local_codes) + probs = cls._symmray_distribution(weights, state) + choice_index, choice_prob = cls._symmray_draw(probs, state, rng) + local_code = local_codes[choice_index] + config.append(site_state["codes"][local_code]) + probability = probability * choice_prob + if site < len(state["sites"]) - 1: + if not cls._symmray_positive(weights[choice_index], state): + raise ValueError( + "MPS sampler selected a zero-norm conditional branch." + ) + boundary = candidates[choice_index] / cls._symmray_sqrt( + weights[choice_index], + state, + ) + return config, probability + + @classmethod + def _symmray_sample_arrays_serial(cls, state, n_samples, seed, *, to_numpy): + """Sample independently, using constant block-sparse boundary memory.""" + rng = cls._symmray_rng(state, seed) + configs = [] + probabilities = [] + stats = { + "strategy": "serial", + "n_samples": int(n_samples), + "conditional_evaluations": 0, + "candidate_contractions": 0, + "static_pruned_branches": 0, + "charge_pruned_branches": 0, + "cached_local_slices": True, + "max_active_prefix_groups": 1, + "serial_fallback": False, + "adaptive_serial_fallback": False, + } + for _ in range(int(n_samples)): + config, probability = cls._symmray_sample_from_prefix( + state, + rng, + [], + cls._symmray_one(state), + None, + 0, + stats, + ) + configs.append(config) + probabilities.append(probability) + return cls._symmray_finalize_samples( + configs, + probabilities, + state, + stats, + to_numpy=to_numpy, + ) + + @staticmethod + def _symmray_boundary_storage_cost(boundary): + """Estimate block-resident boundary storage without densifying it.""" + blocks = getattr(boundary, "blocks", None) + if not isinstance(blocks, dict): + return 1 + cost = 0 + for block in blocks.values(): + size = getattr(block, "size", None) + if callable(size): + size = size() + if size is None or not np.isscalar(size): + size = int(np.prod(getattr(block, "shape", (1,)))) + cost += int(size) + return max(cost, 1) + + @classmethod + def _symmray_select_prefix_groups( + cls, + branches, + *, + max_prefix_groups, + adaptive, + ): + """Keep prefix boundaries only while they amortize their storage. + + A group with one walker cannot share a future conditional, so the + auto strategy finishes it serially immediately. For the remaining + groups, ``max_prefix_groups`` is a hard count cap and also defines a + per-level block-storage budget relative to the median boundary size. + This avoids treating a large multi-sector boundary as equal to a tiny + one-sector boundary. + """ + if not branches: + return (), (), None + + costs = [cls._symmray_boundary_storage_cost(branch[1]) for branch in branches] + if max_prefix_groups is None: + count_limit = None + storage_budget = None + else: + count_limit = int(max_prefix_groups) + baseline = int(np.median(costs)) + storage_budget = max(count_limit * max(baseline, 1), 1) + + kept_ids = set() + used_storage = 0 + # Retain the groups with the greatest prospective reuse first. Ties + # favor smaller block boundaries, then preserve sample order. + ranked = sorted( + range(len(branches)), + key=lambda index: ( + -len(branches[index][0]), + costs[index], + int(branches[index][0][0]), + ), + ) + for index in ranked: + positions = branches[index][0] + if adaptive and len(positions) == 1: + continue + if count_limit is not None and len(kept_ids) >= count_limit: + continue + if ( + storage_budget is not None + and kept_ids + and used_storage + costs[index] > storage_budget + ): + continue + kept_ids.add(index) + used_storage += costs[index] + + kept = tuple(branch for index, branch in enumerate(branches) if index in kept_ids) + dropped = tuple(branch for index, branch in enumerate(branches) if index not in kept_ids) + return kept, dropped, storage_budget + + @classmethod + def _symmray_sample_arrays_prefix( + cls, + state, + n_samples, + seed, + *, + max_prefix_groups, + adaptive, + to_numpy, + ): + """Share boundaries between prefixes, with an optional memory bound.""" + rng = cls._symmray_rng(state, seed) + n_samples = int(n_samples) + configs = np.empty((n_samples, len(state["sites"])), dtype=np.int64) + probabilities = [None] * n_samples + positions = np.arange(n_samples, dtype=np.int64) + stats = { + "strategy": "prefix", + "n_samples": n_samples, + "conditional_evaluations": 0, + "candidate_contractions": 0, + "static_pruned_branches": 0, + "charge_pruned_branches": 0, + "cached_local_slices": True, + "max_active_prefix_groups": 1, + "serial_fallback": False, + "adaptive_serial_fallback": False, + "max_prefix_storage_budget": None, + } + # A group stores the shared selected prefix, its normalized boundary, + # and its probability. Only current-depth groups are retained. + groups = [(positions, None, cls._symmray_one(state))] + + for site in range(len(state["sites"])): + stats["max_active_prefix_groups"] = max( + stats["max_active_prefix_groups"], + len(groups), + ) + + is_final_site = site == len(state["sites"]) - 1 + branches = [] + for group_positions, boundary, prefix_probability in groups: + site_state, local_codes, candidates, weights = cls._symmray_candidates( + state, + site, + boundary, + ) + stats["conditional_evaluations"] += 1 + cls._symmray_note_candidates(stats, site_state, local_codes) + probs = cls._symmray_distribution(weights, state) + choices = cls._symmray_draw_many( + probs, + len(group_positions), + state, + rng, + ) + for choice_index, local_code in enumerate(local_codes): + selected = group_positions[choices == choice_index] + if not len(selected): + continue + code = site_state["codes"][local_code] + configs[selected, site] = code + selected_probability = prefix_probability * probs[choice_index] + if is_final_site: + for sample in selected: + probabilities[int(sample)] = selected_probability + continue + if not cls._symmray_positive(weights[choice_index], state): + raise ValueError( + "MPS sampler selected a zero-norm conditional branch." + ) + next_boundary = candidates[choice_index] / cls._symmray_sqrt( + weights[choice_index], + state, + ) + branches.append((selected, next_boundary, selected_probability)) + + if is_final_site: + groups = () + continue + + groups, serial_branches, storage_budget = cls._symmray_select_prefix_groups( + branches, + max_prefix_groups=max_prefix_groups, + adaptive=adaptive, + ) + if storage_budget is not None: + previous_budget = stats["max_prefix_storage_budget"] + stats["max_prefix_storage_budget"] = ( + storage_budget + if previous_budget is None + else max(previous_budget, storage_budget) + ) + if serial_branches: + stats["serial_fallback"] = True + for selected, next_boundary, selected_probability in serial_branches: + if adaptive and len(selected) == 1: + stats["adaptive_serial_fallback"] = True + for sample in selected: + config, probability = cls._symmray_sample_from_prefix( + state, + rng, + configs[int(sample), : site + 1].tolist(), + selected_probability, + next_boundary, + site + 1, + stats, + ) + configs[int(sample)] = config + probabilities[int(sample)] = probability + + if any(probability is None for probability in probabilities): # pragma: no cover + raise RuntimeError("Symmray prefix sampler did not assign every shot.") + return cls._symmray_finalize_samples( + configs, + probabilities, + state, + stats, + to_numpy=to_numpy, + ) + + @classmethod + def _symmray_finalize_samples( + cls, + configs, + probabilities, + state, + stats, + *, + to_numpy, + ): + configs = cls._symmray_stack(configs, state, integer=True) + probabilities = cls._symmray_stack(probabilities, state) + if to_numpy: + configs = _backend_array_to_numpy(configs) + probabilities = _backend_array_to_numpy(probabilities) + return configs, probabilities, stats + + @classmethod + def _symmray_sample_arrays( + cls, + state, + n_samples, + seed, + *, + strategy, + max_prefix_groups, + to_numpy, + ): + if strategy == "serial": + return cls._symmray_sample_arrays_serial( + state, + n_samples, + seed, + to_numpy=to_numpy, + ) + return cls._symmray_sample_arrays_prefix( + state, + n_samples, + seed, + max_prefix_groups=max_prefix_groups, + adaptive=(strategy == "auto"), + to_numpy=to_numpy, + ) + + @staticmethod + def _symmray_config_rows(configs, *, L): + """Validate discrete configurations for Symmray MPS evaluation.""" + configs = np.asarray(_backend_array_to_numpy(configs), dtype=np.int64) + if configs.ndim != 2 or configs.shape[1] != int(L): + raise ValueError( + f"configs must have shape (batch, L={int(L)}); " + f"got {tuple(configs.shape)}." + ) + return configs + + @classmethod + def _symmray_amplitude_one(cls, state, config): + """Contract one selected configuration without densifying the MPS.""" + boundary = None + for site, code in enumerate(config): + site_state = state["sites"][site] + try: + local_code = site_state["code_to_local"][int(code)] + except KeyError: + return cls._symmray_one(state, complex_value=True) * 0.0 + local = site_state["locals"][local_code] + if boundary is None: + boundary = local + else: + boundary = state["sr"].tensordot( + boundary, + local, + axes=((0,), (site_state["left_axis"],)), + ) + return cls._symmray_scalar(boundary) + + @classmethod + def _symmray_probability_one(cls, state, config): + """Evaluate one Born probability with canonical slice-and-absorb.""" + boundary = None + probability = cls._symmray_one(state) + for site, code in enumerate(config): + site_state = state["sites"][site] + try: + choice = site_state["code_to_local"][int(code)] + except KeyError: + return cls._symmray_one(state) * 0.0 + + site_state, local_codes, candidates, weights = cls._symmray_candidates( + state, + site, + boundary, + ) + try: + choice_index = local_codes.index(choice) + except ValueError: + return cls._symmray_one(state) * 0.0 + probs = cls._symmray_distribution(weights, state) + probability = probability * probs[choice_index] + if site < len(state["sites"]) - 1: + if not cls._symmray_positive(weights[choice_index], state): + return cls._symmray_one(state) * 0.0 + boundary = candidates[choice_index] / cls._symmray_sqrt( + weights[choice_index], + state, + ) + return probability + + def _symmray_amplitudes(self, configs): + state = self._require_symmray_state() + rows = self._symmray_config_rows(configs, L=len(state["sites"])) + values = [self._symmray_amplitude_one(state, row) for row in rows] + return self._symmray_stack(values, state, complex_value=True) + + def _symmray_probabilities(self, configs): + state = self._require_symmray_state() + rows = self._symmray_config_rows(configs, L=len(state["sites"])) + values = [self._symmray_probability_one(state, row) for row in rows] + return self._symmray_stack(values, state) + + def _require_symmray_state(self): + if self._symmray_state is None: # pragma: no cover - internal guard + raise RuntimeError("Symmray sampler state has not been initialized.") + return self._symmray_state + + def fermion_configuration_encoding(self, fermion=None) -> FermionConfigurationEncoding: + """Return the physical-code/occupation contract for a fermionic MPS. + + This is the explicit bridge from MPS Born samples to a VMC walker or + local-estimator configuration. It is intentionally derived from the + source physical-sector maps retained by the Symmray sampler, rather + than assuming a dense-basis or VMC-specific code convention. + """ + fermion = self._resolve_fermion(fermion) + state = self._require_symmray_state() + state_symmetry = str(state["sites"][0]["data"].symmetry).upper() + symmetry = str(fermion.symmetry).upper() + if state_symmetry != symmetry: + raise ValueError( + "The Fermion symmetry must match the sampled Symmray MPS; " + f"got {symmetry!r} for a {state_symmetry!r} state." + ) + + expected_dim = 4 if bool(fermion.spinful) else 2 + tables = [] + for site, code_map in enumerate(state["physical_code_maps"]): + if tuple(code_map) != tuple(range(len(code_map))): + raise ValueError( + "Symmray physical codes must be contiguous at site " + f"{site}; got {sorted(code_map)!r}." + ) + if len(code_map) != expected_dim: + raise ValueError( + "The sampled Symmray MPS physical dimension is incompatible " + f"with this {'spinful' if fermion.spinful else 'spinless'} " + "Fermion." + ) + tables.append(tuple( + _fermion_symmray_occupations(charge, offset, fermion) + for charge, offset in code_map.values() + )) + return FermionConfigurationEncoding( + symmetry=symmetry, + spinful=bool(fermion.spinful), + code_to_occupations=tuple(tables), + ) + + @staticmethod + def _normalize_fermion_diagonal_observable(observable): + aliases = { + "n": "occupation", + "number": "occupation", + "occupation": "occupation", + "density": "occupation", + "total_charge": "total_charge", + "total_number": "total_charge", + "total_occupation": "total_charge", + "doublon": "doublon", + "double": "doublon", + "double_occupancy": "doublon", + "density_correlation": "density_correlation", + "density_correlator": "density_correlation", + "ninj": "density_correlation", + "n_i_n_j": "density_correlation", + } + key = str(observable).strip().lower().replace("-", "_") + try: + return aliases[key] + except KeyError as exc: + allowed = ", ".join(sorted(set(aliases.values()))) + raise ValueError( + "Unknown fermion diagonal observable " + f"{observable!r}. Expected one of: {allowed}." + ) from exc + + def _normalize_fermion_sites(self, sites): + if sites is None: + return tuple(range(self._L)) + if isinstance(sites, (int, np.integer)): + sites = (int(sites),) + else: + try: + sites = tuple(int(site) for site in sites) + except TypeError as exc: + raise TypeError("sites must be an integer or an iterable of integers.") from exc + if not sites: + raise ValueError("sites must contain at least one physical site.") + if len(set(sites)) != len(sites): + raise ValueError("sites must not contain duplicates.") + invalid = [site for site in sites if not 0 <= site < self._L] + if invalid: + raise ValueError( + f"sites contain values outside the MPS range 0..{self._L - 1}: " + f"{invalid!r}." + ) + return sites + + def _normalize_fermion_pairs(self, pairs): + if pairs is None: + raise ValueError("density_correlation requires pairs=((i, j), ...).") + if ( + isinstance(pairs, tuple) + and len(pairs) == 2 + and all(isinstance(site, (int, np.integer)) for site in pairs) + ): + pairs = (pairs,) + try: + pairs = tuple(tuple(int(site) for site in pair) for pair in pairs) + except TypeError as exc: + raise TypeError("pairs must be an (i, j) pair or iterable of pairs.") from exc + if not pairs: + raise ValueError("pairs must contain at least one physical pair.") + for pair in pairs: + if len(pair) != 2: + raise ValueError("Each density-correlation pair must have two sites.") + if pair[0] == pair[1]: + raise ValueError("Density-correlation pairs must contain distinct sites.") + self._normalize_fermion_sites(pair) + return pairs + + @staticmethod + def _fermion_symmray_diagonal_values(charge, offset, fermion): + """Decode standard Fermion occupations from one Symmray sector code.""" + occupations = _fermion_symmray_occupations(charge, offset, fermion) + if not bool(fermion.spinful): + return float(occupations[0]), 0.0 + return float(sum(occupations)), float(occupations == (1, 1)) + + def _fermion_diagonal_tables(self, fermion): + """Build physical-code lookup tables for occupation and doublon values.""" + fermion = self._resolve_fermion(fermion) + + if self._symmray_state is None: + try: + number = np.real( + np.diag(_backend_array_to_numpy(fermion.dense_operator("number"))) + ).astype(float, copy=False) + double = ( + np.real( + np.diag(_backend_array_to_numpy(fermion.dense_operator("double"))) + ).astype(float, copy=False) + if bool(fermion.spinful) + else np.zeros_like(number, dtype=float) + ) + except AttributeError as exc: + raise TypeError( + "fermion must expose dense_operator(name) for dense MPS sampling." + ) from exc + return tuple((number, double) for _ in range(self._L)) + + state_symmetry = str(self._symmray_state["sites"][0]["data"].symmetry) + fermion_symmetry = str(fermion.symmetry) + if state_symmetry != fermion_symmetry: + raise ValueError( + "The Fermion symmetry must match the sampled Symmray MPS; " + f"got {fermion_symmetry!r} for a {state_symmetry!r} state." + ) + expected_dim = 4 if bool(fermion.spinful) else 2 + tables = [] + for code_map in self._symmray_state["physical_code_maps"]: + if len(code_map) != expected_dim: + raise ValueError( + "The sampled Symmray MPS physical dimension is incompatible " + f"with this {'spinful' if fermion.spinful else 'spinless'} Fermion." + ) + number = np.empty(len(code_map), dtype=float) + double = np.empty(len(code_map), dtype=float) + for code, (charge, offset) in code_map.items(): + number[code], double[code] = self._fermion_symmray_diagonal_values( + charge, + offset, + fermion, + ) + tables.append((number, double)) + return tuple(tables) + + def fermion_diagonal_values( + self, + configs, + fermion=None, + observable=None, + *, + sites=None, + pairs=None, + ): + """Evaluate a diagonal fermionic observable on configurations. + + This supports spinful and spinless :class:`pepsy.tensors.Fermion` + conventions. ``"occupation"`` returns the mean local occupation on + ``sites`` (all sites by default), ``"total_charge"`` its sum, + ``"doublon"`` the mean ``n_up n_down`` on ``sites``, and + ``"density_correlation"`` the mean ``n_i n_j`` across ``pairs``. + Symmray physical codes are decoded from the source charge map, so the + spinful Z2 even-sector ordering remains correct. + """ + if observable is None: + observable, fermion = fermion, None + if observable is None: + raise TypeError("observable is required.") + fermion = self._resolve_fermion(fermion) + observable = self._normalize_fermion_diagonal_observable(observable) + configs = self._symmray_config_rows(configs, L=self._L) + tables = self._fermion_diagonal_tables(fermion) + occupations = np.empty(configs.shape, dtype=float) + doublons = np.empty(configs.shape, dtype=float) + for site, (number, double) in enumerate(tables): + codes = configs[:, site] + invalid = (codes < 0) | (codes >= len(number)) + if np.any(invalid): + raise ValueError( + f"configs contain invalid physical index for site {site}." + ) + occupations[:, site] = number[codes] + doublons[:, site] = double[codes] + + if observable == "density_correlation": + if sites is not None: + raise ValueError("density_correlation uses pairs= rather than sites=.") + pairs = self._normalize_fermion_pairs(pairs) + return np.mean( + [occupations[:, i] * occupations[:, j] for i, j in pairs], + axis=0, + ) + + if pairs is not None: + raise ValueError(f"{observable} does not accept pairs=.") + sites = self._normalize_fermion_sites(sites) + if observable == "occupation": + return occupations[:, sites].mean(axis=1) + if observable == "total_charge": + return occupations[:, sites].sum(axis=1) + if not bool(fermion.spinful): + raise ValueError("doublon requires a spinful Fermion.") + return doublons[:, sites].mean(axis=1) + + def estimate_fermion_diagonal( + self, + fermion=None, + observable=None, + n_samples: int = 1024, + seed: int | None = None, + *, + sites=None, + pairs=None, + ) -> MpsDiagonalEstimate: + """Estimate a diagonal fermionic observable from Born samples. + + The returned uncertainty uses the unbiased sample variance. This + sampler deliberately covers diagonal observables only; + hopping, pairing, and spin-flip observables require a fermionic local + estimator based on amplitude ratios. + """ + if observable is None: + observable, fermion = fermion, None + if observable is None: + raise TypeError("observable is required.") + fermion = self._resolve_fermion(fermion) + configs, _ = self.sample_arrays( + n_samples, + seed=seed, + to_numpy=True, + ) + observable = self._normalize_fermion_diagonal_observable(observable) + values = self.fermion_diagonal_values( + configs, + fermion, + observable, + sites=sites, + pairs=pairs, + ) + n_samples = int(values.size) + standard_error = ( + float(np.std(values, ddof=1) / math.sqrt(n_samples)) + if n_samples > 1 + else math.nan + ) + return MpsDiagonalEstimate( + mean=float(np.mean(values)), + standard_error=standard_error, + n_samples=n_samples, + observable=observable, + sites=( + () + if observable == "density_correlation" + else self._normalize_fermion_sites(sites) + ), + pairs=( + self._normalize_fermion_pairs(pairs) + if observable == "density_correlation" + else () + ), + ) + @staticmethod def _torch_site_ops(arrays): import torch # pylint: disable=import-outside-toplevel @@ -891,8 +2625,15 @@ def amplitudes(self, configs, *, to_numpy: bool = True): ``configs`` should have shape ``(batch, L)``. Dense NumPy, Torch, and CuPy MPS tensors are contracted in one batched backend-native pass. - Set ``to_numpy=False`` to keep Torch/CuPy outputs on their device. + Symmray MPSs use block-sparse contractions on the underlying NumPy, + Torch, or CuPy backend. Set ``to_numpy=False`` to keep Torch/CuPy + outputs on their device. """ + if self._symmray_state is not None: + out = self._symmray_amplitudes(configs) + backend = self._symmray_state["array_backend"] + return self._to_numpy_backend_array(out, backend) if to_numpy else out + backend, site_data = self._get_evaluation_site_ops() if backend == "torch": out = self._torch_amplitudes(site_data, configs, L=self._L) @@ -909,9 +2650,16 @@ def probabilities(self, configs, *, to_numpy: bool = True): """Return normalized Born probabilities for batched ``configs``. This follows the same conditional-probability sweep as sampling, but - with user-supplied physical indices. It avoids looping over - configurations and runs on Torch/CuPy when the MPS tensors do. + with user-supplied physical indices. Dense MPSs avoid looping over + configurations and run on Torch/CuPy when the tensors do. Symmray MPSs + use charge-aware block-sparse conditionals without densifying state + tensors. """ + if self._symmray_state is not None: + out = self._symmray_probabilities(configs) + backend = self._symmray_state["array_backend"] + return self._to_numpy_backend_array(out, backend) if to_numpy else out + backend, site_data = self._get_evaluation_site_ops() if backend == "torch": out = self._torch_probabilities(site_data, configs, L=self._L) @@ -934,8 +2682,10 @@ def sample_arrays( ): """Draw samples and return raw ``(configs, probs)`` arrays. - With ``backend="native"``, this returns backend-native arrays by - default: Torch tensors stay on Torch and CuPy arrays stay on CuPy. + With ``backend="native"`` or ``backend="symmray"``, this returns + backend-native arrays by default: Torch tensors stay on Torch and CuPy + arrays stay on CuPy. Symmray inputs retain block-sparse state tensors; + only the returned discrete configurations and probabilities are dense. The returned ``configs`` have shape ``(n_samples, L)`` and ``probs`` has shape ``(n_samples,)``. Set ``to_numpy=True`` to force CPU NumPy arrays, matching the legacy :meth:`sample` result conversion. @@ -946,6 +2696,34 @@ def sample_arrays( raise ValueError("n_samples must be a positive integer.") if not isinstance(track_grad, (bool, np.bool_)): raise TypeError("track_grad must be a boolean.") + if self._symmray_state is not None: + if ( + self._symmray_state["array_backend"] == "torch" + and not track_grad + ): + 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, + ) + 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, + ) + self._last_symmray_sampling_stats = stats + return configs, probs if self._native_arrays is not None: return self._native_sample_arrays( int(n_samples), @@ -968,6 +2746,7 @@ def sample_batch( *, to_numpy: bool = False, track_grad: bool = False, + fermion=None, ) -> MpsBatchSampleResult: """Draw samples and return a named batched result. @@ -975,7 +2754,10 @@ def sample_batch( ``backend="native"`` and ``to_numpy=False``, Torch/CuPy arrays stay on their current device. Use :meth:`sample_arrays` when tuple unpacking is more convenient, or :meth:`sample` when the legacy Python-list/grid - result is needed. + result is needed. Pass ``fermion=...`` for a Symmray fermionic MPS to + attach a :class:`FermionConfigurationEncoding`; downstream code can + then call :meth:`MpsBatchSampleResult.occupations` without guessing the + physical-code convention. """ configs, probs = self.sample_arrays( n_samples, @@ -985,8 +2767,12 @@ def sample_batch( ) backend = ( "numpy" - if to_numpy or self._native_arrays is None - else self.resolved_backend + if to_numpy + else ( + self._symmray_state["array_backend"] + if self._symmray_state is not None + else (self.resolved_backend if self._native_arrays is not None else "numpy") + ) ) return MpsBatchSampleResult( configs=configs, @@ -995,6 +2781,13 @@ def sample_batch( Ly=self.Ly, one_d_to_two_d=dict(self.one_d_to_two_d), backend=backend, + configuration_encoding=( + self.fermion_configuration_encoding( + self.fermion if fermion is None else fermion + ) + if (self.fermion is not None or fermion is not None) + else None + ), ) def sample( @@ -1006,9 +2799,11 @@ def sample( ) -> MpsSampleResult: """Draw ``n_samples`` configurations from the MPS. - The native backend uses batched conditional contractions on the MPS - tensor device. The quimb backend uses ``MatrixProductState.sample()``, - which internally right-canonicalizes the MPS and sweeps left-to-right. + The dense native backend uses batched conditional contractions on the + MPS tensor device. The Symmray backend caches a right-canonical copy + and sweeps block-sparse physical slices left-to-right. The quimb + backend uses ``MatrixProductState.sample()``, which internally + right-canonicalizes the MPS and sweeps left-to-right. Returns ------- @@ -1134,14 +2929,39 @@ class PepsBpSampler: proposal probability ``omega(x)``, contracts the projected PEPS amplitude ``p(x)``, and returns the ingredients for the PEPS norm estimator ``E_q[|p(x)|^2 / omega(x)]``. + + Quimb's public D2BP sampler currently samples binary output indices. For a + four-state spinful PEPS this class transparently samples two binary + occupation legs per site and maps them back to the requested local + fermion encoding. Symmray inputs are densified only in the private BP + proposal copy; the original network remains block-sparse. """ - def __init__(self, tn, *, optimizer=None, sample_kwargs: dict[str, Any] | None = None): - self.tn = tn - self.Lx = tn.Lx - self.Ly = tn.Ly + def __init__( + self, + tn, + *, + optimizer=None, + sample_kwargs: dict[str, Any] | None = None, + encoding=None, + site_order=None, + ): + self.tn = getattr(tn, "tn", tn) + self.Lx = self.tn.Lx + self.Ly = self.tn.Ly self.optimizer = optimizer self.sample_kwargs = dict(sample_kwargs or {}) + self.encoding = encoding + ( + self._bp_tn, + self._split_inds, + self.site_order, + self._code_order, + ) = _prepare_bp_binary_network( + self.tn, + site_order=site_order, + encoding=encoding, + ) @staticmethod def mantissa_exponent10(w: float) -> tuple[float, int]: @@ -1171,6 +2991,33 @@ def _get_optimizer(self): return self.optimizer def _config_list(self, config: dict[str, Any]) -> list[int]: + if self._split_inds: + out = [] + code_order = self._code_order + # Invert the (up, down) -> physical-code map built above. + code_from_bits = { + (0, 0): code_order[0], + (0, 1): code_order[1], + (1, 0): code_order[2], + (1, 1): code_order[3], + } + for site in self.site_order: + up_ind, down_ind = self._split_inds[site] + bits = (int(config[up_ind]), int(config[down_ind])) + out.append(code_from_bits[bits]) + return out + + if self.site_order: + out = [] + for site in self.site_order: + try: + site_ind = self.tn.site_ind(site) + except (AttributeError, KeyError): + site_ind = f"k{site[0]},{site[1]}" + out.append(int(config[site_ind])) + return out + + # Keep the small dummy-network/testing protocol backwards compatible. out = [None] * (self.Lx * self.Ly) for i in range(self.Lx): for j in range(self.Ly): @@ -1200,7 +3047,7 @@ def _sample_d2bp(self, sample_seed, bp_kwargs=None): sample_d2bp = _sample_d2bp - return sample_d2bp(self.tn, **kwargs) + return sample_d2bp(self._bp_tn, **kwargs) def _contract_sample( self, @@ -1214,31 +3061,77 @@ def _contract_sample( ): optimizer = self._get_optimizer() + def scaled_is_finite(value): + if isinstance(value, (tuple, list)) and len(value) == 2: + return bool(np.isfinite(value[0]) and np.isfinite(value[1])) + return bool(np.isfinite(value)) + + def as_scaled(value): + if isinstance(value, (tuple, list)) and len(value) == 2: + return value + return self.mantissa_exponent10(value) + if method == "mps": - return tn_flat.contract_boundary( + opts = { + "optimize": optimizer, + "strip_exponent": True, + } + result = tn_flat.contract_boundary( max_bond=int(chi), mode="mps", - final_contract_opts={"optimize": optimizer, "strip_exponent": True}, + final_contract_opts=opts, max_separation=max_separation, cutoff=cutoff, sequence=["xmin", "xmax", "ymin", "ymax"], equalize_norms=equalize_norms, progbar=False, ) + if not scaled_is_finite(result): + opts["strip_exponent"] = False + result = tn_flat.contract_boundary( + max_bond=int(chi), + mode="mps", + final_contract_opts=opts, + max_separation=max_separation, + cutoff=cutoff, + sequence=["xmin", "xmax", "ymin", "ymax"], + equalize_norms=equalize_norms, + progbar=False, + ) + return as_scaled(result) if method == "ctmrg": - return tn_flat.contract_ctmrg( + opts = { + "optimize": optimizer, + "strip_exponent": True, + } + result = tn_flat.contract_ctmrg( max_bond=int(chi), - final_contract_opts={"optimize": optimizer, "strip_exponent": True}, + final_contract_opts=opts, max_separation=max_separation, cutoff=cutoff, inplace=False, equalize_norms=equalize_norms, progbar=False, ) + if not scaled_is_finite(result): + opts["strip_exponent"] = False + result = tn_flat.contract_ctmrg( + max_bond=int(chi), + final_contract_opts=opts, + max_separation=max_separation, + cutoff=cutoff, + inplace=False, + equalize_norms=equalize_norms, + progbar=False, + ) + return as_scaled(result) if method == "exact": - return tn_flat.contract(all, optimize=optimizer, strip_exponent=True) + result = tn_flat.contract(all, optimize=optimizer, strip_exponent=True) + if not scaled_is_finite(result): + result = tn_flat.contract(all, optimize=optimizer, strip_exponent=False) + return as_scaled(result) raise ValueError(f"Unknown contraction method: {method!r}") diff --git a/src/pepsy/tensors/README.md b/src/pepsy/tensors/README.md index 3cd92cf..18a461c 100644 --- a/src/pepsy/tensors/README.md +++ b/src/pepsy/tensors/README.md @@ -30,11 +30,13 @@ modes include `snake`, `snake-row-major`, `row-major`, `col-major`, `hilbert`, Constructors create common tensor-network states and operators: -- `ps_to_mps`, `ps_to_peps`, `ps_to_3dpeps` +- `ps_to_mps` (bond-one product states), `ps_to_ttn`, `ps_to_peps`, `ps_to_3dpeps` - `ps_to_mpo`, `ps_to_pepo` - `id_to_mpo`, `id_to_pepo` - `haar_random_state`, `random_haar_qubit` -- `hrps_to_mps`, `hrps_to_peps` +- `hrs_to_mps` / `hrps_to_mps` (random MPS states), `hrs_to_ttn` / + `hrps_to_ttn` (random tree states), `hrs_to_peps` / `hrps_to_peps` + (direct Symmray random PEPS states) Contraction helpers include: diff --git a/src/pepsy/tensors/__init__.py b/src/pepsy/tensors/__init__.py index 132aab1..a53858c 100644 --- a/src/pepsy/tensors/__init__.py +++ b/src/pepsy/tensors/__init__.py @@ -4,6 +4,7 @@ from .validation import validate_tensor_network_tags from .symmetric import ( + FermionLatticeSetup, SymGateStream, SymHamiltonian, SymMPS, @@ -55,14 +56,19 @@ get_default_array_backend, get_default_grad_backend, haar_random_state, + hrs_to_mps, + hrs_to_peps, + hrs_to_ttn, hrps_to_mps, hrps_to_peps, + hrps_to_ttn, id_to_mpo, id_to_pepo, measure_obs, ps_to_3dpeps, ps_to_mpo, ps_to_mps, + ps_to_ttn, ps_to_pepo, ps_to_peps, random_haar_qubit, @@ -88,6 +94,7 @@ __all__ = [ "OneDMap", "Fermion", + "FermionLatticeSetup", "SpinfulFermion", "SpinfulFermionHubbard", "SymmFermions", @@ -123,14 +130,19 @@ "get_default_array_backend", "get_default_grad_backend", "haar_random_state", + "hrs_to_mps", + "hrs_to_peps", + "hrs_to_ttn", "hrps_to_mps", "hrps_to_peps", + "hrps_to_ttn", "id_to_mpo", "id_to_pepo", "measure_obs", "ps_to_3dpeps", "ps_to_mpo", "ps_to_mps", + "ps_to_ttn", "ps_to_pepo", "ps_to_peps", "random_haar_qubit", diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index fc2dfc6..34dcb12 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -3,13 +3,18 @@ from .core import ( add_cycle, haar_random_state, + hrs_to_mps, + hrs_to_peps, + hrs_to_ttn, hrps_to_mps, hrps_to_peps, + hrps_to_ttn, id_to_mpo, id_to_pepo, ps_to_3dpeps, ps_to_mpo, ps_to_mps, + ps_to_ttn, ps_to_pepo, ps_to_peps, random_haar_qubit, @@ -18,13 +23,18 @@ __all__ = [ "add_cycle", "haar_random_state", + "hrs_to_mps", + "hrs_to_peps", + "hrs_to_ttn", "hrps_to_mps", "hrps_to_peps", + "hrps_to_ttn", "id_to_mpo", "id_to_pepo", "ps_to_3dpeps", "ps_to_mpo", "ps_to_mps", + "ps_to_ttn", "ps_to_pepo", "ps_to_peps", "random_haar_qubit", diff --git a/src/pepsy/tensors/core.py b/src/pepsy/tensors/core.py index 78e2fd9..637c1cd 100644 --- a/src/pepsy/tensors/core.py +++ b/src/pepsy/tensors/core.py @@ -2,6 +2,7 @@ import math import warnings +from collections.abc import Mapping from numbers import Integral from string import Formatter from typing import Any @@ -51,12 +52,17 @@ "ps_to_peps", "ps_to_3dpeps", "ps_to_mps", + "ps_to_ttn", "ps_to_pepo", "ps_to_mpo", "haar_random_state", "random_haar_qubit", + "hrs_to_peps", + "hrs_to_mps", + "hrs_to_ttn", "hrps_to_peps", "hrps_to_mps", + "hrps_to_ttn", "id_to_pepo", "add_cycle", ] @@ -1161,7 +1167,7 @@ def build_optimizer( directory=False, hash_method: str = "b", overwrite=False, - on_trial_error: str = "warn", + on_trial_error: str = "ignore", slicing_opts: dict | None = None, reconf_opts: dict | None = None, slicing_reconf_opts: dict | None = None, @@ -1193,7 +1199,9 @@ def build_optimizer( overwrite : bool | str, optional Cache overwrite behavior. on_trial_error : str, optional - How to handle individual trial failures. + How to handle individual trial failures. Defaults to ``"ignore"`` + because one invalid candidate path does not make the reusable search + itself invalid; pass ``"warn"`` or ``"raise"`` to inspect failures. slicing_opts : dict | None, optional Options passed to cotengra slicing heuristics. reconf_opts : dict | None, optional @@ -2128,36 +2136,131 @@ def expec_mpo(mpo, mps, *, contraction_opt=None): return (mps_h | mpo | mps_n).contract(all, optimize=contraction_opt) / divisor -def ps_to_peps(Lx: int, Ly: int, dtype: str = "complex128", theta: float = 0.0, cyclic: bool = False, chi: int = 1, rand_strength: float = 0.0): - """Create a product-state PEPS parameterized by ``theta``. +def ps_to_peps( + Lx: int, + Ly: int | None = None, + dtype: str = "complex128", + theta: float = 0.0, + cyclic: bool = False, + *, + fermion=None, + occupations=None, + site_charge=None, + seed=666, + contraction_opt="auto-hq", + to_backend=None, +): + """Create a product-state PEPS, optionally in a Fermion charge sector. Each site tensor is set so the physical vector is ``[cos(theta), sin(theta)]`` with trivial virtual bonds. + If only ``Lx`` is supplied, a square ``Lx x Lx`` lattice is built. For a + Fermion-aware state, ``occupations`` can be a mapping keyed by PEPS + coordinates or a row-major sequence of ``Lx * Ly`` local charge labels. + The returned object is the underlying quimb PEPS, matching + :func:`ps_to_mps`; its physical tensors are native Symmray arrays. + Parameters ---------- Lx : int Lattice size in x direction. Ly : int - Lattice size in y direction. + Lattice size in y direction. If omitted, use a square lattice. dtype : str, optional Tensor dtype passed to numpy/quimb. theta : float, optional Product-state angle controlling local amplitudes. cyclic : bool, optional - If True, add periodic bonds (bond dimension 1) via :func:`add_cycle`. - chi : int, optional - Target bond dimension. If greater than 1, the bond dimension is - expanded via ``expand_bond_dimension`` after initialization. - rand_strength : float, optional - Random noise strength passed to ``expand_bond_dimension``. + If True, add periodic bonds. Fermion-aware states use the native + Symmray PEPS constructor so periodic fermionic bonds are retained. + fermion : :class:`~pepsy.tensors.Fermion`, optional + If supplied, construct a native fermionic Symmray PEPS using the + model's physical sectors and symmetry. + occupations : mapping or sequence, optional + Local charge labels selecting the product-state sector. Mappings use + ``(x, y)`` keys; sequences are interpreted in row-major coordinate + order. If omitted, the Fermion half-filled pattern is used. + site_charge : callable or mapping, optional + Advanced override for the per-site charge pattern. By default this is + derived from ``occupations``. + seed : int, optional + Random seed for the Fermion-aware and ordinary constructors. + contraction_opt : object, optional + Contraction optimizer stored by the internal symmetric wrapper while + constructing a Fermion-aware PEPS. + to_backend : callable, optional + Backend mapper applied to Fermion-aware Symmray blocks. Returns ------- quimb.tensor.PEPS - Initialized PEPS with bond dimension ``chi``. + Initialized bond-one PEPS. """ - peps = qtn.PEPS.rand(Lx=Lx, Ly=Ly, bond_dim=1, seed=666, dtype=dtype) + if Ly is None: + if isinstance(Lx, (tuple, list)): + if len(Lx) != 2: + raise ValueError("A PEPS shape must contain exactly two dimensions.") + Lx, Ly = Lx + else: + Ly = Lx + Lx = int(Lx) + Ly = int(Ly) + if Lx < 1 or Ly < 1: + raise ValueError("PEPS dimensions must be positive integers.") + + if fermion is not None: + from .symmetric import ( # pylint: disable=import-outside-toplevel + Fermion, + SymPEPS, + site_charge_from_occupations, + ) + + if not isinstance(fermion, Fermion): + raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") + + coordinates = tuple( + (x, y) + for x in range(Lx) + for y in range(Ly) + ) + if occupations is None: + occupation_values = fermion.half_filled_occupations(len(coordinates)) + occupations = dict(zip(coordinates, occupation_values)) + elif isinstance(occupations, Mapping): + occupations = dict(occupations) + else: + occupations = tuple(occupations) + if len(occupations) != len(coordinates): + raise ValueError( + "occupations must contain exactly Lx * Ly charge labels." + ) + occupations = dict(zip(coordinates, occupations)) + if site_charge is None: + site_charge = site_charge_from_occupations(occupations) + + state = SymPEPS.random( + Lx, + Ly, + symmetry=fermion.symmetry, + bond_dim=1, + phys_dim=fermion.physical_sectors, + cyclic=cyclic, + seed=seed, + dtype=dtype, + fermionic=True, + site_charge=site_charge, + contraction_opt=contraction_opt, + to_backend=to_backend, + ) + return state.peps + + 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=...") + + peps = qtn.PEPS.rand(Lx=Lx, Ly=Ly, bond_dim=1, seed=seed, dtype=dtype) local_vec = np.array([math.cos(theta), math.sin(theta)], dtype=dtype) for x in range(Lx): for y in range(Ly): @@ -2173,8 +2276,6 @@ def ps_to_peps(Lx: int, Ly: int, dtype: str = "complex128", theta: float = 0.0, peps.astype_(dtype) if cyclic: peps = add_cycle(peps, bond_dim=1) - if chi > 1: - peps.expand_bond_dimension_(chi, rand_strength=rand_strength) return peps @@ -2247,22 +2348,98 @@ def ps_to_3dpeps( return peps +def _fermionic_site_charge_values(site_charge, n): + """Resolve a Fermion constructor's public site-charge input to a map.""" + if callable(site_charge): + return {site: site_charge(site) for site in range(n)} + try: + return {site: site_charge[site] for site in range(n)} + except (KeyError, TypeError) as exc: + raise TypeError( + "site_charge must be callable or map every site 0 .. n - 1." + ) from exc + + +def _fermionic_product_fock_specs(fermion, n, occupations, site_charge): + """Resolve charge labels to definite local Fermion Fock basis states.""" + requested_charges = ( + None + if site_charge is None + else _fermionic_site_charge_values(site_charge, n) + ) + if occupations is None: + if requested_charges is None: + occupations = tuple(fermion.half_filled_occupations(n)) + else: + occupations = tuple(requested_charges[site] for site in range(n)) + else: + occupations = tuple(occupations) + if len(occupations) != n: + raise ValueError("occupations must contain exactly one label per site.") + + specs = { + site: fermion.local_fock_state(occupation, site=site) + for site, occupation in enumerate(occupations) + } + if requested_charges is not None: + for site, requested in requested_charges.items(): + requested_charge, _ = fermion.local_fock_state(requested, site=site) + if requested_charge != specs[site][0]: + raise ValueError( + f"site_charge at site {site} is incompatible with its " + "requested Fock occupation." + ) + + 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.""" + data = tensor.data + physical_axis = tensor.inds.index(physical_index) + chargemap = data.indices[physical_axis].chargemap + local_index = None + for sector, size in chargemap.items(): + size = int(size) + if sector == charge: + local_index = int(basis_index) + if not 0 <= local_index < size: + raise ValueError( + f"Fock basis index {basis_index} is outside charge sector " + f"{charge!r}." + ) + break + if local_index is None: + raise ValueError(f"physical index has no sector for charge {charge!r}.") + + selected = False + for block_sector, block in data.blocks.items(): + block[...] = 0 + if block_sector[physical_axis] == charge: + entry = [0] * block.ndim + entry[physical_axis] = local_index + block[tuple(entry)] = 1 + selected = True + if not selected: + raise ValueError( + f"no compatible Symmray block exists for physical charge {charge!r}." + ) + tensor.modify(data=data) + + def ps_to_mps( L: int, dtype: str = "complex128", theta: float = 0.0, cyclic: bool = False, - chi: int = 1, - rand_strength: float = 0.0, *, fermion=None, occupations=None, site_charge=None, seed=666, - random_rounds=1000, to_backend=None, ): - """Create a product-state MPS, optionally in a Fermion charge sector. + """Create a bond-one product-state MPS, optionally in a Fermion sector. Each site tensor is set so the physical vector is ``[cos(theta), sin(theta)]`` with trivial virtual bonds. @@ -2277,41 +2454,34 @@ def ps_to_mps( Product-state angle controlling local amplitudes. cyclic : bool, optional If True, create a periodic MPS with bond dimension 1. - chi : int, optional - Target bond dimension. If greater than 1, the bond dimension is - 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 fermionic Symmray MPS using the - model's physical sectors and symmetry. With ``chi=1`` this is a - charge-fixed product state. For ``chi>1`` a charge-preserving random - unitary growth produces a symmetric initial MPS. + model's physical sectors and symmetry. The result always has bond + dimension one and is fixed to the requested Fock basis state. occupations : sequence, optional - Local charge labels used to select the product sector. If omitted, - ``fermion.half_filled_occupations(L)`` is used. + Local Fock occupations. A spinful value can be a scalar particle count + or an explicit ``(n_up, n_down)`` pair. If omitted, + ``fermion.half_filled_occupations(L)`` is used. A scalar spinful + charge-one value selects the deterministic checkerboard spin pattern. site_charge : callable or mapping, optional Advanced override for the per-site charge pattern. By default this is derived from ``occupations``. seed : int, optional Seed for the Fermion-aware construction and the ordinary constructor. - random_rounds : int, optional - Maximum number of charge-preserving random-unitary rounds used when - ``fermion`` is supplied with ``chi > 1``. to_backend : callable, optional Backend mapper applied to Fermion-aware Symmray blocks. Returns ------- quimb.tensor.MatrixProductState - Initialized MPS with bond dimension ``chi``. When ``fermion`` is - supplied, the physical tensors are native fermionic Symmray arrays. + Initialized bond-one MPS. When ``fermion`` is supplied, the physical + tensors are native fermionic Symmray arrays. """ if fermion is not None: from .symmetric import ( # pylint: disable=import-outside-toplevel Fermion, SymMPS, + _apply_to_tensor_network_arrays, site_charge_from_occupations, ) @@ -2319,42 +2489,26 @@ def ps_to_mps( raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") if cyclic: raise ValueError("Fermion-aware ps_to_mps currently requires an open chain.") - if occupations is not None: - occupations = tuple(occupations) - if len(occupations) != int(L): - raise ValueError("occupations must contain exactly L charge labels.") - if occupations is None: - occupations = fermion.half_filled_occupations(L) - if site_charge is None: - site_charge = site_charge_from_occupations(occupations) - chi = int(chi) - if chi < 1: - raise ValueError("chi must be a positive integer.") - if chi == 1: - state = SymMPS.random( - L, - symmetry=fermion.symmetry, - bond_dim=1, - phys_dim=fermion.physical_sectors, - seed=seed, - dtype=dtype, - fermionic=True, - site_charge=site_charge, - to_backend=to_backend, - ) - else: - state = SymMPS.random_unitary_evolution( - L, - symmetry=fermion.symmetry, - bond_dim=chi, - phys_dim=fermion.physical_sectors, - seed=seed, - dtype=dtype, - fermionic=True, - site_charge=site_charge, - rounds=random_rounds, - to_backend=to_backend, + fock_specs, leaf_charges = _fermionic_product_fock_specs( + fermion, int(L), occupations, site_charge, + ) + state = SymMPS.random( + L, + symmetry=fermion.symmetry, + bond_dim=1, + phys_dim=fermion.physical_sectors, + seed=seed, + dtype=dtype, + fermionic=True, + site_charge=site_charge_from_occupations(leaf_charges), + to_backend=None, + ) + for site, (charge, basis_index) in fock_specs.items(): + _set_fermionic_product_leaf( + state.mps[site], state.mps.site_ind(site), + charge=charge, basis_index=basis_index, ) + _apply_to_tensor_network_arrays(state.mps, to_backend) return state.mps if occupations is not None or site_charge is not None: @@ -2384,11 +2538,352 @@ def ps_to_mps( tensor.modify(data=data) mps.astype_(dtype) - if chi > 1: - mps.expand_bond_dimension_(chi, rand_strength=rand_strength) return mps +def ps_to_ttn( + n: int, + dtype: str = "complex128", + theta: float = 0.0, + *, + tree=None, + order=None, + structure="balanced", + max_arity=2, + community_frac=0.35, + star_frac=0.75, + chi: int = 1, + rand_strength: float = 0.0, + seed=666, + site_tag_id="I{}", + site_ind_id="k{}", + node_tag_id="N{}", + fermion=None, + occupations=None, + site_charge=None, + to_backend=None, +): + """Create a product-state tree tensor network. + + This is the TTN counterpart of :func:`ps_to_mps`: every qubit starts in + the local state ``[cos(theta), sin(theta)]`` and the tree virtual bonds + start at dimension one. Supply ``tree=`` with an explicit + :class:`~pepsy.optimizers.tree.TreePlan`, or use ``order=`` and the tree + construction options to choose the geometry. + + Parameters + ---------- + n : int + Number of qubits. Qubit labels are ``0 .. n - 1``. + dtype : str, optional + Tensor dtype passed to the tree tensors. + theta : float, optional + Product-state angle controlling each local amplitude vector. + tree : TreePlan, optional + Explicit rooted tree geometry. When omitted, a plan is built from + ``order`` (or ``range(n)``) using ``structure``. + order : sequence of int, optional + Leaf order used to build a plan when ``tree`` is not supplied. + structure, max_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``. + rand_strength : float, optional + Random noise strength for the newly added bond entries, matching the + corresponding ``ps_to_mps`` option. The resulting TTN is + re-canonicalised around its root. + seed : int, optional + Seed used by Quimb when ``rand_strength`` adds random entries. + fermion : :class:`~pepsy.tensors.Fermion`, optional + Build a native fermionic Symmray TTN using the model's local sectors. + This Fock-fixed product-state path requires ``chi=1``. + occupations : sequence or mapping, optional + Local Fock occupations selecting the product state. A mapping is keyed + by qubit label; omitted occupations use the Fermion half-filled + pattern. A spinful value can be a scalar particle count or an explicit + ``(n_up, n_down)`` pair; scalar charge one selects the deterministic + checkerboard spin representative. + site_charge : callable or mapping, optional + Advanced override for the local charge pattern. + to_backend : callable, optional + Backend mapper applied to the native Symmray blocks. + + Returns + ------- + TreeTensorNetwork + The initialized tree state. + """ + from ..optimizers.tree import TreePlan, TreeTensorNetwork + + try: + n = int(n) + except (TypeError, ValueError) as exc: + raise ValueError("n must be a positive integer.") from exc + if n < 1: + raise ValueError("n must be a positive integer.") + if chi is None: + chi = 1 + try: + chi = int(chi) + except (TypeError, ValueError) as exc: + raise ValueError("chi must be a positive integer.") from exc + if chi < 1: + raise ValueError("chi must be a positive integer.") + + if tree is not None and order is not None: + raise ValueError("pass either tree= or order=, not both.") + if tree is None: + if order is None: + order = range(n) + plan = TreePlan.from_order( + order, + structure=structure, + max_arity=max_arity, + community_frac=community_frac, + star_frac=star_frac, + ) + else: + if not isinstance(tree, TreePlan): + raise TypeError("tree must be a TreePlan.") + plan = tree + if plan.n != n: + raise ValueError( + f"tree contains {plan.n} qubits, but n={n} was requested." + ) + + if fermion is not None: + from .symmetric import ( # pylint: disable=import-outside-toplevel + Fermion, + _apply_to_tensor_network_arrays, + ) + + if not isinstance(fermion, Fermion): + raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") + if chi != 1: + raise ValueError( + "Fermion-aware ps_to_ttn is a charge-fixed product-state " + "constructor and requires chi=1; use hrs_to_ttn for chi > 1." + ) + if theta != 0.0: + raise ValueError("theta is not supported with fermion=...") + if isinstance(occupations, Mapping): + try: + occupations = tuple(occupations[q] for q in range(n)) + except KeyError as exc: + raise ValueError( + "occupations mapping must contain every qubit 0 .. n - 1." + ) from exc + elif occupations is not None: + occupations = tuple(occupations) + fock_specs, leaf_charges = _fermionic_product_fock_specs( + fermion, n, occupations, site_charge, + ) + ttn = TreeTensorNetwork.from_symmray_plan( + plan, + symmetry=fermion.symmetry, + physical_sectors=fermion.physical_sectors, + leaf_charges=leaf_charges, + bond_dim=1, + fermionic=True, + seed=seed, + dtype=dtype, + site_tag_id=site_tag_id, + site_ind_id=site_ind_id, + 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), + charge=charge, basis_index=basis_index, + ) + _apply_to_tensor_network_arrays(ttn, to_backend) + norm_sq = np.asarray( + ar.to_numpy(ttn._fermionic_norm_squared()) + ).item() + norm = math.sqrt(abs(complex(norm_sq))) + if norm == 0.0: + raise ValueError("fermionic TTN product state has zero norm.") + root_tensor = ttn.node_tensor(ttn.root) + root_tensor.modify(data=root_tensor.data / norm) + # The norm readout is cached on native TTNs; this direct final scaling + # changes the state without going through a TTN mutator wrapper. + ttn._invalidate_norm_cache() + return ttn + + 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=...") + + ttn = TreeTensorNetwork.from_plan( + plan, + dtype=dtype, + site_tag_id=site_tag_id, + site_ind_id=site_ind_id, + node_tag_id=node_tag_id, + ) + 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)) + phys_axis = tensor.inds.index(ttn.site_ind(q)) + data = np.zeros_like(tensor.data, dtype=dtype) + slicer = [0] * data.ndim + slicer[phys_axis] = slice(None) + data[tuple(slicer)] = local_vec + tensor.modify(data=data) + + if chi > 1: + if rand_strength: + from quimb import seed_rand + + seed_rand(seed) + ttn.expand_bond_dimension_(chi, rand_strength=rand_strength) + ttn.canonize_around_node_(plan.root) + return ttn.validate() + + +def hrs_to_ttn( + n: int, + dtype: str = "complex128", + *, + tree=None, + order=None, + structure="balanced", + max_arity=2, + community_frac=0.35, + star_frac=0.75, + seed=None, + perturb: float = 0.0, + haar_params=None, + chi: int = 1, + rand_strength: float = 0.0, + fermion=None, + occupations=None, + site_charge=None, + to_backend=None, +): + """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. + """ + from ..optimizers.tree import TreePlan, TreeTensorNetwork + + try: + n = int(n) + except (TypeError, ValueError) as exc: + raise ValueError("n must be a positive integer.") from exc + if n < 1: + raise ValueError("n must be a positive integer.") + try: + chi = int(chi) + except (TypeError, ValueError) as exc: + raise ValueError("chi must be a positive integer.") from exc + if chi < 1: + raise ValueError("chi must be a positive integer.") + if tree is not None and order is not None: + raise ValueError("pass either tree= or order=, not both.") + if tree is None: + plan = TreePlan.from_order( + range(n) if order is None else order, + structure=structure, + max_arity=max_arity, + community_frac=community_frac, + star_frac=star_frac, + ) + 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 fermion is not None: + from .symmetric import ( # pylint: disable=import-outside-toplevel + Fermion, + _apply_to_tensor_network_arrays, + site_charge_from_occupations, + ) + + if not isinstance(fermion, Fermion): + raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") + if haar_params is not None: + raise ValueError("haar_params is not supported with fermion=...") + if occupations is None: + occupations = tuple(fermion.half_filled_occupations(n)) + elif isinstance(occupations, Mapping): + try: + occupations = tuple(occupations[q] for q in range(n)) + except KeyError as exc: + raise ValueError( + "occupations mapping must contain every qubit 0 .. n - 1." + ) from exc + else: + occupations = tuple(occupations) + if len(occupations) != n: + raise ValueError("occupations must contain exactly n charge labels.") + if site_charge is None: + site_charge = site_charge_from_occupations(occupations) + if callable(site_charge): + leaf_charges = {q: site_charge(q) for q in range(n)} + else: + try: + leaf_charges = {q: site_charge[q] for q in range(n)} + except (KeyError, TypeError) as exc: + raise TypeError( + "site_charge must be callable or map every qubit label." + ) from exc + ttn = TreeTensorNetwork.from_symmray_plan( + plan, + symmetry=fermion.symmetry, + physical_sectors=fermion.physical_sectors, + leaf_charges=leaf_charges, + bond_dim=chi, + fermionic=True, + seed=seed, + dtype=dtype, + ) + _apply_to_tensor_network_arrays(ttn, to_backend) + return ttn + + 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=...") + ttn = TreeTensorNetwork.from_plan(plan, dtype=dtype) + if haar_params is not None: + if len(haar_params) != n: + raise ValueError(f"haar_params must have length {n}.") + params = tuple(haar_params) + else: + params = tuple( + random_haar_qubit( + None if seed is None else int(seed) + q, + perturb=perturb, + ) + for q in range(n) + ) + for q, (theta, phi) in enumerate(params): + local_vec = np.array( + [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)) + physical_axis = tensor.inds.index(ttn.site_ind(q)) + data = np.zeros_like(tensor.data, dtype=dtype) + selector = [0] * data.ndim + selector[physical_axis] = slice(None) + data[tuple(selector)] = local_vec + tensor.modify(data=data) + if chi > 1: + ttn.expand_bond_dimension_(chi, rand_strength=rand_strength) + ttn.canonize_around_node_(plan.root) + return ttn.validate() + + def ps_to_pepo( Lx: int, Ly: int, @@ -2531,7 +3026,7 @@ def haar_random_state( """Create a dense Haar-random ``L``-qubit state. This samples a full Hilbert-space state, so the result is generally - entangled. Unlike :func:`hrps_to_mps`, this is not a product-state tensor + entangled. Unlike :func:`hrs_to_mps`, this is not a product-state tensor network: it returns dense amplitudes with ``2**L`` entries. Parameters @@ -2589,9 +3084,9 @@ def haar_random_state( return state -def hrps_to_peps( +def hrs_to_peps( Lx: int, - Ly: int, + Ly: int | None = None, dtype: str = "complex128", cyclic: bool = False, seed=None, @@ -2599,22 +3094,161 @@ def hrps_to_peps( haar_params=None, chi: int = 1, rand_strength: float = 0.0, + *, + fermion=None, + occupations=None, + site_charge=None, + method="direct", + subsizes="maximal", + contraction_opt="auto-hq", + to_backend=None, ): - """Create a PEPS with per-site single-qubit Haar states. + """Create a random product or Fermion-symmetric PEPS. - If ``haar_params`` is omitted, each site uses :func:`random_haar_qubit`. - With ``seed`` set, site ``k`` uses ``seed + k`` for reproducible but - distinct samples. + Without ``fermion``, each site is an independent single-qubit Haar state. + 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 + implemented. + In the fermionic branch, ``haar_params`` and ``perturb`` do not apply. Parameters ---------- + Lx, Ly : int or tuple[int, int] + Lattice dimensions. If only ``Lx`` is supplied, use a square lattice. + dtype : str, optional + Tensor dtype passed to numpy/quimb or Symmray. + cyclic : bool, optional + Whether to create periodic PEPS bonds. + seed : int, optional + Random seed. In the ordinary branch, site ``k`` uses ``seed + k`` for + reproducible but distinct Haar samples. + perturb : float, optional + Perturbation applied to ordinary single-qubit Haar angles. + haar_params : sequence, optional + Explicit ``(theta, phi)`` pairs for the ordinary branch. chi : int, optional - Target bond dimension. If greater than 1, the bond dimension is - expanded via ``expand_bond_dimension`` after initialization. + Target virtual bond dimension. In the Fermion branch this controls + the symmetric random-state construction directly. rand_strength : float, optional - Random noise strength passed to ``expand_bond_dimension``. + Random noise passed to ordinary ``expand_bond_dimension``. + fermion : :class:`~pepsy.tensors.Fermion`, optional + Build a native fermionic Symmray PEPS using this model's symmetry and + physical sectors. + occupations : mapping or sequence, optional + Local charge labels selecting the Fermion sector. Mappings use + ``(x, y)`` keys; sequences use row-major order. If omitted, the + Fermion half-filled pattern is used. + site_charge : callable or mapping, optional + 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. + 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. + + Returns + ------- + quimb.tensor.PEPS + The initialized PEPS. Fermion-aware states use native Symmray arrays. """ - peps = ps_to_peps(Lx=Lx, Ly=Ly, dtype=dtype, theta=0.0, cyclic=cyclic) + if Ly is None: + if isinstance(Lx, (tuple, list)): + if len(Lx) != 2: + raise ValueError("A PEPS shape must contain exactly two dimensions.") + Lx, Ly = Lx + else: + Ly = Lx + Lx = int(Lx) + Ly = int(Ly) + if Lx < 1 or Ly < 1: + raise ValueError("PEPS dimensions must be positive integers.") + + method = str(method).strip().lower().replace("-", "_") + if method not in {"direct", "unitary"}: + raise ValueError("method must be 'direct' or 'unitary'.") + + if fermion is not None: + from .symmetric import ( # pylint: disable=import-outside-toplevel + Fermion, + SymPEPS, + site_charge_from_occupations, + ) + + if not isinstance(fermion, Fermion): + raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") + if method == "unitary": + raise NotImplementedError( + "hrs_to_peps(method='unitary') is not implemented; use " + "method='direct'." + ) + if haar_params is not None: + raise ValueError("haar_params is not supported with fermion=...") + + coordinates = tuple( + (x, y) + for x in range(Lx) + for y in range(Ly) + ) + if occupations is None: + occupation_values = fermion.half_filled_occupations(len(coordinates)) + occupations = dict(zip(coordinates, occupation_values)) + elif isinstance(occupations, Mapping): + occupations = dict(occupations) + missing = set(coordinates).difference(occupations) + if missing: + raise ValueError( + "occupations is missing PEPS coordinates: " + f"{sorted(missing)!r}." + ) + else: + occupations = tuple(occupations) + if len(occupations) != len(coordinates): + raise ValueError( + "occupations must contain exactly Lx * Ly charge labels." + ) + occupations = dict(zip(coordinates, occupations)) + if site_charge is None: + site_charge = site_charge_from_occupations(occupations) + + state = SymPEPS.random( + Lx, + Ly, + symmetry=fermion.symmetry, + bond_dim=chi, + phys_dim=fermion.physical_sectors, + cyclic=cyclic, + seed=seed, + dtype=dtype, + fermionic=True, + site_charge=site_charge, + subsizes=subsizes, + contraction_opt=contraction_opt, + to_backend=None, + ) + state.normalize() + if to_backend is not None: + state.apply_to_arrays(to_backend) + return state.peps + + 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=...") + + peps = ps_to_peps( + Lx=Lx, + Ly=Ly, + dtype=dtype, + theta=0.0, + cyclic=cyclic, + seed=seed, + ) n_sites = Lx * Ly if haar_params is not None: @@ -2651,7 +3285,7 @@ def hrps_to_peps( return peps -def hrps_to_mps( +def hrs_to_mps( L: int, dtype: str = "complex128", cyclic: bool = False, @@ -2660,22 +3294,185 @@ def hrps_to_mps( haar_params=None, chi: int = 1, rand_strength: float = 0.0, + *, + fermion=None, + occupations=None, + site_charge=None, + method="unitary", + subsizes="maximal", + random_rounds=1000, + stall_rounds=8, + cutoff=1e-12, + contraction_opt="auto-hq", + to_backend=None, ): - """Create a MPS with per-site single-qubit Haar states. + """Create a random product or Fermion-symmetric MPS. - If ``haar_params`` is omitted, each site uses :func:`random_haar_qubit`. - With ``seed`` set, site ``k`` uses ``seed + k`` for reproducible but - distinct samples. + Without ``fermion``, each site is an independent single-qubit Haar state. + With ``fermion``, construct a native charge-preserving random MPS instead: + ``method="unitary"`` (the default) starts from a random product state and + applies random charge-preserving two-site unitaries, while + ``method="direct"`` calls Symmray's direct block-filled random-MPS + constructor. In both cases ``chi`` controls the target total bond + dimension and the resulting state is normalized. In the fermionic branch, + ``haar_params`` and ``perturb`` do not apply. Parameters ---------- + L : int + Number of sites. + dtype : str, optional + Tensor dtype passed to numpy/quimb or Symmray. + cyclic : bool, optional + Whether to create periodic bonds. Fermion-aware MPS currently require + an open chain. + seed : int, optional + Random seed. In the ordinary branch, site ``k`` uses ``seed + k`` for + reproducible but distinct Haar samples. + perturb : float, optional + Perturbation applied to ordinary single-qubit Haar angles. + haar_params : sequence, optional + Explicit ``(theta, phi)`` pairs for the ordinary branch. chi : int, optional - Target bond dimension. If greater than 1, the bond dimension is - expanded via ``expand_bond_dimension`` after initialization. + Target virtual bond dimension. In the Fermion branch this controls + the symmetric random-state construction directly. rand_strength : float, optional - Random noise strength passed to ``expand_bond_dimension``. + Random noise passed to ordinary ``expand_bond_dimension``. + fermion : :class:`~pepsy.tensors.Fermion`, optional + Build a native fermionic Symmray MPS using this model's symmetry and + physical sectors. + occupations : sequence, optional + Local charge labels selecting the Fermion sector. If omitted, the + Fermion half-filled pattern is used. + site_charge : callable or mapping, optional + Advanced override for the per-site charge pattern. + method : {"unitary", "direct"}, optional + Fermion-aware random-state construction. ``"unitary"`` grows from a + product state using random neutral two-site unitaries. ``"direct"`` + uses :func:`symmray.MPS_fermionic_rand` to fill allowed random blocks + directly, then right-canonicalizes and normalizes the result. + subsizes : object, optional + Symmray charge-sector sizing policy used by ``method="direct"``. + The default ``"maximal"`` keeps as many charge sectors as possible. + This option is ignored by the unitary and non-fermionic branches. + random_rounds : int, optional + Maximum number of random charge-preserving unitary rounds for + ``chi > 1``. + stall_rounds : int, optional + Stop after this many rounds without increasing the bond dimension. + cutoff : float, optional + Truncation cutoff used during charge-preserving growth. + contraction_opt : object, optional + Contraction optimizer stored by the internal symmetric wrapper. + to_backend : callable, optional + Backend mapper applied to Fermion-aware Symmray blocks. + + Returns + ------- + quimb.tensor.MatrixProductState + The initialized MPS. Fermion-aware states use native Symmray arrays. """ - mps = ps_to_mps(L=L, dtype=dtype, theta=0.0, cyclic=cyclic) + method = str(method).strip().lower().replace("-", "_") + if method not in {"unitary", "direct"}: + raise ValueError("method must be 'unitary' or 'direct'.") + if fermion is None and method == "direct": + raise ValueError("method='direct' requires fermion=... .") + + if fermion is not None: + from .symmetric import ( # pylint: disable=import-outside-toplevel + Fermion, + SymMPS, + site_charge_from_occupations, + ) + + if not isinstance(fermion, Fermion): + raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") + if cyclic: + raise ValueError("Fermion-aware hrs_to_mps currently requires an open chain.") + if haar_params is not None: + raise ValueError("haar_params is not supported with fermion=...") + + if occupations is None: + occupations = fermion.half_filled_occupations(L) + elif isinstance(occupations, Mapping): + try: + occupations = tuple(occupations[i] for i in range(int(L))) + except KeyError as exc: + raise ValueError( + "occupations mapping must contain every site 0 .. L - 1." + ) from exc + else: + occupations = tuple(occupations) + if len(occupations) != int(L): + raise ValueError("occupations must contain exactly L charge labels.") + if site_charge is None: + site_charge = site_charge_from_occupations(occupations) + + chi = int(chi) + if chi < 1: + raise ValueError("chi must be a positive integer.") + + if method == "direct": + try: + import symmray as sr # pylint: disable=import-outside-toplevel + except ImportError as exc: + raise ImportError( + "hrs_to_mps(method='direct') requires the optional " + "dependency 'symmray'." + ) from exc + + constructor = getattr(sr, "MPS_fermionic_rand", None) + if constructor is None: # pragma: no cover - old Symmray fallback + raise ImportError( + "The installed Symmray version does not provide " + "MPS_fermionic_rand." + ) + state = constructor( + fermion.symmetry, + int(L), + bond_dim=chi, + phys_dim=fermion.physical_sectors, + cyclic=False, + seed=seed, + dtype=dtype, + site_charge=site_charge, + subsizes=subsizes, + ) + state.right_canonize() + state.normalize() + if to_backend is not None: + state.apply_to_arrays(to_backend) + return state + + state = SymMPS.random_unitary_evolution( + L, + symmetry=fermion.symmetry, + bond_dim=chi, + phys_dim=fermion.physical_sectors, + seed=seed, + dtype=dtype, + fermionic=True, + site_charge=site_charge, + rounds=random_rounds, + stall_rounds=stall_rounds, + cutoff=cutoff, + contraction_opt=contraction_opt, + to_backend=to_backend, + ) + return state.mps + + 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=...") + + mps = ps_to_mps( + L=L, + dtype=dtype, + theta=0.0, + cyclic=cyclic, + seed=seed, + ) if haar_params is not None: if len(haar_params) != L: @@ -2707,3 +3504,9 @@ def hrps_to_mps( if chi > 1: mps.expand_bond_dimension_(chi, rand_strength=rand_strength) return mps + + +# Backwards-compatible aliases for the original longer spelling. +hrps_to_peps = hrs_to_peps +hrps_to_mps = hrs_to_mps +hrps_to_ttn = hrs_to_ttn diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 54d8c9e..b5dbf7f 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from collections.abc import Mapping from dataclasses import dataclass, field from itertools import product @@ -13,6 +14,7 @@ __all__ = [ "Fermion", + "FermionLatticeSetup", "SpinfulFermion", "fermion_density_param_gen", "fermion_hopping_param_gen", @@ -3715,6 +3717,7 @@ def symm_operator_from_dense( charge=0, fermionic=False, sites=None, + index_maps=None, ): """Convert a dense local operator to a Symmray block-sparse array. @@ -3732,6 +3735,10 @@ def symm_operator_from_dense( ``charge=+/-1`` for U(1) raising/lowering-style operators. sites : int | None Number of local sites acted on. Inferred from rank when omitted. + index_maps : sequence of mappings, optional + Explicit ordered charge maps for the row and column indices. When + omitted for a native fermionic one- or two-site operator, Symmray's + canonical fermion basis ordering is used. """ # Keep user supplied backend arrays intact. In particular, converting a # torch or jax value through ``np.asarray`` either errors for a value that @@ -3762,8 +3769,24 @@ def symm_operator_from_dense( if tuple(int(dim) for dim in ar.shape(arr)) != expected_shape: raise ValueError(f"Operator shape {ar.shape(arr)} does not match expected {expected_shape}.") - index_map = sector_index_map(sectors) - index_maps = tuple(dict(index_map) for _ in range(2 * sites)) + if index_maps is None: + index_map = sector_index_map(sectors) + if fermionic and phys_dim in {2, 4}: + import symmray.fermionic_local_operators as flo # pylint: disable=import-outside-toplevel + + if phys_dim == 2: + charges = flo.get_spinless_charge_indexmap(str(symmetry)) + else: + charges = flo.get_spinful_charge_indexmap(str(symmetry)) + if len(charges) == phys_dim: + index_map = dict(enumerate(charges)) + index_maps = tuple(dict(index_map) for _ in range(2 * sites)) + else: + index_maps = tuple(dict(index_map) for index_map in index_maps) + if len(index_maps) != 2 * sites: + raise ValueError( + "index_maps must contain one map for each row and column index." + ) duals = (False,) * sites + (True,) * sites array_cls = _array_class_for_symmetry(symmetry, fermionic=fermionic) kwargs = {} @@ -3857,6 +3880,28 @@ def repeat(self, steps): ) +@dataclass(frozen=True) +class FermionLatticeSetup: + """Metadata for a spinful half-filled rectangular fermion lattice. + + This container deliberately does not build a PEPS, Hamiltonian, or gate + stream. It only centralizes the lattice sites, edges, symmetry-compatible + occupations, and conserved charge needed by an explicit workflow. + """ + + Lx: int + Ly: int + pattern: str + cyclic: bool + sites: tuple + edges: tuple + occupations: Mapping + spin_occupations: Mapping + target_charge: object + target_particles: int + site_charge: object + + def _sites_from_edges(edges, sites): if sites is not None: out = tuple(sites) @@ -5466,11 +5511,49 @@ 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()) + did_compress = bool(compress and L > 1) if compress and L > 1: compress_opts = {"cutoff": cutoff} if max_bond is not None: compress_opts["max_bond"] = int(max_bond) mpo.compress(**compress_opts) + + requested_max_bond = None if max_bond is None else int(max_bond) + final_max_bond = int(mpo.max_bond()) + report = { + "compressed": did_compress, + "cutoff": cutoff, + "requested_max_bond": requested_max_bond, + "raw_max_bond": raw_max_bond, + "final_max_bond": final_max_bond, + "rank_reduced": final_max_bond < raw_max_bond, + "cap_bound": ( + did_compress + and requested_max_bond is not None + and raw_max_bond > requested_max_bond + ), + "max_bond_exceeded": ( + did_compress + and requested_max_bond is not None + and final_max_bond > requested_max_bond + ), + } + # This record describes MPO construction only; it is not used during + # contraction and can safely travel with the returned MPO as user-facing + # build metadata. + mpo.pepsy_compression_report = report + if report["max_bond_exceeded"]: + warnings.warn( + "SymHamiltonian.to_mpo requested " + f"max_bond={requested_max_bond}, but Symmray compression returned " + f"max bond {final_max_bond}. Tied singular values at the " + "truncation threshold can make this a soft cap; inspect " + "mpo.pepsy_compression_report before relying on a hard memory " + "limit.", + RuntimeWarning, + stacklevel=3, + ) return mpo @@ -6580,7 +6663,12 @@ def norm(self, *, contraction_opt=None): def normalize(self): """Normalize the wrapped tensor network in place.""" - self.psi.normalize() + normalized = self.psi.normalize() + # MPS normalization is in-place and returns the old scalar norm, + # whereas quimb's PEPS normalization returns a new network by + # default. Keep the wrapper's state synchronized with both APIs. + if hasattr(normalized, "tensors"): + self.psi = normalized return self def trotter_gates(self, dt, *, model=None, hamiltonian=None, imaginary=False, order=1, **params): @@ -7062,6 +7150,8 @@ def _fermion_generic_local_modes(spinful, sites): "create_d": (down.dag,), "number_d": (down.dag, down), "double": (up.dag, up, down.dag, down), + "s_plus": (up.dag, down), + "s_minus": (down.dag, up), "pair_create": (up.dag, down.dag), "pair_annihilate": (down, up), } @@ -7231,7 +7321,11 @@ def model(self): """The matching :class:`SymHamiltonian` model name.""" if not self.spinful: return "fermi_hubbard_spinless" - return "fermi_hubbard" if self.symmetry == "U1" else "fermi_hubbard_u1u1" + return ( + "fermi_hubbard" + if self.symmetry in {"U1", "Z2"} + else "fermi_hubbard_u1u1" + ) @property def physical_sectors(self): @@ -7264,6 +7358,86 @@ def pair_annihilation_charge(self): self.symmetry, ) + def lattice_half_filling( + self, + Lx, + Ly=None, + *, + pattern="checkerboard", + cyclic=False, + ): + """Prepare metadata for an explicit half-filled spinful lattice workflow. + + The returned :class:`FermionLatticeSetup` contains the coordinate + sites, nearest-neighbor lattice edges, spin-resolved occupations, and + occupations expressed in this fermion's symmetry. It intentionally + does not construct a PEPS, Hamiltonian, or gate stream, so callers can + keep those steps explicit. + + Parameters + ---------- + Lx, Ly : int + Lattice dimensions. If ``Ly`` is omitted, use a square lattice. + pattern : {"checkerboard", "neel", "neel_like"} + Spin pattern for the one-particle-per-site initial state. + cyclic : bool, optional + Whether to include periodic physical lattice edges. This metadata + flag does not determine the boundary conditions of a PEPS or MPS + state built from the returned setup. + """ + if not self.spinful: + raise ValueError( + "lattice_half_filling is defined for spinful fermions." + ) + if Ly is None: + Ly = Lx + if not isinstance(Lx, Integral) or int(Lx) < 1: + raise ValueError("Lx must be a positive integer.") + if not isinstance(Ly, Integral) or int(Ly) < 1: + raise ValueError("Ly must be a positive integer.") + + pattern_name = str(pattern).lower().replace("-", "_") + if pattern_name not in {"checkerboard", "neel", "neel_like"}: + raise ValueError( + "pattern must be 'checkerboard', 'neel', or 'neel_like'." + ) + + Lx = int(Lx) + Ly = int(Ly) + sites = tuple((x, y) for x in range(Lx) for y in range(Ly)) + spin_occupations = { + site: (1, 0) if (site[0] + site[1]) % 2 == 0 else (0, 1) + for site in sites + } + if self.symmetry == "U1U1": + occupations = dict(spin_occupations) + else: + occupations = { + site: n_up + n_down + for site, (n_up, n_down) in spin_occupations.items() + } + + edges = tuple( + tuple(edge) + for edge in qtn.edges_2d_square(Lx, Ly, cyclic=cyclic) + ) + target_particles = sum( + n_up + n_down for n_up, n_down in spin_occupations.values() + ) + return FermionLatticeSetup( + Lx=Lx, + Ly=Ly, + pattern=pattern_name, + cyclic=bool(cyclic), + sites=sites, + edges=edges, + occupations=occupations, + spin_occupations=spin_occupations, + target_charge=self.total_charge(occupations.values()), + target_particles=target_particles, + site_charge=site_charge_from_occupations(occupations), + ) + def half_filled_occupations(self, L): """Return a half-filled product-state charge pattern of length ``L``.""" if not isinstance(L, Integral) or int(L) < 1: @@ -7275,6 +7449,74 @@ def half_filled_occupations(self, L): return (1,) * L return tuple((1, 0) if site % 2 == 0 else (0, 1) for site in range(L)) + def local_fock_state(self, occupation, *, site=0): + """Return ``(physical_charge, sector_index)`` for one local Fock state. + + The physical basis is ``|0>, |up>, |down>, |up down>`` for spinful + fermions (and ``|0>, |1>`` for spinless fermions). A spin-resolved + ``(n_up, n_down)`` occupation always chooses that basis state exactly. + For spinful ``U1`` and ``Z2`` models, a scalar charge label ``1`` does + not resolve the two one-particle states; it therefore denotes the + deterministic checkerboard representative: ``|up>`` at even sites and + ``|down>`` at odd sites. Pass a pair to select either spin explicitly. + + This distinction matters for product-state constructors: fixing only a + degenerate symmetry sector leaves an arbitrary vector in that sector, + whereas a product state must select a definite Fock basis vector. + """ + if self.spinful: + if isinstance(occupation, (tuple, list, np.ndarray)): + if len(occupation) != 2: + raise ValueError( + "a spinful local occupation must be a scalar or " + "a length-2 (n_up, n_down) pair." + ) + n_up, n_down = (int(value) for value in occupation) + if (n_up, n_down) not in {(0, 0), (0, 1), (1, 0), (1, 1)}: + raise ValueError( + "spinful local occupations must have n_up, n_down in {0, 1}." + ) + else: + number = int(occupation) + if number not in {0, 1, 2}: + raise ValueError( + "a scalar spinful local occupation must be 0, 1, or 2." + ) + if number == 0: + n_up, n_down = 0, 0 + elif number == 2: + n_up, n_down = 1, 1 + elif _site_parity(site): + n_up, n_down = 0, 1 + else: + n_up, n_down = 1, 0 + + charge = ( + n_up + n_down + if self.symmetry in {"U1", "Z2"} + else (n_up, n_down) + ) + charge = _normalize_group_charge(charge, self.symmetry) + if self.symmetry in {"U1", "Z2"}: + fock_charges = (0, 1, 1, 2) + else: + fock_charges = ((0, 0), (1, 0), (0, 1), (1, 1)) + fock_charges = tuple( + _normalize_group_charge(candidate, self.symmetry) + for candidate in fock_charges + ) + fock_index = n_up + 2 * n_down + return charge, sum( + candidate == charge for candidate in fock_charges[:fock_index] + ) + + if isinstance(occupation, (tuple, list, np.ndarray)): + raise ValueError("a spinless local occupation must be the scalar 0 or 1.") + number = int(occupation) + if number not in {0, 1}: + raise ValueError("a spinless local occupation must be 0 or 1.") + return _normalize_group_charge(number, self.symmetry), 0 + def total_charge(self, occupations): """Return the charge sum for an occupation/charge sequence.""" occupations = tuple(occupations) @@ -7300,6 +7542,10 @@ def _local_ops(self): ops["charge"] = ops["number_u"] + ops["number_d"] ops["number"] = ops["charge"] ops["sz"] = 0.5 * (ops["number_u"] - ops["number_d"]) + ops["s_plus"] = ops["create_u"] @ ops["annihilate_d"] + ops["s_minus"] = ops["create_d"] @ ops["annihilate_u"] + ops["sx"] = 0.5 * (ops["s_plus"] + ops["s_minus"]) + ops["sy"] = (-0.5j * ops["s_plus"]) + (0.5j * ops["s_minus"]) ops["pair_create"] = ops["create_u"] @ ops["create_d"] ops["pair_annihilate"] = ops["annihilate_d"] @ ops["annihilate_u"] else: @@ -7323,10 +7569,34 @@ def _operator_name(name): "annihilate_down": "annihilate_d", "doublon": "double", "pair_annihilation": "pair_annihilate", + "spin_plus": "s_plus", + "s_plus": "s_plus", + "spin_minus": "s_minus", + "s_minus": "s_minus", + "spin_x": "sx", + "spin_y": "sy", "spin_z": "sz", } return aliases.get(str(name), str(name)) + @classmethod + def _adjoint_operator_name(cls, name): + """Return the local fermion-operator name for its adjoint.""" + name = cls._operator_name(name) + adjoints = { + "create": "annihilate", + "annihilate": "create", + "create_u": "annihilate_u", + "annihilate_u": "create_u", + "create_d": "annihilate_d", + "annihilate_d": "create_d", + "s_plus": "s_minus", + "s_minus": "s_plus", + "pair_create": "pair_annihilate", + "pair_annihilate": "pair_create", + } + return adjoints.get(name, name) + def dense_operator(self, name): """Return a dense one-site operator in the native basis order. @@ -7347,6 +7617,8 @@ def dense_operator(self, name): def operator_charge(self, name): """Return the Abelian charge carried by ``dense_operator(name)``.""" name = self._operator_name(name) + if name in {"s_plus", "s_minus", "sx", "sy", "sz"} and not self.spinful: + raise ValueError("Spin operators require spinful fermions.") if not self.spinful: if name in {"create", "annihilate"}: charge = 1 if name == "create" else -1 @@ -7362,6 +7634,27 @@ def operator_charge(self, name): if not name.startswith("create"): charge = _neg_charge(charge) return _normalize_group_charge(charge, self.symmetry) + if name in {"s_plus", "s_minus"}: + if name == "s_plus": + charge = _charge_add( + self.operator_charge("create_u"), + self.operator_charge("annihilate_d"), + self.symmetry, + ) + else: + charge = _charge_add( + self.operator_charge("create_d"), + self.operator_charge("annihilate_u"), + self.symmetry, + ) + return _normalize_group_charge(charge, self.symmetry) + if name in {"sx", "sy"}: + if self.symmetry not in {"U1", "Z2"}: + raise ValueError( + f"{name} is not a homogeneous operator under symmetry " + f"{self.symmetry!r}; use symmetry='U1' or 'Z2'." + ) + return self.zero_charge if name == "pair_create": return self.pair_charge if name == "pair_annihilate": @@ -7372,7 +7665,154 @@ def operator(self, name): """Return the cached native Symmray operator for ``name``.""" return self.observable(name) - def operator_term(self, terms, *, sites=None, charge=None, like=None): + def _require_spinful(self, feature): + if not self.spinful: + raise ValueError(f"{feature} requires spinful fermions.") + + def _require_spin_flip_symmetry(self, feature): + self._require_spinful(feature) + if self.symmetry not in {"U1", "Z2"}: + raise ValueError( + f"{feature} requires symmetry='U1' or 'Z2'; " + f"symmetry={self.symmetry!r} keeps up/down charges separate." + ) + + @staticmethod + def _resolve_operator_parameter(value, *, site=None, edge=None): + if site is not None and edge is not None: + raise ValueError("Specify either site= or edge=, not both.") + if edge is not None: + try: + left, right = tuple(edge) + except (TypeError, ValueError) as exc: + raise ValueError("edge must contain exactly two site labels.") from exc + if callable(value) or isinstance(value, Mapping): + return _edge_parameter(value, left, right) + return value + if site is not None: + if callable(value) or isinstance(value, Mapping): + return _node_parameter(value, site) + return value + if callable(value) or isinstance(value, Mapping): + raise ValueError("site= or edge= is required for a site/edge parameter.") + return value + + def _spin_flip_operator(self, name): + self._require_spin_flip_symmetry(f"{name.upper()} operator") + if name == "sx": + terms = [ + (0.5, ((0, "s_plus"),)), + ] + elif name == "sy": + terms = [ + (-0.5j, ((0, "s_plus"),)), + ] + else: # pragma: no cover - private callers pass canonical names. + raise ValueError(f"Unknown spin-flip operator {name!r}.") + return self.operator_term(terms, sites=(0,), add_hc=True) + + def spin_x_operator(self): + """Return the native one-site ``Sx`` operator.""" + return self.observable("sx") + + def spin_y_operator(self): + """Return the native one-site ``Sy`` operator.""" + return self.observable("sy") + + def spin_z_operator(self): + """Return the native one-site ``Sz`` operator.""" + self._require_spinful("Sz operator") + return self.observable("sz") + + def sx_operator(self): + """Alias for :meth:`spin_x_operator`.""" + return self.spin_x_operator() + + def sy_operator(self): + """Alias for :meth:`spin_y_operator`.""" + return self.spin_y_operator() + + def sz_operator(self): + """Alias for :meth:`spin_z_operator`.""" + return self.spin_z_operator() + + def spin_z_correlator(self): + """Return the bare native two-site ``Sz_i Sz_j`` operator.""" + self._require_spinful("Sz-Sz correlators") + return self.operator_term( + [ + (0.25, ((0, "number_u"), (1, "number_u"))), + (-0.25, ((0, "number_u"), (1, "number_d"))), + (-0.25, ((0, "number_d"), (1, "number_u"))), + (0.25, ((0, "number_d"), (1, "number_d"))), + ], + sites=(0, 1), + charge=self.zero_charge, + ) + + def spin_x_correlator(self): + """Return the native two-site ``Sx_i Sx_j`` operator.""" + self._require_spin_flip_symmetry("Sx-Sx correlators") + return self.operator_term( + [ + (0.25, ((0, "s_plus"), (1, "s_plus"))), + (0.25, ((0, "s_plus"), (1, "s_minus"))), + (0.25, ((0, "s_minus"), (1, "s_plus"))), + (0.25, ((0, "s_minus"), (1, "s_minus"))), + ], + sites=(0, 1), + charge=self.zero_charge, + ) + + def spin_y_correlator(self): + """Return the native two-site ``Sy_i Sy_j`` operator.""" + self._require_spin_flip_symmetry("Sy-Sy correlators") + return self.operator_term( + [ + (-0.25, ((0, "s_plus"), (1, "s_plus"))), + (0.25, ((0, "s_plus"), (1, "s_minus"))), + (0.25, ((0, "s_minus"), (1, "s_plus"))), + (-0.25, ((0, "s_minus"), (1, "s_minus"))), + ], + sites=(0, 1), + charge=self.zero_charge, + ) + + def xy_exchange_operator(self): + """Return the native ``Sx_i Sx_j + Sy_i Sy_j`` operator.""" + self._require_spinful("XY exchange operators") + return self.operator_term( + [(0.5, ((0, "s_plus"), (1, "s_minus")))], + sites=(0, 1), + charge=self.zero_charge, + add_hc=True, + ) + + def heisenberg_operator(self): + """Return the native two-site ``S_i dot S_j`` operator.""" + self._require_spinful("Heisenberg operators") + return self.operator_term( + [ + (0.25, ((0, "number_u"), (1, "number_u"))), + (-0.25, ((0, "number_u"), (1, "number_d"))), + (-0.25, ((0, "number_d"), (1, "number_u"))), + (0.25, ((0, "number_d"), (1, "number_d"))), + (0.5, ((0, "s_plus"), (1, "s_minus"))), + (0.5, ((1, "s_plus"), (0, "s_minus"))), + ], + sites=(0, 1), + charge=self.zero_charge, + ) + + def operator_term( + self, + terms, + *, + sites=None, + charge=None, + like=None, + add_hc=False, + ): """Return a native operator made from explicit fermion monomials. Parameters @@ -7389,6 +7829,11 @@ def operator_term(self, terms, *, sites=None, charge=None, like=None): charge : optional Total Abelian operator charge. By default it is inferred and all monomials must have the same charge. + add_hc : bool, optional + Append the Hermitian conjugate of every supplied monomial. The + input must then define a self-conjugate charge sector, as required + for one homogeneous Symmray operator. Fermionic factor order is + reversed when taking the adjoint. This returns the operator itself, not ``exp(-i dt H)``. For example, @@ -7405,7 +7850,6 @@ def operator_term(self, terms, *, sites=None, charge=None, like=None): inferred_sites = [] normalized = [] - term_charges = [] for entry in entries: if not isinstance(entry, (tuple, list)) or len(entry) != 2: raise ValueError("terms must have form (coefficient, operators).") @@ -7421,13 +7865,30 @@ def operator_term(self, terms, *, sites=None, charge=None, like=None): if site not in inferred_sites: inferred_sites.append(site) name = self._operator_name(name) + expanded.append((site, name)) + normalized.append((coefficient, tuple(expanded))) + + if add_hc: + normalized.extend( + ( + ar.do("conj", coefficient), + tuple( + (site, self._adjoint_operator_name(name)) + for site, name in reversed(references) + ), + ) + for coefficient, references in tuple(normalized) + ) + + term_charges = [] + for _, references in normalized: + term_charge = self.zero_charge + for _, name in references: term_charge = _charge_add( term_charge, self.operator_charge(name), self.symmetry, ) - expanded.append((site, name)) - normalized.append((coefficient, tuple(expanded))) term_charges.append(term_charge) if sites is None: @@ -7446,6 +7907,12 @@ def operator_term(self, terms, *, sites=None, charge=None, like=None): inferred_charge = term_charges[0] if any(term_charge != inferred_charge for term_charge in term_charges[1:]): + if add_hc: + raise ValueError( + "add_hc requires a self-conjugate operator charge; " + "charged and Hermitian-conjugate monomials cannot be " + "combined into one homogeneous Symmray operator." + ) raise ValueError( "All monomials in one operator_term must carry the same charge." ) @@ -7480,9 +7947,41 @@ def operator_term(self, terms, *, sites=None, charge=None, like=None): charge=charge, fermionic=True, sites=len(sites), + index_maps=tuple( + { + index: value + for index, value in enumerate(self._local_ops()["index_map"]) + } + for _ in range(2 * len(sites)) + ), ) return _apply_to_array_blocks(operator, self.to_backend) + def eta_pair_operator(self, *, coefficient=1.0): + """Return ``coefficient * Delta_0^dag Delta_1 + h.c.``. + + The returned operator uses canonical two-site locations ``(0, 1)``; + place it on physical sites through ``state.measure(..., where=...)`` + or an explicit VMC term mapping. It is neutral under the selected + spinful symmetry and preserves native fermionic grading. + """ + if not self.spinful: + raise ValueError( + "Eta-pair operators require spinful fermions with up and down " + "modes." + ) + return self.operator_term( + [ + ( + coefficient, + ((0, "pair_create"), (1, "pair_annihilate")), + ) + ], + sites=(0, 1), + charge=self.zero_charge, + add_hc=True, + ) + def _hopping_operator_on_sites( self, left, @@ -7668,14 +8167,24 @@ def observable(self, name): """Return a cached one-site fermionic Symmray operator for ``name``.""" name = self._operator_name(name) if name not in self._observable_cache: - operator = symm_operator_from_dense( - self.dense_operator(name), - self.physical_sectors, - symmetry=self.symmetry, - charge=self.operator_charge(name), - fermionic=True, - sites=1, - ) + if name in {"sx", "sy"}: + operator = self._spin_flip_operator(name) + elif name in {"s_plus", "s_minus"}: + self._require_spinful(f"{name} operators") + operator = self.operator_term( + [(1.0, ((0, name),))], + sites=(0,), + charge=self.operator_charge(name), + ) + else: + operator = symm_operator_from_dense( + self.dense_operator(name), + self.physical_sectors, + symmetry=self.symmetry, + charge=self.operator_charge(name), + fermionic=True, + sites=1, + ) self._observable_cache[name] = _apply_to_array_blocks(operator, self.to_backend) return self._observable_cache[name] @@ -7690,6 +8199,99 @@ def _cached_gate(self, key, build): self._gate_cache[key] = build() return self._gate_cache[key] + def operator_gate(self, operator, theta, *, imaginary=False): + """Exponentiate a native operator as ``exp(-i theta operator)``. + + ``operator`` can be a named one-site observable, one of the built-in + two-site names ``sxx``, ``syy``, ``szz``, ``xy``, or ``heisenberg``, + or an already-built native Symmray operator from :meth:`operator_term`. + Gates require a charge-neutral operator because exponentiation adds the + identity term and must remain in one homogeneous symmetry sector. + """ + if isinstance(operator, str): + name = self._operator_name(operator) + factories = { + "sxx": self.spin_x_correlator, + "syy": self.spin_y_correlator, + "szz": self.spin_z_correlator, + "xy": self.xy_exchange_operator, + "heisenberg": self.heisenberg_operator, + } + if name in factories: + factory = factories[name] + else: + factory = lambda: self.observable(name) + operator_key = ("name", name) + else: + factory = lambda: operator + operator_key = ("term", id(operator)) + + def build(): + term = factory() + charge = getattr(term, "charge", None) + if charge is not None and charge != self.zero_charge: + raise ValueError( + "operator_gate requires a charge-neutral operator; " + f"got charge {charge!r} for symmetry {self.symmetry!r}." + ) + gate = _gate_from_term(term, theta, imaginary=imaginary) + return _apply_to_array_blocks(gate, self.to_backend) + + return self._cached_gate( + ("operator", operator_key, theta, imaginary), + build, + ) + + def spin_x_gate(self, theta, *, site=None, imaginary=False): + """Return the native one-site ``exp(-i theta Sx)`` gate.""" + theta = self._resolve_operator_parameter(theta, site=site) + return self.operator_gate("sx", theta, imaginary=imaginary) + + def spin_y_gate(self, theta, *, site=None, imaginary=False): + """Return the native one-site ``exp(-i theta Sy)`` gate.""" + theta = self._resolve_operator_parameter(theta, site=site) + return self.operator_gate("sy", theta, imaginary=imaginary) + + def spin_z_gate(self, theta, *, site=None, imaginary=False): + """Return the native one-site ``exp(-i theta Sz)`` gate.""" + theta = self._resolve_operator_parameter(theta, site=site) + return self.operator_gate("sz", theta, imaginary=imaginary) + + def spin_x_correlator_gate(self, theta, *, edge=None, imaginary=False): + """Return the native two-site ``exp(-i theta Sx Sx)`` gate.""" + theta = self._resolve_operator_parameter(theta, edge=edge) + return self.operator_gate("sxx", theta, imaginary=imaginary) + + def spin_y_correlator_gate(self, theta, *, edge=None, imaginary=False): + """Return the native two-site ``exp(-i theta Sy Sy)`` gate.""" + theta = self._resolve_operator_parameter(theta, edge=edge) + return self.operator_gate("syy", theta, imaginary=imaginary) + + def spin_z_correlator_gate(self, theta, *, edge=None, imaginary=False): + """Return the native two-site ``exp(-i theta Sz Sz)`` gate.""" + theta = self._resolve_operator_parameter(theta, edge=edge) + return self.operator_gate("szz", theta, imaginary=imaginary) + + def xy_exchange_gate(self, theta, *, edge=None, imaginary=False): + """Return the native two-site XY-exchange gate.""" + theta = self._resolve_operator_parameter(theta, edge=edge) + return self.operator_gate("xy", theta, imaginary=imaginary) + + def heisenberg_gate(self, theta, *, edge=None, imaginary=False): + """Return the native two-site Heisenberg gate.""" + theta = self._resolve_operator_parameter(theta, edge=edge) + return self.operator_gate("heisenberg", theta, imaginary=imaginary) + + # Short gate spellings match the operator aliases and are convenient in + # small native gate streams. + sx_gate = spin_x_gate + sy_gate = spin_y_gate + sz_gate = spin_z_gate + sxx_gate = spin_x_correlator_gate + syy_gate = spin_y_correlator_gate + szz_gate = spin_z_correlator_gate + xy_gate = xy_exchange_gate + def interaction_gate(self, dt, *, site=None, U=None, imaginary=False): """Return the exact onsite interaction gate. @@ -7826,8 +8428,25 @@ def build(): def gate(self, name, dt, *, site=None, where=None, imaginary=False, **params): """Build a named native gate using the model's local conventions.""" + edge = params.pop("edge", where) del where # Gate locations belong to the stream entry, not the tensor. name = str(name).lower().replace("-", "_") + if name in {"sx", "spin_x"}: + return self.spin_x_gate(dt, site=site, imaginary=imaginary) + if name in {"sy", "spin_y"}: + return self.spin_y_gate(dt, site=site, imaginary=imaginary) + if name in {"sz", "spin_z"}: + return self.spin_z_gate(dt, site=site, imaginary=imaginary) + if name in {"sxx", "spin_x_x", "sx_sx"}: + return self.spin_x_correlator_gate(dt, edge=edge, imaginary=imaginary) + if name in {"syy", "spin_y_y", "sy_sy"}: + return self.spin_y_correlator_gate(dt, edge=edge, imaginary=imaginary) + if name in {"szz", "spin_z_z", "sz_sz"}: + return self.spin_z_correlator_gate(dt, edge=edge, imaginary=imaginary) + if name in {"xy", "xy_exchange"}: + return self.xy_exchange_gate(dt, edge=edge, imaginary=imaginary) + if name in {"heisenberg", "heis"}: + return self.heisenberg_gate(dt, edge=edge, imaginary=imaginary) if name in {"onsite", "hubbard_onsite"}: return self.onsite_gate( dt, diff --git a/src/pepsy/vmc/__init__.py b/src/pepsy/vmc/__init__.py index ced36f8..d63ed8f 100644 --- a/src/pepsy/vmc/__init__.py +++ b/src/pepsy/vmc/__init__.py @@ -7,10 +7,27 @@ from importlib import import_module _SYMBOL_MODULES = { + "BackendCapabilityWarning": ".api", + "VMCBackendCapabilityError": ".api", + "ContractionFallbackWarning": ".api", + "ContractionConfig": ".api", + "CompiledOperatorSum": ".api", "FermionSiteEncoding": ".torch", + "LocalMatrixTerm": ".api", + "MCState": ".api", + "NumericalStabilityWarning": ".api", + "OperatorFactor": ".api", + "OperatorSum": ".api", + "OptimizationConfig": ".api", + "ProductTerm": ".api", + "SamplingConfig": ".api", + "SamplingDiagnosticWarning": ".api", + "SpinlessSiteEncoding": ".torch", + "SymmetryFallbackWarning": ".api", "NetKetLocalConfigMap": ".netket", "NetKetChunkSettings": ".netket", "NetKetPEPSVMC": ".netket", + "NetKetVMCSetup": ".netket", "NetKetFermiHubbardVMC": ".netket", "NetKetSparseFermiHubbardVMC": ".netket", "NetKetVMCSettings": ".netket", @@ -18,19 +35,39 @@ "PackedFermionicPEPS": ".netket", "SpinOrbitalColumns": ".netket", "TorchConnections": ".torch", + "TorchFermionVMC": ".torch", + "TorchFermionVMCMetadata": ".torch", + "TorchChainDiagnostics": ".torch", + "TorchMCMCSamples": ".torch", "TorchMetropolisResult": ".torch", + "TorchMetropolisSampler": ".torch", + "TorchBPMetropolisSampler": ".torch", "TorchPEPSAmplitude": ".torch", "TorchPEPSBoundaryAmplitude": ".torch", "TorchVMCDriver": ".torch", + "TorchVMCSetup": ".torch", + "TorchVMCEnergyEstimate": ".torch", + "TorchVMCImportanceEstimate": ".torch", "TorchVMCStepResult": ".torch", "TorchSRResult": ".torch", "TorchSquareLattice": ".torch", + "VMCMeasurement": ".api", + "VMCOptimizationResult": ".api", + "VMC": ".api", + "VMCProblem": ".api", + "VMCSamples": ".api", + "VMCWarning": ".api", "apply_torch_sr_update": ".torch", "build_heisenberg_vmc": ".netket", "build_ising_vmc": ".netket", "build_fermi_hubbard_vmc": ".netket", + "build_fermion_vmc": ".netket", + "build_netket_vmc": ".netket", "build_sparse_fermi_hubbard_vmc": ".netket", "fermionic_peps_rand": ".netket", + "fermion_model_terms": ".netket", + "netket_fermion_operator": ".netket", + "standard_fermion_observables": ".netket", "choose_netket_chunk_size": ".netket", "configure_jax_for_vmc": ".netket", "config_to_phys_indices": ".netket", @@ -45,22 +82,35 @@ "make_netket_sr_preconditioner": ".netket", "make_netket_vmc_driver": ".netket", "make_torch_peps_amplitude_model": ".torch", + "compile_operator_sum_torch": ".torch", + "build_torch_vmc": ".torch", + "compile_operator_sum_netket": ".netket", "netket_spin_orbital_columns": ".netket", "occupation_to_phys_indices": ".netket", "pack_peps_ansatz": ".netket", "pack_fermionic_peps_ansatz": ".netket", + "prepare_fermionic_peps_for_netket": ".netket", "propose_spin_exchange": ".torch", "propose_spinful_exchange_or_hopping": ".torch", + "propose_spinful_u1_exchange_or_hopping": ".torch", + "propose_spinful_z2_exchange_or_hopping": ".torch", + "propose_spinful_z2z2_exchange_or_hopping": ".torch", "random_spin_configs": ".torch", "random_spinful_configs": ".torch", "recommend_netket_vmc_settings": ".netket", "solve_torch_sr": ".torch", "square_lattice_edges": ".netket", "metropolis_exchange_sweep": ".torch", + "metropolis_local_sampler": ".torch", "spinful_fermi_hubbard_connections": ".torch", "torch_log_derivative_matrix": ".torch", "transverse_ising_connections": ".torch", + "torch_hamiltonian_connections": ".torch", + "torch_chain_diagnostics": ".torch", "verify_netket_spin_columns": ".netket", + "VMCOptimizeResult": ".netket", + "warmup_netket_vmc": ".netket", + "normalize_operator_sum": ".api", } __all__ = tuple(_SYMBOL_MODULES) diff --git a/src/pepsy/vmc/api.py b/src/pepsy/vmc/api.py new file mode 100644 index 0000000..18755d3 --- /dev/null +++ b/src/pepsy/vmc/api.py @@ -0,0 +1,958 @@ +"""Backend-neutral data contracts for PEPS variational Monte Carlo. + +The Torch and NetKet integrations intentionally keep different numerical +engines, but they can share the objects that describe an operator problem and +the results returned by sampling, measurement, and optimization. This module +does not require Torch, JAX, Flax, or NetKet. NumPy is used only for small, +backend-neutral result conveniences. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Mapping + +import numpy as np + +__all__ = [ + "BackendCapabilityWarning", + "VMCBackendCapabilityError", + "ContractionFallbackWarning", + "ContractionConfig", + "CompiledOperatorSum", + "LocalMatrixTerm", + "MCState", + "NumericalStabilityWarning", + "OperatorFactor", + "OperatorSum", + "OptimizationConfig", + "ProductTerm", + "SamplingConfig", + "SamplingDiagnosticWarning", + "SymmetryFallbackWarning", + "VMCMeasurement", + "VMCOptimizationResult", + "VMC", + "VMCProblem", + "VMCSamples", + "VMCWarning", + "normalize_operator_sum", +] + + +class VMCWarning(UserWarning): + """Base warning for non-fatal VMC capability or numerical diagnostics.""" + + +class BackendCapabilityWarning(VMCWarning): + """A backend cannot use a requested optional acceleration.""" + + +class VMCBackendCapabilityError(NotImplementedError): + """A requested portable VMC operation is unavailable on one backend. + + The high-level adapters raise this instead of accepting a shared setting + and silently changing its meaning. Native backend APIs can still expose + their additional capabilities directly. + """ + + +class SymmetryFallbackWarning(VMCWarning): + """A symmetry-aware path selected a slower or less specialized fallback.""" + + +class ContractionFallbackWarning(VMCWarning): + """A requested contraction optimizer or route required a fallback.""" + + +class SamplingDiagnosticWarning(VMCWarning): + """Sampling diagnostics indicate insufficient mixing or support.""" + + +class NumericalStabilityWarning(VMCWarning): + """Amplitudes, logarithms, or update diagnostics approach instability.""" + + +_CONTRACTION_ALIASES = { + "exact": "exact", + "hotrg": "hotrg", + "ctmrg": "ctmrg", + "boundary": "boundary", + "mps": "boundary", + "boundary-mps": "boundary", + "contract-boundary": "boundary", +} + + +def _positive_int(name, value, *, allow_none=False): + if value is None and allow_none: + return None + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + suffix = " or None" if allow_none else "" + raise ValueError(f"{name} must be a positive integer{suffix}.") + return int(value) + + +@dataclass(frozen=True) +class ContractionConfig: + """Shared contraction settings consumed by Torch and NetKet adapters.""" + + method: str = "exact" + chi: int | None = None + cutoff: float = 0.0 + options: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self): + key = str(self.method).replace("_", "-").lower() + try: + method = _CONTRACTION_ALIASES[key] + except KeyError as exc: + raise ValueError( + "contraction method must be 'exact', 'hotrg', 'ctmrg', " + "or 'boundary'." + ) from exc + chi = _positive_int("chi", self.chi, allow_none=True) + if method != "exact" and chi is None: + raise ValueError(f"chi is required for contraction={method!r}.") + cutoff = float(self.cutoff) + if cutoff < 0: + raise ValueError("cutoff must be non-negative.") + object.__setattr__(self, "method", method) + object.__setattr__(self, "chi", chi) + object.__setattr__(self, "cutoff", cutoff) + object.__setattr__(self, "options", MappingProxyType(dict(self.options))) + + +def _resolve_contraction_config(contraction, chi=None, cutoff=None, options=None): + """Return legacy constructor values for a common contraction config.""" + if not isinstance(contraction, ContractionConfig): + return contraction, chi, cutoff, options + if chi is not None and contraction.chi is not None and chi != contraction.chi: + raise ValueError( + f"chi={chi} conflicts with ContractionConfig.chi={contraction.chi}." + ) + if cutoff is not None and float(cutoff) != contraction.cutoff: + raise ValueError( + f"cutoff={cutoff} conflicts with ContractionConfig.cutoff=" + f"{contraction.cutoff}." + ) + if options is not None and dict(options) != dict(contraction.options): + raise ValueError("contraction options conflict with ContractionConfig.options.") + return ( + contraction.method, + contraction.chi if chi is None else chi, + contraction.cutoff if cutoff is None else cutoff, + dict(contraction.options) if options is None else options, + ) + + +@dataclass(frozen=True) +class SamplingConfig: + """Shared chain-preserving sampling settings.""" + + 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 + + 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): + raise ValueError("seed must be an integer or None.") + if self.sampler_seed is not None and isinstance(self.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_chains", n_chains) + object.__setattr__(self, "burn_in", int(self.burn_in)) + object.__setattr__(self, "thin", thin) + object.__setattr__(self, "chunk_size", chunk_size) + + @property + def n_samples(self): + """Total requested samples across all chains.""" + return self.n_samples_per_chain * self.n_chains + + def torch_kwargs(self): + """Return the canonical keyword mapping for Torch samplers. + + The Torch sampler accepts a total ``n_samples`` value, so this + conversion deliberately multiplies by ``n_chains``. Keeping that + conversion here prevents the two backends from silently interpreting + the same user setting differently. + """ + if self.seed is not None and self.sampler_seed is not None: + raise ValueError("Pass either seed or sampler_seed, not both.") + return { + "n_samples": self.n_samples, + "n_chains": self.n_chains, + "n_discard_per_chain": self.burn_in, + "n_thin": self.thin, + "seed": self.seed, + "sampler_seed": self.sampler_seed, + } + + def netket_kwargs(self): + """Return the settings understood by ``MCState.sample``. + + NetKet's sampler owns the chain count, therefore ``n_chains`` is + retained in this mapping for validation by the setup façade rather + than passed to ``MCState.sample``. + """ + return { + "n_samples": self.n_samples, + "n_chains": self.n_chains, + "n_discard_per_chain": self.burn_in, + } + + +@dataclass(frozen=True) +class OptimizationConfig: + """Shared energy-optimization settings.""" + + n_steps: int = 1 + method: str = "sr" + learning_rate: float = 1.0e-3 + diag_shift: float = 1.0e-2 + sr_mode: str = "real" + energy_shift: float = 0.0 + per_site: int | None = None + progress: bool = False + warmup: bool = True + + def __post_init__(self): + n_steps = _positive_int("n_steps", self.n_steps) + method = str(self.method).replace("-", "_").lower() + aliases = { + "sgd": "sgd", + "vmc": "sgd", + "sr": "sr", + "minsr": "minsr", + "min_sr": "minsr", + } + if method not in aliases: + raise ValueError("method must be 'sgd', 'sr', or 'minsr'.") + sr_mode = str(self.sr_mode).replace("_", "-").lower() + sr_mode_aliases = { + "real": "real", + "complex": "complex", + "holomorphic": "holomorphic", + "holomorphic-complex": "holomorphic", + "real-imag": "real-imag", + "real-imaginary": "real-imag", + } + if sr_mode not in sr_mode_aliases: + raise ValueError( + "sr_mode must be 'real', 'complex', 'holomorphic', or " + "'real-imag'." + ) + learning_rate = float(self.learning_rate) + diag_shift = float(self.diag_shift) + if learning_rate <= 0: + raise ValueError("learning_rate must be positive.") + if diag_shift < 0: + raise ValueError("diag_shift must be non-negative.") + per_site = _positive_int("per_site", self.per_site, allow_none=True) + object.__setattr__(self, "n_steps", n_steps) + object.__setattr__(self, "method", aliases[method]) + object.__setattr__(self, "sr_mode", sr_mode_aliases[sr_mode]) + object.__setattr__(self, "learning_rate", learning_rate) + object.__setattr__(self, "diag_shift", diag_shift) + object.__setattr__(self, "per_site", per_site) + + +@dataclass(frozen=True) +class CompiledOperatorSum: + """Backend adapter output containing terms plus an identity constant.""" + + backend: str + terms: Any + constant: Any = 0.0 + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not isinstance(self.backend, str) or not self.backend: + raise ValueError("CompiledOperatorSum.backend must be non-empty.") + object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) + + +@dataclass(frozen=True) +class OperatorFactor: + """One local operator factor in a product term. + + ``name`` is deliberately a string rather than a backend operator object. + Common names include ``"x"``, ``"number"``, and ``"fermion"``. Backend + adapters interpret the name and may use ``spin``/``dagger`` to construct + their native operator. Keeping the factor symbolic is what lets the same + Hamiltonian feed both Torch and NetKet. + """ + + site: Any + name: str + spin: str | None = None + dagger: bool | None = None + + def __post_init__(self): + if not isinstance(self.name, str) or not self.name: + raise ValueError("OperatorFactor.name must be a non-empty string.") + if self.spin is not None and not isinstance(self.spin, str): + raise TypeError("OperatorFactor.spin must be a string or None.") + if self.dagger is not None and not isinstance(self.dagger, bool): + raise TypeError("OperatorFactor.dagger must be a bool or None.") + + +@dataclass(frozen=True) +class ProductTerm: + """A coefficient times a product of symbolic local operator factors.""" + + coefficient: Any + factors: tuple[OperatorFactor, ...] + + def __post_init__(self): + factors = tuple(self.factors) + if not factors: + raise ValueError("ProductTerm requires at least one operator factor.") + if not all(isinstance(factor, OperatorFactor) for factor in factors): + raise TypeError("ProductTerm.factors must contain OperatorFactor objects.") + object.__setattr__(self, "factors", factors) + + @property + def support(self): + """Sites touched by the term, retaining factor order.""" + return tuple(factor.site for factor in self.factors) + + +@dataclass(frozen=True) +class LocalMatrixTerm: + """A coefficient times a local operator on an explicit site support. + + ``matrix`` uses output axes followed by input axes. Thus a one-site + operator has shape ``(d, d)`` and a two-site operator has shape + ``(d0, d1, d0, d1)``. This is the convention used by the Torch connection + kernel; adapters flatten the two axis groups only where their native + operator API requires a matrix. + """ + + support: tuple[Any, ...] + matrix: Any + coefficient: Any = 1.0 + basis: tuple[Any, ...] | None = None + + def __post_init__(self): + support = tuple(self.support) + if not support: + raise ValueError("LocalMatrixTerm.support cannot be empty.") + if len(set(support)) != len(support): + raise ValueError("LocalMatrixTerm.support must contain unique sites.") + shape = getattr(self.matrix, "shape", None) + if shape is not None: + shape = tuple(int(size) for size in shape) + n_sites = len(support) + if len(shape) != 2 * n_sites: + raise ValueError( + "LocalMatrixTerm.matrix must have output axes followed by " + f"input axes: expected rank {2 * n_sites} for support " + f"{support!r}, got shape {shape!r}." + ) + if any(size <= 0 for size in shape): + raise ValueError("LocalMatrixTerm.matrix dimensions must be positive.") + if shape[:n_sites] != shape[n_sites:]: + raise ValueError( + "LocalMatrixTerm.matrix output and input dimensions must match." + ) + basis = None if self.basis is None else tuple(self.basis) + object.__setattr__(self, "support", support) + object.__setattr__(self, "basis", basis) + + +@dataclass(frozen=True) +class OperatorSum: + """Backend-neutral finite sum of local or symbolic operator terms.""" + + terms: tuple[ProductTerm | LocalMatrixTerm, ...] = () + constant: Any = 0.0 + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self): + terms = tuple(self.terms) + if not all(isinstance(term, (ProductTerm, LocalMatrixTerm)) for term in terms): + raise TypeError( + "OperatorSum.terms must contain ProductTerm or LocalMatrixTerm objects." + ) + object.__setattr__(self, "terms", terms) + object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) + + @classmethod + def from_terms(cls, terms, *, constant=0.0, metadata=None): + """Construct an operator sum while normalizing a list or generator.""" + return cls( + terms=tuple(terms), + constant=constant, + metadata={} if metadata is None else metadata, + ) + + def __iter__(self): + return iter(self.terms) + + def __len__(self): + return len(self.terms) + + @property + def sites(self): + """Sorted-by-first-use tuple of sites appearing in the sum.""" + result = [] + seen = set() + for term in self.terms: + support = term.support + for site in support: + try: + key = site + already_seen = key in seen + except TypeError: + key = repr(site) + already_seen = key in seen + if not already_seen: + seen.add(key) + result.append(site) + return tuple(result) + + +def normalize_operator_sum(value, *, constant=0.0, metadata=None): + """Normalize common legacy term containers to :class:`OperatorSum`. + + A mapping or iterable of ``(support, operator)`` pairs is interpreted as + the legacy local-matrix form used by Torch VMC. A sequence of + :class:`ProductTerm` and :class:`LocalMatrixTerm` objects is preserved. + Backend-native NetKet operators and other opaque objects are intentionally + rejected here; callers should pass those through their backend adapter. + """ + if isinstance(value, OperatorSum): + if constant == 0.0 and metadata is None: + return value + merged_metadata = dict(value.metadata) + if metadata is not None: + merged_metadata.update(metadata) + return OperatorSum( + terms=value.terms, + constant=value.constant + constant, + metadata=merged_metadata, + ) + + if hasattr(value, "terms") and not isinstance(value, (tuple, list, dict)): + value = value.terms + if hasattr(value, "items"): + entries = tuple(value.items()) + else: + try: + entries = tuple(value) + except TypeError as exc: + raise TypeError( + "terms must be an OperatorSum, a sequence of common term " + "objects, or a mapping/iterable of (support, operator) pairs." + ) from exc + + if all(isinstance(term, (ProductTerm, LocalMatrixTerm)) for term in entries): + return OperatorSum( + terms=entries, + constant=constant, + metadata={} if metadata is None else metadata, + ) + + matrix_terms = [] + for entry in entries: + if not isinstance(entry, (tuple, list)) or len(entry) != 2: + raise TypeError( + "legacy terms must contain (support, operator) pairs or common " + "OperatorTerm objects." + ) + support, operator = entry + if isinstance(support, (tuple, list)): + normalized_support = tuple(support) + else: + normalized_support = (support,) + matrix_terms.append( + LocalMatrixTerm(support=normalized_support, matrix=operator) + ) + return OperatorSum( + terms=tuple(matrix_terms), + constant=constant, + metadata={} if metadata is None else metadata, + ) + + +def _canonical_fermion_name(factor): + """Map a symbolic factor to Pepsy's local fermion operator names.""" + name = str(factor.name).lower() + aliases = { + "n": "number", + "occupation": "number", + "number_up": "number_u", + "n_up": "number_u", + "number_down": "number_d", + "n_down": "number_d", + "create_up": "create_u", + "create_down": "create_d", + "annihilate_up": "annihilate_u", + "annihilate_down": "annihilate_d", + "destroy": "annihilate", + "destroy_up": "annihilate_u", + "destroy_down": "annihilate_d", + "doublon": "double", + } + name = aliases.get(name, name) + if name == "fermion": + if factor.dagger is None: + raise ValueError("fermion factors require dagger=True or False.") + action = "create" if factor.dagger else "annihilate" + if factor.spin is None: + return action + spin = str(factor.spin).lower() + if spin in {"up", "u", "↑", "+"}: + return f"{action}_u" + if spin in {"down", "d", "↓", "-"}: + return f"{action}_d" + raise ValueError("fermion factor spin must be 'up', 'down', or None.") + if name in {"create", "annihilate", "number"} and factor.spin is not None: + spin = str(factor.spin).lower() + suffix = "u" if spin in {"up", "u", "↑", "+"} else ( + "d" if spin in {"down", "d", "↓", "-"} else None + ) + if suffix is None: + raise ValueError("fermion factor spin must be 'up' or 'down'.") + return f"{name}_{suffix}" + return name + + +def _expand_fermion_factor(factor): + """Expand number/double aliases into creation/annihilation factors.""" + name = _canonical_fermion_name(factor) + site = factor.site + expansions = { + "number": ("create", "annihilate"), + "number_u": ("create_u", "annihilate_u"), + "number_d": ("create_d", "annihilate_d"), + "double": ( + "create_u", + "annihilate_u", + "create_d", + "annihilate_d", + ), + } + return tuple((site, item) for item in expansions.get(name, (name,))) + + +@dataclass(frozen=True) +class VMCProblem: + """Compatibility bundle for the earlier portable VMC construction API. + + ``hamiltonian`` and entries in ``observables`` may be an :class:`OperatorSum` + or a backend-native object during the compatibility transition. New code + should prefer :class:`MCState` plus :class:`VMC`, while still using + ``OperatorSum`` so Torch and NetKet receive identical terms. + """ + + peps: Any + hamiltonian: Any + observables: Mapping[str, Any] = field(default_factory=dict) + symmetry: str | None = None + site_order: tuple[Any, ...] | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self): + observables = dict(self.observables) + if any(not isinstance(name, str) or not name for name in observables): + raise ValueError("VMCProblem observable names must be non-empty strings.") + site_order = None if self.site_order is None else tuple(self.site_order) + if site_order is not None and len(set(site_order)) != len(site_order): + raise ValueError("VMCProblem.site_order must contain unique sites.") + object.__setattr__(self, "observables", MappingProxyType(observables)) + object.__setattr__(self, "site_order", site_order) + object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) + + +@dataclass(frozen=True, init=False) +class MCState: + """Portable, NetKet-shaped Monte Carlo variational-state recipe. + + ``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. + """ + + peps: Any + sampling: SamplingConfig + symmetry: str | None + site_order: tuple[Any, ...] | None + metadata: Mapping[str, Any] + + def __init__( + self, + peps, + *, + sampling=None, + n_samples=None, + n_chains=None, + n_discard_per_chain=None, + thin=None, + seed=None, + sampler_seed=None, + chunk_size=None, + proposal=None, + symmetry=None, + site_order=None, + metadata=None, + ): + if sampling is not None: + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + supplied = { + "n_samples": n_samples, + "n_chains": n_chains, + "n_discard_per_chain": n_discard_per_chain, + "thin": thin, + "seed": seed, + "sampler_seed": sampler_seed, + "chunk_size": chunk_size, + "proposal": proposal, + } + conflicting = [name for name, value in supplied.items() if value is not None] + if conflicting: + raise ValueError( + "Pass either sampling=... or direct MCState sampling " + f"options, not both; got {', '.join(conflicting)}." + ) + else: + n_chains = 16 if n_chains is None else _positive_int("n_chains", n_chains) + n_samples = 1024 if n_samples is None else _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." + ) + sampling = SamplingConfig( + n_samples_per_chain=n_samples // n_chains, + 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, + seed=seed, + sampler_seed=sampler_seed, + chunk_size=chunk_size, + proposal=proposal, + ) + + site_order = None if site_order is None else tuple(site_order) + if site_order is not None and len(set(site_order)) != len(site_order): + raise ValueError("MCState.site_order must contain unique sites.") + object.__setattr__(self, "peps", peps) + object.__setattr__(self, "sampling", sampling) + object.__setattr__(self, "symmetry", symmetry) + object.__setattr__(self, "site_order", site_order) + object.__setattr__(self, "metadata", MappingProxyType(dict(metadata or {}))) + + @property + def ansatz(self): + """Alias for the variational tensor-network ansatz.""" + return self.peps + + @property + def model(self): + """NetKet-style alias for the variational ansatz.""" + return self.peps + + @property + def n_samples(self): + """Total retained samples over all chains.""" + return self.sampling.n_samples + + @property + def n_chains(self): + """Number of independent Markov chains requested for this state.""" + return self.sampling.n_chains + + @property + def n_discard_per_chain(self): + """Per-chain burn-in, using NetKet's naming convention.""" + return self.sampling.burn_in + + def to_problem(self, hamiltonian, *, observables=None): + """Make the compatibility :class:`VMCProblem` representation.""" + return VMCProblem( + peps=self.peps, + hamiltonian=hamiltonian, + observables={} if observables is None else observables, + symmetry=self.symmetry, + site_order=self.site_order, + metadata=self.metadata, + ) + + +class VMC: + """Portable NetKet-style VMC driver over a Pepsy ``MCState``. + + The class owns a Hamiltonian and a variational state recipe, then exposes + ``sample()``, ``expect()``, and ``run()``. Backend-specific state and + driver controls remain reachable through :attr:`native`. + """ + + def __init__( + self, + hamiltonian, + variational_state, + *, + backend="torch", + fermion=None, + observables=None, + contraction=None, + **build_kwargs, + ): + if not isinstance(variational_state, MCState): + raise TypeError("variational_state must be an MCState.") + backend = str(backend).lower() + if backend not in {"torch", "netket"}: + raise ValueError("backend must be 'torch' or 'netket'.") + + self.hamiltonian = hamiltonian + self.variational_state = variational_state + self.backend = backend + self.problem = variational_state.to_problem( + hamiltonian, + observables=observables, + ) + if backend == "torch": + from .torch import build_torch_vmc + + self._setup = build_torch_vmc( + self.problem, + fermion=fermion, + contraction=contraction, + sampling=variational_state.sampling, + **build_kwargs, + ) + else: + from .netket import build_netket_vmc + + self._setup = build_netket_vmc( + self.problem, + fermion=fermion, + contraction=contraction, + sampling=variational_state.sampling, + **build_kwargs, + ) + + @property + def state(self): + """Alias for :attr:`variational_state`.""" + return self.variational_state + + @property + def native(self): + """The underlying native Torch or NetKet setup.""" + return self._setup.native + + @property + def setup(self): + """Portable backend setup used by this driver.""" + return self._setup + + def sample(self, sampling=None): + """Collect samples through the selected backend.""" + return self._setup.sample(sampling) + + def measure( + self, + observables=None, + *, + sampling=None, + samples=None, + weights=None, + proposal_log_probs=None, + ): + """Measure energy and optional observables. + + ``samples`` optionally supplies a previously collected or external + batch. Torch accepts fixed ``weights`` or ``proposal_log_probs`` for + self-normalized importance sampling; the latter is preferred when the + batch will also be reused for optimization. + """ + kwargs = {"sampling": sampling} + if samples is not None: + kwargs["samples"] = samples + if weights is not None: + kwargs["weights"] = weights + if proposal_log_probs is not None: + kwargs["proposal_log_probs"] = proposal_log_probs + return self._setup.measure(observables, **kwargs) + + def expect( + self, + observable=None, + *, + sampling=None, + samples=None, + weights=None, + proposal_log_probs=None, + ): + """NetKet-style expectation call returning a common measurement. + + With no argument this measures the Hamiltonian. Supplying one + observable adds it under the ``"expectation"`` key alongside energy. + """ + if observable is None: + return self.measure( + sampling=sampling, + samples=samples, + weights=weights, + proposal_log_probs=proposal_log_probs, + ) + return self.measure( + {"expectation": observable}, + sampling=sampling, + samples=samples, + weights=weights, + proposal_log_probs=proposal_log_probs, + ) + + def optimize(self, optimization=None, *, n_steps=None, **kwargs): + """Optimize the variational state through the selected backend.""" + return self._setup.optimize( + optimization=optimization, + n_steps=n_steps, + **kwargs, + ) + + def run(self, n_iter=None, *, optimization=None, **kwargs): + """NetKet-style optimization entry point returning the run history.""" + if optimization is not None: + if n_iter is not None and n_iter != optimization.n_steps: + raise ValueError("n_iter conflicts with optimization.n_steps.") + return self.optimize(optimization=optimization, **kwargs) + if n_iter is None: + raise TypeError("n_iter is required unless optimization is supplied.") + return self.optimize(n_steps=n_iter, **kwargs) + + +@dataclass(frozen=True) +class VMCSamples: + """Backend-neutral Monte Carlo or externally supplied samples. + + ``weights`` defines a fixed weighted empirical batch. Alternatively, + ``proposal_log_probs`` records ``log q(x)`` so a Torch importance estimate + can recompute ``|psi_theta(x)|**2 / q(x)`` as parameters change. + """ + + configs: Any + amplitudes: Any | None = None + log_amplitudes: Any | None = None + n_samples_per_chain: int | None = None + n_chains: int | None = None + acceptance_rate: float | None = None + diagnostics: Mapping[str, Any] = field(default_factory=dict) + native: Any | None = None + weights: Any | None = None + proposal_log_probs: Any | None = None + + def __post_init__(self): + shape = getattr(self.configs, "shape", None) + if shape is not None and len(shape) not in (2, 3): + raise ValueError( + "VMCSamples.configs must have shape (samples, sites) or " + "(samples_per_chain, chains, sites)." + ) + if self.n_samples_per_chain is not None and self.n_samples_per_chain <= 0: + raise ValueError("n_samples_per_chain must be positive when supplied.") + if self.n_chains is not None and self.n_chains <= 0: + raise ValueError("n_chains must be positive when supplied.") + if self.weights is not None and self.proposal_log_probs is not None: + raise ValueError("Pass either weights or proposal_log_probs, not both.") + object.__setattr__(self, "diagnostics", MappingProxyType(dict(self.diagnostics))) + + @property + def chain_shape(self): + """Return the retained ``(samples_per_chain, chains)`` shape.""" + shape = getattr(self.configs, "shape", None) + if shape is not None and len(shape) == 3: + return tuple(int(value) for value in shape[:2]) + if self.n_samples_per_chain is not None or self.n_chains is not None: + return self.n_samples_per_chain, self.n_chains + return None + + +@dataclass(frozen=True) +class VMCMeasurement: + """Backend-neutral observable or local-energy measurement result.""" + + energy_mean: Any | None = None + energy_variance: Any | None = None + energy_stderr: Any | None = None + observables: Mapping[str, Any] = field(default_factory=dict) + local_values: Any | None = None + effective_sample_size: Any | None = None + diagnostics: Mapping[str, Any] = field(default_factory=dict) + native: Any | None = None + + def __post_init__(self): + object.__setattr__(self, "observables", MappingProxyType(dict(self.observables))) + object.__setattr__(self, "diagnostics", MappingProxyType(dict(self.diagnostics))) + + @property + def energy(self): + """Return the backend-neutral sampled energy mean.""" + return self.energy_mean + + @property + def native_energy(self): + """Return the backend-specific energy estimate, when retained.""" + return self.observables.get("energy") + + +@dataclass(frozen=True) +class VMCOptimizationResult: + """Common optimization history for Torch and NetKet VMC runs.""" + + steps: Any + energies: Any + errors: Any + variances: Any | None = None + final_energy: float | None = None + final_error: float | None = None + energy_shift: float = 0.0 + per_site: int | None = None + diagnostics: Mapping[str, Any] = field(default_factory=dict) + native: Any | None = None + + def __post_init__(self): + object.__setattr__(self, "diagnostics", MappingProxyType(dict(self.diagnostics))) + per_site = _positive_int("per_site", self.per_site, allow_none=True) + object.__setattr__(self, "per_site", per_site) + if self.final_energy is None: + values = np.asarray(self.energies) + if values.size: + object.__setattr__(self, "final_energy", float(np.real(values[-1]))) + if self.final_error is None: + values = np.asarray(self.errors) + if values.size: + object.__setattr__(self, "final_error", float(np.real(values[-1]))) + + @property + def shifted_energies(self): + """Return energies with the display-only energy shift applied.""" + return np.asarray(self.energies, dtype=float) + float(self.energy_shift) + + @property + def displayed_energies(self): + """Return shifted energies, divided by ``per_site`` when requested.""" + values = self.shifted_energies + return values if self.per_site is None else values / self.per_site diff --git a/src/pepsy/vmc/netket.py b/src/pepsy/vmc/netket.py index 6f19485..b17d168 100644 --- a/src/pepsy/vmc/netket.py +++ b/src/pepsy/vmc/netket.py @@ -20,6 +20,7 @@ "NetKetChunkSettings", "NetKetPEPSVMC", "NetKetFermiHubbardVMC", + "NetKetVMCSetup", "NetKetSparseFermiHubbardVMC", "NetKetVMCSettings", "PackedPEPS", @@ -28,8 +29,14 @@ "build_heisenberg_vmc", "build_ising_vmc", "build_fermi_hubbard_vmc", + "build_fermion_vmc", + "build_netket_vmc", "build_sparse_fermi_hubbard_vmc", "fermionic_peps_rand", + "fermion_model_terms", + "netket_fermion_operator", + "compile_operator_sum_netket", + "standard_fermion_observables", "choose_netket_chunk_size", "configure_jax_for_vmc", "config_to_phys_indices", @@ -44,9 +51,12 @@ "occupation_to_phys_indices", "pack_peps_ansatz", "pack_fermionic_peps_ansatz", + "prepare_fermionic_peps_for_netket", "recommend_netket_vmc_settings", "square_lattice_edges", "verify_netket_spin_columns", + "VMCOptimizeResult", + "warmup_netket_vmc", ] @@ -179,6 +189,135 @@ class NetKetVMCSettings: notes: tuple[str, ...] +def _make_progress_bar(*, total=None, desc=None, enabled=True): + """Return a notebook-friendly ``tqdm`` bar, or ``None`` when unavailable.""" + if not enabled: + return None + try: + from tqdm.auto import tqdm # pylint: disable=import-outside-toplevel + except ImportError: # pragma: no cover - tqdm ships with netket + return None + return tqdm(total=total, desc=desc, leave=True, dynamic_ncols=True) + + +@dataclass(frozen=True) +class VMCOptimizeResult: + """Energy-optimization history returned by :meth:`NetKetPEPSVMC.optimize`. + + ``energies``/``errors``/``variances`` are per-step Monte-Carlo estimates + (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. + """ + + steps: Any + energies: Any + errors: Any + variances: Any + final_energy: float + final_error: float + compile_seconds: float | None = None + energy_shift: float = 0.0 + + @property + def shifted_energies(self): + """Energies with ``energy_shift`` added.""" + return np.asarray(self.energies, dtype=float) + float(self.energy_shift) + + def plot(self, ax=None, *, per_site=None, reference=None, reference_label=None): + """Plot a clean energy-vs-iteration curve with a Monte-Carlo error band.""" + import matplotlib.pyplot as plt # pylint: disable=import-outside-toplevel + + if ax is None: + _, ax = plt.subplots(figsize=(7.0, 4.5)) + steps = np.asarray(self.steps) + y = self.shifted_energies + band = np.asarray(self.errors, dtype=float) + scale = 1.0 if not per_site else float(per_site) + y = y / scale + band = band / scale + ax.plot(steps, y, "-", color="#1f77b4", lw=1.8, label="VMC energy") + ax.fill_between(steps, y - band, y + band, color="#1f77b4", alpha=0.2) + if reference is not None: + ax.axhline( + float(reference), + ls="--", + color="k", + lw=1.4, + label=reference_label or "reference", + ) + ax.set_xlabel("iteration", fontsize=13) + ax.set_ylabel("energy / site" if per_site else "energy", fontsize=13) + ax.grid(alpha=0.3) + ax.legend(fontsize=11) + return ax + + +class _VMCProgressCallback: + """NetKet callback: one live ``tqdm`` energy bar that records the history.""" + + def __init__(self, n_iter, *, enabled=True, energy_shift=0.0, per_site=None): + self.n_iter = int(n_iter) + self.enabled = enabled + self.energy_shift = float(energy_shift) + self.per_site = per_site + self._bar = None + self.steps = [] + self.energies = [] + self.errors = [] + self.variances = [] + + def _ensure_bar(self): + if self._bar is None and self.enabled: + self._bar = _make_progress_bar( + total=self.n_iter, desc="VMC energy", enabled=True + ) + return self._bar + + def __call__(self, step, log_data, driver): + name = getattr(driver, "_loss_name", "Energy") + stats = None + if isinstance(log_data, dict): + stats = log_data.get(name) + if stats is None: + stats = getattr(driver, "_loss_stats", None) + if stats is not None: + 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)) + self.steps.append(int(step)) + self.energies.append(mean) + self.errors.append(err) + self.variances.append(var) + bar = self._ensure_bar() + 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}") + bar.update(1) + return True + + def close(self): + if self._bar is not None: + self._bar.close() + self._bar = None + + def result(self, compile_seconds=None): + energies = np.asarray(self.energies, dtype=float) + errors = np.asarray(self.errors, dtype=float) + return VMCOptimizeResult( + steps=np.asarray(self.steps, dtype=int), + energies=energies, + errors=errors, + variances=np.asarray(self.variances, dtype=float), + 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, + energy_shift=self.energy_shift, + ) + + @dataclass(frozen=True) class NetKetPEPSVMC: """Bundle returned by generic PEPS/NetKet VMC builders.""" @@ -214,12 +353,383 @@ def make_driver(self, **kwargs): """ return make_netket_vmc_driver(self, **kwargs) + def warmup(self, *, progress=True, hamiltonian=None): + """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) + + def sample(self, sampling=None): + """Collect samples using the shared :class:`SamplingConfig` contract. + + NetKet stores samples as ``(n_chains, n_samples_per_chain, sites)``; + this façade canonicalizes them to the backend-neutral + ``(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. + """ + 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 requested_chunk is not None: + self.vstate.chunk_size = old_chunk + + shape = getattr(native, "shape", ()) + if len(shape) == 3: + # NetKet's native order is chains first. + configs = native.swapaxes(0, 1) + n_samples_per_chain, n_chains = int(shape[1]), int(shape[0]) + elif len(shape) == 2: + configs = native.reshape((int(shape[0]), 1, int(shape[1]))) + n_samples_per_chain, n_chains = int(shape[0]), 1 + else: + raise ValueError( + "NetKet samples must have shape (chains, samples, sites) or " + f"(samples, sites), got {shape}." + ) + return VMCSamples( + configs=configs, + n_samples_per_chain=n_samples_per_chain, + n_chains=n_chains, + native=native, + ) + + def optimize( + self, + n_iter=None, + *, + optimization=None, + learning_rate=None, + driver=None, + optimizer=None, + progress=None, + warmup=None, + energy_shift=None, + per_site=None, + extra_callbacks=None, + driver_options=None, + **run_kwargs, + ): + """Run VMC energy optimization with a single clean progress bar. + + Builds a driver (see :func:`make_netket_vmc_driver`), optionally warms + up XLA compilation (:meth:`warmup`), then runs ``n_iter`` steps while a + live ``tqdm`` bar shows the current energy. NetKet's own bar is + disabled so there is exactly one bar. + + ``energy_shift``/``per_site`` only affect the displayed and returned + energies (for example ``energy_shift=-U/4`` and ``per_site=n_sites`` + for the Fermi-Hubbard convention); the raw Monte-Carlo means are kept. + Extra NetKet callbacks (early stopping, timeout, ...) can be passed via + ``extra_callbacks``. Returns a :class:`VMCOptimizeResult`. + """ + if optimization is not None: + from .api import OptimizationConfig, VMCBackendCapabilityError + if not isinstance(optimization, OptimizationConfig): + raise TypeError("optimization must be an OptimizationConfig or None.") + if optimization.method == "minsr": + raise VMCBackendCapabilityError( + "OptimizationConfig(method='minsr') is Torch-specific. " + "Use method='sr' for portable stochastic reconfiguration, " + "or call make_netket_vmc_driver(driver='vmc_sr') directly." + ) + if n_iter is not None and n_iter != optimization.n_steps: + raise ValueError("n_iter conflicts with optimization.n_steps.") + n_iter = optimization.n_steps + if learning_rate is None: + learning_rate = optimization.learning_rate + if progress is None: + progress = optimization.progress + if warmup is None: + warmup = optimization.warmup + if energy_shift is None: + energy_shift = optimization.energy_shift + if per_site is None: + per_site = optimization.per_site + if driver is None: + driver = "vmc" if optimization.method == "sgd" else "vmc_sr" + driver_options = {} if driver_options is None else dict(driver_options) + if driver != "vmc": + driver_options.setdefault("sr_mode", optimization.sr_mode) + driver_options.setdefault("sr_diag_shift", optimization.diag_shift) + if n_iter is None: + raise TypeError("n_iter is required unless optimization is supplied.") + if learning_rate is None: + learning_rate = 0.02 + if driver is None: + driver = "vmc" + if progress is None: + progress = True + if warmup is None: + warmup = True + if energy_shift is None: + energy_shift = 0.0 + driver_options = {} if driver_options is None else dict(driver_options) + run_driver = make_netket_vmc_driver( + self, + optimizer=optimizer, + learning_rate=learning_rate, + driver=driver, + **driver_options, + ) + 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) + try: + run_driver.run( + n_iter, + show_progress=False, + callback=callbacks, + **run_kwargs, + ) + finally: + cb.close() + return cb.result(compile_seconds=compile_seconds) + + def measure(self, observables=None): + """Measure observables on the current variational state. + + ``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. + """ + if observables is None: + observables = getattr(self, "observables", None) + if observables is None: + raise ValueError( + "No observables to measure: pass observables=... or build the " + "setup with observables=... (see build_fermion_vmc)." + ) + if isinstance(observables, dict): + return { + name: self.vstate.expect(op) + for name, op in observables.items() + } + return self.vstate.expect(observables) + + +@dataclass(frozen=True) +class NetKetVMCSetup: + """Backend-neutral façade over a native :class:`NetKetPEPSVMC` setup. + + The wrapped setup keeps its NetKet drivers, state, and statistics objects + available through :attr:`native`. This façade is deliberately strict about + shared settings which NetKet cannot honour after construction, rather than + warning and silently changing their meaning. + """ + + setup: NetKetPEPSVMC + problem: Any + + @property + def backend(self): + """Name of the numerical backend behind this setup.""" + return "netket" + + @property + def native(self): + """Return the native NetKet/Flax setup for advanced control.""" + return self.setup + + @property + def n_sites(self): + return self.setup.n_sites + + @property + def n_params(self): + return self.setup.n_params + + def sample(self, sampling=None): + """Collect samples as backend-neutral :class:`VMCSamples`. + + Chain count and seeds belong to an MCState at build time in NetKet. + The portable sampling call therefore rejects requests which would be + ignored by an existing state. + """ + from .api import SamplingConfig, VMCBackendCapabilityError + + if sampling is None: + return self.setup.sample() + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + unsupported = [] + if sampling.thin != 1: + unsupported.append("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: + unsupported.append("proposal") + if unsupported: + joined = ", ".join(unsupported) + raise VMCBackendCapabilityError( + f"NetKet cannot apply SamplingConfig.{joined} after setup " + "construction. Supply it while building the native sampler, " + "or use the native NetKet API explicitly." + ) + return self.setup.sample(sampling) + + def measure( + self, + observables=None, + *, + sampling=None, + samples=None, + weights=None, + proposal_log_probs=None, + ): + """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: + 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." + ) + + samples = self.sample(sampling) if sampling is not None else None + if observables is None: + extra = dict(getattr(self.setup, "observables", None) or {}) + else: + try: + extra = dict(observables) + except (TypeError, ValueError) as exc: + raise TypeError("observables must be a mapping of names to operators.") from exc + if "energy" in extra: + raise ValueError( + "'energy' is reserved for problem.hamiltonian; use a different " + "observable name." + ) + native = self.setup.measure({"energy": self.setup.hamiltonian, **extra}) + energy = native["energy"] + return VMCMeasurement( + energy_mean=getattr(energy, "mean", energy), + energy_variance=getattr(energy, "variance", None), + energy_stderr=getattr(energy, "error_of_mean", None), + observables=native, + effective_sample_size=getattr(energy, "effective_sample_size", None), + diagnostics={ + "backend": self.backend, + "samples": samples, + }, + native=native, + ) + + def optimize(self, optimization=None, *, n_steps=None, **kwargs): + """Optimize and return a backend-neutral VMC history.""" + from .api import ( + OptimizationConfig, + VMCBackendCapabilityError, + VMCOptimizationResult, + ) + + if any( + kwargs.get(name) is not None + for name in ("samples", "weights", "proposal_log_probs") + ): + raise VMCBackendCapabilityError( + "NetKet's portable adapter does not yet optimize from an " + "externally supplied weighted sample batch. Use its MCState " + "sampling path, or use the Torch adapter for importance sampling." + ) + + if optimization is not None and not isinstance(optimization, OptimizationConfig): + raise TypeError("optimization must be an OptimizationConfig or None.") + if optimization is not None and optimization.method == "minsr": + raise VMCBackendCapabilityError( + "OptimizationConfig(method='minsr') is Torch-specific. " + "Use method='sr' for the portable NetKet path." + ) + if optimization is not None: + if n_steps is not None and n_steps != optimization.n_steps: + raise ValueError("n_steps conflicts with optimization.n_steps.") + native = self.setup.optimize(optimization=optimization, **kwargs) + energy_shift = optimization.energy_shift + per_site = optimization.per_site + else: + if n_steps is None: + raise TypeError("n_steps is required unless optimization is supplied.") + native = self.setup.optimize(n_steps, **kwargs) + energy_shift = native.energy_shift + per_site = None + return VMCOptimizationResult( + steps=native.steps, + energies=native.energies, + errors=native.errors, + variances=native.variances, + final_energy=native.final_energy, + final_error=native.final_error, + energy_shift=energy_shift, + per_site=per_site, + diagnostics={ + "backend": self.backend, + "compile_seconds": native.compile_seconds, + }, + native=native, + ) + @dataclass(frozen=True) class NetKetFermiHubbardVMC(NetKetPEPSVMC): """Bundle returned by :func:`build_fermi_hubbard_vmc`.""" columns: SpinOrbitalColumns | None = None + observables: dict | None = None @dataclass(frozen=True) @@ -289,6 +799,72 @@ def energy_estimate(self, configs=None, amplitudes=None): """Return the Monte-Carlo mean of :meth:`local_energy`.""" return self.local_energy(configs=configs, amplitudes=amplitudes).mean() + def sample(self, sampling=None): + """Collect chain-preserving samples with the torch sparse-block path. + + This is the non-jitted counterpart of :meth:`NetKetPEPSVMC.sample`. + It is intended for ``U1``/``U1U1`` Symmray PEPS that cannot use + NetKet's jitted ``MCState`` amplitude kernel. + """ + from .api import BackendCapabilityWarning, SamplingConfig + from .torch import TorchMetropolisSampler + + if sampling is None: + sampling = SamplingConfig( + n_samples_per_chain=128, + n_chains=int(self.configs.shape[0]), + burn_in=0, + thin=1, + ) + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if sampling.proposal is not None: + warnings.warn( + "SamplingConfig.proposal is ignored by the sparse fermion " + "sampler; its symmetry-preserving spinful proposal is fixed.", + BackendCapabilityWarning, + stacklevel=2, + ) + n_chains = sampling.n_chains + if self.configs.shape[0] == 1 and n_chains > 1: + initial_configs = self.configs.expand(n_chains, -1).clone() + initial_amplitudes = self.amplitudes.expand(n_chains).clone() + elif self.configs.shape[0] >= n_chains: + initial_configs = self.configs[:n_chains].clone() + initial_amplitudes = self.amplitudes[:n_chains].clone() + else: + raise ValueError( + "sampling.n_chains exceeds the sparse setup's initial walker " + f"count ({self.configs.shape[0]}). Rebuild with n_samples >= " + "the requested chain count." + ) + if sampling.seed is not None and sampling.sampler_seed is not None: + raise ValueError("Pass either sampling.seed or sampling.sampler_seed, not both.") + sampler = TorchMetropolisSampler( + self.model, + self.torch_graph, + initial_configs, + amplitudes=initial_amplitudes, + n_chains=n_chains, + proposal="spinful", + encoding=self.encoding, + chunk_size=sampling.chunk_size, + seed=( + sampling.seed + if sampling.seed is not None + else sampling.sampler_seed + ), + ) + result = sampler.sample( + n_samples=sampling.n_samples, + n_discard_per_chain=sampling.burn_in, + n_thin=sampling.thin, + ) + object.__setattr__(self, "configs", sampler.configs) + object.__setattr__(self, "amplitudes", sampler.amplitudes) + object.__setattr__(self, "generator", sampler.generator) + return result.to_common() + def sample_sweep( self, configs=None, @@ -430,10 +1006,119 @@ def _validate_contraction(name, contraction, chi): return contraction +def _require_static_cutoff_for_jit(name, contraction, cutoff): + """Reject a nonzero ``cutoff`` on an approximate jitted contraction. + + Approximate boundary/HOTRG/CTMRG contractions keep + ``min(max_bond, rank-above-cutoff)`` singular values. Under + ``jax.jit``/``jax.vmap`` that surviving count is data-dependent, so a + nonzero ``cutoff`` produces dynamic tensor shapes that XLA cannot trace. A + fixed ``max_bond`` with ``cutoff=0.0`` keeps every intermediate shape + static; use the ``jit=False`` builder for adaptive truncation. + """ + if contraction in {"hotrg", "ctmrg", "boundary"} and float(cutoff) != 0.0: + raise ValueError( + f"{name}={contraction!r} under jax.jit requires cutoff=0.0 " + f"(data-dependent truncation is not traceable); got " + f"cutoff={cutoff!r}. Use cutoff=0.0 with a fixed max_bond=chi, or " + "build the amplitude with jit=False for adaptive truncation." + ) + + +def _maybe_register_stable_jax_svd(contraction): + """Install Pepsy's regularized 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 + is a global autoray side effect and a soft no-op when JAX is unavailable. + """ + from .api import ContractionConfig + if isinstance(contraction, ContractionConfig): + contraction = contraction.method + key = str(contraction).replace("_", "-").lower() + if _CONTRACTION_ALIASES.get(key) not in {"hotrg", "ctmrg", "boundary"}: + return + try: + from ..tensors import reg_rel_svd_jax # pylint: disable=import-outside-toplevel + + reg_rel_svd_jax() + except ImportError: + return + + def _contraction_options(contraction_opts): return {} if contraction_opts is None else dict(contraction_opts) +def _resolve_netket_contraction(contraction, chi, cutoff, contraction_opts): + """Resolve a common ContractionConfig for NetKet public builders.""" + from .api import ContractionConfig + if not isinstance(contraction, ContractionConfig): + return contraction, chi, cutoff, contraction_opts + if chi is not None and chi != contraction.chi: + raise ValueError( + f"chi={chi} conflicts with ContractionConfig.chi={contraction.chi}." + ) + if contraction_opts is not None and dict(contraction_opts) != dict(contraction.options): + raise ValueError("contraction options conflict with ContractionConfig.options.") + return ( + contraction.method, + contraction.chi, + contraction.cutoff, + dict(contraction.options), + ) + + +def _resolve_sampling_build_config( + sampling, + *, + n_samples, + n_chains, + n_discard_per_chain, + chunk_size, + seed, + sampler_seed, +): + """Apply a common SamplingConfig before constructing a NetKet sampler.""" + from .api import SamplingConfig + if sampling is None: + return ( + n_samples, + n_chains, + n_discard_per_chain, + chunk_size, + seed, + sampler_seed, + ) + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if seed is not None and sampling.seed is not None and seed != sampling.seed: + raise ValueError("seed conflicts with sampling.seed.") + if ( + sampler_seed is not None + and sampling.sampler_seed is not None + and sampler_seed != sampling.sampler_seed + ): + raise ValueError("sampler_seed conflicts with sampling.sampler_seed.") + if seed is not None and sampling.sampler_seed is not None: + raise ValueError("Pass either seed or sampling.sampler_seed, not both.") + if sampler_seed is not None and sampling.seed is not None: + raise ValueError("Pass either sampler_seed or sampling.seed, not both.") + if sampling.seed is not None and sampling.sampler_seed is not None: + raise ValueError("Pass either sampling.seed or sampling.sampler_seed, not both.") + return ( + sampling.n_samples, + sampling.n_chains, + sampling.burn_in, + 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, + ) + + def _coerce_config_map(config_map): if config_map is None: return NetKetLocalConfigMap.spin_half() @@ -464,6 +1149,105 @@ def _uses_flat_symmray_arrays(tn): return True if symmray_seen else None +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) + + +def _z2_flat_padded(data, sr): + """Flatten a Z2 array whose charge blocks have unequal truncated sizes. + + Symmray's regular ``to_flat`` path stacks charge blocks directly. A PEPS + simple update can leave those blocks with different shapes after an SVD + truncation, so the blocks need zero-padding to their flattened index sizes + before they can share one dense JAX array. Padding is exact: the newly + added entries are outside the original charge-block support. + """ + flat_indices = tuple(index.to_flat() for index in data.indices) + target_shape = tuple(index.charge_size for index in flat_indices) + padded_blocks = {} + for sector, block in data.blocks.items(): + array = _host_array_for_flatten(block) + if len(array.shape) != len(target_shape): + raise ValueError( + "Cannot flatten a Z2 block with rank " + f"{array.ndim}; expected {len(target_shape)}." + ) + if any(got > wanted for got, wanted in zip(array.shape, target_shape)): + raise ValueError( + "Cannot zero-pad a Z2 block larger than its flattened index: " + f"block={array.shape}, target={target_shape}." + ) + if tuple(array.shape) != target_shape: + padded = np.zeros(target_shape, dtype=array.dtype) + padded[tuple(slice(0, size) for size in array.shape)] = array + array = padded + padded_blocks[sector] = array + return sr.Z2FermionicArrayFlat.from_blocks( + padded_blocks, + flat_indices, + phases=data.phases, + label=data.label, + dummy_modes=data.dummy_modes, + symmetry=data.symmetry, + ) + + +def prepare_fermionic_peps_for_netket(peps, *, device=None): + """Prepare an evolved fermionic PEPS for NetKet's jitted JAX model. + + The returned copy uses the natural ``qtn.pack`` representation consumed by + :func:`pack_fermionic_peps_ansatz`. Z2 block-sparse tensors are flattened + here, including the zero-padding required after unequal simple-update SVD + 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 + tensors only go through the JAX backend conversion. + """ + work = peps.copy() + tn = getattr(work, "tn", work) + padded_sites = [] + sr = None + + for site in tuple(tn.sites): + data = tn[site].data + type_name = type(data).__name__ + if ( + _is_symmray_array(data) + and "Z2" in type_name + and "Flat" not in type_name + ): + if sr is None: + sr = _require_symmray() + try: + converted = data.to_flat() + except RuntimeError: + converted = _z2_flat_padded(data, sr) + padded_sites.append(site) + tn[site].modify(data=converted) + + if padded_sites: + warnings.warn( + "Zero-padded unequal Z2 charge blocks while preparing " + f"{len(padded_sites)} PEPS tensor(s) for NetKet.", + RuntimeWarning, + stacklevel=2, + ) + + from ..tensors import backend_jax # pylint: disable=import-outside-toplevel + + if device is None: + device = os.environ.get("PEPSY_FH_JAX_DEVICE") + work.apply_to_arrays(backend_jax(device=device, dtype=None)) + return work + + def _param_leaf_size(leaf): size = getattr(leaf, "size", None) if size is not None: @@ -560,6 +1344,37 @@ def _require_jittable_fermionic_ansatz(ansatz): ) +def _warn_flat_z2_ansatz_fixed_u1u1_sector(ansatz, n_fermions_per_spin): + """Make the supported flat-Z2/fixed-U1U1 NetKet configuration explicit.""" + phys_charges = tuple(getattr(ansatz, "phys_charges", ()) or ()) + is_flat_z2 = ( + getattr(ansatz, "uses_flat_symmray", None) is True + # Current Symmray exposes flat fermionic arrays only for Z2. FlatIndex + # does not retain a ``chargemap``, hence ``phys_charges`` is empty for + # normal flat Z2 PEPS after packing. + and ( + not phys_charges + or all(charge in (0, 1) for charge in phys_charges) + ) + ) + if not is_flat_z2: + return + + n_up, n_down = (int(value) for value in n_fermions_per_spin) + from .api import SymmetryFallbackWarning + + warnings.warn( + "NetKet VMC is using a flat Z2 fermionic PEPS ansatz with the fixed " + f"U1U1 sampling sector (N_up, N_down)=({n_up}, {n_down}). This is " + "intentional and supported: Z2 is the JAX-friendly tensor-storage " + "symmetry, while NetKet's SpinOrbitalFermions Hilbert space and " + "MetropolisFermionHop sampler preserve N_up and N_down separately. " + "The PEPS itself does not carry block-sparse U1U1 charges.", + SymmetryFallbackWarning, + stacklevel=3, + ) + + def _make_torch_generator(seed, device=None): if seed is None: return None @@ -787,6 +1602,161 @@ def site_index(i, j): return tuple(edges) +def _infer_lattice_shape_from_peps(peps): + """Infer a rectangular ``(Lx, Ly)`` shape from a PEPS-like object.""" + tn = getattr(peps, "tn", peps) + shape = (getattr(tn, "Lx", None), getattr(tn, "Ly", None)) + if all(value is not None for value in shape): + return tuple(int(value) for value in shape) + + sites = tuple(getattr(tn, "sites", ())) + if not sites or not all( + isinstance(site, tuple) and len(site) == 2 for site in sites + ): + raise ValueError( + "Could not infer a rectangular PEPS lattice; provide both Lx and Ly." + ) + Lx = max(int(site[0]) for site in sites) + 1 + Ly = max(int(site[1]) for site in sites) + 1 + expected = {(i, j) for i in range(Lx) for j in range(Ly)} + if set(sites) != expected: + raise ValueError( + "PEPS sites are not a complete row-major rectangular lattice; " + "provide both Lx and Ly." + ) + return Lx, Ly + + +def _site_index_for_lattice(site, Lx, Ly): + if isinstance(site, Integral): + site = int(site) + return site if 0 <= site < Lx * Ly else None + try: + i, j = tuple(site) + except (TypeError, ValueError): + return None + i, j = int(i), int(j) + if 0 <= i < Lx and 0 <= j < Ly: + return i * Ly + j + return None + + +def _is_coordinate_edge_key(key): + try: + left, right = tuple(key) + return ( + isinstance(left, (tuple, list)) + and isinstance(right, (tuple, list)) + and len(left) == 2 + and len(right) == 2 + ) + except (TypeError, ValueError): + return False + + +def _edges_from_fermi_terms(terms, Lx, Ly): + """Extract integer graph edges from native coordinate-keyed terms.""" + if terms is None or not hasattr(terms, "keys"): + return () + keys = tuple(terms.keys()) + coordinate_edges = tuple(key for key in keys if _is_coordinate_edge_key(key)) + use_integer_edges = not coordinate_edges + found = [] + for key in keys: + try: + left, right = tuple(key) + except (TypeError, ValueError): + continue + if coordinate_edges and not _is_coordinate_edge_key(key): + continue + if not coordinate_edges and not use_integer_edges: + continue + left_index = _site_index_for_lattice(left, Lx, Ly) + right_index = _site_index_for_lattice(right, Lx, Ly) + if left_index is None or right_index is None or left_index == right_index: + continue + edge = (min(left_index, right_index), max(left_index, right_index)) + if edge not in found: + found.append(edge) + return tuple(found) + + +def _edges_from_operator_sum(operator_sum, Lx, Ly): + """Extract NetKet integer edges from the common operator IR.""" + sites = _row_major_sites(int(Lx), int(Ly)) + site_to_index = {site: index for index, site in enumerate(sites)} + found = [] + for term in operator_sum: + support = tuple(getattr(term, "support", ())) + if len(support) != 2: + continue + mapped = [] + for site in support: + if site in site_to_index: + mapped.append(site_to_index[site]) + elif isinstance(site, Integral) and 0 <= int(site) < len(sites): + mapped.append(int(site)) + else: + mapped = [] + break + if len(mapped) == 2 and mapped[0] != mapped[1]: + edge = tuple(sorted(mapped)) + if edge not in found: + found.append(edge) + return tuple(found) + + +def _infer_pbc_from_fermi_terms(terms, Lx, Ly): + """Infer periodic axes from coordinate-keyed nearest-neighbor terms.""" + if terms is None or not hasattr(terms, "keys"): + return False + pbc_x = False + pbc_y = False + for key in terms.keys(): + if not _is_coordinate_edge_key(key): + continue + (i0, j0), (i1, j1) = key + 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 + return (pbc_x, pbc_y) + + +def _coerce_spin_sector(value): + """Convert common Pepsy sector/setup forms to ``(N_up, N_down)``.""" + if value is None: + return None + if hasattr(value, "spin_occupations"): + value = value.spin_occupations + elif hasattr(value, "n_fermions_per_spin"): + value = value.n_fermions_per_spin + + if isinstance(value, dict): + if {"n_up", "n_down"}.issubset(value): + return int(value["n_up"]), int(value["n_down"]) + if "n_fermions_per_spin" in value: + return _coerce_spin_sector(value["n_fermions_per_spin"]) + values = tuple(value.values()) + if values and all( + isinstance(item, (tuple, list)) and len(item) == 2 + for item in values + ): + return tuple( + int(sum(int(item[spin]) for item in values)) + for spin in (0, 1) + ) + + try: + values = tuple(value) + except TypeError: + return None + if len(values) != 2 or any( + isinstance(item, (tuple, list, dict)) for item in values + ): + return None + return tuple(int(item) for item in values) + + def netket_spin_orbital_columns(hilbert): """Return NetKet occupation columns for ``(orbital, spin)`` modes.""" n_orbitals = int(getattr(hilbert, "n_orbitals")) @@ -839,6 +1809,41 @@ def match_col(op): return columns +def _site_major_to_netket_jw_phase( + occ_rows, + columns, + *, + site_to_orb=None, + xp=np, +): + """Return the fermionic basis phase from PEPS to NetKet mode order. + + A fermionic PEPS contraction uses site-major local modes + (up_0, down_0, up_1, down_1, ...). NetKet's SpinOrbitalFermions stores + configurations as (down_0, ..., down_n, up_0, ..., up_n). Reordering the + occupied creation operators gives the phase + + (-1) ** sum_j(n_down[j] * sum_{i <= j} n_up[i]). + + site_to_orb changes NetKet's orbital order into the packed PEPS site order + before the phase is evaluated. + """ + n_orbitals = len(columns.up) + occ_rows = xp.asarray(occ_rows, dtype=xp.int32).reshape( + (-1, 2 * n_orbitals) + ) + col_up = xp.asarray(columns.up, dtype=xp.int32) + col_down = xp.asarray(columns.down, dtype=xp.int32) + n_up = occ_rows[:, col_up] + n_down = occ_rows[:, col_down] + if site_to_orb is not None: + site_to_orb = xp.asarray(site_to_orb, dtype=xp.int32) + n_up = n_up[:, site_to_orb] + n_down = n_down[:, site_to_orb] + inversions = xp.sum(xp.cumsum(n_up, axis=1) * n_down, axis=1) + return 1 - 2 * (inversions % 2) + + def occupation_to_phys_indices(occ_rows, columns, *, site_to_orb=None, phys_charges=None): """Map NetKet spin-orbital occupations to Symmray spinful physical indices. @@ -952,7 +1957,12 @@ def pack_peps_ansatz(peps, *, lattice_shape=None, config_sites=None): def pack_fermionic_peps_ansatz(peps, *, lattice_shape=None, orbital_sites=None): - """Pack a Symmray/quimb fermionic PEPS for use as a Flax pytree.""" + """Pack a fermionic PEPS using quimb's natural params/skeleton pair. + + This is the package equivalent of ``params, skeleton = qtn.pack(peps)``; + the returned ansatz retains ``skeleton`` so every JAX/Flax evaluation can + reconstruct it with ``qtn.unpack(params, skeleton)``. + """ return _pack_peps_ansatz( peps, lattice_shape=lattice_shape, @@ -974,6 +1984,7 @@ def _make_peps_batched_amplitude_apply( contraction = _validate_contraction("contraction", contraction, chi) if output not in {"log", "amplitude", "mantissa_exponent"}: raise ValueError("output must be 'log', 'amplitude', or 'mantissa_exponent'.") + _require_static_cutoff_for_jit("contraction", contraction, cutoff) if ansatz.uses_flat_symmray is False: warnings.warn( "JAX-jitted Symmray PEPS amplitudes are fastest with flat=True " @@ -1167,6 +2178,12 @@ def make_peps_batched_amplitude_function( jit=True, ): """Return a batched JAX amplitude function for NetKet local configs.""" + contraction, chi, cutoff, contraction_opts = _resolve_netket_contraction( + contraction, + chi, + cutoff, + contraction_opts, + ) jax, _ = _require_jax() config_map = _coerce_config_map(config_map) if jit: @@ -1210,6 +2227,12 @@ def make_peps_log_amplitude_model( param_dtype=None, ): """Build a Flax model returning batched ``log(psi(config_row))``.""" + contraction, chi, cutoff, contraction_opts = _resolve_netket_contraction( + contraction, + chi, + cutoff, + contraction_opts, + ) _validate_contraction("contraction", contraction, chi) config_map = _coerce_config_map(config_map) @@ -1262,6 +2285,7 @@ def _make_fermionic_peps_batched_amplitude_apply( contraction = _validate_contraction("contraction", contraction, chi) if output not in {"log", "amplitude", "mantissa_exponent"}: raise ValueError("output must be 'log', 'amplitude', or 'mantissa_exponent'.") + _require_static_cutoff_for_jit("contraction", contraction, cutoff) method_opts = _contraction_options(contraction_opts) if contraction == "boundary": @@ -1286,7 +2310,13 @@ def occ_rows_to_phys_jax(occ_rows): phys_orb = phys_lut[nu, nd] else: phys_orb = 2 * (nu != nd).astype(jnp.int32) + nd - return phys_orb[:, site_to_orb] + phase = _site_major_to_netket_jw_phase( + occ_rows, + columns, + site_to_orb=site_to_orb, + xp=jnp, + ) + return phys_orb[:, site_to_orb], phase def select_phys(tn, phys): if site_inds: @@ -1335,7 +2365,7 @@ def amplitude_from_mantissa_exponent(mantissa, exponent): def apply(occ_rows, params): tn = qtn.unpack(params, ansatz.skeleton) - phys_rows = occ_rows_to_phys_jax(occ_rows) + phys_rows, phase = occ_rows_to_phys_jax(occ_rows) def evaluate_phys(phys): mantissa, exponent = contract_mantissa_exponent(select_phys(tn, phys)) @@ -1345,7 +2375,13 @@ def evaluate_phys(phys): return amplitude_from_mantissa_exponent(mantissa, exponent) return log_from_mantissa_exponent(mantissa, exponent) - return jax.vmap(evaluate_phys)(phys_rows) + values = jax.vmap(evaluate_phys)(phys_rows) + if output == "mantissa_exponent": + mantissas, exponents = values + return mantissas * phase.astype(mantissas.dtype), exponents + if output == "amplitude": + return values * phase.astype(values.dtype) + return values + jnp.log(phase.astype(complex_dtype)) return apply @@ -1422,20 +2458,30 @@ def evaluate_one(tn, phys): def apply(occ_rows, params): tn = qtn.unpack(params, ansatz.skeleton) + occ_rows = np.asarray(occ_rows) phys_rows = occupation_to_phys_indices( - np.asarray(occ_rows), + occ_rows, columns, site_to_orb=ansatz.site_to_orb, phys_charges=getattr(ansatz, "phys_charges", ()), ) + phase = _site_major_to_netket_jw_phase( + occ_rows, + columns, + site_to_orb=ansatz.site_to_orb, + ) values = [evaluate_one(tn, phys) for phys in phys_rows] if output == "mantissa_exponent": mantissas, exponents = zip(*values) + mantissas = jnp.stack([jnp.asarray(x) for x in mantissas]) return ( - jnp.stack([jnp.asarray(x) for x in mantissas]), + mantissas * jnp.asarray(phase, dtype=mantissas.dtype), jnp.asarray(exponents, dtype=real_dtype), ) - return jnp.stack([jnp.asarray(x) for x in values]) + values = jnp.stack([jnp.asarray(x) for x in values]) + if output == "amplitude": + return values * jnp.asarray(phase, dtype=values.dtype) + return values + jnp.log(jnp.asarray(phase, dtype=complex_dtype)) return apply @@ -1469,6 +2515,12 @@ def make_fermionic_peps_batched_amplitude_function( amplitudes or ``output="amplitude"`` for direct scalar amplitudes on tiny/exact contractions. """ + contraction, chi, cutoff, contraction_opts = _resolve_netket_contraction( + contraction, + chi, + cutoff, + contraction_opts, + ) jax, _ = _require_jax() if jit: _require_jittable_fermionic_ansatz(ansatz) @@ -1512,6 +2564,12 @@ def make_fermionic_peps_log_amplitude_model( param_dtype=None, ): """Build a Flax model returning batched ``log(psi(occupation_row))``.""" + contraction, chi, cutoff, contraction_opts = _resolve_netket_contraction( + contraction, + chi, + cutoff, + contraction_opts, + ) _validate_contraction("contraction", contraction, chi) _require_jittable_fermionic_ansatz(ansatz) @@ -1667,6 +2725,7 @@ def build_ising_vmc( n_chains=16, n_discard_per_chain=32, chunk_size=256, + sampling=None, sampler_chunk_size=None, seed=None, sampler_seed=None, @@ -1680,6 +2739,22 @@ def build_ising_vmc( param_dtype=None, ): """Create NetKet VMC objects for a spin-1/2 transverse-field Ising PEPS.""" + ( + n_samples, + n_chains, + n_discard_per_chain, + chunk_size, + seed, + sampler_seed, + ) = _resolve_sampling_build_config( + sampling, + n_samples=n_samples, + n_chains=n_chains, + n_discard_per_chain=n_discard_per_chain, + chunk_size=chunk_size, + seed=seed, + sampler_seed=sampler_seed, + ) nk = _require_netket() n_sites = int(Lx) * int(Ly) hilbert = nk.hilbert.Spin(s=1 / 2, N=n_sites) @@ -1752,6 +2827,7 @@ def build_heisenberg_vmc( n_chains=16, n_discard_per_chain=32, chunk_size=256, + sampling=None, sampler_chunk_size=None, seed=None, sampler_seed=None, @@ -1765,6 +2841,22 @@ def build_heisenberg_vmc( param_dtype=None, ): """Create NetKet VMC objects for a spin-1/2 Heisenberg PEPS.""" + ( + n_samples, + n_chains, + n_discard_per_chain, + chunk_size, + seed, + sampler_seed, + ) = _resolve_sampling_build_config( + sampling, + n_samples=n_samples, + n_chains=n_chains, + n_discard_per_chain=n_discard_per_chain, + chunk_size=chunk_size, + seed=seed, + sampler_seed=sampler_seed, + ) nk = _require_netket() n_sites = int(Lx) * int(Ly) hilbert = nk.hilbert.Spin(s=1 / 2, N=n_sites, total_sz=total_sz) @@ -1933,25 +3025,356 @@ def fermionic_peps_rand( ) +def _maybe_conserving_fermion_operator(operator, *, strict=False): + """Convert to a particle-number/spin-conserving fermion operator if valid. + + NetKet's ``ParticleNumberAndSpinConservingFermioperator2nd`` exposes far + fewer connected configurations, which makes VMC local energies cheaper. It + only applies to operators that conserve both particle number and spin, so + ``strict=False`` silently returns the input operator when the conversion is + not applicable (for example a spin-flipping observable). + """ + try: + import netket.experimental as nkx # pylint: disable=import-outside-toplevel + + cls = nkx.operator.ParticleNumberAndSpinConservingFermioperator2nd + except (ImportError, AttributeError): + if strict: + raise + return operator + try: + return cls.from_fermionoperator2nd(operator) + except Exception: # pylint: disable=broad-except + if strict: + raise + return operator + + +def netket_fermion_operator(hilbert, terms, *, constant=0.0, conserving=False): + r"""Build a NetKet fermion operator from symbolic second-quantized terms. + + This is the general "build from terms" primitive: it turns a list of + fermionic monomials into a jittable NetKet operator usable both as a VMC + Hamiltonian (``vstate.expect_and_grad``) and as a measurement observable + (``vstate.expect``). + + Parameters + ---------- + hilbert: + A NetKet ``SpinOrbitalFermions`` Hilbert space. + terms: + Iterable of ``(coefficient, ops)`` pairs. ``ops`` is an iterable of + ``(site, sz, dagger)`` tuples where ``site`` is the orbital index, + ``sz`` is ``+1`` (spin up), ``-1`` (spin down), or ``None`` (spinless), + and ``dagger`` is ``True`` for a creation operator + :math:`c^{\dagger}` and ``False`` for an annihilation operator + :math:`c`. Factors in ``ops`` are applied left to right, so a number + operator :math:`n = c^{\dagger} c` is + ``((site, sz, True), (site, sz, False))``. + constant: + Scalar identity shift added to the operator. + conserving: + When ``True`` or ``"auto"``, convert to NetKet's + particle-number/spin-conserving fermion operator (cheaper in VMC). + ``"auto"`` falls back to the plain operator when the conversion is not + applicable; ``True`` raises on failure. + + Returns + ------- + A jittable NetKet fermion operator. + """ + _require_netket() + from netket.operator import ( # pylint: disable=import-outside-toplevel + fermion as nkf, + ) + + total = None + for coeff, ops in terms: + term_op = None + for site, sz, dagger in ops: + builder = nkf.create if dagger else nkf.destroy + factor = builder(hilbert, int(site), sz=sz) + term_op = factor if term_op is None else term_op @ factor + if term_op is None: + continue + term_op = coeff * term_op + total = term_op if total is None else total + term_op + + if total is None: + total = nkf.zero(hilbert) + if constant: + total = total + constant * nkf.identity(hilbert) + + if conserving: + total = _maybe_conserving_fermion_operator( + total, strict=conserving != "auto" + ) + return total + + +def compile_operator_sum_netket(hilbert, terms, *, site_order=None, conserving=False): + """Compile a backend-neutral :class:`OperatorSum` for NetKet. + + 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. + """ + from .api import ( + LocalMatrixTerm, + ProductTerm, + _expand_fermion_factor, + normalize_operator_sum, + ) + + nk = _require_netket() + operator_sum = normalize_operator_sum(terms) + site_order = None if site_order is None else tuple(site_order) + + def map_site(site): + if site_order is None: + return int(site) if isinstance(site, Integral) else site + if site in site_order: + return site_order.index(site) + if isinstance(site, Integral) and 0 <= int(site) < len(site_order): + return int(site) + raise ValueError(f"Term site {site!r} is not present in site_order.") + + symbolic = [] + matrix_ops = [] + for term in operator_sum: + if isinstance(term, ProductTerm): + factors = [] + for factor in term.factors: + for site, name in _expand_fermion_factor(factor): + if name in {"create", "annihilate"}: + sz = None + elif name in {"create_u", "annihilate_u"}: + sz = 1 + elif name in {"create_d", "annihilate_d"}: + sz = -1 + else: + raise ValueError( + "NetKet ProductTerm factors must be fermionic " + "creation, annihilation, or number operators; " + f"got {name!r}." + ) + factors.append( + (map_site(site), sz, name.startswith("create")) + ) + symbolic.append((term.coefficient, tuple(factors))) + continue + if not isinstance(term, LocalMatrixTerm): # pragma: no cover + raise TypeError(f"Unsupported operator term {type(term).__name__}.") + matrix = term.matrix + detach = getattr(matrix, "detach", None) + if callable(detach): + matrix = detach() + cpu = getattr(matrix, "cpu", None) + if callable(cpu): + matrix = cpu() + matrix = np.asarray(matrix) * term.coefficient + n_sites = len(term.support) + # The common term contract uses output axes followed by input axes, + # while NetKet LocalOperator takes a flattened square local matrix. + local_dim = int(np.prod(matrix.shape[:n_sites])) + matrix = matrix.reshape(local_dim, local_dim) + matrix_ops.append( + nk.operator.LocalOperator( + hilbert, + matrix, + acting_on=[map_site(site) for site in term.support], + ) + ) + + total = None + if symbolic: + total = netket_fermion_operator( + hilbert, + symbolic, + constant=operator_sum.constant, + conserving=conserving, + ) + elif operator_sum.constant: + total = nk.operator.LocalOperator( + hilbert, + constant=operator_sum.constant, + ) + for matrix_op in matrix_ops: + total = matrix_op if total is None else total + matrix_op + if total is None: + total = nk.operator.LocalOperator(hilbert, constant=0.0) + return total + + +def fermion_model_terms(fermion, edges, *, 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. + """ + if not getattr(fermion, "spinful", True): + raise NotImplementedError( + "fermion_model_terms supports spinful fermions; build spinless " + "operators directly with netket_fermion_operator." + ) + edges = tuple((int(i), int(j)) for i, j in edges) + if n_sites is 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) + try: + return float(value) + except (TypeError, ValueError) as exc: + raise TypeError( + f"fermion.{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) + + terms = [] + for i, j in edges: + for sz in (1, -1): + terms.append((-t, ((i, sz, True), (j, sz, False)))) + terms.append((-t, ((j, sz, True), (i, sz, False)))) + if U: + for i in sites: + terms.append( + ( + U, + ( + (i, 1, True), + (i, 1, False), + (i, -1, True), + (i, -1, False), + ), + ) + ) + if V: + for i, j in edges: + for sz in (1, -1): + for sz2 in (1, -1): + terms.append( + ( + V, + ( + (i, sz, True), + (i, sz, False), + (j, sz2, True), + (j, sz2, False), + ), + ) + ) + if mu: + for i in sites: + for sz in (1, -1): + terms.append((-mu, ((i, sz, True), (i, sz, False)))) + return terms + + +def standard_fermion_observables(hilbert): + """Return common spinful-fermion observables for :meth:`NetKetPEPSVMC.measure`. + + The returned ``{name: operator}`` mapping contains the total particle + number per spin (``"n_up"``, ``"n_down"``), the total particle number + (``"n_total"``), and the total on-site double occupancy + (``"double_occupancy"``). Each value is a jittable NetKet operator. + """ + n = int(hilbert.n_orbitals) + n_up = [(1.0, ((i, 1, True), (i, 1, False))) for i in range(n)] + n_down = [(1.0, ((i, -1, True), (i, -1, False))) for i in range(n)] + doub = [ + (1.0, ((i, 1, True), (i, 1, False), (i, -1, True), (i, -1, False))) + 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), + } + + +def _normalize_fermion_observables(hilbert, observables, *, site_order=None): + """Resolve an ``{name: operator_or_terms}`` mapping to NetKet operators.""" + from .api import OperatorSum + if observables is None: + return None + if not isinstance(observables, dict): + raise TypeError( + "observables must be a {name: operator_or_terms} mapping." + ) + resolved = {} + for name, spec in observables.items(): + if hasattr(spec, "hilbert"): + resolved[str(name)] = spec + elif isinstance(spec, OperatorSum): + resolved[str(name)] = compile_operator_sum_netket( + hilbert, + spec, + site_order=site_order, + ) + else: + resolved[str(name)] = netket_fermion_operator(hilbert, spec) + return resolved + + +def _looks_like_native_fermion_terms(obj): + """True for a native Pepsy ``SymHamiltonian`` or coordinate-keyed mapping.""" + if obj is None or hasattr(obj, "hilbert"): + return False + terms_map = getattr(obj, "terms", None) + if terms_map is not None and hasattr(terms_map, "keys"): + return True + return hasattr(obj, "keys") + + +def _native_fermion_terms_mapping(obj): + """Return the coordinate/int-keyed terms mapping from a native source.""" + terms_map = getattr(obj, "terms", None) + if terms_map is not None and hasattr(terms_map, "keys"): + return terms_map + return obj + + def build_fermi_hubbard_vmc( peps, *, - Lx, - Ly, - t=1.0, - U=8.0, + Lx=None, + Ly=None, + t=None, + U=None, + fermion=None, + terms=None, n_fermions_per_spin=None, - pbc=False, + sector=None, + pbc=None, edges=None, graph=None, contraction="exact", chi=None, cutoff=0.0, contraction_opts=None, + register_stable_svd=True, n_samples=1024, n_chains=16, n_discard_per_chain=32, chunk_size=256, + sampling=None, sampler_chunk_size=None, seed=None, sampler_seed=None, @@ -1965,11 +3388,76 @@ def build_fermi_hubbard_vmc( param_dtype=None, verify_columns=False, ): - """Create the first-pass NetKet VMC objects for a fermionic PEPS.""" + """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 + 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 + ``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 + ``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 ``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. + """ + ( + n_samples, + n_chains, + n_discard_per_chain, + chunk_size, + seed, + sampler_seed, + ) = _resolve_sampling_build_config( + sampling, + n_samples=n_samples, + n_chains=n_chains, + n_discard_per_chain=n_discard_per_chain, + chunk_size=chunk_size, + seed=seed, + sampler_seed=sampler_seed, + ) nk = _require_netket() + 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) + Lx, Ly = int(Lx), int(Ly) n_sites = int(Lx) * int(Ly) + if sector is not None: + if n_fermions_per_spin is not None: + raise ValueError( + "Pass only one of sector and n_fermions_per_spin." + ) + n_fermions_per_spin = sector + n_fermions_per_spin = _coerce_spin_sector(n_fermions_per_spin) if n_fermions_per_spin is None: n_fermions_per_spin = (n_sites // 2, n_sites // 2) + 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) + t = 1.0 if t is None else t + U = 8.0 if U is None else U + if pbc is None: + pbc = _infer_pbc_from_fermi_terms(terms, Lx, Ly) hilbert = nk.hilbert.SpinOrbitalFermions( n_sites, @@ -1978,7 +3466,9 @@ def build_fermi_hubbard_vmc( ) if graph is None: if edges is None: - edges = square_lattice_edges(Lx, Ly, pbc=pbc) + edges = _edges_from_fermi_terms(terms, Lx, Ly) + if not edges: + edges = square_lattice_edges(Lx, Ly, pbc=pbc) graph = nk.graph.Graph(edges=tuple(edges), n_nodes=n_sites) hamiltonian = nk.operator.FermiHubbardJax( @@ -1992,7 +3482,12 @@ def build_fermi_hubbard_vmc( if verify_columns: verify_netket_spin_columns(hilbert, columns) + if register_stable_svd: + _maybe_register_stable_jax_svd(contraction) + + 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) model = make_fermionic_peps_log_amplitude_model( ansatz, columns, @@ -2047,6 +3542,336 @@ def build_fermi_hubbard_vmc( ) +def build_fermion_vmc( + peps, + *, + fermion=None, + hamiltonian=None, + terms=None, + observables=None, + Lx=None, + Ly=None, + n_fermions_per_spin=None, + sector=None, + pbc=None, + edges=None, + graph=None, + conserving="auto", + contraction="exact", + chi=None, + cutoff=0.0, + contraction_opts=None, + register_stable_svd=True, + n_samples=1024, + n_chains=16, + n_discard_per_chain=32, + chunk_size=256, + sampling=None, + sampler_chunk_size=None, + seed=None, + sampler_seed=None, + use_sr="auto", + max_sr_params=5_000, + sr_diag_shift=0.01, + sr_diag_scale=None, + sr_qgt="auto", + sr_solver=None, + sr_solver_restart=False, + param_dtype=None, + verify_columns=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`. + * ``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 + Torch VMC; it is compiled to a NetKet fermion/local operator. + * ``hamiltonian`` -- an already-built NetKet fermion operator. + + ``edges`` / ``graph`` / ``pbc``, when given explicitly, take precedence over + any geometry inferred from native terms. + + ``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 + sampler / state / SR / contraction options match + :func:`build_fermi_hubbard_vmc`, and the returned setup exposes the same + :meth:`NetKetPEPSVMC.warmup` and :meth:`NetKetPEPSVMC.optimize` helpers. + """ + ( + n_samples, + n_chains, + n_discard_per_chain, + chunk_size, + seed, + sampler_seed, + ) = _resolve_sampling_build_config( + sampling, + n_samples=n_samples, + n_chains=n_chains, + n_discard_per_chain=n_discard_per_chain, + chunk_size=chunk_size, + seed=seed, + sampler_seed=sampler_seed, + ) + nk = _require_netket() + from .api import OperatorSum + common_hamiltonian = hamiltonian if isinstance(hamiltonian, OperatorSum) else None + common_terms = terms if isinstance(terms, OperatorSum) else None + if common_hamiltonian is not None and common_terms is not None: + raise ValueError("Pass the common OperatorSum as hamiltonian or terms, not both.") + common_operator_sum = ( + common_hamiltonian + if common_hamiltonian is not None + else common_terms + ) + if fermion is not None and not getattr(fermion, "spinful", True): + raise NotImplementedError( + "build_fermion_vmc supports spinful fermions; use the sparse/torch " + "path for spinless models." + ) + 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) + Lx, Ly = int(Lx), int(Ly) + n_sites = Lx * Ly + if sector is not None: + if n_fermions_per_spin is not None: + raise ValueError( + "Pass only one of sector and n_fermions_per_spin." + ) + n_fermions_per_spin = sector + n_fermions_per_spin = _coerce_spin_sector(n_fermions_per_spin) + if n_fermions_per_spin is None: + n_fermions_per_spin = (n_sites // 2, n_sites // 2) + 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).") + + hilbert = nk.hilbert.SpinOrbitalFermions( + n_sites, + s=1 / 2, + 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: + pbc = ( + False + if common_operator_sum is not None + else _infer_pbc_from_fermi_terms(native_terms, Lx, Ly) + ) + if graph is None: + if edges is None: + edges = ( + _edges_from_operator_sum(common_operator_sum, Lx, Ly) + if common_operator_sum is not None + else _edges_from_fermi_terms(native_terms, Lx, Ly) + ) + if not edges: + edges = square_lattice_edges(Lx, Ly, pbc=pbc) + graph = nk.graph.Graph(edges=tuple(edges), n_nodes=n_sites) + elif edges is None: + edges = tuple(tuple(edge) for edge in graph.edges()) + + # Build the NetKet Hamiltonian. + if common_operator_sum is not None: + hamiltonian = compile_operator_sum_netket( + hilbert, + common_operator_sum, + site_order=_row_major_sites(Lx, Ly), + conserving=conserving, + ) + elif prebuilt_operator is not None: + hamiltonian = prebuilt_operator + elif symbolic_terms is not None: + 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) + hamiltonian = netket_fermion_operator( + hilbert, model_terms, conserving=conserving + ) + 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." + ) + + observable_ops = _normalize_fermion_observables( + hilbert, + observables, + site_order=_row_major_sites(Lx, Ly), + ) + + if register_stable_svd: + _maybe_register_stable_jax_svd(contraction) + + columns = netket_spin_orbital_columns(hilbert) + if verify_columns: + verify_netket_spin_columns(hilbert, columns) + + 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) + model = make_fermionic_peps_log_amplitude_model( + ansatz, + columns, + contraction=contraction, + chi=chi, + cutoff=cutoff, + contraction_opts=contraction_opts, + param_dtype=param_dtype, + ) + sampler_kwargs = { + "graph": graph, + "n_chains": n_chains, + "spin_symmetric": True, + } + if sampler_chunk_size is not None: + sampler_kwargs["chunk_size"] = _check_positive_int( + "sampler_chunk_size", + sampler_chunk_size, + ) + sampler = nk.sampler.MetropolisFermionHop(hilbert, **sampler_kwargs) + vstate = nk.vqs.MCState( + sampler, + model, + 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, + ) + preconditioner = _maybe_make_sr_preconditioner( + ansatz, + use_sr=use_sr, + max_sr_params=max_sr_params, + sr_diag_shift=sr_diag_shift, + sr_diag_scale=sr_diag_scale, + sr_qgt=sr_qgt, + sr_solver=sr_solver, + sr_solver_restart=sr_solver_restart, + ) + return NetKetFermiHubbardVMC( + hilbert=hilbert, + graph=graph, + hamiltonian=hamiltonian, + sampler=sampler, + vstate=vstate, + model=model, + ansatz=ansatz, + config_map=None, + preconditioner=preconditioner, + columns=columns, + observables=observable_ops, + ) + + +def build_netket_vmc( + problem, + *, + fermion=None, + contraction=None, + sampling=None, + **kwargs, +): + """Build the portable NetKet façade from a :class:`VMCProblem`. + + The portable fermion path targets NetKet's fixed ``U1U1`` spin-orbital + Hilbert space. Native builders remain available for spin models and for + backend-specific sampler/driver choices. + """ + from .api import ( + ContractionConfig, + SamplingConfig, + VMCBackendCapabilityError, + VMCProblem, + ) + + if not isinstance(problem, VMCProblem): + raise TypeError("problem must be a VMCProblem.") + if sampling is not None and not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if problem.symmetry not in {None, "U1U1"}: + raise VMCBackendCapabilityError( + "The portable NetKet fermion setup currently supports only the " + "fixed U1U1 sector; use build_torch_vmc for " + f"symmetry={problem.symmetry!r}." + ) + if problem.site_order is not None: + raise VMCBackendCapabilityError( + "build_netket_vmc uses NetKet's fixed row-major site order. " + "Use build_fermion_vmc directly for an explicitly remapped model." + ) + if sampling is not None: + unsupported = [] + if sampling.thin != 1: + unsupported.append("thin") + if sampling.proposal is not None: + unsupported.append("proposal") + if unsupported: + raise VMCBackendCapabilityError( + "NetKet cannot honour portable SamplingConfig." + f"{', '.join(unsupported)}; configure its native sampler directly." + ) + if contraction is None: + contraction = ContractionConfig() + setup = build_fermion_vmc( + problem.peps, + fermion=fermion, + hamiltonian=problem.hamiltonian, + observables=problem.observables, + contraction=contraction, + sampling=sampling, + **kwargs, + ) + return NetKetVMCSetup(setup=setup, problem=problem) + + def build_sparse_fermi_hubbard_vmc( peps, *, @@ -2342,3 +4167,79 @@ def make_netket_vmc_driver( variational_state=setup.vstate, **kwargs, ) + + +def warmup_netket_vmc(setup, *, hamiltonian=None, progress=True, verbose=None): + """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. + + Parameters + ---------- + setup: + A NetKet VMC setup (for example from :func:`build_fermi_hubbard_vmc`) + exposing ``vstate`` and ``hamiltonian`` attributes, or a bare NetKet + variational state. + 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. + 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. + + Returns + ------- + float + Elapsed wall-clock seconds spent compiling and evaluating once. + """ + import time # pylint: disable=import-outside-toplevel + + vstate = getattr(setup, "vstate", setup) + if hamiltonian is None: + hamiltonian = getattr(setup, "hamiltonian", None) + if hamiltonian is None: + raise ValueError( + "warmup_netket_vmc needs a Hamiltonian: pass hamiltonian=... or a " + "setup exposing a .hamiltonian attribute." + ) + bar = _make_progress_bar( + total=2, desc="Compiling VMC kernels", enabled=progress + ) + if verbose is None: + verbose = bar is None + if verbose: + print( + "Compiling NetKet VMC kernels (sampler, amplitude, energy)...", + flush=True, + ) + 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() + if verbose: + print( + f"Compiled and warmed up in {elapsed:.1f} s; " + "progress-bar ETA is now meaningful.", + flush=True, + ) + return elapsed diff --git a/src/pepsy/vmc/torch.py b/src/pepsy/vmc/torch.py index 7314609..c554ebd 100644 --- a/src/pepsy/vmc/torch.py +++ b/src/pepsy/vmc/torch.py @@ -7,17 +7,35 @@ from __future__ import annotations -from dataclasses import dataclass +from itertools import product +import math +import os +import sysconfig +import time +import warnings +from dataclasses import dataclass, replace from numbers import Integral from typing import Any +import numpy as np + __all__ = [ "FermionSiteEncoding", + "SpinlessSiteEncoding", + "TorchFermionVMCMetadata", "TorchPEPSAmplitude", "TorchPEPSBoundaryAmplitude", "TorchConnections", "TorchMetropolisResult", + "TorchMCMCSamples", + "TorchChainDiagnostics", + "TorchMetropolisSampler", + "TorchBPMetropolisSampler", "TorchVMCDriver", + "TorchFermionVMC", + "TorchVMCSetup", + "TorchVMCEnergyEstimate", + "TorchVMCImportanceEstimate", "TorchVMCStepResult", "TorchSRResult", "TorchSquareLattice", @@ -25,16 +43,24 @@ "count_spinful_particles", "heisenberg_connections", "local_energy_from_connections", + "torch_chain_diagnostics", + "metropolis_local_sampler", "metropolis_exchange_sweep", "propose_spin_exchange", "propose_spinful_exchange_or_hopping", + "propose_spinful_u1_exchange_or_hopping", + "propose_spinful_z2_exchange_or_hopping", + "propose_spinful_z2z2_exchange_or_hopping", "random_spin_configs", "random_spinful_configs", "make_torch_peps_amplitude_model", + "compile_operator_sum_torch", + "build_torch_vmc", "solve_torch_sr", "spinful_fermi_hubbard_connections", "torch_log_derivative_matrix", "transverse_ising_connections", + "torch_hamiltonian_connections", ] @@ -94,6 +120,24 @@ def _site_value(value, site): "boundary-mps": "boundary", } +_PROPOSAL_BATCHING_MODES = {"auto", "cache", "vmap"} + +# ``amplitude_batching`` controls the independent-configuration path. This +# is intentionally separate from ``proposal_batching``: local Metropolis +# proposals can use boundary-environment reuse even when a PEPS's ordinary +# amplitude path must remain serial (for example, native U1/U1U1 Symmray). +_AMPLITUDE_BATCHING_MODES = {"auto", "serial", "vmap"} + +# Boundary-environment reuse is useful for small connected sets, while the +# vmapped full-boundary path is substantially faster once many off-diagonal +# configurations are measured together. +_BOUNDARY_VMAP_CONNECTION_THRESHOLD = 64 + +# Only pure, fixed-shape tensor bookkeeping is eligible for ``torch.compile``. +# PEPS/Symmray selection and contraction deliberately remain eager. +_COMPILED_CHEAP_TORCH_KERNELS = {} +_FAILED_CHEAP_TORCH_KERNELS = set() + def _validate_contraction(contraction, chi): key = str(contraction).replace("_", "-").lower() @@ -133,6 +177,654 @@ def _find_symmray_tensors(tn): ] +def _graded_torch_index_map(index): + """Return the dense charge label at every position of ``index``.""" + return tuple( + charge + for charge, size in index.chargemap.items() + for _ in range(int(size)) + ) + + +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()) + else: + dense = np.zeros(shape, dtype=float) + + selectors = [] + for index, label in zip(array.indices, labels): + positions = [] + for charge in _graded_torch_index_map(index): + positions.append(full_maps[label].index(charge)) + selectors.append(positions) + + if not labels: + return dense.reshape(()) + + full = np.zeros(shape, dtype=dense.dtype) + if all(selectors): + full[np.ix_(*selectors)] = dense + return full + + +def _graded_torch_pad(array, labels, full_maps): + """Pad a Symmray result while retaining its graded metadata.""" + dense = _graded_torch_embed_dense(array, labels, full_maps) + return type(array).from_dense( + dense, + [full_maps[label] for label in labels], + duals=array.duals, + charge=array.charge, + symmetry=array.symmetry, + invalid_sectors="ignore", + dummy_modes=(), + ) + + +def _graded_torch_dense(array, labels, full_maps): + """Get a padded dense probe, including the empty-array case.""" + return _graded_torch_embed_dense(array, labels, full_maps) + + +def _graded_torch_sign_mask(left, right): + """Extract a +/-1 phase mask from two equal-layout dense arrays.""" + left = np.asarray(left) + right = np.asarray(right) + if left.shape != right.shape: + raise ValueError( + "Symmray graded projector produced incompatible dense shapes: " + f"{left.shape} and {right.shape}." + ) + mask = np.ones_like(left, dtype=float) + nonzero = np.abs(right) > 1.0e-12 + if np.any(nonzero): + ratio = np.real(left[nonzero] / right[nonzero]) + if np.max(np.abs(np.abs(ratio) - 1.0)) > 1.0e-7: + raise ValueError( + "Could not compile a fixed graded Torch phase mask: " + "the Symmray phase was not a +/-1 sector phase; " + f"observed ratios {np.unique(ratio)!r}." + ) + mask[nonzero] = ratio + return mask + + +def _graded_torch_unit_probe( + array_cls, + reference, + labels, + charge, + full_maps, + duals, + rng, +): + """Construct a random symmetric probe for static phase compilation.""" + shape = tuple(len(full_maps[label]) for label in labels) + dense = rng.normal(size=shape) + return array_cls.from_dense( + dense, + [full_maps[label] for label in labels], + duals=duals, + charge=charge, + symmetry=reference.symmetry, + invalid_sectors="ignore", + dummy_modes=(), + ) + + +def _graded_torch_prepare_pair(a, b, axes): + """Use Symmray's public graded operation with a narrow compatibility shim. + + Symmray currently exposes the complete preparation logic through the + internal helper used by its public ``tensordot`` implementation. Keeping + this call in one adapter makes the dependency/version assumption explicit, + rather than reproducing the upstream Koszul-phase rules here. + """ + prepare = getattr(a, "_prepare_for_tensordot_fermionic", None) + if prepare is None: + raise TypeError( + "The installed Symmray version does not expose the graded " + "tensordot preparation needed by graded_torch." + ) + return prepare(b, axes) + + +def _graded_torch_contraction_mask(a, b, axes, perm_b): + """Keep only equal-charge positions on contracted dense axes.""" + mask = np.ones(tuple(b.shape[axis] for axis in perm_b), dtype=float) + for axis_a, axis_b in zip(*axes): + charges_a = _graded_torch_index_map(a.indices[axis_a]) + charges_b = _graded_torch_index_map(b.indices[axis_b]) + if len(charges_a) != len(charges_b): + raise ValueError( + "Symmray graded projector found mismatched contracted index " + f"sizes {len(charges_a)} and {len(charges_b)}." + ) + valid = np.asarray( + [charge_a == charge_b for charge_a, charge_b in zip(charges_a, charges_b)], + dtype=float, + ) + axis_b_new = perm_b.index(axis_b) + shape = [1] * len(perm_b) + shape[axis_b_new] = len(valid) + mask *= valid.reshape(shape) + return mask + + +def _graded_torch_compile_pair( + a, + b, + labels_a, + labels_b, + axes, + full_maps, + rng, +): + """Compile dense input/output masks for one graded contraction.""" + left_axes = tuple(k for k in range(a.ndim) if k not in axes[0]) + right_axes = tuple(k for k in range(b.ndim) if k not in axes[1]) + perm_a = (*left_axes, *axes[0]) + perm_b = (*axes[1], *right_axes) + aa, bb, new_axes_a, new_axes_b = _graded_torch_prepare_pair(a, b, axes) + + a_dense = a.to_dense() + b_dense = b.to_dense() + mask_a = _graded_torch_sign_mask( + aa.to_dense(), np.transpose(a_dense, perm_a) + ) + mask_b = _graded_torch_sign_mask( + bb.to_dense(), np.transpose(b_dense, perm_b) + ) + mask_b = mask_b * _graded_torch_contraction_mask(a, b, axes, perm_b) + raw = np.tensordot( + np.transpose(a_dense, perm_a) * mask_a, + np.transpose(b_dense, perm_b) * mask_b, + axes=(new_axes_a, new_axes_b), + ) + + c = a.tensordot(b, axes=axes, preserve_array=True) + labels_c = tuple( + labels_a[k] for k in left_axes + ) + tuple(labels_b[k] for k in right_axes) + c_padded = _graded_torch_pad(c, labels_c, full_maps) + c_dense = _graded_torch_dense(c_padded, labels_c, full_maps) + try: + mask_c = _graded_torch_sign_mask(c_dense, raw) + except ValueError as exc: + raise ValueError( + "Failed compiling graded Torch pair " + f"{labels_a!r} x {labels_b!r}, axes={axes!r}, " + f"charges={a.charge!r},{b.charge!r}, duals={a.duals!r},{b.duals!r}." + ) from exc + return c_padded, mask_a, mask_b, mask_c + + +@dataclass(frozen=True) +class _GradedTorchPair: + """One static dense contraction in the graded Torch projector.""" + + left: int + right: int + axes_a: tuple[int, ...] + axes_b: tuple[int, ...] + perm_a: tuple[int, ...] + perm_b: tuple[int, ...] + ncon: int + charge_masks_a: tuple[np.ndarray, ...] + charge_masks_b: tuple[np.ndarray, ...] + output_masks: tuple[np.ndarray, ...] + output_charge_ids: np.ndarray + n_charge_b: int + left_sites: tuple[int, ...] + right_sites: tuple[int, ...] + + +class _GradedTorchProjector: + """Compile a fixed-shape, graded dense projector for native Symmray PEPS. + + The dense values remain differentiable Torch tensors. Symmray is used at + construction time to compile charge-sector transitions and the exact + fermionic phase masks for the chosen contraction tree. This avoids the + dynamic sparse ``isel`` path, whose scalar charge lookup is not vmap-safe. + """ + + def __init__(self, tn, sites, *, contraction_opts=None): + import quimb.tensor as qtn + + self.sites = tuple(sites) + self.n_sites = len(self.sites) + self.full_maps = {} + self.array_cls = None + self.reference = None + self.leaf_labels = [] + self.leaf_duals = [] + self.physical_axes = [] + self.leaf_charge_values = [] + self.leaf_charge_ids = [] + self.leaf_dummy_modes = [] + self.leaf_dummy_parities = [] + self._mode_objects = {} + + for site in self.sites: + tensor = tn[site] + data = getattr(tensor, "data", None) + if not _is_symmray_data(data) or not getattr(data, "fermionic", False): + raise TypeError( + "graded_torch requires native Symmray fermionic arrays." + ) + if not hasattr(data, "_prepare_for_tensordot_fermionic"): + raise TypeError( + "graded_torch requires a Symmray sparse fermionic backend." + ) + if self.array_cls is None: + self.array_cls = type(data) + self.reference = data + if str(data.symmetry).upper() != "U1U1": + raise NotImplementedError( + "graded_torch currently targets native U1U1 fermionic " + "arrays; use the established native path for other " + "Symmray symmetries." + ) + + physical_ind = tn.site_ind(site) + try: + physical_axis = tuple(tensor.inds).index(physical_ind) + except ValueError as exc: + raise ValueError( + f"Could not locate physical index for PEPS site {site!r}." + ) from exc + self.physical_axes.append(physical_axis) + labels = tuple( + ind for axis, ind in enumerate(tensor.inds) + if axis != physical_axis + ) + self.leaf_labels.append(labels) + virtual_indices = tuple( + index for axis, index in enumerate(data.indices) + if axis != physical_axis + ) + self.leaf_duals.append(tuple(index.dual for index in virtual_indices)) + for label, index in zip(labels, virtual_indices): + index_map = list(_graded_torch_index_map(index)) + old_map = self.full_maps.setdefault(label, index_map) + if old_map != index_map: + raise ValueError( + f"Symmray index {label!r} has inconsistent charge order." + ) + + charge_values = [] + dummy_modes = [] + dummy_parities = [] + physical_dim = int(data.shape[physical_axis]) + for physical_value in range(physical_dim): + selected = data.isel(physical_axis, physical_value) + charge = self._charge_key(selected.charge) + if charge not in charge_values: + charge_values.append(charge) + modes = tuple(selected.dummy_modes) + if any(mode.dual for mode in modes): + raise NotImplementedError( + "graded_torch currently supports ket PEPS with " + "non-dual dummy fermion modes only." + ) + mode_keys = tuple(self._register_mode(mode) for mode in modes) + dummy_modes.append(mode_keys) + dummy_parities.append(sum(mode.parity for mode in modes) % 2) + self.leaf_charge_values.append(tuple(charge_values)) + self.leaf_charge_ids.append( + tuple(charge_values.index(self._charge_key(data.isel( + physical_axis, value + ).charge)) for value in range(physical_dim)) + ) + self.leaf_dummy_modes.append(tuple(dummy_modes)) + self.leaf_dummy_parities.append(tuple(dummy_parities)) + + self._mode_order = self._build_mode_order() + self._dummy_inversions = self._build_dummy_inversions() + + path_tensors = [ + qtn.Tensor( + np.zeros(tuple(len(self.full_maps[label]) for label in labels)), + inds=labels, + ) + for labels in self.leaf_labels + ] + optimize = "auto-hq" + if contraction_opts: + final_opts = contraction_opts.get("final_contract_opts") or {} + optimize = final_opts.get("optimize", optimize) + path = qtn.TensorNetwork(path_tensors).contract( + all, get="path", optimize=optimize + ) + + rng = np.random.default_rng(928371) + active_labels = list(self.leaf_labels) + active_duals = list(self.leaf_duals) + active_charges = [list(values) for values in self.leaf_charge_values] + active_sites = [(site,) for site in range(self.n_sites)] + self.pairs = [] + + for left, right in path: + labels_a = active_labels[left] + labels_b = active_labels[right] + left_sites = active_sites[left] + right_sites = active_sites[right] + axes_a = tuple( + axis for axis, label in enumerate(labels_a) + if label in labels_b + ) + axes_b = tuple(labels_b.index(labels_a[axis]) for axis in axes_a) + if len(axes_a) != len(axes_b): + raise ValueError("Invalid contraction path for graded_torch.") + charge_masks_a = [] + charge_masks_b = [] + output_masks = [] + output_charges = [] + for charge_a in active_charges[left]: + for charge_b in active_charges[right]: + a = _graded_torch_unit_probe( + self.array_cls, + self.reference, + labels_a, + charge_a, + self.full_maps, + active_duals[left], + rng, + ) + b = _graded_torch_unit_probe( + self.array_cls, + self.reference, + labels_b, + charge_b, + self.full_maps, + active_duals[right], + rng, + ) + if not getattr(a, "num_blocks", 0) or not getattr( + b, "num_blocks", 0 + ): + labels_c = tuple( + labels_a[axis] + for axis in range(len(labels_a)) + if axis not in axes_a + ) + tuple( + labels_b[axis] + for axis in range(len(labels_b)) + if axis not in axes_b + ) + duals_c = tuple( + active_duals[left][axis] + for axis in range(len(labels_a)) + if axis not in axes_a + ) + tuple( + active_duals[right][axis] + for axis in range(len(labels_b)) + if axis not in axes_b + ) + charge_c = self._charge_key( + self.reference.symmetry.combine(charge_a, charge_b) + ) + shape_c = tuple( + len(self.full_maps[label]) for label in labels_c + ) + zero = np.zeros(shape_c, dtype=float) + c = self.array_cls.from_dense( + zero, + [self.full_maps[label] for label in labels_c], + duals=duals_c, + charge=charge_c, + symmetry=self.reference.symmetry, + invalid_sectors="ignore", + dummy_modes=(), + ) + mask_a = np.ones(a.shape, dtype=float) + mask_b = np.ones(b.shape, dtype=float) + mask_c = np.ones(shape_c, dtype=float) + else: + c, mask_a, mask_b, mask_c = _graded_torch_compile_pair( + a, + b, + labels_a, + labels_b, + (axes_a, axes_b), + self.full_maps, + rng, + ) + charge_c = self._charge_key(c.charge) + charge_masks_a.append(mask_a) + charge_masks_b.append(mask_b) + output_masks.append(mask_c) + output_charges.append(charge_c) + + n_charge_b = len(active_charges[right]) + charge_ids_c = tuple(dict.fromkeys(output_charges)) + output_charge_ids = np.asarray( + [charge_ids_c.index(charge) for charge in output_charges], + dtype=np.int64, + ).reshape(len(active_charges[left]), n_charge_b) + labels_c = tuple( + labels_a[axis] for axis in range(len(labels_a)) + if axis not in axes_a + ) + tuple( + labels_b[axis] for axis in range(len(labels_b)) + if axis not in axes_b + ) + active_labels[left] = labels_c + active_duals[left] = tuple( + active_duals[left][axis] for axis in range(len(labels_a)) + if axis not in axes_a + ) + tuple( + active_duals[right][axis] for axis in range(len(labels_b)) + if axis not in axes_b + ) + active_charges[left] = list(charge_ids_c) + active_sites[left] = active_sites[left] + active_sites[right] + active_labels.pop(right) + active_duals.pop(right) + active_charges.pop(right) + active_sites.pop(right) + + self.pairs.append( + _GradedTorchPair( + left=left, + right=right, + axes_a=axes_a, + axes_b=axes_b, + perm_a=tuple( + axis for axis in range(len(labels_a)) if axis not in axes_a + ) + axes_a, + perm_b=axes_b + tuple( + axis for axis in range(len(labels_b)) if axis not in axes_b + ), + ncon=len(axes_a), + charge_masks_a=tuple(charge_masks_a), + charge_masks_b=tuple(charge_masks_b), + output_masks=tuple(output_masks), + output_charge_ids=output_charge_ids, + n_charge_b=n_charge_b, + left_sites=left_sites, + right_sites=right_sites, + ) + ) + + self._torch_cache = {} + + @staticmethod + def _charge_key(charge): + if isinstance(charge, np.ndarray): + charge = charge.tolist() + if isinstance(charge, (list, tuple)): + return tuple(int(value) for value in charge) + return int(charge) + + def _register_mode(self, mode): + key = (repr(mode.label), bool(mode.dual), int(mode.parity)) + self._mode_objects.setdefault(key, mode) + return key + + def _build_mode_order(self): + keys = list(self._mode_objects) + try: + keys.sort(key=lambda key: self._mode_objects[key]) + except TypeError as exc: + raise NotImplementedError( + "graded_torch could not order Symmray dummy fermion modes." + ) from exc + return {key: position for position, key in enumerate(keys)} + + def _build_dummy_inversions(self): + inversions = [] + for left in range(self.n_sites): + row = [] + for right in range(self.n_sites): + table = np.zeros( + ( + len(self.leaf_dummy_modes[left]), + len(self.leaf_dummy_modes[right]), + ), + dtype=np.int64, + ) + for code_a, modes_a in enumerate(self.leaf_dummy_modes[left]): + for code_b, modes_b in enumerate(self.leaf_dummy_modes[right]): + table[code_a, code_b] = sum( + self._mode_order[mode_b] < self._mode_order[mode_a] + for mode_a in modes_a + for mode_b in modes_b + ) % 2 + row.append(table) + inversions.append(row) + return inversions + + def _torch_tables(self, reference): + torch = _require_torch() + key = (reference.device, reference.dtype) + cached = self._torch_cache.get(key) + if cached is not None: + return cached + charge_luts = tuple( + torch.as_tensor(ids, dtype=torch.long, device=reference.device) + for ids in self.leaf_charge_ids + ) + parity_luts = tuple( + torch.as_tensor(values, dtype=torch.long, device=reference.device) + for values in self.leaf_dummy_parities + ) + pairs = [] + for pair in self.pairs: + pairs.append( + ( + torch.stack([ + torch.as_tensor(mask, dtype=reference.dtype, device=reference.device) + for mask in pair.charge_masks_a + ]), + torch.stack([ + torch.as_tensor(mask, dtype=reference.dtype, device=reference.device) + for mask in pair.charge_masks_b + ]), + torch.stack([ + torch.as_tensor(mask, dtype=reference.dtype, device=reference.device) + for mask in pair.output_masks + ]), + torch.as_tensor( + pair.output_charge_ids.reshape(-1), + dtype=torch.long, + device=reference.device, + ), + ) + ) + cached = charge_luts, parity_luts, pairs + self._torch_cache[key] = cached + return cached + + def _dummy_phase(self, config, left_sites, right_sites, parity_luts): + torch = _require_torch() + def lookup_1d(table, index): + return torch.index_select(table, 0, index.reshape(1)).squeeze(0) + + left_parity = sum( + lookup_1d(parity_luts[site], config[site]) for site in left_sites + ) % 2 + right_parity = sum( + lookup_1d(parity_luts[site], config[site]) for site in right_sites + ) % 2 + exponent = left_parity * right_parity + for site_a in left_sites: + for site_b in right_sites: + table = torch.as_tensor( + self._dummy_inversions[site_a][site_b], + dtype=torch.long, + device=config.device, + ) + row = lookup_1d(table, config[site_a]) + exponent = exponent + lookup_1d(row, config[site_b]) + return torch.where( + exponent % 2 == 0, + torch.ones((), dtype=parity_luts[0].dtype, device=config.device), + torch.full((), -1, dtype=parity_luts[0].dtype, device=config.device), + ) + + def __call__(self, dense_leaves, configs, *, use_vmap=True): + torch = _require_torch() + reference = dense_leaves[0] + charge_luts, parity_luts, pairs = self._torch_tables(reference) + + def select(data, axis, value): + return torch.index_select(data, axis, value.reshape(1)).squeeze(axis) + + def lookup_1d(table, index): + return torch.index_select(table, 0, index.reshape(1)).squeeze(0) + + def evaluate(config): + arrays = [ + select(data, axis, config[site]) + for site, (data, axis) in enumerate( + zip(dense_leaves, self.physical_axes) + ) + ] + charge_ids = [ + lookup_1d(charge_luts[site], config[site]) + for site in range(self.n_sites) + ] + for pair, (mask_a, mask_b, mask_c, output_ids) in zip( + self.pairs, pairs + ): + lookup = charge_ids[pair.left] * pair.n_charge_b + charge_ids[ + pair.right + ] + a = arrays[pair.left].permute(pair.perm_a) * lookup_1d( + mask_a, lookup + ) + b = arrays[pair.right].permute(pair.perm_b) * lookup_1d( + mask_b, lookup + ) + raw = torch.tensordot( + a, + b, + dims=( + tuple(range(a.ndim - pair.ncon, a.ndim)), + tuple(range(pair.ncon)), + ), + ) + phase = self._dummy_phase( + config, + pair.left_sites, + pair.right_sites, + parity_luts, + ).to(dtype=raw.dtype) + arrays[pair.left] = raw * lookup_1d(mask_c, lookup) * phase + charge_ids[pair.left] = lookup_1d(output_ids, lookup) + arrays.pop(pair.right) + charge_ids.pop(pair.right) + return arrays[0] + + if use_vmap: + return torch.vmap(evaluate)(configs) + return torch.stack([evaluate(config) for config in configs]) + + def _as_torch_scalar(value, reference): torch = _require_torch() if isinstance(value, torch.Tensor): @@ -148,6 +840,30 @@ def _normalize_chunk_size(chunk_size, *, name="chunk_size"): return _check_positive_int(name, chunk_size) +def _normalize_proposal_batching(proposal_batching): + """Normalize the boundary-proposal batching policy.""" + mode = str(proposal_batching).replace("_", "-").lower() + if mode not in _PROPOSAL_BATCHING_MODES: + choices = ", ".join( + repr(choice) for choice in sorted(_PROPOSAL_BATCHING_MODES) + ) + raise ValueError(f"proposal_batching must be one of {choices}.") + return mode + + +def _normalize_amplitude_batching(amplitude_batching): + """Normalize the independent-amplitude batching policy.""" + mode = str(amplitude_batching).replace("_", "-").lower() + aliases = {"loop": "serial"} + mode = aliases.get(mode, mode) + if mode not in _AMPLITUDE_BATCHING_MODES: + choices = ", ".join( + repr(choice) for choice in sorted(_AMPLITUDE_BATCHING_MODES) + ) + raise ValueError(f"amplitude_batching must be one of {choices}.") + return mode + + def _check_nonnegative_int(name, value): if isinstance(value, bool) or not isinstance(value, Integral) or value < 0: raise ValueError(f"{name} must be a non-negative integer.") @@ -172,6 +888,61 @@ def _call_amplitude_fn(amplitude_fn, configs, *, chunk_size=None): return torch.cat(pieces, dim=0) +def _resolve_log_amplitude_fn(amplitude_fn, log_amplitude_fn=None): + """Resolve an optional ``(phase, log_abs)`` amplitude interface.""" + if log_amplitude_fn is False: + return None + if log_amplitude_fn is not None: + if not callable(log_amplitude_fn): + raise TypeError("log_amplitude_fn must be callable or False.") + return log_amplitude_fn + for name in ("forward_log", "log_amplitude"): + candidate = getattr(amplitude_fn, name, None) + if callable(candidate): + return candidate + return None + + +def _call_log_amplitude_fn(log_amplitude_fn, configs, *, chunk_size=None): + """Evaluate a log-amplitude function in optional fixed-size chunks.""" + torch = _require_torch() + configs = _as_long_matrix(configs) + chunk_size = _normalize_chunk_size(chunk_size) + if chunk_size is None or configs.shape[0] <= chunk_size: + chunks = (configs,) + else: + chunks = ( + configs[start:min(start + chunk_size, configs.shape[0])] + for start in range(0, configs.shape[0], chunk_size) + ) + + phases = [] + log_abs = [] + for chunk in chunks: + result = log_amplitude_fn(chunk) + if not isinstance(result, (tuple, list)) or len(result) != 2: + raise ValueError( + "log_amplitude_fn must return a (phase, log_abs) pair." + ) + phase, chunk_log_abs = result + phase = torch.as_tensor(phase, device=chunk.device) + chunk_log_abs = torch.as_tensor(chunk_log_abs, device=chunk.device) + if phase.ndim != 1 or chunk_log_abs.ndim != 1: + raise ValueError( + "log_amplitude_fn must return one phase and one log magnitude " + "per configuration." + ) + if phase.shape[0] != chunk.shape[0] or ( + chunk_log_abs.shape[0] != chunk.shape[0] + ): + raise ValueError( + "log_amplitude_fn outputs must have one entry per configuration." + ) + phases.append(phase) + log_abs.append(chunk_log_abs.real) + return torch.cat(phases, dim=0), torch.cat(log_abs, dim=0) + + def _diagonal_connection_mask(configs, connections): torch = _require_torch() if connections.configs.numel() == 0: @@ -180,6 +951,68 @@ def _diagonal_connection_mask(configs, connections): return torch.all(connections.configs == parents, dim=1) +def _connection_key_rows(batch_ids, configs): + """Pack fixed-width connection keys before eager unique grouping.""" + torch = _require_torch() + return torch.cat((batch_ids.reshape(-1, 1), configs), dim=1) + + +def _coalesce_connections(connections, *, device=None, compile_kernels=False): + """Merge duplicate ``(batch_id, connected_config)`` rows. + + Local Hamiltonians assembled from several terms can emit the same target + configuration more than once. Summing those coefficients before amplitude + evaluation avoids redundant contractions while preserving the exact local + energy. + """ + torch = _require_torch() + if connections.configs.numel() == 0: + return connections + + target_device = connections.configs.device if device is None else device + configs = connections.configs.to(device=target_device, dtype=torch.long) + batch_ids = connections.batch_ids.to(device=target_device, dtype=torch.long) + coeffs = connections.coeffs.to(device=target_device) + keys = _run_cheap_torch_kernel( + "connection-key-rows", + _connection_key_rows, + batch_ids, + configs, + compile_kernels=compile_kernels, + ) + unique_keys, inverse = torch.unique( + keys, + dim=0, + return_inverse=True, + sorted=False, + ) + unique_coeffs = torch.zeros( + unique_keys.shape[0], + dtype=coeffs.dtype, + device=target_device, + ) + unique_coeffs.index_add_(0, inverse, coeffs) + nonzero = unique_coeffs != 0 + return TorchConnections( + configs=unique_keys[nonzero, 1:], + coeffs=unique_coeffs[nonzero], + batch_ids=unique_keys[nonzero, 0], + ) + + +def _unique_config_rows(configs): + """Return unique configuration rows and an inverse scatter index.""" + torch = _require_torch() + if configs.shape[0] <= 1: + return configs, None + return torch.unique( + configs, + dim=0, + return_inverse=True, + sorted=False, + ) + + def _default_connected_amplitudes( configs, amplitudes, @@ -219,11 +1052,21 @@ def _default_connected_amplitudes( out[diag] = amplitudes[connections.batch_ids[diag]] offdiag = ~diag if bool(torch.any(offdiag)): - out[offdiag] = _call_amplitude_fn( + unique_configs, inverse = _unique_config_rows( + connections.configs[offdiag] + ) + unique_amplitudes = _call_amplitude_fn( amplitude_fn, - connections.configs[offdiag], + unique_configs, chunk_size=chunk_size, - ).to(dtype=out.dtype, device=out.device) + ) + if inverse is None: + out[offdiag] = unique_amplitudes + else: + out[offdiag] = unique_amplitudes[inverse].to( + dtype=out.dtype, + device=out.device, + ) return out @@ -242,6 +1085,19 @@ class TorchPEPSAmplitude: through ``quimb.tensor.pack`` / ``unpack``. For Symmray, this preserves the array's own pytree metadata, including fermionic phases and charge sectors, while replacing numeric block leaves with torch trainable parameters. + + Set ``graded_torch=True`` for native sparse U1U1 PEPS to use the exact + fixed-shape Torch projector. That opt-in path compiles Symmray's graded + charge and phase rules once, then performs the per-configuration dense + contractions under ``torch.vmap``. + + ``amplitude_batching`` controls independent configuration batches. The + default ``"auto"`` probes ``torch.vmap`` once and permanently falls back + to the serial contraction path if the selected PEPS backend is not + vmappable. Use ``"vmap"`` for flat Z2 Symmray PEPS when the fast path is + known to be supported, or ``"serial"`` for native U1/U1U1 PEPS and other + dynamic block-sparse contractions. A failed explicit ``"vmap"`` request + still falls back safely rather than changing numerical semantics. """ def __init__( @@ -250,19 +1106,33 @@ def __init__( *, contraction="exact", chi=None, - cutoff=0.0, + cutoff=None, contraction_opts=None, dtype=None, device=None, site_order=None, + graded_torch=False, + amplitude_batching="auto", ): torch = _require_torch() import quimb as qu import quimb.tensor as qtn + from .api import _resolve_contraction_config + contraction, chi, cutoff, contraction_opts = _resolve_contraction_config( + contraction, + chi, + cutoff, + contraction_opts, + ) + self.contraction = _validate_contraction(contraction, chi) self.chi = None if chi is None else int(chi) - self.cutoff = float(cutoff) + self.cutoff = ( + 0.0 + if cutoff is None and self.contraction == "exact" + else 1.0e-10 if cutoff is None else float(cutoff) + ) self.contraction_opts = _as_contraction_options(contraction_opts) if self.contraction == "boundary": self.contraction_opts.setdefault("mode", "mps") @@ -276,6 +1146,30 @@ def __init__( if missing: raise ValueError(f"site_order contains site(s) not in PEPS: {missing!r}") self.site_inds = tuple(tn.site_ind(site) for site in self.sites) + self.cutoff_fallbacks = 0 + self.final_optimizer_fallbacks = 0 + self._warned_final_optimizer_fallback = False + self.graded_torch = bool(graded_torch) + self.amplitude_batching = _normalize_amplitude_batching( + amplitude_batching + ) + self.last_amplitude_batching = None + if self.graded_torch: + if self.contraction != "exact": + raise ValueError( + "graded_torch currently supports contraction='exact' only." + ) + if not self.symmray_tensor_ids: + raise TypeError( + "graded_torch requires a native Symmray fermionic PEPS." + ) + self._graded_torch_projector = _GradedTorchProjector( + tn, + self.sites, + contraction_opts=self.contraction_opts, + ) + else: + self._graded_torch_projector = None params, skeleton = qtn.pack(tn) flat_params, params_pytree = qu.utils.tree_flatten(params, get_ref=True) @@ -286,6 +1180,25 @@ def __init__( self.params = torch.nn.ParameterList(leaves) self.params_pytree = params_pytree self.skeleton = skeleton + if ( + self.contraction == "ctmrg" + 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" + # ``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. + has_vmap = callable(getattr(torch, "vmap", None)) + self._vmap_forward_enabled = has_vmap + self._vmap_log_enabled = has_vmap @property def is_symmray(self): @@ -367,23 +1280,108 @@ def _select_config(self, tn, config): ) return tn.isel({ind: config[i] for i, ind in enumerate(self.site_inds)}) + def _final_contraction_options(self, *, strip_exponent=None): + """Return the caller's path options for the exact scalar closure.""" + options = dict(self.contraction_opts.get("final_contract_opts") or {}) + if strip_exponent is not None: + options.setdefault("strip_exponent", strip_exponent) + 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: + final_opts = self._final_contraction_options() + try: + return tn.contract(*args, **final_opts) + except KeyError as exc: + # cotengra's ReusableHyperOptimizer can raise this after every + # trial fails, leaving no ``best['tree']``. Preserve a long VMC + # run by falling back only for this known optimizer failure. + if exc.args != ("tree",) or final_opts.get("optimize") in ( + None, + "auto-hq", + ): + raise + self.final_optimizer_fallbacks += 1 + if not self._warned_final_optimizer_fallback: + warnings.warn( + "The supplied final contraction optimizer produced no " + "cotengra tree; retrying affected VMC scalar closures " + "with optimize='auto-hq'.", + RuntimeWarning, + stacklevel=2, + ) + self._warned_final_optimizer_fallback = True + fallback_opts = dict(final_opts) + fallback_opts["optimize"] = "auto-hq" + return tn.contract(*args, **fallback_opts) + + def _contract_approximate(self, fn, *args, close_final=False, **kwargs): + """Contract with the requested cutoff, retrying empty sparse sectors.""" + kwargs = dict(kwargs) + kwargs["cutoff"] = self.cutoff + if close_final: + # Close the final small tensor network here rather than inside + # Quimb. This makes ``final_contract_opts`` apply identically to + # full and cached boundary paths, and lets us recover from a + # failed reusable cotengra search at the scalar-closure boundary. + # Environment builders do not accept this option, hence the + # explicit opt-in at amplitude call sites below. + kwargs["final_contract"] = False + + def finish(value): + if not close_final: + return value + # Boundary, CTMRG, and HOTRG now all return the partially + # contracted flat TN here. Amplitude evaluation still needs a + # scalar, so complete that last contraction with the requested + # optimizer and strip-exponent options. + contract = getattr(value, "contract", None) + if not callable(contract): + return value + return self._contract_remaining( + value, + final_opts=self._final_contraction_options( + strip_exponent=kwargs.get("strip_exponent"), + ), + ) + + try: + return finish(fn(*args, **kwargs)) + except Exception: # pragma: no cover - exact upstream exception varies + if not self.symmray_tensor_ids or self.cutoff <= 0.0: + raise + retry_kwargs = dict(kwargs) + retry_kwargs["cutoff"] = 0.0 + compress_opts = retry_kwargs.get("compress_opts") + if isinstance(compress_opts, dict) and compress_opts.get("method") == "cholesky": + retry_kwargs["compress_opts"] = { + **compress_opts, + "method": "svd", + } + self.cutoff_fallbacks += 1 + return finish(fn(*args, **retry_kwargs)) + def _contract_value(self, tnx, reference=None): if self.contraction == "hotrg": - value = tnx.contract_hotrg( + value = self._contract_approximate( + tnx.contract_hotrg, max_bond=self.chi, - cutoff=self.cutoff, + close_final=True, **self.contraction_opts, ) elif self.contraction == "ctmrg": - value = tnx.contract_ctmrg( + value = self._contract_approximate( + tnx.contract_ctmrg, max_bond=self.chi, - cutoff=self.cutoff, + close_final=True, **self.contraction_opts, ) elif self.contraction == "boundary": - value = tnx.contract_boundary( + value = self._contract_approximate( + tnx.contract_boundary, max_bond=self.chi, - cutoff=self.cutoff, + close_final=True, **self.contraction_opts, ) else: @@ -393,24 +1391,27 @@ def _contract_value(self, tnx, reference=None): def _contract_log_parts(self, tnx, reference=None): torch = _require_torch() if self.contraction == "hotrg": - mantissa, exponent_10 = tnx.contract_hotrg( + mantissa, exponent_10 = self._contract_approximate( + tnx.contract_hotrg, max_bond=self.chi, - cutoff=self.cutoff, strip_exponent=True, + close_final=True, **self.contraction_opts, ) elif self.contraction == "ctmrg": - mantissa, exponent_10 = tnx.contract_ctmrg( + mantissa, exponent_10 = self._contract_approximate( + tnx.contract_ctmrg, max_bond=self.chi, - cutoff=self.cutoff, strip_exponent=True, + close_final=True, **self.contraction_opts, ) elif self.contraction == "boundary": - mantissa, exponent_10 = tnx.contract_boundary( + mantissa, exponent_10 = self._contract_approximate( + tnx.contract_boundary, max_bond=self.chi, - cutoff=self.cutoff, strip_exponent=True, + close_final=True, **self.contraction_opts, ) else: @@ -449,19 +1450,88 @@ def _contract_log_parts(self, tnx, reference=None): ) return phase, log_abs + def _graded_torch_forward(self, configs, *, params=None): + """Evaluate configs through the fixed graded Torch projector.""" + if self._graded_torch_projector is None: # pragma: no cover - guard + raise RuntimeError("graded_torch projector is not initialized.") + tn = self._unpack_tn(params) + dense_leaves = [tn[site].data.to_dense() for site in self.sites] + return self._graded_torch_projector( + dense_leaves, + configs, + use_vmap=self.amplitude_batching != "serial", + ) + def amplitude(self, config, params=None): """Evaluate a single configuration amplitude.""" config = _as_long_matrix(config).reshape(-1) + if self.graded_torch: + return self._graded_torch_forward(config.reshape(1, -1), params=params)[0] tn = self._unpack_tn(params) reference = self._reference_tensor(params) return self._contract_value(self._select_config(tn, config), reference) + def _try_vmapped_forward(self, configs, *, params=None): + """Attempt a native batched amplitude contraction. + + Symmray fermionic arrays can support ``torch.vmap`` directly when the + selected contraction route is pure and all required block operations + have batching rules. A failed trace disables only this optional fast + path; the established serial route remains numerically authoritative. + """ + torch = _require_torch() + if self.graded_torch: + return None + if self.amplitude_batching == "serial": + return None + if not self._vmap_forward_enabled: + return None + try: + return torch.vmap( + lambda config: self.amplitude(config, params=params) + )(configs) + except (AttributeError, IndexError, KeyError, NotImplementedError, + RuntimeError, TypeError, ValueError): + self._vmap_forward_enabled = False + return None + + def _try_vmapped_forward_log(self, configs, *, params=None): + """Attempt a vmapped phase/log-magnitude contraction. + + Stable-log sampling must not silently turn the flat-Z2 fast path back + into a scalar loop. Exact contractions use the same pure selected-TN + closure as :meth:`amplitude`, while approximate boundary/CTMRG paths + simply decline this optional route and retain their serial fallback. + """ + torch = _require_torch() + if self.graded_torch or self.amplitude_batching == "serial": + return None + if not self._vmap_log_enabled: + return None + tn = self._unpack_tn(params) + reference = self._reference_tensor(params) + + def evaluate(config): + return self._contract_log_parts( + self._select_config(tn, config), + reference, + ) + + try: + phases, log_abs = torch.vmap(evaluate)(configs) + except (AttributeError, IndexError, KeyError, NotImplementedError, + RuntimeError, TypeError, ValueError): + self._vmap_log_enabled = False + return None + return phases, log_abs + def forward(self, configs, params=None, *, chunk_size=None): """Evaluate a batch of configuration amplitudes.""" + torch = _require_torch() configs = _as_long_matrix(configs) chunk_size = _normalize_chunk_size(chunk_size) if chunk_size is not None and configs.shape[0] > chunk_size: - return _require_torch().cat([ + return torch.cat([ self.forward( configs[start:start + chunk_size], params=params, @@ -469,9 +1539,24 @@ def forward(self, configs, params=None, *, chunk_size=None): ) for start in range(0, configs.shape[0], chunk_size) ]) + + if self.graded_torch: + self.last_amplitude_batching = ( + "graded-vmap" + if self.amplitude_batching != "serial" + else "graded-serial" + ) + return self._graded_torch_forward(configs, params=params) + + vmapped = self._try_vmapped_forward(configs, params=params) + if vmapped is not None: + self.last_amplitude_batching = "vmap" + return vmapped + + self.last_amplitude_batching = "serial" tn = self._unpack_tn(params) reference = self._reference_tensor(params) - return _require_torch().stack([ + return torch.stack([ self._contract_value(self._select_config(tn, row), reference) for row in configs ]) @@ -479,6 +1564,22 @@ def forward(self, configs, params=None, *, chunk_size=None): def forward_log(self, configs, params=None): """Return ``(phase, log_abs)`` for a batch of configurations.""" configs = _as_long_matrix(configs) + if self.graded_torch: + torch = _require_torch() + amp = self.forward(configs, params=params) + abs_amp = amp.abs() + tiny = _torch_finfo_tiny(abs_amp.dtype) + phase = torch.where( + abs_amp > 0, + amp / abs_amp.to(dtype=amp.dtype), + torch.zeros_like(amp), + ) + return phase, torch.log(abs_amp.clamp_min(tiny)) + vmapped = self._try_vmapped_forward_log(configs, params=params) + if vmapped is not None: + self.last_amplitude_batching = "log-vmap" + return vmapped + self.last_amplitude_batching = "serial" tn = self._unpack_tn(params) reference = self._reference_tensor(params) phases = [] @@ -541,12 +1642,17 @@ def __init__( *, contraction="boundary", chi=None, - cutoff=0.0, + cutoff=None, contraction_opts=None, dtype=None, device=None, site_order=None, + graded_torch=False, + amplitude_batching="auto", environment_radius=0, + boundary_cache_size=128, + proposal_batching="auto", + proposal_vmap_min_batch=8, ): super().__init__( peps, @@ -557,60 +1663,435 @@ def __init__( dtype=dtype, device=device, site_order=site_order, + graded_torch=graded_torch, + amplitude_batching=amplitude_batching, ) self.environment_radius = _check_nonnegative_int( "environment_radius", environment_radius, ) + self.boundary_cache_size = _check_positive_int( + "boundary_cache_size", + boundary_cache_size, + ) + self.proposal_batching = _normalize_proposal_batching(proposal_batching) + self.proposal_vmap_min_batch = _check_positive_int( + "proposal_vmap_min_batch", + proposal_vmap_min_batch, + ) self._boundary_geometry = self._infer_boundary_geometry(self._unpack_tn()) self.last_connected_reuse_stats = None + self.last_proposal_cache_stats = None + self.last_amplitude_cache_stats = None + self._boundary_cache_token = None + self._boundary_environment_cache = {} + self._boundary_transition_cache = {} + self._boundary_amplitude_cache = {} + # Parent-selected strip templates are much smaller than a full PEPS + # contraction and let connected local estimators replace only their + # changed physical projectors. + self._boundary_strip_cache = {} + + def _parameter_cache_token(self): + """Return a cheap token that changes when torch leaves are updated.""" + return tuple(int(getattr(param, "_version", 0)) for param in self.params) + + def clear_boundary_cache(self): + """Clear cached boundary environments and proposal transitions.""" + self._boundary_environment_cache.clear() + self._boundary_transition_cache.clear() + self._boundary_strip_cache.clear() + self._boundary_amplitude_cache.clear() + self._boundary_cache_token = self._parameter_cache_token() + self.last_connected_reuse_stats = None + self.last_proposal_cache_stats = None + self.last_amplitude_cache_stats = None + return self - def _infer_boundary_geometry(self, tn): - try: - Lx = getattr(tn, "Lx", None) - Ly = getattr(tn, "Ly", None) - if Lx is None: - Lx = getattr(tn, "_Lx") - if Ly is None: - Ly = getattr(tn, "_Ly") - Lx = int(Lx) - Ly = int(Ly) - view_kwargs = { - "site_tag_id": getattr(tn, "_site_tag_id"), - "x_tag_id": getattr(tn, "_x_tag_id"), - "y_tag_id": getattr(tn, "_y_tag_id"), - "site_ind_id": getattr(tn, "_site_ind_id"), - "Lx": Lx, - "Ly": Ly, - } - except AttributeError: - return None + def to(self, *args, **kwargs): + super().to(*args, **kwargs) + self.clear_boundary_cache() + return self - coords = [] - for site in self.sites: - if not isinstance(site, tuple) or len(site) != 2: - return None - try: - x, y = int(site[0]), int(site[1]) - except (TypeError, ValueError): - return None - if not (0 <= x < Lx and 0 <= y < Ly): - return None - coords.append((x, y)) + def _ensure_boundary_cache_current(self): + token = self._parameter_cache_token() + if token != self._boundary_cache_token: + self._boundary_environment_cache.clear() + self._boundary_transition_cache.clear() + self._boundary_strip_cache.clear() + self._boundary_amplitude_cache.clear() + self._boundary_cache_token = token + + @staticmethod + def _configuration_key(config): + return tuple(int(value) for value in config.detach().cpu().tolist()) + + def _cache_put(self, cache, key, value, *, max_size=None): + cache[key] = value + if max_size is None: + max_size = self.boundary_cache_size + while len(cache) > max_size: + cache.pop(next(iter(cache))) - return { - "Lx": Lx, - "Ly": Ly, - "coords": tuple(coords), - "view_kwargs": view_kwargs, - } + def forward(self, configs, params=None, *, chunk_size=None): + """Evaluate amplitudes, caching serial boundary contractions safely. - def _changed_axis_window(self, parent_config, target_config): + Boundary amplitudes are cached only for detached/no-grad calls using + the model's own parameters. Gradient-enabled contractions and custom + parameter pytrees always use the base implementation, so the cache + never retains an autograd graph or returns stale derivatives. + """ torch = _require_torch() - changed = torch.nonzero(parent_config != target_config, as_tuple=True)[0] - if changed.numel() == 0: + configs = _as_long_matrix(configs) + vmap_preferred = ( + self.amplitude_batching != "serial" + and self._vmap_forward_enabled + ) + if ( + params is not None + or torch.is_grad_enabled() + or self.contraction != "boundary" + or self.graded_torch + or vmap_preferred + ): + return super().forward(configs, params=params, chunk_size=chunk_size) + if configs.shape[0] == 0: + return super().forward(configs, params=params, chunk_size=chunk_size) + + self._ensure_boundary_cache_current() + unique_configs, inverse = _unique_config_rows(configs) + if inverse is None: + inverse = torch.zeros( + 1, + dtype=torch.long, + device=configs.device, + ) + cached_values = [None] * int(unique_configs.shape[0]) + missing_indices = [] + num_hits = 0 + for index, config in enumerate(unique_configs): + key = self._configuration_key(config) + try: + cached_values[index] = self._boundary_amplitude_cache[key] + except KeyError: + missing_indices.append(index) + else: + num_hits += 1 + + if missing_indices: + missing = torch.as_tensor( + missing_indices, + dtype=torch.long, + device=configs.device, + ) + computed = TorchPEPSAmplitude.forward( + self, + unique_configs[missing], + params=params, + chunk_size=chunk_size, + ) + for offset, index in enumerate(missing_indices): + value = computed[offset].detach() + cached_values[index] = value + self._cache_put( + self._boundary_amplitude_cache, + self._configuration_key(unique_configs[index]), + value, + ) + + unique_amplitudes = torch.stack([ + value.to(device=configs.device) + for value in cached_values + ]) + self.last_amplitude_batching = "serial" + self.last_amplitude_cache_stats = { + "num_requests": int(configs.shape[0]), + "num_unique_requests": int(unique_configs.shape[0]), + "num_hits": num_hits, + "num_misses": len(missing_indices), + } + return unique_amplitudes[inverse] + + def _cached_boundary_environments( + self, + tn, + parent_config, + axis, + *, + parent_tn=None, + ): + """Get one walker's MPS environments, retaining them across sweeps.""" + self._ensure_boundary_cache_current() + key = (axis, self._configuration_key(parent_config)) + try: + return self._boundary_environment_cache[key], True + except KeyError: + if parent_tn is None: + parent_tn = self._select_config(tn, parent_config) + envs = self._compute_boundary_environments(parent_tn, axis) + self._cache_put(self._boundary_environment_cache, key, envs) + return envs, False + + def _cached_boundary_strip( + self, + tn, + parent_config, + axis, + indices, + *, + parent_tn=None, + ): + """Get a parent-selected strip template for local impurity updates.""" + self._ensure_boundary_cache_current() + parent_key = self._configuration_key(parent_config) + key = (axis, tuple(indices), parent_key) + try: + return self._boundary_strip_cache[key], True + except KeyError: + if parent_tn is None: + parent_tn = self._select_config(tn, parent_config) + tags = ( + [tn.x_tag(index) for index in indices] + if axis == "x" + else [tn.y_tag(index) for index in indices] + ) + strip_tn = parent_tn.select(tags, which="any") + # Keep this LRU smaller than the environment cache: a long-range + # observable can otherwise retain many selected PEPS strips. + self._cache_put( + self._boundary_strip_cache, + key, + strip_tn, + max_size=min(self.boundary_cache_size, 32), + ) + return strip_tn, False + + def _boundary_transition_amplitude( + self, + tn, + parent_config, + target_config, + reference, + ): + """Evaluate one local proposal using the parent's cached boundaries.""" + if self._boundary_geometry is None: + return self._contract_value( + self._select_config(tn, target_config), + reference, + ), 0, 0, False + + windows = self._changed_axis_windows(parent_config, target_config) + if not windows: + return self._contract_value( + self._select_config(tn, target_config), + reference, + ), 0, 0, False + + # A periodic Hamiltonian edge is local in the transverse boundary + # sweep direction even when its endpoints are far apart in the + # longitudinal coordinate. Try that short strip first. Some upstream + # tensor backends reject a particular sweep direction, so use the + # other cached boundary direction before a full contraction fallback. + num_environment_hits = 0 + num_environment_builds = 0 + for window_index, (axis, indices) in enumerate(windows): + try: + envs, reused = self._cached_boundary_environments( + tn, + parent_config, + axis, + ) + value = self._contract_axis_window( + tn, + target_config, + axis, + indices, + envs, + reference, + ) + except Exception: # pragma: no cover - upstream exceptions vary + continue + if reused: + num_environment_hits += 1 + else: + num_environment_builds += 1 + return ( + value, + num_environment_hits, + num_environment_builds, + window_index > 0, + ) + + return self._contract_value( + self._select_config(tn, target_config), + reference, + ), num_environment_hits, num_environment_builds, False + + def _should_vmap_proposals(self, *, n_changed, device): + """Whether this proposal batch should prefer full vmapped amplitudes.""" + if not self._vmap_forward_enabled: + return False + if self.proposal_batching == "cache": + return False + if self.proposal_batching == "vmap": + return True + return ( + device.type == "cuda" + and n_changed >= self.proposal_vmap_min_batch + ) + + def proposal_amplitudes( + self, + parent_configs, + target_configs, + current_amplitudes, + *, + chunk_size=None, + ): + """Evaluate local Metropolis proposals with cached MPS environments. + + The cache is deliberately attached to the amplitude model rather than + the sampler. It therefore survives burn-in/thinning sweeps, while + :meth:`clear_boundary_cache` invalidates it when VMC parameters change. + Unsupported geometries fall back to ordinary batched amplitudes. + """ + torch = _require_torch() + parent_configs = _as_long_matrix(parent_configs) + target_configs = _as_long_matrix(target_configs) + if parent_configs.shape != target_configs.shape: + raise ValueError( + "parent_configs and target_configs must have matching shapes." + ) + current_amplitudes = torch.as_tensor( + current_amplitudes, + device=parent_configs.device, + ) + if current_amplitudes.shape != (parent_configs.shape[0],): + raise ValueError( + "current_amplitudes must have one value per proposal." + ) + if self.contraction != "boundary" or self._boundary_geometry is None: + return _call_amplitude_fn( + self, + target_configs, + chunk_size=chunk_size, + ) + + self._ensure_boundary_cache_current() + out = current_amplitudes.clone() + stats = { + "num_requests": int(parent_configs.shape[0]), + "num_vmapped": 0, + "num_transition_cache_hits": 0, + "num_environment_cache_hits": 0, + "num_environment_builds": 0, + "num_alternative_axis_reused": 0, + "num_fallback": 0, + } + changed = torch.any(parent_configs != target_configs, dim=1) + n_changed = int(changed.sum().item()) + if n_changed and self._should_vmap_proposals( + n_changed=n_changed, + device=parent_configs.device, + ): + vmapped = self._try_vmapped_forward(target_configs[changed]) + if vmapped is not None: + out[changed] = vmapped.to(dtype=out.dtype, device=out.device) + stats["num_vmapped"] = n_changed + stats["num_fallback"] = int( + (~torch.isfinite(vmapped)).sum().item() + ) + self.last_proposal_cache_stats = stats + return out + + tn = self._unpack_tn() + reference = self._reference_tensor() + for index in range(parent_configs.shape[0]): + parent_config = parent_configs[index] + target_config = target_configs[index] + if torch.equal(parent_config, target_config): + continue + parent_key = self._configuration_key(parent_config) + target_key = self._configuration_key(target_config) + cache_key = (parent_key, target_key) + try: + value = self._boundary_transition_cache[cache_key] + stats["num_transition_cache_hits"] += 1 + except KeyError: + ( + value, + num_environment_hits, + num_environment_builds, + alternative_axis, + ) = self._boundary_transition_amplitude( + tn, + parent_config, + target_config, + reference, + ) + stats["num_environment_cache_hits"] += num_environment_hits + stats["num_environment_builds"] += num_environment_builds + if alternative_axis: + stats["num_alternative_axis_reused"] += 1 + self._cache_put(self._boundary_transition_cache, cache_key, value) + if not torch.isfinite(torch.as_tensor(value)).all(): + stats["num_fallback"] += 1 + out[index] = value + self.last_proposal_cache_stats = stats + return out + + def _infer_boundary_geometry(self, tn): + try: + Lx = getattr(tn, "Lx", None) + Ly = getattr(tn, "Ly", None) + if Lx is None: + Lx = getattr(tn, "_Lx") + if Ly is None: + Ly = getattr(tn, "_Ly") + Lx = int(Lx) + Ly = int(Ly) + view_kwargs = { + "site_tag_id": getattr(tn, "_site_tag_id"), + "x_tag_id": getattr(tn, "_x_tag_id"), + "y_tag_id": getattr(tn, "_y_tag_id"), + "site_ind_id": getattr(tn, "_site_ind_id"), + "Lx": Lx, + "Ly": Ly, + } + except AttributeError: return None + coords = [] + for site in self.sites: + if not isinstance(site, tuple) or len(site) != 2: + return None + try: + x, y = int(site[0]), int(site[1]) + except (TypeError, ValueError): + return None + if not (0 <= x < Lx and 0 <= y < Ly): + return None + coords.append((x, y)) + + return { + "Lx": Lx, + "Ly": Ly, + "coords": tuple(coords), + "view_kwargs": view_kwargs, + } + + def _changed_axis_windows(self, parent_config, target_config): + """Return cached boundary windows ordered by estimated work. + + Both directions are available as a safe fallback. In particular, a + wrap-around Hamiltonian bond has two distant endpoints in one lattice + direction but only one changed plane in the transverse direction. + """ + torch = _require_torch() + changed = torch.nonzero(parent_config != target_config, as_tuple=True)[0] + if changed.numel() == 0: + return () + geometry = self._boundary_geometry coords = [geometry["coords"][int(i)] for i in changed.detach().cpu()] xs = [coord[0] for coord in coords] @@ -621,33 +2102,91 @@ def _changed_axis_window(self, parent_config, target_config): y0 = max(0, min(ys) - radius) y1 = min(geometry["Ly"], max(ys) + radius + 1) - if (x1 - x0) <= (y1 - y0): - return "x", tuple(range(x0, x1)) - return "y", tuple(range(y0, y1)) + windows = ( + ("x", tuple(range(x0, x1))), + ("y", tuple(range(y0, y1))), + ) + # A complete plane has the other lattice extent. This is equivalent to + # comparing window widths on a fixed geometry, but makes the choice + # explicit and deterministic for separated/PBC updates. + return tuple(sorted( + windows, + key=lambda item: ( + len(item[1]) * ( + geometry["Ly"] if item[0] == "x" else geometry["Lx"] + ), + item[0], + ), + )) + + def _changed_axis_window(self, parent_config, target_config): + """Return the preferred cached boundary window for compatibility.""" + windows = self._changed_axis_windows(parent_config, target_config) + return None if not windows else windows[0] + + def _boundary_environment_options(self): + """Adapt full-boundary options to Quimb environment builders.""" + options = dict(self.contraction_opts) + # These belong to ``contract_boundary``'s final scalar contraction, + # not to ``compute_*_environments`` or ``contract_boundary_from_*``. + for key in ( + "final_contract", + "final_contract_opts", + "sequence", + "inplace", + "progbar", + "optimize", + "max_separation", + ): + options.pop(key, None) + return options def _compute_boundary_environments(self, parent_tn, axis): + options = self._boundary_environment_options() if axis == "x": - return parent_tn.compute_x_environments( + return self._contract_approximate( + parent_tn.compute_x_environments, max_bond=self.chi, - cutoff=self.cutoff, - **self.contraction_opts, + **options, ) - return parent_tn.compute_y_environments( + return self._contract_approximate( + parent_tn.compute_y_environments, max_bond=self.chi, - cutoff=self.cutoff, - **self.contraction_opts, + **options, ) - def _contract_axis_window(self, tn, target_config, axis, indices, envs, reference): + def _replace_strip_projectors( + self, + tn, + strip_tn, + parent_config, + target_config, + ): + """Copy a parent strip and replace only its changed physical tensors.""" + torch = _require_torch() + changed = torch.nonzero(parent_config != target_config, as_tuple=True)[0] + target_strip = strip_tn.copy() + for config_index in changed.detach().cpu().tolist(): + site = self.sites[int(config_index)] + site_tag = tn.site_tag(site) + physical_index = tn.site_ind(site) + selected_tensor = tn[site_tag].isel({ + physical_index: int(target_config[config_index].item()), + }) + # ``strip_tn.copy()`` clones tensor objects while sharing immutable + # parent data, so modifying this tensor leaves the cached template + # untouched for the next connected configuration. + target_strip[site_tag].modify(data=selected_tensor.data) + return target_strip + + def _contract_axis_strip(self, tn, strip_tn, axis, indices, envs, reference): import quimb.tensor as qtn - target_tn = self._select_config(tn, target_config) + options = self._boundary_environment_options() first = indices[0] last = indices[-1] if axis == "x": - tags = [tn.x_tag(i) for i in indices] - window_tn = target_tn.select(tags, which="any") - reuse_tn = envs[("xmin", first)] | window_tn | envs[("xmax", last)] + reuse_tn = envs[("xmin", first)] | strip_tn | envs[("xmax", last)] reuse_tn.view_as_( qtn.PEPS, **self._boundary_geometry["view_kwargs"], @@ -656,12 +2195,10 @@ def _contract_axis_window(self, tn, target_config, axis, indices, envs, referenc xrange=[first, last + 1], max_bond=self.chi, cutoff=self.cutoff, - **self.contraction_opts, + **options, ) else: - tags = [tn.y_tag(i) for i in indices] - window_tn = target_tn.select(tags, which="any") - reuse_tn = envs[("ymin", first)] | window_tn | envs[("ymax", last)] + reuse_tn = envs[("ymin", first)] | strip_tn | envs[("ymax", last)] reuse_tn.view_as_( qtn.PEPS, **self._boundary_geometry["view_kwargs"], @@ -670,9 +2207,57 @@ def _contract_axis_window(self, tn, target_config, axis, indices, envs, referenc yrange=[first, last + 1], max_bond=self.chi, cutoff=self.cutoff, - **self.contraction_opts, + **options, ) - return _as_torch_scalar(reuse_tn.contract(all), reference) + return _as_torch_scalar( + self._contract_remaining(reuse_tn, all), + reference, + ) + + def _contract_cached_axis_window( + self, + tn, + parent_config, + target_config, + axis, + indices, + envs, + strip_tn, + reference, + ): + """Contract a cached parent strip with its local target impurities.""" + target_strip = self._replace_strip_projectors( + tn, + strip_tn, + parent_config, + target_config, + ) + return self._contract_axis_strip( + tn, + target_strip, + axis, + indices, + envs, + reference, + ) + + def _contract_axis_window(self, tn, target_config, axis, indices, envs, reference): + """Compatibility contraction path that rebuilds the selected strip.""" + target_tn = self._select_config(tn, target_config) + tags = ( + [tn.x_tag(index) for index in indices] + if axis == "x" + else [tn.y_tag(index) for index in indices] + ) + strip_tn = target_tn.select(tags, which="any") + return self._contract_axis_strip( + tn, + strip_tn, + axis, + indices, + envs, + reference, + ) def connected_amplitudes( self, @@ -691,6 +2276,11 @@ def connected_amplitudes( self.last_connected_reuse_stats = { "num_diagonal": 0, "num_reused": 0, + "num_groups": 0, + "num_grouped_connections": 0, + "num_strip_cache_hits": 0, + "num_strip_builds": 0, + "num_alternative_axis_reused": 0, "num_fallback": 0, } return torch.empty(0, dtype=amplitudes.dtype, device=configs.device) @@ -711,15 +2301,15 @@ def connected_amplitudes( self.last_connected_reuse_stats = { "num_diagonal": num_diagonal, "num_reused": 0, + "num_groups": 0, + "num_grouped_connections": 0, + "num_strip_cache_hits": 0, + "num_strip_builds": 0, + "num_alternative_axis_reused": 0, "num_fallback": int(connections.configs.shape[0]) - num_diagonal, } return result - out = torch.empty( - connections.configs.shape[0], - dtype=amplitudes.dtype, - device=configs.device, - ) diag = ( _diagonal_connection_mask(configs, connections) if reuse_diagonal @@ -729,71 +2319,228 @@ def connected_amplitudes( device=configs.device, ) ) + offdiag = (~diag).nonzero(as_tuple=True)[0] + if ( + self.amplitude_batching != "serial" + and self._vmap_forward_enabled + and offdiag.numel() >= _BOUNDARY_VMAP_CONNECTION_THRESHOLD + ): + result = super().connected_amplitudes( + configs, + amplitudes, + connections, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + ) + self.last_connected_reuse_stats = { + "num_diagonal": int(diag.sum().item()), + "num_reused": 0, + "num_batched": int(offdiag.numel()), + "num_groups": 0, + "num_grouped_connections": 0, + "num_strip_cache_hits": 0, + "num_strip_builds": 0, + "num_alternative_axis_reused": 0, + "num_fallback": 0, + } + return result + + out = torch.empty( + connections.configs.shape[0], + dtype=amplitudes.dtype, + device=configs.device, + ) if bool(torch.any(diag)): out[diag] = amplitudes[connections.batch_ids[diag]] + self._ensure_boundary_cache_current() tn = self._unpack_tn() reference = self._reference_tensor() - env_cache = {} stats = { "num_diagonal": int(diag.sum().item()), "num_reused": 0, + "num_groups": 0, + "num_grouped_connections": 0, + "num_environment_cache_hits": 0, + "num_environment_builds": 0, + "num_strip_cache_hits": 0, + "num_strip_builds": 0, + "num_alternative_axis_reused": 0, "num_fallback": 0, } - offdiag = (~diag).nonzero(as_tuple=True)[0] + # Group by the first/cheapest boundary strip. Local Hamiltonian terms + # often share a parent and a single transverse plane, especially for + # PBC edges. They can then reuse one environment pair and one selected + # parent-strip template while only their changed projectors differ. + groups = {} for conn_idx_tensor in offdiag: conn_idx = int(conn_idx_tensor) parent_idx = int(connections.batch_ids[conn_idx].item()) parent_config = configs[parent_idx] target_config = connections.configs[conn_idx] - window = self._changed_axis_window(parent_config, target_config) - if window is None: + 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, ) stats["num_fallback"] += 1 continue + axis, indices = windows[0] + groups.setdefault((parent_idx, axis, indices), []).append( + (conn_idx, windows) + ) - axis, indices = window - cache_key = (parent_idx, axis) - try: - envs = env_cache[cache_key] - except KeyError: - parent_tn = self._select_config(tn, parent_config) - envs = self._compute_boundary_environments(parent_tn, axis) - env_cache[cache_key] = envs + stats["num_groups"] = len(groups) + stats["num_grouped_connections"] = sum( + len(entries) for entries in groups.values() + ) + parent_tns = {} + contexts = {} + def get_context(parent_idx, parent_config, axis, indices): + cache_key = (parent_idx, axis, indices) + if cache_key in contexts: + return contexts[cache_key] try: - out[conn_idx] = self._contract_axis_window( - tn, - target_config, - axis, - indices, - envs, - reference, - ) - stats["num_reused"] += 1 - except ( - AttributeError, - KeyError, - NotImplementedError, - TypeError, - ValueError, - ): - out[conn_idx] = self._contract_value( - self._select_config(tn, target_config), - reference, - ) - stats["num_fallback"] += 1 + parent_key = self._configuration_key(parent_config) + environment_key = (axis, parent_key) + strip_key = (axis, tuple(indices), parent_key) + envs = self._boundary_environment_cache.get(environment_key) + strip_tn = self._boundary_strip_cache.get(strip_key) + environment_reused = envs is not None + strip_reused = strip_tn is not None + if not (environment_reused and strip_reused): + parent_tn = parent_tns.get(parent_idx) + if parent_tn is None: + parent_tn = self._select_config(tn, parent_config) + parent_tns[parent_idx] = parent_tn + if not environment_reused: + envs, environment_reused = ( + self._cached_boundary_environments( + tn, + parent_config, + axis, + parent_tn=parent_tn, + ) + ) + if not strip_reused: + strip_tn, strip_reused = self._cached_boundary_strip( + tn, + parent_config, + axis, + indices, + parent_tn=parent_tn, + ) + except Exception: # pragma: no cover - upstream exceptions vary + contexts[cache_key] = None + return None + stats[ + "num_environment_cache_hits" if environment_reused + else "num_environment_builds" + ] += 1 + stats["num_strip_cache_hits" if strip_reused else "num_strip_builds"] += 1 + contexts[cache_key] = (envs, strip_tn) + return contexts[cache_key] + + for (parent_idx, axis, indices), entries in groups.items(): + parent_config = configs[parent_idx] + primary_context = get_context( + parent_idx, + parent_config, + axis, + 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 + + if value_found: + continue + + for window_index, (alt_axis, alt_indices) in enumerate( + windows[1:], + start=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, + 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), + reference, + ) + stats["num_fallback"] += 1 self.last_connected_reuse_stats = stats return out def make_torch_peps_amplitude_model(peps, **kwargs): - """Build a :class:`TorchPEPSAmplitude` from a quimb PEPS-like object.""" + """Build the appropriate torch amplitude model for ``contraction``.""" + from .api import _resolve_contraction_config + contraction, chi, cutoff, contraction_opts = _resolve_contraction_config( + kwargs.get("contraction", "exact"), + kwargs.get("chi"), + kwargs.get("cutoff"), + kwargs.get("contraction_opts"), + ) + kwargs = dict(kwargs) + kwargs.update( + contraction=contraction, + chi=chi, + cutoff=cutoff, + contraction_opts=contraction_opts, + ) + contraction = _validate_contraction( + contraction, + chi, + ) + if contraction == "boundary": + return TorchPEPSBoundaryAmplitude(peps, **kwargs) return TorchPEPSAmplitude(peps, **kwargs) @@ -831,68 +2578,301 @@ def _flatten_torch_tensors(tensors, refs): return torch.cat(pieces) if pieces else torch.empty(0) -def torch_log_derivative_matrix( +def _log_derivative_denominator(amplitudes, amplitude_floor): + """Build stable per-sample denominators for log-amplitude derivatives.""" + torch = _require_torch() + amplitudes = torch.as_tensor(amplitudes) + amplitude_abs = amplitudes.detach().abs() + if amplitude_floor is None: + if bool(torch.any(amplitude_abs == 0)): + raise ZeroDivisionError( + "Encountered a zero amplitude while forming log derivatives." + ) + return amplitudes + + floor = torch.as_tensor( + amplitude_floor, + dtype=( + amplitudes.real.dtype + if torch.is_complex(amplitudes) + else amplitudes.dtype + ), + device=amplitudes.device, + ) + if torch.is_complex(amplitudes): + phase = torch.where( + amplitude_abs > 0, + amplitudes / amplitude_abs.to(dtype=amplitudes.dtype), + torch.ones_like(amplitudes), + ) + return torch.where( + amplitude_abs < floor, + phase * floor.to(dtype=amplitudes.dtype), + amplitudes, + ) + + sign = torch.where( + amplitudes.detach() >= 0, + torch.ones_like(amplitudes), + -torch.ones_like(amplitudes), + ) + return torch.where(amplitude_abs < floor, sign * floor, amplitudes) + + +def _batched_model_log_derivatives( model, configs, *, - amplitude_floor=None, - create_graph=False, + amplitude_floor, + create_graph, + complex_parameter_mode, ): - """Return per-sample log-amplitude derivatives for a torch model. + """Evaluate PEPS log derivatives with one batched Jacobian graph. - The returned matrix has shape ``(n_samples, n_params)`` and entries - ``d psi(config) / d theta / psi(config)``. It is intended for real-valued - PEPS amplitudes; complex-amplitude SR needs an explicit real/imaginary - parameter convention and is not silently guessed here. + ``TorchPEPSAmplitude`` accepts an explicit parameter tuple, which lets + ``torch.autograd.functional.jacobian`` differentiate all walker + amplitudes together without mutating the model's registered parameters. + The function deliberately targets that parameterized PEPS interface; + generic amplitude models use the compatibility loop instead. """ torch = _require_torch() + if not callable(getattr(model, "_params_pytree", None)): + raise TypeError("model does not expose the functional PEPS parameter API.") configs = _as_long_matrix(configs) - params = _torch_model_parameters(model) - rows = [] + params = tuple(_torch_model_parameters(model)) + if not params: + raise ValueError("model must expose at least one trainable parameter.") - for config in configs: - amp = model(config.reshape(1, -1)) - amp = torch.as_tensor(amp).reshape(-1) - if amp.numel() != 1: - raise ValueError("model(config) must return one amplitude per row.") - amp = amp[0] - if torch.is_complex(amp): - raise NotImplementedError( - "torch_log_derivative_matrix currently supports real scalar " - "amplitudes. Split complex parameters/amplitudes explicitly " - "before using torch SR." - ) - if not amp.requires_grad: - raise RuntimeError("model amplitude does not require gradients.") + mode = str(complex_parameter_mode).replace("_", "-").lower() + if mode not in { + "holomorphic", + "holomorphic-complex", + "real-imag", + "real-imaginary", + }: + raise ValueError( + "complex_parameter_mode must be 'holomorphic' or 'real-imag'." + ) + real_imag_mode = mode in {"real-imag", "real-imaginary"} - amp_abs = amp.detach().abs() - if amplitude_floor is None: - if amp_abs.item() == 0: - raise ZeroDivisionError( - "Encountered a zero amplitude while forming log derivatives." - ) - denom = amp - else: - floor = torch.as_tensor( - amplitude_floor, - dtype=amp.dtype, - device=amp.device, - ) - sign = torch.where( - amp.detach() >= 0, - torch.ones_like(amp), - -torch.ones_like(amp), + def amplitudes_with_params(*values): + amplitudes = torch.as_tensor(model(configs, params=values)).reshape(-1) + if amplitudes.numel() != configs.shape[0]: + raise ValueError( + "model must return one amplitude per configuration row." ) - denom = torch.where(amp_abs < floor, sign * floor, amp) + return amplitudes + + amplitudes = amplitudes_with_params(*params) + denominator = _log_derivative_denominator(amplitudes, amplitude_floor) + complex_output = torch.is_complex(amplitudes) + complex_parameters = tuple(torch.is_complex(param) for param in params) + need_real_and_imag = complex_output and ( + real_imag_mode or not all(complex_parameters) + ) - grads = torch.autograd.grad( - amp, + if need_real_and_imag: + jacobian = torch.autograd.functional.jacobian( + lambda *values: torch.view_as_real(amplitudes_with_params(*values)), params, - retain_graph=create_graph, + create_graph=create_graph, + strict=False, + vectorize=True, + ) + jacobian_real = tuple(value[:, 0, ...] for value in jacobian) + jacobian_imag = tuple(value[:, 1, ...] for value in jacobian) + else: + jacobian_real = torch.autograd.functional.jacobian( + lambda *values: ( + amplitudes_with_params(*values).real + if complex_output + else amplitudes_with_params(*values) + ), + params, + create_graph=create_graph, + strict=False, + vectorize=True, + ) + jacobian_real = tuple(jacobian_real) + jacobian_imag = (None,) * len(params) + + pieces = [] + for param, real_grad, imag_grad in zip( + params, + jacobian_real, + jacobian_imag, + strict=True, + ): + real_grad = real_grad.reshape(configs.shape[0], -1) + if complex_output: + if torch.is_complex(param): + if real_imag_mode: + imag_grad = imag_grad.reshape(configs.shape[0], -1) + derivative_real = real_grad.real + 1j * imag_grad.real + derivative_imag = real_grad.imag + 1j * imag_grad.imag + pieces.append( + torch.stack((derivative_real, derivative_imag), dim=-1) + .reshape(configs.shape[0], -1) + ) + else: + # For a holomorphic f(z), autograd's real-output + # derivative is conjugated before forming df / f. + pieces.append(real_grad.conj()) + else: + imag_grad = imag_grad.reshape(configs.shape[0], -1) + pieces.append(real_grad + 1j * imag_grad) + elif real_imag_mode and torch.is_complex(param): + pieces.append( + torch.stack((real_grad.real, real_grad.imag), dim=-1) + .reshape(configs.shape[0], -1) + ) + else: + pieces.append(real_grad) + + result = torch.cat(pieces, dim=1) / denominator.reshape(-1, 1) + return result if create_graph else result.detach() + + +def _torch_log_derivative_matrix_loop( + model, + configs, + *, + amplitude_floor=None, + create_graph=False, + complex_parameter_mode="holomorphic", +): + """Return per-sample log-amplitude derivatives for a torch model. + + The returned matrix has shape ``(n_samples, n_params)`` and entries + ``d psi(config) / d theta / psi(config)``. Real parameters use the + ordinary real derivative, while complex parameters use the explicitly + selected ``complex_parameter_mode``. The default ``"holomorphic"`` mode + is appropriate for packed PEPS amplitudes, which are holomorphic in their + complex tensor entries, and returns one complex derivative per complex + parameter. In ``"real-imag"`` mode, each complex parameter contributes + two interleaved columns, ``d log(psi) / d Re(theta)`` and + ``d log(psi) / d Im(theta)``. + + Complex parameters are not treated as real parameters implicitly. The + holomorphic convention is used by :func:`TorchVMCDriver.step` and by + :func:`apply_torch_sr_update` for complex PEPS tensors. + """ + torch = _require_torch() + configs = _as_long_matrix(configs) + params = _torch_model_parameters(model) + parameter_mode = str(complex_parameter_mode).replace("_", "-").lower() + if parameter_mode not in { + "holomorphic", + "holomorphic-complex", + "real-imag", + "real-imaginary", + }: + raise ValueError( + "complex_parameter_mode must be 'holomorphic' or 'real-imag'." + ) + real_imag_mode = parameter_mode in {"real-imag", "real-imaginary"} + rows = [] + + for config in configs: + amp = model(config.reshape(1, -1)) + amp = torch.as_tensor(amp).reshape(-1) + if amp.numel() != 1: + raise ValueError("model(config) must return one amplitude per row.") + amp = amp[0] + if not amp.requires_grad: + raise RuntimeError("model amplitude does not require gradients.") + + amp_abs = amp.detach().abs() + if amplitude_floor is None: + if amp_abs.item() == 0: + raise ZeroDivisionError( + "Encountered a zero amplitude while forming log derivatives." + ) + denom = amp + else: + floor = torch.as_tensor( + amplitude_floor, + dtype=amp.real.dtype if torch.is_complex(amp) else amp.dtype, + device=amp.device, + ) + if torch.is_complex(amp): + phase = torch.where( + amp_abs > 0, + amp / amp_abs.to(dtype=amp.dtype), + torch.ones_like(amp), + ) + denom = torch.where( + amp_abs < floor, + phase * floor.to(dtype=amp.dtype), + amp, + ) + else: + sign = torch.where( + amp.detach() >= 0, + torch.ones_like(amp), + -torch.ones_like(amp), + ) + denom = torch.where(amp_abs < floor, sign * floor, amp) + + grad_real = torch.autograd.grad( + amp.real if torch.is_complex(amp) else amp, + params, + retain_graph=True, create_graph=create_graph, allow_unused=True, ) - row = _flatten_torch_tensors(grads, params) / denom + if torch.is_complex(amp) and amp.imag.requires_grad: + grad_imag = torch.autograd.grad( + amp.imag, + params, + retain_graph=True, + create_graph=create_graph, + allow_unused=True, + ) + else: + grad_imag = (None,) * len(params) + + derivative_pieces = [] + for param, real_grad, imag_grad in zip( + params, + grad_real, + grad_imag, + strict=True, + ): + if real_grad is None: + real_grad = torch.zeros_like(param) + if torch.is_complex(amp): + if imag_grad is None: + imag_grad = torch.zeros_like(param) + if torch.is_complex(param) and not real_imag_mode: + # For a holomorphic f(z), torch's gradient of Re[f] with + # respect to z is conjugate(df / dz). + derivative = real_grad.conj() + elif torch.is_complex(param): + # PyTorch encodes the coordinate gradients of a real + # scalar with real and imaginary parts of its complex + # gradient. Recover both output components explicitly. + derivative = ( + real_grad.real + 1j * imag_grad.real, + real_grad.imag + 1j * imag_grad.imag, + ) + else: + # Real parameters need both output components to recover + # the complex derivative along the real parameter axis. + derivative = real_grad + 1j * imag_grad + else: + if real_imag_mode and torch.is_complex(param): + derivative = (real_grad.real, real_grad.imag) + else: + derivative = real_grad + if isinstance(derivative, tuple): + derivative_pieces.append( + torch.stack(derivative, dim=-1).reshape(-1) + ) + else: + derivative_pieces.append(derivative.reshape(-1)) + + row = torch.cat(derivative_pieces) / denom if not create_graph: row = row.detach() rows.append(row) @@ -900,6 +2880,56 @@ def torch_log_derivative_matrix( return torch.stack(rows, dim=0) +def torch_log_derivative_matrix( + model, + configs, + *, + amplitude_floor=None, + create_graph=False, + complex_parameter_mode="holomorphic", + derivative_backend="auto", +): + """Return per-sample log-amplitude derivatives for a torch model. + + The returned matrix has shape ``(n_samples, n_params)`` and entries + ``d psi(config) / d theta / psi(config)``. ``derivative_backend="auto"`` + uses one batched Jacobian for functional PEPS amplitude models and falls + back to the original per-sample autograd loop for generic models or + unsupported contraction transformations. Use ``"loop"`` to force the + compatibility path or ``"batched"`` to require the fast PEPS path. + + Complex parameters use the explicitly selected + ``complex_parameter_mode``. In ``"holomorphic"`` mode, one complex + derivative is returned per complex parameter. In ``"real-imag"`` mode, + each complex parameter contributes interleaved real and imaginary + coordinate derivatives. + """ + backend = str(derivative_backend).replace("_", "-").lower() + if backend not in {"auto", "batched", "loop", "scalar"}: + raise ValueError( + "derivative_backend must be 'auto', 'batched', 'loop', or 'scalar'." + ) + if backend in {"auto", "batched"}: + try: + return _batched_model_log_derivatives( + model, + configs, + amplitude_floor=amplitude_floor, + create_graph=create_graph, + complex_parameter_mode=complex_parameter_mode, + ) + except (AttributeError, NotImplementedError, RuntimeError, TypeError, ValueError): + if backend == "batched": + raise + return _torch_log_derivative_matrix_loop( + model, + configs, + amplitude_floor=amplitude_floor, + create_graph=create_graph, + complex_parameter_mode=complex_parameter_mode, + ) + + def _promote_sr_tensors(log_derivatives, local_energies): torch = _require_torch() log_derivatives = torch.as_tensor(log_derivatives) @@ -924,21 +2954,92 @@ def _promote_sr_tensors(log_derivatives, local_energies): return log_derivatives.to(dtype), local_energies.to(dtype) -def _torch_solve_linear(matrix, rhs): +def _resolve_sr_diag_shift(diag_shift, *, step): + """Resolve a constant or step-indexed SR diagonal shift.""" + value = diag_shift(step) if callable(diag_shift) else diag_shift + try: + value = float(value) + except (TypeError, ValueError) as exc: + raise ValueError( + "diag_shift must be a non-negative number or a callable returning one." + ) from exc + if not math.isfinite(value) or value < 0.0: + raise ValueError( + "diag_shift must be a non-negative finite number or schedule value." + ) + return value + + +def _torch_solve_linear(matrix, rhs, *, pinv_rtol=None): + """Solve a Hermitian SR system, falling back to a stable pseudoinverse.""" torch = _require_torch() try: - return torch.linalg.solve(matrix, rhs), "solve" + factor, info = torch.linalg.cholesky_ex(matrix, check_errors=False) + if bool(torch.all(info == 0)): + solution = torch.cholesky_solve( + rhs.reshape(-1, 1), + factor, + ).reshape_as(rhs) + if bool(torch.isfinite(solution).all()): + return solution, "cholesky" except RuntimeError: - return torch.linalg.lstsq(matrix, rhs).solution, "lstsq" + pass + + pinv_kwargs = {"hermitian": True} + if pinv_rtol is not None: + pinv_kwargs["rtol"] = pinv_rtol + solution = torch.linalg.pinv(matrix, **pinv_kwargs) @ rhs + if not bool(torch.isfinite(solution).all()): + raise RuntimeError( + "The SR pseudoinverse fallback produced non-finite values. " + "Increase diag_shift or inspect the local-energy samples." + ) + return solution, "pinv" + + +def _spring_complement(metric_source, previous_direction, *, pinv_rtol=None): + """Return the previous SR update outside the current sampled tangent span.""" + torch = _require_torch() + previous_direction = torch.as_tensor( + previous_direction, + dtype=metric_source.dtype, + device=metric_source.device, + ).reshape(-1) + if previous_direction.shape[0] != metric_source.shape[1]: + raise ValueError( + "previous_direction must have one entry per SR parameter." + ) + tangent = ( + metric_source.transpose(0, 1) + if not torch.is_complex(metric_source) + else metric_source.conj().transpose(0, 1) + ) + if tangent.shape[1] == 0: + return previous_direction + solve_kwargs = {} + if pinv_rtol is not None: + solve_kwargs["rcond"] = pinv_rtol + coefficients = torch.linalg.lstsq( + tangent, + previous_direction, + **solve_kwargs, + ).solution + return previous_direction - tangent @ coefficients def solve_torch_sr( log_derivatives, local_energies, *, + sample_weights=None, diag_shift=1.0e-4, method="auto", center=True, + parameter_mode="holomorphic", + step=0, + pinv_rtol=None, + momentum=None, + previous_direction=None, ): """Solve direct SR or sample-space minSR for a torch VMC batch. @@ -946,11 +3047,45 @@ def solve_torch_sr( ``method="minsr"`` solves the equivalent sample-space system, which is preferable when the number of PEPS parameters is much larger than the number of Monte Carlo samples. ``method="auto"`` picks minSR when - ``n_samples < n_params``. + ``n_samples < n_params``. Complex derivatives use the Hermitian covariance + ``centered.conj().T @ centered`` and return a complex SR direction under + the holomorphic parameter convention. ``parameter_mode="real-imag"`` + instead solves a real SR system for the explicit real and imaginary + parameter coordinates returned by :func:`torch_log_derivative_matrix`. + ``diag_shift`` can be a callable of the non-negative integer ``step``. + The Hermitian system uses a Cholesky solve when possible and otherwise a + pseudoinverse fallback. Passing ``momentum`` together with a previous + direction applies a SPRING-style complement: only the part of the prior + update outside the current sampled tangent span is retained. Pass + normalized or relative ``sample_weights`` to form the corresponding + weighted energy and tangent-space covariances. """ torch = _require_torch() - if diag_shift < 0: - raise ValueError("diag_shift must be non-negative.") + step = _check_nonnegative_int("step", step) + diag_shift = _resolve_sr_diag_shift(diag_shift, step=step) + if pinv_rtol is not None: + try: + pinv_rtol = float(pinv_rtol) + except (TypeError, ValueError) as exc: + raise ValueError("pinv_rtol must be a positive finite number or None.") from exc + if not math.isfinite(pinv_rtol) or pinv_rtol <= 0.0: + raise ValueError("pinv_rtol must be a positive finite number or None.") + if momentum is not None: + try: + momentum = float(momentum) + except (TypeError, ValueError) as exc: + raise ValueError("momentum must be in [0, 1).") from exc + if not math.isfinite(momentum) or not 0.0 <= momentum < 1.0: + raise ValueError("momentum must be in [0, 1).") + parameter_mode = str(parameter_mode).replace("_", "-").lower() + if parameter_mode not in { + "holomorphic", + "holomorphic-complex", + "real-imag", + "real-imaginary", + }: + raise ValueError("parameter_mode must be 'holomorphic' or 'real-imag'.") + real_imag_mode = parameter_mode in {"real-imag", "real-imaginary"} log_derivatives, local_energies = _promote_sr_tensors( log_derivatives, local_energies, @@ -967,44 +3102,135 @@ def solve_torch_sr( if method_key == "sr": method_key = "direct" - energy_mean = local_energies.mean() + if sample_weights is None: + weights = torch.full( + (n_samples,), + 1.0 / n_samples, + dtype=local_energies.real.dtype, + device=local_energies.device, + ) + else: + weights = torch.as_tensor(sample_weights, device=local_energies.device) + if weights.ndim != 1 or weights.shape[0] != n_samples: + raise ValueError("sample_weights must have one entry per sample.") + if torch.is_complex(weights): + raise ValueError("sample_weights must be real, finite, and non-negative.") + if not torch.is_floating_point(weights): + weights = weights.to(local_energies.real.dtype) + if not bool(torch.isfinite(weights).all()) or bool(torch.any(weights < 0)): + raise ValueError("sample_weights must be real, finite, and non-negative.") + total_weight = weights.sum() + if not bool(torch.isfinite(total_weight)) or bool(total_weight <= 0): + raise ValueError("sample_weights must have a positive finite sum.") + weights = weights / total_weight + + energy_mean = (weights.to(local_energies.dtype) * local_energies).sum() energy_residual = local_energies - energy_mean if center else local_energies centered = ( - log_derivatives - log_derivatives.mean(dim=0, keepdim=True) + log_derivatives + - (weights.to(log_derivatives.dtype).reshape(-1, 1) * log_derivatives) + .sum(dim=0, keepdim=True) if center else log_derivatives ) - force = centered.conj().transpose(0, 1) @ energy_residual / n_samples + sqrt_weights = weights.sqrt().reshape(-1, 1) + if real_imag_mode: + # For real coordinates, write C = A + i B and solve with the real + # design matrix [A; B]. This gives Re(C^H C) and Re(C^H E) while + # retaining an exact direct/minSR equivalence. + centered_imag = ( + centered.imag if torch.is_complex(centered) else torch.zeros_like(centered) + ) + energy_imag = ( + energy_residual.imag + if torch.is_complex(energy_residual) + else torch.zeros_like(energy_residual) + ) + design = torch.cat( + (sqrt_weights * centered.real, sqrt_weights * centered_imag), + dim=0, + ) + solve_energy = torch.cat( + (sqrt_weights.reshape(-1) * energy_residual.real, + sqrt_weights.reshape(-1) * energy_imag) + ) + force = design.transpose(0, 1) @ solve_energy + metric_source = design + sr_dtype = design.dtype + else: + metric_source = sqrt_weights.to(log_derivatives.dtype) * centered + solve_energy = ( + sqrt_weights.reshape(-1).to(local_energies.dtype) * energy_residual + ) + force = metric_source.conj().transpose(0, 1) @ solve_energy + sr_dtype = log_derivatives.dtype shift = torch.as_tensor( diag_shift, - dtype=log_derivatives.dtype, + dtype=sr_dtype, device=log_derivatives.device, ) if method_key == "direct": eye = torch.eye( n_params, - dtype=log_derivatives.dtype, + dtype=sr_dtype, device=log_derivatives.device, ) - sr_matrix = centered.conj().transpose(0, 1) @ centered / n_samples - direction, solver = _torch_solve_linear(sr_matrix + shift * eye, force) + if real_imag_mode: + sr_matrix = metric_source.transpose(0, 1) @ metric_source + else: + sr_matrix = metric_source.conj().transpose(0, 1) @ metric_source + system = sr_matrix + shift * eye + direction, solver = _torch_solve_linear( + system, + force, + pinv_rtol=pinv_rtol, + ) + solve_vector = direction + solve_rhs = force matrix_shape = tuple(sr_matrix.shape) else: + n_system = metric_source.shape[0] eye = torch.eye( - n_samples, - dtype=log_derivatives.dtype, + n_system, + dtype=sr_dtype, device=log_derivatives.device, ) - gram = centered @ centered.conj().transpose(0, 1) + if real_imag_mode: + gram = metric_source @ metric_source.transpose(0, 1) + else: + gram = metric_source @ metric_source.conj().transpose(0, 1) + system = gram + shift * eye alpha, solver = _torch_solve_linear( - gram + n_samples * shift * eye, - energy_residual, + system, + solve_energy, + pinv_rtol=pinv_rtol, ) - direction = centered.conj().transpose(0, 1) @ alpha + if real_imag_mode: + direction = metric_source.transpose(0, 1) @ alpha + else: + direction = metric_source.conj().transpose(0, 1) @ alpha + solve_vector = alpha + solve_rhs = solve_energy matrix_shape = tuple(gram.shape) - energy_variance = energy_residual.abs().square().mean() + spring_complement_norm = None + if momentum is not None and momentum > 0.0 and previous_direction is not None: + spring_complement = _spring_complement( + metric_source, + previous_direction, + pinv_rtol=pinv_rtol, + ) + direction = direction + momentum * spring_complement + spring_complement_norm = float(spring_complement.norm().detach().cpu()) + + energy_variance = (weights * energy_residual.abs().square()).sum().real + residual = system @ solve_vector - solve_rhs + residual_norm = residual.norm() + rhs_norm = solve_rhs.norm() + relative_residual = residual_norm / rhs_norm.clamp_min( + torch.finfo(rhs_norm.real.dtype).tiny + ) return TorchSRResult( direction=direction, energy_mean=energy_mean, @@ -1013,15 +3239,48 @@ def solve_torch_sr( centered_log_derivatives=centered, method=method_key, diag_shift=float(diag_shift), - info={"solver": solver, "matrix_shape": matrix_shape}, + info={ + "solver": solver, + "matrix_shape": matrix_shape, + "residual_norm": float(residual_norm.detach().cpu()), + "relative_residual": float(relative_residual.detach().cpu()), + "step": step, + "pinv_rtol": pinv_rtol, + "momentum": momentum, + "spring_complement_norm": spring_complement_norm, + "effective_sample_size": float((1.0 / weights.square().sum()).detach().cpu()), + }, ) -def apply_torch_sr_update(model, direction, *, learning_rate=1.0): - """Apply ``theta <- theta - learning_rate * direction`` in place.""" +def apply_torch_sr_update( + model, + direction, + *, + learning_rate=1.0, + parameter_mode="holomorphic", +): + """Apply an SR direction in place. + + ``parameter_mode="holomorphic"`` applies one complex direction per + complex parameter. ``parameter_mode="real-imag"`` consumes interleaved + real and imaginary coordinate updates for each complex parameter. + """ torch = _require_torch() params = _torch_model_parameters(model) - n_params = sum(param.numel() for param in params) + parameter_mode = str(parameter_mode).replace("_", "-").lower() + if parameter_mode not in { + "holomorphic", + "holomorphic-complex", + "real-imag", + "real-imaginary", + }: + raise ValueError("parameter_mode must be 'holomorphic' or 'real-imag'.") + real_imag_mode = parameter_mode in {"real-imag", "real-imaginary"} + n_params = sum( + (2 if real_imag_mode and torch.is_complex(param) else 1) * param.numel() + for param in params + ) direction = torch.as_tensor(direction) if direction.numel() != n_params: raise ValueError( @@ -1032,7 +3291,17 @@ def apply_torch_sr_update(model, direction, *, learning_rate=1.0): with torch.no_grad(): for param in params: size = param.numel() - update = direction[offset:offset + size].reshape_as(param) + if real_imag_mode and torch.is_complex(param): + coordinate_updates = direction[ + offset:offset + 2 * size + ].reshape(-1, 2) + real_update = coordinate_updates[:, 0].reshape_as(param.real) + imag_update = coordinate_updates[:, 1].reshape_as(param.real) + update = real_update + 1j * imag_update + offset += 2 * size + else: + update = direction[offset:offset + size].reshape_as(param) + offset += size if torch.is_complex(update) and not torch.is_complex(param): if update.imag.abs().max().item() > 1.0e-12: raise ValueError( @@ -1041,7 +3310,6 @@ def apply_torch_sr_update(model, direction, *, learning_rate=1.0): update = update.real update = update.to(dtype=param.dtype, device=param.device) param.sub_(learning_rate * update) - offset += size return model @@ -1049,10 +3317,10 @@ def apply_torch_sr_update(model, direction, *, learning_rate=1.0): class FermionSiteEncoding: """Four-state spinful-fermion on-site encoding. - Pepsy's Symmray physical-index convention is ``0=empty, 1=double, - 2=up, 3=down``. The reference ``vmc_torch`` code often uses - ``0=empty, 1=down, 2=up, 3=double``. Use the class constructors to make - that choice explicit. + Native torch VMC and Pepsy's four-sector fermionic PEPS use + ``0=empty, 1=down, 2=up, 3=double``. The ``symmray`` constructor is kept + for callers that explicitly use the alternate legacy labels. Use the + class constructors to make the choice explicit at interop boundaries. """ empty: int = 0 @@ -1075,6 +3343,48 @@ def vmc_torch(cls): """Return the convention used by ``sjdu10/vmc_torch``.""" return cls(empty=0, double=3, up=2, down=1) + @classmethod + def from_fermion(cls, fermion, *, physical_charges=None): + """Return the PEPS physical-index encoding for a native ``Fermion``. + + When the PEPS exposes four resolved ``U1U1`` physical charges, their + ordered charge map is authoritative. For charge-collapsed spinful + ``U1`` data, the conventional four-state PEPS order is used. This + keeps the physical-index contract distinct from the dense local basis + used internally while constructing native Fermion operators. + """ + if not bool(getattr(fermion, "spinful", False)): + raise ValueError("FermionSiteEncoding.from_fermion requires spinful=True.") + if physical_charges: + try: + return cls.from_physical_charges(physical_charges) + except ValueError: + pass + return cls.vmc_torch() + + @classmethod + def from_physical_charges(cls, physical_charges): + """Return an encoding from four resolved ``(n_up, n_down)`` charges.""" + lookup = {} + for code, charge in enumerate(tuple(physical_charges)): + if not isinstance(charge, tuple) or len(charge) != 2: + raise ValueError("Physical charges must be two-component tuples.") + charge = tuple(int(value) for value in charge) + if charge in lookup: + raise ValueError("PEPS physical charges must be unique.") + lookup[charge] = code + required = {(0, 0), (0, 1), (1, 0), (1, 1)} + if set(lookup) != required: + raise ValueError( + "Physical charges must contain exactly the four spinful states." + ) + return cls( + empty=lookup[(0, 0)], + down=lookup[(0, 1)], + up=lookup[(1, 0)], + double=lookup[(1, 1)], + ) + @property def max_code(self): return max(self.empty, self.double, self.up, self.down) @@ -1121,6 +3431,53 @@ def encode(self, n_up, n_down): return code +@dataclass(frozen=True) +class SpinlessSiteEncoding: + """Binary local encoding for spinless fermions.""" + + empty: int = 0 + occupied: int = 1 + + def __post_init__(self): + if self.empty == self.occupied or min(self.empty, self.occupied) < 0: + raise ValueError("Spinless site codes must be distinct and non-negative.") + + @classmethod + def from_physical_charges(cls, physical_charges): + """Infer the binary code order from an ordered charge map.""" + charges = tuple(physical_charges) + if set(charges) != {0, 1}: + raise ValueError("Spinless physical charges must be exactly {0, 1}.") + return cls(empty=charges.index(0), occupied=charges.index(1)) + + @property + def max_code(self): + return max(self.empty, self.occupied) + + def validate(self, configs): + torch = _require_torch() + valid = (configs == self.empty) | (configs == self.occupied) + if not torch.all(valid): + bad = torch.unique(configs[~valid]).detach().cpu().tolist() + raise ValueError(f"Unknown spinless fermion site code(s): {bad!r}.") + + def decode(self, configs): + """Return binary occupation tensors.""" + torch = _require_torch() + configs = torch.as_tensor(configs, dtype=torch.long) + self.validate(configs) + return (configs == self.occupied).long() + + def encode(self, occupied): + torch = _require_torch() + occupied = torch.as_tensor(occupied, dtype=torch.long) + return torch.where( + occupied == 1, + torch.as_tensor(self.occupied, device=occupied.device), + torch.as_tensor(self.empty, device=occupied.device), + ) + + @dataclass(frozen=True) class TorchSquareLattice: """Square-lattice nearest-neighbor graph with grouped row/column edges.""" @@ -1170,394 +3527,1114 @@ def edges(self): return tuple(edges) -@dataclass(frozen=True) -class TorchConnections: - """Batched Hamiltonian connections. +def _peps_physical_axis(tn, site): + """Return the physical tensor axis and dimension for ``site``.""" + tensor = tn[site] + try: + axis = tuple(tensor.inds).index(tn.site_ind(site)) + except (AttributeError, ValueError) as exc: + raise ValueError( + f"Could not locate the physical index for PEPS site {site!r}." + ) from exc + shape = getattr(tensor, "shape", None) + if shape is None: + shape = getattr(getattr(tensor, "data", None), "shape", None) + if shape is None: + raise ValueError(f"Could not determine the physical dimension at site {site!r}.") + return axis, int(shape[axis]) + + +def _peps_physical_charges(tn, site): + """Return ordered Symmray physical charges, when available.""" + tensor = tn[site] + data = getattr(tensor, "data", None) + if not _is_symmray_data(data): + return () + try: + axis, _ = _peps_physical_axis(tn, site) + index = data.indices[axis] + chargemap = getattr(index, "chargemap", None) + except (AttributeError, IndexError, TypeError, ValueError): + return () + if chargemap is None: + return () + return tuple(chargemap.keys()) + + +def _peps_symmetry(tn, site_order): + """Return the named Symmray symmetry carried by the PEPS, if present.""" + for site in site_order: + symmetry = getattr(getattr(tn[site], "data", None), "symmetry", None) + if symmetry is not None: + return str(symmetry).upper() + return None + + +def _resolve_peps_pbc(tn, pbc): + """Resolve PBC axes from an explicit value or PEPS cyclic metadata.""" + if pbc is None: + axes = [] + for name in ("is_cyclic_x", "is_cyclic_y"): + checker = getattr(tn, name, None) + if checker is None: + axes.append(False) + continue + try: + value = checker() if callable(checker) else checker + except (AttributeError, TypeError, ValueError): + value = False + axes.append(bool(value)) + return tuple(axes) + if isinstance(pbc, bool): + return (pbc, pbc) + try: + axes = tuple(pbc) + except TypeError as exc: + raise ValueError("pbc must be a bool, None, or a two-entry tuple.") from exc + if len(axes) != 2: + raise ValueError("pbc must be a bool, None, or a two-entry tuple.") + return tuple(bool(axis) for axis in axes) + + +def _peps_lattice_edges(site_order, Lx, Ly, *, pbc=False): + """Infer coordinate-labelled nearest-neighbor edges from PEPS metadata.""" + site_order = tuple(site_order) + if not all( + isinstance(site, tuple) + and len(site) == 2 + and all(isinstance(value, Integral) for value in site) + for site in site_order + ): + raise ValueError( + "PEPS sites must be coordinate labels to infer lattice edges; " + "pass edges explicitly for non-coordinate site labels." + ) + site_order = tuple((int(site[0]), int(site[1])) for site in site_order) + by_coord = {site: site for site in site_order} + expected = {(x, y) for x in range(Lx) for y in range(Ly)} + if set(by_coord) != expected: + raise ValueError( + "PEPS coordinate sites do not form the inferred rectangular grid." + ) - ``configs[k]`` is connected to source sample ``batch_ids[k]`` with - coefficient ``coeffs[k]``. - """ + if isinstance(pbc, bool): + pbc_x = pbc_y = pbc + else: + try: + pbc_x, pbc_y = pbc + except (TypeError, ValueError) as exc: + raise ValueError("pbc must be a bool or a two-entry tuple.") from exc + + edges = [] + for x in range(Lx): + for y in range(Ly - 1): + edges.append(((x, y), (x, y + 1))) + if pbc_y and Ly > 2: + edges.append(((x, Ly - 1), (x, 0))) + for y in range(Ly): + for x in range(Lx - 1): + edges.append(((x, y), (x + 1, y))) + if pbc_x and Lx > 2: + edges.append(((Lx - 1, y), (0, y))) + return tuple(edges) + + +def _term_support_edges(terms, site_order): + """Extract unique two-site supports from explicit local terms.""" + if terms is None: + return () + + from .api import OperatorSum + common_terms = terms if isinstance(terms, OperatorSum) else None + + site_order = tuple(site_order) + positions = {site: index for index, site in enumerate(site_order)} + support_edges = [] + seen = set() + + def map_site(site): + if site in positions: + return site + if ( + isinstance(site, Integral) + and not isinstance(site, bool) + and 0 <= int(site) < len(site_order) + ): + return site_order[int(site)] + raise ValueError( + f"Hamiltonian term site {site!r} is not present in the PEPS." + ) - configs: Any - coeffs: Any - batch_ids: Any + if common_terms is not None: + entries = tuple( + (term.support, term) + for term in common_terms + if len(term.support) == 2 + ) + else: + entries = _term_items(terms) + for where, operator in entries: + if common_terms is not None: + shape = (2, 2, 2, 2) + else: + shape = getattr(operator, "shape", None) + if shape is None: + shape = getattr(_term_dense_array(operator), "shape", ()) + if len(shape) != 4: + continue + try: + left, right = tuple(where) + except (TypeError, ValueError) as exc: + raise ValueError( + "A two-site Hamiltonian term location must contain two sites." + ) from exc + left = map_site(left) + right = map_site(right) + left_position = positions[left] + right_position = positions[right] + if left_position == right_position: + continue + key = frozenset((left_position, right_position)) + if key in seen: + continue + seen.add(key) + if left_position > right_position: + left, right = right, left + support_edges.append((left, right)) + return tuple(support_edges) + + +def _coerce_labelled_edges(edges, site_order): + """Normalize explicit edges to labels in ``site_order``.""" + site_order = tuple(site_order) + positions = {site: i for i, site in enumerate(site_order)} + normalized = [] + for edge in tuple(edges): + try: + left, right = tuple(edge) + except (TypeError, ValueError) as exc: + raise ValueError("Each edge must contain exactly two site labels.") from exc + if left in positions and right in positions: + normalized.append((left, right)) + continue + if ( + isinstance(left, Integral) + and not isinstance(left, bool) + and isinstance(right, Integral) + and not isinstance(right, bool) + and 0 <= int(left) < len(site_order) + and 0 <= int(right) < len(site_order) + ): + normalized.append((site_order[int(left)], site_order[int(right)])) + continue + raise ValueError( + f"Edge {(left, right)!r} contains a site not present in the PEPS." + ) + return tuple(normalized) -@dataclass(frozen=True) -class TorchMetropolisResult: - """Result of one Metropolis sweep.""" - configs: Any - amplitudes: Any - n_proposed: int - n_accepted: int +def _sum_site_charges(tn, site_order): + """Infer a fixed global charge from Symmray tensor charge metadata.""" + charges = [] + for site in site_order: + charge = getattr(getattr(tn[site], "data", None), "charge", None) + if charge is None: + return None + if isinstance(charge, tuple): + charge = tuple(int(value) for value in charge) + else: + charge = int(charge) + charges.append(charge) + if not charges: + return None + first = charges[0] + if isinstance(first, tuple): + if not all( + isinstance(charge, tuple) and len(charge) == len(first) + for charge in charges + ): + return None + return tuple( + sum(charge[axis] for charge in charges) + for axis in range(len(first)) + ) + if any(isinstance(charge, tuple) for charge in charges): + return None + return sum(charges) + + +def _coerce_fermion_sector(sector, symmetry): + """Normalize a requested physical sector for the supported spinful modes.""" + symmetry = str(symmetry).upper() + if symmetry == "Z2": + if isinstance(sector, bool) or not isinstance(sector, Integral): + raise ValueError("A spinful Z2 sector must be parity 0 or 1.") + return int(sector) % 2 + if symmetry == "U1": + if isinstance(sector, bool) or not isinstance(sector, Integral): + raise ValueError( + "A spinful U1 sector must be an integer total particle number." + ) + return int(sector) + if symmetry == "U1U1": + try: + sector = tuple(sector) + except TypeError as exc: + raise ValueError( + "A spinful U1U1 sector must be (N_up, N_down)." + ) from exc + if len(sector) != 2 or any( + isinstance(value, bool) or not isinstance(value, Integral) + for value in sector + ): + raise ValueError("A spinful U1U1 sector must be (N_up, N_down).") + return tuple(int(value) for value in sector) + if symmetry == "Z2Z2": + try: + sector = tuple(sector) + except TypeError as exc: + raise ValueError( + "A spinful Z2Z2 sector must be (parity_up, parity_down)." + ) from exc + if len(sector) != 2 or any( + isinstance(value, bool) or not isinstance(value, Integral) + for value in sector + ): + raise ValueError("A spinful Z2Z2 sector must be (parity_up, parity_down).") + return tuple(int(value) % 2 for value in sector) + raise NotImplementedError( + f"Automatic Torch Fermion VMC does not support {symmetry!r}." + ) - @property - def acceptance_rate(self): - if self.n_proposed == 0: - return 0.0 - return self.n_accepted / self.n_proposed +def _validate_fermion_sector(sector, symmetry, n_sites, *, spinful=True): + sector = _coerce_fermion_sector(sector, symmetry) + if symmetry in {"Z2", "Z2Z2"}: + return sector + if symmetry == "U1": + max_particles = 2 * n_sites if spinful else n_sites + if not 0 <= sector <= max_particles: + raise ValueError( + f"U1 total particle sector must be between 0 and {max_particles}." + ) + elif any(value < 0 or value > n_sites for value in sector): + raise ValueError( + f"U1U1 sector entries must each be between 0 and {n_sites}." + ) + return sector -def _iter_edges(graph): - if hasattr(graph, "edges"): - edges = graph.edges - if callable(edges): - edges = edges() - else: - edges = graph - return tuple((int(i), int(j)) for i, j in edges) +@dataclass(frozen=True) +class TorchFermionVMCMetadata: + """Validated PEPS/Fermion metadata used by :class:`TorchFermionVMC`.""" -def _empty_connections(configs): - torch = _require_torch() - return TorchConnections( - configs=configs.new_empty((0, configs.shape[1])), - coeffs=torch.empty(0, dtype=torch.float64, device=configs.device), - batch_ids=torch.empty(0, dtype=torch.long, device=configs.device), - ) + site_order: tuple[Any, ...] + edges: tuple[tuple[Any, Any], ...] + graph_edges: tuple[tuple[int, int], ...] + Lx: int + Ly: int + physical_dim: int + symmetry: str + spinful: bool + encoding: Any + sector: int | tuple[int, int] | None + physical_charges: tuple[Any, ...] = () + pbc: tuple[bool, bool] = (False, False) + @property + def n_sites(self): + return len(self.site_order) -def count_spinful_particles(configs, *, encoding=None): - """Return per-sample ``(n_up, n_down)`` counts.""" - encoding = FermionSiteEncoding.symmray() if encoding is None else encoding - configs = _as_long_matrix(configs) - n_up, n_down = encoding.decode(configs) - return n_up.sum(dim=-1), n_down.sum(dim=-1) - - -def propose_spin_exchange(i, j, configs): - """Propose spin exchange on one edge for binary spin configs.""" - configs = _as_long_matrix(configs) - proposed = configs.clone() - si = configs[:, i] - sj = configs[:, j] - changed = si != sj - proposed[changed, i] = sj[changed] - proposed[changed, j] = si[changed] - return proposed, changed + @property + def graph(self): + """Return the integer graph consumed by the Torch sampler.""" + return self.graph_edges -def propose_spinful_exchange_or_hopping( - i, - j, - configs, +def _infer_torch_fermion_metadata( + peps, + fermion, *, - hopping_rate=0.25, - encoding=None, - generator=None, + sector=None, + edges=None, + pbc=None, + site_order=None, + terms=None, ): - """Propose spinful Hubbard exchange/hopping moves on one edge. + """Infer and validate all static metadata for native spinful PEPS VMC.""" + tn = getattr(peps, "tn", peps) + if not hasattr(tn, "sites"): + raise TypeError("peps must be a quimb PEPS-like object with sites.") + site_order = tuple(tn.sites if site_order is None else site_order) + if not site_order: + raise ValueError("The PEPS must contain at least one physical site.") + if len(set(site_order)) != len(site_order): + raise ValueError("PEPS site_order must contain unique site labels.") + missing = [site for site in site_order if site not in tn.sites] + if missing: + raise ValueError(f"site_order contains site(s) not in PEPS: {missing!r}") + + Lx = getattr(tn, "Lx", None) + Ly = getattr(tn, "Ly", None) + if Lx is None: + Lx = getattr(tn, "_Lx", None) + if Ly is None: + Ly = getattr(tn, "_Ly", None) + if Lx is None or Ly is None: + if all( + isinstance(site, tuple) + and len(site) == 2 + and all(isinstance(value, Integral) for value in site) + for site in site_order + ): + Lx = max(int(site[0]) for site in site_order) + 1 + Ly = max(int(site[1]) for site in site_order) + 1 + else: + raise ValueError( + "Could not infer PEPS Lx/Ly; use coordinate PEPS sites or " + "pass explicit site_order and edges." + ) + Lx = _check_positive_int("Lx", Lx) + Ly = _check_positive_int("Ly", Ly) + if len(site_order) != Lx * Ly: + raise ValueError( + f"PEPS has {len(site_order)} sites but inferred geometry is {Lx}x{Ly}." + ) - The proposal preserves ``N_up`` and ``N_down``. With probability - ``1 - hopping_rate`` it swaps the two local site states. Otherwise it uses - local hopping-style moves over ``empty/up/down/double`` states, following - the sampling options in ``sjdu10/vmc_torch``. - """ - torch = _require_torch() - encoding = FermionSiteEncoding.symmray() if encoding is None else encoding - configs = _as_long_matrix(configs) - proposed = configs.clone() - device = configs.device - batch = configs.shape[0] + pbc_axes = _resolve_peps_pbc(tn, pbc) + if edges is None: + edges = list(_peps_lattice_edges(site_order, Lx, Ly, pbc=pbc_axes)) + else: + edges = list(_coerce_labelled_edges(edges, site_order)) + + # The proposal graph must also contain non-nearest-neighbor supports from + # explicit Hamiltonian terms. This keeps exchange/hopping Metropolis moves + # able to traverse the same long-range geometry used by the estimator. + positions = {site: index for index, site in enumerate(site_order)} + edge_keys = { + frozenset((positions[left], positions[right])) + for left, right in edges + if left != right + } + for left, right in _term_support_edges(terms, site_order): + key = frozenset((positions[left], positions[right])) + if key not in edge_keys: + edges.append((left, right)) + edge_keys.add(key) + positions = {site: i for i, site in enumerate(site_order)} + graph_edges = tuple((positions[left], positions[right]) for left, right in edges) + + dimensions = [] + physical_charges = [] + for site in site_order: + _, dimension = _peps_physical_axis(tn, site) + dimensions.append(dimension) + charges = _peps_physical_charges(tn, site) + if charges: + physical_charges.append(charges) + if len(set(dimensions)) != 1: + raise ValueError(f"PEPS physical dimensions are inconsistent: {dimensions!r}.") + physical_dim = dimensions[0] + + peps_symmetry = _peps_symmetry(tn, site_order) + spinful = True if fermion is None else bool(getattr(fermion, "spinful", False)) + if fermion is None and physical_dim == 2: + spinful = False + if fermion is None: + symmetry = peps_symmetry + if symmetry is None: + raise ValueError( + "Cannot infer Fermion symmetry from this PEPS. Pass fermion=... " + "or use a Symmray PEPS with symmetry metadata." + ) + else: + symmetry = str(getattr(fermion, "symmetry", "")).upper() + if peps_symmetry is not None and peps_symmetry != symmetry: + raise ValueError( + f"PEPS symmetry {peps_symmetry!r} does not match Fermion " + f"symmetry {symmetry!r}." + ) + if symmetry not in {"U1", "U1U1", "Z2", "Z2Z2"}: + raise NotImplementedError( + "TorchFermionVMC currently supports U1, U1U1, Z2, and Z2Z2, " + f"not {symmetry!r}." + ) + expected_dim = 4 if spinful else 2 + if physical_dim != expected_dim: + raise ValueError( + f"{'Spinful' if spinful else 'Spinless'} Fermion VMC requires PEPS " + f"physical dimension {expected_dim}, got {physical_dim}." + ) + if fermion is None: + if spinful: + sectors = { + "U1": {0: 1, 1: 2, 2: 1}, + "U1U1": {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}, + "Z2": {0: 2, 1: 2}, + "Z2Z2": {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}, + }[symmetry] + else: + sectors = {0: 1, 1: 1} + else: + sectors = getattr(fermion, "physical_sectors", None) + if sectors is None or sum(int(size) for size in sectors.values()) != physical_dim: + raise ValueError("Fermion and PEPS physical dimensions/sectors are incompatible.") + + if physical_charges: + first_charges = physical_charges[0] + if any(charges != first_charges for charges in physical_charges[1:]): + raise ValueError("PEPS physical charge ordering differs between sites.") + expected_charges = tuple(sectors) + if not spinful and first_charges != expected_charges: + raise ValueError( + "PEPS and Fermion spinless physical charge orders differ; " + "refusing to apply an implicit local basis permutation." + ) + if symmetry == "U1U1" and first_charges != expected_charges: + raise ValueError( + "PEPS and Fermion U1U1 physical charge orders differ; refusing " + "to apply an implicit local basis permutation." + ) + if symmetry in {"Z2", "Z2Z2"} and first_charges != expected_charges: + raise ValueError( + "PEPS and Fermion parity physical charge sectors differ; refusing " + "to apply an implicit local basis permutation." + ) + physical_charges = first_charges + else: + physical_charges = () - ci = configs[:, i] - cj = configs[:, j] - changed = ci != cj - if not torch.any(changed): - return proposed, changed + if not spinful: + encoding = SpinlessSiteEncoding.from_physical_charges( + physical_charges or tuple(sectors) + ) + elif symmetry == "Z2": + encoding = FermionSiteEncoding.symmray() + elif symmetry == "Z2Z2" and physical_charges: + encoding = FermionSiteEncoding.from_physical_charges(physical_charges) + elif symmetry == "Z2Z2": + encoding = FermionSiteEncoding.vmc_torch() + elif fermion is None: + if symmetry == "U1U1" and physical_charges: + encoding = FermionSiteEncoding.from_physical_charges(physical_charges) + else: + encoding = FermionSiteEncoding.vmc_torch() + else: + encoding = FermionSiteEncoding.from_fermion( + fermion, + physical_charges=physical_charges, + ) + if sector is None: + sector = _sum_site_charges(tn, site_order) + if sector is not None: + sector = _validate_fermion_sector( + sector, + symmetry, + len(site_order), + spinful=spinful, + ) + return TorchFermionVMCMetadata( + site_order=site_order, + edges=tuple(edges), + graph_edges=graph_edges, + Lx=Lx, + Ly=Ly, + physical_dim=physical_dim, + symmetry=symmetry, + spinful=spinful, + encoding=encoding, + sector=sector, + physical_charges=tuple(physical_charges), + pbc=pbc_axes, + ) - n_up, n_down = encoding.decode(configs) - ni = n_up[:, i] + n_down[:, i] - nj = n_up[:, j] + n_down[:, j] - delta_n = (ni - nj).abs() - rand = torch.rand(batch, device=device, generator=generator) - is_exchange = (rand < (1.0 - hopping_rate)) & changed - is_hopping = (~is_exchange) & changed +@dataclass(frozen=True) +class TorchConnections: + """Batched Hamiltonian connections. - swap_mask = is_exchange | (is_hopping & (delta_n == 1)) - proposed[swap_mask, i] = cj[swap_mask] - proposed[swap_mask, j] = ci[swap_mask] + ``configs[k]`` is connected to source sample ``batch_ids[k]`` with + coefficient ``coeffs[k]``. + """ - mask_d0 = is_hopping & (delta_n == 0) - if torch.any(mask_d0): - bits = torch.randint( - 0, - 2, - (batch,), - device=device, - dtype=torch.long, - generator=generator, - ).bool() - proposed[mask_d0, i] = torch.where( - bits[mask_d0], - torch.as_tensor(encoding.double, device=device), - torch.as_tensor(encoding.empty, device=device), - ) - proposed[mask_d0, j] = torch.where( - bits[mask_d0], - torch.as_tensor(encoding.empty, device=device), - torch.as_tensor(encoding.double, device=device), - ) + configs: Any + coeffs: Any + batch_ids: Any - mask_d2 = is_hopping & (delta_n == 2) - if torch.any(mask_d2): - bits = torch.randint( - 0, - 2, - (batch,), - device=device, - dtype=torch.long, - generator=generator, - ).bool() - proposed[mask_d2, i] = torch.where( - bits[mask_d2], - torch.as_tensor(encoding.down, device=device), - torch.as_tensor(encoding.up, device=device), - ) - proposed[mask_d2, j] = torch.where( - bits[mask_d2], - torch.as_tensor(encoding.up, device=device), - torch.as_tensor(encoding.down, device=device), - ) - return proposed, changed +@dataclass(frozen=True) +class TorchMetropolisResult: + """Result of one Metropolis sweep.""" + configs: Any + amplitudes: Any + n_proposed: int + n_accepted: int + log_abs_amplitudes: Any = None + nonzero_amplitudes: Any = None + proposal_stats: Any = None -def _safe_metropolis_ratio(proposed_amps, current_amps): - torch = _require_torch() - numerator = proposed_amps.abs().square() - denominator = current_amps.abs().square() - zero = torch.zeros_like(numerator) - inf = torch.full_like(numerator, float("inf")) - return torch.where( - denominator > 0, - numerator / denominator, - torch.where(numerator > 0, inf, zero), - ) + @property + def acceptance_rate(self): + if self.n_proposed == 0: + return 0.0 + return self.n_accepted / self.n_proposed -def metropolis_exchange_sweep( - configs, - amplitude_fn, - graph, - *, - current_amplitudes=None, - proposal="spinful", - hopping_rate=0.25, - encoding=None, - generator=None, - chunk_size=None, -): - """Run one nearest-neighbor Metropolis sweep. +@dataclass(frozen=True) +class TorchMCMCSamples: + """Chain-preserving samples and diagnostics from a torch sampler. - ``amplitude_fn`` should accept a ``(batch, n_sites)`` torch integer tensor - and return a batch of amplitudes. The sampler evaluates only changed - proposals when possible. ``chunk_size`` caps proposal-amplitude batch size - without changing the Markov chain. + ``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``. """ - torch = _require_torch() - configs = _as_long_matrix(configs).clone() - current = ( - _call_amplitude_fn(amplitude_fn, configs, chunk_size=chunk_size) - if current_amplitudes is None - else current_amplitudes - ) - current = torch.as_tensor(current, device=configs.device) - n_proposed = 0 - n_accepted = 0 - for i, j in _iter_edges(graph): - if proposal in {"spin", "spin_exchange", "heisenberg"}: - proposed, flags = propose_spin_exchange(i, j, configs) - elif proposal in {"spinful", "hubbard", "spinful_exchange_hopping"}: - proposed, flags = propose_spinful_exchange_or_hopping( - i, - j, - configs, - hopping_rate=hopping_rate, - encoding=encoding, - generator=generator, - ) - else: - raise ValueError( - "proposal must be 'spin' or 'spinful_exchange_hopping'." - ) + configs: Any + amplitudes: Any + n_samples: int + n_samples_per_chain: int + n_chains: int + n_discard_per_chain: int + sweep_size: int + acceptance_rate: float + n_proposed: int + n_accepted: int + elapsed_seconds: float + samples_per_second: float + log_abs_amplitudes: Any = None + proposal_stats: Any = None - if not torch.any(flags): - continue + def diagnostics(self, values=None, *, max_lag=None): + """Compute chain diagnostics for a scalar observable. - n_changed = int(flags.sum().item()) - n_proposed += n_changed - proposed_amps = current.clone() - proposed_amps[flags] = _call_amplitude_fn( - amplitude_fn, - proposed[flags], - chunk_size=chunk_size, - ) - ratio = _safe_metropolis_ratio(proposed_amps, current) - accept = flags & ( - torch.rand(configs.shape[0], device=configs.device, generator=generator) - < ratio + If ``values`` is omitted, the sampled ``|psi|**2`` values are used as + a generic mixing diagnostic. For VMC convergence, pass local + observable values with shape ``(n_samples_per_chain, n_chains)``. + """ + if values is None: + values = self.amplitudes.abs().square() + return torch_chain_diagnostics(values, max_lag=max_lag) + + def to_common(self): + """Convert to the backend-neutral :class:`pepsy.vmc.VMCSamples`.""" + from .api import VMCSamples + return VMCSamples( + configs=self.configs, + amplitudes=self.amplitudes, + log_amplitudes=self.log_abs_amplitudes, + 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, + }, + native=self, ) - if torch.any(accept): - n_accept = int(accept.sum().item()) - n_accepted += n_accept - configs[accept] = proposed[accept] - current[accept] = proposed_amps[accept] - return TorchMetropolisResult( - configs=configs, - amplitudes=current, - n_proposed=n_proposed, - n_accepted=n_accepted, - ) +@dataclass(frozen=True) +class TorchChainDiagnostics: + """MCMC convergence diagnostics for chain-shaped scalar values.""" + r_hat: Any + integrated_autocorrelation_time: Any + effective_sample_size: Any + n_samples_per_chain: int + n_chains: int -def _mode_occupations(n_up, n_down, *, order): + @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_autocorrelation_time + + +def _make_torch_generator(seed, *, device=None): + """Construct a reproducible torch generator for a target device.""" torch = _require_torch() - if order in {"down-up", "du", "symmray"}: - return torch.stack((n_down, n_up), dim=-1).reshape(n_up.shape[0], -1) - if order in {"up-down", "ud", "netket"}: - return torch.stack((n_up, n_down), dim=-1).reshape(n_up.shape[0], -1) - raise ValueError("mode_order must be 'down-up' or 'up-down'.") + if seed is None: + return None + try: + generator = torch.Generator(device=device) + except (RuntimeError, TypeError, ValueError): + generator = torch.Generator() + generator.manual_seed(int(seed)) + return generator -def _mode_index(site, spin, *, order): - if order in {"down-up", "du", "symmray"}: - offset = 1 if spin == "up" else 0 - elif order in {"up-down", "ud", "netket"}: - offset = 0 if spin == "up" else 1 +def _iter_edges(graph): + if hasattr(graph, "edges"): + edges = graph.edges + if callable(edges): + edges = edges() else: - raise ValueError("mode_order must be 'down-up' or 'up-down'.") - return 2 * site + offset + edges = graph + return tuple((int(i), int(j)) for i, j in edges) -def spinful_fermi_hubbard_connections( - configs, - graph, - *, - t=1.0, - U=8.0, - encoding=None, - mode_order="down-up", -): - """Return batched spinful Fermi-Hubbard connected configurations. +def _term_items(terms): + """Return ``(where, operator)`` pairs from common Hamiltonian containers.""" + if hasattr(terms, "terms"): + terms = terms.terms + if hasattr(terms, "items"): + return tuple(terms.items()) + try: + return tuple(terms) + except TypeError as exc: + raise TypeError( + "terms must be a mapping, a SymHamiltonian-like object with " + "`.terms`, or an iterable of (where, operator) pairs." + ) from exc - The local state encoding is configurable. Fermion signs use a site-major - mode order; ``mode_order='down-up'`` matches Symmray/vmc_torch convention. + +def _term_dense_array(operator): + """Convert a local operator to a dense array without changing its backend.""" + if hasattr(operator, "to_dense"): + operator = operator.to_dense() + elif hasattr(operator, "data") and not hasattr(operator, "shape"): + operator = operator.data + return operator + + +def _is_fermionic_operator(operator): + """Return whether ``operator`` carries Symmray fermionic grading.""" + return bool(getattr(operator, "fermionic", False)) or ( + "FermionicArray" in type(operator).__name__ + ) + + +def _expanded_operator_charges(index, *, symmetry=None): + """Expand a Symmray index charge map into linear-index order. + + Sparse Symmray indices expose ``chargemap`` directly. Flat indices do + not, but retain either an explicit private linear map or enough physical + fermion metadata to recover the standard spinless/spinful map. """ - torch = _require_torch() - encoding = FermionSiteEncoding.symmray() if encoding is None else encoding - configs = _as_long_matrix(configs) - batch, n_sites = configs.shape - device = configs.device - n_up, n_down = encoding.decode(configs) - modes = _mode_occupations(n_up, n_down, order=mode_order) + chargemap = getattr(index, "chargemap", None) + if chargemap is not None: + # A spinful Z2 index has two states in each sector. ``BlockIndex`` + # stores those sectors contiguously, while the canonical fermion + # physical basis is ``empty, down, up, double`` and therefore has + # charge map ``(0, 1, 1, 0)``. The block sizes alone cannot recover + # that ordering, so use Symmray's explicit physical map here. + if str(symmetry) == "Z2" and sum(int(size) for size in chargemap.values()) == 4: + import symmray.fermionic_local_operators as flo # noqa: PLC0415 + + return tuple(flo.get_spinful_charge_indexmap("Z2")) + charges = [] + for charge, size in chargemap.items(): + charges.extend([charge] * int(size)) + return tuple(charges) + + # ``FlatIndex.linearmap`` is currently not exposed as a public property + # by Symmray, but the constructor retains it when a non-default physical + # ordering is supplied. Respect it when present. + linearmap = getattr(index, "_linearmap", None) + if linearmap is not None: + return tuple(entry[0] for entry in linearmap) + + size = getattr(index, "size_total", None) + if size is None or symmetry is None: + raise TypeError( + "Fermionic VMC terms require Symmray indices with charge maps " + "or recognizable flat fermion physical indices." + ) + import symmray.fermionic_local_operators as flo # noqa: PLC0415 - all_etas = [] - all_coeffs = [] - all_bids = [] + if int(size) == 2: + charges = flo.get_spinless_charge_indexmap(str(symmetry)) + elif int(size) == 4: + charges = flo.get_spinful_charge_indexmap(str(symmetry)) + else: + raise TypeError( + "Cannot infer fermionic charges for a flat physical index of " + f"dimension {size}; provide a sparse Symmray index instead." + ) + return tuple(charges) - for edge in _iter_edges(graph): - i, j = edge - coeff = float(_edge_value(t, edge)) - if coeff == 0.0: - continue - for spin, occ in (("up", n_up), ("down", n_down)): - valid = occ[:, i] != occ[:, j] - if not torch.any(valid): - continue - idx = valid.nonzero(as_tuple=True)[0] - new_up = n_up[idx].clone() - new_down = n_down[idx].clone() - target = new_up if spin == "up" else new_down - tmp = target[:, i].clone() - target[:, i] = target[:, j] - target[:, j] = tmp - p = _mode_index(i, spin, order=mode_order) - q = _mode_index(j, spin, order=mode_order) - if p > q: - p, q = q, p - between = modes[idx, p + 1:q].sum(dim=-1) % 2 - phase = 1.0 - 2.0 * between.to(torch.float64) +def _charge_parity(charge): + """Return the fermion parity of an Abelian charge.""" + if isinstance(charge, tuple): + return sum(int(value) for value in charge) % 2 + return int(charge) % 2 - all_etas.append(encoding.encode(new_up, new_down)) - all_coeffs.append(-coeff * phase) - all_bids.append(idx) - for site in range(n_sites): - coeff = float(_site_value(U, site)) - if coeff == 0.0: - continue - valid = (n_up[:, site] == 1) & (n_down[:, site] == 1) - if not torch.any(valid): - continue - idx = valid.nonzero(as_tuple=True)[0] - all_etas.append(configs[idx].clone()) - all_coeffs.append(torch.full( - (idx.numel(),), - coeff, - dtype=torch.float64, - device=device, - )) - all_bids.append(idx) +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) - if not all_etas: - return _empty_connections(configs) - return TorchConnections( - configs=torch.cat(all_etas, dim=0), - coeffs=torch.cat(all_coeffs, dim=0), - batch_ids=torch.cat(all_bids, dim=0), +@dataclass(frozen=True) +class _CompiledFermionicTerm: + """Static sparse data for one native fermionic operator term.""" + + operator: Any + rank: int + local_sites: tuple[int, ...] + local_dims: tuple[int, ...] + transitions: tuple[tuple[Any, ...], ...] + input_parity: tuple[np.ndarray, ...] + between: tuple[int, ...] + parity_sites: tuple[int, ...] + + +# Native operators are fixed observables during VMC optimization. Keep a +# strong reference in each value so an ``id``-based key cannot become stale +# through Python object-id reuse. The cap prevents a long-lived driver that +# creates many temporary terms from growing this process-global cache forever. +_FERMION_COMPILED_TERM_CACHE = {} +_FERMION_COMPILED_TERM_CACHE_MAXSIZE = 1024 + + +def _fermionic_operator_shape(operator): + """Return a native operator shape without materializing dense data.""" + shape = getattr(operator, "shape", None) + if shape is None: + indices = getattr(operator, "indices", None) + if indices is None: + return tuple(_operator_dense_numpy(operator).shape) + shape = tuple(getattr(index, "size_total", 0) for index in indices) + try: + return tuple(int(dim) for dim in shape) + except (TypeError, ValueError) as exc: + raise TypeError("Fermionic operator shape must be a finite sequence.") from exc + + +def _compile_fermionic_operator( + operator, + local_sites, + *, + coefficient_cutoff, +): + """Compile a native fermionic operator into configuration transitions.""" + shape = _fermionic_operator_shape(operator) + rank = len(shape) + if rank not in (2, 4): + raise ValueError( + f"Hamiltonian operators must have rank 2 or 4, got rank {rank}." + ) + split = rank // 2 + if shape[:split] != shape[split:]: + raise ValueError( + "Fermionic Hamiltonian terms must have square input/output " + f"dimensions, got shape {shape}." + ) + + indices = getattr(operator, "indices", None) + if indices is None or len(indices) != rank: + raise TypeError( + "Fermionic VMC terms require native Symmray operator indices." + ) + symmetry = getattr(operator, "symmetry", None) + output_charges = tuple( + _expanded_operator_charges(index, symmetry=symmetry) + for index in indices[:split] + ) + input_charges = tuple( + _expanded_operator_charges(index, symmetry=symmetry) + for index in indices[split:] ) + if any(len(values) != shape[axis] for axis, values in enumerate(output_charges)): + raise ValueError("Fermionic operator output charge maps do not match its shape.") + if any( + len(values) != shape[split + axis] + for axis, values in enumerate(input_charges) + ): + raise ValueError("Fermionic operator input charge maps do not match its shape.") + dense = _operator_dense_numpy(operator) + if dense.shape != shape: + raise ValueError( + "Fermionic operator dense data shape does not match its native " + f"shape: {dense.shape} != {shape}." + ) -def heisenberg_connections(configs, graph, *, J=1.0): - """Return batched spin-1/2 Heisenberg connected configurations.""" + # The VMC configuration order follows the physical Symmray index order. + # A labelled term can be supplied in reverse site order; transpose both + # physical sides before compiling the left-to-right parity string. + if rank == 4 and local_sites[0] > local_sites[1]: + dense = dense.transpose(1, 0, 3, 2) + local_sites = (local_sites[1], local_sites[0]) + output_charges = (output_charges[1], output_charges[0]) + input_charges = (input_charges[1], input_charges[0]) + + transitions = [] + if rank == 2: + nonzero = np.argwhere(np.abs(dense) > float(coefficient_cutoff)) + output_parity = np.asarray( + [_charge_parity(charge) for charge in output_charges[0]], + dtype=np.int64, + ) + input_parity = np.asarray( + [_charge_parity(charge) for charge in input_charges[0]], + dtype=np.int64, + ) + for output, input_ in nonzero: + output = int(output) + input_ = int(input_) + transitions.append( + ( + output, + input_, + dense[output, input_], + int(output_parity[output] ^ input_parity[input_]), + ) + ) + return _CompiledFermionicTerm( + operator=operator, + rank=rank, + local_sites=tuple(local_sites), + local_dims=(shape[0],), + transitions=tuple(transitions), + input_parity=(input_parity,), + between=(), + parity_sites=( + tuple(range(local_sites[0])) + if any(transition[3] for transition in transitions) + else () + ), + ) + + # Symmray's raw two-site tensor uses fermionic tensor-product ordering. + # Convert its endpoint crossing phase once, at compile time. The dynamic + # part left for each batch is only the parity string on intermediate sites. + input_parity = tuple( + np.asarray( + [_charge_parity(charge) for charge in charges], + dtype=np.int64, + ) + for charges in input_charges + ) + crossing = np.ones((shape[2], shape[3]), dtype=np.int8) + crossing[np.ix_(input_parity[0].astype(bool), input_parity[1].astype(bool))] = -1 + dense = dense * crossing[None, None, :, :] + + output_left_parity = np.asarray( + [_charge_parity(charge) for charge in output_charges[0]], + dtype=np.int64, + ) + input_left_parity = input_parity[0] + nonzero = np.argwhere(np.abs(dense) > float(coefficient_cutoff)) + for output_left, output_right, input_left, input_right in nonzero: + output_left = int(output_left) + output_right = int(output_right) + input_left = int(input_left) + input_right = int(input_right) + transitions.append( + ( + output_left, + output_right, + input_left, + input_right, + dense[output_left, output_right, input_left, input_right], + int(output_left_parity[output_left] ^ input_left_parity[input_left]), + ) + ) + + left, right = local_sites + return _CompiledFermionicTerm( + operator=operator, + rank=rank, + local_sites=tuple(local_sites), + local_dims=(shape[0], shape[1]), + transitions=tuple(transitions), + input_parity=input_parity, + between=tuple(range(left + 1, right)), + parity_sites=tuple(range(left + 1, right)), + ) + + +def _get_compiled_fermionic_operator( + operator, + where, + *, + site_order, + n_sites, + coefficient_cutoff, +): + """Get or create the static compilation for one native term.""" + rank = len(_fermionic_operator_shape(operator)) + raw_sites = _term_site_indices( + where, + rank, + site_order=site_order, + n_sites=n_sites, + ) + if any(site < 0 or site >= n_sites for site in raw_sites): + raise ValueError( + f"Hamiltonian term at {where!r} resolves outside the supplied " + f"configuration width {n_sites}." + ) + if rank == 4 and raw_sites[0] == raw_sites[1]: + raise ValueError( + "A native two-site fermionic operator must act on two distinct " + "configuration sites." + ) + key = (id(operator), tuple(raw_sites), int(n_sites), float(coefficient_cutoff)) + compiled = _FERMION_COMPILED_TERM_CACHE.get(key) + if compiled is not None and compiled.operator is operator: + return compiled + + compiled = _compile_fermionic_operator( + operator, + raw_sites, + coefficient_cutoff=coefficient_cutoff, + ) + if len(_FERMION_COMPILED_TERM_CACHE) >= _FERMION_COMPILED_TERM_CACHE_MAXSIZE: + _FERMION_COMPILED_TERM_CACHE.pop(next(iter(_FERMION_COMPILED_TERM_CACHE))) + _FERMION_COMPILED_TERM_CACHE[key] = compiled + return compiled + + +def _fermionic_operator_connections( + configs, + where, + operator, + *, + site_order, + coefficient_cutoff=0.0, +): + """Build connections for one native graded fermionic operator. + + ``FermionicArray.to_dense`` exposes Symmray's raw tensor data. Treating + that data as an ordinary matrix loses the crossing phase and the + Jordan-Wigner parity string between separated sites. This routine applies + the same conversion used by Pepsy's generic fermionic MPO path, but keeps + the result as sparse configuration transitions for VMC. + """ torch = _require_torch() configs = _as_long_matrix(configs) - batch = configs.shape[0] - device = configs.device - batch_ids = torch.arange(batch, dtype=torch.long, device=device) + compiled = _get_compiled_fermionic_operator( + operator, + where, + site_order=site_order, + n_sites=configs.shape[1], + coefficient_cutoff=coefficient_cutoff, + ) + rank = compiled.rank + local_sites = compiled.local_sites + if configs.numel() and int(configs.min()) < 0: + raise ValueError("Fermionic configurations must use non-negative local codes.") + if configs.numel() and any( + int(configs[:, site].max()) >= local_dim + for site, local_dim in zip(local_sites, compiled.local_dims) + ): + raise ValueError( + "Fermionic Hamiltonian term has a local dimension too small for " + "the supplied configurations." + ) all_etas = [] all_coeffs = [] all_bids = [] - for edge in _iter_edges(graph): - i, j = edge - coeff = float(_edge_value(J, edge)) - if coeff == 0.0: - continue - diff = configs[:, i] != configs[:, j] - if torch.any(diff): - idx = diff.nonzero(as_tuple=True)[0] - eta = configs[idx].clone() - tmp = eta[:, i].clone() - eta[:, i] = eta[:, j] - eta[:, j] = tmp + if rank == 2: + site = local_sites[0] + if compiled.parity_sites: + config_parity = torch.as_tensor( + compiled.input_parity[0], + dtype=torch.long, + device=configs.device, + ) + prefix_parity = ( + config_parity[configs[:, compiled.parity_sites]].sum(dim=-1) % 2 + ) + else: + prefix_parity = None + for output, input_, coefficient, transfer_parity in compiled.transitions: + mask = configs[:, site] == input_ + batch_ids = mask.nonzero(as_tuple=True)[0] + if batch_ids.numel() == 0: + continue + eta = configs[batch_ids].clone() + eta[:, site] = output all_etas.append(eta) - all_coeffs.append(torch.full( - (idx.numel(),), - 0.5 * coeff, - dtype=torch.float64, - device=device, - )) - all_bids.append(idx) - - diag_sign = 1.0 - 2.0 * ((configs[:, i] - configs[:, j]).abs() % 2).to( - torch.float64 - ) - all_etas.append(configs.clone()) - all_coeffs.append(0.25 * coeff * diag_sign) - all_bids.append(batch_ids.clone()) + coefficient = torch.as_tensor(coefficient, device=configs.device) + if prefix_parity is not None and transfer_parity: + phase = torch.where( + prefix_parity[batch_ids] == 1, + torch.as_tensor(-1.0, device=configs.device), + torch.as_tensor(1.0, device=configs.device), + ) + coefficient = coefficient * phase + all_coeffs.append(coefficient.expand(batch_ids.numel())) + all_bids.append(batch_ids) + else: + # Symmray's raw tensor uses fermionic tensor-product ordering. The + # crossing phase appears when the two endpoint ket indices are both + # odd. A separated term additionally carries the parity string on + # sites strictly between its endpoints. + left, right = local_sites + if compiled.between: + config_parity = torch.as_tensor( + compiled.input_parity[0], + dtype=torch.long, + device=configs.device, + ) + between_parity = config_parity[configs[:, compiled.between]].sum(dim=-1) % 2 + else: + between_parity = None + + for ( + output_left, + output_right, + input_left, + input_right, + coefficient, + transfer_parity, + ) in compiled.transitions: + mask = (configs[:, left] == input_left) & ( + configs[:, right] == input_right + ) + batch_ids = mask.nonzero(as_tuple=True)[0] + if batch_ids.numel() == 0: + continue + eta = configs[batch_ids].clone() + eta[:, left] = output_left + eta[:, right] = output_right + if between_parity is None: + phase = 1.0 + else: + if transfer_parity: + phase = torch.where( + between_parity[batch_ids] == 1, + torch.as_tensor(-1.0, device=configs.device), + torch.as_tensor(1.0, device=configs.device), + ) + else: + phase = 1.0 + coefficient = torch.as_tensor( + coefficient, + device=configs.device, + ) + if not isinstance(phase, float): + coefficient = coefficient * phase + all_etas.append(eta) + all_coeffs.append(coefficient.expand(batch_ids.numel())) + all_bids.append(batch_ids) if not all_etas: return _empty_connections(configs) - return TorchConnections( configs=torch.cat(all_etas, dim=0), coeffs=torch.cat(all_coeffs, dim=0), @@ -1565,277 +4642,5191 @@ def heisenberg_connections(configs, graph, *, J=1.0): ) -def transverse_ising_connections(configs, graph, *, J=1.0, h=1.0): - """Return batched transverse-field Ising connected configurations.""" +def _term_site_indices(where, rank, *, site_order, n_sites): + """Resolve a one- or two-site term location to config-column indices.""" + n_local_sites = rank // 2 + if n_local_sites not in (1, 2): + raise ValueError( + "Hamiltonian operators must have rank 2 or 4, with output axes " + "followed by input axes." + ) + if n_local_sites == 1: + where = (where,) + elif isinstance(where, (str, bytes)): + raise ValueError("A two-site operator location must contain two sites.") + else: + try: + where = tuple(where) + except TypeError as exc: + raise ValueError( + "A two-site operator location must contain two sites." + ) from exc + if len(where) != 2: + raise ValueError("A two-site operator location must contain two sites.") + + if site_order is None: + site_order = tuple(range(n_sites)) + position = {site: i for i, site in enumerate(site_order)} + missing = [site for site in where if site not in position] + if missing: + raise ValueError( + f"Hamiltonian term site(s) {missing!r} are not in site_order. " + "Pass site_order matching the PEPS physical-site order." + ) + return tuple(position[site] for site in where) + + +def torch_hamiltonian_connections( + configs, + terms, + *, + site_order=None, + coefficient_cutoff=0.0, + constant=0.0, +): + """Build connected configurations from explicit local Hamiltonian terms. + + ``terms`` can be a :class:`SymHamiltonian`, its ``.terms`` mapping, or an + iterable of ``(where, operator)`` pairs. Ordinary dense operators expose + output axes followed by input axes. Native Symmray fermionic operators + use the same axis convention, but their graded crossing phase and the + parity string between separated sites are preserved before connections + are emitted. ``constant`` adds a diagonal identity contribution for every + parent configuration. This lets torch VMC measure arbitrary supplied terms, + including spinless operators, without guessing ``t``, ``U``, or a + model-specific connection function. + """ torch = _require_torch() configs = _as_long_matrix(configs) batch, n_sites = configs.shape device = configs.device - batch_ids = torch.arange(batch, dtype=torch.long, device=device) all_etas = [] all_coeffs = [] all_bids = [] - for edge in _iter_edges(graph): - i, j = edge - coeff = float(_edge_value(J, edge)) - if coeff == 0.0: + for where, operator in _term_items(terms): + if _is_fermionic_operator(operator): + term_connections = _fermionic_operator_connections( + configs, + where, + operator, + site_order=site_order, + coefficient_cutoff=coefficient_cutoff, + ) + if term_connections.configs.shape[0]: + all_etas.append(term_connections.configs) + all_coeffs.append(term_connections.coeffs) + all_bids.append(term_connections.batch_ids) continue - diag_sign = 1.0 - 2.0 * ((configs[:, i] - configs[:, j]).abs() % 2).to( - torch.float64 + + dense = torch.as_tensor(_term_dense_array(operator), device=device) + if dense.ndim not in (2, 4): + raise ValueError( + f"Hamiltonian term at {where!r} has rank {dense.ndim}; " + "only one- and two-site terms are supported." + ) + if dense.shape[: dense.ndim // 2] != dense.shape[dense.ndim // 2 :]: + raise ValueError( + f"Hamiltonian term at {where!r} must have square input/output " + f"dimensions, got shape {tuple(dense.shape)}." + ) + local_sites = _term_site_indices( + where, + dense.ndim, + site_order=site_order, + n_sites=n_sites, ) - all_etas.append(configs.clone()) - all_coeffs.append(0.25 * coeff * diag_sign) - all_bids.append(batch_ids.clone()) + local_dim = int(dense.shape[0]) + if any(int(configs[:, site].max()) >= local_dim for site in local_sites): + raise ValueError( + f"Hamiltonian term at {where!r} has local dimension {local_dim}, " + "which is too small for the supplied configurations." + ) - for site in range(n_sites): - coeff = float(_site_value(h, site)) - if coeff == 0.0: - continue - eta = configs.clone() - eta[:, site] = 1 - eta[:, site] - all_etas.append(eta) - all_coeffs.append(torch.full( - (batch,), - 0.5 * coeff, - dtype=torch.float64, - device=device, - )) - all_bids.append(batch_ids.clone()) + nonzero = torch.nonzero( + dense.abs() > float(coefficient_cutoff), + as_tuple=False, + ) + n_local_sites = len(local_sites) + for entry in nonzero: + entry = tuple(int(x) for x in entry.tolist()) + outputs = entry[:n_local_sites] + inputs = entry[n_local_sites:] + mask = torch.ones(batch, dtype=torch.bool, device=device) + for site, value in zip(local_sites, inputs): + mask &= configs[:, site] == value + batch_ids = mask.nonzero(as_tuple=True)[0] + if batch_ids.numel() == 0: + continue + eta = configs[batch_ids].clone() + for site, value in zip(local_sites, outputs): + eta[:, site] = value + all_etas.append(eta) + all_coeffs.append( + dense[entry].expand(batch_ids.numel()).to(device=device) + ) + all_bids.append(batch_ids) + if constant: + identity = configs.clone() + all_etas.append(identity) + all_coeffs.append( + torch.as_tensor(constant, device=device).expand(batch) + ) + all_bids.append(torch.arange(batch, device=device, dtype=torch.long)) if not all_etas: return _empty_connections(configs) - + coefficient_dtype = all_coeffs[0].dtype + for values in all_coeffs[1:]: + coefficient_dtype = torch.promote_types(coefficient_dtype, values.dtype) return TorchConnections( configs=torch.cat(all_etas, dim=0), - coeffs=torch.cat(all_coeffs, dim=0), + coeffs=torch.cat( + [values.to(dtype=coefficient_dtype) for values in all_coeffs], + dim=0, + ), batch_ids=torch.cat(all_bids, dim=0), ) -def local_energy_from_connections( - configs, - amplitudes, - connections, - amplitude_fn, - *, - chunk_size=None, - reuse_diagonal=True, -): - """Accumulate local energies from connected configs and amplitudes. +def _driver_terms_connections(configs, graph, *, terms, site_order=None, **kwargs): + """Adapt explicit-term connections to the driver's connection signature.""" + del graph + return torch_hamiltonian_connections( + configs, + terms, + site_order=site_order, + constant=kwargs.get("constant", 0.0), + ) - If ``amplitude_fn`` exposes ``connected_amplitudes(...)`` that method is - used. Otherwise diagonal connections can reuse the supplied parent - amplitudes and off-diagonal amplitudes are evaluated in optional chunks. + +def compile_operator_sum_torch(terms, *, fermion=None, site_order=None): + """Compile a backend-neutral operator sum for Torch connections. + + Matrix terms are passed through unchanged. Symbolic fermion terms are + lowered with ``Fermion.operator_term`` so native Symmray grading and + Jordan--Wigner parity remain in the existing Torch connection compiler. + The returned :class:`CompiledOperatorSum` keeps an identity constant + separate because connection tables represent the non-constant terms. """ - torch = _require_torch() - configs = _as_long_matrix(configs) - amplitudes = torch.as_tensor(amplitudes, device=configs.device) - if connections.configs.numel() == 0: - return torch.zeros( - configs.shape[0], - dtype=amplitudes.dtype, - device=configs.device, - ) + from .api import ( + CompiledOperatorSum, + LocalMatrixTerm, + ProductTerm, + _expand_fermion_factor, + normalize_operator_sum, + ) - connected_amplitudes = getattr(amplitude_fn, "connected_amplitudes", None) - if callable(connected_amplitudes): - conn_amps = connected_amplitudes( + operator_sum = normalize_operator_sum(terms) + site_order = None if site_order is None else tuple(site_order) + + def map_site(site): + if site_order is None: + return site + if site in site_order: + return site + if isinstance(site, Integral) and 0 <= int(site) < len(site_order): + return site_order[int(site)] + raise ValueError(f"Term site {site!r} is not present in site_order.") + + def native_fermion_name(name): + """Map public VMC spin labels to the native local-mode order. + + ``Fermion.operator_term`` stores its spinful local modes in the + opposite charge-axis order to :class:`FermionSiteEncoding.vmc_torch`. + Keep that implementation detail inside this adapter so a common + ``OperatorFactor(..., spin=\"up\")`` has the same meaning in Torch + and NetKet. + """ + return { + "create_u": "create_d", + "annihilate_u": "annihilate_d", + "create_d": "create_u", + "annihilate_d": "annihilate_u", + }.get(name, name) + + compiled = [] + for term in operator_sum: + if isinstance(term, LocalMatrixTerm): + operator = term.matrix + if term.coefficient != 1: + operator = operator * term.coefficient + mapped_support = tuple(map_site(site) for site in term.support) + where = mapped_support[0] if len(mapped_support) == 1 else mapped_support + compiled.append((where, operator)) + continue + if not isinstance(term, ProductTerm): # pragma: no cover - guarded by IR + raise TypeError(f"Unsupported operator term {type(term).__name__}.") + if fermion is None: + raise ValueError( + "Symbolic ProductTerm entries require fermion=... when compiling " + "for Torch." + ) + references = [] + for factor in term.factors: + references.extend( + (map_site(site), native_fermion_name(name)) + for site, name in _expand_fermion_factor(factor) + ) + # A symbolic local density/doublon expands into several native + # creation/annihilation factors at the same site. ``operator_term`` + # needs the operator support only once, while ``references`` retains + # the full product (and its order). + mapped_support = tuple(dict.fromkeys(site for site, _ in references)) + # Keep the native tensor legs in configuration order. In particular, + # the Hermitian-conjugate hopping monomial has reversed *factor* + # order, but it must share the same two-site graded tensor layout as + # its forward partner. ``references`` above intentionally remains + # untouched, so this is not an operator reordering. + if site_order is not None: + positions = {site: position for position, site in enumerate(site_order)} + mapped_support = tuple( + sorted(mapped_support, key=lambda site: positions[site]) + ) + elif all(isinstance(site, Integral) for site in mapped_support): + mapped_support = tuple(sorted(mapped_support)) + operator = fermion.operator_term( + [(term.coefficient, tuple(references))], + sites=mapped_support, + ) + where = mapped_support[0] if len(mapped_support) == 1 else mapped_support + compiled.append((where, operator)) + return CompiledOperatorSum( + backend="torch", + terms=tuple(compiled), + constant=operator_sum.constant, + metadata=operator_sum.metadata, + ) + + +def _normalize_terms_site_labels(terms, site_order): + """Map positional integer term labels onto PEPS site labels when needed.""" + site_order = tuple(site_order) + positions = {site: i for i, site in enumerate(site_order)} + + def map_site(site): + if site in positions: + return site + if ( + isinstance(site, Integral) + and not isinstance(site, bool) + and 0 <= int(site) < len(site_order) + ): + return site_order[int(site)] + return site + + normalized = {} + for where, operator in _term_items(terms): + dense = _term_dense_array(operator) + rank = getattr(dense, "ndim", None) + if rank is None: + rank = len(getattr(dense, "shape", ())) + n_local_sites = int(rank) // 2 + if n_local_sites == 1: + normalized[map_site(where)] = operator + elif n_local_sites == 2: + try: + left, right = tuple(where) + except (TypeError, ValueError) as exc: + raise ValueError( + "A two-site observable location must contain two sites." + ) from exc + normalized[(map_site(left), map_site(right))] = operator + else: + normalized[where] = operator + return normalized + + +def _empty_connections(configs): + torch = _require_torch() + return TorchConnections( + configs=configs.new_empty((0, configs.shape[1])), + coeffs=torch.empty(0, dtype=torch.float64, device=configs.device), + batch_ids=torch.empty(0, dtype=torch.long, device=configs.device), + ) + + +def count_spinful_particles(configs, *, encoding=None): + """Return per-sample ``(n_up, n_down)`` counts.""" + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding + configs = _as_long_matrix(configs) + n_up, n_down = encoding.decode(configs) + return n_up.sum(dim=-1), n_down.sum(dim=-1) + + +def _run_cheap_torch_kernel(name, fn, *args, compile_kernels=False): + """Optionally compile pure tensor bookkeeping with an eager fallback.""" + if not compile_kernels or name in _FAILED_CHEAP_TORCH_KERNELS: + return fn(*args) + torch = _require_torch() + compiled = _COMPILED_CHEAP_TORCH_KERNELS.get(name) + if compiled is None: + compile_fn = getattr(torch, "compile", None) + include_dir = sysconfig.get_config_var("INCLUDEPY") + has_python_headers = ( + include_dir is not None + and os.path.isfile(os.path.join(include_dir, "Python.h")) + ) + if not callable(compile_fn) or not has_python_headers: + _FAILED_CHEAP_TORCH_KERNELS.add(name) + return fn(*args) + try: + compiled = compile_fn(fn, dynamic=True) + except (AttributeError, RuntimeError, TypeError, ValueError): + _FAILED_CHEAP_TORCH_KERNELS.add(name) + return fn(*args) + _COMPILED_CHEAP_TORCH_KERNELS[name] = compiled + try: + return compiled(*args) + except (AttributeError, RuntimeError, TypeError, ValueError): + # The feature is opt-in. Sparse/device-specific compiler gaps should + # never alter a VMC trajectory or make PEPS evaluation unavailable. + _FAILED_CHEAP_TORCH_KERNELS.add(name) + _COMPILED_CHEAP_TORCH_KERNELS.pop(name, None) + return fn(*args) + + +def _proposal_sites(i, j, configs): + """Build fixed-shape edge indices for proposal kernels.""" + torch = _require_torch() + sites = torch.as_tensor((i, j), dtype=torch.long, device=configs.device) + return sites.reshape(1, 2).expand(configs.shape[0], -1) + + +def _spin_exchange_kernel(configs, sites): + torch = _require_torch() + endpoints = torch.gather(configs, 1, sites) + changed = endpoints[:, 0] != endpoints[:, 1] + values = torch.stack( + ( + torch.where(changed, endpoints[:, 1], endpoints[:, 0]), + torch.where(changed, endpoints[:, 0], endpoints[:, 1]), + ), + dim=1, + ) + return configs.scatter(1, sites, values), changed + + +_PROPOSAL_MOVE_NAMES = ( + "exchange", + "hopping", + "spin_flip", + "pair_toggle", +) +_MOVE_EXCHANGE, _MOVE_HOPPING, _MOVE_SPIN_FLIP, _MOVE_PAIR_TOGGLE = range(4) + + +def _spinful_exchange_hopping_kernel( + configs, + sites, + hopping_rate, + encoding_codes, + branch_random, + d0_random, + d2_random, +): + """Branch-free U1U1 proposal core suitable for ``torch.compile``.""" + torch = _require_torch() + endpoints = torch.gather(configs, 1, sites) + ci, cj = endpoints[:, 0], endpoints[:, 1] + empty, double, up, down = encoding_codes.unbind() + changed = ci != cj + + n_up_i = ((ci == up) | (ci == double)).to(torch.long) + n_up_j = ((cj == up) | (cj == double)).to(torch.long) + n_down_i = ((ci == down) | (ci == double)).to(torch.long) + n_down_j = ((cj == down) | (cj == double)).to(torch.long) + delta_n = ((n_up_i + n_down_i) - (n_up_j + n_down_j)).abs() + + exchange = (branch_random < (1.0 - hopping_rate)) & changed + hopping = (~exchange) & changed + swap = exchange | (hopping & (delta_n == 1)) + next_i = torch.where(swap, cj, ci) + next_j = torch.where(swap, ci, cj) + + d0 = hopping & (delta_n == 0) + d0_i = torch.where(d0_random, double, empty) + d0_j = torch.where(d0_random, empty, double) + next_i = torch.where(d0, d0_i, next_i) + next_j = torch.where(d0, d0_j, next_j) + + d2 = hopping & (delta_n == 2) + d2_i = torch.where(d2_random, down, up) + d2_j = torch.where(d2_random, up, down) + next_i = torch.where(d2, d2_i, next_i) + next_j = torch.where(d2, d2_j, next_j) + move_codes = torch.where( + branch_random < (1.0 - hopping_rate), + torch.full_like(branch_random, _MOVE_EXCHANGE, dtype=torch.long), + torch.full_like(branch_random, _MOVE_HOPPING, dtype=torch.long), + ) + return ( + configs.scatter(1, sites, torch.stack((next_i, next_j), dim=1)), + changed, + move_codes, + ) + + +def propose_spin_exchange( + i, + j, + configs, + *, + compile_kernels=False, + _return_move_codes=False, +): + """Propose spin exchange on one edge for binary spin configs.""" + configs = _as_long_matrix(configs) + proposed, changed = _run_cheap_torch_kernel( + "spin-exchange-proposal", + _spin_exchange_kernel, + configs, + _proposal_sites(i, j, configs), + compile_kernels=compile_kernels, + ) + if _return_move_codes: + torch = _require_torch() + return ( + proposed, + changed, + torch.full( + (configs.shape[0],), + _MOVE_EXCHANGE, + dtype=torch.long, + device=configs.device, + ), + ) + return proposed, changed + + +def propose_spinful_exchange_or_hopping( + i, + j, + configs, + *, + hopping_rate=0.25, + encoding=None, + generator=None, + compile_kernels=False, + _return_move_codes=False, +): + """Propose spinful Hubbard exchange/hopping moves on one edge. + + The proposal preserves ``N_up`` and ``N_down``. With probability + ``1 - hopping_rate`` it swaps the two local site states. Otherwise it uses + local hopping-style moves over ``empty/up/down/double`` states, following + the sampling options in ``sjdu10/vmc_torch``. + """ + torch = _require_torch() + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding + configs = _as_long_matrix(configs) + device = configs.device + batch = configs.shape[0] + + if compile_kernels: + encoding_codes = torch.as_tensor( + (encoding.empty, encoding.double, encoding.up, encoding.down), + dtype=configs.dtype, + device=device, + ) + randoms = torch.rand( + (3, batch), + device=device, + generator=generator, + ) + result = _run_cheap_torch_kernel( + "spinful-exchange-hopping-proposal", + _spinful_exchange_hopping_kernel, configs, - amplitudes, - connections, - chunk_size=chunk_size, - reuse_diagonal=reuse_diagonal, + _proposal_sites(i, j, configs), + float(hopping_rate), + encoding_codes, + randoms[0], + randoms[1] < 0.5, + randoms[2] < 0.5, + compile_kernels=True, + ) + if _return_move_codes: + return result + return result[:2] + + proposed = configs.clone() + + ci = configs[:, i] + cj = configs[:, j] + changed = ci != cj + if not torch.any(changed): + if _return_move_codes: + rand = torch.rand(batch, device=device, generator=generator) + move_codes = torch.where( + rand < (1.0 - hopping_rate), + torch.full_like(rand, _MOVE_EXCHANGE, dtype=torch.long), + torch.full_like(rand, _MOVE_HOPPING, dtype=torch.long), + ) + return proposed, changed, move_codes + return proposed, changed + + n_up, n_down = encoding.decode(configs) + ni = n_up[:, i] + n_down[:, i] + nj = n_up[:, j] + n_down[:, j] + delta_n = (ni - nj).abs() + + rand = torch.rand(batch, device=device, generator=generator) + is_exchange = (rand < (1.0 - hopping_rate)) & changed + is_hopping = (~is_exchange) & changed + move_codes = torch.where( + rand < (1.0 - hopping_rate), + torch.full_like(rand, _MOVE_EXCHANGE, dtype=torch.long), + torch.full_like(rand, _MOVE_HOPPING, dtype=torch.long), + ) + + swap_mask = is_exchange | (is_hopping & (delta_n == 1)) + proposed[swap_mask, i] = cj[swap_mask] + proposed[swap_mask, j] = ci[swap_mask] + + mask_d0 = is_hopping & (delta_n == 0) + if torch.any(mask_d0): + bits = torch.randint( + 0, + 2, + (batch,), + device=device, + dtype=torch.long, + generator=generator, + ).bool() + proposed[mask_d0, i] = torch.where( + bits[mask_d0], + torch.as_tensor(encoding.double, device=device), + torch.as_tensor(encoding.empty, device=device), + ) + proposed[mask_d0, j] = torch.where( + bits[mask_d0], + torch.as_tensor(encoding.empty, device=device), + torch.as_tensor(encoding.double, device=device), + ) + + mask_d2 = is_hopping & (delta_n == 2) + if torch.any(mask_d2): + bits = torch.randint( + 0, + 2, + (batch,), + device=device, + dtype=torch.long, + generator=generator, + ).bool() + proposed[mask_d2, i] = torch.where( + bits[mask_d2], + torch.as_tensor(encoding.down, device=device), + torch.as_tensor(encoding.up, device=device), + ) + proposed[mask_d2, j] = torch.where( + bits[mask_d2], + torch.as_tensor(encoding.up, device=device), + torch.as_tensor(encoding.down, device=device), ) + + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + +def propose_spinful_u1_exchange_or_hopping( + i, + j, + configs, + *, + hopping_rate=0.25, + spin_flip_rate=0.25, + encoding=None, + generator=None, + compile_kernels=False, + _return_move_codes=False, +): + """Propose moves that preserve total spinful particle number only. + + In addition to the ``U1U1``-safe exchange and hopping moves, this rule + includes single-site ``up <-> down`` flips. Those flips allow a ``U1`` + walker to move between different spin-resolved sectors while preserving + ``N_up + N_down``. The selected edge and endpoint are fixed before the + local state is inspected, so no proposal-probability correction is needed + for the spin-flip branch. + """ + torch = _require_torch() + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding + configs = _as_long_matrix(configs) + proposal_result = propose_spinful_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=_return_move_codes, + ) + if _return_move_codes: + proposed, changed, move_codes = proposal_result else: - conn_amps = _default_connected_amplitudes( - configs, - amplitudes, - connections, - amplitude_fn, - chunk_size=chunk_size, - reuse_diagonal=reuse_diagonal, + proposed, changed = proposal_result + + spin_flip_rate = float(spin_flip_rate) + if not 0.0 <= spin_flip_rate <= 1.0: + raise ValueError("spin_flip_rate must be between 0 and 1.") + if spin_flip_rate == 0.0: + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + device = configs.device + ci = configs[:, i] + cj = configs[:, j] + flip_branch = torch.rand( + configs.shape[0], + device=device, + generator=generator, + ) < spin_flip_rate + if _return_move_codes: + move_codes = torch.where( + flip_branch, + torch.full_like(move_codes, _MOVE_SPIN_FLIP), + move_codes, ) - conn_amps = torch.as_tensor(conn_amps, device=configs.device) - ratios = conn_amps / amplitudes[connections.batch_ids] - contrib = connections.coeffs.to(dtype=ratios.dtype) * ratios - energy = torch.zeros( + choose_i = torch.rand( configs.shape[0], - dtype=contrib.dtype, - device=configs.device, + device=device, + generator=generator, + ) < 0.5 + target = torch.where(choose_i, ci, cj) + valid = (target == encoding.up) | (target == encoding.down) + flip = flip_branch & valid + proposed[flip_branch] = configs[flip_branch] + if not torch.any(flip): + changed = changed & ~flip_branch + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + flipped = torch.where( + target == encoding.up, + torch.as_tensor(encoding.down, device=device), + torch.as_tensor(encoding.up, device=device), + ) + proposed[flip & choose_i, i] = flipped[flip & choose_i] + proposed[flip & ~choose_i, j] = flipped[flip & ~choose_i] + changed = torch.where(flip_branch, flip, changed) + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + +def propose_spinful_z2_exchange_or_hopping( + i, + j, + configs, + *, + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + encoding=None, + generator=None, + compile_kernels=False, + _return_move_codes=False, +): + """Propose moves that preserve spinful total fermion parity. + + The U1-preserving exchange, hopping, and spin-flip moves are augmented by + an ``empty <-> double`` toggle on a randomly selected endpoint. The latter + changes particle number by two, allowing the chain to explore the full + fixed-parity sector rather than remaining in one fixed-number sector. + """ + torch = _require_torch() + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding + configs = _as_long_matrix(configs) + proposal_result = propose_spinful_u1_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + spin_flip_rate=spin_flip_rate, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=_return_move_codes, + ) + if _return_move_codes: + proposed, changed, move_codes = proposal_result + else: + proposed, changed = proposal_result + + pair_toggle_rate = float(pair_toggle_rate) + if not 0.0 <= pair_toggle_rate <= 1.0: + raise ValueError("pair_toggle_rate must be between 0 and 1.") + if pair_toggle_rate == 0.0: + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + device = configs.device + ci = configs[:, i] + cj = configs[:, j] + pair_branch = torch.rand( + configs.shape[0], + device=device, + generator=generator, + ) < pair_toggle_rate + if _return_move_codes: + move_codes = torch.where( + pair_branch, + torch.full_like(move_codes, _MOVE_PAIR_TOGGLE), + move_codes, + ) + choose_i = torch.rand( + configs.shape[0], + device=device, + generator=generator, + ) < 0.5 + target = torch.where(choose_i, ci, cj) + valid = (target == encoding.empty) | (target == encoding.double) + pair_toggle = pair_branch & valid + proposed[pair_branch] = configs[pair_branch] + if not torch.any(pair_toggle): + changed = changed & ~pair_branch + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + toggled = torch.where( + target == encoding.empty, + torch.as_tensor(encoding.double, device=device), + torch.as_tensor(encoding.empty, device=device), + ) + proposed[pair_toggle & choose_i, i] = toggled[pair_toggle & choose_i] + proposed[pair_toggle & ~choose_i, j] = toggled[pair_toggle & ~choose_i] + changed = torch.where(pair_branch, pair_toggle, changed) + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + +def propose_spinful_z2z2_exchange_or_hopping( + i, + j, + configs, + *, + hopping_rate=0.25, + pair_toggle_rate=0.25, + encoding=None, + generator=None, + compile_kernels=False, + _return_move_codes=False, +): + """Propose moves preserving spin-resolved parity ``Z2 x Z2``. + + Spin flips are deliberately disabled because they change both resolved + parities. Exchange, species-preserving hopping, and empty/double toggles + preserve each parity independently. + """ + torch = _require_torch() + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding + configs = _as_long_matrix(configs) + proposal_result = propose_spinful_u1_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + spin_flip_rate=0.0, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=_return_move_codes, ) - energy.index_add_(0, connections.batch_ids, contrib) - return energy + if _return_move_codes: + proposed, changed, move_codes = proposal_result + else: + proposed, changed = proposal_result + pair_toggle_rate = float(pair_toggle_rate) + if not 0.0 <= pair_toggle_rate <= 1.0: + raise ValueError("pair_toggle_rate must be between 0 and 1.") + if pair_toggle_rate == 0.0: + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + ci = configs[:, i] + cj = configs[:, j] + pair_branch = torch.rand( + configs.shape[0], + device=configs.device, + generator=generator, + ) < pair_toggle_rate + if _return_move_codes: + move_codes = torch.where( + pair_branch, + torch.full_like(move_codes, _MOVE_PAIR_TOGGLE), + move_codes, + ) + valid_empty_double = (ci == encoding.empty) & (cj == encoding.empty) + valid_double_empty = (ci == encoding.double) & (cj == encoding.double) + valid_up_up = (ci == encoding.up) & (cj == encoding.up) + valid_down_down = (ci == encoding.down) & (cj == encoding.down) + valid = ( + valid_empty_double + | valid_double_empty + | valid_up_up + | valid_down_down + ) + pair_move = pair_branch & valid + proposed[pair_branch] = configs[pair_branch] + proposed[pair_move & valid_empty_double, i] = encoding.double + proposed[pair_move & valid_empty_double, j] = encoding.double + proposed[pair_move & valid_double_empty, i] = encoding.empty + proposed[pair_move & valid_double_empty, j] = encoding.empty + proposed[pair_move & valid_up_up, i] = encoding.down + proposed[pair_move & valid_up_up, j] = encoding.down + proposed[pair_move & valid_down_down, i] = encoding.up + proposed[pair_move & valid_down_down, j] = encoding.up + changed = torch.where(pair_branch, pair_move, changed) + if _return_move_codes: + return proposed, changed, move_codes + return proposed, changed + + +def _empty_proposal_stats(): + """Create the move-wise counters used by optional sampler diagnostics.""" + return { + name: { + "selected": 0, + "no_op": 0, + "proposed": 0, + "accepted": 0, + } + for name in _PROPOSAL_MOVE_NAMES + } + + +def _accumulate_proposal_stats(stats, move_codes, changed, accepted=None): + """Accumulate selected, no-op, proposed, and accepted move counts.""" + torch = _require_torch() + + def add_counts(mask, field): + if not torch.any(mask): + return + counts = torch.bincount( + move_codes[mask], + minlength=len(_PROPOSAL_MOVE_NAMES), + ).tolist() + for name, count in zip(_PROPOSAL_MOVE_NAMES, counts): + stats[name][field] += int(count) + + add_counts(torch.ones_like(changed, dtype=torch.bool), "selected") + add_counts(~changed, "no_op") + add_counts(changed, "proposed") + if accepted is not None: + add_counts(accepted, "accepted") + return stats + + +def _accumulate_accepted_proposal_stats(stats, move_codes, accepted): + """Add acceptances after a proposal's Metropolis decision.""" + torch = _require_torch() + if not torch.any(accepted): + return stats + counts = torch.bincount( + move_codes[accepted], + minlength=len(_PROPOSAL_MOVE_NAMES), + ).tolist() + for name, count in zip(_PROPOSAL_MOVE_NAMES, counts): + stats[name]["accepted"] += int(count) + return stats + + +def _merge_proposal_stats(total, update): + """Merge independently collected proposal diagnostics.""" + if update is None: + return total + if total is None: + total = _empty_proposal_stats() + for name in _PROPOSAL_MOVE_NAMES: + for field in total[name]: + total[name][field] += int(update[name][field]) + return total + + +_PROPOSAL_MIX_FAMILIES = { + "spinful": ("hopping_rate",), + "hubbard": ("hopping_rate",), + "spinful_exchange_hopping": ("hopping_rate",), + "spinful_u1": ("hopping_rate", "spin_flip_rate"), + "u1_spinful": ("hopping_rate", "spin_flip_rate"), + "spinful_total": ("hopping_rate", "spin_flip_rate"), + "spinful_total_exchange_hopping": ("hopping_rate", "spin_flip_rate"), + "spinful_z2": ( + "hopping_rate", + "spin_flip_rate", + "pair_toggle_rate", + ), + "z2_spinful": ("hopping_rate", "spin_flip_rate", "pair_toggle_rate"), + "spinful_parity": ( + "hopping_rate", + "spin_flip_rate", + "pair_toggle_rate", + ), + "spinful_parity_exchange_hopping": ( + "hopping_rate", + "spin_flip_rate", + "pair_toggle_rate", + ), + "spinful_z2z2": ("hopping_rate", "pair_toggle_rate"), + "z2z2_spinful": ("hopping_rate", "pair_toggle_rate"), + "spinful_resolved_parity": ("hopping_rate", "pair_toggle_rate"), +} + + +def _proposal_move_score(stats, names): + selected = sum(stats[name]["selected"] for name in names) + if selected == 0: + return None + accepted = sum(stats[name]["accepted"] for name in names) + return accepted / selected + + +def _adapt_proposal_mix_rate( + owner, + attribute, + candidate_moves, + reference_moves, + stats, + *, + adaptation_rate, + min_probability, + max_probability, +): + """Update one conditional move probability from whole-sweep statistics.""" + candidate_score = _proposal_move_score(stats, candidate_moves) + reference_score = _proposal_move_score(stats, reference_moves) + current = float(getattr(owner, attribute)) + if ( + candidate_score is None + or reference_score is None + or not 0.0 < current < 1.0 + ): + return False + + # The statistic is the fraction of selections which led to an accepted + # configuration change, so invalid/no-op branches are penalized too. + logit = math.log(current / (1.0 - current)) + logit += adaptation_rate * (candidate_score - reference_score) + if logit >= 0.0: + updated = 1.0 / (1.0 + math.exp(-logit)) + else: + exp_logit = math.exp(logit) + updated = exp_logit / (1.0 + exp_logit) + setattr( + owner, + attribute, + min(max(updated, min_probability), max_probability), + ) + return True + + +def _proposal_mix_rates(owner): + return { + "hopping_rate": float(owner.hopping_rate), + "spin_flip_rate": float(owner.spin_flip_rate), + "pair_toggle_rate": float(owner.pair_toggle_rate), + } + + +def _warmup_proposal_mix( + owner, + *, + n_sweeps, + adaptation_rate, + min_probability, + max_probability, + progress, +): + """Tune symmetric proposal weights between warm-up sweeps only.""" + n_sweeps = _check_positive_int("n_sweeps", n_sweeps) + adaptation_rate = float(adaptation_rate) + min_probability = float(min_probability) + max_probability = float(max_probability) + if not math.isfinite(adaptation_rate) or adaptation_rate <= 0.0: + raise ValueError("adaptation_rate must be a finite positive number.") + if not 0.0 <= min_probability < max_probability <= 1.0: + raise ValueError( + "Require 0 <= min_probability < max_probability <= 1." + ) + + supported_rates = _PROPOSAL_MIX_FAMILIES.get(str(owner.proposal), ()) + total_stats = _empty_proposal_stats() + history = [] + bar = _make_progress( + progress, + total=n_sweeps, + desc="Torch VMC proposal warm-up", + unit="sweep", + ) + try: + for sweep in range(1, n_sweeps + 1): + result = owner.sample_sweep( + n_sweeps=1, + track_proposal_stats=True, + ) + if result is None: + raise ValueError("Cannot tune a proposal mix on an empty graph.") + stats = result.proposal_stats + _merge_proposal_stats(total_stats, stats) + + if "hopping_rate" in supported_rates: + _adapt_proposal_mix_rate( + owner, + "hopping_rate", + ("hopping",), + ("exchange",), + stats, + adaptation_rate=adaptation_rate, + min_probability=min_probability, + max_probability=max_probability, + ) + if "spin_flip_rate" in supported_rates: + _adapt_proposal_mix_rate( + owner, + "spin_flip_rate", + ("spin_flip",), + ("exchange", "hopping"), + stats, + adaptation_rate=adaptation_rate, + min_probability=min_probability, + max_probability=max_probability, + ) + if "pair_toggle_rate" in supported_rates: + _adapt_proposal_mix_rate( + owner, + "pair_toggle_rate", + ("pair_toggle",), + ("exchange", "hopping", "spin_flip"), + stats, + adaptation_rate=adaptation_rate, + min_probability=min_probability, + max_probability=max_probability, + ) + history.append( + { + "sweep": sweep, + "rates": _proposal_mix_rates(owner), + "proposal_stats": stats, + } + ) + if bar is not None: + bar.update(1) + set_postfix = getattr(bar, "set_postfix", None) + if callable(set_postfix): + set_postfix(_proposal_mix_rates(owner)) + finally: + if bar is not None: + bar.close() + + summary = { + "n_sweeps": n_sweeps, + "rates": _proposal_mix_rates(owner), + "proposal_stats": total_stats, + "history": tuple(history), + } + owner.last_proposal_tuning = summary + return summary + + +def _safe_metropolis_ratio(proposed_amps, current_amps): + torch = _require_torch() + numerator = proposed_amps.abs().square() + denominator = current_amps.abs().square() + zero = torch.zeros_like(numerator) + inf = torch.full_like(numerator, float("inf")) + return torch.where( + denominator > 0, + numerator / denominator, + torch.where(numerator > 0, inf, zero), + ) + + +def _safe_metropolis_log_ratio( + proposed_log_abs, + current_log_abs, + *, + proposed_nonzero=None, + current_nonzero=None, +): + """Return clipped Metropolis ratios from log magnitudes.""" + torch = _require_torch() + log_ratio = 2.0 * (proposed_log_abs - current_log_abs) + # ``inf - inf`` can occur for user-provided log amplitudes. The explicit + # support masks below decide those zero-amplitude cases, so make the + # finite-ratio branch harmless instead of propagating NaNs into RNG tests. + log_ratio = torch.where( + torch.isnan(log_ratio), + torch.zeros_like(log_ratio), + log_ratio, + ) + ratio = torch.exp(torch.minimum(log_ratio, torch.zeros_like(log_ratio))) + if proposed_nonzero is None or current_nonzero is None: + return ratio + zero = torch.zeros_like(ratio) + one = torch.ones_like(ratio) + return torch.where( + current_nonzero, + torch.where(proposed_nonzero, ratio, zero), + torch.where(proposed_nonzero, one, zero), + ) + + +def metropolis_exchange_sweep( + configs, + amplitude_fn, + graph, + *, + current_amplitudes=None, + current_log_abs=None, + current_nonzero=None, + log_amplitude_fn=None, + proposal="spinful", + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + encoding=None, + generator=None, + chunk_size=None, + compile_kernels=False, + track_proposal_stats=False, +): + """Run one nearest-neighbor Metropolis sweep. + + ``amplitude_fn`` should accept a ``(batch, n_sites)`` torch integer tensor + and return a batch of amplitudes. The sampler evaluates only changed + proposals when possible. ``chunk_size`` caps proposal-amplitude batch size + without changing the Markov chain. Set ``track_proposal_stats=True`` to + retain move-wise selected, no-op, proposed, and accepted counts. + """ + torch = _require_torch() + configs = _as_long_matrix(configs).clone() + current = ( + _call_amplitude_fn(amplitude_fn, configs, chunk_size=chunk_size) + if current_amplitudes is None + else current_amplitudes + ) + current = torch.as_tensor(current, device=configs.device) + log_amplitude_fn = _resolve_log_amplitude_fn( + amplitude_fn, + log_amplitude_fn, + ) + if log_amplitude_fn is not None: + try: + if current_log_abs is None or current_nonzero is None: + current_phase, computed_log_abs = _call_log_amplitude_fn( + log_amplitude_fn, + configs, + chunk_size=chunk_size, + ) + if current_log_abs is None: + current_log_abs = computed_log_abs + if current_nonzero is None: + current_nonzero = current_phase.abs() > 0 + current_log_abs = torch.as_tensor( + current_log_abs, + dtype=torch.float64, + device=configs.device, + ).clone() + current_nonzero = torch.as_tensor( + current_nonzero, + dtype=torch.bool, + device=configs.device, + ).clone() + except (AttributeError, IndexError, KeyError, NotImplementedError, + RuntimeError, TypeError, ValueError): + # Some approximate sparse contractions expose ``forward_log`` but + # cannot represent every intermediate charge sector. Keep the + # raw-amplitude sampler path usable in that case. + log_amplitude_fn = None + current_log_abs = None + current_nonzero = None + n_proposed = 0 + n_accepted = 0 + proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + + for i, j in _iter_edges(graph): + if proposal in {"spin", "spin_exchange", "heisenberg"}: + proposal_result = propose_spin_exchange( + i, + j, + configs, + compile_kernels=compile_kernels, + _return_move_codes=track_proposal_stats, + ) + elif proposal in {"spinful", "hubbard", "spinful_exchange_hopping"}: + proposal_result = propose_spinful_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=track_proposal_stats, + ) + elif proposal in { + "spinful_u1", + "u1_spinful", + "spinful_total", + "spinful_total_exchange_hopping", + }: + proposal_result = propose_spinful_u1_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + spin_flip_rate=spin_flip_rate, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=track_proposal_stats, + ) + elif proposal in { + "spinful_z2", + "z2_spinful", + "spinful_parity", + "spinful_parity_exchange_hopping", + }: + proposal_result = propose_spinful_z2_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + spin_flip_rate=spin_flip_rate, + pair_toggle_rate=pair_toggle_rate, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=track_proposal_stats, + ) + elif proposal in { + "spinful_z2z2", + "z2z2_spinful", + "spinful_resolved_parity", + }: + proposal_result = propose_spinful_z2z2_exchange_or_hopping( + i, + j, + configs, + hopping_rate=hopping_rate, + pair_toggle_rate=pair_toggle_rate, + encoding=encoding, + generator=generator, + compile_kernels=compile_kernels, + _return_move_codes=track_proposal_stats, + ) + else: + raise ValueError( + "proposal must be 'spin', 'spinful_exchange_hopping', or " + "'spinful_u1', 'spinful_z2', or 'spinful_z2z2'." + ) + + if track_proposal_stats: + proposed, flags, move_codes = proposal_result + _accumulate_proposal_stats(proposal_stats, move_codes, flags) + else: + proposed, flags = proposal_result + + if not torch.any(flags): + continue + + n_changed = int(flags.sum().item()) + n_proposed += n_changed + proposed_amps = current.clone() + proposal_amplitude_fn = getattr( + amplitude_fn, + "proposal_amplitudes", + None, + ) + if callable(proposal_amplitude_fn): + proposed_amps[flags] = proposal_amplitude_fn( + configs[flags], + proposed[flags], + current[flags], + chunk_size=chunk_size, + ) + else: + proposed_amps[flags] = _call_amplitude_fn( + amplitude_fn, + proposed[flags], + chunk_size=chunk_size, + ) + if log_amplitude_fn is None: + ratio = _safe_metropolis_ratio(proposed_amps, current) + else: + try: + proposed_phase, proposed_log_abs_values = ( + _call_log_amplitude_fn( + log_amplitude_fn, + proposed[flags], + chunk_size=chunk_size, + ) + ) + proposed_log_abs = current_log_abs.clone() + proposed_log_abs[flags] = proposed_log_abs_values + proposed_nonzero = current_nonzero.clone() + proposed_nonzero[flags] = proposed_phase.abs() > 0 + ratio = _safe_metropolis_log_ratio( + proposed_log_abs, + current_log_abs, + proposed_nonzero=proposed_nonzero, + current_nonzero=current_nonzero, + ) + except (AttributeError, IndexError, KeyError, NotImplementedError, + RuntimeError, TypeError, ValueError): + log_amplitude_fn = None + current_log_abs = None + current_nonzero = None + ratio = _safe_metropolis_ratio(proposed_amps, current) + accept = flags & ( + torch.rand(configs.shape[0], device=configs.device, generator=generator) + < ratio + ) + if track_proposal_stats: + _accumulate_accepted_proposal_stats( + proposal_stats, + move_codes, + accept, + ) + + if torch.any(accept): + n_accept = int(accept.sum().item()) + n_accepted += n_accept + configs[accept] = proposed[accept] + current[accept] = proposed_amps[accept] + if log_amplitude_fn is not None: + current_log_abs[accept] = proposed_log_abs[accept] + current_nonzero[accept] = proposed_nonzero[accept] + + return TorchMetropolisResult( + configs=configs, + amplitudes=current, + n_proposed=n_proposed, + n_accepted=n_accepted, + log_abs_amplitudes=( + current_log_abs if log_amplitude_fn is not None else None + ), + nonzero_amplitudes=( + current_nonzero if log_amplitude_fn is not None else None + ), + proposal_stats=proposal_stats, + ) + + +class TorchMetropolisSampler: + """Stateful batched Metropolis sampler for torch amplitude models. + + The first configuration axis represents independent chains. Sampling + retains that axis and returns arrays shaped as + ``(n_samples_per_chain, n_chains, n_sites)``. ``sweep_size`` is measured + in the graph sweeps performed by :func:`metropolis_exchange_sweep`; the + ``n_thin`` spelling is accepted as a convenience alias. + """ + + def __init__( + self, + amplitude_fn, + graph, + configs, + *, + amplitudes=None, + n_chains=None, + proposal="spinful", + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + encoding=None, + chunk_size=None, + compile_kernels=False, + generator=None, + seed=None, + n_sites=None, + log_amplitude_fn=None, + log_abs_amplitudes=None, + nonzero_amplitudes=None, + ): + if generator is not None and seed is not None: + raise ValueError("Pass either generator=... or seed=..., not both.") + configs = _as_long_matrix(configs).clone() + if n_sites is not None: + n_sites = _check_positive_int("n_sites", n_sites) + if configs.shape[1] != n_sites: + raise ValueError( + f"n_sites={n_sites} does not match configs with " + f"{configs.shape[1]} sites." + ) + if n_chains is None: + n_chains = int(configs.shape[0]) + n_chains = _check_positive_int("n_chains", n_chains) + if configs.shape[0] == 1 and n_chains > 1: + configs = configs.expand(n_chains, -1).clone() + elif configs.shape[0] != n_chains: + raise ValueError( + "configs must contain exactly one initial configuration per " + f"chain: expected {n_chains}, got {configs.shape[0]}." + ) + + self.amplitude_fn = amplitude_fn + self.graph = graph + self.configs = configs + self.n_chains = n_chains + self.proposal = proposal + self.hopping_rate = float(hopping_rate) + self.spin_flip_rate = float(spin_flip_rate) + self.pair_toggle_rate = float(pair_toggle_rate) + self.encoding = encoding + self.chunk_size = _normalize_chunk_size(chunk_size) + self.compile_kernels = bool(compile_kernels) + self.last_proposal_stats = None + self.last_proposal_tuning = None + self.log_amplitude_fn = _resolve_log_amplitude_fn( + amplitude_fn, + log_amplitude_fn, + ) + self.generator = ( + _make_torch_generator(seed, device=configs.device) + if seed is not None + else generator + ) + if amplitudes is None: + self.refresh_amplitudes() + else: + amplitudes = _require_torch().as_tensor( + amplitudes, + device=configs.device, + ) + if amplitudes.numel() == 1 and n_chains > 1: + amplitudes = amplitudes.reshape(1).expand(n_chains).clone() + if amplitudes.shape != (n_chains,): + raise ValueError( + "amplitudes must have one value per chain, got " + f"shape {tuple(amplitudes.shape)}." + ) + self.amplitudes = amplitudes + self._refresh_log_amplitudes( + log_abs_amplitudes=log_abs_amplitudes, + nonzero_amplitudes=nonzero_amplitudes, + ) + + @property + def n_sites(self): + """Number of physical sites in each chain configuration.""" + return int(self.configs.shape[1]) + + def refresh_amplitudes(self): + """Recompute the amplitudes at the current chain positions.""" + with _require_torch().no_grad(): + self.amplitudes = _call_amplitude_fn( + self.amplitude_fn, + self.configs, + chunk_size=self.chunk_size, + ) + self._refresh_log_amplitudes() + return self.amplitudes + + def _refresh_log_amplitudes( + self, + *, + log_abs_amplitudes=None, + nonzero_amplitudes=None, + ): + if self.log_amplitude_fn is None: + self.log_abs_amplitudes = None + self.nonzero_amplitudes = None + return + if log_abs_amplitudes is not None and nonzero_amplitudes is not None: + torch = _require_torch() + self.log_abs_amplitudes = torch.as_tensor( + log_abs_amplitudes, + dtype=torch.float64, + device=self.configs.device, + ) + self.nonzero_amplitudes = torch.as_tensor( + nonzero_amplitudes, + dtype=torch.bool, + device=self.configs.device, + ) + if self.log_abs_amplitudes.shape != (self.n_chains,): + raise ValueError( + "log_abs_amplitudes must have one value per chain." + ) + if self.nonzero_amplitudes.shape != (self.n_chains,): + raise ValueError( + "nonzero_amplitudes must have one value per chain." + ) + return + try: + phase, self.log_abs_amplitudes = _call_log_amplitude_fn( + self.log_amplitude_fn, + self.configs, + chunk_size=self.chunk_size, + ) + self.nonzero_amplitudes = phase.abs() > 0 + except (AttributeError, IndexError, KeyError, NotImplementedError, + RuntimeError, TypeError, ValueError): + self.log_amplitude_fn = None + self.log_abs_amplitudes = None + self.nonzero_amplitudes = None + + def reset(self, configs=None, *, amplitudes=None): + """Reset chain positions, optionally supplying their amplitudes.""" + if configs is None: + raise ValueError("reset requires explicit configs.") + configs = _as_long_matrix(configs).clone() + if tuple(configs.shape) != tuple(self.configs.shape): + raise ValueError( + "reset configs must have shape " + f"{tuple(self.configs.shape)}, got {tuple(configs.shape)}." + ) + self.configs = configs + if amplitudes is None: + self.refresh_amplitudes() + else: + amplitudes = _require_torch().as_tensor( + amplitudes, + device=configs.device, + ) + if amplitudes.shape != (self.n_chains,): + raise ValueError("reset amplitudes must have one value per chain.") + self.amplitudes = amplitudes + self._refresh_log_amplitudes() + return self + + def sample_sweep(self, *, n_sweeps=1, track_proposal_stats=False): + """Advance all chains by one or more graph sweeps.""" + n_sweeps = _check_positive_int("n_sweeps", n_sweeps) + result = None + n_proposed = 0 + n_accepted = 0 + proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + with _require_torch().no_grad(): + for _ in range(n_sweeps): + result = metropolis_exchange_sweep( + self.configs, + self.amplitude_fn, + self.graph, + current_amplitudes=self.amplitudes, + current_log_abs=self.log_abs_amplitudes, + current_nonzero=self.nonzero_amplitudes, + log_amplitude_fn=( + self.log_amplitude_fn + if self.log_amplitude_fn is not None + else False + ), + 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, + generator=self.generator, + chunk_size=self.chunk_size, + compile_kernels=self.compile_kernels, + track_proposal_stats=track_proposal_stats, + ) + self.configs = result.configs + self.amplitudes = result.amplitudes + self.log_abs_amplitudes = result.log_abs_amplitudes + self.nonzero_amplitudes = result.nonzero_amplitudes + if result.log_abs_amplitudes is None: + self.log_amplitude_fn = None + n_proposed += result.n_proposed + n_accepted += result.n_accepted + if track_proposal_stats: + _merge_proposal_stats(proposal_stats, result.proposal_stats) + if result is not None: + result = replace( + result, + n_proposed=n_proposed, + n_accepted=n_accepted, + proposal_stats=proposal_stats, + ) + if track_proposal_stats: + self.last_proposal_stats = proposal_stats + return result + + def burn_in( + self, + n_sweeps=32, + *, + progress=False, + track_proposal_stats=False, + ): + """Equilibrate local walkers before fixed-kernel VMC work. + + This is the canonical convenience method for ordinary fixed-rate + burn-in. Use :meth:`warmup_proposal_mix` first when the local move + weights should be tuned; its adaptive samples are deliberately kept + separate from this fixed-kernel stage. + """ + n_sweeps = _check_positive_int("n_sweeps", n_sweeps) + if not progress: + sweep_kwargs = {"n_sweeps": n_sweeps} + if track_proposal_stats: + sweep_kwargs["track_proposal_stats"] = True + return self.sample_sweep(**sweep_kwargs) + + bar = _make_progress( + True, + total=n_sweeps, + desc="Torch VMC burn-in", + unit="sweep", + ) + result = None + n_proposed = 0 + n_accepted = 0 + proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + try: + for _ in range(n_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 + if track_proposal_stats: + _merge_proposal_stats(proposal_stats, result.proposal_stats) + bar.update(1) + _set_vmc_progress_postfix( + bar, + result, + n_sites=self.n_sites, + include_energy=False, + ) + finally: + bar.close() + + result = replace( + result, + n_proposed=n_proposed, + n_accepted=n_accepted, + proposal_stats=proposal_stats, + ) + if track_proposal_stats: + self.last_proposal_stats = proposal_stats + return result + + def warmup_proposal_mix( + self, + *, + n_sweeps=32, + adaptation_rate=1.0, + min_probability=0.05, + max_probability=0.95, + progress=False, + ): + """Tune move weights during warm-up, then leave them fixed. + + Each adaptation follows a completed graph sweep, never an individual + Metropolis transition. Discard these warm-up configurations before + collecting production samples; normal :meth:`sample` and + :meth:`sample_sweep` calls do not adapt rates. + """ + return _warmup_proposal_mix( + self, + n_sweeps=n_sweeps, + adaptation_rate=adaptation_rate, + min_probability=min_probability, + max_probability=max_probability, + progress=progress, + ) + + def sample( + self, + *, + n_samples=1024, + n_discard_per_chain=None, + n_discard=None, + sweep_size=None, + n_thin=None, + progress=False, + track_proposal_stats=False, + ): + """Discard and collect chain-preserving Metropolis samples. + + ``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. + """ + torch = _require_torch() + n_samples = _check_positive_int("n_samples", n_samples) + if n_discard_per_chain is not None and n_discard is not None: + raise ValueError( + "Pass either n_discard_per_chain=... or n_discard=..., not both." + ) + if sweep_size is not None and n_thin is not None: + raise ValueError("Pass either sweep_size=... or n_thin=..., not both.") + if n_discard_per_chain is None: + n_discard_per_chain = 32 if n_discard is None else n_discard + if sweep_size is None: + sweep_size = self.n_sites if n_thin is None else n_thin + n_discard_per_chain = _check_nonnegative_int( + "n_discard_per_chain", + n_discard_per_chain, + ) + sweep_size = _check_positive_int("sweep_size", sweep_size) + n_samples_per_chain = ( + n_samples + self.n_chains - 1 + ) // self.n_chains + total_sweeps = ( + n_discard_per_chain + n_samples_per_chain + ) * sweep_size + bar = _make_progress( + progress, + total=total_sweeps, + desc="Torch Metropolis", + ) + start = time.perf_counter() + n_proposed = 0 + n_accepted = 0 + proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + + def advance_one_sweep(): + nonlocal n_proposed, n_accepted + 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 + if track_proposal_stats: + _merge_proposal_stats(proposal_stats, result.proposal_stats) + if bar is not None: + bar.update(1) + + for _ in range(n_discard_per_chain * sweep_size): + advance_one_sweep() + + configs = [] + amplitudes = [] + log_abs_amplitudes = [] if self.log_abs_amplitudes is not None else None + for _ in range(n_samples_per_chain): + for _ in range(sweep_size): + advance_one_sweep() + configs.append(self.configs.clone()) + amplitudes.append(self.amplitudes.clone()) + if log_abs_amplitudes is not None: + if self.log_abs_amplitudes is None: + # A sparse/approximate contraction can expose a + # forward_log method that fails for one proposed charge + # sector. The sweep then falls back to raw amplitudes; + # discard the optional log cache rather than appending + # from a state that no longer has one. + log_abs_amplitudes = None + else: + log_abs_amplitudes.append(self.log_abs_amplitudes.clone()) + if bar is not None: + bar.close() + + configs = torch.stack(configs, dim=0) + amplitudes = torch.stack(amplitudes, dim=0) + if log_abs_amplitudes is not None: + log_abs_amplitudes = torch.stack(log_abs_amplitudes, dim=0) + actual_samples = int(configs.shape[0] * configs.shape[1]) + elapsed = time.perf_counter() - start + return TorchMCMCSamples( + configs=configs, + amplitudes=amplitudes, + n_samples=actual_samples, + n_samples_per_chain=int(configs.shape[0]), + n_chains=self.n_chains, + n_discard_per_chain=n_discard_per_chain, + sweep_size=sweep_size, + acceptance_rate=( + n_accepted / n_proposed if n_proposed else 0.0 + ), + n_proposed=n_proposed, + n_accepted=n_accepted, + elapsed_seconds=elapsed, + samples_per_second=( + actual_samples / elapsed if elapsed > 0 else float("inf") + ), + log_abs_amplitudes=log_abs_amplitudes, + proposal_stats=proposal_stats, + ) + + +class TorchBPMetropolisSampler(TorchMetropolisSampler): + """Independence Metropolis sampler driven by a BP proposal. + + ``proposal_sampler`` should return ``configs`` and BP proposal + probabilities in ``omegas``, as :class:`pepsy.sampling.PepsBpSampler` + does. Initial chains are drawn from that proposal, so the proposal + probability of every current chain is known. Later proposals are accepted + with the exact independence Metropolis-Hastings ratio. + + ``symmetry`` and ``sector`` optionally filter spinful fermion proposals. + This is important for approximate BP distributions, which can assign + probability to configurations outside a globally fixed charge sector. + """ + + def __init__( + self, + amplitude_fn, + graph, + proposal_sampler, + configs=None, + *, + amplitudes=None, + initial_log_q=None, + n_chains=None, + sample_kwargs=None, + symmetry=None, + sector=None, + encoding=None, + valid_config_fn=None, + amplitude_floor=0.0, + max_init_attempts=32, + chunk_size=None, + generator=None, + seed=None, + device=None, + log_amplitude_fn=None, + ): + torch = _require_torch() + if generator is not None and seed is not None: + raise ValueError("Pass either generator=... or seed=..., not both.") + if amplitude_floor < 0: + raise ValueError("amplitude_floor must be non-negative.") + max_init_attempts = _check_positive_int( + "max_init_attempts", + max_init_attempts, + ) + if configs is None: + if n_chains is None: + raise ValueError( + "n_chains is required when BP initializes the chains." + ) + n_chains = _check_positive_int("n_chains", n_chains) + else: + configs = _as_long_matrix(configs) + if n_chains is None: + n_chains = int(configs.shape[0]) + n_chains = _check_positive_int("n_chains", n_chains) + + self.amplitude_fn = amplitude_fn + self.proposal_sampler = proposal_sampler + self.proposal_sample_kwargs = dict(sample_kwargs or {}) + self.symmetry = None if symmetry is None else str(symmetry).upper() + self.sector = sector + self.fermion_encoding = encoding + self.valid_config_fn = valid_config_fn + self.amplitude_floor = float(amplitude_floor) + self.chunk_size = _normalize_chunk_size(chunk_size) + self._initial_log_fn = _resolve_log_amplitude_fn( + amplitude_fn, + log_amplitude_fn, + ) + if device is None: + try: + device = next(amplitude_fn.parameters()).device + except (AttributeError, StopIteration, TypeError): + device = None + self._proposal_device = ( + torch.device(device) if device is not None else None + ) + + if configs is None: + configs, amplitudes, initial_log_q = self._draw_initial_chains( + n_chains, + max_attempts=max_init_attempts, + ) + else: + configs = configs.clone() + if self._proposal_device is not None: + configs = configs.to(device=self._proposal_device) + if configs.shape[0] == 1 and n_chains > 1: + configs = configs.expand(n_chains, -1).clone() + elif configs.shape[0] != n_chains: + raise ValueError( + "configs must contain exactly one initial configuration " + f"per chain: expected {n_chains}, got {configs.shape[0]}." + ) + if initial_log_q is None: + raise ValueError( + "initial_log_q is required for explicit initial configs; " + "omit configs to initialize chains from BP." + ) + initial_log_q = torch.as_tensor( + initial_log_q, + dtype=torch.float64, + device=configs.device, + ) + if initial_log_q.shape != (n_chains,): + raise ValueError( + "initial_log_q must have one value per initial chain." + ) + if amplitudes is None: + with torch.no_grad(): + amplitudes = _call_amplitude_fn( + amplitude_fn, + configs, + chunk_size=self.chunk_size, + ) + + super().__init__( + amplitude_fn, + graph, + configs, + amplitudes=amplitudes, + n_chains=n_chains, + # The parent stores amplitude/log-amplitude state. Its local + # proposal is not used because this class overrides the sweep. + proposal="spin", + encoding=encoding, + chunk_size=self.chunk_size, + generator=generator, + seed=seed, + log_amplitude_fn=log_amplitude_fn, + ) + self.log_proposal_probabilities = torch.as_tensor( + initial_log_q, + dtype=torch.float64, + device=self.configs.device, + ).clone() + if self.log_proposal_probabilities.shape != (self.n_chains,): + raise ValueError( + "initial_log_q must have one value per initial chain." + ) + if not bool(torch.isfinite(self.log_proposal_probabilities).all()): + raise ValueError("Initial BP proposal probabilities must be positive and finite.") + self._validate_current_support() + + def _proposal_sample(self, n_samples): + """Draw configurations and decode their BP log probabilities.""" + kwargs = dict(self.proposal_sample_kwargs) + kwargs["samples"] = int(n_samples) + kwargs.setdefault("progbar", False) + try: + proposed = self.proposal_sampler.sample(**kwargs) + except TypeError: + kwargs.pop("progbar", None) + proposed = self.proposal_sampler.sample(**kwargs) + configs = _as_long_matrix(proposed.configs, name="proposal configs") + if self._proposal_device is not None: + configs = configs.to(device=self._proposal_device) + log_q = _proposal_log_probabilities( + proposed.omegas, + device=configs.device, + allow_zero=True, + ) + n_samples = int(n_samples) + if configs.shape[0] != n_samples or log_q.shape != (n_samples,): + raise ValueError( + "The BP proposal must return exactly one config and omega per " + f"requested sample ({n_samples})." + ) + return configs, log_q + + def _sector_mask(self, configs): + """Return the requested symmetry-sector mask for configurations.""" + torch = _require_torch() + if self.valid_config_fn is not None: + mask = torch.as_tensor( + self.valid_config_fn(configs), + dtype=torch.bool, + device=configs.device, + ) + if mask.shape != (configs.shape[0],): + raise ValueError( + "valid_config_fn must return one boolean per configuration." + ) + return mask + if self.symmetry is None or self.sector is None: + return torch.ones( + configs.shape[0], + dtype=torch.bool, + device=configs.device, + ) + n_up, n_down = count_spinful_particles( + configs, + encoding=self.fermion_encoding, + ) + if self.symmetry == "U1": + return n_up + n_down == int(self.sector) + if self.symmetry == "Z2": + return (n_up + n_down) % 2 == int(self.sector) % 2 + try: + sector = tuple(int(value) for value in self.sector) + except (TypeError, ValueError) as exc: + raise ValueError( + f"{self.symmetry} sectors must contain two integer charges." + ) from exc + if len(sector) != 2: + raise ValueError(f"{self.symmetry} sectors must contain two charges.") + if self.symmetry == "U1U1": + return (n_up == sector[0]) & (n_down == sector[1]) + if self.symmetry == "Z2Z2": + return ((n_up % 2) == sector[0] % 2) & ( + (n_down % 2) == sector[1] % 2 + ) + raise ValueError( + "symmetry must be one of U1, U1U1, Z2, or Z2Z2 when sector " + "filtering is enabled." + ) + + def _draw_initial_chains(self, n_chains, *, max_attempts): + """Draw nonzero, sector-valid initial chains from BP.""" + torch = _require_torch() + configs_out = [] + amplitudes_out = [] + log_q_out = [] + n_kept = 0 + for _ in range(max_attempts): + configs, log_q = self._proposal_sample(n_chains) + keep = self._sector_mask(configs) & torch.isfinite(log_q) + if not bool(torch.any(keep)): + continue + with torch.no_grad(): + amplitudes = _call_amplitude_fn( + self.amplitude_fn, + configs[keep], + chunk_size=self.chunk_size, + ) + if self._initial_log_fn is not None: + try: + phase, log_abs = _call_log_amplitude_fn( + self._initial_log_fn, + configs[keep], + chunk_size=self.chunk_size, + ) + support = ( + torch.isfinite(log_abs) + & (phase.abs() > 0) + & ( + log_abs + > ( + -torch.inf + if self.amplitude_floor == 0.0 + else float(torch.log(torch.tensor(self.amplitude_floor))) + ) + ) + ) + except ( + AttributeError, + IndexError, + KeyError, + NotImplementedError, + RuntimeError, + TypeError, + ValueError, + ): + self._initial_log_fn = None + support = ( + torch.isfinite(amplitudes.abs()) + & (amplitudes.abs() > self.amplitude_floor) + ) + else: + support = ( + torch.isfinite(amplitudes.abs()) + & (amplitudes.abs() > self.amplitude_floor) + ) + if not bool(torch.any(support)): + continue + configs_out.append(configs[keep][support]) + amplitudes_out.append(amplitudes[support]) + log_q_out.append(log_q[keep][support]) + n_kept += int(support.sum().item()) + if n_kept >= n_chains: + break + if n_kept < n_chains: + raise RuntimeError( + "Could not initialize enough nonzero BP configurations in the " + "requested Fermion sector. Check the BP encoding/sector or " + "increase max_init_attempts." + ) + return ( + torch.cat(configs_out, dim=0)[:n_chains], + torch.cat(amplitudes_out, dim=0)[:n_chains], + torch.cat(log_q_out, dim=0)[:n_chains], + ) + + def _validate_current_support(self): + """Reject undefined initial states rather than creating 0/0 ratios.""" + torch = _require_torch() + valid = self._sector_mask(self.configs) + if self.log_amplitude_fn is not None: + valid &= self.nonzero_amplitudes + valid &= torch.isfinite(self.log_abs_amplitudes) + if self.amplitude_floor: + valid &= self.log_abs_amplitudes > float( + torch.log(torch.tensor(self.amplitude_floor)) + ) + else: + valid &= torch.isfinite(self.amplitudes.abs()) + valid &= self.amplitudes.abs() > self.amplitude_floor + if not bool(torch.all(valid)): + raise ValueError( + "Initial BP Metropolis walkers must be finite, nonzero, and " + "inside the requested symmetry sector." + ) + + def reset(self, configs=None, *, amplitudes=None, log_proposal_probabilities=None): + """Reset chains and proposal probabilities together.""" + if log_proposal_probabilities is None: + raise ValueError( + "log_proposal_probabilities is required when resetting a BP " + "Metropolis sampler." + ) + super().reset(configs, amplitudes=amplitudes) + torch = _require_torch() + log_q = torch.as_tensor( + log_proposal_probabilities, + dtype=torch.float64, + device=self.configs.device, + ) + if log_q.shape != (self.n_chains,) or not bool(torch.isfinite(log_q).all()): + raise ValueError( + "log_proposal_probabilities must be finite with one value per chain." + ) + self.log_proposal_probabilities = log_q.clone() + self._validate_current_support() + return self + + def sample_sweep(self, *, n_sweeps=1): + """Advance all chains with BP independence proposals.""" + torch = _require_torch() + n_sweeps = _check_positive_int("n_sweeps", n_sweeps) + result = None + with torch.no_grad(): + for _ in range(n_sweeps): + proposed, proposed_log_q = self._proposal_sample(self.n_chains) + proposed = proposed.to(device=self.configs.device) + proposed_log_q = proposed_log_q.to(device=self.configs.device) + proposal_valid = self._sector_mask(proposed) & torch.isfinite( + proposed_log_q + ) + proposed_amplitudes = self.amplitudes.clone() + if bool(torch.any(proposal_valid)): + proposed_amplitudes[proposal_valid] = _call_amplitude_fn( + self.amplitude_fn, + proposed[proposal_valid], + chunk_size=self.chunk_size, + ) + + if self.log_amplitude_fn is not None and bool(torch.any(proposal_valid)): + try: + phase, proposed_log_valid = _call_log_amplitude_fn( + self.log_amplitude_fn, + proposed[proposal_valid], + chunk_size=self.chunk_size, + ) + proposed_log_abs = self.log_abs_amplitudes.clone() + proposed_nonzero = self.nonzero_amplitudes.clone() + proposed_log_abs[proposal_valid] = proposed_log_valid + proposed_nonzero[proposal_valid] = ( + (phase.abs() > 0) + & torch.isfinite(proposed_log_valid) + & ( + proposed_log_valid + > ( + -torch.inf + if self.amplitude_floor == 0.0 + else float(torch.log(torch.tensor(self.amplitude_floor))) + ) + ) + ) + except ( + AttributeError, + IndexError, + KeyError, + NotImplementedError, + RuntimeError, + TypeError, + ValueError, + ): + self.log_amplitude_fn = None + self.log_abs_amplitudes = None + self.nonzero_amplitudes = None + elif self.log_amplitude_fn is not None: + proposed_log_abs = self.log_abs_amplitudes.clone() + proposed_nonzero = torch.zeros_like(self.nonzero_amplitudes) + + if self.log_amplitude_fn is None: + current_abs = self.amplitudes.abs() + proposed_abs = proposed_amplitudes.abs() + current_log_abs = torch.where( + current_abs > 0, + current_abs.to(dtype=torch.float64).log(), + torch.full_like(current_abs, -torch.inf, dtype=torch.float64), + ) + proposed_log_abs = torch.where( + proposed_abs > 0, + proposed_abs.to(dtype=torch.float64).log(), + torch.full_like(proposed_abs, -torch.inf, dtype=torch.float64), + ) + current_nonzero = ( + torch.isfinite(current_abs) + & (current_abs > self.amplitude_floor) + ) + proposed_nonzero = ( + proposal_valid + & torch.isfinite(proposed_abs) + & (proposed_abs > self.amplitude_floor) + ) + else: + current_log_abs = self.log_abs_amplitudes + current_nonzero = self.nonzero_amplitudes + proposed_nonzero &= proposal_valid + + log_ratio = ( + 2.0 * (proposed_log_abs - current_log_abs) + + self.log_proposal_probabilities + - proposed_log_q + ) + log_ratio = torch.where( + torch.isnan(log_ratio), + torch.zeros_like(log_ratio), + log_ratio, + ) + log_ratio = torch.minimum(log_ratio, torch.zeros_like(log_ratio)) + uniform = torch.rand( + self.n_chains, + device=self.configs.device, + generator=self.generator, + ) + accept = ( + proposal_valid + & current_nonzero + & proposed_nonzero + & ( + torch.log( + uniform.clamp_min(torch.finfo(torch.float64).tiny) + ) + < log_ratio + ) + ) + n_accepted = int(accept.sum().item()) + self.configs[accept] = proposed[accept] + self.amplitudes[accept] = proposed_amplitudes[accept] + self.log_proposal_probabilities[accept] = proposed_log_q[accept] + if self.log_amplitude_fn is not None: + self.log_abs_amplitudes[accept] = proposed_log_abs[accept] + self.nonzero_amplitudes[accept] = proposed_nonzero[accept] + result = TorchMetropolisResult( + configs=self.configs, + amplitudes=self.amplitudes, + n_proposed=self.n_chains, + n_accepted=n_accepted, + log_abs_amplitudes=( + self.log_abs_amplitudes + if self.log_amplitude_fn is not None + else None + ), + nonzero_amplitudes=( + self.nonzero_amplitudes + if self.log_amplitude_fn is not None + else None + ), + ) + return result + + +def metropolis_local_sampler( + configs, + amplitude_fn, + graph, + *, + n_sites=None, + n_samples=1024, + n_chains=None, + n_discard_per_chain=None, + n_discard=None, + sweep_size=None, + n_thin=None, + proposal="spinful", + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + encoding=None, + chunk_size=None, + amplitudes=None, + generator=None, + seed=None, + progress=False, + log_amplitude_fn=None, + compile_kernels=False, +): + """Run a convenience batched Metropolis sampling call. + + ``n_sites`` is optional validation only; it is inferred from ``configs``. + Use :class:`TorchMetropolisSampler` when the chain state must be retained + for multiple sampling calls. + """ + sampler = TorchMetropolisSampler( + amplitude_fn, + graph, + configs, + amplitudes=amplitudes, + n_chains=n_chains, + proposal=proposal, + hopping_rate=hopping_rate, + spin_flip_rate=spin_flip_rate, + pair_toggle_rate=pair_toggle_rate, + encoding=encoding, + chunk_size=chunk_size, + generator=generator, + seed=seed, + n_sites=n_sites, + log_amplitude_fn=log_amplitude_fn, + compile_kernels=compile_kernels, + ) + return sampler.sample( + n_samples=n_samples, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + ) + + +def _mode_occupations(n_up, n_down, *, order): + torch = _require_torch() + if order in {"down-up", "du", "symmray"}: + return torch.stack((n_down, n_up), dim=-1).reshape(n_up.shape[0], -1) + if order in {"up-down", "ud", "netket"}: + return torch.stack((n_up, n_down), dim=-1).reshape(n_up.shape[0], -1) + raise ValueError("mode_order must be 'down-up' or 'up-down'.") + + +def _mode_index(site, spin, *, order): + if order in {"down-up", "du", "symmray"}: + offset = 1 if spin == "up" else 0 + elif order in {"up-down", "ud", "netket"}: + offset = 0 if spin == "up" else 1 + else: + raise ValueError("mode_order must be 'down-up' or 'up-down'.") + return 2 * site + offset + + +def spinful_fermi_hubbard_connections( + configs, + graph, + *, + t=1.0, + U=8.0, + encoding=None, + mode_order="down-up", +): + """Return batched spinful Fermi-Hubbard connected configurations. + + The local state encoding is configurable. Fermion signs use a site-major + mode order; ``mode_order='down-up'`` matches Symmray/vmc_torch convention. + """ + torch = _require_torch() + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding + configs = _as_long_matrix(configs) + batch, n_sites = configs.shape + device = configs.device + n_up, n_down = encoding.decode(configs) + modes = _mode_occupations(n_up, n_down, order=mode_order) + + all_etas = [] + all_coeffs = [] + all_bids = [] + + for edge in _iter_edges(graph): + i, j = edge + coeff = float(_edge_value(t, edge)) + if coeff == 0.0: + continue + for spin, occ in (("up", n_up), ("down", n_down)): + valid = occ[:, i] != occ[:, j] + if not torch.any(valid): + continue + idx = valid.nonzero(as_tuple=True)[0] + new_up = n_up[idx].clone() + new_down = n_down[idx].clone() + target = new_up if spin == "up" else new_down + tmp = target[:, i].clone() + target[:, i] = target[:, j] + target[:, j] = tmp + + p = _mode_index(i, spin, order=mode_order) + q = _mode_index(j, spin, order=mode_order) + if p > q: + p, q = q, p + between = modes[idx, p + 1:q].sum(dim=-1) % 2 + phase = 1.0 - 2.0 * between.to(torch.float64) + + all_etas.append(encoding.encode(new_up, new_down)) + all_coeffs.append(-coeff * phase) + all_bids.append(idx) + + for site in range(n_sites): + coeff = float(_site_value(U, site)) + if coeff == 0.0: + continue + valid = (n_up[:, site] == 1) & (n_down[:, site] == 1) + if not torch.any(valid): + continue + idx = valid.nonzero(as_tuple=True)[0] + all_etas.append(configs[idx].clone()) + all_coeffs.append(torch.full( + (idx.numel(),), + coeff, + dtype=torch.float64, + device=device, + )) + all_bids.append(idx) + + if not all_etas: + return _empty_connections(configs) + + return TorchConnections( + configs=torch.cat(all_etas, dim=0), + coeffs=torch.cat(all_coeffs, dim=0), + batch_ids=torch.cat(all_bids, dim=0), + ) + + +def heisenberg_connections(configs, graph, *, J=1.0): + """Return batched spin-1/2 Heisenberg connected configurations.""" + torch = _require_torch() + configs = _as_long_matrix(configs) + batch = configs.shape[0] + device = configs.device + batch_ids = torch.arange(batch, dtype=torch.long, device=device) + all_etas = [] + all_coeffs = [] + all_bids = [] + + for edge in _iter_edges(graph): + i, j = edge + coeff = float(_edge_value(J, edge)) + if coeff == 0.0: + continue + diff = configs[:, i] != configs[:, j] + if torch.any(diff): + idx = diff.nonzero(as_tuple=True)[0] + eta = configs[idx].clone() + tmp = eta[:, i].clone() + eta[:, i] = eta[:, j] + eta[:, j] = tmp + all_etas.append(eta) + all_coeffs.append(torch.full( + (idx.numel(),), + 0.5 * coeff, + dtype=torch.float64, + device=device, + )) + all_bids.append(idx) + + diag_sign = 1.0 - 2.0 * ((configs[:, i] - configs[:, j]).abs() % 2).to( + torch.float64 + ) + all_etas.append(configs.clone()) + all_coeffs.append(0.25 * coeff * diag_sign) + all_bids.append(batch_ids.clone()) + + if not all_etas: + return _empty_connections(configs) + + return TorchConnections( + configs=torch.cat(all_etas, dim=0), + coeffs=torch.cat(all_coeffs, dim=0), + batch_ids=torch.cat(all_bids, dim=0), + ) + + +def transverse_ising_connections(configs, graph, *, J=1.0, h=1.0): + """Return batched transverse-field Ising connected configurations.""" + torch = _require_torch() + configs = _as_long_matrix(configs) + batch, n_sites = configs.shape + device = configs.device + batch_ids = torch.arange(batch, dtype=torch.long, device=device) + all_etas = [] + all_coeffs = [] + all_bids = [] + + for edge in _iter_edges(graph): + i, j = edge + coeff = float(_edge_value(J, edge)) + if coeff == 0.0: + continue + diag_sign = 1.0 - 2.0 * ((configs[:, i] - configs[:, j]).abs() % 2).to( + torch.float64 + ) + all_etas.append(configs.clone()) + all_coeffs.append(0.25 * coeff * diag_sign) + all_bids.append(batch_ids.clone()) + + for site in range(n_sites): + coeff = float(_site_value(h, site)) + if coeff == 0.0: + continue + eta = configs.clone() + eta[:, site] = 1 - eta[:, site] + all_etas.append(eta) + all_coeffs.append(torch.full( + (batch,), + 0.5 * coeff, + dtype=torch.float64, + device=device, + )) + all_bids.append(batch_ids.clone()) + + if not all_etas: + return _empty_connections(configs) + + return TorchConnections( + configs=torch.cat(all_etas, dim=0), + coeffs=torch.cat(all_coeffs, dim=0), + batch_ids=torch.cat(all_bids, dim=0), + ) + + +def _connected_amplitudes_for_connections( + configs, + amplitudes, + connections, + amplitude_fn, + *, + chunk_size=None, + reuse_diagonal=True, +): + """Evaluate a deduplicated connected-configuration batch.""" + connected_amplitudes = getattr(amplitude_fn, "connected_amplitudes", None) + if callable(connected_amplitudes): + return connected_amplitudes( + configs, + amplitudes, + connections, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + ) + return _default_connected_amplitudes( + configs, + amplitudes, + connections, + amplitude_fn, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + ) + + +def _connection_contributions(connections, ratios): + """Multiply operator coefficients by amplitude ratios without dtype loss.""" + torch = _require_torch() + coeffs = connections.coeffs.to(device=ratios.device) + if torch.is_complex(coeffs) and not torch.is_complex(ratios): + # Fermionic operator builders commonly store mathematically real + # coefficients in a complex container. Retain the real VMC path in + # that exact-zero case, but never silently discard a physical phase. + if bool(torch.all(coeffs.imag == 0).item()): + coeffs = coeffs.real + dtype = torch.promote_types(coeffs.dtype, ratios.dtype) + return coeffs.to(dtype=dtype) * ratios.to(dtype=dtype) + + +def _local_energy_scatter_kernel(batch_ids, contributions, n_configs): + """Accumulate fixed-shape local-estimator contributions by walker.""" + torch = _require_torch() + energy = torch.zeros( + n_configs, + dtype=contributions.dtype, + device=contributions.device, + ) + return energy.index_add(0, batch_ids, contributions) + + +def _local_energy_from_connected_amplitudes( + configs, + amplitudes, + connections, + connected_amplitudes, + *, + compile_kernels=False, +): + """Accumulate one observable after its connected amplitudes are known.""" + torch = _require_torch() + if connections.configs.numel() == 0: + return torch.zeros( + configs.shape[0], + dtype=amplitudes.dtype, + device=configs.device, + ) + connected_amplitudes = torch.as_tensor( + connected_amplitudes, + device=configs.device, + ) + ratios = connected_amplitudes / amplitudes[connections.batch_ids] + contrib = _connection_contributions(connections, ratios) + return _run_cheap_torch_kernel( + "local-energy-scatter", + _local_energy_scatter_kernel, + connections.batch_ids, + contrib, + configs.shape[0], + compile_kernels=compile_kernels, + ) + + +def _connected_amplitudes_with_target_dedup( + configs, + amplitudes, + connections, + amplitude_fn, + *, + chunk_size=None, + reuse_diagonal=True, + deduplicate_targets=False, +): + """Evaluate connected amplitudes, optionally sharing target rows globally.""" + torch = _require_torch() + if not deduplicate_targets or connections.configs.shape[0] <= 1: + return _connected_amplitudes_for_connections( + configs, + amplitudes, + connections, + amplitude_fn, + chunk_size=chunk_size, + 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( + 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. + order = torch.argsort(target_inverse) + sorted_inverse = target_inverse[order] + first = torch.ones( + sorted_inverse.shape[0], + dtype=torch.bool, + device=sorted_inverse.device, + ) + 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 = _connected_amplitudes_for_connections( + configs, + amplitudes, + unique_connections, + amplitude_fn, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + ) + return unique_amplitudes[target_inverse] + + +def _local_energies_from_connection_map( + configs, + amplitudes, + connection_map, + amplitude_fn, + *, + chunk_size=None, + reuse_diagonal=True, + deduplicate=True, + deduplicate_targets=False, + compile_kernels=False, +): + """Evaluate several observables while sharing connected amplitudes. + + Connections are coalesced within each observable first, preserving their + individual operator coefficients. Their ``(walker, configuration)`` + targets are then merged across observables, so energy and a correlator can + reuse both ordinary amplitudes and PEPS boundary environments. When + ``deduplicate_targets=True``, identical target configurations are also + merged across parent walkers before connected-amplitude evaluation. + """ + torch = _require_torch() + configs = _as_long_matrix(configs) + amplitudes = torch.as_tensor(amplitudes, device=configs.device) + items = tuple(connection_map.items()) + if not items: + raise ValueError("connection_map must contain at least one observable.") + + prepared = {} + keys = [] + lengths = [] + for name, connections in items: + if deduplicate: + connections = _coalesce_connections( + connections, + device=configs.device, + compile_kernels=compile_kernels, + ) + prepared[name] = connections + length = int(connections.configs.shape[0]) + lengths.append(length) + if length: + keys.append(_run_cheap_torch_kernel( + "connection-key-rows", + _connection_key_rows, + connections.batch_ids, + connections.configs, + compile_kernels=compile_kernels, + )) + + if not keys: + return { + name: torch.zeros( + configs.shape[0], + dtype=amplitudes.dtype, + device=configs.device, + ) + for name, _ in items + } + + unique_keys, inverse = torch.unique( + torch.cat(keys, dim=0), + dim=0, + return_inverse=True, + ) + shared_connections = TorchConnections( + configs=unique_keys[:, 1:], + coeffs=torch.ones( + unique_keys.shape[0], + dtype=amplitudes.dtype, + device=configs.device, + ), + batch_ids=unique_keys[:, 0], + ) + shared_amplitudes = _connected_amplitudes_with_target_dedup( + configs, + amplitudes, + shared_connections, + amplitude_fn, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + deduplicate_targets=deduplicate_targets, + ) + + results = {} + offset = 0 + for (name, _), length in zip(items, lengths): + connections = prepared[name] + if length: + result_amplitudes = shared_amplitudes[inverse[offset:offset + length]] + results[name] = _local_energy_from_connected_amplitudes( + configs, + amplitudes, + connections, + result_amplitudes, + compile_kernels=compile_kernels, + ) + offset += length + else: + results[name] = torch.zeros( + configs.shape[0], + dtype=amplitudes.dtype, + device=configs.device, + ) + return results + + +def local_energy_from_connections( + configs, + amplitudes, + connections, + amplitude_fn, + *, + chunk_size=None, + reuse_diagonal=True, + deduplicate=True, + deduplicate_targets=False, + compile_kernels=False, +): + """Accumulate local energies from connected configs and amplitudes. + + If ``amplitude_fn`` exposes ``connected_amplitudes(...)`` that method is + used. Otherwise diagonal connections can reuse the supplied parent + amplitudes and off-diagonal amplitudes are evaluated in optional chunks. + By default, duplicate ``(walker, configuration)`` connections are + coalesced. Set ``deduplicate_targets=True`` to share identical target + configurations across parent walkers as well. Set ``deduplicate=False`` + for compatibility diagnostics. + """ + torch = _require_torch() + configs = _as_long_matrix(configs) + amplitudes = torch.as_tensor(amplitudes, device=configs.device) + if deduplicate: + connections = _coalesce_connections( + connections, + device=configs.device, + compile_kernels=compile_kernels, + ) + if connections.configs.numel() == 0: + return torch.zeros( + configs.shape[0], + dtype=amplitudes.dtype, + device=configs.device, + ) + + conn_amps = _connected_amplitudes_with_target_dedup( + configs, + amplitudes, + connections, + amplitude_fn, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + deduplicate_targets=deduplicate_targets, + ) + return _local_energy_from_connected_amplitudes( + configs, + amplitudes, + connections, + conn_amps, + compile_kernels=compile_kernels, + ) + + +def _energy_mean_and_variance(local_energies): + energy_mean = local_energies.mean() + centered = local_energies - energy_mean + variance = centered.abs().square().mean() + return energy_mean, variance.real + + +def _flat_sample_values(values, *, n_steps, n_chains, device, name): + """Validate scalar per-sample data and return it as one flat tensor.""" + torch = _require_torch() + values = torch.as_tensor(values, device=device) + expected_shape = (n_steps, n_chains) + n_samples = n_steps * n_chains + if tuple(values.shape) == expected_shape: + return values.reshape(-1) + if values.ndim == 1 and values.shape[0] == n_samples: + return values + raise ValueError( + f"{name} must have shape {expected_shape} or ({n_samples},), got " + f"{tuple(values.shape)}." + ) + + +def _normalized_sample_weights(weights, *, n_steps, n_chains, device): + """Return finite, non-negative supplied sample weights normalized to one.""" + torch = _require_torch() + weights = _flat_sample_values( + weights, + n_steps=n_steps, + n_chains=n_chains, + device=device, + name="weights", + ) + if torch.is_complex(weights): + raise ValueError("weights must be real, finite, and non-negative.") + if not torch.is_floating_point(weights): + weights = weights.to(torch.float64) + if not bool(torch.isfinite(weights).all()) or bool(torch.any(weights < 0)): + raise ValueError("weights must be real, finite, and non-negative.") + total = weights.sum() + if not bool(torch.isfinite(total)) or bool(total <= 0): + raise ValueError("weights must have a positive finite sum.") + # Sampling/importance probabilities are estimator data, never a + # differentiable model output. Detaching also keeps result diagnostics + # safe to convert to NumPy after an optimization step. + return (weights / total).detach() + + +def _importance_weights_from_log_probs( + amplitudes, + proposal_log_probs, + *, + n_steps, + n_chains, +): + """Return self-normalized ``|psi|**2 / q`` weights for a sample batch.""" + torch = _require_torch() + amplitudes = torch.as_tensor(amplitudes).reshape(-1) + log_q = _flat_sample_values( + proposal_log_probs, + n_steps=n_steps, + n_chains=n_chains, + device=amplitudes.device, + name="proposal_log_probs", + ) + if torch.is_complex(log_q): + raise ValueError("proposal_log_probs must be real and finite.") + if not torch.is_floating_point(log_q): + log_q = log_q.to(torch.float64) + if not bool(torch.isfinite(log_q).all()): + raise ValueError("proposal_log_probs must be real and finite.") + amplitude_abs = amplitudes.abs() + log_weights = 2.0 * amplitude_abs.log() - log_q + valid = torch.isfinite(log_weights) + if not bool(torch.any(valid)): + raise ValueError( + "The supplied proposal batch has no configuration with finite " + "non-zero model amplitude." + ) + max_log_weight = log_weights[valid].max() + weights = torch.where( + valid, + torch.exp(log_weights - max_log_weight), + torch.zeros_like(log_weights), + ) + return _normalized_sample_weights( + weights, + n_steps=n_steps, + n_chains=n_chains, + device=amplitudes.device, + ) + + +def _weighted_energy_statistics(local_energies, weights): + """Return mean, variance, standard error, and ESS for normalized weights.""" + torch = _require_torch() + local_energies = torch.as_tensor(local_energies).reshape(-1) + weights = torch.as_tensor(weights, device=local_energies.device).reshape(-1) + if local_energies.shape[0] != weights.shape[0]: + raise ValueError("weights must have one entry per local-energy sample.") + energy_mean = (weights.to(local_energies.dtype) * local_energies).sum() + energy_variance = ( + weights * (local_energies - energy_mean).abs().square() + ).sum().real + effective_sample_size = 1.0 / weights.square().sum() + energy_stderr = torch.sqrt(energy_variance / effective_sample_size) + energy_stderr_naive = torch.sqrt( + energy_variance / max(int(local_energies.numel()), 1) + ) + return ( + energy_mean, + energy_variance, + energy_stderr, + energy_stderr_naive, + effective_sample_size, + ) + + +def torch_chain_diagnostics(values, *, max_lag=None): + """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. + """ + torch = _require_torch() + values = torch.as_tensor(values) + if values.ndim != 2: + raise ValueError( + "values must have shape (n_samples_per_chain, n_chains)." + ) + n_steps, n_chains = (int(value) for value in values.shape) + if n_steps < 2: + raise ValueError("At least two samples per chain are required.") + if n_chains < 2: + raise ValueError("At least two chains are required.") + if not torch.is_floating_point(values) and not torch.is_complex(values): + values = values.to(torch.float64) + values = values.real if values.is_complex() else values + if values.dtype != torch.float64: + 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)) + + if max_lag is None: + max_lag = n_steps - 1 + else: + max_lag = _check_nonnegative_int("max_lag", max_lag) + max_lag = min(max_lag, n_steps - 1) + + centered = values - chain_means + fft_size = 1 << (2 * n_steps - 1).bit_length() + spectrum = torch.fft.rfft(centered, n=fft_size, dim=0) + autocovariance = torch.fft.irfft( + spectrum.conj() * spectrum, + n=fft_size, + dim=0, + )[:n_steps] + variances = autocovariance[0].real + normalized = torch.where( + variances > 0, + autocovariance.real / variances, + torch.zeros_like(autocovariance.real), + ) + rho = normalized.mean(dim=1) + tau = torch.ones((), dtype=values.dtype, device=values.device) + for lag in range(1, max_lag + 1): + if bool(rho[lag] <= 0): + break + tau = tau + 2 * rho[lag] + tau = torch.clamp(tau, min=1.0) + total_samples = n_steps * n_chains + effective_sample_size = torch.as_tensor( + total_samples, + dtype=values.dtype, + device=values.device, + ) / tau + return TorchChainDiagnostics( + r_hat=r_hat, + integrated_autocorrelation_time=tau, + effective_sample_size=effective_sample_size, + n_samples_per_chain=n_steps, + n_chains=n_chains, + ) + + +def _observable_statistics(chain_values): + """Compute an observable estimate with an autocorrelation-aware error.""" + torch = _require_torch() + chain_values = torch.as_tensor(chain_values) + if chain_values.ndim != 2: + raise ValueError( + "chain_values must have shape (n_samples_per_chain, n_chains)." + ) + local_values = chain_values.reshape(-1) + energy_mean, energy_variance = _energy_mean_and_variance(local_values) + n_samples = int(local_values.numel()) + naive_stderr = torch.sqrt(energy_variance / max(n_samples, 1)) + + chain_diagnostics = None + if chain_values.shape[0] >= 2 and chain_values.shape[1] >= 2: + chain_diagnostics = torch_chain_diagnostics(chain_values) + effective_sample_size = chain_diagnostics.effective_sample_size + else: + effective_sample_size = torch.as_tensor( + n_samples, + dtype=energy_variance.dtype, + device=energy_variance.device, + ) + autocorrelation_stderr = torch.sqrt( + energy_variance / torch.clamp(effective_sample_size, min=1.0) + ) + return ( + energy_mean, + energy_variance, + autocorrelation_stderr, + naive_stderr, + effective_sample_size, + chain_diagnostics, + ) + + +def _adaptive_measurement_options( + target_effective_sample_size, + *, + max_measurements, + min_measurements, + ess_check_interval, + rhat_threshold, + auto_thin, +): + """Validate optional ESS-targeted measurement controls.""" + if target_effective_sample_size is None: + return None + try: + target_effective_sample_size = float(target_effective_sample_size) + except (TypeError, ValueError) as exc: + raise ValueError( + "target_effective_sample_size must be a positive finite number." + ) from exc + if not math.isfinite(target_effective_sample_size) or target_effective_sample_size <= 0: + raise ValueError( + "target_effective_sample_size must be a positive finite number." + ) + min_measurements = _check_positive_int( + "min_measurements", + min_measurements, + ) + ess_check_interval = _check_positive_int( + "ess_check_interval", + ess_check_interval, + ) + if min_measurements < 2: + raise ValueError( + "min_measurements must be at least 2 for chain diagnostics." + ) + if min_measurements > max_measurements: + raise ValueError( + "min_measurements cannot exceed n_measurements when targeting ESS." + ) + 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.") + return { + "target_effective_sample_size": target_effective_sample_size, + "min_measurements": min_measurements, + "ess_check_interval": ess_check_interval, + "rhat_threshold": rhat_threshold, + "auto_thin": bool(auto_thin), + } + + +def _diagnostics_meet_target(diagnostics, options): + """Check ESS and optional R-hat stopping conditions.""" + torch = _require_torch() + if diagnostics is None: + return False + if not bool( + diagnostics.effective_sample_size + >= options["target_effective_sample_size"] + ): + return False + rhat_threshold = options["rhat_threshold"] + if rhat_threshold is None: + return True + return bool( + torch.isfinite(diagnostics.r_hat) + & (diagnostics.r_hat <= rhat_threshold) + ) + + +def _adaptive_thinning_interval(diagnostics, baseline): + """Choose a conservative next measurement spacing from chain mixing.""" + if diagnostics is None: + return baseline + tau = float(diagnostics.integrated_autocorrelation_time.detach().cpu()) + if not math.isfinite(tau): + return baseline + return max(baseline, int(math.ceil(tau))) + + +def _resolve_connection_fn(connection_fn): + if callable(connection_fn): + return None, connection_fn + key = str(connection_fn).replace("-", "_").lower() + aliases = { + "fermi_hubbard": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), + "fh": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), + "hubbard": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), + "spinful": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), + "spinful_fermi_hubbard": ( + "spinful_fermi_hubbard", + spinful_fermi_hubbard_connections, + ), + "heisenberg": ("heisenberg", heisenberg_connections), + "heis": ("heisenberg", heisenberg_connections), + "transverse_ising": ("transverse_ising", transverse_ising_connections), + "tfim": ("transverse_ising", transverse_ising_connections), + "ising": ("transverse_ising", transverse_ising_connections), + } + try: + return aliases[key] + except KeyError as exc: + allowed = ", ".join(sorted(aliases)) + raise ValueError( + f"Unknown torch VMC connection_fn {connection_fn!r}. " + f"Expected a callable or one of: {allowed}." + ) from exc + + +@dataclass(frozen=True) +class TorchVMCStepResult: + """Result of one :class:`TorchVMCDriver` step.""" + + configs: Any + amplitudes: Any + local_energies: Any + energy_mean: Any + energy_variance: Any + acceptance_rate: float + n_proposed: int + n_accepted: int + sr: Any = None + profile: Any = None + proposal_stats: Any = None + importance_weights: Any = None + effective_sample_size: Any = None + sample_source: str = "metropolis" + + +@dataclass(frozen=True) +class TorchVMCEnergyEstimate: + """Observable estimate and sampling diagnostics from a torch VMC run. + + ``chain_diagnostics`` is populated when the estimate retained at least + two samples from each of at least two chains. + """ + + configs: Any + amplitudes: Any + local_energies: Any + energy_mean: Any + energy_variance: Any + energy_stderr: Any + acceptance_rate: float + n_proposed: int + n_accepted: int + n_samples: int + n_measurements: int + elapsed_seconds: float + samples_per_second: float + chain_diagnostics: Any = None + profile: Any = None + energy_stderr_naive: Any = None + effective_sample_size: Any = None + importance_weights: Any = None + proposal_log_probs: Any = None + + +@dataclass(frozen=True) +class TorchVMCImportanceEstimate: + """Energy estimate from an external proposal distribution.""" + + configs: Any + amplitudes: Any + local_energies: Any + weights: Any + energy_mean: Any + energy_variance: Any + energy_stderr: Any + effective_sample_size: Any + n_samples: int + n_valid: int + elapsed_seconds: float + samples_per_second: float + + +def _make_progress(progress, *, total, desc, unit=None): + """Create an optional tqdm progress iterator without making tqdm required.""" + if not progress: + return None + try: + from tqdm.auto import tqdm + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "progress=True requires optional dependency 'tqdm'." + ) from exc + kwargs = {"total": total, "desc": desc, "dynamic_ncols": True} + if unit is not None: + kwargs["unit"] = unit + return tqdm(**kwargs) + + +def _proposal_no_op_rate(proposal_stats): + """Return the selected-move no-op fraction for optional diagnostics.""" + if not proposal_stats: + return None + selected = sum(move["selected"] for move in proposal_stats.values()) + if selected == 0: + return None + return sum(move["no_op"] for move in proposal_stats.values()) / selected + + +def _progress_scalar(value): + """Convert a scalar tensor to a display-only Python float.""" + try: + value = value.detach() + is_complex = getattr(value, "is_complex", False) + if callable(is_complex): + is_complex = is_complex() + if is_complex: + value = value.real + return float(value.item()) + except (AttributeError, TypeError, ValueError): + 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.""" + if bar is None: + return + 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: + postfix["E/site"] = ( + f"{_progress_scalar(result.energy_mean) / n_sites:+.6f}" + ) + sr_result = getattr(result, "sr", None) + if sr_result is not None: + solver = sr_result.info.get("solver") + if solver is not None: + postfix["SR"] = solver + 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 = {} + for name, attribute in ( + ("connected", "last_connected_reuse_stats"), + ("proposal", "last_proposal_cache_stats"), + ("amplitude", "last_amplitude_cache_stats"), + ): + value = getattr(model, attribute, None) + if value is not None: + snapshot[name] = dict(value) + if hasattr(model, "cutoff_fallbacks"): + snapshot["cutoff_fallbacks"] = int(model.cutoff_fallbacks) + return snapshot + + +def _accumulate_cache_profile(total, snapshot): + """Accumulate per-call cache counters without retaining every sample.""" + for name, value in snapshot.items(): + if isinstance(value, dict): + destination = total.setdefault(name, {}) + for key, count in value.items(): + if isinstance(count, Integral): + destination[key] = destination.get(key, 0) + int(count) + elif isinstance(value, Integral): + total[name] = int(value) + return total + + +def _proposal_log_probabilities(omegas, *, device, allow_zero=False): + """Decode ``PepsBpSampler`` mantissa/exponent proposal probabilities.""" + torch = _require_torch() + if not isinstance(omegas, (tuple, list)) or len(omegas) != 2: + raise ValueError( + "proposal samples must expose omegas as (mantissas, exponents)." + ) + mantissas = torch.as_tensor(omegas[0], dtype=torch.float64, device=device) + exponents = torch.as_tensor(omegas[1], dtype=torch.float64, device=device) + if mantissas.ndim != 1 or exponents.shape != mantissas.shape: + raise ValueError("proposal mantissas and exponents must be one-dimensional.") + if torch.any(mantissas < 0): + raise ValueError("proposal probabilities cannot have negative mantissas.") + if not allow_zero and torch.any(mantissas <= 0): + raise ValueError("proposal probabilities must have positive mantissas.") + positive = mantissas > 0 + log_prob = torch.where( + positive, + mantissas.log() + exponents * torch.log( + torch.as_tensor(10.0, dtype=torch.float64, device=device) + ), + torch.full_like(mantissas, -torch.inf), + ) + return log_prob + + +class TorchVMCDriver: + """Small PyTorch-native VMC loop around Pepsy's torch kernels. + + The driver keeps walker configurations and amplitudes in sync, runs + Metropolis exchange/hopping sweeps, evaluates local energies with optional + chunking/diagonal reuse, and can apply one SR/minSR update per step. + """ + + def __init__( + self, + model, + graph, + configs, + connection_fn=None, + *, + terms=None, + site_order=None, + connection_kwargs=None, + term_constant=0.0, + amplitudes=None, + proposal="spinful", + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + encoding=None, + chunk_size=None, + generator=None, + compile_kernels=False, + log_amplitude_fn=None, + ): + self.model = model + self.graph = graph + self.configs = _as_long_matrix(configs) + from .api import CompiledOperatorSum, OperatorSum + if isinstance(terms, CompiledOperatorSum): + if terms.backend != "torch": + raise ValueError( + f"Compiled terms target backend {terms.backend!r}, not 'torch'." + ) + term_constant = terms.constant + terms = terms.terms + elif isinstance(terms, OperatorSum): + compiled = compile_operator_sum_torch(terms) + term_constant = compiled.constant + terms = compiled.terms + self.term_constant = term_constant + if terms is not None and connection_fn is not None: + raise ValueError( + "Pass either terms=... or connection_fn=..., not both." + ) + self.terms = terms + if terms is not None: + self.connection_name = "terms" + if site_order is None and hasattr(graph, "Lx") and hasattr(graph, "Ly"): + site_order = tuple( + (x, y) + for x in range(int(graph.Lx)) + for y in range(int(graph.Ly)) + ) + self.site_order = None if site_order is None else tuple(site_order) + self.connection_fn = _driver_terms_connections + self.connection_kwargs = { + "terms": terms, + "site_order": self.site_order, + "constant": self.term_constant, + } + else: + if connection_fn is None: + connection_fn = "spinful_fermi_hubbard" + self.connection_name, self.connection_fn = _resolve_connection_fn( + connection_fn + ) + self.site_order = None if site_order is None else tuple(site_order) + self.connection_kwargs = ( + {} if connection_kwargs is None else dict(connection_kwargs) + ) + self.proposal = proposal + self.hopping_rate = float(hopping_rate) + self.spin_flip_rate = float(spin_flip_rate) + self.pair_toggle_rate = float(pair_toggle_rate) + self.encoding = encoding + self.chunk_size = _normalize_chunk_size(chunk_size) + self.compile_kernels = bool(compile_kernels) + self.last_proposal_stats = None + self.last_proposal_tuning = None + self.generator = generator + self.log_amplitude_fn = _resolve_log_amplitude_fn( + self.model, + log_amplitude_fn, + ) + self.log_abs_amplitudes = None + self.nonzero_amplitudes = None + self._sr_step = 0 + self._sr_previous_direction = None + + if ( + self.connection_name == "spinful_fermi_hubbard" + and encoding is not None + and "encoding" not in self.connection_kwargs + ): + self.connection_kwargs["encoding"] = encoding + + if amplitudes is None: + self.refresh_amplitudes() + else: + self.amplitudes = _require_torch().as_tensor( + amplitudes, + device=self.configs.device, + ) + self._refresh_log_amplitudes() + + @property + def n_walkers(self): + """Number of active walkers.""" + return int(self.configs.shape[0]) + + @property + def n_sites(self): + """Number of sites in each walker configuration.""" + return int(self.configs.shape[1]) + + @property + def sr_step(self): + """Number of completed SR updates, used by shift schedules.""" + return self._sr_step + + def reset_sr_state(self): + """Forget SR momentum and restart callable shift schedules at zero.""" + self._sr_step = 0 + self._sr_previous_direction = None + return self + + def refresh_amplitudes(self): + """Recompute current walker amplitudes from the current model.""" + clear_cache = getattr(self.model, "clear_boundary_cache", None) + if callable(clear_cache): + clear_cache() + with _require_torch().no_grad(): + self.amplitudes = _call_amplitude_fn( + self.model, + self.configs, + chunk_size=self.chunk_size, + ) + self._refresh_log_amplitudes() + return self.amplitudes + + def _refresh_log_amplitudes(self): + if self.log_amplitude_fn is None: + self.log_abs_amplitudes = None + self.nonzero_amplitudes = None + return + try: + phase, self.log_abs_amplitudes = _call_log_amplitude_fn( + self.log_amplitude_fn, + self.configs, + chunk_size=self.chunk_size, + ) + self.nonzero_amplitudes = phase.abs() > 0 + except (AttributeError, IndexError, KeyError, NotImplementedError, + RuntimeError, TypeError, ValueError): + self.log_amplitude_fn = None + self.log_abs_amplitudes = None + self.nonzero_amplitudes = None + + def make_sampler( + self, + *, + configs=None, + amplitudes=None, + n_chains=None, + seed=None, + sampler_seed=None, + proposal=None, + chunk_size=None, + ): + """Create a stateful sampler initialized from the current driver.""" + if seed is not None and sampler_seed is not None: + raise ValueError("Pass either seed=... or sampler_seed=..., not both.") + if configs is None: + configs = self.configs + if amplitudes is None: + amplitudes = self.amplitudes + return TorchMetropolisSampler( + self.model, + self.graph, + configs, + amplitudes=amplitudes, + n_chains=n_chains, + proposal=self.proposal if proposal is None else 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 + if chunk_size is None + else _normalize_chunk_size(chunk_size) + ), + log_amplitude_fn=( + self.log_amplitude_fn + if self.log_amplitude_fn is not None + else False + ), + log_abs_amplitudes=( + self.log_abs_amplitudes + if configs is self.configs + else None + ), + nonzero_amplitudes=( + self.nonzero_amplitudes + if configs is self.configs + else None + ), + compile_kernels=self.compile_kernels, + generator=( + self.generator + if seed is None and sampler_seed is None + else None + ), + seed=seed if seed is not None else sampler_seed, + ) + + def make_bp_sampler( + self, + proposal_sampler=None, + *, + n_chains=None, + sample_kwargs=None, + symmetry=None, + sector=None, + encoding=None, + amplitude_floor=0.0, + max_init_attempts=32, + seed=None, + sampler_seed=None, + chunk_size=None, + ): + """Create a BP independence sampler for this amplitude model. + + The base driver requires an explicit ``proposal_sampler``. The + :class:`TorchFermionVMC` specialization creates a compatible + :class:`pepsy.sampling.PepsBpSampler` automatically from its PEPS. + """ + if proposal_sampler is None: + raise ValueError( + "proposal_sampler is required for TorchVMCDriver; use " + "TorchFermionVMC to infer PepsBpSampler from a PEPS." + ) + if seed is not None and sampler_seed is not None: + raise ValueError("Pass either seed=... or sampler_seed=..., not both.") + n_chains = self.n_walkers if n_chains is None else n_chains + return TorchBPMetropolisSampler( + self.model, + self.graph, + proposal_sampler, + n_chains=n_chains, + sample_kwargs=sample_kwargs, + symmetry=symmetry, + sector=sector, + encoding=encoding, + amplitude_floor=amplitude_floor, + max_init_attempts=max_init_attempts, + chunk_size=( + self.chunk_size + if chunk_size is None + else _normalize_chunk_size(chunk_size) + ), + generator=( + self.generator + if seed is None and sampler_seed is None + else None + ), + seed=seed if seed is not None else sampler_seed, + device=_model_device(self.model), + log_amplitude_fn=( + self.log_amplitude_fn + if self.log_amplitude_fn is not None + else False + ), + ) + + def sample_bp( + self, + proposal_sampler=None, + *, + sampling=None, + n_samples=1024, + n_chains=None, + n_discard_per_chain=None, + n_discard=None, + sweep_size=None, + n_thin=None, + progress=False, + sample_kwargs=None, + symmetry=None, + sector=None, + encoding=None, + amplitude_floor=0.0, + max_init_attempts=32, + seed=None, + sampler_seed=None, + ): + """Collect chain-preserving samples with BP independence proposals.""" + if sampling is not None: + from .api import BackendCapabilityWarning, SamplingConfig + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if sampling.proposal is not None: + warnings.warn( + "SamplingConfig.proposal is ignored by the BP sampler; " + "pass proposal_sampler for BP proposals.", + BackendCapabilityWarning, + stacklevel=2, + ) + config_kwargs = sampling.torch_kwargs() + 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.") + 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.") + n_discard_per_chain = config_kwargs["n_discard_per_chain"] + n_thin = config_kwargs["n_thin"] + seed = config_kwargs["seed"] + sampler_seed = config_kwargs["sampler_seed"] + sampling_chunk_size = sampling.chunk_size + else: + sampling_chunk_size = None + sampler = self.make_bp_sampler( + proposal_sampler, + n_chains=n_chains, + sample_kwargs=sample_kwargs, + symmetry=symmetry, + sector=sector, + encoding=encoding, + amplitude_floor=amplitude_floor, + max_init_attempts=max_init_attempts, + seed=seed, + sampler_seed=sampler_seed, + chunk_size=sampling_chunk_size, + ) + result = sampler.sample( + n_samples=n_samples, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + ) + self.configs = sampler.configs + self.amplitudes = sampler.amplitudes + self.log_abs_amplitudes = sampler.log_abs_amplitudes + self.nonzero_amplitudes = sampler.nonzero_amplitudes + self.generator = sampler.generator + self._bp_sampler = sampler + return result + + def sample( + self, + *, + sampling=None, + n_samples=1024, + n_chains=None, + n_discard_per_chain=None, + n_discard=None, + sweep_size=None, + n_thin=None, + progress=False, + seed=None, + sampler_seed=None, + track_proposal_stats=False, + ): + """Collect chain-preserving samples and update the driver state.""" + sampling_chunk_size = None + sampling_proposal = None + if sampling is not None: + from .api import SamplingConfig + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + config_kwargs = sampling.torch_kwargs() + 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.") + 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.") + n_discard_per_chain = config_kwargs["n_discard_per_chain"] + n_thin = config_kwargs["n_thin"] + seed = config_kwargs["seed"] + sampler_seed = config_kwargs["sampler_seed"] + sampling_chunk_size = sampling.chunk_size + sampling_proposal = sampling.proposal + sampler = self.make_sampler( + n_chains=n_chains, + seed=seed, + sampler_seed=sampler_seed, + proposal=sampling_proposal, + chunk_size=sampling_chunk_size, + ) + result = sampler.sample( + n_samples=n_samples, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + track_proposal_stats=track_proposal_stats, + ) + self.configs = sampler.configs + self.amplitudes = sampler.amplitudes + self.generator = sampler.generator + if track_proposal_stats: + self.last_proposal_stats = result.proposal_stats + return result + + def make_connections(self, configs=None, *, terms=None): + """Build connected configurations for ``configs``. + + Passing ``terms`` compiles a one-off native operator mapping with this + driver's lattice/site order. This is useful for measuring energy and + correlators from the same Markov samples. + """ + configs = self.configs if configs is None else _as_long_matrix(configs) + if terms is not None: + from .api import CompiledOperatorSum + term_constant = 0.0 + if isinstance(terms, CompiledOperatorSum): + if terms.backend != "torch": + raise ValueError( + f"Compiled terms target backend {terms.backend!r}, not 'torch'." + ) + term_constant = terms.constant + terms = terms.terms + return _driver_terms_connections( + configs, + self.graph, + terms=terms, + site_order=self.site_order, + constant=term_constant, + ) + return self.connection_fn(configs, self.graph, **self.connection_kwargs) + + 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) + result = None + n_proposed = 0 + n_accepted = 0 + proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + with _require_torch().no_grad(): + for _ in range(n_sweeps): + result = metropolis_exchange_sweep( + self.configs, + self.model, + self.graph, + current_amplitudes=self.amplitudes, + current_log_abs=self.log_abs_amplitudes, + current_nonzero=self.nonzero_amplitudes, + log_amplitude_fn=( + self.log_amplitude_fn + if self.log_amplitude_fn is not None + else False + ), + 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, + generator=self.generator, + chunk_size=self.chunk_size, + compile_kernels=self.compile_kernels, + track_proposal_stats=track_proposal_stats, + ) + self.configs = result.configs + self.amplitudes = result.amplitudes + self.log_abs_amplitudes = result.log_abs_amplitudes + self.nonzero_amplitudes = result.nonzero_amplitudes + if result.log_abs_amplitudes is None: + self.log_amplitude_fn = None + n_proposed += result.n_proposed + n_accepted += result.n_accepted + if track_proposal_stats: + _merge_proposal_stats(proposal_stats, result.proposal_stats) + if result is not None: + result = replace( + result, + n_proposed=n_proposed, + n_accepted=n_accepted, + proposal_stats=proposal_stats, + ) + if track_proposal_stats: + self.last_proposal_stats = proposal_stats + return result + + def burn_in( + self, + n_sweeps=32, + *, + progress=False, + track_proposal_stats=False, + ): + """Equilibrate local walkers before fixed-kernel VMC work. + + This is the canonical convenience method for ordinary fixed-rate + burn-in. Use :meth:`warmup_proposal_mix` first when the local move + weights should be tuned; its adaptive samples are deliberately kept + separate from this fixed-kernel stage. + """ + n_sweeps = _check_positive_int("n_sweeps", n_sweeps) + if not progress: + return self.sample_sweep( + n_sweeps=n_sweeps, + track_proposal_stats=track_proposal_stats, + ) + + bar = _make_progress( + True, + total=n_sweeps, + desc="Torch VMC burn-in", + unit="sweep", + ) + result = None + n_proposed = 0 + n_accepted = 0 + proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + try: + for _ in range(n_sweeps): + result = self.sample_sweep( + n_sweeps=1, + track_proposal_stats=track_proposal_stats, + ) + n_proposed += result.n_proposed + n_accepted += result.n_accepted + if track_proposal_stats: + _merge_proposal_stats(proposal_stats, result.proposal_stats) + bar.update(1) + _set_vmc_progress_postfix( + bar, + result, + n_sites=self.n_sites, + include_energy=False, + ) + finally: + bar.close() + + result = replace( + result, + n_proposed=n_proposed, + n_accepted=n_accepted, + proposal_stats=proposal_stats, + ) + if track_proposal_stats: + self.last_proposal_stats = proposal_stats + return result + def warmup_proposal_mix( + self, + *, + n_sweeps=32, + adaptation_rate=1.0, + min_probability=0.05, + max_probability=0.95, + progress=False, + ): + """Tune move weights during warm-up, then leave them fixed. -def _energy_mean_and_variance(local_energies): - energy_mean = local_energies.mean() - centered = local_energies - energy_mean - variance = centered.abs().square().mean() - return energy_mean, variance.real + Adaptation occurs only between complete graph sweeps. The returned + counters describe warm-up only; call :meth:`sample`, + :meth:`step`, or an estimator afterwards for fixed-kernel production + sampling. + """ + return _warmup_proposal_mix( + self, + n_sweeps=n_sweeps, + adaptation_rate=adaptation_rate, + min_probability=min_probability, + max_probability=max_probability, + progress=progress, + ) + def local_energies(self, *, connections=None): + """Evaluate local energies for the current walkers.""" + connections = self.make_connections() if connections is None else connections + with _require_torch().no_grad(): + return local_energy_from_connections( + self.configs, + self.amplitudes, + connections, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) -def _resolve_connection_fn(connection_fn): - if callable(connection_fn): - return None, connection_fn - key = str(connection_fn).replace("-", "_").lower() - aliases = { - "fermi_hubbard": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), - "fh": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), - "hubbard": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), - "spinful": ("spinful_fermi_hubbard", spinful_fermi_hubbard_connections), - "spinful_fermi_hubbard": ( - "spinful_fermi_hubbard", - spinful_fermi_hubbard_connections, - ), - "heisenberg": ("heisenberg", heisenberg_connections), - "heis": ("heisenberg", heisenberg_connections), - "transverse_ising": ("transverse_ising", transverse_ising_connections), - "tfim": ("transverse_ising", transverse_ising_connections), - "ising": ("transverse_ising", transverse_ising_connections), - } - try: - return aliases[key] - except KeyError as exc: - allowed = ", ".join(sorted(aliases)) - raise ValueError( - f"Unknown torch VMC connection_fn {connection_fn!r}. " - f"Expected a callable or one of: {allowed}." - ) from exc + def local_observables( + self, + observables, + *, + configs=None, + amplitudes=None, + ): + """Evaluate named native-term observables with shared amplitudes. + ``observables`` maps names to term mappings accepted by + :func:`torch_hamiltonian_connections`. A value of ``None`` reuses the + observable configured on this driver. Matching connected target + configurations are contracted once across all names. + """ + configs = self.configs if configs is None else _as_long_matrix(configs) + if amplitudes is None: + if configs is self.configs: + amplitudes = self.amplitudes + else: + with _require_torch().no_grad(): + amplitudes = _call_amplitude_fn( + self.model, + configs, + chunk_size=self.chunk_size, + ) + amplitudes = _require_torch().as_tensor( + amplitudes, + device=configs.device, + ) + connection_map = { + name: self.make_connections(configs, terms=terms) + if terms is not None + else self.make_connections(configs) + for name, terms in observables.items() + } + with _require_torch().no_grad(): + return _local_energies_from_connection_map( + configs, + amplitudes, + connection_map, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) -@dataclass(frozen=True) -class TorchVMCStepResult: - """Result of one :class:`TorchVMCDriver` step.""" + def measure_samples( + self, + samples, + *, + observables=None, + amplitudes=None, + weights=None, + proposal_log_probs=None, + profile=False, + deduplicate=True, + ): + """Measure saved chain samples without running another sampler. + + ``samples`` can be a :class:`TorchMCMCSamples` instance or an integer + 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. + + With ``observables=None`` the driver's configured connection function + is measured and one :class:`TorchVMCEnergyEstimate` is returned. A + mapping of names to native term mappings returns one estimate per + name, sharing connected-target amplitudes and boundary environments. + The returned configurations and amplitudes always retain their chain + shape, so chain diagnostics are computed without resampling. ``weights`` + supplies a fixed non-negative weighted empirical batch. Alternatively, + pass ``proposal_log_probs`` for the proposal density ``log q(x)``; + the method then computes self-normalized importance weights + ``|psi(x)|**2 / q(x)`` at the current parameters. Passing both is an + error. A :class:`pepsy.vmc.VMCSamples` can carry either value directly. + + Weighted estimates report the importance effective sample size and do + not report MCMC R-hat/autocorrelation diagnostics, since those assume + identically weighted chain 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. + """ + torch = _require_torch() + start = time.perf_counter() + model_device = _model_device(self.model) + + sample_object = samples if hasattr(samples, "configs") else None + raw_configs = ( + getattr(sample_object, "configs", None) + if sample_object is not None + else samples + ) + if raw_configs is None: + raise TypeError( + "samples must be a TorchMCMCSamples instance or an integer " + "tensor of configurations." + ) + raw_configs = torch.as_tensor(raw_configs, dtype=torch.long) + if raw_configs.ndim == 2: + chain_configs = raw_configs.reshape(1, *raw_configs.shape) + elif raw_configs.ndim == 3: + chain_configs = raw_configs + else: + raise ValueError( + "samples must have shape (n_samples_per_chain, n_chains, " + "n_sites) or (n_chains, n_sites)." + ) + chain_configs = chain_configs.to(device=model_device) + n_steps, n_chains, n_sites = (int(value) for value in chain_configs.shape) + 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) + unique_parent_count = ( + int(_unique_config_rows(flat_configs)[0].shape[0]) + if deduplicate + else int(flat_configs.shape[0]) + ) - configs: Any - amplitudes: Any - local_energies: Any - energy_mean: Any - energy_variance: Any - acceptance_rate: float - n_proposed: int - n_accepted: int - sr: Any = None + if amplitudes is None and sample_object is not None: + amplitudes = getattr(sample_object, "amplitudes", None) + if amplitudes is None: + 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, + ) + flat_amplitudes = unique_amplitudes[inverse] + else: + flat_amplitudes = _call_amplitude_fn( + 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) + if amplitudes.ndim == 1: + if tuple(amplitudes.shape) != (n_chains,): + raise ValueError( + "one-dimensional amplitudes must have one value per " + "chain." + ) + chain_amplitudes = amplitudes.reshape(1, n_chains) + if n_steps != 1: + raise ValueError( + "one-dimensional amplitudes are only valid for one " + "sample per chain." + ) + elif amplitudes.ndim == 2: + if tuple(amplitudes.shape) != (n_steps, n_chains): + raise ValueError( + "amplitudes must match the first two sample dimensions: " + f"expected {(n_steps, n_chains)}, got " + f"{tuple(amplitudes.shape)}." + ) + chain_amplitudes = amplitudes + else: + raise ValueError( + "amplitudes must have shape (n_samples_per_chain, n_chains)." + ) + flat_amplitudes = chain_amplitudes.reshape(-1) + + if weights is None and sample_object is not None: + weights = getattr(sample_object, "weights", None) + if proposal_log_probs is None and sample_object is not None: + proposal_log_probs = getattr(sample_object, "proposal_log_probs", None) + if weights is not None and proposal_log_probs is not None: + raise ValueError("Pass either weights or proposal_log_probs, not both.") + if proposal_log_probs is not None: + importance_weights = _importance_weights_from_log_probs( + flat_amplitudes, + proposal_log_probs, + n_steps=n_steps, + n_chains=n_chains, + ) + elif weights is not None: + importance_weights = _normalized_sample_weights( + weights, + n_steps=n_steps, + n_chains=n_chains, + device=model_device, + ) + else: + importance_weights = None + if observables is None: + observable_items = (("observable", None),) + return_mapping = False + 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.") + return_mapping = True + + 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 + } + connection_elapsed = time.perf_counter() - connection_start + + local_start = time.perf_counter() + with torch.no_grad(): + flat_values = _local_energies_from_connection_map( + flat_configs, + flat_amplitudes, + connection_map, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=deduplicate, + compile_kernels=self.compile_kernels, + ) + local_elapsed = time.perf_counter() - local_start + elapsed = time.perf_counter() - start + + acceptance_rate = float( + getattr(sample_object, "acceptance_rate", 0.0) + ) 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 + profile_data = None + if profile: + profile_data = { + "sampling_seconds": 0.0, + "connection_seconds": connection_elapsed, + "local_estimator_seconds": local_elapsed, + "postprocess_seconds": max( + elapsed - connection_elapsed - local_elapsed, + 0.0, + ), + "total_seconds": elapsed, + "cache": _cache_profile_snapshot(self.model), + "samples_only": True, + "deduplicate": bool(deduplicate), + "num_samples": int(flat_configs.shape[0]), + "num_unique_samples": unique_parent_count, + "weighted": importance_weights is not None, + } -class TorchVMCDriver: - """Small PyTorch-native VMC loop around Pepsy's torch kernels. + results = {} + for name, _ in observable_items: + local_values = flat_values[name].reshape(n_steps, n_chains) + ( + energy_mean, + energy_variance, + energy_stderr, + energy_stderr_naive, + effective_sample_size, + chain_diagnostics, + ) = ( + _observable_statistics(local_values) + if importance_weights is None + else (*_weighted_energy_statistics(flat_values[name], importance_weights), None) + ) + result_profile = None + if profile_data is not None: + result_profile = dict(profile_data) + result_profile["observable"] = name + results[name] = TorchVMCEnergyEstimate( + configs=chain_configs, + amplitudes=chain_amplitudes, + local_energies=local_values, + energy_mean=energy_mean, + energy_variance=energy_variance, + energy_stderr=energy_stderr, + acceptance_rate=acceptance_rate, + n_proposed=n_proposed, + n_accepted=n_accepted, + n_samples=int(local_values.numel()), + n_measurements=n_steps, + elapsed_seconds=elapsed, + samples_per_second=( + int(local_values.numel()) / elapsed + if elapsed > 0 + else float("inf") + ), + chain_diagnostics=chain_diagnostics, + profile=result_profile, + energy_stderr_naive=energy_stderr_naive, + effective_sample_size=effective_sample_size, + importance_weights=( + None + if importance_weights is None + else importance_weights.reshape(n_steps, n_chains) + ), + proposal_log_probs=( + None + if proposal_log_probs is None + else _flat_sample_values( + proposal_log_probs, + n_steps=n_steps, + n_chains=n_chains, + device=model_device, + name="proposal_log_probs", + ).reshape(n_steps, n_chains) + ), + ) - The driver keeps walker configurations and amplitudes in sync, runs - Metropolis exchange/hopping sweeps, evaluates local energies with optional - chunking/diagonal reuse, and can apply one SR/minSR update per step. - """ + return results if return_mapping else results["observable"] - def __init__( + def energy_estimate(self): + """Return ``(mean, variance, local_energies)`` for current walkers.""" + local_energies = self.local_energies() + energy_mean, energy_variance = _energy_mean_and_variance(local_energies) + return energy_mean, energy_variance, local_energies + + def estimate_observable( self, - model, - graph, - configs, - connection_fn="spinful_fermi_hubbard", *, - connection_kwargs=None, - amplitudes=None, - proposal="spinful", - hopping_rate=0.25, - encoding=None, - chunk_size=None, - generator=None, + burn_in=0, + n_measurements=1, + sweeps_between=1, + progress=False, + n_samples=None, + n_chains=None, + n_discard_per_chain=None, + n_discard=None, + sweep_size=None, + n_thin=None, + sampler=None, + seed=None, + sampler_seed=None, + profile=False, + target_effective_sample_size=None, + min_measurements=8, + ess_check_interval=1, + rhat_threshold=1.05, + auto_thin=False, ): - self.model = model - self.graph = graph - self.configs = _as_long_matrix(configs) - self.connection_name, self.connection_fn = _resolve_connection_fn( - connection_fn + """Run burn-in and estimate the configured local observable. + + The driver keeps the configured walkers and collects all walker local + observable values after each measurement. ``n_samples`` therefore equals + ``n_walkers * n_measurements``. The returned ``samples_per_second`` + measures completed observable samples, including sampling and + contraction time. + + The result retains the historical ``TorchVMCEnergyEstimate`` type and + ``energy_*`` field names for compatibility. They describe the + observable encoded by this driver's configured ``terms`` or connection + function and are not restricted to a Hamiltonian. + + Set ``target_effective_sample_size`` to stop the legacy sweep-based + loop once the requested ESS (and optionally ``rhat_threshold``) is + reached. In that mode, ``n_measurements`` is a hard cap rather than a + fixed count. ``auto_thin=True`` increases later sweep spacing to the + estimated integrated autocorrelation time. This is currently available + for the legacy ``burn_in``/``n_measurements`` interface only. + + Set ``profile=True`` to attach phase timings and the latest available + PEPS boundary-cache counters to the result. Profiling is deliberately + opt-in so normal short VMC loops keep their existing overhead. + """ + profile = bool(profile) + modern_sampling = ( + sampler is not None + or n_samples is not None + or n_chains is not None + or n_discard_per_chain is not None + or n_discard is not None + or sweep_size is not None + or n_thin is not None + or seed is not None + or sampler_seed is not None ) - self.connection_kwargs = ( - {} if connection_kwargs is None else dict(connection_kwargs) + if modern_sampling: + if target_effective_sample_size is not None: + raise ValueError( + "target_effective_sample_size is currently supported with " + "burn_in/n_measurements sampling; omit n_samples and " + "sampler controls." + ) + if sampler is None: + samples = self.sample( + n_samples=1024 if n_samples is None else n_samples, + n_chains=n_chains, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + seed=seed, + sampler_seed=sampler_seed, + ) + else: + if any( + value is not None + for value in ( + n_chains, + seed, + sampler_seed, + ) + ): + raise ValueError( + "n_chains and sampler seeds must be configured on an " + "explicit sampler." + ) + samples = sampler.sample( + n_samples=1024 if n_samples is None else n_samples, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + ) + estimator_start = time.perf_counter() + sample_configs = samples.configs + sample_amplitudes = samples.amplitudes + flat_configs = sample_configs.reshape(-1, self.n_sites) + flat_amplitudes = sample_amplitudes.reshape(-1) + connection_start = time.perf_counter() + connections = self.make_connections(flat_configs) + connection_elapsed = time.perf_counter() - connection_start + local_start = time.perf_counter() + with _require_torch().no_grad(): + local_values = local_energy_from_connections( + flat_configs, + flat_amplitudes, + connections, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) + local_elapsed = time.perf_counter() - local_start + n_actual = int(local_values.numel()) + chain_values = local_values.reshape(sample_configs.shape[:-1]) + ( + observable_mean, + observable_variance, + observable_stderr, + observable_stderr_naive, + effective_sample_size, + chain_diagnostics, + ) = _observable_statistics(chain_values) + estimator_elapsed = time.perf_counter() - estimator_start + elapsed = estimator_elapsed + samples.elapsed_seconds + profile_data = None + if profile: + profile_data = { + "sampling_seconds": samples.elapsed_seconds, + "connection_seconds": connection_elapsed, + "local_estimator_seconds": local_elapsed, + "postprocess_seconds": max( + estimator_elapsed - connection_elapsed - local_elapsed, + 0.0, + ), + "total_seconds": elapsed, + "cache": _cache_profile_snapshot(self.model), + } + return TorchVMCEnergyEstimate( + configs=sample_configs, + amplitudes=sample_amplitudes, + local_energies=chain_values, + energy_mean=observable_mean, + energy_variance=observable_variance, + energy_stderr=observable_stderr, + acceptance_rate=samples.acceptance_rate, + n_proposed=samples.n_proposed, + n_accepted=samples.n_accepted, + n_samples=n_actual, + n_measurements=samples.n_samples_per_chain, + elapsed_seconds=elapsed, + samples_per_second=( + n_actual / elapsed if elapsed > 0 else float("inf") + ), + chain_diagnostics=chain_diagnostics, + profile=profile_data, + energy_stderr_naive=observable_stderr_naive, + effective_sample_size=effective_sample_size, + ) + + burn_in = int(burn_in) + n_measurements = _check_positive_int("n_measurements", n_measurements) + sweeps_between = _check_positive_int("sweeps_between", sweeps_between) + if burn_in < 0: + raise ValueError("burn_in must be non-negative.") + adaptive_options = _adaptive_measurement_options( + target_effective_sample_size, + max_measurements=n_measurements, + min_measurements=min_measurements, + ess_check_interval=ess_check_interval, + rhat_threshold=rhat_threshold, + auto_thin=auto_thin, ) - self.proposal = proposal - self.hopping_rate = float(hopping_rate) - self.encoding = encoding - self.chunk_size = _normalize_chunk_size(chunk_size) - self.generator = generator - if ( - self.connection_name == "spinful_fermi_hubbard" - and encoding is not None - and "encoding" not in self.connection_kwargs + total_sweeps = burn_in + n_measurements * sweeps_between + bar = _make_progress( + progress, + total=( + None + if adaptive_options is not None and adaptive_options["auto_thin"] + else total_sweeps + ), + desc="Torch VMC", + ) + start = time.perf_counter() + n_proposed = 0 + n_accepted = 0 + sampling_elapsed = 0.0 + connection_elapsed = 0.0 + local_elapsed = 0.0 + cache_profile = {} + + def run_sweeps(count): + nonlocal n_proposed, n_accepted, sampling_elapsed + for _ in range(count): + sweep_start = time.perf_counter() + sample = self.sample_sweep(n_sweeps=1) + sampling_elapsed += time.perf_counter() - sweep_start + n_proposed += sample.n_proposed + n_accepted += sample.n_accepted + proposal_stats = getattr( + self.model, + "last_proposal_cache_stats", + None, + ) + if profile and proposal_stats is not None: + _accumulate_cache_profile( + cache_profile, + {"proposal": proposal_stats}, + ) + if bar is not None: + bar.update(1) + + run_sweeps(burn_in) + measurements = [] + current_sweeps_between = sweeps_between + stop_reason = "max_measurements" + for _ in range(n_measurements): + run_sweeps(current_sweeps_between) + if profile: + connection_start = time.perf_counter() + connections = self.make_connections() + connection_elapsed += time.perf_counter() - connection_start + local_start = time.perf_counter() + with _require_torch().no_grad(): + local_values = local_energy_from_connections( + self.configs, + self.amplitudes, + connections, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) + local_elapsed += time.perf_counter() - local_start + connected_stats = getattr( + self.model, + "last_connected_reuse_stats", + None, + ) + if connected_stats is not None: + _accumulate_cache_profile( + cache_profile, + {"connected": connected_stats}, + ) + else: + local_values = self.local_energies() + measurements.append(local_values.detach()) + if ( + adaptive_options is not None + and len(measurements) >= adaptive_options["min_measurements"] + and ( + len(measurements) - adaptive_options["min_measurements"] + ) + % adaptive_options["ess_check_interval"] + == 0 + ): + diagnostics = _observable_statistics( + _require_torch().stack(measurements, dim=0) + )[-1] + if adaptive_options["auto_thin"]: + current_sweeps_between = _adaptive_thinning_interval( + diagnostics, + sweeps_between, + ) + if _diagnostics_meet_target(diagnostics, adaptive_options): + stop_reason = "target_effective_sample_size" + break + if bar is not None: + bar.close() + + chain_values = _require_torch().stack(measurements, dim=0) + local_energies = chain_values.reshape(-1) + ( + energy_mean, + energy_variance, + energy_stderr, + energy_stderr_naive, + effective_sample_size, + chain_diagnostics, + ) = _observable_statistics(chain_values) + n_samples = int(local_energies.numel()) + elapsed = time.perf_counter() - start + acceptance = n_accepted / n_proposed if n_proposed else 0.0 + profile_data = None + if profile: + _accumulate_cache_profile( + cache_profile, + {"cutoff_fallbacks": getattr(self.model, "cutoff_fallbacks", 0)}, + ) + profile_data = { + "sampling_seconds": sampling_elapsed, + "connection_seconds": connection_elapsed, + "local_estimator_seconds": local_elapsed, + "postprocess_seconds": max( + elapsed - sampling_elapsed - connection_elapsed - local_elapsed, + 0.0, + ), + "total_seconds": elapsed, + "cache": cache_profile, + } + if adaptive_options is not None: + profile_data["adaptive_sampling"] = { + "target_effective_sample_size": adaptive_options[ + "target_effective_sample_size" + ], + "measurements_collected": len(measurements), + "final_sweeps_between": current_sweeps_between, + "stop_reason": stop_reason, + } + return TorchVMCEnergyEstimate( + configs=self.configs, + amplitudes=self.amplitudes, + local_energies=local_energies, + energy_mean=energy_mean, + energy_variance=energy_variance, + energy_stderr=energy_stderr, + acceptance_rate=acceptance, + n_proposed=n_proposed, + n_accepted=n_accepted, + n_samples=n_samples, + n_measurements=len(measurements), + elapsed_seconds=elapsed, + samples_per_second=n_samples / elapsed if elapsed > 0 else float("inf"), + chain_diagnostics=chain_diagnostics, + profile=profile_data, + energy_stderr_naive=energy_stderr_naive, + effective_sample_size=effective_sample_size, + ) + + def estimate_observables( + self, + observables, + *, + burn_in=0, + n_measurements=1, + sweeps_between=1, + progress=False, + n_samples=None, + n_chains=None, + n_discard_per_chain=None, + n_discard=None, + sweep_size=None, + n_thin=None, + sampler=None, + seed=None, + sampler_seed=None, + profile=False, + target_effective_sample_size=None, + min_measurements=8, + ess_check_interval=1, + rhat_threshold=1.05, + auto_thin=False, + ): + """Estimate several native-term observables from the same samples. + + ``observables`` maps result names to native term mappings. Use + ``None`` as a value to reuse the observable configured on this driver. + The returned dictionary maps every name to a + :class:`TorchVMCEnergyEstimate`. It shares Markov samples, connected + target amplitudes, and boundary environments across all observables. + When ``target_effective_sample_size`` is set, the legacy sweep-based + path stops only after every requested observable satisfies the ESS + target and optional R-hat threshold; ``n_measurements`` is then a + hard cap. The modern sampler interface retains its fixed sample count. + """ + try: + observable_items = tuple(observables.items()) + except AttributeError as exc: + raise TypeError("observables must be a mapping of names to terms.") from exc + if not observable_items: + raise ValueError("observables must contain at least one entry.") + profile = bool(profile) + + def make_connection_map(configs): + return { + name: self.make_connections(configs, terms=terms) + if terms is not None + else self.make_connections(configs) + for name, terms in observable_items + } + + def make_results( + *, + sample_configs, + sample_amplitudes, + local_values, + acceptance_rate, + n_proposed, + n_accepted, + n_measurements_result, + elapsed, + profile_data, ): - self.connection_kwargs["encoding"] = encoding + results = {} + n_actual = int(sample_configs.shape[0] * sample_configs.shape[1]) + for name, _ in observable_items: + values = local_values[name] + ( + observable_mean, + observable_variance, + observable_stderr, + observable_stderr_naive, + effective_sample_size, + chain_diagnostics, + ) = _observable_statistics(values) + result_profile = None + if profile_data is not None: + result_profile = dict(profile_data) + result_profile["observable"] = name + results[name] = TorchVMCEnergyEstimate( + configs=sample_configs, + amplitudes=sample_amplitudes, + local_energies=values, + energy_mean=observable_mean, + energy_variance=observable_variance, + energy_stderr=observable_stderr, + acceptance_rate=acceptance_rate, + n_proposed=n_proposed, + n_accepted=n_accepted, + n_samples=n_actual, + n_measurements=n_measurements_result, + elapsed_seconds=elapsed, + samples_per_second=( + n_actual / elapsed if elapsed > 0 else float("inf") + ), + chain_diagnostics=chain_diagnostics, + profile=result_profile, + energy_stderr_naive=observable_stderr_naive, + effective_sample_size=effective_sample_size, + ) + return results + + modern_sampling = ( + sampler is not None + or n_samples is not None + or n_chains is not None + or n_discard_per_chain is not None + or n_discard is not None + or sweep_size is not None + or n_thin is not None + or seed is not None + or sampler_seed is not None + ) + if modern_sampling: + if target_effective_sample_size is not None: + raise ValueError( + "target_effective_sample_size is currently supported with " + "burn_in/n_measurements sampling; omit n_samples and " + "sampler controls." + ) + if sampler is None: + samples = self.sample( + n_samples=1024 if n_samples is None else n_samples, + n_chains=n_chains, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + seed=seed, + sampler_seed=sampler_seed, + ) + else: + if any( + value is not None + for value in (n_chains, seed, sampler_seed) + ): + raise ValueError( + "n_chains and sampler seeds must be configured on an " + "explicit sampler." + ) + samples = sampler.sample( + n_samples=1024 if n_samples is None else n_samples, + n_discard_per_chain=n_discard_per_chain, + n_discard=n_discard, + sweep_size=sweep_size, + n_thin=n_thin, + progress=progress, + ) - if amplitudes is None: - self.refresh_amplitudes() - else: - self.amplitudes = _require_torch().as_tensor( - amplitudes, - device=self.configs.device, + phase_bar = _make_progress( + progress, + total=3, + desc="Torch VMC evaluation", + unit="phase", ) + observable_names = ", ".join(name for name, _ in observable_items) - @property - def n_walkers(self): - """Number of active walkers.""" - return int(self.configs.shape[0]) + def set_phase(stage): + if phase_bar is not None: + phase_bar.set_postfix({"stage": stage}) - @property - def n_sites(self): - """Number of sites in each walker configuration.""" - return int(self.configs.shape[1]) + try: + estimator_start = time.perf_counter() + sample_configs = samples.configs + sample_amplitudes = samples.amplitudes + flat_configs = sample_configs.reshape(-1, self.n_sites) + flat_amplitudes = sample_amplitudes.reshape(-1) + + set_phase("building shared connections") + connection_start = time.perf_counter() + connection_map = make_connection_map(flat_configs) + connection_elapsed = time.perf_counter() - connection_start + if phase_bar is not None: + phase_bar.update(1) + + set_phase(f"contracting {observable_names}") + local_start = time.perf_counter() + with _require_torch().no_grad(): + flat_values = _local_energies_from_connection_map( + flat_configs, + flat_amplitudes, + connection_map, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) + local_elapsed = time.perf_counter() - local_start + if phase_bar is not None: + phase_bar.update(1) + + set_phase("computing statistics") + local_values = { + name: values.reshape(sample_configs.shape[:-1]) + for name, values in flat_values.items() + } + estimator_elapsed = time.perf_counter() - estimator_start + elapsed = samples.elapsed_seconds + estimator_elapsed + profile_data = None + if profile: + profile_data = { + "sampling_seconds": samples.elapsed_seconds, + "connection_seconds": connection_elapsed, + "local_estimator_seconds": local_elapsed, + "postprocess_seconds": max( + estimator_elapsed - connection_elapsed - local_elapsed, + 0.0, + ), + "total_seconds": elapsed, + "shared_observables": tuple( + name for name, _ in observable_items + ), + "cache": _cache_profile_snapshot(self.model), + } + result = make_results( + sample_configs=sample_configs, + sample_amplitudes=sample_amplitudes, + local_values=local_values, + acceptance_rate=samples.acceptance_rate, + n_proposed=samples.n_proposed, + n_accepted=samples.n_accepted, + n_measurements_result=samples.n_samples_per_chain, + elapsed=elapsed, + profile_data=profile_data, + ) + if phase_bar is not None: + phase_bar.update(1) + return result + finally: + if phase_bar is not None: + phase_bar.close() + + burn_in = int(burn_in) + n_measurements = _check_positive_int("n_measurements", n_measurements) + sweeps_between = _check_positive_int("sweeps_between", sweeps_between) + if burn_in < 0: + raise ValueError("burn_in must be non-negative.") + adaptive_options = _adaptive_measurement_options( + target_effective_sample_size, + max_measurements=n_measurements, + min_measurements=min_measurements, + ess_check_interval=ess_check_interval, + rhat_threshold=rhat_threshold, + auto_thin=auto_thin, + ) - def refresh_amplitudes(self): - """Recompute current walker amplitudes from the current model.""" - with _require_torch().no_grad(): - self.amplitudes = _call_amplitude_fn( + total_sweeps = burn_in + n_measurements * sweeps_between + bar = _make_progress( + progress, + total=( + None + if adaptive_options is not None and adaptive_options["auto_thin"] + else total_sweeps + ), + desc="Torch VMC", + ) + start = time.perf_counter() + n_proposed = 0 + n_accepted = 0 + sampling_elapsed = 0.0 + connection_elapsed = 0.0 + local_elapsed = 0.0 + cache_profile = {} + + def run_sweeps(count): + nonlocal n_proposed, n_accepted, sampling_elapsed + for _ in range(count): + sweep_start = time.perf_counter() + sample = self.sample_sweep(n_sweeps=1) + sampling_elapsed += time.perf_counter() - sweep_start + n_proposed += sample.n_proposed + n_accepted += sample.n_accepted + proposal_stats = getattr( + self.model, + "last_proposal_cache_stats", + None, + ) + if profile and proposal_stats is not None: + _accumulate_cache_profile( + cache_profile, + {"proposal": proposal_stats}, + ) + if bar is not None: + bar.update(1) + + run_sweeps(burn_in) + measurements = {name: [] for name, _ in observable_items} + sample_config_records = [] + sample_amplitude_records = [] + current_sweeps_between = sweeps_between + stop_reason = "max_measurements" + for _ in range(n_measurements): + run_sweeps(current_sweeps_between) + sample_config_records.append(self.configs.detach().clone()) + sample_amplitude_records.append(self.amplitudes.detach().clone()) + connection_start = time.perf_counter() + connection_map = make_connection_map(self.configs) + connection_elapsed += time.perf_counter() - connection_start + local_start = time.perf_counter() + with _require_torch().no_grad(): + values = _local_energies_from_connection_map( + self.configs, + self.amplitudes, + connection_map, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) + local_elapsed += time.perf_counter() - local_start + for name, value in values.items(): + measurements[name].append(value.detach()) + connected_stats = getattr( self.model, - self.configs, + "last_connected_reuse_stats", + None, + ) + if profile and connected_stats is not None: + _accumulate_cache_profile( + cache_profile, + {"connected": connected_stats}, + ) + n_collected = len(sample_config_records) + if ( + adaptive_options is not None + and n_collected >= adaptive_options["min_measurements"] + and ( + n_collected - adaptive_options["min_measurements"] + ) + % adaptive_options["ess_check_interval"] + == 0 + ): + diagnostics_by_name = { + name: _observable_statistics( + _require_torch().stack(values, dim=0) + )[-1] + for name, values in measurements.items() + } + if adaptive_options["auto_thin"]: + current_sweeps_between = max( + _adaptive_thinning_interval(diagnostics, sweeps_between) + for diagnostics in diagnostics_by_name.values() + ) + if all( + _diagnostics_meet_target(diagnostics, adaptive_options) + for diagnostics in diagnostics_by_name.values() + ): + stop_reason = "target_effective_sample_size" + break + if bar is not None: + bar.close() + + elapsed = time.perf_counter() - start + local_values = { + name: _require_torch().stack(values, dim=0) + for name, values in measurements.items() + } + acceptance = n_accepted / n_proposed if n_proposed else 0.0 + profile_data = None + if profile: + _accumulate_cache_profile( + cache_profile, + {"cutoff_fallbacks": getattr(self.model, "cutoff_fallbacks", 0)}, + ) + profile_data = { + "sampling_seconds": sampling_elapsed, + "connection_seconds": connection_elapsed, + "local_estimator_seconds": local_elapsed, + "postprocess_seconds": max( + elapsed - sampling_elapsed - connection_elapsed - local_elapsed, + 0.0, + ), + "total_seconds": elapsed, + "shared_observables": tuple(name for name, _ in observable_items), + "cache": cache_profile, + } + if adaptive_options is not None: + profile_data["adaptive_sampling"] = { + "target_effective_sample_size": adaptive_options[ + "target_effective_sample_size" + ], + "measurements_collected": len(sample_config_records), + "final_sweeps_between": current_sweeps_between, + "stop_reason": stop_reason, + } + sample_configs = _require_torch().stack(sample_config_records, dim=0) + sample_amplitudes = _require_torch().stack( + sample_amplitude_records, + dim=0, + ) + return make_results( + sample_configs=sample_configs, + sample_amplitudes=sample_amplitudes, + local_values=local_values, + acceptance_rate=acceptance, + n_proposed=n_proposed, + n_accepted=n_accepted, + n_measurements_result=len(sample_config_records), + elapsed=elapsed, + profile_data=profile_data, + ) + + def estimate_energy( + self, + *, + burn_in=0, + n_measurements=1, + sweeps_between=1, + progress=False, + profile=False, + target_effective_sample_size=None, + min_measurements=8, + ess_check_interval=1, + rhat_threshold=1.05, + auto_thin=False, + ): + """Compatibility wrapper for :meth:`estimate_observable`.""" + return self.estimate_observable( + burn_in=burn_in, + n_measurements=n_measurements, + sweeps_between=sweeps_between, + progress=progress, + profile=profile, + target_effective_sample_size=target_effective_sample_size, + min_measurements=min_measurements, + ess_check_interval=ess_check_interval, + rhat_threshold=rhat_threshold, + auto_thin=auto_thin, + ) + + def importance_energy_estimate( + self, + proposal_sampler, + *, + n_samples=128, + sample_kwargs=None, + amplitude_floor=0.0, + progress=False, + ): + """Measure the driver Hamiltonian using an external proposal sampler. + + ``proposal_sampler`` should expose ``sample(samples=..., progbar=...)`` + and return a PEPS-BP-style result with ``configs`` and ``omegas``. + The sampler proposes configurations; torch evaluates their PEPS + amplitudes and local energies. The returned self-normalized weights are + ``|psi(x)|**2 / q(x)`` and include an effective sample-size diagnostic. + """ + torch = _require_torch() + n_samples = _check_positive_int("n_samples", n_samples) + if amplitude_floor < 0: + raise ValueError("amplitude_floor must be non-negative.") + sample_kwargs = dict(sample_kwargs or {}) + sample_kwargs.setdefault("samples", n_samples) + sample_kwargs.setdefault("progbar", bool(progress)) + start = time.perf_counter() + try: + proposed = proposal_sampler.sample(**sample_kwargs) + except TypeError: + # Small custom proposal samplers often don't expose ``progbar``. + sample_kwargs.pop("progbar", None) + proposed = proposal_sampler.sample(**sample_kwargs) + + device = self.configs.device + configs = _as_long_matrix(proposed.configs, name="proposal configs") + configs = configs.to(device=device) + if configs.shape[0] != n_samples: + n_samples = int(configs.shape[0]) + log_q = _proposal_log_probabilities(proposed.omegas, device=device) + if log_q.shape[0] != configs.shape[0]: + raise ValueError("proposal probabilities must match proposal configs.") + + with torch.no_grad(): + amplitudes = _call_amplitude_fn( + self.model, + configs, chunk_size=self.chunk_size, ) - return self.amplitudes - - def make_connections(self, configs=None): - """Build Hamiltonian-connected configurations for ``configs``.""" - configs = self.configs if configs is None else _as_long_matrix(configs) - return self.connection_fn(configs, self.graph, **self.connection_kwargs) - - def sample_sweep(self, *, n_sweeps=1): - """Run one or more Metropolis sweeps and update driver state.""" - n_sweeps = _check_positive_int("n_sweeps", n_sweeps) - result = None - with _require_torch().no_grad(): - for _ in range(n_sweeps): - result = metropolis_exchange_sweep( - self.configs, - self.model, - self.graph, - current_amplitudes=self.amplitudes, - proposal=self.proposal, - hopping_rate=self.hopping_rate, - encoding=self.encoding, - generator=self.generator, - chunk_size=self.chunk_size, + amp_abs = amplitudes.abs() + valid = torch.isfinite(amp_abs) & (amp_abs > float(amplitude_floor)) + if not torch.any(valid): + raise ValueError( + "The proposal produced no configurations with non-zero " + "torch PEPS amplitude." ) - self.configs = result.configs - self.amplitudes = result.amplitudes - return result - - def local_energies(self, *, connections=None): - """Evaluate local energies for the current walkers.""" - connections = self.make_connections() if connections is None else connections - with _require_torch().no_grad(): - return local_energy_from_connections( - self.configs, - self.amplitudes, + valid_configs = configs[valid] + valid_amplitudes = amplitudes[valid] + connections = self.make_connections(valid_configs) + local_energies = local_energy_from_connections( + valid_configs, + valid_amplitudes, connections, self.model, chunk_size=self.chunk_size, reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, ) - - def energy_estimate(self): - """Return ``(mean, variance, local_energies)`` for current walkers.""" - local_energies = self.local_energies() - energy_mean, energy_variance = _energy_mean_and_variance(local_energies) - return energy_mean, energy_variance, local_energies + log_weights = 2.0 * valid_amplitudes.abs().log() - log_q[valid] + log_weights = log_weights - log_weights.max() + weights = torch.exp(log_weights) + weights = weights / weights.sum() + energy_mean = (weights.to(local_energies.dtype) * local_energies).sum() + energy_variance = ( + weights * (local_energies - energy_mean).abs().square() + ).sum().real + effective_sample_size = 1.0 / weights.square().sum() + + n_valid = int(valid.sum().item()) + elapsed = time.perf_counter() - start + return TorchVMCImportanceEstimate( + configs=valid_configs, + amplitudes=valid_amplitudes, + local_energies=local_energies, + weights=weights, + energy_mean=energy_mean, + energy_variance=energy_variance, + energy_stderr=torch.sqrt(energy_variance / effective_sample_size), + effective_sample_size=effective_sample_size, + n_samples=n_samples, + n_valid=n_valid, + elapsed_seconds=elapsed, + samples_per_second=n_valid / elapsed if elapsed > 0 else float("inf"), + ) def step( self, @@ -1845,36 +9836,228 @@ def step( learning_rate=1.0, sr_diag_shift=1.0e-4, sr_method="auto", + sr_parameter_mode="holomorphic", + sr_pinv_rtol=None, + sr_momentum=None, amplitude_floor=None, + derivative_backend="auto", + samples=None, + weights=None, + proposal_log_probs=None, + profile=False, + track_proposal_stats=False, ): - """Run sampling, estimate energy, and optionally apply an SR update.""" - sample = self.sample_sweep(n_sweeps=sample_sweeps) - local_energies = self.local_energies() - energy_mean, energy_variance = _energy_mean_and_variance(local_energies) + """Run sampling, estimate energy, and optionally update parameters. + + ``sr_parameter_mode="holomorphic"`` is the explicit convention for + complex PEPS tensor parameters. It returns one complex derivative per + complex tensor entry and applies a complex SR direction in place. + Use ``"real-imag"`` to optimize explicit real and imaginary tensor + coordinates instead. ``derivative_backend="auto"`` uses the batched + PEPS Jacobian path when available and retains the scalar autograd loop + as a compatibility fallback. ``sr_diag_shift`` may be a callable of + the SR update number. ``sr_pinv_rtol`` controls the fallback + pseudoinverse, and ``sr_momentum`` enables a SPRING-style retained + complement of the previous SR direction. Set ``profile=True`` to + attach sampling, estimator, SR, and boundary-cache timings to the + result. By default this advances the internal Metropolis chains. + Passing ``samples=`` skips Metropolis and evaluates that supplied two- + or three-dimensional configuration batch instead. + Such a batch may carry fixed ``weights=`` or ``proposal_log_probs=``; + the latter recomputes ``|psi_theta|**2 / q`` at every update and is + therefore the correct reusable importance-sampling input. + """ + profile = bool(profile) + if samples is None and (weights is not None or proposal_log_probs is not None): + raise ValueError( + "weights and proposal_log_probs require an explicitly supplied " + "samples batch." + ) + if weights is not None and proposal_log_probs is not None: + raise ValueError("Pass either weights or proposal_log_probs, not both.") + total_start = time.perf_counter() + sampling_start = time.perf_counter() + sample = None + if samples is None: + sample = self.sample_sweep( + n_sweeps=sample_sweeps, + track_proposal_stats=track_proposal_stats, + ) + batch_configs = self.configs + batch_amplitudes = self.amplitudes + importance_weights = None + sample_source = "metropolis" + else: + torch = _require_torch() + sample_object = samples if hasattr(samples, "configs") else None + raw_configs = ( + getattr(sample_object, "configs", None) + if sample_object is not None + else samples + ) + if raw_configs is None: + raise TypeError("samples must provide an integer configs batch.") + raw_configs = torch.as_tensor(raw_configs, dtype=torch.long) + if raw_configs.ndim == 2: + n_steps, n_chains, n_sites = 1, *raw_configs.shape + batch_configs = raw_configs + elif raw_configs.ndim == 3: + n_steps, n_chains, n_sites = raw_configs.shape + batch_configs = raw_configs.reshape(-1, n_sites) + else: + raise ValueError( + "samples must have shape (n_samples, n_sites) or " + "(n_samples_per_chain, n_chains, n_sites)." + ) + batch_configs = batch_configs.to(device=_model_device(self.model)) + n_steps, n_chains, n_sites = ( + int(n_steps), int(n_chains), int(n_sites) + ) + if n_steps <= 0 or n_chains <= 0 or n_sites != self.n_sites: + raise ValueError( + f"samples must contain configurations with {self.n_sites} sites." + ) + if weights is None and sample_object is not None: + weights = getattr(sample_object, "weights", None) + if proposal_log_probs is None and sample_object is not None: + proposal_log_probs = getattr(sample_object, "proposal_log_probs", None) + if weights is not None and proposal_log_probs is not None: + raise ValueError("Pass either weights or proposal_log_probs, not both.") + with torch.no_grad(): + batch_amplitudes = _call_amplitude_fn( + self.model, + batch_configs, + chunk_size=self.chunk_size, + ) + if proposal_log_probs is not None: + importance_weights = _importance_weights_from_log_probs( + batch_amplitudes, + proposal_log_probs, + n_steps=n_steps, + n_chains=n_chains, + ) + elif weights is not None: + importance_weights = _normalized_sample_weights( + weights, + n_steps=n_steps, + n_chains=n_chains, + device=batch_configs.device, + ) + else: + importance_weights = None + sample_source = "provided" + sampling_elapsed = time.perf_counter() - sampling_start + connection_elapsed = 0.0 + local_elapsed = 0.0 + if profile: + connection_start = time.perf_counter() + connections = self.make_connections(batch_configs) + connection_elapsed = time.perf_counter() - connection_start + local_start = time.perf_counter() + with _require_torch().no_grad(): + local_energies = local_energy_from_connections( + batch_configs, + batch_amplitudes, + connections, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) + local_elapsed = time.perf_counter() - local_start + else: + connections = self.make_connections(batch_configs) + with _require_torch().no_grad(): + local_energies = local_energy_from_connections( + batch_configs, + batch_amplitudes, + connections, + self.model, + chunk_size=self.chunk_size, + reuse_diagonal=True, + deduplicate_targets=True, + compile_kernels=self.compile_kernels, + ) + if importance_weights is None: + energy_mean, energy_variance = _energy_mean_and_variance(local_energies) + effective_sample_size = _require_torch().as_tensor( + int(local_energies.numel()), + dtype=energy_variance.dtype, + device=energy_variance.device, + ) + else: + ( + energy_mean, + energy_variance, + _, + _, + effective_sample_size, + ) = _weighted_energy_statistics(local_energies, importance_weights) + cache_snapshot = _cache_profile_snapshot(self.model) if profile else None sr_result = None + sr_elapsed = 0.0 + refresh_elapsed = 0.0 if sr: + sr_start = time.perf_counter() log_derivatives = torch_log_derivative_matrix( self.model, - self.configs, + batch_configs, amplitude_floor=amplitude_floor, + complex_parameter_mode=sr_parameter_mode, + derivative_backend=derivative_backend, ) sr_result = solve_torch_sr( log_derivatives, local_energies, + sample_weights=importance_weights, method=sr_method, diag_shift=sr_diag_shift, + parameter_mode=sr_parameter_mode, + step=self._sr_step, + pinv_rtol=sr_pinv_rtol, + momentum=sr_momentum, + previous_direction=self._sr_previous_direction, ) apply_torch_sr_update( self.model, sr_result.direction, learning_rate=learning_rate, + parameter_mode=sr_parameter_mode, ) + self._sr_previous_direction = sr_result.direction.detach().clone() + self._sr_step += 1 + sr_elapsed = time.perf_counter() - sr_start + refresh_start = time.perf_counter() self.refresh_amplitudes() + refresh_elapsed = time.perf_counter() - refresh_start + + profile_data = None + if profile: + total_elapsed = time.perf_counter() - total_start + profile_data = { + "sampling_seconds": sampling_elapsed, + "connection_seconds": connection_elapsed, + "local_estimator_seconds": local_elapsed, + "sr_seconds": sr_elapsed, + "refresh_seconds": refresh_elapsed, + "postprocess_seconds": max( + total_elapsed + - sampling_elapsed + - connection_elapsed + - local_elapsed + - sr_elapsed + - refresh_elapsed, + 0.0, + ), + "total_seconds": total_elapsed, + "cache": cache_snapshot, + } return TorchVMCStepResult( - configs=self.configs, - amplitudes=self.amplitudes, + configs=batch_configs if samples is not None else self.configs, + amplitudes=batch_amplitudes if samples is not None else self.amplitudes, local_energies=local_energies, energy_mean=energy_mean, energy_variance=energy_variance, @@ -1882,12 +10065,843 @@ def step( n_proposed=0 if sample is None else sample.n_proposed, n_accepted=0 if sample is None else sample.n_accepted, sr=sr_result, + profile=profile_data, + proposal_stats=None if sample is None else sample.proposal_stats, + importance_weights=importance_weights, + effective_sample_size=effective_sample_size, + sample_source=sample_source, ) - def run(self, n_steps, **step_kwargs): - """Run ``n_steps`` VMC steps and return their result records.""" + def optimize( + self, + n_steps=None, + *, + optimization=None, + progress=None, + progress_desc="Torch VMC optimization", + **step_kwargs, + ): + """Run repeated VMC/SR updates and return one result per update. + + Set ``progress=True`` for an internal notebook/terminal progress bar. + Its live postfix reports energy per site, Metropolis acceptance, the + optional no-op rate, and the SR solver. ``step_kwargs`` are forwarded + unchanged to :meth:`step`. + """ + if optimization is not None: + from .api import OptimizationConfig + if not isinstance(optimization, OptimizationConfig): + raise TypeError("optimization must be an OptimizationConfig or None.") + if n_steps is not None and n_steps != optimization.n_steps: + raise ValueError("n_steps conflicts with optimization.n_steps.") + n_steps = optimization.n_steps + if progress is None: + progress = optimization.progress + if optimization.energy_shift != 0.0 or optimization.per_site is not None: + from .api import BackendCapabilityWarning + warnings.warn( + "TorchVMCDriver.optimize returns raw energy tensors and " + "does not apply OptimizationConfig.energy_shift/per_site.", + BackendCapabilityWarning, + stacklevel=2, + ) + step_kwargs.setdefault("learning_rate", optimization.learning_rate) + if optimization.method == "sgd": + step_kwargs.setdefault("sr", False) + else: + step_kwargs.setdefault("sr", True) + step_kwargs.setdefault("sr_diag_shift", optimization.diag_shift) + step_kwargs.setdefault( + "sr_method", + "minsr" if optimization.method == "minsr" else "auto", + ) + mode = str(optimization.sr_mode).replace("_", "-").lower() + if mode == "real": + mode = "real-imag" + elif mode in {"complex", "holomorphic-complex"}: + mode = "holomorphic" + step_kwargs.setdefault("sr_parameter_mode", mode) + if n_steps is None: + raise TypeError("n_steps is required unless optimization is supplied.") + if progress is None: + progress = False n_steps = _check_positive_int("n_steps", n_steps) - return [self.step(**step_kwargs) for _ in range(n_steps)] + bar = _make_progress( + progress, + total=n_steps, + desc=progress_desc, + unit="step", + ) + results = [] + try: + for _ in range(n_steps): + result = self.step(**step_kwargs) + results.append(result) + if bar is not None: + bar.update(1) + _set_vmc_progress_postfix( + bar, + result, + n_sites=self.n_sites, + ) + finally: + if bar is not None: + bar.close() + return results + + def run(self, n_steps, *, progress=False, **step_kwargs): + """Compatibility alias for :meth:`optimize`.""" + return self.optimize( + n_steps, + progress=progress, + **step_kwargs, + ) + + +def _fermion_sector_from_configs(configs, metadata): + """Return the unique conserved sector represented by ``configs``.""" + if not metadata.spinful: + torch = _require_torch() + occupations = metadata.encoding.decode(configs) + values = occupations.sum(dim=-1) + if metadata.symmetry == "Z2": + values = values % 2 + sector = int(values[0].item()) + if not bool(torch.all(values == values[0])): + raise ValueError("All initial walkers must have the same spinless sector.") + return sector + n_up, n_down = count_spinful_particles( + configs, + encoding=metadata.encoding, + ) + if metadata.symmetry == "U1": + values = n_up + n_down + sector = int(values[0].item()) + elif metadata.symmetry == "Z2": + values = n_up + n_down + sector = int((values[0] % 2).item()) + elif metadata.symmetry == "Z2Z2": + values = list( + zip( + (n_up % 2).detach().cpu().tolist(), + (n_down % 2).detach().cpu().tolist(), + ) + ) + sector = tuple(int(value) for value in values[0]) + else: + values = list( + zip(n_up.detach().cpu().tolist(), n_down.detach().cpu().tolist()) + ) + sector = tuple(int(value) for value in values[0]) + if metadata.symmetry == "U1": + if not bool((values == sector).all()): + raise ValueError("All initial walkers must have the same U1 sector.") + elif metadata.symmetry == "Z2": + if not bool(((values % 2) == sector).all()): + raise ValueError("All initial walkers must have the same Z2 sector.") + elif metadata.symmetry == "Z2Z2": + if any(tuple(value) != sector for value in values): + raise ValueError("All initial walkers must have the same Z2Z2 sector.") + elif any(tuple(value) != sector for value in values): + raise ValueError("All initial walkers must have the same U1U1 sector.") + return sector + + +def _fermion_sector_counts(sector, symmetry, n_sites): + """Choose spin-resolved counts for a spinful initial configuration batch.""" + if symmetry == "U1U1": + return tuple(sector) + if symmetry == "Z2": + total = n_sites if n_sites % 2 == sector else n_sites - 1 + n_up = total // 2 + return n_up, total - n_up + if symmetry == "Z2Z2": + n_up = n_sites if n_sites % 2 == int(sector[0]) else n_sites - 1 + n_down = n_sites if n_sites % 2 == int(sector[1]) else n_sites - 1 + return n_up, n_down + total = int(sector) + n_up = total // 2 + return n_up, total - n_up + + +def _fermion_sector_mask(configs, metadata): + """Return a boolean mask selecting walkers in ``metadata.sector``.""" + if not metadata.spinful: + values = metadata.encoding.decode(configs).sum(dim=-1) + if metadata.symmetry == "Z2": + values = values % 2 + return values == metadata.sector + n_up, n_down = count_spinful_particles( + configs, + encoding=metadata.encoding, + ) + if metadata.symmetry == "U1": + return n_up + n_down == metadata.sector + if metadata.symmetry == "Z2": + return (n_up + n_down) % 2 == metadata.sector + if metadata.symmetry == "Z2Z2": + return (n_up % 2 == metadata.sector[0]) & ( + n_down % 2 == metadata.sector[1] + ) + return (n_up == metadata.sector[0]) & (n_down == metadata.sector[1]) + + +def _model_device(model, device=None): + torch = _require_torch() + if device is not None: + return torch.device(device) + try: + return next(model.parameters()).device + except (AttributeError, StopIteration, TypeError): + return torch.device("cpu") + + +def _initial_fermion_walkers( + model, + metadata, + n_walkers, + *, + device, + generator=None, + amplitude_floor=0.0, + max_attempts=32, + max_states=100_000, +): + """Find nonzero PEPS amplitudes inside the requested conserved sector.""" + torch = _require_torch() + n_walkers = _check_positive_int("n_walkers", n_walkers) + max_attempts = _check_positive_int("init_max_attempts", max_attempts) + max_states = _check_positive_int("init_max_states", max_states) + if amplitude_floor < 0: + raise ValueError("amplitude_floor must be non-negative.") + + if metadata.spinful: + n_up, n_down = _fermion_sector_counts( + metadata.sector, + metadata.symmetry, + metadata.n_sites, + ) + else: + n_particles = int(metadata.sector) + kept_configs = [] + kept_amplitudes = [] + + def keep(candidate): + with torch.no_grad(): + candidate_amplitudes = _call_amplitude_fn(model, candidate) + valid = ( + torch.isfinite(candidate_amplitudes.abs()) + & (candidate_amplitudes.abs() > float(amplitude_floor)) + ) + if bool(torch.any(valid)): + kept_configs.append(candidate[valid]) + kept_amplitudes.append(candidate_amplitudes[valid]) + return int(valid.sum().item()) + + n_kept = 0 + for _ in range(max_attempts): + if metadata.spinful: + candidate = random_spinful_configs( + n_walkers, + metadata.n_sites, + n_up, + n_down, + encoding=metadata.encoding, + device=device, + generator=generator, + ) + else: + candidate = random_spin_configs( + n_walkers, + metadata.n_sites, + n_particles, + device=device, + generator=generator, + ) + n_kept += keep(candidate) + if n_kept >= n_walkers: + break + + dense_states = (4 if metadata.spinful else 2) ** metadata.n_sites + if n_kept == 0 and dense_states <= max_states: + candidate = torch.as_tensor( + tuple( + product( + range(4 if metadata.spinful else 2), + repeat=metadata.n_sites, + ) + ), + dtype=torch.long, + device=device, + ) + candidate = candidate[_fermion_sector_mask(candidate, metadata)] + if candidate.numel(): + n_kept += keep(candidate) + + if n_kept == 0: + raise RuntimeError( + "Could not find a nonzero PEPS amplitude in the requested Fermion " + "sector. Pass valid configs or increase init_max_attempts." + ) + + configs = torch.cat(kept_configs, dim=0) + amplitudes = torch.cat(kept_amplitudes, dim=0) + if configs.shape[0] < n_walkers: + choice = torch.randint( + configs.shape[0], + (n_walkers,), + device=device, + generator=generator, + ) + else: + # The first candidates are often the same ordered sector pattern. + # Randomly selecting without replacement prevents all chains from + # starting at one configuration when the PEPS has broad support. + choice = torch.randperm( + configs.shape[0], + device=device, + generator=generator, + )[:n_walkers] + return configs[choice], amplitudes[choice] + + +class TorchFermionVMC(TorchVMCDriver): + """Automatic native spinful Fermion VMC around a Quimb PEPS. + + The constructor derives the PEPS lattice, physical dimension, local basis, + charge sector, periodic axes, and default native Hamiltonian. When explicit + ``terms`` are supplied, their two-site supports are added to the Metropolis + proposal graph so long-range terms remain traversable. ``fermion`` can be omitted + when explicit ``terms`` are supplied and the PEPS exposes Symmray symmetry + metadata. The lower-level + :class:`TorchVMCDriver` remains available when callers need full manual + control over configurations or connection functions. + """ + + def __init__( + self, + peps, + fermion=None, + terms=None, + *, + hamiltonian=None, + observables=None, + edges=None, + pbc=None, + site_order=None, + sector=None, + configs=None, + n_walkers=128, + contraction="boundary", + chi=4, + cutoff=None, + contraction_opts=None, + dtype=None, + device=None, + proposal=None, + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + graded_torch=False, + amplitude_batching="auto", + encoding=None, + chunk_size=None, + compile_kernels=False, + log_amplitude_fn=None, + proposal_batching="auto", + proposal_vmap_min_batch=8, + generator=None, + seed=None, + amplitude_floor=0.0, + 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; " + "terms is a compatibility alias for hamiltonian." + ) + 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 + metadata = _infer_torch_fermion_metadata( + peps, + fermion, + sector=sector, + edges=edges, + pbc=pbc, + site_order=site_order, + terms=terms, + ) + if encoding is not None and encoding != metadata.encoding: + raise ValueError( + "The supplied encoding does not match the native Fermion local " + "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: + 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 + elif isinstance(terms, OperatorSum): + hamiltonian = terms + terms = compile_operator_sum_torch( + terms, + fermion=fermion, + site_order=metadata.site_order, + ) + else: + hamiltonian = terms + terms = _normalize_terms_site_labels(terms, metadata.site_order) + + if configs is None: + if metadata.sector is None: + raise ValueError( + "Could not infer the PEPS charge sector. Pass sector=... or " + "provide initial configs in the target sector." + ) + configs, amplitudes = _initial_fermion_walkers( + model, + metadata, + n_walkers, + device=model_device, + generator=generator, + amplitude_floor=amplitude_floor, + max_attempts=init_max_attempts, + max_states=init_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]}." + ) + 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}." + ) + if metadata.sector is None: + metadata = replace(metadata, sector=actual_sector) + with torch.no_grad(): + amplitudes = _call_amplitude_fn(model, configs) + valid = ( + torch.isfinite(amplitudes.abs()) + & (amplitudes.abs() > float(amplitude_floor)) + ) + if not bool(torch.all(valid)): + raise ValueError( + "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, + 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, + encoding=metadata.encoding, + chunk_size=chunk_size, + compile_kernels=compile_kernels, + log_amplitude_fn=log_amplitude_fn, + generator=generator, + ) + + @property + def Lx(self): + return self.metadata.Lx + + @property + def Ly(self): + return self.metadata.Ly + + def make_bp_sampler( + self, + proposal_sampler=None, + *, + n_chains=None, + sample_kwargs=None, + bp_sampler_kwargs=None, + amplitude_floor=0.0, + max_init_attempts=32, + seed=None, + sampler_seed=None, + ): + """Create a symmetry-aware BP independence sampler from this PEPS.""" + if proposal_sampler is None: + from ..sampling import PepsBpSampler # pylint: disable=import-outside-toplevel + + proposal_sampler = PepsBpSampler( + self.peps, + encoding=self.metadata.encoding, + site_order=self.metadata.site_order, + sample_kwargs=bp_sampler_kwargs, + ) + return super().make_bp_sampler( + proposal_sampler, + n_chains=n_chains, + sample_kwargs=sample_kwargs, + symmetry=self.metadata.symmetry, + sector=self.metadata.sector, + encoding=self.metadata.encoding, + amplitude_floor=amplitude_floor, + max_init_attempts=max_init_attempts, + seed=seed, + sampler_seed=sampler_seed, + ) + + @property + def sector(self): + return self.metadata.sector + + def _compile_observables(self, observables): + """Compile supplemental observables without changing the Hamiltonian. + + ``TorchVMCDriver`` already owns the configured Hamiltonian connection + path. Keeping extra observables in a separate mapping lets the + backend-neutral façade measure energy and correlators from the same + samples, and avoids the historical ``observables=``/``terms=`` + ambiguity in this constructor. + """ + if observables is None: + return {} + try: + entries = tuple(observables.items()) + except AttributeError as exc: + raise TypeError("observables must be a mapping of names to operators.") from exc + + from .api import CompiledOperatorSum, OperatorSum + + compiled = {} + for name, value in entries: + if not isinstance(name, str) or not name: + raise ValueError("observable names must be non-empty strings.") + if isinstance(value, OperatorSum): + compiled[name] = compile_operator_sum_torch( + value, + fermion=self.fermion, + site_order=self.metadata.site_order, + ) + elif isinstance(value, CompiledOperatorSum): + if value.backend != "torch": + raise ValueError( + f"Observable {name!r} targets backend {value.backend!r}, not 'torch'." + ) + compiled[name] = value + else: + raw_terms = getattr(value, "terms", value) + compiled[name] = _normalize_terms_site_labels( + raw_terms, + self.metadata.site_order, + ) + return compiled + + +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) + if array.size != 1: + raise ValueError("Expected a scalar VMC result.") + return float(np.real(array.reshape(-1)[0])) + + +@dataclass(frozen=True) +class TorchVMCSetup: + """Backend-neutral façade over a native :class:`TorchFermionVMC`. + + The native driver deliberately retains its existing result classes and + detailed performance controls. This setup is the small portable surface: + it consumes shared configuration objects and returns common result + contracts while retaining every native value through ``.native``. + """ + + driver: TorchFermionVMC + problem: Any + sampling: Any = None + + @property + def backend(self): + """Name of the numerical backend behind this setup.""" + return "torch" + + @property + def native(self): + """Return the native stateful driver for backend-specific controls.""" + return self.driver + + @property + def n_sites(self): + return self.driver.n_sites + + @property + def n_params(self): + return sum(parameter.numel() for parameter in self.driver.model.parameters()) + + def sample(self, sampling=None): + """Collect samples as backend-neutral :class:`VMCSamples`.""" + sampling = self.sampling if sampling is None else sampling + native = ( + self.driver.sample() + if sampling is None + else self.driver.sample(sampling=sampling) + ) + return native.to_common() + + def _measurement_terms(self, observables): + if observables is None: + return dict(self.driver.observables) + try: + entries = dict(observables) + except (TypeError, ValueError) as exc: + raise TypeError("observables must be a mapping of names to operators.") from exc + return self.driver._compile_observables(entries) + + def measure( + self, + observables=None, + *, + sampling=None, + samples=None, + weights=None, + proposal_log_probs=None, + ): + """Measure energy and optional observables from one shared sample set. + + Passing ``samples`` avoids an additional Metropolis run. The supplied + batch may be a common :class:`VMCSamples`, native Torch samples, or a + configuration tensor. See :meth:`TorchVMCDriver.measure_samples` for + weighted and proposal-density semantics. + """ + from .api import VMCMeasurement + + if samples is not None and sampling is not None: + raise ValueError("Pass either sampling or samples, not both.") + if samples is None: + samples = self.sample(sampling) + native_samples = getattr(samples, "native", None) or samples + if weights is None: + weights = getattr(samples, "weights", None) + if proposal_log_probs is None: + proposal_log_probs = getattr(samples, "proposal_log_probs", None) + extra_terms = self._measurement_terms(observables) + if "energy" in extra_terms: + raise ValueError( + "'energy' is reserved for problem.hamiltonian; use a different " + "observable name." + ) + if extra_terms: + estimates = self.driver.measure_samples( + native_samples, + observables={"energy": None, **extra_terms}, + weights=weights, + proposal_log_probs=proposal_log_probs, + ) + energy = estimates["energy"] + else: + energy = self.driver.measure_samples( + native_samples, + weights=weights, + proposal_log_probs=proposal_log_probs, + ) + estimates = {"energy": energy} + return VMCMeasurement( + energy_mean=energy.energy_mean, + energy_variance=energy.energy_variance, + energy_stderr=energy.energy_stderr, + observables=estimates, + local_values=energy.local_energies, + effective_sample_size=energy.effective_sample_size, + diagnostics={ + "backend": self.backend, + "samples": samples, + "chain_diagnostics": energy.chain_diagnostics, + "acceptance_rate": energy.acceptance_rate, + }, + native=estimates, + ) + + def optimize(self, optimization=None, *, n_steps=None, **kwargs): + """Optimize and return a backend-neutral history. + + Display-only energy shifting and per-site scaling belong to the common + result object, not to the native Torch update loop. + """ + from .api import OptimizationConfig, VMCOptimizationResult + + if optimization is not None and not isinstance(optimization, OptimizationConfig): + raise TypeError("optimization must be an OptimizationConfig or None.") + supplied_samples = kwargs.get("samples") + if supplied_samples is not None: + if kwargs.get("weights") is None: + kwargs["weights"] = getattr(supplied_samples, "weights", None) + if kwargs.get("proposal_log_probs") is None: + kwargs["proposal_log_probs"] = getattr( + supplied_samples, + "proposal_log_probs", + None, + ) + if optimization is not None: + if n_steps is not None and n_steps != optimization.n_steps: + raise ValueError("n_steps conflicts with optimization.n_steps.") + native_config = replace( + optimization, + energy_shift=0.0, + per_site=None, + ) + history = self.driver.optimize( + optimization=native_config, + **kwargs, + ) + energy_shift = optimization.energy_shift + per_site = optimization.per_site + else: + if n_steps is None: + raise TypeError("n_steps is required unless optimization is supplied.") + history = self.driver.optimize(n_steps, **kwargs) + energy_shift = 0.0 + per_site = None + + energies = np.asarray( + [_vmc_result_scalar(result.energy_mean) for result in history], + dtype=float, + ) + variances = np.asarray( + [_vmc_result_scalar(result.energy_variance) for result in history], + dtype=float, + ) + errors = np.sqrt(np.maximum(variances, 0.0) / self.driver.n_walkers) + return VMCOptimizationResult( + steps=np.arange(1, len(history) + 1, dtype=int), + energies=energies, + errors=errors, + variances=variances, + energy_shift=energy_shift, + per_site=per_site, + diagnostics={ + "backend": self.backend, + "error_estimate": "naive walker standard error per update", + }, + native=tuple(history), + ) + + +def build_torch_vmc( + problem, + *, + fermion=None, + contraction=None, + sampling=None, + **kwargs, +): + """Build the portable Torch VMC façade from a :class:`VMCProblem`. + + This leaves :class:`TorchFermionVMC` untouched as the native integration + seam. The shared builder standardizes the problem, contraction, and chain + configuration without hiding Torch-specific options accepted via + ``**kwargs``. + """ + from .api import ContractionConfig, SamplingConfig, VMCProblem + + if not isinstance(problem, VMCProblem): + raise TypeError("problem must be a VMCProblem.") + if sampling is not None and not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if contraction is None: + contraction = ContractionConfig() + if "n_walkers" not in kwargs and sampling is not None: + kwargs["n_walkers"] = sampling.n_chains + if "proposal" not in kwargs and sampling is not None and sampling.proposal is not None: + kwargs["proposal"] = sampling.proposal + if "site_order" not in kwargs and problem.site_order is not None: + kwargs["site_order"] = problem.site_order + if sampling is not None: + if sampling.seed is not None and sampling.sampler_seed is not None: + raise ValueError("Pass either sampling.seed or sampling.sampler_seed, not both.") + if "seed" not in kwargs: + kwargs["seed"] = ( + sampling.seed + if sampling.seed is not None + else sampling.sampler_seed + ) + driver = TorchFermionVMC( + problem.peps, + fermion=fermion, + hamiltonian=problem.hamiltonian, + observables=problem.observables, + contraction=contraction, + **kwargs, + ) + return TorchVMCSetup(driver=driver, problem=problem, sampling=sampling) def random_spin_configs(n_walkers, n_sites, n_up, *, device=None, generator=None): @@ -1916,7 +10930,7 @@ def random_spinful_configs( ): """Generate spinful fermion configs with fixed ``N_up`` and ``N_down``.""" torch = _require_torch() - encoding = FermionSiteEncoding.symmray() if encoding is None else encoding + encoding = FermionSiteEncoding.vmc_torch() if encoding is None else encoding n_walkers = _check_positive_int("n_walkers", n_walkers) n_sites = _check_positive_int("n_sites", n_sites) if n_up < 0 or n_up > n_sites: diff --git a/tests/test_core_seed.py b/tests/test_core_seed.py index bf8368a..b0e44e3 100644 --- a/tests/test_core_seed.py +++ b/tests/test_core_seed.py @@ -72,6 +72,7 @@ def __init__(self, **kwargs): assert isinstance(out, DummyOpt) assert "seed" not in captured + assert captured["on_trial_error"] == "ignore" def test_build_optimizer_forwards_slicing_related_options(monkeypatch): diff --git a/tests/test_gate.py b/tests/test_gate.py index 68124fa..bd07d75 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 ps_to_3dpeps, ps_to_peps +from pepsy import hrs_to_peps, ps_to_3dpeps, ps_to_peps from pepsy.operators.gates import ( build_mpo_from_gates, build_pepo_from_gates, @@ -1231,7 +1231,7 @@ def test_gate_simple_2d_routes_swaps_with_real_torch_dtype(): """Internal SWAP tensors should match real torch PEPS dtype.""" torch = pytest.importorskip("torch") - peps = ps_to_peps(2, 3, dtype="float64", chi=2) + peps = hrs_to_peps(2, 3, dtype="float64", chi=2) peps.apply_to_arrays(lambda x: torch.as_tensor(x, dtype=torch.float64)) gate = torch.eye(4, dtype=torch.float64).reshape(2, 2, 2, 2) gauges = {} diff --git a/tests/test_optimize_energy_mps.py b/tests/test_optimize_energy_mps.py index b63a84e..94becbd 100644 --- a/tests/test_optimize_energy_mps.py +++ b/tests/test_optimize_energy_mps.py @@ -63,6 +63,21 @@ def test_mps_energy_loss_calls_exact_local_expectation_with_expected_options(): assert kwargs["progbar"] is True +def test_mps_energy_accepts_explicit_terms_keyword(): + """The constructor should expose local terms without overloading `hamiltonian`.""" + state = _FakeMps(value=4.0) + terms = {"edge": object()} + + opt = MpsEnergyOptimizer(state, terms=terms, energy_per_site=False) + + assert opt.hamiltonian is terms + assert opt.terms is terms + assert opt.energy().energy == pytest.approx(4.0) + + with pytest.raises(TypeError, match="either hamiltonian or terms"): + MpsEnergyOptimizer(state, terms, terms=terms) + + def test_mps_energy_compute_kwargs_override_direct_progbar(): """compute_kwargs should still be able to override the convenience flag.""" state = _FakeMps() diff --git a/tests/test_optimize_energy_peps.py b/tests/test_optimize_energy_peps.py index a0da160..d8eafd2 100644 --- a/tests/test_optimize_energy_peps.py +++ b/tests/test_optimize_energy_peps.py @@ -7,6 +7,7 @@ import pepsy from pepsy.optimizers import EnergyEstimate, PepsEnergyOptimizer from pepsy.optimizers.energy.peps import PepsEnergyOptimizer as ModulePepsEnergyOptimizer +from pepsy.tensors import SymPEPS class _FakePeps: @@ -195,6 +196,56 @@ def test_peps_energy_exact_loss_matches_direct_quimb_exact_expectation(): assert complex(opt.loss()) == pytest.approx(complex(direct)) +def test_peps_energy_exact_native_symmray_terms_support_torch_gradient(): + """Native exact energy avoids dense reduced-density-matrix fusion.""" + torch = pytest.importorskip("torch") + pytest.importorskip("symmray") + + to_backend = pepsy.backend_torch(dtype=torch.float64) + occupations = { + (0, 0): (1, 0), + (0, 1): (0, 1), + (1, 0): (0, 1), + (1, 1): (1, 0), + } + state = SymPEPS.for_model( + "fermi_hubbard_u1u1", + 2, + 2, + bond_dim=2, + site_charge=pepsy.site_charge_from_occupations(occupations), + seed=25, + dtype="float64", + to_backend=to_backend, + ) + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=4.0, + to_backend=to_backend, + ) + hamiltonian = fermion.hamiltonian(state.edges) + expected = state.energy(hamiltonian) + + opt = PepsEnergyOptimizer( + state.peps, + hamiltonian.terms, + boundary_mode="exact", + energy_per_site=False, + contraction_opt="auto-hq", + ) + assert float(opt.loss()) == pytest.approx(float(expected)) + + tnopt = opt.make_tn_optimizer( + optimizer="L-BFGS-B", + autodiff_backend="torch", + progbar=False, + ) + _, gradient = tnopt.vectorized_value_and_grad(tnopt.vectorizer.vector.copy()) + assert np.isfinite(gradient).all() + + def test_peps_energy_make_tn_optimizer_and_optimize(monkeypatch): """TNOptimizer construction should receive terms as constants and update state.""" calls = [] diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index 0d5ec94..b0eba62 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -100,6 +100,91 @@ def test_mps_optimizer_accepts_svd_mode(): assert opt.mode == "svd" +def test_mps_optimizer_can_disable_infidelity_tracking(monkeypatch): + """Disabled tracking skips diagnostic target and retained-norm work.""" + p0 = qtn.MPS_computational_state("0000", dtype="complex128") + opt = py.MpsOptimizer( + p0, + gates=[(qu.CNOT(), (0, 3))], + chi=2, + mode="svd", + track_infidelity=False, + ) + + def fail_diagnostic(*_args, **_kwargs): + raise AssertionError("infidelity diagnostics should be disabled") + + monkeypatch.setattr(py.MpsOptimizer, "_raw_state_norm", fail_diagnostic) + monkeypatch.setattr( + py.MpsOptimizer, + "_canonical_span_norm", + fail_diagnostic, + ) + + out = opt.run(progbar=False, cutoff=1.0e-12) + + assert out.max_bond() <= 2 + assert opt.get_infidelities() == [0.0] + assert opt.get_infidelity_samples() == [] + + +def test_mps_optimizer_run_can_override_infidelity_tracking(): + """A run-level override should enable diagnostics without reconstruction.""" + p0 = qtn.MPS_computational_state("0000", dtype="complex128") + opt = py.MpsOptimizer( + p0, + gates=[(qu.CNOT(), (0, 3))], + chi=2, + mode="svd", + track_infidelity=False, + ) + + opt.run(progbar=False, cutoff=1.0e-12, track_infidelity=True) + + assert opt.track_infidelity is True + assert len(opt.get_infidelity_samples()) == 1 + + +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") + + backend = py.backend_torch(dtype=torch.float64, device="cpu") + fermion = py.Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=8.0, + dtype="float64", + ) + state = py.hrs_to_mps( + 4, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + chi=4, + seed=1, + dtype="float64", + cyclic=False, + ) + state.apply_to_arrays(backend) + fermion.to_backend = backend + hopping = fermion.hopping_gate(0.001, imaginary=True) + + optimizer = py.MpsOptimizer( + state, + gates=[(hopping, (0, 3))], + chi=4, + mode="su", + track_infidelity=False, + inplace=True, + ) + out = optimizer.run(progbar=False, cutoff=1.0e-10, non_unitary=True) + + assert type(out[0].data).__name__ == "U1U1FermionicArray" + assert out.max_bond() <= 4 + assert len(optimizer.gauges) == out.L - 1 + + 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") @@ -2174,6 +2259,54 @@ def counting(self, *args, **kwargs): assert calls[0]["info"] is opt.info_c +def test_mps_optimizer_expectation_converts_operator_to_state_backend(monkeypatch): + """The local expectation operator should pass through backend conversion.""" + opt = py.MpsOptimizer( + qtn.MPS_rand_state(4, 2, seed=8), gates=[], chi=4, mode="mpo" + ) + converted = [] + original = opt._to_state_backend + + def recording(array): + converted.append(array) + return original(array) + + monkeypatch.setattr(opt, "_to_state_backend", recording) + opt._state_expectation("Z", (1,)) # pylint: disable=protected-access + + assert len(converted) == 1 + assert converted[0].shape == (2, 2) + + +def test_mps_optimizer_local_expectation_uses_torch_state_backend(monkeypatch): + """Torch-backed control expectations should stay on the Torch backend.""" + torch = pytest.importorskip("torch") + vector = torch.tensor([1.0, 1.0], dtype=torch.complex128) + vector = vector / torch.linalg.vector_norm(vector) + opt = py.MpsOptimizer( + qtn.MPS_product_state([vector.clone() for _ in range(3)]), + gates=[], + chi=4, + mode="mpo", + ) + observed_operators = [] + original = qtn.MatrixProductState.local_expectation_canonical + + def recording(self, operator, *args, **kwargs): + observed_operators.append(operator) + return original(self, operator, *args, **kwargs) + + monkeypatch.setattr( + qtn.MatrixProductState, + "local_expectation_canonical", + recording, + ) + assert opt._state_expectation("Z", (1,)) == pytest.approx(0.0) + + assert isinstance(observed_operators[0], torch.Tensor) + assert all(isinstance(tensor.data, torch.Tensor) for tensor in opt.p.tensors) + + @pytest.mark.parametrize( ("pauli", "where"), [("X", (4,)), ("YZ", (1, 4))], diff --git a/tests/test_optimize_peps.py b/tests/test_optimize_peps.py index 4acd1e1..761f8ca 100644 --- a/tests/test_optimize_peps.py +++ b/tests/test_optimize_peps.py @@ -54,21 +54,33 @@ class _FakeSymmrayArray: shape = (2, 2) dtype = "complex128" + def __init__(self, backend="torch"): + self.backend = backend + _FakeSymmrayArray.__module__ = "symmray.fake" class _FakeTensor: - def __init__(self): - self.data = _FakeSymmrayArray() + def __init__(self, backend="torch"): + self.data = _FakeSymmrayArray(backend=backend) class SymmrayDummyState(DummyState): """Dummy state carrying Symmray-looking tensor data.""" - def __init__(self, bond=1, name="state"): + def __init__(self, bond=1, name="state", backend="torch"): super().__init__(bond=bond, name=name) - self.tensor_map = {"site": _FakeTensor()} + self.backend = backend + self.tensor_map = {"site": _FakeTensor(backend=backend)} + + def copy(self): + other = type(self)(self.bond, f"{self.name}.copy", backend=self.backend) + other.applied = list(self.applied) + other.normalized = self.normalized + other.normalize_kwargs = [dict(opts) for opts in self.normalize_kwargs] + other.mangled = self.mangled + return other def _install_fake_gate(monkeypatch): @@ -216,6 +228,27 @@ def _fake_infidelity(*args, **kwargs): assert infidelity_calls[0]["mode_"] == "mps" +def test_peps_optimizer_rejects_non_torch_symmray_input(): + """Symmray PEPS cleanup is deliberately Torch/autograd-only.""" + with pytest.raises(TypeError, match="Torch-backed Symmray"): + PepsOptimizer( + SymmrayDummyState(bond=1, backend="numpy"), + [], + chi=3, + normalize_initial=False, + ) + + +def test_peps_optimizer_rejects_mixed_symmray_and_dense_inputs(): + """A malformed gate target must not silently drop Symmray structure.""" + with pytest.raises(TypeError, match="do not mix Symmray and dense"): + PepsOptimizer._require_torch_symmray_backend( + SymmrayDummyState(bond=1), + DummyState(bond=1), + role="state and target", + ) + + def test_peps_optimizer_runs_sweep_and_records_geometric_fidelity(monkeypatch): """Large warm-start infidelity should hand off to SweepOptimizer.""" gate_calls = _install_fake_gate(monkeypatch) @@ -954,6 +987,8 @@ def run(self): assert captured["kwargs"]["normalize_kwargs"]["mode_"] == "mps" assert captured["kwargs"]["normalize_kwargs"]["balance_bonds"] is False assert captured["optimize_kwargs"]["env_n_iter"] == 10 + assert captured["optimize_kwargs"]["optimizer"] == "nlopt" + assert captured["optimize_kwargs"]["optimizer_options"]["algorithm"] == "LD_LBFGS" def test_peps_optimizer_explicit_quimb_boundary_engine_forwards_options(monkeypatch): diff --git a/tests/test_optimize_sweep_plot.py b/tests/test_optimize_sweep_plot.py index e98bf4e..1a54a52 100644 --- a/tests/test_optimize_sweep_plot.py +++ b/tests/test_optimize_sweep_plot.py @@ -2,8 +2,11 @@ from types import SimpleNamespace +import numpy as np +import quimb as qu import quimb.tensor as qtn import pytest +import torch import pepsy.optimizers.sweep.optimizer as sweep_mod from pepsy.optimizers.sweep.environments import QuimbMpsBoundaryStore @@ -71,6 +74,31 @@ def test_scaled_overlap_fidelity_uses_mantissa_exponent_pairs(): assert fidelity == pytest.approx(1.0) +def test_scaled_overlap_fidelity_accepts_quimb_scalar_wrapper(): + """A stripped Quimb scalar should not be sent to nonexistent quimb.abs.""" + fidelity = SweepOptimizer._scaled_overlap_fidelity( + overlap=(qu.qarray(1.0 + 2.0j), 0.0), + norm=(qu.qarray(5.0), 0.0), + target_norm=(qu.qarray(1.0), 0.0), + ) + + assert float(fidelity) == pytest.approx(1.0) + + +def test_scaled_overlap_fidelity_unwraps_scalar_tensor_and_preserves_autograd(): + """Zero-index Quimb Tensor wrappers must expose their Torch scalar data.""" + overlap = torch.tensor(2.0, dtype=torch.float64, requires_grad=True) + fidelity = SweepOptimizer._scaled_overlap_fidelity( + overlap=(qtn.Tensor(data=overlap, inds=()), 0.0), + norm=(qtn.Tensor(data=overlap * overlap, inds=()), 0.0), + target_norm=(qtn.Tensor(data=torch.tensor(1.0), inds=()), 0.0), + ) + fidelity.backward() + + assert fidelity.item() == pytest.approx(1.0) + assert overlap.grad.item() == pytest.approx(0.0) + + def test_set_target_norm_accepts_scaled_pair(): """target_norm can be stored as a stripped ``(mantissa, exponent)`` pair.""" opt = object.__new__(SweepOptimizer) @@ -1483,22 +1511,145 @@ def test_sweep_quimb_mps_refreshes_real_quimb_environments(): assert opt.bdy_overlap.update_count == 1 -def test_sweep_quimb_mps_half_sweep_refreshes_per_slice(monkeypatch): - """Quimb sweep branch should rebuild environments before each slice.""" +@pytest.mark.parametrize( + ("axis", "update_side", "seed_index"), + ( + ("y", "left", 0), + ("y", "right", 2), + ("x", "left", 0), + ("x", "right", 1), + ), +) +def test_sweep_quimb_mps_incremental_boundaries_match_fresh_environments( + axis, + update_side, + seed_index, +): + """Directional cache seeds must agree with a full Quimb recomputation.""" + peps = qtn.PEPS.rand(Lx=2, Ly=3, bond_dim=2, seed=41, dtype="complex128") + target = qtn.PEPS.rand(Lx=2, Ly=3, bond_dim=2, seed=42, dtype="complex128") + opt = SweepOptimizer( + peps, + target, + chi=8, + boundary_engine="quimb-mps", + renormalize_state=False, + ) + norm_tn, _ = opt._prepare_current_double_layers() + cached = QuimbMpsBoundaryStore(chi=8) + fresh = QuimbMpsBoundaryStore(chi=8) + + cached.start_sweep(norm_tn, axis, update_side) + cached.advance_sweep( + norm_tn, + seed_index, + axis=axis, + update_side=update_side, + ) + fresh.update_axis(norm_tn, axis) + + for key, cached_env in cached.envs.items(): + fresh_env = fresh.envs[key] + if not cached_env.tensor_map: + assert not fresh_env.tensor_map + continue + np.testing.assert_allclose( + cached_env.contract(all).data, + fresh_env.contract(all).data, + rtol=1.0e-10, + atol=1.0e-11, + ) + + +def test_sweep_quimb_mps_numpy_slice_uses_finite_differences_without_backend_leak(): + """Dense NumPy Quimb sweeps must not inject Torch slice data into environments.""" + pytest.importorskip("torch") + peps = qtn.PEPS.rand(Lx=2, Ly=2, bond_dim=1, seed=31, dtype="complex128") + target = qtn.PEPS.rand(Lx=2, Ly=2, bond_dim=1, seed=32, dtype="complex128") + opt = SweepOptimizer( + peps, + target, + chi=4, + boundary_engine="quimb-mps", + renormalize_state=False, + ) + norm_tn, overlap_tn = opt._prepare_current_double_layers() + opt.bdy.start_sweep(norm_tn, "y", "left") + opt.bdy_overlap.start_sweep(overlap_tn, "y", "left") + + run = opt._optimize_axis_slice_with_current_env( + 0, + axis="y", + solver="torch-adam", + solver_options={"n_steps": 1, "fd_method": "forward", "fd_eps": 1.0e-4}, + ) + + assert np.isfinite(run["loss_initial"]) + assert np.isfinite(run["loss_final"]) + assert all( + type(tensor.data).__module__.split(".", 1)[0] == "numpy" + for tensor in opt.state.tensor_map.values() + ) + + +def test_sweep_quimb_mps_numpy_round_trip_keeps_boundaries_and_state_numpy(): + """The former NumPy/Torch boundary crash must stay fixed across a round trip.""" + pytest.importorskip("torch") + state = qtn.PEPS.rand(Lx=2, Ly=2, bond_dim=1, seed=81, dtype="complex128") + target = qtn.PEPS.rand(Lx=2, Ly=2, bond_dim=1, seed=82, dtype="complex128") + target.mangle_inner_() + opt = SweepOptimizer( + state, + target, + chi=4, + boundary_engine="quimb-mps", + renormalize_state=False, + ) + + runs = opt.optimize_axis( + "y", + n_round_trips=1, + solver="torch-adam", + solver_options={"n_steps": 1, "fd_method": "forward", "fd_eps": 1.0e-4}, + renormalize=False, + ) + + assert [(run["sweep"], run["index"]) for run in runs] == [ + ("forward", 0), + ("forward", 1), + ("backward", 0), + ("forward", 1), + ] + assert all( + type(tensor.data).__module__.split(".", 1)[0] == "numpy" + for tensor in opt.state.tensor_map.values() + ) + + +def test_sweep_quimb_mps_half_sweep_advances_cached_boundary(monkeypatch): + """Quimb sweep branch should build once and move its boundary per slice.""" + + class _FakeStore: + def __init__(self): + self.norm = 1.0 + self.mps_b = {} + self.starts = [] + self.advances = [] + + def start_sweep(self, tn, axis, update_side, *, progress=False): + self.starts.append((tn, axis, update_side, progress)) + + def advance_sweep(self, tn, index, *, axis=None, update_side=None): + self.advances.append((tn, index, axis, update_side)) + opt = object.__new__(SweepOptimizer) opt.boundary_engine = "quimb-mps" - opt.bdy = SimpleNamespace(norm=1.0) - opt.bdy_overlap = SimpleNamespace(norm=1.0) - - refreshes = [] + opt.bdy = _FakeStore() + opt.bdy_overlap = _FakeStore() def _fake_prepare(): return object(), object() - def _fake_refresh(norm_tn, overlap_tn, axis, *, progress=False): - refreshes.append((norm_tn, overlap_tn, axis, progress)) - return {"norm": None, "overlap": None} - def _fake_optimize(index, *, axis, solver, solver_options): return { "axis": axis, @@ -1507,13 +1658,13 @@ def _fake_optimize(index, *, axis, solver, solver_options): "history": [0.2, 0.1 + index], } - def _boom_make_comp(*args, **kwargs): # pylint: disable=unused-argument - raise AssertionError("CompBdy path should not run for quimb-mps") - monkeypatch.setattr(opt, "_prepare_current_double_layers", _fake_prepare) - monkeypatch.setattr(opt, "_refresh_quimb_axis_boundaries", _fake_refresh) monkeypatch.setattr(opt, "_optimize_axis_slice_with_current_env", _fake_optimize) - monkeypatch.setattr(opt, "_make_comp_pair", _boom_make_comp) + monkeypatch.setattr( + opt, + "_update_quimb_double_layer_slice", + lambda *args: None, + ) runs = SweepOptimizer._run_axis_half_sweep( opt, @@ -1526,5 +1677,70 @@ def _boom_make_comp(*args, **kwargs): # pylint: disable=unused-argument ) assert [run["index"] for run in runs] == [0, 1] - assert len(refreshes) == 2 - assert all(item[2] == "y" for item in refreshes) + assert len(opt.bdy.starts) == 1 + assert len(opt.bdy_overlap.starts) == 1 + assert len(opt.bdy.advances) == 2 + assert len(opt.bdy_overlap.advances) == 2 + assert all(item[1] == index for item, index in zip(opt.bdy.advances, (0, 1))) + + +def test_sweep_quimb_mps_backward_and_return_forward_seed_boundaries(monkeypatch): + """Round-trip directions must seed their already-updated endpoint once.""" + + class _FakeStore: + def __init__(self): + self.norm = 1.0 + self.mps_b = {} + self.starts = [] + self.advances = [] + + def start_sweep(self, tn, axis, update_side, *, progress=False): + self.starts.append((tn, axis, update_side, progress)) + + def advance_sweep(self, tn, index, *, axis=None, update_side=None): + self.advances.append((tn, index, axis, update_side)) + + opt = object.__new__(SweepOptimizer) + opt.boundary_engine = "quimb-mps" + opt.bdy = _FakeStore() + opt.bdy_overlap = _FakeStore() + opt.Lx = 2 + opt.Ly = 2 + + monkeypatch.setattr(opt, "_prepare_current_double_layers", lambda: (object(), object())) + monkeypatch.setattr( + opt, + "_optimize_axis_slice_with_current_env", + lambda index, **kwargs: { + "axis": kwargs["axis"], + "index": index, + "loss_final": 0.1, + "history": [0.1], + }, + ) + monkeypatch.setattr(opt, "_update_quimb_double_layer_slice", lambda *args: None) + + SweepOptimizer._run_axis_half_sweep( + opt, + range(0, -1, -1), + axis="y", + update_side="right", + sweep_name="backward", + solver="scipy", + solver_options=None, + ) + assert [item[1] for item in opt.bdy.advances] == [1, 0] + + opt.bdy.advances.clear() + opt.bdy_overlap.advances.clear() + SweepOptimizer._run_axis_half_sweep( + opt, + range(1, 2), + axis="y", + update_side="left", + sweep_name="return-forward", + solver="scipy", + solver_options=None, + ) + assert [item[1] for item in opt.bdy.advances] == [0, 1] + assert [call[2] for call in opt.bdy.starts] == ["right", "left"] diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index e6d9af0..0ef4dca 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1,7 +1,12 @@ """Tests for the tree-tensor-network gate simulator (:class:`TreeOptimizer`).""" +import sys +import types + import numpy as np import pytest +import quimb.tensor as qtn +import pepsy from pepsy.optimizers.tree import ( TreeLayoutFinder, @@ -53,6 +58,33 @@ def _random_stream(n, ngates, rng, two_qubit_frac=0.5): return stream +def _two_branch_flip_submpo(*, L, sites, targets, w0=0.7, w1=0.3): + """Return ``w0 * I + w1 * prod(X_targets)`` as a sparse-site MPO.""" + eye = np.eye(2, dtype=complex) + flip = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + sites = tuple(sites) + targets = set(targets) + branch0 = [eye.copy() for _site in sites] + branch1 = [flip.copy() if site in targets else eye.copy() for site in sites] + branch0[0] *= w0 + branch1[0] *= w1 + mpo0 = qtn.MPO_product_operator( + branch0, + sites=sites, + L=L, + upper_ind_id="k{}", + lower_ind_id="b{}", + ) + mpo1 = qtn.MPO_product_operator( + branch1, + sites=sites, + L=L, + upper_ind_id="k{}", + lower_ind_id="b{}", + ) + return mpo0.add_MPO(mpo1) + + def _exact_state(stream, n): psi = np.zeros(2**n, dtype=complex) psi[0] = 1.0 @@ -94,6 +126,109 @@ def test_tree_matches_statevector(n, seed): assert _fidelity(psi, opt.to_dense()) > 1 - 1e-8 +def test_tree_two_site_direct_and_mpo_modes_agree(): + """Dense direct threading and gate-to-MPO routing are equivalent.""" + rng = np.random.default_rng(918) + n = 6 + stream = _random_stream(n, 24, rng) + exact = _exact_state(stream, n) + direct = TreeOptimizer( + stream, n=n, chi=128, cutoff=0.0, mode="direct", + ) + mpo = TreeOptimizer( + stream, n=n, chi=128, cutoff=0.0, mode="mpo", + ) + + assert _fidelity(direct.to_dense(), mpo.to_dense()) > 1 - 1e-10 + assert _fidelity(direct.to_dense(), exact) > 1 - 1e-9 + assert _fidelity(mpo.to_dense(), exact) > 1 - 1e-9 + + +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) + theta = 1.0e-7 + gate = ( + np.cos(theta) * np.eye(4, dtype=complex) + - 1.0j * np.sin(theta) * np.kron(x, x) + ) + direct = TreeOptimizer( + [(gate, (0, 1))], n=2, chi=4, cutoff=0.0, mode="direct", + ) + mpo = TreeOptimizer( + [(gate, (0, 1))], n=2, chi=4, cutoff=0.0, mode="mpo", + ) + expected = np.array([np.cos(theta), 0.0, 0.0, -1.0j * np.sin(theta)]) + + np.testing.assert_allclose(direct.to_dense(), expected, atol=1e-13) + np.testing.assert_allclose(mpo.to_dense(), expected, atol=1e-13) + np.testing.assert_allclose(mpo.to_dense(), direct.to_dense(), atol=1e-13) + + +def test_tree_mpo_and_direct_share_path_compression_diagnostics(): + """Two-site MPO mode uses the same routed-factor kernel as direct mode.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], + [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex, + ) + direct = TreeOptimizer( + [(cnot, (0, 3))], n=4, chi=1, cutoff=0.0, mode="direct", + track_truncation=True, + ) + mpo = TreeOptimizer( + [(cnot, (0, 3))], n=4, chi=1, cutoff=0.0, mode="mpo", + track_truncation=True, + ) + + fields = ("kind", "edge", "before_bond", "after_bond", "max_bond", "cutoff") + direct_trace = [ + tuple(event[field] for field in fields) + for event in direct.truncation_history + ] + mpo_trace = [ + tuple(event[field] for field in fields) + for event in mpo.truncation_history + ] + assert mpo_trace == direct_trace + assert all(event["max_bond"] == 1 for event in mpo.truncation_history) + + +def test_tree_multisite_submpo_qr_routes_before_one_subtree_sweep(): + """A 3-site MPO transports its virtual legs without routing SVDs.""" + gate = _rand_unitary(3, np.random.default_rng(51)) + mpo = qtn.MatrixProductOperator.from_dense( + gate.reshape((2,) * 6), + dims=(2, 2, 2), + sites=(0, 2, 4), + L=5, + max_bond=None, + cutoff=0.0, + ) + opt = TreeOptimizer( + None, n=5, chi=1, cutoff=0.0, track_truncation=True, run=False, + ) + opt.apply_submpo(mpo, (0, 2, 4)) + + assert opt.truncation_history + assert all(event["kind"] != "split" for event in opt.truncation_history) + assert all(event["max_bond"] == 1 for event in opt.truncation_history) + + +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) + assert opt.mode == "direct" + opt.run(mode="mpo") + assert opt.mode == "mpo" + # Existing shared-front-end spelling remains a deprecated no-op. + with pytest.warns(DeprecationWarning, match="deprecated no-op"): + opt.run(mode="tree") + assert opt.mode == "mpo" + with pytest.warns(DeprecationWarning, match="two_site_mode"): + legacy = TreeOptimizer(None, n=2, two_site_mode="direct", run=False) + assert legacy.mode == "direct" + + def test_single_qubit_stream(): """A one-qubit tree replays single-qubit gates correctly.""" rng = np.random.default_rng(3) @@ -103,22 +238,9 @@ def test_single_qubit_stream(): assert _fidelity(psi, opt.to_dense()) > 1 - 1e-10 -def test_local_expectation_matches_exact(): - """Single-site expectations from the tree match the dense state.""" - rng = np.random.default_rng(4) - n = 6 - stream = _random_stream(n, 40, rng) - opt = TreeOptimizer(stream, n=n, chi=128) - psi = _exact_state(stream, n) - psi /= np.linalg.norm(psi) - - z = np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex) - x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) - for q in range(n): - for op in (z, x): - exact = np.vdot(psi, _sv_apply_1q(psi, op, q, n)).real - got = opt.local_expectation(op, q).real - assert abs(got - exact) < 1e-8 +def test_tree_optimizer_has_no_local_expectation_api(): + """Observable contraction is owned by the TTN state, not its optimizer.""" + assert not hasattr(TreeOptimizer(None, n=1, run=False), "local_expectation") def test_chi_truncation_caps_bond(): @@ -131,6 +253,50 @@ def test_chi_truncation_caps_bond(): assert opt.max_bond() <= chi +def test_tree_truncation_infidelity_compatibility_trace(): + """Tracked tree compression exposes the MPS-style diagnostic readout.""" + 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, + track_truncation=True, + ) + + assert opt.get_infidelities()[0] == 0.0 + assert len(opt.get_infidelities()) >= 2 + assert len(opt.get_infidelity_samples()) == len(opt.get_infidelities()) - 1 + assert 0.0 <= opt.get_infidelities()[-1] <= 1.0 + assert opt.get_normalizations() == [] + + +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) + 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 len(opt.get_normalizations()) == 2 + with pytest.raises(ValueError, match="non_unitary"): + opt.run(normalize_every=True) + + +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) + assert opt.qubits == [0, 1, 2, 3] + assert opt.logical_order == [0, 1, 2, 3] + assert [opt.logical_site(i) for i in range(4)] == [0, 1, 2, 3] + assert [opt.position(i) for i in range(4)] == [0, 1, 2, 3] + sample = np.array([0, 1, 0, 1]) + assert np.array_equal(opt.remap_sample(sample), sample) + assert opt.restore_qubit_order() is opt.tn + + @pytest.mark.parametrize("seed", [0, 1, 2, 3, 4]) def test_truncated_fidelity_improves_with_chi(seed): """Threading the whole gate before truncating yields high truncated fidelity. @@ -180,11 +346,11 @@ def test_user_supplied_plan_runs(): def test_layout_finder_builds_valid_tree(): - """The layout finder returns a rooted binary tree over all qubits.""" + """With max_arity=2 the finder returns a rooted binary tree over all qubits.""" rng = np.random.default_rng(8) n = 8 stream = _random_stream(n, 60, rng) - plan = TreeLayoutFinder(stream, n=n).run() + plan = TreeLayoutFinder(stream, n=n, max_arity=2).run() assert plan.n == n assert set(plan.leaf_of_qubit) == set(range(n)) # every internal node has exactly two children @@ -214,6 +380,478 @@ def test_quality_layout_not_worse_than_balanced(): assert finder.score(quality) <= finder.score(balanced) +def test_congestion_layout_uses_operator_schmidt_edge_load(): + """The load-aware diagnostic predicts the product of crossed ranks.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + rng = np.random.default_rng(101) + generic = rng.standard_normal((4, 4)) + 1j * rng.standard_normal((4, 4)) + finder = TreeLayoutFinder( + [(cnot, (0, 3)), (generic, (0, 3))], + n=4, + objective="congestion", + ) + plan = TreePlan.from_order(range(4), structure="balanced") + loads = finder.edge_loads(plan) + path = plan.node_path(plan.leaf_of_qubit[0], plan.leaf_of_qubit[3]) + path_edges = { + (u, v) if plan.parent.get(v) == u else (v, u) + for u, v in zip(path, path[1:]) + } + + assert path_edges + assert all(loads[edge] == pytest.approx(3.0) for edge in path_edges) + assert max(loads.values()) == pytest.approx(3.0) + report = finder.report(plan) + assert report["objective"] == "congestion" + assert report["peak_bond_growth"] == pytest.approx(8.0) + + +def test_tree_edge_loads_match_full_edge_reference(): + """Steiner-only edge scanning preserves the full congestion calculation.""" + rng = np.random.default_rng(109) + n = 9 + stream = _random_stream(n, 25, rng, two_qubit_frac=0.8) + finder = TreeLayoutFinder(stream, n=n, objective="congestion") + plan = TreePlan.from_order(range(n), structure="balanced") + + got = finder.edge_loads(plan) + below = plan.subtree_qubit_masks() + expected = {edge: 0.0 for edge in got} + for payload, support, event_type in zip( + finder.payloads, finder.supports, finder.event_types + ): + support = tuple(dict.fromkeys(support)) + if len(support) < 2 or event_type in { + "measure", "reset", "measure_reset", "cap", + }: + continue + 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)) + rank = finder._schmidt_rank(payload, support, left) + expected[edge] += np.log2(rank) + + assert got == pytest.approx(expected) + + +def test_tree_layout_reuses_dense_gate_schmidt_rank_across_labels(): + """One gate matrix needs one rank calculation for each wire partition.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + finder = TreeLayoutFinder( + [(cnot, (0, 1)), (cnot, (2, 3))], n=4, + objective="congestion", + ) + + assert finder._schmidt_rank(cnot, (0, 1), (0,)) == 2 + assert finder._schmidt_rank(cnot, (2, 3), (2,)) == 2 + assert len(finder._schmidt_rank_cache) == 1 + + +def test_optimizer_exposes_congestion_layout_objective(): + """TreeOptimizer can request the rank-aware automatic layout.""" + rng = np.random.default_rng(102) + stream = _random_stream(6, 20, rng, two_qubit_frac=0.8) + opt = TreeOptimizer( + stream, + n=6, + max_arity=2, + layout_objective="congestion", + layout_weight_mode="operator_schmidt", + run=False, + ) + + assert opt.layout_objective == "congestion" + assert opt.plan.n == 6 + assert opt.plan.is_binary() + + +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) + stream = _random_stream(8, 30, rng, two_qubit_frac=0.8) + finder = TreeLayoutFinder(stream, n=8, objective="congestion") + + recommendation = finder.recommend_arities((2, 3)) + assert recommendation["recommended_max_arity"] in (2, 3) + assert len(recommendation["candidates"]) == 2 + assert all("max_virtual_degree" in item + for item in recommendation["candidates"]) + report = finder.report(recommendation["plan"]) + assert report["arity_histogram"] + assert report["max_arity"] in (2, 3) + + +def test_layout_finder_layered_direct_block_size(): + """`layered` builds a valid fixed layered tree for a chosen block_size.""" + rng = np.random.default_rng(107) + stream = _random_stream(12, 40, rng, two_qubit_frac=0.7) + finder = TreeLayoutFinder(stream, n=12, weight_mode="operator_schmidt") + + plan = finder.layered(block_size=4) + assert isinstance(plan, TreePlan) + assert plan.n == 12 + # Top tensor is fixed ternary once there are at least three blocks. + assert len(plan.children[plan.root]) == 3 + # The blocking layer groups block_size leaves per blocking node. + assert 4 in {len(ch) for ch in plan.children.values() if ch} + # Direct build matches recommend_layered's plan for the same block_size. + recommended = finder.recommend_layered(block_sizes=(4,)) + assert plan.children == recommended["plan"].children + + +def test_layered_greedy_refinement_improves_without_changing_tree_shape(): + """Planning swaps leaf labels but preserves the immutable TTN topology.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + finder = TreeLayoutFinder([(cnot, (0, 2))], n=4, max_arity=2) + initial = TreePlan.build_layered(range(4), block_size=2) + + choice = finder.recommend_layered( + block_sizes=(2,), + order=range(4), + refine="greedy", + refine_budget=3, + ) + candidate = choice["candidates"][0] + refined = choice["plan"] + + assert choice["refine"] == "greedy" + assert refined.children == initial.children + assert finder.score(refined) < finder.score(initial) + assert refined.tree_distance(0, 2) == 2 + assert candidate["planning"]["refinement"]["accepted_moves"] >= 1 + + +def test_hybrid_layout_objective_reports_normalized_combined_cost(): + """Hybrid selection combines path and operator-Schmidt edge-load costs.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + finder = TreeLayoutFinder( + [(cnot, (0, 3)), (cnot, (1, 2))], + n=4, + objective="hybrid", + hybrid_weights={"path": 1.0, "max_edge_load": 2.0}, + ) + plan = finder.recommend_arities((2, 3))["plan"] + report = finder.report(plan) + + assert report["objective"] == "hybrid" + assert report["hybrid_weights"] == (1.0, 2.0, 0.0) + assert np.isfinite(report["hybrid_cost"]) + assert report["total_edge_load"] is not None + + +def test_layered_nevergrad_search_is_optional_and_pre_simulation(monkeypatch): + """Nevergrad refines only a returned plan and is optional at import time.""" + class FakeArray: + def __init__(self, *, init): + self.value = np.asarray(init) + + def set_bounds(self, _lower, _upper): + return self + + class FakeOptimizer: + def __init__(self, *, parametrization, budget): + self.parametrization = parametrization + self.budget = budget + + def minimize(self, loss): + loss(self.parametrization.value) + proposal = np.array([0.0, 2.0, 1.0, 3.0]) + loss(proposal) + return types.SimpleNamespace(value=proposal) + + fake_nevergrad = types.SimpleNamespace( + p=types.SimpleNamespace(Array=FakeArray), + optimizers=types.SimpleNamespace(registry={"OnePlusOne": FakeOptimizer}), + ) + monkeypatch.setitem(sys.modules, "nevergrad", fake_nevergrad) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + finder = TreeLayoutFinder([(cnot, (0, 2))], n=4, max_arity=2) + initial = TreePlan.build_layered(range(4), block_size=2) + + choice = finder.recommend_layered( + block_sizes=(2,), + order=range(4), + search="nevergrad", + search_budget=4, + seed=17, + ) + search_info = choice["candidates"][0]["planning"]["search"] + + assert choice["search"] == "nevergrad" + assert choice["plan"].children == initial.children + assert finder.score(choice["plan"]) <= finder.score(initial) + assert search_info["method"] == "nevergrad" + assert search_info["evaluations"] == 2 + + +def test_nevergrad_layout_search_explains_missing_optional_dependency(monkeypatch): + """The optional Nevergrad dependency fails with an actionable message.""" + monkeypatch.setitem(sys.modules, "nevergrad", None) + finder = TreeLayoutFinder([], n=4, max_arity=2) + + with pytest.raises(ImportError, match=r"pepsy\[layout\]"): + finder.recommend_layered( + block_sizes=(2,), order=range(4), search="nevergrad" + ) + + +def test_tree_qubit_order_honors_configured_dense_limit(monkeypatch): + """The layered spectral order uses the finder's dense-size policy.""" + import pepsy.optimizers.tree.layout as tree_layout + + seen = {} + + def ordered(sites, _weights, *, dense_max): + seen["dense_max"] = dense_max + return list(sites) + + monkeypatch.setattr(tree_layout, "_gate_stream_spectral_order", ordered) + finder = TreeLayoutFinder([], n=6, dense_max=17) + + assert finder.qubit_order() == list(range(6)) + assert seen["dense_max"] == 17 + + +def test_treeplan_max_bond_cut_is_structural(): + """`max_bond_cut` is the widest qubit bipartition, set by shape alone.""" + order = list(range(16)) + # block_size=3 groups into an even 2+2+2 ternary top -> widest cut is 6. + assert TreePlan.build_layered(order, block_size=3).max_bond_cut() == 6 + # block_size=4 forces a {1,1,2}-block top -> one 8-vs-8 cut. + assert TreePlan.build_layered(order, block_size=4).max_bond_cut() == 8 + # A single leaf has no bonds. + assert TreePlan.build_layered([0]).max_bond_cut() == 0 + + +def test_recommend_layered_chi_aware_prefers_exact_structure(): + """With ``chi`` set, the layered recommendation avoids chi-overflow bonds.""" + rng = np.random.default_rng(211) + stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) + finder = TreeLayoutFinder(stream, n=16, weight_mode="operator_schmidt") + + # chi-blind candidates always carry ``max_bond_cut`` but no chi fields. + blind = finder.recommend_layered(block_sizes=(3, 4)) + assert "chi" in blind and blind["chi"] is None + for c in blind["candidates"]: + assert "max_bond_cut" in c + assert "chi_overflow" not in c and "exact_at_chi" not in c + + # chi=64 fits a 6-qubit cut exactly (2**6) but not an 8-qubit cut (2**8): + # block_size=4 (cut 8) overflows, block_size=3 (cut 6) is exact. + aware = finder.recommend_layered(block_sizes=(3, 4), chi=64) + assert aware["chi"] == 64 + assert aware["recommended_block_size"] == 3 + assert aware["plan"].max_bond_cut() <= 6 + by_bs = {c["block_size"]: c for c in aware["candidates"]} + assert by_bs[3]["exact_at_chi"] and by_bs[3]["chi_overflow"] == 0.0 + assert not by_bs[4]["exact_at_chi"] and by_bs[4]["chi_overflow"] == 2.0 + # The recommended candidate always has the minimum chi_overflow. + overflows = [c["chi_overflow"] for c in aware["candidates"]] + assert by_bs[aware["recommended_block_size"]]["chi_overflow"] == min(overflows) + + +def test_recommend_layered_inherits_finder_chi_unless_overridden(): + """The direct layered recommendation follows the finder's chi policy.""" + rng = np.random.default_rng(216) + stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) + finder = TreeLayoutFinder( + stream, n=16, chi=64, weight_mode="operator_schmidt" + ) + + inherited = finder.recommend_layered(block_sizes=(3, 4)) + assert inherited["chi"] == 64 + assert inherited["recommended_block_size"] == 3 + + blind = finder.recommend_layered(block_sizes=(3, 4), chi=None) + assert blind["chi"] is None + + +def test_recommend_arities_chi_aware_minimizes_overflow(): + """``chi``-aware arity search prefers a structure exact at ``chi``.""" + rng = np.random.default_rng(212) + stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) + finder = TreeLayoutFinder(stream, n=16, objective="congestion") + + rec = finder.recommend_arities((2, 3, 4), chi=64) + assert rec["chi"] == 64 + for c in rec["candidates"]: + assert {"max_bond_cut", "chi_overflow", "exact_at_chi"} <= set(c) + recommended = next( + c for c in rec["candidates"] + if c["max_arity"] == rec["recommended_max_arity"] + ) + overflows = [c["chi_overflow"] for c in rec["candidates"]] + assert recommended["chi_overflow"] == min(overflows) + # recommend_layout forwards chi unchanged. + assert finder.recommend_layout((2, 3, 4), chi=64)["chi"] == 64 + + +def test_recommend_layered_rejects_bad_chi(): + """A non-positive ``chi`` budget is rejected.""" + finder = TreeLayoutFinder([], n=4, structure="balanced") + with pytest.raises(ValueError, match="chi must be a positive integer"): + 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.""" + 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.chi is None + # run() returns the objective-best arity from the chi-blind recommendation. + searched = finder.run() + 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, + weight_mode="operator_schmidt") + assert fixed.arity_candidates is None + assert fixed.run().is_binary() + + +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) + stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) + finder = TreeLayoutFinder(stream, n=16, chi=64, + weight_mode="operator_schmidt") + + assert finder.chi == 64 + searched = finder.run() + aware = finder.recommend_arities((2, 3, 4), chi=64) + assert searched.children == aware["plan"].children + # 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.""" + 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, + weight_mode="operator_schmidt") + assert opt.plan.children == finder.run().children + + # A scalar max_arity=2 forces a fixed binary tree through the optimizer. + fixed = TreeOptimizer(stream, n=16, chi=64, max_arity=2, + layout_weight_mode="operator_schmidt", run=False) + assert fixed.plan.is_binary() + + +def test_layout_and_entangled_state_handoff_is_explicit(): + """A layout finder and an entangled TTN can be handed off safely.""" + plan_finder = TreeLayoutFinder([], n=4, structure="balanced") + state_plan = plan_finder.run() + state = TreeTensorNetwork.rand(state_plan, D=2, seed=104) + before = state.to_statevector() + + opt = TreeOptimizer( + None, + layout=plan_finder, + state=state, + run=False, + ) + + assert _fidelity(before, opt.to_dense()) > 1 - 1e-10 + assert opt.layout_report()["is_binary"] + with pytest.raises(TypeError, match="TreePlan"): + TreeOptimizer(None, tree=state, n=4, run=False) + with pytest.raises(TypeError, match="not a TreeTensorNetwork"): + TreeLayoutFinder(state) + + +def test_product_ttn_is_remounted_exactly_on_a_requested_new_layout(): + """A product TTN can safely move to a different tree geometry.""" + source_plan = TreePlan.from_order(range(4), structure="balanced") + 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() + + 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 np.allclose( + np.asarray(opt.to_dense()).reshape(-1), + np.asarray(source.to_dense()).reshape(-1), + ) + + +def test_entangled_ttn_relayout_is_rejected_before_any_lossy_conversion(): + """Changing an entangled TTN's geometry requires an explicit conversion.""" + source_plan = TreePlan.from_order(range(4), structure="balanced") + 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) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + source = TreeOptimizer( + [(h, 0), (cnot, (0, 1))], tree=source_plan, chi=8 + ).tn.copy() + assert source.max_bond() > 1 + + with pytest.raises(ValueError, match="potentially lossy relayout"): + TreeOptimizer(None, state=source, tree=target_plan, run=False) + + +def test_product_mps_is_accepted_and_exactly_mounted_on_the_tree(): + """A bond-one MPS is a geometry-neutral product-state input.""" + mps = qtn.MPS_computational_state("010", dtype="complex128") + plan = TreePlan.from_order((2, 0, 1), structure="balanced") + + opt = TreeOptimizer(None, state=mps, tree=plan, run=False) + + assert opt.plan is plan + assert opt.max_bond() == 1 + assert np.allclose( + np.asarray(opt.to_dense()).reshape(-1), + np.asarray(mps.to_dense()).reshape(-1), + ) + + +def test_entangled_mps_initial_state_is_rejected(): + """An MPS with a nontrivial virtual bond cannot be silently tree-remapped.""" + mps = qtn.MPS_rand_state( + 3, bond_dim=2, phys_dim=2, dtype="complex128", seed=45 + ) + assert mps.max_bond() > 1 + + with pytest.raises(TypeError, match=r"max_bond\(\) == 1"): + TreeOptimizer(None, state=mps, run=False) + + def test_public_api_exports_tree_optimizer(): """TreeOptimizer is exposed through the public namespaces.""" import pepsy @@ -255,6 +893,138 @@ def test_bond_report_reflects_chi(): assert rep["n_tensors"] == len(opt.plan.nodes()) +def test_estimate_bonds_uses_crossing_operator_schmidt_ranks(): + """The paper dry-run multiplies ranks only on edges crossed by a gate.""" + plan = TreePlan.from_order(range(4), structure="balanced") + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + rng = np.random.default_rng(0) + generic = rng.standard_normal((4, 4)) + 1j * rng.standard_normal((4, 4)) + opt = TreeOptimizer(None, n=4, tree=plan, chi=4, run=False) + before = opt.to_dense() + + report = opt.estimate_bonds([(cnot, (0, 3)), (generic, (0, 3))]) + path_edges = { + tuple(sorted(edge)) + for edge in zip( + plan.node_path(plan.leaf_of_qubit[0], plan.leaf_of_qubit[3]), + plan.node_path(plan.leaf_of_qubit[0], plan.leaf_of_qubit[3])[1:], + ) + } + + assert report["max_bond"] == 8 # rank(CNOT)=2, rank(generic)=4 + assert report["requires_truncation"] + assert set(report["edge_bonds"]) == { + tuple(sorted((parent, child))) + for parent, children in plan.children.items() + for child in children + } + assert all(report["edge_bonds"][edge] == 8 for edge in path_edges) + assert any(report["edge_bonds"][edge] == 1 for edge in report["edge_bonds"] + if edge not in path_edges) + assert report["events"][0]["crossing_edges"] + assert set(report["events"][0]["crossing_edges"].values()) == {2} + assert set(report["events"][1]["crossing_edges"].values()) == {4} + assert np.allclose(before, opt.to_dense()) # diagnostic is non-mutating + + +def test_estimate_bonds_ignores_single_site_and_control_events(): + """One-site operations and measurements do not grow the dry-run bound.""" + z = np.diag([1.0, -1.0]).astype(complex) + opt = TreeOptimizer( + [(z, 0), ("measure", "Z", 1, +1)], n=3, chi=2, run=False + ) + report = opt.estimate_bonds() + assert report["max_bond"] == 1 + assert all(not event["crossing_edges"] for event in report["events"]) + + +def test_preflight_reports_and_rejects_resource_limits(): + """Preflight protects replay without changing the live state.""" + plan = TreePlan.from_order(range(4), structure="balanced") + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + opt = TreeOptimizer(None, n=4, tree=plan, chi=4, run=False) + report = opt.preflight( + [(cnot, (0, 3))], max_bond=1, raise_on_error=False + ) + assert report["ok"] is False + assert report["violations"] + with pytest.raises(MemoryError, match="max_bond"): + opt.preflight([(cnot, (0, 3))], max_bond=1) + with pytest.raises(MemoryError, match="estimated max bond"): + TreeOptimizer( + [(cnot, (0, 3))], n=4, tree=plan, + max_intermediate_bond=1, + ) + + with pytest.raises(MemoryError, match="max_operator_qubits"): + TreeOptimizer(None, n=3, max_operator_qubits=2).apply_gate( + _rand_unitary(3, np.random.default_rng(1)), (0, 1, 2) + ) + + +def test_truncation_report_tracks_per_edge_discarded_weight(): + """Tracked runs expose local spectra and discarded weights per edge.""" + 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, + ) + plan = TreePlan.from_order(range(4), structure="balanced") + opt = TreeOptimizer( + [(h, 0), (cnot, (0, 3))], + n=4, + tree=plan, + chi=1, + track_truncation=True, + ) + + report = opt.truncation_report() + assert report["track_truncation"] is True + assert report["n_events"] == len(opt.truncation_history) > 0 + assert report["n_tracked"] == report["n_events"] + assert report["max_discarded_fraction"] == pytest.approx(0.5) + assert any(event["kind"] == "compress" for event in report["events"]) + updates = report["updates"] + assert len(updates) == 2 + assert updates[0]["support"] == (0,) + assert updates[0]["edge_count"] == 0 + assert updates[0]["relative_discarded_weight"] == pytest.approx(0.0) + assert updates[1]["support"] == (0, 3) + assert updates[1]["edge_count"] == 4 + assert updates[1]["relative_discarded_weight"] == pytest.approx(0.5) + assert updates[1]["cumulative_relative_discarded_weight"] == pytest.approx(0.5) + for event in report["events"]: + assert event["after_bond"] <= event["before_bond"] + assert event["spectrum_rank"] is not None + assert event["discarded_weight"] >= 0.0 + assert event["discarded_fraction"] >= 0.0 + + +def test_truncation_history_keeps_fast_untracked_path_cheap(): + """The default path records dimensions without probing full spectra.""" + 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, 3))], n=4, chi=1, track_truncation=False + ) + + report = opt.truncation_report() + assert report["track_truncation"] is False + assert report["n_events"] > 0 + assert report["n_tracked"] == 0 + assert report["total_discarded_weight"] is None + assert all(event["discarded_weight"] is None for event in report["events"]) + + def test_convergence_sweep_reports_rising_fidelity(): """convergence_sweep reuses one tree and reports monotone fidelity.""" rng = np.random.default_rng(13) @@ -288,6 +1058,24 @@ def test_convergence_sweep_skips_fidelity_when_large(): assert all(r["fidelity"] is None for r in recs) +def test_convergence_sweep_reuses_generator_stream(): + """A one-shot gate iterator produces the same sweep as a list.""" + rng = np.random.default_rng(141) + stream = _random_stream(4, 8, rng, two_qubit_frac=0.6) + kwargs = dict(n=4, chi_values=(1, 2, 4), dense_cap=0) + from_list = TreeOptimizer.convergence_sweep(stream, **kwargs) + from_generator = TreeOptimizer.convergence_sweep( + (entry for entry in stream), **kwargs + ) + assert [ + (rec["chi"], rec["max_bond"], rec["norm"]) + for rec in from_generator + ] == pytest.approx([ + (rec["chi"], rec["max_bond"], rec["norm"]) + for rec in from_list + ]) + + # -- stability / speed hardening ---------------------------------------------- @@ -348,6 +1136,25 @@ def test_apply_gate_routes_three_qubit_gate_to_subtree(): assert _fidelity(psi, opt.to_dense()) > 1 - 1e-9 +def test_subtree_operator_uses_recursive_pairwise_messages(monkeypatch): + """The tree-MPO path never contracts the whole state subtree at once.""" + rng = np.random.default_rng(211) + opt = TreeOptimizer(None, n=7, chi=64) + calls = [] + tensor_contract = qtn.tensor_contract + + def traced_contract(*tensors, **kwargs): + calls.append(len(tensors)) + return tensor_contract(*tensors, **kwargs) + + monkeypatch.setattr(qtn, "tensor_contract", traced_contract) + opt.apply_subtree_operator(_rand_unitary(3, rng), (0, 3, 5)) + + assert calls + assert max(calls) <= 2 + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_apply_gate_four_qubit_trotter_block_matches_dense(): """A four-qubit block applied in one shot matches the dense reference.""" rng = np.random.default_rng(22) @@ -458,6 +1265,99 @@ def test_copy_is_independent(): assert not np.allclose(base.to_dense(), clone.to_dense()) +def test_copy_preserves_layout_and_history_configuration(): + """copy() keeps the parameters that determine future layout/replay.""" + base = TreeOptimizer( + None, n=5, max_arity=3, community_frac=0.11, star_frac=0.22, + record_history=False, run=False, + ) + clone = base.copy() + assert clone.max_arity == 3 + assert clone.community_frac == pytest.approx(0.11) + assert clone.star_frac == pytest.approx(0.22) + assert clone.record_history is False + + +def test_record_history_can_be_disabled(): + """Large replays can omit retained per-edge/update history.""" + opt = TreeOptimizer( + [(_rand_unitary(2, np.random.default_rng(142)), (0, 3))], + n=4, chi=1, record_history=False, + ) + report = opt.truncation_report() + assert report["n_events"] == 0 + assert report["updates"] == [] + + +def test_copy_rng_is_deterministic_but_independent(): + """Copies derive reproducible but distinct random streams.""" + first = TreeOptimizer(None, n=2, seed=91, run=False) + second = TreeOptimizer(None, n=2, seed=91, run=False) + first_draws = [first.copy().rng.random() for _ in range(3)] + second_draws = [second.copy().rng.random() for _ in range(3)] + + assert np.allclose(first_draws, second_draws) + assert not np.isclose(first_draws[0], first_draws[1]) + + +def test_thread_index_clears_after_failed_two_qubit_update(monkeypatch): + """A failed threaded gate cannot leave its temporary bond index live.""" + plan = TreePlan.from_order(range(4), structure="balanced") + opt = TreeOptimizer(None, n=4, tree=plan, run=False) + + def fail_thread(*_args): + raise RuntimeError("synthetic thread failure") + + monkeypatch.setattr(opt, "_thread_hop", fail_thread) + with pytest.raises(RuntimeError, match="synthetic thread failure"): + opt.apply_2q(np.eye(4, dtype=complex), 0, 3) + assert opt._thread_ind is None + + +def test_tree_run_progbar_reports_norm_infidelity(monkeypatch): + """Tree replay exposes MPS-style progress fields without SVD probes.""" + progress_instances = [] + + class _FakeTqdm: + def __init__(self, **kwargs): + self.total = kwargs["total"] + self.n = 0 + self.postfix_calls = [] + progress_instances.append(self) + + def set_postfix(self, postfix): + self.postfix_calls.append(dict(postfix)) + + def update(self, amount): + self.n += amount + + def close(self): + pass + + monkeypatch.setitem( + sys.modules, "tqdm", types.SimpleNamespace(tqdm=_FakeTqdm) + ) + + 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, + ) + plan = TreePlan.from_order(range(4), structure="balanced") + opt = TreeOptimizer(None, n=4, tree=plan, chi=1, run=False) + opt.set_gates([(h, 0), (cnot, (0, 3))]) + opt.run(progbar=True) + + progress = progress_instances[-1] + assert progress.total == 2 + assert progress.n == 2 + last = progress.postfix_calls[-1] + assert {"2q", "infidelity"} <= set(last) + # kq is only shown when multi-qubit (>2) gates are present. + assert "kq" not in last + assert last["2q"] == 1 + + def test_threads_setting_preserves_result(): """The thread cap is a performance knob only; results are identical.""" rng = np.random.default_rng(18) @@ -500,30 +1400,6 @@ def test_mixed_paths_match_statevector(): assert _fidelity(psi, opt.to_dense()) > 1 - 1e-8 -@pytest.mark.parametrize("seed", [22, 23]) -def test_multisite_local_expectation_matches_exact(seed): - """Canonical multi-site expectations match the dense state exactly.""" - rng = np.random.default_rng(seed) - n = 7 - stream = _random_stream(n, 50, rng) - opt = TreeOptimizer(stream, n=n, chi=128) - psi = _exact_state(stream, n) - psi /= np.linalg.norm(psi) - - z = np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex) - x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) - zz = np.kron(z, z) - zxz = np.kron(np.kron(z, x), z) - for where in [(0, 3), (1, 5), (2, 6), (0, 6)]: - got = opt.local_expectation(zz, where) - exact = _sv_expect(psi, zz, where, n) - assert abs(got - exact) < 1e-9 - for where in [(0, 3, 6), (1, 2, 5)]: - got = opt.local_expectation(zxz, where) - exact = _sv_expect(psi, zxz, where, n) - assert abs(got - exact) < 1e-9 - - def test_measure_born_statistics_and_collapse(): """Measurement samples the Born rule and collapses to a unit-norm state.""" theta = 0.7 @@ -543,10 +1419,9 @@ def test_measure_born_statistics_and_collapse(): def test_reset_forces_ground_state(): """reset() returns a qubit to |0> regardless of its prior value.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) - z = np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex) opt = TreeOptimizer([(x, 1)], n=3, chi=4, seed=1) # qubit 1 in |1> opt.reset(1) - assert opt.local_expectation(z, 1).real > 1 - 1e-9 # = +1 -> |0> + assert _fidelity(opt.to_dense(), np.array([1.0] + [0.0] * 7)) > 1 - 1e-9 assert abs(opt.norm() - 1.0) < 1e-9 @@ -558,6 +1433,405 @@ def test_measure_is_seed_reproducible(): assert a == b +def test_tree_stream_measure_and_reset_match_mps_event_contract(): + """TTN streams accept Pauli measurement/reset events and record outcomes.""" + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt = TreeOptimizer( + [(x, 0), ("measure", "Z", 0, -1), ("reset", 0)], + n=2, + chi=4, + seed=12, + ) + + assert len(opt.measurements) == 1 + pauli, where, outcome, probability = opt.measurements[0] + assert (pauli, where, outcome) == ("Z", (0,), -1) + assert probability == pytest.approx(1.0) + assert _fidelity(opt.to_dense(), np.array([1.0, 0.0, 0.0, 0.0])) > 1 - 1e-12 + assert opt.event_types == ["gate", "measure", "reset"] + + +def test_tree_stream_measure_reset_records_then_prepares_pauli_state(): + """measure_reset records the result and leaves the + Pauli eigenstate.""" + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) + opt = TreeOptimizer( + [(h, 0), TreeOptimizer.measure_reset_event("X", 0, +1)], + n=1, + chi=4, + ) + + assert opt.measurements == [("X", (0,), +1, pytest.approx(1.0))] + assert _fidelity( + opt.to_dense(), np.array([1.0, 1.0], dtype=complex) / np.sqrt(2.0) + ) > 1 - 1e-12 + assert opt.norm() == pytest.approx(1.0) + + +def test_tree_stream_multisite_pauli_measurement(): + """A product-Pauli event can collapse a multi-qubit tree subtree.""" + opt = TreeOptimizer([("measure", "ZZ", (0, 1), +1)], n=2, chi=4) + + assert opt.measurements[0][:3] == ("ZZ", (0, 1), +1) + assert opt.measurements[0][3] == pytest.approx(1.0) + assert opt.norm() == pytest.approx(1.0) + + +def test_multisite_pauli_measurement_preserves_parity_sector_coherence(): + """Parity projection must not collapse the individual Pauli outcomes.""" + 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=2, chi=8) + + opt._measure_pauli("ZZ", (0, 1), +1) + assert _fidelity( + opt.to_dense(), np.array([1.0, 0.0, 0.0, 1.0]) / np.sqrt(2.0) + ) > 1 - 1e-12 + + +def test_wide_pauli_measurement_avoids_dense_projector(): + """A wide product-Pauli event remains factorized past the dense limit.""" + n = 9 + opt = TreeOptimizer( + [("measure", "Z" * n, tuple(range(n)), +1)], n=n, chi=4 + ) + assert opt.measurements[0][2] == +1 + assert opt.norm() == pytest.approx(1.0) + + +def test_default_dense_operator_guard(): + """General dense operators have a finite default support limit.""" + opt = TreeOptimizer(None, n=9, run=False) + assert opt.max_operator_qubits == 8 + with pytest.raises(MemoryError, match="max_operator_qubits"): + opt.apply_gate(np.eye(2**9, dtype=complex), tuple(range(9))) + + +def test_tree_stream_control_mapping_and_cap_event(): + """Mapping controls and MPS-compatible cap events work on a tree.""" + opt = TreeOptimizer( + [{"kind": "measure", "pauli": "Z", "where": 0, "outcome": +1}], + n=1, + ) + assert opt.measurements[0][:3] == ("Z", (0,), +1) + + capped = TreeOptimizer( + [TreeOptimizer.cap_event(0, [1.0, 0.0])], n=2, chi=4 + ) + assert capped.n == capped.plan.n == capped.tn.nqubits == 1 + assert np.allclose(capped.to_dense(), [1.0, 0.0]) + + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + shifted = TreeOptimizer( + [TreeOptimizer.cap_event(1, [1.0, 0.0]), (x, 1)], + n=3, + chi=4, + ) + assert shifted.n == 2 + assert np.argmax(np.abs(shifted.to_dense())) == 1 + + +def test_tree_cap_matches_dense_contraction_and_compacts_plan(): + """Capping an entangled leaf matches dense contraction and keeps a valid TTN.""" + 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, + tree=TreePlan.from_order(range(4), structure="balanced"), + chi=8, + ) + before = opt.to_dense().reshape((2,) * 4) + vec = np.array([1.0, 2.0], dtype=complex) + expected = np.tensordot(vec, before, axes=(0, 1)).reshape(-1) + + opt.cap(1, vec) + + assert opt.n == opt.plan.n == opt.tn.nqubits == 3 + assert np.allclose(opt.to_dense(), expected) + assert opt.tn.validate(check_canonical=True) is opt.tn + assert set(opt.plan.leaf_of_qubit) == {0, 1, 2} + + +def test_tree_public_submpo_and_pauli_backend_operations(): + """Public tree operator primitives cover native MPO and Pauli paths.""" + n = 4 + plan = TreePlan.from_order(range(n), structure="balanced") + mpo = _two_branch_flip_submpo(L=n, sites=(0, 3), targets=(0, 3)) + opt = TreeOptimizer(None, n=n, tree=plan, chi=16, run=False) + + opt.apply_submpo(mpo, (0, 3)) + assert np.allclose( + opt.to_dense(), + 0.7 * np.eye(16, dtype=complex)[:, 0] + + 0.3 * _sv_apply_kq( + np.eye(16, dtype=complex)[:, 0], + np.kron(np.array([[0, 1], [1, 0]], dtype=complex), + np.array([[0, 1], [1, 0]], dtype=complex)), + (0, 3), + n, + ), + ) + + rotated = TreeOptimizer(None, n=n, tree=plan, chi=16, run=False) + theta = 0.37 + rotated.apply_pauli_rotation(theta, "XZ", (0, 3)) + pauli = np.kron( + np.array([[0, 1], [1, 0]], dtype=complex), + np.array([[1, 0], [0, -1]], dtype=complex), + ) + expected = ( + np.cos(theta / 2.0) * np.eye(4, dtype=complex) + - 1j * np.sin(theta / 2.0) * pauli + ) + assert np.allclose( + rotated.to_dense(), + _sv_apply_kq(np.eye(16, dtype=complex)[:, 0], expected, (0, 3), n), + ) + + summed = TreeOptimizer(None, n=n, tree=plan, chi=16, run=False) + summed.apply_pauli_sum([(1.0, {0: "X", 3: "X"})]) + assert np.allclose( + summed.to_dense(), + _sv_apply_kq( + np.eye(16, dtype=complex)[:, 0], + np.kron( + np.array([[0, 1], [1, 0]], dtype=complex), + np.array([[0, 1], [1, 0]], dtype=complex), + ), + (0, 3), + n, + ), + ) + + +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) + 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=2, chi=8) + assert opt.expectation_pauli("ZZ", (0, 1)) == pytest.approx(1.0) + assert opt.expectation_pauli("XX", (0, 1)) == pytest.approx(1.0) + opt.project_pauli("ZZ", (0, 1), +1) + assert opt.expectation_pauli("ZZ", (0, 1)) == pytest.approx(1.0) + + +def test_tree_public_pauli_measurement_returns_probability_and_diagnostics(): + """The public Pauli measurement API exposes Born probability diagnostics.""" + theta = 0.8 + ry = np.array([ + [np.cos(theta / 2.0), -np.sin(theta / 2.0)], + [np.sin(theta / 2.0), np.cos(theta / 2.0)], + ], dtype=complex) + opt = TreeOptimizer([(ry, 0)], n=5, chi=8) + + outcome, probability, diagnostics = opt.measure_pauli( + "Z", 0, outcome=+1, return_diagnostics=True + ) + + assert outcome == +1 + assert probability == pytest.approx(np.cos(theta / 2.0) ** 2) + assert diagnostics["probability"] == pytest.approx(probability) + assert diagnostics["norm_before"] == pytest.approx(1.0) + assert diagnostics["norm_after"] == pytest.approx(1.0) + assert diagnostics["support"] == (0,) + assert diagnostics["span_before"] == diagnostics["span_after"] + assert diagnostics["bonds_before"] == diagnostics["bonds_after"] + assert opt.expectation_pauli("Z", 0) == pytest.approx(1.0) + + +def test_tree_pauli_projection_can_preserve_branch_norm(): + """A non-normalizing Pauli projection retains its physical survival norm.""" + theta = 0.9 + ry = np.array([ + [np.cos(theta / 2.0), -np.sin(theta / 2.0)], + [np.sin(theta / 2.0), np.cos(theta / 2.0)], + ], dtype=complex) + opt = TreeOptimizer([(ry, 0)], n=4, chi=8) + + diagnostics = opt.project_pauli( + "Z", 0, +1, renormalize=False, return_diagnostics=True + ) + expected_probability = np.cos(theta / 2.0) ** 2 + + assert diagnostics["renormalized"] is False + assert diagnostics["norm_after"] == pytest.approx( + np.sqrt(expected_probability) + ) + assert diagnostics["norm_ratio"] == pytest.approx( + np.sqrt(expected_probability) + ) + assert opt.expectation_pauli("Z", 0) == pytest.approx(1.0) + assert opt.get_projection_diagnostics()[-1] is diagnostics + + +def test_tree_sparse_long_pauli_avoids_dense_operator_limit(): + """A long sparse Pauli measurement uses the factorized tree path.""" + n = 17 + where = (0, 3, 7, 11, 16) + opt = TreeOptimizer(None, n=n, chi=8, max_operator_qubits=2, run=False) + + outcome, probability, diagnostics = opt.measure_pauli( + "ZZZZZ", where, outcome=+1, return_diagnostics=True + ) + + assert outcome == +1 + assert probability == pytest.approx(1.0) + assert diagnostics["support"] == where + assert diagnostics["max_bond_after"] <= 8 + assert opt.norm() == pytest.approx(1.0) + + +def test_tree_cap_can_preserve_stable_logical_labels(): + """Stable-label caps compact storage but preserve caller-facing IDs.""" + opt = TreeOptimizer(None, n=4, chi=8, run=False) + + opt.cap(1, [1.0, 0.0], stable_labels=True) + + assert opt.n == 3 + assert opt.qubits == [0, 2, 3] + assert opt.logical_order == [0, 2, 3] + assert opt.position(2) == 1 + assert opt.logical_site(1) == 2 + + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt.apply_1q(x, 3) + assert np.argmax(np.abs(opt.to_dense())) == 1 + + +def test_tree_stable_labels_work_for_pauli_projection_paths(): + """Pauli projection resolves stable labels once, then uses compact sites.""" + opt = TreeOptimizer(None, n=4, chi=8, run=False) + opt.cap(1, [1.0, 0.0], stable_labels=True) + + diagnostics = opt.project_pauli( + "Z", 2, +1, renormalize=False, return_diagnostics=True + ) + + assert diagnostics["support"] == (2,) + assert diagnostics["norm_after"] == pytest.approx(1.0) + assert opt.expectation_pauli("Z", 2) == pytest.approx(1.0) + + +def test_tree_stable_cap_event_supports_later_logical_events(): + """A stable-label cap event keeps later stream labels addressable.""" + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt = TreeOptimizer( + [ + TreeOptimizer.cap_event( + 1, [1.0, 0.0], compact_labels=False + ), + (x, 3), + ], + n=4, + chi=8, + ) + + assert opt.qubits == [0, 2, 3] + assert np.argmax(np.abs(opt.to_dense())) == 1 + + +def test_tree_reset_reuses_an_ancilla_without_disturbing_data(): + """Reset and repeated use of an ancilla leave the data Bell pair intact.""" + 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, + ) + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt = TreeOptimizer([(h, 0), (cnot, (0, 1)), (x, 2)], n=3, chi=8) + + opt.reset(2) + assert opt.expectation_pauli("Z", 2) == pytest.approx(1.0) + opt.apply_1q(x, 2) + opt.reset(2) + + assert opt.expectation_pauli("ZZ", (0, 1)) == pytest.approx(1.0) + assert opt.expectation_pauli("XX", (0, 1)) == pytest.approx(1.0) + assert opt.expectation_pauli("Z", 2) == pytest.approx(1.0) + + +def test_tree_stream_submpo_markers_use_recursive_operator_path(): + """Tuple and mapping sub-MPO markers match their dense support operator.""" + mpo = _two_branch_flip_submpo(L=4, sites=(0, 3), targets=(0, 3)) + dense = np.asarray(mpo.to_dense()) + + event = TreeOptimizer.submpo_event(mpo, (0, 3)) + tuple_opt = TreeOptimizer([event], n=4, chi=8) + mapping_opt = TreeOptimizer( + [{"kind": "submpo", "mpo": mpo, "where": [0, 3]}], + n=4, + chi=8, + ) + expected = _sv_apply_kq(np.eye(16, dtype=complex)[:, 0], dense, (0, 3), 4) + + assert tuple_opt.event_types == ["submpo"] + assert TreeOptimizer.is_submpo_event(event) + assert TreeOptimizer.submpo_event_parts(event)[1] == (0, 3) + assert tuple_opt.update_history[0]["kind"] == "submpo" + assert _fidelity(expected, tuple_opt.to_dense()) > 1 - 1e-10 + assert _fidelity(tuple_opt.to_dense(), mapping_opt.to_dense()) > 1 - 1e-10 + + +def test_tree_stream_submpo_does_not_require_dense_materialization(monkeypatch): + """Native MPO replay and bond estimation avoid ``to_dense`` allocation.""" + mpo = _two_branch_flip_submpo(L=4, sites=(0, 3), targets=(0, 3)) + expected = np.asarray(mpo.to_dense()) + + def fail_to_dense(): + raise AssertionError("sub-MPO was unexpectedly materialized") + + monkeypatch.setattr(mpo, "to_dense", fail_to_dense) + event = TreeOptimizer.submpo_event(mpo, (0, 3)) + opt = TreeOptimizer([event], n=4, chi=8) + report = opt.estimate_bonds([event]) + + reference = _sv_apply_kq( + np.eye(16, dtype=complex)[:, 0], expected, (0, 3), 4 + ) + assert _fidelity(reference, opt.to_dense()) > 1 - 1e-10 + assert report["events"][0]["crossing_edges"] + + +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)) + plan = TreePlan.from_order(range(4), structure="balanced") + opt = TreeOptimizer(None, n=4, tree=plan, chi=1, run=False) + + report = opt.estimate_bonds([("submpo", mpo, (0, 3))]) + + assert report["events"][0]["kind"] == "submpo" + assert set(report["events"][0]["crossing_edges"].values()) == {2} + assert report["max_bond"] == 2 + assert report["requires_truncation"] + with pytest.raises(MemoryError, match="max_operator_qubits"): + opt.preflight( + [("submpo", mpo, (0, 3))], + max_operator_qubits=1, + ) + + +def test_measurement_enforces_max_subtree_nodes_on_streamed_product_pauli(): + """Control-event execution applies the same subtree guard as dense gates.""" + opt = TreeOptimizer(None, n=8, max_subtree_nodes=1, run=False) + with pytest.raises(MemoryError, match="max_subtree_nodes"): + opt.run([{ + "kind": "measure", + "pauli": "ZZ", + "where": [0, 7], + "outcome": +1, + }]) + + # -- TreeTensorNetwork class -------------------------------------------------- @@ -579,6 +1853,29 @@ def test_ttn_from_plan_is_product_state(): assert np.linalg.norm(sv[1:]) < 1e-12 +def test_ps_to_ttn_matches_product_state_constructor_api(): + """The high-level TTN constructor mirrors ``ps_to_mps`` amplitudes.""" + theta = 0.31 + expected = np.array([1.0], dtype="complex128") + local = np.array([np.cos(theta), np.sin(theta)], dtype="complex128") + for _ in range(4): + expected = np.kron(expected, local) + + state = pepsy.ps_to_ttn(4, theta=theta) + assert isinstance(state, TreeTensorNetwork) + assert state.max_bond() == 1 + assert np.allclose(state.to_statevector(), expected) + assert state.is_canonical_form(state.root) + + expanded = pepsy.ps_to_ttn(4, chi=2, rand_strength=0.0) + assert expanded.max_bond() == 2 + assert np.allclose(expanded.to_statevector(), [1.0] + [0.0] * 15) + + plan = TreePlan.from_order(range(4), structure="balanced") + explicit = pepsy.ps_to_ttn(4, tree=plan) + assert explicit.plan is plan + + 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") @@ -592,6 +1889,94 @@ def test_ttn_copy_preserves_geometry_and_type(): assert other.node_tid(2) in other.tensor_map +def test_plain_tensor_network_cast_requires_explicit_plan(): + """A generic Quimb network cannot silently become geometry-owning.""" + plain = qtn.TensorNetwork([ + qtn.Tensor(np.ones(2), inds=("k0",)), + ]) + with pytest.raises(TypeError, match="explicit TreePlan"): + TreeTensorNetwork(plain) + + +def test_layout_rejects_out_of_range_supports_early(): + """Bad interaction supports fail during layout construction, not replay.""" + with pytest.raises(ValueError, match="outside"): + TreeLayoutFinder(supports=[(0, 3)], n=2) + + +def test_three_qubit_torch_operator_is_backend_coerced(): + """Torch operators work with the default NumPy-backed TTN state.""" + torch = pytest.importorskip("torch") + opt = TreeOptimizer(None, n=3, run=False) + with pytest.warns(UserWarning, match="backend-compatible gate"): + opt.apply_subtree_operator( + torch.eye(8, dtype=torch.complex128), (0, 1, 2) + ) + assert np.allclose(opt.to_dense(), [1.0] + [0.0] * 7) + + +def test_tree_torch_state_stays_native_across_public_operations(): + """Tree controls, Pauli helpers, and readout preserve a Torch TTN.""" + torch = pytest.importorskip("torch") + 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)) + h = to_backend( + np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) + ) + cnot = to_backend(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))], state=state, chi=8) + assert opt.backend_info() == { + "backend": "torch", "dtype": "complex128", "device": "cpu", + } + 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)) + opt.project_pauli("ZZ", (0, 1), +1) + assert opt.measure(2, outcome=0) == 0 + assert opt.reset(2) == 0 + opt.cap(2, to_backend(np.array([1.0, 0.0], dtype=complex))) + opt.apply_pauli_rotation(0.2, "XZ", (0, 1)) + opt.apply_pauli_sum([(1.0, {0: "X", 1: "Z"})]) + + assert all(torch.is_tensor(tensor.data) for tensor in opt.tn.tensor_map.values()) + + +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") + 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)) + opt = TreeOptimizer(None, state=state, run=False) + + with pytest.warns(UserWarning, match="backend-compatible gate"): + opt.apply_1q(np.eye(2, dtype=complex), 0) + opt.apply_1q(np.eye(2, dtype=complex), 1) + assert opt.backend_info()["backend"] == "torch" + + +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") + plan = TreePlan.from_order(range(2), structure="balanced") + state = TreeTensorNetwork.from_plan(plan) + leaf = state.leaf_of_qubit(0) + state.node_tensor(leaf).modify( + data=torch.as_tensor(state.node_tensor(leaf).data, dtype=torch.complex128) + ) + + with pytest.raises(TypeError, match="one compatible backend"): + TreeOptimizer(None, state=state, run=False) + + def test_ttn_geometry_helpers_match_plan(): """Geometry delegators agree with the underlying TreePlan.""" plan = TreePlan.from_order(range(6), structure="balanced") @@ -612,6 +1997,29 @@ def test_ttn_geometry_helpers_match_plan(): ttn.bond(la, lb) # non-adjacent +def test_ttn_validate_checks_structure_and_canonicality(): + """TreeTensorNetwork.validate catches malformed physical legs.""" + plan = TreePlan.from_order(range(4), structure="balanced") + ttn = TreeTensorNetwork.from_plan(plan) + assert ttn.validate(check_canonical=True) is ttn + + broken = ttn.copy() + broken.reindex_({broken.site_ind(0): "broken-physical"}) + with pytest.raises(ValueError, match="missing physical index"): + broken.validate() + + +def test_tree_canonize_mps_compatibility_entry_point(): + """Shared coefficient frontends can use the MPS canonicalization name.""" + opt = TreeOptimizer(None, n=4, run=False) + info = {} + assert opt.canonize_mps(opt.p, (0, 3), info=info) == (0, 3) + assert info["cur_orthog"] == (0, 3) + assert opt.is_subtree_canonical_form() + assert opt.canonize_mps(opt.p, 2, info=info) == (2, 2) + assert opt.is_canonical_form(opt.plan.leaf_of_qubit[2]) + + def test_ttn_rand_is_canonical_around_root(): """rand(canonicalize=True) leaves the root tensor as the orthogonality centre.""" import quimb.tensor as qtn @@ -641,6 +2049,28 @@ def test_ttn_gate_and_local_expectation(): assert abs(val + 1.0) < 1e-9 # = -1 after X +def test_ttn_multisite_local_expectation_contracts_only_canonical_subtree(): + """The custom Steiner-subtree readout matches a full dense-TTN overlap.""" + rng = np.random.default_rng(41) + ttn = TreeTensorNetwork.rand(TreePlan.from_order(range(5)), D=3, seed=7) + matrix = rng.standard_normal((4, 4)) + 1.0j * rng.standard_normal((4, 4)) + operator = matrix + matrix.conj().T + where = (1, 4) + operated = qtn.tensor_network_gate_inds( + ttn, + operator, + [ttn.site_ind(site) for site in where], + contract=False, + inplace=False, + tags=[], + ) + expected = (ttn.H | operated).contract(all, optimize="auto") + expected /= (ttn.H | ttn).contract(all, optimize="auto") + + actual = ttn.local_expectation(operator, where, max_bond=None, optimize="auto") + assert actual == pytest.approx(expected) + + def test_optimizer_state_is_a_tree_tensor_network(): """TreeOptimizer builds its state on the TreeTensorNetwork class.""" opt = TreeOptimizer(None, n=4) @@ -673,11 +2103,19 @@ def dim_rows(drawing): assert rows and all(set(ln.split()) <= {"1"} for ln in rows) # dropping bond dims removes the annotation rows but keeps the structure assert not dim_rows(ttn.ascii_tree(bond_dims=False)) - # show() prints the same drawing (+ trailing newline) + # the coloured drawing embeds ANSI escapes but strips back to the plain one + colored = ttn.ascii_tree(color=True) + assert "\x1b[" in colored + import re as _re + assert _re.sub(r"\x1b\[[0-9;]*m", "", colored) == text + # show() prints the coloured drawing by default (+ trailing newline) ttn.show() + assert capsys.readouterr().out.rstrip("\n") == colored + # ...and the plain drawing when colour is disabled + ttn.show(color=False) assert capsys.readouterr().out.rstrip("\n") == text # optimizer delegates to the state's drawing - TreeOptimizer(None, n=4).show() + TreeOptimizer(None, n=4).show(color=False) assert capsys.readouterr().out.rstrip("\n") == text @@ -756,10 +2194,10 @@ def test_binary_defaults_unchanged(): for structure in ("quality", "balanced"): plan = TreePlan.from_order(range(9), structure=structure) assert plan.is_binary() - # the layout finder default is still a binary tree + # a scalar max_arity=2 opts back into a binary layout-finder tree rng = np.random.default_rng(2) stream = _random_stream(8, 60, rng) - assert TreeLayoutFinder(stream, n=8).run().is_binary() + assert TreeLayoutFinder(stream, n=8, max_arity=2).run().is_binary() def test_adaptive_layout_emits_star_for_cliques(): @@ -871,6 +2309,69 @@ def test_center_move_only_touches_geodesic(): assert np.array_equal(ttn.node_tensor(nid).data, snap[nid]) +def test_center_moves_use_lossless_qr_in_quimb(monkeypatch): + """Known-centre moves explicitly select Quimb's non-truncating QR path.""" + ttn = _entangled_ttn(seed=41) + calls = [] + canonize_between = ttn.canonize_between + + def traced_canonize_between(*args, **kwargs): + calls.append(dict(kwargs)) + return canonize_between(*args, **kwargs) + + monkeypatch.setattr(ttn, "canonize_between", traced_canonize_between) + ttn.shift_orthogonality_center(ttn.leaf_of_qubit(5)) + + assert calls + assert all(call["method"] == "qr" for call in calls) + assert all(call["cutoff"] == 0.0 for call in calls) + assert ttn.is_canonical_form() + + +def test_shift_center_recovers_from_multinode_region_locally(): + """A tracked canonical region is reduced by QR without touching its exterior.""" + ttn = _entangled_ttn(seed=42) + region = {ttn.root, *ttn.children(ttn.root)} + ttn.canonize_subtree_(region) + assert ttn.orthogonality_center is None + + outside = [nid for nid in ttn.plan.nodes() if nid not in region] + snapshot = { + nid: np.array(ttn.node_tensor(nid).data) + for nid in outside + } + target = sorted(region - {ttn.root})[0] + + ttn.shift_orthogonality_center(target) + + assert ttn.orthogonality_center == target + assert ttn.is_canonical_form(target) + for nid in outside: + assert np.array_equal(ttn.node_tensor(nid).data, snapshot[nid]) + + +def test_two_qubit_anchor_uses_nearest_endpoint(monkeypatch): + """A non-sibling gate starts from the endpoint nearest the centre.""" + plan = TreePlan.from_order(range(8), structure="balanced") + opt = TreeOptimizer(None, n=8, tree=plan, run=False) + near = plan.leaf_of_qubit[7] + opt.shift_orthogonality_center(near) + + moves = [] + move_center = opt._move_center + + def traced_move_center(target): + moves.append(target) + return move_center(target) + + monkeypatch.setattr(opt, "_move_center", traced_move_center) + x = np.array([[0, 1], [1, 0]], dtype=complex) + opt.apply_2q(np.kron(x, x), 0, 7) + + assert moves == [near] + assert opt.center == near + + def test_shift_center_validates_node(): """Shifting to a non-node raises loudly.""" ttn = _entangled_ttn(seed=5) @@ -1053,8 +2554,262 @@ def test_optimizer_subtree_canonicalisation_api(): assert opt.orthogonality_center == opt.plan.root +def test_live_bond_dimensions_survive_general_threading(): + """Tree edge diagnostics use live bonds after a non-sibling gate.""" + plan = TreePlan.from_order(range(8), structure="balanced") + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + opt = TreeOptimizer(None, n=8, tree=plan, chi=8) + opt.apply_gate(h, 0) + opt.apply_gate(cnot, (0, 6)) + for node in plan.nodes(): + for child in plan.children[node]: + ix = opt.tn.bond(node, child) + assert ix in opt.tn.ind_map + assert opt.tn._bond_dim(node, child) == opt.tn.ind_size(ix) +def test_nonunitary_one_qubit_gate_recenters_state(): + """A non-unitary one-qubit gate cannot leave a stale canonical centre.""" + plan = TreePlan.from_order(range(8), structure="balanced") + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + projector = np.diag([1.0, 0.0]).astype(complex) + opt = TreeOptimizer([(h, 0), (cnot, (0, 7))], n=8, tree=plan) + opt.apply_gate(projector, 7) + assert opt.is_canonical_form() + assert np.isclose(opt.norm(), np.linalg.norm(opt.to_dense())) + + +def test_forced_measurement_validates_outcome_probability(): + """Invalid or impossible forced outcomes raise before collapsing the state.""" + opt = TreeOptimizer(None, n=2) + with pytest.raises(ValueError, match="outcome must be 0 or 1"): + opt.measure(0, outcome=2) + with pytest.raises(ValueError, match="~0 probability"): + opt.measure(0, outcome=1) + assert np.isclose(opt.norm(), 1.0) + + +def test_multinode_region_normalization_preserves_canonicality(): + """Normalization scales inside a multi-node canonical region.""" + plan = TreePlan.from_order(range(8), structure="balanced") + opt = TreeOptimizer(None, n=8, tree=plan) + opt.tn = TreeTensorNetwork.rand(plan, D=2, seed=31) + leaf = plan.leaf_of_qubit[0] + region = {leaf, plan.parent[leaf]} + opt.tn.canonize_subtree_(region) + opt.tn.node_tensor(leaf).modify(data=3.0 * opt.tn.node_tensor(leaf).data) + + opt.normalize() + + assert np.isclose(opt.norm(), 1.0) + assert opt.is_subtree_canonical_form() + + +def test_shift_center_supports_left_absorption_orientation(): + """The optional left-absorption orientation still centres the target.""" + plan = TreePlan.from_order(range(8), structure="balanced") + ttn = TreeTensorNetwork.rand(plan, D=3, seed=32) + leaf = plan.leaf_of_qubit[0] + ttn.shift_orthogonality_center(leaf, absorb="left") + assert ttn.orthogonality_center == leaf + assert ttn.is_canonical_form() + + +def test_tree_plan_rejects_malformed_orders(): + """TreePlan.from_order enforces the same qubit-label contract as from_children.""" + with pytest.raises(ValueError, match="at least one"): + TreePlan.from_order([]) + with pytest.raises(ValueError, match="permutation"): + TreePlan.from_order([0, 0]) + with pytest.raises(ValueError, match="structure"): + TreePlan.from_order(range(2), structure="unknown") + + +def test_tree_gate_queue_set_and_add(): + """TreeOptimizer exposes queue replacement and extension like MpsOptimizer.""" + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt = TreeOptimizer(None, n=2) + assert opt.set_gates([(x, 0)]) is opt + assert opt.add_gates([(x, 1)]) is opt + assert len(opt.G) == 2 + opt.run() + assert _fidelity(opt.to_dense(), np.array([0.0, 0.0, 0.0, 1.0])) > 1 - 1e-12 + + +def test_optimizer_accepts_and_copies_initial_ttn(): + """TreeOptimizer can evolve an arbitrary supplied tree state independently.""" + plan = TreePlan.from_order(range(6), structure="balanced") + state = TreeTensorNetwork.rand(plan, D=2, seed=33) + before = state.to_statevector() + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + + opt = TreeOptimizer([(x, 0)], n=6, tn=state, chi=8) + + expected = _sv_apply_1q(before, x, 0, 6) + assert _fidelity(expected, opt.to_dense()) > 1 - 1e-10 + assert _fidelity(before, state.to_statevector()) > 1 - 1e-12 + assert opt.set_tn(state) is opt + + +def test_tree_native_fermionic_gate_stream_matches_mps(): + """Native (dim-4 Symmray) Fermi-Hubbard gates evolve correctly on a tree. + + The tree gate engine must apply block-sparse fermionic gates without + reshaping them into base-2 sub-legs. At a bond dimension large enough to be + exact, the tree real-time evolution must reproduce the MPS reference + (identical seed) to numerical precision. + """ + pytest.importorskip("symmray") + tensors = pepsy.tensors + + Lx, Ly, L = 2, 2, 4 + 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 + ) + setup = fermion.lattice_half_filling(Lx, Ly, pattern="checkerboard", cyclic=True) + mapper = tensors.OneDMap(Lx, Ly, mode="snake") + _, coo2idx = mapper.build() + edges_1d = [ + tuple(sorted((coo2idx[a], coo2idx[b]))) for a, b in setup.edges + ] + sites = tuple(range(L)) + occ_1d = {coo2idx[coo]: c for coo, c in setup.occupations.items()} + occupations = tuple(occ_1d[p] for p in range(L)) + + def build_stream(): + half = dt / 2 + u_hop = fermion.hopping_gate(half, t=t, imaginary=False) + onsite = [ + (fermion.onsite_gate(half, site=s, U=U, mu=0.0, imaginary=False), s) + for s in sites + ] + layers = fermion.edge_coloring_layers(edges_1d) + fwd = [(u_hop, e) for layer in layers for e in layer] + rev = [(u_hop, e) for layer in reversed(layers) for e in reversed(layer)] + return onsite + fwd + rev + onsite + + gates = build_stream() * 3 + native_hopping = fermion.hopping_gate(dt / 2, t=t, imaginary=False) + native_submpo = qtn.MatrixProductOperator.from_dense( + native_hopping, dims=(4, 4), sites=(0, 2), L=L, + ) + assert all( + type(tensor.data).__name__ == "U1U1FermionicArray" + for tensor in native_submpo.tensors + ) + + seed_mps = pepsy.ps_to_mps( + L, fermion=fermion, occupations=occupations, + seed=1234, dtype=state_dtype, cyclic=False, + ) + plan = TreeLayoutFinder( + [(fermion.hopping_gate(0.1, t=t, imaginary=False), e) for e in edges_1d], + n=L, chi=8, objective="hybrid", + ).recommend_arities((2, 3, 4), seed=0)["plan"] + seed_ttn = pepsy.ps_to_ttn( + L, tree=plan, fermion=fermion, occupations=occupations, dtype=state_dtype + ) + + # The tree and MPS seeds must represent the identical fermionic state. + assert float(tensors.tn_fidelity(seed_mps, seed_ttn)) > 1 - 1e-10 + + # Large-chi references (no truncation for L=4: exact bond is 16). + mps_exact = pepsy.MpsOptimizer( + seed_mps.copy(), gates=gates, chi=256, mode="mpo", inplace=False, + ) + mps_exact.run(cutoff=0.0) + engine = TreeOptimizer( + gates, n=L, tree=plan, state=seed_ttn.copy(), chi=256, cutoff=0.0, + mode="mpo", run=False, + ) + engine.run() + + # The two public modes are algebraically the same gate SVD: both defer + # truncation until the whole path has been updated. They can only differ by + # floating-point roundoff from the extra MPO factorisation/QR gauges. + direct = TreeOptimizer( + None, n=L, tree=plan, state=seed_ttn.copy(), chi=256, cutoff=0.0, + mode="direct", run=False, + ) + direct.run(gates) + + auto = TreeOptimizer( + None, n=L, tree=plan, state=seed_ttn.copy(), chi=256, cutoff=0.0, + mode="auto", run=False, + ) + auto.run(gates) + + assert float(tensors.tn_fidelity(engine.p, direct.p)) > 1 - 1e-8 + assert float(tensors.tn_fidelity(auto.p, direct.p)) > 1 - 1e-10 + assert float(tensors.tn_fidelity(direct.p, mps_exact.p)) > 1 - 1e-9 + assert float(tensors.tn_fidelity(engine.p, mps_exact.p)) > 1 - 1e-8 + + +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) + mpo = qtn.MatrixProductOperator.from_dense( + np.kron(x, x), dims=(2, 2), sites=(2, 3), L=4 + ) + opt = TreeOptimizer(None, n=4, chi=8, run=False) + opt.cap(1, [1.0, 0.0], compact_labels=False) + + monkeypatch.setattr( + mpo, + "to_dense", + lambda: (_ for _ in ()).throw(AssertionError("dense MPO fallback")), + ) + opt.apply_submpo(mpo, (2, 3)) + 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( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + events = [ + TreeOptimizer.cap_event(1, [1.0, 0.0], compact_labels=False), + (cnot, (2, 3)), + ] + opt = TreeOptimizer( + None, + n=4, + tree=TreePlan.from_order(range(4), structure="balanced"), + run=False, + ) + report = opt.estimate_bonds(events) + + assert report["events"][0]["kind"] == "cap" + assert report["events"][1]["support"] == (1, 2) + assert report["events"][1]["crossing_edges"] + + +def test_tree_auto_layout_remaps_supports_after_compact_cap(): + """Automatic layout uses original leaves for post-cap compact labels.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + events = [ + TreeOptimizer.cap_event(1, [1.0, 0.0]), + (cnot, (1, 2)), + ] + opt = TreeOptimizer(events, n=4, run=False) + assert (2, 3) in opt.layout_finder.supports diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py index a5753e9..9680ea0 100644 --- a/tests/test_package_layout.py +++ b/tests/test_package_layout.py @@ -43,8 +43,10 @@ backend_torch, default_physical_sectors, haar_random_state, + hrs_to_ttn, ps_to_3dpeps, ps_to_peps, + ps_to_ttn, reg_complex_qr_torch, reg_complex_svd_jax, reg_complex_svd_torch, @@ -56,13 +58,22 @@ site_charge_from_occupations, ) from pepsy.vmc import ( + ContractionConfig, FermionSiteEncoding, + MCState, + NetKetVMCSetup, + OptimizationConfig, + SamplingConfig, + SpinlessSiteEncoding, TorchPEPSAmplitude, TorchPEPSBoundaryAmplitude, TorchVMCDriver, + TorchVMCSetup, TorchVMCStepResult, TorchSquareLattice, apply_torch_sr_update, + build_netket_vmc, + build_torch_vmc, build_heisenberg_vmc, build_ising_vmc, heisenberg_connections, @@ -74,6 +85,12 @@ spinful_fermi_hubbard_connections, square_lattice_edges, torch_log_derivative_matrix, + VMCBackendCapabilityError, + VMCMeasurement, + VMCOptimizationResult, + VMC, + VMCProblem, + VMCSamples, ) @@ -115,8 +132,10 @@ def test_new_namespace_imports_resolve(): assert callable(default_physical_sectors) assert callable(backend_torch) assert callable(haar_random_state) + assert callable(hrs_to_ttn) assert callable(ps_to_peps) assert callable(ps_to_3dpeps) + assert callable(ps_to_ttn) assert callable(reg_rel_svd_torch) assert callable(reg_real_svd_torch) assert callable(reg_complex_svd_torch) @@ -127,14 +146,29 @@ def test_new_namespace_imports_resolve(): assert callable(reg_complex_svd_jax) assert callable(site_charge_from_occupations) assert FermionSiteEncoding is not None + assert ContractionConfig is not None + assert SamplingConfig is not None + assert OptimizationConfig is not None + assert MCState is not None + assert VMC is not None + assert VMCProblem is not None + assert VMCSamples is not None + assert VMCMeasurement is not None + assert VMCOptimizationResult is not None + assert issubclass(VMCBackendCapabilityError, NotImplementedError) + assert SpinlessSiteEncoding is not None assert TorchPEPSAmplitude is not None assert TorchPEPSBoundaryAmplitude is not None assert TorchVMCDriver is not None + assert TorchVMCSetup is not None + assert NetKetVMCSetup is not None assert TorchVMCStepResult is not None assert TorchSquareLattice is not None assert callable(apply_torch_sr_update) assert callable(build_heisenberg_vmc) assert callable(build_ising_vmc) + assert callable(build_torch_vmc) + assert callable(build_netket_vmc) assert callable(heisenberg_connections) assert callable(make_fermionic_peps_batched_amplitude_function) assert callable(make_peps_batched_amplitude_function) diff --git a/tests/test_prepare_boundary_inputs.py b/tests/test_prepare_boundary_inputs.py index 3de7054..97ecc0f 100644 --- a/tests/test_prepare_boundary_inputs.py +++ b/tests/test_prepare_boundary_inputs.py @@ -1218,6 +1218,49 @@ def fake_build_bra_ket(ket=None, *, bra=None): assert ket.balanced is False +def test_peps_normalize_skips_bond_balancing_for_symmray(monkeypatch): + """Symmray normalization must not invoke Quimb's unsupported balancer.""" + + class _SymmrayLikeData: + blocks = {"q0": object()} + + def apply_to_arrays(self, fn): + _ = fn + + class _Tensor: + data = _SymmrayLikeData() + + class _Ket: + tensor_map = {"I0,0": _Tensor()} + + def __init__(self): + self.divisor = None + + def __itruediv__(self, value): + self.divisor = value + return self + + def balance_bonds_(self): + raise AssertionError("Symmray normalization should not balance bonds.") + + class _Norm: + def contract_boundary(self, **kwargs): + _ = kwargs + return 16.0 + + def fake_build_bra_ket(ket=None, *, bra=None): + assert bra is None + return ket, _Norm() + + monkeypatch.setattr(pepsy.boundary.metrics, "build_bra_ket", fake_build_bra_ket) + + ket = _Ket() + old_norm = pepsy.peps_normalize(ket, chi=7, method="mps") + + assert old_norm == 16.0 + assert ket.divisor == 4.0 + + def test_peps_normalize_retries_stripped_when_full_cost_is_nonfinite(monkeypatch): """normalize should avoid dividing the state by a non-finite full norm.""" diff --git a/tests/test_public_api.py b/tests/test_public_api.py index c9e4fed..38ecf94 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -12,6 +12,14 @@ def test_package_version_available(): assert pepsy.__version__ +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 + + assert pepsy.TreeLayoutFinder is TreeLayoutFinder + assert pepsy.TreeOptimizer is TreeOptimizer + + _EXPECTED_IN_ALL = [ "backends", "boundary", "fitting", "operators", "optimizers", "sampling", "solvers", "tensors", "vmc", @@ -22,12 +30,13 @@ def test_package_version_available(): "pauli", "x", "y", "z", "s", "sdg", "t", "tdg", "h", "hadamard", "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", "ps_to_peps", "ps_to_3dpeps", "expec_mpo", - "id_to_mpo", "id_to_pepo", "ps_to_pepo", "ps_to_mpo", "make_numpy_array_caster", "to_float", "SweepOptimizer", + "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", "FDSolver", "MpsEnergyOptimizer", "MpsOptimizer", "MpoOptimizer", "PepsEnergyOptimizer", "PepsOptimizer", "SimpleUpdateGen", "SymDMRG2", "PEPSSampleResult", - "PepsBpSampler", "MpsSampler", "MpsBatchSampleResult", "MpsSampleResult", "VecSampler", "gate", "gauge_all", "gauge_all_simple", "one_norm_bp", "tn_fidelity", "tn_norm", + "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", + "TreeLayoutFinder", "TreeOptimizer", "TreeTensorNetwork", "DeferredInjectionRecord", "DeferredInjectionReport", "DeferredProjectionRecord", @@ -38,7 +47,7 @@ def test_package_version_available(): "TrajectoryChannel", "TrajectoryEvent", "TrajectoryOutcome", "TrajectoryRecord", "TrajectorySample", "TrajectoryShotResult", "compile_stim_circuit", "run_coalesced_noisy_shots", "run_coalesced_stim_shots", "run_coalesced_trajectory_shots", "run_noisy_shots", "run_stabilizer_mps_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", - "Fermion", "SpinfulFermion", "SpinfulFermionHubbard", "SymmFermions", "SymGateStream", "SymHamiltonian", "SymMPS", "SymPEPS", + "Fermion", "FermionLatticeSetup", "SpinfulFermion", "SpinfulFermionHubbard", "SymmFermions", "SymGateStream", "SymHamiltonian", "SymMPS", "SymPEPS", "default_physical_sectors", "draw_symmray_blocks", "draw_symmray_mps", "draw_symmray_mpo", "draw_symmray_peps", "fermi_hubbard_u1u1_gate_stream", "fermi_hubbard_u1u1_hopping_gate_stream", "fermi_hubbard_u1u1_interaction_gate_stream", "fermi_hubbard_u1u1_light_pulse_gate_stream", @@ -92,17 +101,18 @@ def test_internal_symbol_not_exported(name): "x", "y", "z", "s", "sdg", "t", "tdg", "h", "hadamard", "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", "ps_to_peps", "ps_to_3dpeps", "expec_mpo", - "id_to_mpo", "id_to_pepo", "ps_to_pepo", "ps_to_mpo", "SweepOptimizer", + "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", "DeferredInjectionRecord", "DeferredInjectionReport", "DeferredProjectionRecord", "ImmediateInjectionReport", "ImmediateProjectionRecord", "MeasurementRecord", "NormEventRecord", "StabilizerMpsSettingsAdvice", "StabilizerMpsRunResult", "StreamAnalysisRecord", "PepsEnergyOptimizer", "PepsOptimizer", "SimpleUpdateGen", "SymDMRG2", "PEPSSampleResult", "PepsBpSampler", "compile_stim_circuit", "run_coalesced_noisy_shots", "run_coalesced_stim_shots", "run_coalesced_trajectory_shots", "run_noisy_shots", "run_stabilizer_mps_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", + "TreeLayoutFinder", "TreeOptimizer", "TreeTensorNetwork", "TreeSampler", "TreeBatchSampleResult", "TreeSampleResult", - "tn_fidelity", "tn_norm", "Fermion", "SpinfulFermion", "SpinfulFermionHubbard", "SymmFermions", "SymGateStream", "SymHamiltonian", "SymMPS", "SymPEPS", + "tn_fidelity", "tn_norm", "Fermion", "FermionLatticeSetup", "SpinfulFermion", "SpinfulFermionHubbard", "SymmFermions", "SymGateStream", "SymHamiltonian", "SymMPS", "SymPEPS", "default_physical_sectors", "draw_symmray_blocks", "draw_symmray_mps", "draw_symmray_mpo", "draw_symmray_peps", "fermi_hubbard_u1u1_gate_stream", "fermi_hubbard_u1u1_hopping_gate_stream", "fermi_hubbard_u1u1_interaction_gate_stream", "fermi_hubbard_u1u1_light_pulse_gate_stream", diff --git a/tests/test_sampler.py b/tests/test_sampler.py index 60446fe..4a7e58f 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -1,5 +1,7 @@ """Tests for PEPS BP importance sampling helpers.""" +from itertools import product + import numpy as np import pytest import quimb.tensor as qtn @@ -37,6 +39,22 @@ def contract_ctmrg(self, **kwargs): return 4.5 + self.label, -1 +def _dense_mps_from_symmray(psi): + """Convert an ungraded Symmray MPS to an ordinary dense quimb MPS.""" + arrays = [ + sampler_mod.MpsSampler._site_array_lr_phys_r(psi, site).to_dense() + for site in range(psi.L) + ] + return qtn.MatrixProductState( + [ + arrays[0][0].T, + *(array.transpose(0, 2, 1) for array in arrays[1:-1]), + arrays[-1][:, :, 0], + ], + shape="lrp", + ) + + def test_sampler_exact_method_collects_configs_and_scalars(monkeypatch): """Sampler should collect row-major configs, proposal weights, and amplitudes.""" calls = [] @@ -108,6 +126,8 @@ def test_sampler_public_exports_resolve(): """Sampler helpers should be available from the package namespace.""" assert pepsy.PepsBpSampler is sampler_mod.PepsBpSampler assert pepsy.PEPSSampleResult is sampler_mod.PEPSSampleResult + assert pepsy.FermionConfigurationEncoding is sampler_mod.FermionConfigurationEncoding + assert pepsy.MpsDiagonalEstimate is sampler_mod.MpsDiagonalEstimate assert pepsy.MpsBatchSampleResult is sampler_mod.MpsBatchSampleResult @@ -816,6 +836,747 @@ def test_mps_sampler_quimb_sample_batch_returns_numpy_result(): assert batch.to_sample_result().configs_1d == [[1, 0]] * 3 +def test_mps_sampler_detects_symmray_u1u1_product_without_densifying(monkeypatch): + """A U1U1 product MPS should retain its Symmray representation to sample.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, ps_to_mps + + fermion = Fermion(spinful=True, symmetry="U1U1") + psi = ps_to_mps( + 4, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + ) + source_maps = [ + dict(psi[site].data.indices[psi[site].inds.index(psi.site_ind(site))].chargemap) + for site in range(psi.L) + ] + + def fail_dense(*args, **kwargs): # pylint: disable=unused-argument + raise AssertionError("Symmray MPS sampling must not densify tensor data") + + monkeypatch.setattr(type(psi[0].data), "to_dense", fail_dense) + sampler = sampler_mod.MpsSampler(psi) + configs, probs = sampler.sample_arrays(4, seed=7) + + assert sampler.resolved_backend == "symmray" + np.testing.assert_array_equal(configs, np.array([[2, 1, 2, 1]] * 4)) + np.testing.assert_allclose(probs, np.ones(4)) + assert [ + dict(psi[site].data.indices[psi[site].inds.index(psi.site_ind(site))].chargemap) + for site in range(psi.L) + ] == source_maps + + +@pytest.mark.parametrize("symmetry", ("Z2", "U1", "U1U1")) +def test_mps_sampler_symmray_fermionic_matches_direct_amplitudes(symmetry): + """Entangled fermionic branches agree with direct sparse contraction.""" + 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=8, + seed=3, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs, sampled_probs = sampler.sample_arrays(6, seed=9) + probs = sampler.probabilities(configs) + amplitudes = sampler.amplitudes(configs) + norm = (psi.H @ psi).real + + assert sampler.resolved_backend == "symmray" + np.testing.assert_allclose(probs, sampled_probs, atol=1e-12) + np.testing.assert_allclose(np.abs(amplitudes) ** 2, sampled_probs, atol=1e-12) + for config, probability in zip(configs, sampled_probs): + branch = psi.copy() + branch.isel_({ + branch.site_ind(site): int(code) + for site, code in enumerate(config) + }) + amplitude = np.asarray(branch.contract(all).data).item() + assert probability == pytest.approx(abs(amplitude) ** 2 / norm) + + +@pytest.mark.parametrize("symmetry", ("Z2", "U1", "U1U1")) +def test_mps_sampler_symmray_fermionic_sector_probabilities_are_exhaustive( + symmetry, +): + """All allowed physical codes normalize and forbidden sectors vanish.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + L = 3 + fermion = Fermion(spinful=True, symmetry=symmetry) + psi = SymMPS.random( + L, + symmetry=symmetry, + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(L), + bond_dim=4, + seed=11, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs = np.asarray(list(product(range(4), repeat=L)), dtype=np.int64) + probabilities = sampler.probabilities(configs) + amplitudes = sampler.amplitudes(configs) + charge_maps = sampler.physical_code_maps + target_charge = fermion.total_charge(fermion.half_filled_occupations(L)) + allowed = np.asarray( + [ + fermion.total_charge( + charge_maps[site][int(code)][0] + for site, code in enumerate(config) + ) + == target_charge + for config in configs + ] + ) + + np.testing.assert_allclose(probabilities.sum(), 1.0, atol=1e-12) + np.testing.assert_allclose( + probabilities, + np.abs(amplitudes) ** 2, + atol=1e-12, + ) + np.testing.assert_allclose(probabilities[~allowed], 0.0, atol=1e-12) + assert np.any(probabilities[allowed] > 1e-12) + + +@pytest.mark.parametrize("symmetry", ("U1", "U1U1")) +def test_mps_sampler_symmray_dense_and_quimb_baselines_match(symmetry): + """The benchmark's dense baselines preserve U1 fermionic Born weights.""" + 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=13, + dtype="complex128", + ).mps + sparse_sampler = sampler_mod.MpsSampler(psi, backend="symmray") + dense_psi = _dense_mps_from_symmray(psi) + native_sampler = sampler_mod.MpsSampler(dense_psi, backend="native") + quimb_sampler = sampler_mod.MpsSampler(dense_psi, backend="quimb") + + configs, sparse_probabilities = sparse_sampler.sample_arrays(16, seed=5) + np.testing.assert_allclose( + native_sampler.probabilities(configs), + sparse_probabilities, + atol=1e-12, + ) + quimb_configs, quimb_probabilities = quimb_sampler.sample_arrays(16, seed=5) + np.testing.assert_allclose( + sparse_sampler.probabilities(quimb_configs), + quimb_probabilities, + atol=1e-12, + ) + + +@pytest.mark.parametrize("symmetry", ("Z2", "U1", "U1U1")) +def test_mps_sampler_fermion_diagonal_values_follow_symmray_physical_codes( + symmetry, +): + """Diagonal values use the physical code order of each fermionic symmetry.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + fermion = Fermion(spinful=True, symmetry=symmetry) + psi = SymMPS.random( + 3, + symmetry=symmetry, + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(3), + bond_dim=4, + seed=15, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs = np.repeat(np.arange(4, dtype=np.int64)[:, None], psi.L, axis=1) + expected_occupation = np.diag(fermion.observable("number").to_dense()).real + expected_doublon = np.diag(fermion.observable("double").to_dense()).real + + np.testing.assert_allclose( + sampler.fermion_diagonal_values( + configs, + fermion, + "occupation", + sites=0, + ), + expected_occupation, + ) + np.testing.assert_allclose( + sampler.fermion_diagonal_values( + configs, + fermion, + "doublon", + sites=0, + ), + expected_doublon, + ) + np.testing.assert_allclose( + sampler.fermion_diagonal_values( + configs, + fermion, + "density_correlation", + pairs=(0, 1), + ), + expected_occupation**2, + ) + + +@pytest.mark.parametrize("symmetry", ("Z2", "U1", "U1U1")) +def test_mps_sampler_fermion_diagonal_estimates_match_exact_born_sums(symmetry): + """Sampled fermion observables agree with exact MPS Born sums.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + L = 3 + fermion = Fermion(spinful=True, symmetry=symmetry) + psi = SymMPS.random( + L, + symmetry=symmetry, + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(L), + bond_dim=4, + seed=17, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs = np.asarray(list(product(range(4), repeat=L)), dtype=np.int64) + probabilities = sampler.probabilities(configs) + observables = ( + ("occupation", {"sites": (0, 2)}), + ("total_charge", {}), + ("doublon", {}), + ("density_correlation", {"pairs": ((0, 1), (1, 2))}), + ) + + for observable, kwargs in observables: + exact = float(np.dot( + probabilities, + sampler.fermion_diagonal_values(configs, fermion, observable, **kwargs), + )) + estimate = sampler.estimate_fermion_diagonal( + fermion, + observable, + n_samples=4096, + seed=19, + **kwargs, + ) + + assert isinstance(estimate, sampler_mod.MpsDiagonalEstimate) + assert estimate.n_samples == 4096 + assert estimate.observable == observable + assert abs(estimate.mean - exact) <= 7 * estimate.standard_error + 0.01 + + +@pytest.mark.parametrize( + ("symmetry", "occupations"), + ( + ("Z2", (1, 1, 0, 2)), + ("U1", (1, 1, 0, 2)), + ("U1U1", ((1, 0), (0, 1), (0, 0), (1, 1))), + ("Z2Z2", ((1, 0), (0, 1), (0, 0), (1, 1))), + ), +) +def test_mps_sampler_symmray_supports_spinful_fermion_symmetries( + monkeypatch, + symmetry, + occupations, +): + """Fermionic charge maps, including degenerate sectors, stay sparse.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, ps_to_mps + + fermion = Fermion(spinful=True, symmetry=symmetry) + psi = ps_to_mps(4, fermion=fermion, occupations=occupations) + source_maps = [ + dict(psi[site].data.indices[psi[site].inds.index(psi.site_ind(site))].chargemap) + for site in range(psi.L) + ] + + def fail_dense(*args, **kwargs): # pylint: disable=unused-argument + raise AssertionError("Symmray MPS sampling must not densify tensor data") + + monkeypatch.setattr(type(psi[0].data), "to_dense", fail_dense) + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs, sampled_probs = sampler.sample_arrays(16, seed=4) + + assert sampler.resolved_backend == "symmray" + assert configs.shape == (16, 4) + assert np.all((0 <= configs) & (configs < 4)) + np.testing.assert_allclose( + sampler.probabilities(configs), + sampled_probs, + atol=1e-12, + ) + np.testing.assert_allclose( + np.abs(sampler.amplitudes(configs)) ** 2, + sampled_probs, + atol=1e-12, + ) + assert [ + dict(psi[site].data.indices[psi[site].inds.index(psi.site_ind(site))].chargemap) + for site in range(psi.L) + ] == source_maps + + +@pytest.mark.parametrize("symmetry", ("Z2", "U1")) +def test_mps_sampler_symmray_supports_spinless_fermion_symmetries(symmetry): + """The same physical-code reconstruction handles spinless fermions.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, ps_to_mps + + psi = ps_to_mps( + 4, + fermion=Fermion(spinful=False, symmetry=symmetry), + occupations=(0, 1, 0, 1), + ) + sampler = sampler_mod.MpsSampler(psi) + configs, sampled_probs = sampler.sample_arrays(16, seed=4) + + assert sampler.resolved_backend == "symmray" + np.testing.assert_array_equal(configs, np.array([[0, 1, 0, 1]] * 16)) + np.testing.assert_allclose(sampled_probs, np.ones(16)) + np.testing.assert_allclose(sampler.probabilities(configs), sampled_probs) + + +def test_mps_sampler_symmray_reuses_prefix_boundaries(monkeypatch): + """A batch builds one conditional per distinct sampled prefix.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, ps_to_mps + + psi = ps_to_mps( + 4, + fermion=Fermion(spinful=True, symmetry="U1"), + occupations=(1, 1, 0, 2), + ) + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + real_candidates = sampler_mod.MpsSampler._symmray_candidates.__func__ + calls = [] + + def spy_candidates(cls, state, site, boundary): + calls.append(site) + return real_candidates(cls, state, site, boundary) + + monkeypatch.setattr( + sampler_mod.MpsSampler, + "_symmray_candidates", + classmethod(spy_candidates), + ) + configs, probs = sampler.sample_arrays(32, seed=3) + + assert calls[0] == 0 + assert len(calls) < 32 * psi.L + np.testing.assert_allclose(sampler.probabilities(configs), probs, atol=1e-12) + + +@pytest.mark.parametrize( + ("symmetry", "physical_sectors", "site_charges"), + ( + ("Z2", {0: 1, 1: 1}, (0, 1, 0)), + ("U1", {0: 1, 1: 1, 2: 1}, (1, 1, 1)), + ( + "U1U1", + {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}, + ((1, 0), (0, 1), (1, 0)), + ), + ( + "Z2Z2", + {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}, + ((1, 0), (0, 1), (1, 0)), + ), + ), +) +def test_mps_sampler_symmray_supports_generic_abelian_mps( + monkeypatch, + symmetry, + physical_sectors, + site_charges, +): + """Non-fermionic Symmray MPSs use the same no-densification sampler.""" + pytest.importorskip("symmray") + from pepsy.tensors import SymMPS, site_charge_from_occupations + + psi = SymMPS.random( + 3, + symmetry=symmetry, + fermionic=False, + bond_dim=4, + phys_dim=physical_sectors, + site_charge=site_charge_from_occupations(site_charges), + seed=7, + dtype="complex128", + ).mps + + def fail_dense(*args, **kwargs): # pylint: disable=unused-argument + raise AssertionError("Generic Symmray MPS sampling must not densify data") + + monkeypatch.setattr(type(psi[0].data), "to_dense", fail_dense) + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + local_dim = sum(physical_sectors.values()) + configs = np.asarray(list(product(range(local_dim), repeat=psi.L)), dtype=int) + probabilities = sampler.probabilities(configs) + amplitudes = sampler.amplitudes(configs) + sampled_configs, sampled_probs = sampler.sample_arrays(32, seed=3) + + assert sampler.resolved_backend == "symmray" + assert sampler.physical_code_maps is not None + assert all( + set(code_map) == set(range(local_dim)) + for code_map in sampler.physical_code_maps + ) + np.testing.assert_allclose(np.abs(amplitudes) ** 2, probabilities, atol=1e-12) + np.testing.assert_allclose( + sampler.probabilities(sampled_configs), + sampled_probs, + atol=1e-12, + ) + + for config in sampled_configs[:4]: + branch = psi.copy() + branch.isel_({ + branch.site_ind(site): int(code) + for site, code in enumerate(config) + }) + value = branch.contract(all).data + if hasattr(value, "get_scalar_element"): + value = value.phase_sync().get_scalar_element() + else: + value = np.asarray(value).item() + expected = abs(value) ** 2 / (psi.H @ psi).real + assert sampler.probabilities(config[None, :])[0] == pytest.approx(expected) + + +def test_mps_sampler_symmray_physical_code_maps_retain_degenerate_sectors(): + """Source physical sectors remain interpretable after canonical pruning.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, ps_to_mps + + psi = ps_to_mps( + 3, + fermion=Fermion(spinful=True, symmetry="U1"), + occupations=(1, 1, 0), + ) + sampler = sampler_mod.MpsSampler(psi) + maps = sampler.physical_code_maps + + assert maps == ( + {0: (0, 0), 1: (1, 0), 2: (1, 1), 3: (2, 0)}, + ) * 3 + maps[0].clear() + assert sampler.physical_code_maps[0][2] == (1, 1) + + +@pytest.mark.parametrize( + ("symmetry", "expected"), + ( + ("Z2", ((0, 0), (1, 1), (1, 0), (0, 1))), + ("U1", ((0, 0), (0, 1), (1, 0), (1, 1))), + ("U1U1", ((0, 0), (0, 1), (1, 0), (1, 1))), + ("Z2Z2", ((0, 0), (0, 1), (1, 0), (1, 1))), + ), +) +def test_mps_sampler_fermion_configuration_encoding_is_symmetry_aware( + symmetry, + expected, +): + """Raw MPS codes decode consistently before entering a VMC workflow.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + fermion = Fermion(spinful=True, symmetry=symmetry) + psi = SymMPS.random( + 3, + symmetry=symmetry, + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(3), + bond_dim=4, + seed=37, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + encoding = sampler.fermion_configuration_encoding(fermion) + codes = np.repeat(np.arange(4, dtype=np.int64)[:, None], psi.L, axis=1) + + assert encoding.symmetry == symmetry + assert encoding.spinful + assert encoding.code_to_occupations == (expected,) * psi.L + occupations = encoding.decode(codes) + np.testing.assert_array_equal(occupations[:, 0], np.asarray(expected)) + np.testing.assert_array_equal(encoding.encode(occupations), codes) + + batch = sampler.sample_batch(8, seed=7, to_numpy=True, fermion=fermion) + assert batch.configuration_encoding == encoding + np.testing.assert_array_equal(batch.occupations(), encoding.decode(batch.configs)) + + +def test_mps_sampler_fermion_configuration_encoding_rejects_wrong_symmetry(): + """MPS samples cannot silently reuse a codec from another symmetry.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, ps_to_mps + + psi = ps_to_mps( + 2, + fermion=Fermion(spinful=True, symmetry="U1"), + occupations=(1, 1), + ) + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + + with pytest.raises(ValueError, match="must match"): + sampler.fermion_configuration_encoding( + Fermion(spinful=True, symmetry="Z2") + ) + + +def test_mps_sampler_bound_fermion_sets_the_batch_configuration_contract(): + """A bound Fermion removes VMC code-convention boilerplate per batch.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + fermion = Fermion(spinful=True, symmetry="Z2") + psi = SymMPS.random( + 3, + symmetry="Z2", + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(3), + bond_dim=4, + seed=43, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray", fermion=fermion) + + encoding = sampler.fermion_configuration_encoding() + batch = sampler.sample_batch(8, seed=7, to_numpy=True) + + assert sampler.fermion is fermion + assert batch.configuration_encoding == encoding + np.testing.assert_array_equal(batch.occupations(), encoding.decode(batch.configs)) + + codes = np.repeat(np.arange(4, dtype=np.int64)[:, None], psi.L, axis=1) + np.testing.assert_allclose( + sampler.fermion_diagonal_values(codes, "total_charge"), + (0.0, 6.0, 3.0, 3.0), + ) + + +def test_mps_sampler_symmray_cached_branches_prune_charge_forbidden_work(): + """The sparse path skips impossible charge branches without densifying.""" + 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=41, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs, probs = sampler.sample_arrays(32, seed=5) + stats = sampler.symmray_sampling_stats + + assert stats["cached_local_slices"] + assert stats["charge_pruned_branches"] > 0 + assert stats["candidate_contractions"] < 4 * stats["conditional_evaluations"] + np.testing.assert_allclose(sampler.probabilities(configs), probs, atol=1e-12) + + +def test_mps_sampler_symmray_prefix_controls_bound_high_entropy_batches(): + """Prefix sharing reduces work and the auto cap falls back safely.""" + pytest.importorskip("symmray") + from pepsy.tensors import SymMPS, site_charge_from_occupations + + psi = SymMPS.random( + 4, + symmetry="U1", + fermionic=False, + bond_dim=4, + phys_dim={0: 1, 1: 1, 2: 1}, + site_charge=site_charge_from_occupations((1, 1, 1, 1)), + seed=7, + dtype="complex128", + ).mps + shared = sampler_mod.MpsSampler( + psi, + backend="symmray", + prefix_strategy="prefix", + max_prefix_groups=None, + ) + shared_configs, shared_probs = shared.sample_arrays(64, seed=3) + bounded = sampler_mod.MpsSampler( + psi, + backend="symmray", + prefix_strategy="auto", + max_prefix_groups=1, + ) + bounded_configs, bounded_probs = bounded.sample_arrays(64, seed=3) + serial = sampler_mod.MpsSampler( + psi, + backend="symmray", + prefix_strategy="serial", + ) + serial_configs, serial_probs = serial.sample_arrays(64, seed=3) + adaptive = sampler_mod.MpsSampler( + psi, + backend="symmray", + prefix_strategy="auto", + max_prefix_groups=None, + ) + adaptive_configs, adaptive_probs = adaptive.sample_arrays(64, seed=3) + + assert shared.symmray_sampling_stats["conditional_evaluations"] < 64 * psi.L + assert not shared.symmray_sampling_stats["serial_fallback"] + assert bounded.symmray_sampling_stats["serial_fallback"] + assert bounded.symmray_sampling_stats["max_active_prefix_groups"] <= 1 + assert serial.symmray_sampling_stats["conditional_evaluations"] == 64 * psi.L + assert adaptive.symmray_sampling_stats["adaptive_serial_fallback"] + np.testing.assert_allclose( + shared.probabilities(shared_configs), + shared_probs, + atol=1e-12, + ) + np.testing.assert_allclose( + bounded.probabilities(bounded_configs), + bounded_probs, + atol=1e-12, + ) + np.testing.assert_allclose( + serial.probabilities(serial_configs), + serial_probs, + atol=1e-12, + ) + np.testing.assert_allclose( + adaptive.probabilities(adaptive_configs), + adaptive_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") + + with pytest.raises(ValueError, match="Unknown Symmray prefix strategy"): + sampler_mod.MpsSampler(psi, prefix_strategy="branchy") + with pytest.raises(ValueError, match="max_prefix_groups"): + sampler_mod.MpsSampler(psi, max_prefix_groups=0) + + +def test_mps_sampler_symmray_torch_nonfermionic_blocks_stay_on_torch(): + """Generic Symmray Torch blocks preserve device-resident outputs.""" + pytest.importorskip("symmray") + torch = pytest.importorskip("torch") + from pepsy.tensors import SymMPS, site_charge_from_occupations + + psi = SymMPS.random( + 3, + symmetry="U1", + fermionic=False, + bond_dim=3, + phys_dim={0: 1, 1: 1, 2: 1}, + site_charge=site_charge_from_occupations((1, 1, 1)), + seed=7, + dtype="complex128", + to_backend=torch.as_tensor, + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs, probs = sampler.sample_arrays(8, seed=3) + + assert isinstance(configs, torch.Tensor) + assert isinstance(probs, torch.Tensor) + torch.testing.assert_close( + sampler.probabilities(configs, to_numpy=False), + probs, + ) + + +def test_mps_sampler_symmray_cupy_nonfermionic_blocks_stay_on_cupy(): + """Generic Symmray CuPy blocks preserve device-resident outputs.""" + pytest.importorskip("symmray") + 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}") + from pepsy.tensors import SymMPS, site_charge_from_occupations + + psi = SymMPS.random( + 3, + symmetry="U1", + fermionic=False, + bond_dim=3, + phys_dim={0: 1, 1: 1, 2: 1}, + site_charge=site_charge_from_occupations((1, 1, 1)), + seed=7, + dtype="complex128", + to_backend=cupy.asarray, + ).mps + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs, probs = sampler.sample_arrays(8, seed=3) + + assert isinstance(configs, cupy.ndarray) + assert isinstance(probs, cupy.ndarray) + cupy.testing.assert_allclose( + sampler.probabilities(configs, to_numpy=False), + probs, + ) + + +def test_mps_sampler_symmray_torch_blocks_stay_on_torch(): + """The Symmray path should keep Torch-backed blocks and samples on Torch.""" + pytest.importorskip("symmray") + torch = pytest.importorskip("torch") + from pepsy.tensors import Fermion, ps_to_mps + + fermion = Fermion(spinful=True, symmetry="U1U1") + psi = ps_to_mps( + 3, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0)), + to_backend=lambda array: torch.as_tensor(array), + ) + sampler = sampler_mod.MpsSampler(psi, backend="symmray") + configs, probs = sampler.sample_arrays(3, seed=7) + batch = sampler.sample_batch(3, seed=7, fermion=fermion) + + assert sampler.resolved_backend == "symmray" + assert isinstance(configs, torch.Tensor) + assert isinstance(probs, torch.Tensor) + assert batch.backend == "torch" + assert isinstance(batch.occupations(), torch.Tensor) + assert tuple(batch.occupations().shape) == (3, 3, 2) + torch.testing.assert_close( + configs, + torch.tensor([[2, 1, 2]] * 3, dtype=torch.long, device=configs.device), + ) + torch.testing.assert_close(probs, torch.ones(3, dtype=probs.dtype)) + + def test_vec_sampler_rejects_invalid_vector_size(): """Dense vector length must match the site-map Hilbert-space dimension.""" with pytest.raises(ValueError, match=r"2\*\*L=4"): diff --git a/tests/test_stabilizer_tn.py b/tests/test_stabilizer_tn.py index 08bd004..bff0d80 100644 --- a/tests/test_stabilizer_tn.py +++ b/tests/test_stabilizer_tn.py @@ -512,29 +512,30 @@ def test_simulator_truncation_caps_bond_and_tracks_infidelity(): ) -def test_norm_infidelity_excludes_nonunitary_segments(): +def test_nonunitary_dense_gate_reports_gdagger_g_infidelity(): + """Non-unitary compression is normalized by the exact G-dagger-G norm.""" + gate = np.diag([1.0, 1.0, 1.0, 0.2]).astype(complex) sim = MpsStabOptimizer( 2, chi=1, track_infidelity=True, exact_cooling=False ) - sim.apply([(_I + 0.2 * _X, 0)]) - assert sim.infidelities == [] + sim.apply([("h", 0), ("h", 1)]) + target_norm = np.linalg.norm(gate @ sim.to_statevector()) + assert sim._dense_gate_target_norm(gate, (0, 1)) == pytest.approx( + target_norm, abs=1e-10 + ) + sim.apply([(gate, (0, 1))]) - # A later unitary cannot interpret norm change from the arbitrary map as - # truncation loss, so tracking remains unavailable for this segment. - sim.apply([("rxx", 0.8, 0, 1)]) - assert sim.infidelities == [] + expected = 1.0 - (sim.norm() / target_norm) ** 2 + assert len(sim.infidelities) == 1 + assert sim.infidelities[-1] == pytest.approx(expected, abs=1e-10) + assert sim.norm_diagnostics()["infidelity"] == pytest.approx(expected, abs=1e-10) + assert sim.norm_diagnostics()["current_valid"] is False - # Projection restores a normalized baseline without itself reporting a - # metric. Subsequent unitary compression can be tracked again. + # Measurement restores a normalized physical state without losing the + # preceding non-unitary compression diagnostic. sim.measure("Z", 0, outcome=+1) - assert sim.infidelities == [] - assert sim.norm_events[-1]["valid"] is False - assert sim.norm_diagnostics()["completed_segments"] == 0 - sim.apply([("rxx", 0.8, 0, 1)]) - assert len(sim.infidelities) == 1 - assert sim.infidelities[-1] == pytest.approx( - 1.0 - sim.norm() ** 2, abs=1e-10 - ) + assert sim.norm() == pytest.approx(1.0, abs=1e-10) + assert sim.norm_diagnostics()["infidelity"] == pytest.approx(expected, abs=1e-10) def test_norm_events_close_segment_before_measurement_normalizes_after(): @@ -597,6 +598,11 @@ def test_norm_events_close_segment_before_measurement_normalizes_after(): assert diagnostics["norm_infidelity"] == pytest.approx( 1.0 - expected_total_survival ) + assert diagnostics["infidelity"] == pytest.approx( + 1.0 - expected_total_survival + ) + assert diagnostics["fidelity"] == pytest.approx(expected_total_survival) + assert sim.get_infidelities() is sim.infidelities assert diagnostics["norm"] == pytest.approx(expected_total_survival ** 0.5) assert diagnostics["geometric_mean_norm"] == pytest.approx( expected_total_survival ** 0.5 @@ -685,12 +691,12 @@ def close(self): "measurement", ] last = progress.postfix_calls[-1] - assert sorted(last) == ["norm_infidelity", "part"] - assert last["norm_infidelity"] == ( - "n/a" - if sim.norm_diagnostics()["norm_infidelity"] is None - else f"{sim.norm_diagnostics()['norm_infidelity']:.2e}" + assert sorted(last) == ["infidelity", "norm_infidelity", "part"] + expected = sim._format_progress_infidelity( + sim.norm_diagnostics()["infidelity"] ) + assert last["infidelity"] == expected + assert last["norm_infidelity"] == expected def test_simulator_two_qubit_nonclifford_matrix_supported(): diff --git a/tests/test_sweep_symmray_backend.py b/tests/test_sweep_symmray_backend.py new file mode 100644 index 0000000..916f97e --- /dev/null +++ b/tests/test_sweep_symmray_backend.py @@ -0,0 +1,300 @@ +"""Regression tests for the Symmray/backend handling in the sweep optimizer. + +These cover the helper machinery that keeps finite-difference local objectives +on their original backend and preserves Torch-backed Symmray autograd paths. +""" + +import numpy as np +import pytest +from types import SimpleNamespace +import quimb.tensor as qtn + +import pepsy +from pepsy.optimizers.sweep.environments import QuimbMpsBoundaryStore +from pepsy.optimizers.sweep.optimizer import SweepOptimizer +from pepsy.tensors.symmetric import ( + SymPEPS, + default_physical_sectors, + site_charge_from_occupations, +) + +torch = pytest.importorskip("torch") + + +def test_symmray_fd_solver_maps_autograd_solvers_to_finite_difference(): + # nlopt / torch / jax autograd solvers cannot differentiate through NumPy + # Symmray blocks, so they must be routed to a finite-difference backend. + assert SweepOptimizer._symmray_fd_solver("nlopt", {})[0] == "fd-nlopt" + assert SweepOptimizer._symmray_fd_solver("scipy", {})[0] == "fd-scipy" + assert SweepOptimizer._symmray_fd_solver("torch-adam", {})[0] == "fd-adam" + assert SweepOptimizer._symmray_fd_solver("jax-lbfgs", {})[0] == "fd-adam" + + +def test_symmray_fd_solver_keeps_explicit_fd_solver(): + for name in ("fd-nlopt", "fd-scipy", "fd-adam"): + assert SweepOptimizer._symmray_fd_solver(name, {})[0] == name + + +def test_symmray_fd_solver_drops_lr_option(): + # 'lr' is a first-order autograd step size and is rejected by the LBFGS-style + # finite-difference backends; it must be stripped. + _, opts = SweepOptimizer._symmray_fd_solver("nlopt", {"lr": 1e-2, "maxeval": 5}) + assert "lr" not in opts + assert opts["maxeval"] == 5 + + +def test_symmray_fd_adam_keeps_lr_option(): + _, opts = SweepOptimizer._symmray_fd_solver("torch-adam", {"lr": 1e-2}) + assert opts["lr"] == pytest.approx(1e-2) + + +def test_coerce_param_tree_numpy_converts_torch_leaves(): + tree = { + "a": torch.ones(2, dtype=torch.float64), + "b": [torch.zeros(3, dtype=torch.float64), np.arange(3.0)], + "c": {"d": torch.tensor([1.0, 2.0])}, + } + out = SweepOptimizer._coerce_param_tree_numpy(tree) + assert isinstance(out["a"], np.ndarray) + assert isinstance(out["b"][0], np.ndarray) + assert isinstance(out["b"][1], np.ndarray) + assert isinstance(out["c"]["d"], np.ndarray) + np.testing.assert_allclose(out["a"], np.ones(2)) + np.testing.assert_allclose(out["c"]["d"], np.array([1.0, 2.0])) + + +def test_coerce_leaf_to_numpy_is_noop_for_numpy(): + arr = np.arange(4.0) + assert SweepOptimizer._coerce_leaf_to_numpy(arr) is arr + + +def test_match_leaf_backend_restores_numpy_reference(): + reference = np.zeros(2) + value = torch.ones(2, dtype=torch.float64) + restored = SweepOptimizer._match_leaf_backend(value, reference) + assert isinstance(restored, np.ndarray) + np.testing.assert_allclose(restored, np.ones(2)) + + +def test_match_leaf_backend_noop_when_backends_match(): + reference = np.zeros(2) + value = np.ones(2) + assert SweepOptimizer._match_leaf_backend(value, reference) is value + + +def test_restore_leaf_backends_maps_each_leaf(): + params_ref = {"x": np.zeros(2), "y": np.zeros(3)} + params_opt = { + "x": torch.ones(2, dtype=torch.float64), + "y": torch.full((3,), 2.0, dtype=torch.float64), + } + restored = SweepOptimizer._restore_leaf_backends(params_opt, params_ref) + assert isinstance(restored["x"], np.ndarray) + assert isinstance(restored["y"], np.ndarray) + np.testing.assert_allclose(restored["x"], np.ones(2)) + np.testing.assert_allclose(restored["y"], np.full(3, 2.0)) + + +def test_match_param_tree_backends_restores_nested_numpy_reference(): + reference = {"site": [np.zeros(2), {"block": np.zeros(1)}]} + trial = { + "site": [ + torch.ones(2, dtype=torch.float64), + {"block": torch.ones(1, dtype=torch.float64)}, + ] + } + restored = SweepOptimizer._match_param_tree_backends(trial, reference) + assert isinstance(restored["site"][0], np.ndarray) + assert isinstance(restored["site"][1]["block"], np.ndarray) + + +def test_params_require_finite_differences_only_for_non_torch_trees(): + assert SweepOptimizer._params_require_finite_differences({"x": np.zeros(2)}) + assert not SweepOptimizer._params_require_finite_differences( + {"x": torch.zeros(2)} + ) + + +def test_sweep_rejects_mixed_symmray_and_dense_or_backend_inputs(): + """State and target must have one compatible array representation.""" + + class _FakeSymmrayData: + shape = (1,) + dtype = "complex128" + + def __init__(self, backend): + self.backend = backend + + _FakeSymmrayData.__module__ = "symmray.fake" + + def _symmray_state(backend): + return SimpleNamespace( + tensor_map={"site": SimpleNamespace(data=_FakeSymmrayData(backend))} + ) + + dense = SimpleNamespace( + tensor_map={"site": SimpleNamespace(data=np.zeros(1))} + ) + with pytest.raises(TypeError, match="mix Symmray and dense"): + SweepOptimizer._validate_symmray_input_backends(_symmray_state("torch"), dense) + with pytest.raises(TypeError, match="one common array backend"): + SweepOptimizer._validate_symmray_input_backends( + _symmray_state("torch"), + _symmray_state("numpy"), + ) + + +def test_torch_symmray_u1u1_slice_uses_autograd_without_backend_conversion(): + """A real Torch Symmray local update must remain Torch-backed.""" + pytest.importorskip("symmray") + to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) + charges = { + (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_from_occupations(charges), + bond_dim=1, + seed=71, + dtype="complex128", + to_backend=to_backend, + ).tn + target = state.copy() + target.mangle_inner_() + opt = SweepOptimizer( + state, + target, + chi=4, + boundary_engine="quimb-mps", + renormalize_state=False, + ) + + assert opt._symmray_backends == {"torch"} + assert opt._symmray_torch is True + assert opt._symmray_requires_fd is False + + norm_tn, overlap_tn = opt._prepare_current_double_layers() + opt.bdy.start_sweep(norm_tn, "y", "left") + opt.bdy_overlap.start_sweep(overlap_tn, "y", "left") + run = opt._optimize_axis_slice_with_current_env( + 0, + axis="y", + solver="nlopt", + solver_options={"algorithm": "LD_LBFGS", "n_steps": 1, "maxeval": 1}, + ) + + assert np.isfinite(run["loss_initial"]) + assert np.isfinite(run["loss_final"]) + assert all( + isinstance(block, torch.Tensor) + for tensor in opt.state.tensor_map.values() + for block in tensor.data.blocks.values() + ) + + +def test_torch_symmray_interior_slice_attaches_both_cached_boundaries(): + """The overlap local bra must match the cached double-layer index names.""" + pytest.importorskip("symmray") + to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) + charges = { + (x, y): (1, 0) if (x + y) % 2 == 0 else (0, 1) + for x in range(3) + for y in range(3) + } + state = SymPEPS.random( + 3, + 3, + symmetry="U1U1", + phys_dim=default_physical_sectors(model="fermi_hubbard_u1u1"), + fermionic=True, + site_charge=site_charge_from_occupations(charges), + bond_dim=1, + seed=79, + dtype="complex128", + to_backend=to_backend, + ).tn + target = state.copy() + target.mangle_inner_() + opt = SweepOptimizer( + state, + target, + chi=4, + boundary_engine="quimb-mps", + renormalize_state=False, + ) + + norm_tn, overlap_tn = opt._prepare_current_double_layers() + opt.bdy.start_sweep(norm_tn, "y", "left") + opt.bdy_overlap.start_sweep(overlap_tn, "y", "left") + # Seed the moving ymin boundary exactly as the forward half-sweep does + # after processing row zero, then optimize the interior row. + opt.bdy.advance_sweep(norm_tn, 0, axis="y", update_side="left") + opt.bdy_overlap.advance_sweep(overlap_tn, 0, axis="y", update_side="left") + run = opt._optimize_axis_slice_with_current_env( + 1, + axis="y", + solver="nlopt", + solver_options={"algorithm": "LD_LBFGS", "n_steps": 1, "maxeval": 1}, + ) + + assert np.isfinite(run["loss_initial"]) + assert np.isfinite(run["loss_final"]) + assert all( + isinstance(block, torch.Tensor) + for tensor in opt.state.tensor_map.values() + for block in tensor.data.blocks.values() + ) + + +def test_quimb_boundary_return_move_matches_fresh_ymax_environment(): + """Backward cached moves must retain Quimb's ymax-side compressed row.""" + network = qtn.PEPS.rand(3, 3, bond_dim=2, phys_dim=2, seed=83) + double_layer = network.H & network + store = QuimbMpsBoundaryStore(chi=4, layer_tags=None) + store.start_sweep(double_layer, "y", "right") + # Seed row two, then extend the ymax boundary across row one. + store.advance_sweep(double_layer, 2, axis="y", update_side="right") + store.advance_sweep(double_layer, 1, axis="y", update_side="right") + cached = store.mps_b["Y1_r"] + fresh = double_layer.compute_ymax_environments( + max_bond=4, + cutoff=store.cutoff, + canonize=store.canonize, + mode=store.mode, + dense=store.dense, + equalize_norms=store.equalize_norms, + envs={}, + )["ymax", 0] + + assert set(cached.outer_inds()) == set(fresh.outer_inds()) + assert set(cached.ind_map) == set(fresh.ind_map) + + +def test_nested_param_tree_flatten_roundtrip(): + tree = { + "site": { + "blocks": {(0, 0): np.arange(4.0), (1, 1): np.ones(2)}, + "meta": [np.zeros(1), (np.ones(1),)], + } + } + assert SweepOptimizer._needs_nested_param_flatten(tree) is True + flat, spec = SweepOptimizer._flatten_param_tree(tree) + # All leaves are arrays (no nested containers left). + assert all(isinstance(v, np.ndarray) for v in flat.values()) + rebuilt = SweepOptimizer._unflatten_param_tree(flat, spec) + np.testing.assert_allclose(rebuilt["site"]["blocks"][(0, 0)], np.arange(4.0)) + np.testing.assert_allclose(rebuilt["site"]["meta"][1][0], np.ones(1)) + assert isinstance(rebuilt["site"]["meta"], list) + assert isinstance(rebuilt["site"]["meta"][1], tuple) + + +def test_needs_nested_param_flatten_false_for_flat_mapping(): + flat = {"a": np.zeros(2), "b": np.ones(3)} + assert SweepOptimizer._needs_nested_param_flatten(flat) is False diff --git a/tests/test_sym_dmrg.py b/tests/test_sym_dmrg.py index 32dc0f7..7d67fe2 100644 --- a/tests/test_sym_dmrg.py +++ b/tests/test_sym_dmrg.py @@ -372,6 +372,39 @@ def flat_sweep(direction, canonize=True, **kwargs): assert opt.convergence_diagnostics[-1]["bond_schedule_ready"] is True +def test_symdmrg2_explicit_held_bond_schedule_can_converge(monkeypatch): + """An explicit final-bond hold should not need an extra hidden sweep.""" + pytest.importorskip("symmray") + state, mpo = _fh_u1u1_chain( + 4, + [(1, 0), (0, 1), (1, 0), (0, 1)], + bond_dim=1, + ) + opt = pepsy.SymDMRG2( + mpo, + state, + chi=2, + bond_dims=[1, 2, 2], + cutoff=1e-10, + compute_initial_energy=False, + norm_check="off", + ) + seen_bonds = [] + + def flat_sweep(direction, canonize=True, **kwargs): + del direction, canonize + seen_bonds.append(kwargs["max_bond"]) + return 0.0 + + monkeypatch.setattr(opt, "sweep", flat_sweep) + opt.solve(max_sweeps=3, sweep_sequence="R", tol=1e-10) + + assert seen_bonds == [1, 2, 2] + assert opt.converged is True + assert opt.convergence_diagnostics[-1]["bond_schedule_ready"] is True + assert opt.convergence_diagnostics[-1]["previous_bond_is_final"] is True + + def test_symdmrg2_min_sweeps_blocks_early_energy_convergence(monkeypatch): """Sweep convergence should require a clean comparison window.""" pytest.importorskip("symmray") @@ -965,6 +998,8 @@ def test_symdmrg2_profile_records_phase_timings(): assert all(event["elapsed"] >= 0.0 for event in opt.profile_diagnostics) assert summary["enabled"] assert summary["num_events"] == len(opt.profile_diagnostics) + assert summary["wall_elapsed"] == pytest.approx(summary["phase_totals"]["solve"]) + assert summary["phase_wall_fractions"]["solve"] == pytest.approx(1.0) assert summary["num_matvecs"] >= 1 assert summary["phase_counts"]["matvec"] >= 1 assert opt.summary()["num_profile_diagnostics"] == len(opt.profile_diagnostics) @@ -1507,18 +1542,19 @@ def wrapped(site, input_map): native_a = opt.two_site_matvec_symmray(site, trial_a) dense_b = opt.two_site_matvec_dense_reference(site, trial_b) native_b = opt.two_site_matvec_symmray(site, trial_b) + native_c = opt.two_site_matvec_symmray(site, trial_a) assert calls == [site, site + 1] assert opt.projected_problem_cache_misses == 1 - assert opt.projected_problem_cache_hits == 1 - assert opt.summary()["projected_problem_cache_hits"] == 1 + assert opt.projected_problem_cache_hits == 2 + assert opt.summary()["projected_problem_cache_hits"] == 2 assert opt.profile_summary()["projected_problem_cache_misses"] == 1 problem = opt._last_matvec_projected_problem problem_summary = problem.summary() assert problem_summary["right_contract_compiled_block_plan_builds"] == 1 assert problem_summary["left_contract_compiled_block_plan_builds"] == 1 - assert problem_summary["right_contract_compiled_block_plan_uses"] == 1 - assert problem_summary["left_contract_compiled_block_plan_uses"] == 1 + assert problem_summary["right_contract_compiled_block_plan_uses"] == 2 + assert problem_summary["left_contract_compiled_block_plan_uses"] == 2 assert problem_summary["right_contract_compiled_block_plan_terms"] > 0 assert problem_summary["left_contract_compiled_block_plan_terms"] > 0 assert problem_summary["right_contract_compiled_block_plan_mode"] == ( @@ -1527,9 +1563,22 @@ def wrapped(site, input_map): assert problem_summary["left_contract_compiled_block_plan_mode"] == ( "output_block_matmul" ) + assert problem_summary["right_contract_compiled_block_plan_layout_frozen"] + assert problem_summary["left_contract_compiled_block_plan_layout_frozen"] + assert problem_summary["right_contract_compiled_block_plan_layout_checks"] == 1 + assert problem_summary["left_contract_compiled_block_plan_layout_checks"] == 1 + assert ( + problem_summary["right_contract_compiled_block_plan_layout_fastpath_uses"] + >= 1 + ) + assert ( + problem_summary["left_contract_compiled_block_plan_layout_fastpath_uses"] + >= 1 + ) for sector in theta.data.blocks: assert native_a.data.blocks[sector] == pytest.approx(dense_a.data.blocks[sector]) assert native_b.data.blocks[sector] == pytest.approx(dense_b.data.blocks[sector]) + assert native_c.data.blocks[sector] == pytest.approx(native_a.data.blocks[sector]) def test_symdmrg2_matvec_diagnostics_record_cache_and_projector_stats(): @@ -2563,6 +2612,48 @@ def test_symdmrg2_prunes_projected_dead_variational_blocks_on_mapped_2d_case(): assert set(pruned.data.blocks) < set(theta.data.blocks) +def test_symdmrg2_pruning_skips_norm_scan_when_all_sectors_are_live(monkeypatch): + """Fully supported theta layouts should not scan every block norm.""" + pytest.importorskip("symmray") + state, mpo = _fh_u1u1_chain( + 4, + [(1, 0), (0, 1), (1, 0), (0, 1)], + bond_dim=2, + seed=37, + ) + opt = pepsy.SymDMRG2( + mpo, + state, + chi=4, + cutoff=1e-10, + local_solver="lanczos", + dense_threshold=0, + norm_check="off", + compute_initial_energy=False, + ) + opt._canonize_for_sweep("right") + opt.build_block_environments() + theta = opt.two_site_theta(1) + problem, _ = opt._get_projected_problem(1, theta) + all_sectors = set(theta.data.blocks) + monkeypatch.setattr( + problem, + "structural_output_sectors", + lambda: set(all_sectors), + ) + + def fail_block_norm(_block): + raise AssertionError("fully live layouts should skip block norm scans") + + monkeypatch.setattr(opt, "_block_norm", fail_block_norm) + pruned, diagnostic = opt._prune_theta_to_projected_support(1, theta) + + assert pruned is theta + assert diagnostic["removed_blocks"] == 0 + assert diagnostic["nonzero_input_blocks"] is None + assert diagnostic["nonzero_input_scan_skipped"] is True + + def test_symdmrg2_sector_enrichment_reaches_ed_from_narrow_initial_support(): """Template sector enrichment lets a narrow initial MPS reach FH ED.""" pytest.importorskip("symmray") diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index d7b9061..448aa56 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -8,6 +8,7 @@ from pepsy.tensors import symmetric as symmetric_mod from pepsy.tensors import ( Fermion, + FermionLatticeSetup, OneDMap, SpinfulFermion, SpinfulFermionHubbard, @@ -184,13 +185,15 @@ def test_ps_to_mps_accepts_fermion_and_builds_native_charge_sector(): (1, 0), (0, 1), ] + with pytest.raises(TypeError, match="unexpected keyword argument 'chi'"): + pepsy.ps_to_mps(4, fermion=fermion, chi=2) -def test_ps_to_mps_fermion_custom_occupations_preserve_global_sector(): - """Custom Fermion occupations should select the requested total charge.""" +def test_hrs_to_mps_fermion_custom_occupations_preserve_global_sector(): + """Random symmetric growth should preserve the requested total charge.""" fermion = Fermion(spinful=True, symmetry="U1U1") occupations = ((0, 1), (1, 0), (0, 1), (1, 0)) - mps = pepsy.ps_to_mps( + mps = pepsy.hrs_to_mps( 4, fermion=fermion, occupations=occupations, @@ -204,6 +207,30 @@ def test_ps_to_mps_fermion_custom_occupations_preserve_global_sector(): assert fermion.total_charge(occupations) == (2, 2) +def test_hrs_to_mps_direct_uses_symmray_random_blocks(): + """The direct method should build normalized native random blocks.""" + pytest.importorskip("symmray") + fermion = Fermion(spinful=True, symmetry="U1U1") + occupations = ((1, 0), (0, 1), (1, 0), (0, 1)) + mps = pepsy.hrs_to_mps( + 4, + fermion=fermion, + occupations=occupations, + chi=2, + method="direct", + subsizes="maximal", + seed=37, + dtype="complex128", + ) + + assert mps.max_bond() <= 2 + assert np.real(mps.norm()) == pytest.approx(1.0) + assert all( + type(mps[site].data).__name__ == "U1U1FermionicArray" + for site in range(4) + ) + + 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) @@ -217,7 +244,6 @@ def test_ps_to_peps_accepts_fermion_coordinate_occupations(): (2, 3), fermion=fermion, occupations=occupations, - chi=1, seed=19, dtype="complex128", ) @@ -232,6 +258,114 @@ def test_ps_to_peps_accepts_fermion_coordinate_occupations(): peps[x, y].data.charge for x, y in ((0, 0), (0, 1), (1, 0)) ] == [(1, 0), (0, 1), (0, 1)] + with pytest.raises(TypeError, match="unexpected keyword argument 'chi'"): + pepsy.ps_to_peps((2, 3), fermion=fermion, chi=2) + + +def test_ps_to_peps_accepts_periodic_fermion_state(): + """The native fermionic PEPS constructor should preserve cyclic bonds.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + peps = pepsy.ps_to_peps( + (3, 3), + fermion=fermion, + cyclic=True, + seed=23, + dtype="complex128", + ) + + assert (peps.Lx, peps.Ly) == (3, 3) + assert peps.max_bond() == 1 + assert set(peps[0, 0].inds) & set(peps[0, 2].inds) + assert set(peps[0, 0].inds) & set(peps[2, 0].inds) + assert all( + type(peps[x, y].data).__name__ == "U1U1FermionicArray" + for x in range(3) + for y in range(3) + ) + + +def test_hrs_to_mps_accepts_fermion_sector_and_grows_entanglement(): + """The Haar-random constructor should use chi for symmetric growth.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + occupations = ((1, 0), (0, 1), (1, 0), (0, 1)) + mps = pepsy.hrs_to_mps( + 4, + fermion=fermion, + occupations=occupations, + chi=2, + random_rounds=20, + seed=31, + dtype="complex128", + ) + + assert mps.L == 4 + assert 1 < mps.max_bond() <= 2 + assert all( + type(mps[site].data).__name__ == "U1U1FermionicArray" + for site in range(4) + ) + assert fermion.total_charge(occupations) == (2, 2) + assert pepsy.hrps_to_mps is pepsy.hrs_to_mps + + +def test_hrs_to_peps_accepts_fermion_sector_and_chi(): + """The Haar-random PEPS constructor should build native symmetric PEPS.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + occupations = { + (x, y): (1, 0) if (x + y) % 2 == 0 else (0, 1) + for x in range(2) + for y in range(2) + } + peps = pepsy.hrs_to_peps( + (2, 2), + fermion=fermion, + occupations=occupations, + chi=2, + method="direct", + subsizes="maximal", + seed=32, + dtype="complex128", + ) + + assert (peps.Lx, peps.Ly) == (2, 2) + assert 1 < peps.max_bond() <= 2 + assert np.real((peps.H & peps).contract(all)) == pytest.approx(1.0) + assert all( + type(peps[x, y].data).__name__ == "U1U1FermionicArray" + for x, y in occupations + ) + assert pepsy.hrps_to_peps is pepsy.hrs_to_peps + + +@pytest.mark.parametrize( + ("symmetry", "expected_charge", "expected_occupation"), + [ + ("U1U1", (3, 3), (1, 0)), + ("U1", 6, 1), + ("Z2", 0, 1), + ], +) +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) + + setup = fermion.lattice_half_filling(3, 2, pattern="checkerboard") + + assert isinstance(setup, FermionLatticeSetup) + assert pepsy.FermionLatticeSetup is FermionLatticeSetup + assert (setup.Lx, setup.Ly) == (3, 2) + assert setup.sites == tuple( + (x, y) for x in range(3) for y in range(2) + ) + assert len(setup.edges) == 7 + assert setup.target_charge == expected_charge + assert setup.target_particles == 6 + assert setup.occupations[(0, 0)] == expected_occupation + assert setup.site_charge((0, 0)) == expected_occupation + assert setup.spin_occupations[(0, 0)] == (1, 0) + assert setup.spin_occupations[(0, 1)] == (0, 1) def test_symdmrg_fermionic_state_accepts_raw_fermion_constructor_output(): @@ -261,6 +395,9 @@ 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) + assert fermion.model == ( + "fermi_hubbard" if symmetry == "Z2" else "fermi_hubbard_u1u1" + ) assert fermion.physical_sectors == default_physical_sectors(symmetry, 4) assert fermion.operator_charge("create_up") == ( 1 if symmetry == "Z2" else (0, 1) @@ -343,6 +480,37 @@ def test_fermion_exposes_explicit_native_operator_terms(): 0.6 * density.to_dense(), reference_density.to_dense() ) + forward_up = fermion.operator_term( + [(1.0, ((0, "create_up"), (1, "annihilate_up")))], + sites=(0, 1), + add_hc=True, + ) + np.testing.assert_allclose( + forward_up.to_dense(), fermion.hopping_operator(spin="up").to_dense() + ) + + eta_pair = fermion.eta_pair_operator() + explicit_eta_pair = fermion.operator_term( + [ + (1.0, ((0, "pair_create"), (1, "pair_annihilate"))), + (1.0, ((1, "pair_create"), (0, "pair_annihilate"))), + ], + sites=(0, 1), + charge=(0, 0), + ) + assert eta_pair.charge == (0, 0) + np.testing.assert_allclose(eta_pair.to_dense(), explicit_eta_pair.to_dense()) + + with pytest.raises(ValueError, match="self-conjugate operator charge"): + fermion.operator_term( + [(1.0, ((0, "create_up"),))], + sites=(0,), + add_hc=True, + ) + + with pytest.raises(ValueError, match="require spinful fermions"): + Fermion(spinful=False, symmetry="U1").eta_pair_operator() + explicit_terms = { (0, 1): -1.7 * hopping, 0: 3.0 * fermion.interaction_operator(), @@ -361,6 +529,106 @@ def test_fermion_exposes_explicit_native_operator_terms(): ) +def test_fermion_spin_operator_algebra_and_native_correlators(): + """Native spin helpers should match local algebra and stay symmetric.""" + fermion = Fermion(spinful=True, symmetry="U1") + sx = fermion.dense_operator("sx") + sy = fermion.dense_operator("sy") + sz = fermion.dense_operator("sz") + + np.testing.assert_allclose(sx @ sy - sy @ sx, 1j * sz) + np.testing.assert_allclose( + (sx @ sx + sy @ sy + sz @ sz)[1:3, 1:3], + 0.75 * np.eye(2), + ) + np.testing.assert_allclose( + fermion.observable("sx").to_dense(), + sx, + ) + np.testing.assert_allclose( + fermion.observable("sy").to_dense(), + sy, + ) + assert fermion.spin_x_operator() is fermion.observable("sx") + assert fermion.spin_y_operator() is fermion.observable("sy") + assert fermion.spin_z_operator() is fermion.observable("sz") + + for operator in ( + fermion.spin_z_correlator(), + fermion.spin_x_correlator(), + fermion.spin_y_correlator(), + fermion.xy_exchange_operator(), + fermion.heisenberg_operator(), + ): + assert type(operator).__name__ == "U1FermionicArray" + assert operator.charge == 0 + + +@pytest.mark.parametrize("symmetry", ["U1", "Z2"]) +def test_fermion_spin_gates_are_native_unitary_and_parameterized(symmetry): + """Spin gates should exponentiate natively for total-charge symmetries.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + + def as_matrix(operator): + dense = np.asarray(operator.to_dense()) + dim = int(np.sqrt(dense.size)) + return dense.reshape(dim, dim) + + gates = ( + fermion.sx_gate({3: 0.13}, site=3), + fermion.sy_gate(lambda site: 0.17 + 0.01 * site, site=3), + fermion.sz_gate(0.11), + fermion.sxx_gate(0.07), + fermion.syy_gate({(3, 4): 0.09}, edge=(3, 4)), + fermion.szz_gate(0.05), + fermion.xy_gate(0.03), + fermion.heisenberg_gate(0.02), + ) + for gate_ in gates: + matrix = as_matrix(gate_) + np.testing.assert_allclose( + matrix.conj().T @ matrix, + np.eye(matrix.shape[0]), + atol=1e-12, + ) + + imaginary = as_matrix(fermion.syy_gate(0.2, imaginary=True)) + np.testing.assert_allclose(imaginary, imaginary.conj().T, atol=1e-12) + assert np.all(np.abs(np.linalg.eigvalsh(imaginary)) > 0.0) + + +def test_fermion_spin_flip_restrictions_are_explicit_under_u1u1(): + """Separate spin charges must reject inhomogeneous spin-flip gates.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + + with pytest.raises(ValueError, match="symmetry='U1' or 'Z2'"): + fermion.observable("sx") + with pytest.raises(ValueError, match="symmetry='U1' or 'Z2'"): + fermion.observable("sy") + with pytest.raises(ValueError, match="symmetry='U1' or 'Z2'"): + fermion.syy_gate(0.1) + + assert fermion.observable("sz").charge == (0, 0) + assert fermion.szz_gate(0.1).charge == (0, 0) + assert fermion.xy_gate(0.1).charge == (0, 0) + assert fermion.heisenberg_gate(0.1).charge == (0, 0) + assert fermion.observable("s_plus").charge != (0, 0) + with pytest.raises(ValueError, match="charge-neutral"): + fermion.operator_gate("s_plus", 0.1) + + +def test_fermion_spin_gates_preserve_torch_backend(): + """The generic operator gate should use the configured block backend.""" + torch = pytest.importorskip("torch") + fermion = Fermion( + spinful=True, + symmetry="U1", + to_backend=pepsy.backend_torch(dtype=torch.complex128), + ) + gate_ = fermion.sy_gate(0.1) + assert gate_.backend == "torch" + + 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) @@ -584,7 +852,6 @@ def test_native_hopping_gate_has_correct_long_range_parity_sign( len(initial), fermion=fermion, occupations=initial, - chi=1, seed=123, dtype="complex128", ) @@ -1501,6 +1768,37 @@ def test_symmps_heisenberg_builds_energy_and_imaginary_step(): assert state.norm() == pytest.approx(1.0) +@pytest.mark.parametrize("model", ["heisenberg", "fermi_hubbard"]) +def test_symmps_mps_optimizer_simple_update_preserves_symmray_data(model): + """Simple update should preserve Symmray tensor data under default settings.""" + state = SymMPS.for_model( + model, + 3, + bond_dim=2, + seed=18, + dtype="complex128", + ) + if model == "fermi_hubbard": + hamiltonian = state.build_hamiltonian(t=1.0, U=2.0, mu=0.1) + else: + hamiltonian = state.build_hamiltonian() + gates = hamiltonian.gate_stream(0.001, imaginary=True) + + optimizer = pepsy.MpsOptimizer( + state.tn.copy(), + gates, + chi=4, + mode="su", + ) + out = optimizer.run(progbar=False, cutoff=1.0e-10) + + assert _all_tensor_data_symmray(out) + assert out.max_bond() <= 4 + assert len(optimizer.gauges) == out.L - 1 + assert optimizer.p_ungauged is not None + assert np.isfinite(np.real(optimizer.p_ungauged.norm())) + + def test_symmps_measures_dense_generic_observables(): """SymMPS.measure should convert dense local operators to Symmray arrays.""" state = SymMPS.for_model( @@ -2032,6 +2330,44 @@ def test_spinless_fermi_hubbard_u1_hamiltonian_mpo_compresses_long_range(): assert complex(energy_compressed) == pytest.approx(complex(energy)) +def test_symhamiltonian_mpo_compression_report_warns_for_soft_bond_cap(): + """A tied Symmray cutoff must not silently exceed the requested cap.""" + lx, ly = 4, 3 + length = lx * ly + edges = [] + for x in range(lx): + for y in range(ly): + site = x * ly + y + edges.append((site, ((x + 1) % lx) * ly + y)) + edges.append((site, x * ly + ((y + 1) % ly))) + ham = SymHamiltonian.from_edges( + "fermi_hubbard_u1u1", + "U1U1", + edges, + t=1.0, + U=8.0, + ) + + with pytest.warns(RuntimeWarning, match="requested max_bond=16"): + mpo = ham.to_mpo( + L=length, + compress=True, + max_bond=16, + cutoff=1e-12, + ) + + report = mpo.pepsy_compression_report + assert report["compressed"] is True + assert report["cutoff"] == 1e-12 + assert report["requested_max_bond"] == 16 + assert report["raw_max_bond"] > report["final_max_bond"] + assert report["final_max_bond"] == mpo.max_bond() + assert report["rank_reduced"] is True + assert report["cap_bound"] is True + assert report["max_bond_exceeded"] is True + assert report["final_max_bond"] > report["requested_max_bond"] + + def test_spinless_fermi_hubbard_u1_hamiltonian_mpo_maps_2d_long_range_edge(): """Spinless FH coordinate edges should match their mapped flat edge.""" mapper = OneDMap(2, 2, mode="snake") @@ -2479,6 +2815,68 @@ def test_symmps_mps_optimizer_handles_spinful_fermi_hubbard_dims(): assert np.real(raw_norm) > 0.0 +def test_symmray_mpo_real_time_fermion_stream_preserves_norm_without_truncation(): + """All Symmray MPO gates must use the graded auto-swap application path.""" + torch = pytest.importorskip("torch") + import quimb.tensor as qtn + + backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) + fermion = Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=8.0, + to_backend=backend, + ) + lx, ly = 2, 3 + occupations = { + (x, y): (1, 0) if (x + y) % 2 == 0 else (0, 1) + for x in range(lx) + for y in range(ly) + } + mapper = pepsy.OneDMap(lx, ly, mode="snake") + idx2coo, coo2idx = mapper.build() + state = pepsy.ps_to_mps( + lx * ly, + fermion=fermion, + occupations=tuple(occupations[idx2coo[index]] for index in range(lx * ly)), + seed=101, + dtype="complex128", + to_backend=backend, + ) + state.normalize() + + half_dt = 0.005 / 2.0 + interaction = fermion.interaction_gate(half_dt, U=8.0, imaginary=False) + hopping = fermion.hopping_gate(half_dt, t=1.0, imaginary=False) + edges = tuple(qtn.edges_2d_square(lx, ly, cyclic=False)) + layers = fermion.edge_coloring_layers(edges) + coordinate_stream = ( + [(interaction, site) for site in occupations] + + [(hopping, edge) for layer in layers for edge in layer] + + [(hopping, edge) for layer in reversed(layers) for edge in reversed(layer)] + + [(interaction, site) for site in occupations] + ) + stream = [ + ( + gate, + tuple(coo2idx[site] for site in where) + if isinstance(where, tuple) and len(where) == 2 and isinstance(where[0], tuple) + else coo2idx[where], + ) + for gate, where in coordinate_stream + ] + + out = pepsy.MpsOptimizer(state, stream * 2, chi=64, mode="mpo", inplace=True).run( + progbar=False, + cutoff=0.0, + non_unitary=False, + ) + + assert all(getattr(tensor.data, "backend", None) == "torch" for tensor in out.tensors) + assert float(np.real(pepsy.to_float(out.norm()))) == pytest.approx(1.0, abs=1.0e-10) + + def test_symmps_mps_optimizer_handles_cuda_torch_blocks(): """Symmray canonicalization should not coerce CUDA torch blocks via NumPy.""" torch = pytest.importorskip("torch") @@ -3008,6 +3406,103 @@ def test_sympeps_gate_wrappers_route_nonlocal_symmray_swaps(case_name, method): assert _all_tensor_data_symmray(out) +@pytest.mark.parametrize("method", ["gate", "simple"]) +def test_sympeps_gate_wrappers_route_nonlocal_u1u1_swaps(method): + """Product-symmetry PEPS gates should route through neutral SWAP sectors.""" + charges = {(i, j): (1, 1) for i in range(3) for j in range(3)} + state = SymPEPS.random( + 3, + 3, + symmetry="U1U1", + phys_dim={(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}, + fermionic=True, + site_charge=site_charge_from_occupations(charges), + bond_dim=2, + seed=54, + dtype="complex128", + ) + hamiltonian = SymHamiltonian.from_edges( + "fermi_hubbard_u1u1", + "U1U1", + [((0, 0), (2, 2))], + t=1.0, + U=2.0, + mu=0.1, + ) + gates = hamiltonian.gate_stream(0.0005, imaginary=True) + + if method == "gate": + out = gate( + state.tn.copy(), + gates, + max_bond=4, + cutoff=1.0e-10, + inplace=False, + ) + else: + out = gate_simple( + state.tn.copy(), + gates, + gauges={}, + max_bond=4, + cutoff=1.0e-10, + inplace=False, + ) + + assert out.max_bond() <= 4 + assert _all_tensor_data_symmray(out) + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1", "Z2"]) +def test_sympeps_gate_simple_explicit_reduce_split_fermion_symmetries(symmetry): + """Explicit reduce-split should work for native fermionic PEPS sectors.""" + if symmetry == "U1U1": + phys_dim = {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1} + charges = {(i, j): (1, 1) for i in range(3) for j in range(3)} + model = "fermi_hubbard_u1u1" + elif symmetry == "Z2": + phys_dim = {0: 2, 1: 2} + charges = {(i, j): 1 for i in range(3) for j in range(3)} + model = "fermi_hubbard" + else: + phys_dim = {0: 1, 1: 2, 2: 1} + charges = {(i, j): 1 for i in range(3) for j in range(3)} + model = "fermi_hubbard" + + state = SymPEPS.random( + 3, + 3, + symmetry=symmetry, + phys_dim=phys_dim, + fermionic=True, + site_charge=site_charge_from_occupations(charges), + bond_dim=2, + seed=55, + dtype="complex128", + ) + hamiltonian = SymHamiltonian.from_edges( + model, + symmetry, + [((0, 0), (2, 2))], + t=1.0, + U=2.0, + mu=0.1, + ) + + out = gate_simple( + state.tn.copy(), + hamiltonian.gate_stream(0.0005, imaginary=True), + gauges={}, + max_bond=4, + cutoff=1.0e-10, + contract="reduce-split", + inplace=False, + ) + + assert out.max_bond() <= 4 + assert _all_tensor_data_symmray(out) + + def test_sympeps_gate_method_preserves_pepsy_gate_contract_default(monkeypatch): """SymPEPS method='gate' should not override pepsy.gate's default.""" state = SymPEPS.for_model( diff --git a/tests/test_trajectory_noise.py b/tests/test_trajectory_noise.py index 0cfa72e..24e3e94 100644 --- a/tests/test_trajectory_noise.py +++ b/tests/test_trajectory_noise.py @@ -302,6 +302,28 @@ def test_state_dependent_kraus_branches_are_sampled_from_the_current_state(kind) assert abs(optimizer.p.norm()) == pytest.approx(1.0, abs=1e-8) +def test_tree_state_dependent_kraus_branches_are_sampled_from_the_current_state(): + """Tree trajectories compute Kraus probabilities from copied TTNs.""" + stream = [ + (_X, 0), + pepsy.TrajectoryEvent(pepsy.TrajectoryChannel.amplitude_damping(0.5), 0), + ] + result = pepsy.run_trajectory_shots( + lambda: pepsy.TreeOptimizer(None, n=2, chi=4, run=False), + stream, + shots=12, + seed=6, + ) + + labels = {records[0].label for records in result.records} + assert labels == {"jump", "no_jump"} + assert all(records[0].probability == pytest.approx(0.5) for records in result.records) + for optimizer, records in zip(result.optimizers, result.records): + expected = [1.0, 0.0, 0.0, 0.0] if records[0].label == "jump" else [0.0, 0.0, 1.0, 0.0] + np.testing.assert_allclose(_statevector(optimizer), expected, atol=1e-8) + assert optimizer.norm() == pytest.approx(1.0, abs=1e-8) + + def test_kraus_trajectory_starts_a_fresh_stn_norm_diagnostic_segment(): stream = [ ("rxx", 0.8, 0, 1), @@ -391,6 +413,35 @@ def factory(): assert all(leaf.gate_stream[-1][0] == "measure" for leaf in result.leaves) +def test_coalesced_tree_trajectory_branches_mid_circuit_measurements_by_count(): + """Tree expectations support exact coalesced measurement branching.""" + hadamard = np.array( + [[1.0, 1.0], [1.0, -1.0]], dtype=complex + ) / np.sqrt(2.0) + calls = 0 + + def factory(): + nonlocal calls + calls += 1 + return pepsy.TreeOptimizer(None, n=2, chi=4, run=False) + + result = pepsy.run_coalesced_trajectory_shots( + factory, + [(hadamard, 0), ("measure", "Z", 0)], + shots=64, + seed=9, + ) + + assert calls == 1 + assert result.shots == 64 + assert result.branches == 2 + assert {leaf.measurements[0].outcome for leaf in result.leaves} == {-1, 1} + assert all( + leaf.measurements[0].probability == pytest.approx(0.5) + for leaf in result.leaves + ) + + def test_coalesced_kraus_ensemble_uses_one_copy_per_nonempty_outcome(): """State-dependent channel probabilities are evaluated once per live node.""" result = pepsy.run_coalesced_trajectory_shots( diff --git a/tests/test_tree_symmetric.py b/tests/test_tree_symmetric.py new file mode 100644 index 0000000..595fcbe --- /dev/null +++ b/tests/test_tree_symmetric.py @@ -0,0 +1,372 @@ +"""Symmray-backed tree-state construction tests.""" + +import numpy as np +import pytest + +import pepsy +import quimb.tensor as qtn +from pepsy.optimizers.tree import TreeOptimizer, TreePlan, TreeTensorNetwork +from pepsy.tensors import Fermion + + +pytest.importorskip("symmray") + + +def _is_symmray_data(tensor): + return hasattr(tensor.data, "blocks") and hasattr(tensor.data, "indices") + + +def _nonbinary_u1u1_case(*, chi=64, track_truncation=False): + """Return a small evolved native-fermion state with a branching path.""" + fermion = Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=8.0, + mu=0.0, + dtype="complex128", + ) + plan = TreePlan.from_children( + { + 0: (), + 1: (), + 2: (), + 3: (1, 2), + 4: (), + 5: (), + 6: (), + 7: (5, 6), + 8: (0, 3, 4, 7), + }, + {2: 0, 0: 1, 1: 2, 6: 3, 4: 4, 5: 5}, + root=8, + ) + occupations = fermion.half_filled_occupations(6) + state = pepsy.ps_to_ttn( + 6, + tree=plan, + fermion=fermion, + occupations=occupations, + seed=23, + ) + half = 0.025 + onsite = [ + ( + fermion.onsite_gate( + half, site=site, U=8.0, mu=0.0, imaginary=False + ), + site, + ) + for site in range(6) + ] + hopping = fermion.hopping_gate(half, t=1.0, imaginary=False) + # The final edge crosses two internal branching nodes and exposed both the + # graded-centre norm bug and the Symmray spectrum-conversion bug. + gates = onsite + [ + (hopping, (0, 1)), + (hopping, (2, 3)), + (hopping, (4, 5)), + (hopping, (0, 2)), + ] + engine = TreeOptimizer( + gates, + n=6, + tree=plan, + state=state, + chi=chi, + cutoff=0.0, + mode="direct", + track_truncation=track_truncation, + run=False, + ) + engine.run() + return fermion, engine + + +def _full_norm(state): + return float(abs((state.H & state).contract(all, optimize="auto")) ** 0.5) + + +def test_fermionic_product_ttn_assigns_leaf_charges_and_virtual_sectors(): + """A product TTN has physical fermion sectors only at its leaves.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + occupations = ((1, 0), (0, 1), (1, 0), (0, 1)) + plan = TreePlan.from_order(range(4), structure="balanced") + + ttn = pepsy.ps_to_ttn( + 4, + tree=plan, + fermion=fermion, + occupations=occupations, + chi=1, + seed=11, + ) + + assert ttn.plan is plan + assert ttn.symmetry == "U1U1" + assert ttn.fermionic + assert ttn.max_bond() == 1 + assert all(_is_symmray_data(tensor) for tensor in ttn.tensors) + assert [ + ttn.node_tensor(ttn.leaf_of_qubit(q)).data.charge + for q in range(4) + ] == list(occupations) + assert [ttn.ind_size(ttn.site_ind(q)) for q in range(4)] == [4, 4, 4, 4] + + for parent, children in plan.children.items(): + for child in children: + parent_tensor = ttn.node_tensor(parent) + child_tensor = ttn.node_tensor(child) + bond = next(iter(qtn.bonds(parent_tensor, child_tensor))) + assert ttn.ind_size(bond) == 1 + parent_axis = parent_tensor.inds.index(bond) + child_axis = child_tensor.inds.index(bond) + assert ( + parent_tensor.data.indices[parent_axis].dual + != child_tensor.data.indices[child_axis].dual + ) + + optimizer = TreeOptimizer(None, state=ttn, run=False) + assert optimizer.backend_info()["backend"] == "symmray" + assert _full_norm(ttn) == pytest.approx(1.0) + assert optimizer.norm() == pytest.approx(1.0) + + +@pytest.mark.parametrize("symmetry", ("U1", "Z2", "U1U1", "Z2Z2")) +def test_fermionic_product_mps_and_ttn_select_the_same_fock_state(symmetry): + """Degenerate physical sectors must not leave random product seed vectors.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + occupations = fermion.half_filled_occupations(4) + plan = TreePlan.from_order(range(4), structure="balanced") + + mps = pepsy.ps_to_mps( + 4, fermion=fermion, occupations=occupations, seed=17, + ) + ttn = pepsy.ps_to_ttn( + 4, tree=plan, fermion=fermion, occupations=occupations, seed=31, + ) + + assert float(pepsy.tn_fidelity(mps, ttn)) > 1 - 1e-12 + assert _full_norm(ttn) == pytest.approx(1.0) + + +def test_native_local_expectation_uses_selected_degenerate_fock_basis(): + """The TTN observable path works natively and sees explicit U1 spins.""" + fermion = Fermion(spinful=True, symmetry="U1") + occupations = ((1, 0), (0, 1), (1, 0), (0, 1)) + ttn = pepsy.ps_to_ttn( + 4, + tree=TreePlan.from_order(range(4), structure="balanced"), + fermion=fermion, + occupations=occupations, + ) + + number_up = fermion.observable("number_u") + assert ttn.local_expectation(number_up, (0,), max_bond=None) == pytest.approx(1) + assert ttn.local_expectation(number_up, (1,), max_bond=None) == pytest.approx(0) + with pytest.raises(TypeError, match="native Symmray observable"): + ttn.local_expectation(np.eye(4), (0,), max_bond=None) + + +def test_native_local_expectation_reuses_and_invalidates_norm_cache(): + """Repeated native readouts reuse the denominator until state mutation.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + state = pepsy.ps_to_ttn( + 4, + tree=TreePlan.from_order(range(4), structure="balanced"), + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + ) + number = fermion.observable("number") + other = fermion.observable("number_u") + + state.local_expectation(number, (0,)) + cached = state._fermionic_norm_cache + assert cached is not None + state.local_expectation(other, (1,)) + assert state._fermionic_norm_cache is cached + + state.invalidate_canonical_form() + assert state.orthogonality_center is None + state.local_expectation(number, (0,)) + assert state.orthogonality_center is None + + copied = state.copy() + assert copied._fermionic_norm_cache is None + state.gate_inds_(fermion.spin_z_gate(0.1), [state.site_ind(0)], contract=True) + assert state._fermionic_norm_cache is None + + +def test_dense_local_expectation_preserves_tracked_gauge(): + """Dense readout restores a known centre and preserves an unknown gauge.""" + plan = TreePlan.from_order(range(5), structure="balanced") + state = TreeTensorNetwork.rand(plan, D=3, seed=91) + z = np.diag([1.0, -1.0]).astype(complex) + original_center = state.orthogonality_center + + state.local_expectation(z, (4,)) + assert state.orthogonality_center == original_center + state.local_expectation(np.eye(4), (0, 4), normalized=False) + assert state.orthogonality_center == original_center + + state.invalidate_canonical_form() + assert state.orthogonality_center is None + before = state.to_statevector() + state.local_expectation(z, (4,)) + assert state.orthogonality_center is None + after = state.to_statevector() + fidelity = abs(np.vdot(before, after)) ** 2 / ( + np.vdot(before, before).real * np.vdot(after, after).real + ) + assert float(fidelity) > 1 - 1e-12 + + +def test_native_fermionic_norm_and_observable_use_exact_tree_environment(): + """Native readout remains correct across nonbinary graded centre moves.""" + fermion, engine = _nonbinary_u1u1_case() + state = engine.tn + reference = state.copy() + expected_norm = _full_norm(state) + assert engine._leaf_canonical_norm() == pytest.approx( + expected_norm, abs=1e-12 + ) + + for site in range(6): + state.shift_orthogonality_center(state.leaf_of_qubit(site)) + assert engine.norm() == pytest.approx(expected_norm, abs=1e-12) + assert state._fermionic_center_norm_squared() == pytest.approx( + expected_norm * expected_norm, abs=1e-12 + ) + assert state.is_canonical_form(state.orthogonality_center) + + operator = fermion.interaction_term(site) + operated = qtn.tensor_network_gate_inds( + state, + operator, + [state.site_ind(site)], + contract=False, + tags=[], + inplace=False, + ) + expected = (state.H | operated).contract(all, optimize="auto") + expected /= (state.H | state).contract(all, optimize="auto") + assert state.local_expectation( + operator, site, optimize="auto" + ) == pytest.approx(expected, abs=1e-12) + + assert all(_is_symmray_data(tensor) for tensor in state.tensors) + assert float(pepsy.tn_fidelity(reference, state)) > 1 - 1e-12 + + +def test_native_fermionic_truncation_spectrum_is_blockwise(): + """Tracked native compression handles all Symmray charge blocks.""" + _, engine = _nonbinary_u1u1_case(chi=2, track_truncation=True) + report = engine.truncation_report() + + assert report["n_tracked"] > 0 + assert all( + event["spectrum_rank"] is not None + for event in report["events"] + ) + assert all( + 0.0 <= event["discarded_fraction"] <= 1.0 + for event in report["events"] + ) + assert engine.norm() == pytest.approx(_full_norm(engine.tn), abs=1e-12) + + +def test_native_multisite_submpo_uses_qr_subtree_routing(): + """A native three-site MPO stays Symmray-native through tree routing.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + plan = TreePlan.from_order(range(4), structure="balanced") + state = pepsy.ps_to_ttn(4, tree=plan, fermion=fermion) + identity = fermion.operator_term([(1.0, ())], sites=(0, 1, 2)) + mpo = qtn.MatrixProductOperator.from_dense( + identity, + dims=(4, 4, 4), + sites=(0, 1, 3), + L=4, + max_bond=None, + cutoff=0.0, + ) + engine = TreeOptimizer(None, n=4, state=state.copy(), chi=32, run=False) + + engine.apply_submpo(mpo, (0, 1, 3)) + + assert all(_is_symmray_data(tensor) for tensor in engine.tn.tensors) + assert float(pepsy.tn_fidelity(state, engine.p)) > 1 - 1e-12 + + +def test_native_two_site_gate_eager_preflight_and_subtree_operator(): + """Native rank-four gates survive preflight and the public subtree API.""" + fermion = Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=8.0, + mu=0.0, + dtype="complex128", + ) + plan = TreePlan.from_order(range(2)) + state = pepsy.ps_to_ttn( + 2, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1)), + dtype="complex128", + ) + gate = fermion.hopping_gate(0.05, t=1.0, imaginary=False) + + eager = TreeOptimizer( + [(gate, (0, 1))], n=2, state=state, chi=8, cutoff=0.0 + ) + assert eager.norm() == pytest.approx(1.0, abs=1e-12) + + explicit = TreeOptimizer(None, n=2, state=state, chi=8, run=False) + explicit.apply_subtree_operator(gate, (0, 1)) + assert explicit.norm() == pytest.approx(1.0, abs=1e-12) + + multi_state = pepsy.ps_to_ttn( + 4, + tree=TreePlan.from_order(range(4), structure="balanced"), + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + ) + multisite = fermion.operator_term([(1.0, ())], sites=(0, 1, 2)) + multi = TreeOptimizer(None, n=4, state=multi_state, chi=32, run=False) + multi.apply_subtree_operator(multisite, (0, 1, 3)) + assert multi.norm() == pytest.approx(1.0, abs=1e-12) + + +def test_native_qubit_measurement_helpers_fail_with_actionable_error(): + """Qubit Pauli/measurement helpers do not misinterpret native sites.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + state = pepsy.ps_to_ttn( + 2, + tree=TreePlan.from_order(range(2)), + fermion=fermion, + occupations=((1, 0), (0, 1)), + ) + engine = TreeOptimizer(None, n=2, state=state, run=False) + + with pytest.raises(NotImplementedError, match="dense two-level qubit"): + engine.expectation_pauli("Z", 0) + with pytest.raises(NotImplementedError, match="dense two-level qubit"): + engine.measure(0, outcome=0) + + +def test_fermionic_random_ttn_uses_requested_symmetric_bond_dimension(): + """``hrs_to_ttn`` builds a native random symmetric tree at ``chi``.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + ttn = pepsy.hrs_to_ttn( + 4, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + chi=2, + seed=19, + ) + + assert ttn.fermionic + assert ttn.max_bond() == 2 + assert all(_is_symmray_data(tensor) for tensor in ttn.tensors) + assert pepsy.hrps_to_ttn is pepsy.hrs_to_ttn diff --git a/tests/test_vmc_api.py b/tests/test_vmc_api.py new file mode 100644 index 0000000..3ec4bae --- /dev/null +++ b/tests/test_vmc_api.py @@ -0,0 +1,629 @@ +"""Tests for the backend-neutral VMC API contracts.""" + +import numpy as np +import pytest + +from pepsy.vmc import ( + BackendCapabilityWarning, + ContractionConfig, + LocalMatrixTerm, + MCState, + OperatorFactor, + OperatorSum, + ProductTerm, + SamplingConfig, + OptimizationConfig, + VMCMeasurement, + VMCBackendCapabilityError, + VMCOptimizationResult, + VMC, + VMCProblem, + VMCSamples, + compile_operator_sum_netket, + compile_operator_sum_torch, +) + + +def test_operator_sum_normalizes_symbolic_and_matrix_terms(): + hopping = ProductTerm( + coefficient=-1.0, + factors=( + OperatorFactor(0, "fermion", spin="up", dagger=True), + OperatorFactor(1, "fermion", spin="up", dagger=False), + ), + ) + onsite = LocalMatrixTerm( + support=(0,), + matrix=np.eye(2), + coefficient=0.5, + ) + terms = OperatorSum.from_terms( + (term for term in (hopping, onsite)), + constant=1.25, + metadata={"statistics": "fermion"}, + ) + + assert len(terms) == 2 + assert terms.sites == (0, 1) + assert terms.metadata["statistics"] == "fermion" + assert hopping.support == (0, 1) + + with pytest.raises(TypeError, match="ProductTerm or LocalMatrixTerm"): + OperatorSum(terms=(object(),)) + + with pytest.raises(ValueError, match="expected rank 4"): + LocalMatrixTerm(support=(0, 1), matrix=np.eye(4)) + with pytest.raises(ValueError, match="dimensions must match"): + LocalMatrixTerm(support=(0,), matrix=np.zeros((2, 3))) + + +def test_vmc_problem_freezes_observables_and_site_order(): + problem = VMCProblem( + peps=object(), + hamiltonian=OperatorSum(), + observables={"density": OperatorSum()}, + symmetry="U1U1", + site_order=((0, 0), (0, 1)), + ) + + assert isinstance(problem.observables["density"], OperatorSum) + with pytest.raises(TypeError): + problem.observables["other"] = OperatorSum() + + +def test_mcstate_uses_netket_total_sample_convention_and_bridges_to_problem(): + state = MCState( + object(), + n_samples=12, + n_chains=3, + n_discard_per_chain=4, + symmetry="U1U1", + site_order=(0, 1), + ) + + assert state.n_samples == 12 + assert state.n_chains == 3 + assert state.n_discard_per_chain == 4 + assert state.sampling.n_samples_per_chain == 4 + assert state.ansatz is state.peps + assert state.to_problem(OperatorSum()).site_order == (0, 1) + + with pytest.raises(ValueError, match="divisible"): + MCState(object(), n_samples=5, n_chains=2) + with pytest.raises(ValueError, match="either sampling"): + MCState(object(), sampling=SamplingConfig(), n_samples=8) + + +def test_netket_shaped_vmc_driver_builds_from_mcstate(monkeypatch): + import pepsy.vmc.torch as torch_vmc + + seen = {} + native = object() + + class FakeSetup: + @property + def native(self): + return native + + def sample(self, sampling=None): + return VMCSamples( + configs=np.zeros((2, 2, 1), dtype=np.int64), + n_samples_per_chain=2, + n_chains=2, + ) + + def measure(self, observables=None, *, sampling=None, **kwargs): + seen["measure_kwargs"] = kwargs + values = {"energy": "native-energy"} + if observables: + values.update({name: value for name, value in observables.items()}) + return VMCMeasurement(energy_mean=-1.0, observables=values) + + def optimize(self, optimization=None, *, n_steps=None, **kwargs): + seen["optimization"] = optimization + seen["n_steps"] = n_steps + return VMCOptimizationResult( + steps=np.arange(1, (n_steps or optimization.n_steps) + 1), + energies=np.full(n_steps or optimization.n_steps, -1.0), + errors=np.zeros(n_steps or optimization.n_steps), + ) + + def fake_builder(problem, **kwargs): + seen["problem"] = problem + seen.update(kwargs) + return FakeSetup() + + monkeypatch.setattr(torch_vmc, "build_torch_vmc", fake_builder) + state = MCState(object(), n_samples=4, n_chains=2, symmetry="U1U1") + hamiltonian = OperatorSum() + driver = VMC( + hamiltonian, + state, + backend="torch", + fermion=object(), + observables={"density": OperatorSum()}, + ) + + assert driver.state is state + assert driver.native is native + assert seen["problem"].hamiltonian is hamiltonian + assert seen["problem"].observables["density"] == OperatorSum() + assert seen["sampling"] is state.sampling + assert driver.sample().chain_shape == (2, 2) + assert driver.expect().energy == -1.0 + assert driver.expect(OperatorSum()).observables["expectation"] == OperatorSum() + supplied = VMCSamples( + configs=np.zeros((2, 1), dtype=np.int64), + proposal_log_probs=np.zeros(2), + ) + assert driver.measure(samples=supplied).energy == -1.0 + assert seen["measure_kwargs"]["samples"] is supplied + assert driver.run(3).final_energy == -1.0 + assert seen["n_steps"] == 3 + + +def test_common_results_keep_chain_shape_and_shifted_energy(): + samples = VMCSamples( + configs=np.zeros((4, 2, 3), dtype=np.int64), + n_samples_per_chain=4, + n_chains=2, + ) + measurement = VMCMeasurement( + energy_mean=-1.5, + observables={"density": 0.25}, + ) + history = VMCOptimizationResult( + steps=np.arange(2), + energies=np.array([-1.0, -1.5]), + errors=np.array([0.1, 0.05]), + energy_shift=0.25, + per_site=2, + ) + + assert samples.chain_shape == (4, 2) + assert measurement.energy == -1.5 + assert measurement.observables["density"] == 0.25 + assert history.final_energy == -1.5 + assert history.final_error == 0.05 + assert np.allclose(history.shifted_energies, [-0.75, -1.25]) + assert np.allclose(history.displayed_energies, [-0.375, -0.625]) + + +def test_common_samples_distinguish_fixed_weights_from_proposal_density(): + samples = VMCSamples( + configs=np.zeros((3, 1), dtype=np.int64), + weights=np.array([0.2, 0.3, 0.5]), + ) + assert np.allclose(samples.weights, [0.2, 0.3, 0.5]) + + with pytest.raises(ValueError, match="either weights or proposal_log_probs"): + VMCSamples( + configs=np.zeros((3, 1), dtype=np.int64), + weights=np.ones(3), + proposal_log_probs=np.zeros(3), + ) + + +def test_warning_types_are_backend_neutral(): + assert issubclass(BackendCapabilityWarning, UserWarning) + assert issubclass(VMCBackendCapabilityError, NotImplementedError) + + +def test_netket_portable_adapter_rejects_external_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"): + setup.optimize(n_steps=1, weights=np.ones(2)) + + +def test_shared_configuration_objects_normalize_aliases_and_defaults(): + contraction = ContractionConfig(method="boundary_mps", chi=4, cutoff=1e-8) + sampling = SamplingConfig( + n_samples_per_chain=8, + n_chains=2, + burn_in=3, + thin=2, + ) + optimization = OptimizationConfig(method="min-sr", n_steps=4) + + assert contraction.method == "boundary" + assert contraction.chi == 4 + assert sampling.thin == 2 + assert sampling.n_samples == 16 + assert sampling.torch_kwargs()["n_samples"] == 16 + assert sampling.netket_kwargs()["n_samples"] == 16 + assert optimization.method == "minsr" + + with pytest.raises(ValueError, match="chi is required"): + ContractionConfig(method="ctmrg") + with pytest.raises(ValueError, match="n_samples_per_chain"): + SamplingConfig(n_samples_per_chain=0) + + +def test_torch_driver_consumes_shared_sampling_and_optimization_configs(): + 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) + + 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", + ) + samples = driver.sample( + sampling=SamplingConfig( + n_samples_per_chain=2, + n_chains=2, + burn_in=0, + thin=1, + seed=12, + ) + ) + assert samples.configs.shape == (2, 2, 2) + assert samples.to_common().chain_shape == (2, 2) + + history = driver.optimize( + optimization=OptimizationConfig( + method="sgd", + n_steps=1, + learning_rate=1e-2, + progress=False, + ), + sample_sweeps=1, + ) + assert len(history) == 1 + + +def test_torch_compiler_lowers_common_matrix_terms(): + torch = pytest.importorskip("torch") + from pepsy.vmc import torch_hamiltonian_connections + + common = OperatorSum.from_terms( + [ + LocalMatrixTerm( + support=(0,), + matrix=np.asarray([[1.0, 2.0], [3.0, 4.0]]), + coefficient=0.5, + ) + ], + constant=0.25, + ) + compiled = compile_operator_sum_torch(common) + configs = torch.tensor([[0], [1]], dtype=torch.long) + connections = torch_hamiltonian_connections( + configs, + compiled.terms, + constant=compiled.constant, + ) + + values = {} + for config, coefficient, batch_id in zip( + connections.configs.tolist(), + connections.coeffs.tolist(), + connections.batch_ids.tolist(), + ): + values[(int(batch_id), tuple(config))] = ( + values.get((int(batch_id), tuple(config)), 0.0) + coefficient + ) + assert values[(0, (0,))] == pytest.approx(0.75) + assert values[(0, (1,))] == pytest.approx(1.5) + assert values[(1, (0,))] == pytest.approx(1.0) + assert values[(1, (1,))] == pytest.approx(2.25) + + +def test_torch_driver_consumes_compiled_common_term_constants(): + torch = pytest.importorskip("torch") + from pepsy.vmc import TorchVMCDriver + + class ConstantAmplitude(torch.nn.Module): + def forward(self, configs): + return torch.ones(configs.shape[0], dtype=torch.float64) + + driver = TorchVMCDriver( + ConstantAmplitude(), + [], + torch.tensor([[0]], dtype=torch.long), + terms={0: torch.eye(2, dtype=torch.float64)}, + proposal="spin", + ) + compiled = compile_operator_sum_torch( + OperatorSum.from_terms( + [ + LocalMatrixTerm( + support=(0,), + matrix=np.eye(2), + ) + ], + constant=0.5, + ) + ) + connections = driver.make_connections( + torch.tensor([[0]], dtype=torch.long), + terms=compiled, + ) + assert connections.batch_ids.tolist() == [0, 0] + assert connections.configs.tolist() == [[0], [0]] + assert sum(connections.coeffs.tolist()) == pytest.approx(1.5) + + +def test_netket_facade_rejects_shared_sampling_settings_it_cannot_apply(): + from pepsy.vmc.netket import NetKetVMCSetup + + class FakeNativeSetup: + def sample(self, sampling=None): + return VMCSamples( + configs=np.zeros((2, 2, 1), dtype=np.int64), + n_samples_per_chain=2, + n_chains=2, + ) + + facade = NetKetVMCSetup( + setup=FakeNativeSetup(), + problem=VMCProblem(peps=object(), hamiltonian=object()), + ) + with pytest.raises(VMCBackendCapabilityError, match="thin"): + facade.sample(SamplingConfig(n_samples_per_chain=2, n_chains=2, thin=2)) + with pytest.raises(VMCBackendCapabilityError, match="seed/sampler_seed"): + facade.sample( + SamplingConfig(n_samples_per_chain=2, n_chains=2, seed=3) + ) + + +def test_build_netket_vmc_passes_the_common_problem_to_native_builder(monkeypatch): + import pepsy.vmc.netket as netket_vmc + + native = object() + seen = {} + + def fake_builder(peps, **kwargs): + seen["peps"] = peps + seen.update(kwargs) + return native + + monkeypatch.setattr(netket_vmc, "build_fermion_vmc", fake_builder) + problem = VMCProblem( + peps=object(), + hamiltonian=OperatorSum(), + observables={"density": OperatorSum()}, + symmetry="U1U1", + ) + sampling = SamplingConfig(n_samples_per_chain=2, n_chains=2, burn_in=0) + facade = netket_vmc.build_netket_vmc( + problem, + fermion=object(), + sampling=sampling, + ) + + assert facade.problem is problem + assert facade.native is native + assert seen["peps"] is problem.peps + assert seen["hamiltonian"] is problem.hamiltonian + assert seen["observables"] == problem.observables + assert seen["sampling"] is sampling + + +def test_torch_amplitude_accepts_common_contraction_config(): + qtn = pytest.importorskip("quimb.tensor") + torch = pytest.importorskip("torch") + from pepsy.vmc import TorchPEPSAmplitude + + peps = qtn.PEPS.rand( + Lx=2, + Ly=2, + bond_dim=2, + phys_dim=2, + seed=193, + dtype="float64", + ) + model = TorchPEPSAmplitude( + peps, + contraction=ContractionConfig(method="boundary", chi=4, cutoff=0.0), + dtype=torch.float64, + ) + assert model.contraction == "boundary" + assert model.chi == 4 + + +def test_netket_setup_consumes_shared_sampling_config(): + nk = pytest.importorskip("netket") + from pepsy.vmc.netket import NetKetPEPSVMC + + class Ansatz: + n_sites = 4 + n_params = 4 + + hilbert = nk.hilbert.Spin(s=1 / 2, N=4) + graph = nk.graph.Hypercube(length=4, n_dim=1, pbc=False) + hamiltonian = nk.operator.Ising(hilbert, graph=graph, h=1.0, J=1.0) + sampler = nk.sampler.MetropolisLocal(hilbert, n_chains=2) + vstate = nk.vqs.MCState( + sampler, + nk.models.RBM(alpha=1), + n_samples=8, + n_discard_per_chain=0, + seed=4, + ) + setup = NetKetPEPSVMC( + hilbert, + graph, + hamiltonian, + sampler, + vstate, + vstate.model, + Ansatz(), + None, + None, + ) + samples = setup.sample( + SamplingConfig(n_samples_per_chain=3, n_chains=2, burn_in=0) + ) + assert samples.chain_shape == (3, 2) + assert samples.configs.shape == (3, 2, 4) + + +def test_netket_compiler_lowers_common_fermion_terms(): + nk = pytest.importorskip("netket") + from pepsy.vmc import OperatorFactor + + hilbert = nk.hilbert.SpinOrbitalFermions( + 2, + s=1 / 2, + n_fermions_per_spin=(1, 1), + ) + common = OperatorSum.from_terms( + [ + ProductTerm( + coefficient=-1.0, + factors=( + OperatorFactor(0, "fermion", spin="up", dagger=True), + OperatorFactor(1, "fermion", spin="up", dagger=False), + ), + ) + ], + constant=0.5, + ) + + compiled = compile_operator_sum_netket(hilbert, common) + assert compiled.hilbert is hilbert + + +def test_common_spinful_fermion_operator_has_matching_exact_local_energies(): + """The Torch and NetKet compilers agree in a tiny exact Fock sector. + + This is deliberately an operator/local-energy oracle rather than a VMC + optimization test. It fixes the public ``empty, down, up, double`` Torch + codes, the NetKet spin-orbital columns, and the one fermionic coordinate + phase needed to compare their two-site Fock bases. + """ + torch = pytest.importorskip("torch") + nk = pytest.importorskip("netket") + import pepsy as py + from pepsy.vmc.netket import ( + netket_spin_orbital_columns, + occupation_to_phys_indices, + ) + from pepsy.vmc.torch import torch_hamiltonian_connections + + terms = OperatorSum.from_terms( + ( + ProductTerm( + -1.0, + ( + OperatorFactor(0, "fermion", spin="up", dagger=True), + OperatorFactor(1, "fermion", spin="up", dagger=False), + ), + ), + ProductTerm( + -1.0, + ( + OperatorFactor(1, "fermion", spin="up", dagger=True), + OperatorFactor(0, "fermion", spin="up", dagger=False), + ), + ), + ProductTerm( + -0.4, + ( + OperatorFactor(0, "fermion", spin="down", dagger=True), + OperatorFactor(1, "fermion", spin="down", dagger=False), + ), + ), + ProductTerm( + -0.4, + ( + OperatorFactor(1, "fermion", spin="down", dagger=True), + OperatorFactor(0, "fermion", spin="down", dagger=False), + ), + ), + ProductTerm( + 0.7, + ( + OperatorFactor(0, "number", spin="up"), + OperatorFactor(0, "number", spin="down"), + ), + ), + ProductTerm(0.2, (OperatorFactor(1, "number", spin="up"),)), + ), + constant=0.13, + ) + hilbert = nk.hilbert.SpinOrbitalFermions( + 2, + s=1 / 2, + n_fermions_per_spin=(1, 1), + ) + netket_matrix = np.asarray( + compile_operator_sum_netket(hilbert, terms).to_dense() + ) + rows = np.asarray(hilbert.all_states(), dtype=int) + columns = netket_spin_orbital_columns(hilbert) + configs = occupation_to_phys_indices( + rows, + columns, + phys_charges=((0, 0), (0, 1), (1, 0), (1, 1)), + ) + + compiled = compile_operator_sum_torch( + terms, + fermion=py.Fermion(symmetry="U1U1", spinful=True), + site_order=(0, 1), + ) + connections = torch_hamiltonian_connections( + torch.as_tensor(configs, dtype=torch.long), + compiled.terms, + site_order=(0, 1), + constant=compiled.constant, + ) + config_index = {tuple(config): index for index, config in enumerate(configs)} + torch_matrix = np.zeros((len(configs), len(configs)), dtype=complex) + for eta, coefficient, source in zip( + connections.configs.tolist(), + connections.coeffs.tolist(), + connections.batch_ids.tolist(), + ): + torch_matrix[config_index[tuple(eta)], source] += coefficient + + # NetKet orders modes as (down_0, down_1, up_0, up_1), while the Torch + # connection table uses the site-local (down_0, up_0, down_1, up_1) + # coordinate gauge. The fixed-number sector makes this basis phase exact. + n_up = rows[:, columns.up] + n_down = rows[:, columns.down] + phase = 1 - 2 * n_up[:, 0] * n_down[:, 1] + netket_matrix = phase[:, None] * netket_matrix * phase[None, :] + + assert np.allclose(torch_matrix, torch_matrix.conj().T, atol=1e-7) + assert np.allclose(torch_matrix, netket_matrix, atol=1e-7) + + amplitudes = np.asarray([1.0, 0.6 - 0.2j, -0.4 + 0.5j, 0.3 + 0.7j]) + torch_local_energy = (torch_matrix @ amplitudes) / amplitudes + netket_local_energy = (netket_matrix @ amplitudes) / amplitudes + assert np.allclose(torch_local_energy, netket_local_energy, atol=1e-7) + + +def test_netket_compiler_flattens_common_multisite_matrix_axes(): + nk = pytest.importorskip("netket") + hilbert = nk.hilbert.Spin(s=1 / 2, N=2) + matrix = np.arange(16, dtype=float).reshape(2, 2, 2, 2) + common = OperatorSum.from_terms( + [LocalMatrixTerm(support=(0, 1), matrix=matrix)] + ) + + compiled = compile_operator_sum_netket(hilbert, common) + reference = nk.operator.LocalOperator( + hilbert, + matrix.reshape(4, 4), + acting_on=[0, 1], + ) + assert np.allclose(compiled.to_dense(), reference.to_dense()) diff --git a/tests/test_vmc_netket.py b/tests/test_vmc_netket.py index f30beeb..af2c270 100644 --- a/tests/test_vmc_netket.py +++ b/tests/test_vmc_netket.py @@ -6,6 +6,8 @@ import pytest import quimb.tensor as qtn +from pepsy.vmc import SymmetryFallbackWarning + os.environ.setdefault("JAX_PLATFORMS", "cpu") os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") os.environ.setdefault("NETKET_NO_TIPS", "1") @@ -21,9 +23,11 @@ build_heisenberg_vmc, build_ising_vmc, build_sparse_fermi_hubbard_vmc, + build_fermion_vmc, choose_netket_chunk_size, config_to_phys_indices, fermionic_peps_rand, + fermion_model_terms, make_fermionic_peps_batched_amplitude_function, make_fermionic_peps_log_amplitude_model, make_netket_autochunk_callback, @@ -31,15 +35,22 @@ make_netket_vmc_driver, make_peps_batched_amplitude_function, make_peps_log_amplitude_model, + netket_fermion_operator, netket_spin_orbital_columns, occupation_to_phys_indices, pack_fermionic_peps_ansatz, pack_peps_ansatz, recommend_netket_vmc_settings, square_lattice_edges, + standard_fermion_observables, verify_netket_spin_columns, + VMCOptimizeResult, +) +from pepsy.vmc.netket import ( # noqa: E402 + _site_major_to_netket_jw_phase, + _spinful_phys_lookup, + _warn_flat_z2_ansatz_fixed_u1u1_sector, ) -from pepsy.vmc.netket import _spinful_phys_lookup # noqa: E402 def test_square_lattice_edges_open_boundary_order(): @@ -62,6 +73,18 @@ def test_occupation_to_phys_indices_spinful_ordering(): assert phys.tolist() == [[0, 1, 2, 3]] +def test_site_major_to_netket_jw_phase(): + columns = SpinOrbitalColumns(up=(4, 5, 6, 7), down=(0, 1, 2, 3)) + rows = np.array( + [ + [0, 1, 0, 1, 0, 1, 1, 0], + [1, 0, 1, 0, 1, 0, 0, 1], + ] + ) + phase = _site_major_to_netket_jw_phase(rows, columns) + assert phase.tolist() == [-1, 1] + + def test_config_to_phys_indices_spin_half_ordering(): config_map = NetKetLocalConfigMap.spin_half(up=0, down=1) row = np.array([[1, -1, -1, 1]]) @@ -286,11 +309,14 @@ def test_fermionic_peps_log_model_matches_direct_contraction(): for k, site in enumerate(ansatz.sites) }) direct = tnx.contract(all) + phase = _site_major_to_netket_jw_phase( + np.asarray(row), columns, site_to_orb=ansatz.site_to_orb + )[0] variables = model.init(jax.random.PRNGKey(0), row) log_amp = model.apply(variables, row)[0] amp = jax.block_until_ready(jnp.exp(log_amp)) - assert np.allclose(np.asarray(amp), np.asarray(direct)) + assert np.allclose(np.asarray(amp), phase * np.asarray(direct)) def test_fermionic_peps_batched_amplitude_function_matches_direct_contraction(): @@ -340,7 +366,10 @@ def test_fermionic_peps_batched_amplitude_function_matches_direct_contraction(): }) direct.append(tnx.contract(all)) - assert np.allclose(np.asarray(amps), np.asarray(direct)) + direct = np.asarray(direct) * _site_major_to_netket_jw_phase( + np.asarray(rows), columns, site_to_orb=ansatz.site_to_orb + ) + assert np.allclose(np.asarray(amps), direct) batched_me = make_fermionic_peps_batched_amplitude_function( ansatz, @@ -350,7 +379,7 @@ def test_fermionic_peps_batched_amplitude_function_matches_direct_contraction(): jit=False, ) mantissa, exponent = batched_me(rows) - assert np.allclose(np.asarray(mantissa), np.asarray(direct)) + assert np.allclose(np.asarray(mantissa), direct) assert np.allclose(np.asarray(exponent), np.zeros(2)) batched_hotrg = make_fermionic_peps_batched_amplitude_function( @@ -375,6 +404,49 @@ def test_fermionic_peps_batched_amplitude_function_matches_direct_contraction(): assert np.asarray(exponent).shape == (2,) +def test_fermionic_peps_amplitude_rejects_nonzero_cutoff_under_jit(): + sr = pytest.importorskip("symmray") + pytest.importorskip("jax") + + peps = sr.networks.PEPS_fermionic_rand( + "Z2", + 2, + 2, + 2, + phys_dim=4, + subsizes="equal", + flat=True, + seed=4, + ) + ansatz = pack_fermionic_peps_ansatz(peps, lattice_shape=(2, 2)) + columns = SpinOrbitalColumns(up=(4, 5, 6, 7), down=(0, 1, 2, 3)) + + with pytest.raises(ValueError, match="cutoff=0.0"): + make_fermionic_peps_batched_amplitude_function( + ansatz, + columns, + contraction="ctmrg", + chi=2, + cutoff=1e-8, + ) + + # cutoff is fine on the non-jitted adaptive path. + make_fermionic_peps_batched_amplitude_function( + ansatz, + columns, + contraction="ctmrg", + chi=2, + cutoff=1e-8, + jit=False, + ) + + +def test_warmup_netket_vmc_is_exported(): + import pepsy.vmc as pvmc + + assert callable(pvmc.warmup_netket_vmc) + + def test_build_fermi_hubbard_vmc_tiny_setup(): pytest.importorskip("netket") pytest.importorskip("flax") @@ -412,6 +484,199 @@ def test_build_fermi_hubbard_vmc_tiny_setup(): assert setup.make_driver() is not None +def test_netket_fermion_operator_matches_hubbard(): + nk = pytest.importorskip("netket") + import pepsy as py + + hilbert = nk.hilbert.SpinOrbitalFermions( + 4, s=1 / 2, n_fermions_per_spin=(2, 2) + ) + graph = nk.graph.Graph(edges=[(0, 1), (1, 2), (2, 3)], n_nodes=4) + reference = nk.operator.FermiHubbardJax( + hilbert, graph=graph, t=1.0, U=8.0, dtype=float + ) + + fermion = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0) + edges = [(0, 1), (1, 2), (2, 3)] + terms = fermion_model_terms(fermion, edges, n_sites=4) + general = netket_fermion_operator(hilbert, terms) + assert np.allclose(general.to_dense(), reference.to_dense()) + + conserving = netket_fermion_operator(hilbert, terms, conserving="auto") + assert np.allclose(conserving.to_dense(), reference.to_dense()) + + +def test_build_fermion_vmc_tiny_setup_and_observables(): + nk = pytest.importorskip("netket") + pytest.importorskip("flax") + sr = pytest.importorskip("symmray") + import types + + import pepsy as py + + peps = sr.networks.PEPS_fermionic_rand( + "Z2", + 2, + 2, + 2, + phys_dim=4, + subsizes="equal", + flat=True, + seed=2, + ) + hilbert = nk.hilbert.SpinOrbitalFermions( + 4, s=1 / 2, n_fermions_per_spin=(2, 2) + ) + fermion = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0) + setup = build_fermion_vmc( + peps, + fermion=fermion, + Lx=2, + Ly=2, + observables=standard_fermion_observables(hilbert), + n_samples=16, + n_chains=4, + n_discard_per_chain=0, + chunk_size=8, + seed=1, + sampler_seed=2, + use_sr=False, + ) + assert setup.hilbert.n_states == 36 + assert setup.n_sites == 4 + assert setup.hamiltonian is not None + assert set(setup.observables) == { + "n_up", + "n_down", + "n_total", + "double_occupancy", + } + # Hamiltonian reproduces plain Fermi-Hubbard on the same graph. + reference = nk.operator.FermiHubbardJax( + setup.hilbert, graph=setup.graph, t=1.0, U=8.0, dtype=float + ) + assert np.allclose(setup.hamiltonian.to_dense(), reference.to_dense()) + # measure() with no observables anywhere raises before touching the state. + empty = types.SimpleNamespace(observables=None) + with pytest.raises(ValueError): + NetKetPEPSVMC.measure(empty) + + +def test_build_fermion_vmc_infers_geometry_from_native_terms(): + nk = pytest.importorskip("netket") + pytest.importorskip("flax") + sr = pytest.importorskip("symmray") + import types + + import pepsy as py + + peps = sr.networks.PEPS_fermionic_rand( + "Z2", + 2, + 2, + 2, + phys_dim=4, + subsizes="equal", + flat=True, + seed=2, + ) + fermion = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0) + # Native coordinate-keyed terms with only the two vertical edges present, + # plus on-site keys. Values are irrelevant to geometry inference. + native_terms = { + ((0, 0), (1, 0)): 0, + ((0, 1), (1, 1)): 0, + (0, 0): 0, + (0, 1): 0, + (1, 0): 0, + (1, 1): 0, + } + # (i, j) -> i * Ly + j on a 2x2 lattice: those edges are (0, 2) and (1, 3). + expected_edges = {(0, 2), (1, 3)} + + setup = build_fermion_vmc( + peps, + fermion=fermion, + terms=native_terms, + Lx=2, + Ly=2, + n_samples=16, + n_chains=4, + n_discard_per_chain=0, + chunk_size=8, + seed=1, + sampler_seed=2, + use_sr=False, + ) + got = {tuple(sorted(edge)) for edge in setup.graph.edges()} + assert got == expected_edges + reference = nk.operator.FermiHubbardJax( + setup.hilbert, graph=setup.graph, t=1.0, U=8.0, dtype=float + ) + assert np.allclose(setup.hamiltonian.to_dense(), reference.to_dense()) + + # Passing a native SymHamiltonian-like object as hamiltonian= works too. + ham_like = types.SimpleNamespace(terms=native_terms) + setup2 = build_fermion_vmc( + peps, + fermion=fermion, + hamiltonian=ham_like, + Lx=2, + Ly=2, + n_samples=16, + n_chains=4, + n_discard_per_chain=0, + chunk_size=8, + seed=1, + sampler_seed=2, + use_sr=False, + ) + got2 = {tuple(sorted(edge)) for edge in setup2.graph.edges()} + assert got2 == expected_edges + + +def test_netket_vmc_optimize_returns_history(): + pytest.importorskip("netket") + pytest.importorskip("flax") + + peps = qtn.PEPS.rand( + Lx=2, + Ly=2, + bond_dim=2, + phys_dim=2, + seed=7, + dtype="complex128", + ) + setup = build_ising_vmc( + peps, + Lx=2, + Ly=2, + h=1.0, + J=0.7, + n_samples=16, + n_chains=4, + n_discard_per_chain=0, + chunk_size=8, + seed=1, + sampler_seed=2, + use_sr=False, + ) + result = setup.optimize( + 3, + learning_rate=0.01, + progress=False, + warmup=True, + energy_shift=-2.0, + per_site=4, + ) + assert isinstance(result, VMCOptimizeResult) + assert result.energies.size == 3 + assert result.steps.shape == result.energies.shape + assert np.isfinite(result.final_energy) + assert result.compile_seconds is not None + assert np.allclose(result.shifted_energies, result.energies - 2.0) + + def test_build_ising_vmc_tiny_setup_and_expectation(): pytest.importorskip("netket") pytest.importorskip("flax") @@ -563,6 +828,21 @@ def test_spinful_phys_lookup_u1u1_and_z2_fallback(): assert _spinful_phys_lookup(()) is None +def test_flat_z2_ansatz_warns_about_fixed_u1u1_netket_sector(): + pytest.importorskip("symmray") + peps = fermionic_peps_rand( + "Z2", 2, 2, 2, seed=3 + ) + ansatz = pack_fermionic_peps_ansatz(peps, lattice_shape=(2, 2)) + assert ansatz.uses_flat_symmray is True + + with pytest.warns( + SymmetryFallbackWarning, + match=r"flat Z2 fermionic PEPS ansatz.*fixed U1U1 sampling sector", + ): + _warn_flat_z2_ansatz_fixed_u1u1_sector(ansatz, (1, 1)) + + def test_occupation_fold_switches_on_phys_charges(): columns = SpinOrbitalColumns(up=(2, 3), down=(0, 1)) # one row per (n_up, n_down) for a single orbital pair (2 orbitals here). diff --git a/tests/test_vmc_torch.py b/tests/test_vmc_torch.py index 3aa2f36..7cbbfbe 100644 --- a/tests/test_vmc_torch.py +++ b/tests/test_vmc_torch.py @@ -1,31 +1,66 @@ """Tests for Pepsy's optional torch VMC kernels.""" +import warnings + import pytest import quimb.tensor as qtn +import numpy as np +from itertools import product +from types import SimpleNamespace torch = pytest.importorskip("torch") +from pepsy.tensors import Fermion, hrs_to_peps, ps_to_peps # noqa: E402 +from pepsy.sampling import PepsBpSampler # noqa: E402 +import pepsy.vmc.torch as torch_vmc # noqa: E402 from pepsy.vmc.torch import ( # noqa: E402 FermionSiteEncoding, + SpinlessSiteEncoding, TorchPEPSAmplitude, TorchPEPSBoundaryAmplitude, + TorchFermionVMC, + TorchFermionVMCMetadata, + TorchChainDiagnostics, + TorchConnections, + TorchMCMCSamples, + TorchMetropolisSampler, + TorchBPMetropolisSampler, TorchVMCDriver, + TorchVMCSetup, + TorchVMCEnergyEstimate, + TorchVMCImportanceEstimate, TorchVMCStepResult, TorchSquareLattice, apply_torch_sr_update, + build_torch_vmc, count_spinful_particles, heisenberg_connections, local_energy_from_connections, make_torch_peps_amplitude_model, + metropolis_local_sampler, metropolis_exchange_sweep, propose_spin_exchange, propose_spinful_exchange_or_hopping, + propose_spinful_u1_exchange_or_hopping, + propose_spinful_z2_exchange_or_hopping, + propose_spinful_z2z2_exchange_or_hopping, random_spin_configs, random_spinful_configs, solve_torch_sr, spinful_fermi_hubbard_connections, torch_log_derivative_matrix, transverse_ising_connections, + torch_hamiltonian_connections, + torch_chain_diagnostics, + _observable_statistics, +) +from pepsy.vmc import ( # noqa: E402 + ContractionConfig, + OptimizationConfig, + SamplingConfig, + VMCMeasurement, + VMCOptimizationResult, + VMCProblem, ) @@ -40,6 +75,17 @@ def forward(self, configs): return self.weights[configs].prod(dim=1) +class ComplexProductAmplitude(torch.nn.Module): + def __init__(self): + super().__init__() + self.weights = torch.nn.Parameter( + torch.tensor([1.3 + 0.4j, 0.8 - 0.2j], dtype=torch.complex128) + ) + + def forward(self, configs): + return self.weights[configs].prod(dim=1) + + class CountingAmplitude: def __init__(self): self.calls = [] @@ -49,6 +95,23 @@ def __call__(self, configs): return torch.ones(configs.shape[0], dtype=torch.float64, device=configs.device) +class LogOnlyAmplitude: + """Amplitude whose raw values underflow but whose log values are usable.""" + + def __call__(self, configs): + return torch.zeros(configs.shape[0], dtype=torch.float64) + + def forward_log(self, configs): + high = configs[:, 0] == 1 + log_abs = torch.where( + high, + torch.full((configs.shape[0],), -900.0, dtype=torch.float64), + torch.full((configs.shape[0],), -1000.0, dtype=torch.float64), + ) + phase = torch.ones(configs.shape[0], dtype=torch.complex128) + return phase, log_abs + + def test_fermion_site_encoding_supports_symmray_and_vmc_torch_orders(): symm = FermionSiteEncoding.symmray() vmct = FermionSiteEncoding.vmc_torch() @@ -69,6 +132,80 @@ def test_fermion_site_encoding_supports_symmray_and_vmc_torch_orders(): symm.decode(torch.tensor([[9]])) +def test_fermion_site_encoding_derives_u1u1_peps_charge_order(): + fermion = Fermion(spinful=True, symmetry="U1U1") + encoding = FermionSiteEncoding.from_fermion( + fermion, + physical_charges=((0, 0), (0, 1), (1, 0), (1, 1)), + ) + assert encoding == FermionSiteEncoding.vmc_torch() + assert encoding.encode(torch.tensor([[1, 0]]), torch.tensor([[0, 1]])).tolist() == [[2, 1]] + + +def test_spinful_u1_proposal_preserves_total_and_can_change_spin_sector(): + encoding = FermionSiteEncoding.vmc_torch() + configs = torch.tensor([[encoding.up, encoding.down]]) + before_up, before_down = count_spinful_particles(configs, encoding=encoding) + proposed, changed = propose_spinful_u1_exchange_or_hopping( + 0, + 1, + configs, + spin_flip_rate=1.0, + encoding=encoding, + generator=torch.Generator().manual_seed(7), + ) + after_up, after_down = count_spinful_particles(proposed, encoding=encoding) + assert changed.tolist() == [True] + assert (after_up + after_down).tolist() == (before_up + before_down).tolist() + assert ( + not torch.equal(after_up, before_up) + or not torch.equal(after_down, before_down) + ) + + +def test_spinful_z2_proposal_preserves_parity_and_can_change_number(): + encoding = FermionSiteEncoding.symmray() + configs = torch.tensor([[encoding.empty, encoding.empty]]) + before_up, before_down = count_spinful_particles(configs, encoding=encoding) + proposed, changed = propose_spinful_z2_exchange_or_hopping( + 0, + 1, + configs, + hopping_rate=0.0, + spin_flip_rate=0.0, + pair_toggle_rate=1.0, + encoding=encoding, + generator=torch.Generator().manual_seed(8), + ) + after_up, after_down = count_spinful_particles(proposed, encoding=encoding) + assert changed.tolist() == [True] + assert ((after_up + after_down) % 2).tolist() == ( + (before_up + before_down) % 2 + ).tolist() + assert (after_up + after_down).tolist() == [2] + + +def test_spinful_z2z2_proposal_preserves_resolved_parities(): + encoding = FermionSiteEncoding.vmc_torch() + configs = torch.tensor([[encoding.empty, encoding.empty]]) + before_up, before_down = count_spinful_particles(configs, encoding=encoding) + proposed, changed = propose_spinful_z2z2_exchange_or_hopping( + 0, + 1, + configs, + hopping_rate=0.0, + pair_toggle_rate=1.0, + encoding=encoding, + generator=torch.Generator().manual_seed(8), + ) + after_up, after_down = count_spinful_particles(proposed, encoding=encoding) + assert changed.tolist() == [True] + assert ((after_up % 2).tolist(), (after_down % 2).tolist()) == ( + (before_up % 2).tolist(), + (before_down % 2).tolist(), + ) + + def test_torch_square_lattice_edges_match_row_major_open_boundary(): graph = TorchSquareLattice(2, 3) assert graph.row_edges == { @@ -119,6 +256,148 @@ def test_torch_peps_amplitude_matches_direct_quimb_contraction(): assert model.n_params == sum(p.numel() for p in model.parameters()) +@pytest.mark.parametrize( + ("Lx", "Ly", "bond_dim"), + [ + (3, 4, 2), + (3, 4, 4), + (4, 4, 2), + (4, 4, 4), + ], +) +def test_torch_peps_amplitude_contractions_match_exact_reference( + Lx, + Ly, + bond_dim, +): + """All PEPS contraction modes agree with direct contraction at large chi.""" + peps = qtn.PEPS.rand( + Lx=Lx, + Ly=Ly, + bond_dim=bond_dim, + phys_dim=2, + seed=1000 + 100 * Lx + 10 * Ly + bond_dim, + dtype="float64", + ) + n_sites = Lx * Ly + parity = torch.arange(n_sites, dtype=torch.long) % 2 + rows = torch.stack((parity, 1 - parity, (torch.arange(n_sites) // 2) % 2)) + + reference = torch.as_tensor( + [ + peps.isel({ + peps.site_ind(site): int(row[i]) + for i, site in enumerate(peps.sites) + }).contract(all) + for row in rows + ], + dtype=torch.float64, + ) + + # Use zero cutoff and a generous chi so these are contraction-path + # consistency checks. Production runs can deliberately choose a smaller + # chi and should then report the resulting approximation error. + chi = 16 if bond_dim == 2 else 64 + for contraction in ("exact", "boundary", "ctmrg", "hotrg"): + kwargs = {"contraction": contraction, "dtype": torch.float64} + if contraction != "exact": + kwargs.update(chi=chi, cutoff=0.0) + model = TorchPEPSAmplitude(peps, **kwargs) + amplitudes = model(rows) + + assert torch.isfinite(amplitudes).all() + assert torch.allclose( + amplitudes, + reference, + rtol=2.0e-10, + atol=1.0e-8, + ), contraction + + phase, log_abs = model.forward_log(rows) + assert torch.allclose( + phase * torch.exp(log_abs), + amplitudes, + rtol=2.0e-10, + atol=1.0e-8, + ), contraction + + +@pytest.mark.parametrize("bond_dim", [3, 4]) +def test_torch_peps_4x6_contractions_honor_production_options(bond_dim): + """Boundary and CTMRG options remain accurate on a larger flat PEPS TN.""" + Lx, Ly = 4, 6 + peps = qtn.PEPS.rand( + Lx=Lx, + Ly=Ly, + bond_dim=bond_dim, + phys_dim=2, + seed=1400 + bond_dim, + dtype="float64", + ) + n_sites = Lx * Ly + parity = torch.arange(n_sites, dtype=torch.long) % 2 + rows = torch.stack((parity, 1 - parity, (torch.arange(n_sites) // 2) % 2)) + reference = torch.as_tensor( + [ + peps.isel({ + peps.site_ind(site): int(row[i]) + for i, site in enumerate(peps.sites) + }).contract(all) + for row in rows + ], + dtype=torch.float64, + ) + + cutoff = 1.0e-10 + chi_contract = 32 if bond_dim == 3 else 64 + contract_opt = "auto-hq" + boundary_opts = { + "mode": "mps", + "final_contract": True, + "final_contract_opts": {"optimize": contract_opt}, + "sequence": ["xmin", "xmax", "ymin", "ymax"], + "equalize_norms": False, + "progbar": False, + } + ctmrg_opts = { + "final_contract": False, + "final_contract_opts": {"optimize": contract_opt}, + "max_separation": 1, + "inplace": False, + "equalize_norms": False, + "progbar": False, + } + + for contraction, contraction_opts in ( + ("boundary", boundary_opts), + ("ctmrg", ctmrg_opts), + ): + model = TorchPEPSAmplitude( + peps, + contraction=contraction, + chi=chi_contract, + cutoff=cutoff, + contraction_opts=contraction_opts, + dtype=torch.float64, + ) + amplitudes = model(rows) + assert torch.isfinite(amplitudes).all() + assert torch.allclose( + amplitudes, + reference, + rtol=2.0e-10, + atol=1.0e-6, + ), (bond_dim, contraction) + + phase, log_abs = model.forward_log(rows) + assert torch.allclose( + phase * torch.exp(log_abs), + amplitudes, + rtol=2.0e-10, + atol=1.0e-6, + ), (bond_dim, contraction) + + def test_torch_peps_amplitude_supports_torch_optimizer_step(): peps = qtn.PEPS.rand( Lx=2, @@ -166,6 +445,112 @@ def test_torch_log_derivative_matrix_matches_product_model_manual_values(): assert torch.allclose(log_derivatives, expected) +def test_torch_log_derivative_matrix_supports_complex_holomorphic_parameters(): + model = ComplexProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]]) + + log_derivatives = torch_log_derivative_matrix( + model, + configs, + complex_parameter_mode="holomorphic", + ) + expected = model.weights.detach().reciprocal().expand(2, -1) + assert torch.allclose(log_derivatives, expected, rtol=1.0e-10, atol=1.0e-12) + + # Check both real and imaginary parameter directions against finite + # differences. This catches conjugation mistakes in complex autograd. + epsilon = 1.0e-6 + amplitudes = model(configs).detach() + weights = model.weights.detach() + for parameter in range(weights.numel()): + plus = weights.clone() + minus = weights.clone() + plus[parameter] += epsilon + minus[parameter] -= epsilon + finite_real = ( + plus[configs].prod(dim=1) - minus[configs].prod(dim=1) + ) / (2.0 * epsilon) / amplitudes + + plus = weights.clone() + minus = weights.clone() + plus[parameter] += 1j * epsilon + minus[parameter] -= 1j * epsilon + finite_imag = ( + plus[configs].prod(dim=1) - minus[configs].prod(dim=1) + ) / (2.0j * epsilon) / amplitudes + + assert torch.allclose( + finite_real, + log_derivatives[:, parameter], + rtol=1.0e-8, + atol=1.0e-9, + ) + assert torch.allclose( + finite_imag, + log_derivatives[:, parameter], + rtol=1.0e-8, + atol=1.0e-9, + ) + + +def test_torch_log_derivative_matrix_supports_complex_real_imag_parameters(): + model = ComplexProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]]) + log_derivatives = torch_log_derivative_matrix( + model, + configs, + complex_parameter_mode="real-imag", + ) + weights = model.weights.detach() + expected_row = torch.stack( + ( + weights[0].reciprocal(), + 1j * weights[0].reciprocal(), + weights[1].reciprocal(), + 1j * weights[1].reciprocal(), + ) + ) + assert log_derivatives.shape == (2, 4) + assert torch.allclose( + log_derivatives, + expected_row.expand(2, -1), + rtol=1.0e-10, + atol=1.0e-12, + ) + + epsilon = 1.0e-6 + amplitudes = model(configs).detach() + for parameter in range(weights.numel()): + plus = weights.clone() + minus = weights.clone() + plus[parameter] += epsilon + minus[parameter] -= epsilon + finite_real = ( + plus[configs].prod(dim=1) - minus[configs].prod(dim=1) + ) / (2.0 * epsilon) / amplitudes + + plus = weights.clone() + minus = weights.clone() + plus[parameter] += 1j * epsilon + minus[parameter] -= 1j * epsilon + finite_imag = ( + plus[configs].prod(dim=1) - minus[configs].prod(dim=1) + ) / (2.0 * epsilon) / amplitudes + + assert torch.allclose( + finite_real, + log_derivatives[:, 2 * parameter], + rtol=1.0e-8, + atol=1.0e-9, + ) + assert torch.allclose( + finite_imag, + log_derivatives[:, 2 * parameter + 1], + rtol=1.0e-8, + atol=1.0e-9, + ) + + def test_solve_torch_sr_direct_matches_minsr(): generator = torch.Generator().manual_seed(13) log_derivatives = torch.randn( @@ -202,6 +587,204 @@ def test_solve_torch_sr_direct_matches_minsr(): assert torch.allclose(auto.direction, minsr.direction, atol=1.0e-10) +def test_solve_torch_sr_weighted_direct_matches_minsr(): + generator = torch.Generator().manual_seed(133) + log_derivatives = torch.randn( + 5, + 8, + dtype=torch.float64, + generator=generator, + ) + local_energies = torch.randn(5, dtype=torch.float64, generator=generator) + weights = torch.tensor([0.05, 0.10, 0.15, 0.25, 0.45], dtype=torch.float64) + + direct = solve_torch_sr( + log_derivatives, + local_energies, + sample_weights=weights, + method="direct", + diag_shift=1.0e-3, + ) + minsr = solve_torch_sr( + log_derivatives, + local_energies, + sample_weights=weights, + method="minsr", + diag_shift=1.0e-3, + ) + + assert torch.allclose(direct.direction, minsr.direction, atol=1.0e-10) + assert direct.info["effective_sample_size"] == pytest.approx( + float(1.0 / weights.square().sum()) + ) + + +def test_solve_torch_sr_complex_matches_minsr_and_reports_residual(): + generator = torch.Generator().manual_seed(130) + log_derivatives = torch.randn( + 6, + 4, + dtype=torch.complex128, + generator=generator, + ) + 1j * torch.randn(6, 4, dtype=torch.complex128, generator=generator) + local_energies = torch.randn( + 6, + dtype=torch.complex128, + generator=generator, + ) + 1j * torch.randn(6, dtype=torch.complex128, generator=generator) + + direct = solve_torch_sr( + log_derivatives, + local_energies, + method="direct", + diag_shift=1.0e-3, + ) + minsr = solve_torch_sr( + log_derivatives, + local_energies, + method="minsr", + diag_shift=1.0e-3, + ) + + assert torch.allclose(direct.direction, minsr.direction, atol=1.0e-10) + assert direct.info["relative_residual"] < 1.0e-10 + assert minsr.info["relative_residual"] < 1.0e-10 + + model = ComplexProductAmplitude() + before = model.weights.detach().clone() + direction = torch.tensor([0.2 - 0.1j, -0.3 + 0.4j], dtype=torch.complex128) + apply_torch_sr_update(model, direction, learning_rate=0.25) + assert torch.allclose(model.weights, before - 0.25 * direction) + + +def test_solve_torch_sr_real_imag_matches_minsr_and_applies_coordinates(): + generator = torch.Generator().manual_seed(131) + log_derivatives = torch.randn( + 6, + 4, + dtype=torch.complex128, + generator=generator, + ) + 1j * torch.randn(6, 4, dtype=torch.complex128, generator=generator) + local_energies = torch.randn( + 6, + dtype=torch.complex128, + generator=generator, + ) + 1j * torch.randn(6, dtype=torch.complex128, generator=generator) + + direct = solve_torch_sr( + log_derivatives, + local_energies, + method="direct", + diag_shift=1.0e-3, + parameter_mode="real-imag", + ) + minsr = solve_torch_sr( + log_derivatives, + local_energies, + method="minsr", + diag_shift=1.0e-3, + parameter_mode="real-imag", + ) + + assert not torch.is_complex(direct.direction) + assert torch.allclose(direct.direction, minsr.direction, atol=1.0e-10) + assert direct.info["relative_residual"] < 1.0e-10 + assert minsr.info["relative_residual"] < 1.0e-10 + + model = ComplexProductAmplitude() + before = model.weights.detach().clone() + direction = torch.tensor([0.2, -0.1, -0.3, 0.4], dtype=torch.float64) + apply_torch_sr_update( + model, + direction, + learning_rate=0.25, + parameter_mode="real-imag", + ) + expected = before - 0.25 * torch.tensor( + [0.2 - 0.1j, -0.3 + 0.4j], + dtype=torch.complex128, + ) + assert torch.allclose(model.weights, expected) + + +def test_torch_sr_uses_scheduled_shift_and_cholesky_when_well_conditioned(): + log_derivatives = torch.tensor( + [[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]], + dtype=torch.float64, + ) + local_energies = torch.tensor([1.0, -1.0, 0.5, -0.5]) + calls = [] + + def schedule(step): + calls.append(step) + return 0.25 / (step + 1) + + result = solve_torch_sr( + log_derivatives, + local_energies, + method="direct", + diag_shift=schedule, + step=1, + ) + + assert calls == [1] + assert result.diag_shift == 0.125 + assert result.info["step"] == 1 + assert result.info["solver"] == "cholesky" + + +def test_torch_sr_falls_back_to_pseudoinverse_for_rank_deficient_metric(): + log_derivatives = torch.tensor( + [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]], + dtype=torch.float64, + ) + local_energies = torch.tensor([1.0, 0.0, -1.0], dtype=torch.float64) + + result = solve_torch_sr( + log_derivatives, + local_energies, + method="direct", + diag_shift=0.0, + ) + + assert result.info["solver"] == "pinv" + assert torch.isfinite(result.direction).all() + + +def test_torch_sr_spring_momentum_keeps_only_unsampled_complement(): + log_derivatives = torch.tensor( + [[1.0, 0.0], [-1.0, 0.0], [1.0, 0.0], [-1.0, 0.0]], + dtype=torch.float64, + ) + local_energies = torch.tensor([1.0, -1.0, 1.0, -1.0]) + previous_direction = torch.tensor([0.0, 2.0], dtype=torch.float64) + + without_momentum = solve_torch_sr( + log_derivatives, + local_energies, + method="direct", + diag_shift=1.0e-2, + ) + with_momentum = solve_torch_sr( + log_derivatives, + local_energies, + method="direct", + diag_shift=1.0e-2, + momentum=0.5, + previous_direction=previous_direction, + ) + + assert torch.allclose( + with_momentum.direction[0], + without_momentum.direction[0], + ) + assert torch.allclose( + with_momentum.direction[1], + torch.tensor(1.0, dtype=torch.float64), + ) + assert with_momentum.info["spring_complement_norm"] == pytest.approx(2.0) + + def test_apply_torch_sr_update_changes_model_parameters(): model = ProductAmplitude() before = model.weights.detach().clone() @@ -238,6 +821,37 @@ def test_torch_peps_amplitude_supports_sr_kernel(): assert torch.isfinite(result.direction).all() +@pytest.mark.parametrize("parameter_mode", ("holomorphic", "real-imag")) +def test_torch_peps_batched_log_derivatives_match_loop(parameter_mode): + """The batched PEPS Jacobian must preserve the legacy SR derivatives.""" + peps = qtn.PEPS.rand( + Lx=2, + Ly=2, + bond_dim=2, + phys_dim=2, + seed=16, + dtype="complex128", + ) + model = TorchPEPSAmplitude(peps, contraction="exact", dtype=torch.complex128) + rows = torch.tensor([[0, 1, 0, 1], [1, 0, 1, 0]]) + + batched = torch_log_derivative_matrix( + model, + rows, + complex_parameter_mode=parameter_mode, + derivative_backend="batched", + ) + loop = torch_log_derivative_matrix( + model, + rows, + complex_parameter_mode=parameter_mode, + derivative_backend="loop", + ) + + assert batched.shape == loop.shape + assert torch.allclose(batched, loop, rtol=1.0e-8, atol=1.0e-10) + + def test_torch_peps_boundary_amplitude_reuses_connected_environments(): peps = qtn.PEPS.rand( Lx=3, @@ -278,71 +892,1765 @@ def test_torch_peps_boundary_amplitude_reuses_connected_environments(): assert model.last_connected_reuse_stats["num_fallback"] == 0 -def test_local_energy_reuses_diagonal_connections_and_chunks_offdiagonal(): - configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) - amps = torch.ones(2, dtype=torch.float64) - conn = heisenberg_connections(configs, [(0, 1)], J=1.0) - amplitude = CountingAmplitude() +def test_torch_peps_boundary_cached_closure_uses_final_contraction_options(): + """Cached local estimators must retain the caller's scalar path choice.""" + peps = qtn.PEPS.rand( + Lx=2, + Ly=2, + bond_dim=2, + phys_dim=2, + seed=150, + dtype="float64", + ) + optimizer = object() + model = TorchPEPSBoundaryAmplitude( + peps, + chi=4, + cutoff=0.0, + contraction_opts={"final_contract_opts": {"optimize": optimizer}}, + dtype=torch.float64, + ) - energy = local_energy_from_connections( - configs, - amps, - conn, - amplitude, - chunk_size=1, + class ReusedBoundary: + def __or__(self, other): + return self + + def view_as_(self, *args, **kwargs): + self.view_args = args + self.view_kwargs = kwargs + return self + + def contract_boundary_from_xmin_(self, **kwargs): + self.boundary_kwargs = kwargs + + def contract(self, *args, **kwargs): + self.contract_args = args + self.contract_kwargs = kwargs + return torch.tensor(1.0, dtype=torch.float64) + + reused = ReusedBoundary() + envs = {("xmin", 0): reused, ("xmax", 0): reused} + result = model._contract_axis_strip( + None, + object(), + "x", + (0,), + envs, + torch.tensor(1.0, dtype=torch.float64), ) - assert torch.allclose(energy, torch.tensor([0.25, 0.25], dtype=torch.float64)) - assert [tuple(call.shape) for call in amplitude.calls] == [(1, 2), (1, 2)] + assert result.item() == 1.0 + assert reused.contract_args == (all,) + assert reused.contract_kwargs == {"optimize": optimizer} -def test_torch_vmc_driver_runs_sampling_and_energy_estimate(): - model = ProductAmplitude() - configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) - driver = TorchVMCDriver( - model, +def test_torch_peps_final_contraction_retries_empty_reusable_tree(): + peps = qtn.PEPS.rand( + Lx=2, + Ly=2, + bond_dim=2, + phys_dim=2, + seed=1501, + dtype="float64", + ) + optimizer = object() + model = TorchPEPSAmplitude( + peps, + contraction="boundary", + chi=4, + contraction_opts={"final_contract_opts": {"optimize": optimizer}}, + dtype=torch.float64, + ) + + class EmptyReusableTree: + def __init__(self): + self.optimizers = [] + + def contract(self, *args, **kwargs): + self.optimizers.append(kwargs["optimize"]) + if kwargs["optimize"] is optimizer: + raise KeyError("tree") + return torch.tensor(1.0, dtype=torch.float64) + + closure = EmptyReusableTree() + contract_kwargs = {} + + def approximate_contract(**kwargs): + contract_kwargs.update(kwargs) + return closure + + with pytest.warns(RuntimeWarning, match="produced no cotengra tree"): + result = model._contract_approximate( + approximate_contract, + close_final=True, + final_contract=True, + final_contract_opts={"optimize": optimizer}, + ) + + assert result.item() == 1.0 + assert contract_kwargs["final_contract"] is False + assert closure.optimizers == [optimizer, "auto-hq"] + assert model.final_optimizer_fallbacks == 1 + + +def test_torch_peps_boundary_tries_alternative_axis_before_fallback(monkeypatch): + """A separated PBC-style update should retain a boundary reuse route.""" + peps = qtn.PEPS.rand( + Lx=3, + Ly=3, + bond_dim=2, + phys_dim=2, + seed=151, + dtype="float64", + ) + model = TorchPEPSBoundaryAmplitude( + peps, + chi=4, + cutoff=0.0, + dtype=torch.float64, + ) + parent = torch.tensor([[0, 1, 0, 1, 0, 1, 0, 1, 0]]) + target = parent.clone() + # The endpoints differ in both coordinates. The short x window is tried + # first; y is the retained alternative boundary path. + first = model.sites.index((0, 0)) + second = model.sites.index((1, 2)) + target[0, first] = 1 - target[0, first] + target[0, second] = 1 - target[0, second] + current = model(parent) + connections = TorchConnections( + configs=target, + coeffs=torch.ones(1, dtype=torch.float64), + batch_ids=torch.zeros(1, dtype=torch.long), + ) + + original = model._contract_cached_axis_window + + def reject_x(*args, **kwargs): + if args[3] == "x": + raise RuntimeError("exercise the alternative boundary axis") + return original(*args, **kwargs) + + monkeypatch.setattr(model, "_contract_cached_axis_window", reject_x) + connected = model.connected_amplitudes(parent, current, connections) + fresh = model(target) + + assert torch.allclose(connected, fresh, rtol=1.0e-7, atol=1.0e-8) + assert model.last_connected_reuse_stats["num_alternative_axis_reused"] == 1 + assert model.last_connected_reuse_stats["num_fallback"] == 0 + + +def test_torch_peps_boundary_groups_parent_strips_and_caches_templates(monkeypatch): + """Several targets in one plane should share one parent strip template.""" + peps = qtn.PEPS.rand( + Lx=3, + Ly=3, + bond_dim=2, + phys_dim=2, + seed=153, + dtype="float64", + ) + model = TorchPEPSBoundaryAmplitude( + peps, + chi=4, + cutoff=0.0, + dtype=torch.float64, + ) + parent = torch.tensor([[0, 1, 0, 1, 0, 1, 0, 1, 0]]) + pairs = (((0, 0), (0, 1)), ((0, 1), (0, 2)), ((0, 0), (0, 2))) + targets = [] + for left, right in pairs: + target = parent.clone() + for site in (left, right): + index = model.sites.index(site) + target[0, index] = 1 - target[0, index] + targets.append(target[0]) + connections = TorchConnections( + configs=torch.stack(targets), + coeffs=torch.ones(len(targets), dtype=torch.float64), + batch_ids=torch.zeros(len(targets), dtype=torch.long), + ) + amplitudes = model(parent) + + select_calls = [] + original_select = model._select_config + + def count_select(*args, **kwargs): + select_calls.append(1) + return original_select(*args, **kwargs) + + monkeypatch.setattr(model, "_select_config", count_select) + grouped = model.connected_amplitudes(parent, amplitudes, connections) + first_stats = model.last_connected_reuse_stats + + assert first_stats["num_groups"] == 1 + assert first_stats["num_grouped_connections"] == len(targets) + assert first_stats["num_strip_builds"] == 1 + assert first_stats["num_fallback"] == 0 + # One parent selection serves all three local targets; the target strips + # are cloned from the cached template and only their physical data change. + assert len(select_calls) == 1 + + fresh = model(connections.configs) + assert torch.allclose(grouped, fresh, rtol=1.0e-7, atol=1.0e-8) + + select_calls.clear() + repeated = model.connected_amplitudes(parent, amplitudes, connections) + repeated_stats = model.last_connected_reuse_stats + assert torch.allclose(repeated, fresh, rtol=1.0e-7, atol=1.0e-8) + assert repeated_stats["num_strip_cache_hits"] == 1 + assert not select_calls + + +def test_torch_peps_boundary_amplitude_batches_large_connected_sets(): + """Large boundary local-energy batches use the vmapped contraction path.""" + peps = qtn.PEPS.rand( + Lx=4, + Ly=4, + bond_dim=4, + phys_dim=2, + seed=16, + dtype="float64", + ) + model = TorchPEPSBoundaryAmplitude( + peps, + chi=64, + cutoff=0.0, + dtype=torch.float64, + ) + configs = torch.randint( + 0, + 2, + (8, 16), + generator=torch.Generator().manual_seed(17), + ) + graph = TorchSquareLattice(4, 4) + amplitudes = model(configs) + connections = heisenberg_connections(configs, graph, J=1.0) + + connected = model.connected_amplitudes( + configs, + amplitudes, + connections, + chunk_size=32, + ) + reference = model(connections.configs, chunk_size=32) + + assert torch.allclose(connected, reference, rtol=2.0e-10, atol=1.0e-8) + assert model.last_connected_reuse_stats["num_batched"] >= 64 + assert model.last_connected_reuse_stats["num_fallback"] == 0 + + +def test_torch_peps_boundary_amplitude_can_vmap_proposal_batch(): + peps = qtn.PEPS.rand( + Lx=3, + Ly=3, + bond_dim=2, + phys_dim=2, + seed=167, + dtype="float64", + ) + model = TorchPEPSBoundaryAmplitude( + peps, + chi=4, + cutoff=0.0, + dtype=torch.float64, + proposal_batching="vmap", + ) + parent = torch.tensor([ + (0, 1, 0, 1, 0, 1, 0, 1, 0), + (1, 0, 1, 0, 1, 0, 1, 0, 1), + ]) + target = parent.clone() + target[:, 0], target[:, 1] = parent[:, 1], parent[:, 0] + current = model(parent) + + proposed = model.proposal_amplitudes(parent, target, current) + + assert torch.allclose(proposed, model(target), rtol=1.0e-7, atol=1.0e-8) + assert model.last_proposal_cache_stats["num_vmapped"] == 2 + assert model.last_proposal_cache_stats["num_environment_builds"] == 0 + + +def test_torch_peps_boundary_amplitude_vmaps_flat_symmray_proposals(): + peps = _symmray_fermionic_peps("Z2", flat=True) + model = TorchPEPSBoundaryAmplitude( + peps, + chi=4, + cutoff=0.0, + dtype=torch.float64, + proposal_batching="vmap", + ) + parent = torch.tensor([(0, 0, 0, 0), (0, 0, 0, 0)]) + target = torch.tensor([(1, 1, 0, 0), (0, 0, 1, 1)]) + current = model(parent) + + proposed = model.proposal_amplitudes(parent, target, current) + + assert model.is_symmray + assert torch.allclose(proposed, model(target), rtol=1.0e-7, atol=1.0e-8) + assert model.last_proposal_cache_stats["num_vmapped"] == 2 + assert model.last_proposal_cache_stats["num_environment_builds"] == 0 + + +def test_torch_peps_boundary_factory_selects_cached_model(): + peps = qtn.PEPS.rand( + Lx=2, + Ly=2, + bond_dim=2, + phys_dim=2, + seed=161, + dtype="float64", + ) + + model = make_torch_peps_amplitude_model( + peps, + contraction="boundary", + chi=4, + cutoff=0.0, + dtype=torch.float64, + ) + + assert isinstance(model, TorchPEPSBoundaryAmplitude) + + +def test_torch_peps_boundary_caches_local_proposals_and_invalidates(): + peps = qtn.PEPS.rand( + Lx=3, + Ly=3, + bond_dim=2, + phys_dim=2, + seed=162, + dtype="float64", + ) + model = TorchPEPSBoundaryAmplitude( + peps, + chi=4, + cutoff=0.0, + dtype=torch.float64, + ) + parent = torch.tensor([[0, 1, 0, 1, 0, 1, 0, 1, 0]]) + target = parent.clone() + target[0, 0], target[0, 1] = target[0, 1], target[0, 0] + current = model(parent) + + cached = model.proposal_amplitudes(parent, target, current) + fresh = model(target) + assert torch.allclose(cached, fresh, rtol=1.0e-7, atol=1.0e-8) + assert model.last_proposal_cache_stats["num_environment_builds"] == 1 + + model.proposal_amplitudes(parent, target, current) + assert model.last_proposal_cache_stats["num_transition_cache_hits"] == 1 + + with torch.no_grad(): + model.params[0].add_(0.01) + model.proposal_amplitudes(parent, target, current) + assert model.last_proposal_cache_stats["num_transition_cache_hits"] == 0 + + +@pytest.mark.parametrize( + ("symmetry", "spinful", "shape", "occupations"), + ( + ("U1", False, (1, 2), (1, 0)), + ( + "U1U1", + True, + (2, 2), + [(1, 0), (0, 1), (1, 0), (0, 1)], + ), + ), +) +def test_torch_peps_boundary_serial_amplitude_cache_deduplicates_and_invalidates( + symmetry, spinful, shape, occupations +): + """Serial native-symmetry amplitudes reuse values and track updates.""" + fermion = Fermion( + symmetry=symmetry, + spinful=spinful, + t=1.0, + U=2.0, + ) + peps = ps_to_peps( + *shape, + fermion=fermion, + occupations=occupations, + dtype="complex128", + seed=163, + ) + vmc = TorchFermionVMC( + peps, + fermion, + n_walkers=2, + dtype=torch.complex128, + amplitude_batching="serial", + seed=164, + ) + model = vmc.model + rows = torch.cat((vmc.configs, vmc.configs), dim=0) + model.clear_boundary_cache() + + with torch.no_grad(): + first = model(rows) + first_stats = dict(model.last_amplitude_cache_stats) + second = model(rows) + second_stats = dict(model.last_amplitude_cache_stats) + + assert torch.allclose(first, second) + assert first_stats == { + "num_requests": rows.shape[0], + "num_unique_requests": 1, + "num_hits": 0, + "num_misses": 1, + } + assert second_stats == { + "num_requests": rows.shape[0], + "num_unique_requests": 1, + "num_hits": 1, + "num_misses": 0, + } + + with torch.no_grad(): + model.params[0].add_(0.01) + refreshed = model(rows) + refreshed_stats = dict(model.last_amplitude_cache_stats) + direct = TorchPEPSAmplitude.forward(model, rows) + + assert refreshed_stats["num_hits"] == 0 + assert refreshed_stats["num_misses"] == 1 + assert torch.allclose(refreshed, direct) + + +def test_metropolis_exchange_sweep_uses_cached_proposal_hook(): + class ProposalAwareAmplitude: + def __init__(self): + self.calls = 0 + + def __call__(self, configs): + return torch.ones(configs.shape[0], dtype=torch.float64) + + def proposal_amplitudes( + self, + parent_configs, + target_configs, + current_amplitudes, + *, + chunk_size=None, + ): + del parent_configs, target_configs, chunk_size + self.calls += 1 + return current_amplitudes + + amplitude = ProposalAwareAmplitude() + result = metropolis_exchange_sweep( + torch.tensor([[0, 1], [1, 0]]), + amplitude, + [(0, 1)], + proposal="spin", + generator=torch.Generator().manual_seed(163), + ) + + assert amplitude.calls == 1 + assert result.n_accepted == result.n_proposed == 2 + + +def test_local_energy_reuses_diagonal_connections_and_chunks_offdiagonal(): + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + amps = torch.ones(2, dtype=torch.float64) + conn = heisenberg_connections(configs, [(0, 1)], J=1.0) + amplitude = CountingAmplitude() + + energy = local_energy_from_connections( + configs, + amps, + conn, + amplitude, + chunk_size=1, + ) + + assert torch.allclose(energy, torch.tensor([0.25, 0.25], dtype=torch.float64)) + assert [tuple(call.shape) for call in amplitude.calls] == [(1, 2), (1, 2)] + + +def test_local_energy_coalesces_connections_and_batches_unique_targets(): + configs = torch.tensor([[0], [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.tensor([1.0, 2.0, 4.0, 5.0], dtype=torch.float64), + batch_ids=torch.tensor([0, 0, 1, 1], dtype=torch.long), + ) + + class ConnectionAwareAmplitude: + def __init__(self): + self.connection_sizes = [] + + def __call__(self, rows): + return torch.ones(rows.shape[0], dtype=torch.float64) + + def connected_amplitudes( + self, + _configs, + _amplitudes, + connected, + *, + chunk_size=None, + reuse_diagonal=True, + ): + del chunk_size, reuse_diagonal + self.connection_sizes.append(int(connected.configs.shape[0])) + return torch.ones(connected.configs.shape[0], dtype=torch.float64) + + amplitude = ConnectionAwareAmplitude() + energy = local_energy_from_connections( + configs, + amplitudes, + connections, + amplitude, + ) + + assert torch.allclose(energy, torch.tensor([3.0, 9.0], dtype=torch.float64)) + assert amplitude.connection_sizes == [3] + + # The two off-diagonal rows target the same configuration, so the + # ordinary amplitude path should contract that target only once. + counting = CountingAmplitude() + local_energy_from_connections( + configs, + amplitudes, + connections, + counting, + ) + assert [tuple(call.shape) for call in counting.calls] == [(1, 1)] + + +def test_local_energy_handles_complex_coefficients_with_real_amplitudes(): + configs = torch.tensor([[0], [1]], dtype=torch.long) + amps = torch.ones(2, dtype=torch.float64) + amplitude = CountingAmplitude() + + exactly_real = TorchConnections( + configs=configs.clone(), + coeffs=torch.tensor([1.0 + 0.0j, -2.0 + 0.0j], dtype=torch.complex128), + batch_ids=torch.tensor([0, 1]), + ) + with warnings.catch_warnings(): + warnings.simplefilter("error") + energy = local_energy_from_connections( + configs, + amps, + exactly_real, + amplitude, + ) + assert energy.dtype == torch.float64 + assert torch.allclose(energy, torch.tensor([1.0, -2.0], dtype=torch.float64)) + + complex_phase = TorchConnections( + configs=configs.clone(), + coeffs=torch.tensor([1.0 + 0.5j, -2.0 - 0.25j], dtype=torch.complex128), + batch_ids=torch.tensor([0, 1]), + ) + energy = local_energy_from_connections(configs, amps, complex_phase, amplitude) + assert energy.dtype == torch.complex128 + assert torch.allclose( + energy, + torch.tensor([1.0 + 0.5j, -2.0 - 0.25j], dtype=torch.complex128), + ) + + +def test_torch_vmc_driver_runs_sampling_and_energy_estimate(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + connection_fn="heisenberg", + proposal="spin", + chunk_size=1, + generator=torch.Generator().manual_seed(3), + ) + + result = driver.step() + + assert isinstance(result, TorchVMCStepResult) + assert result.configs.shape == configs.shape + assert result.amplitudes.shape == (2,) + assert result.local_energies.shape == (2,) + assert result.n_proposed == 2 + assert result.n_accepted == 2 + assert result.acceptance_rate == 1.0 + assert torch.isfinite(result.energy_mean) + assert torch.isfinite(result.energy_variance) + + +def test_torch_vmc_driver_can_disable_stable_log_sampling(): + class CountingLogProductAmplitude(ProductAmplitude): + def __init__(self): + super().__init__() + self.log_calls = 0 + + def forward_log(self, configs): + self.log_calls += 1 + amplitudes = self(configs) + return ( + torch.ones_like(amplitudes, dtype=torch.complex128), + amplitudes.abs().log(), + ) + + model = CountingLogProductAmplitude() + driver = TorchVMCDriver( + model, + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + connection_fn="heisenberg", + proposal="spin", + log_amplitude_fn=False, + generator=torch.Generator().manual_seed(41), + ) + + result = driver.sample_sweep() + + assert result.log_abs_amplitudes is None + assert driver.log_amplitude_fn is None + assert model.log_calls == 0 + + +def test_torch_vmc_driver_can_apply_sr_update(): + model = ProductAmplitude() + configs = torch.tensor([[0, 0], [0, 1]], dtype=torch.long) + before = model.weights.detach().clone() + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + connection_fn="transverse_ising", + connection_kwargs={"J": 1.0, "h": 1.0}, + proposal="spin", + chunk_size=1, + generator=torch.Generator().manual_seed(4), + ) + + result = driver.step(sr=True, learning_rate=0.05, sr_diag_shift=1.0e-2) + + assert result.sr is not None + assert result.sr.direction.shape == (2,) + assert not torch.allclose(model.weights.detach(), before) + assert torch.allclose(driver.amplitudes, model(driver.configs)) + + +def test_torch_vmc_measures_and_optimizes_supplied_importance_batch(): + model = ProductAmplitude() + driver = TorchVMCDriver( + model, + [(0, 1)], + torch.tensor([[0, 0], [0, 1]], dtype=torch.long), + connection_fn="transverse_ising", + connection_kwargs={"J": 1.0, "h": 1.0}, + proposal="spin", + generator=torch.Generator().manual_seed(46), + ) + supplied = torch.tensor( + [[0, 0], [0, 1], [1, 0], [1, 1]], + dtype=torch.long, + ) + proposal_log_probs = torch.log( + torch.tensor([0.10, 0.20, 0.30, 0.40], dtype=torch.float64) + ) + + estimate = driver.measure_samples( + supplied, + proposal_log_probs=proposal_log_probs, + ) + expected_weights = ( + model(supplied).abs().square() / proposal_log_probs.exp() + ) + expected_weights = expected_weights / expected_weights.sum() + assert torch.allclose(estimate.importance_weights.reshape(-1), expected_weights) + assert estimate.chain_diagnostics is None + assert torch.allclose( + estimate.energy_mean, + (expected_weights * estimate.local_energies.reshape(-1)).sum(), + ) + + before = model.weights.detach().clone() + current_walkers = driver.configs.detach().clone() + result = driver.step( + samples=supplied, + proposal_log_probs=proposal_log_probs, + sr=True, + learning_rate=0.05, + sr_diag_shift=1.0e-2, + ) + assert result.sample_source == "provided" + assert result.acceptance_rate == 0.0 + assert result.sr is not None + assert torch.allclose(result.importance_weights, expected_weights) + assert float(result.effective_sample_size) == pytest.approx( + float(1.0 / expected_weights.square().sum()) + ) + assert torch.equal(driver.configs, current_walkers) + assert not torch.allclose(model.weights.detach(), before) + + +def test_torch_vmc_sr_optimization_updates_from_supplied_weighted_batch(): + model = ProductAmplitude() + driver = TorchVMCDriver( + model, + [(0, 1)], + torch.tensor([[0, 0], [0, 1]], dtype=torch.long), + connection_fn="transverse_ising", + connection_kwargs={"J": 1.0, "h": 1.0}, + proposal="spin", + ) + before = model.weights.detach().clone() + history = driver.optimize( + optimization=OptimizationConfig(method="sr", n_steps=1, learning_rate=0.01), + samples=torch.tensor([[0, 0], [0, 1], [1, 0]], dtype=torch.long), + weights=torch.tensor([0.2, 0.3, 0.5], dtype=torch.float64), + ) + + assert history[0].sr is not None + assert not torch.allclose(model.weights.detach(), before) + + +def test_torch_vmc_driver_tracks_sr_schedule_and_spring_state(): + model = ProductAmplitude() + driver = TorchVMCDriver( + model, + [(0, 1)], + torch.tensor([[0, 0], [0, 1]], dtype=torch.long), + connection_fn="transverse_ising", + connection_kwargs={"J": 1.0, "h": 1.0}, + proposal="spin", + generator=torch.Generator().manual_seed(44), + ) + steps = [] + + def shift_schedule(step): + steps.append(step) + return 1.0e-2 + + first = driver.step( + sr=True, + learning_rate=1.0e-3, + sr_diag_shift=shift_schedule, + sr_momentum=0.5, + ) + second = driver.step( + sr=True, + learning_rate=1.0e-3, + sr_diag_shift=shift_schedule, + sr_momentum=0.5, + ) + + assert steps == [0, 1] + assert driver.sr_step == 2 + assert first.sr.info["momentum"] == 0.5 + assert second.sr.info["momentum"] == 0.5 + driver.reset_sr_state() + assert driver.sr_step == 0 + + +def test_torch_vmc_compile_safe_kernels_preserve_proposals_and_estimators(): + encoding = FermionSiteEncoding.vmc_torch() + configs = torch.tensor( + [ + [encoding.empty, encoding.double], + [encoding.up, encoding.down], + [encoding.up, encoding.empty], + ], + dtype=torch.long, + ) + before_up, before_down = count_spinful_particles(configs, encoding=encoding) + proposed, changed = propose_spinful_exchange_or_hopping( + 0, + 1, + configs, + hopping_rate=0.5, + encoding=encoding, + generator=torch.Generator().manual_seed(45), + compile_kernels=True, + ) + after_up, after_down = count_spinful_particles(proposed, encoding=encoding) + + assert torch.equal(after_up, before_up) + assert torch.equal(after_down, before_down) + assert torch.equal((proposed != configs).any(dim=1), changed) + + connections = heisenberg_connections( + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + [(0, 1)], + ) + configurations = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + amplitudes = torch.ones(2, dtype=torch.float64) + eager = local_energy_from_connections( + configurations, + amplitudes, + connections, + CountingAmplitude(), + ) + compiled = local_energy_from_connections( + configurations, + amplitudes, + connections, + CountingAmplitude(), + compile_kernels=True, + ) + assert torch.allclose(compiled, eager) + + driver = TorchVMCDriver( + ProductAmplitude(), + [(0, 1)], + configurations, + connection_fn="heisenberg", + proposal="spin", + compile_kernels=True, + generator=torch.Generator().manual_seed(46), + ) + result = driver.step() + assert driver.compile_kernels is True + assert torch.isfinite(result.energy_mean) + + +def test_torch_hamiltonian_connections_accept_explicit_local_terms(): + configs = torch.tensor([[0], [1]], dtype=torch.long) + term = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64) + + conn = torch_hamiltonian_connections(configs, {0: term}) + rows = [ + (tuple(row.tolist()), float(coeff), int(batch_id)) + for row, coeff, batch_id in zip( + conn.configs, + conn.coeffs, + conn.batch_ids, + ) + ] + + assert ((0,), 1.0, 0) in rows + assert ((1,), 3.0, 0) in rows + assert ((0,), 2.0, 1) in rows + assert ((1,), 4.0, 1) in rows + + +def test_torch_hamiltonian_connections_preserve_native_fermion_grading(): + fermion = Fermion(symmetry="U1U1", spinful=True, t=1.0, U=8.0) + sites = ((0, 0), (0, 1), (1, 0), (1, 1)) + edges = ( + ((0, 0), (0, 1)), + ((0, 0), (1, 0)), + ((0, 1), (1, 1)), + ((1, 0), (1, 1)), + ) + configs = torch.tensor( + list(product(range(4), repeat=4)), + dtype=torch.long, + ) + + generic = torch_hamiltonian_connections( + configs, + fermion.hamiltonian(edges), + site_order=sites, + ) + specialized = spinful_fermi_hubbard_connections( + configs, + ((0, 1), (0, 2), (1, 3), (2, 3)), + t=1.0, + U=8.0, + encoding=FermionSiteEncoding.vmc_torch(), + mode_order="down-up", + ) + + def as_map(connections): + values = {} + for eta, coeff, batch_id in zip( + connections.configs.tolist(), + connections.coeffs.tolist(), + connections.batch_ids.tolist(), + ): + key = (int(batch_id), tuple(eta)) + values[key] = values.get(key, 0.0) + coeff + return values + + left = as_map(generic) + right = as_map(specialized) + assert set(left) == set(right) + assert all(abs(left[key] - right[key]) < 1.0e-12 for key in left) + + +@pytest.mark.parametrize("spinful", [False, True]) +def test_torch_hamiltonian_connections_support_flat_native_fermion_terms(spinful): + fermion = Fermion(symmetry="Z2", spinful=spinful, t=1.0, U=8.0) + sparse = fermion.hamiltonian(((0, 2),), flat=False).terms[(0, 2)] + flat = fermion.hamiltonian(((0, 2),), flat=True).terms[(0, 2)] + dimension = 4 if spinful else 2 + configs = torch.tensor( + list(product(range(dimension), repeat=3)), + dtype=torch.long, + ) + + def as_map(operator): + connections = torch_hamiltonian_connections( + configs, + {(0, 2): operator}, + site_order=(0, 1, 2), + ) + values = {} + for eta, coefficient, batch_id in zip( + connections.configs.tolist(), + connections.coeffs.tolist(), + connections.batch_ids.tolist(), + ): + key = (int(batch_id), tuple(eta)) + values[key] = values.get(key, 0.0) + coefficient + return values + + assert as_map(flat) == pytest.approx(as_map(sparse)) + + +def test_torch_hamiltonian_connections_cache_native_compilation(monkeypatch): + import pepsy.vmc.torch as vmc_torch + + fermion = Fermion(symmetry="U1U1", spinful=True, t=1.0, U=8.0) + operator = fermion.hopping_operator() + configs = torch.tensor( + list(product(range(4), repeat=3)), + dtype=torch.long, + ) + vmc_torch._FERMION_COMPILED_TERM_CACHE.clear() + original = vmc_torch._operator_dense_numpy + calls = [] + + def counted(value): + calls.append(value) + return original(value) + + monkeypatch.setattr(vmc_torch, "_operator_dense_numpy", counted) + first = torch_hamiltonian_connections( + configs, + {(0, 2): operator}, + site_order=(0, 1, 2), + ) + second = torch_hamiltonian_connections( + configs, + {(0, 2): operator}, + site_order=(0, 1, 2), + ) + assert len(calls) == 1 + assert torch.equal(first.configs, second.configs) + assert torch.equal(first.batch_ids, second.batch_ids) + assert torch.equal(first.coeffs, second.coeffs) + + # The cutoff changes the compiled sparsity pattern and therefore gets a + # separate cache entry rather than reusing an incompatible table. + torch_hamiltonian_connections( + configs, + {(0, 2): operator}, + site_order=(0, 1, 2), + coefficient_cutoff=1.0e-12, + ) + assert len(calls) == 2 + + +def test_torch_hamiltonian_connections_reject_invalid_native_fermion_sites(): + fermion = Fermion(symmetry="U1", spinful=False) + operator = fermion.hopping_operator() + configs = torch.zeros((2, 3), dtype=torch.long) + + with pytest.raises(ValueError, match="distinct"): + torch_hamiltonian_connections( + configs, + {(1, 1): operator}, + site_order=(0, 1, 2), + ) + with pytest.raises(ValueError, match="non-negative"): + torch_hamiltonian_connections( + torch.tensor([[-1, 0, 0]], dtype=torch.long), + {(0, 2): operator}, + site_order=(0, 1, 2), + ) + + +def test_torch_hamiltonian_connections_handle_odd_local_jw_prefix(): + fermion = Fermion(symmetry="U1", spinful=False) + create = fermion.operator("create") + configs = torch.tensor([[0, 0], [1, 0]], dtype=torch.long) + + connections = torch_hamiltonian_connections( + configs, + {1: create}, + site_order=(0, 1), + ) + values = { + int(batch_id): complex(coefficient) + for coefficient, batch_id in zip(connections.coeffs, connections.batch_ids) + } + assert values == {0: 1.0 + 0.0j, 1: -1.0 + 0.0j} + + +def test_torch_hamiltonian_connections_respect_reversed_native_term_order(): + fermion = Fermion(symmetry="U1", spinful=False) + forward = fermion.operator_term( + [(1.0, ((0, "create"), (1, "annihilate")))], + sites=(0, 1), + ) + reverse = fermion.operator_term( + [(1.0, ((2, "create"), (0, "annihilate")))], + sites=(0, 2), + ) + configs = torch.tensor(list(product(range(2), repeat=3)), dtype=torch.long) + + def as_map(operator, where): + connections = torch_hamiltonian_connections( + configs, + {where: operator}, + site_order=(0, 1, 2), + ) + return { + (int(batch_id), tuple(eta)): complex(coefficient) + for eta, coefficient, batch_id in zip( + connections.configs.tolist(), + connections.coeffs, + connections.batch_ids, + ) + } + + assert as_map(forward, (2, 0)) == as_map(reverse, (0, 2)) + + +def test_torch_fermion_vmc_supports_spinless_terms_and_sampling(): + fermion = Fermion(symmetry="U1", spinful=False, t=1.0, V=2.0) + peps = ps_to_peps( + 1, + 2, + fermion=fermion, + occupations=(1, 0), + dtype="complex128", + ) + vmc = TorchFermionVMC( + peps, + fermion, + n_walkers=2, + dtype=torch.complex128, + seed=19, + ) + + assert isinstance(vmc.model, TorchPEPSBoundaryAmplitude) + assert vmc.model.contraction == "boundary" + assert vmc.model.chi == 4 + assert isinstance(vmc.encoding, SpinlessSiteEncoding) + assert vmc.encoding == SpinlessSiteEncoding() + assert vmc.sector == 1 + assert vmc.proposal == "spin" + assert torch.isfinite(vmc.local_energies()).all() + result = vmc.sample_sweep() + assert torch.equal( + vmc.encoding.decode(result.configs).sum(dim=-1), + torch.ones(2, dtype=torch.long), + ) + + +def test_torch_vmc_driver_accepts_explicit_terms_and_estimates_sampling_rate(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + term = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: term}, + proposal="spin", + generator=torch.Generator().manual_seed(5), + ) + + result = driver.estimate_observable( + burn_in=1, + n_measurements=2, + sweeps_between=1, + ) + + assert isinstance(result, TorchVMCEnergyEstimate) + assert result.n_samples == 4 + assert result.n_measurements == 2 + assert result.elapsed_seconds > 0.0 + assert result.samples_per_second > 0.0 + assert torch.isfinite(result.energy_mean) + assert torch.isfinite(result.energy_stderr) + assert result.energy_stderr_naive is not None + assert result.effective_sample_size is not None + assert result.chain_diagnostics is not None + assert result.chain_diagnostics.n_chains == 2 + + legacy = driver.estimate_energy( + burn_in=0, + n_measurements=1, + sweeps_between=1, + ) + assert isinstance(legacy, TorchVMCEnergyEstimate) + + +def test_torch_vmc_driver_shares_connected_amplitudes_across_observables(): + model = CountingAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: flip}, + proposal="spin", + ) + model.calls.clear() + + values = driver.local_observables({ + "one": {0: flip}, + "two": {0: 2.0 * flip}, + }) + + assert len(model.calls) == 1 + assert torch.allclose(values["two"], 2.0 * values["one"]) + assert torch.allclose(values["one"], torch.ones(2, dtype=torch.float64)) + + +def test_torch_vmc_driver_estimates_observables_with_shared_samples_and_profile(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: flip}, + proposal="spin", + ) + + results = driver.estimate_observables( + {"one": {0: flip}, "two": {0: 2.0 * flip}}, + n_samples=4, + n_discard=0, + n_thin=1, + seed=152, + profile=True, + ) + + assert set(results) == {"one", "two"} + assert results["one"].configs.shape == (2, 2, 2) + assert torch.allclose( + results["two"].local_energies, + 2.0 * results["one"].local_energies, + ) + profile = results["one"].profile + assert profile["shared_observables"] == ("one", "two") + assert profile["sampling_seconds"] >= 0.0 + assert profile["local_estimator_seconds"] >= 0.0 + + +def test_torch_vmc_driver_measures_saved_samples_without_sampling(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: flip}, + proposal="spin", + ) + samples = driver.sample( + n_samples=4, + n_chains=2, + n_discard=0, + n_thin=1, + seed=153, + ) + + result = driver.measure_samples(samples, profile=True) + external = driver.measure_samples(samples.configs, amplitudes=None) + expected = driver.local_observables( + {"energy": {0: flip}}, + configs=samples.configs.reshape(-1, 2), + amplitudes=samples.amplitudes.reshape(-1), + )["energy"].reshape(2, 2) + + assert isinstance(result, TorchVMCEnergyEstimate) + assert result.configs.shape == (2, 2, 2) + assert result.local_energies.shape == (2, 2) + assert result.n_samples == 4 + assert result.n_measurements == 2 + assert result.chain_diagnostics is not None + assert torch.allclose(result.local_energies, expected) + assert torch.allclose(external.local_energies, result.local_energies) + assert result.profile["samples_only"] is True + assert result.profile["sampling_seconds"] == 0.0 + + +def test_torch_vmc_driver_measures_multiple_saved_observables_shared_targets(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: flip}, + proposal="spin", + ) + samples = driver.sample( + n_samples=4, + n_chains=2, + n_discard=0, + n_thin=1, + seed=154, + ) + + results = driver.measure_samples( + samples, + observables={"one": {0: flip}, "two": {0: 2.0 * flip}}, + ) + + assert set(results) == {"one", "two"} + assert results["one"].chain_diagnostics is not None + assert torch.allclose( + results["two"].local_energies, + 2.0 * results["one"].local_energies, + ) + + +def test_torch_vmc_driver_deduplicates_saved_parent_configurations(): + model = CountingAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, [(0, 1)], configs, + terms={0: flip}, + proposal="spin", + ) + saved_configs = torch.tensor( + [ + [[0, 1], [0, 1]], + [[1, 0], [1, 0]], + [[0, 1], [0, 1]], + ], + dtype=torch.long, + ) + + model.calls.clear() + result = driver.measure_samples(saved_configs, profile=True) + deduplicated_first_batch = model.calls[0] + + model.calls.clear() + driver.measure_samples(saved_configs, deduplicate=False) + non_deduplicated_first_batch = model.calls[0] + + assert result.profile["num_samples"] == 6 + assert result.profile["num_unique_samples"] == 2 + assert deduplicated_first_batch.shape[0] == 2 + assert non_deduplicated_first_batch.shape[0] == 6 + + +def test_torch_fermion_vmc_boundary_measurement_dedup_matches_serial(): + fermion = Fermion(symmetry="U1", spinful=False, t=1.0, V=2.0) + peps = ps_to_peps( + 1, + 2, + fermion=fermion, + occupations=(1, 0), + dtype="complex128", + ) + vmc = TorchFermionVMC( + peps, + fermion, + n_walkers=2, + dtype=torch.complex128, + seed=155, + ) + saved_configs = torch.stack((vmc.configs, vmc.configs), dim=0) + + deduplicated = vmc.measure_samples(saved_configs, deduplicate=True) + serial = vmc.measure_samples(saved_configs, deduplicate=False) + + assert vmc.model.last_amplitude_batching == "serial" + assert torch.allclose( + deduplicated.local_energies, + serial.local_energies, + atol=1.0e-10, + rtol=1.0e-10, + ) + + +def test_torch_vmc_step_profile_reports_phase_timings(): + model = ProductAmplitude() + driver = TorchVMCDriver( + model, + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + terms={0: torch.eye(2, dtype=torch.float64)}, + proposal="spin", + ) + + result = driver.step(profile=True) + + assert isinstance(result, TorchVMCStepResult) + assert result.profile["sampling_seconds"] >= 0.0 + assert result.profile["connection_seconds"] >= 0.0 + assert result.profile["local_estimator_seconds"] >= 0.0 + assert result.profile["total_seconds"] > 0.0 + + +def test_torch_vmc_driver_has_progress_aware_burnin_and_optimization(monkeypatch): + class RecordingProgress: + def __init__(self, total, desc, unit): + self.total = total + self.desc = desc + self.unit = unit + self.updates = 0 + self.postfixes = [] + self.closed = False + + def update(self, amount): + self.updates += amount + + def set_postfix(self, values): + self.postfixes.append(dict(values)) + + def close(self): + self.closed = True + + bars = [] + + def make_progress(progress, *, total, desc, unit=None): + if not progress: + return None + bar = RecordingProgress(total, desc, unit) + bars.append(bar) + return bar + + monkeypatch.setattr(torch_vmc, "_make_progress", make_progress) + driver = TorchVMCDriver( + ProductAmplitude(), + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + connection_fn="heisenberg", + proposal="spin", + ) + + burn_in = driver.burn_in( + 3, + progress=True, + track_proposal_stats=True, + ) + history = driver.optimize( + 2, + progress=True, + track_proposal_stats=True, + ) + compatibility_history = driver.run(1) + + assert burn_in.n_proposed == burn_in.n_accepted == 6 + assert burn_in.proposal_stats["exchange"]["selected"] == 6 + assert len(history) == 2 + assert len(compatibility_history) == 1 + assert all(isinstance(result, TorchVMCStepResult) for result in history) + assert [(bar.total, bar.unit, bar.updates, bar.closed) for bar in bars] == [ + (3, "sweep", 3, True), + (2, "step", 2, True), + ] + assert bars[-1].postfixes[-1]["E/site"] + assert bars[-1].postfixes[-1]["accept"] == "1.000" + assert bars[-1].postfixes[-1]["no-op"] == "0.000" + + +def test_torch_vmc_sample_sweep_aggregates_multiple_sweeps(): + driver = TorchVMCDriver( + ProductAmplitude(), + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), connection_fn="heisenberg", proposal="spin", - chunk_size=1, - generator=torch.Generator().manual_seed(3), ) - result = driver.step() + result = driver.sample_sweep(n_sweeps=3, track_proposal_stats=True) + + assert result.n_proposed == result.n_accepted == 6 + assert result.proposal_stats["exchange"]["selected"] == 6 + + +def test_torch_metropolis_sampler_preserves_chains_and_accepts_netket_aliases(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + sampler = TorchMetropolisSampler( + model, + [(0, 1)], + configs, + proposal="spin", + seed=11, + ) + + result = sampler.sample( + n_samples=5, + n_discard=2, + n_thin=2, + ) + + assert isinstance(result, TorchMCMCSamples) + assert result.configs.shape == (3, 2, 2) + assert result.amplitudes.shape == (3, 2) + assert result.n_samples == 6 + assert result.n_samples_per_chain == 3 + assert result.n_chains == 2 + assert result.n_discard_per_chain == 2 + assert result.sweep_size == 2 + assert result.n_proposed >= result.n_accepted >= 0 + assert result.elapsed_seconds > 0.0 + + +def test_torch_metropolis_sampler_drops_failed_optional_log_cache(): + class FlakyLogAmplitude: + def __init__(self): + self.log_calls = 0 + + def __call__(self, configs): + return torch.ones(configs.shape[0], dtype=torch.float64) + + def forward_log(self, configs): + self.log_calls += 1 + if self.log_calls > 1: + raise RuntimeError("unsupported proposed sector") + return ( + torch.ones(configs.shape[0], dtype=torch.complex128), + torch.zeros(configs.shape[0], dtype=torch.float64), + ) + + sampler = TorchMetropolisSampler( + FlakyLogAmplitude(), + [(0, 1)], + torch.tensor([[0, 1], [1, 0]]), + proposal="spin", + seed=164, + ) + result = sampler.sample(n_samples=2, n_discard=0, n_thin=1) + + assert result.log_abs_amplitudes is None + assert result.amplitudes.shape == (1, 2) + + +def test_metropolis_exchange_sweep_uses_log_amplitude_ratio_when_available(): + result = metropolis_exchange_sweep( + torch.tensor([[0, 1]], dtype=torch.long), + LogOnlyAmplitude(), + [(0, 1)], + proposal="spin", + generator=torch.Generator().manual_seed(0), + ) + + assert result.configs.tolist() == [[1, 0]] + assert result.log_abs_amplitudes.tolist() == [-900.0] + assert result.nonzero_amplitudes.tolist() == [True] + + +def test_torch_chain_diagnostics_reports_rhat_tau_and_effective_sample_size(): + values = torch.arange(8, dtype=torch.float64).reshape(8, 1).repeat(1, 4) + diagnostics = torch_chain_diagnostics(values) + + assert isinstance(diagnostics, TorchChainDiagnostics) + assert diagnostics.r_hat >= 1.0 + assert diagnostics.integrated_autocorrelation_time >= 1.0 + assert 0.0 < diagnostics.effective_sample_size <= 32.0 + assert diagnostics.rhat == diagnostics.r_hat + assert diagnostics.tau == diagnostics.integrated_autocorrelation_time + + +def test_torch_observable_statistics_uses_effective_sample_size_for_stderr(): + values = torch.arange(8, dtype=torch.float64).reshape(8, 1).repeat(1, 4) + ( + _, + variance, + autocorrelation_stderr, + naive_stderr, + effective_sample_size, + diagnostics, + ) = _observable_statistics(values) + + assert diagnostics is not None + assert effective_sample_size <= values.numel() + assert torch.allclose( + autocorrelation_stderr, + torch.sqrt(variance / effective_sample_size), + ) + assert autocorrelation_stderr >= naive_stderr + + +def test_torch_vmc_driver_can_stop_at_target_effective_sample_size(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: flip}, + proposal="spin", + generator=torch.Generator().manual_seed(122), + ) + + result = driver.estimate_observable( + n_measurements=6, + sweeps_between=1, + target_effective_sample_size=1.0, + min_measurements=2, + rhat_threshold=None, + auto_thin=True, + profile=True, + ) + + assert result.n_measurements == 2 + assert result.effective_sample_size >= 1.0 + assert result.profile["adaptive_sampling"]["stop_reason"] == ( + "target_effective_sample_size" + ) + assert result.profile["adaptive_sampling"]["measurements_collected"] == 2 + + +def test_torch_vmc_shared_observables_stop_when_all_reach_target_ess(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + flip = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float64) + driver = TorchVMCDriver( + model, + [(0, 1)], + configs, + terms={0: flip}, + proposal="spin", + generator=torch.Generator().manual_seed(123), + ) + + results = driver.estimate_observables( + {"one": {0: flip}, "two": {0: 2.0 * flip}}, + n_measurements=6, + sweeps_between=1, + target_effective_sample_size=1.0, + min_measurements=2, + rhat_threshold=None, + auto_thin=True, + profile=True, + ) + + assert {result.n_measurements for result in results.values()} == {2} + assert all(result.effective_sample_size >= 1.0 for result in results.values()) + assert all( + result.profile["adaptive_sampling"]["stop_reason"] + == "target_effective_sample_size" + for result in results.values() + ) + + +def test_metropolis_local_sampler_infers_sites_and_chain_count(): + model = ProductAmplitude() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + + result = metropolis_local_sampler( + configs, + model, + [(0, 1)], + n_sites=2, + n_samples=4, + n_chains=2, + n_discard_per_chain=0, + sweep_size=1, + proposal="spin", + seed=12, + ) + + assert result.configs.shape == (2, 2, 2) + assert result.n_samples == 4 + + +def test_torch_vmc_driver_exposes_chain_preserving_sampling(): + model = ProductAmplitude() + driver = TorchVMCDriver( + model, + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + terms={0: torch.eye(2, dtype=torch.float64)}, + proposal="spin", + ) + + result = driver.sample( + n_samples=3, + n_discard=1, + n_thin=1, + seed=13, + ) - assert isinstance(result, TorchVMCStepResult) - assert result.configs.shape == configs.shape - assert result.amplitudes.shape == (2,) - assert result.local_energies.shape == (2,) - assert result.n_proposed == 2 - assert result.n_accepted == 2 - assert result.acceptance_rate == 1.0 - assert torch.isfinite(result.energy_mean) - assert torch.isfinite(result.energy_variance) + assert result.configs.shape == (2, 2, 2) + assert driver.configs.shape == (2, 2) + assert torch.allclose(driver.amplitudes, model(driver.configs)) + + estimate = driver.estimate_observable( + n_samples=3, + n_discard=1, + n_thin=1, + seed=14, + ) + assert estimate.configs.shape == (2, 2, 2) + assert estimate.local_energies.shape == (2, 2) + assert estimate.n_samples == 4 + assert torch.isfinite(estimate.energy_mean) + + with pytest.raises(ValueError, match="target_effective_sample_size"): + driver.estimate_observable( + n_samples=3, + target_effective_sample_size=2, + ) -def test_torch_vmc_driver_can_apply_sr_update(): +def test_torch_vmc_driver_importance_estimate_uses_proposal_weights(): model = ProductAmplitude() - configs = torch.tensor([[0, 0], [0, 1]], dtype=torch.long) - before = model.weights.detach().clone() + configs = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) driver = TorchVMCDriver( model, [(0, 1)], configs, - connection_fn="transverse_ising", - connection_kwargs={"J": 1.0, "h": 1.0}, + connection_fn="heisenberg", proposal="spin", - chunk_size=1, - generator=torch.Generator().manual_seed(4), ) - result = driver.step(sr=True, learning_rate=0.05, sr_diag_shift=1.0e-2) + class Proposal: + def sample(self, *, samples, progbar=False): + assert samples == 2 + assert progbar is False + return SimpleNamespace( + configs=[[0, 1], [1, 0]], + omegas=([1.0, 1.0], [0, 0]), + ) + + result = driver.importance_energy_estimate(Proposal(), n_samples=2) + + assert isinstance(result, TorchVMCImportanceEstimate) + assert result.n_samples == 2 + assert result.n_valid == 2 + assert torch.allclose( + result.effective_sample_size, + torch.tensor(2.0, dtype=torch.float64), + ) + assert torch.isfinite(result.energy_mean) + assert torch.isfinite(result.energy_stderr) - assert result.sr is not None - assert result.sr.direction.shape == (2,) - assert not torch.allclose(model.weights.detach(), before) - assert torch.allclose(driver.amplitudes, model(driver.configs)) + +@pytest.mark.parametrize( + ("symmetry", "encoding"), + [ + ("Z2", FermionSiteEncoding.symmray()), + ("U1", FermionSiteEncoding.vmc_torch()), + ("Z2Z2", FermionSiteEncoding.vmc_torch()), + ("U1U1", FermionSiteEncoding.vmc_torch()), + ], +) +def test_torch_bp_metropolis_filters_fermion_symmetry_sectors(symmetry, encoding): + class Amplitude: + def __call__(self, rows): + return torch.ones(rows.shape[0], dtype=torch.float64) + + if symmetry == "U1": + valid, invalid, sector = [0, 3], [0, 0], 2 + elif symmetry == "U1U1": + valid, invalid, sector = [2, 1], [0, 0], (1, 1) + elif symmetry == "Z2": + valid, invalid, sector = [0, 0], [2, 0], 0 + else: + valid, invalid, sector = [2, 1], [0, 0], (1, 1) + + class Proposal: + def __init__(self): + self.calls = 0 + + def sample(self, *, samples, progbar=False): + self.calls += 1 + rows = [valid] * samples if self.calls == 1 else [invalid] * samples + return SimpleNamespace( + configs=rows, + omegas=([1.0] * samples, [0] * samples), + ) + + sampler = TorchBPMetropolisSampler( + Amplitude(), + [], + Proposal(), + n_chains=2, + symmetry=symmetry, + sector=sector, + encoding=encoding, + seed=7, + ) + initial = sampler.configs.clone() + result = sampler.sample_sweep() + + assert result.n_proposed == 2 + assert result.n_accepted == 0 + assert torch.equal(sampler.configs, initial) + + +def test_torch_bp_metropolis_acceptance_uses_independence_log_ratio(): + class Amplitude: + def __call__(self, rows): + return torch.where( + rows[:, 0] == 0, + torch.ones(rows.shape[0], dtype=torch.float64), + torch.full((rows.shape[0],), 2.0, dtype=torch.float64), + ) + + class Proposal: + def __init__(self): + self.calls = 0 + + def sample(self, *, samples, progbar=False): + self.calls += 1 + rows = [[0]] * samples if self.calls == 1 else [[1]] * samples + # q(0) = 3/4 and q(1) = 1/4, so the move 0 -> 1 has ratio + # |2|^2 * (3/4) / (|1|^2 * (1/4)) = 12 and must accept. + probs = ([0.75] * samples, [0] * samples) + if self.calls > 1: + probs = ([0.25] * samples, [0] * samples) + return SimpleNamespace(configs=rows, omegas=probs) + + sampler = TorchBPMetropolisSampler( + Amplitude(), + [], + Proposal(), + n_chains=2, + seed=3, + ) + result = sampler.sample_sweep() + assert result.n_accepted == 2 + assert sampler.configs.tolist() == [[1], [1]] + + +@pytest.mark.parametrize( + ("symmetry", "encoding"), + [ + ("Z2", FermionSiteEncoding.symmray()), + ("U1", FermionSiteEncoding.vmc_torch()), + ("Z2Z2", FermionSiteEncoding.vmc_torch()), + ("U1U1", FermionSiteEncoding.vmc_torch()), + ], +) +def test_peps_bp_sampler_supports_four_state_symmray_fermions(symmetry, encoding): + peps = _symmray_fermionic_peps(symmetry) + result = PepsBpSampler(peps, encoding=encoding).sample( + samples=1, + method="exact", + seed=2, + bp_kwargs={"max_iterations": 3}, + ) + + assert len(result.configs) == 1 + assert len(result.configs[0]) == 4 + assert all(value in {0, 1, 2, 3} for value in result.configs[0]) + assert np.isfinite(np.asarray(result.omegas[0])).all() + assert np.isfinite(np.asarray(result.omegas[1])).all() + assert np.isfinite(np.asarray(result.ps[0])).all() + assert np.isfinite(np.asarray(result.ps[1])).all() def _symmray_site_charge(symmetry): @@ -433,6 +2741,113 @@ def test_torch_peps_amplitude_supports_symmray_flat_z2_vmap(): assert torch.allclose(vmapped, loop) +def test_torch_peps_amplitude_flat_z2_stable_logs_use_vmap(): + peps = _symmray_fermionic_peps("Z2", flat=True) + model = TorchPEPSAmplitude( + peps, + contraction="exact", + dtype=torch.float64, + amplitude_batching="vmap", + ) + rows = torch.tensor([(0, 0, 0, 0), (1, 1, 1, 1)], dtype=torch.long) + + phase, log_abs = model.forward_log(rows) + assert model.last_amplitude_batching == "log-vmap" + amplitudes = model(rows) + expected_phase = torch.where( + amplitudes.abs() > 0, + amplitudes / amplitudes.abs(), + torch.zeros_like(amplitudes), + ) + + assert model.last_amplitude_batching == "vmap" + assert torch.allclose(phase, expected_phase) + assert torch.allclose(log_abs, amplitudes.abs().log()) + + +@pytest.mark.parametrize("batching", ["auto", "vmap", "serial"]) +def test_torch_peps_amplitude_explicit_batching_modes_preserve_flat_z2_values( + batching, +): + peps = _symmray_fermionic_peps("Z2", flat=True) + rows = torch.tensor([(0, 0, 0, 0), (1, 1, 1, 1)], dtype=torch.long) + model = TorchPEPSAmplitude( + peps, + contraction="exact", + dtype=torch.float64, + amplitude_batching=batching, + ) + + amplitudes = model(rows) + reference = _direct_peps_amplitudes(peps, rows) + + assert torch.allclose(amplitudes, reference) + if batching == "serial": + assert model.last_amplitude_batching == "serial" + else: + assert model.last_amplitude_batching == "vmap" + + +def test_torch_peps_amplitude_graded_torch_projects_u1u1_under_vmap(): + peps = _symmray_fermionic_peps("U1U1") + model = TorchPEPSAmplitude( + peps, + contraction="exact", + dtype=torch.float64, + graded_torch=True, + ) + rows = torch.tensor( + list(product(range(4), repeat=4)), + dtype=torch.long, + ) + + amplitudes = model(rows) + direct = _direct_peps_amplitudes(peps, rows) + + assert torch.allclose(amplitudes, direct, atol=1.0e-10, rtol=1.0e-10) + model.zero_grad() + amplitudes.square().sum().backward() + assert all( + param.grad is not None and torch.isfinite(param.grad).all() + for param in model.parameters() + ) + + +def test_torch_peps_amplitude_graded_torch_can_use_serial_batches(): + peps = _symmray_fermionic_peps("U1U1") + model = TorchPEPSAmplitude( + peps, + contraction="exact", + dtype=torch.float64, + graded_torch=True, + amplitude_batching="serial", + ) + rows = torch.tensor([(2, 0, 1, 3), (0, 2, 3, 1)], dtype=torch.long) + + amplitudes = model(rows) + direct = _direct_peps_amplitudes(peps, rows) + + assert model.last_amplitude_batching == "graded-serial" + assert torch.allclose(amplitudes, direct, atol=1.0e-10, rtol=1.0e-10) + + +def test_torch_peps_amplitude_symmray_ctmrg_uses_direct_boundary_compression(): + peps = _symmray_fermionic_peps("U1U1") + model = TorchPEPSAmplitude( + peps, + contraction="ctmrg", + chi=8, + cutoff=1.0e-10, + dtype=torch.float64, + ) + rows = torch.tensor([(2, 0, 1, 3)], dtype=torch.long) + + amplitudes = model(rows) + + assert model.contraction_opts["mode"] == "direct" + assert torch.isfinite(amplitudes).all() + + def test_symmray_fermionic_peps_feeds_hubbard_local_energy_kernel(): peps = _symmray_fermionic_peps("U1U1") model = TorchPEPSAmplitude(peps, contraction="exact", dtype=torch.float64) @@ -454,6 +2869,440 @@ def test_symmray_fermionic_peps_feeds_hubbard_local_energy_kernel(): assert torch.isfinite(energy).all() +def test_torch_fermion_vmc_infers_geometry_sector_encoding_and_terms(): + fermion = Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=2.0, + ) + peps = ps_to_peps( + 2, + 2, + fermion=fermion, + occupations=[(1, 0), (0, 1), (1, 0), (0, 1)], + seed=21, + dtype="complex128", + ) + vmc = TorchFermionVMC( + peps, + fermion, + n_walkers=2, + dtype=torch.complex128, + seed=22, + ) + + assert isinstance(vmc.metadata, TorchFermionVMCMetadata) + assert (vmc.Lx, vmc.Ly) == (2, 2) + assert vmc.site_order == tuple(peps.sites) + assert vmc.sector == (2, 2) + assert vmc.encoding == FermionSiteEncoding.vmc_torch() + assert vmc.metadata.graph_edges == ((0, 1), (2, 3), (0, 2), (1, 3)) + assert torch.isfinite(vmc.local_energies()).all() + + sample = vmc.sample_sweep() + n_up, n_down = count_spinful_particles(sample.configs, encoding=vmc.encoding) + assert n_up.tolist() == [2, 2] + assert n_down.tolist() == [2, 2] + + +def test_torch_fermion_vmc_keeps_hamiltonian_and_observables_separate(): + fermion = Fermion(spinful=True, symmetry="U1U1", t=1.0, U=2.0) + peps = ps_to_peps( + 2, + 2, + fermion=fermion, + occupations=[(1, 0), (0, 1), (1, 0), (0, 1)], + seed=216, + dtype="complex128", + ) + hamiltonian = fermion.hamiltonian( + ( + ((0, 0), (0, 1)), + ((0, 0), (1, 0)), + ((0, 1), (1, 1)), + ((1, 0), (1, 1)), + ) + ) + observables = { + "density": {(0, 0): fermion.observable("number")}, + } + vmc = TorchFermionVMC( + peps, + fermion=fermion, + hamiltonian=hamiltonian, + observables=observables, + n_walkers=2, + contraction="exact", + dtype=torch.complex128, + seed=217, + ) + + assert vmc.hamiltonian is hamiltonian + assert set(vmc.observables) == {"density"} + + problem = VMCProblem( + peps=peps, + hamiltonian=hamiltonian, + observables=observables, + symmetry="U1U1", + ) + setup = build_torch_vmc( + problem, + fermion=fermion, + contraction=ContractionConfig(), + sampling=SamplingConfig( + n_samples_per_chain=1, + n_chains=2, + burn_in=0, + seed=218, + ), + dtype=torch.complex128, + ) + samples = setup.sample() + assert isinstance(setup, TorchVMCSetup) + assert samples.chain_shape == (1, 2) + assert samples.native is not None + + measurement = setup.measure() + assert isinstance(measurement, VMCMeasurement) + assert set(measurement.observables) == {"energy", "density"} + history = setup.optimize( + OptimizationConfig(n_steps=1, method="sgd", learning_rate=1.0e-3), + sample_sweeps=1, + ) + assert isinstance(history, VMCOptimizationResult) + assert history.energies.shape == (1,) + assert history.native[0].sr is None + + +@pytest.mark.parametrize("cyclic", [False, True]) +def test_torch_fermion_vmc_infers_peps_pbc_and_long_range_term_graph(cyclic): + fermion = Fermion(spinful=True, symmetry="U1U1") + occupations = { + (x, y): (1, 0) if (x + y) % 2 == 0 else (0, 1) + for x in range(3) + for y in range(3) + } + peps = ps_to_peps( + 3, + 3, + fermion=fermion, + occupations=occupations, + cyclic=cyclic, + seed=29, + dtype="complex128", + ) + encoding = FermionSiteEncoding.vmc_torch() + configs = torch.tensor( + [[ + encoding.up if (x + y) % 2 == 0 else encoding.down + for x in range(3) + for y in range(3) + ]], + dtype=torch.long, + ) + terms = { + ((0, 0), (2, 2)): fermion.hopping_operator(), + } + + vmc = TorchFermionVMC( + peps, + fermion=fermion, + terms=terms, + configs=configs, + contraction="exact", + dtype=torch.complex128, + seed=30, + ) + + assert vmc.metadata.pbc == (cyclic, cyclic) + positions = {site: i for i, site in enumerate(vmc.site_order)} + graph_edges = { + frozenset((left, right)) for left, right in vmc.metadata.graph_edges + } + long_range_edge = frozenset((positions[(0, 0)], positions[(2, 2)])) + assert long_range_edge in graph_edges + wrap_edge = frozenset((positions[(0, 0)], positions[(0, 2)])) + assert (wrap_edge in graph_edges) is cyclic + + +def test_torch_fermion_vmc_4x6_d4_complex_modes_sector_and_signs(): + """Check the complex Hubbard VMC path on a production-sized PEPS. + + The random PEPS is dense so the same physical state can be evaluated by + all three Quimb contraction modes. The native ``Fermion`` object still + supplies the U1U1 Hubbard terms, including their graded hopping signs. + """ + fermion = Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=8.0, + ) + peps = qtn.PEPS.rand( + Lx=4, + Ly=6, + bond_dim=4, + phys_dim=4, + seed=1404, + dtype="complex128", + ) + configs = torch.tensor( + [[ + 2 if (x + y) % 2 == 0 else 1 + for x in range(4) + for y in range(6) + ]], + dtype=torch.long, + ) + boundary_opts = { + "mode": "mps", + "final_contract": True, + "final_contract_opts": {"optimize": "auto-hq"}, + "sequence": ["xmin", "xmax", "ymin", "ymax"], + "equalize_norms": False, + "progbar": False, + } + ctmrg_opts = { + "final_contract": False, + "final_contract_opts": {"optimize": "auto-hq"}, + "max_separation": 1, + "inplace": False, + "equalize_norms": False, + "progbar": False, + } + mode_opts = { + "exact": {"contraction": "exact"}, + "boundary": { + "contraction": "boundary", + "chi": 64, + "cutoff": 1.0e-10, + "contraction_opts": boundary_opts, + }, + "ctmrg": { + "contraction": "ctmrg", + "chi": 64, + "cutoff": 1.0e-10, + "contraction_opts": ctmrg_opts, + }, + } + + vmcs = { + name: TorchFermionVMC( + peps, + fermion, + configs=configs, + n_walkers=1, + dtype=torch.complex128, + **options, + ) + for name, options in mode_opts.items() + } + + exact = vmcs["exact"] + assert exact.metadata.physical_dim == 4 + assert exact.sector == (12, 12) + connections = exact.make_connections() + n_up, n_down = count_spinful_particles( + connections.configs, + encoding=exact.encoding, + ) + assert torch.equal(n_up, torch.full_like(n_up, 12)) + assert torch.equal(n_down, torch.full_like(n_down, 12)) + + hopping_coeffs = connections.coeffs.real + assert torch.any(torch.isclose(hopping_coeffs, hopping_coeffs.new_tensor(1.0))) + assert torch.any(torch.isclose(hopping_coeffs, hopping_coeffs.new_tensor(-1.0))) + + energies = { + name: vmc.local_energies() + for name, vmc in vmcs.items() + } + assert all(torch.isfinite(values).all() for values in energies.values()) + assert all(torch.is_complex(values) for values in energies.values()) + for name in ("boundary", "ctmrg"): + assert torch.allclose( + energies[name], + energies["exact"], + rtol=2.0e-10, + atol=1.0e-8, + ) + + +@pytest.mark.parametrize( + ("contraction", "contraction_opts"), + [ + ( + "boundary", + { + "mode": "mps", + "final_contract": True, + "final_contract_opts": {"optimize": "auto-hq"}, + "equalize_norms": False, + "progbar": False, + }, + ), + ( + "ctmrg", + { + "final_contract": False, + "max_separation": 1, + "inplace": False, + "equalize_norms": False, + "progbar": False, + }, + ), + ], + ids=("boundary", "ctmrg"), +) +def test_torch_fermion_vmc_approximate_contractions_support_sr( + contraction, contraction_opts +): + """Approximate PEPS contractions should support one finite SR update.""" + fermion = Fermion(spinful=True, symmetry="U1U1", t=1.0, U=2.0) + occupations = { + (x, y): (1, 0) if (x + y) % 2 == 0 else (0, 1) + for x in range(2) + for y in range(2) + } + peps = hrs_to_peps( + (2, 2), + fermion=fermion, + occupations=occupations, + chi=2, + seed=3, + dtype="complex128", + ) + vmc = TorchFermionVMC( + peps, + fermion=fermion, + n_walkers=4, + init_max_attempts=256, + contraction=contraction, + chi=4, + cutoff=1.0e-10, + contraction_opts=contraction_opts, + dtype=torch.complex128, + seed=32, + ) + + result = vmc.step( + sample_sweeps=1, + sr=True, + learning_rate=1.0e-3, + sr_method="minsr", + sr_parameter_mode="real-imag", + ) + + assert result.sr is not None + assert torch.isfinite(result.energy_mean) + assert torch.isfinite(result.sr.direction).all() + assert torch.isfinite(vmc.amplitudes).all() + if contraction == "boundary": + assert vmc.model.last_amplitude_batching == "serial" + + +def test_torch_fermion_vmc_accepts_z2_and_explicit_observable_terms(): + fermion = Fermion(spinful=True, symmetry="Z2", t=1.0, U=2.0) + peps = ps_to_peps(1, 2, fermion=fermion, seed=23, dtype="complex128") + terms = { + (0, 1): -fermion.hopping_operator(), + 0: fermion.onsite_term(0), + 1: fermion.onsite_term(1), + } + vmc = TorchFermionVMC( + peps, + terms=terms, + configs=[[2, 3]], + dtype=torch.complex128, + seed=24, + ) + + assert vmc.proposal == "spinful_z2" + assert vmc.sector == 0 + assert vmc.encoding == FermionSiteEncoding.symmray() + assert torch.isfinite(vmc.local_energies()).all() + + sample = vmc.sample_sweep() + n_up, n_down = count_spinful_particles(sample.configs, encoding=vmc.encoding) + assert ((n_up + n_down) % 2).tolist() == [0] + + +def test_torch_fermion_vmc_flat_z2_vmap_covers_sampling_measurement_and_sr(): + sr = pytest.importorskip("symmray") + peps = sr.networks.PEPS_fermionic_rand( + "Z2", + 2, + 2, + 2, + phys_dim=4, + subsizes="equal", + flat=True, + seed=1, + ) + fermion = Fermion(spinful=True, symmetry="Z2", t=1.0, U=2.0) + vmc = TorchFermionVMC( + peps, + fermion, + configs=[ + [2, 0, 1, 3], + [0, 2, 3, 1], + [0, 0, 0, 0], + [1, 1, 1, 1], + ], + n_walkers=4, + contraction="exact", + amplitude_batching="vmap", + dtype=torch.complex128, + seed=2, + ) + + assert vmc.model.last_amplitude_batching == "log-vmap" + sample = vmc.sample_sweep() + assert sample.configs.shape == (4, 4) + assert torch.isfinite(vmc.local_energies()).all() + + result = vmc.step( + sr=True, + learning_rate=1.0e-4, + sr_method="minsr", + sr_parameter_mode="real-imag", + ) + + assert result.sr is not None + assert torch.isfinite(result.local_energies).all() + assert vmc.model.last_amplitude_batching == "log-vmap" + + +def test_torch_fermion_vmc_accepts_z2z2_and_bp_sampling(): + fermion = Fermion(spinful=True, symmetry="Z2Z2", t=1.0, U=2.0) + peps = ps_to_peps(2, 2, fermion=fermion, seed=25, dtype="complex128") + vmc = TorchFermionVMC( + peps, + fermion, + n_walkers=2, + dtype=torch.complex128, + seed=26, + ) + + assert vmc.proposal == "spinful_z2z2" + assert vmc.sector == (0, 0) + sampler = vmc.make_bp_sampler( + n_chains=2, + bp_sampler_kwargs={"max_iterations": 3}, + sample_kwargs={"method": "exact"}, + seed=27, + ) + result = sampler.sample(n_samples=2, n_discard=1, n_thin=1) + n_up, n_down = count_spinful_particles( + result.configs.reshape(-1, vmc.n_sites), + encoding=vmc.encoding, + ) + assert ((n_up % 2) == 0).all() + assert ((n_down % 2) == 0).all() + + def test_spinful_exchange_hopping_proposal_preserves_particle_counts(): encoding = FermionSiteEncoding.symmray() configs = torch.tensor([ @@ -582,6 +3431,95 @@ def amplitude_fn(rows): assert result.acceptance_rate == 1.0 +def test_metropolis_proposal_stats_record_movewise_noops_and_acceptance(): + encoding = FermionSiteEncoding.vmc_torch() + + def amplitude_fn(rows): + return torch.ones(rows.shape[0], dtype=torch.float64) + + configs = torch.tensor([ + [encoding.empty, encoding.empty], + [encoding.empty, encoding.double], + [encoding.up, encoding.down], + [encoding.up, encoding.up], + ]) + result = metropolis_exchange_sweep( + configs, + amplitude_fn, + [(0, 1)], + proposal="spinful_z2", + hopping_rate=0.4, + spin_flip_rate=0.3, + pair_toggle_rate=0.5, + encoding=encoding, + generator=torch.Generator().manual_seed(37), + track_proposal_stats=True, + ) + + stats = result.proposal_stats + assert sum(move["selected"] for move in stats.values()) == len(configs) + assert sum(move["no_op"] for move in stats.values()) + sum( + move["proposed"] for move in stats.values() + ) == len(configs) + assert sum(move["proposed"] for move in stats.values()) == result.n_proposed + assert sum(move["accepted"] for move in stats.values()) == result.n_accepted + assert result.n_accepted == result.n_proposed + + +def test_torch_metropolis_warmup_adapts_rates_only_between_sweeps(): + encoding = FermionSiteEncoding.vmc_torch() + + def amplitude_fn(rows): + return torch.ones(rows.shape[0], dtype=torch.float64) + + sampler = TorchMetropolisSampler( + amplitude_fn, + [(0, 1)], + torch.full((128, 2), encoding.empty, dtype=torch.long), + proposal="spinful_z2", + hopping_rate=0.25, + spin_flip_rate=0.25, + pair_toggle_rate=0.25, + encoding=encoding, + seed=38, + ) + summary = sampler.warmup_proposal_mix( + n_sweeps=4, + adaptation_rate=1.5, + ) + + assert summary["n_sweeps"] == 4 + assert len(summary["history"]) == 4 + assert sum( + move["selected"] for move in summary["proposal_stats"].values() + ) == 4 * sampler.n_chains + assert summary["rates"]["spin_flip_rate"] < 0.25 + assert summary["rates"]["pair_toggle_rate"] > 0.25 + rates_after_warmup = dict(summary["rates"]) + sampler.sample_sweep() + assert { + "hopping_rate": sampler.hopping_rate, + "spin_flip_rate": sampler.spin_flip_rate, + "pair_toggle_rate": sampler.pair_toggle_rate, + } == rates_after_warmup + + +def test_torch_vmc_step_can_return_proposal_stats(): + driver = TorchVMCDriver( + ProductAmplitude(), + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + connection_fn="heisenberg", + proposal="spin", + ) + result = driver.step(track_proposal_stats=True) + + assert result.proposal_stats["exchange"]["selected"] == 2 + assert sum( + move["accepted"] for move in result.proposal_stats.values() + ) == result.n_accepted + + def test_random_sector_initializers_fix_particle_numbers(): spin = random_spin_configs( 4, diff --git a/vmc_run.jsonl b/vmc_run.jsonl new file mode 100644 index 0000000..6d9f847 --- /dev/null +++ b/vmc_run.jsonl @@ -0,0 +1,8 @@ +{"bond_dim": 4, "chi": null, "contraction": "exact", "cutoff": null, "kind": "start", "learning_rate": 0.002, "n_down": [12], "n_up": [12], "parameter_mode": "real-imag", "sector": [12, 12], "sector_ok": true, "shape": [4, 6], "sr_shift": 0.1, "timestamp": 1784672511.7933125, "walkers": 32} +{"acceptance_rate": 0.5309168443496801, "elapsed_seconds": 3.3106695500009664, "energy_mean": {"imag": 0.22778990794491039, "real": 46.521226825450086}, "energy_variance": 138.70378119719007, "finite_amplitudes": true, "kind": "step", "n_accepted": 498, "n_down": [12], "n_proposed": 938, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.26408440190611, "sr_method": "minsr", "sr_relative_residual": 2.4938287293541705e-16, "sr_solver": "solve", "step": 1, "timestamp": 1784672515.1043494} +{"acceptance_rate": 0.465625, "elapsed_seconds": 3.151786615000674, "energy_mean": {"imag": 0.724822871792401, "real": 45.20427038963628}, "energy_variance": 118.10042169509175, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 960, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.238563806597119, "sr_method": "minsr", "sr_relative_residual": 2.397772405188539e-16, "sr_solver": "solve", "step": 2, "timestamp": 1784672518.2566302} +{"acceptance_rate": 0.4720168954593453, "elapsed_seconds": 3.1743878300003416, "energy_mean": {"imag": -0.7437444363303103, "real": 46.29342148666581}, "energy_variance": 133.0351707344367, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 947, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.684107045175134, "sr_method": "minsr", "sr_relative_residual": 2.3595563568137667e-16, "sr_solver": "solve", "step": 3, "timestamp": 1784672521.4314444} +{"acceptance_rate": 0.5063694267515924, "elapsed_seconds": 3.0805508520006697, "energy_mean": {"imag": -1.0529867112447682, "real": 48.65619948919418}, "energy_variance": 136.51477170477895, "finite_amplitudes": true, "kind": "step", "n_accepted": 477, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.362278262656965, "sr_method": "minsr", "sr_relative_residual": 1.8282704776270593e-16, "sr_solver": "solve", "step": 4, "timestamp": 1784672524.5124533} +{"acceptance_rate": 0.4968152866242038, "elapsed_seconds": 3.2315187829990464, "energy_mean": {"imag": -0.10371245332633083, "real": 46.601914695472296}, "energy_variance": 117.54407195445287, "finite_amplitudes": true, "kind": "step", "n_accepted": 468, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 14.60727702783389, "sr_method": "minsr", "sr_relative_residual": 2.696026844089721e-16, "sr_solver": "solve", "step": 5, "timestamp": 1784672527.7443886} +{"acceptance_rate": 0.47257383966244726, "elapsed_seconds": 3.1237765490004676, "energy_mean": {"imag": -0.5855757114470148, "real": 48.02046940897043}, "energy_variance": 95.97787730909245, "finite_amplitudes": true, "kind": "step", "n_accepted": 448, "n_down": [12], "n_proposed": 948, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 12.695614221088631, "sr_method": "minsr", "sr_relative_residual": 2.6102958772648315e-16, "sr_solver": "solve", "step": 6, "timestamp": 1784672530.8686056} +{"acceptance_rate": 0.4920488894729925, "effective_sample_size": 128.0, "elapsed_seconds": 38.01418741899943, "energy_mean": {"imag": -0.23091432487365787, "real": 48.89689962587832}, "energy_stderr": 1.0109631548273583, "energy_variance": 130.82195205356612, "finite_amplitudes": true, "kind": "final_diagnostic", "n_down": [12], "n_samples": 128, "n_up": [12], "r_hat": 1.2469052698629284, "samples_per_second": 7.1675457132685425, "sector": [12, 12], "sector_ok": true, "tau": 1.0, "timestamp": 1784672548.7274194} diff --git a/vmc_run.log b/vmc_run.log new file mode 100644 index 0000000..6d9f847 --- /dev/null +++ b/vmc_run.log @@ -0,0 +1,8 @@ +{"bond_dim": 4, "chi": null, "contraction": "exact", "cutoff": null, "kind": "start", "learning_rate": 0.002, "n_down": [12], "n_up": [12], "parameter_mode": "real-imag", "sector": [12, 12], "sector_ok": true, "shape": [4, 6], "sr_shift": 0.1, "timestamp": 1784672511.7933125, "walkers": 32} +{"acceptance_rate": 0.5309168443496801, "elapsed_seconds": 3.3106695500009664, "energy_mean": {"imag": 0.22778990794491039, "real": 46.521226825450086}, "energy_variance": 138.70378119719007, "finite_amplitudes": true, "kind": "step", "n_accepted": 498, "n_down": [12], "n_proposed": 938, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.26408440190611, "sr_method": "minsr", "sr_relative_residual": 2.4938287293541705e-16, "sr_solver": "solve", "step": 1, "timestamp": 1784672515.1043494} +{"acceptance_rate": 0.465625, "elapsed_seconds": 3.151786615000674, "energy_mean": {"imag": 0.724822871792401, "real": 45.20427038963628}, "energy_variance": 118.10042169509175, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 960, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.238563806597119, "sr_method": "minsr", "sr_relative_residual": 2.397772405188539e-16, "sr_solver": "solve", "step": 2, "timestamp": 1784672518.2566302} +{"acceptance_rate": 0.4720168954593453, "elapsed_seconds": 3.1743878300003416, "energy_mean": {"imag": -0.7437444363303103, "real": 46.29342148666581}, "energy_variance": 133.0351707344367, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 947, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.684107045175134, "sr_method": "minsr", "sr_relative_residual": 2.3595563568137667e-16, "sr_solver": "solve", "step": 3, "timestamp": 1784672521.4314444} +{"acceptance_rate": 0.5063694267515924, "elapsed_seconds": 3.0805508520006697, "energy_mean": {"imag": -1.0529867112447682, "real": 48.65619948919418}, "energy_variance": 136.51477170477895, "finite_amplitudes": true, "kind": "step", "n_accepted": 477, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.362278262656965, "sr_method": "minsr", "sr_relative_residual": 1.8282704776270593e-16, "sr_solver": "solve", "step": 4, "timestamp": 1784672524.5124533} +{"acceptance_rate": 0.4968152866242038, "elapsed_seconds": 3.2315187829990464, "energy_mean": {"imag": -0.10371245332633083, "real": 46.601914695472296}, "energy_variance": 117.54407195445287, "finite_amplitudes": true, "kind": "step", "n_accepted": 468, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 14.60727702783389, "sr_method": "minsr", "sr_relative_residual": 2.696026844089721e-16, "sr_solver": "solve", "step": 5, "timestamp": 1784672527.7443886} +{"acceptance_rate": 0.47257383966244726, "elapsed_seconds": 3.1237765490004676, "energy_mean": {"imag": -0.5855757114470148, "real": 48.02046940897043}, "energy_variance": 95.97787730909245, "finite_amplitudes": true, "kind": "step", "n_accepted": 448, "n_down": [12], "n_proposed": 948, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 12.695614221088631, "sr_method": "minsr", "sr_relative_residual": 2.6102958772648315e-16, "sr_solver": "solve", "step": 6, "timestamp": 1784672530.8686056} +{"acceptance_rate": 0.4920488894729925, "effective_sample_size": 128.0, "elapsed_seconds": 38.01418741899943, "energy_mean": {"imag": -0.23091432487365787, "real": 48.89689962587832}, "energy_stderr": 1.0109631548273583, "energy_variance": 130.82195205356612, "finite_amplitudes": true, "kind": "final_diagnostic", "n_down": [12], "n_samples": 128, "n_up": [12], "r_hat": 1.2469052698629284, "samples_per_second": 7.1675457132685425, "sector": [12, 12], "sector_ok": true, "tau": 1.0, "timestamp": 1784672548.7274194} diff --git a/vmc_run_clean.jsonl b/vmc_run_clean.jsonl new file mode 100644 index 0000000..78a2ede --- /dev/null +++ b/vmc_run_clean.jsonl @@ -0,0 +1,8 @@ +{"bond_dim": 4, "chi": null, "contraction": "exact", "cutoff": null, "kind": "start", "learning_rate": 0.002, "n_down": [12], "n_up": [12], "parameter_mode": "real-imag", "sector": [12, 12], "sector_ok": true, "shape": [4, 6], "sr_shift": 0.1, "timestamp": 1784672583.8676605, "walkers": 32} +{"acceptance_rate": 0.5309168443496801, "elapsed_seconds": 3.308725337999931, "energy_mean": {"imag": 0.22778990794491039, "real": 46.521226825450086}, "energy_variance": 138.70378119719007, "finite_amplitudes": true, "kind": "step", "n_accepted": 498, "n_down": [12], "n_proposed": 938, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.26408440190611, "sr_method": "minsr", "sr_relative_residual": 2.4938287293541705e-16, "sr_solver": "solve", "step": 1, "timestamp": 1784672587.1767519} +{"acceptance_rate": 0.465625, "elapsed_seconds": 3.2370258560004004, "energy_mean": {"imag": 0.724822871792401, "real": 45.20427038963628}, "energy_variance": 118.10042169509175, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 960, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.238563806597119, "sr_method": "minsr", "sr_relative_residual": 2.397772405188539e-16, "sr_solver": "solve", "step": 2, "timestamp": 1784672590.4142272} +{"acceptance_rate": 0.4720168954593453, "elapsed_seconds": 3.2005471480006236, "energy_mean": {"imag": -0.7437444363303103, "real": 46.29342148666581}, "energy_variance": 133.0351707344367, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 947, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.684107045175134, "sr_method": "minsr", "sr_relative_residual": 2.3595563568137667e-16, "sr_solver": "solve", "step": 3, "timestamp": 1784672593.615232} +{"acceptance_rate": 0.5063694267515924, "elapsed_seconds": 3.237805417000345, "energy_mean": {"imag": -1.0529867112447682, "real": 48.65619948919418}, "energy_variance": 136.51477170477895, "finite_amplitudes": true, "kind": "step", "n_accepted": 477, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.362278262656965, "sr_method": "minsr", "sr_relative_residual": 1.8282704776270593e-16, "sr_solver": "solve", "step": 4, "timestamp": 1784672596.853507} +{"acceptance_rate": 0.4968152866242038, "elapsed_seconds": 3.27815523100071, "energy_mean": {"imag": -0.10371245332633083, "real": 46.601914695472296}, "energy_variance": 117.54407195445287, "finite_amplitudes": true, "kind": "step", "n_accepted": 468, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 14.60727702783389, "sr_method": "minsr", "sr_relative_residual": 2.696026844089721e-16, "sr_solver": "solve", "step": 5, "timestamp": 1784672600.1321368} +{"acceptance_rate": 0.47257383966244726, "elapsed_seconds": 3.2129157520012086, "energy_mean": {"imag": -0.5855757114470148, "real": 48.02046940897043}, "energy_variance": 95.97787730909245, "finite_amplitudes": true, "kind": "step", "n_accepted": 448, "n_down": [12], "n_proposed": 948, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 12.695614221088631, "sr_method": "minsr", "sr_relative_residual": 2.6102958772648315e-16, "sr_solver": "solve", "step": 6, "timestamp": 1784672603.3454819} +{"acceptance_rate": 0.4920488894729925, "effective_sample_size": 128.0, "elapsed_seconds": 39.083517668999775, "energy_mean": {"imag": -0.23091432487365787, "real": 48.89689962587832}, "energy_stderr": 1.0109631548273583, "energy_variance": 130.82195205356612, "finite_amplitudes": true, "kind": "final_diagnostic", "n_down": [12], "n_samples": 128, "n_up": [12], "r_hat": 1.2469052698629284, "samples_per_second": 7.049433909829366, "sector": [12, 12], "sector_ok": true, "tau": 1.0, "timestamp": 1784672621.503566} diff --git a/vmc_run_clean.log b/vmc_run_clean.log new file mode 100644 index 0000000..78a2ede --- /dev/null +++ b/vmc_run_clean.log @@ -0,0 +1,8 @@ +{"bond_dim": 4, "chi": null, "contraction": "exact", "cutoff": null, "kind": "start", "learning_rate": 0.002, "n_down": [12], "n_up": [12], "parameter_mode": "real-imag", "sector": [12, 12], "sector_ok": true, "shape": [4, 6], "sr_shift": 0.1, "timestamp": 1784672583.8676605, "walkers": 32} +{"acceptance_rate": 0.5309168443496801, "elapsed_seconds": 3.308725337999931, "energy_mean": {"imag": 0.22778990794491039, "real": 46.521226825450086}, "energy_variance": 138.70378119719007, "finite_amplitudes": true, "kind": "step", "n_accepted": 498, "n_down": [12], "n_proposed": 938, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.26408440190611, "sr_method": "minsr", "sr_relative_residual": 2.4938287293541705e-16, "sr_solver": "solve", "step": 1, "timestamp": 1784672587.1767519} +{"acceptance_rate": 0.465625, "elapsed_seconds": 3.2370258560004004, "energy_mean": {"imag": 0.724822871792401, "real": 45.20427038963628}, "energy_variance": 118.10042169509175, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 960, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.238563806597119, "sr_method": "minsr", "sr_relative_residual": 2.397772405188539e-16, "sr_solver": "solve", "step": 2, "timestamp": 1784672590.4142272} +{"acceptance_rate": 0.4720168954593453, "elapsed_seconds": 3.2005471480006236, "energy_mean": {"imag": -0.7437444363303103, "real": 46.29342148666581}, "energy_variance": 133.0351707344367, "finite_amplitudes": true, "kind": "step", "n_accepted": 447, "n_down": [12], "n_proposed": 947, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 15.684107045175134, "sr_method": "minsr", "sr_relative_residual": 2.3595563568137667e-16, "sr_solver": "solve", "step": 3, "timestamp": 1784672593.615232} +{"acceptance_rate": 0.5063694267515924, "elapsed_seconds": 3.237805417000345, "energy_mean": {"imag": -1.0529867112447682, "real": 48.65619948919418}, "energy_variance": 136.51477170477895, "finite_amplitudes": true, "kind": "step", "n_accepted": 477, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 16.362278262656965, "sr_method": "minsr", "sr_relative_residual": 1.8282704776270593e-16, "sr_solver": "solve", "step": 4, "timestamp": 1784672596.853507} +{"acceptance_rate": 0.4968152866242038, "elapsed_seconds": 3.27815523100071, "energy_mean": {"imag": -0.10371245332633083, "real": 46.601914695472296}, "energy_variance": 117.54407195445287, "finite_amplitudes": true, "kind": "step", "n_accepted": 468, "n_down": [12], "n_proposed": 942, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 14.60727702783389, "sr_method": "minsr", "sr_relative_residual": 2.696026844089721e-16, "sr_solver": "solve", "step": 5, "timestamp": 1784672600.1321368} +{"acceptance_rate": 0.47257383966244726, "elapsed_seconds": 3.2129157520012086, "energy_mean": {"imag": -0.5855757114470148, "real": 48.02046940897043}, "energy_variance": 95.97787730909245, "finite_amplitudes": true, "kind": "step", "n_accepted": 448, "n_down": [12], "n_proposed": 948, "n_up": [12], "sector": [12, 12], "sector_ok": true, "sr_direction_norm": 12.695614221088631, "sr_method": "minsr", "sr_relative_residual": 2.6102958772648315e-16, "sr_solver": "solve", "step": 6, "timestamp": 1784672603.3454819} +{"acceptance_rate": 0.4920488894729925, "effective_sample_size": 128.0, "elapsed_seconds": 39.083517668999775, "energy_mean": {"imag": -0.23091432487365787, "real": 48.89689962587832}, "energy_stderr": 1.0109631548273583, "energy_variance": 130.82195205356612, "finite_amplitudes": true, "kind": "final_diagnostic", "n_down": [12], "n_samples": 128, "n_up": [12], "r_hat": 1.2469052698629284, "samples_per_second": 7.049433909829366, "sector": [12, 12], "sector_ok": true, "tau": 1.0, "timestamp": 1784672621.503566}