PARU Class for Scientific Visualization of LMS.

flowmap.py

### Imports ###

# Python

import os
import logging
import math
import sys
import itertools
import random
import time
import bisect

# Libraries

import h5py
import lxml.objectify
import numpy as np
import cairo

# Local

from geom import *

# Shortcuts

pi = math.pi
sqrt = math.sqrt
inf = float('inf')


### Classes ###

class PARU(object):
    
    def __init__(self, xmf_path):
        self.xmf_path = xmf_path
        self.h5files = {}
        self.parse()
        
    def parse(self):
        """ Parse the XMF and h5 files.
        """
        xmf = lxml.objectify.parse(self.xmf_path)
        root = xmf.getroot()
        self.data = {}
        for grid in root.Domain.Grid.Grid:
            # Get geometry and time data
            geometry = self.read_h5_data(grid.Geometry.DataItem)
            time = int(grid.Time.get('Value'))
            # Get attributes
            attributes = dict((a.get('Name'), self.read_h5_data(a.DataItem)) for a in grid.Attribute)
            # Store data
            self.data[time] = {'geometry': geometry,
                               'attributes': attributes}
        self.timesteps = sorted(self.data.keys())
        # Compute bounds
        geometries = [d['geometry'] for (time, d) in self.data.items()]
        mins = np.array([geom.value.min(0) for geom in geometries])
        maxs = np.array([geom.value.max(0) for geom in geometries])
        xmin, ymin, _ = mins.min(0)
        xmax, ymax, _ = maxs.max(0)
        self.bounds = (xmin, ymin, xmax, ymax)
        
        paths = []
        for t in self.timesteps:
            paths.append(self.scale_range(list(self.data[t]['geometry'].value[:,0:2])))
        self.paths = zip(*paths)
    
    def scale_range(self, points):
        (xmin, ymin, xmax, ymax) = self.bounds
        return [((x - xmin) / float(xmax - xmin), (y - ymin) / float(ymax - ymin)) for (x,y) in points]
    
    def read_h5_data(self, info):
        """ Get data for a h5 file belonging to an
            xmf node.
        """
        fname, h5path = str(info).split(':')
        if fname in self.h5files:
            f = self.h5files[fname]
        else:
            f = h5py.File(os.path.join(os.path.dirname(self.xmf_path),fname))
            self.h5files[fname] = f
        return f[h5path]
        
    def extract_paths(self):
        return 
    
    def get_paths(self, steps=1):
        for i in range(steps):
            yield [path[i:len(self.timesteps)-steps+i+1] for path in self.paths]

# Graph model

