使用 Scipy 的 RectBivariateSpline 进行插值时，如何使用 MyGrad 正确跟踪梯度？

Question

我正在开发一个项目，需要使用

scipy.interpolate.RectBivariateSpline

插值焓值，然后使用

mygrad

执行自动微分。但是，我遇到了一个问题，即在插值过程中根本没有跟踪渐变。这是我的代码的简化版本：

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}")

这些只是

RectBivariateSpline

还是

mygrad

的问题？这可以与其他代数微分库一起使用吗？除了样条线之外我还应该使用其他东西吗？

Answer 1

这里的问题是MyGrad不知道如何区分这个操作。您可以通过定义具有向后传递的自定义操作来解决此问题。 MyGrad 文档解释了这一点here。

为了实现向后传递，您需要能够计算样条曲线的偏导数。 SciPy 文档在here对此进行了解释。 （请参阅 dx 和 dy 参数。）

将两者结合起来，你会得到这个：

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)

您可以像这样使用此操作：

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

然后你可以区分这个插值操作。

输出：

Result: 176061599645.87317
Gradient: -0.02227965104792612