We need to implement a focus-map pipeline that estimates a global Z-offset surface over XY and applies it automatically during scans. Add a FocusLockManager to persist map(s) and params, extend FocusLockController with map acquisition and lock/settle signaling, and integrate with ExperimentController and PositionerManager/ESP32StageManager to apply Z corrections on the fly. Provide API and socket signals for UI/automation. Include a watchdog that aborts experiments when focus lock goes out of bounds.

Goals

• Record focus at multiple XY points (autofocus or manual) → build a surface (initially 3-point plane; optional denser fit).
• Persist focus maps and focus-lock params (disk + in-memory cache).
• Apply Z-offset automatically during XY moves (pre-move correction + runtime lock).
• Provide lock state, setpoint, current error, and settled status via global commchannel signals.
• Expose map/params via @APIExport and sockets for visualization and control.
• Add a watchdog to stop an experiment if focus lock is unhealthy/out-of-band.
• Load/save focus-lock and map config on boot.
• Provide a Mock focus-lock when hardware/controller is absent.

Out of Scope (v1)

• Advanced surfaces (RBF/Splines, regional fits) → v2.
• UI widgets for map editing/visual overlay → minimal hooks & APIs only => we are going to implement this in the frontend anyway
• Laser autofocus specifics (can be consumed later via same interface - actually already implemented in the Focuslockcontroller)

Architecture Changes

New

• FocusLockManager

  • Responsibilities: manage state, persistence, interpolation, parameter store.
  • Storage: JSON file(s) under the current profile (e.g., ~/.imswitch/focus/focusmap_<profile>.json).
  • Offers: get_z_offset(x, y), get_params(), save_map(), load_map(), clear_map(), interpolate(point).

Extend

• FocusLockController

  • New: map acquisition flow (grid or manual points), emit signals, calculate/set settled. (interface should have min/max x/y positions, ngrid x/y)

  • New signals (global/commchannel):

    • focuslock/state → {locked: bool, enabled: bool}
    • focuslock/setpoint → {z_ref_um: float}
    • focuslock/error → {error_um: float, abs_error_um: float, pct: float}
    • focuslock/settled → {settled: bool, band_um: float, timeout_ms: int}
    • focusmap/updated → {n_points: int, method: "plane", ts: iso8601}
    • focuslock/watchdog → {status: "ok"|"warn"|"abort", reason: str}

  • New APIExports:

    • start_focus_map_acquisition(grid_rows:int, grid_cols:int, margin:float=0.0)
    • add_focus_point(x_um:float, y_um:float, z_um:float=None, autofocus:bool=True)
    • fit_focus_map(method:str="plane")
    • clear_focus_map()
    • get_focus_map()
    • get_focus_params() / set_focus_params({...})
    • lock(enable: bool) (explicit lock/unlock)

  • Settling logic (band or timeout) computed inside controller (not UI).

• ExperimentController

=> it should also work in case the focuslockcontroller is not avialble =>default values

• Reads FocusLockManager for each XY site:

  • If map exists: pre-move Z to Z_ref + Z_map(x,y) + channel_offset.
  • Optionally enable live focus lock mode for continuous correction.

• Subscribes to focuslock/settled and focuslock/watchdog to gate acquisitions and abort on fault.

• New config flags (per protocol):

  • use_focus_map: bool
  • use_focus_lock_live: bool
  • focus_settle_band_um: float (default e.g., 1.0)
  • focus_settle_timeout_ms: int (default e.g., 1500)
  • channel_z_offsets: {channel_name: float_um}

• PositionerManager / ESP32StageManager (optional)

  • Hook: pre-move callback to fetch Z offset and add channel offset.
  • Provide atomic XY→Z move sequence or staged movement: move Z first or last as needed (configurable).

• Mock (MockFocusLockController)

  • Returns default zeros and settled=True, allows testing without hardware.

Data & Persistence

Focus Map JSON (v1, plane)

{
  "profile": "default",
  "method": "plane",
  "points": [
    {"x_um": 0.0, "y_um": 0.0, "z_um": 12.3},
    {"x_um": 1000.0, "y_um": 0.0, "z_um": 12.9},
    {"x_um": 0.0, "y_um": 1000.0, "z_um": 11.8}
  ],
  "fit": {
    "plane": {"a": 0.0004, "b": -0.0005, "c": 12.3}  // z = a*x + b*y + c
  },
  "channel_z_offsets": {"DAPI": 0.0, "FITC": 0.8, "TRITC": 1.2},
  "created_at": "2025-09-18T06:40:00Z",
  "updated_at": "2025-09-18T06:45:12Z"
}

Focus-Lock Params JSON

{
  "lock_enabled": true,
  "z_ref_um": 12345.6,
  "settle_band_um": 1.0,
  "settle_timeout_ms": 1500,
  "watchdog": {
    "max_abs_error_um": 5.0,
    "max_time_without_settle_ms": 5000,
    "action": "abort"
  }
}

APIs & Signals

@APIExport (FocusLockController)

• start_focus_map_acquisition(rows:int, cols:int, margin:float=0.0) -> dict
• add_focus_point(x_um:float, y_um:float, z_um:Optional[float]=None, autofocus:bool=True) -> dict
• fit_focus_map(method:str="plane") -> dict # returns coefficients
• clear_focus_map() -> dict
• get_focus_map() -> dict
• set_focus_params(params:dict) -> dict
• get_focus_params() -> dict
• lock(enable:bool) -> dict
• status() -> dict # {locked, settled, error_um, pct, setpoint}