class Graph(object):
    def __init__(self):
        self.nodes = []
        self.edges = []
        
    def load_paths(self, paths, contents=None):
        """ Convert a list of paths to a graph structure
        """
        if contents is None:
            contents = range(len(paths))
        for path, c in zip(paths, contents):
            n1 = Node(path[0][0],path[0][1],contents=[c])
            self.add_node(n1)
            for p in path[1:]:
                n2 = n1
                n1 = Node(p[0],p[1],contents=[c])
                edge = Edge(n2,n1,weight=1,contents=[c])
                self.add_node(n1)
                self.add_edge(edge)
        return self
    
    def add_node(self, node):
        node.graph = self
        self.nodes.append(node)
        
    def add_edge(self, edge):
        edge.graph = self
        self.edges.append(edge)
        
    def clean(self):
        self.nodes = filter(lambda n: not n.dead, self.nodes)
        self.edges = filter(lambda n: not n.dead, self.edges)
        for node in self.nodes:
            node.changed = False
        for edge in self.edges:
            edge.changed = False
        
        
    def simplify(self, tolerance=0.1*pi):
        """ Removes nodes that lie on an edge between two
            other nodes, thus simplifying the graph.
        """
        print_procedure_start("simplify",self.stats())
        done = False
        while not done:
            to_remove = []
            for i,node in enumerate(self.nodes):
                if len(node.edges_in) == 1 and len(node.edges_out) == 1:
                    edge_in, edge_out = node.edges_in[0], node.edges_out[0]
                    # Skip if edges don't have the same weight
                    if edge_in.weight != edge_out.weight:
                        continue
                    prev = edge_in.from_node.p()
                    next = edge_out.to_node.p()
                    this = node.p()
                    # Deal with colinear points
                    if point_dist(prev,this) == 0:
                        to_remove.append((0,node))
                        continue
                    if point_dist(this,next) == 0:
                        to_remove.append((0,node))
                        continue
                    if point_dist(prev,next) == 0:
                        continue
                    # Check the angle between the line from the previous to the
                    # next node, and the current node.
                    a = abs(angle_fix(line_angle(prev,this,next)))
                    # If the distance is small, mark point
                    if a < tolerance:
                        to_remove.append((a,node))
            to_remove.sort()
            done = True
            for a, node in to_remove:
                if node.changed:
                    done = False
                else:
                    node.remove()
            self.clean()
        print_procedure_done(self.stats())
        return len(to_remove) > 0
        
    def merge_close_nodes(self, max_dist=0.015):
        """ Merges nodes that are close together. 
        """
        print_procedure_start("merge_close_nodes",self.stats())
        # Main loop
        done = False
        iteration = 0
        while(not done):
            iteration += 1
            node_count = len(self.nodes)
            print "    Iteration %d. (%d nodes left)... "%(iteration,node_count)
            # Reset
            nodes_to_merge = []
            node_count = len(self.nodes)
            # Sort nodes by x
            self.nodes.sort(key=lambda n: n.x)
            for i,node1 in enumerate(self.nodes):
                # Sweep and prune
                for node2 in self.nodes[i+1:]:
                    if node2.x > node1.x + max_dist:
                        break
                    dist = ((node1.x - node2.x)**2 + (node1.y - node2.y)**2)**0.5
                    if dist <= max_dist:
                        nodes_to_merge.append((dist, node1, node2))
            # Merge closest pairs first
            print_procedure_status("merging %d pairs of nodes."%len(nodes_to_merge))
            nodes_to_merge.sort()
            done = True
            for i,(dist, node1, node2) in enumerate(nodes_to_merge):
                # If either node was merged before, skip this step.
                if node1.changed or node2.changed:
                    done = False
                else:
                    node1.merge(node2)
                    node1.changed = True
                    node2.changed = True
            self.clean()
        print_procedure_done(self.stats())
        return len(nodes_to_merge) > 0
        
    def unify_identical_edges(self):
        """ Unifies edges that come from and go to the same nodes. 
            Returns whether something changed.
        """
        changed = False
        for j,node in enumerate(self.nodes):
            targets = {}
            for i,edge in enumerate(node.edges_out):
                if edge.to_node not in targets:
                    targets[edge.to_node] = []
                targets[edge.to_node].append(edge)
            # Do the actual merging
            for to_node, edges in targets.items():
                for edge in edges[1:]:
                    changed = True
                    edges[0].unify(edge)
        self.clean()
        return changed
    
    def absorb_nodes_into_edges(self, tolerance=0.1*pi):
        print_procedure_start("absorb_nodes_into_edges")
        # Sort nodes by x
        nodes_by_x = [(n.x, n) for n in self.nodes]
        nodes_by_x.sort()
        done = False
        iteration = 0
        while not done:
            iteration += 1
            print_procedure_status('iteration %d. %d edges'%(iteration, len(self.edges)))
            to_merge = []
            for edge in self.edges:
                l = edge.length()
                f, t = edge.from_node, edge.to_node
                bounds = (l/2)#/math.tan(pi-tolerance)
                # Find bounds
                xmin = min(edge.from_node.x, edge.to_node.x) - bounds
                xmax = max(edge.from_node.x, edge.to_node.x) + bounds
                # Find the first node that is close enough (in x-dimension)
                i = bisect.bisect_left(nodes_by_x,(xmin,None))
                # Sweep and prune
                for node in self.nodes[i:]:
                    if node.x > xmax:
                        break
                    if node == f or node == t:
                        continue
                    a = abs(angle_fix(line_angle(f.p(), node.p(), t.p())))
                    # dist, u = line_point_dist(f.p(), t.p(), node.p())
                    # if 0 < u < 1 and dist < tol:
                    if a < tolerance:
                        to_merge.append((a, edge, node))
            to_merge.sort()
            done = True
            for (d, edge, node) in to_merge:
                if edge.changed or node.changed:
                    done = False
                else:
                    edge.absorb(node)
            self.clean()
        print_procedure_done(self.stats())
        return len(to_merge) > 0
        
    def reduce_complete(self, tolerance = 0.4*pi, max_node_dist=0.05, draw_steps=False):
        if draw_steps:
            self.draw("0-before.png")
        self.simplify(tolerance=tolerance)
        self.merge_close_nodes(max_dist=max_node_dist)
        self.unify_identical_edges()
        changed = True
        iteration = 1
        while (changed):
            changed = False
            changed |= self.absorb_nodes_into_edges(tolerance=tolerance)
            changed |= self.unify_identical_edges()
            changed |= self.simplify(tolerance=tolerance)
            if draw_steps:
                self.draw("%d-after.png"%iteration)
            iteration += 1

    def draw(self, filename="graph.png", w=1000, h=1000, background=(1,1,1,1)):
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, w,h)
        ctx = cairo.Context(surface)
        ctx.set_source_rgba(*background)
        ctx.rectangle(0,0,w,h)
        ctx.fill()
        for edge in self.edges:
            ctx.set_line_width(edge.weight/2.0)
            ctx.set_source_rgba(0,0,0,1)
            ctx.set_line_cap(cairo.LINE_CAP_ROUND); 
            ctx.move_to(edge.from_node.x*w,h-edge.from_node.y*h)
            ctx.line_to(edge.to_node.x*w,h-edge.to_node.y*h)
            ctx.stroke()
        for node in self.nodes:
            ctx.set_source_rgba(0,1,0,0.5)
            ctx.arc(node.x*w, h-node.y*h, math.log(node.weight)+1, 0, 2 * pi)
            ctx.fill()
        surface.write_to_png(filename)
        
    def stats(self):
        return "%d nodes / %d edges"%(len(self.nodes), len(self.edges))

