AdjointDiffusion

AdjointDiffusion is a new method for structural optimization using diffusion models. It is a physics-guided and fabrication-aware structural optimization using diffusion models augmented with adjoint gradient. By combining powerful generative models with adjoint sensitivity analysis, this approach can more efficiently discover complex, high-performance designs than the traditional methods.

The codes are provided following the paper named Physics-guided and fabrication-aware structural optimization using diffusion models

---

Table of Contents

1. TL;DR
2. Intuitive Explanation of Diffusion Models
3. Installation
4. Quick Start
5. Customize Your Simulation
6. Experiment Logging with Weights & Biases
7. Results
8. Code Organization
9. Citation

TL;DR

✨ Integrating adjoint sensitivity analysis with diffusion models can generate high-performance and interesting structures!

Key features:

• Adjoint Sensitivity Integration: Seamlessly incorporates adjoint gradients into the diffusion process.
• Fabrication Constraints: Accounts for manufacturability, ensuring real-world feasibility.
• Extensibility: Users can use their own datasets or simulations.
• Experiment Tracking & Visualization: Integrates with Weights & Biases.

---

Intuitive Explanation of Diffusion Models

Imagine an ink drop falling into water — it slowly spreads and dissolves. Diffusion models mimic this process in reverse: they start from noise and slowly form meaningful structures. By guiding this "reverse diffusion" with gradients from an adjoint method, we ensure the final designs are optimized and fabrication-ready.

---

Installation

This setup ensures compatibility between Meep and PyTorch. If you find any alternatives, feel free to contribute improvements via pull requests!

1. Clone the repository

git clone https://github.com/dongjin-seo2020/AdjointDiffusion.git
cd AdjointDiffusion

2. Set up a Python environment (recommended)

Using conda

To create and activate the recommended environment with necessary dependencies:

conda create -n adjoint_diffusion -c conda-forge pymeep pymeep-extras python=3.9 
conda activate adjoint_diffusion

3. Installation of other libraries

Install torch following the command (recommended):

pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117

Note: It has been observed that newer NVIDIA GPUs (e.g., RTX 5090) may not be compatible with this specific version of PyTorch. If you encounter issues, please refer to the official PyTorch installation guide to find a version compatible with your hardware: https://pytorch.org/get-started/locally/

Install the required packages listed in requirements.txt:

pip install -r requirements.txt

4. Error Handling

If you encounter permission-related issues when trying to run the training script, make sure it is executable by running:

chmod +x 01-train.sh

The, you can execute it with:

./01-train.sh

If you encounter errors while installing mpi4py, try the following steps:

apt --fix-broken install
apt install mpich
pip install mpi4py

Make sure you have root access when using apt.

---

Quick Start

1. Generate a dataset:

python dataset_generation.py

• The data will be saved at datasets/<n>/sigma<k>/struct/, where n is the structure dimension (e.g., n=64 generates 64×64 binary structures) and k is the variance of the Gaussian filter (a larger k increases the minimum feature size).
• Note: To reproduce the condition in the paper, run the code for k=2, k=5, and k=8 (three times). Or you can download our pretrained network from https://zenodo.org/records/15399997
• Note: You can also use your own dataset here! Provide your fabrication-satisfying image dataset and train the diffusion model with it!

2. Update the training and sampling scripts to specify the appropriate output directories.

• For example, for train.sh, you should specify the variables as:

DATA_DIR=/path/to/datasets
LOG_DIR=path/to/experiments
GPU_ID=0

• Or, you can set environment variables (Linux/macOS):

export DATA_DIR=/path/to/datasets
export LOG_DIR=/path/to/experiments
export GPU_ID=0

• For detailed usage examples, including training and sampling with actual settings, see:
  • training_examples.ipynb
  • sampling_example.ipynb

3. Train a diffusion model:

./01-train.sh

• Alternatively: run 02-train.ipynb
• Note: Set --class_cond to False if your dataset contains only a single structural condition (i.e., no class conditioning needed). If you have multiple structural conditions (e.g., different fabrication constraints), set it to True to enable class-conditional training.
• Note: The training process will continue indefinitely unless manually stopped. In our setup, training for around 25,000 steps produced satisfactory results, though fewer steps may also be sufficient. If you are using a customized dataset, the optimal number of steps may vary.

