From ac9f6b8d74c50d92cf2b0550e6d563b09e0768ff Mon Sep 17 00:00:00 2001
From: mat10d <mdiberna@wi.mit.edu>
Date: Fri, 26 Jun 2026 09:19:48 -0400
Subject: [PATCH 1/5] feat(merge): expose hardcoded fast-merge levers as
 backward-compatible config params

Parameterize the formerly-hardcoded fast-merge knobs so they can be tuned per screen:
- evaluate_match: threshold_triangle, threshold_point, threshold_region, ransac_kwargs
- initial_alignment / multistep_alignment: evaluate_kwargs (+ batch_size)
- refine_local_warp: degree, iterations, min_correspondences via warp_kwargs,
  threaded through merge_triangle_hash / merge_sbs_phenotype
Wired through scripts/merge/{fast_alignment,fast_merge}.py and rules/merge.smk
(all config.get default None -> lib literal). Defaults unchanged => existing
screens byte-identical.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SAWdQnZAy3WU6bdxy41EgP
---
 workflow/lib/merge/fast_merge.py         | 81 +++++++++++++++++++++++-
 workflow/lib/merge/hash.py               | 29 +++++++--
 workflow/rules/merge.smk                 | 13 ++++
 workflow/scripts/merge/fast_alignment.py | 43 ++++++++++++-
 workflow/scripts/merge/fast_merge.py     | 14 ++++
 5 files changed, 168 insertions(+), 12 deletions(-)

diff --git a/workflow/lib/merge/fast_merge.py b/workflow/lib/merge/fast_merge.py
index 28a4c394..439c6233 100644
--- a/workflow/lib/merge/fast_merge.py
+++ b/workflow/lib/merge/fast_merge.py
@@ -4,9 +4,12 @@
 import numpy as np
 from scipy.spatial.distance import cdist
 from sklearn.linear_model import LinearRegression
+from sklearn.preprocessing import PolynomialFeatures
 
 
-def merge_triangle_hash(hash_df_0, hash_df_1, alignment, threshold=2):
+def merge_triangle_hash(
+    hash_df_0, hash_df_1, alignment, threshold=2, local_refinement=None, warp_kwargs=None
+):
     """Merges two DataFrames using triangle hashing after images at different magnifications have been hashed together.
 
     Args:
@@ -14,6 +17,12 @@ def merge_triangle_hash(hash_df_0, hash_df_1, alignment, threshold=2):
         hash_df_1 (pandas.DataFrame): The second DataFrame.
         alignment (dict): Alignment parameters containing rotation and translation.
         threshold (int): The threshold value. Defaults to 2.
+        local_refinement (str | bool | None): If truthy (e.g. "polynomial"), refine the
+            global affine with a local non-rigid warp before matching, to correct
+            residual within-tile distortion (e.g. two-scope acquisitions). Defaults None
+            (off) — behaviour is then identical to the plain affine merge.
+        warp_kwargs (dict | None): Keyword args forwarded to `refine_local_warp`
+            (degree, iterations, min_correspondences) when local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: The merged DataFrame.
@@ -25,7 +34,56 @@ def merge_triangle_hash(hash_df_0, hash_df_1, alignment, threshold=2):
     model = build_linear_model(alignment["rotation"], alignment["translation"])
 
     # Merge dataframes using triangle hashing
-    return merge_sbs_phenotype(hash_df_0, hash_df_1, model, threshold=threshold)
+    return merge_sbs_phenotype(
+        hash_df_0,
+        hash_df_1,
+        model,
+        threshold=threshold,
+        local_refinement=local_refinement,
+        warp_kwargs=warp_kwargs,
+    )
+
+
+def refine_local_warp(
+    X, Y, Y_pred, threshold, degree=2, iterations=2, min_correspondences=30
+):
+    """Refine a global affine alignment with a local non-rigid (polynomial) warp.
+
+    Corrects residual within-tile distortion that a single affine cannot capture (the
+    failure mode when SBS and phenotype come from differently-configured microscopes).
+    The warp is fit ONLY on high-confidence correspondences — points already matched
+    within `threshold` under the current prediction — so it cannot be pulled by spurious
+    loose matches. Each iteration the warp improves and more points come into range.
+
+    Degrades gracefully: with fewer than `min_correspondences` confident matches it
+    returns the input prediction unchanged (so sparse tiles behave like plain affine).
+
+    Args:
+        X (numpy.ndarray): Source (phenotype) coordinates, shape (n, 2).
+        Y (numpy.ndarray): Target (SBS) coordinates, shape (m, 2).
+        Y_pred (numpy.ndarray): Current predicted source coords in target space (n, 2),
+            i.e. the global-affine prediction.
+        threshold (float): Match distance defining a high-confidence correspondence.
+        degree (int): Polynomial degree of the warp. Defaults to 2.
+        iterations (int): Number of refine-and-rematch passes. Defaults to 2.
+        min_correspondences (int): Minimum confident matches required to fit. Defaults 30.
+
+    Returns:
+        numpy.ndarray: Refined predicted source coordinates in target space, shape (n, 2).
+    """
+    pf = PolynomialFeatures(degree)
+    refined = Y_pred
+    for _ in range(iterations):
+        distances = cdist(Y, refined, metric="sqeuclidean")
+        nearest = distances.argmin(axis=1)
+        within = np.sqrt(distances.min(axis=1)) < threshold
+        if within.sum() < min_correspondences:
+            break
+        X_corr = X[nearest[within]]
+        Y_corr = Y[within]
+        reg = LinearRegression().fit(pf.fit_transform(X_corr), Y_corr)
+        refined = reg.predict(pf.transform(X))
+    return refined
 
 
 def build_linear_model(rotation, translation):
@@ -45,7 +103,14 @@ def build_linear_model(rotation, translation):
     return m  # Return the linear regression model
 
 
-def merge_sbs_phenotype(cell_locations_0, cell_locations_1, model, threshold=2):
+def merge_sbs_phenotype(
+    cell_locations_0,
+    cell_locations_1,
+    model,
+    threshold=2,
+    local_refinement=None,
+    warp_kwargs=None,
+):
     """Perform fine alignment of one (tile, site) match found using `multistep_alignment`.
 
     Args:
@@ -61,6 +126,11 @@ def merge_sbs_phenotype(cell_locations_0, cell_locations_1, model, threshold=2):
             matrix determined in `multistep_alignment`.
         threshold (float, optional): Maximum Euclidean distance allowed between matching
             points. Defaults to 2.
+        local_refinement (str | bool | None, optional): If truthy, refine the affine
+            prediction with a local polynomial warp (`refine_local_warp`) before matching.
+            Defaults None (off) — prediction is then the plain affine, identical to before.
+        warp_kwargs (dict | None, optional): Keyword args forwarded to `refine_local_warp`
+            (degree, iterations, min_correspondences) when local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: Table of merged identities of cell labels from cell_locations_0 and cell_locations_1.
@@ -98,6 +168,11 @@ def merge_sbs_phenotype(cell_locations_0, cell_locations_1, model, threshold=2):
     # Predict coordinates for dataset 0 using the alignment model
     Y_pred = model.predict(X)
 
+    # Optional local non-rigid refinement of the affine prediction (corrects residual
+    # within-tile distortion; off by default so existing screens are unaffected).
+    if local_refinement:
+        Y_pred = refine_local_warp(X, Y, Y_pred, threshold, **(warp_kwargs or {}))
+
     # Calculate squared Euclidean distances between predicted coordinates and dataset 1 coordinates
     distances = cdist(Y, Y_pred, metric="sqeuclidean")
 
diff --git a/workflow/lib/merge/hash.py b/workflow/lib/merge/hash.py
index d63200be..c924543b 100644
--- a/workflow/lib/merge/hash.py
+++ b/workflow/lib/merge/hash.py
@@ -185,7 +185,9 @@ def nine_edge_hash(dt, i):
     return segments, vector
 
 
