Spatial Mouse Placentation Run

Spatial Mouse Placentation Quick Run¶

1. This notebook provides a compact, end-to-end workflow for Spatial Mouse Placentation based on the main STEER quickstart pipeline. The current settings match the configuration used for the main figure. We set training_mode="full" to leverage the full capacity of the model. The complete run may take about 2 hours on a single NVIDIA L20 GPU (with approximately 40 GB of GPU memory).

2. If GPU memory is limited, velo_batch_size is the key parameter for reducing memory usage. During kinetics training, it randomly samples batches of cells. Empirically, a velocity loss below 0.1 usually indicates that the model has trained well.

All main-figure results associated with this project are available on Zenodo: 10.5281/zenodo.18713189.

Input¶

The recommended input is an AnnData object containing spliced, unspliced, and spatial coordinates. If the file only provides obsm['spatial'], the preparation step below will standardize it to X_spatial automatically.

In [1]:

import os
import random
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import scvelo as scv
import seaborn as sns
import torch
import torch.backends.cudnn as cudnn

import steer
from steer.prior.prior import PriorInferenceManager

warnings.filterwarnings("ignore")
%matplotlib inline

print("=== STEER Quick Start Environment ===")
print(f"PyTorch version: {torch.__version__}")
print(f"Scanpy version: {sc.__version__}")
print(f"scVelo version: {scv.__version__}")
print(f"STEER version: {getattr(steer, '__version__', 'dev build')}")
print("====================================")

import os import random import warnings import matplotlib.pyplot as plt import numpy as np import pandas as pd import scanpy as sc import scvelo as scv import seaborn as sns import torch import torch.backends.cudnn as cudnn import steer from steer.prior.prior import PriorInferenceManager warnings.filterwarnings("ignore") %matplotlib inline print("=== STEER Quick Start Environment ===") print(f"PyTorch version: {torch.__version__}") print(f"Scanpy version: {sc.__version__}") print(f"scVelo version: {scv.__version__}") print(f"STEER version: {getattr(steer, '__version__', 'dev build')}") print("====================================")

=== STEER Quick Start Environment ===
PyTorch version: 2.8.0+cu128
Scanpy version: 1.10.2
scVelo version: 0.3.3
STEER version: 2.2.2
====================================

Config¶

1. This section includes only the settings that users are most likely to modify. The current default setup uses training_mode="full", which matches the configuration used to generate the results shown in this notebook.

• "full": uses the default STEER training arguments without overrides and runs the complete training schedule.

• "fast": uses a lightweight, quickstart-style parameter set for faster demos and iteration.

• For advanced usage, you can further customize the "fast" setting in get_stage1_training_overrides and get_stage2_training_overrides. For example, you may remove the epochs and patience settings to fall back to the same defaults as "full" mode, while still adjusting velo_batch_size to reduce GPU memory usage. This can provide more thorough training than the default "fast" mode while requiring less GPU memory than "full" mode.

2. Only the most important settings are exposed here:

• data_name / input_dir / result_dir: control input and output locations

• For a first run on a new dataset, it is usually enough to update the data path and training_mode.

In [2]:

class Config:
    # Data
    data_name = "E8.5_S1.MPSTA"
    input_dir = "/nvme/users/liuzhy/Review_Files/6_Stereoseq_MousePlacentation"
    result_dir = "./results"
    seed = 618

    # Core model and graph settings
    expert = 10  # Number of experts. You can adjust this based on biological prior knowledge
                 # or an estimated cluster number (see quickstart.ipynb).
    smooth_neigh = 50  # Number of neighbors used for smoothing.
                       # Usually 50-100 works well; increase for noisier data.
    spatial_neighbors = 8  # Number of spatial neighbors. We recommend keeping this below 30.
    
    # Training profile
    training_mode = "full"  # Choose from "full" or "fast".
    # finetune_epochs = 5000  # Only used when training_mode == "fast".

    # Advanced settings
    graph = "union"  # Combine expression and spatial graphs to define cellular context.
    corr_mode = "u"  # Use unspliced counts for temporal supervision.
                     # If time inference is unsatisfactory and unspliced counts are low,
                     # try setting this to "s" and rerun.
    neighbor_metric = "cosine"  # Neighbor metric. You can also use Euclidean distance.
    use_us = True  # Use both unspliced and spliced counts as input features.
    use_filter = True  # Filter genes before kinetics learning to keep informative genes.

    # Time constraint
    # Default setting: constrain time inference within each expert.
    fine_method = "none"
    target_size = 300
    direction_base = "expert"

    # Optional setting for fewer experts (<5):
    # use a more flexible constraint based on fine-grained clusters.
    # fine_method = "hierarchical"
    # target_size = 300
    # direction_base = "fine_cluster"

cfg = Config()

INPUT_FILE = os.path.join(cfg.input_dir, f"{cfg.data_name}.h5ad")
RESULT_PATH = os.path.join(cfg.result_dir, f"{cfg.data_name}_quickstart")
os.makedirs(RESULT_PATH, exist_ok=True)


def get_stage1_training_overrides(mode: str) -> dict:
    if mode == "full":
        return {}
    if mode == "fast":
        return {
            "expert_mode": "slim",
            "pretrain_epochs": 500,
            "cluster_epochs": 200,
        }
    raise ValueError(f"Unsupported training_mode: {mode}")


def get_stage2_training_overrides(mode: str, finetune_epochs=None) -> dict:
    if mode == "full":
        return {}
    if mode == "fast":
        if finetune_epochs is None:
            raise ValueError("cfg.finetune_epochs must be set when training_mode='fast'.")
        return {
            "expert_mode": "slim",
            "pretrain_epochs": 500,
            "cluster_epochs": 200,
            "velo_batch_size": 512,
            "MIN_IMPRO": 0.01,
            "PATIENCE": 100,
            "num_epochs": finetune_epochs,
        }
    raise ValueError(f"Unsupported training_mode: {mode}")


print(f"Input file: {INPUT_FILE}")
print(f"Output directory: {RESULT_PATH}")
print(f"Training mode: {cfg.training_mode}")

class Config: # Data data_name = "E8.5_S1.MPSTA" input_dir = "/nvme/users/liuzhy/Review_Files/6_Stereoseq_MousePlacentation" result_dir = "./results" seed = 618 # Core model and graph settings expert = 10 # Number of experts. You can adjust this based on biological prior knowledge # or an estimated cluster number (see quickstart.ipynb). smooth_neigh = 50 # Number of neighbors used for smoothing. # Usually 50-100 works well; increase for noisier data. spatial_neighbors = 8 # Number of spatial neighbors. We recommend keeping this below 30. # Training profile training_mode = "full" # Choose from "full" or "fast". # finetune_epochs = 5000 # Only used when training_mode == "fast". # Advanced settings graph = "union" # Combine expression and spatial graphs to define cellular context. corr_mode = "u" # Use unspliced counts for temporal supervision. # If time inference is unsatisfactory and unspliced counts are low, # try setting this to "s" and rerun. neighbor_metric = "cosine" # Neighbor metric. You can also use Euclidean distance. use_us = True # Use both unspliced and spliced counts as input features. use_filter = True # Filter genes before kinetics learning to keep informative genes. # Time constraint # Default setting: constrain time inference within each expert. fine_method = "none" target_size = 300 direction_base = "expert" # Optional setting for fewer experts (<5): # use a more flexible constraint based on fine-grained clusters. # fine_method = "hierarchical" # target_size = 300 # direction_base = "fine_cluster" cfg = Config() INPUT_FILE = os.path.join(cfg.input_dir, f"{cfg.data_name}.h5ad") RESULT_PATH = os.path.join(cfg.result_dir, f"{cfg.data_name}_quickstart") os.makedirs(RESULT_PATH, exist_ok=True) def get_stage1_training_overrides(mode: str) -> dict: if mode == "full": return {} if mode == "fast": return { "expert_mode": "slim", "pretrain_epochs": 500, "cluster_epochs": 200, } raise ValueError(f"Unsupported training_mode: {mode}") def get_stage2_training_overrides(mode: str, finetune_epochs=None) -> dict: if mode == "full": return {} if mode == "fast": if finetune_epochs is None: raise ValueError("cfg.finetune_epochs must be set when training_mode='fast'.") return { "expert_mode": "slim", "pretrain_epochs": 500, "cluster_epochs": 200, "velo_batch_size": 512, "MIN_IMPRO": 0.01, "PATIENCE": 100, "num_epochs": finetune_epochs, } raise ValueError(f"Unsupported training_mode: {mode}") print(f"Input file: {INPUT_FILE}") print(f"Output directory: {RESULT_PATH}") print(f"Training mode: {cfg.training_mode}")

Input file: /nvme/users/liuzhy/Review_Files/6_Stereoseq_MousePlacentation/E8.5_S1.MPSTA.h5ad
Output directory: ./results/E8.5_S1.MPSTA_quickstart
Training mode: full

Seed And Device¶

This cell fixes the random seed for reproducibility and automatically selects GPU or CPU.

In [3]:

def setup_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    cudnn.deterministic = True
    cudnn.benchmark = False

setup_seed(cfg.seed)
# you need to use your device number
if torch.cuda.is_available():
    device = torch.device("cuda:0")
print(f"Using device: {device}")

def setup_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) cudnn.deterministic = True cudnn.benchmark = False setup_seed(cfg.seed) # you need to use your device number if torch.cuda.is_available(): device = torch.device("cuda:0") print(f"Using device: {device}")

Using device: cuda:2

Load Data¶

Load the .h5ad file and check that the required layers and spatial coordinates are present. When users switch to their own data, this is usually the first place where formatting issues appear.

In [4]:

adata = sc.read_h5ad(INPUT_FILE)
adata = steer.prepare_spatial_adata(
    adata,
    obs_copy_map={"celltype": "celltype1"},
    na_clear_keys=["subregion"],
    celltype_key="celltype",
    remove_celltypes=[],
)

print(adata)
print("\nAvailable layers:", list(adata.layers.keys()))
print("Available obsm keys:", list(adata.obsm.keys()))
print(f"Number of cells: {adata.n_obs}")
print(f"Number of genes: {adata.n_vars}")

required_layers = ["spliced", "unspliced"]
missing_layers = [layer for layer in required_layers if layer not in adata.layers]
if missing_layers:
    raise ValueError(f"Missing required layers: {missing_layers}")
if "X_spatial" not in adata.obsm:
    raise ValueError("Spatial coordinates were not standardized into adata.obsm['X_spatial'].")

adata = sc.read_h5ad(INPUT_FILE) adata = steer.prepare_spatial_adata( adata, obs_copy_map={"celltype": "celltype1"}, na_clear_keys=["subregion"], celltype_key="celltype", remove_celltypes=[], ) print(adata) print("\nAvailable layers:", list(adata.layers.keys())) print("Available obsm keys:", list(adata.obsm.keys())) print(f"Number of cells: {adata.n_obs}") print(f"Number of genes: {adata.n_vars}") required_layers = ["spliced", "unspliced"] missing_layers = [layer for layer in required_layers if layer not in adata.layers] if missing_layers: raise ValueError(f"Missing required layers: {missing_layers}") if "X_spatial" not in adata.obsm: raise ValueError("Spatial coordinates were not standardized into adata.obsm['X_spatial'].")

AnnData object with n_obs × n_vars = 13258 × 17646
    obs: 'x', 'y', 'sample', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt', 'subregion', 'interface', 'celltype1', 'subclusters_5_last', 'subregion5', 'celltype'
    var: 'n_cells_by_counts'
    obsm: 'X_pca', 'X_umap', 'spatial', 'X_spatial'
    layers: 'ambiguous', 'counts', 'matrix', 'spliced', 'unspliced'

Available layers: ['ambiguous', 'counts', 'matrix', 'spliced', 'unspliced']
Available obsm keys: ['X_pca', 'X_umap', 'spatial', 'X_spatial']
Number of cells: 13258
Number of genes: 17646

In [5]:

adata

adata

Out[5]:

AnnData object with n_obs × n_vars = 13258 × 17646
    obs: 'x', 'y', 'sample', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt', 'subregion', 'interface', 'celltype1', 'subclusters_5_last', 'subregion5', 'celltype'
    var: 'n_cells_by_counts'
    obsm: 'X_pca', 'X_umap', 'spatial', 'X_spatial'
    layers: 'ambiguous', 'counts', 'matrix', 'spliced', 'unspliced'

Preprocess¶

Run the basic normalization step and build the dataframe, adjacency matrix, and processed AnnData object required by STEER.

In [6]:

scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=2000)

df, adjacency_matrix, adata = steer.preprocess_anndata_spatial(
    adata,
    npc=30,
    NUM_AD_NEIGH=30,
    spatial_neighbors=cfg.spatial_neighbors,
    SMOOTH_NEIGH=cfg.smooth_neigh,
    moments_adj=True,
    combine_mode=cfg.graph,
    neighbor_metric=cfg.neighbor_metric,
    spatial_key="X_spatial",
    use_us=cfg.use_us,
)

print(df.head())
print("\nAdjacency matrix shape:", adjacency_matrix.shape)

scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=2000) df, adjacency_matrix, adata = steer.preprocess_anndata_spatial( adata, npc=30, NUM_AD_NEIGH=30, spatial_neighbors=cfg.spatial_neighbors, SMOOTH_NEIGH=cfg.smooth_neigh, moments_adj=True, combine_mode=cfg.graph, neighbor_metric=cfg.neighbor_metric, spatial_key="X_spatial", use_us=cfg.use_us, ) print(df.head()) print("\nAdjacency matrix shape:", adjacency_matrix.shape)

Filtered out 12279 genes that are detected 20 counts (shared).
Normalized count data: X, spliced, unspliced.
Extracted 2000 highly variable genes.
Logarithmized X.

/nvme/users/liuzhy/miniconda3/envs/steer_dev/lib/python3.10/site-packages/scvelo/preprocessing/utils.py:705: DeprecationWarning: `log1p` is deprecated since scVelo v0.3.0 and will be removed in a future version. Please use `log1p` from `scanpy.pp` instead.
  log1p(adata)

computing moments based on connectivities
    finished (0:00:03) --> added 
    'Ms' and 'Mu', moments of un/spliced abundances (adata.layers)
        cellID gene_name  unsplice    splice  orig_unsplice  orig_splice  \
0   11450_9600     Syap1  0.000000  0.162575            0.0     0.000000   
1   11550_8900     Syap1  0.236257  0.237767            0.0     0.000000   
2   12750_9650     Syap1  0.181738  0.190650            0.0     1.941131   
3  10850_10150     Syap1  0.000000  0.069758            0.0     0.000000   
4  13000_12500     Syap1  0.103906  0.080809            0.0     0.000000   

         Mu        Ms  
0  0.000000  0.104713  
1  0.029140  0.153144  
2  0.022416  0.122796  
3  0.000000  0.044931  
4  0.012816  0.052049  

Adjacency matrix shape: (13258, 13258)

Expert Number¶

The expert number used in this workflow was originally chosen from an estimation step rather than fixed arbitrarily. For reproduce the results, cfg.expert is directly set to 10, and the optional helper below can be used to re-estimate the expert number directly from the processed embedding before training.

In [7]:

cfg.expert = 10

cfg.expert = 10

Build PyG Input¶

Package the current dataset into the PyTorch Geometric objects required by STEER training.

In [8]:

dataset = steer.preload_datasets_all_genes_anndata(df=df, MODEL_MODE="pretrain", adata=adata)
pyg_data = steer.create_pyg_data(dataset, adjacency_matrix, normalize=True)

print(pyg_data)

dataset = steer.preload_datasets_all_genes_anndata(df=df, MODEL_MODE="pretrain", adata=adata) pyg_data = steer.create_pyg_data(dataset, adjacency_matrix, normalize=True) print(pyg_data)

Data(x=[13258, 4000], edge_index=[2, 593116], type_features=[13258, 2000], orig_features=[13258, 4000], cell_ids=[13258], adj=[13258, 13258])

Stage 1 Training¶

Train the first-stage model to obtain the initial representation and cluster structure.

• full mode: do not override training epochs or related controls, so STEER uses its built-in defaults
• fast mode: apply a lightweight quickstart-style configuration for a faster tutorial run

In [9]:

stage1_kwargs = dict(
    device=device,
    device2=device,
    pyg_data=pyg_data,
    MODEL_MODE="pretrain",
    adata=adata,
    NUM_LOSS_NEIGH=30,
    corr_mode=cfg.corr_mode,
    max_n_cluster=cfg.expert,
    path=RESULT_PATH,
)
stage1_kwargs.update(get_stage1_training_overrides(cfg.training_mode))

result_adata = steer.model_training_share_neighbor_adata(**stage1_kwargs)

stage1_kwargs = dict( device=device, device2=device, pyg_data=pyg_data, MODEL_MODE="pretrain", adata=adata, NUM_LOSS_NEIGH=30, corr_mode=cfg.corr_mode, max_n_cluster=cfg.expert, path=RESULT_PATH, ) stage1_kwargs.update(get_stage1_training_overrides(cfg.training_mode)) result_adata = steer.model_training_share_neighbor_adata(**stage1_kwargs)

Epoch 0, Loss 0.2706538140773773
Epoch 50, Loss 0.15260574221611023
Epoch 100, Loss 0.1374153047800064
Epoch 150, Loss 0.12543688714504242
Epoch 200, Loss 0.11626759171485901
Epoch 250, Loss 0.10848373174667358
Epoch 300, Loss 0.100981704890728
Epoch 350, Loss 0.09443806111812592
Epoch 400, Loss 0.08860418945550919
Epoch 450, Loss 0.08336362987756729
Epoch 500, Loss 0.07843828201293945
Epoch 550, Loss 0.07374273985624313
Epoch 600, Loss 0.06948872655630112
Epoch 650, Loss 0.06516654044389725
Epoch 700, Loss 0.06151680275797844
Epoch 750, Loss 0.058430690318346024
Epoch 800, Loss 0.05520620942115784
Epoch 850, Loss 0.05212832987308502
Epoch 900, Loss 0.04941016063094139
Epoch 950, Loss 0.04706642031669617
Epoch 1000, Loss 0.044855229556560516
Epoch 1050, Loss 0.042713817209005356
Epoch 1100, Loss 0.04085223376750946
Epoch 1150, Loss 0.03914935886859894
Epoch 1200, Loss 0.03755848854780197
Epoch 1250, Loss 0.0359717532992363
Epoch 1300, Loss 0.0346527025103569
Epoch 1350, Loss 0.033247727900743484
Epoch 1400, Loss 0.03166535869240761
Epoch 1450, Loss 0.030503254383802414
Epoch 1500, Loss 1.1945350170135498
Epoch 1550, Loss 0.3980165123939514
Epoch 1600, Loss 0.13172152638435364
Epoch 1650, Loss 0.1212276890873909
Epoch 1700, Loss 0.11661934852600098
Epoch 1750, Loss 0.11365920305252075
Epoch 1800, Loss 0.11185658723115921
Epoch 1850, Loss 0.10991236567497253
Epoch 1900, Loss 0.10799113661050797
Epoch 1950, Loss 0.10684845596551895
Epoch 2000, Loss 0.10494072735309601
Epoch 2050, Loss 0.10388772934675217
Epoch 2100, Loss 0.10310031473636627
Epoch 2150, Loss 0.10278385877609253
Epoch 2200, Loss 0.10520020872354507
Epoch 2250, Loss 0.10025681555271149
Epoch 2300, Loss 0.099912628531456
Epoch 2350, Loss 0.09965579211711884
Epoch 2400, Loss 0.09919999539852142
Epoch 2450, Loss 0.09795272350311279
Epoch 2500, Loss 0.09779015183448792
Epoch 2550, Loss 0.09711004793643951
Epoch 2600, Loss 0.09680501371622086
Epoch 2650, Loss 0.0965099036693573
Epoch 2700, Loss 0.09606146812438965
Epoch 2750, Loss 0.09554947912693024
Epoch 2800, Loss 0.095319002866745
Epoch 2850, Loss 0.09526750445365906
Epoch 2900, Loss 0.09512843936681747
Epoch 2950, Loss 0.09434902667999268

Optional mclust¶

If R and rpy2 are available, this step runs mclust once for optional clustering refinement. Otherwise, it is skipped automatically.

In [10]:

try:
    result_adata = steer.mclust_R(result_adata, num_cluster=cfg.expert)
    print("Initial clustering with mclust completed.")
except Exception as e:
    print("mclust step skipped.")
    print(f"Reason: {e}")

torch.cuda.empty_cache()

try: result_adata = steer.mclust_R(result_adata, num_cluster=cfg.expert) print("Initial clustering with mclust completed.") except Exception as e: print("mclust step skipped.") print(f"Reason: {e}") torch.cuda.empty_cache()

