Unexpected use of slice when evaluating a custom stencil

[cdelavegamartin]

I'm not sure if this is a bug or something I'm not understanding correctly

I'm simply trying to implement a first order finite difference scheme (I'm building an upwind scheme). When given a slice to evaluate on, I can't get it to work as I would expect
Code:

from findiff.stencils import Stencil
import numpy as np

def calculate_derivative_first_order(level_set: np.ndarray, axis: int, spacings: tuple[float,float,float], mode: str = 'fwd') -> np.ndarray:
    """Compute forward or backward difference derivative along a given axis."""
    
    ndims = level_set.ndim
    offsets_0 = [0] * ndims
    offsets_1 = [0] * ndims
    if mode == 'fwd':
        offsets_1[axis] = 1
    elif mode == 'bwd':
        offsets_1[axis] = -1
    else:
        raise NotImplementedError(f"Mode {mode} not implemented.")
    
    offsets = [tuple(offsets_0), tuple(offsets_1)]
    par_der = [0] * ndims
    par_der[axis] = 1
    partials = {tuple(par_der): 1,}
    

    stencil = Stencil(
        offsets=offsets,
        partials=partials,
        spacings=spacings,
    )
    slices = [slice(None, level_set.shape[0],None), slice(None, level_set.shape[1],None), slice(None, level_set.shape[2],None)]
    if mode == 'fwd':
        slices[axis] = slice(0,-1)
    elif mode == 'bwd':
        slices[axis] = slice(1,None)
    print(f"slices: {slices}")
    print(f"level_set shape: {level_set.shape}")
    print(f"Applying stencil along axis {axis} with mode {mode}")
    print(f"Stencil offsets: {stencil.offsets}")
    print(f"Stencil partials: {stencil.partials}")
    print(f"Stencil spacings: {stencil.spacings}")
    print(f"Level set slice shape: {level_set[tuple(slices)].shape}")
    print("################# Applying stencil #################")
    derivative = stencil(level_set, on=tuple(slices))
    # derivative = stencil(level_set, at=(2,2,2))  # Example point for debugging
    return derivative

def test_calculate_derivative_first_order():
    """Test first order derivative calculation."""
    # Create a simple 3D grid
    grid_shape = (3,3,3)
    x = y = z = np.linspace(0, 1, 3)
    grid_vectors = (x, y, z)

    # Create a level set that varies linearly along x
    level_set = np.zeros(grid_shape)
    for i in range(3):
        level_set[i, :, :] = x[i]

    print("Level Set:\n", level_set)
    dx = x[1] - x[0]
    dy = y[1] - y[0]
    dz = z[1] - z[0]
    spacings = (dx, dy, dz)
    print("Spacings:", spacings)

    # Calculate forward derivative along x
    derivative_fwd = calculate_derivative_first_order(
        level_set, axis=0, spacings=spacings, mode='fwd'
    )
    
    print("Forward Derivative:\n", derivative_fwd)
    

    # Expected derivative is 1
    expected_derivative = 1.0

    # Check values (excluding boundaries)
    assert np.allclose(derivative_fwd[:-1, :, :], expected_derivative)

if __name__ == "__main__":
    test_calculate_derivative_first_order()

I have also modified Stencil._apply_on_multi_slice in findiff.stencils to print out extra information:

def _apply_on_multi_slice(self, f, on):
        result = np.zeros_like(f)
        base_mslice = [self._canonic_slice(sl, f.shape[axis]) for axis, sl in enumerate(on)]
        print("Applying stencil on multi-slice:")
        print("Base mslice:", base_mslice)
        for off, coeff in self.values.items():
            print("Offset:", off)
            print("Coefficient:", coeff)
            off_mslice = list(base_mslice)
            print("Initial off_mslice:", off_mslice)
            for axis, off_ in enumerate(off):
                print("  Axis:", axis, "Offset value:", off_)
                start = base_mslice[axis].start + off_
                stop = base_mslice[axis].stop + off_
                print("    Start:", start, "Stop:", stop)
                off_mslice[axis] = slice(start, stop)
                print("    Updated off_mslice:", off_mslice)
            print("Final off_mslice:", off_mslice)
            print("Adding to result at", tuple(base_mslice), "the value coeff * f at", tuple(off_mslice))
            print("Value to add: \n", coeff * f[tuple(off_mslice)])
            print("Value before addition: \n", result[tuple(base_mslice)])
            result[tuple(base_mslice)] += coeff * f[tuple(off_mslice)]
        return result

Output:

Level Set:
 [[[0.  0.  0. ]
  [0.  0.  0. ]
  [0.  0.  0. ]]

 [[0.5 0.5 0.5]
  [0.5 0.5 0.5]
  [0.5 0.5 0.5]]

 [[1.  1.  1. ]
  [1.  1.  1. ]
  [1.  1.  1. ]]]