-def initial_alignment(well_triangles_0, well_triangles_1, initial_sites=8):
+def initial_alignment(
+    well_triangles_0, well_triangles_1, initial_sites=8, evaluate_kwargs=None
+):
     """Identifies matching tiles from two acquisitions with similar Delaunay triangulations within the same well.
 
     Matches tiles from two datasets based on Delaunay triangulations, assuming minimal cell movement between acquisitions and equivalent segmentations.
@@ -194,14 +196,16 @@ def initial_alignment(well_triangles_0, well_triangles_1, initial_sites=8):
         well_triangles_0 (pandas.DataFrame): Hashed Delaunay triangulation for all tiles in dataset 0. Produced by concatenating outputs of `find_triangles` for individual tiles of a single well. Must include a `tile` column.
         well_triangles_1 (pandas.DataFrame): Hashed Delaunay triangulation for all sites in dataset 1. Produced by concatenating outputs of `find_triangles` for individual sites of a single well. Must include a `site` column.
         initial_sites (int | list[tuple[int, int]], optional): If an integer, specifies the number of sites sampled from `df_1` for initial brute-force matching of tiles to build the alignment model. If a list of 2-tuples, represents known (tile, site) matches to initialize the alignment model. At least 5 pairs are recommended.
+        evaluate_kwargs (dict, optional): Keyword args forwarded to `evaluate_match` (threshold_triangle, threshold_point, threshold_region, ransac_kwargs). Defaults to None.
 
     Returns:
         pandas.DataFrame: Table of possible (tile, site) matches, including rotation and translation transformations. Includes all tested matches, which should be filtered by `score` and `determinant` to retain valid matches.
     """
+    evaluate_kwargs = evaluate_kwargs or {}
 
     # Define a function to work on individual (tile,site) pairs
     def work_on(df_t, df_s):
