deepmodeling/jax-fem

Differentiable Finite Element Method with JAX

jax-fem's Issues

Infinite while loop for non-linear case

From tianjuxue/jax-am#28

Rename Cauchy BC to Robin BC

• According to Tianju, these are named wrong

Boundary conditions using mesh properties.

When defining boundary conditions for nodes, is there a way the interface can be extended to not only take in the node coordinates, but also the node index? Something like this:

def top_nodes(point, index):
    top = jnp.isclose(point[2], upper_surface_height, atol=TOLERANCE    # Nodes on the top surface of my mesh
    active = jnp.any(jnp.in1d(point.index, active_nodes))    # Nodes in my predefined active list
    return top & active

I am implementing a layer-by-layer model with element activation. I can deactivate elements above a certain height based on their centroid, but when applying the boundary condition, I cannot select only the nodes connected to my active elements with the XYZ position alone.

Any suggestions welcome!

Question about stresses for Linear_Elasticity demo

Hi, I noticed there was a recent commit made by @xwpken to the Linear_Elasticity demo that implemented post-processing for stresses in the beam. I'm just confused on how to interpret the data that has been computed at the end of the script (Sigma and sigma_average). I don't think they get saved into the VTU file, so I can't use Paraview to display it.

Could anyone help me out on how I can visualise them? Thank you!

Problem in the readme file.

Can you please change this in the readme (here)?

git clone [email protected]:tianjuxue/jax-fem.git

You can replace this with the HTTP cloning.

git clone https://github.com/tianjuxue/jax-fem.git

Otherwise people will get access denied.

Remove from jax.config import config, from fe.py

Hello everyone,
Is it possible to remove "from jax.config import config" from the files ?
I have to create a dockerfile with JAX-FEM, and it will not be working because of that.
It works well in a conda venv without "from jax.config import config", otherwise I get an error because of the update.
Thanks a lot

JAX-FEM Refactor

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

issue with petsc4py

I am trying to install the jax-fem on a windows platform. I am unable to install Petsc4py using pip install. I am not able to install using condo as well. It says package not found in the conda-forge channels. Is JAX-FEM not meant to work on windows? Please help as this package will be extremely useful for my research.

How to do parallel solver

Hi, I'm trying to create an hybrid model, that use your FEM-solver and a neural network.
To do that, I need to solve the same equations with different parameters (e.g the heat diffusion, where the diffusion coefficient and ic are different for each data).
I can't use vmap because the solver is using scipy and numpy, which isn't compatible.
Do you think the solver can be adapted so it can managed batches or can be pass into vmap ?
Thanks in advance for any idea !

Docs / Tutorial Notebook

It seems like there is currently no proper docs / tutorial for this package. I'm interested in applying this to some problems in my research, but it'd be very helpful to have some sort of tutorial for extending this to new problems.

Cheers!

"pip install -e ." Throws an error: subprocess-exited-with-error

Installing jax-fem with pip install -e . throws the error in the screenshot.

Updated README file and basic documentation

The current README and documentation largely refers to jax-am

Memory leakge in PETSc

from tianjuxue/jax-am#30

Need a direct solver

• PETSc has direct solvers
• We can easily implement them

Clean-up demos and applications

The purpose of these two different sub-directories are not very clear. It is likely that only one of them is necessary

Initial task list

• General cleanup + Flake8 file -> Surya
• Enable pip install in editable mode [Setup] -> Surya
• Environment file -> Surya
• All tests in one location -> Surya
• Baseline performance - Memory, Speed & accuracy - benchmark -> Igor
• GH workflow setup -> Igor

Question about the demo thermal_mechanical_full

Hello everyone,

I've been exploring this code with a keen interest in solving coupling problems. The code successfully passed the benchmark tests and have been able to run demos.hyperelasticity.example without any issues. However, I've encountered a stumbling block while trying to work with the thermal_mechanical_full demo. I'm facing a FileNotFoundError as detailed below:

FileNotFoundError: [Errno 2] No such file or directory: '/home/myname/Desktop/jax-fem/demos/thermal_mechanical_full/input/numpy/points.npy'

I'm wondering if there's something I might miss. Any guidance or suggestions would be greatly appreciated!

Create a uniform approach/template for post processing

Currently the post processing steps are a bit unorganized, and the templates for that aren't easy to find (example). It would be a big help to have a general template (or say a utils submodule) made available so that users can use/ modify them as per their need.

Save output files at user specified intervals

In cases of hyperelasticity where we need incremental load/displacement application, or in the case of rate or history dependent problems, the user would like to save output files at specific intervals. The current implementation 1cc776c makes this difficult as the solver handles the incremental loading internally and the solution saving step gets only the final solution. The user should be able to specify an output save (to disk) frequency and get the desired outputs.

Allow analytical sensitivities for known objectives

• To make it easier to run smaller problems, it might be beneficial to implement analytical formulas for simple objectives
  like compliance

Multi-step loading + auto-diff

• In the plasticity example, the load is applied in smaller steps. This is beneficial for easy convergence of iterative solvers for other cases as well.
• Solution:
  • Just as ad_wrapper performs implicit differentiation through a non-linear solve, we should allow either implicit differentiation through load steps!

PETSc solver not working with JAX transforms

• Currently, the solver [when using petsc] cannot be combined with any of Jax's transforms
• This prevents vmapping and higher order derivatives
• Solution:
  • Wrap petsc with Lineax or
  • Wrap it a pure-callback + custom-jvp