Implemented z-shift estimation module

docs/user_guide/3d_volumes.md

@@ -2,7 +2,12 @@
 
 This guide covers stack alignment for 3D volumetric data acquired through z-stack scanning in 2-photon microscopy.
 
-PyFlowReg performs **2D frame-by-frame motion correction**. It does not perform true 3D volumetric registration. For z-stack alignment, PyFlowReg uses an adaptive reference approach where the reference frame is updated slice-by-slice as you move through the stack.
+PyFlowReg performs **2D frame-by-frame motion correction**. It does not perform true
+3D volumetric registration. For z-stack alignment, PyFlowReg uses an adaptive
+reference approach where the reference frame is updated slice-by-slice as you move
+through the stack. The separate `pyflowreg.z_align` workflow can then estimate
+z-shifts against a reference stack and optionally write a z-corrected signal plus
+a z-shift-only simulation video.
 
 ## Z-Stack Acquisition Strategy
 
@@ -285,7 +290,8 @@ PyFlowReg's z-stack approach vs. true 3D volumetric registration:
 **PyFlowReg (2D + adaptive reference)**:
 - Registers each frame as a 2D image
 - Adapts reference slice-by-slice
-- Cannot correct through-plane (z-axis) motion
+- The core 2D workflow does not correct through-plane (z-axis) motion; use
+  `pyflowreg.z_align` for reference-stack z-shift estimation/correction
 - Fast and memory-efficient
 - Works well when z-motion is minimal and xy-motion dominates
 

examples/z_align_config.toml

@@ -0,0 +1,62 @@
+# Example z-align configuration matching the MATLAB z-shift scripts.
+# Assumes these files are present in `root`:
+# - compensated.tiff
+# - file_00004_00001.tif
+
+# === Data Location ===
+root = "."                             # Run from folder containing MATLAB-style inputs
+input_file = "compensated.tiff"        # Recording to z-correct
+volume_input_file = "file_00004_00001.tif"  # Input for reference volume creation
+reference_source_file = "compensated.tiff"  # Used to compute Stage-1 reference image
+
+# MATLAB snippet: get_video_file_reader("compensated.tiff", 10, 20); mean(first 2000)
+reference_source_frames = 2000
+reference_source_buffer_size = 10
+reference_source_bin_size = 20
+
+# === Output Paths ===
+# Kept identical to MATLAB output names/locations.
+output_root = "."
+volume_output_dir = "aligned_stack"
+recording_prealigned_output_dir = "prealigned_recording"
+z_shift_file = "z_shift.HDF5"
+corrected_output_file = "compensated_shift_corrected.tif"
+simulated_output_file = "simulated_from_z.tif"
+
+# === Pipeline Controls ===
+resume = true
+prealign_stack = true     # 2D-align the reference stack before z estimation
+prealign_recording = false  # optionally 2D-align input_file before z estimation
+write_corrected = true    # direct z-corrected signal
+write_simulated = true    # z-shift-only video interpolated from the volume
+
+# === Stage 1: Reference Volume Build (compensate_recording) ===
+stage1_alpha = 5.0
+stage1_quality_setting = "quality"
+stage1_buffer_size = 500
+stage1_bin_size = 1
+stage1_update_reference = true
+# Set to the number of scans per z slice to process one slice per batch.
+# When set, Stage 1 forces update_reference=true and uses this as volume_bin_size.
+# stack_scans_per_slice = 9
+flow_backend = "flowreg"
+
+# === Stage 2: Patch-Based z Estimation ===
+input_buffer_size = 50
+input_bin_size = 1
+volume_buffer_size = 500
+volume_bin_size = 1
+
+win_half = 10
+patch_size = 128
+overlap = 0.75
+
+spatial_sigma = 1.5
+temporal_sigma = 1.5
+z_smooth_sigma_spatial = 5.0
+z_smooth_sigma_temporal = 1.5
+parabolic_tau_scale = 1e-3
+
+output_dtype = "uint16"
+
+[backend_params]

examples/z_shift_demo.py

