在具有自定义数据类型的 python 包装 C 函数上使用 autograd

Question

Can autograd in principle work on python wrapped C functions?

The C function I'd like to differentiate expects arguments with a REAL8 data type, and I can successfully call it in python by giving it either float or np.float64 arguments.

Upon closer inspection with my debugger, I found that I get a TypeError when autograd changes the float or np.float64 to an ArrayBox object to when it attempts to evaluate the gradient.

Is there a way to allow autograd to proceed in this scenario?

Is there another strategy that might be better for my situation?

Are there any primitives I can code up to fix this?

Background:

Bilby is a newer python library that wraps an older code (LALSIMULATION) written in C. It provides extensive precoded gravitational wave models, and I would like to use these pre-coded models in my research. My first task is to figure out how to calculate accurate Hessian's and gradients of these models. I'd like to use automatic differentiation to solve this problem due to its infamous numerical accuracy, but I'm stuck.

import autograd.numpy as np
from autograd import grad
import BILBY_TAYLORF2 # My own class that wraps Bilby

model = BILBY_TAYLORF2()

gradient_likelihood = grad(model.logLikelihood)

gradient_likelihood(np.array([10., 10.]))

TypeError: in method 'SimInspiralChooseFDWaveform', argument 3 of type 'REAL8'

SimInspiralChooseFDWaveform is the first call to C code for reference.

Answer 1

I'm afraid this answer is probably far too late for the original poster, but in case others come across it here's some information.

Autograd relies on functions being constructed of other functions for which analytical derivatives have been defined. In autograd's case it knows the derivatives for standard operators and it also redefines many numpy functions to include their analytical derivatives (or vector-Jacobian products ). It keeps track of all the functions/operations that have been put together and performs automatic differentiation by repeated use of the chain rule.

The swig wrapped lalsimulation functions just return a waveform. They do not have a gradient function defined (or a gradient for the dot product of the waveform with the data and itself, as would be required for the likelihood calculation in bilby), so autograd doesn't know what to do with them. The autograd tutorial shows's how you could define a primative function, define the vector-Jacobian product and tell autograd to link the two with the defvjp function. So, you would have to create your own functions that wrap the likelihood function and provide the gradient function. As you are treating the lalsimuation/likelihood functions as black boxes you don't know an analytical form of the gradient function, so it would just have to use a simple method like finite difference to get the gradient (maybe with some iteration). An example of this sort of thing, ie, defining a gradient for a black box function, is given in this example for PyMC3, which could probably be adapted for autograd.