/NeurOLight

A Neural Operator-based Integrated Photonic Device Simulation Framework, NeurOLight NeurIPS 2022

Primary LanguagePythonMIT LicenseMIT

NeurOLight

By Jiaqi Gu, Zhengqi Gao, Chenghao Feng, Hanqing Zhu, Ray T. Chen, Duane S. Boning and David Z. Pan.

This repo is the official implementation of "NeurOLight: A Physics-Agnostic Neural Operator Enabling Parametric Photonic Device Simulation".

MIT License


Introduction

NeurOLight is a physics-agnostic neural operator-based framework to enable ultra-fast parametric integrated photonic device simulation.

NeurOLight learns a family of frequency-domain Maxwell PDEs to predict the light field insides integrated photonic devices. We discretize different devices into a unified domain, represent parametric PDEs with a compact wave prior, and encode the incident light via masked source modeling. NeurOLight is designed with parameter-efficient cross-shaped NeurOLight blocks and adopts superposition-based augmentation for data-efficient learning. flow

NeurOLight generalizes to a large space of unseen simulation settings, demonstrates 2-orders-of-magnitude faster simulation speed than numerical solvers, and outperforms prior neural network models by ~54% lower prediction error with ~44% fewer parameters.

Dependencies

  • Python >= 3.6
  • pyutils >= 0.0.1. See pyutils for installation.
  • pytorch-onn >= 0.0.5. See pytorch-onn for installation.
  • Python libraries listed in requirements.txt
  • NVIDIA GPUs and CUDA >= 10.2

Structures

  • configs/: configuration files
  • core/
    • models/
      • layers/
        • neurolight_conv2d: NeurOLight block definition
        • fno_conv2d: FNO block definition
        • factorfno_conv2d: Factorized-FNO block definition
      • neurolight_cnn.py: NeurOLight model definition
      • fno_cnn.py: FNO model definition
      • factorfno_cnn.py: Factorized-FNO (F-FNO) model definition
      • unet.py: UNet-2d model definition
      • pde_base.py: base model definition
      • utils.py: utility functions
      • constant.py: constant definition
    • builder.py: build training utilities
    • utils.py: customized loss function definition
  • scripts/: contains experiment scripts
  • data/: MMI simulation dataset
  • train.py: training logic
  • test.py: inference logic
  • tune.py: transfer and finetune logic

Usage

  • 3x3 tunable MMI dataset simulation and raw dataset generation (Optional). DATASET need to be replaced by one option from (mmi3x3, mmi4x4, mmi5x5, mmi3x3_etched)
    data/mmi/> python3 make_dataset.py DATASET

  • Train MODEL on 3x3 tunable MMI datasets, replace MODEL with one option from (neurolight, fno, factorfno, unet),
    > python3 scripts/mmi/MODEL/train_random_tunable.py

  • Or, train MODEL on 3x3 etched MMI datasets, replace MODEL with one option from (neurolight, fno, factorfno, unet),
    > python3 scripts/mmi/MODEL/train_random_etched.py

  • Test the trained model on test simulation instances with the single-shot multi-source mode. (Need to modify the path to trained checkpoint in the script). python3 scripts/mmi/MODEL/test_main_results.py

  • Checkpoints for 12-layer NeurOLight on tunable MMIs and 16-layer NeurOLight on etched MMIs are available at link

  • Simulation dataset for Tunable MMI and etched MMI are available at link. Unzip it and place all *.pt under data/mmi/raw/*.pt.

Citing NeurOLight

@inproceedings{gu2022NeurOLight,
  title={NeurOLight: A Physics-Agnostic Neural Operator Enabling Parametric Photonic Device Simulation},
  author={Jiaqi Gu and Zhengqi Gao and Chenghao Feng and Hanqing Zhu and Ray T. Chen and Duane S. Boning and David Z. Pan},
  booktitle={Conference on Neural Information Processing Systems (NeurIPS)},
  year={2022}
}