"""Unit tests for the DSL operators. Each operator gets a tiny hand-checkable case; no network, no parquet I/O. """ from __future__ import annotations import numpy as np import pandas as pd import pytest from dsl import ( Apply, Associate, Cohort, Effect, Fit, Reduce, Search, Select, Split, ) # --- fixtures --------------------------------------------------------------- def _tiny_cohort() -> Cohort: """Six samples, five features, with known structure. `geneA` is high in samples 4..6 and low in 1..3; `geneB` is the reverse. `msi_status` is MSI-H in 4..6 and MSS in 1..3 so a single-gene rule should perfectly separate the two groups. """ idx = pd.Index([f"s{i}" for i in range(1, 7)], name="sample_id") expr = pd.DataFrame( { "geneA": [1.0, 1.0, 1.0, 9.0, 9.0, 9.0], "geneB": [9.0, 9.0, 9.0, 1.0, 1.0, 1.0], "geneC": [5.0, 5.0, 5.0, 5.0, 5.0, 5.0], "geneD": [2.0, 3.0, 4.0, 5.0, 6.0, 7.0], "geneE": [7.0, 6.0, 5.0, 4.0, 3.0, 2.0], }, index=idx, ) clinical = pd.DataFrame( { "stage": ["I", "II", "III", "II", "III", "I"], "age": [55.0, 60.0, 65.0, 70.0, 50.0, 75.0], }, index=idx, ) labels = pd.DataFrame( {"msi_status": ["MSS"] * 3 + ["MSI-H"] * 3}, index=idx, ) return Cohort(expression=expr, clinical=clinical, labels=labels) # --- Select / Reduce ------------------------------------------------------- def test_select_picks_columns_in_order(): c = _tiny_cohort() out = Select(c.expression, ["geneB", "geneA"]) assert list(out.columns) == ["geneB", "geneA"] assert out.shape == (6, 2) def test_select_missing_column_raises(): c = _tiny_cohort() with pytest.raises(KeyError): Select(c.expression, ["geneA", "geneZ"]) def test_reduce_mean_collapses_rows(): c = _tiny_cohort() score = Reduce(Select(c.expression, ["geneA", "geneB"]), agg="mean") # (1+9)/2 = 5 for every row. assert list(score.values) == pytest.approx([5.0] * 6) def test_reduce_supports_full_agg_vocab(): c = _tiny_cohort() for agg in ("mean", "median", "max", "min", "var"): s = Reduce(c.expression, agg=agg) assert list(s.index) == list(c.expression.index) def test_reduce_rejects_unknown_agg(): c = _tiny_cohort() with pytest.raises(ValueError): Reduce(c.expression, agg="prod") # --- Split ----------------------------------------------------------------- def test_split_partitions_by_predicate(): c = _tiny_cohort() a, b = Split(c, lambda co: co.labels["msi_status"] == "MSI-H") assert list(a.sample_ids) == ["s4", "s5", "s6"] assert list(b.sample_ids) == ["s1", "s2", "s3"] # --- Associate ------------------------------------------------------------- def test_associate_pearson_perfect_positive(): s = pd.Series([1.0, 2.0, 3.0, 4.0]) assert Associate(s, s * 2 + 7, kind="pearson") == pytest.approx(1.0) def test_associate_pearson_perfect_negative(): s = pd.Series([1.0, 2.0, 3.0, 4.0]) assert Associate(s, -s, kind="pearson") == pytest.approx(-1.0) def test_associate_spearman_monotone_only(): # Strictly monotone but non-linear -> spearman 1, pearson < 1. s = pd.Series([1.0, 2.0, 3.0, 4.0]) t = s ** 3 assert Associate(s, t, kind="spearman") == pytest.approx(1.0) assert Associate(s, t, kind="pearson") < 1.0 def test_associate_unknown_kind_raises(): s = pd.Series([1.0, 2.0]) with pytest.raises(ValueError): Associate(s, s, kind="kendall") # --- Effect ---------------------------------------------------------------- def test_effect_partial_correlation_removes_confounder(): """y caused entirely by z; x also caused by z but no direct link to y. Unadjusted x-vs-y correlation should be ~1, partial correlation given z ~0. """ rng = np.random.default_rng(0) n = 400 z = rng.standard_normal(n) x = z + 0.01 * rng.standard_normal(n) y = z + 0.01 * rng.standard_normal(n) idx = pd.Index([f"s{i}" for i in range(n)], name="sample_id") res = Effect( pd.Series(x, index=idx), pd.Series(y, index=idx), pd.DataFrame({"z": z}, index=idx), ) assert res.unadjusted > 0.95 assert abs(res.partial_corr) < 0.2 assert res.n_used == n def test_effect_drops_nans_and_reports_n_used(): idx = pd.Index([f"s{i}" for i in range(5)], name="sample_id") cause = pd.Series([1.0, 2.0, 3.0, 4.0, np.nan], index=idx) effect = pd.Series([2.0, 4.0, 6.0, np.nan, 10.0], index=idx) adj = pd.DataFrame({"stage": ["I", "II", "II", "III", "I"]}, index=idx) res = Effect(cause, effect, adj) assert res.n_used == 3 # --- Search ---------------------------------------------------------------- def test_search_rejects_non_opaque_columns(): df = pd.DataFrame({"GENE1": [1.0], "GENE2": [2.0]}) with pytest.raises(ValueError, match="opaque IDs"): Search(df, lambda s: float(s.sum()), 1) def test_search_ranks_by_objective_and_takes_top_k(): df = pd.DataFrame({ "g00001": [1.0, 1.0, 1.0], "g00002": [1.0, 2.0, 3.0], "g00003": [0.0, 0.0, 9.0], }) # objective = variance — g00003 highest, g00002 next, g00001 zero. top = Search(df, lambda s: float(s.var()), 2) assert top == ["g00003", "g00002"] # --- Fit / Apply ----------------------------------------------------------- def test_fit_separable_features_score_above_random(): """Two linearly separable classes -> AUROC = 1.0, balanced acc = 1.0.""" rng = np.random.default_rng(0) n = 60 x_pos = rng.standard_normal(n) + 3.0 x_neg = rng.standard_normal(n) - 3.0 state = pd.Series(np.concatenate([x_pos, x_neg]), name="x") y = np.array([1] * n + [0] * n) res = Fit(state, y, random_state=0) assert res.auroc == pytest.approx(1.0) assert res.balanced_acc == pytest.approx(1.0) assert res.n_train + res.n_test == 2 * n def test_apply_returns_probabilities_indexed_like_input(): rng = np.random.default_rng(0) n = 40 state = pd.Series( np.concatenate([rng.standard_normal(n) + 2, rng.standard_normal(n) - 2]), index=[f"s{i}" for i in range(2 * n)], name="x", ) y = np.array([1] * n + [0] * n) res = Fit(state, y, random_state=0) proba = Apply(res, state) assert list(proba.index) == list(state.index) assert ((proba >= 0.0) & (proba <= 1.0)).all()