diff --git a/.claude/skills/add-core/SKILL.md b/.claude/skills/add-core/SKILL.md new file mode 100644 index 000000000..42fb7eb99 --- /dev/null +++ b/.claude/skills/add-core/SKILL.md @@ -0,0 +1,97 @@ +--- +name: add-core +description: Scaffold a new *Core subsystem (settings + parameters + observers + api) and wire it through ModelCore, ViewCore, and ControllerCore. Use when the user says "add a new analysis/processing feature", "add a new model component", or is introducing a subsystem that needs its own settings group and GUI panel. +--- + +# add-core + +Every ptychodus subsystem follows the same shape: a `*Core` class in `src/ptychodus/model//` owns its `SettingsRegistry` group, typed parameters, sub-components, and API; `model/core.py::ModelCore` composes it; and the view + controller layers mirror the layout. This skill encodes that shape. + +## Layer boundaries (do not cross) + +- `src/ptychodus/api/` — pure domain, no dependencies on model/view/controller. +- `src/ptychodus/model/` — logic; depends on `api/` only. +- `src/ptychodus/view/` — PyQt5 widgets; no logic. +- `src/ptychodus/controller/` — mediates view ↔ model. + +Run the `check-layers` skill after wiring to confirm you haven't accidentally introduced a model→view import. + +## Steps + +### 1. Settings + +Create `src/ptychodus/model//settings.py`. Follow the pattern in [src/ptychodus/model/analysis/settings.py](../../src/ptychodus/model/analysis/settings.py): + +```python +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.settings import SettingsRegistry + + +class Settings(Observable, Observer): + def __init__(self, registry: SettingsRegistry) -> None: + super().__init__() + self._group = registry.create_group('') + self._group.add_observer(self) + + # Typed parameters: + self.num_iterations = self._group.create_integer_parameter( + 'NumberOfIterations', 1000, minimum=1 + ) + # create_real_parameter, create_boolean_parameter, create_string_parameter, + # create_path_parameter — see api/parametric.py for the full menu. + + def _update(self, observable: Observable) -> None: + if observable is self._group: + self.notify_observers() +``` + +Preserve unit suffixes on physical quantities (e.g. `probe_energy_eV`, `pixel_width_m`) — CLAUDE.md notes these are accepted via `# noqa: N815` and reviewers expect them. + +### 2. Core + +Create `src/ptychodus/model//core.py`: + +```python +class Core: + def __init__( + self, + settings_registry: SettingsRegistry, + # ... other dependencies from ModelCore (repositories, other Cores) + ) -> None: + self.settings = Settings(settings_registry) + # ... construct sub-components, wire observers +``` + +For an example that composes many sub-components + `VisualizationEngine`s, see [src/ptychodus/model/analysis/core.py](../../src/ptychodus/model/analysis/core.py). + +### 3. Compose in ModelCore + +In [src/ptychodus/model/core.py](../../src/ptychodus/model/core.py), import your `Core` and construct it in `ModelCore.__init__` in the correct dependency order — after everything it needs, before anything that needs it. `ModelCore` is the single composition root; do not construct components anywhere else. + +### 4. View + +Create `src/ptychodus/view//` mirroring the model layout. PyQt5 widgets only — no logic, no imports from `model/` or `controller/`. + +### 5. Controller + +Create `src/ptychodus/controller//` with the `*ViewController` that bridges the view widgets to `Core`. + +### 6. Wire into ViewCore and ControllerCore + +`ViewCore` and `ControllerCore` (in `src/ptychodus/view/core.py` and `src/ptychodus/controller/core.py`) are composition roots. Register the new panel in both. **The navigation toolbar order in `ViewCore` is the source of truth for the left/right stacked-panel indexes** (per CLAUDE.md) — pick the toolbar position deliberately and ensure `ControllerCore` uses the matching stacked-widget index. + +## PluginChooser binding for settings-driven selection + +If your feature has a "current selected implementation" (e.g. current file reader), bind a `PluginChooser` to a `StringParameter` in your settings: + +```python +self.file_reader_chooser.synchronize_with_parameter(self.settings.file_type) +``` + +`synchronize_with_parameter` in [src/ptychodus/api/plugins.py](../../src/ptychodus/api/plugins.py) sets up the two-way binding: parameter value ↔ current plugin. + +## Verify + +- `check-layers` — no boundary violations. +- `pre-push` — full CI gate green. +- Launch the GUI (`uv run ptychodus`) and confirm the new panel shows up in the correct toolbar position. diff --git a/.claude/skills/add-nav-icon/SKILL.md b/.claude/skills/add-nav-icon/SKILL.md new file mode 100644 index 000000000..4c3ae2865 --- /dev/null +++ b/.claude/skills/add-nav-icon/SKILL.md @@ -0,0 +1,103 @@ +--- +name: add-nav-icon +description: Add a navigation-bar icon to ptychodus in both the PyQt GUI and the ptychodus_store web UI. Covers SVG placement, Qt resource compilation with the manual typing fix that pyrcc5 clobbers, ViewCore wiring, nav.ts wiring, and the tsc rebuild. Use when the user says "add a nav icon", "add a toolbar icon", "wire up an icon for the new X panel", or is following up on add-core / a new panel. +--- + +# add-nav-icon + +Ptychodus has two navigation surfaces that share icon assets: the PyQt5 toolbar in `ViewCore` and the browser nav bar in `ptychodus_store`. SVGs live once in [`src/ptychodus_store/ui/icons/`](../../src/ptychodus_store/ui/icons/); the PyQt side pulls them in via a Qt resource file compiled to `resources.py`, and the web UI pulls them straight from the FastAPI static mount. This skill wires a new icon through both. + +## Step 1 — Add the SVG + +Drop the SVG into [`src/ptychodus_store/ui/icons/`](../../src/ptychodus_store/ui/icons/). Do not put icons anywhere else — the `.qrc` file uses a relative path (`../../ptychodus_store/ui/icons/…`) to reach exactly this directory. + +- Kebab-case filename (e.g. `my-icon.svg`). +- `fill="currentColor"` so CSS themes work. +- Font Awesome 7.1.0 (CC BY 4.0) is the house style. See [`icons/Font-Awesome-LICENSE.txt`](../../src/ptychodus_store/ui/icons/Font-Awesome-LICENSE.txt) for attribution and [`icons/README.md`](../../src/ptychodus_store/ui/icons/README.md) for the update procedure. + +## Step 2 — Register in the Qt resource file + +Add one line to [`src/ptychodus/view/resources.qrc`](../../src/ptychodus/view/resources.qrc), keeping the list alphabetized by `alias`: + +```xml +../../ptychodus_store/ui/icons/my-icon.svg +``` + +The `alias` is what code references as `:/icons/my-feature`. Convention: alias matches the panel concept, not the filename (e.g. `atom.svg` is aliased `fluorescence`). + +## Step 3 — Regenerate `resources.py` and reapply the typing fix + +`resources.py` is auto-generated binary. Regenerate from inside the venv: + +```sh +cd src/ptychodus/view +uv run pyrcc5 -o resources.py resources.qrc +``` + +Fresh `pyrcc5` output writes the two module functions **unannotated and camelCase**, which fails both `mypy` (missing return type) and `ruff` (`N802`, function name should be lowercase). Reapply the manual typing fix — patch the tail of the file so the two functions look like this: + +```python +def qInitResources() -> None: # noqa: N802 + QtCore.qRegisterResourceData( + rcc_version, qt_resource_struct, qt_resource_name, qt_resource_data + ) + + +def qCleanupResources() -> None: # noqa: N802 + QtCore.qUnregisterResourceData( + rcc_version, qt_resource_struct, qt_resource_name, qt_resource_data + ) + + +qInitResources() +``` + +Two edits per function: append `-> None` to the signature and `# noqa: N802` to the same line. Leave the `WARNING! All changes made in this file will be lost!` header — the noqa comments and return annotations are the only manual patches, and this pattern must survive every regeneration. (The import site at [`view/core.py:30`](../../src/ptychodus/view/core.py#L30) uses `from . import resources # noqa` for the same reason — it's imported for side effects.) + +## Step 4 — Wire it into `ViewCore` + +Add a call to `self.navigation.add_panel(...)` in [`src/ptychodus/view/core.py`](../../src/ptychodus/view/core.py) `ViewCore.__init__`. Order matters — **the sequence of `add_panel` calls is the source of truth for the left/right stacked-panel indexes** (per CLAUDE.md). Mirror the fluorescence pattern at [`view/core.py:275-280`](../../src/ptychodus/view/core.py#L275-L280): + +```python +self.my_feature_view = MyFeatureView() +self.my_feature_image_view = MyFeatureImageView() +self.my_feature_action = self.navigation.add_panel( + QIcon(':/icons/my-feature'), + 'My Feature', + left=self.my_feature_view, + right=self.my_feature_image_view, +) +``` + +If the new panel belongs nested under a parent (Products or Processing today), add its action to the existing `add_subview_group` call at [`view/core.py:336-352`](../../src/ptychodus/view/core.py#L336-L352). Otherwise it renders as a top-level toolbar button. + +`ControllerCore` uses stacked-widget indexes matching `ViewCore`'s `add_panel` order — inserting a panel in the middle shifts every downstream index. Prefer appending, or run the app afterward and confirm nothing underneath shifted. + +## Step 5 — Wire it into the web UI + +Add one entry to the `NAV` array in [`src/ptychodus_store/ui/src/nav.ts`](../../src/ptychodus_store/ui/src/nav.ts): + +```typescript +{ route: 'my-feature', label: 'My Label', icon: 'my-icon.svg' }, +``` + +- `route` — page identifier used by the front-end router. +- `label` — tooltip / text under the button. +- `icon` — **filename only**; the render code prefixes `/ui/icons/`. + +No backend router change is needed. FastAPI already serves `src/ptychodus_store/ui/` at `/ui/` via the `StaticFiles` mount in [`src/ptychodus_store/app.py`](../../src/ptychodus_store/app.py), so a new SVG in `ui/icons/` is automatically reachable at `/ui/icons/my-icon.svg`. + +## Step 6 — Rebuild the TypeScript + +```sh +cd src/ptychodus_store/ui +tsc +``` + +No bundler, no framework — plain `tsc` compiles `src/**/*.ts` into `dist/` (see `tsconfig.json`). Restart the store server so the fresh compiled JS is loaded; `tsc --watch` is fine for iteration. + +## Verify + +- **PyQt:** `uv run ptychodus` — the new icon appears at the correct toolbar position (top-level or nested), highlights on click, and swaps in the associated left/right panels. +- **Web UI:** run the `store-dev` skill, load the browser UI, confirm the icon renders in the nav bar (network tab should show a 200 for `/ui/icons/my-icon.svg`). +- **CI gate:** run the `pre-push` skill. If `pyrcc5` wiped the typing fix, `ruff` will fail with `N802` on `qInitResources` / `qCleanupResources` — reapply Step 3 and re-run. diff --git a/.claude/skills/add-plugin/SKILL.md b/.claude/skills/add-plugin/SKILL.md new file mode 100644 index 000000000..39f689db0 --- /dev/null +++ b/.claude/skills/add-plugin/SKILL.md @@ -0,0 +1,69 @@ +--- +name: add-plugin +description: Scaffold a new file-format plugin (diffraction / probe / probe-position / object / product / fluorescence reader or writer) under src/ptychodus/plugins/. Use when the user says "add a plugin for X format", "add a reader/writer for Y", or is integrating a new beamline data format. +--- + +# add-plugin + +Plugins are auto-discovered by `PluginRegistry.load_plugins()` in [src/ptychodus/api/plugins.py](../../src/ptychodus/api/plugins.py). Every module under `src/ptychodus/plugins/` that defines a module-level `register_plugins(registry)` function is imported and registered at startup; modules whose imports fail are logged and skipped (this is how optional-dependency plugins silently disable themselves). + +## Steps + +1. **Identify the abstract interface.** Pick the right one from `ptychodus.api`: + + | You want to read/write... | Interface | Registry chooser | + | --- | --- | --- | + | Raw diffraction patterns | `DiffractionFileReader/Writer` (`api/diffraction.py`) | `diffraction_file_readers` / `_writers` | + | Probe positions | `ProbePositionFileReader/Writer` (`api/probe_positions.py`) | `probe_position_file_readers` / `_writers` | + | Probes | `ProbeFileReader/Writer` (`api/probe.py`) | `probe_file_readers` / `_writers` | + | Objects | `ObjectFileReader/Writer` (`api/object.py`) | `object_file_readers` / `_writers` | + | A full Product (probe + positions + object bundled) | `ProductFileReader/Writer` (`api/product.py`) | see step 4 | + | Fluorescence / bad pixels | `Fluorescence*` / `BadPixelsFileReader` | corresponding choosers | + +2. **Create the module** at `src/ptychodus/plugins/__file.py`. Follow the naming convention already in use — e.g. `csv_probe_file.py`, `aps33id_velociprobe/`, `aps19id_isn_diffraction_file.py`. Beamline plugins may be a subpackage if they need more than one file. + +3. **Implement Reader and/or Writer** by subclassing the interface and providing the `read(path) -> ...` / `write(path, data) -> None` method. Keep the imports minimal — put any optional dependency import at module top so a `ModuleNotFoundError` cleanly disables the plugin. + +4. **Register.** Add `register_plugins(registry: PluginRegistry)` at module level. + + - **For a product reader**, prefer `register_product_file_reader_with_adapters` so the same reader is also reachable as a probe/probe-position/object reader without extra boilerplate: + + ```python + def register_plugins(registry: PluginRegistry) -> None: + registry.register_product_file_reader_with_adapters( + MyProductReader(), + display_name='My Format (*.h5)', + simple_name='MyFormat', + ) + ``` + + - **For everything else**, register on the appropriate chooser with both `display_name` (human-readable, includes the glob) and `simple_name` (short token used by the settings string): + + ```python + def register_plugins(registry: PluginRegistry) -> None: + registry.probe_file_readers.register_plugin( + MyProbeReader(), + simple_name='MyFormat', + display_name='My Format Files (*.myext)', + ) + ``` + + `simple_name` defaults to a stripped alphanumeric form of `display_name`; provide it explicitly if you want a stable settings token. + +## Reference examples + +Read one of these before writing new code — pick the closest match: + +- Simple single-file reader+writer pair: [src/ptychodus/plugins/csv_probe_file.py](../../src/ptychodus/plugins/csv_probe_file.py) +- HDF5-backed diffraction: [src/ptychodus/plugins/h5_diffraction_file.py](../../src/ptychodus/plugins/h5_diffraction_file.py) +- Product reader with adapters: [src/ptychodus/plugins/cxi_file.py](../../src/ptychodus/plugins/cxi_file.py), [src/ptychodus/plugins/h5_product_file.py](../../src/ptychodus/plugins/h5_product_file.py) +- Beamline subpackage: [src/ptychodus/plugins/aps33id_velociprobe/](../../src/ptychodus/plugins/aps33id_velociprobe/) +- Optional-dep plugin (fails cleanly if the dep is missing): [src/ptychodus/plugins/lcls_file_readers.py](../../src/ptychodus/plugins/lcls_file_readers.py) + +## Index preservation + +Diffraction patterns and probe positions are paired by **integer scan index**, not array order (see CLAUDE.md "Index-based pattern/position association"). A new diffraction reader must set `DiffractionArray.get_indexes()`; a new probe-position reader must set `ProbePosition.index`. A round-trip test that preserves indexes is worth writing. + +## Verify + +After creating the plugin, run the `check-plugin-load` skill to confirm it registers cleanly, then run `pre-push` before committing. diff --git a/.claude/skills/add-reconstructor/SKILL.md b/.claude/skills/add-reconstructor/SKILL.md new file mode 100644 index 000000000..9a5fd77ed --- /dev/null +++ b/.claude/skills/add-reconstructor/SKILL.md @@ -0,0 +1,111 @@ +--- +name: add-reconstructor +description: Scaffold a new reconstructor backend library (a *ReconstructorLibrary that exposes Reconstructor / TrainableReconstructor implementations to ProcessingCore). Use when the user says "add a new reconstruction backend", "integrate as a reconstructor", or "add a new ptychi/ptychopinn-style backend". +--- + +# add-reconstructor + +Reconstructor backends are optional dependencies. Each lives under `src/ptychodus/model//` and exposes a `*ReconstructorLibrary` that `ModelCore` composes and hands to `ProcessingCore`. The library must degrade gracefully when the backend package isn't installed — the GUI should still open. + +## Steps + +### 1. Directory layout + +Create `src/ptychodus/model//` with: + +- `__init__.py` — re-exports the `*ReconstructorLibrary` class only. +- `core.py` — the `*ReconstructorLibrary` class. +- `settings.py` — one or more `*Settings` classes (see the `add-core` skill for the settings pattern). +- `reconstructor.py` (or several) — the actual `Reconstructor` / `TrainableReconstructor` implementations, importing the backend package at module top so a missing dep produces a clean `ModuleNotFoundError`. + +### 2. Library class + +Model on [src/ptychodus/model/ptychopinn/core.py](../../src/ptychodus/model/ptychopinn/core.py): + +```python +from collections.abc import Iterator +from importlib.metadata import PackageNotFoundError, version +import logging + +from ptychodus.api.reconstructor import ( + NullReconstructor, Reconstructor, ReconstructorLibrary, TrainableReconstructor, +) +from ptychodus.api.settings import SettingsRegistry + +from .settings import Settings + +logger = logging.getLogger(__name__) + + +class ReconstructorLibrary(ReconstructorLibrary): + def __init__( + self, settings_registry: SettingsRegistry, is_developer_mode_enabled: bool + ) -> None: + super().__init__('') + self.settings = Settings(settings_registry) + self._reconstructors: list[Reconstructor] = [] # or list[TrainableReconstructor] + + try: + from .reconstructor import Reconstructor + except ModuleNotFoundError: + logger.info(' not found.') + if is_developer_mode_enabled: + for name in ('Algorithm1', 'Algorithm2'): + self._reconstructors.append(NullReconstructor(name)) + else: + try: + pkg_version = version('') + except PackageNotFoundError: + pkg_version = 'unknown' + logger.info(f' {pkg_version}') + # Instantiate real reconstructors here. + self._reconstructors.append(Reconstructor('Algorithm1', self.settings)) + + @property + def name(self) -> str: + return '' + + def __iter__(self) -> Iterator[Reconstructor]: + return iter(self._reconstructors) +``` + +Key rules: + +- **The `try / except ModuleNotFoundError` for the backend import goes inside `__init__`**, gating the concrete reconstructor construction. Never make the top-level `core.py` fail on a missing backend. +- **In developer mode**, populate `NullReconstructor` stubs so the GUI still shows the algorithm names for testing. Outside developer mode, leave the list empty when the backend is missing. +- **Add `` as an optional extra** in `pyproject.toml` under `[project.optional-dependencies]`. + +### 3. Implement the reconstructor + +In `reconstructor.py`, subclass `Reconstructor` or `TrainableReconstructor` from [src/ptychodus/api/reconstructor.py](../../src/ptychodus/api/reconstructor.py). The core contract: + +- `reconstruct(parameters) -> Iterator[ReconstructOutput]` — yield one `ReconstructOutput` per iteration so the GUI can stream progress. +- For `TrainableReconstructor`, also implement `train(...)` and the ingest/export hooks. + +### 4. Wire into ModelCore + +In [src/ptychodus/model/core.py](../../src/ptychodus/model/core.py): + +1. Import `ReconstructorLibrary`. +2. Construct it in `ModelCore.__init__` after `settings_registry` is available. +3. Include it in the list passed to `ProcessingCore(...)`. + +Follow the existing pattern for `PtyChiReconstructorLibrary`, `PtychoNNReconstructorLibrary`, `PtychoPINNReconstructorLibrary`, `PtychoPINNTorchReconstructorLibrary`. + +### 5. Optional-dependency plumbing + +Add the backend to `pyproject.toml`: + +```toml +[project.optional-dependencies] + = [">="] +``` + +Document the install command in the backend's `core.py` docstring if the install isn't a plain `uv sync --extra ` (e.g. requires a sibling checkout). + +## Verify + +- `uv sync --extra ` — install the backend. +- `uv run ptychodus` — launch the GUI; the new backend should appear in the reconstructor list. Toggle `--log-level 10` to see the "found / not found" message. +- Without the extra (`uv sync` without `--extra `), the GUI still opens and other reconstructors still work; if developer mode is on, `NullReconstructor` stubs appear. +- `pre-push` — full CI gate green. diff --git a/.claude/skills/add-store-endpoint/SKILL.md b/.claude/skills/add-store-endpoint/SKILL.md new file mode 100644 index 000000000..2100e0428 --- /dev/null +++ b/.claude/skills/add-store-endpoint/SKILL.md @@ -0,0 +1,106 @@ +--- +name: add-store-endpoint +description: Add a matched pair of a FastAPI REST route and an MCP tool in the ptychodus_store service. Use when the user says "add an endpoint to the store", "expose X in the API", or "add an MCP tool for Y" — because REST and MCP surfaces must stay in sync. +--- + +# add-store-endpoint + +`ptychodus_store` exposes the same read-only surface through two channels: FastAPI REST under `/api/v1/*` (routers in `src/ptychodus_store/routers/`) and MCP tools at `/mcp` ([src/ptychodus_store/mcp_server.py](../../src/ptychodus_store/mcp_server.py)). They must return the same data for the same query. Adding one without the other silently drifts. + +## Steps + +### 1. Choose or create a router module + +- Existing resource? Extend an existing file in [src/ptychodus_store/routers/](../../src/ptychodus_store/routers/) (e.g. `diffraction.py`, `product.py`, `campaign.py`). +- New resource? Create `src/ptychodus_store/routers/.py`. + +### 2. REST route + +Follow the pattern in [src/ptychodus_store/routers/diffraction.py](../../src/ptychodus_store/routers/diffraction.py): + +```python +from fastapi import APIRouter, HTTPException, Query +from ptychodus_store.db import repositories as repo +from ptychodus_store.routers._convert import _to_read +from ptychodus_store.routers.deps import SessionDep, LayoutDep +from ptychodus_store.routers.schemas import Read, Page +from ptychodus_store.storage.manifest import ResourceKind + +router = APIRouter(prefix='/', tags=['']) + + +@router.get('', response_model=Page[Read]) +async def list_( + session: SessionDep, + limit: int = Query(50, ge=1, le=200), + offset: int = Query(0, ge=0), + # ... filter params +) -> Page[Read]: + ... +``` + +Use `SessionDep` for DB access and `LayoutDep` when you need the on-disk storage layout. Filtering uses SQLAlchemy `where` clauses built up from optional query params — see the diffraction router for the canonical shape. + +### 3. Register the router + +In [src/ptychodus_store/app.py](../../src/ptychodus_store/app.py) inside `create_app()`, add: + +```python +app.include_router(.router, prefix=api_prefix) +``` + +Match the position/style of the existing `app.include_router(...)` calls (health, campaign, diffraction, product, fluorescence, lineage, admin, visualization). + +### 4. Schemas + +Add pydantic response models to [src/ptychodus_store/routers/schemas.py](../../src/ptychodus_store/routers/schemas.py) if a new `*Read` shape is needed. Reuse `Page[T]` for paginated list responses. + +### 5. Converters + +If your ORM row → pydantic conversion is non-trivial (or needs a session query for related data), add a helper in [src/ptychodus_store/routers/_convert.py](../../src/ptychodus_store/routers/_convert.py) — this is what keeps REST and MCP in sync, because both channels call the same converter. + +### 6. Matching MCP tool + +Immediately add the corresponding tool in [src/ptychodus_store/mcp_server.py](../../src/ptychodus_store/mcp_server.py) inside `create_mcp_server()`. Every REST route must have an MCP counterpart with the same behavior: + +```python +@mcp.tool() +async def list_( + limit: int = 50, + offset: int = 0, + # ... same filter params as REST route +) -> Page[Read]: + """""" + async with _session() as session: + # ... same repo calls, same converter, same response +``` + +Rules: + +- MCP tools take `uuid: str` (they lack FastAPI's UUID coercion) — convert with `UUID(uuid_str)` inside the tool. +- Use the same `_convert` helper as the REST route. Do not duplicate business logic. +- Raise `ToolError` for MCP-side errors (not `HTTPException`). + +### 7. DB model (only if adding a new resource type) + +If this is a genuinely new resource (not a new query over an existing one), you also need: + +- SQLAlchemy model in `src/ptychodus_store/db/models.py` (UUID PK, `ingest_state`, metadata fields). +- New `ResourceKind` enum variant in `src/ptychodus_store/storage/manifest.py`. +- Ingestion path in `src/ptychodus_store/ingest/` (reconciler + watcher will pick it up automatically once the kind is registered). + +## Testing + +Health check after launching (see `store-dev` skill): + +```sh +curl -sf http://localhost:8000/api/v1/ +curl -sf http://localhost:8000/api/v1// +``` + +For the MCP tool, `fastmcp inspect http://localhost:8000/mcp` (or invoke via an MCP client) should list the new tool and return matching data. + +## Do not + +- Do not import from `ptychodus.model` or `ptychodus.view` — `ptychodus_store` is read-only from `ptychodus.api` only. +- Do not add write endpoints without user confirmation — the store surface is intentionally read-only. diff --git a/.claude/skills/check-layers/SKILL.md b/.claude/skills/check-layers/SKILL.md new file mode 100644 index 000000000..01dff0ff4 --- /dev/null +++ b/.claude/skills/check-layers/SKILL.md @@ -0,0 +1,76 @@ +--- +name: check-layers +description: Verify api/model/view+controller layer boundary compliance across the ptychodus source tree — no upward imports (api must not depend on model/view/controller; model must not depend on view/controller; ptychodus_store must not depend on ptychodus.model or ptychodus.view). Use when refactoring, before merging structural changes, or when the user asks to audit the architecture. +--- + +# check-layers + +CLAUDE.md documents a strict three-layer separation: + +- `src/ptychodus/api/` — pure domain, depends on nothing else in ptychodus. +- `src/ptychodus/model/` — application logic, depends only on `api/`. +- `src/ptychodus/view/` and `src/ptychodus/controller/` — GUI; may depend on `api/` and `model/`. +- `src/ptychodus_store/` — separate package; reads from `ptychodus.api` only, never from `ptychodus.model` or `ptychodus.view`. + +This skill greps for violations and reports file:line for anything out of place. + +## Steps + +Run each check. Report `OK` for each clean rule and file:line for each violation. + +### 1. api/ must not import model/, view/, or controller/ + +```sh +grep -rn --include='*.py' -E "^(from|import)[[:space:]]+ptychodus\.(model|view|controller)" src/ptychodus/api/ +``` + +Expected: no output. + +### 2. model/ must not import view/ or controller/ + +```sh +grep -rn --include='*.py' -E "^(from|import)[[:space:]]+ptychodus\.(view|controller)" src/ptychodus/model/ +``` + +Expected: no output. + +### 3. api/ and model/ must not import PyQt5 + +PyQt5 is a view-layer concern. Its presence in api/ or model/ is a smell even if no `ptychodus.view.*` import is used. + +```sh +grep -rn --include='*.py' -E "^(from|import)[[:space:]]+PyQt5" src/ptychodus/api/ src/ptychodus/model/ +``` + +Expected: no output. If found, the code should move to `view/` or `controller/`. + +### 4. ptychodus_store/ must not import ptychodus.model or ptychodus.view + +```sh +grep -rn --include='*.py' -E "^(from|import)[[:space:]]+ptychodus\.(model|view|controller)" src/ptychodus_store/ +``` + +Expected: no output. + +### 5. Circular-check inside model/ (informational) + +Not a hard rule, but flag any `model//` file importing from `model//` where `B` is composed *after* `A` in `ModelCore.__init__` — that's an ordering bug waiting to happen. + +```sh +grep -rn --include='*.py' -E "^from[[:space:]]+ptychodus\.model\.[a-z_]+[[:space:]]+import" src/ptychodus/model/ +``` + +Report the cross-subpackage imports and ask the user to verify the composition order in `src/ptychodus/model/core.py` still respects them. + +## Reporting + +- If every check is clean: report "All layer boundaries OK" and stop. +- If any check finds violations: report each violation as a bullet with `file:line` and the offending line's content. Suggest the correct home for the code (e.g. "Move `X` from `model/foo/bar.py:42` to `controller/foo/bar.py` — this uses `PyQt5.QtCore.QObject` which is view-layer"). +- Do not auto-fix. Layer violations often reflect a design choice the user needs to make (e.g. move code, extract an interface into api/, invert a dependency). + +## When this catches things + +- After a big refactor that moves code between layers. +- After adding a new subsystem via the `add-core` skill. +- Before merging structural PRs. +- When a new contributor's first PR touches multiple layers. diff --git a/.claude/skills/check-plugin-load/SKILL.md b/.claude/skills/check-plugin-load/SKILL.md new file mode 100644 index 000000000..fbfe433e2 --- /dev/null +++ b/.claude/skills/check-plugin-load/SKILL.md @@ -0,0 +1,70 @@ +--- +name: check-plugin-load +description: Diagnose why a ptychodus plugin isn't showing up in the file-format dropdowns (import errors, missing register_plugins, silent optional-dep skip). Use when the user says "my plugin isn't loading", "the reader/writer doesn't appear", or "plugin X is missing from the GUI". +--- + +# check-plugin-load + +`PluginRegistry.load_plugins()` in [src/ptychodus/api/plugins.py](../../src/ptychodus/api/plugins.py) walks every module under `src/ptychodus/plugins/`, imports it, and calls `register_plugins(registry)`. Failures are logged at WARNING and skipped — so a broken plugin just silently vanishes from the GUI. This skill surfaces the warnings and then confirms registration. + +## Steps + +### 1. Ask which plugin + +If the user hasn't named the plugin, ask: "Which plugin file or format is missing?" — the module name under `src/ptychodus/plugins/` is what you need. + +### 2. Load plugins with debug logging + +Run this — it triggers the same discovery walk the app does and prints every warning to stdout: + +```sh +uv run python -c " +import logging +logging.basicConfig(level=logging.DEBUG, format='%(levelname)s %(name)s: %(message)s') +from ptychodus.api.plugins import PluginRegistry +r = PluginRegistry.load_plugins() +" +``` + +Scan the output for: + +- **`ModuleNotFoundError`** on the plugin's own module name → the plugin's imports are broken. Look at the top of the plugin file for a missing optional dep, then either install it or wrap the import cleanly. +- **`ModuleNotFoundError`** on a *different* module the plugin imports → optional dependency missing; install it or move the import inside a `try/except`. +- **`Failed to register `** with `AttributeError` → the plugin file was imported but has no top-level `register_plugins(registry)` function. Confirm the function is at module level (not nested), spelled exactly `register_plugins`, and takes one argument. +- **No mention of the plugin at all** → the file isn't under `src/ptychodus/plugins/` or isn't a valid Python module (missing `.py`, in a subdirectory without `__init__.py`, or the name isn't a valid identifier). + +### 3. Confirm the module exists and defines the hook + +```sh +ls src/ptychodus/plugins/ | grep -i +grep -n "^def register_plugins" src/ptychodus/plugins/.py +``` + +If `grep` finds nothing, the plugin needs a `register_plugins(registry: PluginRegistry) -> None` at module level. + +### 4. Confirm it's registered + +If discovery succeeded (a `DEBUG` line said `Registered ptychodus.plugins.`), verify the plugin actually ended up in the right chooser: + +```sh +uv run python -c " +from ptychodus.api.plugins import PluginRegistry +r = PluginRegistry.load_plugins() +# Pick the chooser matching the plugin category: +for p in r.probe_file_readers: + print(p.simple_name, '—', p.display_name) +" +``` + +Categories on `PluginRegistry` (see [src/ptychodus/api/plugins.py](../../src/ptychodus/api/plugins.py)): +`bad_pixels_file_readers`, `diffraction_file_readers` / `_writers`, `probe_position_file_readers` / `_writers`, `fresnel_zone_plates`, `probe_file_readers` / `_writers`, `object_file_readers` / `_writers`, `product_file_readers` / `_writers`, `file_based_workflows`, `fluorescence_file_readers` / `_writers`, `upscaling_strategies`, `deconvolution_strategies`. + +### 5. Report findings + +Summarize: what the plugin should have registered, what actually happened, and the smallest-scope fix. Do not fix silently — hand the diagnosis back to the user with the suggested fix and let them confirm. + +## Common gotchas + +- Product readers not showing up as probe/position/object readers → the plugin registered on `product_file_readers` directly instead of using `registry.register_product_file_reader_with_adapters(...)`. The adapters are the reason product readers work universally. +- Plugin appears twice or with wrong casing → `simple_name` collision or missing `simple_name` (auto-derived from `display_name` may collide). Pass `simple_name` explicitly. +- Plugin loads but crashes at read time → not a load-failure; investigate the reader's `read()` method against a sample file, not with this skill. diff --git a/.claude/skills/pre-push/SKILL.md b/.claude/skills/pre-push/SKILL.md new file mode 100644 index 000000000..785053672 --- /dev/null +++ b/.claude/skills/pre-push/SKILL.md @@ -0,0 +1,48 @@ +--- +name: pre-push +description: Run the ptychodus CI gate locally before pushing — ruff check, ruff format --check, mypy, pytest, and the Markdown linter, matching what .github/workflows/python-package.yml runs on PR. Use when the user says "check before push", "run CI locally", "pre-push", or after a batch of code changes when they're about to open/update a PR. +--- + +# pre-push + +Runs the four commands `.github/workflows/python-package.yml` runs on every PR, in the same order, so a green run here means CI will pass — plus a Markdown lint pass that CI does not run but the project's documentation style depends on. + +## Steps + +Run these sequentially. Stop at the first failure and report clearly; do not proceed to the next step until the current one is clean. + +```sh +uv run ruff check . +uv run ruff format --check . +uv run mypy src/ptychodus scripts +uv run pytest +uv run pymarkdown scan $(git ls-files '*.md') +``` + +## Handling failures + +- **`ruff check` failure:** show the violations. If they are auto-fixable (`--fix` would resolve them), ask the user before running `uv run ruff check --fix .`. Never auto-fix `N` (naming) violations — those often involve deliberate `# noqa: N…` on physical-quantity names per CLAUDE.md. +- **`ruff format --check` failure:** offer to run `uv run ruff format .` and re-run the check. This is almost always safe to apply. +- **`mypy` failure:** show the errors with file:line. Do not attempt fixes here — hand back to the user; typing changes often need real thought. +- **`pytest` failure:** show the failing test name(s) and a compact traceback. Ask the user how to proceed before rerunning. +- **`pymarkdown` failure:** show the `file:line: MDxxx` findings and apply the rule each message names — the house Markdown style is codified in CLAUDE.md's Conventions section. Do not relax `[tool.pymarkdown]` in `pyproject.toml` to silence a finding; that config records deliberate carve-outs (`MD013` line length, `MD014` shell prompts) and widening it hides real drift. + +## Scope shortcuts + +If the user's diff only touches one subpackage, you can offer to run the module-scoped variant instead: + +```sh +uv run ruff check src/ptychodus/ +uv run ruff format --check src/ptychodus/ +uv run mypy src/ptychodus/ +uv run pytest tests/test_.py # if a matching test file exists +``` + +Use this only when the user asks for a fast local check on WIP; the full sweep above is what CI actually runs. + +## Notes + +- Ruff rules for this repo: `F, N, NPY`; single-quoted strings; 100-char lines; py311 target (see `pyproject.toml`). +- `mypy` targets `src/ptychodus` and the top-level `scripts/` tree. `src/ptychodus_store` is not currently in CI's mypy job — check `pyproject.toml`/`.github/workflows/python-package.yml` before assuming coverage. +- `pymarkdown` is a local-only gate; there is no Markdown job in CI. Scan via `git ls-files` so `.venv/`, `docs/build/`, and other untracked trees stay out of scope. +- Do not add `--no-verify` or skip hooks to work around a failure; investigate and fix instead. diff --git a/.claude/skills/pre-release/SKILL.md b/.claude/skills/pre-release/SKILL.md new file mode 100644 index 000000000..3f77a7a93 --- /dev/null +++ b/.claude/skills/pre-release/SKILL.md @@ -0,0 +1,288 @@ +--- +name: pre-release +description: Run a comprehensive pre-release verification of the ptychodus repository — API docs coverage, minimal docstring presence, reader-plugin docs, README/CLAUDE.md/pyproject.toml consistency, install-instruction freshness, Markdown hygiene, zero-warning Sphinx build, full CI gate, and entry-point smoke tests. Report every finding and offer per-issue fixes. Use when the user says "pre-release check", "release audit", "before we cut a release", or "verify the repo is release-ready". +--- + +# pre-release + +Comprehensive gate for cutting a release. Runs eight verification sections in order, reports pass/fail for each, and — for every failure — proposes a specific fix and asks the user before applying it. **Never fixes silently. Never commits.** + +## How to run the check + +Execute the sections in order. After each, capture pass/fail and any findings. At the end, print the summary table (last section). Only then walk the user through fixing failures one at a time. + +--- + +### Section 1 — API modules present in `docs/source/api.md` + +For every `src/ptychodus/api/*.py` file (excluding `_*.py` private modules), confirm `docs/source/api.md` contains a matching `.. automodule:: ptychodus.api.` directive. + +```sh +# Modules that should be documented: +ls src/ptychodus/api/*.py | xargs -n1 basename | sed 's/\.py$//' | grep -v '^_' | grep -v '^__' | sort + +# Modules currently documented: +grep -oE '\.\. automodule:: ptychodus\.api\.[a-z_]+' docs/source/api.md | sed 's|.*ptychodus\.api\.||' | sort +``` + +Compute the diff. `PASS` if empty. `FAIL` with the list of missing module names. + +**Suggested fix per missing module**: append this stanza to `docs/source/api.md`, matching the existing pattern (title-case `##` heading, then the autodoc directive as raw reStructuredText inside an `{eval-rst}` block): + +````markdown +## + +```{eval-rst} +.. automodule:: ptychodus.api.<stem> + :members: + :undoc-members: + :show-inheritance: +``` +```` + +Ask the user to confirm the heading text before writing. + +--- + +### Section 2 — Minimal docstring coverage in `src/ptychodus/api/` + +For each non-private module in `src/ptychodus/api/`, verify: + +- Module has a top-of-file docstring. +- Every top-level `class` has a docstring on its first statement. +- Every top-level `def` whose name does NOT start with `_` has a docstring on its first statement. + +Run this AST scan: + +```sh +uv run python -c " +import ast, pathlib +missing = [] +for p in sorted(pathlib.Path('src/ptychodus/api').glob('*.py')): + if p.name.startswith('_'): + continue + tree = ast.parse(p.read_text()) + if not (tree.body and isinstance(tree.body[0], ast.Expr) and isinstance(tree.body[0].value, ast.Constant) and isinstance(tree.body[0].value.value, str)): + missing.append(f'{p}:1: module docstring') + for node in tree.body: + if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith('_'): + continue + if not (node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant) and isinstance(node.body[0].value.value, str)): + kind = type(node).__name__.replace('Def','').lower() + missing.append(f'{p}:{node.lineno}: {kind} {node.name!r}') +for m in missing: + print(m) +print(f'--- {len(missing)} missing ---') +" +``` + +`PASS` if 0 missing. `FAIL` otherwise, with the list of `file:line` findings. + +**Suggested fix per finding**: propose a one-line docstring derived from the identifier name (e.g. for `class DiffractionPattern:` propose `"""A diffraction pattern captured by the detector."""`) and let the user accept, edit, or skip. Do not batch these — one at a time, so the user can improve the wording as you go. + +--- + +### Section 3 — Reader plugins represented in `docs/source/readers.md` + +`docs/source/readers.md` is a curated bullet list, not a 1:1 file mapping. Cross-check by keyword: for every distinct beamline/format hint in `src/ptychodus/plugins/`, ensure the readers doc mentions it. + +Keyword map (edit as new plugins land): + +| Plugin filename substring | Expected mention in readers.md | +| --- | --- | +| `aps02id_` | "2-ID-D Bionanoprobe" or "2-ID-D Microprobe" or "2-ID-E Microprobe" | +| `aps04id_polar_` | "4-ID" and "Polar" (Polarization Modulation Spectroscopy) | +| `aps09id_cssi_` | "9-ID" and "CSSI" | +| `aps12id_` | "12-ID" and "SAXS" | +| `aps19id_isn_` | "19-ID" and "ISN" | +| `aps31id_lynx_` | "31-ID" and "LYNX" | +| `aps33id_velociprobe` | "33-ID" and "Velociprobe" | +| `lcls_` | "LCLS" | +| `max_iv_nanomax_` | "MAX IV" or "NanoMAX" | +| `fold_slice_` | "fold_slice" | +| `cxi_` | "CXI" or "*.cxi" | +| `csv_` | "CSV" or "Comma-Separated" | +| `mda_` | "MDA" or "*.mda" | +| `npy_` | "NumPy" or "*.npy" | +| `delimited_position_` | "Space-Separated" or "*.txt" | + +```sh +# List all plugin file stems: +ls src/ptychodus/plugins/*.py | xargs -n1 basename | sed 's/\.py$//' + +# For each keyword above, verify readers.md contains the expected mention: +for kw in "2-ID-D Bionanoprobe" "4-ID" "9-ID" "12-ID" "19-ID" "31-ID" "33-ID" "LCLS" "NanoMAX" "fold_slice" "CXI" "CSV" "MDA" "NumPy" "Space-Separated"; do + grep -q "$kw" docs/source/readers.md || echo "MISSING: $kw" +done +``` + +Also flag *the reverse*: any bullet in `readers.md` whose beamline has no matching plugin file — that indicates a stale doc entry. + +`PASS` if every plugin has a doc mention and every doc mention has a plugin. `FAIL` with lists of orphans in either direction. + +**Suggested fix per finding**: propose a bullet insertion under the correct facility heading, matching the existing style (`- <Beamline name> (<abbreviation>)`). Ask the user to confirm the heading and wording — the doc uses friendly names, not filenames. + +--- + +### Section 4 — README, CLAUDE.md, pyproject.toml consistency + +**4a. Extras named in `README.md` all exist in `pyproject.toml`.** + +```sh +# Extract extras from README.md install commands. +# Anchor on `ptychodus[...]` so Markdown link labels ([Ptychodus], [uv]) don't match: +grep -oE 'ptychodus\[[a-z,]+\]' README.md | sed 's/.*\[//; s/\]//' | tr ',' '\n' | sort -u + +# Extras declared in pyproject.toml: +uv run python -c " +import tomllib +with open('pyproject.toml','rb') as f: + p = tomllib.load(f) +for k in sorted(p['project']['optional-dependencies']): + print(k) +" +``` + +`PASS` if the README extras are a subset of the pyproject extras. `FAIL` with the missing extras. + +**Suggested fix**: either add the extra to `pyproject.toml` (rare — usually intentional) or remove/rename the extra reference in `README.md`. + +**4b. Every `[project.scripts]` entry is documented or referenced.** + +```sh +uv run python -c " +import tomllib +with open('pyproject.toml','rb') as f: + p = tomllib.load(f) +for name in sorted(p['project'].get('scripts', {})): + print(name) +" +``` + +For each script name, confirm it appears in `CLAUDE.md` OR `docs/source/getting_started.md`. `PASS` if all are referenced. + +**Suggested fix**: add a `uv run <script>` example under the "Common Commands" section of CLAUDE.md (project convention — see the existing block). + +**4c. Python version claim matches pyproject.** + +- CLAUDE.md says "Python ≥3.11". +- Confirm `pyproject.toml` `requires-python` matches. `PASS`/`FAIL` accordingly. + +--- + +### Section 5 — Installation instructions are fresh + +**5a. Dockerfile variants referenced still exist.** Every Dockerfile name mentioned in `docs/source/getting_started.md` and `CLAUDE.md` should be a real file at the repo root: + +```sh +grep -hoE 'Dockerfile\.[a-z]+' docs/source/getting_started.md CLAUDE.md | sort -u | while read f; do + [ -f "$f" ] || echo "MISSING: $f" +done +``` + +`PASS` if empty. `FAIL` lists missing Dockerfiles. + +**5b. Extras named in `docs/source/getting_started.md` install commands exist in pyproject.** + +Same technique as 4a but against getting_started.md. + +**5c. The `uv sync` command in README.md uses currently-supported extras.** + +Parse the `uv sync --extra <x> --extra <y> ...` line in README.md and confirm each extra is in `pyproject.toml`. + +--- + +### Section 5.5 — Markdown hygiene + +Docs are MyST Markdown; no reStructuredText should reappear, and every tracked `.md` must satisfy the house style codified in `CLAUDE.md`. + +```sh +# No stray reStructuredText anywhere in the repo: +git ls-files '*.rst' + +# House style, mechanically enforced: +uv run pymarkdown scan $(git ls-files '*.md') + +# Shell fences use `sh`, not `bash`: +git ls-files '*.md' | xargs grep -l '^```bash' +``` + +`PASS` if the first and third commands print nothing and `pymarkdown` exits 0. `FAIL` otherwise, listing each offending file. + +**Suggested fix**: for `.rst` files, convert to MyST Markdown and update every reference. For linter findings, apply the rule the message names — do not widen the `[tool.pymarkdown]` config in `pyproject.toml` to silence a real violation. + +--- + +### Section 6 — Sphinx build with zero warnings + +```sh +uv sync --extra docs +make -C docs clean +make -C docs html SPHINXOPTS="-W --keep-going" +``` + +`PASS` if exit code 0. `FAIL` on any warning or error. + +**Suggested fix**: hand the first warning to the user. Common causes: undocumented cross-references, autodoc import failures (usually a missing extra — retry with `uv sync --extra docs --extra ptychi --extra globus --extra gui --extra store` if autodoc can't import a module), duplicate labels. Not auto-fixable. + +*Note*: `-W` is strict, and the prose sources are expected to build clean — treat any warning originating in `docs/source/*.md` as a regression introduced by the change under review. Warnings raised from `src/ptychodus/api/*.py` docstrings (reStructuredText syntax errors surfaced by autodoc) are a known pre-existing backlog; report the count and the first offender, but do not fold fixing them into the release gate. + +--- + +### Section 7 — Full CI gate (chain `pre-push`) + +Invoke the `pre-push` skill. All four steps (ruff check, ruff format --check, mypy, pytest) must be green. + +`PASS` if `pre-push` finishes clean. `FAIL` on any failing step; hand the failure back to the user without attempting fixes here (the `pre-push` skill knows how to offer format fixes; deeper fixes belong outside the release gate). + +--- + +### Section 8 — Entry-point smoke test + +For each name in `[project.scripts]`, run `uv run <script> --version`; if the command doesn't implement `--version`, fall back to `--help`. Exit 0 counts as PASS; non-zero or import error counts as FAIL. + +```sh +uv run python -c " +import tomllib +with open('pyproject.toml','rb') as f: + print('\n'.join(tomllib.load(f)['project']['scripts'])) +" | while read script; do + if uv run "$script" --version >/dev/null 2>&1; then + echo "PASS: $script" + elif uv run "$script" --help >/dev/null 2>&1; then + echo "PASS: $script (--help)" + else + echo "FAIL: $script" + fi +done +``` + +`PASS` if every script exits 0. `FAIL` lists the broken entry points — usually caused by import-time errors introduced by an unrelated change. + +--- + +## Final report + +Print a summary in this exact format after all sections have run: + +```text +1. API modules in docs PASS (<covered>/<total>) +2. Docstrings in api PASS|FAIL (<n> missing) +3. Reader plugins in docs PASS|FAIL (<n> orphans) +4. README/CLAUDE/pyproject PASS|FAIL (<one-line summary>) +5. Install instructions PASS|FAIL (<one-line summary>) +5.5 Markdown hygiene PASS|FAIL (<n> rst files, <n> lint findings) +6. Sphinx build (zero warns) PASS|FAIL (first warning if any) +7. Full CI gate PASS|FAIL +8. Entry-point smoke tests PASS|FAIL (<pass>/<total>) +``` + +Then, and only then, walk through each `FAIL` with the user: state the finding, propose the fix, ask "apply it?" per item, and Edit/Write only after they confirm. + +## Do not + +- Do not `git commit` — release commits are the user's call. +- Do not bump the version — that's a deliberate release step, not drift. +- Do not fix warnings in code you didn't write for the release — surface them and let the user decide scope. +- Do not run this skill against a dirty working tree unless the user explicitly wants to include their WIP in the audit. diff --git a/.claude/skills/store-dev/SKILL.md b/.claude/skills/store-dev/SKILL.md new file mode 100644 index 000000000..642a4cd10 --- /dev/null +++ b/.claude/skills/store-dev/SKILL.md @@ -0,0 +1,60 @@ +--- +name: store-dev +description: Bootstrap and launch the ptychodus-store REST + MCP + browser-UI service locally against a storage root. Use when the user says "start the store", "run ptychodus-store", "serve the store", or wants to iterate on ptychodus_store code (routers, MCP tools, or the TypeScript UI in src/ptychodus_store/ui/). +--- + +# store-dev + +Launches `ptychodus-store serve` with the standard local setup. Docs: `src/ptychodus_store/README.md`. + +## Prerequisites + +Ask the user for the storage root path if not already provided. Options, in priority order: + +1. Path passed in the invocation. +2. `PTYCHODUS_STORE_STORAGE_ROOT` env var if already set in the shell. +3. Prompt: "Which storage root should the store serve from?" — do not guess a path. + +## Steps + +1. **Sync the store extra** (skip if already synced this session): + + ```sh + uv sync --extra store + ``` + +2. **Rebuild the SQLite index from disk** (optional — offer this only if the user says the on-disk artifacts changed outside the running service, e.g. after `git pull` or manual file moves): + + ```sh + uv run ptychodus-store rebuild-index + ``` + +3. **Launch the service** — this is a long-running process; run it in the background so you can continue helping the user: + + ```sh + PTYCHODUS_STORE_STORAGE_ROOT=<path> uv run ptychodus-store serve + ``` + + Use `run_in_background: true` on the Bash call. Report the process id and log file path. + +4. **Sanity-check** with a health probe once the server is up: + + ```sh + curl -sf http://localhost:8000/api/v1/health + ``` + + (Port defaults per `src/ptychodus_store/config.py`; if the user configured `PTYCHODUS_STORE_*` env vars for host/port, use those.) + +## Frontend iteration + +The TypeScript UI lives in `src/ptychodus_store/ui/src/` and compiles to `src/ptychodus_store/ui/dist/` via plain `tsc` (no bundler). If the user is editing UI code: + +```sh +cd src/ptychodus_store/ui && tsc --watch +``` + +`dist/` is a build artifact — never commit it. + +## Stopping + +The service is a foreground blocker unless backgrounded. Use `TaskStop` on the background task id when the user is done, or tell them to `Ctrl+C` if they launched it interactively. diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 3d1607821..731f90035 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -86,4 +86,4 @@ jobs: - name: Install ptychodus + mypy run: pip install . mypy pyqt5-stubs types-psutil types-pyyaml types-requests - name: mypy - run: mypy src/ptychodus + run: mypy src/ptychodus scripts diff --git a/CLAUDE.md b/CLAUDE.md index c79583a72..4061c6b84 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,114 +1,88 @@ # CLAUDE.md -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - ## Project -Ptychodus is a ptychography data-analysis application that ingests instrument data, prepares it for processing, and dispatches it through several reconstruction libraries (PtyChi, PtychoPINN, PtychoPINN-Torch). It runs interactively as a PyQt5 GUI, headless via batch CLI, and as a streaming processor inside beamline pipelines (pvapy area-detector). Python ≥3.11. +Ptychodus is a ptychography data-analysis application that ingests instrument data, prepares it for processing, and dispatches it through reconstruction libraries (PtyChi, PtychoPINN, PtychoPINN-Torch). It runs as a PyQt5 GUI, a headless batch CLI, and a streaming processor inside beamline pipelines (pvapy area-detector). Python ≥3.11, `uv` preferred. ## Common Commands -Project uses `uv` (preferred) and a developer install with extras. - ```sh -# Dev install (preferred). Other available extras: docs, ptychopinn -uv sync --extra globus --extra gui --extra ptychi - -# Launch GUI -uv run ptychodus - -# Headless batch -uv run ptychodus -b reconstruct -i <input_dir> -o <output_dir> -uv run ptychodus -b train -i <input_dir> -o <output_dir> -# Batch mode reads <input_dir>/settings.ini, <input_dir>/diffraction.h5, <input_dir>/product-in.h5 -# (see StandardFileLayout in src/ptychodus/api/io.py) - -# Beamline data-prep CLI -uv run ptychodus-bdp --product-name <name> --diffraction-input <h5> \ - --probe-position-input <csv> --output-directory <dir> \ - --settings <ini> - -# Other entry points (see [project.scripts] in pyproject.toml) -uv run convert-to-ptychodus -uv run ptychodus-system-check -uv run ptychodus-iri-tokens # Genesis/IRI auth -uv run ptychodus-transfer-tokens # AmSC data-transfer auth -uv run ptychodus-ptychopinn-tf-test # PtychoPINN TensorFlow smoke check +uv sync --extra globus --extra gui --extra ptychi # dev install (extras: see pyproject.toml) +uv run ptychodus # GUI +uv run ptychodus -b reconstruct -i <in> -o <out> # batch — reads StandardFileLayout in <in>/ +uv run pytest # tests +uv run ruff check . && uv run ruff format --check . # lint + format +uv run mypy src/ptychodus scripts # types ``` -Tests, lint, types: - -```sh -uv run pytest # full suite (tests/) -uv run pytest tests/test_io.py # one file -uv run pytest tests/test_io.py::test_name # one test -uv run ruff check . # lint (rules: F, N, NPY) -uv run ruff format . # format (single quotes, line-length 100) -uv run mypy src/ptychodus # type check (py 3.11) -``` +Container builds: `podman build -f Dockerfile.{cuda,cpu,rocm,xpu} .`. Docs: `make -C docs html`. -CI (`.github/workflows/python-package.yml`) runs four jobs on push/PR to `main`: a `pip install` + `ptychodus --version` smoke test (Py 3.11/3.12/3.13), `pytest tests/`, `ruff check` + `ruff format --check`, and `mypy src/ptychodus`. Run the local equivalents before pushing — CI will block the PR otherwise. +Other entry points (`convert-to-ptychodus`, `ptychodus-bdp`, `ptychodus-store`, `ptychodus-system-check`) are listed in [pyproject.toml](pyproject.toml) `[project.scripts]`; their modules live in [src/ptychodus/cli/](src/ptychodus/cli/). Unpackaged operator tooling — the podman wrapper and the per-facility HPC token/submit helpers — lives in the top-level [scripts/](scripts/) and is run from a checkout, e.g. `python scripts/genesis/ptychodus_iri_tokens.py`. The store service has its own docs: [src/ptychodus_store/README.md](src/ptychodus_store/README.md). -Container & docs: +### Testing -```sh -podman build -f Dockerfile.cuda -t ptychodus:cuda13.0 . # also: Dockerfile.cpu / .rocm / .xpu -docker build -f Dockerfile.cuda -t ptychodus:cuda13.0 . -make -C docs html # Sphinx docs into docs/build/ -``` +Tests cover `api/`, `model/`, and `view/widgets/` only — **do not add tests for the rest of `view/` or for `controller/`**, including controller helper functions. Widget tests depend on PyQt5, which is an optional extra (`--extra gui`); the whole `tests/view/` subtree is dropped in [tests/conftest.py](tests/conftest.py) via `collect_ignore` when `find_spec('PyQt5')` returns `None`, mirroring the `ptychodus_store` gate. `pytest tests/` must pass on a bare `pip install .` (that is what CI runs), so anything reachable from `ptychodus.model` at import time must be free of optional dependencies: gate optional backends with the `find_spec` probe + deferred factory import used in each `model/*/core.py`. Guard a test that needs an optional backend with `pytest.importorskip('ptychi')`; drop a whole optional-extra directory with `collect_ignore` in [tests/conftest.py](tests/conftest.py) — `importorskip` in a conftest is reported as an error, not a skip. -## High-Level Architecture +## Architecture ### Three-layer separation: api / model / view+controller -- **`src/ptychodus/api/`** — pure-Python domain layer with **no** GUI or model dependencies. Defines core data structures (`Product`, `ProbeSequence`, `Object`, `ProbePositionSequence`, `DiffractionDataset`, `AssembledDiffractionData`), abstract interfaces (`DiffractionFileReader/Writer`, `ProductFileReader/Writer`, `Reconstructor`, `TrainableReconstructor`, `WorkflowAPI`), and infrastructure (`Observable`/`Observer`, `Parameter`/`ParameterGroup`, `SettingsRegistry`, `PluginRegistry`/`PluginChooser`). All other layers depend on this; this layer depends on nothing else in ptychodus. -- **`src/ptychodus/model/`** — application logic. Each subpackage (`diffraction/`, `product/`, `processing/`, `reconstructor/`, `analysis/`, `fluorescence/`, `globus/`, `genesis/`, `automation/`, `agent/`, `visualization/`, `ptychi/`, `ptychopinn/`, `ptychopinn_torch/`) exposes a `*Core` class that owns its settings, repositories, and APIs. `model/core.py::ModelCore` is the composition root: it constructs every `*Core` in dependency order and wires them together. Used both by the GUI and by `__main__.py` batch mode. -- **`src/ptychodus/view/`** (PyQt5 widgets, no logic) and **`src/ptychodus/controller/`** (mediates between widgets and model). Mirror the model package layout. `view/core.py::ViewCore` and `controller/core.py::ControllerCore` are the composition roots; the navigation toolbar order in `ViewCore` is the source of truth for the left/right stacked-panel indexes. +- **`src/ptychodus/api/`** — pure-Python domain layer with **no** GUI or model dependencies. Core data structures (`Product`, `ProbeSequence`, `Object`, `ProbePositionSequence`, `DiffractionDataset`, `AssembledDiffractionData`), abstract interfaces (`DiffractionFileReader/Writer`, `ProductFileReader/Writer`, `Reconstructor`, `TrainableReconstructor`, `WorkflowAPI`), and infrastructure (`Observable`/`Observer`, `Parameter`/`ParameterGroup`, `SettingsRegistry`, `PluginRegistry`/`PluginChooser`). This layer depends on nothing else in ptychodus. +- **`src/ptychodus/model/`** — application logic. Each subpackage (`diffraction/`, `product/`, `processing/`, `reconstructor/`, `fluorescence/`, `analysis/`, `globus/`, `genesis/`, `ptychi/`, `ptychopinn/`, and others) exposes a `*Core` class that owns its settings, repositories, and APIs. `model/core.py::ModelCore` is the composition root that constructs every `*Core` in dependency order. Used by the GUI *and* by `__main__.py` batch mode. +- **`src/ptychodus/view/`** (PyQt5 widgets, no logic) and **`src/ptychodus/controller/`** (mediates widgets↔model). Mirror the model layout. `view/core.py::ViewCore` and `controller/core.py::ControllerCore` are composition roots; the navigation toolbar order in `ViewCore` is the source of truth for left/right stacked-panel indexes. -The GUI is optional: `__main__.py` falls back to headless mode if PyQt5 is missing. `ptychodus_stream_processor.py` (the `PtychodusAdImageProcessor`) is the third entry mode and is only imported when `pvapy` is available. +The GUI is optional: `__main__.py` falls back to headless if PyQt5 is missing. `ptychodus_stream_processor.py` (the `PtychodusAdImageProcessor`) is a third entry mode, imported only when `pvapy` is available. ### Index-based pattern/position association -Diffraction patterns and probe positions are paired by **integer scan index**, never by array order. This is what makes streaming and mismatched-rate ingest robust: patterns and positions can arrive from different files or different PV channels, with dropped or extra samples on either side, and the matcher still pairs them correctly. +Diffraction patterns and probe positions are paired by **integer scan index**, never by array order. This is what makes streaming and mismatched-rate ingest robust. -- Producers: `DiffractionArray.get_indexes()` on the pattern side (`api/diffraction.py`); `ProbePosition.index` on the position side (`api/probe_positions.py`). -- Matcher: `AssembledDiffractionData.prepare_reconstruct_input` in `api/reconstructor.py` treats pattern indexes as authoritative — duplicate position indexes are averaged into anchors, pattern indexes inside the anchor range with no exact position are linearly interpolated, and pattern indexes outside the anchor range are dropped (no extrapolation). The `Product` is rebuilt from the resulting per-pattern positions. +- Producers: `DiffractionArray.get_indexes()` (`api/diffraction.py`); `ProbePosition.index` (`api/probe_positions.py`). +- Matcher: `AssembledDiffractionData.prepare_reconstruct_input` in `api/reconstructor.py` treats pattern indexes as authoritative — duplicate position indexes are averaged into anchors, pattern indexes inside the anchor range without an exact position are linearly interpolated, and pattern indexes outside the anchor range are dropped (no extrapolation). - Round-trip: HDF5 and NPZ product writers persist position indexes via `ProductFileKeys.PROBE_POSITION_INDEXES`. ### Plugin system -File-format support is dynamic. `api/plugins.py::PluginRegistry.load_plugins()` walks `ptychodus.plugins.*` with `pkgutil.iter_modules` and calls each module's `register_plugins(registry)` function. Module-load failures are logged and skipped — this is how optional-dependency plugins (e.g., LCLS, NSLS-II) silently disable themselves. - -To add a new file format: create a module under `src/ptychodus/plugins/`, implement the relevant `*FileReader`/`*FileWriter` from `ptychodus.api`, and define `register_plugins(registry)` calling the appropriate `registry.<category>.register_plugin(...)`. For product readers, prefer `register_product_file_reader_with_adapters` so the reader is also reachable as a probe/probe-position/object reader. +`api/plugins.py::PluginRegistry.load_plugins()` walks `ptychodus.plugins.*` with `pkgutil.iter_modules` and calls each module's `register_plugins(registry)`. Module-load failures are logged and skipped — this is how optional-dependency plugins (LCLS, NSLS-II, …) silently disable themselves. Never make plugin loading fatal. ### Settings & observers -Settings flow through `api/parametric.py::Parameter[T]` and `ParameterGroup`. Each `*Core` calls `settings_registry.create_group(name)` and creates typed parameters on it. `SettingsRegistry` serializes the full tree as INI. Components react to settings changes through the `Observer`/`Observable` pattern in `api/observer.py`; `PluginChooser.synchronize_with_parameter` is the typical bridge for "settings string → currently selected plugin." When adding cross-component reactivity, hook observers rather than poll. +Settings flow through `api/parametric.py::Parameter[T]` and `ParameterGroup`. Each `*Core` calls `settings_registry.create_group(name)` and creates typed parameters. `SettingsRegistry` serializes the tree as INI. React to settings changes via `Observer`/`Observable` (`api/observer.py`); `PluginChooser.synchronize_with_parameter` is the typical bridge for "settings string → currently selected plugin." Hook observers rather than poll. -### Reconstructor libraries +### Standard HDF5 layout (external interface — change carefully) -Each reconstructor backend (`model/ptychi/`, `model/ptychopinn/`, `model/ptychopinn_torch/`) exposes a `*ReconstructorLibrary` class. They are constructed in `ModelCore` and passed as a list to `ProcessingCore`. Backends that aren't installed should fail cleanly at import inside their own `__init__.py` — keep the surface a stable `*ReconstructorLibrary` regardless. Reconstructors implement `Reconstructor`/`TrainableReconstructor` from `api/reconstructor.py` and yield `ReconstructOutput` per iteration so the GUI can stream progress. +`api/io.py::StandardFileLayout` is the canonical contract for batch mode and remote workflows: `diffraction.h5`, `product-in.h5`, `product-out.h5`, `fluorescence-in.h5`, `fluorescence-out.h5`, `settings.ini`. Key names live in `DiffractionFileKeys` / `ProductFileKeys` enums — external consumers depend on them. -### Standard HDF5 layout +### Reconstructor libraries -`api/io.py::StandardFileLayout` is the canonical contract for batch mode and remote workflows: `diffraction.h5`, `product-in.h5`, `product-out.h5`, `fluorescence-in.h5`, `fluorescence-out.h5`, `settings.ini`. `load_diffraction_data` / `save_diffraction_data` and `load_product` / `save_product` round-trip these files; their key names live in `DiffractionFileKeys` / `ProductFileKeys` enums — change those carefully, they are an external interface. +Each backend (`model/ptychi/`, `model/ptychopinn/`, `model/ptychopinn_torch/`) exposes a `*ReconstructorLibrary`. Backends whose deps aren't installed must fail cleanly *at import inside their own `__init__.py`* — the `*ReconstructorLibrary` surface stays stable so `ModelCore` can always construct and pass it to `ProcessingCore`. Reconstructors yield `ReconstructOutput` per iteration so the GUI can stream progress. ### Remote compute -Two providers, both gated by optional dependencies and constructed by `ModelCore` even when disabled: +Both providers are optional-dep gated but constructed by `ModelCore` even when disabled: + +- `model/globus/` — Globus Compute (original APS workflow). +- `model/genesis/` — IRI/AmSC HPC via `facility_adapters.py`, with per-facility scripts under `scripts/genesis/{alcf,nersc,olcf}/`. + +`WorkflowAPI` (`api/workflow.py`, `model/workflow.py::ConcreteWorkflowAPI`) is the unified façade GUI, batch, and remote drive. -- `model/globus/` — Globus Compute submission, used by the original APS workflow. -- `model/genesis/` — IRI/AmSC HPC submission with facility adapters in `model/genesis/facility_adapters.py` and per-facility scripts under `src/ptychodus/scripts/genesis/{alcf,nersc,olcf}/`. +### ptychodus_store service -`WorkflowAPI` (see `api/workflow.py` and `model/workflow.py::ConcreteWorkflowAPI`) is the unified façade that GUI, batch mode, and remote workflows all drive. +`src/ptychodus_store/` is a **separate package** adjacent to `ptychodus`, not a subpackage. It reads-only from `ptychodus.api` and **must not import `ptychodus.model` or `ptychodus.view`**. Surface: FastAPI REST under `/api/v1/*`, MCP server at `/mcp`, minimal TypeScript UI at `/ui/`, SQLite metadata cache reconciled by a watchdog observer. Composition root: `create_app()` in `app.py`. Deployment details in [src/ptychodus_store/README.md](src/ptychodus_store/README.md). ## Conventions -- Ruff is configured for **single-quoted** strings, 100-char lines, py311 target. The selected lint rules (F, N, NPY) flag pyflakes errors, PEP-8 naming, and NumPy-specific issues. NumPy/Qt-style names (e.g., `probe_energy_eV`, `set_value_from_string`) are accepted via `# noqa: N802/N806/N815` — preserve the unit suffixes on physical quantities; reviewers expect them. -- Type hints are mandatory; `pyproject.toml` lists modules whose missing stubs are intentionally ignored. Keep new code typed and avoid widening that ignore list. -- `model/core.py::ModelCore.is_developer_mode_enabled` is `True` whenever the effective log level ≤ DEBUG (`--log-level 10`). Some controllers gate features (Agent panel, probe-position analysis) behind it. +- Ruff: **single-quoted** strings, 100-char lines, py311 target. Enabled rules: F, N, NPY. +- Preserve unit suffixes on physical quantities (`probe_energy_eV`, `set_value_from_string`) — silence naming lints with `# noqa: N802/N806/N815`, don't rename. +- Type hints are mandatory. `pyproject.toml` lists modules whose missing stubs are intentionally ignored — keep new code typed and don't widen that list. +- `ModelCore.is_developer_mode_enabled` is `True` when the effective log level ≤ DEBUG (`--log-level 10`); some controllers gate features (Agent panel, probe-position analysis) behind it. +- HTTP-client code uses `httpx` throughout (sync `httpx.Client` for repeated calls, module-level `httpx.get/post` for one-shots). `requests` is not a dependency. +- Prefer affirmative conditionals when both branches do meaningful work — `if x.is_file(): ... else: ...`, not the negated form. Guard clauses that early-return are still idiomatic. +- The `ptychodus_store/ui/` frontend is TypeScript compiled to native ES modules with plain `tsc` — no bundler, no framework, no runtime npm deps. Wheel builds run `tsc` via a `build_py` cmdclass in `setup.py`; interactive dev needs `tsc` on PATH. +- **`cli/` versus `scripts/`.** `src/ptychodus/cli/` holds the modules behind `[project.scripts]` — it ships in the wheel, and `cli/__init__.py` carries the shared argparse helpers (`DirectoryType`, `verify_all_arguments_parsed`). The top-level `scripts/` is **not** packaged: it is checkout-run operator tooling (podman wrapper, per-facility HPC token/submit helpers, demos) invoked as `python scripts/…`. A new console command goes in `cli/` with an entry point; a new one-off or facility script goes in `scripts/` with none. Both trees are type-checked (`mypy src/ptychodus scripts`) and linted. +- **All documentation is Markdown.** `docs/source/*.md` is MyST, parsed by `myst_parser`; there is no reStructuredText left in the repo and no new `.rst` should be added. Sphinx constructs use MyST directive fences — ` ```{note} `, ` ```{toctree} `, ` ```{image} `, ` ```{literalinclude} ` — with `:option: value` lines directly under the opening fence and a blank line before directive content. Autodoc stanzas stay as raw reStructuredText inside ` ```{eval-rst} ` blocks. Roles use MyST syntax: `` {py:class}`…` ``, `` {py:func}`…` ``, `` {ref}`…` ``, `` {kbd}`…` ``. Enabled MyST extensions (`colon_fence`, `deflist`, `fieldlist`) are declared in [docs/source/conf.py](docs/source/conf.py). +- **Markdown style**, enforced by `pymarkdownlnt` (`uv run pymarkdown scan $(git ls-files '*.md')`, config in `pyproject.toml` `[tool.pymarkdown]`): ATX headings only, one H1 per file, `###` max outside `docs/source/` (MyST pages may use `####`); backtick fences never tildes, and the shell tag is `sh` not `bash`; `-` bullets, ordered lists with real incrementing numbers; inline links with repo-relative paths, bare URLs as `<https://…>` autolinks; pipe tables with leading and trailing pipes and an unpadded `| --- |` separator; **no hard wrapping** — one physical line per paragraph, however long; `**bold**` for identifiers, paths, and warnings, backticks for code tokens, em dash `—` as the aside separator; no YAML front matter except `.claude/skills/*/SKILL.md`; LF endings, no trailing whitespace, one terminating newline, a blank line before every heading and fence. Fences and sub-lists nested inside a list item indent to the parent's content column — 3 spaces under an ordered marker, 2 spaces under a dash marker — with a blank line before the fence. -## Repository Notes +## Repository -- The git CI workflow targets `main`; local development branches such as `amsc` are active — confirm the intended target before opening PRs. -- Sample `.h5`/`.npy` data and the `dist/` build output in the working tree are typically **untracked** local artifacts — do not stage them in commits unless explicitly asked. +CI targets `main`; local dev branches (e.g. `amsc`, `webservice`) are active — confirm the intended target before opening PRs. diff --git a/README.md b/README.md new file mode 100644 index 000000000..570a3eae1 --- /dev/null +++ b/README.md @@ -0,0 +1,33 @@ +# Ptychodus + +[Ptychodus](https://github.com/AdvancedPhotonSource/ptychodus) is a ptychography data analysis application that extracts, loads, and transforms instrument data for processing. It integrates several reconstruction libraries for phase retrieval. Ptychodus can be used interactively or integrated into beamline data pipelines. + +## Standard Installation + +To install ptychodus from PyPI with the most common optional dependencies: + +```sh +$ python -m pip install ptychodus[globus,gui,ptychi] +``` + +Instructions for installing in containers, uv, and from conda-forge are provided in the `docs` directory. + +## Developer Installation + +For a developer installation: + +```sh +$ git clone https://github.com/AdvancedPhotonSource/ptychodus.git +$ cd ptychodus +$ uv sync --extra globus --extra gui --extra ptychi +``` + +Launch `ptychodus`: + +```sh +$ uv run ptychodus +``` + +## Reporting Bugs + +Open a bug at <https://github.com/AdvancedPhotonSource/ptychodus/issues>. diff --git a/README.rst b/README.rst deleted file mode 100644 index c7b3f86e1..000000000 --- a/README.rst +++ /dev/null @@ -1,44 +0,0 @@ -Ptychodus -========= - -`Ptychodus <https://github.com/AdvancedPhotonSource/ptychodus>`_ is a -ptychography data analysis application that extracts, loads, and transforms -instrument data for processing. It integrates several reconstruction libraries -for phase retrieval. Ptychodus can be used interactively or integrated into -beamline data pipelines. - -Standard Installation ---------------------- - -To install ptychodus from PyPI with the most common optional dependencies: - -.. code-block:: shell - - $ python -m pip install ptychodus[globus,gui,ptychi] - -Instructions for installing in containers, uv, and from conda-forge are provided in -the ``docs`` directory. - - -Developer Installation ----------------------- - -- For a developer installation: - -.. code-block:: shell - - $ git clone https://github.com/AdvancedPhotonSource/ptychodus.git - $ cd ptychodus - $ uv sync --extra globus --extra gui --extra ptychi - -- Launch `ptychodus`: - -.. code-block:: shell - - $ uv run ptychodus - - -Reporting Bugs --------------- - -Open a bug at https://github.com/AdvancedPhotonSource/ptychodus/issues. diff --git a/docs/source/api.rst b/docs/source/api.md similarity index 71% rename from docs/source/api.rst rename to docs/source/api.md index 825946a35..53bf53f1b 100644 --- a/docs/source/api.rst +++ b/docs/source/api.md @@ -1,223 +1,253 @@ -API Reference -============= +# API Reference -.. toctree:: - :maxdepth: 2 - :caption: Contents - - -Affine ------- +## Affine +```{eval-rst} .. automodule:: ptychodus.api.affine :members: :undoc-members: :show-inheritance: +``` -Common ------- +## Common +```{eval-rst} .. automodule:: ptychodus.api.common :members: :undoc-members: :show-inheritance: +``` -Diffraction ------------ +## Diffraction +```{eval-rst} .. automodule:: ptychodus.api.diffraction :members: :undoc-members: :show-inheritance: +``` + +## Diffraction Preprocessing + +```{eval-rst} +.. automodule:: ptychodus.api.diffraction_prep + :members: + :undoc-members: + :show-inheritance: +``` -Diffraction Generators ----------------------- +## Diffraction Generators +```{eval-rst} .. automodule:: ptychodus.api.diffraction_gen :members: :undoc-members: :show-inheritance: +``` -Fluorescence ------------- +## Fluorescence +```{eval-rst} .. automodule:: ptychodus.api.fluorescence :members: :undoc-members: :show-inheritance: +``` -Geometry --------- +## Geometry +```{eval-rst} .. automodule:: ptychodus.api.geometry :members: :undoc-members: :show-inheritance: +``` -Illumination ------------- +## Illumination +```{eval-rst} .. automodule:: ptychodus.api.illumination :members: :undoc-members: :show-inheritance: +``` -Input/Output (I/O) ------------------- +## Input/Output (I/O) +```{eval-rst} .. automodule:: ptychodus.api.io :members: :undoc-members: :show-inheritance: +``` -Interpolation -------------- +## Interpolation +```{eval-rst} .. automodule:: ptychodus.api.interpolate :members: :undoc-members: :show-inheritance: +``` -Metrics -------- +## Metrics +```{eval-rst} .. automodule:: ptychodus.api.metrics :members: :undoc-members: :show-inheritance: +``` -Object ------- +## Object +```{eval-rst} .. automodule:: ptychodus.api.object :members: :undoc-members: :show-inheritance: +``` -Object Generators ------------------ +## Object Generators +```{eval-rst} .. automodule:: ptychodus.api.object_gen :members: :undoc-members: :show-inheritance: +``` -Observer --------- +## Observer +```{eval-rst} .. automodule:: ptychodus.api.observer :members: :undoc-members: :show-inheritance: +``` -Parametric ----------- +## Parametric +```{eval-rst} .. automodule:: ptychodus.api.parametric :members: :undoc-members: :show-inheritance: +``` -Plugins -------- +## Plugins +```{eval-rst} .. automodule:: ptychodus.api.plugins :members: :undoc-members: :show-inheritance: +``` -Probe ------ +## Probe +```{eval-rst} .. automodule:: ptychodus.api.probe :members: :undoc-members: :show-inheritance: +``` -Probe Generators ----------------- +## Probe Generators +```{eval-rst} .. automodule:: ptychodus.api.probe_gen :members: :undoc-members: :show-inheritance: +``` -Probe Positions ---------------- +## Probe Positions +```{eval-rst} .. automodule:: ptychodus.api.probe_positions :members: :undoc-members: :show-inheritance: +``` -Probe Position Generators -------------------------- +## Probe Position Generators +```{eval-rst} .. automodule:: ptychodus.api.probe_positions_gen :members: :undoc-members: :show-inheritance: +``` -Product -------- +## Product +```{eval-rst} .. automodule:: ptychodus.api.product :members: :undoc-members: :show-inheritance: +``` -Propagator ----------- +## Propagator +```{eval-rst} .. automodule:: ptychodus.api.propagator :members: :undoc-members: :show-inheritance: +``` -Reconstructor -------------- +## Reconstructor +```{eval-rst} .. automodule:: ptychodus.api.reconstructor :members: :undoc-members: :show-inheritance: +``` -Settings --------- +## Settings +```{eval-rst} .. automodule:: ptychodus.api.settings :members: :undoc-members: :show-inheritance: +``` -Tree ----- +## Tree +```{eval-rst} .. automodule:: ptychodus.api.tree :members: :undoc-members: :show-inheritance: +``` -Visualization -------------- +## Visualization +```{eval-rst} .. automodule:: ptychodus.api.visualization :members: :undoc-members: :show-inheritance: +``` -Workflow --------- +## Workflow +```{eval-rst} .. automodule:: ptychodus.api.workflow :members: :undoc-members: :show-inheritance: +``` -XMCD ----- +## XMCD +```{eval-rst} .. automodule:: ptychodus.api.xmcd :members: :undoc-members: :show-inheritance: +``` diff --git a/docs/source/conf.py b/docs/source/conf.py index ce455798b..38e74774c 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -14,6 +14,7 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration extensions = [ + 'myst_parser', 'sphinx.ext.autodoc', 'sphinx_copybutton', 'sphinx.ext.coverage', @@ -28,11 +29,21 @@ 'scipy': ('https://docs.scipy.org/doc/scipy/', None), } -templates_path = ['_templates'] +source_suffix = { + '.md': 'markdown', +} + +myst_enable_extensions = [ + 'colon_fence', # ::: admonitions that can contain ``` fences + 'deflist', # parameter reference lists in initial_guesses.md + 'fieldlist', # :field: value metadata blocks +] + +myst_heading_anchors = 3 + exclude_patterns = [] # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output html_theme = 'sphinx_rtd_theme' -html_static_path = ['_static'] diff --git a/docs/source/getting_started.md b/docs/source/getting_started.md new file mode 100644 index 000000000..0da2cb9b4 --- /dev/null +++ b/docs/source/getting_started.md @@ -0,0 +1,248 @@ +# Installation Instructions + +## Python Package Index (PyPI) + +To install ptychodus with the most common optional dependencies: + +```sh +$ python -m pip install ptychodus[globus,gui,ptychi] +``` + +## uv + +[uv](https://docs.astral.sh/uv/) is a fast Python package and project manager. + +1. Install [uv](https://docs.astral.sh/uv/getting-started/installation/). + +2. Install ptychodus with the most common optional dependencies: + + ```sh + $ uv tool install ptychodus[globus,gui,ptychi] + ``` + +3. Launch ptychodus: + + ```sh + $ ptychodus + ``` + +4. To upgrade ptychodus, use uv tool upgrade: + + ```sh + $ uv tool upgrade ptychodus[globus,gui,ptychi] + ``` + +## Conda-Forge + +1. Install [miniforge](https://github.com/conda-forge/miniforge). + +2. Create the `ptychodus` environment + + - To install `ptychodus` with the GUI and all optional packages: + + ```sh + $ conda create -n ptychodus ptychodus-all + ``` + + - To install `ptychodus` with the GUI and no optional packages: + + ```sh + $ conda create -n ptychodus ptychodus + ``` + + - To install `ptychodus` without the GUI or optional packages: + + ```sh + $ conda create -n ptychodus ptychodus-core + ``` + +3. Activate the `ptychodus` environment + + ```sh + $ conda activate ptychodus + $ ptychodus + ``` + +## Container image variants + +The repository ships one Dockerfile per accelerator family. Pick the variant that matches your hardware and select an explicit file with `-f`: + +| Dockerfile | Use it for | +| --- | --- | +| `Dockerfile.cpu` | CPU-only hosts (no GPU; ptychi runs on CPU torch) | +| `Dockerfile.cuda` | NVIDIA GPUs (e.g. ALCF Polaris, NERSC Perlmutter); CUDA minor version is a build ARG | +| `Dockerfile.rocm` | AMD GPUs (e.g. OLCF Frontier); ROCm is a build ARG | +| `Dockerfile.xpu` | Intel XPU (e.g. ALCF Aurora); base tag is a build ARG | + +The GPU files default to recent versions and expose `--build-arg` knobs to switch: + +- `Dockerfile.cuda`: `CUDA_VERSION` (default `13.0`), `PYTORCH_VERSION`, `CUDNN_VERSION`. The base image is `pytorch/pytorch:${PYTORCH_VERSION}-cuda${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel`; override any args if a given combination isn't published upstream. +- `Dockerfile.rocm`: `ROCM_VERSION` (default `7.2.4`), `UBUNTU_VERSION`, `PYTHON_VERSION`, `PYTORCH_VERSION`. Base image is `rocm/pytorch:rocm${ROCM_VERSION}_ubuntu${UBUNTU_VERSION}_py${PYTHON_VERSION}_pytorch_release_${PYTORCH_VERSION}`. +- `Dockerfile.xpu`: `BASE_TAG` (default `latest`). Base image is `intel/intel-optimized-pytorch:${BASE_TAG}`; pin to a dated tag for reproducibility. + +## Podman + +Build Podman image + +```sh +$ podman build -f Dockerfile.cpu -t ptychodus:cpu . +$ podman build -f Dockerfile.cuda -t ptychodus:cuda13.0 . +$ podman build -f Dockerfile.cuda --build-arg CUDA_VERSION=12.6 -t ptychodus:cuda12.6 . +$ podman build -f Dockerfile.cuda --build-arg CUDA_VERSION=13.2 -t ptychodus:cuda13.2 . +$ podman build -f Dockerfile.rocm -t ptychodus:rocm . +$ podman build -f Dockerfile.xpu -t ptychodus:xpu . +``` + +Run container + +```{note} +GPU access requires CDI (Container Device Interface) to be configured on the host. Run `sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml` once before using `--device nvidia.com/gpu=all`. +``` + +```sh +$ xhost +local:podman +$ podman run -it --rm --env DISPLAY --security-opt label=type:container_runtime_t --network host \ + --device nvidia.com/gpu=all ptychodus:cuda13.0 +$ xhost -local:podman +``` + +## Docker + +Build Docker image + +```sh +$ docker build -f Dockerfile.cuda -t ptychodus:cuda13.0 . +``` + +(Substitute any variant file and tag as in the Podman section above.) + +Run container + +```{note} +GPU access requires [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) to be installed on the host before using `--gpus all`. +``` + +```sh +$ xhost +local:docker +$ docker run -it --rm -e "DISPLAY=$DISPLAY" -v "$HOME/.Xauthority:/root/.Xauthority:ro" --network host \ + --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ptychodus:cuda13.0 +$ xhost -local:docker +``` + +## Apptainer / Singularity + +The images above are OCI-compliant and can be converted to SIF for HPC sites (NERSC, OLCF, ALCF) that prefer Apptainer. After building an OCI image locally, convert it: + +```sh +$ apptainer build ptychodus-cuda13.0.sif docker-daemon://localhost/ptychodus:cuda13.0 +``` + +If the image lives in a registry, pull directly with `docker://` instead. + +## Wrapper Script (Beamline Workstations) + +The repository ships `scripts/podman/ptychodus`, a small bash wrapper that makes the containerized application behave like a native `ptychodus` command. It auto-selects the image variant for the host's GPU, forwards X11 to the user's desktop, bind-mounts `$HOME`, `$PWD`, and the beamline data roots (`/local`, `/gdata`) into the container, and runs rootless so that files written to those mounts are owned by the invoking user. All command line arguments are passed through to `ptychodus` unchanged. + +### Prerequisites + +The wrapper assumes facility IT has already provisioned the host: + +- Rootless podman is installed and works for the beamline account (`podman info` succeeds without `sudo`). The wrapper itself does **not** install podman. +- For NVIDIA hosts, `/etc/cdi/nvidia.yaml` exists (generated by `nvidia-ctk cdi generate` — a one-time root step; see the note above). +- For AMD hosts, the beamline account is a member of the `video` and `render` groups. +- The image variant for the host's GPU has been built into the user's rootless podman storage using the build commands in the **Podman** section above. No root is required for `podman build`. + +### Install (userspace, no sudo) + +Run from the repository root as the beamline account: + +```sh +$ install -Dm 0755 scripts/podman/ptychodus ~/.local/bin/ptychodus +$ which ptychodus # should print ~/.local/bin/ptychodus +``` + +If `which` resolves to a conda environment's launcher instead (for example `~/miniconda3/envs/ptychodus/bin/ptychodus`), prepend `~/.local/bin` to `PATH` in the account's shell rc: + +```sh +$ export PATH="$HOME/.local/bin:$PATH" +``` + +For a site-wide install, copy the script to any directory the shared beamline account can write to (for example `/opt/beamline/bin` if pre-provisioned by IT) and have each user's `PATH` include it. The wrapper does not care about its install location. + +Sanity check: + +```sh +$ ptychodus --version +``` + +stderr will show a one-line notice such as `ptychodus: using image ptychodus:cuda13.0 (detected: nvidia)`; stdout prints the version reported by the in-container `ptychodus`. + +### Usage + +```sh +$ ptychodus # GUI, auto-detect GPU +$ ptychodus -s settings.ini # GUI with settings +$ ptychodus -b reconstruct -i ./input -o ./output # headless batch +$ PTYCHODUS_IMAGE=ptychodus:cpu ptychodus # force CPU image +$ PTYCHODUS_QUIET=1 ptychodus -v # silent, version only +``` + +File paths passed via `-i`, `-o`, `-s`, or as positional arguments resolve naturally as long as they live under `$HOME`, the current directory, or one of the auto-mounted beamline roots (`/local`, `/gdata`). + +### Environment overrides + +| Variable | Effect | +| --- | --- | +| `PTYCHODUS_IMAGE` | Skip GPU auto-detection; use this image tag verbatim. | +| `PTYCHODUS_QUIET` | Suppress the stderr "using image" notice. | + +The image tag map and the list of beamline mounts live near the top of `scripts/podman/ptychodus` and can be edited in place to retag or add sites (for example `/data`, `/nsls2`). + +### Troubleshooting + +- **"image … not found in rootless podman storage"** — run the `podman build` command printed in the error. +- **"cannot open display"** — confirm `$DISPLAY` is set in the shell, then run `xhost +local:` once per X session. Wayland desktops work via XWayland as long as `$DISPLAY` is set (the default on GNOME and KDE). +- **"podman: command not found"** — ask facility IT to install rootless podman. +- **Wrong GPU family detected** — override with `PTYCHODUS_IMAGE=ptychodus:cpu` (or any other tag). +- **Files written by the container are owned by root** — rootless `--userns=keep-id` is not in effect; check `podman info` for `rootless: true`. +- **"permission denied" on a path argument** — the path is outside the bind-mounted set. Run from under `$HOME` or one of the beamline roots, or add the path to `BEAMLINE_MOUNTS` at the top of the wrapper. + +## VS Code Dev Container + +The repository includes a [Dev Container](https://containers.dev/) configuration at `.devcontainer/devcontainer.json` that builds from `Dockerfile.cpu` — the safe default for contributors without a local GPU. GPU users can edit the `dockerfile:` field to point at `Dockerfile.cuda` (or `Dockerfile.rocm`) before reopening in the container. + +1. Install the [Dev Containers](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers) extension for VS Code. + +2. Open the repository folder in VS Code, then choose **Dev Containers: Reopen in Container** from the Command Palette ({kbd}`Ctrl+Shift+P`). + +VS Code will build the image and reopen the workspace inside the container. + +## For Maintainers: Publishing Releases + +### Via pip / build + twine + +From the directory that contains `pyproject.toml`, create a wheel in `./dist/`: + +```sh +$ python -m build +``` + +Upload to PyPI: + +```sh +$ python -m twine upload --verbose dist/* +``` + +### Via uv + +From the directory that contains `pyproject.toml`, create a wheel in `./dist/`: + +```sh +$ uv build --no-sources +``` + +Upload to PyPI: + +```sh +$ uv publish +``` diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst deleted file mode 100644 index 28fb058ed..000000000 --- a/docs/source/getting_started.rst +++ /dev/null @@ -1,336 +0,0 @@ -Installation Instructions -========================= - -Python Package Index (PyPI) ---------------------------- - -To install ptychodus with the most common optional dependencies: - -.. code-block:: shell - - $ python -m pip install ptychodus[globus,gui,ptychi] - - -uv --- - -`uv <https://docs.astral.sh/uv/>`_ is a fast Python package and project manager. - -#. Install `uv <https://docs.astral.sh/uv/getting-started/installation/>`_. - -#. Install ptychodus with the most common optional dependencies: - - .. code-block:: shell - - $ uv tool install ptychodus[globus,gui,ptychi] - -#. Launch ptychodus: - - .. code-block:: shell - - $ ptychodus - -#. To upgrade ptychodus, use uv tool upgrade: - - .. code-block:: shell - - $ uv tool upgrade ptychodus[globus,gui,ptychi] - - -Conda-Forge ------------ - -#. Install `miniforge <https://github.com/conda-forge/miniforge>`_. - -#. Create the ``ptychodus`` environment - - * To install ``ptychodus`` with the GUI and all optional packages: - - .. code-block:: shell - - $ conda create -n ptychodus ptychodus-all - - * To install ``ptychodus`` with the GUI and no optional packages: - - .. code-block:: shell - - $ conda create -n ptychodus ptychodus - - * To install ``ptychodus`` without the GUI or optional packages: - - .. code-block:: shell - - $ conda create -n ptychodus ptychodus-core - -#. Activate the ``ptychodus`` environment - - .. code-block:: shell - - $ conda activate ptychodus - $ ptychodus - - -Container image variants ------------------------- - -The repository ships one Dockerfile per accelerator family. Pick the variant -that matches your hardware and select an explicit file with ``-f``: - -================================ ========================================================================================== -Dockerfile Use it for -================================ ========================================================================================== -``Dockerfile.cpu`` CPU-only hosts (no GPU; ptychi runs on CPU torch) -``Dockerfile.cuda`` NVIDIA GPUs (e.g. ALCF Polaris, NERSC Perlmutter); CUDA minor version is a build ARG -``Dockerfile.rocm`` AMD GPUs (e.g. OLCF Frontier); ROCm is a build ARG -``Dockerfile.xpu`` Intel XPU (e.g. ALCF Aurora); base tag is a build ARG -================================ ========================================================================================== - -The GPU files default to recent versions and expose ``--build-arg`` knobs to -switch: - -* ``Dockerfile.cuda``: ``CUDA_VERSION`` (default ``13.0``), ``PYTORCH_VERSION``, - ``CUDNN_VERSION``. The base image is - ``pytorch/pytorch:${PYTORCH_VERSION}-cuda${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel``; - override any args if a given combination isn't published upstream. -* ``Dockerfile.rocm``: ``ROCM_VERSION`` (default ``7.2.4``), ``UBUNTU_VERSION``, - ``PYTHON_VERSION``, ``PYTORCH_VERSION``. Base image is - ``rocm/pytorch:rocm${ROCM_VERSION}_ubuntu${UBUNTU_VERSION}_py${PYTHON_VERSION}_pytorch_release_${PYTORCH_VERSION}``. -* ``Dockerfile.xpu``: ``BASE_TAG`` (default ``latest``). Base image is - ``intel/intel-optimized-pytorch:${BASE_TAG}``; pin to a dated tag for - reproducibility. - -Podman ------- - -Build Podman image - -.. code-block:: shell - - $ podman build -f Dockerfile.cpu -t ptychodus:cpu . - $ podman build -f Dockerfile.cuda -t ptychodus:cuda13.0 . - $ podman build -f Dockerfile.cuda --build-arg CUDA_VERSION=12.6 -t ptychodus:cuda12.6 . - $ podman build -f Dockerfile.cuda --build-arg CUDA_VERSION=13.2 -t ptychodus:cuda13.2 . - $ podman build -f Dockerfile.rocm -t ptychodus:rocm . - $ podman build -f Dockerfile.xpu -t ptychodus:xpu . - -Run container - -.. note:: - - GPU access requires CDI (Container Device Interface) to be configured on the host. - Run ``sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml`` once before - using ``--device nvidia.com/gpu=all``. - -.. code-block:: shell - - $ xhost +local:podman - $ podman run -it --rm --env DISPLAY --security-opt label=type:container_runtime_t --network host \ - --device nvidia.com/gpu=all ptychodus:cuda13.0 - $ xhost -local:podman - - -Docker ------- - -Build Docker image - -.. code-block:: shell - - $ docker build -f Dockerfile.cuda -t ptychodus:cuda13.0 . - -(Substitute any variant file and tag as in the Podman section above.) - - -Run container - -.. note:: - - GPU access requires `nvidia-container-toolkit - <https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html>`_ - to be installed on the host before using ``--gpus all``. - -.. code-block:: shell - - $ xhost +local:docker - $ docker run -it --rm -e "DISPLAY=$DISPLAY" -v "$HOME/.Xauthority:/root/.Xauthority:ro" --network host \ - --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ptychodus:cuda13.0 - $ xhost -local:docker - - -Apptainer / Singularity ------------------------ - -The images above are OCI-compliant and can be converted to SIF for HPC sites -(NERSC, OLCF, ALCF) that prefer Apptainer. After building an OCI image -locally, convert it: - -.. code-block:: shell - - $ apptainer build ptychodus-cuda13.0.sif docker-daemon://localhost/ptychodus:cuda13.0 - -If the image lives in a registry, pull directly with ``docker://`` instead. - - -Wrapper Script (Beamline Workstations) --------------------------------------- - -The repository ships ``scripts/podman/ptychodus``, a small bash wrapper that -makes the containerized application behave like a native ``ptychodus`` -command. It auto-selects the image variant for the host's GPU, forwards X11 -to the user's desktop, bind-mounts ``$HOME``, ``$PWD``, and the beamline data -roots (``/local``, ``/gdata``) into the container, and runs rootless so that -files written to those mounts are owned by the invoking user. All command -line arguments are passed through to ``ptychodus`` unchanged. - -Prerequisites -^^^^^^^^^^^^^ - -The wrapper assumes facility IT has already provisioned the host: - -* Rootless podman is installed and works for the beamline account - (``podman info`` succeeds without ``sudo``). The wrapper itself does - **not** install podman. -* For NVIDIA hosts, ``/etc/cdi/nvidia.yaml`` exists (generated by - ``nvidia-ctk cdi generate`` — a one-time root step; see the note above). -* For AMD hosts, the beamline account is a member of the ``video`` and - ``render`` groups. -* The image variant for the host's GPU has been built into the user's - rootless podman storage using the build commands in the **Podman** - section above. No root is required for ``podman build``. - -Install (userspace, no sudo) -^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Run from the repository root as the beamline account: - -.. code-block:: shell - - $ install -Dm 0755 scripts/podman/ptychodus ~/.local/bin/ptychodus - $ which ptychodus # should print ~/.local/bin/ptychodus - -If ``which`` resolves to a conda environment's launcher instead (for example -``~/miniconda3/envs/ptychodus/bin/ptychodus``), prepend ``~/.local/bin`` to -``PATH`` in the account's shell rc: - -.. code-block:: shell - - $ export PATH="$HOME/.local/bin:$PATH" - -For a site-wide install, copy the script to any directory the shared -beamline account can write to (for example ``/opt/beamline/bin`` if -pre-provisioned by IT) and have each user's ``PATH`` include it. The -wrapper does not care about its install location. - -Sanity check: - -.. code-block:: shell - - $ ptychodus --version - -stderr will show a one-line notice such as -``ptychodus: using image ptychodus:cuda13.0 (detected: nvidia)``; stdout -prints the version reported by the in-container ``ptychodus``. - -Usage -^^^^^ - -.. code-block:: shell - - $ ptychodus # GUI, auto-detect GPU - $ ptychodus -s settings.ini # GUI with settings - $ ptychodus -b reconstruct -i ./input -o ./output # headless batch - $ PTYCHODUS_IMAGE=ptychodus:cpu ptychodus # force CPU image - $ PTYCHODUS_QUIET=1 ptychodus -v # silent, version only - -File paths passed via ``-i``, ``-o``, ``-s``, or as positional arguments -resolve naturally as long as they live under ``$HOME``, the current -directory, or one of the auto-mounted beamline roots (``/local``, -``/gdata``). - -Environment overrides -^^^^^^^^^^^^^^^^^^^^^ - -================================ =============================================================== -Variable Effect -================================ =============================================================== -``PTYCHODUS_IMAGE`` Skip GPU auto-detection; use this image tag verbatim. -``PTYCHODUS_QUIET`` Suppress the stderr "using image" notice. -================================ =============================================================== - -The image tag map and the list of beamline mounts live near the top of -``scripts/podman/ptychodus`` and can be edited in place to retag or add -sites (for example ``/data``, ``/nsls2``). - -Troubleshooting -^^^^^^^^^^^^^^^ - -* **"image … not found in rootless podman storage"** — run the - ``podman build`` command printed in the error. -* **"cannot open display"** — confirm ``$DISPLAY`` is set in the shell, then - run ``xhost +local:`` once per X session. Wayland desktops work via - XWayland as long as ``$DISPLAY`` is set (the default on GNOME and KDE). -* **"podman: command not found"** — ask facility IT to install rootless - podman. -* **Wrong GPU family detected** — override with - ``PTYCHODUS_IMAGE=ptychodus:cpu`` (or any other tag). -* **Files written by the container are owned by root** — rootless - ``--userns=keep-id`` is not in effect; check ``podman info`` for - ``rootless: true``. -* **"permission denied" on a path argument** — the path is outside the - bind-mounted set. Run from under ``$HOME`` or one of the beamline roots, - or add the path to ``BEAMLINE_MOUNTS`` at the top of the wrapper. - - -VS Code Dev Container ---------------------- - -The repository includes a `Dev Container <https://containers.dev/>`_ -configuration at ``.devcontainer/devcontainer.json`` that builds from -``Dockerfile.cpu`` — the safe default for contributors without a local GPU. -GPU users can edit the ``dockerfile:`` field to point at ``Dockerfile.cuda`` -(or ``Dockerfile.rocm``) before reopening in the container. - -#. Install the `Dev Containers - <https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers>`_ - extension for VS Code. - -#. Open the repository folder in VS Code, then choose - **Dev Containers: Reopen in Container** from the Command Palette - (:kbd:`Ctrl+Shift+P`). - -VS Code will build the image and reopen the workspace inside the container. - - -For Maintainers: Publishing Releases -------------------------------------- - -Via pip / build + twine -^^^^^^^^^^^^^^^^^^^^^^^^ - -From the directory that contains ``pyproject.toml``, create a wheel in ``./dist/``: - -.. code-block:: shell - - $ python -m build - -Upload to PyPI: - -.. code-block:: shell - - $ python -m twine upload --verbose dist/* - - -Via uv -^^^^^^ - -From the directory that contains ``pyproject.toml``, create a wheel in ``./dist/``: - -.. code-block:: shell - - $ uv build --no-sources - -Upload to PyPI: - -.. code-block:: shell - - $ uv publish diff --git a/docs/source/globus.md b/docs/source/globus.md new file mode 100644 index 000000000..616cbce94 --- /dev/null +++ b/docs/source/globus.md @@ -0,0 +1,55 @@ +# Globus Compute Workflow + +Perform steps 1-3 on the local and remote computers. + +1. Install [miniforge](https://github.com/conda-forge/miniforge). + +2. Make conda available in your current shell environment + + ```sh + $ eval "$(~/miniforge3/bin/conda shell.bash hook)" + ``` + +3. Create and activate the `ptychodus` conda environment + + ```sh + $ conda create -c conda-forge -n ptychodus ptychodus-all + $ conda activate ptychodus + ``` + +4. Install a [Globus compute endpoint](https://globus-compute.readthedocs.io/en/stable/quickstart.html#deploying-an-endpoint) into the `ptychodus` environment on the remote computer and configure it. An example configuration file for ALCF Polaris is bundled with the Ptychodus source distribution. + + ```sh + $ python -m pip install globus-compute-endpoint + $ globus-compute-endpoint configure + ``` + +5. Start the Globus compute endpoint + + ```sh + $ globus-compute-endpoint start <ENDPOINT_NAME> + ``` + +6. For data transfer, use a [guest collection](https://docs.globus.org/how-to/guest-collection-share-and-access) on a Globus Connect Server or use a [Globus Connect Personal endpoint](https://www.globus.org/globus-connect-personal) on your local computer. + +7. On the local computer, launch the reconstruction tasks from the "Workflow" view. + +8. On the remote computer, watch the queue (use qstat on Polaris) and Globus compute endpoint logs + + ```sh + $ tail -f ~/.globus-compute/default/endpoint.log + ``` + +9. When the demo is done, stop the Globus compute endpoint on the remote computer + + ```sh + $ globus-compute-endpoint stop + ``` + +## Example Globus Compute Endpoint Configuration + +Here is an example Globus compute endpoint `config.yaml` for ALCF Polaris: + +```{literalinclude} polaris.yaml +:language: yaml +``` diff --git a/docs/source/globus.rst b/docs/source/globus.rst deleted file mode 100644 index 972796bb7..000000000 --- a/docs/source/globus.rst +++ /dev/null @@ -1,58 +0,0 @@ -Globus Compute Workflow -======================= - -Perform steps 1-3 on the local and remote computers. - -#. Install `miniforge <https://github.com/conda-forge/miniforge>`_. - -#. Make conda available in your current shell environment - -.. code-block:: shell - - $ eval "$(~/miniforge3/bin/conda shell.bash hook)" - -#. Create and activate the ``ptychodus`` conda environment - -.. code-block:: shell - - $ conda create -c conda-forge -n ptychodus ptychodus-all - $ conda activate ptychodus - -#. Install a `Globus compute endpoint <https://globus-compute.readthedocs.io/en/stable/quickstart.html#deploying-an-endpoint>`_ - into the ``ptychodus`` environment on the remote computer and configure it. An - example configuration file for ALCF Polaris is bundled with the Ptychodus - source distribution. - -.. code-block:: shell - - $ python -m pip install globus-compute-endpoint - $ globus-compute-endpoint configure - -#. Start the Globus compute endpoint - -.. code-block:: shell - - $ globus-compute-endpoint start <ENDPOINT_NAME> - -#. For data transfer, use a `guest collection <https://docs.globus.org/how-to/guest-collection-share-and-access>`_ - on a Globus Connect Server or use a `Globus Connect Personal endpoint <https://www.globus.org/globus-connect-personal>`_ on your local computer. -#. On the local computer, launch the reconstruction tasks from the "Workflow" view. -#. On the remote computer, watch the queue (use qstat on Polaris) and Globus compute endpoint logs - -.. code-block:: shell - - $ tail -f ~/.globus-compute/default/endpoint.log - -#. When the demo is done, stop the Globus compute endpoint on the remote computer - -.. code-block:: shell - - $ globus-compute-endpoint stop - - -**Example Globus Compute Endpoint Configuration** - -Here is an example Globus compute endpoint ``config.yaml`` for ALCF Polaris: - -.. literalinclude:: polaris.yaml - :language: yaml diff --git a/docs/source/index.md b/docs/source/index.md new file mode 100644 index 000000000..68de7d6d3 --- /dev/null +++ b/docs/source/index.md @@ -0,0 +1,43 @@ +# Ptychodus documentation + +```{image} ptychodus.svg +:alt: Ptychodus Logo +:align: center +:width: 10em +``` + +[Ptychodus](https://github.com/AdvancedPhotonSource/ptychodus) is a ptychography data analysis application that reads instrument data, prepares the data for processing, and supports calling several reconstruction libraries for phase retrieval. Ptychodus can be used interactively or integrated into a data pipeline. + +```{toctree} +:maxdepth: 2 +:caption: Contents: + +getting_started +initial_guesses +readers +api +globus +pvapy +``` + +## Python API Example + +```python +from pathlib import Path +from ptychodus.model import ModelCore + + +def main() -> int: + settings_file = Path('path/to/settings.ini') + + with ModelCore(settings_file) as model: + input_product_api = model.workflow_api.create_product('new_product_name') + output_product_api = input_product_api.reconstruct_local() + output_product_api.save_product('/path/to/file.h5', file_type='HDF5') +``` + +## Indices and tables + +- {ref}`genindex` +- {ref}`modindex` +- {ref}`search` diff --git a/docs/source/index.rst b/docs/source/index.rst deleted file mode 100644 index c71a1dc0d..000000000 --- a/docs/source/index.rst +++ /dev/null @@ -1,50 +0,0 @@ -Ptychodus documentation -======================= - -.. image:: ptychodus.svg - :alt: Ptychodus Logo - :align: center - :width: 10em - - -`Ptychodus <https://github.com/AdvancedPhotonSource/ptychodus>`_ -is a ptychography data analysis application that reads instrument data, -prepares the data for processing, and supports calling several reconstruction -libraries for phase retrieval. Ptychodus can be used interactively or -integrated into a data pipeline. - - -.. toctree:: - :maxdepth: 2 - :caption: Contents: - - getting_started - initial_guesses - readers - api - globus - pvapy - - -Python API Example ------------------- - -.. code-block:: python - - from pathlib import Path - from ptychodus.model import ModelCore - - def main() -> int: - settings_file = Path("path/to/settings.ini") - - with ModelCore(settings_file) as model: - input_product_api = model.workflow_api.create_product("new_product_name") - output_product_api = input_product_api.reconstruct_local() - output_product_api.save_product("/path/to/file.h5", file_type="HDF5") - - -Indices and tables -================== -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/docs/source/initial_guesses.md b/docs/source/initial_guesses.md new file mode 100644 index 000000000..3fe92e512 --- /dev/null +++ b/docs/source/initial_guesses.md @@ -0,0 +1,524 @@ +# Initial Guess Generation + +Ptychodus can create initial data products from existing files or from generator routines. The easiest public entry point is the workflow API: + +```python +from pathlib import Path + +from ptychodus.model import ModelCore + +with ModelCore(Path('settings.ini')) as model: + product_api = model.workflow_api.create_product( + 'initial_guess', + detector_distance_m=1.0, + probe_energy_eV=10_000.0, + probe_photon_count=1.0e6, + ) + + product_api.generate_probe_positions( + 'rectangular_raster', + { + 'num_points_x': 20, + 'num_points_y': 20, + 'step_size_x_m': 75e-9, + 'step_size_y_m': 75e-9, + }, + ) + + product_api.generate_probe( + 'disk', + { + 'diameter_m': 1.0e-6, + 'defocus_distance_m': 0.0, + 'num_incoherent_modes': 1, + 'num_coherent_modes': 1, + }, + ) + + product_api.generate_object( + 'random', + { + 'amplitude_mean': 1.0, + 'amplitude_deviation': 0.0, + 'phase_deviation_turns': 0.1, + 'blur_deviation_px': 0.0, + }, + ) + + product = product_api.get_product() + product_api.save_product(Path('initial_guess.h5'), file_type='HDF5') +``` + +The product contains four main pieces of initial data: + +- metadata: detector distance, probe energy, photon count, exposure time, etc.; +- probe positions: a {py:class}`ptychodus.api.probe_positions.ProbePositionSequence`; +- probe: a {py:class}`ptychodus.api.probe.ProbeSequence`; +- object: a {py:class}`ptychodus.api.object.Object`. + +The workflow methods select generator builders by name. Passing no name uses the current settings. Passing a name plus a parameter mapping creates a builder for that product and immediately rebuilds the corresponding data. + +## Data Shapes + +Ptychodus stores complex arrays with explicit leading mode or layer dimensions. + +Probe +: A single {py:class}`ptychodus.api.probe.Probe` stores its array as `(n_incoherent_modes, height, width)`. A two-dimensional complex input is promoted to `(1, height, width)`. + +Probe sequence +: {py:class}`ptychodus.api.probe.ProbeSequence` stores its array as `(n_coherent_modes, n_incoherent_modes, height, width)`. The first dimension is the OPR/coherent-mode dimension. The second dimension is the mutually incoherent probe-mode dimension. A two-dimensional complex input is promoted to `(1, 1, height, width)`; a three-dimensional input is promoted to `(1, n_incoherent_modes, height, width)`. + +OPR weights +: If OPR/coherent modes are active, `ProbeSequence` may also carry `opr_weights` with shape `(n_scan_points, n_coherent_modes)`. When a scan point is indexed from a `ProbeSequence`, Ptychodus forms the position-specific dominant incoherent mode as the weighted sum of the coherent/OPR modes and then returns a regular `Probe`. + +Object +: {py:class}`ptychodus.api.object.Object` stores its complex transmission array as `(n_layers, height, width)`. A two-dimensional complex input is promoted to `(1, height, width)`. For multislice objects, the `layer_spacing_m` sequence must contain `n_layers - 1` spacings. + +Probe positions +: Each {py:class}`ptychodus.api.probe_positions.ProbePosition` stores an integer scan index and physical `x`/`y` coordinates in meters. Internally, `ProbePositionSequence` keeps the coordinate array in `(y, x)` order, but the public object properties are named `coordinate_x_m` and `coordinate_y_m`. + +## Probe Position Generators + +Use `WorkflowProductAPI.generate_probe_positions(name, parameters)`. + +Supported generator names are: + +`rectangular_raster` +: Cartesian grid, row-major order. + +`rectangular_snake` +: Cartesian grid with alternating scan direction on each row. + +`triangular_raster` and `triangular_snake` +: Cartesian grid with staggered rows. + +`square_raster` and `square_snake` +: Equilateral Cartesian grid with equal `x`/`y` step sizes. + +`hexagonal_raster` and `hexagonal_snake` +: Equilateral staggered grid. The row spacing is derived from the `x` step by multiplying by `sqrt(3/4)`. + +`concentric` +: Concentric shell scan. + +`spiral` +: Spiral scan. + +`lissajous` +: Lissajous scan. + +All position builders share these transform parameters: + +`affine00`, `affine01`, `affine02`, `affine10`, `affine11`, `affine12` +: Affine transform applied to generated positions. + +`jitter_radius_m` +: Optional random jitter radius in meters. A value of zero disables jitter. + +The Cartesian builders use: + +`num_points_x`, `num_points_y` +: Number of grid points along each axis. + +`step_size_x_m`, `step_size_y_m` +: Physical step size in meters. Equilateral variants derive the effective `y` step from `step_size_x_m`. + +## Probe Generators + +Use `WorkflowProductAPI.generate_probe(name, parameters)`. All generated probe builders first create a base complex probe, rescale its total intensity to the product metadata field `probe_photon_count`, then expand it into incoherent and OPR/coherent modes according to the mode settings described below. + +`disk` +: Generates a binary circular aperture with parameter `diameter_m`. The aperture is optionally propagated by `defocus_distance_m` with the angular-spectrum propagator. This is the default probe builder. + +`rectangular` +: Generates a binary rectangular aperture with `width_m` and `height_m`. It is also optionally propagated by `defocus_distance_m` with the angular-spectrum propagator. + +`super_gaussian` +: Generates a super-Gaussian amplitude profile. Parameters are `annular_radius_m`, `full_width_at_half_maximum_m` (or the builder parameter name `fwhm_m` in the low-level function), and `order_parameter`. `annular_radius_m = 0` gives a Gaussian-like spot; a positive `annular_radius_m` gives a ring or donut-like profile. This is not a propagated zone-plate simulation; it is an analytic amplitude profile. + +`fresnel_zone_plate` +: Simulates a Fresnel zone plate optic and propagates it to the sample plane with the Fresnel-transform propagator. Parameters are `zone_plate_diameter_m`, `outermost_zone_width_m`, `central_beamstop_diameter_m`, and `defocus_distance_m`. The central beamstop creates the donut-like zone-plate aperture. The focal length is computed as `zone_plate_diameter_m * outermost_zone_width_m / probe_wavelength_m` and the propagation distance is `focal_length_m + defocus_distance_m`. + +`average_pattern` +: Estimates a probe from diffraction data by taking the square root of the mean assembled diffraction pattern and back-propagating it by `detector_distance_m` with the Fresnel-transform propagator. This requires assembled diffraction data to already be available in the model. + +`zernike` +: Generates a probe as a superposition of Zernike polynomial modes inside a disk with parameter `diameter_m`. Through the builder object, callers can set the Zernike order and individual coefficients before rebuilding. The workflow parameter mapping covers the common builder parameters, but coefficient-level editing is done on the `ZernikeProbeBuilder` object. + +### Zone Plate Presets + +Ptychodus registers Fresnel zone plate presets as plugins. The current presets are: + +- `2-ID-D`: `160e-6` m diameter, `70e-9` m outermost zone width, `60e-6` m central beamstop. +- `HXN`: `160e-6` m diameter, `30e-9` m outermost zone width, `80e-6` m central beamstop. +- `LYNX`: `114.8e-6` m diameter, `60e-9` m outermost zone width, `40e-6` m central beamstop. +- `PtychoProbe`: `180e-6` m diameter, `15e-9` m outermost zone width, `15e-6` m central beamstop. +- `Velociprobe`: `180e-6` m diameter, `50e-9` m outermost zone width, `60e-6` m central beamstop. + +The GUI builder exposes these as presets. Programmatic workflow use can pass the physical values directly. + +## Probe Mode Generation + +Every generated probe type uses the shared `ProbeSequenceBuilder._build_probe_modes` path: + +1. Generate one base {py:class}`ptychodus.api.probe.Probe`. +2. Expand to `num_incoherent_modes` with {py:func}`ptychodus.api.probe_gen.generate_incoherent_probe_modes`. +3. Expand to `num_coherent_modes` OPR modes with {py:func}`ptychodus.api.probe_gen.generate_coherent_probe_modes`. +4. Store the final result as a `ProbeSequence` with array shape `(num_coherent_modes, num_incoherent_modes, height, width)`. + +### Incoherent Modes + +The incoherent-mode controls are: + +`num_incoherent_modes` +: Number of mutually incoherent modes to store in axis 1 of the final `ProbeSequence`. + +`orthogonalize_incoherent_modes` +: If true and more than one incoherent mode is requested, the generated modes are orthogonalized with `scipy.linalg.orth`. + +`incoherent_mode_decay_type` +: `none`, `polynomial`, or `exponential`. + +`incoherent_mode_decay_ratio` +: Relative strength of later modes. The builder converts this to a list of unnormalized mode weights. + +The low-level `generate_incoherent_probe_modes` routine preserves any existing modes in the input probe. If more modes are requested than already exist, it creates each additional mode by duplicating the 0-th incoherent mode and applying random separable phase wraps along the horizontal and vertical axes. In other words, the added modes start with the same amplitude structure as the dominant mode, but receive different random linear phase ramps in `x` and `y`. After optional orthogonalization, the modes are rescaled so their intensities follow the requested relative weights and their summed intensity matches the original base-probe intensity. + +The weight rules are: + +`none` +: `[1.0, 0.0, 0.0, ...]`. + +`polynomial` +: For mode index `n`, weight `(n + 1) ** b` where `b = log(decay_ratio) / log(2)`. + +`exponential` +: For mode index `n`, weight `(1 / decay_ratio) ** -n`. + +Because the weights are normalized before intensity rescaling, only their relative values matter. + +### OPR / Coherent Modes + +Ptychodus uses the term `coherent modes` for the leading `ProbeSequence` dimension. In the reconstruction context this is the OPR mode dimension. The control is: + +`num_coherent_modes` +: Number of OPR/coherent modes to store in axis 0 of the final `ProbeSequence`. + +For `num_coherent_modes == 1`, no OPR weights are stored and the probe sequence length is one. + +For `num_coherent_modes > 1`: + +- `opr_weights` is initialized with shape `(num_diffraction_patterns, num_coherent_modes)`. +- The first OPR weight is set to `1.0` for every scan point. +- Other weights are initialized as small random normal values with default scale `1.0e-6`. +- OPR mode 0 copies the full incoherent-mode probe. +- Additional OPR modes are initialized as random complex arrays only in incoherent mode 0. Other incoherent modes in those OPR modes are left zero. +- The random OPR mode images are normalized by their mean intensity when `normalize_cmodes` is true. + +### Examples + +#### Generate a multimode probe from scratch + +This example generates only the probe component of a product: a defocused Fresnel zone plate probe with three incoherent modes and two OPR/coherent modes. The product's current scan length is still used internally to size the OPR weights, but no object or explicit scan-position setup is needed for this probe-only example. + +```python +from pathlib import Path + +from ptychodus.model import ModelCore + +with ModelCore(Path('settings.ini')) as model: + product_api = model.workflow_api.create_product( + 'generated_multimode_probe', + detector_distance_m=1.0, + probe_energy_eV=10_000.0, + probe_photon_count=1.0e6, + ) + + product_api.generate_probe( + 'fresnel_zone_plate', + { + 'zone_plate_diameter_m': 180e-6, + 'outermost_zone_width_m': 50e-9, + 'central_beamstop_diameter_m': 60e-6, + 'defocus_distance_m': 10e-6, + 'num_incoherent_modes': 3, + 'orthogonalize_incoherent_modes': True, + 'incoherent_mode_decay_type': 'polynomial', + 'incoherent_mode_decay_ratio': 0.25, + 'num_coherent_modes': 2, + }, + ) + + product = product_api.get_product() + probe_array = product.probes.get_array() + opr_weights = product.probes.get_opr_weights() + + assert probe_array.shape[:2] == (2, 3) + assert opr_weights.shape == (len(product.probe_positions), 2) + + product_api.save_product(Path('generated_multimode_probe.h5'), file_type='HDF5') +``` + +The same mode parameters work with other generated probe types such as `disk`, `rectangular`, `super_gaussian`, and `zernike`. + +#### Add incoherent and OPR modes to an existing single-mode probe + +The workflow API can load an existing probe with `load_probe()`, but the current high-level `generate_probe()` call replaces the probe with a named generator rather than augmenting the loaded probe in place. To augment an already loaded single-mode probe, use the low-level mode-expansion routines on the loaded `ProbeSequence` and then register the updated product back with the workflow API. + +```python +from pathlib import Path + +from ptychodus.api.product import Product +from ptychodus.api.probe_gen import ( + generate_coherent_probe_modes, + generate_incoherent_probe_modes, +) +from ptychodus.model import ModelCore + +with ModelCore(Path('settings.ini')) as model: + product_api = model.workflow_api.create_product( + 'loaded_single_mode_probe', + detector_distance_m=1.0, + probe_energy_eV=10_000.0, + probe_photon_count=1.0e6, + ) + + product_api.load_probe(Path('single_mode_probe.npy'), file_type='NPY') + + product = product_api.get_product() + + # Collapse any OPR handling in the loaded ProbeSequence and start from + # the first coherent/OPR mode as a regular Probe. + base_probe = product.probes.get_probe_no_opr() + + probe_with_incoherent_modes = generate_incoherent_probe_modes( + model.rng, + base_probe, + imode_weights=[1.0, 0.25, 0.1], + orthogonalize=True, + ) + + expanded_probe_sequence = generate_coherent_probe_modes( + model.rng, + probe_with_incoherent_modes, + num_cmodes=2, + num_diffraction_patterns=len(product.probe_positions), + ) + + expanded_product = Product( + metadata=product.metadata, + probe_positions=product.probe_positions, + probes=expanded_probe_sequence, + object_=product.object_, + losses=product.losses, + ) + + expanded_product_api = model.workflow_api.register_product(expanded_product) + expanded_product_api.rename_product('loaded_probe_with_added_modes') + expanded_product_api.save_product( + Path('loaded_probe_with_added_modes.h5'), + file_type='HDF5', + ) +``` + +This replaces the probe with a new `ProbeSequence` whose shape is `(2, 3, height, width)` and whose OPR weights have shape `(n_scan_points, 2)`. + +## Object Generators + +Use `WorkflowProductAPI.generate_object(name, parameters)`. Object geometry is derived from the current probe geometry and scan bounding box, then optional extra padding is applied by `extra_padding_x` and `extra_padding_y`. + +`random` +: Generates a complex object from Gaussian amplitude and phase fields. Parameters are `amplitude_mean`, `amplitude_deviation`, `phase_deviation_turns`, and `blur_deviation_px`. The generated field is `amplitude * exp(2j * pi * phase_turns)`. + +`grf` +: Generates a complex Gaussian random field by spectral synthesis. Parameter: `correlation_length_px`. + +`fractal_noise` +: Generates a complex fractal-noise object by summing multiple octaves of simplex noise. Parameters are `grid_scale_px`, `num_octaves`, `gain`, and `lacunarity`. + +`dead_leaves` +: Generates a layered random disk texture. Parameters include `leaf_radius_lower_px`, `leaf_radius_upper_px`, `leaf_radius_power_law_exponent`, `leaf_amplitude_lower`, `leaf_amplitude_upper`, `leaf_phase_lower_tr`, and `leaf_phase_upper_tr`. + +`stxm` +: Generates an STXM-like object from assembled diffraction data by interpolating per-position diffraction counts onto the object grid. This requires assembled diffraction data and probe positions. + +### Multislice Objects + +Object builders call `generate_layers` after creating a single-slice object. If `object_layer_spacing_m` is empty, the object remains single-slice. If spacings are supplied, the requested number of slices is `len(object_layer_spacing_m) + 1`. When expanding from one slice to several, Ptychodus distributes the amplitude and unwrapped phase across slices as: + +- amplitude: `abs(object) ** (1 / n_slices)`; +- phase: `unwrap_phase(object) / n_slices`. + +## Low-Level Generator API + +The workflow API is usually simpler, but the low-level functions can be used directly when the caller already has geometry objects and wants arrays without registering a product. + +Typical low-level imports are: + +```python +import numpy + +from ptychodus.api.object import ObjectGeometry +from ptychodus.api.object_gen import generate_random_object +from ptychodus.api.probe import ProbeGeometry +from ptychodus.api.probe_gen import ( + FresnelZonePlate, + generate_coherent_probe_modes, + generate_disk_probe, + generate_fresnel_zone_plate_probe, + generate_incoherent_probe_modes, + rescale_probe_intensity, +) +from ptychodus.api.probe_positions_gen import generate_cartesian_probe_positions + +rng = numpy.random.default_rng(0) + +positions = list( + generate_cartesian_probe_positions( + num_points_x=20, + num_points_y=20, + step_size_x=75e-9, + step_size_y=75e-9, + ) +) + +probe_geometry = ProbeGeometry( + width_px=128, + height_px=128, + pixel_width_m=10e-9, + pixel_height_m=10e-9, +) + +base_probe = generate_fresnel_zone_plate_probe( + probe_geometry, + FresnelZonePlate( + zone_plate_diameter_m=180e-6, + outermost_zone_width_m=50e-9, + central_beamstop_diameter_m=60e-6, + ), + probe_wavelength_m=1.239841984e-10, # 10 keV + defocus_distance_m=0.0, +) +base_probe = rescale_probe_intensity(base_probe, 1.0e6) + +probe_with_imodes = generate_incoherent_probe_modes( + rng, + base_probe, + imode_weights=[1.0, 0.25, 0.1], + orthogonalize=True, +) +probe_sequence = generate_coherent_probe_modes( + rng, + probe_with_imodes, + num_cmodes=2, + num_diffraction_patterns=len(positions), +) + +object_geometry = ObjectGeometry( + width_px=512, + height_px=512, + pixel_width_m=10e-9, + pixel_height_m=10e-9, + center_x_m=0.0, + center_y_m=0.0, +) +object_guess = generate_random_object( + rng, + object_geometry, + amplitude_mean=1.0, + amplitude_deviation=0.0, + phase_mean=0.0, + phase_deviation_tr=0.1, + blur_deviation_px=0.0, +) +``` + +## Using Generated Guesses With Pty-Chi + +Ptychodus has a Pty-Chi adapter in `ptychodus.model.ptychi`. When using the normal Ptychodus reconstruction path, the adapter receives a {py:class}`ptychodus.api.reconstructor.ReconstructInput` and builds Pty-Chi task options from the product: + +- object initial guess: `product.object_.get_array()`; +- probe initial guess: `product.probes.get_array()`; +- OPR weights: `product.probes.get_opr_weights()` if available, otherwise a one-dimensional default weight vector with first entry `1.0`; +- probe positions: physical Ptychodus positions are mapped through the Ptychodus object geometry to object pixel coordinates before being passed to Pty-Chi. + +### Shape Compatibility + +Pty-Chi's documented probe convention is `(n_opr_modes, n_modes, height, width)`: + +- `n_opr_modes` is the OPR/eigenmode dimension; +- `n_modes` is the mutually incoherent probe-mode dimension. + +This matches Ptychodus `ProbeSequence.get_array()`: + +- Ptychodus `num_coherent_modes` maps to Pty-Chi `n_opr_modes`; +- Ptychodus `num_incoherent_modes` maps to Pty-Chi `n_modes`. + +Pty-Chi requires a four-dimensional probe initial guess. Therefore, when constructing Pty-Chi options manually, pass `product.probes.get_array()` rather than indexing the probe sequence down to a single `Probe`. + +Pty-Chi's planar object convention is `(n_slices, height, width)`, which matches Ptychodus `Object.get_array()`. Pty-Chi also expects multislice `slice_spacings_m` to contain `n_slices - 1` spacings, matching Ptychodus `Object.layer_spacing_m`. + +### OPR Weights + +Pty-Chi accepts OPR weights as either: + +- `(n_scan_points, n_opr_modes)`: per-position weights; or +- `(n_opr_modes,)`: one weight vector broadcast to every scan point. + +Ptychodus generated OPR weights use `(n_scan_points, n_coherent_modes)` when `num_coherent_modes > 1`. This is directly compatible with Pty-Chi. + +If the probe has more than one OPR mode, Pty-Chi requires initial OPR weights. Ptychodus satisfies this when the OPR modes were generated by `generate_coherent_probe_modes`. If a custom `ProbeSequence` is created manually with multiple leading modes, provide matching `opr_weights`. + +### Positions and Object Origin + +Ptychodus stores probe positions in meters. Pty-Chi stores probe positions in object pixel units, with order `(y, x)`. The Ptychodus Pty-Chi helper converts positions by calling `object_geometry.map_coordinates_probe_to_object(scan_point)` and then passing the resulting pixel `x` and `y` arrays into Pty-Chi probe-position options. + +The current Ptychodus helper chooses Pty-Chi's `SPECIFIED` position-origin mode and returns an origin coordinate of zero. This works together with the absolute object-pixel coordinates produced by the mapping above. If bypassing the Ptychodus helper and creating Pty-Chi options directly, make sure the position origin convention is consistent with the positions you pass: + +- Use object-pixel coordinates in `(y, x)` order. +- Keep the object large enough for every extracted patch: `position_range + probe_shape` must fit inside the object support. +- If using zero-centered positions with Pty-Chi's `SUPPORT` origin mode, ensure `-positions.min()` is approximately `positions.max()` along both axes. + +### Minimal Manual Pty-Chi Mapping + +When bypassing the Ptychodus reconstructor and calling Pty-Chi manually, the data needed from a generated Ptychodus product is: + +```python +product = product_api.get_product() + +object_initial_guess = product.object_.get_array() +probe_initial_guess = product.probes.get_array() + +try: + opr_initial_weights = product.probes.get_opr_weights() +except ValueError: + opr_initial_weights = None + +object_pixel_size_m = product.object_.get_pixel_geometry().width_m +slice_spacings_m = product.object_.layer_spacing_m or None + +object_geometry = product.object_.get_geometry() +position_y_px = [] +position_x_px = [] +for scan_point in product.probe_positions: + object_point = object_geometry.map_coordinates_probe_to_object(scan_point) + position_y_px.append(object_point.coordinate_y_px) + position_x_px.append(object_point.coordinate_x_px) +``` + +Then configure Pty-Chi with: + +- `object_options.initial_guess = object_initial_guess`; +- `object_options.pixel_size_m = object_pixel_size_m`; +- `object_options.slice_spacings_m = slice_spacings_m` for multislice; +- `probe_options.initial_guess = probe_initial_guess`; +- `probe_position_options.position_y_px = position_y_px`; +- `probe_position_options.position_x_px = position_x_px`; +- `opr_mode_weight_options.initial_weights = opr_initial_weights` when `probe_initial_guess.shape[0] > 1`. + +## Practical Notes + +- The generated probe pixel size is derived from detector geometry, detector distance, and probe energy through the product geometry. For physically meaningful generated probes, create or load the diffraction metadata before generating the product. +- `average_pattern` and `stxm` depend on assembled diffraction data. Analytic generators such as `disk`, `super_gaussian`, `fresnel_zone_plate`, `random`, `grf`, `fractal_noise`, and `dead_leaves` do not require measured patterns. +- `super_gaussian` is the closest built-in Gaussian-like probe. Ptychodus does not currently register a separate builder literally named `gaussian`. +- In workflow parameter mappings, use the builder parameter names, not the settings-file names. For example use `diameter_m` rather than `DiskDiameterInMeters`. diff --git a/docs/source/initial_guesses.rst b/docs/source/initial_guesses.rst deleted file mode 100644 index 11eec2ba3..000000000 --- a/docs/source/initial_guesses.rst +++ /dev/null @@ -1,685 +0,0 @@ -Initial Guess Generation -======================== - -Ptychodus can create initial data products from existing files or from -generator routines. The easiest public entry point is the workflow API: - -.. code-block:: python - - from pathlib import Path - - from ptychodus.model import ModelCore - - with ModelCore(Path("settings.ini")) as model: - product_api = model.workflow_api.create_product( - "initial_guess", - detector_distance_m=1.0, - probe_energy_eV=10_000.0, - probe_photon_count=1.0e6, - ) - - product_api.generate_probe_positions("rectangular_raster", { - "num_points_x": 20, - "num_points_y": 20, - "step_size_x_m": 75e-9, - "step_size_y_m": 75e-9, - }) - - product_api.generate_probe("disk", { - "diameter_m": 1.0e-6, - "defocus_distance_m": 0.0, - "num_incoherent_modes": 1, - "num_coherent_modes": 1, - }) - - product_api.generate_object("random", { - "amplitude_mean": 1.0, - "amplitude_deviation": 0.0, - "phase_deviation_turns": 0.1, - "blur_deviation_px": 0.0, - }) - - product = product_api.get_product() - product_api.save_product(Path("initial_guess.h5"), file_type="HDF5") - -The product contains four main pieces of initial data: - -* metadata: detector distance, probe energy, photon count, exposure time, etc.; -* probe positions: a :class:`ptychodus.api.probe_positions.ProbePositionSequence`; -* probe: a :class:`ptychodus.api.probe.ProbeSequence`; -* object: a :class:`ptychodus.api.object.Object`. - -The workflow methods select generator builders by name. Passing no name uses -the current settings. Passing a name plus a parameter mapping creates a builder -for that product and immediately rebuilds the corresponding data. - -Data Shapes ------------ - -Ptychodus stores complex arrays with explicit leading mode or layer -dimensions. - -Probe - A single :class:`ptychodus.api.probe.Probe` stores its array as - ``(n_incoherent_modes, height, width)``. A two-dimensional complex input is - promoted to ``(1, height, width)``. - -Probe sequence - :class:`ptychodus.api.probe.ProbeSequence` stores its array as - ``(n_coherent_modes, n_incoherent_modes, height, width)``. The first - dimension is the OPR/coherent-mode dimension. The second dimension is the - mutually incoherent probe-mode dimension. A two-dimensional complex input - is promoted to ``(1, 1, height, width)``; a three-dimensional input is - promoted to ``(1, n_incoherent_modes, height, width)``. - -OPR weights - If OPR/coherent modes are active, ``ProbeSequence`` may also carry - ``opr_weights`` with shape ``(n_scan_points, n_coherent_modes)``. When a - scan point is indexed from a ``ProbeSequence``, Ptychodus forms the - position-specific dominant incoherent mode as the weighted sum of the - coherent/OPR modes and then returns a regular ``Probe``. - -Object - :class:`ptychodus.api.object.Object` stores its complex transmission array - as ``(n_layers, height, width)``. A two-dimensional complex input is - promoted to ``(1, height, width)``. For multislice objects, the - ``layer_spacing_m`` sequence must contain ``n_layers - 1`` spacings. - -Probe positions - Each :class:`ptychodus.api.probe_positions.ProbePosition` stores an integer - scan index and physical ``x``/``y`` coordinates in meters. Internally, - ``ProbePositionSequence`` keeps the coordinate array in ``(y, x)`` order, - but the public object properties are named ``coordinate_x_m`` and - ``coordinate_y_m``. - -Probe Position Generators -------------------------- - -Use ``WorkflowProductAPI.generate_probe_positions(name, parameters)``. - -Supported generator names are: - -``rectangular_raster`` - Cartesian grid, row-major order. - -``rectangular_snake`` - Cartesian grid with alternating scan direction on each row. - -``triangular_raster`` and ``triangular_snake`` - Cartesian grid with staggered rows. - -``square_raster`` and ``square_snake`` - Equilateral Cartesian grid with equal ``x``/``y`` step sizes. - -``hexagonal_raster`` and ``hexagonal_snake`` - Equilateral staggered grid. The row spacing is derived from the ``x`` step - by multiplying by ``sqrt(3/4)``. - -``concentric`` - Concentric shell scan. - -``spiral`` - Spiral scan. - -``lissajous`` - Lissajous scan. - -All position builders share these transform parameters: - -``affine00``, ``affine01``, ``affine02``, ``affine10``, ``affine11``, ``affine12`` - Affine transform applied to generated positions. - -``jitter_radius_m`` - Optional random jitter radius in meters. A value of zero disables jitter. - -The Cartesian builders use: - -``num_points_x``, ``num_points_y`` - Number of grid points along each axis. - -``step_size_x_m``, ``step_size_y_m`` - Physical step size in meters. Equilateral variants derive the effective - ``y`` step from ``step_size_x_m``. - -Probe Generators ----------------- - -Use ``WorkflowProductAPI.generate_probe(name, parameters)``. All generated -probe builders first create a base complex probe, rescale its total intensity -to the product metadata field ``probe_photon_count``, then expand it into -incoherent and OPR/coherent modes according to the mode settings described -below. - -``disk`` - Generates a binary circular aperture with parameter ``diameter_m``. The - aperture is optionally propagated by ``defocus_distance_m`` with the - angular-spectrum propagator. This is the default probe builder. - -``rectangular`` - Generates a binary rectangular aperture with ``width_m`` and ``height_m``. - It is also optionally propagated by ``defocus_distance_m`` with the - angular-spectrum propagator. - -``super_gaussian`` - Generates a super-Gaussian amplitude profile. Parameters are - ``annular_radius_m``, ``full_width_at_half_maximum_m`` (or the builder - parameter name ``fwhm_m`` in the low-level function), and - ``order_parameter``. ``annular_radius_m = 0`` gives a Gaussian-like spot; - a positive ``annular_radius_m`` gives a ring or donut-like profile. This - is not a propagated zone-plate simulation; it is an analytic amplitude - profile. - -``fresnel_zone_plate`` - Simulates a Fresnel zone plate optic and propagates it to the sample plane - with the Fresnel-transform propagator. Parameters are - ``zone_plate_diameter_m``, ``outermost_zone_width_m``, - ``central_beamstop_diameter_m``, and ``defocus_distance_m``. The central - beamstop creates the donut-like zone-plate aperture. The focal length is - computed as ``zone_plate_diameter_m * outermost_zone_width_m / - probe_wavelength_m`` and the propagation distance is ``focal_length_m + - defocus_distance_m``. - -``average_pattern`` - Estimates a probe from diffraction data by taking the square root of the - mean assembled diffraction pattern and back-propagating it by - ``detector_distance_m`` with the Fresnel-transform propagator. This - requires assembled diffraction data to already be available in the model. - -``zernike`` - Generates a probe as a superposition of Zernike polynomial modes inside a - disk with parameter ``diameter_m``. Through the builder object, callers - can set the Zernike order and individual coefficients before rebuilding. - The workflow parameter mapping covers the common builder parameters, but - coefficient-level editing is done on the ``ZernikeProbeBuilder`` object. - -Zone Plate Presets -~~~~~~~~~~~~~~~~~~ - -Ptychodus registers Fresnel zone plate presets as plugins. The current -presets are: - -* ``2-ID-D``: ``160e-6`` m diameter, ``70e-9`` m outermost zone width, - ``60e-6`` m central beamstop. -* ``HXN``: ``160e-6`` m diameter, ``30e-9`` m outermost zone width, - ``80e-6`` m central beamstop. -* ``LYNX``: ``114.8e-6`` m diameter, ``60e-9`` m outermost zone width, - ``40e-6`` m central beamstop. -* ``PtychoProbe``: ``180e-6`` m diameter, ``15e-9`` m outermost zone width, - ``15e-6`` m central beamstop. -* ``Velociprobe``: ``180e-6`` m diameter, ``50e-9`` m outermost zone width, - ``60e-6`` m central beamstop. - -The GUI builder exposes these as presets. Programmatic workflow use can pass -the physical values directly. - -Probe Mode Generation ---------------------- - -Every generated probe type uses the shared -``ProbeSequenceBuilder._build_probe_modes`` path: - -1. Generate one base :class:`ptychodus.api.probe.Probe`. -2. Expand to ``num_incoherent_modes`` with - :func:`ptychodus.api.probe_gen.generate_incoherent_probe_modes`. -3. Expand to ``num_coherent_modes`` OPR modes with - :func:`ptychodus.api.probe_gen.generate_coherent_probe_modes`. -4. Store the final result as a ``ProbeSequence`` with array shape - ``(num_coherent_modes, num_incoherent_modes, height, width)``. - -Incoherent Modes -~~~~~~~~~~~~~~~~ - -The incoherent-mode controls are: - -``num_incoherent_modes`` - Number of mutually incoherent modes to store in axis 1 of the final - ``ProbeSequence``. - -``orthogonalize_incoherent_modes`` - If true and more than one incoherent mode is requested, the generated modes - are orthogonalized with ``scipy.linalg.orth``. - -``incoherent_mode_decay_type`` - ``none``, ``polynomial``, or ``exponential``. - -``incoherent_mode_decay_ratio`` - Relative strength of later modes. The builder converts this to a list of - unnormalized mode weights. - -The low-level ``generate_incoherent_probe_modes`` routine preserves any -existing modes in the input probe. If more modes are requested than already -exist, it creates each additional mode by duplicating the 0-th incoherent mode -and applying random separable phase wraps along the horizontal and vertical -axes. In other words, the added modes start with the same amplitude structure -as the dominant mode, but receive different random linear phase ramps in ``x`` -and ``y``. After optional orthogonalization, the modes are rescaled so their -intensities follow the requested relative weights and their summed intensity -matches the original base-probe intensity. - -The weight rules are: - -``none`` - ``[1.0, 0.0, 0.0, ...]``. - -``polynomial`` - For mode index ``n``, weight ``(n + 1) ** b`` where - ``b = log(decay_ratio) / log(2)``. - -``exponential`` - For mode index ``n``, weight ``(1 / decay_ratio) ** -n``. - -Because the weights are normalized before intensity rescaling, only their -relative values matter. - -OPR / Coherent Modes -~~~~~~~~~~~~~~~~~~~~ - -Ptychodus uses the term ``coherent modes`` for the leading -``ProbeSequence`` dimension. In the reconstruction context this is the OPR -mode dimension. The control is: - -``num_coherent_modes`` - Number of OPR/coherent modes to store in axis 0 of the final - ``ProbeSequence``. - -For ``num_coherent_modes == 1``, no OPR weights are stored and the probe -sequence length is one. - -For ``num_coherent_modes > 1``: - -* ``opr_weights`` is initialized with shape - ``(num_diffraction_patterns, num_coherent_modes)``. -* The first OPR weight is set to ``1.0`` for every scan point. -* Other weights are initialized as small random normal values with default - scale ``1.0e-6``. -* OPR mode 0 copies the full incoherent-mode probe. -* Additional OPR modes are initialized as random complex arrays only in - incoherent mode 0. Other incoherent modes in those OPR modes are left zero. -* The random OPR mode images are normalized by their mean intensity when - ``normalize_cmodes`` is true. - -Examples -~~~~~~~~ - -Generate a multimode probe from scratch -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -This example generates only the probe component of a product: a defocused -Fresnel zone plate probe with three incoherent modes and two OPR/coherent -modes. The product's current scan length is still used internally to size the -OPR weights, but no object or explicit scan-position setup is needed for this -probe-only example. - -.. code-block:: python - - from pathlib import Path - - from ptychodus.model import ModelCore - - with ModelCore(Path("settings.ini")) as model: - product_api = model.workflow_api.create_product( - "generated_multimode_probe", - detector_distance_m=1.0, - probe_energy_eV=10_000.0, - probe_photon_count=1.0e6, - ) - - product_api.generate_probe("fresnel_zone_plate", { - "zone_plate_diameter_m": 180e-6, - "outermost_zone_width_m": 50e-9, - "central_beamstop_diameter_m": 60e-6, - "defocus_distance_m": 10e-6, - "num_incoherent_modes": 3, - "orthogonalize_incoherent_modes": True, - "incoherent_mode_decay_type": "polynomial", - "incoherent_mode_decay_ratio": 0.25, - "num_coherent_modes": 2, - }) - - product = product_api.get_product() - probe_array = product.probes.get_array() - opr_weights = product.probes.get_opr_weights() - - assert probe_array.shape[:2] == (2, 3) - assert opr_weights.shape == (len(product.probe_positions), 2) - - product_api.save_product(Path("generated_multimode_probe.h5"), file_type="HDF5") - -The same mode parameters work with other generated probe types such as -``disk``, ``rectangular``, ``super_gaussian``, and ``zernike``. - -Add incoherent and OPR modes to an existing single-mode probe -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The workflow API can load an existing probe with ``load_probe()``, but the -current high-level ``generate_probe()`` call replaces the probe with a named -generator rather than augmenting the loaded probe in place. To augment an -already loaded single-mode probe, use the low-level mode-expansion routines on -the loaded ``ProbeSequence`` and then register the updated product back with -the workflow API. - -.. code-block:: python - - from pathlib import Path - - from ptychodus.api.product import Product - from ptychodus.api.probe_gen import ( - generate_coherent_probe_modes, - generate_incoherent_probe_modes, - ) - from ptychodus.model import ModelCore - - with ModelCore(Path("settings.ini")) as model: - product_api = model.workflow_api.create_product( - "loaded_single_mode_probe", - detector_distance_m=1.0, - probe_energy_eV=10_000.0, - probe_photon_count=1.0e6, - ) - - product_api.load_probe(Path("single_mode_probe.npy"), file_type="NPY") - - product = product_api.get_product() - - # Collapse any OPR handling in the loaded ProbeSequence and start from - # the first coherent/OPR mode as a regular Probe. - base_probe = product.probes.get_probe_no_opr() - - probe_with_incoherent_modes = generate_incoherent_probe_modes( - model.rng, - base_probe, - imode_weights=[1.0, 0.25, 0.1], - orthogonalize=True, - ) - - expanded_probe_sequence = generate_coherent_probe_modes( - model.rng, - probe_with_incoherent_modes, - num_cmodes=2, - num_diffraction_patterns=len(product.probe_positions), - ) - - expanded_product = Product( - metadata=product.metadata, - probe_positions=product.probe_positions, - probes=expanded_probe_sequence, - object_=product.object_, - losses=product.losses, - ) - - expanded_product_api = model.workflow_api.register_product(expanded_product) - expanded_product_api.rename_product("loaded_probe_with_added_modes") - expanded_product_api.save_product( - Path("loaded_probe_with_added_modes.h5"), - file_type="HDF5", - ) - -This replaces the probe with a new ``ProbeSequence`` whose shape is -``(2, 3, height, width)`` and whose OPR weights have shape -``(n_scan_points, 2)``. - -Object Generators ------------------ - -Use ``WorkflowProductAPI.generate_object(name, parameters)``. Object geometry -is derived from the current probe geometry and scan bounding box, then optional -extra padding is applied by ``extra_padding_x`` and ``extra_padding_y``. - -``random`` - Generates a complex object from Gaussian amplitude and phase fields. - Parameters are ``amplitude_mean``, ``amplitude_deviation``, - ``phase_deviation_turns``, and ``blur_deviation_px``. The generated field - is ``amplitude * exp(2j * pi * phase_turns)``. - -``grf`` - Generates a complex Gaussian random field by spectral synthesis. Parameter: - ``correlation_length_px``. - -``fractal_noise`` - Generates a complex fractal-noise object by summing multiple octaves of - simplex noise. Parameters are ``grid_scale_px``, ``num_octaves``, - ``gain``, and ``lacunarity``. - -``dead_leaves`` - Generates a layered random disk texture. Parameters include - ``leaf_radius_lower_px``, ``leaf_radius_upper_px``, - ``leaf_radius_power_law_exponent``, ``leaf_amplitude_lower``, - ``leaf_amplitude_upper``, ``leaf_phase_lower_tr``, and - ``leaf_phase_upper_tr``. - -``stxm`` - Generates an STXM-like object from assembled diffraction data by - interpolating per-position diffraction counts onto the object grid. This - requires assembled diffraction data and probe positions. - -Multislice Objects -~~~~~~~~~~~~~~~~~~ - -Object builders call ``generate_layers`` after creating a single-slice object. -If ``object_layer_spacing_m`` is empty, the object remains single-slice. If -spacings are supplied, the requested number of slices is -``len(object_layer_spacing_m) + 1``. When expanding from one slice to several, -Ptychodus distributes the amplitude and unwrapped phase across slices as: - -* amplitude: ``abs(object) ** (1 / n_slices)``; -* phase: ``unwrap_phase(object) / n_slices``. - -Low-Level Generator API ------------------------ - -The workflow API is usually simpler, but the low-level functions can be used -directly when the caller already has geometry objects and wants arrays without -registering a product. - -Typical low-level imports are: - -.. code-block:: python - - import numpy - - from ptychodus.api.object import ObjectGeometry - from ptychodus.api.object_gen import generate_random_object - from ptychodus.api.probe import ProbeGeometry - from ptychodus.api.probe_gen import ( - FresnelZonePlate, - generate_coherent_probe_modes, - generate_disk_probe, - generate_fresnel_zone_plate_probe, - generate_incoherent_probe_modes, - rescale_probe_intensity, - ) - from ptychodus.api.probe_positions_gen import generate_cartesian_probe_positions - - rng = numpy.random.default_rng(0) - - positions = list(generate_cartesian_probe_positions( - num_points_x=20, - num_points_y=20, - step_size_x=75e-9, - step_size_y=75e-9, - )) - - probe_geometry = ProbeGeometry( - width_px=128, - height_px=128, - pixel_width_m=10e-9, - pixel_height_m=10e-9, - ) - - base_probe = generate_fresnel_zone_plate_probe( - probe_geometry, - FresnelZonePlate( - zone_plate_diameter_m=180e-6, - outermost_zone_width_m=50e-9, - central_beamstop_diameter_m=60e-6, - ), - probe_wavelength_m=1.239841984e-10, # 10 keV - defocus_distance_m=0.0, - ) - base_probe = rescale_probe_intensity(base_probe, 1.0e6) - - probe_with_imodes = generate_incoherent_probe_modes( - rng, - base_probe, - imode_weights=[1.0, 0.25, 0.1], - orthogonalize=True, - ) - probe_sequence = generate_coherent_probe_modes( - rng, - probe_with_imodes, - num_cmodes=2, - num_diffraction_patterns=len(positions), - ) - - object_geometry = ObjectGeometry( - width_px=512, - height_px=512, - pixel_width_m=10e-9, - pixel_height_m=10e-9, - center_x_m=0.0, - center_y_m=0.0, - ) - object_guess = generate_random_object( - rng, - object_geometry, - amplitude_mean=1.0, - amplitude_deviation=0.0, - phase_mean=0.0, - phase_deviation_tr=0.1, - blur_deviation_px=0.0, - ) - -Using Generated Guesses With Pty-Chi ------------------------------------- - -Ptychodus has a Pty-Chi adapter in ``ptychodus.model.ptychi``. When using the -normal Ptychodus reconstruction path, the adapter receives a -:class:`ptychodus.api.reconstructor.ReconstructInput` and builds Pty-Chi task -options from the product: - -* object initial guess: ``product.object_.get_array()``; -* probe initial guess: ``product.probes.get_array()``; -* OPR weights: ``product.probes.get_opr_weights()`` if available, otherwise a - one-dimensional default weight vector with first entry ``1.0``; -* probe positions: physical Ptychodus positions are mapped through the - Ptychodus object geometry to object pixel coordinates before being passed to - Pty-Chi. - -Shape Compatibility -~~~~~~~~~~~~~~~~~~~ - -Pty-Chi's documented probe convention is -``(n_opr_modes, n_modes, height, width)``: - -* ``n_opr_modes`` is the OPR/eigenmode dimension; -* ``n_modes`` is the mutually incoherent probe-mode dimension. - -This matches Ptychodus ``ProbeSequence.get_array()``: - -* Ptychodus ``num_coherent_modes`` maps to Pty-Chi ``n_opr_modes``; -* Ptychodus ``num_incoherent_modes`` maps to Pty-Chi ``n_modes``. - -Pty-Chi requires a four-dimensional probe initial guess. Therefore, when -constructing Pty-Chi options manually, pass ``product.probes.get_array()`` -rather than indexing the probe sequence down to a single ``Probe``. - -Pty-Chi's planar object convention is ``(n_slices, height, width)``, which -matches Ptychodus ``Object.get_array()``. Pty-Chi also expects multislice -``slice_spacings_m`` to contain ``n_slices - 1`` spacings, matching -Ptychodus ``Object.layer_spacing_m``. - -OPR Weights -~~~~~~~~~~~ - -Pty-Chi accepts OPR weights as either: - -* ``(n_scan_points, n_opr_modes)``: per-position weights; or -* ``(n_opr_modes,)``: one weight vector broadcast to every scan point. - -Ptychodus generated OPR weights use ``(n_scan_points, n_coherent_modes)`` when -``num_coherent_modes > 1``. This is directly compatible with Pty-Chi. - -If the probe has more than one OPR mode, Pty-Chi requires initial OPR weights. -Ptychodus satisfies this when the OPR modes were generated by -``generate_coherent_probe_modes``. If a custom ``ProbeSequence`` is created -manually with multiple leading modes, provide matching ``opr_weights``. - -Positions and Object Origin -~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -Ptychodus stores probe positions in meters. Pty-Chi stores probe positions in -object pixel units, with order ``(y, x)``. The Ptychodus Pty-Chi helper -converts positions by calling -``object_geometry.map_coordinates_probe_to_object(scan_point)`` and then -passing the resulting pixel ``x`` and ``y`` arrays into Pty-Chi probe-position -options. - -The current Ptychodus helper chooses Pty-Chi's ``SPECIFIED`` position-origin -mode and returns an origin coordinate of zero. This works together with the -absolute object-pixel coordinates produced by the mapping above. If bypassing -the Ptychodus helper and creating Pty-Chi options directly, make sure the -position origin convention is consistent with the positions you pass: - -* Use object-pixel coordinates in ``(y, x)`` order. -* Keep the object large enough for every extracted patch: - ``position_range + probe_shape`` must fit inside the object support. -* If using zero-centered positions with Pty-Chi's ``SUPPORT`` origin mode, - ensure ``-positions.min()`` is approximately ``positions.max()`` along both - axes. - -Minimal Manual Pty-Chi Mapping -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -When bypassing the Ptychodus reconstructor and calling Pty-Chi manually, the -data needed from a generated Ptychodus product is: - -.. code-block:: python - - product = product_api.get_product() - - object_initial_guess = product.object_.get_array() - probe_initial_guess = product.probes.get_array() - - try: - opr_initial_weights = product.probes.get_opr_weights() - except ValueError: - opr_initial_weights = None - - object_pixel_size_m = product.object_.get_pixel_geometry().width_m - slice_spacings_m = product.object_.layer_spacing_m or None - - object_geometry = product.object_.get_geometry() - position_y_px = [] - position_x_px = [] - for scan_point in product.probe_positions: - object_point = object_geometry.map_coordinates_probe_to_object(scan_point) - position_y_px.append(object_point.coordinate_y_px) - position_x_px.append(object_point.coordinate_x_px) - -Then configure Pty-Chi with: - -* ``object_options.initial_guess = object_initial_guess``; -* ``object_options.pixel_size_m = object_pixel_size_m``; -* ``object_options.slice_spacings_m = slice_spacings_m`` for multislice; -* ``probe_options.initial_guess = probe_initial_guess``; -* ``probe_position_options.position_y_px = position_y_px``; -* ``probe_position_options.position_x_px = position_x_px``; -* ``opr_mode_weight_options.initial_weights = opr_initial_weights`` when - ``probe_initial_guess.shape[0] > 1``. - -Practical Notes ---------------- - -* The generated probe pixel size is derived from detector geometry, detector - distance, and probe energy through the product geometry. For physically - meaningful generated probes, create or load the diffraction metadata before - generating the product. -* ``average_pattern`` and ``stxm`` depend on assembled diffraction data. - Analytic generators such as ``disk``, ``super_gaussian``, - ``fresnel_zone_plate``, ``random``, ``grf``, ``fractal_noise``, and - ``dead_leaves`` do not require measured patterns. -* ``super_gaussian`` is the closest built-in Gaussian-like probe. Ptychodus - does not currently register a separate builder literally named ``gaussian``. -* In workflow parameter mappings, use the builder parameter names, not the - settings-file names. For example use ``diameter_m`` rather than - ``DiskDiameterInMeters``. diff --git a/docs/source/ptychodus.svg b/docs/source/ptychodus.svg index ff223b986..43d254c49 120000 --- a/docs/source/ptychodus.svg +++ b/docs/source/ptychodus.svg @@ -1 +1 @@ -../../ptychodus.svg \ No newline at end of file +../../src/ptychodus_store/ui/icons/ptychodus.svg \ No newline at end of file diff --git a/docs/source/pvapy.md b/docs/source/pvapy.md new file mode 100644 index 000000000..e1530bc23 --- /dev/null +++ b/docs/source/pvapy.md @@ -0,0 +1,44 @@ +# PvaPy Streaming Workflow + +- To install the [PvaPy](https://github.com/epics-base/pvaPy) + + ```sh + $ conda install -n ptychodus -c apsu pvapy + ``` + +- In Terminal 1: + + ```sh + $ pvapy-hpc-consumer \ + --input-channel pvapy:image \ + --control-channel consumer:*:control \ + --status-channel consumer:*:status \ + --output-channel consumer:*:output \ + --processor-class ptychodus.PtychodusAdImageProcessor \ + --processor-args '{ "settingsFilePath": "/path/to/ptychodus.ini", "reconstructFrameId": 1000 }' \ + --report-period 10 \ + --log-level debug + ``` + +- In Terminal 2: + + ```sh + # application status + $ pvget consumer:1:status + + # configure application + $ pvput consumer:1:control '{"command" : "configure", "args" : "{\"nPatternsTotal\": 1000}"}' + + # get last command status + $ pvget consumer:1:control + + # start area detector sim server + $ pvapy-ad-sim-server -cn pvapy:image -if /path/to/fly001.npy -rt 120 -fps 1000 + ``` + +- At the end of the demo, + + ```sh + # shutdown consumer process + pvput consumer:1:control '{"command" : "stop"}' + ``` diff --git a/docs/source/pvapy.rst b/docs/source/pvapy.rst deleted file mode 100644 index e7a6c52dc..000000000 --- a/docs/source/pvapy.rst +++ /dev/null @@ -1,45 +0,0 @@ -PvaPy Streaming Workflow -======================== - -* To install the `PvaPy <https://github.com/epics-base/pvaPy>`_ - -.. code-block:: shell - - $ conda install -n ptychodus -c apsu pvapy - -* In Terminal 1: - -.. code-block:: shell - - $ pvapy-hpc-consumer \ - --input-channel pvapy:image \ - --control-channel consumer:*:control \ - --status-channel consumer:*:status \ - --output-channel consumer:*:output \ - --processor-class ptychodus.PtychodusAdImageProcessor \ - --processor-args '{ "settingsFilePath": "/path/to/ptychodus.ini", "reconstructFrameId": 1000 }' \ - --report-period 10 \ - --log-level debug - -* In Terminal 2: - -.. code-block:: shell - - # application status - $ pvget consumer:1:status - - # configure application - $ pvput consumer:1:control '{"command" : "configure", "args" : "{\"nPatternsTotal\": 1000}"}' - - # get last command status - $ pvget consumer:1:control - - # start area detector sim server - $ pvapy-ad-sim-server -cn pvapy:image -if /path/to/fly001.npy -rt 120 -fps 1000 - -* At the end of the demo, - -.. code-block:: shell - - # shutdown consumer process - pvput consumer:1:control '{"command" : "stop"}' diff --git a/docs/source/readers.md b/docs/source/readers.md new file mode 100644 index 000000000..fc3867fde --- /dev/null +++ b/docs/source/readers.md @@ -0,0 +1,40 @@ +# Available File Readers + +File readers are implemented using a Python namespace plugin system. We would be happy to add file readers to support more ptychography instruments. + +- Advanced Photon Source (APS) + - 2-ID-D Bionanoprobe (BNP) + - 2-ID-D Microprobe + - 2-ID-E Microprobe + - 4-ID-B,G,H Polarization Modulation Spectroscopy (Polar) + - 9-ID-D Coherent Surface Scattering Imaging (CSSI) + - 12-ID-E Ptycho-SAXS + - 19-ID-E In-Situ Nanoprobe (ISN) + - 26-ID-C CNM/APS Hard X-ray Nanoprobe (HXN) + - 31-ID-E LYNX + - 33-ID-C PtychoProbe + - 33-ID-C Velociprobe + - 34-ID-C Microdiffraction, Coherent X-ray Scattering +- Linac Coherent Light Source (LCLS) + - Hutch 1.3: X-ray Pump Probe (XPP) + - SLAC NumPy Zipped Archive (`*.npz`) +- MAX IV + - NanoMAX Diffraction Endstation (`*.h5`) +- National Synchrotron Light Source II (NSLS-II) + - 3-ID Hard X-ray Nanoprobe (HXN) +- Swiss Light Source (SLS) + - X12SA: Coherent Small-Angle X-ray Scattering (cSAXS) +- Common File Formats + - Coherent X-ray Imaging (`*.cxi`) + - Comma-Separated Values (`*.csv`) + - EPICS Multi-Dimensional Archive (`*.mda`) + - fold_slice (`*.mat`, `*.h5`) + - NumPy Binary Files (`*.npy`, `*.npz`) + - Ptychodus Diffraction Patterns (`*.h5`, `*.npz`) + - Ptychodus Product (`*.h5`, `*.npz`) + - Space-Separated Values (`*.txt`) + - Tagged Image File Format (`*.tif`, `*.tiff`) + +## Good/Bad Pixel Masks + +Currently there are two numpy (NPY) file formats that can be used to indicate detector pixels that are usable ("good pixels") or unusable ("bad pixels") for processing. Both file types contain a 2-D boolean array with the same dimensions as an unprocessed detector frame. For the "good pixels" format, True indicates a usable pixel and False indicates an unusable pixel. For the "bad pixels" format, True indicates an unusable pixel and False indicates a usable pixel. When one of these files is provided, Ptychodus will zero bad pixels and provide the mask to processing algorithms that support pixel masks. When one of these files is not provided, Ptychodus assumes that all pixels should be used for processing. diff --git a/docs/source/readers.rst b/docs/source/readers.rst deleted file mode 100644 index 4c586614d..000000000 --- a/docs/source/readers.rst +++ /dev/null @@ -1,50 +0,0 @@ -Available File Readers -====================== - -File readers are implemented using a Python namespace plugin system. We would -be happy to add file readers to support more ptychography instruments. - -- Advanced Photon Source (APS) - - 2-ID-D Bionanoprobe (BNP) - - 2-ID-D Microprobe - - 2-ID-E Microprobe - - 4-ID-B,G,H Polarization Modulation Spectroscopy (Polar) - - 9-ID-D Coherent Surface Scattering Imaging (CSSI) - - 12-ID-E Ptycho-SAXS - - 19-ID-E In-Situ Nanoprobe (ISN) - - 26-ID-C CNM/APS Hard X-ray Nanoprobe (HXN) - - 31-ID-E LYNX - - 33-ID-C PtychoProbe - - 33-ID-C Velociprobe -- Advanced Light Source (ALS) - - Coherent X-ray Imaging (``*.cxi``) -- Linac Coherent Light Source (LCLS) - - Hutch 1.3: X-ray Pump Probe (XPP) -- MAX IV - - NanoMAX Diffraction Endstation (``*.h5``) -- National Synchrotron Light Source II (NSLS-II) - - 3-ID Hard X-ray Nanoprobe (HXN) -- Swiss Light Source (SLS) - - X12SA: Coherent Small-Angle X-ray Scattering (cSAXS) -- Common File Formats - - Comma-Separated Values (``*.csv``) - - EPICS Multi-Dimensional Archive (``*.mda``) - - NumPy Binary Files (``*.npy``, ``*.npz``) - - PtychoShelves (``*.mat``, ``*.h5``) - - Ptychodus Diffraction Patterns (``*.h5``, ``*.npz``) - - Ptychodus Product (``*.h5``, ``*.npz``) - - Space-Separated Values (``*.txt``) - - Tagged Image File Format (``*.tif``, ``*.tiff``) - -Good/Bad Pixel Masks -==================== - -Currently there are two numpy (NPY) file formats that can be used to indicate -detector pixels that are usable ("good pixels") or unusable ("bad pixels") for -processing. Both file types contain a 2-D boolean array with the same dimensions -as an unprocessed detector frame. For the "good pixels" format, True indicates a -usable pixel and False indicates an unusable pixel. For the "bad pixels" format, -True indicates an unusable pixel and False indicates a usable pixel. When one of -these files is provided, Ptychodus will zero bad pixels and provide the mask to -processing algorithms that support pixel masks. When one of these files is not -provided, Ptychodus assumes that all pixels should be used for processing. diff --git a/pyproject.toml b/pyproject.toml index d01f3a641..5b142f8a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "ptychodus" description = "Ptychodus is a ptychography data pipeline application." -readme = "README.rst" +readme = "README.md" requires-python = ">=3.11" license = "BSD-3-Clause" license-files = ["LICENSE"] @@ -13,13 +13,14 @@ dependencies = [ "colorcet", "h5py>=3", "hdf5plugin", + "httpx", "matplotlib", "mlflow-skinny", "numpy", "psutil", "pydantic", + "pydantic-ai>=1.0,<3", "pyyaml", - "requests", "scikit-image", "scipy", "tifffile", @@ -28,26 +29,36 @@ dependencies = [ dynamic = ["version"] [project.scripts] -convert-to-ptychodus = "ptychodus.scripts.convert_to_ptychodus:main" +convert-to-ptychodus = "ptychodus.cli.convert_to_ptychodus:main" ptychodus = "ptychodus.__main__:main" -ptychodus-bdp = "ptychodus.scripts.ptychodus_bdp:main" -ptychodus-ptychopinn-tf-test = "ptychodus.scripts.ptychopinn_tf_test:main" -ptychodus-iri-tokens = "ptychodus.scripts.genesis.ptychodus_iri_tokens:main" -ptychodus-transfer-tokens = "ptychodus.scripts.genesis.ptychodus_transfer_tokens:main" -ptychodus-system-check = "ptychodus.scripts.ptychodus_system_check:main" +ptychodus-bdp = "ptychodus.cli.ptychodus_bdp:main" +ptychodus-system-check = "ptychodus.cli.ptychodus_system_check:main" +ptychodus-store = "ptychodus_store.cli:main" [project.optional-dependencies] -docs = ["sphinx", "sphinx-copybutton", "sphinx-rtd-theme"] +docs = ["myst-parser", "sphinx", "sphinx-copybutton", "sphinx-rtd-theme"] globus = ["globus-compute-sdk", "globus-sdk>=4.5.0"] gui = ["PyQt5"] #ptychonn = ["ptychonn==0.3.*,>=0.3.7"] +ptycho-fm = ["ptycho-vit"] ptychi = ["ptychi==1.5.*"] ptychopinn = ["ptychopinn"] ptychozoon = ["ptychozoon"] +store = [ + "fastapi>=0.115", + "fastmcp>=3", + "uvicorn[standard]>=0.30", + "sqlalchemy[asyncio]>=2.0", + "aiosqlite>=0.20", + "pillow>=10", + "pydantic-settings>=2.5", + "python-multipart", +] +xraydb = ["xraydb>=4.5"] [tool.uv.sources] #ptychi = { path = "../pty-chi", editable = true } -#ptychopinn = { path = "../PtychoPINN", editable = true } +ptychopinn = { path = "../PtychoPINN", editable = true } #ptychopinn = { git = "https://github.com/hoidn/PtychoPINN", branch = "main" } #[[tool.uv.index]] @@ -57,15 +68,32 @@ ptychozoon = ["ptychozoon"] [tool.setuptools.package-data] "ptychodus" = ["py.typed"] +"ptychodus_store" = [ + "ui/index.html", + "ui/styles.css", + "ui/dist/**/*.js", + "ui/dist/**/*.js.map", + "ui/icons/*.svg", + "ui/icons/README.md", + "ui/icons/Font-Awesome-LICENSE.txt", +] [tool.setuptools.packages.find] where = ["src"] [tool.setuptools_scm] +[tool.mypy] +# The unpackaged scripts/ tree has same-named modules in sibling directories +# (three genesis submit_job.py). Deriving module names from the repo root keeps +# them distinct without adding __init__.py to a tree that is never imported. +explicit_package_bases = true +mypy_path = "src" + [[tool.mypy.overrides]] module = [ "colorcet", + "fastmcp.*", "globus_compute_sdk.*", "globus_sdk.*", "h5py", @@ -74,6 +102,7 @@ module = [ "ptychi.*", "ptycho.*", "ptycho_torch.*", + "ptycho_vit.*", "ptychonn.*", "ptychozoon.*", "pvaccess", @@ -81,6 +110,8 @@ module = [ "scipy.*", "tifffile", "torch", + "torch.*", + "xraydb.*", ] ignore_missing_imports = true @@ -98,16 +129,28 @@ select = [ "NPY", ] +[tool.pymarkdown] +plugins.md003.style = "atx" # ATX headings only, never setext +plugins.md004.style = "dash" # dash bullets only +plugins.md029.style = "ordered" # real incrementing ordered-list numbers +plugins.md013.enabled = false # line length — house style is no hard wrapping +plugins.md014.enabled = false # `$ ` shell prompts are the house convention +extensions.front-matter.enabled = true + [tool.pyright] pythonVersion = "3.11" [dependency-groups] dev = [ "mypy", + "pymarkdownlnt", "pyqt5-stubs", "pytest", + "pytest-asyncio", "ruff", "types-psutil", "types-pyyaml", - "types-requests", ] + +[tool.pytest.ini_options] +asyncio_mode = "auto" diff --git a/scripts/genesis/README.md b/scripts/genesis/README.md new file mode 100644 index 000000000..4cf3d2b9c --- /dev/null +++ b/scripts/genesis/README.md @@ -0,0 +1,43 @@ +# Genesis Facility Scripts + +14 April 2025 + +These scripts are not packaged; run them from a checkout with `ptychodus` installed: + +```sh +python scripts/genesis/ptychodus_iri_tokens.py --help # IRI facility API tokens +python scripts/genesis/ptychodus_transfer_tokens.py --help # Globus transfer tokens +python scripts/genesis/<facility>/submit_job.py # per-facility job submission +``` + +Per-facility setup is documented in [alcf/README.md](alcf/README.md), [nersc/README.md](nersc/README.md), and [olcf/README.md](olcf/README.md). + +## AmSC Data Transfer API + +The [Demo Data Transfer APIs for AmSC website](https://amsc-data-api.nersc.gov/docs) links to a [script (generate_token.py)](https://gist.github.com/tylern4/924b19e58d75046e593e0db2d87f6c5c) that gets a Globus bearer token for testing: + +```sh +python generate_token.py login \ + --mapped-collections 05d2c76a-e867-4f67-aa57-76edeb0beda0 \ + --mapped-collections 9d6d994a-6d04-11e5-ba46-22000b92c6ec +``` + +## Ptychodus Installation + +Install uv + +```sh +curl -LsSf https://astral.sh/uv/install.sh | sh +``` + +Install Python + +```sh +uv python install 3.11 +``` + +Install Ptychodus + +```sh +uv tool install ptychodus[globus,gui,ptychi] +``` diff --git a/src/ptychodus/scripts/genesis/alcf/README.md b/scripts/genesis/alcf/README.md similarity index 61% rename from src/ptychodus/scripts/genesis/alcf/README.md rename to scripts/genesis/alcf/README.md index e297dd831..096cdbf22 100644 --- a/src/ptychodus/scripts/genesis/alcf/README.md +++ b/scripts/genesis/alcf/README.md @@ -1,23 +1,25 @@ -Get Tokens -========== +# ALCF -Get ALCF IRI API (https://api.alcf.anl.gov) access tokens using instructions and scripts here: +## Get Tokens -https://github.com/argonne-lcf/alcf-facility-api-token +Get [ALCF IRI API](https://api.alcf.anl.gov) access tokens using the instructions and scripts at <https://github.com/argonne-lcf/alcf-facility-api-token>. The `globus_access_token.py` script helps you authenticate with Globus and obtain access tokens for filesystem operations. Authenticate with your ALCF account: -```bash + +```sh python globus_access_token.py authenticate ``` You can view your access token with: -```bash + +```sh python globus_access_token.py get_access_token ``` -Access Polaris -============== +## Access Polaris +```sh ssh USERNAME@polaris.alcf.anl.gov +``` Use passcode only. diff --git a/src/ptychodus/scripts/genesis/alcf/globus_access_token.py b/scripts/genesis/alcf/globus_access_token.py similarity index 100% rename from src/ptychodus/scripts/genesis/alcf/globus_access_token.py rename to scripts/genesis/alcf/globus_access_token.py diff --git a/src/ptychodus/scripts/genesis/alcf/submit_job.py b/scripts/genesis/alcf/submit_job.py similarity index 100% rename from src/ptychodus/scripts/genesis/alcf/submit_job.py rename to scripts/genesis/alcf/submit_job.py diff --git a/src/ptychodus/scripts/genesis/claude_token.py b/scripts/genesis/claude_token.py similarity index 100% rename from src/ptychodus/scripts/genesis/claude_token.py rename to scripts/genesis/claude_token.py diff --git a/src/ptychodus/scripts/genesis/generate_token.py b/scripts/genesis/generate_token.py similarity index 93% rename from src/ptychodus/scripts/genesis/generate_token.py rename to scripts/genesis/generate_token.py index 8962a400e..072b3eb4e 100644 --- a/src/ptychodus/scripts/genesis/generate_token.py +++ b/scripts/genesis/generate_token.py @@ -11,7 +11,7 @@ import globus_sdk from globus_sdk.scopes import TransferScopes from globus_sdk.gare import GlobusAuthorizationParameters -import requests +import httpx import json @@ -93,15 +93,22 @@ def test_transfer(source_url: str = None, destination_url: str = None): 'label': 'Test Transfer on Tutorial Endpoints', } - r = requests.post('https://amsc-data-api.nersc.gov/transfer/globus', json=payload, headers=auth) + r = httpx.post( + 'https://amsc-data-api.nersc.gov/transfer/globus', + json=payload, + headers=auth, + timeout=30.0, + ) print(json.dumps(r.json())) @app.command() def check_transfer(transfer_uuid: str): auth = {'Authorization': f'Bearer {globus_app.get_authorizer(RESOURCE_SERVER).access_token}'} - r = requests.get( - f'https://amsc-data-api.nersc.gov/transfer/globus/{transfer_uuid}', headers=auth + r = httpx.get( + f'https://amsc-data-api.nersc.gov/transfer/globus/{transfer_uuid}', + headers=auth, + timeout=30.0, ) print(json.dumps(r.json())) diff --git a/src/ptychodus/scripts/genesis/nersc/README.md b/scripts/genesis/nersc/README.md similarity index 54% rename from src/ptychodus/scripts/genesis/nersc/README.md rename to scripts/genesis/nersc/README.md index 7eb9974b5..372a1440d 100644 --- a/src/ptychodus/scripts/genesis/nersc/README.md +++ b/scripts/genesis/nersc/README.md @@ -1,24 +1,25 @@ -Get Tokens -========== +# NERSC -The script (get_globus_token.py) and instructions (get_globus_token.md) are from here: +## Get Tokens -https://github.com/NERSC/iri-api-get-globus-token +The script (`get_globus_token.py`) and instructions ([get_globus_token.md](get_globus_token.md)) are from <https://github.com/NERSC/iri-api-get-globus-token>. Run the script and follow instructions to input the auth code: -```bash + +```sh python get_globus_token.py ``` + Token JSON is saved to `~/.globus/auth_tokens.json`. -Provided Instructions -===================== +## Provided Instructions Use account "amsc013". qos name for GPU "express_amsc_g" and for CPU "express_amsc". -Access Perlmutter -================= +## Access Perlmutter +```sh ssh USERNAME@perlmutter.nersc.gov +``` Use password + passcode. diff --git a/src/ptychodus/scripts/genesis/nersc/get_globus_token.md b/scripts/genesis/nersc/get_globus_token.md similarity index 98% rename from src/ptychodus/scripts/genesis/nersc/get_globus_token.md rename to scripts/genesis/nersc/get_globus_token.md index 36c634b81..a21bb353e 100644 --- a/src/ptychodus/scripts/genesis/nersc/get_globus_token.md +++ b/scripts/genesis/nersc/get_globus_token.md @@ -28,7 +28,7 @@ This document explains how to use: 1. Python 3.9+ (recommended). 2. `globus-sdk` installed: -```bash +```sh pip install globus-sdk ``` @@ -36,7 +36,7 @@ pip install globus-sdk Run: -```bash +```sh python get_globus_token.py ``` @@ -53,13 +53,13 @@ What happens: By default, the script requests tokens for both facilities: -```bash +```sh python get_globus_token.py --facilities nersc alcf ``` You can limit token acquisition to a subset: -```bash +```sh python get_globus_token.py --facilities nersc python get_globus_token.py --facilities alcf ``` @@ -71,7 +71,7 @@ Supported facility names are: ## Print token to terminal (optional) -```bash +```sh python get_globus_token.py --print-token ``` @@ -80,7 +80,7 @@ The printed tokens correspond to the selected facilities. By default that means ## Force a new interactive login -```bash +```sh python get_globus_token.py --force-login ``` @@ -88,7 +88,7 @@ This skips refresh and always performs browser auth. ## Force a fresh IdP login prompt -```bash +```sh python get_globus_token.py --force-login --prompt-login ``` @@ -97,7 +97,7 @@ This is useful when the server side shows an empty `session_info.authentications ## Refresh saved tokens only -```bash +```sh python get_globus_token.py --refresh-only ``` @@ -107,13 +107,13 @@ If refresh is not possible, or if refresh does not return all requested facility ## Validate the NERSC IRI token -```bash +```sh python get_globus_token.py --validate-iri ``` This calls: -```bash +```sh GET https://api.iri.nersc.gov/api/v1/account/projects ``` @@ -124,7 +124,7 @@ If the response includes `session_info.authentications: {}`, the script treats t You can combine validation with token printing or refresh-only mode: -```bash +```sh python get_globus_token.py --refresh-only --validate-iri --print-token python get_globus_token.py --force-login --prompt-login --validate-iri python get_globus_token.py --facilities nersc --validate-iri --print-token @@ -132,7 +132,7 @@ python get_globus_token.py --facilities nersc --validate-iri --print-token ## Use a custom token file path -```bash +```sh python get_globus_token.py --token-file /path/to/auth_tokens.json ``` diff --git a/src/ptychodus/scripts/genesis/nersc/get_globus_token.py b/scripts/genesis/nersc/get_globus_token.py similarity index 100% rename from src/ptychodus/scripts/genesis/nersc/get_globus_token.py rename to scripts/genesis/nersc/get_globus_token.py diff --git a/src/ptychodus/scripts/genesis/nersc/submit_job.py b/scripts/genesis/nersc/submit_job.py similarity index 100% rename from src/ptychodus/scripts/genesis/nersc/submit_job.py rename to scripts/genesis/nersc/submit_job.py diff --git a/scripts/genesis/olcf/README.md b/scripts/genesis/olcf/README.md new file mode 100644 index 000000000..524ad0a2a --- /dev/null +++ b/scripts/genesis/olcf/README.md @@ -0,0 +1,17 @@ +# OLCF + +## Get Tokens + +Generate a token at <https://docs.olcf.ornl.gov/services_and_applications/s3m/overview.html#generate-a-token>. + +Use open enclave account along with the CSC682 project when you generate the token. + +## Access Odo + +Instructions: <https://docs.olcf.ornl.gov/systems/odo_user_guide.html> + +```sh +ssh USERNAME@login1.odo.olcf.ornl.gov +``` + +Use password only. diff --git a/src/ptychodus/scripts/genesis/olcf/check_token.py b/scripts/genesis/olcf/check_token.py similarity index 84% rename from src/ptychodus/scripts/genesis/olcf/check_token.py rename to scripts/genesis/olcf/check_token.py index 69dcf3986..e835d66d7 100755 --- a/src/ptychodus/scripts/genesis/olcf/check_token.py +++ b/scripts/genesis/olcf/check_token.py @@ -4,7 +4,7 @@ import logging import sys -import requests +import httpx from ptychodus.model.genesis.iri.client import get_iri_tokens_file from ptychodus.model.genesis.tokens import create_headers, load_tokens @@ -29,14 +29,18 @@ def main() -> None: headers = create_headers(olcf_token) logger.info('Testing token validity...') - token_response = requests.get( - 'https://s3m.olcf.ornl.gov/olcf/v1/token/ctls/introspect', headers=headers + token_response = httpx.get( + 'https://s3m.olcf.ornl.gov/olcf/v1/token/ctls/introspect', + headers=headers, + timeout=30.0, ) print(json.dumps(token_response.json(), indent=2)) logger.info('Testing resource availability...') - resource_response = requests.get( - 'https://amsc-open.s3m.olcf.ornl.gov/api/v1/status/resources/odo', headers=headers + resource_response = httpx.get( + 'https://amsc-open.s3m.olcf.ornl.gov/api/v1/status/resources/odo', + headers=headers, + timeout=30.0, ) print(json.dumps(resource_response.json(), indent=2)) diff --git a/src/ptychodus/scripts/genesis/olcf/example.sh b/scripts/genesis/olcf/example.sh similarity index 100% rename from src/ptychodus/scripts/genesis/olcf/example.sh rename to scripts/genesis/olcf/example.sh diff --git a/src/ptychodus/scripts/genesis/olcf/submit_job.py b/scripts/genesis/olcf/submit_job.py similarity index 100% rename from src/ptychodus/scripts/genesis/olcf/submit_job.py rename to scripts/genesis/olcf/submit_job.py diff --git a/src/ptychodus/scripts/genesis/ptychodus_iri_tokens.py b/scripts/genesis/ptychodus_iri_tokens.py similarity index 97% rename from src/ptychodus/scripts/genesis/ptychodus_iri_tokens.py rename to scripts/genesis/ptychodus_iri_tokens.py index b03077062..2e4a90a0c 100755 --- a/src/ptychodus/scripts/genesis/ptychodus_iri_tokens.py +++ b/scripts/genesis/ptychodus_iri_tokens.py @@ -5,7 +5,7 @@ import json import logging -import requests +import httpx from ptychodus.api.settings import SettingsRegistry from ptychodus.model.genesis.core import create_facility_adapters @@ -53,7 +53,7 @@ def check_tokens() -> None: try: projects = client.account.get_projects() - except requests.HTTPError as exc: + except httpx.HTTPStatusError as exc: logger.error(f'"{name}" token error: {exc}') else: data = [project.model_dump(mode='json') for project in projects] @@ -77,7 +77,7 @@ def list_resources() -> None: facility = client.facility.get_facility() sites = client.facility.get_sites() resources = client.status.get_resources() - except requests.HTTPError as exc: + except httpx.HTTPStatusError as exc: logger.error(f'"{name}" token error: {exc}') else: data[name] = { diff --git a/src/ptychodus/scripts/genesis/ptychodus_transfer_tokens.py b/scripts/genesis/ptychodus_transfer_tokens.py similarity index 97% rename from src/ptychodus/scripts/genesis/ptychodus_transfer_tokens.py rename to scripts/genesis/ptychodus_transfer_tokens.py index 628c6c8d4..104b9459d 100755 --- a/src/ptychodus/scripts/genesis/ptychodus_transfer_tokens.py +++ b/scripts/genesis/ptychodus_transfer_tokens.py @@ -4,7 +4,7 @@ import json import logging -import requests +import httpx from ptychodus.model.genesis.core import create_globus_transfer_providers from ptychodus.model.genesis.tokens import GenesisAccessTokens, save_tokens @@ -32,7 +32,7 @@ def check_tokens() -> None: try: data = client.check_auth_token() - except requests.HTTPError as exc: + except httpx.HTTPStatusError as exc: logger.error(f'"{name}" token error: {exc}') else: logger.info(f'"{name}" token response:' + json.dumps(data, indent=4)) diff --git a/scripts/npz_dump.py b/scripts/npz_dump.py new file mode 100644 index 000000000..f52ed43aa --- /dev/null +++ b/scripts/npz_dump.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python +"""Print a one-line summary (name, dtype, shape) of every array in an .npz file. + +Example usage: + + npz-dump path/to/archive.npz +""" + +from pathlib import Path +import argparse +import sys + +import numpy + + +def main() -> int: + prog = Path(__file__).stem.lower() + parser = argparse.ArgumentParser( + prog=prog, + description='List the arrays inside an .npz file with their dtype and shape.', + ) + parser.add_argument( + 'file', + metavar='NPZ_FILE', + type=argparse.FileType('rb'), + help='Path to the .npz file.', + ) + args = parser.parse_args() + + with numpy.load(args.file.name) as npz: + names = list(npz.files) + name_width = max((len(name) for name in names), default=0) + for name in names: + array = npz[name] + line = f'{name:<{name_width}} {array.dtype!s:<10} {array.shape}' + if array.ndim == 0 or array.size == 1: + line += f' = {array.item()!r}' + print(line) + + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/src/ptychodus/scripts/podman/ptychodus b/scripts/podman/ptychodus similarity index 99% rename from src/ptychodus/scripts/podman/ptychodus rename to scripts/podman/ptychodus index 6c438f6fd..7f6c287b2 100755 --- a/src/ptychodus/scripts/podman/ptychodus +++ b/scripts/podman/ptychodus @@ -127,7 +127,7 @@ fi # --- Prerequisite checks ----------------------------------------------------- command -v podman >/dev/null 2>&1 \ - || die "podman not installed; see docs/source/getting_started.rst" + || die "podman not installed; see docs/source/getting_started.md" if ! podman image exists "$IMAGE"; then build_family="$family" diff --git a/scripts/ptychopinn_demo.py b/scripts/ptychopinn_demo.py new file mode 100644 index 000000000..942d975cd --- /dev/null +++ b/scripts/ptychopinn_demo.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python +"""Generate comparison montages of probeGuess, objectGuess, and mean diffraction +patterns across the experimental and synthetic ptychopinn_demo datasets. + +Produces three PNG files in the output directory: + +* ``probeGuess_montage.png`` and ``objectGuess_montage.png``: 2x4 grids that + render the complex probe and object using ptychodus' HSV color model + (hue=phase, value=amplitude). An HSV color-wheel legend occupies one of the + empty top-row cells. +* ``diff3d_mean_montage.png``: 2x4 grid of the per-pattern mean diffraction + intensity. Display is sqrt-compressed (matplotlib ``PowerNorm(gamma=0.5)``) + with per-panel intensity ranges and a colorbar next to each panel. + +Each montage places the (single) experimental dataset on the top-left and the +four synthetic datasets across the bottom row. +""" + +from __future__ import annotations +from pathlib import Path +from typing import Protocol +import argparse +import sys + +import matplotlib + +matplotlib.use('Agg') + +import matplotlib.pyplot as plt # noqa: E402 +import numpy # noqa: E402 +from matplotlib.colors import PowerNorm # noqa: E402 + +from ptychodus.api.geometry import PixelGeometry # noqa: E402 +from ptychodus.api.visualization import ( # noqa: E402 + CylindricalColorModel, + ScalarTransformation, + visualize_complex_values, +) + +DEFAULT_DATA_DIR = Path('/home/beams0/SHENKE/workspace/ptychopinn_demo') +EXPERIMENTAL_SUBDIR = 'experimental_training_dataset_ic2' +SYNTHETIC_SUBDIR = 'synthetic_training_dataset_ic2' +COMPLEX_KEYS = ('probeGuess', 'objectGuess') +DIFFRACTION_KEY = 'diff3d' +DIFFRACTION_CMAP = 'inferno' +DIFFRACTION_GAMMA = 0.5 # sqrt scaling via PowerNorm +N_SYNTHETIC_COLS = 4 +COLORWHEEL_COL = 1 # blank cell in row 0 to host the HSV legend +UNIT_PIXEL_GEOMETRY = PixelGeometry(width_m=1.0, height_m=1.0) + + +class Renderer(Protocol): + def __call__(self, ax: plt.Axes, values: numpy.ndarray, caption: str) -> None: ... + + +def _render_complex(ax: plt.Axes, values: numpy.ndarray, caption: str) -> None: + """Draw a complex 2D array onto ``ax`` using the ptychodus HSV color model.""" + product = visualize_complex_values( + values, + UNIT_PIXEL_GEOMETRY, + CylindricalColorModel.HSV_VALUE, + amplitude_transform=ScalarTransformation.IDENTITY, + ) + ax.imshow(product.get_image_rgba(), interpolation='nearest') + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_title(caption, fontsize=10) + + +def _draw_colorwheel(ax: plt.Axes, *, size: int = 256) -> None: + """Render the HSV_VALUE legend: a unit-disk color wheel with phase tick labels.""" + y, x = numpy.mgrid[-1.0 : 1.0 : size * 1j, -1.0 : 1.0 : size * 1j] + r = numpy.hypot(x, y) + theta = numpy.arctan2(y, x) + z = (r * numpy.exp(1j * theta)).astype(numpy.complex64) + product = visualize_complex_values( + z, + UNIT_PIXEL_GEOMETRY, + CylindricalColorModel.HSV_VALUE, + amplitude_transform=ScalarTransformation.IDENTITY, + ) + rgba = product.get_image_rgba().copy() + rgba[..., 3] = (r <= 1.0).astype(rgba.dtype) + ax.imshow(rgba, interpolation='bilinear', extent=(-1.0, 1.0, -1.0, 1.0), origin='lower') + + for label, angle_deg in ( + ('0', 0), + (r'$+\pi/2$', 90), + (r'$\pm\pi$', 180), + (r'$-\pi/2$', 270), + ): + angle = numpy.deg2rad(angle_deg) + ax.text( + 1.18 * numpy.cos(angle), + 1.18 * numpy.sin(angle), + label, + ha='center', + va='center', + fontsize=9, + ) + + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_aspect('equal') + ax.set_xlim(-1.5, 1.5) + ax.set_ylim(-1.5, 1.5) + ax.set_title('HSV legend\nhue=phase, value=amplitude', fontsize=10) + for spine in ax.spines.values(): + spine.set_visible(False) + + +def _render_real( + ax: plt.Axes, values: numpy.ndarray, caption: str, *, cmap: str = DIFFRACTION_CMAP +) -> None: + """Draw a real-valued 2D array with a sqrt-compressed colormap and a colorbar.""" + vmin = float(values.min()) + vmax = float(values.max()) + norm = PowerNorm(gamma=DIFFRACTION_GAMMA, vmin=max(vmin, 0.0), vmax=vmax) + img = ax.imshow(values, cmap=cmap, norm=norm, interpolation='nearest') + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_title(caption, fontsize=10) + ax.figure.colorbar(img, ax=ax, fraction=0.046, pad=0.04) + + +def _load_key(path: Path, key: str) -> numpy.ndarray: + with numpy.load(path) as npz: + return numpy.asarray(npz[key]) + + +def _mean_intensity(diff3d: numpy.ndarray) -> numpy.ndarray: + """Mean intensity across the scan axis, clamped to non-negative values.""" + return numpy.maximum(diff3d.mean(axis=0), 0.0) + + +def _populate_row( + axes: numpy.ndarray, + row: int, + datasets: list[tuple[str, numpy.ndarray]], + caption_prefix: str, + render: Renderer, +) -> None: + """Render ``datasets`` along ``axes[row, :]``, turning off any leftover cells.""" + for col in range(N_SYNTHETIC_COLS): + if col < len(datasets): + name, array = datasets[col] + render(axes[row, col], array, f'{caption_prefix}: {name}') + else: + axes[row, col].axis('off') + + +def _build_complex_montage( + array_name: str, + experimental: list[tuple[str, numpy.ndarray]], + synthetic: list[tuple[str, numpy.ndarray]], +) -> plt.Figure: + fig, axes = plt.subplots(2, N_SYNTHETIC_COLS, figsize=(16, 8), squeeze=False) + fig.suptitle(f'{array_name}: experimental (top) vs synthetic (bottom)', fontsize=14) + + _populate_row(axes, 0, experimental, 'experimental', _render_complex) + for col in range(len(experimental), N_SYNTHETIC_COLS): + if col == COLORWHEEL_COL: + _draw_colorwheel(axes[0, col]) + else: + axes[0, col].axis('off') + + _populate_row(axes, 1, synthetic, 'synthetic', _render_complex) + + fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.96)) + return fig + + +def _build_diffraction_montage( + experimental: list[tuple[str, numpy.ndarray]], + synthetic: list[tuple[str, numpy.ndarray]], +) -> plt.Figure: + fig, axes = plt.subplots(2, N_SYNTHETIC_COLS, figsize=(16, 8), squeeze=False) + fig.suptitle( + f'{DIFFRACTION_KEY} mean intensity (sqrt-compressed, per-panel range): ' + 'experimental (top) vs synthetic (bottom)', + fontsize=14, + ) + _populate_row(axes, 0, experimental, 'experimental', _render_real) + _populate_row(axes, 1, synthetic, 'synthetic', _render_real) + fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.96)) + return fig + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + 'Generate probeGuess, objectGuess, and mean diffraction-pattern ' + 'comparison montages from the ptychopinn_demo training datasets.' + ) + ) + parser.add_argument( + '--data-dir', + type=Path, + default=DEFAULT_DATA_DIR, + help=f'Root directory containing the dataset subfolders (default: {DEFAULT_DATA_DIR}).', + ) + parser.add_argument( + '--output-dir', + type=Path, + default=Path.cwd(), + help='Directory to write PNG figures into (default: current working directory).', + ) + parser.add_argument( + '--dpi', + type=int, + default=300, + help='Resolution of saved PNG files (default: 300).', + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + + experimental_dir = args.data_dir / EXPERIMENTAL_SUBDIR + synthetic_dir = args.data_dir / SYNTHETIC_SUBDIR + + experimental_paths = sorted(experimental_dir.glob('*.npz')) + synthetic_paths = sorted(synthetic_dir.glob('synthetic_*.npz')) + + if not experimental_paths: + print(f'No experimental .npz files found in {experimental_dir}', file=sys.stderr) + return 1 + if not synthetic_paths: + print(f'No synthetic .npz files found in {synthetic_dir}', file=sys.stderr) + return 1 + + synthetic_paths = synthetic_paths[:N_SYNTHETIC_COLS] + args.output_dir.mkdir(parents=True, exist_ok=True) + + for key in COMPLEX_KEYS: + experimental = [(p.stem, _load_key(p, key)) for p in experimental_paths] + synthetic = [(p.stem, _load_key(p, key)) for p in synthetic_paths] + fig = _build_complex_montage(key, experimental, synthetic) + out_path = args.output_dir / f'{key}_montage.png' + fig.savefig(out_path, dpi=args.dpi, bbox_inches='tight') + plt.close(fig) + print(f'Wrote {out_path}') + + experimental = [ + (p.stem, _mean_intensity(_load_key(p, DIFFRACTION_KEY))) for p in experimental_paths + ] + synthetic = [(p.stem, _mean_intensity(_load_key(p, DIFFRACTION_KEY))) for p in synthetic_paths] + fig = _build_diffraction_montage(experimental, synthetic) + out_path = args.output_dir / f'{DIFFRACTION_KEY}_mean_montage.png' + fig.savefig(out_path, dpi=args.dpi, bbox_inches='tight') + plt.close(fig) + print(f'Wrote {out_path}') + + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/src/ptychodus/scripts/ptychopinn_tf_test.py b/scripts/ptychopinn_tf_check.py similarity index 97% rename from src/ptychodus/scripts/ptychopinn_tf_test.py rename to scripts/ptychopinn_tf_check.py index 573227448..eaabf731b 100755 --- a/src/ptychodus/scripts/ptychopinn_tf_test.py +++ b/scripts/ptychopinn_tf_check.py @@ -50,7 +50,7 @@ def _validate_npz_keys(file_path: Path, required_keys: set[str], label: str) -> def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - prog='ptychodus-ptychopinn-tf-test', + prog=Path(__file__).stem.lower(), description='Run the PtychoPINN TF train/load/infer workflow via ptychodus.', ) parser.add_argument( @@ -177,6 +177,7 @@ def main() -> int: probe_energy_eV=args.probe_energy_ev, probe_photon_count=max_pattern_counts, exposure_time_s=args.exposure_time_s, + diffraction=workflow_diffraction_api, ) input_product_api.load_probe_positions(dataset_path, file_type='SLAC_NPZ') input_product_api.load_probe(dataset_path, file_type='SLAC_NPZ') @@ -204,7 +205,9 @@ def main() -> int: inference_settings.n_samples.set_value(args.test_samples) input_product_api.train_reconstructor_local( - train_dir, model_out_dir, algorithm=args.reconstructor + train_dir, + model_out_dir, + algorithm=args.reconstructor, ) model_file = model_out_dir / 'wts.h5.zip' diff --git a/setup.py b/setup.py new file mode 100644 index 000000000..5aed2e46e --- /dev/null +++ b/setup.py @@ -0,0 +1,64 @@ +"""Setuptools shim that builds the ptychodus_store frontend before packaging. + +The main project configuration lives in pyproject.toml; this file only exists to +register a build_py subclass that runs `tsc` in src/ptychodus_store/ui/ so wheel +builds ship the compiled UI without requiring maintainers to remember a +pre-build step. The ui/dist/ output stays gitignored (see .gitignore and +CLAUDE.md Repository Notes). +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +from pathlib import Path + +from setuptools import setup # type: ignore[import-untyped] +from setuptools.command.build_py import build_py # type: ignore[import-untyped] + +UI_DIR = Path(__file__).parent / 'src' / 'ptychodus_store' / 'ui' +DIST_ENTRY = UI_DIR / 'dist' / 'main.js' +SKIP_ENV = 'PTYCHODUS_STORE_SKIP_UI_BUILD' + + +class BuildFrontend(build_py): + """Run `tsc` before the standard build_py step so ui/dist/ ships in the wheel.""" + + def run(self) -> None: + self._build_ui() + super().run() + + def _build_ui(self) -> None: + if not (UI_DIR / 'tsconfig.json').is_file(): + return # ui subpackage absent (e.g. sdist without ui/); nothing to build + + if DIST_ENTRY.is_file(): + src_dir = UI_DIR / 'src' + latest_src = max( + (p.stat().st_mtime for p in src_dir.rglob('*.ts')), + default=0.0, + ) + if DIST_ENTRY.stat().st_mtime >= latest_src: + return + + tsc = shutil.which('tsc') + if tsc is None: + sys.stderr.write( + 'ptychodus_store frontend build requires `tsc` on PATH.\n' + ' Install: `npm install -g typescript`, or use nodeenv:\n' + ' uv tool install nodeenv\n' + ' nodeenv --node=lts --prebuilt ~/.local/node-lts\n' + ' export PATH="$HOME/.local/node-lts/bin:$PATH"\n' + ' npm install -g typescript\n' + f' To skip the frontend build (ships an empty UI), set {SKIP_ENV}=1.\n' + ) + if os.environ.get(SKIP_ENV): + return + raise SystemExit(2) + + subprocess.check_call([tsc], cwd=str(UI_DIR)) + + +setup(cmdclass={'build_py': BuildFrontend}) diff --git a/src/ptychodus/api/common.py b/src/ptychodus/api/common.py index 710783430..ded9355a8 100644 --- a/src/ptychodus/api/common.py +++ b/src/ptychodus/api/common.py @@ -1,5 +1,7 @@ """Common type aliases, physical constants, and utility functions used throughout the API.""" +from __future__ import annotations + from dataclasses import dataclass from pathlib import Path from typing import Any, Final, TypeAlias, overload @@ -27,19 +29,38 @@ def get_ptychodus_dir() -> Path: + """Return the user's Ptychodus configuration directory (``~/.ptychodus``).""" return Path.home() / '.ptychodus' @overload -def lerp(lower: float, upper: float, frac: float) -> float: ... +def lerp(lower: float, upper: float, frac: float) -> float: + """Linearly interpolate between *lower* and *upper* by fraction *frac* in [0, 1].""" + ... + + @overload -def lerp(lower: complex, upper: complex, frac: float) -> complex: ... +def lerp(lower: complex, upper: complex, frac: float) -> complex: + """Linearly interpolate between *lower* and *upper* by fraction *frac* in [0, 1].""" + ... + + @overload -def lerp(lower: RealArrayType, upper: RealArrayType, frac: float) -> RealArrayType: ... +def lerp(lower: RealArrayType, upper: RealArrayType, frac: float) -> RealArrayType: + """Linearly interpolate between *lower* and *upper* by fraction *frac* in [0, 1].""" + ... + + @overload -def lerp(lower: RealArrayType, upper: RealArrayType, frac: RealArrayType) -> RealArrayType: ... +def lerp(lower: RealArrayType, upper: RealArrayType, frac: RealArrayType) -> RealArrayType: + """Linearly interpolate between *lower* and *upper* by fraction *frac* in [0, 1].""" + ... + + @overload -def lerp(lower: float, upper: float, frac: RealArrayType) -> RealArrayType: ... +def lerp(lower: float, upper: float, frac: RealArrayType) -> RealArrayType: + """Linearly interpolate between *lower* and *upper* by fraction *frac* in [0, 1].""" + ... def lerp( @@ -53,11 +74,13 @@ def lerp( @dataclass(frozen=True) class NoiseFloor: + """Robust noise-floor estimate: background value and its median absolute deviation.""" + background_value: float median_absolute_deviation: float @classmethod - def from_values(cls, values: RealArrayType) -> 'NoiseFloor': + def from_values(cls, values: RealArrayType) -> NoiseFloor: background_value = numpy.median(values) absolute_deviation = numpy.abs(values - background_value) return cls( diff --git a/src/ptychodus/api/diffraction.py b/src/ptychodus/api/diffraction.py index 756a10163..93b70da27 100644 --- a/src/ptychodus/api/diffraction.py +++ b/src/ptychodus/api/diffraction.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass +from enum import StrEnum from pathlib import Path from typing import overload, Any, TypeAlias @@ -200,29 +201,51 @@ def get_patterns(self) -> DiffractionPatterns: return self._patterns +class Polarization(StrEnum): + """Beam polarization state (used by XMCD analysis). + + Stored as its string value for stable, human-readable HDF5/INI round-trip. + """ + + LEFT_CIRCULAR = 'left_circular' + RIGHT_CIRCULAR = 'right_circular' + + @dataclass(frozen=True) class DiffractionMetadata: """Metadata describing a diffraction dataset (geometry, energy, file path, etc.).""" num_patterns_per_array: Sequence[int] pattern_dtype: DiffractionPatternDType + detector_extent: ImageExtent detector_distance_m: float | None = None - detector_extent: ImageExtent | None = None detector_pixel_geometry: PixelGeometry | None = None crop_center: CropCenter | None = None probe_energy_eV: float | None = None # noqa: N815 probe_photon_count: int | None = None exposure_time_s: float | None = None tomography_angle_deg: float | None = None + tilt_angle_deg: float | None = None + polarization: Polarization | None = None file_path: Path | None = None @classmethod def create_null(cls, file_path: Path | None = None) -> DiffractionMetadata: - return cls([], numpy.dtype(numpy.ubyte), file_path=file_path) + return cls( + num_patterns_per_array=[], + pattern_dtype=numpy.dtype(numpy.ubyte), + detector_extent=ImageExtent(width_px=0, height_px=0), + file_path=file_path, + ) class DiffractionDataset(Sequence[DiffractionArray], ABC): - """A sequence of DiffractionArrays with shared metadata and bad-pixel mask.""" + """A sequence of DiffractionArrays with shared metadata and bad-pixel mask. + + Every dataset owns a bad-pixel mask; callers can always rely on + ``get_bad_pixels()`` returning a real array, defaulting to all-good pixels + when the source data did not include a mask. + """ @abstractmethod def get_metadata(self) -> DiffractionMetadata: @@ -233,7 +256,7 @@ def get_layout(self) -> SimpleTreeNode: pass @abstractmethod - def get_bad_pixels(self) -> BadPixels | None: + def get_bad_pixels(self) -> BadPixels: pass @@ -251,6 +274,17 @@ def __init__( self._metadata = metadata self._contents_tree = contents_tree self._array_list = array_list + + extent = metadata.detector_extent + + if bad_pixels is None: + bad_pixels = numpy.zeros((extent.height_px, extent.width_px), dtype=numpy.bool_) + elif bad_pixels.shape != extent.get_shape(): + raise ValueError( + f'Bad pixels shape {bad_pixels.shape} does not match ' + f'detector extent {extent.get_shape()}.' + ) + self._bad_pixels = bad_pixels @classmethod @@ -266,7 +300,7 @@ def get_metadata(self) -> DiffractionMetadata: def get_layout(self) -> SimpleTreeNode: return self._contents_tree - def get_bad_pixels(self) -> BadPixels | None: + def get_bad_pixels(self) -> BadPixels: return self._bad_pixels @overload diff --git a/src/ptychodus/api/diffraction_prep.py b/src/ptychodus/api/diffraction_prep.py new file mode 100644 index 000000000..5dc98cc36 --- /dev/null +++ b/src/ptychodus/api/diffraction_prep.py @@ -0,0 +1,246 @@ +"""Diffraction pattern preprocessing pipeline. + +Steps share a single `apply(data) -> ndarray` interface and infer whether the +input is a 3-D pattern stack or a 2-D boolean bad-pixel mask from +`data.dtype`. The order of operations is encoded in +`DiffractionPrepPipeline.steps`; the canonical order emitted by the model-layer +factory is filter → crop → binning → padding → hflip → vflip → transpose. +""" + +from __future__ import annotations +from abc import abstractmethod +from typing import Annotated, Literal, TypeAlias + +import numpy +from pydantic import BaseModel, ConfigDict, Discriminator, Field + +from .diffraction import ( + BadPixels, + CropCenter, + DiffractionArray, + DiffractionPatterns, + SimpleDiffractionArray, +) +from .geometry import ImageExtent, PixelGeometry + + +def _is_mask(data: numpy.ndarray) -> bool: + return data.dtype == numpy.bool_ + + +class DiffractionPrepStep(BaseModel): + """Abstract base for a single preprocessing step. + + To add a new step: + + - Declare `type: Literal['<unique_tag>'] = '<unique_tag>'`. The default is required so + callers never pass the tag and `model_dump_json` emits it automatically. + - Add the class to `DiffractionPrepStepUnion`, or it is unreachable on deserialize. + - Override `apply_to_extent` / `apply_to_pixel_geometry` only where the step is not + identity in that dimension; the defaults below pass the input through unchanged. + - Branch on `_is_mask(data)` in `apply` if pattern and mask behavior differ, as + `BinningStep` does (sum for patterns, logical-AND for masks). + """ + + model_config = ConfigDict(frozen=True) + + @abstractmethod + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + """Apply this step. Mask vs pattern behavior is inferred from `data.dtype`.""" + + def apply_to_extent(self, extent: ImageExtent) -> ImageExtent: + """Return the extent this step would produce given `extent`. Default: identity.""" + return extent + + def apply_to_pixel_geometry(self, geometry: PixelGeometry) -> PixelGeometry: + """Return the pixel geometry this step would produce. Default: identity.""" + return geometry + + +class FilterValuesStep(DiffractionPrepStep): + """Zero pattern values outside `[lower_bound, upper_bound)`. No-op on masks.""" + + type: Literal['filter_values'] = 'filter_values' + lower_bound: int | None = None + upper_bound: int | None = None + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + if _is_mask(data): + return data + + if self.lower_bound is None and self.upper_bound is None: + return data + + out = data.copy() + + if self.lower_bound is not None: + out[out < self.lower_bound] = 0 + + if self.upper_bound is not None: + out[out >= self.upper_bound] = 0 + + return out + + +class CropStep(DiffractionPrepStep): + """Center-crop the last two axes to `extent` about `center` (pixel coords).""" + + type: Literal['crop'] = 'crop' + center: CropCenter + extent: ImageExtent + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + radius_x = self.extent.width_px // 2 + slice_x = slice(self.center.position_x_px - radius_x, self.center.position_x_px + radius_x) + radius_y = self.extent.height_px // 2 + slice_y = slice(self.center.position_y_px - radius_y, self.center.position_y_px + radius_y) + leading = (slice(None),) * (data.ndim - 2) + return data[(*leading, slice_y, slice_x)] + + def apply_to_extent(self, extent: ImageExtent) -> ImageExtent: + return self.extent + + +class BinningStep(DiffractionPrepStep): + """Reduce each `bin_size_y × bin_size_x` block. Sum for patterns; logical-AND for masks.""" + + type: Literal['binning'] = 'binning' + bin_size_x: int = Field(gt=0) + bin_size_y: int = Field(gt=0) + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + binned_height = data.shape[-2] // self.bin_size_y + binned_width = data.shape[-1] // self.bin_size_x + shape = data.shape[:-2] + (binned_height, self.bin_size_y, binned_width, self.bin_size_x) + reshaped = data.reshape(shape) + if _is_mask(data): + return numpy.logical_and.reduce(reshaped, axis=(-3, -1), keepdims=False) + return numpy.sum(reshaped, axis=(-3, -1), keepdims=False) + + def apply_to_extent(self, extent: ImageExtent) -> ImageExtent: + return ImageExtent( + width_px=extent.width_px // self.bin_size_x, + height_px=extent.height_px // self.bin_size_y, + ) + + def apply_to_pixel_geometry(self, geometry: PixelGeometry) -> PixelGeometry: + return PixelGeometry( + width_m=geometry.width_m * self.bin_size_x, + height_m=geometry.height_m * self.bin_size_y, + ) + + +class PaddingStep(DiffractionPrepStep): + """Symmetrically pad the last two axes. Fill 0 for patterns; False for masks.""" + + type: Literal['padding'] = 'padding' + pad_x: int = Field(ge=0) + pad_y: int = Field(ge=0) + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + leading_pad = ((0, 0),) * (data.ndim - 2) + pad_width = (*leading_pad, (self.pad_y, self.pad_y), (self.pad_x, self.pad_x)) + fill = False if _is_mask(data) else 0 + return numpy.pad(data, pad_width, mode='constant', constant_values=fill) + + def apply_to_extent(self, extent: ImageExtent) -> ImageExtent: + return ImageExtent( + width_px=extent.width_px + 2 * self.pad_x, + height_px=extent.height_px + 2 * self.pad_y, + ) + + +class HorizontalFlipStep(DiffractionPrepStep): + """Flip the last axis.""" + + type: Literal['hflip'] = 'hflip' + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + return numpy.flip(data, axis=-1) + + +class VerticalFlipStep(DiffractionPrepStep): + """Flip the second-to-last axis.""" + + type: Literal['vflip'] = 'vflip' + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + return numpy.flip(data, axis=-2) + + +class TransposeStep(DiffractionPrepStep): + """Swap the last two axes.""" + + type: Literal['transpose'] = 'transpose' + + def apply(self, data: numpy.ndarray) -> numpy.ndarray: + axes = tuple(range(data.ndim - 2)) + (data.ndim - 1, data.ndim - 2) + return numpy.transpose(data, axes=axes) + + def apply_to_extent(self, extent: ImageExtent) -> ImageExtent: + return ImageExtent(width_px=extent.height_px, height_px=extent.width_px) + + def apply_to_pixel_geometry(self, geometry: PixelGeometry) -> PixelGeometry: + return PixelGeometry(width_m=geometry.height_m, height_m=geometry.width_m) + + +# `Discriminator('type')` is what keeps the field-less steps distinguishable: HorizontalFlipStep, +# VerticalFlipStep, and TransposeStep all serialize to the same empty JSON object, so shape-based +# union resolution would silently deserialize one as another. The explicit tag is mandatory here, +# not stylistic. +DiffractionPrepStepUnion: TypeAlias = Annotated[ + FilterValuesStep + | CropStep + | BinningStep + | PaddingStep + | HorizontalFlipStep + | VerticalFlipStep + | TransposeStep, + Discriminator('type'), +] + + +class DiffractionPrepPipeline(BaseModel): + """Ordered chain of preprocessing steps applied to patterns and the bad-pixel mask.""" + + model_config = ConfigDict(frozen=True) + + steps: tuple[DiffractionPrepStepUnion, ...] = () + + def _apply(self, data: numpy.ndarray) -> numpy.ndarray: + for step in self.steps: + data = step.apply(data) + return data + + def apply_to_patterns(self, patterns: DiffractionPatterns) -> DiffractionPatterns: + """Run every step over a pattern stack. A single 2-D pattern is promoted to a stack.""" + if patterns.ndim == 2: + patterns = patterns[numpy.newaxis, ...] + elif patterns.ndim != 3: + raise ValueError(f'Invalid diffraction pattern dimensions! (shape={patterns.shape})') + + return self._apply(patterns) + + def apply_to_mask(self, bad_pixels: BadPixels) -> BadPixels: + """Run every step over a 2-D boolean bad-pixel mask.""" + if bad_pixels.ndim != 2: + raise ValueError(f'Invalid bad_pixel dimensions! (shape={bad_pixels.shape})') + + return self._apply(bad_pixels) + + def __call__(self, array: DiffractionArray) -> DiffractionArray: + """Return a new array with the pipeline applied, preserving label and scan indexes.""" + patterns = self.apply_to_patterns(array.get_patterns()) + return SimpleDiffractionArray(array.get_label(), array.get_indexes(), patterns) + + def compute_output_extent(self, extent: ImageExtent) -> ImageExtent: + """Return the extent the pipeline would produce, without touching pattern data.""" + for step in self.steps: + extent = step.apply_to_extent(extent) + return extent + + def compute_output_pixel_geometry(self, geometry: PixelGeometry) -> PixelGeometry: + """Return the pixel geometry the pipeline would produce, without touching pattern data.""" + for step in self.steps: + geometry = step.apply_to_pixel_geometry(geometry) + return geometry diff --git a/src/ptychodus/api/geometry.py b/src/ptychodus/api/geometry.py index eb3a25df1..1b1799bd4 100644 --- a/src/ptychodus/api/geometry.py +++ b/src/ptychodus/api/geometry.py @@ -52,6 +52,10 @@ class PixelGeometry: def is_square(self) -> bool: return self.width_m == self.height_m + @property + def is_valid(self) -> bool: + return self.width_m > 0.0 and self.height_m > 0.0 + def get_area_m2(self) -> float: return self.width_m * self.height_m diff --git a/src/ptychodus/api/io.py b/src/ptychodus/api/io.py index b95722cd3..331115296 100644 --- a/src/ptychodus/api/io.py +++ b/src/ptychodus/api/io.py @@ -8,7 +8,8 @@ import h5py import numpy -from .diffraction import zero_bad_pixels +from .diffraction import Polarization, zero_bad_pixels +from .fluorescence import ElementMap, FluorescenceDataset from .geometry import PixelGeometry from .object import Object, ObjectCenter from .probe import ProbeSequence @@ -17,10 +18,13 @@ from .reconstructor import AssembledDiffractionData, LossValue, ReconstructInput __all__ = [ + 'FluorescenceFileKeys', 'StandardFileLayout', 'load_diffraction_data', + 'load_fluorescence_data', 'load_product', 'save_diffraction_data', + 'save_fluorescence_data', 'save_product', 'save_ptychopinn_training_data', ] @@ -32,6 +36,7 @@ class StandardFileLayout(StrEnum): """Conventional file names used in the ptychodus standard HDF5 workflow directory.""" DIFFRACTION = 'diffraction.h5' + FLUORESCENCE = 'fluorescence.h5' FLUORESCENCE_IN = 'fluorescence-in.h5' FLUORESCENCE_OUT = 'fluorescence-out.h5' # Stem only; the per-backend extension comes from @@ -121,6 +126,8 @@ class ProductFileKeys(StrEnum): EXPOSURE_TIME = 'exposure_time_s' MASS_ATTENUATION = 'mass_attenuation_m2_kg' TOMOGRAPHY_ANGLE = 'tomography_angle_deg' + TILT_ANGLE = 'tilt_angle_deg' + POLARIZATION = 'polarization' PROBE_ARRAY = 'probe' OPR_WEIGHTS = 'opr_weights' PROBE_PIXEL_HEIGHT = 'pixel_height_m' @@ -149,6 +156,21 @@ def load_product(file: Path) -> Product: exposure_time_s = float(h5_file.attrs.get(ProductFileKeys.EXPOSURE_TIME, 0.0)) mass_attenuation_m2_kg = float(h5_file.attrs.get(ProductFileKeys.MASS_ATTENUATION, 0.0)) tomography_angle_deg = float(h5_file.attrs.get(ProductFileKeys.TOMOGRAPHY_ANGLE, 0.0)) + tilt_angle_deg = float(h5_file.attrs.get(ProductFileKeys.TILT_ANGLE, 0.0)) + + polarization: Polarization | None = None + if ProductFileKeys.POLARIZATION in h5_file.attrs: + raw_polarization = h5_file.attrs[ProductFileKeys.POLARIZATION] + if isinstance(raw_polarization, bytes): + raw_polarization = raw_polarization.decode('utf-8', errors='replace') + try: + polarization = Polarization(str(raw_polarization)) + except ValueError: + logger.warning( + 'Unknown polarization %r in %s; setting polarization=None.', + raw_polarization, + file, + ) metadata = ProductMetadata( name=name, @@ -159,6 +181,8 @@ def load_product(file: Path) -> Product: exposure_time_s=exposure_time_s, mass_attenuation_m2_kg=mass_attenuation_m2_kg, tomography_angle_deg=tomography_angle_deg, + tilt_angle_deg=tilt_angle_deg, + polarization=polarization, ) h5_object = h5_file[ProductFileKeys.OBJECT_ARRAY] @@ -292,6 +316,10 @@ def save_product(file: Path, product: Product) -> None: h5_file.attrs[ProductFileKeys.PROBE_PHOTON_COUNT] = metadata.probe_photon_count h5_file.attrs[ProductFileKeys.EXPOSURE_TIME] = metadata.exposure_time_s h5_file.attrs[ProductFileKeys.MASS_ATTENUATION] = metadata.mass_attenuation_m2_kg + h5_file.attrs[ProductFileKeys.TOMOGRAPHY_ANGLE] = metadata.tomography_angle_deg + h5_file.attrs[ProductFileKeys.TILT_ANGLE] = metadata.tilt_angle_deg + if metadata.polarization is not None: + h5_file.attrs[ProductFileKeys.POLARIZATION] = metadata.polarization.value h5_file.create_dataset(ProductFileKeys.PROBE_POSITION_INDEXES, data=scan_indexes) h5_file.create_dataset(ProductFileKeys.PROBE_POSITION_X, data=scan_x_m) @@ -331,6 +359,86 @@ def save_product(file: Path, product: Product) -> None: h5_file.create_dataset(ProductFileKeys.LOSS_VALUES, data=loss_values) +class FluorescenceFileKeys(StrEnum): + """HDF5 paths for the primary XRF-Maps NNLS layout used by :func:`save_fluorescence_data`.""" + + COUNTS_PER_SECOND = '/MAPS/XRF_Analyzed/NNLS/Counts_Per_Sec' + CHANNEL_NAMES = '/MAPS/XRF_Analyzed/NNLS/Channel_Names' + + +_FLUORESCENCE_LOAD_PATHS: tuple[tuple[str, str], ...] = ( + (FluorescenceFileKeys.COUNTS_PER_SECOND, FluorescenceFileKeys.CHANNEL_NAMES), + ('/MAPS/XRF_Analyzed/Fitted/Counts_Per_Sec', '/MAPS/XRF_Analyzed/Fitted/Channel_Names'), + ('/MAPS/XRF_fits', '/MAPS/channel_names'), +) + + +def _locate_fluorescence_datasets(h5_file: h5py.File) -> tuple[h5py.Dataset, h5py.Dataset]: + for cps_path, names_path in _FLUORESCENCE_LOAD_PATHS: + cps = h5_file.get(cps_path) + names = h5_file.get(names_path) + if isinstance(cps, h5py.Dataset) and isinstance(names, h5py.Dataset): + return cps, names + + tried = ', '.join(cps for cps, _ in _FLUORESCENCE_LOAD_PATHS) + raise KeyError(f'No known fluorescence layout in file (tried: {tried}).') + + +def _split_h5_path(data_path: str) -> tuple[str, str]: + stripped = data_path.strip('/') + parts = stripped.rsplit('/', 1) + if len(parts) == 1: + return '/', parts[0] + return '/' + parts[0], parts[1] + + +def load_fluorescence_data(file: Path) -> FluorescenceDataset: + """Load a fluorescence dataset from an XRF-Maps HDF5 file. + + Recognises the v10 NNLS layout (``/MAPS/XRF_Analyzed/NNLS/...``, preferred), + the v10 iterative-matrix-fitting layout (``/MAPS/XRF_Analyzed/Fitted/...``), + and the legacy v9 layout (``/MAPS/XRF_fits`` + ``/MAPS/channel_names``). + """ + element_maps: list[ElementMap] = [] + + with h5py.File(file, 'r') as h5_file: + h5_counts_per_second, h5_channel_names = _locate_fluorescence_datasets(h5_file) + + counts_per_second = h5_counts_per_second[...] + channel_names = h5_channel_names[...] + + for bname, cps in zip(channel_names, counts_per_second): + if isinstance(bname, bytes): + name = bname.decode('utf-8', errors='replace') + else: + name = str(bname) + element_maps.append(ElementMap(name, cps)) + + counts_per_second_path = h5_counts_per_second.name + channel_names_path = h5_channel_names.name + + return FluorescenceDataset( + element_maps=element_maps, + counts_per_second_path=counts_per_second_path, + channel_names_path=channel_names_path, + ) + + +def save_fluorescence_data(file: Path, dataset: FluorescenceDataset) -> None: + """Write a fluorescence dataset to an HDF5 file at the paths the dataset carries.""" + counts_group_path, counts_ds_name = _split_h5_path(dataset.counts_per_second_path) + names_group_path, names_ds_name = _split_h5_path(dataset.channel_names_path) + + channel_names = [emap.name for emap in dataset.element_maps] + counts_per_second = [emap.counts_per_second for emap in dataset.element_maps] + + with h5py.File(file, 'w') as h5_file: + counts_group = h5_file.require_group(counts_group_path) + counts_group.create_dataset(counts_ds_name, data=numpy.stack(counts_per_second)) + names_group = h5_file.require_group(names_group_path) + names_group.create_dataset(names_ds_name, data=channel_names, dtype='S256') + + def save_ptychopinn_training_data( file_path: Path, parameters: ReconstructInput, diff --git a/src/ptychodus/api/metrics.py b/src/ptychodus/api/metrics.py index eb70ba5a0..1324bda4a 100644 --- a/src/ptychodus/api/metrics.py +++ b/src/ptychodus/api/metrics.py @@ -40,22 +40,21 @@ class ObjectComparison: 3. Flattens multi-layer objects via :meth:`Object.get_layers_flattened`. 4. Promotes both arrays to a common complex dtype. - Attributes: - reference_complex: 2D complex array, the reference object's flattened layers, - promoted to the common dtype. - test_complex: 2D complex array, the standardized + aligned test object's - flattened layers. Same shape and dtype as ``reference_complex``. - pixel_geometry: Shared pixel geometry of both reconstructions (validated - equal by the upstream primitives). - ambiguities: The ambiguities removed from the test side, useful as - provenance (e.g. for reporting "how much ramp/scale was removed" - alongside the metric value). """ reference_complex: ComplexArrayType + """2D complex array, the reference object's flattened layers, promoted to the common dtype.""" test_complex: ComplexArrayType + """2D complex array, the standardized + aligned test object's flattened layers. + + Same shape and dtype as ``reference_complex``. + """ pixel_geometry: PixelGeometry + """Shared pixel geometry of both reconstructions (validated equal by the upstream + primitives).""" ambiguities: ReconstructionAmbiguities + """The ambiguities removed from the test side, useful as provenance (e.g. for reporting + "how much ramp/scale was removed" alongside the metric value).""" @classmethod def from_products( @@ -133,6 +132,8 @@ def test_phase(self) -> RealArrayType: @dataclass(frozen=True) class FourierRingCorrelation: + """Per-ring Fourier ring correlation between two complex images, with resolution estimators.""" + spatial_frequency_per_m: RealArrayType correlation: RealArrayType pixels_per_ring: IntegerArrayType @@ -491,32 +492,33 @@ class ReconstructionResiduals: as errors in predicted intensity. These maps therefore quantify detector-domain data-fit quality and are *not* phase-blind. - Attributes: - real_space_error_map: 2D array on the object grid. For each object pixel, the - probe-footprint-weighted aggregation of per-frame amplitude residuals over every - frame whose probe touched that pixel, divided by the same probe-weighted - aggregation of measured amplitudes. Each frame's probe-intensity patch is - normalized to sum to 1 before splatting, so frames contribute equally regardless - of probe power (relevant for variable-probe reconstructions). Scan-density - invariant: doubling the number of frames covering a pixel doubles both splats, - leaving the ratio unchanged. NaN where no R-factor is defined: un-illuminated - pixels (no frame contributed) and object regions touched only by frames with zero - measured signal (``Σ √I_meas = 0``). - object_pixel_geometry: Pixel geometry of ``real_space_error_map``. - object_center: Real-space origin of ``real_space_error_map``. - reciprocal_space_error_map: 2D array on the detector grid. Each pixel is - ``Σ_n |√I_meas,n − √I_pred,n| / Σ_n √I_meas,n``, summed across frames. NaN where - no R-factor is defined: bad pixels and detector pixels with no measured signal - across any frame. - detector_pixel_geometry: Pixel geometry of ``reciprocal_space_error_map`` (derived from - the forward propagator). """ real_space_error_map: RealArrayType + """2D array on the object grid. + + For each object pixel, the probe-footprint-weighted aggregation of per-frame amplitude + residuals over every frame whose probe touched that pixel, divided by the same + probe-weighted aggregation of measured amplitudes. Each frame's probe-intensity patch is + normalized to sum to 1 before splatting, so frames contribute equally regardless of probe + power (relevant for variable-probe reconstructions). Scan-density invariant: doubling the + number of frames covering a pixel doubles both splats, leaving the ratio unchanged. NaN + where no R-factor is defined: un-illuminated pixels (no frame contributed) and object + regions touched only by frames with zero measured signal (``Σ √I_meas = 0``). + """ object_pixel_geometry: PixelGeometry + """Pixel geometry of ``real_space_error_map``.""" object_center: ObjectCenter + """Real-space origin of ``real_space_error_map``.""" reciprocal_space_error_map: RealArrayType + """2D array on the detector grid. + + Each pixel is ``Σ_n |√I_meas,n − √I_pred,n| / Σ_n √I_meas,n``, summed across frames. NaN + where no R-factor is defined: bad pixels and detector pixels with no measured signal across + any frame. + """ detector_pixel_geometry: PixelGeometry + """Pixel geometry of ``reciprocal_space_error_map`` (derived from the forward propagator).""" def compute_reconstruction_residuals( diff --git a/src/ptychodus/api/plugins.py b/src/ptychodus/api/plugins.py index 929ce65b9..5d1801634 100644 --- a/src/ptychodus/api/plugins.py +++ b/src/ptychodus/api/plugins.py @@ -27,7 +27,7 @@ ) from .object import ObjectFileReader, ObjectFileWriter, Object from .observer import Observable, Observer -from .parametric import StringParameter +from .parametric import Parameter, StringParameter from .probe import ProbeFileReader, ProbeFileWriter, ProbeSequence from .probe_gen import FresnelZonePlate from .probe_positions import ( @@ -40,6 +40,7 @@ __all__ = [ 'PluginChooser', + 'PluginChooserParameter', 'PluginRegistry', ] @@ -93,14 +94,18 @@ class Plugin(Generic[T]): display_name: str -class PluginChooser(Iterable[Plugin[T]], Observable, Observer): - """Observable list of typed plugins with a tracked current selection.""" +class PluginChooser(Iterable[Plugin[T]], Observable): + """Observable list of typed plugins with a tracked current selection. + + The chooser knows nothing about settings persistence. Bind it to a settings + parameter by constructing a :class:`PluginChooserParameter` over it, which owns + the translation between the two name spaces in both directions. + """ def __init__(self) -> None: super().__init__() self._registered_plugins: list[Plugin[T]] = list() self._current_index = 0 - self._parameter: StringParameter | None = None def stringify_plugin_names(self) -> str: """Return a sorted, comma-separated list of registered plugin simple names.""" @@ -111,41 +116,66 @@ def register_plugin(self, strategy: T, *, display_name: str, simple_name: str = if not simple_name: simple_name = re.sub(r'\W+', '', display_name) - plugin = Plugin[T](strategy, simple_name, display_name) - self._registered_plugins.append(plugin) + # The list is kept sorted by display name, so a registration can move the + # selected plugin to a different index. Track it by identity across the sort + # rather than by value: Plugin is a frozen dataclass, so equality compares the + # strategy field and an array-valued strategy would break the comparison. + current = ( + self._registered_plugins[self._current_index] if self._registered_plugins else None + ) + + self._registered_plugins.append(Plugin[T](strategy, simple_name, display_name)) self._registered_plugins.sort(key=lambda x: x.display_name) - self.notify_observers() - def get_current_plugin(self) -> Plugin[T]: - """Return the currently selected plugin.""" - if not self._registered_plugins: - raise LookupError('No plugins registered') - return self._registered_plugins[self._current_index] + if current is not None: + self._current_index = next( + index for index, plugin in enumerate(self._registered_plugins) if plugin is current + ) - def set_current_plugin(self, name: str) -> None: - """Select the plugin matching *name* (case-insensitive simple or display name); warn if none match.""" + self.notify_observers() + + def _find_index(self, name: str) -> int | None: namecf = name.casefold() for index, plugin in enumerate(self._registered_plugins): if namecf == plugin.simple_name.casefold() or namecf == plugin.display_name.casefold(): - if self._current_index != index: - self._current_index = index - - if self._parameter is not None: - self._parameter.set_value(self.get_current_plugin().simple_name) + return index - self.notify_observers() + return None - return + def find_plugin(self, name: str) -> Plugin[T] | None: + """Return the plugin matching *name* (case-insensitive simple or display name), or None.""" + index = self._find_index(name) + return None if index is None else self._registered_plugins[index] - registered_plugins = ', '.join(f'"{pi.simple_name}"' for pi in self._registered_plugins) - logger.warning(f'Invalid plugin name "{name}". Registered plugins: {registered_plugins}.') + def get_current_plugin(self) -> Plugin[T]: + """Return the currently selected plugin.""" + if not self._registered_plugins: + raise LookupError('No plugins registered') + return self._registered_plugins[self._current_index] - def synchronize_with_parameter(self, parameter: StringParameter) -> None: - """Bind selection to *parameter*: the chooser tracks the parameter's value and vice versa.""" - self._parameter = parameter - self.set_current_plugin(parameter.get_value()) - self._parameter.add_observer(self) + def set_current_plugin(self, name: str) -> None: + """Select the plugin matching *name* (case-insensitive simple or display name). + + An unrecognized name logs a warning and leaves the selection unchanged, but + still notifies observers so that a bound view resynchronizes to the selection + the chooser actually holds. + """ + index = self._find_index(name) + + if index is None: + registered_plugins = ', '.join( + f'"{plugin.simple_name}"' for plugin in self._registered_plugins + ) + logger.warning( + f'Invalid plugin name "{name}". Registered plugins: {registered_plugins}.' + ) + self.notify_observers() + return + + if index != self._current_index: + self._current_index = index + self.notify_observers() def __iter__(self) -> Iterator[Plugin[T]]: for plugin in self._registered_plugins: @@ -154,9 +184,128 @@ def __iter__(self) -> Iterator[Plugin[T]]: def __bool__(self) -> bool: return bool(self._registered_plugins) + +class PluginChooserParameter(Parameter[str], Observer, Generic[T]): + """Parameter[str] view of a PluginChooser whose value space is plugin display names. + + A chooser has two name spaces: the human-readable ``display_name`` shown in the + GUI and the ``simple_name`` persisted to settings. This adapter is the single + place that translates between them. Its own value space is display names, so a + plain combo box bound to it round trips correctly. + + Passing *settings* additionally binds the chooser to that settings parameter: the + persisted name selects a plugin at construction and is rewritten to its canonical + simple name, and every later selection change writes the simple name back. A + persisted name that matches no registered plugin is left alone rather than + reconciled, because a plugin can be missing merely because its optional + dependency failed to import this run. + + Note that :meth:`get_value` raises ``LookupError`` when no plugins are registered, + so this must not be read before ``PluginRegistry.load_plugins()`` has run. + """ + + def __init__(self, chooser: PluginChooser[T], settings: StringParameter | None = None) -> None: + super().__init__() + self._chooser = chooser + self._settings = settings + self._selected: Plugin[T] | None = None + self._suppress_notify = False + + if settings is not None: + self._apply_settings() + settings.add_observer(self) + + self._selected = chooser.get_current_plugin() if chooser else None + chooser.add_observer(self) + + def choices(self) -> Iterator[str]: + """Yield the display names to populate a combo box with.""" + for plugin in self._chooser: + yield plugin.display_name + + def get_strategy(self) -> T: + return self._chooser.get_current_plugin().strategy + + def get_value(self) -> str: + return self._chooser.get_current_plugin().display_name + + def set_value(self, value: str, *, notify: bool = True) -> None: + # The chooser notifies us back through _update; suppress that relay rather + # than notifying here, so a no-op selection stays silent either way. Only the + # relay is suppressed: persistence must not depend on a view flag. + self._suppress_notify = not notify + + try: + self._chooser.set_current_plugin(value) + finally: + self._suppress_notify = False + + def get_value_as_string(self) -> str: + return self.get_value() + + def set_value_from_string(self, value: str) -> None: + self.set_value(value) + + def copy(self) -> Parameter[str]: + """Return an unbound view of the same chooser; the copy does not persist.""" + return PluginChooserParameter(self._chooser) + + def _apply_settings(self) -> None: + settings = self._settings + + if settings is None: + return + + name = settings.get_value() + plugin = self._chooser.find_plugin(name) + self._chooser.set_current_plugin(name) + + if plugin is not None: + # Unconditional, so that a display name normalizes to its simple name even + # when it already resolves to the current selection. ParameterBase guards + # against no-op writes, so this does not notify unless the value changed. + settings.set_value(plugin.simple_name) + + def _reconcile(self) -> None: + """Settle the chooser and the settings parameter against each other. + + Observable notifications carry no payload, so the cached selection is what + distinguishes "the selection moved" from "a plugin was registered". Only the + former may write back; the latter must leave an unresolved setting intact. + """ + settings = self._settings + plugin = self._chooser.get_current_plugin() if self._chooser else None + + if plugin is None or settings is None: + self._selected = plugin + return + + if self._selected is None: + # The chooser just gained its first plugin: the selection came into + # existence rather than moving, so the fallback must not be persisted. + self._selected = plugin + + if plugin is not self._selected: + self._selected = plugin + settings.set_value(plugin.simple_name) + return + + # The selection held, so this was a registration. If it has just made the + # persisted name resolvable, honor it now — set_current_plugin re-enters + # here through the chooser's notification to finish the write-back. + desired = self._chooser.find_plugin(settings.get_value()) + + if desired is not None and desired is not plugin: + self._apply_settings() + def _update(self, observable: Observable) -> None: - if self._parameter is not None and observable is self._parameter: - self.set_current_plugin(self._parameter.get_value()) + if observable is self._settings: + self._apply_settings() + elif observable is self._chooser: + self._reconcile() + + if not self._suppress_notify: + self.notify_observers() class PluginRegistry: diff --git a/src/ptychodus/api/probe.py b/src/ptychodus/api/probe.py index 58b81ec76..508134522 100644 --- a/src/ptychodus/api/probe.py +++ b/src/ptychodus/api/probe.py @@ -66,6 +66,8 @@ class ProbeEntropyMetrics: @dataclass(frozen=True) class ProbeSizeMetrics: + """Probe size metrics: principal-axis tilt, FWHM and RMS extents, and encircled-energy diameter.""" + major_axis_tilt_rad: float minor_axis_tilt_rad: float @@ -468,7 +470,7 @@ def get_intensity(self) -> RealArrayType: return numpy.sum(intensity(self._array), axis=-3) def get_power_spectrum(self) -> RealArrayType: - """Incoherent-sum power spectrum |FFT(psi)|^2 over the mode axis. + """Incoherent-sum power spectrum ``|FFT(psi)|^2`` over the mode axis. No fftshift is applied: Shannon entropy is permutation-invariant, so the frequency ordering is irrelevant for entropy calculations. @@ -603,7 +605,9 @@ def __getitem__(self, index: slice) -> Sequence[Probe]: ... def __getitem__(self, index: int | slice) -> Probe | Sequence[Probe]: if isinstance(index, slice): - return [self[idx] for idx in range(index.start, index.stop, index.step)] + # slice.indices normalizes implicit bounds, negative indexes, and + # out-of-range values against the sequence length. + return [self[idx] for idx in range(*index.indices(len(self)))] array = self._array[0, :, :, :].copy() diff --git a/src/ptychodus/api/probe_positions.py b/src/ptychodus/api/probe_positions.py index acf23e7e1..6f17cd456 100644 --- a/src/ptychodus/api/probe_positions.py +++ b/src/ptychodus/api/probe_positions.py @@ -76,7 +76,13 @@ def __getitem__(self, index: slice) -> Sequence[ProbePosition]: ... def __getitem__(self, index: int | slice) -> ProbePosition | Sequence[ProbePosition]: if isinstance(index, slice): - return [self[idx] for idx in range(index.start, index.stop, index.step)] + # Slice the backing arrays directly rather than materializing a + # ProbePosition per element. Basic numpy slicing returns views; that + # is safe because this class exposes no mutators. + seq = ProbePositionSequence() + seq._indexes = self._indexes[index] + seq._coordinates_m = self._coordinates_m[index, :] + return seq return ProbePosition( index=self._indexes[index], diff --git a/src/ptychodus/api/product.py b/src/ptychodus/api/product.py index 85b28276f..f5bc19c2c 100644 --- a/src/ptychodus/api/product.py +++ b/src/ptychodus/api/product.py @@ -7,6 +7,7 @@ from sys import getsizeof from .common import ELECTRON_VOLT_J, PLANCK_CONSTANT_J_PER_HZ, LIGHT_SPEED_M_PER_S +from .diffraction import Polarization from .object import Object from .probe import Probe, ProbeSequence from .probe_positions import ProbePosition, ProbePositionSequence @@ -24,6 +25,8 @@ class ProductMetadata: exposure_time_s: float mass_attenuation_m2_kg: float tomography_angle_deg: float + tilt_angle_deg: float = 0.0 + polarization: Polarization | None = None @property def probe_energy_J(self) -> float: # noqa: N802 @@ -48,6 +51,8 @@ def nbytes(self) -> int: sz += getsizeof(self.exposure_time_s) sz += getsizeof(self.mass_attenuation_m2_kg) sz += getsizeof(self.tomography_angle_deg) + sz += getsizeof(self.tilt_angle_deg) + sz += getsizeof(self.polarization) return sz diff --git a/src/ptychodus/api/propagator.py b/src/ptychodus/api/propagator.py index 84a044226..fcd22df80 100644 --- a/src/ptychodus/api/propagator.py +++ b/src/ptychodus/api/propagator.py @@ -12,12 +12,14 @@ def intensity(wavefield: ComplexArrayType) -> RealArrayType: - """Return the element-wise intensity (|wavefield|²) of a complex array.""" + """Return the element-wise intensity (``|wavefield|²``) of a complex array.""" return numpy.square(numpy.absolute(wavefield)) @dataclass(frozen=True) class PropagatorParameters: + """Geometric parameters for a wavefield propagator: wavelength, extent, pixel size, and distance.""" + wavelength_m: float """Illumination wavelength in meters.""" width_px: int diff --git a/src/ptychodus/api/reconstructor.py b/src/ptychodus/api/reconstructor.py index 883563024..3dc74c400 100644 --- a/src/ptychodus/api/reconstructor.py +++ b/src/ptychodus/api/reconstructor.py @@ -30,11 +30,12 @@ @dataclass(frozen=True) class ReconstructInput: - """All data required to start a reconstruction: patterns, bad-pixel mask, and initial product.""" + """All data required to start a reconstruction: patterns, bad-pixel mask, initial product, and detector pixel geometry.""" diffraction_patterns: DiffractionPatterns bad_pixels: BadPixels product: Product + pixel_geometry: PixelGeometry @dataclass(frozen=True) @@ -263,6 +264,12 @@ def get_pattern(self, index: int) -> DiffractionPattern: def get_pixel_geometry(self) -> PixelGeometry: return self._pixel_geometry + def set_pixel_geometry(self, pixel_geometry: PixelGeometry) -> None: + # Views produced by assemble() keep their creation-time snapshot; they are + # only used for per-array display (average pattern, counts) and not by + # reconstruction, so leaving them stale is acceptable. + self._pixel_geometry = pixel_geometry + def get_bad_pixels(self) -> BadPixels: return self._bad_pixels @@ -417,7 +424,7 @@ def prepare_reconstruct_input( losses=product.losses, ) - return ReconstructInput(patterns, self._bad_pixels, product) + return ReconstructInput(patterns, self._bad_pixels, product, self._pixel_geometry) def __str__(self) -> str: number, height, width = self._patterns.shape @@ -442,17 +449,16 @@ class ReconstructionAmbiguities: unchanged. See :meth:`standardize_product` for what is and is not exactly preserved. - Attributes: - object_scale_factor: ``scale`` above. Must be non-zero and finite. - phase_offset_rad: ``phi`` above, in radians. - phase_ramp_x_rad_per_m: ``k_x`` above, in rad/m. - phase_ramp_y_rad_per_m: ``k_y`` above, in rad/m. """ object_scale_factor: float + """``scale`` above. Must be non-zero and finite.""" phase_offset_rad: float + """``phi`` above, in radians.""" phase_ramp_x_rad_per_m: float + """``k_x`` above, in rad/m.""" phase_ramp_y_rad_per_m: float + """``k_y`` above, in rad/m.""" def __post_init__(self) -> None: for f in fields(self): diff --git a/src/ptychodus/api/visualization.py b/src/ptychodus/api/visualization.py index e5b49e0be..74036693a 100644 --- a/src/ptychodus/api/visualization.py +++ b/src/ptychodus/api/visualization.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from enum import Enum, auto from typing import Any -from typing import Final +from typing import Final, cast import logging from matplotlib.colors import Colormap, hsv_to_rgb @@ -21,12 +21,33 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class DisplayValues: + """Real-valued array actually shown by a renderer, with its axis label. + + A renderer that maps a complex array to a single scalar component contributes one entry; + the cylindrical color model, which encodes amplitude and phase simultaneously, contributes + two. + """ + + label: str + values: RealArrayType + + +@dataclass(frozen=True) +class LineCutSeries: + """1D profile of a single displayed component, with its axis label.""" + + label: str + value: Sequence[float] + + @dataclass(frozen=True) class LineCut: - """1D profile sampled along a line: distances in meters and corresponding pixel values.""" + """1D profile sampled along a line: distances in meters and one series per component.""" distance_m: Sequence[float] - value: Sequence[float | complex] + series: Sequence[LineCutSeries] @dataclass(frozen=True) @@ -50,6 +71,7 @@ def __init__( rgba: RealArrayType, pixel_geometry: PixelGeometry, color_value_range: Interval[float], + display_values: Sequence[DisplayValues] | None = None, ) -> None: if values.ndim != 2: raise ValueError(f'Values must be a 2-dimensional ndarray (actual={values.ndim}).') @@ -63,8 +85,28 @@ def __init__( if values.shape[0] != rgba.shape[0] or values.shape[1] != rgba.shape[1]: raise ValueError(f'Shape mismatch (values={values.shape} and rgba={rgba.shape}).') + if display_values is None: + # No renderer told us which component is on screen; fall back to amplitude for + # complex input and to the values themselves otherwise. + fallback = numpy.absolute(values) if numpy.iscomplexobj(values) else values + display_values = [DisplayValues(value_label, cast(RealArrayType, fallback))] + + for dv in display_values: + if dv.values.ndim != 2: + raise ValueError( + f'Display values "{dv.label}" must be a 2-dimensional ndarray ' + f'(actual={dv.values.ndim}).' + ) + + if dv.values.shape[0] != rgba.shape[0] or dv.values.shape[1] != rgba.shape[1]: + raise ValueError( + f'Shape mismatch (display values "{dv.label}"={dv.values.shape} ' + f'and rgba={rgba.shape}).' + ) + self._value_label = value_label self._values = values + self._display_values = display_values self._rgba = rgba self._pixel_width_m = pixel_geometry.width_m self._pixel_height_m = pixel_geometry.height_m @@ -77,6 +119,10 @@ def get_value_label(self) -> str: def get_values(self) -> NumberArrayType: return self._values + def get_display_values(self) -> Sequence[DisplayValues]: + """Return the real-valued arrays the renderer put on screen, in display order.""" + return self._display_values + def get_image_rgba(self) -> RealArrayType: return self._rgba @@ -150,7 +196,7 @@ def get_info_text(self, x: float, y: float) -> str: iy = min(iy, self._values.shape[-2]) value = self._values[iy, ix] - if numpy.iscomplex(value): + if numpy.iscomplexobj(value): amplitude = numpy.absolute(value) phase = numpy.angle(value) return f'{x=:.1f} {y=:.1f} {amplitude=:6g} {phase=:6g}' @@ -165,34 +211,41 @@ def get_line_cut(self, line: Line2D) -> LineCut: line_length = numpy.hypot(dx, dy) distances: list[float] = list() - values: list[float] = list() + values: list[list[float]] = [list() for dv in self._display_values] for alpha_l, alpha_r in zip(intersections[:-1], intersections[1:]): alpha = (alpha_l + alpha_r) / 2.0 point = line.lerp(alpha) - value = self._values[int(point.y), int(point.x)] + iy = int(point.y) + ix = int(point.x) distances.append(alpha * line_length) - values.append(value) - return LineCut(distances, values) + for value_list, dv in zip(values, self._display_values): + value_list.append(dv.values[iy, ix].item()) + + series = [ + LineCutSeries(dv.label, value_list) + for dv, value_list in zip(self._display_values, values) + ] + return LineCut(distances, series) def estimate_kernel_density(self, box: Box2D) -> KernelDensityEstimate: - x_range = Interval[int](0, self._values.shape[-1]) + # A histogram has a single value axis, so only the primary displayed component is + # estimated; the cylindrical color model contributes amplitude first for this reason. + display_values = self._display_values[0].values + + x_range = Interval[int](0, display_values.shape[-1]) x_begin = x_range.clamp(int(box.x_begin)) x_end = x_range.clamp(int(box.x_end) + 1) - y_range = Interval[int](0, self._values.shape[-2]) + y_range = Interval[int](0, display_values.shape[-2]) y_begin = y_range.clamp(int(box.y_begin)) y_end = y_range.clamp(int(box.y_end) + 1) - values = self._values[..., y_begin:y_end, x_begin:x_end] + values = display_values[..., y_begin:y_end, x_begin:x_end] values = values.reshape(values.shape[-3], -1) if values.ndim > 2 else values.reshape(-1) - if numpy.iscomplexobj(values): - # TODO improve KDE for complex values - values = numpy.absolute(values) - return KernelDensityEstimate(values.min(), values.max(), gaussian_kde(values)) @@ -290,7 +343,7 @@ def cyclic_colormap_names() -> Iterator[str]: def get_colormap_by_name(name: str) -> Colormap: - """Return the colorcet Colormap for the given short name (prefixed with 'cet_').""" + """Return the colorcet Colormap for the given short name (prefixed with ``cet_``).""" return matplotlib.colormaps[f'cet_{name}'] @@ -344,12 +397,14 @@ def visualize_real_values( values_transformed, value_min=value_min, value_max=value_max, clip=clip ) cmap = colormap if isinstance(colormap, Colormap) else get_colormap_by_name(colormap) + decorated_label = transform.decorate_text(value_label) return VisualizationProduct( - value_label=transform.decorate_text(value_label), + value_label=decorated_label, values=values, rgba=cmap(values_normalized), pixel_geometry=pixel_geometry, color_value_range=color_value_range, + display_values=[DisplayValues(decorated_label, values_transformed)], ) @@ -381,6 +436,7 @@ def visualize_complex_component( rgba=product.get_image_rgba(), pixel_geometry=product.get_pixel_geometry(), color_value_range=product.get_color_value_range(), + display_values=product.get_display_values(), ) @@ -488,10 +544,15 @@ def visualize_complex_values( ) phase_rad = ComplexComponent.PHASE_RAD.extract_component(values) hue = (phase_rad + numpy.pi) / (2 * numpy.pi) + amplitude_label = amplitude_transform.decorate_text(amplitude_component.name.title()) return VisualizationProduct( - value_label=amplitude_transform.decorate_text(amplitude_component.name.title()), + value_label=amplitude_label, values=values, rgba=model.render_rgba(hue, amplitude_normalized), pixel_geometry=pixel_geometry, color_value_range=color_value_range, + display_values=[ + DisplayValues(amplitude_label, amplitude_transformed), + DisplayValues('Phase [rad]', phase_rad), + ], ) diff --git a/src/ptychodus/api/workflow.py b/src/ptychodus/api/workflow.py index b0f629157..dbed7f22c 100644 --- a/src/ptychodus/api/workflow.py +++ b/src/ptychodus/api/workflow.py @@ -7,7 +7,7 @@ from enum import Enum, auto from typing import Any -from ptychodus.api.diffraction import CropCenter +from ptychodus.api.diffraction import CropCenter, Polarization from ptychodus.api.geometry import AffineTransform, ImageExtent from ptychodus.api.product import Product from ptychodus.api.reconstructor import AssembledDiffractionData, ReconstructInput @@ -15,13 +15,20 @@ class RemoteComputeProvider(Enum): + """Supported remote-compute providers for dispatching workflows off-machine.""" + GLOBUS = auto() GENESIS = auto() -class WorkflowDiffractionAPI(ABC): +class DiffractionWorkflowAPI(ABC): """Act on an assembled diffraction dataset within a workflow.""" + @abstractmethod + def get_dataset_index(self) -> int: + """Return the repository index this handle was bound to at load time.""" + pass + @abstractmethod def get_assembled_data(self) -> AssembledDiffractionData: """Return the assembled diffraction data.""" @@ -33,7 +40,7 @@ def save_assembled_data(self, file_path: Path) -> None: pass -class WorkflowProductAPI(ABC): +class ProductWorkflowAPI(ABC): """Act on a data product within a workflow.""" @abstractmethod @@ -110,7 +117,7 @@ def reconstruct_local( algorithm: str | None = None, output_product_file: Path | None = None, block: bool = False, - ) -> WorkflowProductAPI: + ) -> ProductWorkflowAPI: """Run reconstruction locally; returns a handle to the output product. Blocks until completion when block is True, otherwise returns immediately. @@ -183,11 +190,6 @@ def enhance_fluorescence_local( class WorkflowAPI(ABC): """Top-level API for loading data, managing products, and running reconstructions.""" - @abstractmethod - def load_bad_pixels(self, file_path: Path, *, file_type: str | None = None) -> None: - """Load a bad-pixel mask from file, uses format from settings when file_type is None.""" - pass - @abstractmethod def load_diffraction_data( self, @@ -196,29 +198,45 @@ def load_diffraction_data( file_type: str | None = None, crop_center: CropCenter | None = None, crop_extent: ImageExtent | None = None, - detector_extent: ImageExtent | None = None, + bad_pixels_file_path: Path | None = None, + bad_pixels_file_type: str | None = None, process_patterns: bool = True, block: bool = False, - ) -> WorkflowDiffractionAPI: + ) -> DiffractionWorkflowAPI: """Load and assemble a diffraction dataset from file. - Blocks until complete when block is True, otherwise returns immediately. + When bad_pixels_file_path is provided, its mask is applied to the loaded + dataset. Blocks until complete when block is True, otherwise returns immediately. """ pass @abstractmethod - def load_assembled_diffraction_data(self, file_path: Path) -> WorkflowDiffractionAPI: + def load_assembled_diffraction_data(self, file_path: Path) -> DiffractionWorkflowAPI: """Load a pre-assembled diffraction dataset from file.""" pass @abstractmethod - def register_product(self, product: Product) -> WorkflowProductAPI: - """Register an existing Product object and return a handle to it.""" + def register_product( + self, product: Product, *, diffraction: DiffractionWorkflowAPI | None = None + ) -> ProductWorkflowAPI: + """Register an existing Product object and return a handle to it. + + The product is bound to the given diffraction dataset at creation time. + """ pass @abstractmethod - def load_product(self, file_path: Path, *, file_type: str | None = None) -> WorkflowProductAPI: - """Load a product from file and return a handle to it.""" + def load_product( + self, + file_path: Path, + *, + file_type: str | None = None, + diffraction: DiffractionWorkflowAPI | None = None, + ) -> ProductWorkflowAPI: + """Load a product from file and return a handle to it. + + The product is bound to the given diffraction dataset at creation time. + """ pass @abstractmethod @@ -233,12 +251,18 @@ def create_product( exposure_time_s: float | None = None, mass_attenuation_m2_kg: float | None = None, tomography_angle_deg: float | None = None, - ) -> WorkflowProductAPI: - """Create a new product with optional metadata overrides and return a handle to it.""" + tilt_angle_deg: float | None = None, + polarization: Polarization | None = None, + diffraction: DiffractionWorkflowAPI | None = None, + ) -> ProductWorkflowAPI: + """Create a new product with optional metadata overrides and return a handle to it. + + The product is bound to the given diffraction dataset at creation time. + """ pass @abstractmethod - def get_product(self, product_index: int) -> WorkflowProductAPI: + def get_product(self, product_index: int) -> ProductWorkflowAPI: """Return a handle to an already-registered product by index.""" pass diff --git a/src/ptychodus/api/xmcd.py b/src/ptychodus/api/xmcd.py index ed6aee72b..ad792bdea 100644 --- a/src/ptychodus/api/xmcd.py +++ b/src/ptychodus/api/xmcd.py @@ -1,7 +1,7 @@ """XMCD (X-ray Magnetic Circular Dichroism) decomposition math. The decomposition models the helicity-dependent complex transmission of a -magnetic sample as +magnetic sample as:: O_+ (RCP) = O_struct * M O_- (LCP) = O_struct / M @@ -9,13 +9,13 @@ where ``O_struct`` is the helicity-independent (structural) transmission and ``M = |M| * exp(i * alpha)`` is the magnetic factor (carrying both absorption and phase contributions). Given the two reconstructed objects ``O_+`` and -``O_-``, three helicity products are formed: +``O_-``, three helicity products are formed:: parallel_helicity_product = O_+ * O_- = O_struct**2 parallel_helicity_ratio = O_+ / O_- = M**2 cross_helicity_product = O_+ * conj(O_-) = |O_struct|**2 * exp(2i * alpha) -from which the structural and magnetic objects are recovered as +from which the structural and magnetic objects are recovered as:: structural_object = sqrt(|parallel_helicity_product|) * exp(0.5i * angle(parallel_helicity_product)) magnetic_object = sqrt(|parallel_helicity_ratio|) * exp(0.5i * angle(cross_helicity_product)) diff --git a/src/ptychodus/cli.py b/src/ptychodus/cli/__init__.py similarity index 100% rename from src/ptychodus/cli.py rename to src/ptychodus/cli/__init__.py diff --git a/src/ptychodus/scripts/convert_to_ptychodus.py b/src/ptychodus/cli/convert_to_ptychodus.py similarity index 100% rename from src/ptychodus/scripts/convert_to_ptychodus.py rename to src/ptychodus/cli/convert_to_ptychodus.py diff --git a/src/ptychodus/scripts/ptychodus_bdp.py b/src/ptychodus/cli/ptychodus_bdp.py similarity index 96% rename from src/ptychodus/scripts/ptychodus_bdp.py rename to src/ptychodus/cli/ptychodus_bdp.py index 525885f86..a2700eefd 100755 --- a/src/ptychodus/scripts/ptychodus_bdp.py +++ b/src/ptychodus/cli/ptychodus_bdp.py @@ -166,13 +166,16 @@ def main() -> int: if args.defocus_distance_m is not None: logger.warning('Defocus distance is not implemented yet!') # TODO + bad_pixels_path = ( + Path(args.bad_pixels_input.name) if args.bad_pixels_input is not None else None + ) + with ModelCore(Path(args.settings.name), log_level=args.log_level) as model: - if args.bad_pixels_input is not None: - model.workflow_api.load_bad_pixels(Path(args.bad_pixels_input.name)) workflow_diffraction_api = model.workflow_api.load_diffraction_data( Path(args.diffraction_input.name), crop_center=crop_center, crop_extent=crop_extent, + bad_pixels_file_path=bad_pixels_path, block=True, ) workflow_product_api = model.workflow_api.create_product( @@ -183,6 +186,7 @@ def main() -> int: probe_photon_count=args.probe_photon_count, exposure_time_s=args.exposure_time_s, tomography_angle_deg=args.tomography_angle_deg, + diffraction=workflow_diffraction_api, ) workflow_product_api.load_probe_positions(Path(args.probe_position_input.name)) workflow_product_api.generate_probe() diff --git a/src/ptychodus/scripts/ptychodus_system_check.py b/src/ptychodus/cli/ptychodus_system_check.py similarity index 100% rename from src/ptychodus/scripts/ptychodus_system_check.py rename to src/ptychodus/cli/ptychodus_system_check.py diff --git a/src/ptychodus/controller/agent/core.py b/src/ptychodus/controller/agent/core.py index 0af17c1dc..1d318e9a7 100644 --- a/src/ptychodus/controller/agent/core.py +++ b/src/ptychodus/controller/agent/core.py @@ -1,27 +1,35 @@ +from collections.abc import Callable, Iterable + from PyQt5.QtCore import QEvent, QModelIndex, QObject, Qt from PyQt5.QtGui import QKeyEvent from PyQt5.QtWidgets import ( QAbstractItemView, + QComboBox, QFormLayout, QGroupBox, - QInputDialog, + QHBoxLayout, QListView, QPushButton, QVBoxLayout, + QWidget, ) +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.parametric import StringParameter + from ...model.agent import ( - AgentPresenter, - ArgoSettings, - ChatHistory, + AgentSettings, ChatMessage, - ChatObserver, + ChatTerminal, + ConversationObserver, + ConversationRepository, + ModelCatalog, ) from ...view.agent import AgentChatView, AgentInputView, AgentView from ..parametric import ( - ComboBoxParameterViewController, DecimalSliderParameterViewController, LineEditParameterViewController, + ParameterViewController, SpinBoxParameterViewController, ) from .item_delegate import ChatBubbleItemDelegate @@ -31,17 +39,18 @@ class AgentInputController(QObject): - def __init__(self, presenter: AgentPresenter, view: AgentInputView) -> None: + def __init__(self, terminal: ChatTerminal, view: AgentInputView) -> None: super().__init__() - self._presenter = presenter + self._terminal = terminal self._view = view view.text_edit.installEventFilter(self) view.send_button.clicked.connect(self._send_message) + view.clear_button.clicked.connect(terminal.clear_conversation) def _send_message(self) -> None: text = self._view.text_edit.toPlainText() - self._presenter.send_message(text) + self._terminal.send_message(text) self._view.text_edit.clear() def eventFilter(self, a0: QObject, a1: QEvent) -> bool: # noqa: N802 @@ -56,97 +65,152 @@ def eventFilter(self, a0: QObject, a1: QEvent) -> bool: # noqa: N802 return super().eventFilter(a0, a1) -class AgentChatController(ChatObserver): +class AgentChatController(ConversationObserver): def __init__( - self, history: ChatHistory, presenter: AgentPresenter, view: AgentChatView + self, + repository: ConversationRepository, + terminal: ChatTerminal, + view: AgentChatView, ) -> None: super().__init__() - self._history = history - self._presenter = presenter + self._repository = repository + self._terminal = terminal self._view = view - self._message_list_model = AgentMessageListModel(history) - self._input_controller = AgentInputController(presenter, view.input_view) + self._message_list_model = AgentMessageListModel(repository) + self._input_controller = AgentInputController(terminal, view.input_view) view.message_list_view.setModel(self._message_list_model) view.message_list_view.setItemDelegate(ChatBubbleItemDelegate()) view.message_list_view.setResizeMode(QListView.ResizeMode.Adjust) view.message_list_view.setVerticalScrollMode(QAbstractItemView.ScrollMode.ScrollPerPixel) - history.add_observer(self) + repository.add_observer(self) - def handle_new_message(self, message: ChatMessage, index: int) -> None: + def handle_message_appended(self, message: ChatMessage, index: int) -> None: parent = QModelIndex() self._message_list_model.beginInsertRows(parent, index, index) self._message_list_model.endInsertRows() - def handle_chat_cleared(self) -> None: + def handle_conversation_cleared(self) -> None: self._message_list_model.beginResetModel() self._message_list_model.endResetModel() -class AgentController: - def __init__(self, settings: ArgoSettings, presenter: AgentPresenter, view: AgentView) -> None: +class _RefreshableComboBoxParameterViewController(ParameterViewController, Observer): + """ComboBox bound to a StringParameter with a Refresh button that repopulates the items.""" + + def __init__( + self, + parameter: StringParameter, + on_refresh: Callable[[], Iterable[str]], + *, + tool_tip: str = '', + ) -> None: + super().__init__() + self._parameter = parameter + self._on_refresh = on_refresh + + self._combo = QComboBox() + self._refresh_button = QPushButton('Refresh') + + if tool_tip: + self._combo.setToolTip(tool_tip) + + layout = QHBoxLayout() + layout.setContentsMargins(0, 0, 0, 0) + layout.addWidget(self._combo, 1) + layout.addWidget(self._refresh_button) + + self._widget = QWidget() + self._widget.setLayout(layout) + + self._combo.textActivated.connect(parameter.set_value) + self._refresh_button.clicked.connect(self._handle_refresh_clicked) + parameter.add_observer(self) + self.__sync_model_to_view() + + def get_widget(self) -> QWidget: + return self._widget + + def populate(self, items: Iterable[str]) -> None: + self._combo.blockSignals(True) + self._combo.clear() + for item in items: + self._combo.addItem(item) + self._combo.blockSignals(False) + self.__sync_model_to_view() + + def _handle_refresh_clicked(self) -> None: + self.populate(self._on_refresh()) + + def __sync_model_to_view(self) -> None: + self._combo.setCurrentText(self._parameter.get_value()) + + def _update(self, observable: Observable) -> None: + if observable is self._parameter: + self.__sync_model_to_view() + + +class AgentController(QObject): + def __init__(self, settings: AgentSettings, catalog: ModelCatalog, view: AgentView) -> None: + super().__init__() self._settings = settings - self._presenter = presenter + self._catalog = catalog self._view = view - - self._user_view_controller = LineEditParameterViewController(settings.user) - self._chat_endpoint_url_view_controller = LineEditParameterViewController( - settings.chat_endpoint_url, tool_tip='The chat endpoint URL.' + self._models_loaded = False + + self._base_url_view_controller = LineEditParameterViewController( + settings.base_url, + tool_tip=( + 'OpenAI-compatible base URL. pydantic-ai appends /chat/completions and /models. ' + 'OPENAI_API_KEY env var is sent as a Bearer token; for Argo set this to your ' + 'Argonne username.' + ), ) - self._chat_model_view_controller = ComboBoxParameterViewController( - settings.chat_model, - presenter.get_available_chat_models(), - tool_tip='The chat model to use.', + self._model_view_controller = _RefreshableComboBoxParameterViewController( + settings.model, + catalog.refresh, + tool_tip='The chat model to use. Click Refresh to re-fetch the list from /models.', + ) + self._system_prompt_view_controller = LineEditParameterViewController( + settings.system_prompt, + tool_tip='System prompt sent at the start of every conversation.', ) self._temperature_view_controller = DecimalSliderParameterViewController( settings.temperature, - tool_tip='What sampling temperature to use, between 0 and 2. Higher values mean the model takes more risks.', + tool_tip=( + 'Sampling temperature between 0 and 2. Higher values mean the model takes ' + 'more risks.' + ), ) self._top_p_view_controller = DecimalSliderParameterViewController( settings.top_p, - tool_tip='An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass.', + tool_tip=( + 'Nucleus sampling: the model considers tokens with top_p probability mass. ' + 'Alternative to temperature.' + ), ) self._max_tokens_view_controller = SpinBoxParameterViewController( settings.max_tokens, - tool_tip='The maximum number of tokens that can be generated in the chat completion.', - ) - self._max_completion_tokens_view_controller = SpinBoxParameterViewController( - settings.max_completion_tokens, - tool_tip='An upper bound for the number of tokens that can be generated for a completion, including visible output tokens and reasoning tokens.', + tool_tip='Maximum number of tokens generated in the chat completion.', ) - self._embeddings_endpoint_url_view_controller = LineEditParameterViewController( - settings.embeddings_endpoint_url, tool_tip='The embeddings endpoint URL.' + self._mcp_server_url_view_controller = LineEditParameterViewController( + settings.mcp_server_url, + tool_tip='Optional MCP server URL (e.g. ptychodus_store at http://localhost:8000/mcp).', ) - self._embeddings_model_view_controller = ComboBoxParameterViewController( - settings.embeddings_model, - presenter.get_available_embeddings_models(), - tool_tip='The embeddings model to use.', - ) - self._embed_button = QPushButton('Embed Text') - self._embed_button.clicked.connect(self._embed_text) group_box_layout = QFormLayout() - group_box_layout.addRow('User:', self._user_view_controller.get_widget()) - group_box_layout.addRow( - 'Chat Endpoint URL:', self._chat_endpoint_url_view_controller.get_widget() - ) - group_box_layout.addRow('Chat Model:', self._chat_model_view_controller.get_widget()) + group_box_layout.addRow('Base URL:', self._base_url_view_controller.get_widget()) + group_box_layout.addRow('Model:', self._model_view_controller.get_widget()) + group_box_layout.addRow('System Prompt:', self._system_prompt_view_controller.get_widget()) group_box_layout.addRow('Temperature:', self._temperature_view_controller.get_widget()) group_box_layout.addRow('Top P:', self._top_p_view_controller.get_widget()) group_box_layout.addRow('Max Tokens:', self._max_tokens_view_controller.get_widget()) group_box_layout.addRow( - 'Max Completion Tokens:', self._max_completion_tokens_view_controller.get_widget() - ) - group_box_layout.addRow( - 'Embeddings Endpoint URL:', self._embeddings_endpoint_url_view_controller.get_widget() + 'MCP Server URL:', self._mcp_server_url_view_controller.get_widget() ) - group_box_layout.addRow( - 'Embeddings Model:', self._embeddings_model_view_controller.get_widget() - ) - group_box_layout.addRow(self._embed_button) - group_box = QGroupBox('Argo') + group_box = QGroupBox('Agent') group_box.setLayout(group_box_layout) layout = QVBoxLayout() @@ -154,10 +218,10 @@ def __init__(self, settings: ArgoSettings, presenter: AgentPresenter, view: Agen layout.addStretch() view.setLayout(layout) - def _embed_text(self) -> None: - title = 'Embed Text' - label = 'Enter text to embed:' - text, ok_pressed = QInputDialog.getMultiLineText(self._view, title, label, text='') + view.installEventFilter(self) - if ok_pressed: - self._presenter.embed_text(text.splitlines()) + def eventFilter(self, a0: QObject, a1: QEvent) -> bool: # noqa: N802 + if a0 is self._view and a1.type() == QEvent.Type.Show and not self._models_loaded: + self._models_loaded = True + self._model_view_controller.populate(self._catalog.get_available_models()) + return super().eventFilter(a0, a1) diff --git a/src/ptychodus/controller/agent/list_model.py b/src/ptychodus/controller/agent/list_model.py index 96063fbba..8936a6b0b 100644 --- a/src/ptychodus/controller/agent/list_model.py +++ b/src/ptychodus/controller/agent/list_model.py @@ -3,7 +3,7 @@ from PyQt5.QtCore import QAbstractListModel, QModelIndex, QObject, Qt from PyQt5.QtGui import QColor -from ...model.agent import ChatHistory, ChatRole +from ...model.agent import ChatRole, ConversationRepository class AgentMessageListModel(QAbstractListModel): @@ -12,13 +12,13 @@ class AgentMessageListModel(QAbstractListModel): DARK_GREEN: Final[QColor] = QColor('#00894d') LIGHT_GREEN: Final[QColor] = QColor('#78ca2a') - def __init__(self, history: ChatHistory, parent: QObject | None = None) -> None: + def __init__(self, repository: ConversationRepository, parent: QObject | None = None) -> None: super().__init__(parent) - self._history = history + self._repository = repository def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: if index.isValid(): - message = self._history[index.row()] + message = self._repository[index.row()] match role: case Qt.ItemDataRole.DisplayRole: @@ -35,4 +35,4 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A return self.DARK_BLUE if message.role == ChatRole.USER else self.DARK_GREEN def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 - return len(self._history) + return len(self._repository) diff --git a/src/ptychodus/controller/core.py b/src/ptychodus/controller/core.py index 6e0d23eec..85103550b 100644 --- a/src/ptychodus/controller/core.py +++ b/src/ptychodus/controller/core.py @@ -10,6 +10,8 @@ from .automation import AutomationController from .data import FileDialogFactory from .diffraction import DiffractionController +from .fluorescence import FluorescenceController +from .fluorescence.enhance_dialog import FluorescenceEnhanceDialogController from .genesis import GenesisController from .globus import GlobusController from .image import ImageController @@ -17,9 +19,12 @@ from .object import ObjectController from .probe import ProbeController from .probe_positions import ProbePositionsController +from .helpers import create_brush_for_editable_cell from .processing import ProcessingController from .product import ProductController +from .product.core import ProductRepositoryTableModel from .product.visualization import ProductVisualizationController +from .ptycho_fm import PtychoFMViewControllerFactory from .ptychi import PtyChiViewControllerFactory from .ptychonn import PtychoNNViewControllerFactory from .ptychopinn import PtychoPINNViewControllerFactory @@ -55,40 +60,55 @@ def __init__( model.ptychopinn_torch_reconstructor_library, self._file_dialog_factory, ) + self._ptycho_fm_view_controller_factory = PtychoFMViewControllerFactory( + model.ptycho_fm_reconstructor_library, + self._file_dialog_factory, + ) + # Shared product-repository table model. Constructing it here (before + # any consumer) lets SettingsController, DiffractionController, + # ProductController, and ProcessingController all bind to the same + # instance — one observer registration on the repository serves every + # widget that shows product rows. + self._product_table_model = ProductRepositoryTableModel( + model.product_core.product_repository, + create_brush_for_editable_cell(view.product_view.table_view), + ) self._settings_controller = SettingsController( model.settings_registry, model.product_core.product_repository, + self._product_table_model, view.settings_view, view.settings_table_view, self._file_dialog_factory, ) - self._patterns_image_controller = ImageController( + self._diffraction_image_controller = ImageController( model.pattern_visualization_engine, - view.patterns_image_view.image_view, + view.diffraction_image_view.image_view, self._status_bar, self._file_dialog_factory, ) - self._patterns_controller = DiffractionController( + self._diffraction_controller = DiffractionController( model.diffraction_core.detector_settings, model.diffraction_core.diffraction_settings, + model.product_core.settings, model.diffraction_core.pattern_sizer, - model.diffraction_core.detector, model.diffraction_core.diffraction_api, - model.diffraction_core.dataset, + model.diffraction_core.repository, model.diffraction_core.task_monitor, - model.metadata_presenter, model.product_core.product_repository, + self._product_table_model, model.analysis_core.diffraction_simulator, model.analysis_core.diffraction_simulator_settings, - view.patterns_view, - view.patterns_image_view.status_view, - self._patterns_image_controller, + view.datasets_view, + view.diffraction_image_view.status_view, + self._diffraction_image_controller, self._file_dialog_factory, ) self._product_controller = ProductController.create_instance( - model.diffraction_core.diffraction_api, model.product_core.product_repository, model.product_core.product_api, + model.diffraction_core.repository, + self._product_table_model, view.product_view, self._file_dialog_factory, ) @@ -125,13 +145,6 @@ def __init__( model.analysis_core.probe_propagator_visualization_engine, model.analysis_core.illumination_mapper, model.analysis_core.illumination_visualization_engine, - model.fluorescence_core.fluorescence_api, - model.fluorescence_core.enhancer_chooser, - model.fluorescence_core.two_step_enhancer, - model.fluorescence_core.vspi_enhancer, - model.fluorescence_core.ptychozoon_enhancer, - model.fluorescence_core.task_monitor, - model.fluorescence_core.visualization_engine, view.probe_view, self._file_dialog_factory, ) @@ -141,6 +154,25 @@ def __init__( self._status_bar, self._file_dialog_factory, ) + self._fluorescence_image_controller = ImageController( + model.fluorescence_core.visualization_engine, + view.fluorescence_image_view, + self._status_bar, + self._file_dialog_factory, + ) + self._fluorescence_enhance_dialog_controller = FluorescenceEnhanceDialogController( + model.fluorescence_core, + has_ptychozoon=model.fluorescence_core.ptychozoon_enhancer is not None, + ) + self._fluorescence_controller = FluorescenceController( + model.fluorescence_core.repository, + model.fluorescence_core.fluorescence_api, + model.product_core.product_repository, + view.fluorescence_view, + self._fluorescence_image_controller, + self._fluorescence_enhance_dialog_controller, + self._file_dialog_factory, + ) self._object_controller = ObjectController( model.product_core.object_repository, model.product_core.object_api, @@ -159,6 +191,7 @@ def __init__( model.processing_core.algorithm_parameter, model.processing_core.processing_api, model.product_core.product_repository, + self._product_table_model, model.globus_core, model.genesis_core, view.processing_view, @@ -169,6 +202,7 @@ def __init__( self._ptychopinn_torch_view_controller_factory, self._ptychopinn_view_controller_factory, self._ptychonn_view_controller_factory, + self._ptycho_fm_view_controller_factory, ], ) self._globus_controller = GlobusController( @@ -195,10 +229,10 @@ def __init__( self._file_dialog_factory, ) self._agent_controller = AgentController( - model.agent_core.settings, model.agent_core.presenter, view.agent_view + model.agent_core.settings, model.agent_core.catalog, view.agent_view ) self._agent_chat_controller = AgentChatController( - model.agent_core.chat_history, model.agent_core.presenter, view.agent_chat_view + model.agent_core.repository, model.agent_core.terminal, view.agent_chat_view ) self._one_second_counter = 0 @@ -211,8 +245,8 @@ def __init__( view.genesis_action: model.genesis_core.is_supported, } - self._swap_central_widgets(view.patterns_action, animated=False) - view.patterns_action.setChecked(True) + self._swap_central_widgets(view.datasets_action, animated=False) + view.datasets_action.setChecked(True) view.navigation.action_group.triggered.connect( lambda action: self._swap_central_widgets(action) ) diff --git a/src/ptychodus/controller/diffraction/core.py b/src/ptychodus/controller/diffraction/core.py index 579a3f281..6713c0991 100644 --- a/src/ptychodus/controller/diffraction/core.py +++ b/src/ptychodus/controller/diffraction/core.py @@ -1,150 +1,42 @@ import logging - from PyQt5.QtCore import QModelIndex from PyQt5.QtWidgets import ( QAbstractItemView, QDialog, - QFormLayout, - QHBoxLayout, - QLineEdit, QMessageBox, - QPushButton, - QWidget, ) from ptychodus.api.observer import Observable, Observer -from ptychodus.api.parametric import PathParameter, StringParameter from ...model.analysis import DiffractionSimulator, DiffractionSimulatorSettings from ...model.diffraction import ( AssembledDiffractionDataset, - Detector, DetectorSettings, DiffractionAPI, DiffractionDatasetObserver, + DiffractionDatasetRepository, + DiffractionDatasetRepositoryObserver, DiffractionSettings, DiffractionTaskMonitor, PatternSizer, ) -from ...model.metadata import MetadataPresenter -from ...model.product import ProductRepository -from ...view.diffraction import DetectorView, DiffractionStatusView, PatternsView +from .detector_extent import DetectorExtentSource +from ...model.product import ProductRepository, ProductSettings +from ...view.diffraction import DatasetsView, DiffractionStatusView from ...view.widgets import ExceptionDialog, ProgressBarItemDelegate from ..data import FileDialogFactory from ..helpers import connect_triggered_signal from ..image import ImageController -from ..parametric import ( - CheckBoxParameterViewController, - LengthWidgetParameterViewController, - ParameterViewController, - SpinBoxParameterViewController, -) -from ..product.list_model import ProductRepositoryListModel +from ..parametric import CheckBoxParameterViewController +from ..product.core import ProductRepositoryComboProxyModel, ProductRepositoryTableModel from .dataset import DatasetTreeModel -from .dataset_layout import DatasetLayoutViewController +from .dataset_editor import DatasetEditorViewController from .wizard import OpenDatasetWizardController logger = logging.getLogger(__name__) -class BadPixelsViewController(ParameterViewController, Observer): - def __init__( - self, - bad_pixels_file_path: PathParameter, - bad_pixels_file_type: StringParameter, - detector: Detector, - diffraction_api: DiffractionAPI, - file_dialog_factory: FileDialogFactory, - ) -> None: - super().__init__() - self._bad_pixels_file_path = bad_pixels_file_path - self._bad_pixels_file_type = bad_pixels_file_type - self._detector = detector - self._diffraction_api = diffraction_api - self._file_dialog_factory = file_dialog_factory - - self._line_edit = QLineEdit() - self._line_edit.setReadOnly(True) - self._browse_button = QPushButton('Browse...') - self._browse_button.clicked.connect(self._open_bad_pixels) - self._clear_button = QPushButton('Clear') - self._clear_button.clicked.connect(self._diffraction_api.clear_bad_pixels) - self._widget = QWidget() - - layout = QHBoxLayout() - layout.setContentsMargins(0, 0, 0, 0) - layout.addWidget(self._line_edit) - layout.addWidget(self._browse_button) - layout.addWidget(self._clear_button) - self._widget.setLayout(layout) - - self._sync_model_to_view() - detector.add_observer(self) - - def _open_bad_pixels(self) -> None: - file_reader_chooser = self._diffraction_api.get_bad_pixels_file_reader_chooser() - current_plugin = file_reader_chooser.get_current_plugin() - file_path, name_filter = self._file_dialog_factory.get_open_file_path( - self._widget, - 'Open Bad Pixels File', - name_filters=[plugin.display_name for plugin in file_reader_chooser], - selected_name_filter=current_plugin.simple_name, - ) - - if file_path: - try: - self._diffraction_api.open_bad_pixels(file_path, file_type=name_filter) - except Exception as exc: - logger.exception(exc) - ExceptionDialog.show_exception('Bad Pixels File Reader', exc) - - def get_widget(self) -> QWidget: - return self._widget - - def _sync_model_to_view(self) -> None: - num_bad_pixels = self._detector.get_num_bad_pixels() - self._line_edit.setText(str(num_bad_pixels)) - - def _update(self, observable: Observable) -> None: - if observable is self._detector: - self._sync_model_to_view() - - -class DetectorController: - def __init__( - self, - settings: DetectorSettings, - detector: Detector, - diffraction_api: DiffractionAPI, - view: DetectorView, - file_dialog_factory: FileDialogFactory, - ) -> None: - self._width_px_view_controller = SpinBoxParameterViewController(settings.width_px) - self._height_px_view_controller = SpinBoxParameterViewController(settings.height_px) - self._pixel_width_view_controller = LengthWidgetParameterViewController( - settings.pixel_width_m - ) - self._pixel_height_view_controller = LengthWidgetParameterViewController( - settings.pixel_height_m - ) - self._bad_pixels_view_controller = BadPixelsViewController( - settings.bad_pixels_file_path, - settings.bad_pixels_file_type, - detector, - diffraction_api, - file_dialog_factory, - ) - - layout = QFormLayout() - layout.addRow('Detector Width [px]:', self._width_px_view_controller.get_widget()) - layout.addRow('Detector Height [px]:', self._height_px_view_controller.get_widget()) - layout.addRow('Pixel Width:', self._pixel_width_view_controller.get_widget()) - layout.addRow('Pixel Height:', self._pixel_height_view_controller.get_widget()) - layout.addRow('Bad Pixels:', self._bad_pixels_view_controller.get_widget()) - view.setLayout(layout) - - class DiffractionStatusController(Observer): def __init__( self, @@ -178,21 +70,21 @@ def _update(self, observable: Observable) -> None: self._sync_model_to_view() -class DiffractionController(DiffractionDatasetObserver): +class DiffractionController(DiffractionDatasetRepositoryObserver, Observer): def __init__( self, detector_settings: DetectorSettings, diffraction_settings: DiffractionSettings, + product_settings: ProductSettings, pattern_sizer: PatternSizer, - detector: Detector, diffraction_api: DiffractionAPI, - dataset: AssembledDiffractionDataset, + repository: DiffractionDatasetRepository, task_monitor: DiffractionTaskMonitor, - metadata_presenter: MetadataPresenter, product_repository: ProductRepository, + product_table_model: ProductRepositoryTableModel, diffraction_simulator: DiffractionSimulator, diffraction_simulator_settings: DiffractionSimulatorSettings, - view: PatternsView, + view: DatasetsView, status_view: DiffractionStatusView, image_controller: ImageController, file_dialog_factory: FileDialogFactory, @@ -200,28 +92,32 @@ def __init__( super().__init__() self._pattern_sizer = pattern_sizer self._diffraction_api = diffraction_api - self._dataset = dataset - self._product_list_model = ProductRepositoryListModel(product_repository) + self._repository = repository + self._detector_extent_source = DetectorExtentSource() + self._current_dataset_index = -1 + self._product_combo_model = ProductRepositoryComboProxyModel( + product_table_model, product_repository + ) self._diffraction_simulator = diffraction_simulator self._view = view self._image_controller = image_controller self._file_dialog_factory = file_dialog_factory - self._detector_controller = DetectorController( - detector_settings, - detector, - diffraction_api, - view.detector_view, - file_dialog_factory, - ) + self._per_dataset_observers: dict[int, _PerDatasetObserver] = {} self._wizard_controller = OpenDatasetWizardController( diffraction_settings, - pattern_sizer, + detector_settings, + product_settings, + self._detector_extent_source, diffraction_api, - metadata_presenter, + repository, file_dialog_factory, ) self._status_controller = DiffractionStatusController(task_monitor, status_view) - self._tree_model = DatasetTreeModel() + self._tree_model = DatasetTreeModel(pattern_sizer, repository) + + # Sizer changes (binning toggle / bin size / transpose) shift the processed + # pixel geometry for every dataset, so refresh those two columns on notify. + pattern_sizer.add_observer(self) view.tree_view.setModel(self._tree_model) view.tree_view.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows) @@ -234,22 +130,25 @@ def __init__( if selection_model is None: raise ValueError('selection_model is None!') else: - selection_model.currentChanged.connect(self._update_view) + selection_model.currentChanged.connect(self._on_tree_selection_changed) + selection_model.currentChanged.connect(self._update_enabled_buttons) - self._update_view(QModelIndex(), QModelIndex()) + self._image_controller.clear_array() - open_dataset_action = view.button_box.load_menu.addAction('Open File...') + open_dataset_action = view.button_box.insert_menu.addAction('Open File...') connect_triggered_signal(open_dataset_action, self._wizard_controller.open_dataset) + simulate_action = view.button_box.insert_menu.addAction('Simulate...') + connect_triggered_signal(simulate_action, self._choose_product_for_simulation) - view.button_box.save_button.clicked.connect(self._save_dataset) - view.button_box.close_button.clicked.connect(self._close_dataset) + save_file_action = view.button_box.save_menu.addAction('Save File...') + connect_triggered_signal(save_file_action, self._save_dataset) + sync_to_settings_action = view.button_box.save_menu.addAction('Sync To Settings') + connect_triggered_signal(sync_to_settings_action, self._sync_current_to_settings) - dataset_layout_action = view.button_box.analyze_menu.addAction('Dataset Layout...') - connect_triggered_signal(dataset_layout_action, self._show_dataset_layout) - simulate_action = view.button_box.analyze_menu.addAction('Simulate Diffraction...') - connect_triggered_signal(simulate_action, self._choose_product_for_simulation) + view.button_box.edit_button.clicked.connect(self._edit_current_dataset) + view.button_box.remove_button.clicked.connect(self._remove_selected_dataset) - view.simulate_dialog.product_combo_box.setModel(self._product_list_model) + view.simulate_dialog.product_combo_box.setModel(self._product_combo_model) self._poisson_view_controller = CheckBoxParameterViewController( diffraction_simulator_settings.add_poisson_noise, 'Add Poisson Noise', @@ -257,72 +156,232 @@ def __init__( view.simulate_dialog.form_layout.insertRow(1, self._poisson_view_controller.get_widget()) view.simulate_dialog.finished.connect(self._simulate_diffraction) - dataset.add_observer(self) - self._sync_model_to_view() + self._update_enabled_buttons(QModelIndex(), QModelIndex()) + repository.add_observer(self) + + # Populate the tree for datasets already present at construction time. + for index in range(len(repository)): + self._on_dataset_inserted(index, repository[index]) + + self._update_info_text() - def _update_view(self, current: QModelIndex, previous: QModelIndex) -> None: + def _current_dataset(self) -> AssembledDiffractionDataset | None: + index = self._current_dataset_index + return self._repository[index] if 0 <= index < len(self._repository) else None + + def _set_current_dataset_index(self, index: int) -> None: + self._current_dataset_index = index + + def _on_tree_selection_changed(self, current: QModelIndex, previous: QModelIndex) -> None: + # Update the image preview based on the selected tree node. if current.isValid(): node = current.internalPointer() data = node.get_data() - pixel_geometry = self._pattern_sizer.get_processed_pixel_geometry() - self._image_controller.set_array(data, pixel_geometry) + if data is not None: + dataset_row = self._tree_model.dataset_row_for_index(current) + dataset = ( + self._repository[dataset_row] + if dataset_row is not None and 0 <= dataset_row < len(self._repository) + else None + ) + if dataset is not None: + pixel_geometry = self._pattern_sizer.get_processed_pixel_geometry( + dataset.get_raw_pixel_geometry() + ) + self._image_controller.set_array(data, pixel_geometry) + else: + self._image_controller.clear_array() + else: + self._image_controller.clear_array() else: self._image_controller.clear_array() + # And track the containing dataset as the panel's current dataset. + dataset_row = self._tree_model.dataset_row_for_index(current) + self._set_current_dataset_index(dataset_row if dataset_row is not None else -1) + def _save_dataset(self) -> None: - file_writer_chooser = self._diffraction_api.get_file_writer_chooser() + dataset_index = self._current_dataset_index + if dataset_index < 0: + return + file_path, name_filter = self._file_dialog_factory.get_save_file_path( self._view, 'Save Diffraction File', - name_filters=[plugin.display_name for plugin in file_writer_chooser], - selected_name_filter=file_writer_chooser.get_current_plugin().display_name, + name_filters=[nf for nf in self._diffraction_api.get_save_file_filters()], + selected_name_filter=self._diffraction_api.get_save_file_filter(), ) if file_path: try: - self._diffraction_api.save_patterns(file_path, name_filter) + self._diffraction_api.save_patterns( + file_path, name_filter, dataset_index=dataset_index + ) except Exception as exc: logger.exception(exc) ExceptionDialog.show_exception('File Writer', exc) - def _show_dataset_layout(self) -> None: - DatasetLayoutViewController.show_dialog(self._dataset, self._view) + def _edit_current_dataset(self) -> None: + current = self._current_dataset() + if current is not None: + DatasetEditorViewController.edit_dataset(current, self._view) + + def _sync_current_to_settings(self) -> None: + current = self._current_dataset() + if current is not None: + current.sync_pixel_geometry_to_settings() def _choose_product_for_simulation(self) -> None: - self._product_list_model.beginResetModel() - self._product_list_model.endResetModel() self._view.simulate_dialog.open() def _simulate_diffraction(self, result: int) -> None: - if result == QDialog.DialogCode.Accepted: - item_index = self._view.simulate_dialog.product_combo_box.currentIndex() - self._diffraction_simulator.simulate(item_index) + if result != QDialog.DialogCode.Accepted: + return + item_index = self._view.simulate_dialog.product_combo_box.currentIndex() + if item_index < 0: + logger.warning('Cannot simulate diffraction: no product selected.') + return + self._diffraction_simulator.simulate(item_index) + + def _remove_selected_dataset(self) -> None: + dataset_index = self._current_dataset_index + if dataset_index < 0: + return - def _close_dataset(self) -> None: button = QMessageBox.question( self._view, - 'Confirm Close', - 'This will free the diffraction data from memory. Do you want to continue?', + 'Confirm Remove', + 'Remove the selected diffraction dataset from memory?', ) if button == QMessageBox.StandardButton.Yes: - self._diffraction_api.close_patterns() - self._update_view(QModelIndex(), QModelIndex()) + self._diffraction_api.close_patterns(dataset_index) + self._image_controller.clear_array() - def _sync_model_to_view(self) -> None: - self._tree_model.clear() + def _update_info_text(self) -> None: + self._view.info_label.setText(self._repository.get_info_text()) + + def _update_enabled_buttons(self, current: QModelIndex, previous: QModelIndex) -> None: + dataset = self._current_dataset() + ready = dataset is not None and not dataset.is_load_in_progress() + + self._view.button_box.save_button.setEnabled(ready) + self._view.button_box.edit_button.setEnabled(ready) + # Remove is always safe: it drops the row whether pending, failed, or ready. + self._view.button_box.remove_button.setEnabled(dataset is not None) - for index, array in enumerate(self._dataset): - self._tree_model.insert_array(index, array) # type: ignore + def _select_dataset_row(self, dataset_row: int) -> None: + selection_model = self._view.tree_view.selectionModel() + if selection_model is None: + return + model_index = self._tree_model.index(dataset_row, 0) + if model_index.isValid(): + selection_model.setCurrentIndex( + model_index, + selection_model.SelectionFlag.ClearAndSelect | selection_model.SelectionFlag.Rows, + ) + + def _on_dataset_inserted(self, index: int, dataset: AssembledDiffractionDataset) -> None: + self._tree_model.insert_dataset(index, dataset) + + observer = _PerDatasetObserver(self, dataset) + dataset.add_observer(observer) + self._per_dataset_observers[id(dataset)] = observer + + # Populate any arrays already present on the dataset. + for array_index in range(len(dataset)): + self._tree_model.insert_array(index, array_index, dataset[array_index]) + + self._update_info_text() + + def handle_dataset_inserted(self, index: int, dataset: AssembledDiffractionDataset) -> None: + self._on_dataset_inserted(index, dataset) + # Publish the newly-loaded detector extent to the wizard's crop/bin bounds. + self._detector_extent_source.set_extent(dataset.get_metadata().detector_extent) + # Make a freshly loaded dataset the panel's current dataset. + self._select_dataset_row(index) + + def handle_dataset_removed(self, index: int, dataset: AssembledDiffractionDataset) -> None: + observer = self._per_dataset_observers.pop(id(dataset), None) + if observer is not None: + dataset.remove_observer(observer) + self._tree_model.remove_dataset(index) + self._update_info_text() + + # Re-sync the panel's current dataset with the (possibly changed) tree selection. + selection_model = self._view.tree_view.selectionModel() + current = selection_model.currentIndex() if selection_model is not None else QModelIndex() + dataset_row = self._tree_model.dataset_row_for_index(current) + self._set_current_dataset_index(dataset_row if dataset_row is not None else -1) + self._update_enabled_buttons(QModelIndex(), QModelIndex()) + + def _dataset_row(self, dataset: AssembledDiffractionDataset) -> int | None: + try: + return list(self._repository).index(dataset) + except ValueError: + return None + + def _handle_array_inserted_for_dataset( + self, dataset: AssembledDiffractionDataset, array_row: int + ) -> None: + dataset_row = self._dataset_row(dataset) + if dataset_row is None: + return + self._tree_model.insert_array(dataset_row, array_row, dataset[array_row]) + self._update_info_text() + # A landing array may flip the current dataset out of the loading state, + # re-enabling Save / Edit. + if dataset_row == self._current_dataset_index: + self._update_enabled_buttons(QModelIndex(), QModelIndex()) + + def _handle_array_changed_for_dataset( + self, dataset: AssembledDiffractionDataset, array_row: int + ) -> None: + dataset_row = self._dataset_row(dataset) + if dataset_row is None: + return + self._tree_model.refresh_array(dataset_row, array_row) + + def _handle_dataset_reloaded_for_dataset(self, dataset: AssembledDiffractionDataset) -> None: + dataset_row = self._dataset_row(dataset) + if dataset_row is None: + return + self._tree_model.refresh_dataset(dataset_row) + # Refresh the wizard's crop/bin bounds when a dataset is reloaded (e.g. via a + # streaming context or a programmatic re-open with a different file). + self._detector_extent_source.set_extent(dataset.get_metadata().detector_extent) + + def _handle_pixel_geometry_changed_for_dataset( + self, dataset: AssembledDiffractionDataset + ) -> None: + dataset_row = self._dataset_row(dataset) + if dataset_row is None: + return + self._tree_model.refresh_dataset(dataset_row) - info_text = self._dataset.get_info_text() - self._view.info_label.setText(info_text) + def _update(self, observable: Observable) -> None: + if observable is self._pattern_sizer: + self._tree_model.refresh_processed_columns() + + +class _PerDatasetObserver(DiffractionDatasetObserver): + """Per-dataset observer wrapper that captures the dataset ref at registration.""" + + def __init__( + self, controller: DiffractionController, dataset: AssembledDiffractionDataset + ) -> None: + super().__init__() + self._controller = controller + self._dataset = dataset def handle_array_inserted(self, index: int) -> None: - self._tree_model.insert_array(index, self._dataset[index]) + self._controller._handle_array_inserted_for_dataset(self._dataset, index) def handle_array_changed(self, index: int) -> None: - self._tree_model.refresh_array(index) + self._controller._handle_array_changed_for_dataset(self._dataset, index) def handle_dataset_reloaded(self) -> None: - self._sync_model_to_view() + self._controller._handle_dataset_reloaded_for_dataset(self._dataset) + + def handle_pixel_geometry_changed(self) -> None: + self._controller._handle_pixel_geometry_changed_for_dataset(self._dataset) diff --git a/src/ptychodus/controller/diffraction/dataset.py b/src/ptychodus/controller/diffraction/dataset.py index d7f65ba4c..623c90440 100644 --- a/src/ptychodus/controller/diffraction/dataset.py +++ b/src/ptychodus/controller/diffraction/dataset.py @@ -1,124 +1,294 @@ from __future__ import annotations +import logging from typing import Any, overload +import numpy + from PyQt5.QtCore import Qt, QAbstractItemModel, QModelIndex, QObject from ptychodus.api.common import BYTES_PER_MEGABYTE from ptychodus.api.diffraction import DiffractionPattern +from ptychodus.api.geometry import ImageExtent, PixelGeometry -from ptychodus.model.diffraction import AssembledDiffractionArray +from ptychodus.model.diffraction import ( + AssembledDiffractionArray, + AssembledDiffractionDataset, + DiffractionDatasetRepository, + PatternSizer, +) __all__ = ['DatasetTreeModel'] +logger = logging.getLogger(__name__) -class DatasetTreeNode: - def __init__( - self, - parent_node: DatasetTreeNode | None, - array: AssembledDiffractionArray, - frame_index: int, - ) -> None: +_COL_LABEL = 0 +_COL_COUNTS = 1 +_COL_FRAMES = 2 +_COL_SIZE_MB = 3 +_COL_WIDTH_PX = 4 +_COL_HEIGHT_PX = 5 +_COL_PHYSICAL_PIXEL_WIDTH_UM = 6 +_COL_PHYSICAL_PIXEL_HEIGHT_UM = 7 +_COL_PROCESSED_PIXEL_WIDTH_UM = 8 +_COL_PROCESSED_PIXEL_HEIGHT_UM = 9 +_COL_NUM_BAD_PIXELS = 10 + + +class _TreeNode: + """Base tree node — root, dataset, array, or frame.""" + + def __init__(self, parent_node: _TreeNode | None) -> None: self.parent_node = parent_node - self._array = array - self._frame_index = frame_index - self.child_nodes: list[DatasetTreeNode] = list() + self.child_nodes: list[_TreeNode] = [] - @classmethod - def create_root(cls) -> DatasetTreeNode: - return cls(None, AssembledDiffractionArray.create_null(), -1) + def get_label(self) -> str: + return '' - def insert_child(self, pos: int, array: AssembledDiffractionArray) -> DatasetTreeNode: - child = DatasetTreeNode(self, array, -1) + def get_counts(self) -> int: + return 0 + + def get_nframes(self) -> int: + return sum(child.get_nframes() for child in self.child_nodes) + + def get_nbytes(self) -> int: + return sum(child.get_nbytes() for child in self.child_nodes) + + def get_data(self) -> DiffractionPattern | None: + return None + + def get_row(self) -> int: + return 0 if self.parent_node is None else self.parent_node.child_nodes.index(self) - for frame_index in range(array.get_num_patterns()): - grandchild = DatasetTreeNode(child, array, frame_index) - child.child_nodes.append(grandchild) - self.child_nodes.insert(pos, child) - return child +class _DatasetTreeNode(_TreeNode): + def __init__(self, parent_node: _TreeNode, dataset: AssembledDiffractionDataset) -> None: + super().__init__(parent_node) + self._dataset = dataset + + def get_dataset(self) -> AssembledDiffractionDataset: + return self._dataset def get_label(self) -> str: - return self._array.get_label() if self._frame_index < 0 else f'Frame {self._frame_index}' + return self._dataset.get_name() + + def get_counts(self) -> int: + if not self.child_nodes: + return 0 + return sum(child.get_counts() for child in self.child_nodes) // len(self.child_nodes) + + def get_data(self) -> DiffractionPattern | None: + return self._dataset.get_average_pattern() + + def get_detector_extent(self) -> ImageExtent: + return self._dataset.get_metadata().detector_extent + + def get_raw_pixel_geometry(self) -> PixelGeometry: + return self._dataset.get_raw_pixel_geometry() + + def get_processed_pixel_geometry(self, sizer: PatternSizer) -> PixelGeometry: + return sizer.get_processed_pixel_geometry(self._dataset.get_raw_pixel_geometry()) + + def get_num_bad_pixels(self) -> int: + return int(numpy.count_nonzero(self._dataset.get_bad_pixels())) + + +class _ArrayTreeNode(_TreeNode): + def __init__(self, parent_node: _TreeNode, array: AssembledDiffractionArray) -> None: + super().__init__(parent_node) + self._array = array + for frame_index in range(array.get_num_patterns()): + self.child_nodes.append(_FrameTreeNode(self, array, frame_index)) + + def get_label(self) -> str: + return self._array.get_label() + + def get_counts(self) -> int: + return int(self._array.get_mean_pattern_counts()) + + def get_max_counts(self) -> int: + return int(self._array.get_max_pattern_counts()) + + def get_nframes(self) -> int: + return len(self.child_nodes) + + def get_nbytes(self) -> int: + return self._array.get_patterns().nbytes def get_data(self) -> DiffractionPattern: - return ( - self._array.get_average_pattern() - if self._frame_index < 0 - else self._array.get_pattern(self._frame_index) - ) + return self._array.get_average_pattern() + + +class _FrameTreeNode(_TreeNode): + def __init__( + self, + parent_node: _TreeNode, + array: AssembledDiffractionArray, + frame_index: int, + ) -> None: + super().__init__(parent_node) + self._array = array + self._frame_index = frame_index + + def get_label(self) -> str: + return f'Frame {self._frame_index}' def get_counts(self) -> int: - return ( - int(self._array.get_mean_pattern_counts()) - if self._frame_index < 0 - else int(self._array.get_pattern_counts(self._frame_index)) - ) + return int(self._array.get_pattern_counts(self._frame_index)) def get_nframes(self) -> int: - return len(self.child_nodes) if self._frame_index < 0 else 1 + return 1 def get_nbytes(self) -> int: - return ( - self._array.get_patterns().nbytes - if self._frame_index < 0 - else self._array.get_pattern(self._frame_index).nbytes - ) + return self._array.get_pattern(self._frame_index).nbytes - def get_row(self) -> int: - return 0 if self.parent_node is None else self.parent_node.child_nodes.index(self) + def get_data(self) -> DiffractionPattern: + return self._array.get_pattern(self._frame_index) + + +def _find_containing_dataset_row(node: _TreeNode) -> int | None: + """Walk up until the dataset node is found; return its row within the root. None for the root.""" + current: _TreeNode | None = node + while current is not None: + if isinstance(current, _DatasetTreeNode): + return current.get_row() + current = current.parent_node + return None class DatasetTreeModel(QAbstractItemModel): - def __init__(self, parent: QObject | None = None) -> None: + """Three-level tree: root → dataset → array → frame.""" + + def __init__( + self, + sizer: PatternSizer, + repository: DiffractionDatasetRepository, + parent: QObject | None = None, + ) -> None: super().__init__(parent) - self._nodes = DatasetTreeNode.create_root() + self._sizer = sizer + self._repository = repository + self._root = _TreeNode(None) self._max_counts = 1 - self._header = ['Label', 'Counts', 'Frames', 'Size [MB]'] + self._header = [ + 'Label', + 'Counts', + 'Frames', + 'Size [MB]', + 'Width\n[px]', + 'Height\n[px]', + 'Physical Pixel\nWidth [µm]', + 'Physical Pixel\nHeight [µm]', + 'Processed Pixel\nWidth [µm]', + 'Processed Pixel\nHeight [µm]', + 'Num Bad\nPixels', + ] def clear(self) -> None: self.beginResetModel() - self._nodes = DatasetTreeNode.create_root() + self._root = _TreeNode(None) self._max_counts = 1 self.endResetModel() - def insert_array(self, row: int, array: AssembledDiffractionArray) -> None: - max_counts = array.get_max_pattern_counts() + def insert_dataset(self, row: int, dataset: AssembledDiffractionDataset) -> None: + self.beginInsertRows(QModelIndex(), row, row) + dataset_node = _DatasetTreeNode(self._root, dataset) + self._root.child_nodes.insert(row, dataset_node) + self.endInsertRows() + + def remove_dataset(self, row: int) -> None: + if not 0 <= row < len(self._root.child_nodes): + return + self.beginRemoveRows(QModelIndex(), row, row) + del self._root.child_nodes[row] + self.endRemoveRows() + + def _dataset_node(self, dataset_row: int) -> _DatasetTreeNode | None: + if not 0 <= dataset_row < len(self._root.child_nodes): + return None + node = self._root.child_nodes[dataset_row] + assert isinstance(node, _DatasetTreeNode) + return node + + def insert_array( + self, dataset_row: int, array_row: int, array: AssembledDiffractionArray + ) -> None: + dataset_node = self._dataset_node(dataset_row) + if dataset_node is None: + return + max_counts = int(array.get_max_pattern_counts()) if self._max_counts < max_counts: self._max_counts = max_counts - num_rows = self.rowCount() + self._rebroadcast_counts() - top_left = self.index(0, 1) - bottom_right = self.index(num_rows - 1, 1) - self.dataChanged.emit(top_left, bottom_right) - - for row2 in range(num_rows): - parent_index = self.index(row2, 0) - num_rows2 = self.rowCount(parent_index) - - child_top_left = self.index(0, 1, parent_index) - child_bottom_right = self.index(num_rows2 - 1, 1, parent_index) - self.dataChanged.emit(child_top_left, child_bottom_right) - - self.beginInsertRows(QModelIndex(), row, row) - child_node = self._nodes.insert_child(row, array) + dataset_index = self.index(dataset_row, 0) + self.beginInsertRows(dataset_index, array_row, array_row) + array_node = _ArrayTreeNode(dataset_node, array) + dataset_node.child_nodes.insert(array_row, array_node) self.endInsertRows() - index = self.index(row, 0) - self.beginInsertRows(index, 0, len(child_node.child_nodes)) - self.endInsertRows() + # Also announce the frame grandchildren. + array_index = self.index(array_row, 0, dataset_index) + num_frames = len(array_node.child_nodes) + if num_frames > 0: + self.beginInsertRows(array_index, 0, num_frames - 1) + self.endInsertRows() - def refresh_array(self, row: int) -> None: - top_left = self.index(row, 0) - bottom_right = self.index(row, self.columnCount() - 1) + def refresh_array(self, dataset_row: int, array_row: int) -> None: + dataset_index = self.index(dataset_row, 0) + if not dataset_index.isValid(): + return + + top_left = self.index(array_row, 0, dataset_index) + bottom_right = self.index(array_row, self.columnCount() - 1, dataset_index) self.dataChanged.emit(top_left, bottom_right) num_rows = self.rowCount(top_left) num_cols = self.columnCount(top_left) + if num_rows > 0: + child_top_left = self.index(0, 0, top_left) + child_bottom_right = self.index(num_rows - 1, num_cols - 1, top_left) + self.dataChanged.emit(child_top_left, child_bottom_right) + + def refresh_dataset(self, dataset_row: int) -> None: + dataset_index = self.index(dataset_row, 0) + if not dataset_index.isValid(): + return + bottom_right = self.index(dataset_row, self.columnCount() - 1) + self.dataChanged.emit(dataset_index, bottom_right) + + def _rebroadcast_counts(self) -> None: + num_rows = self.rowCount() + if num_rows == 0: + return + top_left = self.index(0, 1) + bottom_right = self.index(num_rows - 1, 1) + self.dataChanged.emit(top_left, bottom_right) - child_top_left = self.index(0, 0, top_left) - child_bottom_right = self.index(num_rows - 1, num_cols - 1, top_left) - self.dataChanged.emit(child_top_left, child_bottom_right) + for dataset_row in range(num_rows): + dataset_index = self.index(dataset_row, 0) + num_arrays = self.rowCount(dataset_index) + if num_arrays == 0: + continue + array_top_left = self.index(0, 1, dataset_index) + array_bottom_right = self.index(num_arrays - 1, 1, dataset_index) + self.dataChanged.emit(array_top_left, array_bottom_right) + + for array_row in range(num_arrays): + array_index = self.index(array_row, 0, dataset_index) + num_frames = self.rowCount(array_index) + if num_frames == 0: + continue + frame_top_left = self.index(0, 1, array_index) + frame_bottom_right = self.index(num_frames - 1, 1, array_index) + self.dataChanged.emit(frame_top_left, frame_bottom_right) + + def dataset_row_for_index(self, index: QModelIndex) -> int | None: + """Return the dataset row that contains the given tree index, or None.""" + if not index.isValid(): + return None + node = index.internalPointer() + return _find_containing_dataset_row(node) @overload def parent(self, child: QModelIndex) -> QModelIndex: ... @@ -134,7 +304,7 @@ def parent(self, child: QModelIndex | None = None) -> QModelIndex | QObject: child_node = child.internalPointer() parent_node = child_node.parent_node - if parent_node is not self._nodes: + if parent_node is not None and parent_node is not self._root: return self.createIndex(parent_node.get_row(), 0, parent_node) return QModelIndex() @@ -149,26 +319,112 @@ def headerData( # noqa: N802 return self._header[section] def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: - if index.isValid(): - node = index.internalPointer() - - if role == Qt.ItemDataRole.DisplayRole: - match index.column(): - case 0: - return node.get_label() - case 1: - return str(node.get_counts()) - case 2: - return node.get_nframes() - case 3: - return f'{node.get_nbytes() / BYTES_PER_MEGABYTE:.2f}' - elif role == Qt.ItemDataRole.UserRole: - if index.column() == 1: - return int(100 * node.get_counts()) // int(self._max_counts) + if not index.isValid(): + return None + + node = index.internalPointer() + column = index.column() + + if role == Qt.ItemDataRole.DisplayRole: + match column: + case 0: + return node.get_label() + case 1: + return str(node.get_counts()) + case 2: + return node.get_nframes() + case 3: + return f'{node.get_nbytes() / BYTES_PER_MEGABYTE:.2f}' + case _: + return self._dataset_column_display(node, column) + elif role == Qt.ItemDataRole.EditRole: + if column == _COL_LABEL and isinstance(node, _DatasetTreeNode): + return node.get_label() + elif role == Qt.ItemDataRole.UserRole: + if column == _COL_COUNTS: + return int(100 * node.get_counts()) // int(self._max_counts) + return None + + def _dataset_column_display(self, node: _TreeNode, column: int) -> Any: + # Columns 4..10 apply only to dataset rows; array/frame nodes render blank. + if not isinstance(node, _DatasetTreeNode): + return None + + match column: + case 4: + return node.get_detector_extent().width_px + case 5: + return node.get_detector_extent().height_px + case 6: + return f'{node.get_raw_pixel_geometry().width_m * 1e6:.4g}' + case 7: + return f'{node.get_raw_pixel_geometry().height_m * 1e6:.4g}' + case 8: + return f'{node.get_processed_pixel_geometry(self._sizer).width_m * 1e6:.4g}' + case 9: + return f'{node.get_processed_pixel_geometry(self._sizer).height_m * 1e6:.4g}' + case 10: + return node.get_num_bad_pixels() + return None + + def flags(self, index: QModelIndex) -> Qt.ItemFlags: # noqa: N802 + base = super().flags(index) + if not index.isValid(): + return base + if index.column() != _COL_LABEL: + return base + node = index.internalPointer() + if not isinstance(node, _DatasetTreeNode): + return base + return base | Qt.ItemFlag.ItemIsEditable + + def setData( # noqa: N802 + self, + index: QModelIndex, + value: Any, + role: int = Qt.ItemDataRole.EditRole, + ) -> bool: + if role != Qt.ItemDataRole.EditRole or not index.isValid(): + return False + if index.column() != _COL_LABEL: + return False + node = index.internalPointer() + if not isinstance(node, _DatasetTreeNode): + return False + + new_name = str(value).strip() + if not new_name: + return False + + dataset = node.get_dataset() + if new_name == dataset.get_name(): + return False + + unique_name = self._repository.create_unique_name(new_name) + dataset.set_name(unique_name) + row = node.get_row() + top_left = self.index(row, _COL_LABEL) + self.dataChanged.emit(top_left, top_left, [Qt.ItemDataRole.DisplayRole]) + return True + + def refresh_processed_columns(self) -> None: + """Emit dataChanged for the processed pixel-size columns of every dataset row. + + Called when the PatternSizer notifies (binning / transpose toggled or bin + size edited) — since the processed columns are derived from the sizer + transforms applied to each dataset's raw geometry, they all need to refresh + even though no dataset itself changed. + """ + num_rows = self.rowCount() + if num_rows == 0: + return + top_left = self.index(0, _COL_PROCESSED_PIXEL_WIDTH_UM) + bottom_right = self.index(num_rows - 1, _COL_PROCESSED_PIXEL_HEIGHT_UM) + self.dataChanged.emit(top_left, bottom_right) def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex: if self.hasIndex(row, column, parent): - parent_node = parent.internalPointer() if parent.isValid() else self._nodes + parent_node = parent.internalPointer() if parent.isValid() else self._root child_node = parent_node.child_nodes[row] if child_node: @@ -177,7 +433,7 @@ def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> Q return QModelIndex() def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 - node = parent.internalPointer() if parent.isValid() else self._nodes + node = parent.internalPointer() if parent.isValid() else self._root return len(node.child_nodes) def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 diff --git a/src/ptychodus/controller/diffraction/dataset_editor.py b/src/ptychodus/controller/diffraction/dataset_editor.py new file mode 100644 index 000000000..8041fa196 --- /dev/null +++ b/src/ptychodus/controller/diffraction/dataset_editor.py @@ -0,0 +1,253 @@ +from __future__ import annotations +import logging +import math +from typing import Any, overload + +from PyQt5.QtCore import Qt, QAbstractItemModel, QAbstractTableModel, QModelIndex, QObject +from PyQt5.QtGui import QBrush +from PyQt5.QtWidgets import QWidget + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.tree import SimpleTreeNode + +from ...model.diffraction import AssembledDiffractionDataset, DiffractionDatasetObserver +from ...view.diffraction import DatasetEditorDialog +from ..helpers import create_brush_for_editable_cell + +logger = logging.getLogger(__name__) + + +class SimpleTreeModel(QAbstractItemModel): + def __init__(self, root_node: SimpleTreeNode, parent: QObject | None = None) -> None: + super().__init__(parent) + self._root_node = root_node + + def set_root_node(self, root_node: SimpleTreeNode) -> None: + self.beginResetModel() + self._root_node = root_node + self.endResetModel() + + @overload + def parent(self, child: QModelIndex) -> QModelIndex: ... + + @overload + def parent(self) -> QObject: ... + + def parent(self, child: QModelIndex | None = None) -> QModelIndex | QObject: + if child is None: + return super().parent() + else: + value = QModelIndex() + + if child.isValid(): + child_item = child.internalPointer() + parent_item = child_item.parent_item + + if parent_item is self._root_node: + value = QModelIndex() + else: + value = self.createIndex(parent_item.row(), 0, parent_item) + + return value + + def headerData( # noqa: N802 + self, + section: int, + orientation: Qt.Orientation, + role: int = Qt.ItemDataRole.DisplayRole, + ) -> Any: + if orientation == Qt.Orientation.Horizontal and role == Qt.ItemDataRole.DisplayRole: + return self._root_node.data(section) + + def flags(self, index: QModelIndex) -> Qt.ItemFlags: + return super().flags(index) + + def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex: + value = QModelIndex() + + if self.hasIndex(row, column, parent): + parent_item = parent.internalPointer() if parent.isValid() else self._root_node + child_item = parent_item.child_items[row] + + if child_item: + value = self.createIndex(row, column, child_item) + + return value + + def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: + if index.isValid() and role == Qt.ItemDataRole.DisplayRole: + node = index.internalPointer() + return node.data(index.column()) + + def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + if parent.column() > 0: + return 0 + + node = self._root_node + + if parent.isValid(): + node = parent.internalPointer() + + return len(node.child_items) + + def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + node = self._root_node + + if parent.isValid(): + node = parent.internalPointer() + + return len(node.item_data) + + +_ROW_PIXEL_WIDTH_UM = 0 +_ROW_PIXEL_HEIGHT_UM = 1 + + +class DatasetPropertyTableModel(QAbstractTableModel): + """Two-row properties table over an AssembledDiffractionDataset: pixel width / height in µm. + + Setting either row calls set_pixel_geometry_override on the dataset. The dataset then + notifies observers via handle_pixel_geometry_changed(); the editor controller listens + and calls beginResetModel/endResetModel to refresh the displayed values. + """ + + def __init__( + self, + dataset: AssembledDiffractionDataset, + editable_item_brush: QBrush, + parent: QObject | None = None, + ) -> None: + super().__init__(parent) + self._dataset = dataset + self._editable_item_brush = editable_item_brush + self._header = ['Property', 'Value'] + self._properties = [ + 'Physical Pixel Width [µm]', + 'Physical Pixel Height [µm]', + ] + + def flags(self, index: QModelIndex) -> Qt.ItemFlags: + value = super().flags(index) + + if index.isValid() and index.column() == 1: + if not self._dataset.is_load_in_progress(): + value |= Qt.ItemFlag.ItemIsEditable + + return value + + def headerData( # noqa: N802 + self, + section: int, + orientation: Qt.Orientation, + role: int = Qt.ItemDataRole.DisplayRole, + ) -> Any: + if orientation == Qt.Orientation.Horizontal and role == Qt.ItemDataRole.DisplayRole: + return self._header[section] + + def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: + if not index.isValid(): + return None + + if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole: + geometry = self._dataset.get_raw_pixel_geometry() + match (index.column(), index.row()): + case (0, row): + return self._properties[row] + case (1, r) if r == _ROW_PIXEL_WIDTH_UM: + return f'{geometry.width_m * 1e6:.4g}' + case (1, r) if r == _ROW_PIXEL_HEIGHT_UM: + return f'{geometry.height_m * 1e6:.4g}' + elif role == Qt.ItemDataRole.BackgroundRole: + if index.flags() & Qt.ItemFlag.ItemIsEditable: + return self._editable_item_brush + + def setData(self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.EditRole) -> bool: # noqa: N802 + if role != Qt.ItemDataRole.EditRole or not index.isValid() or index.column() != 1: + return False + + try: + new_um = float(value) + except (TypeError, ValueError): + return False + if not math.isfinite(new_um) or new_um <= 0.0: + return False + new_m = new_um * 1e-6 + + current = self._dataset.get_raw_pixel_geometry() + if index.row() == _ROW_PIXEL_WIDTH_UM: + new_geometry = PixelGeometry(width_m=new_m, height_m=current.height_m) + elif index.row() == _ROW_PIXEL_HEIGHT_UM: + new_geometry = PixelGeometry(width_m=current.width_m, height_m=new_m) + else: + return False + + self._dataset.set_pixel_geometry_override(new_geometry) + return True + + def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + return len(self._properties) + + def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + return len(self._header) + + +class DatasetEditorViewController(DiffractionDatasetObserver): + def __init__( + self, + dataset: AssembledDiffractionDataset, + tree_model: SimpleTreeModel, + property_model: DatasetPropertyTableModel, + dialog: DatasetEditorDialog, + ) -> None: + super().__init__() + self._dataset = dataset + self._tree_model = tree_model + self._property_model = property_model + self._dialog = dialog + + @classmethod + def edit_dataset(cls, dataset: AssembledDiffractionDataset, parent: QWidget) -> None: + dialog = DatasetEditorDialog(parent) + dialog.setWindowTitle(f'Edit Dataset: {dataset.get_name()}') + + tree_model = SimpleTreeModel(dataset.get_layout()) + dialog.tree_view.setModel(tree_model) + tree_header = dialog.tree_view.header() + if tree_header is not None: + tree_header.setSectionResizeMode(tree_header.ResizeMode.ResizeToContents) + + editable_item_brush = create_brush_for_editable_cell(dialog.table_view) + property_model = DatasetPropertyTableModel(dataset, editable_item_brush) + dialog.table_view.setModel(property_model) + vertical_header = dialog.table_view.verticalHeader() + if vertical_header is not None: + vertical_header.hide() + table_header = dialog.table_view.horizontalHeader() + if table_header is not None: + table_header.setSectionResizeMode(table_header.ResizeMode.ResizeToContents) + dialog.table_view.resizeRowsToContents() + + controller = cls(dataset, tree_model, property_model, dialog) + dataset.add_observer(controller) + + dialog.finished.connect(controller._finish) + dialog.open() + dialog.adjustSize() + + def _finish(self, result: int) -> None: + self._dataset.remove_observer(self) + + def handle_array_inserted(self, index: int) -> None: + pass + + def handle_array_changed(self, index: int) -> None: + pass + + def handle_dataset_reloaded(self) -> None: + self._tree_model.set_root_node(self._dataset.get_layout()) + self._property_model.beginResetModel() + self._property_model.endResetModel() + + def handle_pixel_geometry_changed(self) -> None: + self._property_model.beginResetModel() + self._property_model.endResetModel() diff --git a/src/ptychodus/controller/diffraction/dataset_layout.py b/src/ptychodus/controller/diffraction/dataset_layout.py deleted file mode 100644 index 41c8429cd..000000000 --- a/src/ptychodus/controller/diffraction/dataset_layout.py +++ /dev/null @@ -1,125 +0,0 @@ -from typing import Any, overload - -from PyQt5.QtWidgets import QWidget -from PyQt5.QtCore import Qt, QAbstractItemModel, QModelIndex, QObject - -from ptychodus.api.tree import SimpleTreeNode - -from ...model.diffraction import AssembledDiffractionDataset, DiffractionDatasetObserver -from ...view.diffraction import DatasetFileLayoutDialog - - -class SimpleTreeModel(QAbstractItemModel): - def __init__(self, root_node: SimpleTreeNode, parent: QObject | None = None) -> None: - super().__init__(parent) - self._root_node = root_node - - def set_root_node(self, root_node: SimpleTreeNode) -> None: - self.beginResetModel() - self._root_node = root_node - self.endResetModel() - - @overload - def parent(self, child: QModelIndex) -> QModelIndex: ... - - @overload - def parent(self) -> QObject: ... - - def parent(self, child: QModelIndex | None = None) -> QModelIndex | QObject: - if child is None: - return super().parent() - else: - value = QModelIndex() - - if child.isValid(): - child_item = child.internalPointer() - parent_item = child_item.parent_item - - if parent_item is self._root_node: - value = QModelIndex() - else: - value = self.createIndex(parent_item.row(), 0, parent_item) - - return value - - def headerData( # noqa: N802 - self, - section: int, - orientation: Qt.Orientation, - role: int = Qt.ItemDataRole.DisplayRole, - ) -> Any: - if orientation == Qt.Orientation.Horizontal and role == Qt.ItemDataRole.DisplayRole: - return self._root_node.data(section) - - def flags(self, index: QModelIndex) -> Qt.ItemFlags: - return super().flags(index) - - def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex: - value = QModelIndex() - - if self.hasIndex(row, column, parent): - parent_item = parent.internalPointer() if parent.isValid() else self._root_node - child_item = parent_item.child_items[row] - - if child_item: - value = self.createIndex(row, column, child_item) - - return value - - def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: - if index.isValid() and role == Qt.ItemDataRole.DisplayRole: - node = index.internalPointer() - return node.data(index.column()) - - def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 - if parent.column() > 0: - return 0 - - node = self._root_node - - if parent.isValid(): - node = parent.internalPointer() - - return len(node.child_items) - - def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 - node = self._root_node - - if parent.isValid(): - node = parent.internalPointer() - - return len(node.item_data) - - -class DatasetLayoutViewController(DiffractionDatasetObserver): - def __init__(self, dataset: AssembledDiffractionDataset, tree_model: SimpleTreeModel) -> None: - super().__init__() - self._dataset = dataset - self._tree_model = tree_model - - @classmethod - def show_dialog(cls, dataset: AssembledDiffractionDataset, parent: QWidget) -> None: - tree_model = SimpleTreeModel(dataset.get_layout()) - controller = cls(dataset, tree_model) - dataset.add_observer(controller) - - dialog = DatasetFileLayoutDialog(parent) - dialog.setWindowTitle('Dataset File Layout') - dialog.tree_view.setModel(tree_model) - header = dialog.tree_view.header() - header.setSectionResizeMode(header.ResizeMode.ResizeToContents) - - controller._sync_model_to_view() - dialog.open() - - def _sync_model_to_view(self) -> None: - self._tree_model.set_root_node(self._dataset.get_layout()) - - def handle_array_inserted(self, index: int) -> None: - pass - - def handle_array_changed(self, index: int) -> None: - pass - - def handle_dataset_reloaded(self) -> None: - self._sync_model_to_view() diff --git a/src/ptychodus/controller/diffraction/detector_extent.py b/src/ptychodus/controller/diffraction/detector_extent.py new file mode 100644 index 000000000..079d95cce --- /dev/null +++ b/src/ptychodus/controller/diffraction/detector_extent.py @@ -0,0 +1,31 @@ +from ptychodus.api.geometry import ImageExtent +from ptychodus.api.observer import Observable + + +class DetectorExtentSource(Observable): + """Controller-owned observable holder for the current detector's pixel extent. + + Purely a UI-side signalling channel for the diffraction wizard's crop/bin spin + boxes, which need to redraw their bounds when a new dataset's extent becomes + known. Populated by the diffraction controller in response to dataset repository + inserts and by the wizard when it opens a new file. Not persisted; a freshly + launched session starts with ``None`` until data supplies an extent. + + Not accessed from the model layer — dataset processing derives the extent from + each dataset's own metadata (see ``AssembledDiffractionDataset.reload`` and + ``PatternSizer.get_prep_pipeline``), and product geometry derives it from the + dataset paired via ``ProductRepositoryItem.bind_dataset`` / ``unbind_dataset``. + """ + + def __init__(self) -> None: + super().__init__() + self._extent: ImageExtent | None = None + + def get_extent(self) -> ImageExtent | None: + return self._extent + + def set_extent(self, extent: ImageExtent | None) -> None: + if extent == self._extent: + return + self._extent = extent + self.notify_observers() diff --git a/src/ptychodus/controller/diffraction/wizard/bad_pixels.py b/src/ptychodus/controller/diffraction/wizard/bad_pixels.py new file mode 100644 index 000000000..fc456f6f1 --- /dev/null +++ b/src/ptychodus/controller/diffraction/wizard/bad_pixels.py @@ -0,0 +1,81 @@ +from PyQt5.QtWidgets import QFormLayout, QWizardPage + +from ptychodus.api.observer import Observable, Observer + +from ....model.diffraction import DetectorSettings, DiffractionAPI +from ....view.diffraction import OpenDatasetWizardBadPixelsPage +from ...data import FileDialogFactory +from .files import ( + OpenDatasetWizardBreadcrumbsViewController, + OpenDatasetWizardFilePathViewController, + OpenDatasetWizardFileTypeViewController, + OpenDatasetWizardLocationViewController, +) + + +class OpenDatasetWizardBadPixelsViewController(Observer): + """Bad-pixel file chooser — mirrors the Files page layout. + + Binds to :class:`DetectorSettings` for cross-session persistence (same + convention as :class:`DiffractionSettings.file_path` for the Files page). + The selection is applied by the wizard's next-button dispatcher on the + Bad Pixels → Processing transition via + :meth:`DiffractionAPI.apply_bad_pixels`. + + The page is always complete: leaving the file path invalid or unset is a + supported "no bad-pixel mask" workflow. + """ + + def __init__( + self, + detector_settings: DetectorSettings, + api: DiffractionAPI, + file_dialog_factory: FileDialogFactory, + ) -> None: + super().__init__() + self._detector_settings = detector_settings + self._file_dialog_factory = file_dialog_factory + + self._breadcrumbs_view_controller = OpenDatasetWizardBreadcrumbsViewController( + file_dialog_factory + ) + self._location_view_controller = OpenDatasetWizardLocationViewController( + detector_settings.bad_pixels_file_path, file_dialog_factory + ) + self._file_path_view_controller = OpenDatasetWizardFilePathViewController( + detector_settings.bad_pixels_file_path, file_dialog_factory + ) + self._file_type_view_controller = OpenDatasetWizardFileTypeViewController( + api.get_bad_pixels_file_reader_parameter() + ) + self._file_type_view_controller.get_parameter().add_observer(self) + + layout = QFormLayout() + layout.addRow(self._breadcrumbs_view_controller.get_widget()) + layout.addRow('Location:', self._location_view_controller.get_widget()) + layout.addRow(self._file_path_view_controller.get_widget()) + layout.addRow('File Type:', self._file_type_view_controller.get_widget()) + + self._page = OpenDatasetWizardBadPixelsPage() + self._page.setTitle('Choose Bad Pixels File') + self._page.setLayout(layout) + + self._handle_file_type_changed() + + def get_widget(self) -> QWizardPage: + return self._page + + def restart(self) -> None: + """Focus the file dialog on the current settings path if it points to a file.""" + current = self._detector_settings.bad_pixels_file_path.get_value() + if current.exists(): + self._file_dialog_factory.set_open_working_directory(current) + self._handle_file_type_changed() + + def _handle_file_type_changed(self) -> None: + name_filters = self._file_type_view_controller.get_name_filters() + self._file_path_view_controller.set_name_filters(name_filters) + + def _update(self, observable: Observable) -> None: + if observable is self._file_type_view_controller.get_parameter(): + self._handle_file_type_changed() diff --git a/src/ptychodus/controller/diffraction/wizard/core.py b/src/ptychodus/controller/diffraction/wizard/core.py index 9fd53f055..c38679b35 100644 --- a/src/ptychodus/controller/diffraction/wizard/core.py +++ b/src/ptychodus/controller/diffraction/wizard/core.py @@ -2,14 +2,22 @@ from PyQt5.QtWidgets import QWizard -from ....model.metadata import MetadataPresenter -from ....model.diffraction import DiffractionSettings, PatternSizer, DiffractionAPI +from ....api.diffraction import DiffractionMetadata +from ....model.diffraction import ( + DetectorSettings, + DiffractionAPI, + DiffractionDatasetRepository, + DiffractionSettings, +) +from ....model.product import ProductSettings from ....view.widgets import ExceptionDialog from ...data import FileDialogFactory +from ..detector_extent import DetectorExtentSource +from .bad_pixels import OpenDatasetWizardBadPixelsViewController from .files import OpenDatasetWizardFilesViewController from .metadata import OpenDatasetWizardMetadataViewController -from .patterns import OpenDatasetWizardPatternsViewController +from .processing import OpenDatasetWizardProcessingViewController logger = logging.getLogger(__name__) @@ -18,25 +26,39 @@ class OpenDatasetWizardController: def __init__( self, settings: DiffractionSettings, - sizer: PatternSizer, + detector_settings: DetectorSettings, + product_settings: ProductSettings, + extent_source: DetectorExtentSource, api: DiffractionAPI, - metadata_presenter: MetadataPresenter, + repository: DiffractionDatasetRepository, file_dialog_factory: FileDialogFactory, ) -> None: self._api = api + self._detector_settings = detector_settings + self._repository = repository + self._pending_dataset_index = -1 self._file_view_controller = OpenDatasetWizardFilesViewController( settings, api, file_dialog_factory ) - self._metadata_view_controller = OpenDatasetWizardMetadataViewController(metadata_presenter) - self._patterns_view_controller = OpenDatasetWizardPatternsViewController( - settings, sizer, file_dialog_factory + self._metadata_view_controller = OpenDatasetWizardMetadataViewController( + detector_settings, + settings, + product_settings, + self._get_pending_metadata, + ) + self._bad_pixels_view_controller = OpenDatasetWizardBadPixelsViewController( + detector_settings, api, file_dialog_factory + ) + self._processing_view_controller = OpenDatasetWizardProcessingViewController( + settings, extent_source, file_dialog_factory ) self._wizard = QWizard() self._wizard.setWindowTitle('Open Dataset') self._wizard.addPage(self._file_view_controller.get_widget()) self._wizard.addPage(self._metadata_view_controller.get_widget()) - self._wizard.addPage(self._patterns_view_controller.get_widget()) + self._wizard.addPage(self._bad_pixels_view_controller.get_widget()) + self._wizard.addPage(self._processing_view_controller.get_widget()) next_button = self._wizard.button(QWizard.WizardButton.NextButton) @@ -52,22 +74,49 @@ def __init__( else: finish_button.clicked.connect(self._execute_finish_button_action) + def _get_pending_metadata(self) -> DiffractionMetadata: + if self._pending_dataset_index < 0: + return DiffractionMetadata.create_null() + return self._repository[self._pending_dataset_index].get_metadata() + def _execute_next_button_action(self) -> None: + # Handlers fire AFTER Qt has advanced the wizard, so currentPage() is the + # page the user just arrived on. The branches below therefore describe + # actions taken on the "leaving X → arriving Y" transition: page = self._wizard.currentPage() if page is self._metadata_view_controller.get_widget(): - self._file_view_controller.open_dataset() - elif page is self._patterns_view_controller.get_widget(): + # Files → Metadata: read the dataset file so metadata checkboxes + # populate for the page that just became visible. + self._pending_dataset_index = self._file_view_controller.open_dataset() + self._metadata_view_controller.refresh() + elif page is self._bad_pixels_view_controller.get_widget(): + # Metadata → Bad Pixels: apply the metadata-import selections. Seed + # the bad-pixels file browser to focus on the current settings path + # if it points to a valid file. self._metadata_view_controller.import_metadata() + self._bad_pixels_view_controller.restart() + elif page is self._processing_view_controller.get_widget(): + # Bad Pixels → Processing: load and apply the bad-pixels mask (if + # any) to the pending dataset before the user configures processing. + if self._pending_dataset_index >= 0: + self._api.apply_bad_pixels( + self._pending_dataset_index, + self._detector_settings.bad_pixels_file_path.get_value(), + self._detector_settings.bad_pixels_file_type.get_value(), + ) def _execute_finish_button_action(self) -> None: + if self._pending_dataset_index < 0: + return try: - self._api.load_all_arrays() + self._api.load_all_arrays(dataset_index=self._pending_dataset_index) except Exception as exc: logger.exception(exc) ExceptionDialog.show_exception('Open Dataset', exc) def open_dataset(self) -> None: + self._pending_dataset_index = -1 self._wizard.restart() self._file_view_controller.restart() self._wizard.show() diff --git a/src/ptychodus/controller/diffraction/wizard/files.py b/src/ptychodus/controller/diffraction/wizard/files.py index 59080bf79..f0876beaa 100644 --- a/src/ptychodus/controller/diffraction/wizard/files.py +++ b/src/ptychodus/controller/diffraction/wizard/files.py @@ -1,5 +1,6 @@ from collections.abc import Sequence from pathlib import Path +from typing import Any import logging import re @@ -7,7 +8,6 @@ from PyQt5.QtWidgets import ( QAbstractItemView, QButtonGroup, - QComboBox, QFileSystemModel, QFormLayout, QHBoxLayout, @@ -21,12 +21,14 @@ from ptychodus.api.observer import Observable, Observer from ptychodus.api.parametric import PathParameter +from ptychodus.api.plugins import PluginChooserParameter from ....model.diffraction import DiffractionAPI, DiffractionSettings from ....view.diffraction import OpenDatasetWizardPage from ....view.widgets import ExceptionDialog from ...data import FileDialogFactory from ...helpers import connect_current_changed_signal +from ...parametric import ComboBoxParameterViewController logger = logging.getLogger(__name__) @@ -219,37 +221,29 @@ def _update(self, observable: Observable) -> None: self._sync_model_to_view() -class OpenDatasetWizardFileTypeViewController(Observable, Observer): - def __init__(self, api: DiffractionAPI) -> None: - super().__init__() - self._file_reader_chooser = api.get_file_reader_chooser() - self._file_reader_chooser.add_observer(self) - self._combo_box = QComboBox() +class OpenDatasetWizardFileTypeViewController: + """Binds a file-reader chooser to a combo box and derives the file-name filters. - for plugin in self._file_reader_chooser: - self._combo_box.addItem(plugin.display_name) + Observe :meth:`get_parameter` to react to the selection changing; the parameter + is the chooser's display-name view, which is also the string the glob patterns + are parsed out of. + """ - self._sync_model_to_view() - self._combo_box.textActivated.connect(self._handle_text_activated) + def __init__(self, file_reader_parameter: PluginChooserParameter[Any]) -> None: + self._file_reader_parameter = file_reader_parameter + self._view_controller = ComboBoxParameterViewController( + file_reader_parameter, file_reader_parameter.choices() + ) + + def get_parameter(self) -> PluginChooserParameter[Any]: + return self._file_reader_parameter def get_name_filters(self) -> Sequence[str]: - text = self._combo_box.currentText() - z = re.search(r'\((.+)\)', text) + z = re.search(r'\((.+)\)', self._file_reader_parameter.get_value()) return z.group(1).split() if z else [] - def _handle_text_activated(self, text: str) -> None: - self._file_reader_chooser.set_current_plugin(text) - - def _sync_model_to_view(self) -> None: - self._combo_box.setCurrentText(self._file_reader_chooser.get_current_plugin().display_name) - def get_widget(self) -> QWidget: - return self._combo_box - - def _update(self, observable: Observable) -> None: - if observable is self._file_reader_chooser: - self._sync_model_to_view() - self.notify_observers() + return self._view_controller.get_widget() class OpenDatasetWizardFilesViewController(Observer): @@ -273,8 +267,10 @@ def __init__( self._file_path_view_controller = OpenDatasetWizardFilePathViewController( settings.file_path, file_dialog_factory ) - self._file_type_view_controller = OpenDatasetWizardFileTypeViewController(api) - self._file_type_view_controller.add_observer(self) + self._file_type_view_controller = OpenDatasetWizardFileTypeViewController( + api.get_file_reader_parameter() + ) + self._file_type_view_controller.get_parameter().add_observer(self) layout = QFormLayout() layout.addRow(self._breadcrumbs_view_controller.get_widget()) @@ -289,16 +285,17 @@ def __init__( self._sync_model_to_view() settings.file_path.add_observer(self) - def open_dataset(self) -> None: + def open_dataset(self) -> int: file_reader_chooser = self._api.get_file_reader_chooser() file_type = file_reader_chooser.get_current_plugin().simple_name file_path = self._settings.file_path.get_value() try: - self._api.open_patterns(file_path, file_type=file_type) + return self._api.open_patterns(file_path, file_type=file_type) except Exception as err: logger.exception(err) ExceptionDialog.show_exception('Open Dataset', err) + return -1 def get_widget(self) -> QWizardPage: return self._page @@ -325,5 +322,5 @@ def _sync_model_to_view(self) -> None: def _update(self, observable: Observable) -> None: if observable is self._settings.file_path: self._check_if_complete() - elif observable is self._file_type_view_controller: + elif observable is self._file_type_view_controller.get_parameter(): self._handle_file_type_changed() diff --git a/src/ptychodus/controller/diffraction/wizard/metadata.py b/src/ptychodus/controller/diffraction/wizard/metadata.py index 24f41d148..dc4435f2c 100644 --- a/src/ptychodus/controller/diffraction/wizard/metadata.py +++ b/src/ptychodus/controller/diffraction/wizard/metadata.py @@ -1,81 +1,306 @@ -from PyQt5.QtWidgets import QWizardPage +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Any -from ptychodus.api.observer import Observable, Observer +from PyQt5.QtCore import QAbstractTableModel, QModelIndex, QObject, Qt +from PyQt5.QtWidgets import QHeaderView, QWizardPage -from ....model.metadata import MetadataPresenter +from ptychodus.api.diffraction import DiffractionMetadata + +from ....model.diffraction import DetectorSettings, DiffractionSettings +from ....model.product import ProductSettings from ....view.diffraction import OpenDatasetWizardMetadataPage -class OpenDatasetWizardMetadataViewController(Observer): - def __init__(self, presenter: MetadataPresenter) -> None: - super().__init__() - self._presenter = presenter - self._page = OpenDatasetWizardMetadataPage() +@dataclass(frozen=True) +class _MetadataRow: + name: str + is_present: Callable[[DiffractionMetadata], bool] + format_value: Callable[[DiffractionMetadata], str] + apply: Callable[[DiffractionMetadata], None] - presenter.add_observer(self) - self._sync_model_to_view() - self._page._set_complete(True) - def import_metadata(self) -> None: - if self._page.detector_extent_check_box.isChecked(): - self._presenter.sync_detector_extent() +class MetadataTableModel(QAbstractTableModel): + _HEADER = ('Metadata', 'Sync', 'Value') - if self._page.detector_pixel_size_check_box.isChecked(): - self._presenter.sync_detector_pixel_size() + def __init__(self, parent: QObject | None = None) -> None: + super().__init__(parent) + self._rows: list[_MetadataRow] = [] + self._metadata = DiffractionMetadata.create_null() + self._checked_rows: set[int] = set() - if self._page.detector_distance_check_box.isChecked(): - self._presenter.sync_detector_distance() + def set_rows(self, rows: Sequence[_MetadataRow], metadata: DiffractionMetadata) -> None: + self.beginResetModel() + self._rows = list(rows) + self._metadata = metadata + self._checked_rows = set(range(len(self._rows))) + self.endResetModel() - self._presenter.sync_pattern_crop( - sync_center=self._page.pattern_crop_center_check_box.isChecked(), - sync_extent=self._page.pattern_crop_extent_check_box.isChecked(), - ) + def checked_rows(self) -> list[_MetadataRow]: + return [self._rows[i] for i in sorted(self._checked_rows)] + + def headerData( # noqa: N802 + self, + section: int, + orientation: Qt.Orientation, + role: int = Qt.ItemDataRole.DisplayRole, + ) -> Any: + if role == Qt.ItemDataRole.DisplayRole and orientation == Qt.Orientation.Horizontal: + return self._HEADER[section] + + def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: + if not index.isValid(): + return None - if self._page.probe_energy_check_box.isChecked(): - self._presenter.sync_probe_energy() + row = self._rows[index.row()] + column = index.column() - if self._page.probe_photon_count_check_box.isChecked(): - self._presenter.sync_probe_photon_count() + if role == Qt.ItemDataRole.DisplayRole: + if column == 0: + return row.name + elif column == 2: + return row.format_value(self._metadata) + elif role == Qt.ItemDataRole.CheckStateRole and column == 1: + return ( + Qt.CheckState.Checked + if index.row() in self._checked_rows + else Qt.CheckState.Unchecked + ) - if self._page.exposure_time_check_box.isChecked(): - self._presenter.sync_exposure_time() + return None - def _sync_model_to_view(self) -> None: - can_sync_detector_extent = self._presenter.can_sync_detector_extent() - self._page.detector_extent_check_box.setVisible(can_sync_detector_extent) - self._page.detector_extent_check_box.setChecked(can_sync_detector_extent) + def flags(self, index: QModelIndex) -> Qt.ItemFlags: + value = super().flags(index) - can_sync_detector_pixel_size = self._presenter.can_sync_detector_pixel_size() - self._page.detector_pixel_size_check_box.setVisible(can_sync_detector_pixel_size) - self._page.detector_pixel_size_check_box.setChecked(can_sync_detector_pixel_size) + if index.isValid() and index.column() == 1: + value |= Qt.ItemFlag.ItemIsUserCheckable - can_sync_detector_distance = self._presenter.can_sync_detector_distance() - self._page.detector_distance_check_box.setVisible(can_sync_detector_distance) - self._page.detector_distance_check_box.setChecked(can_sync_detector_distance) + return value - can_sync_pattern_crop_center = self._presenter.can_sync_pattern_crop_center() - self._page.pattern_crop_center_check_box.setVisible(can_sync_pattern_crop_center) - self._page.pattern_crop_center_check_box.setChecked(can_sync_pattern_crop_center) + def setData( # noqa: N802 + self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.EditRole + ) -> bool: + if index.isValid() and index.column() == 1 and role == Qt.ItemDataRole.CheckStateRole: + if value == Qt.CheckState.Checked: + self._checked_rows.add(index.row()) + else: + self._checked_rows.discard(index.row()) + self.dataChanged.emit(index, index) + return True + return False + + def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + return len(self._rows) + + def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + return len(self._HEADER) + + +class OpenDatasetWizardMetadataViewController: + def __init__( + self, + detector_settings: DetectorSettings, + diffraction_settings: DiffractionSettings, + product_settings: ProductSettings, + get_metadata: Callable[[], DiffractionMetadata], + ) -> None: + self._detector_settings = detector_settings + self._diffraction_settings = diffraction_settings + self._product_settings = product_settings + self._get_metadata = get_metadata + self._page = OpenDatasetWizardMetadataPage() + self._table_model = MetadataTableModel() + self._all_rows: tuple[_MetadataRow, ...] = self._build_rows() - can_sync_pattern_crop_extent = self._presenter.can_sync_pattern_crop_extent() - self._page.pattern_crop_extent_check_box.setVisible(can_sync_pattern_crop_extent) - self._page.pattern_crop_extent_check_box.setChecked(can_sync_pattern_crop_extent) + self._page.table_view.setModel(self._table_model) - can_sync_probe_energy = self._presenter.can_sync_probe_energy() - self._page.probe_energy_check_box.setVisible(can_sync_probe_energy) - self._page.probe_energy_check_box.setChecked(can_sync_probe_energy) + horizontal_header = self._page.table_view.horizontalHeader() + if horizontal_header is not None: + horizontal_header.setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents) + horizontal_header.setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents) + horizontal_header.setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch) - can_sync_probe_photon_count = self._presenter.can_sync_probe_photon_count() - self._page.probe_photon_count_check_box.setVisible(can_sync_probe_photon_count) - self._page.probe_photon_count_check_box.setChecked(can_sync_probe_photon_count) + vertical_header = self._page.table_view.verticalHeader() + if vertical_header is not None: + vertical_header.setVisible(False) - can_sync_exposure_time = self._presenter.can_sync_exposure_time() - self._page.exposure_time_check_box.setVisible(can_sync_exposure_time) - self._page.exposure_time_check_box.setChecked(can_sync_exposure_time) + self.refresh() + + def _build_rows(self) -> tuple[_MetadataRow, ...]: + return ( + _MetadataRow( + name='Detector Pixel Size', + is_present=lambda m: m.detector_pixel_geometry is not None, + format_value=self._format_pixel_size, + apply=self._apply_pixel_size, + ), + _MetadataRow( + name='Detector Distance', + is_present=lambda m: m.detector_distance_m is not None, + format_value=self._format_detector_distance, + apply=self._apply_detector_distance, + ), + _MetadataRow( + name='Pattern Crop Center', + is_present=lambda m: m.crop_center is not None or m.detector_extent is not None, + format_value=self._format_crop_center, + apply=self._apply_crop_center, + ), + _MetadataRow( + name='Pattern Crop Extent', + is_present=lambda m: m.detector_extent is not None, + format_value=self._format_crop_extent, + apply=self._apply_crop_extent, + ), + _MetadataRow( + name='Probe Energy', + is_present=lambda m: m.probe_energy_eV is not None, + format_value=self._format_probe_energy, + apply=self._apply_probe_energy, + ), + _MetadataRow( + name='Probe Photon Count', + is_present=lambda m: m.probe_photon_count is not None, + format_value=self._format_probe_photon_count, + apply=self._apply_probe_photon_count, + ), + _MetadataRow( + name='Exposure Time', + is_present=lambda m: m.exposure_time_s is not None, + format_value=self._format_exposure_time, + apply=self._apply_exposure_time, + ), + ) + + def import_metadata(self) -> None: + metadata = self._get_metadata() + + for row in self._table_model.checked_rows(): + row.apply(metadata) + + def refresh(self) -> None: + metadata = self._get_metadata() + visible = [row for row in self._all_rows if row.is_present(metadata)] + self._table_model.set_rows(visible, metadata) def get_widget(self) -> QWizardPage: return self._page - def _update(self, observable: Observable) -> None: - if observable is self._presenter: - self._sync_model_to_view() + # --- Detector pixel size --- + + @staticmethod + def _format_pixel_size(metadata: DiffractionMetadata) -> str: + pixel_geometry = metadata.detector_pixel_geometry + if pixel_geometry is None: + return '—' + return f'{pixel_geometry.width_m * 1e6:.3f} × {pixel_geometry.height_m * 1e6:.3f} µm' + + def _apply_pixel_size(self, metadata: DiffractionMetadata) -> None: + pixel_geometry = metadata.detector_pixel_geometry + if pixel_geometry is not None: + self._detector_settings.pixel_width_m.set_value(pixel_geometry.width_m) + self._detector_settings.pixel_height_m.set_value(pixel_geometry.height_m) + + # --- Detector distance --- + + @staticmethod + def _format_detector_distance(metadata: DiffractionMetadata) -> str: + distance_m = metadata.detector_distance_m + return '—' if distance_m is None else f'{distance_m:.3f} m' + + def _apply_detector_distance(self, metadata: DiffractionMetadata) -> None: + distance_m = metadata.detector_distance_m + if distance_m: + self._product_settings.detector_distance_m.set_value(distance_m) + + # --- Crop center --- + + @staticmethod + def _format_crop_center(metadata: DiffractionMetadata) -> str: + crop_center = metadata.crop_center + if crop_center is not None: + return f'({crop_center.position_x_px}, {crop_center.position_y_px}) px' + extent = metadata.detector_extent + if extent is not None: + return f'({int(extent.width_px) // 2}, {int(extent.height_px) // 2}) px' + return '—' + + def _apply_crop_center(self, metadata: DiffractionMetadata) -> None: + crop_center = metadata.crop_center + if crop_center is not None: + self._diffraction_settings.crop_center_x_px.set_value(crop_center.position_x_px) + self._diffraction_settings.crop_center_y_px.set_value(crop_center.position_y_px) + elif metadata.detector_extent is not None: + self._diffraction_settings.crop_center_x_px.set_value( + int(metadata.detector_extent.width_px) // 2 + ) + self._diffraction_settings.crop_center_y_px.set_value( + int(metadata.detector_extent.height_px) // 2 + ) + + # --- Crop extent --- + + @staticmethod + def _format_crop_extent(metadata: DiffractionMetadata) -> str: + extent = metadata.detector_extent + if extent is None: + return '—' + return f'{int(extent.width_px)} × {int(extent.height_px)} px' + + def _apply_crop_extent(self, metadata: DiffractionMetadata) -> None: + extent = metadata.detector_extent + if extent is None: + return + + center_x = self._diffraction_settings.crop_center_x_px.get_value() + center_y = self._diffraction_settings.crop_center_y_px.get_value() + + extent_x = int(extent.width_px) + extent_y = int(extent.height_px) + + max_radius_x = min(center_x, extent_x - center_x) + max_radius_y = min(center_y, extent_y - center_y) + max_radius = min(max_radius_x, max_radius_y) + crop_diameter = 1 + + while crop_diameter < max_radius: + crop_diameter <<= 1 + + self._diffraction_settings.crop_width_px.set_value(crop_diameter) + self._diffraction_settings.crop_height_px.set_value(crop_diameter) + + # --- Probe energy --- + + @staticmethod + def _format_probe_energy(metadata: DiffractionMetadata) -> str: + energy_eV = metadata.probe_energy_eV # noqa: N806 + return '—' if energy_eV is None else f'{energy_eV:.3f} eV' + + def _apply_probe_energy(self, metadata: DiffractionMetadata) -> None: + energy_eV = metadata.probe_energy_eV # noqa: N806 + if energy_eV: + self._product_settings.probe_energy_eV.set_value(energy_eV) + + # --- Probe photon count --- + + @staticmethod + def _format_probe_photon_count(metadata: DiffractionMetadata) -> str: + count = metadata.probe_photon_count + return '—' if count is None else f'{count:g}' + + def _apply_probe_photon_count(self, metadata: DiffractionMetadata) -> None: + count = metadata.probe_photon_count + if count: + self._product_settings.probe_photon_count.set_value(count) + + # --- Exposure time --- + + @staticmethod + def _format_exposure_time(metadata: DiffractionMetadata) -> str: + exposure_time_s = metadata.exposure_time_s + return '—' if exposure_time_s is None else f'{exposure_time_s * 1e3:.3f} ms' + + def _apply_exposure_time(self, metadata: DiffractionMetadata) -> None: + exposure_time_s = metadata.exposure_time_s + if exposure_time_s: + self._product_settings.exposure_time_s.set_value(exposure_time_s) diff --git a/src/ptychodus/controller/diffraction/wizard/patterns.py b/src/ptychodus/controller/diffraction/wizard/patterns.py deleted file mode 100644 index bdd3be6d2..000000000 --- a/src/ptychodus/controller/diffraction/wizard/patterns.py +++ /dev/null @@ -1,296 +0,0 @@ -from typing import Final - -from PyQt5.QtWidgets import ( - QFormLayout, - QGridLayout, - QGroupBox, - QLabel, - QSpinBox, - QVBoxLayout, - QWidget, - QWizardPage, -) - -from ptychodus.api.observer import Observable - -from ....model.diffraction import DiffractionSettings, PatternSizer -from ....view.diffraction import OpenDatasetWizardPage - -from ...data import FileDialogFactory -from ...parametric import ( - CheckBoxParameterViewController, - CheckableGroupBoxParameterViewController, - ParameterViewController, - PathParameterViewController, - SpinBoxParameterViewController, -) - - -class PatternMemoryMapViewController(CheckableGroupBoxParameterViewController): - def __init__( - self, settings: DiffractionSettings, file_dialog_factory: FileDialogFactory - ) -> None: - super().__init__(settings.memmap_enabled, 'Memory Map Diffraction Data') - self._view_controller = PathParameterViewController.create_directory_chooser( - settings.scratch_directory, file_dialog_factory - ) - - layout = QFormLayout() - layout.addRow('Scratch Directory:', self._view_controller.get_widget()) - self.get_widget().setLayout(layout) - - -class PatternCropViewController(CheckableGroupBoxParameterViewController): - def __init__( - self, - settings: DiffractionSettings, - sizer: PatternSizer, - ) -> None: - super().__init__(settings.crop_enabled, 'Crop') - self._settings = settings - self._sizer = sizer - - self._center_x_spin_box = QSpinBox() - self._center_y_spin_box = QSpinBox() - self._width_spin_box = QSpinBox() - self._height_spin_box = QSpinBox() - - layout = QGridLayout() - layout.addWidget(QLabel('Center:'), 0, 0) - layout.addWidget(self._center_x_spin_box, 0, 1) - layout.addWidget(self._center_y_spin_box, 0, 2) - layout.addWidget(QLabel('Extent:'), 1, 0) - layout.addWidget(self._width_spin_box, 1, 1) - layout.addWidget(self._height_spin_box, 1, 2) - layout.setColumnStretch(1, 1) - layout.setColumnStretch(2, 1) - self.get_widget().setLayout(layout) - - self._sync_model_to_view() - - self._center_x_spin_box.valueChanged.connect(settings.crop_center_x_px.set_value) - self._center_y_spin_box.valueChanged.connect(settings.crop_center_y_px.set_value) - self._width_spin_box.valueChanged.connect(settings.crop_width_px.set_value) - self._height_spin_box.valueChanged.connect(settings.crop_height_px.set_value) - - sizer.add_observer(self) - - def _sync_model_to_view(self) -> None: - center_x = self._sizer.axis_x.get_crop_center() - center_y = self._sizer.axis_y.get_crop_center() - width = self._sizer.axis_x.get_crop_size() - height = self._sizer.axis_y.get_crop_size() - - center_x_limits = self._sizer.axis_x.get_crop_center_limits() - center_y_limits = self._sizer.axis_y.get_crop_center_limits() - width_limits = self._sizer.axis_x.get_crop_size_limits() - height_limits = self._sizer.axis_y.get_crop_size_limits() - - self._center_x_spin_box.blockSignals(True) - self._center_x_spin_box.setRange(center_x_limits.lower, center_x_limits.upper) - self._center_x_spin_box.setValue(center_x) - self._center_x_spin_box.blockSignals(False) - - self._center_y_spin_box.blockSignals(True) - self._center_y_spin_box.setRange(center_y_limits.lower, center_y_limits.upper) - self._center_y_spin_box.setValue(center_y) - self._center_y_spin_box.blockSignals(False) - - self._width_spin_box.blockSignals(True) - self._width_spin_box.setRange(width_limits.lower, width_limits.upper) - self._width_spin_box.setValue(width) - self._width_spin_box.blockSignals(False) - - self._height_spin_box.blockSignals(True) - self._height_spin_box.setRange(height_limits.lower, height_limits.upper) - self._height_spin_box.setValue(height) - self._height_spin_box.blockSignals(False) - - def _update(self, observable: Observable) -> None: - if observable is self._sizer: - self._sync_model_to_view() - else: - super()._update(observable) - - -class PatternBinningViewController(CheckableGroupBoxParameterViewController): - def __init__( - self, - settings: DiffractionSettings, - sizer: PatternSizer, - ) -> None: - super().__init__(settings.binning_enabled, 'Bin Pixels') - self._settings = settings - self._sizer = sizer - - self._bin_size_x_spin_box = QSpinBox() - self._bin_size_y_spin_box = QSpinBox() - - layout = QGridLayout() - layout.addWidget(QLabel('Bin Size:'), 0, 0) - layout.addWidget(self._bin_size_x_spin_box, 0, 1) - layout.addWidget(self._bin_size_y_spin_box, 0, 2) - layout.setColumnStretch(1, 1) - layout.setColumnStretch(2, 1) - self.get_widget().setLayout(layout) - - self._sync_model_to_view() - - self._bin_size_x_spin_box.valueChanged.connect(settings.bin_size_x.set_value) - self._bin_size_y_spin_box.valueChanged.connect(settings.bin_size_y.set_value) - - sizer.add_observer(self) - - def _sync_model_to_view(self) -> None: - bin_size_x = self._sizer.axis_x.get_bin_size() - bin_size_y = self._sizer.axis_y.get_bin_size() - - bin_size_x_limits = self._sizer.axis_x.get_bin_size_limits() - bin_size_y_limits = self._sizer.axis_y.get_bin_size_limits() - - self._bin_size_x_spin_box.blockSignals(True) - self._bin_size_x_spin_box.setRange(bin_size_x_limits.lower, bin_size_x_limits.upper) - self._bin_size_x_spin_box.setValue(bin_size_x) - self._bin_size_x_spin_box.blockSignals(False) - - self._bin_size_y_spin_box.blockSignals(True) - self._bin_size_y_spin_box.setRange(bin_size_y_limits.lower, bin_size_y_limits.upper) - self._bin_size_y_spin_box.setValue(bin_size_y) - self._bin_size_y_spin_box.blockSignals(False) - - def _update(self, observable: Observable) -> None: - if observable is self._sizer: - self._sync_model_to_view() - else: - super()._update(observable) - - -class PatternPaddingViewController(CheckableGroupBoxParameterViewController): - MAX_INT: Final[int] = 0x7FFFFFFF - - def __init__( - self, - settings: DiffractionSettings, - sizer: PatternSizer, - ) -> None: - super().__init__(settings.padding_enabled, 'Pad') - self._settings = settings - self._sizer = sizer - - self._pad_x_spin_box = QSpinBox() - self._pad_y_spin_box = QSpinBox() - - layout = QGridLayout() - layout.addWidget(QLabel('Padding:'), 0, 0) - layout.addWidget(self._pad_x_spin_box, 0, 1) - layout.addWidget(self._pad_y_spin_box, 0, 2) - layout.setColumnStretch(1, 1) - layout.setColumnStretch(2, 1) - self.get_widget().setLayout(layout) - - self._sync_model_to_view() - - self._pad_x_spin_box.valueChanged.connect(settings.pad_x.set_value) - self._pad_y_spin_box.valueChanged.connect(settings.pad_y.set_value) - - sizer.add_observer(self) - - def _sync_model_to_view(self) -> None: - pad_x = self._sizer.axis_x.get_pad_size() - pad_y = self._sizer.axis_y.get_pad_size() - - self._pad_x_spin_box.blockSignals(True) - self._pad_x_spin_box.setRange(0, self.MAX_INT) - self._pad_x_spin_box.setValue(pad_x) - self._pad_x_spin_box.blockSignals(False) - - self._pad_y_spin_box.blockSignals(True) - self._pad_y_spin_box.setRange(0, self.MAX_INT) - self._pad_y_spin_box.setValue(pad_y) - self._pad_y_spin_box.blockSignals(False) - - def _update(self, observable: Observable) -> None: - if observable is self._sizer: - self._sync_model_to_view() - else: - super()._update(observable) - - -class PatternTransformViewController: - def __init__( - self, settings: DiffractionSettings, file_dialog_factory: FileDialogFactory - ) -> None: - self._hflip_view_controller = CheckBoxParameterViewController( - settings.hflip, 'Flip Horizontal' - ) - self._vflip_view_controller = CheckBoxParameterViewController( - settings.vflip, 'Flip Vertical' - ) - self._transpose_view_controller = CheckBoxParameterViewController( - settings.transpose, 'Transpose' - ) - self._lower_bound_enabled_view_controller = CheckBoxParameterViewController( - settings.value_lower_bound_enabled, 'Value Lower Bound:' - ) - self._lower_bound_view_controller = SpinBoxParameterViewController( - settings.value_lower_bound - ) - self._upper_bound_enabled_view_controller = CheckBoxParameterViewController( - settings.value_upper_bound_enabled, 'Value upper Bound:' - ) - self._upper_bound_view_controller = SpinBoxParameterViewController( - settings.value_upper_bound - ) - - layout = QGridLayout() - layout.addWidget(QLabel('Axes:'), 0, 0) - layout.addWidget(self._hflip_view_controller.get_widget(), 0, 1) - layout.addWidget(self._vflip_view_controller.get_widget(), 0, 2) - layout.addWidget(self._transpose_view_controller.get_widget(), 0, 3) - layout.addWidget(self._lower_bound_enabled_view_controller.get_widget(), 1, 0) - layout.addWidget(self._lower_bound_view_controller.get_widget(), 1, 1, 1, 3) - layout.addWidget(self._upper_bound_enabled_view_controller.get_widget(), 2, 0) - layout.addWidget(self._upper_bound_view_controller.get_widget(), 2, 1, 1, 3) - layout.setColumnStretch(1, 1) - layout.setColumnStretch(2, 1) - layout.setColumnStretch(3, 1) - - self._widget = QGroupBox('Transform') - self._widget.setLayout(layout) - - def get_widget(self) -> QWidget: - return self._widget - - -class OpenDatasetWizardPatternsViewController(ParameterViewController): - def __init__( - self, - settings: DiffractionSettings, - sizer: PatternSizer, - file_dialog_factory: FileDialogFactory, - ) -> None: - self._memory_map_view_controller = PatternMemoryMapViewController( - settings, file_dialog_factory - ) - self._crop_view_controller = PatternCropViewController(settings, sizer) - self._binning_view_controller = PatternBinningViewController(settings, sizer) - self._padding_view_controller = PatternPaddingViewController(settings, sizer) - self._transform_view_controller = PatternTransformViewController( - settings, file_dialog_factory - ) - - layout = QVBoxLayout() - layout.addWidget(self._memory_map_view_controller.get_widget()) - layout.addWidget(self._crop_view_controller.get_widget()) - layout.addWidget(self._binning_view_controller.get_widget()) - layout.addWidget(self._padding_view_controller.get_widget()) - layout.addWidget(self._transform_view_controller.get_widget()) - layout.addStretch() - - self._page = OpenDatasetWizardPage() - self._page.setTitle('Pattern Processing') - self._page._set_complete(True) - self._page.setLayout(layout) - - def get_widget(self) -> QWizardPage: - return self._page diff --git a/src/ptychodus/controller/diffraction/wizard/processing.py b/src/ptychodus/controller/diffraction/wizard/processing.py new file mode 100644 index 000000000..be534889c --- /dev/null +++ b/src/ptychodus/controller/diffraction/wizard/processing.py @@ -0,0 +1,378 @@ +from PyQt5.QtWidgets import ( + QFormLayout, + QFrame, + QGridLayout, + QGroupBox, + QLabel, + QSpinBox, + QVBoxLayout, + QWidget, + QWizardPage, +) + +from ptychodus.api.geometry import Interval +from ptychodus.api.observer import Observable + +from ....model.diffraction import DiffractionSettings +from ....view.diffraction import OpenDatasetWizardPage + +from ...data import FileDialogFactory +from ..detector_extent import DetectorExtentSource +from ...parametric import ( + CheckBoxParameterViewController, + CheckableGroupBoxParameterViewController, + ParameterViewController, + PathParameterViewController, + SpinBoxParameterViewController, +) + + +def _crop_size_limits(det_size_px: int) -> Interval[int]: + return Interval[int](1, det_size_px) + + +def _crop_center_limits(det_size_px: int) -> Interval[int]: + return Interval[int](1, det_size_px) + + +def _effective_crop_size(det_size_px: int, requested: int, *, crop_enabled: bool) -> int: + """The crop dimension actually used by the pipeline: clamped to detector when cropping is on; + otherwise the full detector width/height.""" + return _crop_size_limits(det_size_px).clamp(requested) if crop_enabled else det_size_px + + +def _bin_size_limits(effective_crop_px: int) -> Interval[int]: + return Interval[int](1, effective_crop_px) + + +def _set_spin_box(box: QSpinBox, limits: Interval[int], value: int) -> None: + box.blockSignals(True) + box.setRange(limits.lower, limits.upper) + box.setValue(value) + box.blockSignals(False) + + +class StorageViewController(CheckableGroupBoxParameterViewController): + def __init__( + self, settings: DiffractionSettings, file_dialog_factory: FileDialogFactory + ) -> None: + super().__init__(settings.memmap_enabled, 'Memory Map Diffraction Data') + self._view_controller = PathParameterViewController.create_directory_chooser( + settings.scratch_directory, file_dialog_factory + ) + + layout = QFormLayout() + layout.addRow('Scratch Directory:', self._view_controller.get_widget()) + self.get_widget().setLayout(layout) + + +class CropViewController(CheckableGroupBoxParameterViewController): + def __init__( + self, + diffraction_settings: DiffractionSettings, + extent_source: DetectorExtentSource, + ) -> None: + super().__init__(diffraction_settings.crop_enabled, 'Crop') + self._diffraction_settings = diffraction_settings + self._extent_source = extent_source + + self._center_x_spin_box = QSpinBox() + self._center_y_spin_box = QSpinBox() + self._width_spin_box = QSpinBox() + self._height_spin_box = QSpinBox() + + layout = QGridLayout() + layout.addWidget(QLabel('Center:'), 0, 0) + layout.addWidget(self._center_x_spin_box, 0, 1) + layout.addWidget(self._center_y_spin_box, 0, 2) + layout.addWidget(QLabel('Extent:'), 1, 0) + layout.addWidget(self._width_spin_box, 1, 1) + layout.addWidget(self._height_spin_box, 1, 2) + layout.setColumnStretch(1, 1) + layout.setColumnStretch(2, 1) + self.get_widget().setLayout(layout) + + self._observed = ( + diffraction_settings.crop_center_x_px, + diffraction_settings.crop_center_y_px, + diffraction_settings.crop_width_px, + diffraction_settings.crop_height_px, + ) + + self._sync_model_to_view() + + self._center_x_spin_box.valueChanged.connect( + diffraction_settings.crop_center_x_px.set_value + ) + self._center_y_spin_box.valueChanged.connect( + diffraction_settings.crop_center_y_px.set_value + ) + self._width_spin_box.valueChanged.connect(diffraction_settings.crop_width_px.set_value) + self._height_spin_box.valueChanged.connect(diffraction_settings.crop_height_px.set_value) + + for parameter in self._observed: + parameter.add_observer(self) + extent_source.add_observer(self) + + def _sync_model_to_view(self) -> None: + extent = self._extent_source.get_extent() + spin_boxes = ( + self._center_x_spin_box, + self._center_y_spin_box, + self._width_spin_box, + self._height_spin_box, + ) + + if extent is None: + for box in spin_boxes: + box.setEnabled(False) + return + + for box in spin_boxes: + box.setEnabled(True) + + det_w = extent.width_px + det_h = extent.height_px + + _set_spin_box( + self._center_x_spin_box, + _crop_center_limits(det_w), + _crop_center_limits(det_w).clamp( + self._diffraction_settings.crop_center_x_px.get_value() + ), + ) + _set_spin_box( + self._center_y_spin_box, + _crop_center_limits(det_h), + _crop_center_limits(det_h).clamp( + self._diffraction_settings.crop_center_y_px.get_value() + ), + ) + _set_spin_box( + self._width_spin_box, + _crop_size_limits(det_w), + _crop_size_limits(det_w).clamp(self._diffraction_settings.crop_width_px.get_value()), + ) + _set_spin_box( + self._height_spin_box, + _crop_size_limits(det_h), + _crop_size_limits(det_h).clamp(self._diffraction_settings.crop_height_px.get_value()), + ) + + def _update(self, observable: Observable) -> None: + if observable in self._observed or observable is self._extent_source: + self._sync_model_to_view() + else: + super()._update(observable) + + +class BinningViewController(CheckableGroupBoxParameterViewController): + def __init__( + self, + diffraction_settings: DiffractionSettings, + extent_source: DetectorExtentSource, + ) -> None: + super().__init__(diffraction_settings.binning_enabled, 'Bin Pixels') + self._diffraction_settings = diffraction_settings + self._extent_source = extent_source + + self._bin_size_x_spin_box = QSpinBox() + self._bin_size_y_spin_box = QSpinBox() + + layout = QGridLayout() + layout.addWidget(QLabel('Bin Size:'), 0, 0) + layout.addWidget(self._bin_size_x_spin_box, 0, 1) + layout.addWidget(self._bin_size_y_spin_box, 0, 2) + layout.setColumnStretch(1, 1) + layout.setColumnStretch(2, 1) + self.get_widget().setLayout(layout) + + # Also observe crop_enabled + crop extents because the effective bin-size upper bound + # depends on the crop settings. + self._observed = ( + diffraction_settings.bin_size_x, + diffraction_settings.bin_size_y, + diffraction_settings.crop_enabled, + diffraction_settings.crop_width_px, + diffraction_settings.crop_height_px, + ) + + self._sync_model_to_view() + + self._bin_size_x_spin_box.valueChanged.connect(diffraction_settings.bin_size_x.set_value) + self._bin_size_y_spin_box.valueChanged.connect(diffraction_settings.bin_size_y.set_value) + + for parameter in self._observed: + parameter.add_observer(self) + extent_source.add_observer(self) + + def _sync_model_to_view(self) -> None: + extent = self._extent_source.get_extent() + if extent is None: + self._bin_size_x_spin_box.setEnabled(False) + self._bin_size_y_spin_box.setEnabled(False) + return + + self._bin_size_x_spin_box.setEnabled(True) + self._bin_size_y_spin_box.setEnabled(True) + + det_w = extent.width_px + det_h = extent.height_px + crop_enabled = self._diffraction_settings.crop_enabled.get_value() + binning_enabled = self._diffraction_settings.binning_enabled.get_value() + + effective_w = _effective_crop_size( + det_w, + self._diffraction_settings.crop_width_px.get_value(), + crop_enabled=crop_enabled, + ) + effective_h = _effective_crop_size( + det_h, + self._diffraction_settings.crop_height_px.get_value(), + crop_enabled=crop_enabled, + ) + + bin_x_limits = _bin_size_limits(effective_w) + bin_y_limits = _bin_size_limits(effective_h) + + bin_x_value = ( + bin_x_limits.clamp(self._diffraction_settings.bin_size_x.get_value()) + if binning_enabled + else 1 + ) + bin_y_value = ( + bin_y_limits.clamp(self._diffraction_settings.bin_size_y.get_value()) + if binning_enabled + else 1 + ) + + _set_spin_box(self._bin_size_x_spin_box, bin_x_limits, bin_x_value) + _set_spin_box(self._bin_size_y_spin_box, bin_y_limits, bin_y_value) + + def _update(self, observable: Observable) -> None: + if observable in self._observed or observable is self._extent_source: + self._sync_model_to_view() + else: + super()._update(observable) + + +class PaddingViewController(CheckableGroupBoxParameterViewController): + def __init__(self, diffraction_settings: DiffractionSettings) -> None: + super().__init__(diffraction_settings.padding_enabled, 'Pad') + self._pad_x_view_controller = SpinBoxParameterViewController(diffraction_settings.pad_x) + self._pad_y_view_controller = SpinBoxParameterViewController(diffraction_settings.pad_y) + + layout = QGridLayout() + layout.addWidget(QLabel('Padding:'), 0, 0) + layout.addWidget(self._pad_x_view_controller.get_widget(), 0, 1) + layout.addWidget(self._pad_y_view_controller.get_widget(), 0, 2) + layout.setColumnStretch(1, 1) + layout.setColumnStretch(2, 1) + self.get_widget().setLayout(layout) + + +class ValueFilterViewController: + """FilterValuesStep — zero pattern values outside [lower_bound, upper_bound).""" + + def __init__(self, settings: DiffractionSettings) -> None: + self._lower_bound_enabled_view_controller = CheckBoxParameterViewController( + settings.value_lower_bound_enabled, 'Value Lower Bound:' + ) + self._lower_bound_view_controller = SpinBoxParameterViewController( + settings.value_lower_bound + ) + self._upper_bound_enabled_view_controller = CheckBoxParameterViewController( + settings.value_upper_bound_enabled, 'Value Upper Bound:' + ) + self._upper_bound_view_controller = SpinBoxParameterViewController( + settings.value_upper_bound + ) + + layout = QGridLayout() + layout.addWidget(self._lower_bound_enabled_view_controller.get_widget(), 0, 0) + layout.addWidget(self._lower_bound_view_controller.get_widget(), 0, 1) + layout.addWidget(self._upper_bound_enabled_view_controller.get_widget(), 1, 0) + layout.addWidget(self._upper_bound_view_controller.get_widget(), 1, 1) + layout.setColumnStretch(1, 1) + + self._widget = QGroupBox('Value Filter') + self._widget.setLayout(layout) + + def get_widget(self) -> QWidget: + return self._widget + + +class TransformViewController: + """HorizontalFlipStep + VerticalFlipStep + TransposeStep.""" + + def __init__(self, settings: DiffractionSettings) -> None: + self._hflip_view_controller = CheckBoxParameterViewController( + settings.hflip, 'Flip Horizontal' + ) + self._vflip_view_controller = CheckBoxParameterViewController( + settings.vflip, 'Flip Vertical' + ) + self._transpose_view_controller = CheckBoxParameterViewController( + settings.transpose, 'Transpose' + ) + + layout = QGridLayout() + layout.addWidget(self._hflip_view_controller.get_widget(), 0, 0) + layout.addWidget(self._vflip_view_controller.get_widget(), 0, 1) + layout.addWidget(self._transpose_view_controller.get_widget(), 0, 2) + layout.setColumnStretch(0, 1) + layout.setColumnStretch(1, 1) + layout.setColumnStretch(2, 1) + + self._widget = QGroupBox('Transform') + self._widget.setLayout(layout) + + def get_widget(self) -> QWidget: + return self._widget + + +class OpenDatasetWizardProcessingViewController(ParameterViewController): + """Processing wizard page. Groups are laid out top-to-bottom in the + DiffractionPrepPipeline execution order (see api/diffraction_prep.py): + filter → crop → binning → padding → transform (hflip → vflip → transpose). + Storage (memory map) is not part of the pipeline but is retained here as a + load-time concern; the horizontal separator between it and Value Filter + marks that boundary visually. + """ + + def __init__( + self, + diffraction_settings: DiffractionSettings, + extent_source: DetectorExtentSource, + file_dialog_factory: FileDialogFactory, + ) -> None: + self._storage_view_controller = StorageViewController( + diffraction_settings, file_dialog_factory + ) + self._value_filter_view_controller = ValueFilterViewController(diffraction_settings) + self._crop_view_controller = CropViewController(diffraction_settings, extent_source) + self._binning_view_controller = BinningViewController(diffraction_settings, extent_source) + self._padding_view_controller = PaddingViewController(diffraction_settings) + self._transform_view_controller = TransformViewController(diffraction_settings) + + separator = QFrame() + separator.setFrameShape(QFrame.Shape.HLine) + separator.setFrameShadow(QFrame.Shadow.Sunken) + + layout = QVBoxLayout() + layout.addWidget(self._storage_view_controller.get_widget()) + layout.addWidget(separator) + layout.addWidget(self._value_filter_view_controller.get_widget()) + layout.addWidget(self._crop_view_controller.get_widget()) + layout.addWidget(self._binning_view_controller.get_widget()) + layout.addWidget(self._padding_view_controller.get_widget()) + layout.addWidget(self._transform_view_controller.get_widget()) + layout.addStretch() + + self._page = OpenDatasetWizardPage() + self._page.setTitle('Processing') + self._page._set_complete(True) + self._page.setLayout(layout) + + def get_widget(self) -> QWizardPage: + return self._page diff --git a/src/ptychodus/controller/fluorescence/__init__.py b/src/ptychodus/controller/fluorescence/__init__.py new file mode 100644 index 000000000..ffce320a5 --- /dev/null +++ b/src/ptychodus/controller/fluorescence/__init__.py @@ -0,0 +1,5 @@ +from .core import FluorescenceController + +__all__ = [ + 'FluorescenceController', +] diff --git a/src/ptychodus/controller/fluorescence/core.py b/src/ptychodus/controller/fluorescence/core.py new file mode 100644 index 000000000..587e25832 --- /dev/null +++ b/src/ptychodus/controller/fluorescence/core.py @@ -0,0 +1,301 @@ +from __future__ import annotations +import logging + +from PyQt5.QtCore import QModelIndex +from PyQt5.QtWidgets import QInputDialog + +from ptychodus.api.observer import Observable, Observer + +from ...model.fluorescence import ( + FluorescenceAPI, + FluorescenceItemState, + FluorescenceRepository, + FluorescenceRepositoryItem, + FluorescenceRepositoryObserver, +) +from ...model.product import ProductRepository +from ...view.fluorescence import FluorescenceView +from ...view.widgets import ExceptionDialog +from ..data import FileDialogFactory +from ..image import ImageController +from .enhance_dialog import FluorescenceEnhanceDialogController +from .repository_tree_model import FluorescenceRepositoryTreeModel, DisplayMode + +logger = logging.getLogger(__name__) + + +class FluorescenceController(FluorescenceRepositoryObserver, Observer): + """Top-level controller for the promoted Fluorescence subview. + + Owns the left-pane dataset browser, the right-pane element viewer, and + the modal enhance dialog controller. + + The dataset tree expands to element leaves (see + :class:`FluorescenceRepositoryTreeModel`); selecting an item renders a + summary of all element maps, selecting an element renders that map. A + Measured/Enhanced radio pair below the tree selects which display drives + Counts + rendering. + """ + + def __init__( + self, + repository: FluorescenceRepository, + api: FluorescenceAPI, + product_repository: ProductRepository, + view: FluorescenceView, + image_controller: ImageController, + enhance_dialog_controller: FluorescenceEnhanceDialogController, + file_dialog_factory: FileDialogFactory, + ) -> None: + super().__init__() + self._repository = repository + self._api = api + self._product_repository = product_repository + self._view = view + self._image_controller = image_controller + self._enhance_dialog_controller = enhance_dialog_controller + self._file_dialog_factory = file_dialog_factory + self._task_monitor = api.get_task_monitor() + + self._tree_model = FluorescenceRepositoryTreeModel(repository) + view.tree_view.setModel(self._tree_model) + header = view.tree_view.header() + header.setSectionResizeMode(header.ResizeMode.ResizeToContents) + + selection_model = view.tree_view.selectionModel() + if selection_model is None: + raise ValueError('tree_view selection model is None!') + selection_model.currentChanged.connect(self._on_tree_selection_changed) + + view.measured_radio_button.toggled.connect(self._on_display_toggled) + view.enhanced_radio_button.toggled.connect(self._on_display_toggled) + + view.button_box.load_button.clicked.connect(self._load) + view.button_box.enhance_button.clicked.connect(self._enhance) + view.button_box.save_button.clicked.connect(self._save) + view.button_box.remove_button.clicked.connect(self._remove) + + repository.add_observer(self) + self._task_monitor.add_observer(self) + + self._sync_display_toggle_enabled() + self._sync_buttons() + + # ------------------------------------------------------------------ + # Selection & rendering + # ------------------------------------------------------------------ + + def _current_display(self) -> DisplayMode: + if self._view.enhanced_radio_button.isChecked(): + return DisplayMode.ENHANCED + return DisplayMode.MEASURED + + def _current_top_level_row(self) -> int: + return self._tree_model.item_row_for_index(self._view.tree_view.currentIndex()) + + def _current_item(self) -> FluorescenceRepositoryItem | None: + row = self._current_top_level_row() + if row < 0: + return None + try: + return self._repository[row] + except IndexError: + return None + + def _on_tree_selection_changed(self, current: QModelIndex, previous: QModelIndex) -> None: + self._sync_display_toggle_enabled() + self._render_current() + self._sync_buttons() + + def _on_display_toggled(self, checked: bool) -> None: + # Signals fire twice per toggle (button loses check + button gains + # check). Handle only the gain to avoid duplicate re-renders. + if not checked: + return + self._tree_model.set_display(self._current_display()) + self._render_current() + + def _render_current(self) -> None: + current = self._view.tree_view.currentIndex() + if not current.isValid(): + self._image_controller.clear_array() + return + node = current.internalPointer() + item = self._current_item() + if node is None or item is None: + self._image_controller.clear_array() + return + + display = self._current_display() + array = node.get_data(item, display) + if array is None: + self._image_controller.clear_array() + return + + pixel_geometry = item.get_product().get_geometry().get_object_plane_pixel_geometry() + self._image_controller.set_array(array, pixel_geometry) + + # ------------------------------------------------------------------ + # Display toggle enable-state + # ------------------------------------------------------------------ + + def _sync_display_toggle_enabled(self) -> None: + item = self._current_item() + has_enhanced = item is not None and item.get_enhanced() is not None + self._view.enhanced_radio_button.setEnabled(has_enhanced) + # If the current item can't be shown enhanced, fall back to measured + # (and let toggled -> re-render handle the redraw). + if not has_enhanced and self._view.enhanced_radio_button.isChecked(): + self._view.measured_radio_button.setChecked(True) + + # ------------------------------------------------------------------ + # Button actions + # ------------------------------------------------------------------ + + def _choose_product(self, title: str) -> tuple[bool, int]: + """Prompt the user to pick the target product for a new fluorescence dataset. + + Returns (accepted, product_index). product_index is -1 when the user + cancels or when the product repository is empty (message shown). + Mirrors the pattern in ProductController._choose_dataset. + """ + names = [ + self._product_repository[i].get_name() for i in range(len(self._product_repository)) + ] + + if not names: + ExceptionDialog.show_exception( + title, + ValueError('No products loaded — create or open a product first.'), + ) + return False, -1 + + label, accepted = QInputDialog.getItem( + self._view, + title, + 'Target Product:', + names, + 0, + False, + ) + + if not accepted: + return False, -1 + + return True, names.index(label) + + def _load(self) -> None: + title = 'Open Measured Fluorescence Dataset' + # Pick product first so cancelling doesn't waste a file selection. + accepted, product_index = self._choose_product(title) + if not accepted: + return + + file_path, name_filter = self._file_dialog_factory.get_open_file_path( + self._view, + title, + name_filters=list(self._api.get_open_file_filters()), + selected_name_filter=self._api.get_open_file_filter(), + ) + if not file_path: + return + try: + self._api.open_measured_dataset(file_path, product_index, file_type=name_filter) + except Exception as err: + logger.exception(err) + ExceptionDialog.show_exception(title, err) + + def _enhance(self) -> None: + row = self._current_top_level_row() + if row < 0: + return + self._enhance_dialog_controller.launch(row) + + def _save(self) -> None: + row = self._current_top_level_row() + if row < 0: + return + try: + item = self._repository[row] + except IndexError: + return + if item.get_enhanced() is None: + return + title = 'Save Enhanced Fluorescence Dataset' + file_path, name_filter = self._file_dialog_factory.get_save_file_path( + self._view, + title, + name_filters=list(self._api.get_save_file_filters()), + selected_name_filter=self._api.get_save_file_filter(), + ) + if not file_path: + return + try: + self._api.save_enhanced_dataset(row, file_path, file_type=name_filter) + except Exception as err: + logger.exception(err) + ExceptionDialog.show_exception(title, err) + + def _remove(self) -> None: + row = self._current_top_level_row() + if row < 0: + return + self._api.remove_item(row) + + def _sync_buttons(self) -> None: + item = self._current_item() + has_item = item is not None + has_enhanced = item is not None and item.get_enhanced() is not None + can_enhance = ( + item is not None + and item.get_state() is FluorescenceItemState.READY + and not self._task_monitor.is_processing + ) + + self._view.button_box.load_button.setEnabled(True) + self._view.button_box.enhance_button.setEnabled(can_enhance) + self._view.button_box.save_button.setEnabled(has_enhanced) + self._view.button_box.remove_button.setEnabled(has_item) + + # ------------------------------------------------------------------ + # Observers + # ------------------------------------------------------------------ + + def handle_item_inserted(self, index: int, item: FluorescenceRepositoryItem) -> None: + if not self._view.tree_view.currentIndex().isValid(): + self._view.tree_view.setCurrentIndex(self._tree_model.index(index, 0)) + self._sync_display_toggle_enabled() + self._sync_buttons() + + def handle_item_removed(self, index: int, item: FluorescenceRepositoryItem) -> None: + if not self._view.tree_view.currentIndex().isValid(): + row_count = self._tree_model.rowCount() + if row_count > 0: + target = min(index, row_count - 1) + self._view.tree_view.setCurrentIndex(self._tree_model.index(target, 0)) + else: + self._render_current() + self._sync_display_toggle_enabled() + self._sync_buttons() + + def handle_metadata_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + if index == self._current_top_level_row(): + self._sync_buttons() + + def handle_enhanced_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + if index == self._current_top_level_row(): + self._sync_display_toggle_enabled() + self._render_current() + self._sync_buttons() + + def handle_state_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + if index == self._current_top_level_row(): + if item.get_state() is FluorescenceItemState.FAILED: + logger.warning(f'Enhancement failed for item "{item.get_label()}"') + elif item.get_state() is FluorescenceItemState.ORPHANED: + logger.info(f'Fluorescence item "{item.get_label()}" orphaned by product removal') + self._sync_buttons() + + def _update(self, observable: Observable) -> None: + if observable is self._task_monitor: + self._sync_buttons() diff --git a/src/ptychodus/controller/fluorescence/enhance_dialog.py b/src/ptychodus/controller/fluorescence/enhance_dialog.py new file mode 100644 index 000000000..817c0f1bf --- /dev/null +++ b/src/ptychodus/controller/fluorescence/enhance_dialog.py @@ -0,0 +1,216 @@ +from __future__ import annotations +import logging + +from PyQt5.QtWidgets import QWidget + +from ptychodus.api.observer import Observable, Observer + +from ...model.fluorescence import ( + FluorescenceCore, + FluorescenceItemState, + FluorescenceSettings, + FluorescenceTaskMonitor, + PtychozoonFluorescenceEnhancer, + TwoStepFluorescenceEnhancer, + VSPIFluorescenceEnhancer, +) +from ...view.fluorescence import ( + FluorescenceEnhanceDialog, + FluorescenceStatusView, +) +from ...view.widgets import ExceptionDialog +from ..parametric import ComboBoxParameterViewController, ParameterViewBuilder + +logger = logging.getLogger(__name__) + + +def _build_two_step_widget(core: FluorescenceCore) -> QWidget: + builder = ParameterViewBuilder() + builder.add_combo_box( + core.upscaling_strategy_parameter, + core.upscaling_strategy_parameter.choices(), + 'Upscaling Strategy:', + ) + builder.add_combo_box( + core.deconvolution_strategy_parameter, + core.deconvolution_strategy_parameter.choices(), + 'Deconvolution Strategy:', + ) + return _build_page(builder) + + +def _build_vspi_widget(settings: FluorescenceSettings) -> QWidget: + builder = ParameterViewBuilder() + builder.add_decimal_line_edit(settings.vspi_damping_factor, 'Damping Factor:') + builder.add_spin_box(settings.vspi_max_iterations, 'Max Iterations:') + return _build_page(builder) + + +def _build_ptychozoon_widget(settings: FluorescenceSettings) -> QWidget: + builder = ParameterViewBuilder() + builder.add_decimal_line_edit(settings.ptychozoon_damping_factor, 'Damping Factor:') + builder.add_decimal_line_edit(settings.ptychozoon_gradient_smoothness, 'Gradient Smoothness:') + builder.add_spin_box(settings.ptychozoon_max_iterations, 'Max Iterations:') + builder.add_decimal_line_edit(settings.ptychozoon_atol, 'A Tolerance:') + builder.add_decimal_line_edit(settings.ptychozoon_btol, 'B Tolerance:') + builder.add_spin_box(settings.ptychozoon_checkpoint_interval, 'Checkpoint Interval:') + builder.add_check_box(settings.ptychozoon_use_gpu, 'Use GPU:') + builder.add_spin_box(settings.ptychozoon_gpu_device_index, 'CUDA Device Index:') + return _build_page(builder) + + +def _build_page(builder: ParameterViewBuilder) -> QWidget: + """Build a stacked-widget page, matching the zero margins the old views used.""" + widget = builder.build_widget() + layout = widget.layout() + + if layout is not None: + layout.setContentsMargins(0, 0, 0, 0) + + return widget + + +class FluorescenceStatusController(Observer): + def __init__( + self, + task_monitor: FluorescenceTaskMonitor, + view: FluorescenceStatusView, + ) -> None: + super().__init__() + self._monitor = task_monitor + self._view = view + + view.stop_button.clicked.connect(self._monitor.stop_processing) + + self._sync_model_to_view() + self._monitor.add_observer(self) + + def _sync_model_to_view(self) -> None: + log_handler = self._monitor.get_log_handler() + + for text in log_handler.messages(): + self._view.text_edit.appendPlainText(text) + + progress_goal = self._monitor.get_progress_goal() + progress_bar = self._view.progress_bar + + if self._monitor.is_processing and progress_goal > 0: + progress_bar.show() + progress_bar.setRange(0, progress_goal) + progress_bar.setValue(self._monitor.get_progress()) + self._view.stop_button.show() + else: + progress_bar.hide() + self._view.stop_button.hide() + + def _update(self, observable: Observable) -> None: + if observable is self._monitor: + self._sync_model_to_view() + + +class FluorescenceEnhanceDialogController(Observer): + """Owns the modal enhance dialog and its algorithm parameter sub-controllers. + + Launched from the top-level panel via ``launch(item_index)``; the target + product is looked up through the fluorescence item itself (which holds a + ProductRepositoryItem reference from load time). While the dialog is open, + the Run button submits an enhancement task through FluorescenceAPI. The + dialog stays open during enhancement so the user can watch the status log; + the panel handles saving the enhanced result. + """ + + def __init__( + self, + core: FluorescenceCore, + *, + has_ptychozoon: bool, + ) -> None: + super().__init__() + task_monitor = core.task_monitor + self._api = core.fluorescence_api + self._enhancer_parameter = core.enhancer_parameter + self._task_monitor = task_monitor + self._algorithm_view_controller = ComboBoxParameterViewController( + core.enhancer_parameter, core.enhancer_parameter.choices() + ) + self._dialog = FluorescenceEnhanceDialog(self._algorithm_view_controller.get_widget()) + self._item_index = -1 + + self._status_controller = FluorescenceStatusController( + task_monitor, self._dialog.status_view + ) + + # Keyed by display name rather than by index: the chooser sorts its plugins by + # display name, so page insertion order is not the combo-box order. + parameters_view = self._dialog.parameters_view + self._pages: dict[str, QWidget] = { + TwoStepFluorescenceEnhancer.DISPLAY_NAME: _build_two_step_widget(core), + VSPIFluorescenceEnhancer.DISPLAY_NAME: _build_vspi_widget(core.settings), + } + + # Only present when ptychozoon is installed. + if has_ptychozoon: + self._pages[PtychozoonFluorescenceEnhancer.DISPLAY_NAME] = _build_ptychozoon_widget( + core.settings + ) + + for page in self._pages.values(): + parameters_view.stacked_widget.addWidget(page) + + self._dialog.run_button.clicked.connect(self._run) + + core.enhancer_parameter.add_observer(self) + task_monitor.add_observer(self) + self._sync_algorithm_page() + self._sync_run_button_enabled() + + def launch(self, item_index: int) -> None: + self._item_index = item_index + try: + item = self._api.get_item(item_index) + except IndexError: + logger.warning(f'Missing fluorescence item {item_index}') + return + product_name = item.get_product().get_name() + self._dialog.setWindowTitle( + f'Enhance Fluorescence "{item.get_label()}" against "{product_name}"' + ) + self._sync_run_button_enabled() + self._dialog.open() + + def _run(self) -> None: + if self._item_index < 0: + ExceptionDialog.show_exception( + 'Enhance Fluorescence', + ValueError('Fluorescence item is not set'), + ) + return + try: + self._api.enhance(self._item_index) + except Exception as err: + logger.exception(err) + ExceptionDialog.show_exception('Enhance Fluorescence', err) + + def _sync_algorithm_page(self) -> None: + page = self._pages.get(self._enhancer_parameter.get_value()) + + if page is not None: + self._dialog.parameters_view.stacked_widget.setCurrentWidget(page) + + def _sync_run_button_enabled(self) -> None: + enabled = not self._task_monitor.is_processing and self._item_index >= 0 + if enabled: + try: + item = self._api.get_item(self._item_index) + except IndexError: + enabled = False + else: + if item.get_state() is not FluorescenceItemState.READY: + enabled = False + self._dialog.run_button.setEnabled(enabled) + + def _update(self, observable: Observable) -> None: + if observable is self._enhancer_parameter: + self._sync_algorithm_page() + elif observable is self._task_monitor: + self._sync_run_button_enabled() diff --git a/src/ptychodus/controller/fluorescence/repository_tree_model.py b/src/ptychodus/controller/fluorescence/repository_tree_model.py new file mode 100644 index 000000000..6765d3ae6 --- /dev/null +++ b/src/ptychodus/controller/fluorescence/repository_tree_model.py @@ -0,0 +1,413 @@ +from __future__ import annotations +from enum import Enum +from typing import Any, cast, overload + +from PyQt5.QtCore import Qt, QAbstractItemModel, QModelIndex, QObject +from PyQt5.QtGui import QBrush, QFont + +from ptychodus.api.common import RealArrayType +from ptychodus.api.fluorescence import ElementMap, FluorescenceDataset + +from ...model.fluorescence import ( + FluorescenceItemState, + FluorescenceRepository, + FluorescenceRepositoryItem, + FluorescenceRepositoryObserver, +) + + +class DisplayMode(Enum): + MEASURED = 'measured' + ENHANCED = 'enhanced' + + +_COL_NAME = 0 +_COL_COUNTS = 3 + + +def _select_display_quantity( + item: FluorescenceRepositoryItem, display_mode: DisplayMode +) -> FluorescenceDataset: + """Return the requested display quantity, falling back to measured when enhanced is absent. + + The tree walks measured element names as the canonical identity, so this + fallback keeps leaf rows renderable even when the enhancement hasn't been + run yet or when the enhanced dataset happens to be missing an element by + name. + """ + if display_mode is DisplayMode.ENHANCED: + enhanced = item.get_enhanced() + if enhanced is not None: + return enhanced + return item.get_measured() + + +def _lookup_display_element( + item: FluorescenceRepositoryItem, element_name: str, display: DisplayMode +) -> ElementMap | None: + dataset = _select_display_quantity(item, display) + for element_map in dataset.element_maps: + if element_map.name == element_name: + return element_map + return None + + +class _TreeNode: + """Base tree node — item root or element leaf.""" + + def __init__(self, parent: _TreeNode | None) -> None: + self.parent = parent + self.children: list[_TreeNode] = [] + + def row(self) -> int: + return 0 if self.parent is None else self.parent.children.index(self) + + def get_data( + self, item: FluorescenceRepositoryItem, display: DisplayMode + ) -> RealArrayType | None: + return None + + def get_counts(self, item: FluorescenceRepositoryItem, display: DisplayMode) -> float | None: + return None + + +class _ItemNode(_TreeNode): + def __init__(self, parent: _TreeNode) -> None: + super().__init__(parent) + + def get_data( + self, item: FluorescenceRepositoryItem, display: DisplayMode + ) -> RealArrayType | None: + if display is DisplayMode.ENHANCED: + enhanced = item.get_enhanced_summary() + if enhanced is not None: + return enhanced + return item.get_measured_summary() + return item.get_measured_summary() + + def get_counts(self, item: FluorescenceRepositoryItem, display: DisplayMode) -> float | None: + summary = self.get_data(item, display) + if summary is None: + return None + return float(summary.sum()) + + +class _ElementNode(_TreeNode): + def __init__(self, parent: _ItemNode, element_index: int) -> None: + super().__init__(parent) + self.element_index = element_index + + def _measured_element(self, item: FluorescenceRepositoryItem) -> ElementMap | None: + maps = item.get_measured().element_maps + if 0 <= self.element_index < len(maps): + return maps[self.element_index] + return None + + def element_name(self, item: FluorescenceRepositoryItem) -> str: + measured = self._measured_element(item) + return measured.name if measured is not None else '' + + def _resolve(self, item: FluorescenceRepositoryItem, display: DisplayMode) -> ElementMap | None: + measured = self._measured_element(item) + if measured is None: + return None + if display is DisplayMode.MEASURED: + return measured + # Enhanced by name — fall back to measured when the enhanced dataset + # is absent or lacks a matching name. + return _lookup_display_element(item, measured.name, DisplayMode.ENHANCED) or measured + + def get_data( + self, item: FluorescenceRepositoryItem, display: DisplayMode + ) -> RealArrayType | None: + element = self._resolve(item, display) + return None if element is None else element.counts_per_second + + def get_counts(self, item: FluorescenceRepositoryItem, display: DisplayMode) -> float | None: + element = self._resolve(item, display) + if element is None: + return None + return float(element.counts_per_second.sum()) + + +class FluorescenceRepositoryTreeModel(QAbstractItemModel): + """Two-level tree over FluorescenceRepository. + + - Level 1: `_ItemNode` per repository item (top-level rows). + - Level 2: `_ElementNode` per element in the item's measured dataset + (leaves). + + A single `display` field (measured/enhanced) drives the Counts column and + the array returned by ``get_data``. The controller flips it via + ``set_display`` when the user toggles the panel-level radio. + + Duck-typed against FluorescenceRepositoryObserver — inheriting the ABC + would clash with sip's wrappertype metaclass on QAbstractItemModel. + """ + + _HEADER = ('Name', 'Product', 'Elements', 'Counts') + + def __init__( + self, + repository: FluorescenceRepository, + parent: QObject | None = None, + ) -> None: + super().__init__(parent) + self._repository = repository + self._root = _TreeNode(None) + self._display = DisplayMode.MEASURED + + for item in repository: + self._root.children.append(self._build_item_node(item)) + + repository.add_observer(cast(FluorescenceRepositoryObserver, self)) + + # ------------------------------------------------------------------ + # Public helpers + # ------------------------------------------------------------------ + + def set_display(self, display: DisplayMode) -> None: + if display is self._display: + return + self._display = display + self._broadcast_counts_and_state() + + def get_display(self) -> DisplayMode: + return self._display + + def item_row_for_index(self, index: QModelIndex) -> int: + """Repository row index for whichever node ``index`` belongs to; -1 if none.""" + if not index.isValid(): + return -1 + node = index.internalPointer() + if isinstance(node, _ItemNode): + return node.row() + if isinstance(node, _ElementNode): + parent = node.parent + return parent.row() if parent is not None else -1 + return -1 + + # ------------------------------------------------------------------ + # Node construction + # ------------------------------------------------------------------ + + def _build_item_node(self, item: FluorescenceRepositoryItem) -> _ItemNode: + item_node = _ItemNode(self._root) + for element_index in range(len(item.get_measured().element_maps)): + item_node.children.append(_ElementNode(item_node, element_index)) + return item_node + + def _broadcast_counts_and_state(self) -> None: + """Refresh Counts across the whole tree after a display switch.""" + num_rows = self.rowCount() + if num_rows == 0: + return + # Top-level rows: Counts follows the display. + top_left = self.index(0, _COL_COUNTS) + bottom_right = self.index(num_rows - 1, _COL_COUNTS) + self.dataChanged.emit(top_left, bottom_right) + # Element leaves: Counts likewise. + for item_row in range(num_rows): + parent_index = self.index(item_row, 0) + num_children = self.rowCount(parent_index) + if num_children == 0: + continue + leaf_top = self.index(0, _COL_COUNTS, parent_index) + leaf_bottom = self.index(num_children - 1, _COL_COUNTS, parent_index) + self.dataChanged.emit(leaf_top, leaf_bottom) + + # ------------------------------------------------------------------ + # Repository observer callbacks + # ------------------------------------------------------------------ + + def handle_item_inserted(self, index: int, item: FluorescenceRepositoryItem) -> None: + self.beginInsertRows(QModelIndex(), index, index) + self._root.children.insert(index, self._build_item_node(item)) + self.endInsertRows() + + def handle_item_removed(self, index: int, item: FluorescenceRepositoryItem) -> None: + self.beginRemoveRows(QModelIndex(), index, index) + del self._root.children[index] + self.endRemoveRows() + + def handle_metadata_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + self._emit_row_changed(index) + + def handle_enhanced_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + # Structural change: none (the enhanced dataset is a display, not a + # separate branch). Refresh Counts for the item row and its leaves, + # so the current display's numbers stay in sync. + self._emit_row_changed(index) + parent_idx = self.index(index, 0) + num_children = self.rowCount(parent_idx) + if num_children > 0: + top_left = self.index(0, _COL_COUNTS, parent_idx) + bottom_right = self.index(num_children - 1, _COL_COUNTS, parent_idx) + self.dataChanged.emit(top_left, bottom_right) + + def handle_state_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + self._emit_row_changed( + index, + [ + Qt.ItemDataRole.DisplayRole, + Qt.ItemDataRole.FontRole, + Qt.ItemDataRole.ForegroundRole, + ], + ) + + def _emit_row_changed(self, index: int, roles: list[int] | None = None) -> None: + top_left = self.index(index, 0) + bottom_right = self.index(index, self.columnCount() - 1) + if roles is None: + self.dataChanged.emit(top_left, bottom_right) + else: + self.dataChanged.emit(top_left, bottom_right, roles) + + # ------------------------------------------------------------------ + # QAbstractItemModel plumbing + # ------------------------------------------------------------------ + + def headerData( # noqa: N802 + self, + section: int, + orientation: Qt.Orientation, + role: int = Qt.ItemDataRole.DisplayRole, + ) -> Any: + if orientation == Qt.Orientation.Horizontal and role == Qt.ItemDataRole.DisplayRole: + return self._HEADER[section] + + @overload + def parent(self, child: QModelIndex) -> QModelIndex: ... + + @overload + def parent(self) -> QObject: ... + + def parent(self, child: QModelIndex | None = None) -> QModelIndex | QObject: + if child is None: + return super().parent() + if not child.isValid(): + return QModelIndex() + + node = child.internalPointer() + if node is None or node.parent is None or node.parent is self._root: + return QModelIndex() + return self.createIndex(node.parent.row(), 0, node.parent) + + def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex: + if not self.hasIndex(row, column, parent): + return QModelIndex() + parent_node = parent.internalPointer() if parent.isValid() else self._root + try: + node = parent_node.children[row] + except IndexError: + return QModelIndex() + return self.createIndex(row, column, node) + + def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + if parent.column() > 0: + return 0 + parent_node = parent.internalPointer() if parent.isValid() else self._root + return len(parent_node.children) + + def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + return len(self._HEADER) + + def flags(self, index: QModelIndex) -> Qt.ItemFlags: + value = super().flags(index) + if not index.isValid(): + return value + + node = index.internalPointer() + if not isinstance(node, _ItemNode): + return value + if index.column() != _COL_NAME: + return value + + try: + item = self._repository[index.row()] + except IndexError: + return value + # ORPHANED items can still be renamed — label is metadata, not tied to + # product lifetime; ENHANCING items lock the label so the user doesn't + # rename a row mid-task. + if item.get_state() is not FluorescenceItemState.ENHANCING: + value |= Qt.ItemFlag.ItemIsEditable + return value + + def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: + if not index.isValid(): + return None + node = index.internalPointer() + if node is None: + return None + + item_row = self.item_row_for_index(index) + if item_row < 0: + return None + try: + item = self._repository[item_row] + except IndexError: + return None + + state = item.get_state() + + if isinstance(node, _ElementNode): + if role == Qt.ItemDataRole.DisplayRole: + match index.column(): + case 0: + return node.element_name(item) + case 3: + counts = node.get_counts(item, self._display) + return _format_counts(counts) + case _: + return None + elif role == Qt.ItemDataRole.FontRole: + font = QFont() + font.setItalic(True) + return font + return None + + # Item row. + if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole: + match index.column(): + case 0: + return item.get_label() + case 1: + return item.get_product().get_name() + case 2: + return len(item.get_measured().element_maps) + case 3: + counts = node.get_counts(item, self._display) + return _format_counts(counts) + elif role == Qt.ItemDataRole.FontRole: + if state is FluorescenceItemState.ENHANCING: + font = QFont() + font.setItalic(True) + return font + if state is FluorescenceItemState.FAILED or state is FluorescenceItemState.ORPHANED: + font = QFont() + font.setStrikeOut(True) + return font + elif role == Qt.ItemDataRole.ForegroundRole: + if state is not FluorescenceItemState.READY: + return QBrush(Qt.GlobalColor.gray) + return None + + def setData(self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.EditRole) -> bool: # noqa: N802 + if not index.isValid() or role != Qt.ItemDataRole.EditRole: + return False + node = index.internalPointer() + if not isinstance(node, _ItemNode) or index.column() != _COL_NAME: + return False + try: + item = self._repository[index.row()] + except IndexError: + return False + item.set_label(str(value)) + return True + + +def _format_counts(counts: float | None) -> str: + if counts is None: + return '—' + return f'{counts:.4g}' diff --git a/src/ptychodus/controller/image.py b/src/ptychodus/controller/image.py index e7f7c0798..29b69a6d5 100644 --- a/src/ptychodus/controller/image.py +++ b/src/ptychodus/controller/image.py @@ -128,11 +128,8 @@ def __init__( self._sync_model_to_view() engine.add_observer(self) - view.min_display_value_slider.value_changed.connect( - lambda value: engine.set_min_display_value(float(value)) - ) - view.max_display_value_slider.value_changed.connect( - lambda value: engine.set_max_display_value(float(value)) + view.display_range_slider.selection_changed.connect( + lambda sel: engine.set_display_value_range(float(sel.lower), float(sel.upper)) ) view.auto_button.clicked.connect(self._auto_display_range) view.edit_button.clicked.connect(self._display_range_dialog.open) @@ -173,13 +170,13 @@ def _sync_model_to_view(self) -> None: self._display_range_dialog.min_value_line_edit.set_value(min_value) self._display_range_dialog.max_value_line_edit.set_value(max_value) + selection = Interval[Decimal](min_value, max_value) if self._display_range_is_locked: - self._view.min_display_value_slider.set_value(min_value) - self._view.max_display_value_slider.set_value(max_value) + self._view.display_range_slider.set_selection(selection) else: - display_range_limits = Interval[Decimal](min_value, max_value) - self._view.min_display_value_slider.set_value_and_range(min_value, display_range_limits) - self._view.max_display_value_slider.set_value_and_range(max_value, display_range_limits) + self._view.display_range_slider.set_selection_and_bounds( + selection, selection, block_signal=True + ) self._sync_color_legend_to_view() diff --git a/src/ptychodus/controller/object/core.py b/src/ptychodus/controller/object/core.py index 2c4179666..43436d60c 100644 --- a/src/ptychodus/controller/object/core.py +++ b/src/ptychodus/controller/object/core.py @@ -18,7 +18,7 @@ from .editor_factory import ObjectEditorViewControllerFactory from .fourier import FourierAnalysisViewController from .frc import FourierRingCorrelationViewController -from .tree_model import ObjectTreeModel +from .tree_model import ObjectTreeModel, try_get_object from .xmcd import XMCDViewController logger = logging.getLogger(__name__) @@ -245,7 +245,11 @@ def _update_view(self, current: QModelIndex, previous: QModelIndex) -> None: except IndexError: logger.warning('Unable to access item for visualization!') else: - object_ = item.get_object() + object_ = try_get_object(item) + if object_ is None: + # Null sentinel until the dataset binds; observer chain will re-fire. + self._image_controller.clear_array() + return array = ( object_.get_layer(current.row()) if current.parent().isValid() diff --git a/src/ptychodus/controller/object/editor_factory.py b/src/ptychodus/controller/object/editor_factory.py index 241e2c2dc..082e22370 100644 --- a/src/ptychodus/controller/object/editor_factory.py +++ b/src/ptychodus/controller/object/editor_factory.py @@ -5,6 +5,7 @@ from ...model.product.object import ( DeadLeavesObjectBuilder, FractalNoiseObjectBuilder, + FromFileObjectBuilder, GaussianRandomFieldObjectBuilder, ObjectRepositoryItem, PaganinObjectBuilder, @@ -20,6 +21,10 @@ def __init__(self, item: ObjectRepositoryItem) -> None: self._item = item self._parameter = item.layer_spacing_m self._widget = QSpinBox() + self._widget.setToolTip( + 'Layers are only ever added, never removed. An object that already has' + ' more layers than this keeps the ones it has.' + ) self._sync_model_to_view() self._widget.valueChanged.connect(self._sync_view_to_model) @@ -219,6 +224,18 @@ def create_editor_dialog( group=additional_layers_group, ) return dialog_builder.build_dialog(title, parent) + elif isinstance(object_builder, FromFileObjectBuilder): + # A from-file object is conditioned on the way in, but only the layer + # spacing applies -- the extra padding is generation-only, so its spin + # boxes are deliberately absent here. A from-memory object falls + # through to the message box below, since nothing about it is editable. + dialog_builder = ParameterViewBuilder() + dialog_builder.add_view_controller( + MultisliceViewController(item), + 'Number of Layers:', + group=additional_layers_group, + ) + return dialog_builder.build_dialog(title, parent) return QMessageBox( QMessageBox.Icon.Information, diff --git a/src/ptychodus/controller/object/tree_model.py b/src/ptychodus/controller/object/tree_model.py index 436bdbe5c..52c4fce94 100644 --- a/src/ptychodus/controller/object/tree_model.py +++ b/src/ptychodus/controller/object/tree_model.py @@ -5,19 +5,39 @@ from PyQt5.QtGui import QBrush from ptychodus.api.common import BYTES_PER_MEGABYTE +from ptychodus.api.object import Object from ...model.product import ObjectAPI, ObjectRepository from ...model.product.object import ObjectRepositoryItem +def try_get_object(item: ObjectRepositoryItem) -> Object | None: + # Returns None when the object is the null sentinel (no dataset bound yet + # — see ObjectRepositoryItem.__init__). Callers skip work rather than + # letting the ValueError from Object.get_pixel_geometry() escape. + object_ = item.get_object() + try: + object_.get_pixel_geometry() + except ValueError: + return None + return object_ + + class ObjectTreeNode: def __init__(self, parent: ObjectTreeNode | None = None) -> None: self.parent = parent self.children: list[ObjectTreeNode] = list() - def insert_node(self, index: int = -1) -> ObjectTreeNode: + def insert_node(self, index: int | None = None) -> ObjectTreeNode: node = ObjectTreeNode(self) - self.children.insert(index, node) + + if index is None: + # list.insert(-1, ...) inserts *before* the last child rather than + # appending, which scrambles the order the callers below rely on. + self.children.append(node) + else: + self.children.insert(index, node) + return node def remove_node(self, index: int = -1) -> ObjectTreeNode: @@ -68,8 +88,13 @@ def insert_item(self, index: int, item: ObjectRepositoryItem) -> None: self.endInsertRows() def update_item(self, index: int, item: ObjectRepositoryItem) -> None: + # Qt row/column bounds are inclusive, so the last valid column is + # len(self._header) - 1. Asking index() for one past the end fails + # hasIndex and yields an invalid QModelIndex, which makes the + # dataChanged range malformed and the signal a no-op. + last_column = len(self._header) - 1 top_left = self.index(index, 0) - bottom_right = self.index(index, len(self._header)) + bottom_right = self.index(index, last_column) self.dataChanged.emit(top_left, bottom_right) node = self._tree_root.children[index] @@ -77,23 +102,24 @@ def update_item(self, index: int, item: ObjectRepositoryItem) -> None: num_layers_new = item.get_object().num_layers if num_layers_old < num_layers_new: - self.beginInsertRows(top_left, num_layers_old, num_layers_new) + self.beginInsertRows(top_left, num_layers_old, num_layers_new - 1) while len(node.children) < num_layers_new: node.insert_node() self.endInsertRows() elif num_layers_old > num_layers_new: - self.beginRemoveRows(top_left, num_layers_new, num_layers_old) + self.beginRemoveRows(top_left, num_layers_new, num_layers_old - 1) while len(node.children) > num_layers_new: node.remove_node() self.endRemoveRows() - child_top_left = self.index(0, 0, top_left) - child_bottom_right = self.index(num_layers_new, len(self._header), top_left) - self.dataChanged.emit(child_top_left, child_bottom_right) + if num_layers_new > 0: + child_top_left = self.index(0, 0, top_left) + child_bottom_right = self.index(num_layers_new - 1, last_column, top_left) + self.dataChanged.emit(child_top_left, child_bottom_right) def remove_item(self, index: int, item: ObjectRepositoryItem) -> None: self.beginRemoveRows(QModelIndex(), index, index) @@ -158,35 +184,42 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A try: return item.layer_spacing_m[index.row()] except IndexError: - return float('inf') + # N layers have N-1 spacings, so the last layer has + # no distance to the next one. flags() already makes + # that cell read-only; show it blank rather than inf. + return None elif role == Qt.ItemDataRole.BackgroundRole: if index.flags() & Qt.ItemFlag.ItemIsEditable: return self._editable_item_brush else: item = self._repository[index.row()] - object_ = item.get_object() - pixel_geometry = object_.get_pixel_geometry() + raw_object = item.get_object() + object_ = try_get_object(item) + # None when the object is not yet built (see ObjectRepositoryItem + # null sentinel): pixel-geometry- and array-dependent columns return + # None; name/thickness/builder/size still show. + pixel_geometry = object_.get_pixel_geometry() if object_ is not None else None if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole: match index.column(): case 0: return self._repository.get_name(index.row()) case 1: - return object_.get_total_thickness_m() + return raw_object.get_total_thickness_m() case 2: return item.get_builder().get_name() case 3: - return str(object_.dtype) + return str(object_.dtype) if object_ is not None else None case 4: - return object_.width_px + return object_.width_px if object_ is not None else None case 5: - return object_.height_px + return object_.height_px if object_ is not None else None case 6: - return f'{pixel_geometry.width_m * 1e9:.4g}' + return f'{pixel_geometry.width_m * 1e9:.4g}' if pixel_geometry else None case 7: - return f'{pixel_geometry.height_m * 1e9:.4g}' + return f'{pixel_geometry.height_m * 1e9:.4g}' if pixel_geometry else None case 8: - return f'{object_.nbytes / BYTES_PER_MEGABYTE:.2f}' + return f'{raw_object.nbytes / BYTES_PER_MEGABYTE:.2f}' elif role == Qt.ItemDataRole.BackgroundRole: if index.flags() & Qt.ItemFlag.ItemIsEditable: return self._editable_item_brush @@ -223,7 +256,7 @@ def setData(self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.Ed return False item.layer_spacing_m[index.row()] = distance_m - return False + return True else: if index.column() == 0: self._repository.set_name(index.row(), str(value)) diff --git a/src/ptychodus/controller/parametric.py b/src/ptychodus/controller/parametric.py index c4707065f..fdcd1cb7e 100644 --- a/src/ptychodus/controller/parametric.py +++ b/src/ptychodus/controller/parametric.py @@ -653,7 +653,7 @@ def add_check_box( def add_combo_box( self, - parameter: StringParameter, + parameter: Parameter[str], items: Iterable[str], label: str, *, diff --git a/src/ptychodus/controller/probe/core.py b/src/ptychodus/controller/probe/core.py index 12c1db02c..e2b6ad720 100644 --- a/src/ptychodus/controller/probe/core.py +++ b/src/ptychodus/controller/probe/core.py @@ -6,17 +6,7 @@ from ptychodus.api.observer import SequenceObserver -from ptychodus.api.fluorescence import FluorescenceEnhancer -from ptychodus.api.plugins import PluginChooser - from ...model.analysis import IlluminationMapper, ProbePropagatorSettings, ProbePropagator -from ...model.fluorescence import ( - FluorescenceAPI, - FluorescenceTaskMonitor, - PtychozoonFluorescenceEnhancer, - TwoStepFluorescenceEnhancer, - VSPIFluorescenceEnhancer, -) from ...model.product import ProbeAPI, ProbeRepository from ...model.product.probe import ProbeRepositoryItem from ...model.visualization import VisualizationEngine @@ -30,10 +20,9 @@ from ..helpers import connect_triggered_signal, create_brush_for_editable_cell from ..image import ImageController from .editor_factory import ProbeEditorViewControllerFactory -from .fluorescence import FluorescenceViewController from .illumination import IlluminationViewController from .propagator import ProbePropagationViewController -from .tree_model import ProbeTreeModel +from .tree_model import ProbeTreeModel, try_get_probe logger = logging.getLogger(__name__) @@ -49,13 +38,6 @@ def __init__( propagator_visualization_engine: VisualizationEngine, illumination_mapper: IlluminationMapper, illumination_visualization_engine: VisualizationEngine, - fluorescence_api: FluorescenceAPI, - fluorescence_enhancer_chooser: PluginChooser[FluorescenceEnhancer], - fluorescence_two_step_enhancer: TwoStepFluorescenceEnhancer, - fluorescence_vspi_enhancer: VSPIFluorescenceEnhancer, - fluorescence_ptychozoon_enhancer: PtychozoonFluorescenceEnhancer | None, - fluorescence_task_monitor: FluorescenceTaskMonitor, - fluorescence_visualization_engine: VisualizationEngine, view: RepositoryTreeView, file_dialog_factory: FileDialogFactory, ) -> None: @@ -76,16 +58,6 @@ def __init__( illumination_visualization_engine, file_dialog_factory, ) - self._fluorescence_view_controller = FluorescenceViewController( - fluorescence_api, - fluorescence_enhancer_chooser, - fluorescence_two_step_enhancer, - fluorescence_vspi_enhancer, - fluorescence_ptychozoon_enhancer, - fluorescence_task_monitor, - fluorescence_visualization_engine, - file_dialog_factory, - ) # TODO figure out good fix when saving NPY file without suffix (numpy adds suffix) repository.add_observer(self) @@ -136,9 +108,6 @@ def __init__( illumination_action = view.button_box.analyze_menu.addAction('Map Illumination...') connect_triggered_signal(illumination_action, self._map_illumination) - fluorescence_action = view.button_box.analyze_menu.addAction('Enhance Fluorescence...') - connect_triggered_signal(fluorescence_action, self._enhance_fluorescence) - def _get_current_item_index(self) -> int: model_index = self._view.tree_view.currentIndex() @@ -243,14 +212,6 @@ def _map_illumination(self) -> None: else: self._illumination_view_controller.map(item_index) - def _enhance_fluorescence(self) -> None: - item_index = self._get_current_item_index() - - if item_index < 0: - logger.warning('No current item!') - else: - self._fluorescence_view_controller.launch(item_index) - def _update_view(self, current: QModelIndex, previous: QModelIndex) -> None: enabled = current.isValid() self._view.button_box.load_button.setEnabled(enabled) @@ -268,7 +229,11 @@ def _update_view(self, current: QModelIndex, previous: QModelIndex) -> None: except IndexError: logger.warning('Unable to access item for visualization!') else: - probe = item.get_probes().get_probe_no_opr() # TODO OPR + probe = try_get_probe(item) # TODO OPR + if probe is None: + # Null sentinel until the dataset binds; observer chain will re-fire. + self._image_controller.clear_array() + return array = ( probe.get_incoherent_mode(current.row()) if current.parent().isValid() diff --git a/src/ptychodus/controller/probe/editor_factory.py b/src/ptychodus/controller/probe/editor_factory.py index 23ed77121..6521398ce 100644 --- a/src/ptychodus/controller/probe/editor_factory.py +++ b/src/ptychodus/controller/probe/editor_factory.py @@ -19,6 +19,7 @@ AveragePatternProbeBuilder, DiskProbeBuilder, FresnelZonePlateProbeBuilder, + FromFileProbeBuilder, HermiteProbeBuilder, ProbeModeDecayType, ProbeRepositoryItem, @@ -356,10 +357,16 @@ def _append_additional_modes( probe_builder: ProbeSequenceBuilder, dialog_builder: ParameterViewBuilder, ) -> None: + expand_only_tool_tip = ( + 'Modes are only ever added, never removed. A probe that already has' + ' more modes than this keeps the ones it has.' + ) + incoherent_modes_group = 'Incoherent (Mixed State) Modes' dialog_builder.add_spin_box( probe_builder.num_incoherent_modes, 'Number of Modes:', + tool_tip=expand_only_tool_tip, group=incoherent_modes_group, ) dialog_builder.add_check_box( @@ -382,6 +389,7 @@ def _append_additional_modes( dialog_builder.add_spin_box( probe_builder.num_coherent_modes, 'Number of Modes:', + tool_tip=expand_only_tool_tip, group=coherent_modes_group, ) @@ -395,7 +403,14 @@ def create_editor_dialog( dialog_builder = ParameterViewBuilder() dialog_builder.add_view_controller_to_top(ProbeMetricsViewController('Probe Metrics', item)) - if self._append_primary_mode(probe_builder, dialog_builder): + # A from-file probe is conditioned on the way in, so the mode parameters + # do something and belong in the dialog. A from-memory probe holds an + # already-conditioned probe -- reconstruction output, or a product loaded + # from file -- whose mode parameters are deliberately inert, so offering + # controls that would silently do nothing is worse than offering none. + has_primary_mode = self._append_primary_mode(probe_builder, dialog_builder) + + if has_primary_mode or isinstance(probe_builder, FromFileProbeBuilder): self._append_additional_modes(probe_builder, dialog_builder) else: dialog_builder.add_view_controller_to_bottom( diff --git a/src/ptychodus/controller/probe/fluorescence.py b/src/ptychodus/controller/probe/fluorescence.py deleted file mode 100644 index 0988fd112..000000000 --- a/src/ptychodus/controller/probe/fluorescence.py +++ /dev/null @@ -1,476 +0,0 @@ -from decimal import Decimal -from typing import Any, Final -import logging - -from PyQt5.QtCore import Qt, QAbstractListModel, QModelIndex, QObject, QStringListModel -from PyQt5.QtWidgets import QWidget - -from ptychodus.api.fluorescence import ElementMap, FluorescenceDataset -from ptychodus.api.observer import Observable, Observer -from ptychodus.api.plugins import PluginChooser - -from ...model.fluorescence import ( - FluorescenceAPI, - FluorescenceTaskMonitor, - PtychozoonFluorescenceEnhancer, - TwoStepFluorescenceEnhancer, - VSPIFluorescenceEnhancer, -) -from ...model.visualization import VisualizationEngine -from ...view.probe import ( - FluorescenceDialog, - FluorescencePtychozoonParametersView, - FluorescenceStatusView, - FluorescenceTwoStepParametersView, - FluorescenceVSPIParametersView, -) -from ...view.widgets import ExceptionDialog -from ..data import FileDialogFactory -from ..helpers import connect_current_changed_signal -from ..visualization import ( - VisualizationParametersController, - VisualizationWidgetController, -) - -logger = logging.getLogger(__name__) - - -class FluorescenceChannelListModel(QAbstractListModel): - def __init__( - self, controller: 'FluorescenceViewController', parent: QObject | None = None - ) -> None: - super().__init__(parent) - self._controller = controller - - def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: - # TODO make this a table model and show measured/enhanced count statistics - if index.isValid() and role == Qt.ItemDataRole.DisplayRole: - emap = self._controller.get_measured_element_map(index.row()) - - if emap is not None: - return emap.name - - def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 - return self._controller.get_num_channels() - - -class FluorescenceTwoStepViewController(Observer): - def __init__(self, enhancer: TwoStepFluorescenceEnhancer) -> None: - super().__init__() - self._enhancer = enhancer - self._view = FluorescenceTwoStepParametersView() - - self._upscaling_model = QStringListModel() - self._upscaling_model.setStringList(self._enhancer.get_upscaling_strategies()) - self._view.upscaling_strategy_combo_box.setModel(self._upscaling_model) - self._view.upscaling_strategy_combo_box.textActivated.connect( - enhancer.set_upscaling_strategy - ) - - self._deconvolution_model = QStringListModel() - self._deconvolution_model.setStringList(self._enhancer.get_deconvolution_strategies()) - self._view.deconvolution_strategy_combo_box.setModel(self._deconvolution_model) - self._view.deconvolution_strategy_combo_box.textActivated.connect( - enhancer.set_deconvolution_strategy - ) - - self._sync_model_to_view() - enhancer.add_observer(self) - - def get_widget(self) -> QWidget: - return self._view - - def _sync_model_to_view(self) -> None: - self._view.upscaling_strategy_combo_box.setCurrentText( - self._enhancer.get_upscaling_strategy() - ) - self._view.deconvolution_strategy_combo_box.setCurrentText( - self._enhancer.get_deconvolution_strategy() - ) - - def _update(self, observable: Observable) -> None: - if observable is self._enhancer: - self._sync_model_to_view() - - -class FluorescenceVSPIViewController(Observer): - MAX_INT: Final[int] = 0x7FFFFFFF - - def __init__(self, enhancer: VSPIFluorescenceEnhancer) -> None: - super().__init__() - self._enhancer = enhancer - self._view = FluorescenceVSPIParametersView() - - self._view.damping_factor_line_edit.value_changed.connect( - self._sync_damping_factor_to_model - ) - self._view.max_iterations_spin_box.setRange(1, self.MAX_INT) - self._view.max_iterations_spin_box.valueChanged.connect(enhancer.set_max_iterations) - - enhancer.add_observer(self) - self._sync_model_to_view() - - def get_widget(self) -> QWidget: - return self._view - - def _sync_damping_factor_to_model(self, value: Decimal) -> None: - self._enhancer.set_damping_factor(float(value)) - - def _sync_model_to_view(self) -> None: - self._view.damping_factor_line_edit.set_value( - Decimal(repr(self._enhancer.get_damping_factor())) - ) - self._view.max_iterations_spin_box.setValue(self._enhancer.get_max_iterations()) - - def _update(self, observable: Observable) -> None: - if observable is self._enhancer: - self._sync_model_to_view() - - -class FluorescencePtychozoonViewController(Observer): - MAX_INT: Final[int] = 0x7FFFFFFF - - def __init__(self, enhancer: PtychozoonFluorescenceEnhancer) -> None: - super().__init__() - self._enhancer = enhancer - self._view = FluorescencePtychozoonParametersView() - - self._view.damping_factor_line_edit.value_changed.connect( - self._sync_damping_factor_to_model - ) - self._view.gradient_smoothness_line_edit.value_changed.connect( - self._sync_gradient_smoothness_to_model - ) - self._view.max_iterations_spin_box.setRange(1, self.MAX_INT) - self._view.max_iterations_spin_box.valueChanged.connect(enhancer.set_max_iterations) - self._view.atol_line_edit.value_changed.connect(self._sync_atol_to_model) - self._view.btol_line_edit.value_changed.connect(self._sync_btol_to_model) - self._view.checkpoint_interval_spin_box.setRange(1, self.MAX_INT) - self._view.checkpoint_interval_spin_box.valueChanged.connect( - enhancer.set_checkpoint_interval - ) - self._view.use_gpu_check_box.toggled.connect(enhancer.set_gpu_enabled) - self._view.gpu_device_index_spin_box.setRange(0, self.MAX_INT) - self._view.gpu_device_index_spin_box.valueChanged.connect(enhancer.set_gpu_device_index) - - enhancer.add_observer(self) - self._sync_model_to_view() - - def get_widget(self) -> QWidget: - return self._view - - def _sync_damping_factor_to_model(self, value: Decimal) -> None: - self._enhancer.set_damping_factor(float(value)) - - def _sync_gradient_smoothness_to_model(self, value: Decimal) -> None: - self._enhancer.set_gradient_smoothness(float(value)) - - def _sync_atol_to_model(self, value: Decimal) -> None: - self._enhancer.set_atol(float(value)) - - def _sync_btol_to_model(self, value: Decimal) -> None: - self._enhancer.set_btol(float(value)) - - def _sync_model_to_view(self) -> None: - self._view.damping_factor_line_edit.set_value( - Decimal(repr(self._enhancer.get_damping_factor())) - ) - self._view.gradient_smoothness_line_edit.set_value( - Decimal(repr(self._enhancer.get_gradient_smoothness())) - ) - self._view.max_iterations_spin_box.setValue(self._enhancer.get_max_iterations()) - self._view.atol_line_edit.set_value(Decimal(repr(self._enhancer.get_atol()))) - self._view.btol_line_edit.set_value(Decimal(repr(self._enhancer.get_btol()))) - self._view.checkpoint_interval_spin_box.setValue(self._enhancer.get_checkpoint_interval()) - self._view.use_gpu_check_box.setChecked(self._enhancer.is_gpu_enabled()) - self._view.gpu_device_index_spin_box.setValue(self._enhancer.get_gpu_device_index()) - - def _update(self, observable: Observable) -> None: - if observable is self._enhancer: - self._sync_model_to_view() - - -class FluorescenceStatusController(Observer): - def __init__( - self, - task_monitor: FluorescenceTaskMonitor, - view: FluorescenceStatusView, - ) -> None: - super().__init__() - self._monitor = task_monitor - self._view = view - - view.stop_button.clicked.connect(self._monitor.stop_processing) - - self._sync_model_to_view() - self._monitor.add_observer(self) - - def _sync_model_to_view(self) -> None: - log_handler = self._monitor.get_log_handler() - - for text in log_handler.messages(): - self._view.text_edit.appendPlainText(text) - - progress_goal = self._monitor.get_progress_goal() - progress_bar = self._view.progress_bar - - if self._monitor.is_processing and progress_goal > 0: - progress_bar.show() - progress_bar.setRange(0, progress_goal) - progress_bar.setValue(self._monitor.get_progress()) - self._view.stop_button.show() - else: - progress_bar.hide() - self._view.stop_button.hide() - - def _update(self, observable: Observable) -> None: - if observable is self._monitor: - self._sync_model_to_view() - - -class FluorescenceViewController(Observer): - def __init__( - self, - fluorescence_api: FluorescenceAPI, - enhancer_chooser: PluginChooser, - two_step_enhancer: TwoStepFluorescenceEnhancer, - vspi_enhancer: VSPIFluorescenceEnhancer, - ptychozoon_enhancer: PtychozoonFluorescenceEnhancer | None, - task_monitor: FluorescenceTaskMonitor, - engine: VisualizationEngine, - file_dialog_factory: FileDialogFactory, - ) -> None: - super().__init__() - self._fluorescence_api = fluorescence_api - self._enhancer_chooser = enhancer_chooser - self._dataset_emitter = task_monitor.get_dataset_emitter() - self._engine = engine - self._file_dialog_factory = file_dialog_factory - self._dialog = FluorescenceDialog() - self._product_index = -1 - self._measured: FluorescenceDataset | None = None - self._enhanced: FluorescenceDataset | None = None - self._status_controller = FluorescenceStatusController( - task_monitor, - self._dialog.fluorescence_status_view, - ) - self._enhancement_model = QStringListModel() - self._enhancement_model.setStringList([plugin.display_name for plugin in enhancer_chooser]) - self._channel_list_model = FluorescenceChannelListModel(self) - - self._dialog.fluorescence_parameters_view.open_button.clicked.connect( - self._open_measured_dataset - ) - - two_step_view_controller = FluorescenceTwoStepViewController(two_step_enhancer) - self._dialog.fluorescence_parameters_view.algorithm_combo_box.addItem( - TwoStepFluorescenceEnhancer.DISPLAY_NAME, - self._dialog.fluorescence_parameters_view.algorithm_combo_box.count(), - ) - self._dialog.fluorescence_parameters_view.stacked_widget.addWidget( - two_step_view_controller.get_widget() - ) - - vspi_view_controller = FluorescenceVSPIViewController(vspi_enhancer) - self._dialog.fluorescence_parameters_view.algorithm_combo_box.addItem( - VSPIFluorescenceEnhancer.DISPLAY_NAME, - self._dialog.fluorescence_parameters_view.algorithm_combo_box.count(), - ) - self._dialog.fluorescence_parameters_view.stacked_widget.addWidget( - vspi_view_controller.get_widget() - ) - - # Registered last, matching the enhancer_chooser order, so the combo-box - # index selects the correct stacked page. Only present when ptychozoon is - # installed (enhancer is None otherwise). - self._ptychozoon_view_controller: FluorescencePtychozoonViewController | None = None - - if ptychozoon_enhancer is not None: - self._ptychozoon_view_controller = FluorescencePtychozoonViewController( - ptychozoon_enhancer - ) - self._dialog.fluorescence_parameters_view.algorithm_combo_box.addItem( - PtychozoonFluorescenceEnhancer.DISPLAY_NAME, - self._dialog.fluorescence_parameters_view.algorithm_combo_box.count(), - ) - self._dialog.fluorescence_parameters_view.stacked_widget.addWidget( - self._ptychozoon_view_controller.get_widget() - ) - - self._dialog.fluorescence_parameters_view.algorithm_combo_box.textActivated.connect( - enhancer_chooser.set_current_plugin - ) - self._dialog.fluorescence_parameters_view.algorithm_combo_box.currentIndexChanged.connect( - self._dialog.fluorescence_parameters_view.stacked_widget.setCurrentIndex - ) - self._dialog.fluorescence_parameters_view.algorithm_combo_box.setModel( - self._enhancement_model - ) - - self._dialog.fluorescence_parameters_view.enhance_button.clicked.connect( - self._enhance_fluorescence - ) - self._dialog.fluorescence_parameters_view.save_button.clicked.connect( - self._save_enhanced_dataset - ) - - self._dialog.fluorescence_channel_list_view.setModel(self._channel_list_model) - connect_current_changed_signal( - self._dialog.fluorescence_channel_list_view, self._update_view - ) - - self._measured_widget_controller = VisualizationWidgetController( - engine, - self._dialog.measured_widget, - self._dialog.status_bar, - file_dialog_factory, - ) - self._enhanced_widget_controller = VisualizationWidgetController( - engine, - self._dialog.enhanced_widget, - self._dialog.status_bar, - file_dialog_factory, - ) - self._visualization_parameters_controller = VisualizationParametersController( - engine, self._dialog.visualization_parameters_view - ) - - enhancer_chooser.add_observer(self) - self._dataset_emitter.add_observer(self) - self._sync_algorithm_combo_box() - - def get_num_channels(self) -> int: - return 0 if self._measured is None else len(self._measured.element_maps) - - def get_measured_element_map(self, channel_index: int) -> ElementMap | None: - if self._measured is None: - return None - - return self._measured.element_maps[channel_index] - - def get_enhanced_element_map(self, channel_index: int) -> ElementMap | None: - if self._enhanced is not None: - return self._enhanced.element_maps[channel_index] - - return self.get_measured_element_map(channel_index) - - def _reset_channel_list(self) -> None: - self._channel_list_model.beginResetModel() - self._channel_list_model.endResetModel() - - def _open_measured_dataset(self) -> None: - title = 'Open Measured Fluorescence Dataset' - file_path, name_filter = self._file_dialog_factory.get_open_file_path( - self._dialog, - title, - name_filters=[nf for nf in self._fluorescence_api.get_open_file_filters()], - selected_name_filter=self._fluorescence_api.get_open_file_filter(), - ) - - if file_path: - try: - self._measured = self._fluorescence_api.load_measured_dataset( - file_path, file_type=name_filter - ) - except Exception as err: - logger.exception(err) - ExceptionDialog.show_exception(title, err) - else: - self._enhanced = None - self._reset_channel_list() - - def _enhance_fluorescence(self) -> None: - if self._measured is None: - ExceptionDialog.show_exception( - 'Enhance Fluorescence', ValueError('Fluorescence dataset not loaded!') - ) - return - - try: - self._fluorescence_api.enhance(self._product_index, self._measured) - except Exception as err: - logger.exception(err) - ExceptionDialog.show_exception('Enhance Fluorescence', err) - - def launch(self, product_index: int) -> None: - self._product_index = product_index - self._measured = None - self._enhanced = None - self._reset_channel_list() - - try: - item_name = self._fluorescence_api.get_product_name(product_index) - except Exception as err: - logger.exception(err) - ExceptionDialog.show_exception('Launch', err) - else: - self._dialog.setWindowTitle(f'Enhance Fluorescence: {item_name}') - self._dialog.open() - - def _save_enhanced_dataset(self) -> None: - title = 'Save Enhanced Fluorescence Dataset' - - if self._enhanced is None: - ExceptionDialog.show_exception(title, ValueError('Fluorescence dataset not enhanced!')) - return - - file_path, name_filter = self._file_dialog_factory.get_save_file_path( - self._dialog, - title, - name_filters=[nf for nf in self._fluorescence_api.get_save_file_filters()], - selected_name_filter=self._fluorescence_api.get_save_file_filter(), - ) - - if file_path: - try: - self._fluorescence_api.save_enhanced_dataset( - self._enhanced, file_path, file_type=name_filter - ) - except Exception as err: - logger.exception(err) - ExceptionDialog.show_exception(title, err) - - def _sync_algorithm_combo_box(self) -> None: - self._dialog.fluorescence_parameters_view.algorithm_combo_box.setCurrentText( - self._enhancer_chooser.get_current_plugin().display_name - ) - - def _update_view(self, current: QModelIndex, previous: QModelIndex) -> None: - if not current.isValid(): - self._measured_widget_controller.clear_array() - self._enhanced_widget_controller.clear_array() - return - - try: - pixel_geometry = self._fluorescence_api.get_pixel_geometry(self._product_index) - except Exception as err: - logger.exception(err) - self._measured_widget_controller.clear_array() - self._enhanced_widget_controller.clear_array() - ExceptionDialog.show_exception('Render Element Map', err) - return - - emap_measured = self.get_measured_element_map(current.row()) - - if emap_measured is None: - self._measured_widget_controller.clear_array() - else: - self._measured_widget_controller.set_array( - emap_measured.counts_per_second, pixel_geometry - ) - - emap_enhanced = self.get_enhanced_element_map(current.row()) - - if emap_enhanced is None: - self._enhanced_widget_controller.clear_array() - else: - self._enhanced_widget_controller.set_array( - emap_enhanced.counts_per_second, pixel_geometry - ) - - def _update(self, observable: Observable) -> None: - if observable is self._enhancer_chooser: - self._sync_algorithm_combo_box() - elif observable is self._dataset_emitter: - self._enhanced = self._dataset_emitter.get_latest_enhanced() - self._reset_channel_list() diff --git a/src/ptychodus/controller/probe/tree_model.py b/src/ptychodus/controller/probe/tree_model.py index a91bf4eff..28182d3b5 100644 --- a/src/ptychodus/controller/probe/tree_model.py +++ b/src/ptychodus/controller/probe/tree_model.py @@ -47,6 +47,16 @@ def calc_coherent_percent(probe: Probe) -> int: return int(100.0 * coherence) if numpy.isfinite(coherence) else -1 +def try_get_probe(item: ProbeRepositoryItem) -> Probe | None: + # Returns None when the probe sequence is the null sentinel (no dataset + # bound yet — see ProbeRepositoryItem._rebuild). Same pattern used by + # ProbeRepositoryItem.get_size_metrics / get_entropy_metrics. + try: + return item.get_probes().get_probe_no_opr() + except ValueError: + return None + + class ProbeTreeModel(QAbstractItemModel): def __init__( self, @@ -175,47 +185,57 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A case 0: return f'Mode {index.row() + 1}' case 1: - probe = item.get_probes().get_probe_no_opr() # TODO OPR + probe = try_get_probe(item) + if probe is None: + return None power_percent = calc_relative_power_percent(probe, index.row()) return f'{power_percent}%' elif role == Qt.ItemDataRole.BackgroundRole: if index.flags() & Qt.ItemFlag.ItemIsEditable: return self._editable_item_brush elif role == Qt.ItemDataRole.UserRole and index.column() == 1: - probe = item.get_probes().get_probe_no_opr() # TODO OPR + probe = try_get_probe(item) + if probe is None: + return None return calc_relative_power_percent(probe, index.row()) else: item = self._repository[index.row()] probes = item.get_probes() - probe = probes.get_probe_no_opr() # TODO OPR - pixel_geometry = probe.get_pixel_geometry() + probe = try_get_probe(item) + # None when the probe is not yet built (see ProbeRepositoryItem._rebuild + # guard): probe-dependent columns return None; name/builder/size still show. + pixel_geometry = probe.get_pixel_geometry() if probe is not None else None if role == Qt.ItemDataRole.DisplayRole: match index.column(): case 0: return self._repository.get_name(index.row()) case 1: + if probe is None: + return None coherent_percent = calc_coherent_percent(probe) return f'{coherent_percent}%' case 2: return item.get_builder().get_name() case 3: - return str(probe.dtype) + return str(probe.dtype) if probe is not None else None case 4: - return probe.width_px + return probe.width_px if probe is not None else None case 5: - return probe.height_px + return probe.height_px if probe is not None else None case 6: - return f'{pixel_geometry.width_m * 1e9:.4g}' + return f'{pixel_geometry.width_m * 1e9:.4g}' if pixel_geometry else None case 7: - return f'{pixel_geometry.height_m * 1e9:.4g}' + return f'{pixel_geometry.height_m * 1e9:.4g}' if pixel_geometry else None case 8: return f'{probes.nbytes / BYTES_PER_MEGABYTE:.2f}' elif role == Qt.ItemDataRole.BackgroundRole: if index.flags() & Qt.ItemFlag.ItemIsEditable: return self._editable_item_brush elif role == Qt.ItemDataRole.UserRole and index.column() == 1: - probe = item.get_probes().get_probe_no_opr() # TODO OPR + probe = try_get_probe(item) + if probe is None: + return None return calc_coherent_percent(probe) def flags(self, index: QModelIndex) -> Qt.ItemFlags: diff --git a/src/ptychodus/controller/probe_positions/editor_factory.py b/src/ptychodus/controller/probe_positions/editor_factory.py index b8b3717fe..692ea2d25 100644 --- a/src/ptychodus/controller/probe_positions/editor_factory.py +++ b/src/ptychodus/controller/probe_positions/editor_factory.py @@ -129,9 +129,33 @@ class ProbePositionsEditorViewControllerFactory: def _append_common_controls( self, dialog_builder: ParameterViewBuilder, item: ProbePositionsRepositoryItem ) -> None: - dialog_builder.add_view_controller_to_bottom( - ProbePositionsTransformViewController(item.get_builder()) - ) + builder = item.get_builder() + + # A from-memory builder holds already-conditioned positions (reconstruction + # output, or a product loaded from file), so its trim and transform + # parameters are deliberately inert; do not offer controls that would do + # nothing. The bounding box lives on the item and only widens the object + # canvas, so it stays available for every builder. + if not isinstance(builder, FromMemoryProbePositionsBuilder): + trim_group = 'Trim Probe Positions' + dialog_builder.add_spin_box( + builder.num_discard_at_start, + 'Discard at Start:', + tool_tip='Number of probe positions to discard from the beginning' + ' of the scan, in acquisition order.', + group=trim_group, + ) + dialog_builder.add_spin_box( + builder.num_discard_at_end, + 'Discard at End:', + tool_tip='Number of probe positions to discard from the end' + ' of the scan, in acquisition order.', + group=trim_group, + ) + dialog_builder.add_view_controller_to_bottom( + ProbePositionsTransformViewController(builder) + ) + dialog_builder.add_view_controller_to_bottom(ScanBoundingBoxViewController(item)) def create_editor_dialog( diff --git a/src/ptychodus/controller/processing/core.py b/src/ptychodus/controller/processing/core.py index 2c5c957e9..18da01cbd 100644 --- a/src/ptychodus/controller/processing/core.py +++ b/src/ptychodus/controller/processing/core.py @@ -18,12 +18,14 @@ from ...model.genesis import GenesisCore from ...model.globus import GlobusCore from ...model.processing import ProcessingAPI, ProcessingAlgorithmParameter +from ...model.processing.subprocess_reconstructor import SubprocessReconstructor from ...model.product import ProductRepository from ...view.processing import ProcessingActionsView, ProcessingStatusView from ...view.widgets import ExceptionDialog from ..data import FileDialogFactory from ..helpers import connect_triggered_signal from ..parametric import ComboBoxParameterViewController +from ..product.core import ProductRepositoryTableModel from .parameters import ( ProcessingStatusController, ProductParameterViewController, @@ -50,6 +52,7 @@ def __init__( algorithm_parameter: ProcessingAlgorithmParameter, processing_api: ProcessingAPI, product_repository: ProductRepository, + product_table_model: ProductRepositoryTableModel, globus: GlobusCore, genesis: GenesisCore, view: QWidget, @@ -85,7 +88,7 @@ def __init__( product_repository, processing_api.task_monitor, status_view ) self._product_view_controller = ProductParameterViewController( - product_repository, self._status_controller + product_repository, product_table_model, self._status_controller ) self._compute_view_controller = ComputeParameterViewController( globus_supported=globus.is_supported, @@ -119,9 +122,37 @@ def __init__( ) connect_triggered_signal(self._export_training_data_action, self._export_training_data) + combo = self._product_view_controller.get_widget() + combo.currentIndexChanged.connect(self._sync_action_buttons_to_selection) + combo_model = combo.model() + combo_model.dataChanged.connect(lambda *_: self._sync_action_buttons_to_selection()) + combo_model.rowsInserted.connect(lambda *_: self._sync_action_buttons_to_selection()) + combo_model.rowsRemoved.connect(lambda *_: self._sync_action_buttons_to_selection()) + self._sync_action_buttons_to_selection() + self._sync_model_to_view() algorithm_parameter.add_observer(self) + def _sync_action_buttons_to_selection(self) -> None: + index = self._product_view_controller.get_widget().currentIndex() + can_act = False + + if index >= 0: + try: + item = self._product_repository[index] + except IndexError: + pass + else: + can_act = ( + not item.is_pending() + and not item.is_failed() + and item.get_dataset() is not None + ) + + self._actions_view.reconstruct_button.setEnabled(can_act) + self._actions_view.train_button.setEnabled(can_act) + self._export_training_data_action.setEnabled(can_act) + def _populate_stacked_widget( self, view_controller_factories: Iterable[ReconstructorViewControllerFactory] ) -> None: @@ -140,12 +171,31 @@ def _populate_stacked_widget( self._stacked_widget.addWidget(widget) + def _has_dataset_or_warn(self, product_index: int, action: str) -> bool: + try: + item = self._product_repository[product_index] + except IndexError: + logger.warning(f'Cannot {action}: no product at index {product_index}.') + return False + + if item.get_dataset() is None: + logger.warning( + f'Cannot {action}: product "{item.get_name()}" ' + 'has no associated diffraction dataset.' + ) + return False + + return True + def _reconstruct(self) -> None: input_product_index = self._product_view_controller.get_widget().currentIndex() if input_product_index < 0: return + if not self._has_dataset_or_warn(input_product_index, 'reconstruct'): + return + if self._compute_view_controller.is_globus_button_checked(): try: self._globus.executor.reconstruct(input_product_index) @@ -160,7 +210,9 @@ def _reconstruct(self) -> None: ExceptionDialog.show_exception('Reconstruct Remote (Genesis)', exc) else: # local try: - output_product_index = self._processing_api.reconstruct(input_product_index) + output_product_index = self._processing_api.reconstruct( + product_index=input_product_index + ) except Exception as exc: logger.exception(exc) ExceptionDialog.show_exception('Reconstruct Local', exc) @@ -173,6 +225,9 @@ def _train(self) -> None: if product_index < 0: return + if not self._has_dataset_or_warn(product_index, 'train'): + return + if self._compute_view_controller.is_globus_button_checked(): try: self._globus.executor.train(product_index) @@ -194,7 +249,11 @@ def _train(self) -> None: return try: - self._processing_api.train(product_index, data_path, data_path) + self._processing_api.train( + data_path, + data_path, + product_index=product_index, + ) except Exception as exc: logger.exception(exc) ExceptionDialog.show_exception('Train Local', exc) @@ -238,8 +297,11 @@ def _export_training_data(self) -> None: if not file_path: return + if not self._has_dataset_or_warn(input_product_index, 'export training data'): + return + try: - self._processing_api.export_training_data(file_path, input_product_index) + self._processing_api.export_training_data(file_path, product_index=input_product_index) except Exception as exc: logger.exception(exc) ExceptionDialog.show_exception('Export Training Data', exc) @@ -249,10 +311,13 @@ def _sync_model_to_view(self) -> None: self._algorithm_view_controller.get_widget().currentIndex() ) reconstructor = self._algorithm_parameter.get_current_reconstructor() - is_trainable = False + is_trainable = isinstance(reconstructor, TrainableReconstructor) + + if isinstance(reconstructor, SubprocessReconstructor): + is_trainable = reconstructor.is_trainable - if isinstance(reconstructor, TrainableReconstructor): - is_trainable = True + if is_trainable: + assert isinstance(reconstructor, TrainableReconstructor) is_model_loaded = reconstructor.is_model_loaded() self._actions_view.reconstruct_button.setText('Infer') self._actions_view.reconstruct_button.setEnabled(is_model_loaded) diff --git a/src/ptychodus/controller/processing/parameters.py b/src/ptychodus/controller/processing/parameters.py index 7a5d2c89f..ee05edba2 100644 --- a/src/ptychodus/controller/processing/parameters.py +++ b/src/ptychodus/controller/processing/parameters.py @@ -2,7 +2,7 @@ from collections.abc import Sequence import logging -from PyQt5.QtCore import Qt, QModelIndex +from PyQt5.QtCore import QModelIndex from PyQt5.QtWidgets import ( QButtonGroup, QComboBox, @@ -22,7 +22,7 @@ from ...model.product.probe_positions import ProbePositionsRepositoryItem from ...view.processing import ProcessingStatusView from ..parametric import ParameterViewController -from ..product.list_model import ProductRepositoryListModel +from ..product.core import ProductRepositoryComboProxyModel, ProductRepositoryTableModel logger = logging.getLogger(__name__) @@ -93,17 +93,25 @@ def _update(self, observable: Observable) -> None: class ProductParameterViewController(ParameterViewController, ProductRepositoryObserver): + """Combobox for choosing the product to reconstruct on. + + The Qt model (a proxy over the shared ProductRepositoryTableModel) keeps + the combobox in sync with the repository automatically; the only reason + this class still observes the repository is to refresh the loss plot when + the currently-selected product finishes an epoch. + """ + def __init__( self, repository: ProductRepository, + product_table_model: ProductRepositoryTableModel, status_controller: ProcessingStatusController, *, tool_tip: str = '', ) -> None: super().__init__() - self._repository = repository self._status_controller = status_controller - self._model = ProductRepositoryListModel(repository) + self._model = ProductRepositoryComboProxyModel(product_table_model, repository) self._widget = QComboBox() if tool_tip: @@ -112,20 +120,33 @@ def __init__( self._widget.setModel(self._model) self._widget.currentIndexChanged.connect(status_controller.plot_losses) + self._model.rowsInserted.connect(lambda *_: self._auto_select_first_if_empty()) + self._model.rowsRemoved.connect(self._on_rows_removed) + self._auto_select_first_if_empty() + repository.add_observer(self) def get_widget(self) -> QComboBox: return self._widget + def _auto_select_first_if_empty(self) -> None: + if self._widget.currentIndex() >= 0: + return + if self._widget.count() > 0: + self._widget.setCurrentIndex(0) + + def _on_rows_removed(self, parent: QModelIndex, first: int, last: int) -> None: + if self._widget.currentIndex() >= 0: + return + row_count = self._widget.count() + if row_count > 0: + self._widget.setCurrentIndex(min(first, row_count - 1)) + def handle_item_inserted(self, index: int, item: ProductRepositoryItem) -> None: - parent = QModelIndex() - self._model.beginInsertRows(parent, index, index) - self._model.endInsertRows() + pass def handle_metadata_changed(self, index: int, item: MetadataRepositoryItem) -> None: - top_left = self._model.index(index, 0) - bottom_right = self._model.index(index, 0) - self._model.dataChanged.emit(top_left, bottom_right, [Qt.ItemDataRole.DisplayRole]) + pass def handle_probe_positions_changed( self, index: int, item: ProbePositionsRepositoryItem @@ -139,15 +160,17 @@ def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: pass def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: - current_index = self._widget.currentIndex() - - if index == current_index: + if index == self._widget.currentIndex(): self._status_controller.plot_losses(index) + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: - parent = QModelIndex() - self._model.beginRemoveRows(parent, index, index) - self._model.endRemoveRows() + pass class ComputeParameterViewController(ParameterViewController): diff --git a/src/ptychodus/controller/product/core.py b/src/ptychodus/controller/product/core.py index 491a97da3..b5869414d 100644 --- a/src/ptychodus/controller/product/core.py +++ b/src/ptychodus/controller/product/core.py @@ -1,28 +1,29 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Any +from typing import Any, cast import logging from PyQt5.QtCore import ( QAbstractTableModel, + QIdentityProxyModel, QModelIndex, QObject, QSortFilterProxyModel, Qt, ) -from PyQt5.QtGui import QBrush -from PyQt5.QtWidgets import QAbstractItemView, QAction +from PyQt5.QtGui import QBrush, QFont +from PyQt5.QtWidgets import QAbstractItemView, QAction, QInputDialog from ptychodus.api.common import BYTES_PER_MEGABYTE from ptychodus.api.product import LossValue +from ...model.diffraction import AssembledDiffractionDataset, DiffractionDatasetRepository from ...model.product import ( ProductAPI, ProductRepository, ProductRepositoryItem, ProductRepositoryObserver, ) -from ...model.diffraction import DiffractionAPI from ...model.product.metadata import MetadataRepositoryItem from ...model.product.object import ObjectRepositoryItem from ...model.product.probe import ProbeRepositoryItem @@ -30,17 +31,21 @@ from ...view.product import ProductView from ...view.widgets import ExceptionDialog from ..data import FileDialogFactory -from ..helpers import ( - connect_current_changed_signal, - connect_triggered_signal, - create_brush_for_editable_cell, -) +from ..helpers import connect_triggered_signal from .editor import ProductEditorViewController logger = logging.getLogger(__name__) class ProductRepositoryTableModel(QAbstractTableModel): + """Table model over ProductRepository. + + Registers itself as a ProductRepositoryObserver and translates repository + change callbacks into Qt structural / dataChanged signals. Duck-typed + against ProductRepositoryObserver — inheriting the ABC would clash with + sip's wrappertype metaclass on QAbstractTableModel. + """ + def __init__( self, repository: ProductRepository, @@ -59,12 +64,20 @@ def __init__( 'Pixel Height\n[nm]', 'Size\n[MB]', ] + # Duck-typed; see class docstring on the ABC / sip metaclass conflict. + repository.add_observer(cast(ProductRepositoryObserver, self)) def flags(self, index: QModelIndex) -> Qt.ItemFlags: value = super().flags(index) if index.isValid() and index.column() < 4: - value |= Qt.ItemFlag.ItemIsEditable + try: + item = self._repository[index.row()] + except IndexError: + return value + + if not item.is_pending() and not item.is_failed(): + value |= Qt.ItemFlag.ItemIsEditable return value @@ -87,26 +100,54 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A metadata_item = item.get_metadata_item() geometry = item.get_geometry() + pending = item.is_pending() + failed = item.is_failed() if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole: match index.column(): case 0: return metadata_item.name.get_value() case 1: + if pending or failed: + return '—' return f'{metadata_item.detector_distance_m.get_value():.4g}' case 2: + if pending or failed: + return '—' return f'{metadata_item.probe_energy_eV.get_value() / 1e3:.4g}' case 3: + if pending or failed: + return '—' return f'{metadata_item.probe_photon_count.get_value():.4g}' case 4: - return f'{geometry.object_plane_pixel_width_m * 1e9:.4g}' + if pending or failed: + return '—' + return f'{geometry.get_object_plane_pixel_geometry().width_m * 1e9:.4g}' case 5: - return f'{geometry.object_plane_pixel_height_m * 1e9:.4g}' + if pending or failed: + return '—' + return f'{geometry.get_object_plane_pixel_geometry().height_m * 1e9:.4g}' case 6: + if pending or failed: + return '—' product = item.get_product() return f'{product.nbytes / BYTES_PER_MEGABYTE:.2f}' + elif role == Qt.ItemDataRole.FontRole: + if pending or failed: + font = QFont() + font.setItalic(pending) + font.setStrikeOut(failed) + return font + elif role == Qt.ItemDataRole.ToolTipRole: + if pending: + return 'Loading…' + if failed: + return 'Load failed' + elif role == Qt.ItemDataRole.ForegroundRole: + if pending or failed: + return QBrush(Qt.GlobalColor.gray) elif role == Qt.ItemDataRole.BackgroundRole: - if index.flags() & Qt.ItemFlag.ItemIsEditable: + if not (pending or failed) and (index.flags() & Qt.ItemFlag.ItemIsEditable): return self._editable_item_brush def setData(self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.EditRole) -> bool: # noqa: N802 @@ -155,13 +196,95 @@ def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 return len(self._header) + def _emit_row_changed(self, index: int, roles: list[int]) -> None: + top_left = self.index(index, 0) + bottom_right = self.index(index, self.columnCount() - 1) + self.dataChanged.emit(top_left, bottom_right, roles) + + def handle_item_inserted(self, index: int, item: ProductRepositoryItem) -> None: + self.beginInsertRows(QModelIndex(), index, index) + self.endInsertRows() + + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: + self.beginRemoveRows(QModelIndex(), index, index) + self.endRemoveRows() + + def handle_metadata_changed(self, index: int, item: MetadataRepositoryItem) -> None: + self._emit_row_changed(index, [Qt.ItemDataRole.DisplayRole]) + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + self._emit_row_changed( + index, + [ + Qt.ItemDataRole.DisplayRole, + Qt.ItemDataRole.ForegroundRole, + Qt.ItemDataRole.FontRole, + Qt.ItemDataRole.ToolTipRole, + ], + ) + + def handle_probe_positions_changed( + self, index: int, item: ProbePositionsRepositoryItem + ) -> None: + pass + + def handle_probe_changed(self, index: int, item: ProbeRepositoryItem) -> None: + pass + + def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: + pass + + def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: + pass + + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + +class ProductRepositoryComboProxyModel(QIdentityProxyModel): + """Presents ProductRepositoryTableModel as a single-column, non-editable + list suitable for a QComboBox: strips ItemIsEditable and disables pending / + failed items so they cannot be selected. + """ + + def __init__( + self, + source_model: ProductRepositoryTableModel, + repository: ProductRepository, + parent: QObject | None = None, + ) -> None: + super().__init__(parent) + self._repository = repository + self.setSourceModel(source_model) + + def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 + return 1 + + def flags(self, index: QModelIndex) -> Qt.ItemFlags: + base = super().flags(index) + + if not index.isValid(): + return base + + base &= ~Qt.ItemFlag.ItemIsEditable + + try: + item = self._repository[index.row()] + except IndexError: + return base + + if item.is_pending() or item.is_failed(): + return Qt.ItemFlags(Qt.NoItemFlags) + + return base + class ProductController(ProductRepositoryObserver): def __init__( self, - diffraction_api: DiffractionAPI, repository: ProductRepository, api: ProductAPI, + diffraction_repository: DiffractionDatasetRepository, view: ProductView, file_dialog_factory: FileDialogFactory, duplicate_action: QAction, @@ -169,9 +292,9 @@ def __init__( table_proxy_model: QSortFilterProxyModel, ) -> None: super().__init__() - self._diffraction_api = diffraction_api self._repository = repository self._api = api + self._diffraction_repository = diffraction_repository self._view = view self._file_dialog_factory = file_dialog_factory self._duplicate_action = duplicate_action @@ -181,9 +304,10 @@ def __init__( @classmethod def create_instance( cls, - diffraction_api: DiffractionAPI, repository: ProductRepository, api: ProductAPI, + diffraction_repository: DiffractionDatasetRepository, + table_model: ProductRepositoryTableModel, view: ProductView, file_dialog_factory: FileDialogFactory, ) -> ProductController: @@ -193,16 +317,13 @@ def create_instance( save_file_action = view.button_box.save_menu.addAction('Save File...') sync_to_settings_action = view.button_box.save_menu.addAction('Sync To Settings') - editable_item_brush = create_brush_for_editable_cell(view.table_view) - table_model = ProductRepositoryTableModel(repository, editable_item_brush) - table_proxy_model = QSortFilterProxyModel() table_proxy_model.setSourceModel(table_model) controller = cls( - diffraction_api, repository, api, + diffraction_repository, view, file_dialog_factory, duplicate_action, @@ -223,9 +344,18 @@ def create_instance( view.table_view.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows) header = view.table_view.horizontalHeader() header.setSectionResizeMode(header.ResizeMode.ResizeToContents) - connect_current_changed_signal(view.table_view, controller._update_enabled_buttons) + header.setHighlightSections(False) + + selection_model = view.table_view.selectionModel() + if selection_model is not None: + selection_model.currentChanged.connect(controller._update_enabled_buttons) controller._update_enabled_buttons(QModelIndex(), QModelIndex()) - controller._ensure_selection() + + # Auto-select first row on insert if nothing is currently selected; + # re-select a neighbor on removal if the current row went away. + table_proxy_model.rowsInserted.connect(controller._on_rows_inserted) + table_proxy_model.rowsRemoved.connect(controller._on_rows_removed) + controller._auto_select_first_row_if_empty() connect_triggered_signal(open_file_action, controller._open_product_from_file) connect_triggered_signal(create_new_action, controller._create_new_product) @@ -262,6 +392,30 @@ def _get_current_item_index(self) -> int: return item_index + def _choose_dataset(self, title: str) -> tuple[bool, AssembledDiffractionDataset | None]: + """Prompt the user to bind the new product to a diffraction dataset. + + Returns (accepted, dataset). ``dataset`` is None when the user selects "None". + """ + datasets = list(self._diffraction_repository) + labels = ['None', *(dataset.get_name() for dataset in datasets)] + + label, accepted = QInputDialog.getItem( + self._view, + title, + 'Diffraction Dataset:', + labels, + 0, + False, + ) + + if not accepted: + return False, None + + row = labels.index(label) - 1 + dataset = datasets[row] if 0 <= row < len(datasets) else None + return True, dataset + def _open_product_from_file(self) -> None: file_path, name_filter = self._file_dialog_factory.get_open_file_path( self._view, @@ -271,14 +425,26 @@ def _open_product_from_file(self) -> None: ) if file_path: + accepted, dataset = self._choose_dataset('Open Product') + + if not accepted: + return + try: - self._api.open_product(file_path, file_type=name_filter) + self._api.open_product( + file_path, file_type=name_filter, dataset=dataset, block=False + ) except Exception as err: logger.exception(err) ExceptionDialog.show_exception('File Reader', err) def _create_new_product(self) -> None: - self._api.insert_new_product() + accepted, dataset = self._choose_dataset('Create Product') + + if not accepted: + return + + self._api.insert_new_product(dataset=dataset, block=False) def _save_current_product_to_file(self) -> None: current = self._table_proxy_model.mapToSource(self._view.table_view.currentIndex()) @@ -314,7 +480,9 @@ def _duplicate_current_product(self) -> None: if current.isValid(): like_item = self._repository[current.row()] - self._api.insert_product(like_item.get_product()) + self._api.insert_product( + like_item.get_product(), dataset=like_item.get_dataset(), block=False + ) else: logger.error('No current item!') @@ -323,7 +491,9 @@ def _edit_current_product(self) -> None: if current.isValid(): product = self._repository[current.row()] - ProductEditorViewController.edit_product(self._diffraction_api, product, self._view) + ProductEditorViewController.edit_product( + self._diffraction_repository, product, self._view + ) else: logger.error('No current item!') @@ -336,26 +506,30 @@ def _remove_current_product(self) -> None: logger.error('No current item!') def _update_enabled_buttons(self, current: QModelIndex, previous: QModelIndex) -> None: - enabled = current.isValid() - self._duplicate_action.setEnabled(enabled) - self._view.button_box.save_button.setEnabled(enabled) - self._view.button_box.edit_button.setEnabled(enabled) + source_index = ( + self._table_proxy_model.mapToSource(current) if current.isValid() else current + ) + enabled = source_index.isValid() + + ready = False + if enabled: + try: + item = self._repository[source_index.row()] + except IndexError: + ready = False + else: + ready = not item.is_pending() + + self._duplicate_action.setEnabled(ready) + self._view.button_box.save_button.setEnabled(ready) + self._view.button_box.edit_button.setEnabled(ready) + # Remove is always safe: it drops the row whether pending, failed, or ready. self._view.button_box.remove_button.setEnabled(enabled) def _update_info_text(self) -> None: info_text = self._repository.get_info_text() self._view.info_label.setText(info_text) - def _ensure_selection(self) -> None: - if self._view.table_view.currentIndex().isValid(): - return - - if self._table_model.rowCount() > 0: - source_index = self._table_model.index(0, 0) - self._view.table_view.setCurrentIndex( - self._table_proxy_model.mapFromSource(source_index) - ) - def _current_source_row(self) -> int: proxy_index = self._view.table_view.currentIndex() @@ -364,23 +538,33 @@ def _current_source_row(self) -> int: return self._table_proxy_model.mapToSource(proxy_index).row() - def _select_source_row(self, row: int) -> None: - source_index = self._table_model.index(row, 0) - self._view.table_view.setCurrentIndex(self._table_proxy_model.mapFromSource(source_index)) + def _auto_select_first_row_if_empty(self) -> None: + if self._view.table_view.currentIndex().isValid(): + return + + if self._table_proxy_model.rowCount() > 0: + self._view.table_view.setCurrentIndex(self._table_proxy_model.index(0, 0)) + + def _on_rows_inserted(self, parent: QModelIndex, first: int, last: int) -> None: + # Auto-select the newly-inserted row if nothing is currently selected. + self._auto_select_first_row_if_empty() + + def _on_rows_removed(self, parent: QModelIndex, first: int, last: int) -> None: + # Qt clears the current index when the current row is removed; pick a + # neighbor at the same proxy position so the user is never left without + # a selection while rows remain. + if self._view.table_view.currentIndex().isValid(): + return + + row_count = self._table_proxy_model.rowCount() + if row_count > 0: + target = min(first, row_count - 1) + self._view.table_view.setCurrentIndex(self._table_proxy_model.index(target, 0)) def handle_item_inserted(self, index: int, item: ProductRepositoryItem) -> None: - parent = QModelIndex() - self._table_model.beginInsertRows(parent, index, index) - self._table_model.endInsertRows() self._update_info_text() - if not self._view.table_view.currentIndex().isValid(): - self._select_source_row(index) - def handle_metadata_changed(self, index: int, item: MetadataRepositoryItem) -> None: - top_left = self._table_model.index(index, 0) - bottom_right = self._table_model.index(index, self._table_model.columnCount() - 1) - self._table_model.dataChanged.emit(top_left, bottom_right, [Qt.ItemDataRole.DisplayRole]) self._update_info_text() def handle_probe_positions_changed( @@ -397,15 +581,15 @@ def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: self._update_info_text() - def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: - was_current = self._current_source_row() == index - parent = QModelIndex() - self._table_model.beginRemoveRows(parent, index, index) - self._table_model.endRemoveRows() + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: self._update_info_text() - if was_current: - row_count = self._table_model.rowCount() + if self._current_source_row() == index: + current = self._view.table_view.currentIndex() + self._update_enabled_buttons(current, QModelIndex()) - if row_count > 0: - self._select_source_row(min(index, row_count - 1)) + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: + self._update_info_text() diff --git a/src/ptychodus/controller/product/editor.py b/src/ptychodus/controller/product/editor.py index 9fdcbf3de..bd4389a90 100644 --- a/src/ptychodus/controller/product/editor.py +++ b/src/ptychodus/controller/product/editor.py @@ -1,6 +1,8 @@ from typing import Any +import logging from PyQt5.QtCore import ( + QAbstractItemModel, QAbstractTableModel, QModelIndex, QObject, @@ -8,16 +10,22 @@ Qt, ) from PyQt5.QtGui import QBrush -from PyQt5.QtWidgets import QWidget +from PyQt5.QtWidgets import QComboBox, QStyledItemDelegate, QStyleOptionViewItem, QWidget -from ptychodus.api.diffraction import estimate_probe_photon_count +from ptychodus.api.diffraction import Polarization, estimate_probe_photon_count from ptychodus.api.observer import Observable, Observer -from ...model.diffraction import DiffractionAPI +from ...model.diffraction import ( + AssembledDiffractionDataset, + DiffractionDatasetRepository, + DiffractionDatasetRepositoryObserver, +) from ...model.product import ProductRepositoryItem from ...view.product import ProductEditorDialog from ..helpers import create_brush_for_editable_cell +logger = logging.getLogger(__name__) + class ProductPropertyTableModel(QAbstractTableModel): def __init__( @@ -41,6 +49,8 @@ def __init__( 'Exposure Time [s]', 'Mass Attenuation [m\u00b2/kg]', 'Tomography Angle [deg]', + 'Tilt Angle [deg]', + 'Polarization', 'Fresnel Number', 'Detector Numerical Aperture', 'Depth of Field [nm]', @@ -49,7 +59,7 @@ def __init__( def flags(self, index: QModelIndex) -> Qt.ItemFlags: value = super().flags(index) - if index.isValid() and index.row() in (7, 8, 9): + if index.isValid() and index.row() in (7, 8, 9, 10, 11): value |= Qt.ItemFlag.ItemIsEditable return value @@ -85,9 +95,9 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A case 4: return f'{geometry.probe_power_W:.4g}' case 5: - return f'{geometry.object_plane_pixel_width_m * 1e9:.4g}' + return f'{geometry.get_object_plane_pixel_geometry().width_m * 1e9:.4g}' case 6: - return f'{geometry.object_plane_pixel_height_m * 1e9:.4g}' + return f'{geometry.get_object_plane_pixel_geometry().height_m * 1e9:.4g}' case 7: return f'{metadata_item.exposure_time_s.get_value():.4g}' case 8: @@ -95,16 +105,21 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A case 9: return f'{metadata_item.tomography_angle_deg.get_value():.4g}' case 10: + return f'{metadata_item.tilt_angle_deg.get_value():.4g}' + case 11: + raw = metadata_item.polarization.get_value() + return raw if raw else '(unset)' + case 12: try: return f'{geometry.fresnel_number:.4g}' except ZeroDivisionError: return 'inf' - case 11: + case 13: try: return f'{geometry.detector_numerical_aperture:.4g}' except ZeroDivisionError: return 'inf' - case 12: + case 14: try: return f'{geometry.depth_of_field_m * 1e9:.4g}' except ZeroDivisionError: @@ -142,6 +157,25 @@ def setData(self, index: QModelIndex, value: Any, role: int = Qt.ItemDataRole.Ed metadata_item.tomography_angle_deg.set_value(tomography_angle_deg) return True + case 10: + try: + tilt_angle_deg = float(value) + except ValueError: + return False + + metadata_item.tilt_angle_deg.set_value(tilt_angle_deg) + return True + case 11: + text = str(value) + if text in ('', '(unset)'): + metadata_item.polarization.set_value('') + return True + try: + parsed = Polarization(text) + except ValueError: + return False + metadata_item.polarization.set_value(parsed.value) + return True return False @@ -152,23 +186,75 @@ def columnCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 return len(self._header) -class ProductEditorViewController(Observer): +class _PolarizationDelegate(QStyledItemDelegate): + """Combobox editor for the polarization row of the product-property table. + + Attached with ``setItemDelegateForColumn(1, ...)`` so column sorting through + the proxy model cannot mis-target the row. ``mapToSource`` recovers the + row in the underlying ProductPropertyTableModel — everything except the + polarization row falls through to the base delegate. + """ + + POLARIZATION_SOURCE_ROW = 11 + + def _is_polarization_cell(self, index: QModelIndex) -> bool: + if index.column() != 1: + return False + model = index.model() + source_index = ( + model.mapToSource(index) if isinstance(model, QSortFilterProxyModel) else index + ) + return source_index.row() == self.POLARIZATION_SOURCE_ROW + + def createEditor( # noqa: N802 + self, parent: QWidget, option: QStyleOptionViewItem, index: QModelIndex + ) -> QWidget: + if not self._is_polarization_cell(index): + return super().createEditor(parent, option, index) + combo = QComboBox(parent) + combo.addItem('(unset)', '') + for member in Polarization: + combo.addItem(member.value, member.value) + return combo + + def setEditorData(self, editor: QWidget, index: QModelIndex) -> None: # noqa: N802 + if not isinstance(editor, QComboBox): + super().setEditorData(editor, index) + return + current = index.data(Qt.ItemDataRole.DisplayRole) or '' + lookup = '' if current == '(unset)' else str(current) + target = editor.findData(lookup) + editor.setCurrentIndex(max(0, target)) + + def setModelData( # noqa: N802 + self, editor: QWidget, model: QAbstractItemModel, index: QModelIndex + ) -> None: + if not isinstance(editor, QComboBox): + super().setModelData(editor, model, index) + return + model.setData(index, editor.currentData(), Qt.ItemDataRole.EditRole) + + +class ProductEditorViewController(Observer, DiffractionDatasetRepositoryObserver): def __init__( self, - diffraction_api: DiffractionAPI, + diffraction_repository: DiffractionDatasetRepository, product: ProductRepositoryItem, table_model: ProductPropertyTableModel, dialog: ProductEditorDialog, ) -> None: super().__init__() - self._diffraction_api = diffraction_api + self._diffraction_repository = diffraction_repository self._product = product self._table_model = table_model self._dialog = dialog @classmethod def edit_product( - cls, diffraction_api: DiffractionAPI, product: ProductRepositoryItem, parent: QWidget + cls, + diffraction_repository: DiffractionDatasetRepository, + product: ProductRepositoryItem, + parent: QWidget, ) -> None: dialog = ProductEditorDialog(parent) dialog.setWindowTitle(f'Edit Product: {product.get_name()}') @@ -181,6 +267,7 @@ def edit_product( dialog.table_view.setModel(table_proxy_model) dialog.table_view.setSortingEnabled(True) + dialog.table_view.setItemDelegateForColumn(1, _PolarizationDelegate(dialog.table_view)) vertical_header = dialog.table_view.verticalHeader() if vertical_header is not None: @@ -190,8 +277,9 @@ def edit_product( header.setSectionResizeMode(header.ResizeMode.ResizeToContents) dialog.table_view.resizeRowsToContents() - view_controller = cls(diffraction_api, product, table_model, dialog) + view_controller = cls(diffraction_repository, product, table_model, dialog) product.add_observer(view_controller) + diffraction_repository.add_observer(view_controller) dialog.text_edit.textChanged.connect(view_controller._sync_view_to_model) view_controller._sync_model_to_view() @@ -214,9 +302,17 @@ def _sync_model_to_view(self) -> None: metadata = self._product.get_metadata_item() self._dialog.text_edit.setPlainText(metadata.comments.get_value()) + self._dialog.actions_view.estimate_probe_photon_count_button.setEnabled( + self._product.get_dataset() is not None + ) + def _estimate_probe_photon_count(self) -> None: metadata = self._product.get_metadata_item() - assembled_data = self._diffraction_api.get_assembled_data() + dataset = self._product.get_dataset() + if dataset is None: + logger.warning('Cannot estimate probe photon count: no diffraction dataset selected.') + return + assembled_data = dataset.get_assembled_data() photon_count = estimate_probe_photon_count( assembled_data.get_patterns(), assembled_data.get_bad_pixels() ) @@ -227,7 +323,14 @@ def _estimate_probe_photon_count(self) -> None: def _finish(self, result: int) -> None: self._product.remove_observer(self) + self._diffraction_repository.remove_observer(self) def _update(self, observable: Observable) -> None: if observable is self._product: self._sync_model_to_view() + + def handle_dataset_inserted(self, index: int, dataset: AssembledDiffractionDataset) -> None: + pass + + def handle_dataset_removed(self, index: int, dataset: AssembledDiffractionDataset) -> None: + pass diff --git a/src/ptychodus/controller/product/list_model.py b/src/ptychodus/controller/product/list_model.py deleted file mode 100644 index 037735b0f..000000000 --- a/src/ptychodus/controller/product/list_model.py +++ /dev/null @@ -1,30 +0,0 @@ -from typing import Any -import logging - -from PyQt5.QtCore import QAbstractListModel, QModelIndex, QObject, Qt - -from ...model.product import ProductRepository - -logger = logging.getLogger(__name__) - - -class ProductRepositoryListModel(QAbstractListModel): - def __init__( - self, - repository: ProductRepository, - parent: QObject | None = None, - ) -> None: - super().__init__(parent) - self._repository = repository - - def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: - if index.isValid() and role == Qt.ItemDataRole.DisplayRole: - try: - item = self._repository[index.row()] - except IndexError as err: - logger.exception(err) - else: - return item.get_name() - - def rowCount(self, parent: QModelIndex = QModelIndex()) -> int: # noqa: N802 - return len(self._repository) diff --git a/src/ptychodus/controller/ptycho_fm.py b/src/ptychodus/controller/ptycho_fm.py new file mode 100644 index 000000000..d7f683207 --- /dev/null +++ b/src/ptychodus/controller/ptycho_fm.py @@ -0,0 +1,290 @@ +from PyQt5.QtWidgets import QWidget + +from ..model.ptycho_fm.core import PtychoFMReconstructorLibrary +from .data import FileDialogFactory +from .parametric import ParameterViewBuilder +from .processing import ReconstructorViewControllerFactory + + +class PtychoFMViewControllerFactory(ReconstructorViewControllerFactory): + def __init__( + self, + model: PtychoFMReconstructorLibrary, + file_dialog_factory: FileDialogFactory, + ) -> None: + super().__init__() + self._model = model + self._file_dialog_factory = file_dialog_factory + + @property + def name(self) -> str: + return 'PtychoFM' + + def create_view_controller(self, reconstructor_name: str) -> QWidget: + builder = ParameterViewBuilder(self._file_dialog_factory) + enumerators = self._model.enumerators + + # Data + data_group = 'Data' + data_settings = self._model.data_settings + builder.add_decimal_line_edit( + data_settings.scale, + 'Diffraction Scale:', + tool_tip='Scale factor applied to diffraction intensities before sqrt.', + group=data_group, + ) + builder.add_decimal_line_edit( + data_settings.default_normalization, + 'Default Normalization:', + group=data_group, + ) + builder.add_check_box( + data_settings.packed, + 'Packed Dataset', + tool_tip='Use pre-packed dataset shards instead of paired HDF5 files.', + group=data_group, + ) + builder.add_check_box( + data_settings.cache_object, + 'Cache Object In Memory', + group=data_group, + ) + builder.add_integer_line_edit( + data_settings.max_probe_modes, + 'Max Probe Modes:', + tool_tip='Probes are zero-padded to this many mixed-state modes.', + group=data_group, + ) + builder.add_integer_line_edit( + data_settings.target_size, + 'Target Pattern Size:', + group=data_group, + ) + builder.add_decimal_slider( + data_settings.train_split, + 'Train Split:', + tool_tip='Fraction of dataset used for training; remainder is validation.', + group=data_group, + ) + builder.add_integer_line_edit( + data_settings.random_seed, + 'Random Seed:', + group=data_group, + ) + builder.add_combo_box( + data_settings.sharding_strategy, + enumerators.get_sharding_strategies(), + 'Sharding Strategy:', + group=data_group, + ) + builder.add_integer_line_edit( + data_settings.max_files, + 'Max Files (0 = no cap):', + group=data_group, + ) + builder.add_integer_line_edit( + data_settings.num_workers, + 'Dataloader Workers:', + group=data_group, + ) + builder.add_integer_line_edit( + data_settings.prefetch_factor, + 'Prefetch Factor:', + group=data_group, + ) + builder.add_check_box( + data_settings.use_cuda_prefetcher, + 'Use CUDA Prefetcher', + group=data_group, + ) + + # Model + model_group = 'Model' + model_settings = self._model.model_settings + builder.add_combo_box( + model_settings.encoder_type, + enumerators.get_encoder_types(), + 'Encoder Type:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.img_size, + 'Image Size:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.patch_size, + 'Patch Size:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.embed_dim, + 'Embed Dimension:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.depth, + 'Depth:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.num_heads, + 'Attention Heads:', + group=model_group, + ) + builder.add_decimal_line_edit( + model_settings.mlp_ratio, + 'MLP Ratio:', + group=model_group, + ) + builder.add_check_box( + model_settings.use_cls_token, + 'Use CLS Token', + group=model_group, + ) + builder.add_decimal_slider( + model_settings.dropout, + 'Dropout:', + group=model_group, + ) + builder.add_decimal_slider( + model_settings.attn_dropout, + 'Attention Dropout:', + group=model_group, + ) + builder.add_line_edit( + model_settings.timm_model_name, + 'TIMM Model Name:', + tool_tip="Only used when encoder_type is 'pretrained'.", + group=model_group, + ) + + # Decoder + builder.add_integer_line_edit( + model_settings.decoder_base_channels, + 'Decoder Base Channels:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.decoder_latent_dim, + 'Decoder Latent Dim:', + group=model_group, + ) + builder.add_integer_line_edit( + model_settings.decoder_num_stages, + 'Decoder Upsample Stages:', + group=model_group, + ) + builder.add_check_box( + model_settings.decoder_use_batchnorm, + 'Decoder Batch Norm', + group=model_group, + ) + builder.add_decimal_slider( + model_settings.decoder_dropout, + 'Decoder Dropout:', + group=model_group, + ) + + # Init + builder.add_check_box( + model_settings.init_enabled, + 'Custom Weight Init', + group=model_group, + ) + builder.add_combo_box( + model_settings.init_method, + enumerators.get_init_methods(), + 'Init Method:', + group=model_group, + ) + + # Training + training_group = 'Training' + training_settings = self._model.training_settings + builder.add_combo_box( + training_settings.mode, + enumerators.get_training_modes(), + 'Mode:', + group=training_group, + ) + builder.add_integer_line_edit( + training_settings.epochs, + 'Epochs:', + group=training_group, + ) + builder.add_integer_line_edit( + training_settings.batch_size, + 'Batch Size:', + group=training_group, + ) + builder.add_decimal_line_edit( + training_settings.learning_rate, + 'Learning Rate:', + group=training_group, + ) + builder.add_combo_box( + training_settings.loss_function, + enumerators.get_loss_functions(), + 'Loss Function:', + group=training_group, + ) + builder.add_combo_box( + training_settings.weighted_loss_type, + enumerators.get_weighted_loss_types(), + 'Weighted Loss Type:', + tool_tip="Only used when loss_function is 'weighted'.", + group=training_group, + ) + builder.add_decimal_line_edit( + training_settings.weighted_loss_threshold, + 'Weighted Loss Threshold:', + group=training_group, + ) + builder.add_decimal_line_edit( + training_settings.weighted_loss_alpha, + 'Weighted Loss Alpha:', + group=training_group, + ) + builder.add_integer_line_edit( + training_settings.validation_plot_freq, + 'Validation Plot Freq:', + group=training_group, + ) + builder.add_integer_line_edit( + training_settings.checkpoint_freq, + 'Checkpoint Freq:', + group=training_group, + ) + builder.add_check_box( + training_settings.save_epoch_models, + 'Save Per-Epoch Models', + group=training_group, + ) + builder.add_check_box( + training_settings.resume_from_checkpoint, + 'Resume From Checkpoint', + group=training_group, + ) + + # Inference + inference_group = 'Inference' + inference_settings = self._model.inference_settings + builder.add_integer_line_edit( + inference_settings.central_crop, + 'Central Crop:', + tool_tip='Pixels cropped from each border of every patch before stitching.', + group=inference_group, + ) + builder.add_integer_line_edit( + inference_settings.pad, + 'Fourier Shift Pad:', + group=inference_group, + ) + builder.add_integer_line_edit( + inference_settings.batch_size, + 'Batch Size:', + group=inference_group, + ) + + return builder.build_widget() diff --git a/src/ptychodus/controller/ptychopinn_torch.py b/src/ptychodus/controller/ptychopinn_torch.py index 572d0cdda..2221ce2c8 100644 --- a/src/ptychodus/controller/ptychopinn_torch.py +++ b/src/ptychodus/controller/ptychopinn_torch.py @@ -334,6 +334,35 @@ def create_view_controller(self, reconstructor_name: str) -> QWidget: tool_tip='Device to train on ("cuda", "cpu", etc.)', group=training_group, ) # TODO improve + builder.add_integer_line_edit( + training_settings.n_devices, + 'Number of GPUs:', + tool_tip=( + 'Number of CUDA devices to use for training. Lightning fans out to one ' + 'process per device via the configured distributed strategy.' + ), + group=training_group, + ) + builder.add_line_edit( + training_settings.distributed_strategy, + 'Distributed Strategy:', + tool_tip=( + "Lightning distributed strategy. Recommended: 'ddp_spawn' for " + "single-node multi-GPU. 'ddp' requires torchrun and is not wired in " + "phase 1. 'auto' lets Lightning decide." + ), + group=training_group, + ) + builder.add_line_edit( + training_settings.visible_gpu_indices, + 'Visible GPU Indices:', + tool_tip=( + 'Comma-separated CUDA device indices exposed to the training ' + 'subprocess (via CUDA_VISIBLE_DEVICES). Leave blank to inherit ' + "ptychodus's environment." + ), + group=training_group, + ) builder.add_combo_box( training_settings.learning_rate_scheduler, enumerators.get_learning_rate_schedulers(), diff --git a/src/ptychodus/controller/settings.py b/src/ptychodus/controller/settings.py index 46c27c45b..f7a5368c2 100644 --- a/src/ptychodus/controller/settings.py +++ b/src/ptychodus/controller/settings.py @@ -12,7 +12,7 @@ from ..model.product import ProductRepository from ..view.settings import SettingsView, SyncProductToSettingsDialog from .data import FileDialogFactory -from .product.list_model import ProductRepositoryListModel +from .product.core import ProductRepositoryComboProxyModel, ProductRepositoryTableModel logger = logging.getLogger(__name__) @@ -60,6 +60,7 @@ def __init__( self, settings_registry: SettingsRegistry, product_repository: ProductRepository, + product_table_model: ProductRepositoryTableModel, settings_view: SettingsView, settings_table_view: QTableView, file_dialog_factory: FileDialogFactory, @@ -73,7 +74,9 @@ def __init__( self._settings_list_model = QStringListModel() self._settings_table_model = SettingsTableModel() - self._product_list_model = ProductRepositoryListModel(product_repository) + self._product_combo_model = ProductRepositoryComboProxyModel( + product_table_model, product_repository + ) self._sync_dialog = SyncProductToSettingsDialog(settings_view) settings_registry.add_observer(self) @@ -95,7 +98,7 @@ def __init__( save_button = settings_view.button_box.button(QDialogButtonBox.StandardButton.Save) save_button.clicked.connect(self._sync_dialog.open) - self._sync_dialog.product_combo_box.setModel(self._product_list_model) + self._sync_dialog.product_combo_box.setModel(self._product_combo_model) self._sync_dialog.finished.connect(self._save_settings) self._sync_model_to_view() diff --git a/src/ptychodus/controller/visualization/controller.py b/src/ptychodus/controller/visualization/controller.py index dc1658f17..2721e7b9a 100644 --- a/src/ptychodus/controller/visualization/controller.py +++ b/src/ptychodus/controller/visualization/controller.py @@ -108,15 +108,37 @@ def _analyze_line_cut(self, line: QLineF) -> None: logger.warning('No visualization product!') return - value_label = product.get_value_label() line_cut = product.get_line_cut(line2d) - ax = self._line_cut_dialog.axes - ax.clear() - ax.plot(line_cut.distance_m, line_cut.value, '.-', linewidth=1.5) - ax.set_xlabel('Distance [m]') - ax.set_ylabel(value_label) - ax.grid(True) + if not line_cut.series: + logger.warning('Line-cut has no series!') + return + + axes = self._line_cut_dialog.prepare_axes(len(line_cut.series)) + artists = list() + + for index, (ax, series) in enumerate(zip(axes, line_cut.series)): + # Each twinned axis restarts the property cycle, so assign colors explicitly. + color = f'C{index}' + (artist,) = ax.plot( + line_cut.distance_m, + series.value, + '.-', + linewidth=1.5, + color=color, + label=series.label, + ) + artists.append(artist) + ax.set_ylabel(series.label, color=color) + ax.tick_params(axis='y', labelcolor=color) + + primary_axis = axes[0] + primary_axis.set_xlabel('Distance [m]') + primary_axis.grid(True) + + if len(artists) > 1: + primary_axis.legend(artists, [series.label for series in line_cut.series]) + self._line_cut_dialog.figure_canvas.draw() self._line_cut_dialog.open() diff --git a/src/ptychodus/model/agent/__init__.py b/src/ptychodus/model/agent/__init__.py index 939e1a1f6..c2ea85a32 100644 --- a/src/ptychodus/model/agent/__init__.py +++ b/src/ptychodus/model/agent/__init__.py @@ -1,13 +1,17 @@ -from .chat import ChatMessage, ChatHistory, ChatObserver, ChatRole -from .core import AgentCore, AgentPresenter -from .settings import ArgoSettings +from .core import AgentCore +from .model_catalog import ModelCatalog +from .models import ChatMessage, ChatRole +from .repository import ConversationObserver, ConversationRepository +from .settings import AgentSettings +from .terminal import ChatTerminal __all__ = [ 'AgentCore', - 'AgentPresenter', - 'ArgoSettings', - 'ChatHistory', + 'AgentSettings', 'ChatMessage', - 'ChatObserver', 'ChatRole', + 'ChatTerminal', + 'ConversationObserver', + 'ConversationRepository', + 'ModelCatalog', ] diff --git a/src/ptychodus/model/agent/argo.py b/src/ptychodus/model/agent/argo.py deleted file mode 100644 index edaeae630..000000000 --- a/src/ptychodus/model/agent/argo.py +++ /dev/null @@ -1,73 +0,0 @@ -from collections.abc import Sequence -import logging -import requests - -from .chat import ChatHistory, ChatMessage, ChatRole, ChatTerminal -from .settings import ArgoSettings - -logger = logging.getLogger(__name__) - - -class ArgoChatTerminal(ChatTerminal): - def __init__(self, settings: ArgoSettings, history: ChatHistory) -> None: - self._settings = settings - self._history = history - - def send_message(self, content: str, stop: Sequence[str] = []) -> None: - if not content: - return - - messages = [ - ChatMessage( - role=ChatRole.SYSTEM, content='You are a large language model with the name Argo.' - ) - ] - - for line in content.splitlines(): - message = ChatMessage(role=ChatRole.USER, content=line) - messages.append(message) - self._history.add_message(message) - - logger.debug(f'{messages=}') - - url = self._settings.chat_endpoint_url.get_value() - payload = { - 'user': self._settings.user.get_value(), - 'model': self._settings.chat_model.get_value(), - 'messages': [m.to_dict() for m in messages], - 'stop': stop, - 'temperature': self._settings.temperature.get_value(), - 'top_p': self._settings.top_p.get_value(), - 'max_tokens': self._settings.max_tokens.get_value(), - 'max_completion_tokens': self._settings.max_completion_tokens.get_value(), - } - headers = {'Content-Type': 'application/json'} - response = requests.post(url, json=payload, headers=headers) - - logger.debug(f'{response=}') - logger.debug(f'Status Code: {response.status_code}') - response_json = response.json() - logger.debug(f'JSON Response: {response_json}') - response.raise_for_status() - - response_message = ChatMessage(role=ChatRole.AGENT, content=response_json['response']) - self._history.add_message(response_message) - - def embed_texts(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: - """Generates embeddings for a list of strings.""" - url = self._settings.embeddings_endpoint_url.get_value() - payload = { - 'user': self._settings.user.get_value(), - 'model': self._settings.embeddings_model.get_value(), - 'prompt': texts, - } - headers = {'Content-Type': 'application/json'} - response = requests.post(url, json=payload, headers=headers) - - logger.debug(response) - logger.debug(f'Status Code: {response.status_code}') - response_json = response.json() - logger.debug(f'JSON Response: {response_json}') - response.raise_for_status() - - return response_json['embedding'] diff --git a/src/ptychodus/model/agent/chat.py b/src/ptychodus/model/agent/chat.py deleted file mode 100644 index ffd415a07..000000000 --- a/src/ptychodus/model/agent/chat.py +++ /dev/null @@ -1,81 +0,0 @@ -from abc import ABC, abstractmethod -from collections.abc import Sequence -from dataclasses import dataclass -from enum import Enum, auto -from typing import overload - - -class ChatRole(Enum): - SYSTEM = auto() - USER = auto() - AGENT = auto() - - -@dataclass(frozen=True) -class ChatMessage: - role: ChatRole - content: str - - def to_dict(self) -> dict[str, str]: - return { - 'role': self.role.name.lower(), - 'content': self.content, - } - - def __str__(self) -> str: - return str(self.to_dict()) - - -class ChatTerminal(ABC): - @abstractmethod - def send_message(self, content: str) -> None: - pass - - @abstractmethod - def embed_texts(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: - pass - - -class ChatObserver(ABC): - @abstractmethod - def handle_new_message(self, message: ChatMessage, index: int) -> None: - pass - - @abstractmethod - def handle_chat_cleared(self) -> None: - pass - - -class ChatHistory(Sequence[ChatMessage]): - def __init__(self) -> None: - self._messages: list[ChatMessage] = [] - self._observers: list[ChatObserver] = [] - - @overload - def __getitem__(self, index: int) -> ChatMessage: ... - - @overload - def __getitem__(self, index: slice) -> Sequence[ChatMessage]: ... - - def __getitem__(self, index: int | slice) -> ChatMessage | Sequence[ChatMessage]: - return self._messages[index] - - def __len__(self) -> int: - return len(self._messages) - - def add_observer(self, observer: ChatObserver) -> None: - if observer not in self._observers: - self._observers.append(observer) - - def add_message(self, message: ChatMessage) -> None: - index = len(self._messages) - self._messages.append(message) - - for observer in self._observers: - observer.handle_new_message(message, index) - - def clear(self) -> None: - self._messages.clear() - - for observer in self._observers: - observer.handle_chat_cleared() diff --git a/src/ptychodus/model/agent/core.py b/src/ptychodus/model/agent/core.py index 8e5ddee58..dfd39e22d 100644 --- a/src/ptychodus/model/agent/core.py +++ b/src/ptychodus/model/agent/core.py @@ -1,48 +1,14 @@ -from collections.abc import Iterator, Sequence -import logging - from ptychodus.api.settings import SettingsRegistry -from .argo import ArgoChatTerminal -from .chat import ChatHistory, ChatTerminal -from .settings import ArgoSettings - -logger = logging.getLogger(__name__) - - -class AgentPresenter: - def __init__(self, terminal: ChatTerminal) -> None: - self._terminal = terminal - - def get_available_chat_models(self) -> Iterator[str]: - for model in [ - 'gpt35', - 'gpt35large', - 'gpt4', - 'gpt4large', - 'gpt4o', - 'gpt4olatestgpt4turbo', - 'gpto1', - 'gpto1mini', - 'gpto3mini', - ]: - yield model - - def send_message(self, content: str) -> None: - if self._terminal is not None: - self._terminal.send_message(content) - - def get_available_embeddings_models(self) -> Iterator[str]: - for model in ['ada002', 'v3large', 'v3small']: - yield model - - def embed_text(self, texts: Sequence[str]) -> Sequence[Sequence[float]]: - return [[]] if self._terminal is None else self._terminal.embed_texts(texts) +from .model_catalog import ModelCatalog +from .repository import ConversationRepository +from .settings import AgentSettings +from .terminal import ChatTerminal class AgentCore: - def __init__(self, settings_registry: SettingsRegistry): - self.settings = ArgoSettings(settings_registry) - self.chat_history = ChatHistory() - self._terminal = ArgoChatTerminal(self.settings, self.chat_history) - self.presenter = AgentPresenter(self._terminal) + def __init__(self, settings_registry: SettingsRegistry) -> None: + self.settings = AgentSettings(settings_registry) + self.repository = ConversationRepository(self.settings.database_path.get_value()) + self.terminal = ChatTerminal(self.settings, self.repository) + self.catalog = ModelCatalog(self.settings) diff --git a/src/ptychodus/model/agent/model_catalog.py b/src/ptychodus/model/agent/model_catalog.py new file mode 100644 index 000000000..fc992e40c --- /dev/null +++ b/src/ptychodus/model/agent/model_catalog.py @@ -0,0 +1,36 @@ +import logging +import os + +import httpx + +from .settings import AgentSettings + +logger = logging.getLogger(__name__) + + +class ModelCatalog: + """Fetches and caches the list of chat models advertised by the configured endpoint.""" + + def __init__(self, settings: AgentSettings) -> None: + self._settings = settings + self._cached_models: list[str] | None = None + + def get_available_models(self) -> list[str]: + if self._cached_models is None: + return self.refresh() + return self._cached_models + + def refresh(self) -> list[str]: + base_url = self._settings.base_url.get_value().rstrip('/') + api_key = os.environ.get('OPENAI_API_KEY', '') + headers = {'Authorization': f'Bearer {api_key}'} if api_key else {} + try: + response = httpx.get(f'{base_url}/models', timeout=10.0, headers=headers) + response.raise_for_status() + data = response.json()['data'] + models = [str(item['id']) for item in data] + except (httpx.HTTPError, KeyError, ValueError, TypeError) as exc: + logger.error(f'Failed to fetch chat models from {base_url}/models: {exc}') + models = [] + self._cached_models = models + return models diff --git a/src/ptychodus/model/agent/models.py b/src/ptychodus/model/agent/models.py new file mode 100644 index 000000000..16ed02168 --- /dev/null +++ b/src/ptychodus/model/agent/models.py @@ -0,0 +1,16 @@ +from dataclasses import dataclass +from datetime import datetime +from enum import Enum + + +class ChatRole(Enum): + SYSTEM = 'system' + USER = 'user' + ASSISTANT = 'assistant' + + +@dataclass(frozen=True) +class ChatMessage: + role: ChatRole + content: str + created_at: datetime diff --git a/src/ptychodus/model/agent/properties.py b/src/ptychodus/model/agent/properties.py deleted file mode 100644 index 3f58b4be6..000000000 --- a/src/ptychodus/model/agent/properties.py +++ /dev/null @@ -1,50 +0,0 @@ -from collections.abc import Sequence -from dataclasses import dataclass - - -@dataclass(frozen=True) -class ChatModelProperties: - name: str - - @property - def is_o_series_model(self) -> bool: - return self.name.startswith('gpto') - - @property - def accepts_system_prompt(self) -> bool: - return not self.is_o_series_model - - @property - def accepts_stop_sequence(self) -> bool: - return not self.is_o_series_model - - @property - def accepts_temperature(self) -> bool: - return not self.is_o_series_model - - @property - def accepts_top_p(self) -> bool: - return not self.is_o_series_model - - @property - def accepts_max_tokens(self) -> bool: - return not self.is_o_series_model - - @property - def accepts_max_completion_tokens(self) -> bool: - return self.is_o_series_model - - -def list_argo_model_properties() -> Sequence[ChatModelProperties]: - return [ - ChatModelProperties(name='gpt35'), - ChatModelProperties(name='gpt35large'), - ChatModelProperties(name='gpt4'), - ChatModelProperties(name='gpt4large'), - ChatModelProperties(name='gpt4o'), - ChatModelProperties(name='gpt4olatest'), - ChatModelProperties(name='gpt4turbo'), - ChatModelProperties(name='gpto1'), - ChatModelProperties(name='gpto1mini'), - ChatModelProperties(name='gpto3mini'), - ] diff --git a/src/ptychodus/model/agent/repository.py b/src/ptychodus/model/agent/repository.py new file mode 100644 index 000000000..475116e40 --- /dev/null +++ b/src/ptychodus/model/agent/repository.py @@ -0,0 +1,149 @@ +import logging +import sqlite3 +from abc import ABC, abstractmethod +from collections.abc import Sequence +from datetime import datetime +from typing import overload + +from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter + +from .models import ChatMessage + +logger = logging.getLogger(__name__) + + +_SCHEMA = ( + """ + CREATE TABLE IF NOT EXISTS conversations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + created_at TEXT NOT NULL, + model_messages_json TEXT NOT NULL DEFAULT '[]' + ) + """, + """ + CREATE TABLE IF NOT EXISTS messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id INTEGER NOT NULL REFERENCES conversations(id), + role TEXT NOT NULL, + content TEXT NOT NULL, + created_at TEXT NOT NULL + ) + """, + 'CREATE INDEX IF NOT EXISTS idx_messages_conversation ON messages(conversation_id, id)', +) + + +class ConversationObserver(ABC): + @abstractmethod + def handle_message_appended(self, message: ChatMessage, index: int) -> None: + pass + + @abstractmethod + def handle_conversation_cleared(self) -> None: + pass + + +class ConversationRepository(Sequence[ChatMessage]): + """SQLite-backed conversation store. + + Holds one logical "current conversation" per app launch, created lazily on the + first append. Mirrors the current conversation's messages in memory so Qt list + reads stay O(1); writes hit SQLite. The connection is shared across the Qt + thread and the asyncio.run-driven terminal; we never call concurrently because + asyncio.run blocks the caller. + """ + + def __init__(self, database_path: str) -> None: + self._conn = sqlite3.connect( + database_path or ':memory:', + check_same_thread=False, + isolation_level=None, + ) + self._conn.execute('PRAGMA foreign_keys = ON') + for statement in _SCHEMA: + self._conn.execute(statement) + + self._current_conversation_id: int | None = None + self._cached_messages: list[ChatMessage] = [] + self._observers: list[ConversationObserver] = [] + + @overload + def __getitem__(self, index: int) -> ChatMessage: ... + + @overload + def __getitem__(self, index: slice) -> Sequence[ChatMessage]: ... + + def __getitem__(self, index: int | slice) -> ChatMessage | Sequence[ChatMessage]: + return self._cached_messages[index] + + def __len__(self) -> int: + return len(self._cached_messages) + + def add_observer(self, observer: ConversationObserver) -> None: + if observer not in self._observers: + self._observers.append(observer) + + def append(self, message: ChatMessage) -> None: + if self._current_conversation_id is None: + cursor = self._conn.execute( + 'INSERT INTO conversations(created_at) VALUES (?)', + (datetime.now().astimezone().isoformat(),), + ) + self._current_conversation_id = int(cursor.lastrowid or 0) + + self._conn.execute( + 'INSERT INTO messages(conversation_id, role, content, created_at) VALUES (?, ?, ?, ?)', + ( + self._current_conversation_id, + message.role.value, + message.content, + message.created_at.isoformat(), + ), + ) + + index = len(self._cached_messages) + self._cached_messages.append(message) + + for observer in self._observers: + observer.handle_message_appended(message, index) + + def clear(self) -> None: + if self._current_conversation_id is not None: + self._conn.execute( + 'DELETE FROM messages WHERE conversation_id = ?', + (self._current_conversation_id,), + ) + self._conn.execute( + 'DELETE FROM conversations WHERE id = ?', + (self._current_conversation_id,), + ) + + self._current_conversation_id = None + self._cached_messages.clear() + + for observer in self._observers: + observer.handle_conversation_cleared() + + def load_model_messages(self) -> list[ModelMessage]: + if self._current_conversation_id is None: + return [] + row = self._conn.execute( + 'SELECT model_messages_json FROM conversations WHERE id = ?', + (self._current_conversation_id,), + ).fetchone() + if row is None: + return [] + return list(ModelMessagesTypeAdapter.validate_json(row[0])) + + def save_model_messages(self, messages: list[ModelMessage]) -> None: + if self._current_conversation_id is None: + logger.warning('save_model_messages called with no current conversation; ignoring') + return + blob = ModelMessagesTypeAdapter.dump_json(messages).decode() + self._conn.execute( + 'UPDATE conversations SET model_messages_json = ? WHERE id = ?', + (blob, self._current_conversation_id), + ) + + def close(self) -> None: + self._conn.close() diff --git a/src/ptychodus/model/agent/settings.py b/src/ptychodus/model/agent/settings.py index 42f7eeac8..e7c57efd1 100644 --- a/src/ptychodus/model/agent/settings.py +++ b/src/ptychodus/model/agent/settings.py @@ -1,20 +1,20 @@ -import getpass - from ptychodus.api.observer import Observable, Observer from ptychodus.api.settings import SettingsRegistry -class ArgoSettings(Observable, Observer): +class AgentSettings(Observable, Observer): def __init__(self, registry: SettingsRegistry) -> None: super().__init__() - self._group = registry.create_group('Argo') + self._group = registry.create_group('Agent') self._group.add_observer(self) - self.user = self._group.create_string_parameter('User', getpass.getuser()) - self.chat_endpoint_url = self._group.create_string_parameter( - 'ChatEndpointURL', 'https://apps.inside.anl.gov/argoapi/api/v1/resource/chat/' + self.base_url = self._group.create_string_parameter( + 'BaseURL', 'https://apps.inside.anl.gov/argoapi/v1' + ) + self.model = self._group.create_string_parameter('Model', 'GPT-4o') + self.system_prompt = self._group.create_string_parameter( + 'SystemPrompt', 'You are a helpful assistant.' ) - self.chat_model = self._group.create_string_parameter('ChatModel', 'gpt35') self.temperature = self._group.create_real_parameter( 'Temperature', 0.1, minimum=0.0, maximum=2.0 ) @@ -22,13 +22,8 @@ def __init__(self, registry: SettingsRegistry) -> None: self.max_tokens = self._group.create_integer_parameter( 'MaxTokens', 1000, minimum=0, maximum=128000 ) - self.max_completion_tokens = self._group.create_integer_parameter( - 'MaxCompletionTokens', 1000, minimum=0, maximum=128000 - ) - self.embeddings_endpoint_url = self._group.create_string_parameter( - 'EmbeddingsEndpointURL', 'https://apps.inside.anl.gov/argoapi/api/v1/resource/embed/' - ) - self.embeddings_model = self._group.create_string_parameter('EmbeddingsModel', 'ada002') + self.mcp_server_url = self._group.create_string_parameter('MCPServerURL', '') + self.database_path = self._group.create_string_parameter('DatabasePath', '') def _update(self, observable: Observable) -> None: if observable is self._group: diff --git a/src/ptychodus/model/agent/terminal.py b/src/ptychodus/model/agent/terminal.py new file mode 100644 index 000000000..6f7e76b4b --- /dev/null +++ b/src/ptychodus/model/agent/terminal.py @@ -0,0 +1,86 @@ +import asyncio +import logging +import os +from datetime import datetime + +from pydantic_ai import Agent +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.openai import OpenAIProvider +from pydantic_ai.settings import ModelSettings + +from .models import ChatMessage, ChatRole +from .repository import ConversationRepository +from .settings import AgentSettings + +logger = logging.getLogger(__name__) + + +class ChatTerminal: + def __init__(self, settings: AgentSettings, repository: ConversationRepository) -> None: + self._settings = settings + self._repository = repository + + def clear_conversation(self) -> None: + self._repository.clear() + + def send_message(self, content: str) -> None: + if not content.strip(): + return + + self._repository.append( + ChatMessage( + role=ChatRole.USER, + content=content, + created_at=datetime.now().astimezone(), + ) + ) + + # NOTE: pydantic-ai uses Authorization: Bearer headers via the OpenAI SDK. + # The Argo OpenAPI spec instead declares ?authorization= as a query param; + # if the live endpoint rejects the Bearer header, pass a custom + # httpx.AsyncClient(params={'authorization': key}) to OpenAIProvider. + provider = OpenAIProvider( + base_url=self._settings.base_url.get_value(), + api_key=os.environ.get('OPENAI_API_KEY', ''), + ) + model = OpenAIChatModel(self._settings.model.get_value(), provider=provider) + model_settings = ModelSettings( + temperature=self._settings.temperature.get_value(), + top_p=self._settings.top_p.get_value(), + max_tokens=self._settings.max_tokens.get_value(), + ) + # TODO(mcp-followup): when settings.mcp_server_url is non-empty, pass + # toolsets=[MCPServerStreamableHTTP(url=...)] so the agent can call + # ptychodus_store tools (see src/ptychodus_store/mcp_server.py). + agent = Agent( + model, + system_prompt=self._settings.system_prompt.get_value(), + model_settings=model_settings, + ) + + prior = self._repository.load_model_messages() + + try: + # asyncio.run blocks the Qt event loop for the duration of the call, + # same as the previous requests.post. Move to a QThread if it stops + # being acceptable for this developer-mode panel. + result = asyncio.run(agent.run(content, message_history=prior)) + except Exception as exc: + logger.exception('Chat request failed') + self._repository.append( + ChatMessage( + role=ChatRole.ASSISTANT, + content=f'[error] {exc}', + created_at=datetime.now().astimezone(), + ) + ) + return + + self._repository.save_model_messages(list(result.all_messages())) + self._repository.append( + ChatMessage( + role=ChatRole.ASSISTANT, + content=str(result.output), + created_at=datetime.now().astimezone(), + ) + ) diff --git a/src/ptychodus/model/analysis/core.py b/src/ptychodus/model/analysis/core.py index 1c7981d34..351eedd95 100644 --- a/src/ptychodus/model/analysis/core.py +++ b/src/ptychodus/model/analysis/core.py @@ -2,7 +2,7 @@ from ptychodus.api.settings import SettingsRegistry -from ..diffraction import AssembledDiffractionDataset +from ..diffraction import DiffractionDatasetRepository from ..product import ProbePositionsRepository, ProductRepository from ..visualization import VisualizationEngine from .affine import AffineTransformEstimator @@ -25,7 +25,7 @@ def __init__( self, rng: numpy.random.Generator, settings_registry: SettingsRegistry, - dataset: AssembledDiffractionDataset, + diffraction_repository: DiffractionDatasetRepository, product_repository: ProductRepository, probe_positions_repository: ProbePositionsRepository, ) -> None: @@ -38,7 +38,7 @@ def __init__( self.diffraction_simulator_settings = DiffractionSimulatorSettings(settings_registry) self.diffraction_simulator = DiffractionSimulator( - rng, self.diffraction_simulator_settings, dataset, product_repository + rng, self.diffraction_simulator_settings, diffraction_repository, product_repository ) self.fourier_analyzer = FourierAnalyzer(product_repository) @@ -54,7 +54,7 @@ def __init__( self.probe_propagator = ProbePropagator(self.probe_propagator_settings, product_repository) self.probe_propagator_visualization_engine = VisualizationEngine(is_complex=False) - self.residual_analyzer = ResidualAnalyzer(product_repository, dataset) + self.residual_analyzer = ResidualAnalyzer(product_repository) self.residual_real_space_visualization_engine = VisualizationEngine(is_complex=False) self.residual_reciprocal_space_visualization_engine = VisualizationEngine(is_complex=False) diff --git a/src/ptychodus/model/analysis/diffraction.py b/src/ptychodus/model/analysis/diffraction.py index a5cebfc62..59d075740 100644 --- a/src/ptychodus/model/analysis/diffraction.py +++ b/src/ptychodus/model/analysis/diffraction.py @@ -5,7 +5,7 @@ from ptychodus.api.diffraction_gen import generate_diffraction_data -from ..diffraction import AssembledDiffractionDataset +from ..diffraction import DiffractionDatasetRepository from ..product import ProductRepository from .settings import DiffractionSimulatorSettings @@ -17,16 +17,17 @@ def __init__( self, rng: numpy.random.Generator, settings: DiffractionSimulatorSettings, - dataset: AssembledDiffractionDataset, + diffraction_repository: DiffractionDatasetRepository, repository: ProductRepository, ) -> None: self._rng = rng self._settings = settings - self._dataset = dataset + self._diffraction_repository = diffraction_repository self._repository = repository - def simulate(self, product_index: int) -> None: - product = self._repository[product_index].get_product() + def simulate(self, product_index: int) -> int: + product_item = self._repository[product_index] + product = product_item.get_product() rng = self._rng if self._settings.add_poisson_noise.get_value() else None logger.info('Computing diffraction data...') @@ -35,4 +36,7 @@ def simulate(self, product_index: int) -> None: toc = time.perf_counter() logger.info(f'Computed diffraction data in {toc - tic:.4f} seconds.') - self._dataset.set_assembled_patterns(data) + dataset = self._diffraction_repository.create_dataset(product_item.get_name()) + dataset_index = self._diffraction_repository.insert_dataset(dataset) + dataset.set_assembled_patterns(data) + return dataset_index diff --git a/src/ptychodus/model/analysis/residuals.py b/src/ptychodus/model/analysis/residuals.py index f863c1a8f..a16b68c57 100644 --- a/src/ptychodus/model/analysis/residuals.py +++ b/src/ptychodus/model/analysis/residuals.py @@ -8,7 +8,6 @@ compute_reconstruction_residuals, ) -from ..diffraction import AssembledDiffractionDataset from ..product import ProductRepository __all__ = ['ReconstructionResiduals', 'ResidualAnalyzer'] @@ -20,14 +19,20 @@ class ResidualAnalyzer: def __init__( self, repository: ProductRepository, - dataset: AssembledDiffractionDataset, ) -> None: self._repository = repository - self._dataset = dataset def analyze(self, product_index: int) -> ReconstructionResiduals: - product = self._repository[product_index].get_product() - recon_input = self._dataset.get_assembled_data().prepare_reconstruct_input(product) + item = self._repository[product_index] + dataset = item.get_dataset() + + if dataset is None: + raise RuntimeError( + f'Product "{item.get_name()}" has no associated diffraction dataset.' + ) + + product = item.get_product() + recon_input = dataset.get_assembled_data().prepare_reconstruct_input(product) logger.info('Computing reconstruction residuals...') tic = time.perf_counter() diff --git a/src/ptychodus/model/automation/core.py b/src/ptychodus/model/automation/core.py index 2203b08cc..0f2469022 100644 --- a/src/ptychodus/model/automation/core.py +++ b/src/ptychodus/model/automation/core.py @@ -2,7 +2,7 @@ from collections.abc import Iterator from ptychodus.api.observer import Observable, Observer -from ptychodus.api.plugins import Plugin, PluginChooser +from ptychodus.api.plugins import Plugin, PluginChooser, PluginChooserParameter from ptychodus.api.settings import SettingsRegistry from ptychodus.api.workflow import FileBasedWorkflow, WorkflowAPI @@ -27,7 +27,7 @@ def __init__( self._dataset_buffer = dataset_buffer self._workflow_chooser = workflow_chooser - workflow_chooser.synchronize_with_parameter(self._settings.workflow) + self.workflow_parameter = PluginChooserParameter(workflow_chooser, self._settings.workflow) watcher.add_observer(self) diff --git a/src/ptychodus/model/core.py b/src/ptychodus/model/core.py index d5bbede1d..338860465 100644 --- a/src/ptychodus/model/core.py +++ b/src/ptychodus/model/core.py @@ -30,11 +30,11 @@ from .genesis import GenesisCore from .globus import GlobusCore from .memory import MemoryPresenter -from .metadata import MetadataPresenter from .processing import ProcessingCore from .product import PositionsStreamingContext, ProductCore from .ptychi import PtyChiReconstructorLibrary from .ptychonn import PtychoNNReconstructorLibrary +from .ptycho_fm import PtychoFMReconstructorLibrary from .ptychopinn import PtychoPINNReconstructorLibrary from .ptychopinn_torch import PtychoPINNTorchReconstructorLibrary from .task_manager import TaskManager @@ -117,7 +117,6 @@ def __init__( self.settings_registry, self.diffraction_core.pattern_sizer, self.diffraction_core.diffraction_api, - self.diffraction_core.dataset, self.plugin_registry.probe_position_file_readers, self.plugin_registry.probe_position_file_writers, self.plugin_registry.fresnel_zone_plates, @@ -128,14 +127,8 @@ def __init__( self.plugin_registry.product_file_readers, self.plugin_registry.product_file_writers, self.settings_registry, + self._task_manager, ) - self.metadata_presenter = MetadataPresenter( - self.diffraction_core.detector_settings, - self.diffraction_core.diffraction_settings, - self.diffraction_core.dataset, - self.product_core.settings, - ) - self.pattern_visualization_engine = VisualizationEngine(is_complex=False) self.probe_visualization_engine = VisualizationEngine(is_complex=True) self.object_visualization_engine = VisualizationEngine(is_complex=True) @@ -154,22 +147,26 @@ def __init__( self.ptychopinn_torch_reconstructor_library = PtychoPINNTorchReconstructorLibrary( self.settings_registry, self.is_developer_mode_enabled ) + self.ptycho_fm_reconstructor_library = PtychoFMReconstructorLibrary( + self.settings_registry, self.is_developer_mode_enabled + ) self.processing_core = ProcessingCore( self._task_manager, self.settings_registry, - self.diffraction_core.diffraction_api, self.product_core.product_api, [ self.ptychi_reconstructor_library, self.ptychonn_reconstructor_library, self.ptychopinn_reconstructor_library, self.ptychopinn_torch_reconstructor_library, + self.ptycho_fm_reconstructor_library, ], ) self.fluorescence_core = FluorescenceCore( self._task_manager, self.settings_registry, self.product_core.product_api, + self.product_core.product_repository, self.plugin_registry.upscaling_strategies, self.plugin_registry.deconvolution_strategies, self.plugin_registry.fluorescence_file_readers, @@ -178,20 +175,18 @@ def __init__( self.analysis_core = AnalysisCore( self.rng, self.settings_registry, - self.diffraction_core.dataset, + self.diffraction_core.repository, self.product_core.product_repository, self.product_core.probe_positions_repository, ) self.globus_core = GlobusCore( self.settings_registry, - self.diffraction_core.diffraction_api, self.product_core.product_api, self.processing_core.processing_api, ) self.genesis_core = GenesisCore( self._task_manager, self.settings_registry, - self.diffraction_core.diffraction_api, self.product_core.product_api, self.processing_core.processing_api, ) @@ -262,7 +257,7 @@ def _batch_mode_reconstruct(self, input_directory: Path, output_directory: Path) diffraction_path = input_directory / StandardFileLayout.DIFFRACTION if diffraction_path.is_file(): - self.workflow_api.load_assembled_diffraction_data(diffraction_path) + diffraction_handle = self.workflow_api.load_assembled_diffraction_data(diffraction_path) else: logger.error('Diffraction data is not a file!') return -1 @@ -289,10 +284,17 @@ def _batch_mode_reconstruct(self, input_directory: Path, output_directory: Path) if product_out_path.is_file(): logger.warning('Output product file will be overwritten!') - input_product_api = self.workflow_api.load_product(product_in_path) - output_product_api = input_product_api.reconstruct_local( - output_product_file=product_out_path, block=True + input_product_api = self.workflow_api.load_product( + product_in_path, diffraction=diffraction_handle ) + try: + output_product_api = input_product_api.reconstruct_local( + output_product_file=product_out_path, + block=True, + ) + except Exception as exc: + logger.error(f'Reconstruction failed: {type(exc).__name__}: {exc}') + return 1 else: logger.error('Input product is not a file!') return -1 @@ -318,13 +320,27 @@ def _batch_mode_reconstruct(self, input_directory: Path, output_directory: Path) return 0 def _batch_mode_train(self, input_directory: Path, output_directory: Path) -> int: + diffraction_path = input_directory / StandardFileLayout.DIFFRACTION + + if diffraction_path.is_file(): + diffraction_handle = self.workflow_api.load_assembled_diffraction_data(diffraction_path) + else: + logger.error('Diffraction data is not a file!') + return -1 + product_in_path = input_directory / StandardFileLayout.PRODUCT_IN if product_in_path.is_file(): - input_product_api = self.workflow_api.load_product(product_in_path) - input_product_api.train_reconstructor_local( - input_directory, output_directory, block=True + input_product_api = self.workflow_api.load_product( + product_in_path, diffraction=diffraction_handle ) + try: + input_product_api.train_reconstructor_local( + input_directory, output_directory, block=True + ) + except Exception as exc: + logger.error(f'Training failed: {type(exc).__name__}: {exc}') + return 1 return 0 else: logger.error('Input product is not a file!') diff --git a/src/ptychodus/model/diffraction/__init__.py b/src/ptychodus/model/diffraction/__init__.py index 45136f1e7..bc137d968 100644 --- a/src/ptychodus/model/diffraction/__init__.py +++ b/src/ptychodus/model/diffraction/__init__.py @@ -5,19 +5,20 @@ AssembledDiffractionDataset, DiffractionDatasetObserver, ) -from .detector import Detector from .monitor import DiffractionTaskMonitor +from .repository import DiffractionDatasetRepository, DiffractionDatasetRepositoryObserver from .settings import DetectorSettings, DiffractionSettings from .sizer import PatternSizer __all__ = [ 'AssembledDiffractionArray', 'AssembledDiffractionDataset', - 'Detector', 'DetectorSettings', 'DiffractionAPI', 'DiffractionCore', 'DiffractionDatasetObserver', + 'DiffractionDatasetRepository', + 'DiffractionDatasetRepositoryObserver', 'DiffractionSettings', 'DiffractionTaskMonitor', 'PatternSizer', diff --git a/src/ptychodus/model/diffraction/_loader.py b/src/ptychodus/model/diffraction/_loader.py index 280472075..4e84a6ff6 100644 --- a/src/ptychodus/model/diffraction/_loader.py +++ b/src/ptychodus/model/diffraction/_loader.py @@ -9,13 +9,14 @@ BadPixels, DiffractionArray, SimpleDiffractionArray, + zero_bad_pixels, ) +from ptychodus.api.diffraction_prep import DiffractionPrepPipeline from ptychodus.api.geometry import PixelGeometry from ptychodus.api.io import AssembledDiffractionData from ..task_manager import BackgroundTask, ForegroundTask, ForegroundTaskManager from ..task_monitor import TaskProgressMonitor -from .processor import DiffractionPatternProcessor logger = logging.getLogger(__name__) @@ -43,44 +44,50 @@ def __init__( array_index: int, array: DiffractionArray, pixel_geometry: PixelGeometry, - bad_pixels: BadPixels, - processor: DiffractionPatternProcessor | None, + *, + raw_bad_pixels: BadPixels, + processed_bad_pixels: BadPixels, + pipeline: DiffractionPrepPipeline | None, assembler: ArrayAssembler, ) -> None: super().__init__() self._array_index = array_index self._array = array self._pixel_geometry = pixel_geometry - self._bad_pixels = bad_pixels - self._processor = processor + self._raw_bad_pixels = raw_bad_pixels + self._processed_bad_pixels = processed_bad_pixels + self._pipeline = pipeline self._assembler = assembler def __call__(self) -> ForegroundTask | None: label = self._array.get_label() try: - loaded_array = SimpleDiffractionArray( - label, - self._array.get_indexes(), - self._array.get_patterns(), - ) + raw_patterns = self._array.get_patterns() except FileNotFoundError: logger.warning(f'File not found for "{label}"!') - else: - processed_array = ( - loaded_array if self._processor is None else self._processor(loaded_array) - ) - data = AssembledDiffractionData( - indexes=processed_array.get_indexes(), - patterns=processed_array.get_patterns(), - pixel_geometry=self._pixel_geometry, - bad_pixels=self._bad_pixels, - ) - self._assembler.assemble_array( - self._array_index, - label, - data, - ) + return None + + # Zero bad pixels in raw detector coords before crop/bin/pad/flip/transpose so + # saturated pixel values can't leak into neighboring bins downstream. + repaired_patterns = zero_bad_pixels(raw_patterns, self._raw_bad_pixels) + loaded_array = SimpleDiffractionArray( + label, + self._array.get_indexes(), + repaired_patterns, + ) + processed_array = loaded_array if self._pipeline is None else self._pipeline(loaded_array) + data = AssembledDiffractionData( + indexes=processed_array.get_indexes(), + patterns=processed_array.get_patterns(), + pixel_geometry=self._pixel_geometry, + bad_pixels=self._processed_bad_pixels, + ) + self._assembler.assemble_array( + self._array_index, + label, + data, + ) return None @@ -100,6 +107,7 @@ def __init__( self._task_monitor = task_monitor self._process_patterns = False self._finished_event = threading.Event() + self._error: BaseException | None = None def enable_pattern_processing(self) -> None: self._process_patterns = True @@ -107,6 +115,10 @@ def enable_pattern_processing(self) -> None: def get_finished_event(self) -> threading.Event: return self._finished_event + def get_error(self) -> BaseException | None: + """First per-array exception observed during __call__, or None on clean load.""" + return self._error + def __call__(self) -> None: try: with self._task_monitor as monitor: @@ -134,6 +146,8 @@ def __call__(self) -> None: pass except Exception as ex: logger.warning(ex) + if self._error is None: + self._error = ex else: if task is not None: self._foreground_task_manager.put_foreground_task(task) @@ -148,5 +162,9 @@ def __call__(self) -> None: f'{num_cancelled} pending arrays. In-flight reads will finish.' ) cancelled = True + except BaseException as ex: + if self._error is None: + self._error = ex + raise finally: self._finished_event.set() diff --git a/src/ptychodus/model/diffraction/api.py b/src/ptychodus/model/diffraction/api.py index f8808ec53..42b08f6d3 100644 --- a/src/ptychodus/model/diffraction/api.py +++ b/src/ptychodus/model/diffraction/api.py @@ -1,3 +1,4 @@ +from collections.abc import Iterator from pathlib import Path import logging @@ -12,13 +13,13 @@ DiffractionArray, SimpleDiffractionDataset, ) -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.reconstructor import AssembledDiffractionData from ptychodus.api.tree import SimpleTreeNode from .dataset import AssembledDiffractionDataset -from .detector import Detector -from .settings import DetectorSettings, DiffractionSettings +from .repository import DiffractionDatasetRepository +from .settings import DiffractionSettings logger = logging.getLogger(__name__) @@ -45,53 +46,42 @@ class DiffractionAPI: def __init__( self, diffraction_settings: DiffractionSettings, - detector_settings: DetectorSettings, - detector: Detector, - dataset: AssembledDiffractionDataset, + repository: DiffractionDatasetRepository, bad_pixels_file_reader_chooser: PluginChooser[BadPixelsFileReader], file_reader_chooser: PluginChooser[DiffractionFileReader], file_writer_chooser: PluginChooser[DiffractionFileWriter], + bad_pixels_file_reader_parameter: PluginChooserParameter[BadPixelsFileReader], + file_reader_parameter: PluginChooserParameter[DiffractionFileReader], ) -> None: super().__init__() self._diffraction_settings = diffraction_settings - self._detector_settings = detector_settings - self._detector = detector - self._dataset = dataset + self._repository = repository self._bad_pixels_file_reader_chooser = bad_pixels_file_reader_chooser self._file_reader_chooser = file_reader_chooser self._file_writer_chooser = file_writer_chooser + self._bad_pixels_file_reader_parameter = bad_pixels_file_reader_parameter + self._file_reader_parameter = file_reader_parameter + + def get_repository(self) -> DiffractionDatasetRepository: + return self._repository def create_streaming_context(self, metadata: DiffractionMetadata) -> PatternsStreamingContext: - return PatternsStreamingContext(self._dataset, metadata) + dataset = self._repository.create_dataset('stream') + self._repository.insert_dataset(dataset) + return PatternsStreamingContext(dataset, metadata) def get_bad_pixels_file_reader_chooser(self) -> PluginChooser[BadPixelsFileReader]: return self._bad_pixels_file_reader_chooser - def open_bad_pixels(self, file_path: Path, *, file_type: str | None = None) -> None: - if file_path.is_file(): - if file_type is not None: - self._bad_pixels_file_reader_chooser.set_current_plugin(file_type) - - plugin = self._bad_pixels_file_reader_chooser.get_current_plugin() - logger.debug(f'Reading "{file_path}" as "{plugin.simple_name}"') - - try: - bad_pixels = plugin.strategy.read(file_path) - except Exception as exc: - raise RuntimeError(f'Failed to read "{file_path}"') from exc - else: - self._detector.set_bad_pixels(bad_pixels) - self._detector_settings.bad_pixels_file_path.set_value(file_path) - else: - logger.warning(f'Refusing to read invalid file path {file_path}') - - def clear_bad_pixels(self) -> None: - self._detector.set_bad_pixels(None) - self._detector_settings.bad_pixels_file_path.set_value(Path()) - def get_file_reader_chooser(self) -> PluginChooser[DiffractionFileReader]: return self._file_reader_chooser + def get_bad_pixels_file_reader_parameter(self) -> PluginChooserParameter[BadPixelsFileReader]: + return self._bad_pixels_file_reader_parameter + + def get_file_reader_parameter(self) -> PluginChooserParameter[DiffractionFileReader]: + return self._file_reader_parameter + def open_patterns( self, file_path: Path, @@ -99,10 +89,15 @@ def open_patterns( file_type: str | None = None, crop_center: CropCenter | None = None, crop_extent: ImageExtent | None = None, - detector_extent: ImageExtent | None = None, + bad_pixels_file_path: Path | None = None, + bad_pixels_file_type: str | None = None, process_patterns: bool = True, block: bool = False, ) -> int: + if not file_path.is_file(): + logger.warning(f'Refusing to read invalid file path {file_path}') + return -1 + if crop_center is not None: self._diffraction_settings.crop_center_x_px.set_value(crop_center.position_x_px) self._diffraction_settings.crop_center_y_px.set_value(crop_center.position_y_px) @@ -111,54 +106,123 @@ def open_patterns( self._diffraction_settings.crop_width_px.set_value(crop_extent.width_px) self._diffraction_settings.crop_height_px.set_value(crop_extent.height_px) - if detector_extent is not None: - self._detector_settings.width_px.set_value(detector_extent.width_px) - self._detector_settings.height_px.set_value(detector_extent.height_px) + if file_type is not None: + self._file_reader_chooser.set_current_plugin(file_type) - if file_path.is_file(): - if file_type is not None: - self._file_reader_chooser.set_current_plugin(file_type) + plugin = self._file_reader_chooser.get_current_plugin() + logger.debug(f'Reading "{file_path}" as "{plugin.simple_name}"') - plugin = self._file_reader_chooser.get_current_plugin() - logger.debug(f'Reading "{file_path}" as "{plugin.simple_name}"') + try: + source_dataset = plugin.strategy.read(file_path) + except Exception as exc: + raise RuntimeError(f'Failed to read "{file_path}"') from exc - try: - dataset = plugin.strategy.read(file_path) - except Exception as exc: - raise RuntimeError(f'Failed to read "{file_path}"') from exc - else: - self._dataset.reload(dataset) + dataset = self._repository.create_dataset(file_path.stem) + dataset_index = self._repository.insert_dataset(dataset) + dataset.reload(source_dataset) - if block: - self._dataset.load_all_arrays(process_patterns=process_patterns, block=True) + if bad_pixels_file_path is not None: + self._apply_bad_pixels_from_file(dataset, bad_pixels_file_path, bad_pixels_file_type) - return 0 - else: - logger.warning(f'Refusing to read invalid file path {file_path}') + if block: + dataset.load_all_arrays(process_patterns=process_patterns, block=True) - return -1 + return dataset_index - def load_all_arrays(self, *, process_patterns: bool = True, block: bool = False) -> None: - self._dataset.load_all_arrays(process_patterns=process_patterns, block=block) + def _apply_bad_pixels_from_file( + self, + dataset: AssembledDiffractionDataset, + file_path: Path, + file_type: str | None, + ) -> None: + if not file_path.is_file(): + logger.warning(f'Refusing to read invalid bad pixels file path {file_path}') + return + + if file_type is not None: + self._bad_pixels_file_reader_chooser.set_current_plugin(file_type) + bad_pixels_plugin = self._bad_pixels_file_reader_chooser.get_current_plugin() + logger.debug(f'Reading "{file_path}" as "{bad_pixels_plugin.simple_name}"') + try: + bad_pixels = bad_pixels_plugin.strategy.read(file_path) + except Exception: + logger.warning(f'Failed to load bad pixels from "{file_path}"') + return + + try: + dataset.set_bad_pixels(bad_pixels) + except ValueError as exc: + logger.warning(f'Ignoring bad pixels from "{file_path}": {exc}') + + def apply_bad_pixels( + self, + dataset_index: int, + file_path: Path, + file_type: str | None = None, + ) -> None: + """Load a bad-pixel mask from disk and apply it to an already-open dataset.""" + try: + dataset = self._repository[dataset_index] + except IndexError: + logger.warning(f'Cannot apply bad pixels: no dataset at index {dataset_index}') + return + self._apply_bad_pixels_from_file(dataset, file_path, file_type) + + def load_all_arrays( + self, *, dataset_index: int, process_patterns: bool = True, block: bool = False + ) -> None: + try: + dataset = self._repository[dataset_index] + except IndexError: + logger.warning(f'Cannot load arrays: no dataset at index {dataset_index}') + else: + dataset.load_all_arrays(process_patterns=process_patterns, block=block) + + def close_patterns(self, dataset_index: int) -> None: + self._repository.remove_dataset(dataset_index) - def close_patterns(self) -> None: - self._dataset.clear() + def close_all_patterns(self) -> None: + self._repository.clear() def get_file_writer_chooser(self) -> PluginChooser[DiffractionFileWriter]: return self._file_writer_chooser - def save_patterns(self, file_path: Path, file_type: str) -> None: - self._file_writer_chooser.set_current_plugin(file_type) - file_type = self._file_writer_chooser.get_current_plugin().simple_name - logger.debug(f'Writing "{file_path}" as "{file_type}"') - writer = self._file_writer_chooser.get_current_plugin().strategy - writer.write(file_path, self._dataset) + def get_save_file_filters(self) -> Iterator[str]: + for plugin in self._file_writer_chooser: + yield plugin.display_name - def get_assembled_data(self) -> AssembledDiffractionData: - return self._dataset.get_assembled_data() + def get_save_file_filter(self) -> str: + return self._file_writer_chooser.get_current_plugin().display_name - def import_assembled_patterns(self, file_path: Path) -> None: - self._dataset.import_assembled_patterns(file_path) + def save_patterns(self, file_path: Path, file_type: str, *, dataset_index: int) -> None: + try: + dataset = self._repository[dataset_index] + except IndexError: + logger.warning(f'Cannot save patterns: no dataset at index {dataset_index}') + return - def export_assembled_patterns(self, file_path: Path) -> None: - self._dataset.export_assembled_patterns(file_path) + self._file_writer_chooser.set_current_plugin(file_type) + plugin = self._file_writer_chooser.get_current_plugin() + logger.debug(f'Writing "{file_path}" as "{plugin.simple_name}"') + plugin.strategy.write(file_path, dataset) + + def get_assembled_data(self, dataset_index: int) -> AssembledDiffractionData: + return self._repository[dataset_index].get_assembled_data() + + def import_assembled_patterns(self, file_path: Path) -> int: + if not file_path.is_file(): + logger.warning(f'Refusing to read invalid file path {file_path}') + return -1 + + dataset = self._repository.create_dataset(file_path.stem) + dataset_index = self._repository.insert_dataset(dataset) + dataset.import_assembled_patterns(file_path) + return dataset_index + + def export_assembled_patterns(self, file_path: Path, *, dataset_index: int) -> None: + try: + dataset = self._repository[dataset_index] + except IndexError: + logger.warning(f'Cannot export patterns: no dataset at index {dataset_index}') + return + dataset.export_assembled_patterns(file_path) diff --git a/src/ptychodus/model/diffraction/core.py b/src/ptychodus/model/diffraction/core.py index 4073cb96a..70b0f42ad 100644 --- a/src/ptychodus/model/diffraction/core.py +++ b/src/ptychodus/model/diffraction/core.py @@ -6,14 +6,13 @@ DiffractionFileReader, DiffractionFileWriter, ) -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.settings import SettingsRegistry from ..task_manager import TaskManager from .api import DiffractionAPI -from .dataset import AssembledDiffractionDataset -from .detector import Detector from .monitor import DiffractionTaskMonitor +from .repository import DiffractionDatasetRepository, build_default_factory from .settings import DetectorSettings, DiffractionSettings from .sizer import PatternSizer @@ -33,30 +32,38 @@ def __init__( super().__init__() self.detector_settings = DetectorSettings(settings_registry) self.diffraction_settings = DiffractionSettings(settings_registry) - self.pattern_sizer = PatternSizer(self.detector_settings, self.diffraction_settings) - self.detector = Detector(self.detector_settings) + self.pattern_sizer = PatternSizer(self.diffraction_settings) self.task_monitor = DiffractionTaskMonitor(task_manager) - self.dataset = AssembledDiffractionDataset( - self.diffraction_settings, - self.pattern_sizer, - self.detector, - task_manager, - self.task_monitor, + self.repository = DiffractionDatasetRepository( + factory=build_default_factory( + self.diffraction_settings, + self.pattern_sizer, + self.detector_settings, + task_manager, + self.task_monitor, + ) + ) + # Display-name views of each reader chooser, bound to the settings parameter + # that persists the selection. These are what the GUI binds combo boxes to. + self.bad_pixels_file_reader_parameter = PluginChooserParameter( + bad_pixels_file_reader_chooser, self.detector_settings.bad_pixels_file_type + ) + self.file_reader_parameter = PluginChooserParameter( + file_reader_chooser, self.diffraction_settings.file_type ) + self.diffraction_api = DiffractionAPI( self.diffraction_settings, - self.detector_settings, - self.detector, - self.dataset, + self.repository, bad_pixels_file_reader_chooser, file_reader_chooser, file_writer_chooser, + self.bad_pixels_file_reader_parameter, + self.file_reader_parameter, ) - bad_pixels_file_reader_chooser.synchronize_with_parameter( - self.detector_settings.bad_pixels_file_type - ) - file_reader_chooser.synchronize_with_parameter(self.diffraction_settings.file_type) + # Deliberately unbound: the writer shares file_type with the reader above, so + # binding both would make them fight over the same parameter. file_writer_chooser.set_current_plugin(self.diffraction_settings.file_type.get_value()) self._reinit_observable = reinit_observable @@ -64,12 +71,12 @@ def __init__( def _update(self, observable: Observable) -> None: if observable is self._reinit_observable: - self.diffraction_api.open_bad_pixels( - file_path=self.detector_settings.bad_pixels_file_path.get_value(), - file_type=self.detector_settings.bad_pixels_file_type.get_value(), - ) - self.diffraction_api.open_patterns( + dataset_index = self.diffraction_api.open_patterns( file_path=self.diffraction_settings.file_path.get_value(), file_type=self.diffraction_settings.file_type.get_value(), + bad_pixels_file_path=self.detector_settings.bad_pixels_file_path.get_value(), + bad_pixels_file_type=self.detector_settings.bad_pixels_file_type.get_value(), ) - self.diffraction_api.load_all_arrays() + + if dataset_index >= 0: + self.diffraction_api.load_all_arrays(dataset_index=dataset_index) diff --git a/src/ptychodus/model/diffraction/dataset.py b/src/ptychodus/model/diffraction/dataset.py index 020fef356..d63e2e6ac 100644 --- a/src/ptychodus/model/diffraction/dataset.py +++ b/src/ptychodus/model/diffraction/dataset.py @@ -6,6 +6,7 @@ from typing import IO, overload import logging import tempfile +import threading import numpy @@ -26,9 +27,8 @@ from ..task_manager import BackgroundTask, TaskManager from ._loader import ArrayAssembler, LoadAllArrays, LoadArray -from .detector import Detector from .monitor import DiffractionTaskMonitor -from .settings import DiffractionSettings +from .settings import DetectorSettings, DiffractionSettings from .sizer import PatternSizer logger = logging.getLogger(__name__) @@ -47,6 +47,10 @@ def handle_array_changed(self, index: int) -> None: def handle_dataset_reloaded(self) -> None: pass + @abstractmethod + def handle_pixel_geometry_changed(self) -> None: + pass + class AssembledDiffractionArray(DiffractionArray): def __init__( @@ -101,14 +105,17 @@ def __init__( self, settings: DiffractionSettings, sizer: PatternSizer, - detector: Detector, + detector_settings: DetectorSettings, task_manager: TaskManager, task_monitor: DiffractionTaskMonitor, + *, + name: str = 'default', ) -> None: super().__init__() + self._name = name self._settings = settings self._sizer = sizer - self._detector = detector + self._detector_settings = detector_settings self._task_manager = task_manager self._task_monitor = task_monitor self._observer_list: list[DiffractionDatasetObserver] = [] @@ -118,8 +125,62 @@ def __init__( self._array_list: list[AssembledDiffractionArray] = list() self._array_counter = 0 self._array_loader: LoadAllArrays | None = None + # Retained after dispatch so callers can detect \"still loading\" and read + # any error the loader stored (see is_load_in_progress / get_last_load_error). + self._last_array_loader: LoadAllArrays | None = None self._scratch_tempfile: IO[bytes] | None = None + # Raw (pre-processing) bad-pixel mask; starts as an empty (0, 0) placeholder + # and is always overwritten by reload() or by load_all_arrays() before any + # array is processed, so the shape here only matters when nothing is loaded. + self._bad_pixels = self._create_default_bad_pixels() + + # Per-dataset override for the raw detector pixel geometry. When set, takes + # priority over metadata and DetectorSettings in get_raw_pixel_geometry(). + # Cleared by clear()/reload() so freshly-read metadata is the new source of truth. + self._pixel_geometry_override: PixelGeometry | None = None + + def _create_default_bad_pixels(self) -> BadPixels: + extent = self._dataset.get_metadata().detector_extent + return numpy.zeros((extent.height_px, extent.width_px), dtype=numpy.bool_) + + def get_name(self) -> str: + return self._name + + def set_name(self, name: str) -> None: + """Set the dataset's display name. Callers must ensure uniqueness themselves + (typically by routing the candidate through DiffractionDatasetRepository.create_unique_name). + """ + self._name = name + + def sync_pixel_geometry_to_settings(self) -> None: + """Promote this dataset's effective raw pixel geometry to the global fallback. + + Writes the current raw geometry (override > metadata > current fallback) into + DetectorSettings.pixel_width_m / pixel_height_m so freshly loaded datasets that + lack pixel metadata pick it up. Leaves this dataset's override in place. + """ + geometry = self.get_raw_pixel_geometry() + self._detector_settings.pixel_width_m.set_value(geometry.width_m) + self._detector_settings.pixel_height_m.set_value(geometry.height_m) + + def set_bad_pixels(self, bad_pixels: BadPixels) -> None: + if bad_pixels.ndim != 2: + raise ValueError(f'Bad pixels array must be 2D, got {bad_pixels.ndim}D.') + + extent = self._dataset.get_metadata().detector_extent + + if bad_pixels.shape != extent.get_shape(): + raise ValueError( + f'Bad pixels shape {bad_pixels.shape} does not match ' + f'loaded detector extent {extent.get_shape()}.' + ) + + self._bad_pixels = bad_pixels + + def reset_bad_pixels(self) -> None: + self._bad_pixels = self._create_default_bad_pixels() + def add_observer(self, observer: DiffractionDatasetObserver) -> None: if observer not in self._observer_list: self._observer_list.append(observer) @@ -136,15 +197,74 @@ def get_metadata(self) -> DiffractionMetadata: def get_layout(self) -> SimpleTreeNode: return self._dataset.get_layout() - def _get_pixel_geometry(self) -> PixelGeometry: - return self._detector.get_pixel_geometry() + def get_raw_pixel_geometry(self) -> PixelGeometry: + """Resolve the raw (pre-processing) detector pixel geometry for this dataset. + + Priority: user override > metadata > global DetectorSettings fallback. + """ + if self._pixel_geometry_override is not None: + return self._pixel_geometry_override + + metadata_geometry = self._dataset.get_metadata().detector_pixel_geometry + if metadata_geometry is not None: + return metadata_geometry + + return PixelGeometry( + width_m=self._detector_settings.pixel_width_m.get_value(), + height_m=self._detector_settings.pixel_height_m.get_value(), + ) + + def set_pixel_geometry_override(self, geometry: PixelGeometry | None) -> None: + """Set (or clear when None) the per-dataset raw pixel geometry override. - def get_bad_pixels(self) -> BadPixels | None: - return self._dataset.get_bad_pixels() + Also mutates the assembled-data snapshot so consumers reading + AssembledDiffractionData.get_pixel_geometry() see the update, and notifies + observers so downstream views (tree columns, bound products) can refresh. + """ + self._pixel_geometry_override = geometry + self._data.set_pixel_geometry(self.get_raw_pixel_geometry()) + + for observer in self._observer_list: + observer.handle_pixel_geometry_changed() + + def get_bad_pixels(self) -> BadPixels: + return self._bad_pixels def get_assembled_data(self) -> AssembledDiffractionData: return self._data + def is_load_in_progress(self) -> bool: + """True while a LoadAllArrays task is queued but has not finished. + + Note: this does not track per-array append_array() streams, which are + used only by the pvapy streaming path; products are created there only + after the stream stops, so the streaming case doesn't rely on this. + """ + if self._array_loader is not None: + return True + + loader = self._last_array_loader + return loader is not None and not loader.get_finished_event().is_set() + + def get_last_load_error(self) -> BaseException | None: + """First exception raised by the most recent LoadAllArrays run, or None.""" + loader = self._last_array_loader + return loader.get_error() if loader is not None else None + + def get_last_load_finished_event(self) -> threading.Event | None: + """The finished_event of the most recent LoadAllArrays task, if any.""" + loader = self._last_array_loader + return loader.get_finished_event() if loader is not None else None + + def get_average_pattern(self) -> DiffractionPattern | None: + if not self._array_list: + return None + weights = numpy.array( + [array.get_num_patterns() for array in self._array_list], dtype=numpy.float64 + ) + averages = numpy.stack([array.get_average_pattern() for array in self._array_list]) + return numpy.average(averages, axis=0, weights=weights) + @overload def __getitem__(self, index: int) -> AssembledDiffractionArray: ... @@ -164,22 +284,22 @@ def create_array_loader( ) -> BackgroundTask: """Build a loader task for one array. Loaders are assigned a monotonic array_index; arrays may complete out of order and are sorted on insertion via bisect.""" - bad_pixels = self._dataset.get_bad_pixels() - - if bad_pixels is None: - raise RuntimeError('Cannot load array without bad pixel map!') - array_index = self._array_counter self._array_counter += 1 - processor = self._sizer.get_processor() + detector_extent = self._dataset.get_metadata().detector_extent + pipeline = self._sizer.get_prep_pipeline(detector_extent) if process_patterns else None + processed_bad_pixels = ( + pipeline.apply_to_mask(self._bad_pixels) if pipeline is not None else self._bad_pixels + ) return LoadArray( array_index, array, - self._get_pixel_geometry(), - bad_pixels, - processor if process_patterns else None, - self, + self.get_raw_pixel_geometry(), + raw_bad_pixels=self._bad_pixels, + processed_bad_pixels=processed_bad_pixels, + pipeline=pipeline, + assembler=self, ) def append_array(self, array: DiffractionArray, *, process_patterns: bool = True) -> None: @@ -216,6 +336,9 @@ def clear(self) -> None: self._array_list.clear() self._array_counter = 0 self._array_loader = None + self._last_array_loader = None + self._bad_pixels = self._create_default_bad_pixels() + self._pixel_geometry_override = None if self._scratch_tempfile is not None: self._scratch_tempfile.close() @@ -226,7 +349,9 @@ def clear(self) -> None: def reload(self, dataset: DiffractionDataset) -> None: self.clear() - self._dataset = SimpleDiffractionDataset(dataset.get_metadata(), dataset.get_layout(), []) + metadata = dataset.get_metadata() + self._dataset = SimpleDiffractionDataset(metadata, dataset.get_layout(), []) + self._bad_pixels = dataset.get_bad_pixels() self._array_loader = LoadAllArrays(dataset, self, self._task_manager, self._task_monitor) for observer in self._observer_list: @@ -239,20 +364,13 @@ def load_all_arrays(self, *, process_patterns: bool, block: bool) -> None: metadata = self._dataset.get_metadata() - if metadata.detector_extent is not None: - self._detector.set_extent(metadata.detector_extent) - - bad_pixels = self._detector.get_bad_pixels() + bad_pixels = self._bad_pixels if process_patterns: - processor = self._sizer.get_processor() - bad_pixels = processor.process_bad_pixels(bad_pixels) + pipeline = self._sizer.get_prep_pipeline(metadata.detector_extent) + bad_pixels = pipeline.apply_to_mask(bad_pixels) self._array_loader.enable_pattern_processing() - self._dataset = SimpleDiffractionDataset( - metadata, self._dataset.get_layout(), [], bad_pixels - ) - num_patterns_total = sum(metadata.num_patterns_per_array) indexes = -numpy.ones(num_patterns_total, dtype=int) @@ -276,15 +394,17 @@ def load_all_arrays(self, *, process_patterns: bool, block: bool) -> None: logger.debug(f'{patterns.nbytes / BYTES_PER_MEGABYTE:.2f}MB allocated for patterns') self._data = AssembledDiffractionData( - indexes, patterns, self._get_pixel_geometry(), bad_pixels + indexes, patterns, self.get_raw_pixel_geometry(), bad_pixels ) for observer in self._observer_list: observer.handle_dataset_reloaded() # load all arrays in background - finished_event = self._array_loader.get_finished_event() - self._task_manager.put_background_task(self._array_loader) + loader = self._array_loader + finished_event = loader.get_finished_event() + self._last_array_loader = loader + self._task_manager.put_background_task(loader) self._array_loader = None if block: @@ -306,9 +426,8 @@ def _generate_dataset_for_assembled_data(self, file_path: Path | None = None) -> label='In-Memory' if file_path is None else file_path.stem, data=self._data, ) - bad_pixels = self._data.get_bad_pixels() - - self._dataset = SimpleDiffractionDataset(metadata, contents_tree, [], bad_pixels) + self._bad_pixels = self._data.get_bad_pixels() + self._dataset = SimpleDiffractionDataset(metadata, contents_tree, [], self._bad_pixels) self._array_list = [array] self._array_counter = 1 diff --git a/src/ptychodus/model/diffraction/detector.py b/src/ptychodus/model/diffraction/detector.py deleted file mode 100644 index 43f43baba..000000000 --- a/src/ptychodus/model/diffraction/detector.py +++ /dev/null @@ -1,52 +0,0 @@ -import logging - -import numpy - -from ptychodus.api.diffraction import BadPixels -from ptychodus.api.geometry import ImageExtent, PixelGeometry -from ptychodus.api.observer import Observable, Observer -from .settings import DetectorSettings - -logger = logging.getLogger(__name__) - - -class Detector(Observable, Observer): - def __init__(self, settings: DetectorSettings) -> None: - super().__init__() - self._settings = settings - self._bad_pixels: BadPixels | None = None - - settings.add_observer(self) - - def set_extent(self, extent: ImageExtent) -> None: - logger.debug(f'Detector {extent=}') - self._settings.height_px.set_value(extent.height_px) - self._settings.width_px.set_value(extent.width_px) - - def get_pixel_geometry(self) -> PixelGeometry: - return PixelGeometry( - width_m=self._settings.pixel_width_m.get_value(), - height_m=self._settings.pixel_height_m.get_value(), - ) - - def set_bad_pixels(self, bad_pixels: BadPixels | None) -> None: - if bad_pixels is not None and bad_pixels.ndim != 2: - raise ValueError(f'Bad pixels array must be 2D, got {bad_pixels.ndim}D.') - - self._bad_pixels = bad_pixels - self.notify_observers() - - def get_bad_pixels(self) -> BadPixels: - if self._bad_pixels is None: - detector_height_px = self._settings.height_px.get_value() - detector_width_px = self._settings.width_px.get_value() - return numpy.full((detector_height_px, detector_width_px), False) - - return self._bad_pixels - - def get_num_bad_pixels(self) -> int: - return 0 if self._bad_pixels is None else int(numpy.count_nonzero(self._bad_pixels)) - - def _update(self, observable: Observable) -> None: - if observable is self._settings: - self.notify_observers() diff --git a/src/ptychodus/model/diffraction/processor.py b/src/ptychodus/model/diffraction/processor.py deleted file mode 100644 index 321fafeb0..000000000 --- a/src/ptychodus/model/diffraction/processor.py +++ /dev/null @@ -1,131 +0,0 @@ -from __future__ import annotations -from dataclasses import dataclass - -import numpy - -from ptychodus.api.geometry import ImageExtent -from ptychodus.api.diffraction import ( - BadPixels, - CropCenter, - DiffractionArray, - DiffractionPatterns, - SimpleDiffractionArray, -) - - -@dataclass(frozen=True) -class DiffractionPatternFilterValues: - lower_bound: int | None - upper_bound: int | None - - def apply(self, data: DiffractionPatterns) -> DiffractionPatterns: - if self.lower_bound is None and self.upper_bound is None: - return data - - out = data.copy() - - if self.lower_bound is not None: - out[out < self.lower_bound] = 0 - - if self.upper_bound is not None: - out[out >= self.upper_bound] = 0 - - return out - - -class DiffractionPatternCrop: - def __init__(self, center: CropCenter, extent: ImageExtent) -> None: - center_x = center.position_x_px - radius_x = extent.width_px // 2 - self.slice_x = slice(center_x - radius_x, center_x + radius_x) - - center_y = center.position_y_px - radius_y = extent.height_px // 2 - self.slice_y = slice(center_y - radius_y, center_y + radius_y) - - def apply(self, data: numpy.ndarray, *, is_mask: bool = False) -> numpy.ndarray: - leading = (slice(None),) * (data.ndim - 2) - return data[(*leading, self.slice_y, self.slice_x)] - - -@dataclass(frozen=True) -class DiffractionPatternBinning: - bin_size_x: int - bin_size_y: int - - def apply(self, data: numpy.ndarray, *, is_mask: bool = False) -> numpy.ndarray: - binned_height = data.shape[-2] // self.bin_size_y - binned_width = data.shape[-1] // self.bin_size_x - shape = data.shape[:-2] + (binned_height, self.bin_size_y, binned_width, self.bin_size_x) - reshaped = data.reshape(shape) - if is_mask: - return numpy.logical_and.reduce(reshaped, axis=(-3, -1), keepdims=False) - return numpy.sum(reshaped, axis=(-3, -1), keepdims=False) - - -@dataclass(frozen=True) -class DiffractionPatternPadding: - pad_x: int - pad_y: int - - def apply(self, data: numpy.ndarray, *, is_mask: bool = False) -> numpy.ndarray: - leading_pad = ((0, 0),) * (data.ndim - 2) - pad_width = (*leading_pad, (self.pad_y, self.pad_y), (self.pad_x, self.pad_x)) - fill = False if is_mask else 0 - return numpy.pad(data, pad_width, mode='constant', constant_values=fill) - - -@dataclass(frozen=True) -class DiffractionPatternProcessor: - crop: DiffractionPatternCrop | None - filter_values: DiffractionPatternFilterValues | None - binning: DiffractionPatternBinning | None - padding: DiffractionPatternPadding | None - hflip: bool - vflip: bool - transpose: bool - - def _apply_geometric(self, data: numpy.ndarray, *, is_mask: bool) -> numpy.ndarray: - """Run the geometric pipeline (crop → bin → pad → flips → transpose) on a 2-D mask - or a 3-D pattern stack. Order matches __call__; mirror changes in both paths.""" - if self.crop is not None: - data = self.crop.apply(data, is_mask=is_mask) - - if self.binning is not None: - data = self.binning.apply(data, is_mask=is_mask) - - if self.padding is not None: - data = self.padding.apply(data, is_mask=is_mask) - - if self.hflip: - data = numpy.flip(data, axis=-1) - - if self.vflip: - data = numpy.flip(data, axis=-2) - - if self.transpose: - axes = tuple(range(data.ndim - 2)) + (data.ndim - 1, data.ndim - 2) - data = numpy.transpose(data, axes=axes) - - return data - - def process_bad_pixels(self, bad_pixels: BadPixels) -> BadPixels: - if bad_pixels.ndim != 2: - raise ValueError(f'Invalid bad_pixel dimensions! (shape={bad_pixels.shape})') - - return self._apply_geometric(bad_pixels, is_mask=True) - - def __call__(self, array: DiffractionArray) -> DiffractionArray: - patterns = array.get_patterns() - - if patterns.ndim == 2: - patterns = patterns[numpy.newaxis, ...] - elif patterns.ndim != 3: - raise ValueError(f'Invalid diffraction pattern dimensions! (shape={patterns.shape})') - - if self.filter_values is not None: - patterns = self.filter_values.apply(patterns) - - patterns = self._apply_geometric(patterns, is_mask=False) - - return SimpleDiffractionArray(array.get_label(), array.get_indexes(), patterns) diff --git a/src/ptychodus/model/diffraction/repository.py b/src/ptychodus/model/diffraction/repository.py new file mode 100644 index 000000000..b4fd27bc3 --- /dev/null +++ b/src/ptychodus/model/diffraction/repository.py @@ -0,0 +1,129 @@ +from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence +from typing import overload +import logging +import sys + +from ptychodus.api.common import BYTES_PER_MEGABYTE + +from ..task_manager import TaskManager +from .dataset import AssembledDiffractionDataset +from .monitor import DiffractionTaskMonitor +from .settings import DetectorSettings, DiffractionSettings +from .sizer import PatternSizer + +logger = logging.getLogger(__name__) + + +def build_default_factory( + diffraction_settings: DiffractionSettings, + pattern_sizer: PatternSizer, + detector_settings: DetectorSettings, + task_manager: TaskManager, + task_monitor: DiffractionTaskMonitor, +) -> Callable[[str], AssembledDiffractionDataset]: + def _factory(name: str) -> AssembledDiffractionDataset: + return AssembledDiffractionDataset( + diffraction_settings, + pattern_sizer, + detector_settings, + task_manager, + task_monitor, + name=name, + ) + + return _factory + + +class DiffractionDatasetRepositoryObserver(ABC): + @abstractmethod + def handle_dataset_inserted(self, index: int, dataset: AssembledDiffractionDataset) -> None: + pass + + @abstractmethod + def handle_dataset_removed(self, index: int, dataset: AssembledDiffractionDataset) -> None: + pass + + +class DiffractionDatasetRepository(Sequence[AssembledDiffractionDataset]): + def __init__( + self, + factory: Callable[[str], AssembledDiffractionDataset] | None = None, + ) -> None: + super().__init__() + self._factory = factory + self._dataset_list: list[AssembledDiffractionDataset] = [] + self._observer_list: list[DiffractionDatasetRepositoryObserver] = [] + + def create_dataset(self, name: str) -> AssembledDiffractionDataset: + if self._factory is None: + raise RuntimeError( + 'DiffractionDatasetRepository was constructed without a factory; ' + 'cannot build new datasets.' + ) + unique_name = self.create_unique_name(name) + return self._factory(unique_name) + + @overload + def __getitem__(self, index: int) -> AssembledDiffractionDataset: ... + + @overload + def __getitem__(self, index: slice) -> Sequence[AssembledDiffractionDataset]: ... + + def __getitem__( + self, index: int | slice + ) -> AssembledDiffractionDataset | Sequence[AssembledDiffractionDataset]: + return self._dataset_list[index] + + def __len__(self) -> int: + return len(self._dataset_list) + + def add_observer(self, observer: DiffractionDatasetRepositoryObserver) -> None: + if observer not in self._observer_list: + self._observer_list.append(observer) + + def remove_observer(self, observer: DiffractionDatasetRepositoryObserver) -> None: + try: + self._observer_list.remove(observer) + except ValueError: + pass + + def create_unique_name(self, candidate_name: str) -> str: + reserved_names = {dataset.get_name() for dataset in self._dataset_list} + name = candidate_name or 'Unnamed' + match = 0 + + while name in reserved_names: + match += 1 + name = f'{candidate_name}-{match}' + + return name + + def insert_dataset(self, dataset: AssembledDiffractionDataset) -> int: + index = len(self._dataset_list) + self._dataset_list.append(dataset) + + for observer in self._observer_list: + observer.handle_dataset_inserted(index, dataset) + + return index + + def remove_dataset(self, index: int) -> None: + try: + dataset = self._dataset_list.pop(index) + except IndexError: + logger.debug(f'Failed to remove dataset {index}!') + return + + dataset.clear() + + for observer in self._observer_list: + observer.handle_dataset_removed(index, dataset) + + def clear(self) -> None: + for idx in range(len(self._dataset_list) - 1, -1, -1): + self.remove_dataset(idx) + + def get_info_text(self) -> str: + size_MB = sum(sys.getsizeof(ds) for ds in self._dataset_list) / BYTES_PER_MEGABYTE # noqa: N806 + return f'Total: {len(self)} [{size_MB:.2f}MB]' diff --git a/src/ptychodus/model/diffraction/settings.py b/src/ptychodus/model/diffraction/settings.py index 1414029d4..4a7501e99 100644 --- a/src/ptychodus/model/diffraction/settings.py +++ b/src/ptychodus/model/diffraction/settings.py @@ -11,11 +11,9 @@ def __init__(self, registry: SettingsRegistry) -> None: self._group = registry.create_group('Detector') self._group.add_observer(self) - self.width_px = self._group.create_integer_parameter('WidthInPixels', 1024, minimum=1) self.pixel_width_m = self._group.create_real_parameter( 'PixelWidthInMeters', 75e-6, minimum=0.0 ) - self.height_px = self._group.create_integer_parameter('HeightInPixels', 1024, minimum=1) self.pixel_height_m = self._group.create_real_parameter( 'PixelHeightInMeters', 75e-6, minimum=0.0 ) diff --git a/src/ptychodus/model/diffraction/sizer.py b/src/ptychodus/model/diffraction/sizer.py index 9c1c7a8f7..803d5c9e4 100644 --- a/src/ptychodus/model/diffraction/sizer.py +++ b/src/ptychodus/model/diffraction/sizer.py @@ -1,23 +1,25 @@ +from ptychodus.api.diffraction import CropCenter +from ptychodus.api.diffraction_prep import ( + BinningStep, + CropStep, + DiffractionPrepPipeline, + DiffractionPrepStepUnion, + FilterValuesStep, + HorizontalFlipStep, + PaddingStep, + TransposeStep, + VerticalFlipStep, +) from ptychodus.api.geometry import ImageExtent, Interval, PixelGeometry from ptychodus.api.observer import Observable, Observer -from ptychodus.api.parametric import BooleanParameter, IntegerParameter, RealParameter -from ptychodus.api.diffraction import CropCenter +from ptychodus.api.parametric import BooleanParameter, IntegerParameter -from .processor import ( - DiffractionPatternBinning, - DiffractionPatternCrop, - DiffractionPatternFilterValues, - DiffractionPatternPadding, - DiffractionPatternProcessor, -) -from .settings import DetectorSettings, DiffractionSettings +from .settings import DiffractionSettings class PatternAxisSizer(Observable, Observer): def __init__( self, - detector_size: IntegerParameter, - detector_pixel_size_m: RealParameter, crop_enabled: BooleanParameter, crop_size: IntegerParameter, crop_center: IntegerParameter, @@ -27,8 +29,6 @@ def __init__( pad_size: IntegerParameter, ) -> None: super().__init__() - self._detector_size = detector_size - self._detector_pixel_size_m = detector_pixel_size_m self._crop_enabled = crop_enabled self._crop_size = crop_size self._crop_center = crop_center @@ -37,8 +37,6 @@ def __init__( self._padding_enabled = padding_enabled self._pad_size = pad_size - detector_size.add_observer(self) - detector_pixel_size_m.add_observer(self) crop_enabled.add_observer(self) crop_size.add_observer(self) crop_center.add_observer(self) @@ -47,46 +45,38 @@ def __init__( padding_enabled.add_observer(self) pad_size.add_observer(self) - def get_detector_size(self) -> int: - return self._detector_size.get_value() - - def get_crop_size_limits(self) -> Interval[int]: - return Interval[int](1, self.get_detector_size()) - - def get_crop_size(self) -> int: + def get_crop_size(self, detector_size: int | None) -> int: if self._crop_enabled.get_value(): - limits = self.get_crop_size_limits() - return limits.clamp(self._crop_size.get_value()) - - return self.get_detector_size() - - def get_safe_crop_center(self) -> int: - """Crop center clamped so the configured crop window fits inside the detector.""" - xmin = (self.get_crop_size() + 1) // 2 - xmax = self.get_detector_size() - 1 - xmin - limits = Interval[int](xmin, xmax) - return limits.clamp(self._crop_center.get_value()) - - def get_crop_center_limits(self) -> Interval[int]: - return Interval[int](1, self.get_detector_size()) - - def get_crop_center(self) -> int: - limits = self.get_crop_center_limits() - return limits.clamp(self._crop_center.get_value()) - - def get_bin_size_limits(self) -> Interval[int]: - return Interval[int](1, self.get_crop_size()) - - def get_bin_size(self) -> int: + requested = self._crop_size.get_value() + if detector_size is None: + return max(1, requested) + return Interval[int](1, detector_size).clamp(requested) + # No crop: fall back to whatever the detector reports. Callers that need + # a concrete extent (pipeline construction) must ensure a dataset is loaded. + return detector_size if detector_size is not None else 0 + + def get_safe_crop_center(self, detector_size: int | None) -> int: + """Crop center clamped so the configured crop window fits inside the detector. + + CropStep slices ``[center - radius, center + radius)`` with + ``radius = crop_size // 2``, so valid centers satisfy + ``radius <= center <= det_size - radius``. + """ + radius = self.get_crop_size(detector_size) // 2 + if detector_size is None: + return max(radius, self._crop_center.get_value()) + return Interval[int](radius, detector_size - radius).clamp(self._crop_center.get_value()) + + def get_bin_size(self, detector_size: int | None) -> int: if self._binning_enabled.get_value(): - limits = self.get_bin_size_limits() - return limits.clamp(self._bin_size.get_value()) - + return Interval[int](1, self.get_crop_size(detector_size)).clamp( + self._bin_size.get_value() + ) return 1 - def validate_bin_size(self) -> None: - crop_size = self.get_crop_size() - bin_size = self.get_bin_size() + def validate_bin_size(self, detector_size: int | None) -> None: + crop_size = self.get_crop_size(detector_size) + bin_size = self.get_bin_size(detector_size) if crop_size % bin_size != 0: raise ValueError(f'Invalid binning size! ({crop_size=}, {bin_size=})') @@ -94,42 +84,21 @@ def validate_bin_size(self) -> None: def get_pad_size(self) -> int: if self._padding_enabled.get_value(): return self._pad_size.get_value() - return 0 - def get_processed_size(self) -> int: - return self.get_crop_size() // self.get_bin_size() + 2 * self.get_pad_size() - - def get_processed_pixel_size_m(self) -> float: - return self.get_bin_size() * self._detector_pixel_size_m.get_value() - - def get_processed_size_m(self) -> float: - return self.get_processed_size() * self.get_processed_pixel_size_m() - def _update(self, observable: Observable) -> None: - if observable in ( - self._detector_size, - self._detector_pixel_size_m, - self._crop_enabled, - self._crop_size, - self._crop_center, - self._binning_enabled, - self._bin_size, - self._padding_enabled, - self._pad_size, - ): - self.notify_observers() + self.notify_observers() class PatternSizer(Observable, Observer): def __init__( - self, detector_settings: DetectorSettings, diffraction_settings: DiffractionSettings + self, + diffraction_settings: DiffractionSettings, ) -> None: super().__init__() self._diffraction_settings = diffraction_settings - self.axis_x = PatternAxisSizer( - detector_settings.width_px, - detector_settings.pixel_width_m, + + self._axis_x = PatternAxisSizer( diffraction_settings.crop_enabled, diffraction_settings.crop_width_px, diffraction_settings.crop_center_x_px, @@ -138,9 +107,7 @@ def __init__( diffraction_settings.padding_enabled, diffraction_settings.pad_x, ) - self.axis_y = PatternAxisSizer( - detector_settings.height_px, - detector_settings.pixel_height_m, + self._axis_y = PatternAxisSizer( diffraction_settings.crop_enabled, diffraction_settings.crop_height_px, diffraction_settings.crop_center_y_px, @@ -150,87 +117,111 @@ def __init__( diffraction_settings.pad_y, ) - self.axis_x.add_observer(self) - self.axis_y.add_observer(self) - - def get_detector_extent(self) -> ImageExtent: - return ImageExtent( - width_px=self.axis_x.get_detector_size(), - height_px=self.axis_y.get_detector_size(), - ) - - def get_processed_width_m(self) -> float: - return self.axis_x.get_processed_size_m() - - def get_processed_height_m(self) -> float: - return self.axis_y.get_processed_size_m() - - def get_processed_image_extent(self) -> ImageExtent: - return ImageExtent( - width_px=self.axis_x.get_processed_size(), - height_px=self.axis_y.get_processed_size(), - ) - - def get_processed_pixel_geometry(self) -> PixelGeometry: - return PixelGeometry( - width_m=self.axis_x.get_processed_pixel_size_m(), - height_m=self.axis_y.get_processed_pixel_size_m(), + self._axis_x.add_observer(self) + self._axis_y.add_observer(self) + + # Whole-image parameters that don't decompose per axis. Register directly so + # get_prep_pipeline()/get_processed_*() consumers wake up on these edits. + diffraction_settings.hflip.add_observer(self) + diffraction_settings.vflip.add_observer(self) + diffraction_settings.transpose.add_observer(self) + diffraction_settings.value_lower_bound_enabled.add_observer(self) + diffraction_settings.value_lower_bound.add_observer(self) + diffraction_settings.value_upper_bound_enabled.add_observer(self) + diffraction_settings.value_upper_bound.add_observer(self) + + def get_processed_image_extent(self, detector_extent: ImageExtent | None = None) -> ImageExtent: + if detector_extent is None: + return ImageExtent(width_px=0, height_px=0) + return self.get_prep_pipeline(detector_extent).compute_output_extent(detector_extent) + + def get_processed_pixel_geometry(self, raw_pixel_geometry: PixelGeometry) -> PixelGeometry: + # Pixel geometry only depends on binning and transpose (see + # DiffractionPrepStep.apply_to_pixel_geometry overrides); crop, filter, padding, + # and flips are identity. Compute directly from the raw settings so this method + # works without knowing the detector extent. + geometry = raw_pixel_geometry + if self._diffraction_settings.binning_enabled.get_value(): + geometry = BinningStep( + bin_size_x=self._diffraction_settings.bin_size_x.get_value(), + bin_size_y=self._diffraction_settings.bin_size_y.get_value(), + ).apply_to_pixel_geometry(geometry) + if self._diffraction_settings.transpose.get_value(): + geometry = TransposeStep().apply_to_pixel_geometry(geometry) + return geometry + + def get_prep_pipeline( + self, detector_extent: ImageExtent | None = None + ) -> DiffractionPrepPipeline: + """Snapshot live settings as an ordered preprocessing pipeline. + + When ``detector_extent`` is ``None``, axis clamping degrades gracefully — the + pipeline can still be constructed but crop/bin bounds may not match a real + detector. Callers that will feed real patterns through the pipeline must pass + an extent. + + Canonical order: filter → crop → binning → padding → hflip → vflip → transpose. + """ + det_w = detector_extent.width_px if detector_extent is not None else None + det_h = detector_extent.height_px if detector_extent is not None else None + + steps: list[DiffractionPrepStepUnion] = [] + + lower_bound = ( + self._diffraction_settings.value_lower_bound.get_value() + if self._diffraction_settings.value_lower_bound_enabled.get_value() + else None ) - - def get_processor(self) -> DiffractionPatternProcessor: - value_lower_bound: int | None = None - value_upper_bound: int | None = None - crop: DiffractionPatternCrop | None = None - binning: DiffractionPatternBinning | None = None - padding: DiffractionPatternPadding | None = None - - if self._diffraction_settings.value_lower_bound_enabled.get_value(): - value_lower_bound = self._diffraction_settings.value_lower_bound.get_value() - - if self._diffraction_settings.value_upper_bound_enabled.get_value(): - value_upper_bound = self._diffraction_settings.value_upper_bound.get_value() - - filter_values = DiffractionPatternFilterValues( - lower_bound=value_lower_bound, - upper_bound=value_upper_bound, + upper_bound = ( + self._diffraction_settings.value_upper_bound.get_value() + if self._diffraction_settings.value_upper_bound_enabled.get_value() + else None ) + if lower_bound is not None or upper_bound is not None: + steps.append(FilterValuesStep(lower_bound=lower_bound, upper_bound=upper_bound)) if self._diffraction_settings.crop_enabled.get_value(): - crop = DiffractionPatternCrop( - center=CropCenter( - self.axis_x.get_safe_crop_center(), - self.axis_y.get_safe_crop_center(), - ), - extent=ImageExtent( - self.axis_x.get_crop_size(), - self.axis_y.get_crop_size(), - ), + steps.append( + CropStep( + center=CropCenter( + self._axis_x.get_safe_crop_center(det_w), + self._axis_y.get_safe_crop_center(det_h), + ), + extent=ImageExtent( + width_px=self._axis_x.get_crop_size(det_w), + height_px=self._axis_y.get_crop_size(det_h), + ), + ) ) if self._diffraction_settings.binning_enabled.get_value(): - self.axis_x.validate_bin_size() - self.axis_y.validate_bin_size() - binning = DiffractionPatternBinning( - bin_size_x=self.axis_x.get_bin_size(), - bin_size_y=self.axis_y.get_bin_size(), + self._axis_x.validate_bin_size(det_w) + self._axis_y.validate_bin_size(det_h) + steps.append( + BinningStep( + bin_size_x=self._axis_x.get_bin_size(det_w), + bin_size_y=self._axis_y.get_bin_size(det_h), + ) ) if self._diffraction_settings.padding_enabled.get_value(): - padding = DiffractionPatternPadding( - pad_x=self.axis_x.get_pad_size(), - pad_y=self.axis_y.get_pad_size(), + steps.append( + PaddingStep( + pad_x=self._axis_x.get_pad_size(), + pad_y=self._axis_y.get_pad_size(), + ) ) - return DiffractionPatternProcessor( - filter_values=filter_values, - crop=crop, - binning=binning, - padding=padding, - hflip=self._diffraction_settings.hflip.get_value(), - vflip=self._diffraction_settings.vflip.get_value(), - transpose=self._diffraction_settings.transpose.get_value(), - ) + if self._diffraction_settings.hflip.get_value(): + steps.append(HorizontalFlipStep()) + + if self._diffraction_settings.vflip.get_value(): + steps.append(VerticalFlipStep()) + + if self._diffraction_settings.transpose.get_value(): + steps.append(TransposeStep()) + + return DiffractionPrepPipeline(steps=tuple(steps)) def _update(self, observable: Observable) -> None: - if observable in (self.axis_x, self.axis_y): - self.notify_observers() + self.notify_observers() diff --git a/src/ptychodus/model/fluorescence/__init__.py b/src/ptychodus/model/fluorescence/__init__.py index 4ab9acfd8..2dfd75287 100644 --- a/src/ptychodus/model/fluorescence/__init__.py +++ b/src/ptychodus/model/fluorescence/__init__.py @@ -1,11 +1,15 @@ from .api import FluorescenceAPI from .core import FluorescenceCore -from .monitor import ( - EnhanceFluorescenceBackgroundTask, - FluorescenceDatasetEmitter, - FluorescenceTaskMonitor, -) +from .monitor import EnhanceFluorescenceBackgroundTask, FluorescenceTaskMonitor from .ptychozoon import PtychozoonFluorescenceEnhancer +from .repository import ( + FluorescenceItemState, + FluorescenceRepository, + FluorescenceRepositoryItem, + FluorescenceRepositoryItemObserver, + FluorescenceRepositoryObserver, +) +from .settings import FluorescenceSettings from .two_step import TwoStepFluorescenceEnhancer from .vspi import VSPIFluorescenceEnhancer @@ -13,7 +17,12 @@ 'EnhanceFluorescenceBackgroundTask', 'FluorescenceAPI', 'FluorescenceCore', - 'FluorescenceDatasetEmitter', + 'FluorescenceItemState', + 'FluorescenceRepository', + 'FluorescenceRepositoryItem', + 'FluorescenceRepositoryItemObserver', + 'FluorescenceRepositoryObserver', + 'FluorescenceSettings', 'FluorescenceTaskMonitor', 'PtychozoonFluorescenceEnhancer', 'TwoStepFluorescenceEnhancer', diff --git a/src/ptychodus/model/fluorescence/_payload.py b/src/ptychodus/model/fluorescence/_payload.py new file mode 100644 index 000000000..b696e68cd --- /dev/null +++ b/src/ptychodus/model/fluorescence/_payload.py @@ -0,0 +1,41 @@ +"""Payload dataclass for the ptychozoon VSPI subprocess entry point. + +``ptychozoon`` is an optional extra, so it is imported under +:data:`~typing.TYPE_CHECKING` only: this module has ``from __future__ import +annotations`` and :class:`PtychozoonPayload` is a dataclass, whose field +annotations are never evaluated. That keeps ``ptychodus.model`` importable +without the extra installed. The enhancer defers the matching runtime imports +into :meth:`~.ptychozoon.PtychozoonFluorescenceEnhancer._build_payload`. + +Only ``ptychozoon.data_structures`` and ``ptychozoon.settings`` may be reached +from the parent: they pull in just ``numpy``, ``dataclasses``, ``enum``, and +``typing`` (no CuPy). Anything else under ``ptychozoon.*`` is GPU-tainted and +must stay inside the child; see +[tests/test_no_gpu_context.py](../../../../tests/test_no_gpu_context.py) for +the allow-list that pins that guarantee. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ptychozoon.data_structures import ( + FluorescenceDataset, + PtychographyProduct, + ) + from ptychozoon.settings import DeconvolutionEnhancementSettings + +__all__ = [ + 'PtychozoonPayload', +] + + +@dataclass(frozen=True) +class PtychozoonPayload: + """Picklable inputs for one ptychozoon enhancement run.""" + + product: PtychographyProduct + dataset: FluorescenceDataset + settings: DeconvolutionEnhancementSettings diff --git a/src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py b/src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py deleted file mode 100644 index e045a77b4..000000000 --- a/src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Subprocess worker that runs the ptychozoon (GPU VSPI) fluorescence enhancement. - -This module is executed in a freshly ``spawn``ed process so that the CuPy GPU -context is created cleanly and fully released when the process exits. It must not -import ``ptychozoon`` (or CuPy) at module load time; the import happens inside the -worker function so that the parent process never pulls GPU libraries into its own -interpreter. - -All data crosses the process boundary as plain numpy arrays and scalars (see -:class:`PtychozoonPayload`). Messages are streamed back over a single queue as -tagged tuples: - -- ``('result', iteration, [(element_name, counts_per_second_array), ...])`` per checkpoint -- ``('log', levelno, message)`` for each captured log record from the ptychozoon logger -- ``('error', traceback_str)`` on failure - -followed by a ``None`` sentinel. Because a single child thread produces every -message, ordering is preserved (a log line arrives before the result it precedes). -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from multiprocessing.queues import Queue -from typing import Any - -import numpy - -__all__ = [ - 'PtychozoonPayload', - 'run_vspi_enhancement', -] - - -class _QueueLogHandler(logging.Handler): - """Logging handler that forwards formatted records to the parent over the result queue.""" - - def __init__(self, result_queue: 'Queue[Any]') -> None: - super().__init__() - self._result_queue = result_queue - - def emit(self, record: logging.LogRecord) -> None: - try: - self._result_queue.put(('log', record.levelno, self.format(record))) - except Exception: - # Never let logging failures break the enhancement worker. - pass - - -@dataclass(frozen=True) -class PtychozoonPayload: - """Picklable inputs for one ptychozoon enhancement run.""" - - # Ptychography product (plain arrays / scalars) - probe_positions_m: numpy.ndarray # (N, 2) float [y, x] meters - probe: numpy.ndarray # (n_opr, modes, height, width) complex - object_array: numpy.ndarray # (height, width) complex - pixel_size_m: tuple[float, float] # (pixel_height_m, pixel_width_m) - object_center_m: tuple[float, float] # (center_y_m, center_x_m) - opr_mode_weights: numpy.ndarray | None # (n_opr, N) float or None - - # Fluorescence element maps - element_maps: list[tuple[str, numpy.ndarray]] # [(name, counts_per_second), ...] - - # Solver settings - damping_factor: float - gradient_smoothness: float - max_iterations: int - atol: float - btol: float - checkpoint_interval: int - use_gpu: bool - gpu_device_index: int - - # Effective log level of the parent's ptychozoon logger, so the child only - # forwards records that the parent would actually surface. - log_level: int - - -def run_vspi_enhancement(payload: PtychozoonPayload, result_queue: 'Queue[Any]') -> None: - """Run ptychozoon VSPI enhancement and stream checkpoint results over a queue. - - Intended as the target of a ``spawn``ed ``multiprocessing.Process``. Puts - ``('result', iteration, [(name, cps), ...])`` tuples per checkpoint and - ``('log', levelno, message)`` tuples for ptychozoon log records, then ``None`` - when complete, or ``('error', traceback_str)`` on failure. - """ - # Forward ptychozoon's own log output to the parent so it appears in the - # Enhance Fluorescence status view (this process has no ptychodus handlers). - log_handler = _QueueLogHandler(result_queue) - log_handler.setFormatter(logging.Formatter('%(message)s')) - ptychozoon_logger = logging.getLogger('ptychozoon') - ptychozoon_logger.addHandler(log_handler) - ptychozoon_logger.setLevel(payload.log_level) - - try: - from ptychozoon.data_structures import ( - ElementMap, - FluorescenceDataset, - PtychographyProduct, - ) - from ptychozoon.settings import ( - DeconvolutionEnhancementSettings, - InterpolationTypes, - ) - from ptychozoon.vspi_enhance import VSPIFluorescenceEnhancingAlgorithm - - product = PtychographyProduct( - probe_positions=payload.probe_positions_m, - probe=payload.probe, - object_array=payload.object_array, - pixel_size_m=payload.pixel_size_m, - object_center_m=payload.object_center_m, - opr_mode_weights=payload.opr_mode_weights, - ) - dataset = FluorescenceDataset( - element_maps=[ - ElementMap(name=name, counts_per_second=cps) for name, cps in payload.element_maps - ] - ) - - settings = DeconvolutionEnhancementSettings() - settings.lsmr.damping_factor = payload.damping_factor - settings.lsmr.gradient_smoothness = payload.gradient_smoothness - settings.lsmr.max_iter = payload.max_iterations - settings.lsmr.atol = payload.atol - settings.lsmr.btol = payload.btol - settings.lsmr.checkpoint_interval = payload.checkpoint_interval - settings.gpu.enabled = payload.use_gpu - settings.gpu.index = payload.gpu_device_index - # Fourier interpolation requires the GPU; fall back to Barycentric on CPU. - settings._interpolation = ( - InterpolationTypes.FOURIER if payload.use_gpu else InterpolationTypes.BARYCENTRIC - ) - - algorithm = VSPIFluorescenceEnhancingAlgorithm() - - for enhanced_dataset, iteration in algorithm.enhance(dataset, product, settings=settings): - result_queue.put( - ( - 'result', - int(iteration), - [ - (emap.name, numpy.asarray(emap.counts_per_second)) - for emap in enhanced_dataset.element_maps - ], - ) - ) - except Exception: - import traceback - - result_queue.put(('error', traceback.format_exc())) - finally: - ptychozoon_logger.removeHandler(log_handler) - result_queue.put(None) diff --git a/src/ptychodus/model/fluorescence/_subprocess.py b/src/ptychodus/model/fluorescence/_subprocess.py new file mode 100644 index 000000000..77a0b2c20 --- /dev/null +++ b/src/ptychodus/model/fluorescence/_subprocess.py @@ -0,0 +1,52 @@ +"""Subprocess worker that runs the ptychozoon (GPU VSPI) fluorescence enhancement. + +This module is executed in a freshly ``spawn``ed process so that the CuPy GPU +context is created cleanly and fully released when the process exits. The +GPU-tainted ``ptychozoon.vspi_enhance`` import is deferred to the worker body +so the parent process never pulls CuPy into its own interpreter. + +The payload carries fully constructed ``ptychozoon`` objects (see +:mod:`._payload`); this file just hands them to the enhancement algorithm and +streams checkpoint results back. Log forwarding and error marshaling are +handled by the shared subprocess protocol in +:mod:`ptychodus.model.processing._subprocess_protocol`. +""" + +from __future__ import annotations + +from multiprocessing.queues import Queue +from typing import Any + +import numpy + +from ._payload import PtychozoonPayload + +__all__ = [ + 'run_vspi_enhancement', +] + + +def run_vspi_enhancement(payload: PtychozoonPayload, result_queue: Queue[Any]) -> None: + """Run ptychozoon VSPI enhancement and stream checkpoint results over a queue. + + Intended as the ``entry_point`` for + :func:`ptychodus.model.processing._subprocess_protocol.run_subprocess`. + Puts ``('result', iteration, [(name, cps), ...])`` tuples per checkpoint. + """ + from ptychozoon.vspi_enhance import VSPIFluorescenceEnhancingAlgorithm + + algorithm = VSPIFluorescenceEnhancingAlgorithm() + + for enhanced_dataset, iteration in algorithm.enhance( + payload.dataset, payload.product, settings=payload.settings + ): + result_queue.put( + ( + 'result', + int(iteration), + [ + (emap.name, numpy.asarray(emap.counts_per_second)) + for emap in enhanced_dataset.element_maps + ], + ) + ) diff --git a/src/ptychodus/model/fluorescence/api.py b/src/ptychodus/model/fluorescence/api.py index 354aa8ee0..a4269c437 100644 --- a/src/ptychodus/model/fluorescence/api.py +++ b/src/ptychodus/model/fluorescence/api.py @@ -3,20 +3,19 @@ import logging from ptychodus.api.fluorescence import ( - FluorescenceDataset, FluorescenceEnhancer, FluorescenceFileReader, FluorescenceFileWriter, ) -from ptychodus.api.geometry import PixelGeometry from ptychodus.api.plugins import PluginChooser from ..product import ProductAPI from ..task_manager import TaskManager -from .monitor import ( - EnhanceFluorescenceBackgroundTask, - FluorescenceDatasetEmitter, - FluorescenceTaskMonitor, +from .monitor import EnhanceFluorescenceBackgroundTask, FluorescenceTaskMonitor +from .repository import ( + FluorescenceItemState, + FluorescenceRepository, + FluorescenceRepositoryItem, ) from .settings import FluorescenceSettings @@ -29,6 +28,7 @@ def __init__( task_manager: TaskManager, settings: FluorescenceSettings, product_api: ProductAPI, + repository: FluorescenceRepository, enhancer_chooser: PluginChooser[FluorescenceEnhancer], task_monitor: FluorescenceTaskMonitor, file_reader_chooser: PluginChooser[FluorescenceFileReader], @@ -37,20 +37,17 @@ def __init__( self._task_manager = task_manager self._settings = settings self._product_api = product_api + self._repository = repository self._enhancer_chooser = enhancer_chooser self._task_monitor = task_monitor self._file_reader_chooser = file_reader_chooser self._file_writer_chooser = file_writer_chooser - def get_dataset_emitter(self) -> FluorescenceDatasetEmitter: - return self._task_monitor.get_dataset_emitter() + def get_task_monitor(self) -> FluorescenceTaskMonitor: + return self._task_monitor - def get_product_name(self, product_index: int) -> str: - return self._product_api.get_item(product_index).get_name() - - def get_pixel_geometry(self, product_index: int) -> PixelGeometry: - item = self._product_api.get_item(product_index) - return item.get_geometry().get_object_plane_pixel_geometry() + def get_item(self, item_index: int) -> FluorescenceRepositoryItem: + return self._repository[item_index] def get_open_file_filters(self) -> Iterator[str]: for plugin in self._file_reader_chooser: @@ -66,14 +63,18 @@ def get_save_file_filters(self) -> Iterator[str]: def get_save_file_filter(self) -> str: return self._file_writer_chooser.get_current_plugin().display_name - def load_measured_dataset( - self, file_path: Path, *, file_type: str | None = None - ) -> FluorescenceDataset: + def open_measured_dataset( + self, file_path: Path, product_index: int, *, file_type: str | None = None + ) -> int: if not file_path.is_file(): raise FileNotFoundError( f'Refusing to load dataset from invalid file path "{file_path}"' ) + # Resolve the product first — a bad product_index should fail before we + # touch the file, so an orphaned item is never inserted. + product_item = self._product_api.get_item(product_index) + if file_type is not None: self._file_reader_chooser.set_current_plugin(file_type) @@ -87,11 +88,33 @@ def load_measured_dataset( raise RuntimeError(f'Failed to read "{file_path}"') from exc self._settings.file_path.set_value(file_path) - return dataset + + label = self._repository.create_unique_name(file_path.stem) + item = FluorescenceRepositoryItem( + self._repository, + label=label, + product=product_item, + measured=dataset, + source_path=file_path, + source_file_type=resolved_type, + ) + return self._repository.insert_item(item) + + def remove_item(self, item_index: int) -> None: + self._repository.remove_item(item_index) def save_enhanced_dataset( - self, dataset: FluorescenceDataset, file_path: Path, *, file_type: str | None = None + self, item_index: int, file_path: Path, *, file_type: str | None = None ) -> None: + try: + item = self._repository[item_index] + except IndexError as exc: + raise ValueError(f'No fluorescence item at index {item_index}') from exc + + dataset = item.get_enhanced() + if dataset is None: + raise ValueError(f'Item "{item.get_label()}" has no enhanced dataset to save') + if file_type is not None: self._file_writer_chooser.set_current_plugin(file_type) @@ -102,14 +125,26 @@ def save_enhanced_dataset( def enhance( self, - product_index: int, - dataset: FluorescenceDataset, + item_index: int, *, algorithm: str | None = None, output_file_path: Path | None = None, output_file_type: str | None = None, block: bool = False, ) -> None: + try: + item = self._repository[item_index] + except IndexError as exc: + raise ValueError(f'No fluorescence item at index {item_index}') from exc + + if item.get_state() is FluorescenceItemState.ORPHANED: + raise RuntimeError( + f'Cannot enhance item "{item.get_label()}": target product was removed' + ) + + if self._task_monitor.is_processing: + raise RuntimeError('Fluorescence enhancement is already in progress') + if algorithm is not None: self._enhancer_chooser.set_current_plugin(algorithm) @@ -121,12 +156,15 @@ def enhance( output_file_writer = self._file_writer_chooser.get_current_plugin().strategy - product = self._product_api.get_item(product_index).get_product() + product = item.get_product().get_product() + item.set_state(FluorescenceItemState.ENHANCING) + background_task = EnhanceFluorescenceBackgroundTask( self._task_monitor, self._enhancer_chooser.get_current_plugin().strategy, - dataset, + item.get_measured(), product, + item, output_file_path, output_file_writer, ) @@ -140,24 +178,3 @@ def enhance( ): self._task_manager.run_foreground_tasks() break - - def enhance_local( - self, - product_index: int, - input_file_path: Path, - output_file_path: Path, - *, - input_file_type: str | None = None, - output_file_type: str | None = None, - algorithm: str | None = None, - block: bool = False, - ) -> None: - dataset = self.load_measured_dataset(input_file_path, file_type=input_file_type) - self.enhance( - product_index, - dataset, - algorithm=algorithm, - output_file_path=output_file_path, - output_file_type=output_file_type, - block=block, - ) diff --git a/src/ptychodus/model/fluorescence/core.py b/src/ptychodus/model/fluorescence/core.py index 72a049b04..50c728684 100644 --- a/src/ptychodus/model/fluorescence/core.py +++ b/src/ptychodus/model/fluorescence/core.py @@ -7,15 +7,16 @@ FluorescenceFileWriter, UpscalingStrategy, ) -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.settings import SettingsRegistry -from ..product import ProductAPI +from ..product import ProductAPI, ProductRepository from ..task_manager import TaskManager from ..visualization import VisualizationEngine from .api import FluorescenceAPI from .monitor import FluorescenceTaskMonitor from .ptychozoon import PtychozoonFluorescenceEnhancer +from .repository import FluorescenceRepository from .settings import FluorescenceSettings from .two_step import TwoStepFluorescenceEnhancer from .vspi import VSPIFluorescenceEnhancer @@ -29,16 +30,19 @@ def __init__( task_manager: TaskManager, settings_registry: SettingsRegistry, product_api: ProductAPI, + product_repository: ProductRepository, upscaling_strategy_chooser: PluginChooser[UpscalingStrategy], deconvolution_strategy_chooser: PluginChooser[DeconvolutionStrategy], file_reader_chooser: PluginChooser[FluorescenceFileReader], file_writer_chooser: PluginChooser[FluorescenceFileWriter], ) -> None: - self._settings = FluorescenceSettings(settings_registry) + self.settings = FluorescenceSettings(settings_registry) + self.upscaling_strategy_chooser = upscaling_strategy_chooser + self.deconvolution_strategy_chooser = deconvolution_strategy_chooser self.two_step_enhancer = TwoStepFluorescenceEnhancer( - self._settings, upscaling_strategy_chooser, deconvolution_strategy_chooser + upscaling_strategy_chooser, deconvolution_strategy_chooser ) - self.vspi_enhancer = VSPIFluorescenceEnhancer(self._settings) + self.vspi_enhancer = VSPIFluorescenceEnhancer(self.settings) self.enhancer_chooser = PluginChooser[FluorescenceEnhancer]() self.enhancer_chooser.register_plugin( @@ -53,10 +57,9 @@ def __init__( ) # The GPU VSPI enhancer is optional: register it only when ptychozoon is - # installed so it silently disappears otherwise. Register it last so its - # combo-box entry and stacked GUI page stay aligned by index. Importing - # the top-level package is cheap and does not pull in CuPy (that happens - # only inside the spawned subprocess), so probing it here is safe. + # installed so it silently disappears otherwise. Importing the top-level + # package is cheap and does not pull in CuPy (that happens only inside the + # spawned subprocess), so probing it here is safe. self.ptychozoon_enhancer: PtychozoonFluorescenceEnhancer | None = None try: @@ -64,24 +67,40 @@ def __init__( except ModuleNotFoundError: logger.info('ptychozoon not found.') else: - self.ptychozoon_enhancer = PtychozoonFluorescenceEnhancer(self._settings) + self.ptychozoon_enhancer = PtychozoonFluorescenceEnhancer(self.settings) self.enhancer_chooser.register_plugin( self.ptychozoon_enhancer, simple_name=PtychozoonFluorescenceEnhancer.SIMPLE_NAME, display_name=PtychozoonFluorescenceEnhancer.DISPLAY_NAME, ) - self.enhancer_chooser.synchronize_with_parameter(self._settings.algorithm) + # Display-name views of each chooser, bound to the settings parameter that + # persists the selection. These are what the GUI binds combo boxes to. + self.enhancer_parameter = PluginChooserParameter( + self.enhancer_chooser, self.settings.algorithm + ) + self.upscaling_strategy_parameter = PluginChooserParameter( + upscaling_strategy_chooser, self.settings.upscaling_strategy + ) + self.deconvolution_strategy_parameter = PluginChooserParameter( + deconvolution_strategy_chooser, self.settings.deconvolution_strategy + ) + self.file_reader_parameter = PluginChooserParameter( + file_reader_chooser, self.settings.file_type + ) - file_reader_chooser.synchronize_with_parameter(self._settings.file_type) - file_writer_chooser.set_current_plugin(self._settings.file_type.get_value()) + # Deliberately unbound: the writer shares file_type with the reader above, so + # binding both would make them fight over the same parameter. + file_writer_chooser.set_current_plugin(self.settings.file_type.get_value()) self.visualization_engine = VisualizationEngine(is_complex=False) self.task_monitor = FluorescenceTaskMonitor(task_manager) + self.repository = FluorescenceRepository(product_repository) self.fluorescence_api = FluorescenceAPI( task_manager, - self._settings, + self.settings, product_api, + self.repository, self.enhancer_chooser, self.task_monitor, file_reader_chooser, diff --git a/src/ptychodus/model/fluorescence/monitor.py b/src/ptychodus/model/fluorescence/monitor.py index ef0a5f01d..645e6a123 100644 --- a/src/ptychodus/model/fluorescence/monitor.py +++ b/src/ptychodus/model/fluorescence/monitor.py @@ -2,7 +2,6 @@ from dataclasses import dataclass from pathlib import Path import logging -import threading import time from ptychodus.api.fluorescence import ( @@ -11,49 +10,62 @@ FluorescenceEnhancerInput, FluorescenceFileWriter, ) -from ptychodus.api.observer import Observable from ptychodus.api.product import Product -from ..task_monitor import NotifyObserversTask, TaskProgressMonitor -from ..task_manager import ForegroundTaskManager +from ..task_monitor import TaskProgressMonitor +from .repository import FluorescenceItemState, FluorescenceRepositoryItem __all__ = [ 'EnhanceFluorescenceBackgroundTask', - 'FluorescenceDatasetEmitter', 'FluorescenceTaskMonitor', ] logger = logging.getLogger(__name__) -class FluorescenceDatasetEmitter(Observable): - def __init__(self, foreground_task_manager: ForegroundTaskManager) -> None: - super().__init__() - self._foreground_task_manager = foreground_task_manager - self._lock = threading.Lock() - self._latest_enhanced: FluorescenceDataset | None = None +class UpdateEnhancedTask: + """Foreground-scheduled application of a new enhanced dataset onto its item.""" - def get_latest_enhanced(self) -> FluorescenceDataset | None: - with self._lock: - return self._latest_enhanced + def __init__(self, item: FluorescenceRepositoryItem, dataset: FluorescenceDataset) -> None: + self._item = item + self._dataset = dataset - def _publish(self, dataset: FluorescenceDataset) -> None: - with self._lock: - self._latest_enhanced = dataset + def __call__(self) -> None: + self._item.set_enhanced(self._dataset) - self._foreground_task_manager.put_foreground_task(NotifyObserversTask(self)) +class SetItemStateTask: + """Foreground-scheduled state transition on a FluorescenceRepositoryItem. -class FluorescenceTaskMonitor(TaskProgressMonitor): - def __init__(self, foreground_task_manager: ForegroundTaskManager) -> None: - super().__init__(foreground_task_manager) - self._dataset_emitter = FluorescenceDatasetEmitter(foreground_task_manager) + Skips the transition if the item is already ORPHANED — the product was + removed mid-enhancement and the orphan flag must win over the task's own + READY/FAILED verdict. + """ - def get_dataset_emitter(self) -> FluorescenceDatasetEmitter: - return self._dataset_emitter + def __init__( + self, item: FluorescenceRepositoryItem, target_state: FluorescenceItemState + ) -> None: + self._item = item + self._target_state = target_state + + def __call__(self) -> None: + if self._item.get_state() is FluorescenceItemState.ORPHANED: + return + self._item.set_state(self._target_state) + + +class FluorescenceTaskMonitor(TaskProgressMonitor): + def update_enhanced( + self, item: FluorescenceRepositoryItem, dataset: FluorescenceDataset + ) -> None: + task = UpdateEnhancedTask(item, dataset) + self._foreground_task_manager.put_foreground_task(task) - def update_enhanced(self, dataset: FluorescenceDataset) -> None: - self._dataset_emitter._publish(dataset) + def transition_item_state( + self, item: FluorescenceRepositoryItem, target_state: FluorescenceItemState + ) -> None: + task = SetItemStateTask(item, target_state) + self._foreground_task_manager.put_foreground_task(task) @dataclass(frozen=True) @@ -62,33 +74,44 @@ class EnhanceFluorescenceBackgroundTask: enhancer: FluorescenceEnhancer dataset: FluorescenceDataset product: Product + item: FluorescenceRepositoryItem output_file_path: Path | None output_file_writer: FluorescenceFileWriter | None def __call__(self) -> None: - with self.task_monitor as monitor: - progress_goal = self.enhancer.get_progress_goal() - monitor.update_progress(0, progress_goal) - tic = time.perf_counter() - last_dataset: FluorescenceDataset | None = None - parameters = FluorescenceEnhancerInput(dataset=self.dataset, product=self.product) - - for output in self.enhancer.enhance(parameters): - monitor.update_progress(output.progress, progress_goal) - monitor.update_enhanced(output.dataset) - last_dataset = output.dataset - - if monitor.is_stopping: - break - - toc = time.perf_counter() - logger.info(f'Enhancement time {toc - tic:.4f} seconds.') - - if ( - last_dataset is not None - and not monitor.is_stopping - and self.output_file_path is not None - and self.output_file_writer is not None - ): - logger.debug(f'Writing enhanced fluorescence to "{self.output_file_path}"') - self.output_file_writer.write(self.output_file_path, last_dataset) + try: + with self.task_monitor as monitor: + progress_goal = self.enhancer.get_progress_goal() + monitor.update_progress(0, progress_goal) + tic = time.perf_counter() + last_dataset: FluorescenceDataset | None = None + parameters = FluorescenceEnhancerInput(dataset=self.dataset, product=self.product) + + for output in self.enhancer.enhance(parameters): + monitor.update_progress(output.progress, progress_goal) + monitor.update_enhanced(self.item, output.dataset) + last_dataset = output.dataset + + if monitor.is_stopping: + break + + toc = time.perf_counter() + logger.info(f'Enhancement time {toc - tic:.4f} seconds.') + + if ( + last_dataset is not None + and not monitor.is_stopping + and self.output_file_path is not None + and self.output_file_writer is not None + ): + logger.debug(f'Writing enhanced fluorescence to "{self.output_file_path}"') + self.output_file_writer.write(self.output_file_path, last_dataset) + except Exception: + # Foreground state update runs even on failure so the UI reflects + # FAILED before the exception propagates through TaskManager. + self.task_monitor.transition_item_state(self.item, FluorescenceItemState.FAILED) + raise + else: + # Natural completion or user stop both land here — user-stop is + # considered a graceful end of the enhancement, not a failure. + self.task_monitor.transition_item_state(self.item, FluorescenceItemState.READY) diff --git a/src/ptychodus/model/fluorescence/ptychozoon.py b/src/ptychodus/model/fluorescence/ptychozoon.py index 9907305c9..a7bfe709c 100644 --- a/src/ptychodus/model/fluorescence/ptychozoon.py +++ b/src/ptychodus/model/fluorescence/ptychozoon.py @@ -1,16 +1,17 @@ """GPU-accelerated VSPI fluorescence enhancer backed by the ptychozoon package. The heavy CuPy computation runs in a freshly ``spawn``ed subprocess (see -:mod:`._ptychozoon_subprocess`) so each run gets a clean GPU context and all GPU -memory is released when the run finishes or is stopped. This module never imports -ptychozoon or CuPy directly. +:mod:`._subprocess`) so each run gets a clean GPU context and all GPU memory +is released when the run finishes or is stopped. This module only imports the +CPU-safe ``ptychozoon.data_structures`` and ``ptychozoon.settings`` submodules +(needed to construct the payload); CuPy-linked submodules stay inside the +child. """ from __future__ import annotations from collections.abc import Iterator from typing import Final import logging -import multiprocessing import numpy @@ -21,12 +22,14 @@ FluorescenceEnhancerInput, FluorescenceEnhancerOutput, ) -from ptychodus.api.observer import Observable, Observer from ptychodus.api.product import Product -from ._ptychozoon_subprocess import PtychozoonPayload, run_vspi_enhancement +from ..processing._subprocess_protocol import run_subprocess +from ._payload import PtychozoonPayload from .settings import FluorescenceSettings +_ENTRY_POINT: Final[str] = 'ptychodus.model.fluorescence._subprocess:run_vspi_enhancement' + logger = logging.getLogger(__name__) __all__ = [ @@ -34,7 +37,7 @@ ] -class PtychozoonFluorescenceEnhancer(FluorescenceEnhancer, Observable, Observer): +class PtychozoonFluorescenceEnhancer(FluorescenceEnhancer): SIMPLE_NAME: Final[str] = 'VSPI-GPU' DISPLAY_NAME: Final[str] = 'Virtual Single Pixel Imaging (GPU)' @@ -42,15 +45,6 @@ def __init__(self, settings: FluorescenceSettings) -> None: super().__init__() self._settings = settings - settings.ptychozoon_damping_factor.add_observer(self) - settings.ptychozoon_gradient_smoothness.add_observer(self) - settings.ptychozoon_max_iterations.add_observer(self) - settings.ptychozoon_atol.add_observer(self) - settings.ptychozoon_btol.add_observer(self) - settings.ptychozoon_checkpoint_interval.add_observer(self) - settings.ptychozoon_use_gpu.add_observer(self) - settings.ptychozoon_gpu_device_index.add_observer(self) - @property def name(self) -> str: return self.DISPLAY_NAME @@ -59,6 +53,16 @@ def get_progress_goal(self) -> int: return self._settings.ptychozoon_max_iterations.get_value() def _build_payload(self, parameters: FluorescenceEnhancerInput) -> PtychozoonPayload: + # Imported here, not at module scope, so this module stays importable without + # the optional ptychozoon extra. FluorescenceCore gates registration on + # availability, but the class itself must always import -- the GUI enhance + # dialog reads DISPLAY_NAME off it. Only the two CPU-safe submodules are + # touched; see the allow-list in tests/test_no_gpu_context.py. + from ptychozoon.data_structures import ElementMap as PtychozoonElementMap + from ptychozoon.data_structures import FluorescenceDataset as PtychozoonFluorescenceDataset + from ptychozoon.data_structures import PtychographyProduct + from ptychozoon.settings import DeconvolutionEnhancementSettings, InterpolationTypes + product: Product = parameters.product dataset = parameters.dataset object_geometry = product.object_.get_geometry() @@ -71,25 +75,44 @@ def _build_payload(self, parameters: FluorescenceEnhancerInput) -> PtychozoonPay opr_weights = product.probes.get_opr_weights_or_none() opr_mode_weights = None if opr_weights is None else numpy.ascontiguousarray(opr_weights.T) - return PtychozoonPayload( - probe_positions_m=probe_positions_m, + ptychozoon_product = PtychographyProduct( + probe_positions=probe_positions_m, probe=product.probes.get_array(), object_array=product.object_.get_layers_flattened(), pixel_size_m=(object_geometry.pixel_height_m, object_geometry.pixel_width_m), object_center_m=(object_geometry.center_y_m, object_geometry.center_x_m), opr_mode_weights=opr_mode_weights, + ) + ptychozoon_dataset = PtychozoonFluorescenceDataset( element_maps=[ - (emap.name, numpy.asarray(emap.counts_per_second)) for emap in dataset.element_maps - ], - damping_factor=self._settings.ptychozoon_damping_factor.get_value(), - gradient_smoothness=self._settings.ptychozoon_gradient_smoothness.get_value(), - max_iterations=self._settings.ptychozoon_max_iterations.get_value(), - atol=self._settings.ptychozoon_atol.get_value(), - btol=self._settings.ptychozoon_btol.get_value(), - checkpoint_interval=self._settings.ptychozoon_checkpoint_interval.get_value(), - use_gpu=self._settings.ptychozoon_use_gpu.get_value(), - gpu_device_index=self._settings.ptychozoon_gpu_device_index.get_value(), - log_level=logger.getEffectiveLevel(), + PtychozoonElementMap(name=emap.name, counts_per_second=emap.counts_per_second) + for emap in dataset.element_maps + ] + ) + + settings = DeconvolutionEnhancementSettings() + settings.lsmr.damping_factor = self._settings.ptychozoon_damping_factor.get_value() + settings.lsmr.gradient_smoothness = ( + self._settings.ptychozoon_gradient_smoothness.get_value() + ) + settings.lsmr.max_iter = self._settings.ptychozoon_max_iterations.get_value() + settings.lsmr.atol = self._settings.ptychozoon_atol.get_value() + settings.lsmr.btol = self._settings.ptychozoon_btol.get_value() + settings.lsmr.checkpoint_interval = ( + self._settings.ptychozoon_checkpoint_interval.get_value() + ) + use_gpu = self._settings.ptychozoon_use_gpu.get_value() + settings.gpu.enabled = use_gpu + settings.gpu.index = self._settings.ptychozoon_gpu_device_index.get_value() + # Fourier interpolation requires the GPU; fall back to Barycentric on CPU. + settings._interpolation = ( + InterpolationTypes.FOURIER if use_gpu else InterpolationTypes.BARYCENTRIC + ) + + return PtychozoonPayload( + product=ptychozoon_product, + dataset=ptychozoon_dataset, + settings=settings, ) def enhance( @@ -99,110 +122,19 @@ def enhance( payload = self._build_payload(parameters) # A fresh spawned process gives ptychozoon/CuPy a clean GPU context and - # releases all GPU memory when it exits. - ctx = multiprocessing.get_context('spawn') - result_queue = ctx.Queue() - process = ctx.Process(target=run_vspi_enhancement, args=(payload, result_queue)) - process.start() - - try: - while True: - item = result_queue.get() - - if item is None: - break - - tag = item[0] - - if tag == 'log': - _, levelno, message = item - # Re-emit through this process's logger so it reaches the - # fluorescence status view via the registered handler. - logger.log(levelno, message) - continue - - if tag == 'error': - raise RuntimeError(f'ptychozoon enhancement failed:\n{item[1]}') - - if tag == 'result': - _, iteration, enhanced_maps = item - element_maps = [ElementMap(name, cps) for name, cps in enhanced_maps] - yield FluorescenceEnhancerOutput( - dataset=FluorescenceDataset( - element_maps=element_maps, - counts_per_second_path=dataset.counts_per_second_path, - channel_names_path=dataset.channel_names_path, - ), - progress=iteration, - ) + # releases all GPU memory when it exits. Log forwarding and error + # marshaling live in the shared subprocess protocol. + with run_subprocess(_ENTRY_POINT, payload) as events: + for event in events: + if event[0] != 'result': continue - finally: - if process.is_alive(): - process.terminate() - - process.join(timeout=10.0) - - if process.is_alive(): - process.kill() - process.join() - - def get_damping_factor(self) -> float: - return self._settings.ptychozoon_damping_factor.get_value() - - def set_damping_factor(self, factor: float) -> None: - self._settings.ptychozoon_damping_factor.set_value(factor) - - def get_gradient_smoothness(self) -> float: - return self._settings.ptychozoon_gradient_smoothness.get_value() - - def set_gradient_smoothness(self, value: float) -> None: - self._settings.ptychozoon_gradient_smoothness.set_value(value) - - def get_max_iterations(self) -> int: - return self._settings.ptychozoon_max_iterations.get_value() - - def set_max_iterations(self, number: int) -> None: - self._settings.ptychozoon_max_iterations.set_value(number) - - def get_atol(self) -> float: - return self._settings.ptychozoon_atol.get_value() - - def set_atol(self, value: float) -> None: - self._settings.ptychozoon_atol.set_value(value) - - def get_btol(self) -> float: - return self._settings.ptychozoon_btol.get_value() - - def set_btol(self, value: float) -> None: - self._settings.ptychozoon_btol.set_value(value) - - def get_checkpoint_interval(self) -> int: - return self._settings.ptychozoon_checkpoint_interval.get_value() - - def set_checkpoint_interval(self, number: int) -> None: - self._settings.ptychozoon_checkpoint_interval.set_value(number) - - def is_gpu_enabled(self) -> bool: - return self._settings.ptychozoon_use_gpu.get_value() - - def set_gpu_enabled(self, enabled: bool) -> None: - self._settings.ptychozoon_use_gpu.set_value(enabled) - - def get_gpu_device_index(self) -> int: - return self._settings.ptychozoon_gpu_device_index.get_value() - - def set_gpu_device_index(self, index: int) -> None: - self._settings.ptychozoon_gpu_device_index.set_value(index) - - def _update(self, observable: Observable) -> None: - if observable in ( - self._settings.ptychozoon_damping_factor, - self._settings.ptychozoon_gradient_smoothness, - self._settings.ptychozoon_max_iterations, - self._settings.ptychozoon_atol, - self._settings.ptychozoon_btol, - self._settings.ptychozoon_checkpoint_interval, - self._settings.ptychozoon_use_gpu, - self._settings.ptychozoon_gpu_device_index, - ): - self.notify_observers() + _, iteration, enhanced_maps = event + element_maps = [ElementMap(name, cps) for name, cps in enhanced_maps] + yield FluorescenceEnhancerOutput( + dataset=FluorescenceDataset( + element_maps=element_maps, + counts_per_second_path=dataset.counts_per_second_path, + channel_names_path=dataset.channel_names_path, + ), + progress=iteration, + ) diff --git a/src/ptychodus/model/fluorescence/repository.py b/src/ptychodus/model/fluorescence/repository.py new file mode 100644 index 000000000..a27198f33 --- /dev/null +++ b/src/ptychodus/model/fluorescence/repository.py @@ -0,0 +1,306 @@ +from __future__ import annotations +from abc import ABC, abstractmethod +from collections.abc import Sequence +from enum import Enum +from pathlib import Path +from typing import overload +import logging + +from ptychodus.api.common import RealArrayType +from ptychodus.api.fluorescence import FluorescenceDataset + +from ..product import ProductRepository, ProductRepositoryItem, ProductRepositoryObserver +from ..product.metadata import MetadataRepositoryItem +from ..product.object import ObjectRepositoryItem +from ..product.probe import ProbeRepositoryItem +from ..product.probe_positions import ProbePositionsRepositoryItem +from ptychodus.api.product import LossValue + +logger = logging.getLogger(__name__) + + +class FluorescenceItemState(Enum): + READY = 'ready' + ENHANCING = 'enhancing' + FAILED = 'failed' + ORPHANED = 'orphaned' + + +class FluorescenceRepositoryItemObserver(ABC): + @abstractmethod + def handle_metadata_changed(self, item: FluorescenceRepositoryItem) -> None: + pass + + @abstractmethod + def handle_enhanced_changed(self, item: FluorescenceRepositoryItem) -> None: + pass + + @abstractmethod + def handle_state_changed(self, item: FluorescenceRepositoryItem) -> None: + pass + + +class FluorescenceRepositoryItem: + def __init__( + self, + parent: FluorescenceRepositoryItemObserver, + *, + label: str, + product: ProductRepositoryItem, + measured: FluorescenceDataset, + source_path: Path | None = None, + source_file_type: str | None = None, + ) -> None: + self._parent = parent + self._label = label + self._product = product + self._measured = measured + self._source_path = source_path + self._source_file_type = source_file_type + self._enhanced: FluorescenceDataset | None = None + self._measured_summary_cache: RealArrayType | None = None + self._enhanced_summary_cache: RealArrayType | None = None + self._state = FluorescenceItemState.READY + self._index = -1 # used by FluorescenceRepository + + def get_label(self) -> str: + return self._label + + def set_label(self, label: str) -> None: + if self._label != label: + self._label = label + self._parent.handle_metadata_changed(self) + + def get_product(self) -> ProductRepositoryItem: + return self._product + + def get_source_path(self) -> Path | None: + return self._source_path + + def get_source_file_type(self) -> str | None: + return self._source_file_type + + def get_measured(self) -> FluorescenceDataset: + return self._measured + + def get_enhanced(self) -> FluorescenceDataset | None: + return self._enhanced + + def set_enhanced(self, dataset: FluorescenceDataset) -> None: + self._enhanced = dataset + # Enhanced was just (re)set — drop the cached summary so the next read + # recomputes from the fresh element maps. + self._enhanced_summary_cache = None + self._parent.handle_enhanced_changed(self) + + @staticmethod + def _sum_element_maps(dataset: FluorescenceDataset) -> RealArrayType | None: + maps = dataset.element_maps + if not maps: + return None + total = maps[0].counts_per_second.copy() + for element_map in maps[1:]: + total += element_map.counts_per_second + return total + + def get_measured_summary(self) -> RealArrayType | None: + """Elementwise sum across the measured element maps (cached, immutable).""" + if self._measured_summary_cache is None: + self._measured_summary_cache = self._sum_element_maps(self._measured) + return self._measured_summary_cache + + def get_enhanced_summary(self) -> RealArrayType | None: + """Elementwise sum across the enhanced element maps, or None if not enhanced.""" + if self._enhanced is None: + return None + if self._enhanced_summary_cache is None: + self._enhanced_summary_cache = self._sum_element_maps(self._enhanced) + return self._enhanced_summary_cache + + def get_state(self) -> FluorescenceItemState: + return self._state + + def set_state(self, state: FluorescenceItemState) -> None: + if self._state is not state: + self._state = state + self._parent.handle_state_changed(self) + + def mark_orphaned(self) -> None: + # Idempotent; once orphaned, an item cannot un-orphan even if the same + # product were re-added — the stored reference is to the old instance. + if self._state is not FluorescenceItemState.ORPHANED: + self._state = FluorescenceItemState.ORPHANED + self._parent.handle_state_changed(self) + + +class FluorescenceRepositoryObserver(ABC): + @abstractmethod + def handle_item_inserted(self, index: int, item: FluorescenceRepositoryItem) -> None: + pass + + @abstractmethod + def handle_item_removed(self, index: int, item: FluorescenceRepositoryItem) -> None: + pass + + @abstractmethod + def handle_metadata_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + pass + + @abstractmethod + def handle_enhanced_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + pass + + @abstractmethod + def handle_state_changed(self, index: int, item: FluorescenceRepositoryItem) -> None: + pass + + +class _ProductRemovalAdapter(ProductRepositoryObserver): + """Bridges ProductRepositoryObserver → FluorescenceRepository._on_product_removed. + + The only signal we care about is product removal; all other callbacks are + no-ops. Kept separate so FluorescenceRepository doesn't have to implement + every method of the ProductRepositoryObserver ABC just to catch removals. + """ + + def __init__(self, repository: FluorescenceRepository) -> None: + super().__init__() + self._repository = repository + + def handle_item_inserted(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_metadata_changed(self, index: int, item: MetadataRepositoryItem) -> None: + pass + + def handle_probe_positions_changed( + self, index: int, item: ProbePositionsRepositoryItem + ) -> None: + pass + + def handle_probe_changed(self, index: int, item: ProbeRepositoryItem) -> None: + pass + + def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: + pass + + def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: + pass + + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: + self._repository._on_product_removed(item) + + +class FluorescenceRepository( + Sequence[FluorescenceRepositoryItem], FluorescenceRepositoryItemObserver +): + def __init__(self, product_repository: ProductRepository) -> None: + super().__init__() + self._product_repository = product_repository + self._item_list: list[FluorescenceRepositoryItem] = [] + self._observer_list: list[FluorescenceRepositoryObserver] = [] + # Adapter kept as an instance attribute so the observer registration + # survives for the lifetime of the repository. + self._product_removal_adapter = _ProductRemovalAdapter(self) + product_repository.add_observer(self._product_removal_adapter) + + @overload + def __getitem__(self, index: int) -> FluorescenceRepositoryItem: ... + + @overload + def __getitem__(self, index: slice) -> Sequence[FluorescenceRepositoryItem]: ... + + def __getitem__( + self, index: int | slice + ) -> FluorescenceRepositoryItem | Sequence[FluorescenceRepositoryItem]: + return self._item_list[index] + + def __len__(self) -> int: + return len(self._item_list) + + def create_unique_name(self, candidate_name: str) -> str: + reserved_names = {item.get_label() for item in self._item_list} + name = candidate_name or 'Unnamed' + match = 0 + + while name in reserved_names: + match += 1 + name = f'{candidate_name}-{match}' + + return name + + def _update_indexes(self) -> None: + for index, item in enumerate(self._item_list): + item._index = index + + def insert_item(self, item: FluorescenceRepositoryItem) -> int: + index = len(self._item_list) + self._item_list.append(item) + self._update_indexes() + + for observer in self._observer_list: + observer.handle_item_inserted(index, item) + + return index + + def remove_item(self, index: int) -> None: + try: + item = self._item_list.pop(index) + except IndexError: + logger.debug(f'Failed to remove fluorescence item {index}!') + return + + self._update_indexes() + + for observer in self._observer_list: + observer.handle_item_removed(index, item) + + def _on_product_removed(self, removed_product: ProductRepositoryItem) -> None: + """Mark every fluorescence item bound to the removed product as orphaned. + + Items retain their strong reference to the (now-detached) product, so + rendering and saving keep working; only re-enhancement is blocked. + """ + for item in self._item_list: + if item.get_product() is removed_product: + item.mark_orphaned() + + def add_observer(self, observer: FluorescenceRepositoryObserver) -> None: + if observer not in self._observer_list: + self._observer_list.append(observer) + + def remove_observer(self, observer: FluorescenceRepositoryObserver) -> None: + try: + self._observer_list.remove(observer) + except ValueError: + pass + + def handle_metadata_changed(self, item: FluorescenceRepositoryItem) -> None: + index = item._index + if index < 0: + logger.warning(f'Failed to look up index for "{item.get_label()}"!') + return + for observer in self._observer_list: + observer.handle_metadata_changed(index, item) + + def handle_enhanced_changed(self, item: FluorescenceRepositoryItem) -> None: + index = item._index + if index < 0: + logger.warning(f'Failed to look up index for "{item.get_label()}"!') + return + for observer in self._observer_list: + observer.handle_enhanced_changed(index, item) + + def handle_state_changed(self, item: FluorescenceRepositoryItem) -> None: + index = item._index + if index < 0: + logger.warning(f'Failed to look up index for "{item.get_label()}"!') + return + for observer in self._observer_list: + observer.handle_state_changed(index, item) diff --git a/src/ptychodus/model/fluorescence/settings.py b/src/ptychodus/model/fluorescence/settings.py index b67249d99..7a784a3e3 100644 --- a/src/ptychodus/model/fluorescence/settings.py +++ b/src/ptychodus/model/fluorescence/settings.py @@ -21,7 +21,7 @@ def __init__(self, registry: SettingsRegistry) -> None: ) self.upscaling_strategy = self._group.create_string_parameter('UpscalingStrategy', 'Linear') self.deconvolution_strategy = self._group.create_string_parameter( - 'DeconvolutionStrategy', 'Richardson-Lucy' + 'DeconvolutionStrategy', 'RichardsonLucy' ) # Ptychozoon (GPU VSPI) settings diff --git a/src/ptychodus/model/fluorescence/two_step.py b/src/ptychodus/model/fluorescence/two_step.py index deb931604..0cd6e6ecd 100644 --- a/src/ptychodus/model/fluorescence/two_step.py +++ b/src/ptychodus/model/fluorescence/two_step.py @@ -13,11 +13,8 @@ FluorescenceEnhancerOutput, UpscalingStrategy, ) -from ptychodus.api.observer import Observable, Observer from ptychodus.api.plugins import PluginChooser -from .settings import FluorescenceSettings - logger = logging.getLogger(__name__) __all__ = [ @@ -25,13 +22,12 @@ ] -class TwoStepFluorescenceEnhancer(FluorescenceEnhancer, Observable, Observer): +class TwoStepFluorescenceEnhancer(FluorescenceEnhancer): SIMPLE_NAME: Final[str] = 'TwoStep' DISPLAY_NAME: Final[str] = 'Upscale and Deconvolve' def __init__( self, - settings: FluorescenceSettings, upscaling_strategy_chooser: PluginChooser[UpscalingStrategy], deconvolution_strategy_chooser: PluginChooser[DeconvolutionStrategy], ) -> None: @@ -39,12 +35,6 @@ def __init__( self._upscaling_strategy_chooser = upscaling_strategy_chooser self._deconvolution_strategy_chooser = deconvolution_strategy_chooser - upscaling_strategy_chooser.synchronize_with_parameter(settings.upscaling_strategy) - upscaling_strategy_chooser.add_observer(self) - - deconvolution_strategy_chooser.synchronize_with_parameter(settings.deconvolution_strategy) - deconvolution_strategy_chooser.add_observer(self) - @property def name(self) -> str: return self.DISPLAY_NAME @@ -78,29 +68,3 @@ def enhance( ), progress=len(element_maps), ) - - def get_upscaling_strategies(self) -> Iterator[str]: - for plugin in self._upscaling_strategy_chooser: - yield plugin.display_name - - def get_upscaling_strategy(self) -> str: - return self._upscaling_strategy_chooser.get_current_plugin().display_name - - def set_upscaling_strategy(self, name: str) -> None: - self._upscaling_strategy_chooser.set_current_plugin(name) - - def get_deconvolution_strategies(self) -> Iterator[str]: - for plugin in self._deconvolution_strategy_chooser: - yield plugin.display_name - - def get_deconvolution_strategy(self) -> str: - return self._deconvolution_strategy_chooser.get_current_plugin().display_name - - def set_deconvolution_strategy(self, name: str) -> None: - self._deconvolution_strategy_chooser.set_current_plugin(name) - - def _update(self, observable: Observable) -> None: - if observable is self._upscaling_strategy_chooser: - self.notify_observers() - elif observable is self._deconvolution_strategy_chooser: - self.notify_observers() diff --git a/src/ptychodus/model/fluorescence/vspi.py b/src/ptychodus/model/fluorescence/vspi.py index ef993e961..690cf6e2a 100644 --- a/src/ptychodus/model/fluorescence/vspi.py +++ b/src/ptychodus/model/fluorescence/vspi.py @@ -16,7 +16,6 @@ FluorescenceEnhancerOutput, ) from ptychodus.api.object import ObjectPosition -from ptychodus.api.observer import Observable, Observer from ptychodus.api.product import Product from .settings import FluorescenceSettings @@ -120,7 +119,7 @@ def _rmatvec(self, x: RealArrayType) -> RealArrayType: # noqa: N803 return HX -class VSPIFluorescenceEnhancer(FluorescenceEnhancer, Observable, Observer): +class VSPIFluorescenceEnhancer(FluorescenceEnhancer): SIMPLE_NAME: Final[str] = 'VSPI' DISPLAY_NAME: Final[str] = 'Virtual Single Pixel Imaging' @@ -128,9 +127,6 @@ def __init__(self, settings: FluorescenceSettings) -> None: super().__init__() self._settings = settings - settings.vspi_damping_factor.add_observer(self) - settings.vspi_max_iterations.add_observer(self) - @property def name(self) -> str: return self.DISPLAY_NAME @@ -176,21 +172,3 @@ def enhance( ), progress=len(element_maps), ) - - def get_damping_factor(self) -> float: - return self._settings.vspi_damping_factor.get_value() - - def set_damping_factor(self, factor: float) -> None: - self._settings.vspi_damping_factor.set_value(factor) - - def get_max_iterations(self) -> int: - return self._settings.vspi_max_iterations.get_value() - - def set_max_iterations(self, number: int) -> None: - self._settings.vspi_max_iterations.set_value(number) - - def _update(self, observable: Observable) -> None: - if observable is self._settings.vspi_damping_factor: - self.notify_observers() - elif observable is self._settings.vspi_max_iterations: - self.notify_observers() diff --git a/src/ptychodus/model/genesis/core.py b/src/ptychodus/model/genesis/core.py index 7f3672642..299e33527 100644 --- a/src/ptychodus/model/genesis/core.py +++ b/src/ptychodus/model/genesis/core.py @@ -3,10 +3,9 @@ import queue import threading -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.settings import SettingsRegistry -from ..diffraction import DiffractionAPI from ..processing import ProcessingAPI from ..product import ProductAPI from ..task_manager import TaskManager @@ -168,7 +167,6 @@ def __init__( self, task_manager: TaskManager, settings_registry: SettingsRegistry, - diffraction_api: DiffractionAPI, product_api: ProductAPI, processing_api: ProcessingAPI, ) -> None: @@ -184,12 +182,22 @@ def __init__( for name, provider in create_globus_transfer_providers().items(): self._transfer_client_chooser.register_plugin(provider, display_name=name) + # Both choosers are empty when the IRI tokens file is absent. Binding an empty + # chooser is harmless but would log a spurious "invalid plugin name" warning, + # so stay quiet in that case. + self.facility_parameter: PluginChooserParameter[IRIFacilityAdapter] | None = None + self.transfer_client_parameter: PluginChooserParameter[AmSCGlobusTransferClient] | None = ( + None + ) + if self._facility_chooser: - self._facility_chooser.synchronize_with_parameter(self.settings.facility) + self.facility_parameter = PluginChooserParameter( + self._facility_chooser, self.settings.facility + ) if self._transfer_client_chooser: - self._transfer_client_chooser.synchronize_with_parameter( - self.settings.globus_transfer_provider + self.transfer_client_parameter = PluginChooserParameter( + self._transfer_client_chooser, self.settings.globus_transfer_provider ) status_q: queue.Queue[GenesisStatus] = queue.Queue() @@ -197,7 +205,6 @@ def __init__( self.executor = GenesisExecutor( task_manager, settings_registry, - diffraction_api, product_api, processing_api, self.settings, diff --git a/src/ptychodus/model/genesis/executor.py b/src/ptychodus/model/genesis/executor.py index e7dd207a6..55c1544f4 100644 --- a/src/ptychodus/model/genesis/executor.py +++ b/src/ptychodus/model/genesis/executor.py @@ -9,7 +9,6 @@ from ptychodus.api.plugins import PluginChooser from ptychodus.api.settings import SettingsRegistry -from ..diffraction import DiffractionAPI from ..processing import ProcessingAPI from ..product import ProductAPI from ..task_manager import BackgroundTaskManager @@ -59,7 +58,6 @@ def __init__( self, task_manager: BackgroundTaskManager, settings_registry: SettingsRegistry, - diffraction_api: DiffractionAPI, product_api: ProductAPI, processing_api: ProcessingAPI, settings: GenesisSettings, @@ -72,7 +70,6 @@ def __init__( self._task_manager = task_manager self._settings = settings self._settings_registry = settings_registry - self._diffraction_api = diffraction_api self._product_api = product_api self._processing_api = processing_api self._facility_chooser = facility_chooser @@ -87,6 +84,13 @@ def populate_input_directory(self, input_product_index: int) -> WorkflowDirector logger.exception(f'Failed access product for flow ({input_product_index=})!') raise + dataset = product_item.get_dataset() + + if dataset is None: + raise RuntimeError( + f'Product "{product_item.get_name()}" has no associated diffraction dataset.' + ) + local_dir_struct = WorkflowDirectoryStructure( self._settings.local_collection_posix_path.get_value() / product_item.get_name() ) @@ -100,7 +104,7 @@ def populate_input_directory(self, input_product_index: int) -> WorkflowDirector self._settings_registry.save_settings( local_dir_struct.input_directory / StandardFileLayout.SETTINGS ) - self._diffraction_api.export_assembled_patterns( + dataset.export_assembled_patterns( local_dir_struct.input_directory / StandardFileLayout.DIFFRACTION ) self._product_api.save_product( diff --git a/src/ptychodus/model/genesis/iri/account.py b/src/ptychodus/model/genesis/iri/account.py index 3cc76fff0..c40652598 100644 --- a/src/ptychodus/model/genesis/iri/account.py +++ b/src/ptychodus/model/genesis/iri/account.py @@ -3,7 +3,7 @@ import logging from pydantic import BaseModel -import requests +import httpx from ..tokens import create_headers @@ -23,21 +23,21 @@ class IRIAccountClient: # See https://api.iri.nersc.gov/#/account def __init__(self, api_base_url: str, access_token: str) -> None: - self._base_url = api_base_url.rstrip('/') + '/api/v1/account' - self._headers = create_headers(access_token) + self._client = httpx.Client( + base_url=api_base_url.rstrip('/') + '/api/v1/account', + headers=create_headers(access_token), + timeout=30.0, + ) def get_projects(self) -> Sequence[Project]: - response = requests.get( - f'{self._base_url}/projects', - headers=self._headers, - ) + response = self._client.get('/projects') response.raise_for_status() return [Project.model_validate(item) for item in response.json()] def get_project(self, project_id: str) -> Project: - response = requests.get( - f'{self._base_url}/projects/{project_id}', - headers=self._headers, - ) + response = self._client.get(f'/projects/{project_id}') response.raise_for_status() return Project.model_validate(response.json()) + + def close(self) -> None: + self._client.close() diff --git a/src/ptychodus/model/genesis/iri/client.py b/src/ptychodus/model/genesis/iri/client.py index f8c2fe711..5e715863b 100644 --- a/src/ptychodus/model/genesis/iri/client.py +++ b/src/ptychodus/model/genesis/iri/client.py @@ -1,7 +1,7 @@ from pathlib import Path import json -import requests +import httpx from ptychodus.api.common import get_ptychodus_dir @@ -25,7 +25,7 @@ def get_api_base_url(self) -> str: return self._api_base_url def print_openapi_specification(self) -> None: - response = requests.get(f'{self._api_base_url}/openapi.json') + response = httpx.get(f'{self._api_base_url}/openapi.json', timeout=30.0) response.raise_for_status() print(json.dumps(response.json(), indent=2)) diff --git a/src/ptychodus/model/genesis/iri/compute.py b/src/ptychodus/model/genesis/iri/compute.py index 60e589006..b29984681 100644 --- a/src/ptychodus/model/genesis/iri/compute.py +++ b/src/ptychodus/model/genesis/iri/compute.py @@ -4,7 +4,7 @@ import logging from pydantic import BaseModel, ConfigDict, Field, StrictBool -import requests +import httpx from ..tokens import create_headers @@ -88,23 +88,24 @@ class IRIComputeClient: # See https://api.iri.nersc.gov/#/compute def __init__(self, api_base_url: str, access_token: str) -> None: - self._base_url = api_base_url.rstrip('/') + '/api/v1/compute' - self._headers = create_headers(access_token) + self._client = httpx.Client( + base_url=api_base_url.rstrip('/') + '/api/v1/compute', + headers=create_headers(access_token), + timeout=30.0, + ) def submit_job(self, resource_id: str, spec: JobSpecification) -> JobResponse: - response = requests.post( - f'{self._base_url}/job/{resource_id}', + response = self._client.post( + f'/job/{resource_id}', json=spec.model_dump(mode='json'), - headers=self._headers, ) response.raise_for_status() return JobResponse.model_validate(response.json()) def update_job(self, resource_id: str, job_id: str, spec: JobSpecification) -> JobResponse: - response = requests.put( - f'{self._base_url}/job/{resource_id}/{job_id}', + response = self._client.put( + f'/job/{resource_id}/{job_id}', json=spec.model_dump(mode='json'), - headers=self._headers, ) response.raise_for_status() return JobResponse.model_validate(response.json()) @@ -116,10 +117,9 @@ def get_job_status( historical: bool = False, include_spec: bool = False, ) -> JobResponse: - response = requests.get( - f'{self._base_url}/status/{resource_id}/{job_id}', + response = self._client.get( + f'/status/{resource_id}/{job_id}', params={'historical': historical, 'include_spec': include_spec}, - headers=self._headers, ) response.raise_for_status() return JobResponse.model_validate(response.json()) @@ -132,22 +132,22 @@ def get_job_statuses( historical: bool = False, include_spec: bool = False, ) -> Sequence[JobResponse]: - response = requests.post( - f'{self._base_url}/status/{resource_id}', + response = self._client.post( + f'/status/{resource_id}', params={ 'offset': offset, 'limit': limit, 'historical': historical, 'include_spec': include_spec, }, - headers=self._headers, ) response.raise_for_status() return [JobResponse.model_validate(item) for item in response.json()] def cancel_job(self, resource_id: str, job_id: str) -> bool: - response = requests.delete( - f'{self._base_url}/cancel/{resource_id}/{job_id}', headers=self._headers - ) + response = self._client.delete(f'/cancel/{resource_id}/{job_id}') response.raise_for_status() return response.status_code == 204 + + def close(self) -> None: + self._client.close() diff --git a/src/ptychodus/model/genesis/iri/facility.py b/src/ptychodus/model/genesis/iri/facility.py index 38e29e322..fdffd9cc4 100644 --- a/src/ptychodus/model/genesis/iri/facility.py +++ b/src/ptychodus/model/genesis/iri/facility.py @@ -3,7 +3,7 @@ import logging from pydantic import BaseModel -import requests +import httpx from ..tokens import create_headers @@ -45,8 +45,11 @@ class IRIFacilityClient: # See https://api.iri.nersc.gov/#/facility def __init__(self, api_base_url: str, access_token: str) -> None: - self._base_url = api_base_url.rstrip('/') + '/api/v1/facility' - self._headers = create_headers(access_token) + self._client = httpx.Client( + base_url=api_base_url.rstrip('/') + '/api/v1/facility', + headers=create_headers(access_token), + timeout=30.0, + ) def get_facility(self, modified_since: datetime | None = None) -> Facility: params: dict = {} @@ -54,11 +57,7 @@ def get_facility(self, modified_since: datetime | None = None) -> Facility: if modified_since is not None: params['modified_since'] = modified_since.isoformat() - response = requests.get( - self._base_url, - params=params, - headers=self._headers, - ) + response = self._client.get('', params=params) response.raise_for_status() return Facility.model_validate(response.json()) @@ -81,11 +80,7 @@ def get_sites( if short_name is not None: params['short_name'] = short_name - response = requests.get( - f'{self._base_url}/sites', - params=params, - headers=self._headers, - ) + response = self._client.get('/sites', params=params) response.raise_for_status() return [Site.model_validate(item) for item in response.json()] @@ -95,10 +90,9 @@ def get_site(self, site_id: str, modified_since: datetime | None = None) -> Site if modified_since is not None: params['modified_since'] = modified_since.isoformat() - response = requests.get( - f'{self._base_url}/sites/{site_id}', - params=params, - headers=self._headers, - ) + response = self._client.get(f'/sites/{site_id}', params=params) response.raise_for_status() return Site.model_validate(response.json()) + + def close(self) -> None: + self._client.close() diff --git a/src/ptychodus/model/genesis/iri/status.py b/src/ptychodus/model/genesis/iri/status.py index 36d4f6c00..154d453b4 100644 --- a/src/ptychodus/model/genesis/iri/status.py +++ b/src/ptychodus/model/genesis/iri/status.py @@ -4,7 +4,7 @@ import logging from pydantic import BaseModel -import requests +import httpx from ..tokens import create_headers @@ -51,8 +51,11 @@ class IRIStatusClient: # See https://api.iri.nersc.gov/#/status def __init__(self, api_base_url: str, access_token: str) -> None: - self._base_url = api_base_url.rstrip('/') + '/api/v1/status' - self._headers = create_headers(access_token) + self._client = httpx.Client( + base_url=api_base_url.rstrip('/') + '/api/v1/status', + headers=create_headers(access_token), + timeout=30.0, + ) def get_resources( self, @@ -81,18 +84,14 @@ def get_resources( params['current_status'] = current_status if capability is not None: params['capability'] = [c.value for c in capability] - response = requests.get( - f'{self._base_url}/resources', - params=params, - headers=self._headers, - ) + response = self._client.get('/resources', params=params) response.raise_for_status() return [Resource.model_validate(item) for item in response.json()] def get_resource(self, resource_id: str) -> Resource: - response = requests.get( - f'{self._base_url}/resources/{resource_id}', - headers=self._headers, - ) + response = self._client.get(f'/resources/{resource_id}') response.raise_for_status() return Resource.model_validate(response.json()) + + def close(self) -> None: + self._client.close() diff --git a/src/ptychodus/model/genesis/tasks.py b/src/ptychodus/model/genesis/tasks.py index 85f92cbd7..be451f52b 100644 --- a/src/ptychodus/model/genesis/tasks.py +++ b/src/ptychodus/model/genesis/tasks.py @@ -5,7 +5,7 @@ import threading import time -import requests +import httpx from .iri import IRIComputeClient, JobSpecification, JobState from .transfer import AmSCGlobusTransferClient, GlobusTransferInputs, TransferStatus @@ -98,7 +98,7 @@ def compute_task( try: job_response = client.submit_job(resource_id, spec) - except requests.HTTPError as exc: + except httpx.HTTPStatusError as exc: if exc.response.status_code == 404: if attempt == max_retries: yield GenesisStatus( @@ -135,7 +135,7 @@ def compute_task( while not stop_event.is_set(): try: response = client.get_job_status(resource_id, job_id) - except requests.HTTPError as exc: + except httpx.HTTPStatusError as exc: if exc.response.status_code == 400: logger.warning('Invalid request parameters.') break diff --git a/src/ptychodus/model/genesis/transfer.py b/src/ptychodus/model/genesis/transfer.py index 5595267a9..bad41944f 100644 --- a/src/ptychodus/model/genesis/transfer.py +++ b/src/ptychodus/model/genesis/transfer.py @@ -6,7 +6,7 @@ import json from pydantic import BaseModel -import requests +import httpx from ptychodus.api.common import get_ptychodus_dir @@ -47,43 +47,43 @@ class AmSCGlobusTransferClient: def __init__(self, api_base_url: str, access_token: str) -> None: self._api_base_url = api_base_url.rstrip('/') - self._headers = create_headers(access_token) + self._client = httpx.Client( + base_url=self._api_base_url, + headers=create_headers(access_token), + timeout=30.0, + ) def check_auth_token(self) -> Mapping[str, Any]: - response = requests.get(f'{self._api_base_url}/movement/auth/globus', headers=self._headers) + response = self._client.get('/movement/auth/globus') response.raise_for_status() return response.json() def start_transfer(self, inputs: GlobusTransferInputs) -> GlobusTransferResult: - response = requests.post( - f'{self._api_base_url}/movement/transfer/globus', + response = self._client.post( + '/movement/transfer/globus', json=inputs.model_dump(mode='json'), - headers=self._headers, ) response.raise_for_status() return GlobusTransferResult.model_validate(response.json()) def get_transfer(self, transfer_id: str) -> GlobusTransferResult: - response = requests.get( - f'{self._api_base_url}/movement/transfer/globus/{transfer_id}', - headers=self._headers, - ) + response = self._client.get(f'/movement/transfer/globus/{transfer_id}') response.raise_for_status() return GlobusTransferResult.model_validate(response.json()) def delete_transfer(self, transfer_id: str) -> GlobusTransferResult: - response = requests.delete( - f'{self._api_base_url}/movement/transfer/globus/{transfer_id}', - headers=self._headers, - ) + response = self._client.delete(f'/movement/transfer/globus/{transfer_id}') response.raise_for_status() return GlobusTransferResult.model_validate(response.json()) def print_openapi_specification(self) -> None: - response = requests.get(f'{self._api_base_url}/openapi.json') + response = httpx.get(f'{self._api_base_url}/openapi.json', timeout=30.0) response.raise_for_status() print(json.dumps(response.json(), indent=2)) + def close(self) -> None: + self._client.close() + def get_amsc_transfer_api_url() -> str: return 'https://amsc-data-api.nersc.gov' diff --git a/src/ptychodus/model/globus/core.py b/src/ptychodus/model/globus/core.py index 3ae5719f6..e6fe3318c 100644 --- a/src/ptychodus/model/globus/core.py +++ b/src/ptychodus/model/globus/core.py @@ -4,7 +4,6 @@ from ptychodus.api.settings import SettingsRegistry -from ..diffraction import DiffractionAPI from ..processing import ProcessingAPI from ..product import ProductAPI from .authorizer import GlobusAuthorizer @@ -20,7 +19,6 @@ class GlobusCore: def __init__( self, settings_registry: SettingsRegistry, - diffraction_api: DiffractionAPI, product_api: ProductAPI, processing_api: ProcessingAPI, ) -> None: @@ -42,7 +40,6 @@ def __init__( self.executor = GlobusExecutor( self.settings, settings_registry, - diffraction_api, product_api, processing_api, self._client, diff --git a/src/ptychodus/model/globus/executor.py b/src/ptychodus/model/globus/executor.py index 81fd8fb03..1aaf13430 100644 --- a/src/ptychodus/model/globus/executor.py +++ b/src/ptychodus/model/globus/executor.py @@ -4,7 +4,6 @@ from ptychodus.api.io import StandardFileLayout from ptychodus.api.settings import SettingsRegistry -from ..diffraction import DiffractionAPI from ..product import ProductAPI from ..processing import ProcessingAPI from .client import GlobusClient, GlobusJob @@ -18,7 +17,6 @@ def __init__( self, settings: GlobusSettings, settings_registry: SettingsRegistry, - diffraction_api: DiffractionAPI, product_api: ProductAPI, processing_api: ProcessingAPI, client: GlobusClient, @@ -26,7 +24,6 @@ def __init__( super().__init__() self._settings = settings self._settings_registry = settings_registry - self._diffraction_api = diffraction_api self._product_api = product_api self._processing_api = processing_api self._client = client @@ -38,6 +35,13 @@ def populate_input_directory(self, input_product_index: int) -> Path: logger.exception(f'Failed access product for flow ({input_product_index=})!') raise + dataset = product_item.get_dataset() + + if dataset is None: + raise RuntimeError( + f'Product "{product_item.get_name()}" has no associated diffraction dataset.' + ) + input_directory = ( self._settings.input_collection_posix_path.get_value() / product_item.get_name() ) @@ -49,9 +53,7 @@ def populate_input_directory(self, input_product_index: int) -> Path: raise self._settings_registry.save_settings(input_directory / StandardFileLayout.SETTINGS) - self._diffraction_api.export_assembled_patterns( - input_directory / StandardFileLayout.DIFFRACTION - ) + dataset.export_assembled_patterns(input_directory / StandardFileLayout.DIFFRACTION) self._product_api.save_product( input_product_index, input_directory / StandardFileLayout.PRODUCT_IN, diff --git a/src/ptychodus/model/metadata.py b/src/ptychodus/model/metadata.py deleted file mode 100644 index dea9ef8e6..000000000 --- a/src/ptychodus/model/metadata.py +++ /dev/null @@ -1,137 +0,0 @@ -from __future__ import annotations - -from ptychodus.api.observer import Observable -from ptychodus.api.diffraction import DiffractionMetadata - -from .diffraction import ( - DetectorSettings, - DiffractionDatasetObserver, - AssembledDiffractionDataset, - DiffractionSettings, -) -from .product import ProductSettings - - -class MetadataPresenter(Observable, DiffractionDatasetObserver): - def __init__( - self, - detector_settings: DetectorSettings, - diffraction_settings: DiffractionSettings, - dataset: AssembledDiffractionDataset, - product_settings: ProductSettings, - ) -> None: - super().__init__() - self._detector_settings = detector_settings - self._diffraction_settings = diffraction_settings - self._dataset = dataset - self._product_settings = product_settings - - dataset.add_observer(self) - - @property - def _metadata(self) -> DiffractionMetadata: - return self._dataset.get_metadata() - - def can_sync_detector_extent(self) -> bool: - return self._metadata.detector_extent is not None - - def sync_detector_extent(self) -> None: - detector_extent = self._metadata.detector_extent - - if detector_extent: - self._detector_settings.width_px.set_value(detector_extent.width_px) - self._detector_settings.height_px.set_value(detector_extent.height_px) - - def can_sync_detector_pixel_size(self) -> bool: - return self._metadata.detector_pixel_geometry is not None - - def sync_detector_pixel_size(self) -> None: - pixel_geometry = self._metadata.detector_pixel_geometry - - if pixel_geometry: - self._detector_settings.pixel_width_m.set_value(pixel_geometry.width_m) - self._detector_settings.pixel_height_m.set_value(pixel_geometry.height_m) - - def can_sync_pattern_crop_center(self) -> bool: - return self._metadata.crop_center is not None or self._metadata.detector_extent is not None - - def can_sync_pattern_crop_extent(self) -> bool: - return self._metadata.detector_extent is not None - - def sync_pattern_crop(self, sync_center: bool, sync_extent: bool) -> None: - if sync_center: - crop_center = self._metadata.crop_center - - if crop_center: - self._diffraction_settings.crop_center_x_px.set_value(crop_center.position_x_px) - self._diffraction_settings.crop_center_y_px.set_value(crop_center.position_y_px) - elif self._metadata.detector_extent: - self._diffraction_settings.crop_center_x_px.set_value( - int(self._metadata.detector_extent.width_px) // 2 - ) - self._diffraction_settings.crop_center_y_px.set_value( - int(self._metadata.detector_extent.height_px) // 2 - ) - - if sync_extent and self._metadata.detector_extent: - center_x = self._diffraction_settings.crop_center_x_px.get_value() - center_y = self._diffraction_settings.crop_center_y_px.get_value() - - extent_x = int(self._metadata.detector_extent.width_px) - extent_y = int(self._metadata.detector_extent.height_px) - - max_radius_x = min(center_x, extent_x - center_x) - max_radius_y = min(center_y, extent_y - center_y) - max_radius = min(max_radius_x, max_radius_y) - crop_diameter = 1 - - while crop_diameter < max_radius: - crop_diameter <<= 1 - - self._diffraction_settings.crop_width_px.set_value(crop_diameter) - self._diffraction_settings.crop_height_px.set_value(crop_diameter) - - def can_sync_probe_energy(self) -> bool: - return self._metadata.probe_energy_eV is not None - - def sync_probe_energy(self) -> None: - energy_eV = self._metadata.probe_energy_eV # noqa: N806 - - if energy_eV: - self._product_settings.probe_energy_eV.set_value(energy_eV) - - def can_sync_probe_photon_count(self) -> bool: - return self._metadata.probe_photon_count is not None - - def sync_probe_photon_count(self) -> None: - photon_count = self._metadata.probe_photon_count - - if photon_count: - self._product_settings.probe_photon_count.set_value(photon_count) - - def can_sync_exposure_time(self) -> bool: - return self._metadata.exposure_time_s is not None - - def sync_exposure_time(self) -> None: - exposure_time_s = self._metadata.exposure_time_s - - if exposure_time_s: - self._product_settings.exposure_time_s.set_value(exposure_time_s) - - def can_sync_detector_distance(self) -> bool: - return self._metadata.detector_distance_m is not None - - def sync_detector_distance(self) -> None: - distance_m = self._metadata.detector_distance_m - - if distance_m: - self._product_settings.detector_distance_m.set_value(distance_m) - - def handle_array_inserted(self, index: int) -> None: - pass - - def handle_array_changed(self, index: int) -> None: - pass - - def handle_dataset_reloaded(self) -> None: - self.notify_observers() diff --git a/src/ptychodus/model/processing/_subprocess_protocol.py b/src/ptychodus/model/processing/_subprocess_protocol.py new file mode 100644 index 000000000..59fda1ff7 --- /dev/null +++ b/src/ptychodus/model/processing/_subprocess_protocol.py @@ -0,0 +1,279 @@ +"""Shared machinery for running GPU-touching code in transient spawned subprocesses. + +The parent-side ptychodus process must NEVER acquire a GPU context (no +``torch.cuda.*`` / ``tf.config.experimental.*`` calls, no tensor-on-GPU +allocations, no ``ptychi.api.task.PtychographyTask`` construction, no +``tf.keras.Model`` fit, no Lightning ``Trainer`` construction). All GPU work +happens inside a freshly ``spawn``ed child process that lives only for the +duration of one call (reconstruct / train / enhance) and dies immediately +after. This gives each call a clean GPU context and fully releases GPU +memory between calls, which lets ptychodus mix reconstructor backends built +on different GPU frameworks without driver-state interference. + +The parent MAY import a GPU framework if doing so is required to construct +picklable configuration objects the child needs — e.g. importing +``ptychi.api.options.*`` or ``ptycho_torch.config_params`` pulls torch in for +its type annotations, but no CUDA runtime is initialised until a tensor is +placed on a GPU. Keep this to the minimum needed for the payload, and prefer +importing from inside the payload builder over module scope so the cost lands +on the first call rather than at startup. + +Wire protocol +------------- + +Messages travel over a ``multiprocessing.Queue`` as tagged tuples. Because a +single child thread produces every message, ordering is preserved (a log line +arrives before the result it precedes). + +- ``('log', levelno, logger_name, formatted_message)`` -- child log record +- ``('error', traceback_str, exception_type_name, pickled_exc_or_none)`` +- ``None`` -- end-of-stream sentinel (always the final message) + +Any other tag is passed through to the consumer verbatim; consumers are free +to define their own tags (e.g. ``'output'``, ``'progress'``, ``'result'``, +``'settings_sync'``, ``'model_saved'``). + +Child entry points are addressed by ``'dotted.module.path:function_name'`` +strings; the child imports the module lazily and calls +``function(payload, queue)``. The parent never imports the entry-point +module. +""" + +from __future__ import annotations + +import importlib +import logging +import multiprocessing +import pickle +import traceback +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from multiprocessing.context import SpawnProcess +from multiprocessing.queues import Queue +from typing import Any + +logger = logging.getLogger(__name__) + +__all__ = [ + 'TAG_ERROR', + 'TAG_LOG', + 'ChildError', + 'SubprocessLogHandler', + 'install_child_log_forwarder', + 'run_subprocess', + 'send_error', +] + + +TAG_LOG = 'log' +TAG_ERROR = 'error' + + +class ChildError(RuntimeError): + """Raised in the parent when the child subprocess reports an unhandled exception. + + The child's original traceback is available via ``child_traceback`` and, + when the exception class is available parent-side, the original exception + instance via ``child_exception``. + """ + + def __init__( + self, + message: str, + *, + child_traceback: str, + child_exception_type: str, + child_exception: BaseException | None, + ) -> None: + super().__init__(message) + self.child_traceback = child_traceback + self.child_exception_type = child_exception_type + self.child_exception = child_exception + + +class SubprocessLogHandler(logging.Handler): + """Logging handler that forwards formatted records to the parent over a queue. + + Attach to the child's root logger so records from any GPU-framework logger + (torch, ptychi, ptycho, ptychozoon, ...) propagate back to the parent's + logging tree unchanged. + """ + + def __init__(self, result_queue: Queue[Any]) -> None: + super().__init__() + self._result_queue = result_queue + + def emit(self, record: logging.LogRecord) -> None: + try: + self._result_queue.put((TAG_LOG, record.levelno, record.name, self.format(record))) + except Exception: + # Never let logging failures break the worker. + pass + + +def install_child_log_forwarder( + result_queue: Queue[Any], level: int = logging.INFO +) -> SubprocessLogHandler: + """Install a queue-forwarding handler on the child's root logger. + + Returns the handler so it can be removed in a ``finally`` block. The + formatter emits ``'name: message'`` because the parent-side dispatcher + already tags each record with the child's logger name. + """ + handler = SubprocessLogHandler(result_queue) + handler.setFormatter(logging.Formatter('%(message)s')) + root = logging.getLogger() + root.addHandler(handler) + root.setLevel(level) + return handler + + +def send_error(result_queue: Queue[Any], exc: BaseException) -> None: + """Marshal an unhandled child-side exception onto the queue. + + Pickles the exception when possible so the parent can re-raise the + original class; otherwise the parent falls back to :class:`ChildError` + carrying the traceback text. + """ + tb = traceback.format_exc() + try: + pickled_exc: bytes | None = pickle.dumps(exc) + except Exception: + pickled_exc = None + try: + result_queue.put((TAG_ERROR, tb, type(exc).__name__, pickled_exc)) + except Exception: + # Queue might already be closed; nothing we can do. + pass + + +def _child_main( + entry_point: str, + payload: Any, + result_queue: Queue[Any], + log_level: int, +) -> None: + """Default target for the spawned :class:`multiprocessing.Process`. + + Runs in the child. Installs the log forwarder, imports the user entry + point lazily (so GPU libraries load here, not in the parent), calls + ``entry(payload, queue)``, and always sends the sentinel. + """ + handler = install_child_log_forwarder(result_queue, level=log_level) + try: + module_path, _, func_name = entry_point.partition(':') + if not func_name: + raise ValueError( + f"Entry point {entry_point!r} must be formatted as 'module.path:function_name'" + ) + module = importlib.import_module(module_path) + func = getattr(module, func_name) + func(payload, result_queue) + except BaseException as exc: # noqa: BLE001 - want to marshal every failure + send_error(result_queue, exc) + finally: + try: + logging.getLogger().removeHandler(handler) + except Exception: + pass + try: + result_queue.put(None) + except Exception: + pass + + +@dataclass +class _RunState: + process: SpawnProcess + queue: Queue[Any] + + +@contextmanager +def run_subprocess( + entry_point: str, + payload: Any, + *, + log_level: int | None = None, + terminate_grace_sec: float = 10.0, +) -> Iterator[Iterator[tuple[Any, ...]]]: + """Spawn a child process and yield an iterator over its queue messages. + + Usage:: + + with run_subprocess('ptychodus.model.ptychi._child:run_reconstruct', payload) as events: + for event in events: + tag = event[0] + ... + + Messages tagged ``'log'`` are dispatched to the parent's logging tree by + logger name and NOT re-yielded. Messages tagged ``'error'`` raise + :class:`ChildError` (with the child's original exception attached when + unpicklable-safe). ``None`` terminates the iterator normally. + + On ``__exit__``, the child is terminated (SIGTERM), joined with + ``terminate_grace_sec`` seconds of grace, then killed (SIGKILL) if it did + not exit. This mirrors the ptychozoon shutdown discipline. + """ + if log_level is None: + log_level = logging.getLogger().getEffectiveLevel() + + ctx = multiprocessing.get_context('spawn') + result_queue: Queue[Any] = ctx.Queue() + process = ctx.Process( + target=_child_main, + args=(entry_point, payload, result_queue, log_level), + ) + process.start() + + state = _RunState(process=process, queue=result_queue) + try: + yield _iter_events(state) + finally: + _shutdown(state, terminate_grace_sec) + + +def _iter_events(state: _RunState) -> Iterator[tuple[Any, ...]]: + while True: + item = state.queue.get() + if item is None: + return + + tag = item[0] + + if tag == TAG_LOG: + _, levelno, logger_name, message = item + logging.getLogger(logger_name).log(levelno, message) + continue + + if tag == TAG_ERROR: + _, tb, type_name, pickled_exc = item + child_exc: BaseException | None = None + if pickled_exc is not None: + try: + child_exc = pickle.loads(pickled_exc) + except Exception: + child_exc = None + raise ChildError( + f'Subprocess raised {type_name}:\n{tb}', + child_traceback=tb, + child_exception_type=type_name, + child_exception=child_exc, + ) + + yield item + + +def _shutdown(state: _RunState, grace_sec: float) -> None: + process = state.process + if process.is_alive(): + process.terminate() + process.join(timeout=grace_sec) + if process.is_alive(): + process.kill() + process.join() + try: + state.queue.close() + state.queue.join_thread() + except Exception: + pass diff --git a/src/ptychodus/model/processing/api.py b/src/ptychodus/model/processing/api.py index 37c6383aa..77097872e 100644 --- a/src/ptychodus/model/processing/api.py +++ b/src/ptychodus/model/processing/api.py @@ -15,7 +15,6 @@ from ptychodus.api.observer import Observable, Observer from ptychodus.api.parametric import Parameter, StringParameter -from ..diffraction import DiffractionAPI from ..product import ProductAPI from ..task_manager import TaskManager from .monitor import ( @@ -103,13 +102,11 @@ class ProcessingAPI: def __init__( self, task_manager: TaskManager, - diffraction_api: DiffractionAPI, product_api: ProductAPI, algorithm: ProcessingAlgorithmParameter, task_monitor: ProcessingTaskMonitor, ) -> None: self._task_manager = task_manager - self._diffraction_api = diffraction_api self._product_api = product_api self._algorithm_parameter = algorithm self._task_monitor = task_monitor @@ -120,14 +117,22 @@ def task_monitor(self) -> ProcessingTaskMonitor: def get_reconstruct_input( self, - product_index: int, *, + product_index: int, index_filter: PositionIndexFilter = PositionIndexFilter.ALL, ) -> ReconstructInput: - product = self._product_api.get_item(product_index).get_product() + product_item = self._product_api.get_item(product_index) + dataset = product_item.get_dataset() + + if dataset is None: + raise RuntimeError( + f'Product "{product_item.get_name()}" has no associated diffraction dataset.' + ) + + product = product_item.get_product() logger.info(f'Preparing input data for {product.metadata.name}...') tic = time.perf_counter() - assembled_data = self._diffraction_api.get_assembled_data() + assembled_data = dataset.get_assembled_data() reconstruct_input = assembled_data.prepare_reconstruct_input( product, index_filter=index_filter ) @@ -137,8 +142,8 @@ def get_reconstruct_input( def reconstruct( self, - input_product_index: int, *, + product_index: int, algorithm: str | None = None, index_filter: PositionIndexFilter = PositionIndexFilter.ALL, output_product_suffix: str = '', @@ -146,8 +151,10 @@ def reconstruct( block: bool = False, ) -> int: self.set_reconstructor_if_provided(algorithm) - input_product_item = self._product_api.get_item(input_product_index) - output_product_index = self._product_api.insert_product(input_product_item.get_product()) + input_product_item = self._product_api.get_item(product_index) + output_product_index = self._product_api.insert_product( + input_product_item.get_product(), dataset=input_product_item.get_dataset() + ) output_product_item = self._product_api.get_item(output_product_index) output_product_name = ( f'{input_product_item.get_name()}_{self._algorithm_parameter.get_value()}' @@ -158,7 +165,8 @@ def reconstruct( output_product_item.set_name(output_product_name) reconstruct_input = self.get_reconstruct_input( - output_product_index, index_filter=index_filter + product_index=output_product_index, + index_filter=index_filter, ) background_task = ReconstructBackgroundTask( self._task_monitor, @@ -177,20 +185,21 @@ def reconstruct( ): self._task_manager.run_foreground_tasks() break + self._task_monitor.raise_if_failed() return output_product_index def reconstruct_split( - self, input_product_index: int, *, algorithm: str | None = None + self, *, product_index: int, algorithm: str | None = None ) -> tuple[int, int]: output_product_index_odd = self.reconstruct( - input_product_index, + product_index=product_index, algorithm=algorithm, index_filter=PositionIndexFilter.ODD, output_product_suffix='odd', ) output_product_index_even = self.reconstruct( - input_product_index, + product_index=product_index, algorithm=algorithm, index_filter=PositionIndexFilter.EVEN, output_product_suffix='even', @@ -227,6 +236,7 @@ def save_model_to_file(self, file_path: Path, algorithm: str | None = None) -> N def export_training_data( self, file_path: Path, + *, product_index: int, algorithm: str | None = None, index_filter: PositionIndexFilter = PositionIndexFilter.ALL, @@ -235,7 +245,10 @@ def export_training_data( trainer = self._algorithm_parameter.get_current_reconstructor() if isinstance(trainer, TrainableReconstructor): - reconstruct_input = self.get_reconstruct_input(product_index, index_filter=index_filter) + reconstruct_input = self.get_reconstruct_input( + product_index=product_index, + index_filter=index_filter, + ) logger.info('Exporting...') tic = time.perf_counter() @@ -247,10 +260,10 @@ def export_training_data( def train( self, - product_index: int, input_path: Path, output_path: Path, *, + product_index: int, algorithm: str | None = None, block: bool = False, ) -> None: @@ -258,7 +271,7 @@ def train( trainer = self._algorithm_parameter.get_current_reconstructor() if isinstance(trainer, TrainableReconstructor): - reconstruct_input = self.get_reconstruct_input(product_index) + reconstruct_input = self.get_reconstruct_input(product_index=product_index) background_task = TrainBackgroundTask( self._task_monitor, trainer, @@ -276,6 +289,7 @@ def train( ): self._task_manager.run_foreground_tasks() break + self._task_monitor.raise_if_failed() else: logger.warning('Algorithm is not trainable!') diff --git a/src/ptychodus/model/processing/core.py b/src/ptychodus/model/processing/core.py index cd9bc12af..d36e1dd1b 100644 --- a/src/ptychodus/model/processing/core.py +++ b/src/ptychodus/model/processing/core.py @@ -3,7 +3,6 @@ from ptychodus.api.reconstructor import ReconstructorLibrary from ptychodus.api.settings import SettingsRegistry -from ..diffraction import DiffractionAPI from ..product import ProductAPI from ..task_manager import TaskManager from .api import ProcessingAPI, ProcessingAlgorithmParameter @@ -16,7 +15,6 @@ def __init__( self, task_manager: TaskManager, settings_registry: SettingsRegistry, - diffraction_api: DiffractionAPI, product_api: ProductAPI, algorithm_libraries: Sequence[ReconstructorLibrary], ) -> None: @@ -27,7 +25,6 @@ def __init__( self._task_monitor = ProcessingTaskMonitor(task_manager) self.processing_api = ProcessingAPI( task_manager, - diffraction_api, product_api, self.algorithm_parameter, self._task_monitor, diff --git a/src/ptychodus/model/processing/subprocess_reconstructor.py b/src/ptychodus/model/processing/subprocess_reconstructor.py new file mode 100644 index 000000000..4ac98c048 --- /dev/null +++ b/src/ptychodus/model/processing/subprocess_reconstructor.py @@ -0,0 +1,224 @@ +"""Generic reconstructor adapter that runs the underlying backend in a subprocess. + +Wraps any backend that satisfies the ``Reconstructor`` / +``TrainableReconstructor`` shape (as a child-side entry point) into a +parent-safe object that ptychodus can dispatch to without ever importing a +GPU framework. + +The parent side of the adapter has zero GPU imports. Every +``reconstruct()`` / ``train()`` call spawns a fresh ``spawn``-context +subprocess via :mod:`._subprocess_protocol`, streams outputs back, and dies +at end-of-iteration. Consumers of the reconstructor iterator (see +``ReconstructBackgroundTask``) see the same per-iteration ``ReconstructOutput`` +they did in the in-process version. + +Trainable-model lifecycle +------------------------- + +``load_model_from_file`` on the parent-side adapter records the path only -- +the fresh inference child does the real load per call. ``save_model`` copies +or archives the loaded-from path to the destination (using a backend-supplied +callback so ptychopinn's zip-a-bundle-dir semantic and ptychopinn_torch's +copy-the-.ckpt semantic can share the same adapter). ``export_training_data`` +runs in the parent (all current implementations are pure-numpy and touch no +GPU framework). +""" + +from __future__ import annotations + +import logging +import pickle +import shutil +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Any + +from ptychodus.api.reconstructor import ( + ReconstructInput, + ReconstructOutput, + TrainableReconstructor, + TrainOutput, +) + +from ._subprocess_protocol import run_subprocess + +logger = logging.getLogger(__name__) + +__all__ = [ + 'SubprocessReconstructor', +] + + +# Tags emitted by child entry points and consumed by this adapter. +TAG_OUTPUT = 'output' # payload = pickle.dumps(ReconstructOutput) +TAG_TRAIN_OUTPUT = 'train_output' # payload = pickle.dumps(TrainOutput) +TAG_MODEL_SAVED = 'model_saved' # payload = str path where child wrote checkpoint +TAG_SETTINGS_SYNC = 'settings_sync' # payload = {group_name: {param_name: value_str}} + + +# Type aliases for the callbacks each backend supplies. +BuildReconstructPayload = Callable[[ReconstructInput, 'Path | None'], Any] +"""(reconstruct_input, loaded_model_path_or_none) -> pickled payload for the child.""" + +BuildTrainPayload = Callable[[Path, Path], Any] +"""(input_dir, output_dir) -> pickled payload for the child.""" + +ExportTrainingData = Callable[[Path, ReconstructInput], None] +"""(file_path, reconstruct_input) -> None. Runs parent-side; must not touch GPU.""" + +SaveModel = Callable[[Path, Path], None] +"""(source_path_loaded_by_child, destination_path) -> None. Runs parent-side.""" + +ApplySettingsSync = Callable[[dict[str, dict[str, str]]], None] +"""Optional handler for ('settings_sync', dict) messages the child emits.""" + + +class SubprocessReconstructor(TrainableReconstructor): + """Parent-side adapter that runs any reconstructor in a fresh subprocess. + + Implements the full :class:`TrainableReconstructor` interface. When + ``is_trainable`` is False, the trainable methods raise + :class:`NotImplementedError` so the object still slots into ``ProcessingCore``'s + list of ``Reconstructor`` instances without special-casing. + + The parent-side object holds no GPU state. It stores: + - the child entry-point strings (dotted module paths), + - backend-supplied callbacks that build payloads and export training data, + - the loaded-model path (as recorded by ``load_model_from_file`` / + remembered from the last ``train`` call). + + Each ``reconstruct()`` / ``train()`` call spawns exactly one child; the + child dies before the call returns. + """ + + def __init__( + self, + *, + name: str, + reconstruct_entry_point: str, + progress_goal_fn: Callable[[], int], + build_reconstruct_payload: BuildReconstructPayload, + is_trainable: bool = False, + train_entry_point: str | None = None, + build_train_payload: BuildTrainPayload | None = None, + model_file_filter: str = '', + model_file_extension: str = '', + training_data_file_filter: str = '', + export_training_data: ExportTrainingData | None = None, + save_model: SaveModel | None = None, + apply_settings_sync: ApplySettingsSync | None = None, + terminate_grace_sec: float = 10.0, + ) -> None: + super().__init__() + self._name = name + self._reconstruct_entry_point = reconstruct_entry_point + self._progress_goal_fn = progress_goal_fn + self._build_reconstruct_payload = build_reconstruct_payload + self._is_trainable = is_trainable + self._train_entry_point = train_entry_point + self._build_train_payload = build_train_payload + self._model_file_filter = model_file_filter + self._model_file_extension = model_file_extension + self._training_data_file_filter = training_data_file_filter + self._export_training_data = export_training_data + self._save_model_fn = save_model + self._apply_settings_sync = apply_settings_sync + self._terminate_grace_sec = terminate_grace_sec + + self._loaded_model_path: Path | None = None + + @property + def name(self) -> str: + return self._name + + @property + def is_trainable(self) -> bool: + return self._is_trainable + + def get_progress_goal(self) -> int: + return self._progress_goal_fn() + + def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: + payload = self._build_reconstruct_payload(parameters, self._loaded_model_path) + + with run_subprocess( + self._reconstruct_entry_point, + payload, + terminate_grace_sec=self._terminate_grace_sec, + ) as events: + for event in events: + tag = event[0] + if tag == TAG_OUTPUT: + yield pickle.loads(event[1]) + elif tag == TAG_SETTINGS_SYNC: + if self._apply_settings_sync is not None: + try: + self._apply_settings_sync(event[1]) + except Exception: + logger.exception('Failed to apply settings sync from subprocess.') + else: + logger.debug( + f'{self._name}: dropping unrecognized subprocess message tag {tag!r}' + ) + + def is_model_loaded(self) -> bool: + return self._loaded_model_path is not None + + def get_model_file_filter(self) -> str: + return self._model_file_filter + + def load_model_from_file(self, file_path: Path) -> None: + # Parent-side: just remember the path. The child does the actual load + # (which touches GPU frameworks) per reconstruct/train call. + self._loaded_model_path = file_path + + def get_model_file_extension(self) -> str: + return self._model_file_extension + + def save_model(self, file_path: Path) -> None: + if self._loaded_model_path is None: + raise RuntimeError( + f'Cannot save {self._name} model: no model has been loaded or trained.' + ) + if self._save_model_fn is not None: + self._save_model_fn(self._loaded_model_path, file_path) + else: + shutil.copyfile(self._loaded_model_path, file_path) + + def get_training_data_file_filter(self) -> str: + return self._training_data_file_filter + + def export_training_data(self, file_path: Path, parameters: ReconstructInput) -> None: + if self._export_training_data is None: + raise NotImplementedError(f'{self._name} does not support exporting training data.') + self._export_training_data(file_path, parameters) + + def train(self, input_path: Path, output_path: Path) -> Iterator[TrainOutput]: + if not self._is_trainable: + raise NotImplementedError(f'{self._name} does not support training.') + if self._train_entry_point is None or self._build_train_payload is None: + raise NotImplementedError(f'{self._name} does not support training.') + + payload = self._build_train_payload(input_path, output_path) + + with run_subprocess( + self._train_entry_point, + payload, + terminate_grace_sec=self._terminate_grace_sec, + ) as events: + for event in events: + tag = event[0] + if tag == TAG_TRAIN_OUTPUT: + yield pickle.loads(event[1]) + elif tag == TAG_MODEL_SAVED: + self._loaded_model_path = Path(event[1]) + elif tag == TAG_SETTINGS_SYNC: + if self._apply_settings_sync is not None: + try: + self._apply_settings_sync(event[1]) + except Exception: + logger.exception('Failed to apply settings sync from subprocess.') + else: + logger.debug( + f'{self._name}: dropping unrecognized subprocess message tag {tag!r}' + ) diff --git a/src/ptychodus/model/product/api.py b/src/ptychodus/model/product/api.py index 3386ddd39..83f19de19 100644 --- a/src/ptychodus/model/product/api.py +++ b/src/ptychodus/model/product/api.py @@ -1,12 +1,15 @@ -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from pathlib import Path from typing import Any import logging +from ptychodus.api.diffraction import Polarization from ptychodus.api.plugins import PluginChooser from ptychodus.api.product import Product, ProductFileReader, ProductFileWriter -from .item import ProductRepositoryItem +from ..diffraction import AssembledDiffractionDataset +from ..task_manager import TaskManager +from .item import ProductRepositoryItem, ProductState from .item_factory import ProductRepositoryItemFactory from .object.builder_factory import ObjectBuilderFactory from .object.settings import ObjectSettings @@ -176,7 +179,10 @@ def builder_names(self) -> Iterator[str]: return iter(self._builder_factory) def build_probe( - self, index: int, builder_name: str, builder_parameters: Mapping[str, Any] | None = None + self, + index: int, + builder_name: str, + builder_parameters: Mapping[str, Any] | None = None, ) -> None: try: item = self._repository[index] @@ -184,10 +190,12 @@ def build_probe( logger.warning(f'Failed to access item {index}!') return + dataset = self._repository.get_dataset(index) + try: - builder = self._builder_factory.create(builder_name) - except KeyError: - logger.warning(f'Failed to create builder {builder_name}!') + builder = self._builder_factory.create(builder_name, dataset=dataset) + except (KeyError, RuntimeError) as exc: + logger.warning(f'Failed to create builder {builder_name}: {exc}') return if builder_parameters is not None: @@ -211,10 +219,12 @@ def build_probe_from_settings(self, index: int) -> None: logger.warning(f'Failed to access item {index}!') return + dataset = self._repository.get_dataset(index) + try: - builder = self._builder_factory.create_from_settings() - except KeyError: - logger.warning('Failed to create builder from settings!') + builder = self._builder_factory.create_from_settings(dataset=dataset) + except (KeyError, RuntimeError) as exc: + logger.warning(f'Failed to create builder from settings: {exc}') return item.set_builder(builder) @@ -285,7 +295,10 @@ def builder_names(self) -> Iterator[str]: return iter(self._builder_factory) def build_object( - self, index: int, builder_name: str, builder_parameters: Mapping[str, Any] | None = None + self, + index: int, + builder_name: str, + builder_parameters: Mapping[str, Any] | None = None, ) -> None: try: item = self._repository[index] @@ -293,10 +306,12 @@ def build_object( logger.warning(f'Failed to access item {index}!') return + dataset = self._repository.get_dataset(index) + try: - builder = self._builder_factory.create(builder_name) - except KeyError: - logger.warning(f'Failed to create builder {builder_name}!') + builder = self._builder_factory.create(builder_name, dataset=dataset) + except (KeyError, RuntimeError) as exc: + logger.warning(f'Failed to create builder {builder_name}: {exc}') return if builder_parameters is not None: @@ -320,10 +335,12 @@ def build_object_from_settings(self, index: int) -> None: logger.warning(f'Failed to access item {index}!') return + dataset = self._repository.get_dataset(index) + try: - builder = self._builder_factory.create_from_settings() - except KeyError: - logger.warning('Failed to create builder from settings!') + builder = self._builder_factory.create_from_settings(dataset=dataset) + except (KeyError, RuntimeError) as exc: + logger.warning(f'Failed to create builder from settings: {exc}') return item.set_builder(builder) @@ -387,12 +404,72 @@ def __init__( item_factory: ProductRepositoryItemFactory, file_reader_chooser: PluginChooser[ProductFileReader], file_writer_chooser: PluginChooser[ProductFileWriter], + task_manager: TaskManager, ) -> None: self._settings = settings self._repository = repository self._item_factory = item_factory self._file_reader_chooser = file_reader_chooser self._file_writer_chooser = file_writer_chooser + self._task_manager = task_manager + + def _insert_via_queue( + self, + dataset: AssembledDiffractionDataset | None, + build: Callable[[], ProductRepositoryItem], + *, + block: bool, + stub_name: str, + ) -> int: + """If dataset is None or already loaded, run build() synchronously and insert. + Otherwise, insert a pending stub, then enqueue construction to run after the + dataset's LoadAllArrays finishes (via the shared FIFO background worker).""" + if dataset is None or not dataset.is_load_in_progress(): + item = build() + return self._repository.insert_product(item) + + finished_event = dataset.get_last_load_finished_event() + + if block: + if finished_event is not None: + while not self._task_manager.is_stopping: + if finished_event.wait(timeout=TaskManager.WAIT_TIME_S): + break + error = dataset.get_last_load_error() + if error is not None: + raise RuntimeError('Diffraction dataset failed to load') from error + item = build() + return self._repository.insert_product(item) + + stub = self._item_factory.create_pending_stub(name=stub_name) + index = self._repository.insert_product(stub) + + def background_finalize() -> Callable[[], None]: + error = dataset.get_last_load_error() + + def foreground_finalize() -> None: + if error is not None: + logger.error( + f'Cancelling queued product {index} because dataset ' + f'failed to load: {error!r}' + ) + stub.set_state(ProductState.FAILED) + return + + try: + real = build() + except Exception: + logger.exception(f'Queued product {index} construction failed') + stub.set_state(ProductState.FAILED) + return + + stub.copy_contents_from(real) + stub.set_state(ProductState.READY) + + return foreground_finalize + + self._task_manager.put_background_task(background_finalize) + return index def insert_new_product( self, @@ -405,26 +482,50 @@ def insert_new_product( exposure_time_s: float | None = None, mass_attenuation_m2_kg: float | None = None, tomography_angle_deg: float | None = None, + tilt_angle_deg: float | None = None, + polarization: Polarization | None = None, + dataset: AssembledDiffractionDataset | None = None, + block: bool = True, ) -> int: - item = self._item_factory.create_from_values( - name=name, - comments=comments, - detector_distance_m=detector_distance_m, - probe_energy_eV=probe_energy_eV, - probe_photon_count=probe_photon_count, - exposure_time_s=exposure_time_s, - mass_attenuation_m2_kg=mass_attenuation_m2_kg, - tomography_angle_deg=tomography_angle_deg, - ) - return self._repository.insert_product(item) + def build() -> ProductRepositoryItem: + return self._item_factory.create_from_values( + name=name, + comments=comments, + detector_distance_m=detector_distance_m, + probe_energy_eV=probe_energy_eV, + probe_photon_count=probe_photon_count, + exposure_time_s=exposure_time_s, + mass_attenuation_m2_kg=mass_attenuation_m2_kg, + tomography_angle_deg=tomography_angle_deg, + tilt_angle_deg=tilt_angle_deg, + polarization=polarization, + dataset=dataset, + ) + + return self._insert_via_queue(dataset, build, block=block, stub_name=name) + + def insert_product( + self, + product: Product, + *, + dataset: AssembledDiffractionDataset | None = None, + block: bool = True, + ) -> int: + def build() -> ProductRepositoryItem: + return self._item_factory.create_from_product(product, dataset=dataset) - def insert_product(self, product: Product) -> int: - item = self._item_factory.create_from_product(product) - return self._repository.insert_product(item) + return self._insert_via_queue(dataset, build, block=block, stub_name=product.metadata.name) + + def insert_product_from_settings( + self, + *, + dataset: AssembledDiffractionDataset | None = None, + block: bool = True, + ) -> int: + def build() -> ProductRepositoryItem: + return self._item_factory.create_from_settings(dataset=dataset) - def insert_product_from_settings(self) -> int: - item = self._item_factory.create_from_settings() - return self._repository.insert_product(item) + return self._insert_via_queue(dataset, build, block=block, stub_name='Unnamed') def get_item(self, product_index: int) -> ProductRepositoryItem: return self._repository[product_index] @@ -436,7 +537,14 @@ def get_open_file_filters(self) -> Iterator[str]: def get_open_file_filter(self) -> str: return self._file_reader_chooser.get_current_plugin().display_name - def open_product(self, file_path: Path, *, file_type: str | None = None) -> int: + def open_product( + self, + file_path: Path, + *, + file_type: str | None = None, + dataset: AssembledDiffractionDataset | None = None, + block: bool = True, + ) -> int: if file_path.is_file(): if file_type is not None: self._file_reader_chooser.set_current_plugin(file_type) @@ -450,7 +558,7 @@ def open_product(self, file_path: Path, *, file_type: str | None = None) -> int: except Exception as exc: raise RuntimeError(f'Failed to read "{file_path}"') from exc else: - return self.insert_product(product) + return self.insert_product(product, dataset=dataset, block=block) else: logger.warning(f'Refusing to create product with invalid file path "{file_path}"') diff --git a/src/ptychodus/model/product/core.py b/src/ptychodus/model/product/core.py index 8dbcfc45e..3e2ae76f4 100644 --- a/src/ptychodus/model/product/core.py +++ b/src/ptychodus/model/product/core.py @@ -2,14 +2,20 @@ from ptychodus.api.object import ObjectFileReader, ObjectFileWriter from ptychodus.api.observer import Observable, Observer -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.probe import ProbeFileReader, ProbeFileWriter from ptychodus.api.probe_gen import FresnelZonePlate from ptychodus.api.probe_positions import ProbePositionFileReader, ProbePositionFileWriter from ptychodus.api.product import ProductFileReader, ProductFileWriter from ptychodus.api.settings import SettingsRegistry -from ..diffraction import AssembledDiffractionDataset, DiffractionAPI, PatternSizer +from ..diffraction import ( + AssembledDiffractionDataset, + DiffractionAPI, + DiffractionDatasetRepositoryObserver, + PatternSizer, +) +from ..task_manager import TaskManager from .api import ObjectAPI, ProbeAPI, ProductAPI, ProbePositionsAPI from .item_factory import ProductRepositoryItemFactory from .object import ObjectBuilderFactory, ObjectRepositoryItemFactory, ObjectSettings @@ -26,6 +32,21 @@ from .settings import ProductSettings +class _DatasetOrphanObserver(DiffractionDatasetRepositoryObserver): + """Clears product references to a diffraction dataset when it leaves the repository.""" + + def __init__(self, product_repository: ProductRepository) -> None: + self._product_repository = product_repository + + def handle_dataset_inserted(self, index: int, dataset: AssembledDiffractionDataset) -> None: + pass + + def handle_dataset_removed(self, index: int, dataset: AssembledDiffractionDataset) -> None: + for item in self._product_repository: + if item.get_dataset() is dataset: + item.unbind_dataset() + + class ProductCore(Observer): def __init__( self, @@ -33,7 +54,6 @@ def __init__( settings_registry: SettingsRegistry, pattern_sizer: PatternSizer, diffraction_api: DiffractionAPI, - diffraction_dataset: AssembledDiffractionDataset, scan_file_reader_chooser: PluginChooser[ProbePositionFileReader], scan_file_writer_chooser: PluginChooser[ProbePositionFileWriter], fresnel_zone_plate_chooser: PluginChooser[FresnelZonePlate], @@ -44,6 +64,7 @@ def __init__( product_file_reader_chooser: PluginChooser[ProductFileReader], product_file_writer_chooser: PluginChooser[ProductFileWriter], reinit_observable: Observable, + task_manager: TaskManager, ) -> None: super().__init__() self.settings = ProductSettings(settings_registry) @@ -60,7 +81,6 @@ def __init__( self._probe_builder_factory = ProbeBuilderFactory( rng, self._probe_settings, - diffraction_api, fresnel_zone_plate_chooser, probe_file_reader_chooser, probe_file_writer_chooser, @@ -73,7 +93,6 @@ def __init__( self._object_builder_factory = ObjectBuilderFactory( rng, self._object_settings, - diffraction_api, object_file_reader_chooser, object_file_writer_chooser, ) @@ -85,7 +104,6 @@ def __init__( self._item_factory = ProductRepositoryItemFactory( self.settings, pattern_sizer, - diffraction_dataset, self._scan_repository_item_factory, self._probe_repository_item_factory, self._object_repository_item_factory, @@ -98,7 +116,11 @@ def __init__( self._item_factory, product_file_reader_chooser, product_file_writer_chooser, + task_manager, ) + self._diffraction_api = diffraction_api + self._dataset_orphan_observer = _DatasetOrphanObserver(self.product_repository) + diffraction_api.get_repository().add_observer(self._dataset_orphan_observer) self.probe_positions_repository = ProbePositionsRepository(self.product_repository) self.probe_positions_api = ProbePositionsAPI( self._scan_settings, self.probe_positions_repository, self._scan_builder_factory @@ -113,13 +135,21 @@ def __init__( ) # TODO vvv refactor vvv - product_file_reader_chooser.synchronize_with_parameter(self.settings.file_type) + self.product_file_reader_parameter = PluginChooserParameter( + product_file_reader_chooser, self.settings.file_type + ) product_file_writer_chooser.set_current_plugin(self.settings.file_type.get_value()) - scan_file_reader_chooser.synchronize_with_parameter(self._scan_settings.file_type) + self.scan_file_reader_parameter = PluginChooserParameter( + scan_file_reader_chooser, self._scan_settings.file_type + ) scan_file_writer_chooser.set_current_plugin(self._scan_settings.file_type.get_value()) - probe_file_reader_chooser.synchronize_with_parameter(self._probe_settings.file_type) + self.probe_file_reader_parameter = PluginChooserParameter( + probe_file_reader_chooser, self._probe_settings.file_type + ) probe_file_writer_chooser.set_current_plugin(self._probe_settings.file_type.get_value()) - object_file_reader_chooser.synchronize_with_parameter(self._object_settings.file_type) + self.object_file_reader_parameter = PluginChooserParameter( + object_file_reader_chooser, self._object_settings.file_type + ) object_file_writer_chooser.set_current_plugin(self._object_settings.file_type.get_value()) # TODO ^^^^^^^^^^^^^^^^ @@ -128,4 +158,9 @@ def __init__( def _update(self, observable: Observable) -> None: if observable is self._reinit_observable: - self.product_api.insert_product_from_settings() + # Depends on DiffractionCore being registered as a reinit observer + # before ProductCore, so that open_patterns has already inserted the + # settings-driven dataset into the repository by the time we run. + repo = self._diffraction_api.get_repository() + dataset = repo[-1] if len(repo) > 0 else None + self.product_api.insert_product_from_settings(dataset=dataset, block=False) diff --git a/src/ptychodus/model/product/geometry.py b/src/ptychodus/model/product/geometry.py index d1002d9da..08419bbd5 100644 --- a/src/ptychodus/model/product/geometry.py +++ b/src/ptychodus/model/product/geometry.py @@ -2,7 +2,7 @@ import numpy -from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.geometry import ImageExtent, PixelGeometry from ptychodus.api.object import ObjectGeometry, ObjectGeometryProvider from ptychodus.api.observer import Observable, Observer from ptychodus.api.probe import ProbeGeometry, ProbeGeometryProvider @@ -29,11 +29,28 @@ def __init__( self._pattern_sizer = pattern_sizer self._metadata_item = metadata_item self._scan_item = scan_item + # Set via set_detector_extent()/set_detector_pixel_geometry() when a dataset + # is bound (see ProductRepositoryItem.bind_dataset / unbind_dataset). Derived + # quantities that need these degenerate to zero-sized while unbound. + self._detector_extent: ImageExtent | None = None + self._raw_pixel_geometry: PixelGeometry | None = None self._pattern_sizer.add_observer(self) self._metadata_item.add_observer(self) self._scan_item.add_observer(self) + def set_detector_extent(self, extent: ImageExtent | None) -> None: + if extent == self._detector_extent: + return + self._detector_extent = extent + self.notify_observers() + + def set_detector_pixel_geometry(self, geometry: PixelGeometry | None) -> None: + if geometry == self._raw_pixel_geometry: + return + self._raw_pixel_geometry = geometry + self.notify_observers() + @property def probe_photon_count(self) -> float: return self._metadata_item.probe_photon_count.get_value() @@ -84,35 +101,51 @@ def detector_distance_m(self) -> float: def _lambda_z_m2(self) -> float: return self.probe_wavelength_m * self.detector_distance_m - @property - def object_plane_pixel_width_m(self) -> float: - return self._lambda_z_m2 / self._pattern_sizer.get_processed_width_m() - - @property - def object_plane_pixel_height_m(self) -> float: - return self._lambda_z_m2 / self._pattern_sizer.get_processed_height_m() + def _processed_pixel_geometry(self) -> PixelGeometry: + # No dataset bound yet: degrade to a zero-sized geometry so downstream + # divisions bail out gracefully (they already handle ZeroDivisionError). + raw = self._raw_pixel_geometry + if raw is None: + return PixelGeometry(width_m=0.0, height_m=0.0) + return self._pattern_sizer.get_processed_pixel_geometry(raw) def get_detector_pixel_geometry(self): - return self._pattern_sizer.get_processed_pixel_geometry() + return self._processed_pixel_geometry() def get_object_plane_pixel_geometry(self) -> PixelGeometry: - return PixelGeometry( - width_m=self.object_plane_pixel_width_m, - height_m=self.object_plane_pixel_height_m, - ) + extent = self._pattern_sizer.get_processed_image_extent(self._detector_extent) + detector_pixel_geometry = self._processed_pixel_geometry() + lambda_z = self._lambda_z_m2 + try: + return PixelGeometry( + width_m=lambda_z / (extent.width_px * detector_pixel_geometry.width_m), + height_m=lambda_z / (extent.height_px * detector_pixel_geometry.height_m), + ) + except ZeroDivisionError: + return PixelGeometry(width_m=0.0, height_m=0.0) @property def fresnel_number(self) -> float: - width_m = self._pattern_sizer.get_processed_width_m() - height_m = self._pattern_sizer.get_processed_height_m() + extent = self._pattern_sizer.get_processed_image_extent(self._detector_extent) + pixel_geometry = self._processed_pixel_geometry() + width_m = extent.width_px * pixel_geometry.width_m + height_m = extent.height_px * pixel_geometry.height_m area_m2 = width_m * height_m - return area_m2 / self._lambda_z_m2 + try: + return area_m2 / self._lambda_z_m2 + except ZeroDivisionError: + return 0.0 @property def _detector_numerical_aperture_sq(self) -> float: - two_z_m = 2 * self.detector_distance_m - NA_x = self._pattern_sizer.get_processed_width_m() / two_z_m # noqa: N806 - NA_y = self._pattern_sizer.get_processed_height_m() / two_z_m # noqa: N806 + extent = self._pattern_sizer.get_processed_image_extent(self._detector_extent) + pixel_geometry = self._processed_pixel_geometry() + try: + two_z_m = 2 * self.detector_distance_m + NA_x = (extent.width_px * pixel_geometry.width_m) / two_z_m # noqa: N806 + NA_y = (extent.height_px * pixel_geometry.height_m) / two_z_m # noqa: N806 + except ZeroDivisionError: + return 0.0 return NA_x * NA_y @property @@ -124,19 +157,20 @@ def depth_of_field_m(self) -> float: return self.probe_wavelength_m / self._detector_numerical_aperture_sq def get_probe_geometry(self) -> ProbeGeometry: - extent = self._pattern_sizer.get_processed_image_extent() + extent = self._pattern_sizer.get_processed_image_extent(self._detector_extent) + pixel_geometry = self.get_object_plane_pixel_geometry() return ProbeGeometry( width_px=extent.width_px, height_px=extent.height_px, - pixel_width_m=self.object_plane_pixel_width_m, - pixel_height_m=self.object_plane_pixel_height_m, + pixel_width_m=pixel_geometry.width_m, + pixel_height_m=pixel_geometry.height_m, ) def is_probe_geometry_valid(self, geometry: ProbeGeometry) -> bool: expected = self.get_probe_geometry() - width_is_valid = geometry.pixel_width_m > 0.0 and geometry.width_m == expected.width_m - height_is_valid = geometry.pixel_height_m > 0.0 and geometry.height_m == expected.height_m - return width_is_valid and height_is_valid + if not geometry.get_pixel_geometry().is_valid: + return False + return geometry.width_m == expected.width_m and geometry.height_m == expected.height_m def get_probe_positions(self) -> Sequence[ProbePosition]: return self._scan_item.get_probe_positions() @@ -156,25 +190,26 @@ def get_object_geometry(self) -> ObjectGeometry: center_x_m = scan_bbox.center_x_m center_y_m = scan_bbox.center_y_m - pixel_width_m = self.object_plane_pixel_width_m - width_px = width_m / pixel_width_m if pixel_width_m > 0.0 else 0.0 - - pixel_height_m = self.object_plane_pixel_height_m - height_px = height_m / pixel_height_m if pixel_height_m > 0.0 else 0.0 + pixel_geometry = self.get_object_plane_pixel_geometry() + if pixel_geometry.is_valid: + width_px = width_m / pixel_geometry.width_m + height_px = height_m / pixel_geometry.height_m + else: + width_px = 0.0 + height_px = 0.0 return ObjectGeometry( width_px=int(numpy.ceil(width_px)), height_px=int(numpy.ceil(height_px)), - pixel_width_m=self.object_plane_pixel_width_m, - pixel_height_m=self.object_plane_pixel_height_m, + pixel_width_m=pixel_geometry.width_m, + pixel_height_m=pixel_geometry.height_m, center_x_m=center_x_m, center_y_m=center_y_m, ) def is_object_geometry_valid(self, geometry: ObjectGeometry) -> bool: expected_geometry = self.get_object_geometry() - pixel_size_is_valid = geometry.pixel_width_m > 0.0 and geometry.pixel_height_m > 0.0 - return pixel_size_is_valid and geometry.contains(expected_geometry) + return geometry.get_pixel_geometry().is_valid and geometry.contains(expected_geometry) def _update(self, observable: Observable) -> None: if observable is self._metadata_item: diff --git a/src/ptychodus/model/product/item.py b/src/ptychodus/model/product/item.py index d508064f2..aa410485b 100644 --- a/src/ptychodus/model/product/item.py +++ b/src/ptychodus/model/product/item.py @@ -1,12 +1,14 @@ from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import Sequence +from enum import Enum import logging from ptychodus.api.observer import Observable from ptychodus.api.parametric import ParameterGroup from ptychodus.api.product import LossValue, Product +from ..diffraction import AssembledDiffractionDataset, DiffractionDatasetObserver from .geometry import ProductGeometry from .metadata import MetadataRepositoryItem, UniqueNameFactory from .object import ObjectRepositoryItem @@ -16,6 +18,12 @@ logger = logging.getLogger(__name__) +class ProductState(Enum): + READY = 'ready' + PENDING = 'pending' + FAILED = 'failed' + + class ProductRepositoryItemObserver(UniqueNameFactory): @abstractmethod def handle_metadata_changed(self, item: ProductRepositoryItem) -> None: @@ -37,6 +45,14 @@ def handle_object_changed(self, item: ProductRepositoryItem) -> None: def handle_losses_changed(self, item: ProductRepositoryItem) -> None: pass + @abstractmethod + def handle_dataset_changed(self, item: ProductRepositoryItem) -> None: + pass + + @abstractmethod + def handle_state_changed(self, item: ProductRepositoryItem) -> None: + pass + class ProductRepositoryItem(ParameterGroup): def __init__( @@ -48,6 +64,8 @@ def __init__( probe_item: ProbeRepositoryItem, object_item: ObjectRepositoryItem, losses: Sequence[LossValue], + dataset: AssembledDiffractionDataset | None = None, + state: ProductState = ProductState.READY, ) -> None: super().__init__() self._parent = parent @@ -57,6 +75,9 @@ def __init__( self._probe_item = probe_item self._object_item = object_item self._losses = list(losses) + self._dataset: AssembledDiffractionDataset | None = None + self._dataset_observer: _BoundDatasetObserver | None = None + self._state: ProductState = state self._add_group('metadata', self._metadata_item, observe=True) self._add_group('probe_positions', self._probe_positions_item, observe=True) @@ -65,6 +86,12 @@ def __init__( self._index = -1 # used by ProductRepository + # Bind the geometry's detector extent + pixel geometry to the initial + # dataset (if any) so downstream probe/object sizes are correct from the + # start. + if dataset is not None: + self._bind_dataset(dataset) + def assign(self, product: Product) -> None: self._metadata_item.assign(product.metadata) self._probe_positions_item.assign(product.probe_positions) @@ -73,6 +100,30 @@ def assign(self, product: Product) -> None: self._losses = list(product.losses) self._parent.handle_losses_changed(self) + def copy_contents_from( + self, + source: ProductRepositoryItem, + ) -> None: + """Copy inner state from a freshly-built source item into this stub. + + Uses each subgroup's assign_item so the stub's subgroup identities are + preserved (peripheral scan/probe/object repositories continue to observe + the same subgroup instances they registered at insert time). The item's + index in the ProductRepository never changes. + """ + self._metadata_item.assign(source._metadata_item.get_metadata()) + # Bind the dataset (and thus the detector extent on the geometry) BEFORE + # rebuilding probe/object subgroups — their _rebuild() otherwise sees an + # invalid pixel geometry and silently no-ops, leaving them empty. + # _insert_via_queue only routes here when source has a bound dataset. + assert source._dataset is not None + self.bind_dataset(source._dataset) + self._probe_positions_item.assign_item(source._probe_positions_item) + self._probe_item.assign_item(source._probe_item) + self._object_item.assign_item(source._object_item) + self._losses = list(source._losses) + self._parent.handle_losses_changed(self) + def sync_to_settings(self) -> None: self._metadata_item.sync_to_settings() self._probe_positions_item.sync_to_settings() @@ -100,6 +151,64 @@ def get_probe_item(self) -> ProbeRepositoryItem: def get_object_item(self) -> ObjectRepositoryItem: return self._object_item + def get_dataset(self) -> AssembledDiffractionDataset | None: + """Return the diffraction dataset this product is associated with, if any. + + Model-only reference; not part of the persisted Product. + """ + return self._dataset + + def bind_dataset(self, dataset: AssembledDiffractionDataset) -> None: + if self._dataset is not dataset: + self._bind_dataset(dataset) + self._parent.handle_dataset_changed(self) + + def unbind_dataset(self) -> None: + if self._dataset is not None: + self._bind_dataset(None) + self._parent.handle_dataset_changed(self) + + def _bind_dataset(self, dataset: AssembledDiffractionDataset | None) -> None: + # Detach the previous dataset's observer before rebinding. + if self._dataset is not None and self._dataset_observer is not None: + self._dataset.remove_observer(self._dataset_observer) + self._dataset_observer = None + + self._dataset = dataset + + if dataset is None: + self._geometry.set_detector_extent(None) + self._geometry.set_detector_pixel_geometry(None) + return + + self._geometry.set_detector_extent(dataset.get_metadata().detector_extent) + self._geometry.set_detector_pixel_geometry(dataset.get_raw_pixel_geometry()) + + # Mirror future edits (pixel geometry, reload) from the dataset back into + # the geometry so probe/object sizes stay in sync. + self._dataset_observer = _BoundDatasetObserver(self) + dataset.add_observer(self._dataset_observer) + + def _sync_geometry_from_dataset(self) -> None: + if self._dataset is None: + return + self._geometry.set_detector_extent(self._dataset.get_metadata().detector_extent) + self._geometry.set_detector_pixel_geometry(self._dataset.get_raw_pixel_geometry()) + + def get_state(self) -> ProductState: + return self._state + + def is_pending(self) -> bool: + return self._state is ProductState.PENDING + + def is_failed(self) -> bool: + return self._state is ProductState.FAILED + + def set_state(self, state: ProductState) -> None: + if self._state != state: + self._state = state + self._parent.handle_state_changed(self) + def _invalidate_losses(self) -> None: self._losses = list() self._parent.handle_losses_changed(self) @@ -160,6 +269,34 @@ def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: pass + @abstractmethod + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + @abstractmethod + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + @abstractmethod def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: pass + + +class _BoundDatasetObserver(DiffractionDatasetObserver): + """Mirrors dataset changes back into the bound product's geometry.""" + + def __init__(self, item: ProductRepositoryItem) -> None: + super().__init__() + self._item = item + + def handle_array_inserted(self, index: int) -> None: + pass + + def handle_array_changed(self, index: int) -> None: + pass + + def handle_dataset_reloaded(self) -> None: + self._item._sync_geometry_from_dataset() + + def handle_pixel_geometry_changed(self) -> None: + self._item._sync_geometry_from_dataset() diff --git a/src/ptychodus/model/product/item_factory.py b/src/ptychodus/model/product/item_factory.py index 9c7dae102..54aa386ef 100644 --- a/src/ptychodus/model/product/item_factory.py +++ b/src/ptychodus/model/product/item_factory.py @@ -1,11 +1,12 @@ import logging +from ptychodus.api.diffraction import Polarization from ptychodus.api.plugins import PluginChooser from ptychodus.api.product import Product, ProductFileReader from ..diffraction import AssembledDiffractionDataset, PatternSizer from .geometry import ProductGeometry -from .item import ProductRepositoryItem +from .item import ProductRepositoryItem, ProductState from .metadata import MetadataRepositoryItem from .object import ObjectRepositoryItemFactory from .probe import ProbeRepositoryItemFactory @@ -21,7 +22,6 @@ def __init__( self, settings: ProductSettings, pattern_sizer: PatternSizer, - dataset: AssembledDiffractionDataset, scan_item_factory: ProbePositionsRepositoryItemFactory, probe_item_factory: ProbeRepositoryItemFactory, object_item_factory: ObjectRepositoryItemFactory, @@ -31,13 +31,27 @@ def __init__( super().__init__() self._settings = settings self._pattern_sizer = pattern_sizer - self._dataset = dataset self._scan_item_factory = scan_item_factory self._probe_item_factory = probe_item_factory self._object_item_factory = object_item_factory self._repository = repository self._file_reader_chooser = file_reader_chooser + @staticmethod + def _bind_dataset_geometry( + geometry: ProductGeometry, dataset: AssembledDiffractionDataset | None + ) -> None: + """Push the dataset's detector extent and raw pixel geometry into ``geometry`` + so probe & object items built next see a valid geometry inside their own + __init__ rebuild. This keeps the observer-triggered rebuild that fires later + (from ProductRepositoryItem._bind_dataset) a no-op — the setters short-circuit + on unchanged values, avoiding a spurious rebuild during ProductRepositoryItem + construction (which would fire index<0 warnings from the repository).""" + if dataset is None: + return + geometry.set_detector_extent(dataset.get_metadata().detector_extent) + geometry.set_detector_pixel_geometry(dataset.get_raw_pixel_geometry()) + def create_from_values( self, *, @@ -49,6 +63,9 @@ def create_from_values( exposure_time_s: float | None = None, mass_attenuation_m2_kg: float | None = None, tomography_angle_deg: float | None = None, + tilt_angle_deg: float | None = None, + polarization: Polarization | None = None, + dataset: AssembledDiffractionDataset | None = None, ) -> ProductRepositoryItem: metadata_item = MetadataRepositoryItem( self._settings, @@ -61,15 +78,17 @@ def create_from_values( exposure_time_s=exposure_time_s, mass_attenuation_m2_kg=mass_attenuation_m2_kg, tomography_angle_deg=tomography_angle_deg, + tilt_angle_deg=tilt_angle_deg, + polarization=polarization, ) - if metadata_item.probe_photon_count.get_value() <= 0: - assembled_data = self._dataset.get_assembled_data() - max_pattern_counts = assembled_data.get_pattern_counts().max() - metadata_item.probe_photon_count.set_value(max_pattern_counts) + # probe_photon_count auto-estimation from diffraction data now lives in the + # controller layer (see ProductEditorViewController._estimate_probe_photon_count). + # This factory takes whatever value the caller supplied. scan_item = self._scan_item_factory.create() geometry = ProductGeometry(self._pattern_sizer, metadata_item, scan_item) + self._bind_dataset_geometry(geometry, dataset) probe_item = self._probe_item_factory.create(geometry) object_item = self._object_item_factory.create(geometry) @@ -81,9 +100,12 @@ def create_from_values( probe_item=probe_item, object_item=object_item, losses=list(), + dataset=dataset, ) - def create_from_product(self, product: Product) -> ProductRepositoryItem: + def create_from_product( + self, product: Product, *, dataset: AssembledDiffractionDataset | None = None + ) -> ProductRepositoryItem: metadata_item = MetadataRepositoryItem( self._settings, self._repository, @@ -95,10 +117,13 @@ def create_from_product(self, product: Product) -> ProductRepositoryItem: exposure_time_s=product.metadata.exposure_time_s, mass_attenuation_m2_kg=product.metadata.mass_attenuation_m2_kg, tomography_angle_deg=product.metadata.tomography_angle_deg, + tilt_angle_deg=product.metadata.tilt_angle_deg, + polarization=product.metadata.polarization, ) scan_item = self._scan_item_factory.create(product.probe_positions) geometry = ProductGeometry(self._pattern_sizer, metadata_item, scan_item) + self._bind_dataset_geometry(geometry, dataset) probe_item = self._probe_item_factory.create(geometry, product.probes) object_item = self._object_item_factory.create(geometry, product.object_) @@ -110,9 +135,34 @@ def create_from_product(self, product: Product) -> ProductRepositoryItem: probe_item=probe_item, object_item=object_item, losses=product.losses, + dataset=dataset, + ) + + def create_pending_stub(self, name: str = 'Unnamed') -> ProductRepositoryItem: + """Build a fresh ProductRepositoryItem in the 'pending' state with default + subgroups and no dataset. Its inner content is replaced later via + ProductRepositoryItem.copy_contents_from once the source dataset finishes + loading.""" + metadata_item = MetadataRepositoryItem(self._settings, self._repository, name=name) + scan_item = self._scan_item_factory.create() + geometry = ProductGeometry(self._pattern_sizer, metadata_item, scan_item) + probe_item = self._probe_item_factory.create(geometry) + object_item = self._object_item_factory.create(geometry) + return ProductRepositoryItem( + parent=self._repository, + metadata_item=metadata_item, + probe_positions_item=scan_item, + geometry=geometry, + probe_item=probe_item, + object_item=object_item, + losses=list(), + dataset=None, + state=ProductState.PENDING, ) - def create_from_settings(self) -> ProductRepositoryItem: + def create_from_settings( + self, *, dataset: AssembledDiffractionDataset | None = None + ) -> ProductRepositoryItem: file_path = self._settings.file_path.get_value() if file_path.is_file(): @@ -125,13 +175,14 @@ def create_from_settings(self) -> ProductRepositoryItem: except Exception as exc: raise RuntimeError(f'Failed to read "{file_path}"') from exc else: - return self.create_from_product(product) + return self.create_from_product(product, dataset=dataset) metadata_item = MetadataRepositoryItem(self._settings, self._repository) scan_item = self._scan_item_factory.create_from_settings() geometry = ProductGeometry(self._pattern_sizer, metadata_item, scan_item) - probe_item = self._probe_item_factory.create_from_settings(geometry) - object_item = self._object_item_factory.create_from_settings(geometry) + self._bind_dataset_geometry(geometry, dataset) + probe_item = self._probe_item_factory.create_from_settings(geometry, dataset=dataset) + object_item = self._object_item_factory.create_from_settings(geometry, dataset=dataset) item = ProductRepositoryItem( parent=self._repository, @@ -141,6 +192,7 @@ def create_from_settings(self) -> ProductRepositoryItem: probe_item=probe_item, object_item=object_item, losses=list(), + dataset=dataset, ) logger.debug(f'Created product from settings: {item.get_name()}') return item diff --git a/src/ptychodus/model/product/metadata.py b/src/ptychodus/model/product/metadata.py index 9ca47871f..3bfabc7a3 100644 --- a/src/ptychodus/model/product/metadata.py +++ b/src/ptychodus/model/product/metadata.py @@ -2,6 +2,7 @@ from abc import abstractmethod, ABC import logging +from ptychodus.api.diffraction import Polarization from ptychodus.api.parametric import Parameter, ParameterGroup from ptychodus.api.product import ProductMetadata @@ -61,6 +62,8 @@ def __init__( exposure_time_s: float | None = None, mass_attenuation_m2_kg: float | None = None, tomography_angle_deg: float | None = None, + tilt_angle_deg: float | None = None, + polarization: Polarization | None = None, ) -> None: super().__init__() self._settings = settings @@ -112,6 +115,22 @@ def __init__( self._add_parameter('tomography_angle_deg', self.tomography_angle_deg) + self.tilt_angle_deg = settings.tilt_angle_deg.copy() + + if tilt_angle_deg is not None: + self.tilt_angle_deg.set_value(tilt_angle_deg) + + self._add_parameter('tilt_angle_deg', self.tilt_angle_deg) + + # Polarization is stored as a string parameter for INI round-trip; + # empty string means "unset" (i.e. None on the ProductMetadata side). + self.polarization = settings.polarization.copy() + + if polarization is not None: + self.polarization.set_value(polarization.value) + + self._add_parameter('polarization', self.polarization) + def assign(self, metadata: ProductMetadata) -> None: self.name.set_value(metadata.name) self.comments.set_value(metadata.comments) @@ -121,6 +140,10 @@ def assign(self, metadata: ProductMetadata) -> None: self.exposure_time_s.set_value(metadata.exposure_time_s) self.mass_attenuation_m2_kg.set_value(metadata.mass_attenuation_m2_kg) self.tomography_angle_deg.set_value(metadata.tomography_angle_deg) + self.tilt_angle_deg.set_value(metadata.tilt_angle_deg) + self.polarization.set_value( + metadata.polarization.value if metadata.polarization is not None else '' + ) def sync_to_settings(self) -> None: for parameter in self.parameters().values(): @@ -136,4 +159,16 @@ def get_metadata(self) -> ProductMetadata: exposure_time_s=self.exposure_time_s.get_value(), mass_attenuation_m2_kg=self.mass_attenuation_m2_kg.get_value(), tomography_angle_deg=self.tomography_angle_deg.get_value(), + tilt_angle_deg=self.tilt_angle_deg.get_value(), + polarization=self._parse_polarization(self.polarization.get_value()), ) + + @staticmethod + def _parse_polarization(raw: str) -> Polarization | None: + if not raw: + return None + try: + return Polarization(raw) + except ValueError: + logger.warning('Unknown polarization value in settings: %r; treating as None.', raw) + return None diff --git a/src/ptychodus/model/product/object/__init__.py b/src/ptychodus/model/product/object/__init__.py index a47b7fda7..663e58b3a 100644 --- a/src/ptychodus/model/product/object/__init__.py +++ b/src/ptychodus/model/product/object/__init__.py @@ -1,4 +1,4 @@ -from .builder import ObjectBuilder +from .builder import FromFileObjectBuilder, FromMemoryObjectBuilder, ObjectBuilder from .builder_factory import ObjectBuilderFactory from .dead_leaves import DeadLeavesObjectBuilder from .fractal_noise import FractalNoiseObjectBuilder @@ -13,6 +13,8 @@ __all__ = [ 'DeadLeavesObjectBuilder', 'FractalNoiseObjectBuilder', + 'FromFileObjectBuilder', + 'FromMemoryObjectBuilder', 'GaussianRandomFieldObjectBuilder', 'ObjectBuilder', 'ObjectBuilderFactory', diff --git a/src/ptychodus/model/product/object/builder.py b/src/ptychodus/model/product/object/builder.py index 0ddb5c13d..f6708374b 100644 --- a/src/ptychodus/model/product/object/builder.py +++ b/src/ptychodus/model/product/object/builder.py @@ -37,26 +37,99 @@ def copy(self) -> ObjectBuilder: pass @abstractmethod + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: + """Return the raw, unconditioned object. + + Implementations must NOT generate layers; `build` owns the conditioning + pipeline. Generative implementations should return + `self._pad_object(object_)` so the extra padding is applied to the canvas + they just sized against the scan geometry. + """ + pass + def build( self, geometry_provider: ObjectGeometryProvider, layer_spacing_m: Sequence[float], ) -> Object: - pass + """Return the conditioned object: slice it to the requested layer spacing. - def _create_object( + Overriding this method is reserved for builders whose object is already + conditioned; see `FromMemoryObjectBuilder`. Every builder that generates + or ingests a raw object must leave it alone and implement `_build_raw` + instead. + """ + return self._condition_object(self._build_raw(geometry_provider), layer_spacing_m) + + def _condition_object( self, object_: Object, layer_spacing_m: Sequence[float], ) -> Object: + """Slice the object into the requested number of layers, never destroying + layers it already has. + + `generate_layers` is only non-destructive when the object has a single + layer. Given fewer layers than the input it truncates; given more it keeps + layer zero and throws the rest away before splitting. So a multi-layer + input whose layer count does not already match the request is left alone + -- otherwise a converged multislice result loaded from file would collapse + to one layer under the default empty spacing. + + Note the padding is applied earlier, in `_build_raw`, so the order here is + pad-then-layers rather than the layers-then-pad of the original + `_create_object`. The two do not commute, because `generate_layers` + unwraps the phase of layer zero and a zero-amplitude border changes that + unwrapping. Only generated multislice objects can tell the difference; for + the single-layer default `generate_layers` is a no-op. + """ + num_layers_requested = 1 + len(layer_spacing_m) + num_layers_actual = object_.num_layers + + if num_layers_actual > 1 and num_layers_actual != num_layers_requested: + logger.info( + f'Object already has {num_layers_actual} layer(s);' + f' keeping them rather than re-slicing to {num_layers_requested}.' + ) + return object_ + + return generate_layers(object_, layer_spacing_m) + + def _pad_object(self, object_: Object) -> Object: + """Widen a freshly generated canvas by the extra padding. + + Only generative builders call this. `pad_object` is strictly additive -- + N applications add 2*N*pad pixels per dimension -- and it leaves no trace + in the array, so there is no way to detect an already-padded object and + skip it. The size of a file-supplied object is already fixed by the file, + so the parameter has no coherent meaning there; applying it would grow + every warm-start object on every load/save round trip, unbounded. + """ return pad_object( - generate_layers(object_, layer_spacing_m), + object_, self.extra_padding_x.get_value(), self.extra_padding_y.get_value(), ) class FromMemoryObjectBuilder(ObjectBuilder): + """An object that has already been conditioned. + + Two things produce these. Reconstruction output, which `ProcessingTaskMonitor` + re-assigns to the output product item on every reconstructor iteration (see + `model/processing/monitor.py`), and products loaded from HDF5/NPZ. In both + cases the layer structure and the canvas size are already what the + reconstructor solved for. `generate_layers` would truncate a multislice result + back to whatever the item's `layer_spacing_m` parameter happens to say, and + `pad_object` is strictly additive, so it would grow the array by twice the + padding in each dimension on every iteration. `build` therefore deliberately + bypasses the conditioning pipeline. + + The requested `layer_spacing_m` is likewise ignored in favor of the spacing + the object actually has, which is why `ObjectRepositoryItem.set_num_layers` + is inert for from-memory items. + """ + def __init__(self, settings: ObjectSettings, object_: Object) -> None: super().__init__(settings, 'from_memory') self._settings = settings @@ -70,11 +143,7 @@ def copy(self) -> FromMemoryObjectBuilder: return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_geometry = geometry_provider.get_object_geometry() try: @@ -94,8 +163,26 @@ def build( self._object.layer_spacing_m, ) + def build( + self, + geometry_provider: ObjectGeometryProvider, + layer_spacing_m: Sequence[float], + ) -> Object: + return self._build_raw(geometry_provider) + class FromFileObjectBuilder(ObjectBuilder): + """An object read from file, conditioned on the way in. + + Unlike `FromMemoryObjectBuilder` this is an ingest path, so the layer spacing + does apply -- slicing a two-dimensional object into a multislice warm start is + a real workflow, and before the conditioning pipeline existed the setting was + silently ignored here. `_condition_object` keeps whatever layers the file + already carries. + + The extra padding is deliberately not applied; see `ObjectBuilder._pad_object`. + """ + def __init__( self, settings: ObjectSettings, @@ -119,11 +206,7 @@ def copy(self) -> FromFileObjectBuilder: return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: file_path = self.file_path.get_value() file_type = self.file_type.get_value() logger.debug(f'Reading "{file_path}" as "{file_type}"') diff --git a/src/ptychodus/model/product/object/builder_factory.py b/src/ptychodus/model/product/object/builder_factory.py index cce91a1fb..2531471be 100644 --- a/src/ptychodus/model/product/object/builder_factory.py +++ b/src/ptychodus/model/product/object/builder_factory.py @@ -7,7 +7,7 @@ from ptychodus.api.object import Object, ObjectFileReader, ObjectFileWriter from ptychodus.api.plugins import PluginChooser -from ...diffraction import DiffractionAPI +from ...diffraction import AssembledDiffractionDataset from .builder import FromFileObjectBuilder, ObjectBuilder from .dead_leaves import DeadLeavesObjectBuilder from .fractal_noise import FractalNoiseObjectBuilder @@ -25,37 +25,53 @@ def __init__( self, rng: numpy.random.Generator, settings: ObjectSettings, - diffraction_api: DiffractionAPI, file_reader_chooser: PluginChooser[ObjectFileReader], file_writer_chooser: PluginChooser[ObjectFileWriter], ) -> None: + self._rng = rng self._settings = settings self._file_reader_chooser = file_reader_chooser self._file_writer_chooser = file_writer_chooser - self._builders: Mapping[str, Callable[[], ObjectBuilder]] = { + self._non_diffraction_builders: Mapping[str, Callable[[], ObjectBuilder]] = { 'random': lambda: RandomObjectBuilder(rng, settings), 'dead_leaves': lambda: DeadLeavesObjectBuilder(rng, settings), 'fractal_noise': lambda: FractalNoiseObjectBuilder(rng, settings), 'grf': lambda: GaussianRandomFieldObjectBuilder(rng, settings), - 'stxm': lambda: STXMObjectBuilder(settings, diffraction_api), - 'paganin': lambda: PaganinObjectBuilder(settings, diffraction_api), + } + self._diffraction_builders: Mapping[ + str, Callable[[AssembledDiffractionDataset], ObjectBuilder] + ] = { + 'stxm': lambda dataset: STXMObjectBuilder(settings, dataset), + 'paganin': lambda dataset: PaganinObjectBuilder(settings, dataset), } def __iter__(self) -> Iterator[str]: - return iter(self._builders) + yield from self._non_diffraction_builders + yield from self._diffraction_builders + + def create( + self, name: str, *, dataset: AssembledDiffractionDataset | None = None + ) -> ObjectBuilder: + diffraction_factory = self._diffraction_builders.get(name) + if diffraction_factory is not None: + if dataset is None: + raise RuntimeError( + f'Object builder "{name}" requires an associated diffraction dataset.' + ) + return diffraction_factory(dataset) - def create(self, name: str) -> ObjectBuilder: try: - factory = self._builders[name] + factory = self._non_diffraction_builders[name] except KeyError as exc: raise KeyError(f'Unknown object builder "{name}"!') from exc - return factory() def create_default(self) -> ObjectBuilder: - return next(iter(self._builders.values()))() + return next(iter(self._non_diffraction_builders.values()))() - def create_from_settings(self) -> ObjectBuilder: + def create_from_settings( + self, *, dataset: AssembledDiffractionDataset | None = None + ) -> ObjectBuilder: name = self._settings.builder.get_value() name_repaired = name.casefold() @@ -65,7 +81,7 @@ def create_from_settings(self) -> ObjectBuilder: self._settings.file_type.get_value(), ) - return self.create(name_repaired) + return self.create(name_repaired, dataset=dataset) def get_open_file_filters(self) -> Iterator[str]: for plugin in self._file_reader_chooser: diff --git a/src/ptychodus/model/product/object/dead_leaves.py b/src/ptychodus/model/product/object/dead_leaves.py index 455a1d7bc..e74609483 100644 --- a/src/ptychodus/model/product/object/dead_leaves.py +++ b/src/ptychodus/model/product/object/dead_leaves.py @@ -1,5 +1,4 @@ from __future__ import annotations -from collections.abc import Sequence import logging import numpy @@ -44,11 +43,7 @@ def copy(self) -> DeadLeavesObjectBuilder: return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_ = generate_dead_leaves_object( self._rng, geometry_provider.get_object_geometry(), @@ -60,4 +55,4 @@ def build( leaf_phase_lower_tr=self.leaf_phase_lower_tr.get_value(), leaf_phase_upper_tr=self.leaf_phase_upper_tr.get_value(), ) - return self._create_object(object_, layer_spacing_m) + return self._pad_object(object_) diff --git a/src/ptychodus/model/product/object/fractal_noise.py b/src/ptychodus/model/product/object/fractal_noise.py index dab805edf..46ee50d14 100644 --- a/src/ptychodus/model/product/object/fractal_noise.py +++ b/src/ptychodus/model/product/object/fractal_noise.py @@ -1,5 +1,4 @@ from __future__ import annotations -from collections.abc import Sequence import numpy @@ -36,11 +35,7 @@ def copy(self) -> FractalNoiseObjectBuilder: return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_ = generate_fractal_noise_object( self._rng, geometry_provider.get_object_geometry(), @@ -49,4 +44,4 @@ def build( gain=self.gain.get_value(), lacunarity=self.lacunarity.get_value(), ) - return self._create_object(object_, layer_spacing_m) + return self._pad_object(object_) diff --git a/src/ptychodus/model/product/object/grf.py b/src/ptychodus/model/product/object/grf.py index a259b6b53..585f1f06a 100644 --- a/src/ptychodus/model/product/object/grf.py +++ b/src/ptychodus/model/product/object/grf.py @@ -1,5 +1,4 @@ from __future__ import annotations -from collections.abc import Sequence import numpy @@ -27,14 +26,10 @@ def copy(self) -> GaussianRandomFieldObjectBuilder: return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_ = generate_gaussian_random_field_object( self._rng, geometry_provider.get_object_geometry(), correlation_length_px=self.correlation_length_px.get_value(), ) - return self._create_object(object_, layer_spacing_m) + return self._pad_object(object_) diff --git a/src/ptychodus/model/product/object/item.py b/src/ptychodus/model/product/object/item.py index 39af652d4..326d51382 100644 --- a/src/ptychodus/model/product/object/item.py +++ b/src/ptychodus/model/product/object/item.py @@ -29,6 +29,8 @@ def __init__( self._add_parameter('layer_spacing_m', self.layer_spacing_m) self._add_group('builder', builder, observe=True) + if isinstance(geometry_provider, Observable): + geometry_provider.add_observer(self) self.rebuild() def assign_item(self, item: ObjectRepositoryItem) -> None: @@ -82,6 +84,10 @@ def set_builder(self, builder: ObjectBuilder) -> None: self.rebuild() def rebuild(self, *, recenter: bool = False) -> None: + if not self._geometry_provider.get_object_geometry().get_pixel_geometry().is_valid: + # Geometry not yet bound; the observer wired in __init__ will re-run + # rebuild when the geometry becomes valid. + return try: object_ = self._builder.build(self._geometry_provider, self.layer_spacing_m.get_value()) except Exception: @@ -105,5 +111,7 @@ def rebuild(self, *, recenter: bool = False) -> None: def _update(self, observable: Observable) -> None: if observable is self._builder: self.rebuild() + elif observable is self._geometry_provider: + self.rebuild() else: super()._update(observable) diff --git a/src/ptychodus/model/product/object/item_factory.py b/src/ptychodus/model/product/object/item_factory.py index 1e18dd9be..d65ed22ca 100644 --- a/src/ptychodus/model/product/object/item_factory.py +++ b/src/ptychodus/model/product/object/item_factory.py @@ -4,6 +4,7 @@ from ptychodus.api.object import Object, ObjectGeometryProvider +from ...diffraction import AssembledDiffractionDataset from .builder import FromMemoryObjectBuilder from .builder_factory import ObjectBuilderFactory from .item import ObjectRepositoryItem @@ -23,21 +24,49 @@ def __init__( self._settings = settings self._builder_factory = builder_factory + def _warn_if_conditioning_ignored(self) -> None: + """Note that conditioning settings do not apply to in-memory objects. + + An object supplied in memory comes from reconstruction output or a product + loaded from file, so its layer structure and canvas size are already what + the reconstructor solved for. Batch mode reads product-in.h5 through this + path, where a user who sets a layer spacing in settings.ini would + otherwise see it silently do nothing. Set it on the run that produces the + object instead -- the from-file builder does apply the layer spacing. + """ + settings = self._settings + is_conditioning_requested = ( + # Note the padding parameters default to 1, not 0. + settings.extra_padding_x.get_value() != 1 + or settings.extra_padding_y.get_value() != 1 + or len(settings.object_layer_spacing_m.get_value()) != 0 + ) + + if is_conditioning_requested: + logger.info( + 'Objects supplied in memory are already conditioned;' + ' ignoring the extra padding and layer spacing settings.' + ) + def create( self, geometry_provider: ObjectGeometryProvider, object_: Object | None = None ) -> ObjectRepositoryItem: if object_ is None: builder = self._builder_factory.create_default() else: + self._warn_if_conditioning_ignored() builder = FromMemoryObjectBuilder(self._settings, object_) return ObjectRepositoryItem(geometry_provider, self._settings, builder) def create_from_settings( - self, geometry_provider: ObjectGeometryProvider + self, + geometry_provider: ObjectGeometryProvider, + *, + dataset: AssembledDiffractionDataset | None = None, ) -> ObjectRepositoryItem: try: - builder = self._builder_factory.create_from_settings() + builder = self._builder_factory.create_from_settings(dataset=dataset) except Exception as exc: logger.error(''.join(exc.args)) builder = self._builder_factory.create_default() diff --git a/src/ptychodus/model/product/object/paganin.py b/src/ptychodus/model/product/object/paganin.py index 4a140509b..2b352ee9d 100644 --- a/src/ptychodus/model/product/object/paganin.py +++ b/src/ptychodus/model/product/object/paganin.py @@ -1,12 +1,11 @@ from __future__ import annotations -from collections.abc import Sequence import logging from ptychodus.api.object import Object, ObjectGeometryProvider from ptychodus.api.object_gen import generate_paganin_object -from ...diffraction import DiffractionAPI +from ...diffraction import AssembledDiffractionDataset from .builder import ObjectBuilder from .settings import ObjectSettings @@ -17,11 +16,11 @@ class PaganinObjectBuilder(ObjectBuilder): def __init__( self, settings: ObjectSettings, - diffraction_api: DiffractionAPI, + dataset: AssembledDiffractionDataset, ) -> None: super().__init__(settings, 'paganin') self._settings = settings - self._diffraction_api = diffraction_api + self._dataset = dataset self.probe_wavelength_m = settings.paganin_probe_wavelength_m.copy() self._add_parameter('probe_wavelength_m', self.probe_wavelength_m) @@ -31,24 +30,20 @@ def __init__( self._add_parameter('delta_over_beta', self.delta_over_beta) def copy(self) -> PaganinObjectBuilder: - builder = PaganinObjectBuilder(self._settings, self._diffraction_api) + builder = PaganinObjectBuilder(self._settings, self._dataset) for key, value in self.parameters().items(): builder.parameters()[key].set_value(value.get_value()) return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_ = generate_paganin_object( geometry_provider.get_object_geometry(), - self._diffraction_api.get_assembled_data(), + self._dataset.get_assembled_data(), geometry_provider.get_probe_positions(), probe_wavelength_m=self.probe_wavelength_m.get_value(), propagation_distance_m=self.propagation_distance_m.get_value(), delta_over_beta=self.delta_over_beta.get_value(), ) - return self._create_object(object_, layer_spacing_m) + return self._pad_object(object_) diff --git a/src/ptychodus/model/product/object/random.py b/src/ptychodus/model/product/object/random.py index 3a59ee206..06b9c7143 100644 --- a/src/ptychodus/model/product/object/random.py +++ b/src/ptychodus/model/product/object/random.py @@ -1,5 +1,4 @@ from __future__ import annotations -from collections.abc import Sequence import numpy @@ -33,11 +32,7 @@ def copy(self) -> RandomObjectBuilder: return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_ = generate_random_object( self._rng, geometry_provider.get_object_geometry(), @@ -47,4 +42,4 @@ def build( phase_deviation_tr=self.phase_deviation_tr.get_value(), blur_deviation_px=self.blur_deviation_px.get_value(), ) - return self._create_object(object_, layer_spacing_m) + return self._pad_object(object_) diff --git a/src/ptychodus/model/product/object/stxm.py b/src/ptychodus/model/product/object/stxm.py index 4600968f9..5b07b46d0 100644 --- a/src/ptychodus/model/product/object/stxm.py +++ b/src/ptychodus/model/product/object/stxm.py @@ -1,12 +1,11 @@ from __future__ import annotations -from collections.abc import Sequence import logging from ptychodus.api.object import Object, ObjectGeometryProvider from ptychodus.api.object_gen import generate_stxm_object -from ...diffraction import DiffractionAPI +from ...diffraction import AssembledDiffractionDataset from .builder import ObjectBuilder from .settings import ObjectSettings @@ -17,28 +16,24 @@ class STXMObjectBuilder(ObjectBuilder): def __init__( self, settings: ObjectSettings, - diffraction_api: DiffractionAPI, + dataset: AssembledDiffractionDataset, ) -> None: super().__init__(settings, 'stxm') self._settings = settings - self._diffraction_api = diffraction_api + self._dataset = dataset def copy(self) -> STXMObjectBuilder: - builder = STXMObjectBuilder(self._settings, self._diffraction_api) + builder = STXMObjectBuilder(self._settings, self._dataset) for key, value in self.parameters().items(): builder.parameters()[key].set_value(value.get_value()) return builder - def build( - self, - geometry_provider: ObjectGeometryProvider, - layer_spacing_m: Sequence[float], - ) -> Object: + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: object_ = generate_stxm_object( geometry_provider.get_object_geometry(), - self._diffraction_api.get_assembled_data(), + self._dataset.get_assembled_data(), geometry_provider.get_probe_positions(), ) - return self._create_object(object_, layer_spacing_m) + return self._pad_object(object_) diff --git a/src/ptychodus/model/product/object_repository.py b/src/ptychodus/model/product/object_repository.py index a7bf36699..57bbce3ac 100644 --- a/src/ptychodus/model/product/object_repository.py +++ b/src/ptychodus/model/product/object_repository.py @@ -5,6 +5,7 @@ from ptychodus.api.observer import ObservableSequence from ptychodus.api.product import LossValue +from ..diffraction import AssembledDiffractionDataset from .item import ProductRepositoryItem, ProductRepositoryObserver from .metadata import MetadataRepositoryItem from .object import ObjectRepositoryItem @@ -27,6 +28,9 @@ def get_name(self, index: int) -> str: def set_name(self, index: int, name: str) -> None: self._repository[index].set_name(name) + def get_dataset(self, index: int) -> AssembledDiffractionDataset | None: + return self._repository[index].get_dataset() + @overload def __getitem__(self, index: int) -> ObjectRepositoryItem: ... @@ -64,5 +68,11 @@ def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: pass + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: self.notify_observers_item_removed(index, item.get_object_item()) diff --git a/src/ptychodus/model/product/probe/__init__.py b/src/ptychodus/model/product/probe/__init__.py index 645c00e68..f98aa0498 100644 --- a/src/ptychodus/model/product/probe/__init__.py +++ b/src/ptychodus/model/product/probe/__init__.py @@ -1,5 +1,10 @@ from .average_pattern import AveragePatternProbeBuilder -from .builder import ProbeModeDecayType, ProbeSequenceBuilder +from .builder import ( + FromFileProbeBuilder, + FromMemoryProbeBuilder, + ProbeModeDecayType, + ProbeSequenceBuilder, +) from .builder_factory import ProbeBuilderFactory from .disk import DiskProbeBuilder from .fzp import FresnelZonePlateProbeBuilder @@ -15,6 +20,8 @@ 'AveragePatternProbeBuilder', 'DiskProbeBuilder', 'FresnelZonePlateProbeBuilder', + 'FromFileProbeBuilder', + 'FromMemoryProbeBuilder', 'HermiteProbeBuilder', 'ProbeBuilderFactory', 'ProbeModeDecayType', diff --git a/src/ptychodus/model/product/probe/average_pattern.py b/src/ptychodus/model/product/probe/average_pattern.py index d20b81e31..a86c414be 100644 --- a/src/ptychodus/model/product/probe/average_pattern.py +++ b/src/ptychodus/model/product/probe/average_pattern.py @@ -3,9 +3,9 @@ import numpy from ptychodus.api.probe import ProbeSequence, ProbeGeometryProvider -from ptychodus.api.probe_gen import generate_average_pattern_probe, rescale_probe_intensity +from ptychodus.api.probe_gen import generate_average_pattern_probe -from ...diffraction import DiffractionAPI +from ...diffraction import AssembledDiffractionDataset from .builder import ProbeSequenceBuilder from .settings import ProbeSettings @@ -15,29 +15,27 @@ def __init__( self, rng: numpy.random.Generator, settings: ProbeSettings, - diffraction_api: DiffractionAPI, + dataset: AssembledDiffractionDataset, ) -> None: - super().__init__(settings, 'average_pattern') - self._rng = rng + super().__init__(rng, settings, 'average_pattern') self._settings = settings - self._diffraction_api = diffraction_api + self._dataset = dataset def copy(self) -> AveragePatternProbeBuilder: - builder = AveragePatternProbeBuilder(self._rng, self._settings, self._diffraction_api) + builder = AveragePatternProbeBuilder(self._rng, self._settings, self._dataset) for key, value in self.parameters().items(): builder.parameters()[key].set_value(value.get_value()) return builder - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: - probe = rescale_probe_intensity( + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._rescale_to_photon_count( generate_average_pattern_probe( geometry_provider.get_probe_geometry(), - self._diffraction_api.get_assembled_data(), + self._dataset.get_assembled_data(), probe_wavelength_m=geometry_provider.probe_wavelength_m, detector_distance_m=geometry_provider.detector_distance_m, ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe/builder.py b/src/ptychodus/model/product/probe/builder.py index 3c1825afc..df094b3f5 100644 --- a/src/ptychodus/model/product/probe/builder.py +++ b/src/ptychodus/model/product/probe/builder.py @@ -7,7 +7,11 @@ import numpy from ptychodus.api.parametric import ParameterGroup -from ptychodus.api.probe_gen import generate_coherent_probe_modes, generate_incoherent_probe_modes +from ptychodus.api.probe_gen import ( + generate_coherent_probe_modes, + generate_incoherent_probe_modes, + rescale_probe_intensity, +) from ptychodus.api.probe import ( Probe, ProbeSequence, @@ -38,8 +42,10 @@ def get_weights(self, num_modes: int, decay_ratio: float) -> Sequence[float]: class ProbeSequenceBuilder(ParameterGroup): - def __init__(self, settings: ProbeSettings, name: str) -> None: + def __init__(self, rng: numpy.random.Generator, settings: ProbeSettings, name: str) -> None: super().__init__() + self._rng = rng + self._name = settings.builder.copy() self._name.set_value(name) self._add_parameter('name', self._name) @@ -71,9 +77,45 @@ def copy(self) -> ProbeSequenceBuilder: pass @abstractmethod - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + """Return the raw, unconditioned probe. + + Implementations must NOT expand the incoherent or coherent (OPR) modes; + `build` owns the conditioning pipeline. Generative implementations should + return `self._rescale_to_photon_count(probe, geometry_provider)`, which + normalizes the intensity and widens the 3-D `Probe` they generated to the + 4-D `ProbeSequence` this method returns. + """ pass + def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + """Return the conditioned probe: incoherent modes, then coherent (OPR) modes. + + Overriding this method is reserved for builders whose probe is already + conditioned; see `FromMemoryProbeBuilder`. Every builder that generates or + ingests a raw probe must leave it alone and implement `_build_raw` + instead. + """ + return self._condition_probe(self._build_raw(geometry_provider), geometry_provider) + + def _rescale_to_photon_count( + self, probe: Probe, geometry_provider: ProbeGeometryProvider + ) -> ProbeSequence: + """Normalize a freshly generated probe to the expected photon count and + widen it to the `ProbeSequence` that `_build_raw` returns. + + Only generative builders call this. A probe read from file already carries + the intensity it was reconstructed at; rescaling it would silently decouple + it from a matching from-file object, because the data constrains the + product of probe and object, not either one alone. + """ + rescaled = rescale_probe_intensity(probe, geometry_provider.probe_photon_count) + return ProbeSequence( + array=rescaled.get_array(), + opr_weights=None, + pixel_geometry=rescaled.get_pixel_geometry(), + ) + def _get_imode_weights(self) -> Sequence[float]: imode_decay_ratio = self.incoherent_mode_decay_ratio.get_value() imode_decay_type_text = self.incoherent_mode_decay_type.get_value() @@ -88,41 +130,117 @@ def _get_imode_weights(self) -> Sequence[float]: num_imodes = self.num_incoherent_modes.get_value() return imode_decay_type.get_weights(num_imodes, imode_decay_ratio) - def _build_probe_modes( - self, rng: numpy.random.Generator, probe: Probe, num_diffraction_patterns: int + def _condition_probe( + self, probe_seq: ProbeSequence, geometry_provider: ProbeGeometryProvider ) -> ProbeSequence: - probe_with_imodes = generate_incoherent_probe_modes( - rng, - probe, - self._get_imode_weights(), - orthogonalize=self.orthogonalize_incoherent_modes.get_value(), - ) - probe_seq = generate_coherent_probe_modes( - rng, - probe_with_imodes, - num_cmodes=self.num_coherent_modes.get_value(), - num_diffraction_patterns=num_diffraction_patterns, - ) - array = probe_seq.get_array() - logger.debug(f'Multimodal probe {array.shape=}') + """Expand the probe to the requested mode structure, never shrinking it. + + Every step is expand-only, so the pipeline is idempotent: conditioning an + already-conditioned probe returns it unchanged. That matters because the + generators in `ptychodus.api.probe_gen` are not safe to re-apply. + `generate_incoherent_probe_modes` re-orthogonalizes and renormalizes every + incoherent mode to the decay profile, and `generate_coherent_probe_modes` + fills the whole output with fresh Gaussian noise, keeps only coherent mode + zero of its input, and regenerates the OPR weights from scratch. Run + either one on a converged probe and the reconstruction is gone. + + The guards are data-driven rather than provenance-driven, so they live + here rather than in the ingesting subclasses. Generative builders always + emit a single coherent, single incoherent mode, which makes every guard + inert on that path. + """ + num_imodes_requested = self.num_incoherent_modes.get_value() + num_cmodes_requested = self.num_coherent_modes.get_value() + + if probe_seq.num_coherent_modes > 1 or probe_seq.get_opr_weights_or_none() is not None: + # There is no non-destructive way to extend a solved OPR basis, so + # leave the whole mode structure alone. + if ( + num_cmodes_requested > probe_seq.num_coherent_modes + or num_imodes_requested > probe_seq.num_incoherent_modes + ): + logger.info( + 'Probe already has an OPR mode basis' + f' ({probe_seq.num_coherent_modes} coherent,' + f' {probe_seq.num_incoherent_modes} incoherent);' + ' leaving its mode structure unchanged.' + ) + + return probe_seq + + probe = probe_seq.get_probe_no_opr() + num_imodes_actual = probe.num_incoherent_modes + + if num_imodes_actual < num_imodes_requested: + probe = generate_incoherent_probe_modes( + self._rng, + probe, + self._get_imode_weights(), + orthogonalize=self.orthogonalize_incoherent_modes.get_value(), + ) + elif num_imodes_actual > num_imodes_requested: + logger.info( + f'Probe has {num_imodes_actual} incoherent mode(s);' + f' keeping them rather than discarding down to {num_imodes_requested}.' + ) + + if num_cmodes_requested > 1: + probe_seq = generate_coherent_probe_modes( + self._rng, + probe, + num_cmodes=num_cmodes_requested, + num_diffraction_patterns=geometry_provider.num_scan_points, + ) + else: + probe_seq = ProbeSequence( + array=probe.get_array(), + opr_weights=None, + pixel_geometry=probe.get_pixel_geometry(), + ) + + logger.debug(f'Conditioned probe {probe_seq.get_array().shape=}') return probe_seq class FromMemoryProbeBuilder(ProbeSequenceBuilder): - def __init__(self, settings: ProbeSettings, probe: ProbeSequence) -> None: - super().__init__(settings, 'from_memory') + """A probe that has already been conditioned. + + Two things produce these. Reconstruction output, which `ProcessingTaskMonitor` + re-assigns to the output product item on every reconstructor iteration (see + `model/processing/monitor.py`), and products loaded from HDF5/NPZ, whose probe + was conditioned before it was saved. In both cases the incoherent and coherent + (OPR) mode structure is already what the reconstructor solved for, so + re-running the mode generators would be catastrophic rather than merely lossy: + `generate_incoherent_probe_modes` re-orthogonalizes and renormalizes every + incoherent mode to the decay profile, and `generate_coherent_probe_modes` + replaces every coherent mode but the first with fresh Gaussian noise and + regenerates the OPR weights from scratch -- once per iteration. `build` + therefore deliberately bypasses the conditioning pipeline. + + The expand-only guards in `_condition_probe` would in fact catch most of this + on their own, but the bypass is explicit so that the invariant does not depend + on them. + """ + + def __init__( + self, + rng: numpy.random.Generator, + settings: ProbeSettings, + probe: ProbeSequence, + ) -> None: + super().__init__(rng, settings, 'from_memory') self._settings = settings self._probe = probe.copy() def copy(self) -> FromMemoryProbeBuilder: - builder = FromMemoryProbeBuilder(self._settings, self._probe) + builder = FromMemoryProbeBuilder(self._rng, self._settings, self._probe) for key, value in self.parameters().items(): builder.parameters()[key].set_value(value.get_value()) return builder - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: probe_geometry = geometry_provider.get_probe_geometry() try: @@ -130,22 +248,37 @@ def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: except ValueError: pixel_geometry = probe_geometry.get_pixel_geometry() - try: - opr_weights = self._probe.get_opr_weights() - except ValueError: - opr_weights = None - # TODO regrid probe as needed based on probe geometry from file/provider return ProbeSequence( self._probe.get_array(), - opr_weights, + self._probe.get_opr_weights_or_none(), pixel_geometry, ) + def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._build_raw(geometry_provider) + class FromFileProbeBuilder(ProbeSequenceBuilder): - def __init__(self, settings: ProbeSettings, file_reader: ProbeFileReader) -> None: - super().__init__(settings, 'from_file') + """A probe read from file, conditioned on the way in. + + Unlike `FromMemoryProbeBuilder` this is an ingest path, so the mode settings + do apply -- warm-starting a mixed-state run from a single-mode probe file is a + real workflow, and before the conditioning pipeline existed those settings + were silently ignored here. `_condition_probe` is expand-only, so a file that + already carries more modes, or an OPR basis, keeps what it has. + + The photon-count rescale is deliberately not applied; see + `ProbeSequenceBuilder._rescale_to_photon_count`. + """ + + def __init__( + self, + rng: numpy.random.Generator, + settings: ProbeSettings, + file_reader: ProbeFileReader, + ) -> None: + super().__init__(rng, settings, 'from_file') self._settings = settings self._file_reader = file_reader @@ -156,14 +289,14 @@ def __init__(self, settings: ProbeSettings, file_reader: ProbeFileReader) -> Non self._add_parameter('file_type', self.file_type) def copy(self) -> FromFileProbeBuilder: - builder = FromFileProbeBuilder(self._settings, self._file_reader) + builder = FromFileProbeBuilder(self._rng, self._settings, self._file_reader) for key, value in self.parameters().items(): builder.parameters()[key].set_value(value.get_value()) return builder - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: file_path = self.file_path.get_value() file_type = self.file_type.get_value() logger.debug(f'Reading "{file_path}" as "{file_type}"') @@ -180,14 +313,9 @@ def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: except ValueError: pixel_geometry = probe_geometry.get_pixel_geometry() - try: - opr_weights = probe_from_file.get_opr_weights() - except ValueError: - opr_weights = None - # TODO regrid probe as needed based on probe geometry from file/provider return ProbeSequence( probe_from_file.get_array(), - opr_weights, + probe_from_file.get_opr_weights_or_none(), pixel_geometry, ) diff --git a/src/ptychodus/model/product/probe/builder_factory.py b/src/ptychodus/model/product/probe/builder_factory.py index 00918b59f..2655b35d2 100644 --- a/src/ptychodus/model/product/probe/builder_factory.py +++ b/src/ptychodus/model/product/probe/builder_factory.py @@ -8,7 +8,7 @@ from ptychodus.api.probe import ProbeFileReader, ProbeFileWriter, ProbeSequence from ptychodus.api.probe_gen import FresnelZonePlate -from ...diffraction import DiffractionAPI +from ...diffraction import AssembledDiffractionDataset from .average_pattern import AveragePatternProbeBuilder from .builder import FromFileProbeBuilder, ProbeSequenceBuilder from .disk import DiskProbeBuilder @@ -27,7 +27,6 @@ def __init__( self, rng: numpy.random.Generator, settings: ProbeSettings, - diffraction_api: DiffractionAPI, fresnel_zone_plate_chooser: PluginChooser[FresnelZonePlate], file_reader_chooser: PluginChooser[ProbeFileReader], file_writer_chooser: PluginChooser[ProbeFileWriter], @@ -35,35 +34,50 @@ def __init__( super().__init__() self._rng = rng self._settings = settings - self._diffraction_api = diffraction_api self._fresnel_zone_plate_chooser = fresnel_zone_plate_chooser self._file_reader_chooser = file_reader_chooser self._file_writer_chooser = file_writer_chooser - self._builders: Mapping[str, Callable[[], ProbeSequenceBuilder]] = { + self._non_diffraction_builders: Mapping[str, Callable[[], ProbeSequenceBuilder]] = { 'disk': lambda: DiskProbeBuilder(rng, settings), - 'average_pattern': self._create_average_pattern_builder, 'fresnel_zone_plate': self._create_fresnel_zone_plate_builder, 'hermite': lambda: HermiteProbeBuilder(rng, settings), 'rectangular': lambda: RectangularProbeBuilder(rng, settings), 'super_gaussian': lambda: SuperGaussianProbeBuilder(rng, settings), 'zernike': lambda: ZernikeProbeBuilder(rng, settings), } + self._diffraction_builders: Mapping[ + str, Callable[[AssembledDiffractionDataset], ProbeSequenceBuilder] + ] = { + 'average_pattern': lambda dataset: AveragePatternProbeBuilder(rng, settings, dataset), + } def __iter__(self) -> Iterator[str]: - return iter(self._builders) + yield from self._non_diffraction_builders + yield from self._diffraction_builders + + def create( + self, name: str, *, dataset: AssembledDiffractionDataset | None = None + ) -> ProbeSequenceBuilder: + diffraction_factory = self._diffraction_builders.get(name) + if diffraction_factory is not None: + if dataset is None: + raise RuntimeError( + f'Probe builder "{name}" requires an associated diffraction dataset.' + ) + return diffraction_factory(dataset) - def create(self, name: str) -> ProbeSequenceBuilder: try: - factory = self._builders[name] + factory = self._non_diffraction_builders[name] except KeyError as exc: raise KeyError(f'Unknown probe builder "{name}"!') from exc - return factory() def create_default(self) -> ProbeSequenceBuilder: - return next(iter(self._builders.values()))() + return next(iter(self._non_diffraction_builders.values()))() - def create_from_settings(self) -> ProbeSequenceBuilder: + def create_from_settings( + self, *, dataset: AssembledDiffractionDataset | None = None + ) -> ProbeSequenceBuilder: name = self._settings.builder.get_value() name_repaired = name.casefold() @@ -73,10 +87,7 @@ def create_from_settings(self) -> ProbeSequenceBuilder: self._settings.file_type.get_value(), ) - return self.create(name_repaired) - - def _create_average_pattern_builder(self) -> ProbeSequenceBuilder: - return AveragePatternProbeBuilder(self._rng, self._settings, self._diffraction_api) + return self.create(name_repaired, dataset=dataset) def _create_fresnel_zone_plate_builder(self) -> ProbeSequenceBuilder: return FresnelZonePlateProbeBuilder( @@ -94,7 +105,7 @@ def create_probe_from_file(self, file_path: Path, file_filter: str) -> ProbeSequ self._file_reader_chooser.set_current_plugin(file_filter) file_reader = self._file_reader_chooser.get_current_plugin().strategy - builder = FromFileProbeBuilder(self._settings, file_reader) + builder = FromFileProbeBuilder(self._rng, self._settings, file_reader) builder.file_path.set_value(file_path) builder.file_type.set_value(self._file_reader_chooser.get_current_plugin().simple_name) return builder diff --git a/src/ptychodus/model/product/probe/disk.py b/src/ptychodus/model/product/probe/disk.py index 828612e8f..f1753a5b7 100644 --- a/src/ptychodus/model/product/probe/disk.py +++ b/src/ptychodus/model/product/probe/disk.py @@ -3,7 +3,7 @@ import numpy from ptychodus.api.probe import ProbeSequence, ProbeGeometryProvider -from ptychodus.api.probe_gen import defocus_probe, generate_disk_probe, rescale_probe_intensity +from ptychodus.api.probe_gen import defocus_probe, generate_disk_probe from .builder import ProbeSequenceBuilder from .settings import ProbeSettings @@ -11,8 +11,7 @@ class DiskProbeBuilder(ProbeSequenceBuilder): def __init__(self, rng: numpy.random.Generator, settings: ProbeSettings) -> None: - super().__init__(settings, 'disk') - self._rng = rng + super().__init__(rng, settings, 'disk') self._settings = settings self.diameter_m = settings.disk_diameter_m.copy() @@ -30,8 +29,8 @@ def copy(self) -> DiskProbeBuilder: return builder - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: - probe = rescale_probe_intensity( + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._rescale_to_photon_count( defocus_probe( generate_disk_probe( geometry_provider.get_probe_geometry(), @@ -40,6 +39,5 @@ def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: probe_wavelength_m=geometry_provider.probe_wavelength_m, defocus_distance_m=self.defocus_distance_m.get_value(), ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe/fzp.py b/src/ptychodus/model/product/probe/fzp.py index da2727ba8..eb1764b3e 100644 --- a/src/ptychodus/model/product/probe/fzp.py +++ b/src/ptychodus/model/product/probe/fzp.py @@ -8,7 +8,6 @@ from ptychodus.api.probe_gen import ( FresnelZonePlate, generate_fresnel_zone_plate_probe, - rescale_probe_intensity, ) from .builder import ProbeSequenceBuilder @@ -22,8 +21,7 @@ def __init__( settings: ProbeSettings, fresnel_zone_plate_chooser: PluginChooser[FresnelZonePlate], ) -> None: - super().__init__(settings, 'fresnel_zone_plate') - self._rng = rng + super().__init__(rng, settings, 'fresnel_zone_plate') self._settings = settings self._fresnel_zone_plate_chooser = fresnel_zone_plate_chooser @@ -61,19 +59,18 @@ def apply_presets(self, display_name: str) -> None: self.outermost_zone_width_m.set_value(fzp.outermost_zone_width_m) self.central_beamstop_diameter_m.set_value(fzp.central_beamstop_diameter_m) - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: zone_plate = FresnelZonePlate( zone_plate_diameter_m=self.zone_plate_diameter_m.get_value(), outermost_zone_width_m=self.outermost_zone_width_m.get_value(), central_beamstop_diameter_m=self.central_beamstop_diameter_m.get_value(), ) - probe = rescale_probe_intensity( + return self._rescale_to_photon_count( generate_fresnel_zone_plate_probe( geometry=geometry_provider.get_probe_geometry(), zone_plate=zone_plate, probe_wavelength_m=geometry_provider.probe_wavelength_m, defocus_distance_m=self.defocus_distance_m.get_value(), ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe/hermite.py b/src/ptychodus/model/product/probe/hermite.py index 095907683..2ccb1d58d 100644 --- a/src/ptychodus/model/product/probe/hermite.py +++ b/src/ptychodus/model/product/probe/hermite.py @@ -5,7 +5,7 @@ from ptychodus.api.geometry import HermiteMode from ptychodus.api.probe import ProbeSequence, ProbeGeometryProvider -from ptychodus.api.probe_gen import generate_hermite_probe, rescale_probe_intensity +from ptychodus.api.probe_gen import generate_hermite_probe from .builder import ProbeSequenceBuilder from .settings import ProbeSettings @@ -15,8 +15,7 @@ class HermiteProbeBuilder(ProbeSequenceBuilder): def __init__(self, rng: numpy.random.Generator, settings: ProbeSettings) -> None: - super().__init__(settings, 'hermite') - self._rng = rng + super().__init__(rng, settings, 'hermite') self._settings = settings self._polynomial: list[HermiteMode] = list() @@ -90,14 +89,13 @@ def get_mode(self, idx: int) -> HermiteMode: def __len__(self) -> int: return len(self._polynomial) - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: - probe = rescale_probe_intensity( + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._rescale_to_photon_count( generate_hermite_probe( geometry_provider.get_probe_geometry(), self._polynomial, width_m=self.width_m.get_value(), height_m=self.height_m.get_value(), ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe/item.py b/src/ptychodus/model/product/probe/item.py index 23d7240ec..1ee724a55 100644 --- a/src/ptychodus/model/product/probe/item.py +++ b/src/ptychodus/model/product/probe/item.py @@ -1,6 +1,8 @@ from __future__ import annotations import logging +import numpy + from ptychodus.api.observer import Observable from ptychodus.api.parametric import ParameterGroup from ptychodus.api.probe import ( @@ -21,17 +23,21 @@ class ProbeRepositoryItem(ParameterGroup): def __init__( self, + rng: numpy.random.Generator, geometry_provider: ProbeGeometryProvider, settings: ProbeSettings, builder: ProbeSequenceBuilder, ) -> None: super().__init__() + self._rng = rng self._geometry_provider = geometry_provider self._settings = settings self._builder = builder self._probe_seq = ProbeSequence(array=None, opr_weights=None, pixel_geometry=None) self._add_group('builder', builder, observe=True) + if isinstance(geometry_provider, Observable): + geometry_provider.add_observer(self) self._rebuild() def assign_item(self, item: ProbeRepositoryItem) -> None: @@ -39,7 +45,7 @@ def assign_item(self, item: ProbeRepositoryItem) -> None: self._rebuild() def assign(self, probe: ProbeSequence) -> None: - builder = FromMemoryProbeBuilder(self._settings, probe) + builder = FromMemoryProbeBuilder(self._rng, self._settings, probe) self.set_builder(builder) def sync_to_settings(self) -> None: @@ -99,6 +105,10 @@ def set_builder(self, builder: ProbeSequenceBuilder) -> None: self._rebuild() def _rebuild(self) -> None: + if not self._geometry_provider.get_probe_geometry().get_pixel_geometry().is_valid: + # Geometry not yet bound; the observer wired in __init__ will re-run + # _rebuild when the geometry becomes valid. + return try: probe_seq = self._builder.build(self._geometry_provider) except Exception: @@ -110,5 +120,7 @@ def _rebuild(self) -> None: def _update(self, observable: Observable) -> None: if observable is self._builder: self._rebuild() + elif observable is self._geometry_provider: + self._rebuild() else: super()._update(observable) diff --git a/src/ptychodus/model/product/probe/item_factory.py b/src/ptychodus/model/product/probe/item_factory.py index cbe571fac..ed63f2992 100644 --- a/src/ptychodus/model/product/probe/item_factory.py +++ b/src/ptychodus/model/product/probe/item_factory.py @@ -4,6 +4,7 @@ from ptychodus.api.probe import ProbeSequence, ProbeGeometryProvider +from ...diffraction import AssembledDiffractionDataset from .builder import FromMemoryProbeBuilder from .builder_factory import ProbeBuilderFactory from .item import ProbeRepositoryItem @@ -23,21 +24,53 @@ def __init__( self._settings = settings self._builder_factory = builder_factory + def _warn_if_conditioning_ignored(self) -> None: + """Note that the mode settings do not apply to in-memory probes. + + A probe supplied in memory comes from reconstruction output or a product + loaded from file, so its mode structure is already what the reconstructor + solved for. Batch mode reads product-in.h5 through this path, where a user + who sets a mode count in settings.ini would otherwise see it silently do + nothing. Set the mode counts on the run that produces the probe instead -- + ptychodus-bdp reads probes through the from-file builder, which does + condition them. + """ + settings = self._settings + # The decay parameters and the orthogonalization flag are inert at a + # single incoherent mode, so gate on the two counts alone; including them + # would fire on default settings. + is_conditioning_requested = ( + settings.num_incoherent_modes.get_value() != 1 + or settings.num_coherent_modes.get_value() != 1 + ) + + if is_conditioning_requested: + logger.info( + 'Probes supplied in memory are already conditioned;' + ' ignoring the incoherent and coherent mode settings.' + ) + def create( self, geometry_provider: ProbeGeometryProvider, probe: ProbeSequence | None = None ) -> ProbeRepositoryItem: if probe is None: builder = self._builder_factory.create_default() else: - builder = FromMemoryProbeBuilder(self._settings, probe) + self._warn_if_conditioning_ignored() + builder = FromMemoryProbeBuilder(self._rng, self._settings, probe) - return ProbeRepositoryItem(geometry_provider, self._settings, builder) + return ProbeRepositoryItem(self._rng, geometry_provider, self._settings, builder) - def create_from_settings(self, geometry_provider: ProbeGeometryProvider) -> ProbeRepositoryItem: + def create_from_settings( + self, + geometry_provider: ProbeGeometryProvider, + *, + dataset: AssembledDiffractionDataset | None = None, + ) -> ProbeRepositoryItem: try: - builder = self._builder_factory.create_from_settings() + builder = self._builder_factory.create_from_settings(dataset=dataset) except Exception as exc: logger.error(''.join(exc.args)) builder = self._builder_factory.create_default() - return ProbeRepositoryItem(geometry_provider, self._settings, builder) + return ProbeRepositoryItem(self._rng, geometry_provider, self._settings, builder) diff --git a/src/ptychodus/model/product/probe/rect.py b/src/ptychodus/model/product/probe/rect.py index 4bee26b27..41a06c880 100644 --- a/src/ptychodus/model/product/probe/rect.py +++ b/src/ptychodus/model/product/probe/rect.py @@ -6,7 +6,6 @@ from ptychodus.api.probe_gen import ( defocus_probe, generate_rectangular_probe, - rescale_probe_intensity, ) from .builder import ProbeSequenceBuilder @@ -15,8 +14,7 @@ class RectangularProbeBuilder(ProbeSequenceBuilder): def __init__(self, rng: numpy.random.Generator, settings: ProbeSettings) -> None: - super().__init__(settings, 'rectangular') - self._rng = rng + super().__init__(rng, settings, 'rectangular') self._settings = settings self.width_m = settings.rectangle_width_m.copy() @@ -37,8 +35,8 @@ def copy(self) -> RectangularProbeBuilder: return builder - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: - probe = rescale_probe_intensity( + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._rescale_to_photon_count( defocus_probe( generate_rectangular_probe( geometry_provider.get_probe_geometry(), @@ -48,6 +46,5 @@ def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: probe_wavelength_m=geometry_provider.probe_wavelength_m, defocus_distance_m=self.defocus_distance_m.get_value(), ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe/super_gaussian.py b/src/ptychodus/model/product/probe/super_gaussian.py index 1319a9010..b9b6330b1 100644 --- a/src/ptychodus/model/product/probe/super_gaussian.py +++ b/src/ptychodus/model/product/probe/super_gaussian.py @@ -3,7 +3,7 @@ import numpy from ptychodus.api.probe import ProbeSequence, ProbeGeometryProvider -from ptychodus.api.probe_gen import rescale_probe_intensity, generate_super_gaussian_probe +from ptychodus.api.probe_gen import generate_super_gaussian_probe from .builder import ProbeSequenceBuilder from .settings import ProbeSettings @@ -11,8 +11,7 @@ class SuperGaussianProbeBuilder(ProbeSequenceBuilder): def __init__(self, rng: numpy.random.Generator, settings: ProbeSettings) -> None: - super().__init__(settings, 'super_gaussian') - self._rng = rng + super().__init__(rng, settings, 'super_gaussian') self._settings = settings self.annular_radius_m = settings.super_gaussian_annular_radius_m.copy() @@ -32,14 +31,13 @@ def copy(self) -> SuperGaussianProbeBuilder: return builder - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: - probe = rescale_probe_intensity( + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._rescale_to_photon_count( generate_super_gaussian_probe( geometry_provider.get_probe_geometry(), annular_radius_m=self.annular_radius_m.get_value(), fwhm_m=self.fwhm_m.get_value(), order_parameter=self.order_parameter.get_value(), ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe/zernike.py b/src/ptychodus/model/product/probe/zernike.py index f6c09de9b..7011e84ef 100644 --- a/src/ptychodus/model/product/probe/zernike.py +++ b/src/ptychodus/model/product/probe/zernike.py @@ -5,7 +5,7 @@ from ptychodus.api.geometry import ZernikeMode from ptychodus.api.probe import ProbeSequence, ProbeGeometryProvider -from ptychodus.api.probe_gen import generate_zernike_probe, rescale_probe_intensity +from ptychodus.api.probe_gen import generate_zernike_probe from .builder import ProbeSequenceBuilder from .settings import ProbeSettings @@ -15,8 +15,7 @@ class ZernikeProbeBuilder(ProbeSequenceBuilder): def __init__(self, rng: numpy.random.Generator, settings: ProbeSettings) -> None: - super().__init__(settings, 'zernike') - self._rng = rng + super().__init__(rng, settings, 'zernike') self._settings = settings self._polynomial: list[ZernikeMode] = list() self._order = 0 @@ -67,13 +66,12 @@ def get_mode(self, idx: int) -> ZernikeMode: def __len__(self) -> int: return len(self._polynomial) - def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: - probe = rescale_probe_intensity( + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + return self._rescale_to_photon_count( generate_zernike_probe( geometry_provider.get_probe_geometry(), self._polynomial, radius_m=self.diameter_m.get_value() / 2.0, ), - geometry_provider.probe_photon_count, + geometry_provider, ) - return self._build_probe_modes(self._rng, probe, geometry_provider.num_scan_points) diff --git a/src/ptychodus/model/product/probe_positions/builder.py b/src/ptychodus/model/product/probe_positions/builder.py index b7a104b1b..8cfb7977c 100644 --- a/src/ptychodus/model/product/probe_positions/builder.py +++ b/src/ptychodus/model/product/probe_positions/builder.py @@ -1,6 +1,6 @@ from __future__ import annotations from abc import abstractmethod -from collections.abc import Iterable, Iterator, Sequence +from collections.abc import Iterator, Sequence import logging import numpy @@ -51,6 +51,12 @@ def __init__( self.jitter_radius_m = settings.jitter_radius_m.copy() self._add_parameter('jitter_radius_m', self.jitter_radius_m) + self.num_discard_at_start = settings.num_discard_at_start.copy() + self._add_parameter('num_discard_at_start', self.num_discard_at_start) + + self.num_discard_at_end = settings.num_discard_at_end.copy() + self._add_parameter('num_discard_at_end', self.num_discard_at_end) + def get_name(self) -> str: return self._name.get_value() @@ -125,21 +131,74 @@ def copy(self) -> ProbePositionsBuilder: pass @abstractmethod - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> Sequence[ProbePosition]: + """Return the raw, unconditioned probe positions in acquisition order. + + Implementations must NOT apply the trim, affine transform, or jitter; + `build` owns the conditioning pipeline. + """ pass - def _create_position_sequence( - self, positions: Iterable[ProbePosition] - ) -> ProbePositionSequence: + def build(self) -> ProbePositionSequence: + """Return the conditioned probe positions: trim, then affine, then jitter. + + Overriding this method is reserved for builders whose positions are + already conditioned; see `FromMemoryProbePositionsBuilder`. Every builder + that ingests raw instrument coordinates must leave it alone and implement + `_build_raw` instead. + """ + return self._condition_positions(self._build_raw()) + + def _condition_positions(self, positions: Sequence[ProbePosition]) -> ProbePositionSequence: + trimmed = self._trim_positions(positions) transform = self.get_transform() jitter_radius_m = self.jitter_radius_m.get_value() rng = self._rng if jitter_radius_m > 0.0 else None return ProbePositionSequence( - [*transform_probe_positions(positions, transform, rng, jitter_radius_m)] + [*transform_probe_positions(trimmed, transform, rng, jitter_radius_m)] ) + def _trim_positions(self, positions: Sequence[ProbePosition]) -> Sequence[ProbePosition]: + """Discard points from each end of the scan, in acquisition order. + + Surviving points keep their original scan indexes. Diffraction patterns + whose index falls outside the trimmed range are dropped downstream by + `AssembledDiffractionData.prepare_reconstruct_input`, which never + extrapolates beyond the position-index anchors. + """ + num_discard_at_start = self.num_discard_at_start.get_value() + num_discard_at_end = self.num_discard_at_end.get_value() + + if num_discard_at_start == 0 and num_discard_at_end == 0: + return positions + + num_positions = len(positions) + stop = num_positions - num_discard_at_end + + if stop <= num_discard_at_start: + logger.warning( + f'Discarding {num_discard_at_start} probe position(s) at the start and' + f' {num_discard_at_end} at the end leaves nothing of {num_positions}!' + ) + return [] + + return positions[num_discard_at_start:stop] + class FromMemoryProbePositionsBuilder(ProbePositionsBuilder): + """Probe positions that have already been conditioned. + + Two things produce these. Reconstruction output, which `ProcessingTaskMonitor` + re-assigns to the output product item on every reconstructor iteration (see + `model/processing/monitor.py`), and products loaded from HDF5/NPZ, whose + positions were conditioned before they were saved. In both cases the trim, + affine transform, and jitter have already been applied upstream. Re-applying + them here would corrupt position-corrected output a little more on every + iteration, and would move the positions out of the coordinate frame the + reconstructed object was solved in. `build` therefore deliberately bypasses + the conditioning pipeline. + """ + def __init__( self, rng: numpy.random.Generator, @@ -151,9 +210,6 @@ def __init__( self._settings = settings self._position_seq = ProbePositionSequence(position_seq) - # set identity transformation - self.assign_preset_transform(0) - def copy(self) -> FromMemoryProbePositionsBuilder: builder = FromMemoryProbePositionsBuilder(self._rng, self._settings, self._position_seq) @@ -162,9 +218,12 @@ def copy(self) -> FromMemoryProbePositionsBuilder: return builder - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> ProbePositionSequence: return self._position_seq + def build(self) -> ProbePositionSequence: + return self._build_raw() + class FromFileProbePositionsBuilder(ProbePositionsBuilder): def __init__( @@ -192,7 +251,7 @@ def copy(self) -> FromFileProbePositionsBuilder: return builder - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> ProbePositionSequence: file_path = self.file_path.get_value() file_type = self.file_type.get_value() logger.debug(f'Reading "{file_path}" as "{file_type}"') diff --git a/src/ptychodus/model/product/probe_positions/cartesian.py b/src/ptychodus/model/product/probe_positions/cartesian.py index fecf52cb1..11883e2bb 100644 --- a/src/ptychodus/model/product/probe_positions/cartesian.py +++ b/src/ptychodus/model/product/probe_positions/cartesian.py @@ -1,9 +1,10 @@ from __future__ import annotations +from collections.abc import Sequence from enum import IntEnum import numpy -from ptychodus.api.probe_positions import ProbePositionSequence +from ptychodus.api.probe_positions import ProbePosition from ptychodus.api.probe_positions_gen import generate_cartesian_probe_positions from .builder import ProbePositionsBuilder @@ -69,7 +70,7 @@ def copy(self) -> CartesianProbePositionsBuilder: def is_equilateral(self) -> bool: return self._variant.is_equilateral - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> Sequence[ProbePosition]: step_size_x_m = self.step_size_x_m.get_value() if self._variant.is_equilateral: @@ -88,4 +89,4 @@ def build(self) -> ProbePositionSequence: snake=self._variant.is_snaked, stagger=self._variant.is_staggered, ) - return self._create_position_sequence(positions) + return [*positions] diff --git a/src/ptychodus/model/product/probe_positions/concentric.py b/src/ptychodus/model/product/probe_positions/concentric.py index bb39f790b..35183af13 100644 --- a/src/ptychodus/model/product/probe_positions/concentric.py +++ b/src/ptychodus/model/product/probe_positions/concentric.py @@ -1,8 +1,9 @@ from __future__ import annotations +from collections.abc import Sequence import numpy -from ptychodus.api.probe_positions import ProbePositionSequence +from ptychodus.api.probe_positions import ProbePosition from ptychodus.api.probe_positions_gen import generate_concentric_probe_positions from .builder import ProbePositionsBuilder @@ -31,10 +32,10 @@ def copy(self) -> ConcentricProbePositionsBuilder: return builder - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> Sequence[ProbePosition]: positions = generate_concentric_probe_positions( self.radial_step_size_m.get_value(), self.num_shells.get_value(), self.num_points_1st_shell.get_value(), ) - return self._create_position_sequence(positions) + return [*positions] diff --git a/src/ptychodus/model/product/probe_positions/item.py b/src/ptychodus/model/product/probe_positions/item.py index 9ba4d7a96..5692e4a3c 100644 --- a/src/ptychodus/model/product/probe_positions/item.py +++ b/src/ptychodus/model/product/probe_positions/item.py @@ -108,7 +108,11 @@ def _rebuild(self) -> None: logger.exception('Failed to rebuild scan!') return - self._probe_positions = ProbePositionSequence(probe_positions) + # build() always returns a ProbePositionSequence, and the class has no + # mutators, so there is nothing to defend against by copying. This path + # runs once per reconstructor iteration; the old round-trip through + # Python dataclasses was O(N) every time. + self._probe_positions = probe_positions self._geometry = calculate_scan_geometry(probe_positions) self.notify_observers() diff --git a/src/ptychodus/model/product/probe_positions/item_factory.py b/src/ptychodus/model/product/probe_positions/item_factory.py index 4d4316964..e2a1c331a 100644 --- a/src/ptychodus/model/product/probe_positions/item_factory.py +++ b/src/ptychodus/model/product/probe_positions/item_factory.py @@ -23,12 +23,42 @@ def __init__( self._settings = settings self._builder_factory = builder_factory + def _warn_if_conditioning_ignored(self) -> None: + """Note that conditioning settings do not apply to in-memory positions. + + Positions supplied in memory come from reconstruction output or a product + loaded from file, so they are already conditioned. Batch mode reads + product-in.h5 through this path, where a user who sets a trim in + settings.ini would otherwise see it silently do nothing. Trim at ingest + instead -- ptychodus-bdp reads raw probe positions through the from-file + builder, which does condition them. + """ + settings = self._settings + is_conditioning_requested = ( + settings.num_discard_at_start.get_value() != 0 + or settings.num_discard_at_end.get_value() != 0 + or settings.jitter_radius_m.get_value() != 0.0 + or settings.affine00.get_value() != 1.0 + or settings.affine01.get_value() != 0.0 + or settings.affine02.get_value() != 0.0 + or settings.affine10.get_value() != 0.0 + or settings.affine11.get_value() != 1.0 + or settings.affine12.get_value() != 0.0 + ) + + if is_conditioning_requested: + logger.info( + 'Probe positions supplied in memory are already conditioned;' + ' ignoring the trim, affine transform, and jitter settings.' + ) + def create( self, position_seq: ProbePositionSequence | None = None ) -> ProbePositionsRepositoryItem: if position_seq is None: builder = self._builder_factory.create_default() else: + self._warn_if_conditioning_ignored() builder = FromMemoryProbePositionsBuilder(self._rng, self._settings, position_seq) return ProbePositionsRepositoryItem(self._rng, self._settings, builder) diff --git a/src/ptychodus/model/product/probe_positions/lissajous.py b/src/ptychodus/model/product/probe_positions/lissajous.py index 1ae42b536..0cdcad0cc 100644 --- a/src/ptychodus/model/product/probe_positions/lissajous.py +++ b/src/ptychodus/model/product/probe_positions/lissajous.py @@ -1,8 +1,9 @@ from __future__ import annotations +from collections.abc import Sequence import numpy -from ptychodus.api.probe_positions import ProbePositionSequence +from ptychodus.api.probe_positions import ProbePosition from ptychodus.api.probe_positions_gen import generate_lissajous_probe_positions from .builder import ProbePositionsBuilder @@ -48,7 +49,7 @@ def copy(self) -> LissajousProbePositionsBuilder: return builder - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> Sequence[ProbePosition]: positions = generate_lissajous_probe_positions( self.num_points.get_value(), self.amplitude_x_m.get_value(), @@ -57,4 +58,4 @@ def build(self) -> ProbePositionSequence: self.angular_step_y_turns.get_value(), self.angular_shift_turns.get_value(), ) - return self._create_position_sequence(positions) + return [*positions] diff --git a/src/ptychodus/model/product/probe_positions/settings.py b/src/ptychodus/model/product/probe_positions/settings.py index 63f6b2a50..c1b4eb5fd 100644 --- a/src/ptychodus/model/product/probe_positions/settings.py +++ b/src/ptychodus/model/product/probe_positions/settings.py @@ -23,6 +23,12 @@ def __init__(self, registry: SettingsRegistry) -> None: self.jitter_radius_m = self._group.create_real_parameter( 'JitterRadiusInMeters', 0.0, minimum=0.0 ) + self.num_discard_at_start = self._group.create_integer_parameter( + 'NumberOfPointsToDiscardAtStart', 0, minimum=0 + ) + self.num_discard_at_end = self._group.create_integer_parameter( + 'NumberOfPointsToDiscardAtEnd', 0, minimum=0 + ) self.expand_bbox = self._group.create_boolean_parameter('ExpandBoundingBox', False) self.expand_bbox_xmin_m = self._group.create_real_parameter( diff --git a/src/ptychodus/model/product/probe_positions/spiral.py b/src/ptychodus/model/product/probe_positions/spiral.py index b0b0f991c..f7ac24ff7 100644 --- a/src/ptychodus/model/product/probe_positions/spiral.py +++ b/src/ptychodus/model/product/probe_positions/spiral.py @@ -1,8 +1,9 @@ from __future__ import annotations +from collections.abc import Sequence import numpy -from ptychodus.api.probe_positions import ProbePositionSequence +from ptychodus.api.probe_positions import ProbePosition from ptychodus.api.probe_positions_gen import generate_spiral_probe_positions from .builder import ProbePositionsBuilder @@ -38,8 +39,8 @@ def copy(self) -> SpiralProbePositionsBuilder: return builder - def build(self) -> ProbePositionSequence: + def _build_raw(self) -> Sequence[ProbePosition]: positions = generate_spiral_probe_positions( self.num_points.get_value(), self.radius_scalar_m.get_value() ) - return self._create_position_sequence(positions) + return [*positions] diff --git a/src/ptychodus/model/product/probe_positions/streaming.py b/src/ptychodus/model/product/probe_positions/streaming.py index 202dc201b..416ffd872 100644 --- a/src/ptychodus/model/product/probe_positions/streaming.py +++ b/src/ptychodus/model/product/probe_positions/streaming.py @@ -1,33 +1,50 @@ +from __future__ import annotations from collections.abc import Sequence import numpy # TODO from pvaccess import Channel, PvObjectQueue -from ptychodus.api.probe_positions import ProbePositionSequence, ProbePosition +from ptychodus.api.probe_positions import ProbePosition from .builder import ProbePositionsBuilder from .settings import ProbePositionsSettings class StreamingScanBuilder(ProbePositionsBuilder): + # TODO The "discard at end" trim chases a moving tail while the stream is + # still growing, so each build drops a different set of trailing points. + # Decide on the semantics (most likely: honor the head trim, ignore the tail + # trim until the stream is marked complete) before wiring up the pvaccess + # path below. + def __init__( self, rng: numpy.random.Generator, settings: ProbePositionsSettings, point_seq: Sequence[ProbePosition], ) -> None: - super().__init__(rng, settings, 'Streaming') + super().__init__(rng, settings, 'streaming') + self._settings = settings self._point_list = list(point_seq) + def copy(self) -> StreamingScanBuilder: + builder = StreamingScanBuilder(self._rng, self._settings, self._point_list) + + for key, value in self.parameters().items(): + builder.parameters()[key].set_value(value.get_value()) + + return builder + def append(self, point: ProbePosition) -> None: self._point_list.append(point) def extend(self, point_seq: Sequence[ProbePosition]) -> None: self._point_list.extend(point_seq) - def build(self) -> ProbePositionSequence: - return ProbePositionSequence(self._point_list) + def _build_raw(self) -> Sequence[ProbePosition]: + # Snapshot the list so a concurrent append cannot tear the trim. + return [*self._point_list] # TODO def echo(self, value: int = 125) -> None: diff --git a/src/ptychodus/model/product/probe_repository.py b/src/ptychodus/model/product/probe_repository.py index c2e833381..d6b3db148 100644 --- a/src/ptychodus/model/product/probe_repository.py +++ b/src/ptychodus/model/product/probe_repository.py @@ -5,6 +5,7 @@ from ptychodus.api.observer import ObservableSequence from ptychodus.api.product import LossValue +from ..diffraction import AssembledDiffractionDataset from .item import ProductRepositoryItem, ProductRepositoryObserver from .metadata import MetadataRepositoryItem from .object import ObjectRepositoryItem @@ -27,6 +28,9 @@ def get_name(self, index: int) -> str: def set_name(self, index: int, name: str) -> None: self._repository[index].set_name(name) + def get_dataset(self, index: int) -> AssembledDiffractionDataset | None: + return self._repository[index].get_dataset() + @overload def __getitem__(self, index: int) -> ProbeRepositoryItem: ... @@ -64,5 +68,11 @@ def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: pass + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: self.notify_observers_item_removed(index, item.get_probe_item()) diff --git a/src/ptychodus/model/product/repository.py b/src/ptychodus/model/product/repository.py index a1ebb359a..34bc9891e 100644 --- a/src/ptychodus/model/product/repository.py +++ b/src/ptychodus/model/product/repository.py @@ -136,3 +136,23 @@ def handle_losses_changed(self, item: ProductRepositoryItem) -> None: for observer in self._observer_list: observer.handle_losses_changed(index, losses) + + def handle_dataset_changed(self, item: ProductRepositoryItem) -> None: + index = item._index + + if index < 0: + logger.warning(f'Failed to look up index for "{item.get_name()}"!') + return + + for observer in self._observer_list: + observer.handle_dataset_changed(index, item) + + def handle_state_changed(self, item: ProductRepositoryItem) -> None: + index = item._index + + if index < 0: + logger.warning(f'Failed to look up index for "{item.get_name()}"!') + return + + for observer in self._observer_list: + observer.handle_state_changed(index, item) diff --git a/src/ptychodus/model/product/scan_repository.py b/src/ptychodus/model/product/scan_repository.py index 1f608e363..dcda1a11a 100644 --- a/src/ptychodus/model/product/scan_repository.py +++ b/src/ptychodus/model/product/scan_repository.py @@ -66,5 +66,11 @@ def handle_object_changed(self, index: int, item: ObjectRepositoryItem) -> None: def handle_losses_changed(self, index: int, losses: Sequence[LossValue]) -> None: pass + def handle_dataset_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + + def handle_state_changed(self, index: int, item: ProductRepositoryItem) -> None: + pass + def handle_item_removed(self, index: int, item: ProductRepositoryItem) -> None: self.notify_observers_item_removed(index, item.get_probe_positions_item()) diff --git a/src/ptychodus/model/product/settings.py b/src/ptychodus/model/product/settings.py index a1ddb94c5..4c6af7cbb 100644 --- a/src/ptychodus/model/product/settings.py +++ b/src/ptychodus/model/product/settings.py @@ -31,6 +31,8 @@ def __init__(self, registry: SettingsRegistry) -> None: self.tomography_angle_deg = self._group.create_real_parameter( 'TomographyAngleInDegrees', 0.0 ) + self.tilt_angle_deg = self._group.create_real_parameter('TiltAngleInDegrees', 0.0) + self.polarization = self._group.create_string_parameter('Polarization', '') def _update(self, observable: Observable) -> None: if observable is self._group: diff --git a/src/ptychodus/model/ptychi/_payload.py b/src/ptychodus/model/ptychi/_payload.py new file mode 100644 index 000000000..a2f38e61b --- /dev/null +++ b/src/ptychodus/model/ptychi/_payload.py @@ -0,0 +1,31 @@ +"""Payload dataclass for the PtyChi subprocess entry point. + +Parent-safe to import: pulls ``PtychographyTaskOptions`` from ``ptychi.api``, +which imports torch for its type annotations but does not acquire a GPU +context — that happens only when the child instantiates ``PtychographyTask``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ptychi.api.options.task import PtychographyTaskOptions + +from ptychodus.api.reconstructor import ReconstructInput + + +@dataclass(frozen=True) +class PtyChiPayload: + """Everything the child needs to acquire a GPU context and run one reconstruction. + + ``task_options`` is a fully-populated algorithm-specific subclass of + :class:`PtychographyTaskOptions` (``DMOptions``, ``PIEOptions``, + ``LSQMLOptions``, ...) that the parent built via the option-helper classes + in :mod:`.helper` and the per-algorithm classes in :mod:`.dm`, :mod:`.pie`, + etc. It carries all diffraction data, positions, and initial guesses; the + child hands it straight to ``PtychographyTask``. + """ + + task_options: PtychographyTaskOptions + num_sync_epochs: int + reconstruct_input: ReconstructInput diff --git a/src/ptychodus/model/ptychi/_subprocess.py b/src/ptychodus/model/ptychi/_subprocess.py new file mode 100644 index 000000000..2bd20a4d9 --- /dev/null +++ b/src/ptychodus/model/ptychi/_subprocess.py @@ -0,0 +1,131 @@ +"""Child-side subprocess entry points for the PtyChi backend. + +Every function in this module runs INSIDE a spawned subprocess and touches +the GPU. This is the only place in the ptychodus tree that instantiates +``ptychi.api.task.PtychographyTask`` — that's the step that acquires a CUDA +context. Everything else (option translation, settings reads) already ran +parent-side; the child receives a finished :class:`PtychographyTaskOptions` +and streams outputs back. + +The device-probing entry point is a one-shot: it enumerates devices and +exits. Parent-side callers spawn it via :mod:`.device` when populating +:class:`PtyChiDeviceRepository`. +""" + +from __future__ import annotations + +import logging +import pickle +from collections.abc import Iterator +from multiprocessing.queues import Queue +from typing import Any + +import numpy + +from ptychodus.api.object import Object, ObjectPosition +from ptychodus.api.probe import ProbeSequence +from ptychodus.api.probe_positions import ProbePosition, ProbePositionSequence +from ptychodus.api.product import LossValue, Product +from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput + +from ..processing.subprocess_reconstructor import TAG_OUTPUT +from ._payload import PtyChiPayload + +logger = logging.getLogger(__name__) + + +def reconstruct_with_ptychi( + parameters: ReconstructInput, + payload: PtyChiPayload, +) -> Iterator[ReconstructOutput]: + """Instantiate ``PtychographyTask`` and yield a ``ReconstructOutput`` every + ``num_sync_epochs`` epochs. The ``PtychographyTask`` import is deferred to + call time so parent-side test collection of this module (if it ever + happens) does not acquire a GPU context.""" + from ptychi.api.task import PtychographyTask + + task_options = payload.task_options + num_sync_epochs = payload.num_sync_epochs + num_epochs = task_options.reconstructor_options.num_epochs + task = PtychographyTask(task_options) + + with task: + epoch = 0 + + task_reconstructor = task.reconstructor + + if task_reconstructor is None: + raise RuntimeError('Task reconstructor is None!') + + loss_tracker = task_reconstructor.loss_tracker + + while epoch < num_epochs: + step_epochs = min(num_sync_epochs, num_epochs - epoch) + task.run(step_epochs) + + losses: list[LossValue] = list() + epoch_array = loss_tracker.table['epoch'].to_numpy() + loss_array = loss_tracker.table['loss'].to_numpy() + + for e, loss in zip(epoch_array.flat, loss_array.flat): + losses.append(LossValue(epoch=e, value=loss.item())) + + product_in = parameters.product + object_in = product_in.object_ + object_out = Object( + array=numpy.array(task.get_data_to_cpu('object', as_numpy=True)), + layer_spacing_m=object_in.layer_spacing_m, + pixel_geometry=object_in.get_pixel_geometry(), + center=object_in.get_center(), + ) + probe_out = ProbeSequence( + array=numpy.array(task.get_data_to_cpu('probe', as_numpy=True)), + opr_weights=numpy.array(task.get_data_to_cpu('opr_mode_weights', as_numpy=True)), + pixel_geometry=product_in.probes.get_pixel_geometry(), + ) + + position_x_px = task.get_probe_positions_x(as_numpy=True) + position_y_px = task.get_probe_positions_y(as_numpy=True) + object_geometry = object_in.get_geometry() + corrected_scan_points: list[ProbePosition] = list() + + for uncorrected_point, pos_x_px, pos_y_px in zip( + product_in.probe_positions, position_x_px, position_y_px + ): + object_point = ObjectPosition( + index=uncorrected_point.index, + coordinate_x_px=float(pos_x_px), + coordinate_y_px=float(pos_y_px), + ) + scan_point = object_geometry.map_coordinates_object_to_probe(object_point) + corrected_scan_points.append(scan_point) + + product = Product( + metadata=product_in.metadata, + probe_positions=ProbePositionSequence(corrected_scan_points), + probes=probe_out, + object_=object_out, + losses=losses, + ) + + epoch += step_epochs + + yield ReconstructOutput(product=product, progress=epoch) + + +def run_reconstruct(payload: PtyChiPayload, queue: Queue[Any]) -> None: + """Child entry point. Acquire a GPU context via PtychographyTask, stream outputs.""" + for output in reconstruct_with_ptychi(payload.reconstruct_input, payload): + queue.put((TAG_OUTPUT, pickle.dumps(output))) + + +def probe_device_list() -> list[str]: + """One-shot device enumeration for the parent-side device repository.""" + import ptychi + + return [f'{d.name} ({d.torch_device})' for d in ptychi.list_available_devices()] + + +def probe_devices(_payload: Any, queue: Queue[Any]) -> None: + """Spawn-safe entry point that emits the device list on the queue and exits.""" + queue.put((TAG_OUTPUT, pickle.dumps(probe_device_list()))) diff --git a/src/ptychodus/model/ptychi/autodiff.py b/src/ptychodus/model/ptychi/autodiff.py index ad87dffe6..967fc8284 100644 --- a/src/ptychodus/model/ptychi/autodiff.py +++ b/src/ptychodus/model/ptychi/autodiff.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -12,12 +11,10 @@ ForwardModels, LossFunctions, ) -from ptychi.api.task import PtychographyTask - from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -26,18 +23,12 @@ logger = logging.getLogger(__name__) -class AutodiffReconstructor(Reconstructor): +class AutodiffReconstructor: def __init__( self, options_helper: PtyChiOptionsHelper, settings: PtyChiAutodiffSettings ) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'Autodiff' def _create_reconstructor_options(self) -> AutodiffPtychographyReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -176,49 +167,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> AutodiffPtychogr ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptychi/bh.py b/src/ptychodus/model/ptychi/bh.py index 6849d6f0b..85f436e50 100644 --- a/src/ptychodus/model/ptychi/bh.py +++ b/src/ptychodus/model/ptychi/bh.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -10,12 +9,10 @@ BHProbePositionOptions, BHReconstructorOptions, ) -from ptychi.api.task import PtychographyTask - from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -24,16 +21,10 @@ logger = logging.getLogger(__name__) -class BHReconstructor(Reconstructor): +class BHReconstructor: def __init__(self, options_helper: PtyChiOptionsHelper, settings: PtyChiBHSettings) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'BH' def _create_reconstructor_options(self) -> BHReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -147,49 +138,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> BHOptions: ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptychi/core.py b/src/ptychodus/model/ptychi/core.py index aa8bb7a03..b5d9a5911 100644 --- a/src/ptychodus/model/ptychi/core.py +++ b/src/ptychodus/model/ptychi/core.py @@ -1,5 +1,6 @@ from collections.abc import Iterator -from importlib.metadata import version +from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec import logging from ptychodus.api.reconstructor import ( @@ -28,6 +29,11 @@ logger = logging.getLogger(__name__) +def _ptychi_available() -> bool: + """Return True iff ``ptychi`` is importable, without importing it.""" + return find_spec('ptychi') is not None + + class PtyChiReconstructorLibrary(ReconstructorLibrary): def __init__( self, @@ -53,41 +59,43 @@ def __init__( ) self.reconstructor_list: list[Reconstructor] = list() - try: - from .autodiff import AutodiffReconstructor - from .bh import BHReconstructor - from .dm import DMReconstructor - from .epie import EPIEReconstructor - from .helper import PtyChiOptionsHelper - from .lsqml import LSQMLReconstructor - from .pie import PIEReconstructor - from .rpie import RPIEReconstructor - except ModuleNotFoundError: + if not _ptychi_available(): logger.info('pty-chi not found.') if is_developer_mode_enabled: for reconstructor in ('DM', 'PIE', 'ePIE', 'rPIE', 'LSQML', 'Autodiff', 'BH'): self.reconstructor_list.append(NullReconstructor(reconstructor)) - else: - logger.info('Pty-Chi ' + version('ptychi')) + return + + try: + ptychi_version = version('ptychi') + except PackageNotFoundError: + ptychi_version = 'unknown' + logger.info(f'Pty-Chi {ptychi_version}') - options_helper = PtyChiOptionsHelper( + # Parent-side factory. Imports ptychi.api transitively (via .helper and + # per-algorithm modules) — that pulls torch but does not acquire a GPU + # context; see the invariant note in _subprocess_protocol.py. + from .reconstructor import PtyChiSettingsBundle, build_reconstructor_list + + bundle = PtyChiSettingsBundle( + dm=self.dm_settings, + pie=self.pie_settings, + lsqml=self.lsqml_settings, + autodiff=self.autodiff_settings, + bh=self.bh_settings, + ) + self.reconstructor_list.extend( + build_reconstructor_list( self.settings, self.object_settings, self.probe_settings, self.probe_position_settings, self.opr_settings, + bundle, pattern_sizer, ) - self.reconstructor_list.append(DMReconstructor(options_helper, self.dm_settings)) - self.reconstructor_list.append(PIEReconstructor(options_helper, self.pie_settings)) - self.reconstructor_list.append(EPIEReconstructor(options_helper, self.pie_settings)) - self.reconstructor_list.append(RPIEReconstructor(options_helper, self.pie_settings)) - self.reconstructor_list.append(LSQMLReconstructor(options_helper, self.lsqml_settings)) - self.reconstructor_list.append( - AutodiffReconstructor(options_helper, self.autodiff_settings) - ) - self.reconstructor_list.append(BHReconstructor(options_helper, self.bh_settings)) + ) @property def name(self) -> str: diff --git a/src/ptychodus/model/ptychi/device.py b/src/ptychodus/model/ptychi/device.py index 1518c80cd..09d2381ec 100644 --- a/src/ptychodus/model/ptychi/device.py +++ b/src/ptychodus/model/ptychi/device.py @@ -1,23 +1,60 @@ +"""Parent-safe device enumeration for the PtyChi backend. + +Historically this module called ``ptychi.list_available_devices()`` at import +time inside the parent process, which pulled ptychi (and hence torch/CuPy) +into the parent's ``sys.modules``. The refactor moves the probe into a +one-shot spawned subprocess whose result is cached on the repository +instance. +""" + +from __future__ import annotations + from collections.abc import Sequence -from typing import overload +from importlib.util import find_spec +from typing import Any, overload import logging +import pickle + +from ..processing._subprocess_protocol import ChildError, run_subprocess logger = logging.getLogger(__name__) +_PROBE_ENTRY = 'ptychodus.model.ptychi._subprocess:probe_devices' + + +def _probe_devices_via_subprocess() -> list[str]: + """Spawn a child that calls ``ptychi.list_available_devices()`` and return the list. + + Any failure -- ptychi not installed, subprocess crash, timeout -- is + logged and an empty list is returned. This keeps device probing + non-fatal for parents whose GPU stack is temporarily broken. + """ + try: + with run_subprocess(_PROBE_ENTRY, None, terminate_grace_sec=5.0) as events: + for event in events: + if event[0] == 'output': + devices = pickle.loads(event[1]) + if isinstance(devices, list): + return list(devices) + except ChildError as exc: + logger.warning('Device probe subprocess failed: %s', exc.child_exception_type) + except Exception: + logger.exception('Device probe subprocess raised in the parent.') + return [] + + class PtyChiDeviceRepository(Sequence[str]): def __init__(self, *, is_developer_mode_enabled: bool) -> None: self._devices: list[str] = list() - try: - import ptychi - except ModuleNotFoundError: + if find_spec('ptychi') is None: if is_developer_mode_enabled: self._devices.extend(f'gpu:{n}' for n in range(4)) else: - for device in ptychi.list_available_devices(): + for device in _probe_devices_via_subprocess(): logger.info(device) - self._devices.append(f'{device.name} ({device.torch_device})') + self._devices.append(device) if not self._devices: logger.info('No devices found!') @@ -33,3 +70,15 @@ def __getitem__(self, index: int | slice) -> str | Sequence[str]: def __len__(self) -> int: return len(self._devices) + + # Kept for callers that want to force a re-enumeration after e.g. plugging + # in a new device. Not used inside ptychodus today. + def refresh(self) -> None: + if find_spec('ptychi') is None: + return + self._devices = _probe_devices_via_subprocess() + + +# `_PROBE_PAYLOAD` is None; keeping this alias documents that the entry point +# ignores its payload argument. +_PROBE_PAYLOAD: Any = None diff --git a/src/ptychodus/model/ptychi/dm.py b/src/ptychodus/model/ptychi/dm.py index 3ba40982c..bb3146d2c 100644 --- a/src/ptychodus/model/ptychi/dm.py +++ b/src/ptychodus/model/ptychi/dm.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -10,12 +9,11 @@ DMProbePositionOptions, DMReconstructorOptions, ) -from ptychi.api.task import PtychographyTask from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -24,16 +22,10 @@ logger = logging.getLogger(__name__) -class DMReconstructor(Reconstructor): +class DMReconstructor: def __init__(self, options_helper: PtyChiOptionsHelper, settings: PtyChiDMSettings) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'DM' def _create_reconstructor_options(self) -> DMReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -149,49 +141,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> DMOptions: ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptychi/epie.py b/src/ptychodus/model/ptychi/epie.py index bb5c7be57..8711cb989 100644 --- a/src/ptychodus/model/ptychi/epie.py +++ b/src/ptychodus/model/ptychi/epie.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -10,12 +9,10 @@ PIEProbeOptions, PIEProbePositionOptions, ) -from ptychi.api.task import PtychographyTask - from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -24,16 +21,10 @@ logger = logging.getLogger(__name__) -class EPIEReconstructor(Reconstructor): +class EPIEReconstructor: def __init__(self, options_helper: PtyChiOptionsHelper, settings: PtyChiPIESettings) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'ePIE' def _create_reconstructor_options(self) -> EPIEReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -146,49 +137,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> EPIEOptions: ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptychi/helper.py b/src/ptychodus/model/ptychi/helper.py index 394a161d5..c1fe69ffc 100644 --- a/src/ptychodus/model/ptychi/helper.py +++ b/src/ptychodus/model/ptychi/helper.py @@ -1,7 +1,18 @@ -from collections.abc import Sequence +"""Parent-safe pty-chi option builders. + +Every class here reads ptychodus settings via ``.get_value()`` and constructs +pydantic ``*Options`` dataclasses from ``ptychi.api``. None of it acquires a +GPU context — importing ``ptychi.api`` pulls torch in for its type +annotations, but no CUDA runtime is initialised until a ``PtychographyTask`` +is actually constructed. That happens child-side in ``_subprocess.py``. + +The parent-side factory in ``reconstructor.py`` uses these builders to +assemble a fully-populated ``PtychographyTaskOptions`` and ships it as the +subprocess payload. +""" + import logging -import torch import numpy import math @@ -48,10 +59,10 @@ ) from ptychodus.api.common import ComplexArrayType, RealArrayType -from ptychodus.api.object import Object, ObjectGeometry, ObjectPosition +from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.probe_positions import ProbePositionSequence, ProbePosition -from ptychodus.api.product import LossValue, Product, ProductMetadata +from ptychodus.api.probe_positions import ProbePositionSequence +from ptychodus.api.product import ProductMetadata from ptychodus.api.reconstructor import ReconstructInput from ..diffraction import PatternSizer @@ -742,7 +753,7 @@ def __init__( def create_data_options(self, parameters: ReconstructInput) -> PtychographyDataOptions: metadata = parameters.product.metadata - pixel_geometry = self._pattern_sizer.get_processed_pixel_geometry() + pixel_geometry = self._pattern_sizer.get_processed_pixel_geometry(parameters.pixel_geometry) free_space_propagation_distance_m = ( numpy.inf if self._reconstructor_settings.use_far_field_propagation @@ -758,54 +769,6 @@ def create_data_options(self, parameters: ReconstructInput) -> PtychographyDataO save_data_on_device=self._reconstructor_settings.save_data_on_device.get_value(), ) - def create_product( - self, - product: Product, - position_x_px: torch.Tensor | numpy.ndarray, - position_y_px: torch.Tensor | numpy.ndarray, - probe_array: torch.Tensor | numpy.ndarray, - object_array: torch.Tensor | numpy.ndarray, - opr_weights: torch.Tensor | numpy.ndarray, - losses: Sequence[LossValue], - ) -> Product: - object_in = product.object_ - object_out = Object( - array=numpy.array(object_array), - layer_spacing_m=object_in.layer_spacing_m, - pixel_geometry=object_in.get_pixel_geometry(), - center=object_in.get_center(), - ) - - probe_out = ProbeSequence( - array=numpy.array(probe_array), - opr_weights=numpy.array(opr_weights), - pixel_geometry=product.probes.get_pixel_geometry(), - ) - - corrected_scan_points: list[ProbePosition] = list() - object_geometry = object_in.get_geometry() - - for uncorrected_point, pos_x_px, pos_y_px in zip( - product.probe_positions, position_x_px, position_y_px - ): - object_point = ObjectPosition( - index=uncorrected_point.index, - coordinate_x_px=float(pos_x_px), - coordinate_y_px=float(pos_y_px), - ) - scan_point = object_geometry.map_coordinates_object_to_probe(object_point) - corrected_scan_points.append(scan_point) - - scan_out = ProbePositionSequence(corrected_scan_points) - - return Product( - metadata=product.metadata, - probe_positions=scan_out, - probes=probe_out, - object_=object_out, - losses=losses, - ) - @property def num_epochs(self) -> int: return self.reconstructor_helper.num_epochs diff --git a/src/ptychodus/model/ptychi/lsqml.py b/src/ptychodus/model/ptychi/lsqml.py index eea1cbfc9..08e0afa44 100644 --- a/src/ptychodus/model/ptychi/lsqml.py +++ b/src/ptychodus/model/ptychi/lsqml.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -11,12 +10,10 @@ LSQMLReconstructorOptions, NoiseModels, ) -from ptychi.api.task import PtychographyTask - from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -25,16 +22,10 @@ logger = logging.getLogger(__name__) -class LSQMLReconstructor(Reconstructor): +class LSQMLReconstructor: def __init__(self, options_helper: PtyChiOptionsHelper, settings: PtyChiLSQMLSettings) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'LSQML' def _create_reconstructor_options(self) -> LSQMLReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -187,49 +178,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> LSQMLOptions: ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptychi/pie.py b/src/ptychodus/model/ptychi/pie.py index a94e3ecaf..8a338064c 100644 --- a/src/ptychodus/model/ptychi/pie.py +++ b/src/ptychodus/model/ptychi/pie.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -10,12 +9,10 @@ PIEProbePositionOptions, PIEReconstructorOptions, ) -from ptychi.api.task import PtychographyTask - from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -24,16 +21,10 @@ logger = logging.getLogger(__name__) -class PIEReconstructor(Reconstructor): +class PIEReconstructor: def __init__(self, options_helper: PtyChiOptionsHelper, settings: PtyChiPIESettings) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'PIE' def _create_reconstructor_options(self) -> PIEReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -146,49 +137,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> PIEOptions: ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptychi/reconstructor.py b/src/ptychodus/model/ptychi/reconstructor.py new file mode 100644 index 000000000..34e8f49e2 --- /dev/null +++ b/src/ptychodus/model/ptychi/reconstructor.py @@ -0,0 +1,153 @@ +"""Parent-side factory that builds :class:`SubprocessReconstructor`s for PtyChi. + +This module imports ``ptychi.api`` (via the per-algorithm modules and +:mod:`.helper`) to construct the ``*Options`` dataclasses that make up a +``PtychographyTaskOptions``. That import chain pulls torch in — but torch's +CUDA runtime is lazy, so no GPU context is acquired here. The context is +acquired only in the spawned child when it instantiates ``PtychographyTask`` +(see :mod:`._subprocess`). +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from ptychodus.api.reconstructor import ReconstructInput + +from ..diffraction import PatternSizer +from ..processing.subprocess_reconstructor import SubprocessReconstructor +from ._payload import PtyChiPayload +from .autodiff import AutodiffReconstructor +from .bh import BHReconstructor +from .dm import DMReconstructor +from .epie import EPIEReconstructor +from .helper import PtyChiOptionsHelper +from .lsqml import LSQMLReconstructor +from .pie import PIEReconstructor +from .rpie import RPIEReconstructor +from .settings import ( + PtyChiAutodiffSettings, + PtyChiBHSettings, + PtyChiDMSettings, + PtyChiLSQMLSettings, + PtyChiOPRSettings, + PtyChiObjectSettings, + PtyChiPIESettings, + PtyChiProbePositionSettings, + PtyChiProbeSettings, + PtyChiSettings, +) + +__all__ = ['build_reconstructor_list'] + + +_RECONSTRUCT_ENTRY = 'ptychodus.model.ptychi._subprocess:run_reconstruct' + + +# Callable that turns a ReconstructInput into a fully-built pty-chi +# ``PtychographyTaskOptions`` (algorithm-specific subclass). ``Any`` because the +# concrete return type lives in ``ptychi.api.options.task`` and typing it here +# would drag that module into non-ptychi environments. +_TaskOptionsBuilder = Callable[[ReconstructInput], Any] + + +@dataclass(frozen=True) +class _AlgorithmSpec: + """One row of the algorithm dispatch table used by the factory.""" + + display_name: str + make_option_factory: Callable[[PtyChiOptionsHelper, PtyChiSettingsBundle], _TaskOptionsBuilder] + + +@dataclass(frozen=True) +class PtyChiSettingsBundle: + """The algorithm-specific settings groups held by :class:`PtyChiReconstructorLibrary`.""" + + dm: PtyChiDMSettings + pie: PtyChiPIESettings + lsqml: PtyChiLSQMLSettings + autodiff: PtyChiAutodiffSettings + bh: PtyChiBHSettings + + +_ALGORITHMS: tuple[_AlgorithmSpec, ...] = ( + _AlgorithmSpec( + 'DM', lambda helper, bundle: DMReconstructor(helper, bundle.dm)._create_task_options + ), + _AlgorithmSpec( + 'PIE', lambda helper, bundle: PIEReconstructor(helper, bundle.pie)._create_task_options + ), + _AlgorithmSpec( + 'ePIE', lambda helper, bundle: EPIEReconstructor(helper, bundle.pie)._create_task_options + ), + _AlgorithmSpec( + 'rPIE', lambda helper, bundle: RPIEReconstructor(helper, bundle.pie)._create_task_options + ), + _AlgorithmSpec( + 'LSQML', + lambda helper, bundle: LSQMLReconstructor(helper, bundle.lsqml)._create_task_options, + ), + _AlgorithmSpec( + 'Autodiff', + lambda helper, bundle: AutodiffReconstructor(helper, bundle.autodiff)._create_task_options, + ), + _AlgorithmSpec( + 'BH', lambda helper, bundle: BHReconstructor(helper, bundle.bh)._create_task_options + ), +) + + +def build_reconstructor_list( + reconstructor_settings: PtyChiSettings, + object_settings: PtyChiObjectSettings, + probe_settings: PtyChiProbeSettings, + probe_position_settings: PtyChiProbePositionSettings, + opr_settings: PtyChiOPRSettings, + bundle: PtyChiSettingsBundle, + pattern_sizer: PatternSizer, +) -> list[SubprocessReconstructor]: + """Build one :class:`SubprocessReconstructor` per pty-chi algorithm.""" + options_helper = PtyChiOptionsHelper( + reconstructor_settings, + object_settings, + probe_settings, + probe_position_settings, + opr_settings, + pattern_sizer, + ) + + def num_epochs() -> int: + return reconstructor_settings.num_epochs.get_value() + + def num_sync_epochs() -> int: + return options_helper.num_sync_epochs + + reconstructors: list[SubprocessReconstructor] = [] + + for spec in _ALGORITHMS: + build_task_options = spec.make_option_factory(options_helper, bundle) + + def build_payload( + parameters: ReconstructInput, + _loaded_model_path: Path | None, + _build: _TaskOptionsBuilder = build_task_options, + ) -> PtyChiPayload: + return PtyChiPayload( + task_options=_build(parameters), + num_sync_epochs=num_sync_epochs(), + reconstruct_input=parameters, + ) + + reconstructors.append( + SubprocessReconstructor( + name=spec.display_name, + reconstruct_entry_point=_RECONSTRUCT_ENTRY, + progress_goal_fn=num_epochs, + build_reconstruct_payload=build_payload, + ) + ) + + return reconstructors diff --git a/src/ptychodus/model/ptychi/rpie.py b/src/ptychodus/model/ptychi/rpie.py index db956e021..9d125f524 100644 --- a/src/ptychodus/model/ptychi/rpie.py +++ b/src/ptychodus/model/ptychi/rpie.py @@ -1,4 +1,3 @@ -from collections.abc import Iterator import logging @@ -10,12 +9,10 @@ RPIEOptions, RPIEReconstructorOptions, ) -from ptychi.api.task import PtychographyTask - from ptychodus.api.object import Object, ObjectGeometry from ptychodus.api.probe import ProbeSequence -from ptychodus.api.product import LossValue, ProductMetadata -from ptychodus.api.reconstructor import ReconstructInput, ReconstructOutput, Reconstructor +from ptychodus.api.product import ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput from ptychodus.api.probe_positions import ProbePositionSequence from .helper import PtyChiOptionsHelper @@ -24,16 +21,10 @@ logger = logging.getLogger(__name__) -class RPIEReconstructor(Reconstructor): +class RPIEReconstructor: def __init__(self, options_helper: PtyChiOptionsHelper, settings: PtyChiPIESettings) -> None: - super().__init__() self._options_helper = options_helper self._settings = settings - self._epoch = 0 - - @property - def name(self) -> str: - return 'rPIE' def _create_reconstructor_options(self) -> RPIEReconstructorOptions: helper = self._options_helper.reconstructor_helper @@ -146,49 +137,3 @@ def _create_task_options(self, parameters: ReconstructInput) -> RPIEOptions: ), opr_mode_weight_options=self._create_opr_mode_weight_options(product.probes), ) - - def get_progress_goal(self) -> int: - return self._options_helper.num_epochs - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - task_options = self._create_task_options(parameters) - num_epochs = task_options.reconstructor_options.num_epochs - - task = PtychographyTask(task_options) - - with task: - self._epoch = 0 - step_epochs = self._options_helper.num_sync_epochs - - task_reconstructor = task.reconstructor - - if task_reconstructor is None: - raise RuntimeError('Task reconstructor is None!') - - loss_tracker = task_reconstructor.loss_tracker - - while self._epoch < num_epochs: - task.run(step_epochs) - - losses: list[LossValue] = list() - epoch_array = loss_tracker.table['epoch'].to_numpy() - loss_array = loss_tracker.table['loss'].to_numpy() - - for epoch, loss in zip(epoch_array.flat, loss_array.flat): - loss_value = LossValue(epoch=epoch, value=loss.item()) - losses.append(loss_value) - - product = self._options_helper.create_product( - product=parameters.product, - position_x_px=task.get_probe_positions_x(as_numpy=True), - position_y_px=task.get_probe_positions_y(as_numpy=True), - probe_array=task.get_data_to_cpu('probe', as_numpy=True), - object_array=task.get_data_to_cpu('object', as_numpy=True), - opr_weights=task.get_data_to_cpu('opr_mode_weights', as_numpy=True), - losses=losses, - ) - - self._epoch += step_epochs - step_epochs = min(step_epochs, num_epochs - self._epoch) - - yield ReconstructOutput(product=product, progress=self._epoch) diff --git a/src/ptychodus/model/ptycho_fm/__init__.py b/src/ptychodus/model/ptycho_fm/__init__.py new file mode 100644 index 000000000..fa97a24c5 --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/__init__.py @@ -0,0 +1,5 @@ +from .core import PtychoFMReconstructorLibrary + +__all__ = [ + 'PtychoFMReconstructorLibrary', +] diff --git a/src/ptychodus/model/ptycho_fm/_payload.py b/src/ptychodus/model/ptycho_fm/_payload.py new file mode 100644 index 000000000..f44ec30b6 --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/_payload.py @@ -0,0 +1,61 @@ +"""Payload dataclasses for the PtychoFM (ptycho-vit) subprocess entry points. + +``ptycho_vit`` is an optional extra, so nothing in this module imports it. The +config the child needs is carried as a plain nested ``dict[str, Any]`` -- the +same shape ptycho_vit's own ``config.yaml`` produces -- assembled parent-side +by the factory. Dicts are picklable and pull in no framework, so the parent +never touches torch just to build a payload. + +The training mode (``'Unsupervised'`` / ``'Supervised'``) is carried in the +``name`` field, matching PtychoPINN's ``model_type`` convention. The child +warns if it disagrees with the mode baked into the training config. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from ptychodus.api.reconstructor import ReconstructInput + +__all__ = [ + 'ReconstructPayload', + 'TrainPayload', +] + + +@dataclass(frozen=True) +class ReconstructPayload: + """Everything the child needs to load a checkpoint and run inference once.""" + + # 'Unsupervised' or 'Supervised' — becomes the display name and is logged + # for the checkpoint-mode consistency warning. + name: str + + # ptycho_vit config as a nested dict (data/model/training/inference sections), + # matching config.yaml. Assembled parent-side from the settings groups. + config: dict[str, Any] + + # Path to the .pth weights file recorded by + # SubprocessReconstructor.load_model_from_file() parent-side. Required for + # inference; the child raises RuntimeError if None. + model_path: Path | None + + reconstruct_input: ReconstructInput + + +@dataclass(frozen=True) +class TrainPayload: + """Everything the child needs to run one training session. + + Ptycho-vit training reads its dataset from a directory of ``.hdf5`` files, + so ``input_path`` should be a directory (not a single file). ``output_path`` + is where the child writes ``best.pth`` and any per-epoch snapshots. + """ + + name: str + config: dict[str, Any] + + input_path: Path + output_path: Path diff --git a/src/ptychodus/model/ptycho_fm/_subprocess.py b/src/ptychodus/model/ptycho_fm/_subprocess.py new file mode 100644 index 000000000..deb9fe01b --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/_subprocess.py @@ -0,0 +1,485 @@ +"""Child-side subprocess entry points for the PtychoFM (ptycho-vit) backend. + +Runs INSIDE a spawned subprocess. This is the only ptychodus module allowed to +import ``torch`` or ``ptycho_vit``. The parent never imports it -- it reaches +only the payload dataclasses in :mod:`._payload` and the settings-to-dict +translator in :mod:`.reconstructor`, neither of which touches torch. + +Two entry points are exposed: + +- :func:`run_reconstruct` -- load a ``.pth`` checkpoint, run one inference pass + over the diffraction stack (batched through :class:`PtychoViT`), stitch the + per-patch amplitude/phase outputs into a full-object array via + ``place_patches_fourier_shift``, and stream back a single + :class:`ReconstructOutput`. +- :func:`run_train` -- run one single-device training session (no DDP, no + mlflow, no wandb) driven by :class:`ptycho_vit.model.model.PtychoViT` and a + minimal train/validate loop that mirrors ``ptycho_vit/train.py`` stripped of + its distributed machinery. Emits a :class:`TrainOutput` after each epoch and + a final ``TAG_MODEL_SAVED`` with the path of the best checkpoint. +""" + +from __future__ import annotations + +import logging +import os +import pickle +from collections.abc import Sequence +from multiprocessing.queues import Queue +from pathlib import Path +from typing import Any + +import numpy + +from ptychodus.api.diffraction import zero_bad_pixels +from ptychodus.api.object import Object +from ptychodus.api.product import LossValue, Product +from ptychodus.api.reconstructor import ReconstructOutput, TrainOutput + +from ..processing.subprocess_reconstructor import ( + TAG_MODEL_SAVED, + TAG_OUTPUT, + TAG_TRAIN_OUTPUT, +) +from ._payload import ReconstructPayload, TrainPayload + +logger = logging.getLogger(__name__) + + +def _select_device() -> Any: + """Return the best available torch device. + + Kept in its own helper so both entry points share the selection rule and + the import of torch is confined to the child. + """ + import torch + + if torch.cuda.is_available(): + return torch.device('cuda') + return torch.device('cpu') + + +def _pad_probe_to_modes(probe: numpy.ndarray, target_modes: int) -> numpy.ndarray: + """Zero-pad a complex probe ``(N_modes, H, W)`` up to ``target_modes`` along axis 0. + + Mirrors :meth:`ptycho_vit.data.PtychographyDataset._pad_probe` but operates + on the ``(N, H, W)`` layout ptychodus hands us (rather than the + ``(1, N, H, W)`` layout the dataset uses internally). + """ + current_modes = probe.shape[0] + if current_modes >= target_modes: + return probe + padding = numpy.zeros( + (target_modes - current_modes, probe.shape[1], probe.shape[2]), + dtype=probe.dtype, + ) + return numpy.concatenate([probe, padding], axis=0) + + +def _zero_pad_2d_to(image: numpy.ndarray, target_size: int) -> numpy.ndarray: + """Center a 2D array inside a ``(target_size, target_size)`` zero-padded canvas.""" + h, w = image.shape + if h == target_size and w == target_size: + return image + if h > target_size or w > target_size: + raise ValueError( + f'Image size ({h}, {w}) exceeds target size {target_size}; refusing to crop.' + ) + pad_h = target_size - h + pad_w = target_size - w + pad_top = pad_h // 2 + pad_left = pad_w // 2 + return numpy.pad( + image, + ((pad_top, pad_h - pad_top), (pad_left, pad_w - pad_left)), + mode='constant', + constant_values=0, + ) + + +def _build_positions_top_left(parameters: Any) -> numpy.ndarray: + """Convert ptychodus probe positions to top-left-origin object-pixel coords. + + Returns an ``(N, 2)`` float32 array of ``[y_px, x_px]`` centres, matching + what ``place_patches_fourier_shift`` expects. + """ + object_geometry = parameters.product.object_.get_geometry() + coords: list[float] = [] + for position in parameters.product.probe_positions: + object_point = object_geometry.map_coordinates_probe_to_object(position) + coords.append(object_point.coordinate_y_px) + coords.append(object_point.coordinate_x_px) + return numpy.asarray(coords, dtype=numpy.float32).reshape(-1, 2) + + +def run_reconstruct(payload: ReconstructPayload, queue: Queue[Any]) -> None: + """Child entry point for one inference pass. Streams a single ReconstructOutput. + + Loads the ``.pth`` state dict with ``weights_only=True`` (safe: nothing to + execute is expected in a plain ptycho_vit checkpoint), rebuilds the model + from the payload's config, then runs the same batch + stitch loop as + ``ptycho_vit/scripts/run_inference_and_stitch.py``. + """ + if payload.model_path is None: + raise RuntimeError('Cannot reconstruct: no model checkpoint has been loaded.') + + import torch + from ptycho_vit.model.model import PtychoViT + from ptycho_vit.utils.ptychi_utils import place_patches_fourier_shift + + device = _select_device() + config = payload.config + data_config = config['data'] + model_config = config['model'] + inference_config = config['inference'] + + model = PtychoViT(config=model_config) + state = torch.load(payload.model_path, map_location=device, weights_only=True) + model.load_state_dict(state) + model.to(device) + model.eval() + + parameters = payload.reconstruct_input + + diff_intensity = zero_bad_pixels(parameters.diffraction_patterns, parameters.bad_pixels) + diff_intensity = numpy.asarray(diff_intensity, dtype=numpy.float32) + if diff_intensity.ndim != 3: + raise ValueError( + f'Expected diffraction patterns with shape (N, H, W); got {diff_intensity.shape}.' + ) + + # Match PtychographyDataset preprocessing: normalise by dataset max, scale, + # then take sqrt (the model's ``x`` input is amplitude-domain). + normalization_value = float(diff_intensity.max()) + scale_value = float(data_config['scale']) + if normalization_value <= 0.0: + raise ValueError('Diffraction stack max is non-positive; cannot normalise for inference.') + diff_amp = numpy.sqrt(diff_intensity / normalization_value * scale_value) + + target_size = int(data_config['target_size']) + if diff_amp.shape[-1] != target_size or diff_amp.shape[-2] != target_size: + diff_amp = numpy.stack([_zero_pad_2d_to(p, target_size) for p in diff_amp], axis=0) + + probe_array = parameters.product.probes.get_probe_no_opr().get_array() + if probe_array.ndim != 3: + raise ValueError(f'Expected probe with shape (N_modes, H, W); got {probe_array.shape}.') + max_modes = int(data_config['max_probe_modes']) + probe_padded = _pad_probe_to_modes(probe_array, max_modes) + # ptycho_vit expects the probe input as a real view with the mode-count + # index in the third position: (B, 1, N_modes, H, W, 2). We build a single + # (1, 1, N_modes, H, W) complex tensor and broadcast per batch below. + probe_complex = torch.from_numpy(numpy.ascontiguousarray(probe_padded)).to( + dtype=torch.complex64, device=device + ) + probe_real_view = torch.view_as_real(probe_complex).unsqueeze(0).unsqueeze(0) + + positions_np = _build_positions_top_left(parameters) + positions = torch.from_numpy(positions_np) + + object_in = parameters.product.object_ + object_array = object_in.get_array() + object_shape = (object_array.shape[-2], object_array.shape[-1]) + + pred_amp_object = torch.zeros(object_shape, dtype=torch.float32) + pred_ph_object = torch.zeros(object_shape, dtype=torch.float32) + buffer = torch.zeros(object_shape, dtype=torch.float32) + + central_crop = int(inference_config['central_crop']) + pad = int(inference_config['pad']) + batch_size = max(1, int(inference_config['batch_size'])) + + diff_amp_tensor = torch.from_numpy(diff_amp).unsqueeze(1) # (N, 1, H, W) + n_patterns = diff_amp_tensor.shape[0] + + with torch.no_grad(): + for start in range(0, n_patterns, batch_size): + end = min(start + batch_size, n_patterns) + actual_bs = end - start + + input_diff = diff_amp_tensor[start:end].to(device) + input_probe = probe_real_view.expand(actual_bs, -1, -1, -1, -1, -1) + input_norm = torch.full( + (actual_bs, 1), normalization_value, dtype=torch.float32, device=device + ) + input_scale = torch.full( + (actual_bs, 1), scale_value, dtype=torch.float32, device=device + ) + + _pred_diff, output_amp, output_ph = model( + input_diff, input_probe, input_norm, input_scale + ) + + output_amp = output_amp.squeeze(1).detach().cpu() + output_ph = output_ph.squeeze(1).detach().cpu() + + amp_patches = output_amp[:, central_crop:-central_crop, central_crop:-central_crop] + ph_patches = output_ph[:, central_crop:-central_crop, central_crop:-central_crop] + + batch_positions = positions[start:end] + + pred_amp_object = place_patches_fourier_shift( + pred_amp_object, + batch_positions, + amp_patches, + op='add', + adjoint_mode=False, + pad=pad, + ) + pred_ph_object = place_patches_fourier_shift( + pred_ph_object, + batch_positions, + ph_patches, + op='add', + adjoint_mode=False, + pad=pad, + ) + buffer = place_patches_fourier_shift( + buffer, + batch_positions, + torch.ones_like(ph_patches), + op='add', + adjoint_mode=False, + pad=pad, + ) + + divisor = torch.clip(buffer, min=1.0) + pred_amp_object = pred_amp_object / divisor + pred_ph_object = pred_ph_object / divisor + + # Fold amplitude + phase back into a complex layer. Preserve the input + # object's outer shape (layers, H, W) by using layer 0 only. + complex_layer = (pred_amp_object.numpy() * numpy.exp(1j * pred_ph_object.numpy())).astype( + numpy.complex64 + ) + if object_array.ndim == 3: + object_out_array = numpy.zeros_like(object_array) + object_out_array[0] = complex_layer + else: + object_out_array = complex_layer + + object_out = Object( + array=object_out_array, + layer_spacing_m=object_in.layer_spacing_m, + pixel_geometry=object_in.get_pixel_geometry(), + center=object_in.get_center(), + ) + + losses: Sequence[LossValue] = [] + product = Product( + metadata=parameters.product.metadata, + probe_positions=parameters.product.probe_positions, + probes=parameters.product.probes, + object_=object_out, + losses=losses, + ) + + queue.put((TAG_OUTPUT, pickle.dumps(ReconstructOutput(product=product, progress=1)))) + + +def _build_criterion(training_config: dict[str, Any]) -> Any: + """Instantiate the loss module named in ``training_config['loss_function']``. + + Mirrors the branch at ``ptycho_vit/train.py`` line ~728. + """ + import torch.nn as nn + + from ptycho_vit.custom_loss import WeightedLoss + + name = training_config['loss_function'] + if name == 'smooth_l1': + return nn.SmoothL1Loss() + if name == 'mse': + return nn.MSELoss() + if name == 'l1': + return nn.L1Loss() + if name == 'poisson_nll': + return nn.PoissonNLLLoss(log_input=False, full=False) + if name == 'weighted': + w = training_config['weighted_loss'] + return WeightedLoss(loss_type=w['loss_type'], threshold=w['threshold'], alpha=w['alpha']) + raise ValueError(f'Unknown loss function: {name!r}') + + +def run_train(payload: TrainPayload, queue: Queue[Any]) -> None: + """Child entry point for one training session. + + Single-device, no DDP, no mlflow, no wandb. Builds a CombinedDataset over + ``payload.input_path`` (a directory of paired ``*_dp.hdf5`` / ``*_para.hdf5`` + files -- ptycho_vit's own training format) and runs a lightweight + train/validate loop directly against :class:`PtychoViT`, emitting one + :class:`TrainOutput` per epoch. The best checkpoint by validation loss is + written to ``payload.output_path/best.pth`` and its path streamed back via + ``TAG_MODEL_SAVED``. + """ + # Guard the training loop against picking up an in-flight environment: a + # user may have limited devices via CUDA_VISIBLE_DEVICES already; we + # respect that and never override it here. + _visible = os.environ.get('CUDA_VISIBLE_DEVICES') + if _visible is not None: + logger.info('Training with CUDA_VISIBLE_DEVICES=%s', _visible) + + import torch + import torch.optim as optim + from torch.utils.data import DataLoader, Subset, random_split + + from ptycho_vit.data import CombinedDataset + from ptycho_vit.model.model import PtychoViT + + device = _select_device() + config = payload.config + data_config = config['data'] + training_config = config['training'] + model_config = config['model'] + + if not payload.input_path.exists(): + raise FileNotFoundError(f'Training input directory does not exist: {payload.input_path}') + if not payload.input_path.is_dir(): + raise NotADirectoryError( + 'ptycho_vit training expects a directory of paired *_dp.hdf5 / ' + f'*_para.hdf5 files; got file: {payload.input_path}' + ) + + payload.output_path.mkdir(parents=True, exist_ok=True) + + dataset = CombinedDataset( + file_paths=str(payload.input_path), + rank=0, + world_size=1, + scale=data_config['scale'], + normalization_dict_path=None, + default_normalization=data_config['default_normalization'], + apply_noise=False, + cache_object=data_config['cache_object'], + max_probe_modes=data_config['max_probe_modes'], + target_size=data_config['target_size'], + max_files=data_config['max_files'], + debug=False, + ) + total_size = len(dataset) + if total_size < 2: + raise ValueError( + f'Training dataset has {total_size} sample(s); need at least 2 for a train/val split.' + ) + + train_split = float(data_config['train_split']) + train_size = max(1, int(total_size * train_split)) + val_size = max(1, total_size - train_size) + if train_size + val_size > total_size: + train_size = total_size - val_size + generator = torch.Generator().manual_seed(int(data_config['random_seed'])) + train_subset: Subset[Any] + val_subset: Subset[Any] + train_subset, val_subset = random_split(dataset, [train_size, val_size], generator=generator) + + num_workers = int(data_config['num_workers']) + batch_size = int(training_config['batch_size']) + loader_kwargs: dict[str, Any] = { + 'batch_size': batch_size, + 'num_workers': num_workers, + 'pin_memory': False, + } + if num_workers > 0: + loader_kwargs['prefetch_factor'] = int(data_config['prefetch_factor']) + train_loader = DataLoader(train_subset, shuffle=True, drop_last=False, **loader_kwargs) + val_loader = DataLoader(val_subset, shuffle=False, drop_last=False, **loader_kwargs) + + model = PtychoViT(config=model_config).to(device) + + lr = float(training_config['learning_rate']) + param_groups = [ + {'params': model.encoder.parameters(), 'lr': lr, 'name': 'encoder'}, + {'params': model.amp_decoder.parameters(), 'lr': lr, 'name': 'amp_decoder'}, + {'params': model.ph_decoder.parameters(), 'lr': lr, 'name': 'ph_decoder'}, + ] + optimizer = optim.Adam(param_groups) + + criterion = _build_criterion(training_config) + + def _forward(batch: tuple[Any, ...]) -> Any: + diff_amp, amp_patch, ph_patch, probe, _probe_pos, norm, scale = batch + input_diff = diff_amp.to(device) + input_probe = torch.view_as_real(probe.clone().detach()).to(device) + input_norm = norm.to(device) + input_scale = scale.to(device) + pred_diff, pred_amp, pred_ph = model(input_diff, input_probe, input_norm, input_scale) + target_amp = amp_patch.to(device) + target_ph = ph_patch.to(device) + # Compose an amplitude+phase loss. Matches ptycho_vit's default target + # (the model's amp/ph decoders drive the loss, not the reconstructed + # diffraction), stripped of the wandb-driven auxiliary terms. + amp_loss = criterion(pred_amp, target_amp) + ph_loss = criterion(pred_ph, target_ph) + return amp_loss + ph_loss + + training_losses: list[LossValue] = [] + validation_losses: list[LossValue] = [] + best_val = float('inf') + best_path = payload.output_path / 'best.pth' + save_epoch_models = bool(training_config['save_epoch_models']) + epochs = int(training_config['epochs']) + + for epoch in range(epochs): + model.train() + running_train = 0.0 + n_train_batches = 0 + for batch in train_loader: + optimizer.zero_grad(set_to_none=True) + loss = _forward(batch) + loss.backward() + optimizer.step() + running_train += float(loss.detach().cpu().item()) + n_train_batches += 1 + train_loss = running_train / max(n_train_batches, 1) + training_losses.append(LossValue(epoch=epoch, value=train_loss)) + + model.eval() + running_val = 0.0 + n_val_batches = 0 + with torch.no_grad(): + for batch in val_loader: + loss = _forward(batch) + running_val += float(loss.detach().cpu().item()) + n_val_batches += 1 + val_loss = running_val / max(n_val_batches, 1) + validation_losses.append(LossValue(epoch=epoch, value=val_loss)) + + if val_loss < best_val: + best_val = val_loss + _atomic_save_state_dict(model.state_dict(), best_path) + if save_epoch_models: + epoch_path = payload.output_path / f'model_epoch_{epoch + 1:03d}.pth' + _atomic_save_state_dict(model.state_dict(), epoch_path) + + queue.put( + ( + TAG_TRAIN_OUTPUT, + pickle.dumps( + TrainOutput( + training_loss=list(training_losses), + validation_loss=list(validation_losses), + progress=epoch + 1, + ) + ), + ) + ) + + if not best_path.exists(): + raise FileNotFoundError( + f'Training finished but no best checkpoint was written at {best_path}.' + ) + queue.put((TAG_MODEL_SAVED, str(best_path))) + + +def _atomic_save_state_dict(state_dict: Any, destination: Path) -> None: + """Save a torch state dict atomically: write to a temp path, then rename. + + ``model.state_dict()`` may reference CUDA tensors; we move them to CPU so + reloading does not require the same device layout. + """ + import torch + + cpu_state = {k: v.detach().cpu() if hasattr(v, 'detach') else v for k, v in state_dict.items()} + tmp = destination.with_suffix(destination.suffix + '.tmp') + tmp.parent.mkdir(parents=True, exist_ok=True) + torch.save(cpu_state, tmp) + os.replace(tmp, destination) diff --git a/src/ptychodus/model/ptycho_fm/core.py b/src/ptychodus/model/ptycho_fm/core.py new file mode 100644 index 000000000..b927948b0 --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/core.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import logging +from collections.abc import Iterator +from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec + +from ...api.reconstructor import ( + NullReconstructor, + Reconstructor, + ReconstructorLibrary, + TrainableReconstructor, +) +from ...api.settings import SettingsRegistry +from .enums import PtychoFMEnumerators +from .settings import ( + PtychoFMDataSettings, + PtychoFMInferenceSettings, + PtychoFMModelSettings, + PtychoFMTrainingSettings, +) + +logger = logging.getLogger(__name__) + + +def _ptycho_fm_available() -> bool: + """Return True iff ptycho_vit and torch are importable, without importing them.""" + return all(find_spec(mod) is not None for mod in ('ptycho_vit', 'torch')) + + +class PtychoFMReconstructorLibrary(ReconstructorLibrary): + def __init__( + self, settings_registry: SettingsRegistry, is_developer_mode_enabled: bool + ) -> None: + super().__init__('ptycho-fm') + self.data_settings = PtychoFMDataSettings(settings_registry) + self.model_settings = PtychoFMModelSettings(settings_registry) + self.training_settings = PtychoFMTrainingSettings(settings_registry) + self.inference_settings = PtychoFMInferenceSettings(settings_registry) + self.enumerators = PtychoFMEnumerators() + self._reconstructors: list[TrainableReconstructor] = list() + + if not _ptycho_fm_available(): + logger.info('PtychoFM (ptycho-vit) not found.') + + if is_developer_mode_enabled: + for reconstructor in ('Unsupervised', 'Supervised'): + self._reconstructors.append(NullReconstructor(reconstructor)) + return + + try: + ptycho_fm_version = version('ptycho-vit') + except PackageNotFoundError: + ptycho_fm_version = 'unknown' + logger.info(f'PtychoFM (ptycho-vit) {ptycho_fm_version}') + + # Lazy import: keeps this module's cost small in headless mode and + # ensures the parent never pulls torch in just because ptycho-vit is + # installed. + from .reconstructor import build_reconstructor + + for mode in ('Unsupervised', 'Supervised'): + self._reconstructors.append( + build_reconstructor( + mode, + self.data_settings, + self.model_settings, + self.inference_settings, + self.training_settings, + ) + ) + + @property + def name(self) -> str: + return 'PtychoFM' + + def __iter__(self) -> Iterator[Reconstructor]: + return iter(self._reconstructors) diff --git a/src/ptychodus/model/ptycho_fm/enums.py b/src/ptychodus/model/ptycho_fm/enums.py new file mode 100644 index 000000000..32e97c3f8 --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/enums.py @@ -0,0 +1,35 @@ +from collections.abc import Iterator, Sequence + + +class PtychoFMEnumerators: + def __init__(self) -> None: + self._sharding_strategies: Sequence[str] = ['static', 'dynamic'] + self._encoder_types: Sequence[str] = ['custom', 'pretrained'] + self._init_methods: Sequence[str] = ['trunc_normal', 'kaiming'] + self._loss_functions: Sequence[str] = [ + 'smooth_l1', + 'mse', + 'l1', + 'poisson_nll', + 'weighted', + ] + self._weighted_loss_types: Sequence[str] = ['mse', 'mae'] + self._training_modes: Sequence[str] = ['unsupervised', 'supervised'] + + def get_sharding_strategies(self) -> Iterator[str]: + return iter(self._sharding_strategies) + + def get_encoder_types(self) -> Iterator[str]: + return iter(self._encoder_types) + + def get_init_methods(self) -> Iterator[str]: + return iter(self._init_methods) + + def get_loss_functions(self) -> Iterator[str]: + return iter(self._loss_functions) + + def get_weighted_loss_types(self) -> Iterator[str]: + return iter(self._weighted_loss_types) + + def get_training_modes(self) -> Iterator[str]: + return iter(self._training_modes) diff --git a/src/ptychodus/model/ptycho_fm/reconstructor.py b/src/ptychodus/model/ptycho_fm/reconstructor.py new file mode 100644 index 000000000..0f2ae634c --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/reconstructor.py @@ -0,0 +1,222 @@ +"""Parent-side factory that builds a :class:`SubprocessReconstructor` for PtychoFM. + +Zero torch imports. All GPU work runs inside a spawned child; see +:mod:`._subprocess` for the child entry points. + +PtychoFM's own ``config.yaml`` is a nested dict of scalars, so ``_build_config`` +produces a plain :class:`dict` from the ptychodus settings groups -- pickleable +without pulling any ptycho_vit module in. The child feeds that dict directly +into ``PtychoViT(config)`` and the training loop. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import h5py +import numpy + +from ptychodus.api.diffraction import zero_bad_pixels +from ptychodus.api.reconstructor import ReconstructInput + +from ..processing.subprocess_reconstructor import SubprocessReconstructor +from ._payload import ReconstructPayload, TrainPayload +from .settings import ( + PtychoFMDataSettings, + PtychoFMInferenceSettings, + PtychoFMModelSettings, + PtychoFMTrainingSettings, +) + +__all__ = [ + 'build_reconstructor', +] + + +_RECONSTRUCT_ENTRY = 'ptychodus.model.ptycho_fm._subprocess:run_reconstruct' +_TRAIN_ENTRY = 'ptychodus.model.ptycho_fm._subprocess:run_train' + + +def _build_config( + data_settings: PtychoFMDataSettings, + model_settings: PtychoFMModelSettings, + training_settings: PtychoFMTrainingSettings, + inference_settings: PtychoFMInferenceSettings, +) -> dict[str, Any]: + """Translate ptychodus settings into a nested ``ptycho_vit`` config dict. + + Mirrors the top-level shape of ``ptycho_vit/config.yaml``: ``data``, + ``model`` (with ``encoder`` / ``decoder`` / ``init`` sub-sections), + ``training`` (with ``weighted_loss`` sub-section), and ``inference``. The + child fills in any missing paths (``data_path`` / model save path) from + the training payload at call time. + """ + max_files = data_settings.max_files.get_value() + data_config: dict[str, Any] = { + 'scale': data_settings.scale.get_value(), + 'default_normalization': data_settings.default_normalization.get_value(), + 'packed': data_settings.packed.get_value(), + 'cache_object': data_settings.cache_object.get_value(), + 'max_probe_modes': data_settings.max_probe_modes.get_value(), + 'target_size': data_settings.target_size.get_value(), + 'train_split': data_settings.train_split.get_value(), + 'random_seed': data_settings.random_seed.get_value(), + 'sharding_strategy': data_settings.sharding_strategy.get_value(), + # 0 in settings means "no cap" (null in the YAML). + 'max_files': max_files if max_files > 0 else None, + 'num_workers': data_settings.num_workers.get_value(), + 'prefetch_factor': data_settings.prefetch_factor.get_value(), + 'use_cuda_prefetcher': data_settings.use_cuda_prefetcher.get_value(), + } + + model_config: dict[str, Any] = { + 'encoder_type': model_settings.encoder_type.get_value(), + 'encoder': { + 'img_size': model_settings.img_size.get_value(), + 'patch_size': model_settings.patch_size.get_value(), + 'in_channels': 1, + 'embed_dim': model_settings.embed_dim.get_value(), + 'depth': model_settings.depth.get_value(), + 'num_heads': model_settings.num_heads.get_value(), + 'mlp_ratio': model_settings.mlp_ratio.get_value(), + 'use_cls_token': model_settings.use_cls_token.get_value(), + 'dropout': model_settings.dropout.get_value(), + 'attn_dropout': model_settings.attn_dropout.get_value(), + 'timm_model_name': model_settings.timm_model_name.get_value(), + }, + 'decoder': { + 'base_channels': model_settings.decoder_base_channels.get_value(), + 'latent_dim': model_settings.decoder_latent_dim.get_value(), + 'num_stages': model_settings.decoder_num_stages.get_value(), + 'use_batchnorm': model_settings.decoder_use_batchnorm.get_value(), + 'dropout': model_settings.decoder_dropout.get_value(), + }, + 'init': { + 'enabled': model_settings.init_enabled.get_value(), + 'method': model_settings.init_method.get_value(), + }, + } + + training_config: dict[str, Any] = { + 'mode': training_settings.mode.get_value(), + 'batch_size': training_settings.batch_size.get_value(), + 'learning_rate': training_settings.learning_rate.get_value(), + 'epochs': training_settings.epochs.get_value(), + 'loss_function': training_settings.loss_function.get_value(), + 'weighted_loss': { + 'loss_type': training_settings.weighted_loss_type.get_value(), + 'threshold': training_settings.weighted_loss_threshold.get_value(), + 'alpha': training_settings.weighted_loss_alpha.get_value(), + }, + 'validation_plot_freq': training_settings.validation_plot_freq.get_value(), + 'checkpoint_freq': training_settings.checkpoint_freq.get_value(), + 'save_epoch_models': training_settings.save_epoch_models.get_value(), + 'resume_from_checkpoint': training_settings.resume_from_checkpoint.get_value(), + } + + inference_config: dict[str, Any] = { + 'central_crop': inference_settings.central_crop.get_value(), + 'pad': inference_settings.pad.get_value(), + 'batch_size': inference_settings.batch_size.get_value(), + } + + return { + 'data': data_config, + 'model': model_config, + 'training': training_config, + 'inference': inference_config, + } + + +def build_reconstructor( + name: str, + data_settings: PtychoFMDataSettings, + model_settings: PtychoFMModelSettings, + inference_settings: PtychoFMInferenceSettings, + training_settings: PtychoFMTrainingSettings, +) -> SubprocessReconstructor: + """Build a :class:`SubprocessReconstructor` for one PtychoFM mode. + + ``name`` is 'Unsupervised' or 'Supervised' and becomes the reconstructor's + display name. + """ + + def build_reconstruct_payload( + parameters: ReconstructInput, loaded_model_path: Path | None + ) -> ReconstructPayload: + return ReconstructPayload( + name=name, + config=_build_config( + data_settings, model_settings, training_settings, inference_settings + ), + model_path=loaded_model_path, + reconstruct_input=parameters, + ) + + def build_train_payload(input_path: Path, output_path: Path) -> TrainPayload: + return TrainPayload( + name=name, + config=_build_config( + data_settings, model_settings, training_settings, inference_settings + ), + input_path=input_path, + output_path=output_path, + ) + + def export_training_data(file_path: Path, parameters: ReconstructInput) -> None: + # ptycho_vit.CombinedDataset expects one directory per scan holding + # two HDF5 files sharing a common stem: + # <stem>_dp.hdf5 dataset 'dp' (N, H, W) float + # <stem>_para.hdf5 dataset 'object' (1, H, W) complex + # 'probe' (1, N_modes, H, W) complex + # 'probe_position_x_m' (N,) float64 + # 'probe_position_y_m' (N,) float64 + # 'object' carries a 'pixel_height_m' attr; the + # loader auto-detects meters vs pixels from range. + # file_path is treated as a stem, so the picked filename itself is + # never created; the two derived files land next to it. + stem = file_path.stem + dp_path = file_path.with_name(f'{stem}_dp.hdf5') + para_path = file_path.with_name(f'{stem}_para.hdf5') + + dp = zero_bad_pixels(parameters.diffraction_patterns, parameters.bad_pixels) + obj = parameters.product.object_ + object_layer = obj.get_layer(0) + pixel_geometry = obj.get_pixel_geometry() + probe_array = parameters.product.probes.get_probe_no_opr().get_array() + + pos_x_m: list[float] = [] + pos_y_m: list[float] = [] + for point in parameters.product.probe_positions: + pos_x_m.append(point.coordinate_x_m) + pos_y_m.append(point.coordinate_y_m) + + with h5py.File(dp_path, 'w') as h5_dp: + h5_dp.create_dataset('dp', data=dp) + + with h5py.File(para_path, 'w') as h5_para: + obj_ds = h5_para.create_dataset('object', data=object_layer[numpy.newaxis, :, :]) + obj_ds.attrs['pixel_height_m'] = pixel_geometry.height_m + obj_ds.attrs['pixel_width_m'] = pixel_geometry.width_m + h5_para.create_dataset('probe', data=probe_array[numpy.newaxis, :, :, :]) + h5_para.create_dataset( + 'probe_position_x_m', data=numpy.asarray(pos_x_m, dtype=numpy.float64) + ) + h5_para.create_dataset( + 'probe_position_y_m', data=numpy.asarray(pos_y_m, dtype=numpy.float64) + ) + + return SubprocessReconstructor( + name=name, + reconstruct_entry_point=_RECONSTRUCT_ENTRY, + progress_goal_fn=lambda: training_settings.epochs.get_value(), + build_reconstruct_payload=build_reconstruct_payload, + is_trainable=True, + train_entry_point=_TRAIN_ENTRY, + build_train_payload=build_train_payload, + model_file_filter='PyTorch Checkpoint (*.pth *.pt)', + model_file_extension='.pth', + training_data_file_filter='PtychoFM Training Pair (*.hdf5)', + export_training_data=export_training_data, + ) diff --git a/src/ptychodus/model/ptycho_fm/settings.py b/src/ptychodus/model/ptycho_fm/settings.py new file mode 100644 index 000000000..0d0f1b2a3 --- /dev/null +++ b/src/ptychodus/model/ptycho_fm/settings.py @@ -0,0 +1,128 @@ +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.settings import SettingsRegistry + + +class PtychoFMDataSettings(Observable, Observer): + def __init__(self, registry: SettingsRegistry) -> None: + super().__init__() + self._group = registry.create_group('PtychoFMData') + self._group.add_observer(self) + + self.scale = self._group.create_real_parameter('scale', 10000.0, minimum=0.0) + self.default_normalization = self._group.create_real_parameter( + 'default_normalization', 100000.0, minimum=0.0 + ) + self.packed = self._group.create_boolean_parameter('packed', True) + self.cache_object = self._group.create_boolean_parameter('cache_object', True) + self.max_probe_modes = self._group.create_integer_parameter( + 'max_probe_modes', 10, minimum=1 + ) + self.target_size = self._group.create_integer_parameter('target_size', 256, minimum=32) + self.train_split = self._group.create_real_parameter( + 'train_split', 0.80, minimum=0.0, maximum=1.0 + ) + self.random_seed = self._group.create_integer_parameter('random_seed', 8) + self.sharding_strategy = self._group.create_string_parameter('sharding_strategy', 'dynamic') + # 0 → treat as "use all files" (null in the YAML). + self.max_files = self._group.create_integer_parameter('max_files', 0, minimum=0) + self.num_workers = self._group.create_integer_parameter('num_workers', 4, minimum=0) + self.prefetch_factor = self._group.create_integer_parameter('prefetch_factor', 2, minimum=0) + self.use_cuda_prefetcher = self._group.create_boolean_parameter('use_cuda_prefetcher', True) + + def _update(self, observable: Observable) -> None: + if observable is self._group: + self.notify_observers() + + +class PtychoFMModelSettings(Observable, Observer): + def __init__(self, registry: SettingsRegistry) -> None: + super().__init__() + self._group = registry.create_group('PtychoFMModel') + self._group.add_observer(self) + + self.encoder_type = self._group.create_string_parameter('encoder_type', 'custom') + self.img_size = self._group.create_integer_parameter('img_size', 256, minimum=32) + self.patch_size = self._group.create_integer_parameter('patch_size', 16, minimum=1) + self.embed_dim = self._group.create_integer_parameter('embed_dim', 512, minimum=1) + self.depth = self._group.create_integer_parameter('depth', 12, minimum=1) + self.num_heads = self._group.create_integer_parameter('num_heads', 8, minimum=1) + self.mlp_ratio = self._group.create_real_parameter('mlp_ratio', 4.0, minimum=0.0) + self.use_cls_token = self._group.create_boolean_parameter('use_cls_token', False) + self.dropout = self._group.create_real_parameter('dropout', 0.1, minimum=0.0, maximum=1.0) + self.attn_dropout = self._group.create_real_parameter( + 'attn_dropout', 0.0, minimum=0.0, maximum=1.0 + ) + self.timm_model_name = self._group.create_string_parameter( + 'timm_model_name', 'vit_large_patch32_224' + ) + + self.decoder_base_channels = self._group.create_integer_parameter( + 'decoder_base_channels', 64, minimum=1 + ) + self.decoder_latent_dim = self._group.create_integer_parameter( + 'decoder_latent_dim', 512, minimum=1 + ) + self.decoder_num_stages = self._group.create_integer_parameter( + 'decoder_num_stages', 4, minimum=1 + ) + self.decoder_use_batchnorm = self._group.create_boolean_parameter( + 'decoder_use_batchnorm', True + ) + self.decoder_dropout = self._group.create_real_parameter( + 'decoder_dropout', 0.1, minimum=0.0, maximum=1.0 + ) + + self.init_enabled = self._group.create_boolean_parameter('init_enabled', False) + self.init_method = self._group.create_string_parameter('init_method', 'trunc_normal') + + def _update(self, observable: Observable) -> None: + if observable is self._group: + self.notify_observers() + + +class PtychoFMTrainingSettings(Observable, Observer): + def __init__(self, registry: SettingsRegistry) -> None: + super().__init__() + self._group = registry.create_group('PtychoFMTraining') + self._group.add_observer(self) + + # ``mode`` is kept independent from the per-reconstructor mode split so + # that a checkpoint trained supervised can be evaluated as unsupervised + # (and vice-versa) without touching the file. + self.mode = self._group.create_string_parameter('mode', 'unsupervised') + self.batch_size = self._group.create_integer_parameter('batch_size', 64, minimum=1) + self.learning_rate = self._group.create_real_parameter('learning_rate', 1.0e-5, minimum=0.0) + self.epochs = self._group.create_integer_parameter('epochs', 11, minimum=1) + self.loss_function = self._group.create_string_parameter('loss_function', 'weighted') + self.weighted_loss_type = self._group.create_string_parameter('weighted_loss_type', 'mse') + self.weighted_loss_threshold = self._group.create_real_parameter( + 'weighted_loss_threshold', 0.0 + ) + self.weighted_loss_alpha = self._group.create_real_parameter('weighted_loss_alpha', 1.0) + self.validation_plot_freq = self._group.create_integer_parameter( + 'validation_plot_freq', 1, minimum=1 + ) + self.checkpoint_freq = self._group.create_integer_parameter('checkpoint_freq', 1, minimum=0) + self.save_epoch_models = self._group.create_boolean_parameter('save_epoch_models', True) + self.resume_from_checkpoint = self._group.create_boolean_parameter( + 'resume_from_checkpoint', False + ) + + def _update(self, observable: Observable) -> None: + if observable is self._group: + self.notify_observers() + + +class PtychoFMInferenceSettings(Observable, Observer): + def __init__(self, registry: SettingsRegistry) -> None: + super().__init__() + self._group = registry.create_group('PtychoFMInference') + self._group.add_observer(self) + + self.central_crop = self._group.create_integer_parameter('central_crop', 64, minimum=1) + self.pad = self._group.create_integer_parameter('pad', 32, minimum=0) + self.batch_size = self._group.create_integer_parameter('batch_size', 256, minimum=1) + + def _update(self, observable: Observable) -> None: + if observable is self._group: + self.notify_observers() diff --git a/src/ptychodus/model/ptychonn/_payload.py b/src/ptychodus/model/ptychonn/_payload.py new file mode 100644 index 000000000..7269d34a1 --- /dev/null +++ b/src/ptychodus/model/ptychonn/_payload.py @@ -0,0 +1,53 @@ +"""Payload dataclasses for the PtychoNN subprocess entry points. + +Parent-safe: no ptychonn / torch / lightning imports. + +The subprocess carries a small pydantic ``BaseModel`` (rather than a serialized +ptychodus ``SettingsRegistry``) so the child does not need to import +``PtychoNNModelSettings`` / ``PtychoNNTrainingSettings`` or rehydrate a +registry just to read a handful of scalars. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from pydantic import BaseModel, ConfigDict + +from ptychodus.api.reconstructor import ReconstructInput + + +class PtychoNNReconstructConfig(BaseModel): + """Minimal scalar config the child needs to build a ``LitReconSmallModel``.""" + + model_config = ConfigDict(frozen=True) + + enable_amplitude: bool + num_convolution_kernels: int + use_batch_normalization: bool + max_learning_rate: float + min_learning_rate: float + + +class PtychoNNTrainConfig(PtychoNNReconstructConfig): + """Adds the four train-only scalars ``ptychonn.train`` reads.""" + + batch_size: int + training_epochs: int + status_interval_in_epochs: int + validation_set_fractional_size: float + + +@dataclass(frozen=True) +class ReconstructPayload: + config: PtychoNNReconstructConfig + model_path: Path | None + reconstruct_input: ReconstructInput + + +@dataclass(frozen=True) +class TrainPayload: + config: PtychoNNTrainConfig + input_path: Path + output_path: Path diff --git a/src/ptychodus/model/ptychonn/_subprocess.py b/src/ptychodus/model/ptychonn/_subprocess.py new file mode 100644 index 000000000..03b00bbaf --- /dev/null +++ b/src/ptychodus/model/ptychonn/_subprocess.py @@ -0,0 +1,160 @@ +"""Child-side subprocess entry points for the PtychoNN backend. + +This module runs INSIDE a spawned subprocess. It is the only place in the +ptychodus tree that is allowed to import ptychonn / torch / lightning. + +The child receives a small pydantic config on the payload and reads the +scalars it needs directly — it does not touch ``SettingsRegistry`` or the +parent's settings-class layout. +""" + +from __future__ import annotations + +import logging +import pickle +from collections.abc import Sequence +from multiprocessing.queues import Queue +from typing import Any + +import numpy + +from ptychodus.api.interpolate import BarycentricArrayStitcher +from ptychodus.api.object import Object +from ptychodus.api.product import LossValue, Product +from ptychodus.api.reconstructor import ReconstructOutput, TrainOutput + +from ..processing.subprocess_reconstructor import ( + TAG_MODEL_SAVED, + TAG_OUTPUT, + TAG_TRAIN_OUTPUT, +) +from ._payload import ( + PtychoNNReconstructConfig, + ReconstructPayload, + TrainPayload, +) + +logger = logging.getLogger(__name__) + + +PATCHES_KEY = 'real' +PATTERNS_KEY = 'reciprocal' + + +def _build_model( + config: PtychoNNReconstructConfig, + *, + checkpoint_path: Any = None, +) -> Any: + import ptychonn + + if checkpoint_path is not None: + return ptychonn.LitReconSmallModel.load_from_checkpoint(checkpoint_path) + return ptychonn.LitReconSmallModel( + nconv=config.num_convolution_kernels, + use_batch_norm=config.use_batch_normalization, + enable_amplitude=config.enable_amplitude, + max_lr=config.max_learning_rate, + min_lr=config.min_learning_rate, + ) + + +def run_reconstruct(payload: ReconstructPayload, queue: Queue[Any]) -> None: + """Child entry point for one PtychoNN inference pass.""" + import ptychonn + + model = _build_model(payload.config, checkpoint_path=payload.model_path) + + parameters = payload.reconstruct_input + data = parameters.diffraction_patterns + data_size = data.shape[-1] + if data_size != data.shape[-2]: + raise ValueError('PtychoNN expects square diffraction data!') + is_data_size_pow2 = data_size & (data_size - 1) == 0 and data_size > 0 + if not is_data_size_pow2: + raise ValueError('PtychoNN expects that the diffraction data size is a power of two!') + + logger.debug('Inferring...') + object_patches = ptychonn.infer(data=data.astype(numpy.float32), model=model) + + logger.debug('Stitching...') + object_array = parameters.product.object_.get_array() + object_geometry = parameters.product.object_.get_geometry() + stitcher = BarycentricArrayStitcher( + upper=numpy.zeros_like(object_array), + lower=numpy.zeros_like(object_array, dtype=float), + ) + for scan_point, object_patch_channels in zip( + parameters.product.probe_positions, object_patches + ): + patch_array = numpy.exp(1j * object_patch_channels[0]) + if object_patch_channels.shape[0] == 2: + patch_array *= object_patch_channels[1] + else: + patch_array *= 0.5 + object_point = object_geometry.map_coordinates_probe_to_object(scan_point) + stitcher.add_patch(object_point.coordinate_x_px, object_point.coordinate_y_px, patch_array) + + object_ = Object( + array=stitcher.stitch(), + pixel_geometry=object_geometry.get_pixel_geometry(), + center=object_geometry.get_center(), + layer_spacing_m=parameters.product.object_.layer_spacing_m, + ) + losses: Sequence[LossValue] = list() + + product = Product( + metadata=parameters.product.metadata, + probe_positions=parameters.product.probe_positions, + probes=parameters.product.probes, + object_=object_, + losses=losses, + ) + queue.put((TAG_OUTPUT, pickle.dumps(ReconstructOutput(product)))) + + +def run_train(payload: TrainPayload, queue: Queue[Any]) -> None: + """Child entry point for one PtychoNN training session.""" + import ptychonn + + config = payload.config + + logger.debug(f'Reading "{payload.input_path}" as "NPZ"') + training_data = numpy.load(payload.input_path) + + model = _build_model(config) + training_set_fractional_size = 1 - config.validation_set_fractional_size + trainer, trainer_log = ptychonn.train( + model=model, + batch_size=config.batch_size, + out_dir=None, + X_train=training_data[PATTERNS_KEY], + Y_train=training_data[PATCHES_KEY], + epochs=config.training_epochs, + training_fraction=training_set_fractional_size, + log_frequency=config.status_interval_in_epochs, + strategy='ddp_notebook', + ) + + training_loss: list[LossValue] = [] + validation_loss: list[LossValue] = [] + for epoch, entry in enumerate(trainer_log.logs): + try: + tloss = LossValue(epoch, entry['training_loss']) + vloss = LossValue(epoch, entry['validation_loss']) + except KeyError: + pass + else: + training_loss.append(tloss) + validation_loss.append(vloss) + + checkpoint_path = payload.output_path + trainer.save_checkpoint(checkpoint_path) + + queue.put( + ( + TAG_TRAIN_OUTPUT, + pickle.dumps(TrainOutput(training_loss=training_loss, validation_loss=validation_loss)), + ) + ) + queue.put((TAG_MODEL_SAVED, str(checkpoint_path))) diff --git a/src/ptychodus/model/ptychonn/core.py b/src/ptychodus/model/ptychonn/core.py index 1b7cd874d..9104e3f43 100644 --- a/src/ptychodus/model/ptychonn/core.py +++ b/src/ptychodus/model/ptychonn/core.py @@ -1,5 +1,7 @@ from __future__ import annotations from collections.abc import Iterator +from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec import logging from ptychodus.api.reconstructor import ( @@ -15,6 +17,11 @@ logger = logging.getLogger(__name__) +def _ptychonn_available() -> bool: + """Return True iff ptychonn and lightning are importable, without importing them.""" + return all(find_spec(mod) is not None for mod in ('ptychonn', 'lightning')) + + class PtychoNNReconstructorLibrary(ReconstructorLibrary): def __init__( self, settings_registry: SettingsRegistry, is_developer_mode_enabled: bool @@ -24,33 +31,38 @@ def __init__( self.training_settings = PtychoNNTrainingSettings(settings_registry) self._reconstructors: list[TrainableReconstructor] = list() - try: - from .model import PtychoNNModelProvider - from .reconstructor import PtychoNNTrainableReconstructor - except ModuleNotFoundError: + if not _ptychonn_available(): logger.info('PtychoNN not found.') if is_developer_mode_enabled: self._reconstructors.append(NullReconstructor('PhaseOnly')) self._reconstructors.append(NullReconstructor('AmplitudePhase')) - else: - phase_only_model_provider = PtychoNNModelProvider( - self.model_settings, self.training_settings, enable_amplitude=False - ) - amplitude_phase_model_provider = PtychoNNModelProvider( - self.model_settings, self.training_settings, enable_amplitude=True - ) + return + + try: + ptychonn_version = version('ptychonn') + except PackageNotFoundError: + ptychonn_version = 'unknown' + logger.info(f'PtychoNN {ptychonn_version}') + + from .reconstructor import build_reconstructor - self._reconstructors.append( - PtychoNNTrainableReconstructor( - self.model_settings, self.training_settings, phase_only_model_provider - ) + self._reconstructors.append( + build_reconstructor( + 'PhaseOnly', + enable_amplitude=False, + model_settings=self.model_settings, + training_settings=self.training_settings, ) - self._reconstructors.append( - PtychoNNTrainableReconstructor( - self.model_settings, self.training_settings, amplitude_phase_model_provider - ) + ) + self._reconstructors.append( + build_reconstructor( + 'AmplitudePhase', + enable_amplitude=True, + model_settings=self.model_settings, + training_settings=self.training_settings, ) + ) @property def name(self) -> str: diff --git a/src/ptychodus/model/ptychonn/reconstructor.py b/src/ptychodus/model/ptychonn/reconstructor.py index 0942712ad..cf463421d 100644 --- a/src/ptychodus/model/ptychonn/reconstructor.py +++ b/src/ptychodus/model/ptychonn/reconstructor.py @@ -1,221 +1,161 @@ -from collections.abc import Iterator, Sequence -from importlib.metadata import version -from pathlib import Path -from typing import Final +"""Parent-side factory that builds :class:`SubprocessReconstructor`s for PtychoNN. + +Zero ptychonn / torch / lightning imports. All GPU work runs inside a +spawned child; see :mod:`._subprocess` for the child entry points. + +Training-data export runs parent-side because it is pure-numpy (barycentric +interpolation + numpy.savez); the ptychonn / lightning stack is not touched. +""" + +from __future__ import annotations + import logging +from pathlib import Path import numpy -import ptychonn -from ptychodus.api.common import ComplexArrayType from ptychodus.api.geometry import ImageExtent -from ptychodus.api.interpolate import BarycentricArrayInterpolator, BarycentricArrayStitcher -from ptychodus.api.object import Object -from ptychodus.api.product import Product -from ptychodus.api.reconstructor import ( - LossValue, - ReconstructInput, - ReconstructOutput, - TrainOutput, - TrainableReconstructor, +from ptychodus.api.interpolate import BarycentricArrayInterpolator +from ptychodus.api.reconstructor import ReconstructInput + +from ..processing.subprocess_reconstructor import SubprocessReconstructor +from ._payload import ( + PtychoNNReconstructConfig, + PtychoNNTrainConfig, + ReconstructPayload, + TrainPayload, ) - -from .model import PtychoNNModelProvider from .settings import PtychoNNModelSettings, PtychoNNTrainingSettings -logger = logging.getLogger(__name__) - - -class CenterBoxMeanPhaseCenteringStrategy: # TODO USE - def __call__(self, array: ComplexArrayType) -> ComplexArrayType: - one_third_height = array.shape[-2] // 3 - one_third_width = array.shape[-1] // 3 - - amplitude = numpy.absolute(array) - phase = numpy.angle(array) - - center_box_mean_phase = phase[ - one_third_height : one_third_height * 2, one_third_width : one_third_width * 2 - ].mean() - - return amplitude * numpy.exp(1j * (phase - center_box_mean_phase)) - - -class PtychoNNTrainableReconstructor(TrainableReconstructor): - MODEL_FILE_FILTER: Final[str] = 'PyTorch Lightning Checkpoint Files (*.ckpt)' - MODEL_FILE_EXTENSION: Final[str] = '.ckpt' - TRAINING_DATA_FILE_FILTER: Final[str] = 'NumPy Zipped Archive (*.npz)' - PATCHES_KEY: Final[str] = 'real' - PATTERNS_KEY: Final[str] = 'reciprocal' - - def __init__( - self, - model_settings: PtychoNNModelSettings, - training_settings: PtychoNNTrainingSettings, - model_provider: PtychoNNModelProvider, - ) -> None: - self._model_settings = model_settings - self._training_settings = training_settings - self._model_provider = model_provider - - ptychonn_version = version('ptychonn') - logger.info(f'\tPtychoNN {ptychonn_version}') - - @property - def name(self) -> str: - return self._model_provider.get_model_name() - - def get_progress_goal(self) -> int: - return 0 +__all__ = [ + 'build_reconstructor', +] - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - # TODO data size/shape requirements to GUI - data = parameters.diffraction_patterns - data_size = data.shape[-1] - - if data_size != data.shape[-2]: - raise ValueError('PtychoNN expects square diffraction data!') - - is_data_size_pow2 = data_size & (data_size - 1) == 0 and data_size > 0 - - if not is_data_size_pow2: - raise ValueError('PtychoNN expects that the diffraction data size is a power of two!') +logger = logging.getLogger(__name__) - model = self._model_provider.get_model() - logger.debug('Inferring...') - object_patches = ptychonn.infer( - data=data.astype(numpy.float32), - model=model, +_RECONSTRUCT_ENTRY = 'ptychodus.model.ptychonn._subprocess:run_reconstruct' +_TRAIN_ENTRY = 'ptychodus.model.ptychonn._subprocess:run_train' +_PATCHES_KEY = 'real' +_PATTERNS_KEY = 'reciprocal' + + +def _export_training_data( + file_path: Path, parameters: ReconstructInput, *, num_channels: int +) -> None: + object_geometry = parameters.product.object_.get_geometry() + interpolator = BarycentricArrayInterpolator(parameters.product.object_.get_array()) + probe_extent = ImageExtent( + width_px=parameters.product.probes.width_px, + height_px=parameters.product.probes.height_px, + ) + patches = numpy.zeros( + (len(parameters.product.probe_positions), num_channels, *probe_extent.get_shape()), + dtype=numpy.float32, + ) + + for index, scan_point in enumerate(parameters.product.probe_positions): + object_point = object_geometry.map_coordinates_probe_to_object(scan_point) + patch = interpolator.get_patch( + object_point.coordinate_x_px, + object_point.coordinate_y_px, + probe_extent.width_px, + probe_extent.height_px, ) - - logger.debug('Stitching...') - object_array = parameters.product.object_.get_array() - object_geometry = parameters.product.object_.get_geometry() - stitcher = BarycentricArrayStitcher( - upper=numpy.zeros_like(object_array), lower=numpy.zeros_like(object_array, dtype=float) + patches[index, 0, :, :] = numpy.angle(patch) + if num_channels > 1: + patches[index, 1, :, :] = numpy.absolute(patch) + + logger.debug(f'Writing "{file_path}" as "NPZ"') + contents = { + _PATTERNS_KEY: parameters.diffraction_patterns.astype(numpy.float32), + _PATCHES_KEY: patches, + } + numpy.savez_compressed(file_path, allow_pickle=False, **contents) + + +def _build_reconstruct_config( + *, + enable_amplitude: bool, + model_settings: PtychoNNModelSettings, + training_settings: PtychoNNTrainingSettings, +) -> PtychoNNReconstructConfig: + return PtychoNNReconstructConfig( + enable_amplitude=enable_amplitude, + num_convolution_kernels=model_settings.num_convolution_kernels.get_value(), + use_batch_normalization=model_settings.use_batch_normalization.get_value(), + max_learning_rate=float(training_settings.max_learning_rate.get_value()), + min_learning_rate=float(training_settings.min_learning_rate.get_value()), + ) + + +def _build_train_config( + *, + enable_amplitude: bool, + model_settings: PtychoNNModelSettings, + training_settings: PtychoNNTrainingSettings, +) -> PtychoNNTrainConfig: + return PtychoNNTrainConfig( + enable_amplitude=enable_amplitude, + num_convolution_kernels=model_settings.num_convolution_kernels.get_value(), + use_batch_normalization=model_settings.use_batch_normalization.get_value(), + max_learning_rate=float(training_settings.max_learning_rate.get_value()), + min_learning_rate=float(training_settings.min_learning_rate.get_value()), + batch_size=model_settings.batch_size.get_value(), + training_epochs=training_settings.training_epochs.get_value(), + status_interval_in_epochs=training_settings.status_interval_in_epochs.get_value(), + validation_set_fractional_size=float( + training_settings.validation_set_fractional_size.get_value() + ), + ) + + +def build_reconstructor( + display_name: str, + *, + enable_amplitude: bool, + model_settings: PtychoNNModelSettings, + training_settings: PtychoNNTrainingSettings, +) -> SubprocessReconstructor: + num_channels = 2 if enable_amplitude else 1 + + def build_reconstruct_payload( + parameters: ReconstructInput, loaded_model_path: Path | None + ) -> ReconstructPayload: + return ReconstructPayload( + config=_build_reconstruct_config( + enable_amplitude=enable_amplitude, + model_settings=model_settings, + training_settings=training_settings, + ), + model_path=loaded_model_path, + reconstruct_input=parameters, ) - for scan_point, object_patch_channels in zip( - parameters.product.probe_positions, object_patches - ): - patch_array = numpy.exp(1j * object_patch_channels[0]) - - if object_patch_channels.shape[0] == 2: - patch_array *= object_patch_channels[1] - else: - patch_array *= 0.5 - - object_point = object_geometry.map_coordinates_probe_to_object(scan_point) - stitcher.add_patch( - object_point.coordinate_x_px, object_point.coordinate_y_px, patch_array - ) - - object_ = Object( - array=stitcher.stitch(), - pixel_geometry=object_geometry.get_pixel_geometry(), - center=object_geometry.get_center(), - layer_spacing_m=parameters.product.object_.layer_spacing_m, + def build_train_payload(input_path: Path, output_path: Path) -> TrainPayload: + return TrainPayload( + config=_build_train_config( + enable_amplitude=enable_amplitude, + model_settings=model_settings, + training_settings=training_settings, + ), + input_path=input_path, + output_path=output_path, ) - losses: Sequence[LossValue] = list() - - product = Product( - metadata=parameters.product.metadata, - probe_positions=parameters.product.probe_positions, - probes=parameters.product.probes, - object_=object_, - losses=losses, - ) - - yield ReconstructOutput(product) - - def is_model_loaded(self): - return True # TODO - - def get_model_file_filter(self) -> str: - return self.MODEL_FILE_FILTER - - def get_model_file_extension(self) -> str: - return self.MODEL_FILE_EXTENSION - def load_model_from_file(self, file_path: Path) -> None: - self._model_provider.load_model_from_file(file_path) - - def save_model(self, file_path: Path) -> None: - self._model_provider.save_model(file_path) - - def get_training_data_file_filter(self) -> str: - return self.TRAINING_DATA_FILE_FILTER - - def export_training_data(self, file_path: Path, parameters: ReconstructInput) -> None: - object_geometry = parameters.product.object_.get_geometry() - interpolator = BarycentricArrayInterpolator(parameters.product.object_.get_array()) - num_channels = self._model_provider.get_num_channels() - probe_extent = ImageExtent( - width_px=parameters.product.probes.width_px, - height_px=parameters.product.probes.height_px, - ) - patches = numpy.zeros( - (len(parameters.product.probe_positions), num_channels, *probe_extent.get_shape()), - dtype=numpy.float32, - ) - - for index, scan_point in enumerate(parameters.product.probe_positions): - object_point = object_geometry.map_coordinates_probe_to_object(scan_point) - patch = interpolator.get_patch( - object_point.coordinate_x_px, - object_point.coordinate_y_px, - probe_extent.width_px, - probe_extent.height_px, - ) - patches[index, 0, :, :] = numpy.angle(patch) - - if num_channels > 1: - patches[index, 1, :, :] = numpy.absolute(patch) - - logger.debug(f'Writing "{file_path}" as "NPZ"') - contents = { - self.PATTERNS_KEY: parameters.diffraction_patterns.astype(numpy.float32), - self.PATCHES_KEY: patches, - } - numpy.savez_compressed(file_path, allow_pickle=False, **contents) - - def train(self, input_path: Path, output_path: Path) -> Iterator[TrainOutput]: - logger.debug(f'Reading "{input_path}" as "NPZ"') - training_data = numpy.load(input_path) - - model = self._model_provider.get_model() - logger.debug('Training...') - training_set_fractional_size = ( - 1 - self._training_settings.validation_set_fractional_size.get_value() - ) - trainer, trainer_log = ptychonn.train( - model=model, - batch_size=self._model_settings.batch_size.get_value(), - out_dir=None, - X_train=training_data[self.PATTERNS_KEY], - Y_train=training_data[self.PATCHES_KEY], - epochs=self._training_settings.training_epochs.get_value(), - training_fraction=float(training_set_fractional_size), - log_frequency=self._training_settings.status_interval_in_epochs.get_value(), - strategy='ddp_notebook', - ) - self._model_provider.set_trainer(trainer) - - training_loss: list[LossValue] = [] - validation_loss: list[LossValue] = [] - - for epoch, entry in enumerate(trainer_log.logs): - try: - tloss = LossValue(epoch, entry['training_loss']) - vloss = LossValue(epoch, entry['validation_loss']) - except KeyError: - pass - else: - training_loss.append(tloss) - training_loss.append(vloss) - - yield TrainOutput( - training_loss=training_loss, - validation_loss=validation_loss, - ) + def export_training_data(file_path: Path, parameters: ReconstructInput) -> None: + _export_training_data(file_path, parameters, num_channels=num_channels) + + return SubprocessReconstructor( + name=display_name, + reconstruct_entry_point=_RECONSTRUCT_ENTRY, + progress_goal_fn=lambda: 0, + build_reconstruct_payload=build_reconstruct_payload, + is_trainable=True, + train_entry_point=_TRAIN_ENTRY, + build_train_payload=build_train_payload, + model_file_filter='PyTorch Lightning Checkpoint Files (*.ckpt)', + model_file_extension='.ckpt', + training_data_file_filter='NumPy Zipped Archive (*.npz)', + export_training_data=export_training_data, + ) diff --git a/src/ptychodus/model/ptychopinn/_payload.py b/src/ptychodus/model/ptychopinn/_payload.py new file mode 100644 index 000000000..086588068 --- /dev/null +++ b/src/ptychodus/model/ptychopinn/_payload.py @@ -0,0 +1,71 @@ +"""Payload dataclasses for the PtychoPINN (TensorFlow) subprocess entry points. + +``ptycho`` is an optional extra, so it is imported under +:data:`~typing.TYPE_CHECKING` only: this module has ``from __future__ import +annotations`` and both payloads are dataclasses, whose field annotations are +never evaluated. That keeps ``ptychodus.model`` importable without the extra +installed. The factory defers the matching runtime imports into its config +builders. + +Only ``ptycho.config.config`` may be reached from the parent: it pulls in just +``dataclasses``, ``enum``, and ``typing`` -- no TensorFlow, as evidenced by the +fact that it imports cleanly in an environment where TensorFlow is absent +entirely. Anything else under ``ptycho.*`` (``raw_data``, ``probe``, +``tf_helper``, ``loader``) needs TensorFlow and must stay inside the child; see +[tests/test_no_gpu_context.py](../../../../tests/test_no_gpu_context.py). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Final + +from ptychodus.api.reconstructor import ReconstructInput + +if TYPE_CHECKING: + from ptycho.config.config import InferenceConfig, TrainingConfig + +__all__ = [ + 'MODEL_FILE_NAME', + 'ReconstructPayload', + 'TrainPayload', +] + + +MODEL_FILE_NAME: Final[str] = 'wts.h5.zip' +"""Name ``ptycho.model_manager`` gives the weights archive inside a bundle.""" + + +@dataclass(frozen=True) +class ReconstructPayload: + """Everything the child needs to load a bundle and run inference once.""" + + # Fully populated by the parent, including the nested ModelConfig whose + # ``N`` is the (already validated square) diffraction pattern size. + inference_config: InferenceConfig + + # Directory the child loads the model from. The parent has already + # unpacked an outer .zip if there was one. + model_bundle_dir: Path + + # Arguments to RawData.generate_grouped_data, which only the child can call. + n_nearest_neighbors: int + n_samples: int + + reconstruct_input: ReconstructInput + + +@dataclass(frozen=True) +class TrainPayload: + """Everything the child needs to run one training session. + + ``training_config.model.N`` is left at its default: the pattern size is + known only once the child loads ``train_data.npz``, so the child fills it + in before handing the config to ``run_cdi_example``. + """ + + training_config: TrainingConfig + + input_path: Path + output_path: Path diff --git a/src/ptychodus/model/ptychopinn/_subprocess.py b/src/ptychodus/model/ptychopinn/_subprocess.py new file mode 100644 index 000000000..a671016df --- /dev/null +++ b/src/ptychodus/model/ptychopinn/_subprocess.py @@ -0,0 +1,147 @@ +"""Child-side subprocess entry points for the PtychoPINN (TensorFlow) backend. + +This module runs INSIDE a spawned subprocess. It is the only place in the +ptychodus tree that is allowed to import tensorflow or the TensorFlow-backed +``ptycho`` submodules. The parent-side ptychodus process never imports this +module. + +The payload carries finished ``ptycho`` config objects and an already-unpacked +bundle directory (see :mod:`._payload`), so the child reads no ptychodus +settings and does no archive handling -- it converts the input product to +``RawData``, runs the model, and converts the result back. +""" + +from __future__ import annotations + +import logging +import pickle +from collections.abc import Sequence +from multiprocessing.queues import Queue +from typing import Any + +import numpy + +from ptychodus.api.object import Object +from ptychodus.api.product import LossValue, Product +from ptychodus.api.reconstructor import ReconstructOutput, TrainOutput + +from ..processing.subprocess_reconstructor import ( + TAG_MODEL_SAVED, + TAG_OUTPUT, + TAG_TRAIN_OUTPUT, +) +from ._payload import MODEL_FILE_NAME, ReconstructPayload, TrainPayload + +logger = logging.getLogger(__name__) + + +def _create_raw_data(parameters: Any) -> Any: + from ptycho.raw_data import RawData + + object_geometry = parameters.product.object_.get_geometry() + position_x_px: list[float] = list() + position_y_px: list[float] = list() + + for scan_point in parameters.product.probe_positions: + object_point = object_geometry.map_coordinates_probe_to_object(scan_point) + position_x_px.append(object_point.coordinate_x_px) + position_y_px.append(object_point.coordinate_y_px) + + return RawData.from_coords_without_pc( + xcoords=numpy.array(position_x_px), + ycoords=numpy.array(position_y_px), + diff3d=parameters.diffraction_patterns, + probeGuess=parameters.product.probes.get_probe_no_opr().get_incoherent_mode(0), + scan_index=numpy.zeros(len(parameters.product.probe_positions), dtype=int), + objectGuess=parameters.product.object_.get_layer(0), + ) + + +def run_reconstruct(payload: ReconstructPayload, queue: Queue[Any]) -> None: + """Child entry point for inference.""" + from ptycho.config.config import update_legacy_dict + from ptycho.workflows.components import load_inference_bundle + import ptycho.loader + import ptycho.params + import ptycho.probe + import ptycho.tf_helper + + inference_config = payload.inference_config + update_legacy_dict(ptycho.params.cfg, inference_config) + + model_obj, _config = load_inference_bundle(payload.model_bundle_dir) + + parameters = payload.reconstruct_input + test_raw_data = _create_raw_data(parameters) + ptycho.probe.set_probe_guess(None, test_raw_data.probeGuess) + + test_dataset = test_raw_data.generate_grouped_data( + inference_config.model.N, + K=payload.n_nearest_neighbors, + nsamples=payload.n_samples, + ) + test_data_container = ptycho.loader.load( + lambda: test_dataset, test_raw_data.probeGuess, which=None, create_split=False + ) + + try: + intensity_scale = ptycho.params.get('intensity_scale') + except KeyError as exc: + raise RuntimeError('Missing intensity_scale in ptycho.params.cfg') from exc + + obj_tensor_full = model_obj.predict( + [test_data_container.X * intensity_scale, test_data_container.local_offsets] + ) + object_out_array = ptycho.tf_helper.reassemble_position( + obj_tensor_full, test_data_container.global_offsets, M=20 + ) + + object_in = parameters.product.object_ + object_out = Object( + array=numpy.squeeze(object_out_array), + layer_spacing_m=object_in.layer_spacing_m, + pixel_geometry=object_in.get_pixel_geometry(), + center=object_in.get_center(), + ) + losses: Sequence[LossValue] = list() + product = Product( + metadata=parameters.product.metadata, + probe_positions=parameters.product.probe_positions, + probes=parameters.product.probes, + object_=object_out, + losses=losses, + ) + + queue.put((TAG_OUTPUT, pickle.dumps(ReconstructOutput(product)))) + + +def run_train(payload: TrainPayload, queue: Queue[Any]) -> None: + """Child entry point for training.""" + from ptycho.config.config import update_legacy_dict + from ptycho.raw_data import RawData + from ptycho.workflows.components import run_cdi_example, save_outputs + import ptycho.model_manager + import ptycho.params + + test_raw_data = RawData.from_file(payload.input_path / 'test_data.npz') + train_raw_data = RawData.from_file(payload.input_path / 'train_data.npz') + + model_size = train_raw_data.diff3d.shape[-1] + if train_raw_data.diff3d.shape[-2] != model_size: + raise ValueError('Model requires square diffraction patterns!') + + # Only the child sees the training arrays, so it is the only place that can + # resolve the model size the parent left unset. + training_config = payload.training_config + training_config.model.N = model_size + update_legacy_dict(ptycho.params.cfg, training_config) + + recon_amp, recon_phase, train_results = run_cdi_example( + train_raw_data, test_raw_data, training_config + ) + model_path = payload.output_path / MODEL_FILE_NAME + ptycho.model_manager.save(payload.output_path) + save_outputs(recon_amp, recon_phase, train_results, str(payload.output_path)) + + queue.put((TAG_TRAIN_OUTPUT, pickle.dumps(TrainOutput()))) + queue.put((TAG_MODEL_SAVED, str(model_path))) diff --git a/src/ptychodus/model/ptychopinn/core.py b/src/ptychodus/model/ptychopinn/core.py index b7b1f3565..177edcbb9 100644 --- a/src/ptychodus/model/ptychopinn/core.py +++ b/src/ptychodus/model/ptychopinn/core.py @@ -1,5 +1,6 @@ from collections.abc import Iterator from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec import logging from ptychodus.api.reconstructor import ( @@ -20,6 +21,11 @@ logger = logging.getLogger(__name__) +def _ptychopinn_available() -> bool: + """Return True iff the ``ptycho`` package is importable, without importing it.""" + return find_spec('ptycho') is not None + + class PtychoPINNReconstructorLibrary(ReconstructorLibrary): def __init__( self, settings_registry: SettingsRegistry, is_developer_mode_enabled: bool @@ -31,34 +37,30 @@ def __init__( self.enumerators = PtychoPINNEnumerators() self._reconstructors: list[TrainableReconstructor] = list() - try: - from .reconstructor import PtychoPINNTrainableReconstructor - except ModuleNotFoundError: + if not _ptychopinn_available(): logger.info('PtychoPINN not found.') if is_developer_mode_enabled: for reconstructor in ('PINN', 'Supervised'): self._reconstructors.append(NullReconstructor(reconstructor)) - else: + return + + try: + ptychopinn_version = version('ptychopinn') + except PackageNotFoundError: try: - ptychopinn_version = version('ptychopinn') - except PackageNotFoundError: ptychopinn_version = version('ptycho') + except PackageNotFoundError: + ptychopinn_version = 'unknown' - logger.info(f'PtychoPINN {ptychopinn_version}') + logger.info(f'PtychoPINN {ptychopinn_version}') + from .reconstructor import build_reconstructor + + for mode in ('PINN', 'Supervised'): self._reconstructors.append( - PtychoPINNTrainableReconstructor( - 'PINN', - self.model_settings, - self.inference_settings, - self.training_settings, - is_developer_mode_enabled=is_developer_mode_enabled, - ) - ) - self._reconstructors.append( - PtychoPINNTrainableReconstructor( - 'Supervised', + build_reconstructor( + mode, self.model_settings, self.inference_settings, self.training_settings, diff --git a/src/ptychodus/model/ptychopinn/reconstructor.py b/src/ptychodus/model/ptychopinn/reconstructor.py index 9926072b0..b4dadc6f8 100644 --- a/src/ptychodus/model/ptychopinn/reconstructor.py +++ b/src/ptychodus/model/ptychopinn/reconstructor.py @@ -1,31 +1,28 @@ -from collections.abc import Iterator, Sequence -from pathlib import Path -from typing import Any, Final +"""Parent-side factory that builds a :class:`SubprocessReconstructor` for PtychoPINN. + +Zero tensorflow imports. All GPU work runs inside a spawned child; see +:mod:`._subprocess` for the child entry points. + +This module does reach ``ptycho.config.config`` -- but only from inside the +config builders, so the import happens on the first reconstruct/train call +rather than at composition-root time. That subpackage is TensorFlow-free; see +the note in :mod:`._payload` for why it is allowed parent-side. +""" + +from __future__ import annotations + import logging import shutil import tempfile import zipfile - -import numpy - -from ptycho.config.config import InferenceConfig, ModelConfig, TrainingConfig, update_legacy_dict -from ptycho.raw_data import RawData -from ptycho.workflows.components import load_inference_bundle -import ptycho.loader -import ptycho.model_manager -import ptycho.params +from pathlib import Path +from typing import Any from ptychodus.api.io import save_ptychopinn_training_data -from ptychodus.api.object import Object -from ptychodus.api.product import Product -from ptychodus.api.reconstructor import ( - LossValue, - ReconstructInput, - ReconstructOutput, - TrainOutput, - TrainableReconstructor, -) +from ptychodus.api.reconstructor import ReconstructInput +from ..processing.subprocess_reconstructor import SubprocessReconstructor +from ._payload import MODEL_FILE_NAME, ReconstructPayload, TrainPayload from .settings import ( PtychoPINNInferenceSettings, PtychoPINNModelSettings, @@ -33,236 +30,195 @@ ) __all__ = [ - 'PtychoPINNTrainableReconstructor', + 'build_reconstructor', ] logger = logging.getLogger(__name__) -def create_raw_data(parameters: ReconstructInput) -> RawData: - object_geometry = parameters.product.object_.get_geometry() - position_x_px: list[float] = list() - position_y_px: list[float] = list() - - for scan_point in parameters.product.probe_positions: - object_point = object_geometry.map_coordinates_probe_to_object(scan_point) - position_x_px.append(object_point.coordinate_x_px) - position_y_px.append(object_point.coordinate_y_px) - - return RawData.from_coords_without_pc( - xcoords=numpy.array(position_x_px), - ycoords=numpy.array(position_y_px), - diff3d=parameters.diffraction_patterns, - probeGuess=parameters.product.probes.get_probe_no_opr().get_incoherent_mode(0), - # assume that all patches are from the same object - scan_index=numpy.zeros(len(parameters.product.probe_positions), dtype=int), - objectGuess=parameters.product.object_.get_layer(0), - ) - - -class PtychoPINNTrainableReconstructor(TrainableReconstructor): - MODEL_FILE_NAME: Final[str] = 'wts.h5.zip' - - def __init__( - self, - name: str, - model_settings: PtychoPINNModelSettings, - inference_settings: PtychoPINNInferenceSettings, - training_settings: PtychoPINNTrainingSettings, - *, - is_developer_mode_enabled: bool, - ) -> None: - super().__init__() - self._name = name - self._model_settings = model_settings - self._inference_settings = inference_settings - self._training_settings = training_settings - self.__model: Any = None - self._config: dict[str, Any] = dict() - self._is_developer_mode_enabled = is_developer_mode_enabled - self._model_bundle_dir: Path | None = None - - def _create_model_config(self, model_size: int) -> ModelConfig: - return ModelConfig( - N=model_size, - gridsize=self._model_settings.gridsize.get_value(), - n_filters_scale=self._model_settings.n_filters_scale.get_value(), - model_type=self._name.lower(), - amp_activation=self._model_settings.amp_activation.get_value(), - object_big=self._model_settings.object_big.get_value(), - probe_big=self._model_settings.probe_big.get_value(), - probe_mask=self._model_settings.probe_mask.get_value(), - pad_object=self._model_settings.pad_object.get_value(), - probe_scale=self._model_settings.probe_scale.get_value(), - gaussian_smoothing_sigma=self._model_settings.gaussian_smoothing_sigma.get_value(), - ) - - @property - def name(self) -> str: - return self._name - - def get_progress_goal(self) -> int: - return 0 +_RECONSTRUCT_ENTRY = 'ptychodus.model.ptychopinn._subprocess:run_reconstruct' +_TRAIN_ENTRY = 'ptychodus.model.ptychopinn._subprocess:run_train' - @property - def _model(self) -> Any: # TODO tensorflow.keras.Model | None - if self.__model is None: - raise RuntimeError('Model not loaded!') - return self.__model +def _build_model_config(model_settings: PtychoPINNModelSettings, name: str, model_size: int) -> Any: + from ptycho.config.config import ModelConfig - def _reconstruct_image(self, test_data: ptycho.loader.PtychoDataContainer) -> Any: - try: - intensity_scale = ptycho.params.get('intensity_scale') - except KeyError as exc: - raise RuntimeError('Missing intensity_scale in ptycho.params.cfg') from exc - return self._model.predict([test_data.X * intensity_scale, test_data.local_offsets]) + return ModelConfig( + N=model_size, + gridsize=model_settings.gridsize.get_value(), + n_filters_scale=model_settings.n_filters_scale.get_value(), + model_type=name.lower(), + amp_activation=model_settings.amp_activation.get_value(), + object_big=model_settings.object_big.get_value(), + probe_big=model_settings.probe_big.get_value(), + probe_mask=model_settings.probe_mask.get_value(), + pad_object=model_settings.pad_object.get_value(), + probe_scale=model_settings.probe_scale.get_value(), + gaussian_smoothing_sigma=model_settings.gaussian_smoothing_sigma.get_value(), + ) - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - model_size = parameters.diffraction_patterns.shape[-1] - if parameters.diffraction_patterns.shape[-2] != model_size: - raise ValueError('Model requires square diffraction patterns!') - - model_config = self._create_model_config(model_size) - inference_config = InferenceConfig( - model=model_config, - model_path=Path(), # not used - test_data_file=Path(), # not used - debug=self._is_developer_mode_enabled, - output_dir=Path(), # not used - ) +def _build_inference_config( + model_settings: PtychoPINNModelSettings, + name: str, + model_size: int, + *, + is_developer_mode_enabled: bool, +) -> Any: + from ptycho.config.config import InferenceConfig + + return InferenceConfig( + model=_build_model_config(model_settings, name, model_size), + model_path=Path(), # not used + test_data_file=Path(), # not used + debug=is_developer_mode_enabled, + output_dir=Path(), # not used + ) - # Update global params with new-style config - update_legacy_dict(ptycho.params.cfg, inference_config) - # Create RawData - test_raw_data = create_raw_data(parameters) - ptycho.probe.set_probe_guess(None, test_raw_data.probeGuess) +def _build_training_config( + model_settings: PtychoPINNModelSettings, + training_settings: PtychoPINNTrainingSettings, + name: str, +) -> Any: + from ptycho.config.config import TrainingConfig + + # ``N`` is unknown here: it comes from the training data the child loads. + # The child overwrites ``model.N`` before using this config. + return TrainingConfig( + model=_build_model_config(model_settings, name, model_size=0), + train_data_file=Path(), + test_data_file=None, + batch_size=training_settings.batch_size.get_value(), + nepochs=training_settings.nepochs.get_value(), + mae_weight=training_settings.mae_weight.get_value(), + nll_weight=training_settings.nll_weight.get_value(), + realspace_mae_weight=training_settings.realspace_mae_weight.get_value(), + realspace_weight=training_settings.realspace_weight.get_value(), + nphotons=training_settings.nphotons.get_value(), + positions_provided=training_settings.positions_provided.get_value(), + probe_trainable=training_settings.probe_trainable.get_value(), + intensity_scale_trainable=training_settings.intensity_scale_trainable.get_value(), + output_dir=Path(), + ) - # Group overlapping scan positions - test_dataset = test_raw_data.generate_grouped_data( - model_config.N, - K=self._inference_settings.n_nearest_neighbors.get_value(), - nsamples=self._inference_settings.n_samples.get_value(), - ) - # Create PtychoDataContainer - test_data_container = ptycho.loader.load( - lambda: test_dataset, test_raw_data.probeGuess, which=None, create_split=False - ) +def _extract_bundle_dir(model_bundle_path: Path) -> Path: + """Resolve the recorded model path to a directory the child can load from. - # Perform reconstruction - obj_tensor_full = self._reconstruct_image(test_data_container) + Accepts the bundle directory itself, a ``wts.h5.zip`` inside a bundle + directory, or an outer zip archive to unpack into a fresh tempdir. Runs + parent-side because it is pure zipfile/tempfile work with no GPU + involvement. + """ + if model_bundle_path.name == MODEL_FILE_NAME: + return model_bundle_path.parent - # Process the reconstructed image - object_out_array = ptycho.tf_helper.reassemble_position( - obj_tensor_full, test_data_container.global_offsets, M=20 - ) + if model_bundle_path.suffix == '.zip': + bundle_dir = Path(tempfile.mkdtemp(prefix='ptychopinn-bundle-')) + logger.debug(f'Extracting bundle "{model_bundle_path}" -> "{bundle_dir}"') - object_in = parameters.product.object_ - object_out = Object( - array=numpy.squeeze(object_out_array), - layer_spacing_m=object_in.layer_spacing_m, - pixel_geometry=object_in.get_pixel_geometry(), - center=object_in.get_center(), - ) - losses: Sequence[LossValue] = list() - - product = Product( - metadata=parameters.product.metadata, - probe_positions=parameters.product.probe_positions, - probes=parameters.product.probes, - object_=object_out, - losses=losses, - ) + with zipfile.ZipFile(model_bundle_path) as archive: + archive.extractall(bundle_dir) - yield ReconstructOutput(product) - - def is_model_loaded(self): - return True # TODO - - def get_model_file_filter(self) -> str: - return 'Zipped Archive (*.zip)' - - def get_model_file_extension(self) -> str: - return '.zip' - - def load_model_from_file(self, file_path: Path) -> None: - # TODO model path to/from settings - self._inference_settings.model_path.set_value(file_path) - - if file_path.name == self.MODEL_FILE_NAME: - # Loose bundle directory containing wts.h5.zip + auxiliary files. - bundle_dir = file_path.parent - elif file_path.suffix == '.zip': - # Zipped bundle (e.g. produced by save_model). Unpack to a tempdir - # that lives for the process lifetime; load_inference_bundle may - # keep lazy file handles open against this directory. - bundle_dir = Path(tempfile.mkdtemp(prefix='ptychopinn-bundle-')) - with zipfile.ZipFile(file_path) as archive: - archive.extractall(bundle_dir) - else: - logger.warning(f"PtychoPINN expects the file name '{self.MODEL_FILE_NAME}'.") - bundle_dir = file_path.parent - - # global config (ptycho.params.cfg) updated during load - self.__model, self._config = load_inference_bundle(bundle_dir) - self._model_bundle_dir = bundle_dir - # TODO sync ptycho.params.cfg with settings after load - - def save_model(self, file_path: Path) -> None: - if self._model_bundle_dir is None: - raise RuntimeError('Cannot save PtychoPINN model: model is not loaded.') - archive_stem = str(file_path.with_suffix('')) - logger.debug(f'Archiving bundle "{self._model_bundle_dir}" -> "{file_path}"') - shutil.make_archive(archive_stem, 'zip', root_dir=self._model_bundle_dir) - - def get_training_data_file_filter(self) -> str: - return 'NumPy Zipped Archive (*.npz)' - - def export_training_data(self, file_path: Path, parameters: ReconstructInput) -> None: - save_ptychopinn_training_data(file_path, parameters, multimodal_probe=False) + return bundle_dir - def train(self, input_path: Path, output_path: Path) -> Iterator[TrainOutput]: - test_raw_data = RawData.from_file(input_path / 'test_data.npz') # TODO RawData | None - train_raw_data = RawData.from_file(input_path / 'train_data.npz') + logger.warning( + f'PtychoPINN expected the file name {MODEL_FILE_NAME!r}; got {model_bundle_path.name!r}.' + ) + return model_bundle_path.parent + + +def _save_model_bundle(loaded_from: Path, dest: Path) -> None: + """Archive the loaded bundle directory (or copy the bundle .zip) to ``dest``. + + The child records either a bundle directory (right after training) or a + ``wts.h5.zip`` file (right after inference load). This function normalises + to a zip archive at ``dest``. + """ + if loaded_from.is_dir(): + archive_stem = str(dest.with_suffix('')) + logger.debug(f'Archiving bundle {loaded_from!r} -> {dest!r}') + shutil.make_archive(archive_stem, 'zip', root_dir=loaded_from) + return + + if loaded_from.suffix == '.zip': + shutil.copyfile(loaded_from, dest) + return + + # wts.h5.zip inside a bundle dir — archive the parent dir. + if loaded_from.name.endswith('.zip'): + parent = loaded_from.parent + archive_stem = str(dest.with_suffix('')) + shutil.make_archive(archive_stem, 'zip', root_dir=parent) + return + + raise RuntimeError(f'Cannot save PtychoPINN model: unrecognized source path {loaded_from!r}.') + + +def build_reconstructor( + name: str, + model_settings: PtychoPINNModelSettings, + inference_settings: PtychoPINNInferenceSettings, + training_settings: PtychoPINNTrainingSettings, + *, + is_developer_mode_enabled: bool, +) -> SubprocessReconstructor: + # Source model path -> extracted bundle directory, so repeated reconstructs + # against the same model unpack the archive only once. + bundle_dir_cache: dict[Path, Path] = dict() + + def build_reconstruct_payload( + parameters: ReconstructInput, loaded_model_path: Path | None + ) -> ReconstructPayload: + if loaded_model_path is None: + raise RuntimeError('Cannot reconstruct: no PtychoPINN model has been loaded.') - model_size = train_raw_data.diff3d.shape[-1] + model_size = parameters.diffraction_patterns.shape[-1] - if train_raw_data.diff3d.shape[-2] != model_size: + if parameters.diffraction_patterns.shape[-2] != model_size: raise ValueError('Model requires square diffraction patterns!') - model_config = self._create_model_config(model_size) - training_config = TrainingConfig( - model=model_config, - train_data_file=Path(), # not used - test_data_file=None, # not used - batch_size=self._training_settings.batch_size.get_value(), - nepochs=self._training_settings.nepochs.get_value(), - mae_weight=self._training_settings.mae_weight.get_value(), - nll_weight=self._training_settings.nll_weight.get_value(), - realspace_mae_weight=self._training_settings.realspace_mae_weight.get_value(), - realspace_weight=self._training_settings.realspace_weight.get_value(), - nphotons=self._training_settings.nphotons.get_value(), # TODO get from product - positions_provided=self._training_settings.positions_provided.get_value(), - probe_trainable=self._training_settings.probe_trainable.get_value(), - intensity_scale_trainable=self._training_settings.intensity_scale_trainable.get_value(), - output_dir=Path(), # not used + try: + bundle_dir = bundle_dir_cache[loaded_model_path] + except KeyError: + bundle_dir = _extract_bundle_dir(loaded_model_path) + bundle_dir_cache[loaded_model_path] = bundle_dir + + return ReconstructPayload( + inference_config=_build_inference_config( + model_settings, + name, + model_size, + is_developer_mode_enabled=is_developer_mode_enabled, + ), + model_bundle_dir=bundle_dir, + n_nearest_neighbors=inference_settings.n_nearest_neighbors.get_value(), + n_samples=inference_settings.n_samples.get_value(), + reconstruct_input=parameters, ) - # Update global params with new-style config - update_legacy_dict(ptycho.params.cfg, training_config) - - from ptycho.workflows.components import run_cdi_example, save_outputs - - recon_amp, recon_phase, train_results = run_cdi_example( - train_raw_data, test_raw_data, training_config + def build_train_payload(input_path: Path, output_path: Path) -> TrainPayload: + return TrainPayload( + training_config=_build_training_config(model_settings, training_settings, name), + input_path=input_path, + output_path=output_path, ) - model_path = output_path / self.MODEL_FILE_NAME - ptycho.model_manager.save(output_path) - self._model_bundle_dir = output_path - save_outputs(recon_amp, recon_phase, train_results, str(output_path)) - self.load_model_from_file(model_path) - yield TrainOutput() # TODO yield losses & progress + def export_training_data(file_path: Path, parameters: ReconstructInput) -> None: + save_ptychopinn_training_data(file_path, parameters, multimodal_probe=False) + + return SubprocessReconstructor( + name=name, + reconstruct_entry_point=_RECONSTRUCT_ENTRY, + progress_goal_fn=lambda: 0, + build_reconstruct_payload=build_reconstruct_payload, + is_trainable=True, + train_entry_point=_TRAIN_ENTRY, + build_train_payload=build_train_payload, + model_file_filter='Zipped Archive (*.zip)', + model_file_extension='.zip', + training_data_file_filter='NumPy Zipped Archive (*.npz)', + export_training_data=export_training_data, + save_model=_save_model_bundle, + ) diff --git a/src/ptychodus/model/ptychopinn_torch/_payload.py b/src/ptychodus/model/ptychopinn_torch/_payload.py new file mode 100644 index 000000000..f5d35921d --- /dev/null +++ b/src/ptychodus/model/ptychopinn_torch/_payload.py @@ -0,0 +1,81 @@ +"""Payload dataclasses for the PtychoPINN-Torch subprocess entry points. + +``ptycho_torch`` is an optional extra, so it is imported under +:data:`~typing.TYPE_CHECKING` only: this module has ``from __future__ import +annotations`` and both payloads are dataclasses, whose field annotations are +never evaluated. That keeps ``ptychodus.model`` importable without the extra +installed. The factory defers the matching runtime import into +``_build_configs``. + +Unlike ``ptycho.config.config``, ``ptycho_torch.config_params`` does pull torch +into the parent (``ptycho_torch/__init__.py`` imports it unconditionally). That +is permitted -- importing torch acquires no GPU context -- and it is the same +bargain ptychi makes for ``PtychographyTaskOptions``. The cost lands on the +first training call rather than at startup because the import lives inside the +payload builder. Inference needs no config at all, so it never pays it. + +Everything here is picklable: the configs are plain dataclasses of scalars, and +:class:`ReconstructInput` is a frozen dataclass of numpy arrays and a +:class:`Product` (also all numpy), so shipping it verbatim is correct. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from ptychodus.api.reconstructor import ReconstructInput + +if TYPE_CHECKING: + from ptycho_torch.config_params import ( + DataConfig, + DatagenConfig, + InferenceConfig, + ModelConfig, + TrainingConfig, + ) + +__all__ = [ + 'ReconstructPayload', + 'TrainPayload', +] + + +@dataclass(frozen=True) +class ReconstructPayload: + """Everything the child needs to load a checkpoint and run inference once. + + No configs: the child rebuilds its ``ConfigManager`` from the ones saved + inside the checkpoint, so ptychodus settings play no part in inference. + """ + + # Which of ('Unsupervised', 'Supervised') this reconstructor targets. The + # child warns if the checkpoint disagrees. + model_training_mode: str + + # Path to the checkpoint .ckpt file recorded by the parent-side + # SubprocessReconstructor.load_model_from_file() call. Required for + # inference (child raises RuntimeError if None). + model_path: Path | None + + reconstruct_input: ReconstructInput + + +@dataclass(frozen=True) +class TrainPayload: + """Everything the child needs to run one training session.""" + + data_config: DataConfig + model_config: ModelConfig + training_config: TrainingConfig + inference_config: InferenceConfig + datagen_config: DatagenConfig + + input_path: Path + output_path: Path + + # Comma-separated GPU indices, or '' to inherit. The child exports this as + # CUDA_VISIBLE_DEVICES before touching the CUDA runtime, then clamps + # ``training_config.n_devices`` to what actually became visible. + visible_gpu_indices: str diff --git a/src/ptychodus/model/ptychopinn_torch/_subprocess.py b/src/ptychodus/model/ptychopinn_torch/_subprocess.py new file mode 100644 index 000000000..7fd6d7a97 --- /dev/null +++ b/src/ptychodus/model/ptychopinn_torch/_subprocess.py @@ -0,0 +1,277 @@ +"""Child-side subprocess entry points for the PtychoPINN-Torch backend. + +This module runs INSIDE a spawned subprocess. It is the only place in the +ptychodus tree that is allowed to import lightning or the GPU-side +``ptycho_torch`` packages (``api.base_api``, ``model``, ``lightning_utils``). +The parent-side ptychodus process never imports this module -- it reaches only +``ptycho_torch.config_params``, from inside :func:`.reconstructor._build_configs`. + +Two entry points are exposed: + +- :func:`run_reconstruct` -- load a checkpoint, run one inference pass, + stream back a single :class:`ReconstructOutput`. +- :func:`run_train` -- run one Lightning training session (which may itself + fan out to ``n_devices`` DDP ranks via ``strategy='ddp_spawn'``), save the + best checkpoint, and stream back the final :class:`TrainOutput` plus the + saved-checkpoint path. + +Neither entry point reads ptychodus settings. Training receives finished +``ptycho_torch`` config objects on the payload and only assembles them into a +``ConfigManager``; inference rebuilds its configs from the checkpoint. +""" + +from __future__ import annotations + +import logging +import os +import pickle +from collections.abc import Sequence +from multiprocessing.queues import Queue +from pathlib import Path +from typing import Any + +import numpy + +from ptychodus.api.diffraction import zero_bad_pixels +from ptychodus.api.object import Object +from ptychodus.api.product import LossValue, Product +from ptychodus.api.reconstructor import ReconstructOutput, TrainOutput + +from ..processing.subprocess_reconstructor import ( + TAG_MODEL_SAVED, + TAG_OUTPUT, + TAG_TRAIN_OUTPUT, +) +from ._payload import ReconstructPayload, TrainPayload + +logger = logging.getLogger(__name__) + + +def _load_ptycho_model(model_path: Path) -> tuple[Any, Any]: + """Load a PtychoModel from a .ckpt file. Returns (model, config_manager).""" + from ptycho_torch.api.base_api import ConfigManager, PtychoModel + from ptycho_torch.model import PtychoPINN_Lightning + + data_config, model_config, training_config, inference_config = ( + PtychoModel._extract_configs_from_checkpoint(str(model_path)) + ) + if any(c is None for c in (data_config, model_config, training_config, inference_config)): + raise ValueError( + f'Checkpoint at {model_path} is missing one or more saved configs ' + f'(data/model/training/inference).' + ) + + ptycho_model = PtychoModel( + model_config=model_config, + data_config=data_config, + training_config=training_config, + inference_config=inference_config, + ) + ptycho_model.model = PtychoPINN_Lightning.load_from_checkpoint( + model_path, + model_config=model_config, + data_config=data_config, + training_config=training_config, + inference_config=inference_config, + ) + + config_manager = ConfigManager.from_loaded_model(ptycho_model) + config_manager.validate_arch_compatibility(ptycho_model) + return ptycho_model, config_manager + + +def run_reconstruct(payload: ReconstructPayload, queue: Queue[Any]) -> None: + """Child entry point for one inference pass. Streams a single ReconstructOutput.""" + if payload.model_path is None: + raise RuntimeError('Cannot reconstruct: no model checkpoint has been loaded.') + + from ptycho_torch.api.base_api import InferenceEngine, PtychoDataLoader + + ptycho_model, config_manager = _load_ptycho_model(payload.model_path) + + if config_manager.model_config.mode != payload.model_training_mode: + logger.warning( + 'Loaded checkpoint mode %r does not match reconstructor mode %r; ' + 'predictions may be inconsistent.', + config_manager.model_config.mode, + payload.model_training_mode, + ) + + inference_engine = InferenceEngine(config_manager=config_manager, ptycho_model=ptycho_model) + + parameters = payload.reconstruct_input + object_geometry = parameters.product.object_.get_geometry() + positions_px: list[float] = list() + + for position in parameters.product.probe_positions: + object_point = object_geometry.map_coordinates_probe_to_object(position) + positions_px.append(object_point.coordinate_y_px) + positions_px.append(object_point.coordinate_x_px) + + diff_patterns = zero_bad_pixels(parameters.diffraction_patterns, parameters.bad_pixels) + data_loader = PtychoDataLoader.from_np( + diff_patterns=diff_patterns, + probe=parameters.product.probes.get_probe_no_opr().get_array(), + positions=numpy.reshape(positions_px, (-1, 2)), + config_manager=config_manager, + ) + object_out_array = numpy.asarray(inference_engine.predict_and_stitch(data_loader)) + + object_in = parameters.product.object_ + object_out = Object( + array=object_out_array, + layer_spacing_m=object_in.layer_spacing_m, + pixel_geometry=object_in.get_pixel_geometry(), + center=object_in.get_center(), + ) + + losses: Sequence[LossValue] = [] + product = Product( + metadata=parameters.product.metadata, + probe_positions=parameters.product.probe_positions, + probes=parameters.product.probes, + object_=object_out, + losses=losses, + ) + + queue.put((TAG_OUTPUT, pickle.dumps(ReconstructOutput(product=product, progress=1)))) + + +def _apply_visible_devices(visible_gpu_indices: str) -> None: + """Set CUDA_VISIBLE_DEVICES before the first CUDA call in this process. + + Torch has already been imported by the time this runs -- unpickling the + payload pulls in ``ptycho_torch.config_params``, which imports it. That is + fine: importing torch reads no device list. The driver resolves + CUDA_VISIBLE_DEVICES when the runtime is first touched, so masking works as + long as nothing has called ``torch.cuda.*`` yet. Keep this the first + statement of :func:`run_train`. + """ + trimmed = visible_gpu_indices.strip() + if trimmed: + os.environ['CUDA_VISIBLE_DEVICES'] = trimmed + + +def _clamp_n_devices(requested: int) -> int: + """Return min(requested, torch.cuda.device_count()); warn on clamp.""" + import torch + + available = torch.cuda.device_count() + if available == 0: + logger.warning('No CUDA devices available; training will fall back to CPU.') + return 1 + if requested > available: + logger.warning( + 'Requested n_devices=%d but only %d CUDA device(s) available; clamping.', + requested, + available, + ) + return available + return requested + + +def run_train(payload: TrainPayload, queue: Queue[Any]) -> None: + """Child entry point for one training session. + + Masks the visible GPUs, clamps ``n_devices`` to what's actually visible, + then runs the Lightning training loop with the configured DDP strategy. + The configs themselves were built parent-side; all this does with them is + assemble the ``ConfigManager``. + """ + _apply_visible_devices(payload.visible_gpu_indices) + + from lightning.pytorch.callbacks import Callback + from ptycho_torch.api.base_api import ( + ConfigManager, + DataloaderFormats, + PtychoDataLoader, + PtychoModel, + Trainer, + ) + from ptycho_torch.lightning_utils import find_best_checkpoint + from ptycho_torch.model import PtychoPINN_Lightning + + training_config = payload.training_config + training_config.n_devices = _clamp_n_devices(training_config.n_devices) + + config_manager = ConfigManager( + data_config=payload.data_config, + model_config=payload.model_config, + training_config=training_config, + inference_config=payload.inference_config, + datagen_config=payload.datagen_config, + ) + + class _LossCollectorCallback(Callback): + """Collects per-epoch train and validation losses from the Lightning Trainer.""" + + def __init__(self, train_metric_name: str, val_metric_name: str) -> None: + super().__init__() + self._train_metric_name = train_metric_name + self._val_metric_name = val_metric_name + self.training_loss: list[LossValue] = [] + self.validation_loss: list[LossValue] = [] + self.epochs_completed = 0 + + def on_train_epoch_end(self, trainer, pl_module) -> None: # noqa: ANN001 + value = trainer.callback_metrics.get(self._train_metric_name) + if value is not None: + self.training_loss.append( + LossValue(epoch=trainer.current_epoch, value=float(value)) + ) + self.epochs_completed = trainer.current_epoch + 1 + + def on_validation_epoch_end(self, trainer, pl_module) -> None: # noqa: ANN001 + if trainer.sanity_checking: + return + value = trainer.callback_metrics.get(self._val_metric_name) + if value is not None: + self.validation_loss.append( + LossValue(epoch=trainer.current_epoch, value=float(value)) + ) + + data_loader = PtychoDataLoader( + data_dir=payload.input_path, + config_manager=config_manager, + data_format=DataloaderFormats('lightning_only_module'), + output_dir=payload.output_path, + ) + model = PtychoModel._new_model(model=PtychoPINN_Lightning, config_manager=config_manager) + trainer = Trainer._from_lightning( + model=model, + dataloader=data_loader, + orchestration='lightning', + config_manager=config_manager, + ) + + loss_collector = _LossCollectorCallback( + train_metric_name=model.model.loss_name, + val_metric_name=model.model.val_loss_name, + ) + trainer._trainer.callbacks.append(loss_collector) + + trainer.train( + orchestration='lightning', + experiment_name='', + ) + # PtychoDataLoader appends a `run_<timestamp>` segment to output_dir, + # and the checkpoint callback writes there — not at output_path. + run_dir = Path(data_loader.output_dir) + checkpoint_path = find_best_checkpoint(run_dir) + + if checkpoint_path is None: + raise FileNotFoundError(f'No checkpoints found in {run_dir} after training.') + + queue.put( + ( + TAG_TRAIN_OUTPUT, + pickle.dumps( + TrainOutput( + training_loss=loss_collector.training_loss, + validation_loss=loss_collector.validation_loss, + progress=loss_collector.epochs_completed, + ) + ), + ) + ) + queue.put((TAG_MODEL_SAVED, str(checkpoint_path))) diff --git a/src/ptychodus/model/ptychopinn_torch/core.py b/src/ptychodus/model/ptychopinn_torch/core.py index 53150e455..18b11358f 100644 --- a/src/ptychodus/model/ptychopinn_torch/core.py +++ b/src/ptychodus/model/ptychopinn_torch/core.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterator -from importlib.metadata import version +from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec import logging from ...api.reconstructor import ( @@ -21,6 +22,11 @@ logger = logging.getLogger(__name__) +def _ptychopinn_torch_available() -> bool: + """Return True iff ptycho_torch and lightning are importable, without importing them.""" + return all(find_spec(mod) is not None for mod in ('ptycho_torch', 'lightning')) + + class PtychoPINNTorchReconstructorLibrary(ReconstructorLibrary): def __init__( self, settings_registry: SettingsRegistry, is_developer_mode_enabled: bool @@ -33,36 +39,34 @@ def __init__( self.enumerators = PtychoPINNTorchEnumerators() self._reconstructors: list[TrainableReconstructor] = list() - try: - from .reconstructor import PtychoPINNTorchTrainableReconstructor - except ModuleNotFoundError: + if not _ptychopinn_torch_available(): logger.info('PtychoPINN-Torch not found.') if is_developer_mode_enabled: for reconstructor in ('PINN', 'Supervised'): self._reconstructors.append(NullReconstructor(reconstructor)) - else: + return + + # find_spec succeeded above, but the metadata query still hits importlib + # metadata (not the module itself), so no torch/lightning import happens. + try: ptychopinn_torch_version = version('ptychopinn') - logger.info(f'PtychoPINN-Torch {ptychopinn_torch_version}') + except PackageNotFoundError: + ptychopinn_torch_version = 'unknown' + logger.info(f'PtychoPINN-Torch {ptychopinn_torch_version}') + # Import the parent-side factory lazily so that this module's import + # cost stays small even in headless mode. + from .reconstructor import build_reconstructor + + for mode in ('Unsupervised', 'Supervised'): self._reconstructors.append( - PtychoPINNTorchTrainableReconstructor( - 'Unsupervised', - self.data_settings, - self.model_settings, - self.inference_settings, - self.training_settings, - is_developer_mode_enabled=is_developer_mode_enabled, - ) - ) - self._reconstructors.append( - PtychoPINNTorchTrainableReconstructor( - 'Supervised', + build_reconstructor( + mode, self.data_settings, self.model_settings, self.inference_settings, self.training_settings, - is_developer_mode_enabled=is_developer_mode_enabled, ) ) diff --git a/src/ptychodus/model/ptychopinn_torch/reconstructor.py b/src/ptychodus/model/ptychopinn_torch/reconstructor.py index d9155aaef..4a1028e81 100644 --- a/src/ptychodus/model/ptychopinn_torch/reconstructor.py +++ b/src/ptychodus/model/ptychopinn_torch/reconstructor.py @@ -1,43 +1,26 @@ -from __future__ import annotations -from collections.abc import Iterator, Sequence -from pathlib import Path -import logging -import shutil +"""Parent-side factory that builds a :class:`SubprocessReconstructor` for PtychoPINN-Torch. -import numpy +Zero lightning imports, and no GPU context is ever acquired here. All GPU work +runs inside a spawned child; see :mod:`._subprocess` for the child entry +points. -from lightning.pytorch.callbacks import Callback +:func:`_build_configs` does import ``ptycho_torch.config_params`` (and +transitively torch) so the parent can hand the child finished config objects +instead of an INI blob. It is called only from ``build_train_payload``, so that +import happens on the first training run -- never at composition-root time, and +never at all for inference. See :mod:`._payload` for why this is allowed. +""" -from ptycho_torch.api.base_api import ( - ConfigManager, - DataloaderFormats, - InferenceEngine, - PtychoDataLoader, - PtychoModel, - Trainer, -) -from ptycho_torch.config_params import ( - DataConfig, - DatagenConfig, - InferenceConfig, - ModelConfig, - TrainingConfig, -) -from ptycho_torch.lightning_utils import find_best_checkpoint -from ptycho_torch.model import PtychoPINN_Lightning +from __future__ import annotations + +from pathlib import Path +from typing import Any -from ptychodus.api.diffraction import zero_bad_pixels from ptychodus.api.io import save_ptychopinn_training_data -from ptychodus.api.object import Object -from ptychodus.api.product import Product -from ptychodus.api.reconstructor import ( - LossValue, - ReconstructInput, - ReconstructOutput, - TrainOutput, - TrainableReconstructor, -) +from ptychodus.api.reconstructor import ReconstructInput +from ..processing.subprocess_reconstructor import SubprocessReconstructor +from ._payload import ReconstructPayload, TrainPayload from .settings import ( PtychoPINNTorchDataSettings, PtychoPINNTorchInferenceSettings, @@ -46,473 +29,216 @@ ) __all__ = [ - 'PtychoPINNTorchTrainableReconstructor', + 'build_reconstructor', ] -logger = logging.getLogger(__name__) - - -class _LossCollectorCallback(Callback): - """Collects per-epoch train and validation losses from the Lightning Trainer.""" - - def __init__(self, train_metric_name: str, val_metric_name: str) -> None: - super().__init__() - self._train_metric_name = train_metric_name - self._val_metric_name = val_metric_name - self.training_loss: list[LossValue] = [] - self.validation_loss: list[LossValue] = [] - self.epochs_completed = 0 - - def on_train_epoch_end(self, trainer, pl_module) -> None: # noqa: ANN001 - value = trainer.callback_metrics.get(self._train_metric_name) - if value is not None: - self.training_loss.append(LossValue(epoch=trainer.current_epoch, value=float(value))) - self.epochs_completed = trainer.current_epoch + 1 - - def on_validation_epoch_end(self, trainer, pl_module) -> None: # noqa: ANN001 - if trainer.sanity_checking: - return - value = trainer.callback_metrics.get(self._val_metric_name) - if value is not None: - self.validation_loss.append(LossValue(epoch=trainer.current_epoch, value=float(value))) - - -class PtychoPINNTorchTrainableReconstructor(TrainableReconstructor): - def __init__( - self, - model_training_mode: str, - data_settings: PtychoPINNTorchDataSettings, - model_settings: PtychoPINNTorchModelSettings, - inference_settings: PtychoPINNTorchInferenceSettings, - training_settings: PtychoPINNTorchTrainingSettings, - *, - is_developer_mode_enabled: bool, - ) -> None: - super().__init__() - self._model_training_mode = model_training_mode - self._data_settings = data_settings - self._model_settings = model_settings - self._inference_settings = inference_settings - self._training_settings = training_settings - self._is_developer_mode_enabled = is_developer_mode_enabled - - self._inference_engine: InferenceEngine | None = None - self._inference_config_manager: ConfigManager | None = None - self._loaded_from: Path | None = None - - def _create_config_from_settings(self) -> ConfigManager: - grid_size = ( - self._data_settings.grid_size_y.get_value(), - self._data_settings.grid_size_x.get_value(), - ) - x_bounds = ( - self._data_settings.x_lower_bound.get_value(), - self._data_settings.x_upper_bound.get_value(), - ) - y_bounds = ( - self._data_settings.y_lower_bound.get_value(), - self._data_settings.y_upper_bound.get_value(), - ) - data_config = DataConfig( - N=self._data_settings.model_size.get_value(), - C=self._data_settings.num_channels.get_value(), - normalize=self._data_settings.data_normalization_mode.get_value(), - neighbor_function=self._data_settings.neighbor_lookup_method.get_value(), - scan_pattern=self._data_settings.scan_pattern.get_value(), - probe_normalize=self._data_settings.normalize_probe.get_value(), - x_bounds=x_bounds, - y_bounds=y_bounds, - min_neighbor_distance=self._data_settings.min_neighbor_distance.get_value(), - max_neighbor_distance=self._data_settings.max_neighbor_distance.get_value(), - K_quadrant=self._data_settings.num_nearest_neighbors_for_quadrant_lookup.get_value(), - n_subsample=self._data_settings.coordinate_subsampling_factor.get_value(), - probe_scale=self._data_settings.probe_scale.get_value(), - K=self._data_settings.num_nearest_neighbors_for_lookup.get_value(), - grid_size=grid_size, - probe_ramp_removal=self._data_settings.probe_ramp_removal.get_value(), - data_scaling=self._data_settings.data_scaling_method.get_value(), - phase_subtraction=self._data_settings.subtract_mean_phase.get_value(), - ) - - amp_loss = self._model_settings.auxiliary_amplitude_loss.get_value() - phase_loss = self._model_settings.auxiliary_phase_loss.get_value() - - model_config = ModelConfig( - mode=self._model_training_mode, - object_big=self._model_settings.object_big.get_value(), - probe_big=self._model_settings.probe_big.get_value(), - loss_function=self._model_settings.loss_function.get_value(), - amp_activation=self._model_settings.amplitude_activation_function.get_value(), - cbam_encoder=self._model_settings.cbam_encoder.get_value(), - decoder_last_amp_channels=self._data_settings.num_channels.get_value(), - use_shared_decoder=self._model_settings.use_shared_decoder.get_value(), - intensity_scale_trainable=self._model_settings.intensity_scale_trainable.get_value(), - intensity_scale=self._model_settings.intensity_scale.get_value(), - max_position_jitter=self._model_settings.max_position_jitter.get_value(), - num_datasets=self._model_settings.num_datasets.get_value(), - C_model=self._data_settings.num_channels.get_value(), - C_forward=self._data_settings.num_channels.get_value(), - amp_loss=None if amp_loss.casefold() == 'none' else amp_loss, - phase_loss=None if phase_loss.casefold() == 'none' else phase_loss, - amp_loss_coeff=self._model_settings.auxiliary_amplitude_loss_coeff.get_value(), - phase_loss_coeff=self._model_settings.auxiliary_phase_loss_coeff.get_value(), - n_filters_scale=self._model_settings.num_filters_scale.get_value(), - probe_mask=None, - eca_decoder=self._model_settings.eca_decoder.get_value(), - batch_norm=self._model_settings.use_batch_normalization.get_value(), - edge_pad=self._model_settings.edge_pad.get_value(), - decoder_last_c_outer_fraction=self._model_settings.decoder_last_c_outer_fraction.get_value(), - cbam_bottleneck=self._model_settings.cbam_bottleneck.get_value(), - cbam_decoder=self._model_settings.cbam_decoder.get_value(), - spatial_decoder=self._model_settings.spatial_decoder.get_value(), - decoder_spatial_kernel=self._model_settings.decoder_spatial_kernel.get_value(), - eca_encoder=self._model_settings.eca_encoder.get_value(), - offset=self._model_settings.offset.get_value(), - probe_reference_coeff=self._model_settings.probe_reference_loss_coeff.get_value(), - amplitude_variance_loss=self._model_settings.amplitude_variance_loss.get_value(), - amplitude_variance_coeff=self._model_settings.amplitude_variance_coeff.get_value(), - ) - gradient_clip_val = self._training_settings.gradient_clip_val.get_value() - training_config = TrainingConfig( - epochs=self._training_settings.epochs.get_value(), - batch_size=self._training_settings.batch_size.get_value(), - learning_rate=self._training_settings.learning_rate.get_value(), - n_devices=1, # TODO "auto" - num_workers=self._training_settings.num_dataloader_workers.get_value(), - accum_steps=self._training_settings.gradient_accumulation_steps.get_value(), - epochs_fine_tune=self._training_settings.epochs_finetune.get_value(), - fine_tune_gamma=self._training_settings.finetune_gamma.get_value(), - gradient_clip_val=gradient_clip_val if gradient_clip_val > 0.0 else None, - nll=self._training_settings.use_negative_log_likelihood_loss.get_value(), - device=self._training_settings.device.get_value(), - strategy='ddp_spawn', - framework='Lightning', - orchestrator='Lightning', - scheduler=self._training_settings.learning_rate_scheduler.get_value(), - warmup_epochs=self._training_settings.learning_rate_warmup_epochs.get_value(), - min_lr_ratio=self._training_settings.minimum_learning_rate_ratio.get_value(), - notes=self._training_settings.notes.get_value(), - model_name=self._training_settings.model_name.get_value(), - enable_staged_finetuning=self._training_settings.enable_staged_finetuning.get_value(), - finetune_stage1_epochs=self._training_settings.finetune_stage1_epochs.get_value(), - finetune_stage2_epochs=self._training_settings.finetune_stage2_epochs.get_value(), - finetune_stage3_epochs=self._training_settings.finetune_stage3_epochs.get_value(), - finetune_stage1_lr_decoder=self._training_settings.finetune_stage1_lr_decoder.get_value(), - finetune_stage2_lr_encoder_top=self._training_settings.finetune_stage2_lr_encoder_top.get_value(), - finetune_stage2_lr_decoder=self._training_settings.finetune_stage2_lr_decoder.get_value(), - finetune_stage2_lr_phase_head=self._training_settings.finetune_stage2_lr_phase_head.get_value(), - finetune_stage3_lr_encoder_bottom=self._training_settings.finetune_stage3_lr_encoder_bottom.get_value(), - finetune_stage3_lr_encoder_top=self._training_settings.finetune_stage3_lr_encoder_top.get_value(), - finetune_stage3_lr_decoder=self._training_settings.finetune_stage3_lr_decoder.get_value(), - finetune_stage3_lr_phase_head=self._training_settings.finetune_stage3_lr_phase_head.get_value(), - finetune_skip_stage3=self._training_settings.finetune_skip_stage3.get_value(), - finetune_early_stop_patience=self._training_settings.finetune_early_stop_patience.get_value(), - finetune_val_split=self._training_settings.finetune_validation_split.get_value(), - ) - inference_config = InferenceConfig( - batch_size=self._inference_settings.batch_size.get_value(), - middle_trim=self._inference_settings.middle_trim.get_value(), - experiment_number=self._inference_settings.experiment_number.get_value(), - patch_weighting=self._inference_settings.patch_weighting_method.get_value(), - pad_eval=self._inference_settings.pad_eval.get_value(), - window=self._inference_settings.window.get_value(), - ) - datagen_config = DatagenConfig() - return ConfigManager( - data_config=data_config, - model_config=model_config, - training_config=training_config, - inference_config=inference_config, - datagen_config=datagen_config, - ) - - def _sync_config_to_settings(self, config_manager: ConfigManager) -> None: - """Update ptychodus settings to match the configs from the loaded checkpoint. - - Mirrors the field mapping in `_create_config_from_settings` in reverse so - that what the user sees in the UI accurately reflects what the loaded - model was trained with. - """ - data_config = config_manager.data_config - model_config = config_manager.model_config - training_config = config_manager.training_config - inference_config = config_manager.inference_config - - if model_config.mode != self._model_training_mode: - logger.warning( - 'Loaded checkpoint mode %r does not match reconstructor mode %r; ' - 'predictions may be inconsistent.', - model_config.mode, - self._model_training_mode, +_RECONSTRUCT_ENTRY = 'ptychodus.model.ptychopinn_torch._subprocess:run_reconstruct' +_TRAIN_ENTRY = 'ptychodus.model.ptychopinn_torch._subprocess:run_train' + + +def _build_configs( + model_training_mode: str, + data_settings: PtychoPINNTorchDataSettings, + model_settings: PtychoPINNTorchModelSettings, + training_settings: PtychoPINNTorchTrainingSettings, + inference_settings: PtychoPINNTorchInferenceSettings, +) -> tuple[Any, Any, Any, Any, Any]: + """Translate ptychodus settings into the five ptycho_torch config objects. + + Returns ``(data, model, training, inference, datagen)``. The child + assembles them into a ``ConfigManager``; that step needs the GPU-side + ``ptycho_torch.api`` package and stays child-side. + """ + from ptycho_torch.config_params import ( + DataConfig, + DatagenConfig, + InferenceConfig, + ModelConfig, + TrainingConfig, + ) + + grid_size = ( + data_settings.grid_size_y.get_value(), + data_settings.grid_size_x.get_value(), + ) + x_bounds = ( + data_settings.x_lower_bound.get_value(), + data_settings.x_upper_bound.get_value(), + ) + y_bounds = ( + data_settings.y_lower_bound.get_value(), + data_settings.y_upper_bound.get_value(), + ) + data_config = DataConfig( + N=data_settings.model_size.get_value(), + C=data_settings.num_channels.get_value(), + normalize=data_settings.data_normalization_mode.get_value(), + neighbor_function=data_settings.neighbor_lookup_method.get_value(), + scan_pattern=data_settings.scan_pattern.get_value(), + probe_normalize=data_settings.normalize_probe.get_value(), + x_bounds=x_bounds, + y_bounds=y_bounds, + min_neighbor_distance=data_settings.min_neighbor_distance.get_value(), + max_neighbor_distance=data_settings.max_neighbor_distance.get_value(), + K_quadrant=data_settings.num_nearest_neighbors_for_quadrant_lookup.get_value(), + n_subsample=data_settings.coordinate_subsampling_factor.get_value(), + probe_scale=data_settings.probe_scale.get_value(), + K=data_settings.num_nearest_neighbors_for_lookup.get_value(), + grid_size=grid_size, + probe_ramp_removal=data_settings.probe_ramp_removal.get_value(), + data_scaling=data_settings.data_scaling_method.get_value(), + phase_subtraction=data_settings.subtract_mean_phase.get_value(), + ) + + amp_loss = model_settings.auxiliary_amplitude_loss.get_value() + phase_loss = model_settings.auxiliary_phase_loss.get_value() + + model_config = ModelConfig( + mode=model_training_mode, + object_big=model_settings.object_big.get_value(), + probe_big=model_settings.probe_big.get_value(), + loss_function=model_settings.loss_function.get_value(), + amp_activation=model_settings.amplitude_activation_function.get_value(), + cbam_encoder=model_settings.cbam_encoder.get_value(), + decoder_last_amp_channels=data_settings.num_channels.get_value(), + use_shared_decoder=model_settings.use_shared_decoder.get_value(), + intensity_scale_trainable=model_settings.intensity_scale_trainable.get_value(), + intensity_scale=model_settings.intensity_scale.get_value(), + max_position_jitter=model_settings.max_position_jitter.get_value(), + num_datasets=model_settings.num_datasets.get_value(), + C_model=data_settings.num_channels.get_value(), + C_forward=data_settings.num_channels.get_value(), + amp_loss=None if amp_loss.casefold() == 'none' else amp_loss, + phase_loss=None if phase_loss.casefold() == 'none' else phase_loss, + amp_loss_coeff=model_settings.auxiliary_amplitude_loss_coeff.get_value(), + phase_loss_coeff=model_settings.auxiliary_phase_loss_coeff.get_value(), + n_filters_scale=model_settings.num_filters_scale.get_value(), + probe_mask=None, + eca_decoder=model_settings.eca_decoder.get_value(), + batch_norm=model_settings.use_batch_normalization.get_value(), + edge_pad=model_settings.edge_pad.get_value(), + decoder_last_c_outer_fraction=model_settings.decoder_last_c_outer_fraction.get_value(), + cbam_bottleneck=model_settings.cbam_bottleneck.get_value(), + cbam_decoder=model_settings.cbam_decoder.get_value(), + spatial_decoder=model_settings.spatial_decoder.get_value(), + decoder_spatial_kernel=model_settings.decoder_spatial_kernel.get_value(), + eca_encoder=model_settings.eca_encoder.get_value(), + offset=model_settings.offset.get_value(), + probe_reference_coeff=model_settings.probe_reference_loss_coeff.get_value(), + amplitude_variance_loss=model_settings.amplitude_variance_loss.get_value(), + amplitude_variance_coeff=model_settings.amplitude_variance_coeff.get_value(), + ) + + gradient_clip_val = training_settings.gradient_clip_val.get_value() + training_config = TrainingConfig( + epochs=training_settings.epochs.get_value(), + batch_size=training_settings.batch_size.get_value(), + learning_rate=training_settings.learning_rate.get_value(), + n_devices=training_settings.n_devices.get_value(), + num_workers=training_settings.num_dataloader_workers.get_value(), + accum_steps=training_settings.gradient_accumulation_steps.get_value(), + epochs_fine_tune=training_settings.epochs_finetune.get_value(), + fine_tune_gamma=training_settings.finetune_gamma.get_value(), + gradient_clip_val=gradient_clip_val if gradient_clip_val > 0.0 else None, + nll=training_settings.use_negative_log_likelihood_loss.get_value(), + device=training_settings.device.get_value(), + strategy=training_settings.distributed_strategy.get_value(), + framework='Lightning', + orchestrator='Lightning', + scheduler=training_settings.learning_rate_scheduler.get_value(), + warmup_epochs=training_settings.learning_rate_warmup_epochs.get_value(), + min_lr_ratio=training_settings.minimum_learning_rate_ratio.get_value(), + notes=training_settings.notes.get_value(), + model_name=training_settings.model_name.get_value(), + enable_staged_finetuning=training_settings.enable_staged_finetuning.get_value(), + finetune_stage1_epochs=training_settings.finetune_stage1_epochs.get_value(), + finetune_stage2_epochs=training_settings.finetune_stage2_epochs.get_value(), + finetune_stage3_epochs=training_settings.finetune_stage3_epochs.get_value(), + finetune_stage1_lr_decoder=training_settings.finetune_stage1_lr_decoder.get_value(), + finetune_stage2_lr_encoder_top=training_settings.finetune_stage2_lr_encoder_top.get_value(), + finetune_stage2_lr_decoder=training_settings.finetune_stage2_lr_decoder.get_value(), + finetune_stage2_lr_phase_head=training_settings.finetune_stage2_lr_phase_head.get_value(), + finetune_stage3_lr_encoder_bottom=( + training_settings.finetune_stage3_lr_encoder_bottom.get_value() + ), + finetune_stage3_lr_encoder_top=training_settings.finetune_stage3_lr_encoder_top.get_value(), + finetune_stage3_lr_decoder=training_settings.finetune_stage3_lr_decoder.get_value(), + finetune_stage3_lr_phase_head=training_settings.finetune_stage3_lr_phase_head.get_value(), + finetune_skip_stage3=training_settings.finetune_skip_stage3.get_value(), + finetune_early_stop_patience=training_settings.finetune_early_stop_patience.get_value(), + finetune_val_split=training_settings.finetune_validation_split.get_value(), + ) + + inference_config = InferenceConfig( + batch_size=inference_settings.batch_size.get_value(), + middle_trim=inference_settings.middle_trim.get_value(), + experiment_number=inference_settings.experiment_number.get_value(), + patch_weighting=inference_settings.patch_weighting_method.get_value(), + pad_eval=inference_settings.pad_eval.get_value(), + window=inference_settings.window.get_value(), + ) + + return data_config, model_config, training_config, inference_config, DatagenConfig() + + +def build_reconstructor( + model_training_mode: str, + data_settings: PtychoPINNTorchDataSettings, + model_settings: PtychoPINNTorchModelSettings, + inference_settings: PtychoPINNTorchInferenceSettings, + training_settings: PtychoPINNTorchTrainingSettings, +) -> SubprocessReconstructor: + """Build a :class:`SubprocessReconstructor` for one PtychoPINN-Torch mode. + + ``model_training_mode`` is 'Unsupervised' or 'Supervised' and becomes the + reconstructor's display name plus the config-manager mode string used + inside the child. + """ + + def build_reconstruct_payload( + parameters: ReconstructInput, loaded_model_path: Path | None + ) -> ReconstructPayload: + return ReconstructPayload( + model_training_mode=model_training_mode, + model_path=loaded_model_path, + reconstruct_input=parameters, + ) + + def build_train_payload(input_path: Path, output_path: Path) -> TrainPayload: + data_config, model_config, training_config, inference_config, datagen_config = ( + _build_configs( + model_training_mode, + data_settings, + model_settings, + training_settings, + inference_settings, ) - - d = self._data_settings - d.model_size.set_value(data_config.N) - d.num_channels.set_value(data_config.C) - d.data_normalization_mode.set_value(data_config.normalize) - d.neighbor_lookup_method.set_value(data_config.neighbor_function) - d.scan_pattern.set_value(data_config.scan_pattern) - d.normalize_probe.set_value(data_config.probe_normalize) - d.x_lower_bound.set_value(data_config.x_bounds[0]) - d.x_upper_bound.set_value(data_config.x_bounds[1]) - d.y_lower_bound.set_value(data_config.y_bounds[0]) - d.y_upper_bound.set_value(data_config.y_bounds[1]) - d.min_neighbor_distance.set_value(data_config.min_neighbor_distance) - d.max_neighbor_distance.set_value(data_config.max_neighbor_distance) - d.num_nearest_neighbors_for_quadrant_lookup.set_value(data_config.K_quadrant) - d.coordinate_subsampling_factor.set_value(data_config.n_subsample) - d.probe_scale.set_value(data_config.probe_scale) - d.num_nearest_neighbors_for_lookup.set_value(data_config.K) - d.grid_size_y.set_value(data_config.grid_size[0]) - d.grid_size_x.set_value(data_config.grid_size[1]) - d.probe_ramp_removal.set_value(data_config.probe_ramp_removal) - d.data_scaling_method.set_value(data_config.data_scaling) - d.subtract_mean_phase.set_value(data_config.phase_subtraction) - - m = self._model_settings - m.object_big.set_value(model_config.object_big) - m.probe_big.set_value(model_config.probe_big) - m.loss_function.set_value(model_config.loss_function) - m.amplitude_activation_function.set_value(model_config.amp_activation) - m.cbam_encoder.set_value(model_config.cbam_encoder) - m.use_shared_decoder.set_value(model_config.use_shared_decoder) - m.intensity_scale_trainable.set_value(model_config.intensity_scale_trainable) - m.intensity_scale.set_value(model_config.intensity_scale) - m.max_position_jitter.set_value(model_config.max_position_jitter) - m.num_datasets.set_value(model_config.num_datasets) - m.auxiliary_amplitude_loss.set_value( - 'None' if model_config.amp_loss is None else model_config.amp_loss ) - m.auxiliary_phase_loss.set_value( - 'None' if model_config.phase_loss is None else model_config.phase_loss - ) - m.auxiliary_amplitude_loss_coeff.set_value(model_config.amp_loss_coeff) - m.auxiliary_phase_loss_coeff.set_value(model_config.phase_loss_coeff) - m.num_filters_scale.set_value(model_config.n_filters_scale) - m.eca_decoder.set_value(model_config.eca_decoder) - m.use_batch_normalization.set_value(model_config.batch_norm) - m.edge_pad.set_value(model_config.edge_pad) - m.decoder_last_c_outer_fraction.set_value(model_config.decoder_last_c_outer_fraction) - m.cbam_bottleneck.set_value(model_config.cbam_bottleneck) - m.cbam_decoder.set_value(model_config.cbam_decoder) - m.spatial_decoder.set_value(model_config.spatial_decoder) - m.decoder_spatial_kernel.set_value(model_config.decoder_spatial_kernel) - m.eca_encoder.set_value(model_config.eca_encoder) - m.offset.set_value(model_config.offset) - m.probe_reference_loss_coeff.set_value(model_config.probe_reference_coeff) - m.amplitude_variance_loss.set_value(model_config.amplitude_variance_loss) - m.amplitude_variance_coeff.set_value(model_config.amplitude_variance_coeff) - - t = self._training_settings - t.epochs.set_value(training_config.epochs) - t.batch_size.set_value(training_config.batch_size) - t.learning_rate.set_value(training_config.learning_rate) - t.num_dataloader_workers.set_value(training_config.num_workers) - t.gradient_accumulation_steps.set_value(training_config.accum_steps) - t.epochs_finetune.set_value(training_config.epochs_fine_tune) - t.finetune_gamma.set_value(training_config.fine_tune_gamma) - t.gradient_clip_val.set_value( - training_config.gradient_clip_val - if training_config.gradient_clip_val is not None - else 0.0 - ) - t.use_negative_log_likelihood_loss.set_value(training_config.nll) - t.device.set_value(training_config.device) - t.learning_rate_scheduler.set_value(training_config.scheduler) - t.learning_rate_warmup_epochs.set_value(training_config.warmup_epochs) - t.minimum_learning_rate_ratio.set_value(training_config.min_lr_ratio) - t.notes.set_value(training_config.notes) - t.model_name.set_value(training_config.model_name) - t.enable_staged_finetuning.set_value(training_config.enable_staged_finetuning) - t.finetune_stage1_epochs.set_value(training_config.finetune_stage1_epochs) - t.finetune_stage2_epochs.set_value(training_config.finetune_stage2_epochs) - t.finetune_stage3_epochs.set_value(training_config.finetune_stage3_epochs) - t.finetune_stage1_lr_decoder.set_value(training_config.finetune_stage1_lr_decoder) - t.finetune_stage2_lr_encoder_top.set_value(training_config.finetune_stage2_lr_encoder_top) - t.finetune_stage2_lr_decoder.set_value(training_config.finetune_stage2_lr_decoder) - t.finetune_stage2_lr_phase_head.set_value(training_config.finetune_stage2_lr_phase_head) - t.finetune_stage3_lr_encoder_bottom.set_value( - training_config.finetune_stage3_lr_encoder_bottom - ) - t.finetune_stage3_lr_encoder_top.set_value(training_config.finetune_stage3_lr_encoder_top) - t.finetune_stage3_lr_decoder.set_value(training_config.finetune_stage3_lr_decoder) - t.finetune_stage3_lr_phase_head.set_value(training_config.finetune_stage3_lr_phase_head) - t.finetune_skip_stage3.set_value(training_config.finetune_skip_stage3) - t.finetune_early_stop_patience.set_value(training_config.finetune_early_stop_patience) - t.finetune_validation_split.set_value(training_config.finetune_val_split) - - i = self._inference_settings - i.batch_size.set_value(inference_config.batch_size) - i.middle_trim.set_value(inference_config.middle_trim) - i.experiment_number.set_value(inference_config.experiment_number) - i.patch_weighting_method.set_value(inference_config.patch_weighting) - i.pad_eval.set_value(inference_config.pad_eval) - i.window.set_value(inference_config.window) - - @property - def name(self) -> str: - return self._model_training_mode - - def get_progress_goal(self) -> int: - # Reflects the main training fit only; finetuning (epochs_fine_tune, - # staged finetuning) runs on separate Lightning Trainer instances we - # don't observe. - return self._training_settings.epochs.get_value() - - def reconstruct(self, parameters: ReconstructInput) -> Iterator[ReconstructOutput]: - if self._inference_engine is None or self._inference_config_manager is None: - raise RuntimeError('Model must be loaded before reconstruction.') - - object_geometry = parameters.product.object_.get_geometry() - positions_px: list[float] = list() - - for position in parameters.product.probe_positions: - object_point = object_geometry.map_coordinates_probe_to_object(position) - positions_px.append(object_point.coordinate_y_px) - positions_px.append(object_point.coordinate_x_px) - - diff_patterns = zero_bad_pixels(parameters.diffraction_patterns, parameters.bad_pixels) - data_loader = PtychoDataLoader.from_np( - diff_patterns=diff_patterns, - probe=parameters.product.probes.get_probe_no_opr().get_array(), - positions=numpy.reshape(positions_px, (-1, 2)), - config_manager=self._inference_config_manager, - ) - object_out_array = numpy.asarray(self._inference_engine.predict_and_stitch(data_loader)) - - # predict_and_stitch returns a 2D (H, W) array. Object accepts 2- or - # 3-D; pass it through and let Object validate against layer_spacing_m - # rather than blanket-squeezing (which would also collapse legitimate - # singleton spatial axes). - object_in = parameters.product.object_ - object_out = Object( - array=object_out_array, - layer_spacing_m=object_in.layer_spacing_m, - pixel_geometry=object_in.get_pixel_geometry(), - center=object_in.get_center(), - ) - - # TODO: Fourier error - losses: Sequence[LossValue] = [] - product = Product( - metadata=parameters.product.metadata, - probe_positions=parameters.product.probe_positions, - probes=parameters.product.probes, - object_=object_out, - losses=losses, - ) - - yield ReconstructOutput(product) - - def is_model_loaded(self) -> bool: - return self._inference_engine is not None - - def get_model_file_filter(self) -> str: - return 'PyTorch Lightning Checkpoint Files (*.ckpt)' - - def load_model_from_file(self, file_path: Path) -> None: - # The checkpoint is the authority for architecture-critical configs. - # Read the saved hyperparameters directly from the .ckpt rather than - # rebuilding configs from current settings (which may have drifted). - data_config, model_config, training_config, inference_config = ( - PtychoModel._extract_configs_from_checkpoint(str(file_path)) - ) - if any(c is None for c in (data_config, model_config, training_config, inference_config)): - raise ValueError( - f'Checkpoint at {file_path} is missing one or more saved configs ' - f'(data/model/training/inference).' - ) - - ptycho_model = PtychoModel( - model_config=model_config, + return TrainPayload( data_config=data_config, - training_config=training_config, - inference_config=inference_config, - ) - ptycho_model.model = PtychoPINN_Lightning.load_from_checkpoint( - file_path, model_config=model_config, - data_config=data_config, training_config=training_config, inference_config=inference_config, + datagen_config=datagen_config, + input_path=input_path, + output_path=output_path, + visible_gpu_indices=training_settings.visible_gpu_indices.get_value(), ) - # Build the canonical ConfigManager from the loaded model (frozen-mode), - # then audit field-by-field to catch any silent drift. - config_manager = ConfigManager.from_loaded_model(ptycho_model) - config_manager.validate_arch_compatibility(ptycho_model) - - # Mirror the loaded configs back into ptychodus settings so the UI stays - # consistent with what the model was actually trained with. - self._sync_config_to_settings(config_manager) - - self._inference_engine = InferenceEngine( - config_manager=config_manager, ptycho_model=ptycho_model - ) - # Cache the canonical ConfigManager for reuse in reconstruct(); InferenceEngine - # only stores the four config dataclasses, not the manager itself. - self._inference_config_manager = config_manager - self._loaded_from = file_path - - def get_model_file_extension(self) -> str: - return '.ckpt' - - def save_model(self, file_path: Path) -> None: - if self._inference_engine is None: - raise RuntimeError('Cannot save PtychoPINN_Torch model: model is not loaded.') - if self._loaded_from is None or not self._loaded_from.is_file(): - raise RuntimeError( - 'Cannot save PtychoPINN_Torch model: no source checkpoint to copy from. ' - 'Train or load a model first.' - ) - logger.debug(f'Copying loaded checkpoint "{self._loaded_from}" -> "{file_path}"') - shutil.copyfile(self._loaded_from, file_path) - - def get_training_data_file_filter(self) -> str: - return 'NumPy Zipped Archive (*.npz)' - - def export_training_data(self, file_path: Path, parameters: ReconstructInput) -> None: + def export_training_data(file_path: Path, parameters: ReconstructInput) -> None: save_ptychopinn_training_data(file_path, parameters, multimodal_probe=True) - def train(self, input_path: Path, output_path: Path) -> Iterator[TrainOutput]: - config_manager = self._create_config_from_settings() - data_loader = PtychoDataLoader( - data_dir=input_path, - config_manager=config_manager, - data_format=DataloaderFormats('lightning_only_module'), - output_dir=output_path, - ) - model = PtychoModel._new_model(model=PtychoPINN_Lightning, config_manager=config_manager) - trainer = Trainer._from_lightning( - model=model, - dataloader=data_loader, - orchestration='lightning', - config_manager=config_manager, - ) - - loss_collector = _LossCollectorCallback( - train_metric_name=model.model.loss_name, - val_metric_name=model.model.val_loss_name, - ) - trainer._trainer.callbacks.append(loss_collector) - - trainer.train( - orchestration='lightning', - experiment_name='', # unused on the Lightning orchestration path - ) - # PtychoDataLoader appends a `run_<timestamp>` segment to output_dir, - # and the checkpoint callback writes there — not at output_path. - run_dir = Path(data_loader.output_dir) - checkpoint_path = find_best_checkpoint(run_dir) - - if checkpoint_path is None: - raise FileNotFoundError(f'No checkpoints found in {run_dir} after training.') - else: - self.load_model_from_file(checkpoint_path) - - yield TrainOutput( - training_loss=loss_collector.training_loss, - validation_loss=loss_collector.validation_loss, - progress=loss_collector.epochs_completed, - ) + return SubprocessReconstructor( + name=model_training_mode, + reconstruct_entry_point=_RECONSTRUCT_ENTRY, + progress_goal_fn=lambda: training_settings.epochs.get_value(), + build_reconstruct_payload=build_reconstruct_payload, + is_trainable=True, + train_entry_point=_TRAIN_ENTRY, + build_train_payload=build_train_payload, + model_file_filter='PyTorch Lightning Checkpoint Files (*.ckpt)', + model_file_extension='.ckpt', + training_data_file_filter='NumPy Zipped Archive (*.npz)', + export_training_data=export_training_data, + ) diff --git a/src/ptychodus/model/ptychopinn_torch/settings.py b/src/ptychodus/model/ptychopinn_torch/settings.py index 66c6bd931..612729712 100644 --- a/src/ptychodus/model/ptychopinn_torch/settings.py +++ b/src/ptychodus/model/ptychopinn_torch/settings.py @@ -165,6 +165,18 @@ def __init__(self, registry: SettingsRegistry) -> None: 'use_negative_log_likelihood_loss', True ) self.device = self._group.create_string_parameter('device', 'cuda') + # Number of GPU devices to use for training. Lightning's `Trainer` + # spawns one process per device via the configured strategy. + self.n_devices = self._group.create_integer_parameter('n_devices', 1, minimum=1) + # Valid values: 'ddp_spawn' (default, safest inside our subprocess), + # 'ddp' (script re-exec — requires a self-contained entry point, + # not wired in phase 1), 'auto'. + self.distributed_strategy = self._group.create_string_parameter( + 'distributed_strategy', 'ddp_spawn' + ) + # Comma-separated CUDA device indices to expose to the training subprocess + # via CUDA_VISIBLE_DEVICES, or empty to inherit the parent's setting. + self.visible_gpu_indices = self._group.create_string_parameter('visible_gpu_indices', '') self.learning_rate_scheduler = self._group.create_string_parameter( 'learning_rate_scheduler', 'Default' ) diff --git a/src/ptychodus/model/task_monitor.py b/src/ptychodus/model/task_monitor.py index 730518d33..a37434f70 100644 --- a/src/ptychodus/model/task_monitor.py +++ b/src/ptychodus/model/task_monitor.py @@ -57,6 +57,7 @@ def __init__(self, foreground_task_manager: ForegroundTaskManager) -> None: self._is_processing = False self._is_stopping = False self._completed_runs = 0 + self._last_error: BaseException | None = None self._progress_lock = threading.Lock() self._progress_goal = 0 @@ -114,6 +115,18 @@ def wait_for_completion_after(self, snapshot: int, timeout: float | None = None) lambda: self._completed_runs > snapshot, timeout=timeout ) + def get_last_error(self) -> BaseException | None: + """Return the exception captured by the most recent task run, if any.""" + with self._state_condition: + return self._last_error + + def raise_if_failed(self) -> None: + """Re-raise the exception captured by the most recent task run, if any.""" + with self._state_condition: + error = self._last_error + if error is not None: + raise error + def _notify_observers_foreground(self) -> None: task = NotifyObserversTask(self) self._foreground_task_manager.put_foreground_task(task) @@ -122,6 +135,7 @@ def __enter__(self) -> Self: with self._state_condition: self._is_processing = True self._is_stopping = False + self._last_error = None self._state_condition.notify_all() self._notify_observers_foreground() return self @@ -146,5 +160,7 @@ def __exit__( with self._state_condition: self._is_processing = False self._completed_runs += 1 + if exception_value is not None: + self._last_error = exception_value self._state_condition.notify_all() self._notify_observers_foreground() diff --git a/src/ptychodus/model/visualization/color_model.py b/src/ptychodus/model/visualization/color_model.py index 88a60267d..7a09b4a24 100644 --- a/src/ptychodus/model/visualization/color_model.py +++ b/src/ptychodus/model/visualization/color_model.py @@ -1,48 +1,21 @@ -from collections.abc import Iterator - -from ptychodus.api.observer import Observable, Observer from ptychodus.api.parametric import Parameter -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.visualization import CylindricalColorModel -class CylindricalColorModelParameter(Parameter[str], Observer): +class CylindricalColorModelParameter(PluginChooserParameter[CylindricalColorModel]): def __init__(self) -> None: - super().__init__() - self._chooser = PluginChooser[CylindricalColorModel]() + chooser = PluginChooser[CylindricalColorModel]() for model in CylindricalColorModel: - self._chooser.register_plugin( + chooser.register_plugin( model, simple_name=model.simple_name, display_name=model.display_name ) - self._chooser.set_current_plugin('HSV-V') - self._chooser.add_observer(self) - - def choices(self) -> Iterator[str]: - for plugin in self._chooser: - yield plugin.display_name - - def get_value(self) -> str: - return self._chooser.get_current_plugin().display_name - - def set_value(self, value: str, *, notify: bool = True) -> None: - self._chooser.set_current_plugin(value) - - def get_value_as_string(self) -> str: - return self.get_value() - - def set_value_from_string(self, value: str) -> None: - self.set_value(value) + super().__init__(chooser) + self.set_value('HSV-V') def copy(self) -> Parameter[str]: parameter = CylindricalColorModelParameter() parameter.set_value(self.get_value()) return parameter - - def get_strategy(self) -> CylindricalColorModel: - return self._chooser.get_current_plugin().strategy - - def _update(self, observable: Observable) -> None: - if observable is self._chooser: - self.notify_observers() diff --git a/src/ptychodus/model/visualization/colormap.py b/src/ptychodus/model/visualization/colormap.py index c4294dcd0..e7016b0dd 100644 --- a/src/ptychodus/model/visualization/colormap.py +++ b/src/ptychodus/model/visualization/colormap.py @@ -1,10 +1,7 @@ -from collections.abc import Iterator - from matplotlib.colors import Colormap -from ptychodus.api.observer import Observable, Observer from ptychodus.api.parametric import Parameter -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.visualization import ( cyclic_colormap_names, get_colormap_by_name, @@ -12,44 +9,20 @@ ) -class ColormapParameter(Parameter[str], Observer): +class ColormapParameter(PluginChooserParameter[Colormap]): def __init__(self, *, is_cyclic: bool) -> None: - super().__init__() self._is_cyclic = is_cyclic - self._chooser = PluginChooser[Colormap]() + chooser = PluginChooser[Colormap]() cmap_name_it = cyclic_colormap_names() if is_cyclic else linear_colormap_names() for name in sorted(cmap_name_it): cmap = get_colormap_by_name(name) - self._chooser.register_plugin(cmap, display_name=name) + chooser.register_plugin(cmap, display_name=name) + super().__init__(chooser) self.set_value('colorwheel' if is_cyclic else 'gray') - self._chooser.add_observer(self) - - def choices(self) -> Iterator[str]: - for plugin in self._chooser: - yield plugin.display_name - - def get_value(self) -> str: - return self._chooser.get_current_plugin().display_name - - def set_value(self, value: str, *, notify: bool = True) -> None: - self._chooser.set_current_plugin(value) - - def get_value_as_string(self) -> str: - return self.get_value() - - def set_value_from_string(self, value: str) -> None: - self.set_value(value) def copy(self) -> Parameter[str]: parameter = ColormapParameter(is_cyclic=self._is_cyclic) parameter.set_value(self.get_value()) return parameter - - def get_strategy(self) -> Colormap: - return self._chooser.get_current_plugin().strategy - - def _update(self, observable: Observable) -> None: - if observable is self._chooser: - self.notify_observers() diff --git a/src/ptychodus/model/visualization/transformation.py b/src/ptychodus/model/visualization/transformation.py index 1bb763b0c..042c180a4 100644 --- a/src/ptychodus/model/visualization/transformation.py +++ b/src/ptychodus/model/visualization/transformation.py @@ -1,66 +1,40 @@ -from collections.abc import Iterator - -from ptychodus.api.observer import Observable, Observer from ptychodus.api.parametric import Parameter -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter from ptychodus.api.visualization import ScalarTransformation -class ScalarTransformationParameter(Parameter[str], Observer): +class ScalarTransformationParameter(PluginChooserParameter[ScalarTransformation]): def __init__(self) -> None: - super().__init__() - self._chooser = PluginChooser[ScalarTransformation]() - self._chooser.register_plugin( + chooser = PluginChooser[ScalarTransformation]() + chooser.register_plugin( ScalarTransformation.IDENTITY, display_name='Identity', ) - self._chooser.register_plugin( + chooser.register_plugin( ScalarTransformation.SQRT, simple_name='sqrt', display_name='Square Root', ) - self._chooser.register_plugin( + chooser.register_plugin( ScalarTransformation.LOG2, simple_name='log2', display_name='Logarithm (Base 2)', ) - self._chooser.register_plugin( + chooser.register_plugin( ScalarTransformation.LOG, simple_name='ln', display_name='Natural Logarithm', ) - self._chooser.register_plugin( + chooser.register_plugin( ScalarTransformation.LOG10, simple_name='log10', display_name='Logarithm (Base 10)', ) - self.set_value('Identity') - self._chooser.add_observer(self) - - def choices(self) -> Iterator[str]: - for plugin in self._chooser: - yield plugin.display_name - - def get_value(self) -> str: - return self._chooser.get_current_plugin().display_name - - def set_value(self, value: str, *, notify: bool = True) -> None: - self._chooser.set_current_plugin(value) - def get_value_as_string(self) -> str: - return self.get_value() - - def set_value_from_string(self, value: str) -> None: - self.set_value(value) + super().__init__(chooser) + self.set_value('Identity') def copy(self) -> Parameter[str]: parameter = ScalarTransformationParameter() parameter.set_value(self.get_value()) return parameter - - def get_strategy(self) -> ScalarTransformation: - return self._chooser.get_current_plugin().strategy - - def _update(self, observable: Observable) -> None: - if observable is self._chooser: - self.notify_observers() diff --git a/src/ptychodus/model/workflow.py b/src/ptychodus/model/workflow.py index 6fb857fb3..242f12d8b 100644 --- a/src/ptychodus/model/workflow.py +++ b/src/ptychodus/model/workflow.py @@ -4,19 +4,19 @@ from typing import Any import logging -from ptychodus.api.diffraction import CropCenter +from ptychodus.api.diffraction import CropCenter, Polarization from ptychodus.api.geometry import AffineTransform, ImageExtent from ptychodus.api.product import Product from ptychodus.api.reconstructor import AssembledDiffractionData, ReconstructInput from ptychodus.api.settings import PathPrefixChange, SettingsRegistry from ptychodus.api.workflow import ( + DiffractionWorkflowAPI, + ProductWorkflowAPI, RemoteComputeProvider, WorkflowAPI, - WorkflowDiffractionAPI, - WorkflowProductAPI, ) -from .diffraction import DiffractionAPI +from .diffraction import AssembledDiffractionDataset, DiffractionAPI from .fluorescence import FluorescenceAPI from .genesis import GenesisExecutor from .globus import GlobusExecutor @@ -26,18 +26,24 @@ logger = logging.getLogger(__name__) -class ConcreteWorkflowDiffractionAPI(WorkflowDiffractionAPI): - def __init__(self, diffraction_api: DiffractionAPI) -> None: +class ConcreteDiffractionWorkflowAPI(DiffractionWorkflowAPI): + def __init__(self, diffraction_api: DiffractionAPI, dataset_index: int) -> None: self._diffraction_api = diffraction_api + self._dataset_index = dataset_index + + def get_dataset_index(self) -> int: + return self._dataset_index def get_assembled_data(self) -> AssembledDiffractionData: - return self._diffraction_api.get_assembled_data() + return self._diffraction_api.get_assembled_data(self._dataset_index) def save_assembled_data(self, file_path: Path) -> None: - self._diffraction_api.export_assembled_patterns(file_path) + self._diffraction_api.export_assembled_patterns( + file_path, dataset_index=self._dataset_index + ) -class ConcreteWorkflowProductAPI(WorkflowProductAPI): +class ConcreteProductWorkflowAPI(ProductWorkflowAPI): def __init__( self, product_api: ProductAPI, @@ -119,7 +125,7 @@ def generate_object( self._object_api.build_object(self._product_index, generator_name, generator_parameters) def get_reconstruct_input(self) -> ReconstructInput: - return self._processing_api.get_reconstruct_input(self._product_index) + return self._processing_api.get_reconstruct_input(product_index=self._product_index) def reconstruct_local( self, @@ -127,15 +133,14 @@ def reconstruct_local( algorithm: str | None = None, output_product_file: Path | None = None, block: bool = False, - ) -> WorkflowProductAPI: + ) -> ProductWorkflowAPI: output_product_index = self._processing_api.reconstruct( - self._product_index, + product_index=self._product_index, algorithm=algorithm, output_product_file=output_product_file, block=block, ) - - return ConcreteWorkflowProductAPI( + return ConcreteProductWorkflowAPI( self._product_api, self._probe_positions_api, self._probe_api, @@ -171,7 +176,11 @@ def train_reconstructor_local( ) -> None: # TODO mlflow self._processing_api.train( - self._product_index, input_path, output_path, algorithm=algorithm, block=block + input_path, + output_path, + product_index=self._product_index, + algorithm=algorithm, + block=block, ) def train_reconstructor_remote( @@ -191,7 +200,9 @@ def train_reconstructor_remote( def export_training_data(self, file_path: Path, *, algorithm: str | None = None) -> None: self._processing_api.export_training_data( - file_path, self._product_index, algorithm=algorithm + file_path, + product_index=self._product_index, + algorithm=algorithm, ) def save_product(self, file_path: Path, *, file_type: str | None = None) -> None: @@ -207,13 +218,14 @@ def enhance_fluorescence_local( algorithm: str | None = None, block: bool = False, ) -> None: - self._fluorescence_api.enhance_local( - self._product_index, - input_path, - output_path, - input_file_type=input_file_type, - output_file_type=output_file_type, + item_index = self._fluorescence_api.open_measured_dataset( + input_path, self._product_index, file_type=input_file_type + ) + self._fluorescence_api.enhance( + item_index, algorithm=algorithm, + output_file_path=output_path, + output_file_type=output_file_type, block=block, ) @@ -243,8 +255,19 @@ def __init__( self._globus_executor = globus_executor self._genesis_executor = genesis_executor - def load_bad_pixels(self, file_path: Path, *, file_type: str | None = None) -> None: - self._diffraction_api.open_bad_pixels(file_path, file_type=file_type) + def _fetch_dataset( + self, diffraction: DiffractionWorkflowAPI | None + ) -> AssembledDiffractionDataset | None: + if diffraction is None: + return None + + repository = self._diffraction_api.get_repository() + dataset_index = diffraction.get_dataset_index() + + if 0 <= dataset_index < len(repository): + return repository[dataset_index] + + return None def load_diffraction_data( self, @@ -253,33 +276,35 @@ def load_diffraction_data( file_type: str | None = None, crop_center: CropCenter | None = None, crop_extent: ImageExtent | None = None, - detector_extent: ImageExtent | None = None, + bad_pixels_file_path: Path | None = None, + bad_pixels_file_type: str | None = None, process_patterns: bool = True, block: bool = False, - ) -> WorkflowDiffractionAPI: - self._diffraction_api.open_patterns( + ) -> DiffractionWorkflowAPI: + dataset_index = self._diffraction_api.open_patterns( file_path, file_type=file_type, crop_center=crop_center, crop_extent=crop_extent, - detector_extent=detector_extent, + bad_pixels_file_path=bad_pixels_file_path, + bad_pixels_file_type=bad_pixels_file_type, process_patterns=process_patterns, block=block, ) - return ConcreteWorkflowDiffractionAPI(self._diffraction_api) + return ConcreteDiffractionWorkflowAPI(self._diffraction_api, dataset_index) def available_reconstructors(self) -> Iterator[str]: return self._processing_api.available_reconstructors() - def load_assembled_diffraction_data(self, file_path: Path) -> WorkflowDiffractionAPI: - self._diffraction_api.import_assembled_patterns(file_path) - return ConcreteWorkflowDiffractionAPI(self._diffraction_api) + def load_assembled_diffraction_data(self, file_path: Path) -> DiffractionWorkflowAPI: + dataset_index = self._diffraction_api.import_assembled_patterns(file_path) + return ConcreteDiffractionWorkflowAPI(self._diffraction_api, dataset_index) - def get_product(self, product_index: int) -> WorkflowProductAPI: + def get_product(self, product_index: int) -> ProductWorkflowAPI: if product_index < 0: raise ValueError(f'Bad product index ({product_index=})!') - return ConcreteWorkflowProductAPI( + return ConcreteProductWorkflowAPI( self._product_api, self._probe_positions_api, self._probe_api, @@ -291,12 +316,24 @@ def get_product(self, product_index: int) -> WorkflowProductAPI: product_index, ) - def register_product(self, product: Product) -> WorkflowProductAPI: - product_index = self._product_api.insert_product(product) + def register_product( + self, product: Product, *, diffraction: DiffractionWorkflowAPI | None = None + ) -> ProductWorkflowAPI: + product_index = self._product_api.insert_product( + product, dataset=self._fetch_dataset(diffraction) + ) return self.get_product(product_index) - def load_product(self, file_path: Path, *, file_type: str | None = None) -> WorkflowProductAPI: - product_index = self._product_api.open_product(file_path, file_type=file_type) + def load_product( + self, + file_path: Path, + *, + file_type: str | None = None, + diffraction: DiffractionWorkflowAPI | None = None, + ) -> ProductWorkflowAPI: + product_index = self._product_api.open_product( + file_path, file_type=file_type, dataset=self._fetch_dataset(diffraction) + ) if product_index < 0: raise RuntimeError(f'Failed to open product "{file_path}"!') @@ -314,7 +351,10 @@ def create_product( exposure_time_s: float | None = None, mass_attenuation_m2_kg: float | None = None, tomography_angle_deg: float | None = None, - ) -> WorkflowProductAPI: + tilt_angle_deg: float | None = None, + polarization: Polarization | None = None, + diffraction: DiffractionWorkflowAPI | None = None, + ) -> ProductWorkflowAPI: product_index = self._product_api.insert_new_product( name, comments=comments, @@ -324,6 +364,9 @@ def create_product( exposure_time_s=exposure_time_s, mass_attenuation_m2_kg=mass_attenuation_m2_kg, tomography_angle_deg=tomography_angle_deg, + tilt_angle_deg=tilt_angle_deg, + polarization=polarization, + dataset=self._fetch_dataset(diffraction), ) return self.get_product(product_index) diff --git a/src/ptychodus/plugins/aps12id_diffraction_file.py b/src/ptychodus/plugins/aps12id_diffraction_file.py index 90c0f48f0..88fcafb82 100644 --- a/src/ptychodus/plugins/aps12id_diffraction_file.py +++ b/src/ptychodus/plugins/aps12id_diffraction_file.py @@ -62,13 +62,11 @@ def read(self, file_path: Path) -> DiffractionDataset: try: h5_data = h5_file[self.DATA_PATH] - except KeyError: - logger.warning(f'File "{file_path}" is not an APS 12-ID data file.') - return SimpleDiffractionDataset.create_null(file_path) + except KeyError as exc: + raise ValueError(f'File "{file_path}" is not an APS 12-ID data file.') from exc if not isinstance(h5_data, h5py.Dataset): - logger.warning(f'Data path "{self.DATA_PATH}" in "{file_path}" is not a dataset.') - return SimpleDiffractionDataset.create_null(file_path) + raise ValueError(f'Data path "{self.DATA_PATH}" in "{file_path}" is not a dataset.') data_shape = h5_data.shape data_dtype = h5_data.dtype @@ -114,10 +112,9 @@ def read(self, file_path: Path) -> DiffractionDataset: ) return SimpleDiffractionDataset(metadata, contents_tree, array_list) - logger.warning( + raise ValueError( f'Data path "{self.DATA_PATH}" in "{file_path}" has unsupported shape {data_shape}.' ) - return SimpleDiffractionDataset.create_null(file_path) def register_plugins(registry: PluginRegistry) -> None: diff --git a/src/ptychodus/plugins/aps33id_velociprobe/aps33id_velociprobe_diffraction_file.py b/src/ptychodus/plugins/aps33id_velociprobe/aps33id_velociprobe_diffraction_file.py index 46d16fd43..a647ba2cd 100644 --- a/src/ptychodus/plugins/aps33id_velociprobe/aps33id_velociprobe_diffraction_file.py +++ b/src/ptychodus/plugins/aps33id_velociprobe/aps33id_velociprobe_diffraction_file.py @@ -16,7 +16,6 @@ DiffractionFileReader, DiffractionMetadata, DiffractionArray, - SimpleDiffractionDataset, ) from ptychodus.api.tree import SimpleTreeNode @@ -198,8 +197,9 @@ def get_metadata(self) -> DiffractionMetadata: def get_layout(self) -> SimpleTreeNode: return self._contents_tree - def get_bad_pixels(self) -> BadPixels | None: - return None + def get_bad_pixels(self) -> BadPixels: + extent = self._metadata.detector_extent + return numpy.zeros((extent.height_px, extent.width_px), dtype=numpy.bool_) @overload def __getitem__(self, index: int) -> DiffractionArray: ... @@ -221,8 +221,6 @@ def __init__(self) -> None: self.stage_rotation_deg = 0.0 # TODO This is a hack; remove when able! def read(self, file_path: Path) -> DiffractionDataset: - dataset: DiffractionDataset = SimpleDiffractionDataset.create_null(file_path) - with h5py.File(file_path, 'r') as h5_file: h5_dataset = h5_file['/entry/data/data_000001'] num_patterns_per_array = h5_dataset.shape[0] diff --git a/src/ptychodus/plugins/cxi_file.py b/src/ptychodus/plugins/cxi_file.py index 95913d0fb..5e4d2e487 100644 --- a/src/ptychodus/plugins/cxi_file.py +++ b/src/ptychodus/plugins/cxi_file.py @@ -1,120 +1,653 @@ +"""Read and write CXI files (Coherent X-ray Imaging Data Bank format, spec v1.6). + +Writers produce split ``diffraction.cxi`` / ``product.cxi`` files that follow the +CXI spec where a mapping exists, and stash ptychodus-specific fields (OPR weights, +loss history, layer spacing, etc.) under ``/entry_1/ptychodus/``. Readers accept +either split files or a single combined file that carries both raw detector data +and reconstructed image_N groups. +""" + +from datetime import datetime, timezone from pathlib import Path -from typing import Final +from typing import Any, Final import logging import h5py import numpy -from .h5_diffraction_file import H5DiffractionPatternArray, H5DiffractionFileTreeBuilder -from ptychodus.api.geometry import ImageExtent, PixelGeometry +from ptychodus import __version__ as ptychodus_version +from ptychodus.api.common import ELECTRON_VOLT_J from ptychodus.api.diffraction import ( + BadPixels, DiffractionDataset, DiffractionFileReader, + DiffractionFileWriter, DiffractionMetadata, SimpleDiffractionDataset, ) +from ptychodus.api.geometry import ImageExtent, PixelGeometry +from ptychodus.api.object import Object, ObjectCenter from ptychodus.api.plugins import PluginRegistry -from ptychodus.api.probe import ProbeSequence, ProbeFileReader -from ptychodus.api.probe_positions import ( - ProbePositionSequence, - ProbePositionFileReader, - ProbePosition, +from ptychodus.api.probe import ProbeSequence +from ptychodus.api.probe_positions import ProbePosition, ProbePositionSequence +from ptychodus.api.product import ( + LossValue, + Product, + ProductFileReader, + ProductFileWriter, + ProductMetadata, ) -from ptychodus.api.common import ComplexArrayType, ELECTRON_VOLT_J + +from .h5_diffraction_file import H5DiffractionPatternArray, H5DiffractionFileTreeBuilder logger = logging.getLogger(__name__) +CXI_VERSION: Final[int] = 160 # spec version 1.6 stored as int * 100 +CXI_PIXEL_IS_INVALID: Final[int] = 0x00000001 # spec Table 7 / Table 11 + + +class _P: + """CXI HDF5 path constants used by both readers and writers.""" + + ENTRY: Final[str] = '/entry_1' + START_TIME: Final[str] = '/entry_1/start_time' + TITLE: Final[str] = '/entry_1/title' + EXPERIMENT_DESCRIPTION: Final[str] = '/entry_1/experiment_description' + + INSTRUMENT: Final[str] = '/entry_1/instrument_1' + SOURCE: Final[str] = '/entry_1/instrument_1/source_1' + SOURCE_ENERGY: Final[str] = '/entry_1/instrument_1/source_1/energy' + ILLUMINATION: Final[str] = '/entry_1/instrument_1/source_1/illumination' # legacy probe + + DETECTOR: Final[str] = '/entry_1/instrument_1/detector_1' + DETECTOR_DATA: Final[str] = '/entry_1/instrument_1/detector_1/data' + DETECTOR_DISTANCE: Final[str] = '/entry_1/instrument_1/detector_1/distance' + DETECTOR_X_PIXEL_SIZE: Final[str] = '/entry_1/instrument_1/detector_1/x_pixel_size' + DETECTOR_Y_PIXEL_SIZE: Final[str] = '/entry_1/instrument_1/detector_1/y_pixel_size' + DETECTOR_MASK: Final[str] = '/entry_1/instrument_1/detector_1/mask' + + SAMPLE: Final[str] = '/entry_1/sample_1' + SAMPLE_NAME: Final[str] = '/entry_1/sample_1/name' + SAMPLE_GEOMETRY: Final[str] = '/entry_1/sample_1/geometry_1' + TRANSLATION: Final[str] = '/entry_1/sample_1/geometry_1/translation' + + DATA_GROUP: Final[str] = '/entry_1/data_1' + DATA_DATA: Final[str] = '/entry_1/data_1/data' + DATA_TRANSLATION: Final[str] = '/entry_1/data_1/translation' # legacy positions + + PROBE_IMAGE: Final[str] = '/entry_1/image_1' + PROBE_IMAGE_DATA: Final[str] = '/entry_1/image_1/data' + OBJECT_IMAGE: Final[str] = '/entry_1/image_2' + OBJECT_IMAGE_DATA: Final[str] = '/entry_1/image_2/data' + + PROCESS: Final[str] = '/entry_1/process_1' + + # ptychodus extension namespace --------------------------------------------------- + PTYCHODUS: Final[str] = '/entry_1/ptychodus' + PT_PROBE_PHOTON_COUNT: Final[str] = '/entry_1/ptychodus/probe_photon_count' + PT_EXPOSURE_TIME: Final[str] = '/entry_1/ptychodus/exposure_time_s' + PT_MASS_ATTENUATION: Final[str] = '/entry_1/ptychodus/mass_attenuation_m2_kg' + PT_TOMOGRAPHY_ANGLE: Final[str] = '/entry_1/ptychodus/tomography_angle_deg' + PT_POSITION_INDEXES: Final[str] = '/entry_1/ptychodus/probe_position_indexes' + PT_PROBE_PIXEL_WIDTH: Final[str] = '/entry_1/ptychodus/probe_pixel_width_m' + PT_PROBE_PIXEL_HEIGHT: Final[str] = '/entry_1/ptychodus/probe_pixel_height_m' + PT_OPR_WEIGHTS: Final[str] = '/entry_1/ptychodus/opr_weights' + PT_OBJECT_PIXEL_WIDTH: Final[str] = '/entry_1/ptychodus/object_pixel_width_m' + PT_OBJECT_PIXEL_HEIGHT: Final[str] = '/entry_1/ptychodus/object_pixel_height_m' + PT_OBJECT_CENTER_X: Final[str] = '/entry_1/ptychodus/object_center_x_m' + PT_OBJECT_CENTER_Y: Final[str] = '/entry_1/ptychodus/object_center_y_m' + PT_OBJECT_LAYER_SPACING: Final[str] = '/entry_1/ptychodus/object_layer_spacing_m' + PT_LOSS_EPOCHS: Final[str] = '/entry_1/ptychodus/loss_epochs' + PT_LOSS_VALUES: Final[str] = '/entry_1/ptychodus/loss_values' + + +def _iso_now() -> str: + return datetime.now(timezone.utc).replace(microsecond=0).isoformat() + + +def _read_scalar(h5_file: h5py.File, path: str) -> Any | None: + """Return the value at *path*, or None if the path does not exist.""" + if path in h5_file: + return h5_file[path][()] + return None + + +def _read_string(h5_file: h5py.File, path: str) -> str: + """Decode a possibly-bytes string dataset at *path*; return '' if absent.""" + value = _read_scalar(h5_file, path) + if value is None: + return '' + if isinstance(value, bytes): + return value.decode('utf-8', errors='replace') + return str(value) + + class CXIDiffractionFileReader(DiffractionFileReader): + """Read a diffraction dataset from a CXI file.""" + def __init__(self) -> None: - self._data_path = '/entry_1/data_1/data' self._tree_builder = H5DiffractionFileTreeBuilder() def read(self, file_path: Path) -> DiffractionDataset: with h5py.File(file_path, 'r') as h5_file: contents_tree = self._tree_builder.build(h5_file) - data = h5_file[self._data_path] - - if isinstance(data, h5py.Dataset): - num_patterns, detector_height, detector_width = data.shape + # Canonical data location (soft link in files we write); fall back to + # the detector data path if the data_1 group is absent. + data_path = _P.DATA_DATA if _P.DATA_DATA in h5_file else _P.DETECTOR_DATA + data = h5_file.get(data_path) - detector_extent = ImageExtent(detector_width, detector_height) - detector_distance_m = float( - h5_file['/entry_1/instrument_1/detector_1/distance'][()] - ) - detector_pixel_geometry = PixelGeometry( - float(h5_file['/entry_1/instrument_1/detector_1/x_pixel_size'][()]), - float(h5_file['/entry_1/instrument_1/detector_1/y_pixel_size'][()]), - ) - probe_energy_J = float(h5_file['/entry_1/instrument_1/source_1/energy'][()]) # noqa: N806 - probe_energy_eV = probe_energy_J / ELECTRON_VOLT_J # noqa: N806 - - # TODO load detector mask; zeros are good pixels - # /entry_1/instrument_1/detector_1/mask Dataset {512, 512} - - metadata = DiffractionMetadata( - num_patterns_per_array=[num_patterns], - pattern_dtype=data.dtype, - detector_distance_m=detector_distance_m, - detector_extent=detector_extent, - detector_pixel_geometry=detector_pixel_geometry, - probe_energy_eV=probe_energy_eV, - file_path=file_path, + if not isinstance(data, h5py.Dataset): + raise ValueError( + f'Expected diffraction pattern dataset at "{data_path}" in {file_path};' + f' got {type(data).__name__}.' ) - array = H5DiffractionPatternArray( - label=file_path.stem, - indexes=numpy.arange(num_patterns), - file_path=file_path, - data_path=self._data_path, + if data.ndim != 3: + raise ValueError( + f'Diffraction pattern dataset at "{data_path}" must be 3D (N,H,W);' + f' got shape {data.shape}.' ) - return SimpleDiffractionDataset(metadata, contents_tree, [array]) - else: - raise ValueError(f'Expected dataset at {self._data_path}, got {type(data)}.') + num_patterns, detector_height, detector_width = data.shape + detector_extent = ImageExtent(detector_width, detector_height) + detector_distance_m: float | None = None + distance = _read_scalar(h5_file, _P.DETECTOR_DISTANCE) + if distance is not None: + detector_distance_m = float(distance) -class CXIPositionFileReader(ProbePositionFileReader): - def read(self, file_path: Path) -> ProbePositionSequence: - point_list: list[ProbePosition] = list() + detector_pixel_geometry: PixelGeometry | None = None + x_pixel = _read_scalar(h5_file, _P.DETECTOR_X_PIXEL_SIZE) + y_pixel = _read_scalar(h5_file, _P.DETECTOR_Y_PIXEL_SIZE) + if x_pixel is not None and y_pixel is not None: + detector_pixel_geometry = PixelGeometry(float(x_pixel), float(y_pixel)) - with h5py.File(file_path, 'r') as h5_file: - xyz_m = h5_file['/entry_1/data_1/translation'][()] + probe_energy_eV: float | None = None # noqa: N806 + energy_J = _read_scalar(h5_file, _P.SOURCE_ENERGY) # noqa: N806 + if energy_J is not None: + probe_energy_eV = float(energy_J) / ELECTRON_VOLT_J # noqa: N806 + + probe_photon_count: int | None = None + photon_count = _read_scalar(h5_file, _P.PT_PROBE_PHOTON_COUNT) + if photon_count is not None: + probe_photon_count = int(photon_count) + + exposure_time_s: float | None = None + exposure = _read_scalar(h5_file, _P.PT_EXPOSURE_TIME) + if exposure is not None: + exposure_time_s = float(exposure) + + tomography_angle_deg: float | None = None + angle = _read_scalar(h5_file, _P.PT_TOMOGRAPHY_ANGLE) + if angle is not None: + tomography_angle_deg = float(angle) + + bad_pixels: BadPixels | None = None + if _P.DETECTOR_MASK in h5_file: + mask = h5_file[_P.DETECTOR_MASK][()] + bad_pixels = numpy.asarray(mask != 0, dtype=bool) + + metadata = DiffractionMetadata( + num_patterns_per_array=[num_patterns], + pattern_dtype=data.dtype, + detector_distance_m=detector_distance_m, + detector_extent=detector_extent, + detector_pixel_geometry=detector_pixel_geometry, + probe_energy_eV=probe_energy_eV, + probe_photon_count=probe_photon_count, + exposure_time_s=exposure_time_s, + tomography_angle_deg=tomography_angle_deg, + file_path=file_path, + ) + + array = H5DiffractionPatternArray( + label=file_path.stem, + indexes=numpy.arange(num_patterns), + file_path=file_path, + data_path=data_path, + ) + + return SimpleDiffractionDataset(metadata, contents_tree, [array], bad_pixels) + + +class CXIDiffractionFileWriter(DiffractionFileWriter): + """Write a diffraction dataset to a CXI file.""" + + def write(self, file_path: Path, dataset: DiffractionDataset) -> None: + patterns = numpy.concatenate([array.get_patterns() for array in dataset]) + metadata = dataset.get_metadata() + bad_pixels = dataset.get_bad_pixels() - for idx, (x, y, z) in enumerate(xyz_m): - point = ProbePosition(idx, x, y) - point_list.append(point) + with h5py.File(file_path, 'w') as h5_file: + h5_file.create_dataset('cxi_version', data=CXI_VERSION) - return ProbePositionSequence(point_list) + entry = h5_file.create_group(_P.ENTRY) + entry.create_dataset('start_time', data=_iso_now()) + instrument = h5_file.create_group(_P.INSTRUMENT) + source = instrument.create_group('source_1') + detector = instrument.create_group('detector_1') -class CXIProbeFileReader(ProbeFileReader): - def read(self, file_path: Path) -> ProbeSequence: - array: ComplexArrayType | None = None + if metadata.probe_energy_eV is not None: + source.create_dataset('energy', data=metadata.probe_energy_eV * ELECTRON_VOLT_J) + if metadata.detector_distance_m is not None: + detector.create_dataset('distance', data=metadata.detector_distance_m) + + pixel_geometry = metadata.detector_pixel_geometry + if pixel_geometry is not None: + detector.create_dataset('x_pixel_size', data=pixel_geometry.width_m) + detector.create_dataset('y_pixel_size', data=pixel_geometry.height_m) + + detector.create_dataset('data', data=patterns, compression='lzf') + + if bad_pixels is not None: + mask = numpy.where(bad_pixels, CXI_PIXEL_IS_INVALID, 0).astype(numpy.uint32) + detector.create_dataset('mask', data=mask, compression='lzf') + + data_group = h5_file.create_group(_P.DATA_GROUP) + data_group['data'] = h5py.SoftLink(_P.DETECTOR_DATA) + + pt_extras: dict[str, Any] = {} + if metadata.probe_photon_count is not None: + pt_extras[_P.PT_PROBE_PHOTON_COUNT] = int(metadata.probe_photon_count) + if metadata.exposure_time_s is not None: + pt_extras[_P.PT_EXPOSURE_TIME] = float(metadata.exposure_time_s) + if metadata.tomography_angle_deg is not None: + pt_extras[_P.PT_TOMOGRAPHY_ANGLE] = float(metadata.tomography_angle_deg) + + if pt_extras: + h5_file.create_group(_P.PTYCHODUS) + for path, value in pt_extras.items(): + h5_file.create_dataset(path, data=value) + + +class CXIProductFileIO(ProductFileReader, ProductFileWriter): + """Read and write ptychodus data products in CXI format.""" + + SIMPLE_NAME: Final[str] = 'CXI' + DISPLAY_NAME: Final[str] = 'Coherent X-ray Imaging Files (*.cxi)' + + def read(self, file_path: Path) -> Product: with h5py.File(file_path, 'r') as h5_file: - array = h5_file['/entry_1/instrument_1/source_1/illumination'][()] + metadata = self._read_metadata(h5_file) + probe = self._read_probe(h5_file) + object_ = self._read_object(h5_file) + positions = self._read_positions(h5_file) + losses = self._read_losses(h5_file) + + return Product( + metadata=metadata, + probe_positions=positions, + probes=probe, + object_=object_, + losses=losses, + ) + + def write(self, file_path: Path, product: Product) -> None: + metadata = product.metadata + probe = product.probes + object_ = product.object_ + object_geometry = object_.get_geometry() + probe_pixel_geometry = probe.get_pixel_geometry() + + with h5py.File(file_path, 'w') as h5_file: + h5_file.create_dataset('cxi_version', data=CXI_VERSION) - return ProbeSequence(array=array, opr_weights=None, pixel_geometry=None) + entry = h5_file.create_group(_P.ENTRY) + entry.create_dataset('start_time', data=_iso_now()) + + if metadata.name: + entry.create_dataset('title', data=metadata.name) + if metadata.comments: + entry.create_dataset('experiment_description', data=metadata.comments) + + sample = h5_file.create_group(_P.SAMPLE) + if metadata.name: + sample.create_dataset('name', data=metadata.name) + + positions_xyz = numpy.zeros((len(product.probe_positions), 3), dtype=numpy.float64) + position_indexes = numpy.empty(len(product.probe_positions), dtype=numpy.int64) + for i, point in enumerate(product.probe_positions): + position_indexes[i] = point.index + positions_xyz[i, 0] = point.coordinate_x_m + positions_xyz[i, 1] = point.coordinate_y_m + + sample_geometry = sample.create_group('geometry_1') + sample_geometry.create_dataset('translation', data=positions_xyz) + + data_group = h5_file.create_group(_P.DATA_GROUP) + data_group['translation'] = h5py.SoftLink(_P.TRANSLATION) + + instrument = h5_file.create_group(_P.INSTRUMENT) + source = instrument.create_group('source_1') + source.create_dataset('energy', data=metadata.probe_energy_J) + detector = instrument.create_group('detector_1') + detector.create_dataset('distance', data=metadata.detector_distance_m) + + self._write_probe_image(h5_file, probe, probe_pixel_geometry) + source['illumination'] = h5py.SoftLink(_P.PROBE_IMAGE_DATA) + + self._write_object_image(h5_file, object_, object_geometry) + + process = h5_file.create_group(_P.PROCESS) + process.create_dataset('program', data='Ptychodus') + process.create_dataset('version', data=ptychodus_version) + process.create_dataset('date', data=_iso_now()) + + self._write_ptychodus_extras( + h5_file, + metadata, + probe, + probe_pixel_geometry, + object_, + object_geometry, + position_indexes, + product.losses, + ) + + # -- read helpers --------------------------------------------------------------- + + def _read_metadata(self, h5_file: h5py.File) -> ProductMetadata: + name = _read_string(h5_file, _P.TITLE) or _read_string(h5_file, _P.SAMPLE_NAME) + comments = _read_string(h5_file, _P.EXPERIMENT_DESCRIPTION) + + detector_distance_m = 0.0 + distance = _read_scalar(h5_file, _P.DETECTOR_DISTANCE) + if distance is not None: + detector_distance_m = float(distance) + + probe_energy_eV = 0.0 # noqa: N806 + energy_J = _read_scalar(h5_file, _P.SOURCE_ENERGY) # noqa: N806 + if energy_J is not None: + probe_energy_eV = float(energy_J) / ELECTRON_VOLT_J # noqa: N806 + + probe_photon_count = 0.0 + photon_count = _read_scalar(h5_file, _P.PT_PROBE_PHOTON_COUNT) + if photon_count is not None: + probe_photon_count = float(photon_count) + + exposure_time_s = 0.0 + exposure = _read_scalar(h5_file, _P.PT_EXPOSURE_TIME) + if exposure is not None: + exposure_time_s = float(exposure) + + mass_attenuation_m2_kg = 0.0 + attenuation = _read_scalar(h5_file, _P.PT_MASS_ATTENUATION) + if attenuation is not None: + mass_attenuation_m2_kg = float(attenuation) + + tomography_angle_deg = 0.0 + angle = _read_scalar(h5_file, _P.PT_TOMOGRAPHY_ANGLE) + if angle is not None: + tomography_angle_deg = float(angle) + + return ProductMetadata( + name=name, + comments=comments, + detector_distance_m=detector_distance_m, + probe_energy_eV=probe_energy_eV, + probe_photon_count=probe_photon_count, + exposure_time_s=exposure_time_s, + mass_attenuation_m2_kg=mass_attenuation_m2_kg, + tomography_angle_deg=tomography_angle_deg, + ) + + def _read_probe(self, h5_file: h5py.File) -> ProbeSequence: + if _P.PROBE_IMAGE_DATA in h5_file: + probe_path = _P.PROBE_IMAGE_DATA + image_group_path = _P.PROBE_IMAGE + elif _P.ILLUMINATION in h5_file: + probe_path = _P.ILLUMINATION + image_group_path = '' + else: + raise ValueError( + f'CXI file has no probe: expected {_P.PROBE_IMAGE_DATA} or {_P.ILLUMINATION}.' + ) + + array = h5_file[probe_path][()] + + pixel_geometry: PixelGeometry | None = None + pw = _read_scalar(h5_file, _P.PT_PROBE_PIXEL_WIDTH) + ph = _read_scalar(h5_file, _P.PT_PROBE_PIXEL_HEIGHT) + if pw is not None and ph is not None: + pixel_geometry = PixelGeometry(width_m=float(pw), height_m=float(ph)) + elif image_group_path: + # Fall back to computing from image_size / array shape. + image_size_path = f'{image_group_path}/image_size' + if image_size_path in h5_file: + image_size = numpy.asarray(h5_file[image_size_path][()], dtype=float) + if image_size.size >= 2 and array.ndim >= 2: + height_px, width_px = array.shape[-2], array.shape[-1] + if height_px > 0 and width_px > 0: + pixel_geometry = PixelGeometry( + width_m=float(image_size[0]) / width_px, + height_m=float(image_size[1]) / height_px, + ) + + opr_weights = None + if _P.PT_OPR_WEIGHTS in h5_file: + opr_weights = h5_file[_P.PT_OPR_WEIGHTS][()] + + return ProbeSequence(array=array, opr_weights=opr_weights, pixel_geometry=pixel_geometry) + + def _read_object(self, h5_file: h5py.File) -> Object: + if _P.OBJECT_IMAGE_DATA not in h5_file: + raise ValueError(f'CXI file has no object at {_P.OBJECT_IMAGE_DATA}.') + + array = h5_file[_P.OBJECT_IMAGE_DATA][()] + + pixel_geometry: PixelGeometry | None = None + pw = _read_scalar(h5_file, _P.PT_OBJECT_PIXEL_WIDTH) + ph = _read_scalar(h5_file, _P.PT_OBJECT_PIXEL_HEIGHT) + if pw is not None and ph is not None: + pixel_geometry = PixelGeometry(width_m=float(pw), height_m=float(ph)) + else: + image_size_path = f'{_P.OBJECT_IMAGE}/image_size' + if image_size_path in h5_file: + image_size = numpy.asarray(h5_file[image_size_path][()], dtype=float) + if image_size.size >= 2 and array.ndim >= 2: + height_px, width_px = array.shape[-2], array.shape[-1] + if height_px > 0 and width_px > 0: + pixel_geometry = PixelGeometry( + width_m=float(image_size[0]) / width_px, + height_m=float(image_size[1]) / height_px, + ) + + center: ObjectCenter | None = None + cx = _read_scalar(h5_file, _P.PT_OBJECT_CENTER_X) + cy = _read_scalar(h5_file, _P.PT_OBJECT_CENTER_Y) + if cx is not None and cy is not None: + center = ObjectCenter(coordinate_x_m=float(cx), coordinate_y_m=float(cy)) + + layer_spacing_m: list[float] = [] + if _P.PT_OBJECT_LAYER_SPACING in h5_file: + spacing = h5_file[_P.PT_OBJECT_LAYER_SPACING][()] + layer_spacing_m = [float(v) for v in numpy.atleast_1d(spacing)] + + return Object( + array=array, + pixel_geometry=pixel_geometry, + center=center, + layer_spacing_m=layer_spacing_m, + ) + + def _read_positions(self, h5_file: h5py.File) -> ProbePositionSequence: + if _P.TRANSLATION in h5_file: + translation_path = _P.TRANSLATION + elif _P.DATA_TRANSLATION in h5_file: + translation_path = _P.DATA_TRANSLATION + else: + logger.warning( + 'CXI file has no probe positions at %s or %s; returning empty sequence.', + _P.TRANSLATION, + _P.DATA_TRANSLATION, + ) + return ProbePositionSequence() + + translation = numpy.asarray(h5_file[translation_path][()]) + if translation.ndim != 2 or translation.shape[1] < 2: + raise ValueError( + f'Expected translation dataset at "{translation_path}" to have shape (N, >=2);' + f' got {translation.shape}.' + ) + + num_positions = translation.shape[0] + + if _P.PT_POSITION_INDEXES in h5_file: + indexes = numpy.asarray(h5_file[_P.PT_POSITION_INDEXES][()], dtype=int) + else: + indexes = numpy.arange(num_positions, dtype=int) + + if len(indexes) != num_positions: + raise ValueError( + f'Position index count {len(indexes)} does not match' + f' translation count {num_positions}.' + ) + + points = [ + ProbePosition(int(idx), float(translation[i, 0]), float(translation[i, 1])) + for i, idx in enumerate(indexes) + ] + return ProbePositionSequence(points) + + def _read_losses(self, h5_file: h5py.File) -> list[LossValue]: + if _P.PT_LOSS_VALUES not in h5_file: + return [] + + values = numpy.atleast_1d(h5_file[_P.PT_LOSS_VALUES][()]) + + if _P.PT_LOSS_EPOCHS in h5_file: + epochs = numpy.atleast_1d(h5_file[_P.PT_LOSS_EPOCHS][()]) + else: + epochs = numpy.arange(len(values)) + + return [LossValue(int(epoch), float(value)) for epoch, value in zip(epochs, values)] + + # -- write helpers -------------------------------------------------------------- + + def _write_probe_image( + self, h5_file: h5py.File, probe: ProbeSequence, pixel_geometry: PixelGeometry + ) -> None: + image = h5_file.create_group(_P.PROBE_IMAGE) + data = image.create_dataset('data', data=probe.get_array()) + data.attrs['axes'] = 'coherent:incoherent:y:x' + image.create_dataset('data_type', data='amplitude') + image.create_dataset('data_space', data='real') + image.create_dataset( + 'image_size', + data=numpy.array( + [ + pixel_geometry.width_m * probe.width_px, + pixel_geometry.height_m * probe.height_px, + 0.0, + ], + dtype=numpy.float64, + ), + ) + image['source_1'] = h5py.SoftLink(_P.SOURCE) + image['detector_1'] = h5py.SoftLink(_P.DETECTOR) + + def _write_object_image(self, h5_file: h5py.File, object_: Object, geometry: Any) -> None: + image = h5_file.create_group(_P.OBJECT_IMAGE) + data = image.create_dataset('data', data=object_.get_array()) + data.attrs['axes'] = 'layer:y:x' + image.create_dataset('data_type', data='electron density') + image.create_dataset('data_space', data='real') + image.create_dataset( + 'image_size', + data=numpy.array( + [ + geometry.pixel_width_m * geometry.width_px, + geometry.pixel_height_m * geometry.height_px, + 0.0, + ], + dtype=numpy.float64, + ), + ) + image['source_1'] = h5py.SoftLink(_P.SOURCE) + image['detector_1'] = h5py.SoftLink(_P.DETECTOR) + + def _write_ptychodus_extras( + self, + h5_file: h5py.File, + metadata: ProductMetadata, + probe: ProbeSequence, + probe_pixel_geometry: PixelGeometry, + object_: Object, + object_geometry: Any, + position_indexes: numpy.ndarray, + losses: Any, + ) -> None: + h5_file.create_group(_P.PTYCHODUS) + + h5_file.create_dataset(_P.PT_POSITION_INDEXES, data=position_indexes) + h5_file.create_dataset(_P.PT_PROBE_PIXEL_WIDTH, data=probe_pixel_geometry.width_m) + h5_file.create_dataset(_P.PT_PROBE_PIXEL_HEIGHT, data=probe_pixel_geometry.height_m) + + try: + opr_weights = probe.get_opr_weights() + except ValueError: + pass + else: + h5_file.create_dataset(_P.PT_OPR_WEIGHTS, data=opr_weights) + + h5_file.create_dataset(_P.PT_OBJECT_PIXEL_WIDTH, data=object_geometry.pixel_width_m) + h5_file.create_dataset(_P.PT_OBJECT_PIXEL_HEIGHT, data=object_geometry.pixel_height_m) + h5_file.create_dataset(_P.PT_OBJECT_CENTER_X, data=object_geometry.center_x_m) + h5_file.create_dataset(_P.PT_OBJECT_CENTER_Y, data=object_geometry.center_y_m) + + layer_spacing = list(object_.layer_spacing_m) + if layer_spacing: + h5_file.create_dataset( + _P.PT_OBJECT_LAYER_SPACING, data=numpy.asarray(layer_spacing, dtype=numpy.float64) + ) + + if metadata.probe_photon_count: + h5_file.create_dataset(_P.PT_PROBE_PHOTON_COUNT, data=metadata.probe_photon_count) + if metadata.exposure_time_s: + h5_file.create_dataset(_P.PT_EXPOSURE_TIME, data=metadata.exposure_time_s) + if metadata.mass_attenuation_m2_kg: + h5_file.create_dataset(_P.PT_MASS_ATTENUATION, data=metadata.mass_attenuation_m2_kg) + if metadata.tomography_angle_deg: + h5_file.create_dataset(_P.PT_TOMOGRAPHY_ANGLE, data=metadata.tomography_angle_deg) + + loss_epochs = [loss.epoch for loss in losses] + loss_values = [loss.value for loss in losses] + if loss_epochs: + h5_file.create_dataset( + _P.PT_LOSS_EPOCHS, data=numpy.asarray(loss_epochs, dtype=numpy.int64) + ) + h5_file.create_dataset( + _P.PT_LOSS_VALUES, data=numpy.asarray(loss_values, dtype=numpy.float64) + ) def register_plugins(registry: PluginRegistry) -> None: - SIMPLE_NAME: Final[str] = 'CXI' # noqa: N806 - DISPLAY_NAME: Final[str] = 'Coherent X-ray Imaging Files (*.cxi)' # noqa: N806 + display_name: Final[str] = CXIProductFileIO.DISPLAY_NAME + simple_name: Final[str] = CXIProductFileIO.SIMPLE_NAME registry.diffraction_file_readers.register_plugin( CXIDiffractionFileReader(), - simple_name=SIMPLE_NAME, - display_name=DISPLAY_NAME, + simple_name=simple_name, + display_name=display_name, + ) + registry.diffraction_file_writers.register_plugin( + CXIDiffractionFileWriter(), + simple_name=simple_name, + display_name=display_name, ) - registry.probe_position_file_readers.register_plugin( - CXIPositionFileReader(), - simple_name=SIMPLE_NAME, - display_name=DISPLAY_NAME, + + cxi_product_io = CXIProductFileIO() + registry.register_product_file_reader_with_adapters( + cxi_product_io, + simple_name=simple_name, + display_name=display_name, ) - registry.probe_file_readers.register_plugin( - CXIProbeFileReader(), - simple_name=SIMPLE_NAME, - display_name=DISPLAY_NAME, + registry.product_file_writers.register_plugin( + cxi_product_io, + simple_name=simple_name, + display_name=display_name, ) diff --git a/src/ptychodus/plugins/npz_diffraction_file.py b/src/ptychodus/plugins/npz_diffraction_file.py index 018b9dc57..bfe8b2d0a 100644 --- a/src/ptychodus/plugins/npz_diffraction_file.py +++ b/src/ptychodus/plugins/npz_diffraction_file.py @@ -28,14 +28,12 @@ class NPZDiffractionFileIO(DiffractionFileReader, DiffractionFileWriter): BAD_PIXELS: Final[str] = 'bad_pixels' def read(self, file_path: Path) -> DiffractionDataset: - dataset = SimpleDiffractionDataset.create_null(file_path) contents = numpy.load(file_path) try: patterns = contents[self.PATTERNS] - except KeyError: - logger.warning(f'Failed to read patterns in "{file_path}".') - return dataset + except KeyError as exc: + raise ValueError(f'Failed to read patterns in "{file_path}".') from exc num_patterns, detector_height, detector_width = patterns.shape diff --git a/src/ptychodus/plugins/tiff_diffraction_file.py b/src/ptychodus/plugins/tiff_diffraction_file.py index 53f8e9b02..af7d31f56 100644 --- a/src/ptychodus/plugins/tiff_diffraction_file.py +++ b/src/ptychodus/plugins/tiff_diffraction_file.py @@ -57,7 +57,7 @@ def _get_file_series(self, file_path: Path) -> tuple[Mapping[int, Path], str]: z = re.match(file_pattern, fp.name) if z: - index = int(z.group(1).lstrip('0')) + index = int(z.group(1)) file_path_dict[index] = fp return file_path_dict, file_pattern diff --git a/src/ptychodus/plugins/workflow.py b/src/ptychodus/plugins/workflow.py index 80ca33dc5..e69e8f929 100644 --- a/src/ptychodus/plugins/workflow.py +++ b/src/ptychodus/plugins/workflow.py @@ -35,8 +35,8 @@ def execute(self, api: WorkflowAPI, file_path: Path) -> None: scan_id = int(re.findall(r'\d+', scan_name)[-1]) diffraction_file_path = file_path.parents[1] / 'raw_data' / f'scan{scan_id}_master.h5' - api.load_diffraction_data(diffraction_file_path) - product_api = api.create_product(f'scan{scan_id}') + diffraction_api = api.load_diffraction_data(diffraction_file_path) + product_api = api.create_product(f'scan{scan_id}', diffraction=diffraction_api) product_api.load_probe_positions(file_path) product_api.generate_probe() product_api.generate_object() @@ -61,8 +61,8 @@ def execute(self, api: WorkflowAPI, file_path: Path) -> None: digits = int(re.findall(r'\d+', diffraction_file_path.stem)[-1]) if digits == 0: - api.load_diffraction_data(diffraction_file_path) - product_api = api.create_product(f'scan_{scan_id}') + diffraction_api = api.load_diffraction_data(diffraction_file_path) + product_api = api.create_product(f'scan_{scan_id}', diffraction=diffraction_api) product_api.load_probe_positions(file_path) product_api.generate_probe() product_api.generate_object() @@ -142,8 +142,10 @@ def execute(self, api: WorkflowAPI, file_path: Path) -> None: logger.warning(f'Failed to locate metadata for {scan_num}!') else: product_name = f'scan{scan_num:05d}_' + metadata.label - api.load_diffraction_data(file_path) - input_product_api = api.create_product(product_name, comments=str(metadata)) + diffraction_api = api.load_diffraction_data(file_path) + input_product_api = api.create_product( + product_name, comments=str(metadata), diffraction=diffraction_api + ) input_product_api.load_probe_positions(scan_file) input_product_api.generate_probe() input_product_api.generate_object() diff --git a/src/ptychodus/plugins/xrf_maps_file.py b/src/ptychodus/plugins/xrf_maps_file.py index 09c3dccf8..e21753b17 100644 --- a/src/ptychodus/plugins/xrf_maps_file.py +++ b/src/ptychodus/plugins/xrf_maps_file.py @@ -1,16 +1,14 @@ from pathlib import Path from typing import Final -import h5py import numpy -from ptychodus.api.common import RealArrayType from ptychodus.api.fluorescence import ( - ElementMap, FluorescenceDataset, FluorescenceFileReader, FluorescenceFileWriter, ) +from ptychodus.api.io import load_fluorescence_data, save_fluorescence_data from ptychodus.api.plugins import PluginRegistry @@ -18,68 +16,11 @@ class XRFMapsFileIO(FluorescenceFileReader, FluorescenceFileWriter): SIMPLE_NAME: Final[str] = 'XRF-Maps' DISPLAY_NAME: Final[str] = 'XRF-Maps Fluorescence Dataset (*.h5 *.h5*)' - @staticmethod - def _split_path(data_path: str) -> tuple[str, str]: - parts = data_path.split('/') - return '/'.join(parts[:-1]), parts[-1] - def read(self, file_path: Path) -> FluorescenceDataset: - element_maps: list[ElementMap] = list() - counts_per_second_path = str() - channel_names_path = str() - - with h5py.File(file_path, 'r') as h5_file: - # try to see if v10 layout, Non Negative Lease squares fitting tech was used - h5_counts_per_second = h5_file['/MAPS/XRF_Analyzed/NNLS/Counts_Per_Sec'] - h5_channel_names = h5_file['/MAPS/XRF_Analyzed/NNLS/Channel_Names'] - - if h5_counts_per_second is None: - # try to see if v10 layout, iterative matrix fitting tech was used - h5_counts_per_second = h5_file['/MAPS/XRF_Analyzed/Fitted/Counts_Per_Sec'] - h5_channel_names = h5_file['/MAPS/XRF_Analyzed/Fitted/Channel_Names'] - - if h5_counts_per_second is None: - # try to see if was saved in v9 layout - h5_counts_per_second = h5_file['/MAPS/XRF_fits'] - h5_channel_names = h5_file['/MAPS/channel_names'] - - if h5_counts_per_second is not None: - # Counts_Per_Sec is an N x H x W - # where N is number of elements, use channel_names to find what element index - counts_per_second = h5_counts_per_second[...] - channel_names = h5_channel_names[...] - - for bname, cps in zip(channel_names, counts_per_second): - string_info = h5py.check_string_dtype(bname.dtype) - name = bname.decode(string_info.encoding) - emap = ElementMap(name, cps) - element_maps.append(emap) - - counts_per_second_path = h5_counts_per_second.name - channel_names_path = h5_channel_names.name - - return FluorescenceDataset( - element_maps=element_maps, - counts_per_second_path=counts_per_second_path, - channel_names_path=channel_names_path, - ) + return load_fluorescence_data(file_path) def write(self, file_path: Path, dataset: FluorescenceDataset) -> None: - channel_names: list[str] = list() - counts_per_sec: list[RealArrayType] = list() - - for emap in dataset.element_maps: - channel_names.append(emap.name) - counts_per_sec.append(emap.counts_per_second) - - cps_group_path, cps_dataset_name = self._split_path(dataset.counts_per_second_path) - ch_group_path, ch_dataset_name = self._split_path(dataset.channel_names_path) - - with h5py.File(file_path, 'w') as h5_file: - cps_group = h5_file.require_group(cps_group_path) - cps_group.create_dataset(cps_dataset_name, data=numpy.stack(counts_per_sec)) - ch_group = h5_file.require_group(ch_group_path) - ch_group.create_dataset(ch_dataset_name, data=channel_names, dtype='S256') + save_fluorescence_data(file_path, dataset) class NPZFluorescenceFileWriter(FluorescenceFileWriter): diff --git a/src/ptychodus/ptychodus_stream_processor.py b/src/ptychodus/ptychodus_stream_processor.py index d315e5a88..5d61ad327 100644 --- a/src/ptychodus/ptychodus_stream_processor.py +++ b/src/ptychodus/ptychodus_stream_processor.py @@ -15,6 +15,8 @@ from ptychodus.model import ModelCore import ptychodus import ptychodus.api +import ptychodus.api.diffraction +import ptychodus.api.geometry class ReconstructionThread(threading.Thread): @@ -104,10 +106,16 @@ def configure(self, config_dict: dict[str, Any]) -> None: num_arrays = config_dict['num_arrays'] num_patterns_per_array = config_dict.get('num_patterns_per_array', 1) pattern_dtype = config_dict.get('pattern_dtype', 'uint16') + detector_width_px = int(config_dict['detector_width_px']) + detector_height_px = int(config_dict['detector_height_px']) metadata = ptychodus.api.diffraction.DiffractionMetadata( num_patterns_per_array=[int(num_patterns_per_array)] * int(num_arrays), pattern_dtype=numpy.dtype(pattern_dtype), + detector_extent=ptychodus.api.geometry.ImageExtent( + width_px=detector_width_px, + height_px=detector_height_px, + ), ) self._ptychodus_streaming_context = self._ptychodus.create_streaming_context(metadata) self._ptychodus_streaming_context.start() # TODO clean up diff --git a/src/ptychodus/scripts/__init__.py b/src/ptychodus/scripts/__init__.py deleted file mode 100644 index 0c2458287..000000000 --- a/src/ptychodus/scripts/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Script entry points for ptychodus.""" diff --git a/src/ptychodus/scripts/genesis/README.md b/src/ptychodus/scripts/genesis/README.md deleted file mode 100644 index 298d07632..000000000 --- a/src/ptychodus/scripts/genesis/README.md +++ /dev/null @@ -1,34 +0,0 @@ -14 April 2025 - -AmSC Data Transfer API -====================== - -The Demo Data Transfer APIs for AmSC website (https://amsc-data-api.nersc.gov/docs) -links to a script (generate_token.py) that gets a Globus bearer token for testing -(https://gist.github.com/tylern4/924b19e58d75046e593e0db2d87f6c5c): - -```bash -python generate_token.py login \ - --mapped-collections 05d2c76a-e867-4f67-aa57-76edeb0beda0 \ - --mapped-collections 9d6d994a-6d04-11e5-ba46-22000b92c6ec -``` - -Ptychodus Installation -====================== - -Install uv - -```bash -curl -LsSf https://astral.sh/uv/install.sh | sh -``` - -Install Python -```bash -uv python install 3.11 -``` - -Install Ptychodus - -```bash -uv tool install ptychodus[globus,gui,ptychi] -``` diff --git a/src/ptychodus/scripts/genesis/olcf/README.md b/src/ptychodus/scripts/genesis/olcf/README.md deleted file mode 100644 index cba952f41..000000000 --- a/src/ptychodus/scripts/genesis/olcf/README.md +++ /dev/null @@ -1,17 +0,0 @@ -Get Tokens -========== - -Generate a token here: - -https://docs.olcf.ornl.gov/services_and_applications/s3m/overview.html#generate-a-token - -Use open enclave account along with the CSC682 project when you generate the token. - -Access Odo -========== - -Instructions: https://docs.olcf.ornl.gov/systems/odo_user_guide.html - -ssh USERNAME@login1.odo.olcf.ornl.gov - -Use password only. diff --git a/src/ptychodus/view/agent.py b/src/ptychodus/view/agent.py index e81926278..182a31aee 100644 --- a/src/ptychodus/view/agent.py +++ b/src/ptychodus/view/agent.py @@ -22,16 +22,17 @@ def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) self.text_edit = QPlainTextEdit() - send_button_size_policy = QSizePolicy( - QSizePolicy.Policy.Preferred, QSizePolicy.Policy.Preferred - ) + button_size_policy = QSizePolicy(QSizePolicy.Policy.Preferred, QSizePolicy.Policy.Preferred) self.send_button = QPushButton(QIcon(':/icons/send'), 'Send') - self.send_button.setSizePolicy(send_button_size_policy) + self.send_button.setSizePolicy(button_size_policy) + self.clear_button = QPushButton('Clear chat') + self.clear_button.setSizePolicy(button_size_policy) layout = QHBoxLayout() layout.setContentsMargins(0, 0, 0, 0) layout.addWidget(self.text_edit) layout.addWidget(self.send_button) + layout.addWidget(self.clear_button) self.setLayout(layout) diff --git a/src/ptychodus/view/core.py b/src/ptychodus/view/core.py index e707dd36c..4f7d944d3 100644 --- a/src/ptychodus/view/core.py +++ b/src/ptychodus/view/core.py @@ -29,7 +29,8 @@ from . import resources # noqa from .agent import AgentView, AgentChatView -from .diffraction import PatternsImageView, PatternsView +from .diffraction import DatasetsView, DiffractionImageView +from .fluorescence import FluorescenceView from .image import ImageView from .product import ProductView, ProductVisualizationView from .processing import ProcessingStatusView @@ -220,13 +221,13 @@ def __init__( right=self.settings_table_view, ) - self.patterns_view = PatternsView() - self.patterns_image_view = PatternsImageView() - self.patterns_action = self.navigation.add_panel( + self.datasets_view = DatasetsView() + self.diffraction_image_view = DiffractionImageView() + self.datasets_action = self.navigation.add_panel( QIcon(':/icons/patterns'), - 'Patterns', - left=self.patterns_view, - right=self.patterns_image_view, + 'Diffraction', + left=self.datasets_view, + right=self.diffraction_image_view, ) self.product_view = ProductView() @@ -305,6 +306,15 @@ def __init__( right=self.automation_widget, ) + self.fluorescence_view = FluorescenceView() + self.fluorescence_image_view = ImageView() + self.fluorescence_action = self.navigation.add_panel( + QIcon(':/icons/fluorescence'), + 'Fluorescence', + left=self.fluorescence_view, + right=self.fluorescence_image_view, + ) + self.agent_view = AgentView() self.agent_chat_view = AgentChatView() self.agent_action = self.navigation.add_panel( @@ -325,14 +335,18 @@ def __init__( self.navigation.add_subview_group( parent_action=self.product_action, - child_actions=(self.positions_action, self.probe_action, self.object_action), + child_actions=( + self.positions_action, + self.probe_action, + self.object_action, + ), insert_before=self.processing_action, child_icon_size=QSize(24, 24), ) self.navigation.add_subview_group( parent_action=self.processing_action, child_actions=(self.globus_action, self.genesis_action, self.automation_action), - insert_before=self.agent_action, + insert_before=self.fluorescence_action, child_icon_size=QSize(24, 24), ) diff --git a/src/ptychodus/view/diffraction.py b/src/ptychodus/view/diffraction.py index 827beae38..9e9415374 100644 --- a/src/ptychodus/view/diffraction.py +++ b/src/ptychodus/view/diffraction.py @@ -2,7 +2,6 @@ from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( - QCheckBox, QComboBox, QDialog, QDialogButtonBox, @@ -14,6 +13,7 @@ QMenu, QProgressBar, QPushButton, + QTableView, QTreeView, QVBoxLayout, QWidget, @@ -23,30 +23,25 @@ from .image import ImageView -class DetectorView(QGroupBox): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__('Detector', parent) - - -class PatternsButtonBox(QWidget): +class DatasetsButtonBox(QWidget): def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) - self.load_button = QPushButton('Load') - self.load_menu = QMenu() + self.insert_menu = QMenu() + self.insert_button = QPushButton('Insert') + self.save_menu = QMenu() self.save_button = QPushButton('Save') - self.close_button = QPushButton('Close') - self.analyze_button = QPushButton('Analyze') - self.analyze_menu = QMenu() + self.edit_button = QPushButton('Edit') + self.remove_button = QPushButton('Remove') - self.load_button.setMenu(self.load_menu) - self.analyze_button.setMenu(self.analyze_menu) + self.insert_button.setMenu(self.insert_menu) + self.save_button.setMenu(self.save_menu) layout = QHBoxLayout() layout.setContentsMargins(0, 0, 0, 0) - layout.addWidget(self.load_button) + layout.addWidget(self.insert_button) layout.addWidget(self.save_button) - layout.addWidget(self.close_button) - layout.addWidget(self.analyze_button) + layout.addWidget(self.edit_button) + layout.addWidget(self.remove_button) self.setLayout(layout) @@ -65,38 +60,31 @@ def _set_complete(self, complete: bool) -> None: self.completeChanged.emit() +class OpenDatasetWizardBadPixelsPage(OpenDatasetWizardPage): + """Bad-pixels chooser page — always complete; layout populated by the controller.""" + + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) + self._set_complete(True) + + class OpenDatasetWizardMetadataPage(OpenDatasetWizardPage): def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) - self.detector_extent_check_box = QCheckBox('Detector Extent') - self.detector_pixel_size_check_box = QCheckBox('Detector Pixel Size') - self.detector_distance_check_box = QCheckBox('Detector Distance') - self.pattern_crop_center_check_box = QCheckBox('Pattern Crop Center') - self.pattern_crop_extent_check_box = QCheckBox('Pattern Crop Extent') - self.probe_energy_check_box = QCheckBox('Probe Energy') - self.probe_photon_count_check_box = QCheckBox('Probe Photon Count') - self.exposure_time_check_box = QCheckBox('Exposure Time') + self.table_view = QTableView() self.setTitle('Import Metadata') layout = QVBoxLayout() - layout.addWidget(self.detector_extent_check_box) - layout.addWidget(self.detector_pixel_size_check_box) - layout.addWidget(self.detector_distance_check_box) - layout.addWidget(self.pattern_crop_center_check_box) - layout.addWidget(self.pattern_crop_extent_check_box) - layout.addWidget(self.probe_energy_check_box) - layout.addWidget(self.probe_photon_count_check_box) - layout.addWidget(self.exposure_time_check_box) - layout.addStretch() + layout.addWidget(self.table_view) self.setLayout(layout) + self._set_complete(True) -class DatasetFileLayoutDialog(QDialog): +class DatasetEditorLayoutView(QGroupBox): def __init__(self, parent: QWidget | None = None) -> None: - super().__init__(parent) + super().__init__('Layout') self.tree_view = QTreeView() - self.button_box = QDialogButtonBox() tree_header = self.tree_view.header() @@ -104,14 +92,48 @@ def __init__(self, parent: QWidget | None = None) -> None: tree_header.setDefaultAlignment(Qt.AlignmentFlag.AlignCenter) tree_header.setSectionResizeMode(QHeaderView.ResizeMode.ResizeToContents) + layout = QVBoxLayout() + layout.addWidget(self.tree_view) + self.setLayout(layout) + + +class DatasetEditorPropertiesView(QGroupBox): + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__('Properties') + self.table_view = QTableView() + + layout = QVBoxLayout() + layout.addWidget(self.table_view) + self.setLayout(layout) + + +class DatasetEditorDialog(QDialog): + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) + self.layout_view = DatasetEditorLayoutView() + self.properties_view = DatasetEditorPropertiesView() + self.button_box = QDialogButtonBox() + self.button_box.addButton(QDialogButtonBox.StandardButton.Ok) self.button_box.accepted.connect(self.accept) + top_layout = QHBoxLayout() + top_layout.addWidget(self.layout_view) + top_layout.addWidget(self.properties_view) + layout = QVBoxLayout() - layout.addWidget(self.tree_view) + layout.addLayout(top_layout) layout.addWidget(self.button_box) self.setLayout(layout) + @property + def tree_view(self) -> QTreeView: + return self.layout_view.tree_view + + @property + def table_view(self) -> QTableView: + return self.properties_view.table_view + class SimulateDiffractionDialog(QDialog): def __init__(self, parent: QWidget | None = None) -> None: @@ -145,13 +167,12 @@ def __init__(self, parent: QWidget | None = None) -> None: self.setLayout(layout) -class PatternsView(QWidget): +class DatasetsView(QWidget): def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) - self.detector_view = DetectorView() self.tree_view = QTreeView() self.info_label = QLabel() - self.button_box = PatternsButtonBox() + self.button_box = DatasetsButtonBox() self.simulate_dialog = SimulateDiffractionDialog(self) tree_view_header = self.tree_view.header() @@ -160,14 +181,13 @@ def __init__(self, parent: QWidget | None = None) -> None: tree_view_header.setDefaultAlignment(Qt.AlignmentFlag.AlignCenter) layout = QVBoxLayout() - layout.addWidget(self.detector_view) layout.addWidget(self.tree_view) layout.addWidget(self.info_label) layout.addWidget(self.button_box) self.setLayout(layout) -class PatternsImageView(QWidget): +class DiffractionImageView(QWidget): def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) self.image_view = ImageView() diff --git a/src/ptychodus/view/fluorescence.py b/src/ptychodus/view/fluorescence.py new file mode 100644 index 000000000..22e74e6f0 --- /dev/null +++ b/src/ptychodus/view/fluorescence.py @@ -0,0 +1,143 @@ +from PyQt5.QtWidgets import ( + QButtonGroup, + QDialog, + QDialogButtonBox, + QFormLayout, + QGroupBox, + QHBoxLayout, + QMenu, + QPlainTextEdit, + QProgressBar, + QPushButton, + QRadioButton, + QStackedWidget, + QTreeView, + QVBoxLayout, + QWidget, +) + + +class FluorescenceEnhanceParametersView(QGroupBox): + """Lean enhancement parameter form: algorithm chooser + per-algorithm parameter stack. + + Compared with the older ``FluorescenceParametersView`` this drops the + Open/Save buttons — measured-dataset loading and enhanced-dataset saving + now live in the top-level fluorescence panel; this widget is only shown + inside the modal enhance dialog. + """ + + def __init__(self, algorithm_widget: QWidget, parent: QWidget | None = None) -> None: + super().__init__('Enhancement Strategy', parent) + self.stacked_widget = QStackedWidget() + + stacked_widget_layout = self.stacked_widget.layout() + + if stacked_widget_layout is not None: + stacked_widget_layout.setContentsMargins(0, 0, 0, 0) + + layout = QFormLayout() + layout.addRow('Algorithm:', algorithm_widget) + layout.addRow(self.stacked_widget) + self.setLayout(layout) + + +class FluorescenceStatusView(QGroupBox): + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__('Status', parent) + self.text_edit = QPlainTextEdit() + self.text_edit.setReadOnly(True) + self.progress_bar = QProgressBar() + self.stop_button = QPushButton('Stop') + + progress_layout = QHBoxLayout() + progress_layout.addWidget(self.progress_bar) + progress_layout.addWidget(self.stop_button) + + layout = QVBoxLayout() + layout.addWidget(self.text_edit) + layout.addLayout(progress_layout) + self.setLayout(layout) + + +class FluorescenceEnhanceDialog(QDialog): + """Modal enhancement dialog: algorithm + params + status + Run/Close. + + Visualization, element selection, and save affordances live in the parent + panel; this dialog is intentionally minimal. + """ + + def __init__(self, algorithm_widget: QWidget, parent: QWidget | None = None) -> None: + super().__init__(parent) + self.parameters_view = FluorescenceEnhanceParametersView(algorithm_widget) + self.status_view = FluorescenceStatusView() + self.run_button = QPushButton('Run') + self.button_box = QDialogButtonBox(QDialogButtonBox.StandardButton.Close) + self.button_box.rejected.connect(self.reject) + + run_layout = QHBoxLayout() + run_layout.addStretch() + run_layout.addWidget(self.run_button) + + layout = QVBoxLayout() + layout.addWidget(self.parameters_view) + layout.addLayout(run_layout) + layout.addWidget(self.status_view, 1) + layout.addWidget(self.button_box) + self.setLayout(layout) + + +class FluorescenceButtonBox(QWidget): + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) + self.load_button = QPushButton('Load') + self.enhance_menu = QMenu() + self.enhance_button = QPushButton('Enhance') + self.save_button = QPushButton('Save') + self.remove_button = QPushButton('Remove') + + self.enhance_button.setMenu(self.enhance_menu) + + layout = QHBoxLayout() + layout.setContentsMargins(0, 0, 0, 0) + layout.addWidget(self.load_button) + layout.addWidget(self.enhance_button) + layout.addWidget(self.save_button) + layout.addWidget(self.remove_button) + self.setLayout(layout) + + +class FluorescenceView(QWidget): + """Left-pane dataset browser for the top-level Fluorescence subview. + + Layout (top → bottom): dataset tree (expandable to element leaves), a + measured/enhanced variant selector, and the button box. Each fluorescence + dataset is bound to a target product at load time, so the panel carries + no global product picker — the tree's Product column shows the bound + product per item. + """ + + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) + self.tree_view = QTreeView() + self.tree_view.setRootIsDecorated(True) + self.tree_view.setUniformRowHeights(True) + self.measured_radio_button = QRadioButton('Measured') + self.enhanced_radio_button = QRadioButton('Enhanced') + self.measured_radio_button.setChecked(True) + self.variant_button_group = QButtonGroup(self) + self.variant_button_group.setExclusive(True) + self.variant_button_group.addButton(self.measured_radio_button) + self.variant_button_group.addButton(self.enhanced_radio_button) + self.button_box = FluorescenceButtonBox() + + variant_group = QGroupBox('Variant') + variant_layout = QHBoxLayout() + variant_layout.addWidget(self.measured_radio_button) + variant_layout.addWidget(self.enhanced_radio_button) + variant_group.setLayout(variant_layout) + + layout = QVBoxLayout() + layout.addWidget(self.tree_view, 1) + layout.addWidget(variant_group) + layout.addWidget(self.button_box) + self.setLayout(layout) diff --git a/src/ptychodus/view/genesis.svg b/src/ptychodus/view/genesis.svg deleted file mode 120000 index 993e8410e..000000000 --- a/src/ptychodus/view/genesis.svg +++ /dev/null @@ -1 +0,0 @@ -../../../genesis.svg \ No newline at end of file diff --git a/src/ptychodus/view/globus.svg b/src/ptychodus/view/globus.svg deleted file mode 120000 index b729917a6..000000000 --- a/src/ptychodus/view/globus.svg +++ /dev/null @@ -1 +0,0 @@ -../../../globus.svg \ No newline at end of file diff --git a/src/ptychodus/view/image.py b/src/ptychodus/view/image.py index 49bae5826..20ef81daa 100644 --- a/src/ptychodus/view/image.py +++ b/src/ptychodus/view/image.py @@ -22,7 +22,7 @@ from ptychodus.api.common import RealArrayType from .visualization import VisualizationView -from .widgets import BottomTitledGroupBox, DecimalLineEdit, DecimalSlider +from .widgets import BottomTitledGroupBox, DecimalLineEdit, DecimalRangeSlider class ImageDisplayRangeDialog(QDialog): @@ -124,14 +124,12 @@ def __init__(self, parent: QWidget | None = None) -> None: class ImageDataRangeGroupBox(BottomTitledGroupBox): def __init__(self, parent: QWidget | None = None) -> None: super().__init__('Data Range', parent) - self.min_display_value_slider = DecimalSlider.create_instance(Qt.Orientation.Horizontal) - self.max_display_value_slider = DecimalSlider.create_instance(Qt.Orientation.Horizontal) + self.display_range_slider = DecimalRangeSlider.create_instance(Qt.Orientation.Horizontal) self.auto_button = QPushButton('Auto') self.edit_button = QPushButton('Edit') self.color_legend_button = QPushButton('Color Legend') - self.min_display_value_slider.setToolTip('Minimum Display Value') - self.max_display_value_slider.setToolTip('Maximum Display Value') + self.display_range_slider.setToolTip('Display Value Range') self.auto_button.setToolTip('Rescale to Data Range') self.edit_button.setToolTip('Rescale to Custom Range') self.color_legend_button.setToolTip('Toggle Color Legend Visibility') @@ -142,11 +140,10 @@ def __init__(self, parent: QWidget | None = None) -> None: button_layout.addWidget(self.edit_button) button_layout.addWidget(self.color_legend_button) - layout = QFormLayout() + layout = QVBoxLayout() layout.setContentsMargins(10, 10, 10, 35) - layout.addRow('Min:', self.min_display_value_slider) - layout.addRow('Max:', self.max_display_value_slider) - layout.addRow(button_layout) + layout.addWidget(self.display_range_slider) + layout.addLayout(button_layout) self.setLayout(layout) self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Preferred) diff --git a/src/ptychodus/view/make_qrc.sh b/src/ptychodus/view/make_qrc.sh index 02403d191..d604504ba 100755 --- a/src/ptychodus/view/make_qrc.sh +++ b/src/ptychodus/view/make_qrc.sh @@ -1,6 +1,13 @@ #!/bin/sh - -wget https://github.com/FortAwesome/Font-Awesome/archive/7.1.0.tar.gz -O font-awesome.tar.gz -tar xf font-awesome.tar.gz +# Regenerate resources.py from resources.qrc. +# All referenced icons live in ../../ptychodus_store/ui/icons/ (see the README there +# for how to add or update icons). +set -e pyrcc5 resources.qrc -o resources.py -rm -i font-awesome.tar.gz +# pyrcc5 emits camelCase function names that ruff flags as N802 and formatting that +# ruff would rewrite. Apply the same fixups the tracked file uses so lint stays clean. +sed -i \ + -e 's/^def qInitResources():$/def qInitResources() -> None: # noqa: N802/' \ + -e 's/^def qCleanupResources():$/def qCleanupResources() -> None: # noqa: N802/' \ + resources.py +ruff format resources.py >/dev/null diff --git a/src/ptychodus/view/probe.py b/src/ptychodus/view/probe.py index cf244150d..e1e851937 100644 --- a/src/ptychodus/view/probe.py +++ b/src/ptychodus/view/probe.py @@ -1,22 +1,15 @@ from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( - QCheckBox, - QComboBox, QDialog, QFormLayout, QGridLayout, QGroupBox, QHBoxLayout, QLabel, - QListView, - QPlainTextEdit, - QProgressBar, QPushButton, QRadioButton, QScrollArea, QSlider, - QSpinBox, - QStackedWidget, QStatusBar, QVBoxLayout, QWidget, @@ -206,123 +199,3 @@ def __init__(self, parent: QWidget | None = None) -> None: layout.addLayout(contents_layout) layout.addWidget(self.status_bar) self.setLayout(layout) - - -class FluorescenceVSPIParametersView(QWidget): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__(parent) - self.damping_factor_line_edit = DecimalLineEdit.create_instance() - self.max_iterations_spin_box = QSpinBox() - - layout = QFormLayout() - layout.setContentsMargins(0, 0, 0, 0) - layout.addRow('Damping Factor:', self.damping_factor_line_edit) - layout.addRow('Max Iterations:', self.max_iterations_spin_box) - self.setLayout(layout) - - -class FluorescencePtychozoonParametersView(QWidget): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__(parent) - self.damping_factor_line_edit = DecimalLineEdit.create_instance() - self.gradient_smoothness_line_edit = DecimalLineEdit.create_instance() - self.max_iterations_spin_box = QSpinBox() - self.atol_line_edit = DecimalLineEdit.create_instance() - self.btol_line_edit = DecimalLineEdit.create_instance() - self.checkpoint_interval_spin_box = QSpinBox() - self.use_gpu_check_box = QCheckBox('Use GPU') - self.gpu_device_index_spin_box = QSpinBox() - - layout = QFormLayout() - layout.setContentsMargins(0, 0, 0, 0) - layout.addRow('Damping Factor:', self.damping_factor_line_edit) - layout.addRow('Gradient Smoothness:', self.gradient_smoothness_line_edit) - layout.addRow('Max Iterations:', self.max_iterations_spin_box) - layout.addRow('A Tolerance:', self.atol_line_edit) - layout.addRow('B Tolerance:', self.btol_line_edit) - layout.addRow('Checkpoint Interval:', self.checkpoint_interval_spin_box) - layout.addRow(self.use_gpu_check_box) - layout.addRow('CUDA Device Index:', self.gpu_device_index_spin_box) - self.setLayout(layout) - - -class FluorescenceTwoStepParametersView(QWidget): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__(parent) - self.upscaling_strategy_combo_box = QComboBox() - self.deconvolution_strategy_combo_box = QComboBox() - - layout = QFormLayout() - layout.setContentsMargins(0, 0, 0, 0) - layout.addRow('Upscaling Strategy:', self.upscaling_strategy_combo_box) - layout.addRow('Deconvolution Strategy:', self.deconvolution_strategy_combo_box) - self.setLayout(layout) - - -class FluorescenceParametersView(QGroupBox): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__('Enhancement Strategy', parent) - self.open_button = QPushButton('Open Measured Dataset') - self.algorithm_combo_box = QComboBox() - self.stacked_widget = QStackedWidget() - self.enhance_button = QPushButton('Enhance') - self.save_button = QPushButton('Save Enhanced Dataset') - - stacked_widget_layout = self.stacked_widget.layout() - - if stacked_widget_layout is not None: - stacked_widget_layout.setContentsMargins(0, 0, 0, 0) - - layout = QFormLayout() - layout.addRow(self.open_button) - layout.addRow('Algorithm:', self.algorithm_combo_box) - layout.addRow(self.stacked_widget) - layout.addRow(self.enhance_button) - layout.addRow(self.save_button) - self.setLayout(layout) - - -class FluorescenceStatusView(QGroupBox): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__('Status', parent) - self.text_edit = QPlainTextEdit() - self.text_edit.setReadOnly(True) - self.progress_bar = QProgressBar() - self.stop_button = QPushButton('Stop') - - progress_layout = QHBoxLayout() - progress_layout.addWidget(self.progress_bar) - progress_layout.addWidget(self.stop_button) - - layout = QVBoxLayout() - layout.addWidget(self.text_edit) - layout.addLayout(progress_layout) - self.setLayout(layout) - - -class FluorescenceDialog(QDialog): - def __init__(self, parent: QWidget | None = None) -> None: - super().__init__(parent) - self.measured_widget = VisualizationWidget('Measured') - self.enhanced_widget = VisualizationWidget('Enhanced') - self.fluorescence_parameters_view = FluorescenceParametersView() - self.fluorescence_channel_list_view = QListView() - self.visualization_parameters_view = VisualizationParametersView() - self.fluorescence_status_view = FluorescenceStatusView() - self.status_bar = QStatusBar() - - parameter_layout = QVBoxLayout() - parameter_layout.addWidget(self.fluorescence_parameters_view) - parameter_layout.addWidget(self.fluorescence_channel_list_view, 1) - parameter_layout.addWidget(self.visualization_parameters_view) - parameter_layout.addWidget(self.fluorescence_status_view, 1) - - contents_layout = QHBoxLayout() - contents_layout.addWidget(self.measured_widget, 1) - contents_layout.addWidget(self.enhanced_widget, 1) - contents_layout.addLayout(parameter_layout) - - layout = QVBoxLayout() - layout.addLayout(contents_layout) - layout.addWidget(self.status_bar) - self.setLayout(layout) diff --git a/src/ptychodus/view/ptychodus.svg b/src/ptychodus/view/ptychodus.svg deleted file mode 120000 index a314ea958..000000000 --- a/src/ptychodus/view/ptychodus.svg +++ /dev/null @@ -1 +0,0 @@ -../../../ptychodus.svg \ No newline at end of file diff --git a/src/ptychodus/view/resources.py b/src/ptychodus/view/resources.py index 28b87f68e..982c78b2a 100644 --- a/src/ptychodus/view/resources.py +++ b/src/ptychodus/view/resources.py @@ -369,6 +369,98 @@ \x2e\x37\x20\x31\x38\x2e\x37\x20\x34\x39\x2e\x31\x20\x30\x20\x36\ \x37\x2e\x39\x4c\x32\x30\x39\x2e\x31\x20\x35\x31\x36\x2e\x32\x7a\ \x22\x2f\x3e\x3c\x2f\x73\x76\x67\x3e\ +\x00\x00\x05\x97\ +\x3c\ +\x73\x76\x67\x20\x78\x6d\x6c\x6e\x73\x3d\x22\x68\x74\x74\x70\x3a\ +\x2f\x2f\x77\x77\x77\x2e\x77\x33\x2e\x6f\x72\x67\x2f\x32\x30\x30\ +\x30\x2f\x73\x76\x67\x22\x20\x76\x69\x65\x77\x42\x6f\x78\x3d\x22\ +\x30\x20\x30\x20\x34\x34\x38\x20\x35\x31\x32\x22\x3e\x3c\x21\x2d\ +\x2d\x21\x20\x46\x6f\x6e\x74\x20\x41\x77\x65\x73\x6f\x6d\x65\x20\ +\x46\x72\x65\x65\x20\x37\x2e\x31\x2e\x30\x20\x62\x79\x20\x40\x66\ +\x6f\x6e\x74\x61\x77\x65\x73\x6f\x6d\x65\x20\x2d\x20\x68\x74\x74\ +\x70\x73\x3a\x2f\x2f\x66\x6f\x6e\x74\x61\x77\x65\x73\x6f\x6d\x65\ +\x2e\x63\x6f\x6d\x20\x4c\x69\x63\x65\x6e\x73\x65\x20\x2d\x20\x68\ +\x74\x74\x70\x73\x3a\x2f\x2f\x66\x6f\x6e\x74\x61\x77\x65\x73\x6f\ +\x6d\x65\x2e\x63\x6f\x6d\x2f\x6c\x69\x63\x65\x6e\x73\x65\x2f\x66\ +\x72\x65\x65\x20\x28\x49\x63\x6f\x6e\x73\x3a\x20\x43\x43\x20\x42\ +\x59\x20\x34\x2e\x30\x2c\x20\x46\x6f\x6e\x74\x73\x3a\x20\x53\x49\ +\x4c\x20\x4f\x46\x4c\x20\x31\x2e\x31\x2c\x20\x43\x6f\x64\x65\x3a\ +\x20\x4d\x49\x54\x20\x4c\x69\x63\x65\x6e\x73\x65\x29\x20\x43\x6f\ +\x70\x79\x72\x69\x67\x68\x74\x20\x32\x30\x32\x35\x20\x46\x6f\x6e\ +\x74\x69\x63\x6f\x6e\x73\x2c\x20\x49\x6e\x63\x2e\x20\x2d\x2d\x3e\ +\x3c\x70\x61\x74\x68\x20\x66\x69\x6c\x6c\x3d\x22\x63\x75\x72\x72\ +\x65\x6e\x74\x43\x6f\x6c\x6f\x72\x22\x20\x64\x3d\x22\x4d\x32\x32\ +\x34\x20\x33\x39\x38\x2e\x38\x63\x2d\x31\x31\x2e\x38\x20\x35\x2e\ +\x31\x2d\x32\x33\x2e\x34\x20\x39\x2e\x37\x2d\x33\x34\x2e\x39\x20\ +\x31\x33\x2e\x35\x20\x31\x36\x2e\x37\x20\x33\x33\x2e\x38\x20\x33\ +\x31\x20\x33\x35\x2e\x37\x20\x33\x34\x2e\x39\x20\x33\x35\x2e\x37\ +\x73\x31\x38\x2e\x31\x2d\x31\x2e\x39\x20\x33\x34\x2e\x39\x2d\x33\ +\x35\x2e\x37\x63\x2d\x31\x31\x2e\x34\x2d\x33\x2e\x39\x2d\x32\x33\ +\x2e\x31\x2d\x38\x2e\x34\x2d\x33\x34\x2e\x39\x2d\x31\x33\x2e\x35\ +\x7a\x4d\x34\x31\x34\x20\x32\x35\x36\x63\x33\x33\x20\x34\x35\x2e\ +\x32\x20\x34\x34\x2e\x33\x20\x39\x30\x2e\x39\x20\x32\x33\x2e\x36\ +\x20\x31\x32\x38\x2d\x32\x30\x2e\x32\x20\x33\x36\x2e\x33\x2d\x36\ +\x32\x2e\x35\x20\x34\x39\x2e\x33\x2d\x31\x31\x35\x2e\x32\x20\x34\ +\x33\x2e\x32\x2d\x32\x32\x20\x35\x32\x2e\x31\x2d\x35\x35\x2e\x37\ +\x20\x38\x34\x2e\x38\x2d\x39\x38\x2e\x34\x20\x38\x34\x2e\x38\x73\ +\x2d\x37\x36\x2e\x34\x2d\x33\x32\x2e\x37\x2d\x39\x38\x2e\x34\x2d\ +\x38\x34\x2e\x38\x43\x37\x32\x2e\x39\x20\x34\x33\x33\x2e\x33\x20\ +\x33\x30\x2e\x36\x20\x34\x32\x30\x2e\x33\x20\x31\x30\x2e\x34\x20\ +\x33\x38\x34\x2d\x31\x30\x2e\x33\x20\x33\x34\x36\x2e\x39\x20\x31\ +\x20\x33\x30\x31\x2e\x32\x20\x33\x34\x20\x32\x35\x36\x20\x31\x20\ +\x32\x31\x30\x2e\x38\x2d\x31\x30\x2e\x33\x20\x31\x36\x35\x2e\x31\ +\x20\x31\x30\x2e\x34\x20\x31\x32\x38\x20\x33\x30\x2e\x36\x20\x39\ +\x31\x2e\x37\x20\x37\x32\x2e\x39\x20\x37\x38\x2e\x37\x20\x31\x32\ +\x35\x2e\x36\x20\x38\x34\x2e\x38\x20\x31\x34\x37\x2e\x36\x20\x33\ +\x32\x2e\x37\x20\x31\x38\x31\x2e\x32\x20\x30\x20\x32\x32\x34\x20\ +\x30\x73\x37\x36\x2e\x34\x20\x33\x32\x2e\x37\x20\x39\x38\x2e\x34\ +\x20\x38\x34\x2e\x38\x63\x35\x32\x2e\x37\x2d\x36\x2e\x31\x20\x39\ +\x35\x20\x36\x2e\x38\x20\x31\x31\x35\x2e\x32\x20\x34\x33\x2e\x32\ +\x20\x32\x30\x2e\x37\x20\x33\x37\x2e\x31\x20\x39\x2e\x34\x20\x38\ +\x32\x2e\x38\x2d\x32\x33\x2e\x36\x20\x31\x32\x38\x7a\x6d\x2d\x36\ +\x35\x2e\x38\x20\x36\x37\x2e\x34\x63\x2d\x31\x2e\x37\x20\x31\x34\ +\x2e\x32\x2d\x33\x2e\x39\x20\x32\x38\x2d\x36\x2e\x37\x20\x34\x31\ +\x2e\x32\x20\x33\x31\x2e\x38\x20\x31\x2e\x34\x20\x33\x38\x2e\x36\ +\x2d\x38\x2e\x37\x20\x34\x30\x2e\x32\x2d\x31\x31\x2e\x37\x20\x32\ +\x2e\x33\x2d\x34\x2e\x32\x20\x37\x2d\x31\x37\x2e\x39\x2d\x31\x31\ +\x2e\x39\x2d\x34\x38\x2e\x31\x2d\x36\x2e\x38\x20\x36\x2e\x33\x2d\ +\x31\x34\x20\x31\x32\x2e\x35\x2d\x32\x31\x2e\x36\x20\x31\x38\x2e\ +\x36\x7a\x6d\x2d\x36\x2e\x37\x2d\x31\x37\x35\x2e\x39\x63\x32\x2e\ +\x38\x20\x31\x33\x2e\x31\x20\x35\x20\x32\x36\x2e\x39\x20\x36\x2e\ +\x37\x20\x34\x31\x2e\x32\x20\x37\x2e\x36\x20\x36\x2e\x31\x20\x31\ +\x34\x2e\x38\x20\x31\x32\x2e\x33\x20\x32\x31\x2e\x36\x20\x31\x38\ +\x2e\x36\x20\x31\x38\x2e\x39\x2d\x33\x30\x2e\x32\x20\x31\x34\x2e\ +\x32\x2d\x34\x34\x20\x31\x31\x2e\x39\x2d\x34\x38\x2e\x31\x2d\x31\ +\x2e\x36\x2d\x32\x2e\x39\x2d\x38\x2e\x34\x2d\x31\x33\x2d\x34\x30\ +\x2e\x32\x2d\x31\x31\x2e\x37\x7a\x4d\x32\x35\x38\x2e\x39\x20\x39\ +\x39\x2e\x37\x43\x32\x34\x32\x2e\x31\x20\x36\x35\x2e\x39\x20\x32\ +\x32\x37\x2e\x39\x20\x36\x34\x20\x32\x32\x34\x20\x36\x34\x73\x2d\ +\x31\x38\x2e\x31\x20\x31\x2e\x39\x2d\x33\x34\x2e\x39\x20\x33\x35\ +\x2e\x37\x63\x31\x31\x2e\x34\x20\x33\x2e\x39\x20\x32\x33\x2e\x31\ +\x20\x38\x2e\x34\x20\x33\x34\x2e\x39\x20\x31\x33\x2e\x35\x20\x31\ +\x31\x2e\x38\x2d\x35\x2e\x31\x20\x32\x33\x2e\x34\x2d\x39\x2e\x37\ +\x20\x33\x34\x2e\x39\x2d\x31\x33\x2e\x35\x7a\x6d\x2d\x31\x35\x39\ +\x20\x38\x38\x2e\x39\x63\x31\x2e\x37\x2d\x31\x34\x2e\x33\x20\x33\ +\x2e\x39\x2d\x32\x38\x20\x36\x2e\x37\x2d\x34\x31\x2e\x32\x2d\x33\ +\x31\x2e\x38\x2d\x31\x2e\x34\x2d\x33\x38\x2e\x36\x20\x38\x2e\x37\ +\x2d\x34\x30\x2e\x32\x20\x31\x31\x2e\x37\x2d\x32\x2e\x33\x20\x34\ +\x2e\x32\x2d\x37\x20\x31\x37\x2e\x39\x20\x31\x31\x2e\x39\x20\x34\ +\x38\x2e\x31\x20\x36\x2e\x38\x2d\x36\x2e\x33\x20\x31\x34\x2d\x31\ +\x32\x2e\x35\x20\x32\x31\x2e\x36\x2d\x31\x38\x2e\x36\x7a\x4d\x37\ +\x38\x2e\x32\x20\x33\x30\x34\x2e\x38\x63\x2d\x31\x38\x2e\x39\x20\ +\x33\x30\x2e\x32\x2d\x31\x34\x2e\x32\x20\x34\x34\x2d\x31\x31\x2e\ +\x39\x20\x34\x38\x2e\x31\x20\x31\x2e\x36\x20\x32\x2e\x39\x20\x38\ +\x2e\x34\x20\x31\x33\x20\x34\x30\x2e\x32\x20\x31\x31\x2e\x37\x2d\ +\x32\x2e\x38\x2d\x31\x33\x2e\x31\x2d\x35\x2d\x32\x36\x2e\x39\x2d\ +\x36\x2e\x37\x2d\x34\x31\x2e\x32\x2d\x37\x2e\x36\x2d\x36\x2e\x31\ +\x2d\x31\x34\x2e\x38\x2d\x31\x32\x2e\x33\x2d\x32\x31\x2e\x36\x2d\ +\x31\x38\x2e\x36\x7a\x4d\x33\x30\x34\x20\x32\x35\x36\x61\x38\x30\ +\x20\x38\x30\x20\x30\x20\x31\x20\x30\x20\x2d\x31\x36\x30\x20\x30\ +\x20\x38\x30\x20\x38\x30\x20\x30\x20\x31\x20\x30\x20\x31\x36\x30\ +\x20\x30\x7a\x6d\x2d\x38\x30\x2d\x33\x32\x61\x33\x32\x20\x33\x32\ +\x20\x30\x20\x31\x20\x31\x20\x30\x20\x36\x34\x20\x33\x32\x20\x33\ +\x32\x20\x30\x20\x31\x20\x31\x20\x30\x2d\x36\x34\x7a\x22\x2f\x3e\ +\x3c\x2f\x73\x76\x67\x3e\ \x00\x00\x05\x4b\ \x3c\ \x73\x76\x67\x20\x78\x6d\x6c\x6e\x73\x3d\x22\x68\x74\x74\x70\x3a\ @@ -1819,6 +1911,10 @@ \x00\x79\xc2\xc2\ \x00\x72\ \x00\x75\x00\x6c\x00\x65\x00\x72\ +\x00\x0c\ +\x01\x3e\x49\x95\ +\x00\x66\ +\x00\x6c\x00\x75\x00\x6f\x00\x72\x00\x65\x00\x73\x00\x63\x00\x65\x00\x6e\x00\x63\x00\x65\ \x00\x08\ \x06\x89\x2d\x83\ \x00\x73\ @@ -1879,7 +1975,7 @@ qt_resource_struct_v1 = b'\ \x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\x00\x00\x00\x01\ -\x00\x00\x00\x00\x00\x02\x00\x00\x00\x15\x00\x00\x00\x02\ +\x00\x00\x00\x00\x00\x02\x00\x00\x00\x16\x00\x00\x00\x02\ \x00\x00\x00\x10\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\ \x00\x00\x00\x1e\x00\x00\x00\x00\x00\x01\x00\x00\x02\xa0\ \x00\x00\x00\x2c\x00\x00\x00\x00\x00\x01\x00\x00\x06\x72\ @@ -1888,71 +1984,74 @@ \x00\x00\x00\x5e\x00\x00\x00\x00\x00\x01\x00\x00\x0e\x63\ \x00\x00\x00\x6e\x00\x00\x00\x00\x00\x01\x00\x00\x12\x51\ \x00\x00\x00\x7e\x00\x00\x00\x00\x00\x01\x00\x00\x15\x8f\ -\x00\x00\x00\x94\x00\x00\x00\x00\x00\x01\x00\x00\x1a\xde\ -\x00\x00\x00\xa6\x00\x00\x00\x00\x00\x01\x00\x00\x27\xb6\ -\x00\x00\x00\xb8\x00\x00\x00\x00\x00\x01\x00\x00\x2b\x64\ -\x00\x00\x00\xce\x00\x00\x00\x00\x00\x01\x00\x00\x2e\x0c\ -\x00\x00\x00\xe4\x00\x00\x00\x00\x00\x01\x00\x00\x31\xb1\ -\x00\x00\x00\xfe\x00\x00\x00\x00\x00\x01\x00\x00\x36\x30\ -\x00\x00\x01\x16\x00\x00\x00\x00\x00\x01\x00\x00\x39\xd3\ -\x00\x00\x01\x2e\x00\x00\x00\x00\x00\x01\x00\x00\x3d\xdf\ -\x00\x00\x01\x46\x00\x00\x00\x00\x00\x01\x00\x00\x41\x27\ -\x00\x00\x01\x5c\x00\x00\x00\x00\x00\x01\x00\x00\x45\x2c\ -\x00\x00\x01\x72\x00\x00\x00\x00\x00\x01\x00\x00\x49\x6f\ -\x00\x00\x01\x86\x00\x00\x00\x00\x00\x01\x00\x00\x4b\x5b\ -\x00\x00\x01\x9a\x00\x00\x00\x00\x00\x01\x00\x00\x53\x5b\ +\x00\x00\x00\x9c\x00\x00\x00\x00\x00\x01\x00\x00\x1b\x2a\ +\x00\x00\x00\xb2\x00\x00\x00\x00\x00\x01\x00\x00\x20\x79\ +\x00\x00\x00\xc4\x00\x00\x00\x00\x00\x01\x00\x00\x2d\x51\ +\x00\x00\x00\xd6\x00\x00\x00\x00\x00\x01\x00\x00\x30\xff\ +\x00\x00\x00\xec\x00\x00\x00\x00\x00\x01\x00\x00\x33\xa7\ +\x00\x00\x01\x02\x00\x00\x00\x00\x00\x01\x00\x00\x37\x4c\ +\x00\x00\x01\x1c\x00\x00\x00\x00\x00\x01\x00\x00\x3b\xcb\ +\x00\x00\x01\x34\x00\x00\x00\x00\x00\x01\x00\x00\x3f\x6e\ +\x00\x00\x01\x4c\x00\x00\x00\x00\x00\x01\x00\x00\x43\x7a\ +\x00\x00\x01\x64\x00\x00\x00\x00\x00\x01\x00\x00\x46\xc2\ +\x00\x00\x01\x7a\x00\x00\x00\x00\x00\x01\x00\x00\x4a\xc7\ +\x00\x00\x01\x90\x00\x00\x00\x00\x00\x01\x00\x00\x4f\x0a\ +\x00\x00\x01\xa4\x00\x00\x00\x00\x00\x01\x00\x00\x50\xf6\ +\x00\x00\x01\xb8\x00\x00\x00\x00\x00\x01\x00\x00\x58\xf6\ ' qt_resource_struct_v2 = b'\ \x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\x00\x00\x00\x01\ \x00\x00\x00\x00\x00\x00\x00\x00\ -\x00\x00\x00\x00\x00\x02\x00\x00\x00\x15\x00\x00\x00\x02\ +\x00\x00\x00\x00\x00\x02\x00\x00\x00\x16\x00\x00\x00\x02\ \x00\x00\x00\x00\x00\x00\x00\x00\ \x00\x00\x00\x10\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x4e\ \x00\x00\x00\x1e\x00\x00\x00\x00\x00\x01\x00\x00\x02\xa0\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x57\ \x00\x00\x00\x2c\x00\x00\x00\x00\x00\x01\x00\x00\x06\x72\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x8c\ \x00\x00\x00\x3a\x00\x00\x00\x00\x00\x01\x00\x00\x09\x4e\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x7c\ \x00\x00\x00\x48\x00\x00\x00\x00\x00\x01\x00\x00\x0b\xc4\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x53\ \x00\x00\x00\x5e\x00\x00\x00\x00\x00\x01\x00\x00\x0e\x63\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x68\ \x00\x00\x00\x6e\x00\x00\x00\x00\x00\x01\x00\x00\x12\x51\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ +\x00\x00\x01\x9f\x80\xee\xf3\x78\ \x00\x00\x00\x7e\x00\x00\x00\x00\x00\x01\x00\x00\x15\x8f\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x00\x94\x00\x00\x00\x00\x00\x01\x00\x00\x1a\xde\ +\x00\x00\x01\x9f\x80\xee\xf3\x89\ +\x00\x00\x00\x9c\x00\x00\x00\x00\x00\x01\x00\x00\x1b\x2a\ +\x00\x00\x01\x9f\x80\xee\xf3\x85\ +\x00\x00\x00\xb2\x00\x00\x00\x00\x00\x01\x00\x00\x20\x79\ \x00\x00\x01\x9a\x5a\x95\xb8\xc7\ -\x00\x00\x00\xa6\x00\x00\x00\x00\x00\x01\x00\x00\x27\xb6\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x00\xb8\x00\x00\x00\x00\x00\x01\x00\x00\x2b\x64\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x00\xce\x00\x00\x00\x00\x00\x01\x00\x00\x2e\x0c\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x00\xe4\x00\x00\x00\x00\x00\x01\x00\x00\x31\xb1\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x00\xfe\x00\x00\x00\x00\x00\x01\x00\x00\x36\x30\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x01\x16\x00\x00\x00\x00\x00\x01\x00\x00\x39\xd3\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x01\x2e\x00\x00\x00\x00\x00\x01\x00\x00\x3d\xdf\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x01\x46\x00\x00\x00\x00\x00\x01\x00\x00\x41\x27\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x01\x5c\x00\x00\x00\x00\x00\x01\x00\x00\x45\x2c\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x01\x72\x00\x00\x00\x00\x00\x01\x00\x00\x49\x6f\ -\x00\x00\x01\x99\xa1\x32\xb9\x00\ -\x00\x00\x01\x86\x00\x00\x00\x00\x00\x01\x00\x00\x4b\x5b\ -\x00\x00\x01\x9d\x3f\xe5\x73\x03\ -\x00\x00\x01\x9a\x00\x00\x00\x00\x00\x01\x00\x00\x53\x5b\ +\x00\x00\x00\xc4\x00\x00\x00\x00\x00\x01\x00\x00\x2d\x51\ +\x00\x00\x01\x9f\x80\xee\xf3\x5b\ +\x00\x00\x00\xd6\x00\x00\x00\x00\x00\x01\x00\x00\x30\xff\ +\x00\x00\x01\x9f\x80\xee\xf3\x5f\ +\x00\x00\x00\xec\x00\x00\x00\x00\x00\x01\x00\x00\x33\xa7\ +\x00\x00\x01\x9f\x80\xee\xf3\x70\ +\x00\x00\x01\x02\x00\x00\x00\x00\x00\x01\x00\x00\x37\x4c\ +\x00\x00\x01\x9f\x80\xee\xf3\x6c\ +\x00\x00\x01\x1c\x00\x00\x00\x00\x00\x01\x00\x00\x3b\xcb\ +\x00\x00\x01\x9f\x80\xee\xf3\x64\ +\x00\x00\x01\x34\x00\x00\x00\x00\x00\x01\x00\x00\x3f\x6e\ +\x00\x00\x01\x9f\x80\xee\xf3\x74\ +\x00\x00\x01\x4c\x00\x00\x00\x00\x00\x01\x00\x00\x43\x7a\ +\x00\x00\x01\x9f\x80\xee\xf3\x47\ +\x00\x00\x01\x64\x00\x00\x00\x00\x00\x01\x00\x00\x46\xc2\ +\x00\x00\x01\x9f\x80\xee\xf3\x42\ +\x00\x00\x01\x7a\x00\x00\x00\x00\x00\x01\x00\x00\x4a\xc7\ +\x00\x00\x01\x9f\x80\xee\xf3\x80\ +\x00\x00\x01\x90\x00\x00\x00\x00\x00\x01\x00\x00\x4f\x0a\ +\x00\x00\x01\x9f\x80\xee\xf3\x4b\ +\x00\x00\x01\xa4\x00\x00\x00\x00\x00\x01\x00\x00\x50\xf6\ +\x00\x00\x01\x9d\x6d\xf6\x87\xf8\ +\x00\x00\x01\xb8\x00\x00\x00\x00\x00\x01\x00\x00\x58\xf6\ \x00\x00\x01\x97\x54\xf1\xdc\x55\ ' -qt_version = [int(v) for v in QtCore.qVersion().split('.')] # type: ignore[union-attr] +qt_version = [int(v) for v in QtCore.qVersion().split('.')] if qt_version < [5, 8, 0]: rcc_version = 1 qt_resource_struct = qt_resource_struct_v1 diff --git a/src/ptychodus/view/resources.qrc b/src/ptychodus/view/resources.qrc index 77946235e..81e8d38d9 100644 --- a/src/ptychodus/view/resources.qrc +++ b/src/ptychodus/view/resources.qrc @@ -1,26 +1,27 @@ <!DOCTYPE RCC> <RCC version="1.0"> <qresource prefix="icons"> - <file alias="automate">Font-Awesome-7.1.0/svgs/solid/robot.svg</file> - <file alias="autoscale">Font-Awesome-7.1.0/svgs/solid/arrows-left-right-to-line.svg</file> - <file alias="fourier">Font-Awesome-7.1.0/svgs/solid/f.svg</file> - <file alias="genesis">genesis.svg</file> - <file alias="globus">globus.svg</file> - <file alias="home">Font-Awesome-7.1.0/svgs/solid/house-chimney.svg</file> - <file alias="line-cut">Font-Awesome-7.1.0/svgs/solid/chart-line.svg</file> - <file alias="move">Font-Awesome-7.1.0/svgs/solid/arrows-up-down-left-right.svg</file> - <file alias="object">Font-Awesome-7.1.0/svgs/solid/layer-group.svg</file> - <file alias="patterns">Font-Awesome-7.1.0/svgs/solid/table-cells.svg</file> - <file alias="positions">Font-Awesome-7.1.0/svgs/solid/route.svg</file> - <file alias="probe">Font-Awesome-7.1.0/svgs/solid/circle-radiation.svg</file> - <file alias="processing">Font-Awesome-7.1.0/svgs/solid/microchip.svg</file> - <file alias="products">Font-Awesome-7.1.0/svgs/solid/list.svg</file> - <file alias="ptychodus">ptychodus.svg</file> - <file alias="rectangle">Font-Awesome-7.1.0/svgs/solid/object-group.svg</file> - <file alias="ruler">Font-Awesome-7.1.0/svgs/solid/ruler.svg</file> - <file alias="save">Font-Awesome-7.1.0/svgs/regular/floppy-disk.svg</file> - <file alias="send">Font-Awesome-7.1.0/svgs/solid/paper-plane.svg</file> - <file alias="settings">Font-Awesome-7.1.0/svgs/solid/gear.svg</file> - <file alias="sparkles">Font-Awesome-7.1.0/svgs/solid/wand-magic-sparkles.svg</file> + <file alias="automate">../../ptychodus_store/ui/icons/robot.svg</file> + <file alias="autoscale">../../ptychodus_store/ui/icons/arrows-left-right-to-line.svg</file> + <file alias="fluorescence">../../ptychodus_store/ui/icons/atom.svg</file> + <file alias="fourier">../../ptychodus_store/ui/icons/f.svg</file> + <file alias="genesis">../../ptychodus_store/ui/icons/genesis.svg</file> + <file alias="globus">../../ptychodus_store/ui/icons/globus.svg</file> + <file alias="home">../../ptychodus_store/ui/icons/house-chimney.svg</file> + <file alias="line-cut">../../ptychodus_store/ui/icons/chart-line.svg</file> + <file alias="move">../../ptychodus_store/ui/icons/arrows-up-down-left-right.svg</file> + <file alias="object">../../ptychodus_store/ui/icons/layer-group.svg</file> + <file alias="patterns">../../ptychodus_store/ui/icons/table-cells.svg</file> + <file alias="positions">../../ptychodus_store/ui/icons/route.svg</file> + <file alias="probe">../../ptychodus_store/ui/icons/circle-radiation.svg</file> + <file alias="processing">../../ptychodus_store/ui/icons/microchip.svg</file> + <file alias="products">../../ptychodus_store/ui/icons/list.svg</file> + <file alias="ptychodus">../../ptychodus_store/ui/icons/ptychodus.svg</file> + <file alias="rectangle">../../ptychodus_store/ui/icons/object-group.svg</file> + <file alias="ruler">../../ptychodus_store/ui/icons/ruler.svg</file> + <file alias="save">../../ptychodus_store/ui/icons/floppy-disk.svg</file> + <file alias="send">../../ptychodus_store/ui/icons/paper-plane.svg</file> + <file alias="settings">../../ptychodus_store/ui/icons/gear.svg</file> + <file alias="sparkles">../../ptychodus_store/ui/icons/wand-magic-sparkles.svg</file> </qresource> </RCC> diff --git a/src/ptychodus/view/visualization.py b/src/ptychodus/view/visualization.py index bbc5ff4f5..0ff27b79d 100644 --- a/src/ptychodus/view/visualization.py +++ b/src/ptychodus/view/visualization.py @@ -1,4 +1,5 @@ from __future__ import annotations +from collections.abc import Sequence from enum import auto, Enum import logging @@ -26,6 +27,7 @@ QWidget, ) +from matplotlib.axes import Axes from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg from matplotlib.backends.backend_qt import NavigationToolbar2QT as NavigationToolbar from matplotlib.figure import Figure @@ -292,6 +294,19 @@ def __init__(self, parent: QWidget | None = None) -> None: self.setWindowTitle('Line-Cut Dialog') + def prepare_axes(self, count: int) -> Sequence[Axes]: + """Clear the figure and return `count` axes: a primary plus twinned y-axes. + + The figure is rebuilt rather than cleared in place because `twinx` adds a new axis + every call, which would accumulate across line cuts. + """ + if count < 1: + raise ValueError(f'Axes count must be positive (actual={count}).') + + self.figure.clear() + self.axes = self.figure.add_subplot(111) + return [self.axes, *(self.axes.twinx() for _ in range(count - 1))] + class RectangleView(QGroupBox): @staticmethod diff --git a/src/ptychodus/view/widgets/__init__.py b/src/ptychodus/view/widgets/__init__.py index fa49ff741..5075e0b5a 100644 --- a/src/ptychodus/view/widgets/__init__.py +++ b/src/ptychodus/view/widgets/__init__.py @@ -2,6 +2,7 @@ from .group_box import BottomTitledGroupBox, GroupBoxWithPresets from .combo_box_item_delegate import ComboBoxItemDelegate from .decimal_line_edit import DecimalLineEdit +from .decimal_range_slider import DecimalRangeSlider, Handle from .decimal_slider import DecimalSlider from .exception_dialog import ExceptionDialog from .length_widget import LengthWidget @@ -14,9 +15,11 @@ 'BottomTitledGroupBox', 'ComboBoxItemDelegate', 'DecimalLineEdit', + 'DecimalRangeSlider', 'DecimalSlider', 'ExceptionDialog', 'GroupBoxWithPresets', + 'Handle', 'LengthWidget', 'PowerTwoSpinBox', 'ProgressBarItemDelegate', diff --git a/src/ptychodus/view/widgets/decimal_range_slider.py b/src/ptychodus/view/widgets/decimal_range_slider.py new file mode 100644 index 000000000..59386b1f4 --- /dev/null +++ b/src/ptychodus/view/widgets/decimal_range_slider.py @@ -0,0 +1,342 @@ +from __future__ import annotations +from decimal import Decimal +from enum import Enum + +import numpy + +from PyQt5.QtCore import QPoint, QRect, Qt, pyqtSignal +from PyQt5.QtGui import ( + QBrush, + QKeyEvent, + QMouseEvent, + QPainter, + QPaintEvent, + QPen, +) +from PyQt5.QtWidgets import QHBoxLayout, QLabel, QSizePolicy, QWidget + +from ptychodus.api.geometry import Interval + + +class Handle(Enum): + LOWER = 'lower' + UPPER = 'upper' + + +class _RangeSliderPaintArea(QWidget): + """Interactive paint surface for `DecimalRangeSlider`. + + Owns paint and mouse handling; all state and the keyboard/signal surface + live on the outer `DecimalRangeSlider`. + """ + + _HANDLE_RADIUS = 7 + _GROOVE_HEIGHT = 4 + _TICK_HEIGHT = 5 + _TICK_GAP = 2 + _HIT_TOLERANCE = 2 + + def __init__(self, owner: DecimalRangeSlider) -> None: + super().__init__(owner) + self._owner = owner + height = 2 * self._HANDLE_RADIUS + self._TICK_GAP + self._TICK_HEIGHT + 4 + self.setMinimumHeight(height) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed) + + def _groove_x_range(self) -> tuple[int, int]: + pad = self._HANDLE_RADIUS + 1 + return pad, self.width() - pad - 1 + + def _groove_y(self) -> int: + return self._HANDLE_RADIUS + 1 + + def _handle_x(self, handle: Handle) -> int: + selection = self._owner._selection + value = selection.lower if handle is Handle.LOWER else selection.upper + return self._value_to_x(value) + + def _value_to_x(self, value: Decimal) -> int: + bounds = self._owner._bounds + left, right = self._groove_x_range() + span = bounds.upper - bounds.lower + if span == 0: + return left + alpha = (value - bounds.lower) / span + alpha_f = max(0.0, min(1.0, float(alpha))) + return int(numpy.rint(left + alpha_f * (right - left))) + + def _x_to_value(self, x: int) -> Decimal: + bounds = self._owner._bounds + left, right = self._groove_x_range() + span_px = right - left + if span_px <= 0: + return bounds.lower + alpha_f = max(0.0, min(1.0, (x - left) / span_px)) + num_ticks = self._owner._num_ticks + tick = int(numpy.rint(alpha_f * num_ticks)) + alpha = Decimal(tick) / Decimal(num_ticks) + return bounds.lower + alpha * (bounds.upper - bounds.lower) + + def _hit_test(self, pos: QPoint) -> Handle | None: + lower_x = self._handle_x(Handle.LOWER) + upper_x = self._handle_x(Handle.UPPER) + y_center = self._groove_y() + reach = self._HANDLE_RADIUS + self._HIT_TOLERANCE + + def within(hx: int) -> bool: + return (pos.x() - hx) ** 2 + (pos.y() - y_center) ** 2 <= reach * reach + + lower_hit = within(lower_x) + upper_hit = within(upper_x) + + if lower_hit and upper_hit: + return Handle.LOWER if pos.x() <= lower_x else Handle.UPPER + if lower_hit: + return Handle.LOWER + if upper_hit: + return Handle.UPPER + return None + + def mousePressEvent(self, event: QMouseEvent) -> None: # noqa: N802 + if event.button() != Qt.MouseButton.LeftButton: + super().mousePressEvent(event) + return + handle = self._hit_test(event.pos()) + if handle is None: + super().mousePressEvent(event) + return + self._owner._active_handle = handle + self._owner._focused_handle = handle + self._owner.setFocus(Qt.FocusReason.MouseFocusReason) + self.update() + event.accept() + + def mouseMoveEvent(self, event: QMouseEvent) -> None: # noqa: N802 + if self._owner._active_handle is None: + super().mouseMoveEvent(event) + return + value = self._x_to_value(event.pos().x()) + self._owner._drive_handle(self._owner._active_handle, value) + event.accept() + + def mouseReleaseEvent(self, event: QMouseEvent) -> None: # noqa: N802 + if event.button() != Qt.MouseButton.LeftButton: + super().mouseReleaseEvent(event) + return + if self._owner._active_handle is not None: + self._owner._active_handle = None + self.update() + event.accept() + else: + super().mouseReleaseEvent(event) + + def paintEvent(self, event: QPaintEvent) -> None: # noqa: N802 + painter = QPainter(self) + painter.setRenderHint(QPainter.RenderHint.Antialiasing) + palette = self.palette() + + left, right = self._groove_x_range() + y = self._groove_y() + groove_rect = QRect(left, y - self._GROOVE_HEIGHT // 2, right - left, self._GROOVE_HEIGHT) + painter.setPen(QPen(palette.dark().color(), 1)) + painter.setBrush(QBrush(palette.mid().color())) + painter.drawRoundedRect(groove_rect, 2, 2) + + lower_x = self._handle_x(Handle.LOWER) + upper_x = self._handle_x(Handle.UPPER) + if upper_x > lower_x: + fill_rect = QRect( + lower_x, y - self._GROOVE_HEIGHT // 2, upper_x - lower_x, self._GROOVE_HEIGHT + ) + painter.setPen(Qt.PenStyle.NoPen) + painter.setBrush(QBrush(palette.highlight().color())) + painter.drawRect(fill_rect) + + tick_top = y + self._GROOVE_HEIGHT // 2 + self._TICK_GAP + tick_bottom = tick_top + self._TICK_HEIGHT + painter.setPen(QPen(palette.dark().color(), 1)) + for i in range(11): + tx = int(numpy.rint(left + (i / 10.0) * (right - left))) + painter.drawLine(tx, tick_top, tx, tick_bottom) + + has_focus = self._owner.hasFocus() + for handle, hx in ((Handle.LOWER, lower_x), (Handle.UPPER, upper_x)): + focused = has_focus and handle is self._owner._focused_handle + painter.setPen( + QPen(palette.highlight().color(), 2) if focused else QPen(palette.dark().color(), 1) + ) + painter.setBrush(QBrush(palette.button().color())) + painter.drawEllipse(QPoint(hx, y), self._HANDLE_RADIUS, self._HANDLE_RADIUS) + + painter.end() + + +class DecimalRangeSlider(QWidget): + """Two-handle range slider over `Interval[Decimal]`. + + Public surface parallels the behaviors of :class:`DecimalSlider` + (clamping, edge-triggered signal, `ValueError` on inverted bounds) but + uses range-slider names: `get_selection` / `set_selection` / + `set_selection_and_bounds` and the `selection_changed` signal. + """ + + selection_changed = pyqtSignal(Interval) + + def __init__( + self, + parent: QWidget | None, + *, + num_ticks: int, + ) -> None: + super().__init__(parent) + self._num_ticks = num_ticks + self._bounds = Interval[Decimal](Decimal(0), Decimal(1)) + self._selection = Interval[Decimal](Decimal(0), Decimal(1)) + self._active_handle: Handle | None = None + self._focused_handle: Handle = Handle.LOWER + + self._min_label = QLabel() + self._max_label = QLabel() + self._paint_area = _RangeSliderPaintArea(self) + + layout = QHBoxLayout() + layout.setContentsMargins(0, 0, 0, 0) + layout.addWidget(self._min_label) + layout.addWidget(self._paint_area, stretch=1) + layout.addWidget(self._max_label) + self.setLayout(layout) + + self.setFocusPolicy(Qt.FocusPolicy.StrongFocus) + self._update_labels() + + @classmethod + def create_instance( + cls, + orientation: Qt.Orientation, + parent: QWidget | None = None, + *, + num_ticks: int = 1000, + ) -> DecimalRangeSlider: + if orientation != Qt.Orientation.Horizontal: + raise NotImplementedError('DecimalRangeSlider only supports horizontal orientation.') + return cls(parent, num_ticks=num_ticks) + + def get_selection(self) -> Interval[Decimal]: + return Interval[Decimal](self._selection.lower, self._selection.upper) + + def get_bounds(self) -> Interval[Decimal]: + return Interval[Decimal](self._bounds.lower, self._bounds.upper) + + def set_selection(self, selection: Interval[Decimal]) -> None: + if selection.upper < selection.lower: + raise ValueError(f'upper < lower ({selection.upper} < {selection.lower})') + new_lower = self._bounds.clamp(selection.lower) + new_upper = self._bounds.clamp(selection.upper) + if self._apply_selection(new_lower, new_upper): + self._emit_selection_changed() + + def set_selection_and_bounds( + self, + selection: Interval[Decimal], + bounds: Interval[Decimal], + block_signal: bool = False, + ) -> None: + if bounds.upper <= bounds.lower: + raise ValueError(f'maximum <= minimum ({bounds.upper} <= {bounds.lower})') + if selection.upper < selection.lower: + raise ValueError(f'upper < lower ({selection.upper} < {selection.lower})') + + self._bounds = Interval[Decimal](bounds.lower, bounds.upper) + new_lower = self._bounds.clamp(selection.lower) + new_upper = self._bounds.clamp(selection.upper) + selection_changed = self._apply_selection(new_lower, new_upper) + + self._paint_area.update() + + if selection_changed and not block_signal: + self._emit_selection_changed() + + def _apply_selection(self, new_lower: Decimal, new_upper: Decimal) -> bool: + if new_upper < new_lower: + new_upper = new_lower + changed = new_lower != self._selection.lower or new_upper != self._selection.upper + if changed: + self._selection = Interval[Decimal](new_lower, new_upper) + self._update_labels() + self._paint_area.update() + return changed + + def _drive_handle(self, handle: Handle, candidate: Decimal) -> None: + candidate = self._bounds.clamp(candidate) + if handle is Handle.LOWER: + new_lower = min(candidate, self._selection.upper) + new_upper = self._selection.upper + else: + new_lower = self._selection.lower + new_upper = max(candidate, self._selection.lower) + if self._apply_selection(new_lower, new_upper): + self._emit_selection_changed() + + def _tick_step(self) -> Decimal: + span = self._bounds.upper - self._bounds.lower + if self._num_ticks <= 0 or span == 0: + return Decimal(0) + return span / Decimal(self._num_ticks) + + def keyPressEvent(self, event: QKeyEvent) -> None: # noqa: N802 + key = event.key() + + if key in (Qt.Key.Key_BracketLeft, Qt.Key.Key_BracketRight): + new_focus = Handle.LOWER if key == Qt.Key.Key_BracketLeft else Handle.UPPER + if new_focus is not self._focused_handle: + self._focused_handle = new_focus + self._paint_area.update() + event.accept() + return + + step = self._tick_step() + current = ( + self._selection.lower if self._focused_handle is Handle.LOWER else self._selection.upper + ) + + if key in (Qt.Key.Key_Left, Qt.Key.Key_Down): + candidate = current - step + elif key in (Qt.Key.Key_Right, Qt.Key.Key_Up): + candidate = current + step + elif key == Qt.Key.Key_PageDown: + candidate = current - 10 * step + elif key == Qt.Key.Key_PageUp: + candidate = current + 10 * step + elif key == Qt.Key.Key_Home: + candidate = ( + self._bounds.lower + if self._focused_handle is Handle.LOWER + else self._selection.lower + ) + elif key == Qt.Key.Key_End: + candidate = ( + self._selection.upper + if self._focused_handle is Handle.LOWER + else self._bounds.upper + ) + else: + super().keyPressEvent(event) + return + + self._drive_handle(self._focused_handle, candidate) + event.accept() + + def focusInEvent(self, event) -> None: # noqa: N802, ANN001 + super().focusInEvent(event) + self._paint_area.update() + + def focusOutEvent(self, event) -> None: # noqa: N802, ANN001 + super().focusOutEvent(event) + self._paint_area.update() + + def _update_labels(self) -> None: + self._min_label.setText(f'{self._selection.lower:.3f}') + self._max_label.setText(f'{self._selection.upper:.3f}') + + def _emit_selection_changed(self) -> None: + self.selection_changed.emit(self.get_selection()) diff --git a/src/ptychodus/view/widgets/decimal_slider.py b/src/ptychodus/view/widgets/decimal_slider.py index f2194c2af..2f29e2eb9 100644 --- a/src/ptychodus/view/widgets/decimal_slider.py +++ b/src/ptychodus/view/widgets/decimal_slider.py @@ -58,23 +58,15 @@ def set_value_and_range( range_: Interval[Decimal], block_value_changed_signal: bool = False, ) -> None: - should_emit = False - if range_.upper <= range_.lower: raise ValueError(f'maximum <= minimum ({range_.upper} <= {range_.lower})') - if range_.lower != self._minimum: - self._minimum = range_.lower - should_emit = True + self._minimum = range_.lower + self._maximum = range_.upper - if range_.upper != self._maximum: - self._maximum = range_.upper - should_emit = True - - if self._set_value_to_slider(value): - should_emit = True + value_changed = self._set_value_to_slider(value) - if not block_value_changed_signal and should_emit: + if not block_value_changed_signal and value_changed: self._emit_value_changed() def _set_value_from_slider(self) -> None: diff --git a/src/ptychodus_store/README.md b/src/ptychodus_store/README.md new file mode 100644 index 000000000..850e76da3 --- /dev/null +++ b/src/ptychodus_store/README.md @@ -0,0 +1,167 @@ +# ptychodus-store + +FastAPI service that indexes on-disk ptychodus artifacts (campaigns, diffraction datasets, reconstruction products, fluorescence maps), exposes them through a REST API, an MCP server, and a minimal browser UI, and keeps a SQLite metadata cache in sync with the storage root via a filesystem watcher. + +- **REST**: `/api/v1/*` — list / get / render endpoints per resource kind +- **MCP**: `/mcp` — read-only tools mirroring the REST surface +- **Browser UI**: `/ui/` — six-page shell (Diffraction, Products, Positions, Probe, Object, Fluorescence) served from compiled TypeScript +- **OpenAPI**: `/openapi.json` and interactive docs at `/docs` + +## Scope of the browser UI + +The browser UI is intentionally **read-only** for this release. It browses artifacts already ingested into the storage root, previews them with the same colormap defaults as the PyQt desktop app, and offers `.h5` file downloads from each detail view. Reconstruction, settings editing, dataset ingestion, remote-compute (Globus, Genesis), fluorescence enhancement, and the automation / agent panels are only available in the desktop app (`uv run ptychodus`). Any writes to the storage root happen out of band — via the desktop app, batch runs (`uv run ptychodus -b reconstruct ...`), or the streaming processor. + +## Install + +```sh +uv sync --extra store # required +uv sync --extra store --extra xraydb # optional X-ray reference-data MCP sub-server +``` + +The shipped wheel already contains the compiled frontend. Rebuilding the UI in place is only needed when editing `ui/src/*.ts` — see [Rebuild the frontend](#rebuild-the-frontend). + +## Storage layout + +`PTYCHODUS_STORE_STORAGE_ROOT` should point at a directory laid out like this: + +```text +<storage_root>/ + campaign/<uuid>/manifest.json + diffraction/<uuid>/ + manifest.json + diffraction.h5 + product/<uuid>/ + manifest.json + product.h5 + fluorescence/<uuid>/ + manifest.json + fluorescence.h5 +``` + +Per-kind subdirectories are created on first start via `layout.ensure_kind_dirs()`. The service watches for changes to `manifest.json` files and reconciles them into the SQLite cache. Canonical definitions live in [storage/layout.py](storage/layout.py) and [storage/manifest.py](storage/manifest.py). + +## Configuration + +All settings come from `PTYCHODUS_STORE_*` environment variables (or a `.env` file in the working directory — pydantic-settings loads it automatically). See [config.py](config.py). + +| Env var | Type | Default | Purpose | +| --- | --- | --- | --- | +| `PTYCHODUS_STORE_STORAGE_ROOT` | path | *(required)* | Root of the on-disk artifact tree | +| `PTYCHODUS_STORE_DATABASE_URL` | str | `sqlite+aiosqlite:///:memory:` | Async SQLAlchemy URL — use an on-disk path for durable state | +| `PTYCHODUS_STORE_HOST` | str | `127.0.0.1` | Bind address (`0.0.0.0` to expose beyond localhost) | +| `PTYCHODUS_STORE_PORT` | int | `8000` | Bind port | +| `PTYCHODUS_STORE_LOG_LEVEL` | str | `INFO` | Python + uvicorn log level | +| `PTYCHODUS_STORE_API_PREFIX` | str | `/api/v1` | REST route prefix | +| `PTYCHODUS_STORE_MCP_MOUNT_PATH` | str | `/mcp` | MCP HTTP mount path | +| `PTYCHODUS_STORE_POLLING_INTERVAL_S` | float | `2.0` | Watchdog observer poll interval | +| `PTYCHODUS_STORE_DEBOUNCE_WINDOW_S` | float | `1.0` | Manifest-change debounce window | +| `PTYCHODUS_STORE_AUTO_RECONCILE_ON_STARTUP` | bool | `true` | Run a full rescan at boot | + +## Start + +Foreground (dev / interactive): + +```sh +PTYCHODUS_STORE_STORAGE_ROOT=/data/ptycho-store uv run ptychodus-store serve +``` + +The service prints `Uvicorn running on http://127.0.0.1:8000` when ready. Open `http://127.0.0.1:8000/` in a browser — the root redirects to `/ui/`. + +Background (quick and dirty): + +```sh +nohup uv run ptychodus-store serve > store.log 2>&1 & +echo $! > store.pid +``` + +For real deployments, use a systemd unit (see [Running in production](#running-in-production)). + +## Stop + +- **Foreground**: `Ctrl+C`. +- **Background (nohup)**: `kill $(cat store.pid)` or `pkill -f 'ptychodus-store serve'`. +- **systemd**: `systemctl stop ptychodus-store`. + +## Health check + +```sh +curl -s http://127.0.0.1:8000/api/v1/health/ +# → {"status":"ok","db":"ok","watcher":"alive"} +``` + +Wire this to your uptime monitor. `status: degraded` means one of `db` or `watcher` is not `ok`/`alive`. + +## Reindex + +If manifests were added or moved out of band, or the watcher missed events: + +```sh +PTYCHODUS_STORE_STORAGE_ROOT=/data/ptycho-store uv run ptychodus-store rebuild-index +``` + +This runs the same `full_rescan` the watcher invokes at startup. Safe to run while the service is up. + +## Logs + +Both application and uvicorn logs go to stdout, formatted as: + +```text +2026-07-20 14:23:43,857 INFO ptychodus_store.ingest.watcher: manifest watcher started on /data/ptycho-store +``` + +Raise or lower volume with `PTYCHODUS_STORE_LOG_LEVEL=DEBUG` / `WARNING`. For file logging, redirect stdout (systemd captures stdout via journald automatically). + +## Rebuild the frontend + +The wheel build (`python -m build` or `pip install .`) runs `tsc` automatically via a `build_py` hook in `setup.py`. For interactive UI development, run `tsc` directly. + +One-time setup on hosts without Node.js: + +```sh +uv tool install nodeenv +nodeenv --node=lts --prebuilt ~/.local/node-lts +export PATH="$HOME/.local/node-lts/bin:$PATH" +npm install -g typescript +``` + +Then: + +```sh +cd src/ptychodus_store/ui && tsc # one-shot +cd src/ptychodus_store/ui && tsc --watch # incremental during dev +``` + +The compiled output at `src/ptychodus_store/ui/dist/` is git-ignored. + +## Running in production + +- **Use a durable database.** The default `sqlite+aiosqlite:///:memory:` loses all state on restart. Point at a file, e.g.: + + ```sh + PTYCHODUS_STORE_DATABASE_URL=sqlite+aiosqlite:////var/lib/ptychodus-store/store.db + ``` + +- **Sample systemd unit** (`/etc/systemd/system/ptychodus-store.service`): + + ```ini + [Unit] + Description=ptychodus-store HTTP+MCP service + After=network.target + + [Service] + Type=simple + User=ptycho + WorkingDirectory=/opt/ptychodus + EnvironmentFile=/etc/ptychodus/store.env + ExecStart=/opt/ptychodus/.venv/bin/ptychodus-store serve + Restart=on-failure + RestartSec=5 + + [Install] + WantedBy=multi-user.target + ``` + + Put the `PTYCHODUS_STORE_*` variables in `/etc/ptychodus/store.env`. +- **No built-in TLS.** Front with nginx or Caddy if you're exposing beyond `127.0.0.1`. To bind all interfaces set `PTYCHODUS_STORE_HOST=0.0.0.0`. +- **Storage-root permissions.** The service user must be able to create per-kind subdirectories under `PTYCHODUS_STORE_STORAGE_ROOT` on first start. +- **No auth today.** The current service is unauthenticated; keep it behind a reverse proxy that enforces auth if the data warrants it. diff --git a/src/ptychodus_store/__init__.py b/src/ptychodus_store/__init__.py new file mode 100644 index 000000000..073622664 --- /dev/null +++ b/src/ptychodus_store/__init__.py @@ -0,0 +1,5 @@ +"""ptychodus-store: filesystem-watched catalog service for ptychodus artifacts.""" + +__all__ = ['__version__'] + +__version__ = '0.1.0' diff --git a/src/ptychodus_store/__main__.py b/src/ptychodus_store/__main__.py new file mode 100644 index 000000000..c807e3929 --- /dev/null +++ b/src/ptychodus_store/__main__.py @@ -0,0 +1,4 @@ +from ptychodus_store.cli import main + +if __name__ == '__main__': + main() diff --git a/src/ptychodus_store/app.py b/src/ptychodus_store/app.py new file mode 100644 index 000000000..0f9ef769c --- /dev/null +++ b/src/ptychodus_store/app.py @@ -0,0 +1,114 @@ +"""FastAPI app factory: wires DB, watcher, routers, and MCP server.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path + +from fastapi import FastAPI +from fastapi.responses import RedirectResponse +from fastapi.staticfiles import StaticFiles + +from ptychodus_store.config import Settings, get_settings +from ptychodus_store.db.session import SessionProvider, create_engine, create_schema +from ptychodus_store.ingest.reconciler import full_rescan +from ptychodus_store.ingest.watcher import ManifestWatcher +from ptychodus_store.mcp_server import bind_layout, bind_session_provider, create_mcp_server +from ptychodus_store.routers import ( + admin, + campaign, + diffraction, + fluorescence, + health, + lineage, + visualization, +) +from ptychodus_store.routers import product as product_router +from ptychodus_store.storage.layout import StoreLayout + +logger = logging.getLogger(__name__) + + +def _configure_logging(level: str) -> None: + logging.basicConfig( + level=getattr(logging, level.upper(), logging.INFO), + format='%(asctime)s %(levelname)s %(name)s: %(message)s', + ) + + +def create_app(settings: Settings | None = None) -> FastAPI: + settings = settings or get_settings() + _configure_logging(settings.log_level) + + layout = StoreLayout(settings.storage_root) + engine = create_engine(settings.database_url) + session_provider = SessionProvider(engine) + mcp = create_mcp_server() + + @asynccontextmanager + async def lifespan(app: FastAPI) -> AsyncIterator[None]: + await create_schema(engine) + bind_session_provider(session_provider) + bind_layout(layout) + layout.ensure_kind_dirs() + + watcher: ManifestWatcher | None = None + if settings.auto_reconcile_on_startup: + async with session_provider.session_factory() as session: + counts = await full_rescan(session, layout) + logger.info('startup reconcile counts: %s', counts) + + loop = asyncio.get_running_loop() + watcher = ManifestWatcher( + layout, + session_provider.session_factory, + loop, + polling_interval_s=settings.polling_interval_s, + debounce_window_s=settings.debounce_window_s, + ) + watcher.start() + app.state.watcher = watcher + + try: + yield + finally: + if watcher is not None: + watcher.stop() + await session_provider.dispose() + + app = FastAPI(title='ptychodus-store', version='0.1.0', lifespan=lifespan) + app.state.settings = settings + app.state.layout = layout + app.state.session_provider = session_provider + + api_prefix = settings.api_prefix + app.include_router(health.router, prefix=api_prefix) + app.include_router(campaign.router, prefix=api_prefix) + app.include_router(diffraction.router, prefix=api_prefix) + app.include_router(product_router.router, prefix=api_prefix) + app.include_router(fluorescence.router, prefix=api_prefix) + app.include_router(lineage.router, prefix=api_prefix) + app.include_router(admin.router, prefix=api_prefix) + app.include_router(visualization.router, prefix=api_prefix) + + # Mount the fastmcp HTTP app at the configured path + try: + mcp_app = mcp.http_app(path='/') + app.mount(settings.mcp_mount_path, mcp_app) + except Exception: # noqa: BLE001 + logger.exception('failed to mount MCP server; continuing without it') + + ui_dir = Path(__file__).parent / 'ui' + if ui_dir.is_dir(): + app.mount('/ui', StaticFiles(directory=ui_dir, html=True), name='ui') + + @app.get('/', include_in_schema=False) + async def _root_redirect() -> RedirectResponse: + return RedirectResponse(url='/ui/') + else: + logger.warning('ui directory not found at %s; skipping /ui mount', ui_dir) + + return app diff --git a/src/ptychodus_store/cli.py b/src/ptychodus_store/cli.py new file mode 100644 index 000000000..1663caab7 --- /dev/null +++ b/src/ptychodus_store/cli.py @@ -0,0 +1,69 @@ +"""Argparse-based CLI: serve | rebuild-index.""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import sys + +from ptychodus_store.config import get_settings + + +def _serve(_args: argparse.Namespace) -> None: + import uvicorn + + settings = get_settings() + uvicorn.run( + 'ptychodus_store.app:create_app', + host=settings.host, + port=settings.port, + factory=True, + log_level=settings.log_level.lower(), + ) + + +def _rebuild_index(_args: argparse.Namespace) -> None: + from ptychodus_store.db.session import SessionProvider, create_engine, create_schema + from ptychodus_store.ingest.reconciler import full_rescan + from ptychodus_store.storage.layout import StoreLayout + + settings = get_settings() + logging.basicConfig( + level=getattr(logging, settings.log_level.upper(), logging.INFO), + format='%(asctime)s %(levelname)s %(name)s: %(message)s', + ) + layout = StoreLayout(settings.storage_root) + engine = create_engine(settings.database_url) + provider = SessionProvider(engine) + + async def _run() -> None: + await create_schema(engine) + layout.ensure_kind_dirs() + async with provider.session_factory() as session: + counts = await full_rescan(session, layout) + await provider.dispose() + print(f'reconcile counts: {counts}') + + asyncio.run(_run()) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog='ptychodus-store') + subparsers = parser.add_subparsers(dest='cmd', required=True) + + serve = subparsers.add_parser('serve', help='run the HTTP + MCP server') + serve.set_defaults(func=_serve) + + rebuild = subparsers.add_parser( + 'rebuild-index', help='full reconciliation of the DB cache from disk' + ) + rebuild.set_defaults(func=_rebuild_index) + + args = parser.parse_args(argv if argv is not None else sys.argv[1:]) + args.func(args) + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/src/ptychodus_store/config.py b/src/ptychodus_store/config.py new file mode 100644 index 000000000..9107ec280 --- /dev/null +++ b/src/ptychodus_store/config.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + model_config = SettingsConfigDict( + env_prefix='PTYCHODUS_STORE_', + env_file='.env', + extra='ignore', + ) + + storage_root: Path = Field(...) + database_url: str = 'sqlite+aiosqlite:///:memory:' + polling_interval_s: float = 2.0 + debounce_window_s: float = 1.0 + log_level: str = 'INFO' + host: str = '127.0.0.1' + port: int = 8000 + mcp_mount_path: str = '/mcp' + api_prefix: str = '/api/v1' + auto_reconcile_on_startup: bool = True + + +@lru_cache(maxsize=1) +def get_settings() -> Settings: + return Settings() # type: ignore[call-arg] diff --git a/src/ptychodus_store/db/__init__.py b/src/ptychodus_store/db/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/ptychodus_store/db/base.py b/src/ptychodus_store/db/base.py new file mode 100644 index 000000000..c294e4d6b --- /dev/null +++ b/src/ptychodus_store/db/base.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from enum import StrEnum + +from sqlalchemy import MetaData +from sqlalchemy.orm import DeclarativeBase + +NAMING_CONVENTION = { + 'ix': 'ix_%(column_0_label)s', + 'uq': 'uq_%(table_name)s_%(column_0_name)s', + 'ck': 'ck_%(table_name)s_%(constraint_name)s', + 'fk': 'fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s', + 'pk': 'pk_%(table_name)s', +} + + +class Base(DeclarativeBase): + metadata = MetaData(naming_convention=NAMING_CONVENTION) + + +class IngestState(StrEnum): + DISCOVERED = 'DISCOVERED' + VALID = 'VALID' + INVALID = 'INVALID' + MISSING_FILES = 'MISSING_FILES' + ORPHANED = 'ORPHANED' diff --git a/src/ptychodus_store/db/models.py b/src/ptychodus_store/db/models.py new file mode 100644 index 000000000..e32343783 --- /dev/null +++ b/src/ptychodus_store/db/models.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +from datetime import datetime +from uuid import UUID + +from sqlalchemy import JSON, DateTime, Float, ForeignKey, Index, Integer, String, Uuid, func +from sqlalchemy.orm import Mapped, mapped_column + +from ptychodus_store.db.base import Base, IngestState + + +def _ts() -> Mapped[datetime]: + return mapped_column( + DateTime(timezone=True), server_default=func.current_timestamp(), nullable=False + ) + + +def _updated_ts() -> Mapped[datetime]: + return mapped_column( + DateTime(timezone=True), + server_default=func.current_timestamp(), + onupdate=func.current_timestamp(), + nullable=False, + ) + + +class Campaign(Base): + __tablename__ = 'campaign' + + uuid: Mapped[UUID] = mapped_column(Uuid(as_uuid=True, native_uuid=False), primary_key=True) + label: Mapped[str] = mapped_column(String, default='', nullable=False) + comments: Mapped[str] = mapped_column(String, default='', nullable=False) + sample_name: Mapped[str] = mapped_column(String, default='', nullable=False) + sample_description: Mapped[str] = mapped_column(String, default='', nullable=False) + tags: Mapped[list[str]] = mapped_column(JSON, default=list, nullable=False) + + folder_path: Mapped[str] = mapped_column(String, nullable=False) + manifest_mtime: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + created_from_manifest_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + ingest_state: Mapped[IngestState] = mapped_column( + String, default=IngestState.DISCOVERED, nullable=False + ) + error_message: Mapped[str | None] = mapped_column(String, nullable=True) + created_at: Mapped[datetime] = _ts() + updated_at: Mapped[datetime] = _updated_ts() + + __table_args__ = ( + Index('ix_campaign_ingest_state', 'ingest_state'), + Index('ix_campaign_sample_name', 'sample_name'), + ) + + +class Diffraction(Base): + __tablename__ = 'diffraction' + + uuid: Mapped[UUID] = mapped_column(Uuid(as_uuid=True, native_uuid=False), primary_key=True) + label: Mapped[str] = mapped_column(String, default='', nullable=False) + comments: Mapped[str] = mapped_column(String, default='', nullable=False) + campaign_uuid: Mapped[UUID | None] = mapped_column( + Uuid(as_uuid=True, native_uuid=False), + ForeignKey('campaign.uuid', ondelete='SET NULL'), + nullable=True, + ) + + # Manifest-supplied + detector_distance_m: Mapped[float | None] = mapped_column(Float, nullable=True) + probe_energy_eV: Mapped[float | None] = mapped_column(Float, nullable=True) # noqa: N815 + probe_photon_count: Mapped[int | None] = mapped_column(Integer, nullable=True) + exposure_time_s: Mapped[float | None] = mapped_column(Float, nullable=True) + tomography_angle_deg: Mapped[float | None] = mapped_column(Float, nullable=True) + tilt_angle_deg: Mapped[float | None] = mapped_column(Float, nullable=True) + polarization: Mapped[str | None] = mapped_column(String, nullable=True) + crop_center_x_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + crop_center_y_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + + # HDF5-derived + pattern_dtype: Mapped[str | None] = mapped_column(String, nullable=True) + pattern_height_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + pattern_width_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + num_patterns_total: Mapped[int | None] = mapped_column(Integer, nullable=True) + detector_pixel_width_m: Mapped[float | None] = mapped_column(Float, nullable=True) + detector_pixel_height_m: Mapped[float | None] = mapped_column(Float, nullable=True) + + # Bookkeeping + folder_path: Mapped[str] = mapped_column(String, nullable=False) + manifest_mtime: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + created_from_manifest_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + ingest_state: Mapped[IngestState] = mapped_column( + String, default=IngestState.DISCOVERED, nullable=False + ) + error_message: Mapped[str | None] = mapped_column(String, nullable=True) + created_at: Mapped[datetime] = _ts() + updated_at: Mapped[datetime] = _updated_ts() + + __table_args__ = ( + Index('ix_diffraction_campaign_uuid', 'campaign_uuid'), + Index('ix_diffraction_probe_energy_eV', 'probe_energy_eV'), + Index('ix_diffraction_tomography_angle_deg', 'tomography_angle_deg'), + Index('ix_diffraction_tilt_angle_deg', 'tilt_angle_deg'), + Index('ix_diffraction_ingest_state', 'ingest_state'), + ) + + +class Product(Base): + __tablename__ = 'product' + + uuid: Mapped[UUID] = mapped_column(Uuid(as_uuid=True, native_uuid=False), primary_key=True) + + # HDF5-derived (product.h5 is the source of truth for everything below) + name: Mapped[str | None] = mapped_column(String, nullable=True) + comments: Mapped[str | None] = mapped_column(String, nullable=True) + detector_distance_m: Mapped[float | None] = mapped_column(Float, nullable=True) + probe_energy_eV: Mapped[float | None] = mapped_column(Float, nullable=True) # noqa: N815 + probe_photon_count: Mapped[int | None] = mapped_column(Integer, nullable=True) + exposure_time_s: Mapped[float | None] = mapped_column(Float, nullable=True) + mass_attenuation_m2_kg: Mapped[float | None] = mapped_column(Float, nullable=True) + tomography_angle_deg: Mapped[float | None] = mapped_column(Float, nullable=True) + tilt_angle_deg: Mapped[float | None] = mapped_column(Float, nullable=True) + polarization: Mapped[str | None] = mapped_column(String, nullable=True) + object_layers: Mapped[int | None] = mapped_column(Integer, nullable=True) + object_height_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + object_width_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + object_pixel_width_m: Mapped[float | None] = mapped_column(Float, nullable=True) + object_pixel_height_m: Mapped[float | None] = mapped_column(Float, nullable=True) + probe_modes: Mapped[int | None] = mapped_column(Integer, nullable=True) + probe_height_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + probe_width_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + num_scan_points: Mapped[int | None] = mapped_column(Integer, nullable=True) + num_loss_epochs: Mapped[int | None] = mapped_column(Integer, nullable=True) + + # Bookkeeping + folder_path: Mapped[str] = mapped_column(String, nullable=False) + manifest_mtime: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + created_from_manifest_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + ingest_state: Mapped[IngestState] = mapped_column( + String, default=IngestState.DISCOVERED, nullable=False + ) + error_message: Mapped[str | None] = mapped_column(String, nullable=True) + created_at: Mapped[datetime] = _ts() + updated_at: Mapped[datetime] = _updated_ts() + + __table_args__ = ( + Index('ix_product_probe_energy_eV', 'probe_energy_eV'), + Index('ix_product_ingest_state', 'ingest_state'), + ) + + +class Fluorescence(Base): + __tablename__ = 'fluorescence' + + uuid: Mapped[UUID] = mapped_column(Uuid(as_uuid=True, native_uuid=False), primary_key=True) + label: Mapped[str] = mapped_column(String, default='', nullable=False) + comments: Mapped[str] = mapped_column(String, default='', nullable=False) + + # HDF5-derived + element_names: Mapped[list[str]] = mapped_column(JSON, default=list, nullable=False) + map_height_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + map_width_px: Mapped[int | None] = mapped_column(Integer, nullable=True) + + # Bookkeeping + folder_path: Mapped[str] = mapped_column(String, nullable=False) + manifest_mtime: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + created_from_manifest_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + ingest_state: Mapped[IngestState] = mapped_column( + String, default=IngestState.DISCOVERED, nullable=False + ) + error_message: Mapped[str | None] = mapped_column(String, nullable=True) + created_at: Mapped[datetime] = _ts() + updated_at: Mapped[datetime] = _updated_ts() + + __table_args__ = (Index('ix_fluorescence_ingest_state', 'ingest_state'),) + + +class DerivationEdge(Base): + """Flattens manifest `derived_from` lists for fast lineage queries. + + No DB-level FK enforcement: target spans 4 tables. Integrity is checked at + ingest time and reflected in the source row's `ingest_state`. + """ + + __tablename__ = 'derivation_edge' + + source_uuid: Mapped[UUID] = mapped_column( + Uuid(as_uuid=True, native_uuid=False), primary_key=True + ) + target_uuid: Mapped[UUID] = mapped_column( + Uuid(as_uuid=True, native_uuid=False), primary_key=True + ) + source_kind: Mapped[str] = mapped_column(String, nullable=False) + target_kind: Mapped[str] = mapped_column(String, nullable=False) + + __table_args__ = ( + Index('ix_derivation_edge_source', 'source_kind', 'source_uuid'), + Index('ix_derivation_edge_target', 'target_kind', 'target_uuid'), + ) diff --git a/src/ptychodus_store/db/repositories.py b/src/ptychodus_store/db/repositories.py new file mode 100644 index 000000000..21fcaec25 --- /dev/null +++ b/src/ptychodus_store/db/repositories.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any +from uuid import UUID + +from sqlalchemy import delete, func, select +from sqlalchemy.dialects.sqlite import insert as sqlite_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import ( + Campaign, + DerivationEdge, + Diffraction, + Fluorescence, + Product, +) +from ptychodus_store.storage.manifest import ResourceKind + +KIND_TO_MODEL: dict[str, type[Campaign | Diffraction | Product | Fluorescence]] = { + ResourceKind.CAMPAIGN: Campaign, + ResourceKind.DIFFRACTION: Diffraction, + ResourceKind.PRODUCT: Product, + ResourceKind.FLUORESCENCE: Fluorescence, +} + + +async def upsert_row(session: AsyncSession, kind: str, values: dict[str, Any]) -> None: + """Insert or update one row keyed by `values['uuid']` on the table for `kind`.""" + model = KIND_TO_MODEL[kind] + stmt = sqlite_insert(model).values(**values) + update_cols = {k: stmt.excluded[k] for k in values if k != 'uuid'} + stmt = stmt.on_conflict_do_update(index_elements=['uuid'], set_=update_cols) + await session.execute(stmt) + + +async def delete_row(session: AsyncSession, kind: str, uuid: UUID) -> None: + model = KIND_TO_MODEL[kind] + await session.execute(delete(model).where(model.uuid == uuid)) + + +async def get_row( + session: AsyncSession, kind: str, uuid: UUID +) -> Campaign | Diffraction | Product | Fluorescence | None: + model = KIND_TO_MODEL[kind] + return await session.get(model, uuid) + + +async def list_rows( + session: AsyncSession, + kind: str, + *, + limit: int, + offset: int, + where: Sequence[Any] = (), +) -> tuple[list[Campaign | Diffraction | Product | Fluorescence], int]: + model = KIND_TO_MODEL[kind] + stmt = select(model) + for clause in where: + stmt = stmt.where(clause) + stmt = stmt.order_by(model.created_at.desc()).limit(limit).offset(offset) + raw_items = (await session.execute(stmt)).scalars().all() + + count_stmt = select(func.count()).select_from(model) + for clause in where: + count_stmt = count_stmt.where(clause) + total = (await session.execute(count_stmt)).scalar_one() + items: list[Campaign | Diffraction | Product | Fluorescence] = list(raw_items) # type: ignore[arg-type] + return items, int(total) + + +async def replace_edges( + session: AsyncSession, + source_kind: str, + source_uuid: UUID, + edges: Sequence[tuple[str, UUID]], +) -> None: + """Clear all outgoing edges for `source_uuid` and reinsert from `edges`.""" + await session.execute(delete(DerivationEdge).where(DerivationEdge.source_uuid == source_uuid)) + if not edges: + return + rows = [ + { + 'source_uuid': source_uuid, + 'target_uuid': target_uuid, + 'source_kind': source_kind, + 'target_kind': target_kind, + } + for (target_kind, target_uuid) in edges + ] + await session.execute(sqlite_insert(DerivationEdge).values(rows)) + + +async def edges_referencing(session: AsyncSession, target_uuid: UUID) -> list[DerivationEdge]: + stmt = select(DerivationEdge).where(DerivationEdge.target_uuid == target_uuid) + return list((await session.execute(stmt)).scalars().all()) + + +async def outgoing_edges(session: AsyncSession, source_uuid: UUID) -> list[DerivationEdge]: + stmt = select(DerivationEdge).where(DerivationEdge.source_uuid == source_uuid) + return list((await session.execute(stmt)).scalars().all()) + + +async def row_exists(session: AsyncSession, kind: str, uuid: UUID) -> bool: + return (await get_row(session, kind, uuid)) is not None + + +async def update_state(session: AsyncSession, kind: str, uuid: UUID, state: IngestState) -> None: + row = await get_row(session, kind, uuid) + if row is not None: + row.ingest_state = state + + +async def find_kind_for_uuid(session: AsyncSession, uuid: UUID) -> str | None: + """Look up which table holds a row with this uuid (None if absent).""" + for kind, model in KIND_TO_MODEL.items(): + if await session.get(model, uuid) is not None: + return kind + return None diff --git a/src/ptychodus_store/db/session.py b/src/ptychodus_store/db/session.py new file mode 100644 index 000000000..a91b46320 --- /dev/null +++ b/src/ptychodus_store/db/session.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator + +from sqlalchemy import event +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) + +from ptychodus_store.db.base import Base + + +def _enable_sqlite_fk_pragma(engine: AsyncEngine) -> None: + """Ensure FKs are enforced on every SQLite connection.""" + + sync_engine = engine.sync_engine + if not sync_engine.dialect.name.startswith('sqlite'): + return + + @event.listens_for(sync_engine, 'connect') + def _set_pragma(dbapi_connection, _connection_record): # type: ignore[no-untyped-def] + cursor = dbapi_connection.cursor() + try: + cursor.execute('PRAGMA foreign_keys=ON') + finally: + cursor.close() + + +def create_engine(database_url: str) -> AsyncEngine: + # For in-memory SQLite we want a shared connection across the engine, not the + # default `:memory:` which gives each connection its own scratch DB. Using + # StaticPool with `check_same_thread=False` keeps the in-memory schema alive + # across requests in the same process. + connect_args: dict[str, object] = {} + engine_kwargs: dict[str, object] = {'future': True} + if database_url.startswith('sqlite+aiosqlite:///:memory:'): + from sqlalchemy.pool import StaticPool + + engine_kwargs['poolclass'] = StaticPool + connect_args['check_same_thread'] = False + engine = create_async_engine(database_url, connect_args=connect_args, **engine_kwargs) + _enable_sqlite_fk_pragma(engine) + return engine + + +def create_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]: + return async_sessionmaker(engine, expire_on_commit=False, autoflush=False, class_=AsyncSession) + + +async def create_schema(engine: AsyncEngine) -> None: + """Create all tables. Safe to call multiple times — idempotent.""" + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + +class SessionProvider: + """Holds the engine + session factory so FastAPI deps can grab a session per request.""" + + def __init__(self, engine: AsyncEngine) -> None: + self.engine = engine + self.session_factory = create_session_factory(engine) + + async def session(self) -> AsyncIterator[AsyncSession]: + async with self.session_factory() as session: + yield session + + async def dispose(self) -> None: + await self.engine.dispose() diff --git a/src/ptychodus_store/ingest/__init__.py b/src/ptychodus_store/ingest/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/ptychodus_store/ingest/pipeline.py b/src/ptychodus_store/ingest/pipeline.py new file mode 100644 index 000000000..e8f20faa3 --- /dev/null +++ b/src/ptychodus_store/ingest/pipeline.py @@ -0,0 +1,286 @@ +"""Parse → validate → upsert pipeline shared by the watcher and the reconciler.""" + +from __future__ import annotations + +import logging +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from uuid import UUID + +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.storage import h5_introspect +from ptychodus_store.storage.layout import LayoutError, StoreLayout +from ptychodus_store.storage.manifest import ( + CampaignManifest, + DiffractionManifest, + FluorescenceManifest, + ManifestLoadError, + ProductManifest, + ResourceKind, + load_manifest, +) + +logger = logging.getLogger(__name__) + + +def _mtime(path: Path) -> datetime | None: + try: + return datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc) + except OSError: + return None + + +def _common_bookkeeping( + *, + uuid: UUID, + folder: Path, + manifest_path: Path, + manifest_created_at: datetime, + state: IngestState, + error_message: str | None, +) -> dict[str, Any]: + return { + 'uuid': uuid, + 'folder_path': str(folder), + 'manifest_mtime': _mtime(manifest_path), + 'created_from_manifest_at': manifest_created_at, + 'ingest_state': state, + 'error_message': error_message, + } + + +def _values_for_campaign(m: CampaignManifest, **bookkeeping: Any) -> dict[str, Any]: + return { + **bookkeeping, + 'label': m.label, + 'comments': m.comments, + 'sample_name': m.sample_name, + 'sample_description': m.sample_description, + 'tags': list(m.tags), + } + + +def _values_for_diffraction( + m: DiffractionManifest, h5: dict[str, Any] | None, **bookkeeping: Any +) -> dict[str, Any]: + h5 = h5 or {} + pattern_shape = h5.get('pattern_shape') + return { + **bookkeeping, + 'label': m.label, + 'comments': m.comments, + 'campaign_uuid': m.campaign_uuid, + 'detector_distance_m': m.detector_distance_m, + 'probe_energy_eV': m.probe_energy_eV, + 'probe_photon_count': m.probe_photon_count, + 'exposure_time_s': m.exposure_time_s, + 'tomography_angle_deg': m.tomography_angle_deg, + 'tilt_angle_deg': m.tilt_angle_deg, + 'polarization': m.polarization.value if m.polarization is not None else None, + 'crop_center_x_px': m.crop_center_x_px, + 'crop_center_y_px': m.crop_center_y_px, + 'pattern_dtype': h5.get('pattern_dtype'), + 'pattern_height_px': pattern_shape[0] if pattern_shape else None, + 'pattern_width_px': pattern_shape[1] if pattern_shape else None, + 'num_patterns_total': h5.get('num_patterns_total'), + 'detector_pixel_width_m': h5.get('detector_pixel_width_m'), + 'detector_pixel_height_m': h5.get('detector_pixel_height_m'), + } + + +def _values_for_product( + m: ProductManifest, h5: dict[str, Any] | None, **bookkeeping: Any +) -> dict[str, Any]: + h5 = h5 or {} + obj_shape = h5.get('object_shape') or (None, None, None) + probe_shape = h5.get('probe_shape') or (None, None, None) + return { + **bookkeeping, + 'name': h5.get('name'), + 'comments': h5.get('comments'), + 'detector_distance_m': h5.get('detector_distance_m'), + 'probe_energy_eV': h5.get('probe_energy_eV'), + 'probe_photon_count': h5.get('probe_photon_count'), + 'exposure_time_s': h5.get('exposure_time_s'), + 'mass_attenuation_m2_kg': h5.get('mass_attenuation_m2_kg'), + 'tomography_angle_deg': h5.get('tomography_angle_deg'), + 'tilt_angle_deg': h5.get('tilt_angle_deg'), + 'polarization': h5.get('polarization'), + 'object_layers': obj_shape[0], + 'object_height_px': obj_shape[1], + 'object_width_px': obj_shape[2], + 'object_pixel_width_m': h5.get('object_pixel_width_m'), + 'object_pixel_height_m': h5.get('object_pixel_height_m'), + 'probe_modes': probe_shape[0], + 'probe_height_px': probe_shape[1], + 'probe_width_px': probe_shape[2], + 'num_scan_points': h5.get('num_scan_points'), + 'num_loss_epochs': h5.get('num_loss_epochs'), + } + + +def _values_for_fluorescence( + m: FluorescenceManifest, h5: dict[str, Any] | None, **bookkeeping: Any +) -> dict[str, Any]: + h5 = h5 or {} + map_shape = h5.get('map_shape') + return { + **bookkeeping, + 'label': m.label, + 'comments': m.comments, + 'element_names': list(h5.get('element_names') or []), + 'map_height_px': map_shape[0] if map_shape else None, + 'map_width_px': map_shape[1] if map_shape else None, + } + + +def _introspect( + kind: str, folder: Path, files: dict[str, str] +) -> tuple[dict[str, Any] | None, IngestState, str | None]: + """Open declared HDF5 file(s) and pull HDF5-derived metadata. + + Returns (introspected_dict, state, error_message). + """ + try: + if kind == ResourceKind.DIFFRACTION: + target = folder / files.get('diffraction', 'diffraction.h5') + if not target.is_file(): + return None, IngestState.MISSING_FILES, f'missing file: {target.name}' + return h5_introspect.introspect_diffraction(target), IngestState.VALID, None + + if kind == ResourceKind.PRODUCT: + target = folder / files.get('product', 'product.h5') + if not target.is_file(): + return None, IngestState.MISSING_FILES, f'missing file: {target.name}' + return h5_introspect.introspect_product(target), IngestState.VALID, None + + if kind == ResourceKind.FLUORESCENCE: + target = folder / files.get('fluorescence', 'fluorescence.h5') + if not target.is_file(): + return None, IngestState.MISSING_FILES, f'missing file: {target.name}' + return h5_introspect.introspect_fluorescence(target), IngestState.VALID, None + + return None, IngestState.VALID, None + except h5_introspect.IntrospectionError as exc: + return None, IngestState.INVALID, str(exc) + + +async def _reevaluate_orphan(session: AsyncSession, kind: str, uuid: UUID) -> None: + """Flip VALID↔ORPHANED for one row based on whether its outgoing edges resolve.""" + row = await repo.get_row(session, kind, uuid) + if row is None or row.ingest_state in (IngestState.INVALID, IngestState.MISSING_FILES): + return + edges = await repo.outgoing_edges(session, uuid) + has_unresolved = False + for edge in edges: + if not await repo.row_exists(session, edge.target_kind, edge.target_uuid): + has_unresolved = True + break + desired = IngestState.ORPHANED if has_unresolved else IngestState.VALID + if row.ingest_state != desired: + row.ingest_state = desired + + +async def _propagate_orphan_to_referrers(session: AsyncSession, uuid: UUID) -> None: + """When a row's existence changes, re-evaluate every row that points at it.""" + for edge in await repo.edges_referencing(session, uuid): + await _reevaluate_orphan(session, edge.source_kind, edge.source_uuid) + + +async def ingest_manifest(session: AsyncSession, layout: StoreLayout, manifest_path: Path) -> None: + """Parse a manifest at `manifest_path`, validate, introspect HDF5, upsert the row.""" + try: + location = layout.parse_manifest_path(manifest_path) + except LayoutError as exc: + logger.warning('skipping %s: %s', manifest_path, exc) + return + + try: + manifest = load_manifest( + manifest_path, expected_kind=location.kind, expected_uuid=location.uuid + ) + except ManifestLoadError as exc: + logger.warning('invalid manifest %s: %s', manifest_path, exc) + bookkeeping = _common_bookkeeping( + uuid=location.uuid, + folder=location.folder, + manifest_path=manifest_path, + manifest_created_at=datetime.now(timezone.utc), + state=IngestState.INVALID, + error_message=str(exc), + ) + # Insert a minimal placeholder so the error is visible via the API. + await _upsert_minimal(session, location.kind, bookkeeping) + return + + files = getattr(manifest, 'files', {}) or {} + h5, state, error_message = _introspect(location.kind, location.folder, files) + + bookkeeping = _common_bookkeeping( + uuid=location.uuid, + folder=location.folder, + manifest_path=manifest_path, + manifest_created_at=manifest.created_at, + state=state, + error_message=error_message, + ) + + if isinstance(manifest, CampaignManifest): + values = _values_for_campaign(manifest, **bookkeeping) + elif isinstance(manifest, DiffractionManifest): + values = _values_for_diffraction(manifest, h5, **bookkeeping) + elif isinstance(manifest, ProductManifest): + values = _values_for_product(manifest, h5, **bookkeeping) + elif isinstance(manifest, FluorescenceManifest): + values = _values_for_fluorescence(manifest, h5, **bookkeeping) + else: # pragma: no cover — discriminated union exhausted + raise AssertionError(f'unexpected manifest type: {type(manifest)!r}') + + await repo.upsert_row(session, location.kind, values) + + # Rewrite outgoing edges from manifest.derived_from + derived_from = list(getattr(manifest, 'derived_from', []) or []) + edges = [(ref.kind, ref.uuid) for ref in derived_from] + await repo.replace_edges(session, location.kind, location.uuid, edges) + + # Re-evaluate orphan state for this row and anyone pointing at it + if state == IngestState.VALID: + await _reevaluate_orphan(session, location.kind, location.uuid) + await _propagate_orphan_to_referrers(session, location.uuid) + + +async def _upsert_minimal(session: AsyncSession, kind: str, bookkeeping: dict[str, Any]) -> None: + """Upsert just the bookkeeping fields for an invalid manifest, so it's visible.""" + if kind == ResourceKind.CAMPAIGN: + values: dict[str, Any] = { + **bookkeeping, + 'label': '', + 'comments': '', + 'sample_name': '', + 'sample_description': '', + 'tags': [], + } + elif kind == ResourceKind.DIFFRACTION: + values = {**bookkeeping, 'label': '', 'comments': ''} + elif kind == ResourceKind.PRODUCT: + values = {**bookkeeping} + elif kind == ResourceKind.FLUORESCENCE: + values = {**bookkeeping, 'label': '', 'comments': '', 'element_names': []} + else: + return + await repo.upsert_row(session, kind, values) + + +async def delete_manifest(session: AsyncSession, layout: StoreLayout, manifest_path: Path) -> None: + """Handle a manifest delete event: drop the row, clear edges, propagate orphan state.""" + try: + location = layout.parse_manifest_path(manifest_path) + except LayoutError: + return + await repo.delete_row(session, location.kind, location.uuid) + await repo.replace_edges(session, location.kind, location.uuid, []) + await _propagate_orphan_to_referrers(session, location.uuid) diff --git a/src/ptychodus_store/ingest/reconciler.py b/src/ptychodus_store/ingest/reconciler.py new file mode 100644 index 000000000..e971130cd --- /dev/null +++ b/src/ptychodus_store/ingest/reconciler.py @@ -0,0 +1,55 @@ +"""Full-store rescan: ingest every manifest, drop rows whose folder is gone.""" + +from __future__ import annotations + +import logging +from uuid import UUID + +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.models import Campaign, DerivationEdge, Diffraction, Fluorescence, Product +from ptychodus_store.ingest.pipeline import ingest_manifest +from ptychodus_store.storage.layout import StoreLayout +from ptychodus_store.storage.manifest import ResourceKind + +logger = logging.getLogger(__name__) + + +async def full_rescan(session: AsyncSession, layout: StoreLayout) -> dict[str, int]: + """Walk the store, upsert every manifest, then drop rows whose folder no longer exists.""" + + found_uuids: dict[str, set[UUID]] = { + ResourceKind.CAMPAIGN: set(), + ResourceKind.DIFFRACTION: set(), + ResourceKind.PRODUCT: set(), + ResourceKind.FLUORESCENCE: set(), + } + + for manifest_path in layout.iter_manifest_paths(): + try: + location = layout.parse_manifest_path(manifest_path) + except Exception as exc: # noqa: BLE001 — defensive + logger.warning('skipping %s: %s', manifest_path, exc) + continue + await ingest_manifest(session, layout, manifest_path) + found_uuids[location.kind].add(location.uuid) + + counts = {kind: len(uuids) for kind, uuids in found_uuids.items()} + + # Sweep DB-side rows whose folder vanished + for kind, model in ( + (ResourceKind.CAMPAIGN, Campaign), + (ResourceKind.DIFFRACTION, Diffraction), + (ResourceKind.PRODUCT, Product), + (ResourceKind.FLUORESCENCE, Fluorescence), + ): + existing_uuids = set((await session.execute(select(model.uuid))).scalars().all()) + stale = existing_uuids - found_uuids[kind] + for uuid in stale: + await repo.delete_row(session, kind, uuid) + await session.execute(delete(DerivationEdge).where(DerivationEdge.source_uuid == uuid)) + + await session.commit() + return counts diff --git a/src/ptychodus_store/ingest/watcher.py b/src/ptychodus_store/ingest/watcher.py new file mode 100644 index 000000000..fd81f3007 --- /dev/null +++ b/src/ptychodus_store/ingest/watcher.py @@ -0,0 +1,162 @@ +"""PollingObserver-based watchdog thread with per-path debounce on `manifest.json`.""" + +from __future__ import annotations + +import asyncio +import logging +from pathlib import Path +from threading import Lock + +from sqlalchemy.ext.asyncio import async_sessionmaker +from watchdog.events import FileSystemEvent, FileSystemEventHandler +from watchdog.observers.polling import PollingObserver + +from ptychodus_store.ingest.pipeline import delete_manifest, ingest_manifest +from ptychodus_store.storage.layout import StoreLayout +from ptychodus_store.storage.manifest import MANIFEST_FILENAME + +logger = logging.getLogger(__name__) + + +class _ManifestEventHandler(FileSystemEventHandler): + """Collapses bursts of manifest events into a single debounced ingest per path.""" + + def __init__(self, on_upsert, on_delete, debounce_window_s: float) -> None: # type: ignore[no-untyped-def] + super().__init__() + self._on_upsert = on_upsert + self._on_delete = on_delete + self._debounce_window_s = debounce_window_s + self._pending_upserts: set[Path] = set() + self._pending_deletes: set[Path] = set() + self._lock = Lock() + + def _is_manifest(self, src_path: str) -> bool: + return Path(src_path).name == MANIFEST_FILENAME + + def _schedule_upsert(self, path: Path) -> None: + with self._lock: + self._pending_upserts.add(path) + self._pending_deletes.discard(path) + self._on_upsert(path, self._debounce_window_s) + + def _schedule_delete(self, path: Path) -> None: + with self._lock: + self._pending_deletes.add(path) + self._pending_upserts.discard(path) + self._on_delete(path, self._debounce_window_s) + + def on_created(self, event: FileSystemEvent) -> None: + if event.is_directory or not self._is_manifest(str(event.src_path)): + return + self._schedule_upsert(Path(str(event.src_path))) + + def on_modified(self, event: FileSystemEvent) -> None: + if event.is_directory or not self._is_manifest(str(event.src_path)): + return + self._schedule_upsert(Path(str(event.src_path))) + + def on_moved(self, event: FileSystemEvent) -> None: # type: ignore[override] + if event.is_directory: + return + src = str(event.src_path) + dest = str(getattr(event, 'dest_path', '') or '') + if self._is_manifest(src): + self._schedule_delete(Path(src)) + if dest and self._is_manifest(dest): + self._schedule_upsert(Path(dest)) + + def on_deleted(self, event: FileSystemEvent) -> None: + if event.is_directory or not self._is_manifest(str(event.src_path)): + return + self._schedule_delete(Path(str(event.src_path))) + + +class ManifestWatcher: + """Owns the watchdog thread and bridges file events to async ingestion.""" + + def __init__( + self, + layout: StoreLayout, + session_factory: async_sessionmaker, + loop: asyncio.AbstractEventLoop, + *, + polling_interval_s: float = 2.0, + debounce_window_s: float = 1.0, + ) -> None: + self._layout = layout + self._session_factory = session_factory + self._loop = loop + self._polling_interval_s = polling_interval_s + self._debounce_window_s = debounce_window_s + self._observer = PollingObserver(timeout=polling_interval_s) + self._handler = _ManifestEventHandler( + self._schedule_upsert, self._schedule_delete, debounce_window_s + ) + self._pending_upsert_handles: dict[Path, asyncio.TimerHandle] = {} + self._pending_delete_handles: dict[Path, asyncio.TimerHandle] = {} + self._started = False + + @property + def is_alive(self) -> bool: + return self._observer.is_alive() + + def start(self) -> None: + if self._started: + return + self._layout.ensure_kind_dirs() + self._observer.schedule(self._handler, str(self._layout.root), recursive=True) + self._observer.start() + self._started = True + logger.info('manifest watcher started on %s', self._layout.root) + + def stop(self) -> None: + if not self._started: + return + self._observer.stop() + self._observer.join(timeout=self._polling_interval_s * 2 + 1) + self._started = False + logger.info('manifest watcher stopped') + + # --- debounce + asyncio bridge --- + + def _schedule_upsert(self, path: Path, debounce: float) -> None: + self._loop.call_soon_threadsafe(self._debounce_upsert, path, debounce) + + def _schedule_delete(self, path: Path, debounce: float) -> None: + self._loop.call_soon_threadsafe(self._debounce_delete, path, debounce) + + def _debounce_upsert(self, path: Path, debounce: float) -> None: + existing = self._pending_upsert_handles.pop(path, None) + if existing is not None: + existing.cancel() + handle = self._loop.call_later( + debounce, lambda: asyncio.ensure_future(self._run_upsert(path)) + ) + self._pending_upsert_handles[path] = handle + + def _debounce_delete(self, path: Path, debounce: float) -> None: + existing = self._pending_delete_handles.pop(path, None) + if existing is not None: + existing.cancel() + handle = self._loop.call_later( + debounce, lambda: asyncio.ensure_future(self._run_delete(path)) + ) + self._pending_delete_handles[path] = handle + + async def _run_upsert(self, path: Path) -> None: + self._pending_upsert_handles.pop(path, None) + try: + async with self._session_factory() as session: + await ingest_manifest(session, self._layout, path) + await session.commit() + except Exception: # noqa: BLE001 + logger.exception('error ingesting %s', path) + + async def _run_delete(self, path: Path) -> None: + self._pending_delete_handles.pop(path, None) + try: + async with self._session_factory() as session: + await delete_manifest(session, self._layout, path) + await session.commit() + except Exception: # noqa: BLE001 + logger.exception('error deleting %s', path) diff --git a/src/ptychodus_store/mcp_server.py b/src/ptychodus_store/mcp_server.py new file mode 100644 index 000000000..5349ee8d6 --- /dev/null +++ b/src/ptychodus_store/mcp_server.py @@ -0,0 +1,568 @@ +"""fastmcp server with read-only tools mirroring the REST surface.""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path +from uuid import UUID + +from fastmcp import FastMCP +from fastmcp.exceptions import ToolError +from fastmcp.utilities.types import Image as MCPImage +from sqlalchemy import String, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus.api.diffraction import Polarization +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.io import load_diffraction_data, load_fluorescence_data, load_product +from ptychodus.api.visualization import ( + ComplexComponent, + CylindricalColorModel, + ScalarTransformation, + cyclic_colormap_names, + linear_colormap_names, +) + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import Campaign, DerivationEdge, Diffraction, Fluorescence, Product +from ptychodus_store.db.session import SessionProvider +from ptychodus_store.rendering import ( + OptionsRead, + RenderParamsError, + build_visualization_complex, + build_visualization_real, + coerce_render_params, + product_to_png_bytes, +) +from ptychodus_store.rendering.params import InvalidRenderParamError +from ptychodus_store.routers._convert import ( + campaign_to_read, + diffraction_to_read, + fluorescence_to_read, + product_to_read, +) +from ptychodus_store.routers.schemas import ( + CampaignRead, + DiffractionRead, + FluorescenceRead, + LineageNode, + LineageRead, + Page, + ProductRead, + StoreStats, +) +from ptychodus_store.storage.layout import StoreLayout +from ptychodus_store.storage.manifest import ResourceKind + +logger = logging.getLogger(__name__) + +_DEFAULT_FLUORESCENCE_PIXEL_GEOMETRY = PixelGeometry(width_m=1e-6, height_m=1e-6) + + +class _Ctx: + provider: SessionProvider | None = None + layout: StoreLayout | None = None + + +def bind_session_provider(provider: SessionProvider) -> None: + _Ctx.provider = provider + + +def bind_layout(layout: StoreLayout) -> None: + _Ctx.layout = layout + + +@asynccontextmanager +async def _session() -> AsyncIterator[AsyncSession]: + if _Ctx.provider is None: + raise RuntimeError('SessionProvider not bound; call bind_session_provider() first') + async with _Ctx.provider.session_factory() as session: + yield session + + +def _require_layout() -> StoreLayout: + if _Ctx.layout is None: + raise ToolError('StoreLayout not bound; call bind_layout() first') + return _Ctx.layout + + +def _resolve_resource_file(kind: str, uuid: UUID, filename: str) -> Path: + layout = _require_layout() + path = layout.resource_folder(kind, uuid) / filename + if not path.is_file(): + raise ToolError(f'{filename} for {kind} {uuid} not present on disk') + return path + + +async def _ensure_row_exists(kind: str, uuid: UUID) -> None: + async with _session() as session: + row = await repo.get_row(session, kind, uuid) + if row is None: + raise ToolError(f'{kind} {uuid} not found') + + +def _coerce_render_params_or_error( + colormap: str, + transform: str, + component: str | None, + color_model: str | None, + value_min: float | None, + value_max: float | None, + clip: bool, +): # -> RenderParams (typing.TYPE_CHECKING avoided to keep import surface small) + try: + return coerce_render_params( + colormap=colormap, + transform=transform, + component=component, + color_model=color_model, + value_min=value_min, + value_max=value_max, + clip=clip, + ) + except InvalidRenderParamError as exc: + raise ToolError(str(exc)) from exc + + +def create_mcp_server() -> FastMCP: + mcp = FastMCP(name='ptychodus-store') + + @mcp.tool() + async def list_campaign( + limit: int = 50, + offset: int = 0, + sample_name: str | None = None, + ingest_state: str | None = None, + ) -> Page[CampaignRead]: + """List campaigns.""" + async with _session() as session: + where = [] + if sample_name is not None: + where.append(Campaign.sample_name == sample_name) + if ingest_state is not None: + where.append(Campaign.ingest_state == ingest_state) + items, total = await repo.list_rows( + session, ResourceKind.CAMPAIGN, limit=limit, offset=offset, where=where + ) + return Page( + items=[campaign_to_read(i) for i in items], + total=total, + limit=limit, + offset=offset, + ) + + @mcp.tool() + async def get_campaign(uuid: str) -> CampaignRead | None: + """Get a campaign by UUID.""" + async with _session() as session: + row = await repo.get_row(session, ResourceKind.CAMPAIGN, UUID(uuid)) + return campaign_to_read(row) if row is not None else None + + @mcp.tool() + async def list_diffraction( + limit: int = 50, + offset: int = 0, + campaign_uuid: str | None = None, + derived_from_uuid: str | None = None, + probe_energy_eV_min: float | None = None, # noqa: N803 + probe_energy_eV_max: float | None = None, # noqa: N803 + tilt_angle_deg_min: float | None = None, + tilt_angle_deg_max: float | None = None, + polarization: str | None = None, + ingest_state: str | None = None, + ) -> Page[DiffractionRead]: + """List diffraction datasets.""" + async with _session() as session: + where = [] + if campaign_uuid is not None: + where.append(Diffraction.campaign_uuid == UUID(campaign_uuid)) + if ingest_state is not None: + where.append(Diffraction.ingest_state == ingest_state) + if probe_energy_eV_min is not None: + where.append(Diffraction.probe_energy_eV >= probe_energy_eV_min) + if probe_energy_eV_max is not None: + where.append(Diffraction.probe_energy_eV <= probe_energy_eV_max) + if tilt_angle_deg_min is not None: + where.append(Diffraction.tilt_angle_deg >= tilt_angle_deg_min) + if tilt_angle_deg_max is not None: + where.append(Diffraction.tilt_angle_deg <= tilt_angle_deg_max) + if polarization is not None: + try: + parsed_pol = Polarization(polarization) + except ValueError as exc: + raise ToolError( + f'Unknown polarization {polarization!r}; ' + f'expected one of {[p.value for p in Polarization]}.' + ) from exc + where.append(Diffraction.polarization == parsed_pol.value) + if derived_from_uuid is not None: + target = UUID(derived_from_uuid) + where.append( + Diffraction.uuid.in_( + select(DerivationEdge.source_uuid).where( + DerivationEdge.target_uuid == target + ) + ) + ) + items, total = await repo.list_rows( + session, ResourceKind.DIFFRACTION, limit=limit, offset=offset, where=where + ) + reads = [await diffraction_to_read(session, i) for i in items] # type: ignore[arg-type] + return Page(items=reads, total=total, limit=limit, offset=offset) + + @mcp.tool() + async def get_diffraction(uuid: str) -> DiffractionRead | None: + """Get a diffraction dataset by UUID.""" + async with _session() as session: + row = await repo.get_row(session, ResourceKind.DIFFRACTION, UUID(uuid)) + return await diffraction_to_read(session, row) if row is not None else None # type: ignore[arg-type] + + @mcp.tool() + async def list_product( + limit: int = 50, + offset: int = 0, + derived_from_uuid: str | None = None, + ingest_state: str | None = None, + ) -> Page[ProductRead]: + """List reconstruction products.""" + async with _session() as session: + where = [] + if ingest_state is not None: + where.append(Product.ingest_state == ingest_state) + if derived_from_uuid is not None: + target = UUID(derived_from_uuid) + where.append( + Product.uuid.in_( + select(DerivationEdge.source_uuid).where( + DerivationEdge.target_uuid == target + ) + ) + ) + items, total = await repo.list_rows( + session, ResourceKind.PRODUCT, limit=limit, offset=offset, where=where + ) + reads = [await product_to_read(session, i) for i in items] # type: ignore[arg-type] + return Page(items=reads, total=total, limit=limit, offset=offset) + + @mcp.tool() + async def get_product(uuid: str) -> ProductRead | None: + """Get a product by UUID.""" + async with _session() as session: + row = await repo.get_row(session, ResourceKind.PRODUCT, UUID(uuid)) + return await product_to_read(session, row) if row is not None else None # type: ignore[arg-type] + + @mcp.tool() + async def list_fluorescence( + limit: int = 50, + offset: int = 0, + derived_from_uuid: str | None = None, + element: str | None = None, + ) -> Page[FluorescenceRead]: + """List fluorescence datasets.""" + async with _session() as session: + where = [] + if derived_from_uuid is not None: + target = UUID(derived_from_uuid) + where.append( + Fluorescence.uuid.in_( + select(DerivationEdge.source_uuid).where( + DerivationEdge.target_uuid == target + ) + ) + ) + if element is not None: + where.append( + func.instr(func.cast(Fluorescence.element_names, String), f'"{element}"') > 0 # type: ignore[arg-type] + ) + items, total = await repo.list_rows( + session, ResourceKind.FLUORESCENCE, limit=limit, offset=offset, where=where + ) + reads = [await fluorescence_to_read(session, i) for i in items] # type: ignore[arg-type] + return Page(items=reads, total=total, limit=limit, offset=offset) + + @mcp.tool() + async def get_fluorescence(uuid: str) -> FluorescenceRead | None: + """Get a fluorescence dataset by UUID.""" + async with _session() as session: + row = await repo.get_row(session, ResourceKind.FLUORESCENCE, UUID(uuid)) + return await fluorescence_to_read(session, row) if row is not None else None # type: ignore[arg-type] + + @mcp.tool() + async def get_lineage(uuid: str) -> LineageRead | None: + """Walk the derivation DAG up and down from a node.""" + from ptychodus_store.routers.lineage import ( + _find_campaign, + _label_for, + _resolve_node, + _walk_ancestors, + _walk_descendants, + ) + + target = UUID(uuid) + async with _session() as session: + resolved = await _resolve_node(session, target) + if resolved is None: + return None + kind, row = resolved + node = LineageNode(kind=kind, uuid=target, label=_label_for(row)) # type: ignore[arg-type] + ancestors = await _walk_ancestors(session, target) + descendants = await _walk_descendants(session, target) + campaign = await _find_campaign(session, target, ancestors) + return LineageRead( + node=node, ancestors=ancestors, descendants=descendants, campaign=campaign + ) + + @mcp.tool() + async def get_store_stats() -> StoreStats: + """Return resource counts and invalid-row count for the store.""" + async with _session() as session: + + async def _count(model): # type: ignore[no-untyped-def] + return int( + (await session.execute(select(func.count()).select_from(model))).scalar_one() + ) + + invalid_total = 0 + for model in (Campaign, Diffraction, Product, Fluorescence): + stmt = ( + select(func.count()) + .select_from(model) + .where( + model.ingest_state.in_( + [IngestState.INVALID, IngestState.MISSING_FILES, IngestState.ORPHANED] + ) + ) + ) + invalid_total += int((await session.execute(stmt)).scalar_one()) + + return StoreStats( + campaign_count=await _count(Campaign), + diffraction_count=await _count(Diffraction), + product_count=await _count(Product), + fluorescence_count=await _count(Fluorescence), + invalid_count=invalid_total, + ) + + @mcp.tool() + async def get_visualization_options() -> OptionsRead: + """Enumerate valid choices for colormap / transform / component / color_model.""" + return OptionsRead( + colormaps_linear=list(linear_colormap_names()), + colormaps_cyclic=list(cyclic_colormap_names()), + transforms=[member.name.lower() for member in ScalarTransformation], + components=[member.name.lower() for member in ComplexComponent], + color_models=[member.name.lower() for member in CylindricalColorModel], + ) + + @mcp.tool() + async def render_diffraction_pattern( + uuid: str, + index: int, + colormap: str = 'gray', + transform: str = 'identity', + value_min: float | None = None, + value_max: float | None = None, + clip: bool = False, + ) -> MCPImage: + """Render a single diffraction pattern from an assembled dataset.""" + target = UUID(uuid) + await _ensure_row_exists(ResourceKind.DIFFRACTION, target) + path = _resolve_resource_file(ResourceKind.DIFFRACTION, target, 'diffraction.h5') + data = load_diffraction_data(path) + num_patterns = data.get_patterns_shape()[0] + if not 0 <= index < num_patterns: + raise ToolError(f'pattern index {index} out of range [0, {num_patterns})') + params = _coerce_render_params_or_error( + colormap, transform, None, None, value_min, value_max, clip + ) + try: + vp = build_visualization_real( + data.get_pattern(index), data.get_pixel_geometry(), params, value_label='Counts' + ) + except RenderParamsError as exc: + raise ToolError(str(exc)) from exc + return MCPImage(data=product_to_png_bytes(vp), format='png') + + @mcp.tool() + async def render_diffraction_aggregate( + uuid: str, + colormap: str = 'gray', + transform: str = 'identity', + value_min: float | None = None, + value_max: float | None = None, + clip: bool = False, + ) -> MCPImage: + """Render the mean pattern across an assembled diffraction dataset.""" + target = UUID(uuid) + await _ensure_row_exists(ResourceKind.DIFFRACTION, target) + path = _resolve_resource_file(ResourceKind.DIFFRACTION, target, 'diffraction.h5') + data = load_diffraction_data(path) + params = _coerce_render_params_or_error( + colormap, transform, None, None, value_min, value_max, clip + ) + try: + vp = build_visualization_real( + data.get_average_pattern(), + data.get_pixel_geometry(), + params, + value_label='Mean Counts', + ) + except RenderParamsError as exc: + raise ToolError(str(exc)) from exc + return MCPImage(data=product_to_png_bytes(vp), format='png') + + @mcp.tool() + async def render_probe( + uuid: str, + incoherent: int = 0, + colormap: str = 'gray', + transform: str = 'identity', + component: str | None = None, + color_model: str | None = None, + value_min: float | None = None, + value_max: float | None = None, + clip: bool = False, + ) -> MCPImage: + """Render a single incoherent probe mode. coherent axis is fixed at 0.""" + target = UUID(uuid) + await _ensure_row_exists(ResourceKind.PRODUCT, target) + path = _resolve_resource_file(ResourceKind.PRODUCT, target, 'product.h5') + product = load_product(path) + probe = product.probes.get_probe_no_opr() + if not 0 <= incoherent < probe.num_incoherent_modes: + raise ToolError( + f'incoherent mode {incoherent} out of range [0, {probe.num_incoherent_modes})' + ) + params = _coerce_render_params_or_error( + colormap, transform, component, color_model, value_min, value_max, clip + ) + try: + vp = build_visualization_complex( + probe.get_incoherent_mode(incoherent), + product.probes.get_pixel_geometry(), + params, + ) + except RenderParamsError as exc: + raise ToolError(str(exc)) from exc + return MCPImage(data=product_to_png_bytes(vp), format='png') + + @mcp.tool() + async def render_probe_modes( + uuid: str, + colormap: str = 'gray', + transform: str = 'identity', + component: str | None = None, + color_model: str | None = None, + value_min: float | None = None, + value_max: float | None = None, + clip: bool = False, + ) -> MCPImage: + """Render all incoherent probe modes tiled horizontally into a single image.""" + target = UUID(uuid) + await _ensure_row_exists(ResourceKind.PRODUCT, target) + path = _resolve_resource_file(ResourceKind.PRODUCT, target, 'product.h5') + product = load_product(path) + probe = product.probes.get_probe_no_opr() + params = _coerce_render_params_or_error( + colormap, transform, component, color_model, value_min, value_max, clip + ) + try: + vp = build_visualization_complex( + probe.get_incoherent_modes_flattened(), + product.probes.get_pixel_geometry(), + params, + ) + except RenderParamsError as exc: + raise ToolError(str(exc)) from exc + return MCPImage(data=product_to_png_bytes(vp), format='png') + + @mcp.tool() + async def render_object_layer( + uuid: str, + layer: int, + colormap: str = 'gray', + transform: str = 'identity', + component: str | None = None, + color_model: str | None = None, + value_min: float | None = None, + value_max: float | None = None, + clip: bool = False, + ) -> MCPImage: + """Render a single object layer as a complex-valued image.""" + target = UUID(uuid) + await _ensure_row_exists(ResourceKind.PRODUCT, target) + path = _resolve_resource_file(ResourceKind.PRODUCT, target, 'product.h5') + product = load_product(path) + if not 0 <= layer < product.object_.num_layers: + raise ToolError(f'object layer {layer} out of range [0, {product.object_.num_layers})') + params = _coerce_render_params_or_error( + colormap, transform, component, color_model, value_min, value_max, clip + ) + try: + vp = build_visualization_complex( + product.object_.get_layer(layer), + product.object_.get_pixel_geometry(), + params, + ) + except RenderParamsError as exc: + raise ToolError(str(exc)) from exc + return MCPImage(data=product_to_png_bytes(vp), format='png') + + @mcp.tool() + async def render_fluorescence_element( + uuid: str, + name: str, + product_uuid: str | None = None, + colormap: str = 'gray', + transform: str = 'identity', + value_min: float | None = None, + value_max: float | None = None, + clip: bool = False, + ) -> MCPImage: + """Render one element map from a fluorescence dataset.""" + target = UUID(uuid) + await _ensure_row_exists(ResourceKind.FLUORESCENCE, target) + path = _resolve_resource_file(ResourceKind.FLUORESCENCE, target, 'fluorescence.h5') + dataset = load_fluorescence_data(path) + matches = [emap for emap in dataset.element_maps if emap.name == name] + if not matches: + available = ', '.join(emap.name for emap in dataset.element_maps) + raise ToolError(f'element {name!r} not found; available: [{available}]') + emap = matches[0] + + if product_uuid is not None: + product_target = UUID(product_uuid) + await _ensure_row_exists(ResourceKind.PRODUCT, product_target) + product_path = _resolve_resource_file( + ResourceKind.PRODUCT, product_target, 'product.h5' + ) + pixel_geometry = load_product(product_path).object_.get_pixel_geometry() + else: + pixel_geometry = _DEFAULT_FLUORESCENCE_PIXEL_GEOMETRY + + params = _coerce_render_params_or_error( + colormap, transform, None, None, value_min, value_max, clip + ) + try: + vp = build_visualization_real( + emap.counts_per_second, + pixel_geometry, + params, + value_label=f'{emap.name} counts/s', + ) + except RenderParamsError as exc: + raise ToolError(str(exc)) from exc + return MCPImage(data=product_to_png_bytes(vp), format='png') + + try: + from ptychodus_store.mcp_tools.xraydb import create_xraydb_mcp + + mcp.mount(create_xraydb_mcp(), namespace='xraydb') + logger.info('mounted xraydb MCP sub-server (tools namespaced as xraydb_*)') + except ImportError: + logger.debug('xraydb not installed; xraydb MCP tools disabled') + + return mcp diff --git a/src/ptychodus_store/mcp_tools/__init__.py b/src/ptychodus_store/mcp_tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/ptychodus_store/mcp_tools/xraydb.py b/src/ptychodus_store/mcp_tools/xraydb.py new file mode 100644 index 000000000..f12ca550d --- /dev/null +++ b/src/ptychodus_store/mcp_tools/xraydb.py @@ -0,0 +1,385 @@ +"""Optional MCP sub-server exposing xraydb reference data. + +Mounted with namespace 'xraydb' by :func:`ptychodus_store.mcp_server.create_mcp_server` +when the ``xraydb`` optional extra is installed. Import-guarded: ``import xraydb`` at +module top means this module fails to import when the extra is absent, and the parent +skips mounting. +""" + +from __future__ import annotations + +import asyncio +from typing import Literal + +import xraydb +from fastmcp import FastMCP +from fastmcp.exceptions import ToolError +from pydantic import BaseModel + + +class ElementInfo(BaseModel): + symbol: str + atomic_number: int + atomic_mass_amu: float + density_g_per_cm3: float + + +class XrayEdgeInfo(BaseModel): + edge: str + energy_eV: float # noqa: N815 + fluorescence_yield: float + jump_ratio: float + + +class XrayLineInfo(BaseModel): + siegbahn_name: str + energy_eV: float # noqa: N815 + intensity: float + initial_level: str + final_level: str + + +class AnomalousFactorPoint(BaseModel): + energy_eV: float # noqa: N815 + f1_electrons: float + f2_electrons: float + + +class MassAttenuationPoint(BaseModel): + energy_eV: float # noqa: N815 + mu_over_rho_cm2_per_g: float + + +class LinearAttenuationPoint(BaseModel): + energy_eV: float # noqa: N815 + mu_per_m: float + + +class RefractiveIndexInfo(BaseModel): + delta: float + beta: float + attenuation_length_m: float + + +class FluorescenceYieldInfo(BaseModel): + effective_yield: float + weighted_energy_eV: float # noqa: N815 + net_probability: float + + +class ElementIdentification(BaseModel): + symbol: str + edge: str + + +class NamedMaterialInfo(BaseModel): + name: str + formula: str + density_g_per_cm3: float + categories: list[str] + + +def _resolve_element(element: str | int) -> str: + try: + z = xraydb.atomic_number(str(element)) if isinstance(element, str) else int(element) + return xraydb.atomic_symbol(z) + except (ValueError, KeyError) as exc: + raise ToolError(f'unknown element: {element!r}') from exc + + +async def _run(func, /, *args, **kwargs): # type: ignore[no-untyped-def] + return await asyncio.to_thread(func, *args, **kwargs) + + +def create_xraydb_mcp() -> FastMCP: + """Build the xraydb MCP sub-server. Caller mounts it with ``namespace='xraydb'``.""" + mcp: FastMCP = FastMCP(name='xraydb') + + @mcp.tool() + async def element_info(element: str | int) -> ElementInfo: + """Return atomic number, symbol, atomic mass (AMU), and elemental density (g/cm^3). + + Args: + element: Chemical symbol (case-sensitive, e.g. 'Fe') or atomic number. + """ + symbol = _resolve_element(element) + z = await _run(xraydb.atomic_number, symbol) + mass = await _run(xraydb.atomic_mass, symbol) + density = await _run(xraydb.atomic_density, symbol) + return ElementInfo( + symbol=symbol, + atomic_number=int(z), + atomic_mass_amu=float(mass), + density_g_per_cm3=float(density), + ) + + @mcp.tool() + async def xray_absorption_edges(element: str | int) -> list[XrayEdgeInfo]: + """List all tabulated absorption edges (K, L1-L3, M1-M5, ...) for an element. + + Each entry gives edge energy in eV, fluorescence yield (0-1), and edge-jump ratio. + + Args: + element: Chemical symbol or atomic number. + """ + symbol = _resolve_element(element) + edges = await _run(xraydb.xray_edges, symbol) + return [ + XrayEdgeInfo( + edge=name, + energy_eV=float(row.energy), + fluorescence_yield=float(row.fyield), + jump_ratio=float(row.jump_ratio), + ) + for name, row in edges.items() + ] + + @mcp.tool() + async def xray_emission_lines( + element: str | int, + excitation_energy_eV: float | None = None, # noqa: N803 + ) -> list[XrayLineInfo]: + """List characteristic X-ray fluorescence lines for an element. + + Returns Siegbahn-labelled lines (Ka1, Kb1, La1, ...) with emission energy in eV + and normalized intensity within their initial-level manifold. If + ``excitation_energy_eV`` is provided, only lines that can be excited at or below + that energy are returned. + + Args: + element: Chemical symbol or atomic number. + excitation_energy_eV: Optional excitation energy in eV. + """ + symbol = _resolve_element(element) + lines = await _run(xraydb.xray_lines, symbol, None, excitation_energy_eV) + return [ + XrayLineInfo( + siegbahn_name=name, + energy_eV=float(row.energy), + intensity=float(row.intensity), + initial_level=str(row.initial_level), + final_level=str(row.final_level), + ) + for name, row in lines.items() + ] + + @mcp.tool() + async def anomalous_scattering_factors( + element: str | int, + energies_eV: list[float], # noqa: N803 + ) -> list[AnomalousFactorPoint]: + """Real (f1) and imaginary (f2) anomalous atomic scattering factors, in electrons. + + Uses the Chantler tables. Useful for anomalous / resonant ptychography and for + contrast calculations near an absorption edge. + + Args: + element: Chemical symbol or atomic number. + energies_eV: List of photon energies in eV. Must be non-empty. + """ + if not energies_eV: + raise ToolError('energies_eV must contain at least one entry') + symbol = _resolve_element(element) + f1 = await _run(xraydb.f1_chantler, symbol, energies_eV) + f2 = await _run(xraydb.f2_chantler, symbol, energies_eV) + return [ + AnomalousFactorPoint( + energy_eV=float(e), + f1_electrons=float(a), + f2_electrons=float(b), + ) + for e, a, b in zip(energies_eV, f1, f2, strict=True) + ] + + @mcp.tool() + async def mass_attenuation_coefficient( + element: str | int, + energies_eV: list[float], # noqa: N803 + kind: Literal['total', 'photo', 'coh', 'incoh'] = 'total', + ) -> list[MassAttenuationPoint]: + """Mass attenuation coefficient mu/rho in cm^2/g for a pure element (Elam tables). + + Args: + element: Chemical symbol or atomic number. + energies_eV: List of photon energies in eV. Must be non-empty. + kind: Cross-section component: 'total' (default), 'photo' (photoabsorption), + 'coh' (coherent scatter), or 'incoh' (incoherent scatter). + """ + if not energies_eV: + raise ToolError('energies_eV must contain at least one entry') + symbol = _resolve_element(element) + mu = await _run(xraydb.mu_elam, symbol, energies_eV, kind) + return [ + MassAttenuationPoint(energy_eV=float(e), mu_over_rho_cm2_per_g=float(v)) + for e, v in zip(energies_eV, mu, strict=True) + ] + + @mcp.tool() + async def material_refractive_index( + formula: str, + density_g_per_cm3: float, + energy_eV: float, # noqa: N803 + ) -> RefractiveIndexInfo: + """X-ray refractive index n = 1 - delta - i*beta at a single photon energy. + + Also returns the 1/e X-ray attenuation length in meters. Use this for phase-object + design and thickness estimates. + + Args: + formula: Chemical formula (case-sensitive, e.g. 'SiO2', 'C22H10N2O5') or a + named material from ``list_named_materials``. + density_g_per_cm3: Bulk density of the material in g/cm^3. + energy_eV: Photon energy in eV. Scalar only (xraydb limitation). + """ + try: + delta, beta, atlen_cm = await _run( + xraydb.xray_delta_beta, formula, density_g_per_cm3, energy_eV + ) + except (ValueError, KeyError) as exc: + raise ToolError(f'xray_delta_beta failed for {formula!r}: {exc}') from exc + return RefractiveIndexInfo( + delta=float(delta), + beta=float(beta), + attenuation_length_m=float(atlen_cm) * 1e-2, + ) + + @mcp.tool() + async def material_attenuation( + name_or_formula: str, + energies_eV: list[float], # noqa: N803 + density_g_per_cm3: float | None = None, + kind: Literal['total', 'photo'] = 'total', + ) -> list[LinearAttenuationPoint]: + """Linear X-ray attenuation coefficient mu in 1/m for a compound or named material. + + The 1/e attenuation length is 1 / mu_per_m (meters). + + Args: + name_or_formula: Chemical formula (case-sensitive) or a named material from + ``list_named_materials`` (case-insensitive). For named materials, density + may be omitted. + energies_eV: List of photon energies in eV. Must be non-empty. + density_g_per_cm3: Bulk density in g/cm^3. Optional if a named material + supplies its own density. + kind: 'total' (default) or 'photo' (photoabsorption only). + """ + if not energies_eV: + raise ToolError('energies_eV must contain at least one entry') + try: + mu = await _run( + xraydb.material_mu, name_or_formula, energies_eV, density_g_per_cm3, kind + ) + except (ValueError, KeyError) as exc: + raise ToolError(f'material_mu failed for {name_or_formula!r}: {exc}') from exc + return [ + LinearAttenuationPoint(energy_eV=float(e), mu_per_m=float(v) * 100.0) + for e, v in zip(energies_eV, mu, strict=True) + ] + + @mcp.tool() + async def effective_fluorescence_yield( + element: str | int, + edge: str, + line: str, + excitation_energy_eV: float, # noqa: N803 + ) -> FluorescenceYieldInfo: + """Effective fluorescence yield for one emission line under a given excitation. + + Accounts for Coster-Kronig cascades within the L or M manifolds. Returns the + effective yield, an intensity-weighted mean line energy in eV, and the net + probability that the excitation produces this line. + + Args: + element: Chemical symbol or atomic number. + edge: IUPAC edge name (e.g. 'K', 'L3'). + line: Siegbahn line name (e.g. 'Ka1', 'La1'). + excitation_energy_eV: Excitation photon energy in eV. Must be at or above the + edge energy for a non-zero result. + """ + symbol = _resolve_element(element) + try: + fyield, weighted_energy, net_prob = await _run( + xraydb.fluor_yield, symbol, edge, line, excitation_energy_eV + ) + except (ValueError, KeyError) as exc: + raise ToolError( + f'fluor_yield failed for {symbol!r} edge={edge!r} line={line!r}: {exc}' + ) from exc + return FluorescenceYieldInfo( + effective_yield=float(fyield), + weighted_energy_eV=float(weighted_energy), + net_probability=float(net_prob), + ) + + @mcp.tool() + async def identify_element_by_edge_energy( + observed_energy_eV: float, # noqa: N803 + edges: list[str] | None = None, + ) -> ElementIdentification: + """Best-guess element and absorption edge for an observed edge energy. + + Useful for XANES / EXAFS feature identification. Not intended for identifying + fluorescence emission lines - use ``xray_emission_lines`` for that. + + Args: + observed_energy_eV: Approximate edge energy in eV. + edges: Edges to consider. Defaults to ('K', 'L3', 'L2', 'L1', 'M5'). + """ + edges_tuple = tuple(edges) if edges else ('K', 'L3', 'L2', 'L1', 'M5') + try: + symbol, edge = await _run(xraydb.guess_edge, observed_energy_eV, edges_tuple) + except (ValueError, KeyError) as exc: + raise ToolError(f'guess_edge failed: {exc}') from exc + return ElementIdentification(symbol=str(symbol), edge=str(edge)) + + @mcp.tool() + async def core_hole_width(element: str | int, edge: str | None = None) -> float: + """Natural core-hole linewidth in eV. Returns the K-edge width if edge omitted. + + Args: + element: Chemical symbol or atomic number. + edge: Optional IUPAC edge name (e.g. 'L3'). If omitted, returns the K width. + """ + symbol = _resolve_element(element) + try: + width = await _run(xraydb.core_width, symbol, edge) + except (ValueError, KeyError) as exc: + raise ToolError(f'core_width failed for {symbol!r}: {exc}') from exc + return float(width) + + @mcp.tool() + async def parse_chemical_formula(formula: str) -> dict[str, float]: + """Parse a chemical formula into an {element_symbol: stoichiometric_count} map. + + Example: 'SiO2' -> {'Si': 1.0, 'O': 2.0}. Case-sensitive. + + Args: + formula: Chemical formula string. + """ + try: + parsed = await _run(xraydb.chemparse, formula) + except (ValueError, KeyError) as exc: + raise ToolError(f'chemparse failed for {formula!r}: {exc}') from exc + return {str(k): float(v) for k, v in parsed.items()} + + @mcp.tool() + async def list_named_materials() -> list[NamedMaterialInfo]: + """Enumerate xraydb's built-in named materials (e.g. 'kapton', 'water', 'silicon'). + + The returned names are valid inputs to ``material_refractive_index`` and + ``material_attenuation``. Includes chemical formula, density, and categories. + """ + from xraydb.materials import get_materials + + mats = await _run(get_materials) + return [ + NamedMaterialInfo( + name=str(name), + formula=str(mat.formula), + density_g_per_cm3=float(mat.density), + categories=list(mat.categories), + ) + for name, mat in mats.items() + ] + + return mcp diff --git a/src/ptychodus_store/rendering/__init__.py b/src/ptychodus_store/rendering/__init__.py new file mode 100644 index 000000000..1fc5563a0 --- /dev/null +++ b/src/ptychodus_store/rendering/__init__.py @@ -0,0 +1,32 @@ +"""Reusable helpers for the visualization endpoints exposed by ptychodus_store.""" + +from __future__ import annotations + +from ptychodus_store.rendering.params import ( + RenderParams, + coerce_render_params, + render_params_dep, +) +from ptychodus_store.rendering.render import ( + RenderParamsError, + build_visualization_complex, + build_visualization_real, + product_to_png_bytes, + render_complex, + render_real, +) +from ptychodus_store.rendering.schemas import OptionsRead, RenderedImage + +__all__ = [ + 'OptionsRead', + 'RenderParams', + 'RenderParamsError', + 'RenderedImage', + 'build_visualization_complex', + 'build_visualization_real', + 'coerce_render_params', + 'product_to_png_bytes', + 'render_complex', + 'render_params_dep', + 'render_real', +] diff --git a/src/ptychodus_store/rendering/params.py b/src/ptychodus_store/rendering/params.py new file mode 100644 index 000000000..f8a705f22 --- /dev/null +++ b/src/ptychodus_store/rendering/params.py @@ -0,0 +1,127 @@ +"""Shared query-parameter dependency for the visualization endpoints.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Annotated, Any, TypeVar + +from fastapi import Depends, HTTPException, Query + +from ptychodus.api.visualization import ( + ComplexComponent, + CylindricalColorModel, + ScalarTransformation, + cyclic_colormap_names, + linear_colormap_names, +) + +E = TypeVar('E', bound=Enum) + + +class InvalidRenderParamError(ValueError): + """Raised when a raw render-param string cannot be coerced to its enum / colormap.""" + + def __init__(self, message: str, detail: dict[str, Any]) -> None: + super().__init__(message) + self.detail = detail + + +@dataclass(frozen=True) +class RenderParams: + """Parsed and validated visualization parameters shared by every render endpoint.""" + + colormap: str + transform: ScalarTransformation + component: ComplexComponent | None + color_model: CylindricalColorModel | None + value_min: float | None + value_max: float | None + clip: bool + + +def _coerce_enum(name: str, enum_cls: type[E], arg_name: str) -> E: + normalized = name.strip().upper().replace('-', '_') + try: + return enum_cls[normalized] + except KeyError as exc: + valid = [member.name.lower() for member in enum_cls] + raise InvalidRenderParamError( + f'invalid {arg_name}: {name!r}', + detail={'error': f'invalid {arg_name}: {name!r}', 'valid': valid}, + ) from exc + + +def _coerce_colormap(name: str) -> str: + valid_linear = set(linear_colormap_names()) + valid_cyclic = set(cyclic_colormap_names()) + if name in valid_linear or name in valid_cyclic: + return name + raise InvalidRenderParamError( + f'invalid colormap: {name!r}', + detail={ + 'error': f'invalid colormap: {name!r}', + 'valid_linear': sorted(valid_linear), + 'valid_cyclic': sorted(valid_cyclic), + }, + ) + + +def coerce_render_params( + colormap: str, + transform: str, + component: str | None, + color_model: str | None, + value_min: float | None, + value_max: float | None, + clip: bool, +) -> RenderParams: + """Coerce raw values into a :class:`RenderParams`; raises :class:`InvalidRenderParamError`.""" + return RenderParams( + colormap=_coerce_colormap(colormap), + transform=_coerce_enum(transform, ScalarTransformation, 'transform'), + component=_coerce_enum(component, ComplexComponent, 'component') + if component is not None + else None, + color_model=_coerce_enum(color_model, CylindricalColorModel, 'color_model') + if color_model is not None + else None, + value_min=value_min, + value_max=value_max, + clip=clip, + ) + + +def render_params_dep( + colormap: Annotated[ + str, Query(description='Colormap name from /visualization/options.') + ] = 'gray', + transform: Annotated[str, Query(description='Scalar transformation enum name.')] = 'identity', + component: Annotated[ + str | None, Query(description='ComplexComponent enum name; for complex arrays only.') + ] = None, + color_model: Annotated[ + str | None, + Query(description='CylindricalColorModel enum name; for complex arrays only.'), + ] = None, + value_min: Annotated[float | None, Query(description='Color-axis lower bound.')] = None, + value_max: Annotated[float | None, Query(description='Color-axis upper bound.')] = None, + clip: Annotated[ + bool, Query(description='Clip out-of-range values to the axis bounds.') + ] = False, +) -> RenderParams: + try: + return coerce_render_params( + colormap=colormap, + transform=transform, + component=component, + color_model=color_model, + value_min=value_min, + value_max=value_max, + clip=clip, + ) + except InvalidRenderParamError as exc: + raise HTTPException(status_code=400, detail=exc.detail) from exc + + +RenderParamsDep = Annotated[RenderParams, Depends(render_params_dep)] diff --git a/src/ptychodus_store/rendering/render.py b/src/ptychodus_store/rendering/render.py new file mode 100644 index 000000000..a21311815 --- /dev/null +++ b/src/ptychodus_store/rendering/render.py @@ -0,0 +1,142 @@ +"""Adapters from `ptychodus.api.visualization` to the store's HTTP and MCP response shapes.""" + +from __future__ import annotations + +import base64 +from io import BytesIO + +import numpy +from fastapi import HTTPException +from PIL import Image + +from ptychodus.api.common import ComplexArrayType, NumberArrayType, RealArrayType +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.visualization import ( + ComplexComponent, + VisualizationProduct, + visualize_complex_component, + visualize_complex_values, + visualize_real_values, +) + +from ptychodus_store.rendering.params import RenderParams +from ptychodus_store.rendering.schemas import RenderedImage + + +class RenderParamsError(ValueError): + """Raised when the RenderParams are inconsistent with the array being rendered.""" + + +def build_visualization_real( + values: NumberArrayType, + pixel_geometry: PixelGeometry, + params: RenderParams, + value_label: str, +) -> VisualizationProduct: + """Build a VisualizationProduct from a real-valued 2D array.""" + if params.component is not None or params.color_model is not None: + raise RenderParamsError( + 'component and color_model are only valid for complex-valued arrays.' + ) + + values_real: RealArrayType = numpy.asarray(values, dtype=numpy.float64) + return visualize_real_values( + value_label=value_label, + values=values_real, + pixel_geometry=pixel_geometry, + colormap=params.colormap, + transform=params.transform, + value_min=params.value_min, + value_max=params.value_max, + clip=params.clip, + ) + + +def build_visualization_complex( + values: ComplexArrayType, + pixel_geometry: PixelGeometry, + params: RenderParams, +) -> VisualizationProduct: + """Build a VisualizationProduct from a complex-valued 2D array. + + Dispatches on the provided render params: `color_model` → cylindrical encoding; + `component` → single scalar component; neither → defaults to amplitude. + """ + if params.component is not None and params.color_model is not None: + raise RenderParamsError('component and color_model are mutually exclusive.') + + if params.color_model is not None: + return visualize_complex_values( + values=values, + pixel_geometry=pixel_geometry, + model=params.color_model, + amplitude_transform=params.transform, + value_min=params.value_min, + value_max=params.value_max, + clip=params.clip, + ) + + component = params.component or ComplexComponent.AMPLITUDE + return visualize_complex_component( + values=values, + pixel_geometry=pixel_geometry, + component=component, + colormap=params.colormap, + transform=params.transform, + value_min=params.value_min, + value_max=params.value_max, + clip=params.clip, + ) + + +def product_to_png_bytes(vp: VisualizationProduct) -> bytes: + """Encode a VisualizationProduct's RGBA image as PNG bytes.""" + rgba_uint8 = numpy.clip(vp.get_image_rgba() * 255.0, 0.0, 255.0).astype(numpy.uint8) + image = Image.fromarray(rgba_uint8, mode='RGBA') + buf = BytesIO() + image.save(buf, format='PNG') + return buf.getvalue() + + +def _product_to_response(vp: VisualizationProduct) -> RenderedImage: + png_bytes = product_to_png_bytes(vp) + pixel_geometry = vp.get_pixel_geometry() + color_range = vp.get_color_value_range() + rgba = vp.get_image_rgba() + return RenderedImage( + png_base64=base64.b64encode(png_bytes).decode('ascii'), + value_label=vp.get_value_label(), + color_value_min=float(color_range.lower), + color_value_max=float(color_range.upper), + pixel_width_m=float(pixel_geometry.width_m), + pixel_height_m=float(pixel_geometry.height_m), + shape_h_px=int(rgba.shape[0]), + shape_w_px=int(rgba.shape[1]), + ) + + +def render_real( + values: NumberArrayType, + pixel_geometry: PixelGeometry, + params: RenderParams, + value_label: str, +) -> RenderedImage: + """Render a real-valued (or integer) 2D array. Rejects component / color_model (400).""" + try: + vp = build_visualization_real(values, pixel_geometry, params, value_label) + except RenderParamsError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return _product_to_response(vp) + + +def render_complex( + values: ComplexArrayType, + pixel_geometry: PixelGeometry, + params: RenderParams, +) -> RenderedImage: + """Render a complex-valued 2D array. Rejects component+color_model together (400).""" + try: + vp = build_visualization_complex(values, pixel_geometry, params) + except RenderParamsError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return _product_to_response(vp) diff --git a/src/ptychodus_store/rendering/schemas.py b/src/ptychodus_store/rendering/schemas.py new file mode 100644 index 000000000..e8f88657a --- /dev/null +++ b/src/ptychodus_store/rendering/schemas.py @@ -0,0 +1,35 @@ +"""Pydantic response models for the visualization endpoints.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + + +class RenderedImage(BaseModel): + """A rendered image plus the metadata a front-end needs to draw a colorbar.""" + + png_base64: str = Field(description='Base64-encoded PNG bytes.') + mime_type: Literal['image/png'] = 'image/png' + value_label: str = Field(description='LaTeX-decorated label from the underlying transform.') + color_value_min: float = Field( + description='Lower bound of the color axis, in transformed units.' + ) + color_value_max: float = Field( + description='Upper bound of the color axis, in transformed units.' + ) + pixel_width_m: float = Field(description='Physical pixel width in meters.') + pixel_height_m: float = Field(description='Physical pixel height in meters.') + shape_h_px: int = Field(description='Rendered image height in pixels.') + shape_w_px: int = Field(description='Rendered image width in pixels.') + + +class OptionsRead(BaseModel): + """Enumeration of valid choices a client may pass to visualization endpoints.""" + + colormaps_linear: list[str] + colormaps_cyclic: list[str] + transforms: list[str] + components: list[str] + color_models: list[str] diff --git a/src/ptychodus_store/routers/__init__.py b/src/ptychodus_store/routers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/ptychodus_store/routers/_convert.py b/src/ptychodus_store/routers/_convert.py new file mode 100644 index 000000000..f274a37f2 --- /dev/null +++ b/src/ptychodus_store/routers/_convert.py @@ -0,0 +1,44 @@ +"""Helpers to convert ORM rows + edges into API read models.""" + +from __future__ import annotations + +from uuid import UUID + +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.models import Diffraction, Fluorescence, Product +from ptychodus_store.routers.schemas import ( + CampaignRead, + DerivedFromEdge, + DiffractionRead, + FluorescenceRead, + ProductRead, +) + + +async def _edges_for(session: AsyncSession, uuid: UUID) -> list[DerivedFromEdge]: + edges = await repo.outgoing_edges(session, uuid) + return [DerivedFromEdge(kind=e.target_kind, uuid=e.target_uuid) for e in edges] # type: ignore[arg-type] + + +async def diffraction_to_read(session: AsyncSession, row: Diffraction) -> DiffractionRead: + data = DiffractionRead.model_validate(row) + data.derived_from = await _edges_for(session, row.uuid) + return data + + +async def product_to_read(session: AsyncSession, row: Product) -> ProductRead: + data = ProductRead.model_validate(row) + data.derived_from = await _edges_for(session, row.uuid) + return data + + +async def fluorescence_to_read(session: AsyncSession, row: Fluorescence) -> FluorescenceRead: + data = FluorescenceRead.model_validate(row) + data.derived_from = await _edges_for(session, row.uuid) + return data + + +def campaign_to_read(row: object) -> CampaignRead: + return CampaignRead.model_validate(row) diff --git a/src/ptychodus_store/routers/admin.py b/src/ptychodus_store/routers/admin.py new file mode 100644 index 000000000..39c99639c --- /dev/null +++ b/src/ptychodus_store/routers/admin.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio +import logging +import uuid as uuid_lib +from uuid import UUID + +from fastapi import APIRouter, BackgroundTasks, HTTPException + +from ptychodus_store.ingest.pipeline import ingest_manifest +from ptychodus_store.ingest.reconciler import full_rescan +from ptychodus_store.routers.deps import LayoutDep, get_session_provider +from ptychodus_store.routers.schemas import ReindexResponse, StoreStats +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import Campaign, Diffraction, Fluorescence, Product +from ptychodus_store.routers.deps import SessionDep +from sqlalchemy import func, or_, select +from fastapi import Request + +logger = logging.getLogger(__name__) +router = APIRouter(prefix='/admin', tags=['admin']) + + +async def _do_rescan(session_provider, layout) -> None: # type: ignore[no-untyped-def] + try: + async with session_provider.session_factory() as session: + await full_rescan(session, layout) + except Exception: # noqa: BLE001 + logger.exception('background reindex failed') + + +@router.post('/reindex', response_model=ReindexResponse) +async def reindex( + request: Request, background_tasks: BackgroundTasks, layout: LayoutDep +) -> ReindexResponse: + job_id = str(uuid_lib.uuid4()) + session_provider = get_session_provider(request) + background_tasks.add_task(asyncio.create_task, _do_rescan(session_provider, layout)) + return ReindexResponse(status='accepted', job_id=job_id) + + +@router.post('/rescan/{kind}/{uuid}') +async def rescan_one( + kind: str, uuid: UUID, session: SessionDep, layout: LayoutDep +) -> dict[str, str]: + folder = layout.resource_folder(kind, uuid) + manifest_path = folder / 'manifest.json' + if not manifest_path.is_file(): + raise HTTPException(status_code=404, detail=f'no manifest at {manifest_path}') + await ingest_manifest(session, layout, manifest_path) + await session.commit() + return {'status': 'rescanned', 'kind': kind, 'uuid': str(uuid)} + + +@router.get('/stats', response_model=StoreStats) +async def stats(session: SessionDep) -> StoreStats: + async def _count(model) -> int: # type: ignore[no-untyped-def] + return int((await session.execute(select(func.count()).select_from(model))).scalar_one()) + + async def _invalid_count() -> int: + total = 0 + for model in (Campaign, Diffraction, Product, Fluorescence): + stmt = ( + select(func.count()) + .select_from(model) + .where( + or_( + model.ingest_state == IngestState.INVALID, + model.ingest_state == IngestState.MISSING_FILES, + model.ingest_state == IngestState.ORPHANED, + ) + ) + ) + total += int((await session.execute(stmt)).scalar_one()) + return total + + return StoreStats( + campaign_count=await _count(Campaign), + diffraction_count=await _count(Diffraction), + product_count=await _count(Product), + fluorescence_count=await _count(Fluorescence), + invalid_count=await _invalid_count(), + ) diff --git a/src/ptychodus_store/routers/campaign.py b/src/ptychodus_store/routers/campaign.py new file mode 100644 index 000000000..44a59bb91 --- /dev/null +++ b/src/ptychodus_store/routers/campaign.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from uuid import UUID + +from fastapi import APIRouter, HTTPException, Query +from sqlalchemy import String, func + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import Campaign +from ptychodus_store.routers._convert import campaign_to_read +from ptychodus_store.routers.deps import SessionDep +from ptychodus_store.routers.schemas import CampaignRead, Page +from ptychodus_store.storage.manifest import ResourceKind + +router = APIRouter(prefix='/campaign', tags=['campaign']) + + +@router.get('', response_model=Page[CampaignRead]) +async def list_campaigns( + session: SessionDep, + limit: int = Query(50, ge=1, le=200), + offset: int = Query(0, ge=0), + sample_name: str | None = None, + tag: str | None = None, + ingest_state: IngestState | None = None, +) -> Page[CampaignRead]: + where = [] + if sample_name is not None: + where.append(Campaign.sample_name == sample_name) + if ingest_state is not None: + where.append(Campaign.ingest_state == ingest_state) + if tag is not None: + # Match if the tag string appears as an element of the JSON array. + # tags is stored as a JSON-encoded TEXT column, so we look for the + # quoted JSON form of the value. + where.append(func.instr(func.cast(Campaign.tags, String), f'"{tag}"') > 0) + + items, total = await repo.list_rows( + session, ResourceKind.CAMPAIGN, limit=limit, offset=offset, where=where + ) + return Page(items=[campaign_to_read(i) for i in items], total=total, limit=limit, offset=offset) + + +@router.get('/{uuid}', response_model=CampaignRead) +async def get_campaign(uuid: UUID, session: SessionDep) -> CampaignRead: + row = await repo.get_row(session, ResourceKind.CAMPAIGN, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'campaign {uuid} not found') + return campaign_to_read(row) diff --git a/src/ptychodus_store/routers/deps.py b/src/ptychodus_store/routers/deps.py new file mode 100644 index 000000000..6ca5ec5ba --- /dev/null +++ b/src/ptychodus_store/routers/deps.py @@ -0,0 +1,37 @@ +"""FastAPI dependency providers.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Annotated + +from fastapi import Depends, Request +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.config import Settings +from ptychodus_store.db.session import SessionProvider +from ptychodus_store.storage.layout import StoreLayout + + +def get_settings_dep(request: Request) -> Settings: + return request.app.state.settings # type: ignore[no-any-return] + + +def get_layout(request: Request) -> StoreLayout: + return request.app.state.layout # type: ignore[no-any-return] + + +def get_session_provider(request: Request) -> SessionProvider: + return request.app.state.session_provider # type: ignore[no-any-return] + + +async def get_session( + provider: Annotated[SessionProvider, Depends(get_session_provider)], +) -> AsyncIterator[AsyncSession]: + async for session in provider.session(): + yield session + + +SettingsDep = Annotated[Settings, Depends(get_settings_dep)] +LayoutDep = Annotated[StoreLayout, Depends(get_layout)] +SessionDep = Annotated[AsyncSession, Depends(get_session)] diff --git a/src/ptychodus_store/routers/diffraction.py b/src/ptychodus_store/routers/diffraction.py new file mode 100644 index 000000000..afcddba73 --- /dev/null +++ b/src/ptychodus_store/routers/diffraction.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from uuid import UUID + +from fastapi import APIRouter, HTTPException, Query +from fastapi.responses import FileResponse +from sqlalchemy import exists, select + +from ptychodus.api.diffraction import Polarization +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.io import load_diffraction_data +from ptychodus.api.reconstructor import AssembledDiffractionData + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import DerivationEdge, Diffraction +from ptychodus_store.rendering import RenderedImage, render_real +from ptychodus_store.rendering.params import RenderParamsDep +from ptychodus_store.routers._convert import diffraction_to_read +from ptychodus_store.routers.deps import LayoutDep, SessionDep +from ptychodus_store.routers.schemas import DiffractionRead, Page +from ptychodus_store.storage.manifest import ResourceKind + +router = APIRouter(prefix='/diffraction', tags=['diffraction']) + + +@router.get('', response_model=Page[DiffractionRead]) +async def list_diffraction( + session: SessionDep, + limit: int = Query(50, ge=1, le=200), + offset: int = Query(0, ge=0), + campaign_uuid: UUID | None = None, + derived_from_uuid: UUID | None = None, + ingest_state: IngestState | None = None, + probe_energy_eV_min: float | None = Query(None, alias='probe_energy_eV_min'), # noqa: N803 + probe_energy_eV_max: float | None = Query(None, alias='probe_energy_eV_max'), # noqa: N803 + tomography_angle_deg_min: float | None = None, + tomography_angle_deg_max: float | None = None, + tilt_angle_deg_min: float | None = None, + tilt_angle_deg_max: float | None = None, + polarization: Polarization | None = None, +) -> Page[DiffractionRead]: + where = [] + if campaign_uuid is not None: + where.append(Diffraction.campaign_uuid == campaign_uuid) + if ingest_state is not None: + where.append(Diffraction.ingest_state == ingest_state) + if probe_energy_eV_min is not None: + where.append(Diffraction.probe_energy_eV >= probe_energy_eV_min) + if probe_energy_eV_max is not None: + where.append(Diffraction.probe_energy_eV <= probe_energy_eV_max) + if tomography_angle_deg_min is not None: + where.append(Diffraction.tomography_angle_deg >= tomography_angle_deg_min) + if tomography_angle_deg_max is not None: + where.append(Diffraction.tomography_angle_deg <= tomography_angle_deg_max) + if tilt_angle_deg_min is not None: + where.append(Diffraction.tilt_angle_deg >= tilt_angle_deg_min) + if tilt_angle_deg_max is not None: + where.append(Diffraction.tilt_angle_deg <= tilt_angle_deg_max) + if polarization is not None: + where.append(Diffraction.polarization == polarization.value) + if derived_from_uuid is not None: + edge_subq = select(DerivationEdge.source_uuid).where( + DerivationEdge.source_uuid == Diffraction.uuid, + DerivationEdge.target_uuid == derived_from_uuid, + ) + where.append(exists(edge_subq)) + + items, total = await repo.list_rows( + session, ResourceKind.DIFFRACTION, limit=limit, offset=offset, where=where + ) + reads = [await diffraction_to_read(session, i) for i in items] # type: ignore[arg-type] + return Page(items=reads, total=total, limit=limit, offset=offset) + + +@router.get('/{uuid}', response_model=DiffractionRead) +async def get_diffraction(uuid: UUID, session: SessionDep) -> DiffractionRead: + row = await repo.get_row(session, ResourceKind.DIFFRACTION, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'diffraction {uuid} not found') + return await diffraction_to_read(session, row) # type: ignore[arg-type] + + +@router.get('/{uuid}/files/diffraction') +async def get_diffraction_file(uuid: UUID, session: SessionDep, layout: LayoutDep) -> FileResponse: + row = await repo.get_row(session, ResourceKind.DIFFRACTION, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'diffraction {uuid} not found') + path = layout.resource_folder(ResourceKind.DIFFRACTION, uuid) / 'diffraction.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='diffraction.h5 not present on disk') + return FileResponse(path, media_type='application/x-hdf5', filename=path.name) + + +async def _load_diffraction_or_404( + uuid: UUID, session: SessionDep, layout: LayoutDep +) -> tuple[AssembledDiffractionData, PixelGeometry]: + row = await repo.get_row(session, ResourceKind.DIFFRACTION, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'diffraction {uuid} not found') + path = layout.resource_folder(ResourceKind.DIFFRACTION, uuid) / 'diffraction.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='diffraction.h5 not present on disk') + data = load_diffraction_data(path) + return data, data.get_pixel_geometry() + + +@router.get('/{uuid}/patterns/aggregate/image', response_model=RenderedImage) +async def get_diffraction_aggregate_image( + uuid: UUID, + session: SessionDep, + layout: LayoutDep, + params: RenderParamsDep, +) -> RenderedImage: + data, pixel_geometry = await _load_diffraction_or_404(uuid, session, layout) + return render_real( + data.get_average_pattern(), pixel_geometry, params, value_label='Mean Counts' + ) + + +@router.get('/{uuid}/patterns/{index}/image', response_model=RenderedImage) +async def get_diffraction_pattern_image( + uuid: UUID, + index: int, + session: SessionDep, + layout: LayoutDep, + params: RenderParamsDep, +) -> RenderedImage: + data, pixel_geometry = await _load_diffraction_or_404(uuid, session, layout) + num_patterns = data.get_patterns_shape()[0] + if not 0 <= index < num_patterns: + raise HTTPException( + status_code=404, + detail=f'pattern index {index} out of range [0, {num_patterns})', + ) + return render_real(data.get_pattern(index), pixel_geometry, params, value_label='Counts') diff --git a/src/ptychodus_store/routers/fluorescence.py b/src/ptychodus_store/routers/fluorescence.py new file mode 100644 index 000000000..1ebf67253 --- /dev/null +++ b/src/ptychodus_store/routers/fluorescence.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +from uuid import UUID + +from fastapi import APIRouter, HTTPException, Query +from fastapi.responses import FileResponse +from sqlalchemy import String, exists, func, select + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.io import load_fluorescence_data, load_product + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import DerivationEdge, Fluorescence +from ptychodus_store.rendering import RenderedImage, render_real +from ptychodus_store.rendering.params import RenderParamsDep +from ptychodus_store.routers._convert import fluorescence_to_read +from ptychodus_store.routers.deps import LayoutDep, SessionDep +from ptychodus_store.routers.schemas import FluorescenceRead, Page +from ptychodus_store.storage.manifest import ResourceKind + +_DEFAULT_FLUORESCENCE_PIXEL_GEOMETRY = PixelGeometry(width_m=1e-6, height_m=1e-6) + +router = APIRouter(prefix='/fluorescence', tags=['fluorescence']) + + +@router.get('', response_model=Page[FluorescenceRead]) +async def list_fluorescence( + session: SessionDep, + limit: int = Query(50, ge=1, le=200), + offset: int = Query(0, ge=0), + derived_from_uuid: UUID | None = None, + element: str | None = None, + ingest_state: IngestState | None = None, +) -> Page[FluorescenceRead]: + where = [] + if ingest_state is not None: + where.append(Fluorescence.ingest_state == ingest_state) + if element is not None: + # JSON column contains the element string — use LIKE on the serialized JSON + where.append(func.instr(func.cast(Fluorescence.element_names, String), f'"{element}"') > 0) + if derived_from_uuid is not None: + edge_subq = select(DerivationEdge.source_uuid).where( + DerivationEdge.source_uuid == Fluorescence.uuid, + DerivationEdge.target_uuid == derived_from_uuid, + ) + where.append(exists(edge_subq)) + + items, total = await repo.list_rows( + session, ResourceKind.FLUORESCENCE, limit=limit, offset=offset, where=where + ) + reads = [await fluorescence_to_read(session, i) for i in items] # type: ignore[arg-type] + return Page(items=reads, total=total, limit=limit, offset=offset) + + +@router.get('/{uuid}', response_model=FluorescenceRead) +async def get_fluorescence(uuid: UUID, session: SessionDep) -> FluorescenceRead: + row = await repo.get_row(session, ResourceKind.FLUORESCENCE, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'fluorescence {uuid} not found') + return await fluorescence_to_read(session, row) # type: ignore[arg-type] + + +@router.get('/{uuid}/files/fluorescence') +async def get_fluorescence_file(uuid: UUID, session: SessionDep, layout: LayoutDep) -> FileResponse: + row = await repo.get_row(session, ResourceKind.FLUORESCENCE, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'fluorescence {uuid} not found') + path = layout.resource_folder(ResourceKind.FLUORESCENCE, uuid) / 'fluorescence.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='fluorescence.h5 not present on disk') + return FileResponse(path, media_type='application/x-hdf5', filename=path.name) + + +async def _resolve_product_pixel_geometry( + product_uuid: UUID | None, session: SessionDep, layout: LayoutDep +) -> PixelGeometry: + """Look up a paired product's object pixel geometry; fall back to a 1 µm placeholder.""" + if product_uuid is None: + return _DEFAULT_FLUORESCENCE_PIXEL_GEOMETRY + row = await repo.get_row(session, ResourceKind.PRODUCT, product_uuid) + if row is None: + raise HTTPException( + status_code=404, detail=f'product {product_uuid} not found for pixel geometry' + ) + path = layout.resource_folder(ResourceKind.PRODUCT, product_uuid) / 'product.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='product.h5 not present on disk') + return load_product(path).object_.get_pixel_geometry() + + +@router.get('/{uuid}/elements/{name}/image', response_model=RenderedImage) +async def get_fluorescence_element_image( + uuid: UUID, + name: str, + session: SessionDep, + layout: LayoutDep, + params: RenderParamsDep, + product_uuid: UUID | None = Query( + None, + description='Optional paired product for physical pixel geometry.', + ), +) -> RenderedImage: + row = await repo.get_row(session, ResourceKind.FLUORESCENCE, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'fluorescence {uuid} not found') + path = layout.resource_folder(ResourceKind.FLUORESCENCE, uuid) / 'fluorescence.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='fluorescence.h5 not present on disk') + + dataset = load_fluorescence_data(path) + matches = [emap for emap in dataset.element_maps if emap.name == name] + if not matches: + raise HTTPException( + status_code=404, + detail={ + 'error': f'element {name!r} not found', + 'available': [emap.name for emap in dataset.element_maps], + }, + ) + emap = matches[0] + + pixel_geometry = await _resolve_product_pixel_geometry(product_uuid, session, layout) + return render_real( + emap.counts_per_second, pixel_geometry, params, value_label=f'{emap.name} counts/s' + ) diff --git a/src/ptychodus_store/routers/health.py b/src/ptychodus_store/routers/health.py new file mode 100644 index 000000000..a3904d812 --- /dev/null +++ b/src/ptychodus_store/routers/health.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from fastapi import APIRouter, Request +from sqlalchemy import text + +from ptychodus_store.routers.deps import SessionDep +from ptychodus_store.routers.schemas import HealthRead + +router = APIRouter(tags=['health']) + + +@router.get('/health', response_model=HealthRead) +async def health(request: Request, session: SessionDep) -> HealthRead: + db_state: str = 'ok' + try: + await session.execute(text('SELECT 1')) + except Exception: # noqa: BLE001 + db_state = 'down' + + watcher = getattr(request.app.state, 'watcher', None) + if watcher is None: + watcher_state = 'disabled' + elif watcher.is_alive: + watcher_state = 'alive' + else: + watcher_state = 'dead' + + status = 'ok' if db_state == 'ok' and watcher_state != 'dead' else 'degraded' + return HealthRead(status=status, db=db_state, watcher=watcher_state) # type: ignore[arg-type] diff --git a/src/ptychodus_store/routers/lineage.py b/src/ptychodus_store/routers/lineage.py new file mode 100644 index 000000000..728128916 --- /dev/null +++ b/src/ptychodus_store/routers/lineage.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from collections import deque +from uuid import UUID + +from fastapi import APIRouter, HTTPException +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.models import Campaign, DerivationEdge, Diffraction +from ptychodus_store.routers._convert import campaign_to_read +from ptychodus_store.routers.deps import SessionDep +from ptychodus_store.routers.schemas import LineageNode, LineageRead + +router = APIRouter(tags=['lineage']) + + +def _label_for(row: object) -> str: + return str(getattr(row, 'label', '') or getattr(row, 'name', '') or '') + + +async def _resolve_node(session: AsyncSession, uuid: UUID) -> tuple[str, object] | None: + kind = await repo.find_kind_for_uuid(session, uuid) + if kind is None: + return None + row = await repo.get_row(session, kind, uuid) + if row is None: + return None + return kind, row + + +async def _walk_ancestors(session: AsyncSession, root: UUID) -> list[LineageNode]: + """BFS over outgoing edges (child -> parent), cycle-guarded by visited set.""" + out: list[LineageNode] = [] + visited: set[UUID] = {root} + queue: deque[UUID] = deque([root]) + while queue: + current = queue.popleft() + edges = await repo.outgoing_edges(session, current) + for edge in edges: + if edge.target_uuid in visited: + continue + visited.add(edge.target_uuid) + resolved = await _resolve_node(session, edge.target_uuid) + if resolved is None: + continue + kind, row = resolved + out.append(LineageNode(kind=kind, uuid=edge.target_uuid, label=_label_for(row))) # type: ignore[arg-type] + queue.append(edge.target_uuid) + return out + + +async def _walk_descendants(session: AsyncSession, root: UUID) -> list[LineageNode]: + """BFS over incoming edges (parent -> child), cycle-guarded.""" + out: list[LineageNode] = [] + visited: set[UUID] = {root} + queue: deque[UUID] = deque([root]) + while queue: + current = queue.popleft() + rows = ( + ( + await session.execute( + select(DerivationEdge).where(DerivationEdge.target_uuid == current) + ) + ) + .scalars() + .all() + ) + for edge in rows: + if edge.source_uuid in visited: + continue + visited.add(edge.source_uuid) + resolved = await _resolve_node(session, edge.source_uuid) + if resolved is None: + continue + kind, row = resolved + out.append(LineageNode(kind=kind, uuid=edge.source_uuid, label=_label_for(row))) # type: ignore[arg-type] + queue.append(edge.source_uuid) + return out + + +async def _find_campaign(session: AsyncSession, root_uuid: UUID, ancestors: list[LineageNode]): + """Look at the root node and its ancestors for any diffraction → return its campaign.""" + candidates: list[UUID] = [root_uuid] + [a.uuid for a in ancestors] + for uuid in candidates: + diff = await session.get(Diffraction, uuid) + if diff is not None and diff.campaign_uuid is not None: + campaign = await session.get(Campaign, diff.campaign_uuid) + if campaign is not None: + return campaign_to_read(campaign) + return None + + +@router.get('/lineage/{uuid}', response_model=LineageRead) +async def get_lineage(uuid: UUID, session: SessionDep) -> LineageRead: + resolved = await _resolve_node(session, uuid) + if resolved is None: + raise HTTPException(status_code=404, detail=f'no resource with uuid {uuid}') + kind, row = resolved + node = LineageNode(kind=kind, uuid=uuid, label=_label_for(row)) # type: ignore[arg-type] + + ancestors = await _walk_ancestors(session, uuid) + descendants = await _walk_descendants(session, uuid) + campaign = await _find_campaign(session, uuid, ancestors) + return LineageRead(node=node, ancestors=ancestors, descendants=descendants, campaign=campaign) diff --git a/src/ptychodus_store/routers/product.py b/src/ptychodus_store/routers/product.py new file mode 100644 index 000000000..767b86ebd --- /dev/null +++ b/src/ptychodus_store/routers/product.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +from base64 import b64encode +from io import BytesIO +from uuid import UUID + +import numpy +from fastapi import APIRouter, HTTPException, Query +from fastapi.responses import FileResponse +from PIL import Image, ImageDraw +from sqlalchemy import exists, select + +from ptychodus.api.io import load_product +from ptychodus.api.product import Product as ProductAggregate + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.db.models import DerivationEdge, Product +from ptychodus_store.rendering import RenderedImage, render_complex +from ptychodus_store.rendering.params import RenderParamsDep +from ptychodus_store.routers._convert import product_to_read +from ptychodus_store.routers.deps import LayoutDep, SessionDep +from ptychodus_store.routers.schemas import Page, ProductRead +from ptychodus_store.storage.manifest import ResourceKind + +router = APIRouter(prefix='/product', tags=['product']) + + +@router.get('', response_model=Page[ProductRead]) +async def list_product( + session: SessionDep, + limit: int = Query(50, ge=1, le=200), + offset: int = Query(0, ge=0), + derived_from_uuid: UUID | None = None, + ingest_state: IngestState | None = None, + probe_energy_eV_min: float | None = Query(None, alias='probe_energy_eV_min'), # noqa: N803 + probe_energy_eV_max: float | None = Query(None, alias='probe_energy_eV_max'), # noqa: N803 +) -> Page[ProductRead]: + where = [] + if ingest_state is not None: + where.append(Product.ingest_state == ingest_state) + if probe_energy_eV_min is not None: + where.append(Product.probe_energy_eV >= probe_energy_eV_min) + if probe_energy_eV_max is not None: + where.append(Product.probe_energy_eV <= probe_energy_eV_max) + if derived_from_uuid is not None: + edge_subq = select(DerivationEdge.source_uuid).where( + DerivationEdge.source_uuid == Product.uuid, + DerivationEdge.target_uuid == derived_from_uuid, + ) + where.append(exists(edge_subq)) + + items, total = await repo.list_rows( + session, ResourceKind.PRODUCT, limit=limit, offset=offset, where=where + ) + reads = [await product_to_read(session, i) for i in items] # type: ignore[arg-type] + return Page(items=reads, total=total, limit=limit, offset=offset) + + +@router.get('/{uuid}', response_model=ProductRead) +async def get_product(uuid: UUID, session: SessionDep) -> ProductRead: + row = await repo.get_row(session, ResourceKind.PRODUCT, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'product {uuid} not found') + return await product_to_read(session, row) # type: ignore[arg-type] + + +@router.get('/{uuid}/files/product') +async def get_product_file(uuid: UUID, session: SessionDep, layout: LayoutDep) -> FileResponse: + row = await repo.get_row(session, ResourceKind.PRODUCT, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'product {uuid} not found') + path = layout.resource_folder(ResourceKind.PRODUCT, uuid) / 'product.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='product.h5 not present on disk') + return FileResponse(path, media_type='application/x-hdf5', filename=path.name) + + +async def _load_product_or_404( + uuid: UUID, session: SessionDep, layout: LayoutDep +) -> ProductAggregate: + row = await repo.get_row(session, ResourceKind.PRODUCT, uuid) + if row is None: + raise HTTPException(status_code=404, detail=f'product {uuid} not found') + path = layout.resource_folder(ResourceKind.PRODUCT, uuid) / 'product.h5' + if not path.is_file(): + raise HTTPException(status_code=404, detail='product.h5 not present on disk') + return load_product(path) + + +@router.get('/{uuid}/probe/image', response_model=RenderedImage) +async def get_probe_image( + uuid: UUID, + session: SessionDep, + layout: LayoutDep, + params: RenderParamsDep, + incoherent: int = Query(0, ge=0, description='Incoherent probe mode index.'), +) -> RenderedImage: + product = await _load_product_or_404(uuid, session, layout) + probe = product.probes.get_probe_no_opr() + if not 0 <= incoherent < probe.num_incoherent_modes: + raise HTTPException( + status_code=404, + detail=(f'incoherent mode {incoherent} out of range [0, {probe.num_incoherent_modes})'), + ) + values = probe.get_incoherent_mode(incoherent) + return render_complex(values, product.probes.get_pixel_geometry(), params) + + +@router.get('/{uuid}/probe/modes/image', response_model=RenderedImage) +async def get_probe_modes_image( + uuid: UUID, + session: SessionDep, + layout: LayoutDep, + params: RenderParamsDep, +) -> RenderedImage: + product = await _load_product_or_404(uuid, session, layout) + probe = product.probes.get_probe_no_opr() + values = probe.get_incoherent_modes_flattened() + return render_complex(values, product.probes.get_pixel_geometry(), params) + + +@router.get('/{uuid}/object/{layer}/image', response_model=RenderedImage) +async def get_object_layer_image( + uuid: UUID, + layer: int, + session: SessionDep, + layout: LayoutDep, + params: RenderParamsDep, +) -> RenderedImage: + product = await _load_product_or_404(uuid, session, layout) + if not 0 <= layer < product.object_.num_layers: + raise HTTPException( + status_code=404, + detail=f'object layer {layer} out of range [0, {product.object_.num_layers})', + ) + values = product.object_.get_layer(layer) + return render_complex(values, product.object_.get_pixel_geometry(), params) + + +@router.get('/{uuid}/positions/image', response_model=RenderedImage) +async def get_positions_image( + uuid: UUID, + session: SessionDep, + layout: LayoutDep, + canvas_px: int = Query(512, ge=64, le=2048, description='Square output resolution in pixels.'), + connect_path: bool = Query(True, description='Draw a polyline through successive scan points.'), + margin_frac: float = Query( + 0.05, ge=0.0, le=0.5, description='Blank margin around the bounding box, as a fraction.' + ), +) -> RenderedImage: + product = await _load_product_or_404(uuid, session, layout) + positions = product.probe_positions + n_points = len(positions) + if n_points == 0: + raise HTTPException(status_code=404, detail='product has no probe positions') + + xs = numpy.array([positions[i].coordinate_x_m for i in range(n_points)], dtype=float) + ys = numpy.array([positions[i].coordinate_y_m for i in range(n_points)], dtype=float) + + x_min, x_max = float(xs.min()), float(xs.max()) + y_min, y_max = float(ys.min()), float(ys.max()) + width_m = max(x_max - x_min, 1e-9) + height_m = max(y_max - y_min, 1e-9) + range_m = max(width_m, height_m) * (1.0 + 2.0 * margin_frac) + x_center = 0.5 * (x_min + x_max) + y_center = 0.5 * (y_min + y_max) + + scale = canvas_px / range_m + px_x = (xs - x_center) * scale + canvas_px / 2.0 + px_y = canvas_px / 2.0 - (ys - y_center) * scale # flip y so +y is up + pts = [(int(round(px)), int(round(py))) for px, py in zip(px_x, px_y)] + + img = Image.new('RGB', (canvas_px, canvas_px), (0, 0, 0)) + draw = ImageDraw.Draw(img) + if connect_path and n_points > 1: + draw.line(pts, fill=(80, 80, 80), width=1) + radius = max(1, canvas_px // 200) + denom = max(1, n_points - 1) + for i, (px, py) in enumerate(pts): + t = i / denom + color = (int(255 * t), 80, int(255 * (1.0 - t))) + draw.ellipse((px - radius, py - radius, px + radius, py + radius), fill=color) + + buf = BytesIO() + img.save(buf, format='PNG') + pixel_m = range_m / canvas_px + return RenderedImage( + png_base64=b64encode(buf.getvalue()).decode('ascii'), + value_label='Scan Index', + color_value_min=0.0, + color_value_max=float(max(0, n_points - 1)), + pixel_width_m=pixel_m, + pixel_height_m=pixel_m, + shape_h_px=canvas_px, + shape_w_px=canvas_px, + ) diff --git a/src/ptychodus_store/routers/schemas.py b/src/ptychodus_store/routers/schemas.py new file mode 100644 index 000000000..20ab6e542 --- /dev/null +++ b/src/ptychodus_store/routers/schemas.py @@ -0,0 +1,145 @@ +"""Pydantic response models for the REST API.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Generic, Literal, TypeVar +from uuid import UUID + +from pydantic import BaseModel, ConfigDict + +from ptychodus.api.diffraction import Polarization +from ptychodus_store.db.base import IngestState + + +class _RowBase(BaseModel): + model_config = ConfigDict(from_attributes=True) + + uuid: UUID + folder_path: str + manifest_mtime: datetime | None + created_from_manifest_at: datetime | None + ingest_state: IngestState + error_message: str | None + created_at: datetime + updated_at: datetime + + +class CampaignRead(_RowBase): + label: str + comments: str + sample_name: str + sample_description: str + tags: list[str] + + +class DerivedFromEdge(BaseModel): + kind: Literal['diffraction', 'product', 'fluorescence'] + uuid: UUID + + +class DiffractionRead(_RowBase): + label: str + comments: str + campaign_uuid: UUID | None + derived_from: list[DerivedFromEdge] = [] + + detector_distance_m: float | None + probe_energy_eV: float | None # noqa: N815 + probe_photon_count: int | None + exposure_time_s: float | None + tomography_angle_deg: float | None + tilt_angle_deg: float | None = None + polarization: Polarization | None = None + crop_center_x_px: int | None + crop_center_y_px: int | None + + pattern_dtype: str | None + pattern_height_px: int | None + pattern_width_px: int | None + num_patterns_total: int | None + detector_pixel_width_m: float | None + detector_pixel_height_m: float | None + + +class ProductRead(_RowBase): + derived_from: list[DerivedFromEdge] = [] + + name: str | None + comments: str | None + detector_distance_m: float | None + probe_energy_eV: float | None # noqa: N815 + probe_photon_count: int | None + exposure_time_s: float | None + mass_attenuation_m2_kg: float | None + tomography_angle_deg: float | None + tilt_angle_deg: float | None = None + polarization: Polarization | None = None + + object_layers: int | None + object_height_px: int | None + object_width_px: int | None + object_pixel_width_m: float | None + object_pixel_height_m: float | None + probe_modes: int | None + probe_height_px: int | None + probe_width_px: int | None + num_scan_points: int | None + num_loss_epochs: int | None + + +class FluorescenceRead(_RowBase): + label: str + comments: str + derived_from: list[DerivedFromEdge] = [] + + element_names: list[str] + map_height_px: int | None + map_width_px: int | None + + +ItemT = TypeVar('ItemT', bound=BaseModel) + + +class Page(BaseModel, Generic[ItemT]): + items: list[ItemT] + total: int + limit: int + offset: int + + +class ResourceRef(BaseModel): + kind: Literal['campaign', 'diffraction', 'product', 'fluorescence'] + uuid: UUID + + +class LineageNode(BaseModel): + kind: Literal['campaign', 'diffraction', 'product', 'fluorescence'] + uuid: UUID + label: str = '' + + +class LineageRead(BaseModel): + node: LineageNode + ancestors: list[LineageNode] + descendants: list[LineageNode] + campaign: CampaignRead | None + + +class StoreStats(BaseModel): + campaign_count: int + diffraction_count: int + product_count: int + fluorescence_count: int + invalid_count: int + + +class HealthRead(BaseModel): + status: Literal['ok', 'degraded'] + db: Literal['ok', 'down'] + watcher: Literal['alive', 'dead', 'disabled'] + + +class ReindexResponse(BaseModel): + status: Literal['accepted'] + job_id: str diff --git a/src/ptychodus_store/routers/visualization.py b/src/ptychodus_store/routers/visualization.py new file mode 100644 index 000000000..f1d694bb6 --- /dev/null +++ b/src/ptychodus_store/routers/visualization.py @@ -0,0 +1,29 @@ +"""Options-enumeration endpoint for the visualization API.""" + +from __future__ import annotations + +from fastapi import APIRouter + +from ptychodus.api.visualization import ( + ComplexComponent, + CylindricalColorModel, + ScalarTransformation, + cyclic_colormap_names, + linear_colormap_names, +) + +from ptychodus_store.rendering.schemas import OptionsRead + +router = APIRouter(prefix='/visualization', tags=['visualization']) + + +@router.get('/options', response_model=OptionsRead) +async def get_visualization_options() -> OptionsRead: + """Enumerate the valid choices a client may pass to render endpoints.""" + return OptionsRead( + colormaps_linear=list(linear_colormap_names()), + colormaps_cyclic=list(cyclic_colormap_names()), + transforms=[member.name.lower() for member in ScalarTransformation], + components=[member.name.lower() for member in ComplexComponent], + color_models=[member.name.lower() for member in CylindricalColorModel], + ) diff --git a/src/ptychodus_store/storage/__init__.py b/src/ptychodus_store/storage/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/ptychodus_store/storage/h5_introspect.py b/src/ptychodus_store/storage/h5_introspect.py new file mode 100644 index 000000000..a02fb6167 --- /dev/null +++ b/src/ptychodus_store/storage/h5_introspect.py @@ -0,0 +1,173 @@ +"""Read HDF5 attrs / dataset shapes for the DB cache without loading array data.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import h5py + +from ptychodus.api.io import DiffractionFileKeys, ProductFileKeys, load_fluorescence_data + + +class IntrospectionError(Exception): + """Raised when an HDF5 file cannot be read or lacks expected structure.""" + + +def _attr(group: h5py.HLObject, key: str, *, cast: type) -> Any: + raw = group.attrs.get(key) + if raw is None: + return None + try: + return cast(raw) + except (TypeError, ValueError): + return None + + +def introspect_diffraction(path: Path) -> dict[str, Any]: + """Read scalar attrs and dataset shape/dtype from a diffraction.h5 file. + + Returns a dict of HDF5-derived fields: + * pattern_dtype: str + * pattern_shape: tuple[int, int] | None — (height, width) + * num_patterns_total: int | None + * detector_pixel_width_m, detector_pixel_height_m: float | None + """ + try: + with h5py.File(path, 'r') as f: + patterns = f.get(DiffractionFileKeys.PATTERNS) + if not isinstance(patterns, h5py.Dataset): + raise IntrospectionError( + f'{path}: missing {DiffractionFileKeys.PATTERNS!r} dataset' + ) + + shape = tuple(int(x) for x in patterns.shape) + num_patterns_total = shape[0] if len(shape) >= 1 else None + pattern_shape = (shape[1], shape[2]) if len(shape) == 3 else None + pattern_dtype = str(patterns.dtype) + + pixel_width = _attr(patterns, DiffractionFileKeys.DETECTOR_PIXEL_WIDTH, cast=float) + pixel_height = _attr(patterns, DiffractionFileKeys.DETECTOR_PIXEL_HEIGHT, cast=float) + + return { + 'pattern_dtype': pattern_dtype, + 'pattern_shape': pattern_shape, + 'num_patterns_total': num_patterns_total, + 'detector_pixel_width_m': pixel_width, + 'detector_pixel_height_m': pixel_height, + } + except (OSError, KeyError) as exc: + raise IntrospectionError(f'{path}: {exc}') from exc + + +def introspect_product(path: Path) -> dict[str, Any]: + """Read root-level attrs and probe/object dataset shapes from product.h5. + + Returns a dict of HDF5-derived fields: + * name, comments + * detector_distance_m, probe_energy_eV, probe_photon_count, exposure_time_s, + mass_attenuation_m2_kg, tomography_angle_deg, tilt_angle_deg, polarization + * object_shape: tuple[int, int, int] | None — (layers, h, w) + * object_pixel_width_m, object_pixel_height_m: float | None + * probe_shape: tuple[int, int, int] | None — (modes, h, w) + * num_scan_points: int | None + * num_loss_epochs: int + """ + try: + with h5py.File(path, 'r') as f: + name = str(f.attrs.get(ProductFileKeys.NAME, '')) + comments = str(f.attrs.get(ProductFileKeys.COMMENTS, '')) + + def _root_attr(key: str, cast: type) -> Any: + raw = f.attrs.get(key) + if raw is None: + return None + try: + return cast(raw) + except (TypeError, ValueError): + return None + + def _root_str_attr(key: str) -> str | None: + raw = f.attrs.get(key) + if raw is None: + return None + if isinstance(raw, bytes): + raw = raw.decode('utf-8', errors='replace') + text = str(raw) + return text or None + + obj = f.get(ProductFileKeys.OBJECT_ARRAY) + probe = f.get(ProductFileKeys.PROBE_ARRAY) + positions = f.get(ProductFileKeys.PROBE_POSITION_INDEXES) + loss_epochs = f.get(ProductFileKeys.LOSS_EPOCHS) + + object_shape: tuple[int, int, int] | None = None + object_pixel_width_m: float | None = None + object_pixel_height_m: float | None = None + if isinstance(obj, h5py.Dataset) and len(obj.shape) == 3: + object_shape = (int(obj.shape[0]), int(obj.shape[1]), int(obj.shape[2])) + object_pixel_width_m = _attr(obj, ProductFileKeys.OBJECT_PIXEL_WIDTH, cast=float) + object_pixel_height_m = _attr(obj, ProductFileKeys.OBJECT_PIXEL_HEIGHT, cast=float) + + probe_shape: tuple[int, int, int] | None = None + if isinstance(probe, h5py.Dataset): + shape = tuple(int(x) for x in probe.shape) + if len(shape) == 3: + probe_shape = (shape[0], shape[1], shape[2]) + elif len(shape) == 4: + # (coherent, incoherent, h, w) — collapse coherent x incoherent into modes + probe_shape = (shape[0] * shape[1], shape[2], shape[3]) + + num_scan_points: int | None = None + if isinstance(positions, h5py.Dataset): + num_scan_points = int(positions.shape[0]) + + num_loss_epochs = 0 + if isinstance(loss_epochs, h5py.Dataset): + num_loss_epochs = int(loss_epochs.shape[0]) + + return { + 'name': name, + 'comments': comments, + 'detector_distance_m': _root_attr(ProductFileKeys.DETECTOR_OBJECT_DISTANCE, float), + 'probe_energy_eV': _root_attr(ProductFileKeys.PROBE_ENERGY, float), + 'probe_photon_count': _root_attr(ProductFileKeys.PROBE_PHOTON_COUNT, int), + 'exposure_time_s': _root_attr(ProductFileKeys.EXPOSURE_TIME, float), + 'mass_attenuation_m2_kg': _root_attr(ProductFileKeys.MASS_ATTENUATION, float), + 'tomography_angle_deg': _root_attr(ProductFileKeys.TOMOGRAPHY_ANGLE, float), + 'tilt_angle_deg': _root_attr(ProductFileKeys.TILT_ANGLE, float), + 'polarization': _root_str_attr(ProductFileKeys.POLARIZATION), + 'object_shape': object_shape, + 'object_pixel_width_m': object_pixel_width_m, + 'object_pixel_height_m': object_pixel_height_m, + 'probe_shape': probe_shape, + 'num_scan_points': num_scan_points, + 'num_loss_epochs': num_loss_epochs, + } + except (OSError, KeyError) as exc: + raise IntrospectionError(f'{path}: {exc}') from exc + + +def introspect_fluorescence(path: Path) -> dict[str, Any]: + """Read element names and map shape from a fluorescence.h5 file (XRF-Maps layout). + + Returns a dict with `element_names: list[str]` and `map_shape: tuple[int, int] | None`. + Delegates to :func:`ptychodus.api.io.load_fluorescence_data`, which recognises the + v10 NNLS/Fitted and legacy v9 layouts. + """ + try: + dataset = load_fluorescence_data(path) + except (OSError, KeyError, ValueError) as exc: + raise IntrospectionError(f'{path}: {exc}') from exc + + element_names = [emap.name for emap in dataset.element_maps] + map_shape: tuple[int, int] | None = None + if dataset.element_maps: + cps = dataset.element_maps[0].counts_per_second + if cps.ndim == 2: + map_shape = (int(cps.shape[0]), int(cps.shape[1])) + + return { + 'element_names': element_names, + 'map_shape': map_shape, + } diff --git a/src/ptychodus_store/storage/layout.py b/src/ptychodus_store/storage/layout.py new file mode 100644 index 000000000..4c9b21375 --- /dev/null +++ b/src/ptychodus_store/storage/layout.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from uuid import UUID + +from ptychodus_store.storage.manifest import MANIFEST_FILENAME, ResourceKind + +_KIND_VALUES: frozenset[str] = frozenset(k.value for k in ResourceKind) + + +class LayoutError(ValueError): + """Raised when a path cannot be mapped to the store layout.""" + + +@dataclass(frozen=True) +class ResourceLocation: + """A resolved <store_root>/<kind>/<uuid>/ folder.""" + + kind: str + uuid: UUID + folder: Path + + @property + def manifest_path(self) -> Path: + return self.folder / MANIFEST_FILENAME + + +class StoreLayout: + """Resolves and validates paths within the store root.""" + + def __init__(self, store_root: Path) -> None: + self._root = store_root.resolve() + + @property + def root(self) -> Path: + return self._root + + def kind_dir(self, kind: str) -> Path: + if kind not in _KIND_VALUES: + raise LayoutError(f'unknown kind {kind!r}') + return self._root / kind + + def resource_folder(self, kind: str, uuid: UUID) -> Path: + return self.kind_dir(kind) / str(uuid) + + def manifest_path(self, kind: str, uuid: UUID) -> Path: + return self.resource_folder(kind, uuid) / MANIFEST_FILENAME + + def ensure_kind_dirs(self) -> None: + """Create the four top-level kind directories if they do not exist.""" + for kind in ResourceKind: + (self._root / kind).mkdir(parents=True, exist_ok=True) + + def parse_manifest_path(self, manifest_path: Path) -> ResourceLocation: + """Map a `manifest.json` path to its (kind, uuid, folder). + + Expected shape: `<store_root>/<kind>/<uuid>/manifest.json`. + """ + path = manifest_path.resolve() + if path.name != MANIFEST_FILENAME: + raise LayoutError(f'not a manifest file: {manifest_path}') + + folder = path.parent + kind_dir = folder.parent + + try: + kind_dir.relative_to(self._root) + except ValueError as exc: + raise LayoutError(f'{manifest_path} is outside store root {self._root}') from exc + + if kind_dir.parent != self._root: + raise LayoutError( + f'{manifest_path} is not at the expected depth ' + f'<store_root>/<kind>/<uuid>/manifest.json' + ) + + kind = kind_dir.name + if kind not in _KIND_VALUES: + raise LayoutError(f'{manifest_path} sits under unknown kind {kind!r}') + + try: + uuid = UUID(folder.name) + except ValueError as exc: + raise LayoutError( + f'{manifest_path} parent folder name {folder.name!r} is not a UUID' + ) from exc + + return ResourceLocation(kind=kind, uuid=uuid, folder=folder) + + def iter_manifest_paths(self) -> list[Path]: + """Return all `manifest.json` paths under the store, campaigns first.""" + paths: list[Path] = [] + for kind in ( + ResourceKind.CAMPAIGN, + ResourceKind.DIFFRACTION, + ResourceKind.PRODUCT, + ResourceKind.FLUORESCENCE, + ): + kind_dir = self._root / kind + if not kind_dir.is_dir(): + continue + for child in sorted(kind_dir.iterdir()): + if not child.is_dir(): + continue + m = child / MANIFEST_FILENAME + if m.is_file(): + paths.append(m) + return paths diff --git a/src/ptychodus_store/storage/manifest.py b/src/ptychodus_store/storage/manifest.py new file mode 100644 index 000000000..90085038f --- /dev/null +++ b/src/ptychodus_store/storage/manifest.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import json +from datetime import datetime +from enum import StrEnum +from pathlib import Path +from typing import Annotated, Any, Literal, Union +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator, model_validator + +from ptychodus.api.diffraction import Polarization + +MANIFEST_FILENAME = 'manifest.json' + + +class ResourceKind(StrEnum): + CAMPAIGN = 'campaign' + DIFFRACTION = 'diffraction' + PRODUCT = 'product' + FLUORESCENCE = 'fluorescence' + + +# Kinds that may appear as a derived_from target. Campaign is excluded by design +# (campaign is context, not derivation). +DERIVATION_TARGET_KINDS: frozenset[str] = frozenset( + {ResourceKind.DIFFRACTION, ResourceKind.PRODUCT, ResourceKind.FLUORESCENCE} +) + +# HDF5-owned keys per kind: a manifest MUST NOT carry these (single-source rule). +# Values listed here are read from the companion HDF5 file at reconciliation time. +HDF5_OWNED_KEYS: dict[str, frozenset[str]] = { + ResourceKind.CAMPAIGN: frozenset(), + ResourceKind.DIFFRACTION: frozenset( + { + 'pattern_dtype', + 'pattern_shape', + 'num_patterns_total', + 'detector_pixel_width_m', + 'detector_pixel_height_m', + } + ), + ResourceKind.PRODUCT: frozenset( + { + 'name', + 'comments', + 'detector_distance_m', + 'probe_energy_eV', + 'probe_photon_count', + 'exposure_time_s', + 'mass_attenuation_m2_kg', + 'tomography_angle_deg', + 'tilt_angle_deg', + 'polarization', + 'object_shape', + 'object_pixel_width_m', + 'object_pixel_height_m', + 'probe_shape', + 'num_scan_points', + 'num_loss_epochs', + } + ), + ResourceKind.FLUORESCENCE: frozenset({'element_names', 'map_shape'}), +} + + +class DerivedFromRef(BaseModel): + """A typed pointer to another resource that this node was derived from.""" + + model_config = ConfigDict(extra='forbid') + + kind: Literal['diffraction', 'product', 'fluorescence'] + uuid: UUID + + +class _ManifestBase(BaseModel): + """Fields common to every kind.""" + + model_config = ConfigDict(extra='forbid') + + schema_version: Literal[1] = 1 + uuid: UUID + created_at: datetime + extra: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode='after') + def _reject_hdf5_owned_keys_in_extra(self) -> _ManifestBase: + owned = HDF5_OWNED_KEYS.get(self.kind, frozenset()) # type: ignore[attr-defined] + clashes = sorted(owned & set(self.extra.keys())) + if clashes: + raise ValueError( + f'manifest extra carries HDF5-owned key(s) {clashes}; ' + 'these values must live in the HDF5 file only (single source of truth).' + ) + return self + + +class CampaignManifest(_ManifestBase): + kind: Literal['campaign'] = 'campaign' + label: str = '' + comments: str = '' + sample_name: str = '' + sample_description: str = '' + tags: list[str] = Field(default_factory=list) + + +class DiffractionManifest(_ManifestBase): + kind: Literal['diffraction'] = 'diffraction' + label: str = '' + comments: str = '' + campaign_uuid: UUID | None = None + derived_from: list[DerivedFromRef] = Field(default_factory=list) + detector_distance_m: float | None = None + probe_energy_eV: float | None = None # noqa: N815 + probe_photon_count: int | None = None + exposure_time_s: float | None = None + tomography_angle_deg: float | None = None + tilt_angle_deg: float | None = None + polarization: Polarization | None = None + crop_center_x_px: int | None = None + crop_center_y_px: int | None = None + files: dict[str, str] = Field(default_factory=lambda: {'diffraction': 'diffraction.h5'}) + + @field_validator('derived_from') + @classmethod + def _no_self_ref(cls, v: list[DerivedFromRef]) -> list[DerivedFromRef]: + return v + + +class ProductManifest(_ManifestBase): + kind: Literal['product'] = 'product' + derived_from: list[DerivedFromRef] = Field(default_factory=list) + files: dict[str, str] = Field(default_factory=lambda: {'product': 'product.h5'}) + + +class FluorescenceManifest(_ManifestBase): + kind: Literal['fluorescence'] = 'fluorescence' + label: str = '' + comments: str = '' + derived_from: list[DerivedFromRef] = Field(default_factory=list) + files: dict[str, str] = Field(default_factory=lambda: {'fluorescence': 'fluorescence.h5'}) + + +Manifest = Annotated[ + Union[CampaignManifest, DiffractionManifest, ProductManifest, FluorescenceManifest], + Field(discriminator='kind'), +] + + +_KIND_TO_MODEL: dict[str, type[_ManifestBase]] = { + ResourceKind.CAMPAIGN: CampaignManifest, + ResourceKind.DIFFRACTION: DiffractionManifest, + ResourceKind.PRODUCT: ProductManifest, + ResourceKind.FLUORESCENCE: FluorescenceManifest, +} + + +class ManifestLoadError(Exception): + """Raised when a manifest fails to parse, validate, or pass cross-checks.""" + + +def load_manifest(path: Path, *, expected_kind: str, expected_uuid: UUID) -> _ManifestBase: + """Read and validate a manifest.json from `path` against the expected kind and folder UUID. + + The manifest's `kind` must match `expected_kind` (derived from the folder layout), + its `uuid` must equal `expected_uuid` (the folder name), and `derived_from` entries + must not self-reference. + """ + try: + raw = json.loads(path.read_text(encoding='utf-8')) + except (OSError, json.JSONDecodeError) as exc: + raise ManifestLoadError(f'failed to read or parse {path}: {exc}') from exc + + declared_kind = raw.get('kind') + if declared_kind != expected_kind: + raise ManifestLoadError( + f'{path}: kind={declared_kind!r} does not match folder kind {expected_kind!r}' + ) + + model_cls = _KIND_TO_MODEL[expected_kind] + try: + manifest = model_cls.model_validate(raw) + except ValidationError as exc: + raise ManifestLoadError(f'{path}: validation error: {exc}') from exc + + if manifest.uuid != expected_uuid: + raise ManifestLoadError( + f'{path}: uuid={manifest.uuid} does not match folder name {expected_uuid}' + ) + + derived = getattr(manifest, 'derived_from', None) + if derived: + for ref in derived: + if ref.uuid == manifest.uuid: + raise ManifestLoadError( + f'{path}: derived_from entry self-references uuid {ref.uuid}' + ) + + return manifest diff --git a/src/ptychodus_store/ui/icons/Font-Awesome-LICENSE.txt b/src/ptychodus_store/ui/icons/Font-Awesome-LICENSE.txt new file mode 100644 index 000000000..45063c11e --- /dev/null +++ b/src/ptychodus_store/ui/icons/Font-Awesome-LICENSE.txt @@ -0,0 +1,165 @@ +Fonticons, Inc. (https://fontawesome.com) + +-------------------------------------------------------------------------------- + +Font Awesome Free License + +Font Awesome Free is free, open source, and GPL friendly. You can use it for +commercial projects, open source projects, or really almost whatever you want. +Full Font Awesome Free license: https://fontawesome.com/license/free. + +-------------------------------------------------------------------------------- + +# Icons: CC BY 4.0 License (https://creativecommons.org/licenses/by/4.0/) + +The Font Awesome Free download is licensed under a Creative Commons +Attribution 4.0 International License and applies to all icons packaged +as SVG and JS file types. + +-------------------------------------------------------------------------------- + +# Fonts: SIL OFL 1.1 License + +In the Font Awesome Free download, the SIL OFL license applies to all icons +packaged as web and desktop font files. + +Copyright (c) 2025 Fonticons, Inc. (https://fontawesome.com) +with Reserved Font Name: "Font Awesome". + +This Font Software is licensed under the SIL Open Font License, Version 1.1. +This license is copied below, and is also available with a FAQ at: +http://scripts.sil.org/OFL + +SIL OPEN FONT LICENSE +Version 1.1 - 26 February 2007 + +PREAMBLE +The goals of the Open Font License (OFL) are to stimulate worldwide +development of collaborative font projects, to support the font creation +efforts of academic and linguistic communities, and to provide a free and +open framework in which fonts may be shared and improved in partnership +with others. + +The OFL allows the licensed fonts to be used, studied, modified and +redistributed freely as long as they are not sold by themselves. The +fonts, including any derivative works, can be bundled, embedded, +redistributed and/or sold with any software provided that any reserved +names are not used by derivative works. The fonts and derivatives, +however, cannot be released under any other type of license. The +requirement for fonts to remain under this license does not apply +to any document created using the fonts or their derivatives. + +DEFINITIONS +"Font Software" refers to the set of files released by the Copyright +Holder(s) under this license and clearly marked as such. This may +include source files, build scripts and documentation. + +"Reserved Font Name" refers to any names specified as such after the +copyright statement(s). + +"Original Version" refers to the collection of Font Software components as +distributed by the Copyright Holder(s). + +"Modified Version" refers to any derivative made by adding to, deleting, +or substituting — in part or in whole — any of the components of the +Original Version, by changing formats or by porting the Font Software to a +new environment. + +"Author" refers to any designer, engineer, programmer, technical +writer or other person who contributed to the Font Software. + +PERMISSION & CONDITIONS +Permission is hereby granted, free of charge, to any person obtaining +a copy of the Font Software, to use, study, copy, merge, embed, modify, +redistribute, and sell modified and unmodified copies of the Font +Software, subject to the following conditions: + +1) Neither the Font Software nor any of its individual components, +in Original or Modified Versions, may be sold by itself. + +2) Original or Modified Versions of the Font Software may be bundled, +redistributed and/or sold with any software, provided that each copy +contains the above copyright notice and this license. These can be +included either as stand-alone text files, human-readable headers or +in the appropriate machine-readable metadata fields within text or +binary files as long as those fields can be easily viewed by the user. + +3) No Modified Version of the Font Software may use the Reserved Font +Name(s) unless explicit written permission is granted by the corresponding +Copyright Holder. This restriction only applies to the primary font name as +presented to the users. + +4) The name(s) of the Copyright Holder(s) or the Author(s) of the Font +Software shall not be used to promote, endorse or advertise any +Modified Version, except to acknowledge the contribution(s) of the +Copyright Holder(s) and the Author(s) or with their explicit written +permission. + +5) The Font Software, modified or unmodified, in part or in whole, +must be distributed entirely under this license, and must not be +distributed under any other license. The requirement for fonts to +remain under this license does not apply to any document created +using the Font Software. + +TERMINATION +This license becomes null and void if any of the above conditions are +not met. + +DISCLAIMER +THE FONT SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO ANY WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT +OF COPYRIGHT, PATENT, TRADEMARK, OR OTHER RIGHT. IN NO EVENT SHALL THE +COPYRIGHT HOLDER BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, +INCLUDING ANY GENERAL, SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL +DAMAGES, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF THE USE OR INABILITY TO USE THE FONT SOFTWARE OR FROM +OTHER DEALINGS IN THE FONT SOFTWARE. + +-------------------------------------------------------------------------------- + +# Code: MIT License (https://opensource.org/licenses/MIT) + +In the Font Awesome Free download, the MIT license applies to all non-font and +non-icon files. + +Copyright 2025 Fonticons, Inc. + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in the +Software without restriction, including without limitation the rights to use, copy, +modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, +and to permit persons to whom the Software is furnished to do so, subject to the +following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, +INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT +HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +-------------------------------------------------------------------------------- + +# Attribution + +Attribution is required by MIT, SIL OFL, and CC BY licenses. Downloaded Font +Awesome Free files already contain embedded comments with sufficient +attribution, so you shouldn't need to do anything additional when using these +files normally. + +We've kept attribution comments terse, so we ask that you do not actively work +to remove them from files, especially code. They're a great way for folks to +learn about Font Awesome. + +-------------------------------------------------------------------------------- + +# Brand Icons + +All brand icons are trademarks of their respective owners. The use of these +trademarks does not indicate endorsement of the trademark holder by Font +Awesome, nor vice versa. **Please do not use brand logos for any purpose except +to represent the company, product, or service to which they refer.** diff --git a/src/ptychodus_store/ui/icons/README.md b/src/ptychodus_store/ui/icons/README.md new file mode 100644 index 000000000..a6a824d62 --- /dev/null +++ b/src/ptychodus_store/ui/icons/README.md @@ -0,0 +1,26 @@ +# UI icons + +Single source of truth for icon SVGs used by both the ptychodus web UI (`src/ptychodus_store/ui/`) and the PyQt GUI (via `src/ptychodus/view/resources.qrc`, which uses relative paths that point back here). + +## Files + +- `ptychodus.svg`, `genesis.svg`, `globus.svg` — project-owned SVGs. +- Everything else — a subset of [Font-Awesome 7.1.0](https://fontawesome.com/) (Free), CC BY 4.0. See `Font-Awesome-LICENSE.txt`. + +Filenames are the original Font-Awesome names, with no `solid/` vs `regular/` prefix (the subset in use has no name collisions). Serving is flat: `/ui/icons/<name>.svg`. + +## Adding a new icon + +1. Drop the SVG file into this directory (from Font-Awesome or elsewhere). +2. Update `src/ptychodus_store/ui/src/nav.ts` (or whichever component uses it) to reference `/ui/icons/<name>.svg`. +3. If the PyQt GUI also needs it, add a line to `src/ptychodus/view/resources.qrc`: + + ```xml + <file alias="<alias>">../../ptychodus_store/ui/icons/<name>.svg</file> + ``` + + then run `src/ptychodus/view/make_qrc.sh` to regenerate `resources.py`. + +## Updating Font-Awesome + +Fetch a newer release from <https://github.com/FortAwesome/Font-Awesome>, replace the SVGs referenced above by name (Font-Awesome preserves filenames across minor releases), and refresh `Font-Awesome-LICENSE.txt`. Rerun `make_qrc.sh` if `resources.qrc` references changed files. diff --git a/src/ptychodus_store/ui/icons/arrows-left-right-to-line.svg b/src/ptychodus_store/ui/icons/arrows-left-right-to-line.svg new file mode 100644 index 000000000..23da5710b --- /dev/null +++ b/src/ptychodus_store/ui/icons/arrows-left-right-to-line.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 576 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M32 96C14.3 96 0 110.3 0 128L0 384c0 17.7 14.3 32 32 32s32-14.3 32-32l0-256c0-17.7-14.3-32-32-32zM390.6 342.6l64-64c12.5-12.5 12.5-32.8 0-45.3l-64-64c-12.5-12.5-32.8-12.5-45.3 0s-12.5 32.8 0 45.3l9.4 9.4-133.5 0 9.4-9.4c12.5-12.5 12.5-32.8 0-45.3s-32.8-12.5-45.3 0l-64 64c-6 6-9.4 14.1-9.4 22.6s3.4 16.6 9.4 22.6l64 64c12.5 12.5 32.8 12.5 45.3 0s12.5-32.8 0-45.3l-9.4-9.4 133.5 0-9.4 9.4c-12.5 12.5-12.5 32.8 0 45.3s32.8 12.5 45.3 0zM576 128c0-17.7-14.3-32-32-32s-32 14.3-32 32l0 256c0 17.7 14.3 32 32 32s32-14.3 32-32l0-256z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/arrows-up-down-left-right.svg b/src/ptychodus_store/ui/icons/arrows-up-down-left-right.svg new file mode 100644 index 000000000..f610fab26 --- /dev/null +++ b/src/ptychodus_store/ui/icons/arrows-up-down-left-right.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M278.6 9.4c-12.5-12.5-32.8-12.5-45.3 0l-64 64c-12.5 12.5-12.5 32.8 0 45.3s32.8 12.5 45.3 0l9.4-9.4 0 114.7-114.7 0 9.4-9.4c12.5-12.5 12.5-32.8 0-45.3s-32.8-12.5-45.3 0l-64 64c-12.5 12.5-12.5 32.8 0 45.3l64 64c12.5 12.5 32.8 12.5 45.3 0s12.5-32.8 0-45.3l-9.4-9.4 114.7 0 0 114.7-9.4-9.4c-12.5-12.5-32.8-12.5-45.3 0s-12.5 32.8 0 45.3l64 64c12.5 12.5 32.8 12.5 45.3 0l64-64c12.5-12.5 12.5-32.8 0-45.3s-32.8-12.5-45.3 0l-9.4 9.4 0-114.7 114.7 0-9.4 9.4c-12.5 12.5-12.5 32.8 0 45.3s32.8 12.5 45.3 0l64-64c12.5-12.5 12.5-32.8 0-45.3l-64-64c-12.5-12.5-32.8-12.5-45.3 0s-12.5 32.8 0 45.3l9.4 9.4-114.7 0 0-114.7 9.4 9.4c12.5 12.5 32.8 12.5 45.3 0s12.5-32.8 0-45.3l-64-64z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/atom.svg b/src/ptychodus_store/ui/icons/atom.svg new file mode 100644 index 000000000..b665517f9 --- /dev/null +++ b/src/ptychodus_store/ui/icons/atom.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M224 398.8c-11.8 5.1-23.4 9.7-34.9 13.5 16.7 33.8 31 35.7 34.9 35.7s18.1-1.9 34.9-35.7c-11.4-3.9-23.1-8.4-34.9-13.5zM414 256c33 45.2 44.3 90.9 23.6 128-20.2 36.3-62.5 49.3-115.2 43.2-22 52.1-55.7 84.8-98.4 84.8s-76.4-32.7-98.4-84.8C72.9 433.3 30.6 420.3 10.4 384-10.3 346.9 1 301.2 34 256 1 210.8-10.3 165.1 10.4 128 30.6 91.7 72.9 78.7 125.6 84.8 147.6 32.7 181.2 0 224 0s76.4 32.7 98.4 84.8c52.7-6.1 95 6.8 115.2 43.2 20.7 37.1 9.4 82.8-23.6 128zm-65.8 67.4c-1.7 14.2-3.9 28-6.7 41.2 31.8 1.4 38.6-8.7 40.2-11.7 2.3-4.2 7-17.9-11.9-48.1-6.8 6.3-14 12.5-21.6 18.6zm-6.7-175.9c2.8 13.1 5 26.9 6.7 41.2 7.6 6.1 14.8 12.3 21.6 18.6 18.9-30.2 14.2-44 11.9-48.1-1.6-2.9-8.4-13-40.2-11.7zM258.9 99.7C242.1 65.9 227.9 64 224 64s-18.1 1.9-34.9 35.7c11.4 3.9 23.1 8.4 34.9 13.5 11.8-5.1 23.4-9.7 34.9-13.5zm-159 88.9c1.7-14.3 3.9-28 6.7-41.2-31.8-1.4-38.6 8.7-40.2 11.7-2.3 4.2-7 17.9 11.9 48.1 6.8-6.3 14-12.5 21.6-18.6zM78.2 304.8c-18.9 30.2-14.2 44-11.9 48.1 1.6 2.9 8.4 13 40.2 11.7-2.8-13.1-5-26.9-6.7-41.2-7.6-6.1-14.8-12.3-21.6-18.6zM304 256a80 80 0 1 0 -160 0 80 80 0 1 0 160 0zm-80-32a32 32 0 1 1 0 64 32 32 0 1 1 0-64z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/chart-line.svg b/src/ptychodus_store/ui/icons/chart-line.svg new file mode 100644 index 000000000..17a89ee76 --- /dev/null +++ b/src/ptychodus_store/ui/icons/chart-line.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M64 64c0-17.7-14.3-32-32-32S0 46.3 0 64L0 400c0 44.2 35.8 80 80 80l400 0c17.7 0 32-14.3 32-32s-14.3-32-32-32L80 416c-8.8 0-16-7.2-16-16L64 64zm406.6 86.6c12.5-12.5 12.5-32.8 0-45.3s-32.8-12.5-45.3 0L320 210.7 262.6 153.4c-12.5-12.5-32.8-12.5-45.3 0l-96 96c-12.5 12.5-12.5 32.8 0 45.3s32.8 12.5 45.3 0l73.4-73.4 57.4 57.4c12.5 12.5 32.8 12.5 45.3 0l128-128z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/circle-radiation.svg b/src/ptychodus_store/ui/icons/circle-radiation.svg new file mode 100644 index 000000000..3d614c17d --- /dev/null +++ b/src/ptychodus_store/ui/icons/circle-radiation.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M0 256a256 256 0 1 1 512 0 256 256 0 1 1 -512 0zm80 0l64.3 0c8.7 0 15.7-7.1 17.3-15.6 4.4-24.4 18.1-45.5 37.2-59.7 7.4-5.5 10.6-15.6 6-23.6l-32.5-56.3c-4.3-7.5-13.9-10.3-21.2-5.5-48.2 31.5-81.3 84.2-86.3 144.8-.7 8.8 6.5 16 15.3 16zm137.9 89.8c-8.5-3.7-18.8-1.4-23.5 6.6l-31 53.8c-4.3 7.5-1.9 17.2 5.8 21.1 26.1 13.2 55.5 20.7 86.8 20.7s60.7-7.5 86.8-20.7c7.7-3.9 10.1-13.6 5.8-21.1l-31-53.8c-4.6-8-15-10.3-23.5-6.6-11.7 5-24.5 7.8-38.1 7.8s-26.4-2.8-38.1-7.8zM350.4 240.4c1.6 8.6 8.5 15.6 17.3 15.6l64.3 0c8.8 0 16.1-7.2 15.3-16-5-60.6-38.1-113.2-86.3-144.8-7.3-4.8-16.8-2-21.2 5.5L307.3 157c-4.6 8-1.4 18.1 6 23.6 19.1 14.2 32.7 35.4 37.2 59.7zM256 305.7a48 48 0 1 0 0-96 48 48 0 1 0 0 96z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/f.svg b/src/ptychodus_store/ui/icons/f.svg new file mode 100644 index 000000000..868c897b9 --- /dev/null +++ b/src/ptychodus_store/ui/icons/f.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 320 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M32 32C14.3 32 0 46.3 0 64L0 448c0 17.7 14.3 32 32 32s32-14.3 32-32l0-160 160 0c17.7 0 32-14.3 32-32s-14.3-32-32-32l-160 0 0-128 224 0c17.7 0 32-14.3 32-32s-14.3-32-32-32L32 32z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/floppy-disk.svg b/src/ptychodus_store/ui/icons/floppy-disk.svg new file mode 100644 index 000000000..9d0cf2019 --- /dev/null +++ b/src/ptychodus_store/ui/icons/floppy-disk.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M64 80c-8.8 0-16 7.2-16 16l0 320c0 8.8 7.2 16 16 16l320 0c8.8 0 16-7.2 16-16l0-242.7c0-4.2-1.7-8.3-4.7-11.3L320 86.6 320 176c0 17.7-14.3 32-32 32l-160 0c-17.7 0-32-14.3-32-32l0-96-32 0zm80 0l0 80 128 0 0-80-128 0zM0 96C0 60.7 28.7 32 64 32l242.7 0c17 0 33.3 6.7 45.3 18.7L429.3 128c12 12 18.7 28.3 18.7 45.3L448 416c0 35.3-28.7 64-64 64L64 480c-35.3 0-64-28.7-64-64L0 96zM160 320a64 64 0 1 1 128 0 64 64 0 1 1 -128 0z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/gear.svg b/src/ptychodus_store/ui/icons/gear.svg new file mode 100644 index 000000000..33b418c82 --- /dev/null +++ b/src/ptychodus_store/ui/icons/gear.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M195.1 9.5C198.1-5.3 211.2-16 226.4-16l59.8 0c15.2 0 28.3 10.7 31.3 25.5L332 79.5c14.1 6 27.3 13.7 39.3 22.8l67.8-22.5c14.4-4.8 30.2 1.2 37.8 14.4l29.9 51.8c7.6 13.2 4.9 29.8-6.5 39.9L447 233.3c.9 7.4 1.3 15 1.3 22.7s-.5 15.3-1.3 22.7l53.4 47.5c11.4 10.1 14 26.8 6.5 39.9l-29.9 51.8c-7.6 13.1-23.4 19.2-37.8 14.4l-67.8-22.5c-12.1 9.1-25.3 16.7-39.3 22.8l-14.4 69.9c-3.1 14.9-16.2 25.5-31.3 25.5l-59.8 0c-15.2 0-28.3-10.7-31.3-25.5l-14.4-69.9c-14.1-6-27.2-13.7-39.3-22.8L73.5 432.3c-14.4 4.8-30.2-1.2-37.8-14.4L5.8 366.1c-7.6-13.2-4.9-29.8 6.5-39.9l53.4-47.5c-.9-7.4-1.3-15-1.3-22.7s.5-15.3 1.3-22.7L12.3 185.8c-11.4-10.1-14-26.8-6.5-39.9L35.7 94.1c7.6-13.2 23.4-19.2 37.8-14.4l67.8 22.5c12.1-9.1 25.3-16.7 39.3-22.8L195.1 9.5zM256.3 336a80 80 0 1 0 -.6-160 80 80 0 1 0 .6 160z"/></svg> \ No newline at end of file diff --git a/genesis.svg b/src/ptychodus_store/ui/icons/genesis.svg similarity index 100% rename from genesis.svg rename to src/ptychodus_store/ui/icons/genesis.svg diff --git a/globus.svg b/src/ptychodus_store/ui/icons/globus.svg similarity index 100% rename from globus.svg rename to src/ptychodus_store/ui/icons/globus.svg diff --git a/src/ptychodus_store/ui/icons/house-chimney.svg b/src/ptychodus_store/ui/icons/house-chimney.svg new file mode 100644 index 000000000..a69ec6cfe --- /dev/null +++ b/src/ptychodus_store/ui/icons/house-chimney.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M234.2 8.6c12.3-11.4 31.3-11.4 43.5 0L368 92.3 368 80c0-17.7 14.3-32 32-32l32 0c17.7 0 32 14.3 32 32l0 101.5 37.8 35.1c9.6 9 12.8 22.9 8 35.1S493.2 272 480 272l-16 0 0 176c0 35.3-28.7 64-64 64l-288 0c-35.3 0-64-28.7-64-64l0-176-16 0c-13.2 0-25-8.1-29.8-20.3s-1.6-26.2 8-35.1l224-208zM240 320c-26.5 0-48 21.5-48 48l0 96 128 0 0-96c0-26.5-21.5-48-48-48l-32 0z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/layer-group.svg b/src/ptychodus_store/ui/icons/layer-group.svg new file mode 100644 index 000000000..de2e0950e --- /dev/null +++ b/src/ptychodus_store/ui/icons/layer-group.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M232.5 5.2c14.9-6.9 32.1-6.9 47 0l218.6 101c8.5 3.9 13.9 12.4 13.9 21.8s-5.4 17.9-13.9 21.8l-218.6 101c-14.9 6.9-32.1 6.9-47 0L13.9 149.8C5.4 145.8 0 137.3 0 128s5.4-17.9 13.9-21.8L232.5 5.2zM48.1 218.4l164.3 75.9c27.7 12.8 59.6 12.8 87.3 0l164.3-75.9 34.1 15.8c8.5 3.9 13.9 12.4 13.9 21.8s-5.4 17.9-13.9 21.8l-218.6 101c-14.9 6.9-32.1 6.9-47 0L13.9 277.8C5.4 273.8 0 265.3 0 256s5.4-17.9 13.9-21.8l34.1-15.8zM13.9 362.2l34.1-15.8 164.3 75.9c27.7 12.8 59.6 12.8 87.3 0l164.3-75.9 34.1 15.8c8.5 3.9 13.9 12.4 13.9 21.8s-5.4 17.9-13.9 21.8l-218.6 101c-14.9 6.9-32.1 6.9-47 0L13.9 405.8C5.4 401.8 0 393.3 0 384s5.4-17.9 13.9-21.8z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/list.svg b/src/ptychodus_store/ui/icons/list.svg new file mode 100644 index 000000000..81c434fd0 --- /dev/null +++ b/src/ptychodus_store/ui/icons/list.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M40 48C26.7 48 16 58.7 16 72l0 48c0 13.3 10.7 24 24 24l48 0c13.3 0 24-10.7 24-24l0-48c0-13.3-10.7-24-24-24L40 48zM192 64c-17.7 0-32 14.3-32 32s14.3 32 32 32l288 0c17.7 0 32-14.3 32-32s-14.3-32-32-32L192 64zm0 160c-17.7 0-32 14.3-32 32s14.3 32 32 32l288 0c17.7 0 32-14.3 32-32s-14.3-32-32-32l-288 0zm0 160c-17.7 0-32 14.3-32 32s14.3 32 32 32l288 0c17.7 0 32-14.3 32-32s-14.3-32-32-32l-288 0zM16 232l0 48c0 13.3 10.7 24 24 24l48 0c13.3 0 24-10.7 24-24l0-48c0-13.3-10.7-24-24-24l-48 0c-13.3 0-24 10.7-24 24zM40 368c-13.3 0-24 10.7-24 24l0 48c0 13.3 10.7 24 24 24l48 0c13.3 0 24-10.7 24-24l0-48c0-13.3-10.7-24-24-24l-48 0z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/microchip.svg b/src/ptychodus_store/ui/icons/microchip.svg new file mode 100644 index 000000000..0d6b73e70 --- /dev/null +++ b/src/ptychodus_store/ui/icons/microchip.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M176 24c0-13.3-10.7-24-24-24s-24 10.7-24 24l0 40c-35.3 0-64 28.7-64 64l-40 0c-13.3 0-24 10.7-24 24s10.7 24 24 24l40 0 0 56-40 0c-13.3 0-24 10.7-24 24s10.7 24 24 24l40 0 0 56-40 0c-13.3 0-24 10.7-24 24s10.7 24 24 24l40 0c0 35.3 28.7 64 64 64l0 40c0 13.3 10.7 24 24 24s24-10.7 24-24l0-40 56 0 0 40c0 13.3 10.7 24 24 24s24-10.7 24-24l0-40 56 0 0 40c0 13.3 10.7 24 24 24s24-10.7 24-24l0-40c35.3 0 64-28.7 64-64l40 0c13.3 0 24-10.7 24-24s-10.7-24-24-24l-40 0 0-56 40 0c13.3 0 24-10.7 24-24s-10.7-24-24-24l-40 0 0-56 40 0c13.3 0 24-10.7 24-24s-10.7-24-24-24l-40 0c0-35.3-28.7-64-64-64l0-40c0-13.3-10.7-24-24-24s-24 10.7-24 24l0 40-56 0 0-40c0-13.3-10.7-24-24-24s-24 10.7-24 24l0 40-56 0 0-40zM160 128l192 0c17.7 0 32 14.3 32 32l0 192c0 17.7-14.3 32-32 32l-192 0c-17.7 0-32-14.3-32-32l0-192c0-17.7 14.3-32 32-32zm16 48l0 160 160 0 0-160-160 0z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/object-group.svg b/src/ptychodus_store/ui/icons/object-group.svg new file mode 100644 index 000000000..8118782ff --- /dev/null +++ b/src/ptychodus_store/ui/icons/object-group.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 576 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M32 119.4C12.9 108.4 0 87.7 0 64 0 28.7 28.7 0 64 0 87.7 0 108.4 12.9 119.4 32l337.1 0c11.1-19.1 31.7-32 55.4-32 35.3 0 64 28.7 64 64 0 23.7-12.9 44.4-32 55.4l0 273.1c19.1 11.1 32 31.7 32 55.4 0 35.3-28.7 64-64 64-23.7 0-44.4-12.9-55.4-32l-337.1 0c-11.1 19.1-31.7 32-55.4 32-35.3 0-64-28.7-64-64 0-23.7 12.9-44.4 32-55.4l0-273.1zm448 0c-9.7-5.6-17.8-13.7-23.4-23.4L119.4 96c-5.6 9.7-13.7 17.8-23.4 23.4l0 273.1c9.7 5.6 17.8 13.7 23.4 23.4l337.1 0c5.6-9.7 13.7-17.8 23.4-23.4l0-273.1zM144 176c0-17.7 14.3-32 32-32l112 0c17.7 0 32 14.3 32 32l0 64c0 17.7-14.3 32-32 32l-112 0c-17.7 0-32-14.3-32-32l0-64zM256 320l32 0c44.2 0 80-35.8 80-80l32 0c17.7 0 32 14.3 32 32l0 64c0 17.7-14.3 32-32 32l-112 0c-17.7 0-32-14.3-32-32l0-16z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/paper-plane.svg b/src/ptychodus_store/ui/icons/paper-plane.svg new file mode 100644 index 000000000..4e7f49691 --- /dev/null +++ b/src/ptychodus_store/ui/icons/paper-plane.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 576 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M536.4-26.3c9.8-3.5 20.6-1 28 6.3s9.8 18.2 6.3 28l-178 496.9c-5 13.9-18.1 23.1-32.8 23.1-14.2 0-27-8.6-32.3-21.7l-64.2-158c-4.5-11-2.5-23.6 5.2-32.6l94.5-112.4c5.1-6.1 4.7-15-.9-20.6s-14.6-6-20.6-.9L229.2 276.1c-9.1 7.6-21.6 9.6-32.6 5.2L38.1 216.8c-13.1-5.3-21.7-18.1-21.7-32.3 0-14.7 9.2-27.8 23.1-32.8l496.9-178z"/></svg> \ No newline at end of file diff --git a/ptychodus.svg b/src/ptychodus_store/ui/icons/ptychodus.svg similarity index 100% rename from ptychodus.svg rename to src/ptychodus_store/ui/icons/ptychodus.svg diff --git a/src/ptychodus_store/ui/icons/robot.svg b/src/ptychodus_store/ui/icons/robot.svg new file mode 100644 index 000000000..dc5b5f9fd --- /dev/null +++ b/src/ptychodus_store/ui/icons/robot.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 640 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M352 0c0-17.7-14.3-32-32-32S288-17.7 288 0l0 64-96 0c-53 0-96 43-96 96l0 224c0 53 43 96 96 96l256 0c53 0 96-43 96-96l0-224c0-53-43-96-96-96l-96 0 0-64zM160 368c0-13.3 10.7-24 24-24l32 0c13.3 0 24 10.7 24 24s-10.7 24-24 24l-32 0c-13.3 0-24-10.7-24-24zm120 0c0-13.3 10.7-24 24-24l32 0c13.3 0 24 10.7 24 24s-10.7 24-24 24l-32 0c-13.3 0-24-10.7-24-24zm120 0c0-13.3 10.7-24 24-24l32 0c13.3 0 24 10.7 24 24s-10.7 24-24 24l-32 0c-13.3 0-24-10.7-24-24zM224 176a48 48 0 1 1 0 96 48 48 0 1 1 0-96zm144 48a48 48 0 1 1 96 0 48 48 0 1 1 -96 0zM64 224c0-17.7-14.3-32-32-32S0 206.3 0 224l0 96c0 17.7 14.3 32 32 32s32-14.3 32-32l0-96zm544-32c-17.7 0-32 14.3-32 32l0 96c0 17.7 14.3 32 32 32s32-14.3 32-32l0-96c0-17.7-14.3-32-32-32z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/route.svg b/src/ptychodus_store/ui/icons/route.svg new file mode 100644 index 000000000..075e16968 --- /dev/null +++ b/src/ptychodus_store/ui/icons/route.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M512 96c0 50.2-59.1 125.1-84.6 155-3.8 4.4-9.4 6.1-14.5 5L320 256c-17.7 0-32 14.3-32 32s14.3 32 32 32l96 0c53 0 96 43 96 96s-43 96-96 96l-276.4 0c8.7-9.9 19.3-22.6 30-36.8 6.3-8.4 12.8-17.6 19-27.2L416 448c17.7 0 32-14.3 32-32s-14.3-32-32-32l-96 0c-53 0-96-43-96-96s43-96 96-96l39.8 0c-21-31.5-39.8-67.7-39.8-96 0-53 43-96 96-96s96 43 96 96zM117.1 489.1c-3.8 4.3-7.2 8.1-10.1 11.3l-1.8 2-.2-.2c-6 4.6-14.6 4-20-1.8-25.2-27.4-85-97.9-85-148.4 0-53 43-96 96-96s96 43 96 96c0 30-21.1 67-43.5 97.9-10.7 14.7-21.7 28-30.8 38.5l-.6 .7zM128 352a32 32 0 1 0 -64 0 32 32 0 1 0 64 0zM416 128a32 32 0 1 0 0-64 32 32 0 1 0 0 64z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/ruler.svg b/src/ptychodus_store/ui/icons/ruler.svg new file mode 100644 index 000000000..fec5348ee --- /dev/null +++ b/src/ptychodus_store/ui/icons/ruler.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 576 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M209.1 516.2c-18.7 18.7-49.1 18.7-67.9 0L28.1 403.1c-18.7-18.7-18.7-49.1 0-67.9l17-17 73.5 73.5c9.4 9.4 24.6 9.4 33.9 0s9.4-24.6 0-33.9l-73.5-73.5 33.9-33.9 50.9 50.9c9.4 9.4 24.6 9.4 33.9 0s9.4-24.6 0-33.9l-50.9-50.9 33.9-33.9 73.5 73.5c9.4 9.4 24.6 9.4 33.9 0s9.4-24.6 0-33.9l-73.5-73.5 33.9-33.9 50.9 50.9c9.4 9.4 24.6 9.4 33.9 0s9.4-24.6 0-33.9l-50.9-50.9 33.9-33.9 73.5 73.5c9.4 9.4 24.6 9.4 33.9 0s9.4-24.6 0-33.9l-73.5-73.5 17-17c18.7-18.7 49.1-18.7 67.9 0L548.5 108.9c18.7 18.7 18.7 49.1 0 67.9L209.1 516.2z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/table-cells.svg b/src/ptychodus_store/ui/icons/table-cells.svg new file mode 100644 index 000000000..6f3e4d005 --- /dev/null +++ b/src/ptychodus_store/ui/icons/table-cells.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M384 96l0 64-64 0 0-64 64 0zm0 128l0 64-64 0 0-64 64 0zm0 128l0 64-64 0 0-64 64 0zM256 288l-64 0 0-64 64 0 0 64zm-64 64l64 0 0 64-64 0 0-64zm-64-64l-64 0 0-64 64 0 0 64zM64 352l64 0 0 64-64 0 0-64zm0-192l0-64 64 0 0 64-64 0zm128 0l0-64 64 0 0 64-64 0zM64 32C28.7 32 0 60.7 0 96L0 416c0 35.3 28.7 64 64 64l320 0c35.3 0 64-28.7 64-64l0-320c0-35.3-28.7-64-64-64L64 32z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/icons/wand-magic-sparkles.svg b/src/ptychodus_store/ui/icons/wand-magic-sparkles.svg new file mode 100644 index 000000000..481bbb90a --- /dev/null +++ b/src/ptychodus_store/ui/icons/wand-magic-sparkles.svg @@ -0,0 +1 @@ +<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 576 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc. --><path fill="currentColor" d="M263.4-27L278.2 9.8 315 24.6c3 1.2 5 4.2 5 7.4s-2 6.2-5 7.4L278.2 54.2 263.4 91c-1.2 3-4.2 5-7.4 5s-6.2-2-7.4-5L233.8 54.2 197 39.4c-3-1.2-5-4.2-5-7.4s2-6.2 5-7.4L233.8 9.8 248.6-27c1.2-3 4.2-5 7.4-5s6.2 2 7.4 5zM110.7 41.7l21.5 50.1 50.1 21.5c5.9 2.5 9.7 8.3 9.7 14.7s-3.8 12.2-9.7 14.7l-50.1 21.5-21.5 50.1c-2.5 5.9-8.3 9.7-14.7 9.7s-12.2-3.8-14.7-9.7L59.8 164.2 9.7 142.7C3.8 140.2 0 134.4 0 128s3.8-12.2 9.7-14.7L59.8 91.8 81.3 41.7C83.8 35.8 89.6 32 96 32s12.2 3.8 14.7 9.7zM464 304c6.4 0 12.2 3.8 14.7 9.7l21.5 50.1 50.1 21.5c5.9 2.5 9.7 8.3 9.7 14.7s-3.8 12.2-9.7 14.7l-50.1 21.5-21.5 50.1c-2.5 5.9-8.3 9.7-14.7 9.7s-12.2-3.8-14.7-9.7l-21.5-50.1-50.1-21.5c-5.9-2.5-9.7-8.3-9.7-14.7s3.8-12.2 9.7-14.7l50.1-21.5 21.5-50.1c2.5-5.9 8.3-9.7 14.7-9.7zM460 0c11 0 21.6 4.4 29.5 12.2l42.3 42.3C539.6 62.4 544 73 544 84s-4.4 21.6-12.2 29.5l-88.2 88.2-101.3-101.3 88.2-88.2C438.4 4.4 449 0 460 0zM44.2 398.5L308.4 134.3 409.7 235.6 145.5 499.8C137.6 507.6 127 512 116 512s-21.6-4.4-29.5-12.2L44.2 457.5C36.4 449.6 32 439 32 428s4.4-21.6 12.2-29.5z"/></svg> \ No newline at end of file diff --git a/src/ptychodus_store/ui/index.html b/src/ptychodus_store/ui/index.html new file mode 100644 index 000000000..3a997c902 --- /dev/null +++ b/src/ptychodus_store/ui/index.html @@ -0,0 +1,15 @@ +<!DOCTYPE html> +<html lang="en"> + <head> + <meta charset="utf-8" /> + <meta name="viewport" content="width=device-width, initial-scale=1" /> + <title>ptychodus + + + + + +
+ + + diff --git a/src/ptychodus_store/ui/src/api.ts b/src/ptychodus_store/ui/src/api.ts new file mode 100644 index 000000000..7ed8f60f4 --- /dev/null +++ b/src/ptychodus_store/ui/src/api.ts @@ -0,0 +1,132 @@ +const API = '/api/v1'; + +export type IngestState = 'valid' | 'invalid' | 'pending'; + +export type Polarization = 'left_circular' | 'right_circular'; + +export interface DerivedFromEdge { + kind: 'diffraction' | 'product' | 'fluorescence'; + uuid: string; +} + +export interface Page { + items: T[]; + total: number; + limit: number; + offset: number; +} + +export interface DiffractionRead { + uuid: string; + label: string; + comments: string; + ingest_state: IngestState; + campaign_uuid: string | null; + derived_from: DerivedFromEdge[]; + probe_energy_eV: number | null; + probe_photon_count: number | null; + tomography_angle_deg: number | null; + tilt_angle_deg: number | null; + polarization: Polarization | null; + pattern_dtype: string | null; + pattern_height_px: number | null; + pattern_width_px: number | null; + num_patterns_total: number | null; + detector_pixel_width_m: number | null; + detector_pixel_height_m: number | null; +} + +export interface ProductRead { + uuid: string; + ingest_state: IngestState; + derived_from: DerivedFromEdge[]; + name: string | null; + comments: string | null; + detector_distance_m: number | null; + probe_energy_eV: number | null; + probe_photon_count: number | null; + tomography_angle_deg: number | null; + tilt_angle_deg: number | null; + polarization: Polarization | null; + object_layers: number | null; + object_height_px: number | null; + object_width_px: number | null; + object_pixel_width_m: number | null; + object_pixel_height_m: number | null; + probe_modes: number | null; + probe_height_px: number | null; + probe_width_px: number | null; + num_scan_points: number | null; +} + +export interface FluorescenceRead { + uuid: string; + label: string; + comments: string; + ingest_state: IngestState; + derived_from: DerivedFromEdge[]; + element_names: string[]; + map_height_px: number | null; + map_width_px: number | null; +} + +export interface RenderedImage { + png_base64: string; + mime_type: 'image/png'; + value_label: string; + color_value_min: number; + color_value_max: number; + pixel_width_m: number; + pixel_height_m: number; + shape_h_px: number; + shape_w_px: number; +} + +async function get(path: string): Promise { + const res = await fetch(`${API}${path}`, { headers: { accept: 'application/json' } }); + if (!res.ok) { + let body = ''; + try { + body = await res.text(); + } catch { + // ignore + } + throw new Error(`GET ${path} → ${res.status} ${res.statusText}: ${body}`); + } + return (await res.json()) as T; +} + +export const api = { + listDiffraction: (limit = 200, offset = 0) => + get>(`/diffraction?limit=${limit}&offset=${offset}`), + getDiffraction: (uuid: string) => get(`/diffraction/${uuid}`), + diffractionAggregateImage: (uuid: string) => + get(`/diffraction/${uuid}/patterns/aggregate/image`), + diffractionPatternImage: (uuid: string, index: number) => + get(`/diffraction/${uuid}/patterns/${index}/image`), + diffractionFileUrl: (uuid: string) => `${API}/diffraction/${uuid}/files/diffraction`, + + listProduct: (limit = 200, offset = 0) => + get>(`/product?limit=${limit}&offset=${offset}`), + getProduct: (uuid: string) => get(`/product/${uuid}`), + productObjectImage: (uuid: string, layer: number, colorModel = 'hsv_value') => + get( + `/product/${uuid}/object/${layer}/image?color_model=${encodeURIComponent(colorModel)}` + ), + productProbeModesImage: (uuid: string, colorModel = 'hsv_value') => + get( + `/product/${uuid}/probe/modes/image?color_model=${encodeURIComponent(colorModel)}` + ), + productPositionsImage: (uuid: string, canvasPx = 512, connectPath = true) => + get( + `/product/${uuid}/positions/image?canvas_px=${canvasPx}&connect_path=${connectPath}` + ), + productFileUrl: (uuid: string) => `${API}/product/${uuid}/files/product`, + + listFluorescence: (limit = 200, offset = 0) => + get>(`/fluorescence?limit=${limit}&offset=${offset}`), + getFluorescence: (uuid: string) => get(`/fluorescence/${uuid}`), + fluorescenceElementImage: (uuid: string, name: string) => + get(`/fluorescence/${uuid}/elements/${encodeURIComponent(name)}/image`), + fluorescenceFileUrl: (uuid: string) => `${API}/fluorescence/${uuid}/files/fluorescence`, +}; diff --git a/src/ptychodus_store/ui/src/components/download_bar.ts b/src/ptychodus_store/ui/src/components/download_bar.ts new file mode 100644 index 000000000..8a71bdb79 --- /dev/null +++ b/src/ptychodus_store/ui/src/components/download_bar.ts @@ -0,0 +1,11 @@ +export function createDownloadBar(href: string, filename: string): HTMLElement { + const bar = document.createElement('div'); + bar.className = 'download-bar'; + const a = document.createElement('a'); + a.className = 'download-btn'; + a.href = href; + a.setAttribute('download', filename); + a.textContent = 'Download .h5'; + bar.appendChild(a); + return bar; +} diff --git a/src/ptychodus_store/ui/src/components/image_panel.ts b/src/ptychodus_store/ui/src/components/image_panel.ts new file mode 100644 index 000000000..854af56d1 --- /dev/null +++ b/src/ptychodus_store/ui/src/components/image_panel.ts @@ -0,0 +1,75 @@ +import type { RenderedImage } from '../api.js'; + +export interface ImagePanel { + el: HTMLElement; + setEmpty: (message?: string) => void; + setLoading: (label: string) => void; + setError: (err: Error) => void; + setImage: (img: RenderedImage, title: string) => void; +} + +export function createImagePanel(): ImagePanel { + const el = document.createElement('div'); + el.className = 'image-panel'; + + const wrap = document.createElement('div'); + wrap.className = 'image-wrap'; + const empty = document.createElement('div'); + empty.className = 'empty'; + empty.textContent = 'Select an item to preview.'; + wrap.appendChild(empty); + const caption = document.createElement('div'); + caption.className = 'caption'; + el.append(wrap, caption); + + return { + el, + setEmpty(msg = 'Select an item to preview.') { + wrap.replaceChildren(makeMessage('empty', msg)); + caption.textContent = ''; + }, + setLoading(label) { + wrap.replaceChildren(makeMessage('status', `Loading ${label}…`)); + caption.textContent = ''; + }, + setError(err) { + wrap.replaceChildren(makeMessage('error', err.message)); + caption.textContent = ''; + }, + setImage(img, title) { + const el = document.createElement('img'); + el.src = `data:${img.mime_type};base64,${img.png_base64}`; + el.alt = title; + wrap.replaceChildren(el); + const pw_um = img.pixel_width_m * 1e6; + const ph_um = img.pixel_height_m * 1e6; + caption.innerHTML = ` +
${escapeHtml(title)}
+
${escapeHtml(img.value_label)} — range [${fmt(img.color_value_min)}, ${fmt(img.color_value_max)}]
+
${img.shape_w_px} × ${img.shape_h_px} px — pixel ${pw_um.toFixed(3)} × ${ph_um.toFixed(3)} µm
+ `; + }, + }; +} + +function makeMessage(cls: string, text: string): HTMLElement { + const div = document.createElement('div'); + div.className = cls; + div.textContent = text; + return div; +} + +function fmt(x: number): string { + if (!Number.isFinite(x)) return String(x); + const abs = Math.abs(x); + if (abs !== 0 && (abs >= 1e4 || abs < 1e-2)) return x.toExponential(3); + return x.toPrecision(4); +} + +function escapeHtml(s: string): string { + return s + .replace(/&/g, '&') + .replace(//g, '>') + .replace(/"/g, '"'); +} diff --git a/src/ptychodus_store/ui/src/components/product_picker.ts b/src/ptychodus_store/ui/src/components/product_picker.ts new file mode 100644 index 000000000..a17359de2 --- /dev/null +++ b/src/ptychodus_store/ui/src/components/product_picker.ts @@ -0,0 +1,52 @@ +import { api, type ProductRead } from '../api.js'; +import { createTable } from './table.js'; + +export interface ProductPicker { + el: HTMLElement; + load: () => Promise; +} + +export function createProductPicker(onSelect: (p: ProductRead) => void): ProductPicker { + const host = document.createElement('div'); + + const load = async (): Promise => { + let items: ProductRead[] = []; + try { + const listing = await api.listProduct(); + items = listing.items; + } catch (err) { + host.replaceChildren(errorBlock(err as Error)); + return []; + } + if (items.length === 0) { + host.replaceChildren(emptyBlock('No products in store.')); + return []; + } + const table = createTable( + [{ header: 'Name', render: (p) => p.name ?? p.uuid.slice(0, 8) }], + (row) => onSelect(row) + ); + host.replaceChildren(table.el); + table.setRows(items); + return items; + }; + + return { el: host, load }; +} + +function errorBlock(err: Error): HTMLElement { + const el = document.createElement('div'); + el.style.color = '#ff8080'; + el.style.padding = '1rem'; + el.textContent = err.message; + return el; +} + +function emptyBlock(msg: string): HTMLElement { + const el = document.createElement('div'); + el.style.color = 'var(--fg-muted)'; + el.style.padding = '1rem'; + el.style.fontStyle = 'italic'; + el.textContent = msg; + return el; +} diff --git a/src/ptychodus_store/ui/src/components/table.ts b/src/ptychodus_store/ui/src/components/table.ts new file mode 100644 index 000000000..2f56d54a5 --- /dev/null +++ b/src/ptychodus_store/ui/src/components/table.ts @@ -0,0 +1,57 @@ +export interface Column { + header: string; + render: (row: T) => string; +} + +export interface DataTable { + el: HTMLElement; + setRows: (rows: T[]) => void; + setSelected: (index: number | null) => void; +} + +export function createTable(columns: Column[], onRowClick: (row: T, index: number) => void): DataTable { + const el = document.createElement('table'); + el.className = 'data-table'; + const thead = document.createElement('thead'); + const trHead = document.createElement('tr'); + for (const col of columns) { + const th = document.createElement('th'); + th.textContent = col.header; + trHead.appendChild(th); + } + thead.appendChild(trHead); + const tbody = document.createElement('tbody'); + el.append(thead, tbody); + + let selectedIndex: number | null = null; + + function setSelected(idx: number | null): void { + selectedIndex = idx; + tbody.querySelectorAll('tr.selected').forEach((n) => n.classList.remove('selected')); + if (idx !== null) { + const tr = tbody.children[idx] as HTMLElement | undefined; + tr?.classList.add('selected'); + } + } + + function setRows(rows: T[]): void { + tbody.replaceChildren(); + rows.forEach((row, i) => { + const tr = document.createElement('tr'); + for (const col of columns) { + const td = document.createElement('td'); + td.textContent = col.render(row); + tr.appendChild(td); + } + tr.addEventListener('click', () => { + setSelected(i); + onRowClick(row, i); + }); + tbody.appendChild(tr); + }); + if (selectedIndex !== null && selectedIndex < rows.length) setSelected(selectedIndex); + else selectedIndex = null; + } + + return { el, setRows, setSelected }; +} diff --git a/src/ptychodus_store/ui/src/components/tree.ts b/src/ptychodus_store/ui/src/components/tree.ts new file mode 100644 index 000000000..ef45bf6ef --- /dev/null +++ b/src/ptychodus_store/ui/src/components/tree.ts @@ -0,0 +1,106 @@ +export interface TreeNode { + id: string; + label: string; + loadChildren?: () => Promise; + children?: TreeNode[]; + onSelect?: () => void; +} + +export interface TreeRoot { + el: HTMLElement; + setSelected: (id: string | null) => void; + setNodes: (nodes: TreeNode[]) => void; +} + +export function createTree(): TreeRoot { + const el = document.createElement('ul'); + el.className = 'tree'; + let selectedId: string | null = null; + + function render(nodes: TreeNode[]): void { + el.replaceChildren(); + for (const node of nodes) el.appendChild(nodeToLi(node)); + } + + function nodeToLi(node: TreeNode): HTMLElement { + const li = document.createElement('li'); + if (node.loadChildren || (node.children && node.children.length > 0)) { + const details = document.createElement('details'); + const summary = document.createElement('summary'); + summary.textContent = node.label; + details.appendChild(summary); + const inner = document.createElement('ul'); + details.appendChild(inner); + if (node.children) { + for (const child of node.children) inner.appendChild(nodeToLi(child)); + } + let loaded = !node.loadChildren; + details.addEventListener( + 'toggle', + () => { + if (details.open && !loaded && node.loadChildren) { + loaded = true; + inner.replaceChildren(makeLoading()); + node + .loadChildren() + .then((kids) => { + inner.replaceChildren(); + for (const child of kids) inner.appendChild(nodeToLi(child)); + }) + .catch((err: Error) => { + inner.replaceChildren(makeError(err)); + }); + } + }, + { passive: true } + ); + li.appendChild(details); + } else { + const leaf = document.createElement('div'); + leaf.className = 'leaf'; + leaf.dataset.nodeId = node.id; + leaf.textContent = node.label; + if (node.id === selectedId) leaf.classList.add('selected'); + leaf.addEventListener('click', () => { + setSelected(node.id); + node.onSelect?.(); + }); + li.appendChild(leaf); + } + return li; + } + + function setSelected(id: string | null): void { + selectedId = id; + el.querySelectorAll('.leaf.selected').forEach((n) => n.classList.remove('selected')); + if (id !== null) { + const found = el.querySelector(`.leaf[data-node-id="${cssEscape(id)}"]`); + found?.classList.add('selected'); + } + } + + return { el, setSelected, setNodes: render }; +} + +function makeLoading(): HTMLElement { + const li = document.createElement('li'); + const leaf = document.createElement('div'); + leaf.className = 'leaf'; + leaf.textContent = 'Loading…'; + li.appendChild(leaf); + return li; +} + +function makeError(err: Error): HTMLElement { + const li = document.createElement('li'); + const leaf = document.createElement('div'); + leaf.className = 'leaf'; + leaf.style.color = '#ff8080'; + leaf.textContent = err.message; + li.appendChild(leaf); + return li; +} + +function cssEscape(s: string): string { + return (window as unknown as { CSS: { escape: (s: string) => string } }).CSS.escape(s); +} diff --git a/src/ptychodus_store/ui/src/layout.ts b/src/ptychodus_store/ui/src/layout.ts new file mode 100644 index 000000000..b88d9cbd3 --- /dev/null +++ b/src/ptychodus_store/ui/src/layout.ts @@ -0,0 +1,92 @@ +const NARROW_QUERY = '(max-width: 56.25rem)'; +const REM_PX = () => parseFloat(getComputedStyle(document.documentElement).fontSize) || 16; + +export interface PageLayout { + page: HTMLElement; + left: HTMLElement; + right: HTMLElement; + setActiveTab: (which: 'left' | 'right') => void; +} + +const activeTab = new Map(); + +export function buildPageLayout(pageKey: string): PageLayout { + const page = document.createElement('div'); + page.className = 'page'; + + const tabs = document.createElement('div'); + tabs.className = 'tabs'; + const tabLeft = document.createElement('button'); + tabLeft.type = 'button'; + tabLeft.textContent = 'Browse'; + const tabRight = document.createElement('button'); + tabRight.type = 'button'; + tabRight.textContent = 'View'; + tabs.append(tabLeft, tabRight); + + const left = document.createElement('div'); + left.className = 'panel-left'; + const divider = document.createElement('div'); + divider.className = 'divider'; + divider.setAttribute('role', 'separator'); + divider.setAttribute('aria-orientation', 'vertical'); + const right = document.createElement('div'); + right.className = 'panel-right'; + + page.append(tabs, left, divider, right); + + const mq = window.matchMedia(NARROW_QUERY); + const apply = () => { + const isNarrow = mq.matches; + if (isNarrow) { + const active = activeTab.get(pageKey) ?? 'left'; + left.classList.toggle('hidden', active !== 'left'); + right.classList.toggle('hidden', active !== 'right'); + tabLeft.classList.toggle('active', active === 'left'); + tabRight.classList.toggle('active', active === 'right'); + } else { + left.classList.remove('hidden'); + right.classList.remove('hidden'); + } + }; + mq.addEventListener('change', apply); + apply(); + + const setActiveTab = (which: 'left' | 'right') => { + activeTab.set(pageKey, which); + apply(); + }; + tabLeft.addEventListener('click', () => setActiveTab('left')); + tabRight.addEventListener('click', () => setActiveTab('right')); + + wireSplitterDrag(divider); + + return { page, left, right, setActiveTab }; +} + +function wireSplitterDrag(divider: HTMLElement): void { + divider.addEventListener('pointerdown', (ev) => { + ev.preventDefault(); + divider.classList.add('dragging'); + divider.setPointerCapture(ev.pointerId); + const rem = REM_PX(); + const contentRect = divider.parentElement!.getBoundingClientRect(); + const navWidth = parseFloat(getComputedStyle(document.documentElement).getPropertyValue('--nav-w')) * rem; + + const move = (e: PointerEvent) => { + const xInContent = e.clientX - contentRect.left - navWidth; + const clampedPx = Math.max(rem * 12, Math.min(rem * 60, xInContent)); + document.documentElement.style.setProperty('--left-w', `${(clampedPx / rem).toFixed(3)}rem`); + }; + const up = (e: PointerEvent) => { + divider.classList.remove('dragging'); + divider.releasePointerCapture(e.pointerId); + divider.removeEventListener('pointermove', move); + divider.removeEventListener('pointerup', up); + divider.removeEventListener('pointercancel', up); + }; + divider.addEventListener('pointermove', move); + divider.addEventListener('pointerup', up); + divider.addEventListener('pointercancel', up); + }); +} diff --git a/src/ptychodus_store/ui/src/main.ts b/src/ptychodus_store/ui/src/main.ts new file mode 100644 index 000000000..92a74766b --- /dev/null +++ b/src/ptychodus_store/ui/src/main.ts @@ -0,0 +1,56 @@ +import { NAV, renderNav } from './nav.js'; +import { mountDiffraction } from './pages/diffraction.js'; +import { mountFluorescence } from './pages/fluorescence.js'; +import { mountObject } from './pages/object.js'; +import { mountPositions } from './pages/positions.js'; +import { mountProbe } from './pages/probe.js'; +import { mountProduct } from './pages/product.js'; + +type Mount = (root: HTMLElement) => void | Promise; + +const PAGES: Record = { + diffraction: mountDiffraction, + product: mountProduct, + positions: mountPositions, + probe: mountProbe, + object: mountObject, + fluorescence: mountFluorescence, +}; + +function currentRoute(): string { + const hash = window.location.hash.replace(/^#/, ''); + if (hash && Object.hasOwn(PAGES, hash)) return hash; + return NAV[0]!.route; +} + +function navigate(route: string): void { + if (window.location.hash !== `#${route}`) { + window.location.hash = `#${route}`; + return; + } + render(route); +} + +function render(route: string): void { + const navEl = document.getElementById('nav'); + const contentEl = document.getElementById('content'); + if (!navEl || !contentEl) return; + renderNav(navEl, route, navigate); + contentEl.replaceChildren(); + const mount = PAGES[route]; + if (mount) { + const result = mount(contentEl); + if (result instanceof Promise) result.catch((err: Error) => showFatalError(contentEl, err)); + } +} + +function showFatalError(root: HTMLElement, err: Error): void { + const div = document.createElement('div'); + div.style.padding = '1.5rem'; + div.style.color = '#ff8080'; + div.textContent = err.message; + root.replaceChildren(div); +} + +window.addEventListener('hashchange', () => render(currentRoute())); +render(currentRoute()); diff --git a/src/ptychodus_store/ui/src/nav.ts b/src/ptychodus_store/ui/src/nav.ts new file mode 100644 index 000000000..1a95bc3a1 --- /dev/null +++ b/src/ptychodus_store/ui/src/nav.ts @@ -0,0 +1,106 @@ +export interface NavEntry { + route: string; + label: string; + icon: string; +} + +export const NAV: NavEntry[] = [ + { route: 'diffraction', label: 'Diffraction', icon: 'table-cells.svg' }, + { route: 'product', label: 'Products', icon: 'list.svg' }, + { route: 'positions', label: 'Positions', icon: 'route.svg' }, + { route: 'probe', label: 'Probe', icon: 'circle-radiation.svg' }, + { route: 'object', label: 'Object', icon: 'layer-group.svg' }, + { route: 'fluorescence', label: 'Fluorescence', icon: 'atom.svg' }, +]; + +interface AboutLink { + href: string; + label: string; +} + +const ABOUT_LINKS: AboutLink[] = [ + { href: 'https://github.com/AdvancedPhotonSource/ptychodus', label: 'GitHub' }, + { href: 'https://ptychodus.readthedocs.io/', label: 'Documentation' }, +]; + +export function renderNav(root: HTMLElement, current: string, onSelect: (route: string) => void): void { + root.replaceChildren(); + for (const entry of NAV) { + const btn = document.createElement('button'); + btn.type = 'button'; + btn.title = entry.label; + btn.setAttribute('aria-label', entry.label); + if (entry.route === current) btn.classList.add('active'); + const img = document.createElement('img'); + img.src = `/ui/icons/${entry.icon}`; + img.alt = ''; + btn.appendChild(img); + btn.addEventListener('click', () => onSelect(entry.route)); + root.appendChild(btn); + } + + const spacer = document.createElement('div'); + spacer.className = 'nav-spacer'; + root.appendChild(spacer); + + root.appendChild(buildAboutMenu()); +} + +function buildAboutMenu(): HTMLElement { + const wrap = document.createElement('div'); + wrap.className = 'nav-about'; + + const btn = document.createElement('button'); + btn.type = 'button'; + btn.className = 'nav-logo'; + btn.title = 'About ptychodus'; + btn.setAttribute('aria-label', 'About ptychodus'); + btn.setAttribute('aria-haspopup', 'menu'); + btn.setAttribute('aria-expanded', 'false'); + const img = document.createElement('img'); + img.src = '/ui/icons/ptychodus.svg'; + img.alt = ''; + btn.appendChild(img); + + const menu = document.createElement('div'); + menu.className = 'nav-menu'; + menu.setAttribute('role', 'menu'); + menu.hidden = true; + for (const link of ABOUT_LINKS) { + const a = document.createElement('a'); + a.href = link.href; + a.textContent = link.label; + a.target = '_blank'; + a.rel = 'noopener noreferrer'; + a.setAttribute('role', 'menuitem'); + menu.appendChild(a); + } + + const close = () => { + menu.hidden = true; + btn.setAttribute('aria-expanded', 'false'); + document.removeEventListener('pointerdown', onOutside, true); + document.removeEventListener('keydown', onKey, true); + }; + const onOutside = (e: Event) => { + if (!wrap.contains(e.target as Node)) close(); + }; + const onKey = (e: KeyboardEvent) => { + if (e.key === 'Escape') { close(); btn.focus(); } + }; + btn.addEventListener('click', () => { + const open = menu.hidden; + menu.hidden = !open; + btn.setAttribute('aria-expanded', String(open)); + if (open) { + document.addEventListener('pointerdown', onOutside, true); + document.addEventListener('keydown', onKey, true); + } else { + document.removeEventListener('pointerdown', onOutside, true); + document.removeEventListener('keydown', onKey, true); + } + }); + + wrap.append(btn, menu); + return wrap; +} diff --git a/src/ptychodus_store/ui/src/pages/diffraction.ts b/src/ptychodus_store/ui/src/pages/diffraction.ts new file mode 100644 index 000000000..3ff0fcfa3 --- /dev/null +++ b/src/ptychodus_store/ui/src/pages/diffraction.ts @@ -0,0 +1,116 @@ +import { api, type DiffractionRead } from '../api.js'; +import { createDownloadBar } from '../components/download_bar.js'; +import { createImagePanel, type ImagePanel } from '../components/image_panel.js'; +import { createTree, type TreeNode } from '../components/tree.js'; +import { buildPageLayout } from '../layout.js'; + +const PATTERN_PAGE = 200; + +export async function mountDiffraction(root: HTMLElement): Promise { + const { page, left, right, setActiveTab } = buildPageLayout('diffraction'); + root.replaceChildren(page); + + const tree = createTree(); + left.replaceChildren(tree.el); + const downloadHost = document.createElement('div'); + const image = createImagePanel(); + right.replaceChildren(downloadHost, image.el); + image.setEmpty(); + + let items: DiffractionRead[] = []; + try { + const listing = await api.listDiffraction(); + items = listing.items; + } catch (err) { + left.replaceChildren(errorBlock(err as Error)); + return; + } + + if (items.length === 0) { + left.replaceChildren(emptyBlock('No diffraction datasets in store.')); + return; + } + + const nodes: TreeNode[] = items.map((d) => ({ + id: `diff:${d.uuid}`, + label: d.label || d.uuid.slice(0, 8), + loadChildren: () => loadPatternNodes(d, image, downloadHost, setActiveTab), + })); + tree.setNodes(nodes); +} + +async function loadPatternNodes( + d: DiffractionRead, + image: ImagePanel, + downloadHost: HTMLElement, + setActiveTab: (which: 'left' | 'right') => void +): Promise { + const detail = await api.getDiffraction(d.uuid); + downloadHost.replaceChildren( + createDownloadBar(api.diffractionFileUrl(d.uuid), `${detail.label || d.uuid}.h5`) + ); + const total = detail.num_patterns_total ?? 0; + const children: TreeNode[] = [ + { + id: `diff:${d.uuid}:agg`, + label: 'Aggregate (mean)', + onSelect: () => { + setActiveTab('right'); + image.setLoading('aggregate pattern'); + api + .diffractionAggregateImage(d.uuid) + .then((img) => image.setImage(img, `${detail.label || d.uuid} — aggregate`)) + .catch((err: Error) => image.setError(err)); + }, + }, + ]; + const shown = Math.min(total, PATTERN_PAGE); + for (let i = 0; i < shown; i++) { + children.push(patternLeaf(d, detail, i, image, setActiveTab)); + } + if (total > PATTERN_PAGE) { + children.push({ + id: `diff:${d.uuid}:more`, + label: `… ${total - PATTERN_PAGE} more (open a specific index directly)`, + }); + } + return children; +} + +function patternLeaf( + d: DiffractionRead, + detail: DiffractionRead, + index: number, + image: ImagePanel, + setActiveTab: (which: 'left' | 'right') => void +): TreeNode { + return { + id: `diff:${d.uuid}:${index}`, + label: `Pattern ${index}`, + onSelect: () => { + setActiveTab('right'); + image.setLoading(`pattern ${index}`); + api + .diffractionPatternImage(d.uuid, index) + .then((img) => image.setImage(img, `${detail.label || d.uuid} — pattern ${index}`)) + .catch((err: Error) => image.setError(err)); + }, + }; +} + +function errorBlock(err: Error): HTMLElement { + const el = document.createElement('div'); + el.style.color = '#ff8080'; + el.style.padding = '1rem'; + el.textContent = err.message; + return el; +} + +function emptyBlock(msg: string): HTMLElement { + const el = document.createElement('div'); + el.style.color = 'var(--fg-muted)'; + el.style.padding = '1rem'; + el.style.fontStyle = 'italic'; + el.textContent = msg; + return el; +} diff --git a/src/ptychodus_store/ui/src/pages/fluorescence.ts b/src/ptychodus_store/ui/src/pages/fluorescence.ts new file mode 100644 index 000000000..3bed3bc3c --- /dev/null +++ b/src/ptychodus_store/ui/src/pages/fluorescence.ts @@ -0,0 +1,83 @@ +import { api, type FluorescenceRead } from '../api.js'; +import { createDownloadBar } from '../components/download_bar.js'; +import { createImagePanel, type ImagePanel } from '../components/image_panel.js'; +import { createTree, type TreeNode } from '../components/tree.js'; +import { buildPageLayout } from '../layout.js'; + +export async function mountFluorescence(root: HTMLElement): Promise { + const { page, left, right, setActiveTab } = buildPageLayout('fluorescence'); + root.replaceChildren(page); + + const tree = createTree(); + left.replaceChildren(tree.el); + const downloadHost = document.createElement('div'); + const image = createImagePanel(); + right.replaceChildren(downloadHost, image.el); + image.setEmpty(); + + let items: FluorescenceRead[] = []; + try { + const listing = await api.listFluorescence(); + items = listing.items; + } catch (err) { + left.replaceChildren(errorBlock(err as Error)); + return; + } + + if (items.length === 0) { + left.replaceChildren(emptyBlock('No fluorescence datasets in store.')); + return; + } + + const nodes: TreeNode[] = items.map((f) => ({ + id: `flu:${f.uuid}`, + label: f.label || f.uuid.slice(0, 8), + loadChildren: () => loadElementNodes(f, image, downloadHost, setActiveTab), + })); + tree.setNodes(nodes); +} + +async function loadElementNodes( + f: FluorescenceRead, + image: ImagePanel, + downloadHost: HTMLElement, + setActiveTab: (which: 'left' | 'right') => void +): Promise { + const detail = await api.getFluorescence(f.uuid); + downloadHost.replaceChildren( + createDownloadBar(api.fluorescenceFileUrl(f.uuid), `${detail.label || f.uuid}.h5`) + ); + const elements = detail.element_names ?? []; + if (elements.length === 0) { + return [{ id: `flu:${f.uuid}:none`, label: '(no elements)' }]; + } + return elements.map((name) => ({ + id: `flu:${f.uuid}:${name}`, + label: name, + onSelect: () => { + setActiveTab('right'); + image.setLoading(`element ${name}`); + api + .fluorescenceElementImage(f.uuid, name) + .then((img) => image.setImage(img, `${detail.label || f.uuid} — ${name}`)) + .catch((err: Error) => image.setError(err)); + }, + })); +} + +function errorBlock(err: Error): HTMLElement { + const el = document.createElement('div'); + el.style.color = '#ff8080'; + el.style.padding = '1rem'; + el.textContent = err.message; + return el; +} + +function emptyBlock(msg: string): HTMLElement { + const el = document.createElement('div'); + el.style.color = 'var(--fg-muted)'; + el.style.padding = '1rem'; + el.style.fontStyle = 'italic'; + el.textContent = msg; + return el; +} diff --git a/src/ptychodus_store/ui/src/pages/object.ts b/src/ptychodus_store/ui/src/pages/object.ts new file mode 100644 index 000000000..3393bc126 --- /dev/null +++ b/src/ptychodus_store/ui/src/pages/object.ts @@ -0,0 +1,82 @@ +import { api, type ProductRead } from '../api.js'; +import { createImagePanel } from '../components/image_panel.js'; +import { createProductPicker } from '../components/product_picker.js'; +import { buildPageLayout } from '../layout.js'; + +export async function mountObject(root: HTMLElement): Promise { + const { page, left, right, setActiveTab } = buildPageLayout('object'); + root.replaceChildren(page); + + const image = createImagePanel(); + const picker = document.createElement('div'); + picker.className = 'sub-picker'; + right.replaceChildren(picker, image.el); + image.setEmpty('Select a product to view its reconstructed object.'); + + let selected: ProductRead | null = null; + let layer = 0; + + const productPicker = createProductPicker((product) => selectProduct(product)); + left.replaceChildren(productPicker.el); + await productPicker.load(); + + function selectProduct(product: ProductRead): void { + selected = product; + layer = 0; + rebuildPicker(); + render(); + } + + function rebuildPicker(): void { + picker.replaceChildren(); + if (!selected) return; + const layers = selected.object_layers ?? 1; + if (layers <= 1) return; + picker.appendChild(indexInput(layer, layers - 1, (v) => { + layer = v; + render(); + })); + } + + function render(): void { + if (!selected) return; + setActiveTab('right'); + const clamped = clampInt(layer, 0, (selected.object_layers ?? 1) - 1); + image.setLoading(`object layer ${clamped}`); + api + .productObjectImage(selected.uuid, clamped) + .then((img) => image.setImage(img, `${labelFor(selected!)} — object[${clamped}]`)) + .catch((err: Error) => image.setError(err)); + } +} + +function labelFor(p: ProductRead): string { + return p.name ?? p.uuid.slice(0, 8); +} + +function clampInt(v: number, lo: number, hi: number): number { + if (hi < lo) return lo; + return Math.max(lo, Math.min(hi, v | 0)); +} + +function indexInput(value: number, max: number, onChange: (v: number) => void): HTMLElement { + const wrap = document.createElement('label'); + wrap.style.display = 'inline-flex'; + wrap.style.alignItems = 'center'; + wrap.style.gap = '0.25rem'; + const label = document.createElement('span'); + label.textContent = `layer (0–${max})`; + label.style.color = 'var(--fg-muted)'; + label.style.fontSize = '0.85em'; + const input = document.createElement('input'); + input.type = 'number'; + input.min = '0'; + input.max = String(max); + input.value = String(value); + input.addEventListener('change', () => { + const v = parseInt(input.value, 10); + if (Number.isFinite(v)) onChange(v); + }); + wrap.append(label, input); + return wrap; +} diff --git a/src/ptychodus_store/ui/src/pages/positions.ts b/src/ptychodus_store/ui/src/pages/positions.ts new file mode 100644 index 000000000..7acc66da0 --- /dev/null +++ b/src/ptychodus_store/ui/src/pages/positions.ts @@ -0,0 +1,30 @@ +import { api, type ProductRead } from '../api.js'; +import { createImagePanel } from '../components/image_panel.js'; +import { createProductPicker } from '../components/product_picker.js'; +import { buildPageLayout } from '../layout.js'; + +export async function mountPositions(root: HTMLElement): Promise { + const { page, left, right, setActiveTab } = buildPageLayout('positions'); + root.replaceChildren(page); + + const image = createImagePanel(); + right.replaceChildren(image.el); + image.setEmpty('Select a product to view its scan positions.'); + + const picker = createProductPicker((product) => selectProduct(product)); + left.replaceChildren(picker.el); + await picker.load(); + + function selectProduct(product: ProductRead): void { + setActiveTab('right'); + image.setLoading('probe positions'); + api + .productPositionsImage(product.uuid) + .then((img) => image.setImage(img, `${labelFor(product)} — positions`)) + .catch((err: Error) => image.setError(err)); + } +} + +function labelFor(p: ProductRead): string { + return p.name ?? p.uuid.slice(0, 8); +} diff --git a/src/ptychodus_store/ui/src/pages/probe.ts b/src/ptychodus_store/ui/src/pages/probe.ts new file mode 100644 index 000000000..446e6949d --- /dev/null +++ b/src/ptychodus_store/ui/src/pages/probe.ts @@ -0,0 +1,30 @@ +import { api, type ProductRead } from '../api.js'; +import { createImagePanel } from '../components/image_panel.js'; +import { createProductPicker } from '../components/product_picker.js'; +import { buildPageLayout } from '../layout.js'; + +export async function mountProbe(root: HTMLElement): Promise { + const { page, left, right, setActiveTab } = buildPageLayout('probe'); + root.replaceChildren(page); + + const image = createImagePanel(); + right.replaceChildren(image.el); + image.setEmpty('Select a product to view its probe modes.'); + + const picker = createProductPicker((product) => selectProduct(product)); + left.replaceChildren(picker.el); + await picker.load(); + + function selectProduct(product: ProductRead): void { + setActiveTab('right'); + image.setLoading('probe modes'); + api + .productProbeModesImage(product.uuid) + .then((img) => image.setImage(img, `${labelFor(product)} — probe modes`)) + .catch((err: Error) => image.setError(err)); + } +} + +function labelFor(p: ProductRead): string { + return p.name ?? p.uuid.slice(0, 8); +} diff --git a/src/ptychodus_store/ui/src/pages/product.ts b/src/ptychodus_store/ui/src/pages/product.ts new file mode 100644 index 000000000..22cde0cc6 --- /dev/null +++ b/src/ptychodus_store/ui/src/pages/product.ts @@ -0,0 +1,127 @@ +import { api, type ProductRead } from '../api.js'; +import { createDownloadBar } from '../components/download_bar.js'; +import { createTable } from '../components/table.js'; +import { buildPageLayout } from '../layout.js'; + +export async function mountProduct(root: HTMLElement): Promise { + const { page, left, right, setActiveTab } = buildPageLayout('product'); + root.replaceChildren(page); + + const detail = document.createElement('div'); + detail.className = 'detail-panel'; + right.replaceChildren(detail); + showEmpty(detail); + + let items: ProductRead[] = []; + try { + const listing = await api.listProduct(); + items = listing.items; + } catch (err) { + left.replaceChildren(errorBlock(err as Error)); + return; + } + + if (items.length === 0) { + left.replaceChildren(emptyBlock('No products in store.')); + return; + } + + const table = createTable( + [ + { header: 'Name', render: (p) => p.name ?? p.uuid.slice(0, 8) }, + { header: 'Detector-Object\nDistance [m]', render: (p) => fmt(p.detector_distance_m) }, + { header: 'Probe Energy\n[keV]', render: (p) => fmt(scale(p.probe_energy_eV, 1e-3)) }, + { header: 'Probe Photon\nCount', render: (p) => fmt(p.probe_photon_count) }, + { header: 'Pixel Width\n[nm]', render: (p) => fmt(scale(p.object_pixel_width_m, 1e9)) }, + { header: 'Pixel Height\n[nm]', render: (p) => fmt(scale(p.object_pixel_height_m, 1e9)) }, + { header: 'State', render: (p) => p.ingest_state }, + ], + (row) => { + setActiveTab('right'); + showDetail(detail, row); + } + ); + left.replaceChildren(table.el); + table.setRows(items); +} + +function showEmpty(host: HTMLElement): void { + host.replaceChildren(emptyBlock('Select a product row.')); +} + +function showDetail(host: HTMLElement, p: ProductRead): void { + host.replaceChildren(); + host.appendChild(createDownloadBar(api.productFileUrl(p.uuid), `${p.name ?? p.uuid}.h5`)); + + const title = document.createElement('h2'); + title.textContent = p.name ?? p.uuid; + host.appendChild(title); + + const dl = document.createElement('dl'); + dl.className = 'detail-list'; + const rows: [string, string][] = [ + ['UUID', p.uuid], + ['State', p.ingest_state], + ['Detector-Object Distance [m]', fmt(p.detector_distance_m)], + ['Probe Energy [keV]', fmt(scale(p.probe_energy_eV, 1e-3))], + ['Probe Photon Count', fmt(p.probe_photon_count)], + ['Probe Modes', fmt(p.probe_modes)], + ['Probe Shape [px]', shape(p.probe_height_px, p.probe_width_px)], + ['Object Layers', fmt(p.object_layers)], + ['Object Shape [px]', shape(p.object_height_px, p.object_width_px)], + ['Object Pixel Width [nm]', fmt(scale(p.object_pixel_width_m, 1e9))], + ['Object Pixel Height [nm]', fmt(scale(p.object_pixel_height_m, 1e9))], + ['Scan Points', fmt(p.num_scan_points)], + ['Tomography Angle [deg]', fmt(p.tomography_angle_deg)], + ['Tilt Angle [deg]', fmt(p.tilt_angle_deg)], + ['Polarization', p.polarization ?? '—'], + ]; + for (const [k, v] of rows) { + const dt = document.createElement('dt'); + dt.textContent = k; + const dd = document.createElement('dd'); + dd.textContent = v; + dl.append(dt, dd); + } + host.appendChild(dl); + + if (p.comments) { + const h3 = document.createElement('h3'); + h3.textContent = 'Comments'; + const pre = document.createElement('pre'); + pre.className = 'detail-comments'; + pre.textContent = p.comments; + host.append(h3, pre); + } +} + +function scale(x: number | null, factor: number): number | null { + return x === null ? null : x * factor; +} + +function fmt(x: number | null): string { + if (x === null || !Number.isFinite(x)) return '—'; + return Number(x).toPrecision(4); +} + +function shape(h: number | null, w: number | null): string { + if (h === null || w === null) return '—'; + return `${h} × ${w}`; +} + +function errorBlock(err: Error): HTMLElement { + const el = document.createElement('div'); + el.style.color = '#ff8080'; + el.style.padding = '1rem'; + el.textContent = err.message; + return el; +} + +function emptyBlock(msg: string): HTMLElement { + const el = document.createElement('div'); + el.style.color = 'var(--fg-muted)'; + el.style.padding = '1rem'; + el.style.fontStyle = 'italic'; + el.textContent = msg; + return el; +} diff --git a/src/ptychodus_store/ui/styles.css b/src/ptychodus_store/ui/styles.css new file mode 100644 index 000000000..ca8d3c403 --- /dev/null +++ b/src/ptychodus_store/ui/styles.css @@ -0,0 +1,314 @@ +:root { + --bg: #1e1f22; + --bg-panel: #252629; + --bg-panel-2: #2b2c31; + --fg: #e6e6e6; + --fg-muted: #9ea0a6; + --accent: #4ea1ff; + --border: #3a3b40; + --nav-w: 3.5rem; + --left-w: 20rem; + --left-min: 14rem; + --divider: 4px; + font-family: system-ui, -apple-system, "Segoe UI", sans-serif; + font-size: 14px; + color-scheme: dark; +} + +* { box-sizing: border-box; } + +html, body { + margin: 0; + height: 100%; + background: var(--bg); + color: var(--fg); +} + +body { + display: grid; + grid-template-columns: var(--nav-w) 1fr; + grid-template-rows: 100vh; +} + +/* -- nav rail -- */ +.nav { + background: var(--bg-panel); + border-right: 1px solid var(--border); + display: flex; + flex-direction: column; + align-items: center; + padding: 0.5rem 0; + gap: 0.25rem; +} + +.nav button { + width: 2.5rem; + height: 2.5rem; + padding: 0.375rem; + border: none; + border-left: 3px solid transparent; + background: transparent; + color: var(--fg-muted); + cursor: pointer; + display: flex; + align-items: center; + justify-content: center; + border-radius: 0 4px 4px 0; +} +.nav button:hover { background: var(--bg-panel-2); color: var(--fg); } +.nav button.active { + border-left-color: var(--accent); + background: var(--bg-panel-2); + color: var(--fg); +} +.nav button img { + width: 1.25rem; + height: 1.25rem; + filter: invert(100%) brightness(85%); + pointer-events: none; +} +.nav button.active img { filter: invert(100%) brightness(100%); } + +.nav-spacer { flex: 1; } + +.nav-about { position: relative; } +.nav-about .nav-logo { + width: 2.5rem; + height: 2.5rem; + padding: 0.25rem; + border: none; + border-left: 3px solid transparent; + background: transparent; + cursor: pointer; + display: flex; + align-items: center; + justify-content: center; + border-radius: 0 4px 4px 0; +} +.nav-about .nav-logo:hover { background: var(--bg-panel-2); } +.nav-about .nav-logo img { + width: 1.75rem; + height: 1.75rem; + filter: invert(100%) brightness(90%); + pointer-events: none; +} +.nav-about .nav-menu { + position: absolute; + left: calc(100% + 0.375rem); + bottom: 0; + min-width: 10rem; + background: var(--bg-panel-2); + border: 1px solid var(--border); + border-radius: 4px; + padding: 0.25rem; + display: flex; + flex-direction: column; + gap: 0.125rem; + box-shadow: 0 4px 12px rgba(0, 0, 0, 0.4); + z-index: 10; +} +.nav-about .nav-menu[hidden] { display: none; } +.nav-about .nav-menu a { + padding: 0.375rem 0.625rem; + color: var(--fg); + text-decoration: none; + border-radius: 3px; + font-size: 0.875em; +} +.nav-about .nav-menu a:hover { background: var(--bg-panel); color: var(--accent); } + +/* -- content area -- */ +.content { + display: grid; + grid-template-columns: minmax(var(--left-min), var(--left-w)) var(--divider) 1fr; + overflow: hidden; +} + +.panel-left, .panel-right { + overflow: auto; + background: var(--bg); + padding: 0.75rem; +} +.panel-left { border-right: 1px solid var(--border); } + +.divider { + cursor: col-resize; + background: var(--border); + user-select: none; +} +.divider:hover, .divider.dragging { background: var(--accent); } + +/* -- tabs (only visible on narrow) -- */ +.tabs { display: none; } + +/* -- tree / list / table -- */ +.tree { list-style: none; margin: 0; padding: 0; } +.tree ul { list-style: none; margin: 0; padding-left: 1rem; } +.tree details { padding: 0.125rem 0; } +.tree summary { cursor: pointer; padding: 0.125rem 0.25rem; border-radius: 3px; } +.tree summary:hover { background: var(--bg-panel-2); } +.tree .leaf { + padding: 0.125rem 0.25rem 0.125rem 1.125rem; + cursor: pointer; + border-radius: 3px; + font-size: 0.875em; +} +.tree .leaf:hover { background: var(--bg-panel-2); } +.tree .leaf.selected { background: var(--accent); color: #0a0a0a; } +.tree .load-more { + padding: 0.25rem 0.5rem; + color: var(--fg-muted); + cursor: pointer; + font-size: 0.85em; +} + +.data-table { width: 100%; border-collapse: collapse; font-size: 0.875em; } +.data-table th, .data-table td { + padding: 0.375rem 0.5rem; + text-align: left; + border-bottom: 1px solid var(--border); +} +.data-table th { color: var(--fg-muted); font-weight: 600; white-space: pre-line; vertical-align: bottom; } +.data-table tbody tr { cursor: pointer; } +.data-table tbody tr:hover { background: var(--bg-panel-2); } +.data-table tbody tr.selected { background: var(--accent); color: #0a0a0a; } +.data-table tbody tr.selected td { color: #0a0a0a; } + +.sub-picker { display: flex; gap: 0.25rem; margin-bottom: 0.75rem; flex-wrap: wrap; } +.sub-picker button { + padding: 0.25rem 0.625rem; + border: 1px solid var(--border); + background: var(--bg-panel); + color: var(--fg); + border-radius: 3px; + cursor: pointer; + font-size: 0.85em; +} +.sub-picker button:hover { border-color: var(--accent); } +.sub-picker button.active { background: var(--accent); color: #0a0a0a; border-color: var(--accent); } +.sub-picker input { + padding: 0.25rem 0.375rem; + border: 1px solid var(--border); + background: var(--bg-panel); + color: var(--fg); + border-radius: 3px; + width: 5rem; + font-size: 0.85em; +} + +/* -- image panel -- */ +.image-panel { display: flex; flex-direction: column; height: 100%; } +.image-panel .image-wrap { + flex: 1; + display: flex; + align-items: center; + justify-content: center; + background: #000; + border: 1px solid var(--border); + overflow: hidden; + min-height: 0; +} +.image-panel img { + max-width: 100%; + max-height: 100%; + image-rendering: pixelated; + object-fit: contain; +} +.image-panel .caption { + padding-top: 0.5rem; + color: var(--fg-muted); + font-size: 0.8em; + line-height: 1.5; +} +.image-panel .caption strong { color: var(--fg); font-weight: 600; } +.image-panel .empty { + color: var(--fg-muted); + font-style: italic; + padding: 2rem; + text-align: center; +} +.image-panel .status { padding: 1rem; color: var(--fg-muted); font-size: 0.85em; } +.image-panel .error { color: #ff8080; padding: 1rem; font-size: 0.85em; } + +/* -- responsive collapse: single column with tab strip -- */ +@media (max-width: 56.25rem) { + .content { + grid-template-columns: 1fr; + grid-template-rows: auto 1fr; + } + .divider { display: none; } + .panel-left, .panel-right { border-right: none; } + .panel-left.hidden, .panel-right.hidden { display: none; } + .tabs { + display: flex; + gap: 0; + background: var(--bg-panel); + border-bottom: 1px solid var(--border); + } + .tabs button { + flex: 1; + padding: 0.5rem; + border: none; + border-bottom: 2px solid transparent; + background: transparent; + color: var(--fg-muted); + cursor: pointer; + font-size: 0.9em; + } + .tabs button.active { + color: var(--fg); + border-bottom-color: var(--accent); + } +} + +.page { display: contents; } + +/* -- download bar (right-panel header) -- */ +.download-bar { margin-bottom: 0.5rem; } +.download-btn { + display: inline-block; + padding: 0.25rem 0.625rem; + border: 1px solid var(--border); + background: var(--bg-panel); + color: var(--fg); + border-radius: 3px; + cursor: pointer; + font-size: 0.85em; + text-decoration: none; +} +.download-btn:hover { border-color: var(--accent); color: var(--accent); } + +/* -- detail panel (product metadata view) -- */ +.detail-panel h2 { + margin: 0.25rem 0 0.75rem; + font-size: 1.1em; + font-weight: 600; + color: var(--fg); + word-break: break-all; +} +.detail-panel h3 { + margin: 1rem 0 0.375rem; + font-size: 0.95em; + font-weight: 600; + color: var(--fg-muted); +} +.detail-list { + margin: 0; + display: grid; + grid-template-columns: max-content 1fr; + column-gap: 0.75rem; + row-gap: 0.25rem; + font-size: 0.875em; +} +.detail-list dt { color: var(--fg-muted); } +.detail-list dd { margin: 0; color: var(--fg); font-variant-numeric: tabular-nums; } +.detail-comments { + margin: 0; + padding: 0.5rem; + background: var(--bg-panel); + border: 1px solid var(--border); + border-radius: 3px; + font-size: 0.875em; + white-space: pre-wrap; + font-family: inherit; +} diff --git a/src/ptychodus_store/ui/tsconfig.json b/src/ptychodus_store/ui/tsconfig.json new file mode 100644 index 000000000..aab9f0aee --- /dev/null +++ b/src/ptychodus_store/ui/tsconfig.json @@ -0,0 +1,21 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ES2022", + "moduleResolution": "bundler", + "strict": true, + "noUncheckedIndexedAccess": true, + "noImplicitOverride": true, + "noFallthroughCasesInSwitch": true, + "forceConsistentCasingInFileNames": true, + "esModuleInterop": true, + "isolatedModules": true, + "lib": ["ES2022", "DOM", "DOM.Iterable"], + "outDir": "./dist", + "rootDir": "./src", + "sourceMap": true, + "declaration": false, + "removeComments": false + }, + "include": ["src/**/*.ts"] +} diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..47ff74e22 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +"""Skip optional-extra test packages when their dependencies are absent.""" + +from importlib.util import find_spec + +# tests/ptychodus_store/ needs the "store" extra plus pytest-asyncio. Its own conftest +# imports sqlalchemy eagerly, so the directory has to be dropped before pytest descends +# into it -- pytest.importorskip in a conftest is reported as an error, not a skip. +_STORE_TEST_DEPS = ( + 'aiosqlite', + 'fastapi', + 'fastmcp', + 'pydantic_settings', + 'pytest_asyncio', + 'sqlalchemy', +) + +collect_ignore = [] + +if any(find_spec(name) is None for name in _STORE_TEST_DEPS): + collect_ignore.append('ptychodus_store') + +if find_spec('PyQt5') is None: + collect_ignore.append('view') diff --git a/tests/ptychodus_store/conftest.py b/tests/ptychodus_store/conftest.py new file mode 100644 index 000000000..9d3999b29 --- /dev/null +++ b/tests/ptychodus_store/conftest.py @@ -0,0 +1,274 @@ +"""Shared fixtures for ptychodus_store tests.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator, Callable +from datetime import datetime, timezone +from pathlib import Path +from uuid import UUID, uuid4 + +import h5py +import numpy as np +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient +from sqlalchemy.ext.asyncio import AsyncSession + +from ptychodus_store.config import Settings +from ptychodus_store.db.session import SessionProvider, create_engine, create_schema +from ptychodus_store.storage.layout import StoreLayout + + +@pytest.fixture +def tmp_storage_root(tmp_path: Path) -> Path: + root = tmp_path / 'store' + layout = StoreLayout(root) + layout.ensure_kind_dirs() + return root + + +@pytest.fixture +def layout(tmp_storage_root: Path) -> StoreLayout: + return StoreLayout(tmp_storage_root) + + +@pytest_asyncio.fixture +async def db_engine() -> AsyncIterator: + engine = create_engine('sqlite+aiosqlite:///:memory:') + await create_schema(engine) + try: + yield engine + finally: + await engine.dispose() + + +@pytest_asyncio.fixture +async def session_provider(db_engine) -> AsyncIterator[SessionProvider]: # type: ignore[no-untyped-def] + yield SessionProvider(db_engine) + + +@pytest_asyncio.fixture +async def db_session(session_provider: SessionProvider) -> AsyncIterator[AsyncSession]: + async with session_provider.session_factory() as session: + yield session + + +# -------------------------------------------------------------------------------------- +# Seed helpers — write a manifest + minimal HDF5 file +# -------------------------------------------------------------------------------------- + + +def _write_manifest(folder: Path, manifest: dict) -> None: + (folder / 'manifest.json').write_text(json.dumps(manifest)) + + +def _write_diffraction_h5(path: Path, *, num_patterns: int = 4, h: int = 8, w: int = 12) -> None: + with h5py.File(path, 'w') as f: + ds = f.create_dataset('patterns', data=np.zeros((num_patterns, h, w), dtype=np.uint16)) + ds.attrs['detector_pixel_width_m'] = 55e-6 + ds.attrs['detector_pixel_height_m'] = 55e-6 + f.create_dataset('indexes', data=np.arange(num_patterns)) + f.create_dataset('bad_pixels', data=np.zeros((h, w), dtype=bool)) + + +def _write_product_h5(path: Path) -> None: + with h5py.File(path, 'w') as f: + f.attrs['name'] = 'test-product' + f.attrs['comments'] = '' + f.attrs['detector_object_distance_m'] = 1.5 + f.attrs['probe_energy_eV'] = 9000.0 + f.attrs['probe_photon_count'] = 1_000_000 + f.attrs['exposure_time_s'] = 0.1 + f.attrs['mass_attenuation_m2_kg'] = 0.0 + f.attrs['tomography_angle_deg'] = 0.0 + obj = f.create_dataset('object', data=np.zeros((1, 16, 16), dtype=np.complex64)) + obj.attrs['pixel_width_m'] = 1e-9 + obj.attrs['pixel_height_m'] = 1e-9 + obj.attrs['center_x_m'] = 0.0 + obj.attrs['center_y_m'] = 0.0 + probe = f.create_dataset('probe', data=np.zeros((1, 8, 8), dtype=np.complex64)) + probe.attrs['pixel_width_m'] = 1e-9 + probe.attrs['pixel_height_m'] = 1e-9 + f.create_dataset('probe_position_indexes', data=np.arange(4)) + f.create_dataset('probe_position_x_m', data=np.zeros(4)) + f.create_dataset('probe_position_y_m', data=np.zeros(4)) + f.create_dataset('loss_epochs', data=np.array([1, 2, 3])) + f.create_dataset('loss_values', data=np.array([0.3, 0.2, 0.1])) + + +def _write_fluorescence_h5(path: Path, elements: list[str], h: int = 6, w: int = 10) -> None: + """Write the XRF-Maps v10 NNLS layout expected by :func:`load_fluorescence_data`.""" + with h5py.File(path, 'w') as f: + group = f.require_group('/MAPS/XRF_Analyzed/NNLS') + group.create_dataset( + 'Counts_Per_Sec', data=np.zeros((len(elements), h, w), dtype=np.float32) + ) + group.create_dataset('Channel_Names', data=np.array(elements, dtype='S16')) + + +@pytest.fixture +def seed_campaign(tmp_storage_root: Path) -> Callable[..., UUID]: + def _seed( + uuid: UUID | None = None, + *, + label: str = 'campaign-a', + sample_name: str = 'sample-x', + tags: list[str] | None = None, + ) -> UUID: + uuid = uuid or uuid4() + folder = tmp_storage_root / 'campaign' / str(uuid) + folder.mkdir(parents=True) + _write_manifest( + folder, + { + 'schema_version': 1, + 'kind': 'campaign', + 'uuid': str(uuid), + 'created_at': datetime.now(timezone.utc).isoformat(), + 'label': label, + 'sample_name': sample_name, + 'tags': tags or [], + }, + ) + return uuid + + return _seed + + +@pytest.fixture +def seed_diffraction(tmp_storage_root: Path) -> Callable[..., UUID]: + def _seed( + uuid: UUID | None = None, + *, + campaign_uuid: UUID | None = None, + derived_from: list[dict] | None = None, + probe_energy_eV: float | None = 8000.0, # noqa: N803 + write_h5: bool = True, + ) -> UUID: + uuid = uuid or uuid4() + folder = tmp_storage_root / 'diffraction' / str(uuid) + folder.mkdir(parents=True) + if write_h5: + _write_diffraction_h5(folder / 'diffraction.h5') + manifest = { + 'schema_version': 1, + 'kind': 'diffraction', + 'uuid': str(uuid), + 'created_at': datetime.now(timezone.utc).isoformat(), + 'probe_energy_eV': probe_energy_eV, + } + if campaign_uuid is not None: + manifest['campaign_uuid'] = str(campaign_uuid) + if derived_from is not None: + manifest['derived_from'] = derived_from + _write_manifest(folder, manifest) + return uuid + + return _seed + + +@pytest.fixture +def seed_product(tmp_storage_root: Path) -> Callable[..., UUID]: + def _seed( + uuid: UUID | None = None, + *, + derived_from: list[dict] | None = None, + write_h5: bool = True, + ) -> UUID: + uuid = uuid or uuid4() + folder = tmp_storage_root / 'product' / str(uuid) + folder.mkdir(parents=True) + if write_h5: + _write_product_h5(folder / 'product.h5') + manifest = { + 'schema_version': 1, + 'kind': 'product', + 'uuid': str(uuid), + 'created_at': datetime.now(timezone.utc).isoformat(), + } + if derived_from is not None: + manifest['derived_from'] = derived_from + _write_manifest(folder, manifest) + return uuid + + return _seed + + +@pytest.fixture +def seed_fluorescence(tmp_storage_root: Path) -> Callable[..., UUID]: + def _seed( + uuid: UUID | None = None, + *, + derived_from: list[dict] | None = None, + elements: list[str] | None = None, + ) -> UUID: + uuid = uuid or uuid4() + folder = tmp_storage_root / 'fluorescence' / str(uuid) + folder.mkdir(parents=True) + _write_fluorescence_h5(folder / 'fluorescence.h5', elements or ['Fe', 'Cu']) + manifest = { + 'schema_version': 1, + 'kind': 'fluorescence', + 'uuid': str(uuid), + 'created_at': datetime.now(timezone.utc).isoformat(), + } + if derived_from is not None: + manifest['derived_from'] = derived_from + _write_manifest(folder, manifest) + return uuid + + return _seed + + +# -------------------------------------------------------------------------------------- +# FastAPI client fixture +# -------------------------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def app_client( + tmp_storage_root: Path, + db_engine, # type: ignore[no-untyped-def] +) -> AsyncIterator[AsyncClient]: + from fastapi import FastAPI + + from ptychodus_store.db.session import SessionProvider + from ptychodus_store.mcp_server import bind_session_provider + from ptychodus_store.routers import ( + admin, + campaign, + diffraction, + fluorescence, + health, + lineage, + visualization, + ) + from ptychodus_store.routers import product as product_router + + layout = StoreLayout(tmp_storage_root) + provider = SessionProvider(db_engine) + bind_session_provider(provider) + settings = Settings( # type: ignore[call-arg] + storage_root=tmp_storage_root, + database_url='sqlite+aiosqlite:///:memory:', + auto_reconcile_on_startup=False, + ) + + app = FastAPI() + app.state.settings = settings + app.state.layout = layout + app.state.session_provider = provider + api_prefix = '/api/v1' + app.include_router(health.router, prefix=api_prefix) + app.include_router(campaign.router, prefix=api_prefix) + app.include_router(diffraction.router, prefix=api_prefix) + app.include_router(product_router.router, prefix=api_prefix) + app.include_router(fluorescence.router, prefix=api_prefix) + app.include_router(lineage.router, prefix=api_prefix) + app.include_router(admin.router, prefix=api_prefix) + app.include_router(visualization.router, prefix=api_prefix) + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url='http://testserver') as client: + yield client diff --git a/tests/ptychodus_store/test_h5_introspect.py b/tests/ptychodus_store/test_h5_introspect.py new file mode 100644 index 000000000..0dcdb003d --- /dev/null +++ b/tests/ptychodus_store/test_h5_introspect.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from pathlib import Path + +import h5py +import numpy as np +import pytest + +from ptychodus_store.storage.h5_introspect import ( + IntrospectionError, + introspect_diffraction, + introspect_fluorescence, + introspect_product, +) + + +def test_introspect_diffraction(tmp_path: Path) -> None: + path = tmp_path / 'diffraction.h5' + with h5py.File(path, 'w') as f: + ds = f.create_dataset('patterns', data=np.zeros((7, 64, 128), dtype=np.uint16)) + ds.attrs['detector_pixel_width_m'] = 1.5e-5 + ds.attrs['detector_pixel_height_m'] = 1.5e-5 + f.create_dataset('indexes', data=np.arange(7)) + f.create_dataset('bad_pixels', data=np.zeros((64, 128), dtype=bool)) + + result = introspect_diffraction(path) + assert result['num_patterns_total'] == 7 + assert result['pattern_shape'] == (64, 128) + assert result['pattern_dtype'] == 'uint16' + assert result['detector_pixel_width_m'] == pytest.approx(1.5e-5) + + +def test_introspect_diffraction_missing_dataset(tmp_path: Path) -> None: + path = tmp_path / 'bad.h5' + with h5py.File(path, 'w') as f: + f.create_dataset('something_else', data=np.zeros(1)) + with pytest.raises(IntrospectionError): + introspect_diffraction(path) + + +def test_introspect_product(tmp_path: Path) -> None: + path = tmp_path / 'product.h5' + with h5py.File(path, 'w') as f: + f.attrs['name'] = 'p1' + f.attrs['comments'] = 'hello' + f.attrs['detector_object_distance_m'] = 2.0 + f.attrs['probe_energy_eV'] = 8500.0 + f.attrs['probe_photon_count'] = 12345 + f.attrs['exposure_time_s'] = 0.05 + f.attrs['mass_attenuation_m2_kg'] = 0.0 + f.attrs['tomography_angle_deg'] = 30.0 + f.attrs['tilt_angle_deg'] = 12.5 + f.attrs['polarization'] = 'left_circular' + obj = f.create_dataset('object', data=np.zeros((2, 32, 48), dtype=np.complex64)) + obj.attrs['pixel_width_m'] = 1e-9 + obj.attrs['pixel_height_m'] = 1e-9 + probe = f.create_dataset('probe', data=np.zeros((3, 16, 16), dtype=np.complex64)) + probe.attrs['pixel_width_m'] = 1e-9 + probe.attrs['pixel_height_m'] = 1e-9 + f.create_dataset('probe_position_indexes', data=np.arange(11)) + f.create_dataset('loss_epochs', data=np.array([1, 2])) + + result = introspect_product(path) + assert result['name'] == 'p1' + assert result['comments'] == 'hello' + assert result['probe_energy_eV'] == 8500.0 + assert result['object_shape'] == (2, 32, 48) + assert result['probe_shape'] == (3, 16, 16) + assert result['num_scan_points'] == 11 + assert result['num_loss_epochs'] == 2 + assert result['tomography_angle_deg'] == 30.0 + assert result['tilt_angle_deg'] == 12.5 + assert result['polarization'] == 'left_circular' + + +def test_introspect_fluorescence(tmp_path: Path) -> None: + path = tmp_path / 'fluorescence.h5' + with h5py.File(path, 'w') as f: + group = f.require_group('/MAPS/XRF_Analyzed/NNLS') + group.create_dataset('Counts_Per_Sec', data=np.zeros((2, 6, 9), dtype=np.float32)) + group.create_dataset('Channel_Names', data=np.array(['Fe', 'Cu'], dtype='S8')) + + result = introspect_fluorescence(path) + assert result['element_names'] == ['Fe', 'Cu'] + assert result['map_shape'] == (6, 9) diff --git a/tests/ptychodus_store/test_layout.py b/tests/ptychodus_store/test_layout.py new file mode 100644 index 000000000..f36f889ea --- /dev/null +++ b/tests/ptychodus_store/test_layout.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path +from uuid import uuid4 + +import pytest + +from ptychodus_store.storage.layout import LayoutError, StoreLayout + + +def test_resource_folder_resolution(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path) + uuid = uuid4() + folder = layout.resource_folder('product', uuid) + assert folder == tmp_path.resolve() / 'product' / str(uuid) + + +def test_manifest_path_resolution(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path) + uuid = uuid4() + m = layout.manifest_path('fluorescence', uuid) + assert m.name == 'manifest.json' + assert m.parent == tmp_path.resolve() / 'fluorescence' / str(uuid) + + +def test_parse_manifest_path_valid(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path) + layout.ensure_kind_dirs() + uuid = uuid4() + folder = tmp_path / 'campaign' / str(uuid) + folder.mkdir() + m = folder / 'manifest.json' + m.touch() + loc = layout.parse_manifest_path(m) + assert loc.kind == 'campaign' + assert loc.uuid == uuid + assert loc.folder == folder.resolve() + + +def test_parse_manifest_rejects_outside_root(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path / 'inside') + elsewhere = tmp_path / 'outside' / 'campaign' / str(uuid4()) / 'manifest.json' + elsewhere.parent.mkdir(parents=True) + elsewhere.touch() + with pytest.raises(LayoutError): + layout.parse_manifest_path(elsewhere) + + +def test_parse_manifest_rejects_non_uuid_folder(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path) + layout.ensure_kind_dirs() + bad = tmp_path / 'product' / 'not-a-uuid' / 'manifest.json' + bad.parent.mkdir() + bad.touch() + with pytest.raises(LayoutError): + layout.parse_manifest_path(bad) + + +def test_parse_manifest_rejects_unknown_kind(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path) + layout.ensure_kind_dirs() + bad = tmp_path / 'mystery' / str(uuid4()) / 'manifest.json' + bad.parent.mkdir(parents=True) + bad.touch() + with pytest.raises(LayoutError): + layout.parse_manifest_path(bad) + + +def test_iter_manifest_paths_orders_campaigns_first(tmp_path: Path) -> None: + layout = StoreLayout(tmp_path) + layout.ensure_kind_dirs() + + for kind in ('product', 'diffraction', 'campaign', 'fluorescence'): + folder = tmp_path / kind / str(uuid4()) + folder.mkdir() + (folder / 'manifest.json').touch() + + paths = layout.iter_manifest_paths() + kinds = [p.parent.parent.name for p in paths] + # Campaign first; then diffraction, product, fluorescence in any order — the + # spec orders them as campaign, diffraction, product, fluorescence. + assert kinds[0] == 'campaign' + assert kinds == ['campaign', 'diffraction', 'product', 'fluorescence'] diff --git a/tests/ptychodus_store/test_manifest_schema.py b/tests/ptychodus_store/test_manifest_schema.py new file mode 100644 index 000000000..165696bc8 --- /dev/null +++ b/tests/ptychodus_store/test_manifest_schema.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from uuid import uuid4 + +import pytest +from pydantic import ValidationError + +from ptychodus_store.storage.manifest import ( + HDF5_OWNED_KEYS, + CampaignManifest, + DerivedFromRef, + DiffractionManifest, + FluorescenceManifest, + ManifestLoadError, + ProductManifest, + ResourceKind, + load_manifest, +) + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _write(path: Path, data: dict) -> None: + path.write_text(json.dumps(data)) + + +def test_diffraction_manifest_default_files() -> None: + m = DiffractionManifest(uuid=uuid4(), created_at=datetime.now(timezone.utc)) + assert m.files == {'diffraction': 'diffraction.h5'} + assert m.derived_from == [] + assert m.campaign_uuid is None + + +def test_product_manifest_minimal_fields() -> None: + m = ProductManifest(uuid=uuid4(), created_at=datetime.now(timezone.utc)) + assert m.files == {'product': 'product.h5'} + assert m.derived_from == [] + + +def test_fluorescence_manifest_defaults() -> None: + m = FluorescenceManifest(uuid=uuid4(), created_at=datetime.now(timezone.utc)) + assert m.files == {'fluorescence': 'fluorescence.h5'} + + +def test_derived_from_ref_rejects_campaign() -> None: + with pytest.raises(ValidationError): + DerivedFromRef(kind='campaign', uuid=uuid4()) # type: ignore[arg-type] + + +def test_manifest_extra_rejects_hdf5_owned_keys() -> None: + # detector_pixel_width_m is HDF5-owned for diffraction; supplying it via extra must raise + with pytest.raises(ValidationError): + DiffractionManifest( + uuid=uuid4(), + created_at=datetime.now(timezone.utc), + extra={'detector_pixel_width_m': 1.0}, + ) + + +def test_hdf5_owned_disjoint_from_model_fields() -> None: + """Sanity: HDF5-owned keys must NOT also appear as manifest fields.""" + for kind, model_cls in ( + (ResourceKind.CAMPAIGN, CampaignManifest), + (ResourceKind.DIFFRACTION, DiffractionManifest), + (ResourceKind.PRODUCT, ProductManifest), + (ResourceKind.FLUORESCENCE, FluorescenceManifest), + ): + manifest_fields = set(model_cls.model_fields.keys()) + clash = HDF5_OWNED_KEYS[kind] & manifest_fields + assert not clash, f'{kind}: manifest fields overlap HDF5-owned keys: {clash}' + + +def test_load_manifest_rejects_uuid_mismatch(tmp_path: Path) -> None: + folder_uuid = uuid4() + other_uuid = uuid4() + path = tmp_path / 'manifest.json' + _write( + path, + { + 'schema_version': 1, + 'kind': 'diffraction', + 'uuid': str(other_uuid), + 'created_at': _now(), + }, + ) + with pytest.raises(ManifestLoadError, match='does not match folder name'): + load_manifest(path, expected_kind='diffraction', expected_uuid=folder_uuid) + + +def test_load_manifest_rejects_kind_mismatch(tmp_path: Path) -> None: + uuid = uuid4() + path = tmp_path / 'manifest.json' + _write( + path, + { + 'schema_version': 1, + 'kind': 'product', + 'uuid': str(uuid), + 'created_at': _now(), + }, + ) + with pytest.raises(ManifestLoadError, match='does not match folder kind'): + load_manifest(path, expected_kind='diffraction', expected_uuid=uuid) + + +def test_load_manifest_rejects_self_reference(tmp_path: Path) -> None: + uuid = uuid4() + path = tmp_path / 'manifest.json' + _write( + path, + { + 'schema_version': 1, + 'kind': 'product', + 'uuid': str(uuid), + 'created_at': _now(), + 'derived_from': [{'kind': 'product', 'uuid': str(uuid)}], + }, + ) + with pytest.raises(ManifestLoadError, match='self-references'): + load_manifest(path, expected_kind='product', expected_uuid=uuid) + + +def test_load_manifest_happy_path(tmp_path: Path) -> None: + uuid = uuid4() + path = tmp_path / 'manifest.json' + _write( + path, + { + 'schema_version': 1, + 'kind': 'diffraction', + 'uuid': str(uuid), + 'created_at': _now(), + 'label': 'foo', + 'probe_energy_eV': 8000.0, + }, + ) + m = load_manifest(path, expected_kind='diffraction', expected_uuid=uuid) + assert isinstance(m, DiffractionManifest) + assert m.label == 'foo' + assert m.probe_energy_eV == 8000.0 diff --git a/tests/ptychodus_store/test_mcp_tools.py b/tests/ptychodus_store/test_mcp_tools.py new file mode 100644 index 000000000..9937d17e5 --- /dev/null +++ b/tests/ptychodus_store/test_mcp_tools.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import base64 +from io import BytesIO + +import pytest +from fastmcp.exceptions import ToolError +from PIL import Image + +from ptychodus_store.ingest.pipeline import ingest_manifest +from ptychodus_store.mcp_server import bind_layout, bind_session_provider, create_mcp_server + +pytestmark = pytest.mark.asyncio + + +def _decode_image_content(result, expected_h: int, expected_w: int) -> Image.Image: # type: ignore[no-untyped-def] + assert len(result.content) == 1 + item = result.content[0] + assert item.type == 'image' + assert item.mimeType == 'image/png' + png_bytes = base64.b64decode(item.data) + image = Image.open(BytesIO(png_bytes)) + image.load() + assert image.size == (expected_w, expected_h) + return image + + +async def test_mcp_tools_registered_and_callable( # type: ignore[no-untyped-def] + session_provider, + layout, + seed_campaign, + seed_diffraction, + seed_product, +): + # Seed a tiny graph + c = seed_campaign(sample_name='alpha') + d = seed_diffraction(campaign_uuid=c) + p = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d)}]) + async with session_provider.session_factory() as session: + await ingest_manifest(session, layout, layout.manifest_path('campaign', c)) + await ingest_manifest(session, layout, layout.manifest_path('diffraction', d)) + await ingest_manifest(session, layout, layout.manifest_path('product', p)) + await session.commit() + + bind_session_provider(session_provider) + bind_layout(layout) + mcp = create_mcp_server() + tools_list = await mcp.list_tools() + names = {t.name for t in tools_list} + expected = { + 'list_campaign', + 'get_campaign', + 'list_diffraction', + 'get_diffraction', + 'list_product', + 'get_product', + 'list_fluorescence', + 'get_fluorescence', + 'get_lineage', + 'get_store_stats', + 'get_visualization_options', + 'render_diffraction_pattern', + 'render_diffraction_aggregate', + 'render_probe', + 'render_probe_modes', + 'render_object_layer', + 'render_fluorescence_element', + } + assert expected.issubset(names), f'missing MCP tools: {expected - names}' + + stats_result = await mcp.call_tool('get_store_stats', {}) + payload = stats_result.structured_content + assert payload['campaign_count'] == 1 + assert payload['diffraction_count'] == 1 + assert payload['product_count'] == 1 + + lineage_result = await mcp.call_tool('get_lineage', {'uuid': str(p)}) + payload = lineage_result.structured_content + # Optional return types are wrapped as {'result': ...} by fastmcp + lineage_payload = payload['result'] if 'result' in payload else payload + assert lineage_payload is not None + ancestor_uuids = {a['uuid'] for a in lineage_payload['ancestors']} + assert str(d) in ancestor_uuids + + +@pytest.fixture +def bound_mcp(session_provider, layout): # type: ignore[no-untyped-def] + bind_session_provider(session_provider) + bind_layout(layout) + return create_mcp_server() + + +async def _ingest(session_provider, layout, manifest_path) -> None: # type: ignore[no-untyped-def] + async with session_provider.session_factory() as session: + await ingest_manifest(session, layout, manifest_path) + await session.commit() + + +async def test_mcp_get_visualization_options(bound_mcp) -> None: # type: ignore[no-untyped-def] + result = await bound_mcp.call_tool('get_visualization_options', {}) + payload = result.structured_content + assert 'amplitude' in payload['components'] + assert 'hsv_value' in payload['color_models'] + assert payload['transforms'] == ['identity', 'sqrt', 'log2', 'log', 'log10'] + + +async def test_mcp_render_diffraction_pattern( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(session_provider, layout, layout.manifest_path('diffraction', d)) + result = await bound_mcp.call_tool( + 'render_diffraction_pattern', {'uuid': str(d), 'index': 0, 'colormap': 'gray'} + ) + _decode_image_content(result, expected_h=8, expected_w=12) + + +async def test_mcp_render_diffraction_aggregate( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(session_provider, layout, layout.manifest_path('diffraction', d)) + result = await bound_mcp.call_tool('render_diffraction_aggregate', {'uuid': str(d)}) + _decode_image_content(result, expected_h=8, expected_w=12) + + +async def test_mcp_render_probe( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_product +) -> None: + p = seed_product() + await _ingest(session_provider, layout, layout.manifest_path('product', p)) + result = await bound_mcp.call_tool('render_probe', {'uuid': str(p), 'component': 'amplitude'}) + _decode_image_content(result, expected_h=8, expected_w=8) + + +async def test_mcp_render_probe_cylindrical( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_product +) -> None: + p = seed_product() + await _ingest(session_provider, layout, layout.manifest_path('product', p)) + result = await bound_mcp.call_tool('render_probe', {'uuid': str(p), 'color_model': 'hsv_value'}) + _decode_image_content(result, expected_h=8, expected_w=8) + + +async def test_mcp_render_probe_modes( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_product +) -> None: + p = seed_product() + await _ingest(session_provider, layout, layout.manifest_path('product', p)) + result = await bound_mcp.call_tool('render_probe_modes', {'uuid': str(p)}) + _decode_image_content(result, expected_h=8, expected_w=8) + + +async def test_mcp_render_object_layer( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_product +) -> None: + p = seed_product() + await _ingest(session_provider, layout, layout.manifest_path('product', p)) + result = await bound_mcp.call_tool( + 'render_object_layer', {'uuid': str(p), 'layer': 0, 'component': 'phase_rad'} + ) + _decode_image_content(result, expected_h=16, expected_w=16) + + +async def test_mcp_render_fluorescence_element( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_fluorescence +) -> None: + f = seed_fluorescence(elements=['Fe', 'Cu']) + await _ingest(session_provider, layout, layout.manifest_path('fluorescence', f)) + result = await bound_mcp.call_tool( + 'render_fluorescence_element', {'uuid': str(f), 'name': 'Fe'} + ) + _decode_image_content(result, expected_h=6, expected_w=10) + + +async def test_mcp_render_object_layer_out_of_range( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_product +) -> None: + p = seed_product() + await _ingest(session_provider, layout, layout.manifest_path('product', p)) + with pytest.raises(ToolError, match='out of range'): + await bound_mcp.call_tool('render_object_layer', {'uuid': str(p), 'layer': 99}) + + +async def test_mcp_render_bad_colormap( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(session_provider, layout, layout.manifest_path('diffraction', d)) + with pytest.raises(ToolError, match='invalid colormap'): + await bound_mcp.call_tool( + 'render_diffraction_pattern', + {'uuid': str(d), 'index': 0, 'colormap': 'not-a-colormap'}, + ) + + +async def test_mcp_render_probe_component_and_color_model_rejected( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_product +) -> None: + p = seed_product() + await _ingest(session_provider, layout, layout.manifest_path('product', p)) + with pytest.raises(ToolError, match='mutually exclusive'): + await bound_mcp.call_tool( + 'render_probe', + {'uuid': str(p), 'component': 'amplitude', 'color_model': 'hsv_value'}, + ) + + +async def test_mcp_render_fluorescence_missing_element( # type: ignore[no-untyped-def] + bound_mcp, session_provider, layout, seed_fluorescence +) -> None: + f = seed_fluorescence(elements=['Fe', 'Cu']) + await _ingest(session_provider, layout, layout.manifest_path('fluorescence', f)) + with pytest.raises(ToolError, match='not found'): + await bound_mcp.call_tool('render_fluorescence_element', {'uuid': str(f), 'name': 'Au'}) diff --git a/tests/ptychodus_store/test_pipeline.py b/tests/ptychodus_store/test_pipeline.py new file mode 100644 index 000000000..aca5032ad --- /dev/null +++ b/tests/ptychodus_store/test_pipeline.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import json + +import pytest + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.ingest.pipeline import delete_manifest, ingest_manifest + + +pytestmark = pytest.mark.asyncio + + +async def test_ingest_diffraction_valid(db_session, layout, seed_diffraction): # type: ignore[no-untyped-def] + uuid = seed_diffraction() + manifest_path = layout.manifest_path('diffraction', uuid) + await ingest_manifest(db_session, layout, manifest_path) + await db_session.commit() + + row = await repo.get_row(db_session, 'diffraction', uuid) + assert row is not None + assert row.ingest_state == IngestState.VALID + assert row.probe_energy_eV == 8000.0 + assert row.num_patterns_total == 4 + assert row.pattern_height_px == 8 + assert row.pattern_width_px == 12 + assert row.detector_pixel_width_m is not None + + +async def test_ingest_diffraction_missing_h5(db_session, layout, seed_diffraction): # type: ignore[no-untyped-def] + uuid = seed_diffraction(write_h5=False) + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', uuid)) + await db_session.commit() + row = await repo.get_row(db_session, 'diffraction', uuid) + assert row is not None + assert row.ingest_state == IngestState.MISSING_FILES + + +async def test_ingest_bad_json(db_session, layout, tmp_storage_root): # type: ignore[no-untyped-def] + from uuid import uuid4 + + uuid = uuid4() + folder = tmp_storage_root / 'product' / str(uuid) + folder.mkdir(parents=True) + (folder / 'manifest.json').write_text('not json at all') + await ingest_manifest(db_session, layout, folder / 'manifest.json') + await db_session.commit() + + row = await repo.get_row(db_session, 'product', uuid) + assert row is not None + assert row.ingest_state == IngestState.INVALID + assert row.error_message is not None + + +async def test_derived_from_creates_edge( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction, seed_product +): + d_uuid = seed_diffraction() + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d_uuid)) + p_uuid = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d_uuid)}]) + await ingest_manifest(db_session, layout, layout.manifest_path('product', p_uuid)) + await db_session.commit() + + edges = await repo.outgoing_edges(db_session, p_uuid) + assert len(edges) == 1 + assert edges[0].target_uuid == d_uuid + assert edges[0].target_kind == 'diffraction' + + p_row = await repo.get_row(db_session, 'product', p_uuid) + assert p_row is not None + assert p_row.ingest_state == IngestState.VALID + + +async def test_orphan_then_resolves( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction, seed_product +): + # Ingest a product whose parent doesn't yet exist + from uuid import uuid4 + + missing_parent = uuid4() + p_uuid = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(missing_parent)}]) + await ingest_manifest(db_session, layout, layout.manifest_path('product', p_uuid)) + await db_session.commit() + + p = await repo.get_row(db_session, 'product', p_uuid) + assert p is not None + assert p.ingest_state == IngestState.ORPHANED + + # Now the parent shows up under the same UUID + d_uuid = seed_diffraction(uuid=missing_parent) + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d_uuid)) + await db_session.commit() + + p_again = await repo.get_row(db_session, 'product', p_uuid) + assert p_again is not None + assert p_again.ingest_state == IngestState.VALID + + +async def test_delete_manifest_removes_row_and_propagates( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction, seed_product +): + d_uuid = seed_diffraction() + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d_uuid)) + p_uuid = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d_uuid)}]) + await ingest_manifest(db_session, layout, layout.manifest_path('product', p_uuid)) + await db_session.commit() + + await delete_manifest(db_session, layout, layout.manifest_path('diffraction', d_uuid)) + await db_session.commit() + + assert await repo.get_row(db_session, 'diffraction', d_uuid) is None + p = await repo.get_row(db_session, 'product', p_uuid) + assert p is not None + assert p.ingest_state == IngestState.ORPHANED + + +async def test_multi_parent_edges( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction, seed_product +): + d1 = seed_diffraction() + d2 = seed_diffraction() + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d1)) + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d2)) + + p_uuid = seed_product( + derived_from=[ + {'kind': 'diffraction', 'uuid': str(d1)}, + {'kind': 'diffraction', 'uuid': str(d2)}, + ] + ) + await ingest_manifest(db_session, layout, layout.manifest_path('product', p_uuid)) + await db_session.commit() + + edges = await repo.outgoing_edges(db_session, p_uuid) + assert {e.target_uuid for e in edges} == {d1, d2} + + +async def test_rewrite_manifest_shrinks_edges( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction, seed_product +): + d1 = seed_diffraction() + d2 = seed_diffraction() + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d1)) + await ingest_manifest(db_session, layout, layout.manifest_path('diffraction', d2)) + + p_uuid = seed_product( + derived_from=[ + {'kind': 'diffraction', 'uuid': str(d1)}, + {'kind': 'diffraction', 'uuid': str(d2)}, + ] + ) + p_manifest = layout.manifest_path('product', p_uuid) + await ingest_manifest(db_session, layout, p_manifest) + await db_session.commit() + assert len(await repo.outgoing_edges(db_session, p_uuid)) == 2 + + # Shrink derived_from to a single parent + from datetime import datetime, timezone + + p_manifest.write_text( + json.dumps( + { + 'schema_version': 1, + 'kind': 'product', + 'uuid': str(p_uuid), + 'created_at': datetime.now(timezone.utc).isoformat(), + 'derived_from': [{'kind': 'diffraction', 'uuid': str(d1)}], + } + ) + ) + await ingest_manifest(db_session, layout, p_manifest) + await db_session.commit() + + edges = await repo.outgoing_edges(db_session, p_uuid) + assert len(edges) == 1 + assert edges[0].target_uuid == d1 diff --git a/tests/ptychodus_store/test_reconciler.py b/tests/ptychodus_store/test_reconciler.py new file mode 100644 index 000000000..8e9cea849 --- /dev/null +++ b/tests/ptychodus_store/test_reconciler.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import shutil + +import pytest + +from ptychodus_store.db import repositories as repo +from ptychodus_store.db.base import IngestState +from ptychodus_store.ingest.reconciler import full_rescan + +pytestmark = pytest.mark.asyncio + + +async def test_full_rescan_counts( # type: ignore[no-untyped-def] + db_session, layout, seed_campaign, seed_diffraction, seed_product, seed_fluorescence +): + c = seed_campaign() + d = seed_diffraction(campaign_uuid=c) + p = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d)}]) + seed_fluorescence(derived_from=[{'kind': 'product', 'uuid': str(p)}]) + + counts = await full_rescan(db_session, layout) + assert counts == { + 'campaign': 1, + 'diffraction': 1, + 'product': 1, + 'fluorescence': 1, + } + + +async def test_full_rescan_deletes_stale( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction +): + d1 = seed_diffraction() + d2 = seed_diffraction() + await full_rescan(db_session, layout) + assert await repo.get_row(db_session, 'diffraction', d1) is not None + assert await repo.get_row(db_session, 'diffraction', d2) is not None + + # Remove d2 from disk + shutil.rmtree(layout.resource_folder('diffraction', d2)) + await full_rescan(db_session, layout) + assert await repo.get_row(db_session, 'diffraction', d1) is not None + assert await repo.get_row(db_session, 'diffraction', d2) is None + + +async def test_rescan_resolves_forward_refs( # type: ignore[no-untyped-def] + db_session, layout, seed_diffraction, seed_product +): + d = seed_diffraction() + p = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d)}]) + await full_rescan(db_session, layout) + + row = await repo.get_row(db_session, 'product', p) + assert row is not None + assert row.ingest_state == IngestState.VALID diff --git a/tests/ptychodus_store/test_routers.py b/tests/ptychodus_store/test_routers.py new file mode 100644 index 000000000..ed2edbd3e --- /dev/null +++ b/tests/ptychodus_store/test_routers.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import pytest + +from ptychodus_store.ingest.pipeline import ingest_manifest + +pytestmark = pytest.mark.asyncio + + +async def _ingest(client, db_engine, layout, manifest_path) -> None: # type: ignore[no-untyped-def] + # Use a fresh session from the same engine the app is wired to. + from sqlalchemy.ext.asyncio import async_sessionmaker + + factory = async_sessionmaker(db_engine, expire_on_commit=False) + async with factory() as session: + await ingest_manifest(session, layout, manifest_path) + await session.commit() + + +async def test_health(app_client) -> None: # type: ignore[no-untyped-def] + resp = await app_client.get('/api/v1/health') + assert resp.status_code == 200 + body = resp.json() + assert body['db'] == 'ok' + assert body['watcher'] == 'disabled' + + +async def test_campaign_list_and_get( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_campaign +): + c = seed_campaign(sample_name='alpha', tags=['benchmark']) + await _ingest(app_client, db_engine, layout, layout.manifest_path('campaign', c)) + + resp = await app_client.get('/api/v1/campaign') + assert resp.status_code == 200 + body = resp.json() + assert body['total'] == 1 + assert body['items'][0]['uuid'] == str(c) + assert body['items'][0]['sample_name'] == 'alpha' + + one = await app_client.get(f'/api/v1/campaign/{c}') + assert one.status_code == 200 + assert one.json()['uuid'] == str(c) + + miss = await app_client.get('/api/v1/campaign/00000000-0000-0000-0000-000000000000') + assert miss.status_code == 404 + + +async def test_diffraction_filter_by_campaign( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_campaign, seed_diffraction +): + c = seed_campaign() + other = seed_campaign() + d1 = seed_diffraction(campaign_uuid=c) + d2 = seed_diffraction(campaign_uuid=other) + for m in ( + layout.manifest_path('campaign', c), + layout.manifest_path('campaign', other), + layout.manifest_path('diffraction', d1), + layout.manifest_path('diffraction', d2), + ): + await _ingest(app_client, db_engine, layout, m) + + resp = await app_client.get('/api/v1/diffraction', params={'campaign_uuid': str(c)}) + assert resp.status_code == 200 + body = resp.json() + assert body['total'] == 1 + assert body['items'][0]['uuid'] == str(d1) + + +async def test_product_derived_from_filter( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction, seed_product +): + d = seed_diffraction() + p_match = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d)}]) + p_other = seed_product() + for m in ( + layout.manifest_path('diffraction', d), + layout.manifest_path('product', p_match), + layout.manifest_path('product', p_other), + ): + await _ingest(app_client, db_engine, layout, m) + + resp = await app_client.get('/api/v1/product', params={'derived_from_uuid': str(d)}) + assert resp.status_code == 200 + body = resp.json() + assert body['total'] == 1 + assert body['items'][0]['uuid'] == str(p_match) + + +async def test_lineage_dag_walk( # type: ignore[no-untyped-def] + app_client, + db_engine, + layout, + seed_campaign, + seed_diffraction, + seed_product, + seed_fluorescence, +): + c = seed_campaign() + d = seed_diffraction(campaign_uuid=c) + p = seed_product(derived_from=[{'kind': 'diffraction', 'uuid': str(d)}]) + f = seed_fluorescence(derived_from=[{'kind': 'product', 'uuid': str(p)}]) + for m in ( + layout.manifest_path('campaign', c), + layout.manifest_path('diffraction', d), + layout.manifest_path('product', p), + layout.manifest_path('fluorescence', f), + ): + await _ingest(app_client, db_engine, layout, m) + + # Walk from the fluorescence node — ancestors should reach the diffraction + resp = await app_client.get(f'/api/v1/lineage/{f}') + assert resp.status_code == 200 + body = resp.json() + ancestor_uuids = {a['uuid'] for a in body['ancestors']} + assert str(p) in ancestor_uuids + assert str(d) in ancestor_uuids + assert body['campaign'] is not None + assert body['campaign']['uuid'] == str(c) + + # Walk from the diffraction node — descendants should include product and fluorescence + resp2 = await app_client.get(f'/api/v1/lineage/{d}') + body2 = resp2.json() + desc_uuids = {x['uuid'] for x in body2['descendants']} + assert str(p) in desc_uuids + assert str(f) in desc_uuids + + +async def test_file_download( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +): + d = seed_diffraction() + await _ingest(app_client, db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get(f'/api/v1/diffraction/{d}/files/diffraction') + assert resp.status_code == 200 + assert resp.headers['content-type'] == 'application/x-hdf5' + assert len(resp.content) > 0 + + +async def test_admin_stats( # type: ignore[no-untyped-def] + app_client, + db_engine, + layout, + seed_campaign, + seed_diffraction, +): + c = seed_campaign() + d = seed_diffraction() + await _ingest(app_client, db_engine, layout, layout.manifest_path('campaign', c)) + await _ingest(app_client, db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get('/api/v1/admin/stats') + assert resp.status_code == 200 + body = resp.json() + assert body['campaign_count'] == 1 + assert body['diffraction_count'] == 1 + assert body['product_count'] == 0 + assert body['fluorescence_count'] == 0 diff --git a/tests/ptychodus_store/test_visualization_endpoints.py b/tests/ptychodus_store/test_visualization_endpoints.py new file mode 100644 index 000000000..30c1679fe --- /dev/null +++ b/tests/ptychodus_store/test_visualization_endpoints.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import base64 +from io import BytesIO + +import pytest +from PIL import Image + +from ptychodus_store.ingest.pipeline import ingest_manifest + +pytestmark = pytest.mark.asyncio + + +async def _ingest(db_engine, layout, manifest_path) -> None: # type: ignore[no-untyped-def] + from sqlalchemy.ext.asyncio import async_sessionmaker + + factory = async_sessionmaker(db_engine, expire_on_commit=False) + async with factory() as session: + await ingest_manifest(session, layout, manifest_path) + await session.commit() + + +def _decode_png(body: dict, expected_h: int, expected_w: int) -> Image.Image: + assert body['mime_type'] == 'image/png' + assert body['shape_h_px'] == expected_h + assert body['shape_w_px'] == expected_w + png_bytes = base64.b64decode(body['png_base64']) + image = Image.open(BytesIO(png_bytes)) + image.load() + assert image.size == (expected_w, expected_h) + return image + + +async def test_visualization_options(app_client) -> None: # type: ignore[no-untyped-def] + resp = await app_client.get('/api/v1/visualization/options') + assert resp.status_code == 200 + body = resp.json() + assert body['transforms'] == ['identity', 'sqrt', 'log2', 'log', 'log10'] + assert 'amplitude' in body['components'] + assert 'phase_rad' in body['components'] + assert 'hsv_value' in body['color_models'] + assert len(body['colormaps_linear']) > 0 + assert 'gray' in body['colormaps_linear'] or 'gray' in body['colormaps_cyclic'] + + +async def test_diffraction_pattern_image( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get( + f'/api/v1/diffraction/{d}/patterns/0/image', + params={'colormap': 'gray', 'transform': 'identity'}, + ) + assert resp.status_code == 200 + body = resp.json() + _decode_png(body, expected_h=8, expected_w=12) + assert body['pixel_width_m'] == pytest.approx(55e-6) + + +async def test_diffraction_pattern_out_of_range( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get(f'/api/v1/diffraction/{d}/patterns/999/image') + assert resp.status_code == 404 + + +async def test_diffraction_aggregate_image( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get(f'/api/v1/diffraction/{d}/patterns/aggregate/image') + assert resp.status_code == 200 + _decode_png(resp.json(), expected_h=8, expected_w=12) + + +async def test_probe_image( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get( + f'/api/v1/product/{p}/probe/image', + params={'component': 'amplitude', 'colormap': 'gray'}, + ) + assert resp.status_code == 200 + _decode_png(resp.json(), expected_h=8, expected_w=8) + + +async def test_probe_image_cylindrical( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get( + f'/api/v1/product/{p}/probe/image', params={'color_model': 'hsv_value'} + ) + assert resp.status_code == 200 + + +async def test_probe_image_component_and_color_model_rejected( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get( + f'/api/v1/product/{p}/probe/image', + params={'component': 'amplitude', 'color_model': 'hsv_value'}, + ) + assert resp.status_code == 400 + + +async def test_probe_incoherent_out_of_range( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get(f'/api/v1/product/{p}/probe/image', params={'incoherent': 42}) + assert resp.status_code == 404 + + +async def test_probe_modes_image( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get(f'/api/v1/product/{p}/probe/modes/image') + assert resp.status_code == 200 + # Fixture writes a single incoherent mode of shape (8, 8); tiled width == 8. + _decode_png(resp.json(), expected_h=8, expected_w=8) + + +async def test_object_layer_image( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get( + f'/api/v1/product/{p}/object/0/image', params={'component': 'phase_rad'} + ) + assert resp.status_code == 200 + _decode_png(resp.json(), expected_h=16, expected_w=16) + + +async def test_object_layer_out_of_range( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product +) -> None: + p = seed_product() + await _ingest(db_engine, layout, layout.manifest_path('product', p)) + + resp = await app_client.get(f'/api/v1/product/{p}/object/99/image') + assert resp.status_code == 404 + + +async def test_fluorescence_element_image( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_fluorescence +) -> None: + f = seed_fluorescence(elements=['Fe', 'Cu']) + await _ingest(db_engine, layout, layout.manifest_path('fluorescence', f)) + + resp = await app_client.get(f'/api/v1/fluorescence/{f}/elements/Fe/image') + assert resp.status_code == 200 + _decode_png(resp.json(), expected_h=6, expected_w=10) + + +async def test_fluorescence_element_missing( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_fluorescence +) -> None: + f = seed_fluorescence(elements=['Fe', 'Cu']) + await _ingest(db_engine, layout, layout.manifest_path('fluorescence', f)) + + resp = await app_client.get(f'/api/v1/fluorescence/{f}/elements/Au/image') + assert resp.status_code == 404 + detail = resp.json()['detail'] + assert 'Fe' in detail['available'] + + +async def test_fluorescence_with_product_pixel_geometry( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_product, seed_fluorescence +) -> None: + p = seed_product() + f = seed_fluorescence(elements=['Fe']) + for m in ( + layout.manifest_path('product', p), + layout.manifest_path('fluorescence', f), + ): + await _ingest(db_engine, layout, m) + + resp = await app_client.get( + f'/api/v1/fluorescence/{f}/elements/Fe/image', params={'product_uuid': str(p)} + ) + assert resp.status_code == 200 + body = resp.json() + # Product fixture writes object pixel geometry = 1e-9 m + assert body['pixel_width_m'] == pytest.approx(1e-9) + + +async def test_bad_colormap_returns_400( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get( + f'/api/v1/diffraction/{d}/patterns/0/image', params={'colormap': 'not-a-colormap'} + ) + assert resp.status_code == 400 + detail = resp.json()['detail'] + assert 'valid_linear' in detail + assert 'valid_cyclic' in detail + + +async def test_bad_transform_returns_400( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get( + f'/api/v1/diffraction/{d}/patterns/0/image', params={'transform': 'cbrt'} + ) + assert resp.status_code == 400 + + +async def test_component_on_real_endpoint_rejected( # type: ignore[no-untyped-def] + app_client, db_engine, layout, seed_diffraction +) -> None: + d = seed_diffraction() + await _ingest(db_engine, layout, layout.manifest_path('diffraction', d)) + + resp = await app_client.get( + f'/api/v1/diffraction/{d}/patterns/0/image', params={'component': 'amplitude'} + ) + assert resp.status_code == 400 diff --git a/tests/ptychodus_store/test_watcher.py b/tests/ptychodus_store/test_watcher.py new file mode 100644 index 000000000..645d930b1 --- /dev/null +++ b/tests/ptychodus_store/test_watcher.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import asyncio +import json +import shutil +from datetime import datetime, timezone +from pathlib import Path +from uuid import uuid4 + +import h5py +import numpy as np +import pytest + +from ptychodus_store.db import repositories as repo +from ptychodus_store.ingest.watcher import ManifestWatcher + +pytestmark = pytest.mark.asyncio + + +async def _wait_until(predicate, *, timeout: float = 8.0, interval: float = 0.2): + """Poll `predicate` (async callable) until truthy or timeout.""" + deadline = asyncio.get_event_loop().time() + timeout + while True: + result = await predicate() + if result: + return result + if asyncio.get_event_loop().time() > deadline: + return result + await asyncio.sleep(interval) + + +def _write_diffraction_folder(root: Path, uuid_str: str) -> Path: + folder = root / 'diffraction' / uuid_str + folder.mkdir(parents=True) + with h5py.File(folder / 'diffraction.h5', 'w') as f: + ds = f.create_dataset('patterns', data=np.zeros((2, 4, 4), dtype=np.uint16)) + ds.attrs['detector_pixel_width_m'] = 1e-5 + ds.attrs['detector_pixel_height_m'] = 1e-5 + f.create_dataset('indexes', data=np.arange(2)) + f.create_dataset('bad_pixels', data=np.zeros((4, 4), dtype=bool)) + manifest = { + 'schema_version': 1, + 'kind': 'diffraction', + 'uuid': uuid_str, + 'created_at': datetime.now(timezone.utc).isoformat(), + 'label': 'live', + } + (folder / 'manifest.json').write_text(json.dumps(manifest)) + return folder + + +async def test_watcher_picks_up_new_manifest( # type: ignore[no-untyped-def] + tmp_storage_root, layout, session_provider +): + loop = asyncio.get_running_loop() + watcher = ManifestWatcher( + layout, + session_provider.session_factory, + loop, + polling_interval_s=0.3, + debounce_window_s=0.2, + ) + watcher.start() + try: + uuid = uuid4() + _write_diffraction_folder(tmp_storage_root, str(uuid)) + + async def check(): + async with session_provider.session_factory() as session: + return await repo.get_row(session, 'diffraction', uuid) + + row = await _wait_until(check, timeout=8.0) + assert row is not None, 'watcher did not ingest the new manifest in time' + finally: + watcher.stop() + + +async def test_watcher_handles_manifest_delete( # type: ignore[no-untyped-def] + tmp_storage_root, layout, session_provider +): + loop = asyncio.get_running_loop() + watcher = ManifestWatcher( + layout, + session_provider.session_factory, + loop, + polling_interval_s=0.3, + debounce_window_s=0.2, + ) + watcher.start() + try: + uuid = uuid4() + folder = _write_diffraction_folder(tmp_storage_root, str(uuid)) + + async def appeared(): + async with session_provider.session_factory() as session: + return await repo.get_row(session, 'diffraction', uuid) + + row = await _wait_until(appeared, timeout=8.0) + assert row is not None + + shutil.rmtree(folder) + + async def gone(): + async with session_provider.session_factory() as session: + return (await repo.get_row(session, 'diffraction', uuid)) is None + + assert await _wait_until(gone, timeout=8.0) + finally: + watcher.stop() diff --git a/tests/subprocess_child_fixtures.py b/tests/subprocess_child_fixtures.py new file mode 100644 index 000000000..c4f8da77c --- /dev/null +++ b/tests/subprocess_child_fixtures.py @@ -0,0 +1,63 @@ +"""Child-side entry points used by ``test_subprocess_reconstructor``. + +The ``spawn``-context child imports this module by dotted path, so it must be +resolvable on the child's ``sys.path``; the test module puts the tests +directory there before spawning. Not collected by pytest (the filename does +not match ``test_*.py``). Uses no GPU frameworks. Do not import from +application code. +""" + +from __future__ import annotations + +import logging +import pickle +import time +from typing import Any + +from ptychodus.api.reconstructor import ReconstructOutput, TrainOutput +from ptychodus.model.processing.subprocess_reconstructor import ( + TAG_MODEL_SAVED, + TAG_OUTPUT, + TAG_SETTINGS_SYNC, + TAG_TRAIN_OUTPUT, +) + +logger = logging.getLogger(__name__) + + +def yield_n_outputs(payload: Any, queue: Any) -> None: + """Emit ``payload['n']`` :class:`ReconstructOutput`s with product=None.""" + n = int(payload['n']) + for i in range(n): + output = ReconstructOutput(product=payload['product'], progress=i + 1) + queue.put((TAG_OUTPUT, pickle.dumps(output))) + + +def raise_immediately(payload: Any, queue: Any) -> None: + """Raise a :class:`ValueError` before emitting anything.""" + raise ValueError(payload['message']) + + +def hang_forever(payload: Any, queue: Any) -> None: + """Sleep so the parent must terminate us.""" + time.sleep(3600.0) + + +def emit_log_then_output(payload: Any, queue: Any) -> None: + """Log a message, then emit one output; parent must see log before output.""" + logger.warning(payload['log_message']) + output = ReconstructOutput(product=payload['product'], progress=1) + queue.put((TAG_OUTPUT, pickle.dumps(output))) + + +def emit_settings_sync_then_output(payload: Any, queue: Any) -> None: + """Emit a settings-sync message, then an output.""" + queue.put((TAG_SETTINGS_SYNC, payload['settings'])) + output = ReconstructOutput(product=payload['product'], progress=1) + queue.put((TAG_OUTPUT, pickle.dumps(output))) + + +def train_and_save(payload: Any, queue: Any) -> None: + """Emit one :class:`TrainOutput` and a model-saved path.""" + queue.put((TAG_TRAIN_OUTPUT, pickle.dumps(TrainOutput(progress=1)))) + queue.put((TAG_MODEL_SAVED, payload['saved_path'])) diff --git a/tests/test_diffraction_loader.py b/tests/test_diffraction_loader.py new file mode 100644 index 000000000..2e8275ecf --- /dev/null +++ b/tests/test_diffraction_loader.py @@ -0,0 +1,179 @@ +"""Regression tests for the per-array LoadArray pipeline. + +Covers the invariant that when pattern processing is enabled, the bad-pixel +mask handed to the per-array AssembledDiffractionData matches the processed +patterns' shape, and that bad-pixel repair (zeroing) is applied to raw +patterns before crop/bin/pad/flip/transpose runs. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import numpy + +from ptychodus.api.diffraction import ( + DiffractionMetadata, + SimpleDiffractionArray, + SimpleDiffractionDataset, +) +from ptychodus.api.geometry import ImageExtent +from ptychodus.api.settings import SettingsRegistry +from ptychodus.api.tree import SimpleTreeNode +from ptychodus.model.diffraction.dataset import AssembledDiffractionDataset +from ptychodus.model.diffraction.settings import DetectorSettings, DiffractionSettings +from ptychodus.model.diffraction.sizer import PatternSizer + + +def _make_dataset( + detector_height: int, + detector_width: int, + crop_center_y: int, + crop_center_x: int, + crop_height: int, + crop_width: int, +) -> AssembledDiffractionDataset: + registry = SettingsRegistry() + detector_settings = DetectorSettings(registry) + diffraction_settings = DiffractionSettings(registry) + + diffraction_settings.crop_enabled.set_value(True) + diffraction_settings.crop_center_y_px.set_value(crop_center_y) + diffraction_settings.crop_center_x_px.set_value(crop_center_x) + diffraction_settings.crop_height_px.set_value(crop_height) + diffraction_settings.crop_width_px.set_value(crop_width) + + sizer = PatternSizer(diffraction_settings) + task_manager = MagicMock() + task_monitor = MagicMock() + + dataset = AssembledDiffractionDataset( + diffraction_settings, + sizer, + detector_settings, + task_manager, + task_monitor, + ) + # Give the dataset a metadata source with the intended detector extent so the + # per-dataset extent flows through reload() -> the sizer's pipeline builder. + metadata = DiffractionMetadata( + num_patterns_per_array=[0], + pattern_dtype=numpy.dtype(numpy.uint16), + detector_extent=ImageExtent(width_px=detector_width, height_px=detector_height), + ) + contents_tree = SimpleTreeNode.create_root(['Name', 'Type', 'Details']) + source = SimpleDiffractionDataset(metadata, contents_tree, []) + dataset.reload(source) + return dataset + + +def test_load_array_matches_shapes_and_repairs_bad_pixels_before_crop() -> None: + """Reproduces the fly001.ini warning: raw 40x60 → crop to 12x16. + + Without the fix, LoadArray hands AssembledDiffractionData mismatched shapes + (patterns (12,16) vs bad_pixels (40,60)) and ValueError propagates. With + the fix, shapes match, the raw bad pixel appears at its cropped location, + and its raw huge value has been zeroed rather than surviving into the crop. + """ + dataset = _make_dataset( + detector_height=40, + detector_width=60, + crop_center_y=20, + crop_center_x=30, + crop_height=12, + crop_width=16, + ) + + # Bad pixel at raw (18, 30). Crop [y=14..26, x=22..38] → cropped (4, 8). + raw_bad = numpy.zeros((40, 60), dtype=bool) + raw_bad[18, 30] = True + dataset.set_bad_pixels(raw_bad) + + raw_patterns = numpy.ones((3, 40, 60), dtype=numpy.uint16) + raw_patterns[:, 18, 30] = 65535 # saturation that must not leak into the crop + array = SimpleDiffractionArray('test', numpy.arange(3, dtype=numpy.intp), raw_patterns) + + captured: dict[str, object] = {} + + def stub_assemble_array(array_index, label, data): # type: ignore[no-untyped-def] + captured['data'] = data + + dataset.assemble_array = stub_assemble_array # type: ignore[method-assign] + + task = dataset.create_array_loader(array, process_patterns=True) + task() + + data = captured['data'] + assert data.get_patterns_shape() == (3, 12, 16) # type: ignore[attr-defined] + assert data.get_bad_pixels().shape == (12, 16) # type: ignore[attr-defined] + assert data.get_bad_pixels()[4, 8] # type: ignore[attr-defined] + # Repair happened before crop: the saturated raw value was zeroed. + assert int(data.get_pattern(0)[4, 8]) == 0 # type: ignore[attr-defined] + # Sanity: an unmasked pixel elsewhere in the crop still holds its raw value. + assert int(data.get_pattern(0)[0, 0]) == 1 # type: ignore[attr-defined] + + +def test_load_all_arrays_preserves_raw_metadata_extent() -> None: + """load_all_arrays(process_patterns=True) must not conflate raw detector geometry + with the processed bad-pixel mask shape. + + Regression for the fly001.ini failure: the old code rebuilt the internal + SimpleDiffractionDataset with raw metadata but the *processed* bad_pixels, + tripping SimpleDiffractionDataset's shape-mismatch check. The processed + mask belongs on self._data (AssembledDiffractionData), not on the metadata + holder. + """ + dataset = _make_dataset( + detector_height=40, + detector_width=60, + crop_center_y=20, + crop_center_x=30, + crop_height=12, + crop_width=16, + ) + + dataset.load_all_arrays(process_patterns=True, block=True) + + # Metadata still reports the raw detector extent. + extent = dataset.get_metadata().detector_extent + assert extent.width_px == 60 + assert extent.height_px == 40 + # The assembled data holds the processed bad-pixel mask matching the crop output. + assert dataset.get_assembled_data().get_bad_pixels().shape == (12, 16) + + +def test_load_array_without_processing_keeps_raw_shapes() -> None: + """process_patterns=False should skip the processor and keep raw shapes intact.""" + dataset = _make_dataset( + detector_height=40, + detector_width=60, + crop_center_y=20, + crop_center_x=30, + crop_height=12, + crop_width=16, + ) + + raw_bad = numpy.zeros((40, 60), dtype=bool) + raw_bad[5, 5] = True + dataset.set_bad_pixels(raw_bad) + + raw_patterns = numpy.full((2, 40, 60), 7, dtype=numpy.uint16) + raw_patterns[:, 5, 5] = 999 + array = SimpleDiffractionArray('raw', numpy.arange(2, dtype=numpy.intp), raw_patterns) + + captured: dict[str, object] = {} + + def stub_assemble_array(array_index, label, data): # type: ignore[no-untyped-def] + captured['data'] = data + + dataset.assemble_array = stub_assemble_array # type: ignore[method-assign] + + task = dataset.create_array_loader(array, process_patterns=False) + task() + + data = captured['data'] + assert data.get_patterns_shape() == (2, 40, 60) # type: ignore[attr-defined] + assert data.get_bad_pixels().shape == (40, 60) # type: ignore[attr-defined] + # Repair still runs because the raw mask is meaningful even without geometric processing. + assert int(data.get_pattern(0)[5, 5]) == 0 # type: ignore[attr-defined] + assert int(data.get_pattern(0)[0, 0]) == 7 # type: ignore[attr-defined] diff --git a/tests/test_diffraction_processor.py b/tests/test_diffraction_processor.py index f21c9ce38..fc4de3689 100644 --- a/tests/test_diffraction_processor.py +++ b/tests/test_diffraction_processor.py @@ -1,24 +1,28 @@ -"""Unit tests for the diffraction model: processor ops, sizer, and the settings → processor wiring. +"""Unit tests for the diffraction preprocessing pipeline and the settings → pipeline factory. -These tests lock the intended behavior of `DiffractionPatternProcessor` and `PatternSizer` -so regressions in the op pipeline (crop, binning, padding, transpose, value filtering) and -the processed-extent math are caught at the unit level. +These tests lock the intended behavior of `DiffractionPrepPipeline` step ops and +`PatternSizer.get_prep_pipeline()` so regressions in the op chain (crop, binning, padding, +transpose, value filtering) and the processed-extent math are caught at the unit level. """ import numpy import pytest from ptychodus.api.diffraction import CropCenter, SimpleDiffractionArray -from ptychodus.api.geometry import ImageExtent +from ptychodus.api.diffraction_prep import ( + BinningStep, + CropStep, + DiffractionPrepPipeline, + DiffractionPrepStepUnion, + FilterValuesStep, + HorizontalFlipStep, + PaddingStep, + TransposeStep, + VerticalFlipStep, +) +from ptychodus.api.geometry import ImageExtent, PixelGeometry from ptychodus.api.settings import SettingsRegistry -from ptychodus.model.diffraction.processor import ( - DiffractionPatternBinning, - DiffractionPatternCrop, - DiffractionPatternFilterValues, - DiffractionPatternPadding, - DiffractionPatternProcessor, -) from ptychodus.model.diffraction.settings import DetectorSettings, DiffractionSettings from ptychodus.model.diffraction.sizer import PatternSizer @@ -32,13 +36,13 @@ def _zeros_patterns(shape: tuple[int, int, int]) -> numpy.ndarray: def test_filter_lower_bound_zeros_below() -> None: data = numpy.array([[[0, 1, 5, 10]]], dtype=numpy.int32) - out = DiffractionPatternFilterValues(lower_bound=3, upper_bound=None).apply(data.copy()) + out = FilterValuesStep(lower_bound=3, upper_bound=None).apply(data.copy()) assert out.tolist() == [[[0, 0, 5, 10]]] def test_filter_upper_bound_zeros_at_or_above() -> None: data = numpy.array([[[0, 1, 5, 10]]], dtype=numpy.int32) - out = DiffractionPatternFilterValues(lower_bound=None, upper_bound=5).apply(data.copy()) + out = FilterValuesStep(lower_bound=None, upper_bound=5).apply(data.copy()) assert out.tolist() == [[[0, 1, 0, 0]]] @@ -46,16 +50,23 @@ def test_filter_does_not_mutate_input() -> None: """B5: filter must not scribble on the caller's buffer.""" original = numpy.array([[[0, 1, 5, 10]]], dtype=numpy.int32) snapshot = original.copy() - DiffractionPatternFilterValues(lower_bound=3, upper_bound=8).apply(original) + FilterValuesStep(lower_bound=3, upper_bound=8).apply(original) assert numpy.array_equal(original, snapshot) +def test_filter_is_noop_on_mask() -> None: + """Value filtering is meaningless for boolean masks; step must short-circuit.""" + mask = numpy.array([[True, False], [True, True]], dtype=bool) + out = FilterValuesStep(lower_bound=1, upper_bound=2).apply(mask) + assert numpy.array_equal(out, mask) + + # ---------- Crop ---------- def test_crop_apply_reduces_shape_around_center() -> None: data = numpy.arange(2 * 8 * 8, dtype=numpy.uint16).reshape(2, 8, 8) - crop = DiffractionPatternCrop(CropCenter(position_x_px=4, position_y_px=4), ImageExtent(4, 4)) + crop = CropStep(center=CropCenter(position_x_px=4, position_y_px=4), extent=ImageExtent(4, 4)) out = crop.apply(data) assert out.shape == (2, 4, 4) # Center-crop of 8x8 around (4,4) with radius 2 = rows 2:6, cols 2:6 @@ -64,8 +75,8 @@ def test_crop_apply_reduces_shape_around_center() -> None: def test_crop_apply_mask_reduces_shape() -> None: data = numpy.ones((8, 8), dtype=bool) - crop = DiffractionPatternCrop(CropCenter(position_x_px=4, position_y_px=4), ImageExtent(4, 4)) - assert crop.apply(data, is_mask=True).shape == (4, 4) + crop = CropStep(center=CropCenter(position_x_px=4, position_y_px=4), extent=ImageExtent(4, 4)) + assert crop.apply(data).shape == (4, 4) # ---------- Binning ---------- @@ -73,7 +84,7 @@ def test_crop_apply_mask_reduces_shape() -> None: def test_binning_apply_sums_blocks() -> None: data = numpy.ones((1, 4, 4), dtype=numpy.uint16) - out = DiffractionPatternBinning(bin_size_x=2, bin_size_y=2).apply(data) + out = BinningStep(bin_size_x=2, bin_size_y=2).apply(data) assert out.shape == (1, 2, 2) assert (out == 4).all() @@ -81,7 +92,7 @@ def test_binning_apply_sums_blocks() -> None: def test_binning_apply_mask_logical_and() -> None: data = numpy.ones((4, 4), dtype=bool) data[0, 0] = False # one True-cell of the (0,0) 2x2 block becomes False - out = DiffractionPatternBinning(bin_size_x=2, bin_size_y=2).apply(data, is_mask=True) + out = BinningStep(bin_size_x=2, bin_size_y=2).apply(data) assert out.shape == (2, 2) assert out[0, 0] == False # logical AND of the block assert out[0, 1] == True @@ -89,13 +100,18 @@ def test_binning_apply_mask_logical_and() -> None: assert out[1, 1] == True +def test_binning_rejects_zero_bin_size() -> None: + with pytest.raises(ValueError): + BinningStep(bin_size_x=0, bin_size_y=1) + + # ---------- Padding (B1) ---------- def test_padding_apply_3d_produces_correct_shape() -> None: """B1: pad_width must broadcast to (ndim, 2); flat tuples raise ValueError.""" data = _zeros_patterns((2, 4, 4)) - out = DiffractionPatternPadding(pad_x=1, pad_y=1).apply(data) + out = PaddingStep(pad_x=1, pad_y=1).apply(data) assert out.shape == (2, 6, 6) assert (out == 0).all() @@ -103,7 +119,7 @@ def test_padding_apply_3d_produces_correct_shape() -> None: def test_padding_apply_mask_2d_produces_correct_shape() -> None: """B1 mirror: bad-pixels padding must not raise either.""" data = numpy.ones((4, 4), dtype=bool) - out = DiffractionPatternPadding(pad_x=1, pad_y=1).apply(data, is_mask=True) + out = PaddingStep(pad_x=1, pad_y=1).apply(data) assert out.shape == (6, 6) # Edges are padded with False; interior preserved. assert out[0, 0] == False @@ -112,103 +128,195 @@ def test_padding_apply_mask_2d_produces_correct_shape() -> None: def test_padding_asymmetric_pad_x_pad_y() -> None: data = _zeros_patterns((1, 4, 6)) - out = DiffractionPatternPadding(pad_x=2, pad_y=1).apply(data) + out = PaddingStep(pad_x=2, pad_y=1).apply(data) assert out.shape == (1, 6, 10) -# ---------- Processor.__call__ ---------- - - -def _processor( - *, - crop: DiffractionPatternCrop | None = None, - filter_values: DiffractionPatternFilterValues | None = None, - binning: DiffractionPatternBinning | None = None, - padding: DiffractionPatternPadding | None = None, - hflip: bool = False, - vflip: bool = False, - transpose: bool = False, -) -> DiffractionPatternProcessor: - return DiffractionPatternProcessor( - crop=crop, - filter_values=filter_values, - binning=binning, - padding=padding, - hflip=hflip, - vflip=vflip, - transpose=transpose, - ) +def test_padding_rejects_negative() -> None: + with pytest.raises(ValueError): + PaddingStep(pad_x=-1, pad_y=0) + + +# ---------- Pipeline.__call__ ---------- -def test_processor_promotes_2d_input_to_3d() -> None: +def _pipeline(*steps: DiffractionPrepStepUnion) -> DiffractionPrepPipeline: + return DiffractionPrepPipeline(steps=steps) + + +def test_pipeline_promotes_2d_input_to_3d() -> None: array = SimpleDiffractionArray( 'a', numpy.zeros(1, dtype=int), numpy.zeros((8, 8), dtype=numpy.uint16) ) - out = _processor()(array) + out = _pipeline()(array) assert out.get_patterns().shape == (1, 8, 8) -def test_processor_rejects_4d_input() -> None: +def test_pipeline_rejects_4d_input() -> None: array = SimpleDiffractionArray( 'a', numpy.zeros(1, dtype=int), numpy.zeros((1, 2, 4, 4), dtype=numpy.uint16) ) with pytest.raises(ValueError, match='Invalid diffraction pattern dimensions'): - _processor()(array) + _pipeline()(array) -def test_processor_padding_in_full_pipeline() -> None: - """Padding inside a processor stack must succeed (regression for B1).""" +def test_pipeline_padding_in_full_pipeline() -> None: + """Padding inside a pipeline stack must succeed (regression for B1).""" array = SimpleDiffractionArray( 'a', numpy.zeros(1, dtype=int), numpy.ones((1, 4, 4), dtype=numpy.uint16) ) - proc = _processor(padding=DiffractionPatternPadding(pad_x=1, pad_y=1)) - assert proc(array).get_patterns().shape == (1, 6, 6) + assert _pipeline(PaddingStep(pad_x=1, pad_y=1))(array).get_patterns().shape == (1, 6, 6) -def test_processor_transpose_swaps_spatial_axes() -> None: +def test_pipeline_transpose_swaps_spatial_axes() -> None: patterns = numpy.zeros((1, 3, 5), dtype=numpy.uint16) array = SimpleDiffractionArray('a', numpy.zeros(1, dtype=int), patterns) - assert _processor(transpose=True)(array).get_patterns().shape == (1, 5, 3) + assert _pipeline(TransposeStep())(array).get_patterns().shape == (1, 5, 3) -# ---------- Processor.process_bad_pixels (B2) ---------- +def test_pipeline_hflip_flips_last_axis() -> None: + patterns = numpy.arange(6, dtype=numpy.uint16).reshape(1, 2, 3) + array = SimpleDiffractionArray('a', numpy.zeros(1, dtype=int), patterns) + out = _pipeline(HorizontalFlipStep())(array).get_patterns() + assert numpy.array_equal(out[0], numpy.flip(patterns[0], axis=-1)) -def test_process_bad_pixels_requires_2d() -> None: +def test_pipeline_vflip_flips_second_to_last_axis() -> None: + patterns = numpy.arange(6, dtype=numpy.uint16).reshape(1, 2, 3) + array = SimpleDiffractionArray('a', numpy.zeros(1, dtype=int), patterns) + out = _pipeline(VerticalFlipStep())(array).get_patterns() + assert numpy.array_equal(out[0], numpy.flip(patterns[0], axis=-2)) + + +# ---------- Pipeline.apply_to_mask (B2) ---------- + + +def test_apply_to_mask_requires_2d() -> None: with pytest.raises(ValueError, match='Invalid bad_pixel dimensions'): - _processor().process_bad_pixels(numpy.zeros((1, 4, 4), dtype=bool)) + _pipeline().apply_to_mask(numpy.zeros((1, 4, 4), dtype=bool)) -def test_process_bad_pixels_transpose_does_not_crash() -> None: +def test_apply_to_mask_transpose_does_not_crash() -> None: """B2: transpose used axes=(0,2,1) on 2D, raising 'axes don't match array'.""" bad = numpy.zeros((3, 5), dtype=bool) bad[0, 4] = True - out = _processor(transpose=True).process_bad_pixels(bad) + out = _pipeline(TransposeStep()).apply_to_mask(bad) assert out.shape == (5, 3) assert out[4, 0] == True -def test_process_bad_pixels_padding_does_not_crash() -> None: +def test_apply_to_mask_padding_does_not_crash() -> None: """B1: padding flow on bad pixels must not raise.""" bad = numpy.ones((4, 4), dtype=bool) - out = _processor(padding=DiffractionPatternPadding(pad_x=1, pad_y=1)).process_bad_pixels(bad) + out = _pipeline(PaddingStep(pad_x=1, pad_y=1)).apply_to_mask(bad) assert out.shape == (6, 6) -def test_process_bad_pixels_full_pipeline() -> None: +def test_apply_to_mask_full_pipeline() -> None: bad = numpy.zeros((8, 8), dtype=bool) bad[4, 4] = True - out = _processor( - crop=DiffractionPatternCrop( - CropCenter(position_x_px=4, position_y_px=4), ImageExtent(4, 4) - ), - binning=DiffractionPatternBinning(bin_size_x=2, bin_size_y=2), - padding=DiffractionPatternPadding(pad_x=1, pad_y=1), - ).process_bad_pixels(bad) + out = _pipeline( + CropStep(center=CropCenter(position_x_px=4, position_y_px=4), extent=ImageExtent(4, 4)), + BinningStep(bin_size_x=2, bin_size_y=2), + PaddingStep(pad_x=1, pad_y=1), + ).apply_to_mask(bad) # 8x8 → crop to 4x4 (rows 2:6, cols 2:6, bad[4,4] inside) → bin 2x2 to 2x2 (logical AND so False) → pad to 4x4 assert out.shape == (4, 4) +# ---------- Step extent / pixel-geometry ---------- + + +def test_crop_apply_to_extent_returns_configured_extent() -> None: + step = CropStep(center=CropCenter(position_x_px=4, position_y_px=4), extent=ImageExtent(4, 6)) + out = step.apply_to_extent(ImageExtent(64, 64)) + assert (out.width_px, out.height_px) == (4, 6) + + +def test_binning_apply_to_extent_floor_divides() -> None: + step = BinningStep(bin_size_x=2, bin_size_y=4) + out = step.apply_to_extent(ImageExtent(9, 12)) + assert (out.width_px, out.height_px) == (4, 3) + + +def test_binning_apply_to_pixel_geometry_multiplies() -> None: + step = BinningStep(bin_size_x=2, bin_size_y=4) + out = step.apply_to_pixel_geometry(PixelGeometry(width_m=1e-5, height_m=2e-5)) + assert out.width_m == 2e-5 + assert out.height_m == 8e-5 + + +def test_padding_apply_to_extent_adds_double_pad() -> None: + step = PaddingStep(pad_x=1, pad_y=2) + out = step.apply_to_extent(ImageExtent(4, 4)) + assert (out.width_px, out.height_px) == (6, 8) + + +def test_transpose_apply_to_extent_swaps_dimensions() -> None: + step = TransposeStep() + out = step.apply_to_extent(ImageExtent(3, 5)) + assert (out.width_px, out.height_px) == (5, 3) + + +def test_transpose_apply_to_pixel_geometry_swaps() -> None: + step = TransposeStep() + out = step.apply_to_pixel_geometry(PixelGeometry(width_m=1e-5, height_m=2e-5)) + assert out.width_m == 2e-5 + assert out.height_m == 1e-5 + + +def test_identity_steps_do_not_change_extent_or_geometry() -> None: + extent = ImageExtent(4, 6) + geometry = PixelGeometry(width_m=1e-5, height_m=2e-5) + for step in ( + FilterValuesStep(lower_bound=1, upper_bound=99), + HorizontalFlipStep(), + VerticalFlipStep(), + ): + assert step.apply_to_extent(extent) == extent + assert step.apply_to_pixel_geometry(geometry) == geometry + + +def test_pipeline_compute_output_extent_composes_all_shape_steps() -> None: + pipeline = _pipeline( + CropStep(center=CropCenter(position_x_px=32, position_y_px=32), extent=ImageExtent(16, 16)), + BinningStep(bin_size_x=2, bin_size_y=2), + PaddingStep(pad_x=1, pad_y=1), + TransposeStep(), + ) + # 64x64 → crop 16x16 → bin 2x2 → 8x8 → pad → 10x10 → transpose → 10x10 + assert pipeline.compute_output_extent(ImageExtent(64, 64)) == ImageExtent(10, 10) + + +def test_pipeline_compute_output_pixel_geometry_composes_binning_and_transpose() -> None: + pipeline = _pipeline( + BinningStep(bin_size_x=2, bin_size_y=4), + TransposeStep(), + ) + out = pipeline.compute_output_pixel_geometry(PixelGeometry(width_m=1e-5, height_m=2e-5)) + # bin: width 2e-5, height 8e-5; transpose swaps → width 8e-5, height 2e-5 + assert out.width_m == 8e-5 + assert out.height_m == 2e-5 + + +# ---------- Serialization ---------- + + +def test_pipeline_serializes_and_round_trips() -> None: + """Pydantic tagged-union round-trip so `ptychodus_store` can persist a pipeline.""" + original = _pipeline( + FilterValuesStep(lower_bound=1, upper_bound=99), + CropStep(center=CropCenter(position_x_px=4, position_y_px=4), extent=ImageExtent(4, 4)), + BinningStep(bin_size_x=2, bin_size_y=2), + PaddingStep(pad_x=1, pad_y=1), + HorizontalFlipStep(), + VerticalFlipStep(), + TransposeStep(), + ) + round_tripped = DiffractionPrepPipeline.model_validate_json(original.model_dump_json()) + assert round_tripped == original + + # ---------- Sizer (B3, B4) ---------- @@ -223,8 +331,6 @@ def test_sizer_processed_size_accounts_for_double_sided_padding( ) -> None: """B4: padding is applied on both sides; processed size adds 2 * pad.""" diff, det = settings - det.width_px.set_value(64) - det.height_px.set_value(64) diff.crop_enabled.set_value(True) diff.crop_width_px.set_value(32) diff.crop_height_px.set_value(32) @@ -235,16 +341,48 @@ def test_sizer_processed_size_accounts_for_double_sided_padding( diff.pad_x.set_value(4) diff.pad_y.set_value(4) - sizer = PatternSizer(det, diff) - assert sizer.axis_x.get_processed_size() == 32 + 2 * 4 - assert sizer.axis_y.get_processed_size() == 32 + 2 * 4 + sizer = PatternSizer(diff) + detector_extent = ImageExtent(width_px=64, height_px=64) + extent = sizer.get_processed_image_extent(detector_extent) + assert extent.width_px == 32 + 2 * 4 + assert extent.height_px == 32 + 2 * 4 - # Cross-check against the processor's actual output shape. + # Cross-check against the pipeline's actual output shape. array = SimpleDiffractionArray( 'a', numpy.zeros(1, dtype=int), numpy.zeros((1, 64, 64), dtype=numpy.uint16) ) - out_shape = sizer.get_processor()(array).get_patterns().shape - assert out_shape == (1, sizer.axis_y.get_processed_size(), sizer.axis_x.get_processed_size()) + out_shape = sizer.get_prep_pipeline(detector_extent)(array).get_patterns().shape + assert out_shape == (1, extent.height_px, extent.width_px) + + +def test_sizer_processed_extent_reflects_transpose( + settings: tuple[DiffractionSettings, DetectorSettings], +) -> None: + """Regression: transpose must swap width/height in the processed extent + pixel geometry.""" + diff, _ = settings + diff.transpose.set_value(True) + + sizer = PatternSizer(diff) + detector_extent = ImageExtent(width_px=64, height_px=32) + extent = sizer.get_processed_image_extent(detector_extent) + assert (extent.width_px, extent.height_px) == (32, 64) + + geo = sizer.get_processed_pixel_geometry(PixelGeometry(width_m=1e-5, height_m=2e-5)) + assert (geo.width_m, geo.height_m) == (2e-5, 1e-5) + + # Cross-check: pipeline's actual output stack shape agrees with the sizer. + array = SimpleDiffractionArray( + 'a', numpy.zeros(1, dtype=int), numpy.zeros((1, 32, 64), dtype=numpy.uint16) + ) + out_shape = sizer.get_prep_pipeline(detector_extent)(array).get_patterns().shape + assert out_shape == (1, extent.height_px, extent.width_px) + + +def _filter_step(pipeline: DiffractionPrepPipeline) -> FilterValuesStep | None: + for step in pipeline.steps: + if isinstance(step, FilterValuesStep): + return step + return None def test_sizer_lower_bound_filter_uses_its_own_toggle( @@ -256,10 +394,11 @@ def test_sizer_lower_bound_filter_uses_its_own_toggle( diff.value_lower_bound.set_value(7) diff.value_upper_bound_enabled.set_value(False) - proc = PatternSizer(det, diff).get_processor() - assert proc.filter_values is not None - assert proc.filter_values.lower_bound == 7 - assert proc.filter_values.upper_bound is None + pipeline = PatternSizer(diff).get_prep_pipeline() + step = _filter_step(pipeline) + assert step is not None + assert step.lower_bound == 7 + assert step.upper_bound is None def test_sizer_upper_bound_filter_uses_its_own_toggle( @@ -270,10 +409,11 @@ def test_sizer_upper_bound_filter_uses_its_own_toggle( diff.value_upper_bound_enabled.set_value(True) diff.value_upper_bound.set_value(1234) - proc = PatternSizer(det, diff).get_processor() - assert proc.filter_values is not None - assert proc.filter_values.lower_bound is None - assert proc.filter_values.upper_bound == 1234 + pipeline = PatternSizer(diff).get_prep_pipeline() + step = _filter_step(pipeline) + assert step is not None + assert step.lower_bound is None + assert step.upper_bound == 1234 def test_sizer_both_filter_bounds_independent( @@ -285,18 +425,118 @@ def test_sizer_both_filter_bounds_independent( diff.value_upper_bound_enabled.set_value(True) diff.value_upper_bound.set_value(99) - proc = PatternSizer(det, diff).get_processor() - assert proc.filter_values is not None - assert proc.filter_values.lower_bound == 3 - assert proc.filter_values.upper_bound == 99 + pipeline = PatternSizer(diff).get_prep_pipeline() + step = _filter_step(pipeline) + assert step is not None + assert step.lower_bound == 3 + assert step.upper_bound == 99 -def test_sizer_no_filter_bounds(settings: tuple[DiffractionSettings, DetectorSettings]) -> None: +def test_sizer_no_filter_bounds( + settings: tuple[DiffractionSettings, DetectorSettings], +) -> None: + """Both filter toggles off → no FilterValuesStep in the pipeline (avoid a no-op step).""" diff, det = settings diff.value_lower_bound_enabled.set_value(False) diff.value_upper_bound_enabled.set_value(False) - proc = PatternSizer(det, diff).get_processor() - assert proc.filter_values is not None - assert proc.filter_values.lower_bound is None - assert proc.filter_values.upper_bound is None + pipeline = PatternSizer(diff).get_prep_pipeline() + assert _filter_step(pipeline) is None + + +# ---------- Safe crop center ---------- + + +def _crop_step(pipeline: DiffractionPrepPipeline) -> CropStep | None: + for step in pipeline.steps: + if isinstance(step, CropStep): + return step + return None + + +@pytest.mark.parametrize( + 'user_center, expected', + [ + (6, 6), # max valid center (previous code clamped to 5) + (2, 2), # min valid center + (0, 2), # below min → clamped up to radius + (100, 6), # above max → clamped down to det_size - radius + (4, 4), # in-range identity + ], +) +def test_sizer_safe_crop_center_matches_cropstep_bounds( + settings: tuple[DiffractionSettings, DetectorSettings], + user_center: int, + expected: int, +) -> None: + """The clamped center must equal the CropStep radius-based bounds (see CropStep.apply).""" + diff, det = settings + diff.crop_enabled.set_value(True) + diff.crop_width_px.set_value(4) + diff.crop_height_px.set_value(4) + diff.crop_center_x_px.set_value(user_center) + diff.crop_center_y_px.set_value(user_center) + + detector_extent = ImageExtent(width_px=8, height_px=8) + step = _crop_step(PatternSizer(diff).get_prep_pipeline(detector_extent)) + assert step is not None + assert step.center.position_x_px == expected + assert step.center.position_y_px == expected + + +def test_sizer_safe_crop_center_produces_in_bounds_slice( + settings: tuple[DiffractionSettings, DetectorSettings], +) -> None: + """Cross-check: a max-valid center must yield a slice fully inside the detector.""" + diff, det = settings + diff.crop_enabled.set_value(True) + diff.crop_width_px.set_value(4) + diff.crop_height_px.set_value(4) + diff.crop_center_x_px.set_value(6) # max valid; radius = 2 → slice [4:8] + diff.crop_center_y_px.set_value(6) + + array = SimpleDiffractionArray( + 'a', numpy.zeros(1, dtype=int), numpy.zeros((1, 8, 8), dtype=numpy.uint16) + ) + detector_extent = ImageExtent(width_px=8, height_px=8) + out = PatternSizer(diff).get_prep_pipeline(detector_extent)(array).get_patterns() + assert out.shape == (1, 4, 4) # no truncation from an out-of-bounds slice + + +# ---------- Observer notifications ---------- + + +class _CountingObserver: + def __init__(self) -> None: + self.n_updates = 0 + + def _update(self, observable: object) -> None: + self.n_updates += 1 + + +@pytest.mark.parametrize( + 'attr', + [ + 'hflip', + 'vflip', + 'transpose', + 'value_lower_bound_enabled', + 'value_lower_bound', + 'value_upper_bound_enabled', + 'value_upper_bound', + ], +) +def test_sizer_notifies_on_whole_image_parameter_change( + settings: tuple[DiffractionSettings, DetectorSettings], attr: str +) -> None: + """Whole-image settings (flips, transpose, filter bounds) must wake up sizer observers.""" + diff, det = settings + sizer = PatternSizer(diff) + observer = _CountingObserver() + sizer.add_observer(observer) # type: ignore[arg-type] + + parameter = getattr(diff, attr) + current = parameter.get_value() + parameter.set_value(current + 1 if isinstance(current, int) else not current) + + assert observer.n_updates >= 1 diff --git a/tests/test_diffraction_repository.py b/tests/test_diffraction_repository.py new file mode 100644 index 000000000..02b8bcff7 --- /dev/null +++ b/tests/test_diffraction_repository.py @@ -0,0 +1,166 @@ +"""Unit tests for DiffractionDatasetRepository.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from ptychodus.model.diffraction.repository import ( + DiffractionDatasetRepository, + DiffractionDatasetRepositoryObserver, +) + + +def _make_dataset(name: str) -> MagicMock: + dataset = MagicMock() + dataset.get_name.return_value = name + return dataset + + +class _RecordingObserver(DiffractionDatasetRepositoryObserver): + def __init__(self) -> None: + self.inserted: list[tuple[int, str]] = [] + self.removed: list[tuple[int, str]] = [] + + def handle_dataset_inserted(self, index, dataset) -> None: # noqa: ANN001 + self.inserted.append((index, dataset.get_name())) + + def handle_dataset_removed(self, index, dataset) -> None: # noqa: ANN001 + self.removed.append((index, dataset.get_name())) + + +def test_empty_repository_is_empty() -> None: + repo = DiffractionDatasetRepository() + assert len(repo) == 0 + + +def test_insert_dataset_appends_and_returns_index() -> None: + repo = DiffractionDatasetRepository() + a = _make_dataset('a') + b = _make_dataset('b') + + assert repo.insert_dataset(a) == 0 + assert repo.insert_dataset(b) == 1 + assert len(repo) == 2 + assert repo[0] is a + assert repo[1] is b + + +def test_insert_notifies_observers() -> None: + repo = DiffractionDatasetRepository() + observer = _RecordingObserver() + repo.add_observer(observer) + + a = _make_dataset('a') + b = _make_dataset('b') + repo.insert_dataset(a) + repo.insert_dataset(b) + + assert observer.inserted == [(0, 'a'), (1, 'b')] + + +def test_remove_dataset_pops_and_calls_clear() -> None: + repo = DiffractionDatasetRepository() + a = _make_dataset('a') + b = _make_dataset('b') + repo.insert_dataset(a) + repo.insert_dataset(b) + + repo.remove_dataset(0) + + assert len(repo) == 1 + assert repo[0] is b + a.clear.assert_called_once() + + +def test_remove_notifies_observers() -> None: + repo = DiffractionDatasetRepository() + observer = _RecordingObserver() + repo.add_observer(observer) + + a = _make_dataset('a') + repo.insert_dataset(a) + repo.remove_dataset(0) + + assert observer.removed == [(0, 'a')] + + +def test_remove_out_of_range_is_noop() -> None: + repo = DiffractionDatasetRepository() + repo.insert_dataset(_make_dataset('a')) + repo.remove_dataset(5) + assert len(repo) == 1 + + +def test_clear_removes_all_and_fires_per_remove_events() -> None: + repo = DiffractionDatasetRepository() + observer = _RecordingObserver() + repo.add_observer(observer) + + repo.insert_dataset(_make_dataset('a')) + repo.insert_dataset(_make_dataset('b')) + repo.insert_dataset(_make_dataset('c')) + + repo.clear() + + assert len(repo) == 0 + # Removed bottom-up so the caller sees stable indexes during iteration. + assert observer.removed == [(2, 'c'), (1, 'b'), (0, 'a')] + + +def test_create_unique_name_returns_input_when_free() -> None: + repo = DiffractionDatasetRepository() + assert repo.create_unique_name('foo') == 'foo' + + +def test_create_unique_name_suffixes_collisions() -> None: + repo = DiffractionDatasetRepository() + repo.insert_dataset(_make_dataset('foo')) + repo.insert_dataset(_make_dataset('foo-1')) + assert repo.create_unique_name('foo') == 'foo-2' + + +def test_create_unique_name_maps_empty_to_unnamed() -> None: + repo = DiffractionDatasetRepository() + assert repo.create_unique_name('') == 'Unnamed' + + +def test_remove_observer_stops_notifications() -> None: + repo = DiffractionDatasetRepository() + observer = _RecordingObserver() + repo.add_observer(observer) + repo.remove_observer(observer) + + repo.insert_dataset(_make_dataset('a')) + assert observer.inserted == [] + + +def test_getitem_slice_returns_sequence() -> None: + repo = DiffractionDatasetRepository() + a = _make_dataset('a') + b = _make_dataset('b') + repo.insert_dataset(a) + repo.insert_dataset(b) + + tail = repo[1:] + assert list(tail) == [b] + + +def test_create_dataset_without_factory_raises() -> None: + repo = DiffractionDatasetRepository() + import pytest + + with pytest.raises(RuntimeError): + repo.create_dataset('foo') + + +def test_create_dataset_with_factory_uses_unique_name() -> None: + called_with: list[str] = [] + + def _factory(name: str) -> MagicMock: + called_with.append(name) + return _make_dataset(name) + + repo = DiffractionDatasetRepository(factory=_factory) + repo.insert_dataset(_make_dataset('foo')) + repo.create_dataset('foo') + assert called_with == ['foo-1'] diff --git a/tests/test_diffraction_two_datasets.py b/tests/test_diffraction_two_datasets.py new file mode 100644 index 000000000..bc2490b56 --- /dev/null +++ b/tests/test_diffraction_two_datasets.py @@ -0,0 +1,229 @@ +"""Two-dataset integration test for the diffraction model. + +Exercises the DiffractionDatasetRepository together with real +AssembledDiffractionDataset instances (built via the repository's factory) +to verify per-dataset bad-pixels ownership and stable index-based routing. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import numpy +import pytest + +from ptychodus.api.diffraction import DiffractionMetadata, SimpleDiffractionDataset +from ptychodus.api.geometry import ImageExtent, PixelGeometry +from ptychodus.api.io import AssembledDiffractionData +from ptychodus.api.settings import SettingsRegistry +from ptychodus.api.tree import SimpleTreeNode +from ptychodus.model.diffraction.dataset import ( + AssembledDiffractionArray, + AssembledDiffractionDataset, +) +from ptychodus.model.diffraction.repository import DiffractionDatasetRepository +from ptychodus.model.diffraction.settings import DetectorSettings, DiffractionSettings +from ptychodus.model.diffraction.sizer import PatternSizer + + +def _make_repository() -> DiffractionDatasetRepository: + registry = SettingsRegistry() + detector_settings = DetectorSettings(registry) + diffraction_settings = DiffractionSettings(registry) + sizer = PatternSizer(diffraction_settings) + task_manager = MagicMock() + task_monitor = MagicMock() + + def _factory(name: str) -> AssembledDiffractionDataset: + return AssembledDiffractionDataset( + diffraction_settings, + sizer, + detector_settings, + task_manager, + task_monitor, + name=name, + ) + + return DiffractionDatasetRepository(factory=_factory) + + +def _reload_with_extent(dataset: AssembledDiffractionDataset, extent: ImageExtent) -> None: + metadata = DiffractionMetadata( + num_patterns_per_array=[0], + pattern_dtype=numpy.dtype(numpy.uint16), + detector_extent=extent, + ) + contents_tree = SimpleTreeNode.create_root(['Name', 'Type', 'Details']) + source = SimpleDiffractionDataset(metadata, contents_tree, []) + dataset.reload(source) + + +def test_two_datasets_end_up_at_stable_indexes() -> None: + repo = _make_repository() + + a = repo.create_dataset('scan_a') + index_a = repo.insert_dataset(a) + + b = repo.create_dataset('scan_b') + index_b = repo.insert_dataset(b) + + assert index_a == 0 + assert index_b == 1 + assert repo[0].get_name() == 'scan_a' + assert repo[1].get_name() == 'scan_b' + + +def test_bad_pixels_are_per_dataset() -> None: + repo = _make_repository() + + a = repo.create_dataset('scan_a') + repo.insert_dataset(a) + _reload_with_extent(a, ImageExtent(width_px=16, height_px=16)) + b = repo.create_dataset('scan_b') + repo.insert_dataset(b) + _reload_with_extent(b, ImageExtent(width_px=8, height_px=8)) + + # Each dataset's default mask is sized to *its own* detector extent — regression + # against the earlier singleton-source behavior where both would have shared shape. + assert repo[0].get_bad_pixels().shape == (16, 16) + assert repo[1].get_bad_pixels().shape == (8, 8) + assert not repo[0].get_bad_pixels().any() + assert not repo[1].get_bad_pixels().any() + + # Set a distinctive mask on dataset A only. + custom_mask_a = numpy.zeros_like(repo[0].get_bad_pixels()) + custom_mask_a[0, 0] = True + repo[0].set_bad_pixels(custom_mask_a) + + # A picks it up; B is untouched. + assert repo[0].get_bad_pixels()[0, 0] + assert not repo[1].get_bad_pixels()[0, 0] + + # And they are not aliased. + assert repo[0].get_bad_pixels() is not repo[1].get_bad_pixels() + + +def test_removing_first_dataset_shifts_indexes() -> None: + repo = _make_repository() + + repo.insert_dataset(repo.create_dataset('scan_a')) + b = repo.create_dataset('scan_b') + repo.insert_dataset(b) + + repo.remove_dataset(0) + + assert len(repo) == 1 + assert repo[0] is b + assert repo[0].get_name() == 'scan_b' + + +def test_reset_bad_pixels_restores_default() -> None: + repo = _make_repository() + a = repo.create_dataset('scan_a') + repo.insert_dataset(a) + _reload_with_extent(a, ImageExtent(width_px=16, height_px=16)) + + mask = numpy.zeros_like(a.get_bad_pixels()) + mask[1, 1] = True + a.set_bad_pixels(mask) + assert a.get_bad_pixels()[1, 1] + + a.reset_bad_pixels() + assert a.get_bad_pixels().shape == (16, 16) + assert not a.get_bad_pixels().any() + + +def test_set_bad_pixels_rejects_shape_mismatch_after_reload() -> None: + repo = _make_repository() + a = repo.create_dataset('scan_a') + repo.insert_dataset(a) + _reload_with_extent(a, ImageExtent(width_px=16, height_px=16)) + + wrong_shape = numpy.zeros((8, 8), dtype=numpy.bool_) + with pytest.raises(ValueError, match='does not match loaded detector extent'): + a.set_bad_pixels(wrong_shape) + + # The original default mask is untouched. + assert a.get_bad_pixels().shape == (16, 16) + + +def test_set_bad_pixels_rejects_shape_mismatch_before_reload() -> None: + """Pre-reload, the metadata extent is (0, 0); any non-empty mask is rejected.""" + repo = _make_repository() + a = repo.create_dataset('scan_a') + repo.insert_dataset(a) + + mask = numpy.zeros((32, 64), dtype=numpy.bool_) + with pytest.raises(ValueError, match='does not match loaded detector extent'): + a.set_bad_pixels(mask) + + +def test_simple_diffraction_dataset_rejects_bad_pixels_shape_mismatch() -> None: + metadata = DiffractionMetadata( + num_patterns_per_array=[0], + pattern_dtype=numpy.dtype(numpy.uint16), + detector_extent=ImageExtent(width_px=16, height_px=16), + ) + contents_tree = SimpleTreeNode.create_root(['Name', 'Type', 'Details']) + wrong_shape = numpy.zeros((8, 8), dtype=numpy.bool_) + + with pytest.raises(ValueError, match='does not match detector extent'): + SimpleDiffractionDataset(metadata, contents_tree, [], wrong_shape) + + +def test_simple_diffraction_dataset_default_bad_pixels_matches_extent() -> None: + metadata = DiffractionMetadata( + num_patterns_per_array=[0], + pattern_dtype=numpy.dtype(numpy.uint16), + detector_extent=ImageExtent(width_px=16, height_px=8), + ) + contents_tree = SimpleTreeNode.create_root(['Name', 'Type', 'Details']) + dataset = SimpleDiffractionDataset(metadata, contents_tree, []) + assert dataset.get_bad_pixels().shape == (8, 16) + assert not dataset.get_bad_pixels().any() + + +def test_create_unique_name_prevents_collision_after_insert() -> None: + repo = _make_repository() + repo.insert_dataset(repo.create_dataset('scan')) + repo.insert_dataset(repo.create_dataset('scan')) + assert [ds.get_name() for ds in repo] == ['scan', 'scan-1'] + + +def _make_array( + array_index: int, + label: str, + fill_value: float, + num_patterns: int, + detector_shape: tuple[int, int], + index_offset: int, +) -> AssembledDiffractionArray: + height, width = detector_shape + patterns = numpy.full((num_patterns, height, width), fill_value, dtype=numpy.float64) + indexes = numpy.arange(index_offset, index_offset + num_patterns, dtype=numpy.intp) + data = AssembledDiffractionData( + indexes=indexes, + patterns=patterns, + pixel_geometry=PixelGeometry(width_m=1.0, height_m=1.0), + bad_pixels=numpy.zeros(detector_shape, dtype=numpy.bool_), + ) + return AssembledDiffractionArray(array_index=array_index, label=label, data=data) + + +def test_dataset_average_pattern_is_weighted_mean_across_arrays() -> None: + """Selecting a dataset node previews the average across all its patterns.""" + repo = _make_repository() + dataset = repo.create_dataset('scan') + repo.insert_dataset(dataset) + + # No arrays yet -> no preview to show. + assert dataset.get_average_pattern() is None + + # Two arrays of different sizes with distinct uniform fills — the correct + # weighted mean is (3*2 + 7*6) / (3 + 7) = 4.8 everywhere. + dataset._insert_array(_make_array(0, 'a', 2.0, 3, (4, 4), index_offset=0)) + dataset._insert_array(_make_array(1, 'b', 6.0, 7, (4, 4), index_offset=3)) + + result = dataset.get_average_pattern() + assert result is not None + numpy.testing.assert_allclose(result, numpy.full((4, 4), 4.8)) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 31c683859..e31f8e291 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -670,6 +670,25 @@ def test_pixel_geometry_equality() -> None: assert PixelGeometry(1e-6, 1e-6) != PixelGeometry(1e-6, 2e-6) +def test_pixel_geometry_is_valid_true_for_strictly_positive() -> None: + assert PixelGeometry(width_m=1e-6, height_m=1e-6).is_valid + + +def test_pixel_geometry_is_valid_false_when_either_dimension_is_zero() -> None: + """Zero on either axis is the 'not ready' sentinel returned by ProductGeometry + before a dataset binds — must be rejected.""" + assert not PixelGeometry(width_m=0.0, height_m=1e-6).is_valid + assert not PixelGeometry(width_m=1e-6, height_m=0.0).is_valid + assert not PixelGeometry(width_m=0.0, height_m=0.0).is_valid + + +def test_pixel_geometry_is_valid_false_for_negative_dimensions() -> None: + """Negative pixel sizes are physically meaningless — the predicate treats + them the same as zero.""" + assert not PixelGeometry(width_m=-1e-6, height_m=1e-6).is_valid + assert not PixelGeometry(width_m=1e-6, height_m=-1e-6).is_valid + + # --------------------------------------------------------------------------- # ImageExtent # --------------------------------------------------------------------------- diff --git a/tests/test_io.py b/tests/test_io.py index 4cafa378f..96aadba0a 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -2,14 +2,18 @@ from __future__ import annotations +from dataclasses import replace from pathlib import Path +import h5py import numpy import numpy.testing import pytest +from ptychodus.api.diffraction import Polarization from ptychodus.api.geometry import PixelGeometry from ptychodus.api.io import ( + ProductFileKeys, StandardFileLayout, load_diffraction_data, load_product, @@ -280,6 +284,9 @@ def _assert_metadata_equal(self, a: ProductMetadata, b: ProductMetadata) -> None assert a.probe_photon_count == pytest.approx(b.probe_photon_count) assert a.exposure_time_s == pytest.approx(b.exposure_time_s) assert a.mass_attenuation_m2_kg == pytest.approx(b.mass_attenuation_m2_kg) + assert a.tomography_angle_deg == pytest.approx(b.tomography_angle_deg) + assert a.tilt_angle_deg == pytest.approx(b.tilt_angle_deg) + assert a.polarization == b.polarization def test_basic_round_trip(self, tmp_path: Path) -> None: original = _make_product() @@ -421,3 +428,66 @@ def test_metadata_optional_fields_default(self, tmp_path: Path) -> None: loaded = load_product(file) assert loaded.metadata.name == 'Unnamed' assert loaded.metadata.comments == '' + + def test_tomography_angle_round_trip(self, tmp_path: Path) -> None: + original = _make_product() + original = Product( + metadata=replace(original.metadata, tomography_angle_deg=42.5), + probe_positions=original.probe_positions, + probes=original.probes, + object_=original.object_, + losses=original.losses, + ) + file = tmp_path / 'product.h5' + + save_product(file, original) + loaded = load_product(file) + + assert loaded.metadata.tomography_angle_deg == pytest.approx(42.5) + + def test_tilt_and_polarization_round_trip(self, tmp_path: Path) -> None: + original = _make_product() + original = Product( + metadata=replace( + original.metadata, + tilt_angle_deg=12.5, + polarization=Polarization.LEFT_CIRCULAR, + ), + probe_positions=original.probe_positions, + probes=original.probes, + object_=original.object_, + losses=original.losses, + ) + file = tmp_path / 'product.h5' + + save_product(file, original) + loaded = load_product(file) + + assert loaded.metadata.tilt_angle_deg == pytest.approx(12.5) + assert loaded.metadata.polarization is Polarization.LEFT_CIRCULAR + + def test_polarization_absent_reads_none(self, tmp_path: Path) -> None: + original = _make_product() + file = tmp_path / 'product.h5' + + save_product(file, original) + loaded = load_product(file) + + assert loaded.metadata.polarization is None + assert loaded.metadata.tilt_angle_deg == pytest.approx(0.0) + + def test_polarization_invalid_string_falls_back_to_none( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + original = _make_product() + file = tmp_path / 'product.h5' + save_product(file, original) + + with h5py.File(file, 'a') as f: + f.attrs[ProductFileKeys.POLARIZATION] = 'bogus_value' + + with caplog.at_level('WARNING'): + loaded = load_product(file) + + assert loaded.metadata.polarization is None + assert 'Unknown polarization' in caplog.text diff --git a/tests/test_no_gpu_context.py b/tests/test_no_gpu_context.py new file mode 100644 index 000000000..a09e095fb --- /dev/null +++ b/tests/test_no_gpu_context.py @@ -0,0 +1,203 @@ +"""Discipline gate for the GPU subprocess isolation refactor. + +The invariant the refactor buys us is that the parent ptychodus process never +acquires a GPU context. Acquiring one is what pins driver state and GPU memory +to the long-lived parent, and it is what stops ptychodus from mixing backends +built on different frameworks. Context acquisition means things like +constructing ``ptychi.api.task.PtychographyTask``, placing a tensor on a +device, ``tf.keras.Model.fit``, constructing a Lightning ``Trainer``, or +allocating a CuPy array -- all of which must happen inside a freshly spawned +child that dies at end-of-call. + +*Importing* a GPU framework in the parent is fine and expected. ``ptychi.api`` +pulls torch in for its type annotations so the parent can build the picklable +``PtychographyTaskOptions`` the child needs, and no CUDA runtime is +initialised by that. See the invariant note at the top of +``ptychodus.model.processing._subprocess_protocol`` for the canonical +statement. + +So this module checks two things, in one spawned probe process that imports +``ptychodus`` plus the main composition roots: + +1. No module that would acquire a context merely by being imported (or that + simply has no parent-side use) ended up in ``sys.modules``. +2. No framework that *is* allowed parent-side is holding a live GPU context. + +Runs in a spawned subprocess so it is not tainted by whatever the pytest +worker previously loaded (e.g. earlier tests that touched a GPU-backed +plugin). +""" + +from __future__ import annotations + +from typing import Any +import multiprocessing +import sys +import textwrap + +import pytest + +# Frameworks the parent is allowed to import -- it needs them to build the +# picklable options/payload objects it hands to the child -- but which must +# never be holding a GPU context parent-side. Enforced by +# ``_gpu_context_offenders`` rather than by an import ban. +CONTEXT_CAPABLE_MODULES = ( + 'torch', + 'tensorflow', + 'ptychi', # covers ptychi.api etc. + 'lightning', + 'pytorch_lightning', +) + +# Modules banned from the parent AT IMPORT TIME: importing them acquires a +# context, or they are child-side backend packages the composition roots have +# no reason to touch. +# +# ``ptycho`` and ``ptycho_torch`` are a softer case than ``cupy``. Their config +# subpackages -- ``ptycho.config.config`` and ``ptycho_torch.config_params`` -- +# are deliberately reachable from the parent AT CALL TIME, because the parent +# builds the backend config objects it ships to the child (see +# ``ptychopinn/_payload.py`` and ``ptychopinn_torch/_payload.py``). Neither +# acquires a GPU context: the former pulls no heavy modules at all, the latter +# pulls torch, which is allowed above. Both imports live inside the payload +# builders, so they do not run until the first reconstruct/train call and this +# import-time probe never sees them. Keep it that way -- hoisting either to +# module scope in a ``reconstructor.py`` is what this list is here to catch. +CHILD_ONLY_MODULES = ( + 'cupy', # links and initialises the CUDA runtime at import; no non-invasive probe + 'ptycho', # ptychopinn TensorFlow package; raw_data/probe/tf_helper pull TensorFlow + 'ptycho_torch', # ptychopinn_torch backend; child-side entry point only +) + +# ``ptychozoon.data_structures`` and ``ptychozoon.settings`` are CPU-only -- +# they pull in only numpy, dataclasses, enum, and typing. The parent-side +# fluorescence factory imports them so it can construct the ptychozoon +# payload directly instead of duplicating its fields. Any *other* +# ``ptychozoon.*`` module (notably ``vspi_enhance``) pulls in CuPy and is a +# regression if it lands in the parent's ``sys.modules``. +ALLOWED_PTYCHOZOON_MODULES = frozenset( + { + 'ptychozoon', + 'ptychozoon.data_structures', + 'ptychozoon.settings', + } +) + + +def _is_child_only(module_name: str) -> bool: + if module_name == 'ptychozoon' or module_name.startswith('ptychozoon.'): + return module_name not in ALLOWED_PTYCHOZOON_MODULES + return any(module_name == m or module_name.startswith(m + '.') for m in CHILD_ONLY_MODULES) + + +def _gpu_context_offenders() -> tuple[list[str], list[str]]: + """Report any live GPU context held by an already-imported framework. + + Returns ``(offenders, notes)``. Every probe below only *observes* state -- + none of them create a context as a side effect, so calling this cannot + itself break the invariant. A probe that fails (framework internals moved + between versions) lands in ``notes`` instead of silently passing. + """ + offenders: list[str] = [] + notes: list[str] = [] + + torch = sys.modules.get('torch') + if torch is not None: + for backend_name in ('cuda', 'xpu'): + backend = getattr(torch, backend_name, None) + is_initialized = getattr(backend, 'is_initialized', None) + if is_initialized is None: + continue # e.g. torch.xpu is absent on older torch builds + try: + if is_initialized(): + offenders.append(f'torch.{backend_name} context is initialized') + except Exception as exc: # noqa: BLE001 + notes.append(f'could not probe torch.{backend_name}: {type(exc).__name__}: {exc}') + + if 'tensorflow' in sys.modules: + try: + from tensorflow.python.eager import context as tf_context + + # context_safe() returns the eager context, or None when none has + # been created -- unlike context(), it does not create one. + if tf_context.context_safe() is not None: + offenders.append('tensorflow eager context is initialized') + except Exception as exc: # noqa: BLE001 + notes.append(f'could not probe tensorflow eager context: {type(exc).__name__}: {exc}') + + return offenders, notes + + +def _check_in_subprocess(result_queue: multiprocessing.Queue[dict[str, list[str]] | str]) -> None: + try: + # Import the top-level package plus the main composition roots. + import ptychodus # noqa: F401 + import ptychodus.model.core # noqa: F401 + import ptychodus.model.processing.subprocess_reconstructor # noqa: F401 + import ptychodus.model.processing._subprocess_protocol # noqa: F401 + + context_offenders, notes = _gpu_context_offenders() + result_queue.put( + { + 'child_only': sorted(m for m in sys.modules if _is_child_only(m)), + 'context': context_offenders, + 'notes': notes, + } + ) + except BaseException as exc: # noqa: BLE001 + result_queue.put(f'{type(exc).__name__}: {exc}') + + +@pytest.fixture(scope='module') +def probe_result() -> dict[str, list[str]]: + """Run the import/context probe once and share it across both tests.""" + ctx = multiprocessing.get_context('spawn') + result_queue: multiprocessing.Queue[dict[str, list[str]] | str] = ctx.Queue() + process = ctx.Process(target=_check_in_subprocess, args=(result_queue,)) + process.start() + process.join(timeout=60.0) + + assert not process.is_alive(), 'GPU-isolation probe subprocess did not exit in 60s.' + + result: Any = result_queue.get(timeout=1.0) + + if isinstance(result, str): + raise AssertionError(f'Probe subprocess raised: {result}') + + return result + + +def test_parent_imports_no_child_only_modules(probe_result: dict[str, list[str]]) -> None: + offenders = probe_result['child_only'] + assert offenders == [], textwrap.dedent( + f"""\ + Parent ptychodus process imported modules that belong to a child only. + Offending sys.modules entries: {offenders} + + These modules acquire a GPU context just by being imported, or are + backend packages the parent has no reason to touch. Importing a GPU + framework parent-side to build a picklable payload is fine -- see + CONTEXT_CAPABLE_MODULES -- but these are not in that category. Wire the + backend through SubprocessReconstructor with a child-side entry point + instead. + """ + ) + + +def test_parent_holds_no_gpu_context(probe_result: dict[str, list[str]]) -> None: + offenders = probe_result['context'] + notes = probe_result['notes'] + assert offenders == [], textwrap.dedent( + f"""\ + Parent ptychodus process is holding a live GPU context. + Offenders: {offenders} + Probe notes: {notes or 'none'} + + Importing torch/tensorflow/ptychi in the parent is allowed; acquiring a + context is not. Something at import time placed a tensor on a device, + constructed a PtychographyTask/Trainer, or otherwise initialised the + runtime. Move that work into a child entry point dispatched through + SubprocessReconstructor (see + ptychodus.model.processing._subprocess_protocol). + """ + ) diff --git a/tests/test_object_builder.py b/tests/test_object_builder.py new file mode 100644 index 000000000..bbbe26e9d --- /dev/null +++ b/tests/test_object_builder.py @@ -0,0 +1,257 @@ +"""Regression tests for the object conditioning pipeline (extra padding -> layers). + +Two invariants carry the weight here. + +First, the two operations sit on opposite sides of a split. generate_layers is +conditioning: it applies to generated and file-loaded objects alike, guarded so +it never destroys layers the input already has, because it truncates when asked +for fewer than it is given. pad_object is generation-only: it is strictly +additive and leaves no trace in the array, so there is no way to detect an +already-padded object and skip it. Applying it to a file-loaded object would grow +that object on every load/reconstruct/save/load round trip, unbounded -- and the +padding defaults to 1, so it would happen to users who never touched the setting. + +Second, FromMemoryObjectBuilder must never condition. It holds an object that is +already conditioned -- reconstruction output, which ProcessingTaskMonitor +re-assigns to the output product item on every reconstructor iteration, and +products loaded from HDF5/NPZ. Padding there would grow the array once per +iteration, and generate_layers would collapse a converged multislice result back +to whatever the item's layer spacing happens to say. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path + +import numpy +import pytest + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.object import ( + Object, + ObjectCenter, + ObjectFileReader, + ObjectGeometry, + ObjectGeometryProvider, +) +from ptychodus.api.probe_positions import ProbePosition +from ptychodus.api.settings import SettingsRegistry +from ptychodus.model.product.object.builder import ( + FromFileObjectBuilder, + FromMemoryObjectBuilder, + ObjectBuilder, +) +from ptychodus.model.product.object.item import ObjectRepositoryItem +from ptychodus.model.product.object.random import RandomObjectBuilder +from ptychodus.model.product.object.settings import ObjectSettings + +EXTENT_PX = 8 +PIXEL_SIZE_M = 1.0e-6 +LAYER_SPACING_M = 1.0e-6 + + +def _make_settings() -> ObjectSettings: + return ObjectSettings(SettingsRegistry()) + + +def _make_rng() -> numpy.random.Generator: + return numpy.random.default_rng(42) + + +class _StubObjectGeometryProvider(ObjectGeometryProvider): + def get_probe_positions(self) -> Sequence[ProbePosition]: + return () + + def get_object_geometry(self) -> ObjectGeometry: + return ObjectGeometry( + width_px=EXTENT_PX, + height_px=EXTENT_PX, + pixel_width_m=PIXEL_SIZE_M, + pixel_height_m=PIXEL_SIZE_M, + center_x_m=0.0, + center_y_m=0.0, + ) + + +def _make_object(num_layers: int) -> Object: + """A deterministic, non-degenerate object with the requested layer count.""" + rng = numpy.random.default_rng(7) + shape = (num_layers, EXTENT_PX, EXTENT_PX) + # Keep the amplitude away from zero so the phase unwrapping in + # generate_layers is well conditioned. + array = (1.0 + 0.1 * rng.normal(size=shape)) * numpy.exp(1j * 0.1 * rng.normal(size=shape)) + return Object( + array=array.astype(complex), + pixel_geometry=PixelGeometry(width_m=PIXEL_SIZE_M, height_m=PIXEL_SIZE_M), + center=ObjectCenter(coordinate_x_m=0.0, coordinate_y_m=0.0), + layer_spacing_m=[LAYER_SPACING_M] * (num_layers - 1), + ) + + +class _StubObjectFileReader(ObjectFileReader): + def __init__(self, object_: Object) -> None: + self._object = object_ + + def read(self, file_path: Path) -> Object: + return self._object + + +def _make_from_file_builder(settings: ObjectSettings, object_: Object) -> FromFileObjectBuilder: + return FromFileObjectBuilder(settings, _StubObjectFileReader(object_)) + + +def test_generator_pads_canvas() -> None: + """The padding moved from the shared pipeline into _build_raw; the generative + path must still honor it.""" + settings = _make_settings() + builder = RandomObjectBuilder(_make_rng(), settings) + builder.extra_padding_x.set_value(3) + builder.extra_padding_y.set_value(2) + + object_ = builder.build(_StubObjectGeometryProvider(), []) + + assert object_.width_px == EXTENT_PX + 6 + assert object_.height_px == EXTENT_PX + 4 + + +def test_generator_generates_layers() -> None: + settings = _make_settings() + builder = RandomObjectBuilder(_make_rng(), settings) + + object_ = builder.build(_StubObjectGeometryProvider(), [LAYER_SPACING_M] * 2) + + assert object_.num_layers == 3 + assert list(object_.layer_spacing_m) == [LAYER_SPACING_M] * 2 + + +def test_from_file_builder_generates_layers() -> None: + """FromFileObjectBuilder.build() used to return the reader's output verbatim, + so the layer spacing was silently ignored for every file-loaded object.""" + settings = _make_settings() + builder = _make_from_file_builder(settings, _make_object(1)) + + object_ = builder.build(_StubObjectGeometryProvider(), [LAYER_SPACING_M] * 2) + + assert object_.num_layers == 3 + assert list(object_.layer_spacing_m) == [LAYER_SPACING_M] * 2 + + +def test_from_file_builder_keeps_existing_layers(caplog: pytest.LogCaptureFixture) -> None: + """The destructive case: generate_layers truncates when asked for fewer layers + than it is given, so the default empty spacing would collapse a converged + four-layer warm start to one layer.""" + settings = _make_settings() + from_file = _make_object(4) + builder = _make_from_file_builder(settings, from_file) + + with caplog.at_level('INFO'): + object_ = builder.build(_StubObjectGeometryProvider(), []) + + assert object_.num_layers == 4 + assert 'keeping them rather than re-slicing to 1' in caplog.text + + +def test_from_file_builder_does_not_pad() -> None: + """Padding is generation-only. It is strictly additive and undetectable after + the fact, so applying it here would grow a warm-start object by twice the + padding on every load/save round trip -- with the default padding of 1, for + users who never touched the setting.""" + settings = _make_settings() + from_file = _make_object(1) + builder = _make_from_file_builder(settings, from_file) + + assert builder.extra_padding_x.get_value() == 1 + assert builder.extra_padding_y.get_value() == 1 + + object_ = builder.build(_StubObjectGeometryProvider(), []) + + assert object_.width_px == from_file.width_px + assert object_.height_px == from_file.height_px + + +def test_from_file_conditioning_is_idempotent() -> None: + """Several rebuild paths -- a geometry-provider notification, a + builder-parameter edit -- can re-run build() on an already-conditioned + object.""" + settings = _make_settings() + provider = _StubObjectGeometryProvider() + layer_spacing_m = [LAYER_SPACING_M] * 2 + + conditioned = _make_from_file_builder(settings, _make_object(1)).build( + provider, layer_spacing_m + ) + reconditioned = _make_from_file_builder(settings, conditioned).build(provider, layer_spacing_m) + + assert numpy.array_equal(reconditioned.get_array(), conditioned.get_array()) + assert list(reconditioned.layer_spacing_m) == list(conditioned.layer_spacing_m) + + +def test_from_memory_builder_ignores_conditioning() -> None: + """Guards reconstruction output: the from-memory builder must return its + object verbatim no matter what the conditioning parameters say.""" + settings = _make_settings() + raw = _make_object(3) + builder = FromMemoryObjectBuilder(settings, raw) + builder.extra_padding_x.set_value(5) + builder.extra_padding_y.set_value(5) + + object_ = builder.build(_StubObjectGeometryProvider(), []) + + assert numpy.array_equal(object_.get_array(), raw.get_array()) + assert list(object_.layer_spacing_m) == list(raw.layer_spacing_m) + + +def test_repeated_from_memory_builds_are_idempotent() -> None: + """The reconstruct loop rebuilds the output item's object once per iteration. + pad_object is strictly additive, so conditioning here would grow the array + without bound.""" + settings = _make_settings() + settings.extra_padding_x.set_value(4) + settings.extra_padding_y.set_value(4) + + provider = _StubObjectGeometryProvider() + expected = _make_object(3) + object_ = expected + + for _ in range(3): + builder = FromMemoryObjectBuilder(settings, object_) + object_ = builder.build(provider, []) + + assert object_.get_array().shape == expected.get_array().shape + assert numpy.array_equal(object_.get_array(), expected.get_array()) + assert list(object_.layer_spacing_m) == list(expected.layer_spacing_m) + + +def test_item_adopts_layer_spacing_actually_produced() -> None: + """ObjectRepositoryItem.rebuild overwrites its own layer_spacing_m parameter + from whatever the builder returned, so the parameter tracks the layers the + object really has rather than the ones that were asked for.""" + settings = _make_settings() + builder = _make_from_file_builder(settings, _make_object(4)) + + item = ObjectRepositoryItem(_StubObjectGeometryProvider(), settings, builder) + + assert item.get_object().num_layers == 4 + assert list(item.layer_spacing_m.get_value()) == [LAYER_SPACING_M] * 3 + + +@pytest.mark.parametrize('builder_name', ['random', 'from_file']) +def test_copy_preserves_padding_parameters(builder_name: str) -> None: + """copy() iterates parameters() generically, so the padding rides along + without any per-subclass change.""" + settings = _make_settings() + builder: ObjectBuilder + + if builder_name == 'random': + builder = RandomObjectBuilder(_make_rng(), settings) + else: + builder = _make_from_file_builder(settings, _make_object(1)) + + builder.extra_padding_x.set_value(2) + builder.extra_padding_y.set_value(3) + + duplicate = builder.copy() + + assert duplicate.extra_padding_x.get_value() == 2 + assert duplicate.extra_padding_y.get_value() == 3 diff --git a/tests/test_object_item.py b/tests/test_object_item.py new file mode 100644 index 000000000..858a2fb13 --- /dev/null +++ b/tests/test_object_item.py @@ -0,0 +1,104 @@ +"""Regression tests for ObjectRepositoryItem's rebuild guard. + +ObjectRepositoryItem must not build an Object while the geometry provider +still reports zero-valued pixel dimensions, and must rebuild once the provider +notifies that real dimensions have arrived. See CLAUDE fly001.ini bug report. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.object import Object, ObjectCenter, ObjectGeometry, ObjectGeometryProvider +from ptychodus.api.observer import Observable +from ptychodus.api.probe_positions import ProbePosition +from ptychodus.api.settings import SettingsRegistry +from ptychodus.model.product.object.builder import ObjectBuilder +from ptychodus.model.product.object.item import ObjectRepositoryItem +from ptychodus.model.product.object.settings import ObjectSettings + + +class _ObservableObjectProvider(ObjectGeometryProvider, Observable): + """Test double: an Observable + ObjectGeometryProvider. set_geometry() + mutates the returned geometry and fires notify_observers, mimicking what + ProductGeometry.set_detector_extent does in production.""" + + def __init__(self, geometry: ObjectGeometry) -> None: + Observable.__init__(self) + self._geometry = geometry + + def set_geometry(self, geometry: ObjectGeometry) -> None: + self._geometry = geometry + self.notify_observers() + + def get_probe_positions(self) -> Sequence[ProbePosition]: + return () + + def get_object_geometry(self) -> ObjectGeometry: + return self._geometry + + +class _RecordingObjectBuilder(ObjectBuilder): + """Minimal builder that records build() calls and returns a canned Object.""" + + def __init__(self, settings: ObjectSettings, canned: Object) -> None: + super().__init__(settings, 'recording') + self._settings = settings + self._canned = canned + self.build_calls: list[ObjectGeometryProvider] = [] + + def copy(self) -> _RecordingObjectBuilder: + return _RecordingObjectBuilder(self._settings, self._canned) + + def _build_raw(self, geometry_provider: ObjectGeometryProvider) -> Object: + self.build_calls.append(geometry_provider) + return self._canned + + def build( + self, + geometry_provider: ObjectGeometryProvider, + layer_spacing_m: Sequence[float], + ) -> Object: + # These tests exercise rebuild's geometry guard, not the conditioning + # pipeline, so bypass it and hand back the canned object by identity. + return self._build_raw(geometry_provider) + + +def _make_object_geometry(pixel_width_m: float, pixel_height_m: float) -> ObjectGeometry: + return ObjectGeometry( + width_px=8, + height_px=8, + pixel_width_m=pixel_width_m, + pixel_height_m=pixel_height_m, + center_x_m=0.0, + center_y_m=0.0, + ) + + +def test_rebuild_fires_on_geometry_observer_notification() -> None: + """When the geometry provider is Observable, ObjectRepositoryItem should + register itself and re-run rebuild each time notify_observers fires + (matches the ProductGeometry.set_detector_extent path in production). + Also verifies the is_valid guard blocks the initial rebuild when the + provider reports zero-valued pixel dimensions. + """ + registry = SettingsRegistry() + settings = ObjectSettings(registry) + provider = _ObservableObjectProvider(_make_object_geometry(0.0, 0.0)) + canned = Object( + array=numpy.zeros((1, 4, 4), dtype=complex), + pixel_geometry=PixelGeometry(width_m=1e-6, height_m=1e-6), + center=ObjectCenter(coordinate_x_m=0.0, coordinate_y_m=0.0), + ) + builder = _RecordingObjectBuilder(settings, canned) + + item = ObjectRepositoryItem(provider, settings, builder) + assert builder.build_calls == [] # guard blocks initial rebuild + + provider.set_geometry(_make_object_geometry(1e-6, 1e-6)) + + assert len(builder.build_calls) == 1 + assert item.get_object().get_array() is canned.get_array() diff --git a/tests/test_parent_config_builders.py b/tests/test_parent_config_builders.py new file mode 100644 index 000000000..706f28edc --- /dev/null +++ b/tests/test_parent_config_builders.py @@ -0,0 +1,218 @@ +"""Discipline gate for the parent-side backend config builders. + +``ptychopinn`` and ``ptychopinn_torch`` build their backend config objects in +the parent so the child receives finished objects and does nothing but GPU +work. That means the parent reaches ``ptycho.config.config`` / +``ptycho_torch.config_params`` at call time -- which +``tests/test_no_gpu_context.py`` cannot observe, because its probe only imports +the composition roots and never dispatches a reconstruction. + +This module closes that gap. For each installed backend it spawns a probe that +calls the builders exactly as ``build_*_payload`` does, then asserts: + +1. the configs construct and survive the pickle round-trip the spawn transport + requires, and +2. no GPU context was acquired in the process that built them. + +Each probe runs in its own spawned process so a framework loaded by an earlier +test cannot mask a failure here. +""" + +from __future__ import annotations + +from importlib.util import find_spec +from typing import Any +import multiprocessing +import sys +import textwrap + +import pytest + + +def _gpu_context_offenders() -> list[str]: + """Report any live GPU context. Observation only -- creates no context.""" + offenders: list[str] = [] + + torch = sys.modules.get('torch') + if torch is not None: + for backend_name in ('cuda', 'xpu'): + backend = getattr(torch, backend_name, None) + is_initialized = getattr(backend, 'is_initialized', None) + if is_initialized is None: + continue # e.g. torch.xpu is absent on older torch builds + try: + if is_initialized(): + offenders.append(f'torch.{backend_name} context is initialized') + except Exception as exc: # noqa: BLE001 + offenders.append(f'could not probe torch.{backend_name}: {type(exc).__name__}') + + if 'tensorflow' in sys.modules: + try: + from tensorflow.python.eager import context as tf_context + + if tf_context.context_safe() is not None: + offenders.append('tensorflow eager context is initialized') + except Exception as exc: # noqa: BLE001 + offenders.append(f'could not probe tensorflow: {type(exc).__name__}') + + return offenders + + +def _build_ptychopinn_configs() -> list[Any]: + from ptychodus.api.settings import SettingsRegistry + from ptychodus.model.ptychopinn.reconstructor import ( + _build_inference_config, + _build_training_config, + ) + from ptychodus.model.ptychopinn.settings import ( + PtychoPINNModelSettings, + PtychoPINNTrainingSettings, + ) + + registry = SettingsRegistry() + model_settings = PtychoPINNModelSettings(registry) + training_settings = PtychoPINNTrainingSettings(registry) + + return [ + _build_inference_config( + model_settings, 'PINN', model_size=64, is_developer_mode_enabled=False + ), + _build_training_config(model_settings, training_settings, 'PINN'), + ] + + +def _build_ptychopinn_torch_configs() -> list[Any]: + from ptychodus.api.settings import SettingsRegistry + from ptychodus.model.ptychopinn_torch.reconstructor import _build_configs + from ptychodus.model.ptychopinn_torch.settings import ( + PtychoPINNTorchDataSettings, + PtychoPINNTorchInferenceSettings, + PtychoPINNTorchModelSettings, + PtychoPINNTorchTrainingSettings, + ) + + registry = SettingsRegistry() + + return list( + _build_configs( + 'Unsupervised', + PtychoPINNTorchDataSettings(registry), + PtychoPINNTorchModelSettings(registry), + PtychoPINNTorchTrainingSettings(registry), + PtychoPINNTorchInferenceSettings(registry), + ) + ) + + +def _build_ptycho_fm_configs() -> list[Any]: + from ptychodus.api.settings import SettingsRegistry + from ptychodus.model.ptycho_fm.reconstructor import _build_config + from ptychodus.model.ptycho_fm.settings import ( + PtychoFMDataSettings, + PtychoFMInferenceSettings, + PtychoFMModelSettings, + PtychoFMTrainingSettings, + ) + + registry = SettingsRegistry() + + return [ + _build_config( + PtychoFMDataSettings(registry), + PtychoFMModelSettings(registry), + PtychoFMTrainingSettings(registry), + PtychoFMInferenceSettings(registry), + ) + ] + + +_BUILDERS = { + 'ptychopinn': _build_ptychopinn_configs, + 'ptychopinn_torch': _build_ptychopinn_torch_configs, + 'ptycho_fm': _build_ptycho_fm_configs, +} + +_REQUIRED_MODULES: dict[str, tuple[str, ...]] = { + 'ptychopinn': ('ptycho',), + 'ptychopinn_torch': ('ptycho_torch',), + # ptycho_fm's _build_config is a pure-Python dict factory; it does not + # import ptycho_vit or torch, so this test can run everywhere. + 'ptycho_fm': (), +} + + +def _probe(backend: str, result_queue: multiprocessing.Queue[dict[str, Any] | str]) -> None: + """Build the configs, then report the outcome AND the GPU-context state. + + A build failure is reported rather than raised, because the context check + is meaningful either way: by the time a builder has failed it has already + imported whatever framework it was going to import. + """ + import pickle + + result: dict[str, Any] = {'count': 0, 'types': [], 'build_error': None} + + try: + configs = _BUILDERS[backend]() + round_tripped = [pickle.loads(pickle.dumps(config)) for config in configs] + result['count'] = len(configs) + result['types'] = [type(config).__name__ for config in round_tripped] + except BaseException as exc: # noqa: BLE001 + result['build_error'] = f'{type(exc).__name__}: {exc}' + + try: + result['context'] = _gpu_context_offenders() + except BaseException as exc: # noqa: BLE001 + result_queue.put(f'context probe failed: {type(exc).__name__}: {exc}') + return + + result_queue.put(result) + + +@pytest.mark.parametrize('backend', sorted(_BUILDERS)) +def test_parent_config_builders_acquire_no_gpu_context(backend: str) -> None: + for module_name in _REQUIRED_MODULES[backend]: + if find_spec(module_name) is None: + pytest.skip(f'{module_name} is not installed') + + ctx = multiprocessing.get_context('spawn') + result_queue: multiprocessing.Queue[dict[str, Any] | str] = ctx.Queue() + process = ctx.Process(target=_probe, args=(backend, result_queue)) + process.start() + process.join(timeout=180.0) + + assert not process.is_alive(), f'{backend} config-builder probe did not exit in 180s.' + + result: Any = result_queue.get(timeout=1.0) + + if isinstance(result, str): + raise AssertionError(f'{backend} config-builder probe raised: {result}') + + # The invariant this module exists to pin, checked whether or not the + # builder itself succeeded. + assert result['context'] == [], textwrap.dedent( + f"""\ + Building {backend} configs parent-side acquired a GPU context. + Offenders: {result['context']} + Configs built: {result['types']} + + The parent is allowed to import a GPU framework to construct the + picklable configs it ships to the child, but not to acquire a context + doing it. Something in the builder placed a tensor on a device or + queried the runtime. See the invariant note in + ptychodus.model.processing._subprocess_protocol. + """ + ) + + if result['build_error'] is not None: + # The installed backend's config schema does not match what the + # settings mapping targets. That is an environment/version mismatch + # rather than a ptychodus defect -- the mapping is the same one the + # child used before it moved parent-side -- so report it loudly + # instead of failing the suite against an arbitrary local checkout. + pytest.skip( + f'{backend} config schema has drifted from the installed backend: ' + f'{result["build_error"]}' + ) + + assert result['count'] > 0, f'{backend} built no configs.' diff --git a/tests/test_plugin_chooser_parameter.py b/tests/test_plugin_chooser_parameter.py new file mode 100644 index 000000000..d85f68aab --- /dev/null +++ b/tests/test_plugin_chooser_parameter.py @@ -0,0 +1,176 @@ +"""Tests for PluginChooserParameter, the display-name view of a PluginChooser. + +A PluginChooser carries two name spaces: the human-readable ``display_name`` +shown in the GUI and the ``simple_name`` persisted to settings. A combo box +bound directly to the settings parameter therefore mis-restores at startup, +because it would look up a simple name among display-name items. +PluginChooserParameter exists to close that gap, and it is also the sole owner +of the settings binding, so these tests pin the round trip in both directions +and confirm persistence still stores simple names. + +No Qt is required — this is pure model-layer behavior. +""" + +from __future__ import annotations + +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.parametric import ParameterGroup, StringParameter +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter + +# Display names deliberately chosen so that simple_name != display_name for +# most of them: re.sub(r'\W+', '', ...) strips the space and the hyphen. +DISPLAY_NAMES = ['Identity', 'Richardson-Lucy', 'Unsupervised Wiener', 'Wiener'] + + +class _Counter(Observer): + def __init__(self) -> None: + self.count = 0 + + def _update(self, observable: Observable) -> None: + self.count += 1 + + +def _make_settings(default: str) -> StringParameter: + group = ParameterGroup() + return group.create_string_parameter('DeconvolutionStrategy', default) + + +def _build( + default: str = 'Richardson-Lucy', +) -> tuple[PluginChooserParameter[str], StringParameter]: + settings = _make_settings(default) + chooser = PluginChooser[str]() + + for display_name in DISPLAY_NAMES: + chooser.register_plugin(display_name, display_name=display_name) + + return PluginChooserParameter(chooser, settings), settings + + +def test_value_is_display_name_while_settings_hold_simple_name() -> None: + chooser_parameter, settings = _build() + + assert chooser_parameter.get_value() == 'Richardson-Lucy' + assert settings.get_value() == 'RichardsonLucy' + + +def test_restores_display_name_from_persisted_simple_name() -> None: + """The regression this adapter exists to prevent. + + Reading back a simple name from the INI must still surface the display name, + otherwise a combo box populated with display names silently falls back to + its first item. + """ + chooser_parameter, _ = _build(default='UnsupervisedWiener') + + assert chooser_parameter.get_value() == 'Unsupervised Wiener' + + +def test_set_value_by_display_name_moves_chooser_and_settles() -> None: + chooser_parameter, settings = _build() + counter = _Counter() + chooser_parameter.add_observer(counter) + + chooser_parameter.set_value('Unsupervised Wiener') + + assert chooser_parameter.get_value() == 'Unsupervised Wiener' + assert settings.get_value() == 'UnsupervisedWiener' + # One notification, not a cascade: the chooser's index guard breaks the loop. + assert counter.count == 1 + + +def test_selecting_the_current_plugin_is_silent() -> None: + chooser_parameter, _ = _build() + counter = _Counter() + chooser_parameter.add_observer(counter) + + chooser_parameter.set_value('Richardson-Lucy') + + assert counter.count == 0 + + +def test_set_value_honors_notify_false() -> None: + chooser_parameter, _ = _build() + counter = _Counter() + chooser_parameter.add_observer(counter) + + chooser_parameter.set_value('Wiener', notify=False) + + assert chooser_parameter.get_value() == 'Wiener' + assert counter.count == 0 + + # The suppression must not leak into the next assignment. + chooser_parameter.set_value('Identity') + assert counter.count == 1 + + +def test_notify_false_still_persists() -> None: + """Suppression is a view concern; it must never skip the settings write-back.""" + chooser_parameter, settings = _build() + + chooser_parameter.set_value('Unsupervised Wiener', notify=False) + + assert settings.get_value() == 'UnsupervisedWiener' + + +def test_external_chooser_change_notifies() -> None: + """A selection made elsewhere (e.g. batch mode) must reach GUI observers.""" + chooser_parameter, _ = _build() + counter = _Counter() + chooser_parameter.add_observer(counter) + + chooser_parameter.set_value('Wiener') + + assert chooser_parameter.get_value() == 'Wiener' + assert counter.count == 1 + + +def test_settings_change_moves_the_selection() -> None: + """Loading an INI mutates the settings parameter; the selection must follow.""" + chooser_parameter, settings = _build() + + settings.set_value('Wiener') + + assert chooser_parameter.get_value() == 'Wiener' + + +def test_string_conversion_round_trips_display_names() -> None: + chooser_parameter, _ = _build() + + chooser_parameter.set_value_from_string('Unsupervised Wiener') + + assert chooser_parameter.get_value_as_string() == 'Unsupervised Wiener' + + +def test_unknown_name_leaves_selection_unchanged() -> None: + chooser_parameter, _ = _build() + + chooser_parameter.set_value('Nonexistent Strategy') + + assert chooser_parameter.get_value() == 'Richardson-Lucy' + + +def test_choices_are_display_names_in_chooser_order() -> None: + chooser_parameter, _ = _build() + + assert list(chooser_parameter.choices()) == DISPLAY_NAMES + + +def test_get_strategy_returns_the_selected_plugin_strategy() -> None: + chooser_parameter, _ = _build() + + chooser_parameter.set_value('Wiener') + + assert chooser_parameter.get_strategy() == 'Wiener' + + +def test_copy_is_an_unbound_view() -> None: + """A copy tracks the same chooser but must not write to the original's settings.""" + chooser_parameter, settings = _build() + copied = chooser_parameter.copy() + + copied.set_value('Wiener') + + assert chooser_parameter.get_value() == 'Wiener' + # The original adapter is still bound, so it persists the change it observed. + assert settings.get_value() == 'Wiener' diff --git a/tests/test_plugins.py b/tests/test_plugins.py index e2375f68e..e97c96e4a 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -1,10 +1,36 @@ -"""Unit tests for ptychodus.api.plugins.PluginChooser.""" +"""Unit tests for ptychodus.api.plugins.PluginChooser and its settings binding. + +The chooser is a pure registry plus a tracked selection; PluginChooserParameter +is the only thing that knows about settings persistence. These tests pin the +selection semantics that the two classes agree on at that seam: + +- a name resolves against either name space, but only the canonical simple name + is ever written back to settings, including when it resolves to index 0; +- an unrecognized name holds the selection *and* the persisted value, but still + notifies so a bound view resynchronizes; +- registering a plugin re-sorts the list without repointing the selection. +""" from __future__ import annotations import pytest -from ptychodus.api.plugins import PluginChooser +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.parametric import ParameterGroup, StringParameter +from ptychodus.api.plugins import PluginChooser, PluginChooserParameter + + +class _Counter(Observer): + def __init__(self) -> None: + self.count = 0 + + def _update(self, observable: Observable) -> None: + self.count += 1 + + +def _make_settings(default: str) -> StringParameter: + group = ParameterGroup() + return group.create_string_parameter('FileType', default) def test_get_current_plugin_empty_raises_lookup_error() -> None: @@ -35,3 +61,144 @@ def test_populated_chooser_is_truthy() -> None: chooser.register_plugin('s', display_name='S') assert chooser + + +def test_find_plugin_matches_either_name_space_case_insensitively() -> None: + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('s', display_name='Richardson-Lucy') + + by_display = chooser.find_plugin('richardson-lucy') + by_simple = chooser.find_plugin('RICHARDSONLUCY') + + assert by_display is not None + assert by_simple is by_display + assert by_display.simple_name == 'RichardsonLucy' + + +def test_find_plugin_returns_none_for_unknown_name() -> None: + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('s', display_name='Richardson-Lucy') + + assert chooser.find_plugin('Nonexistent') is None + + +def test_binding_normalizes_a_persisted_display_name() -> None: + settings = _make_settings('Richardson-Lucy') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('i', display_name='Identity') + chooser.register_plugin('r', display_name='Richardson-Lucy') + + PluginChooserParameter(chooser, settings) + + assert chooser.get_current_plugin().simple_name == 'RichardsonLucy' + assert settings.get_value() == 'RichardsonLucy' + + +def test_binding_leaves_a_persisted_simple_name_alone() -> None: + settings = _make_settings('RichardsonLucy') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('i', display_name='Identity') + chooser.register_plugin('r', display_name='Richardson-Lucy') + + PluginChooserParameter(chooser, settings) + + assert settings.get_value() == 'RichardsonLucy' + + +def test_binding_normalizes_even_at_index_zero() -> None: + """The selection does not move, but the persisted value must still be canonical. + + 'Alpha-Plugin' sorts first, so binding to it leaves _current_index at 0. An + earlier implementation gated the write-back on the index changing, which left + a display name sitting in the INI forever. + """ + settings = _make_settings('Alpha-Plugin') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('a', display_name='Alpha-Plugin') + chooser.register_plugin('z', display_name='Zeta-Plugin') + + PluginChooserParameter(chooser, settings) + + assert chooser.get_current_plugin().simple_name == 'AlphaPlugin' + assert settings.get_value() == 'AlphaPlugin' + + +def test_unknown_name_holds_selection_and_setting_but_notifies() -> None: + """A name can be unresolvable because an optional-dependency plugin is missing. + + Overwriting the setting would discard the user's choice for good, so only the + view is resynchronized. + """ + settings = _make_settings('Identity') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('i', display_name='Identity') + chooser.register_plugin('w', display_name='Wiener') + chooser_parameter = PluginChooserParameter(chooser, settings) + + settings.set_value('MissingPlugin') + counter = _Counter() + chooser_parameter.add_observer(counter) + chooser.set_current_plugin('MissingPlugin') + + assert chooser.get_current_plugin().display_name == 'Identity' + assert settings.get_value() == 'MissingPlugin' + assert counter.count == 1 + + +def test_registration_does_not_repoint_the_selection() -> None: + """register_plugin re-sorts by display name; the selected plugin must survive it.""" + settings = _make_settings('Zeta') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('m', display_name='Middle') + chooser.register_plugin('z', display_name='Zeta') + PluginChooserParameter(chooser, settings) + + chooser.register_plugin('a', display_name='Alpha') + + assert chooser.get_current_plugin().display_name == 'Zeta' + assert settings.get_value() == 'Zeta' + + +def test_registration_does_not_clobber_an_unresolved_setting() -> None: + settings = _make_settings('MissingPlugin') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('m', display_name='Middle') + PluginChooserParameter(chooser, settings) + + chooser.register_plugin('a', display_name='Alpha') + + assert settings.get_value() == 'MissingPlugin' + + +def test_late_registration_settles_an_empty_chooser_onto_the_persisted_name() -> None: + """Binding before plugins are registered must still honor the persisted choice.""" + settings = _make_settings('Zeta') + chooser: PluginChooser[str] = PluginChooser() + chooser_parameter = PluginChooserParameter(chooser, settings) + + chooser.register_plugin('a', display_name='Alpha') + chooser.register_plugin('z', display_name='Zeta') + + assert chooser_parameter.get_value() == 'Zeta' + assert settings.get_value() == 'Zeta' + + +def test_two_adapters_over_one_chooser_both_track_it() -> None: + """The single-binding slot that synchronize_with_parameter had is gone. + + A second binding used to silently replace the first and leave it diverging; + now each adapter owns its own settings parameter and both stay in step. + """ + settings_a = _make_settings('Alpha') + settings_b = _make_settings('Alpha') + chooser: PluginChooser[str] = PluginChooser() + chooser.register_plugin('a', display_name='Alpha') + chooser.register_plugin('z', display_name='Zeta') + + parameter_a = PluginChooserParameter(chooser, settings_a) + parameter_b = PluginChooserParameter(chooser, settings_b) + parameter_a.set_value('Zeta') + + assert parameter_b.get_value() == 'Zeta' + assert settings_a.get_value() == 'Zeta' + assert settings_b.get_value() == 'Zeta' diff --git a/tests/test_probe.py b/tests/test_probe.py index e16867488..11971718b 100644 --- a/tests/test_probe.py +++ b/tests/test_probe.py @@ -7,6 +7,7 @@ from ptychodus.api.geometry import PixelGeometry from ptychodus.api.probe import ( Probe, + ProbeSequence, ProbeSizeMetrics, compute_shannon_entropy, estimate_probe_entropy, @@ -229,3 +230,64 @@ def test_fourier_duality_direction(self) -> None: assert narrow.real_space_intensity_entropy < broad.real_space_intensity_entropy assert narrow.spectral_entropy > broad.spectral_entropy + + +class TestProbeSequenceSlicing: + """Regression tests for ProbeSequence.__getitem__. + + The slice path used to build its indices from + ``range(index.start, index.stop, index.step)``, so every slice with an + implicit bound raised TypeError. Only a fully-specified positive-step slice + such as ``seq[0:2:1]`` worked. + """ + + def _sequence(self, num_positions: int = 5) -> ProbeSequence: + """Two coherent modes over *num_positions* scan positions, so len() > 1.""" + array = numpy.arange(2 * 1 * 4 * 4, dtype=numpy.complex128).reshape(2, 1, 4, 4) + opr_weights = numpy.arange(num_positions * 2, dtype=float).reshape(num_positions, 2) + return ProbeSequence(array, opr_weights, PIXEL_GEOMETRY) + + @pytest.mark.parametrize( + ('index', 'expected_length'), + [ + (slice(None, 2), 2), + (slice(1, 3), 2), + (slice(None), 5), + (slice(None, None, 2), 3), + (slice(-2, None), 2), + (slice(None, 99), 5), + ], + ) + def test_slice_returns_expected_number_of_probes( + self, index: slice, expected_length: int + ) -> None: + probes = self._sequence()[index] + assert len(probes) == expected_length + assert all(isinstance(probe, Probe) for probe in probes) + + def test_empty_slice_returns_empty_list(self) -> None: + assert len(self._sequence()[3:1]) == 0 + + def test_negative_step_reverses(self) -> None: + sequence = self._sequence() + reversed_probes = sequence[::-1] + assert len(reversed_probes) == 5 + + for offset, probe in enumerate(reversed_probes): + expected = sequence[4 - offset] + numpy.testing.assert_array_equal(probe.get_array(), expected.get_array()) + + def test_slice_agrees_with_integer_indexing(self) -> None: + sequence = self._sequence() + + for offset, probe in enumerate(sequence[1:4]): + expected = sequence[1 + offset] + numpy.testing.assert_array_equal(probe.get_array(), expected.get_array()) + + def test_slice_without_opr_weights(self) -> None: + """__len__ reports 1 when there are no OPR weights, so [:] yields one probe.""" + array = numpy.ones((1, 1, 4, 4), dtype=numpy.complex128) + sequence = ProbeSequence(array, None, PIXEL_GEOMETRY) + + assert len(sequence) == 1 + assert len(sequence[:]) == 1 diff --git a/tests/test_probe_builder.py b/tests/test_probe_builder.py new file mode 100644 index 000000000..91c3b7d03 --- /dev/null +++ b/tests/test_probe_builder.py @@ -0,0 +1,326 @@ +"""Regression tests for the probe conditioning pipeline (incoherent modes -> OPR modes). + +Two invariants carry the weight here. + +First, conditioning is expand-only and therefore idempotent. The generators in +ptychodus.api.probe_gen are not safe to re-apply: generate_incoherent_probe_modes +re-orthogonalizes and renormalizes every mode to the decay profile, and +generate_coherent_probe_modes fills its output with fresh Gaussian noise, keeps +only coherent mode zero of its input, and regenerates the OPR weights from +scratch. The guards in ProbeSequenceBuilder._condition_probe are what stop a +converged OPR basis being replaced with noise. + +Second, FromMemoryProbeBuilder must never condition. It holds a probe that is +already conditioned -- reconstruction output, which ProcessingTaskMonitor +re-assigns to the output product item on every reconstructor iteration, and +products loaded from HDF5/NPZ. Re-running the mode generators there would destroy +the reconstruction, once per iteration. + +The photon-count rescale sits on the other side of the split: it is +generation-only, because a probe read from file already carries the intensity it +was reconstructed at. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy +import pytest + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.probe import ( + ProbeFileReader, + ProbeGeometry, + ProbeGeometryProvider, + ProbeSequence, +) +from ptychodus.api.settings import SettingsRegistry +from ptychodus.model.product.probe.builder import ( + FromFileProbeBuilder, + FromMemoryProbeBuilder, + ProbeSequenceBuilder, +) +from ptychodus.model.product.probe.disk import DiskProbeBuilder +from ptychodus.model.product.probe.settings import ProbeSettings + +NUM_SCAN_POINTS = 7 +PROBE_PHOTON_COUNT = 1.0e6 +# Fine enough that the default 1 um disk covers several pixels; at a coarser +# pixel size the generated probe is empty and rescale_probe_intensity bails. +PIXEL_SIZE_M = 1.0e-7 +PROBE_EXTENT_PX = 64 + + +def _make_settings() -> ProbeSettings: + return ProbeSettings(SettingsRegistry()) + + +def _make_rng() -> numpy.random.Generator: + return numpy.random.default_rng(42) + + +class _StubProbeGeometryProvider(ProbeGeometryProvider): + """A ready geometry provider, so the builders never hit the not-yet-bound guard.""" + + def __init__(self, *, probe_photon_count: float = PROBE_PHOTON_COUNT) -> None: + self._probe_photon_count = probe_photon_count + + @property + def detector_distance_m(self) -> float: + return 1.0 + + @property + def probe_photon_count(self) -> float: + return self._probe_photon_count + + @property + def probe_wavelength_m(self) -> float: + return 1.0e-10 + + @property + def probe_power_W(self) -> float: # noqa: N802 + return 1.0 + + @property + def num_scan_points(self) -> int: + return NUM_SCAN_POINTS + + def get_detector_pixel_geometry(self) -> PixelGeometry: + return PixelGeometry(width_m=PIXEL_SIZE_M, height_m=PIXEL_SIZE_M) + + def get_probe_geometry(self) -> ProbeGeometry: + return ProbeGeometry( + width_px=PROBE_EXTENT_PX, + height_px=PROBE_EXTENT_PX, + pixel_width_m=PIXEL_SIZE_M, + pixel_height_m=PIXEL_SIZE_M, + ) + + +def _make_probe_seq( + num_cmodes: int, num_imodes: int, *, with_opr_weights: bool = False +) -> ProbeSequence: + """A deterministic, non-degenerate probe of the requested mode structure.""" + rng = numpy.random.default_rng(7) + shape = (num_cmodes, num_imodes, 8, 8) + array = (rng.normal(size=shape) + 1j * rng.normal(size=shape)).astype(complex) + + opr_weights = None + + if with_opr_weights: + opr_weights = rng.normal(size=(NUM_SCAN_POINTS, num_cmodes)) + opr_weights[:, 0] = 1.0 + + return ProbeSequence( + array=array, + opr_weights=opr_weights, + pixel_geometry=PixelGeometry(width_m=PIXEL_SIZE_M, height_m=PIXEL_SIZE_M), + ) + + +class _StubProbeFileReader(ProbeFileReader): + def __init__(self, probe_seq: ProbeSequence) -> None: + self._probe_seq = probe_seq + + def read(self, file_path: Path) -> ProbeSequence: + return self._probe_seq + + +def _make_from_file_builder( + settings: ProbeSettings, probe_seq: ProbeSequence +) -> FromFileProbeBuilder: + return FromFileProbeBuilder(_make_rng(), settings, _StubProbeFileReader(probe_seq)) + + +def _total_intensity(probe_seq: ProbeSequence) -> float: + return float(numpy.sum(numpy.abs(probe_seq.get_array()) ** 2)) + + +def test_generator_expands_incoherent_modes() -> None: + """The refactor moved mode generation out of each generator's tail and into + the base pipeline; generators must still come back multimodal.""" + settings = _make_settings() + builder = DiskProbeBuilder(_make_rng(), settings) + builder.num_incoherent_modes.set_value(4) + + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert probe_seq.num_incoherent_modes == 4 + assert probe_seq.num_coherent_modes == 1 + + +def test_generator_expands_coherent_modes() -> None: + settings = _make_settings() + builder = DiskProbeBuilder(_make_rng(), settings) + builder.num_coherent_modes.set_value(3) + + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert probe_seq.get_array().shape[:2] == (3, 1) + + opr_weights = probe_seq.get_opr_weights_or_none() + assert opr_weights is not None + assert opr_weights.shape == (NUM_SCAN_POINTS, 3) + + +def test_generator_rescales_to_photon_count() -> None: + """rescale_probe_intensity was duplicated in all seven generator tails and is + now a single base helper; the generative path must still be normalized.""" + settings = _make_settings() + builder = DiskProbeBuilder(_make_rng(), settings) + + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert _total_intensity(probe_seq) == pytest.approx(PROBE_PHOTON_COUNT) + + +def test_from_file_builder_expands_modes() -> None: + """FromFileProbeBuilder.build() used to return the reader's output verbatim, + so the mode settings were silently ignored for every file-loaded probe even + though the reconstructor honored them.""" + settings = _make_settings() + builder = _make_from_file_builder(settings, _make_probe_seq(1, 1)) + builder.num_incoherent_modes.set_value(3) + + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert probe_seq.num_incoherent_modes == 3 + + +def test_from_file_builder_does_not_rescale_intensity() -> None: + """The photon-count rescale is generation-only. A file probe carries the + intensity it was reconstructed at, and rescaling it without applying the + reciprocal to a matching from-file object would break the product P*O.""" + settings = _make_settings() + from_file = _make_probe_seq(1, 1) + builder = _make_from_file_builder(settings, from_file) + + probe_seq = builder.build(_StubProbeGeometryProvider(probe_photon_count=1.0e12)) + + assert _total_intensity(probe_seq) == pytest.approx(_total_intensity(from_file)) + + +def test_from_file_builder_keeps_extra_incoherent_modes( + caplog: pytest.LogCaptureFixture, +) -> None: + """Expand-only: asking for fewer modes than the file carries must not discard + the converged ones.""" + settings = _make_settings() + builder = _make_from_file_builder(settings, _make_probe_seq(1, 4)) + builder.num_incoherent_modes.set_value(1) + + with caplog.at_level('INFO'): + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert probe_seq.num_incoherent_modes == 4 + assert 'keeping them rather than discarding down to 1' in caplog.text + + +def test_from_file_builder_preserves_opr_basis(caplog: pytest.LogCaptureFixture) -> None: + """The catastrophic case. generate_coherent_probe_modes would replace every + coherent mode but the first with Gaussian noise and regenerate the OPR + weights, so a solved OPR basis must never reach it.""" + settings = _make_settings() + from_file = _make_probe_seq(3, 2, with_opr_weights=True) + builder = _make_from_file_builder(settings, from_file) + builder.num_coherent_modes.set_value(5) + builder.num_incoherent_modes.set_value(6) + + with caplog.at_level('INFO'): + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert numpy.array_equal(probe_seq.get_array(), from_file.get_array()) + assert numpy.array_equal( + probe_seq.get_opr_weights(), + from_file.get_opr_weights(), + ) + assert 'leaving its mode structure unchanged' in caplog.text + + +@pytest.mark.parametrize(('num_imodes', 'num_cmodes'), [(1, 1), (3, 1), (1, 3), (3, 2)]) +def test_conditioning_is_idempotent(num_imodes: int, num_cmodes: int) -> None: + """Conditioning an already-conditioned probe must be a no-op. Several rebuild + paths -- a geometry-provider notification, a builder-parameter edit -- can + re-run build() on a probe that has already been through the pipeline.""" + settings = _make_settings() + provider = _StubProbeGeometryProvider() + + first = _make_from_file_builder(settings, _make_probe_seq(1, 1)) + first.num_incoherent_modes.set_value(num_imodes) + first.num_coherent_modes.set_value(num_cmodes) + conditioned = first.build(provider) + + second = _make_from_file_builder(settings, conditioned) + second.num_incoherent_modes.set_value(num_imodes) + second.num_coherent_modes.set_value(num_cmodes) + reconditioned = second.build(provider) + + assert numpy.array_equal(reconditioned.get_array(), conditioned.get_array()) + + weights = conditioned.get_opr_weights_or_none() + reweights = reconditioned.get_opr_weights_or_none() + + if weights is None: + assert reweights is None + else: + assert reweights is not None + assert numpy.array_equal(reweights, weights) + + +def test_from_memory_builder_ignores_conditioning() -> None: + """Guards reconstruction output: the from-memory builder must return its + probe verbatim no matter what the mode parameters say.""" + settings = _make_settings() + raw = _make_probe_seq(2, 3, with_opr_weights=True) + builder = FromMemoryProbeBuilder(_make_rng(), settings, raw) + builder.num_incoherent_modes.set_value(8) + builder.num_coherent_modes.set_value(8) + + probe_seq = builder.build(_StubProbeGeometryProvider()) + + assert numpy.array_equal(probe_seq.get_array(), raw.get_array()) + assert numpy.array_equal(probe_seq.get_opr_weights(), raw.get_opr_weights()) + + +def test_repeated_from_memory_builds_are_idempotent() -> None: + """The reconstruct loop rebuilds the output item's probe once per iteration; + conditioning must not accumulate across those rebuilds.""" + settings = _make_settings() + settings.num_incoherent_modes.set_value(5) + settings.num_coherent_modes.set_value(4) + + provider = _StubProbeGeometryProvider() + expected = _make_probe_seq(2, 3, with_opr_weights=True) + probe_seq = expected + + for _ in range(3): + builder = FromMemoryProbeBuilder(_make_rng(), settings, probe_seq) + probe_seq = builder.build(provider) + + assert numpy.array_equal(probe_seq.get_array(), expected.get_array()) + assert numpy.array_equal(probe_seq.get_opr_weights(), expected.get_opr_weights()) + + +@pytest.mark.parametrize('builder_name', ['disk', 'from_file']) +def test_copy_preserves_mode_parameters(builder_name: str) -> None: + """copy() iterates parameters() generically and now also has to carry the rng + hoisted into the base, so the copy must still build.""" + settings = _make_settings() + builder: ProbeSequenceBuilder + + if builder_name == 'disk': + builder = DiskProbeBuilder(_make_rng(), settings) + else: + builder = _make_from_file_builder(settings, _make_probe_seq(1, 1)) + + builder.num_incoherent_modes.set_value(3) + builder.num_coherent_modes.set_value(2) + + duplicate = builder.copy() + + assert duplicate.num_incoherent_modes.get_value() == 3 + assert duplicate.num_coherent_modes.get_value() == 2 + + probe_seq = duplicate.build(_StubProbeGeometryProvider()) + assert probe_seq.get_array().shape[:2] == (2, 3) diff --git a/tests/test_probe_item.py b/tests/test_probe_item.py new file mode 100644 index 000000000..ba76ecc3c --- /dev/null +++ b/tests/test_probe_item.py @@ -0,0 +1,197 @@ +"""Regression tests for ProbeRepositoryItem._rebuild. + +The critical invariant: when the geometry provider has not yet bound to a +dataset (its ProbeGeometry has zero-valued pixel dimensions), _rebuild must +NOT invoke the builder — otherwise builders like FZP that divide by +`geometry.width_m` crash with ZeroDivisionError. See CLAUDE fly001.ini bug +report. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import numpy + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.observer import Observable +from ptychodus.api.probe import ProbeGeometry, ProbeGeometryProvider, ProbeSequence +from ptychodus.api.settings import SettingsRegistry +from ptychodus.model.product.probe.builder import ProbeSequenceBuilder +from ptychodus.model.product.probe.item import ProbeRepositoryItem +from ptychodus.model.product.probe.settings import ProbeSettings + + +def _make_rng() -> numpy.random.Generator: + return numpy.random.default_rng(42) + + +class _RecordingBuilder(ProbeSequenceBuilder): + """Minimal builder that records build() invocations without touching numpy math.""" + + def __init__(self, settings: ProbeSettings, probe_seq: ProbeSequence) -> None: + super().__init__(_make_rng(), settings, 'recording') + self._settings = settings + self._probe_seq = probe_seq + self.build_calls: list[ProbeGeometryProvider] = [] + + def copy(self) -> _RecordingBuilder: + return _RecordingBuilder(self._settings, self._probe_seq) + + def _build_raw(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + self.build_calls.append(geometry_provider) + return self._probe_seq + + def build(self, geometry_provider: ProbeGeometryProvider) -> ProbeSequence: + # These tests exercise _rebuild's geometry guard, not the conditioning + # pipeline, so bypass it and hand back the canned sequence by identity. + return self._build_raw(geometry_provider) + + +def _make_provider(pixel_width_m: float, pixel_height_m: float) -> MagicMock: + provider = MagicMock(spec=ProbeGeometryProvider) + provider.get_probe_geometry.return_value = ProbeGeometry( + width_px=64, + height_px=64, + pixel_width_m=pixel_width_m, + pixel_height_m=pixel_height_m, + ) + return provider + + +def _make_probe_seq(pixel_size_m: float) -> ProbeSequence: + array = numpy.zeros((1, 4, 4), dtype=numpy.complex64) + return ProbeSequence( + array=array, + opr_weights=None, + pixel_geometry=PixelGeometry(width_m=pixel_size_m, height_m=pixel_size_m), + ) + + +def test_rebuild_skips_when_geometry_not_ready() -> None: + """Pre-dataset startup: pixel dimensions are zero. Builder must NOT be called; + the initial null ProbeSequence stays in place; no exception escapes.""" + registry = SettingsRegistry() + settings = ProbeSettings(registry) + provider = _make_provider(pixel_width_m=0.0, pixel_height_m=0.0) + canned = _make_probe_seq(pixel_size_m=1e-6) + builder = _RecordingBuilder(settings, canned) + + item = ProbeRepositoryItem(_make_rng(), provider, settings, builder) + + assert builder.build_calls == [] + # The null sentinel from ProbeRepositoryItem.__init__ has size 0; the canned + # replacement would be shape (1, 4, 4). Same-size check would let a silent + # overwrite slip past. + assert item.get_probes().get_array().size == 0 + + +def test_rebuild_fires_when_geometry_becomes_ready() -> None: + """Once the provider reports a valid pixel geometry and the item is nudged + (e.g. via set_builder from a settings change), the builder runs and its + ProbeSequence replaces the null sentinel.""" + registry = SettingsRegistry() + settings = ProbeSettings(registry) + provider = _make_provider(pixel_width_m=0.0, pixel_height_m=0.0) + canned = _make_probe_seq(pixel_size_m=2e-6) + builder = _RecordingBuilder(settings, canned) + + item = ProbeRepositoryItem(_make_rng(), provider, settings, builder) + assert builder.build_calls == [] + + # Provider becomes ready (dataset would bind in production). + provider.get_probe_geometry.return_value = ProbeGeometry( + width_px=64, + height_px=64, + pixel_width_m=2e-6, + pixel_height_m=2e-6, + ) + # A settings change would normally re-fire _rebuild via the observer chain; + # set_builder is the shortest public path that triggers a rebuild. + replacement = _RecordingBuilder(settings, canned) + item.set_builder(replacement) + + assert len(replacement.build_calls) == 1 + assert item.get_probes().get_array() is canned.get_array() + + +def test_rebuild_skips_when_only_one_dimension_is_zero() -> None: + """The guard uses PixelGeometry.is_valid, which requires BOTH dims positive. + An asymmetric zero (e.g. width provided, height missing) still blocks the + rebuild — matches PixelGeometry.is_valid semantics.""" + registry = SettingsRegistry() + settings = ProbeSettings(registry) + provider = _make_provider(pixel_width_m=1e-6, pixel_height_m=0.0) + canned = _make_probe_seq(pixel_size_m=1e-6) + builder = _RecordingBuilder(settings, canned) + + item = ProbeRepositoryItem(_make_rng(), provider, settings, builder) + + assert builder.build_calls == [] + assert item.get_probes().get_array().size == 0 + + +class _ObservableProbeProvider(ProbeGeometryProvider, Observable): + """Test double: an Observable + ProbeGeometryProvider. Only get_probe_geometry + is exercised by ProbeRepositoryItem's rebuild guard; the other abstract + properties are stubbed with sensible defaults. set_geometry() mutates and + fires notify_observers, mimicking what ProductGeometry.set_detector_extent + does in production.""" + + def __init__(self, geometry: ProbeGeometry) -> None: + Observable.__init__(self) + self._geometry = geometry + + def set_geometry(self, geometry: ProbeGeometry) -> None: + self._geometry = geometry + self.notify_observers() + + @property + def detector_distance_m(self) -> float: + return 1.0 + + @property + def probe_photon_count(self) -> float: + return 1.0 + + @property + def probe_wavelength_m(self) -> float: + return 1e-10 + + @property + def probe_power_W(self) -> float: # noqa: N802 + return 1.0 + + @property + def num_scan_points(self) -> int: + return 1 + + def get_detector_pixel_geometry(self) -> PixelGeometry: + return PixelGeometry(width_m=1e-6, height_m=1e-6) + + def get_probe_geometry(self) -> ProbeGeometry: + return self._geometry + + +def test_rebuild_fires_on_geometry_observer_notification() -> None: + """When the geometry provider is Observable, ProbeRepositoryItem should + register itself and re-run _rebuild each time notify_observers fires + (matches the ProductGeometry.set_detector_extent path in production). + """ + registry = SettingsRegistry() + settings = ProbeSettings(registry) + provider = _ObservableProbeProvider( + ProbeGeometry(width_px=64, height_px=64, pixel_width_m=0.0, pixel_height_m=0.0), + ) + canned = _make_probe_seq(pixel_size_m=1e-6) + builder = _RecordingBuilder(settings, canned) + + item = ProbeRepositoryItem(_make_rng(), provider, settings, builder) + assert builder.build_calls == [] # guard blocks initial rebuild + + provider.set_geometry( + ProbeGeometry(width_px=64, height_px=64, pixel_width_m=1e-6, pixel_height_m=1e-6), + ) + + assert len(builder.build_calls) == 1 + assert item.get_probes().get_array() is canned.get_array() diff --git a/tests/test_probe_positions_builder.py b/tests/test_probe_positions_builder.py new file mode 100644 index 000000000..014eb0a86 --- /dev/null +++ b/tests/test_probe_positions_builder.py @@ -0,0 +1,211 @@ +"""Regression tests for the probe-position conditioning pipeline (trim -> affine -> jitter). + +Two invariants carry the weight here. + +First, trimming discards points by acquisition order but must leave the surviving +points' scan indexes untouched. AssembledDiffractionData.prepare_reconstruct_input +joins diffraction patterns to positions by index and refuses to extrapolate beyond +the position-index anchors, so renumbering the survivors would silently pair every +pattern with the wrong position. + +Second, FromMemoryProbePositionsBuilder must never condition. It holds positions +that are already conditioned -- reconstruction output, which ProcessingTaskMonitor +re-assigns to the output product item on every reconstructor iteration, and +products loaded from HDF5/NPZ. Re-applying the trim or the affine transform there +would corrupt position-corrected output a little more on every iteration. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy +import pytest + +from ptychodus.api.probe_positions import ( + ProbePosition, + ProbePositionFileReader, + ProbePositionSequence, +) +from ptychodus.api.settings import SettingsRegistry +from ptychodus.model.product.probe_positions.builder import ( + FromFileProbePositionsBuilder, + FromMemoryProbePositionsBuilder, +) +from ptychodus.model.product.probe_positions.cartesian import ( + CartesianProbePositionsBuilder, + CartesianProbePositionsVariant, +) +from ptychodus.model.product.probe_positions.settings import ProbePositionsSettings +from ptychodus.model.product.probe_positions.streaming import StreamingScanBuilder + + +def _make_settings() -> ProbePositionsSettings: + return ProbePositionsSettings(SettingsRegistry()) + + +def _make_rng() -> numpy.random.Generator: + return numpy.random.default_rng(42) + + +def _make_line(num_points: int) -> ProbePositionSequence: + """A horizontal line of positions whose index equals its ordinal.""" + return ProbePositionSequence( + [ + ProbePosition(index=idx, coordinate_x_m=float(idx), coordinate_y_m=0.0) + for idx in range(num_points) + ] + ) + + +class _StubPositionFileReader(ProbePositionFileReader): + def __init__(self, positions: ProbePositionSequence) -> None: + self._positions = positions + + def read(self, file_path: Path) -> ProbePositionSequence: + return self._positions + + +def _make_cartesian_line_builder( + settings: ProbePositionsSettings, num_points: int +) -> CartesianProbePositionsBuilder: + """A single-row raster, so acquisition order is unambiguous.""" + builder = CartesianProbePositionsBuilder( + CartesianProbePositionsVariant.RECTANGULAR_RASTER, _make_rng(), settings + ) + builder.num_points_x.set_value(num_points) + builder.num_points_y.set_value(1) + return builder + + +def test_slice_returns_subsequence() -> None: + """ProbePositionSequence.__getitem__ used to raise TypeError on any slice + with an implicit bound, because it built a range() from the raw slice + attributes. The trim needs slicing to work.""" + seq = _make_line(5) + + assert [p.index for p in seq[:2]] == [0, 1] + assert [p.index for p in seq[1:3]] == [1, 2] + assert [p.index for p in seq[-1:]] == [4] + assert len(seq[5:5]) == 0 + assert [p.coordinate_x_m for p in seq[1:3]] == [1.0, 2.0] + + +def test_generator_builder_trims_by_acquisition_order() -> None: + settings = _make_settings() + builder = _make_cartesian_line_builder(settings, 5) + builder.num_discard_at_start.set_value(1) + builder.num_discard_at_end.set_value(1) + + assert len(builder.build()) == 3 + + +def test_trim_preserves_original_indexes() -> None: + """The load-bearing invariant for prepare_reconstruct_input's index anchoring: + surviving points keep the indexes they were acquired with.""" + settings = _make_settings() + builder = _make_cartesian_line_builder(settings, 5) + builder.num_discard_at_start.set_value(1) + builder.num_discard_at_end.set_value(1) + + assert [p.index for p in builder.build()] == [1, 2, 3] + + +def test_over_trim_yields_empty_sequence_and_warns( + caplog: pytest.LogCaptureFixture, +) -> None: + settings = _make_settings() + builder = _make_cartesian_line_builder(settings, 5) + builder.num_discard_at_start.set_value(3) + builder.num_discard_at_end.set_value(3) + + with caplog.at_level('WARNING'): + positions = builder.build() + + assert len(positions) == 0 + assert 'leaves nothing of 5' in caplog.text + + +def test_from_file_builder_applies_affine_and_trim() -> None: + """FromFileProbePositionsBuilder.build() used to return the reader's output + verbatim, so the affine transform and jitter were silently ignored for every + file-loaded scan even though the GUI offered the transform editor.""" + settings = _make_settings() + reader = _StubPositionFileReader(_make_line(5)) + builder = FromFileProbePositionsBuilder(_make_rng(), settings, reader) + builder.affine00.set_value(-1.0) + builder.num_discard_at_start.set_value(1) + builder.num_discard_at_end.set_value(1) + + positions = builder.build() + + assert [p.index for p in positions] == [1, 2, 3] + assert [p.coordinate_x_m for p in positions] == [-1.0, -2.0, -3.0] + + +def test_from_memory_builder_ignores_conditioning() -> None: + """Guards reconstruction output: the from-memory builder must return its + positions verbatim no matter what the conditioning parameters say.""" + settings = _make_settings() + raw = _make_line(5) + builder = FromMemoryProbePositionsBuilder(_make_rng(), settings, raw) + builder.affine00.set_value(-1.0) + builder.num_discard_at_start.set_value(2) + builder.num_discard_at_end.set_value(2) + builder.jitter_radius_m.set_value(1e-6) + + positions = builder.build() + + assert [p.index for p in positions] == [0, 1, 2, 3, 4] + assert [p.coordinate_x_m for p in positions] == [0.0, 1.0, 2.0, 3.0, 4.0] + + +def test_repeated_from_memory_builds_are_idempotent() -> None: + """The reconstruct loop rebuilds the output item's positions once per + iteration; conditioning must not accumulate across those rebuilds.""" + settings = _make_settings() + settings.num_discard_at_start.set_value(1) + settings.affine00.set_value(-1.0) + + positions = _make_line(5) + + for _ in range(3): + builder = FromMemoryProbePositionsBuilder(_make_rng(), settings, positions) + positions = builder.build() + + assert [p.index for p in positions] == [0, 1, 2, 3, 4] + assert [p.coordinate_x_m for p in positions] == [0.0, 1.0, 2.0, 3.0, 4.0] + + +@pytest.mark.parametrize('builder_name', ['cartesian', 'from_file']) +def test_copy_preserves_trim_parameters(builder_name: str) -> None: + """copy() iterates parameters() generically, so the new counts should ride + along without any per-subclass change.""" + settings = _make_settings() + + if builder_name == 'cartesian': + builder = _make_cartesian_line_builder(settings, 5) + else: + reader = _StubPositionFileReader(_make_line(5)) + builder = FromFileProbePositionsBuilder(_make_rng(), settings, reader) # type: ignore[assignment] + + builder.num_discard_at_start.set_value(2) + builder.num_discard_at_end.set_value(3) + + duplicate = builder.copy() + + assert duplicate.num_discard_at_start.get_value() == 2 + assert duplicate.num_discard_at_end.get_value() == 3 + + +def test_streaming_builder_is_instantiable_and_copyable() -> None: + """StreamingScanBuilder never implemented the abstract copy(), so it could + not be instantiated at all.""" + settings = _make_settings() + builder = StreamingScanBuilder(_make_rng(), settings, _make_line(5)) + builder.num_discard_at_start.set_value(1) + + duplicate = builder.copy() + + assert duplicate.num_discard_at_start.get_value() == 1 + assert [p.index for p in builder.build()] == [1, 2, 3, 4] diff --git a/tests/test_product_core_reinit.py b/tests/test_product_core_reinit.py new file mode 100644 index 000000000..a5842fce9 --- /dev/null +++ b/tests/test_product_core_reinit.py @@ -0,0 +1,67 @@ +"""Tests for the settings-reinit wiring in ProductCore. + +The bug this guards against: on `-s settings.ini` startup, ProductCore._update +must pass the just-inserted diffraction dataset to insert_product_from_settings, +otherwise the settings-driven product ends up with no associated dataset and +reconstruction is refused with a "no associated diffraction dataset" warning. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from ptychodus.api.observer import Observable +from ptychodus.model.product.core import ProductCore + + +def _make_core( + *, repo_datasets: list, product_api: MagicMock, reinit_observable: Observable +) -> ProductCore: + """Build a ProductCore instance with only the fields _update reads set.""" + core = ProductCore.__new__(ProductCore) + diffraction_api = MagicMock() + diffraction_api.get_repository.return_value = repo_datasets + core._diffraction_api = diffraction_api + core.product_api = product_api + core._reinit_observable = reinit_observable + reinit_observable.add_observer(core) + return core + + +def test_update_passes_last_dataset_and_block_false_when_repo_nonempty() -> None: + reinit = Observable() + product_api = MagicMock() + dataset_a = MagicMock(name='dataset_a') + dataset_b = MagicMock(name='dataset_b') + _make_core( + repo_datasets=[dataset_a, dataset_b], + product_api=product_api, + reinit_observable=reinit, + ) + + reinit.notify_observers() + + product_api.insert_product_from_settings.assert_called_once_with(dataset=dataset_b, block=False) + + +def test_update_passes_dataset_none_when_repo_empty() -> None: + reinit = Observable() + product_api = MagicMock() + _make_core(repo_datasets=[], product_api=product_api, reinit_observable=reinit) + + reinit.notify_observers() + + product_api.insert_product_from_settings.assert_called_once_with(dataset=None, block=False) + + +def test_update_ignores_other_observables() -> None: + reinit = Observable() + other = Observable() + product_api = MagicMock() + core = _make_core( + repo_datasets=[MagicMock()], product_api=product_api, reinit_observable=reinit + ) + + core._update(other) + + product_api.insert_product_from_settings.assert_not_called() diff --git a/tests/test_product_creation_queue.py b/tests/test_product_creation_queue.py new file mode 100644 index 000000000..f5c825772 --- /dev/null +++ b/tests/test_product_creation_queue.py @@ -0,0 +1,181 @@ +"""Unit tests for the async product-creation queue in ProductAPI._insert_via_queue.""" + +from __future__ import annotations + +import threading +from unittest.mock import MagicMock + +import pytest + +from ptychodus.model.product.api import ProductAPI +from ptychodus.model.product.item import ProductState + + +class _StubTaskManager: + """Minimal TaskManager stand-in: records enqueued background tasks and + exposes the standard is_stopping / WAIT_TIME_S attributes ProductAPI reads.""" + + is_stopping = False + WAIT_TIME_S = 0.01 + + def __init__(self) -> None: + self.background_tasks: list = [] + + def put_background_task(self, task) -> None: # noqa: ANN001 + self.background_tasks.append(task) + + +def _make_api( + task_manager: _StubTaskManager, + stub_item: MagicMock, + real_item: MagicMock, +) -> tuple[ProductAPI, MagicMock, MagicMock]: + repository = MagicMock() + inserted: list = [] + + def insert_product(item): # noqa: ANN001 + inserted.append(item) + return len(inserted) - 1 + + repository.insert_product.side_effect = insert_product + + item_factory = MagicMock() + item_factory.create_pending_stub.return_value = stub_item + item_factory.create_from_values.return_value = real_item + item_factory.create_from_product.return_value = real_item + item_factory.create_from_settings.return_value = real_item + + api = ProductAPI( + settings=MagicMock(), + repository=repository, + item_factory=item_factory, + file_reader_chooser=MagicMock(), + file_writer_chooser=MagicMock(), + task_manager=task_manager, # type: ignore[arg-type] + ) + return api, repository, item_factory + + +def _dataset( + *, + in_progress: bool, + error: BaseException | None = None, + event_set: bool | None = None, +) -> MagicMock: + dataset = MagicMock() + dataset.is_load_in_progress.return_value = in_progress + dataset.get_last_load_error.return_value = error + event = threading.Event() + # By default, event is set iff the load is not in progress. Individual tests + # can pass event_set=True to simulate \"load just finished with an error\". + should_set = (not in_progress) if event_set is None else event_set + if should_set: + event.set() + dataset.get_last_load_finished_event.return_value = event + return dataset + + +def test_sync_path_when_dataset_is_none() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, repo, factory = _make_api(tm, stub, real) + + index = api.insert_new_product(dataset=None, block=False) + + assert index == 0 + factory.create_from_values.assert_called_once() + factory.create_pending_stub.assert_not_called() + repo.insert_product.assert_called_once_with(real) + assert tm.background_tasks == [] + + +def test_sync_path_when_dataset_already_loaded() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, repo, factory = _make_api(tm, stub, real) + + dataset = _dataset(in_progress=False) + index = api.insert_new_product(dataset=dataset, block=False) + + assert index == 0 + factory.create_from_values.assert_called_once() + factory.create_pending_stub.assert_not_called() + assert tm.background_tasks == [] + + +def test_enqueues_stub_when_mid_load_and_not_blocking() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, repo, factory = _make_api(tm, stub, real) + + dataset = _dataset(in_progress=True) + index = api.insert_new_product(name='p', dataset=dataset, block=False) + + assert index == 0 + factory.create_pending_stub.assert_called_once_with(name='p') + repo.insert_product.assert_called_once_with(stub) + assert len(tm.background_tasks) == 1 + factory.create_from_values.assert_not_called() + + +def test_background_task_finalizes_stub_on_success() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, repo, factory = _make_api(tm, stub, real) + + dataset = _dataset(in_progress=True) + api.insert_new_product(dataset=dataset, block=False) + + background_task = tm.background_tasks[0] + foreground_task = background_task() + foreground_task() + + factory.create_from_values.assert_called_once() + stub.copy_contents_from.assert_called_once_with(real) + stub.set_state.assert_called_once_with(ProductState.READY) + + +def test_background_task_marks_stub_failed_on_load_error() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, repo, factory = _make_api(tm, stub, real) + + err = RuntimeError('boom') + dataset = _dataset(in_progress=True, error=err) + api.insert_new_product(dataset=dataset, block=False) + + background_task = tm.background_tasks[0] + foreground_task = background_task() + foreground_task() + + factory.create_from_values.assert_not_called() + stub.set_state.assert_called_once_with(ProductState.FAILED) + + +def test_blocking_call_raises_on_load_error() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, _repo, _factory = _make_api(tm, stub, real) + + err = RuntimeError('boom') + # Simulate load that has just finished (event set) with an error stored. + dataset = _dataset(in_progress=True, error=err, event_set=True) + + with pytest.raises(RuntimeError): + api.insert_new_product(dataset=dataset, block=True) + + +def test_blocking_call_returns_index_after_wait() -> None: + tm = _StubTaskManager() + stub, real = MagicMock(), MagicMock() + api, repo, factory = _make_api(tm, stub, real) + + # in_progress=True but event already set (immediate wake), no error. + dataset = _dataset(in_progress=True, event_set=True) + + index = api.insert_new_product(dataset=dataset, block=True) + + assert index == 0 + factory.create_from_values.assert_called_once() + repo.insert_product.assert_called_once_with(real) + assert tm.background_tasks == [] diff --git a/tests/test_product_dataset.py b/tests/test_product_dataset.py new file mode 100644 index 000000000..ae758d3df --- /dev/null +++ b/tests/test_product_dataset.py @@ -0,0 +1,172 @@ +"""Unit tests for per-product diffraction dataset association.""" + +from __future__ import annotations + +from collections.abc import Sequence +from unittest.mock import MagicMock + +import pytest + +from ptychodus.model.processing.api import ProcessingAPI +from ptychodus.model.product.core import _DatasetOrphanObserver +from ptychodus.model.product.item import ProductRepositoryItem +from ptychodus.model.product.repository import ProductRepository +from ptychodus.model.product.item import ProductRepositoryObserver +from ptychodus.model.workflow import ConcreteWorkflowAPI + + +class _RecordingObserver(ProductRepositoryObserver): + def __init__(self) -> None: + self.dataset_changed: list[int] = [] + + def handle_item_inserted(self, index, item) -> None: # noqa: ANN001 + pass + + def handle_metadata_changed(self, index, item) -> None: # noqa: ANN001 + pass + + def handle_probe_positions_changed(self, index, item) -> None: # noqa: ANN001 + pass + + def handle_probe_changed(self, index, item) -> None: # noqa: ANN001 + pass + + def handle_object_changed(self, index, item) -> None: # noqa: ANN001 + pass + + def handle_losses_changed(self, index, losses) -> None: # noqa: ANN001 + pass + + def handle_dataset_changed(self, index, item) -> None: # noqa: ANN001 + self.dataset_changed.append(index) + + def handle_state_changed(self, index, item) -> None: # noqa: ANN001 + pass + + def handle_item_removed(self, index, item) -> None: # noqa: ANN001 + pass + + +def test_repository_fans_out_dataset_changed() -> None: + repo = ProductRepository() + observer = _RecordingObserver() + repo.add_observer(observer) + + item = MagicMock(spec=ProductRepositoryItem) + item._index = 2 + item.get_name.return_value = 'product' + + repo.handle_dataset_changed(item) + + assert observer.dataset_changed == [2] + + +def test_repository_ignores_dataset_changed_for_unregistered_item() -> None: + repo = ProductRepository() + observer = _RecordingObserver() + repo.add_observer(observer) + + item = MagicMock(spec=ProductRepositoryItem) + item._index = -1 + item.get_name.return_value = 'orphan' + + repo.handle_dataset_changed(item) + + assert observer.dataset_changed == [] + + +def test_orphan_observer_clears_only_matching_products() -> None: + dataset = MagicMock() + other_dataset = MagicMock() + + matching = MagicMock() + matching.get_dataset.return_value = dataset + unrelated = MagicMock() + unrelated.get_dataset.return_value = other_dataset + + product_repository: Sequence[MagicMock] = [matching, unrelated] + observer = _DatasetOrphanObserver(product_repository) # type: ignore[arg-type] + + observer.handle_dataset_removed(0, dataset) + + matching.unbind_dataset.assert_called_once_with() + unrelated.unbind_dataset.assert_not_called() + + +def _make_workflow_api(diffraction_repository: MagicMock) -> ConcreteWorkflowAPI: + diffraction_api = MagicMock() + diffraction_api.get_repository.return_value = diffraction_repository + return ConcreteWorkflowAPI( + MagicMock(), # settings_registry + diffraction_api, + MagicMock(), # product_api + MagicMock(), # probe_positions_api + MagicMock(), # probe_api + MagicMock(), # object_api + MagicMock(), # processing_api + MagicMock(), # fluorescence_api + MagicMock(), # globus_executor + MagicMock(), # genesis_executor + ) + + +def test_fetch_dataset_resolves_object_by_index() -> None: + dataset = MagicMock() + diffraction_repository = MagicMock() + diffraction_repository.__len__.return_value = 3 + diffraction_repository.__getitem__.return_value = dataset + + api = _make_workflow_api(diffraction_repository) + handle = MagicMock() + handle.get_dataset_index.return_value = 1 + + assert api._fetch_dataset(handle) is dataset + diffraction_repository.__getitem__.assert_called_once_with(1) + + +def test_fetch_dataset_clears_on_out_of_range() -> None: + diffraction_repository = MagicMock() + diffraction_repository.__len__.return_value = 2 + + api = _make_workflow_api(diffraction_repository) + handle = MagicMock() + handle.get_dataset_index.return_value = 5 + + assert api._fetch_dataset(handle) is None + + +def test_fetch_dataset_returns_none_without_handle() -> None: + api = _make_workflow_api(MagicMock()) + + assert api._fetch_dataset(None) is None + + +def _make_processing_api(item: MagicMock) -> ProcessingAPI: + product_api = MagicMock() + product_api.get_item.return_value = item + return ProcessingAPI(MagicMock(), product_api, MagicMock(), MagicMock()) + + +def test_get_reconstruct_input_raises_without_dataset() -> None: + item = MagicMock() + item.get_dataset.return_value = None + item.get_name.return_value = 'product' + api = _make_processing_api(item) + + with pytest.raises(RuntimeError): + api.get_reconstruct_input(product_index=0) + + +def test_get_reconstruct_input_uses_product_dataset() -> None: + dataset = MagicMock() + product = MagicMock() + item = MagicMock() + item.get_dataset.return_value = dataset + item.get_product.return_value = product + api = _make_processing_api(item) + + result = api.get_reconstruct_input(product_index=0) + + assembled = dataset.get_assembled_data.return_value + assembled.prepare_reconstruct_input.assert_called_once() + assert result is assembled.prepare_reconstruct_input.return_value diff --git a/tests/test_product_item_copy_contents.py b/tests/test_product_item_copy_contents.py new file mode 100644 index 000000000..4083149b0 --- /dev/null +++ b/tests/test_product_item_copy_contents.py @@ -0,0 +1,74 @@ +"""Tests for ProductRepositoryItem.copy_contents_from ordering. + +Guards against the regression where the stub's probe/object subgroups were +assigned before the dataset was bound: their _rebuild() saw an invalid pixel +geometry (detector_extent still None) and silently no-op'd, leaving the +finalized product with empty probe/object arrays. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from ptychodus.model.product.item import ProductRepositoryItem, ProductState + + +def _make_source_and_stub() -> tuple[ProductRepositoryItem, ProductRepositoryItem, MagicMock]: + parent = MagicMock() + + # Bypass __init__ so we don't have to satisfy every dependency of ProductGeometry. + stub = ProductRepositoryItem.__new__(ProductRepositoryItem) + stub._parent = parent + stub._metadata_item = MagicMock() + stub._probe_positions_item = MagicMock() + stub._probe_item = MagicMock() + stub._object_item = MagicMock() + stub._geometry = MagicMock() + stub._losses = [] + stub._dataset = None + stub._state = ProductState.PENDING + + source = ProductRepositoryItem.__new__(ProductRepositoryItem) + source._parent = MagicMock() + source._metadata_item = MagicMock() + source._probe_positions_item = MagicMock() + source._probe_item = MagicMock() + source._object_item = MagicMock() + source._geometry = MagicMock() + source._losses = ['loss-value'] + source._dataset = MagicMock() + + return source, stub, parent + + +def test_copy_contents_from_binds_dataset_before_assigning_probe_and_object() -> None: + source, stub, _parent = _make_source_and_stub() + + manager = MagicMock() + manager.attach_mock(stub._geometry.set_detector_extent, 'set_detector_extent') + manager.attach_mock(stub._probe_item.assign_item, 'probe_assign') + manager.attach_mock(stub._object_item.assign_item, 'object_assign') + manager.attach_mock(stub._probe_positions_item.assign_item, 'positions_assign') + + stub.copy_contents_from(source) + + call_names = [call[0] for call in manager.mock_calls] + # set_detector_extent must precede both probe and object assign_item. + assert call_names.index('set_detector_extent') < call_names.index('probe_assign') + assert call_names.index('set_detector_extent') < call_names.index('object_assign') + + +def test_copy_contents_from_copies_all_state() -> None: + source, stub, parent = _make_source_and_stub() + + stub.copy_contents_from(source) + + stub._metadata_item.assign.assert_called_once_with( + source._metadata_item.get_metadata.return_value + ) + stub._probe_positions_item.assign_item.assert_called_once_with(source._probe_positions_item) + stub._probe_item.assign_item.assert_called_once_with(source._probe_item) + stub._object_item.assign_item.assert_called_once_with(source._object_item) + assert stub._losses == source._losses + assert stub._dataset is source._dataset + parent.handle_losses_changed.assert_called_once_with(stub) diff --git a/tests/test_ptychi_options.py b/tests/test_ptychi_options.py index d49ecc953..8854362a6 100644 --- a/tests/test_ptychi_options.py +++ b/tests/test_ptychi_options.py @@ -27,6 +27,19 @@ from ptychodus.api.settings import SettingsRegistry from ptychodus.model.ptychi.core import PtyChiReconstructorLibrary +# Concrete (child-side) reconstructor classes plus the options helper. This test +# is an in-process validator that must reach into pty-chi's Pydantic constructors +# directly, so it deliberately does the imports the parent avoids. This is safe +# because the file already does ``pytest.importorskip('ptychi')`` above. +from ptychodus.model.ptychi.autodiff import AutodiffReconstructor # noqa: E402 +from ptychodus.model.ptychi.bh import BHReconstructor # noqa: E402 +from ptychodus.model.ptychi.dm import DMReconstructor # noqa: E402 +from ptychodus.model.ptychi.epie import EPIEReconstructor # noqa: E402 +from ptychodus.model.ptychi.helper import PtyChiOptionsHelper # noqa: E402 +from ptychodus.model.ptychi.lsqml import LSQMLReconstructor # noqa: E402 +from ptychodus.model.ptychi.pie import PIEReconstructor # noqa: E402 +from ptychodus.model.ptychi.rpie import RPIEReconstructor # noqa: E402 + PIXEL_M = 1.0e-9 OBJ_HEIGHT_PX = 32 OBJ_WIDTH_PX = 40 @@ -38,8 +51,8 @@ class _StubPatternSizer: """Minimal stand-in — the options helper only reads the processed pixel geometry.""" - def get_processed_pixel_geometry(self) -> PixelGeometry: - return PixelGeometry(width_m=1.0e-6, height_m=1.0e-6) + def get_processed_pixel_geometry(self, raw_pixel_geometry: PixelGeometry) -> PixelGeometry: + return raw_pixel_geometry def _make_reconstruct_input() -> ReconstructInput: @@ -85,7 +98,12 @@ def _make_reconstruct_input() -> ReconstructInput: ) patterns = rng.random((NUM_PATTERNS, PROBE_HEIGHT_PX, PROBE_WIDTH_PX)).astype(numpy.float32) bad_pixels = numpy.zeros((PROBE_HEIGHT_PX, PROBE_WIDTH_PX), dtype=numpy.bool_) - return ReconstructInput(diffraction_patterns=patterns, bad_pixels=bad_pixels, product=product) + return ReconstructInput( + diffraction_patterns=patterns, + bad_pixels=bad_pixels, + product=product, + pixel_geometry=PixelGeometry(width_m=1.0e-6, height_m=1.0e-6), + ) def _make_library() -> PtyChiReconstructorLibrary: @@ -96,8 +114,35 @@ def _make_library() -> PtyChiReconstructorLibrary: ) +def _make_concrete_reconstructors(library: PtyChiReconstructorLibrary) -> list: + """Instantiate the child-side algorithm classes directly against ``library``'s settings. + + The parent-side ``library.reconstructor_list`` now holds + :class:`SubprocessReconstructor` shells that hide ``_create_task_options`` + behind a process boundary. To validate defaults / bounds we need the + concrete classes, so build them here using the same wiring the child does. + """ + helper = PtyChiOptionsHelper( + library.settings, + library.object_settings, + library.probe_settings, + library.probe_position_settings, + library.opr_settings, + _StubPatternSizer(), # type: ignore[arg-type] + ) + return [ + DMReconstructor(helper, library.dm_settings), + PIEReconstructor(helper, library.pie_settings), + EPIEReconstructor(helper, library.pie_settings), + RPIEReconstructor(helper, library.pie_settings), + LSQMLReconstructor(helper, library.lsqml_settings), + AutodiffReconstructor(helper, library.autodiff_settings), + BHReconstructor(helper, library.bh_settings), + ] + + def _build_all_task_options(library: PtyChiReconstructorLibrary, parameters: ReconstructInput): - for reconstructor in library.reconstructor_list: + for reconstructor in _make_concrete_reconstructors(library): # Every ptychi reconstructor exposes ``_create_task_options``; building it # runs pty-chi's Pydantic validators over all sub-option objects. reconstructor._create_task_options(parameters) # type: ignore[attr-defined] @@ -158,7 +203,7 @@ def test_hard_limits_are_serialized_as_lists() -> None: obj.constrain_hard_limits_enable_abs.set_value(True) obj.constrain_hard_limits_enable_phase.set_value(True) - options = library.reconstructor_list[0]._create_task_options( # type: ignore[attr-defined] + options = _make_concrete_reconstructors(library)[0]._create_task_options( # type: ignore[attr-defined] _make_reconstruct_input() ) hard_limits = options.object_options.hard_limits_magnitude_phase @@ -175,7 +220,7 @@ def test_compact_mode_clustering_stride_is_at_least_one() -> None: """Disabled compact-mode clustering must still yield a stride >= 1 for pty-chi.""" library = _make_library() # Default (disabled) value is 0; pty-chi's stride field is now ge=1. - options = library.reconstructor_list[0]._create_task_options( # type: ignore[attr-defined] + options = _make_concrete_reconstructors(library)[0]._create_task_options( # type: ignore[attr-defined] _make_reconstruct_input() ) assert options.reconstructor_options.compact_mode_update_clustering_stride >= 1 diff --git a/tests/test_subprocess_reconstructor.py b/tests/test_subprocess_reconstructor.py new file mode 100644 index 000000000..6afe3c0d0 --- /dev/null +++ b/tests/test_subprocess_reconstructor.py @@ -0,0 +1,239 @@ +"""Unit tests for the generic :class:`SubprocessReconstructor` adapter. + +These tests spawn real subprocesses (so they exercise the pickling, log +forwarding, sentinel/error paths, and cleanup) but use fake entry points that +touch no GPU framework. See ``tests/subprocess_child_fixtures.py``. +""" + +from __future__ import annotations + +import logging +import sys +from pathlib import Path +from typing import Any + +import numpy +import pytest + +from ptychodus.api.geometry import PixelGeometry +from ptychodus.api.object import Object, ObjectCenter +from ptychodus.api.probe import ProbeSequence +from ptychodus.api.probe_positions import ProbePositionSequence +from ptychodus.api.product import Product, ProductMetadata +from ptychodus.api.reconstructor import ReconstructInput +from ptychodus.model.processing._subprocess_protocol import ChildError +from ptychodus.model.processing.subprocess_reconstructor import ( + SubprocessReconstructor, +) + + +# The spawned child resolves the fixtures module by dotted path against the +# sys.path it inherits from this process, so the tests directory must be on it +# before any child is spawned. pytest's default import mode already does this; +# the insert keeps it true under other invocations. +_TESTS_DIR = str(Path(__file__).parent) +if _TESTS_DIR not in sys.path: + sys.path.insert(0, _TESTS_DIR) + +FIXTURES_MODULE = 'subprocess_child_fixtures' + + +def _minimal_product() -> Product: + metadata = ProductMetadata( + name='fake', + comments='', + detector_distance_m=1.0, + probe_energy_eV=10_000.0, + probe_photon_count=1.0, + exposure_time_s=1.0, + mass_attenuation_m2_kg=0.0, + tomography_angle_deg=0.0, + ) + pixel_geometry = PixelGeometry(1.0e-9, 1.0e-9) + center = ObjectCenter(0.0, 0.0) + return Product( + metadata=metadata, + probe_positions=ProbePositionSequence([]), + probes=ProbeSequence( + array=numpy.zeros((1, 1, 4, 4), dtype=numpy.complex64), + opr_weights=None, + pixel_geometry=pixel_geometry, + ), + object_=Object( + array=numpy.zeros((1, 4, 4), dtype=numpy.complex64), + layer_spacing_m=[], + pixel_geometry=pixel_geometry, + center=center, + ), + losses=[], + ) + + +def _minimal_reconstruct_input() -> ReconstructInput: + product = _minimal_product() + return ReconstructInput( + diffraction_patterns=numpy.zeros((1, 4, 4), dtype=numpy.float32), + bad_pixels=numpy.zeros((4, 4), dtype=numpy.bool_), + product=product, + pixel_geometry=PixelGeometry(1.0e-9, 1.0e-9), + ) + + +def test_reconstruct_streams_all_outputs_in_order() -> None: + product = _minimal_product() + + def build_payload(parameters: ReconstructInput, _loaded: Path | None) -> Any: + return {'n': 3, 'product': product} + + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:yield_n_outputs', + progress_goal_fn=lambda: 3, + build_reconstruct_payload=build_payload, + ) + + outputs = list(adapter.reconstruct(_minimal_reconstruct_input())) + + assert len(outputs) == 3 + assert [o.progress for o in outputs] == [1, 2, 3] + + +def test_child_exception_propagates_as_child_error() -> None: + def build_payload(parameters: ReconstructInput, _loaded: Path | None) -> Any: + return {'message': 'boom'} + + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:raise_immediately', + progress_goal_fn=lambda: 0, + build_reconstruct_payload=build_payload, + ) + + with pytest.raises(ChildError) as excinfo: + list(adapter.reconstruct(_minimal_reconstruct_input())) + + assert excinfo.value.child_exception_type == 'ValueError' + # The original exception is picklable, so the parent got it back. + assert isinstance(excinfo.value.child_exception, ValueError) + assert 'boom' in str(excinfo.value.child_exception) + + +def test_hanging_child_is_terminated_on_iterator_close() -> None: + def build_payload(parameters: ReconstructInput, _loaded: Path | None) -> Any: + return {} + + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:hang_forever', + progress_goal_fn=lambda: 0, + build_reconstruct_payload=build_payload, + terminate_grace_sec=1.0, + ) + + # Get one message (there is none coming) - close the iterator to trigger cleanup. + iterator = adapter.reconstruct(_minimal_reconstruct_input()) + # Close immediately; the context manager's finally must terminate the child. + iterator.close() + + +def test_child_log_is_forwarded_to_parent_logger(caplog: pytest.LogCaptureFixture) -> None: + product = _minimal_product() + + def build_payload(parameters: ReconstructInput, _loaded: Path | None) -> Any: + return {'product': product, 'log_message': 'hello from the child'} + + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:emit_log_then_output', + progress_goal_fn=lambda: 1, + build_reconstruct_payload=build_payload, + ) + + with caplog.at_level(logging.WARNING, logger=FIXTURES_MODULE): + outputs = list(adapter.reconstruct(_minimal_reconstruct_input())) + + assert len(outputs) == 1 + assert any('hello from the child' in rec.message for rec in caplog.records) + + +def test_settings_sync_message_invokes_callback() -> None: + product = _minimal_product() + seen: list[dict[str, dict[str, str]]] = [] + + def build_payload(parameters: ReconstructInput, _loaded: Path | None) -> Any: + return {'product': product, 'settings': {'GroupA': {'p': '42'}}} + + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:emit_settings_sync_then_output', + progress_goal_fn=lambda: 1, + build_reconstruct_payload=build_payload, + apply_settings_sync=seen.append, + ) + + outputs = list(adapter.reconstruct(_minimal_reconstruct_input())) + + assert len(outputs) == 1 + assert seen == [{'GroupA': {'p': '42'}}] + + +def test_train_records_model_saved_path(tmp_path: Path) -> None: + saved_path = tmp_path / 'ckpt.bin' + saved_path.write_bytes(b'x') + + def build_train_payload(input_path: Path, output_path: Path) -> Any: + return {'saved_path': str(saved_path)} + + def build_reconstruct_payload(_p: ReconstructInput, _loaded: Path | None) -> Any: + return {} + + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:hang_forever', + progress_goal_fn=lambda: 1, + build_reconstruct_payload=build_reconstruct_payload, + is_trainable=True, + train_entry_point=f'{FIXTURES_MODULE}:train_and_save', + build_train_payload=build_train_payload, + model_file_extension='.bin', + ) + + assert not adapter.is_model_loaded() + + outputs = list(adapter.train(tmp_path, tmp_path)) + + assert len(outputs) == 1 + assert outputs[0].progress == 1 + assert adapter.is_model_loaded() + + # save_model must copy the recorded path. + dest = tmp_path / 'copied.bin' + adapter.save_model(dest) + assert dest.read_bytes() == b'x' + + +def test_non_trainable_train_raises() -> None: + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:hang_forever', + progress_goal_fn=lambda: 1, + build_reconstruct_payload=lambda p, m: {}, + ) + + with pytest.raises(NotImplementedError): + list(adapter.train(Path('/tmp'), Path('/tmp'))) + + +def test_save_without_load_or_train_raises() -> None: + adapter = SubprocessReconstructor( + name='FAKE', + reconstruct_entry_point=f'{FIXTURES_MODULE}:hang_forever', + progress_goal_fn=lambda: 1, + build_reconstruct_payload=lambda p, m: {}, + is_trainable=True, + train_entry_point=f'{FIXTURES_MODULE}:train_and_save', + build_train_payload=lambda i, o: {}, + ) + + with pytest.raises(RuntimeError): + adapter.save_model(Path('/tmp/x')) diff --git a/tests/test_visualization.py b/tests/test_visualization.py index 93d24bcdc..9c761e5bc 100644 --- a/tests/test_visualization.py +++ b/tests/test_visualization.py @@ -10,6 +10,7 @@ from ptychodus.api.visualization import ( ComplexComponent, CylindricalColorModel, + DisplayValues, KernelDensityEstimate, LineCut, ScalarTransformation, @@ -83,6 +84,37 @@ def test_rejects_shape_mismatch(self): with pytest.raises(ValueError, match='Shape mismatch'): VisualizationProduct('x', values, rgba, _pixel_geo(), Interval[float](0.0, 1.0)) + def test_rejects_1d_display_values(self): + values = numpy.ones((4, 4)) + rgba = numpy.ones((4, 4, 4)) + display_values = [DisplayValues('bad', numpy.ones(4))] + with pytest.raises(ValueError, match='2-dimensional'): + VisualizationProduct( + 'x', values, rgba, _pixel_geo(), Interval[float](0.0, 1.0), display_values + ) + + def test_rejects_display_values_shape_mismatch(self): + values = numpy.ones((4, 4)) + rgba = numpy.ones((4, 4, 4)) + display_values = [DisplayValues('bad', numpy.ones((3, 4)))] + with pytest.raises(ValueError, match='Shape mismatch'): + VisualizationProduct( + 'x', values, rgba, _pixel_geo(), Interval[float](0.0, 1.0), display_values + ) + + def test_default_display_values_for_real_input(self): + vp = _make_product() + (dv,) = vp.get_display_values() + assert dv.label == 'test' + numpy.testing.assert_array_equal(dv.values, vp.get_values()) + + def test_default_display_values_for_complex_input(self): + values = numpy.ones((4, 4), dtype=complex) * (3 + 4j) + rgba = numpy.ones((4, 4, 4), dtype=numpy.float32) + vp = VisualizationProduct('c', values, rgba, _pixel_geo(), Interval[float](0.0, 1.0)) + (dv,) = vp.get_display_values() + numpy.testing.assert_allclose(dv.values, 5.0) + # --------------------------------------------------------------------------- # VisualizationProduct accessors @@ -130,6 +162,16 @@ def test_complex_value(self): assert 'amplitude=' in text assert 'phase=' in text + def test_complex_value_with_zero_imaginary_part(self): + # A complex-dtype pixel whose imaginary part is exactly zero must still take the + # amplitude/phase branch; formatting a complex with '6g' raises TypeError. + values = numpy.zeros((4, 4), dtype=complex) + rgba = numpy.ones((4, 4, 4), dtype=numpy.float32) + vp = VisualizationProduct('c', values, rgba, _pixel_geo(), Interval[float](0.0, 1.0)) + text = vp.get_info_text(1.0, 1.0) + assert 'amplitude=' in text + assert 'phase=' in text + def test_clamps_negative_coords(self): vp = _make_product() text = vp.get_info_text(-5.0, -5.0) @@ -155,6 +197,14 @@ def test_returns_linecut(self): lc = vp.get_line_cut(line) assert isinstance(lc, LineCut) + def test_single_series_for_real_values(self): + vp = self._make_gradient() + line = Line2D(Point2D(0.0, 0.0), Point2D(3.0, 0.0)) + lc = vp.get_line_cut(line) + assert len(lc.series) == 1 + assert lc.series[0].label == 'grad' + assert len(lc.series[0].value) == len(lc.distance_m) + def test_distances_are_nonneg(self): vp = self._make_gradient() line = Line2D(Point2D(0.0, 0.0), Point2D(3.0, 0.0)) @@ -179,6 +229,83 @@ def test_zero_length_line(self): assert isinstance(lc, LineCut) +# --------------------------------------------------------------------------- +# VisualizationProduct.get_line_cut over complex arrays +# --------------------------------------------------------------------------- + + +def _complex_gradient() -> numpy.ndarray: + """4x4 complex array with varying amplitude and phase, and no wrapped-phase ambiguity.""" + amplitude = numpy.arange(1.0, 17.0, dtype=numpy.float32).reshape(4, 4) + phase_rad = numpy.linspace(-1.5, 1.5, 16, dtype=numpy.float32).reshape(4, 4) + return amplitude * numpy.exp(1j * phase_rad) + + +# A horizontal line across the top row samples pixels (0, 0..3) exactly once each. +_TOP_ROW_LINE = Line2D(Point2D(0.0, 0.5), Point2D(4.0, 0.5)) + + +class TestGetLineCutComplex: + @pytest.mark.parametrize('component', list(ComplexComponent)) + def test_matches_selected_component(self, component: ComplexComponent): + # Regression: the line cut used to sample the original complex array, so every + # component silently plotted the real part. + arr = _complex_gradient() + vp = visualize_complex_component(arr, _pixel_geo(), component) + lc = vp.get_line_cut(_TOP_ROW_LINE) + expected = component.extract_component(arr)[0, :] + assert len(lc.series) == 1 + numpy.testing.assert_allclose(lc.series[0].value, expected, rtol=1e-6) + + @pytest.mark.parametrize('component', list(ComplexComponent)) + def test_series_label_matches_value_label(self, component: ComplexComponent): + arr = _complex_gradient() + vp = visualize_complex_component(arr, _pixel_geo(), component) + lc = vp.get_line_cut(_TOP_ROW_LINE) + assert lc.series[0].label == vp.get_value_label() + + def test_components_differ_from_each_other(self): + arr = _complex_gradient() + cuts = { + component: tuple( + visualize_complex_component(arr, _pixel_geo(), component) + .get_line_cut(_TOP_ROW_LINE) + .series[0] + .value + ) + for component in ComplexComponent + } + # REAL and IMAGINARY must not coincide; if they did, the sampling would be ignoring + # the component selection entirely. + assert cuts[ComplexComponent.REAL] != cuts[ComplexComponent.IMAGINARY] + assert cuts[ComplexComponent.AMPLITUDE] != cuts[ComplexComponent.REAL] + + def test_applies_scalar_transform(self): + arr = _complex_gradient() + vp = visualize_complex_component( + arr, _pixel_geo(), ComplexComponent.AMPLITUDE, transform=ScalarTransformation.SQRT + ) + lc = vp.get_line_cut(_TOP_ROW_LINE) + expected = numpy.sqrt(ComplexComponent.AMPLITUDE.extract_component(arr)[0, :]) + numpy.testing.assert_allclose(lc.series[0].value, expected, rtol=1e-6) + + def test_cylindrical_model_yields_amplitude_and_phase(self): + arr = _complex_gradient() + vp = visualize_complex_values( + arr, + _pixel_geo(), + CylindricalColorModel.HSV_VALUE, + amplitude_transform=ScalarTransformation.IDENTITY, + ) + lc = vp.get_line_cut(_TOP_ROW_LINE) + assert len(lc.series) == 2 + amplitude_series, phase_series = lc.series + assert amplitude_series.label == 'Amplitude' + assert phase_series.label == 'Phase [rad]' + numpy.testing.assert_allclose(amplitude_series.value, numpy.absolute(arr)[0, :], rtol=1e-6) + numpy.testing.assert_allclose(phase_series.value, numpy.angle(arr)[0, :], rtol=1e-5) + + # --------------------------------------------------------------------------- # VisualizationProduct.estimate_kernel_density # --------------------------------------------------------------------------- @@ -215,6 +342,36 @@ def test_box_clamped_to_image(self): kde = vp.estimate_kernel_density(box) assert kde.value_lower <= kde.value_upper + def test_phase_component_spans_phase_range(self): + # Regression: the histogram used to hardcode amplitude for any complex array, so a + # phase histogram silently showed amplitudes. + rng = numpy.random.default_rng(11) + amplitude = rng.random((8, 8)) + 10.0 + phase_rad = rng.uniform(-numpy.pi, numpy.pi, (8, 8)) + arr = amplitude * numpy.exp(1j * phase_rad) + vp = visualize_complex_component(arr, _pixel_geo(), ComplexComponent.PHASE_RAD) + box = Box2D(x=0, y=0, width=8, height=8) + kde = vp.estimate_kernel_density(box) + assert kde.value_lower == pytest.approx(phase_rad.min(), rel=1e-4) + assert kde.value_upper == pytest.approx(phase_rad.max(), rel=1e-4) + + def test_cylindrical_model_uses_amplitude(self): + # A histogram has one value axis, so the Complex renderer keeps showing amplitude. + rng = numpy.random.default_rng(13) + amplitude = rng.random((8, 8)) + 1.0 + phase_rad = rng.uniform(-numpy.pi, numpy.pi, (8, 8)) + arr = amplitude * numpy.exp(1j * phase_rad) + vp = visualize_complex_values( + arr, + _pixel_geo(), + CylindricalColorModel.HSV_VALUE, + amplitude_transform=ScalarTransformation.IDENTITY, + ) + box = Box2D(x=0, y=0, width=8, height=8) + kde = vp.estimate_kernel_density(box) + assert kde.value_lower == pytest.approx(amplitude.min(), rel=1e-4) + assert kde.value_upper == pytest.approx(amplitude.max(), rel=1e-4) + # --------------------------------------------------------------------------- # ComplexComponent @@ -457,6 +614,15 @@ def test_values_preserved(self): vp = visualize_real_values('I', values, _pixel_geo()) numpy.testing.assert_array_equal(vp.get_values(), values) + def test_display_values_are_transformed(self): + # The value label is decorated with the transform, so the displayed values must be + # transformed too or the axis label lies. + values = numpy.arange(1, 17, dtype=numpy.float32).reshape(4, 4) + vp = visualize_real_values('I', values, _pixel_geo(), transform=ScalarTransformation.SQRT) + (dv,) = vp.get_display_values() + assert dv.label == vp.get_value_label() + numpy.testing.assert_allclose(dv.values, numpy.sqrt(values), rtol=1e-6) + def test_with_transform(self): values = numpy.array([[1.0, 4.0], [9.0, 16.0]], dtype=numpy.float32) vp = visualize_real_values('I', values, _pixel_geo(), transform=ScalarTransformation.SQRT) diff --git a/tests/test_visualization_parameters.py b/tests/test_visualization_parameters.py new file mode 100644 index 000000000..bef312fe0 --- /dev/null +++ b/tests/test_visualization_parameters.py @@ -0,0 +1,116 @@ +"""Tests for the model/visualization parameters built on PluginChooserParameter. + +ScalarTransformationParameter, ColormapParameter, and +CylindricalColorModelParameter were each a hand-written copy of the same +chooser-to-Parameter adapter, and each silently dropped the ``notify`` keyword +that ``Parameter.set_value`` declares. They are now thin subclasses of +PluginChooserParameter, so these tests pin the contract they used to break plus +the accessors callers rely on. + +No Qt is required — this is pure model-layer behavior. +""" + +from __future__ import annotations + +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.visualization import CylindricalColorModel, ScalarTransformation +from ptychodus.model.visualization.color_model import CylindricalColorModelParameter +from ptychodus.model.visualization.colormap import ColormapParameter +from ptychodus.model.visualization.transformation import ScalarTransformationParameter + + +class _Counter(Observer): + def __init__(self) -> None: + self.count = 0 + + def _update(self, observable: Observable) -> None: + self.count += 1 + + +def test_transformation_defaults_to_identity() -> None: + parameter = ScalarTransformationParameter() + + assert parameter.get_value() == 'Identity' + assert parameter.get_strategy() is ScalarTransformation.IDENTITY + + +def test_transformation_set_value_honors_notify_false() -> None: + parameter = ScalarTransformationParameter() + counter = _Counter() + parameter.add_observer(counter) + + parameter.set_value('Square Root', notify=False) + + assert parameter.get_value() == 'Square Root' + assert parameter.get_strategy() is ScalarTransformation.SQRT + assert counter.count == 0 + + # The suppression must not leak into the next assignment. + parameter.set_value('Natural Logarithm') + assert counter.count == 1 + + +def test_transformation_resolves_simple_names() -> None: + """Simple names differ from display names here ('log2' vs 'Logarithm (Base 2)').""" + parameter = ScalarTransformationParameter() + + parameter.set_value('log2') + + assert parameter.get_value() == 'Logarithm (Base 2)' + assert parameter.get_strategy() is ScalarTransformation.LOG2 + + +def test_transformation_copy_is_independent() -> None: + parameter = ScalarTransformationParameter() + parameter.set_value('Square Root') + + copied = parameter.copy() + copied.set_value('Identity') + + assert parameter.get_value() == 'Square Root' + assert copied.get_value() == 'Identity' + + +def test_colormap_defaults_by_cyclicity() -> None: + assert ColormapParameter(is_cyclic=False).get_value() == 'gray' + assert ColormapParameter(is_cyclic=True).get_value() == 'colorwheel' + + +def test_colormap_choices_are_non_empty_and_contain_the_default() -> None: + parameter = ColormapParameter(is_cyclic=False) + + choices = list(parameter.choices()) + + assert 'gray' in choices + assert len(choices) > 1 + + +def test_colormap_copy_preserves_cyclicity() -> None: + parameter = ColormapParameter(is_cyclic=True) + + copied = parameter.copy() + + assert copied.get_value() == 'colorwheel' + + +def test_color_model_default_resolves_the_simple_name() -> None: + """The default is given as 'HSV-V', a simple name; the value space is display names.""" + parameter = CylindricalColorModelParameter() + + assert parameter.get_strategy() is CylindricalColorModel.HSV_VALUE + assert parameter.get_value() == 'HSV Value' + + +def test_color_model_set_value_honors_notify_false() -> None: + parameter = CylindricalColorModelParameter() + counter = _Counter() + parameter.add_observer(counter) + + parameter.set_value('HLS Lightness', notify=False) + + assert parameter.get_strategy() is CylindricalColorModel.HLS_LIGHTNESS + assert counter.count == 0 + + # The suppression must not leak into the next assignment. + parameter.set_value('HSV Alpha') + assert counter.count == 1 diff --git a/tests/view/conftest.py b/tests/view/conftest.py new file mode 100644 index 000000000..4d8a62bdc --- /dev/null +++ b/tests/view/conftest.py @@ -0,0 +1,20 @@ +"""Shared PyQt5 fixtures for widget tests. + +The outer `tests/conftest.py` drops this whole subtree when PyQt5 is missing, +so importing Qt at module top-level is safe here. +""" + +from __future__ import annotations + +import pytest + +from PyQt5.QtWidgets import QApplication + + +@pytest.fixture(scope='session') +def qapp() -> QApplication: + """A single QApplication shared across all widget tests in the session.""" + app = QApplication.instance() + if app is None: + app = QApplication([]) + return app diff --git a/tests/view/widgets/test_decimal_range_slider.py b/tests/view/widgets/test_decimal_range_slider.py new file mode 100644 index 000000000..0c110d7b9 --- /dev/null +++ b/tests/view/widgets/test_decimal_range_slider.py @@ -0,0 +1,226 @@ +from __future__ import annotations +from decimal import Decimal + +import pytest + +from PyQt5.QtCore import QPoint, Qt +from PyQt5.QtTest import QTest + +from ptychodus.api.geometry import Interval +from ptychodus.view.widgets import DecimalRangeSlider, Handle + + +def _iv(lower, upper) -> Interval[Decimal]: + return Interval[Decimal](Decimal(str(lower)), Decimal(str(upper))) + + +@pytest.fixture +def slider(qapp) -> DecimalRangeSlider: + del qapp + widget = DecimalRangeSlider.create_instance(Qt.Orientation.Horizontal) + widget.resize(400, 40) + widget.show() + return widget + + +def _capture(widget: DecimalRangeSlider) -> list[Interval]: + emissions: list[Interval] = [] + widget.selection_changed.connect(emissions.append) + return emissions + + +def _endpoints(interval: Interval) -> tuple[Decimal, Decimal]: + return interval.lower, interval.upper + + +def test_default_state(slider: DecimalRangeSlider) -> None: + assert _endpoints(slider.get_bounds()) == (Decimal(0), Decimal(1)) + assert _endpoints(slider.get_selection()) == (Decimal(0), Decimal(1)) + + +def test_default_construction_emits_no_signal(qapp) -> None: + del qapp + emissions: list[Interval] = [] + widget = DecimalRangeSlider.create_instance(Qt.Orientation.Horizontal) + widget.selection_changed.connect(emissions.append) + assert emissions == [] + + +def test_only_horizontal_supported(qapp) -> None: + del qapp + with pytest.raises(NotImplementedError): + DecimalRangeSlider.create_instance(Qt.Orientation.Vertical) + + +def test_set_selection_clamps_into_bounds(slider: DecimalRangeSlider) -> None: + emissions = _capture(slider) + slider.set_selection(_iv(-5, 5)) + assert _endpoints(slider.get_selection()) == (Decimal(0), Decimal(1)) + assert len(emissions) == 0 # selection was already (0, 1); clamped result matches + + +def test_set_selection_clamped_change_emits_once(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.3', '0.7')) + emissions = _capture(slider) + slider.set_selection(_iv(-5, 5)) + assert len(emissions) == 1 + assert _endpoints(emissions[0]) == (Decimal(0), Decimal(1)) + + +def test_set_selection_rejects_inverted_interval(slider: DecimalRangeSlider) -> None: + with pytest.raises(ValueError, match='upper < lower'): + slider.set_selection(_iv('0.7', '0.3')) + + +def test_set_selection_no_op_when_unchanged(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.3', '0.7')) + emissions = _capture(slider) + slider.set_selection(_iv('0.3', '0.7')) + assert emissions == [] + + +def test_set_selection_and_bounds_widens_bounds_and_updates_selection( + slider: DecimalRangeSlider, +) -> None: + slider.set_selection_and_bounds(_iv(-5, 5), _iv(-10, 10)) + assert _endpoints(slider.get_bounds()) == (Decimal(-10), Decimal(10)) + assert _endpoints(slider.get_selection()) == (Decimal(-5), Decimal(5)) + + +def test_set_selection_and_bounds_rejects_inverted_bounds(slider: DecimalRangeSlider) -> None: + with pytest.raises(ValueError, match='maximum <= minimum'): + slider.set_selection_and_bounds(_iv(0, 1), _iv(1, 0)) + + +def test_set_selection_and_bounds_rejects_degenerate_bounds(slider: DecimalRangeSlider) -> None: + with pytest.raises(ValueError, match='maximum <= minimum'): + slider.set_selection_and_bounds(_iv(0, 0), _iv(0, 0)) + + +def test_block_signal_suppresses_emission(slider: DecimalRangeSlider) -> None: + emissions = _capture(slider) + slider.set_selection_and_bounds(_iv(-5, 5), _iv(-10, 10), block_signal=True) + assert emissions == [] + assert _endpoints(slider.get_selection()) == (Decimal(-5), Decimal(5)) + + +def test_no_emit_on_bounds_only_change(slider: DecimalRangeSlider) -> None: + """Widen bounds while keeping the selection: no signal.""" + slider.set_selection(_iv('0.25', '0.75')) + emissions = _capture(slider) + slider.set_selection_and_bounds(_iv('0.25', '0.75'), _iv(-2, 2)) + assert emissions == [] + + +def test_signal_payload_is_interval_with_ordered_endpoints(slider: DecimalRangeSlider) -> None: + emissions = _capture(slider) + slider.set_selection(_iv('0.3', '0.7')) + assert len(emissions) == 1 + payload = emissions[0] + assert isinstance(payload, Interval) + assert payload.lower <= payload.upper + + +def _tick(bounds: Interval[Decimal], num_ticks: int = 1000) -> Decimal: + return (bounds.upper - bounds.lower) / Decimal(num_ticks) + + +def test_keyboard_arrow_moves_focused_handle_one_tick(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.3', '0.7')) + emissions = _capture(slider) + slider.setFocus() + QTest.keyClick(slider, Qt.Key.Key_Right) + step = _tick(slider.get_bounds()) + assert slider.get_selection().lower == Decimal('0.3') + step + assert len(emissions) == 1 + + +def test_keyboard_pageup_moves_ten_ticks(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.3', '0.7')) + slider.setFocus() + QTest.keyClick(slider, Qt.Key.Key_PageUp) + step = _tick(slider.get_bounds()) + assert slider.get_selection().lower == Decimal('0.3') + 10 * step + + +def test_keyboard_home_focused_lower_goes_to_bounds_lower(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.3', '0.7')) + slider.setFocus() + QTest.keyClick(slider, Qt.Key.Key_Home) + assert slider.get_selection().lower == Decimal(0) + assert slider.get_selection().upper == Decimal('0.7') + + +def test_keyboard_end_focused_lower_stops_at_upper(slider: DecimalRangeSlider) -> None: + """End on the lower handle must not push past the upper handle.""" + slider.set_selection(_iv('0.3', '0.7')) + slider.setFocus() + QTest.keyClick(slider, Qt.Key.Key_End) + assert slider.get_selection().lower == Decimal('0.7') + assert slider.get_selection().upper == Decimal('0.7') + + +def test_keyboard_handle_switch_via_bracket_keys(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.3', '0.7')) + emissions = _capture(slider) + slider.setFocus() + QTest.keyClick(slider, Qt.Key.Key_BracketRight) + QTest.keyClick(slider, Qt.Key.Key_End) + assert slider.get_selection().upper == Decimal(1) + QTest.keyClick(slider, Qt.Key.Key_BracketLeft) + QTest.keyClick(slider, Qt.Key.Key_Home) + assert slider.get_selection().lower == Decimal(0) + # bracket keys alone do not emit — only the End/Home keys after do + assert len(emissions) == 2 + + +def test_arrow_on_lower_stops_at_upper(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.5', '0.5')) + slider.setFocus() + QTest.keyClick(slider, Qt.Key.Key_Right) + assert slider.get_selection().lower == Decimal('0.5') + assert slider.get_selection().upper == Decimal('0.5') + + +def _paint_area(widget: DecimalRangeSlider): + return widget._paint_area # noqa: SLF001 - test accesses internal paint surface + + +def test_mouse_press_selects_lower_handle(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.25', '0.75')) + pa = _paint_area(slider) + x = pa._handle_x(Handle.LOWER) # noqa: SLF001 + y = pa._groove_y() # noqa: SLF001 + QTest.mousePress(pa, Qt.MouseButton.LeftButton, Qt.KeyboardModifier.NoModifier, QPoint(x, y)) + assert slider._active_handle is Handle.LOWER # noqa: SLF001 + QTest.mouseRelease(pa, Qt.MouseButton.LeftButton, Qt.KeyboardModifier.NoModifier, QPoint(x, y)) + assert slider._active_handle is None # noqa: SLF001 + + +def test_mouse_press_selects_upper_handle(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.25', '0.75')) + pa = _paint_area(slider) + x = pa._handle_x(Handle.UPPER) # noqa: SLF001 + y = pa._groove_y() # noqa: SLF001 + QTest.mousePress(pa, Qt.MouseButton.LeftButton, Qt.KeyboardModifier.NoModifier, QPoint(x, y)) + assert slider._active_handle is Handle.UPPER # noqa: SLF001 + + +def test_mouse_press_stacked_handles_tiebreak_by_side(slider: DecimalRangeSlider) -> None: + slider.set_selection(_iv('0.5', '0.5')) + pa = _paint_area(slider) + x = pa._handle_x(Handle.LOWER) # noqa: SLF001 + y = pa._groove_y() # noqa: SLF001 + # click just to the left of the stacked pair -> lower + QTest.mousePress( + pa, Qt.MouseButton.LeftButton, Qt.KeyboardModifier.NoModifier, QPoint(x - 2, y) + ) + assert slider._active_handle is Handle.LOWER # noqa: SLF001 + QTest.mouseRelease( + pa, Qt.MouseButton.LeftButton, Qt.KeyboardModifier.NoModifier, QPoint(x - 2, y) + ) + # click just to the right of the stacked pair -> upper + QTest.mousePress( + pa, Qt.MouseButton.LeftButton, Qt.KeyboardModifier.NoModifier, QPoint(x + 2, y) + ) + assert slider._active_handle is Handle.UPPER # noqa: SLF001 diff --git a/tests/view/widgets/test_decimal_slider.py b/tests/view/widgets/test_decimal_slider.py new file mode 100644 index 000000000..6f4dc7ed9 --- /dev/null +++ b/tests/view/widgets/test_decimal_slider.py @@ -0,0 +1,93 @@ +from __future__ import annotations +from decimal import Decimal + +import pytest + +from PyQt5.QtCore import Qt + +from ptychodus.api.geometry import Interval +from ptychodus.view.widgets import DecimalSlider + + +@pytest.fixture +def slider(qapp) -> DecimalSlider: + del qapp # fixture presence ensures QApplication exists + return DecimalSlider.create_instance(Qt.Orientation.Horizontal) + + +def _capture(widget: DecimalSlider) -> list[Decimal]: + emissions: list[Decimal] = [] + widget.value_changed.connect(emissions.append) + return emissions + + +def test_default_state(slider: DecimalSlider) -> None: + assert slider.get_value() == Decimal('0.5') + + +def test_set_value_stores_input(slider: DecimalSlider) -> None: + slider.set_value(Decimal('0.25')) + assert slider.get_value() == Decimal('0.25') + + +def test_set_value_clamps_below_minimum(slider: DecimalSlider) -> None: + slider.set_value(Decimal('-5')) + assert slider.get_value() == Decimal(0) + + +def test_set_value_clamps_above_maximum(slider: DecimalSlider) -> None: + slider.set_value(Decimal('5')) + assert slider.get_value() == Decimal(1) + + +def test_set_value_no_op_when_unchanged(slider: DecimalSlider) -> None: + slider.set_value(Decimal('0.25')) + emissions = _capture(slider) + slider.set_value(Decimal('0.25')) + assert emissions == [] + + +def test_set_value_emits_on_change(slider: DecimalSlider) -> None: + emissions = _capture(slider) + slider.set_value(Decimal('0.25')) + assert emissions == [Decimal('0.25')] + + +def test_set_value_and_range_rejects_inverted_bounds(slider: DecimalSlider) -> None: + with pytest.raises(ValueError, match='maximum <= minimum'): + slider.set_value_and_range(Decimal(0), Interval[Decimal](Decimal(1), Decimal(0))) + + +def test_set_value_and_range_rejects_degenerate_bounds(slider: DecimalSlider) -> None: + with pytest.raises(ValueError, match='maximum <= minimum'): + slider.set_value_and_range(Decimal(0), Interval[Decimal](Decimal(0), Decimal(0))) + + +def test_set_value_and_range_updates_bounds(slider: DecimalSlider) -> None: + slider.set_value_and_range(Decimal('5'), Interval[Decimal](Decimal(0), Decimal(10))) + assert slider.get_value() == Decimal('5') + + +def test_block_value_changed_signal_suppresses_emission(slider: DecimalSlider) -> None: + emissions = _capture(slider) + slider.set_value_and_range( + Decimal('7'), + Interval[Decimal](Decimal(0), Decimal(10)), + block_value_changed_signal=True, + ) + assert emissions == [] + assert slider.get_value() == Decimal('7') + + +def test_no_emit_on_bounds_only_change(slider: DecimalSlider) -> None: + """`set_value_and_range` must not fire `value_changed` when only bounds move.""" + slider.set_value(Decimal('0.5')) + emissions = _capture(slider) + slider.set_value_and_range(Decimal('0.5'), Interval[Decimal](Decimal(0), Decimal(2))) + assert emissions == [] + + +def test_emit_on_value_change_with_bounds_change(slider: DecimalSlider) -> None: + emissions = _capture(slider) + slider.set_value_and_range(Decimal('1.5'), Interval[Decimal](Decimal(0), Decimal(2))) + assert emissions == [Decimal('1.5')]