deepmodeling / deepmodeling/jax-fem

JAX-FEM in the Mosaic benchmark suite

Open
#92 0 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
754
Forks
133
PR merge metrics
No merged PRs in 30d

Description

We just released **Mosaic** ([github.com/pasteurlabs/mosaic](https://github.com/pasteurlabs/mosaic)), a benchmark for *differentiable* physics solvers. Each solver runs in a container behind a uniform `apply` / `vjp` interface, so one harness compares gradient accuracy and cost across JAX, PyTorch, Julia, and C++-adjoint codes on shared tasks.

**JAX-FEM is one of the benchmarked solvers** on the structural-mesh and thermal-mesh domains ([results](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results.html) if you're interested).

We did our best to do it justice, but we don't know your code as well as you do and have surely made suboptimal calls. For example, we solve the (linear) elasticity and heat problems with the default direct UMFPACK solver and take the implicit-function-theorem adjoint JAX-FEM exposes via `ad_wrapper`. If a different linear solver, or a different way of invoking the adjoint, is what you'd recommend, we'd love to know.

The setup is self-contained under [`mosaic/tesseracts/structural-mesh/jax-fem/`](https://github.com/pasteurlabs/mosaic/tree/main/mosaic/tesseracts/structural-mesh/jax-fem/) and [`mosaic/tesseracts/thermal-mesh/jax-fem/`](https://github.com/pasteurlabs/mosaic/tree/main/mosaic/tesseracts/thermal-mesh/jax-fem/). If something looks odd, corrections or a PR are very welcome, here on GitHub or on the [Tesseract Forum](https://si-tesseract.discourse.group/).

Mostly we wanted you to know your work is represented, and to get it right.

Cheers,
Dion on behalf of the Mosaic authors

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.