smftools 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

smftools/preprocessing/calculate_position_Youden.py

@@ -1,69 +1,95 @@
 ## calculate_position_Youden
 ## Calculating and applying position level thresholds for methylation calls to binarize the SMF data
-def calculate_position_Youden(adata, 
-                              positive_control_sample=None, 
-                              negative_control_sample=None, 
-                              J_threshold=0.5, 
-                              ref_column='Reference_strand', 
-                              sample_column='Sample_names',
-                              infer_on_percentile=True, 
-                              inference_variable='Raw_modification_signal', 
-                              save=False, 
-                              output_directory=''):
-    """
-    Adds new variable metadata to each position indicating whether the position provides reliable SMF methylation calls. Also outputs plots of the positional ROC curves.
+from __future__ import annotations
 
-    Parameters:
-        adata (AnnData): An AnnData object.
-        positive_control_sample (str): string representing the sample name corresponding to the Plus MTase control sample.
-        negative_control_sample (str): string representing the sample name corresponding to the Minus MTase control sample.
-        J_threshold (float): A float indicating the J-statistic used to indicate whether a position passes QC for methylation calls.
-        obs_column (str): The category to iterate over.
-        infer_on_perdentile (bool | int): If False, use defined postive and negative control samples. If an int (0 < int < 100) is passed, this uses the top and bottom int percentile of methylated reads based on metric in inference_variable column.
-        inference_variable (str): If infer_on_percentile has an integer value passed, use the AnnData observation column name passed by this string as the metric.
-        save (bool): Whether to save the ROC plots.
-        output_directory (str): String representing the path to the output directory to output the ROC curves.
-    
-    Returns: 
-        None
-    """
-    import numpy as np
-    import pandas as pd
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from smftools.logging_utils import get_logger
+
+if TYPE_CHECKING:
     import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def calculate_position_Youden(
+    adata: "ad.AnnData",
+    positive_control_sample: str | None = None,
+    negative_control_sample: str | None = None,
+    J_threshold: float = 0.5,
+    ref_column: str = "Reference_strand",
+    sample_column: str = "Sample_names",
+    infer_on_percentile: bool | int = True,
+    inference_variable: str = "Raw_modification_signal",
+    save: bool = False,
+    output_directory: str | Path = "",
+) -> None:
+    """Add position-level Youden thresholds and optional ROC plots.
+
+    Args:
+        adata: AnnData object.
+        positive_control_sample: Sample name for the plus MTase control.
+        negative_control_sample: Sample name for the minus MTase control.
+        J_threshold: J-statistic threshold for QC.
+        ref_column: Obs column for reference/strand categories.
+        sample_column: Obs column for sample identifiers.
+        infer_on_percentile: If ``False``, use provided controls. If an int in ``(0, 100)``,
+            use percentile-based inference from ``inference_variable``.
+        inference_variable: Obs column used for percentile inference.
+        save: Whether to save ROC plots to disk.
+        output_directory: Output directory for ROC plots.
+    """
     import matplotlib.pyplot as plt
-    from sklearn.metrics import roc_curve, roc_auc_score
+    import numpy as np
+    from sklearn.metrics import roc_curve
 
     control_samples = [positive_control_sample, negative_control_sample]
-    references = adata.obs[ref_column].cat.categories 
+    references = adata.obs[ref_column].cat.categories
     # Iterate over each category in the specified obs_column
     for ref in references:
-        print(f"Calculating position Youden statistics for {ref}")
+        logger.info("Calculating position Youden statistics for %s", ref)
         # Subset to keep only reads associated with the category
         ref_subset = adata[adata.obs[ref_column] == ref]
         # Iterate over positive and negative control samples
         for i, control in enumerate(control_samples):
-            # Initialize a dictionary for the given control sample. This will be keyed by dataset and position to point to a tuple of coordinate position and an array of methylation probabilities
-            adata.uns[f'{ref}_position_methylation_dict_{control}'] = {}
             # If controls are not passed and infer on percentile is True, infer thresholds based on top and bottom percentile windows for a given obs column metric.
             if infer_on_percentile and not control:
