Sampling many dimensions, with different PDF in each dimension

Questions version agnostic
1

Dear Community,

I’ve been trying, without much success, to create a model in which each data point corresponds to a 2-dimensional array, and each element comes from a different PDF.

For instance, in my mind I would imagine doing something like this:

basic_model = pm.Model()
with basic_model:

x0 = pm.Normal("x0", mu=0, sigma=10)
x1  = pm.HalfNormal("x1", sigma=2)

Y_obs = [ pm.Normal("Y0", mu=x0, sigma=x1, observed=[x[0] for x in data]) , pm.Exponential("Y1", lam=x1, observed= [x[1] for x in data) ]

That is, the sampling space has dimension=2, and the two dimensions have different PDFs, but connected through some variable (in this case x1),

The above is the most simplified example I could imagine of a broader problem that is what I’m trying to solve. To be more complete, I am trying to reproduce [2210.07358] Bayesian inference to study a signal with two or more decaying particles in a non-resonant background results using PYMC (in the paper is done through nested sampling). Hence, the full problem corresponds to a mixture model of 2 classes:

Class 0: Direct product of 2 Normal distributions
Class 1: Direct product of 2 Exponential distributions

Data corresponds to N data points, each consisting in a pair of numbers, [mbb, mAA)]. And one needs to infer all the parameters of the mixture model (2 parameters for each Normal, 1 for each Exponential, and the mixture weight, which yields 7 parameters).

I know how to do a mixture in PYMC, but I don’t know how to do a mixture in which each class is a direct product of two known distributions.

Thank you for illuminating me with your answers! Ezequeil.

2

You can have multiple likelihoods in a model and they can share parameters just fine

Just pass each vector as observed like you did without trying to put it in a list.

Pseudo-code:

obs1 = pm.Normal("obs1", ... x, observed=data[..., 0])
obs2 = pm.Exponential("obs2", ... x, observed=data[..., 1])
3

Hi Ricardo,

That’s great, I wasn’t aware of it, thank you! I’ve tried it and works great!

What I cannot realize how to do it is the following. I need to do a model that is a mixture of 2 classes, and in each class I have the situation that you have just solved. That is, each class has as a likelihood the product of 2 likelihoods (or a direct product of 2 random variables which have their likelihoods, as I prefer to understand it).

Would you have an idea of how to implement such a model? Because the pm.Mixture() receives the input ‘comp_dists’ which is a vector with 1 distribution in each element. I’d like to put 2 distributions in each element, or some work-around in the line of what you’ve smartly proposed!

Thank you very much! Ezequiel.

4

Hi Ezequiel,

I am not sure what you mean by product of variables.

You can use mixtures as the components of a mixture (mixture of mixtures) but I am not sure if that’s what you’re looking for.

Can you perhaps explain your model or data generating process in more detail?

5

Ola Ricardo!

Thank you, sorry for not being clear enough.

My data consists of N data points, where each data point consists of 2 numbers, say x0, x1. These two numbers at each data point could be sampled from either

class_0, in which case
x0 ~ Normal(mu00, sigma00)
x1 ~ Normal(mu01, sigma01)
(i.e. different normals)

or from class_1, in which case
x0 ~ Exponential(lambda10)
x1 ~ Exponential(lambda11)
(again, they are different exponentials)

I hope this is understandable. If not, I can copy paste the code to generate the data,

Thank you again, best, Ezequiel.

6

Sounds like x0~Mixture(..., comp_dists=[Normal(mu00, sigma00), Exponential(lambda10)])
x1~Mixture(..., comp_dists=[Normal(mu01, sigma01), Exponential(lambda11)]).

Which you can model as two separate Mixture likelihoods, each with one Normal and Exponential components with the appropriate parameters.

7

Hi Ricardo

Many thanks again!

I think that your proposal would not work because at each data point x1 and x2 could belong to different classes. In the probabilistic model, at each data point both x1 and x2 should belong to either class_1, or either class_2.

In my understanding of your response, this feature is not included. Or am I missing something?

By the way, this is a typical problem in physics, and I guess in other fields as well.

