Add unit tests to wsmf.selectors; Improve code quality

experiments_engine/data.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import numpy as np
 import pandas as pd
 from openml import OpenMLTask
@@ -14,7 +16,7 @@ def move_target_to_last_column(
     return df
 
 
-def is_eligible_task(task: OpenMLTask) -> bool:
+def is_eligible_task(task: OpenMLTask) -> Any:
     if task is None:
         return False
     if (

experiments_engine/data_utils.py

@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-from typing import Tuple
+from typing import Any, Tuple
 
 import numpy as np
 import pandas as pd
@@ -18,15 +18,15 @@ def get_dataset_from_path(path: Path) -> Tuple[Tensor, Tensor]:
     dataset = pd.read_parquet(path / "test.parquet")
     X, y = dataset.iloc[:, :-1], dataset.iloc[:, -1]
     pipeline = DataUtils.get_preprocessing_pipeline()
-    X = pipeline.fit_transform(X)  # type: ignore
+    X = pipeline.fit_transform(X)
     X = Tensor(X.values).cuda()
     y = Tensor(y.values).reshape(-1, 1).cuda()
     return X, y
 
 
 def load_datasets_with_landmarkers(
     reduce_landmarkers_dimensionality: bool = False,
-):
+) -> Tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]]:
     train_datasets_names = list(
         sorted(
             [
@@ -91,7 +91,7 @@ def load_datasets_with_landmarkers(
 
 def __project_landmarkers_to_smaller_space(
     train_landmarkers: dict[str, Tensor], val_landmarkers: dict[str, Tensor]
-):
+) -> Tuple[dict[str, Tensor], dict[str, Tensor]]:
     projection_train_data = np.stack(
         list([item.cpu().numpy() for item in train_landmarkers.values()])
     )
@@ -100,15 +100,15 @@ def __project_landmarkers_to_smaller_space(
     scaling = StandardScaler().fit(scaling_train_data)
     out_train_landmarkers = {
         name: Tensor(
-            scaling.transform(  # type: ignore
+            scaling.transform(
                 projection.transform(value.cpu().numpy().reshape(1, -1))
             )[0]
         )
         for name, value in train_landmarkers.items()
     }
     out_val_landmarkers = {
         name: Tensor(
-            scaling.transform(  # type: ignore
+            scaling.transform(
                 projection.transform(value.cpu().numpy().reshape(1, -1))
             )[0]
         )

experiments_engine/hpo.py

@@ -4,6 +4,7 @@
 import optuna
 import pandas as pd
 from dataset2vec.utils import DataUtils
+from numpy.typing import NDArray
 from sklearn.base import BaseEstimator
 from sklearn.metrics import roc_auc_score
 from xgboost import XGBClassifier
@@ -15,24 +16,24 @@ def __init__(
         self,
         df_train: pd.DataFrame,
         df_test: pd.DataFrame,
-        metric: Callable = roc_auc_score,
+        metric: Callable[[NDArray[Any], NDArray[Any]], float] = roc_auc_score,
     ):
         self.X_train, self.y_train = (
             df_train.iloc[:, :-1],
             df_train.iloc[:, -1],
         )
         self.X_test, self.y_test = df_test.iloc[:, :-1], df_test.iloc[:, -1]
         self.pipeline = DataUtils.get_preprocessing_pipeline()
-        self.X_train = self.pipeline.fit_transform(self.X_train)  # type: ignore # noqa
-        self.X_test = self.pipeline.transform(self.X_test)  # type: ignore
+        self.X_train = self.pipeline.fit_transform(self.X_train)
+        self.X_test = self.pipeline.transform(self.X_test)
         self.metric = metric
 
     def __call__(self, trial: optuna.Trial) -> float:
         X_train, y_train, X_test, y_test = self.get_data()
         model = self.get_model(trial)
-        model.fit(X_train, y_train)  # type: ignore
-        probas = model.predict_proba(X_test)[:, 1]  # type: ignore
-        return self.metric(y_test, probas)  # type: ignore
+        model.fit(X_train, y_train)
+        probas = model.predict_proba(X_test)[:, 1]
+        return self.metric(y_test, probas)
 
     @abstractmethod
     def get_model(self, trial: optuna.Trial) -> BaseEstimator:
@@ -83,7 +84,7 @@ def perform_study(
     return study
 
 
-def get_best_study_params(study: optuna.Study) -> dict:
+def get_best_study_params(study: optuna.Study) -> dict[str, Any]:
     return study.best_params
 
 

experiments_engine/portfolio_selection.py

@@ -1,20 +1,22 @@
+from typing import Any
+
 import numpy as np
 from numpy.typing import NDArray
 from scipy.stats import rankdata
 
 
 def extract_best_configuration_idx_from_cluster_eval_results(
-    datasets_inside_clusters_performances: NDArray,
+    datasets_inside_clusters_performances: NDArray[Any],
 ) -> int:
     return get_ranks_of_hp_configurations(
         datasets_inside_clusters_performances
     )[0]
 
 
-def get_ranks_of_hp_configurations(hp_performances: NDArray) -> list[int]:
+def get_ranks_of_hp_configurations(hp_performances: NDArray[Any]) -> list[int]:
     ranks_per_dataset = np.array(
         [rankdata(-row, method="dense") for row in hp_performances]
     )
     average_ranks_per_configuration = ranks_per_dataset.mean(axis=0)
     final_ranks = np.argsort(average_ranks_per_configuration)
-    return final_ranks.tolist()
+    return list(final_ranks)

experiments_engine/utils.py

@@ -1,12 +1,13 @@
 import json
 from pathlib import Path
+from typing import Any
 
 import torch
 
-device = "cuda" if torch.cuda.is_available else "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-def read_json(path: Path) -> dict:
+def read_json(path: Path) -> Any:
     with open(path, "r") as f:
         return json.load(f)
 

experiments_engine/warmstart_utils.py

@@ -37,7 +37,7 @@ def perform_ground_truth_warm_start_experiment(
     )
     if n_initial_trials > 0:
         initial_trials = warm_starter.propose_configurations(
-            objective_landmarkers,  # type: ignore
+            objective_landmarkers,
             n_initial_trials,
         )
         for trial in initial_trials:

requirements_dev.txt

@@ -0,0 +1,15 @@
+dataset2vec==1.0.0
+openml==0.14.2
+loguru==0.7.2
+optuna==3.6.1
+xgboost==2.0.3
+scikit-learn==1.4.2
+pytest==8.2.1
+pymfe==0.4.3
+seaborn==0.13.2
+tensorboard==2.17.0
+numpy==1.26.4
+isort==5.13.2
+black==24.8.0
+flake8==7.1.0
+mypy==1.11.1

scripts/check_code.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+set -e
+
+export PATHS_TO_CHECK="wsmf experiments_engine test"
+
+echo "Running isort"
+isort --profile=black --line-length=79 $PATHS_TO_CHECK
+
+echo "Running black"
+black --line-length=79 $PATHS_TO_CHECK
+
+echo "Running flake8"
+flake8 --ignore=W605,W503 --exclude experiments_engine/cd_plot.py $PATHS_TO_CHECK
+
+echo "Running mypy"
+mypy \
+    --install-types \
+    --non-interactive \
+    --ignore-missing-imports \
+    --strict \
+    --namespace-packages \
+    --exclude experiments_engine/cd_plot.py \
+    $PATHS_TO_CHECK

test/experiments_engine/test_data.py

@@ -1,4 +1,4 @@
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import numpy as np
 import pandas as pd
@@ -10,7 +10,7 @@
 )
 
 
-def test_move_target_to_last_column_when_target_last():
+def test_move_target_to_last_column_when_target_last() -> None:
     # Given
     df = pd.DataFrame({"col1": [1, 2], "col2": [1, 2]})
 
@@ -21,7 +21,7 @@ def test_move_target_to_last_column_when_target_last():
     assert (actual_df == df).all(axis=None)
 
 
-def test_move_target_to_last_column_when_target_not_last():
+def test_move_target_to_last_column_when_target_not_last() -> None:
     # Given
     df = pd.DataFrame({"col1": [1, 2], "col2": [1, 2]})
 
@@ -32,7 +32,7 @@ def test_move_target_to_last_column_when_target_not_last():
     assert (actual_df == df[["col2", "col1"]]).all(axis=None)
 
 
-def test_remove_unwanted_columns():
+def test_remove_unwanted_columns() -> None:
     # Given
     df = pd.DataFrame({"id_1": [1], "2id": [2], "3_id": [3], "col": [4]})
 
@@ -43,7 +43,7 @@ def test_remove_unwanted_columns():
     assert (actual_df == df[["2id", "col"]]).all(axis=None)
 
 
-def test_clean_and_binarize_classification_multiple_classes():
+def test_clean_and_binarize_classification_multiple_classes() -> None:
     # Given
     df = pd.DataFrame({"col": [1, 2, 3, 4], "target": [1, 2, 3, 1]})
 
@@ -56,7 +56,7 @@ def test_clean_and_binarize_classification_multiple_classes():
     assert actual_df["target"].max() == 1
 
 
-def test_clean_and_binarize_classification_text_classes():
+def test_clean_and_binarize_classification_text_classes() -> None:
     # Given
     df = pd.DataFrame({"col": [1, 2, 3, 4], "target": ["A", "B", "C", "A"]})
 
@@ -71,8 +71,8 @@ def test_clean_and_binarize_classification_text_classes():
 
 @patch("numpy.random.uniform")
 def test_clean_and_binarize_classification_when_subset_only_zeros(
-    uniform_mock,
-):
+    uniform_mock: Mock,
+) -> None:
     uniform_mock.return_value = np.array([0.1, 0.2, 0.3])
     # Given
     df = pd.DataFrame({"col": [1, 2, 3, 4], "target": ["A", "B", "C", "A"]})
@@ -88,8 +88,8 @@ def test_clean_and_binarize_classification_when_subset_only_zeros(
 
 @patch("numpy.random.uniform")
 def test_clean_and_binarize_classification_when_subset_only_ones(
-    uniform_mock,
-):
+    uniform_mock: Mock,
+) -> None:
     uniform_mock.return_value = np.array([0.6, 0.7, 0.8])
     # Given
     df = pd.DataFrame({"col": [1, 2, 3, 4], "target": ["A", "B", "C", "A"]})

test/experiments_engine/test_portfolio_choice.py

@@ -5,7 +5,7 @@
 )
 
 
-def test_extract_best_configuration_idx_from_cluster_eval_results():
+def test_extract_best_configuration_idx_from_cluster_eval_results() -> None:
     # Given
     datasets_inside_clusters_performances = np.array(
         [

test/wsmf/metamodels/data/test_dataset.py

@@ -4,7 +4,7 @@
 from wsmf.metamodels.data import EncoderHpoDataset
 
 
-def test_d2v_hpo_dataset_has_proper_length():
+def test_d2v_hpo_dataset_has_proper_length() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])
@@ -26,7 +26,7 @@ def test_d2v_hpo_dataset_has_proper_length():
     assert len(d2v_hpo_dataset) == 2
 
 
-def test_d2v_hpo_dataset_has_proper_dataset_names():
+def test_d2v_hpo_dataset_has_proper_dataset_names() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])
@@ -48,7 +48,7 @@ def test_d2v_hpo_dataset_has_proper_dataset_names():
     assert d2v_hpo_dataset.dataset_names == ["dataset1", "dataset2"]
 
 
-def test_d2v_hpo_dataset_returns_proper_data_on_index():
+def test_d2v_hpo_dataset_returns_proper_data_on_index() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])
@@ -75,7 +75,7 @@ def test_d2v_hpo_dataset_returns_proper_data_on_index():
     assert (actual_landmarkers == Tensor([-1, -2, -3])).all()
 
 
-def test_d2v_hpo_dataset_fail_when_inconsistent_data_sizes():
+def test_d2v_hpo_dataset_fail_when_inconsistent_data_sizes() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])