+                logger.info(
+                    "Inferring methylation control reads for %s based on %s percentiles of %s",
+                    ref,
+                    infer_on_percentile,
+                    inference_variable,
+                )
                 sorted_column = ref_subset.obs[inference_variable].sort_values(ascending=False)
                 if i == 0:
-                    control == 'positive'
+                    logger.info("Using top %s percentile for positive control", infer_on_percentile)
+                    control = "positive"
                     positive_control_sample = control
                     threshold = np.percentile(sorted_column, 100 - infer_on_percentile)
                     control_subset = ref_subset[ref_subset.obs[inference_variable] >= threshold, :]
                 else:
-                    control == 'negative'
+                    logger.info(
+                        "Using bottom %s percentile for negative control", infer_on_percentile
+                    )
+                    control = "negative"
                     negative_control_sample = control
                     threshold = np.percentile(sorted_column, infer_on_percentile)
-                    control_subset = ref_subset[ref_subset.obs[inference_variable] <= threshold, :]   
+                    control_subset = ref_subset[ref_subset.obs[inference_variable] <= threshold, :]
             elif not infer_on_percentile and not control:
-                print("Can not threshold Anndata on Youden threshold. Need to either provide control samples or set infer_on_percentile to True")
+                logger.error(
+                    "Can not threshold Anndata on Youden threshold. Need to either provide control samples or set infer_on_percentile to True"
+                )
                 return
             else:
+                logger.info("Using provided control sample: %s", control)
                 # get the current control subset on the given category
                 filtered_obs = ref_subset.obs[ref_subset.obs[sample_column] == control]
                 control_subset = ref_subset[filtered_obs.index]
+
+            # Initialize a dictionary for the given control sample. This will be keyed by dataset and position to point to a tuple of coordinate position and an array of methylation probabilities
+            adata.uns[f"{ref}_position_methylation_dict_{control}"] = {}
+
             # Iterate through every position in the control subset
             for position in range(control_subset.shape[1]):
                 # Get the coordinate name associated with that position