class Node(object):
    def __init__(self,x,y,weight=1,contents=[]):
        self.x = x;
        self.y = y;
        self.weight = weight
        self.contents = set(contents)
        self.edges_in = []
        self.edges_out = []
        self.changed = False
        self.dead = False
        
    def remove(self):
        """ Removes this node, and joins the edges that
            connect to it.
        """
        if len(self.edges_in) != 1 or len(self.edges_out) != 1:
            raise Exception("Can only remove node with 1 incoming and 1 outgoing edge.")
        prev = self.edges_in[0].from_node
        next = self.edges_out[0].to_node
        prev.changed = True
        self.changed = True
        next.changed = True
        self.edges_in[0].reconnect(prev, next)
        self.edges_out[0].delete()
        self.dead = True        
        
    def delete(self):
        self.dead = True
        
    def merge(self, other):
        """ Merge two nodes. New position is averaged, 
            and new edges are both nodes' edges.
        """
        tw = float(self.weight+other.weight)
        self.x = (self.x*self.weight + other.x*other.weight)/tw
        self.y = (self.y*self.weight + other.y*other.weight)/tw
        self.weight = tw
        self.contents |= other.contents
        # Rewire edges
        for edge in other.edges_in:
            if edge.from_node == self:
                edge.delete()
            else:
                edge.to_node = self
                self.edges_in.append(edge)
        for edge in other.edges_out:
            if edge.to_node == self:
                edge.delete()
            else:
                edge.from_node = self
                self.edges_out.append(edge)
        # Remove the other node
        other.delete()
        
    def p(self):
        """ Return node position as a point
        """
        return self.x, self.y

    

