Solver choices

[jorgensd]

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]

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]

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.