From c3f14eb570f770de2b897d74f783cc233380f534 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Fri, 27 Mar 2026 11:55:49 +0100
Subject: [PATCH 1/9] knn_impute with 3d, its tests, manually assigned feature
 types for the edata_mini_3D_missing_values

---
 ehrapy/preprocessing/_imputation.py    | 28 +++++++++++++++++++-------
 tests/conftest.py                      |  4 +++-
 tests/preprocessing/test_imputation.py | 13 ++++++++----
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index 1a4db5b2..a97075ef 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -268,7 +268,6 @@ def simple_impute(
 
 @_check_feature_types
 @use_ehrdata(deprecated_after="1.0.0")
-@function_2D_only()
 @spinner("Performing KNN impute")
 def knn_impute(
     edata: EHRData | AnnData,
@@ -388,14 +387,29 @@ def _knn_impute(
             "var_names parameter or perform an encoding of your data."
         )
     X = edata.X if layer is None else edata.layers[layer]
-    complete_numerical_columns = np.array(numerical_indices)[~np.isnan(X[:, numerical_indices]).any(axis=0)].tolist()
-    imputer_data_indices = var_indices + [i for i in complete_numerical_columns if i not in var_indices]
-    imputer_x = X[::, imputer_data_indices].astype("float64")
 
-    if layer is None:
-        edata.X[::, imputer_data_indices] = imputer.fit_transform(imputer_x)
+    if X.ndim == 3:
+        if backend == "scikit-learn":
+            # Since scikit-learn is relatively slow even for 2D, if we want to implement it we can do it via slicing per timepoint
+            raise NotImplementedError(
+                "The 'scikit-learn' backend does not support 3D data. Use backend='faiss' instead."
+            )
+        # complete columns depend on the timepoint, so computing it here would be wrong
+        imputer_x = X[:, var_indices, :].astype("float64")
+        if layer is None:
+            edata.X[:, var_indices, :] = imputer.fit_transform(imputer_x)
+        else:
+            edata.layers[layer][:, var_indices, :] = imputer.fit_transform(imputer_x)
     else:
-        edata.layers[layer][::, imputer_data_indices] = imputer.fit_transform(imputer_x)
+        complete_numerical_columns = np.array(numerical_indices)[
+            ~np.isnan(X[:, numerical_indices]).any(axis=0)
+        ].tolist()
+        imputer_data_indices = var_indices + [i for i in complete_numerical_columns if i not in var_indices]
+        imputer_x = X[:, imputer_data_indices].astype("float64")
+        if layer is None:
+            edata.X[:, imputer_data_indices] = imputer.fit_transform(imputer_x)
+        else:
+            edata.layers[layer][:, imputer_data_indices] = imputer.fit_transform(imputer_x)
 
 
 @use_ehrdata(deprecated_after="1.0.0")
diff --git a/tests/conftest.py b/tests/conftest.py
index 0c29527b..80d8c1b3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -196,7 +196,9 @@ def edata_mini_3D_missing_values():
         dtype=object,
     )
     n_obs, n_vars, _ = tiny_mixed_array.shape
-    return ed.EHRData(shape=(n_obs, n_vars), layers={DEFAULT_TEM_LAYER_NAME: tiny_mixed_array})
+    edata = ed.EHRData(shape=(n_obs, n_vars), layers={DEFAULT_TEM_LAYER_NAME: tiny_mixed_array})
+    edata.var[FEATURE_TYPE_KEY] = [NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG, CATEGORICAL_TAG, CATEGORICAL_TAG]
+    return edata
 
 
 @pytest.fixture
diff --git a/tests/preprocessing/test_imputation.py b/tests/preprocessing/test_imputation.py
index a0144c6f..2ee43522 100644
--- a/tests/preprocessing/test_imputation.py
+++ b/tests/preprocessing/test_imputation.py
@@ -278,10 +278,15 @@ def test_simple_impute_invalid_strategy(impute_edata):
         simple_impute(impute_edata, strategy="invalid_strategy", copy=True)  # type: ignore
 
 
-def test_knn_impute_3D_edata(edata_blob_small):
-    knn_impute(edata_blob_small, layer="layer_2")
-    with pytest.raises(ValueError, match=r"only supports 2D data"):
-        knn_impute(edata_blob_small, layer=DEFAULT_TEM_LAYER_NAME)
+def test_knn_impute_3d_numerical(edata_mini_3D_missing_values):
+    edata = edata_mini_3D_missing_values[:, :4].copy()
+    edata_imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, copy=True)
+    _base_check_imputation(
+        edata,
+        edata_imputed,
+        before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+        after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+    )
 
 
 def test_knn_impute_check_backend(impute_num_edata):

