"Dynamic" N-dimensional finite difference in Python along an axis

Question

I have a function to compute the finite difference of a 1d np.array and I want to extrapolate to a n-d array.

The function is like this:

def fpp_fourth_order_term(U):
    """Returns the second derivative of fourth order term without the interval multiplier."""
    # U-slices
    fm2 = values[ :-4]
    fm1 = values[1:-3]
    fc0 = values[2:-2]
    fp1 = values[3:-1]
    fp2 = values[4:  ]

    return -fm2 + 16*(fm1+fp1) - 30*fc0 - fp2

It is missing the 4th order multiplier (1/(12*h**2)), but that is ok, because I will multiply when grouping the terms.

I would love to extend it as a N-dimensional. For that I would do the following changes:

def fpp_fourth_order_term(U, axis=0):
    """Returns the second derivative of fourth order term along an axis without the interval multiplier."""
    # U-slices

But here is the issue

    fm2 = values[ :-4]
    fm1 = values[1:-3]
    fc0 = values[2:-2]
    fp1 = values[3:-1]
    fp2 = values[4:  ]

This works fine in 1D, if is 2D along first axis for example I would have to change for something like:

    fm2 = values[:-4,:]
    fm1 = values[1:-3,:]
    fc0 = values[2:-2,:]
    fp1 = values[3:-1,:]
    fp2 = values[4:,:]

But along the second axis would be:

    fm2 = values[:,:-4]
    fm1 = values[:,1:-3]
    fc0 = values[:,2:-2]
    fp1 = values[:,3:-1]
    fp2 = values[:,4:]

The same applies to 3d, but has 3 possibilities and goes on and on. The return always works if the neighbors are correctly set.

    return -fm2 + 16*(fm1+fp1) - 30*fc0 - fp2

Of course axis cannot be larger than len(U.shape)-1 (I call this the dimension, is there any way to extract instead this snippet?

How can I do a elegant and pythonic approach for this coding problem?

Is there a better way to do it?

PS: Regarding np.diff and np.gradient, those do not work since the first one is first order and the second one is second order, I'm doing a fourth order approximation. In fact soon I finish this problem I will also generalize the order. But yes, I want to be able to do in any axis as np.gradient do.

Answer 1

A simple and effective solution is to use swapaxes at the very beginning and end of your procedure:

import numpy as np

def f(values, axis=-1):
    values = values.swapaxes(0, axis)

    fm2 = values[ :-4]
    fm1 = values[1:-3]
    fc0 = values[2:-2]
    fp1 = values[3:-1]
    fp2 = values[4:  ]

    return (-fm2 + 16*(fm1+fp1) - 30*fc0 - fp2).swapaxes(0, axis)

a = (np.arange(4*7*8)**3).reshape(4,7,8)
res = f(a, axis=1)
print(res)
print(res.flags)

Output:

# [[[ 73728  78336  82944  87552  92160  96768 101376 105984]
#   [110592 115200 119808 124416 129024 133632 138240 142848]
#   [147456 152064 156672 161280 165888 170496 175104 179712]]

#  [[331776 336384 340992 345600 350208 354816 359424 364032]
#   [368640 373248 377856 382464 387072 391680 396288 400896]
#   [405504 410112 414720 419328 423936 428544 433152 437760]]

#  [[589824 594432 599040 603648 608256 612864 617472 622080]
#   [626688 631296 635904 640512 645120 649728 654336 658944]
#   [663552 668160 672768 677376 681984 686592 691200 695808]]

#  [[847872 852480 857088 861696 866304 870912 875520 880128]
#   [884736 889344 893952 898560 903168 907776 912384 916992]
#   [921600 926208 930816 935424 940032 944640 949248 953856]]]

The result is even contiguous.

#   C_CONTIGUOUS : True
#   F_CONTIGUOUS : False
#   OWNDATA : False
#   WRITEABLE : True
#   ALIGNED : True
#   UPDATEIFCOPY : False