Automatic Optimization of a PDE Solver

I would like to test whether automatic AI scientist can improve a PDE optimizer.

The code for the project is here: https://github.com/meliao/jaxhps. It's a python package written in JAX. The associated paper is here: https://arxiv.org/abs/2503.17535.

Here are initial thoughts about improvement requests we could phrase to an LLM:

1. Improving the parallelization for solves on adaptive grids. This is the big one. Currently, all of the code on uniform grids is able to parallelize to some extent during the most computationally-intensive part of the algorithm called the "merge step", but code on adaptive grids is not parallelized. This is because I was unable to figure out how to apply the JAX vmap transformation on a collection of problems of different sizes. This part of the algorithm is embarrassingly parallel, so it's really just an implementation issue blocking us from doing this.

2. Kernel fusion. Right now, the code is separated into "local solve" and "merge" stages. This separation is convenient for writing the code and paper, but it may be a bit faster to fuse these operations together a la FlashAttention.

3. Improving interpolation from spectral grids to uniform grids. Currently, the code computes a solution to the PDE on a "spectral grid" and for downstream uses, like visualization for instance, we usually want the solution on regularly-spaced grid points. Currently this interpolation step from the spectral grid to the uniform grid is implemented with dense matrix vector products, but it could be done with fast Fourier transforms; we could write down that algorithm succinctly. This would save a lot of memory during this interpolation step, and I believe it would be faster too.

Also try tools like OpenEvolve: https://github.com/algorithmicsuperintelligence/openevolve and ShinkaEvolve: https://github.com/SakanaAI/ShinkaEvolve.