Add cupy & cuML support for GPU-native imputation (#23)

Signed-off-by: Lukas Heumos <lukas.heumos@posteo.net>

Author: Zethson

.github/workflows/test-gpu.yaml

@@ -40,7 +40,7 @@ jobs:
         run: nvidia-smi
 
       - name: Install Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: "3.12"
 
@@ -50,7 +50,7 @@ jobs:
           cache-dependency-glob: pyproject.toml
 
       - name: Install fknni
-        run: uv pip install --system -e ".[test,faissgpu]"
+        run: uv pip install --system -e ".[test,faissgpu,rapids12]"
       - name: Pip list
         run: pip list
 

ci/environment.yml

@@ -28,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-04-25T13:28:24.303873023Z",
@@ -40,7 +40,7 @@
    "source": [
     "import numpy as np\n",
     "import pandas as pd\n",
-    "from fknni import FaissImputer\n",
+    "from fknni import FastKNNImputer\n",
     "from sklearn.impute import KNNImputer"
    ]
   },
@@ -795,7 +795,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-04-25T13:28:24.316976982Z",
@@ -822,7 +822,7 @@
     }
    ],
    "source": [
-    "faiss_imputer = FaissImputer(n_neighbors=5, strategy=\"weighted\")\n",
+    "faiss_imputer = FastKNNImputer(n_neighbors=5, strategy=\"weighted\")\n",
     "\n",
     "df_imputed_faiss = faiss_imputer.fit_transform(df_missing)\n",
     "df_imputed_faiss"
@@ -965,7 +965,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-04-25T13:29:06.633118482Z",
@@ -998,7 +998,7 @@
     "X.values[np.unravel_index(missing_indices, X.shape)] = np.nan\n",
     "\n",
     "knn_imputer = KNNImputer(n_neighbors=5)\n",
-    "faiss_imputer = FaissImputer(n_neighbors=5)\n",
+    "faiss_imputer = FastKNNImputer(n_neighbors=5)\n",
     "\n",
     "start_time = time.time()\n",
     "knn_imputed = knn_imputer.fit_transform(X)\n",
@@ -1009,7 +1009,7 @@
     "faiss_time = time.time() - start_time\n",
     "\n",
     "times = [knn_time, faiss_time]\n",
-    "labels = [\"scikit-learn KNNImputer\", \"FaissImputer\"]\n",
+    "labels = [\"scikit-learn KNNImputer\", \"FastKNNImputer\"]\n",
     "plt.bar(labels, times, color=[\"blue\", \"green\"])\n",
     "plt.ylabel(\"Time in seconds\")\n",
     "plt.title(\"Imputation Time Comparison for 10000 samples and 50 features with 10% missing rate\")\n",

docs/notebooks/faiss.ipynb

@@ -21,7 +21,7 @@ classifiers = [
   "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-  "lamin-utils",
+  "array-api-compat",
   "pandas",
   "scikit-learn",
 ]

pyproject.toml

@@ -1,7 +1,7 @@
 from importlib.metadata import version
 
-__all__ = ["faiss"]
+__version__ = version("fknni")
 
-from .faiss import FaissImputer
+from .knn import FastKNNImputer
 
-__version__ = version("fknni")
+__all__ = ["FastKNNImputer", "FaissImputer"]

src/fknni/__init__.py

@@ -0,0 +1 @@
+from .knn import FastKNNImputer