biolmnet/training.py

from __future__ import annotations

import copy
import random
from dataclasses import asdict, dataclass
from typing import Callable

import numpy as np
import pandas as pd
import torch
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from .data import PreparedWorkspace
from .model import BioLMNet


ProgressCallback = Callable[[float, str], None]


@dataclass(frozen=True)
class Hyperparameters:
    epochs: int = 50
    batch_size: int = 16
    learning_rate: float = 0.001
    weight_decay: float = 0.01
    dropout: float = 0.3
    projection_dim: int = 64
    fusion_dim: int = 12
    validation_fraction: float = 0.2
    optimizer: str = "Adam"
    class_weighting: bool = True
    early_stopping_patience: int = 12
    seed: int = 42

    def validate(self) -> None:
        if not 1 <= self.epochs <= 1000:
            raise ValueError("Epochs must be between 1 and 1,000.")
        if not 2 <= self.batch_size <= 1024:
            raise ValueError("Batch size must be between 2 and 1,024.")
        if not 0 < self.learning_rate <= 1:
            raise ValueError("Learning rate must be in (0, 1].")
        if not 0 <= self.dropout < 1:
            raise ValueError("Dropout must be in [0, 1).")
        if not 0.05 <= self.validation_fraction <= 0.5:
            raise ValueError("Validation fraction must be between 0.05 and 0.5.")
        if self.optimizer.lower() not in {"adam", "sgd"}:
            raise ValueError("Optimizer must be Adam or SGD.")


@dataclass
class ModelBundle:
    model: BioLMNet
    gene_features: list[str]
    dna_features: list[str]
    label_names: list[str]
    gene_mean: np.ndarray
    gene_scale: np.ndarray
    dna_mean: np.ndarray
    dna_scale: np.ndarray
    config: dict
    metrics: dict
    history: list[dict[str, float]]


@dataclass
class TrainingResult:
    bundle: ModelBundle
    validation_predictions: pd.DataFrame
    confusion: np.ndarray


