Merge pull request #380 from QunaSys/rounding_sampling

Rounding sampling

packages/core/quri_parts/core/sampling/__init__.py

@@ -438,7 +438,11 @@ def sample_from_probibility_distribution(
 ) -> MeasurementCounts:
     """Sample from a probibility distribution."""
     rng = np.random.default_rng()
-    counts = rng.multinomial(n_sample, np.round(probibility_distribution, 12))
+    rounded_prob = np.round(probibility_distribution, 12)
+    norm = np.sum(rounded_prob)
+    assert np.isclose(norm, 1.0), "Probabilty does not sum to 1.0"
+    rounded_prob = rounded_prob / norm
+    counts = rng.multinomial(n_sample, rounded_prob)
     return Counter(dict(((i, count) for i, count in enumerate(counts) if count > 0)))
 
 

packages/core/tests/core/sampling/test_create_sampler.py

@@ -33,6 +33,7 @@
     create_parametric_sampler_from_sampler,
     create_parametric_state_sampler_from_state_sampler,
     create_sampler_from_sampling_backend,
+    sample_from_probibility_distribution,
 )
 from quri_parts.core.state import (
     CircuitQuantumState,
@@ -43,6 +44,39 @@
 )
 
 
+def test_sample_from_probibility_distribution() -> None:
+    norm = 1.0
+    p1 = 0.4
+    prob = np.array([p1, norm - p1])
+    cnts = sample_from_probibility_distribution(1000, prob)
+    assert sum(list(cnts.values())) == 1000
+
+    norm = 1.000000000002
+    p1 = 0.4
+    prob = np.array([p1, norm - p1])
+    cnts = sample_from_probibility_distribution(1000, prob)
+    assert sum(list(cnts.values())) == 1000
+
+    norm = 1.001
+    p1 = 0.4
+    prob = np.array([p1, norm - p1])
+    with pytest.raises(AssertionError, match="Probabilty does not sum to 1.0"):
+        cnts = sample_from_probibility_distribution(1000, prob)
+
+    # Potentially dangerous case: large amount of small probabilities (p < 1e-12)
+    # When they are rounded away, sum of the rest of the probabilities slightly > 1.
+    n = 2**12
+    n_small = 2000
+    small_prob = -np.ones(n_small) * 1e-14
+    big_prob = np.ones(n - n_small) * (1 - np.sum(small_prob)) / (n - n_small)
+    prob = np.hstack([big_prob, small_prob])
+    with pytest.raises(ValueError):
+        rng = np.random.default_rng()
+        rng.multinomial(1000, prob.round(12))
+    cnts = sample_from_probibility_distribution(1000, prob)
+    assert sum(list(cnts.values())) == 1000
+
+
 def fake_sampler(circuit: NonParametricQuantumCircuit, shot: int) -> MeasurementCounts:
     cnt = 0.0
     for g in circuit.gates: