Multivariate Bernoulli Distribution: Bayesian Analysis with PyMC3

import matplotlib.pyplot as plt
import pymc3 as pm
import aesara.tensor as at
import arviz as az
import scipy.stats as st
import numpy as np

az.style.use('arviz-whitegrid')

Multivariate bernoulli distribution

Sources

1. https://doi.org/10.1016/0047-259X(90)90084-U

2. https://arxiv.org/pdf/1206.1874

Motivation

I have recently been introduced to Bayesian framework for statistical analysis with PyMC3 package for Probabilistic Programming in Python 3.

I came across to a data given to me by my friend, which had multiple columns of binary data (responses to various questions with answers in binary format, yes or no). Each column could be represented by bernoulli process with an unknown probability of getting answer “yes”, which is what I wished to estimate. In addition to the above, the data also contains possible association between the columns, which I so dearly wanted to estimate. After a lot of literature search I came across the work done by B Dai in 2012 - 2013, which have been cited above in Sources.

The following post aims to translate the log likelihood function, joint probability distribution, marginal probability distribution and conditions for independence described above into PyMC3 code, which can be used by others to solve related problems.

Formulation

Let $Y=(Y_1,Y_2,\dots ,Y_K)$ be a $K$ dimensional random vector of possibly correlated Bernoulli random variables and let $y=(y_1,y_2,\dots ,y_K)$ be a realisation of $Y$. The more general form of $p(y_1,y_2,\dots ,y_K)$ of joint probability density is

$$P(Y_1=y_1,\dots ,Y_K=y_K)=p_{0,0,..,0}^{{\prod }_{j=1}^K(1-y_j)}{p}_{1,0,..,0}^{y_1{\prod }_{j=2}^K(1-y_j)}{p}_{0,1,..,0}^{y_1{y}_2{\prod }_{j=3}^K(1-y_j)}{p}_{1,1,..,1}^{{\prod }_{j=}^K{y}_j}$$

To simplify the notation, denote quantity $S$ to be

$$S^{j_1{j}_2\dots j_r}=\sum _{1\le s\le r}{f}^{j_s}+\sum _{1\le s<t\le r}{f}^{j_s{j}_t}+\dots +f^{j_1{j}_2\dots j_r}$$

Also, we will define the interaction function $B$ as

$$B^{j_1{j}_2\dots j_r}\left(y\right)=y_{j_1}{y}_{j_2}\dots y_{j_r}$$

The log-linear formulation of multivariate Bernoulli distribution will be

$$l(y,f)=-log\left[p\right(y\left)\right]=-\left[\sum _{r=1}^K\left(\sum _{1\le j_1<j_2<\dots <j_r\le K}{f}^{j_1{j}_2\dots j_r}{B}^{j_1{j}_2\dots j_r}\left(y\right)\right)-b\left(f\right)\right]$$

where, $f=(f^1,f^2,\dots ,f^{12\dots K}{)}^T$, is the vector of natural parameters for the multivariate Bernoulli distribution.

The normalising factor $b\left(f\right)=p_{0,0,\dots ,0}$ is defined as

$$b\left(f\right)=log\sum _{r=1}^K\left[1+\left(\sum _{1\le j_1<j_2<\dots <j_r\le K}{e}^{S^{j_1{j}_2\dots j_r}}\right)\right]$$

PyMC3 implementation

from itertools import chain, combinations

def powerset(iterable):
    '''
    [0,1,2] -> [(0,), (1,), (2,), (0, 1), (0, 2), (1, 2), (0, 1, 2)]
    '''
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(1, len(s) + 1))

def bernoulli_mv(f, b_f, subsets):
    '''
    Multivariate Bernoulli Distribution
    ===================================
    f: vector of f
    subsets: list of tuples of indices like [0,1,2] -> [(0,), (1,), (2,), (0, 1), (0, 2), (1, 2), (0, 1, 2)]
    '''
    
    def logp_(y):
        res = 0
        for i, sub in enumerate(subsets):
            res += pm.math.sum(f[i] * pm.math.prod(y[:, sub], axis = 1))
        
        res = res - (b_f * y.shape[0])
        
        return res
    
    return logp_

y = np.random.choice([0, 1], (50,5))

subsets = list((powerset(range(y.shape[1]))))

subsets

[(0,),
 (1,),
 (2,),
 (3,),
 (4,),
 (0, 1),
 (0, 2),
 (0, 3),
 (0, 4),
 (1, 2),
 (1, 3),
 (1, 4),
 (2, 3),
 (2, 4),
 (3, 4),
 (0, 1, 2),
 (0, 1, 3),
 (0, 1, 4),
 (0, 2, 3),
 (0, 2, 4),
 (0, 3, 4),
 (1, 2, 3),
 (1, 2, 4),
 (1, 3, 4),
 (2, 3, 4),
 (0, 1, 2, 3),
 (0, 1, 2, 4),
 (0, 1, 3, 4),
 (0, 2, 3, 4),
 (1, 2, 3, 4),
 (0, 1, 2, 3, 4)]

