Package Organization

Overview

Janssen is a typed and tested JAX based library for optical microscopy and ptychography, to utilize two of JAX’s capabilities - multi-device computation and autodifferentiation. The package is organized into nine main modules: types for PyTree data structures, utils for common math utilities, optics for optical elements, prop for propagation algorithms, scopes for microscope forward models, lenses for lens implementations, models for test pattern generation, cohere for coherence modeling, and invert for reconstruction algorithms.

Module Structure

janssen.types

PyTree data structures and factory functions for type-safe JAX programming. Includes OpticalWavefront, CoherentModeSet, MixedStatePtychoData, and other core types.

janssen.utils

Common mathematical utilities including Bessel functions, Fourier shifting, Wirtinger gradients, and distributed computing helpers.

janssen.optics

Optical elements and transformations for simulating light propagation through various optical components.

janssen.cohere

Coherence modeling for partially coherent optical fields. Includes Gaussian Schell-model mode generation, spatial coherence kernels, and coherent mode propagation.

janssen.scopes

Microscope forward models for simulating image formation and diffraction patterns in optical microscopy.

janssen.prop

Propagation algorithms for simulating optical field propagation through different media. Includes free-space propagation methods (angular spectrum, Fresnel, Fraunhofer) and material propagation for volumetric optical media.

janssen.lenses

Lens element definitions and optical calculations for modeling various lens types.

janssen.models

Models for generating test patterns and optical beams for testing and validation.

janssen.invert

The reconstruction module containing phase retrieval algorithms, ptychography, and mixed-state reconstruction for partially coherent sources.

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/janssen/
├── __init__.py           # Top-level exports
├── types/
│   ├── __init__.py       # Types module exports
│   ├── scalars.py        # Scalar type definitions
│   ├── wavefronts.py     # OpticalWavefront, PropagatingWavefront
│   ├── coherence_types.py # CoherentModeSet, MixedStatePtychoData
│   └── factory.py        # PyTree factory functions
├── utils/
│   ├── __init__.py       # Utils module exports
│   ├── bessel.py         # Bessel functions (J, K)
│   ├── math.py           # Fourier shift, Wirtinger gradients
│   └── distributed.py    # Multi-device utilities
├── optics/
│   ├── __init__.py       # Optics module exports
│   ├── apertures.py      # Aperture functions
│   ├── elements.py       # Optical elements
│   ├── helper.py         # Utility functions
│   └── zernike.py        # Zernike polynomials
├── cohere/
│   ├── __init__.py       # Coherence module exports
│   ├── modes.py          # Gaussian Schell-model modes
│   ├── spatial.py        # Spatial coherence kernels
│   └── propagation.py    # Coherent mode propagation
├── prop/
│   ├── __init__.py       # Prop module exports
│   ├── free_space_prop.py # Free-space propagation algorithms
│   └── material_prop.py  # Material propagation
├── scopes/
│   ├── __init__.py       # Scopes module exports
│   └── simple_microscopes.py  # Simple microscope forward models
├── models/
│   ├── __init__.py       # Models module exports
│   ├── beams.py          # Optical beam generation
│   └── usaf_pattern.py   # USAF test pattern generation
├── lenses/
│   ├── __init__.py       # Lenses module exports
│   └── lens_elements.py  # Lens element definitions
└── invert/
    ├── __init__.py       # Invert module exports
    ├── engine.py         # PIE reconstruction engine
    ├── ptychography.py   # Ptychographic algorithms
    ├── mixed_state.py    # Mixed-state ptychography
    ├── initialization.py # Reconstruction initialization
    └── loss_functions.py # Loss function definitions

Extension Points

The package is designed to be extensible:

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

  2. New Optimizers: Add optimizers with Wirtinger derivative support

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

  4. Custom Optical Elements: Add new elements in optics.elements

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

Future Directions

The package architecture supports future extensions:

  • Real-time microscopy inversion

  • Adaptive optics simulations

  • Non-linear optics