smftools 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

smftools/preprocessing/clean_NaN.py

@@ -0,0 +1,46 @@
+def clean_NaN(adata, layer=None):
+    """
+    Append layers to adata that contain NaN cleaning strategies.
+
+    Parameters:
+        adata (AnnData): an anndata object
+        layer (str, optional): Name of the layer to fill NaN values in. If None, uses adata.X.
+
+    Modifies:
+        - Adds new layers to `adata.layers` with different NaN-filling strategies.
+    """
+    import numpy as np
+    import pandas as pd
+    import anndata as ad
+    from ..readwrite import adata_to_df 
+
+    # Ensure the specified layer exists
+    if layer and layer not in adata.layers:
+        raise ValueError(f"Layer '{layer}' not found in adata.layers.")
+
+    # Convert to DataFrame
+    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
+
+    # Replace NaN with 0, and 0 with -1
+    print('Making layer: nan0_0minus1')
+    df_nan0_0minus1 = df.replace(0, -1).fillna(0)
+    adata.layers['nan0_0minus1'] = df_nan0_0minus1.values
+
+    # Replace NaN with 1, and 1 with 2
+    print('Making layer: nan1_12')
+    df_nan1_12 = df.replace(1, 2).fillna(1)
+    adata.layers['nan1_12'] = df_nan1_12.values
+
+    # Replace NaN with -1
+    print('Making layer: nan_minus_1')
+    df_nan_minus_1 = df.fillna(-1)
+    adata.layers['nan_minus_1'] = df_nan_minus_1.values
+
+    # Replace NaN with -1
+    print('Making layer: nan_half')
+    df_nan_half = df.fillna(0.5)
+    adata.layers['nan_half'] = df_nan_half.values

smftools/preprocessing/filter_adata_by_nan_proportion.py

@@ -0,0 +1,31 @@
+## filter_adata_by_nan_proportion
+
+def filter_adata_by_nan_proportion(adata, threshold, axis='obs'):
+    """
+    Filters an anndata object on a nan proportion threshold in a given matrix axis.
+
+    Parameters:
+        adata (AnnData):
+        threshold (float): The max np.nan content to allow in the given axis.
+        axis (str): Whether to filter the adata based on obs or var np.nan content
+    Returns:
+        filtered_adata
+    """
+    import numpy as np
+    import anndata as ad
+
+    if axis == 'obs':
+        # Calculate the proportion of NaN values in each read
+        nan_proportion = np.isnan(adata.X).mean(axis=1)
+        # Filter reads to keep reads with less than a certain NaN proportion
+        filtered_indices = np.where(nan_proportion <= threshold)[0]
+        filtered_adata = adata[filtered_indices, :].copy()
+    elif axis == 'var':
+        # Calculate the proportion of NaN values at a given position
+        nan_proportion = np.isnan(adata.X).mean(axis=0)
+        # Filter positions to keep positions with less than a certain NaN proportion
+        filtered_indices = np.where(nan_proportion <= threshold)[0]
+        filtered_adata = adata[:, filtered_indices].copy()
+    else:
+        raise ValueError("Axis must be either 'obs' or 'var'")
+    return filtered_adata

smftools/preprocessing/filter_converted_reads_on_methylation.py

