Model with a loop - scan conversion

Questions v5
1

First time actually posting on a coding forum, so please let me know if anything needs clarified or edited.

I have never used PYMC before, and I am attempting to calibrate a model, however, the model contains a for loop that I am struggling to convert into a scan. I believe the model should look like the below code block.

with pm.Model() as model:
    #priors 
    alpha = pm.Uniform("alpha", 0, 5)
    t_eq  = pm.Uniform("t_eq", -2, 2)
    s0    = pm.Uniform("s0", -150, 150)
    rho0  = pm.Uniform("rho0", -1, 1)
    sigma0 = pm.Uniform("sigma0", 0, 20)

    # sea level generation - Rahmstorf model
    sea_levels = [None]*n
    sea_levels[0] = s0

    for i in range(n-1):
        sea_levels[i+1] = sea_levels[i] + alpha*(temps[i] - t_eq)

    # AR1 covariance matrix generation
    sigma_ar_1 = (sigma**2) / (1 - rho0**2)
    H_mat = np.power(rho0, abs(np.arange(1, n+1).reshape(-1, 1) - np.arange(1, n+1)))

    cov_matrix_hay = (sigma_ar_1 * H_mat) + hay_stds_diag

    cov_matrix_hay += 0.00001*np.identity(cov_matrix_hay.shape[0])
    cov_matrix_hay = (1/2)*(cov_matrix_hay + cov_matrix_hay.T)

    # likelihood
    likelihood = pm.MvNormal(name = "sea_levels", \
                             mu = sea_levels, \
                             cov = cov_matrix_hay, \
                             observed = synth_levels)

There could be errors in the rest of the code as well, but I am currently trying to convert this section into a scan (as I saw recommended in other threads).

sea_levels = [None]*n
sea_levels[0] = s0

for i in range(n-1):
    sea_levels[i+1] = sea_levels[i] + alpha*(temps[i] - t_eq)

My current best attempt at creating an equivalent scan is…

temps = pt.dvector("temps")
k = pt.iscalar("k")

rahm_steps, updates = pytensor.scan(
                      fn = lambda previous_level, alpha, temps, t_eq: previous_level + alpha*(temps - t_eq),
                      outputs_info = s0,
                      sequences = temps,
                      non_sequences = [alpha, t_eq],
                      n_steps = k
                      )

rahmstorf = pytensor.function(inputs = [alpha, t_eq, s0, temps, k],outputs=rahm_steps, updates=updates)

rahmstorf(alpha, t_eq, s0, hay_temps, len(hay_temps))

However, when I run the following code…



with pm.Model() as model:
    alpha = pm.Normal("alpha", mu=2.5, sigma=0.5)
    t_eq  = pm.Normal("t_eq", mu=-1.5, sigma=0.5)
    s0    = pm.Normal("s0", mu=100, sigma=20)
    rho0  = pm.TruncatedNormal("rho0", mu=0.5, sigma=0.3, lower=-1, upper=1)
    sigma = pm.TruncatedNormal("sigma", mu=8, sigma=3, lower=0, upper=15)\

    temps = pt.dvector("temps")
    k = pt.iscalar("k")

    rahm_steps, updates = pytensor.scan(
                          fn = lambda previous_level, alpha, temps, t_eq: previous_level + alpha*(temps - t_eq),
                          outputs_info = s0,
                          sequences = temps,
                          non_sequences = [alpha, t_eq],
                          n_steps = k
    )

    rahmstorf = pytensor.function(inputs = [alpha, t_eq, s0, temps, k], outputs=rahm_steps, updates=updates)

    rahmstorf(alpha, t_eq, s0, hay_temps, len(hay_temps))

    part1 = (sigma**2) / (1 - rho0**2)
    part2 = np.power(rho0, abs(np.arange(1, n+1).reshape(-1, 1) - np.arange(1, n+1)))

    cov_matrix_hay = (part1 * part2) + hay_stds_diag

    cov_matrix_hay += 0.00001*np.identity(cov_matrix_hay.shape[0])
    cov_matrix_hay = (1/2)*(cov_matrix_hay + cov_matrix_hay.T)

    likelihood = pm.MvNormal(name = "sea_levels", \
                             mu = sea_levels, \
                             cov = cov_matrix_hay, \
                             observed = synth_levels)