Thank you, best, Ezequiel.

9

I see, so each pair comes from either a diagonal MvNormal or a “diagonal” MvExponential. We don’t have a thing like MvExponential, otherwise your model would be

pm.Mixture("llike", w, comp_dists=[
  MvNormal(mu=[mu00, mu01], cov=np.eye(2)*np.array([sigma00, sigma01])**2), 
  MvExponential.dist(...),
  ], observed=data)

I thought about adding a helper that makes a “fake” multivariate distribution from vectors of independent univariate distributions, but didn’t actually do it.

Anyway, you can create a MvExponential with CustomDist: pymc.CustomDist — PyMC dev documentation

Something like this (untested code):

def diag_mv_exponential_dist(lambda1, lambda2, size):
  if size is None:
    size = ()
  return pm.Exponential.dist([lambda1, lambda2], shape=tuple(size) + (2,))

def diag_mv_exponential_logp(value, lambda1, lambda2):
    # Contract logp of the pairs
    return pm.logp(pm.Exponential.dist([lambda1, lambda2]), value).sum(-1)

with pm.Model() as m:
...
  diag_mv_exponential = pm.CustomDist.dist(
    lambda0, lambda1, 
    dist=diag_mv_exponenntial_dist,
    logp=diag_mv_exponential_logp, 
    ndim_supp=1,   
    class_name="diag_mv_exponential",
)

And use that in the Mixture components. You can also make a diag_mv_normal if you don’t want to use MvNormal

10

Hi Ricardo, FANTASTIC! This was what I was looking for. Thanks!

I’ll try to implement it and come back here to let it public for others.

Best, Ezequiel.

11

You’re welcome. One important thing in the CustomDist is to set ndim_supp=1. That’s how you tell PyMC that it’s a multivariate distribution.

I had forgotten about it, but added it to the untested code example above.

12

Hi Ricardo,

It took me some time to correctly implement pymc5 to make CustomDist work. However, once implemented, the script does not work. I don’t make it to understand the output message . Here the info. Thank you very much! (Maybe there is another option simpler than CustomDist?)

Best, Ezequiel.

def diag_mv_exponential_dist(lambda1, lambda2, size):
    if size is None:
        size = ()
    return pm.Exponential.dist([lambda1, lambda2], shape=tuple(size) + (2,))

def diag_mv_exponential_logp(value, lambda1, lambda2):
    # Contract logp of the pairs
    return pm.logp(pm.Exponential.dist([lambda1, lambda2]), value).sum(-1)

with pm.Model() as m:
    w = pm.Dirichlet('w', a=np.array([1, 1]))
    mu0 = pm.Normal("mu0", mu=120, sigma=15)
    mu1 = pm.Normal("mu1", mu=124, sigma=2)
    sigma0 = pm.HalfNormal("sigma0", sigma=15)
    sigma1 = pm.HalfNormal("sigma1", sigma=1)        
    lambda0 = pm.Normal("lambda0", mu=0.05, sigma=0.02)
    lambda1 = pm.Normal("lambda1", mu=0.5, sigma=0.2)
    
    diag_mv_exponential = pm.CustomDist.dist(lambda0, lambda1, dist=diag_mv_exponential_dist, logp=diag_mv_exponential_logp, ndim_supp=1, class_name="diag_mv_exponential")
    diag_mv_normal = pm.MvNormal.dist(mu=[mu0,mu1], cov=[[sigma0,0],[0,sigma1]])
    mymixture = [diag_mv_normal, diag_mv_exponential]

    Y_obs = pm.Mixture("Y_obs", w = w, comp_dists=mymixture, observed=data)

And sends me the following message, which starts with a " UserWarning: CustomDist with dist function is still experimental. Expect bugs!"

