Самый быстрый способ создать разреженную матрицу вида AT * diag(b) * A + C?

Я приближаюсь к 4-кратному ускорению вычисления произведения M.T * D * M из трех диагональных массивов. Если d0 и d1 — главная и верхняя диагональ M, а d — главная диагональ D, то следующий код создает M.T * D * M напрямую:

def make_tridi_bis(d0, d1, d, nc=10):
    d00 = d0*d0*d
    d11 = d1*d1*d
    d01 = d0*d1*d
    len_ = d0.size
    data = np.empty((3*len_ + nc,))
    indices = np.empty((3*len_ + nc,), dtype=np.int)
    # Fill main diagonal
    data[:2*nc:2] = d00[:nc]
    indices[:2*nc:2] = np.arange(nc)
    data[2*nc+1:-2*nc:3] = d00[nc:] + d11[:-nc]
    indices[2*nc+1:-2*nc:3] = np.arange(nc, len_)
    data[-2*nc+1::2] = d11[-nc:]
    indices[-2*nc+1::2] = np.arange(len_, len_ + nc)
    # Fill top diagonal
    data[1:2*nc:2] = d01[:nc]
    indices[1:2*nc:2] = np.arange(nc, 2*nc)
    data[2*nc+2:-2*nc:3] = d01[nc:]
    indices[2*nc+2:-2*nc:3] = np.arange(2*nc, len_+nc)
    # Fill bottom diagonal
    data[2*nc:-2*nc:3] = d01[:-nc]
    indices[2*nc:-2*nc:3] = np.arange(len_ - nc)
    data[-2*nc::2] = d01[-nc:]
    indices[-2*nc::2] = np.arange(len_ - nc ,len_)

    indptr = np.empty((len_ + nc + 1,), dtype=np.int)
    indptr[0] = 0
    indptr[1:nc+1] = 2
    indptr[nc+1:len_+1] = 3
    indptr[-nc:] = 2
    np.cumsum(indptr, out=indptr)

    return sparse.csr_matrix((data, indices, indptr), shape=(len_+nc, len_+nc))

Если ваша матрица M была в формате CSR, вы можете извлечь d0 и d1 как d0 = M.data[::2] и d1 = M.data[1::2], я модифицировал вашу игрушечную процедуру создания данных, чтобы она также возвращала эти массивы, и вот что я получаю:

In [90]: np.allclose((M.T * sparse.diags(d, 0) * M).A, make_tridi_bis(d0, d1, d).A)
Out[90]: True

In [92]: %timeit make_tridi_bis(d0, d1, d)
10 loops, best of 3: 124 ms per loop

In [93]: %timeit M.T * sparse.diags(d, 0) * M
1 loops, best of 3: 501 ms per loop

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

d_0 = np.zeros((len_ + nc,))
d_0[:len_] = d00
d_0[-len_:] += d11

d_1 = d01

Остальной код в этой функции просто напрямую строит трехдиагональную матрицу, так как вызов sparse.diags с приведенными выше данными в несколько раз медленнее.

Я попытался запустить ваш тестовый пример, и у меня возникли проблемы с файлом np.dot(N, M). Я не копался в этом, но я думаю, что у моей комбинации numpy/sparse (оба довольно новые) были проблемы с использованием np.dot на разреженных массивах.

Но H = -Hsig - z*M.T.dot(N.dot(M)) работает нормально. Это использует sparse dot.

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

In [37]: timeit Hsig,M,n,z=make_toy_data()
1 loops, best of 3: 2 s per loop

In [38]: timeit N = sparse.diags(1. / n ** 2, 0, format='csc')
1 loops, best of 3: 377 ms per loop

In [39]: timeit H = -Hsig - z*M.T.dot(N.dot(M))
1 loops, best of 3: 1.55 s per loop

H is a

<2000000x2000000 sparse matrix of type '<type 'numpy.float64'>'
    with 5999980 stored elements in Compressed Sparse Column format>