JAX-FEM Refactor

[itk22]

Below is a tasklist for the refactor of the JAX-FEM core code. The first stage of the refactor will focus on the core.py file. The proposed file structure is detailed on the refactor branch and this is where all the new changes should be merged.

Tasks
Stage 1:

• Fill in the core_fem file - @itk22
• Create a general object to represent kernels and integrate with the core code - @SNMS95 + @itk22 (support)
• Create wrappers for linear solvers based on Lineax - @SNMS95
• Checking all files and ammendments - Gawel