google / google/jax-cfd

JAX-CFD in the Mosaic benchmark suite

Open
#259 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
964
Forks
142
Avg merge
3h 9m
Merged PRs (30d)
1

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-CFD is one of the benchmarked solvers** on the ns-grid domain ([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 use the semi-implicit finite-difference (MAC-grid) stepper with the default advection scheme and spectral pressure projection, and differentiate through it with JAX autodiff. If you'd pick a different advection method, CFL or step size, or projection tolerance for accurate gradients, we'd love to know.

The setup is self-contained under [`mosaic/tesseracts/navier-stokes-grid/jax-cfd/`](https://github.com/pasteurlabs/mosaic/tree/main/mosaic/tesseracts/navier-stokes-grid/jax-cfd/). 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

Open the contributing guide

Research direction

Review the external setup under mosaic/tesseracts/navier-stokes-grid/jax-cfd/ and compare its semi-implicit MAC-grid stepper, advection scheme, CFL or step size, and spectral pressure projection with JAX-CFD's documented usage. Done means identifying a concrete correction or improvement and proposing it as a focused change or PR; the issue names no repository file or test to modify.

Written by the indexing model from the issue text.

Assessment

Tech stack
jupyter-notebook, python
Domain
performance, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.