@@ -79,29 +105,45 @@ def calculate_position_Youden(adata,
                 # Get fraction coverage
                 fraction_coverage = position_coverage / control_subset.shape[0]
                 # Save the position and the position methylation data for the control subset
-                adata.uns[f'{ref}_position_methylation_dict_{control}'][f'{position}'] = (position, position_data, fraction_coverage)
+                adata.uns[f"{ref}_position_methylation_dict_{control}"][f"{position}"] = (
+                    position,
+                    position_data,
+                    fraction_coverage,
+                )
 
     for ref in references:
         fig, ax = plt.subplots(figsize=(6, 4))
-        plt.plot([0, 1], [0, 1], linestyle='--', color='gray')
-        plt.xlabel('False Positive Rate')
-        plt.ylabel('True Positive Rate')
-        ax.spines['right'].set_visible(False)
-        ax.spines['top'].set_visible(False)
+        plt.plot([0, 1], [0, 1], linestyle="--", color="gray")
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        ax.spines["right"].set_visible(False)
+        ax.spines["top"].set_visible(False)
         n_passed_positions = 0
         n_total_positions = 0
         # Initialize a list that will hold the positional thresholds for the category
         probability_thresholding_list = [(np.nan, np.nan)] * adata.shape[1]
-        for i, key in enumerate(adata.uns[f'{ref}_position_methylation_dict_{positive_control_sample}'].keys()):
-            position = int(adata.uns[f'{ref}_position_methylation_dict_{positive_control_sample}'][key][0])
-            positive_position_array = adata.uns[f'{ref}_position_methylation_dict_{positive_control_sample}'][key][1]
-            fraction_coverage = adata.uns[f'{ref}_position_methylation_dict_{positive_control_sample}'][key][2]
+        for i, key in enumerate(
+            adata.uns[f"{ref}_position_methylation_dict_{positive_control_sample}"].keys()
+        ):
+            position = int(
+                adata.uns[f"{ref}_position_methylation_dict_{positive_control_sample}"][key][0]
+            )
+            positive_position_array = adata.uns[
+                f"{ref}_position_methylation_dict_{positive_control_sample}"
+            ][key][1]
+            fraction_coverage = adata.uns[
+                f"{ref}_position_methylation_dict_{positive_control_sample}"
+            ][key][2]
             if fraction_coverage > 0.2:
                 try:
-                    negative_position_array = adata.uns[f'{ref}_position_methylation_dict_{negative_control_sample}'][key][1] 
+                    negative_position_array = adata.uns[
+                        f"{ref}_position_methylation_dict_{negative_control_sample}"
+                    ][key][1]
                     # Combine the negative and positive control data
                     data = np.concatenate([negative_position_array, positive_position_array])
-                    labels = np.array([0] * len(negative_position_array) + [1] * len(positive_position_array))
+                    labels = np.array(
+                        [0] * len(negative_position_array) + [1] * len(positive_position_array)
+                    )
                     # Calculate the ROC curve
                     fpr, tpr, thresholds = roc_curve(labels, data)
                     # Calculate Youden's J statistic
@@ -114,18 +156,24 @@ def calculate_position_Youden(adata,
                     n_total_positions += 1
                     if max_J > J_threshold:
                         n_passed_positions += 1
-                        plt.plot(fpr, tpr, label='ROC curve')
-                except:
+                        plt.plot(fpr, tpr, label="ROC curve")
+                except Exception:
                     probability_thresholding_list[position] = (0.8, np.nan)
-        title = f'ROC Curve for {n_passed_positions} positions with J-stat greater than {J_threshold}\n out of {n_total_positions} total positions on {ref}'
+        title = f"ROC Curve for {n_passed_positions} positions with J-stat greater than {J_threshold}\n out of {n_total_positions} total positions on {ref}"
         plt.title(title)
         save_name = output_directory / f"{title}.png"
         if save:
             plt.savefig(save_name)
             plt.close()
         else:
-            plt.show()   
+            plt.show()
 
-        adata.var[f'{ref}_position_methylation_thresholding_Youden_stats'] = probability_thresholding_list
+        adata.var[f"{ref}_position_methylation_thresholding_Youden_stats"] = (
+            probability_thresholding_list
+        )
         J_max_list = [probability_thresholding_list[i][1] for i in range(adata.shape[1])]
-        adata.var[f'{ref}_position_passed_QC'] = [True if i > J_threshold else False for i in J_max_list]
+        adata.var[f"{ref}_position_passed_Youden_thresholding_QC"] = [
+            True if i > J_threshold else False for i in J_max_list
+        ]
+
+    logger.info("Finished calculating position Youden statistics")

smftools/preprocessing/calculate_read_length_stats.py

@@ -1,45 +1,74 @@
 ## calculate_read_length_stats
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from smftools.logging_utils import get_logger
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
 # Read length QC
-def calculate_read_length_stats(adata, reference_column='', sample_names_col=''):
-    """
-    Append first valid position in a read and last valid position in the read. From this determine and append the read length. 
+def calculate_read_length_stats(
+    adata: "ad.AnnData",
+    reference_column: str = "",
+    sample_names_col: str = "",
+) -> tuple[int, int]:
+    """Calculate per-read length statistics and store them in ``adata.obs``.
+
+    Args:
+        adata: AnnData object.
+        reference_column: Obs column containing reference identifiers.
+        sample_names_col: Obs column containing sample identifiers.
 
-    Parameters:
-        adata (AnnData): An adata object
-        reference_column (str): String representing the name of the Reference column to use
-        sample_names_col (str): String representing the name of the sample name column to use
-    
     Returns:
-        upper_bound (int): last valid position in the dataset
-        lower_bound (int): first valid position in the dataset
+        tuple[int, int]: ``(upper_bound, lower_bound)`` for valid positions in the dataset.
     """
     import numpy as np
-    import anndata as ad
     import pandas as pd
 
-    print('Calculating read length statistics')
+    logger.info("Calculating read length statistics")
 
     references = set(adata.obs[reference_column])
     sample_names = set(adata.obs[sample_names_col])
 
     ## Add basic observation-level (read-level) metadata to the object: first valid position in a read and last valid position in the read. From this determine the read length. Save two new variable which hold the first and last valid positions in the entire dataset
-    print('calculating read length stats')
+    logger.info("Calculating read length stats")
     # Add some basic observation-level (read-level) metadata to the anndata object
-    read_first_valid_position = np.array([int(adata.var_names[i]) for i in np.argmax(~np.isnan(adata.X), axis=1)])
-    read_last_valid_position = np.array([int(adata.var_names[i]) for i in (adata.X.shape[1] - 1 - np.argmax(~np.isnan(adata.X[:, ::-1]), axis=1))])
-    read_length = read_last_valid_position - read_first_valid_position + np.ones(len(read_first_valid_position))
+    read_first_valid_position = np.array(
+        [int(adata.var_names[i]) for i in np.argmax(~np.isnan(adata.X), axis=1)]
+    )
+    read_last_valid_position = np.array(
+        [
+            int(adata.var_names[i])
+            for i in (adata.X.shape[1] - 1 - np.argmax(~np.isnan(adata.X[:, ::-1]), axis=1))
+        ]
+    )
+    read_length = (
+        read_last_valid_position
+        - read_first_valid_position
+        + np.ones(len(read_first_valid_position))
+    )
 
-    adata.obs['first_valid_position'] = pd.Series(read_first_valid_position, index=adata.obs.index, dtype=int)
-    adata.obs['last_valid_position'] = pd.Series(read_last_valid_position, index=adata.obs.index, dtype=int)
-    adata.obs['read_length'] = pd.Series(read_length, index=adata.obs.index, dtype=int)
+    adata.obs["first_valid_position"] = pd.Series(
+        read_first_valid_position, index=adata.obs.index, dtype=int
+    )
+    adata.obs["last_valid_position"] = pd.Series(
+        read_last_valid_position, index=adata.obs.index, dtype=int
+    )
+    adata.obs["read_length"] = pd.Series(read_length, index=adata.obs.index, dtype=int)
 
     # Define variables to hold the first and last valid position in the dataset
-    upper_bound = int(np.nanmax(adata.obs['last_valid_position']))
-    lower_bound = int(np.nanmin(adata.obs['first_valid_position']))
+    upper_bound = int(np.nanmax(adata.obs["last_valid_position"]))
+    lower_bound = int(np.nanmin(adata.obs["first_valid_position"]))
 
     return upper_bound, lower_bound
 
+
 # # Add an unstructured element to the anndata object which points to a dictionary of read lengths keyed by reference and sample name. Points to a tuple containing (mean, median, stdev) of the read lengths of the sample for the given reference strand
 #     ## Plot histogram of read length data and save the median and stdev of the read lengths for each sample.
 #     adata.uns['read_length_dict'] = {}
@@ -70,10 +99,10 @@ def calculate_read_length_stats(adata, reference_column='', sample_names_col='')
 #                 # Add a vertical line at the median
 #                 plt.axvline(median, color='red', linestyle='dashed', linewidth=1)
 #                 # Annotate the median
-#                 plt.xlim(lower_bound - 100, upper_bound + 100) 
+#                 plt.xlim(lower_bound - 100, upper_bound + 100)
 #                 if save_read_length_histogram:
 #                     save_name = output_directory + f'/{readwrite.date_string()} {title}'
 #                     plt.savefig(save_name, bbox_inches='tight', pad_inches=0.1)
 #                     plt.close()
 #                 else:
-#                     plt.show()
+#                     plt.show()

smftools/preprocessing/calculate_read_modification_stats.py

@@ -1,62 +1,92 @@
-def calculate_read_modification_stats(adata, 
-                                      reference_column, 
-                                      sample_names_col, 
-                                      mod_target_bases,
-                                      uns_flag="calculate_read_modification_stats_performed",
-                                      bypass=False,
-                                      force_redo=False
-):
-    """
-    Adds methylation/deamination statistics for each read. 
-    Indicates the read GpC and CpG methylation ratio to other_C methylation (background false positive metric for Cytosine MTase SMF).
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from smftools.logging_utils import get_logger
+
+if TYPE_CHECKING:
+    import anndata as ad
 
-    Parameters:
-        adata (AnnData): An adata object
-        reference_column (str): String representing the name of the Reference column to use
-        sample_names_col (str): String representing the name of the sample name column to use
-        mod_target_bases:
+logger = get_logger(__name__)
 
-    Returns:
-        None 
+
+def calculate_read_modification_stats(
+    adata: "ad.AnnData",
+    reference_column: str,
+    sample_names_col: str,
+    mod_target_bases: list[str],
+    uns_flag: str = "calculate_read_modification_stats_performed",
+    bypass: bool = False,
+    force_redo: bool = False,
+    valid_sites_only: bool = False,
+    valid_site_suffix: str = "_valid_coverage",
+) -> None:
+    """Add methylation/deamination statistics for each read.
+
+    Args:
+        adata: AnnData object.
+        reference_column: Obs column containing reference identifiers.
+        sample_names_col: Obs column containing sample identifiers.
+        mod_target_bases: List of target base contexts (e.g., ``["GpC", "CpG"]``).
+        uns_flag: Flag in ``adata.uns`` indicating prior completion.
+        bypass: Whether to skip processing.
+        force_redo: Whether to rerun even if ``uns_flag`` is set.
+        valid_sites_only: Whether to restrict to valid coverage sites.
+        valid_site_suffix: Suffix used for valid-site matrices.
     """
     import numpy as np
-    import anndata as ad
     import pandas as pd
 
+    if valid_sites_only:
+        if adata.uns.get("calculate_coverage_performed", False):
+            pass
+        else:
+            valid_sites_only = False
+
+    if not valid_sites_only:
+        valid_site_suffix = ""
+
     # Only run if not already performed
     already = bool(adata.uns.get(uns_flag, False))
     if (already and not force_redo) or bypass:
         # QC already performed; nothing to do
         return
 
-    print('Calculating read level Modification statistics')
+    logger.info("Calculating read level Modification statistics")
 
     references = set(adata.obs[reference_column])
     sample_names = set(adata.obs[sample_names_col])
     site_types = []
 
-    if any(base in mod_target_bases for base in ['GpC', 'CpG', 'C']):
-        site_types += ['GpC_site', 'CpG_site', 'ambiguous_GpC_CpG_site', 'other_C_site', 'C_site']
-    
-    if 'A' in mod_target_bases:
-        site_types += ['A_site']
+    if any(base in mod_target_bases for base in ["GpC", "CpG", "C"]):
+        site_types += ["GpC_site", "CpG_site", "ambiguous_GpC_CpG_site", "other_C_site", "C_site"]
 
-    for site_type in site_types:
-        adata.obs[f'Modified_{site_type}_count'] = pd.Series(0, index=adata.obs_names, dtype=int)
-        adata.obs[f'Total_{site_type}_in_read'] = pd.Series(0, index=adata.obs_names, dtype=int)
-        adata.obs[f'Fraction_{site_type}_modified'] = pd.Series(np.nan, index=adata.obs_names, dtype=float)
-        adata.obs[f'Total_{site_type}_in_reference'] = pd.Series(np.nan, index=adata.obs_names, dtype=int)
-        adata.obs[f'Valid_{site_type}_in_read_vs_reference'] = pd.Series(np.nan, index=adata.obs_names, dtype=float)
+    if "A" in mod_target_bases:
+        site_types += ["A_site"]
 
+    for site_type in site_types:
+        adata.obs[f"Modified_{site_type}_count"] = pd.Series(0, index=adata.obs_names, dtype=int)
+        adata.obs[f"Total_{site_type}_in_read"] = pd.Series(0, index=adata.obs_names, dtype=int)
+        adata.obs[f"Fraction_{site_type}_modified"] = pd.Series(
+            np.nan, index=adata.obs_names, dtype=float
+        )
+        adata.obs[f"Total_{site_type}_in_reference"] = pd.Series(
+            np.nan, index=adata.obs_names, dtype=int
+        )
+        adata.obs[f"Valid_{site_type}_in_read_vs_reference"] = pd.Series(
+            np.nan, index=adata.obs_names, dtype=float
+        )
 
     for ref in references:
         ref_subset = adata[adata.obs[reference_column] == ref]
         for site_type in site_types:
-            print(f'Iterating over {ref}_{site_type}')
-            observation_matrix = ref_subset.obsm[f'{ref}_{site_type}']
+            logger.info("Iterating over %s_%s", ref, site_type)
+            observation_matrix = ref_subset.obsm[f"{ref}_{site_type}{valid_site_suffix}"]
             total_positions_in_read = np.nansum(~np.isnan(observation_matrix), axis=1)
             total_positions_in_reference = observation_matrix.shape[1]
-            fraction_valid_positions_in_read_vs_ref = total_positions_in_read / total_positions_in_reference
+            fraction_valid_positions_in_read_vs_ref = (
+                total_positions_in_read / total_positions_in_reference
+            )
             number_mods_in_read = np.nansum(observation_matrix, axis=1)
             fraction_modified = number_mods_in_read / total_positions_in_read
 
@@ -64,38 +94,42 @@ def calculate_read_modification_stats(adata,
                 number_mods_in_read,
                 total_positions_in_read,
                 out=np.full_like(number_mods_in_read, np.nan, dtype=float),
-                where=total_positions_in_read != 0
+                where=total_positions_in_read != 0,
+            )
+
+            temp_obs_data = pd.DataFrame(
+                {
+                    f"Total_{site_type}_in_read": total_positions_in_read,
+                    f"Modified_{site_type}_count": number_mods_in_read,
+                    f"Fraction_{site_type}_modified": fraction_modified,
+                    f"Total_{site_type}_in_reference": total_positions_in_reference,
+                    f"Valid_{site_type}_in_read_vs_reference": fraction_valid_positions_in_read_vs_ref,
+                },
+                index=ref_subset.obs.index,
             )
 
-            temp_obs_data = pd.DataFrame({f'Total_{site_type}_in_read': total_positions_in_read,
-                                        f'Modified_{site_type}_count': number_mods_in_read,
-                                        f'Fraction_{site_type}_modified': fraction_modified,
-                                        f'Total_{site_type}_in_reference': total_positions_in_reference,
-                                        f'Valid_{site_type}_in_read_vs_reference': fraction_valid_positions_in_read_vs_ref}, 
-                                        index=ref_subset.obs.index)
-            
             adata.obs.update(temp_obs_data)
 
-    if any(base in mod_target_bases for base in ['GpC', 'CpG', 'C']):
-        with np.errstate(divide='ignore', invalid='ignore'):
+    if any(base in mod_target_bases for base in ["GpC", "CpG", "C"]):
+        with np.errstate(divide="ignore", invalid="ignore"):
             gpc_to_c_ratio = np.divide(
-                adata.obs[f'Fraction_GpC_site_modified'],
-                adata.obs[f'Fraction_other_C_site_modified'],
-                out=np.full_like(adata.obs[f'Fraction_GpC_site_modified'], np.nan, dtype=float),
-                where=adata.obs[f'Fraction_other_C_site_modified'] != 0
+                adata.obs["Fraction_GpC_site_modified"],
+                adata.obs["Fraction_other_C_site_modified"],
+                out=np.full_like(adata.obs["Fraction_GpC_site_modified"], np.nan, dtype=float),
+                where=adata.obs["Fraction_other_C_site_modified"] != 0,
             )
 
             cpg_to_c_ratio = np.divide(
-                adata.obs[f'Fraction_CpG_site_modified'],
-                adata.obs[f'Fraction_other_C_site_modified'],
-                out=np.full_like(adata.obs[f'Fraction_CpG_site_modified'], np.nan, dtype=float),
-                where=adata.obs[f'Fraction_other_C_site_modified'] != 0
-                )
-            
-        adata.obs['GpC_to_other_C_mod_ratio'] = gpc_to_c_ratio
-        adata.obs['CpG_to_other_C_mod_ratio'] = cpg_to_c_ratio
+                adata.obs["Fraction_CpG_site_modified"],
+                adata.obs["Fraction_other_C_site_modified"],
+                out=np.full_like(adata.obs["Fraction_CpG_site_modified"], np.nan, dtype=float),
+                where=adata.obs["Fraction_other_C_site_modified"] != 0,
+            )
+
+        adata.obs["GpC_to_other_C_mod_ratio"] = gpc_to_c_ratio
+        adata.obs["CpG_to_other_C_mod_ratio"] = cpg_to_c_ratio
 
     # mark as done
     adata.uns[uns_flag] = True
 
-    return
+    return

smftools/preprocessing/clean_NaN.py

@@ -1,23 +1,33 @@
-def clean_NaN(adata, 
-            layer=None,
-            uns_flag='clean_NaN_performed',
-            bypass=False,
-            force_redo=True
-):
-    """
-    Append layers to adata that contain NaN cleaning strategies.
+from __future__ import annotations
 
-    Parameters:
-        adata (AnnData): an anndata object
-        layer (str, optional): Name of the layer to fill NaN values in. If None, uses adata.X.
+from typing import TYPE_CHECKING
 
-    Modifies:
-        - Adds new layers to `adata.layers` with different NaN-filling strategies.
-    """
-    import numpy as np
-    import pandas as pd
+from smftools.logging_utils import get_logger
+
+if TYPE_CHECKING:
     import anndata as ad
-    from ..readwrite import adata_to_df 
+
+logger = get_logger(__name__)
+
+
+def clean_NaN(
+    adata: "ad.AnnData",
+    layer: str | None = None,
+    uns_flag: str = "clean_NaN_performed",
+    bypass: bool = False,
+    force_redo: bool = True,
+) -> None:
+    """Append layers to ``adata`` that contain NaN-cleaning strategies.
+
+    Args:
+        adata: AnnData object.
+        layer: Layer to fill NaN values in. If ``None``, uses ``adata.X``.
+        uns_flag: Flag in ``adata.uns`` indicating prior completion.
+        bypass: Whether to skip processing.
+        force_redo: Whether to rerun even if ``uns_flag`` is set.
+    """
+
+    from ..readwrite import adata_to_df
 
     # Only run if not already performed
     already = bool(adata.uns.get(uns_flag, False))
@@ -33,30 +43,30 @@ def clean_NaN(adata,
     df = adata_to_df(adata, layer=layer)
 
     # Fill NaN with closest SMF value (forward then backward fill)
-    print('Making layer: fill_nans_closest')
-    adata.layers['fill_nans_closest'] = df.ffill(axis=1).bfill(axis=1).values
+    logger.info("Making layer: fill_nans_closest")
+    adata.layers["fill_nans_closest"] = df.ffill(axis=1).bfill(axis=1).values
 
     # Replace NaN with 0, and 0 with -1
-    print('Making layer: nan0_0minus1')
+    logger.info("Making layer: nan0_0minus1")
     df_nan0_0minus1 = df.replace(0, -1).fillna(0)
-    adata.layers['nan0_0minus1'] = df_nan0_0minus1.values
+    adata.layers["nan0_0minus1"] = df_nan0_0minus1.values
 
     # Replace NaN with 1, and 1 with 2
-    print('Making layer: nan1_12')
+    logger.info("Making layer: nan1_12")
     df_nan1_12 = df.replace(1, 2).fillna(1)
-    adata.layers['nan1_12'] = df_nan1_12.values
+    adata.layers["nan1_12"] = df_nan1_12.values
 
     # Replace NaN with -1
-    print('Making layer: nan_minus_1')
+    logger.info("Making layer: nan_minus_1")
     df_nan_minus_1 = df.fillna(-1)
-    adata.layers['nan_minus_1'] = df_nan_minus_1.values
+    adata.layers["nan_minus_1"] = df_nan_minus_1.values
 
     # Replace NaN with -1
-    print('Making layer: nan_half')
+    logger.info("Making layer: nan_half")
     df_nan_half = df.fillna(0.5)
-    adata.layers['nan_half'] = df_nan_half.values
+    adata.layers["nan_half"] = df_nan_half.values
 
     # mark as done
     adata.uns[uns_flag] = True
 
-    return None
+    return None