@@ -0,0 +1,44 @@
+## filter_converted_reads_on_methylation
+
+## Conversion SMF Specific 
+def filter_converted_reads_on_methylation(adata, valid_SMF_site_threshold=0.8, min_SMF_threshold=0.025, max_SMF_threshold=0.975):
+    """
+    Filter adata object using minimum thresholds for valid SMF site fraction in read, as well as minimum methylation content in read.
+
+    Parameters:
+        adata (AnnData): An adata object.
+        valid_SMF_site_threshold (float): A minimum proportion of valid SMF sites that must be present in the read. Default is 0.8
+        min_SMF_threshold (float): A minimum read methylation level. Default is 0.025
+    Returns:
+        Anndata
+    """
+    import numpy as np
+    import anndata as ad
+    import pandas as pd
+
+    if valid_SMF_site_threshold:
+        # Keep reads that have over a given valid GpC site content
+        adata = adata[adata.obs['fraction_valid_GpC_site_in_range'] > valid_SMF_site_threshold].copy()
+
+    if min_SMF_threshold:
+        # Keep reads with SMF methylation over background methylation.
+        below_background = (~adata.obs['GpC_above_other_C']).sum()
+        print(f'Removing {below_background} reads that have GpC conversion below background conversion rate')
+        adata = adata[adata.obs['GpC_above_other_C'] == True].copy()
+        # Keep reads over a defined methylation threshold
+        s0 = adata.shape[0]
+        adata = adata[adata.obs['GpC_site_row_methylation_means'] > min_SMF_threshold].copy()
+        s1 = adata.shape[0]
+        below_threshold = s0 - s1
+        print(f'Removing {below_threshold} reads that have GpC conversion below a minimum threshold conversion rate')
+
+    if max_SMF_threshold:
+        # Keep reads below a defined methylation threshold
+        s0 = adata.shape[0]
+        adata = adata[adata.obs['GpC_site_row_methylation_means'] < max_SMF_threshold].copy()
+        s1 = adata.shape[0]
+        above_threshold = s0 - s1
+        print(f'Removing {above_threshold} reads that have GpC conversion above a maximum threshold conversion rate')
+
+    return adata
+

smftools/preprocessing/filter_reads_on_length.py

@@ -0,0 +1,51 @@
+## filter_reads_on_length
+
+def filter_reads_on_length(adata, filter_on_coordinates=False, min_read_length=2700, max_read_length=3200):
+    """
+    Filters the adata object to keep a defined coordinate window, as well as reads that are over a minimum threshold in length.
+
+    Parameters:
+        adata (AnnData): An adata object.
+        filter_on_coordinates (bool | list): If False, skips filtering. Otherwise, provide a list containing integers representing the lower and upper bound coordinates to filter on. Default is False.
+        min_read_length (int): The minimum read length to keep in the filtered dataset. Default is 2700.
+        max_read_length (int): The maximum query read length to keep in the filtered dataset. Default is 3200.
+
+    Returns:
+        adata
+    """
+    import numpy as np
+    import anndata as ad
+    import pandas as pd
+
+    if filter_on_coordinates:
+        lower_bound, upper_bound = filter_on_coordinates
+        # Extract the position information from the adata object as an np array
+        var_names_arr = adata.var_names.astype(int).to_numpy()
+        # Find the upper bound coordinate that is closest to the specified value
+        closest_end_index = np.argmin(np.abs(var_names_arr - upper_bound))
+        upper_bound = int(adata.var_names[closest_end_index])
+        # Find the lower bound coordinate that is closest to the specified value
+        closest_start_index = np.argmin(np.abs(var_names_arr - lower_bound))
+        lower_bound = int(adata.var_names[closest_start_index])
+        # Get a list of positional indexes that encompass the lower and upper bounds of the dataset
+        position_list = list(range(lower_bound, upper_bound + 1))
+        position_list = [str(pos) for pos in position_list]
+        position_set = set(position_list)
+        print(f'Subsetting adata to keep data between coordinates {lower_bound} and {upper_bound}')
+        adata = adata[:, adata.var_names.isin(position_set)].copy()
+
+    if min_read_length:
+        print(f'Subsetting adata to keep reads longer than {min_read_length}')
+        s0 = adata.shape[0]
+        adata = adata[adata.obs['read_length'] > min_read_length].copy()
+        s1 = adata.shape[0]
+        print(f'Removed {s0-s1} reads less than {min_read_length} basepairs in length')
+
+    if max_read_length:
+        print(f'Subsetting adata to keep reads shorter than {max_read_length}')
+        s0 = adata.shape[0]
+        adata = adata[adata.obs['read_length'] < max_read_length].copy()
+        s1 = adata.shape[0]
+        print(f'Removed {s0-s1} reads greater than {max_read_length} basepairs in length')
+
+    return adata

smftools/preprocessing/flag_duplicate_reads.py