/home/sequi/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:623: UserWarning: CustomDist with dist function is still experimental. Expect bugs!
  warnings.warn(
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/pytensorf.py:436, in floatX(X)
    435 try:
--> 436     return X.astype(pytensor.config.floatX)
    437 except AttributeError:
    438     # Scalar passed

AttributeError: 'list' object has no attribute 'astype'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
TypeError: float() argument must be a string or a real number, not 'TensorVariable'

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[50], line 19
     16 lambda0 = pm.Normal("lambda0", mu=0.05, sigma=0.02)
     17 lambda1 = pm.Normal("lambda1", mu=0.5, sigma=0.2)
---> 19 diag_mv_exponential = pm.CustomDist.dist(lambda0, lambda1, dist=diag_mv_exponential_dist, logp=diag_mv_exponential_logp, ndim_supp=1, class_name="diag_mv_exponential")
     20 diag_mv_normal = pm.MvNormal.dist(mu=[mu0,mu1], cov=[[sigma0,0],[0,sigma1]])
     21 mymixture = [diag_mv_normal, diag_mv_exponential]

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:1023, in CustomDist.dist(cls, class_name, dist, random, logp, logcdf, moment, ndim_supp, ndims_params, dtype, *dist_params, **kwargs)
   1021 cls.check_valid_dist_random(dist, random, dist_params)
   1022 if dist is not None:
-> 1023     return _CustomSymbolicDist.dist(
   1024         *dist_params,
   1025         class_name=class_name,
   1026         dist=dist,
   1027         logp=logp,
   1028         logcdf=logcdf,
   1029         moment=moment,
   1030         ndim_supp=ndim_supp,
   1031         **kwargs,
   1032     )
   1033 else:
   1034     return _CustomDist.dist(
   1035         *dist_params,
   1036         class_name=class_name,
   (...)
   1044         **kwargs,
   1045     )

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:641, in _CustomSymbolicDist.dist(cls, class_name, dist, logp, logcdf, moment, ndim_supp, dtype, *dist_params, **kwargs)
    633 if moment is None:
    634     moment = functools.partial(
    635         default_moment,
    636         rv_name=class_name,
    637         has_fallback=True,
    638         ndim_supp=ndim_supp,
    639     )
--> 641 return super().dist(
    642     dist_params,
    643     class_name=class_name,
    644     logp=logp,
    645     logcdf=logcdf,
    646     dist=dist,
    647     moment=moment,
    648     ndim_supp=ndim_supp,
    649     **kwargs,
    650 )

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:387, in Distribution.dist(cls, dist_params, shape, **kwargs)
    385 ndim_supp = getattr(cls.rv_op, "ndim_supp", None)
    386 if ndim_supp is None:
--> 387     ndim_supp = cls.rv_op(*dist_params, **kwargs).owner.op.ndim_supp
    388 create_size = find_size(shape=shape, size=size, ndim_supp=ndim_supp)
    389 rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:670, in _CustomSymbolicDist.rv_op(cls, class_name, dist, logp, logcdf, moment, size, ndim_supp, *dist_params)
    666 dummy_dist_params = [dist_param.type() for dist_param in dist_params]
    667 with BlockModelAccess(
    668     error_msg_on_access="Model variables cannot be created in the dist function. Use the `.dist` API"
    669 ):
--> 670     dummy_rv = dist(*dummy_dist_params, dummy_size_param)
    671 dummy_params = [dummy_size_param] + dummy_dist_params
    672 dummy_updates_dict = collect_default_updates(dummy_params, (dummy_rv,))

Cell In[50], line 4, in diag_mv_exponential_dist(lambda1, lambda2, size)
      2 if size is None:
      3     size = ()
----> 4 return pm.Exponential.dist([lambda1, lambda2], shape=tuple(size) + (2,))

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/continuous.py:1314, in Exponential.dist(cls, lam, *args, **kwargs)
   1312 @classmethod
   1313 def dist(cls, lam, *args, **kwargs):
-> 1314     lam = at.as_tensor_variable(floatX(lam))
   1316     # PyTensor exponential op is parametrized in terms of mu (1/lam)
   1317     return super().dist([at.reciprocal(lam)], **kwargs)

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/pytensorf.py:439, in floatX(X)
    436     return X.astype(pytensor.config.floatX)
    437 except AttributeError:
    438     # Scalar passed
--> 439     return np.asarray(X, dtype=pytensor.config.floatX)

ValueError: setting an array element with a sequence.
13

Try to call pm.math.concatenate([lam1, lam2]) before passing it into the Exponential distribution.

It’s an issue with floatX being called inside Exponential before converting the input into a PyMC variable

14

Hi Ricardo!

I’ve solve it with your help, now it works! I leave the working code here for the next guy/girl! And a few comments at the end, also a question.

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import graphviz
import seaborn as sns

import pymc as pm
import pymc.sampling_jax
import arviz as az

import gc
import warnings
warnings.simplefilter(action="ignore", category=FutureWarning)

%matplotlib inline
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all" 
from IPython.display import Image
print(f"Running on PyMC v{pm.__version__}")

############

# Create synthetic dara
lam_b = 0.08
lam_A = 0.2
mu_b = 130
sigma_b = 5
mu_A = 125
sigma_A = 1.5

ndata = 5000
data = []

for i in range(ndata):
    if i%1000 == 0: print(i)
    if np.random.binomial(1,.7):
        data.append([100 + np.random.exponential(1/lam_b), 120 + np.random.exponential(1/lam_A)])
    else:
        data.append([np.random.normal(mu_b,sigma_b),np.random.normal(mu_A,sigma_A)])

############

# Define aux functions and the mixture model usgin CustomDist (needs PyMC 5!!)
def diag_mv_exponential_dist(lambda0, lambda1, size):    
    if size == None: 
        size = ()
    return pm.Exponential.dist(pm.math.stack([lambda0, lambda1]), shape=tuple(size) + (2,))

def diag_mv_exponential_logp(value, lambda0, lambda1):
    # Contract logp of the pairs
    return pm.logp(pm.Exponential.dist(pm.math.stack([lambda0, lambda1])), value).sum(-1)

with pm.Model() as mixture_of_direct_product:
    w = pm.Dirichlet('w', a=np.array([1, 1]))
    mu0 = pm.Normal("mu0", mu=120, sigma=15)
    mu1 = pm.Normal("mu1", mu=124, sigma=2)
    sigma0 = pm.HalfNormal("sigma0", sigma=15)
    sigma1 = pm.HalfNormal("sigma1", sigma=1)        
    lambda0 = pm.Normal("lambda0", mu=0.05, sigma=0.02)
    lambda1 = pm.Normal("lambda1", mu=0.5, sigma=0.2)
    
    diag_mv_exponential = pm.CustomDist.dist(lambda0, lambda1, dist=diag_mv_exponential_dist, logp=diag_mv_exponential_logp, ndim_supp=1, class_name="diag_mv_exponential")
    diag_mv_normal = pm.MvNormal.dist(mu=[mu0,mu1], cov=pm.math.stack(([[sigma0,0],[0,sigma1]])))
    mymixture = [diag_mv_normal, diag_mv_exponential]

    Y_obs = pm.Mixture("Y_obs", w = w, comp_dists=mymixture, observed=data)

############

# Sample from the posterior using the pm.sample method and plot results
data_sample = pm.sample(model = mixture_of_direct_product, chains=30)
pm.plot_trace(data_sample)

############

# Use another sampling method:
data_jax = pm.sampling_jax.sample_numpyro_nuts(model = mixture_of_direct_product)
az.plot_posterior(data_jax)

A few comments:

  • I had to use pm.math.stack instead of pm.math.concatenate, I think that this is because the lambda are scalars
  • I had to also modify the mvNormal covariance as a pm.math.stack with double parenthesis to make it work
  • Now works perfect and samples!

However…

Results from the posterior do not match the true values. Something is not working properly, I don’t know, may be it is more Art & Metiers . But at least it is in principle working.

(The problem proposed is not easy, I know that the posterior is sophisticated. (However, I’ve had a very good sampling using Dynesty, just that it took 4hs for 500 datapoints. Here it took 40m for 5k datapoints… but didn’t retrieve the true values. Maybe 30 chains are too few (?).)

Thank you for everything, all the best, Ezequiel.

15

@Ezequiel_Alvarez

I don’t see where in the PyMC model do you include this offset by [100, 120]?