[ENH] Minor script updates

scripts/run_experiments.py

@@ -65,7 +65,7 @@ def _get_parser() -> argparse.ArgumentParser:
         dest="tasks",
         action="append",
         choices=list(TASKS),
-        help="Task to run. Can be repeated.",
+        help=f"Task to run. Can be repeated. Default: {DEFAULT_TASKS}",
     )
     parser.add_argument(
         "--n-rounds",
@@ -153,7 +153,7 @@ def _get_node_data(task: str) -> dict[str, dict]:
     return fds.data() if fds is not None else {}
 
 
-def _get_feature_stats(
+def _get_column_stats(
     analytics: FederatedAnalytics, node_data: dict[str, dict], target_column: str
 ) -> tuple[list[str], list[float], list[float], dict[str, list[str]], float]:
     """Get the feature/excluded columns, some stats, and the target variable's mean."""
@@ -241,7 +241,7 @@ def _get_agg_optimizer(aggregator: str) -> Optimizer | None:
 
 
 def _move_tensorboard_logs(exp: Experiment, dpath_tensorboard: Path) -> None:
-    """Move the per-node logs out of the Fed-BioMed experiment directory."""
+    """Move the Tensorboard logs out of the Fed-BioMed experiment directory."""
     # flush the event files before moving them
     exp.monitor().close_writer()
     dpath_tensorboard.mkdir(parents=True, exist_ok=True)
@@ -300,7 +300,7 @@ def _evaluate_held_out(
         secagg_arguments=secagg_arguments,
         do_training=False,
     ).execute()
-    # the metrics are sent separately, so a failed request would go unnoticed
+    # raise explicit errors, otherwise failed requests are silent
     for node in node_data_held_out:
         if node not in replies:
             raise RuntimeError(f"Held-out node {node} did not reply")
@@ -383,7 +383,7 @@ def _run_task(
             )
 
     # the training plan needs the feature stats to be instantiated
-    columns, mean, std, excluded_columns, target_mean = _get_feature_stats(
+    columns, mean, std, excluded_columns, target_mean = _get_column_stats(
         exp.analytics, exp.training_data().data(), config.target_column
     )
     for reason, cols in excluded_columns.items():
@@ -458,7 +458,8 @@ def run_experiments(
     output_dir = output_dir.resolve()
     settings = locals().copy()
 
-    # held-out nodes for each task, with None for no held-out node
+    # keys: task
+    # values: held-out node (None if not doing CV)
     runs: dict[str, list[str | None]] = {}
     node_data: dict[str, dict[str, dict]] = {}
     for task in tasks:

tests/scripts/test_run_experiments.py

@@ -241,9 +241,9 @@ def test_setup_held_out_secagg(
         assert secagg_arguments == secagg.train_arguments.return_value
 
 
-class TestGetFeatureStats:
+class TestGetColumnStats:
     def test_target_dropped(self, module: ModuleType, analytics: MagicMock):
-        columns, mean, std, excluded_columns, target_mean = module._get_feature_stats(
+        columns, mean, std, excluded_columns, target_mean = module._get_column_stats(
             analytics, NODE_DATA, "nb:Age"
         )
         assert columns == ["Left-Hippocampus", "Left-Thalamus", "nb:Sex"]
@@ -278,7 +278,7 @@ def test_excluded_columns(self, module: ModuleType, analytics: MagicMock):
         }
         # only in the dataset of one node
         node_data["NODE_1"]["dtypes"]["Left-Amygdala"] = "Float64"
-        columns, mean, std, excluded_columns, _ = module._get_feature_stats(
+        columns, mean, std, excluded_columns, _ = module._get_column_stats(
             analytics, node_data, "nb:Age"
         )
         assert columns == ["Left-Hippocampus", "Left-Thalamus", "nb:Sex"]
@@ -297,11 +297,11 @@ def test_excluded_columns(self, module: ModuleType, analytics: MagicMock):
         }
 
     def test_json_serializable(self, module: ModuleType, analytics: MagicMock):
-        json.dumps(module._get_feature_stats(analytics, NODE_DATA, "nb:Sex"))
+        json.dumps(module._get_column_stats(analytics, NODE_DATA, "nb:Sex"))
 
     def test_no_nodes(self, module: ModuleType, analytics: MagicMock):
         with pytest.raises(ValueError, match="No node has a dataset"):
-            module._get_feature_stats(analytics, {}, "nb:Age")
+            module._get_column_stats(analytics, {}, "nb:Age")
         analytics.fetch_stats.assert_not_called()
 
     def test_target_not_in_all_nodes(self, module: ModuleType, analytics: MagicMock):
@@ -310,7 +310,7 @@ def test_target_not_in_all_nodes(self, module: ModuleType, analytics: MagicMock)
             "NODE_2": {"dtypes": {"nb:Sex": "Int64"}},
         }
         with pytest.raises(ValueError, match="not in the datasets of all nodes"):
-            module._get_feature_stats(analytics, node_data, "nb:Age")
+            module._get_column_stats(analytics, node_data, "nb:Age")
         analytics.fetch_stats.assert_not_called()
 
 
@@ -332,7 +332,7 @@ def test_training_data(
         expected_x: list[list[float]],
     ):
         config = module.TASKS[task]
-        columns, mean, std, _, target_mean = module._get_feature_stats(
+        columns, mean, std, _, target_mean = module._get_column_stats(
             analytics, NODE_DATA, target_column
         )
         model_args = module._build_model_args(

training_plans/tabular_classifier.py

@@ -86,7 +86,7 @@ def testing_step(self, data, target):
             "BALANCED_ACCURACY": balanced_accuracy_score(target, prediction),
             "AUROC": roc_auc_score(target, self.model().decision_function(data)),
         }
-        # NaN cannot be saved in the results JSON
+        # omit metric(s) with NaN value since they cannot be saved in the results JSON
         return {
             name: metrics[name]
             for name in self.model_args().get("metrics", self.DEFAULT_METRICS)