EM_optimization/assignment.py at master · zimmerlab/EM_optimization · 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
import pandas as pd


def assign (allItems, itemType, labeling, scoreMtx, support_minScore, assign_minPctSupport, unassign_maxPctSupport,
            allContexts, allTemplates, uniqueTemplateAssignment = False):
    assignment = list ()
    for context in allContexts:
        if context not in labeling:
            continue
        if itemType == "feature" or itemType == "edge":
            score = scoreMtx[:, :, labeling[context]]
            pctSupport = pd.DataFrame ((score > support_minScore).mean (axis = 2).T, columns = allTemplates)
        elif itemType == "sample":
            score = scoreMtx[:, labeling[context], :]
            pctSupport = pd.DataFrame ((score > support_minScore).mean (axis = 1).T, columns = allTemplates)
        else:
            raise ValueError
        templateAssignment = list ()
        for temp in allTemplates:
            supported = pctSupport.loc[(pctSupport[temp] > assign_minPctSupport) &
                                       (pctSupport.drop (temp, axis = 1) < unassign_maxPctSupport).all (axis = 1)]
            if itemType == "feature" or itemType == "edge":
                avgScore = pd.DataFrame (score[:, supported.index, :].mean (axis = 2).T,
                                         index = supported.index, columns = allTemplates)
            elif itemType == "sample":
                avgScore = pd.DataFrame (score[:, :, supported.index].mean (axis = 1).T,
                                         index = supported.index, columns = allTemplates)
            else:
                raise ValueError
            supported = pd.concat ([pd.DataFrame ({itemType: [allItems[x] for x in supported.index]}, index = supported.index),
                                    pd.DataFrame ({"context": context, "template": temp}, index = supported.index),
                                    avgScore.rename (columns = {temp: f"avgScore_{temp}" for temp in allTemplates}).round (3),
                                    supported.rename (columns = {temp: f"pctSupport_{temp}" for temp in allTemplates}).round (3)],
                                   axis = 1)
            if not supported.empty:
                templateAssignment.append (supported.reset_index ())
        if len (templateAssignment) > 0:
            templateAssignment = pd.concat (templateAssignment, axis = 0, ignore_index = True)
            assigned = pd.DataFrame ({"template": templateAssignment["template"],
                                      "by_avgScore": templateAssignment.filter (regex = "^avgScore_").idxmax (axis = 1, skipna = True)\
                                        .str.split ("avgScore_", expand = True)[1],
                                      "by_pctSupport": templateAssignment.filter (regex = "^pctSupport_").idxmax (axis = 1, skipna = True)\
                                        .str.split ("pctSupport_", expand = True)[1]})
            templateAssignment = templateAssignment.loc[(assigned["template"] == assigned["by_avgScore"]) &\
                                                        (assigned["template"] == assigned["by_pctSupport"])]
            numAssignedTemplates = templateAssignment.value_counts ("index")
            uniqueTemp = numAssignedTemplates[numAssignedTemplates == 1].index
            if uniqueTemplateAssignment:
                templateAssignment = templateAssignment.loc[templateAssignment["index"].isin (uniqueTemp)]
            else:
                templateAssignment.loc[~templateAssignment["index"].isin (uniqueTemp), "template"] = "ALL"
                templateAssignment = templateAssignment.drop_duplicates ()
            if not templateAssignment.empty:
                assignment.append (templateAssignment)
    if len (assignment) == 0:
        assignment = pd.concat ([pd.DataFrame ({"index": pd.Series (dtype = int)}),
                                 pd.DataFrame (columns = [itemType, "context", "template"], dtype = str),
                                 pd.DataFrame (columns = [f"avgScore_{temp}" for temp in allTemplates], dtype = float),
                                 pd.DataFrame (columns = [f"pctSupport_{temp}" for temp in allTemplates], dtype = float)],
                                axis = 1)
    else:
        assignment = pd.concat (assignment, axis = 0).sort_values (["index", "context", "template"]).reset_index (drop = True)
    return assignment