Support cell-level metadata input in sccompPy

README.md

@@ -261,6 +261,20 @@ estimate_res = sccompPy.sccomp_estimate(
     count = 'count',
     verbose = False
 )
+```
+
+Cell-level metadata can also be supplied directly. When the selected `count`
+column is absent, `sccompPy` derives counts by grouping rows by `sample` and
+`cell_group`, while retaining sample-level metadata columns for use in formulas:
+
+```python
+estimate_res = sccompPy.sccomp_estimate(
+    data = cell_metadata,
+    formula_composition = '~ 0 + type',
+    sample = 'sample',
+    cell_group = 'cell_group',
+    verbose = False
+)
 ```
 
     17:05:31 - cmdstanpy - INFO - CmdStan start processing

sccompPy/methods.py

@@ -3,6 +3,25 @@
 from .utilities import *
 from .functions_multi_beta_binomia import *
 
+def _aggregate_cell_metadata(data, sample, cell_group, count):
+    """Aggregate one row per cell into the count table expected by sccomp."""
+    groups = data.groupby([sample, cell_group], sort=False, observed=True)
+    aggregated = groups.size().rename(count).reset_index()
+
+    metadata_columns = []
+    for column in data.columns:
+        if column in (sample, cell_group, count):
+            continue
+        if data.groupby(sample, observed=True)[column].nunique(dropna=False).le(1).all():
+            metadata_columns.append(column)
+
+    if metadata_columns:
+        sample_metadata = data[[sample] + metadata_columns].drop_duplicates(sample)
+        aggregated = aggregated.merge(sample_metadata, on=sample, how="left", sort=False)
+
+    return aggregated
+
+
 def sccomp_estimate(
     data,
     formula_composition="~1",
@@ -37,6 +56,15 @@ def sccomp_estimate(
     variational_inference=None,
     **kwargs
 ):
+    sample = sample or "sample"
+    cell_group = cell_group or "cell_group"
+    if count is None:
+        count = count_col
+    count = count or "count"
+
+    if count not in data.columns:
+        data = _aggregate_cell_metadata(data, sample, cell_group, count)
+
     res = sccomp_glm_data_frame_counts(
         data,
         formula_composition=formula_composition,

tests/test_metadata_input.py

@@ -0,0 +1,42 @@
+import pandas as pd
+from pandas.testing import assert_frame_equal
+
+from sccompPy import methods
+
+
+def test_sccomp_estimate_aggregates_cell_metadata(monkeypatch):
+    cell_metadata = pd.DataFrame(
+        {
+            "sample": ["s1", "s1", "s1", "s2"],
+            "cell_group": ["A", "A", "B", "A"],
+            "type": ["control", "control", "control", "treated"],
+            "batch": [1, 1, 1, 2],
+            "cell_id": ["c1", "c2", "c3", "c4"],
+        }
+    )
+    captured = {}
+
+    def fake_sccomp_glm_data_frame_counts(data, **kwargs):
+        captured["data"] = data
+        return {}
+
+    monkeypatch.setattr(methods, "sccomp_glm_data_frame_counts", fake_sccomp_glm_data_frame_counts)
+
+    methods.sccomp_estimate(
+        cell_metadata,
+        sample="sample",
+        cell_group="cell_group",
+        verbose=False,
+    )
+
+    expected = pd.DataFrame(
+        {
+            "sample": ["s1", "s1", "s2"],
+            "cell_group": ["A", "B", "A"],
+            "count": [2, 1, 1],
+            "type": ["control", "control", "treated"],
+            "batch": [1, 1, 2],
+        }
+    )
+    assert captured["data"].columns.tolist() == expected.columns.tolist()
+    assert_frame_equal(captured["data"], expected)