Итак, количество диагоналей — это длина массива diagonals, а смещения основаны на количестве столбцов. Я придумал следующий массив смещений:

offsets = np.arange(len(diagonals))[::-1]-(a.shape[0]-1)

Где a — ваш исходный массив.

Затем вы можете создать диагонали, используя scipy.sparse.diags (это работает с неквадратными диагоналями) и добавить их в массив нулей. Затем окончательный результат необходимо будет перевернуть слева направо (аналогично исходному коду при создании диагоналей).

b = np.zeros_like(a)
for diagonal, offset in zip(diagonals, offsets):
    b += diags(diagonal, offset, shape=a.shape)
b = np.fliplr(b)

Трудно представить себе, что происходит, просто прочитав код и увидев усеченные примеры. Вы уже какое-то время работали над этим, а я (и, вероятно, большинство ваших читателей) — нет. Так что многое из дальнейшего может быть очевидным для вас, но новым для тех, кто не проводил много времени с диагоналями.

Я собираюсь немного упростить ситуацию - начну с квадратного массива и возьму прямые диагонали - без переворота и разворота.

In [78]: arr = np.arange(1,17).reshape(4,4)

In [79]: N, M = arr.shape
    ...: rangeNMbeg =    -M + 2  if N < M  else   -N + 1 
    ...: rangeNMend =      N + 1 if N < M  else    M
    ...: offsets = np.arange(rangeNMbeg, rangeNMend)
    ...: diagonals = [ np.diagonal(arr, offset=i) for i in offsets ]

In [80]: offsets
Out[80]: array([-3, -2, -1,  0,  1,  2,  3])

In [81]: diagonals
Out[81]: 
[array([13]),
 array([ 9, 14]),
 array([ 5, 10, 15]),
 array([ 1,  6, 11, 16]),
 array([ 2,  7, 12]),
 array([3, 8]),
 array([4])]

С помощью этой пары смещений и диагоналей np.diag(d,o) создает массив (4,4), который можно суммировать, чтобы воссоздать оригинал:

In [82]: res = np.zeros_like(arr)
    ...: for o,d in zip(offsets, diagonals):
    ...:     res += np.diag(d,o)
    ...:     

In [83]: res
Out[83]: 
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12],
       [13, 14, 15, 16]])

Следующий вопрос: могу ли я обобщить исходную форму?

Одна из проблем при выполнении того же самого с массивом (5,3) заключается в том, что некоторые из np.diag(d,o) — это (5,5), один — (4,4), а остальные (3,3). np.diag не позволяет нам указать форму 2d-массива. Я недостаточно использую диагональные функции, чтобы знать, подойдет ли какая-нибудь из них. Создание такого массива не составит труда выполнить «реверс-инжиниринг» np.diag.

Внутренне np.diag и np.diagflat используют отображение типа

res[:n-k].flat[i::n+1] = v   # or
res.flat[fi] = v

Итак, если res форма не квадратная, нам нужно проработать детали этого flat отображения.

scipy.sparse

Как отметил @jared, использование scipy.sparse упрощает это. Для случая (5,3) (без переворотов и т. д.):

In [130]: offsets
Out[130]: array([-4, -3, -2, -1,  0,  1,  2])

In [131]: diagonals
Out[131]: 
[array([13]),
 array([10, 14]),
 array([ 7, 11, 15]),
 array([ 4,  8, 12]),
 array([1, 5, 9]),
 array([2, 6]),
 array([3])]

In [132]: from scipy import sparse    
In [133]: sparse.diags(diagonals, offsets, shape=(5,3)).A
Out[133]: 
array([[ 1.,  2.,  3.],
       [ 4.,  5.,  6.],
       [ 7.,  8.,  9.],
       [10., 11., 12.],
       [13., 14., 15.]])

Разреженные матрицы часто имеют строго диагональную структуру (например, в задачах конечных разностей), поэтому этот пакет имеет формат dia_matrix.

dia_matrix((data, offsets), shape=(M, N))

diags — конструктор, создающий правильный data:

In [147]: M = sparse.diags(diagonals, offsets, shape=(5,3))    
In [148]: M
Out[148]: 
<5x3 sparse matrix of type '<class 'numpy.float64'>'
    with 15 stored elements (7 diagonals) in DIAgonal format>

In [149]: M.offsets
Out[149]: array([-4, -3, -2, -1,  0,  1,  2])

