From 9ab6a443c551f76f7006bb2cb0344ffeb98d5619 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Fri, 6 Mar 2026 08:03:26 +0800 Subject: [PATCH] feat: support non-unique sampling in VMC and Guide algorithms - Added 'unique' parameter to VmcConfig and GuideConfig. - Implemented reweighting logic in energy and distribution calculations to support non-unique configurations. - Updated logging to include the 'unique' parameter status. - Cleaned up redundant type hints in MPS network. --- qmp/algorithms/guide.py | 32 +++++++++++++++++++++++--------- qmp/algorithms/vmc.py | 17 +++++++++++++---- qmp/networks/mps.py | 2 -- 3 files changed, 36 insertions(+), 15 deletions(-) diff --git a/qmp/algorithms/guide.py b/qmp/algorithms/guide.py index 8610a00..df3e75f 100644 --- a/qmp/algorithms/guide.py +++ b/qmp/algorithms/guide.py @@ -26,6 +26,8 @@ class GuideConfig: local_step: int = 1000 # The number of steps for the distribution optimizer dist_step: int = 100 + # Generate configuration uniquely + unique: bool = True def main( self, @@ -58,11 +60,12 @@ def main( ) logging.info( - "Arguments Summary: Sampling Count: %d, Relative Count: %d, Local Steps: %d, Dist Steps: %d", + "Arguments Summary: Sampling Count: %d, Relative Count: %d, Local Steps: %d, Dist Steps: %d, Unique: %s", self.sampling_count, self.relative_count, self.local_step, self.dist_step, + self.unique, ) if "guide" not in data: @@ -74,8 +77,16 @@ def main( logging.info("Starting a new optimization cycle") logging.info("Sampling configurations") - configs_src_network, psi_src_network, _, _ = network.generate_unique(self.sampling_count) - configs_src_sampling, psi_src_sampling, _, _ = sampling.generate_unique(self.sampling_count) + if self.unique: + configs_src_network, psi_src_network, _, _ = network.generate_unique(self.sampling_count) + reweight_network = torch.ones_like(psi_src_network, dtype=torch.float64) + configs_src_sampling, psi_src_sampling, _, _ = sampling.generate_unique(self.sampling_count) + reweight_sampling = torch.ones_like(psi_src_sampling, dtype=torch.float64) + else: + configs_src_network, psi_src_network, count_network, _ = network.generate(self.sampling_count) + reweight_network = count_network / psi_src_network.abs() ** 2 + configs_src_sampling, psi_src_sampling, count_sampling, _ = sampling.generate(self.sampling_count) + reweight_sampling = count_sampling / psi_src_sampling.abs() ** 2 logging.info("Calculating relative configurations") if self.relative_count <= len(configs_src_network): @@ -108,8 +119,9 @@ def energy_sampling() -> torch.Tensor: with torch.no_grad(): psi_dst = network(configs_dst) hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src) - num = psi_src.conj() @ hamiltonian_psi_dst - den = psi_src.conj() @ psi_src.detach() + psi_src_reweight = psi_src.conj() * reweight_sampling + num = psi_src_reweight @ hamiltonian_psi_dst + den = psi_src_reweight @ psi_src.detach() energy = num / den energy.real.backward() # type: ignore[no-untyped-call] return energy.real @@ -121,8 +133,9 @@ def energy_network() -> torch.Tensor: psi_src = network(configs_src) psi_dst = network(configs_dst) hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src) - num = psi_src.conj() @ hamiltonian_psi_dst - den = psi_src.conj() @ psi_src.detach() + psi_src_reweight = psi_src.conj() * reweight_network + num = psi_src_reweight @ hamiltonian_psi_dst + den = psi_src_reweight @ psi_src.detach() energy = num / den return energy.real @@ -133,8 +146,9 @@ def distribution() -> torch.Tensor: psi_src = network(configs_src) psi_dst = network(configs_dst) hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src) - num = psi_src.conj() @ hamiltonian_psi_dst - den = psi_src.conj() @ psi_src + psi_src_reweight = psi_src.conj() * reweight_sampling + num = psi_src_reweight @ hamiltonian_psi_dst + den = psi_src_reweight @ psi_src energy = num / den local_energy = hamiltonian_psi_dst / psi_src energy_diff = local_energy - energy diff --git a/qmp/algorithms/vmc.py b/qmp/algorithms/vmc.py index 89b5eba..67e84a5 100644 --- a/qmp/algorithms/vmc.py +++ b/qmp/algorithms/vmc.py @@ -24,6 +24,8 @@ class VmcConfig: relative_count: int = 40000 # The number of steps for the local optimizer local_step: int = 1000 + # Generate configuration uniquely + unique: bool = True def main( self, @@ -46,10 +48,11 @@ def main( ) logging.info( - "Arguments Summary: Sampling Count: %d, Relative Count: %d, Local Steps: %d, ", + "Arguments Summary: Sampling Count: %d, Relative Count: %d, Local Steps: %d, Unique: %s", self.sampling_count, self.relative_count, self.local_step, + self.unique, ) if "vmc" not in data: @@ -61,7 +64,12 @@ def main( logging.info("Starting a new optimization cycle") logging.info("Sampling configurations") - configs_i, psi_i, _, _ = network.generate_unique(self.sampling_count) + if self.unique: + configs_i, psi_i, _, _ = network.generate_unique(self.sampling_count) + reweight = torch.ones_like(psi_i, dtype=torch.float64) + else: + configs_i, psi_i, count, _ = network.generate(self.sampling_count) + reweight = count / psi_i.abs() ** 2 logging.info("Sampling completed, unique configurations count: %d", len(configs_i)) logging.info("Calculating relative configurations") @@ -82,8 +90,9 @@ def closure() -> torch.Tensor: with torch.no_grad(): psi_dst = network(configs_dst) hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src) - num = psi_src.conj() @ hamiltonian_psi_dst - den = psi_src.conj() @ psi_src.detach() + psi_src_reweight = psi_src.conj() * reweight + num = psi_src_reweight @ hamiltonian_psi_dst + den = psi_src_reweight @ psi_src.detach() energy = num / den energy = energy.real energy.backward() # type: ignore[no-untyped-call] diff --git a/qmp/networks/mps.py b/qmp/networks/mps.py index a1c57d3..0b7d4b1 100644 --- a/qmp/networks/mps.py +++ b/qmp/networks/mps.py @@ -320,8 +320,6 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t # M_next is the right-environment density matrix for positions after site i M_next: torch.Tensor = renv[self.sites - i - 1] - prob: torch.Tensor - v_new_all: torch.Tensor prob, v_new_all = self.mps.site_probs(v, i, M_next) # Sample a physical state at site i for each batch element