Skip to content

chiahao3/ptyrad

Repository files navigation

PtyRAD: Ptychographic Reconstruction with Automatic Differentiation

PyPI - Version PyPI Downloads Anaconda-Server Badge Anaconda-Server Badge arXiv DOI

Docs | Install Guide | Quickstart | Paper | Youtube

PtyRAD performs ptychographic reconstruction using an automatic differention (AD) framework powered by PyTorch, which enables flexible and efficient implementation of gradient descent optimization. See our Microscopy and Microanalysis paper and the Zenodo record for more information and demo datasets.

Features

  • Gradient descent algorithms (Adam, SGD, LBFGS, etc.)
  • Mixed-state probe and object
  • Position correction
  • Position-dependent object tilt correction
  • Interoperability with PtychoShelves (fold_slice) and py4DSTEM
  • Streamlined preprocessing of cropping, padding, resampling, adding noises, and many more
  • Hyperparameter tuning
  • Multi-GPU reconstructions
  • JIT compilation with torch.compile

Recommended Tools

We recommend using Miniforge for Python environment management, and
Visual Studio Code for code editing and execution.

Major dependencies

  • Python 3.10 or above
  • PyTorch 2.0 or above
  • While PtyRAD can run on CPU, GPU is strongly suggested for high-speed ptychographic reconstructions.
    • PtyRAD supports both NVIDIA GPUs with CUDA and Apple Silicon (MPS)
  • PtyRAD was tested on Windows, MacOS, and Linux

Installation

We recommend installing PtyRAD using pip inside a fresh conda environment.

First, create and activate a new conda environment (ptyrad) with Python > 3.10:

conda create -n ptyrad python=3.12
conda activate ptyrad

💡 Note: After activating the environment, your terminal prompt should show (ptyrad) at the beginning, indicating that the environment is active.

Then install PtyRAD using:

pip install ptyrad

If you're using Windows with NVIDIA CUDA GPU, you will also need to install the GPU version of PyTorch with:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 --force-reinstall

PtyRAD can also be installed via conda. For detailed instructions on installing PtyRAD on different machines or pinning specific CUDA versions, see the installation guide.

How do I check if my installed PtyRAD has the GPU support?

CUDA version, GPU support, and PyTorch build across platforms can be extremely confusing, so PtyRAD provides handy CLI tools to help check these information for you!

Once you activated (ptyrad) environment and installed PtyRAD via pip install ptyrad, you'll have access to the following command:

# You can run this command anywhere from your terminal, as long as `ptyrad` is installed in the environment
ptyrad check-gpu

This command will print out relevant information of your CUDA information if available.

How do I update my existing PtyRAD installation to a newer release?

Assuming you've activated the (ptyrad) environment and you've installed PtyRAD via pip, you can simply update your PtyRAD installation with:

pip install -U ptyrad

Get Started with the Demo

💡 Note: Download the demo datasets and scripts from our Box folder, which contains demo/ and other tutorial slides and videos.

Before running the demo, please check the following:

  1. Demo datasets are downloaded and placed to the correct location under demo/data/
  2. (ptyrad) environment is created and activated (in VS Code it's the "Select Kernel")
  3. PtyRAD is installed in the (ptyrad) environment

Now you're ready to run a quick demo using one of two interfaces:

  • Interactive Jupyter interface (Recommended)

    Use demo/scripts/run_ptyrad_quick_example.ipynb to quickly reconstruct the demo dataset in a Jupyter notebook

  • Command-line interface (like your Miniforge Prompt terminal)

    # Here we assume working directory is set at `demo/`
    ptyrad run --params_path "params/tBL_WSe2_reconstruct.yml"

Documentation

PtyRAD documentation is available at https://ptyrad.readthedocs.io/en/latest/index.html.

Author

Chia-Hao Lee (cl2696@cornell.edu)

Developed at the Muller Group, Cornell University.

Acknowledgments

Besides great support from the entire Muller group, this package gets inspiration from lots of community efforts, and specifically from the following packages. Some of the functions in PtyRAD are directly translated or modified from these packages as noted in their docstrings/comments to give explicit acknowledgment.

Other resources

  • ptycho-packages lists many available ptychography packages
  • Cornell Box folder compiled by myself that keeps demo data, tutorial recordings, and slides for PtyRAD
  • Argonne Box folder compiled by Dr. Yi Jiang that holds tutorial slides of fold_slice
  • Blog post written by myself that details the algorithms and code structure of PtychoShelves / fold_slice

About

PtyRAD: Ptychographic Reconstruction with Automatic Differentiation

Topics

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Languages