diff --git a/qmp/networks/mps.py b/qmp/networks/mps.py index 6eecc0f..a1c57d3 100644 --- a/qmp/networks/mps.py +++ b/qmp/networks/mps.py @@ -6,9 +6,9 @@ from ..utility.bitspack import pack_int, unpack_int -class WaveFunction(torch.nn.Module): +class MPS(torch.nn.Module): """ - A trivial Matrix Product State (MPS) wave function. + Model-agnostic Matrix Product State (MPS) core. The MPS tensors have shape (physical_dim, virtual_dim, virtual_dim) for each site, with open boundary conditions: the left boundary vector is e_0 = [1, 0, ..., 0] @@ -17,14 +17,23 @@ class WaveFunction(torch.nn.Module): The amplitude is computed as: Ψ(s) = e_L^T @ A[0][s_0] @ A[1][s_1] @ ... @ A[L-1][s_{L-1}] @ e_R where e_L = e_R = [1, 0, ..., 0] (first standard basis vector). + + Parameters + ---------- + sites : int + Number of sites in the MPS. + physical_dim : int + Dimension of the physical (local) Hilbert space at each site. + virtual_dim : int + Bond dimension (virtual dimension) of the MPS. """ def __init__( self, *, - sites: int, # Number of sites in the MPS - physical_dim: int, # Dimension of the physical (local) Hilbert space at each site - virtual_dim: int, # Bond dimension (virtual dimension) of the MPS + sites: int, + physical_dim: int, + virtual_dim: int, ) -> None: super().__init__() self.sites: int = sites @@ -38,6 +47,192 @@ def __init__( torch.randn(sites, physical_dim, virtual_dim, virtual_dim) ) + @torch.jit.export + def amplitude(self, configs: torch.Tensor) -> torch.Tensor: + """ + Compute MPS amplitudes for a batch of configurations. + + Parameters + ---------- + configs : torch.Tensor + Integer configurations of shape (batch_size, sites) with values in [0, physical_dim). + + Returns + ------- + torch.Tensor + Real amplitudes of shape (batch_size,), using the same dtype as the MPS tensors. + """ + batch_size: int = configs.shape[0] + tensors: torch.Tensor = self.tensors + + # v: (batch_size, virtual_dim) -- left environment vector, starting from e_L = [1, 0, ..., 0] + 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 + s_i: torch.Tensor = configs[:, i].to(dtype=torch.int64) + # A_s_i: (batch_size, virtual_dim, virtual_dim) + A_s_i: torch.Tensor = A_i[s_i] + # v = v @ A_s_i: (batch_size, virtual_dim) + v = torch.bmm(v.unsqueeze(1), A_s_i).squeeze(1) + + # Project onto right boundary e_R = [1, 0, ..., 0]: extract the first element + return v[:, 0] + + @torch.jit.export + def build_right_envs(self) -> list[torch.Tensor]: + """ + Pre-compute right-environment density matrices from right to left. + + Used as auxiliary variables for autoregressive sampling. Define M[i] as the + right-environment density matrix that captures all sites from ``i`` to the right + boundary (exclusive): + + - M[sites] = e_R @ e_R^T (right-boundary density matrix, only [0,0] entry is 1) + - M[i] = sum_s A[i][s] @ M[i+1] @ A[i][s]^T for i = sites-1, ..., 1 + + The returned list has length ``sites``; entry ``renv[k] = M[sites - k]``: + + - ``renv[0]`` = M[sites] (right boundary; used when sampling the last site) + - ``renv[k]`` = M[sites - k] (right env to the right of site sites-k-1) + - ``renv[sites-1]`` = M[1] (right env to the right of site 0) + + Note: M[0] (looking right from before site 0) is never needed and is not computed. + + When sampling at site ``i``, pass ``renv[sites - i - 1]`` as the right environment. + + Returns + ------- + list[torch.Tensor] + List of ``sites`` matrices, each of shape (virtual_dim, virtual_dim). + """ + tensors: torch.Tensor = self.tensors + device: torch.device = tensors.device + dtype: torch.dtype = tensors.dtype + + renv: list[torch.Tensor] = [] + M: torch.Tensor = torch.zeros([self.virtual_dim, self.virtual_dim], device=device, dtype=dtype) + M[0, 0] = 1.0 # e_R @ e_R^T (only the [0,0] entry is non-zero) + renv.append(M) + + for i in range(self.sites - 1, 0, -1): + A_i: torch.Tensor = tensors[i] # (physical_dim, virtual_dim, virtual_dim) + # M_new[a,d] = sum_{s,b,c} A_i[s,a,b] * M[b,c] * A_i[s,d,c] + # = sum_s (A_i[s] @ M @ A_i[s]^T)[a,d] + M = torch.einsum("sab,bc,sdc->ad", A_i, M, A_i) + renv.append(M) + + # After the loop renv has length `sites`: + # renv[0] = M[sites], renv[1] = M[sites-1], ..., renv[sites-1] = M[1] + return renv + + @torch.jit.export + def init_left_vec(self, batch_size: int) -> torch.Tensor: + """ + Initialize the left boundary vector for a batch. + + Parameters + ---------- + batch_size : int + Number of samples in the batch. + + Returns + ------- + torch.Tensor + Left boundary vector of shape (batch_size, virtual_dim) with + v[:, 0] = 1 and all other entries 0 (e_L = [1, 0, ..., 0]). + """ + tensors: torch.Tensor = self.tensors + v: torch.Tensor = torch.zeros([batch_size, self.virtual_dim], device=tensors.device, dtype=tensors.dtype) + v[:, 0] = 1.0 + return v + + @torch.jit.export + def site_probs(self, v: torch.Tensor, site: int, right_env: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute unnormalized conditional probabilities at a site for autoregressive sampling. + + Parameters + ---------- + v : torch.Tensor + Current left boundary vector of shape (batch_size, virtual_dim). + site : int + The site index to compute probabilities for. + right_env : torch.Tensor + Right-environment density matrix of shape (virtual_dim, virtual_dim) for the + sites to the right of ``site``. + + Returns + ------- + prob : torch.Tensor + Unnormalized probabilities of shape (batch_size, physical_dim), clamped to [0, ∞). + v_new_all : torch.Tensor + Candidate left vectors of shape (batch_size, physical_dim, virtual_dim), + one for each possible physical state. + """ + A_i: torch.Tensor = self.tensors[site] # (physical_dim, virtual_dim, virtual_dim) + + # v_new_all[b, s, j] = sum_k v[b, k] * A_i[s, k, j] (batch, phys, virt) + v_new_all: torch.Tensor = torch.einsum("bi,sij->bsj", v, A_i) + + # 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) + + return prob, v_new_all + + @torch.jit.export + def update_left_vec(self, v_new_all: torch.Tensor, s_i: torch.Tensor) -> torch.Tensor: + """ + Select the updated left vector for the sampled physical states. + + Parameters + ---------- + v_new_all : torch.Tensor + Candidate left vectors of shape (batch_size, physical_dim, virtual_dim). + s_i : torch.Tensor + Sampled physical-state indices of shape (batch_size,). + + Returns + ------- + torch.Tensor + Updated left vector of shape (batch_size, virtual_dim). + """ + batch_size: int = v_new_all.shape[0] + arange: torch.Tensor = torch.arange(batch_size, device=v_new_all.device) + return v_new_all[arange, s_i, :] + + +class WaveFunction(torch.nn.Module): + """ + A trivial Matrix Product State (MPS) wave function. + + Wraps the model-agnostic ``MPS`` core and provides the packed-config interface + expected by the rest of the framework. Open boundary conditions are used: + the left and right boundary vectors are both e_0 = [1, 0, ..., 0]. + + The amplitude is computed as: + Ψ(s) = e_L^T @ A[0][s_0] @ A[1][s_1] @ ... @ A[L-1][s_{L-1}] @ e_R + where e_L = e_R = [1, 0, ..., 0] (first standard basis vector). + """ + + def __init__( + self, + *, + sites: int, # Number of sites in the MPS + physical_dim: int, # Dimension of the physical (local) Hilbert space at each site + virtual_dim: int, # Bond dimension (virtual dimension) of the MPS + ) -> None: + super().__init__() + self.sites: int = sites + self.physical_dim: int = physical_dim + self.virtual_dim: int = virtual_dim + + self.mps: MPS = MPS(sites=sites, physical_dim=physical_dim, virtual_dim=virtual_dim) + # Dummy Parameter for Device and Dtype Retrieval # This parameter is used to infer the device and dtype of the model. self.dummy_param = torch.nn.Parameter(torch.empty(0)) @@ -72,33 +267,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: torch.Tensor Complex128 amplitudes of shape (batch_size,). """ - device: torch.device = self.dummy_param.device - dtype: torch.dtype = self.dummy_param.dtype - batch_size: int = x.shape[0] - # x_unpacked: batch_size * sites, each value in [0, physical_dim) - x_unpacked: torch.Tensor = unpack_int(x, size=self._bit_size(), last_dim=self.sites).view( - [batch_size, self.sites] - ) - - tensors: torch.Tensor = self.tensors.to(dtype=dtype) - - # v: batch_size * virtual_dim -- left environment vector, starting from e_L = [1, 0, ..., 0] - v: torch.Tensor = torch.zeros([batch_size, self.virtual_dim], device=device, dtype=dtype) - v[:, 0] = 1.0 + # Unpack bits to integer indices: (batch_size, sites) + configs: torch.Tensor = unpack_int(x, size=self._bit_size(), last_dim=self.sites).view([batch_size, self.sites]) - 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: torch.Tensor = x_unpacked[:, i].to(dtype=torch.int64) - # A_s_i: batch_size * virtual_dim * virtual_dim - A_s_i: torch.Tensor = A_i[s_i] - # v = v @ A_s_i: (batch_size, virtual_dim) - v = torch.bmm(v.unsqueeze(1), A_s_i).squeeze(1) - - # Project onto right boundary e_R = [1, 0, ..., 0]: extract the first element - amplitude: torch.Tensor = v[:, 0] + amplitude: torch.Tensor = self.mps.amplitude(configs) # Return as complex128 tensor with zero imaginary part return torch.view_as_complex(torch.stack([amplitude.double(), torch.zeros_like(amplitude).double()], dim=-1)) @@ -132,66 +305,31 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t The counts are a one-dimensional int64 tensor recording how many times each unique configuration was sampled. The last None value is reserved for future use. """ - device: torch.device = self.dummy_param.device - dtype: torch.dtype = self.dummy_param.dtype - - tensors: torch.Tensor = self.tensors.to(dtype=dtype) - # ------------------------------------------------------------------ # # Pre-compute right-environment density matrices from right to left. # - # # - # renv[k] = M[sites - k], where: # - # M[sites] = e_R @ e_R^T (right-boundary density matrix) # - # M[i] = sum_s A[i][s] @ M[i+1] @ A[i][s]^T for i < sites # - # # - # When sampling at site i we need M[i+1] = renv[sites - i - 1]. # # ------------------------------------------------------------------ # - renv: list[torch.Tensor] = [] - M: torch.Tensor = torch.zeros([self.virtual_dim, self.virtual_dim], device=device, dtype=dtype) - M[0, 0] = 1.0 # e_R @ e_R^T (only the [0,0] entry is non-zero) - renv.append(M) - - for i in range(self.sites - 1, 0, -1): - A_i: torch.Tensor = tensors[i] # physical_dim * virtual_dim * virtual_dim - # M_new[a,d] = sum_{s,b,c} A_i[s,a,b] * M[b,c] * A_i[s,d,c] - # = sum_s (A_i[s] @ M @ A_i[s]^T)[a,d] - M = torch.einsum("sab,bc,sdc->ad", A_i, M, A_i) - renv.append(M) - - # After the loop renv has length `sites`: - # renv[0] = M[sites], renv[1] = M[sites-1], ..., renv[sites-1] = M[1] + renv: list[torch.Tensor] = self.mps.build_right_envs() # ------------------------------------------------------------------ # # Autoregressive sampling. # # ------------------------------------------------------------------ # - # v: batch_size * virtual_dim (left state, initialised to e_L) - v: torch.Tensor = torch.zeros([batch_size, self.virtual_dim], device=device, dtype=dtype) - v[:, 0] = 1.0 - - samples: torch.Tensor = torch.zeros([batch_size, self.sites], device=device, dtype=torch.uint8) + v: torch.Tensor = self.mps.init_left_vec(batch_size) + samples: torch.Tensor = torch.zeros([batch_size, self.sites], device=self.dummy_param.device, dtype=torch.uint8) for i in range(self.sites): - A_i = tensors[i] # physical_dim * virtual_dim * virtual_dim - # M_next is the right-environment density matrix for the positions after site i + # M_next is the right-environment density matrix for positions after site i M_next: torch.Tensor = renv[self.sites - i - 1] - # v_new_all[b, s, j] = sum_k v[b, k] * A_i[s, k, j] (batch, phys, virt) - v_new_all: torch.Tensor = torch.einsum("bi,sij->bsj", v, A_i) - - # prob[b, s] = v_new_all[b, s] @ M_next @ v_new_all[b, s]^T (non-negative) - Mv: torch.Tensor = torch.einsum("jk,bsk->bsj", M_next, v_new_all) - prob: torch.Tensor = (v_new_all * Mv).sum(dim=-1) # batch_size * physical_dim - - # Correct any numerical round-off that may produce tiny negative values - prob = prob.clamp(min=0.0) + 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 - s_i: torch.Tensor = torch.multinomial(prob, num_samples=1).squeeze(1) # batch_size + s_i: torch.Tensor = torch.multinomial(prob, num_samples=1).squeeze(1) # (batch_size,) samples[:, i] = s_i.to(torch.uint8) # Advance the left state vector - arange: torch.Tensor = torch.arange(batch_size, device=device) - v = v_new_all[arange, s_i, :] # batch_size * virtual_dim + v = self.mps.update_left_vec(v_new_all, s_i) # ------------------------------------------------------------------ # # Pack, deduplicate, and count. #