Skip to content

Commit ffe104f

Browse files
committed
feat(vmc): add reusable native measurement workflow
1 parent 0b22d90 commit ffe104f

14 files changed

Lines changed: 1763 additions & 131 deletions

File tree

‎.github/skills/pepsy-vmc/SKILL.md‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -196,9 +196,13 @@ For sampled configurations `x`, evaluate
196196

197197
`O_loc(x) = sum_{x'} O[x, x'] * psi(x') / psi(x)`.
198198

199-
Average local estimators over the Markov-chain samples. The canonical sampled
200-
driver method is `estimate_observable(...)`; `estimate_energy(...)` remains a
201-
compatibility wrapper. Return complex means
199+
Average local estimators over the Markov-chain samples. The canonical native
200+
Torch workflow is `samples = vmc.sample(sampling=...)`, followed by
201+
`vmc.measure(samples, observables={...})`. One `measure` call shares stored
202+
parent amplitudes and connected-target contractions across every named
203+
observable. `run(observables=..., sampling=...)` is the convenience form that
204+
warms up, samples, and measures once; `estimate_observables(...)` and
205+
`estimate_energy(...)` remain compatibility wrappers. Return complex means
202206
when appropriate but report Hermitian observables using their real part only
203207
after checking the imaginary residual. Keep energy, diagonal observables,
204208
hopping, correlations, and arbitrary supported Fermion observables on the

‎docs/api/vmc.md‎

Lines changed: 107 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -638,10 +638,10 @@ The estimate result retains the legacy `energy_mean`, `energy_variance`, and
638638
`energy_stderr` field names; they contain statistics for the configured
639639
observable. `estimate_energy(...)` remains as a compatibility alias.
640640

641-
For a coordinate-labelled PEPS, use `TorchFermionVMC` to derive the lattice,
642-
physical charge ordering, initial sector, and sampler rule in one place. Pass
643-
`fermion` to generate the default Hamiltonian, or omit it when supplying
644-
explicit `terms`:
641+
For compatibility with an existing lower-level sweep loop, coordinate-labelled
642+
PEPS can still initialize `TorchFermionVMC` in the constructor. New code
643+
should prefer the first-run recipe below instead. Pass `fermion` to generate
644+
the default Hamiltonian, or omit it when supplying explicit `terms`:
645645

646646
```python
647647
from pepsy import Fermion
@@ -663,6 +663,109 @@ result = vmc.estimate_observable(
663663
)
664664
```
665665

