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

[hzhangxyz]

• 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.

[Copilot]

Pull request overview

Adds an opt-in path to use non-unique (deduplicated-with-counts) sampling in the VMC and Guide optimization loops, with corresponding reweighting in energy and distribution computations.

Changes:

• Added a unique: bool flag to VmcConfig and GuideConfig, and logged its value.
• Implemented count-based reweighting when using network.generate(...) / sampling.generate(...) instead of generate_unique(...).
• Removed redundant local type annotations in MPS sampling code.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
qmp/networks/mps.py	Removes redundant local type hints in WaveFunction.generate() sampling loop.
qmp/algorithms/vmc.py	Adds unique config and introduces reweighted energy computation when sampling via generate().
qmp/algorithms/guide.py	Adds unique config and introduces reweighting in both energy estimators and the distribution-target calculation.

Comments suppressed due to low confidence (1)

qmp/algorithms/vmc.py:99

• RuntimeContext.setup() disables gradient tracking globally (torch.set_grad_enabled(False)), but this closure relies on autograd (energy.backward()). As written, psi_src = network(configs_src) will not require grad and backward() will error. Wrap the optimization step/closure invocation in torch.enable_grad(...) (as done in haar.py) or use a with torch.enable_grad(): block inside the closure so gradients are actually recorded.

                # Optimizing energy
                optimizer.zero_grad()
                psi_src = network(configs_src)
                with torch.no_grad():
                    psi_dst = network(configs_dst)
                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src)
                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]
                return energy

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

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]

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]

RuntimeContext.setup() disables gradient tracking globally (torch.set_grad_enabled(False)), but energy_sampling() calls energy.real.backward() and distribution() calls total_loss.backward(). Without wrapping these closures in torch.enable_grad(...) (like haar.py does), psi_src = network(configs_src) will not require grad and backward() will fail. Please ensure these optimization closures run under torch.enable_grad().

[Copilot]

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]

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