Question about memory and computational efficiency for solving large-scale systems of linear equations.

[DoTulip]

Hello JAX Team,
I am doing a finite element analysis of solid mechanics based on JAX, which needs to solve a linear equation of the form AX=B, and I need to get the corresponding gradient information. But there is a problem that has been bothering me. JAX does not seem to be able to implement automatic differentiation of sparse matrices yet, so full matrices need to be used. The memory and computational overhead brought by this are very expensive. For example, for a simple problem of discretizing by 200x100 mesh, the dimension of the corresponding matrix A reaches 40602x40602, and the memory overhead caused by this is unbearable. Is there any way to reduce memory and computational overhead while obtaining accurate gradient information?

[YouJiacheng]

You can express the sparse matrix A as an operator, i.e. R^n -> R^n.
For example, a tridiagnoal matrix:

a, b, c = ... # tridiagnoal
def operator_A(x):
  return a * x[1:] + b * x + c * x[:-1]

I think most of sparse matrices in FEM can be expressed in stencil computations, thus it is not very hard to re-express them.

[DoTulip]

Thank you for your reply! Forgive my stupidity, I'm sorry I didn't understand you well. Do you mean not assembling matrix A?

[YouJiacheng]

Yes, not assembling matrix A.

[DoTulip]

According to your answer, we do not need to assemble matrix A, but how to solve AX=B?

[YouJiacheng]

IIUC, all sparse linear systems are solved by cg or bicgstab or other similar methods.
cg
bicgstab
These two JAX function accept function A as input.

[DoTulip]

I understand! Very grateful for your help!

[jakevdp]

jax.experimental.sparse does support autodiff of sparse matrix operations. If you can add a small self-contained example of what sort of computation you're hoping to do, I may be able to help you translate it to sparse types.

[DoTulip]

You are so kind!
I have attached an executable code that I hope will reduce the memory and computational overhead of the function "fea_and_obj," which is mainly brought about by assembling s_K. I would like to reduce memory and computational overhead while using AD. In the case of using a full matrix, AD needs to track a very large matrix A, which is also an enormous burden for AD.

import sys
import jax.scipy.linalg as jsl
from jax import jit, value_and_grad
import jax.numpy as jnp
from numpy import float64
import scipy.io as scio
import time
from pathlib import Path  
import matplotlib.pyplot as plt
from functools import partial
from jax.config import config
config.update("jax_enable_x64", True)


def coordinate_and_preFEA(nelx, nely, total_ele, total_nod):
    nodegrd = jnp.array([[i for i in range(1, total_nod+1)]])
    nodegrd = nodegrd.reshape((nely+1, nelx+1), order='F')
    nodelast = nodegrd[1-1:-1, 1-1:-1]
    nodelast = nodelast.reshape((total_ele, 1), order='F')
    edofVec = 2*nodelast.T.reshape(-1, 1)+1
    edofMat = jnp.tile(edofVec, (1, 8))+jnp.tile(jnp.array(
        [0, 1, 2*nely+2, 2*nely+3, 2*nely+0, 2*nely+1, -2, -1]), (total_ele, 1))
    iK = jnp.kron(edofMat, jnp.ones((8, 1))).flatten()
    jK = jnp.kron(edofMat, jnp.ones((1, 8))).flatten()
    iK_int = iK.astype(jnp.int32)
    jK_int = jK.astype(jnp.int32)
    return iK_int-1, jK_int-1


def linear_stiffness_matrix(young, poisson, cc):
    young, poisson, h = young, poisson, cc
    k = jnp.array([1/2-poisson/6, 1/8+poisson/8, -1/4-poisson/12, -1/8+3*poisson/8,
                   -1/4+poisson/12, -1/8-poisson/8, poisson/6, 1/8-3*poisson/8])
    return h/(1-poisson**2)*jnp.array([[k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]],
                                       [k[1], k[0], k[7], k[6],
                                        k[5], k[4], k[3], k[2]],
                                       [k[2], k[7], k[0], k[5],
                                        k[6], k[3], k[4], k[1]],
                                       [k[3], k[6], k[5], k[0],
                                        k[7], k[2], k[1], k[4]],
                                       [k[4], k[5], k[6], k[7],
                                        k[0], k[1], k[2], k[3]],
                                       [k[5], k[4], k[3], k[2],
                                        k[1], k[0], k[7], k[6]],
                                       [k[6], k[3], k[4], k[1],
                                        k[2], k[7], k[0], k[5]],
                                       [k[7], k[2], k[1], k[4],
                                        k[3], k[6], k[5], k[0]]
                                       ], dtype=float64)

@partial(jit, static_argnums=(0, 1, 5))
def fea_and_obj(young, young_min, ke, xPhys_1D, freedofs, penal, iK, jK, fext_1D, s_K_predefined, dispTD_predefined):
    def kuf_solve(fext_1D_free, s_K_free):
        u_free = jsl.solve(s_K_free, fext_1D_free,
                           sym_pos=True, check_finite=False)
        return u_free

    young_modulus = young_min + (xPhys_1D) ** penal * (young-young_min)
    sK = ((ke.flatten()[jnp.newaxis]).T * young_modulus).flatten(order='F') # Stores the data in the stiffness matrix.
    idx = (iK, jK) # The coordinates of the stiffness matrix assembly.
    # Assembly stiffness matrix
    # How to reduce the memory and computational overhead of the following operations ?
    # If we don't use automatic differentiation, we can use scipy.sparse in scipy, so the memory and computational overhead are small.
    # *************************************************************************************
    s_K = s_K_predefined.at[idx].add(sK) # The stiffness matrix, corresponding to A in Ax=B
    s_K_free = s_K[freedofs, :][:, freedofs]
    fext_1D_free = fext_1D[freedofs] # Corresponding to B in Ax=B
    u_free = kuf_solve(fext_1D_free, s_K_free)
    dispTD = dispTD_predefined.at[freedofs].add(u_free) # displacement
    comp = fext_1D.T @ dispTD
    # *************************************************************************************
    return comp