666+
For a native fermionic PEPS measurement, construct `TorchFermionVMC` from the
667+
state and native Fermion terms only. The first `sample` (or `warmup`) owns both
668+
the chain recipe and PEPS contraction recipe:
669+
670+
1. `SamplingConfig` owns the number of chains, retained samples, burn-in,
671+
thinning, and RNG seeds. In the native Torch sampler, `burn_in` counts
672+
discarded thinning intervals, so the `Metropolis` total is
673+
`(burn_in + n_samples_per_chain) * thin` batched sweeps; every batched
674+
sweep advances all chains once.
675+
2. `contraction_opts` is a single mapping with `method`, `chi`, `cutoff`, and
676+
any backend options such as `mode`. It is consumed when the first operation builds
677+
the amplitude model, then remains fixed with the Markov state.
678+
3. `observables` is a name-to-term mapping. `measure(samples, observables=...)`
679+
uses one retained batch for every entry. Include `"energy": terms`
680+
explicitly when the Hamiltonian should be visible in the measurement recipe.
681+
682+
```python
683+
sampling = pvmc.SamplingConfig(
684+
n_samples_per_chain=256,
685+
n_chains=32,
686+
burn_in=64,
687+
thin=2,
688+
seed=7,
689+
)
690+
contraction_opts = {
691+
"method": "boundary",
692+
"chi": 32,
693+
"cutoff": 1e-10,
694+
"mode": "mps",
695+
}
696+
697+
vmc = pvmc.TorchFermionVMC(
698+
peps,
699+
fermion=fermion,
700+
terms=terms, # native Fermi-Hubbard terms; no JW conversion
701+
)
702+
703+
# Optional, non-MCMC warm-up: inspect one valid PEPS amplitude.
704+
warmup = vmc.warmup(
705+
sampling=sampling,
706+
contraction_opts=contraction_opts,
707+
)
708+
709+
# Phase 1: exactly one Metropolis pass. `samples` retains psi(x) for every x.
710+
samples = vmc.sample(sampling=sampling, progress=True)
711+
712+
# Phase 2: no new Metropolis work. Reuse this batch for energy, eta, density, ...
713+
estimates = vmc.measure(
714+
samples,
715+
observables={"energy": terms, "eta": eta_terms},
716+
progress=True,
717+
)
718+
print(warmup.amplitude)
719+
print(estimates["energy"].energy_mean, estimates["eta"].energy_stderr)
720+
```
721+
722+
`estimate_observables({...}, sampling=..., contraction_opts=...)` provides the
723+
same one-batch behavior without retaining a separately named sample batch.
724+
`run(...)` is the one-command convenience form: warm up, sample once, then
725+
measure. `measure_samples(...)` remains the lower-level spelling of
726+
`measure(samples, ...)`. Native sample batches carry the PEPS parameter
727+
versions and contraction signature used to draw them, so `measure` rejects a
728+
batch after either changes; draw fresh samples after an optimization update.
729+
`progress=True` reports optional burn-in sweeps, MCMC
730+
sampling, then the shared connection-building, amplitude-contraction, and
731+
statistics phases. The `Metropolis` bar reports walkers (chains), retained
732+
samples per walker, burn-in/thinning, proposal, contraction method/`chi`,
733+
acceptance, and live boundary-environment cache reuse/build activity. Its
734+
`phase` is `equilibrate` while discarded intervals run, then `retain i/n` as
735+
each retained configuration per walker is recorded. The
736+
`Evaluation` bar reports the shared sample shape, observables, whether parent
737+
amplitudes were stored, connection count, and the diagonal/environment/direct
738+
target-amplitude split. The warm-up amplitude is a representative PEPS
739+
amplitude, not an energy estimate. The legacy constructor-level `n_walkers`,
740+
`contraction`, `chi`, and `cutoff` options remain supported for existing
741+
scripts, but new measurement code should keep them in `SamplingConfig` and
742+
`contraction_opts` as above.
743+
744+
External MPS/BP/tree proposal sampling uses the same explicit two-stage
745+
shape, but has no Metropolis burn-in or thinning. Pass its independent count
746+
as `n_samples`, rather than a `SamplingConfig`:
747+
748+
```python
749+
importance_samples = vmc.sample(
750+
proposal=mps_sampler,
751+
n_samples=512,
752+
fermion=proposal_fermion,
753+
one_d_to_two_d=mps_site_to_peps_coordinate,
754+
)
755+
importance_estimates = vmc.measure(
756+
importance_samples,
757+
observables={"energy": terms, "eta": eta_terms},
758+
)
759+
```
760+
761+
`importance_samples` stores PEPS-code configurations, `log q(x)`, and target
762+
parent amplitudes. The later `measure` call automatically forms the
763+
self-normalized weights `|psi(x)|**2 / q(x)` once and shares the resulting
764+
target-amplitude work across every observable. Unlike a target-Metropolis
765+
batch, an external-proposal batch remains valid after a PEPS update: `measure`
766+
refreshes its target amplitudes while retaining the fixed proposal density.
767+
`measure_from_proposal(...)` remains the one-call compatibility shortcut.
768+
666769
By default, `pbc=None` reads the PEPS cyclic axes through Quimb's
667770
`is_cyclic_x()` and `is_cyclic_y()` metadata; pass `pbc=` or `edges=` to
668771
override that inference. When explicit two-site `terms` are supplied, their

‎src/pepsy/vmc/__init__.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,11 @@
2626
"SymmetryFallbackWarning": ".api",
2727
"NetKetLocalConfigMap": ".netket",
2828
"NetKetChunkSettings": ".netket",
29+
"NetKetBuildTiming": ".netket",
30+
"NetKetAmplitudeTiming": ".netket",
2931
"NetKetPEPSVMC": ".netket",
3032
"NetKetVMCSetup": ".netket",
33+
"NetKetVMCConfig": ".netket",
3134
"NetKetFermiHubbardVMC": ".netket",
3235
"NetKetSparseFermiHubbardVMC": ".netket",
3336
"NetKetVMCSettings": ".netket",
@@ -38,7 +41,9 @@
3841
"TorchFermionVMC": ".torch",
3942
"TorchFermionVMCMetadata": ".torch",
4043
"TorchChainDiagnostics": ".torch",
44+
"TorchImportanceSamples": ".torch",
4145
"TorchMCMCSamples": ".torch",
46+
"TorchSampleProvenance": ".torch",
4247
"TorchMetropolisResult": ".torch",
4348
"TorchMetropolisSampler": ".torch",
4449
"TorchBPMetropolisSampler": ".torch",
@@ -48,7 +53,9 @@
4853
"TorchVMCSetup": ".torch",
4954
"TorchVMCEnergyEstimate": ".torch",
5055
"TorchVMCImportanceEstimate": ".torch",
56+
"TorchVMCMeasurementRun": ".torch",
5157
"TorchVMCStepResult": ".torch",
58+
"TorchVMCWarmupResult": ".torch",
5259
"TorchSRResult": ".torch",
5360
"TorchSquareLattice": ".torch",
5461
"VMCMeasurement": ".api",

