How to optimize and speed up this matrices multiplication in Python

Here is some improvements:

• beta * Q2 is computed twice, J can be used instead the second time.
• J * c is also computed multiple time while it can be done once. The same for I - J.
• B @ (J * c) @ E and B @ (J * c @ E) are mathematically equivalent, but the later is faster in your case and can also be computed once.
• CPython optimizes (almost) nothing and Numpy performs operations eagerly so doing 0.1 * (2 * Matrix / beta_square) actually compute a new matrix M2 = 2 * Matrix, then a new one M3 = M2 / beta_square, then another one M4 = 0.1 * M4. Creating many temporary matrices like this is expensive because it is a memory-bound operation and the memory bandwidth is pretty limited on modern machines (compared to the computing power), not to mention filling new temporary arrays is generally slower than already filled ones (because of virtual memory and more specifically page faults). Thus, it is faster to do (0.1 * 2 / beta_square) * Matrix (since multiplying float is much faster than multiplying big matrices).
• Some basic operations like np.argmax(c + tmp3.flatten(), axis=1) or np.max(np.absolute(Q3 - Q2)) can be easily accelerated using Numba. In fact, most
• In-place operations are generally faster than out-of-place ones (again, because of expensive temporary arrays). You can use them thanks to the out parameter of basic functions (e.g. np.multiply(A, B, out=C)). That being said, the benefit is quite small here since inv takes a significant time.
• Assuming B is not needed in the end of the loop, you can use np.linalg.solve instead as mentioned by Homer512. Solving a system is significantly faster for large matrices (O(n**3) versus O(n**2)) and often more accurate. See Don’t invert that matrix. For example, inv(I-J) @ b can be replaced by solve(I-J, b). The benefit of using solve is not so big though in your specific use-case because of the sparse I-J matrix.
• If B is actually used, in the end of the loop, then this is a bit more complex. Numba can help to write a relatively fast matrix inversion specifically for sparse matrices like in your use-case (since the one of Scipy turns out to be pretty slow).
• np.linalg.matrix_power(tmp0, 2) * c can also be optimized in Numba for sparse matrices.

Here is a (barely tested) implementation using Numba:

@nb.njit('(float64[:,::1], float64[::1])', parallel=True)
def compute_ini(a, b):
    n, m = a.shape
    assert b.size == m and m > 0
    res = np.empty(n, np.int64)
    for i in nb.prange(n):
        max_val, max_pos = a[i, 0] + b[0], 0
        for j in range(1, m):
            val = a[i, j] + b[j]
            if val > max_val:
                max_val = val
                max_pos = j
        res[i] = max_pos
    return res

@nb.njit('(float64[:,::1], float64[:,::1])', parallel=True)
def max_abs_diff(a, b):
    return np.max(np.absolute(a - b))

# Utility function for invert_sparse_matrix
@nb.njit
def invert_sparse_matrix_subkernel(b, out, i1, n, eps):
    for i2 in range(n):
        if i2 != i1:
            scale = b[i2, i1]
            if abs(scale) >= eps:
                for j in range(n):
                    b[i2, j] -= scale * b[i1, j]
                    out[i2, j] -= scale * out[i1, j]

@nb.njit('(float64[:,::1], float64[:,::1])', parallel=True)
def invert_sparse_matrix(a, out):
    eps = 1e-14
    n, m = a.shape
    assert n == m and out.shape == (n,n)

    b = np.empty((n, n))

    for i in nb.prange(n):
        out[i, :i] = 0.0
        out[i, i] = 1.0
        out[i, i+1:] = 0.0
        b[i, :] = a[i, :]

    for i1 in range(n):
        scale = 1.0 / b[i1, i1]
        if abs(scale) < eps:
            b[i1, :].fill(0.0)
            out[i1, :].fill(0.0)
            invert_sparse_matrix_subkernel(b, out, i1, n, eps)
        elif abs(scale-1.0) < eps:
            invert_sparse_matrix_subkernel(b, out, i1, n, eps)
        else:
            b[i1, :] *= scale
            out[i1, :] *= scale
            invert_sparse_matrix_subkernel(b, out, i1, n, eps)

@nb.njit('(float64[:,::1], float64[:,::1], float64[:,::1])', parallel=True)
def sparse_square_premult(a, premult, out):
    eps = 1e-14
    n, m = a.shape
    assert n == m
    assert premult.shape == (n, n)
    assert out.shape == (n, n)

    for i in nb.prange(n):
        out[i, :] = 0.0
        for j in range(n):
            if abs(a[i, j]) >= eps:
                for k in range(n):
                    out[i, k] += a[i, j] * a[j, k]
        out[i, :] *= premult[i, :]

def compute_numba():
    # parameters
    beta     =  0.98 
    alpha    =  0.03
    delta    =  0.1
    T        =  1000
    loop     =  1
    dif      =  1
    tol      =  1e-8

    kss      =  ((1 / beta - (1 - delta)) / alpha)**(1 / (alpha - 1))
    k        =  np.linspace(0.5 * kss, 1.8 * kss, T)

    k_reshaped   =  k.reshape(-1, 1)
    c            =  k_reshaped ** alpha + (1 - delta) * k_reshaped - k
    c[c<0]       =  1e-11
    c            =  np.log(c)
    beta_square  =  beta**2

    # multiplication
    I   =  np.identity(T)
    E   =  np.ones(T)[:,None]
    Q2  =  I

    J = np.empty((T, T))
    tmp0 = np.empty((T, T))
    tmp1 = np.empty((T, T))
    tmp2 = np.empty((T, 1))
    tmp3 = np.empty((T, 1))
    tmp4 = np.empty((T, T))
    B = np.empty((T, T))

    while np.any(dif > tol) and loop < 200:
        np.multiply(beta, Q2, out=J)
        np.subtract(I, J, out=tmp0)
        invert_sparse_matrix(tmp0, B)

        Q3 = np.zeros((T,T))
        np.multiply(J, c, out=tmp1)
        np.matmul(tmp1, E, out=tmp2)
        np.matmul(B, tmp2, out=tmp3)
        ini = compute_ini(c, tmp3.flatten())
        Q3[np.arange(T), ini] = 1

        factor = 0.1 * 2 / beta_square
        sparse_square_premult(tmp0, c, tmp4)
        np.add(B, (factor * tmp3) @ (tmp2 + B @ (tmp4 @ E)).T, out=B)

        dif = max_abs_diff(Q3, Q2)
        kcQ = k[ini]

        Q2 = Q3
        loop += 1

compute_numba()

Note that the invert_sparse_matrix function still takes a significant time (>50% on my machine) though it is about 3x faster than inv and about as fast as solve. It is an improvement of the naive inversion algorithm with few optimizations for very sparse matrices. It can certainly be optimized further (eg. using tiling) but this is certainly not trivial to do (especially for novice programmers).

Note the compilation time takes few seconds.

Overall, this implementation is about 4~5 times faster than the initial one on my i5-9600KF processor (6 cores).