PtyLabX

Warning

This project is under active development. Some existing implementations could be unstable until everything is thoroughly tested. API might also change and would be around making some interface immutable to allow gradient flow/differentiable ptychography work.

Installation | Development and Contribution | Getting Started | PtyLabX Documentation

PtyLabX is an experimental JAX-based ptychographic reconstruction toolbox, forked from PtyLab.py (see the original publication). The JAX backend brings GPU acceleration, JIT, and XLA, aiming to accelerate existing algorithms and enable new models with automatic differentiation and gradient-based optimization.

Installation

From source (pip)

pip install git+https://github.com/ShantanuKodgirwar/PtyLabX.git

With GPU support

CUDA 12

pip install "ptylabx[cuda12]@git+https://github.com/ShantanuKodgirwar/PtyLabX.git"

CUDA 13 (Recommended for newer GPUs)

This is recommended for newer GPUs with SM version 7.5 or newer.

pip install "ptylabx[cuda13]@git+https://github.com/ShantanuKodgirwar/PtyLabX.git"

Development and Contribution

This project uses the (super-fast and very easy to install) package manager uv. Follow the below steps:

git clone https://github.com/ShantanuKodgirwar/PtyLabX.git
cd PtyLabX
uv sync 

The dependencies are installed in a virtual environment which can be activated within your IDE interpreter or in the terminal source .venv/bin/activate.

To install with GPU support:

uv sync --extra cuda13 # alternatively `cuda12` for older GPUs

To test if GPU is detected:

uv run python -c "import jax; print(jax.default_backend())"

Note

Contributions are welcome (new engines, bug fixes, example scripts, documentation). For any new implementations,

1. Add an appropriate test under tests/ using pytest.
2. Run tests locally before opening a PR:

uv run pytest tests/ -v -s

Tests run automatically with CI on every PR and if they fail, please review your changes. Additionally, code formatting based on ruff is done automatically when a PR is opened.

Documentation

The documentation is a work-in-progress, but can be rendered as a webpage:

uv run mkdocs serve

Acknowledgements

Some of the new implementation ideas mainly follow the very recent JAX-based electron ptychography package phaser and the well established pytorch-based ptyrad library. Additionally, differentiable wave optics library chromatix also provided us with new ideas.

Citation

If you use this package in your work, cite us as below.

@article{Loetgering:23,
        author = {Lars Loetgering and Mengqi Du and Dirk Boonzajer Flaes and Tomas Aidukas and Felix Wechsler and Daniel S. Penagos Molina and Max Rose and Antonios Pelekanidis and Wilhelm Eschen and J\"{u}rgen Hess and Thomas Wilhein and Rainer Heintzmann and Jan Rothhardt and Stefan Witte},
        journal = {Opt. Express},
        number = {9},
        pages = {13763--13797},
        publisher = {Optica Publishing Group},
        title = {PtyLab.m/py/jl: a cross-platform, open-source inverse modeling toolbox for conventional and Fourier ptychography},
        volume = {31},
        month = {Apr},
        year = {2023},
        doi = {10.1364/OE.485370},
}