Refactor imputation module

src/acore/imputation_analysis/__init__.py

@@ -252,6 +252,58 @@ def imputation_normal_distribution(
     return df
 
 
+def _select_imputation_targets(
+    data: pd.DataFrame,
+    on_cols: Iterable[str] | None,
+    on_rows: Iterable[str] | None,
+    drop_cols: Iterable[str] | None,
+) -> tuple[pd.DataFrame, list[str], list[str] | None]:
+    """Copy `data`, drop `drop_cols` and resolve which columns and rows to impute.
+
+    :return: tuple of the copied DataFrame (without `drop_cols`), the numeric columns
+             to impute and the row labels found in the index (`None` if `on_rows`
+             is `None`).
+    :raises TypeError: if `on_cols` contains non-numeric columns.
+    """
+    if drop_cols:
+        if on_cols is not None:
+            overlap = set(on_cols) & set(drop_cols)
+            if overlap:
+                logger.warning(
+                    f"Columns in both on_cols and drop_cols will be dropped, not filled: {overlap}"
+                )
+        df = data.drop(columns=drop_cols)  # already a copy
+    else:
+        df = data.copy()
+
+    if on_cols is None:
+        for col in df.columns[df.dtypes == object]:
+            try:
+                df[col] = pd.to_numeric(df[col])
+            except (ValueError, TypeError):
+                pass
+        cols = df.select_dtypes(include="number").columns.tolist()
+        non_numeric = df.select_dtypes(exclude="number").columns.tolist()
+        if non_numeric:
+            logger.warning(f"Non-numeric columns ignored for imputation: {non_numeric}")
+    else:
+        cols = [c for c in on_cols if c in df.columns]
+        non_numeric = [c for c in cols if not pd.api.types.is_numeric_dtype(df[c])]
+        if non_numeric:
+            raise TypeError(f"Non-numeric columns passed to `on_cols`: {non_numeric}")
+
+    rows = None
+    if on_rows is not None:
+        on_rows = list(on_rows)
+        rows = [r for r in on_rows if r in df.index]
+        if len(rows) != len(on_rows):
+            logger.warning(
+                f"Some rows in `on_rows` were not found in the DataFrame index and will be skipped: "
+                f"{set(on_rows) - set(rows)}"
+            )
+    return df, cols, rows
+
+
 def imputation_zeros(
     data: pd.DataFrame,
     on_cols: Iterable[str] | None = None,
@@ -277,33 +329,9 @@ def imputation_zeros(
         result = imputation_zeros(data, on_cols=['featureA', 'featureB'])
         result = imputation_zeros(data, on_rows=['QC1', 'QC2', 'blank1'])
     """
-    df = data.copy()
-
-    if drop_cols is not None and drop_cols:
-        if on_cols is not None:
-            overlap = set(on_cols) & set(drop_cols)
-            if overlap:
-                logger.warning(
-                    f"Columns in both on_cols and drop_cols will be dropped, not filled: {overlap}"
-                )
-        df = df.drop(columns=drop_cols)
-
-    if on_cols is None:
-        for col in df.columns:
-            if df[col].dtype == object:
-                df[col] = pd.to_numeric(df[col], errors="ignore")
-        cols = df.select_dtypes(include="number").columns.tolist()
-        non_numeric = df.select_dtypes(exclude="number").columns.tolist()
-        if non_numeric:
-            logger.warning(f"Non-numeric columns ignored for imputation: {non_numeric}")
-    else:
-        cols = [c for c in on_cols if c in df.columns]
-        non_numeric = [c for c in cols if not pd.api.types.is_numeric_dtype(df[c])]
-        if non_numeric:
-            raise TypeError(f"Non-numeric columns passed to `on_cols`: {non_numeric}")
+    df, cols, rows = _select_imputation_targets(data, on_cols, on_rows, drop_cols)
 
     if on_rows is not None:
-        rows = [r for r in on_rows if r in df.index]
         df.loc[rows, cols] = df.loc[rows, cols].fillna(0)
     else:
         df[cols] = df[cols].fillna(0)
@@ -337,38 +365,9 @@ def imputation_half_minimum(
         result = imputation_half_minimum(data, on_cols=['featureA', 'featureB'])
         result = imputation_half_minimum(data, on_rows=['blank1', 'blank2'])
     """
-    df = data.copy()
-
-    if drop_cols is not None and drop_cols:
-        if on_cols is not None:
-            overlap = set(on_cols) & set(drop_cols)
-            if overlap:
-                logger.warning(
-                    f"Columns in both on_cols and drop_cols will be dropped, not filled: {overlap}"
-                )
-        df = df.drop(columns=drop_cols)
-
-    if on_cols is None:
-        for col in df.columns:
-            if df[col].dtype == object:
-                df[col] = pd.to_numeric(df[col], errors="ignore")
-        cols = df.select_dtypes(include="number").columns.tolist()
-        non_numeric = df.select_dtypes(exclude="number").columns.tolist()
-        if non_numeric:
-            logger.warning(f"Non-numeric columns ignored for imputation: {non_numeric}")
-    else:
-        cols = [c for c in on_cols if c in df.columns]
-        non_numeric = [c for c in cols if not pd.api.types.is_numeric_dtype(df[c])]
-        if non_numeric:
-            raise TypeError(f"Non-numeric columns passed to `on_cols`: {non_numeric}")
+    df, cols, rows = _select_imputation_targets(data, on_cols, on_rows, drop_cols)
 
     if on_rows is not None:
-        rows = [r for r in on_rows if r in df.index]
-        if len(rows) != len(on_rows):
-            logger.warning(
-                f"Some rows in `on_rows` were not found in the DataFrame index and will be skipped: "
-                f"{set(on_rows) - set(rows)}"
-            )
         subset = df.loc[rows, cols]
         all_nan = [c for c in cols if subset[c].isna().all()]
         if all_nan:

tests/test_imputation.py

@@ -449,3 +449,11 @@ def test_imputation_half_minimum_unknown_row_index_warning(caplog):
         result = imputation_half_minimum(df, on_rows=["r1", "r2", "ghost"])
     assert any("ghost" in msg for msg in caplog.messages)
     assert result.loc["r2", "A"] == pytest.approx(1.0)
+
+
+def test_imputation_zeros_unknown_row_index_warning(caplog):
+    df = pd.DataFrame({"A": [2.0, nan]}, index=["r1", "r2"])
+    with caplog.at_level(logging.WARNING, logger="acore.imputation_analysis"):
+        result = imputation_zeros(df, on_rows=["r1", "r2", "ghost"])
+    assert any("ghost" in msg for msg in caplog.messages)
+    assert result.loc["r2", "A"] == 0.0