### Install PHATE and scprep

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Install the PHATE and scprep libraries using pip. A kernel restart may be required after installation.

```python
!pip install --user --upgrade --quiet phate scprep
```

--------------------------------

### Install TrajectoryNet

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/README.rst

Install the TrajectoryNet library using pip. This code was tested with Python 3.7 and 3.8.

```bash
pip install TrajectoryNet
```

--------------------------------

### Load Example Anndata Dataset

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Loads the 'gastrulation_erythroid' dataset from ScVelo. This dataset is used as an example for processing.

```python
adata = scv.datasets.gastrulation_erythroid()
```

--------------------------------

### Load and Configure ScVelo

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Loads the ScVelo library and sets up configuration parameters for verbosity, presenter view, and figure aesthetics. Ensure ScVelo is installed.

```python
%load_ext autoreload
%autoreload 2
%load_ext lab_black
import scvelo as scv
import scanpy as sc

scv.settings.verbosity = 3  # show errors(0), warnings(1), info(2), hints(3)
scv.settings.presenter_view = True  # set max width size for presenter view
scv.settings.set_figure_params("scvelo")  # for beautified visualization
```

--------------------------------

### Trajectory Density Visualization

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Saves trajectory visualizations as image sequences and videos. Requires torch and numpy. Ensure ffmpeg is installed for video conversion.

```python
import torch
import numpy as np
from TrajectoryNet.lib.viz_scrna import (
    save_trajectory,
    save_trajectory_density,
    save_vectors,
    trajectory_to_video
)

# Assume model is trained and loaded
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()

# Define base distribution functions
def base_density(z):
    """Standard normal log probability"""
    import math
    logZ = -0.5 * math.log(2 * math.pi)
    return torch.sum(logZ - z.pow(2) / 2, 1, keepdim=True)

def base_sampler(n, dim):
    """Sample from standard normal"""
    return torch.randn(n, dim)

# Sample some reference data points
data_samples = np.random.randn(2000, 2)

# Save trajectory visualization
save_trajectory(
    prior_logdensity=base_density,
    prior_sampler=base_sampler,
    model=model,
    data_samples=data_samples,
    savedir='./output/trajectory',
    ntimes=25,                    # Number of time steps
    end_times=[0.5, 1.0, 1.5],   # Integration endpoints
    memory=0.1,                   # Memory fraction for batching
    device=device
)

# Save density evolution
save_trajectory_density(
    prior_logdensity=base_density,
    model=model,
    data_samples=data_samples,
    savedir='./output/density',
    ntimes=100,
    end_times=[0.5, 1.0, 1.5],
    memory=0.1,
    device=device
)

# Convert image sequence to video (requires ffmpeg)
trajectory_to_video('./output/trajectory')
trajectory_to_video('./output/density')

# Save vector field visualization
start_points = torch.randn(1000, 2)
save_vectors(
    prior_logdensity=base_density,
    model=model,
    data_samples=start_points,
    full_data=data_samples,
    labels=np.zeros(len(data_samples)),
    savedir='./output/vectors',
    ntimes=100,
    end_times=[0.5, 1.0],
    device=device
)
```

--------------------------------

### Print scVelo Version

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Prints the installed version of scVelo and the Python version. Useful for debugging and ensuring compatibility.

```python
import scvelo as scv
scv.logging.print_version()
```

--------------------------------

### Run TrajectoryNet on SCURVE dataset

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/README.rst

Execute TrajectoryNet on the S Curve example dataset. For custom datasets, provide the path to the .npz file and the embedding name.

```bash
python -m TrajectoryNet.main --dataset SCURVE
```

--------------------------------

### Add sample_labels key to adata.obs

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Assigns integer labels to cell types in adata.obs['sample_labels']. These labels should be increasing integers starting from zero with no gaps for optimal TrajectoryNet performance.

```python
adata.obs["sample_labels"] = adata.obs["celltype"].replace(
    {
        "Blood progenitors 1": 0,
        "Blood progenitors 2": 1,
        "Erythroid1": 2,
        "Erythroid2": 3,
        "Erythroid3": 4,
    }
)
```

--------------------------------

### Backward Pass: Integrate from Observed Data to Base

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Performs a backward integration from observed data back to the base distribution. Use `reverse=False` to start from data and integrate backward.

```python
observed_data = torch.randn(500, input_dim).to(device)
logp_obs = torch.zeros(500, 1).to(device)

with torch.no_grad():
    # reverse=False: start from data, integrate backward to base
    z_base, logp_base = model(observed_data, logp_obs,
                               integration_times=integration_times,
                               reverse=False)
```

--------------------------------

### Display TrajectoryNet Help

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Use this command to list all available parameters and configuration options.

```bash
python -m TrajectoryNet.main --help
```

--------------------------------

### Multi-step Integration Through Multiple Timepoints

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Demonstrates multi-step integration by iterating through a list of target timepoints. Each step integrates from the previous timepoint to the current one.

```python
time_scale = 0.5
int_tps = [0.5, 1.0, 1.5]  # Multiple target times

z_current = z
logp_current = logpz
trajectories = [z_current.cpu().numpy()]

with torch.no_grad():
    for i, tp in enumerate(int_tps):
        prev_tp = 0.0 if i == 0 else int_tps[i-1]
        times = torch.tensor([prev_tp, tp]).to(device)
        z_current, logp_current = model(z_current, logp_current,
                                         integration_times=times,
                                         reverse=True)
        trajectories.append(z_current.cpu().numpy())

trajectories = np.stack(trajectories)  # (n_times+1, n_samples, input_dim)
```