test/wsmf/metamodels/data/test_landmarker_reconstruction.py

@@ -11,7 +11,7 @@
 @patch("numpy.random.choice")
 def test_landmarker_reconstruction_loader_returns_proper_sample(
     choice_mock: Mock,
-):
+) -> None:
     # Given
     choice_mock.return_value = [0, 1]
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
@@ -31,10 +31,10 @@ def test_landmarker_reconstruction_loader_returns_proper_sample(
     # Then
     assert (sample[0] == dataset1_X).all()
     assert (sample[1] == dataset1_y).all()
-    assert (sample[2] == landmarkers["dataset1"]).all()
+    assert (sample[2] == landmarkers["dataset1"]).all()  # type: ignore
 
 
-def test_landmarker_reconstruction_loader_returns_proper_batch_size():
+def test_landmarker_reconstruction_loader_returns_proper_batch_size() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])
@@ -58,7 +58,7 @@ def test_landmarker_reconstruction_loader_returns_proper_batch_size():
     assert len(batch) == 2
 
 
-def test_landmarker_reconstruction_loader_returns_all_datasets():
+def test_landmarker_reconstruction_loader_returns_all_datasets() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])
@@ -81,7 +81,7 @@ def test_landmarker_reconstruction_loader_returns_all_datasets():
     # Then
     assert (returned_datasets[0][0][0] == dataset1_X).all()
     assert (returned_datasets[0][0][1] == dataset1_y).all()