From 39b8416c60314a5fe9d6df6d63579dbd367d7749 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Tue, 31 Mar 2026 10:28:39 +0200
Subject: [PATCH 2/9] tests for 3D knn extended, case when X and layer both
 None handled

---
 ehrapy/preprocessing/_imputation.py    | 11 +++++++----
 tests/preprocessing/test_imputation.py | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index a97075ef..6e64cc48 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -323,6 +323,11 @@ def knn_impute(
     if copy:
         edata = edata.copy()
 
+    if edata.X is None and layer is None:  # if edata is 3D
+        raise ValueError(
+            "3D imputation requires a layer to be specified. Pass the layer containing the full temporal data."
+        )
+
     _warn_imputation_threshold(edata, var_names, threshold=warning_threshold, layer=layer)
 
     if backend not in {"scikit-learn", "faiss"}:
@@ -396,10 +401,8 @@ def _knn_impute(
             )
         # complete columns depend on the timepoint, so computing it here would be wrong
         imputer_x = X[:, var_indices, :].astype("float64")
-        if layer is None:
-            edata.X[:, var_indices, :] = imputer.fit_transform(imputer_x)
-        else:
-            edata.layers[layer][:, var_indices, :] = imputer.fit_transform(imputer_x)
+        edata.layers[layer][:, var_indices, :] = imputer.fit_transform(imputer_x)
+
     else:
         complete_numerical_columns = np.array(numerical_indices)[
             ~np.isnan(X[:, numerical_indices]).any(axis=0)
diff --git a/tests/preprocessing/test_imputation.py b/tests/preprocessing/test_imputation.py
index 2ee43522..b490634a 100644
--- a/tests/preprocessing/test_imputation.py
+++ b/tests/preprocessing/test_imputation.py
@@ -289,6 +289,23 @@ def test_knn_impute_3d_numerical(edata_mini_3D_missing_values):
     )
 
 
+def test_knn_impute_3d_var_names_subset(edata_mini_3D_missing_values):
+    edata = edata_mini_3D_missing_values.copy()
+    imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, var_names=["1", "2"], copy=True)
+    edata_imputed = imputed[:, :2].copy()
+    _base_check_imputation(
+        edata[:, :2],
+        edata_imputed,
+        before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+        after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+    )
+
+
+def test_knn_impute_3d_layer_none(edata_mini_3D_missing_values):
+    with pytest.raises(ValueError, match="requires a layer"):
+        knn_impute(edata_mini_3D_missing_values, copy=True)
+
+
 def test_knn_impute_check_backend(impute_num_edata):
     knn_impute(impute_num_edata, backend="faiss", copy=True)
     knn_impute(impute_num_edata, backend="scikit-learn", copy=True)

From 1a4e39a27f61c7d3894e02b615b65b9c871fc2c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Thu, 9 Apr 2026 15:15:28 +0200
Subject: [PATCH 3/9] 3d support for knn_impute with vertical stacking

---
 ehrapy/preprocessing/_imputation.py | 42 ++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 16 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index 6e64cc48..abdc976f 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -392,27 +392,37 @@ def _knn_impute(
             "var_names parameter or perform an encoding of your data."
         )
     X = edata.X if layer is None else edata.layers[layer]
+    var_indices_original = var_indices
+    is_3d = False
 
+    # if input data is 3D flatten along axis 0 before giving it to the imputer: each timepoint becomes a row
     if X.ndim == 3:
