MvNormal much slower in pymc5 than pymc3? – v5

Hi there! I have a 30-member ensemble of simulations at nt=100 time steps. Each ensemble member is assumed to be the combination of some central response \mu and underlying noise. The noise is uncorrelated across ensemble members and has temporal covariance matrix \Sigma. In this example, I assume AR(1) noise, but I would like to compare this to alternate models using different \Sigma (white noise, long-memory noise, etc). In pymc3, the following code takes ~200 seconds to run, but in pymc 5.3.0 it takes days. Any tips? Thank you very much.

with pm.Model() as from_pymc3:

    σ=pm.HalfNormal("σ",10)
    ϕ=pm.Uniform("ϕ",-1,1)
   
    
    #AR(1) covariance matrix
    δ=np.arange(nt)
    imj=np.abs(np.subtract.outer(δ,δ))
    cov=pt.pow(σ,2)/(1-pt.pow(ϕ,2))*pt.pow(ϕ,imj)

    chol=pt.slinalg.cholesky(cov)
    
  
    #Uncertain mean value across all ensemble members
    μ=pm.Normal("μ",0,10,shape=nt)
    
    F=pm.MvNormal("likelihood_hist",mu=μ,chol=chol,\
                             observed= ensemble)
    pymc3_trace=pm.sample()

The model sampled fine for me in 5.3.0, took about 3 minutes using the following data:


nt = 100
n = 30
innovations = np.random.normal(scale=10, size=(n, nt))
true_rho = np.random.uniform()
ensemble = np.zeros((n, nt))
for t in range(1, nt):
    ensemble[:, t] = true_rho * ensemble[:, t-1] + innovations[:, t]

Are you running it in a clean environment? How is your data generated? In my experience, when a model takes an eternity to run it’s a sign that some unexpected broadcasting is happening. In this case, I also wonder if there’s an environment problem from having multiple backends installed – it’s not recommended to have more than one of theano, aesara, or pytensor in a single environment.

Thank you! I tried creating a new environment using

conda create -c conda-forge -n new_env “pymc>=4”

and the problem persists. After 5 minutes, using the ensemble data above, the samples are 0.28% complete. Neither aesara or theano seems to be installed in this environment. I get the same results on my laptop (Mac Ventura, Apple M2) and a linux box. (And many thanks for your time and your kind response).

I tried the model on colab and while it was slow, it wasn’t as slow as you are reporting. If you’re getting the same error across 3 machines on fresh environments, there must be a problem in the script that is going unnoticed. Again, I’d check that all the shapes are as you expect, especially the F object.

Thank you! Colab is slightly faster but still too slow to be useful (pm.sample(100) takes ~10 minutes). I will try downgrading to pymc3 for the near term.

I sampled 1 chain of 2000 draws in 10 minutes on colab, so again, I think something is wrong with your code.

Almost certainly! For future reference, here is the code that runs extremely slowly:

import numpy as np
import matplotlib.pyplot as plt
import arviz as az
import pytensor.tensor as pt
nt = 100
n = 30

x = np.arange(nt)
#arbitrary good-looking function
true_mu=(50*x**2-.5*x**3-100*x)/1.e4
innovations = np.random.normal(scale=10, size=(n, nt))
true_rho = np.random.uniform()
ensemble = np.zeros((n, nt))
for t in range(1, nt):
    ensemble[:, t] = true_rho * ensemble[:, t-1] + innovations[:, t]+true_mu[t]
with pm.Model() as from_pymc3:

    σ=pm.HalfNormal("σ",10)
    ϕ=pm.Uniform("ϕ",-1,1)
   
    
    #AR(1) covariance matrix
    δ=np.arange(nt)
    imj=np.abs(np.subtract.outer(δ,δ))
    cov=pt.pow(σ,2)/(1-pt.pow(ϕ,2))*pt.pow(ϕ,imj)

    chol=pt.slinalg.cholesky(cov)
    
  
    #Uncertain mean value across all ensemble members
    μ=pm.Normal("μ",0,10,shape=nt)
    
    F=pm.MvNormal("likelihood_hist",mu=μ,chol=chol,\
                             observed= ensemble)
    pymc3_trace=pm.sample()

What are the sampling speeds if you set true_mu = 0 and F=pm.MvNormal("likelihood_hist",mu=0, chol=chol, observed= ensemble) ?

Still slow as hell on my laptop in the clean environment. I believe you that it sampled fine for you, but I am totally baffled- must be something about my environment?

Here’s a notebook showing that it samples very quickly without the mean. If it’s still slow on 3 different machines, each with a fresh install into a fresh env, I have no idea.

Read more here: Source link