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

smftools/tools/cluster_adata_on_methylation.py

@@ -0,0 +1,105 @@
+# cluster_adata_on_methylation
+
+def cluster_adata_on_methylation(adata, obs_columns, method='hierarchical', n_clusters=3, layer=None, site_types = ['GpC_site', 'CpG_site']):
+    """
+    Adds cluster groups to the adata object as an observation column
+
+    Parameters:
+        adata
+        obs_columns
+        method
+        n_clusters
+        layer
+        site_types
+
+    Returns:
+        None
+    """
+    import pandas as pd
+    import numpy as np
+    from . import subset_adata
+    from ..readwrite import adata_to_df
+
+    # Ensure obs_columns are categorical
+    for col in obs_columns:
+        adata.obs[col] = adata.obs[col].astype('category')
+
+    references = adata.obs['Reference'].cat.categories
+
+    # Add subset metadata to the adata
+    subset_adata(adata, obs_columns)
+
+    subgroup_name = '_'.join(obs_columns)
+    subgroups = adata.obs[subgroup_name].cat.categories
+
+    subgroup_to_reference_map = {}
+    for subgroup in subgroups:
+        for reference in references:
+            if reference in subgroup:
+                subgroup_to_reference_map[subgroup] = reference
+            else:
+                pass
+
+    if method == 'hierarchical':
+        for site_type in site_types:
+            adata.obs[f'{site_type}_{layer}_hierarchical_clustering_index_within_{subgroup_name}'] = pd.Series(-1, index=adata.obs_names, dtype=int)
+    elif method == 'kmeans':
+        for site_type in site_types:
+            adata.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}'] = pd.Series(-1, index=adata.obs_names, dtype=int)
+        
+    for subgroup in subgroups:
+        subgroup_subset = adata[adata.obs[subgroup_name] == subgroup].copy()
+        reference = subgroup_to_reference_map[subgroup]
+        for site_type in site_types:
+            site_subset = subgroup_subset[:, np.array(subgroup_subset.var[f'{reference}_{site_type}'])].copy()
+            df = adata_to_df(site_subset, layer=layer)
+            df2 = df.reset_index(drop=True)
+            if method == 'hierarchical':
+                try:
+                    from scipy.cluster.hierarchy import linkage, dendrogram
+                    # Perform hierarchical clustering on rows using the average linkage method and Euclidean metric
+                    row_linkage = linkage(df2.values, method='average', metric='euclidean')
+
+                    # Generate the dendrogram to get the ordered indices
+                    dendro = dendrogram(row_linkage, no_plot=True)
+                    reordered_row_indices = np.array(dendro['leaves']).astype(int)
+
+                    # Get the reordered observation names
+                    reordered_obs_names = [df.index[i] for i in reordered_row_indices]
+
+                    temp_obs_data = pd.DataFrame({f'{site_type}_{layer}_hierarchical_clustering_index_within_{subgroup_name}': np.arange(0, len(reordered_obs_names), 1)}, index=reordered_obs_names, dtype=int)
+                    adata.obs.update(temp_obs_data)
+                except:
+                    print(f'Error found in {subgroup} of {site_type}_{layer}_hierarchical_clustering_index_within_{subgroup_name}')
+            elif method == 'kmeans':
+                try:
+                    from sklearn.cluster import KMeans
+                    kmeans = KMeans(n_clusters=n_clusters)
+                    kmeans.fit(site_subset.layers[layer])
+                    # Get the cluster labels for each data point
+                    cluster_labels = kmeans.labels_
+                    # Add the kmeans cluster data as an observation to the anndata object
+                    site_subset.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}'] = cluster_labels.astype(str)
+                    # Calculate the mean of each observation categoty of each cluster
+                    cluster_means = site_subset.obs.groupby(f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}').mean()
+                    # Sort the cluster indices by mean methylation value
+                    sorted_clusters = cluster_means.sort_values(by=f'{site_type}_row_methylation_means', ascending=False).index
+                    # Create a mapping of the old cluster values to the new cluster values
+                    sorted_cluster_mapping = {old: new for new, old in enumerate(sorted_clusters)}
+                    # Apply the mapping to create a new observation value: kmeans_labels_reordered
+                    site_subset.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}'] = site_subset.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}'].map(sorted_cluster_mapping)
+                    temp_obs_data = pd.DataFrame({f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}': site_subset.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}']}, index=site_subset.obs_names, dtype=int)
+                    adata.obs.update(temp_obs_data)      
+                except:
+                        print(f'Error found in {subgroup} of {site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}')   
+
+    if method == 'hierarchical':             
+        # Ensure that the observation values are type int
+        for site_type in site_types:
+            adata.obs[f'{site_type}_{layer}_hierarchical_clustering_index_within_{subgroup_name}'] = adata.obs[f'{site_type}_{layer}_hierarchical_clustering_index_within_{subgroup_name}'].astype(int)
+    elif method == 'kmeans':
+        # Ensure that the observation values are type int
+        for site_type in site_types:
+            adata.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}'] = adata.obs[f'{site_type}_{layer}_kmeans_clustering_index_within_{subgroup_name}'].astype(int)
+
+    return None   

smftools/tools/data/__init__.py

@@ -0,0 +1,2 @@
+from .anndata_data_module import AnnDataModule
+from .preprocessing import random_fill_nans

smftools/tools/data/anndata_data_module.py

@@ -0,0 +1,90 @@
+import torch
+from torch.utils.data import DataLoader, TensorDataset, random_split
+import pytorch_lightning as pl
+import numpy as np
+import pandas as pd
+
+class AnnDataModule(pl.LightningDataModule):
+    def __init__(self, adata, tensor_source="X", tensor_key=None, label_col="labels", 
+                 batch_size=64, train_frac=0.7, random_seed=42, split_col='train_val_split', split_save_path=None, load_existing_split=False,
+                 inference_mode=False):
+        super().__init__()
+        self.adata = adata # The adata object
+        self.tensor_source = tensor_source # X, layers, obsm
+        self.tensor_key = tensor_key  # name of the layer or obsm key
+        self.label_col = label_col # name of the label column in obs
+        self.batch_size = batch_size
+        self.train_frac = train_frac
+        self.random_seed = random_seed
+        self.split_col = split_col  # Name of obs column to store "train"/"val"
+        self.split_save_path = split_save_path # Where to save the obs_names and train/test split logging
+        self.load_existing_split = load_existing_split # Whether to load from an existing split
+        self.inference_mode = inference_mode # Whether to load the AnnDataModule in inference mode.
+
+    def setup(self, stage=None):
+        # Load feature matrix
+        if self.tensor_source == "X":
+            X = self.adata.X
+        elif self.tensor_source == "layers":
+            assert self.tensor_key in self.adata.layers, f"Layer '{self.tensor_key}' not found."
+            X = self.adata.layers[self.tensor_key]
+        elif self.tensor_source == "obsm":
+            assert self.tensor_key in self.adata.obsm, f"obsm key '{self.tensor_key}' not found."
+            X = self.adata.obsm[self.tensor_key]
+        else:
+            raise ValueError(f"Invalid tensor_source: {self.tensor_source}")
+
+        # Convert to tensor
+        X_tensor = torch.tensor(X, dtype=torch.float32)
+
+        if self.inference_mode:
+            self.infer_dataset = TensorDataset(X_tensor)
+
+        else:
+            # Load and encode labels
+            y = self.adata.obs[self.label_col]
+            if y.dtype.name == 'category':
+                y = y.cat.codes
+            y_tensor = torch.tensor(y.values, dtype=torch.long)
+
+            # Use existing split
+            if self.load_existing_split:
+                split_df = pd.read_csv(self.split_save_path, index_col=0)
+                assert self.split_col in split_df.columns, f"'{self.split_col}' column missing in split file."
+                self.adata.obs[self.split_col] = split_df.loc[self.adata.obs_names][self.split_col].values
+
+            # If no split exists, create one
+            if self.split_col not in self.adata.obs:
+                full_dataset = TensorDataset(X_tensor, y_tensor)
+                n_train = int(self.train_frac * len(full_dataset))
+                n_val = len(full_dataset) - n_train
+                self.train_set, self.val_set = random_split(
+                    full_dataset, [n_train, n_val],
+                    generator=torch.Generator().manual_seed(self.random_seed)
+                )
+                # Assign split labels
+                split_array = np.full(len(self.adata), "val", dtype=object)
+                train_idx = self.train_set.indices if hasattr(self.train_set, "indices") else self.train_set._indices
+                split_array[train_idx] = "train"
+                self.adata.obs[self.split_col] = split_array
+
+                # Save to disk
+                if self.split_save_path:
+                    self.adata.obs[[self.split_col]].to_csv(self.split_save_path)
+            else:
+                split_labels = self.adata.obs[self.split_col].values
+                train_mask = split_labels == "train"
+                val_mask = split_labels == "val"
+                self.train_set = TensorDataset(X_tensor[train_mask], y_tensor[train_mask])
+                self.val_set = TensorDataset(X_tensor[val_mask], y_tensor[val_mask])
+
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size)
+    
+    def predict_dataloader(self):
+        if not self.inference_mode:
+            raise RuntimeError("predict_dataloader only available in inference mode.")
+        return DataLoader(self.infer_dataset, batch_size=self.batch_size)

