Add test for 2D GaussianRamdomWalk logp

[ricardoV94]

Now that the GRW supports batched dimensions (#5298) we should add a test for its logp as well

[danhphan]

Hi @ricardoV94, I'm working on this task. Thanks.

[danhphan]

Hi @ricardoV94, what the 2D mean? Is this relate to this issue #4010?

I run the code here and it work, though not sure if the output shape is correct.

In [59]: with pm.Model() as m1:
    ...:     grw1 = pm.GaussianRandomWalk('grw1', mu=np.arange(4), shape=(3,4,), steps=10)
    ...: 
In [60]: grw1.shape.eval()
Out[60]: array([ 3, 11])
In [61]: grw1.ndim
Out[61]: 2

However, when sampling with pm.sample(), It will fail with shape mismatch error:

In [78]: with pm.Model() as m1:
    ...:     grw1 = pm.GaussianRandomWalk('grw1', mu=np.arange(4), shape=(3,4,), steps=10)
    ...:     trace = pm.sample()

Outputs:

ValueError: shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (3,) and arg 1 with shape (4,).
Apply node that caused the error: normal_rv{0, (0, 0), floatX, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F40F1FB1AC0>), TensorConstant{(1,) of 3}, TensorConstant{11}, TensorConstant{[0. 1. 2. 3.]}, TensorConstant{1.0})
Toposort index: 0
Inputs types: [RandomGeneratorType, TensorType(int64, (1,)), TensorType(int64, ()), TensorType(float64, (4,)), TensorType(float64, ())]
Inputs shapes: ['No shapes', (1,), (), (4,), ()]
Inputs strides: ['No strides', (8,), (), (8,), ()]
Inputs values: [Generator(PCG64) at 0x7F40F1FB1AC0, array([3]), array(11), array([0., 1., 2., 3.]), array(1.)]
Outputs clients: [['output'], [GaussianRandomWalk_rv{1, (0, 0, 0, 0), floatX, True}(RandomStateSharedVariable(<RandomState(PCG64) at 0x7F40ED983C40>), TensorConstant{(1,) of 3}, TensorConstant{11}, TensorConstant{[0. 1. 2. 3.]}, TensorConstant{1.0}, normal_rv{0, (0, 0), floatX, True}.out, TensorConstant{10})]]

I do other tests on the following model:

In [84]: with pm.Model() as m2:
    ...:     mu = pm.Normal("mu", 0, 1, shape=(3,2))
    ...:     grw2 = GaussianRandomWalk("grw2", _mu,  steps=100)
    ...: 
In [85]: mu.ndim
Out[85]: 2
In [86]: mu.shape.eval()
Out[86]: array([3, 2])
In [87]: grw2.ndim
Out[87]: 1
In [88]: grw2.shape.eval()
Out[88]: array([101])

Or this model:

In [95]: with pm.Model() as m3:
    ...:     mu = pm.Normal("mu", 0, 1, shape=(3,2))
    ...:     sigma = pm.HalfNormal("sigma", 5, shape=(3,2))
    ...:     grw3 = GaussianRandomWalk("grw3", _mu, sigma, steps=100)
In [96]: mu.ndim, sigma.ndim, grw3.ndim
Out[96]: (2, 2, 3)
In [97]: mu.shape.eval(), sigma.shape.eval(), grw3.shape.eval()
Out[97]: (array([3, 2]), array([3, 2]), array([  3,   2, 101]))

The grw2 has shape of [101], while qrw3 has shape of [ 3, 2, 101]. Both also have shape mismatch errors when do trace = pm.sample().

ValueError: shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (3, 2, 100) and arg 2 with shape (3, 2).
Apply node that caused the error: GaussianRandomWalk_rv{1, (0, 0, 0, 0), floatX, True}(RandomStateSharedVariable(<RandomState(PCG64) at 0x7F40EA76D340>), TensorConstant{[]}, TensorConstant{11}, mu, TensorConstant{(3, 2) of ..9999999999}, normal_rv{0, (0, 0), floatX, True}.out, TensorConstant{100})
Toposort index: 4
Inputs types: [RandomStateType, TensorType(int64, (0,)), TensorType(int64, ()), TensorType(float64, ()), TensorType(float64, (3, 2)), TensorType(float64, (None, None)), TensorType(int32, ())]
Inputs shapes: ['No shapes', (0,), (), (), (3, 2), (3, 2), ()]
Inputs strides: ['No strides', (8,), (), (), (16, 8), (16, 8), ()]
Inputs values: [RandomState(PCG64) at 0x7F40EA76D340, array([], dtype=int64), array(11), array(-1.13008785), 'not shown', 'not shown', array(100, dtype=int32)]
Outputs clients: [['output'], ['output']]

