Extract common pattern in comparison stat functions to reduce CCN (#267)

Summary:
- [x] Create a generic helper function `_apply_comparison_stat_to_BalanceDF` to reduce code duplication
- [x] Update `_asmd_BalanceDF` to use the new helper function
- [x] Update `_kld_BalanceDF` to use the new helper function
- [x] Update `_emd_BalanceDF` to use the new helper function
- [x] Update `_cvmd_BalanceDF` to use the new helper function
- [x] Update `_ks_BalanceDF` to use the new helper function
- [x] Run tests to ensure all functionality remains unchanged (all 83 tests in test_balancedf.py pass)
- [x] Run code review and security checks (no issues found)
- [x] Add direct test coverage for _kld_BalanceDF, _emd_BalanceDF, _cvmd_BalanceDF, and _ks_BalanceDF (5 new tests added)
- [x] Fix flake8 linting errors (removed trailing whitespace from blank lines)
- [x] Fix ufmt formatting errors (formatted test file with ufmt)

Successfully extracted the common pattern from five similar functions (`_asmd_BalanceDF`, `_kld_BalanceDF`, `_emd_BalanceDF`, `_cvmd_BalanceDF`, `_ks_BalanceDF`) into a single generic helper function `_apply_comparison_stat_to_BalanceDF`.

### Changes Made
- Added `Callable` to the imports from `typing` module
- Created `_apply_comparison_stat_to_BalanceDF` helper function that:
  1. Validates inputs are BalanceDF objects
  2. Extracts df and weights from both objects
  3. Calls the comparison function with the extracted data
- Refactored all five comparison functions to use the helper (reduced from ~30 lines to ~5 lines each)
- Maintains special handling for `_asmd_BalanceDF` which passes `std_type="target"` via kwargs
- **Added comprehensive test coverage** for all comparison methods:
  - `test_BalanceDF__kld_BalanceDF`: Direct test of _kld_BalanceDF method
  - `test_BalanceDF__emd_BalanceDF`: Direct test of _emd_BalanceDF method
  - `test_BalanceDF__cvmd_BalanceDF`: Direct test of _cvmd_BalanceDF method
  - `test_BalanceDF__ks_BalanceDF`: Direct test of _ks_BalanceDF method
  - `test_BalanceDF_comparison_functions_invalid_input`: Tests input validation for all methods
- **Fixed flake8 linting errors**: Removed trailing whitespace from blank lines in test file
- **Fixed ufmt formatting**: Formatted test file according to project standards (black + usort)

### Test Coverage
- All tests now directly exercise the helper function through the four comparison methods
- Tests verify correct Series output with expected keys (a, b, mean(metric))
- Tests verify mathematical properties (non-negativity, bounded ranges)
- Tests verify aggregate_by_main_covar parameter works
- Tests verify proper input validation with clear error messages
- Total tests: 88 (83 original + 5 new), all passing
- **Code quality compliance**: All linting (flake8) and formatting (ufmt) checks pass

### Benefits
- **Reduces Cyclomatic Complexity Number (CCN)** - the original goal of the issue
- **Eliminates code duplication** - DRY principle applied
- **Easier maintenance** - future changes only need to be made in one place
- **Type safety** - added proper type hint for the callable parameter
- **No behavioral changes** - all 83 existing tests pass without modification
- **Comprehensive test coverage** - direct tests for all comparison methods and edge cases
- **Code quality** - passes all linting and formatting checks

<details>

<summary>Original prompt</summary>

>
> ----
>
> *This section details on the original issue you should resolve*
>
> <issue_title>[FEATURE] generalize functions in ‎balance/balancedf_class.py‎</issue_title>
> <issue_description>a bunch of functions in
> ‎balance/balancedf_class.py‎
>
> Follow the exact same pattern (other than one word change:
> staticmethod
>     def _emd_BalanceDF(
>         sample_BalanceDF: "BalanceDF",
>         target_BalanceDF: "BalanceDF",
>         aggregate_by_main_covar: bool = False,
>     ) -> pd.Series:
>         """Run EMD on two BalanceDF objects.
>
>         Prepares the BalanceDF objects by passing them through :func:`_get_df_and_weights`, and
>         then passes the df and weights into :func:`weighted_comparisons_stats.emd`.
>
>         Args:
>             sample_BalanceDF (BalanceDF): Object.
>             target_BalanceDF (BalanceDF): Object.
>             aggregate_by_main_covar (bool, optional): See :func:`weighted_comparisons_stats.emd`. Defaults to False.
>
>         Returns:
>             pd.Series: See :func:`weighted_comparisons_stats.emd`.
>         """
>         BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
>         BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
>
>         sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
>         target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
>
>         return weighted_comparisons_stats.emd(
>             sample_df_values,
>             target_df_values,
>             sample_weights,
>             target_weights,
>             aggregate_by_main_covar=aggregate_by_main_covar,
>         )
>
>
>
> Extract this pattern to a helper function to reduce CCN.
>
> </issue_description>
>
> ## Comments on the Issue (you are copilot in this section)
>
> <comments>
> </comments>
>

</details>

- Fixes #266

 ---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our [2 minute survey](https://gh.io/copilot-coding-agent-survey).

Pull Request resolved: #267

Differential Revision: D90870323

Pulled By: talgalili

fbshipit-source-id: 064e3878506e1f65726ce1302fc692e9ec794676

Author: Copilot

balance/balancedf_class.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, Literal, Tuple
+from typing import Any, Callable, Dict, Literal, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -1110,6 +1110,50 @@ def _get_df_and_weights(
         weights = self._weights.values if (self._weights is not None) else None
         return df_model_matrix, weights
 
+    @staticmethod
+    def _apply_comparison_stat_to_BalanceDF(
+        comparison_func: Callable[..., pd.Series],
+        sample_BalanceDF: "BalanceDF",
+        target_BalanceDF: "BalanceDF",
+        aggregate_by_main_covar: bool = False,
+        **kwargs: Any,
+    ) -> pd.Series:
+        """Generic helper to apply a weighted comparison statistic function to two BalanceDF objects.
+
+        This helper function reduces code duplication across multiple comparison methods
+        (asmd, kld, emd, cvmd, ks) by extracting the common pattern of:
+        1. Validating inputs are BalanceDF objects
+        2. Extracting df and weights from both objects
+        3. Calling the comparison function with the extracted data
+
+        Args:
+            comparison_func (Callable[..., pd.Series]): The comparison function from
+                weighted_comparisons_stats to apply (e.g., asmd, kld, emd, cvmd, ks).
+            sample_BalanceDF (BalanceDF): Sample object.
+            target_BalanceDF (BalanceDF): Target object.
+            aggregate_by_main_covar (bool, optional): Whether to aggregate by main covariate.
+                Defaults to False. Passed to the comparison function.
+            **kwargs: Additional keyword arguments to pass to the comparison function
+                (e.g., std_type for asmd).
+
+        Returns:
+            pd.Series: The result from the comparison function.
+        """
+        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
+        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
+
+        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
+        target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
+
+        return comparison_func(
+            sample_df_values,
+            target_df_values,
+            sample_weights,
+            target_weights,
+            aggregate_by_main_covar=aggregate_by_main_covar,
+            **kwargs,
+        )
+
     @staticmethod
     def _asmd_BalanceDF(
         sample_BalanceDF: "BalanceDF",
@@ -1156,19 +1200,12 @@ def _asmd_BalanceDF(
                     # mean(asmd)    1.756543
                     # dtype: float64
         """
-        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
-        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "target_BalanceDF")
-
-        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
-        target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
-
-        return weighted_comparisons_stats.asmd(
-            sample_df_values,
-            target_df_values,
-            sample_weights,
-            target_weights,
+        return BalanceDF._apply_comparison_stat_to_BalanceDF(
+            weighted_comparisons_stats.asmd,
+            sample_BalanceDF,
+            target_BalanceDF,
+            aggregate_by_main_covar,
             std_type="target",
-            aggregate_by_main_covar=aggregate_by_main_covar,
         )
 
     @staticmethod
@@ -1190,18 +1227,11 @@ def _kld_BalanceDF(
         Returns:
             pd.Series: See :func:`weighted_comparisons_stats.kld`.
         """
-        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
-        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
-
-        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
-        target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
-
-        return weighted_comparisons_stats.kld(
-            sample_df_values,
-            target_df_values,
-            sample_weights,
-            target_weights,
-            aggregate_by_main_covar=aggregate_by_main_covar,
+        return BalanceDF._apply_comparison_stat_to_BalanceDF(
+            weighted_comparisons_stats.kld,
+            sample_BalanceDF,
+            target_BalanceDF,
+            aggregate_by_main_covar,
         )
 
     @staticmethod
@@ -1223,18 +1253,11 @@ def _emd_BalanceDF(
         Returns:
             pd.Series: See :func:`weighted_comparisons_stats.emd`.
         """
-        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
-        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
-
-        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
-        target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
-
-        return weighted_comparisons_stats.emd(
-            sample_df_values,
-            target_df_values,
-            sample_weights,
-            target_weights,
-            aggregate_by_main_covar=aggregate_by_main_covar,
+        return BalanceDF._apply_comparison_stat_to_BalanceDF(
+            weighted_comparisons_stats.emd,
+            sample_BalanceDF,
+            target_BalanceDF,
+            aggregate_by_main_covar,
         )
 
     @staticmethod
@@ -1256,18 +1279,11 @@ def _cvmd_BalanceDF(
         Returns:
             pd.Series: See :func:`weighted_comparisons_stats.cvmd`.
         """
-        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
-        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
-
-        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
-        target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
-
-        return weighted_comparisons_stats.cvmd(
-            sample_df_values,
-            target_df_values,
-            sample_weights,
-            target_weights,
-            aggregate_by_main_covar=aggregate_by_main_covar,
+        return BalanceDF._apply_comparison_stat_to_BalanceDF(
+            weighted_comparisons_stats.cvmd,
+            sample_BalanceDF,
+            target_BalanceDF,
+            aggregate_by_main_covar,
         )
 
     @staticmethod
@@ -1289,18 +1305,11 @@ def _ks_BalanceDF(
         Returns:
             pd.Series: See :func:`weighted_comparisons_stats.ks`.
         """
-        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
-        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
-
-        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights()
-        target_df_values, target_weights = target_BalanceDF._get_df_and_weights()
-
-        return weighted_comparisons_stats.ks(
-            sample_df_values,
-            target_df_values,
-            sample_weights,
-            target_weights,
-            aggregate_by_main_covar=aggregate_by_main_covar,
+        return BalanceDF._apply_comparison_stat_to_BalanceDF(
+            weighted_comparisons_stats.ks,
+            sample_BalanceDF,
+            target_BalanceDF,
+            aggregate_by_main_covar,
         )
 
     def asmd(

tests/test_balancedf.py

@@ -1301,6 +1301,136 @@ def test_BalanceDF_asmd_aggregate_by_main_covar(self) -> None:
         self.assertEqual(outcome_default, expected_default)
         self.assertEqual(outcome_main_covar, expected_main_covar)
 
+    def test_BalanceDF__kld_BalanceDF(self) -> None:
+        """Test _kld_BalanceDF static method directly."""
+        sample = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (1, 2), "b": (-1, 12), "weight": (1, 2)})
+        ).covars()
+
+        target = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (3, 4), "b": (0, 42), "weight": (1, 2)})
+        ).covars()
+
+        result = BalanceDF._kld_BalanceDF(sample, target)
+
+        # Verify result is a Series with expected keys
+        self.assertIsInstance(result, pd.Series)
+        self.assertIn("a", result.index)
+        self.assertIn("b", result.index)
+        self.assertIn("mean(kld)", result.index)
+
+        # Verify all values are non-negative (KLD property)
+        self.assertTrue((result >= 0).all())
+
+        # Test with aggregate_by_main_covar
+        result_agg = BalanceDF._kld_BalanceDF(
+            sample, target, aggregate_by_main_covar=True
+        )
+        self.assertIsInstance(result_agg, pd.Series)
+
+    def test_BalanceDF__emd_BalanceDF(self) -> None:
+        """Test _emd_BalanceDF static method directly."""
+        sample = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (1, 2), "b": (-1, 12), "weight": (1, 2)})
+        ).covars()
+
+        target = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (3, 4), "b": (0, 42), "weight": (1, 2)})
+        ).covars()
+
+        result = BalanceDF._emd_BalanceDF(sample, target)
+
+        # Verify result is a Series with expected keys
+        self.assertIsInstance(result, pd.Series)
+        self.assertIn("a", result.index)
+        self.assertIn("b", result.index)
+        self.assertIn("mean(emd)", result.index)
+
+        # Verify all values are non-negative (EMD property)
+        self.assertTrue((result >= 0).all())
+
+        # Test with aggregate_by_main_covar
+        result_agg = BalanceDF._emd_BalanceDF(
+            sample, target, aggregate_by_main_covar=True
+        )
+        self.assertIsInstance(result_agg, pd.Series)
+
+    def test_BalanceDF__cvmd_BalanceDF(self) -> None:
+        """Test _cvmd_BalanceDF static method directly."""
+        sample = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (1, 2), "b": (-1, 12), "weight": (1, 2)})
+        ).covars()
+
+        target = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (3, 4), "b": (0, 42), "weight": (1, 2)})
+        ).covars()
+
+        result = BalanceDF._cvmd_BalanceDF(sample, target)
+
+        # Verify result is a Series with expected keys
+        self.assertIsInstance(result, pd.Series)
+        self.assertIn("a", result.index)
+        self.assertIn("b", result.index)
+        self.assertIn("mean(cvmd)", result.index)
+
+        # Verify all values are non-negative (CVMD property)
+        self.assertTrue((result >= 0).all())
+
+        # Test with aggregate_by_main_covar
+        result_agg = BalanceDF._cvmd_BalanceDF(
+            sample, target, aggregate_by_main_covar=True
+        )
+        self.assertIsInstance(result_agg, pd.Series)
+
+    def test_BalanceDF__ks_BalanceDF(self) -> None:
+        """Test _ks_BalanceDF static method directly."""
+        sample = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (1, 2), "b": (-1, 12), "weight": (1, 2)})
+        ).covars()
+
+        target = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (3, 4), "b": (0, 42), "weight": (1, 2)})
+        ).covars()
+
+        result = BalanceDF._ks_BalanceDF(sample, target)
+
+        # Verify result is a Series with expected keys
+        self.assertIsInstance(result, pd.Series)
+        self.assertIn("a", result.index)
+        self.assertIn("b", result.index)
+        self.assertIn("mean(ks)", result.index)
+
+        # Verify all values are in [0, 1] (KS property)
+        self.assertTrue((result >= 0).all())
+        self.assertTrue((result <= 1).all())
+
+        # Test with aggregate_by_main_covar
+        result_agg = BalanceDF._ks_BalanceDF(
+            sample, target, aggregate_by_main_covar=True
+        )
+        self.assertIsInstance(result_agg, pd.Series)
+
+    def test_BalanceDF_comparison_functions_invalid_input(self) -> None:
+        """Test that all comparison functions properly validate inputs."""
+        sample = Sample.from_frame(
+            pd.DataFrame({"id": (1, 2), "a": (1, 2), "weight": (1, 2)})
+        ).covars()
+
+        # Test with non-BalanceDF inputs
+        invalid_input = "not a BalanceDF"
+
+        with self.assertRaisesRegex(ValueError, "must be balancedf_class.BalanceDF"):
+            BalanceDF._kld_BalanceDF(invalid_input, sample)  # type: ignore
+
+        with self.assertRaisesRegex(ValueError, "must be balancedf_class.BalanceDF"):
+            BalanceDF._emd_BalanceDF(sample, invalid_input)  # type: ignore
+
+        with self.assertRaisesRegex(ValueError, "must be balancedf_class.BalanceDF"):
+            BalanceDF._cvmd_BalanceDF(invalid_input, sample)  # type: ignore
+
+        with self.assertRaisesRegex(ValueError, "must be balancedf_class.BalanceDF"):
+            BalanceDF._ks_BalanceDF(sample, invalid_input)  # type: ignore
+
 
 class TestBalanceDF_to_download(BalanceTestCase):
     def test_BalanceDF_to_download(self) -> None: