How do I make a jit compiled backward differenced operator?

[HannoSpreeuw]

I would like to make jit compiled backward and forward gradients, using e.g. make_operator as in Numba-accelerated PDEs.

I have read the documentation about make_operator and register_operator, so I may have understood these routines, but the way I currently implement the backward differentiation cannot be jit compiled because it involves ScalarField. Not sure how to work around this.

This is what I got so far:

from pde import CartesianGrid, ScalarField
from pde.grids.operators.cartesian import _make_derivative
import numpy as np

depths = CartesianGrid([[0, 502/1319]], [400], periodic=False)
CA = ScalarField(depths, np.random.random(depths.shape))
CA_grad_back = CA.copy()
CA_grad_back.data[:] = 0
CA0 = 0.6
bc_CA = [{"value": CA0}, {"curvature" : 0}]
CA.set_ghost_cells(bc_CA)
backward_diff = _make_derivative(CA.grid, method="backward")
backward_diff(CA._data_full, out = CA_grad_back.data)

This works, but how to get this jit compiled?
I thought of something like this:

def make_backward_gradient(bcs):
     def backward_grad(arr, out):
           field = ScalarField(CA.grid, arr).copy()
           backward_differenced = np.empty(arr.shape)
           field.set_ghost_cells(bcs)
           backward_diff = _make_derivative(state.grid, method="backward")
           backward_diff(field._data_full, out = backward_differenced)
    return backward_grad

depths.register_operator("back", make_backward_gradient)
gradient_b = CA.grid.make_operator("back", bc=bc_CA)
CA_grad_back = gradient_b(CA.data)

but that gives a Numba problem:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function GridBase.make_operator.<locals>.apply_op at 0x7f6e378070a0>) found for signature:
 
 >>> apply_op(array(float64, 1d, C), out=array(float64, 1d, C), args=none)
 
There are 2 candidate implementations:
  - Of which 1 did not match due to:
  Overload in function 'register_jitable.<locals>.wrap.<locals>.ov_wrap': File: numba/core/extending.py: Line 159.
    With argument(s): '(array(float64, 1d, C), out=array(float64, 1d, C), args=none)':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   Failed in nopython mode pipeline (step: nopython frontend)
   Untyped global name 'ScalarField': Cannot determine Numba type of <class 'abc.ABCMeta'>
   
   File "<ipython-input-27-84ce6a4c9c03>", line 3:
       def backward_grad(arr, out):
           field = ScalarField(CA.grid, arr).copy()

which I understand, I cannot use ScalarField, of course, but how should I then enforce the boundary condition?

Is the way I use make_operator and register_operator okay?
Where would you do the registration, in _make_pde_rhs_numba?

[david-zwicker]

The operator returned by _make_derivative is already numba-compiled, so this should not be an issue.
If you want to register a new operator and later use it with a grid class, you can typically do that during the initialization of your module or script since it only needs to be done once.

Here is an example that registers an operator and then uses it to apply it to a field (after setting BCs) and also to create a numba-compiled function that does the same and can be used in other numba-compiled code:

import pde
from pde.grids.operators.cartesian import _make_derivative

# register the backward gradient operator with the Cartesian grid
pde.CartesianGrid.register_operator("grad_backward", lambda grid: _make_derivative(grid, method="backward"))

# create a test grid
grid = pde.UnitGrid([8])

# apply custom operator to a field
field = pde.ScalarField.from_expression(grid, "x**2")
field_grad1 = field._apply_operator("grad_backward", {"value": 0})
field_grad1.plot()

# make a numba compiled operator that sets BCs and applys operator
op = grid.make_operator("grad_backward", {"value": 0})
plt.plot(op(field.data))

[HannoSpreeuw]

That is a very complete explanation, much appreciated. apply_operator allows me to apply backward differentiation directly on ScalarFields.

[HannoSpreeuw]

Does it make a difference if one registers the operator with an instantiation of CartesianGrid, as in

grid = pde.CartesianGrid([[0, 10]], [100])
grid.register_operator("grad_backward", lambda grid: _make_derivative(grid, method="backward"))

?

[david-zwicker]

No, since register_operator is a classmethod, you can either call it on the class or the instance.