@@ -0,0 +1,149 @@
+import torch
+from tqdm import tqdm
+
+class UnionFind:
+    def __init__(self, size):
+        self.parent = torch.arange(size)
+
+    def find(self, x):
+        while self.parent[x] != x:
+            self.parent[x] = self.parent[self.parent[x]]
+            x = self.parent[x]
+        return x
+
+    def union(self, x, y):
+        root_x = self.find(x)
+        root_y = self.find(y)
+        if root_x != root_y:
+            self.parent[root_y] = root_x
+
+
+def flag_duplicate_reads(adata, var_filters_sets, distance_threshold=0.05, obs_reference_col='Reference_strand'):
+    import numpy as np
+    import pandas as pd
+    import matplotlib.pyplot as plt
+
+    all_hamming_dists = []
+    merged_results = []
+
+    references = adata.obs[obs_reference_col].cat.categories
+
+    for ref in references:
+        print(f'🔹 Processing reference: {ref}')
+
+        ref_mask = adata.obs[obs_reference_col] == ref
+        adata_subset = adata[ref_mask].copy()
+        N = adata_subset.shape[0]
+
+        combined_mask = torch.zeros(len(adata.var), dtype=torch.bool)
+        for var_set in var_filters_sets:
+            if any(ref in v for v in var_set):
+                set_mask = torch.ones(len(adata.var), dtype=torch.bool)
+                for key in var_set:
+                    set_mask &= torch.from_numpy(adata.var[key].values)
+                combined_mask |= set_mask
+
+        selected_cols = adata.var.index[combined_mask.numpy()].to_list()
+        col_indices = [adata.var.index.get_loc(col) for col in selected_cols]
+
+        print(f"Selected {len(col_indices)} columns out of {adata.var.shape[0]} for {ref}")
+
+        X = adata_subset.X
+        if not isinstance(X, np.ndarray):
+            X = X.toarray()
+        X_subset = X[:, col_indices]
+        X_tensor = torch.from_numpy(X_subset.astype(np.float32))
+
+        fwd_hamming_to_next = torch.full((N,), float('nan'))
+        rev_hamming_to_prev = torch.full((N,), float('nan'))
+
+        def cluster_pass(X_tensor, reverse=False, window_size=50, record_distances=False):
+            N_local = X_tensor.shape[0]
+            X_sortable = X_tensor.nan_to_num(-1)
+            sort_keys = X_sortable.tolist()
+            sorted_idx = sorted(range(N_local), key=lambda i: sort_keys[i], reverse=reverse)
+            sorted_X = X_tensor[sorted_idx]
+
+            cluster_pairs = []
+
+            for i in tqdm(range(len(sorted_X)), desc=f"Pass {'rev' if reverse else 'fwd'} ({ref})"):
+                row_i = sorted_X[i]
+                j_range = range(i + 1, min(i + 1 + window_size, len(sorted_X)))
+
+                if len(j_range) > 0:
+                    row_i_exp = row_i.unsqueeze(0)
+                    block_rows = sorted_X[j_range]
+                    valid_mask = (~torch.isnan(row_i_exp)) & (~torch.isnan(block_rows))
+                    valid_counts = valid_mask.sum(dim=1)
+                    diffs = (row_i_exp != block_rows) & valid_mask
+                    hamming_dists = diffs.sum(dim=1) / valid_counts.clamp(min=1)
+                    all_hamming_dists.extend(hamming_dists.cpu().numpy().tolist())
+
+                    matches = (hamming_dists < distance_threshold) & (valid_counts > 0)
+                    for offset_idx, m in zip(j_range, matches):
+                        if m:
+                            cluster_pairs.append((sorted_idx[i], sorted_idx[offset_idx]))
+
+                    if record_distances and i + 1 < len(sorted_X):
+                        next_idx = sorted_idx[i + 1]
+                        valid_mask_pair = (~torch.isnan(row_i)) & (~torch.isnan(sorted_X[i + 1]))
+                        if valid_mask_pair.sum() > 0:
+                            d = (row_i[valid_mask_pair] != sorted_X[i + 1][valid_mask_pair]).sum()
+                            norm_d = d.item() / valid_mask_pair.sum().item()
+                            if reverse:
+                                rev_hamming_to_prev[next_idx] = norm_d
+                            else:
+                                fwd_hamming_to_next[sorted_idx[i]] = norm_d
+
+            return cluster_pairs
+
+        pairs_fwd = cluster_pass(X_tensor, reverse=False, record_distances=True)
+        involved_in_fwd = set([p[0] for p in pairs_fwd] + [p[1] for p in pairs_fwd])
+        mask_for_rev = torch.ones(N, dtype=torch.bool)
+        mask_for_rev[list(involved_in_fwd)] = False
+        pairs_rev = cluster_pass(X_tensor[mask_for_rev], reverse=True, record_distances=True)
+
+        all_pairs = pairs_fwd + [(list(mask_for_rev.nonzero(as_tuple=True)[0])[i], list(mask_for_rev.nonzero(as_tuple=True)[0])[j]) for i, j in pairs_rev]
+
+        uf = UnionFind(N)
+        for i, j in all_pairs:
+            uf.union(i, j)
+
+        merged_cluster = torch.zeros(N, dtype=torch.long)
+        for i in range(N):
+            merged_cluster[i] = uf.find(i)
+
+        cluster_sizes = torch.zeros_like(merged_cluster)
+        for cid in merged_cluster.unique():
+            members = (merged_cluster == cid).nonzero(as_tuple=True)[0]
+            cluster_sizes[members] = len(members)
+
+        is_duplicate = torch.zeros(N, dtype=torch.bool)
+        for cid in merged_cluster.unique():
+            members = (merged_cluster == cid).nonzero(as_tuple=True)[0]
+            if len(members) > 1:
+                is_duplicate[members[1:]] = True
+
+        adata_subset.obs['is_duplicate'] = is_duplicate.numpy()
+        adata_subset.obs['merged_cluster_id'] = merged_cluster.numpy()
+        adata_subset.obs['cluster_size'] = cluster_sizes.numpy()
+        adata_subset.obs['fwd_hamming_to_next'] = fwd_hamming_to_next.numpy()
+        adata_subset.obs['rev_hamming_to_prev'] = rev_hamming_to_prev.numpy()
+
+        merged_results.append(adata_subset.obs)
+
+    merged_obs = pd.concat(merged_results)
+    adata.obs = adata.obs.join(merged_obs[['is_duplicate', 'merged_cluster_id', 'cluster_size', 'fwd_hamming_to_next', 'rev_hamming_to_prev']])
+
+    adata_unique = adata[~adata.obs['is_duplicate']].copy()
+
+    plt.figure(figsize=(5, 4))
+    plt.hist(all_hamming_dists, bins=50, alpha=0.75)
+    plt.axvline(distance_threshold, color="red", linestyle="--", label=f"threshold = {distance_threshold}")
+    plt.xlabel("Hamming Distance")
+    plt.ylabel("Frequency")
+    plt.title("Histogram of Pairwise Hamming Distances")
+    plt.legend()
+    plt.show()
+
+    return adata_unique, adata

