Merge branch 'dev/add_faiss_mode'

# Conflicts:
#	hypex/matching.py
#	tests/test_tutorials.py

Author: tikhomirovd

hypex/extensions/faiss.py

@@ -11,8 +11,9 @@
 
 
 class FaissExtension(MLExtension):
-    def __init__(self, n_neighbors: int = 1):
+    def __init__(self, n_neighbors: int = 1, faiss_mode: Literal["base", "fast", "auto"] = "auto"):
         self.n_neighbors = n_neighbors
+        self.faiss_mode = faiss_mode
         super().__init__()
 
     @staticmethod
@@ -52,7 +53,7 @@ def _calc_pandas(
         X = data.data.values
         if mode in ["auto", "fit"]:
             self.index = faiss.IndexFlatL2(X.shape[1])
-            if len(X) > 1_000_000:
+            if (len(X) > 1_000_000 and self.faiss_mode == "auto") or self.faiss_mode == "fast":
                 self.index = faiss.IndexIVFFlat(self.index, 1, 1000)
                 self.index.train(X)
             self.index.add(X)

hypex/matching.py

@@ -59,14 +59,15 @@ class Matching(ExperimentShell):
 
     @staticmethod
     def _make_experiment(
-        group_match: bool = False,
-        distance: Literal["mahalanobis", "l2"] = "mahalanobis",
-        metric: Literal["atc", "att", "ate"] = "ate",
-        bias_estimation: bool = True,
-        quality_tests: (
-            Literal["smd", "psi", "ks-test", "repeats", "t-test", "auto"]
-            | list[Literal["smd", "psi", "ks-test", "repeats", "t-test", "auto"]]
-        ) = "auto",
+            group_match: bool = False,
+            distance: Literal["mahalanobis", "l2"] = "mahalanobis",
+            metric: Literal["atc", "att", "ate"] = "ate",
+            bias_estimation: bool = True,
+            quality_tests: (
+                    Literal["smd", "psi", "ks-test", "repeats", "t-test", "auto"]
+                    | list[Literal["smd", "psi", "ks-test", "repeats", "t-test", "auto"]]
+            ) = "auto",
+            faiss_mode: Literal["base", "fast", "auto"] = "auto",
     ) -> Experiment:
         """Creates an experiment configuration with specified matching parameters.
 
@@ -109,6 +110,7 @@ def _make_experiment(
                 grouping_role=TreatmentRole(),
                 two_sides=two_sides,
                 test_pairs=test_pairs,
+                faiss_mode=faiss_mode,
             )
         ]
         if bias_estimation:
@@ -160,10 +162,12 @@ def __init__(
             Literal["smd", "psi", "ks-test", "repeats", "t-test", "auto"]
             | list[Literal["smd", "psi", "ks-test", "repeats", "t-test", "auto"]]
         ) = "auto",
+        faiss_mode: Literal["base", "fast", "auto"] = "auto",
+
     ):
         super().__init__(
             experiment=self._make_experiment(
-                group_match, distance, metric, bias_estimation, quality_tests
+                group_match, distance, metric, bias_estimation, quality_tests, faiss_mode
             ),
             output=MatchingOutput(GroupExperiment if group_match else MatchingAnalyzer),
         )

hypex/ml/faiss.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Literal
 
 from ..comparators.distances import MahalanobisDistance
 from ..dataset import (
@@ -24,10 +24,12 @@ def __init__(
         test_pairs: bool = False,
         grouping_role: ABCRole | None = None,
         key: Any = "",
+        faiss_mode: Literal["base", "fast", "auto"] = "auto",
     ):
         self.n_neighbors = n_neighbors
         self.two_sides = two_sides
         self.test_pairs = test_pairs
+        self.faiss_mode = faiss_mode
         super().__init__(
             grouping_role=grouping_role, target_role=FeatureRole(), key=key
         )
@@ -40,13 +42,15 @@ def _execute_inner_function(
         n_neighbors: int | None = None,
         two_sides: bool | None = None,
         test_pairs: bool | None = None,
+        faiss_mode: Literal["base", "fast", "auto"] = "auto",
         **kwargs,
     ) -> dict:
         if test_pairs is not True:
             data = cls._inner_function(
                 data=grouping_data[0][1],
                 test_data=grouping_data[1][1],
                 n_neighbors=n_neighbors or 1,
+                faiss_mode=faiss_mode,
                 **kwargs,
             )
             if two_sides is not True:
@@ -57,13 +61,15 @@ def _execute_inner_function(
                     data=grouping_data[1][1],
                     test_data=grouping_data[0][1],
                     n_neighbors=n_neighbors or 1,
+                    faiss_mode=faiss_mode,
                     **kwargs,
                 ),
             }
         data = cls._inner_function(
             data=grouping_data[1][1],
             test_data=grouping_data[0][1],
             n_neighbors=n_neighbors or 1,
+            faiss_mode=faiss_mode,
             **kwargs,
         )
         if two_sides is not True:
@@ -74,6 +80,7 @@ def _execute_inner_function(
                 data=grouping_data[1][1],
                 test_data=grouping_data[0][1],
                 n_neighbors=n_neighbors or 1,
+                faiss_mode=faiss_mode,
                 **kwargs,
             ),
         }
@@ -85,14 +92,15 @@ def _inner_function(
         test_data: Dataset | None = None,
         target_data: Dataset | None = None,
         n_neighbors: int | None = None,
+        faiss_mode: Literal["base", "fast", "auto"] = "auto",
         **kwargs,
     ) -> Any:
-        return FaissExtension(n_neighbors=n_neighbors or 1).calc(
+        return FaissExtension(n_neighbors=n_neighbors or 1, faiss_mode=faiss_mode).calc(
             data=data, test_data=test_data
         )
 
     def fit(self, X: Dataset, Y: Dataset | None = None) -> MLExecutor:
-        return FaissExtension(self.n_neighbors).fit(X=X, Y=Y)
+        return FaissExtension(self.n_neighbors, self.faiss_mode).fit(X=X, Y=Y)
 
     def predict(self, X: Dataset) -> Dataset:
         return FaissExtension().predict(X)
@@ -114,6 +122,7 @@ def execute(self, data: ExperimentData) -> ExperimentData:
             grouping_data=grouping_data,
             features_fields=features_fields,
             n_neighbors=self.n_neighbors,
+            faiss_mode=self.faiss_mode,
             two_sides=self.two_sides,
             test_pairs=self.test_pairs,
         )

tests/test_tutorials.py

@@ -170,6 +170,8 @@ def test_matchingtest(matching_data):
         "matching-atc": Matching(metric="atc"),
         "matching-att": Matching(metric="att"),
         "matching-l2": Matching(distance="l2", metric="att"),
+        "matching-faiss-auto": Matching(distance="l2", faiss_mode="auto"),
+        "matching-faiss_base": Matching(distance="mahalanobis", faiss_mode="base"),
     }
 
     for test_name in mapping.keys():