Add optional rng argument to Metropolis Proposals

Author: ricardoV94

pymc/step_methods/metropolis.py

@@ -11,7 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 import numpy.random as nr
@@ -50,50 +50,64 @@
 
 
 class Proposal:
-    def __init__(self, s):
+    def __init__(self, s, rng_seed: Optional[int] = None):
         self.s = s
+        self.rng = np.random.default_rng(rng_seed)
 
 
 class NormalProposal(Proposal):
-    def __call__(self):
-        return nr.normal(scale=self.s)
+    def __call__(self, rng: Optional[np.random.Generator] = None):
+        if rng is None:
+            rng = self.rng
+        return rng.normal(scale=self.s)
 
 
 class UniformProposal(Proposal):
-    def __call__(self):
-        return nr.uniform(low=-self.s, high=self.s, size=len(self.s))
+    def __call__(self, rng: Optional[np.random.Generator] = None):
+        if rng is None:
+            rng = self.rng
+        return rng.uniform(low=-self.s, high=self.s, size=len(self.s))
 
 
 class CauchyProposal(Proposal):
-    def __call__(self):
-        return nr.standard_cauchy(size=np.size(self.s)) * self.s
+    def __call__(self, rng: Optional[np.random.Generator] = None):
+        if rng is None:
+            rng = self.rng
+        return rng.standard_cauchy(size=np.size(self.s)) * self.s
 
 
 class LaplaceProposal(Proposal):
-    def __call__(self):
+    def __call__(self, rng: Optional[np.random.Generator] = None):
+        if rng is None:
+            rng = self.rng
         size = np.size(self.s)
-        return (nr.standard_exponential(size=size) - nr.standard_exponential(size=size)) * self.s
+        return (rng.standard_exponential(size=size) - rng.standard_exponential(size=size)) * self.s
 
 
 class PoissonProposal(Proposal):
-    def __call__(self):
-        return nr.poisson(lam=self.s, size=np.size(self.s)) - self.s
+    def __call__(self, rng: Optional[np.random.Generator] = None):
+        if rng is None:
+            rng = self.rng
+        return rng.poisson(lam=self.s, size=np.size(self.s)) - self.s
 
 
 class MultivariateNormalProposal(Proposal):
-    def __init__(self, s):
-        n, m = s.shape
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        n, m = self.s.shape
         if n != m:
             raise ValueError("Covariance matrix is not symmetric.")
         self.n = n
-        self.chol = scipy.linalg.cholesky(s, lower=True)
+        self.chol = scipy.linalg.cholesky(self.s, lower=True)
 
-    def __call__(self, num_draws=None):
+    def __call__(self, num_draws=None, rng: Optional[np.random.Generator] = None):
+        if rng is None:
+            rng = self.rng
         if num_draws is not None:
-            b = np.random.randn(self.n, num_draws)
+            b = rng.normal(size=(self.n, num_draws))
             return np.dot(self.chol, b).T
         else:
-            b = np.random.randn(self.n)
+            b = rng.normal(size=self.n)
             return np.dot(self.chol, b)