sample method now support status

Author: refraction-ray

CHANGELOG.md

@@ -22,6 +22,8 @@
 
 - Add `probability_sample` method for backend as an alternative for `random_choice` since it supports `status` as external randomness format
 
+- Add `status` support for `sample` and `sample_expection_ps` methods
+
 ### Changed
 
 - The inner mechanism for `sample_expectation_ps` is changed to sample representation from count representation for a fast speed

tensorcircuit/backends/abstract_backend.py

@@ -1098,7 +1098,7 @@ def probability_sample(
         p_cuml = self.cumsum(p)
         r = p_cuml[-1] * (1 - self.cast(status, p.dtype))
         ind = self.searchsorted(p_cuml, r)
-        a = self.arange(shots)
+        a = self.arange(self.shape_tuple(p)[0])
         res = self.gather1d(a, ind)
         return res
 

tensorcircuit/basecircuit.py

@@ -510,6 +510,7 @@ def sample(
         readout_error: Optional[Sequence[Any]] = None,
         format: Optional[str] = None,
         random_generator: Optional[Any] = None,
+        status: Optional[Tensor] = None,
     ) -> Any:
         """
         batched sampling from state or circuit tensor network directly
@@ -526,6 +527,9 @@ def sample(
         :type format: Optional[str]
         :param random_generator: random generator,  defaults to None
         :type random_generator: Optional[Any], optional
+        :param status: external randomness given by tensor uniformly from [0, 1],
+            if set, can overwrite random_generator
+        :type status: Optional[Tensor]
         :return: List (if batch) of tuple (binary configuration tensor and correponding probability)
             if the format is None, and consitent with format when given
         :rtype: Any
@@ -578,21 +582,20 @@ def perfect_sampling(key: Any) -> Any:
                 # readout error
                 if readout_error is not None:
                     p = self.readouterror_bs(readout_error, p)
-
-            a_range = backend.arange(2**self._nqubits)
-            if random_generator is None:
-                ch = backend.implicit_randc(a=a_range, shape=[nbatch], p=p)
-            else:
-                ch = backend.stateful_randc(
-                    random_generator, a=a_range, shape=[nbatch], p=p
-                )
+            ch = backend.probability_sample(nbatch, p, status, random_generator)
+            # if random_generator is None:
+            #     ch = backend.implicit_randc(a=a_range, shape=[nbatch], p=p)
+            # else:
+            #     ch = backend.stateful_randc(
+            #         random_generator, a=a_range, shape=[nbatch], p=p
+            #     )
             # confg = backend.mod(
             #     backend.right_shift(
             #         ch[..., None], backend.reverse(backend.arange(self._nqubits))
             #     ),
             #     2,
             # )
-            if format is None:
+            if format is None:  # for backward compatibility
                 confg = sample_int2bin(ch, self._nqubits)
                 prob = backend.gather1d(p, ch)
                 r = list(zip(confg, prob))  # type: ignore
@@ -608,6 +611,7 @@ def sample_expectation_ps(
         z: Optional[Sequence[int]] = None,
         shots: Optional[int] = None,
         random_generator: Optional[Any] = None,
+        status: Optional[Tensor] = None,
         readout_error: Optional[Sequence[Any]] = None,
         **kws: Any,
     ) -> Tensor:
@@ -635,7 +639,10 @@ def sample_expectation_ps(
         :param shots: number of measurement shots, defaults to None, indicating analytical result
         :type shots: Optional[int], optional
         :param random_generator: random_generator, defaults to None
-        :type random_general: Optional[Any]
+        :type random_generator: Optional[Any]
+        :param status: external randomness given by tensor uniformly from [0, 1],
+            if set, can overwrite random_generator
+        :type status: Optional[Tensor]
         :param readout_error: readout_error, defaults to None
         :type readout_error: Optional[Sequence[Any]]. Tensor, List, Tuple
         :return: [description]
@@ -674,6 +681,7 @@ def sample_expectation_ps(
                 counts=shots,
                 format="count_vector",
                 random_generator=random_generator,
+                status=status,
                 jittable=True,
                 is_prob=True,
             )
@@ -684,6 +692,7 @@ def sample_expectation_ps(
                 counts=shots,
                 format="sample_bin",
                 random_generator=random_generator,
+                status=status,
                 jittable=True,
                 is_prob=True,
             )

tensorcircuit/quantum.py

@@ -1999,6 +1999,7 @@ def measurement_counts(
     format: str = "count_vector",
     is_prob: bool = False,
     random_generator: Optional[Any] = None,
+    status: Optional[Tensor] = None,
     jittable: bool = False,
 ) -> Any:
     """
@@ -2048,7 +2049,10 @@ def measurement_counts(
         defaults to be False
     :type is_prob: bool
     :param random_generator: random_generator, defaults to None
-    :type random_general: Optional[Any]
+    :type random_generator: Optional[Any]
+    :param status: external randomness given by tensor uniformly from [0, 1],
+        if set, can overwrite random_generator
+    :type status: Optional[Tensor]
     :param jittable: if True, jax backend try using a jittable count, defaults to False
     :type jittable: bool
     :return: The counts for each bit string measured.
@@ -2066,7 +2070,6 @@ def measurement_counts(
         pi = backend.reshape(pi, [-1])
     d = int(backend.shape_tuple(pi)[0])
     n = int(np.log(d) / np.log(2) + 1e-8)
-    drange = backend.arange(d)
     if (counts is None) or counts <= 0:
         if format == "count_vector":
             return pi
@@ -2081,12 +2084,15 @@ def measurement_counts(
                 "unsupported format %s for analytical measurement" % format
             )
     else:
-        if random_generator is None:
-            raw_counts = backend.implicit_randc(drange, shape=counts, p=pi)
-        else:
-            raw_counts = backend.stateful_randc(
-                random_generator, a=drange, shape=counts, p=pi
-            )
+        raw_counts = backend.probability_sample(
+            counts, pi, status=status, g=random_generator
+        )
+        # if random_generator is None:
+        # raw_counts = backend.implicit_randc(drange, shape=counts, p=pi)
+        # else:
+        # raw_counts = backend.stateful_randc(
+        # random_generator, a=drange, shape=counts, p=pi
+        # )
         return sample2all(raw_counts, n, format=format, jittable=jittable)
 
 

tests/test_backends.py

@@ -291,7 +291,7 @@ def test_backend_methods_2(backend):
     r = tc.backend.probability_sample(10000, p, status=np.random.uniform(size=[10000]))
     _, r = np.unique(r, return_counts=True)
     np.testing.assert_allclose(
-        r - tc.backend.numpy(p) * 10000.0, np.zeros([10]), atol=100, rtol=1
+        r - tc.backend.numpy(p) * 10000.0, np.zeros([10]), atol=200, rtol=1
     )
 
 

tests/test_circuit.py

@@ -1017,6 +1017,14 @@ def test_batch_sample(backend):
             batch=8, allow_state=True, random_generator=tc.backend.get_random_state(42)
         )
     )
+    print(
+        c.sample(
+            batch=8,
+            allow_state=True,
+            status=np.random.uniform(size=[8]),
+            format="sample_bin",
+        )
+    )
 
 
 def test_expectation_y_bug():