Secondary object

tests/small_test_analysis/config/aggregate_combo.tsv

@@ -1,3 +1,5 @@
-cell_class	channel_combo	plate	well
-Interphase	DAPI_COXIV_CENPA_WGA	1	A1
-Interphase	DAPI_COXIV_CENPA_WGA	1	A2
+cell_class	channel_combo	compartment_combo	plate	well
+Interphase	DAPI_COXIV_CENPA_WGA	cell-nucleus-cytoplasm	1	A1
+Interphase	DAPI_COXIV_CENPA_WGA	cell-nucleus-cytoplasm	1	A2
+Interphase	DAPI	nucleus	1	A1
+Interphase	DAPI	nucleus	1	A2

tests/small_test_analysis/config/cluster_combo.tsv

@@ -1,2 +1,3 @@
-cell_class	channel_combo	leiden_resolution
-Interphase	DAPI_COXIV_CENPA_WGA	2
+cell_class	channel_combo	compartment_combo	leiden_resolution
+Interphase	DAPI_COXIV_CENPA_WGA	cell-nucleus-cytoplasm	2
+Interphase	DAPI	nucleus	2

tests/small_test_analysis/config/config.yml

@@ -96,6 +96,7 @@ phenotype:
   mode: null
   nuclei_diameter: 41.97813156768016
   reconcile: contained_in_cells
+  second_obj_detection: false
 merge:
   approach: fast
   merge_combo_fp: config/merge_combo.tsv
@@ -172,6 +173,10 @@ aggregate:
   bootstrap_combinations:
   - cell_class: Interphase
     channel_combo: DAPI_COXIV_CENPA_WGA
+    compartment_combo: cell-nucleus-cytoplasm
+  - cell_class: Interphase
+    channel_combo: DAPI
+    compartment_combo: nucleus
 cluster:
   cluster_combo_fp: config/cluster_combo.tsv
   uniprot_data_fp: config/benchmark_clusters/uniprot_data.tsv

tests/test_compartment_filter.py

@@ -0,0 +1,161 @@
+"""Unit tests for compartment filtering in the aggregate module."""
+
+import sys
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+# Make the workflow's `lib` package importable without installing the package
+WORKFLOW_LIB_DIR = Path(__file__).resolve().parents[1] / "workflow"
+sys.path.insert(0, str(WORKFLOW_LIB_DIR))
+
+from lib.aggregate.cell_data_utils import (
+    COMPARTMENT_PREFIXES,
+    SECOND_OBJ_EXTRA_COLS,
+    compartment_combo_subset,
+)
+
+ALL_COMPARTMENTS_4 = ["cell", "nucleus", "cytoplasm", "second_obj"]
+ALL_COMPARTMENTS_3 = ["cell", "nucleus", "cytoplasm"]
+
+
+def _make_features_4c():
+    return pd.DataFrame(
+        {
+            "cell_DAPI_mean": [1.0],
+            "cell_area": [2.0],
+            "nucleus_DAPI_mean": [3.0],
+            "nucleus_area": [4.0],
+            "cytoplasm_DAPI_mean": [5.0],
+            "cytoplasm_area": [6.0],
+            "second_obj_DAPI_mean": [7.0],
+            "second_obj_area": [8.0],
+            "total_second_obj_area": [10.0],
+            "mean_second_obj_diameter": [11.0],
+            "mean_distance_to_nucleus": [12.0],
+        }
+    )
+
+
+def test_all_compartments_kept():
+    df = _make_features_4c()
+    out = compartment_combo_subset(df, ALL_COMPARTMENTS_4, ALL_COMPARTMENTS_4)
+    assert list(out.columns) == list(df.columns)
+
+
+def test_nucleus_only_drops_other_prefixes():
+    df = _make_features_4c()
+    out = compartment_combo_subset(df, ["nucleus"], ALL_COMPARTMENTS_4)
+    assert set(out.columns) == {"nucleus_DAPI_mean", "nucleus_area"}
+
+
+def test_cell_plus_nucleus():
+    df = _make_features_4c()
+    out = compartment_combo_subset(df, ["cell", "nucleus"], ALL_COMPARTMENTS_4)
+    expected = {"cell_DAPI_mean", "cell_area", "nucleus_DAPI_mean", "nucleus_area"}
+    assert set(out.columns) == expected
+
+
+def test_second_obj_exclusion_drops_extra_cols():
+    df = _make_features_4c()
+    out = compartment_combo_subset(df, ALL_COMPARTMENTS_3, ALL_COMPARTMENTS_4)
+    for c in SECOND_OBJ_EXTRA_COLS:
+        assert c not in out.columns
+    assert "second_obj_DAPI_mean" not in out.columns
+    assert "cell_area" in out.columns
+
+
+def test_second_obj_included_keeps_extra_cols():
+    df = _make_features_4c()
+    out = compartment_combo_subset(df, ["second_obj"], ALL_COMPARTMENTS_4)
+    for c in SECOND_OBJ_EXTRA_COLS:
+        assert c in out.columns
+
+
+def test_compartment_prefixes_constant_shape():
+    assert set(COMPARTMENT_PREFIXES) == {"cell", "nucleus", "cytoplasm", "second_obj"}
+    for name, prefix in COMPARTMENT_PREFIXES.items():
+        assert prefix == f"{name}_"
+
+
+from lib.aggregate.cell_data_utils import resolve_aggregate_combos
+
+
+def test_resolve_defaults_4_when_detection_true():
+    out = resolve_aggregate_combos([{"channels": ["DAPI"]}], second_obj_detection=True)
+    assert out == [
+        {
+            "channels": ["DAPI"],
+            "compartments": ["cell", "nucleus", "cytoplasm", "second_obj"],
+        }
+    ]
+
+
+def test_resolve_defaults_3_when_detection_false():
+    out = resolve_aggregate_combos([{"channels": ["DAPI"]}], second_obj_detection=False)
+    assert out == [
+        {"channels": ["DAPI"], "compartments": ["cell", "nucleus", "cytoplasm"]}
+    ]
+
+
+def test_resolve_rejects_unknown_compartment():
+    with pytest.raises(ValueError, match="unknown compartment"):
+        resolve_aggregate_combos(
+            [{"channels": ["DAPI"], "compartments": ["nucleous"]}],
+            second_obj_detection=True,
+        )
+
+
+def test_resolve_rejects_second_obj_when_detection_false():
+    with pytest.raises(ValueError, match="second_obj_detection"):
+        resolve_aggregate_combos(
+            [{"channels": ["DAPI"], "compartments": ["second_obj"]}],
+            second_obj_detection=False,
+        )
+
+
+def test_resolve_rejects_empty_compartments():
+    with pytest.raises(ValueError, match="at least one compartment"):
+        resolve_aggregate_combos(
+            [{"channels": ["DAPI"], "compartments": []}],
+            second_obj_detection=True,
+        )
+
+
+def test_resolve_rejects_empty_channels():
+    with pytest.raises(ValueError, match="at least one channel"):
+        resolve_aggregate_combos(
+            [{"channels": [], "compartments": ["cell"]}],
+            second_obj_detection=True,
+        )
+
+
+def test_resolve_dedupes_compartments_within_combo():
+    out = resolve_aggregate_combos(
+        [{"channels": ["DAPI"], "compartments": ["cell", "cell", "nucleus"]}],
+        second_obj_detection=True,
+    )
+    assert out[0]["compartments"] == ["cell", "nucleus"]
+
+
+def test_resolve_dedupes_duplicate_combos():
+    out = resolve_aggregate_combos(
+        [
+            {"channels": ["DAPI"], "compartments": ["cell"]},
+            {"channels": ["DAPI"], "compartments": ["cell"]},
+        ],
+        second_obj_detection=True,
+    )
+    assert len(out) == 1
+
+
+def test_resolve_preserves_distinct_combos():
+    out = resolve_aggregate_combos(
+        [
+            {"channels": ["DAPI"], "compartments": ["cell"]},
+            {"channels": ["DAPI"], "compartments": ["nucleus"]},
+        ],
+        second_obj_detection=True,
+    )
+    assert len(out) == 2

