Package Organization

Overview

Ptyrodactyl is a typed and tested JAX-based library for electron microscopy and ptychography, utilizing JAX’s capabilities for multi-device computation and autodifferentiation. The package is organized into four main modules: simul for electron microscopy simulations, invert for reconstruction algorithms, tools for common utilities including optimizers and loss functions, and workflows for high-level meta-functions that orchestrate complete simulation pipelines.

Module Structure

ptyrodactyl.simul

JAX-based electron microscopy simulation toolkit for ptychography and 4D-STEM.

ptyrodactyl.invert

Inverse reconstruction algorithms for electron ptychography including single-slice, position-corrected, and multi-modal reconstruction methods.

ptyrodactyl.tools

Common utilities and shared data structures used throughout the package, including complex-valued optimizers with Wirtinger derivatives.

ptyrodactyl.workflows

High-level workflow functions that combine primitives from simul and data structures from tools into complete end-to-end pipelines. These meta-functions provide convenient interfaces for common tasks like running full 4D-STEM simulations from structure files.

Design Principles

1. JAX-First Architecture

All functions are designed to be:

  • Differentiable: Full support for jax.grad

  • JIT-compilable: Optimized with jax.jit

  • Vectorizable: Compatible with jax.vmap

2. Type Safety

  • Type hints using jaxtyping

  • Runtime type checking with beartype

  • Using PyTrees for data containers

  • PyTrees are loaded with type-checked factory functions.

File Organization

The package structure is organized for clarity and maintainability:

src/ptyrodactyl/
├── __init__.py           # Top-level exports
├── simul/
│   ├── __init__.py       # Simul module exports
│   ├── atom_potentials.py # Atomic potential calculations
│   ├── geometry.py       # Geometric transformations
│   ├── preprocessing.py  # Data preprocessing utilities
│   └── simulations.py    # Forward simulation functions
├── invert/
│   ├── __init__.py       # Invert module exports
│   └── phase_recon.py    # Phase reconstruction algorithms
├── tools/
│   ├── __init__.py       # Tools module exports
│   ├── electron_types.py # Data structures and type definitions
│   ├── loss_functions.py # Loss function definitions
│   ├── optimizers.py     # Complex-valued optimizers
│   └── parallel.py       # Parallel processing utilities
└── workflows/
    ├── __init__.py       # Workflows module exports
    └── stem_4d.py        # End-to-end 4D-STEM workflows

Extension Points

The package is designed to be extensible:

  1. Custom Loss Functions: Implement new loss functions following the pattern in tools.loss_functions

  2. New Optimizers: Add optimizers with Wirtinger derivative support

  3. Additional Reconstructions: Build on base reconstruction algorithms in invert.phase_recon

  4. Custom Workflows: Combine existing functions for specific use cases

Future Directions

The package architecture supports future extensions:

  • Real-time microscopy inversion

  • Additional electron microscopy modalities

  • Machine learning-enhanced reconstructions