Spacings: (np.float64(0.5), np.float64(0.5), np.float64(0.5))
slices: [slice(0, -1, None), slice(None, 3, None), slice(None, 3, None)]
level_set shape: (3, 3, 3)
Applying stencil along axis 0 with mode fwd
Stencil offsets: [(0, 0, 0), (1, 0, 0)]
Stencil partials: {(1, 0, 0): 1}
Stencil spacings: (np.float64(0.5), np.float64(0.5), np.float64(0.5))
Level set slice shape: (2, 3, 3)
################# Applying stencil #################
Applying stencil on multi-slice:
Base mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
Offset: (0, 0, 0)
Coefficient: -2.0
Initial off_mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
  Axis: 0 Offset value: 0
    Start: 0 Stop: 3
    Updated off_mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
  Axis: 1 Offset value: 0
    Start: 0 Stop: 3
    Updated off_mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
  Axis: 2 Offset value: 0
    Start: 0 Stop: 3
    Updated off_mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
Final off_mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
Adding to result at (slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)) the value coeff * f at (slice(0, 3, None), slice(0, 3, None), slice(0, 3, None))
Value to add: 
 [[[-0. -0. -0.]
  [-0. -0. -0.]
  [-0. -0. -0.]]

 [[-1. -1. -1.]
  [-1. -1. -1.]
  [-1. -1. -1.]]

 [[-2. -2. -2.]
  [-2. -2. -2.]
  [-2. -2. -2.]]]
Value before addition: 
 [[[0. 0. 0.]
  [0. 0. 0.]
  [0. 0. 0.]]

 [[0. 0. 0.]
  [0. 0. 0.]
  [0. 0. 0.]]

 [[0. 0. 0.]
  [0. 0. 0.]
  [0. 0. 0.]]]
Offset: (1, 0, 0)
Coefficient: 2.0
Initial off_mslice: [slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)]
  Axis: 0 Offset value: 1
    Start: 1 Stop: 4
    Updated off_mslice: [slice(1, 4, None), slice(0, 3, None), slice(0, 3, None)]
  Axis: 1 Offset value: 0
    Start: 0 Stop: 3
    Updated off_mslice: [slice(1, 4, None), slice(0, 3, None), slice(0, 3, None)]
  Axis: 2 Offset value: 0
    Start: 0 Stop: 3
    Updated off_mslice: [slice(1, 4, None), slice(0, 3, None), slice(0, 3, None)]
Final off_mslice: [slice(1, 4, None), slice(0, 3, None), slice(0, 3, None)]
Adding to result at (slice(0, 3, None), slice(0, 3, None), slice(0, 3, None)) the value coeff * f at (slice(1, 4, None), slice(0, 3, None), slice(0, 3, None))
Value to add: 
 [[[1. 1. 1.]
  [1. 1. 1.]
  [1. 1. 1.]]

 [[2. 2. 2.]
  [2. 2. 2.]
  [2. 2. 2.]]]
Value before addition: 
 [[[ 0.  0.  0.]
  [ 0.  0.  0.]
  [ 0.  0.  0.]]

 [[-1. -1. -1.]
  [-1. -1. -1.]
  [-1. -1. -1.]]

 [[-2. -2. -2.]
  [-2. -2. -2.]
  [-2. -2. -2.]]]
