From 512822a8a3d555fc6846762303a704dd000f8e6b Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Thu, 24 Dec 2020 17:01:45 -0800
Subject: [PATCH 01/15] Implement random for MvGaussianRandomWalk

Implements the random method for MvGaussianRandomWalk,
partially fixing #4337.
---
 pymc3/distributions/timeseries.py | 68 +++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 5174851547..85b554475d 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -461,6 +461,74 @@ def _distr_parameters_for_repr(self):
         return ["mu", "cov"]
 
 
+    def random(self, point=None, size=None):
+        """
+        Draw random values from MvGaussianRandomWalk.
+
+        Parameters
+        ----------
+        point: dict, optional
+            Dict of variable values on which random values are to be
+            conditioned (uses default point if not specified).
+        size: int, optional
+            Desired size of random sample (returns one sample if not
+            specified).
+
+        Returns
+        -------
+        array
+        """
+
+        param_attribute = getattr(self.innov, "chol_cov" if self.innov._cov_type == "chol" else self.innov._cov_type)
+        mu, param = distribution.draw_values([self.innov.mu, param_attribute], point=point, size=size)
+        return distribution.generate_samples(
+            self._random, 
+            size=size,
+            dist_shape=self.shape,
+            not_broadcast_kwargs={
+                "sample_shape": to_tuple(size),
+                "param": param,
+                "mu": mu,
+                "cov_type": self.innov._cov_type
+                }
+        )
+
+    def _random(self, mu, param, size, sample_shape, cov_type):     
+        """
+        Implements the multivariate Gaussian random walk as a cumulative
+        sum of i.i.d. multivariate Gaussians.
+        Assumes that
+        size is of the form (samples, time, dims).
+        """       
+
+        if cov_type == "chol":
+            cov = np.matmul(param, param.transpose()) 
+        elif cov_type == "tau":
+            cov = np.linalg.inv(param) 
+        else:
+            cov = param
+
+        # time axis comes after the sample axis
+        time_axis = len(sample_shape) 
+
+        # spatial axis is last
+        spatial_axis = -1 
+
+        rv = stats.multivariate_normal(mean=mu, cov=cov) 
+
+        # only feed in sample and time dimensions since stats.multivariate_normal
+        # automatically adds back in the spatial dimensions to the end when it samples.
+        data = rv.rvs(size[:spatial_axis]).cumsum(axis=time_axis) 
+
+        # shift the walk to start at zero
+        if len(data.shape) > 2:
+            for i in range(size[0]):
+                data[i] = data[i] - data[i][0]
+        else:
+            data = data - data[0]
+        return data
+
+
 class MvStudentTRandomWalk(MvGaussianRandomWalk):
     r"""
     Multivariate Random Walk with StudentT innovations

From 28591b44ead0af0b566396d8ba06d8236edaa9e9 Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Thu, 24 Dec 2020 17:04:15 -0800
Subject: [PATCH 02/15] Update docstring for GaussianRandomWalk._random()

---
 pymc3/distributions/timeseries.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 85b554475d..c9236268ff 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -275,7 +275,7 @@ def _random(self, sigma, mu, size, sample_shape):
         """Implement a Gaussian random walk as a cumulative sum of normals.
         axis = len(size) - 1 denotes the axis along which cumulative sum would be calculated.
         This might need to be corrected in future when issue #4010 is fixed.
-        Lines 318-322 ties the starting point of each instance of random walk to 0"
+        Lines 291-295 set the starting point of each instance of the random walk to 0."
         """
         if size[len(sample_shape)] == sample_shape:
             axis = len(sample_shape)

From fe33fe5888e210cf62ab3b077bd132a5822e9072 Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Sun, 27 Dec 2020 11:18:48 -0800
Subject: [PATCH 03/15] Add example to random docstring

---
 pymc3/distributions/timeseries.py | 40 ++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index c9236268ff..9769defd5c 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -460,7 +460,6 @@ def logp(self, x):
     def _distr_parameters_for_repr(self):
         return ["mu", "cov"]
 
-
     def random(self, point=None, size=None):
         """
         Draw random values from MvGaussianRandomWalk.
