Эффективный плотный аналог scipy.sparse.diags

Question

scipy.sparse.diags позволяет мне вводить несколько диагональных векторов вместе с их расположением для построения матрицы, такой как

from scipy.sparse import diags
vec = np.ones((5,))
vec2 = vec + 1
diags([vec, vec2], [-2, 2])

Я ищу эффективный способ сделать то же самое, но построить плотную матрицу вместо DIA. np.diag поддерживает только одну диагональ. Какой эффективный способ построить плотную матрицу из нескольких диагональных векторов?

Ожидаемый результат: аналогично np.array(diags([vec, vec2], [-2, 2]).todense())

Answer 1

Один из способов заключается в том, чтобы индексировать в плоский массив вывода с помощью шага N+1:

import numpy as np
from scipy.sparse import diags
from timeit import timeit

def diags_pp(vecs, offs, dtype=float, N=None):
    if N is None:
        N = len(vecs[0]) + abs(offs[0])
    out = np.zeros((N, N), dtype)
    outf = out.reshape(-1)
    for vec, off in zip(vecs, offs):
        if off<0:
            outf[-N*off::N+1] = vec
        else:
            outf[off:N*(N-off):N+1] = vec
    return out

def diags_sp(vecs, offs):
    return diags(vecs, offs).A

for N, k in [(10, 2), (100, 20), (1000, 200)]:
    print(N)
    O = np.arange(-k,k)
    D = [np.arange(1, N+1-abs(o)) for o in O]
    for n, f in list(globals().items()):
        if n.startswith('diags_'):
            print(n.replace('diags_', ''), timeit(lambda: f(D, O), number=10000//N)*N)
            if n != 'diags_sp':
                assert np.all(f(D, O) == diags_sp(D, O))

Пример выполнения:

10
pp 0.06757194991223514
sp 1.9529316504485905
100
pp 0.45834919437766075
sp 4.684177896706387
1000
pp 23.397524026222527
sp 170.66762899048626

Answer 2

В случае с Полом Панцером (10,2)

In [107]: O                                                                     
Out[107]: array([-2, -1,  0,  1])
In [108]: D                                                                     
Out[108]: 
[array([1, 2, 3, 4, 5, 6, 7, 8]),
 array([1, 2, 3, 4, 5, 6, 7, 8, 9]),
 array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10]),
 array([1, 2, 3, 4, 5, 6, 7, 8, 9])]

Диагонали имеют разную длину.

sparse.diags преобразует это в sparse.dia_matrix:

In [109]: M = sparse.diags(D,O)                                                 
In [110]: M                                                                     
Out[110]: 
<10x10 sparse matrix of type '<class 'numpy.float64'>'
    with 36 stored elements (4 diagonals) in DIAgonal format>
In [111]: M.data                                                                
Out[111]: 
array([[ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  0.,  0.],
       [ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9.,  0.],
       [ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.],
       [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9.]])

Здесь рваный список диагоналей был преобразован в дополненный 2d массив.Это может быть удобным способом определения диагоналей, но это не особенно эффективно.Для большинства вычислений его необходимо преобразовать в формат csr:

In [112]: timeit sparse.diags(D,O)                                              
99.8 µs ± 3.65 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [113]: timeit sparse.diags(D,O, format='csr')                                
371 µs ± 155 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Используя np.diag, я могу построить тот же массив с итерацией

np.add.reduce([np.diag(v,k) for v,k in zip(D,O)])

In [117]: timeit np.add.reduce([np.diag(v,k) for v,k in zip(D,O)])              
39.3 µs ± 131 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

и функцией Пола:

In [120]: timeit diags_pp(D,O)                                                  
12.3 µs ± 24.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Ключевой шаг в np.diags - это плоское задание:

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

По сути, это то же самое, что и задания Павла outf.Таким образом, функциональность в основном одинакова, назначая каждую диагональ с помощью среза.Пол оптимизирует его.

Создание массива M.data (Out[111]) также требует копирования массивов D в двумерный массив - но с разными фрагментами.