Transcript-based cell type prediction

import os
from torch_geometric.data import OnDiskDataset
import random
from torch_geometric.data import Data, Batch
from torch_geometric.loader import LinkNeighborLoader, NeighborLoader
from torch_geometric.nn import GraphSAGE, GAT, GATConv, GATv2Conv, AttentionalAggregation
import copy
from tqdm import tqdm
import pandas as pd
import numpy as np
import torch
from torch_geometric import seed_everything
import sys 
import spatialrna.spatialrna_ondisk as spod
from spatialrna import SpatialRNA

## for reproduciable results
seed = 1024
random.seed(seed)  # python random seed
np.random.seed(seed)  # numpy random seed
torch.manual_seed(seed)  # pytorch random seed
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
seed_everything(seed)

!head "../data/spatialRNA_samples_label_pred/VUHD116A/raw/VUHD116A.csv"

#celltype_labeled

unique_ct = ['Capillary', 'Inflammatory FBs', 'Venous', 'Alveolar Macrophages',
       'AT1', 'Alveolar FBs', 'Proliferating Myeloid', 'AT2',
       'Monocytes/MDMs', 'Neutrophils', 'Interstitial Macrophages',
       'Plasma', 'NK/NKT', 'Proliferating AT2', 'SMCs/Pericytes',
       'Migratory DCs', 'CD4+ T-cells', 'Mast', 'Secretory',
       'Transitional AT2', 'Multiciliated', 'SPP1+ Macrophages',
       'Proliferating FBs', 'Arteriole', 'Lymphatic', 'cDCs',
       'CD8+ T-cells', 'Macrophages - IFN-activated', 'RASC',
       'Activated Fibrotic FBs', 'Myofibroblasts', 'Subpleural FBs',
       'Adventitial FBs', 'B cells', 'Tregs', 'Goblet', 'Basal', 'PNEC',
       'Proliferating Airway', 'Proliferating T-cells', 'Basophils',
       'pDCs']

cell_label_dict = {}
for i,ct in enumerate(unique_ct):
    cell_label_dict[ct] = i

print(cell_label_dict)

# reverse the key, item in cell_label_dict
code_to_label = {v: k for k, v in cell_label_dict.items()}
print(code_to_label)

gene_panel = pd.read_csv("../../../case_study_ipf_revision/resources/xenium_gene_panel.csv")
gene_panel.head(3)

gene_list = np.unique(gene_panel.x)
x = torch.tensor(np.arange(gene_list.shape[0]))
one_hot_encoding = dict(zip(gene_list, F.one_hot(x, num_classes=gene_list.shape[0])))

for k in one_hot_encoding.keys():
    one_hot_encoding[k] = one_hot_encoding[k].double()

one_hot_encoding_int = dict()
for key in one_hot_encoding.keys():
    one_hot_encoding_int[key] = one_hot_encoding[key].argmax()

SpatialRNA(
    root="../data/spatialRNA_samples_label_pred/VUHD116A",
    sample_name="VUHD116A",
    radius_r=3.0,
    tile_by_dim="y_location",
    dim_x="x_location",
    dim_y="y_location",
    num_tiles=5,
    force_reload=True,
    feature_col="feature_name",
    force_resample=False,
    process_mode="tile",
    load_type="blank",
    pred_label_col = "final_CT",
    pred_label_map = cell_label_dict,
    pred_label_cell_id_col= "cell_id",
    process_tile_ids=[x for x in range(5)],
    one_hot_encoding=one_hot_encoding_int,
    log=False
)

ls -hl ../data/spatialRNA_samples_label_pred/VUHD116A/processed

import torch
from torch_geometric.loader import DataLoader
from sklearn.model_selection import train_test_split
import torch
from time import time

myod = spod.SpatialRNAOnDiskDataset(root = "../data/spatialRNA_samples_label_pred/",pt_dir="processed")

myod.len()
#myod.multi_get([0,1])

#construct dataset

# Split indices: 80% train, 20% validation
indices = list(range(len(myod)))
#train_idx, val_idx = train_test_split(indices, test_size=0.2, random_state=42)

#train_idx[0:5], val_idx[0:5],len(train_idx),

# Lazy subsets
# PyG supports this
train_dataset = myod.index_select(indices)   
#val_dataset   = myod.index_select(val_idx)
# DataLoaders
train_loader = DataLoader(train_dataset, batch_size=5, shuffle=True, num_workers=2)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_label_pred = GraphSAGE(
    gene_panel.shape[0],
    hidden_channels=50,
    num_layers=3,
    out_channels=42,
).to(device)

model_label_pred

model_label_pred

import spatialrna.train_val as train_val
optimizer = torch.optim.Adam(model_label_pred.parameters(), lr=0.001)
# Early stopping setup
best_acc = 0.0
patience = 3
counter = 0
best_model_wts = copy.deepcopy(model_label_pred.state_dict())
for i in range(1, 50):
    print("Epoch", i)
    loss, acc = train_val.train_label_pred(
        model=model_label_pred,
        device=device,
        train_loader=train_loader,
        num_classes=gene_panel.shape[0],
        batch_size=2048,
        num_train_nodes=200000,
        optimizer=optimizer,
        verbose=True
    )
    print(f"loss={loss:.4f}, acc={acc:.4f}")
    if acc > best_acc:
        best_acc = acc
        counter = 0
        best_model_wts = copy.deepcopy(model_label_pred.state_dict())
    else:
        counter += 1
        if counter >= patience:
            break
