This repository contains the code and experiments for the paper:
This repository implements stochastic Sequential Quadratic Programming (SQP) methods for solving Physics-Informed Neural Network (PINN) problems with nonlinear equality constraints. The code includes implementations of several optimization algorithms including SQP-Adam and SQP-Heavyball, tested on various benchmark problems (Spring, Darcy, Burgers, Chemistry).
The installation of torch_harmonics and neuraloperator may report errors on some systems. We recommend using Google Colab for the most reliable experience. Please follow run_on_colab.ipynb to install and setup experiments.
To run a test, use (Please add ! at the beginning of command if using Google Colab):
python3 solve.py spring_test
A folder result_test will be created with the structure shown below:
result_test/
├── log
│ └── Spring
│ └── test_0.txt
├── mdl
│ └── Spring
│ ├── nn_test_0_1000
│ ├── ...
│ ├── nn_test_0_900
│ ├── optim_test_0_1000.pt
│ ├── ...
│ └── optim_test_0_900.pt
└── plot
└── Spring
├── animation_test_0.gif
├── ...
└── plot_test_0_00001001.png
The configuration for this test run is in the file ./conf/spring_test.yaml. Other config files are also in the ./conf/ directory. The results folder will be created within the SSQPPINN folder.
(Please add ! at the beginning of command if using Google Colab)
python3 run.pyYou may modify the settings in the top lines of run.py to run experiments for a specific problem, algorithm, or learning rate. The configurations for all test problems are set with the settings used in the paper experiments.
Available problems: spring, darcy, burgers, chemistry
Available methods: Adam (unconstrained), Adam (constrained), SQP-Heavyball, SQP-Adam
Training time estimation:
The time estimation is based on a Google Colab L4 GPU single session.
Full batch: Below is the running time estimation of setting 4 (SQP Adam, Full Batch). Setting 1 (Adam unc) runs faster while Settings 2 and 3 spend similar time.
| Problem | Number of Epochs | Time (hour) |
|---|---|---|
| spring | 30000 | 0.17 |
| burgers | 10000 | 0.2 |
| chemistry | 100000 | 1.50 |
| darcy | 4000 | 1.91 |
Mini batch: Below is the running time estimation of setting 8 (SQP Adam, Mini Batch) for ONE random seed. Setting 5 (Adam unc) runs faster while Settings 6 and 7 spend similar time. The experiments in the paper are the result of 5 random seeds.
| Problem | Number of Epochs | Time (hour) |
|---|---|---|
| spring | 30000 | 0.34 |
| burgers | 10000 | 0.77 |
| chemistry | 20000 | 1.52 |
| darcy | 2000 | 4.63 |
Please edit run.py to customize the set of settings, number of epochs, learning rates, and random seeds.
SSQPPINN/
├── conf/ # YAML configuration files for each problem
├── problems/ # Problem definitions (Spring, Darcy, Burgers, Chemistry)
├── nn_architecture.py # Neural network architectures (FCN, FNO)
├── stochasticsqp.py # SQP optimizer implementation
├── solve.py # Main training loop
├── run.py # Script to run all experiments
├── utils.py # Utility functions
└── run_on_colab.ipynb # Google Colab notebook
If you use this code in your research, please cite:
@article{wang2026projected,
title={Projected Stochastic Momentum Methods for Nonlinear Equality-Constrained Optimization for Machine Learning},
author={Qi Wang, Christian Piermarini, Yunlang Zhu, Frank E. Curtis},
journal={arXiv preprint arXiv:2601.11795},
year={2026},
doi={10.48550/arXiv.2601.11795}
}This project is licensed under the MIT License - see the LICENSE file for details.
If you have any questions or issues, please contact the authors: Qi Wang (qiwangqi@umich.edu), and Yulang Zhu (yuza23@lehigh.edu)