Refactor mps.py: extract model-agnostic MPS class from WaveFunction

[Copilot]

WaveFunction conflated MPS tensor mechanics with the framework's packed-config interface, making the contraction logic impossible to reuse independently.

Changes

New MPS class

Model-agnostic core with only sites, physical_dim, virtual_dim parameters. Exposes the standard MPS contraction primitives:

Method	Description
amplitude(configs)	Real amplitudes from unpacked integer configs (batch, sites)
build_right_envs()	Precompute right-environment matrices M[sites]→M[1] for autoregressive sampling
init_left_vec(batch_size)	Initialise left boundary batch e_L = [1, 0, …]
site_probs(v, site, right_env)	Unnormalized per-site probs + candidate left vectors for one AR step
update_left_vec(v_new_all, s_i)	Advance left vector after sampling

Refactored WaveFunction

Now wraps self.mps: MPS and delegates all tensor contractions to it. Public interface is unchanged.

# MPS can now be used standalone for custom AR sampling loops
mps = MPS(sites=8, physical_dim=2, virtual_dim=16)

renv = mps.build_right_envs()          # precompute once
v = mps.init_left_vec(batch_size)      # e_L batch

for i in range(mps.sites):
    prob, v_new_all = mps.site_probs(v, i, renv[mps.sites - i - 1])
    s_i = torch.multinomial(prob, 1).squeeze(1)
    v = mps.update_left_vec(v_new_all, s_i)

Original prompt

---

This section details on the original issue you should resolve

<issue_title>更新 mps.py 可读性和可维护性</issue_title>
<issue_description>将 WaveFunction 中 mps 部分分离, 独立出一个和模型无关的 class MPS , 他的参数只有 virtual dim, physical dim . 支持的功能有:

1. 给定batch of config, 计算amplitude
2. 生成一个用于auto regressive采样的辅助变量(就是把右边的双层mps收缩的矩阵计算出来) , 然后可以让外层使用这个变量来auto regressive采样(batch的), 并更新左侧已经采样后的向量(batch的, 每个batch是virtual dim维度的), in short, 这个是mps收缩的常用技巧, 把它独立到class MPS中

然后外面的WaveFunction使用这个class MPS</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes 更新 mps.py 可读性和可维护性 #202

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

[Copilot]

Pull request overview

This PR extracts model-agnostic MPS tensor contraction logic into a standalone MPS module and rewires WaveFunction to delegate contraction/sampling primitives to it, while keeping WaveFunction’s public (packed-config) interface unchanged.

Changes:

• Added MPS class with core contraction + autoregressive sampling primitives (amplitude, right envs, left vec updates).
• Refactored WaveFunction to wrap self.mps: MPS and delegate contraction/sampling steps.
• Simplified/centralized right-environment and AR sampling logic by reusing MPS helpers.

---

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

[Copilot]

These methods are marked with @torch.jit.export, but the return annotations use PEP 585 built-in generics (list[...], tuple[...]). TorchScript compatibility has historically been stricter here; using typing.List / typing.Tuple (and, if needed, torch.jit.annotate) is more reliably scriptable. Consider switching the annotations to List[torch.Tensor] and Tuple[torch.Tensor, torch.Tensor] to avoid TorchScript compilation failures.

[Copilot]

These methods are marked with @torch.jit.export, but the return annotations use PEP 585 built-in generics (list[...], tuple[...]). TorchScript compatibility has historically been stricter here; using typing.List / typing.Tuple (and, if needed, torch.jit.annotate) is more reliably scriptable. Consider switching the annotations to List[torch.Tensor] and Tuple[torch.Tensor, torch.Tensor] to avoid TorchScript compilation failures.

[Copilot]

The API/docs describe real amplitudes/probabilities, but nothing enforces a real dtype for self.tensors. If self.tensors is ever complex, prob will be complex and .clamp(min=0.0) will raise at runtime. Either (a) explicitly restrict MPS tensors to real dtypes (and raise a clear error if complex), or (b) update the probability computation to a complex-safe nonnegative scalar (e.g., using a real scalar derived from the quadratic form) and adjust docs accordingly.

Suggested change

	# prob[b, s] = v_new_all[b, s] @ right_env @ v_new_all[b, s]^T (non-negative)
	Mv: torch.Tensor = torch.einsum("jk,bsk->bsj", right_env, v_new_all)
	prob: torch.Tensor = (v_new_all * Mv).sum(dim=-1).clamp(min=0.0)
	# prob[b, s] = v_new_all[b, s]^† @ right_env @ v_new_all[b, s] (non-negative real scalar)
	Mv: torch.Tensor = torch.einsum("jk,bsk->bsj", right_env, v_new_all)
	prob_raw: torch.Tensor = (v_new_all.conj() * Mv).sum(dim=-1)
	prob: torch.Tensor = prob_raw.real.clamp(min=0.0)

[Copilot]

amplitude() duplicates the left-boundary initialization logic that already exists in init_left_vec(). Consider calling self.init_left_vec(batch_size) here so the boundary convention is defined in one place (and changes won’t risk diverging implementations).

Suggested change

	v: torch.Tensor = torch.zeros([batch_size, self.virtual_dim], device=tensors.device, dtype=tensors.dtype)
	v[:, 0] = 1.0
	for i in range(self.sites):
	# A_i: (physical_dim, virtual_dim, virtual_dim)
	A_i: torch.Tensor = tensors[i]
	# s_i: (batch_size,) integer indices for the physical state at site i
	v: torch.Tensor = self.init_left_vec(batch_size)
	for i in range(self.sites):
	# A_i: (physical_dim, virtual_dim, virtual_dim)
	A_i: torch.Tensor = tensors[i]
	# s_i: (batch_size,) integer indices for the physical state at site i
	# s_i: (batch_size,) integer indices for the physical state at site i

[Copilot]

Casting configs[:, i] to int64 inside the per-site loop adds repeated overhead. Consider converting configs to torch.int64 once before the loop (or ensuring unpack_int already produces int64) and then indexing without per-iteration .to().