smftools/preprocessing/invert_adata.py

@@ -0,0 +1,30 @@
+## invert_adata
+
+# Optional inversion of the adata
+
+def invert_adata(adata):
+    """
+    Inverts the AnnData object along the column (variable) axis.
+
+    Parameters:
+        adata (AnnData): An AnnData object.
+
+    Returns:
+        AnnData: A new AnnData object with inverted column ordering.
+    """
+    import numpy as np
+    import anndata as ad
+
+    print("🔄 Inverting AnnData along the column axis...")
+
+    # Reverse the order of columns (variables)
+    inverted_adata = adata[:, ::-1].copy()
+
+    # Reassign var_names with new order
+    inverted_adata.var_names = adata.var_names
+
+    # Optional: Store original coordinates for reference
+    inverted_adata.var["Original_var_names"] = adata.var_names[::-1]
+
+    print("✅ Inversion complete!")
+    return inverted_adata

smftools/preprocessing/load_sample_sheet.py

@@ -0,0 +1,38 @@
+def load_sample_sheet(adata, sample_sheet_path, mapping_key_column='obs_names', as_category=True):
+    """
+    Loads a sample sheet CSV and maps metadata into the AnnData object as categorical columns.
+
+    Parameters:
+        adata (AnnData): The AnnData object to append sample information to.
+        sample_sheet_path (str): Path to the CSV file.
+        mapping_key_column (str): Column name in the CSV to map against adata.obs_names or an existing obs column.
+        as_category (bool): If True, added columns will be cast as pandas Categorical.
+
+    Returns:
+        AnnData: Updated AnnData object.
+    """
+    import pandas as pd
+
+    print('🔹 Loading sample sheet...')
+    df = pd.read_csv(sample_sheet_path)
+    df[mapping_key_column] = df[mapping_key_column].astype(str)
+    
+    # If matching against obs_names directly
+    if mapping_key_column == 'obs_names':
+        key_series = adata.obs_names.astype(str)
+    else:
+        key_series = adata.obs[mapping_key_column].astype(str)
+
+    value_columns = [col for col in df.columns if col != mapping_key_column]
+    
+    print(f'🔹 Appending metadata columns: {value_columns}')
+    df = df.set_index(mapping_key_column)
+
+    for col in value_columns:
+        mapped = key_series.map(df[col])
+        if as_category:
+            mapped = mapped.astype('category')
+        adata.obs[col] = mapped
+
+    print('✅ Sample sheet metadata successfully added as categories.' if as_category else '✅ Metadata added.')
+    return adata