4. Sample and optimize structures:

./01-sample.sh

• Alternatively: run 02-sample.ipynb
• Note: Set --class_cond to False if your dataset contains only a single structural condition (i.e., no class conditioning needed). If you have multiple structural conditions (e.g., different fabrication constraints), set it to True to enable class-conditional training.
• Note: We recommend to use the model where the name is ema_0.9999_*.pt. The number in * means the training step number. ema is for Exponential Moving Average.

5. View outputs

• Every output (performance, structure) is logged in wandb.
• Logs and generated structures are saved in ./logs/<run_name>

6. Baseline Algorithms

We provide baseline algorithms in the ./baseline_algorithms directory. These include nlopt methods like MMA for comparison.

---

Customize Your Simulation

If you'd like to integrate a custom physical simulation into the reverse diffusion process, follow these steps:

1. Implement Your simulation class in guided_diffusion/simulation.py.

   Create a class that defines how to compute the figure of merit (FoM) and its corresponding adjoint gradient. For example:

   class YourSimClass:
       def __init__(self, ...):
           ...
       def compute_fom(self, structure):
           ...
       def compute_adjoint(self, structure):
           ...

2. Update the import in guided_diffusion/gaussian_diffusion.py:

   Replace the existing simulation import with your custom class:

   from guided_diffusion.simulation import YourSimClass

3. Plug your simulation into the sampling loop in guided_diffusion/gaussian_diffusion.py In guided_diffusion/gaussian_diffusion.py, locate where simulation_() is called (typically inside the p_sample() function), and replace it with your custom simulation logic. Make sure your class is initialized properly and passed via my_kwargs. For example, for my_kwargs:

my_kwargs = {
    "sim_guided": True,
    "simulation_": YourSimClass(...),
    "eta": 0.05,
    "inter_rate": 25,
    "stoptime": 0.1,
    "guidance_type": "dps",  # or "dds"
    "exp_name": "experiment1",
    ...
}

Now, your custom simulation will be used during the reverse diffusion process.

---

Experiment Logging with Weights & Biases

We use wandb for logging and visualization.

1. Sign up at wandb.ai
2. Log in:

wandb login

3. Run any training/sampling script and it will automatically log data to wandb.

---

Results

We visualize the performance of AdjointDiffusion across different tasks and configurations.

Optimization Convergence and Comparisons - Problem Setup 1 (Waveguide)

Comparison of Generated Structures - Problem Setup 1 (Waveguide)

Optimization Convergence and Comparisons 2 - Problem Setup 2 (Multi-wavelength Splitter)

Comparison of Generated Structures - Problem Setup 2 (Multi-wavelength Splitter)

Result: Waveguide

Result: Multi-wavelength Splitter

---

Code Organization

AdjointDiffusion/
├── dataset_generation.py       # Dataset generation script
├── image_train.py              # Main training script
├── image_sample.py             # Main sampling script
├── requirements.txt            # Python dependencies
├── guided_diffusion/           # Backend of diffusion models
└── baseline_algorithms/        # Baseline algorithms (nlopt, Gradient Ascent)

---

Citation

If you use this code, please cite the following paper:

@article{Seo2026DiffusionPhotonics,
  title   = {Physics-Guided and Fabrication-Aware Inverse Design of Photonic Devices Using Diffusion Models},
  author  = {Seo, Dongjin and Um, Soobin and Lee, Sangbin and Ye, Jong Chul and Chung, Haejun},
  journal = {ACS Photonics},
  year    = {2026},
  doi     = {10.1021/acsphotonics.5c00993},
  url     = {https://doi.org/10.1021/acsphotonics.5c00993}
}

---

Acknowledgements

Parts of this repository are adapted from OpenAI's guided-diffusion, which is licensed under the MIT License.

We thank the OpenAI team for their contribution. Significant modifications have been made to enable adjoint sensitivity integration and fabrication-aware optimization.

---

Happy Diffusing & Optimizing!