Skip to content

Commit fe33fe5

Browse files
committed
Add example to random docstring
1 parent 28591b4 commit fe33fe5

File tree

1 file changed

+26
-14
lines changed

1 file changed

+26
-14
lines changed

pymc3/distributions/timeseries.py

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -460,7 +460,6 @@ def logp(self, x):
460460
def _distr_parameters_for_repr(self):
461461
return ["mu", "cov"]
462462

463-
464463
def random(self, point=None, size=None):
465464
"""
466465
Draw random values from MvGaussianRandomWalk.
@@ -477,48 +476,61 @@ def random(self, point=None, size=None):
477476
Returns
478477
-------
479478
array
479+
480+
481+
Examples
482+
-------
483+
Create one sample from a 2-dimensional Gaussian random walk with 10 timesteps::
484+
485+
mu = np.array([1.0, 0.0])
486+
cov = np.array([[1.0, 0.0], [0.0, 2.0]])
487+
sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=1)
480488
"""
481489

482-
param_attribute = getattr(self.innov, "chol_cov" if self.innov._cov_type == "chol" else self.innov._cov_type)
483-
mu, param = distribution.draw_values([self.innov.mu, param_attribute], point=point, size=size)
490+
param_attribute = getattr(
491+
self.innov, "chol_cov" if self.innov._cov_type == "chol" else self.innov._cov_type
492+
)
493+
mu, param = distribution.draw_values(
494+
[self.innov.mu, param_attribute], point=point, size=size
495+
)
484496
return distribution.generate_samples(
485-
self._random,
497+
self._random,
486498
size=size,
487499
dist_shape=self.shape,
488500
not_broadcast_kwargs={
489501
"sample_shape": to_tuple(size),
490502
"param": param,
491503
"mu": mu,
492-
"cov_type": self.innov._cov_type
493-
}
504+
"cov_type": self.innov._cov_type,
505+
},
494506
)
495507

496-
def _random(self, mu, param, size, sample_shape, cov_type):
508+
def _random(self, mu, param, size, sample_shape, cov_type):
497509
"""
498510
Implements the multivariate Gaussian random walk as a cumulative
499511
sum of i.i.d. multivariate Gaussians.
500512
Assumes that
501513
size is of the form (samples, time, dims).
502-
"""
514+
"""
503515

504516
if cov_type == "chol":
505-
cov = np.matmul(param, param.transpose())
517+
cov = np.matmul(param, param.transpose())
506518
elif cov_type == "tau":
507-
cov = np.linalg.inv(param)
519+
cov = np.linalg.inv(param)
508520
else:
509521
cov = param
510522

511523
# time axis comes after the sample axis
512-
time_axis = len(sample_shape)
524+
time_axis = len(sample_shape)
513525

514526
# spatial axis is last
515-
spatial_axis = -1
527+
spatial_axis = -1
516528

517-
rv = stats.multivariate_normal(mean=mu, cov=cov)
529+
rv = stats.multivariate_normal(mean=mu, cov=cov)
518530

519531
# only feed in sample and time dimensions since stats.multivariate_normal
520532
# automatically adds back in the spatial dimensions to the end when it samples.
521-
data = rv.rvs(size[:spatial_axis]).cumsum(axis=time_axis)
533+
data = rv.rvs(size[:spatial_axis]).cumsum(axis=time_axis)
522534

523535
# shift the walk to start at zero
524536
if len(data.shape) > 2:

0 commit comments

Comments
 (0)