diff --git a/workflow/lib/merge/fast_merge.py b/workflow/lib/merge/fast_merge.py index 28a4c394..11714ac6 100644 --- a/workflow/lib/merge/fast_merge.py +++ b/workflow/lib/merge/fast_merge.py @@ -4,9 +4,17 @@ 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 +22,13 @@ 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, 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 diff --git a/workflow/lib/merge/hash.py b/workflow/lib/merge/hash.py index d63200be..4ed0602a 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, 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,15 +274,16 @@ 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 min_distances = np.sqrt(distances.min(axis=0)) filt = min_distances < threshold_region score = (min_distances[filt] < threshold_point).mean() # Calculate score @@ -343,8 +351,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. @@ -367,21 +375,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( { 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)}") diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk index 1d5e85dc..be266175 100644 --- a/workflow/rules/merge.smk +++ b/workflow/rules/merge.smk @@ -45,9 +45,14 @@ 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"), + ransac_random_state=config.get("merge", {}).get("ransac_random_state"), + seed_optimize=config.get("merge", {}).get("seed_optimize", False), + seed_topk=config.get("merge", {}).get("seed_topk"), script: "../scripts/merge/fast_alignment.py" @@ -72,6 +77,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_smoothing=config.get("merge", {}).get("warp_smoothing"), script: "../scripts/merge/fast_merge.py" diff --git a/workflow/scripts/merge/fast_alignment.py b/workflow/scripts/merge/fast_alignment.py index 876971bc..7f2d87b0 100644 --- a/workflow/scripts/merge/fast_alignment.py +++ b/workflow/scripts/merge/fast_alignment.py @@ -25,8 +25,11 @@ "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' 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 +if metadata_align or any( [ alignment_params["flip_x"], alignment_params["flip_y"], @@ -106,6 +109,29 @@ d0, d1 = snakemake.params.det_range score_thresh = snakemake.params.score + +# 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( + {"random_state": getattr(snakemake.params, "ransac_random_state", None)} +) +evaluate_kwargs = ( + _drop_none( + { + "threshold_triangle": getattr(snakemake.params, "threshold_triangle", None), + "ransac_kwargs": ransac_kwargs or None, + } + ) + or None +) + +# 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 + if initial_sbs_tiles is not None: # Auto-discover initial sites from SBS tiles candidate_pairs = [] @@ -113,10 +139,14 @@ 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 @@ -125,7 +155,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 @@ -133,6 +166,16 @@ "@d0 <= determinant <= @d1 & score > @score_thresh" ) +# 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") + 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( @@ -155,6 +198,7 @@ score=snakemake.params.score, initial_sites=initial_sites, n_jobs=snakemake.threads, + evaluate_kwargs=evaluate_kwargs, ) # Reset index diff --git a/workflow/scripts/merge/fast_merge.py b/workflow/scripts/merge/fast_merge.py index dc005882..b6a10170 100644 --- a/workflow/scripts/merge/fast_merge.py +++ b/workflow/scripts/merge/fast_merge.py @@ -33,6 +33,17 @@ print(f"Total alignments: {len(fast_alignment)}") print(f"Filtered alignments: {len(fast_alignment_filtered)}") +# 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), + "smoothing": getattr(snakemake.params, "warp_smoothing", 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 +61,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)