move dist functions out of simulator class

pymc3/distributions/simulator.py

@@ -59,12 +59,12 @@ def __init__(
         self.epsilon = epsilon
 
         if distance == "gaussian_kernel":
-            self.distance = self.gaussian_kernel
+            self.distance = gaussian_kernel
         elif distance == "wasserstein":
-            self.distance = self.wasserstein
+            self.distance = wasserstein
             sum_stat = "sort"
         elif distance == "energy":
-            self.distance = self.energy
+            self.distance = energy
             sum_stat = "sort"
         elif hasattr(distance, "__call__"):
             self.distance = distance
@@ -107,17 +107,6 @@ def random(self, point=None, size=None):
         else:
             return np.array([self.function(*params) for _ in range(size)])
 
-    def gaussian_kernel(self, obs_data, sim_data):
-        return np.sum(-0.5 * ((obs_data - sim_data) / self.epsilon) ** 2)
-
-    # we are assuming obs_data and sim_data are already sorted!
-    def wasserstein(self, obs_data, sim_data):
-        return np.mean(np.abs((obs_data - sim_data) / self.epsilon))
-
-    # we are assuming obs_data and sim_data are already sorted!
-    def energy(self, obs_data, sim_data):
-        return 1.4142 * np.mean(((obs_data - sim_data) / self.epsilon) ** 2) ** 0.5
-
     def _repr_latex_(self, name=None, dist=None):
         if dist is None:
             dist = self
@@ -130,4 +119,26 @@ def _repr_latex_(self, name=None, dist=None):
 
 
 def identity(x):
+    """identity function, used as a summary statistics"""
     return x
+
+
+def gaussian_kernel(epsilon, obs_data, sim_data):
+    """gaussian distance function"""
+    return np.sum(-0.5 * ((obs_data - sim_data) / epsilon) ** 2)
+
+
+def wasserstein(epsilon, obs_data, sim_data):
+    """wasserstein distance function
+
+    We are assuming obs_data and sim_data are already sorted!
+    """
+    return np.mean(np.abs((obs_data - sim_data) / epsilon))
+
+
+def energy(epsilon, obs_data, sim_data):
+    """Energy distance function
+
+    We are assuming obs_data and sim_data are already sorted!
+    """
+    return 1.4142 * np.mean(((obs_data - sim_data) / epsilon) ** 2) ** 0.5

pymc3/smc/smc.py

@@ -366,4 +366,4 @@ def __call__(self, posterior):
         sim_data = self.function(**func_parameters)
         if self.save:
             self.save_data(sim_data)
-        return self.distance(self.observations, self.sum_stat(sim_data))
+        return self.distance(self.epsilon, self.observations, self.sum_stat(sim_data))

pymc3/tests/test_smc.py

@@ -104,8 +104,8 @@ def normal_sim(a, b):
             return np.random.normal(a, b, 1000)
 
         with pm.Model() as self.SMABC_test:
-            a = pm.Normal("a", mu=0, sd=5)
-            b = pm.HalfNormal("b", sd=2)
+            a = pm.Normal("a", mu=0, sigma=5)
+            b = pm.HalfNormal("b", sigma=2)
             s = pm.Simulator(
                 "s", normal_sim, params=(a, b), sum_stat="sort", epsilon=1, observed=self.data
             )