analysis.py

"""
script to:
    create initial plots of method's embedding space
    get target prioritizations depending on the method 
"""
from library import *

hostname = socket.gethostname()
print(f"The hostname of this server is: {hostname}")

parser = argparse.ArgumentParser()
parser.add_argument("--method", type=str, default="raw", help="which method to use for embedding: raw, umap, scimilarity, geneformer, scgpt, uce")
parser.add_argument("--dataset", type=str, default="full", help="determines which dataset to load and process; full is the entire imru dataset")
opt = parser.parse_args()
print(opt)

if not os.path.isdir(f"outputs/{opt.dataset}"):
    os.makedirs(f"outputs/{opt.dataset}")

adata = sc.read_h5ad(f"data/{opt.method}_imru_{opt.dataset}_processed.h5ad")
print(f"loaded: data/{opt.method}_imru_{opt.dataset}_processed.h5ad")
control_strings = set(["SAFE_TARGET", "NO-TARGET"])
unique_perturbations = set(['SH3PXD2A', 'EMCN', 'RAB11FIP3', 'CKAP5', 'POLR2K'] + ['VWF', 'HNRNPA3', 'TRIM13', 'CORO1C', 'DDX27', 'PACSIN2']) ##set of perturbations that are only present in either treated OR untreated condition, but not both
anchor_state_candidate_treatment_types = [("Untreated_SAFE_TARGET", "Treated")] ##each tuple is anchor_state, desired treatment type of the candidates closest to anchor_state -- for most cases we want the candidates to be the other treatment_type than the anchor; if you want to find perturbations similar to some target of interest (e.g. T1) in the same treatment condition we can do: ("Untreated_T1", "Untreated"); No-target example: ("Untreated_NO-TARGET", "Treated"); reverse example of inflammation induction via knockdown: ("Treated_SAFE_TARGET", "Untreated") 
normalize = True ##DL embeddings are already L2 normalized by default (except STATE)
knockdown_threshold = 500.0 ##e.g. don't filter out any cells based on kd efficiency
min_sample = 0 ##e.g. no filtering targets based on cell count of replicates in treatment condition -- can filter downstream later if wanted
distance_metric = "cosine" #or euclidean (note that these are monotonically related if embeddings are l2 normalized, so rankings will be the same and it doesn't matter)
embedding_type = "method_embedding" #or umap_embedding

##runner calls
#get ranked list of targets for each anchor state
sorted_perts, count_map = get_top_perturbations(adata, method=opt.method, top_n=10, anchor_state_candidate_treatment_types=anchor_state_candidate_treatment_types, min_sample=min_sample, dataset=opt.dataset, knockdown_threshold=knockdown_threshold, normalize=normalize, distance_metric=distance_metric, embedding_type=embedding_type) 
pickle.dump((sorted_perts, count_map), open(f"pickles/{opt.dataset}_{opt.method}_sorted_perts_knockdown_threshold_{knockdown_threshold}_min_sample_{min_sample}.pkl", "wb"))
print(f"wrote: pickles/{opt.dataset}_{opt.method}_sorted_perts_knockdown_threshold_{knockdown_threshold}_min_sample_{min_sample}.pkl")

##umap visualizations
untreated_with_treated_control_umap(adata, control_strings, opt.method, opt.dataset, normalize)
if opt.method == "raw":
    untreated_with_treated_control_umap(adata, control_strings, opt.method, opt.dataset, normalize, just_inflammatory_genes=True)
##plot random selection of perturbations vs control
perturbation_vs_control_umap(adata, control_strings, unique_perturbations, n=50, method=opt.method, dataset=opt.dataset, save_dir=f"outputs/{opt.dataset}")
##plot internal expertise derived (+) control anchors
perturbation_vs_control_umap(adata, control_strings, unique_perturbations, plot_perts=["TNFRSF1A", "TRADD", "JUNB", "JUND", "NFKB1", "NFKB2"], n=None, method=opt.method, direction="to_healthy", dataset=opt.dataset, save_dir=f"outputs/{opt.dataset}")

##assess embedding quality clustering labels by condition and gene taret
score_perturbation_clusters(adata, control_strings, method=opt.method, dataset=opt.dataset, knockdown_threshold=knockdown_threshold, normalize=normalize) ##requires a lot of compute time