{"id":13458556,"url":"https://github.com/deepmodeling/jax-fem","last_synced_at":"2025-05-15T08:05:52.243Z","repository":{"id":197818461,"uuid":"699410096","full_name":"deepmodeling/jax-fem","owner":"deepmodeling","description":"Differentiable Finite Element Method with JAX","archived":false,"fork":false,"pushed_at":"2025-03-18T07:49:13.000Z","size":82898,"stargazers_count":363,"open_issues_count":31,"forks_count":58,"subscribers_count":12,"default_branch":"main","last_synced_at":"2025-04-07T04:09:07.440Z","etag":null,"topics":["differentiable-programming","finite-element-methods","jax","topology-optimization"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"gpl-3.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/deepmodeling.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-10-02T15:20:54.000Z","updated_at":"2025-04-04T11:43:37.000Z","dependencies_parsed_at":"2025-03-24T02:11:38.034Z","dependency_job_id":"d6e64347-0f9e-4c46-9e09-bf721d42dfaa","html_url":"https://github.com/deepmodeling/jax-fem","commit_stats":null,"previous_names":["snms95/jax-fem","deepmodeling/jax-fem"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2Fjax-fem","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2Fjax-fem/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2Fjax-fem/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2Fjax-fem/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/deepmodeling","download_url":"https://codeload.github.com/deepmodeling/jax-fem/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248871132,"owners_count":21175161,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["differentiable-programming","finite-element-methods","jax","topology-optimization"],"created_at":"2024-07-31T09:00:53.804Z","updated_at":"2025-04-14T11:28:17.196Z","avatar_url":"https://github.com/deepmodeling.png","language":"Python","funding_links":[],"categories":["Others","simulation","Libraries"],"sub_categories":[],"readme":"A GPU-accelerated differentiable finite element analysis package based on [JAX](https://github.com/google/jax). Used to be part of the suite of open-source python packages for Additive Manufacturing (AM) research, [JAX-AM](https://github.com/tianjuxue/jax-am).\n\n## Finite Element Method (FEM)\n![Github Star](https://img.shields.io/github/stars/deepmodeling/jax-fem)\n![Github Fork](https://img.shields.io/github/forks/deepmodeling/jax-fem)\n![License](https://img.shields.io/github/license/deepmodeling/jax-fem)\n\nFEM is a powerful tool, where we support the following features\n\n- 2D quadrilateral/triangle elements\n- 3D hexahedron/tetrahedron elements\n- First and second order elements\n- Dirichlet/Neumann/Robin boundary conditions\n- Linear and nonlinear analysis including\n  - Heat equation\n  - Linear elasticity\n  - Hyperelasticity\n  - Plasticity (macro and crystal plasticity)\n- Differentiable simulation for solving inverse/design problems __without__ deriving sensitivities by hand, e.g.,\n  - Topology optimization\n  - Optimal thermal control\n- Integration with PETSc for solver choices\n\n**Updates** (Dec 11, 2023):\n\n- We now support multi-physics problems in the sense that multiple variables can be solved monolithically. For example, consider running  `python -m applications.stokes.example`\n- Weak form is now defined through  volume integral and surface integral. We can now treat body force, \"mass kernel\" and \"Laplace kernel\" in a unified way through volume integral, and treat \"Neumann B.C.\" and \"Robin B.C.\" in a unified way through surface integral. \n\n\u003cp align=\"middle\"\u003e\n  \u003cimg src=\"images/ded.gif\" width=\"600\" /\u003e\n\u003c/p\u003e\n\u003cp align=\"middle\"\u003e\n    \u003cem \u003eThermal profile in direct energy deposition.\u003c/em\u003e\n\u003c/p\u003e\n\n\u003cp align=\"middle\"\u003e\n  \u003cimg src=\"images/von_mises.png\" width=\"400\" /\u003e\n\u003c/p\u003e\n\u003cp align=\"middle\"\u003e\n    \u003cem \u003eLinear static analysis of a bracket.\u003c/em\u003e\n\u003c/p\u003e\n\n\u003cp align=\"middle\"\u003e\n  \u003cimg src=\"images/polycrystal_grain.gif\" width=\"360\" /\u003e\n  \u003cimg src=\"images/polycrystal_stress.gif\" width=\"360\" /\u003e\n\u003c/p\u003e\n\u003cp align=\"middle\"\u003e\n    \u003cem \u003eCrystal plasticity: grain structure (left) and stress-xx (right).\u003c/em\u003e\n\u003c/p\u003e\n\n\u003cp align=\"middle\"\u003e\n  \u003cimg src=\"images/stokes_u.png\" width=\"360\" /\u003e\n  \u003cimg src=\"images/stokes_p.png\" width=\"360\" /\u003e\n\u003c/p\u003e\n\u003cp align=\"middle\"\u003e\n    \u003cem \u003eStokes flow: velocity (left) and pressure(right).\u003c/em\u003e\n\u003c/p\u003e\n\n\u003cp align=\"middle\"\u003e\n  \u003cimg src=\"images/to.gif\" width=\"600\" /\u003e\n\u003c/p\u003e\n\u003cp align=\"middle\"\u003e\n    \u003cem \u003eTopology optimization with differentiable simulation.\u003c/em\u003e\n\u003c/p\u003e\n\n## Installation\n\nCreate a conda environment from the given [`environment.yml`](https://github.com/deepmodeling/jax-fem/blob/main/environment.yml) file and activate it:\n\n```bash\nconda env create -f environment.yml\nconda activate jax-fem-env\n```\n\nInstall JAX\n- See jax installation [instructions](https://github.com/jax-ml/jax?tab=readme-ov-file#installation). Depending on your hardware, you may install the CPU or GPU version of JAX. Both will work, while GPU version usually gives better performance.\n\n\nThen there are two options to continue:\n\n### Option 1\n\nClone the repository:\n\n```bash\ngit clone https://github.com/deepmodeling/jax-fem.git\ncd jax-fem\n```\n\nand install the package locally:\n\n```bash\n\npip install -e .\n```\n\n**Quick tests**: You can check `demos/` for a variety of FEM cases. For example, run\n\n```bash\npython -m demos.hyperelasticity.example\n```\n\nfor hyperelasticity. \n\nAlso, \n\n```bash\npython -m tests.benchmarks\n```\n\nwill execute a set of test cases.\n\n\n### Option 2\n\nInstall the package from the [PyPI release](https://pypi.org/project/jax-fem/) directly:\n\n```bash\npip install jax-fem\n```\n\n**Quick tests**: You can create an `example.py` file and run it:\n\n```bash\npython example.py\n```\n\n```python\nimport jax\nimport jax.numpy as np\nimport os\n\nfrom jax_fem.problem import Problem\nfrom jax_fem.solver import solver\nfrom jax_fem.utils import save_sol\nfrom jax_fem.generate_mesh import get_meshio_cell_type, Mesh, rectangle_mesh\n\nclass Poisson(Problem):\n    def get_tensor_map(self):\n        return lambda x: x\n\n    def get_mass_map(self):\n        def mass_map(u, x):\n            val = -np.array([10*np.exp(-(np.power(x[0] - 0.5, 2) + np.power(x[1] - 0.5, 2)) / 0.02)])\n            return val\n        return mass_map\n\nele_type = 'QUAD4'\ncell_type = get_meshio_cell_type(ele_type)\nLx, Ly = 1., 1.\nmeshio_mesh = rectangle_mesh(Nx=32, Ny=32, domain_x=Lx, domain_y=Ly)\nmesh = Mesh(meshio_mesh.points, meshio_mesh.cells_dict[cell_type])\n\ndef left(point):\n    return np.isclose(point[0], 0., atol=1e-5)\n\ndef right(point):\n    return np.isclose(point[0], Lx, atol=1e-5)\n\ndef bottom(point):\n    return np.isclose(point[1], 0., atol=1e-5)\n\ndef top(point):\n    return np.isclose(point[1], Ly, atol=1e-5)\n\ndef dirichlet_val(point):\n    return 0.\n\nlocation_fns = [left, right, bottom, top]\nvalue_fns = [dirichlet_val]*4\nvecs = [0]*4\ndirichlet_bc_info = [location_fns, vecs, value_fns]\n\nproblem = Poisson(mesh=mesh, vec=1, dim=2, ele_type=ele_type, dirichlet_bc_info=dirichlet_bc_info)\nsol = solver(problem)\n\ndata_dir = os.path.join(os.path.dirname(__file__), 'data')\nvtk_path = os.path.join(data_dir, f'vtk/u.vtu')\nsave_sol(problem.fes[0], sol[0], vtk_path)\n```\n\nBy running the code above and use [Paraview](https://www.paraview.org/) for visualization, you should see the following solution.\n\n\u003cp align=\"middle\"\u003e\n  \u003cimg src=\"images/poisson.png\" width=\"400\" /\u003e\n\u003c/p\u003e\n\u003cp align=\"middle\"\u003e\n    \u003cem \u003eSolution to the Poisson's equation due to a source term.\u003c/em\u003e\n\u003c/p\u003e\n\n## Tutorial\n\n| Example                                                      | Highlight                                                    |\n| ------------------------------------------------------------ | ------------------------------------------------------------ |\n| [poisson](https://github.com/deepmodeling/jax-fem/tree/main/demos/poisson) | $${\\color{green}Basics:}$$  Poisson's equation in a unit square domain with Dirichlet and Neumann boundary conditions, as well as a source term. |\n| [linear_elasticity](https://github.com/deepmodeling/jax-fem/tree/main/demos/linear_elasticity) | $${\\color{green}Basics:}$$  Bending of a linear elastic beam due to Dirichlet and Neumann boundary conditions. Second order tetrahedral element (TET10) is used. |\n| [hyperelasticity](https://github.com/deepmodeling/jax-fem/tree/main/demos/hyperelasticity) | $${\\color{blue}Nonlinear \\space Constitutive \\space Law:}$$ Deformation of a hyperelastic cube due to Dirichlet boundary conditions. |\n| [plasticity](https://github.com/deepmodeling/jax-fem/tree/main/demos/plasticity) | $${\\color{blue}Nonlinear \\space Constitutive \\space Law:}$$ Perfect J2-plasticity model is implemented for small deformation theory. |\n| [phase_field_fracture](https://github.com/deepmodeling/jax-fem/tree/main/demos/phase_field_fracture) | $${\\color{orange}Multi-physics \\space Coupling:}$$ Phase field fracture model is implemented. Staggered scheme is used for two-way coupling of displacement field and damage field. Miehe's model of spectral decomposition is implemented for a 3D case. |\n| [thermal_mechanical](https://github.com/deepmodeling/jax-fem/tree/main/demos/thermal_mechanical) | $${\\color{orange}Multi-physics \\space Coupling:}$$ Thermal-mechanical modeling of metal additive manufacturing process. One-way coupling is implemented (temperature affects displacement). |\n| [thermal_mechanical_full](https://github.com/deepmodeling/jax-fem/tree/main/demos/thermal_mechanical_full) | $${\\color{orange}Multi-physics \\space Coupling:}$$ Thermal-mechanical modeling of 2D plate. Two-way coupling (temperature and displacement) is implemented with a monolithic scheme. |\n| [wave](https://github.com/deepmodeling/jax-fem/tree/main/demos/wave) | $${\\color{lightblue}Time \\space Dependent \\space Problem:}$$ The scalar wave equation is solved with backward difference scheme. |\n| [topology_optimization](https://github.com/deepmodeling/jax-fem/tree/main/demos/topology_optimization) | $${\\color{red}Inverse \\space Problem:}$$ SIMP topology optimization for a 2D beam. Note that sensitivity analysis is done by the program, rather than manual derivation. |\n| [inverse](https://github.com/deepmodeling/jax-fem/tree/main/demos/inverse) | $${\\color{red}Inverse \\space Problem:}$$ Sanity check of how automatic differentiation works. |\n| [plasticity_gradient](https://github.com/deepmodeling/jax-fem/tree/main/applications/plasticity_gradient) | $${\\color{red}Inverse \\space Problem:}$$ Automatic sensitivity analysis involving history variables such as plasticity. |\n\n## License\n\nThis project is licensed under the GNU General Public License v3 - see the [LICENSE](https://www.gnu.org/licenses/) for details.\n\n## Citations\n\nIf you found this library useful in academic or industry work, we appreciate your support if you consider 1) starring the project on Github, and 2) citing relevant papers:\n\n```bibtex\n@article{xue2023jax,\n  title={JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science},\n  author={Xue, Tianju and Liao, Shuheng and Gan, Zhengtao and Park, Chanwook and Xie, Xiaoyu and Liu, Wing Kam and Cao, Jian},\n  journal={Computer Physics Communications},\n  pages={108802},\n  year={2023},\n  publisher={Elsevier}\n}\n```\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdeepmodeling%2Fjax-fem","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdeepmodeling%2Fjax-fem","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdeepmodeling%2Fjax-fem/lists"}