CODEWITHSUNDEEP
pythonnumpyscipy
Nico Schlömer
Nico Schlömer

Reputation: 58791

Solving a large number of small linear systems

I need to solve a large number of 3x3 symmetric, postive-definite systems with Python. So far, I did

res = numpy.zeros(n)
for k, obj in enumerate(data_array):
    # construct A, rhs, idx from obj
    res[idx] += numpy.linalg.solve(A, rhs)

This produces the correct result, however is also quite slow if n is large. (Well... Yeah.) Perhaps 3x3 isn't a problem size where calling solve() makes much sense.

Any hints?

Upvotes: 1

Views: 206

Answers (1)

user2357112
user2357112

Reputation: 280562

In NumPy 1.8 and later, numpy.linalg.solve actually broadcasts. For numpy.linalg.solve(a, b), if b.ndim == a.ndim - 1, it will perform a broadcasted matrix-vector solve; otherwise, it'll do a broadcasted matrix-matrix solve. (This decision criterion isn't documented; I had to look at the source.)

If you can efficiently construct a stack of As and rhss, you can call solve once and avoid a Python loop.

Upvotes: 4

Related Questions