-    assert (returned_datasets[0][0][2] == landmarkers["dataset1"]).all()
+    assert (returned_datasets[0][0][2] == landmarkers["dataset1"]).all()  # type: ignore # noqa E501
     assert (returned_datasets[1][0][0] == dataset2_X).all()
     assert (returned_datasets[1][0][1] == dataset2_y).all()
-    assert (returned_datasets[1][0][2] == landmarkers["dataset2"]).all()
+    assert (returned_datasets[1][0][2] == landmarkers["dataset2"]).all()  # type: ignore # noqa E501

test/wsmf/metamodels/data/test_metric_loader.py

@@ -7,7 +7,9 @@
 
 
 @patch("numpy.random.choice")
-def test_encoder_metric_loader_calculates_sample_properly(choice_mock: Mock):
+def test_encoder_metric_loader_calculates_sample_properly(
+    choice_mock: Mock,
+) -> None:
     # Given
     choice_mock.return_value = [0, 1]
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
@@ -40,7 +42,7 @@ def test_encoder_metric_loader_calculates_sample_properly(choice_mock: Mock):
     assert np.isclose(sample[4], 56 / 3)
 
 
-def test_encoder_metric_loader_returns_proper_number_of_batches():
+def test_encoder_metric_loader_returns_proper_number_of_batches() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])
@@ -64,7 +66,7 @@ def test_encoder_metric_loader_returns_proper_number_of_batches():
     assert len(batches) == 16
 
 
-def test_encoder_metric_loader_returns_batch_with_proper_size():
+def test_encoder_metric_loader_returns_batch_with_proper_size() -> None:
     # Given
     dataset1_X = Tensor([[1, 2, 3], [4, 5, 6]])
     dataset1_y = Tensor([[0], [1]])