--------------------------------

### Prepare and Save Data

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Package the extracted data into a dictionary and save as a compressed NumPy archive.

```python
to_save = {
    'umap' : adata.obsm['X_umap'],
    'delta_umap' : adata.obsm['velocity_umap'],
    'sample_labels' : inverse,
}
```

```python
np.savez('../data/dentategyrus.npz', **to_save)
```

--------------------------------

### Set Download Path

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Sets the download path for datasets. Ensure sufficient disk space for the download.

```python
import os
import scprep
download_path = os.path.expanduser("~")
print(download_path)
```

--------------------------------

### Get Unique Valid Days

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Extracts unique day values from the days_df DataFrame, excluding any NaN values. This provides a list of distinct time points in the dataset.

```python
unique_days = days_df['day'].unique()
unique_days = unique_days[np.isnan(unique_days) == False]
```

--------------------------------

### Train TrajectoryNet Model (CLI)

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Run TrajectoryNet training from the command line. Supports various datasets and configurations, including embeddings, regularization, and growth models.

```bash
# Train on the S-Curve synthetic dataset
python -m TrajectoryNet.main --dataset SCURVE
```

```bash
# Train on the Embryoid Body dataset with PHATE embedding
python -m TrajectoryNet.main --dataset EB
```

```bash
# Train on a custom NPZ dataset with PCA embedding
python -m TrajectoryNet.main \
    --dataset /path/to/your/data.npz \
    --embedding_name pca \
    --max_dim 10 \
    --niters 10000 \
    --batch_size 1000 \
    --lr 1e-3 \
    --save ./results/my_experiment
```

```bash
# Train with velocity regularization (using RNA velocity data)
python -m TrajectoryNet.main \
    --dataset /path/to/data.npz \
    --embedding_name pca \
    --vecint 1.0 \
    --use_magnitude
```

```bash
# Train with growth model for cell division/death
python -m TrajectoryNet.main \
    --dataset EB-PCA \
    --use_growth \
    --max_dim 5
```

```bash
# Train with spectral normalization and batch normalization
python -m TrajectoryNet.main \
    --dataset SCURVE \
    --spectral_norm \
    --batch_norm \
    --bn_lag 0.1
```

--------------------------------

### Configure scVelo Settings

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Sets scVelo's verbosity level and enables presenter view for wider output. Use verbosity levels 0-3 to control the amount of information displayed.

```python
scv.settings.verbosity = 3  # show errors(0), warnings(1), info(2), hints(3)
scv.settings.presenter_view = True  # set max width size for presenter view
scv.set_figure_params('scvelo')  # for beautified visualization
```

--------------------------------

### Load 10X Genomics Data (Sparse Matrix)

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Imports scRNA-seq data from 10X Genomics using a sparse matrix format for memory efficiency. Set `sparse=False` for dense matrix loading. `gene_labels='both'` displays gene symbols and IDs.

```python
sparse=True
T1 = scprep.io.load_10X(os.path.join(download_path, "scRNAseq", "T0_1A"), sparse=sparse, gene_labels='both')
T2 = scprep.io.load_10X(os.path.join(download_path, "scRNAseq", "T2_3B"), sparse=sparse, gene_labels='both')
T3 = scprep.io.load_10X(os.path.join(download_path, "scRNAseq", "T4_5C"), sparse=sparse, gene_labels='both')
T4 = scprep.io.load_10X(os.path.join(download_path, "scRNAseq", "T6_7D"), sparse=sparse, gene_labels='both')
T5 = scprep.io.load_10X(os.path.join(download_path, "scRNAseq", "T8_9E"), sparse=sparse, gene_labels='both')
T1.head()
```

--------------------------------

### Load Data using Dataset Factory (Python)

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Use the `SCData.factory()` method to load various dataset types programmatically. It supports built-in synthetic and biological datasets, as well as custom NPZ and AnnData files.

```python
import numpy as np
from TrajectoryNet import dataset
from argparse import Namespace

# Create args namespace with required parameters
args = Namespace(
    embedding_name='pca',
    max_dim=10,
    whiten=True
)

# Load built-in synthetic datasets
scurve_data = dataset.SCData.factory("SCURVE", args)
moons_data = dataset.SCData.factory("MOONS", args)
circles_data = dataset.SCData.factory("CIRCLES", args)
blobs_data = dataset.SCData.factory("BLOBS", args)

# Load built-in biological datasets
circle3_data = dataset.SCData.factory("CIRCLE3", args)
tree_data = dataset.SCData.factory("TREE", args)
cycle_data = dataset.SCData.factory("CYCLE", args)

# Load from custom NPZ file
# NPZ must contain: embedding array + 'sample_labels' array
custom_data = dataset.SCData.factory("/path/to/custom_data.npz", args)

# Load from AnnData (.h5ad) file
anndata_data = dataset.SCData.factory("/path/to/adata.h5ad", args)

# Access data properties
print(f"Data shape: {custom_data.get_shape()}")
print(f"Number of cells: {custom_data.get_ncells()}")
print(f"Unique timepoints: {custom_data.get_unique_times()}")
print(f"Has velocity: {custom_data.has_velocity()}")

# Get data arrays
embedding = custom_data.get_data()  # (n_cells, n_dims)
timepoints = custom_data.get_times()  # (n_cells,)
if custom_data.has_velocity():
    velocity = custom_data.get_velocity()  # (n_cells, n_dims)
```

