RobroKools's picture
Upload 44 files
e59f78e verified
import scanpy as sc
import numpy as np
from sklearn.model_selection import train_test_split
import os
import warnings
warnings.filterwarnings("ignore", category=FutureWarning, module="anndata")
warnings.filterwarnings("ignore", message="Moving element from .uns")
def process():
os.makedirs("celldreamer/data/processed", exist_ok=True)
adata = sc.read("celldreamer/data/raw/panc8_raw.h5ad")
sc.pp.filter_cells(adata, min_genes=200)
sc.pp.filter_genes(adata, min_cells=3)
print(f"cleaned Shape: {adata.shape}")
print("getting K-nearest nieghbors")
sc.pp.pca(adata, n_comps=50)
sc.pp.neighbors(adata, n_neighbors=30, n_pcs=20)
sc.tl.diffmap(adata)
# find step 0 stem cell
try:
root_candidates = np.where(adata.obs['celltype'].str.contains('ductal', case=False))[0]
adata.uns['iroot'] = root_candidates[0] if len(root_candidates) > 0 else 0
except:
adata.uns['iroot'] = 0
sc.tl.dpt(adata)
# create t,t+1 pairs
print("creating pairs")
graph = adata.obsp['connectivities']
times = adata.obs['dpt_pseudotime'].values
pairs = []
rows, cols = graph.nonzero()
for i, j in zip(rows, cols):
t_i, t_j = times[i], times[j]
# max time diff is 0.1 for ~similar time diffs
if t_j > t_i and (t_j - t_i) < 0.1:
pairs.append([i, j])
pairs = np.array(pairs)
train, temp = train_test_split(pairs, test_size=0.2, random_state=42)
val, test = train_test_split(temp, test_size=0.5, random_state=42)
np.save("celldreamer/data/processed/train_pairs.npy", train)
np.save("celldreamer/data/processed/val_pairs.npy", val)
np.save("celldreamer/data/processed/test_pairs.npy", test)
print(f"Train({len(train)}), Val({len(val)}), Test({len(test)})")
adata.write("celldreamer/data/processed/cleaned.h5ad")
np.save("celldreamer/data/processed/full_set.npy", pairs)