Extending Pty-chi with External Algorithm Packages

Pty-chi already ships with several well-tested reconstructors. When a new reconstruction strategy is still evolving or serves a narrow use case, the preferred approach is to publish it as its own Python package that depends on ptychi. Keeping experimental algorithms outside of this repository keeps the core lean, reduces the review burden on the maintainers, and lets you iterate without waiting for upstream releases.

The guidance below walks through the recommended structure for an extension project that reuses Pty-chi's data structures, tasks, and utilities while introducing new algorithms in a separate codebase.

Why use a companion repository

• Ensure the main project remains focused on widely adopted, well-supported algorithms.
• Ship a faster release cadence for your research code without asking for core maintainer reviews on every iteration.
• Give yourself freedom to add external dependencies, larger models, or data assets that would be out of scope for Pty-chi proper.
• Maintain clear boundaries: bug fixes and API improvements belong upstream, while novel method development can live downstream.

Package layout checklist

1. Create a new Python project with standard packaging metadata.
2. Declare Pty-chi as a dependency (pin to the minimum version that offers the extension points you need).
3. Reuse Pty-chi data structures (PtychographyTask, parameter groups, probes, etc.) and subclass the base Reconstructor and Options types to introduce custom behaviour.
4. Expose a stable entry point (CLI, script, or library function) that instantiates your options and reconstructor.
5. Stay in sync with new Pty-chi releases by running your extension's tests against them.

Example pyproject.toml

[project]
name = "my-ptychi-algorithms"
version = "0.1.0"
dependencies = [
  "ptychi>=0.9",
  "torch>=2.2",  # any extra requirements your algorithm needs
]

Extending options and reconstructor classes

Create tailored options by subclassing the Pty-chi dataclasses. Override or extend fields to capture algorithm-specific hyperparameters, and implement get_reconstructor_type if the base task logic should still create a core reconstructor.

from dataclasses import dataclass

from ptychi.api import enums
from ptychi.api.options.base import ReconstructorOptions
from ptychi.api.options.task import PtychographyTaskOptions


@dataclass
class MyAlgorithmReconstructorOptions(ReconstructorOptions):
    total_variation_weight: float = 5e-4
    diffusion_strength: float = 0.3

    def get_reconstructor_type(self) -> enums.Reconstructors:
        # Return an existing enum if you intend to reuse a built-in reconstructor,
        # otherwise the task subclass below decides how to instantiate yours.
        return enums.Reconstructors.base


@dataclass
class MyAlgorithmTaskOptions(PtychographyTaskOptions):
    reconstructor_options: MyAlgorithmReconstructorOptions = MyAlgorithmReconstructorOptions()

Subclass an appropriate Reconstructor class of Pty-Chi, override relevant methods for your own implementation:

from ptychi.reconstructors.base import IterativePtychographyReconstructor


class MyAlgorithmReconstructor(IterativePtychographyReconstructor):
    def build(self):
        super().build()
        # add checkpoints, priors, or external models here

    def run(self, n_epochs: int | None = None):
        # customise the iteration loop while still calling into the base utilities
        return super().run(n_epochs=n_epochs)

Particularly, consider only overriding run_minibatch, which only processes a batch of data, if your high-level workflow is already abstracted by the run method.

If your reconstructor does not align with one of the built-in enum entries, override build_reconstructor in a task subclass to instantiate it explicitly.

from ptychi.api.task import PtychographyTask
from ptychi.data_structures.parameter_group import PtychographyParameterGroup

from .reconstructors.my_algorithm import MyAlgorithmReconstructor


class MyAlgorithmTask(PtychographyTask):
    def build_reconstructor(self):
        parameter_group = PtychographyParameterGroup(
            object=self.object,
            probe=self.probe,
            probe_positions=self.probe_positions,
            opr_mode_weights=self.opr_mode_weights,
        )
        self.reconstructor = MyAlgorithmReconstructor(
            parameter_group=parameter_group,
            dataset=self.dataset,
            options=self.reconstructor_options,
        )
        self.reconstructor.build()

This arrangement keeps the underlying tensor layouts, buffer management, and device handling provided by Pty-chi while letting you plug in new optimisation strategies.

Documenting and testing your extension

• Mirror the option schema in your README so users know how to configure the new algorithm.
• Provide example scripts that construct MyAlgorithmTaskOptions, load data, and run a reconstruction end-to-end.
• Reuse Pty-chi factories (ptychi.maps.*) wherever possible rather than reimplementing low-level utilities.
• Set up unit tests that import your package alongside the pinned Pty-chi version and cover both successful runs and expected failures.

When to change Pty-chi directly

• Bug fixes, documentation improvements, and new extension hooks belong in this repository.
• Breaking API changes should be discussed with maintainers before opening a pull request.
• Large algorithmic additions should be upstreamed only after they stabilise in an external package and gain broad adoption.

By developing new methods in companion repositories, you help keep Pty-chi maintainable and reliable while still making advanced capabilities available to the community.