--------------------------------

### Import NumPy

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Initial import required for array operations.

```python
import numpy as np
```

--------------------------------

### View MAGIC Execution Logs

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Displays the standard output logs generated during the MAGIC algorithm execution, including timing for KNN search, affinity calculation, and imputation.

```text
    Calculated KNN search in 0.45 seconds.
    Calculating affinities...
    Calculated affinities in 0.41 seconds.
  Calculated graph and diffusion operator in 0.87 seconds.
  Calculating imputation...
  Calculated imputation in 0.02 seconds.
Calculated MAGIC in 0.89 seconds.
Calculating MAGIC...
  Running MAGIC on 3332 cells and 68 genes.
  Calculating graph and diffusion operator...
    Calculating KNN search...
```

```text
    Calculated KNN search in 0.39 seconds.
    Calculating affinities...
    Calculated affinities in 0.40 seconds.
  Calculated graph and diffusion operator in 0.80 seconds.
  Calculating imputation...
  Calculated imputation in 0.02 seconds.
Calculated MAGIC in 0.83 seconds.
Calculating MAGIC...
  Running MAGIC on 3332 cells and 68 genes.
  Calculating graph and diffusion operator...
    Calculating KNN search...
```

```text
    Calculated KNN search in 0.40 seconds.
    Calculating affinities...
    Calculated affinities in 0.41 seconds.
  Calculated graph and diffusion operator in 0.82 seconds.
  Calculating imputation...
  Calculated imputation in 0.02 seconds.
Calculated MAGIC in 0.84 seconds.
```

--------------------------------

### Saving and Loading Model Checkpoints

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Provides functions to save and load trained model checkpoints using PyTorch. Includes saving periodic and best checkpoints, and loading for evaluation.

```python
import torch
import os
from TrajectoryNet.train_misc import (
    build_model_tabular,
    create_regularization_fns,
    set_cnf_options,
    add_spectral_norm
)
from TrajectoryNet.lib import utils

# Save checkpoint during training
def save_checkpoint(model, growth_model, save_dir, epoch):
    utils.makedirs(save_dir)

    checkpoint = {
        'state_dict': model.state_dict(),
        'epoch': epoch,
    }

    if growth_model is not None:
        checkpoint['growth_state_dict'] = growth_model.state_dict()

    # Save periodic checkpoint
    torch.save(checkpoint, os.path.join(save_dir, f'checkpt_{epoch}.pth'))

    # Save best checkpoint (overwrite)
    torch.save(checkpoint, os.path.join(save_dir, 'checkpt.pth'))

# Load checkpoint for evaluation
def load_checkpoint(args, input_dim, checkpoint_path, device):
    # Recreate model architecture
    regularization_fns, _ = create_regularization_fns(args)
    model = build_model_tabular(args, input_dim, regularization_fns)

    if args.spectral_norm:
        add_spectral_norm(model)

    set_cnf_options(args, model)

    # Load weights
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint['state_dict'])
    model = model.to(device)
    model.eval()

    return model

# Example usage
# model = load_checkpoint(args, input_dim=10,
#                         checkpoint_path='./results/checkpt.pth',
#                         device=device)
```

--------------------------------

### Instantiate PHATE Estimator and Transform Data

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Instantiate a PHATE estimator with specified parameters and transform the input data to generate a PHATE embedding. Use `n_jobs` for parallel processing and `random_state` for reproducibility.

```python
phate_operator = phate.PHATE(n_jobs=-2, random_state=42)
Y_phate = phate_operator.fit_transform(EBT_counts)
```

--------------------------------

### Define save dictionary

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Initializes a dictionary for saving data.

```python
to_save = {
```

--------------------------------

### Define end genes and colors

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Initializes lists for end genes and cell populations, and maps them to colors for visualization.

```python
end_genes = ['PDGFRA ', 'HAND1', 'SOX17', 'ONECUT2', ]
end_points = ['Muscle', 'Cardiac', 'Endothelial', 'Neuronal',]

colors = dict(zip(*[end_genes, [plt.get_cmap('tab10')(i+1) for i in range(len(end_genes))]]))
```

--------------------------------

### Run PHATE dimensionality reduction

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Initializes and fits a PHATE operator to the data.

```python
phate_op = phate.PHATE(n_jobs=-2, k=25, random_state=42)
data_phate = phate_op.fit_transform(df)
```

--------------------------------

### Visualize 2D scatter plots

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Various configurations for plotting 2D scatter plots using scprep.

```python
scprep.plot.scatter2d(data_phate[mask],
                      ticks=False,
                      c=days_df['day'])
```

```python
scprep.plot.scatter2d(data_phate, c=gene_set_df['MEF.identity'])
```

```python
scprep.plot.scatter2d(phate_op.graph.data_nu[mask], c=days_df['day'])
```

```python
scprep.plot.scatter2d(phate_op.graph.data_nu[mask], figsize=(16,9), c=days_df['day'])
```

```python
scprep.plot.scatter2d(coord_df[mask], 
                      #figsize=(16,9), 
                      ticks=False,
                      c=days_df['day'])
```

--------------------------------

