Fixes

scripts/create_resources/combine/process_datasets_vizgen_nebius.sh

@@ -12,7 +12,14 @@ cd "$REPO_ROOT"
 
 set -e
 
-input_dir="s3://openproblems-data/resources/datasets"
+# Read the FRESH loader outputs + SC refs from the scratch raw folder, NOT the S3
+# resources/datasets copy. The S3 vizgen loader outputs are stale (Oct 2025, uint16):
+# the old loader rasterized ~700k cells into a uint16 label, so cell IDs > 65535 wrapped
+# and most cells were lost -> the segmentation covered only a partial band. The rebuilt
+# scratch loader outputs are uint32 / full-field. (process_dataset itself is correct; it
+# faithfully cropped the already-broken S3 label, which is why the band appeared.)
+input_dir="/scratch/task_ist_preprocessing/raw"
+#input_dir="s3://openproblems-data/resources/datasets"   # stale (uint16 labels) -- do not use
 #publish_dir="s3://openproblems-data/resources/task_ist_preprocessing/datasets"
 publish_dir='/scratch/task_ist_preprocessing/datasets'
 

src/data_processors/process_dataset/script.py

@@ -36,7 +36,21 @@
 ### VIASH END
 
 
-def get_crop_coords(sdata, max_n_pixels=20000*20000): #50000*50000): 
+def _full_res_image(img):
+    """Full-resolution image DataArray, whether `img` is single-scale (a plain
+    xarray DataArray) or multiscale (a DataTree with scale0/scale1/... nodes).
+
+    spatialdata stores an image either way depending on the loader, and indexing
+    a single-scale DataArray with ["scale0"] raises KeyError (it looks for a
+    coordinate named 'scale0'), so pick the top-res array explicitly.
+    """
+    from xarray import DataArray
+    if isinstance(img, DataArray):
+        return img
+    return img["scale0"].image
+
+
+def get_crop_coords(sdata, max_n_pixels=20000*20000): #50000*50000):
     """Get the crop coordinates to subset the sdata to max_n_pixels
 
     Arguments
@@ -52,7 +66,7 @@ def get_crop_coords(sdata, max_n_pixels=20000*20000): #50000*50000):
         The crop coordinates
     """
 
-    _, h, w = sdata['image']["scale0"].image.shape
+    _, h, w = _full_res_image(sdata['image']).shape
     #h, w = sdata
 
     # Check if the image is already below the maximum number of pixels

src/methods_transcript_assignment/proseg/script.py

@@ -23,7 +23,13 @@
   'samples': 200,
   'burnin_samples': 200,
   'voxel_size': 1.0,
