diff --git a/src/acore/imputation_analysis/__init__.py b/src/acore/imputation_analysis/__init__.py index 3f7eba2..75c5441 100644 --- a/src/acore/imputation_analysis/__init__.py +++ b/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: diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 961cdef..790c34b 100644 --- a/tests/test_imputation.py +++ b/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