def calculate_U(x):
    young = 1.
    young_min = 1e-9*young
    poisson = 0.3
    cc = 1.0
    penal = 3
    nelx = x.shape[1]
    nely = x.shape[0]
    total_ele, total_nod, total_dof = x.size, (nelx+1)*(nely+1), 2*(nelx+1)*(nely+1)
    fext = jnp.zeros(total_dof)
    dofs = jnp.arange(total_dof)
    # BC condition
    loaddof = total_dof-nely-1
    fext = fext.at[loaddof].set(-1.)
    fixeddofs = jnp.array(dofs[0:2*(nely+1):1])
    freedofs = jnp.setdiff1d(dofs, fixeddofs)
    # 
    iK, jK = coordinate_and_preFEA(nelx, nely, total_ele, total_nod) 
    ke = linear_stiffness_matrix(young, poisson, cc)
    xPhys_1D = x.T.flatten()
    s_K_predefined = jnp.zeros([total_dof, total_dof])
    dispTD_predefined = jnp.zeros([total_dof])
    # FEA & Obj
    obj = fea_and_obj(young, young_min, ke, xPhys_1D, freedofs, penal,
                      iK, jK, fext, s_K_predefined, dispTD_predefined)
    return obj


if __name__ == '__main__':

    x = jnp.ones([40, 80], dtype=jnp.float64)
    obj, gradout = value_and_grad(calculate_U)(x)
    print("compliance:", obj)
    print("gradout:", gradout)

[jakevdp]

Thanks - which part of this is sparse?

[DoTulip]

In scipy, when assembling the stiffness matrix s_K, we use the following code

s_K = scipy.sparse.coo_matrix((sK, (idx[0], idx[1])), shape=(Total_dof, Total_dof)).tocsc()
s_K_free = s_K[freedofs, :][:, freedofs]

Next, use the following code to solve the linear equations

u_free = scipy.sparse.linalg.spsolve(fext_1D_free, s_K_free)

In this way, both memory and computational overhead are small. How can I achieve a similar function in JAX and support AD?

[jakevdp]

Here's an example of performing a sparse linear solve with a sparse matrix in scipy and in JAX:

from jax.experimental import sparse
import scipy.sparse
import jax.scipy.sparse
import numpy as np
from functools import partial
import operator


A = scipy.sparse.coo_matrix(np.random.rand(5, 5))
x = np.random.rand(5)

A_jax = sparse.BCOO.from_scipy_sparse(A)
A_jax_op = lambda x: A_jax @ x

print(scipy.sparse.linalg.spsolve(A, x))
# [ 0.29048918  3.1450352  -0.00979041 -3.21662844  1.48401175]
print(jax.scipy.sparse.linalg.bicgstab(A_jax_op, x)[0])
# [ 0.29049     3.1450188  -0.00978914 -3.2166128   1.4840035 ]

I haven't read through your code closely enough to understand exactly how you're constructing your sparse matrix, but I would expect you could use something along these lines. One difference is that the scipy version is a direct solve, where the jax version is an iterative solve. I've been looking into wrapping cusolver_sparse in JAX to allow for direct sparse solves, but that is not yet available.

(Edit: I just opened #10590 which will make it so A_jax can be passed directly to bicgstab and friends instead of having to wrap it in a function)

[DoTulip]

This is really useful! Thank you very much for your patient reply! Looking forward to a more powerful JAX!

[soraros]

One can 'assemble the finite element system' by using A.at[(row, col)].add(data), or s_K = s_K_predefined.at[idx].add(sK) in you example code. This is not desirable as A is dense. You could instead write

from jax.experimental import sparse

def spsolve(row, col, data, b):
  A = sparse.coo.COO((data, row, col), shape=(size,) * 2)
  return jax.scipy.sparse.linalg.cg(partial(sparse.coo.coo_matvec, A), b)

[jakevdp]

I would avoid COO and use BCOO instead, because it has more compatibility with JAX transforms, and more implemented functionality.

[DoTulip]

Thanks @soraros for the answer! This works great for me!

[soraros]

@jakevdp Is that (prefer BCOO over COO) a piece of general advice or does it only reflects the current state of affairs? I'm asking because I see you put a lot of work into making the whole JAX sparse matrix story better, and am very curious about future directions.

[DoTulip]

Thanks to all the participants in this discussion, I really appreciate your help!

[adambomandel]

Hey DoTulip,

This looks like a framework for a topology optimization algorithm. I am currently working on a similar project, mind shooting me a message at [email protected]? Would love to see how you ended up solving this.

br
ADam

[DoTulip]

Hi Adam, sorry for the delay in replying to you (I haven't been using github for a while now, as well as taking a big vacation to deal with other personal matters). The method I ended up using was the one Jake introduced above, using iterative solutions. Recently, I started to study this problem again. It seems that direct solution was adopted in the latest discussion #18452 . I am ready to have a try.

By the way, I do topology optimization, and maybe you're a student of Professor sigmund's?

Best regards