Custom distribution using sympy

Is it possible to implement sympy using pymc3?
For a toy example, I’m trying to make an inference about the parameters of the discretized gamma distribution. But the following code doesn’t seem to work. Is there any example of Theano wrapper that would be appropriate for using sympy in pymc3?

import numpy as np
from sympy import Symbol
from theano.tensor.elemwise import Elemwise
from theano.scalar.basic_sympy import SymPyCCode
data=np.floor(np.random.gamma(3,scale=20,size=200))


def discrete_gamma( k_th,theta_th):
    k = Symbol('k')
    theta = Symbol('theta')
    X = Gamma("x", k, theta)
    y = Symbol('y')
    p=sympy.stats.P(X<y+1)-sympy.stats.P(X<y) # y is integer value
    p_op = Elemwise(SymPyCCode([y, k, theta], p))
    def _logp(y_th):
        p_th = theano.function([y_th, m_th, theta_th], p_op(y_th, m_th, theta_th))
        p=tt.log(p_th)
        return p
    return _logp
with pm.Model() as potential_model:
    q1=pm.Uniform("q1", lower=0,upper=2)
    mu1=pm.Normal('mu1', mu=10, sigma=4)
    mu_sq=tt.sqr(mu1)
    k2=mu_sq/q1
    a = pm.DensityDist('obs',discrete_gamma(k_th=k2,theta_th=mu1/q1),observed=data)
    start = pm.find_MAP()
    trace = pm.sample(start=start)
pm.traceplot(trace)
1 Like

The functions sympy.printing.theanocode.[theano_code, theano_function] are sufficient for converting straightforward SymPy graphs to Theano graphs.

The primary difficulty I’ve noticed involves the re-use of existing Theano variables (e.g. the ones produced by PyMC3).

Here’s an example:

import theano
import theano.tensor as tt

import sympy as sp
import sympy.stats as st

from theano.printing import debugprint as tt_dprint
from sympy.printing.theanocode import theano_code


# theano.config.cxx = ""
# theano.config.mode = "FAST_COMPILE"
# tt.config.compute_test_value = 'ignore'


k_sp, theta_sp, y_sp = sp.symbols(["k", "theta", "y"])

X_sp = st.Gamma("X", k_sp, theta_sp)

gamma_pdf_sp = st.density(X_sp, y_sp)


# Create a Theano graph for the SymPy PDF graph
gamma_pdf_tt = theano_code(gamma_pdf_sp(y_sp))

tt_dprint(gamma_pdf_tt)

The output is

Elemwise{mul,no_inplace} [id A] ''   
 |Elemwise{pow,no_inplace} [id B] ''   
 | |theta [id C]
 | |Elemwise{mul,no_inplace} [id D] ''   
 |   |TensorConstant{-1} [id E]
 |   |k [id F]
 |Elemwise{pow,no_inplace} [id G] ''   
 | |y [id H]
 | |Elemwise{add,no_inplace} [id I] ''   
 |   |TensorConstant{-1} [id E]
 |   |k [id F]
 |Elemwise{pow,no_inplace} [id J] ''   
 | |Elemwise{gamma,no_inplace} [id K] ''   
 | | |k [id F]
 | |TensorConstant{-1} [id E]
 |Elemwise{exp,no_inplace} [id L] ''   
   |Elemwise{mul,no_inplace} [id M] ''   
     |TensorConstant{-1} [id E]
     |y [id H]
     |Elemwise{pow,no_inplace} [id N] ''   
       |theta [id C]
       |TensorConstant{-1} [id E]

Now, let’s say we had our own Theano variable that we wanted to use in place of the SymPy variable k_sp (e.g. the kind produced by pm.Normal, etc.):

def create_cache_key(s, name=None, dtype=None, broadcastable=None):
    """Create a TheanoPrinter cache key for a SymPy variable."""
    if name is None:
        name = s.name
    if dtype is None:
        dtype = 'floatX'
    if broadcastable is None:
        broadcastable = ()
    return (name, type(s), s.args, dtype, broadcastable)



# An existing Theano variable that we want to use in the output Theano graph
k_tt = tt.scalar('k_ours')

gamma_pdf_tt = theano_code(gamma_pdf_sp(y_sp), cache={create_cache_key(k_sp): k_tt})

tt_dprint(gamma_pdf_tt)

Now, the output Theano graph contains our k_ours variable:

Elemwise{mul,no_inplace} [id A] ''   
 |Elemwise{pow,no_inplace} [id B] ''   
 | |theta [id C]
 | |Elemwise{mul,no_inplace} [id D] ''   
 |   |TensorConstant{-1} [id E]
 |   |k_ours [id F]
 |Elemwise{pow,no_inplace} [id G] ''   
 | |y [id H]
 | |Elemwise{add,no_inplace} [id I] ''   
 |   |TensorConstant{-1} [id E]
 |   |k_ours [id F]
 |Elemwise{pow,no_inplace} [id J] ''   
 | |Elemwise{gamma,no_inplace} [id K] ''   
 | | |k_ours [id F]
 | |TensorConstant{-1} [id E]
 |Elemwise{exp,no_inplace} [id L] ''   
   |Elemwise{mul,no_inplace} [id M] ''   
     |TensorConstant{-1} [id E]
     |y [id H]
     |Elemwise{pow,no_inplace} [id N] ''   
       |theta [id C]
       |TensorConstant{-1} [id E]

The same cache={...} approach should work for theano_function, as well.

1 Like