revert two moons overwrite

Author: LarsKue

bayesflow/simulators/benchmark_simulators/two_moons.py

@@ -1,23 +1,61 @@
 import numpy as np
 
-from ..simulator import Simulator
-from bayesflow.types import Shape
+from .benchmark_simulator import BenchmarkSimulator
 
 
-class TwoMoons(Simulator):
-    def __init__(self, lower_bound: float = -1.0, upper_bound: float = 1.0, seed: int = None):
+class TwoMoons(BenchmarkSimulator):
+    def __init__(self, lower_bound: float = -1.0, upper_bound: float = 1.0, rng: np.random.Generator = None):
+        """Two moons simulated benchmark.
+        See: https://arxiv.org/pdf/2101.04653.pdf, Task T.8
+
+        Parameters
+        ----------
+        lower_bound: float, optional, default: -1.0
+            The lower bound of the uniform prior
+        upper_bound: float, optional, default:  1.0
+            The upper bound of the uniform prior
+        rng: np.random.Generator or None, optional, default: None
+            An option random number generator to use
+        """
+
         self.lower_bound = lower_bound
         self.upper_bound = upper_bound
-        self.rng = np.random.default_rng(seed)
+        self.rng = rng
+        if self.rng is None:
+            self.rng = np.random.default_rng()
+
+    def prior(self):
+        """Generates a random draw from a 2-dimensional uniform prior bounded between
+        `lower_bound` and `upper_bound` which represents the two parameters of the two moons simulator.
+
+        Returns
+        -------
+        params: np.ndarray of shape (2, )
+            A single draw from the 2-dimensional uniform prior.
+        """
+
+        return self.rng.uniform(low=self.lower_bound, high=self.upper_bound, size=2)
+
+    def observation_model(self, params: np.ndarray):
+        """Implements data generation from the two-moons model with a bimodal posterior.
+
+        Parameters
+        ----------
+        params: np.ndarray of shape (2, )
+            The vector of two model parameters.
 
-    def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
-        r = self.rng.normal(0.1, 0.01, size=batch_shape + (1,))
-        alpha = self.rng.uniform(-0.5 * np.pi, 0.5 * np.pi, size=batch_shape + (1,))
-        theta = self.rng.uniform(self.lower_bound, self.upper_bound, size=batch_shape + (2,))
+        Returns
+        -------
+        observables: np.ndarray of shape (2, )
+            The 2D vector generated from the two moons simulator.
+        """
 
-        x1 = -np.abs(theta[..., :1] + theta[..., 1:]) / np.sqrt(2.0) + r * np.cos(alpha) + 0.25
-        x2 = (-theta[..., :1] + theta[..., 1:]) / np.sqrt(2.0) + r * np.sin(alpha)
+        # Generate noise
+        alpha = self.rng.uniform(low=-0.5 * np.pi, high=0.5 * np.pi)
+        r = self.rng.normal(loc=0.1, scale=0.01)
 
-        x = np.concatenate([x1, x2], axis=-1)
+        # Forward process
+        rhs1 = np.array([r * np.cos(alpha) + 0.25, r * np.sin(alpha)])
+        rhs2 = np.array([-np.abs(params[0] + params[1]) / np.sqrt(2.0), (-params[0] + params[1]) / np.sqrt(2.0)])
 
-        return dict(r=r, alpha=alpha, theta=theta, x=x)
+        return rhs1 + rhs2