Skip to content

fix(etype): homogenize _etype per entity after every merge - #4

Merged
atuldeshpande merged 3 commits into
mainfrom
fix/etype-homogeneity-after-merge
Jun 13, 2026
Merged

fix(etype): homogenize _etype per entity after every merge#4
atuldeshpande merged 3 commits into
mainfrom
fix/etype-homogeneity-after-merge

Conversation

@atuldeshpande

Copy link
Copy Markdown
Member

Problem

Long found that Stitch was not enforcing the cell–cell merge gate (can_merge is supposed to block merging two entities that both contain a cell).

Root cause: merge stages (phase1_maha_remerge, Stitch) can leave a single resulting entity with heterogeneous _etype across its transcript rows. Downstream, build_entity_table derives has_cell from groupby(entity)["_etype"].first() — which is non-deterministic on a heterogeneous entity. So a merged entity can be silently misclassified, and the cell–cell merge block (can_merge) can be bypassed.

Fix

Enforce a single _etype per entity immediately after each merge, only for the affected entities (O(merges), not O(N)):

  • _etype.py — add homogenize_etype_for_entity() with a fixed priority cell > partial > component > drop > unknown. Idempotent; no-op when an entity is already homogeneous.
  • phase1_rescue.py — after the Maha-remerge DSU union+remap, homogenize each affected root.
  • stitching.py — detect stitched labels that became heterogeneous (groupby(out_col)["_etype"].nunique() > 1) and homogenize each; removes the old non-deterministic .first()-based propagation (and the categorical fast-path early-return that skipped it).

Verification

  • PDAC 500 µm ROI: heterogeneous-_etype entities go 160 → 0 at Maha-remerge and 268 → 0 at Stitch, at every stage; the first=cell count stabilizes (1632 → 1651).
  • No regression where no merges fire: on samples with zero merges at these stages (e.g. MERFISH), output is bit-identical — the change only acts on entities that were actually merged.

🤖 Generated with Claude Code

Closes the silent failure path in stitching.can_merge's cell-cell gate
on the PDAC pipeline. The gate checks `has_cell[i] = (etypes[i] == "cell")`
where `etypes` comes from `summary_df.groupby(entity_col)["_etype"].first()`
— a non-deterministic operator. When upstream stages merge entities of
different types without updating `_etype` on the absorbed rows, the
merged entity carries mixed _etype, `.first()` picks whichever row
happens to come first in the df, and the gate's view of the entity
becomes a coin-flip.

Two stages had this defect:

1. phase1_maha_remerge — performs DSU unions on cell+partial pairs to
   correct mis-assigned tx between overlapping nuclei, but only remaps
   `entity_col`. The 160 merges on PDAC 500µm ROI left 160 mixed-etype
   entities behind; downstream `.first()` would silently miscode 17
   of them at the Stitch gate.