-  'burnin_voxel_size': 2.0 
+  'burnin_voxel_size': 2.0,
+  'cell_compactness': 0.04,
+  'nuclear_reassignment_prob': 0.2,
+  'diffusion_probability': 0.2,
+  'ncomponents': 10,
+  'diffusion_sigma_far': 4.0,
+  'diffusion_sigma_near': 1.0
 }
 meta = {
   'name': 'proseg',
@@ -113,14 +119,14 @@
 
 command_suffix = (
     f'''--nthreads {n_threads} '''
-# these should not need to be changed
-    f'''--ncomponents 10 '''
+# model knobs (exposed as config arguments; defaults == proseg's own CLI defaults)
+    f'''--ncomponents {par['ncomponents']} '''
     # f'''--no-diffusion ''' # should be off by default (?)
-    f'''--diffusion-probability 0.2 '''
-    f'''--diffusion-sigma-far 4 '''
-    f'''--diffusion-sigma-near 1 '''
-    f'''--nuclear-reassignment-prob 0.2 '''
-    f'''--cell-compactness 0.03 '''
+    f'''--diffusion-probability {par['diffusion_probability']} '''
+    f'''--diffusion-sigma-far {par['diffusion_sigma_far']} '''
+    f'''--diffusion-sigma-near {par['diffusion_sigma_near']} '''
+    f'''--nuclear-reassignment-prob {par['nuclear_reassignment_prob']} '''
+    f'''--cell-compactness {par['cell_compactness']} '''
 # these can be changed as arguments
     f'''--voxel-layers {par['voxel_layers']} '''
     f'''--samples {par['samples']} '''

src/methods_transcript_assignment/segger/NOTES.md

@@ -54,14 +54,19 @@ Links: docs https://elihei2.github.io/segger_dev/ · repo https://github.com/dpe
    `cell_id` (prior; background → `"UNASSIGNED"`), plus dummy `qv=40.0` and
    `overlaps_nucleus`. The **prior `cell_id`** per transcript comes from looking up
    the label image at each transcript's integer pixel coords (truncated to int64,
-   then clipped to image bounds). `tx_pd` (clean RangeIndex) is the **canonical
-   transcript frame** over ALL transcripts. **Only the IN-BOUNDS transcripts are written
-   to the parquet** (see the empty-`bd`-batch gotcha): transcripts whose pixel coords fall
-   outside the label image are **excluded from segger's input** and stay unassigned (0) in
-   the output. So the parquet row order == the **in-bounds subset**, and segger's reported
-   `row_index` indexes that subset — `seg_orig_idx = np.nonzero(in_bounds)[0]` maps it back
-   to the full-frame position in step 5. On datasets whose segmentation covers the whole
-   transcript field this is a no-op (`n_oob == 0`, parquet == full frame).
+   then **clamped** to image bounds — a few can round a pixel past the raster edge, same
+   clamp as `basic_transcript_assignment` / `baysor` / `proseg`). `tx_pd` (clean RangeIndex)
+   is the **canonical transcript frame** over ALL transcripts, and **ALL of them are written
+   to the parquet**, so the parquet row order == `tx_pd` and segger's reported `row_index`
+   indexes it directly (step 5). Guards against a total coordinate-frame mismatch by raising
+   if *every* transcript clamps out of bounds (`n_oob == n_tx`).
+   **History:** this step formerly EXCLUDED out-of-bounds transcripts from the parquet and
+   remapped `row_index` through a `seg_orig_idx` array — a workaround for when the
+   process_dataset crop left ~38% of transcripts outside the cropped labels. That crop bug is
+   fixed (`crop_points_by_global_xy` crops transcripts to the same global box as the labels)
+   and a full run on re-processed Xenium + MERFISH confirmed `n_oob == 0`, so the exclusion
+   was a no-op and was removed in favour of the plain clamp. This is INDEPENDENT of the
+   empty-`bd`-batch crash (see that gotcha) — that one is intra-field, not OOB-driven.
 
 4. **Run segger** — `segger segment -i -o --n-epochs --prediction-graph-buffer-ratio
    --prediction-mode --node-representation-dim`, launched via a `run_segger(node_dim)`
@@ -265,30 +270,36 @@ Nextflow runner labels: `[hightime, midcpu, highmem, gpuh100]`.
   guards right after reading the segmentation (`int(seg_arr.data.max()) == 0 → raise` with an
   actionable message), mirroring pciseq's guard. Note the benchmark retries an exit-1 task 3×
   before ignoring it, so a guarded fail still burns 3 retries — re-processing is the cure.
-- **Transcripts outside the label window crash segger's encoder with an empty `bd` batch.**
+- **Empty-`bd`-batch crash — a cell-free tile, NOT out-of-bounds transcripts (STILL OPEN).**
   segger builds a heterogeneous graph with `'tx'` (transcript) and `'bd'` (boundary/cell)
   node types and applies a `Positional2dEmbedder` to **both**. That embedder does
