Best Way to Handle a for Loop in pymc4

Questions v5
, ,
1

Hello everyone,

I’m currently trying to do some Bayesian inference to do parameters estimation on a model. Part of this model (the v variable below) is an iterative map: “v_{t+1} = delta*v_t + gamma”. This iterative map gets plugged into a deterministic equation (shares in the code below). I am trying to obtain parameters alpha/beta/delta (gamma is a known value) in my model. What would be the best way for me to handle this, since I can’t think of any way to get that iterative map value without some sort of for loop? Right now I get the error: “TypeError: Unsupported dtype for TensorType: object”.

Also any other feedback on what I might be doing wrong is appreciated, I am very new to Bayesian inference and pymc.

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import arviz as az

tweet_model = pm.Model()
with tweet_model:
    alpha = pm.Uniform("alpha", lower=-30, upper=700)
    beta  = pm.LogNormal("beta", mu=0, sigma=2)
    delta = pm.Beta("delta", alpha=5, beta=1)
    
    v = np.array([0])
    for i in range(first_index['index']):
        v = np.append(v, delta*v + data.gamma())
        
    shares = alpha*(pm.math.exp(beta*v) - 1/(v+1))
    
    Y_obs = pm.Normal("Y_obs", mu=shares, sigma=1, observed=Y) 
    trace_g = pm.sample(2000, tune=1000, cores=2)

az.summary(trace_g)
chain_count =  trace_g.get_values('delta').shape[0]
y_pred_g = pm.sample_posterior_predictive(trace_g, samples=chain_count, model=model_g)
data_spp = az.from_pymc3(trace=trace_g, posterior_predictive=y_pred_g)
1 Like
2

You need a scan: scan – Looping in Aesara — Aesara 2.7.9+40.ge0d918074.dirty documentation

1 Like
3

Thanks! Can I scan just for the v term and then plug that into the shares equation?

4

Yup

5

Sorry, I’m struggling trying to understand scan. I’m getting this error: ValueError: When compiling the inner function of scan the following error has been encountered: The initial state (outputs_info in scan nomenclature) of variable IncSubtensor{Set;:int64:}.0 (argument number 0) has dtype int64, while the result of the inner function (fn) has dtype float64. This can happen if the inner function of scan results in an upcast or downcast.

I’m not 100% sure how to do my initial values. I’m just trying to set the initial value to 0 and then run that for loop. Lmk if you’d prefer this as a separate question since it’s a different thing.

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import aesara
import aesara.tensor as at

tweet_model = pm.Model()
with tweet_model:
    alpha = pm.Uniform("alpha", lower=-30, upper=700)
    beta  = pm.LogNormal("beta", mu=0, sigma=2)
    delta = pm.Beta("delta", alpha=5, beta=1)
    
    k = at.iscalar("k")
    v = at.iscalar("v")
        
    outputs_info = at.as_tensor_variable(np.asarray(0, first_index['index']))
    result, updates = aesara.scan(fn=lambda prior_result: delta*prior_result * data.gamma(),
                              outputs_info=outputs_info,
                              n_steps=k)
    
    virality = aesara.function(inputs=[v,k], outputs=result, updates=updates)
    shares = alpha*(pm.math.exp(beta*virality) - 1/(virality+1))
    
    Y_obs = pm.Normal("Y_obs", mu=shares, sigma=1, observed=Y) 
    trace_g = pm.sample(2000, tune=1000, cores=2)
6

For your specific error, you need to explicitly set the dtype of the values passed to outputs_info, i.e. at.as_tensor_variable(np.asarray(0, dtype='float64')). Aesara is very strict about dtypes and shapes, and this is doubly true for scan. As the error explains, it will not permit the result of the scan operation (the ‘inner function’) to be of a different datatype than the first value that “kicks off” the loop (outputs_info). I find it good to always be extra explicit in my code when working with scan.

One other note: you shouldn’t compile the aesara function yourself. Instead, let PyMC compile the entire computational graph by itself. In this case, it seems that “result” is a symbolic representation of “virality”, so you can use that directly in your deterministic function. If delta * prior_result * data.gamma() is a scalar, result is k x 1 vector. If it’s a vector, then result is k x v matrix (assuming k = n_rows and v = n_cols in your data)

1 Like
7

Hey thanks for the help! The first part definitely makes sense, but I’m a little confused about the second half. I definitely want a kx1 vector here. Are you saying instead of calling aesara.function, I should just use “results” in my deterministic function? When exactly would I call the function then in order to set my initial values and number of steps? I switched to this code and I’m getting “TypeError: ufunc ‘isnan’ not supported for the input types”.

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import aesara
import aesara.tensor as at

tweet_model = pm.Model()
with tweet_model:
alpha = pm.Uniform("alpha", lower=-30, upper=700)
beta = pm.LogNormal("beta", mu=0, sigma=2)
delta = pm.Beta("delta", alpha=5, beta=1)

k = at.iscalar("k")
v = at.iscalar("v")

outputs_info = at.as_tensor_variable(np.asarray(0, dtype='float64'))
result, updates = aesara.scan(fn=lambda prior_result: delta*prior_result + data.gamma(),
outputs_info=outputs_info,
n_steps=k)