@@ -477,48 +476,61 @@ def random(self, point=None, size=None):
         Returns
         -------
         array
+
+
+        Examples
+        -------
+        Create one sample from a 2-dimensional Gaussian random walk with 10 timesteps::
+
+            mu = np.array([1.0, 0.0])
+            cov = np.array([[1.0, 0.0], [0.0, 2.0]])
+            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=1)
         """
 
-        param_attribute = getattr(self.innov, "chol_cov" if self.innov._cov_type == "chol" else self.innov._cov_type)
-        mu, param = distribution.draw_values([self.innov.mu, param_attribute], point=point, size=size)
+        param_attribute = getattr(
+            self.innov, "chol_cov" if self.innov._cov_type == "chol" else self.innov._cov_type
+        )
+        mu, param = distribution.draw_values(
+            [self.innov.mu, param_attribute], point=point, size=size
+        )
         return distribution.generate_samples(
-            self._random, 
+            self._random,
             size=size,
             dist_shape=self.shape,
             not_broadcast_kwargs={
                 "sample_shape": to_tuple(size),
                 "param": param,
                 "mu": mu,
-                "cov_type": self.innov._cov_type
-                }
+                "cov_type": self.innov._cov_type,
+            },
         )
 
-    def _random(self, mu, param, size, sample_shape, cov_type):     
+    def _random(self, mu, param, size, sample_shape, cov_type):
         """
         Implements the multivariate Gaussian random walk as a cumulative
         sum of i.i.d. multivariate Gaussians.
         Assumes that
         size is of the form (samples, time, dims).
-        """       
+        """
 
         if cov_type == "chol":
-            cov = np.matmul(param, param.transpose()) 
+            cov = np.matmul(param, param.transpose())
         elif cov_type == "tau":
-            cov = np.linalg.inv(param) 
+            cov = np.linalg.inv(param)
         else:
             cov = param
 
         # time axis comes after the sample axis
-        time_axis = len(sample_shape) 
+        time_axis = len(sample_shape)
 
         # spatial axis is last
-        spatial_axis = -1 
+        spatial_axis = -1
 
-        rv = stats.multivariate_normal(mean=mu, cov=cov) 
+        rv = stats.multivariate_normal(mean=mu, cov=cov)
 
         # only feed in sample and time dimensions since stats.multivariate_normal
         # automatically adds back in the spatial dimensions to the end when it samples.
-        data = rv.rvs(size[:spatial_axis]).cumsum(axis=time_axis) 
+        data = rv.rvs(size[:spatial_axis]).cumsum(axis=time_axis)
 
         # shift the walk to start at zero
         if len(data.shape) > 2:

From c3cc70ccae23ba301b7710b68c5631247cfa297c Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Mon, 28 Dec 2020 16:35:41 -0800
Subject: [PATCH 04/15] Modify MvGaussianRandomWalk.random

