Upload mulgit/causal.py with huggingface_hub

mulgit/causal.py

@@ -0,0 +1,307 @@
+"""
+Causal Discovery Module
+
+Identifies causal genetic factors and molecular interactions underlying
+exceptional longevity. Combines structural causal models with deep 
+learning-based causal inference.
+
+Methods implemented:
+  1. Causal Feature Selection via Information Bottleneck (Seq2Exp-inspired)
+     - Learn binary masks that identify causal features from each omics layer
+     - Beta distribution prior for sparsity
+  
+  2. Causal Structure Learning via NOTEARS-inspired DAG constraint
+     - Learn causal graph between molecular features
+     - Differentiable acyclicity constraint
+  
+  3. Causal Mediation Analysis
+     - Identify mediated effects through the central dogma layers
+     - Decompose total effect into direct and indirect (pathway-mediated)
+
+References:
+  - Seq2Exp (arxiv:2502.13991): Causal regulatory element discovery
+  - Avici: Amortized causal structure learning in genomics
+  - NOTEARS: Non-combinatorial Optimization via Trace Exponential 
+    Augmented lagRangian Structure learning
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, List, Tuple, Dict
+
+
+# ─── Causal Feature Selection ────────────────────────────────────────────────
+
+class CausalFeatureMask(nn.Module):
+    """
+    Learns a binary mask over input features identifying causal features.
+    
+    Inspired by Seq2Exp's information bottleneck: uses Beta distribution
+    prior to encourage sparsity. The mask is learned via the 
+    concrete/Gumbel-softmax reparameterization for differentiability.
+    """
+    
+    def __init__(
+        self,
+        num_features: int,
+        prior_alpha: float = 0.1,
+        prior_beta: float = 0.9,
+        temperature: float = 0.5,
+    ):
+        """
+        Args:
+          num_features: number of input features
+          prior_alpha, prior_beta: Beta distribution parameters (skewed
+            toward 0 to encourage sparse selection)
+          temperature: Gumbel-softmax temperature (lower = more discrete)
+        """
+        super().__init__()
+        # Learnable logits for each feature's selection probability
+        self.logit_p = nn.Parameter(torch.zeros(num_features))
+        self.prior_alpha = prior_alpha
+        self.prior_beta = prior_beta
+        self.temperature = temperature
+
+    def forward(self, training: bool = True) -> torch.Tensor:
+        """
+        Returns a soft (training) or hard (inference) binary mask.
+        """
+        if training:
+            # Concrete distribution (Gumbel-softmax)
+            u = torch.rand_like(self.logit_p)
+            gumbel = -torch.log(-torch.log(u + 1e-8) + 1e-8)
+            logits = (self.logit_p + gumbel) / self.temperature
+            mask = torch.sigmoid(logits)
+        else:
+            # Hard threshold at 0.5
+            mask = (torch.sigmoid(self.logit_p) > 0.5).float()
+        return mask
+
+    def sparsity_loss(self) -> torch.Tensor:
+        """
+        Sparsity regularization: penalize large selection probabilities.
+        Uses L1 norm of sigmoid(logit) to encourage zeros.
+        """
+        p = torch.sigmoid(self.logit_p)
+        # L1 penalty: encourages p → 0 for non-causal features
+        return p.mean()
+
+
+class CausalOmicsSelector(nn.Module):
+    """
+    Per-modality causal feature selection.
+    Selects which features from each omics layer are causal for the outcome.
+    """
+    
+    def __init__(
+        self,
+        modality_dims: Dict[str, int],
+        prior_alpha: float = 0.1,
+        prior_beta: float = 0.9,
+    ):
+        super().__init__()
+        self.masks = nn.ModuleDict({
+            name: CausalFeatureMask(dim, prior_alpha, prior_beta)
+            for name, dim in modality_dims.items()
+        })
+        self.modality_dims = modality_dims
+
+    def forward(
+        self, modalities: Dict[str, torch.Tensor], training: bool = True
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Apply causal masks to each modality.
+        
+        Returns:
+          selected: masked features per modality
+          masks: learned masks per modality
+        """
+        selected = {}
+        masks = {}
+        for name, x in modalities.items():
+            mask = self.masks[name](training=training)
+            selected[name] = x * mask.unsqueeze(0)  # broadcast over batch
+            masks[name] = mask
+        return selected, masks
+
+    def total_sparsity_loss(self) -> torch.Tensor:
+        """Sum of sparsity losses across all modalities."""
+        return sum(self.masks[name].sparsity_loss() for name in self.masks)
+
+
+# ─── Causal Graph Structure Learning ────────────────────────────────────────
+
+class CausalGraphLearner(nn.Module):
+    """
+    Learns a causal graph (DAG) between a set of latent variables using
+    a differentiable acyclicity constraint (NOTEARS-inspired).
+    
+    Adapted for molecular features: the learned adjacency matrix represents
+    causal relationships between latent molecular representations.
+    """
+    
+    def __init__(
+        self,
+        num_variables: int,
+        hidden_dim: int = 64,
+        lambda_dag: float = 1.0,
+    ):
+        """
+        Args:
+          num_variables: number of variables in the causal graph
+          hidden_dim: dimension of each variable's representation
+          lambda_dag: weight for the DAG constraint
+        """
+        super().__init__()
+        # Learnable adjacency matrix (causal strengths)
+        self.W = nn.Parameter(torch.zeros(num_variables, num_variables))
+        self.num_variables = num_variables
+        self.lambda_dag = lambda_dag
+        nn.init.xavier_normal_(self.W)
+
+    def forward(self) -> torch.Tensor:
+        """Returns the learned weighted adjacency matrix."""
+        return self.W
+
+    def dag_constraint(self) -> torch.Tensor:
+        """
+        Differentiable DAG constraint (NOTEARS formulation).
+        trace(exp(W * W)) - d = 0 iff W is a DAG.
+        """
+        W = self.W * self.W  # element-wise square for non-negativity
+        M = torch.matrix_exp(W)  # matrix exponential
+        h = torch.trace(M) - self.num_variables
+        return h * h  # squared to ensure non-negative loss
+
+    def causal_effects(self) -> torch.Tensor:
+        """
+        Compute total causal effects using the learned adjacency.
+        For linear SEM: total effect = (I - W)^(-1)
+        """
+        W = self.W
+        I = torch.eye(self.num_variables, device=W.device)
+        total_effects = torch.linalg.inv(I - W)
+        return total_effects
+
+
+# ─── Mediation Analysis ─────────────────────────────────────────────────────
+
+class CausalMediationAnalyzer(nn.Module):
+    """
+    Analyzes causal mediation through the central dogma layers.
+    
+    For the path DNA → RNA → Protein → Phenotype, decomposes the total
+    effect of a DNA feature on longevity into:
+      - Direct effect (DNA → Phenotype, bypassing intermediates)
+      - Indirect effects (DNA → RNA → Phenotype, DNA → RNA → Protein → Phenotype)
+    
+    This maps to the MuLGIT central dogma architecture.
+    """
+    
+    def __init__(
+        self,
+        dna_dim: int,
+        rna_dim: int,
+        protein_dim: int,
+    ):
+        super().__init__()
+        # Path-specific coefficients
+        self.dna_to_phenotype = nn.Linear(dna_dim, 1, bias=False)      # direct
+        self.dna_to_rna = nn.Linear(dna_dim, rna_dim, bias=False)       # path 1
+        self.rna_to_phenotype = nn.Linear(rna_dim, 1, bias=False)       # path 2
+        self.rna_to_protein = nn.Linear(rna_dim, protein_dim, bias=False) # path 3
+        self.protein_to_phenotype = nn.Linear(protein_dim, 1, bias=False) # path 4
+
+    def decompose_effect(
+        self,
+        dna_features: torch.Tensor,
+        rna_features: Optional[torch.Tensor] = None,
+        protein_features: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Decompose total effect into direct and pathway-mediated effects.
+        
+        Returns dict with:
+          total_effect: combined effect on phenotype
+          direct_effect: DNA → Phenotype (bypassing RNA/protein)
+          dna_rna_effect: DNA → RNA → Phenotype
+          dna_rna_protein_effect: DNA → RNA → Protein → Phenotype
+        """
+        direct = self.dna_to_phenotype(dna_features)
+        
+        rna_pred = self.dna_to_rna(dna_features)
+        rna_effect = self.rna_to_phenotype(rna_pred)
+        
+        protein_pred = self.rna_to_protein(rna_pred)
+        protein_effect = self.protein_to_phenotype(protein_pred)
+        
+        total = direct + rna_effect + protein_effect
+        
+        return {
+            "total_effect": total,
+            "direct_effect": direct,
+            "dna_to_rna_effect": rna_effect,
+            "dna_to_rna_to_protein_effect": protein_effect,
+        }
+
+
+# ─── Causal Attribution ─────────────────────────────────────────────────────
+
+def compute_feature_attribution(
+    model: nn.Module,
+    input_modalities: Dict[str, torch.Tensor],
+    target: int = 0,
+    n_steps: int = 20,
+) -> Dict[str, torch.Tensor]:
+    """
+    Integrated Gradients-style causal attribution.
+    
+    Computes the contribution of each feature to the predicted risk score
+    by integrating gradients along the path from baseline (zero) to input.
+    """
+    attributions = {}
+    
+    for name, x in input_modalities.items():
+        baseline = torch.zeros_like(x)
+        integrated_grad = torch.zeros_like(x)
+        
+        for alpha in torch.linspace(0, 1, n_steps):
+            interpolated = baseline + alpha * (x - baseline)
+            interpolated.requires_grad_(True)
+            
+            # Construct full input dict
+            full_input = {k: v for k, v in input_modalities.items()}
+            full_input[name] = interpolated
+            
+            # Forward pass
+            output = model(**full_input)
+            risk = output["risk"]
+            
+            # Gradient of risk w.r.t. interpolated input
+            grad = torch.autograd.grad(risk.sum(), interpolated)[0]
+            integrated_grad += grad.detach()
+        
+        # Average and multiply by (input - baseline)
+        attributions[name] = (x - baseline) * (integrated_grad / n_steps)
+    
+    return attributions
+
+
+def identify_causal_features(
+    attributions: Dict[str, torch.Tensor],
+    top_k: int = 100,
+) -> Dict[str, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Identify top causal features from attribution scores.
+    
+    Returns dict mapping modality name to (top_indices, top_scores).
+    """
+    results = {}
+    for name, attr in attributions.items():
+        # Average attribution across batch
+        mean_attr = attr.abs().mean(dim=0)
+        # Get top-k features
+        top_scores, top_indices = torch.topk(mean_attr, k=min(top_k, mean_attr.shape[0]))
+        results[name] = (top_indices, top_scores)
+    return results