workflow/lib/aggregate/aggregate.py

@@ -16,6 +16,7 @@ def aggregate(
     method="mean",
     ps_probability_threshold=None,
     ps_percentile_threshold=None,
+    carry_cols: list[str] | None = None,
 ) -> tuple[np.ndarray, pd.DataFrame]:
     """Apply mean or median aggregation to replicate embeddings and perturbation scores for each perturbation.
 
@@ -31,6 +32,13 @@ def aggregate(
             Defaults to "mean".
         ps_probability_threshold (float, optional): Threshold for filtering based on perturbation score.
         ps_percentile_threshold (float, optional): Percentile threshold for filtering based on perturbation score.
+        carry_cols (list[str], optional): Metadata columns functionally determined by
+            `pert_col` to preserve (one value per group) in the output. Typical use:
+            when `pert_col` is a construct ID (e.g. `cell_barcode_0`), carry the
+            human-readable gene symbol (`gene_symbol_0`) through so downstream
+            clustering / benchmarking / lookups can match by gene name. Raises
+            ValueError if a carry_col is missing from metadata or has more than one
+            unique value within a group.
 
     Returns:
         tuple:
@@ -48,6 +56,21 @@ def aggregate(
     if aggr_func is None:
         raise ValueError(f"Invalid aggregation method: {method}")
 
+    if carry_cols is None:
+        carry_cols = []
+    else:
+        missing = [c for c in carry_cols if c not in metadata.columns]
+        if missing:
+            raise ValueError(
+                f"carry_cols not found in metadata: {missing}. "
+                f"Available columns: {list(metadata.columns)}"
+            )
+        overlap = [c for c in carry_cols if c == pert_col]
+        if overlap:
+            raise ValueError(
+                f"carry_cols overlap with pert_col: {overlap}. Remove duplicates."
+            )
+
     # filter by ps_probability_threshold; keep NaNs
     if ps_probability_threshold is not None:
         mask = metadata["perturbation_score"].isna() | (
@@ -81,6 +104,16 @@ def aggregate(
         if "perturbation_auc" in metadata.columns:
             agg_meta["perturbation_auc"] = group["perturbation_auc"].iloc[0]
 
+        # Carry through columns that are functionally determined by pert_col.
+        for c in carry_cols:
+            nuniq = group[c].nunique(dropna=False)
+            if nuniq > 1:
+                raise ValueError(
+                    f"carry_col {c!r} has {nuniq} unique values within group "
+                    f"{pert_col}={pert!r}; not functionally determined by {pert_col}."
+                )
+            agg_meta[c] = group[c].iloc[0]
+
         if ps_probability_threshold is not None or ps_percentile_threshold is not None:
             # aggregate perturbation score with same function
             pert_score = (

workflow/lib/aggregate/cell_data_utils.py

@@ -180,3 +180,123 @@ def get_feature_table_cols(feature_cols):
     selected_columns.extend(overlap_cols)
 
     return selected_columns
+
+
+COMPARTMENT_PREFIXES = {
+    "cell": "cell_",
+    "nucleus": "nucleus_",
+    "cytoplasm": "cytoplasm_",
+    "second_obj": "second_obj_",
+}
+
+# Per-cell summary columns added by aggregate_second_obj_data that don't carry
+# the second_obj_ prefix but are derived from second-object data.
+SECOND_OBJ_EXTRA_COLS = frozenset(
+    {
+        "total_second_obj_area",
+        "mean_second_obj_diameter",
+        "mean_distance_to_nucleus",
+    }
+)
+
+
+def compartment_combo_subset(
+    features: pd.DataFrame, compartment_combo: str, all_compartments: list[str]
+) -> pd.DataFrame:
+    """Filter features to keep only columns belonging to the requested compartments.
+
+    Args:
+        features (pd.DataFrame): Feature columns (metadata already split out).
+        compartment_combo (list[str]): Compartments to keep, e.g. ["cell", "nucleus"].
+        all_compartments (list[str]): Compartments present in the input. Used to
+            determine which prefixes to exclude.
+
+    Returns:
+        pd.DataFrame: features without columns belonging to excluded compartments.
+    """
+    excluded = [c for c in all_compartments if c not in compartment_combo]
+    excluded_prefixes = [COMPARTMENT_PREFIXES[c] for c in excluded]
+
+    cols_to_drop = [
+        col
+        for col in features.columns
+        if any(col.startswith(p) for p in excluded_prefixes)
+    ]
+    if "second_obj" in excluded:
+        cols_to_drop += [c for c in SECOND_OBJ_EXTRA_COLS if c in features.columns]
+
+    return features.drop(columns=cols_to_drop)
+
+
+def resolve_aggregate_combos(
+    aggregate_combos: list[dict], second_obj_detection: bool
+) -> list[dict]:
+    """Normalize and validate AGGREGATE_COMBOS entries.
+
+    For each entry:
+    - Fills in default compartments (all 4 if detection on, 3 otherwise) when missing.
+    - Dedupes compartments within the combo (preserving order).
+    - Validates compartment names, non-empty channels/compartments, and
+      second_obj-vs-detection consistency.
+    After normalization, duplicate (channels, compartments) pairs are deduped.
+
+    Args:
+        aggregate_combos (list[dict]): Each dict has "channels" (list[str]) and
+            optionally "compartments" (list[str]).
+        second_obj_detection (bool): From config["phenotype"]["second_obj_detection"].
+
+    Returns:
+        list[dict]: Normalized, validated, and de-duplicated combos.
+
+    Raises:
+        ValueError: On any validation failure.
+    """
+    valid_compartments = set(COMPARTMENT_PREFIXES)
+    default_compartments = (
+        ["cell", "nucleus", "cytoplasm", "second_obj"]
+        if second_obj_detection
+        else ["cell", "nucleus", "cytoplasm"]
+    )
+
+    resolved = []
+    seen = set()
+    for idx, combo in enumerate(aggregate_combos):
+        channels = list(combo.get("channels") or [])
+        if not channels:
+            raise ValueError(
+                f"AGGREGATE_COMBOS[{idx}] must specify at least one channel"
+            )
+
+        comps = combo.get("compartments")
+        if comps is None:
+            comps = list(default_compartments)
+        else:
+            comps = list(comps)
+        if not comps:
+            raise ValueError(
+                f"AGGREGATE_COMBOS[{idx}] must specify at least one compartment"
+            )
+
+        # Dedupe compartments while preserving order
+        deduped = []
+        for c in comps:
+            if c not in valid_compartments:
+                raise ValueError(
+                    f"AGGREGATE_COMBOS[{idx}] has unknown compartment {c!r}; "
+                    f"must be one of {sorted(valid_compartments)}"
+                )
+            if c == "second_obj" and not second_obj_detection:
+                raise ValueError(
+                    f"AGGREGATE_COMBOS[{idx}] lists 'second_obj' but "
+                    f"config['phenotype']['second_obj_detection'] is False"
+                )
+            if c not in deduped:
+                deduped.append(c)
+
+        key = (tuple(channels), tuple(deduped))
+        if key in seen:
+            continue
+        seen.add(key)
+        resolved.append({"channels": channels, "compartments": deduped})
+
+    return resolved