Traceback (most recent call last):
  File "/home/mcarlos/ToffeeX/overhang-fixer/first_order_so_issue.py", line 79, in <module>
    test_calculate_derivative_first_order()
  File "/home/mcarlos/ToffeeX/overhang-fixer/first_order_so_issue.py", line 65, in test_calculate_derivative_first_order
    derivative_fwd = calculate_derivative_first_order(
  File "/home/mcarlos/ToffeeX/overhang-fixer/first_order_so_issue.py", line 41, in calculate_derivative_first_order
    derivative = stencil(level_set, on=tuple(slices))
  File "/home/mcarlos/ToffeeX/overhang-fixer/.venv/lib/python3.10/site-packages/findiff/stencils.py", line 177, in __call__
    return self._apply_on_multi_slice(f, on)
  File "/home/mcarlos/ToffeeX/overhang-fixer/.venv/lib/python3.10/site-packages/findiff/stencils.py", line 216, in _apply_on_multi_slice
    result[tuple(base_mslice)] += coeff * f[tuple(off_mslice)]
ValueError: operands could not be broadcast together with shapes (3,3,3) (2,3,3) (3,3,3) 

If instead I build the slices trying to use what should be selecting the whole array, slices = [slice(None)] * ndims, base_mslice is built by the _canonical_slice method in a way that is not how python slices usually work, as it sets stop to zero

Level Set:
 [[[0.  0.  0. ]
  [0.  0.  0. ]
  [0.  0.  0. ]]

 [[0.5 0.5 0.5]
  [0.5 0.5 0.5]
  [0.5 0.5 0.5]]

 [[1.  1.  1. ]
  [1.  1.  1. ]
  [1.  1.  1. ]]]
Spacings: (np.float64(0.5), np.float64(0.5), np.float64(0.5))
slices: [slice(0, -1, None), slice(None, None, None), slice(None, None, None)]
level_set shape: (3, 3, 3)
Applying stencil along axis 0 with mode fwd
Stencil offsets: [(0, 0, 0), (1, 0, 0)]
Stencil partials: {(1, 0, 0): 1}
Stencil spacings: (np.float64(0.5), np.float64(0.5), np.float64(0.5))
Level set slice shape: (2, 3, 3)
################# Applying stencil #################
Applying stencil on multi-slice:
Base mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
Offset: (0, 0, 0)
Coefficient: -2.0
Initial off_mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
  Axis: 0 Offset value: 0
    Start: 0 Stop: 3
    Updated off_mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
  Axis: 1 Offset value: 0
    Start: 0 Stop: 0
    Updated off_mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
  Axis: 2 Offset value: 0
    Start: 0 Stop: 0
    Updated off_mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
Final off_mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
Adding to result at (slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)) the value coeff * f at (slice(0, 3, None), slice(0, 0, None), slice(0, 0, None))
Value to add: 
 []
Value before addition: 
 []
Offset: (1, 0, 0)
Coefficient: 2.0
Initial off_mslice: [slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)]
  Axis: 0 Offset value: 1
    Start: 1 Stop: 4
    Updated off_mslice: [slice(1, 4, None), slice(0, 0, None), slice(0, 0, None)]
  Axis: 1 Offset value: 0
    Start: 0 Stop: 0
    Updated off_mslice: [slice(1, 4, None), slice(0, 0, None), slice(0, 0, None)]
  Axis: 2 Offset value: 0
    Start: 0 Stop: 0
    Updated off_mslice: [slice(1, 4, None), slice(0, 0, None), slice(0, 0, None)]
Final off_mslice: [slice(1, 4, None), slice(0, 0, None), slice(0, 0, None)]
Adding to result at (slice(0, 3, None), slice(0, 0, None), slice(0, 0, None)) the value coeff * f at (slice(1, 4, None), slice(0, 0, None), slice(0, 0, None))
Value to add: 
 []
Value before addition: 
 []
Traceback (most recent call last):
  File "/home/mcarlos/ToffeeX/overhang-fixer/first_order_so_issue.py", line 81, in <module>
    test_calculate_derivative_first_order()
  File "/home/mcarlos/ToffeeX/overhang-fixer/first_order_so_issue.py", line 67, in test_calculate_derivative_first_order
    derivative_fwd = calculate_derivative_first_order(
  File "/home/mcarlos/ToffeeX/overhang-fixer/first_order_so_issue.py", line 43, in calculate_derivative_first_order
    derivative = stencil(level_set, on=tuple(slices))
  File "/home/mcarlos/ToffeeX/overhang-fixer/.venv/lib/python3.10/site-packages/findiff/stencils.py", line 177, in __call__
    return self._apply_on_multi_slice(f, on)
  File "/home/mcarlos/ToffeeX/overhang-fixer/.venv/lib/python3.10/site-packages/findiff/stencils.py", line 216, in _apply_on_multi_slice
    result[tuple(base_mslice)] += coeff * f[tuple(off_mslice)]
ValueError: operands could not be broadcast together with shapes (3,0,0) (2,0,0) (3,0,0) 

[cdelavegamartin]

If generating the slices as below, the behaviour is as expected,

slices = [slice(None, level_set.shape[0],None), slice(None, level_set.shape[1],None), slice(None, level_set.shape[2],None)]
    if mode == 'fwd':
        slices[axis] = slice(0,level_set.shape[axis]-1)
    elif mode == 'bwd':
        slices[axis] = slice(1,level_set.shape[axis])

Forward Derivative:
 [[[1. 1. 1.]
  [1. 1. 1.]
  [1. 1. 1.]]

 [[1. 1. 1.]
  [1. 1. 1.]
  [1. 1. 1.]]

 [[0. 0. 0.]
  [0. 0. 0.]
  [0. 0. 0.]]]

So this is not a blocking issue, but it is an unusual manipulation of slices

[maroba]

Thanks for reporting this! I fixed the bug and released version 0.12.2 just a few minutes ago.

[cdelavegamartin]

Tested and it's working correctly, thanks a lot