### Build Continuous Normalizing Flow (CNF) Model

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Configures and builds a Continuous Normalizing Flow model for trajectory inference. Allows customization of network architecture, time integration, ODE solver settings, and various regularization options. Requires PyTorch and argparse.Namespace.

```python
import torch
from argparse import Namespace
from TrajectoryNet.train_misc import (
    build_model_tabular,
    create_regularization_fns,
    set_cnf_options,
    add_spectral_norm,
    count_parameters
)

# Configure model architecture
args = Namespace(
    # Network architecture
    dims='64-64-64',           # Hidden layer dimensions
    layer_type='concatsquash', # Options: ignore, concat, concat_v2, squash, concatsquash, blend
    nonlinearity='tanh',       # Options: tanh, relu, softplus, elu, swish
    num_blocks=1,              # Number of stacked CNF blocks

    # Time integration
    time_scale=0.5,            # Integration time scale
    train_T=True,              # Learn integration time

    # ODE solver settings
    solver='dopri5',           # Options: dopri5, bdf, rk4, midpoint, adams
    atol=1e-5,                 # Absolute tolerance
    rtol=1e-5,                 # Relative tolerance
    step_size=None,            # Optional fixed step size

    # Test settings (can differ from training)
    test_solver=None,
    test_atol=None,
    test_rtol=None,

    # Regularization
    residual=False,
    rademacher=False,
    batch_norm=False,
    bn_lag=0.0,

    # Regularization coefficients (None = disabled)
    l1int=None,                # L1 norm of velocity
    l2int=None,                # L2 norm of velocity
    sl2int=None,               # Squared L2 norm
    JFrobint=0.01,             # Jacobian Frobenius norm
    JdiagFrobint=None,         # Diagonal Jacobian Frobenius
    JoffdiagFrobint=None,      # Off-diagonal Jacobian Frobenius
)

# Create regularization functions
regularization_fns, regularization_coeffs = create_regularization_fns(args)

# Build the model (input_dim = embedding dimension)
input_dim = 10
model = build_model_tabular(args, input_dim, regularization_fns)

# Optionally add spectral normalization
# add_spectral_norm(model)

# Configure solver options
set_cnf_options(args, model)

print(f"Model parameters: {count_parameters(model)}")
print(model)

# Move to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
```

--------------------------------

### Define Base Path for Data Files

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Sets the base directory path where the dataset files are located. Ensure this path is correct for your system.

```python
PATH = '/data/lab/DataSets/Schiebinger_2019_OT/'
```

--------------------------------

### Forward Pass: Integrate from Base to Target Time

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Performs a forward integration from a base distribution to a target time. Use `reverse=True` to sample from the base and integrate forward.

```python
with torch.no_grad():
    # reverse=True: sample from base, integrate forward in time
    z_t, logp_t = model(z, logpz, integration_times=integration_times, reverse=True)

print(f"Input shape: {z.shape}")
print(f"Output shape: {z_t.shape}")
print(f"Log probability shape: {logp_t.shape}")
```

--------------------------------

### Import Libraries for Data Loading and Analysis

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Imports necessary Python libraries for data manipulation, plotting, and analysis, including scprep, phate, matplotlib, pandas, numpy, and scanpy.

```python
import scprep
import phate
import matplotlib.pyplot as plt
%matplotlib inline
import pandas as pd
import numpy as np
import scanpy
import os
```

--------------------------------

### Evaluate TrajectoryNet Model (CLI)

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Evaluate trained TrajectoryNet models and compute metrics from the command line. Supports validation with left-out timepoints.

```bash
# Evaluate a trained model
python -m TrajectoryNet.eval \
    --dataset /path/to/data.npz \
    --embedding_name pca \
    --save ./results/my_experiment
```

```bash
# Evaluate with a left-out timepoint for validation
python -m TrajectoryNet.eval \
    --dataset EB-PCA \
    --leaveout_timepoint 2 \
    --save ./results/eb_experiment
```

--------------------------------

### Visualize Velocity Graph

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Generates a visualization of the velocity graph, showing cell-to-cell transition probabilities based on velocity vectors.

```python
scv.pl.velocity_graph(adata)
```

--------------------------------

### Normalize library size

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Corrects for library size differences by dividing by library size and rescaling by the median.

```python
EBT_counts = scprep.normalize.library_size_normalize(EBT_counts)
```

--------------------------------

### Complete Training Loop for TrajectoryNet

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

This script outlines a comprehensive training process for TrajectoryNet, including data loading, model configuration, optimization, and loss calculation with regularization. It requires PyTorch and specific TrajectoryNet modules.

