Vector Field Analysis and Visualization

set.seed(42)
n_cells <- 500
n_genes <- 100

# Simulate trajectory
t_true <- seq(0, 1, length.out = n_cells)

# Create expression data (genes change along trajectory)
counts <- matrix(0, nrow = n_genes, ncol = n_cells)
for (i in 1:n_genes) {
    phase <- runif(1, 0, 2 * pi)
    amplitude <- runif(1, 3, 10)
    counts[i, ] <- rpois(n_cells, lambda = pmax(1, amplitude * (1 + sin(t_true * 2 * pi + phase))))
}
rownames(counts) <- paste0("Gene", seq_len(n_genes))
colnames(counts) <- paste0("Cell", seq_len(n_cells))

# Create Seurat object
srt <- CreateSeuratObject(counts = counts)

# Add UMAP embedding (curved trajectory)
umap1 <- t_true * 10 + rnorm(n_cells, sd = 0.5)
umap2 <- 4 * sin(t_true * 2 * pi) + rnorm(n_cells, sd = 0.3)

umap_embeddings <- matrix(c(umap1, umap2), ncol = 2)
rownames(umap_embeddings) <- colnames(srt)
colnames(umap_embeddings) <- c("UMAP_1", "UMAP_2")

srt[["umap"]] <- CreateDimReducObject(
    embeddings = umap_embeddings,
    key = "UMAP_",
    assay = "RNA"
)

# Add pseudotime
srt$pseudotime <- t_true + rnorm(n_cells, sd = 0.02)
srt$pseudotime <- pmax(0, pmin(1, srt$pseudotime))

# Simulate latent space (normally from model output)
n_latent <- 5
zs <- matrix(rnorm(n_cells * n_latent), ncol = n_latent)

# Make latent space correlate with trajectory
zs[, 1] <- t_true * 2 + rnorm(n_cells, sd = 0.2)
zs[, 2] <- sin(t_true * pi) + rnorm(n_cells, sd = 0.1)
rownames(zs) <- colnames(srt)

# Simulate vector field in latent space
vf <- matrix(0, nrow = n_cells, ncol = n_latent)
vf[, 1] <- 1 + rnorm(n_cells, sd = 0.1)  # Constant velocity in z1
vf[, 2] <- cos(t_true * pi) * pi / 2 + rnorm(n_cells, sd = 0.1)  # Derivative of sin
rownames(vf) <- colnames(srt)

# Store in Seurat object (as CellODE does)
srt@misc$X_zs <- zs
srt@misc$X_VF <- vf

# Compute cosine similarity matrix
T_mat <- cosine_similarity(
    zs = zs,
    vf = vf,
    n_neigh = 20
)

cat("Transition matrix dimensions:", dim(T_mat), "\n")
#> Transition matrix dimensions: 500 500
cat("Non-zero entries:", sum(T_mat != 0), "\n")
#> Non-zero entries: 50067

# Get embedding coordinates
E <- Embeddings(srt, "umap")

# Compute velocity in embedding space
V_embed <- vector_field_embedding(
    T_mat = T_mat,
    E = E,
    scale = 10
)

cat("Embedding velocity dimensions:", dim(V_embed), "\n")
#> Embedding velocity dimensions: 500 2
cat("Mean velocity magnitude:", mean(sqrt(rowSums(V_embed^2))), "\n")
#> Mean velocity magnitude: 0.240518

# Compute grid-based vector field
grid_result <- vector_field_embedding_grid(
    E = E,
    V = V_embed,
    smooth = 0.5,
    stream = FALSE,
    density = 1.0
)

cat("Grid points:", nrow(grid_result$E_grid), "\n")
#> Grid points: 830

# Basic pseudotime plot
p1 <- plot_pseudotime(
    srt, 
    time_key = "pseudotime",
    arrow_show = TRUE,
    arrow_density = 1.5,
    point_palette = "Spectral",
    title = "Pseudotime with Direction Arrows"
)
print(p1)

# Vector field - grid mode
p2 <- plot_vector_field(
    srt,
    plot_type = "grid",
    color_by = "pseudotime",
    arrow_density = 0.8,
    arrow_color = "grey20",
    point_palette = "Spectral",
    title = "Vector Field (Grid Mode)"
)
print(p2)

# Vector field - stream mode with colored arrows
p3 <- plot_vector_field(
    srt,
    plot_type = "stream",
    color_by = "pseudotime",
    arrow_density = 1.2,
    point_palette = "viridis",
    title = "Vector Field (Stream Mode)"
)
print(p3)

# Vector field - raw mode (per-cell arrows)
p4 <- plot_vector_field(
    srt,
    plot_type = "raw",
    color_by = "pseudotime",
    arrow_density = 0.3,  # Subsample for clarity
    arrow_color = "navy",
    point_palette = "magma",
    title = "Vector Field (Raw Mode)"
)
print(p4)

# Custom arrow styling
p_custom <- plot_vector_field(
    srt,
    plot_type = "grid",
    color_by = "pseudotime",
    arrow_type = "closed",  # Filled arrowheads
    arrow_color = "#E63946",  # Red arrows
    arrow_density = 0.6,
    arrow_length = 0.2,
    arrow_width = 0.8,
    arrow_angle = 30,
    point_palette = "RdYlBu",
    point_size = 0.8,
    title = "Custom Arrow Styling"
)
print(p_custom)

# Clean figure for publication
p_pub <- plot_pseudotime(
    srt,
    time_key = "pseudotime",
    arrow_show = TRUE,
    arrow_density = 1.0,
    arrow_type = "open",
    arrow_color = "black",
    point_palette = "viridis",
    point_size = 1.2,
    show_axes = FALSE,
    title = NULL
)
print(p_pub)

# More neighbors = smoother but less detailed
T_smooth <- cosine_similarity(zs, vf, n_neigh = 30)

# Fewer neighbors = more detail but noisier
T_detail <- cosine_similarity(zs, vf, n_neigh = 10)

# Publication figure (fewer arrows)
plot_vector_field(srt, arrow_density = 0.3)

# Exploration (more arrows)
plot_vector_field(srt, arrow_density = 1.0)

# Include self-transition for smoother results
V_with_self <- vector_field_embedding(T_mat, E, self_transition = TRUE)