Spaces:
Sleeping
Sleeping
File size: 5,016 Bytes
914512c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """bioai.models.pinn_fate -- Physics-Informed Neural Network for dsRNA environmental fate.
Predicts the first-order degradation rate ``k`` of a dsRNA in the field, given
8 environmental + sequence features. The physics constraint is the standard
exponential decay ``C(t) = C0 * exp(-k * t)`` so half-life is ``ln(2) / k``.
Wired into the pipeline (Task D deliverable #5): the ranker uses
``DegradationPINN.half_life(features)`` to penalise candidates whose
half-life is under 6 hours (too short to be effective in the field).
Specs:
Input features (8-dim):
temperature_C, pH, UV_index, GC_content, length,
salinity_ppt, soil_clay_pct, humidity_pct
Output: positive rate k (Softplus ensures positivity)
Network: Linear(8,64) -> Tanh -> Linear(64,32) -> Tanh -> Linear(32,1) -> Softplus
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
# Feature names in the expected order. Used by callers to build feature vectors.
PINN_FEATURE_NAMES = (
"temperature_C", # ambient temperature in Celsius
"pH", # soil/leaf pH
"UV_index", # 0..15 typical daylight UV index
"GC_content", # 0..1 fraction G+C of the dsRNA
"length", # dsRNA length in nt (e.g. 200)
"salinity_ppt", # salinity in parts per thousand
"soil_clay_pct", # 0..100 soil clay percentage
"humidity_pct", # 0..100 relative humidity
)
class DegradationPINN(nn.Module):
"""PINN for dsRNA environmental fate.
Forward pass returns the rate ``k`` (positive). Use :meth:`predict_rate`,
:meth:`predict_concentration`, or :meth:`half_life` for the physics-wrapped
quantities.
"""
LN2 = math.log(2.0)
def __init__(self, feature_dim: int = 8):
super().__init__()
self.feature_dim = feature_dim
self.net = nn.Sequential(
nn.Linear(feature_dim, 64), nn.Tanh(),
nn.Linear(64, 32), nn.Tanh(),
nn.Linear(32, 1), nn.Softplus(), # ensures k > 0
)
# ------------------------------------------------------------------ #
def forward(self, features: torch.Tensor) -> torch.Tensor:
"""Raw forward returns the rate ``k`` (positive). Shape ``(B, 1)``."""
if features.dim() == 1:
features = features.unsqueeze(0)
return self.net(features)
# ------------------------------------------------------------------ #
def predict_rate(self, features: torch.Tensor) -> torch.Tensor:
"""Positive degradation rate ``k`` (1/hours). Shape ``(B, 1)``."""
return self.forward(features)
# ------------------------------------------------------------------ #
def predict_concentration(
self,
C0: torch.Tensor | float,
t_points: torch.Tensor,
features: torch.Tensor,
) -> torch.Tensor:
"""Physics-informed concentration trajectory.
Parameters
----------
C0:
Initial concentration (scalar or ``(B,)`` or ``(B, 1)``).
t_points:
1-D tensor of time points (hours), shape ``(T,)``.
features:
``(B, feature_dim)`` env/sequence features.
Returns
-------
torch.Tensor
Concentration ``C(t) = C0 * exp(-k * t)`` for each sample and
each time point. Shape ``(B, T)``.
"""
k = self.predict_rate(features) # (B, 1)
if not torch.is_tensor(C0):
C0 = torch.tensor(float(C0), dtype=k.dtype, device=k.device)
C0 = C0.view(-1, 1) if C0.dim() == 1 else C0.view(-1, 1)
# Broadcast: (B, 1) * (1, T) -> (B, T)
t = t_points.to(k.device).unsqueeze(0)
return C0 * torch.exp(-k * t)
# ------------------------------------------------------------------ #
def half_life(self, features: torch.Tensor) -> torch.Tensor:
"""Half-life (hours) = ln(2) / k. Shape ``(B, 1)``.
Softplus guarantees ``k > 0`` so the division is well-defined.
"""
k = self.predict_rate(features)
# Clamp k from below to avoid Inf half-lives from a near-zero rate.
k_safe = torch.clamp(k, min=1e-6)
return torch.tensor(self.LN2, dtype=k.dtype, device=k.device) / k_safe
# ------------------------------------------------------------------ #
def physics_consistency_loss(
self,
C0: torch.Tensor | float,
t_points: torch.Tensor,
features: torch.Tensor,
C_observed: torch.Tensor,
) -> torch.Tensor:
"""MSE between the PINN-predicted concentration and observed C(t).
Use this as the "physics-consistency" loss term in training. Combined
with a direct MSE on ``k`` (when rate labels are available), it
enforces that the network both fits the data *and* respects the
exponential-decay ODE.
"""
C_pred = self.predict_concentration(C0, t_points, features)
return nn.functional.mse_loss(C_pred, C_observed)
|