d2_path.py 3.38 KB
Newer Older
Yoann Dufresne's avatar
Yoann Dufresne committed
1
2
import networkx as nx

3

Yoann Dufresne's avatar
Yoann Dufresne committed
4
5
""" Represent an udg path into a d2g graph
"""
Yoann Dufresne's avatar
Yoann Dufresne committed
6
class Path(list):
7

Yoann Dufresne's avatar
Yoann Dufresne committed
8
    def __init__(self, d2g):
9
        super(Path, self).__init__()
Yoann Dufresne's avatar
Yoann Dufresne committed
10
        self.penalty = 0
Yoann Dufresne's avatar
Yoann Dufresne committed
11
12
        self.d2g = d2g
        self.covering_variables = {x: 0 for x in self.d2g.barcode_edge_idxs.values()}
13

Yoann Dufresne's avatar
Yoann Dufresne committed
14
    def add_path(self, udgs):
Yoann Dufresne's avatar
Yoann Dufresne committed
15
16
17
        if len(udgs) == 0:
            return

Yoann Dufresne's avatar
Yoann Dufresne committed
18
19
20
21
22
23
        for udg in udgs:
            # Register edges
            for barcode_edge in udg.edges:
                edge_idx = self.d2g.barcode_edge_idxs[barcode_edge]
                self.covering_variables[edge_idx] += 1

Yoann Dufresne's avatar
Yoann Dufresne committed
24
25
26
27
28
29
30
31
32
33
        # Special case for previously empty path
        if len(self) == 0:
            # 4 because it's the ideal case (1 node of difference with same length on 1 shift.
            self.penalty = 4
            self.append(udgs[0])
            udgs = udgs[1:]

        # Add udg one by one
        for udg in udgs:
            # Compute distance
Yoann Dufresne's avatar
Yoann Dufresne committed
34
            dist = self.d2g[str(udg.idx)][str(self[-1].idx)]["distance"]
Yoann Dufresne's avatar
Yoann Dufresne committed
35
36
37
38
            # Update penalty regarding distance
            self.penalty += pow(dist, 2)
            # Add the node
            self.append(udg)
Yoann Dufresne's avatar
Yoann Dufresne committed
39

Yoann Dufresne's avatar
Yoann Dufresne committed
40
41
    def normalized_penalty(self):
        return self.penalty / len(self)
Yoann Dufresne's avatar
Yoann Dufresne committed
42

Yoann Dufresne's avatar
Yoann Dufresne committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
    def covering_score(self):
        covered_num = 0
        for val in self.covering_variables.values():
            if val > 0:
                covered_num += 1

        return covered_num / len(self.covering_variables)

    """ Count the number of variable covered by path that are not self covered. 
    """
    def covering_difference(self, path):
        self_covered = [var for var, num in self.covering_variables.items() if num > 0]
        self_covered = frozenset(self_covered)

        num_covered = 0
        for var, val in path.covering_variables.items():
            if val > 0 and var not in self_covered:
                num_covered += 1

        return num_covered
Yoann Dufresne's avatar
Yoann Dufresne committed
63

Yoann Dufresne's avatar
Yoann Dufresne committed
64
    def save_path(self, filename):
Yoann Dufresne's avatar
Yoann Dufresne committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
        d2p = nx.Graph()
        # Add the nodes
        for udg in self:
            d2p.add_node(udg.idx)
            d2p.nodes[udg.idx]["center"] = udg.center
            d2p.nodes[udg.idx]["udg"] = str(udg)
            d2p.nodes[udg.idx]["score"] = f"{udg.score}/{udg.get_optimal_score()}"

        # add the edges
        for idx in range(len(self)-1):
            udg1 = self[idx]
            udg2 = self[idx+1]

            d2p.add_edge(udg1.idx, udg2.idx)

Yoann Dufresne's avatar
Yoann Dufresne committed
80
81
82
83
84
85
86
87
        nx.write_gexf(d2p, filename)

    def save_path_in_graph(self, filename):
        d2c = self.d2g.clone()
        for idx, udg in enumerate(self):
            d2c.nodes[str(udg.idx)]["path"] = idx
        nx.write_gexf(d2c, filename)

Yoann Dufresne's avatar
Yoann Dufresne committed
88
89

class Unitig(Path):
Yoann Dufresne's avatar
Yoann Dufresne committed
90
91
92

    def __init__(self, d2g):
        super(Unitig, self).__init__(d2g)
Yoann Dufresne's avatar
Yoann Dufresne committed
93

Yoann Dufresne's avatar
Yoann Dufresne committed
94
    def add_right(self, udg):
Yoann Dufresne's avatar
Yoann Dufresne committed
95
        self.add_path([udg])
Yoann Dufresne's avatar
Yoann Dufresne committed
96
97


Yoann Dufresne's avatar
Yoann Dufresne committed
98
99
100
101
102
103
104
105
106
107
108
109
def d2_path_to_barcode_path(path):
    barcode_per_idx = [set(udg.to_list()) for udg in path]
    diff_barcode_per_idx = []
    rev_diff_barcode_per_idx = []
    for idx in range(len(barcode_per_idx)-1):
        diff_barcode_per_idx.append(barcode_per_idx[idx] - barcode_per_idx[idx+1])
        rev_diff_barcode_per_idx.append(barcode_per_idx[idx+1] - barcode_per_idx[idx])
    diff_barcode_per_idx.append(barcode_per_idx[-1] - diff_barcode_per_idx[-1])
    rev_diff_barcode_per_idx.insert(0, barcode_per_idx[0] - rev_diff_barcode_per_idx[0])

    for diff, rev_diff in zip(diff_barcode_per_idx, rev_diff_barcode_per_idx):
        print(diff, rev_diff)
Yoann Dufresne's avatar
Yoann Dufresne committed
110