In [150]: M.data
Out[150]: 
array([[13.,  0.,  0.,  0.],
       [10., 14.,  0.,  0.],
       [ 7., 11., 15.,  0.],
       [ 4.,  8., 12.,  0.],
       [ 1.,  5.,  9.,  0.],
       [ 0.,  2.,  6.,  0.],
       [ 0.,  0.,  3.,  0.]])

Обратите внимание, что data — это квадратный массив с теми же значениями, что и список diagonals, но с некоторым дополнением. Код, который сопоставляет список diagonals с этим дополненным data, является Python, и его можно рассматривать как np.diags [источник].

sparse также имеет код для быстрого преобразования этого dia_matrix формата в csr, который используется для большинства вычислений. Это преобразование также используется при рендеринге матрицы в виде плотного массива.

Numpy 1.25

Исправление первоначального подхода

Во-первых, давайте заставим работать первоначальный подход. Есть две важные части: разделение данных на диагонали и их обратное объединение.

Наименьшее и наибольшее диагональное смещение — -(hight-1) и width-1 соответственно. Таким образом, мы можем упростить разбиение на диагонали следующим образом:

def get_diagonals(array):
    '''Return all the top-right to bottom-left diagonals of an array
    starting from the top-left corner.
    '''
    height, width = array.shape
    fliplr_arr = np.fliplr(array)
    return [fliplr_arr.diagonal(offset).copy()    # Starting in NumPy 1.9 diagonal returns a read-only view on the original array
            for offset in range(width-1, -height, -1)]

Собрав все диагонали в список, пронумеруем их по индексу. Теперь для всех диагоналей с индексом меньше ширины массива индекс элемента внутри диагонали равен его первому индексу в массиве. Для других диагоналей позиция элемента внутри диагонали плюс высота начальной точки диагонали дают высоту элемента в массиве. Начальная точка диагонали равна разнице номера диагонали и последнего возможного индекса во втором измерении массива (см. рисунок ниже).

Итак, мы можем переписать вторую часть кода следующим образом:

def from_diagonals(diagonals, shape, dtype):
    arr = np.empty(shape, dtype)
    height, width = shape
    for h in range(height): 
        arr[h, :] = [
            diagonal[h if d < width else h - (d - (width-1))]
            for d, diagonal in enumerate(diagonals[h:h+width], start=h)
        ]
    return arr

Теперь весь процесс работы выглядит следующим образом:

arr = ...
diagonals = get_diagonals(arr)
for diag in diagonals:
    ...
new_arr = from_diagonals(diagonals, arr.shape, arr.dtype)

Работа с индексами

Мы можем создать функцию для доступа к данным по диагональным индексам следующим образом:

def get_diagonal_function(height, width, base=0, antidiagonal=False):
    from functools import partial
    index = np.mgrid[:height, :width]
    if antidiagonal:
        index = index[..., ::-1]    # flip horizontal indices left-to-right
    _diag = partial(index.diagonal, axis1=1, axis2=2) 
    max_offset = width - 1
    shift = max_offset + base
    diagonal = lambda x: _diag(shift - x)
    diagonal.min = base
    diagonal.max = (height-1) + (width-1) + base
    diagonal.off = _diag
    return diagonal

Пример выделения всех антидиагоналей:

diagonal = get_diagonal_function(*arr.shape, base=1, antidiagonal=True)
diagonals = [arr[*diagonal(x)] for x in range(diagonal.min, 1+diagonal.max)]
main_antidiagonal = diagonal.off(0)

Пример переворота данных по антидиагоналям:

arr = [
    [ 1,  2,  4],
    [ 3,  5,  7],
    [ 6,  8, 10],
    [ 9, 11, 13],
    [12, 14, 15]
]
arr = np.array(arr)
diagonal = get_diagonal_function(*arr.shape, antidiagonal=True)

for x in range(1+diagonal.max):
    index = tuple(diagonal(x))
    arr[index] = arr[index][::-1]

print(arr)

[[ 1  3  6]
 [ 2  5  9]
 [ 4  8 12]
 [ 7 11 14]
 [10 13 15]]

Преобразование координат

Предупреждение: приведенный ниже код не был должным образом протестирован; он не предназначен для быстрых вычислений; это всего лишь специальное доказательство концепции.

В случае антидиагоналей, описанных в исходном посте, преобразование можно выполнить следующим образом (см. рисунок выше):

def diag_to_array(d, x, shape, direction='top-bottom,right-left'):
    height, width = shape
    h = x if d < width else x + (d - (width-1))
    w = d - h
    return h, w

В целом конвертация напоминает аффинное преобразование (с обрезкой по краям, если можно так сказать). Но нам нужно знать направление и порядок диагоналей, которых может быть 16 вариантов. Их можно задать векторами, указывающими направление отсчета диагоналей и элементов вдоль них. Эти векторы также можно использовать для определения нового источника.