def set_reproducible_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _fit_scaler(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    mean = values.mean(axis=0, dtype=np.float64).astype(np.float32)
    scale = values.std(axis=0, dtype=np.float64).astype(np.float32)
    scale[scale < 1e-8] = 1.0
    return mean, scale


def _scale(values: np.ndarray, mean: np.ndarray, scale: np.ndarray) -> np.ndarray:
    return ((values.astype(np.float32) - mean) / scale).astype(np.float32)


def _make_model(
    workspace: PreparedWorkspace, hyperparameters: Hyperparameters
) -> BioLMNet:
    gene = workspace.gene_branch
    dna = workspace.dna_branch
    if (
        gene.embeddings is None
        or dna.embeddings is None
        or gene.pathway_mask is None
        or dna.pathway_mask is None
    ):
        raise ValueError("Workspace priors are incomplete; run data preparation first.")

    return BioLMNet(
        gene_biological_mask=torch.from_numpy(gene.biological_mask),
        dna_biological_mask=torch.from_numpy(dna.biological_mask),
        gene_embeddings=torch.from_numpy(gene.embeddings),
        dna_embeddings=torch.from_numpy(dna.embeddings),
        gene_pathway_mask=torch.from_numpy(gene.pathway_mask),
        dna_pathway_mask=torch.from_numpy(dna.pathway_mask),
        n_classes=len(workspace.label_names),
        projection_dim=hyperparameters.projection_dim,
        fusion_dim=hyperparameters.fusion_dim,
        dropout=hyperparameters.dropout,
    )


def _evaluate(
    model: BioLMNet,
    gene_values: np.ndarray,
    dna_values: np.ndarray,
    labels: np.ndarray,
    device: torch.device,
) -> tuple[float, np.ndarray, np.ndarray]:
    model.eval()
    with torch.no_grad():
        logits = model(
            torch.from_numpy(gene_values).to(device),
            torch.from_numpy(dna_values).to(device),
        )
        loss = nn.functional.cross_entropy(
            logits, torch.from_numpy(labels).to(device)
        ).item()
        probabilities = torch.softmax(logits, dim=1).cpu().numpy()
    predictions = probabilities.argmax(axis=1)
    return float(loss), probabilities, predictions


def train(
    workspace: PreparedWorkspace,
    hyperparameters: Hyperparameters,
    progress: ProgressCallback | None = None,
) -> TrainingResult:
    hyperparameters.validate()
    set_reproducible_seed(hyperparameters.seed)
    labels = workspace.labels
    classes, class_counts = np.unique(labels, return_counts=True)
    if class_counts.min() < 2:
        raise ValueError(
            "Each class needs at least two samples for a stratified train/validation split."
        )

    indices = np.arange(len(labels))
    train_index, validation_index = train_test_split(
        indices,
        test_size=hyperparameters.validation_fraction,
        random_state=hyperparameters.seed,
        stratify=labels,
    )
    gene_mean, gene_scale = _fit_scaler(
        workspace.gene_expression[train_index]
    )
    dna_mean, dna_scale = _fit_scaler(workspace.dna_methylation[train_index])
    gene_train = _scale(
        workspace.gene_expression[train_index], gene_mean, gene_scale
    )
    gene_validation = _scale(
        workspace.gene_expression[validation_index], gene_mean, gene_scale
    )
    dna_train = _scale(
        workspace.dna_methylation[train_index], dna_mean, dna_scale
    )
    dna_validation = _scale(
        workspace.dna_methylation[validation_index], dna_mean, dna_scale
    )
    y_train = labels[train_index]
    y_validation = labels[validation_index]

    dataset = TensorDataset(
        torch.from_numpy(gene_train),
        torch.from_numpy(dna_train),
        torch.from_numpy(y_train),
    )
    generator = torch.Generator().manual_seed(hyperparameters.seed)
    loader = DataLoader(
        dataset,
        batch_size=min(hyperparameters.batch_size, len(dataset)),
        shuffle=True,
        generator=generator,
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = _make_model(workspace, hyperparameters).to(device)
    if hyperparameters.optimizer.lower() == "sgd":
        optimizer = torch.optim.SGD(
            model.parameters(),
            lr=hyperparameters.learning_rate,
            weight_decay=hyperparameters.weight_decay,
        )
    else:
        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=hyperparameters.learning_rate,
            weight_decay=hyperparameters.weight_decay,
        )

    class_weights: torch.Tensor | None = None
    if hyperparameters.class_weighting:
        count_by_class = np.bincount(
            y_train, minlength=len(workspace.label_names)
        ).astype(np.float32)
        weights = len(y_train) / (len(count_by_class) * count_by_class)
        class_weights = torch.from_numpy(weights).to(device)
    criterion = nn.CrossEntropyLoss(weight=class_weights)

    history: list[dict[str, float]] = []
    best_state: dict[str, torch.Tensor] | None = None
    best_loss = float("inf")
    patience = 0
    for epoch in range(hyperparameters.epochs):
        model.train()
        running_loss = 0.0
        seen = 0
        for gene_batch, dna_batch, label_batch in loader:
            gene_batch = gene_batch.to(device)
            dna_batch = dna_batch.to(device)
            label_batch = label_batch.to(device)
            optimizer.zero_grad(set_to_none=True)
            logits = model(gene_batch, dna_batch)
            loss = criterion(logits, label_batch)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * len(label_batch)
            seen += len(label_batch)

        validation_loss, _, validation_predictions = _evaluate(
            model,
            gene_validation,
            dna_validation,
            y_validation,
            device,
        )
        epoch_row = {
            "epoch": float(epoch + 1),
            "training_loss": float(running_loss / max(seen, 1)),
            "validation_loss": validation_loss,
            "validation_accuracy": float(
                accuracy_score(y_validation, validation_predictions)
            ),
            "validation_f1_macro": float(
                f1_score(
                    y_validation,
                    validation_predictions,
                    average="macro",
                    zero_division=0,
                )
            ),
        }
        history.append(epoch_row)
        if progress:
            progress(
                (epoch + 1) / hyperparameters.epochs,
                (
                    f"Epoch {epoch + 1}/{hyperparameters.epochs} · "
                    f"validation F1 {epoch_row['validation_f1_macro']:.3f}"
                ),
            )

        if validation_loss < best_loss - 1e-5:
            best_loss = validation_loss
            best_state = copy.deepcopy(model.state_dict())
            patience = 0
        else:
            patience += 1
            if patience >= hyperparameters.early_stopping_patience:
                break

    if best_state is not None:
        model.load_state_dict(best_state)
    validation_loss, probabilities, predictions = _evaluate(
        model,
        gene_validation,
        dna_validation,
        y_validation,
        device,
    )
    confusion = confusion_matrix(
        y_validation,
        predictions,
        labels=np.arange(len(workspace.label_names)),
    )
    metrics = {
        "validation_loss": validation_loss,
        "accuracy": float(accuracy_score(y_validation, predictions)),
        "f1_macro": float(
            f1_score(
                y_validation, predictions, average="macro", zero_division=0
            )
        ),
        "f1_weighted": float(
            f1_score(
                y_validation, predictions, average="weighted", zero_division=0
            )
        ),
        "precision_macro": float(
            precision_score(
                y_validation, predictions, average="macro", zero_division=0
            )
        ),
        "recall_macro": float(
            recall_score(
                y_validation, predictions, average="macro", zero_division=0
            )
        ),
        "epochs_completed": len(history),
        "training_samples": int(len(train_index)),
        "validation_samples": int(len(validation_index)),
        "device": str(device),
    }
    predictions_frame = pd.DataFrame(
        {
            "sample_row": validation_index,
            "observed": [
                workspace.label_names[value] for value in y_validation
            ],
            "predicted": [
                workspace.label_names[value] for value in predictions
            ],
        }
    )
    for index, label in enumerate(workspace.label_names):
        predictions_frame[f"P({label})"] = probabilities[:, index]

    gene = workspace.gene_branch
    dna = workspace.dna_branch
    config = {
        "format_version": 1,
        "source_name": workspace.source_name,
        "gene_features": gene.input_genes,
        "dna_features": dna.input_genes,
        "label_names": workspace.label_names,
        "gene_hidden_genes": gene.hidden_genes,
        "dna_hidden_genes": dna.hidden_genes,
        "gene_pathways": gene.pathways,
        "dna_pathways": dna.pathways,
        "hyperparameters": asdict(hyperparameters),
        "architecture": {
            "projection_dim": hyperparameters.projection_dim,
            "fusion_dim": hyperparameters.fusion_dim,
            "dropout": hyperparameters.dropout,
            "biological_activation": "relu",
            "projection_activation": "sigmoid",
            "fusion_activation": "tanh",
        },
    }
    bundle = ModelBundle(
        model=model.cpu(),
        gene_features=gene.input_genes,
        dna_features=dna.input_genes,
        label_names=workspace.label_names,
        gene_mean=gene_mean,
        gene_scale=gene_scale,
        dna_mean=dna_mean,
        dna_scale=dna_scale,
        config=config,
        metrics=metrics,
        history=history,
    )
    return TrainingResult(
        bundle=bundle,
        validation_predictions=predictions_frame,
        confusion=confusion,
    )


def validate_prediction_frames(
    gene_frame: pd.DataFrame,
    dna_frame: pd.DataFrame,
    bundle: ModelBundle,
) -> tuple[np.ndarray, np.ndarray]:
    if len(gene_frame) != len(dna_frame):
        raise ValueError(
            "Prediction gene-expression and DNA-methylation files must have "
            "the same number of rows."
        )
    missing_gene = sorted(set(bundle.gene_features) - set(gene_frame.columns))
    missing_dna = sorted(set(bundle.dna_features) - set(dna_frame.columns))
    if missing_gene or missing_dna:
        details = []
        if missing_gene:
            details.append(
                "gene-expression: " + ", ".join(missing_gene[:8])
            )
        if missing_dna:
            details.append("DNA-methylation: " + ", ".join(missing_dna[:8]))
        raise ValueError(
            "Prediction files are missing trained features (" + "; ".join(details) + ")."
        )
    gene_values = gene_frame.loc[:, bundle.gene_features].apply(
        pd.to_numeric, errors="coerce"
    )
    dna_values = dna_frame.loc[:, bundle.dna_features].apply(
        pd.to_numeric, errors="coerce"
    )
    if gene_values.isna().any().any() or dna_values.isna().any().any():
        raise ValueError("Prediction inputs contain missing or non-numeric values.")
    return (
        _scale(gene_values.to_numpy(), bundle.gene_mean, bundle.gene_scale),
        _scale(dna_values.to_numpy(), bundle.dna_mean, bundle.dna_scale),
    )


def predict(
    gene_frame: pd.DataFrame,
    dna_frame: pd.DataFrame,
    bundle: ModelBundle,
) -> pd.DataFrame:
    gene_values, dna_values = validate_prediction_frames(
        gene_frame, dna_frame, bundle
    )
    bundle.model.eval()
    with torch.no_grad():
        logits = bundle.model(
            torch.from_numpy(gene_values), torch.from_numpy(dna_values)
        )
        probabilities = torch.softmax(logits, dim=1).numpy()
    predicted = probabilities.argmax(axis=1)
    output = pd.DataFrame(
        {
            "sample_row": np.arange(len(gene_frame)),
            "predicted_class": [
                bundle.label_names[index] for index in predicted
            ],
            "confidence": probabilities.max(axis=1),
        }
    )
    for index, label in enumerate(bundle.label_names):
        output[f"P({label})"] = probabilities[:, index]
    return output


def pathway_importance(bundle: ModelBundle) -> pd.DataFrame:
    bundle.model.eval()
    with torch.no_grad():
        attention = bundle.model.pathway_attention()
    records: list[dict[str, float | str]] = []
    for branch_name, config_key in (
        ("Gene expression", "gene_pathways"),
        ("DNA methylation", "dna_pathways"),
    ):
        key = "gene_expression" if branch_name == "Gene expression" else "dna_methylation"
        weights = attention[key].cpu().numpy()
        pathways = bundle.config[config_key]
        peak_attention = weights.max(axis=0)
        entropy = -(
            weights * np.log(np.clip(weights, 1e-12, None))
        ).sum(axis=0)
        for pathway, peak, entropy_value in zip(
            pathways, peak_attention, entropy, strict=True
        ):
            records.append(
                {
                    "branch": branch_name,
                    "pathway": pathway,
                    "peak_gene_attention": float(peak),
                    "attention_entropy": float(entropy_value),
                }
            )
    return (
        pd.DataFrame(records)
        .sort_values("peak_gene_attention", ascending=False)
        .reset_index(drop=True)
    )