class Edge(object):
    def __init__(self,from_node,to_node,weight=1,contents=[]):
        if (from_node == to_node):
            raise Exception("Reflexive edges not allowed.")
        self.from_node = from_node
        self.to_node = to_node
        self.from_node.edges_out.append(self)
        self.to_node.edges_in.append(self)
        self.weight   = weight
        self.contents = set(contents)
        self.changed  = False
        self.dead     = False
    
    def angle(self):
        return math.atan2(self.to_node.y - self.from_node.y, self.to_node.x - self.from_node.x)
        
    def length(self):
        return ((self.to_node.x - self.from_node.x)**2 + (self.to_node.y - self.from_node.y)**2)**0.5
        
    def ps(self):
        return (self.from_node.p(), self.to_node.p())
    
    def reconnect(self, new_from_node, new_to_node):
        self.from_node.edges_out.remove(self)
        self.to_node.edges_in.remove(self)
        self.from_node = new_from_node
        self.to_node = new_to_node
        self.from_node.edges_out.append(self)
        self.to_node.edges_in.append(self)
    
    def delete(self, splice=False):
        """ Deletes edge without further operations. """
        self.from_node.edges_out.remove(self)
        self.to_node.edges_in.remove(self)
        self.dead = True
        if splice:
            self.graph.edges.remove(self)
    
    def unify(self, other):
        """ Unifies identical edges. No loss of information. """
        if (self.from_node != other.from_node or 
            self.to_node != other.to_node):
            raise Exception("Edges must come from and go to same nodes.")
        self.weight = self.weight + other.weight
        self.contents |= other.contents
        other.delete()
        
    def absorb(self, node):
        """ Absorbs a node into this edge. 
            Essentially reroutes this edge through given node.
        """
        self.changed = True
        node.changed = True
        if node == self.to_node or node == self.from_node:
            raise Exception("Cannot absorb node that edge already connects.")
        self.graph.add_edge(Edge(node, self.to_node, self.weight, self.contents))
        self.reconnect(self.from_node, node)



### Utils ###


def print_percentage(current, total, intervals=10):
    procs = globals().get('procs',[])
    print_at = [int((i/float(intervals))*total) for i in xrange(intervals+1)]
    prefix = '    '*len(procs)+'- '
    if current in print_at:
        print prefix+'%.1f%%'%(100*current/float(total))
    
def print_procedure_start(proc, message=''):
    procs = globals().get('procs',[])
    print '%s%s started. %s'%('    '*len(procs), proc, message)
    procs.append((proc, time.time()))
    globals()['procs'] = procs

def print_procedure_status(message):
    procs = globals().get('procs',[])
    print '    '*len(procs) + '- ' + message

def print_procedure_done(message=''):
    procs = globals().get('procs',[])
    proc,start_time = procs.pop()
    print '%s%s done in %.1fs. %s'%('    '*len(procs), proc, time.time()-start_time, message)
    
### Procedure ###
    
if __name__ == '__main__':
    # Set up logging
    # Initialization
    h,w = 2000,2000
    paru = PARU('UVA1/UVA_PARU_Out.xmf')
    
    # Get paths and do all the magic
    # paths = random.sample(paru.get_paths(),200)
    
    # print [(len(p1), len(p2)) for (p1, p2) in zip(paths, pathsl)]
    for i,paths in enumerate(paru.get_paths(1)):
        graph = Graph().load_paths(paths)
        graph.reduce_complete()
        graph.draw("step-%03d.png"%i,w=1000,h=700)
    # graph.reduce_complete()

    
    # combine_edges(graph)
    
    

geom.py

import math

# Shortcuts

sqrt  = math.sqrt
pi    = math.pi
atan2 = math.atan2

def point_add(a, b):
    return (a[0] + b[0], a[1] + b[1])

def point_sub(a, b):
    return (a[0] - b[0], a[1] - b[1])
    
def point_mul(a, f):
    return (a[0]*f, a[1]*f)
    
def point_dist(a, b):
    """ Distance between two points. """
    return ((a[0]-b[0]) ** 2 + (a[1]-b[1]) ** 2) ** 0.5

def line_point_dist((ax,ay), (bx,by), (px,py)):
    """ Distance from a point to given line. Also returns u
        with the projection of P on the line: z = l1 + u*(l2-l1)
    """
    lx, ly = bx-ax, by-ay
    if lx == 0 and ly == 0:
        return point_dist((ax,ay),(px,py)),0.5
    u = ((px - ax)*lx + (py - ay)*ly)/float(lx**2+ly**2)
    zx, zy = (ax + u*lx), (ay + u*ly)
    return point_dist((zx,zy),(px,py)), u
    
def line_angle((ax, ay), (bx, by), (cx, cy)):
    return atan2(by-ay, bx-ax) - atan2(cy-by, cx-bx)
    
def angle_fix(theta):
    """ Fixes an angle to a value between -pi and pi.
    """
    return ((theta + pi) % (2*pi)) - pi
    
def angle_average(theta1, theta2, f=0.5):
    diff = angle_fix(theta2-theta1)
    return theta1 + f*diff


if __name__ == '__main__':
    # Testing
    import doctest
    doctest.testmod()