2. apply_stitching_to_transcripts_memory_efficient — has a propagation
   step (lines 2790-2828) but it (a) iterates summary["entity_id"] which
   filters out non-cell entities by default (etype_filter=("cell",)),
   missing heterogeneity that arises from non-cell entities sharing a
   stitched label, AND (b) is bypassed by the early-return on the
   categorical fast path (the runner's default), so the propagation
   never fires in production. Both gaps fixed.

Adds a shared helper `tracer._etype.homogenize_etype_for_entity(df,
entity_label, *, entity_col, etype_col)` that picks the highest-
priority _etype present in the group (cell > component > partial >
drop > unknown) and stamps it on every row. Idempotent; cheap; O(merges).

Verified on PDAC 500µm ROI:
- Pre-fix: 160 het at Phase1-Maha-Remerge, 268 het at Stitch
- Post-fix: 0 het at every stage
- Cell count preserved: first=cell 1632 → 1651 at Stitch (matches the
  pre-Maha-Remerge ground truth)
- 29/29 phase1_maha_remerge + stitch_mahalanobis_rescue + phase1_rerank
  tests still pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR enforces a single, deterministic transcript-level _etype per entity immediately after merge operations (Phase1 Maha-remerge and Stitch) to prevent downstream entity classification from becoming non-deterministic and bypassing the cell–cell merge gate.

Changes:

  • Add a shared _etype homogenization helper with an explicit priority rule (cell > partial > component > drop > unknown).
  • Invoke _etype homogenization after DSU-based remaps in phase1_maha_remerge.
  • Remove Stitch’s categorical-mode early return so post-stitch _etype homogenization can run, and homogenize stitched labels that become heterogeneous.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File Description
src/tracer/_etype.py Adds homogenize_etype_for_entity() and a priority rule for resolving heterogeneous _etype within an entity.
src/tracer/phase1_rescue.py Calls _etype homogenization after Maha-remerge DSU union+remap for affected merged roots.
src/tracer/stitching.py Ensures Stitch doesn’t skip _etype homogenization and homogenizes stitched labels that are heterogeneous post-merge.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread src/tracer/stitching.py Outdated
Comment on lines +2809 to +2813
if "_etype" in df_out.columns:
from ._etype import homogenize_etype_for_entity
# Find stitched labels whose rows have heterogeneous _etype.
_SENT = {"-1", "DROP", "UNASSIGNED", "nan"}
labels = df_out[out_col].astype(str)
Comment thread src/tracer/stitching.py
Comment on lines +2793 to +2799
# Homogenize `_etype` on every stitched label that ended up
# heterogeneous. Without this, rows of merged-in entities keep their
# original (now-stale) etype, and downstream
# `df.groupby(entity_col)["_etype"].first()` becomes non-deterministic
# — silently miscoding a merged entity and breaking the cell-cell
# merge gate in any subsequent stitch pass.
#
Comment thread src/tracer/phase1_rescue.py Outdated
Comment on lines +330 to +336
if "_etype" in df.columns:
from ._etype import homogenize_etype_for_entity
affected_roots = {dsu.find(p["a"]) for p in pairs_rescued}
for root in affected_roots:
homogenize_etype_for_entity(
df, root, entity_col=entity_col, etype_col="_etype",
)
Comment on lines +323 to +329
# Homogenize _etype on every affected merged entity. Without this
# the merged entity carries mixed-_etype rows (cell tx from one
# side, partial tx from the other), and downstream
# `groupby(entity_col)["_etype"].first()` becomes non-deterministic
# — silently miscoding the entity for the cell-cell merge gate in
# Stitch. See `tracer._etype.homogenize_etype_for_entity` for the
# priority rule.
Comment thread src/tracer/_etype.py
Comment on lines +203 to +220
def homogenize_etype_for_entity(
df: pd.DataFrame,
entity_label: str,
*,
entity_col: str = "tracer_id",
etype_col: str = "_etype",
) -> None:
"""In-place: ensure every row of ``entity_label`` shares a single
``_etype`` value (the highest-priority etype present in the group).

No-op when:
- ``etype_col`` is not in df
- no row matches ``entity_label``
- the entity is already homogeneous

Designed to be called right after a DSU union + entity_col remap
by the stage that did the merge. Cheap and idempotent.
"""
Addresses the Copilot review on #4.

Performance:
- `homogenize_etype_for_entities`: resolve many entities in a single pass
  over only the affected rows (O(rows-in-affected) instead of
  O(N x #roots) from per-root full-column rescans). `homogenize_etype_for_entity`
  now delegates to it.
- Phase1 Maha-remerge and Stitch call the batch form once over all
  affected/heterogeneous roots.
- Stitch: gate the O(N) heterogeneity scan behind a `merges_happened`
  flag (derived from the category mapping), so the no-merge categorical
  fast path skips it entirely.

Tests:
- `_etype`: no-op on homogeneous, priority selection (cell>partial>
  component>drop>unknown), absent label, missing column, batch
  independence, empty labels.
- Maha-remerge: cell+partial merge homogenizes the merged entity to cell.
- Stitch: cell+partial merge homogenizes the stitched label to cell
  (the original cell-cell-gate failure mode).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Comment thread src/tracer/_etype.py
Comment on lines +210 to +222
"""In-place batch homogenization: give every entity in
``entity_labels`` a single ``_etype`` (the highest-priority etype
present within that entity) in **one pass over only the affected
rows**.

Cost is O(rows in the affected entities), not O(N · #entities) — the
column is scanned/masked once rather than re-scanned per entity.
Prefer this at merge call sites that touch many roots.

No-op when ``etype_col`` is absent, ``entity_labels`` is empty, or no
row matches. Idempotent: passing already-homogeneous entities (or
re-running) leaves the column unchanged.
"""
Comment thread src/tracer/_etype.py Outdated
Comment on lines +232 to +233
sub_ent = ent.to_numpy()[mask]
sub_et = df[etype_col].astype(str).to_numpy()[mask]
Comment thread src/tracer/stitching.py Outdated
Comment on lines +2815 to +2819
# rows under the same stitched label (when a cell entity's id equals
# a partial's parent prefix). So we homogenize all heterogeneous
# stitched labels rather than only those reachable via summary —
# still O(heterogeneous labels), not O(N). See `tracer._etype.
# homogenize_etype_for_entity` for the priority rule.
Comment thread src/tracer/stitching.py Outdated
Comment on lines 2766 to 2768
df_out[out_col] = pd.Categorical.from_codes(out_codes, categories=new_categories)
if debug_stages and debug_legacy_col != out_col:
df_out[debug_legacy_col] = df_out[out_col].copy()
Comment on lines +332 to +335
affected_roots = {dsu.find(p["a"]) for p in pairs_rescued}
# One pass over only the merged roots' rows (not per-root).
homogenize_etype_for_entities(
df, affected_roots, entity_col=entity_col, etype_col="_etype",
…ments

Second Copilot pass on #4:
- homogenize_etype_for_entities: mask first, then string-cast only the
  affected `_etype` rows (was casting the whole column).
- Correct the docstring/comments to state the true cost: one O(N) scan
  of entity_col to locate affected rows, resolution over only those rows
  (not "one pass over only affected rows").
- Stitch: drop the now-redundant `debug_legacy_col` write in the
  categorical branch (the shared post-mapping block already writes it).
- Stitch: comment now states the heterogeneity detection is an O(N_tx)
  groupby scan gated on `merges_happened`, not "O(het labels)".

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@atuldeshpande
atuldeshpande merged commit c2cc41f into main Jun 13, 2026
8 checks passed
@atuldeshpande
atuldeshpande deleted the fix/etype-homogeneity-after-merge branch June 13, 2026 21:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants