merge: expose fast-merge config levers + TPS warp + optimal-site seeding (backward-compatible)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "BrieFlow"
-version = "1.4.13"
+version = "1.4.14"
 description = "Package for processing data from optical poooled screens."
 readme = "README.md"
 license = { file = "LICENSE" }

workflow/lib/merge/fast_merge.py

@@ -4,16 +4,31 @@
 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:
         hash_df_0 (pandas.DataFrame): The first DataFrame.
         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, 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, smoothing) when local_refinement is on. Defaults None.
 
     Returns:
         pandas.DataFrame: The merged DataFrame.
@@ -25,27 +40,24 @@ 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)
-
-
-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
+    return merge_sbs_phenotype(
+        hash_df_0,
+        hash_df_1,
+        model,
+        threshold=threshold,
+        local_refinement=local_refinement,
+        warp_kwargs=warp_kwargs,
+    )
 
 
-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 +73,12 @@ 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 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, 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.
@@ -98,6 +116,13 @@ 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 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):
+            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")
 
@@ -130,3 +155,84 @@ def merge_sbs_phenotype(cell_locations_0, cell_locations_1, model, threshold=2):
             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,
+    model="polynomial",
+    smoothing=10.0,
+):
+    """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.
+        model (str): "polynomial" (default) or "thin_plate_spline".
+        smoothing (float): Thin-plate-spline regularization. Defaults to 10.0.
+
+    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):
+        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

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, 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,10 @@ 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,
+    ransac_kwargs=None,
 ):
     """Evaluates the match between two sets of vectors and centers.
 
@@ -237,7 +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.
+        ransac_kwargs (dict, optional): Keyword args forwarded to RANSACRegressor (e.g. random_state). Defaults to None (sklearn defaults).
 
     Returns:
         tuple:
@@ -267,19 +274,19 @@ 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
     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")
-    threshold_region = 50  # Threshold for the region to consider
-    filt = np.sqrt(distances.min(axis=0)) < threshold_region
-    score = (
-        np.sqrt(distances.min(axis=0))[filt] < threshold_point
-    ).mean()  # Calculate score
+    min_distances = np.sqrt(distances.min(axis=0))
+    filt = min_distances < threshold_region
+    score = (min_distances[filt] < threshold_point).mean()  # Calculate score
 
     return rotation, translation, score  # Return rotation, translation, and score
 
@@ -337,8 +344,8 @@ def multistep_alignment(
     det_range=(1.125, 1.186),
     score=0.1,
     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.
 
@@ -361,21 +368,26 @@ 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,
+            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:
         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(
             {

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,14 +173,19 @@ 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:
         df_ph (pandas.DataFrame): Phenotype data with 'i', 'j' columns.
         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)}")