CODEWITHSUNDEEP
pythonscipy
Nico Schlömer
Nico Schlömer

Reputation: 58791

vectorized computation of many matrix exponentials

I have many 2-by-2 matrices, A, A.shape == (2, 2, 7324), and I have to compute the matrix exponential of all of those. Unfortunately, scipy.linalg.expm only accepts one matrix at a time such that I'd have to loop over the computations,

import numpy
import scipy.linalg

numpy.random.seed(0)
A = numpy.random.rand(2, 2, 7324)

out = numpy.array([scipy.linalg.expm(A[:, :, k]) for k in range(A.shape[-1])])
out = numpy.moveaxis(out, 0, -1)

Is there a way to avoid this loop?

Edit: The corresponding scipy bug: #12838

Upvotes: 0

Views: 708

Answers (3)

prolyx
prolyx

Reputation: 491

As of SciPy release 1.9.0 (released in July 28, 22), you can pass scipy.linalg.expm an array where the last two dimensions are square—that is, an array with shape (..., n, n)—and it will compute the matrix exponential in a vectorized fashion.

The documentation is here.

Upvotes: 1

R K Rupesh
R K Rupesh

Reputation: 13

here's the code for the pade approximation. This is just for 1 array, you'll have to numpy vectorise/ cythonise it. or just use @njit from numba

def matrixexponen(n,a):
    q = 6
    a2 = a.copy ( )
    a_norm = np.linalg.norm ( a2, np.inf )
    ee = ( int ) ( np.log2 ( a_norm ) ) + 1
    s = max ( 0, ee + 1 )
    a2 = a2 / ( 2.0 ** s )
    x = a2.copy ( )
    c = 0.5
    e = np.eye ( n, dtype = np.complex64 ) + c * a2
    d = np.eye ( n, dtype = np.complex64 ) - c * a2
    p = True

    for k in range ( 2, q + 1 ):
        c = c * float ( q - k + 1 ) / float ( k * ( 2 * q - k + 1 ) )
        x = np.dot ( a2, x )
        e = e + c * x
        
        if ( p ):
            d = d + c * x
        else:
            d = d - c * x

        p = not p
    #  E -> inverse(D) * E
    e = np.linalg.solve ( d, e )
    #  E -> E^(2*S)
    for k in range ( 0, s ):
        e = np.dot ( e, e )
    return e

Upvotes: 0

ev-br
ev-br

Reputation: 26040

From glancing over the code, https://github.com/scipy/scipy/blob/v1.4.1/scipy/sparse/linalg/matfuncs.py#L550-L595 it seems that it really only assumes a single matrix. Therefore, if you want to avoid a python loop, you'll need to extract the relevant parts and either cythonize or numpy-vectorize them yourselves. Only catering to dense matrices can make it easier.

Upvotes: 0

Related Questions