Update tests

Author: jaydu1

tests/test_DR_learner.py

@@ -3,74 +3,63 @@
 import pytest
 from causarray.DR_learner import LFC
 
-# package/tests/test_DR_learner.py
-
 
 @pytest.fixture
 def sample_data():
-    np.random.seed(0)
-    Y = np.random.poisson(5, (100, 10))
-    W = np.random.normal(0, 1, (100, 5))
-    A = np.random.binomial(1, 0.5, 100)
-    Y += np.random.poisson(1, (100, 10)) * A[:, None]
+    rng = np.random.default_rng(0)
+    Y = rng.poisson(5, (100, 10)).astype(float)
+    W = rng.standard_normal((100, 5))
+    A = rng.binomial(1, 0.5, 100)
+    Y += rng.poisson(1, (100, 10)) * A[:, None]
     return Y, W, A
 
-def test_LFC_basic(sample_data):
-    Y, W, A = sample_data
-    result, estimation = LFC(Y, W, A)
-    assert isinstance(result, pd.DataFrame)
-    assert 'tau' in result.columns
-    assert 'std' in result.columns
-    assert 'stat' in result.columns
-    assert 'rej' in result.columns
-    assert 'pvalue' in result.columns
-    assert 'padj' in result.columns
-    assert 'pvalue_emp_null_adj' in result.columns
-    assert 'padj_emp_null_adj' in result.columns
 
-def test_LFC_with_offset(sample_data):
-    Y, W, A = sample_data
-    offset = np.log(np.random.poisson(5, 100))
-    result, estimation = LFC(Y, W, A, offset=offset)
-    assert isinstance(result, pd.DataFrame)
-    assert 'tau' in result.columns
+class TestLFCOutputSchema:
+    def test_output_columns(self, sample_data):
+        Y, W, A = sample_data
+        result, estimation = LFC(Y, W, A)
+        assert isinstance(result, pd.DataFrame)
+        for col in ('tau', 'std', 'stat', 'rej', 'pvalue', 'padj',
+                    'pvalue_emp_null_adj', 'padj_emp_null_adj'):
+            assert col in result.columns
 
-# def test_LFC_with_cross_est(sample_data):
-#     Y, W, A = sample_data
-#     result, estimation = LFC(Y, W, A, cross_est=True)
-#     assert isinstance(result, pd.DataFrame)
-#     assert 'tau' in result.columns
+    def test_with_offset(self, sample_data):
+        Y, W, A = sample_data
+        rng = np.random.default_rng(1)
+        offset = np.log(rng.poisson(5, 100).clip(1))
+        result, estimation = LFC(Y, W, A, offset=offset)
+        assert isinstance(result, pd.DataFrame)
+        assert 'tau' in result.columns
 
-def test_LFC_with_fdx(sample_data):
-    Y, W, A = sample_data
-    result, estimation = LFC(Y, W, A, fdx=True)
-    assert isinstance(result, pd.DataFrame)
-    assert 'tau' in result.columns
+    def test_with_fdx(self, sample_data):
+        Y, W, A = sample_data
+        result, estimation = LFC(Y, W, A, fdx=True)
+        assert isinstance(result, pd.DataFrame)
+        assert 'tau' in result.columns
 
-def test_LFC_with_custom_family(sample_data):
-    Y, W, A = sample_data
-    result, estimation = LFC(Y, W, A, family='poisson')
-    assert isinstance(result, pd.DataFrame)
-    assert 'tau' in result.columns
+    def test_with_custom_family(self, sample_data):
+        Y, W, A = sample_data
+        result, estimation = LFC(Y, W, A, family='poisson')
+        assert isinstance(result, pd.DataFrame)
+        assert 'tau' in result.columns
 
 
-def test_LFC_multi_treatment():
-    """Test LFC with multiple perturbations (exercises per-perturbation fast path)."""
-    np.random.seed(42)
-    n, p, a = 200, 50, 3
-    W = np.random.normal(0, 1, (n, 3))
-    # One-hot treatment: first 50 cells = control, rest split among 3 perturbations
-    A = np.zeros((n, a))
-    A[50:100, 0] = 1
-    A[100:150, 1] = 1
-    A[150:200, 2] = 1
-    Y = np.random.poisson(5, (n, p))
-    # Add treatment effects for perturbation 0
-    Y[50:100] += np.random.poisson(2, (50, p))
+class TestLFCMultiTreatment:
+    def test_multi_treatment(self):
+        """LFC with multiple perturbations exercises the per-perturbation fast path."""
+        np.random.seed(42)
+        n, p, a = 200, 50, 3
+        W = np.random.normal(0, 1, (n, 3))
+        A = np.zeros((n, a))
+        A[50:100, 0] = 1
+        A[100:150, 1] = 1
+        A[150:200, 2] = 1
+        Y = np.random.poisson(5, (n, p))
+        Y[50:100] += np.random.poisson(2, (50, p))
 