Improves the implementation by using MvNormal.random as suggested by
@Sayam753. Also updated its docstring to add more examples.
---
 pymc3/distributions/timeseries.py | 75 ++++++++++++-------------------
 1 file changed, 29 insertions(+), 46 deletions(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 9769defd5c..670fbb8b14 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -435,7 +435,7 @@ def __init__(
 
         self.init = init
         self.innovArgs = (mu, cov, tau, chol, lower)
-        self.innov = multivariate.MvNormal.dist(*self.innovArgs)
+        self.innov = multivariate.MvNormal.dist(*self.innovArgs, shape=self.shape[-1])
         self.mean = tt.as_tensor_variable(0.0)
 
     def logp(self, x):
@@ -469,7 +469,7 @@ def random(self, point=None, size=None):
         point: dict, optional
             Dict of variable values on which random values are to be
             conditioned (uses default point if not specified).
-        size: int, optional
+        size: int or tuple of ints, optional
             Desired size of random sample (returns one sample if not
             specified).
 
@@ -484,61 +484,44 @@ def random(self, point=None, size=None):
 
             mu = np.array([1.0, 0.0])
             cov = np.array([[1.0, 0.0], [0.0, 2.0]])
-            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=1)
-        """
+            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random()
 
-        param_attribute = getattr(
-            self.innov, "chol_cov" if self.innov._cov_type == "chol" else self.innov._cov_type
-        )
-        mu, param = distribution.draw_values(
-            [self.innov.mu, param_attribute], point=point, size=size
-        )
-        return distribution.generate_samples(
-            self._random,
-            size=size,
-            dist_shape=self.shape,
-            not_broadcast_kwargs={
-                "sample_shape": to_tuple(size),
-                "param": param,
-                "mu": mu,
-                "cov_type": self.innov._cov_type,
-            },
-        )
+        Create three samples from a 2-dimensional Gaussian random walk with 10 timesteps::
 
-    def _random(self, mu, param, size, sample_shape, cov_type):
-        """
-        Implements the multivariate Gaussian random walk as a cumulative
-        sum of i.i.d. multivariate Gaussians.
-        Assumes that
-        size is of the form (samples, time, dims).
+            mu = np.array([1.0, 0.0])
+            cov = np.array([[1.0, 0.0], [0.0, 2.0]])
+            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=3)
+
+        Create four samples from a 2-dimensional Gaussian random walk with 10
+        timesteps, indexed with a (2, 2) array::
+
+            mu = np.array([1.0, 0.0])
+            cov = np.array([[1.0, 0.0], [0.0, 2.0]])
+            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=(2, 2))
         """
 
-        if cov_type == "chol":
-            cov = np.matmul(param, param.transpose())
-        elif cov_type == "tau":
-            cov = np.linalg.inv(param)
-        else:
-            cov = param
+        time_steps = self.shape[0]
+        size = to_tuple(size)
 
-        # time axis comes after the sample axis
-        time_axis = len(sample_shape)
+        # for each draw specified by the size input, we need to draw time_steps many
+        # samples from MvNormal.
+        size_time_steps = size + to_tuple(time_steps)
 
-        # spatial axis is last
-        spatial_axis = -1
+        multivariate_samples = self.innov.random(point=point, size=size_time_steps)
+        # this has shape (size, time_steps, MvNormal_shape)
 
-        rv = stats.multivariate_normal(mean=mu, cov=cov)
+        time_axis = len(size)
 
-        # only feed in sample and time dimensions since stats.multivariate_normal
-        # automatically adds back in the spatial dimensions to the end when it samples.
-        data = rv.rvs(size[:spatial_axis]).cumsum(axis=time_axis)
+        multivariate_samples = multivariate_samples.cumsum(axis=time_axis)
 
         # shift the walk to start at zero
-        if len(data.shape) > 2:
-            for i in range(size[0]):
-                data[i] = data[i] - data[i][0]
+        if len(multivariate_samples.shape) > 2:
+            # this for loop covers the case where size is a tuple
+            for idx in np.ndindex(size):
+                multivariate_samples[idx] = multivariate_samples[idx] - multivariate_samples[idx][0]
         else:
-            data = data - data[0]
-        return data
+            multivariate_samples = multivariate_samples - multivariate_samples[0]
+        return multivariate_samples
 
 
 class MvStudentTRandomWalk(MvGaussianRandomWalk):

From 34199692878c03e65ceb98b609db5f955d9ff6fa Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Mon, 28 Dec 2020 16:47:14 -0800
Subject: [PATCH 05/15] Add TestMvGaussianRandomWalk

---
 pymc3/tests/test_distributions_random.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py
index 8d19836174..0dfcf5c57e 100644
--- a/pymc3/tests/test_distributions_random.py
+++ b/pymc3/tests/test_distributions_random.py
@@ -310,6 +310,12 @@ class TestGaussianRandomWalk(BaseTestCases.BaseTestCase):
     default_shape = (1,)
 
 
+class TestMvGaussianRandomWalk(BaseTestCases.BaseTestCase):
+    distribution = pm.MvGaussianRandomWalk
+    params = {"mu": np.array([1.0, 0.0]), "cov": np.array([[1.0, 0.0], [0.0, 2.0]])}
+    default_shape = (10, 2)
+
+
 class TestNormal(BaseTestCases.BaseTestCase):
     distribution = pm.Normal
     params = {"mu": 0.0, "tau": 1.0}

From 111f9c58254f031fa2788067151b0770f116fdab Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Tue, 29 Dec 2020 12:57:10 -0800
Subject: [PATCH 06/15] Pass entire shape to MvNormal

This commit rewrites the random method to pass the entire shape
into MvNormal. MvGaussianRandomWalk also now supports an integer shape
that must match the dimensionality of mu. In this case it is assumed
that there are no time steps, and a random sample from the multivariate
normal distribution will be returned.
---
 pymc3/distributions/timeseries.py | 45 ++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 670fbb8b14..f2ec70d765 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -435,7 +435,7 @@ def __init__(
 
         self.init = init
         self.innovArgs = (mu, cov, tau, chol, lower)
-        self.innov = multivariate.MvNormal.dist(*self.innovArgs, shape=self.shape[-1])
+        self.innov = multivariate.MvNormal.dist(*self.innovArgs, shape=self.shape)
         self.mean = tt.as_tensor_variable(0.0)
 
     def logp(self, x):
@@ -452,6 +452,10 @@ def logp(self, x):
         -------
         TensorVariable
         """
