3D support for mice_forest_impute

ehrapy/preprocessing/_imputation.py

@@ -491,6 +491,7 @@ def miss_forest_impute(
         >>> edata = ed.dt.mimic_2()
         >>> edata = ep.pp.encode(edata, autodetect=True)
         >>> ep.pp.miss_forest_impute(edata)
+
     """
     if copy:
         edata = edata.copy()
@@ -555,7 +556,6 @@ def miss_forest_impute(
 
 @_check_feature_types
 @use_ehrdata(deprecated_after="1.0.0")
-@function_2D_only()
 @spinner("Performing mice-forest impute")
 def mice_forest_impute(
     edata: EHRData | AnnData,
@@ -564,7 +564,6 @@ def mice_forest_impute(
     warning_threshold: int = 70,
     save_all_iterations_data: bool = True,
     random_state: int | None = None,
-    inplace: bool = False,
     iterations: int = 5,
     variable_parameters: dict | None = None,
     verbose: bool = False,
@@ -578,6 +577,9 @@ def mice_forest_impute(
 
     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::
         This function is not supported on MacOS.
 
@@ -588,13 +590,11 @@ def mice_forest_impute(
         save_all_iterations_data: Whether to save all imputed values from all iterations or just the latest.
                              Saving all iterations allows for additional plotting, but may take more memory.
         random_state: The random state ensures script reproducibility.
-        inplace: If True, modify the input data object in-place and return None.
-                 If False, return a copy of the modified data object. Default is False.
         iterations: The number of iterations to run.
         variable_parameters: Model parameters can be specified by variable here.
                              Keys should be variable names or indices, and values should be a dict of parameter which should apply to that variable only.
         verbose: Whether to print information about the imputation process.
-        layer: The layer to impute.
+        layer: The layer to impute. Required when input data is 3D.
         copy: Whether to return a copy of the data object or modify it in-place.
 
     Returns:
@@ -604,13 +604,29 @@ def mice_forest_impute(
     Examples:
         >>> import ehrdata as ed
         >>> import ehrapy as ep
-        >>> edata = ed.dt.mimic_2()
-        >>> edata = ep.pp.encode(edata, autodetect=True)
-        >>> ep.pp.mice_forest_impute(edata)
+        >>> edata = ed.dt.ehrdata_blobs(n_variables=3, n_observations=20, base_timepoints=2, missing_values=0.3)
+        >>> edata_imputed = ep.pp.mice_forest_impute(edata, layer="tem_data", copy=True)
+
+        Example Output:
+
+        >>> edata.layers["tem_data"][0, :, :]
+        [[-11.3735387 , -17.00612946],
+        [         nan,  -3.13348925],
+        [         nan, -10.87061402]]
+        >>> edata_imputed.layers["tem_data"][0, :, :]
+        [[-11.3735387 , -17.00612946],
+        [ -2.29990557,  -3.13348925],
+        [ -6.72812888, -10.87061402]]
+
     """
     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."
+        )
+
     if var_names is None:
         var_names = edata.var_names
     var_indices = edata.var_names.get_indexer(var_names).tolist()
@@ -624,12 +640,12 @@ def mice_forest_impute(
             "Can only impute numerical data. Try to restrict imputation to certain columns using "
             "var_names parameter or perform an encoding of your data."
         )
+
     _miceforest_impute(
         edata,
         var_names,
         save_all_iterations_data,
         random_state,
-        inplace,
         iterations,
         variable_parameters,
         verbose,
@@ -644,50 +660,63 @@ def load_dataframe(arr, columns, index):
     _raise_array_type_not_implemented(load_dataframe, type(arr))
 
 
-@load_dataframe.register
+@load_dataframe.register(np.ndarray)
 def _(arr: np.ndarray, columns, index):
     return pd.DataFrame(arr, columns=columns, index=index)
 
 
 def _miceforest_impute(
-    edata, var_names, save_all_iterations_data, random_state, inplace, iterations, variable_parameters, verbose, layer
+    edata, var_names, save_all_iterations_data, random_state, iterations, variable_parameters, verbose, layer
 ) -> None:
     import miceforest as mf
 
-    data_df = load_dataframe(
-        edata.X if layer is None else edata.layers[layer], columns=edata.var_names, index=edata.obs_names
-    )
-    data_df = data_df.apply(pd.to_numeric, errors="coerce")
-
-    if isinstance(var_names, Iterable) and all(isinstance(item, str) for item in var_names):
-        column_indices = edata.var_names.get_indexer(var_names).tolist()
-        selected_columns = data_df.iloc[:, column_indices]
-        selected_columns = selected_columns.reset_index(drop=True)
+    mtx = edata.X if layer is None else edata.layers[layer]
+    # ensure floating point dtype for 3D flatten, object dtype arrays cause NaN predictions in miceforest's KD-tree
+    input_dtype = mtx.dtype if np.issubdtype(mtx.dtype, np.floating) else np.float64
+    var_indices = edata.var_names.get_indexer(var_names).tolist()
 
-        kernel = mf.ImputationKernel(
-            selected_columns,
-            num_datasets=1,
-            save_all_iterations_data=save_all_iterations_data,
-            random_state=random_state,
+    is_3d = False
+    if mtx.ndim == 3:
+        is_3d = True
+        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))
         )
+        var_names = edata.var_names[var_indices]
+        column_indices = list(range(len(var_indices)))
+    else:
+        column_indices = var_indices
 
-        kernel.mice(iterations=iterations, variable_parameters=variable_parameters or {}, verbose=verbose)
-        data_df.iloc[:, column_indices] = kernel.complete_data(dataset=0, inplace=inplace)
+    idx = pd.RangeIndex(n_obs * n_t) if is_3d else edata.obs_names
 