I get this error…

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[10], line 29
     19 rahm_steps, updates = pytensor.scan(
     20                       fn = lambda previous_level, alpha, temps, t_eq: previous_level + alpha*(temps - t_eq),
     21                       outputs_info = s0,
   (...)     24                       n_steps = k
     25 )
     27 rahmstorf = pytensor.function(inputs = [alpha, t_eq, s0, temps, k], outputs=rahm_steps, updates=updates)
---> 29 rahmstorf(alpha, t_eq, s0, hay_temps, len(hay_temps))
     31 part1 = (sigma**2) / (1 - rho0**2)
     32 part2 = np.power(rho0, abs(np.arange(1, n+1).reshape(-1, 1) - np.arange(1, n+1)))

File c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\pytensor\compile\function\types.py:946, in Function.__call__(self, output_subset, *args, **kwargs)
    944 else:
    945     try:
--> 946         arg_container.storage[0] = arg_container.type.filter(
    947             arg,
    948             strict=arg_container.strict,
    949             allow_downcast=arg_container.allow_downcast,
    950         )
    952     except Exception as e:
    953         i = input_storage.index(arg_container)

File c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\pytensor\tensor\type.py:163, in TensorType.filter(self, data, strict, allow_downcast)
    160 # Explicit error message when one accidentally uses a Variable as
    161 # input (typical mistake, especially with shared variables).
    162 if isinstance(data, Variable):
--> 163     raise TypeError(
    164         "Expected an array-like object, but found a Variable: "
    165         "maybe you are trying to call a function on a (possibly "
    166         "shared) variable instead of a numeric array?"
    167     )
    169 if isinstance(data, np.memmap) and (data.dtype == self.numpy_dtype):
    170     # numpy.memmap is a "safe" subclass of ndarray,
    171     # so we can use it wherever we expect a base ndarray.
    172     # however, casting it would defeat the purpose of not
    173     # loading the whole data into memory
    174     pass

TypeError: Bad input argument with name "alpha" to pytensor function with name "C:\Users\Brocktree\AppData\Local\Temp\ipykernel_38248\101790853.py:27" at index 0 (0-based).  
Backtrace when that variable is created:

  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\IPython\core\async_helpers.py", line 128, in _pseudo_sync_runner
    coro.send(None)
  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\IPython\core\interactiveshell.py", line 3394, in run_cell_async
    has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\IPython\core\interactiveshell.py", line 3639, in run_ast_nodes
    if await self.run_code(code, result, async_=asy):
  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\IPython\core\interactiveshell.py", line 3699, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "C:\Users\Brocktree\AppData\Local\Temp\ipykernel_38248\101790853.py", line 2, in <module>
    alpha = pm.Normal("alpha", mu=2.5, sigma=0.5)
  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\pymc\distributions\distribution.py", line 505, in __new__
    rv_out = cls.dist(*args, **kwargs)
  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\pymc\distributions\continuous.py", line 495, in dist
    return super().dist([mu, sigma], **kwargs)
  File "c:\Users\Brocktree\anaconda3\envs\pymc\Lib\site-packages\pymc\distributions\distribution.py", line 574, in dist
    rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
Expected an array-like object, but found a Variable: maybe you are trying to call a function on a (possibly shared) variable instead of a numeric array?

Many thanks in advance.

2

Hey hey

You’re on the right track, but your mental model of what pytensor is is wrong. Amusingly, my very first post on this forum made the same mistake, so you’re in good company.