```python
import torch
import torch.optim as optim
import numpy as np
from argparse import Namespace
from TrajectoryNet import dataset
from TrajectoryNet.train_misc import (
    build_model_tabular,
    create_regularization_fns,
    set_cnf_options,
    count_parameters,
    get_regularization
)
from TrajectoryNet.lib import utils

# Configuration
args = Namespace(
    # Data
    dataset='SCURVE',
    embedding_name='pca',
    max_dim=2,
    whiten=False,

    # Model
    dims='64-64-64',
    layer_type='concatsquash',
    nonlinearity='tanh',
    num_blocks=1,
    time_scale=0.5,
    train_T=True,

    # Solver
    solver='dopri5',
    atol=1e-5,
    rtol=1e-5,
    step_size=None,
    test_solver=None,
    test_atol=None,
    test_rtol=None,

    # Training
    niters=1000,
    batch_size=500,
    lr=1e-3,
    weight_decay=1e-5,
    training_noise=0.1,

    # Regularization
    residual=False,
    rademacher=False,
    spectral_norm=False,
    batch_norm=False,
    bn_lag=0.0,
    l1int=None,
    l2int=None,
    sl2int=None,
    JFrobint=0.01,
    JdiagFrobint=None,
    JoffdiagFrobint=None,

    # Output
    save='./results/example',
    val_freq=100,
)

# Setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = dataset.SCData.factory(args.dataset, args)
args.timepoints = data.get_unique_times()
args.int_tps = (np.arange(max(args.timepoints) + 1) + 1.0) * args.time_scale

# Build model
regularization_fns, regularization_coeffs = create_regularization_fns(args)
model = build_model_tabular(args, data.get_shape()[0], regularization_fns).to(device)
set_cnf_options(args, model)

print(f"Parameters: {count_parameters(model)}")

# Optimizer
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# Training loop
for itr in range(1, args.niters + 1):
    model.train()
    optimizer.zero_grad()

    # Compute loss (simplified version)
    loss = torch.tensor(0.0, device=device)
    z = None

    for i, (itp, tp) in enumerate(zip(args.int_tps[::-1], args.timepoints[::-1])):
        integration_times = torch.tensor([itp - args.time_scale, itp]).float().to(device)

        idx = data.sample_index(args.batch_size, tp)
        x = torch.from_numpy(data.get_data()[idx]).float().to(device)
        x += torch.randn_like(x) * args.training_noise

        if z is not None:
            x = torch.cat((z, x))

        zero = torch.zeros(x.shape[0], 1).to(device)
        z, delta_logp = model(x, zero, integration_times=integration_times)

    logpz = data.base_density()(z)
    loss = -torch.mean(logpz)

    # Add regularization
    if len(regularization_coeffs) > 0:
        reg_states = get_regularization(model, regularization_coeffs)
        if reg_states:
            for reg_state, coeff in zip(reg_states, regularization_coeffs):
                if coeff != 0:
                    loss += reg_state * coeff

    loss.backward()
    optimizer.step()

    if itr % 100 == 0:
        print(f"Iter {itr}: Loss = {loss.item():.6f}")

# Save final model
utils.makedirs(args.save)
torch.save({'state_dict': model.state_dict()}, f'{args.save}/checkpt.pth')
print(f"Model saved to {args.save}/checkpt.pth")
```

--------------------------------

### Process Sample Labels

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Convert observation labels into numerical arrays for model input.

```python
np.array(adata.obs['age(days)'])
```

```python
unique, inverse = np.unique(np.array(adata.obs['age(days)']), return_inverse=True)
unique, inverse
```

--------------------------------

### Run MAGIC imputation

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Initializes the MAGIC operator using the existing PHATE graph and transforms the counts.

```python
m_op = magic.MAGIC()
m_op.graph = phate_operator.graph
EBT_magic = m_op.transform(EBT_counts, genes=genes_of_interest_full)
```

--------------------------------

### CNF Forward Pass for Sample Transformation

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Performs a forward pass using a trained CNF model to transform samples between time points. Requires a configured model and PyTorch tensors for input data and log probabilities. Samples are transformed from a base distribution (e.g., Gaussian) to a target distribution.

```python
import torch
import numpy as np
from TrajectoryNet.lib.layers import CNF

# Assume model is built and loaded
# model = build_model_tabular(args, input_dim, regularization_fns)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()

# Sample data at a specific timepoint
n_samples = 1000
input_dim = 10
z = torch.randn(n_samples, input_dim).to(device)

# Initialize log probability (zeros for uniform weighting)
logpz = torch.zeros(n_samples, 1).to(device)

# Define integration times (from t=0 to t=0.5)
integration_times = torch.tensor([0.0, 0.5]).to(device)
```

--------------------------------

### Train TrajectoryNet Model

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Executes the training process for the TrajectoryNet model with specified density regularization.

```bash
python main.py --save [SAVE_DIR] --dataset EB-PCA --top_k_reg 0.1 --training_noise 0.0 --max_dim 5
```

--------------------------------

### Inspect unique sample labels

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Retrieves the unique values from the sample labels array to verify dataset composition.

```python
np.unique(to_save['sample_labels'])
```

--------------------------------

### Plot Library Size Distribution

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Visualizes the distribution of library sizes for a given dataset. Use this to understand cell library size variations before filtering.

```python
scprep.plot.plot_library_size(T1, percentile=20)
```

--------------------------------

### Compute Training Loss with Backward Integration

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Calculates the negative log-likelihood loss by integrating backward from the last time step. Includes options for training noise regularization.