smftools/preprocessing/make_dirs.py

@@ -0,0 +1,21 @@
+## make_dirs
+
+# General
+def make_dirs(directories):
+    """
+    Takes a list of file paths and makes new directories if the directory does not already exist.
+
+    Parameters:
+        directories (list): A list of directories to make
+    
+    Returns:
+        None
+    """
+    import os
+
+    for directory in directories:
+        if not os.path.isdir(directory):
+            os.mkdir(directory)
+            print(f"Directory '{directory}' created successfully.")
+        else:
+            print(f"Directory '{directory}' already exists.")

smftools/preprocessing/min_non_diagonal.py

@@ -0,0 +1,25 @@
+## min_non_diagonal
+
+def min_non_diagonal(matrix):
+    """
+    Takes a matrix and returns the smallest value from each row with the diagonal masked.
+
+    Parameters:
+        matrix (ndarray): A 2D ndarray.
+
+    Returns:
+        min_values (list): A list of minimum values from each row of the matrix
+    """
+    import numpy as np
+
+    n = matrix.shape[0]
+    min_values = []
+    for i in range(n):
+        # Mask to exclude the diagonal element
+        row_mask = np.ones(n, dtype=bool)
+        row_mask[i] = False
+        # Extract the row excluding the diagonal element
+        row = matrix[i, row_mask]
+        # Find the minimum value in the row
+        min_values.append(np.min(row))
+    return min_values

smftools/preprocessing/recipes.py