-        if backend == "scikit-learn":
-            # Since scikit-learn is relatively slow even for 2D, if we want to implement it we can do it via slicing per timepoint
-            raise NotImplementedError(
-                "The 'scikit-learn' backend does not support 3D data. Use backend='faiss' instead."
-            )
-        # complete columns depend on the timepoint, so computing it here would be wrong
-        imputer_x = X[:, var_indices, :].astype("float64")
-        edata.layers[layer][:, var_indices, :] = imputer.fit_transform(imputer_x)
-
+        is_3d = True
+        n_obs, n_vars, n_t = X.shape
+        X = X[:, var_indices, :].astype("float64").transpose(0, 2, 1).reshape(n_obs * n_t, len(var_indices))
+        numerical_indices = list(range(len(var_indices)))
+        var_indices = numerical_indices
+
+    # complete columns to be used as anchors
+    complete_numerical_columns = np.array(numerical_indices)[~np.isnan(X[:, numerical_indices]).any(axis=0)].tolist()
+
+    imputer_data_indices = var_indices + [
+        i for i in complete_numerical_columns if i not in var_indices
+    ]  # columns to impute
+    imputer_x = X[:, imputer_data_indices].astype("float64")
+    X_imputed = imputer.fit_transform(imputer_x)
+
+    if is_3d:
+        # slice back to only requested columns and transpose back to n_obs, n_var, n_t
+        X_imputed = (
+            X_imputed[:, : len(var_indices_original)].reshape(n_obs, n_t, len(var_indices_original)).transpose(0, 2, 1)
+        )
+        edata.layers[layer][:, var_indices_original, :] = X_imputed
     else:
-        complete_numerical_columns = np.array(numerical_indices)[
-            ~np.isnan(X[:, numerical_indices]).any(axis=0)
-        ].tolist()
-        imputer_data_indices = var_indices + [i for i in complete_numerical_columns if i not in var_indices]
-        imputer_x = X[:, imputer_data_indices].astype("float64")
         if layer is None:
-            edata.X[:, imputer_data_indices] = imputer.fit_transform(imputer_x)
+            edata.X[:, imputer_data_indices] = X_imputed
         else:
-            edata.layers[layer][:, imputer_data_indices] = imputer.fit_transform(imputer_x)
+            edata.layers[layer][:, imputer_data_indices] = X_imputed
 
 
 @use_ehrdata(deprecated_after="1.0.0")

From 2f6b357dce9b9a87ee4af51dcd6e6a9fe859fe2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Fri, 10 Apr 2026 11:35:27 +0200
Subject: [PATCH 4/9] test with scikit-learn as backend

---
 ehrapy/preprocessing/_imputation.py    |  2 +-
 tests/preprocessing/test_imputation.py | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index abdc976f..45eb51a1 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -295,7 +295,7 @@ def knn_impute(
         var_names: A list of variable names indicating which columns to impute.
                    If `None`, all columns are imputed. Default is `None`.
         n_neighbors: Number of neighbors to use when performing the imputation.
-        layer: The layer to impute.
+        layer: The layer to impute. Required when the input data is 3D.
         copy: Whether to perform the imputation on a copy of the original data object.
               If `True`, the original object remains unmodified.
         backend: The implementation to use for the KNN imputation.
diff --git a/tests/preprocessing/test_imputation.py b/tests/preprocessing/test_imputation.py
index b490634a..b607551c 100644
--- a/tests/preprocessing/test_imputation.py
+++ b/tests/preprocessing/test_imputation.py
@@ -289,6 +289,17 @@ def test_knn_impute_3d_numerical(edata_mini_3D_missing_values):
     )
 
 
+def test_knn_impute_3d_scikit_backend(edata_mini_3D_missing_values):
+    edata = edata_mini_3D_missing_values[:, :4].copy()
+    edata_imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, copy=True, backend="scikit-learn")
+    _base_check_imputation(
+        edata,
+        edata_imputed,
+        before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+        after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+    )
+
+
 def test_knn_impute_3d_var_names_subset(edata_mini_3D_missing_values):
     edata = edata_mini_3D_missing_values.copy()
     imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, var_names=["1", "2"], copy=True)

From c66ac4cb4242d14511233915f3372aaf33514197 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Tue, 14 Apr 2026 09:15:33 +0200
Subject: [PATCH 5/9] comments addressed, not sure of astype() comment

