task_cyto_batch_integration/src/metrics/average_batch_r2/script.py at fc5efe3d33c6c99a8bfe563d9779ef73963e1bcd · openproblems-bio/task_cyto_batch_integration · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import sys

import anndata as ad
import numpy as np

## VIASH START
# The following code has been auto-generated by Viash.
par = {
    "input_unintegrated": "resources_test/task_cyto_batch_integration/mouse_spleen_flow_cytometry_subset/unintegrated.h5ad",
    "input_integrated_split1": "resources_test/task_cyto_batch_integration/mouse_spleen_flow_cytometry_subset/integrated_split1.h5ad",
    "input_integrated_split2": "resources_test/task_cyto_batch_integration/mouse_spleen_flow_cytometry_subset/integrated_split2.h5ad",
    "output": "resources_test/task_cyto_batch_integration/mouse_spleen_flow_cytometry_subset/score.h5ad",
}
meta = {"name": "average_batch_r2"}

## VIASH END

sys.path.append(meta["resources_dir"])
from helper import batch_r2
from helper_functions import (
    get_obs_var_for_integrated,
    remove_unlabelled,
    subset_markers_tocorrect,
    subset_nocontrols,
)

print("Reading input files", flush=True)
integrated_s1 = ad.read_h5ad(par["input_integrated_split1"])
integrated_s2 = ad.read_h5ad(par["input_integrated_split2"])
unintegrated = ad.read_h5ad(par["input_unintegrated"])

print("Formatting input files", flush=True)
integrated_s1, integrated_s2 = get_obs_var_for_integrated(
    integrated_s1, integrated_s2, unintegrated
)

integrated_s1 = subset_nocontrols(integrated_s1)
integrated_s1 = subset_markers_tocorrect(integrated_s1)

integrated_s2 = subset_nocontrols(integrated_s2)
integrated_s2 = subset_markers_tocorrect(integrated_s2)

print("Computing average_batch_r2 per cell type", flush=True)

donor_list = integrated_s1.obs["donor"].unique()

r2_values = []
r2_info = []

for donor in donor_list:
    s1_view = integrated_s1[integrated_s1.obs["donor"] == donor]  # split 1
    s1_view = remove_unlabelled(s1_view)
    s2_view = integrated_s2[integrated_s2.obs["donor"] == donor]  # split 2
    s2_view = remove_unlabelled(s2_view)

    ct_list = s1_view.obs["cell_type"].unique()

    for ct in ct_list:
        s1_view_ct = s1_view[s1_view.obs["cell_type"] == ct]
        s2_view_ct = s2_view[s2_view.obs["cell_type"] == ct]
        if (
            s1_view_ct.shape[0] < 20 or s2_view_ct.shape[0] < 20
        ):  # Skip Rˆ2 calculation if there are less than 20 cells
            print(
                f"Warning: Rˆ2 not computed for donor {donor} cell type {ct}. Too few cells were present: {s1_view_ct.shape[0]} for split 1 and {s2_view_ct.shape[0]} for split 2"
            )
            continue

        r2_list, marker_list = batch_r2(s1_view_ct, s2_view_ct)

        marker_list = [ct + "_" + donor + "_" + x for x in marker_list]
        r2_info = [*r2_info, *marker_list]
        r2_values = [*r2_values, *r2_list]

average_batch_r2_ct = np.mean(r2_values)
r2_collection_ct = dict(zip(r2_info, r2_values))  # kept for debugging only
print("average_batch_r2_ct", average_batch_r2_ct)

print("Write output AnnData to file", flush=True)
output = ad.AnnData(
    uns={
        "dataset_id": integrated_s1.uns["dataset_id"],
        "method_id": integrated_s1.uns["method_id"],
        "metric_ids": ["average_batch_r2_ct"],
        "metric_values": [average_batch_r2_ct],
        "r2_collection_ct": r2_collection_ct,
    }
)
output.write_h5ad(par["output"], compression="gzip")