+
+        if x.ndim == 1:
+            x = x[np.newaxis, :]
+
         x_im1 = x[:-1]
         x_i = x[1:]
 
@@ -500,27 +504,32 @@ def random(self, point=None, size=None):
             sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=(2, 2))
         """
 
-        time_steps = self.shape[0]
-        size = to_tuple(size)
-
         # for each draw specified by the size input, we need to draw time_steps many
         # samples from MvNormal.
-        size_time_steps = size + to_tuple(time_steps)
-
-        multivariate_samples = self.innov.random(point=point, size=size_time_steps)
-        # this has shape (size, time_steps, MvNormal_shape)
 
-        time_axis = len(size)
-
-        multivariate_samples = multivariate_samples.cumsum(axis=time_axis)
+        size = to_tuple(size)
+        multivariate_samples = self.innov.random(point=point, size=size)
+        # this has shape (size, self.shape)
+
+        if len(self.shape) >= 2:
+            # have time dimension in first slot of shape. Therefore the time
+            # component can be accessed with the index equal to the length of size.
+            time_axis = len(size)
+            multivariate_samples = multivariate_samples.cumsum(axis=time_axis)
+
+            if time_axis:
+                # this for loop covers the case where size is a tuple
+                for idx in np.ndindex(size):
+                    multivariate_samples[idx] = (
+                        multivariate_samples[idx] - multivariate_samples[idx][0]
+                    )
+            else:
+                # size was passed as None
+                multivariate_samples = multivariate_samples - multivariate_samples[0]
 
-        # shift the walk to start at zero
-        if len(multivariate_samples.shape) > 2:
-            # this for loop covers the case where size is a tuple
-            for idx in np.ndindex(size):
-                multivariate_samples[idx] = multivariate_samples[idx] - multivariate_samples[idx][0]
-        else:
-            multivariate_samples = multivariate_samples - multivariate_samples[0]
+        # if the above if statement fails, then only a spatial dimension was passed in for self.shape.
+        # Therefore don't subtract off the initial value since otherwise you get all zeros
+        # as your output.
         return multivariate_samples
 
 

From 2c494c90e7889c928328de0285f422b7dbfa5cdd Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Wed, 30 Dec 2020 18:08:34 -0800
Subject: [PATCH 07/15] Fix rebase issue

---
 pymc3/tests/test_distributions_random.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py
index 0dfcf5c57e..d9c986db47 100644
--- a/pymc3/tests/test_distributions_random.py
+++ b/pymc3/tests/test_distributions_random.py
@@ -310,12 +310,6 @@ class TestGaussianRandomWalk(BaseTestCases.BaseTestCase):
     default_shape = (1,)
 
 
-class TestMvGaussianRandomWalk(BaseTestCases.BaseTestCase):
-    distribution = pm.MvGaussianRandomWalk
-    params = {"mu": np.array([1.0, 0.0]), "cov": np.array([[1.0, 0.0], [0.0, 2.0]])}
-    default_shape = (10, 2)
-
-
 class TestNormal(BaseTestCases.BaseTestCase):
     distribution = pm.Normal
     params = {"mu": 0.0, "tau": 1.0}
@@ -1726,3 +1720,17 @@ def test_matrix_normal_random_with_random_variables():
         prior = pm.sample_prior_predictive(2)
 
     assert prior["mu"].shape == (2, D, K)
+
+
+class TestMvGaussianRandomWalk(SeededTest):
+    @pytest.mark.parametrize(
+        ["sample_shape", "dist_shape", "mu_shape", "param"],
+        generate_shapes(include_params=True),
+        ids=str,
+    )
+    def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param):
+        dist = pm.MvGaussianRandomWalk.dist(
+            mu=np.ones(mu_shape), **{param: np.eye(3)}, shape=dist_shape
+        )
+        output_shape = to_tuple(sample_shape) + dist_shape
+        assert dist.random(size=sample_shape).shape == output_shape

From 4da787ea399569857b9b7ecd44fec4a84ef8f8cf Mon Sep 17 00:00:00 2001
From: Andrew Wray <46057140+awray3@users.noreply.github.com>
Date: Wed, 30 Dec 2020 09:16:28 -0800
Subject: [PATCH 08/15] Update pymc3/distributions/timeseries.py

Remove equality since the shape parameter cannot have more than 2 dimensions.

Co-authored-by: Sayam Kumar <sayamkumar753@yahoo.in>
---
 pymc3/distributions/timeseries.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index f2ec70d765..8a0397aaed 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -510,8 +510,7 @@ def random(self, point=None, size=None):
         size = to_tuple(size)
         multivariate_samples = self.innov.random(point=point, size=size)
         # this has shape (size, self.shape)
-
-        if len(self.shape) >= 2:
+        if len(self.shape) == 2
             # have time dimension in first slot of shape. Therefore the time
             # component can be accessed with the index equal to the length of size.
             time_axis = len(size)

From e2f71edbf945fad98bf09b4a00494a9174a599c3 Mon Sep 17 00:00:00 2001
From: Andrew Wray <46057140+awray3@users.noreply.github.com>
Date: Wed, 30 Dec 2020 09:16:48 -0800
Subject: [PATCH 09/15] Update comment pymc3/distributions/timeseries.py

Co-authored-by: Sayam Kumar <sayamkumar753@yahoo.in>
---
 pymc3/distributions/timeseries.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 8a0397aaed..0595bc5a61 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -526,7 +526,7 @@ def random(self, point=None, size=None):
                 # size was passed as None
                 multivariate_samples = multivariate_samples - multivariate_samples[0]
 
-        # if the above if statement fails, then only a spatial dimension was passed in for self.shape.
+        # if the above statement fails, then only a spatial dimension was passed in for self.shape.
         # Therefore don't subtract off the initial value since otherwise you get all zeros
         # as your output.
         return multivariate_samples

From f142842f16df3f9928a438c60921a4933ab4704c Mon Sep 17 00:00:00 2001
From: Andrew Wray <46057140+awray3@users.noreply.github.com>
Date: Wed, 30 Dec 2020 09:17:51 -0800
Subject: [PATCH 10/15] Update pymc3/distributions/timeseries.py

Adds a more explicit conditional on the time_axis statement.

Co-authored-by: Sayam Kumar <sayamkumar753@yahoo.in>
---
 pymc3/distributions/timeseries.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 0595bc5a61..76c5785fd7 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -515,8 +515,7 @@ def random(self, point=None, size=None):
             # component can be accessed with the index equal to the length of size.
             time_axis = len(size)
             multivariate_samples = multivariate_samples.cumsum(axis=time_axis)
-
-            if time_axis:
+            if time_axis != 0:
                 # this for loop covers the case where size is a tuple
                 for idx in np.ndindex(size):
                     multivariate_samples[idx] = (

From f3ec9ca371323b0bb90ea42138454afe9114467d Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Thu, 31 Dec 2020 09:32:06 -0800
Subject: [PATCH 11/15] Fix syntax error

---
 pymc3/distributions/timeseries.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 76c5785fd7..850ef12ad2 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -510,7 +510,7 @@ def random(self, point=None, size=None):
         size = to_tuple(size)
         multivariate_samples = self.innov.random(point=point, size=size)
         # this has shape (size, self.shape)
-        if len(self.shape) == 2
+        if len(self.shape) == 2:
             # have time dimension in first slot of shape. Therefore the time
             # component can be accessed with the index equal to the length of size.
             time_axis = len(size)

From 4b6bb800c5dd4d36c7344492578a84a7324a5d35 Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Thu, 31 Dec 2020 09:57:20 -0800
Subject: [PATCH 12/15] Add tests for MvGaussianRandomWalk with RV params

Added tests for both chol and cov to be random variables.
---
 pymc3/tests/test_distributions_random.py | 36 ++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py
index d9c986db47..5e735b6aed 100644
--- a/pymc3/tests/test_distributions_random.py
+++ b/pymc3/tests/test_distributions_random.py
@@ -1734,3 +1734,39 @@ def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param):
         )
         output_shape = to_tuple(sample_shape) + dist_shape
         assert dist.random(size=sample_shape).shape == output_shape
+
+    @pytest.mark.xfail
+    @pytest.mark.parametrize(
+        ["sample_shape", "dist_shape", "mu_shape"],
+        generate_shapes(include_params=False),
+        ids=str,
+    )
+    def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape):
+        with pm.Model() as model:
+            mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape)
+            sd_dist = pm.Exponential.dist(1.0, shape=3)
+            chol, corr, stds = pm.LKJCholeskyCov(
+                "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True
+            )
+            mv = pm.MvGaussianRandomWalk("mv", mu, chol=chol, shape=dist_shape)
+            prior = pm.sample_prior_predictive(samples=sample_shape)
+
+        assert prior["mv"].shape == to_tuple(sample_shape) + dist_shape
+
+    @pytest.mark.xfail
+    @pytest.mark.parametrize(
+        ["sample_shape", "dist_shape", "mu_shape"],
+        generate_shapes(include_params=False),
+        ids=str,
+    )
+    def test_with_cov_rv(self, sample_shape, dist_shape, mu_shape):
+        with pm.Model() as model:
+            mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape)
+            sd_dist = pm.Exponential.dist(1.0, shape=3)
+            chol, corr, stds = pm.LKJCholeskyCov(
+                "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True
+            )
+            mv = pm.MvGaussianRandomWalk("mv", mu, cov=pm.math.dot(chol, chol.T), shape=dist_shape)
+            prior = pm.sample_prior_predictive(samples=sample_shape)
+
+        assert prior["mv"].shape == to_tuple(sample_shape) + dist_shape

From dbaed341de59030709094458e1018c8ec5e2f5c9 Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Thu, 31 Dec 2020 09:59:40 -0800
Subject: [PATCH 13/15] Improve comment clarity in MvNormal._random

Moves the brittle comment in MvRandom._random's docstring to the actual location in the code for clarity.
---
 pymc3/distributions/timeseries.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 850ef12ad2..7b7f34f2fb 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -275,7 +275,6 @@ def _random(self, sigma, mu, size, sample_shape):
         """Implement a Gaussian random walk as a cumulative sum of normals.
         axis = len(size) - 1 denotes the axis along which cumulative sum would be calculated.
         This might need to be corrected in future when issue #4010 is fixed.
