Allow configuring IPW logistic regression (#138)

Summary:
Changes:
- Added an optional `logistic_regression_kwargs` argument to `ipw` so callers can override LogisticRegression settings while retaining sensible defaults for existing workflows.
- Extended the CLI to parse `--ipw_logistic_regression_kwargs`, persist the parsed options, and forward them through batch processing to IPW adjustments.
- Added a regression test confirming that solver overrides supplied through the new keyword hook reach the fitted LogisticRegression model.

Why?
- Closes #130

Pull Request resolved: #138

Reviewed By: wesleytlee

Differential Revision: D86750632

Pulled By: talgalili

fbshipit-source-id: 12405247db26cc53ba7e932f4e14ac23008bb62d

Author: neuralsorcerer

balance/cli.py

@@ -8,12 +8,13 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import inspect
+import json
 import logging
 
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
 
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import balance
 
@@ -40,10 +41,25 @@ def __init__(self, args) -> None:
             self._lambda_max,
             self._num_lambdas,
             self._weight_trimming_mean_ratio,
+            self._logistic_regression_kwargs,
             self._sample_cls,
             self._sample_package_name,
             self._sample_package_version,
-        ) = (None, None, None, None, None, None, None, None, None, None, None, None)
+        ) = (
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
 
     def check_input_columns(self, columns: Union[List[str], pd.Index]) -> None:
         needed_columns = []
@@ -133,6 +149,24 @@ def standardize_types(self) -> bool:
     def weight_trimming_mean_ratio(self) -> float:
         return self.args.weight_trimming_mean_ratio
 
+    def logistic_regression_kwargs(self) -> Optional[Dict[str, Any]]:
+        raw_kwargs = self.args.ipw_logistic_regression_kwargs
+        if raw_kwargs is None:
+            return None
+        if isinstance(raw_kwargs, dict):
+            return raw_kwargs
+        try:
+            parsed = json.loads(raw_kwargs)
+        except json.JSONDecodeError as exc:
+            raise ValueError(
+                "--ipw_logistic_regression_kwargs must be a JSON object string"
+            ) from exc
+        if not isinstance(parsed, dict):
+            raise ValueError(
+                "--ipw_logistic_regression_kwargs must decode to a JSON object"
+            )
+        return parsed
+
     def split_sample(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
         in_sample = df[self.sample_column()] == 1
         sample_df = df[in_sample]
@@ -152,6 +186,7 @@ def process_batch(
         lambda_max: Optional[float] = 10,
         num_lambdas: Optional[int] = 250,
         weight_trimming_mean_ratio: float = 20,
+        logistic_regression_kwargs: Optional[Dict[str, Any]] = None,
         sample_cls: Type[balance_sample_cls] = balance_sample_cls,
         sample_package_name: str = __package__,
     ) -> Dict[str, pd.DataFrame]:
@@ -215,6 +250,7 @@ def process_batch(
                 lambda_max=lambda_max,
                 num_lambdas=num_lambdas,
                 weight_trimming_mean_ratio=weight_trimming_mean_ratio,
+                logistic_regression_kwargs=logistic_regression_kwargs,
             )
             logger.info("Succeeded with adjusting sample to target")
             logger.info("%s adjusted object: %s" % (sample_package_name, str(adjusted)))
@@ -353,6 +389,7 @@ def update_attributes_for_main_used_by_adjust(self) -> None:
         one_hot_encoding = self.one_hot_encoding()
         max_de = self.max_de()
         weight_trimming_mean_ratio = self.weight_trimming_mean_ratio()
+        logistic_regression_kwargs = self.logistic_regression_kwargs()
         sample_cls, sample_package_name, sample_package_version = (
             balance_sample_cls,
             __package__,
@@ -370,6 +407,7 @@ def update_attributes_for_main_used_by_adjust(self) -> None:
             self._lambda_max,
             self._num_lambdas,
             self._weight_trimming_mean_ratio,
+            self._logistic_regression_kwargs,
             self._sample_cls,
             self._sample_package_name,
             self._sample_package_version,
@@ -383,6 +421,7 @@ def update_attributes_for_main_used_by_adjust(self) -> None:
             lambda_max,
             num_lambdas,
             weight_trimming_mean_ratio,
+            logistic_regression_kwargs,
             sample_cls,
             sample_package_name,
             sample_package_version,
@@ -400,6 +439,7 @@ def main(self) -> None:
             lambda_max,
             num_lambdas,
             weight_trimming_mean_ratio,
+            logistic_regression_kwargs,
             sample_cls,
             sample_package_name,
             sample_package_version,
@@ -413,6 +453,7 @@ def main(self) -> None:
             self._lambda_max,
             self._num_lambdas,
             self._weight_trimming_mean_ratio,
+            self._logistic_regression_kwargs,
             self._sample_cls,
             self._sample_package_name,
             self._sample_package_version,
@@ -434,6 +475,7 @@ def main(self) -> None:
             "lambda_max",
             "num_lambdas",
             "weight_trimming_mean_ratio",
+            "logistic_regression_kwargs",
             "sample_cls",
             "sample_package_name",
             "sample_package_version",
@@ -448,6 +490,7 @@ def main(self) -> None:
             lambda_max,
             num_lambdas,
             weight_trimming_mean_ratio,
+            logistic_regression_kwargs,
             sample_cls,
             sample_package_name,
             sample_package_version,
@@ -476,6 +519,7 @@ def main(self) -> None:
                     lambda_max,
                     num_lambdas,
                     weight_trimming_mean_ratio,
+                    logistic_regression_kwargs,
                     sample_cls,
                     sample_package_name,
                 )
@@ -501,6 +545,7 @@ def main(self) -> None:
                 lambda_max,
                 num_lambdas,
                 weight_trimming_mean_ratio,
+                logistic_regression_kwargs,
                 sample_cls,
                 sample_package_name,
             )
@@ -686,6 +731,16 @@ def add_arguments_to_parser(parser: ArgumentParser) -> ArgumentParser:
             "If not supplied it defaults to 250."
         ),
     )
+    parser.add_argument(
+        "--ipw_logistic_regression_kwargs",
+        type=str,
+        required=False,
+        default=None,
+        help=(
+            "JSON object string with additional keyword arguments passed to "
+            "sklearn.linear_model.LogisticRegression when method is ipw."
+        ),
+    )
     parser.add_argument(
         "--weight_trimming_mean_ratio",
         type=_float_or_none,
@@ -710,15 +765,13 @@ def add_arguments_to_parser(parser: ArgumentParser) -> ArgumentParser:
     # TODO: Ideally we would like transformations argument to be able to get three types of values: None (for no transformations),
     # "default" for default transformations or a dictionary of transformations.
     # However, as a first step I added the option for "default" (which is also the default) and None (for no transformations).
-    (
-        parser.add_argument(
-            "--transformations",
-            default="default",
-            required=False,
-            help=(
-                "Define the transformations for the covariates. Can be set to None for no transformations or"
-                "'default' for default transformations."
-            ),
+    parser.add_argument(
+        "--transformations",
+        default="default",
+        required=False,
+        help=(
+            "Define the transformations for the covariates. Can be set to None for no transformations or"
+            "'default' for default transformations."
         ),
     )
     # TODO: we currently support only the option of a string formula (or None), not a list of formulas.