How to troubleshoot sampling error with custom, not differentiable likelihood?

Questions v3
12

I don’t know of any packages unfortunately. I’m also not really sure how a package like this would look like, there are a lot of choices you can make if you want to include a pde solution into a model.

If I had a problem involving a pde right now, I’d probably start thinking about it along those lines:

  • Does the PDE have a time component? If so, maybe method of lines and then using an ode solver might be an option. sunode could do this, but unfortunately I still didn’t implement support for aesara right-hand-side functions there, so you either have to manually write the right hand side in numba or you have to work with sympy. The sympy right hand side might lead to compile time problems if the dimensionality gets too large though, which might easily happen with method of lines. (or you implement an aesara backend for the right hand side function in sunode and make me very happy with a PR :slight_smile: )

  • How do I want to solve the pde in the first place: Finite elements/finite diff/finite volume? What solver? If the problem is small enough, and there is no time component, maybe the PDE could be discretized and solved directly in aesara? If not, fenics is probably a good choice.

  • How large is the parameter space? If it is small (say <20), maybe an smc sampler could be a good choice? That way you don’t need gradients.

  • If you want to sample with nuts, first make a rough estimate to see if that is even realistic. If you are very lucky, you’ll need 10_000 gradient evaluations, more if you have bad geometry. Is that realistic? If not, you could start to look into surrogate models? Basically, you train some machine learning model to predict the logp and gradient, and use that during sampling. You might have to retrain it for the actual posterior though, so this would be something of a research project. :slight_smile:
    The tuning improvements in nutpie might also reduce the number of gradient evals, but not by orders of magnitude.

  • How do you get at the gradient? There used to be a package (fenics-adjoint I think) that automatically computes derivatives of pde solutions using fenics, I don’t know if that is still well supported though (fenics recently did an API breaking release). Is that fast enough, or is it worth the effort to write and solve the adjoint PDE manually? Then you need to wrap that in aesara.

It seems however that you already have a fenics implementation, including the gradient? Whether that is wrapped using aesara (pymc) or the old theano (pymc3) probably doesn’t make much of a difference, that seems to me to be the easiest part of the whole exercise. To switch to pymc version 4 you should mainly have to change the imports in the implementation of thenao_fem_solver from theano to aesara. :slight_smile:

If you have problems with initial points as you seem to have, it might help to start with really tight priors, reparameterized so that they are centered at zero with std 1, so something like this:

    k_raw = pm.Normal("k_raw", mu=0, sigma=1)
    k_raw = pm.Deterministic("k", tt.exp(k_raw * 4 + np.log(1)))

    A_raw = pm.Normal("A_raw", mu=0, sigma=1)
    A = pm.Deterministic("A", tt.exp(A_raw) * 100 + np.log(9000))
    ...

but probably tighten the priors if that makes sense (I don’t know what the parameters are).

I also removed the shape=(1,), because that looks strange to me. You are saying that those are arrays of length 1, when they really look like they are just scalars (ie shape=()).

3 Likes