srm check finished

Author: dmitrkozlovsk

src/services/analytics/stat_functions.py

@@ -620,7 +620,7 @@ def sample_size_ratio_metric(
 def sample_ratio_mismatch_test(
     observed_counts: list[int] | np.ndarray,
     expected_ratios: list[float] | np.ndarray | None = None,
-    alpha: float = 1e-3, 
+    alpha: float = 1e-3,
 ) -> SRMResult:
     """Sample Ratio Mismatch (SRM) via Pearson's chi-square goodness-of-fit.
 
@@ -652,17 +652,16 @@ def sample_ratio_mismatch_test(
         alloc = np.asarray(expected_ratios, dtype=float)
         if alloc.shape != (k,):
             raise ValueError("expected_ratios must have the same length as observed_counts")
-        if np.any(alloc < 0):
-            raise ValueError("All expected ratios must be non-negative")
-        s = alloc.sum()
+        if np.any(~np.isfinite(alloc)) or np.any(alloc <= 0):
+            raise ValueError("All expected ratios must be finite and strictly positive")
+        s = float(alloc.sum())
         if not np.isfinite(s) or s <= 0:
             raise ValueError("expected_ratios must sum to a positive number")
-        if not np.isclose(alloc.sum(), 1):
-            raise ValueError("expected_ratios must sum to 1")
-        alloc = alloc / s
+        if not np.isclose(s, 1.0, rtol=1e-6, atol=1e-12):
+            raise ValueError("expected_ratios must sum to 1 (within tolerance)")
 
     expected = N * alloc
-    if np.any(expected == 0):
+    if np.any(expected <= 0):
         raise ValueError("Expected counts contain zeros; ensure expected_ratios > 0 and N > 0")
 
     # Pearson's chi-square
@@ -680,4 +679,4 @@ def sample_ratio_mismatch_test(
         observed=obs.astype(int),
         allocation=alloc.astype(float),
         is_srm=is_srm,
-    )
+    )

src/ui/results/elements.py

@@ -152,9 +152,10 @@ def render(observation_cnt: dict[str, Any]) -> None:
         Args:
             observation_cnt: A dictionary with observation counts per group.
         """
-        from src.services.analytics.stat_functions import sample_ratio_mismatch_test
         import pandas as pd
 
+        from src.services.analytics.stat_functions import sample_ratio_mismatch_test
+
         with st.expander("SRM Check", expanded=False):
             placeholder = st.empty()
             if not observation_cnt or len(observation_cnt) < 2:
@@ -167,7 +168,7 @@ def render(observation_cnt: dict[str, Any]) -> None:
                 options=groups_from_page,
                 key="srm_check_control",
                 default=groups_from_page,
-                selection_mode="multi"
+                selection_mode="multi",
             )
             if len(selected_groups) < 2:
                 st.warning("Need at least 2 groups to perform SRM check")
@@ -197,23 +198,26 @@ def render(observation_cnt: dict[str, Any]) -> None:
                         value=series.expected_ratio,
                         step=0.01,
                         key=f"srm_ratio_{series.group}",
-                        label_visibility="collapsed"
+                        label_visibility="collapsed",
                     )
                     expected_ratios.append(ratio_value)
 
-
             if st.button("🔍 Check for SRM", type="primary"):
                 try:
                     result = sample_ratio_mismatch_test(
                         observed_counts=df.counts, expected_ratios=expected_ratios, alpha=1e-3
                     )
                     if result.is_srm:
-                        placeholder.error(f"Sample Ratio Mismatch detected! p-value: {result.p_value:.4f}", icon="🔥")
+                        placeholder.error(
+                            f"Sample Ratio Mismatch detected! p-value: {result.p_value:.4f}", icon="🔥"
+                        )
                     else:
-                        placeholder.success(f"No Sample Ratio Mismatch detected. p-value: {result.p_value:.4f}", icon="✅")
+                        placeholder.success(
+                            f"No Sample Ratio Mismatch detected. p-value: {result.p_value:.4f}", icon="✅"
+                        )
 
                 except Exception as e:
-                    placeholder.error(f"Error running SRM test: {e}")
+                    placeholder.warning(f"Error running SRM test: {e}")
 
 
 class ResultsDataframes:

tests/services/analytics/stat_functions/test_srm.py

@@ -0,0 +1,39 @@
+import pytest
+
+from src.services.analytics.stat_functions import sample_ratio_mismatch_test
+
+
+@pytest.mark.parametrize(
+    "observed_counts, expected_ratios, expected_is_srm, p_value_threshold",
+    [
+        ([1000, 1010], None, False, 0.001),
+        ([1000, 1200], None, True, 0.001),
+        ([100, 205], [1 / 3, 2 / 3], False, 0.001),
+        ([300, 700], [0.25, 0.75], True, 0.001),
+    ],
+)
+def test_sample_ratio_mismatch_happy_path(
+    observed_counts, expected_ratios, expected_is_srm, p_value_threshold
+):
+    result = sample_ratio_mismatch_test(observed_counts, expected_ratios, alpha=p_value_threshold)
+    assert result.is_srm is expected_is_srm
+    if expected_is_srm:
+        assert result.p_value < p_value_threshold
+    else:
+        assert result.p_value >= p_value_threshold
+
+
+@pytest.mark.parametrize(
+    "observed_counts, expected_ratios, error_message",
+    [
+        ([100], None, "observed_counts must contain at least 2 groups"),
+        ([100, -10], None, "All observed counts must be non-negative"),
+        ([100, 100], [0.5, 0.4, 0.1], "expected_ratios must have the same length as observed_counts"),
+        ([100, 100], [0.6, 0.6], "expected_ratios must sum to 1"),
+        ([0, 0], None, "Total count must be positive"),
+    ],
+)
+def test_sample_ratio_mismatch_corner_cases(observed_counts, expected_ratios, error_message):
+    """Tests corner cases and invalid inputs for sample_ratio_mismatch_test."""
+    with pytest.raises(ValueError, match=error_message):
+        sample_ratio_mismatch_test(observed_counts, expected_ratios)