How to create a function which applies lag to a numpy array in python?

Here are 3 numpy arrays with some numbers:

import numpy as np

arr_a = np.array([[0.1, 0.2, 0.3, 0.4],
                  [0.2, 0.3, 0.4, 0.5],
                  [0.3, 0.4, 0.5, 0.6],
                  [0.4, 0.5, 0.6, 0.7],
                  [0.5, 0.6, 0.7, 0.8],
                  [0.6, 0.7, 0.8, 0.9]])

arr_b = np.array([[0.15, 0.25, 0.35, 0.45],
                  [0.35, 0.45, 0.55, 0.65],
                  [0.55, 0.65, 0.75, 0.85],
                  [0.75, 0.85, 0.95, 1.05],
                  [0.95, 1.05, 1.15, 1.25],
                  [1.15, 1.25, 1.35, 1.45]])

arr_c = np.array([[0.3, 0.6, 0.9, 1.2],
                  [0.6, 0.9, 1.2, 1.5],
                  [0.9, 1.2, 1.5, 1.8],
                  [1.2, 1.5, 1.8, 2.1],
                  [1.5, 1.8, 2.1, 2.4],
                  [1.8, 2.1, 2.4, 2.7]])

Each array has a shape of (6, 4). Where rows=person and columns=time(seconds), consider each row to represent a unique person and each column to represent acceleration associated with that person for that particular moment in time.

I would like to create a function called calc_acc_change which computes the change in acceleration according to a given lag value. I would like this function to take in an array and a lag value (default=2) which satisfies the following formula: acc_change(t) = acc(t) - acc(t- lag), where t=time. And I would like the output to remain an array.

I've started my function as follow, but I don't know how to complete it:

def calc_acc_change(array, lag=2):
   ...
   return acc_change

To test the function works, please input arr_a, arr_b and arr_c into the function and print the output.

Any help is much appreciated :)

Solution

Slice your input array based on the given lag and then subtract the sliced arrays.

import numpy as np

def calc_acc_change(array, lag=2):
    acc_change = np.zeros_like(array)
    acc_change[:, lag:] = array[:, lag:] - array[:, :-lag]
    return acc_change

arr_a = np.array([[0.1, 0.2, 0.3, 0.4],
                  [0.2, 0.3, 0.4, 0.5],
                  [0.3, 0.4, 0.5, 0.6],
                  [0.4, 0.5, 0.6, 0.7],
                  [0.5, 0.6, 0.7, 0.8],
                  [0.6, 0.7, 0.8, 0.9]])

arr_b = np.array([[0.15, 0.25, 0.35, 0.45],
                  [0.35, 0.45, 0.55, 0.65],
                  [0.55, 0.65, 0.75, 0.85],
                  [0.75, 0.85, 0.95, 1.05],
                  [0.95, 1.05, 1.15, 1.25],
                  [1.15, 1.25, 1.35, 1.45]])

arr_c = np.array([[0.3, 0.6, 0.9, 1.2],
                  [0.6, 0.9, 1.2, 1.5],
                  [0.9, 1.2, 1.5, 1.8],
                  [1.2, 1.5, 1.8, 2.1],
                  [1.5, 1.8, 2.1, 2.4],
                  [1.8, 2.1, 2.4, 2.7]])

print("Array A:\n", calc_acc_change(arr_a))
print("Array B:\n", calc_acc_change(arr_b))
print("Array C:\n", calc_acc_change(arr_c))