deepspatial.module.DeepSpatialModule.sample

deepspatial.module.DeepSpatialModule.sample#

DeepSpatialModule.sample(batch, mode='ODE', steps=20, reverse_time=False)[source]#

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:

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.

Return type:

dict

Contents