@@ -0,0 +1,96 @@
+"""
+Z-Shift Demo - MATLAB-style z alignment workflow
+
+This example assumes the same input files as the MATLAB scripts:
+- compensated.tiff         (time recording to z-correct)
+- file_00004_00001.tif     (stack/source for reference volume creation)
+
+Outputs (matching MATLAB names) are written to the working directory:
+- aligned_stack/compensated.HDF5
+- z_shift.HDF5
+- compensated_shift_corrected.tif
+- simulated_from_z.tif
+"""
+
+from pathlib import Path
+
+from pyflowreg.z_align import ZAlignConfig, run_all_stages
+
+
+def main():
+    root = Path(".").resolve()
+
+    required = [root / "compensated.tiff", root / "file_00004_00001.tif"]
+    missing = [p.name for p in required if not p.exists()]
+    if missing:
+        raise FileNotFoundError(
+            "Missing required input files in working directory: " + ", ".join(missing)
+        )
+
+    config = ZAlignConfig(
+        root=root,
+        # MATLAB-style inputs
+        input_file="compensated.tiff",
+        volume_input_file="file_00004_00001.tif",
+        reference_source_file="compensated.tiff",
+        # MATLAB script: read first 2000 frames with buffer/bin (10, 20)
+        reference_source_frames=2000,
+        reference_source_buffer_size=10,
+        reference_source_bin_size=20,
+        # Keep MATLAB output paths/names
+        output_root=".",
+        volume_output_dir="aligned_stack",
+        recording_prealigned_output_dir="prealigned_recording",
+        z_shift_file="z_shift.HDF5",
+        corrected_output_file="compensated_shift_corrected.tif",
+        simulated_output_file="simulated_from_z.tif",
+        # Stage toggles:
+        # write_corrected=True  -> direct z-corrected signal
+        # write_simulated=True  -> z-shift-only video interpolated from the volume
+        prealign_stack=True,
+        prealign_recording=False,
+        write_corrected=True,
+        write_simulated=True,
+        resume=True,
+        # Stage 1 (volume build) defaults from MATLAB snippet
+        stage1_alpha=5.0,
+        stage1_quality_setting="quality",
+        stage1_buffer_size=500,
+        stage1_bin_size=1,
+        stage1_update_reference=True,
+        # Set to scans per z slice when the stack stores repeated scans per slice.
+        stack_scans_per_slice=None,
+        # Stage 2 (patch-based z estimation) defaults from MATLAB snippet
+        input_buffer_size=50,
+        input_bin_size=1,
+        win_half=10,
+        patch_size=128,
+        overlap=0.75,
+        spatial_sigma=1.5,
+        temporal_sigma=1.5,
+        z_smooth_sigma_spatial=5.0,
+        z_smooth_sigma_temporal=1.5,
+    )
+
+    print("=" * 60)
+    print("Z-SHIFT DEMO")
+    print("=" * 60)
+    print(f"Root: {root}")
+    print("Input recording: compensated.tiff")
+    print("Volume source: file_00004_00001.tif")
+
+    outputs = run_all_stages(config)
+
+    print("\n" + "=" * 60)
+    print("DEMO COMPLETE")
+    print("=" * 60)
+    print(f"Reference volume: {outputs['volume_path']}")
+    print(f"Z-shift file:      {outputs['z_shift_path']}")
+    if outputs["corrected_path"] is not None:
+        print(f"Corrected signal:  {outputs['corrected_path']}")
+    if outputs["simulated_path"] is not None:
+        print(f"Z-shift simulation:{outputs['simulated_path']}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
 
 [project.scripts]
 pyflowreg-session = "pyflowreg.session.cli:main"
+pyflowreg-z-align = "pyflowreg.z_align.cli:main"
 
 [project.optional-dependencies]
 vis = [

src/pyflowreg/z_align/__init__.py

@@ -0,0 +1,28 @@
+"""
+Z-alignment pipeline for depth-shift correction.
+
+The ``z_align`` module implements a stage-based workflow that mirrors the
+MATLAB prototype used for z-shift estimation/correction:
+
+1. Build/load a reference volume
+2. Estimate per-pixel z shifts and optionally write z-corrected output
+3. Optionally simulate a z-shift-only recording from the estimated z shifts
+"""
+
+from pyflowreg.z_align.config import ZAlignConfig
+from pyflowreg.z_align.pipeline import (
+    run_recording_prealignment,
+    run_stage1,
+    run_stage2,
+    run_stage3,
+    run_all_stages,
+)
+
+__all__ = [
+    "ZAlignConfig",
+    "run_recording_prealignment",
+    "run_stage1",
+    "run_stage2",
+    "run_stage3",
+    "run_all_stages",
+]

src/pyflowreg/z_align/cli.py

@@ -0,0 +1,136 @@
+"""
+Command-line interface for z-alignment workflows.
+
+Provides the ``pyflowreg-z-align`` command.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from pyflowreg.z_align.config import ZAlignConfig
+from pyflowreg.z_align.pipeline import (
+    run_all_stages,
+    run_stage1,
+    run_stage2,
+    run_stage3,
+)
+
+
+def _parse_value(raw: str) -> Any:
+    """Parse CLI override values."""
+    lower = raw.lower()
+    if lower == "true":
+        return True
+    if lower == "false":
+        return False
+
+    for cast in (int, float):
+        try:
+            return cast(raw)
+        except ValueError:
+            pass
+
+    # Optional JSON parsing for lists/dicts
+    if raw.startswith("[") or raw.startswith("{"):
+        try:
+            return json.loads(raw)
+        except json.JSONDecodeError:
+            return raw
+    return raw
+
+
+def _parse_overrides(params: Optional[list[str]]) -> Dict[str, Any]:
+    """Parse KEY=VALUE CLI overrides."""
+    overrides: Dict[str, Any] = {}
+    if not params:
+        return overrides
+
+    for item in params:
+        if "=" not in item:
+            print(f"Warning: ignoring malformed override '{item}' (expected KEY=VALUE)")
+            continue
+        key, value = item.split("=", 1)
+        overrides[key] = _parse_value(value)
+    return overrides
+
+
+def cmd_run(args: argparse.Namespace) -> None:
+    """Handle the ``run`` subcommand."""
+    config_path = Path(args.config)
+    if not config_path.exists():
+        print(f"Error: configuration file not found: {config_path}")
+        sys.exit(1)
+
+    config = ZAlignConfig.from_file(config_path)
+    overrides = _parse_overrides(args.of_params)
+
+    if args.stage == "1":
+        run_stage1(config, overrides or None)
+        return
+
+    if args.stage == "2":
+        run_stage2(config)
+        return
+
+    if args.stage == "3":
+        run_stage3(config)
+        return
+
+    run_all_stages(config, overrides or None)
+
+
+def main() -> None:
+    """CLI entry point."""
+    parser = argparse.ArgumentParser(
+        description="PyFlowReg z-shift alignment pipeline",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Run full z-align workflow
+  pyflowreg-z-align run --config z_align.toml
+
+  # Run only z-shift estimation/correction
+  pyflowreg-z-align run --config z_align.toml --stage 2
+
+  # Override stage-1 OFOptions from CLI
+  pyflowreg-z-align run --config z_align.toml --of-params alpha=8 quality_setting=balanced
+        """,
+    )
+
+    subparsers = parser.add_subparsers(dest="command", help="Command to run")
+
+    run_parser = subparsers.add_parser("run", help="Run z-align processing")
+    run_parser.add_argument(
+        "--config",
+        "-c",
+        required=True,
+        help="Path to z-align config file (.toml/.yaml/.yml)",
+    )
+    run_parser.add_argument(
+        "--stage",
+        "-s",
+        choices=["1", "2", "3"],
+        help="Run only one stage (default: run all applicable stages)",
+    )
+    run_parser.add_argument(
+        "--of-params",
+        nargs="*",
+        metavar="KEY=VALUE",
+        help="Override stage-1 OFOptions parameters",
+    )
+    run_parser.set_defaults(func=cmd_run)
+
+    args = parser.parse_args()
+    if not hasattr(args, "func"):
+        parser.print_help()
+        sys.exit(1)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()