Cleanup + backward-compat for secondary_object: make compartment_combo optional

[mat10d]

Builds on #171 (secondary_object, @akcd1) and targets that branch — it does not rewrite or force-push #171. It bundles some cleanup and fixes to the secondary-object code, plus the backward-compatibility work needed before the feature can land on main. Happy to split any of this out or adjust if you'd prefer.

1. Correctness fixes

• cellpose training: augment in the image's native value range instead of clipping uint8 to [0, 1] (which saturated every nonzero pixel when intensity_scaling is on by default).
• secondary-object segmentation: use the h-maxima transform for peak suppression (was h_minima, which produced poor watershed seeds).
• aggregate: no longer aborts when a well has zero secondary objects; second-object merge-key dtypes are aligned to avoid a silent all-NaN join.

2. compartment_combo made optional (backward-compat)

The feature stamps a compartment_combo dimension onto every aggregate/cluster/bootstrap output path (_CmCo- token, /{combo}/ cluster dir, __{combo}__ bootstrap segment) and requires the column in the combo TSVs. For screens that don't use secondary objects this meant missing-column KeyErrors and a forced full re-run from changed paths.

New config flag aggregate.split_by_compartment (default false):

• Off → output paths are byte-identical to pre-feature, and a missing compartment_combo column is tolerated (defaulted to all compartments at TSV-read time).
• On → compartment-specific paths are emitted exactly as before.

Because column filtering is already a no-op when all compartments are present, default mode is a pure path/plumbing change with no effect on the data. New lib/shared/compartment_utils.py owns the conditional token/dir/suffix emission and a wildcard accessor; the Snakefile initializes the toggle and normalizes the combo tables at read; aggregate/bootstrap/cluster targets, rules, and metrics route through the helpers.

3. Alignment offsets treated as metadata by prefix

The per-cell offset_* columns (offset_y/x, offset_*_stepN, offset_*_cycleN) are alignment-QC diagnostics, not clustering features. Their names depend on each screen's alignment step/cycle count, so enumerating them in cell_data_metadata_cols.tsv is fragile — any unlisted offset_* column silently leaks into the feature space. A shared is_reserved_metadata_col predicate now reserves the offset_ prefix as metadata wherever features are split from metadata (split_cell_data, filter, perturbation_score), and the enumerated offsets are dropped from the small_test fixture (they're now recognized by prefix).

4. Secondary-object code tidy

• Moved the ML/classical parameter routing out of scripts/phenotype/identify_second_objs.py into a lib helper (segment_second_objs_from_config); the script is now a thin entry point (161 → 59 lines).
• segment_second_objs / create_empty_results always return a 3-tuple now (cytoplasm is None when none is given), so the callers' 3-value unpack is always consistent.
• Narrowed a bare except (→ IndexError) so only a sub-3D array falls back to channel 0.
• Post-declump fill_holes only fills background pixels, so a filled hole can't overwrite a neighboring object.
• cellpose_training uses a local RNG instead of the global seed, makes augmentation seedable, and keeps ≥1 training sample for a single-sample input.
• _aggregate_single preserves input cell order.

5. House-keeping

Removed the committed test file, collapsed multi-line comment blocks to single lines, moved to Google-style docstrings with public-first ordering in lib/, and stripped unused imports / duplicate functions / debug prints.

Validation (small_test_analysis)

• Compartment split on: 1221/1221 steps pass (the test config sets split_by_compartment: true, so the fixture exercises the feature, including both bootstrap branches + cluster).
• Default-off aggregate → cluster: 17/17 steps pass, producing legacy _ChCo-*__… filenames and cluster/{channel_combo}/{cell_class} directories; off-mode paths compared byte-for-byte against the parent.
• Ruff + format + py_compile clean.
• Note: the small_test config runs second_obj_detection: false, so the gate does not execute the secondary-object segmentation path. The section-4 changes there are verified by py_compile, ruff, and focused in-memory checks rather than the gate.

Follow-ups (not in this PR)

• A per-cell num_nuclei (multinucleation count) column — full wiring spec is in a follow-up comment below.
• The find_optimal_resolution / analyze_all_resolutions helpers had their signatures returned to the pre-feature shape; analysis notebooks that call them positionally will need a matching update on the analysis side.

🤖 Generated with Claude Code

[mat10d]

Follow-on (deferred from this PR): per-cell num_nuclei