def get_origin(d, e, shape):
    height, width = np.array(shape) - 1    # max coordinate values along dimensions
    match d, e:
        case ((1, x), (y, 1)) | ((x, 1), (1, y)):
            origin =  (0, 0)
        case ((1, x), (y, -1)) | ((x, -1), (1, y)):
            origin =  (0, width)
        case ((-1, x), (y, 1)) | ((x, 1), (-1, y)):
            origin =  (height, 0)
        case ((-1, x), (y, -1)) | ((x, -1), (-1, y)):
            origin =  (height, width)
    return np.array(origin)

Примечания:

• d — один из {(1, 0), (-1, 0), (0, 1), (0, -1)} векторов, показывающий направление диагональной числовой координаты;
• e — один из векторов {(1, 1), (1, -1), (-1, 1), (-1, -1)}, показывающий направление нумерации элементов по диагонали.

Что касается преобразований координат, то их удобнее реализовать в классе:

class Transfomer:
    def __init__(self, d, e, shape):
        self.d = np.array(d)
        self.e = np.array(e)
        self.shape = np.array(shape)
        self.A = np.stack([self.d, self.e])
        self.origin = get_origin(d, e, shape)
        self.edge = abs(self.d @ self.shape) - 1
        
    def diag_to_array(self, ndiagonal, element):
        '''Return the array coordiantes where:
        ndiagonal: the number of a diagonal
        element: the element index on the diagonal
        '''
        if ndiagonal > self.edge:
            element = element + ndiagonal - self.edge
        elif ndiagonal < 0:
            element = element - ndiagonal
        diag_coords = np.array([ndiagonal, element])
        array_coords = diag_coords @ self.A + self.origin
        return array_coords

    def array_to_diag(self, *args, **kwargs):
        raise NotImplementedError

Пример:

d = (0, 1)     # take anti-diagonals from left to right starting from the top-left corner
e = (1, -1)    # and index their elements from top-right to bottom-left
arr = np.arange(5*3).reshape(5, 3)
coords = Transfomer(d, e, arr.shape)

diagonal = get_diagonal_function(*arr.shape, base=1, antidiagonal=True)
diagonals = [arr[*diagonal(x)] for x in range(diagonal.min, 1+diagonal.max)]

for ndiag, element in [(0, 0), (2, 1), (3, 0), (5, 1), (6, 0)]:
    array_point = coords.diag_to_array(ndiag, element)
    try:
        assert diagonals[ndiag][element] == arr[*array_point], f'{ndiag=}, {element=}'
    except AssertionError as e:
        print(e)        

Работа с диагоналями как с возможностью записи

Если структура массива является прямой и простой (т. е. представляет собой непрерывную последовательность данных без каких-либо трюков с шагом или сложных транспозиций), мы можем попытаться (с некоторой осторожностью) сделать диагональное представление доступным для записи. Тогда все изменения будут произведены с исходными данными, и вам не придется заново собирать диагонали. Сделаем зеркальное отображение антидиагоналей, например:

def diag_range(shape, asnumpy=False):
    '''Returns a sequence of numbers arranged diagonally
    in a given shape as separate diagonals
    '''
    diagonals = []
    dmin, dmax = min(shape), max(shape)
    start = 1
    for n in range(dmin):
        end = start + n +1
        diagonals.append([*range(start, end)])
        start = end
    for n in range(dmin, dmax):
        end = start + dmin
        diagonals.append([*range(start, end)])
        start = end
    for n in range(dmin-1, 0, -1):
        end = start + n 
        diagonals.append([*range(start, end)])
        start = end
    if asnumpy:
        from numpy import array
        diagonals = list(map(array, diagonals))
    return diagonals

def from_diagonals(diagonals, shape, dtype):
    '''Returns an array constructed by the given diagonals'''
    arr = np.empty(shape, dtype)
    height, width = shape
    for h in range(height): 
        row = []
        for w, diag in enumerate(diagonals[h:h+width], start=h): 
            index = h - (w >= width) * (w - (width-1))
            row.append(diag[index])
        arr[h, :] = row
    return arr

h, w = shape = (6,3)
arr = from_diagonals(diag_range(shape), shape, 'int')
print('Initial array:'.upper(), arr, '', sep='\n')

fliplr_arr = np.fliplr(arr)
diagonals = [fliplr_arr.diagonal(i) for i in range(-h+1, w)]
for d in diagonals:
    d.flags.writeable = True    # use with caution
    d[:] = d[::-1]
