Improve workflow (#125)

- migrate from gladier and fair-research-login to directly use globus-sdk
- fix probe orthogonalization bug that was introduced in previous release
- move data visualization methods to API
- switch to colorcet for perceptually uniform color maps
- improve forward model for synthetic diffraction pattern generation
- improve API documentation & test coverage

Author: stevehenke

docs/source/api.rst

@@ -22,6 +22,14 @@ Diffraction
    :undoc-members:
    :show-inheritance:
 
+Diffraction Generators
+----------------------
+
+.. automodule:: ptychodus.api.diffraction_gen
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Fluorescence
 ------------
 
@@ -46,6 +54,14 @@ Input/Output (I/O)
    :undoc-members:
    :show-inheritance:
 
+Interpolation
+-------------
+
+.. automodule:: ptychodus.api.interpolate
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Object
 ------
 
@@ -54,6 +70,14 @@ Object
    :undoc-members:
    :show-inheritance:
 
+Object Generators
+-----------------
+
+.. automodule:: ptychodus.api.object_gen
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Observer
 --------
 
@@ -86,6 +110,14 @@ Probe
    :undoc-members:
    :show-inheritance:
 
+Probe Generators
+----------------
+
+.. automodule:: ptychodus.api.probe_gen
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Probe Positions
 ---------------
 
@@ -94,6 +126,14 @@ Probe Positions
    :undoc-members:
    :show-inheritance:
 
+Probe Position Generators
+-------------------------
+
+.. automodule:: ptychodus.api.probe_positions_gen
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Product
 -------
 

genesis.svg

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=80", "setuptools_scm[simple]>=8"]
+requires = ["setuptools>=80", "setuptools-scm[simple]>=9.2"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -10,11 +10,13 @@ requires-python = ">=3.11"
 license = "BSD-3-Clause"
 license-files = ["LICENSE"]
 dependencies = [
+    "colorcet",
     "h5py>=3",
     "hdf5plugin",
     "matplotlib",
     "numpy",
     "psutil",
+    "pydantic",
     "pyyaml",
     "requests",
     "scikit-image",
@@ -31,10 +33,15 @@ ptychodus-bdp = "ptychodus.ptychodus_bdp:main"
 ptychodus-ptychopinn-tf-test = "ptychodus.scripts.ptychopinn_tf_test:main"
 
 [project.optional-dependencies]
-globus = ["gladier", "gladier-tools>=0.5.4"]
+globus = ["globus-compute-sdk", "globus-sdk>=4.5.0"]
 gui = ["PyQt5"]
 #ptychonn = ["ptychonn==0.3.*,>=0.3.7"]
 ptychi = ["ptychi==1.*,>=1.2.0"]
+ptychopinn = ["ptychopinn[torch]"]
+
+[tool.uv.sources]
+#ptychopinn = { path = "../PtychoPINN", editable = true }
+ptychopinn = { git = "https://github.com/hoidn/PtychoPINN", branch = "torchapi-devel" }
 
 [tool.setuptools.package-data]
 "ptychodus" = ["py.typed"]
@@ -46,16 +53,14 @@ where = ["src"]
 
 [[tool.mypy.overrides]]
 module = [
-    "fair_research_login",
-    "funcx_endpoint.*",
-    "gladier.*",
+    "colorcet",
     "globus_sdk.*",
     "h5py",
     "hdf5plugin",
     "lightning",
-    "parsl.*",
     "ptychi.*",
     "ptycho.*",
+    "ptycho_torch.*",
     "ptychonn.*",
     "pvaccess",
     "pvapy.*",

pyproject.toml

@@ -1,3 +1,5 @@
+"""Phase unwrapping via gradient-domain integration for complex 2D wavefields."""
+
 from scipy import ndimage
 from scipy.fft import fft, fft2, fftfreq, ifft, ifft2
 from scipy.signal.windows import gaussian as gaussian_window
@@ -8,6 +10,8 @@
 
 
 class PhaseUnwrapper:
+    """Gradient-domain phase unwrapper for complex 2D images."""
+
     def __init__(
         self,
         fourier_shift_step: float = 0.5,

src/ptychodus/api/_phase_unwrapper.py

@@ -1,5 +1,6 @@
 """Common type aliases, physical constants, and utility functions used throughout the API."""
 
+from pathlib import Path
 from typing import Any, Final, TypeAlias, overload
 
 import numpy
@@ -24,6 +25,10 @@
 PLANCK_CONSTANT_J_PER_HZ: Final[float] = 6.62607015e-34
 
 
+def get_ptychodus_dir() -> Path:
+    return Path.home() / '.ptychodus'
+
+
 @overload
 def lerp(lower: float, upper: float, frac: float) -> float: ...
 @overload

src/ptychodus/api/common.py

@@ -0,0 +1,113 @@
+"""Forward model for simulating diffraction patterns from a ptychography data product."""
+
+import numpy
+
+from scipy.fft import fft2, fftfreq, ifft2
+
+from .diffraction import BadPixels, DiffractionIndexes, DiffractionPatterns
+from .geometry import PixelGeometry
+from .io import AssembledDiffractionData
+from .product import Product
+from .propagator import AngularSpectrumPropagator, FraunhoferPropagator, PropagatorParameters
+
+
+def generate_diffraction_data(
+    product: Product, rng: numpy.random.Generator | None = None
+) -> AssembledDiffractionData:
+    """Simulate diffraction patterns for all scan positions in *product* using a multislice forward model.
+
+    If *rng* is provided, Poisson noise is added to the intensity patterns.
+    """
+    object_ = product.object_
+    probe_geometry = product.probes.get_geometry()
+
+    propagator_parameters = PropagatorParameters(
+        wavelength_m=product.metadata.probe_wavelength_m,
+        width_px=probe_geometry.width_px,
+        height_px=probe_geometry.height_px,
+        pixel_width_m=probe_geometry.pixel_width_m,
+        pixel_height_m=probe_geometry.pixel_height_m,
+        propagation_distance_m=product.metadata.detector_distance_m,
+    )
+
+    # TODO also support near-field propagation
+    propagator = FraunhoferPropagator(propagator_parameters)
+
+    # One angular-spectrum propagator per inter-layer gap
+    interlayer_propagators = [
+        AngularSpectrumPropagator(
+            PropagatorParameters(
+                wavelength_m=product.metadata.probe_wavelength_m,
+                width_px=probe_geometry.width_px,
+                height_px=probe_geometry.height_px,
+                pixel_width_m=probe_geometry.pixel_width_m,
+                pixel_height_m=probe_geometry.pixel_height_m,
+                propagation_distance_m=spacing_m,
+            )
+        )
+        for spacing_m in object_.layer_spacing_m
+    ]
+
+    num_positions = len(product.probe_positions)
+    indexes: DiffractionIndexes = numpy.zeros(num_positions, dtype=int)
+    patterns: DiffractionPatterns = numpy.zeros(
+        (num_positions, probe_geometry.height_px, probe_geometry.width_px),
+        dtype=float,
+    )
+    lambda_z_m2 = propagator_parameters.wavelength_m * propagator_parameters.propagation_distance_m
+    pixel_geometry = PixelGeometry(
+        width_m=lambda_z_m2 / probe_geometry.width_m,
+        height_m=lambda_z_m2 / probe_geometry.height_m,
+    )
+    bad_pixels: BadPixels = numpy.full((probe_geometry.height_px, probe_geometry.width_px), False)
+
+    # Precompute frequency grids for Fourier shifting the probe
+    freq_y = fftfreq(probe_geometry.height_px)
+    freq_x = fftfreq(probe_geometry.width_px)
+    fy, fx = numpy.meshgrid(freq_y, freq_x, indexing='ij')
+
+    object_geometry = object_.get_geometry()
+
+    for index, (probe_position, probe) in enumerate(zip(product.probe_positions, product.probes)):
+        object_position = object_geometry.map_coordinates_probe_to_object(probe_position)
+
+        cx = object_position.coordinate_x_px
+        cy = object_position.coordinate_y_px
+
+        x_lower = int(cx - probe_geometry.width_px / 2)
+        y_lower = int(cy - probe_geometry.height_px / 2)
+
+        # Extract patches from all layers at the same integer position
+        object_patches = [
+            object_.get_layer(ilayer)[
+                y_lower : y_lower + probe_geometry.height_px,
+                x_lower : x_lower + probe_geometry.width_px,
+            ]
+            for ilayer in range(object_.num_layers)
+        ]
+
+        # Subpixel offsets between the true position and the integer extraction position
+        dx = cx - (x_lower + probe_geometry.width_px / 2)
+        dy = cy - (y_lower + probe_geometry.height_px / 2)
+
+        # Shift all incoherent modes at once via a single batched FFT pair
+        shift_phase = numpy.exp(-2j * numpy.pi * (fy * dy + fx * dx))
+        shifted_modes = ifft2(fft2(probe.get_array()) * shift_phase)
+
+        for wavefield in shifted_modes:
+            # Multislice: apply each layer then propagate to the next; last layer has no propagation
+            for object_patch, interlayer_propagator in zip(object_patches, interlayer_propagators):
+                wavefield = wavefield * object_patch
+                wavefield = interlayer_propagator.propagate(wavefield)
+            wavefield = wavefield * object_patches[-1]
+
+            wavefield = propagator.propagate(wavefield)
+            patterns[index] += numpy.square(numpy.abs(wavefield))
+
+        indexes[index] = object_position.index
+
+    if rng is not None:
+        # NOTE: object and probe scaling influence how much noise is added
+        patterns = rng.poisson(patterns).astype(float)
+
+    return AssembledDiffractionData(indexes, patterns, pixel_geometry, bad_pixels)

src/ptychodus/api/diffraction_gen.py

@@ -178,6 +178,8 @@ def __repr__(self) -> str:
 
 @dataclass(frozen=True)
 class ZernikeMonomial:
+    """A single Zernike polynomial term with a complex coefficient, radial degree n, and angular frequency m."""
+
     coefficient: complex
     radial_degree: int  # n
     angular_frequency: int  # m

src/ptychodus/api/geometry.py

@@ -1,3 +1,5 @@
+"""Sub-pixel array interpolation and stitching utilities."""
+
 from typing import Any, Generic, TypeVar
 import logging
 
@@ -18,12 +20,15 @@
 
 
 def calculate_support_frac(x: float, n: int) -> tuple[slice, float]:
+    """Return the integer support slice and sub-pixel fractional offset for a width-n window centered at x."""
     lower = x - n / 2
     whole = int(lower)
     return slice(whole, whole + n + 1), lower - whole
 
 
 class NearestNeighborArrayInterpolator(Generic[InexactDType]):
+    """Extract patches from an array using integer (nearest-neighbor) pixel centering."""
+
     def __init__(self, array: numpy.typing.NDArray[InexactDType]) -> None:
         super().__init__()
         self._array = array
@@ -43,6 +48,8 @@ def get_patch(
 
 
 class BarycentricArrayInterpolator(Generic[InexactDType]):
+    """Extract patches from an array using bilinear (barycentric) sub-pixel interpolation."""
+
     def __init__(self, array: numpy.typing.NDArray[InexactDType]) -> None:
         super().__init__()
         self._array = array
@@ -72,6 +79,8 @@ def get_patch(
 
 
 class BarycentricArrayStitcher(Generic[InexactDType]):
+    """Accumulate weighted patches into a canvas using bilinear sub-pixel spreading, then normalize."""
+
     def __init__(
         self, upper: numpy.typing.NDArray[InexactDType], lower: RealArrayType | None = None
     ) -> None:
@@ -129,6 +138,7 @@ def add_patch(
             lsupport[..., 1:, 1:] += weight11 * weight
 
     def stitch(self) -> numpy.typing.NDArray[InexactDType]:
+        """Return the accumulated canvas divided by the weight array; pixels with zero weight remain zero."""
         if self._lower is None:
             return self._upper
 

…tychodus/model/analysis/interpolators.py src/ptychodus/api/interpolate.pysrc/ptychodus/model/analysis/interpolators.py renamed to src/ptychodus/api/interpolate.py src/ptychodus/model/analysis/interpolators.py renamed to src/ptychodus/api/interpolate.py

@@ -26,6 +26,8 @@
 
 
 class StandardFileLayout(StrEnum):
+    """Conventional file names used in the ptychodus standard HDF5 workflow directory."""
+
     DIFFRACTION = 'diffraction.h5'
     FLUORESCENCE_IN = 'fluorescence-in.h5'
     FLUORESCENCE_OUT = 'fluorescence-out.h5'
@@ -35,6 +37,8 @@ class StandardFileLayout(StrEnum):
 
 
 class DiffractionFileKeys(StrEnum):
+    """HDF5 dataset and attribute names for the assembled diffraction data file."""
+
     PATTERNS = 'patterns'
     DETECTOR_PIXEL_HEIGHT = 'detector_pixel_height_m'
     DETECTOR_PIXEL_WIDTH = 'detector_pixel_width_m'
@@ -43,6 +47,7 @@ class DiffractionFileKeys(StrEnum):
 
 
 def load_diffraction_data(file: Path, *, mmap_file: Path | None = None) -> AssembledDiffractionData:
+    """Load assembled diffraction data from an HDF5 file written by :func:`save_diffraction_data`."""
     if mmap_file is not None:
         raise NotImplementedError('Load to memory map not implemented yet!')
 
@@ -78,6 +83,7 @@ def load_diffraction_data(file: Path, *, mmap_file: Path | None = None) -> Assem
 def save_diffraction_data(
     file: Path, data: AssembledDiffractionData, *, compression: str = 'lzf'
 ) -> None:
+    """Write assembled diffraction data to an HDF5 file."""
     with h5py.File(file, 'w') as h5_file:
         h5_file.create_dataset(
             DiffractionFileKeys.INDEXES, data=data._indexes, compression=compression
@@ -98,6 +104,8 @@ def save_diffraction_data(
 
 
 class ProductFileKeys(StrEnum):
+    """HDF5 dataset and attribute names for the product file."""
+
     NAME = 'name'
     COMMENTS = 'comments'
     DETECTOR_OBJECT_DISTANCE = 'detector_object_distance_m'
@@ -124,6 +132,7 @@ class ProductFileKeys(StrEnum):
 
 
 def load_product(file: Path) -> Product:
+    """Load a data product from an HDF5 file written by :func:`save_product`."""
     point_list: list[ProbePosition] = []
 
     with h5py.File(file, 'r') as h5_file:
@@ -257,6 +266,7 @@ def load_product(file: Path) -> Product:
 
 
 def save_product(file: Path, product: Product) -> None:
+    """Write a data product to an HDF5 file."""
     scan_indexes: list[int] = []
     scan_x_m: list[float] = []
     scan_y_m: list[float] = []