Присвойте значения массиву срез

Question

Я работаю с трехмерными массивами, и мне нужно взять дискретную производную вдоль заданной оси, которая состоит в том, чтобы взять разницу между двумя соседними плоскими срезами массива. До сих пор я сделал три функции, по одной для каждой производной.

Например, я беру производную вместе с x:

def take_derivative_x(data_box, msh):
    data_box_derivative = np.zeros_like(data_box)
    # the boundary derivatives
    data_box_derivative[0, :, :] \
          = (data_box[1, :, :] - data_box[0, :, :]) / msh[0]
    data_box_derivative[-1, :, :] \
          = (data_box[-1, :, :] - data_box[-2, :, :]) / msh[0]
    # the bulk derivatives
    for i in range(np.shape(data_box)[0] - 2):
        data_box_derivative[i+1, :, :] \
         = (data_box[i+2, :, :] - data_box[i, :, :]) / (2*msh[0])
    return data_box_derivative

Теперь я хочу создать функцию, которая делает это для массивы любого размера и принимает ось в качестве входных данных. Моя проблема в том, что после вычисления производной на плоском срезе я не могу найти способ поместить эти значения в новый массив, как это делается с помощью функции выше.

Вот что я сделал до сих пор:

def take_derivative(data_box, ax, msh):
    data_box_derivative = np.zeros_like(data_box)
    # the boundary derivatives
    data_box_derivative[0, :, :] # Problem here: I want [:,...,:, 0, :,...,:] with 0 at axis 'ax'
          = (np.take(data_box, 1, axis=ax) 
             - np.take(data_box, 0, axis=ax)) / msh[ax]
    data_box_derivative[-1, :, :] \ # Problem here
          = (np.take(data_box, -1, axis=ax)
             - np.take(data_box, -2, axis=ax)) / msh[ax]
    # the bulk derivatives
    for i in range(np.shape(data_box)[ax] - 2):
        data_box_derivative[i, :, :] \  # Problem here
              = (np.take(data_box, i+2, axis=ax) 
                 - np.take(data_box, i, axis=ax)) / (2*msh[ax])
    return data_box_derivative

Моя проблема сводится к поиску функции, которая для массивов

[[0 0 0]
 [0 0 0]
 [0 0 0]]

и [a b c], например, для i = 0, ax = 1

[[a 0 0]
 [b 0 0]
 [c 0 0]]

или для i = 1, топор = 0

[[0 0 0]
 [a b c]
 [0 0 0]]