@@ -0,0 +1,127 @@
+# recipes
+
+def recipe_1_Kissiov_and_McKenna_2025(adata, sample_sheet_path, output_directory, mapping_key_column='Sample', reference_column = 'Reference', sample_names_col='Sample_names', invert=True):
+    """
+    The first part of the preprocessing workflow applied to the smf.inform.pod_to_adata() output derived from Kissiov_and_McKenna_2025.
+
+    Performs the following tasks:
+    1) Loads a sample CSV to append metadata mappings to the adata object.
+    2) Appends a boolean indicating whether each position in var_names is within a given reference.
+    3) Appends the cytosine context to each position from each reference.
+    4) Calculate read level methylation statistics.
+    5) Calculates read length statistics (start position, end position, read length).
+    6) Optionally inverts the adata to flip the position coordinate orientation.
+    7) Adds new layers containing NaN replaced variants of adata.X (fill_closest, nan0_0minus1, nan1_12).
+    8) Returns a dictionary to pass the variable namespace to the parent scope.
+
+    Parameters:
+        adata (AnnData): The AnnData object to use as input.
+        sample_sheet_path (str): String representing the path to the sample sheet csv containing the sample metadata.
+        output_directory (str): String representing the path to the output directory for plots.
+        mapping_key_column (str): The column name to use as the mapping keys for applying the sample sheet metadata.
+        reference_column (str): The name of the reference column to use.
+        sample_names_col (str): The name of the sample name column to use.
+        invert (bool): Whether to invert the positional coordinates of the adata object.
+
+    Returns:
+        variables (dict): A dictionary of variables to append to the parent scope.
+    """
+    import anndata as ad
+    import pandas as pd
+    import numpy as np
+    from .load_sample_sheet import load_sample_sheet
+    from .calculate_coverage import calculate_coverage
+    from .append_C_context import append_C_context
+    from .calculate_converted_read_methylation_stats import calculate_converted_read_methylation_stats
+    from .invert_adata import invert_adata
+    from .calculate_read_length_stats import calculate_read_length_stats
+    from .clean_NaN import clean_NaN
+
+    # Clean up some of the Reference metadata and save variable names that point to sets of values in the column.
+    adata.obs[reference_column] = adata.obs[reference_column].astype('category')
+    references = adata.obs[reference_column].cat.categories
+    split_references = [(reference, reference.split('_')[0][1:]) for reference in references]
+    reference_mapping = {k: v for k, v in split_references}
+    adata.obs[f'{reference_column}_short'] = adata.obs[reference_column].map(reference_mapping)
+    short_references = set(adata.obs[f'{reference_column}_short'])
+    binary_layers = list(adata.layers.keys())
+
+    # load sample sheet metadata
+    load_sample_sheet(adata, sample_sheet_path, mapping_key_column)
+
+    # hold sample names set
+    adata.obs[sample_names_col] = adata.obs[sample_names_col].astype('category')
+    sample_names = adata.obs[sample_names_col].cat.categories
+
+    # Add position level metadata
+    calculate_coverage(adata, obs_column=reference_column)
+    adata.var['SNP_position'] = (adata.var[f'N_{reference_column}_with_position'] > 0) & (adata.var[f'N_{reference_column}_with_position'] < len(references)).astype(bool)
+
+    # Append cytosine context to the reference positions based on the conversion strand.
+    append_C_context(adata, obs_column=reference_column, use_consensus=False)
+
+    # Calculate read level methylation statistics. Assess if GpC methylation level is above other_C methylation level as a QC.
+    calculate_converted_read_methylation_stats(adata, reference_column, sample_names_col)
+
+    # Calculate read length statistics
+    upper_bound, lower_bound = calculate_read_length_stats(adata, reference_column, sample_names_col)
+
+    # Invert the adata object (ie flip the strand orientation for visualization)
+    if invert:
+        adata = invert_adata(adata)
+    else:
+        pass
+
+    # NaN replacement strategies stored in additional layers. Having layer=None uses adata.X
+    clean_NaN(adata, layer=None)
+
+    variables = {
+        "short_references": short_references,
+        "binary_layers": binary_layers,
+        "sample_names": sample_names,
+        "upper_bound": upper_bound,
+        "lower_bound": lower_bound,
+        "references": references
+    }
+    return variables
+
+def recipe_2_Kissiov_and_McKenna_2025(adata, output_directory, binary_layers, distance_thresholds={}, reference_column = 'Reference', sample_names_col='Sample_names'):
+    """
+    The second part of the preprocessing workflow applied to the adata that has already been preprocessed by recipe_1_Kissiov_and_McKenna_2025.
+
+    Performs the following tasks:
+    1) Marks putative PCR duplicates using pairwise hamming distance metrics.
+    2) Performs a complexity analysis of the library based on the PCR duplicate detection rate.
+    3) Removes PCR duplicates from the adata.
+    4) Returns two adata object: one for the filtered adata and one for the duplicate adata.
+
+    Parameters:
+        adata (AnnData): The AnnData object to use as input.
+        output_directory (str): String representing the path to the output directory for plots.
+        binary_layers (list): A list of layers to used for the binary encoding of read sequences. Used for duplicate detection.
+        distance_thresholds (dict): A dictionary keyed by obs_column categories that points to a float corresponding to the distance threshold to apply. Default is an empty dict.
+        reference_column (str): The name of the reference column to use.
+        sample_names_col (str): The name of the sample name column to use.
+
+    Returns:
+        filtered_adata (AnnData): An AnnData object containing the filtered reads
+        duplicates (AnnData): An AnnData object containing the duplicate reads
+    """
+    import anndata as ad
+    import pandas as pd
+    import numpy as np
+    from .mark_duplicates import mark_duplicates
+    from .calculate_complexity import calculate_complexity
+    from .remove_duplicates import remove_duplicates
+
+    # Add here a way to remove reads below a given read quality (based on nan content). Need to also add a way to pull from BAM files the read quality from each read
+
+    # Duplicate detection using pairwise hamming distance across reads
+    mark_duplicates(adata, binary_layers, obs_column=reference_column, sample_col=sample_names_col, distance_thresholds=distance_thresholds, method='N_masked_distances')
+
+    # Complexity analysis using the marked duplicates and the lander-watermann algorithm
+    calculate_complexity(adata, output_directory, obs_column=reference_column, sample_col=sample_names_col, plot=True, save_plot=False)
+
+    # Remove duplicate reads and store the duplicate reads in a new AnnData object named duplicates.
+    filtered_adata, duplicates = remove_duplicates(adata)
+    return filtered_adata, duplicates

smftools/preprocessing/subsample_adata.py

@@ -0,0 +1,58 @@
+def subsample_adata(adata, obs_columns=None, max_samples=2000, random_seed=42):
+    """
+    Subsamples an AnnData object so that each unique combination of categories 
+    in the given `obs_columns` has at most `max_samples` observations.
+    If `obs_columns` is None or empty, the function randomly subsamples the entire dataset.
+    
+    Parameters:
+        adata (AnnData): The AnnData object to subsample.
+        obs_columns (list of str, optional): List of observation column names to group by.
+        max_samples (int): The maximum number of observations per category combination.
+        random_seed (int): Random seed for reproducibility.
+
+    Returns:
+        AnnData: A new AnnData object with subsampled observations.
+    """
+    import anndata as ad
+    import numpy as np
+
+    np.random.seed(random_seed)  # Ensure reproducibility
+
+    if not obs_columns:  # If no obs columns are given, sample globally
+        if adata.shape[0] > max_samples:
+            sampled_indices = np.random.choice(adata.obs.index, max_samples, replace=False)
+        else:
+            sampled_indices = adata.obs.index  # Keep all if fewer than max_samples
+        
+        return adata[sampled_indices].copy()
+
+    sampled_indices = []
+
+    # Get unique category combinations from all specified obs columns
+    unique_combinations = adata.obs[obs_columns].drop_duplicates()
+
+    for _, row in unique_combinations.iterrows():
+        # Build filter condition dynamically for multiple columns
+        condition = (adata.obs[obs_columns] == row.values).all(axis=1)
+        
+        # Get indices for the current category combination
+        subset_indices = adata.obs[condition].index.to_numpy()
+
+        # Subsample or take all
+        if len(subset_indices) > max_samples:
+            sampled = np.random.choice(subset_indices, max_samples, replace=False)
+        else:
+            sampled = subset_indices  # Keep all if fewer than max_samples
+
+        sampled_indices.extend(sampled)
+
+    # ⚠ Handle backed mode detection
+    if adata.isbacked:
+        print("⚠ Detected backed mode. Subset will be loaded fully into memory.")
+        subset = adata[sampled_indices]
+        subset = subset.to_memory()
+    else:
+        subset = adata[sampled_indices]
+
+    # Create a new AnnData object with only the selected indices
+    return subset[sampled_indices].copy()