---
 ehrapy/preprocessing/_imputation.py    | 22 ++++++++++++++--------
 tests/conftest.py                      | 25 +++++++++++++++++++++----
 tests/preprocessing/test_imputation.py | 12 +++++++-----
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index 45eb51a1..0ff47550 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -391,25 +391,31 @@ def _knn_impute(
             "Can only impute numerical data. Try to restrict imputation to certain columns using "
             "var_names parameter or perform an encoding of your data."
         )
-    X = edata.X if layer is None else edata.layers[layer]
+    mtx = edata.X if layer is None else edata.layers[layer]
     var_indices_original = var_indices
     is_3d = False
+    input_dtype = mtx.dtype if np.issubdtype(mtx.dtype, np.floating) else np.float64
 
-    # if input data is 3D flatten along axis 0 before giving it to the imputer: each timepoint becomes a row
-    if X.ndim == 3:
+    # if input data is 3D, flatten along axis 0 before passing it to the imputer: each timepoint becomes a row
+    if mtx.ndim == 3:
         is_3d = True
-        n_obs, n_vars, n_t = X.shape
-        X = X[:, var_indices, :].astype("float64").transpose(0, 2, 1).reshape(n_obs * n_t, len(var_indices))
+        n_obs, n_vars, n_t = mtx.shape
+        mtx = (
+            mtx[:, var_indices, :]
+            .astype(input_dtype, copy=True)
+            .transpose(0, 2, 1)
+            .reshape(n_obs * n_t, len(var_indices))
+        )
         numerical_indices = list(range(len(var_indices)))
         var_indices = numerical_indices
 
     # complete columns to be used as anchors
-    complete_numerical_columns = np.array(numerical_indices)[~np.isnan(X[:, numerical_indices]).any(axis=0)].tolist()
+    complete_numerical_columns = np.array(numerical_indices)[~np.isnan(mtx[:, numerical_indices]).any(axis=0)].tolist()
 
     imputer_data_indices = var_indices + [
-        i for i in complete_numerical_columns if i not in var_indices
+        column for column in complete_numerical_columns if column not in var_indices
     ]  # columns to impute
-    imputer_x = X[:, imputer_data_indices].astype("float64")
+    imputer_x = mtx[:, imputer_data_indices].astype(input_dtype, copy=True)
     X_imputed = imputer.fit_transform(imputer_x)
 
     if is_3d:
diff --git a/tests/conftest.py b/tests/conftest.py
index 80d8c1b3..0fb87c4c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -185,7 +185,9 @@ def edata_mini():
 
 
 @pytest.fixture
-def edata_mini_3D_missing_values():
+def edata_mini_3D_missing_values(request):
+    only_numerical = getattr(request, "param", False)
+
     tiny_mixed_array = np.array(
         [
             [[138, 139], [78, np.nan], [77, 76], [1, 2], ["A", "B"], ["Yes", np.nan]],
@@ -195,9 +197,24 @@ def edata_mini_3D_missing_values():
         ],
         dtype=object,
     )
-    n_obs, n_vars, _ = tiny_mixed_array.shape
-    edata = ed.EHRData(shape=(n_obs, n_vars), layers={DEFAULT_TEM_LAYER_NAME: tiny_mixed_array})
-    edata.var[FEATURE_TYPE_KEY] = [NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG, CATEGORICAL_TAG, CATEGORICAL_TAG]
+
+    if only_numerical:
+        layer = tiny_mixed_array[:, :4, :]
+        feature_types = [NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG]
+    else:
+        layer = tiny_mixed_array
+        feature_types = [
+            NUMERIC_TAG,
+            NUMERIC_TAG,
+            NUMERIC_TAG,
+            NUMERIC_TAG,
+            CATEGORICAL_TAG,
+            CATEGORICAL_TAG,
+        ]
+
+    n_obs, n_vars, _ = layer.shape
+    edata = ed.EHRData(shape=(n_obs, n_vars), layers={DEFAULT_TEM_LAYER_NAME: layer})
+    edata.var[FEATURE_TYPE_KEY] = feature_types
     return edata
 
 
