Issues with installing and using jax-am

[tajtac]

Hello, I am having issues with installing and using jax-am and I would highly appreciate any help you can provide.

I was able to fully install jax-am on my Apple silicon macbook (by following the instructions at https://jax-am.readthedocs.io/en/latest/installation.html), but when I try to run the hyperelasticity example demos.fem.hyperelasticity.example, I get the following error:

[2023-07-20 01:01:53,936 - INFO] - Creating sparse matrix with scipy...
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/vt/Library/CloudStorage/OneDrive-purdue.edu/Research/jax-fem/jax-am/demos/fem/hyperelasticity/example.py", line 83, in <module>
    sol = solver(problem, use_petsc=True)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vt/Library/CloudStorage/OneDrive-purdue.edu/Research/jax-fem/jax-am/jax_am/fem/solver.py", line 823, in solver
    return solver_row_elimination(problem, linear, precond, initial_guess,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vt/Library/CloudStorage/OneDrive-purdue.edu/Research/jax-fem/jax-am/jax_am/fem/solver.py", line 358, in solver_row_elimination
    res_vec, A_fn = newton_update_helper(dofs)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vt/Library/CloudStorage/OneDrive-purdue.edu/Research/jax-fem/jax-am/jax_am/fem/solver.py", line 342, in newton_update_helper
    A_fn = get_A_fn(problem, use_petsc)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vt/Library/CloudStorage/OneDrive-purdue.edu/Research/jax-fem/jax-am/jax_am/fem/solver.py", line 300, in get_A_fn
    A = PETSc.Mat().createAIJ(size=A_sp_scipy.shape,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Mat.pyx", line 357, in petsc4py.PETSc.Mat.createAIJ
  File "petsc4py/PETSc/petscmat.pxi", line 836, in petsc4py.PETSc.Mat_AllocAIJ
  File "petsc4py/PETSc/petscmat.pxi", line 804, in petsc4py.PETSc.Mat_AllocAIJ_CSR
  File "petsc4py/PETSc/arraynpy.pxi", line 137, in petsc4py.PETSc.iarray_i
  File "petsc4py/PETSc/arraynpy.pxi", line 130, in petsc4py.PETSc.iarray
TypeError: Cannot cast array data from dtype('int64') to dtype('int32') according to the rule 'safe'

As an alternative, I tried installing jax-am on a linux cluster. However, there I receive another error during the installation:

  note: This error originates from a subprocess, and is likely not a problem with pip.
  ERROR: Failed building wheel for fenics-basix
Failed to build fenics-basix
ERROR: Could not build wheels for fenics-basix, which is required to install pyproject.toml-based projects

I believe I am following the instructions at https://jax-am.readthedocs.io/en/latest/installation.html verbatim. I start with a fresh conda environment, install pip and then do pip install .. Then I install petsc4py using conda. (I have also tried installing petsc4py before jax-am, but that didn't work either).

Thank you for your attention.

[tianjuxue]

petsc4py did a new release on Jun 30, 2023 that broke some code in JAX-AM. We have fixed the issue by @SNMS95. The updates will be merged to main shortly.

[SNMS95]

@tianjuxue Also, doesn't fenics_basix require a conda installation ?

[SNMS95]

@tajtac PR has been merged. PETSC will now run.