smftools/tools/data/preprocessing.py

@@ -0,0 +1,6 @@
+import numpy as np
+
+def random_fill_nans(X):
+    nan_mask = np.isnan(X)
+    X[nan_mask] = np.random.rand(*X[nan_mask].shape)
+    return X

smftools/tools/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/tools/general_tools.py

@@ -0,0 +1,69 @@
+def create_nan_mask_from_X(adata, new_layer_name="nan_mask"):
+    """
+    Generates a nan mask where 1 = NaN in adata.X and 0 = valid value.
+    """
+    import numpy as np
+    nan_mask = np.isnan(adata.X).astype(int)
+    adata.layers[new_layer_name] = nan_mask
+    print(f"✅ Created '{new_layer_name}' layer based on NaNs in adata.X")
+    return adata
+
+def create_nan_or_non_gpc_mask(adata, obs_column, new_layer_name="nan_or_non_gpc_mask"):
+    import numpy as np
+
+    nan_mask = np.isnan(adata.X).astype(int)
+    combined_mask = np.zeros_like(nan_mask)
+
+    for idx, row in enumerate(adata.obs.itertuples()):
+        ref = getattr(row, obs_column)
+        gpc_mask = adata.var[f"{ref}_GpC_site"].astype(int).values
+        combined_mask[idx, :] = 1 - gpc_mask  # non-GpC is 1
+
+    mask = np.maximum(nan_mask, combined_mask)
+    adata.layers[new_layer_name] = mask
+
+    print(f"✅ Created '{new_layer_name}' layer based on NaNs in adata.X and non-GpC regions using {obs_column}")
+    return adata
+
+def combine_layers(adata, input_layers, output_layer, negative_mask=None, values=None, binary_mode=False):
+    """
+    Combines layers into a single layer with specific coding:
+        - Background stays 0
+        - If binary_mode=True: any overlap = 1
+        - If binary_mode=False:
+            - Defaults to [1, 2, 3, ...] if values=None
+            - Later layers take precedence in overlaps
+    
+    Parameters:
+        adata: AnnData object
+        input_layers: list of str
+        output_layer: str, name of the output layer
+        negative_mask: str (optional), binary mask to enforce 0s
+        values: list of ints (optional), values to assign to each input layer
+        binary_mode: bool, if True, creates a simple 0/1 mask regardless of values
+    
+    Returns:
+        Updated AnnData with new layer.
+    """
+    import numpy as np
+    combined = np.zeros_like(adata.layers[input_layers[0]])
+
+    if binary_mode:
+        for layer in input_layers:
+            combined = np.logical_or(combined, adata.layers[layer] > 0)
+        combined = combined.astype(int)
+    else:
+        if values is None:
+            values = list(range(1, len(input_layers) + 1))
+        for i, layer in enumerate(input_layers):
+            arr = adata.layers[layer]
+            combined[arr > 0] = values[i]
+
+    if negative_mask:
+        mask = adata.layers[negative_mask]
+        combined[mask == 0] = 0
+
+    adata.layers[output_layer] = combined
+    print(f"✅ Combined layers into {output_layer} {'(binary)' if binary_mode else f'with values {values}'}")
+
+    return adata