# virality = aesara.function(inputs=[v,k], outputs=result, updates=updates)
shares = alpha*(pm.math.exp(beta*result) - 1/(result+1))

Y_obs = pm.Normal("Y_obs", mu=shares, sigma=1, observed=Y)
trace_g = pm.sample(2000, tune=1000, cores=2)
8

In general you don’t want to do any intermediate compilation of Aesara functions, because Aesara does lots of optimizations and graph re-writes at compile time to make things run faster. If you compile in chunks, you potentially lose out on optimizations between the chunks.

PyMC will compile the entire model, including your scan, and call it internally when it does its MCMC voodoo. So yes, you want to use result directly in your deterministic function. You can look back over the scan tutorial link to get a good sense of exactly what scan gives you back. It’s a stack of scan step results from t=1 to t=T, as a symbolic tensor (note that the value defined in outputs_info, t=0, is NOT included!), with the “batch” or “step” dimension as the first dimension. So if your function returns a scalar, and you have k steps, you will get a k x 1 vector out. You can continue to operate on it like any other symbolic tensor.

Anyway, a consequence of letting PyMC handle it is that the iscalar values you defined will not be included anywhere, so you need to replace these with the actual values you want to be used for your model. If you need these to be variable, you could wrap the model creation in a function, like:

def build_tweet_model(data, k, v, x0 = 0):
    tweet_model = pm.Model()
    
    with tweet_model:
        alpha = pm.Uniform("alpha", lower=-30, upper=700)
        beta = pm.LogNormal("beta", mu=0, sigma=2)
        delta = pm.Beta("delta", alpha=5, beta=1)

        outputs_info = at.as_tensor_variable(np.asarray(x0, dtype='float64'))

        def inner_func(prior_result):
            return delta * prior_result + data.gamma()

        virality, updates = aesara.scan(fn=inner_func,
                                      outputs_info=outputs_info,
                                      n_steps=k)

        shares = alpha*(pm.math.exp(beta*virality) - 1/(virality+1))
        
    return tweet_model

Now if you wanted to fit this model to several different datasets with different values of k and v, it’s quite easy.

Or, if you wanted to estimate any of these values, you would just put prior over them. You could, for example, define outputs_info = pm.Normal('x0', mu=0, sigma=1), and let the model estimate the initial state. You can’t do that for the number of steps, though, because the shape of the scan result needs to be static. You will need to define that up front.

As for the specific error, it seems like np.isnan is being called somewhere else in your code, but it’s not in this snippet, so I can’t say where your error is occurring. Perhaps something to do with data.gamma()? I couldn’t say more without knowing more.

9

Hey Jesse, thanks for your help! I finally had time to get back to work on this. Your code sample is really helpful, and your explanation makes a lot of sense. Here’s where I landed:

import pandas as pd
import matplotlib.pyplot as plt
import pymc as pm
import numpy as np
import aesara
import aesara.tensor as at

def build_tweet_model(data, k, x0 = 0):
    tweet_model = pm.Model()
    
    with tweet_model:
        alpha = pm.Uniform("alpha", lower=-30, upper=700)
        beta = pm.LogNormal("beta", mu=0, sigma=2)
        delta = pm.Beta("delta", alpha=5, beta=1)

        outputs_info = at.as_tensor_variable(np.asarray(x0, dtype='float64'))

        def inner_func(prior_result):
            return delta * prior_result + data.gamma()

        virality, updates = aesara.scan(fn=inner_func,
                                      outputs_info=outputs_info,
                                      n_steps=k)

        shares = alpha*(pm.math.exp(beta*virality) - 1/(virality+1))
        Y_obs = pm.Normal("Y_obs", mu=shares, sigma=1, observed=time_series)

    return tweet_model


tweet_model = build_tweet_model(time_series, len(time_series))
with tweet_model:
    trace_g = pm.sample(2000, tune=1000, cores=2)

The symbolic stuff is definitely making more sense, I did a lot more reading on theano/aesara/other frameworks. In the last 3 lines, is this the proper way to call this function? In my mind it’s just returning a pymc model so I should be able to use it like a normal model when I sample. I’m also curious if adding the observable inside build_tweet_model() is the right idea.

However I’m getting a different error that might need it’s own thread: “TypeError: <class ‘numpy.typing._dtype_like._SupportsDType’> is not a generic class”. The traceback is showing that this originates in the line import pymc as pm. Google searching showed me this: python 3.x - <class 'numpy.typing._dtype_like._SupportsDType'> is not a generic class when importing the plotly.express library - Stack Overflow, but upgrading numpy on my end didn’t help.

Thanks,
Trevor Crupi

10

You’re exactly right with how to call the function. And yeah, supplying the data is good. That way you can re-use the function on different time series. Based on what you have here, I’d even remove the k argument and just define k = len(data) inside the function (as long as k always equals len(data)).

I’m sorry to say I can’t help with the other issue. I can only doing a fresh install into a new environment following the PyMC installation instructions on the wiki (right hand side, you can find instructions for linux/mac/windows) to test if it’s an environment problem or something else.

Also definitely switch from conda to mamba and use that to do the reinstall (if you haven’t already). You can just swap all the conda commands to mamba commands in the install instructions. I think the install guides used to explicitly recommend using mamba, but I don’t see that on the page now.