-  `torch.zeros((batch.max()+1, 2))` — and `batch.max()` throws `RuntimeError: max(): Expected
-  reduction dim to be specified for input.numel() == 0` on an **empty** tensor. A training
-  mini-batch drawn entirely from **tiles that have transcripts but no boundaries** has zero
-  `'bd'` nodes → that empty `batch` → crash. Because the Lightning DataLoader shuffles, this
-  fires **mid-training** (e.g. epoch 8/19), not on batch 0 — so it looks intermittent. Source
-  of the transcript-only tiles: the combine step crops the image/labels to ~20000² but the
-  transcripts span a larger field, so a large fraction (observed **~38%**, 7.6M/19.7M on
-  mouse-brain rep3) sit outside the label. **Fix (`script.py`, in our adapter — no segger
-  patch):** compute an `in_bounds` mask and write **only in-bounds transcripts** to
-  `transcripts.parquet`; OOB transcripts (which can't belong to any segmented cell anyway)
-  stay unassigned (0) in the output, and segger's `row_index` is mapped back through
-  `seg_orig_idx = np.nonzero(in_bounds)[0]`. This also removes the earlier "40% clamped →
-  garbage prior" quality issue. **Fallback if empty-`bd` batches ever persist** (e.g. a
-  sparse in-coverage tile with transcripts but no cells): also patch segger's
-  `ist_encoder.py` `Positional2dEmbedder.forward` to no-op on `batch.numel()==0` (a
-  from_torch-style Docker `sed` shim), but prefer the adapter filter — it stops the degenerate
-  tile from forming rather than tolerating it. Separately worth chasing upstream: **why ~38%
-  of transcripts fall outside the label at all** — `sd.query.bounding_box` in the combine step
-  (`process_dataset/script.py:258`) should crop transcripts and labels together, so this may
-  be a combine/crop inconsistency affecting every assignment method (the basic ones just mark
-  those transcripts background and don't crash).
+  `torch.zeros((batch.max()+1, 2))` (`ist_encoder.py`) — and `batch.max()` throws
+  `RuntimeError: max(): Expected reduction dim to be specified for input.numel() == 0` on an
+  **empty** tensor. A mini-batch drawn from **tiles that have transcripts but no boundaries**
+  has zero `'bd'` nodes → empty `batch` → crash.
+  - **Original hypothesis (now disproven as the cause):** the transcript-only tiles came from
+    the combine step cropping labels to ~20000² while transcripts spanned a larger field
+    (~38%, 7.6M/19.7M on mouse-brain rep3, sat outside the label). The adapter worked around it
+    by EXCLUDING out-of-bounds transcripts from `transcripts.parquet`.
+  - **What actually happened:** the crop bug was fixed at the source (`crop_points_by_global_xy`
+    crops transcripts to the same global box as the labels). A full validation run (Seqera
+    `1HlVGwXktVakGm`) then showed segger **completes on Xenium** (mouse brain rep3, human breast)
+    but **still crashes on MERFISH** (`2022_vizgen_human_breast_cancer_merfish_combined/rep1`) —
+    with `0/19222237 transcripts fall outside the label image` (**n_oob == 0**), training all 80
+    batches, then dying at **validation tile ~21/28**. So the crash is **intra-field**: a tile
+    covering tissue that has transcripts but no *segmented cells* (MERSCOPE fields have acellular
+    / off-tissue gaps + a big extracellular transcript fraction; `custom_segmentation` copies the
+    vendor mask, so transcript-only tiles exist). Xenium coverage is dense → never trips.
+  - Because n_oob is 0, the OOB exclusion was a confirmed no-op and has been **removed**
+    (replaced by the plain edge clamp; see step 3). It never addressed this crash anyway.
+  - **The real fix (still to do — needs a container rebuild):** patch segger's
+    `Positional2dEmbedder.forward` in `ist_encoder.py` to no-op on empty input, e.g. return
+    `pos.new_zeros((pos.shape[0], 2 * self.dim))` when `pos.numel() == 0`, applied as a
+    from_torch-style Docker `sed`/python shim in `config.vsh.yaml`. This is a **segger source
+    patch (image rebuild)**, unlike the OOB simplification which is script-only. Verify the shim
+    on a live `build_main` pod (`debug-component-k8s`) before baking it in — the GPU/RAPIDS
+    rebuild is the expensive feedback loop. Not upstream in segger as of pinned `0233cf62`
+    (PR #74 is a *different* fix — Shapely polygon validity, not this).
 - **No CI test coverage** — GPU-only; can only be exercised on a CUDA host.
 - **Small panels break segger's gene-embedding PCA.** segger's `setup_anndata` does
   `PCA(n_components=in_channels=128)` on a genes×genes correlation matrix built from genes
@@ -352,8 +363,8 @@ So 0.05 is *not* a deviation; it is segger's current default. (`transcripts_key`
 ### Tiers
 
 - **Tier 0 — inputs, not parameters.** The biggest levers aren't in the sweep: the segmentation
-  prior fed in (which segmentation method, how good its boundaries) and the fraction of
-  transcripts that fall inside the label window (OOB transcripts are dropped — see the empty-`bd`
+  prior fed in (which segmentation method, how good its boundaries) and how densely those
+  boundaries cover the transcript field (sparse coverage → cell-free tiles → the empty-`bd`
   gotcha). Both are fixed by the upstream stage here.
 - **Tier 1 — highest impact on *what* gets assigned:**
   - **`prediction_mode`** {`nucleus`,`cell`,`uniform`} — which polygon set drives the prediction

src/methods_transcript_assignment/segger/script.py

@@ -160,35 +160,40 @@
     label_image = sdata_segm["segmentation"]["scale0"].image.to_numpy()
 else:
     label_image = sdata_segm["segmentation"].to_numpy()
-# Transcripts can land outside the label image after the transform (e.g. the combine step
-# crops the image/labels to ~20000x20000 but transcripts span a larger field -> here ~38%
-# of them fall outside). Those OOB transcripts must NOT be fed to segger: segger tiles the
-# transcript field into a heterogeneous graph with 'tx' (transcript) and 'bd' (boundary/
-# cell) node types, and a tile covering a transcripts-only region (no boundaries) produces a
-# training mini-batch with ZERO 'bd' nodes. segger's positional encoder then runs
-# `torch.zeros((batch.max()+1, 2))` on an empty batch -> `RuntimeError: max(): Expected
-# reduction dim to be specified for input.numel() == 0`, which kills training MID-RUN (only
-# when the DataLoader's shuffle happens to draw an all-OOB mini-batch, so it can survive
-# several epochs first). A transcript outside the segmentation can't belong to any segmented
-# cell anyway, so we EXCLUDE it from the segger input and leave it unassigned (cell_id 0) in
-# the final output. Datasets whose segmentation covers the whole transcript field are
-# unaffected (in_bounds all-True, n_oob == 0).
-in_bounds = (
-    (y_coords >= 0) & (y_coords < label_image.shape[0])
-    & (x_coords >= 0) & (x_coords < label_image.shape[1])
-)
-n_oob = int((~in_bounds).sum())
-y_look = np.clip(y_coords, 0, label_image.shape[0] - 1)
-x_look = np.clip(x_coords, 0, label_image.shape[1] - 1)
-prior_cell_id = label_image[y_look, x_look].astype(np.int64)  # 0 == background
-prior_cell_id[~in_bounds] = 0  # OOB -> background (excluded from segger regardless)
-print(f"{n_oob}/{len(in_bounds)} transcripts fall outside the "
-      f"{label_image.shape[0]}x{label_image.shape[1]} label image; excluding them from the "
-      f"segger input (they stay unassigned in the output).", flush=True)
+# Clamp transcript pixel coords to the label-image bounds before the lookup: after the
+# transform a few can round a pixel past the raster edge. Matches the clamp in
+# basic_transcript_assignment / baysor / proseg.
+#
+# HISTORY: this step used to EXCLUDE out-of-bounds transcripts from segger's input (and remap
+# segger's row_index back through a `seg_orig_idx` array), from when the process_dataset crop
+# left a large fraction (~38% on mouse-brain rep3) of transcripts OUTSIDE the cropped labels.
+# That crop bug is fixed -- crop_points_by_global_xy now crops transcripts to the SAME global
+# box as the labels -- and a full run on re-processed Xenium + MERFISH confirmed n_oob == 0, so
+# the exclusion + remap were a confirmed no-op. Removed here in favour of the plain clamp; edge
+# rounding is all that remains, and it's handled below. NOTE: this does NOT touch the separate
+# empty-`bd`-batch crash (a cell-free INTERIOR tile), which is intra-field, not OOB-driven, and
+# still needs the segger-encoder guard -- see NOTES.md.
+n_oob = int(np.count_nonzero(
+    (y_coords < 0) | (y_coords >= label_image.shape[0])
+    | (x_coords < 0) | (x_coords >= label_image.shape[1])
+))
+y_coords = np.clip(y_coords, 0, label_image.shape[0] - 1)
+x_coords = np.clip(x_coords, 0, label_image.shape[1] - 1)
+prior_cell_id = label_image[y_coords, x_coords].astype(np.int64)  # 0 == background
+print(f"Clamped {n_oob}/{len(prior_cell_id)} transcripts outside the "
+      f"{label_image.shape[0]}x{label_image.shape[1]} label image to its edge", flush=True)
+# Every transcript clamping out of bounds means transcripts and segmentation are almost
+# certainly in mismatched coordinate frames -- fail fast rather than assign everything to edge
+# pixels (mirrors the intent of the old zero-in-bounds guard).
+if n_oob == len(prior_cell_id):
+    raise ValueError(
+        "Every transcript falls outside the segmentation label image, so the clamp would map "
+        "them all to edge pixels. This usually means the transcripts and segmentation are in "
+        "mismatched coordinate frames (check the dataset's crop / transforms)."
+    )
 
 # Canonical transcripts frame (native coordinates, clean RangeIndex). Its row order matches
-# prior_cell_id / in_bounds; the IN-BOUNDS subset's order matches the parquet we write and
-# hence segger's reported row_index (mapped back through seg_orig_idx below).
+# prior_cell_id AND the parquet we write below, and hence segger's reported row_index.
 tx_pd = transcripts_reset.compute()
 n_tx = len(tx_pd)
 
@@ -197,30 +202,19 @@
 else:
     transcript_id = np.arange(n_tx, dtype=np.uint64)
 
-# Full-frame positions of the transcripts actually handed to segger. segger's output
-# row_index indexes into the (in-bounds) parquet row order; seg_orig_idx maps it back.
-seg_orig_idx = np.nonzero(in_bounds)[0]
-if seg_orig_idx.size == 0:
-    raise ValueError(
-        "No transcripts fall within the segmentation label image, so segger has nothing to "
-        "assign. This usually means the transcripts and segmentation are in mismatched "
-        "coordinate frames (check the dataset's crop / transforms)."
-    )
-
-n_in = int(in_bounds.sum())
 tx_out = pd.DataFrame({
-    "transcript_id": transcript_id[in_bounds],
-    "x_location": tx_pd["x"].to_numpy().astype(np.float32)[in_bounds],
-    "y_location": tx_pd["y"].to_numpy().astype(np.float32)[in_bounds],
-    "feature_name": tx_pd["feature_name"].astype(str).to_numpy()[in_bounds],
+    "transcript_id": transcript_id,
+    "x_location": tx_pd["x"].to_numpy().astype(np.float32),
+    "y_location": tx_pd["y"].to_numpy().astype(np.float32),
+    "feature_name": tx_pd["feature_name"].astype(str).to_numpy(),
     # segger's Xenium loader keys transcripts to boundaries by cell_id string;
     # background (0) becomes the UNASSIGNED sentinel.
-    "cell_id": np.where(prior_cell_id[in_bounds] > 0, prior_cell_id[in_bounds].astype(str), "UNASSIGNED"),
-    "qv": np.full(n_in, 40.0, dtype=np.float32),
-    "overlaps_nucleus": (prior_cell_id[in_bounds] > 0).astype(np.int8),
+    "cell_id": np.where(prior_cell_id > 0, prior_cell_id.astype(str), "UNASSIGNED"),
+    "qv": np.full(n_tx, 40.0, dtype=np.float32),
+    "overlaps_nucleus": (prior_cell_id > 0).astype(np.int8),
 })
 if "z" in tx_pd.columns:
-    tx_out.insert(3, "z_location", tx_pd["z"].to_numpy().astype(np.float32)[in_bounds])
+    tx_out.insert(3, "z_location", tx_pd["z"].to_numpy().astype(np.float32))
 tx_out.to_parquet(XENIUM_DIR / "transcripts.parquet", index=False)
 del transcripts, y_coords, x_coords, label_image
 
@@ -341,15 +335,15 @@ def run_segger(node_dim):
 seg = seg[keep_mask]
 print(f"segger kept {len(seg)} assignments of {n_tx} transcripts", flush=True)
 
-# segger reports row_index into the transcripts.parquet we wrote, which is the IN-BOUNDS
-# subset (OOB transcripts were excluded above). Map that back to the full-frame position via
-# seg_orig_idx, then factorize the (string) segger_cell_id into contiguous positive integers;
-# unassigned / out-of-bounds transcripts stay 0.
+# segger reports row_index into the transcripts.parquet we wrote. Since we now write ALL
+# transcripts (the full canonical frame, no OOB exclusion), row_index indexes tx_pd directly.
+# Factorize the (string) segger_cell_id into contiguous positive integers; unassigned
+# transcripts stay 0.
 cell_id_per_tx = np.zeros(n_tx, dtype=np.int64)
 if len(seg):
     row_idx = seg["row_index"].to_numpy().astype(np.int64)
     codes, _ = pd.factorize(seg["segger_cell_id"].astype(str), sort=True)
-    cell_id_per_tx[seg_orig_idx[row_idx]] = codes.astype(np.int64) + 1
+    cell_id_per_tx[row_idx] = codes.astype(np.int64) + 1
 
 #############################################
 # Build transcript-assignment output object #