```python
import torch
import numpy as np
from TrajectoryNet import dataset
from argparse import Namespace

def compute_loss(device, model, data, args, batch_size=1000):
    """
    Compute loss by integrating backwards from the last time step.
    """
    timepoints = data.get_unique_times()
    int_tps = (np.arange(max(timepoints) + 1) + 1.0) * args.time_scale

    deltas = []
    zs = []
    z = None

    # Backward pass through timepoints
    for i, (itp, tp) in enumerate(zip(int_tps[::-1], timepoints[::-1])):
        integration_times = torch.tensor([itp - args.time_scale, itp])
        integration_times = integration_times.float().to(device)

        # Sample batch from current timepoint
        idx = data.sample_index(batch_size, tp)
        x = data.get_data()[idx]

        # Add training noise for regularization
        if args.training_noise > 0:
            x += np.random.randn(*x.shape) * args.training_noise

        x = torch.from_numpy(x).float().to(device)

        # Concatenate with previous transformed samples
        if i > 0:
            x = torch.cat((z, x))
            zs.append(z)

        zero = torch.zeros(x.shape[0], 1).to(device)

        # Transform to previous timepoint
        z, delta_logp = model(x, zero, integration_times=integration_times)
        deltas.append(delta_logp)

    # Compute base distribution log probability
    logpz = data.base_density()(z)

    # Accumulate losses across timepoints
    losses = []
    logps = [logpz]

    for i, delta_logp in enumerate(deltas[::-1]):
        logpx = logps[-1] - delta_logp
        logps.append(logpx[:-batch_size])
        losses.append(-torch.mean(logpx[-batch_size:]))

    total_loss = torch.mean(torch.stack(losses))
    return total_loss

# Example usage
args = Namespace(
    time_scale=0.5,
    training_noise=0.1,
    batch_size=1000
)

# loss = compute_loss(device, model, data, args)
# loss.backward()
```

--------------------------------

### Apply square root transformation

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Applies a square root transformation to the data.

```python
df = scprep.transform.sqrt(df)
```

--------------------------------

### Load Custom NPZ Data (Python)

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Load custom trajectory data from NumPy NPZ files using the `CustomData` class. The NPZ file must contain an embedding array and a 'sample_labels' array.

```python
import numpy as np
from TrajectoryNet.dataset import CustomData
from argparse import Namespace

# Prepare data: cells x dimensions embedding with time labels
n_cells = 5000
n_dims = 10
n_timepoints = 4

# Create synthetic data with 4 timepoints
embedding = np.random.randn(n_cells, n_dims)
sample_labels = np.repeat(np.arange(n_timepoints), n_cells // n_timepoints)

# Optional: add velocity information
delta_embedding = np.random.randn(n_cells, n_dims) * 0.1
```

--------------------------------

### Import Libraries for Data Analysis

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Imports necessary libraries for single-cell data analysis, including pandas for DataFrames, numpy for numerical operations, and scprep for data loading.

```python
import pandas as pd
import numpy as np
import phate
import scprep
import magic
import matplotlib.pyplot as plt
import sklearn.preprocessing

# matplotlib settings for Jupyter notebooks only
%matplotlib inline
```

--------------------------------

### Scatter plot with scprep

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Generates a scatter plot using scprep's plotting utilities.

```python
scprep.plot.scatter(coord_df['x'], coord_df['y'],c=gene_set_df['MEF.identity'],)
```

--------------------------------

### Plot library size

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Visualizes the library size distribution of the provided DataFrame.

```python
scprep.plot.plot_library_size(df)
```

--------------------------------

### Generate Backward Trajectories

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Creates the backward_trajectories.npy file from a saved model checkpoint.

```bash
python eval.py --save [SAVE_DIR] --dataset EB-PCA --top_k_reg 0.1 --training_noise 0.0 --max_dim 5
```

--------------------------------

### Load Trajectory Data from AnnData

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Loads trajectory data directly from an AnnData object. Requires 'X_pca' in obsm, 'sample_labels' in obs, and optionally velocity data. Configuration is done via a Namespace object. This method is suitable for single-cell analysis workflows.

```python
import scanpy as sc
import numpy as np
from TrajectoryNet.dataset import CustomAnnData
from argparse import Namespace

# Load or create AnnData object
adata = sc.read_h5ad('your_data.h5ad')

# Ensure required fields exist
# 1. Embedding in obsm (e.g., 'X_pca', 'X_umap', 'X_phate')
if 'X_pca' not in adata.obsm:
    sc.pp.pca(adata, n_comps=50)

# 2. Time labels in obs as 'sample_labels'
adata.obs['sample_labels'] = adata.obs['timepoint'].astype(int).values

# 3. Optional: velocity in obsm (e.g., 'velocity_pca')
# This can come from scvelo or similar RNA velocity tools

# Load into TrajectoryNet
args = Namespace(
    embedding_name='pca',  # Will look for 'X_pca' in obsm
    max_dim=20,
    whiten=True
)

data = CustomAnnData(adata, args)

# Use the data
print(f"Loaded {data.get_ncells()} cells")
print(f"Timepoints: {data.get_unique_times()}")
print(f"Embedding dimensions: {data.get_shape()}")

# Sample for training
indices = data.sample_index(1000, timepoint=0)
batch = data.get_data()[indices]
```

--------------------------------

### Run TrajectoryNet with custom dataset

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/README.rst

Use TrajectoryNet with a custom dataset by specifying the path to the .npz file and the embedding name. This command is for running the main TrajectoryNet process.

```bash
python -m TrajectoryNet.main --dataset [PATH_TO_NPZ_FILE] --embedding_name [EMBEDDING_NAME]
```

--------------------------------

### Visualize PHATE Embedding

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Visualize the 2D PHATE embedding using a scatter plot. Color points by sample labels and customize the appearance with `cmap`, `ticks`, and `label_prefix`.

```python
scprep.plot.scatter2d(Y_phate, c=sample_labels, figsize=(12,8), cmap="Spectral",
                      ticks=False, label_prefix="PHATE")
```

