PyMC3 Sampling with a model function that does not take array input natively

Questions v3
1

I’m curious if any PyMC3 experts can point me the right direction concerning very complicated model functions. I’ve found the PyMC3 and Theano documentation overwhelming and have not yet found either documentation or an example that is close to my situation. Apologies for the length of this question, I’m not sure exactly how to frame the question so I’ll just walk through the situation.

Essentially, I am trying to do Bayesian Logistic Regression on measured data, where the regression model is a very complicated function that cannot be written in a closed form. (Without being too detailed, the model function is actually the result of propagating a set of differential equations, subject to a very nasty nonlinear constraint.)

My general approach to logistic regression has been similar to:

with pm.Model() as model:
    # These are the priors for parameters in the model
    K_prior = pm.Normal('K', Kmid, Kstd_guess) # prior for K parameter
    unc_prior = pm.Uniform('error', 0, 2.0)  # prior for standard deviation of observations

    y_pred = pm.Normal('y_pred', mu = model_function(x_obs, a_param, b_param, K_prior),
                       sd = unc_prior, observed=y_obs)
    
    trace = pm.sample(2000, tune=1000, cores=2, chains=2, return_inferencedata=True)
az.summary(trace)

where model function is the function with parameters (A,B,K), where I am attempting to fit the parameters subject to observations (x_obs, y_obs). In examples that I have followed, this has been rather trivial to implement when the model function was straightforward. But it is failing for my very complicated model.

One problem is that my actual model function includes some Boolean operations, which do not work with Theano. Let’s ignore this for now, because I can build an example without Boolean operations.

The other problem, which I think is more severe, is that my model function does not behave like simpler examples or tutorials when given different types of input parameters. As in, suppose that a function is written as

def example_function(A,B,K):
    <various operations>
    return scalar_result

If example_function is called where all (A,B,K) are scalars, then it returns a scalar. If it is called with, say, K as a numpy array, then Numpy overloads the function and returns an ndarray. This appears to be a needed feature of model_function in the regression, where the sampler evaluates the function using samples drawn from the priors, but where the priors are passed to the function as iterable objects in the line with y_pred. I’ve worked through a handful of examples where this worked just fine; the key characteristic, though, was that model_function could handle scalar or array input natively, without any issues.

Now to my particular situation: my actual model_function is not well-suited to taking numpy arrays as inputs, unless I write a wrapper that catches array input, vectorizes the calculation across the individual terms in the array, then reassembles the results as an ndarray to return. I haven’t gotten this to work (yet) due to the Boolean issue I mentioned earlier. But it seems like a really terrible way to make the code run successfully.

So, without knowing exactly how to frame the question: when I have a model_function that is not suited to taking array input natively and return array output, what is the “right” approach to writing the model function? I’m happy to explain the problem more, even though I may not fully understand the issue.

Help?

2

There’s some overview here that you might find useful: PyMC3 and Theano — PyMC3 3.11.5 documentation

In general, if you can reexpress your model in an “array-like” way that’s the easiest because Theano tried to follow Numpy API where possible. There are some cases like branching and looping that require special Theano constructors like Scan and IfElse.

If you really can’t express your function with Theano primitives you always have the option of wrapping arbitrary (deterministic) code in a Theano Op. This will become a black-box which Theano can’t introspect. The most direct limitations is that it becomes your responsibility to provide optimized code and gradient expressions (if you want to use something that needs it like NUTS)

There is a v3 version of this guide, but I didn’t bother to get its link: Using a “black box” likelihood function (numpy) — PyMC example gallery

3

@ricardoV94 Thank you for the quick reply. I had found some (unclear) discussion that pointed me toward a custom Theano operation, so it is good to have that confirmed. And the examples you link are much better than others that I found. Again, thank you, I’ll see if this works for my model.

1 Like