Surrogate-Based Differentiable Pipeline for Shape Optimization

This repository contains the reference implementation for the paper Surrogate-Based Differentiable Pipeline for Shape Optimization.

Overview

We demonstrate an end-to-end differentiable pipeline for shape optimization, serving as a template for more complex scientific machine learning workflows. Our toy example uses a simple obstacle geometry optimization problem to illustrate how to combine high-fidelity CFD simulations with neural network surrogate models for gradient-based optimization.

The the individual components use OpenFOAM, PyTorch, JAX, PySDF, TriMesh and more. Tesseracts enable composition and autodiff of these heterogeneous components. The optimization problem is intentionally simple but serves as a template for more complex tasks.

Quick Links:

• Getting Started - Installation and running the pipeline
• Technical Details - Problem formulation and methodology
• Citation - How to cite this work

Pipeline Architecture

The toy example implements a three-stage pipeline that demonstrates how to build composable scientific ML workflows. Each stage can run independently or be chained together:

1. Dataset Generation

Samples the design space and runs CFD simulations for each configuration. Creates parametric geometries, generates meshes, and simulates them in OpenFOAM to get steady-state flow solutions.

• Components: dpp-geometry, dpp-mesh, dpp-openfoam, dpp-interpolate-to-grid
• Compute: HPC resources for parallel CFD
• Output: 1000 paired geometry-flow samples (stored on AzureML)

2. Surrogate Model Training

Trains a UNet+Attention graph neural network (Physics-NEMO architecture) to map from geometry to flow fields. The model learns from the CFD dataset to enable fast inference during optimization.

• Components: dpp-unet-trainer
• Compute: Multi-GPU with CUDA (distributed training supported)
• Output: Trained surrogate model weights and configuration

3. Shape Optimization

Runs gradient-based optimization using the trained surrogate. Optimizes geometry parameters to find improved designs.

• Components: dpp-geometry, dpp-unet-inference
• Compute: CPU or GPU
• Output: Optimized geometry parameters and flow predictions

Getting Started

Prerequisites

Install the shared tools for running and composing Tesseracts:

# create a virtual environment
python3 -m venv .venv --prompt shapeopt
source .venv/bin/activate

# install glue code
pip install -e shapeopt_tools/

# This will take a while. You can also build only the necessary tesseracts for
# each pipeline in subsequent steps
bash build-tesseracts.sh

Dataset Generation

Generate parameterized geometries and run OpenFOAM simulations to produce training data:

# if you have not already:
bash build-tesseracts.sh dpp-geometry dpp-mesh dpp-openfoam dpp-interpolate-to-grid

# run data generation
python datagen/run.py datagen/sample-dataset --num-samples 4
python datagen/clean.py datagen/sample-dataset

Model Training

Currently, training runs outside of a Tesseract container to enable distributed training - which tesseracts don't support yet. To install dependencies and train the surrogate model:

pip install -r tesseracts/dpp-unet-trainer/tesseract_requirements.txt
torchrun --nproc_per_node=1 \
  train/run.py train/config.json --data-dir datagen/sample-dataset

Design Optimization

Run the optimization pipeline using the trained surrogate:

# if you have not already:
bash build-tesseracts.sh dpp-geometry dpp-unet-inference

# run design optimization
python optimization/optimize_design.py \
  --config optimization/optimizer_config.json

because building the inference tesseract can take several hours, alternatively you can run the optimization without building it by running:

# if you have not already:
bash build-tesseracts.sh dpp-geometry
pip install -r tesseracts/dpp-unet-inference/tesseract_requirements.txt

# run design optimization
python optimization/optimize_design.py \
  --config optimization/optimizer_config.json \
  --unet-api tesseracts/dpp-unet-inference/tesseract_api.py

For a rendered .mp4 animation of the optimization process, install ffmpeg first.

Citation

If you use this code in your research, please cite:

@article{sygaldry2025,
  title={Surrogate-Based Differentiable Pipeline for Shape Optimization},
  author={[Andrin Rehmann, Nolan Black, Josiah Bjorgaard, Alessandro Angioi, Andrei Paleyes, Niklas Heim, Dion Häfner, Alexander Lavin]},
  journal={[EurIPS - DiffSys Workshop]},
  year={2025},
  url={https://openreview.net/forum?id=TZ3VkOQINw}
}

Additional Resources

Component-specific documentation:

• UNet Trainer - Surrogate model training configuration and usage
• UNet Inference - Surrogate model deployment and inference
• Geometry Generation (dpp-geometry) - Parametric SDF generation
• Meshing (dpp-mesh) - Marching cubes and mesh processing
• OpenFOAM Simulation (dpp-openfoam) - CFD simulation setup and execution
• Grid Interpolation (dpp-interpolate-to-grid) - Solution transfer between meshes