--------------------------------

### Apply square root transformation

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Applies a square root transformation to the data to stabilize variance at zero.

```python
EBT_counts = scprep.transform.sqrt(EBT_counts)
```

--------------------------------

### Display Anndata Object Structure

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Displays the structure of the loaded Anndata object, showing its dimensions and available annotations. This is useful for verifying the data structure.

```python
adata
```

--------------------------------

### Download and Extract scRNAseq Data

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Downloads and extracts the scRNAseq dataset if it has not been downloaded already. The dataset is approximately 746MB.

```python
if not os.path.isdir(os.path.join(download_path, "scRNAseq", "T0_1A")):
    # need to download the data
    scprep.io.download.download_and_extract_zip(
        "https://md-datasets-public-files-prod.s3.eu-west-1.amazonaws.com/"
        "5739738f-d4dd-49f7-b8d1-5841abdbeb1e",
        download_path)
```

--------------------------------

### Load and Inspect Trajectories

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Loads the precomputed trajectory tensor and checks its dimensions.

```python
zs = np.load('../results/fig8_results/backward_trajectories.npy')
```

```python
zs.shape
```

--------------------------------

### Normalize library size

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Performs library size normalization on the input DataFrame.

```python
df = scprep.normalize.library_size_normalize(df)
```

--------------------------------

### Plot full cellular trajectories

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Visualizes the full trajectories by plotting the path of the top cells over the PHATE embedding.

```python
fig, ax = plt.subplots(1,1)
scprep.plot.scatter2d(phate_operator.graph.data_nu, c='Gray', alpha=0.1, ax=ax)
for i, gene in enumerate(end_genes):
    scprep.plot.scatter2d(phate_operator.graph.data_nu[sample_labels=='Day 24-27'][masks[gene]], ax=ax, 
                          label='%s - %s' % (end_points[i], end_genes[i]), c=colors[gene], ticks=[])

for gene in end_genes:
    for g in top_idxs[gene]:
        ax.plot(zss[:,g,0], zss[:,g,1], c=colors[gene])
plt.legend()
```

--------------------------------

### Process and Compute Moments for Velocity Analysis

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Performs filtering, normalization, and computes moments for velocity analysis using ScVelo. This prepares the Anndata object for velocity inference.

```python
scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=1000)
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)
scv.tl.velocity(adata, mode="stochastic")
scv.tl.velocity_graph(adata)
```

--------------------------------

### Visualize Velocity Embedding (Grid)

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Displays velocity vectors on a grid overlay on a specified embedding. Useful for visualizing velocity fields across different cell states. Requires 'velocity' and 'spliced' layers.

```python
scv.pl.velocity_embedding_grid(adata, color='Tmsb10', layer=['velocity', 'spliced'], arrow_size=1.5)
```

--------------------------------

### Visualize Velocity Embedding (Arrows)

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Visualizes velocity vectors on a specified embedding using arrows. Allows customization of arrow length, size, and DPI for clarity.

```python
scv.pl.velocity_embedding(adata, basis='umap', arrow_length=2, arrow_size=2, dpi=150)
```

--------------------------------

### Display adata.obs

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Shows the observation (adata.obs) DataFrame, which contains metadata for each cell, including the newly added 'sample_labels'.

```python
adata.obs
```

--------------------------------

### Load Coordinate and Day Information

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Loads coordinate data and cell day information from tab-separated text files into pandas DataFrames. The 'id' column is used as the index for both DataFrames.

```python
# Path to input files
FLE_COORDS_PATH ='data/fle_coords.txt'
FULL_DS_PATH = 'data/ExprMatrix.h5ad'
VAR_DS_PATH = 'data/ExprMatrix.var.genes.h5ad'
CELL_DAYS_PATH = 'data/cell_days.txt'
GENE_SETS_PATH = 'data/gene_sets.gmx'
GENE_SET_SCORES_PATH = 'data/gene_set_scores.csv'
CELL_SETS_PATH = 'data/cell_sets.gmt'

coord_df = pd.read_csv(PATH+FLE_COORDS_PATH, index_col='id', sep='\t')
days_df = pd.read_csv(PATH+CELL_DAYS_PATH, index_col='id', sep='\t')
```

--------------------------------

### Train TrajectoryNet Model

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/Example_Anndata_to_TrajectoryNet.ipynb

Executes the training process for a specified number of iterations and saves the output to the designated directory.

```python
!python -m TrajectoryNet.main --dataset adata.h5ad \
        --embedding_name "pca" \
        --max_dim 10 \
        --niter 10000 \
        --vecint 1e-4 \
        --save results/tmp
```

--------------------------------

### Load Dentate Gyrus Dataset

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Loads the built-in dentate gyrus dataset into an AnnData object. This dataset is suitable for velocity analysis.

```python
adata = scv.datasets.dentategyrus()
```

--------------------------------

### Plot gene expression patterns with Matplotlib

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Visualizes gene expression changes over time for specific cell subsets. Requires pre-computed inverse trajectories and end gene indices.

```python
fig, ax = plt.subplots(2,2, figsize=(6,6), sharex=True)
ax = ax.flatten()

for i, gene in enumerate(end_genes):
    for j, eg in enumerate(end_genes):
        for cell in top_idxs[eg]:
            ax[i].plot(np.linspace(1,5,100)[::-1], inverse[:,cell,end_gene_indexes[i]], c=colors[eg])
            ax[i].set_title(gene)
            ax[i].set_yticks([])
            ax[i].set_xticks(range(1,6))
```

