arg-dashboard 0.1.19__py3-none-any.whl

arg_dashboard/arg.py

@@ -0,0 +1,1019 @@
+
+import sys
+from math import exp, log
+#import networkx.algorithms.non_randomness
+
+import sys
+import logging
+
+log = logging.getLogger('arg_dashboard.arg')
+# handler = logging.StreamHandler(sys.stderr)
+# handler.setFormatter(logging.Formatter('%(name)s - %(levelname)s - %(message)s'))
+# log.addHandler(handler)
+log.setLevel(logging.INFO)
+import json
+
+import random
+# print("Setting random seed")
+# random.seed(3)
+# numpy.random.seed(3)
+
+from random import random, shuffle, choices
+from functools import partial
+from copy import deepcopy, copy
+# import networkx as nx
+from functools import reduce
+from numpy.random import exponential, sample
+import numpy as np
+from time import sleep
+# import matplotlib.pyplot as plt
+from collections import defaultdict
+
+# from .arg_layout_json import compute_arg_xpos
+# from .arg_layout_force import compute_arg_xpos
+from .arg_layout_mindist import compute_arg_xpos
+
+__all__ = [
+    'Lineage',
+    'Leaf',
+    'get_arg_nodes',
+    'arg2json',
+    'json2arg',
+    ]
+
+class Lineage(object):
+    """
+    Lineages are the edges of our graph
+    """
+    def __init__(self, lineageid=None, down=None, up=None, intervals=None):
+        self.lineageid = lineageid
+        self.down = down # node at bottom of edte
+        self.up = up # node at top of edge
+        self.intervals = intervals
+
+    def __hash__(self):
+        return self.lineageid
+
+    def __eq__(self, other):
+        return hasattr(other, 'lineageid') and self.lineageid == other.lineageid  
+
+    def __repr__(self):
+        return f'{self.lineageid}:{self.intervals}'
+
+    def get_dict(self):
+        d = self.__dict__.copy()
+        if d['up']:
+            d['up'] = d['up'].nodeid
+        d['down'] = d['down'].nodeid
+        return d
+
+    def toJSON(self):
+        return json.dumps(self, default=self.get_dict, sort_keys=True, indent=4)
+
+    def __deepcopy__(self, memo):
+        if self in memo:
+            return memo.get(self)
+        dup = type(self)(
+            lineageid=deepcopy(self.lineageid, memo),
+            # down=deepcopy(self.down, memo),
+            # up=deepcopy(self.up, memo),
+            intervals=deepcopy(self.intervals, memo)
+            )
+        memo[self] = dup
+        dup.down=deepcopy(self.down, memo)
+        dup.up=deepcopy(self.up, memo)
+        return dup
+
+class Node():
+    """
+    Leaf, Coalscent and Recombination are the nodes of the graph
+    """
+    def __hash__(self):
+        return self.nodeid
+
+    def __eq__(self, other):
+        return hasattr(other, 'nodeid') and self.nodeid == other.nodeid  
+
+    def __repr__(self):
+        return f'{self.nodeid}'
+
+class Leaf(Node):
+
+    def __init__(self, nodeid=None, height=None, parent=None, intervals=[(0, 1)], xpos=None):
+        self.nodeid = nodeid
+        self.height = height
+        self.intervals = intervals
+        self.parent = parent
+        self.xpos = xpos
+
+    def get_dict(self):
+        d = self.__dict__.copy()
+        d['parent'] = d['parent'].lineageid
+        return d
+
+    def toJSON(self):
+        return json.dumps(self, default=self.get_dict, sort_keys=True, indent=4)
+
+    def __deepcopy__(self, memo):
+        if self in memo:
+            return memo.get(self)
+        dup = type(self)(
+            nodeid=deepcopy(self.nodeid, memo),
+            height=deepcopy(self.height, memo),
+            # parent=deepcopy(self.parent, memo),
+            intervals=deepcopy(self.intervals, memo),
+            xpos=deepcopy(self.xpos, memo)
+            )
+        memo[self] = dup
+        dup.parent = deepcopy(self.parent, memo)
+        return dup            
+
+class Coalescent(Node):
+
+    def __init__(self, nodeid=None, height=None, children=None, parent=None, xpos=None):
+        self.nodeid = nodeid
+        self.height = height
+        self.children = children
+        self.parent = parent
+        self.xpos = xpos
+
+    def get_dict(self):
+        d = self.__dict__.copy()
+        d['children'] = [c.lineageid for c in d['children']]
+        d['parent'] = d['parent'].lineageid
+        return d
+
+    def toJSON(self):
+        return json.dumps(self, default=self.get_dict, sort_keys=True, indent=4)
+
+    def __deepcopy__(self, memo):
+        if self in memo:
+            return memo.get(self)
+        dup = type(self)(
+            nodeid=deepcopy(self.nodeid, memo),
+            height=deepcopy(self.height, memo),
+            # children=deepcopy(self.children, memo),
+            # parent=deepcopy(self.parent, memo),
+            xpos=deepcopy(self.xpos, memo)
+            )
+        memo[self] = dup
+        dup.children=[deepcopy(c, memo) for c in self.children]
+        dup.parent=deepcopy(self.parent, memo)
+        return dup            
+
+class Recombination(Node):
+
+    def __init__(self, nodeid=None, height=None, child=None, recomb_point=None, left_parent=None, right_parent=None, xpos=None):
+        self.nodeid = nodeid
+        self.height = height
+        self.child = child
+        self.left_parent = left_parent
+        self.right_parent = right_parent
+        self.recomb_point = recomb_point
+        self.xpos = xpos
+
+    def get_dict(self):
+        d = self.__dict__.copy()
+        d['left_parent'] = d['left_parent'].lineageid
+        d['right_parent'] = d['right_parent'].lineageid
+        d['child'] = d['child'].lineageid
+        return d
+
+    def toJSON(self):
+        return json.dumps(self, default=self.get_dict, sort_keys=True, indent=4)
+
+    def __deepcopy__(self, memo):
+        if self in memo:
+            return memo.get(self)
+        dup = type(self)(
+            nodeid=deepcopy(self.nodeid, memo),
+            height=deepcopy(self.height, memo),
+            # child=deepcopy(self.child, memo),
+            recomb_point=deepcopy(self.recomb_point, memo),
+            # left_parent=deepcopy(self.left_parent, memo),
+            # right_parent=deepcopy(self.right_parent, memo),
+            xpos=deepcopy(self.xpos, memo)
+            )  
+        memo[self] = dup
+        dup.child=deepcopy(self.child, memo)
+        dup.left_parent=deepcopy(self.left_parent, memo)
+        dup.right_parent=deepcopy(self.right_parent, memo)
+        return dup                  
+
+def flatten(list_of_tps):
+ 
+    return reduce(lambda ls, ival: ls + list(ival), list_of_tps, [])
+
+
+def unflatten(list_of_endpoints):
+
+    return [ [list_of_endpoints[i], list_of_endpoints[i + 1]]
+          for i in range(0, len(list_of_endpoints) - 1, 2)]
+
+def merge(query, annot, op):
+
+    a_endpoints = flatten(query)
+    b_endpoints = flatten(annot)
+
+    assert a_endpoints == sorted(a_endpoints), "not sorted or non-overlaping"
+    assert b_endpoints == sorted(b_endpoints), "not sorted or non-overlaping"
+
+
+    sentinel = max(a_endpoints[-1], b_endpoints[-1]) + 1
+    a_endpoints += [sentinel]
+    b_endpoints += [sentinel]
+
+    a_index = 0
+    b_index = 0
+
+    res = []
+
+    scan = min(a_endpoints[0], b_endpoints[0])
+    while scan < sentinel:
+        in_a = not ((scan < a_endpoints[a_index]) ^ (a_index % 2))
+        in_b = not ((scan < b_endpoints[b_index]) ^ (b_index % 2))
+        in_res = op(in_a, in_b)
+
+        if in_res ^ (len(res) % 2):
+            res += [scan]
+        if scan == a_endpoints[a_index]: 
+            a_index += 1
+        if scan == b_endpoints[b_index]: 
+            b_index += 1
+        scan = min(a_endpoints[a_index], b_endpoints[b_index])
+
+    return unflatten(res)
+
+def interval_diff(a, b):
+    if not (a and b):
+        return a and a or b
+    return merge(a, b, lambda in_a, in_b: in_a and not in_b)
+
+def interval_union(a, b):
+    if not (a and b):
+        return []
+    return merge(a, b, lambda in_a, in_b: in_a or in_b)
+
+def interval_intersect(a, b):
+    if not (a and b):
+        return []
+    return merge(a, b, lambda in_a, in_b: in_a and in_b)
+    
+def interval_sum(intervals):
+    return sum(e - s for (s, e) in intervals)
+
+def interval_span(intervals):
+    starts, ends = zip(*intervals)
+    return max(ends) - min(starts)
+
+def interval_split(intervals, pos):
+    left, right = list(), list()
+    for s, e in intervals:
+        if pos >= e:
+            left.append((s, e))
+        elif pos >= s and pos < e:
+            left.append((s, pos))
+            right.append((pos, e))
+        else:
+            right.append((s, e))
+    return left, right
+
+def interval_any_shared_borders(a, b):
+    log.info('shared borders')
+    flat = flatten(a) + flatten(b)
+    return len(flat) > len(set(flat))
+
+def _get_neighbors_below(node):
+    """Get nodes connected below this node via child lineages."""
+    if type(node) is Coalescent:
+        return [c.down for c in node.children]
+    if type(node) is Recombination:
+        return [node.child.down]
+    return []
+
+def _get_neighbors_above(node):
+    """Get nodes connected above this node via parent lineages."""
+    result = []
+    if type(node) is Coalescent:
+        if node.parent and node.parent.up:
+            result.append(node.parent.up)
+    elif type(node) is Recombination:
+        if node.left_parent and node.left_parent.up:
+            result.append(node.left_parent.up)
+        if node.right_parent and node.right_parent.up:
+            result.append(node.right_parent.up)
+    elif type(node) is Leaf:
+        if node.parent and node.parent.up:
+            result.append(node.parent.up)
+    return result
+
+def _dfs_leaf_order(node, visited=None):
+    """DFS from root to get a leaf ordering that groups related leaves.
+    Randomizes child visit order at coalescent nodes for multi-start."""
+    if visited is None:
+        visited = set()
+    if node in visited:
+        return []
+    visited.add(node)
+    if type(node) is Leaf:
+        return [node]
+    leaves = []
+    if type(node) is Coalescent:
+        children = list(node.children)
+        shuffle(children)
+        for child in children:
+            leaves.extend(_dfs_leaf_order(child.down, visited))
+    elif type(node) is Recombination:
+        leaves.extend(_dfs_leaf_order(node.child.down, visited))
+    return leaves
+
+####################################################
+
+def segments_crossing(line1, line2):
+    xdiff = (line1[0][0] - line1[1][0], line2[0][0] - line2[1][0])
+    ydiff = (line1[0][1] - line1[1][1], line2[0][1] - line2[1][1])
+
+    def det(a, b):
+        return a[0] * b[1] - a[1] * b[0]
+
+    div = det(xdiff, ydiff)
+    if div == 0:
+        return False
+    #    raise Exception('lines do not intersect')
+
+    d = (det(*line1), det(*line2))
+    x = det(d, xdiff) / div
+    y = det(d, ydiff) / div
+
+    # return x, y
+
+    return \
+        min(line1[0][0], line1[1][0]) < x < max(line1[0][0], line1[1][0]) and \
+        min(line2[0][0], line2[1][0]) < x < max(line2[0][0], line2[1][0]) and \
+        min(line1[0][1], line1[1][1]) < y < max(line1[0][1], line1[1][1]) and \
+        min(line2[0][1], line2[1][1]) < y < max(line2[0][1], line2[1][1])
+
+
+def crossing(a, b):
+    return segments_crossing(
+        ((a.down.xpos, a.down.height), (a.up.xpos, a.up.height)), 
+        ((b.down.xpos, b.down.height), (b.up.xpos, b.up.height))
+        )
+
+def get_all_crossing_pairs(lineages):
+    crossing_pairs = []
+    for i in range(len(lineages)):
+        for j in range(i+1, len(lineages)):
+            if crossing(lineages[i], lineages[j]):
+                crossing_pairs.append([i, j])
+    return crossing_pairs
+
+def reduce_crossovers(nodes):
+    lineages = get_parent_lineages(nodes)
+    lineages = [x for x in lineages if x.up is not None]
+    crossing_pairs = get_all_crossing_pairs(lineages)
+
+    for i, j in crossing_pairs:
+        # swap downs
+        lineages[i].down.xpos, lineages[j].down.xpos = lineages[j].down.xpos, lineages[i].down.xpos
+        new_crossing_pairs = get_all_crossing_pairs(lineages)
+        if not len(new_crossing_pairs) < len(crossing_pairs):
+            # swap downs back
+            lineages[i].down.xpos, lineages[j].down.xpos = lineages[j].down.xpos, lineages[i].down.xpos
+        else:
+            crossing_pairs = new_crossing_pairs
+
+        # swap ups
+        lineages[i].up.xpos, lineages[j].up.xpos = lineages[j].up.xpos, lineages[i].up.xpos
+        new_crossing_pairs = get_all_crossing_pairs(lineages)
+        if not len(new_crossing_pairs) < len(crossing_pairs):
+            # swap ups back
+            lineages[i].up.xpos, lineages[j].up.xpos = lineages[j].up.xpos, lineages[i].up.xpos
+        else:
+            crossing_pairs = new_crossing_pairs
+
+
+        leaf_idx = [i for i, n in enumerate(nodes) if type(n) is Leaf]
+
+        if type(lineages[i].down) is Leaf and type(lineages[j].down) is not Leaf:
+            # try to find new position for leaf i
+
+            for k in leaf_idx:
+                # swap
+                lineages[i].down.xpos, lineages[k].down.xpos = lineages[k].down.xpos, lineages[i].down.xpos
+
+                new_crossing_pairs = get_all_crossing_pairs(lineages)
+                if not len(new_crossing_pairs) < len(crossing_pairs):
+                    # swap back
+                    lineages[i].down.xpos, lineages[k].down.xpos = lineages[k].down.xpos, lineages[i].down.xpos
+                else:
+                    break
+
+
+        if type(lineages[i].down) is Leaf and type(lineages[j].down) is not Leaf:
+            # try to find new position for leaf j
+            for k in leaf_idx:
+                # swap
+                lineages[j].down.xpos, lineages[k].down.xpos = lineages[k].down.xpos, lineages[j].down.xpos
+
+                new_crossing_pairs = get_all_crossing_pairs(lineages)
+                if not len(new_crossing_pairs) < len(crossing_pairs):
+                    # swap back
+                    lineages[j].down.xpos, lineages[k].down.xpos = lineages[k].down.xpos, lineages[j].down.xpos
+                else:
+                    break
+
+
+    # TODO: if cross over involves an internal and an external branch, see if
+    # you can move the leaf to a new position that removes the crossover
+
+
+
+    # leaf_idx = [i for i, n in enumerate(nodes) if type(n) is Leaf]
+    # suffled_leaf_idx = leaf_idx[:]
+    # shuffle(suffled_leaf_idx)
+    # for i, j in sorted(zip(leaf_idx, suffled_leaf_idx)):
+    #     print(i, j)
+    #     nodes[i].xpos = nodes[j].xpos
+
+    # new_crossing_pairs = get_all_crossing_pairs(lineages)
+    # if not len(new_crossing_pairs) < len(crossing_pairs):
+    #     for i, j in sorted(zip(leaf_idx, suffled_leaf_idx), reverse=True):
+    #         nodes[i].xpos = nodes[j].xpos
+
+
+
+
+def branch_length(lineage):
+    return ((lineage.down.xpos - lineage.up.xpos)**2 + (lineage.down.height - lineage.up.height)**2)**0.5
+
+def reduce_total_branch_length(nodes, shift):
+    lineages = get_parent_lineages(nodes)
+    lineages = [x for x in lineages if x.up is not None]
+
+    cur_tot = sum(branch_length(l) for l in lineages)
+    for i in range(len(lineages)):
+        orig_xpos = lineages[i].up.xpos
+        lineages[i].up.xpos += shift
+        tot = sum(branch_length(l) for l in lineages)
+        if tot < cur_tot:
+            cur_tot = tot
+        else:
+            lineages[i].up.xpos = orig_xpos
+        lineages[i].up.xpos -= shift
+        tot = sum(branch_length(l) for l in lineages)
+        if tot < cur_tot:
+            cur_tot = tot
+        else:
+            lineages[i].up.xpos = orig_xpos
+
+def redistribute_leaves(nodes):
+
+    leaves = [n for n in nodes if type(n) is Leaf]
+    # min_x = min(n.xpos for n in leaves)
+    # max_x = max(n.xpos for n in leaves)
+    leaves = sorted(leaves, key=lambda x: x.xpos)
+    min_x = 0
+    max_x = 1
+    for xpos, leaf in zip(np.linspace(min_x, max_x, len(leaves)), leaves):
+        leaf.xpos = xpos
+
+####################################################
+
+def _bottom_up_barycenter(nodes):
+    """Single bottom-up pass: each node at mean of its children's x."""
+    non_leaves = sorted(
+        [n for n in nodes if type(n) is not Leaf],
+        key=lambda n: n.height
+    )
+    for node in non_leaves:
+        below = _get_neighbors_below(node)
+        placed = [n for n in below if n.xpos is not None]
+        if placed:
+            node.xpos = sum(n.xpos for n in placed) / len(placed)
+        else:
+            node.xpos = 0.5
+
+def _count_crossings(nodes):
+    """Count total edge crossings in current layout."""
+    lineages = get_parent_lineages(nodes)
+    lineages = [x for x in lineages if x.up is not None]
+    return len(get_all_crossing_pairs(lineages))
+
+def add_node_x_positions(nodes, n_restarts=10):
+    """
+    Position nodes to minimize crossings using:
+    1. DFS leaf ordering (groups related leaves)
+    2. Bottom-up barycenter (single pass, no funnel)
+    3. Rescale + crossing reduction post-processing
+
+    Multi-start: tries n_restarts random DFS orderings, keeps best.
+    """
+    best_crossings = float('inf')
+    best_positions = None
+
+    for _ in range(n_restarts):
+        # Step 1: DFS leaf ordering (randomized child order)
+        leaf_order = _dfs_leaf_order(nodes[-1])
+        seen = set(id(l) for l in leaf_order)
+        for n in nodes:
+            if type(n) is Leaf and id(n) not in seen:
+                leaf_order.append(n)
+
+        # Step 2: Fix leaves evenly spaced
+        n_leaves = len(leaf_order)
+        for i, leaf in enumerate(leaf_order):
+            leaf.xpos = i / max(n_leaves - 1, 1)
+
+        # Step 3: Bottom-up barycenter
+        _bottom_up_barycenter(nodes)
+
+        # Step 4: Rescale to [0, 1]
+        rescale_positions(nodes)
+
+        # Step 5: Re-even leaves, reduce crossings
+        redistribute_leaves(nodes)
+        for _ in range(5):
+            reduce_crossovers(nodes)
+
+        # Evaluate
+        crossings = _count_crossings(nodes)
+        if crossings < best_crossings:
+            best_crossings = crossings
+            best_positions = {n.nodeid: n.xpos for n in nodes}
+
+        if best_crossings == 0:
+            break
+
+    # Restore best positions
+    if best_positions:
+        for n in nodes:
+            n.xpos = best_positions[n.nodeid]
+
+    # Final refinement
+    shift_list = [0.5, 0.4, 0.3, 0.2, 0.1, 0.05, 0.02]
+    for shift in shift_list:
+        reduce_total_branch_length(nodes, shift)
+    rescale_positions(nodes)
+    for _ in range(5):
+        reduce_crossovers(nodes)
+    redistribute_leaves(nodes)
+
+def get_arg_nodes(n=5, N=10000, r=1e-8, L=5e3, simulation="arg"):
+    """
+    Simulates an ARG
+    """
+    # because we use the sequence interval from 0 to 1
+    r = r * L
+
+#    print('\nSIM:')
+
+
+    assert simulation in ["arg", "smcprime", "smc"]
+
+    nodes = list()
+    live = list() # live lineages
+
+    for i in range(n):
+        leaf = Leaf(i, height=0)
+        nodes.append(leaf)
+        lin = Lineage(lineageid=i, down=leaf, intervals=[(0, 1)])
+        leaf.parent = lin
+        live.append(lin)
+        last_node = i # max node number used so far
+        last_lineage = i
+        
+    while len(live) > 1:
+        shuffle(live)
+
+        coal_prob = (len(live) * len(live)-1) / 2 / (2*N) 
+        tot_ancestral = sum(interval_sum(x.intervals) for x in live)
+        rec_prob = r * tot_ancestral
+        wating_time = exponential(1 / (coal_prob + rec_prob))
+        height = wating_time + nodes[-1].height
+        if random() < coal_prob / (coal_prob + rec_prob):
+            # coalescence
+
+            # if full arg:            
+            a, b = 0, 1
+
+            if simulation == "smcprime":
+                # intervals must overlap or be adjacent:
+                while not (interval_any_shared_borders(live[a].intervals, live[b].intervals) or \
+                    interval_intersect(live[a].intervals, live[b].intervals)):
+                    shuffle(live)
+            elif simulation == "smc":
+                # intervals must overlap:
+                while not interval_intersect(live[a].intervals, live[b].intervals):
+                    shuffle(live)
+    
+            lin_a, lin_b = live.pop(0), live.pop(0)
+
+            intervals = interval_union(lin_a.intervals, lin_b.intervals)
+
+            # new node
+            if lin_a.down is lin_b.down:
+                # make sure diamond is symmetric
+                children=[lin_a.down.left_parent, lin_b.down.right_parent]
+            else:
+                children=[lin_a, lin_b]  
+            node_c = Coalescent(nodeid=last_node+1, height=height, 
+                        children=children)
+            
+            last_node += 1
+            nodes.append(node_c)
+
+            # add node to top of coalescing lineages
+            lin_a.up = node_c
+            lin_b.up = node_c
+
+            # new lineage
+            lin_c = Lineage(lineageid=last_lineage+1, down=node_c, intervals=intervals) # fixme intervals
+            last_lineage += 1
+            live.append(lin_c)
+            node_c.parent = lin_c
+        else:
+            # recombination
+
+            # rec_lin = choices(live, weights=[interval_sum(x.intervals) for x in live], k=1)[0]
+            rec_lin = choices(live, weights=[interval_span(x.intervals) for x in live], k=1)[0]
+            live.remove(rec_lin)
+
+            # total ancestral material
+            total_anc = interval_sum(rec_lin.intervals)
+
+            # recombination point in ancestral material
+            recomb_point_anc = random() * total_anc
+
+            # recombination point in full sequence
+            cum = 0
+            for s, e in rec_lin.intervals:
+                cum += e - s
+                if cum > recomb_point_anc:
+                    recomb_point = e - (cum - recomb_point_anc)
+                    break
+
+            # recombination node
+            rec_node = Recombination(nodeid=last_node+1, height=height, 
+                child=rec_lin, recomb_point=recomb_point)
+            last_node += 1
+            nodes.append(rec_node)
+
+            # two new lineages both refering back to recombination node
+            intervals_a, intervals_b = interval_split(rec_lin.intervals, recomb_point)
+            assert interval_sum(intervals_a) and interval_sum(intervals_b)
+            assert sum(e-s for s, e in intervals_a) + sum(e-s for s, e in intervals_b) == sum(e-s for s, e in rec_lin.intervals)
+            lin_a = Lineage(lineageid=last_lineage+1, down=rec_node, intervals=intervals_a)
+            last_lineage += 1
+            lin_b = Lineage(lineageid=last_lineage+1, down=rec_node, intervals=intervals_b)
+            last_lineage += 1
+            live.append(lin_a)
+            live.append(lin_b)
+
+            # add parents of node
+            rec_node.left_parent = lin_a
+            rec_node.right_parent = lin_b
+
+            # add parent for child
+            rec_lin.up = rec_node
+
+    add_node_x_positions(nodes)
+
+    return nodes
+
+def _rescale_layout(pos, scale=1):
+    # rescale to (0,pscale) in all axes
+
+    # shift origin to (0,0)
+    lim = 0  # max coordinate for all axes
+    for i in range(pos.shape[1]):
+        pos[:, i] -= pos[:, i].min()
+        lim = max(pos[:, i].max(), lim)
+    # rescale to (0,scale) in all directions, preserves aspect
+    for i in range(pos.shape[1]):
+        pos[:, i] *= scale / lim
+    return pos
+
+def get_positions(nodes):
+    """
+    Gets list of x,y positions for nodes
+    """
+    positions = list()
+    for node in nodes:
+        positions.append((node.xpos, node.height))
+    return positions
+
+def get_breakpoints(nodes):
+    """
+    Gets list of recombination break points
+    """    
+    # get list of break points
+    breakpoints = list()
+    for n in nodes:
+        if type(n) is Recombination:
+#        if type(n) is Recombination and all(p.up in nodes for p in n.parents):
+            breakpoints.append(n.recomb_point)
+    return sorted(breakpoints)
+
+def get_parent_lineages(nodes, root=True):
+    """
+    Get list of parent lineages of nodes ordered by id
+    """
+
+    lineages = list()
+    for node in nodes:
+        if type(node) is Coalescent:
+            lineages.append(node.parent)
+        if type(node) is Recombination:
+            lineages.append(node.left_parent)
+            lineages.append(node.right_parent)
+        if type(node) is Leaf:
+            lineages.append(node.parent)
+
+    # get unique and sort them so lineageid matches index
+    lineages = list(set(lineages))
+    lineages.sort(key=lambda x: x.lineageid)
+
+    if not root:
+        # remove dangling root lineage
+        lineages = lineages[:-1]
+
+    return lineages
+
+def get_child_lineages(nodes):
+    """
+    Get list of child lineages of nodes ordered by id
+    """
+
+    lineages = list()
+    for node in nodes:
+        if type(node) is Coalescent:
+            lineages.extend(node.children)
+        if type(node) is Recombination:
+            lineages.append(node.child)
+
+    # get unique and sort them so lineageid matches index
+    lineages = list(set(lineages))
+    lineages.sort(key=lambda x: x.lineageid)
+
+    return lineages
+
+# def traverse_marginal(node, interval):
+#     """
+#     Recursive function for getting marginal tree/ARG
+#     """    
+#     node = deepcopy(node) # TODO: remove if input is a cloned arg.
+#     tree_nodes = set()
+#     if type(node) is Leaf:
+#         tree_nodes.add(node)
+#     if type(node) is Recombination:
+#         if interval_intersect([interval], node.child.intervals):
+#             tree_nodes.add(node)
+#             tree_nodes.update(traverse_marginal(node.child.down, interval))
+#     elif type(node) is Coalescent:
+#         if node.parent is None or interval_intersect([interval], node.parent.intervals):
+#             tree_nodes.add(node)
+#         del_child = None
+#         for i, child in enumerate(node.children):
+#             if interval_intersect([interval], child.intervals):
+#                 tree_nodes.update(traverse_marginal(child.down, interval))
+#             else:
+#                 del_child = i
+#         if del_child is not None:
+#             del node.children[del_child]
+#     return tree_nodes
+
+def _traverse_marginal(node, interval):
+    """
+    Recursive function for getting marginal tree/ARG
+    """    
+    tree_nodes = set()
+    if type(node) is Leaf:
+        tree_nodes.add(node)
+    if type(node) is Recombination:
+        if interval_intersect([interval], node.child.intervals):
+            tree_nodes.add(node)
+            tree_nodes.update(_traverse_marginal(node.child.down, interval))
+    elif type(node) is Coalescent:
+        if node.parent is None or interval_intersect([interval], node.parent.intervals):
+            tree_nodes.add(node)
+        del_child = None
+        for i, child in enumerate(node.children):
+            if interval_intersect([interval], child.intervals):
+                tree_nodes.update(_traverse_marginal(child.down, interval))
+            else:
+                del_child = i
+        if del_child is not None:
+            del node.children[del_child]
+    return tree_nodes       
+
+def traverse_marginal(node, interval):
+    clone = deepcopy(node) # TODO: remove if input is a cloned arg.
+    return _traverse_marginal(clone, interval)
+
+
+def remove_dangling_root(tree_nodes):
+    """
+    Remove the nodes of a marginal tree or ARG 
+    from root to first coalescence
+    """    
+    for_del = list()
+    for i in range(len(tree_nodes)-1, -1, -1):
+        if type(tree_nodes[i]) is Coalescent and len(tree_nodes[i].children) == 2:
+            break
+        for_del.append(i)
+    for i in for_del:
+        del tree_nodes[i]
+
+def marginal_arg(nodes, interval, strip_dangling_root=True):
+    """
+    Gets the marginal ARG given a sequene interval
+    """
+
+    # TODO: 
+    # nodes = clone_arg(nodes)
+
+    # get nodes for marginal
+    marg_nodes = traverse_marginal(nodes[-1], list(interval))
+
+    # set to list
+    marg_nodes = list(marg_nodes)
+
+    # sort on height
+    marg_nodes.sort(key=lambda x: x.height)
+    # prune top path above last coalescence
+    if strip_dangling_root:
+        remove_dangling_root(marg_nodes)
+
+    return marg_nodes
+
+def marginal_trees(nodes, interval, strip_dangling_root=False):
+    """
+    Gets list of marginal trees
+    """
+    tree_list = list()
+    breakpoints = get_breakpoints(nodes)
+
+    borders = [0] + breakpoints + [1]
+    
+    # for interval in zip(borders[:-1], borders[1:]):
+    #     marg_nodes = marginal_arg(nodes, interval)
+    #     tree_list.append(marg_nodes)
+    # return tree_list
+
+    interval_list = list()
+    for interv in zip(borders[:-1], borders[1:]):
+        marg_nodes = marginal_arg(nodes, interv, strip_dangling_root=strip_dangling_root)
+        # TODO: make this correct:...
+        if marg_nodes and interv[1] > interval[0] and interval[1] > interv[0]:
+            tree_list.append(marg_nodes)
+            interval_list.append(interv)
+    return tree_list, interval_list
+
+# def draw_graph(nodes):
+#     """
+#     Draws graph using matplotlib
+#     """    
+#     # make a graph and positions list
+#     positions = list()
+#     arg = nx.Graph()
+#     for node in nodes:
+#         arg.add_node(node.nodeid)
+#         positions.append((node.xpos, node.height))
+#         if isinstance(node, Recombination):
+#             arg.add_edge(node.child.down.nodeid, node.nodeid)
+#         elif isinstance(node, Coalescent):
+#             for child in node.children:
+#                 arg.add_edge(child.down.nodeid, node.nodeid)
+
+#     positions = np.array(positions)
+#     positions = dict(zip(arg.nodes(), positions))
+
+#     #pos = nx.spring_layout(arg)
+#     nx.draw(arg, positions, alpha=0.5, node_size=200, with_labels=True)
+#     #with_labels=True, 
+#     #connectionstyle='arc3, rad = 0.1', 
+#     #arrowstyle='-')
+
+#     # G = nx.DiGraph() #or G = nx.MultiDiGraph()
+#     # G.add_node('A')
+#     # G.add_node('B')
+#     # G.add_edge('A', 'B', length = 2)
+#     # G.add_edge('B', 'A', length = 3)
+#     # pos = nx.spring_layout(G)
+#     # nx.draw(G, pos, with_labels=True, connectionstyle='arc3, rad = 0.1')
+#     # edge_labels=dict([((u,v,),d['length'])
+#     #              for u,v,d in G.edges(data=True)])
+
+#     plt.show()
+
+def rescale_positions(nodes):
+    """
+    Rescales xpos and heights to 0,1 range
+    """
+    max_height = max(n.height for n in nodes)
+    max_xpos = max(n.xpos for n in nodes)
+    min_xpos = min(n.xpos for n in nodes)
+    for node in nodes:
+        node.xpos = (node.xpos - min_xpos) / (max_xpos-min_xpos)
+        node.height = node.height / max_height
+
+def arg2json(nodes):
+
+    lineages = get_parent_lineages(nodes)
+
+    json_data = dict(Coalescent=[], Recombination=[], Leaf=[], 
+                        Lineage=[l.get_dict() for l in lineages])
+    for node in nodes:
+        json_data[node.__class__.__name__].append(node.get_dict())
+
+    json_data = compute_arg_xpos(json_data)
+
+    return json.dumps(json_data, indent=4)
+
+def json2arg(json_str):
+
+    nodes = list()
+
+    data = json.loads(json_str)
+
+    # make lineages (with indexes instead of node refs)
+    lineages = [Lineage(**data) for data in data['Lineage']]
+
+    for node_data in data['Coalescent']:            
+        # make the node
+        node = Coalescent(**node_data)
+        # populate the parent and children with actual lineages
+        node.parent = lineages[node.parent]
+        node.children = [lineages[i] for i in node.children]
+        nodes.append(node)
+
+    for node_data in data['Recombination']:
+        # make the node
+        node = Recombination(**node_data)
+        # populate the parents and child with actual lineages
+        node.left_parent = lineages[node.left_parent]
+        node.right_parent = lineages[node.right_parent]
+        node.child = lineages[node.child]
+        nodes.append(node)
+
+    for node_data in data['Leaf']:
+        # make the node
+        node = Leaf(**node_data)
+        # populate the parent with actual lineages
+        node.parent = lineages[node.parent]
+        nodes.append(node)
+
+    nodes.sort(key=lambda x: x.nodeid)
+
+    # populate up and down with the nodes
+    for lineage in lineages:
+        lineage.down = nodes[lineage.down]
+        if lineage.up is not None:
+            lineage.up = nodes[lineage.up]
+
+    return nodes
+
+if __name__ == '__main__':
+
+    # get arg and add positions
+    nodes = get_arg_nodes()
+
+    print(nodes)
+
+    print(deepcopy(nodes))
+
+
+    json_str = arg2json(nodes)
+    print(json_str)
+    retrieved_nodes = json2arg(json_str)
+
+    print(nodes)
+    print(retrieved_nodes)
+    print(nodes == retrieved_nodes)
+
+
+
+    # get breakpoints
+    breakpoints = get_breakpoints(nodes)
+    print(breakpoints)
+
+    sys.exit()
+
+    # get marginal trees
+    trees = marginal_trees(nodes)
+    print(trees)
+
+    # marginal arg for some consequtive intervals
+    marg_arg = marginal_arg(nodes, [0, breakpoints[1]])
+
+
+    # # draw graphs for testing
+    # draw_graph(nodes)
+    # draw_graph(marg_arg)
+
+    # draw_graph(nodes)
+    # for tree in marginal_trees(nodes):
+    #     print([n.xpos for n in tree])
+    #     draw_graph(tree)
+