import torch
import torch.nn as nn
import pytorch_lightning as pl
from copy import deepcopy
from collections import OrderedDict
from .transport import create_transport, Sampler
[docs]
class DeepSpatialModule(pl.LightningModule):
"""
DeepSpatial Module for Training & Inference.
Parameters
----------
args : dict or argparse.Namespace
Configuration dictionary containing hyperparameters such as learning rate,
path type, and sampling settings.
model : torch.nn.Module
The core neural network architecture (e.g., the GiT model) that predicts
the velocity fields.
"""
[docs]
def __init__(self, args, model):
"""
Initializes an LightningModule instance for DeepSpatial.
"""
super().__init__()
self.save_hyperparameters(args)
# Core Model
self.model = model
# EMA Setup
self.ema_decay = self.hparams.get('ema_decay', 0.999)
self.ema_model = deepcopy(self.model)
self._freeze(self.ema_model)
# Transport & Path Setup
self.transport = create_transport(
path_type=self.hparams.path_type,
prediction=self.hparams.prediction,
train_eps=self.hparams.get('train_eps', 0.02),
sample_eps=self.hparams.get('sample_eps', 0.02),
)
self.sampler = Sampler(self.transport)
def _freeze(self, module):
"""Freeze model parameters for EMA or evaluation."""
for param in module.parameters():
param.requires_grad = False
module.eval()
# ============================================================
# Training & Validation Logic
# ============================================================
def _shared_step(self, batch):
"""
Computes joint Flow Matching losses across spatial and molecular dimensions.
"""
x0, x1 = batch['x0'], batch['x1']
g0, g1 = batch['g0'], batch['g1']
c0, c1 = batch['c0'], batch['c1']
z0, z1 = batch['z0'], batch['z1']
delta_z = batch['delta_z']
# Sample time steps
t, _, _ = self.transport.sample(x1)
# Plan paths (interpolation)
_, xt, ux_t = self.transport.path_sampler.plan(t, x0, x1)
_, gt, ug_t = self.transport.path_sampler.plan(t, g0, g1)
_, ct, uc_t = self.transport.path_sampler.plan(t, c0, c1)
_, zt, _ = self.transport.path_sampler.plan(t, z0, z1)
# 3. Predict velocity fields
vx_pred, vg_pred, vc_pred = self.model(
xt=xt, gt=gt, t=t, zt=zt, delta_z=delta_z, ct=ct
)
# ``loss_fn`` already averages the squared velocity error over each
# modality's feature dimensions. Averaging the resulting per-sample
# losses therefore gives one, and only one, dimensional normalization.
loss_x = self.transport.loss_fn(vx_pred, x0, xt, t, ux_t).mean()
loss_g = self.transport.loss_fn(vg_pred, g0, gt, t, ug_t).mean()
loss_c = self.transport.loss_fn(vc_pred, c0, ct, t, uc_t).mean()
# Weighted total loss
lambda_g = self.hparams.get('lambda_g', 0.1)
lambda_c = self.hparams.get('lambda_c', 10.0)
loss_total = loss_x + (lambda_g * loss_g) + (lambda_c * loss_c)
return {
'loss': loss_total,
'loss_x': loss_x,
'loss_g': loss_g,
'loss_c': loss_c
}
[docs]
def training_step(self, batch, batch_idx):
loss = self._shared_step(batch)
self.log('loss', loss['loss'], prog_bar=True, on_step=True, on_epoch=True)
self.log('loss_x', loss['loss_x'], on_epoch=True)
self.log('loss_g', loss['loss_g'], on_epoch=True)
self.log('loss_c', loss['loss_c'], on_epoch=True)
return loss
# ============================================================
# EMA Update Logic
# ============================================================
def on_train_batch_end(self, outputs, batch, batch_idx):
"""Update EMA parameters after each optimizer step."""
self._update_ema()
@torch.no_grad()
def _update_ema(self):
"""Update EMA model weights with exponential decay."""
for ema_param, model_param in zip(self.ema_model.parameters(), self.model.parameters()):
ema_param.data.mul_(self.ema_decay).add_(model_param.data, alpha=1 - self.ema_decay)
def on_load_checkpoint(self, checkpoint):
"""Ensure EMA model is also loaded from checkpoint."""
self._update_ema() # Warm start EMA with current loaded model params
# ============================================================
# Inference / Sampling Logic
# ============================================================
[docs]
@torch.no_grad()
def sample(self, batch, mode="ODE", steps=20, reverse_time=False):
"""
Integrates the learned flow field to reconstruct intermediate biological states.
Parameters
----------
batch : dict
A dictionary containing the initial states (`x0`, `g0`, `c0`), the physical
Z-depth conditions (`z0`, `z1`, `delta_z`), and other necessary tensors.
mode : str, optional
The integration mode, either `"ODE"` (Ordinary Differential Equation) or
`"SDE"` (Stochastic Differential Equation). By default `"ODE"`.
steps : int, optional
The number of integration steps from the source to the target slice.
Higher values yield more accurate but slower trajectories. By default 20.
reverse_time : bool, optional
If ``True``, start from the terminal state at physical flow time
``t=1`` and integrate the equivalent reparameterized ODE
``dq/du = -v(q, 1-u)`` for solver time ``u`` increasing from 0 to 1.
The depth endpoints and ``delta_z`` must remain in their original
forward order. By default ``False``.
Returns
-------
dict
A dictionary containing the full integration trajectories:
- `'x_traj'` : torch.Tensor of shape `(Steps, Batch, 2)`
- `'g_traj'` : torch.Tensor of shape `(Steps, Batch, Gene_Dim)`
- `'c_traj_discrete'` : torch.Tensor of shape `(Steps, Batch)` obtained by
applying argmax to the integrated continuous cell-type state.
"""
if reverse_time and mode != "ODE":
raise ValueError("reverse_time=True is only supported for ODE sampling.")
self.ema_model.eval()
# Configure the ODE/SDE sampler based on hyperparameters
sample_config = {
'num_steps': steps,
'sampling_method': self.hparams.get('sampling_method', 'dopri5'),
'atol': self.hparams.get('atol', 1e-5),
'rtol': self.hparams.get('rtol', 1e-5)
}
# Instantiate the integration function
sampler_fn = self.sampler.sample_ode(**sample_config) if mode == "ODE" else self.sampler.sample_sde(**sample_config)
# Initial state is at physical t=0 for forward sampling and at physical
# t=1 for reverse-time sampling.
x0, g0, c0 = batch['x0'], batch['g0'], batch['c0']
x_dim, g_dim = x0.shape[-1], g0.shape[-1]
z0, z1, delta_z = batch['z0'], batch['z1'], batch['delta_z']
# Concatenate for joint integration
init_state = torch.cat([x0, g0, c0], dim=-1)
def velocity_field_wrapper(joint_state_t, solver_t):
# Unpack modalities
xt = joint_state_t[..., :x_dim]
gt = joint_state_t[..., x_dim : x_dim + g_dim]
ct = joint_state_t[..., x_dim + g_dim :]
# The numerical solver always advances from u=0 to u=1. For
# reverse-time sampling, physical flow time is t=1-u.
if torch.is_tensor(solver_t):
solver_t_val = solver_t.item() if solver_t.dim() == 0 else solver_t[0].item()
else:
solver_t_val = solver_t
model_t_val = 1.0 - solver_t_val if reverse_time else solver_t_val
model_t = torch.full((xt.shape[0],), model_t_val, device=xt.device, dtype=xt.dtype)
# z0, z1, and delta_z retain their original forward orientation.
_, zt, _ = self.transport.path_sampler.plan(model_t, z0, z1)
# Forward pass through EMA model
vx, vg, vc = self.ema_model(
xt=xt, gt=gt, t=model_t, zt=zt, delta_z=delta_z, ct=ct
)
velocity = torch.cat([vx, vg, vc], dim=-1)
# q(u)=s(1-u), hence dq/du=-v(q, 1-u).
return -velocity if reverse_time else velocity
# Compute trajectory: Shape [steps, batch, dim]
trajectory = sampler_fn(init_state, velocity_field_wrapper)
if isinstance(trajectory, (list, tuple)):
trajectory = torch.stack(trajectory, dim=0)
# Unpack integrated results
x_traj = trajectory[..., :x_dim]
g_traj = trajectory[..., x_dim : x_dim + g_dim]
c_traj_cont = trajectory[..., x_dim + g_dim :]
return {
'x_traj': x_traj,
'g_traj': g_traj,
'c_traj_discrete': torch.argmax(c_traj_cont, dim=-1)
}