--------------------------------

### Save trajectory data to NPZ

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Serializes dictionary data containing PHATE coordinates and embeddings to a compressed NumPy file.

```python
'pcs': phate_op.graph.data_nu[mask],
    'original_embedding': np.array(coord_df)[mask],
    'phate' : data_phate[mask],
    'sample_labels' : np.array(days_df['day']),
}
np.savez('../data/schiebinger.npz', **to_save)
```

--------------------------------

### Visualize Specific Gene Velocities

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/DentateGyrus-Load.ipynb

Plots the velocity of specified genes, optionally with a colorbar. Useful for examining the dynamics of individual genes of interest.

```python
var_names = ['Tmsb10', 'Camk2a', 'Ppp3ca', 'Igfbpl1']
scv.pl.velocity(adata, var_names=var_names, colorbar=True, ncols=2)
```

--------------------------------

### Load Trajectory Data from NPZ

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Loads trajectory data from an NPZ file using CustomData. Requires a Namespace object with configuration arguments like embedding_name and whiten. Allows sampling cells from specific timepoints and leaving out timepoints for validation.

```python
args = Namespace(
    embedding_name='pca',
    max_dim=10,
    whiten=True  # Standardize data
)

data = CustomData('trajectory_data.npz', args)

# Sample cells from specific timepoint
batch_size = 100
timepoint = 2
indices = data.sample_index(batch_size, timepoint)
batch = data.get_data()[indices]

# Leave out a timepoint for validation
data.leaveout_timepoint(2)
print(f"Remaining timepoints: {data.get_unique_times()}")
```

--------------------------------

### Evaluate TrajectoryNet with custom dataset

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/README.rst

Evaluate TrajectoryNet on a custom dataset. This command is used for the evaluation phase, requiring the dataset path and embedding name.

```bash
python -m TrajectoryNet.eval --dataset [PATH_TO_NPZ_FILE] --embedding_name [EMBEDDING_NAME]
```

--------------------------------

### Select and plot mitochondrial genes

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Identifies mitochondrial genes and plots their expression distribution.

```python
mito_genes = scprep.select.get_gene_set(EBT_counts, starts_with="MT-") # Get all mitochondrial genes. There are 14, FYI.
scprep.plot.plot_gene_set_expression(EBT_counts, genes=mito_genes, percentile=90)
```

--------------------------------

### Impute Trajectories with MAGIC

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Applies MAGIC imputation to the gene-space trajectories across all timepoints.

```python
trajectory_eb_magic = np.zeros((100,3332,68))
```

```python
m_op = magic.MAGIC()
for i in range(100):
    trajectory_eb_magic[i,:,:] = m_op.fit_transform(trajectory_eb[i,:,:])
```

--------------------------------

### Calculate inverse and gene indexes

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Computes the inverse transformation and identifies indices for end genes.

```python
inverse = np.dot(zss, phate_operator.graph.data_pca.components_[:5, genes_mask])
end_gene_indexes = [(np.where(genes_of_interest_full == gene)[0][0]) for gene in genes_of_interest_end]
```

--------------------------------

### Optimal Transport Interpolation

Source: https://context7.com/krishnaswamylab/trajectorynet/llms.txt

Interpolates between cell populations at different timepoints using optimal transport maps. Requires numpy and ot libraries.

```python
import numpy as np
import ot as pot
from TrajectoryNet.optimal_transport.emd import interpolate_with_ot

# Two populations at different timepoints
n_cells = 1000
n_dims = 10

population_t0 = np.random.randn(n_cells, n_dims)
population_t1 = np.random.randn(n_cells, n_dims) + 1.0  # Shifted

# Compute optimal transport map
cost_matrix = pot.dist(population_t0, population_t1)
transport_map = pot.emd(
    np.ones(n_cells) / n_cells,  # Uniform weights
    np.ones(n_cells) / n_cells,
    cost_matrix
)

# Interpolate at t=0.5
interp_frac = 0.5
n_interp = 1000

interpolated = interpolate_with_ot(
    population_t0,
    population_t1,
    transport_map,
    interp_frac,
    size=n_interp
)

print(f"Interpolated shape: {interpolated.shape}")  # (1000, 10)

# Multiple interpolation fractions
for t in [0.25, 0.5, 0.75]:
    interp_t = interpolate_with_ot(population_t0, population_t1,
                                    transport_map, t, size=500)
    print(f"t={t}: mean position = {interp_t.mean(axis=0)[:3]}")
```

--------------------------------

### Project Trajectories to Gene Space

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/EmbryoidBody_TrajectoryInference.ipynb

Projects the trajectories from the embedding space back into the original gene space.

```python
trajectory_gene_space = np.dot(zss, phate_operator.graph.data_pca.components_[:5,:])
```

```python
trajectory_gene_space.shape
```

--------------------------------

### Load Annotated Data Matrix

Source: https://github.com/krishnaswamylab/trajectorynet/blob/master/notebooks/WOT-Schiebinger-load.ipynb

Loads the full dataset, likely containing gene expression counts, from an h5ad file using scanpy. This object is typically an AnnData object.

```python
adata = scanpy.read_h5ad(PATH+FULL_DS_PATH)
```