smftools/tools/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/tools/inference/__init__.py

@@ -0,0 +1 @@
+from .lightning_inference import run_lightning_inference

smftools/tools/inference/lightning_inference.py

@@ -0,0 +1,41 @@
+import torch
+import pandas as pd
+import numpy as np
+from pytorch_lightning import Trainer
+
+def run_lightning_inference(
+    adata,
+    model,
+    datamodule,
+    label_col="labels",
+    prefix="model"
+):
+
+    # Get class labels
+    if label_col in adata.obs and pd.api.types.is_categorical_dtype(adata.obs[label_col]):
+        class_labels = adata.obs[label_col].cat.categories.tolist()
+    else:
+        raise ValueError("label_col must be a categorical column in adata.obs")
+
+    # Run predictions
+    trainer = Trainer(accelerator="auto", devices=1, logger=False, enable_checkpointing=False)
+    preds = trainer.predict(model, datamodule=datamodule)
+    probs = torch.cat(preds, dim=0).cpu().numpy()  # (N, C)
+    pred_class_idx = probs.argmax(axis=1)
+    pred_class_labels = [class_labels[i] for i in pred_class_idx]
+    pred_class_probs = probs[np.arange(len(probs)), pred_class_idx]
+
+    # Construct full prefix with label_col
+    full_prefix = f"{prefix}_{label_col}"
+
+    # Store predictions in obs
+    adata.obs[f"{full_prefix}_pred"] = pred_class_idx
+    adata.obs[f"{full_prefix}_pred_label"] = pd.Categorical(pred_class_labels, categories=class_labels)
+    adata.obs[f"{full_prefix}_pred_prob"] = pred_class_probs
+
+    # Per-class probabilities
+    for i, class_name in enumerate(class_labels):
+        adata.obs[f"{full_prefix}_prob_{class_name}"] = probs[:, i]
+
+    # Full probability matrix in obsm
+    adata.obsm[f"{full_prefix}_pred_prob_all"] = probs

smftools/tools/models/__init__.py

@@ -0,0 +1,9 @@
+from .base import BaseTorchModel
+from .mlp import MLPClassifier
+from .cnn import CNNClassifier
+from .rnn import RNNClassifier
+from .transformer import BaseTransformer, TransformerClassifier, DANNTransformerClassifier, MaskedTransformerPretrainer
+from .positional import PositionalEncoding
+from .wrappers import ScaledModel
+from .lightning_base import TorchClassifierWrapper
+from .sklearn_models import SklearnModelWrapper

smftools/tools/models/base.py

@@ -0,0 +1,14 @@
+import torch.nn as nn
+from ..utils.device import detect_device
+
+class BaseTorchModel(nn.Module):
+    """
+    Minimal base class for torch models that:
+    - Stores device
+    - Moves model to detected device on init
+    """
+    def __init__(self, dropout_rate=0.2):
+        super().__init__()
+        self.device = detect_device() # detects available devices
+        self.dropout_rate = dropout_rate # default dropout rate to be used in regularization.
+        self.to(self.device)  # move model to device