Socket/Signals (commchannel topics)

• focuslock/state
• focuslock/setpoint
• focuslock/error
• focuslock/settled
• focusmap/updated
• focuslock/watchdog

Settled rule: emit settled=True if abs_error_um ≤ settle_band_um continuously for ≥ settle_window_ms, else False. Also emit watchdog if abs_error_um > max_abs_error_um for more than max_time_without_settle_ms.

Core Algorithms

Plane fit (minimum viable)

Given N≥3 points (x_i, y_i, z_i), least-squares solve z = a x + b y + c.
Function: FocusLockManager._fit_plane(points) -> (a, b, c)

Interpolation

get_z_offset(x, y) = a*x + b*y + c (add per-channel offset in ExperimentController).

Integration Flow

1. Acquisition/Calibration

   • User triggers map acquisition (grid or manual).

   • For each XY:

     • Move XY, autofocus (or manual), read Z → add_focus_point.

   • Fit plane → fit_focus_map → persisted by manager → focusmap/updated.

2. During Scan

   • On XY site:

     • Query Z_correction = FocusLockManager.get_z_offset(x,y) + channel_offset.
     • Pre-position Z to Z_ref + Z_correction.
     • If use_focus_lock_live: lock(True) and wait for focuslock/settled=True (respect timeout).
     • If not settled in time → watchdog → Experiment abort.

3. Fallback

   • If FocusLockController missing → MockFocusLockController returns zero offsets and settled=True.

Config (load on boot)

focuslock:
  enabled: true
  settle_band_um: 1.0
  settle_timeout_ms: 1500
  settle_window_ms: 200
  watchdog:
    max_abs_error_um: 5.0
    max_time_without_settle_ms: 5000
    action: abort
focusmap:
  method: plane
  path: ~/.imswitch/focus/focusmap_default.json
  use_focus_map: true
experiment:
  use_focus_lock_live: true
  channel_z_offsets:
    DAPI: 0.0
    FITC: 0.8
    TRITC: 1.2

Code Skeletons

focus_lock_manager.py

class FocusLockManager:
    def __init__(self, storage_path: Path):
        self._map = None
        self._params = {}
        self._path = storage_path

    def load_map(self) -> None: ...
    def save_map(self) -> None: ...
    def clear_map(self) -> None: ...
    def add_point(self, x_um: float, y_um: float, z_um: float) -> None: ...
    def fit(self, method: str = "plane") -> dict: ...
    def get_z_offset(self, x_um: float, y_um: float) -> float: ...
    def set_params(self, params: dict) -> None: ...
    def get_params(self) -> dict: ...

focus_lock_controller.py (additions)

class FocusLockController(ImConWidgetController):
    settled = Signal(bool, dict)  # {band_um, timeout_ms}
    state = Signal(dict)
    error = Signal(dict)
    map_updated = Signal(dict)
    watchdog = Signal(dict)

    @APIExport
    def start_focus_map_acquisition(self, rows:int, cols:int, margin:float=0.0): ...

    @APIExport
    def add_focus_point(self, x_um:float, y_um:float, z_um:float=None, autofocus:bool=True): ...

    @APIExport
    def fit_focus_map(self, method:str="plane"): ...

    @APIExport
    def get_focus_map(self): ...

    @APIExport
    def lock(self, enable: bool): ...

    def _update_settled(self): ...
    def _publish_error(self, err_um: float): ...

experiment_controller.py (hooks)

if cfg.focusmap.use_focus_map:
    z_corr = flm.get_z_offset(x_um, y_um) + channel_offsets.get(channel, 0.0)
    self.stage.move_z_to(z_ref + z_corr)

if cfg.experiment.use_focus_lock_live:
    foc.lock(True)
    if not self._wait_for_settle(timeout_ms=cfg.focuslock.settle_timeout_ms):
        self._abort("Focus did not settle")

Watchdog Logic

• Trigger warn if abs_error_um > settle_band_um for > settle_timeout_ms.
• Trigger abort if abs_error_um > max_abs_error_um for > max_time_without_settle_ms or loss of lock signal.
• Experiment listens and stops gracefully (close devices, flush buffers, mark run incomplete with reason).

Acceptance Criteria

• Given a recorded 3-point map, get_z_offset(x,y) returns plane-interpolated Z with <0.2 µm RMS error on synthetic tests.
• During multi-site acquisition with use_focus_map: true, the stage pre-positions Z per site, and images are captured only after focuslock/settled=True or timeout (then abort).
• Losing focus beyond max_abs_error_um for > max_time_without_settle_ms emits focuslock/watchdog: abort and stops the experiment.
• Map and params persist across restarts; APIs return current config and map.
• When FocusLock is absent, Mock path runs without exceptions and settled=True.

Documentation

• New section: “Focus Map & Z-Offset Correction” (workflow, APIs, config).
• Signal reference (topics + payload).
• Example scripts (map acquisition + scan).

Open Questions

• Preferred persistence location under current ImSwitch profile?
• Atomic move order (Z-first vs Z-last) for your stages—default to Z-first?
• Default bands/timeouts (provide sensible defaults; tune on hardware).
• Channel naming source of truth (illumination manager vs acquisition config).