-        Lines 291-295 set the starting point of each instance of the random walk to 0."
         """
         if size[len(sample_shape)] == sample_shape:
             axis = len(sample_shape)
@@ -284,6 +283,8 @@ def _random(self, sigma, mu, size, sample_shape):
         rv = stats.norm(mu, sigma)
         data = rv.rvs(size).cumsum(axis=axis)
         data = np.array(data)
+
+        # the following lines center the random walk to start at the origin.
         if len(data.shape) > 1:
             for i in range(data.shape[0]):
                 data[i] = data[i] - data[i][0]

From 9dc5b7cfacba5c177473f4ab58737581dc208d54 Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Sat, 2 Jan 2021 21:52:22 -0800
Subject: [PATCH 14/15] Update docstring formatting

---
 RELEASE-NOTES.md                  |  1 +
 pymc3/distributions/timeseries.py | 26 ++++++++++++--------------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index ebe696e824..b1172869d8 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -20,6 +20,7 @@ It also brings some dreadfully awaited fixes, so be sure to go through the chang
 - `OrderedProbit` distribution added (see [#4232](https://github.com/pymc-devs/pymc3/pull/4232)).
 - `plot_posterior_predictive_glm` now works with `arviz.InferenceData` as well (see [#4234](https://github.com/pymc-devs/pymc3/pull/4234))
 - Add `logcdf` method to all univariate discrete distributions (see [#4387](https://github.com/pymc-devs/pymc3/pull/4387)).
+- Add `random` method to `MvGaussianRandomWalk` (see [#4388](https://github.com/pymc-devs/pymc3/issues/4337))
 
 ### Maintenance
 - Fixed bug whereby partial traces returns after keyboard interrupt during parallel sampling had fewer draws than would've been available [#4318](https://github.com/pymc-devs/pymc3/pull/4318)
diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py
index 7b7f34f2fb..e3e1aa15bc 100644
--- a/pymc3/distributions/timeseries.py
+++ b/pymc3/distributions/timeseries.py
@@ -485,24 +485,22 @@ def random(self, point=None, size=None):
 
         Examples
         -------
-        Create one sample from a 2-dimensional Gaussian random walk with 10 timesteps::
+        .. code-block:: python
 
-            mu = np.array([1.0, 0.0])
-            cov = np.array([[1.0, 0.0], [0.0, 2.0]])
-            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random()
+            with pm.Model():
+                mu = np.array([1.0, 0.0])
+                cov = np.array([[1.0, 0.0],
+                                [0.0, 2.0]])
 
-        Create three samples from a 2-dimensional Gaussian random walk with 10 timesteps::
+                # draw one sample from a 2-dimensional Gaussian random walk with 10 timesteps
+                sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random()
 
-            mu = np.array([1.0, 0.0])
-            cov = np.array([[1.0, 0.0], [0.0, 2.0]])
-            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=3)
+                # draw three samples from a 2-dimensional Gaussian random walk with 10 timesteps
+                sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=3)
 
-        Create four samples from a 2-dimensional Gaussian random walk with 10
-        timesteps, indexed with a (2, 2) array::
-
-            mu = np.array([1.0, 0.0])
-            cov = np.array([[1.0, 0.0], [0.0, 2.0]])
-            sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=(2, 2))
+                # draw four samples from a 2-dimensional Gaussian random walk with 10 timesteps,
+                # indexed with a (2, 2) array
+                sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=(2, 2))
         """
 
         # for each draw specified by the size input, we need to draw time_steps many