smftools/tools/models/cnn.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+from .base import BaseTorchModel
+
+class CNNClassifier(BaseTorchModel):
+    def __init__(self, input_size, num_classes, **kwargs):
+        super().__init__(**kwargs)
+        # Define convolutional layers
+        self.conv1 = nn.Conv1d(1, 16, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(16, 32, kernel_size=3, padding=1)
+        # Define activation function
+        self.relu = nn.ReLU()
+
+        # Determine the flattened size dynamically
+        dummy_input = torch.zeros(1, 1, input_size).to(self.device)
+        with torch.no_grad():
+            dummy_output = self._forward_conv(dummy_input)
+        flattened_size = dummy_output.view(1, -1).shape[1]
+
+        # Define fully connected layers
+        self.fc1 = nn.Linear(flattened_size, 64)
+        self.fc2 = nn.Linear(64, num_classes)
+
+    def _forward_conv(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        return x
+
+    def forward(self, x):
+        x = x.unsqueeze(1)  # [B, 1, L]
+        x = self._forward_conv(x)
+        x = x.view(x.size(0), -1)  # flatten
+        x = self.relu(self.fc1(x))
+        return self.fc2(x)

smftools/tools/models/lightning_base.py

@@ -0,0 +1,41 @@
+import torch
+import pytorch_lightning as pl
+
+class TorchClassifierWrapper(pl.LightningModule):
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        optimizer_cls=torch.optim.AdamW,
+        optimizer_kwargs=None,
+        criterion_cls=torch.nn.CrossEntropyLoss,
+        criterion_kwargs=None,
+        lr: float = 1e-3,
+    ):
+        super().__init__()
+        self.model = model
+        self.save_hyperparameters(ignore=['model'])  # logs all except actual model instance
+        self.optimizer_cls = optimizer_cls
+        self.optimizer_kwargs = optimizer_kwargs or {}
+        self.criterion = criterion_cls(**(criterion_kwargs or {}))
+        self.lr = lr
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = self.criterion(logits, y)
+        self.log("train_loss", loss, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = self.criterion(logits, y)
+        acc = (logits.argmax(dim=1) == y).float().mean()
+        self.log_dict({"val_loss": loss, "val_acc": acc}, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        return self.optimizer_cls(self.parameters(), lr=self.lr, **self.optimizer_kwargs)

smftools/tools/models/mlp.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+from .base import BaseTorchModel
+
+class MLPClassifier(BaseTorchModel):
+    def __init__(self, input_dim, num_classes, hidden_sizes=(128, 64), **kwargs):
+        super().__init__(**kwargs)
+        layers = []
+        prev = input_dim
+        for h in hidden_sizes:
+            layers.extend([nn.Linear(prev, h), nn.ReLU(), nn.Dropout(self.dropout_rate)])
+            prev = h
+        layers.append(nn.Linear(prev, num_classes))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)

smftools/tools/models/positional.py

@@ -0,0 +1,17 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1).float()
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.pe = pe.unsqueeze(0)  # (1, max_len, d_model)
+
+    def forward(self, x):
+        x = x + self.pe[:, :x.size(1)].to(x.device)
+        return x

smftools/tools/models/rnn.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+from .base import BaseTorchModel
+
+class RNNClassifier(BaseTorchModel):
+    def __init__(self, input_size, hidden_dim, num_classes, **kwargs):
+        super().__init__(**kwargs)
+        # Define LSTM layer
+        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_dim, batch_first=True)
+        # Define fully connected output layer
+        self.fc = nn.Linear(hidden_dim, num_classes)
+
+    def forward(self, x):
+        x = x.unsqueeze(1)  # [B, 1, L] → for LSTM expecting batch_first
+        _, (h_n, _) = self.lstm(x)  # h_n: [1, B, H]
+        return self.fc(h_n.squeeze(0))  # [B, H] → [B, num_classes]

smftools/tools/models/sklearn_models.py

@@ -0,0 +1,40 @@
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.naive_bayes import GaussianNB
+from sklearn.metrics import (
+    roc_curve, precision_recall_curve, auc, f1_score, confusion_matrix
+)
+from sklearn.utils.class_weight import compute_class_weight
+
+import numpy as np
+
+class SklearnModelWrapper:
+    def __init__(self, model):
+        self.model = model
+
+    def fit(self, X_train, y_train):
+        self.model.fit(X_train, y_train)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        return self.model.predict_proba(X)
+
+    def evaluate(self, X_test, y_test):
+        probs = self.predict_proba(X_test)[:, 1]
+        preds = self.predict(X_test)
+
+        fpr, tpr, _ = roc_curve(y_test, probs)
+        precision, recall, _ = precision_recall_curve(y_test, probs)
+        f1 = f1_score(y_test, preds)
+        auc_score = auc(fpr, tpr)
+        pr_auc = auc(recall, precision)
+        cm = confusion_matrix(y_test, preds)
+        pos_freq = np.mean(y_test == 1)
+        pr_auc_norm = pr_auc / pos_freq
+
+        return {
+            "fpr": fpr, "tpr": tpr, "precision": precision, "recall": recall,
+            "f1": f1, "auc": auc_score, "pr_auc": pr_auc,
+            "pr_auc_norm": pr_auc_norm, "confusion_matrix": cm
+        }