-    else:
-        data_df = data_df.reset_index(drop=True)
+    data_df = load_dataframe(mtx, columns=var_names, index=idx)
+    data_df = data_df.apply(pd.to_numeric, errors="coerce")
 
-        kernel = mf.ImputationKernel(
-            data_df, num_datasets=1, save_all_iterations_data=save_all_iterations_data, random_state=random_state
-        )
+    # no need for branching as var_names is always either pd.Index or list of strings (resolved to strings before _miceforest_impute is called)
+    selected_columns = data_df.iloc[:, column_indices]
+    selected_columns = selected_columns.reset_index(drop=True)
 
-        kernel.mice(iterations=iterations, variable_parameters=variable_parameters or {}, verbose=verbose)
-        data_df = kernel.complete_data(dataset=0, inplace=inplace)
+    kernel = mf.ImputationKernel(
+        selected_columns,
+        num_datasets=1,
+        save_all_iterations_data=save_all_iterations_data,
+        random_state=random_state,
+    )
 
-    if layer is None:
-        edata.X = data_df.values
+    kernel.mice(iterations=iterations, variable_parameters=variable_parameters or {}, verbose=verbose)
+    data_df.iloc[:, column_indices] = kernel.complete_data(dataset=0, inplace=False)
+
+    if is_3d:
+        result = data_df.values.reshape(n_obs, n_t, len(var_indices)).transpose(0, 2, 1)
+        edata.layers[layer][:, var_indices, :] = result
     else:
-        edata.layers[layer] = data_df.values
+        if layer is None:
+            edata.X = data_df.values
+        else:
+            edata.layers[layer] = data_df.values
 
 
 def _warn_imputation_threshold(

tests/conftest.py

@@ -186,7 +186,8 @@ def edata_mini():
 
 @pytest.fixture
 def edata_mini_3D_missing_values(request):
-    only_numerical = getattr(request, "param", False)
+    params = getattr(request, "param", (False, False))
+    only_numerical, more_obs = params if isinstance(params, tuple) else (params, False)
 
     tiny_mixed_array = np.array(
         [
@@ -198,6 +199,10 @@ def edata_mini_3D_missing_values(request):
         dtype=object,
     )
 
+    if more_obs:
+        # stack the array to get more observations
+        tiny_mixed_array = np.concatenate([tiny_mixed_array] * 5, axis=0)
+
     if only_numerical:
         layer = tiny_mixed_array[:, :4, :]
         feature_types = [NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG, NUMERIC_TAG]

tests/preprocessing/test_imputation.py

@@ -396,15 +396,6 @@ def test_miceforest_array_types(impute_num_edata, array_type, expected_error):
             mice_forest_impute(impute_num_edata, copy=True)
 
 
-@pytest.mark.skipif(platform.system() == "Darwin", reason="miceforest Imputation not supported by MacOS.")
-def test_miceforest_impute_3D_edata(edata_blob_small):
-    edata_blob_small.X[3:5, 4:6] = np.nan
-    edata_blob_small.layers[DEFAULT_TEM_LAYER_NAME][3:5, 4:6] = np.nan
-    mice_forest_impute(edata_blob_small)
-    with pytest.raises(ValueError, match=r"only supports 2D data"):
-        mice_forest_impute(edata_blob_small, layer=DEFAULT_TEM_LAYER_NAME)
-
-
 @pytest.mark.skipif(platform.system() == "Darwin", reason="miceforest Imputation not supported by MacOS.")
 @pytest.mark.filterwarnings("ignore:invalid value encountered in divide:RuntimeWarning")
 def test_miceforest_impute_no_copy(impute_iris_edata):
@@ -438,6 +429,41 @@ def test_miceforest_impute_numerical_data(impute_iris_edata):
     _base_check_imputation(edata_not_imputed, edata_imputed)
 
 
+@pytest.mark.skipif(platform.system() == "Darwin", reason="miceforest Imputation not supported by MacOS.")
+@pytest.mark.filterwarnings("ignore:invalid value encountered in divide:RuntimeWarning")
+@pytest.mark.parametrize("edata_mini_3D_missing_values", [(True, True)], indirect=True)
+def test_miceforest_impute_3D_edata(edata_mini_3D_missing_values):
+    edata_imputed = mice_forest_impute(edata_mini_3D_missing_values, layer=DEFAULT_TEM_LAYER_NAME, copy=True)
+
+    _base_check_imputation(
+        edata_mini_3D_missing_values,
+        edata_imputed,
+        before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+        after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+    )
+    assert id(edata_mini_3D_missing_values) != id(edata_imputed)
+
+
+@pytest.mark.parametrize("edata_mini_3D_missing_values", [(True, True)], indirect=True)
+def test_miceforest_impute_3D_var_names_subset(edata_mini_3D_missing_values):
+    edata = edata_mini_3D_missing_values.copy()
+    imputed = mice_forest_impute(edata, layer=DEFAULT_TEM_LAYER_NAME, var_names=["1", "2"], copy=True)
+    edata_imputed = imputed[:, :2].copy()
+    _base_check_imputation(
+        edata_mini_3D_missing_values[:, :2],
+        edata_imputed,
+        before_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+        after_imputation_layer=DEFAULT_TEM_LAYER_NAME,
+    )
+
+    assert edata.shape == imputed.shape
+
+
+def test_miceforest_impute_3d_layer_none(edata_mini_3D_missing_values):
+    with pytest.raises(ValueError, match="requires a layer"):
+        mice_forest_impute(edata_mini_3D_missing_values, copy=True)
+
+
 @pytest.mark.parametrize(
     "array_type,expected_error",
     [