From b67b25c6366efc13695609dba1f5428419b973e5 Mon Sep 17 00:00:00 2001
From: Andrew Wray <andrewwray2718@gmail.com>
Date: Sat, 2 Jan 2021 21:57:52 -0800
Subject: [PATCH 15/15] Update RELEASE-NOTES.md

---
 RELEASE-NOTES.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index b1172869d8..a313d6cc33 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -20,7 +20,7 @@ It also brings some dreadfully awaited fixes, so be sure to go through the chang
 - `OrderedProbit` distribution added (see [#4232](https://github.com/pymc-devs/pymc3/pull/4232)).
 - `plot_posterior_predictive_glm` now works with `arviz.InferenceData` as well (see [#4234](https://github.com/pymc-devs/pymc3/pull/4234))
 - Add `logcdf` method to all univariate discrete distributions (see [#4387](https://github.com/pymc-devs/pymc3/pull/4387)).
-- Add `random` method to `MvGaussianRandomWalk` (see [#4388](https://github.com/pymc-devs/pymc3/issues/4337))
+- Add `random` method to `MvGaussianRandomWalk` (see [#4388](https://github.com/pymc-devs/pymc3/pull/4388))
 
 ### Maintenance
 - Fixed bug whereby partial traces returns after keyboard interrupt during parallel sampling had fewer draws than would've been available [#4318](https://github.com/pymc-devs/pymc3/pull/4318)