Variable Lag and slicing

Trying to estimate the lag and the way x affects y. In essence, the model should look like:

lag = pm.DiscreteUniform("lag", 1, 12)
p = pm.Deterministic("p", 1.0 / (1. + np.exp(beta * x[12-lag:-lag] + alpha)))
y = pm.Normal("y", gamma * y[11:-1] + gamma2 * p, sigma=sigma, observed=y[12:])

Unfortunately, Python doesn’t like having a PyMC variable as an index:

TypeError: slice indices must be integers or None or have an index method

Is there a method I should be aware of for slicing series based on a PyMC variable?

If you want to use pytensor tensors as indexers, you need to be indexing into a pytensor tensor. What are the datatypes of x and y in your example code?

The types of x and y are numpy.ndarray

Convert x to a TensorVariable first with x = pt.as_tensor(x) and then you can index it with pymc variables.

import pytensor.tensor as pt

You can also use pm.ConstantData to wrap your x as a tensor if you prefer

2 Likes

I rewrote my example so that’s reproducible:

n_obs = 10000
true_lag = 2
lead = np.random.normal(loc=0, scale=0.5, size=n_obs + true_lag)
y_obs = -2.0 * lead[:-true_lag] + np.random.normal(loc=0, scale=0.4, size=n_obs)
lead = lead[true_lag:]

with pm.Model() as model:
    lead = pm.ConstantData("lead", value=lead)
    alpha = pm.Normal("alpha")
    lag = pm.DiscreteUniform("lag", 1, 3)
    p = pm.Deterministic("p", pm.math.switch(lag <= 1, lead[11:-1],
                                             pm.math.switch(lag <= 2, lead[10:-2], lead[9:-3])))
    sigma = pm.HalfNormal("sigma")
    y = pm.Normal("y", alpha * p, sigma=sigma, observed=y_obs[12:])
    trace = pm.sample(5000, tune=1000)

When I replace the eyesore of the pm.math.switch onion with this:

p = pm.Deterministic("p", lead[12-lag:-lag])

I see the following error:

ParallelSamplingError: Chain 3 failed with: Input dimension mismatch. One other input has shape[0] = 0, but input[2].shape[0] = 9988.
Apply node that caused the error: Elemwise{Composite}(InplaceDimShuffle{x}.0, Subtensor{int64:int64:}.0, y{[-1.432294...12875804]}, InplaceDimShuffle{x}.0, TensorConstant{(1,) of -0.5}, TensorConstant{(1,) of 0...5332046727}, Elemwise{log,no_inplace}.0, Elemwise{gt,no_inplace}.0, TensorConstant{(1,) of -inf})
Toposort index: 23
Inputs types: [TensorType(float64, (1,)), TensorType(float64, (?,)), TensorType(float64, (9988,)), TensorType(float64, (1,)), TensorType(float64, (1,)), TensorType(float64, (1,)), TensorType(float64, (1,)), TensorType(bool, (1,)), TensorType(float32, (1,))]
Inputs shapes: [(1,), (0,), (9988,), (1,), (1,), (1,), (1,), (1,), (1,)]
Inputs strides: [(8,), (8,), (8,), (8,), (8,), (8,), (8,), (1,), (4,)]
Inputs values: [array([0.]), array([], dtype=float64), 'not shown', array([1.]), array([-0.5]), array([0.91893853]), array([0.]), array([ True]), array([-inf], dtype=float32)]
Outputs clients: [[Sum{acc_dtype=float64}(sigma > 0)]]

It’s quite a strange error, because:

  1. model.debug() finds no errors
  2. pm.draw works as expected
  3. the compiled logp and dlopg functions work as expected
  4. pm.find_MAP works as expected

So you really got me. If you use pytensor.dprint and look at the compiled logp graph, the input with zero shape (TensorType(float64, (?,))) is p. My guess is that the data in p is changing, which causes problems for the shape inference of the object. But that doesn’t explain why it only fails when you go to sample.

A priori, the switch way is more correct anyway – nested switches are not uncommon in the codebase. You could beautify it a bit by importing switch from pm.math, and ditching the deterministic (do you need to inspect the trace for p anyway?). But @ricardoV94 for an expert opinion of why the slicing way doesn’t work?

Another approach would be to use an AR(p), with p quite large, together with shrinkage priors for automatic order selection. This is common in the bVAR literature, but no reason why you couldn’t do something similar here.