Deferred while shipping the offset-metadata fix here. Keeping the full wiring spec with this PR so it's ready to pick up. Builds directly on the is_reserved_metadata_col predicate added in this branch.

---

Add a num_nuclei column to the per-cell phenotype table recording how many distinct
nuclei each cell contains, persisted through phenotype → merge → aggregate. Enables
finding multinucleate (e.g. binucleate) cells downstream.

Deferred from the offset-metadata work (#243). It is inert under the default consensus
reconcile (every surviving cell has exactly 1 nucleus), so it only carries information for
screens run with reconcile: contained_in_cells. Filing with full wiring so it's fast to pick up.

Design decisions (agreed)

• num_nuclei is per-cell metadata, not a clustering feature. It rides through merge/aggregate.
• Always present, default 1. Under consensus/1:1 reconcile it is 1; under
  contained_in_cells it is the real count. Stable schema across screens.
• Exact count: count distinct segmented nuclei, so two touching-but-distinct nuclei = 2.
  This requires reading the nuclei labels before reconcile merges/drops them — a
  connected-components count on the saved (post-reconcile) masks would fuse a touching pair, and
  contained_in_cells relabels all of a cell's nuclei to the cell id (so a post-hoc count = 1).

Why the wiring is non-trivial

The raw nuclei labels exist only inside each segmenter, right at its reconcile_nuclei_cells
call. The per-cell table is built later in a separate rule (extract_phenotype_cp) that only
receives the reconciled masks. So the count must be computed at the segmenter and handed across the
segment→phenotype rule boundary via one small per-tile output.

---

Wiring (exact approach)

1. workflow/lib/aggregate/cell_data_utils.py — reserve the name as metadata

Extend the predicate added in #243 so num_nuclei is always treated as metadata (like offset_):

RESERVED_METADATA_PREFIXES = ("offset_",)
RESERVED_METADATA_NAMES = frozenset({"num_nuclei"})

def is_reserved_metadata_col(col):
    return col in RESERVED_METADATA_NAMES or col.startswith(RESERVED_METADATA_PREFIXES)

No other metadata change needed — split_cell_data / filter / perturbation_score already
route through this predicate, so num_nuclei flows through aggregate as metadata automatically.

2. workflow/lib/shared/segmentation_utils.py — counting helper

Uses eroded pre-reconcile nuclei (touching nuclei separate) against the final cells, keyed to the
final cell label. center_pixels + regionprops are already imported here.

def count_nuclei_per_cell(nuclei, cells):
    """Count distinct nuclei contained in each cell.

    Uses pre-reconcile nuclei labels (eroded to separate touching nuclei) against the final
    cell labels, keyed by the final cell label.
    """
    nuclei_eroded = center_pixels(nuclei)
    counts = {}
    for region in regionprops(cells, intensity_image=nuclei_eroded):
        contained = region.intensity_image[region.intensity_image > 0]
        counts[region.label] = int(len(np.unique(contained)))
    return counts

3. The three segmenters — capture raw nuclei, compute, return it

reconcile_nuclei_cells stays untouched. At each reconcile call site, snapshot the raw nuclei
before reconcile and compute the count after. Add nuclei_per_cell to the return in the
return_counts=True / cells=True branch only.

• workflow/lib/shared/segment_cellpose.py — segment_cellpose_rgb (reconcile ~L500) and outer
  segment_cellpose (cells+return_counts branch ~L234):

# inside segment_cellpose_rgb, around the reconcile block:
raw_nuclei = nuclei.copy()
if reconcile:
    nuclei, cells = reconcile_nuclei_cells(nuclei, cells, how=reconcile)
nuclei_per_cell = count_nuclei_per_cell(raw_nuclei, cells)
...
return nuclei, cells, counts, nuclei_per_cell   # return_counts branch

# in outer segment_cellpose (cells + return_counts branch):
nuclei, cells, seg_counts, nuclei_per_cell = segment_cellpose_rgb(..., return_counts=True, ...)
...
return nuclei, cells, counts_df, nuclei_per_cell

• workflow/lib/shared/segment_stardist.py — segment_stardist (reconcile ~L244, return ~L253).
• workflow/lib/shared/segment_watershed.py — segment_watershed (reconcile ~L102, return ~L107).
  Same pattern: raw_nuclei = nuclei.copy() before reconcile, count_nuclei_per_cell after, append
  nuclei_per_cell to the return_counts=True return. (Note: stardist currently returns a bare
  counts dict where segment.py treats it as counts_df — pre-existing; just append consistently.)
  Import count_nuclei_per_cell in all three.

4. workflow/scripts/shared/segment.py — write the new per-tile output

Unpack the 4th value in each cells=True branch and write a {cell, num_nuclei} TSV to a new
snakemake.output[3]. For the cells=False branches, write an empty (header-only) frame.

# cells=True branches:
nuclei_data, cells_data, counts_df, nuclei_per_cell = segment_cellpose(...)
...
# after writing masks + counts:
nuclei_per_cell_df = pd.DataFrame(
    sorted(nuclei_per_cell.items()), columns=["cell", "num_nuclei"]
)
nuclei_per_cell_df.to_csv(snakemake.output[3], index=False, sep="\t")

5. Targets + rule — declare and thread the output

• workflow/targets/phenotype.smk: add a segment output (e.g. nuclei_per_cell TSV) and add it
  to extract_phenotype_cp's inputs.
• workflow/rules/phenotype.smk: add output[3] to rule segment; add the file as an input to
  rule extract_phenotype_cp.

6. workflow/scripts/phenotype/extract_phenotype.py — join by cell label

phenotype_cp is keyed by cell label (regionprops label → cell). Join the count; cells absent
from the map (shouldn't happen) default to 1.

nuclei_per_cell = pd.read_csv(snakemake.input[<n>], sep="\t")
cell_to_n = dict(zip(nuclei_per_cell["cell"], nuclei_per_cell["num_nuclei"]))
phenotype_cp["num_nuclei"] = phenotype_cp["cell"].map(cell_to_n).fillna(1).astype(int)

(Confirm the cell-label column name in phenotype_cp for the active cp_method — cell vs label.)

---

Leaner alternative (if the segmenter surgery isn't wanted)

Compute num_nuclei in segment.py from the saved nuclei+cells masks (connected components),
touching only ~4 files (segment.py + the phenotype threading, no segmenter changes). Exact under
consensus; only undercounts nuclei that physically touch with zero gap under
contained_in_cells. Chosen against because it isn't exact for touching nuclei, but it's the
low-risk option if multinucleation precision isn't critical.

Testing

• small_test_analysis uses cellpose → the cellpose path is gate-covered; stardist/watershed
  are not exercised by the gate, so verify those return-arity changes by hand.
• Add a contained_in_cells fixture (or a targeted unit test on count_nuclei_per_cell) to prove
  a 2-nucleus cell yields num_nuclei == 2 and rides through merge/aggregate as metadata.
• Per house rules, don't commit standalone test files; the gate is the committed guard.

[akcd1]

Code review

Found 1 issue:

1. DEFAULT_METADATA_COLS re-export removed from cell_data_utils.py in f43a39d, but metrics.py still imports it from there. It was a re-export (defined in lib/phenotype/constants.py), so it looked unused in the file it was deleted from. Both consumers use function-local imports, so py_compile and ruff do not catch it — get_all_stats() raises ImportError at call time via get_phenotype_stats(), and _calculate_batch_effects() breaks under include_batch_effects=True.

brieflow/workflow/lib/shared/metrics.py

Lines 145 to 147 in b23a896

	"""
	from lib.aggregate.cell_data_utils import DEFAULT_METADATA_COLS

brieflow/workflow/lib/shared/metrics.py

Lines 473 to 475 in b23a896

	"""Calculate batch effect metrics (pre/post alignment)."""
	from lib.aggregate.cell_data_utils import DEFAULT_METADATA_COLS
	from lib.shared.file_utils import load_parquet_subset

Reached unconditionally from get_all_stats():

brieflow/workflow/lib/shared/metrics.py

Lines 804 to 806 in b23a896

	print("\n[3/6] Gathering phenotype statistics...")
	stats["phenotype"] = get_phenotype_stats(config)
	print("\n PHENOTYPE STATISTICS:")

The constant lives here, so either restore the re-export or repoint both imports at lib.phenotype.constants:

brieflow/workflow/lib/phenotype/constants.py

Lines 2 to 5 in b23a896

	DEFAULT_METADATA_COLS = [
	"plate",
	"well",

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[akcd1]

Follow-up: the zero-secondary-objects guard

18b1b6d ("don't abort when a well has zero secondary objects") sits after the "none" short-circuit, so it only fires when the user explicitly asked for single/all/average:

brieflow/workflow/lib/aggregate/second_obj_utils.py

Lines 48 to 51 in b23a896

	if agg_strategy == "none":
	return cells_df.copy()

brieflow/workflow/lib/aggregate/second_obj_utils.py

Lines 56 to 60 in b23a896

	# No secondary objects in this well: keep all cells, nothing to merge
	if second_objs_df is None or second_objs_df.empty:
	return cells_df.copy()

Two problems with returning early there.

1. It drops the feature columns instead of NaN-filling them, which loses them screen-wide. aggregate_cells_second_objs runs per plate/well, and second_objs_filtered is a boolean slice of the full parquet, so an empty well still carries its columns. _validate_merge_keys passes, and _aggregate_single NaN-fills every cell — which is the documented contract ("NaN for cells with 0 or 2+ objects; no rows are dropped"):

brieflow/workflow/lib/aggregate/second_obj_utils.py

Lines 161 to 164 in b23a896

	# Add NaN columns for other cells
	if feature_cols:
	nan_df = pd.DataFrame(np.nan, index=other_cells.index, columns=feature_cols)
	other_cells = pd.concat([other_cells, nan_df], axis=1)

The early return instead emits a well whose parquet lacks those columns. harmonize_pool_schema keeps only the intersection across wells, so one empty well silently removes the secondary-object features from every well:

brieflow/workflow/lib/aggregate/filter.py

Lines 62 to 65 in b23a896

	per_file_schemas = {p: set(pq.read_schema(p).names) for p in paths}
	union_cols = set().union(*per_file_schemas.values())
	intersection_cols = set.intersection(*per_file_schemas.values())
	schema_mismatch = {

Its own log line anticipates exactly this ("per-file missingness — typically driven by class composition differences across wells"). A screen with a class-pure control well can lose all secondary-object features and still complete successfully.

2. It also swallows the case where the error was correct. A per-well slice retains columns, so the empty-well case was never the aborting one. The abort is reachable when the whole screen has no secondary objects and the upstream parquet has no schema to slice — and there, aborting is right: the user asked to integrate secondary objects into a screen that has none, and should set second_obj_agg_strategy: none. This change turns that misconfiguration into a silent no-op.

Suggested fix, splitting the two cases:

• Drop the per-well early return. Let empty wells flow through the merge and NaN-fill, per the docstring. Every well then emits the same columns and nothing is intersected away.
• Validate once at screen level instead: if second_obj_agg_strategy != "none" and the screen has no secondary objects at all, raise with a message pointing at second_obj_agg_strategy: none.

That keeps the loud failure for genuine misconfiguration without breaking a correctly-configured screen that has a legitimately empty control well.

🤖 Generated with Claude Code

[akcd1]

Docstrings (non-blocking)

Cosmetic, but the PR description claims more than landed.

The Google-style conversion skipped the largest file. 9d65852 converted file_utils.py, cellpose_training.py, and second_obj_utils.py cleanly, but segment_secondary_object.py is untouched in this respect — its only change in that commit was relocating two functions, bodies byte-identical. 12 of its 16 functions are still NumPy-style (segment_second_objs_ml L61, estimate_second_obj_diameter L655, apply_threshold_method L763, create_second_obj_standard_visualization L961, get_feret_diameters L1054, apply_morphological_opening L1087, apply_h_maxima_suppression L1119, apply_declumping L1151, shape_based_declumping L1273, create_empty_results L1356, get_spatial_overlap_candidates L1405, _postprocess_secondary_objects L1448).

brieflow/workflow/lib/phenotype/segment_secondary_object.py

Lines 1124 to 1137 in b23a896

	Parameters
	----------
	peak_map : ndarray
	Distance transform or intensity image
	h_factor : float
	Height threshold factor (0.0-1.0)
	h = h_factor * (peak_map.max() - peak_map.min())
	Higher values = more aggressive suppression
	Returns:
	-------
	filtered_map : ndarray
	Map with weak maxima suppressed

Against the target style, in the same file:

brieflow/workflow/lib/phenotype/segment_secondary_object.py

Lines 565 to 575 in b23a896

	Args:
	image (numpy.ndarray): Phenotype image (CHANNELS, I, J).
	cell_masks (numpy.ndarray): Cell segmentation masks.
	cytoplasm_masks (numpy.ndarray): Cytoplasm segmentation masks.
	second_obj_params (dict): Second-object parameters from config["phenotype"].
	nuclei_centroids (dict, optional): Cell nuclei centroids for distance calculations.
	Returns:
	tuple: (second_obj_masks, cell_second_obj_table, updated_cytoplasm_masks).
	"""

Worth noting why ruff passes: convention = "google" checks header punctuation, not body layout. The Returns:-with-dashes hybrid in the first link is the visible artifact — the colon got added, the NumPy sections underneath did not get rewritten.

Four return descriptions now contradict the code. segment_second_objs / segment_second_objs_ml / create_empty_results return an unconditional 3-tuple after this PR, with None as the third element when no cytoplasm masks are given. The docs still describe it as conditionally present ("if provided", "optionally"), which reads as a 2-tuple:

brieflow/workflow/lib/phenotype/segment_secondary_object.py

Lines 154 to 158 in b23a896

	- second_obj_masks: Labeled mask of secondary objects [height, width]
	- cell_second_obj_table: Dict with 'cell_summary' and 'second_obj_cell_mapping' DataFrames
	- updated_cytoplasm_masks: Cytoplasm masks with secondary objects removed (if provided)
	Raises:

brieflow/workflow/lib/phenotype/segment_secondary_object.py

Lines 403 to 412 in b23a896

	- second_obj_masks (numpy.ndarray): Labeled mask of secondary objects
	- cell_second_obj_table (dict): Dictionary with DataFrames containing associations
	- updated_cytoplasm_masks (numpy.ndarray): Updated cytoplasm masks (if cytoplasm_masks provided)
	If return_threshold_output=True:
	- second_obj_masks (numpy.ndarray): Labeled mask of secondary objects
	- cell_second_obj_table (dict): Dictionary with DataFrames containing associations
	- updated_cytoplasm_masks (numpy.ndarray): Updated cytoplasm masks (if cytoplasm_masks provided)
	- threshold_output (dict): Intermediate thresholding results with keys:
	- 'binary_mask': Binary mask before declumping

brieflow/workflow/lib/phenotype/segment_secondary_object.py

Lines 1368 to 1373 in b23a896

	Returns:
	-------
	tuple
	Empty secondary object masks, cell_second_obj_table dict, and optionally cytoplasm_masks
	"""
	cell_ids = np.unique(cell_masks[cell_masks > 0])

_postprocess_secondary_objects' own docstring already says "(or None)" and is correct — these four should match it.

One stale arg description. harmonize_pool_schema documents metadata_cols as "only those present in every file are returned", but the reserved offset_* columns appended by this PR were never in the caller's list, so the returned set can now include names outside it:

brieflow/workflow/lib/aggregate/filter.py

Lines 45 to 49 in b23a896

	Args:
	paths (list[str]): parquet file paths forming the pool.
	metadata_cols (list[str]): canonical metadata column names (only those
	present in every file are returned).
	drop_cols_threshold (float | None): pool-level NaN proportion threshold

🤖 Generated with Claude Code

[akcd1]

Taking these three myself so they don't sit — pushing to secondary-object-cleanup shortly:

1. The DEFAULT_METADATA_COLS import (repointing both call sites at lib.phenotype.constants).
2. The zero-secondary-objects guard, as proposed above: drop the per-well early return, add the screen-level check in aggregate_cells_second_objs.py where the unfiltered frame is still in scope.
3. The docstring items.

Shout if you're already mid-edit on any of these and I'll hold off.

[akcd1]

Pushed to secondary-object-cleanup (b23a896..69f86e7), one commit per item:

• 9398480 — fix(metrics): repoint both DEFAULT_METADATA_COLS imports at lib.phenotype.constants rather than restoring the re-export. constants.py has no imports of its own, so no cycle.
• 87ccb6a — fix(aggregate): drop the per-well early return so empty wells NaN-fill as documented, and add the screen-level guard in aggregate_cells_second_objs.py, where the unfiltered frame is still in scope. A non-none strategy against a screen with zero secondary objects now fails with a message naming second_obj_agg_strategy.
• 69f86e7 — docs: the 12 remaining NumPy-style docstrings in segment_secondary_object.py, the four stale return-arity entries, and the harmonize_pool_schema arg description.

Verification: ruff check and ruff format --check clean on all five files. For the docs commit, the ASTs with docstrings stripped are byte-identical to b23a896 for both files, so no code moved.

Not run: the small_test_analysis gate. Its config has second_obj_detection: false, so it would not exercise the aggregate guard either way — worth a targeted run against a second_obj_detection: true fixture if you have one handy.