-    result, estimation = LFC(Y, W, A, family='nb')
-    assert isinstance(result, pd.DataFrame)
-    assert 'trt' in result.columns
-    assert len(result) == p * a
-    assert result['tau'].notna().all()
-    assert result['padj'].notna().all()
+        result, estimation = LFC(Y, W, A, family='nb')
+        assert isinstance(result, pd.DataFrame)
+        assert 'trt' in result.columns
+        assert len(result) == p * a
+        assert result['tau'].notna().all()
+        assert result['padj'].notna().all()

tests/test_batch_fitting.py

@@ -48,52 +48,44 @@ def _make_data(n=200, p=40, a=6, intercept=True, seed=0, family='nb'):
 
 
 # ---------------------------------------------------------------------------
-# Tests: subsample_ctrl_cells
+# Tests: subsample utilities
 # ---------------------------------------------------------------------------
 
-def test_subsample_ctrl_cells_smaller_pool():
-    ctrl_idx = np.arange(500)
-    result = subsample_ctrl_cells(ctrl_idx, n_ctrl=2000, random_state=0)
-    np.testing.assert_array_equal(result, ctrl_idx)
-
-
-def test_subsample_ctrl_cells_larger_pool():
-    ctrl_idx = np.arange(5000)
-    result = subsample_ctrl_cells(ctrl_idx, n_ctrl=2000, random_state=0)
-    assert result.shape == (2000,)
-    assert np.all(result[:-1] < result[1:]), "Result should be sorted"
-    assert np.all(np.isin(result, ctrl_idx)), "All indices should be in ctrl_idx"
-
-
-def test_subsample_ctrl_cells_reproducible():
-    ctrl_idx = np.arange(5000)
-    r1 = subsample_ctrl_cells(ctrl_idx, n_ctrl=1000, random_state=42)
-    r2 = subsample_ctrl_cells(ctrl_idx, n_ctrl=1000, random_state=42)
-    np.testing.assert_array_equal(r1, r2)
-
-
-# ---------------------------------------------------------------------------
-# Tests: subsample_pert_cells
-# ---------------------------------------------------------------------------
-
-def test_subsample_pert_cells_within_budget():
-    pert_idx = np.arange(100)
-    result = subsample_pert_cells(pert_idx, max_cells=2000)
-    np.testing.assert_array_equal(result, pert_idx)
-
-
-def test_subsample_pert_cells_over_budget():
-    pert_idx = np.arange(4000)
-    result = subsample_pert_cells(pert_idx, max_cells=2000)
-    assert len(result) == 2000
-    assert np.all(result[:-1] < result[1:])
-
-
-def test_subsample_pert_cells_no_cap():
-    """max_cells=None keeps all pert cells."""
-    pert_idx = np.arange(5000)
-    result = subsample_pert_cells(pert_idx, max_cells=None)
-    np.testing.assert_array_equal(result, pert_idx)
+class TestSubsampleUtils:
+    def test_ctrl_cells_smaller_pool(self):
+        ctrl_idx = np.arange(500)
+        result = subsample_ctrl_cells(ctrl_idx, n_ctrl=2000, random_state=0)
+        np.testing.assert_array_equal(result, ctrl_idx)
+
+    def test_ctrl_cells_larger_pool(self):
+        ctrl_idx = np.arange(5000)
+        result = subsample_ctrl_cells(ctrl_idx, n_ctrl=2000, random_state=0)
+        assert result.shape == (2000,)
+        assert np.all(result[:-1] < result[1:]), "Result should be sorted"
+        assert np.all(np.isin(result, ctrl_idx)), "All indices should be in ctrl_idx"
+
+    def test_ctrl_cells_reproducible(self):
+        ctrl_idx = np.arange(5000)
+        r1 = subsample_ctrl_cells(ctrl_idx, n_ctrl=1000, random_state=42)
+        r2 = subsample_ctrl_cells(ctrl_idx, n_ctrl=1000, random_state=42)
+        np.testing.assert_array_equal(r1, r2)
+
+    def test_pert_cells_within_budget(self):
+        pert_idx = np.arange(100)
+        result = subsample_pert_cells(pert_idx, max_cells=2000)
+        np.testing.assert_array_equal(result, pert_idx)
+
+    def test_pert_cells_over_budget(self):
+        pert_idx = np.arange(4000)
+        result = subsample_pert_cells(pert_idx, max_cells=2000)
+        assert len(result) == 2000
+        assert np.all(result[:-1] < result[1:])
+
+    def test_pert_cells_no_cap(self):
+        """max_cells=None keeps all pert cells."""
+        pert_idx = np.arange(5000)
+        result = subsample_pert_cells(pert_idx, max_cells=None)
+        np.testing.assert_array_equal(result, pert_idx)
 
 
 # ---------------------------------------------------------------------------

