SMC: replace multinomial sampling with systematic sampling

pymc/smc/smc.py

@@ -274,9 +274,7 @@ def update_beta_and_weights(self):
 
     def resample(self):
         """Resample particles based on importance weights"""
-        self.resampling_indexes = self.rng.choice(
-            np.arange(self.draws), size=self.draws, p=self.weights
-        )
+        self.resampling_indexes = systematic_resampling(self.weights, self.rng)
 
         self.tempered_posterior = self.tempered_posterior[self.resampling_indexes]
         self.prior_logp = self.prior_logp[self.resampling_indexes]
@@ -546,6 +544,36 @@ def sample_settings(self):
         return stats
 
 
+def systematic_resampling(weights, rng):
+    """
+    Systematic resampling.
+
+    Parameters
+    ----------
+    weights :
+        The weights should be probabilities and the total sum should be 1.
+
+    Returns
+    -------
+    new_indices: array
+        A vector of indices in the interval 0, ..., len(normalized_weights)
+    """
+    lnw = len(weights)
+    arange = np.arange(lnw)
+    uniform = (rng.random(1) + arange) / lnw
+
+    idx = 0
+    weight_accu = weights[0]
+    new_indices = np.empty(lnw, dtype=int)
+    for i in arange:
+        while uniform[i] > weight_accu:
+            idx += 1
+            weight_accu += weights[idx]
+        new_indices[i] = idx
+
+    return new_indices
+
+
 def _logp_forw(point, out_vars, in_vars, shared):
     """Compile Aesara function of the model and the input and output variables.
 

pymc/tests/test_smc.py

@@ -286,3 +286,11 @@ def test_proposal_dist_shape(self):
                     kernel=pm.smc.MH,
                     return_inferencedata=False,
                 )
+
+
+def test_systematic():
+    rng = np.random.default_rng(seed=34)
+    weights = [0.33, 0.33, 0.33]
+    np.testing.assert_array_equal(pm.smc.systematic_resampling(weights, rng), [0, 1, 2])
+    weights = [0.99, 0.01]
+    np.testing.assert_array_equal(pm.smc.systematic_resampling(weights, rng), [0, 0])