[ricardoV94]

Sorry for the delay. This new specification is incorrect, because the last entry of shape has to be the same as steps+1

In [59]: with pm.Model() as m1:
    ...:     grw1 = pm.GaussianRandomWalk('grw1', mu=np.arange(4), shape=(3,4,), steps=10)

It should be

grw1 = pm.GaussianRandomWalk('grw1', mu=np.arange(4), shape=(3,4,), steps=3)

[danhphan]

Hi thanks, I am going to do this soon, just need to spend sometimes to learn the Gaussian Random Walk model :)

[danhphan]

Hi @ricardoV94,
This one will also not work:

with pm.Model():
    grw1 = pm.GaussianRandomWalk('grw1', mu=np.arange(4), shape=(3,4,), steps=3)
grw1.shape.eval() # array([3, 4])
grw1.eval() # Failed

It seems due to mu only accept as a float. Hence, if we update mu as follow, it will work:

with pm.Model():
    grw1 = pm.GaussianRandomWalk('grw1', mu=1.0, shape=(3,4,), steps=3)
grw1.shape.eval() # array([3, 4])
grw1.eval()

Output:

array([[0.32473023, 0.80932398, 2.81870732, 2.39093408],
       [3.13770902, 5.14466547, 6.06865731, 7.71960022],
       [3.12229395, 3.67718623, 3.94715773, 6.07684834]])

By the way, just want to be clear what do you mean by 2D GaussianRamdomWalk?

My understanding is that 2D mean that the observed data has 2 dimensions (D=2). For example, this will work:

mu, sigma, steps, D = 2, 4, 1000, 2
obs = np.concatenate([np.zeros((1,D)), np.random.normal(mu, sigma, size=(steps, D))]).cumsum(axis=1)
print(obs.shape) # (1001, 2)

with pm.Model() as grw2D:
    _mu = pm.Uniform("mu", -10, 10)
    _sigma = pm.Uniform("sigma", 0, 10)

    obs_data = pm.MutableData("obs_data", obs)
    grw = pm.GaussianRandomWalk("grw", _mu, _sigma, steps=steps, observed=obs_data) # or shape=(D, steps+1)

    trace = pm.sample(chains=1)
    
recovered_mu = trace.posterior["mu"].mean()
recovered_sigma = trace.posterior["sigma"].mean()
print(recovered_mu, recovered_sigma)

Output:

<xarray.DataArray 'mu' ()> array(1.96290314) 
<xarray.DataArray 'sigma' ()> array(4.03082129)

Let's me know if this is acceptable to test 2D GaussianRamdomWalk.

Thanks!

[danhphan]

Also, I seem to have confusion on the shape parameter in the pm.GaussianRandomWalk.

# D=2, steps=1000
obs_data = pm.MutableData("obs_data", obs)
grw = pm.GaussianRandomWalk("grw", _mu, _sigma, shape=(D, steps+1), observed=obs_data)

Here, the grw has shape of (2, 1001), while the obs_data has shape of (1001, 2). It seems not consistent to me.

[danhphan]

Hi, my confusion on shape leads me to this issue #4552, which provides a lot insights.

I am kind of understand more on different kinds of shape (sample, batch, and event shape for distributions). This implementation makes sense to me:

pm.MvNormal(mu=[1,2,3], cov=np.eye(3), shape=(2, 3))
# is the same as
pm.MvNormal(mu=[1,2,3], cov=np.eye(3), shape=(2, ...))
# is the same as
pm.MvNormal(mu=[1,2,3], cov=np.eye(3), size=2)

I am trying to get my head around how the steps parameters in pm.GaussianRandomWalk will suitable into these kinds of shape :)

[ricardoV94]

steps is equivalent to shape[-1] - 1 in the case of the GaussianRandomWalk. This notebook might help you understand the different flavors of shapes in PyMC: https://docs.pymc.io/en/latest/learn/core_notebooks/dimensionality.html

[ricardoV94]

Closing this. After #6072 GRW now relies on Aeppl which is responsible for testing the right logp is generated.