initial MPS sampling and CircuitMPS sampling

quimb/tensor/circuit.py

@@ -4008,7 +4008,7 @@ def apply_gates(self, gates, progbar=False, **gate_opts):
     @property
     def psi(self):
         # no squeeze so that bond dims of 1 preserved
-        return self._psi
+        return self._psi.copy()
 
     @property
     def uni(self):
@@ -4029,6 +4029,23 @@ def get_psi_reverse_lightcone(self, where, keep_psi0=False):
         """
         return self.psi
 
+    def sample(
+        self,
+        C,
+        seed=None,
+    ):
+        """Sample the MPS circuit ``C`` times.
+
+        Parameters
+        ----------
+        C : int
+            The number of samples to generate.
+        seed : None, int, or generator, optional
+            A random seed or generator to use for reproducibility.
+        """
+        for config, _ in self._psi.sample(C, seed=seed):
+            yield "".join(map(str, config))
+
     def fidelity_estimate(self):
         r"""Estimate the fidelity of the current state based on its norm, which
         tracks how much the state has been truncated:
@@ -4122,6 +4139,26 @@ def get_psi_unordered(self):
         """
         return self._psi.copy()
 
+    def sample(self, C, seed=None):
+        """Sample the PermMPS circuit ``C`` times.
+
+        Parameters
+        ----------
+        C : int
+            The number of samples to generate.
+        seed : None, int, or generator, optional
+            A random seed or generator to use for reproducibility.
+
+        Yields
+        ------
+        str
+            The next sample bitstring.
+        """
+        # configuring is in physical order, so need to reorder for sampling
+        ordering = self.calc_qubit_ordering()
+        for config, _ in self._psi.sample(C, seed=seed):
+            yield "".join(str(config[i]) for i in ordering)
+
     @property
     def psi(self):
         # need to reindex and retag the MPS

quimb/tensor/tensor_1d.py

@@ -3191,6 +3191,83 @@ def measure(
 
     measure_ = functools.partialmethod(measure, inplace=True)
 
+    def sample_configuration(self, seed=None, info=None):
+        """Sample a configuration from this MPS.
+
+        Parameters
+        ----------
+        seed : None, int, or np.random.Generator, optional
+            A random seed or generator to use.
+        info : dict, optional
+            If given, will be used to infer and store various extra
+            information. Currently the key "cur_orthog" is used to store the
+            current orthogonality center.
+        """
+        import numpy as np
+
+        # if seed is already a generator this simply returns it
+        rng = np.random.default_rng(seed)
+
+        # right canonicalize
+        psi = self.canonicalize(0, info=info)
+
+        config = []
+        omega = 1.0
+        for i in range(psi.L):
+
+            # form local density matrix
+            ki = psi[i]
+            bi = ki.H
+            ix = psi.site_ind(i)
+            # contract diagonal to get probabilities
+            pi = (ki & bi).contract(output_inds=[ix]).data
+
+            # sample outcome using numpy
+            pi = do("to_numpy", pi).real
+            pi /= pi.sum()
+            xi = rng.choice(pi.size, p=pi)
+            config.append(xi)
+            # track local probability
+            omega *= pi[xi]
+
+            # project outcome
+            psi.isel_({ix: xi})
+            if i < psi.L - 1:
+                # and absorb projected site into next site
+                psi.contract_tags_([psi.site_tag(i), psi.site_tag(i + 1)])
+
+        return config, omega
+
+    def sample(self, C, seed=None, info=None):
+        """Generate ``C`` samples rom this MPS, along with their probabilities.
+
+        Parameters
+        ----------
+        C : int
+            The number of samples to generate.
+        seed : None, int, or np.random.Generator, optional
+            A random seed or generator to use.
+        info : dict, optional
+            If given, will be used to infer and store various extra
+            information. Currently the key "cur_orthog" is used to store the
+            current orthogonality center.
+
+        Yields
+        ------
+        config : sequence of int
+            The sample configuration.
+        omega : float
+            The probability of this configuration.
+        """
+
+        if info is None:
+            info = {}
+
+        # do right canonicalization once (supplying info avoids re-performing)
+        psi0 = self.canonicalize(0, info=info)
+
+        for _ in range(C):
+            yield psi0.sample_configuration(seed=seed, info=info)
 
 class MatrixProductOperator(TensorNetwork1DOperator, TensorNetwork1DFlat):
     """Initialise a matrix product operator, with auto labelling and tagging.