‎src/pepsy/vmc/api.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,15 @@ def _resolve_contraction_config(contraction, chi=None, cutoff=None, options=None
149149

150150
@dataclass(frozen=True)
151151
class SamplingConfig:
152-
"""Shared chain-preserving sampling settings."""
152+
"""Shared chain-preserving sampling settings.
153+
154+
``burn_in`` is the number of discarded *thinning intervals* per chain.
155+
Thus the native Torch sampler advances each chain
156+
``(burn_in + n_samples_per_chain) * thin`` Metropolis sweeps: it discards
157+
``burn_in * thin`` sweeps, then retains one configuration after every
158+
``thin`` further sweeps. The returned batch has shape
159+
``(n_samples_per_chain, n_chains, n_sites)``.
160+
"""
153161

154162
n_samples_per_chain: int = 128
155163
n_chains: int = 16

‎src/pepsy/vmc/torch/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,11 +42,15 @@
4242
)
4343
from .results import (
4444
TorchMetropolisResult,
45+
TorchImportanceSamples,
4546
TorchMCMCSamples,
47+
TorchSampleProvenance,
4648
TorchChainDiagnostics,
4749
TorchVMCEnergyEstimate,
4850
TorchVMCImportanceEstimate,
51+
TorchVMCMeasurementRun,
4952
TorchVMCStepResult,
53+
TorchVMCWarmupResult,
5054
)
5155
from .sampler import TorchBPMetropolisSampler, TorchMetropolisSampler, metropolis_local_sampler
5256
from .sr import TorchSRResult, apply_torch_sr_update, solve_torch_sr, torch_log_derivative_matrix
@@ -59,7 +63,9 @@
5963
"TorchPEPSBoundaryAmplitude",
6064
"TorchConnections",
6165
"TorchMetropolisResult",
66+
"TorchImportanceSamples",
6267
"TorchMCMCSamples",
68+
"TorchSampleProvenance",
6369
"TorchChainDiagnostics",
6470
"TorchMetropolisSampler",
6571
"TorchBPMetropolisSampler",
@@ -68,7 +74,9 @@
6874
"TorchVMCSetup",
6975
"TorchVMCEnergyEstimate",
7076
"TorchVMCImportanceEstimate",
77+
"TorchVMCMeasurementRun",
7178
"TorchVMCStepResult",
79+
"TorchVMCWarmupResult",
7280
"TorchSRResult",
7381
"TorchSquareLattice",
7482
"apply_torch_sr_update",

‎src/pepsy/vmc/torch/_core.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,15 @@
3030
)
3131
from .results import (
3232
TorchChainDiagnostics,
33+
TorchImportanceSamples,
3334
TorchMCMCSamples,
3435
TorchMetropolisResult,
36+
TorchSampleProvenance,
3537
TorchVMCImportanceEstimate,
3638
TorchVMCEnergyEstimate,
39+
TorchVMCMeasurementRun,
3740
TorchVMCStepResult,
41+
TorchVMCWarmupResult,
3842
)
3943
from .metadata import (
4044
TorchFermionVMCMetadata,
@@ -116,7 +120,9 @@
116120
"TorchPEPSBoundaryAmplitude",
117121
"TorchConnections",
118122
"TorchMetropolisResult",
123+
"TorchImportanceSamples",
119124
"TorchMCMCSamples",
125+
"TorchSampleProvenance",
120126
"TorchChainDiagnostics",
121127
"TorchMetropolisSampler",
122128
"TorchBPMetropolisSampler",
@@ -125,7 +131,9 @@
125131
"TorchVMCSetup",
126132
"TorchVMCEnergyEstimate",
127133
"TorchVMCImportanceEstimate",
134+
"TorchVMCMeasurementRun",
128135
"TorchVMCStepResult",
136+
"TorchVMCWarmupResult",
129137
"TorchSRResult",
130138
"TorchSquareLattice",
131139
"apply_torch_sr_update",

0 commit comments

Comments
 (0)