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

smftools/hmm/__init__.py

@@ -0,0 +1,20 @@
+from .apply_hmm_batched import apply_hmm_batched
+from .calculate_distances import calculate_distances
+from .call_hmm_peaks import call_hmm_peaks
+from .display_hmm import display_hmm
+from .hmm_readwrite import load_hmm, save_hmm
+from .nucleosome_hmm_refinement import refine_nucleosome_calls, infer_nucleosomes_in_large_bound
+from .train_hmm import train_hmm
+
+
+__all__ = [
+    "apply_hmm_batched",
+    "calculate_distances",
+    "call_hmm_peaks",
+    "display_hmm",
+    "load_hmm",
+    "refine_nucleosome_calls",
+    "infer_nucleosomes_in_large_bound",
+    "save_hmm",
+    "train_hmm"
+]

smftools/hmm/apply_hmm_batched.py

@@ -0,0 +1,242 @@
+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", "C"], device="cpu", threshold=0.7, deaminase_footprinting=False):
+    """
+    Applies an HMM model to an AnnData object using tensor-based sequence inputs.
+    If multiple methbases are passed, generates a combined feature set.
+    """
+
+    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",
+                "c": ref_subset.var[f"{ref}_any_C_site"] == True,
+                "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 ---
+    coordinates = adata.var_names.astype(int).values
+    
+    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:
+                    start_idx = np.searchsorted(coordinates, start, side="left")
+                    end_idx = np.searchsorted(coordinates, start + length - 1, side="right")
+                    bin_matrix[row_idx, start_idx:end_idx] = 1
+                    counts[row_idx] += 1
+        adata.layers[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.
+    """
+
+    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/hmm/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/hmm/call_hmm_peaks.py

@@ -0,0 +1,106 @@
+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,
+    inplace=False
+):
+    import numpy as np
+    import pandas as pd
+    import matplotlib.pyplot as plt
+    from scipy.signal import find_peaks
+
+    if not inplace:
+        adata = adata.copy()
+
+    # Ensure obs_column is categorical
+    if not isinstance(adata.obs[obs_column].dtype, pd.CategoricalDtype):
+        adata.obs[obs_column] = pd.Categorical(adata.obs[obs_column])
+
+    coordinates = adata.var_names.astype(int).values
+    peak_columns = []
+
+    obs_updates = {}
+
+    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)
+
+        matrix = adata.layers[feature_layer]
+        means = np.mean(matrix, axis=0)
+        peak_indices, _ = find_peaks(means, prominence=peak_prominence, distance=min_distance)
+        peak_centers = coordinates[peak_indices]
+        adata.uns[f'{feature_layer} peak_centers'] = peak_centers.tolist()
+
+        # Plot
+        plt.figure(figsize=(6, 3))
+        plt.plot(coordinates, means)
+        plt.title(f"{feature_layer} with peak calls")
+        plt.xlabel("Genomic position")
+        plt.ylabel("Mean intensity")
+        for i, center in enumerate(peak_centers):
+            start, end = center - peak_width // 2, center + peak_width // 2
+            plt.axvspan(start, end, color='purple', alpha=0.2)
+            plt.axvline(center, color='red', linestyle='--')
+            aligned = [end if i % 2 else start, 'left' if i % 2 else '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()
+
+        feature_peak_columns = []
+        for center in peak_centers:
+            start, end = center - peak_width // 2, center + peak_width // 2
+            colname = f'{feature_layer}_peak_{center}'
+            peak_columns.append(colname)
+            feature_peak_columns.append(colname)
+
+            peak_mask = (coordinates >= start) & (coordinates <= end)
+            adata.var[colname] = peak_mask
+
+            region = matrix[:, peak_mask]
+            obs_updates[f'mean_{feature_layer}_around_{center}'] = np.mean(region, axis=1)
+            obs_updates[f'sum_{feature_layer}_around_{center}'] = np.sum(region, axis=1)
+            obs_updates[f'{feature_layer}_present_at_{center}'] = np.mean(region, axis=1) > peak_threshold
+
+            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
+
+            for ref in adata.obs[obs_column].cat.categories:
+                ref_idx = adata.obs[obs_column] == ref
+                mask_key = f"{ref}_{site_type}"
+                for site_type in site_types:
+                    if mask_key not in adata.var:
+                        continue
+                    site_mask = adata.var[mask_key].values
+                    site_coords = coordinates[site_mask]
+                    region_mask = (site_coords >= start) & (site_coords <= end)
+                    if not region_mask.any():
+                        continue
+                    full_mask = site_mask.copy()
+                    full_mask[site_mask] = region_mask
+                    site_region = adata[ref_idx, full_mask].X
+                    if hasattr(site_region, "A"):
+                        site_region = site_region.A
+                    if site_region.shape[1] > 0:
+                        adata.obs.loc[ref_idx, f'{site_type}_sum_around_{center}'] = np.nansum(site_region, axis=1)
+                        adata.obs.loc[ref_idx, f'{site_type}_mean_around_{center}'] = np.nanmean(site_region, axis=1)
+                    else:
+                        pass
+
+        adata.var[f'is_in_any_{feature_layer}_peak'] = adata.var[feature_peak_columns].any(axis=1)
+        print(f"Annotated {len(peak_centers)} peaks for {feature_layer}")
+
+    adata.var['is_in_any_peak'] = adata.var[peak_columns].any(axis=1)
+    adata.obs = pd.concat([adata.obs, pd.DataFrame(obs_updates, index=adata.obs.index)], axis=1)
+
+    return adata if not inplace else None

smftools/hmm/display_hmm.py

@@ -0,0 +1,18 @@
+def display_hmm(hmm, state_labels=["Non-Methylated", "Methylated"], obs_labels=["0", "1"]):
+    import torch
+    print("\n**HMM Model Overview**")
+    print(hmm)
+
+    print("\n**Transition Matrix**")
+    transition_matrix = torch.exp(hmm.edges).detach().cpu().numpy()
+    for i, row in enumerate(transition_matrix):
+        label = state_labels[i] if state_labels else f"State {i}"
+        formatted_row = ", ".join(f"{p:.6f}" for p in row)
+        print(f"{label}: [{formatted_row}]")
+
+    print("\n**Emission Probabilities**")
+    for i, dist in enumerate(hmm.distributions):
+        label = state_labels[i] if state_labels else f"State {i}"
+        probs = dist.probs.detach().cpu().numpy()
+        formatted_emissions = {obs_labels[j]: probs[j] for j in range(len(probs))}
+        print(f"{label}: {formatted_emissions}")

smftools/hmm/hmm_readwrite.py

@@ -0,0 +1,16 @@
+def load_hmm(model_path, device='cpu'):
+    """
+    Reads in a pretrained HMM.
+    
+    Parameters:
+        model_path (str): Path to a pretrained HMM
+    """
+    import torch
+    # Load model using PyTorch
+    hmm = torch.load(model_path)
+    hmm.to(device) 
+    return hmm
+
+def save_hmm(model, model_path):
+    import torch
+    torch.save(model, model_path)

smftools/hmm/nucleosome_hmm_refinement.py

@@ -0,0 +1,104 @@
+def refine_nucleosome_calls(adata, layer_name, nan_mask_layer, hexamer_size=120, octamer_size=147, max_wiggle=40, device="cpu"):
+    import numpy as np
+
+    nucleosome_layer = adata.layers[layer_name]
+    nan_mask = adata.layers[nan_mask_layer]  # should be binary mask: 1 = nan region, 0 = valid data
+
+    hexamer_layer = np.zeros_like(nucleosome_layer)
+    octamer_layer = np.zeros_like(nucleosome_layer)
+
+    for read_idx, row in enumerate(nucleosome_layer):
+        in_patch = False
+        start_idx = None
+
+        for pos in range(len(row)):
+            if row[pos] == 1 and not in_patch:
+                in_patch = True
+                start_idx = pos
+            if (row[pos] == 0 or pos == len(row) - 1) and in_patch:
+                in_patch = False
+                end_idx = pos if row[pos] == 0 else pos + 1
+
+                # Expand boundaries into NaNs
+                left_expand = 0
+                right_expand = 0
+
+                # Left
+                for i in range(1, max_wiggle + 1):
+                    if start_idx - i >= 0 and nan_mask[read_idx, start_idx - i] == 1:
+                        left_expand += 1
+                    else:
+                        break
+                # Right
+                for i in range(1, max_wiggle + 1):
+                    if end_idx + i < nucleosome_layer.shape[1] and nan_mask[read_idx, end_idx + i] == 1:
+                        right_expand += 1
+                    else:
+                        break
+
+                expanded_start = start_idx - left_expand
+                expanded_end = end_idx + right_expand
+
+                available_size = expanded_end - expanded_start
+            
+                # Octamer placement
+                if available_size >= octamer_size:
+                    center = (expanded_start + expanded_end) // 2
+                    half_oct = octamer_size // 2
+                    octamer_layer[read_idx, center - half_oct: center - half_oct + octamer_size] = 1
+
+                # Hexamer placement
+                elif available_size >= hexamer_size:
+                    center = (expanded_start + expanded_end) // 2
+                    half_hex = hexamer_size // 2
+                    hexamer_layer[read_idx, center - half_hex: center - half_hex + hexamer_size] = 1
+
+    adata.layers[f"{layer_name}_hexamers"] = hexamer_layer
+    adata.layers[f"{layer_name}_octamers"] = octamer_layer
+
+    print(f"Added layers: {layer_name}_hexamers and {layer_name}_octamers")
+    return adata
+
+def infer_nucleosomes_in_large_bound(adata, large_bound_layer, combined_nuc_layer, nan_mask_layer, nuc_size=147, linker_size=50, exclusion_buffer=30, device="cpu"):
+    import numpy as np
+
+    large_bound = adata.layers[large_bound_layer]
+    existing_nucs = adata.layers[combined_nuc_layer]
+    nan_mask = adata.layers[nan_mask_layer]
+
+    inferred_layer = np.zeros_like(large_bound)
+
+    for read_idx, row in enumerate(large_bound):
+        in_patch = False
+        start_idx = None
+
+        for pos in range(len(row)):
+            if row[pos] == 1 and not in_patch:
+                in_patch = True
+                start_idx = pos
+            if (row[pos] == 0 or pos == len(row) - 1) and in_patch:
+                in_patch = False
+                end_idx = pos if row[pos] == 0 else pos + 1
+
+                # Adjust boundaries into flanking NaN regions without getting too close to existing nucleosomes
+                left_expand = start_idx
+                while left_expand > 0 and nan_mask[read_idx, left_expand - 1] == 1 and np.sum(existing_nucs[read_idx, max(0, left_expand - exclusion_buffer):left_expand]) == 0:
+                    left_expand -= 1
+
+                right_expand = end_idx
+                while right_expand < row.shape[0] and nan_mask[read_idx, right_expand] == 1 and np.sum(existing_nucs[read_idx, right_expand:min(row.shape[0], right_expand + exclusion_buffer)]) == 0:
+                    right_expand += 1
+
+                # Phase nucleosomes with linker spacing only
+                region = (left_expand, right_expand)
+                pos_cursor = region[0]
+                while pos_cursor + nuc_size <= region[1]:
+                    if np.all((existing_nucs[read_idx, pos_cursor - exclusion_buffer:pos_cursor + nuc_size + exclusion_buffer] == 0)):
+                        inferred_layer[read_idx, pos_cursor:pos_cursor + nuc_size] = 1
+                        pos_cursor += nuc_size + linker_size 
+                    else:
+                        pos_cursor += 1
+
+    adata.layers[f"{large_bound_layer}_phased_nucleosomes"] = inferred_layer
+    print(f"Added layer: {large_bound_layer}_phased_nucleosomes")
+    return adata

smftools/hmm/train_hmm.py

@@ -0,0 +1,78 @@
+def train_hmm(
+    data,
+    emission_probs=[[0.8, 0.2], [0.2, 0.8]],
+    transitions=[[0.9, 0.1], [0.1, 0.9]],
+    start_probs=[0.5, 0.5],
+    end_probs=[0.5, 0.5],
+    device=None,
+    max_iter=50,
+    verbose=True,
+    tol=50,
+    pad_value=0,
+):
+    """
+    Trains a 2-state DenseHMM model on binary methylation/deamination data.
+
+    Parameters:
+        data (list or np.ndarray): List of sequences (lists) with 0, 1, or NaN.
+        emission_probs (list): List of emission probabilities for two states.
+        transitions (list): Transition matrix between states.
+        start_probs (list): Initial state probabilities.
+        end_probs (list): End state probabilities.
+        device (str or torch.device): "cpu", "mps", "cuda", or None (auto).
+        max_iter (int): Maximum EM iterations.
+        verbose (bool): Verbose output from pomegranate.
+        tol (float): Convergence tolerance.
+        pad_value (int): Value used to pad shorter sequences.
+
+    Returns:
+        hmm: Trained DenseHMM model
+    """
+    import torch
+    from pomegranate.hmm import DenseHMM
+    from pomegranate.distributions import Categorical
+    import numpy as np
+    from tqdm import tqdm
+
+    # Auto device detection
+    if device is None:
+        device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
+    elif isinstance(device, str):
+        device = torch.device(device)
+    print(f"Using device: {device}")
+
+    # Ensure emission probs on correct device
+    dists = [
+        Categorical(torch.tensor([p], device=device))
+        for p in emission_probs
+    ]
+
+    # Create DenseHMM
+    hmm = DenseHMM(
+        distributions=dists,
+        edges=transitions,
+        starts=start_probs,
+        ends=end_probs,
+        verbose=verbose,
+        max_iter=max_iter,
+        tol=tol,
+    ).to(device)
+
+    # Convert data to list if needed
+    if isinstance(data, np.ndarray):
+        data = data.tolist()
+
+    # Preprocess data (replace NaNs + pad)
+    max_length = max(len(seq) for seq in data)
+    processed_data = []
+    for sequence in tqdm(data, desc="Preprocessing Sequences"):
+        cleaned_seq = [int(x) if not np.isnan(x) else np.random.choice([0, 1]) for x in sequence]
+        cleaned_seq += [pad_value] * (max_length - len(cleaned_seq))
+        processed_data.append(cleaned_seq)
+
+    tensor_data = torch.tensor(processed_data, dtype=torch.long, device=device).unsqueeze(-1)
+
+    # Fit HMM
+    hmm.fit(tensor_data)
+
+    return hmm

smftools/informatics/__init__.py

@@ -0,0 +1,14 @@
+from . import helpers
+from .basecall_pod5s import basecall_pod5s
+from .subsample_fasta_from_bed import subsample_fasta_from_bed
+from .subsample_pod5 import subsample_pod5
+from .fast5_to_pod5 import fast5_to_pod5
+
+
+__all__ = [
+    "basecall_pod5s",
+    "subsample_fasta_from_bed",
+    "subsample_pod5",
+    "fast5_to_pod5",
+    "helpers"
+]