How to Properly Track Gradients with MyGrad When Using Scipy's RectBivariateSpline for Interpolation?

I'm working on a project where I need to interpolate enthalpy values using scipy.interpolate.RectBivariateSpline and then perform automatic differentiation using mygrad. However, I'm encountering an issue where the gradient is not tracked at all across the interpolation. Here is a simplified version of my code:

import numpy as np
from scipy.interpolate import RectBivariateSpline
import CoolProp.CoolProp as CP
import mygrad as mg
from mygrad import tensor

# Define the refrigerant
refrigerant = 'R134a'

# Constant temperature (e.g., 20°C)
T = 20 + 273.15  # Convert to Kelvin

# Get saturation pressures
P_sat = CP.PropsSI('P', 'T', T, 'Q', 0, refrigerant)

# Define a pressure range around the saturation pressure
P_min = P_sat * 0.5
P_max = P_sat * 1.5
P_values = np.linspace(P_min, P_max, 100)

# Define a temperature range around the constant temperature
T_min = T - 10
T_max = T + 10
T_values = np.linspace(T_min, T_max, 100)

# Generate enthalpy data
h_values = []

for P in P_values:
    h_row = []
    for T in T_values:
        try:
            h = CP.PropsSI('H', 'P', P, 'T', T, refrigerant)
            h_row.append(h)
        except:
            h_row.append(np.nan)
    h_values.append(h_row)

# Convert lists to arrays
h_values = np.array(h_values)

# Fit spline for enthalpy
h_spline = RectBivariateSpline(P_values, T_values, h_values)

# Function to interpolate enthalpy
def h_interp(P, T):
    return tensor(h_spline.ev(P, T))

# Example function using the interpolated enthalpy with AD
def example_function(P):
    h = h_interp(P, T)
    result = h**2  # Example calculation
    return result

# Define a pressure value for testing
P_test = tensor(P_sat, )

# Compute the example function and its gradient
result = example_function(P_test)
result.backward()

# Print the result and the gradient
print(f"Result: {result.item()}")
print(f"Gradient: {P_test.grad}")

Are these just issues of RectBivariateSpline or mygrad? Would this work with other algebraic differentiation libs? Should I use something else besides splines?

The problem here is that MyGrad doesn't know how to differentiate this operation. You can get around this by defining a custom operation with a backwards pass. The MyGrad docs explain this here.

In order to implement the backward pass, you need to be able to evaluate a partial derivative of the spline. The SciPy docs explain this here. (See the dx and dy arguments.)

Combining the two, you get this:

import numpy as np

import mygrad as mg
from mygrad import execute_op
from mygrad.operation_base import Operation
from mygrad.typing import ArrayLike

# All operations should inherit from Operation, or one of its subclasses
class CustomInterpolate(Operation):
    """ Performs f(x, y) = RectBivariateSpline.ev(x, y) """

    def __call__(self, x: mg.Tensor, y: mg.Tensor, spline) -> np.ndarray:
        # This method defines the "forward pass" of the operation.
        # It must bind the variable tensors to the op and compute
        # the output of the operation as a numpy array

        # All tensors must be bound as a tuple to the `variables`
        # instance variable.
        self.variables = (x, y)
        self.spline = spline

        # The forward pass should be performed using numpy arrays,
        # not the tensors themselves.
        x_arr = x.data
        y_arr = y.data
        return self.spline.ev(x_arr, y_arr)

    def backward_var(self, grad, index, **kwargs):
        """Given ``grad = dℒ/df``, computes ``∂ℒ/∂x`` and ``∂ℒ/∂y``

        ``ℒ`` is assumed to be the terminal node from which ``ℒ.backward()`` was
        called.

        Parameters
        ----------
        grad : numpy.ndarray
            The back-propagated total derivative with respect to the present
            operation: dℒ/df. This will have the same shape as f, the result
            of the forward pass.

        index : Literal[0, 1]
            The index-location of ``var`` in ``self.variables``

        Returns
        -------
        numpy.ndarray
            ∂ℒ/∂x_{i}

        Raises
        ------
        SkipGradient"""
        x, y = self.variables
        x_arr = x.data
        y_arr = y.data

        # The operation need not incorporate specialized logic for
        # broadcasting. The appropriate sum-reductions will be performed
        # by MyGrad's autodiff system.
        if index == 0:  # backprop through a
            return self.spline.ev(x.data, y.data, dx=1)
        elif index == 1:  # backprop through b
            return self.spline.ev(x.data, y.data, dy=1)


# Our function stitches together our operation class with the
# operation arguments via `mygrad.prepare_op`
def custom_interpolate(x: ArrayLike, y: ArrayLike, spline, constant=None) -> mg.Tensor:
    # `execute_op` will take care of:
    #  - casting `x` and `y` to tensors if they are instead array-likes
    #  - propagating 'constant' status to the resulting output based on the inputs
    #  - handling in-place operations (specified via the `out` parameter)
    return execute_op(CustomInterpolate, x, y, op_args=(spline,), constant=constant)

You can use this operation like so:

def h_interp(P, T):
    return custom_interpolate(P, T, h_spline)

And then you can differentiate across this interpolation operation.

Output:

Result: 176061599645.87317
Gradient: -0.02227965104792612