This repository contains the codebase for the methods presented in the paper "Physics Constrained Unsupervised Deep Learning for Rapid, High Resolution Scanning Coherent Diffraction Reconstruction".
PtychoPINN is an unsupervised physics-informed neural network reconstruction method for scanning CDI designed to improve upon the speed of conventional reconstruction methods without sacrificing image quality. Compared to prior NN approaches, the main source of improvements in image quality are its combination of the diffraction forward map with real-space overlap constraints.
- Unsupervised / self-supervised learning: There is no need for extensive labeled training data, making the model more practical to train.
- Resolution: PtychoPINN outperforms existing deep learning models for ptychographic reconstruction in terms of image quality, with a 10 dB PSNR increase and a 3- to 6-fold gain in linear resolution. Generalizability and robustness are also improved.
- Scalability and Speed: PtychoPINN is two or three orders of magnitude as fast as iterative scanning CDI reconstruction.
conda create -n ptycho python=3.10
conda activate ptycho
pip install .
python scripts/training/train.py --train_data_file <train_path.npz> --test_data_file <test_path.npz>
python scripts/inference/inference.py --model_path <model artifact directory path> --test_data <test_path.npz>
See examples and READMEs under scripts/.
For an example of interactive (Jupyter) usage, see notebooks/nongrid_simulations.ipynb. If you don't have inputs in the right .npz format you can simulate data.
non_grid_CDI_example.ipynb shows interactive usage using a dataset that is provided with the repo.