Pytensor works by making a graph of symbolic computation. When you do e.g. z = x + y, z is not a number, it’s a set of instructions describing how x, y, and z are related. You build a computation up bit by bit, using and re-using these computational graphs. When you call pytensor.function, you signal that you’re done build and you want to do actual numeric computation from the inputs to the outputs.

So, your proximate error is caused because you are passing symbolic values like alpha to your compiled function, which is done with symbolic stuff and is now ready for numeric values.

But the ultimately error is that you shouldn’t be doing anything with pytensor.function in the context of a pymc model. PyMC is going to compile some function itself, and it can’t see into or use your compiled function. But you don’t need this – instead, you need to keep everything symbolic from start to finish. For example:

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

n = 100
temp_data = np.random.normal(size=(n, ))

with pm.Model() as model:
    alpha = pm.Normal("alpha", mu=2.5, sigma=0.5)
    t_eq  = pm.Normal("t_eq", mu=-1.5, sigma=0.5)
    s0    = pm.Normal("s0", mu=100, sigma=20)
    rho0  = pm.TruncatedNormal("rho0", mu=0.5, sigma=0.3, lower=-1, upper=1)
    sigma = pm.TruncatedNormal("sigma", mu=8, sigma=3, lower=0, upper=15)

    temps = pm.Data('temps', temp_data)

    def scan_step(previous_level, alpha, temps, t_eq):
        return previous_level + alpha * (temps - t_eq)
    
    sea_levels = pytensor.scan(scan_step,
                               outputs_info = s0,
                               sequences = temps,
                               non_sequences = [alpha, t_eq],
                               return_updates=False)
    
    hay_stds_diag = pt.eye(n)
    
    sigma_ar_1 = sigma **2 / (1 - rho0 ** 2)
    powers = pt.arange(1, n+1)
    cov_matrix_hay = sigma_ar_1 * rho0 ** pt.abs(pt.sub.outer(powers, powers)) + hay_stds_diag
    
    sea_level_hat = pm.MvNormal(name='sea_levels',
                                mu=sea_levels,
                                cov=cov_matrix_hay)

Some other issues though:

You seem to be computing an AR(1) model twice, once in the mean, then again in the covariance. It seems like you’re mixing a model defined once as a generative model (the AR(1) scan) and a likelihood (the MvNormal with AR(1) covariance matrix)

Using a big covariance matrix to define the observed variable will end up being really slow compared to using pm.AR. You’re essentially making a Gaussian Process by hand. Using a GP would give you access to approximations that could speed things up (like HSGP)

Nitpick: What you call “likelihood” isn’t a likelihood. It’s a symbolic random variable that generates draws from your model. You can check that pm.draw(likelihood) does not give back evaluations of the likelihood function. This distribution will be used to compute the likelihood term during estimation, but it’s not there yet.

3

Thank you so much for the thorough response. The model is currently working for me and I will definitely look into the other issues you have pointed out. I wonder if the accidental “double counting” is one of the issues making it difficult to calibrate, because the noise is getting amplified. I will also definitely look into the built-in pm.AR function in-order to speed up the code (I am currently re-writing the model from some Julia code I have, where I am relying on the large covariance matrix, and it is indeed very slow).

Would you be able to expand a bit more on the likelihood bit? Going off of some of the example PYMC models, I assumed pymc.sample combined the priors and the “likelihood” behind the scenes in order to get at the posterior.

4

My point is that if you do likelihood.eval() you will not get a vector of likelihoods, you will get a random draw from your model for the distribution associated with your observed data.

Under the hood, pymc uses pytensor to take a generative graph (that takes input data and rng state and returns draws from distributions) and transforms it into a log probability graph (that takes input data, parameter values, and observed data and returns a scalar).

You can access the logp graph with model.logp(), and you can even see the graphical representations by doing likelihood.dprint() and model.logp().dprint(). Reading these tree representations of your computation take some getting used to, but at minimum you can convince yourself they’re not the same thing.

1 Like