diff --git a/tests/preprocessing/test_imputation.py b/tests/preprocessing/test_imputation.py
index b607551c..11f64474 100644
--- a/tests/preprocessing/test_imputation.py
+++ b/tests/preprocessing/test_imputation.py
@@ -278,22 +278,24 @@ def test_simple_impute_invalid_strategy(impute_edata):
         simple_impute(impute_edata, strategy="invalid_strategy", copy=True)  # type: ignore
 
 
+@pytest.mark.parametrize("edata_mini_3D_missing_values", [True], indirect=True)
 def test_knn_impute_3d_numerical(edata_mini_3D_missing_values):
-    edata = edata_mini_3D_missing_values[:, :4].copy()
+    edata = edata_mini_3D_missing_values.copy()
     edata_imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, copy=True)
     _base_check_imputation(
-        edata,
+        edata_mini_3D_missing_values,
         edata_imputed,
         before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
         after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
     )
 
 
+@pytest.mark.parametrize("edata_mini_3D_missing_values", [True], indirect=True)
 def test_knn_impute_3d_scikit_backend(edata_mini_3D_missing_values):
-    edata = edata_mini_3D_missing_values[:, :4].copy()
+    edata = edata_mini_3D_missing_values.copy()
     edata_imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, copy=True, backend="scikit-learn")
     _base_check_imputation(
-        edata,
+        edata_mini_3D_missing_values,
         edata_imputed,
         before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
         after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
@@ -305,7 +307,7 @@ def test_knn_impute_3d_var_names_subset(edata_mini_3D_missing_values):
     imputed = knn_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, var_names=["1", "2"], copy=True)
     edata_imputed = imputed[:, :2].copy()
     _base_check_imputation(
-        edata[:, :2],
+        edata_mini_3D_missing_values[:, :2],
         edata_imputed,
         before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
         after_imputation_layer=DEFAULT_TEM_LAYER_NAME,

From 7ad2ad8bf513e6c31e8205fe7dc88cc89be20ba2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Tue, 14 Apr 2026 12:02:58 +0200
Subject: [PATCH 6/9] docstring

---
 ehrapy/preprocessing/_imputation.py | 35 ++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index 0ff47550..f28fb190 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -284,6 +284,8 @@ def knn_impute(
     """Imputes missing values in the input data object using K-nearest neighbor imputation.
 
     If required, the data needs to be properly encoded as this imputation requires numerical data only.
+    For 2D data, if layer is `None`, `edata.X` is used directly.
+    For 3D data, the layer is flattened along axis 0 before imputation and reshaped back to 3D afterwards.
 
     .. warning::
         Currently, both `n_neighbours` and `n_neighbors` are accepted as parameters for the number of neighbors.
@@ -316,9 +318,36 @@ def knn_impute(
     Examples:
         >>> import ehrdata as ed
         >>> import ehrapy as ep
-        >>> edata = ed.dt.mimic_2()
-        >>> ed.infer_feature_types(edata)
-        >>> ep.pp.knn_impute(edata)
+        >>> edata_3d = ed.dt.ehrdata_blobs(n_variables=3, n_observations=3, base_timepoints=2, missing_values=0.3)
+        >>> edata_imputed = ep.pp.knn_impute(edata_3d, layer="tem_data", copy=True)
+
+        Example Output:
+        >>> edata_3d.layers["tem_data"]
+        [[[-12.12732884, -18.37304373],
+        [         nan,  -0.91339411],
+        [         nan,  -7.88514984]],
+
+       [[  0.46031019,          nan],
+        [  0.75960393,          nan],
+        [  2.37346143,          nan]],
+
+       [[  1.17615504,   2.21695734],
+        [         nan,          nan],
+        [  3.63259441,          nan]]]
+
+        >>> edata_imputed.layers["tem_data"]
+        [[[-12.12732884, -18.37304373],
+        [ -0.07689509,  -0.91339411],
+        [ -2.75584421,  -7.88514984]],
+
+       [[  0.46031019,  -8.95636677],
+        [  0.75960393,  -0.07689509],
+        [  2.37346143,  -2.75584421]],
+
+       [[  1.17615504,   2.21695734],
+        [ -0.07689509,  -0.07689509],
+        [  3.63259441,  -2.75584421]]]
+
     """
     if copy:
         edata = edata.copy()

From 13e26af8ba64a27e9eb0ea2eac91935318ab0b25 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Tue, 14 Apr 2026 12:09:44 +0200
Subject: [PATCH 7/9] fix

---
 ehrapy/preprocessing/_imputation.py | 30 +++++------------------------
 1 file changed, 5 insertions(+), 25 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index f28fb190..b2e30cd9 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -322,31 +322,11 @@ def knn_impute(
         >>> edata_imputed = ep.pp.knn_impute(edata_3d, layer="tem_data", copy=True)
 
         Example Output:
-        >>> edata_3d.layers["tem_data"]
-        [[[-12.12732884, -18.37304373],
-        [         nan,  -0.91339411],
-        [         nan,  -7.88514984]],
-
-       [[  0.46031019,          nan],
-        [  0.75960393,          nan],
-        [  2.37346143,          nan]],
-
-       [[  1.17615504,   2.21695734],
-        [         nan,          nan],
-        [  3.63259441,          nan]]]
-
-        >>> edata_imputed.layers["tem_data"]
-        [[[-12.12732884, -18.37304373],
-        [ -0.07689509,  -0.91339411],
-        [ -2.75584421,  -7.88514984]],
-
-       [[  0.46031019,  -8.95636677],
-        [  0.75960393,  -0.07689509],
-        [  2.37346143,  -2.75584421]],
-
-       [[  1.17615504,   2.21695734],
-        [ -0.07689509,  -0.07689509],
-        [  3.63259441,  -2.75584421]]]
+        >>> edata_3d.layers["tem_data"][0, :, :]
+        [[-12.12732884, -18.37304373], [         nan,  -0.91339411], [         nan,  -7.88514984]]
+
+        >>> edata_imputed.layers["tem_data"][0, :, :]
+        [[-12.12732884, -18.37304373], [ -0.07689509,  -0.91339411], [ -2.75584421,  -7.88514984]]
 
     """
     if copy:

From 3f972df3288b4ef01f3e751e2e266405b84fada1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Tue, 14 Apr 2026 12:12:25 +0200
Subject: [PATCH 8/9] fix

---
 ehrapy/preprocessing/_imputation.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index b2e30cd9..4db72333 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -324,7 +324,6 @@ def knn_impute(
         Example Output:
         >>> edata_3d.layers["tem_data"][0, :, :]
         [[-12.12732884, -18.37304373], [         nan,  -0.91339411], [         nan,  -7.88514984]]
-
         >>> edata_imputed.layers["tem_data"][0, :, :]
         [[-12.12732884, -18.37304373], [ -0.07689509,  -0.91339411], [ -2.75584421,  -7.88514984]]
 

From 672fba93ec408d7b69f80807271e901eea56863c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96yk=C3=BC=20S=C3=BCoglu?= <oykusuogluu@gmail.com>
Date: Tue, 14 Apr 2026 12:15:36 +0200
Subject: [PATCH 9/9] fix

---
 ehrapy/preprocessing/_imputation.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/ehrapy/preprocessing/_imputation.py b/ehrapy/preprocessing/_imputation.py
index 4db72333..04e890bf 100644
--- a/ehrapy/preprocessing/_imputation.py
+++ b/ehrapy/preprocessing/_imputation.py
@@ -322,10 +322,15 @@ def knn_impute(
         >>> edata_imputed = ep.pp.knn_impute(edata_3d, layer="tem_data", copy=True)
 
         Example Output:
+
         >>> edata_3d.layers["tem_data"][0, :, :]
-        [[-12.12732884, -18.37304373], [         nan,  -0.91339411], [         nan,  -7.88514984]]
+        [[-12.12732884, -18.37304373],
+        [         nan,  -0.91339411],
+        [         nan,  -7.88514984]]
         >>> edata_imputed.layers["tem_data"][0, :, :]
-        [[-12.12732884, -18.37304373], [ -0.07689509,  -0.91339411], [ -2.75584421,  -7.88514984]]
+        [[-12.12732884, -18.37304373],
+        [ -0.07689509,  -0.91339411],
+        [ -2.75584421,  -7.88514984]]
 
     """
     if copy: