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

smftools/tools/apply_hmm_batched.py

@@ -0,0 +1,241 @@
+import numpy as np
+import pandas as pd
+import torch
+from tqdm import tqdm    
+
+def apply_hmm_batched(adata, model, obs_column, layer=None, footprints=True, accessible_patches=False, cpg=False, methbases=["GpC", "CpG", "A"], device="cpu", threshold=0.7):
+    """
+    Applies an HMM model to an AnnData object using tensor-based sequence inputs.
+    If multiple methbases are passed, generates a combined feature set.
+    """
+    import numpy as np
+    import torch
+    from tqdm import tqdm
+
+    model.to(device)
+
+    # --- Feature Definitions ---
+    feature_sets = {}
+    if footprints:
+        feature_sets["footprint"] = {
+            "features": {
+                "small_bound_stretch": [0, 20],
+                "medium_bound_stretch": [20, 50],
+                "putative_nucleosome": [50, 200],
+                "large_bound_stretch": [200, np.inf]
+            },
+            "state": "Non-Methylated"
+        }
+    if accessible_patches:
+        feature_sets["accessible"] = {
+            "features": {
+                "small_accessible_patch": [0, 20],
+                "mid_accessible_patch": [20, 80],
+                "large_accessible_patch": [80, np.inf]
+            },
+            "state": "Methylated"
+        }
+    if cpg:
+        feature_sets["cpg"] = {
+            "features": {
+                "cpg_patch": [0, np.inf]
+            },
+            "state": "Methylated"
+        }
+
+    # --- Init columns ---
+    all_features = []
+    combined_prefix = "Combined"
+    for key, fs in feature_sets.items():
+        if key == 'cpg':
+            all_features += [f"CpG_{f}" for f in fs["features"]]
+            all_features.append(f"CpG_all_{key}_features")
+        else:
+            for methbase in methbases:
+                all_features += [f"{methbase}_{f}" for f in fs["features"]]
+                all_features.append(f"{methbase}_all_{key}_features")
+            if len(methbases) > 1:
+                all_features += [f"{combined_prefix}_{f}" for f in fs["features"]]
+                all_features.append(f"{combined_prefix}_all_{key}_features")
+
+    for feature in all_features:
+        adata.obs[feature] = [[] for _ in range(adata.shape[0])]
+        adata.obs[f"{feature}_distances"] = [None] * adata.shape[0]
+        adata.obs[f"n_{feature}"] = -1
+
+    # --- Main loop ---
+    references = adata.obs[obs_column].cat.categories
+
+    for ref in tqdm(references, desc="Processing References"):
+        ref_subset = adata[adata.obs[obs_column] == ref]
+
+        # Combined methbase mask
+        combined_mask = None
+        for methbase in methbases:
+            mask = {
+                "a": ref_subset.var[f"{ref}_strand_FASTA_base"] == "A",
+                "gpc": ref_subset.var[f"{ref}_GpC_site"] == True,
+                "cpg": ref_subset.var[f"{ref}_CpG_site"] == True
+            }[methbase.lower()]
+            combined_mask = mask if combined_mask is None else combined_mask | mask
+
+            methbase_subset = ref_subset[:, mask]
+            matrix = methbase_subset.layers[layer] if layer else methbase_subset.X
+
+            processed_reads = [[int(x) if not np.isnan(x) else np.random.choice([0, 1]) for x in read] for read in matrix]
+            tensor_batch = torch.tensor(processed_reads, dtype=torch.long, device=device).unsqueeze(-1)
+
+            coords = methbase_subset.var_names
+            for key, fs in feature_sets.items():
+                if key == 'cpg':
+                    continue
+                state_target = fs["state"]
+                feature_map = fs["features"]
+
+                pred_states = model.predict(tensor_batch)
+                probs = model.predict_proba(tensor_batch)
+                classifications = classify_batch(pred_states, probs, coords, feature_map, target_state=state_target)
+
+                for i, idx in enumerate(methbase_subset.obs.index):
+                    for start, length, label, prob in classifications[i]:
+                        adata.obs.at[idx, f"{methbase}_{label}"].append([start, length, prob])
+                        adata.obs.at[idx, f"{methbase}_all_{key}_features"].append([start, length, prob])
+
+        # Combined subset
+        if len(methbases) > 1:
+            combined_subset = ref_subset[:, combined_mask]
+            combined_matrix = combined_subset.layers[layer] if layer else combined_subset.X
+            processed_combined_reads = [[int(x) if not np.isnan(x) else np.random.choice([0, 1]) for x in read] for read in combined_matrix]
+            tensor_combined_batch = torch.tensor(processed_combined_reads, dtype=torch.long, device=device).unsqueeze(-1)
+
+            coords = combined_subset.var_names
+            for key, fs in feature_sets.items():
+                if key == 'cpg':
+                    continue
+                state_target = fs["state"]
+                feature_map = fs["features"]
+
+                pred_states = model.predict(tensor_combined_batch)
+                probs = model.predict_proba(tensor_combined_batch)
+                classifications = classify_batch(pred_states, probs, coords, feature_map, target_state=state_target)
+
+                for i, idx in enumerate(combined_subset.obs.index):
+                    for start, length, label, prob in classifications[i]:
+                        adata.obs.at[idx, f"{combined_prefix}_{label}"].append([start, length, prob])
+                        adata.obs.at[idx, f"{combined_prefix}_all_{key}_features"].append([start, length, prob])
+
+    # --- Special handling for CpG ---
+    if cpg:
+        for ref in tqdm(references, desc="Processing CpG"):
+            ref_subset = adata[adata.obs[obs_column] == ref]
+            mask = (ref_subset.var[f"{ref}_CpG_site"] == True)
+            cpg_subset = ref_subset[:, mask]
+            matrix = cpg_subset.layers[layer] if layer else cpg_subset.X
+
+            processed_reads = [[int(x) if not np.isnan(x) else np.random.choice([0, 1]) for x in read] for read in matrix]
+            tensor_batch = torch.tensor(processed_reads, dtype=torch.long, device=device).unsqueeze(-1)
+
+            coords = cpg_subset.var_names
+            fs = feature_sets['cpg']
+            state_target = fs["state"]
+            feature_map = fs["features"]
+
+            pred_states = model.predict(tensor_batch)
+            probs = model.predict_proba(tensor_batch)
+            classifications = classify_batch(pred_states, probs, coords, feature_map, target_state=state_target)
+
+            for i, idx in enumerate(cpg_subset.obs.index):
+                for start, length, label, prob in classifications[i]:
+                    adata.obs.at[idx, f"CpG_{label}"].append([start, length, prob])
+                    adata.obs.at[idx, f"CpG_all_cpg_features"].append([start, length, prob])
+
+    # --- Binarization + Distance ---
+    for feature in tqdm(all_features, desc="Finalizing Layers"):
+        bin_matrix = np.zeros((adata.shape[0], adata.shape[1]), dtype=int)
+        counts = np.zeros(adata.shape[0], dtype=int)
+        for row_idx, intervals in enumerate(adata.obs[feature]):
+            if not isinstance(intervals, list):
+                intervals = []
+            for start, length, prob in intervals:
+                if prob > threshold:
+                    bin_matrix[row_idx, start:start+length] = 1
+                    counts[row_idx] += 1
+        adata.layers[f"{feature}"] = bin_matrix
+        adata.obs[f"n_{feature}"] = counts
+        adata.obs[f"{feature}_distances"] = calculate_batch_distances(adata.obs[feature].tolist(), threshold)
+
+def calculate_batch_distances(intervals_list, threshold=0.9):
+    """
+    Vectorized calculation of distances across multiple reads.
+
+    Parameters:
+        intervals_list (list of list): Outer list = reads, inner list = intervals [start, length, prob]
+        threshold (float): Minimum probability threshold for filtering
+
+    Returns:
+        List of distance lists per read.
+    """
+    results = []
+    for intervals in intervals_list:
+        if not isinstance(intervals, list) or len(intervals) == 0:
+            results.append([])
+            continue
+        valid = [iv for iv in intervals if iv[2] > threshold]
+        valid = sorted(valid, key=lambda x: x[0])
+        dists = [(valid[i + 1][0] - (valid[i][0] + valid[i][1])) for i in range(len(valid) - 1)]
+        results.append(dists)
+    return results
+
+
+
+def classify_batch(predicted_states_batch, probabilities_batch, coordinates, classification_mapping, target_state="Methylated"):
+    """
+    Classify batch sequences efficiently.
+
+    Parameters:
+        predicted_states_batch: Tensor [batch_size, seq_len]
+        probabilities_batch: Tensor [batch_size, seq_len, n_states]
+        coordinates: list of genomic coordinates
+        classification_mapping: dict of feature bins
+        target_state: state name ("Methylated" or "Non-Methylated")
+
+    Returns:
+        List of classifications for each sequence.
+    """
+    import numpy as np
+
+    state_labels = ["Non-Methylated", "Methylated"]
+    target_idx = state_labels.index(target_state)
+    batch_size = predicted_states_batch.shape[0]
+
+    all_classifications = []
+
+    for b in range(batch_size):
+        predicted_states = predicted_states_batch[b].cpu().numpy()
+        probabilities = probabilities_batch[b].cpu().numpy()
+
+        regions = []
+        current_start, current_length, current_probs = None, 0, []
+
+        for i, state_index in enumerate(predicted_states):
+            state_prob = probabilities[i][state_index]
+            if state_index == target_idx:
+                if current_start is None:
+                    current_start = i
+                current_length += 1
+                current_probs.append(state_prob)
+            elif current_start is not None:
+                regions.append((current_start, current_length, np.mean(current_probs)))
+                current_start, current_length, current_probs = None, 0, []
+
+        if current_start is not None:
+            regions.append((current_start, current_length, np.mean(current_probs)))
+
+        final = []
+        for start, length, prob in regions:
+            feature_length = int(coordinates[start + length - 1]) - int(coordinates[start]) + 1
+            label = next((ftype for ftype, rng in classification_mapping.items() if rng[0] <= feature_length < rng[1]), target_state)
+            final.append((int(coordinates[start]) + 1, feature_length, label, prob))
+        all_classifications.append(final)
+
+    return all_classifications

smftools/tools/archived/classify_methylated_features.py

@@ -0,0 +1,66 @@
+# classify_methylated_features
+
+def classify_methylated_features(read, model, coordinates, classification_mapping={}):
+    """
+    Classifies methylated features (accessible features or CpG methylated features)
+
+    Parameters:
+        read (np.ndarray) : An array of binarized SMF data representing a read
+        model (): a trained pomegranate HMM
+        coordinates (list): a list of postional coordinates corresponding to the positions in the read
+        classification_mapping (dict): A dictionary keyed by classification name that points to a 2-element list containing size boundary constraints for that feature.
+    Returns:
+        final_classifications (list): A list of tuples, where each tuple is an instance of a non-methylated feature in the read. The tuple contains: feature start, feature length, feature classification, and HMM probability.
+    """
+    import numpy as np
+
+    sequence = list(read)
+    # Get the predicted states using the MAP algorithm
+    predicted_states = model.predict(sequence, algorithm='map')
+    
+    # Get the probabilities for each state using the forward-backward algorithm
+    probabilities = model.predict_proba(sequence)
+    
+    # Initialize lists to store the classifications and their probabilities
+    classifications = []
+    current_start = None
+    current_length = 0
+    current_probs = []
+    
+    for i, state_index in enumerate(predicted_states):
+        state_name = model.states[state_index].name
+        state_prob = probabilities[i][state_index]
+        
+        if state_name == "Methylated":
+            if current_start is None:
+                current_start = i
+            current_length += 1
+            current_probs.append(state_prob)
+        else:
+            if current_start is not None:
+                avg_prob = np.mean(current_probs)
+                classifications.append((current_start, current_length, "Methylated", avg_prob))
+                current_start = None
+                current_length = 0
+                current_probs = []
+    
+    if current_start is not None:
+        avg_prob = np.mean(current_probs)
+        classifications.append((current_start, current_length, "Methylated", avg_prob))
+
+    final_classifications = []
+    for start, length, classification, prob in classifications:
+        final_classification = None
+        feature_length = int(coordinates[start + length - 1]) - int(coordinates[start]) + 1
+        for feature_type, size_range in classification_mapping.items():
+            if size_range[0] <= feature_length < size_range[1]:
+                final_classification = feature_type
+                break
+            else:
+                pass
+        if not final_classification:
+            final_classification = classification
+
+        final_classifications.append((int(coordinates[start]) + 1, feature_length, final_classification, prob))
+    
+    return final_classifications

smftools/tools/archived/classify_non_methylated_features.py

@@ -0,0 +1,75 @@
+# classify_non_methylated_features
+
+def classify_non_methylated_features(read, model, coordinates, classification_mapping={}):
+    """
+    Classifies non-methylated features (inaccessible features)
+
+    Parameters:
+        read (np.ndarray) : An array of binarized SMF data representing a read
+        model (): a trained pomegranate HMM
+        coordinates (list): a list of postional coordinates corresponding to the positions in the read
+        classification_mapping (dict): A dictionary keyed by classification name that points to a 2-element list containing size boundary constraints for that feature.
+    Returns:
+        final_classifications (list): A list of tuples, where each tuple is an instance of a non-methylated feature in the read. The tuple contains: feature start, feature length, feature classification, and HMM probability.
+    """
+    import numpy as np
+
+    sequence = list(read)
+    # Get the predicted states using the MAP algorithm
+    predicted_states = model.predict(sequence, algorithm='map')
+    
+    # Get the probabilities for each state using the forward-backward algorithm
+    probabilities = model.predict_proba(sequence)
+    
+    # Initialize lists to store the classifications and their probabilities
+    classifications = []
+    current_start = None
+    current_length = 0
+    current_probs = []
+    
+    for i, state_index in enumerate(predicted_states):
+        state_name = model.states[state_index].name
+        state_prob = probabilities[i][state_index]
+        
+        if state_name == "Non-Methylated":
+            if current_start is None:
+                current_start = i
+            current_length += 1
+            current_probs.append(state_prob)
+        else:
+            if current_start is not None:
+                avg_prob = np.mean(current_probs)
+                classifications.append((current_start, current_length, "Non-Methylated", avg_prob))
+                current_start = None
+                current_length = 0
+                current_probs = []
+    
+    if current_start is not None:
+        avg_prob = np.mean(current_probs)
+        classifications.append((current_start, current_length, "Non-Methylated", avg_prob))
+
+    final_classifications = []
+    for start, length, classification, prob in classifications:
+        final_classification = None
+        feature_length = int(coordinates[start + length - 1]) - int(coordinates[start]) + 1
+        for feature_type, size_range in classification_mapping.items():
+            if size_range[0] <= feature_length < size_range[1]:
+                final_classification = feature_type
+                break
+            else:
+                pass
+        if not final_classification:
+            final_classification = classification
+
+        final_classifications.append((int(coordinates[start]) + 1, feature_length, final_classification, prob))
+    
+    return final_classifications
+
+            # if feature_length < 80:
+            #     final_classification = 'small_bound_stretch'
+            # elif 80 <= feature_length <= 200:
+            #     final_classification = 'Putative_Nucleosome'
+            # elif 200 < feature_length:
+            #     final_classification = 'large_bound_stretch'
+            # else:
+            #     pass

smftools/tools/archived/subset_adata_v1.py

@@ -0,0 +1,32 @@
+# subset_adata
+
+def subset_adata(adata, obs_columns):
+    """
+    Subsets an AnnData object based on categorical values in specified `.obs` columns.
+    
+    Parameters:
+        adata (AnnData): The AnnData object to subset.
+        obs_columns (list of str): List of `.obs` column names to subset by. The order matters.
+
+    Returns:
+        dict: A dictionary where keys are tuples of category values and values are corresponding AnnData subsets.
+    """
+
+    def subset_recursive(adata_subset, columns):
+        if not columns:
+            return {(): adata_subset}
+        
+        current_column = columns[0]
+        categories = adata_subset.obs[current_column].cat.categories
+        
+        subsets = {}
+        for cat in categories:
+            subset = adata_subset[adata_subset.obs[current_column] == cat]
+            subsets.update(subset_recursive(subset, columns[1:]))
+        
+        return subsets
+    
+    # Start the recursive subset process
+    subsets_dict = subset_recursive(adata, obs_columns)
+    
+    return subsets_dict

smftools/tools/archived/subset_adata_v2.py

@@ -0,0 +1,46 @@
+# subset_adata
+
+def subset_adata(adata, columns, cat_type='obs'):
+    """
+    Subsets an AnnData object based on categorical values in specified .obs or .var columns.
+    
+    Parameters:
+        adata (AnnData): The AnnData object to subset.
+        columns (list of str): List of .obs or .var column names to subset by. The order matters.
+        cat_type (str): obs or var. Default is obs
+
+    Returns:
+        dict: A dictionary where keys are tuples of category values and values are corresponding AnnData subsets.
+    """
+
+    def subset_recursive(adata_subset, columns, cat_type, key_prefix=()):
+        # Returns when the bottom of the stack is reached
+        if not columns:
+            # If there's only one column, return the key as a single value, not a tuple
+            if len(key_prefix) == 1:
+                return {key_prefix[0]: adata_subset}
+            return {key_prefix: adata_subset}
+        
+        current_column = columns[0]
+        subsets = {}
+
+        if 'obs' in cat_type:
+            categories = adata_subset.obs[current_column].cat.categories
+            for cat in categories:
+                subset = adata_subset[adata_subset.obs[current_column] == cat].copy()
+                new_key = key_prefix + (cat,)
+                subsets.update(subset_recursive(subset, columns[1:], cat_type, new_key))
+
+        elif 'var' in cat_type:
+            categories = adata_subset.var[current_column].cat.categories
+            for cat in categories:
+                subset = adata_subset[:, adata_subset.var[current_column] == cat].copy()
+                new_key = key_prefix + (cat,)
+                subsets.update(subset_recursive(subset, columns[1:], cat_type, new_key))            
+        
+        return subsets
+    
+    # Start the recursive subset process
+    subsets_dict = subset_recursive(adata, columns, cat_type)
+    
+    return subsets_dict

smftools/tools/calculate_distances.py

@@ -0,0 +1,18 @@
+# calculate_distances
+
+def calculate_distances(intervals, threshold=0.9):
+    """
+    Calculates distance between features in a read.
+    Takes in a list of intervals (start of feature, length of feature)
+    """
+    # Sort intervals by start position
+    intervals = sorted(intervals, key=lambda x: x[0])
+    intervals = [interval for interval in intervals if interval[2] > threshold]
+    
+    # Calculate distances
+    distances = []
+    for i in range(len(intervals) - 1):
+        end_current = intervals[i][0] + intervals[i][1]
+        start_next = intervals[i + 1][0]
+        distances.append(start_next - end_current)
+    return distances

smftools/tools/calculate_umap.py

@@ -0,0 +1,62 @@
+def calculate_umap(adata, layer='nan_half', var_filters=None, n_pcs=15, knn_neighbors=100, overwrite=True, threads=8):
+    import scanpy as sc
+    import numpy as np
+    import os
+    from scipy.sparse import issparse
+
+    os.environ["OMP_NUM_THREADS"] = str(threads)
+
+    # Step 1: Apply var filter
+    if var_filters:
+        subset_mask = np.logical_or.reduce([adata.var[f].values for f in var_filters])
+        adata_subset = adata[:, subset_mask].copy()
+        print(f"🔹 Subsetting adata: Retained {adata_subset.shape[1]} features based on filters {var_filters}")
+    else:
+        adata_subset = adata.copy()
+        print("🔹 No var filters provided. Using all features.")
+
+    # Step 2: NaN handling inside layer
+    if layer:
+        data = adata_subset.layers[layer]
+        if not issparse(data):
+            if np.isnan(data).any():
+                print("⚠ NaNs detected, filling with 0.5 before PCA + neighbors.")
+                data = np.nan_to_num(data, nan=0.5)
+                adata_subset.layers[layer] = data
+            else:
+                print("✅ No NaNs detected.")
+        else:
+            print("✅ Sparse matrix detected; skipping NaN check (sparse formats typically do not store NaNs).")
+
+    # Step 3: PCA + neighbors + UMAP on subset
+    if "X_umap" not in adata_subset.obsm or overwrite:
+        n_pcs = min(adata_subset.shape[1], n_pcs)
+        print(f"Running PCA with n_pcs={n_pcs}")
+        sc.pp.pca(adata_subset, layer=layer)
+        print('Running neighborhood graph')
+        sc.pp.neighbors(adata_subset, use_rep="X_pca", n_pcs=n_pcs, n_neighbors=knn_neighbors)
+        print('Running UMAP')
+        sc.tl.umap(adata_subset)
+
+    # Step 4: Store results in original adata
+    adata.obsm["X_pca"] = adata_subset.obsm["X_pca"]
+    adata.obsm["X_umap"] = adata_subset.obsm["X_umap"]
+    adata.obsp["distances"] = adata_subset.obsp["distances"]
+    adata.obsp["connectivities"] = adata_subset.obsp["connectivities"]
+    adata.uns["neighbors"] = adata_subset.uns["neighbors"]
+
+
+    # Fix varm["PCs"] shape mismatch
+    pc_matrix = np.zeros((adata.shape[1], adata_subset.varm["PCs"].shape[1]))
+    if var_filters:
+        subset_mask = np.logical_or.reduce([adata.var[f].values for f in var_filters])
+        pc_matrix[subset_mask, :] = adata_subset.varm["PCs"]
+    else:
+        pc_matrix = adata_subset.varm["PCs"]  # No subsetting case
+
+    adata.varm["PCs"] = pc_matrix
+
+
+    print(f"✅ Stored: adata.obsm['X_pca'] and adata.obsm['X_umap']")
+
+    return adata

smftools/tools/call_hmm_peaks.py

@@ -0,0 +1,105 @@
+def call_hmm_peaks(adata, feature_configs, obs_column='Reference_strand', site_types=['GpC_site', 'CpG_site'], save_plot=False, output_dir=None, date_tag=None):
+    """
+    Calls peaks from HMM feature layers and annotates them into the AnnData object.
+
+    Parameters:
+        adata : AnnData object with HMM layers (from apply_hmm)
+        feature_configs : dict
+        min_distance : minimum distance between peaks
+        peak_width : window size around peak centers
+        peak_prominence : required peak prominence
+        peak_threshold : threshold for labeling a read as "present" at a peak
+        site_types : list of var site types to aggregate
+        save_plot : whether to save the plot
+        output_dir : path to save the figure if save_plot=True
+        date_tag : optional tag for filename
+    """
+    import matplotlib.pyplot as plt
+    from scipy.signal import find_peaks
+    import os
+    import numpy as np
+
+    peak_columns = []
+        
+    for feature_layer, config in feature_configs.items():
+        min_distance = config.get('min_distance', 200)
+        peak_width = config.get('peak_width', 200)
+        peak_prominence = config.get('peak_prominence', 0.2)
+        peak_threshold = config.get('peak_threshold', 0.8)
+
+        # 1️⃣ Calculate mean intensity profile
+        matrix = adata.layers[feature_layer]
+        means = np.mean(matrix, axis=0)
+        feature_peak_columns = []
+
+        # 2️⃣ Peak calling
+        peak_centers, _ = find_peaks(means, prominence=peak_prominence, distance=min_distance)
+        adata.uns[f'{feature_layer} peak_centers'] = peak_centers
+
+        # 3️⃣ Plot
+        plt.figure(figsize=(6, 3))
+        plt.plot(range(len(means)), means)
+        plt.title(f"{feature_layer} density with peak calls")
+        plt.xlabel("Genomic position")
+        plt.ylabel("Mean feature density")
+        y = max(means) / 2
+        for i, center in enumerate(peak_centers):
+            plus_minus_width = peak_width // 2
+            start = center - plus_minus_width
+            end = center + plus_minus_width
+            plt.axvspan(start, end, color='purple', alpha=0.2)
+            plt.axvline(center, color='red', linestyle='--')
+            if i%2:
+                aligned = [end, 'left']
+            else:
+                aligned = [start, 'right']
+            plt.text(aligned[0], 0, f"Peak {i}\n{center}", color='red', ha=aligned[1])
+
+        if save_plot and output_dir:
+            filename = f"{output_dir}/{date_tag or 'output'}_{feature_layer}_peaks.png"
+            plt.savefig(filename, bbox_inches='tight')
+            print(f"Saved plot to {filename}")
+        else:
+            plt.show()
+
+        # 4️⃣ Annotate peaks back into adata.obs
+        for center in peak_centers:
+            half_width = peak_width // 2
+            start, end = center - half_width, center + half_width
+            colname = f'{feature_layer}_peak_{center}'
+            peak_columns.append(colname)
+            feature_peak_columns.append(colname)
+
+            adata.var[colname] = (
+                (adata.var_names.astype(int) >= start) & 
+                (adata.var_names.astype(int) <= end)
+            )
+
+            # Feature layer intensity around peak
+            mean_values = np.mean(matrix[:, start:end+1], axis=1)
+            sum_values = np.sum(matrix[:, start:end+1], axis=1)
+            adata.obs[f'mean_{feature_layer}_around_{center}'] = mean_values
+            adata.obs[f'sum_{feature_layer}_around_{center}'] = sum_values
+            adata.obs[f'{feature_layer}_present_at_{center}'] = mean_values > peak_threshold
+
+            # Site-type based aggregation
+            for site_type in site_types:
+                adata.obs[f'{site_type}_sum_around_{center}'] = 0
+                adata.obs[f'{site_type}_mean_around_{center}'] = np.nan
+
+            references = adata.obs[obs_column].cat.categories
+            for ref in adata.obs[obs_column].cat.categories:
+                subset = adata[adata.obs[obs_column] == ref]
+                for site_type in site_types:
+                    mask = subset.var.get(f'{ref}_{site_type}', None)
+                    if mask is not None:
+                        region_mask = (subset.var_names[mask].astype(int) >= start) & (subset.var_names[mask].astype(int) <= end)
+                        region = subset[:, mask].X[:, region_mask]
+                        adata.obs.loc[subset.obs.index, f'{site_type}_sum_around_{center}'] = np.nansum(region, axis=1)
+                        adata.obs.loc[subset.obs.index, f'{site_type}_mean_around_{center}'] = np.nanmean(region, axis=1)
+
+        adata.var[f'is_in_any_{feature_layer}_peak'] = adata.var[feature_peak_columns].any(axis=1)
+        print(f"✅ Peak annotation completed for {feature_layer} with {len(peak_centers)} peaks.")
+
+    # Combine all peaks into a single "is_in_any_peak" column
+    adata.var['is_in_any_peak'] = adata.var[peak_columns].any(axis=1)