search

tianjuxue / jax-am

Additive manufacturing simulation with JAX.
https://jax-am.readthedocs.io/en/latest/
GNU General Public License v3.0
259 stars 56 forks source link

Solver choices #7

Open jorgensd opened 1 year ago

jorgensd commented 1 year ago

I looked through your paper: https://arxiv.org/pdf/2212.00964.pdf and it does not seem to mention the solver choices.

https://github.com/tianjuxue/jax-am/blob/main/jax_am/fem/tests/fenicsx_gold.py#L100-L106 uses un-preconditioned https://petsc.org/release/manualpages/KSP/KSPBICG/ with no null-space set, while as far as I can tell, JAX-FEM uses a Jacobi preconditioner with BICGSTAB (as far as I can tell from: https://github.com/tianjuxue/jax-am/blob/main/jax_am/fem/solver.py#L31)

It would be interesting to see more detailed breakdowns of the timers as well, as in:

  1. How much time does it take to solve the problem for the various problems
  2. How much time is spent assembling inside the Newton iteration
tianjuxue commented 1 year ago

Hi Jørgen,

Thanks for your comments. In the paper, when we did the comparison, JAX-FEM used BICGSTAB without preconditioner to be relatively fair. We will mention the solver choice in our next revised version of the paper.

JAX-FEM is now interfaced with PETSc, so I assume that the linear solver should have a similar performance as in FEniCS. For most problems, matrix assembling takes less than 20% of the total computational time (roughly speaking). There are exceptions, e.g., when we do crystal plasticity examples, computing the Jacobian matrix in Newton iteration could take a significant portion of time (e.g., >80% of the total time) due to that certain nonlinear equations occurring in constitutive relationships must be solved at quadrature level.

JAX so far only provides BICGSTAB and GMRES solvers. When GPU is available, their performance can be quite competitive. Otherwise, PETSc solver is preferred.

Best, Tianju

jorgensd commented 1 year ago

There are exceptions, e.g., when we do crystal plasticity examples, computing the Jacobian matrix in Newton iteration could take a significant portion of time (e.g., >80% of the total time) due to that certain nonlinear equations occurring in constitutive relationships must be solved at quadrature level.

It would be interesting to see a breakdown of where most of the time is spent.