zxycxj / helmholtz-fno

Neural operator surrogates for electromagnetic inverse design

Surrogate solvers for electromagnetic field inference and inverse design

This repository contains code for the paper "Neural Operator-Based Surrogate Solver for Free-Form Electromagnetic Inverse Design" (ACS Photonics, arXiv).

Installation

To run the code, set up a conda environment from the provided environment file and activate it:

conda env create -f environment.yml -p your_env_prefix
conda activate your_env_prefix

For data generation, use the environment file env_datagen.yml provided here. Note that the environment files specify strict version requirements - this is simply because everything has been tested using those versions. It is likely that everything works fine for different versions (of PyTorch, for instance), but we can make no guarantees.

Usage

We provide the Slurm batch scripts that were used to perform all the computations in the paper here. While these probably won't run on your particular setup without modification, they should provide enough insight on how to use the data generation, training and inverse design scripts.

FNO & UNet training

To train a surrogate solver, use the file train_surrogates.py, see ./train_surrogates.py --help for a list of parameters. In the simplest case, using just the defaults, you can train an FNO-2D model by simply providing a path to the dataset:

python3 train_surrogates.py --data-dir /path/to/data

Note that the dataloaders used here expect the following folder structure:

data
├── train
│   ├── train_data_0.h5
│   ├── train_data_1.h5
│   ├── etc.
└── test
    ├── test_data_0.h5
    ├── test_data_1.h5
    └── etc.

VAE training

The VAE models can be trained using train_vae.py, see ./train_vae.py --help for a list of parameters. You can simply run the file without providing any arguments, which will train a 2D VAE.

Inverse design

For inverse design, use the file inverse_design_3d.py. We do not currently provide an implementation for 2D inverse design, but it should be straightforward to adapt the current 3D implementation. Please note that this file needs pre-trained FNO and VAE models in TorchScript format.

Data availability

The official research data repository is provided by RADAR4KIT at https://doi.org/10.35097/911.

If you have trouble obtaining the data from the official repository, we additionally provide it via OneDrive here.

If you are downloading the dataset from China, you may download from this link as an alternative option.

Reproducibility

To reproduce results from the paper, please use the commit 73e7e19 (tagged publication).

Citing

If you use this code or associated data for your research, please cite:

@article{augenstein2023neural,
  title = {Neural Operator-Based Surrogate Solver for Free-Form Electromagnetic Inverse Design},
  author = {Augenstein, Yannick and Rep\"{a}n, Taavi and Rockstuhl, Carsten},
  year = 2023,
  journal = {ACS Photonics},
  volume = 10,
  number = 5,
  pages = {1547--1557},
  doi = {10.1021/acsphotonics.3c00156}
}