Module

Module#

The deepspatial.module serves as the training and inference orchestrator. It encapsulates joint continuous Flow Matching, modality-normalized velocity MSE losses, Exponential Moving Average (EMA) updates, and continuous integration solvers. Spatial coordinates, gene expression, and one-hot cell-type states are interpolated along the same linear path; no cross-entropy loss is used for cell-type evolution.

Lightning Module#

The core engine responsible for managing the optimization lifecycle and the generation phase.

DeepSpatialModule

DeepSpatial Module for Training & Inference.

Key Methods#

Fundamental operations managed by the module. sample(..., reverse_time=True) performs explicit reverse-time ODE integration through the equivalent increasing-time parameterization $q(u)=s(1-u)$ and $dq/du=-v(q,1-u)$. Original depth endpoints and the forward-oriented delta_z remain unchanged. Reverse-time sampling is supported for ODE mode only.

DeepSpatialModule.configure_optimizers

Initialize AdamW optimizer with weight decay.

DeepSpatialModule.training_step

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

DeepSpatialModule.sample

Integrates the learned flow field to reconstruct intermediate biological states.