diff --git a/README.md b/README.md index e7f7f4d..22bd346 100644 --- a/README.md +++ b/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 diff --git a/sccompPy/methods.py b/sccompPy/methods.py index 14a28d3..a0d2643 100644 --- a/sccompPy/methods.py +++ b/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, diff --git a/tests/test_metadata_input.py b/tests/test_metadata_input.py new file mode 100644 index 0000000..6e44c80 --- /dev/null +++ b/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)