with pm.Model() as mod:
    fs = pm.Normal('f', 0, 2, shape = len(subsets))

    b_f = 1
    for s in subsets:
        ss = list(powerset(s))
        b_f += pm.math.exp(pm.math.sum(fs[np.array([s in ss for s in subsets])]))
    b_f = pm.Deterministic('b_f', pm.math.log(b_f))
    
    ys = pm.DensityDist('ys', bernoulli_mv(fs, b_f, subsets), observed = {'y': y})

with mod:
    trace = pm.sample(return_inferencedata=True)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
/home/suman/miniconda3/lib/python3.8/site-packages/aesara/graph/fg.py:525: UserWarning: Variable Elemwise{mul,no_inplace}.0 cannot be replaced; it isn't in the FunctionGraph
  warnings.warn(
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [f]

Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 166 seconds.

az.plot_trace(trace)

sumry = az.summary(trace)
sumry.index = subsets + ['b_f']
sumry

f $\to$ Joint Probabilities

$$p(j_1,j_2,\dots ,j_{r}positionsare1,othersare0)=\frac{e^{S^{j_1{j}_2\dots j_r}}}{e^{b\left(f\right)}}$$

We can get all the joint probabilities as follows

p_names = []
for sub in subsets:
    p_name = ['0'] * y.shape[1]
    for i in sub:
        p_name[i] = '1'
    p_names.append(p_name)

p_names = ['p_' + ''.join(name) for name in p_names]

for p_name, s in zip(p_names, subsets):
    ss = list(powerset(s))
    idx_S = np.array([s in ss for s in subsets])
    trace.posterior[p_name] = np.exp(np.sum(trace.posterior['f'][...,idx_S], axis = 2)) / \
                              np.exp(trace.posterior['b_f'])

trace.posterior['p_' + '0' * y.shape[1]] = np.exp(-trace.posterior['b_f'])

az.summary(trace, var_names = ['p_'], filter_vars='like')

Marginal Probabilities

We represent marginal probability as follows

$p_i$ is probability of $i^{th}$ column being 1, where $i\in \{1,2,\dots ,K\}$

Similarly, $p_{ij}$ is the probability of $i^{th}$ and $j^{th}$ columns being 1, where $i,j\in \{1,2,\dots ,K\}$.

$$p_i=\sum _{forall{j}_i=1}{p}_{j_1{j}_2\dots j_k}$$

def marg_prob(idxs, trace, num_columns):
    '''
    idxs: list of indexes which are 1. eg [0,1] means first and second columns are 1
    trace: trace of pymc3 model output
    '''
    assert max(idxs) < num_columns
    keys = [i for i in trace.posterior.keys() if i.startswith('p_')]
    res = []
    for k in keys:
        k_num = k[2:]
        for c, i in enumerate(idxs):
            if k_num[i] != '1':
                break
            if c + 1 == len(idxs):
                res.append('p_' + k_num)
    return np.sum(trace.posterior[res].to_array(), axis = 0)

for i in range(y.shape[1]):
    print(f'Marginal Probability {i + 1} ---')
    print(az.summary(marg_prob([i], trace, y.shape[1]), kind = 'stats'))

Marginal Probability 1 ---
    mean     sd  hdi_3%  hdi_97%
x  0.573  0.067   0.442     0.69
Marginal Probability 2 ---
   mean     sd  hdi_3%  hdi_97%
x   0.4  0.065    0.27    0.511
Marginal Probability 3 ---
   mean     sd  hdi_3%  hdi_97%
x  0.58  0.067    0.46    0.706
Marginal Probability 4 ---
    mean     sd  hdi_3%  hdi_97%
x  0.463  0.071   0.323    0.588
Marginal Probability 5 ---
    mean    sd  hdi_3%  hdi_97%
x  0.477  0.07   0.345    0.603

Independence of outcomes

The outcomes are independent, if

$$S^{j_1{j}_2\dots j_r}-\sum _{k=1}^r{f}^{j_k}=0,\forall r\ge 2$$

def trace_of_independence(idxs, trace, subsets, num_columns):
    '''
        idxs: list of indexes whose independence is to be checked, eg [0,1,2] means that 
        independence of 1st, 2nd and 3rd index needs to be checked
        trace: pymc3 trace, inferenceData
        subsets: list depicting powerset of the indexes
        num_columns: number of outcomes
    '''
    assert len(idxs) > 1
    
    assert max(idxs) < num_columns
    
    ss = [i for i in powerset(idxs) if len(i) > 1]
    idx_S = np.array([s in ss for s in subsets])
    
    return np.sum(trace.posterior['f'][...,idx_S], axis = 2)

Issues

The sampling takes prohibitively long time when the number of variables $K$ exceeds 7 - 8. I am in the process to tackle the issue.

Also, I have to make mechanism by which random samples can be generated.