model_label_pred.load_state_dict(best_model_wts)

torch.save(model_label_pred, "../data/spatialRNA_samples_label_pred/pred_trained.pt")

model_label_pred_load = torch.load("../data/spatialRNA_samples_label_pred/pred_trained.pt",weights_only=False)
model_label_pred_load = model_label_pred_load.to(device)

torch.serialization.add_safe_globals([GraphSAGE,Data])

train_val.inference(model=model_label_pred_load, 
                    device=device,
                    sample_name="VUHD116A",
                    root="../data/spatialRNA_samples_label_pred/VUHD116A/",
                    tile_id=[x for x in range(5)],
                    num_classes=gene_panel.shape[0],
                    num_neighbors=[20,10])

all_npy = [np.load(f"../data/spatialRNA_samples_label_pred/VUHD116A/embedding/VUHD116A_data_tile{x}.npy") for x in range(5)]
all_npy = np.concat(all_npy)

all_ids = [pd.read_csv(f"../data/spatialRNA_samples_label_pred/VUHD116A/embedding/VUHD116A_data_tile{x}input_tx_id.csv") for x in range(5)]
all_ids = pd.concat(all_ids)

all_npy.shape, all_ids.shape

tx_meta = pd.read_csv("../data/spatialRNA_samples_label_pred/VUHD116A/raw/VUHD116A.csv")

tx_meta.head(3)

tx_meta = tx_meta.iloc[all_ids["tx_id"].values]

nuc_poly = pd.read_csv("/mnt/beegfs/mccarthy/backed_up/general/rlyu/Dataset/LFST_2022/GEO_2025/VUHD116A/relabel_output-XETG00048__0003817__VUHD116A__20230308__003730/outs/nucleus_boundaries.csv.gz")
nuc_poly['cell_id'] = "VUHD116A_"+nuc_poly['cell_id']
nuc_poly = nuc_poly[nuc_poly.cell_id.isin(tx_meta.cell_id)]

nuc_poly.head(3)

import tifffile
dapi_auto_focus = "/mnt/beegfs/mccarthy/backed_up/general/rlyu/Dataset/LFST_2022/GEO_2025/VUHD116A/relabel_output-XETG00048__0003817__VUHD116A__20230308__003730/outs/morphology_focus.ome.tif"

with tifffile.TiffFile(
    dapi_auto_focus,
) as tif:
    for tag in tif.pages[0].tags.values():
        if tag.name == "ImageDescription":
            print(tag.name+":", tag.value)

from scipy.spatial.distance import cdist
import matplotlib.pyplot as plt
from matplotlib.colors import to_rgb,to_rgba
#subset node meta and node embedding matrix according to a selected cell id
def get_cell_tx_and_embedding(cell_id, df = tx_meta, out = all_npy):
    df = df.reset_index()
    sub_df = df[df.cell_id == cell_id]
    sub_out = out[sub_df.index,:]
    return sub_df,sub_out
# find the major label of an array
def most_frequent_numpy(arr):
    values, counts = np.unique(arr, return_counts=True)
    return values[np.argmax(counts)]

# find the 20 nearest cells by centroids
def find_nearest_cells(target_cell_id, centroids, n=20):
    # Get the centroid of the target cell
    target_centroid = centroids[centroids['cell_id'] == target_cell_id][['vertex_x', 'vertex_y']].values
    # Calculate distances between the target centroid and all other centroids
    distances = cdist(target_centroid, centroids[['vertex_x', 'vertex_y']], metric='euclidean')[0]
    # Add distances to the centroids DataFrame
    centroids['distance'] = distances
    # Sort by distance and exclude the target cell itself
    nearest_cells = centroids.sort_values(by='distance').head(n)
    return nearest_cells

# Group by cell_id and calculate centroids for cell ids:
cell_centroids = nuc_poly.groupby('cell_id').agg({'vertex_x': 'mean', 'vertex_y': 'mean'}).reset_index()
find_nearest_cells("VUHD116A_aaaadchd-1",centroids=cell_centroids,n =5)

l2_img = tifffile.imread(
    dapi_auto_focus, level=0,
    is_ome=True, aszarr=False)

# Examine shape of array (number of channels, height, width), e.g. (4, 40867, 31318)
l2_img.shape

plot_nearby_nuc_withimage_scale_tx("VUHD116A_aaaaaopa-1",with_confi=True,alpha_by="label",show_dapi=False)

plot_nearby_nuc_withimage_scale_tx("VUHD116A_aaaaaopa-1",alpha_by="label",show_dapi=True)

pio.renderers.default = "plotly_mimetype"
i_plot = plot_nearby_cells_interactive("VUHD116A_aaaaaopa-1",alpha_by="label",nearest=30)

pio.renderers.default = "plotly_mimetype"
i_plot = plot_nearby_cells_interactive("VUHD116A_aaaaaopa-1",alpha_by="prediction",nearest=30)