tests/test_fit_glm_numpy.py

@@ -96,37 +96,33 @@ def test_fit_gcate_near_singular_X():
 
 
 # ---------------------------------------------------------------------------
-# G10 – backend toggle propagates through gcate (module-qualified import)
+# Backend integration via fit_gcate
 # ---------------------------------------------------------------------------
 
-def test_g10_backend_toggle_propagates_through_fit_gcate():
-    """_USE_FAST_BACKEND=False must cause fit_gcate to use statsmodels path."""
-    np.random.seed(0)
-    n, p = 60, 20
-    Y = np.random.negative_binomial(5, 0.5, (n, p))
-    X = np.column_stack([np.ones(n), np.random.randn(n)])
-    A = np.random.binomial(1, 0.3, (n, 1)).astype(float)
-
-    with gcate_glm._backend_override("original"):
-        with patch.object(gcate_glm, "fit_glm", wraps=gcate_glm.fit_glm) as mock_fg:
-            fit_gcate(Y, X, A, r=1, family="nb", offset=False, backend="original")
-            assert mock_fg.called, (
-                "fit_glm (statsmodels) was not called despite backend='original'"
-            )
-
-
-# ---------------------------------------------------------------------------
-# G9 – backend parameter accepted without error
-# ---------------------------------------------------------------------------
-
-def test_g9_fit_gcate_backend_original_small():
-    """fit_gcate(..., backend='original') must complete and return dicts."""
-    np.random.seed(1)
-    n, p = 60, 20
-    Y = np.random.negative_binomial(5, 0.5, (n, p))
-    X = np.column_stack([np.ones(n), np.random.randn(n)])
-    A = np.random.binomial(1, 0.3, (n, 1)).astype(float)
-
-    res_1, res_2 = fit_gcate(Y, X, A, r=1, family="nb", offset=False, backend="original")
-    assert isinstance(res_1, dict)
-    assert isinstance(res_2, dict)
+class TestBackendIntegration:
+    def test_backend_toggle_propagates_through_fit_gcate(self):
+        """_USE_FAST_BACKEND=False must cause fit_gcate to use statsmodels path."""
+        np.random.seed(0)
+        n, p = 60, 20
+        Y = np.random.negative_binomial(5, 0.5, (n, p))
+        X = np.column_stack([np.ones(n), np.random.randn(n)])
+        A = np.random.binomial(1, 0.3, (n, 1)).astype(float)
+
+        with gcate_glm._backend_override("original"):
+            with patch.object(gcate_glm, "fit_glm", wraps=gcate_glm.fit_glm) as mock_fg:
+                fit_gcate(Y, X, A, r=1, family="nb", offset=False, backend="original")
+                assert mock_fg.called, (
+                    "fit_glm (statsmodels) was not called despite backend='original'"
+                )
+
+    def test_backend_original_param_accepted(self):
+        """fit_gcate(..., backend='original') must complete and return dicts."""
+        np.random.seed(1)
+        n, p = 60, 20
+        Y = np.random.negative_binomial(5, 0.5, (n, p))
+        X = np.column_stack([np.ones(n), np.random.randn(n)])
+        A = np.random.binomial(1, 0.3, (n, 1)).astype(float)
+
+        res_1, res_2 = fit_gcate(Y, X, A, r=1, family="nb", offset=False, backend="original")
+        assert isinstance(res_1, dict)
+        assert isinstance(res_2, dict)

tests/test_gcate.py

@@ -1,9 +1,4 @@
-"""
-Unit tests for GCATE confounder estimation quality.
-
-Tests T1–T6 use simulated NB/Poisson data with known latent structure so
-that confounder recovery can be evaluated objectively.
-"""
+"""Unit tests for GCATE confounder estimation quality using simulated NB/Poisson data with known latent structure."""
 
 import numpy as np
 import pytest