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

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