Implement mode “wrap” for numba.stencil

[Socob]

Feature request

Currently, even though numba.stencil is designed in a way that allows different “modes”, it only supports one:

1.11.4.2. mode

The optional mode parameter controls how the border of the output array is handled. Currently, there is only one supported value, "constant".

— https://numba.pydata.org/numba-doc/latest/user/stencil.html#mode

However, other options would certainly be useful. In my case, when dealing with differential equations with periodic boundary conditions, numba.stencil in its current form unfortunately can’t be used. If a mode 'wrap' were implemented, where array accesses outside the valid range would be taken modulo the length of the relevant array dimension, numba.stencil could also handle periodic boundary conditions.

For completeness, perhaps it is worth to draw inspiration from scipy.ndimage (e. g. scipy.ndimage.laplace, scipy.ndimage.convolve) for consistency (considering the identical naming and functionality of cval, perhaps this is already on the radar?):

'reflect' (d c b a | a b c d | d c b a)
The input is extended by reflecting about the edge of the last pixel.

'constant' (k k k k | a b c d | k k k k)
The input is extended by filling all values beyond the edge with the same constant value, defined by the cval parameter.

'nearest' (a a a a | a b c d | d d d d)
The input is extended by replicating the last pixel.

'mirror' (d c b | a b c d | c b a)
The input is extended by reflecting about the center of the last pixel.

'wrap' (a b c d | a b c d | a b c d)
The input is extended by wrapping around to the opposite edge.

[darothen]

Is this currently on the numba roadmap? If not, would it be possible for a high-level breakdown of what this might entail to see if a contributor may be able help out?

[kleinsimon]

I've created the following decorator, which adds the border using numpy pad and removes it afterwards, effectivley allowing the stencil decorator to use all border options of numpy.pad:

from numba import stencil
import numpy as np

def applyBorder(borderWidth, mode="reflect", **kwargs):
    def advStencil_decorator(func):
        def func_wrapper(a, *parms):
            a=np.pad(a, borderWidth, mode=mode, **kwargs)
            res  = func(a, *parms)
            return res[borderWidth:-borderWidth,borderWidth:-borderWidth]
        return func_wrapper
    return advStencil_decorator


@applyBorder(2, "reflect")
@stencil
def test(a):
    return a[-1,0] + a[1,0]


bla = np.arange(20).reshape((4,5))
print(test(bla))

[DrTodd13]

@kleinsimon Please see the original post. The design was to integrate different border handling into the stencil decorator itself over time. A benefit then is that you don't have to specify the borderWidth but it can be computed automatically (or through the neighborhood option.) Also, have you tested this by calling test from within a parallel region? I don't think that is going to work. Can you adapt your approach to work within the stencil decorator and within a parallel context?