1 Like

The problem has to do with how discrete proposals are attempted by PyMC. The metropolis sampler will sometimes propose lags of 0 or -1, and check the logp delta to decide whether to accept or reject that proposal. Under “normal” models, the logp of the proposal will be -inf, and it will be promptly rejected.

But in this case it first leads p to have the wrong shape and the model logp to fail to evaluate at all. Tthe acceptance ratio takes the whole model logp in consideration, even if in this case the prior probability would be enough to rule out such proposal.

You can observe something similar by wrapping lag in a PyTensor assert:

import pymc as pm
import numpy as np
from pytensor.raise_op import assert_op

n_obs = 10000
true_lag = 2
lead = np.random.normal(loc=0, scale=0.5, size=n_obs + true_lag)
y_obs = -2.0 * lead[:-true_lag] + np.random.normal(loc=0, scale=0.4, size=n_obs)
lead = lead[true_lag:]

with pm.Model() as model:
    lead = pm.ConstantData("lead", value=lead)
    alpha = pm.Normal("alpha")
    lag = pm.DiscreteUniform("lag", 1, 3)
    lag = assert_op(lag, lag > 0)
    p = pm.Deterministic("p", lead[12-lag:-lag])
    sigma = pm.HalfNormal("sigma")
    y = pm.Normal("y", alpha * p, sigma=sigma)
    
    trace = pm.sample(chains=1)
AssertionError: PyTensor Assert failed!
Apply node that caused the error: Assert{msg=PyTensor Assert failed!}(Reshape{0}.0, Elemwise{gt,no_inplace}.0)
Toposort index: 21
Inputs types: [TensorType(int64, ()), TensorType(bool, ())]
Inputs shapes: [(), ()]
Inputs strides: [(), ()]
Inputs values: [array(0), array(False)]

Check the Assert failing with the proposed value of 0 array(0), which would otherwise lead to a weird shape in the indexing of lead.

np.arange(100)[12-0:-0]  # [ ]

A solution is to build your model to be robust to stupid proposals:

import pymc as pm
import numpy as np
import pytensor.tensor as pt

n_obs = 10000
true_lag = 2
lead = np.random.normal(loc=0, scale=0.5, size=n_obs + true_lag)
y_obs = -2.0 * lead[:-true_lag] + np.random.normal(loc=0, scale=0.4, size=n_obs)
lead = lead[true_lag:]

with pm.Model() as model:
    lead = pm.ConstantData("lead", value=lead)
    alpha = pm.Normal("alpha")
    lag = pm.DiscreteUniform("lag", 1, 3)
    p = pm.Deterministic("p", lead[12 - pt.max((1, lag)):-pt.max((1, lag))])
    sigma = pm.HalfNormal("sigma")
    y = pm.Normal("y", alpha * p, sigma=sigma)

    trace = pm.sample(chains=1)

We have considered other options, such as implementing specialized samplers for bounded discrete variables or even discrete variable transformations: Implement transforms for discrete variables by ricardoV94 · Pull Request #6102 · pymc-devs/pymc · GitHub

Actually if you had used a Categorical prior you wouldn’t have seen this issue, because it uses a less dumb sampler.

1 Like

By the way, this type of structural discrete priors are a good candidate to marginalization: multiple_changepoint_marginalization.ipynb · GitHub, which also has the side-benefit that impossible values won’t be considered.

1 Like

@ricardoV94 Thank you very much for the explanation and the link. I am trying to run the changepoint marginalization example locally, and here’s an error that I am seeing:

NotImplementedError: Cannot convert Elemwise{mul,no_inplace}.0 to a tensor variable.

here are the relevan lines from my conda list:

aesara                    2.8.12                   pypi_0    pypi
numpy                     1.24.3          py311hb8f3215_0    conda-forge
pymc                      5.3.1                    pypi_0    pypi
pymc-experimental         0+untagged.107.g2e8505a          pypi_0    pypi
pytensor                  2.11.2                   pypi_0    pypi
python                    3.11.3          h1456518_0_cpython    conda-forge
scikit-learn              1.2.2                    pypi_0    pypi
scipy                     1.10.1          py311h0bcca16_1    conda-forge

Make sure you don’t mix Aesara and PyTensor. PyMC 5 should use only PyTensor

1 Like

Most helpful, thank you! Replacing import aesara.tensor as at with import pytensor.tensor as pt solved the issue.