R[write to console]:                    __           __ 
   ____ ___  _____/ /_  _______/ /_
  / __ `__ \/ ___/ / / / / ___/ __/
 / / / / / / /__/ / /_/ (__  ) /_  
/_/ /_/ /_/\___/_/\__,_/____/\__/   version 6.1.1
Type 'citation("mclust")' for citing this R package in publications.

/nvme/users/liuzhy/miniconda3/envs/steer_dev/lib/python3.10/site-packages/rpy2/robjects/numpy2ri.py:241: DeprecationWarning: The global conversion available with activate() is deprecated and will be removed in the next major release. Use a local converter.
  warnings.warn('The global conversion available with activate() '

fitting ...
  |======================================================================| 100%
Initial clustering with mclust completed.

Prior Inference¶

Infer the kinetic prior from the first-stage result for use in the downstream kinetic-learning stage.

In [11]:

prior_manager = PriorInferenceManager(result_adata, df, RESULT_PATH, seed=cfg.seed)
prior_manager.task1_define_fine_clusters(method=cfg.fine_method, target_size=cfg.target_size)
prior_manager.task2_filter_genes(based_on="expert", keep_ngene=1000, use_filter=cfg.use_filter)
prior_manager.task3_calc_convexity(based_on=cfg.direction_base)

result_adata, df_updated, fine_clus_vec_np = prior_manager.finalize_for_training()
print("Prior inference completed.")

prior_manager = PriorInferenceManager(result_adata, df, RESULT_PATH, seed=cfg.seed) prior_manager.task1_define_fine_clusters(method=cfg.fine_method, target_size=cfg.target_size) prior_manager.task2_filter_genes(based_on="expert", keep_ngene=1000, use_filter=cfg.use_filter) prior_manager.task3_calc_convexity(based_on=cfg.direction_base) result_adata, df_updated, fine_clus_vec_np = prior_manager.finalize_for_training() print("Prior inference completed.")

--- Task 1: Generating Fine Clusters (Method: none) ---
  -> Method is 'none', copying Expert clusters to Fine clusters.
--- Task 2: Filtering Genes (Based on: EXPERT, Keep: 1000) ---
--- Task 3: Calculating Direction (Based on: EXPERT) ---
Starting parallel processing with n_jobs=-1...

[Parallel(n_jobs=-1)]: Using backend LokyBackend with 128 concurrent workers.
[Parallel(n_jobs=-1)]: Done  32 tasks      | elapsed:   18.3s
[Parallel(n_jobs=-1)]: Done 194 tasks      | elapsed:   19.7s
[Parallel(n_jobs=-1)]: Done 392 tasks      | elapsed:   20.5s
[Parallel(n_jobs=-1)]: Done 626 tasks      | elapsed:   21.2s
[Parallel(n_jobs=-1)]: Done 896 tasks      | elapsed:   21.7s
[Parallel(n_jobs=-1)]: Done 1202 tasks      | elapsed:   22.3s
[Parallel(n_jobs=-1)]: Done 1544 tasks      | elapsed:   22.9s
[Parallel(n_jobs=-1)]: Done 2000 out of 2000 | elapsed:   23.9s finished

Aggregating results...
Done.
--- Finalizing: Restoring Expert labels & Preparing Fine Cluster Vector ---
Prior inference completed.

Velocity Dataset¶

Subset to velocity genes and rebuild the data objects required for the second training stage.

In [12]:

prior_adata = result_adata.copy()

raw_adata = sc.read_h5ad(INPUT_FILE)
raw_adata = steer.prepare_spatial_adata(
    raw_adata,
    obs_copy_map={"celltype": "celltype1"},
    na_clear_keys=["subregion"],
    celltype_key="celltype",
    remove_celltypes=[],
)

scv.pp.filter_and_normalize(
    raw_adata,
    min_shared_counts=20,
    n_top_genes=2000,
)

assert all(prior_adata.obs_names == raw_adata.obs_names), "Observation names are not aligned!"
assert all(prior_adata.var_names == raw_adata.var_names), "Variable names are not aligned!"

raw_adata.layers["pred_cell_type"] = prior_adata.layers["pred_cell_type"]
raw_adata.obsm["X_pre_embed"] = prior_adata.obsm["X_pre_embed"]
raw_adata.obs["pred_cluster"] = prior_adata.obs["pred_cluster"].astype(int)

velo_adata = raw_adata[:, prior_adata.var["is_velocity_gene"]].copy()

df_fine, adjacency_matrix_fine, velo_adata = steer.preprocess_anndata_spatial(
    velo_adata,
    npc=30,
    NUM_AD_NEIGH=30,
    spatial_neighbors=cfg.spatial_neighbors,
    SMOOTH_NEIGH=cfg.smooth_neigh,
    moments_adj=True,
    neighbor_metric=cfg.neighbor_metric,
    combine_mode=cfg.graph,
    spatial_key="X_spatial",
    use_us=cfg.use_us,
)

dataset_fine = steer.preload_datasets_all_genes_anndata(
    df=df_fine,
    MODEL_MODE="whole",
    adata=velo_adata,
)
pyg_data_fine = steer.create_pyg_data(dataset_fine, adjacency_matrix_fine, normalize=True)
pyg_data_fine.fine_clus_vec = torch.tensor(fine_clus_vec_np, dtype=torch.long, device=device)

print(velo_adata)
print("Velocity-gene subset shape:", velo_adata.shape)

prior_adata = result_adata.copy() raw_adata = sc.read_h5ad(INPUT_FILE) raw_adata = steer.prepare_spatial_adata( raw_adata, obs_copy_map={"celltype": "celltype1"}, na_clear_keys=["subregion"], celltype_key="celltype", remove_celltypes=[], ) scv.pp.filter_and_normalize( raw_adata, min_shared_counts=20, n_top_genes=2000, ) assert all(prior_adata.obs_names == raw_adata.obs_names), "Observation names are not aligned!" assert all(prior_adata.var_names == raw_adata.var_names), "Variable names are not aligned!" raw_adata.layers["pred_cell_type"] = prior_adata.layers["pred_cell_type"] raw_adata.obsm["X_pre_embed"] = prior_adata.obsm["X_pre_embed"] raw_adata.obs["pred_cluster"] = prior_adata.obs["pred_cluster"].astype(int) velo_adata = raw_adata[:, prior_adata.var["is_velocity_gene"]].copy() df_fine, adjacency_matrix_fine, velo_adata = steer.preprocess_anndata_spatial( velo_adata, npc=30, NUM_AD_NEIGH=30, spatial_neighbors=cfg.spatial_neighbors, SMOOTH_NEIGH=cfg.smooth_neigh, moments_adj=True, neighbor_metric=cfg.neighbor_metric, combine_mode=cfg.graph, spatial_key="X_spatial", use_us=cfg.use_us, ) dataset_fine = steer.preload_datasets_all_genes_anndata( df=df_fine, MODEL_MODE="whole", adata=velo_adata, ) pyg_data_fine = steer.create_pyg_data(dataset_fine, adjacency_matrix_fine, normalize=True) pyg_data_fine.fine_clus_vec = torch.tensor(fine_clus_vec_np, dtype=torch.long, device=device) print(velo_adata) print("Velocity-gene subset shape:", velo_adata.shape)

Filtered out 12279 genes that are detected 20 counts (shared).
Normalized count data: X, spliced, unspliced.
Extracted 2000 highly variable genes.
Logarithmized X.

/nvme/users/liuzhy/miniconda3/envs/steer_dev/lib/python3.10/site-packages/scvelo/preprocessing/utils.py:705: DeprecationWarning: `log1p` is deprecated since scVelo v0.3.0 and will be removed in a future version. Please use `log1p` from `scanpy.pp` instead.
  log1p(adata)

computing moments based on connectivities
    finished (0:00:01) --> added 
    'Ms' and 'Mu', moments of un/spliced abundances (adata.layers)
AnnData object with n_obs × n_vars = 13258 × 1000
    obs: 'x', 'y', 'sample', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt', 'subregion', 'interface', 'celltype1', 'subclusters_5_last', 'subregion5', 'celltype', 'initial_size_unspliced', 'initial_size_spliced', 'initial_size', 'n_counts', 'pred_cluster'
    var: 'n_cells_by_counts', 'gene_count_corr', 'means', 'dispersions', 'dispersions_norm', 'highly_variable'
    uns: 'log1p', 'neighbors'
    obsm: 'X_pca', 'X_umap', 'spatial', 'X_spatial', 'X_pre_embed', 'X_pca_combined', 'X_pca_moments'
    layers: 'ambiguous', 'counts', 'matrix', 'spliced', 'unspliced', 'pred_cell_type', 'Ms', 'Mu', 'scale_Mu', 'scale_Ms'
    obsp: 'distances', 'connectivities'
Velocity-gene subset shape: (13258, 1000)

Kinetic Learning¶

Run the main kinetic-learning stage using the prior inferred above.

• full mode: preserve the behavior of your original run by not passing extra training-control arguments and therefore using STEER defaults
• fast mode: use a lightweight quickstart-style configuration, including expert_mode="slim", shorter pretraining/clustering, velo_batch_size=512, and more aggressive early stopping

In [13]:

stage2_kwargs = dict(
    device=device,
    device2=device,
    pyg_data=pyg_data_fine,
    MODEL_MODE="whole",
    adata=velo_adata,
    NUM_LOSS_NEIGH=30,
    max_n_cluster=cfg.expert,
    corr_mode=cfg.corr_mode,
    path=RESULT_PATH,
)
stage2_kwargs.update(
    get_stage2_training_overrides(
        cfg.training_mode,
        getattr(cfg, "finetune_epochs", None),
    )
)

velo_adata = steer.model_training_share_neighbor_adata(**stage2_kwargs)

print("Kinetic-learning stage completed.")

stage2_kwargs = dict( device=device, device2=device, pyg_data=pyg_data_fine, MODEL_MODE="whole", adata=velo_adata, NUM_LOSS_NEIGH=30, max_n_cluster=cfg.expert, corr_mode=cfg.corr_mode, path=RESULT_PATH, ) stage2_kwargs.update( get_stage2_training_overrides( cfg.training_mode, getattr(cfg, "finetune_epochs", None), ) ) velo_adata = steer.model_training_share_neighbor_adata(**stage2_kwargs) print("Kinetic-learning stage completed.")

Using Fine Cluster Vector for Correlation Loss.
Epoch 0, Loss 0.2589128911495209
Epoch 50, Loss 0.16269493103027344
Epoch 100, Loss 0.14793340861797333
Epoch 150, Loss 0.1351746916770935
Epoch 200, Loss 0.12515060603618622
Epoch 250, Loss 0.11659395694732666
Epoch 300, Loss 0.10895542055368423
Epoch 350, Loss 0.10202070325613022
Epoch 400, Loss 0.09567110985517502
Epoch 450, Loss 0.09001993387937546
Epoch 500, Loss 0.0846339762210846
Epoch 550, Loss 0.07984798401594162
Epoch 600, Loss 0.07547919452190399
Epoch 650, Loss 0.07120946794748306
Epoch 700, Loss 0.06731133162975311
Epoch 750, Loss 0.06387172639369965
Epoch 800, Loss 0.06048090383410454
Epoch 850, Loss 0.05737226456403732
Epoch 900, Loss 0.055343613028526306
Epoch 950, Loss 0.05166403204202652
Epoch 1000, Loss 0.04917184263467789
Epoch 1050, Loss 0.04679033160209656
Epoch 1100, Loss 0.044529758393764496
Epoch 1150, Loss 0.04279634356498718
Epoch 1200, Loss 0.04079746827483177
Epoch 1250, Loss 0.039058245718479156
Epoch 1300, Loss 0.03734171763062477
Epoch 1350, Loss 0.040087949484586716
Epoch 1400, Loss 0.035372477024793625
Epoch 1450, Loss 0.03351101279258728
Epoch 1500, Loss 1.196820855140686
Epoch 1550, Loss 0.43179768323898315
Epoch 1600, Loss 0.13766944408416748
Epoch 1650, Loss 0.12595361471176147
Epoch 1700, Loss 0.1224173828959465
Epoch 1750, Loss 0.12003050744533539
Epoch 1800, Loss 0.11823020130395889
Epoch 1850, Loss 0.11621135473251343
Epoch 1900, Loss 0.11506755650043488
Epoch 1950, Loss 0.11379842460155487
Epoch 2000, Loss 0.1127137765288353
Epoch 2050, Loss 0.1121092140674591
Epoch 2100, Loss 0.11080144345760345
Epoch 2150, Loss 0.10985640436410904
Epoch 2200, Loss 0.10872479528188705
Epoch 2250, Loss 0.10754203051328659
Epoch 2300, Loss 0.10670628398656845
Epoch 2350, Loss 0.10669361799955368
Epoch 2400, Loss 0.10572731494903564
Epoch 2450, Loss 0.10473927110433578
Epoch 2500, Loss 0.104276642203331
Epoch 2550, Loss 0.1037045568227768
Epoch 2600, Loss 0.10339640080928802
Epoch 2650, Loss 0.10257221758365631
Epoch 2700, Loss 0.10216522961854935
Epoch 2750, Loss 0.10235632210969925
Epoch 2800, Loss 0.10156235098838806
Epoch 2850, Loss 0.10093402117490768
Epoch 2900, Loss 0.1017594113945961
Epoch 2950, Loss 0.10103590041399002
Epoch 3000, Loss 0.10016175359487534
Epoch 3050, Loss 1.082615613937378
GATE: 0.012044417671859264, Cluster: 0.08827710151672363, Time_cor: 0.970267117023468, Time_smooth: 0.012027065269649029
Epoch 3100, Loss 1.0812548398971558
GATE: 0.012347091920673847, Cluster: 0.08857107162475586, Time_cor: 0.9688102602958679, Time_smooth: 0.011526440270245075
Epoch 3150, Loss 1.0804443359375
GATE: 0.012219026684761047, Cluster: 0.08889162540435791, Time_cor: 0.9681321382522583, Time_smooth: 0.011201534420251846
Epoch 3200, Loss 1.0788326263427734
GATE: 0.011120345443487167, Cluster: 0.08933907747268677, Time_cor: 0.9675549864768982, Time_smooth: 0.010818195529282093
Epoch 3250, Loss 1.3803349733352661
GATE: 0.011094989255070686, Cluster: 0.08863979578018188, Velo: 0.30653777718544006, Time_cor: 0.9634653329849243, Time_smooth: 0.010597116313874722
Epoch 3300, Loss 1.3738303184509277
GATE: 0.01123539824038744, Cluster: 0.08874636888504028, Velo: 0.3042329251766205, Time_cor: 0.9593216776847839, Time_smooth: 0.01029396615922451
Epoch 3350, Loss 1.367795467376709
GATE: 0.010930167511105537, Cluster: 0.08890658617019653, Velo: 0.3031860888004303, Time_cor: 0.9547840356826782, Time_smooth: 0.009988497011363506
Epoch 3400, Loss 1.361010193824768
GATE: 0.010407967492938042, Cluster: 0.08880615234375, Velo: 0.3018128275871277, Time_cor: 0.9502310752868652, Time_smooth: 0.009752143174409866
Epoch 3450, Loss 1.2211984395980835
GATE: 0.010212086141109467, Cluster: 0.08806401491165161, Velo: 0.16594989597797394, Time_cor: 0.9473838806152344, Time_smooth: 0.009588559158146381
Epoch 3500, Loss 1.2118154764175415
GATE: 0.010086567141115665, Cluster: 0.08804678916931152, Velo: 0.1607392579317093, Time_cor: 0.9435639381408691, Time_smooth: 0.009378894232213497
Epoch 3550, Loss 1.209521770477295
GATE: 0.009999600239098072, Cluster: 0.08803069591522217, Velo: 0.1597042828798294, Time_cor: 0.9424842000007629, Time_smooth: 0.009302948601543903
Epoch 3600, Loss 1.1965620517730713
GATE: 0.009900867007672787, Cluster: 0.08825445175170898, Velo: 0.15065324306488037, Time_cor: 0.9386317133903503, Time_smooth: 0.009121719747781754
Epoch 3650, Loss 1.190578579902649
GATE: 0.009708236902952194, Cluster: 0.08801466226577759, Velo: 0.14929789304733276, Time_cor: 0.9347052574157715, Time_smooth: 0.008852426894009113
Epoch 3700, Loss 1.18563711643219
GATE: 0.009592004120349884, Cluster: 0.08800023794174194, Velo: 0.1463208794593811, Time_cor: 0.9329900741577148, Time_smooth: 0.008733856491744518
Epoch 3750, Loss 1.1844679117202759
GATE: 0.009557134471833706, Cluster: 0.08799785375595093, Velo: 0.1455390453338623, Time_cor: 0.9326629042625427, Time_smooth: 0.008711026981472969
Epoch 3800, Loss 1.1757668256759644
GATE: 0.01015174575150013, Cluster: 0.08862680196762085, Velo: 0.13902631402015686, Time_cor: 0.9288412928581238, Time_smooth: 0.009120664559304714
Epoch 3850, Loss 1.1668519973754883
GATE: 0.009345742873847485, Cluster: 0.08832329511642456, Velo: 0.13680201768875122, Time_cor: 0.9237710237503052, Time_smooth: 0.008609897457063198
Epoch 3900, Loss 1.1572538614273071
GATE: 0.009130877442657948, Cluster: 0.08809685707092285, Velo: 0.13205373287200928, Time_cor: 0.9196163415908813, Time_smooth: 0.008356085978448391
Epoch 3950, Loss 1.1492500305175781
GATE: 0.008929635398089886, Cluster: 0.08796972036361694, Velo: 0.12770427763462067, Time_cor: 0.9164651036262512, Time_smooth: 0.008181351236999035
Epoch 4000, Loss 1.1438512802124023
GATE: 0.008734568953514099, Cluster: 0.08795660734176636, Velo: 0.12473256886005402, Time_cor: 0.9143545031547546, Time_smooth: 0.008073066361248493
Epoch 4050, Loss 1.140921711921692
GATE: 0.008637391962110996, Cluster: 0.08794885873794556, Velo: 0.12316607683897018, Time_cor: 0.9131571054458618, Time_smooth: 0.008012274280190468
Epoch 4100, Loss 1.1397737264633179
GATE: 0.008604566566646099, Cluster: 0.08794623613357544, Velo: 0.122579425573349, Time_cor: 0.9126566052436829, Time_smooth: 0.007986941374838352
Epoch 4150, Loss 1.139508843421936
GATE: 0.008597088046371937, Cluster: 0.08794564008712769, Velo: 0.1224483773112297, Time_cor: 0.9125370979309082, Time_smooth: 0.007980701513588428
Epoch 4200, Loss 1.157185673713684
GATE: 0.012622547335922718, Cluster: 0.08936619758605957, Velo: 0.13409006595611572, Time_cor: 0.9105046391487122, Time_smooth: 0.010602226480841637
Epoch 4250, Loss 1.142844796180725
GATE: 0.010127781890332699, Cluster: 0.0894087553024292, Velo: 0.12892132997512817, Time_cor: 0.9048349261283875, Time_smooth: 0.009552056901156902
Epoch 4300, Loss 1.1327553987503052
GATE: 0.009350032545626163, Cluster: 0.0891040563583374, Velo: 0.1255977302789688, Time_cor: 0.8998115658760071, Time_smooth: 0.008892087265849113
Epoch 4350, Loss 1.1224281787872314
GATE: 0.008731761015951633, Cluster: 0.08822035789489746, Velo: 0.12165818363428116, Time_cor: 0.8952330946922302, Time_smooth: 0.008584747090935707
Epoch 4400, Loss 1.1156086921691895
GATE: 0.00841191504150629, Cluster: 0.08808088302612305, Velo: 0.11958674341440201, Time_cor: 0.8911604881286621, Time_smooth: 0.008368715643882751
Epoch 4450, Loss 1.1080775260925293
GATE: 0.008132564835250378, Cluster: 0.08805596828460693, Velo: 0.11624801903963089, Time_cor: 0.8874685168266296, Time_smooth: 0.008172564208507538
Epoch 4500, Loss 1.1032919883728027
GATE: 0.007996812462806702, Cluster: 0.0881463885307312, Velo: 0.1147657185792923, Time_cor: 0.884342610836029, Time_smooth: 0.008040444925427437
Epoch 4550, Loss 1.0978318452835083
GATE: 0.007839338853955269, Cluster: 0.08794772624969482, Velo: 0.11246862262487411, Time_cor: 0.8816572427749634, Time_smooth: 0.007918905466794968
Epoch 4600, Loss 1.0944464206695557
GATE: 0.007739794906228781, Cluster: 0.08793038129806519, Velo: 0.11145562678575516, Time_cor: 0.8794903755187988, Time_smooth: 0.007830262184143066
Epoch 4650, Loss 1.09187912940979
GATE: 0.007657997775822878, Cluster: 0.0879209041595459, Velo: 0.11075300723314285, Time_cor: 0.8777870535850525, Time_smooth: 0.007760196458548307
Epoch 4700, Loss 1.0900304317474365
GATE: 0.007593474350869656, Cluster: 0.08791428804397583, Velo: 0.1102980375289917, Time_cor: 0.876514732837677, Time_smooth: 0.007709831465035677
Epoch 4750, Loss 1.0886799097061157
GATE: 0.007543893530964851, Cluster: 0.08790987730026245, Velo: 0.10992904752492905, Time_cor: 0.8756230473518372, Time_smooth: 0.007674045395106077
Epoch 4800, Loss 1.087944746017456
GATE: 0.007511216215789318, Cluster: 0.08790719509124756, Velo: 0.10982445627450943, Time_cor: 0.8750511407852173, Time_smooth: 0.007650775834918022
Epoch 4850, Loss 1.0874361991882324
GATE: 0.007492555771023035, Cluster: 0.08790552616119385, Velo: 0.10967205464839935, Time_cor: 0.874728798866272, Time_smooth: 0.007637291215360165
Epoch 4900, Loss 1.087217092514038
GATE: 0.007483609952032566, Cluster: 0.08790498971939087, Velo: 0.1096191257238388, Time_cor: 0.8745784759521484, Time_smooth: 0.007630867417901754
Epoch 4950, Loss 1.0870885848999023
GATE: 0.0074798958376049995, Cluster: 0.08790445327758789, Velo: 0.10955894738435745, Time_cor: 0.8745170831680298, Time_smooth: 0.007628169842064381
Epoch 5000, Loss 1.1013144254684448
GATE: 0.009586738422513008, Cluster: 0.08883917331695557, Velo: 0.12169327586889267, Time_cor: 0.8721827864646912, Time_smooth: 0.00901248212903738
Epoch 5050, Loss 1.0866996049880981
GATE: 0.007731691002845764, Cluster: 0.08912616968154907, Velo: 0.11460129916667938, Time_cor: 0.8671166300773621, Time_smooth: 0.008123761974275112
Epoch 5100, Loss 1.078181505203247
GATE: 0.007448095828294754, Cluster: 0.08855944871902466, Velo: 0.11184629052877426, Time_cor: 0.8625697493553162, Time_smooth: 0.007757938001304865
Epoch 5150, Loss 1.0715179443359375
GATE: 0.007221666164696217, Cluster: 0.0883033275604248, Velo: 0.11007972806692123, Time_cor: 0.8583517074584961, Time_smooth: 0.007561408914625645
Epoch 5200, Loss 1.0643658638000488
GATE: 0.006921063642948866, Cluster: 0.08808070421218872, Velo: 0.10766797512769699, Time_cor: 0.8542972803115845, Time_smooth: 0.007398800924420357
Epoch 5250, Loss 1.060773253440857
GATE: 0.007050701417028904, Cluster: 0.08801567554473877, Velo: 0.1077314019203186, Time_cor: 0.8506755828857422, Time_smooth: 0.007299864199012518
Epoch 5300, Loss 1.0558470487594604
GATE: 0.00666264770552516, Cluster: 0.08852708339691162, Velo: 0.10647322237491608, Time_cor: 0.8470093607902527, Time_smooth: 0.007174781523644924
Epoch 5350, Loss 1.0525844097137451
GATE: 0.006647567730396986, Cluster: 0.08822524547576904, Velo: 0.10695832222700119, Time_cor: 0.843681275844574, Time_smooth: 0.007071946747601032
Epoch 5400, Loss 1.047547459602356
GATE: 0.0065012602135539055, Cluster: 0.08797264099121094, Velo: 0.10562252253293991, Time_cor: 0.8404737710952759, Time_smooth: 0.006977146025747061
Epoch 5450, Loss 1.0437800884246826
GATE: 0.00633258419111371, Cluster: 0.08788490295410156, Velo: 0.10518580675125122, Time_cor: 0.837491512298584, Time_smooth: 0.0068852873519063
Epoch 5500, Loss 1.04043710231781
GATE: 0.006242380477488041, Cluster: 0.08790385723114014, Velo: 0.1047382578253746, Time_cor: 0.8347528576850891, Time_smooth: 0.0067998687736690044
Epoch 5550, Loss 1.0378684997558594
GATE: 0.006227951962500811, Cluster: 0.0880555510520935, Velo: 0.10457754135131836, Time_cor: 0.8322846293449402, Time_smooth: 0.006722853519022465
Epoch 5600, Loss 1.0358641147613525
GATE: 0.006239944137632847, Cluster: 0.08801436424255371, Velo: 0.10489521920681, Time_cor: 0.8300444483757019, Time_smooth: 0.006670068949460983
Epoch 5650, Loss 1.0320398807525635
GATE: 0.0060389540158212185, Cluster: 0.08782577514648438, Velo: 0.10379302501678467, Time_cor: 0.8277932405471802, Time_smooth: 0.006588905118405819
Epoch 5700, Loss 1.029962182044983
GATE: 0.005978072062134743, Cluster: 0.08782970905303955, Velo: 0.10375072807073593, Time_cor: 0.8258833289146423, Time_smooth: 0.006520414724946022
Epoch 5750, Loss 1.0278149843215942
GATE: 0.00593440281227231, Cluster: 0.0878250002861023, Velo: 0.10344870388507843, Time_cor: 0.8241484761238098, Time_smooth: 0.006458403542637825
Epoch 5800, Loss 1.0261303186416626
GATE: 0.005905304569751024, Cluster: 0.08780622482299805, Velo: 0.10340522229671478, Time_cor: 0.8226078152656555, Time_smooth: 0.006405672989785671
Epoch 5850, Loss 1.0247693061828613
GATE: 0.005853667389601469, Cluster: 0.08779603242874146, Velo: 0.10352982580661774, Time_cor: 0.8212324380874634, Time_smooth: 0.006357325706630945
Epoch 5900, Loss 1.0233055353164673
GATE: 0.005850446410477161, Cluster: 0.08780771493911743, Velo: 0.10330219566822052, Time_cor: 0.8200313448905945, Time_smooth: 0.006313822697848082
Epoch 5950, Loss 1.0219359397888184
GATE: 0.005800541490316391, Cluster: 0.08778965473175049, Velo: 0.10309317708015442, Time_cor: 0.8189746141433716, Time_smooth: 0.006277864798903465
Epoch 6000, Loss 1.020864725112915
GATE: 0.005770120769739151, Cluster: 0.08777827024459839, Velo: 0.10300976037979126, Time_cor: 0.8180639147758484, Time_smooth: 0.0062426612712442875
Epoch 6050, Loss 1.0199862718582153
GATE: 0.00575246149674058, Cluster: 0.08777529001235962, Velo: 0.1029452383518219, Time_cor: 0.8172997832298279, Time_smooth: 0.006213495507836342
Epoch 6100, Loss 1.0192850828170776
GATE: 0.005739361979067326, Cluster: 0.08777165412902832, Velo: 0.10292348265647888, Time_cor: 0.8166617751121521, Time_smooth: 0.006188767962157726
Epoch 6150, Loss 1.0187125205993652
GATE: 0.005729443393647671, Cluster: 0.087771475315094, Velo: 0.10289962589740753, Time_cor: 0.8161433339118958, Time_smooth: 0.006168588064610958
Epoch 6200, Loss 1.0181814432144165
GATE: 0.005718035623431206, Cluster: 0.08776623010635376, Velo: 0.1028205007314682, Time_cor: 0.8157249689102173, Time_smooth: 0.00615169620141387
Epoch 6250, Loss 1.0177758932113647
GATE: 0.005705184303224087, Cluster: 0.08776432275772095, Velo: 0.1027686595916748, Time_cor: 0.8153986930847168, Time_smooth: 0.006139045115560293
Epoch 6300, Loss 1.0175398588180542
GATE: 0.005695028230547905, Cluster: 0.08776283264160156, Velo: 0.10279858112335205, Time_cor: 0.8151549696922302, Time_smooth: 0.00612839637324214
Epoch 6350, Loss 1.0173245668411255
GATE: 0.005689541343599558, Cluster: 0.08776170015335083, Velo: 0.10277228057384491, Time_cor: 0.8149803280830383, Time_smooth: 0.006120631471276283
Epoch 6400, Loss 1.0172476768493652
GATE: 0.0056862072087824345, Cluster: 0.0877610445022583, Velo: 0.1028229296207428, Time_cor: 0.8148621916770935, Time_smooth: 0.006115361116826534
Epoch 6450, Loss 1.0171583890914917
GATE: 0.00568402511999011, Cluster: 0.08776050806045532, Velo: 0.10281575471162796, Time_cor: 0.8147862553596497, Time_smooth: 0.0061119189485907555
Epoch 6500, Loss 1.0171239376068115
GATE: 0.005682635121047497, Cluster: 0.08776021003723145, Velo: 0.10283230245113373, Time_cor: 0.8147390484809875, Time_smooth: 0.006109725683927536
Epoch 6550, Loss 1.0170707702636719
GATE: 0.005681649781763554, Cluster: 0.0877600908279419, Velo: 0.10281441360712051, Time_cor: 0.8147064447402954, Time_smooth: 0.0061081754975020885
Epoch 6600, Loss 1.0457849502563477
GATE: 0.008375262841582298, Cluster: 0.0890914797782898, Velo: 0.12211146205663681, Time_cor: 0.8155385851860046, Time_smooth: 0.010668143630027771
Epoch 6650, Loss 1.0268605947494507
GATE: 0.006450243294239044, Cluster: 0.08852696418762207, Velo: 0.11061352491378784, Time_cor: 0.811881959438324, Time_smooth: 0.00938780140131712
Epoch 6700, Loss 1.0161387920379639
GATE: 0.005956344306468964, Cluster: 0.08829057216644287, Velo: 0.10436268150806427, Time_cor: 0.8088057637214661, Time_smooth: 0.008723445236682892
Epoch 6750, Loss 1.0108702182769775
GATE: 0.0056641120463609695, Cluster: 0.08864718675613403, Velo: 0.10222592204809189, Time_cor: 0.805975615978241, Time_smooth: 0.008357441984117031
Epoch 6800, Loss 1.0074751377105713
GATE: 0.005453695543110371, Cluster: 0.08873772621154785, Velo: 0.10181617736816406, Time_cor: 0.8033656477928162, Time_smooth: 0.008101857267320156
Epoch 6850, Loss 1.0029737949371338
GATE: 0.005300185643136501, Cluster: 0.08806926012039185, Velo: 0.10081362724304199, Time_cor: 0.8009004592895508, Time_smooth: 0.007890259847044945
Epoch 6900, Loss 0.9993692636489868
GATE: 0.005176184698939323, Cluster: 0.08794909715652466, Velo: 0.10000029951334, Time_cor: 0.7985448837280273, Time_smooth: 0.0076987603679299355
Epoch 6950, Loss 0.9979360103607178
GATE: 0.005089192185550928, Cluster: 0.0880441665649414, Velo: 0.10085313022136688, Time_cor: 0.7963911890983582, Time_smooth: 0.007558345329016447
Epoch 7000, Loss 0.9947862029075623
GATE: 0.005004344042390585, Cluster: 0.08805012702941895, Velo: 0.10000236332416534, Time_cor: 0.7943517565727234, Time_smooth: 0.007377649657428265
Epoch 7050, Loss 0.9926965832710266
GATE: 0.0049346862360835075, Cluster: 0.08800286054611206, Velo: 0.1000981330871582, Time_cor: 0.7924322485923767, Time_smooth: 0.007228698581457138
Epoch 7100, Loss 0.9903712868690491
GATE: 0.004870917182415724, Cluster: 0.08790290355682373, Velo: 0.09987255185842514, Time_cor: 0.79063880443573, Time_smooth: 0.0070860725827515125
Epoch 7150, Loss 0.9895790815353394
GATE: 0.004824119620025158, Cluster: 0.08845144510269165, Velo: 0.10034463554620743, Time_cor: 0.7890077829360962, Time_smooth: 0.0069510736502707005
Epoch 7200, Loss 0.9862426519393921
GATE: 0.00477056298404932, Cluster: 0.08788585662841797, Velo: 0.09923910349607468, Time_cor: 0.7875185608863831, Time_smooth: 0.006828551180660725
Epoch 7250, Loss 0.9851499199867249
GATE: 0.004726836457848549, Cluster: 0.08793431520462036, Velo: 0.09980065375566483, Time_cor: 0.7859762907028198, Time_smooth: 0.006711858324706554
Epoch 7300, Loss 0.9837131500244141
GATE: 0.0046791499480605125, Cluster: 0.08794409036636353, Velo: 0.09971475601196289, Time_cor: 0.7847806811332703, Time_smooth: 0.006594495847821236
Epoch 7350, Loss 0.9817723035812378
GATE: 0.004638052079826593, Cluster: 0.08784276247024536, Velo: 0.09931929409503937, Time_cor: 0.7834752798080444, Time_smooth: 0.006496930029243231
Epoch 7400, Loss 0.9811117053031921
GATE: 0.004614874720573425, Cluster: 0.08793169260025024, Velo: 0.09987549483776093, Time_cor: 0.78228360414505, Time_smooth: 0.006405985448509455
Epoch 7450, Loss 0.9792163372039795
GATE: 0.004567503929138184, Cluster: 0.0877918004989624, Velo: 0.09940885752439499, Time_cor: 0.7811370491981506, Time_smooth: 0.006311125122010708
Epoch 7500, Loss 0.9781901836395264
GATE: 0.004540708381682634, Cluster: 0.08786767721176147, Velo: 0.09936755895614624, Time_cor: 0.7801856398582458, Time_smooth: 0.006228545214980841
Epoch 7550, Loss 0.9767673015594482
GATE: 0.004507418721914291, Cluster: 0.0877676010131836, Velo: 0.0991223156452179, Time_cor: 0.7792171239852905, Time_smooth: 0.006152895279228687
Epoch 7600, Loss 0.9758310317993164
GATE: 0.004476058296859264, Cluster: 0.08774620294570923, Velo: 0.09914430230855942, Time_cor: 0.7783869504928589, Time_smooth: 0.006077541038393974
Epoch 7650, Loss 0.9751962423324585
GATE: 0.0044507356360554695, Cluster: 0.08780640363693237, Velo: 0.09935236722230911, Time_cor: 0.7775810956954956, Time_smooth: 0.006005644798278809
Epoch 7700, Loss 0.9749343991279602
GATE: 0.004427350126206875, Cluster: 0.08793371915817261, Velo: 0.09977144747972488, Time_cor: 0.7768612504005432, Time_smooth: 0.005940631497651339
Epoch 7750, Loss 0.9736799597740173
GATE: 0.004397893324494362, Cluster: 0.08774352073669434, Velo: 0.09949829429388046, Time_cor: 0.7761685848236084, Time_smooth: 0.005871680565178394
Epoch 7800, Loss 0.9726776480674744
GATE: 0.004374772775918245, Cluster: 0.08783948421478271, Velo: 0.09915327280759811, Time_cor: 0.7755023837089539, Time_smooth: 0.005807682406157255
Epoch 7850, Loss 0.9723786115646362
GATE: 0.004352892283350229, Cluster: 0.08769649267196655, Velo: 0.09967965632677078, Time_cor: 0.7749003767967224, Time_smooth: 0.005749234929680824
Epoch 7900, Loss 0.9717591404914856
GATE: 0.004330942407250404, Cluster: 0.08773022890090942, Velo: 0.09960772097110748, Time_cor: 0.7743942141532898, Time_smooth: 0.005696005187928677
Epoch 7950, Loss 0.9711437225341797
GATE: 0.004310221876949072, Cluster: 0.08767622709274292, Velo: 0.09955906867980957, Time_cor: 0.7739584445953369, Time_smooth: 0.005639771930873394
Epoch 8000, Loss 0.9705800414085388
GATE: 0.0042869397439062595, Cluster: 0.08765441179275513, Velo: 0.0996483787894249, Time_cor: 0.7734007239341736, Time_smooth: 0.005589613690972328
Epoch 8050, Loss 0.969946563243866
GATE: 0.004266001284122467, Cluster: 0.0876724123954773, Velo: 0.09952178597450256, Time_cor: 0.772949755191803, Time_smooth: 0.005536619573831558
Epoch 8100, Loss 0.9703080058097839
GATE: 0.004253047052770853, Cluster: 0.087718665599823, Velo: 0.1002044603228569, Time_cor: 0.7726386785507202, Time_smooth: 0.005493152420967817
Epoch 8150, Loss 0.9692478775978088
GATE: 0.004232471343129873, Cluster: 0.087654709815979, Velo: 0.09965667128562927, Time_cor: 0.7722550630569458, Time_smooth: 0.005448993295431137
Epoch 8200, Loss 0.9688740968704224
GATE: 0.004214796237647533, Cluster: 0.08764523267745972, Velo: 0.09969495236873627, Time_cor: 0.7719132900238037, Time_smooth: 0.0054058488458395
Epoch 8250, Loss 0.9683599472045898
GATE: 0.004199066665023565, Cluster: 0.08764547109603882, Velo: 0.0995921790599823, Time_cor: 0.771557629108429, Time_smooth: 0.005365635268390179
Epoch 8300, Loss 0.9678453803062439
GATE: 0.004181047901511192, Cluster: 0.08763772249221802, Velo: 0.09943448752164841, Time_cor: 0.7712644934654236, Time_smooth: 0.005327656865119934
Epoch 8350, Loss 0.9678047895431519
GATE: 0.004168678540736437, Cluster: 0.0876307487487793, Velo: 0.09965796023607254, Time_cor: 0.7710549235343933, Time_smooth: 0.005292469635605812
Epoch 8400, Loss 0.9674119353294373
GATE: 0.004154940135776997, Cluster: 0.0876263976097107, Velo: 0.09961070865392685, Time_cor: 0.7707613110542297, Time_smooth: 0.0052585722878575325
Epoch 8450, Loss 0.9672373533248901
GATE: 0.004142634104937315, Cluster: 0.08761757612228394, Velo: 0.09970902651548386, Time_cor: 0.7705414891242981, Time_smooth: 0.005226606503129005
Epoch 8500, Loss 0.966927707195282
GATE: 0.004132611211389303, Cluster: 0.08760976791381836, Velo: 0.09964275360107422, Time_cor: 0.7703459858894348, Time_smooth: 0.0051965778693556786
Epoch 8550, Loss 0.9667996168136597
GATE: 0.004120145924389362, Cluster: 0.0876050591468811, Velo: 0.09976320713758469, Time_cor: 0.7701445817947388, Time_smooth: 0.0051666260696947575
Epoch 8600, Loss 0.9663020968437195
GATE: 0.0041094934567809105, Cluster: 0.08760350942611694, Velo: 0.09947992861270905, Time_cor: 0.7699694037437439, Time_smooth: 0.005139756947755814
Epoch 8650, Loss 0.9662418961524963
GATE: 0.004100230056792498, Cluster: 0.08759897947311401, Velo: 0.09961529076099396, Time_cor: 0.7698128819465637, Time_smooth: 0.0051145232282578945
Epoch 8700, Loss 0.9659525156021118
GATE: 0.004092263989150524, Cluster: 0.08759880065917969, Velo: 0.09949983656406403, Time_cor: 0.7696707844734192, Time_smooth: 0.005090820603072643
Epoch 8750, Loss 0.9659520983695984
GATE: 0.004084437154233456, Cluster: 0.08759260177612305, Velo: 0.09966809302568436, Time_cor: 0.7695383429527283, Time_smooth: 0.005068597849458456
Epoch 8800, Loss 0.9657763242721558
GATE: 0.004077694844454527, Cluster: 0.08759093284606934, Velo: 0.09963715076446533, Time_cor: 0.7694226503372192, Time_smooth: 0.005047900136560202
Epoch 8850, Loss 0.9657258987426758
GATE: 0.004071352072060108, Cluster: 0.08758807182312012, Velo: 0.09972009807825089, Time_cor: 0.7693174481391907, Time_smooth: 0.005028945859521627
Epoch 8900, Loss 0.9655774235725403
GATE: 0.00406515970826149, Cluster: 0.08758604526519775, Velo: 0.09969264268875122, Time_cor: 0.7692224383354187, Time_smooth: 0.005011124070733786
Epoch 8950, Loss 0.9655265808105469
GATE: 0.004059385973960161, Cluster: 0.08758413791656494, Velo: 0.0997503250837326, Time_cor: 0.7691377997398376, Time_smooth: 0.004994953982532024
Epoch 9000, Loss 0.9653282165527344
GATE: 0.004053900949656963, Cluster: 0.08758246898651123, Velo: 0.09964900463819504, Time_cor: 0.7690629959106445, Time_smooth: 0.004979853983968496
Epoch 9050, Loss 0.9653526544570923
GATE: 0.0040480345487594604, Cluster: 0.08758103847503662, Velo: 0.0997605249285698, Time_cor: 0.76899653673172, Time_smooth: 0.004966516979038715
Epoch 9100, Loss 0.9651688933372498
GATE: 0.004043589346110821, Cluster: 0.08757972717285156, Velo: 0.09965334087610245, Time_cor: 0.7689381241798401, Time_smooth: 0.0049541182816028595
Epoch 9150, Loss 0.9650940299034119
GATE: 0.004039971623569727, Cluster: 0.0875784158706665, Velo: 0.099645234644413, Time_cor: 0.7688872814178467, Time_smooth: 0.004943148232996464
Epoch 9200, Loss 0.9650557637214661
GATE: 0.004036327823996544, Cluster: 0.08757740259170532, Velo: 0.0996655598282814, Time_cor: 0.7688432335853577, Time_smooth: 0.004933238495141268
Epoch 9250, Loss 0.9649985432624817
GATE: 0.0040329936891794205, Cluster: 0.08757656812667847, Velo: 0.09965869039297104, Time_cor: 0.7688056826591492, Time_smooth: 0.004924597218632698
Epoch 9300, Loss 0.9649551510810852
GATE: 0.004030311014503241, Cluster: 0.08757579326629639, Velo: 0.09965816885232925, Time_cor: 0.7687737941741943, Time_smooth: 0.004917107988148928
Epoch 9350, Loss 0.9649384617805481
GATE: 0.00402803486213088, Cluster: 0.08757507801055908, Velo: 0.09967734664678574, Time_cor: 0.7687472701072693, Time_smooth: 0.004910745192319155
Epoch 9400, Loss 0.964871346950531
GATE: 0.004026159644126892, Cluster: 0.0875745415687561, Velo: 0.09963998198509216, Time_cor: 0.7687253355979919, Time_smooth: 0.004905354231595993
Epoch 9450, Loss 0.9648260474205017
GATE: 0.004024612717330456, Cluster: 0.0875740647315979, Velo: 0.09961894899606705, Time_cor: 0.7687074542045593, Time_smooth: 0.0049009788781404495
Epoch 9500, Loss 0.964786946773529
GATE: 0.00402335450053215, Cluster: 0.0875735878944397, Velo: 0.09959957003593445, Time_cor: 0.7686931490898132, Time_smooth: 0.004897290840744972
Epoch 9550, Loss 0.9647733569145203
GATE: 0.004022343084216118, Cluster: 0.08757340908050537, Velo: 0.09960167109966278, Time_cor: 0.7686816453933716, Time_smooth: 0.0048943180590868
Epoch 9600, Loss 0.9647328853607178
GATE: 0.004021534230560064, Cluster: 0.08757305145263672, Velo: 0.09957371652126312, Time_cor: 0.7686726450920105, Time_smooth: 0.004891927819699049
Epoch 9650, Loss 0.9647408723831177
GATE: 0.00402087764814496, Cluster: 0.08757281303405762, Velo: 0.09959176182746887, Time_cor: 0.7686654925346375, Time_smooth: 0.004889989271759987
Epoch 9700, Loss 0.9647354483604431
GATE: 0.00402032770216465, Cluster: 0.08757257461547852, Velo: 0.09959463030099869, Time_cor: 0.7686595320701599, Time_smooth: 0.004888363182544708
Epoch 9750, Loss 0.9647352695465088
GATE: 0.004019826650619507, Cluster: 0.08757239580154419, Velo: 0.09960190951824188, Time_cor: 0.7686542868614197, Time_smooth: 0.004886863753199577
Epoch 9800, Loss 1.0113807916641235
GATE: 0.005942315794527531, Cluster: 0.08935075998306274, Velo: 0.13152064383029938, Time_cor: 0.7731704115867615, Time_smooth: 0.01139668095856905
Epoch 9850, Loss 0.9987475275993347
GATE: 0.005331888794898987, Cluster: 0.08846592903137207, Velo: 0.12266112864017487, Time_cor: 0.7717241048812866, Time_smooth: 0.01056448183953762
Epoch 9900, Loss 0.9907968044281006
GATE: 0.004886934068053961, Cluster: 0.08894503116607666, Velo: 0.11621556431055069, Time_cor: 0.77089524269104, Time_smooth: 0.009854043833911419
Epoch 9950, Loss 0.9823846817016602
GATE: 0.004556606523692608, Cluster: 0.08899080753326416, Velo: 0.10943824797868729, Time_cor: 0.7701858878135681, Time_smooth: 0.009213155135512352
Loading best model state from memory...
Kinetic-learning stage completed.

Spatial Velocity Plot¶

Build the velocity graph and inspect the learned flow field in spatial coordinates.

In [14]:

result_adata = velo_adata.copy()

sc.pp.neighbors(result_adata, use_rep="X_refine_embed_t", n_neighbors=30)
steer.velocity_graph(result_adata, vkey="pred_vs_norm", xkey="model_Ms")

print("obs columns:")
print(sorted(result_adata.obs.columns.tolist()))

print("\nobsm keys:")
print(sorted(result_adata.obsm.keys()))

print("\nlayers:")
print(sorted(result_adata.layers.keys()))

result_adata = velo_adata.copy() sc.pp.neighbors(result_adata, use_rep="X_refine_embed_t", n_neighbors=30) steer.velocity_graph(result_adata, vkey="pred_vs_norm", xkey="model_Ms") print("obs columns:") print(sorted(result_adata.obs.columns.tolist())) print("\nobsm keys:") print(sorted(result_adata.obsm.keys())) print("\nlayers:") print(sorted(result_adata.layers.keys()))

computing velocity graph (using 1/128 cores)

  0%|          | 0/13258 [00:00<?, ?cells/s]

    finished (0:00:15) --> added 
    'pred_vs_norm_graph', sparse matrix with cosine correlations (adata.uns)
obs columns:
['Expert', 'Expert Weight', 'Pred Time', 'celltype', 'celltype1', 'initial_size', 'initial_size_spliced', 'initial_size_unspliced', 'interface', 'n_counts', 'n_genes_by_counts', 'pct_counts_mt', 'pred_vs_norm_self_transition', 'pretrain_cluster', 'sample', 'subclusters_5_last', 'subregion', 'subregion5', 'total_counts', 'total_counts_mt', 'x', 'y']

obsm keys:
['X_alpha', 'X_beta', 'X_gamma', 'X_para', 'X_para_t', 'X_pca', 'X_pca_combined', 'X_pca_moments', 'X_pre_embed', 'X_refine_embed', 'X_refine_embed_t', 'X_spatial', 'X_umap', 'cluster_matrix', 'spatial']

layers:
['counts', 'final_recon_s', 'final_recon_u', 'init_regulate_state', 'model_Ms', 'model_Mu', 'orig_s', 'orig_u', 'pred_time_layer', 'pred_vs', 'pred_vs_norm', 'pred_vu', 'pred_vu_norm', 'recon_alpha', 'recon_alpha_norm', 'recon_beta', 'recon_gamma', 'recon_gamma_norm', 'regulate_state', 'scale_Ms', 'scale_Mu']

In [15]:

result_adata

result_adata

Out[15]:

AnnData object with n_obs × n_vars = 13258 × 1000
    obs: 'x', 'y', 'sample', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt', 'subregion', 'interface', 'celltype1', 'subclusters_5_last', 'subregion5', 'celltype', 'initial_size_unspliced', 'initial_size_spliced', 'initial_size', 'n_counts', 'pretrain_cluster', 'Expert', 'Expert Weight', 'Pred Time', 'pred_vs_norm_self_transition'
    var: 'n_cells_by_counts', 'gene_count_corr', 'means', 'dispersions', 'dispersions_norm', 'highly_variable'
    uns: 'log1p', 'neighbors', 'pred_vs_norm_graph', 'pred_vs_norm_graph_neg', 'pred_vs_norm_params'
    obsm: 'X_pca', 'X_umap', 'spatial', 'X_spatial', 'X_pre_embed', 'X_pca_combined', 'X_pca_moments', 'X_refine_embed', 'cluster_matrix', 'X_alpha', 'X_beta', 'X_gamma', 'X_para', 'X_para_t', 'X_refine_embed_t'
    layers: 'counts', 'scale_Mu', 'scale_Ms', 'recon_alpha', 'recon_beta', 'recon_gamma', 'pred_vu', 'pred_vs', 'pred_vu_norm', 'pred_vs_norm', 'init_regulate_state', 'regulate_state', 'final_recon_s', 'final_recon_u', 'model_Ms', 'model_Mu', 'orig_s', 'orig_u', 'pred_time_layer', 'recon_alpha_norm', 'recon_gamma_norm'
    obsp: 'distances', 'connectivities'

In [16]:

scv.settings.figdir = RESULT_PATH
scv.set_figure_params(style="scvelo", dpi=150, figsize=(5, 4), transparent=True)
if "X_spatial" not in result_adata.obsm and "spatial" in result_adata.obsm:
    result_adata.obsm["X_spatial"] = result_adata.obsm["spatial"].astype(float)
elif "X_spatial" in result_adata.obsm:
    result_adata.obsm["X_spatial"] = result_adata.obsm["X_spatial"].astype(float)
# scv.pl.velocity_embedding_stream(
#     result_adata,
#     basis="spatial",
#     vkey="pred_vs_norm",
#     color=["Expert", "Pred Time","subregion","celltype"],
#     title=["Expert Assignment", "Latent Time","subregion","celltype"],
#     show=True,
# )
scv.pl.velocity_embedding_stream(result_adata,color='subregion',size=66,smooth=0.8,vkey="pred_vs_norm",
                                 alpha=1,groups=['Inner EPC','Outer EPC','Allantois','Labyrinth (LaTP)/Chorion','JZ (P-TGC)',"Reichert's membrane/CP",'Yolk sac'],
                                 legend_loc='right',basis='spatial',add_margin=0.1,xlabel= 'Spatial X',ylabel='Spatial Y',figsize=(4,3), frameon=True, dpi=300,
                                 title='STEER results',save='velo_joint_spatial_part.png')

scv.settings.figdir = RESULT_PATH scv.set_figure_params(style="scvelo", dpi=150, figsize=(5, 4), transparent=True) if "X_spatial" not in result_adata.obsm and "spatial" in result_adata.obsm: result_adata.obsm["X_spatial"] = result_adata.obsm["spatial"].astype(float) elif "X_spatial" in result_adata.obsm: result_adata.obsm["X_spatial"] = result_adata.obsm["X_spatial"].astype(float) # scv.pl.velocity_embedding_stream( # result_adata, # basis="spatial", # vkey="pred_vs_norm", # color=["Expert", "Pred Time","subregion","celltype"], # title=["Expert Assignment", "Latent Time","subregion","celltype"], # show=True, # ) scv.pl.velocity_embedding_stream(result_adata,color='subregion',size=66,smooth=0.8,vkey="pred_vs_norm", alpha=1,groups=['Inner EPC','Outer EPC','Allantois','Labyrinth (LaTP)/Chorion','JZ (P-TGC)',"Reichert's membrane/CP",'Yolk sac'], legend_loc='right',basis='spatial',add_margin=0.1,xlabel= 'Spatial X',ylabel='Spatial Y',figsize=(4,3), frameon=True, dpi=300, title='STEER results',save='velo_joint_spatial_part.png')

computing velocity embedding
    finished (0:00:01) --> added
    'pred_vs_norm_spatial', embedded velocity vectors (adata.obsm)
saving figure to file ./results/E8.5_S1.MPSTA_quickstart/scvelo_velo_joint_spatial_part.png