EM_optimization/maximization.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
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
import numpy as np
import pandas as pd
from assignment import assign


def maximization (itemList, itemLabel, labeling, scoreMtx, cutoffDict, allContexts, allTemplates, markerAssignment = False):
    support_minScore = cutoffDict["minimal_score_for_support"]; assign_minPctSupport = cutoffDict["minimal_percent_for_support"]
    unassign_maxPctSupport = cutoffDict.get ("maximal_percent_for_not_support", np.inf)
    pre_assignment = assign (itemList, itemLabel, labeling, scoreMtx, support_minScore, assign_minPctSupport, unassign_maxPctSupport,
                             allContexts, allTemplates, uniqueTemplateAssignment = True)
    if markerAssignment:
        df = pre_assignment.sort_values (["index", itemLabel])
        df["score"] = [pre_assignment.loc[idx, f"avgScore_{pre_assignment.loc[idx, 'template']}"] for idx in pre_assignment.index]
        occurrence = df.groupby (["index", itemLabel])["template"].value_counts ().reset_index (level = 2)
        df = df.set_index (["index", itemLabel]).sort_index (); tmp_list = list ()
        for idx in set (occurrence.index):
            tmpDF = df.loc[[idx]].sort_values ("score"); tmpOcc = occurrence.loc[[idx]]
            if tmpDF.shape[0] == 1:
                tmp_list.append (tmpDF.reset_index ().drop ("score", axis = 1))
            else:
                for temp in tmpOcc.loc[tmpOcc["count"] == 1, "template"]:
                    if temp == tmpDF["template"].iloc[-1]:
                        tmp_list.append (tmpDF.loc[tmpDF["template"] == temp].reset_index ().drop ("score", axis = 1))
        if len (tmp_list) == 0:
            pre_assignment = pd.concat ([pd.DataFrame ({"index": pd.Series (dtype = int)}),
                                         pd.DataFrame (columns = [itemLabel, "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:
            pre_assignment = pd.concat (tmp_list, axis = 0).sort_values (["index", "context", "template"]).reset_index (drop = True)
    if not pre_assignment.empty:
        tmp = pre_assignment.drop_duplicates (["index", itemLabel])
        with np.errstate (divide = "ignore", invalid = "ignore"):
            avgScore = np.array ([[(scoreMtx[i, :, labeling.get (allContexts[j], list ())]).mean (axis = 0)
                                   for i in range (len (allTemplates))] for j in range (len (allContexts))])
        avgScore = pd.DataFrame (avgScore[:, :, tmp["index"]].max (axis = 1).T, index = tmp[itemLabel].values, columns = allContexts)
        assigned = pre_assignment.pivot (index = itemLabel, columns = "context", values = "template").replace (np.nan, "")
        assigned = (assigned != "").rename_axis (None).rename_axis (None, axis = 1)
        if assigned.shape[1] < len (allContexts):
            assigned = pd.concat ([assigned, pd.DataFrame (False, index = assigned.index, columns = list (set (allContexts) - set (assigned.columns)))], axis = 1)
        assigned = assigned.loc[avgScore.index, avgScore.columns]
        bestScore = pd.DataFrame ({"unassigned": avgScore.mask (assigned).max (axis = 1, skipna = True),
                                   "assigned": avgScore.mask (~assigned).min (axis = 1, skipna = True)})
        items = bestScore.loc[(bestScore["assigned"] > support_minScore) & (bestScore["unassigned"] < support_minScore)].index
        pre_assignment = pre_assignment.loc[pre_assignment[itemLabel].isin (items)]
    assignment = list ()
    for temp in allTemplates:
        sub_assignment = pre_assignment.loc[pre_assignment["template"] == temp]
        maxSize = pd.Series ([25 * len (labeling.get (context, list ())) for context in allContexts], index = allContexts)
        assignment += [sub_assignment.loc[sub_assignment["context"] == C].sort_values ([f"pctSupport_{temp}", f"avgScore_{temp}"],
                                                                                       ascending = False).head (maxSize[C])
                       for C in allContexts]
    if len (assignment) == 0:
        assignment = pd.concat ([pd.DataFrame ({"index": pd.Series (dtype = int)}),
                                 pd.DataFrame (columns = [itemLabel, "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")
    numItems = assignment.value_counts (["context", "template"]); cols = assignment.columns
    assignment = assignment.set_index (["context", "template"]).loc[numItems[numItems > 5].index].reset_index ()
    assignment = assignment[cols].sort_values ("index").reset_index (drop = True)
    for cont in set (assignment["context"]):
        shortCont = cont.replace ("__backup", "")
        if cont != shortCont and (not shortCont in set (assignment["context"])):
            assignment.loc[assignment["context"] == cont, "context"] = shortCont
    newLabeling = assignment.groupby ("context")["index"].agg (list).to_dict (); newContexts = sorted (set (assignment["context"]))
    return assignment, newLabeling, newContexts


def mergeBackup (itemLabel, assigned, assigned_backup, backup):
    merged = assigned.merge (assigned_backup, on = ["index", itemLabel], how = "right", suffixes = ("_assigned", "")).replace (np.nan, "")
    merged = merged.loc[((merged["context"] == merged["context_assigned"]) & (merged["template"] == merged["template_assigned"])) |
                        (merged["context_assigned"] == ""), assigned_backup.columns]
    numItems = merged.value_counts (["context", "template"])
    merged = merged.set_index (["context", "template"]).loc[numItems[numItems > 5].index].reset_index ()[assigned_backup.columns]
    newBackup = backup.copy ()
    for cont in set (backup.keys ()) - set (merged["context"]):
        del newBackup[cont]
    for cont in set (merged["context"]) & set (assigned["context"]):
        merged.loc[merged["context"] == cont, "context"] += "__backup"
        newBackup[f"{cont}__backup"] = newBackup[cont].copy (); del newBackup[cont]
    assignment = pd.concat ([assigned, merged], axis = 0).sort_values (["index", "context", "template"]).reset_index (drop = True)
    itemLabeling = assignment.groupby ("context")["index"].agg (list).to_dict (); newContexts = sorted (set (assignment["context"]))
    return assignment, itemLabeling, newContexts, newBackup