Skip to content

feat: support non-unique sampling in VMC and Guide algorithms #213

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hzhangxyz merged 1 commit into from
Mar 6, 2026
Merged

feat: support non-unique sampling in VMC and Guide algorithms #213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 23 additions & 9 deletions qmp/algorithms/guide.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +29 to +30

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Defaulting unique=True can still cause an immediate runtime failure for networks that don't implement generate_unique (e.g., qmp/networks/mps.py raises NotImplementedError). Consider defaulting to unique=False or implementing a runtime fallback to generate(...) when generate_unique(...) is unavailable.

Suggested change
# Generate configuration uniquely
unique: bool = True
# Generate configuration uniquely (opt-in; some networks lack generate_unique)
unique: bool = False

Copilot uses AI. Check for mistakes.

def main(
self,
Expand Down Expand Up @@ -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:
Expand All @@ -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
Comment on lines 79 to +89

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When unique is false, reweight_* = count_* / psi_src_*.abs() ** 2 is computed from psi_src_network / psi_src_sampling produced by *.generate(...). If sampling is done with gradients enabled (or later fixed to run under torch.enable_grad), these reweight tensors can retain a graph from the sampling forward pass and lead to repeated-backward errors or unintended gradients through sampling. Compute sampling + reweight under torch.no_grad() and/or detach the psi_src_* used for reweighting.

Copilot uses AI. Check for mistakes.

logging.info("Calculating relative configurations")
if self.relative_count <= len(configs_src_network):
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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
Expand Down
17 changes: 13 additions & 4 deletions qmp/algorithms/vmc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Defaulting unique=True can still cause an immediate runtime failure for networks that don't implement generate_unique (e.g., qmp/networks/mps.py raises NotImplementedError). Since this PR introduces unique specifically to support non-unique sampling, consider making the default False, or adding a fallback (try generate_unique, and on NotImplementedError fall back to generate + reweighting).

Suggested change
unique: bool = True
unique: bool = False

Copilot uses AI. Check for mistakes.

def main(
self,
Expand All @@ -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:
Expand All @@ -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
Comment on lines +67 to +72

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When unique is false, reweight = count / psi_i.abs() ** 2 is computed from psi_i returned by network.generate(...). If gradients are enabled for sampling (or enabled later via torch.enable_grad), this will keep an autograd graph from the sampling forward pass and can cause repeated-backward errors or unintended gradients through the sampling step. Compute sampling outputs and reweight under torch.no_grad() and/or explicitly detach psi_i before using it to build reweight.

Suggested change
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
with torch.no_grad():
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)
psi_i_detached = psi_i.detach()
reweight = count / psi_i_detached.abs() ** 2

Copilot uses AI. Check for mistakes.
logging.info("Sampling completed, unique configurations count: %d", len(configs_i))

logging.info("Calculating relative configurations")
Expand All @@ -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]
Expand Down
2 changes: 0 additions & 2 deletions qmp/networks/mps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down