⚠️ Repository Moved

This repository has been replaced by mosaix-pde.

Please use the new repository for all code, documentation, and issues.

The old code here is preserved for reference only.

pat-pde-opt

pat-pde-opt is a package for optimizing pattern forming PDEs that appear in different areas of physics, written in JAX. It has code for PDE optimization and control with gradient-based methods and reinforcement learning. We use diffrax for time stepping and implement system-specific solvers, such as semi-implicit Fourier methods and Strang splitting.

You can find the full documentation on read the docs.

Installation

To install the package, we recommend cloning the github repo and then installing locally:

git clone https://github.com/acoh64/pde-opt.git
cd pde-opt
conda create -y -n pde-opt-env python=3.12
conda activate pde-opt-env
pip install -e .

By default, it will install the CPU version of JAX. To use with GPU, run:

pip install -U "jax[cuda12]"

Usage

Here is an example of solving the Cahn-Hilliard equation in 2D with periodic boundary conditions using a semi-implicit Fourier method:

import jax
import jax.numpy as jnp

from pde_opt import PDEModel
from pde_opt import CahnHilliard2DPeriodic
from pde_opt import SemiImplicitFourierSpectral
from pde_opt import Domain
from pde_opt import PeriodicCNN

Nx = Ny = 128
Lx = Ly = 0.01 * Nx

domain = Domain((Nx, Ny), ((-Lx / 2, Lx / 2), (-Ly / 2, Ly / 2)), "dimensionless")

opt_model = PDEModel(equation_type=CahnHilliard2DPeriodic, domain=domain, solver_type=SemiImplicitFourierSpectral)

params = {"kappa": 0.002, "mu": lambda c: jnp.log(c / (1.0 - c)) + 3.0 * (1.0 - 2.0 * c), "D": lambda c: c * (1. - c)}

solver_params = {"A": 0.5}

key = jax.random.PRNGKey(0)
y0 = jnp.clip(0.01 * jax.random.normal(key, (Nx, Ny)) + 0.5, 0.0, 1.0)
ts = jnp.linspace(0.0, 0.02, 100)

sol = opt_model.solve(params, y0, ts, solver_params, dt0=0.000001, max_steps=1000000)

Next, here is an example of using the previous solution as a dataset to fit a neural network for the chemical potential term:

data = {}
data['ys'] = sol
data['ts'] = ts

model = PeriodicCNN(
    in_channels=1,
    hidden_channels=(32, 64, 64),
    out_channels=1,
    kernel_size=3,
    key=jax.random.PRNGKey(0),
)

init_params = {"mu": model}
static_params = {"kappa": 0.002, "D": lambda c: c * (1. - c)}
solver_parameters = {"A": 0.5}
weights = {"mu": None}
lambda_reg = 0.0

inds = [[30,40,50], [50,60,70], [70,80,90]]

res = opt_model.train(data, inds, init_params, static_params, solver_parameters, weights, lambda_reg, method="mse", max_steps=100)

Current Model Implementations

This package is designed to support pattern-forming PDEs across a wide-range of physical systems. We have currently implemented variants of the following equations:

• Cahn-Hilliard equation
  • 2D with periodic boundary conditions
  • 3D with periodic boundary conditions
  • 2D with smoothed boundary method
• Allen-Cahn equation
  • 2D with periodic boundary conditions
  • 2D with constant current conditions + Butler-Volmer kinetics (for battery applications)
  • 2D with smoothed boundary
  • 2D with smoothed boundar and constant current conditions + Butler-Volmer kinetics (for battery applications)
• Gross-Pitaevskii
  • Reduced 2D with periodic boundary conditions
  • Rotating reduced with 2D periodic boundary conditions

TODO

• Arbitrary boundary conditions
• Implicit time stepping
• Multi-GPU support
• Extend to non-Cartesian domains
• WandB logging and checkpointing

License

This code has been published under the MIT licence.