Increase balance package test coverage (facebookresearch#281)

Summary:

The coverage report showed 98% overall coverage with 7 files having gaps. These tests cover previously untested edge cases including:

- CLI exception handling paths for weighting failures
- Sample class design effect diagnostics and IPW model parameters
- CBPS optimization convergence warnings and constraint violation exceptions
- Plotting functions with missing values, default parameters, and various dist_types
- Distance metrics with empty numeric columns

Differential Revision: D90946146

Author: talgalili

tests/test_cbps.py

@@ -1352,6 +1352,8 @@ def test_cbps_over_method_with_extreme_data_logs_warning(self) -> None:
             transformations=None,
             cbps_method="exact",
         )
+
+    def test_cbps_over_method_logs_warnings(self) -> None:
         """Test CBPS over method logs warnings when optimization fails (lines 713, 747, 765).
 
         Verifies that when optimization algorithms fail to converge, appropriate
@@ -1477,3 +1479,286 @@ def test_cbps_alpha_function_convergence_warning(self) -> None:
                 ),
                 msg=f"Expected exception to contain convergence-related message, got: {e}",
             )
+
+
+class TestCbpsOptimizationConvergenceWithMocking(balance.testutil.BalanceTestCase):
+    """Test CBPS optimization convergence warning branches using mocking (lines 689, 713, 726, 747, 765, 778).
+
+    These tests use unittest.mock to directly control scipy.optimize.minimize return values,
+    ensuring the specific warning and exception branches in _cbps_optimization are executed.
+    """
+
+    def _create_simple_test_data(self) -> tuple:
+        """Create simple test data for CBPS testing."""
+        sample_df = pd.DataFrame({"a": [1.0, 2.0, 3.0, 4.0, 5.0]})
+        target_df = pd.DataFrame({"a": [2.0, 3.0, 4.0, 5.0, 6.0]})
+        sample_weights = pd.Series([1.0] * 5)
+        target_weights = pd.Series([1.0] * 5)
+        return sample_df, sample_weights, target_df, target_weights
+
+    def test_exact_method_constraint_violation_exception(self) -> None:
+        """Test line 726: Exception when exact method constraints can't be satisfied.
+
+        Uses mocking to simulate scipy.optimize.minimize returning success=False
+        with a specific constraint violation message.
+        """
+        from unittest.mock import MagicMock, patch
+
+        sample_df, sample_weights, target_df, target_weights = (
+            self._create_simple_test_data()
+        )
+
+        def mock_minimize(fun, x0, **kwargs):
+            result = MagicMock()
+            # Simulate constraint violation failure
+            result.__getitem__ = lambda self, key: {
+                "success": np.bool_(False),
+                "message": "Did not converge to a solution satisfying the constraints",
+                "x": x0,
+                "fun": 100.0,
+            }[key]
+            return result
+
+        def mock_minimize_scalar(fun, **kwargs):
+            return {
+                "success": np.bool_(True),
+                "message": "Success",
+                "x": np.array([1.0]),
+            }
+
+        with patch(
+            "scipy.optimize.minimize_scalar", side_effect=mock_minimize_scalar
+        ), patch("scipy.optimize.minimize", side_effect=mock_minimize):
+            with self.assertRaises(Exception) as context:
+                balance_cbps.cbps(
+                    sample_df,
+                    sample_weights,
+                    target_df,
+                    target_weights,
+                    transformations=None,
+                    cbps_method="exact",
+                )
+
+            self.assertIn(
+                "no solution satisfying the constraints",
+                str(context.exception).lower(),
+                msg="Expected exception about constraint violation",
+            )
+
+    def test_gmm_loss_glm_init_convergence_failure_warning(self) -> None:
+        """Test line 747: Warning when gmm_loss with gmm_init fails.
+
+        Uses mocking to simulate the first gmm_loss optimization failing.
+        """
+        import logging
+        from unittest.mock import MagicMock, patch
+
+        sample_df, sample_weights, target_df, target_weights = (
+            self._create_simple_test_data()
+        )
+
+        call_count = [0]
+
+        def mock_minimize(fun, x0, **kwargs):
+            call_count[0] += 1
+            result = MagicMock()
+            if call_count[0] == 1:
+                # First call is balance_optimize - succeed
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(True),
+                    "message": "Success",
+                    "x": x0,
+                    "fun": 1.0,
+                }[key]
+            elif call_count[0] == 2:
+                # Second call is gmm_optimize with glm_init - fail
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(False),
+                    "message": "Maximum iterations reached for gmm_init",
+                    "x": x0,
+                    "fun": 100.0,
+                }[key]
+            else:
+                # Third call is gmm_optimize with bal_init - succeed
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(True),
+                    "message": "Success",
+                    "x": x0,
+                    "fun": 0.5,
+                }[key]
+            return result
+
+        def mock_minimize_scalar(fun, **kwargs):
+            return {
+                "success": np.bool_(True),
+                "message": "Success",
+                "x": np.array([1.0]),
+            }
+
+        with patch(
+            "scipy.optimize.minimize_scalar", side_effect=mock_minimize_scalar
+        ), patch("scipy.optimize.minimize", side_effect=mock_minimize), self.assertLogs(
+            "balance.weighting_methods.cbps", level=logging.WARNING
+        ) as log_context:
+            try:
+                balance_cbps.cbps(
+                    sample_df,
+                    sample_weights,
+                    target_df,
+                    target_weights,
+                    transformations=None,
+                    cbps_method="over",
+                )
+            except Exception:
+                pass
+
+            # Check that a warning about gmm_init was logged
+            gmm_warnings = [
+                r
+                for r in log_context.records
+                if "gmm" in r.getMessage().lower()
+                or "convergence" in r.getMessage().lower()
+            ]
+            self.assertTrue(
+                len(gmm_warnings) > 0 or len(log_context.records) > 0,
+                msg="Expected warning about gmm_loss gmm_init convergence failure",
+            )
+
+    def test_gmm_loss_bal_init_convergence_failure_warning(self) -> None:
+        """Test line 765: Warning when gmm_loss with beta_balance fails.
+
+        Uses mocking to simulate the second gmm_loss optimization failing.
+        """
+        import logging
+        from unittest.mock import MagicMock, patch
+
+        sample_df, sample_weights, target_df, target_weights = (
+            self._create_simple_test_data()
+        )
+
+        call_count = [0]
+
+        def mock_minimize(fun, x0, **kwargs):
+            call_count[0] += 1
+            result = MagicMock()
+            if call_count[0] == 1:
+                # First call is balance_optimize - succeed
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(True),
+                    "message": "Success",
+                    "x": x0,
+                    "fun": 1.0,
+                }[key]
+            elif call_count[0] == 2:
+                # Second call is gmm_optimize with glm_init - succeed
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(True),
+                    "message": "Success",
+                    "x": x0,
+                    "fun": 0.5,
+                }[key]
+            else:
+                # Third call is gmm_optimize with bal_init - fail
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(False),
+                    "message": "Maximum iterations reached for beta_balance",
+                    "x": x0,
+                    "fun": 100.0,
+                }[key]
+            return result
+
+        def mock_minimize_scalar(fun, **kwargs):
+            return {
+                "success": np.bool_(True),
+                "message": "Success",
+                "x": np.array([1.0]),
+            }
+
+        with patch(
+            "scipy.optimize.minimize_scalar", side_effect=mock_minimize_scalar
+        ), patch("scipy.optimize.minimize", side_effect=mock_minimize), self.assertLogs(
+            "balance.weighting_methods.cbps", level=logging.WARNING
+        ) as log_context:
+            try:
+                balance_cbps.cbps(
+                    sample_df,
+                    sample_weights,
+                    target_df,
+                    target_weights,
+                    transformations=None,
+                    cbps_method="over",
+                )
+            except Exception:
+                pass
+
+            # Check that a warning about beta_balance was logged
+            bal_warnings = [
+                r
+                for r in log_context.records
+                if "beta_balance" in r.getMessage().lower()
+                or "convergence" in r.getMessage().lower()
+            ]
+            self.assertTrue(
+                len(bal_warnings) > 0 or len(log_context.records) > 0,
+                msg="Expected warning about gmm_loss beta_balance convergence failure",
+            )
+
+    def test_over_method_both_gmm_constraint_violation_exception(self) -> None:
+        """Test line 778: Exception when over method both GMM optimizations fail with constraint violation.
+
+        Uses mocking to simulate both gmm_loss optimizations failing with constraint messages.
+        """
+        from unittest.mock import MagicMock, patch
+
+        sample_df, sample_weights, target_df, target_weights = (
+            self._create_simple_test_data()
+        )
+
+        call_count = [0]
+
+        def mock_minimize(fun, x0, **kwargs):
+            call_count[0] += 1
+            result = MagicMock()
+            if call_count[0] == 1:
+                # First call is balance_optimize - succeed
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(True),
+                    "message": "Success",
+                    "x": x0,
+                    "fun": 1.0,
+                }[key]
+            else:
+                # Both GMM optimizations fail with constraint violation
+                result.__getitem__ = lambda self, key: {
+                    "success": np.bool_(False),
+                    "message": "Did not converge to a solution satisfying the constraints",
+                    "x": x0,
+                    "fun": 100.0,
+                }[key]
+            return result
+
+        def mock_minimize_scalar(fun, **kwargs):
+            return {
+                "success": np.bool_(True),
+                "message": "Success",
+                "x": np.array([1.0]),
+            }
+
+        with patch(
+            "scipy.optimize.minimize_scalar", side_effect=mock_minimize_scalar
+        ), patch("scipy.optimize.minimize", side_effect=mock_minimize):
+            with self.assertRaises(Exception) as context:
+                balance_cbps.cbps(
+                    sample_df,
+                    sample_weights,
+                    target_df,
+                    target_weights,
+                    transformations=None,
+                    cbps_method="over",
+                )
+
+            self.assertIn(
+                "no solution satisfying the constraints",
+                str(context.exception).lower(),
+                msg="Expected exception about constraint violation in over method",
+            )