print('Array with flipped diagonals:'.upper(), arr, sep='\n')

INITIAL ARRAY:
[[ 1  2  4]
 [ 3  5  7]
 [ 6  8 10]
 [ 9 11 13]
 [12 14 16]
 [15 17 18]]

ARRAY WITH FLIPPED DIAGONALS:
[[ 1  3  6]
 [ 2  5  9]
 [ 4  8 12]
 [ 7 11 15]
 [10 14 17]
 [13 16 18]]

Простое решение — получить доступ к диагоналям с помощью индексов. Доступ к записям напрямую с помощью np.digonal невозможен, поскольку np.digonal менялся дважды в версиях 1.7 и 1.9. В дополнение к другим предоставленным ответам эта версия также работает с массивами более высокой размерности, а не только с 2D-массивами. Более того, он не требует изменения параметра флага numpy, который служит определенной цели.

Основная функция решения

import numpy as np

def diagonal_indexes(shape, axis1=0, axis2=1, direction='normal'):
    idx_array = np.arange(np.prod(shape), dtype=int).reshape(shape)
    if direction == 'anti':
        idx_array = np.fliplr(idx_array)
    
    return {i: np.diagonal(idx_array, offset=i, axis1=axis1, axis2=axis2) for i in range(1-idx_array.shape[axis1], idx_array.shape[axis2])}

# diagonal_indexes((2,3))
# -> {-1: array([3]), 0: array([0, 4]), 1: array([1, 5]), 2: array([2])}

• Эта функция возвращает словарь, чтобы позже облегчить доступ к нужным диагоналям.
• Ключи словаря представляют собой смещения для np.diagonal. Следовательно, 0 — это главная диагональ или ведущая диагональ. Этот подход применим к массивам с более чем двумя осями. Следуя соглашениям np.diagonal, axis1 и axis2 указывают двумерные подмассивы, из которых извлекаются диагонали. Кроме того, axis1 и axis2 определяют порядок диагоналей, то есть смещение.
• Я буду использовать normal и anti для диагональных направлений вместо left и right.

Пример использования

Два примера иллюстрируют его функциональность.

Пример 1: Написание диагоналей

array = np.zeros((4,3))
dict_idx = diagonal_indexes(array.shape)


# access a single value: array.flat[dict_idx[idx_diagonal], idx_on_diagonal]
array.flat[dict_idx[0][1]] = 40
array.flat[dict_idx[-1][2]] = 10

# access a whole diagonal
array.flat[dict_idx[1]] = np.arange(len(dict_idx[1]))+1

print(array)
# [[ 0.  1.  0.]
#  [ 0. 40.  2.]
#  [ 0.  0.  0.]
#  [ 0.  0. 10.]]

Пример 2: диагональные индексы

def get_array(*shape, **kwargs):
    return np.arange(np.prod(shape), **kwargs).reshape(shape)

array = get_array(5, 4)

for kwargs in [dict(direction='normal', axis1=0, axis2=1), 
               dict(direction='normal', axis1=1, axis2=0), 
               dict(direction='anti', axis1=0, axis2=1), 
               dict(direction='anti', axis1=1, axis2=0)]:
    for i, idx in diagonal_indexes(array.shape, **kwargs).items():
        array.flat[idx] = i
    
    out_str = str(kwargs)[1:-1].replace(': ','=').replace("'","")
    print(f"diagonal_indexes(array.shape, {out_str})", '\n', array)
    print()

дает:

diagonal_indexes(array.shape, direction=normal, axis1=0, axis2=1) 
 [[ 0  1  2  3]
 [-1  0  1  2]
 [-2 -1  0  1]
 [-3 -2 -1  0]
 [-4 -3 -2 -1]]

diagonal_indexes(array.shape, direction=normal, axis1=1, axis2=0) 
 [[ 0 -1 -2 -3]
 [ 1  0 -1 -2]
 [ 2  1  0 -1]
 [ 3  2  1  0]
 [ 4  3  2  1]]

diagonal_indexes(array.shape, direction=anti, axis1=0, axis2=1) 
 [[ 3  2  1  0]
 [ 2  1  0 -1]
 [ 1  0 -1 -2]
 [ 0 -1 -2 -3]
 [-1 -2 -3 -4]]

diagonal_indexes(array.shape, direction=anti, axis1=1, axis2=0) 
 [[-3 -2 -1  0]
 [-2 -1  0  1]
 [-1  0  1  2]
 [ 0  1  2  3]
 [ 1  2  3  4]]