-        rotation, translation, score = evaluate_match(df_t, df_s)
+        rotation, translation, score = evaluate_match(df_t, df_s, **evaluate_kwargs)
         determinant = None if rotation is None else np.linalg.det(rotation)
         result = pd.Series(
             {
@@ -227,7 +231,12 @@ def work_on(df_t, df_s):
 
 
 def evaluate_match(
-    vec_centers_0, vec_centers_1, threshold_triangle=0.3, threshold_point=2
+    vec_centers_0,
+    vec_centers_1,
+    threshold_triangle=0.3,
+    threshold_point=2,
+    threshold_region=50,
+    ransac_kwargs=None,
 ):
     """Evaluates the match between two sets of vectors and centers.
 
@@ -238,6 +247,8 @@ def evaluate_match(
         vec_centers_1 (pandas.DataFrame): DataFrame containing the second set of vectors and centers.
         threshold_triangle (float, optional): Threshold for matching triangles. Defaults to 0.3.
         threshold_point (float, optional): Threshold for matching points. Defaults to 2.
+        threshold_region (float, optional): Region radius (px) within which a triangle center is scored. Defaults to 50.
+        ransac_kwargs (dict, optional): Keyword args forwarded to RANSACRegressor (e.g. residual_threshold, max_trials, min_samples, random_state). Defaults to None (sklearn defaults).
 
     Returns:
         tuple:
@@ -267,7 +278,7 @@ def evaluate_match(
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore")
         # Use matching triangles to define transformation
-        model = RANSACRegressor()
+        model = RANSACRegressor(**(ransac_kwargs or {}))
         model.fit(X, Y)  # Fit the RANSAC model to the matching centers
 
     rotation = model.estimator_.coef_  # Extract rotation matrix
@@ -275,7 +286,6 @@ def evaluate_match(
 
     # Score transformation based on the triangle centers
     distances = cdist(model.predict(c_0), c_1, metric="sqeuclidean")
-    threshold_region = 50  # Threshold for the region to consider
     min_distances = np.sqrt(distances.min(axis=0))
     filt = min_distances < threshold_region
     score = (min_distances[filt] < threshold_point).mean()  # Calculate score
@@ -345,6 +355,7 @@ def multistep_alignment(
     initial_sites=8,
     batch_size=180,
     n_jobs=None,
+    evaluate_kwargs=None,
 ):
     """Find tiles of two different acquisitions with matching Delaunay triangulations within the same well.
 
@@ -370,18 +381,24 @@ def multistep_alignment(
         batch_size (int, optional): Number of (tile, site) matches to evaluate per batch during global
             alignment model updates. Defaults to 180.
         n_jobs (int, optional): Number of parallel jobs to deploy using joblib. Defaults to None.
+        evaluate_kwargs (dict, optional): Keyword args forwarded to `evaluate_match` (threshold_triangle,
+            threshold_point, threshold_region, ransac_kwargs). Defaults to None.
 
     Returns:
         pandas.DataFrame: Table of possible (tile, site) matches with corresponding rotation and translation
         transformations. All tested matches are included; query based on `score` and `determinant` to filter valid matches.
     """
+    evaluate_kwargs = evaluate_kwargs or {}
+
     # If n_jobs is not provided, set it to one less than the number of CPU cores
     if n_jobs is None:
         n_jobs = multiprocessing.cpu_count() - 1
 
     # Define a function to work on individual (tile,site) pairs
     def work_on(tiles_df, sites_df):
-        rotation, translation, score = evaluate_match(tiles_df, sites_df)
+        rotation, translation, score = evaluate_match(
+            tiles_df, sites_df, **evaluate_kwargs
+        )
         determinant = None if rotation is None else np.linalg.det(rotation)
         result = pd.Series(
             {
diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk
index 1d5e85dc..9c1de07e 100644
--- a/workflow/rules/merge.smk
+++ b/workflow/rules/merge.smk
@@ -45,9 +45,18 @@ if merge_approach == "fast":
             initial_sites=config.get("merge", {}).get("initial_sites"),
             plate=lambda wildcards: wildcards.plate,
             well=lambda wildcards: wildcards.well,
+            metadata_align=config.get("merge", {}).get("metadata_align", False),
             alignment_flip_x=config.get("merge", {}).get("alignment_flip_x"),
             alignment_flip_y=config.get("merge", {}).get("alignment_flip_y"),
             alignment_rotate_90=config.get("merge", {}).get("alignment_rotate_90"),
+            threshold_triangle=config.get("merge", {}).get("threshold_triangle"),
+            threshold_point=config.get("merge", {}).get("threshold_point"),
+            threshold_region=config.get("merge", {}).get("threshold_region"),
+            ransac_residual_threshold=config.get("merge", {}).get("ransac_residual_threshold"),
+            ransac_max_trials=config.get("merge", {}).get("ransac_max_trials"),
+            ransac_min_samples=config.get("merge", {}).get("ransac_min_samples"),
+            ransac_random_state=config.get("merge", {}).get("ransac_random_state"),
+            batch_size=config.get("merge", {}).get("batch_size"),
         script:
             "../scripts/merge/fast_alignment.py"
 
@@ -72,6 +81,10 @@ if merge_approach == "fast":
             det_range=config.get("merge", {}).get("det_range"),
             score=config.get("merge", {}).get("score"),
             threshold=config.get("merge", {}).get("threshold"),
+            local_refinement=config.get("merge", {}).get("local_refinement"),
+            warp_degree=config.get("merge", {}).get("warp_degree"),
+            warp_iterations=config.get("merge", {}).get("warp_iterations"),
+            warp_min_correspondences=config.get("merge", {}).get("warp_min_correspondences"),
         script:
             "../scripts/merge/fast_merge.py"
 
diff --git a/workflow/scripts/merge/fast_alignment.py b/workflow/scripts/merge/fast_alignment.py
index 876971bc..7aaf401c 100644
--- a/workflow/scripts/merge/fast_alignment.py
+++ b/workflow/scripts/merge/fast_alignment.py
@@ -25,8 +25,14 @@
     "rotate_90": snakemake.params.alignment_rotate_90,
 }
 
-# Only apply alignment if at least one transformation is requested
-if any(
+# metadata_align center-aligns the two scopes' stage coordinate frames (translation)
+# even when no flip/rotate is requested — needed when SBS and phenotype were acquired
+# on different microscopes with different coordinate origins. Backward compatible:
+# defaults False, and any requested flip/rotate still triggers alignment as before.
+metadata_align = getattr(snakemake.params, "metadata_align", False)
+
+# Apply alignment if center-alignment is requested or any flip/rotate is set
+if metadata_align or any(
     [
         alignment_params["flip_x"],
         alignment_params["flip_y"],
@@ -106,6 +112,32 @@
 d0, d1 = snakemake.params.det_range
 score_thresh = snakemake.params.score
 
+
+# Optional alignment levers — backward compatible: absent config keys -> None -> drop ->
+# lib falls back to its literal defaults (existing screens unchanged).
+def _drop_none(d):
+    return {k: v for k, v in d.items() if v is not None}
+
+
+ransac_kwargs = _drop_none(
+    {
+        "residual_threshold": getattr(snakemake.params, "ransac_residual_threshold", None),
+        "max_trials": getattr(snakemake.params, "ransac_max_trials", None),
+        "min_samples": getattr(snakemake.params, "ransac_min_samples", None),
+        "random_state": getattr(snakemake.params, "ransac_random_state", None),
+    }
+)
+evaluate_kwargs = _drop_none(
+    {
+        "threshold_triangle": getattr(snakemake.params, "threshold_triangle", None),
+        "threshold_point": getattr(snakemake.params, "threshold_point", None),
+        "threshold_region": getattr(snakemake.params, "threshold_region", None),
+        "ransac_kwargs": ransac_kwargs or None,
+    }
+) or None
+batch_size = getattr(snakemake.params, "batch_size", None)
+multistep_extra = {} if batch_size is None else {"batch_size": batch_size}
+
 if initial_sbs_tiles is not None:
     # Auto-discover initial sites from SBS tiles
     candidate_pairs = []
@@ -125,7 +157,10 @@
 
 # Run initial alignment on candidates (validates both paths)
 initial_alignment_df = initial_alignment(
-    phenotype_info_hash, sbs_info_hash, initial_sites=candidate_pairs
+    phenotype_info_hash,
+    sbs_info_hash,
+    initial_sites=candidate_pairs,
+    evaluate_kwargs=evaluate_kwargs,
 )
 
 # Filter by thresholds
@@ -155,6 +190,8 @@
     score=snakemake.params.score,
     initial_sites=initial_sites,
     n_jobs=snakemake.threads,
+    evaluate_kwargs=evaluate_kwargs,
+    **multistep_extra,
 )
 
 # Reset index
diff --git a/workflow/scripts/merge/fast_merge.py b/workflow/scripts/merge/fast_merge.py
index dc005882..ed35cd8c 100644
--- a/workflow/scripts/merge/fast_merge.py
+++ b/workflow/scripts/merge/fast_merge.py
@@ -33,6 +33,18 @@
 print(f"Total alignments: {len(fast_alignment)}")
 print(f"Filtered alignments: {len(fast_alignment_filtered)}")
 
+# Optional polynomial-warp levers — backward compatible: absent keys -> None -> dropped ->
+# refine_local_warp falls back to its literal defaults (existing screens unchanged).
+warp_kwargs = {
+    k: v
+    for k, v in {
+        "degree": getattr(snakemake.params, "warp_degree", None),
+        "iterations": getattr(snakemake.params, "warp_iterations", None),
+        "min_correspondences": getattr(snakemake.params, "warp_min_correspondences", None),
+    }.items()
+    if v is not None
+} or None
+
 # Merge cells across well
 merge_data = []
 for index, alignment_row in fast_alignment_filtered.iterrows():
@@ -50,6 +62,8 @@
         sbs_info_filtered,
         alignment_row,
         threshold=snakemake.params.threshold,
+        local_refinement=getattr(snakemake.params, "local_refinement", None),
+        warp_kwargs=warp_kwargs,
     )
     merge_data.append(alignment_row_merge)
 

From 98db9fc74150ade7a663d8b4af369a3c043caa2a Mon Sep 17 00:00:00 2001
From: mat10d <mdiberna@wi.mit.edu>
Date: Tue, 7 Jul 2026 14:59:14 -0400
Subject: [PATCH 2/5] feat(merge): thin-plate-spline warp + find-optimal-site
 seeding (gated)

Adds two validated two-scope merge levers, both backward-compatible (defaults
off -> byte-identical to the affine path already on this branch):

- TPS warp: `local_refinement` now doubles as the warp-model selector
  ("polynomial" | "thin_plate_spline"); wire `warp_smoothing` /
  `warp_max_correspondences` through scripts/merge/fast_merge.py + merge.smk.
  On two-scope OPS merges TPS roughly halves the median residual vs polynomial
  while staying stable. String selector is byte-identical to an explicit
  warp_kwargs["model"]; a bare truthy keeps the polynomial default.

- find-optimal-site (`seed_optimize`, `seed_topk`): the stage-nearest phenotype
  tile is not always the best real overlap when the two scopes have a
  metadata-alignment residual. When on, evaluate the top-K nearest tiles per SBS
  seed and keep the highest-scoring one per site. Default off = nearest-tile,
  unchanged.

Verified: default (local_refinement=None, seed_optimize=False) path untouched;
polynomial selector == explicit model kwarg; TPS routes end-to-end and both
warps recover matches an affine drops on distorted data.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SAWdQnZAy3WU6bdxy41EgP
---
 workflow/lib/merge/fast_merge.py         | 73 +++++++++++++++++++-----
 workflow/rules/merge.smk                 |  4 ++
 workflow/scripts/merge/fast_alignment.py | 26 ++++++++-
 workflow/scripts/merge/fast_merge.py     |  8 ++-
 4 files changed, 92 insertions(+), 19 deletions(-)

diff --git a/workflow/lib/merge/fast_merge.py b/workflow/lib/merge/fast_merge.py
index 439c6233..d976ec09 100644
--- a/workflow/lib/merge/fast_merge.py
+++ b/workflow/lib/merge/fast_merge.py
@@ -17,12 +17,14 @@ def merge_triangle_hash(
         hash_df_1 (pandas.DataFrame): The second DataFrame.
         alignment (dict): Alignment parameters containing rotation and translation.
         threshold (int): The threshold value. Defaults to 2.
-        local_refinement (str | bool | None): If truthy (e.g. "polynomial"), refine the
-            global affine with a local non-rigid warp before matching, to correct
-            residual within-tile distortion (e.g. two-scope acquisitions). Defaults None
-            (off) — behaviour is then identical to the plain affine merge.
+        local_refinement (str | bool | None): If truthy, refine the global affine with a
+            local non-rigid warp before matching, to correct residual within-tile
+            distortion (e.g. two-scope acquisitions). The string selects the warp model:
+            "polynomial" or "thin_plate_spline". Defaults None (off) — behaviour is then
+            identical to the plain affine merge.
         warp_kwargs (dict | None): Keyword args forwarded to `refine_local_warp`
-            (degree, iterations, min_correspondences) when local_refinement is on. Defaults None.
+            (degree, iterations, min_correspondences, smoothing, max_correspondences) when
+            local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: The merged DataFrame.
@@ -45,9 +47,10 @@ def merge_triangle_hash(
 
 
 def refine_local_warp(
-    X, Y, Y_pred, threshold, degree=2, iterations=2, min_correspondences=30
+    X, Y, Y_pred, threshold, degree=2, iterations=2, min_correspondences=30,
+    model="polynomial", smoothing=10.0, max_correspondences=700
 ):
-    """Refine a global affine alignment with a local non-rigid (polynomial) warp.
+    """Refine a global affine alignment with a local non-rigid warp.
 
     Corrects residual within-tile distortion that a single affine cannot capture (the
     failure mode when SBS and phenotype come from differently-configured microscopes).
@@ -55,6 +58,16 @@ def refine_local_warp(
     within `threshold` under the current prediction — so it cannot be pulled by spurious
     loose matches. Each iteration the warp improves and more points come into range.
 
+    Two warp models (``model``):
+      - "polynomial" (default): a degree-``degree`` polynomial. Backward-compatible;
+        the original behaviour. Beware high degrees — a degree-5 polynomial oscillates
+        (Runge) and can diverge.
+      - "thin_plate_spline": a smoothed thin-plate spline (scipy RBFInterpolator). It is
+        the smoothest deformation fitting the correspondences, so it captures local
+        distortion without the polynomial's oscillation, regularized by ``smoothing``.
+        Validated to beat the polynomial on two-scope OPS merges (higher single-match
+        rate, ~half the median residual) while remaining stable.
+
     Degrades gracefully: with fewer than `min_correspondences` confident matches it
     returns the input prediction unchanged (so sparse tiles behave like plain affine).
 
@@ -64,9 +77,13 @@ def refine_local_warp(
         Y_pred (numpy.ndarray): Current predicted source coords in target space (n, 2),
             i.e. the global-affine prediction.
         threshold (float): Match distance defining a high-confidence correspondence.
-        degree (int): Polynomial degree of the warp. Defaults to 2.
+        degree (int): Polynomial degree (model="polynomial"). Defaults to 2.
         iterations (int): Number of refine-and-rematch passes. Defaults to 2.
         min_correspondences (int): Minimum confident matches required to fit. Defaults 30.
+        model (str): "polynomial" (default) or "thin_plate_spline".
+        smoothing (float): TPS regularization (model="thin_plate_spline"). Defaults 10.0.
+        max_correspondences (int): TPS fit is capped at this many points (subsampled) for
+            speed; the fit is smooth so a few hundred anchors suffice. Defaults 700.
 
     Returns:
         numpy.ndarray: Refined predicted source coordinates in target space, shape (n, 2).
@@ -79,10 +96,28 @@ def refine_local_warp(
         within = np.sqrt(distances.min(axis=1)) < threshold
         if within.sum() < min_correspondences:
             break
-        X_corr = X[nearest[within]]
         Y_corr = Y[within]
-        reg = LinearRegression().fit(pf.fit_transform(X_corr), Y_corr)
-        refined = reg.predict(pf.transform(X))
+        if model == "thin_plate_spline":
+            from scipy.interpolate import RBFInterpolator
+
+            # TPS is fit on the RESIDUAL (current prediction -> target), not raw
+            # source -> target: its `smoothing` is calibrated for the small residual
+            # scale, so fitting the full transform (which carries the ~0.27x scale)
+            # would mis-regularize. Compose the correction onto the running estimate.
+            src = refined[nearest[within]]
+            if len(src) > max_correspondences:
+                sel = np.random.default_rng(0).choice(
+                    len(src), max_correspondences, replace=False
+                )
+                src, Y_corr = src[sel], Y_corr[sel]
+            rbf = RBFInterpolator(
+                src, Y_corr, kernel="thin_plate_spline", smoothing=smoothing
+            )
+            refined = rbf(refined)
+        else:
+            X_corr = X[nearest[within]]
+            reg = LinearRegression().fit(pf.fit_transform(X_corr), Y_corr)
+            refined = reg.predict(pf.transform(X))
     return refined
 
 
@@ -127,10 +162,12 @@ def merge_sbs_phenotype(
         threshold (float, optional): Maximum Euclidean distance allowed between matching
             points. Defaults to 2.
         local_refinement (str | bool | None, optional): If truthy, refine the affine
-            prediction with a local polynomial warp (`refine_local_warp`) before matching.
-            Defaults None (off) — prediction is then the plain affine, identical to before.
+            prediction with a local warp (`refine_local_warp`) before matching. The string
+            "polynomial" or "thin_plate_spline" selects the model. Defaults None (off) —
+            prediction is then the plain affine, identical to before.
         warp_kwargs (dict | None, optional): Keyword args forwarded to `refine_local_warp`
-            (degree, iterations, min_correspondences) when local_refinement is on. Defaults None.
+            (degree, iterations, min_correspondences, smoothing, max_correspondences) when
+            local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: Table of merged identities of cell labels from cell_locations_0 and cell_locations_1.
@@ -170,8 +207,14 @@ def merge_sbs_phenotype(
 
     # Optional local non-rigid refinement of the affine prediction (corrects residual
     # within-tile distortion; off by default so existing screens are unaffected).
+    # `local_refinement` doubles as the warp-model selector: the string "polynomial" or
+    # "thin_plate_spline" picks the model, while a bare truthy (e.g. True) keeps
+    # refine_local_warp's default (polynomial). An explicit warp_kwargs["model"] wins.
     if local_refinement:
-        Y_pred = refine_local_warp(X, Y, Y_pred, threshold, **(warp_kwargs or {}))
+        wk = dict(warp_kwargs or {})
+        if isinstance(local_refinement, str):
+            wk.setdefault("model", local_refinement)
+        Y_pred = refine_local_warp(X, Y, Y_pred, threshold, **wk)
 
     # Calculate squared Euclidean distances between predicted coordinates and dataset 1 coordinates
     distances = cdist(Y, Y_pred, metric="sqeuclidean")
diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk
index 9c1de07e..bf41d6dc 100644
--- a/workflow/rules/merge.smk
+++ b/workflow/rules/merge.smk
@@ -57,6 +57,8 @@ if merge_approach == "fast":
             ransac_min_samples=config.get("merge", {}).get("ransac_min_samples"),
             ransac_random_state=config.get("merge", {}).get("ransac_random_state"),
             batch_size=config.get("merge", {}).get("batch_size"),
+            seed_optimize=config.get("merge", {}).get("seed_optimize", False),
+            seed_topk=config.get("merge", {}).get("seed_topk"),
         script:
             "../scripts/merge/fast_alignment.py"
 
@@ -85,6 +87,8 @@ if merge_approach == "fast":
             warp_degree=config.get("merge", {}).get("warp_degree"),
             warp_iterations=config.get("merge", {}).get("warp_iterations"),
             warp_min_correspondences=config.get("merge", {}).get("warp_min_correspondences"),
+            warp_smoothing=config.get("merge", {}).get("warp_smoothing"),
+            warp_max_correspondences=config.get("merge", {}).get("warp_max_correspondences"),
         script:
             "../scripts/merge/fast_merge.py"
 
diff --git a/workflow/scripts/merge/fast_alignment.py b/workflow/scripts/merge/fast_alignment.py
index 7aaf401c..d2186c13 100644
--- a/workflow/scripts/merge/fast_alignment.py
+++ b/workflow/scripts/merge/fast_alignment.py
@@ -138,6 +138,13 @@ def _drop_none(d):
 batch_size = getattr(snakemake.params, "batch_size", None)
 multistep_extra = {} if batch_size is None else {"batch_size": batch_size}
 
+# find-optimal-site (gated, default off): the stage-NEAREST phenotype tile is not always
+# the best real overlap when the two scopes have a metadata-alignment residual. When on,
+# offer the top-K nearest tiles per SBS tile as candidates and let the score/det filter
+# keep the single best-scoring one per site.
+seed_optimize = getattr(snakemake.params, "seed_optimize", False)
+seed_topk = getattr(snakemake.params, "seed_topk", None) or 3
+
 if initial_sbs_tiles is not None:
     # Auto-discover initial sites from SBS tiles
     candidate_pairs = []
@@ -145,10 +152,14 @@ def _drop_none(d):
         closest = find_closest_tiles(
             sbs_metadata, phenotype_metadata, sbs_tile, verbose=False
         )
-        best_ph_tile = int(closest.iloc[0]["tile"])
-        candidate_pairs.append([best_ph_tile, sbs_tile])
+        if seed_optimize:
+            for ph_tile in closest.head(seed_topk)["tile"].astype(int):
+                candidate_pairs.append([int(ph_tile), sbs_tile])
+        else:
+            candidate_pairs.append([int(closest.iloc[0]["tile"]), sbs_tile])
     print(
         f"Discovered {len(candidate_pairs)} candidate pairs from {len(initial_sbs_tiles)} SBS tiles"
+        + (f" (seed_optimize: top-{seed_topk} per tile)" if seed_optimize else "")
     )
 else:
     # Use user-provided pairs
@@ -168,6 +179,17 @@ def _drop_none(d):
     "@d0 <= determinant <= @d1 & score > @score_thresh"
 )
 
+# find-optimal-site: keep only the best-scoring phenotype tile per SBS site (higher score
+# == better real overlap), collapsing the top-K candidates expanded above back to one seed.
+if seed_optimize:
+    n_before = len(valid_pairs_df)
+    valid_pairs_df = valid_pairs_df.sort_values("score", ascending=False).drop_duplicates(
+        subset="site", keep="first"
+    )
+    print(
+        f"seed_optimize: kept best-scoring tile per site ({n_before} -> {len(valid_pairs_df)})"
+    )
+
 # Require minimum 5 valid pairs (only if > 5 candidates were provided)
 if len(candidate_pairs) > 5 and len(valid_pairs_df) < 5:
     raise ValueError(
diff --git a/workflow/scripts/merge/fast_merge.py b/workflow/scripts/merge/fast_merge.py
index ed35cd8c..ec3e5c12 100644
--- a/workflow/scripts/merge/fast_merge.py
+++ b/workflow/scripts/merge/fast_merge.py
@@ -33,14 +33,18 @@
 print(f"Total alignments: {len(fast_alignment)}")
 print(f"Filtered alignments: {len(fast_alignment_filtered)}")
 
-# Optional polynomial-warp levers — backward compatible: absent keys -> None -> dropped ->
-# refine_local_warp falls back to its literal defaults (existing screens unchanged).
+# Optional local-warp levers — backward compatible: absent keys -> None -> dropped ->
+# refine_local_warp falls back to its literal defaults (existing screens unchanged). The
+# warp MODEL is selected by `local_refinement` ("polynomial" | "thin_plate_spline");
+# `smoothing`/`max_correspondences` only apply to the thin-plate-spline model.
 warp_kwargs = {
     k: v
     for k, v in {
         "degree": getattr(snakemake.params, "warp_degree", None),
         "iterations": getattr(snakemake.params, "warp_iterations", None),
         "min_correspondences": getattr(snakemake.params, "warp_min_correspondences", None),
+        "smoothing": getattr(snakemake.params, "warp_smoothing", None),
+        "max_correspondences": getattr(snakemake.params, "warp_max_correspondences", None),
     }.items()
     if v is not None
 } or None

From 762c37bd7fe3072162582a3eeba8abfe2df817f3 Mon Sep 17 00:00:00 2001
From: mat10d <mdiberna@wi.mit.edu>
Date: Tue, 7 Jul 2026 15:49:17 -0400
Subject: [PATCH 3/5] style(merge): single-line comments, lib high->low order,
 ruff format

Cleanup of the merge fast-path (no behavioral change; smoke test identical):
- Collapse multiline block comments to single lines across fast_merge.py and
  fast_alignment.py.
- Reorder lib/merge/fast_merge.py high-level -> low-level: merge_triangle_hash,
  merge_sbs_phenotype, then helpers build_linear_model, refine_local_warp.
- Trim refine_local_warp docstring to the file's density; ruff format + check clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SAWdQnZAy3WU6bdxy41EgP
---
 workflow/lib/merge/fast_merge.py         | 188 +++++++++++------------
 workflow/scripts/merge/fast_alignment.py |  45 +++---
 workflow/scripts/merge/fast_merge.py     |  13 +-
 3 files changed, 118 insertions(+), 128 deletions(-)

diff --git a/workflow/lib/merge/fast_merge.py b/workflow/lib/merge/fast_merge.py
index d976ec09..8fa662bb 100644
--- a/workflow/lib/merge/fast_merge.py
+++ b/workflow/lib/merge/fast_merge.py
@@ -8,7 +8,12 @@
 
 
 def merge_triangle_hash(
-    hash_df_0, hash_df_1, alignment, threshold=2, local_refinement=None, warp_kwargs=None
+    hash_df_0,
+    hash_df_1,
+    alignment,
+    threshold=2,
+    local_refinement=None,
+    warp_kwargs=None,
 ):
     """Merges two DataFrames using triangle hashing after images at different magnifications have been hashed together.
 
@@ -46,98 +51,6 @@ def merge_triangle_hash(
     )
 
 
-def refine_local_warp(
-    X, Y, Y_pred, threshold, degree=2, iterations=2, min_correspondences=30,
-    model="polynomial", smoothing=10.0, max_correspondences=700
-):
-    """Refine a global affine alignment with a local non-rigid warp.
-
-    Corrects residual within-tile distortion that a single affine cannot capture (the
-    failure mode when SBS and phenotype come from differently-configured microscopes).
-    The warp is fit ONLY on high-confidence correspondences — points already matched
-    within `threshold` under the current prediction — so it cannot be pulled by spurious
-    loose matches. Each iteration the warp improves and more points come into range.
-
-    Two warp models (``model``):
-      - "polynomial" (default): a degree-``degree`` polynomial. Backward-compatible;
-        the original behaviour. Beware high degrees — a degree-5 polynomial oscillates
-        (Runge) and can diverge.
-      - "thin_plate_spline": a smoothed thin-plate spline (scipy RBFInterpolator). It is
-        the smoothest deformation fitting the correspondences, so it captures local
-        distortion without the polynomial's oscillation, regularized by ``smoothing``.
-        Validated to beat the polynomial on two-scope OPS merges (higher single-match
-        rate, ~half the median residual) while remaining stable.
-
-    Degrades gracefully: with fewer than `min_correspondences` confident matches it
-    returns the input prediction unchanged (so sparse tiles behave like plain affine).
-
-    Args:
-        X (numpy.ndarray): Source (phenotype) coordinates, shape (n, 2).
-        Y (numpy.ndarray): Target (SBS) coordinates, shape (m, 2).
-        Y_pred (numpy.ndarray): Current predicted source coords in target space (n, 2),
-            i.e. the global-affine prediction.
-        threshold (float): Match distance defining a high-confidence correspondence.
-        degree (int): Polynomial degree (model="polynomial"). Defaults to 2.
-        iterations (int): Number of refine-and-rematch passes. Defaults to 2.
-        min_correspondences (int): Minimum confident matches required to fit. Defaults 30.
-        model (str): "polynomial" (default) or "thin_plate_spline".
-        smoothing (float): TPS regularization (model="thin_plate_spline"). Defaults 10.0.
-        max_correspondences (int): TPS fit is capped at this many points (subsampled) for
-            speed; the fit is smooth so a few hundred anchors suffice. Defaults 700.
-
-    Returns:
-        numpy.ndarray: Refined predicted source coordinates in target space, shape (n, 2).
-    """
-    pf = PolynomialFeatures(degree)
-    refined = Y_pred
-    for _ in range(iterations):
-        distances = cdist(Y, refined, metric="sqeuclidean")
-        nearest = distances.argmin(axis=1)
-        within = np.sqrt(distances.min(axis=1)) < threshold
-        if within.sum() < min_correspondences:
-            break
-        Y_corr = Y[within]
-        if model == "thin_plate_spline":
-            from scipy.interpolate import RBFInterpolator
-
-            # TPS is fit on the RESIDUAL (current prediction -> target), not raw
-            # source -> target: its `smoothing` is calibrated for the small residual
-            # scale, so fitting the full transform (which carries the ~0.27x scale)
-            # would mis-regularize. Compose the correction onto the running estimate.
-            src = refined[nearest[within]]
-            if len(src) > max_correspondences:
-                sel = np.random.default_rng(0).choice(
-                    len(src), max_correspondences, replace=False
-                )
-                src, Y_corr = src[sel], Y_corr[sel]
-            rbf = RBFInterpolator(
-                src, Y_corr, kernel="thin_plate_spline", smoothing=smoothing
-            )
-            refined = rbf(refined)
-        else:
-            X_corr = X[nearest[within]]
-            reg = LinearRegression().fit(pf.fit_transform(X_corr), Y_corr)
-            refined = reg.predict(pf.transform(X))
-    return refined
-
-
-def build_linear_model(rotation, translation):
-    """Builds a linear regression model using the provided rotation matrix and translation vector.
-
-    Args:
-        rotation (numpy.ndarray): Rotation matrix for the model.
-        translation (numpy.ndarray): Translation vector for the model.
-
-    Returns:
-        sklearn.linear_model.LinearRegression: Linear regression model with the specified rotation
-        and translation.
-    """
-    m = LinearRegression()
-    m.coef_ = rotation  # Set the rotation matrix as the model's coefficients
-    m.intercept_ = translation  # Set the translation vector as the model's intercept
-    return m  # Return the linear regression model
-
-
 def merge_sbs_phenotype(
     cell_locations_0,
     cell_locations_1,
@@ -205,11 +118,7 @@ def merge_sbs_phenotype(
     # Predict coordinates for dataset 0 using the alignment model
     Y_pred = model.predict(X)
 
-    # Optional local non-rigid refinement of the affine prediction (corrects residual
-    # within-tile distortion; off by default so existing screens are unaffected).
-    # `local_refinement` doubles as the warp-model selector: the string "polynomial" or
-    # "thin_plate_spline" picks the model, while a bare truthy (e.g. True) keeps
-    # refine_local_warp's default (polynomial). An explicit warp_kwargs["model"] wins.
+    # Optional local warp; local_refinement gates it and (as a string) names the model
     if local_refinement:
         wk = dict(warp_kwargs or {})
         if isinstance(local_refinement, str):
@@ -248,3 +157,86 @@ def merge_sbs_phenotype(
             cols_final
         ]  # Assign distance column  # Select final columns
     )
+
+
+def build_linear_model(rotation, translation):
+    """Builds a linear regression model using the provided rotation matrix and translation vector.
+
+    Args:
+        rotation (numpy.ndarray): Rotation matrix for the model.
+        translation (numpy.ndarray): Translation vector for the model.
+
+    Returns:
+        sklearn.linear_model.LinearRegression: Linear regression model with the specified rotation
+        and translation.
+    """
+    m = LinearRegression()
+    m.coef_ = rotation  # Set the rotation matrix as the model's coefficients
+    m.intercept_ = translation  # Set the translation vector as the model's intercept
+    return m  # Return the linear regression model
+
+
+def refine_local_warp(
+    X,
+    Y,
+    Y_pred,
+    threshold,
+    degree=2,
+    iterations=2,
+    min_correspondences=30,
+    model="polynomial",
+    smoothing=10.0,
+    max_correspondences=700,
+):
+    """Refine a global affine alignment with a local non-rigid warp.
+
+    Corrects residual within-tile distortion a single affine cannot capture (e.g. two
+    differently-configured microscopes). The warp is fit only on correspondences already
+    matched within `threshold`, then re-matched each iteration. With fewer than
+    `min_correspondences` confident matches it returns the input prediction unchanged.
+    "polynomial" (default) is backward-compatible; "thin_plate_spline" (scipy
+    RBFInterpolator, regularized by `smoothing`) is smoother and stable on two-scope data.
+
+    Args:
+        X (numpy.ndarray): Source (phenotype) coordinates, shape (n, 2).
+        Y (numpy.ndarray): Target (SBS) coordinates, shape (m, 2).
+        Y_pred (numpy.ndarray): Global-affine prediction of X in target space, shape (n, 2).
+        threshold (float): Match distance defining a high-confidence correspondence.
+        degree (int): Polynomial degree (model="polynomial"). Defaults to 2.
+        iterations (int): Number of refine-and-rematch passes. Defaults to 2.
+        min_correspondences (int): Minimum confident matches required to fit. Defaults to 30.
+        model (str): "polynomial" (default) or "thin_plate_spline".
+        smoothing (float): Thin-plate-spline regularization. Defaults to 10.0.
+        max_correspondences (int): Cap on thin-plate-spline anchor points. Defaults to 700.
+
+    Returns:
+        numpy.ndarray: Refined prediction of X in target space, shape (n, 2).
+    """
+    pf = PolynomialFeatures(degree)
+    refined = Y_pred
+    for _ in range(iterations):
+        distances = cdist(Y, refined, metric="sqeuclidean")
+        nearest = distances.argmin(axis=1)
+        within = np.sqrt(distances.min(axis=1)) < threshold
+        if within.sum() < min_correspondences:
+            break
+        Y_corr = Y[within]
+        if model == "thin_plate_spline":
+            from scipy.interpolate import RBFInterpolator
+
+            # Fit on the residual (prediction -> target): smoothing is tuned for that small scale
+            src = refined[nearest[within]]
+            if len(src) > max_correspondences:
+                sel = np.random.default_rng(0).choice(
+                    len(src), max_correspondences, replace=False
+                )
+                src, Y_corr = src[sel], Y_corr[sel]
+            rbf = RBFInterpolator(
+                src, Y_corr, kernel="thin_plate_spline", smoothing=smoothing
+            )
+            refined = rbf(refined)
+        else:
+            X_corr = X[nearest[within]]
+            reg = LinearRegression().fit(pf.fit_transform(X_corr), Y_corr)
+            refined = reg.predict(pf.transform(X))
+    return refined
diff --git a/workflow/scripts/merge/fast_alignment.py b/workflow/scripts/merge/fast_alignment.py
index d2186c13..cf072287 100644
--- a/workflow/scripts/merge/fast_alignment.py
+++ b/workflow/scripts/merge/fast_alignment.py
@@ -25,10 +25,7 @@
     "rotate_90": snakemake.params.alignment_rotate_90,
 }
 
-# metadata_align center-aligns the two scopes' stage coordinate frames (translation)
-# even when no flip/rotate is requested — needed when SBS and phenotype were acquired
-# on different microscopes with different coordinate origins. Backward compatible:
-# defaults False, and any requested flip/rotate still triggers alignment as before.
+# metadata_align center-aligns the two scopes' coordinate frames (translation only)
 metadata_align = getattr(snakemake.params, "metadata_align", False)
 
 # Apply alignment if center-alignment is requested or any flip/rotate is set
@@ -113,35 +110,36 @@
 score_thresh = snakemake.params.score
 
 
-# Optional alignment levers — backward compatible: absent config keys -> None -> drop ->
-# lib falls back to its literal defaults (existing screens unchanged).
+# Optional alignment levers; absent config keys drop out to lib defaults
 def _drop_none(d):
     return {k: v for k, v in d.items() if v is not None}
 
 
 ransac_kwargs = _drop_none(
     {
-        "residual_threshold": getattr(snakemake.params, "ransac_residual_threshold", None),
+        "residual_threshold": getattr(
+            snakemake.params, "ransac_residual_threshold", None
+        ),
         "max_trials": getattr(snakemake.params, "ransac_max_trials", None),
         "min_samples": getattr(snakemake.params, "ransac_min_samples", None),
         "random_state": getattr(snakemake.params, "ransac_random_state", None),
     }
 )
-evaluate_kwargs = _drop_none(
-    {
-        "threshold_triangle": getattr(snakemake.params, "threshold_triangle", None),
-        "threshold_point": getattr(snakemake.params, "threshold_point", None),
-        "threshold_region": getattr(snakemake.params, "threshold_region", None),
-        "ransac_kwargs": ransac_kwargs or None,
-    }
-) or None
+evaluate_kwargs = (
+    _drop_none(
+        {
+            "threshold_triangle": getattr(snakemake.params, "threshold_triangle", None),
+            "threshold_point": getattr(snakemake.params, "threshold_point", None),
+            "threshold_region": getattr(snakemake.params, "threshold_region", None),
+            "ransac_kwargs": ransac_kwargs or None,
+        }
+    )
+    or None
+)
 batch_size = getattr(snakemake.params, "batch_size", None)
 multistep_extra = {} if batch_size is None else {"batch_size": batch_size}
 
-# find-optimal-site (gated, default off): the stage-NEAREST phenotype tile is not always
-# the best real overlap when the two scopes have a metadata-alignment residual. When on,
-# offer the top-K nearest tiles per SBS tile as candidates and let the score/det filter
-# keep the single best-scoring one per site.
+# find-optimal-site (gated): try top-K nearest PH tiles per SBS seed, keep best per site
 seed_optimize = getattr(snakemake.params, "seed_optimize", False)
 seed_topk = getattr(snakemake.params, "seed_topk", None) or 3
 
@@ -179,13 +177,12 @@ def _drop_none(d):
     "@d0 <= determinant <= @d1 & score > @score_thresh"
 )
 
-# find-optimal-site: keep only the best-scoring phenotype tile per SBS site (higher score
-# == better real overlap), collapsing the top-K candidates expanded above back to one seed.
+# find-optimal-site: collapse top-K candidates to the best-scoring tile per site
 if seed_optimize:
     n_before = len(valid_pairs_df)
-    valid_pairs_df = valid_pairs_df.sort_values("score", ascending=False).drop_duplicates(
-        subset="site", keep="first"
-    )
+    valid_pairs_df = valid_pairs_df.sort_values(
+        "score", ascending=False
+    ).drop_duplicates(subset="site", keep="first")
     print(
         f"seed_optimize: kept best-scoring tile per site ({n_before} -> {len(valid_pairs_df)})"
     )
diff --git a/workflow/scripts/merge/fast_merge.py b/workflow/scripts/merge/fast_merge.py
index ec3e5c12..fc60790e 100644
--- a/workflow/scripts/merge/fast_merge.py
+++ b/workflow/scripts/merge/fast_merge.py
@@ -33,18 +33,19 @@
 print(f"Total alignments: {len(fast_alignment)}")
 print(f"Filtered alignments: {len(fast_alignment_filtered)}")
 
-# Optional local-warp levers — backward compatible: absent keys -> None -> dropped ->
-# refine_local_warp falls back to its literal defaults (existing screens unchanged). The
-# warp MODEL is selected by `local_refinement` ("polynomial" | "thin_plate_spline");
-# `smoothing`/`max_correspondences` only apply to the thin-plate-spline model.
+# Optional warp levers; absent keys fall back to refine_local_warp defaults
 warp_kwargs = {
     k: v
     for k, v in {
         "degree": getattr(snakemake.params, "warp_degree", None),
         "iterations": getattr(snakemake.params, "warp_iterations", None),
-        "min_correspondences": getattr(snakemake.params, "warp_min_correspondences", None),
+        "min_correspondences": getattr(
+            snakemake.params, "warp_min_correspondences", None
+        ),
         "smoothing": getattr(snakemake.params, "warp_smoothing", None),
-        "max_correspondences": getattr(snakemake.params, "warp_max_correspondences", None),
+        "max_correspondences": getattr(
+            snakemake.params, "warp_max_correspondences", None
+        ),
     }.items()
     if v is not None
 } or None

From f587ad7e3850b7725bb8d2919aa3b2248df94598 Mon Sep 17 00:00:00 2001
From: mat10d <mdiberna@wi.mit.edu>
Date: Fri, 10 Jul 2026 10:42:20 -0400
Subject: [PATCH 4/5] refactor(merge): expose only sweep-validated levers;
 hardcode inert knobs

Per review of the merge autoresearch results, keep as config-exposed only the
levers with measured impact (threshold_triangle, ransac_random_state,
local_refinement, warp_degree/iterations/smoothing, seed_optimize/seed_topk).
Revert the inert knobs to hardcoded literals in the lib (threshold_point=2,
threshold_region=50, batch_size=180, warp min/max_correspondences=30/700).
Defaults unchanged => byte-identical. TPS + find-optimal-site kept (validated
on Vaishnavi: median 0.08px).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 workflow/lib/merge/fast_merge.py         | 12 ++++--------
 workflow/lib/merge/hash.py               | 16 ++++++----------
 workflow/rules/merge.smk                 |  8 --------
 workflow/scripts/merge/fast_alignment.py | 14 +-------------
 workflow/scripts/merge/fast_merge.py     |  6 ------
 5 files changed, 11 insertions(+), 45 deletions(-)

diff --git a/workflow/lib/merge/fast_merge.py b/workflow/lib/merge/fast_merge.py
index 8fa662bb..11714ac6 100644
--- a/workflow/lib/merge/fast_merge.py
+++ b/workflow/lib/merge/fast_merge.py
@@ -28,8 +28,7 @@ def merge_triangle_hash(
             "polynomial" or "thin_plate_spline". Defaults None (off) — behaviour is then
             identical to the plain affine merge.
         warp_kwargs (dict | None): Keyword args forwarded to `refine_local_warp`
-            (degree, iterations, min_correspondences, smoothing, max_correspondences) when
-            local_refinement is on. Defaults None.
+            (degree, iterations, smoothing) when local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: The merged DataFrame.
@@ -79,8 +78,7 @@ def merge_sbs_phenotype(
             "polynomial" or "thin_plate_spline" selects the model. Defaults None (off) —
             prediction is then the plain affine, identical to before.
         warp_kwargs (dict | None, optional): Keyword args forwarded to `refine_local_warp`
-            (degree, iterations, min_correspondences, smoothing, max_correspondences) when
-            local_refinement is on. Defaults None.
+            (degree, iterations, smoothing) when local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: Table of merged identities of cell labels from cell_locations_0 and cell_locations_1.
@@ -183,10 +181,8 @@ def refine_local_warp(
     threshold,
     degree=2,
     iterations=2,
-    min_correspondences=30,
     model="polynomial",
     smoothing=10.0,
-    max_correspondences=700,
 ):
     """Refine a global affine alignment with a local non-rigid warp.
 
@@ -204,14 +200,14 @@ def refine_local_warp(
         threshold (float): Match distance defining a high-confidence correspondence.
         degree (int): Polynomial degree (model="polynomial"). Defaults to 2.
         iterations (int): Number of refine-and-rematch passes. Defaults to 2.
-        min_correspondences (int): Minimum confident matches required to fit. Defaults to 30.
         model (str): "polynomial" (default) or "thin_plate_spline".
         smoothing (float): Thin-plate-spline regularization. Defaults to 10.0.
-        max_correspondences (int): Cap on thin-plate-spline anchor points. Defaults to 700.
 
     Returns:
         numpy.ndarray: Refined prediction of X in target space, shape (n, 2).
     """
+    min_correspondences = 30  # min confident matches required to fit the warp
+    max_correspondences = 700  # cap on thin-plate-spline anchor points
     pf = PolynomialFeatures(degree)
     refined = Y_pred
     for _ in range(iterations):
diff --git a/workflow/lib/merge/hash.py b/workflow/lib/merge/hash.py
index c924543b..4ed0602a 100644
--- a/workflow/lib/merge/hash.py
+++ b/workflow/lib/merge/hash.py
@@ -196,7 +196,7 @@ def initial_alignment(
         well_triangles_0 (pandas.DataFrame): Hashed Delaunay triangulation for all tiles in dataset 0. Produced by concatenating outputs of `find_triangles` for individual tiles of a single well. Must include a `tile` column.
         well_triangles_1 (pandas.DataFrame): Hashed Delaunay triangulation for all sites in dataset 1. Produced by concatenating outputs of `find_triangles` for individual sites of a single well. Must include a `site` column.
         initial_sites (int | list[tuple[int, int]], optional): If an integer, specifies the number of sites sampled from `df_1` for initial brute-force matching of tiles to build the alignment model. If a list of 2-tuples, represents known (tile, site) matches to initialize the alignment model. At least 5 pairs are recommended.
-        evaluate_kwargs (dict, optional): Keyword args forwarded to `evaluate_match` (threshold_triangle, threshold_point, threshold_region, ransac_kwargs). Defaults to None.
+        evaluate_kwargs (dict, optional): Keyword args forwarded to `evaluate_match` (threshold_triangle, ransac_kwargs). Defaults to None.
 
     Returns:
         pandas.DataFrame: Table of possible (tile, site) matches, including rotation and translation transformations. Includes all tested matches, which should be filtered by `score` and `determinant` to retain valid matches.
@@ -234,8 +234,6 @@ def evaluate_match(
     vec_centers_0,
     vec_centers_1,
     threshold_triangle=0.3,
-    threshold_point=2,
-    threshold_region=50,
     ransac_kwargs=None,
 ):
     """Evaluates the match between two sets of vectors and centers.
@@ -246,9 +244,7 @@ def evaluate_match(
         vec_centers_0 (pandas.DataFrame): DataFrame containing the first set of vectors and centers.
         vec_centers_1 (pandas.DataFrame): DataFrame containing the second set of vectors and centers.
         threshold_triangle (float, optional): Threshold for matching triangles. Defaults to 0.3.
-        threshold_point (float, optional): Threshold for matching points. Defaults to 2.
-        threshold_region (float, optional): Region radius (px) within which a triangle center is scored. Defaults to 50.
-        ransac_kwargs (dict, optional): Keyword args forwarded to RANSACRegressor (e.g. residual_threshold, max_trials, min_samples, random_state). Defaults to None (sklearn defaults).
+        ransac_kwargs (dict, optional): Keyword args forwarded to RANSACRegressor (e.g. random_state). Defaults to None (sklearn defaults).
 
     Returns:
         tuple:
@@ -285,6 +281,8 @@ def evaluate_match(
     translation = model.estimator_.intercept_  # Extract translation vector
 
     # Score transformation based on the triangle centers
+    threshold_point = 2  # px radius for scoring a matched triangle center
+    threshold_region = 50  # px region considered when scoring
     distances = cdist(model.predict(c_0), c_1, metric="sqeuclidean")
     min_distances = np.sqrt(distances.min(axis=0))
     filt = min_distances < threshold_region
@@ -353,7 +351,6 @@ def multistep_alignment(
     det_range=(1.125, 1.186),
     score=0.1,
     initial_sites=8,
-    batch_size=180,
     n_jobs=None,
     evaluate_kwargs=None,
 ):
@@ -378,17 +375,16 @@ def multistep_alignment(
         initial_sites (int | list[tuple], optional): If int, the number of sites to sample from `well_triangles_1` for initial
             brute force matching to build a global alignment model. If a list of 2-tuples, represents known
             (tile, site) matches to start building the model. Defaults to 8.
-        batch_size (int, optional): Number of (tile, site) matches to evaluate per batch during global
-            alignment model updates. Defaults to 180.
         n_jobs (int, optional): Number of parallel jobs to deploy using joblib. Defaults to None.
         evaluate_kwargs (dict, optional): Keyword args forwarded to `evaluate_match` (threshold_triangle,
-            threshold_point, threshold_region, ransac_kwargs). Defaults to None.
+            ransac_kwargs). Defaults to None.
 
     Returns:
         pandas.DataFrame: Table of possible (tile, site) matches with corresponding rotation and translation
         transformations. All tested matches are included; query based on `score` and `determinant` to filter valid matches.
     """
     evaluate_kwargs = evaluate_kwargs or {}
+    batch_size = 180  # (tile, site) matches evaluated per global-alignment batch
 
     # If n_jobs is not provided, set it to one less than the number of CPU cores
     if n_jobs is None:
diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk
index bf41d6dc..be266175 100644
--- a/workflow/rules/merge.smk
+++ b/workflow/rules/merge.smk
@@ -50,13 +50,7 @@ if merge_approach == "fast":
             alignment_flip_y=config.get("merge", {}).get("alignment_flip_y"),
             alignment_rotate_90=config.get("merge", {}).get("alignment_rotate_90"),
             threshold_triangle=config.get("merge", {}).get("threshold_triangle"),
-            threshold_point=config.get("merge", {}).get("threshold_point"),
-            threshold_region=config.get("merge", {}).get("threshold_region"),
-            ransac_residual_threshold=config.get("merge", {}).get("ransac_residual_threshold"),
-            ransac_max_trials=config.get("merge", {}).get("ransac_max_trials"),
-            ransac_min_samples=config.get("merge", {}).get("ransac_min_samples"),
             ransac_random_state=config.get("merge", {}).get("ransac_random_state"),
-            batch_size=config.get("merge", {}).get("batch_size"),
             seed_optimize=config.get("merge", {}).get("seed_optimize", False),
             seed_topk=config.get("merge", {}).get("seed_topk"),
         script:
@@ -86,9 +80,7 @@ if merge_approach == "fast":
             local_refinement=config.get("merge", {}).get("local_refinement"),
             warp_degree=config.get("merge", {}).get("warp_degree"),
             warp_iterations=config.get("merge", {}).get("warp_iterations"),
-            warp_min_correspondences=config.get("merge", {}).get("warp_min_correspondences"),
             warp_smoothing=config.get("merge", {}).get("warp_smoothing"),
-            warp_max_correspondences=config.get("merge", {}).get("warp_max_correspondences"),
         script:
             "../scripts/merge/fast_merge.py"
 
diff --git a/workflow/scripts/merge/fast_alignment.py b/workflow/scripts/merge/fast_alignment.py
index cf072287..7f2d87b0 100644
--- a/workflow/scripts/merge/fast_alignment.py
+++ b/workflow/scripts/merge/fast_alignment.py
@@ -116,28 +116,17 @@ def _drop_none(d):
 
 
 ransac_kwargs = _drop_none(
-    {
-        "residual_threshold": getattr(
-            snakemake.params, "ransac_residual_threshold", None
-        ),
-        "max_trials": getattr(snakemake.params, "ransac_max_trials", None),
-        "min_samples": getattr(snakemake.params, "ransac_min_samples", None),
-        "random_state": getattr(snakemake.params, "ransac_random_state", None),
-    }
+    {"random_state": getattr(snakemake.params, "ransac_random_state", None)}
 )
 evaluate_kwargs = (
     _drop_none(
         {
             "threshold_triangle": getattr(snakemake.params, "threshold_triangle", None),
-            "threshold_point": getattr(snakemake.params, "threshold_point", None),
-            "threshold_region": getattr(snakemake.params, "threshold_region", None),
             "ransac_kwargs": ransac_kwargs or None,
         }
     )
     or None
 )
-batch_size = getattr(snakemake.params, "batch_size", None)
-multistep_extra = {} if batch_size is None else {"batch_size": batch_size}
 
 # find-optimal-site (gated): try top-K nearest PH tiles per SBS seed, keep best per site
 seed_optimize = getattr(snakemake.params, "seed_optimize", False)
@@ -210,7 +199,6 @@ def _drop_none(d):
     initial_sites=initial_sites,
     n_jobs=snakemake.threads,
     evaluate_kwargs=evaluate_kwargs,
-    **multistep_extra,
 )
 
 # Reset index
diff --git a/workflow/scripts/merge/fast_merge.py b/workflow/scripts/merge/fast_merge.py
index fc60790e..b6a10170 100644
--- a/workflow/scripts/merge/fast_merge.py
+++ b/workflow/scripts/merge/fast_merge.py
@@ -39,13 +39,7 @@
     for k, v in {
         "degree": getattr(snakemake.params, "warp_degree", None),
         "iterations": getattr(snakemake.params, "warp_iterations", None),
-        "min_correspondences": getattr(
-            snakemake.params, "warp_min_correspondences", None
-        ),
         "smoothing": getattr(snakemake.params, "warp_smoothing", None),
-        "max_correspondences": getattr(
-            snakemake.params, "warp_max_correspondences", None
-        ),
     }.items()
     if v is not None
 } or None

From 2138742877f7e5e85af0101f3b90407f6656954e Mon Sep 17 00:00:00 2001
From: mat10d <mdiberna@wi.mit.edu>
Date: Fri, 10 Jul 2026 11:07:09 -0400
Subject: [PATCH 5/5] feat(merge): let the configure-notebook preview apply the
 warp levers

plot_merge_example / fast_merge_example gain local_refinement + warp_kwargs and
apply refine_local_warp to the affine prediction, so the notebook match preview
reflects the levers an operator sets. Notebook-only helper; defaults None => unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 workflow/lib/merge/merge_utils.py | 34 +++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/workflow/lib/merge/merge_utils.py b/workflow/lib/merge/merge_utils.py
index b3a00cc2..ce1d83bc 100644
--- a/workflow/lib/merge/merge_utils.py
+++ b/workflow/lib/merge/merge_utils.py
@@ -20,7 +20,7 @@
 from skimage.measure import regionprops
 
 
-from lib.merge.fast_merge import build_linear_model
+from lib.merge.fast_merge import build_linear_model, refine_local_warp
 
 
 def plot_combined_tile_grid(
@@ -173,7 +173,9 @@ def _estimate_tile_size_from_coords(metadata):
     return min(x_spacing, y_spacing) if pd.notna(x_spacing) else 1000
 
 
-def plot_merge_example(df_ph, df_sbs, alignment_vec, threshold=2):
+def plot_merge_example(
+    df_ph, df_sbs, alignment_vec, threshold=2, local_refinement=None, warp_kwargs=None
+):
     """Visualizes the merge process for a single tile-site pair.
 
     Args:
@@ -181,6 +183,9 @@ def plot_merge_example(df_ph, df_sbs, alignment_vec, threshold=2):
         df_sbs (pandas.DataFrame): SBS data with 'i', 'j' columns.
         alignment_vec (dict): Contains 'rotation' and 'translation' for alignment.
         threshold (float, optional): Distance threshold for matching points. Defaults to 2.
+        local_refinement (str | bool | None, optional): Warp model to apply to the affine
+            prediction before matching, matching the pipeline (`refine_local_warp`). Defaults None.
+        warp_kwargs (dict | None, optional): Keyword args forwarded to `refine_local_warp`. Defaults None.
     """
     # Create figure with three subplots
     fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(30, 10))
@@ -197,6 +202,13 @@ def plot_merge_example(df_ph, df_sbs, alignment_vec, threshold=2):
     model = build_linear_model(alignment_vec["rotation"], alignment_vec["translation"])
     Y_pred = model.predict(X)
 
+    # Optional local warp; mirrors the pipeline so the preview reflects the levers
+    if local_refinement:
+        wk = dict(warp_kwargs or {})
+        if isinstance(local_refinement, str):
+            wk.setdefault("model", local_refinement)
+        Y_pred = refine_local_warp(X, Y, Y_pred, threshold, **wk)
+
     # Calculate distances
     distances = cdist(Y, Y_pred, metric="sqeuclidean")
     ix = distances.argmin(axis=1)
@@ -701,7 +713,14 @@ def find_closest_tiles(sbs_metadata, ph_metadata, sbs_tile_id, verbose=True):
 
 
 def fast_merge_example(
-    ph_tile, sbs_site, alignment_df, phenotype_info, sbs_info, threshold
+    ph_tile,
+    sbs_site,
+    alignment_df,
+    phenotype_info,
+    sbs_info,
+    threshold,
+    local_refinement=None,
+    warp_kwargs=None,
 ):
     """Process and plot PH tile and SBS site pairs."""
     print(f"\nProcessing PH tile {ph_tile} and SBS site {sbs_site}...")
@@ -728,7 +747,14 @@ def fast_merge_example(
 
     # Try plotting
     try:
-        plot_merge_example(phenotype_info, sbs_info, alignment_vec, threshold=threshold)
+        plot_merge_example(
+            phenotype_info,
+            sbs_info,
+            alignment_vec,
+            threshold=threshold,
+            local_refinement=local_refinement,
+            warp_kwargs=warp_kwargs,
+        )
         return True
     except Exception as e:
         print(f"  Error plotting: {str(e)}")