<commit_before>from __future__ import division
import warnings
import numpy as np
import scipy.integrate
from scipy import exp, pi

class ComplexPath(object):
	""" A base class for paths in the complex plane """
	def __init__(self):
		self._integralCache = {}
		self._trapValuesCache = {}

	def trap_values(self, f, k, useCache=True):
		"""
		2**k+1 is the number of required points for the function f to
		be evaluated at.
		"""

		if f in self._trapValuesCache.keys() and useCache:
			vals = self._trapValuesCache[f]
			vals_k = int(np.log2(len(vals)-1))

			if vals_k == k:
				return vals
			elif vals_k > k:
				return vals[::2**(vals_k-k)]
			else:
				t = np.linspace(0, 1, 2**k+1)
				vals = np.empty(2**k+1, dtype=np.complex128)
				vals.fill(np.nan)
				vals[::2**(k-vals_k)] = self._trapValuesCache[f]
				vals[np.isnan(vals)] = f(self(t[np.isnan(vals)]))

				# cache values
				self._trapValuesCache[f] = vals
				return vals

		else:
			t = np.linspace(0, 1, 2**k+1)
			vals = f(self(t))
			if useCache:
				self._trapValuesCache[f] = vals
			return vals

	def plot(self, N=100, linecolor='C0', linestyle='-'):
		"""
		Use matplotlib to plot, but not show, the path as a
		2D plot in the Complex plane.  To show use pyplot.show()

		Parameters
		----------
		N : int, optional
			The number of points to use when plotting the path.  Default is 100
		linecolor : optional
			The colour of the plotted path, passed to the pyplot.plot function
			as the keyword argument of 'color'.
		linestyle : str, optional
			The line style of the plotted path, passed to the pyplot.plot
			function as the keyword argument of 'linestyle'.  Default is '-'
			which corresponds to a solid line.
		"""
		import matplotlib.pyplot as plt
		t = np.linspace(0,1,N)
		path = self(t)
		plt.plot(path.real, path.imag, color=linecolor, linestyle=linestyle)
		plt.xlabel('Re[$z$]', size=16)
		plt.ylabel('Im[$z$]', size=16)
		plt.gca().set_aspect(1)
		plt.tight_layout()

		# add arrow to indicate direction of path
		arrow_direction = (self(0.51) - self(0.5))/abs(self(0.51) - self(0.5))
		arrow_extent = 1e-6*arrow_direction
		ymin, ymax = plt.gca().get_ylim()
		xmin, xmax = plt.gca().get_xlim()
		head_length = max(abs(ymax - ymin), abs(xmax - xmin))/40.
		plt.arrow(self(0.5).real, self(0.5).imag,
				  arrow_extent.real, arrow_extent.imag,
				  head_width=head_length*2/3., head_length=head_length,
				  fc=linecolor, ec=linecolor)

	def show(self, *args, **kwargs):
		""" Shows the path as a 2D plot in the complex plane using
		the same arguments as the plot method """
		import matplotlib.pyplot as plt
		self.plot(*args, **kwargs)
		plt.show()

	def integrate(self, f, absTol=0, relTol=1e-12, divMax=15, intMethod='quad', verbose=False):
		"""
		Integrate the function f along the path using SciPy's Romberg
		algorithm.  The value of the integral is cached and will be
		reused if the method is called with same function f and tol.

		Parameters
		----------
		f : function of a single complex variable
		tol : float, optional
			The absolute tolerance passed to SciPy's Romberg function.
			Default is 1e-8.
		rombergDivMax : int, optional
			The maximum order of extrapolation passed to SciPy's Romberg function

		Returns
		-------
		integral : complex
			The integral of the function f along the path.
		"""

		args = (f, absTol, relTol, divMax, intMethod)
		if args in self._integralCache.keys():
			integral = self._integralCache[args]

		elif hasattr(self, '_reversePath') and args in self._reversePath._integralCache:
			# if we have already computed the reverse of this path
			integral = -self._reversePath._integralCache[args]

		else:
			integrand = lambda t: f(self(t))*self.dzdt(t)

			if intMethod == 'romb':
				integral = scipy.integrate.romberg(integrand, 0, 1, tol=absTol, rtol=relTol, divmax=divMax, show=verbose)
			elif intMethod == 'quad':
				integrand_real = lambda t: np.real(integrand(t))
				integrand_imag = lambda t: np.imag(integrand(t))

				integral_real, abserr_real = scipy.integrate.quad(integrand_real, 0, 1, epsabs=absTol, epsrel=relTol)
				integral_imag, abserr_imag = scipy.integrate.quad(integrand_imag, 0, 1, epsabs=absTol, epsrel=relTol)
				integral = integral_real + 1j*integral_imag
			else:
				raise ValueError("intMethod must be either 'romb' or 'quad'")

			if np.isnan(integral):
				raise RuntimeError('The integral along the segment %s is NaN.\
					\nThis is most likely due to a root being on or very close to the path of integration.'%self)

			self._integralCache[args] = integral

		return integral

class ComplexLine(ComplexPath):
	""" A straight line in the complex plane from a to b """
	def __init__(self, a, b):
		self.a, self.b = a, b
		self.dzdt = lambda t: self.b-self.a
		super(ComplexLine, self).__init__()

	def __str__(self):
		return 'ComplexLine from %.3f+%.3fi to %.3f+%.3fi' % (self.a.real, self.a.imag, self.b.real, self.b.imag)

	def __call__(self, t):
		""" The parameterization of the line in the variable t, where 0 <= t <= 1 """
		return self.a + t*(self.b-self.a)

	def distance(self, z):
		"""
		Distance from the point z to the closest point on the line.

		Parameters
		----------
		z : complex

		Returns
		-------
		float
			The distance from z to the point on the line which is
			closest to z.
		"""
		a, b = self.a, self.b

		# convert complex numbers to vectors
		A = np.array([a.real, a.imag])
		B = np.array([b.real, b.imag])
		Z = np.array([z.real, z.imag])

		# the projection of the point z onto the line a -> b is where
		# the parameter t is
		t = (Z-A).dot(B-A)/abs((B-A).dot(B-A))

		# but the line segment only has 0 <= t <= 1
		t = t.clip(0,1)

		# so the point on the line segment closest to z is
		c = self(t)
		return abs(c-z)

class ComplexArc(ComplexPath):
	""" An arc with center z0, radius R, initial angle t0 and change of angle dt """
	def __init__(self, z0, R, t0, dt):
		self.z0, self.R, self.t0, self.dt = z0, R, t0, dt
		self.dzdt = lambda t: 1j*self.dt*self.R*exp(1j*(self.t0 + t*self.dt))
		super(ComplexArc, self).__init__()

	def __str__(self):
		return 'ComplexArc: z0=%.3f, R=%.3f, t0=%.3f, dt=%.3f' % (self.z0, self.R, self.t0, self.dt)

	def __call__(self, t):
		""" The parameterization of the arc in the variable t, where 0 <= t <= 1 """
		return self.R*exp(1j*(self.t0 + t*self.dt)) + self.z0

	def distance(self, z):
		"""
		Distance from the point z to the closest point on the arc.

		Parameters
		----------
		z : complex

		Returns
		-------
		float
			The distance from z to the point on the arc which is
			closest to z.
		"""
		theta = np.angle(z-self.z0) 				# np.angle maps to (-pi,pi]
		theta = (theta-self.t0)%(2*pi) + self.t0 	# put theta in [t0,t0+2pi)

		if ((self.dt > 0 and self.t0 < theta < self.t0+self.dt)
			or (self.dt < 0 and self.t0+self.dt < theta - 2*pi < self.t0)):
			# the closest point to z lies on the arc
			return abs(self.R*exp(1j*theta) + self.z0 - z)
		else:
			# the closest point to z is one of the endpoints
			return min(abs(self(0)-z), abs(self(1)-z))

<commit_msg>Update line distance<commit_after>from __future__ import division
import warnings
import numpy as np
import scipy.integrate
from scipy import exp, pi

class ComplexPath(object):
	""" A base class for paths in the complex plane """
	def __init__(self):
		self._integralCache = {}
		self._trapValuesCache = {}

	def trap_values(self, f, k, useCache=True):
		"""
		2**k+1 is the number of required points for the function f to
		be evaluated at.
		"""

		if f in self._trapValuesCache.keys() and useCache:
			vals = self._trapValuesCache[f]
			vals_k = int(np.log2(len(vals)-1))

			if vals_k == k:
				return vals
			elif vals_k > k:
				return vals[::2**(vals_k-k)]
			else:
				t = np.linspace(0, 1, 2**k+1)
				vals = np.empty(2**k+1, dtype=np.complex128)
				vals.fill(np.nan)
				vals[::2**(k-vals_k)] = self._trapValuesCache[f]
				vals[np.isnan(vals)] = f(self(t[np.isnan(vals)]))

				# cache values
				self._trapValuesCache[f] = vals
				return vals

		else:
			t = np.linspace(0, 1, 2**k+1)
			vals = f(self(t))
			if useCache:
				self._trapValuesCache[f] = vals
			return vals

	def plot(self, N=100, linecolor='C0', linestyle='-'):
		"""
		Use matplotlib to plot, but not show, the path as a
		2D plot in the Complex plane.  To show use pyplot.show()

		Parameters
		----------
		N : int, optional
			The number of points to use when plotting the path.  Default is 100
		linecolor : optional
			The colour of the plotted path, passed to the pyplot.plot function
			as the keyword argument of 'color'.
		linestyle : str, optional
			The line style of the plotted path, passed to the pyplot.plot
			function as the keyword argument of 'linestyle'.  Default is '-'
			which corresponds to a solid line.
		"""
		import matplotlib.pyplot as plt
		t = np.linspace(0,1,N)
		path = self(t)
		plt.plot(path.real, path.imag, color=linecolor, linestyle=linestyle)
		plt.xlabel('Re[$z$]', size=16)
		plt.ylabel('Im[$z$]', size=16)
		plt.gca().set_aspect(1)
		plt.tight_layout()

		# add arrow to indicate direction of path
		arrow_direction = (self(0.51) - self(0.5))/abs(self(0.51) - self(0.5))
		arrow_extent = 1e-6*arrow_direction
		ymin, ymax = plt.gca().get_ylim()
		xmin, xmax = plt.gca().get_xlim()
		head_length = max(abs(ymax - ymin), abs(xmax - xmin))/40.
		plt.arrow(self(0.5).real, self(0.5).imag,
				  arrow_extent.real, arrow_extent.imag,
				  head_width=head_length*2/3., head_length=head_length,
				  fc=linecolor, ec=linecolor)

	def show(self, *args, **kwargs):
		""" Shows the path as a 2D plot in the complex plane using
		the same arguments as the plot method """
		import matplotlib.pyplot as plt
		self.plot(*args, **kwargs)
		plt.show()

	def integrate(self, f, absTol=0, relTol=1e-12, divMax=15, intMethod='quad', verbose=False):
		"""
		Integrate the function f along the path using SciPy's Romberg
		algorithm.  The value of the integral is cached and will be
		reused if the method is called with same function f and tol.

		Parameters
		----------
		f : function of a single complex variable
		tol : float, optional
			The absolute tolerance passed to SciPy's Romberg function.
			Default is 1e-8.
		rombergDivMax : int, optional
			The maximum order of extrapolation passed to SciPy's Romberg function

		Returns
		-------
		integral : complex
			The integral of the function f along the path.
		"""

		args = (f, absTol, relTol, divMax, intMethod)
		if args in self._integralCache.keys():
			integral = self._integralCache[args]

		elif hasattr(self, '_reversePath') and args in self._reversePath._integralCache:
			# if we have already computed the reverse of this path
			integral = -self._reversePath._integralCache[args]

		else:
			integrand = lambda t: f(self(t))*self.dzdt(t)

			if intMethod == 'romb':
				integral = scipy.integrate.romberg(integrand, 0, 1, tol=absTol, rtol=relTol, divmax=divMax, show=verbose)
			elif intMethod == 'quad':
				integrand_real = lambda t: np.real(integrand(t))
				integrand_imag = lambda t: np.imag(integrand(t))

				integral_real, abserr_real = scipy.integrate.quad(integrand_real, 0, 1, epsabs=absTol, epsrel=relTol)
				integral_imag, abserr_imag = scipy.integrate.quad(integrand_imag, 0, 1, epsabs=absTol, epsrel=relTol)
				integral = integral_real + 1j*integral_imag
			else:
				raise ValueError("intMethod must be either 'romb' or 'quad'")

			if np.isnan(integral):
				raise RuntimeError('The integral along the segment %s is NaN.\
					\nThis is most likely due to a root being on or very close to the path of integration.'%self)

			self._integralCache[args] = integral

		return integral

class ComplexLine(ComplexPath):
	""" A straight line in the complex plane from a to b """
	def __init__(self, a, b):
		self.a, self.b = a, b
		self.dzdt = lambda t: self.b-self.a
		super(ComplexLine, self).__init__()

	def __str__(self):
		return 'ComplexLine from %.3f+%.3fi to %.3f+%.3fi' % (self.a.real, self.a.imag, self.b.real, self.b.imag)

	def __call__(self, t):
		""" The parameterization of the line in the variable t, where 0 <= t <= 1 """
		return self.a + t*(self.b-self.a)

	def distance(self, z):
		"""
		Distance from the point z to the closest point on the line.

		Parameters
		----------
		z : complex

		Returns
		-------
		float
			The distance from z to the point on the line which is
			closest to z.
		"""
		# convert complex numbers to vectors
		A = np.array([self.a.real, self.a.imag])
		B = np.array([self.b.real, self.b.imag])
		Z = np.array([z.real, z.imag])

		# the projection of the point z onto the line a -> b is where
		# the parameter t is
		t = (Z-A).dot(B-A)/abs((B-A).dot(B-A))

		# but the line segment only has 0 <= t <= 1
		t = t.clip(0,1)

		# so the point on the line segment closest to z is
		c = self(t)
		return abs(c-z)

class ComplexArc(ComplexPath):
	""" An arc with center z0, radius R, initial angle t0 and change of angle dt """
	def __init__(self, z0, R, t0, dt):
		self.z0, self.R, self.t0, self.dt = z0, R, t0, dt
		self.dzdt = lambda t: 1j*self.dt*self.R*exp(1j*(self.t0 + t*self.dt))
		super(ComplexArc, self).__init__()

	def __str__(self):
		return 'ComplexArc: z0=%.3f, R=%.3f, t0=%.3f, dt=%.3f' % (self.z0, self.R, self.t0, self.dt)

	def __call__(self, t):
		""" The parameterization of the arc in the variable t, where 0 <= t <= 1 """
		return self.R*exp(1j*(self.t0 + t*self.dt)) + self.z0

	def distance(self, z):
		"""
		Distance from the point z to the closest point on the arc.

		Parameters
		----------
		z : complex

		Returns
		-------
		float
			The distance from z to the point on the arc which is
			closest to z.
		"""
		theta = np.angle(z-self.z0) 				# np.angle maps to (-pi,pi]
		theta = (theta-self.t0)%(2*pi) + self.t0 	# put theta in [t0,t0+2pi)

		if ((self.dt > 0 and self.t0 < theta < self.t0+self.dt)
			or (self.dt < 0 and self.t0+self.dt < theta - 2*pi < self.t0)):
			# the closest point to z lies on the arc
			return abs(self.R*exp(1j*theta) + self.z0 - z)
		else:
			# the closest point to z is one of the endpoints
			return min(abs(self(0)-z), abs(self(1)-z))

<|endoftext|>


<commit_before>import networkx as nx
import heapq
import os
import pickle
from collections import namedtuple

def _get_plot_libs():
    import matplotlib.pyplot as plt
    try:
        from networkx import graphviz_layout
    except ImportError:
        raise ImportError("This program needs Graphviz and either PyGraphviz or Pydot")
    return plt, graphviz_layout

class NewmanGreedy:

    RenameMapping = namedtuple('RenameMapping', ['integer', 'original'])

    def __init__(self, graph, snapshot_size=None, forced_clusters = None):
        self.forced_clusters = set(forced_clusters) if forced_clusters else None
        graph = self.remove_orphans(graph)
        self.rename_map = self.remap(graph)
        nx.relabel_nodes(graph, self.rename_map.integer, copy=False)
        self.orig = graph
        self.super_graph = graph.copy()
        self.dendrogram = nx.Graph()
        self.snapshot_size = snapshot_size
        self.pair_cost_heap = []
        self.quality_history = []
        self.den_num = self.super_graph.number_of_nodes()
        for node in graph.node:
            if (isinstance(node, int)) and node > self.den_num:
                self.den_num = node + 1

        def e_weight(graph, e):
            edge = graph[e[0]][e[1]]
            if isinstance(edge, int):
                return edge
            elif isinstance(edge, dict) and 'weight' in edge:
                return edge['weight']
            else:
                return 1
        num_edges = 2 * sum(e_weight(graph, e) for e in graph.edges_iter()) #2*graph.number_of_edges()

        quality = 0.0
        for cluster_id in graph.nodes_iter():
            #node_degree = graph.degree(cluster_id)
            node_degree = 0
            for tags in graph[cluster_id].itervalues():
                if isinstance(tags, dict) and 'weight' in tags:
                    node_degree += tags['weight']
                else:
                    node_degree += 1

            ai = float(node_degree) / num_edges
            self.super_graph.node[cluster_id] = ai
            self.dendrogram.add_node(cluster_id)
            # From equation (1) in secion II of the Newman paper
            quality -= float(node_degree * node_degree) / (num_edges * num_edges)

        for (cluster_id1, cluster_id2) in graph.edges_iter():
            if isinstance(graph[cluster_id1][cluster_id2], (int, float)):
                old_e_ij = graph[cluster_id1][cluster_id2]
            elif isinstance(graph[cluster_id1][cluster_id2], dict) and 'weight' in graph[cluster_id1][cluster_id2]:
                old_e_ij = graph[cluster_id1][cluster_id2]['weight']
            else:
                old_e_ij = 1.0
            eij = float(old_e_ij) / num_edges
            self.super_graph[cluster_id1][cluster_id2] = eij
            self.super_graph[cluster_id2][cluster_id1] = eij
        # Could be calcuated in the edge for-loop above at a loss of code reuse or efficiency
        #for (id1, id2) in graph.edges_iter():
        #    self.add_pair_to_cost_heap(id1, id2)
        self.reheapify()
        if self.forced_clusters:
            self.build_forced_clusters()
        self.run_greedy_clustering(quality)
        nx.relabel_nodes(self.dendrogram, self.rename_map.original, copy=False)

    def build_forced_clusters(self):
        # create a cluster from "unspecified" nodes
        precluster_nodes = []
        for node in self.rename_map.integer.iterkeys():
            if node in self.forced_clusters:
                precluster_nodes.append(self.rename_map.integer[node])
        while len(precluster_nodes) > 1:
            self.combine_clusters(precluster_nodes[0], precluster_nodes[1])
            precluster_nodes.pop(0)
            precluster_nodes.pop(0)
            while precluster_nodes:
                self.combine_clusters(self.den_num-1, precluster_nodes[0])
                precluster_nodes.pop(0)

    def remove_orphans(self, graph):
        # remove orphan nodes except those in "unspecified" cluster
        orphans = []
        for x in graph.nodes():
            if not graph.degree(x) and x not in self.forced_clusters:
                orphans.append(x)
        graph.remove_nodes_from(orphans)
        return graph

    def remap(self, graph):
        mapping_to_int = {}
        mapping_to_orig = {}
        for node_index, node in enumerate(graph.nodes_iter()):
            mapping_to_int[node]= node_index
            mapping_to_orig[node_index] = node
        return self.RenameMapping(mapping_to_int, mapping_to_orig)

    def reheapify(self):
        del self.pair_cost_heap
        self.pair_cost_heap = []
        for (id1, id2) in self.super_graph.edges_iter():
            self.add_pair_to_cost_heap(id1, id2)

    def run_greedy_clustering(self, quality, reheap_steps=500):
        self.quality_history = [quality]
        if self.snapshot_size == None or self.snapshot_size > self.super_graph.number_of_nodes():
            self.snapshot = self.super_graph.copy()
        last_heapify = self.super_graph.number_of_nodes()
        for num_clusters in xrange(self.super_graph.number_of_nodes(), 1, -1):
            if last_heapify > num_clusters + reheap_steps:
                self.reheapify()
                last_heapify = num_clusters
            if self.pair_cost_heap:
                [id1, id2, qual_diff] = self.combine_top_clusters()
                self.combine_clusters(id1, id2)
                quality += qual_diff
                if self.snapshot_size and len(self.super_graph) == self.snapshot_size:
                    self.snapshot = self.super_graph.copy()
                self.quality_history.append(quality)
        for node in self.super_graph.nodes():
            # Combining nodes in the above loop can create orphan nodes (which
            # may represent other clusters). We create edges to the main cluster
            # constructed above. This is necessary for dendrogram_crawl to find
            # and return these orphan nodes/clusters.
            if node != max(self.super_graph.nodes()):
                self.combine_clusters(x, max(self.super_graph.nodes()))
                self.quality_history.append(quality)

    def combine_top_clusters(self):
        while True:
            (qd, id1, id2) = heapq.heappop(self.pair_cost_heap)
            # Ignore edges that do not exist anymore
            if(self.super_graph.has_node(id1) and self.super_graph.has_node(id2)):
                # Maximize quality difference is negated from add_pair_to_cost_heap
                return [id1, id2, -qd]

    def add_pair_to_cost_heap(self, id1, id2):
        qd = self.quality_difference(id1, id2)
        if(id2 < id1):
            temp = id1
            id1 = id2
            id2 = temp
        # Negate quality difference (to maximize), AND id1 < id2
        heapq.heappush(self.pair_cost_heap, (-qd, id1, id2))

    # The "Change in Q" as described by section II of the Newman paper
    def quality_difference(self, cluster_id1, cluster_id2):
        ai = float(self.super_graph.node[cluster_id1])
        aj = float(self.super_graph.node[cluster_id2])
        eij = float(self.super_graph[cluster_id1][cluster_id2])
        return 2.0*(eij - ai*aj)

    def combine_clusters(self, cluster_id1, cluster_id2):
        combine_id = self.den_num
        self.den_num += 1

        # Add combined node
        c1_con = self.super_graph[cluster_id1]
        c2_con = self.super_graph[cluster_id2]
        c12_nodes = set(c1_con.keys()).union(set(c2_con.keys()))

        self.rename_map.original[combine_id] = combine_id
        self.rename_map.integer[combine_id] = combine_id

        self.super_graph.add_node(combine_id)
        combined_degree = self.super_graph.node[cluster_id1] + self.super_graph.node[cluster_id2]
        self.super_graph.node[combine_id] = combined_degree
        for outer_node in c12_nodes:
            total = 0.0
            # ignore edges between the two clusters
            if(outer_node == cluster_id1 or outer_node == cluster_id2):
                continue
            # sum edge weights to clusters that are reached by both clusters
            if(outer_node in c1_con):
                total += c1_con[outer_node]
            if(outer_node in c2_con):
                total += c2_con[outer_node]
            #self.super_graph.add_edge(combine_id, outer_node)
            #self.super_graph.add_edge(outer_node, combine_id)
            self.super_graph[combine_id][outer_node] = total
            self.super_graph[outer_node][combine_id] = total
            self.add_pair_to_cost_heap(combine_id, outer_node)

        # Remove old nodes
        # TODO the except should be removed and the initial weights bug solved...
        try: self.super_graph.remove_node(cluster_id1)
        except nx.exception.NetworkXError: pass
        try: self.super_graph.remove_node(cluster_id2)
        except nx.exception.NetworkXError: pass

        # Update dendrogram
        self.dendrogram.add_node(combine_id)
        self.dendrogram.add_edge(combine_id, cluster_id1)
        self.dendrogram.add_edge(combine_id, cluster_id2)

    def dendrogram_crawl(self, start, priors=None, max_steps=None):
        if priors == None:
            priors = set()
        fringe = []

        # Helper function to push node neighbors into fringe
        def push_dend_list(fringe, priors, node):
            priors.add(node)
            for inode in self.dendrogram[node]:
                if inode not in priors:
                    heapq.heappush(fringe, -inode)

        priors.add(start)
        heapq.heappush(fringe, -start)

        step = 0
        while len(fringe) > 0 and (max_steps == None or step < max_steps):
            node = -heapq.heappop(fringe)
            push_dend_list(fringe, priors, node)
            step += 1

        return priors, fringe

    def get_clusters(self, num_clusters=None):
        if num_clusters == None:
            index, value = max(enumerate(self.quality_history), key=lambda iv: iv[1])
            num_clusters = len(self.quality_history) - index

        nx.relabel_nodes(self.dendrogram, self.rename_map.integer, copy=False)
        start_node = max(self.dendrogram.nodes())

        priors, fringe = self.dendrogram_crawl(start=start_node,
                                               max_steps=num_clusters-1)

        # Double check we got the right number of values
        if len(fringe) != num_clusters:
            raise ValueError("get_clusters failed to retrieve "+
                "%d clusters correctly (got %d instead)" % (num_clusters, len(fringe)))

        clusters = []
        for neg_clust_start in fringe:
            clust_start = -neg_clust_start
            cprior, cfringe = self.dendrogram_crawl(start=clust_start,
                                                    priors=priors.copy())
            clusters.append(set(self.rename_map.original[n] for n in cprior
                                if n <= clust_start and self.orig.has_node(n)))
        nx.relabel_nodes(self.dendrogram, self.rename_map.original, copy=False)
        return sorted(clusters, key=lambda c: -len(c))

    def get_super_graph(self, size=None):
        if size == None:
            return self.super_graph

        clusters = self.get_clusters(size)
        #TODO recombine nodes...

    def plot_quality_history(self, plot_name, out_file_name, show=True):
        plt, graphviz_layout = _get_plot_libs()
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(range(len(self.quality_history), 0, -1), self.quality_history)
        plt.title('Number of Clusters vs. Modularity for %s' % plot_name)
        ax.set_xlabel('Number of Clusters')
        ax.set_ylabel('Modularity Score')
        plt.savefig(out_file_name)
        if show:
            plt.show()

    def plot_dendrogram(self, filename, fsize=10):
        plt, graphviz_layout = _get_plot_libs()
        pos = graphviz_layout(self.dendrogram, prog='twopi', args='')
        plt.figure(figsize=(10,10))
        nx.draw(self.dendrogram, pos, node_size=10,font_size=fsize, alpha=0.5,
                node_color="blue", with_labels=True)
        plt.axis('equal')
        plt.savefig(filename)
        plt.show()

    @staticmethod
    def build_load(graph, graph_name, regen_clustering=False, snapshot_size=None):
        name, ext = os.path.splitext(graph_name)
        pickleFName = name + '.pkl'
        try:
            if regen_clustering:
                raise Exception("Skip loading")
            with open(pickleFName, 'rb') as pklFile:
                cluster_graph = pickle.load(pklFile)
        except:
            cluster_graph = NewmanGreedy(graph, snapshot_size=snapshot_size)
            with open(pickleFName, 'wb') as pklFile:
                pickle.dump(cluster_graph, pklFile)
        return cluster_graph

def main():
    G = nx.karate_club_graph();
    N = NewmanGreedy(G)
    print N.quality_history
    try:
        N.plot_dendrogram()
        N.plot_quality_history('Karate', os.path.join(os.path.dirname(__file__), 'pics', 'karate'))
    except:
        pass

if __name__ == '__main__':
    main()
<commit_msg>clustering fix, style changes<commit_after>import networkx as nx
import heapq
import os
import pickle
from collections import namedtuple

def _get_plot_libs():
    import matplotlib.pyplot as plt
    try:
        from networkx import graphviz_layout
    except ImportError:
        raise ImportError("This program needs Graphviz and either PyGraphviz or Pydot")
    return plt, graphviz_layout

class NewmanGreedy:

    RenameMapping = namedtuple('RenameMapping', ['integer', 'original'])

    def __init__(self, graph, snapshot_size=None, forced_clusters = None):
        self.forced_clusters = set(forced_clusters) if forced_clusters else None
        graph = self.remove_orphans(graph)
        self.rename_map = self.remap(graph)
        nx.relabel_nodes(graph, self.rename_map.integer, copy=False)
        self.orig = graph
        self.super_graph = graph.copy()
        self.dendrogram = nx.Graph()
        self.snapshot_size = snapshot_size
        self.pair_cost_heap = []
        self.quality_history = []
        self.den_num = self.super_graph.number_of_nodes()
        for node in graph.node:
            if (isinstance(node, int)) and node > self.den_num:
                self.den_num = node + 1

        def e_weight(graph, e):
            edge = graph[e[0]][e[1]]
            if isinstance(edge, int):
                return edge
            elif isinstance(edge, dict) and 'weight' in edge:
                return edge['weight']
            else:
                return 1
        num_edges = 2 * sum(e_weight(graph, e) for e in graph.edges_iter()) #2*graph.number_of_edges()

        quality = 0.0
        for cluster_id in graph.nodes_iter():
            #node_degree = graph.degree(cluster_id)
            node_degree = 0
            for tags in graph[cluster_id].itervalues():
                if isinstance(tags, dict) and 'weight' in tags:
                    node_degree += tags['weight']
                else:
                    node_degree += 1

            ai = float(node_degree) / num_edges
            self.super_graph.node[cluster_id] = ai
            self.dendrogram.add_node(cluster_id)
            # From equation (1) in secion II of the Newman paper
            quality -= float(node_degree * node_degree) / (num_edges * num_edges)

        for (cluster_id1, cluster_id2) in graph.edges_iter():
            if isinstance(graph[cluster_id1][cluster_id2], (int, float)):
                old_e_ij = graph[cluster_id1][cluster_id2]
            elif isinstance(graph[cluster_id1][cluster_id2], dict) and 'weight' in graph[cluster_id1][cluster_id2]:
                old_e_ij = graph[cluster_id1][cluster_id2]['weight']
            else:
                old_e_ij = 1.0
            eij = float(old_e_ij) / num_edges
            self.super_graph[cluster_id1][cluster_id2] = eij
            self.super_graph[cluster_id2][cluster_id1] = eij
        # Could be calcuated in the edge for-loop above at a loss of code reuse or efficiency
        #for (id1, id2) in graph.edges_iter():
        #    self.add_pair_to_cost_heap(id1, id2)
        self.reheapify()
        if self.forced_clusters:
            self.build_forced_clusters()
        self.run_greedy_clustering(quality)
        nx.relabel_nodes(self.dendrogram, self.rename_map.original, copy=False)

    def build_forced_clusters(self):
        # create a cluster from "unspecified" nodes
        precluster_nodes = []
        for node in self.rename_map.integer.iterkeys():
            if node in self.forced_clusters:
                precluster_nodes.append(self.rename_map.integer[node])
        while len(precluster_nodes) > 1:
            self.combine_clusters(precluster_nodes[0], precluster_nodes[1])
            precluster_nodes.pop(0)
            precluster_nodes.pop(0)
            while precluster_nodes:
                self.combine_clusters(self.den_num-1, precluster_nodes[0])
                precluster_nodes.pop(0)

    def remove_orphans(self, graph):
        # remove orphan nodes except those in "unspecified" cluster
        orphans = [node for node in graph if not graph.degree(node) and node not in self.forced_clusters]
        graph.remove_nodes_from(orphans)
        return graph

    def remap(self, graph):
        mapping_to_int = {}
        mapping_to_orig = {}
        for node_index, node in enumerate(graph.nodes_iter()):
            mapping_to_int[node]= node_index
            mapping_to_orig[node_index] = node
        return self.RenameMapping(mapping_to_int, mapping_to_orig)

    def reheapify(self):
        del self.pair_cost_heap
        self.pair_cost_heap = []
        for (id1, id2) in self.super_graph.edges_iter():
            self.add_pair_to_cost_heap(id1, id2)

    def run_greedy_clustering(self, quality, reheap_steps=500):
        self.quality_history = [quality]
        if self.snapshot_size == None or self.snapshot_size > self.super_graph.number_of_nodes():
            self.snapshot = self.super_graph.copy()
        last_heapify = self.super_graph.number_of_nodes()
        for num_clusters in xrange(self.super_graph.number_of_nodes(), 1, -1):
            if last_heapify > num_clusters + reheap_steps:
                self.reheapify()
                last_heapify = num_clusters
            if self.pair_cost_heap:
                [id1, id2, qual_diff] = self.combine_top_clusters()
                self.combine_clusters(id1, id2)
                quality += qual_diff
                if self.snapshot_size and len(self.super_graph) == self.snapshot_size:
                    self.snapshot = self.super_graph.copy()
                self.quality_history.append(quality)
        for node in self.super_graph:
            # Combining nodes in the above loop can create orphan nodes (which
            # may represent other clusters). We create edges to the main cluster
            # constructed above. This is necessary for dendrogram_crawl to find
            # and return these orphan nodes/clusters.
            if node != max(self.super_graph):
                self.combine_clusters(x, max(self.super_graph))
                self.quality_history.append(quality)

    def combine_top_clusters(self):
        while True:
            (qd, id1, id2) = heapq.heappop(self.pair_cost_heap)
            # Ignore edges that do not exist anymore
            if(self.super_graph.has_node(id1) and self.super_graph.has_node(id2)):
                # Maximize quality difference is negated from add_pair_to_cost_heap
                return [id1, id2, -qd]

    def add_pair_to_cost_heap(self, id1, id2):
        qd = self.quality_difference(id1, id2)
        if(id2 < id1):
            temp = id1
            id1 = id2
            id2 = temp
        # Negate quality difference (to maximize), AND id1 < id2
        heapq.heappush(self.pair_cost_heap, (-qd, id1, id2))

    # The "Change in Q" as described by section II of the Newman paper
    def quality_difference(self, cluster_id1, cluster_id2):
        ai = float(self.super_graph.node[cluster_id1])
        aj = float(self.super_graph.node[cluster_id2])
        eij = float(self.super_graph[cluster_id1][cluster_id2])
        return 2.0*(eij - ai*aj)

    def combine_clusters(self, cluster_id1, cluster_id2):
        combine_id = self.den_num
        self.den_num += 1

        # Add combined node
        c1_con = self.super_graph[cluster_id1]
        c2_con = self.super_graph[cluster_id2]
        c12_nodes = set(c1_con.keys()).union(set(c2_con.keys()))

        self.rename_map.original[combine_id] = combine_id
        self.rename_map.integer[combine_id] = combine_id

        self.super_graph.add_node(combine_id)
        combined_degree = self.super_graph.node[cluster_id1] + self.super_graph.node[cluster_id2]
        self.super_graph.node[combine_id] = combined_degree
        for outer_node in c12_nodes:
            total = 0.0
            # ignore edges between the two clusters
            if(outer_node == cluster_id1 or outer_node == cluster_id2):
                continue
            # sum edge weights to clusters that are reached by both clusters
            if(outer_node in c1_con):
                total += c1_con[outer_node]
            if(outer_node in c2_con):
                total += c2_con[outer_node]
            #self.super_graph.add_edge(combine_id, outer_node)
            #self.super_graph.add_edge(outer_node, combine_id)
            self.super_graph[combine_id][outer_node] = total
            self.super_graph[outer_node][combine_id] = total
            self.add_pair_to_cost_heap(combine_id, outer_node)

        # Remove old nodes
        # TODO the except should be removed and the initial weights bug solved...
        try: self.super_graph.remove_node(cluster_id1)
        except nx.exception.NetworkXError: pass
        try: self.super_graph.remove_node(cluster_id2)
        except nx.exception.NetworkXError: pass

        # Update dendrogram
        self.dendrogram.add_node(combine_id)
        self.dendrogram.add_edge(combine_id, cluster_id1)
        self.dendrogram.add_edge(combine_id, cluster_id2)

    def dendrogram_crawl(self, start, priors=None, max_steps=None):
        if priors == None:
            priors = set()
        fringe = []

        # Helper function to push node neighbors into fringe
        def push_dend_list(fringe, priors, node):
            priors.add(node)
            for inode in self.dendrogram[node]:
                if inode not in priors:
                    heapq.heappush(fringe, -inode)

        priors.add(start)
        heapq.heappush(fringe, -start)

        step = 0
        while len(fringe) > 0 and (max_steps == None or step < max_steps):
            node = -heapq.heappop(fringe)
            push_dend_list(fringe, priors, node)
            step += 1

        return priors, fringe

    def get_clusters(self, num_clusters=None):
        if num_clusters == None:
            index, value = max(enumerate(self.quality_history), key=lambda iv: iv[1])
            num_clusters = len(self.quality_history) - index

        nx.relabel_nodes(self.dendrogram, self.rename_map.integer, copy=False)
        start_node = max(self.dendrogram)

        priors, fringe = self.dendrogram_crawl(start=start_node,
                                               max_steps=num_clusters-1)

        # Double check we got the right number of values
        if len(fringe) != num_clusters:
            raise ValueError("get_clusters failed to retrieve "+
                "%d clusters correctly (got %d instead)" % (num_clusters, len(fringe)))

        clusters = []
        for neg_clust_start in fringe:
            clust_start = -neg_clust_start
            cprior, cfringe = self.dendrogram_crawl(start=clust_start,
                                                    priors=priors.copy())
            clusters.append(set(self.rename_map.original[n] for n in cprior
                                if n <= clust_start and self.orig.has_node(n)))
        nx.relabel_nodes(self.dendrogram, self.rename_map.original, copy=False)
        return sorted(clusters, key=lambda c: -len(c))

    def get_super_graph(self, size=None):
        if size == None:
            return self.super_graph

        clusters = self.get_clusters(size)
        #TODO recombine nodes...

    def plot_quality_history(self, plot_name, out_file_name, show=True):
        plt, graphviz_layout = _get_plot_libs()
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(range(len(self.quality_history), 0, -1), self.quality_history)
        plt.title('Number of Clusters vs. Modularity for %s' % plot_name)
        ax.set_xlabel('Number of Clusters')
        ax.set_ylabel('Modularity Score')
        plt.savefig(out_file_name)
        if show:
            plt.show()

    def plot_dendrogram(self, filename, fsize=10):
        plt, graphviz_layout = _get_plot_libs()
        pos = graphviz_layout(self.dendrogram, prog='twopi', args='')
        plt.figure(figsize=(10,10))
        nx.draw(self.dendrogram, pos, node_size=10,font_size=fsize, alpha=0.5,
                node_color="blue", with_labels=True)
        plt.axis('equal')
        plt.savefig(filename)
        plt.show()

    @staticmethod
    def build_load(graph, graph_name, regen_clustering=False, snapshot_size=None):
        name, ext = os.path.splitext(graph_name)
        pickleFName = name + '.pkl'
        try:
            if regen_clustering:
                raise Exception("Skip loading")
            with open(pickleFName, 'rb') as pklFile:
                cluster_graph = pickle.load(pklFile)
        except:
            cluster_graph = NewmanGreedy(graph, snapshot_size=snapshot_size)
            with open(pickleFName, 'wb') as pklFile:
                pickle.dump(cluster_graph, pklFile)
        return cluster_graph

def main():
    G = nx.karate_club_graph();
    N = NewmanGreedy(G)
    print N.quality_history
    try:
        N.plot_dendrogram()
        N.plot_quality_history('Karate', os.path.join(os.path.dirname(__file__), 'pics', 'karate'))
    except:
        pass

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>
__copyright__ = "Copyright 2017-2018, http://radical.rutgers.edu"
__author__    = "Vivek Balasubramanian <vivek.balasubramanian@rutgers.edu>"
__license__   = "MIT"

import json
import pika
import queue
import os
import pickle

import threading       as mt
import multiprocessing as mp

import radical.pilot   as rp

from ...exceptions     import EnTKError
from ...               import states, Task
from ..base            import Base_TaskManager
from .task_processor   import create_td_from_task, create_task_from_rp

# ------------------------------------------------------------------------------
#
class TaskManager(Base_TaskManager):
    """
    A Task Manager takes the responsibility of dispatching tasks it receives
    from a queue for execution on to the available resources using a runtime
    system. In this case, the runtime system being used RADICAL Pilot. Once
    the tasks have completed execution, they are pushed on to another queue for
    other components of EnTK to access.

    :arguments:
        :pending_queue:     (list) List of queue(s) with tasks ready to be
                            executed. Currently, only one queue.
        :completed_queue:   (list) List of queue(s) with tasks that have
                            finished execution. Currently, only one queue.
        :rmgr:              (ResourceManager) Object to be used to access the
                            Pilot where the tasks can be submitted
        :rmq_conn_params:   (pika.connection.ConnectionParameters) object of
                            parameters necessary to connect to RabbitMQ

    Currently, EnTK is configured to work with one pending queue and one
    completed queue. In the future, the number of queues can be varied for
    different throughput requirements at the cost of additional Memory and CPU
    consumption.
    """

    # --------------------------------------------------------------------------
    #
    def __init__(self, sid, pending_queue, completed_queue, rmgr,
                       rmq_conn_params):

        super(TaskManager, self).__init__(sid, pending_queue, completed_queue,
                                          rmgr, rmq_conn_params,
                                          rts='radical.pilot')
        self._rts_runner = None
        self._submitted_tasks = dict()
        self._log.info('Created task manager object: %s', self._uid)
        self._prof.prof('tmgr_create', uid=self._uid)

    # --------------------------------------------------------------------------
    #
    def _tmgr(self, uid, rmgr, pending_queue, completed_queue,
                    rmq_conn_params):
        """
        **Purpose**: This method has 3 purposes: Respond to a heartbeat thread
                     indicating the live-ness of the RTS, receive tasks from the
                     pending_queue, start a new thread that processes these
                     tasks and submits to the RTS.

                     It is important to separate the reception of the tasks from
                     their processing due to RMQ/AMQP design. The channel (i.e.,
                     the thread that holds the channel) that is receiving msgs
                     from the RMQ server needs to be non-blocking as it can
                     interfere with the heartbeat intervals of RMQ.  Processing
                     a large number of tasks can considerable time and can block
                     the communication channel. Hence, the two are separated.

                     The new thread is responsible for pushing completed tasks
                     (returned by the RTS) to the dequeueing queue. It also
                     converts Tasks into TDs and CUs into (partially described)
                     Tasks. This conversion is necessary since the current RTS
                     is RADICAL Pilot. Once Tasks are recovered from a CU, they
                     are then pushed to the completed_queue. At all state
                     transititons, they are synced (blocking) with the
                     AppManager in the master process.

                     In addition the tmgr also receives heartbeat 'request' msgs
                     from the heartbeat-request queue. It responds with
                     a 'response' message to the heartbeart-response queue.

        **Details**: The AppManager can re-invoke the tmgr process with this
                     function if the execution of the workflow is still
                     incomplete. There is also population of a dictionary,
                     `placeholders`, which stores the path of each of the
                     tasks on the remote machine.
        """

        try:

            # ------------------------------------------------------------------
            def heartbeat_response(mq_channel, conn_params):

                channel = mq_channel
                try:

                    # Get request from heartbeat-req for heartbeat response
                    method_frame, props, body = \
                                  channel.basic_get(queue=self._hb_request_q)

                    if not body:
                        return

                    self._log.info('Received heartbeat request')
                    try:
                        nprops = pika.BasicProperties(
                                            correlation_id=props.correlation_id)
                        channel.basic_publish(exchange='',
                                              routing_key=self._hb_response_q,
                                              properties=nprops,
                                              body='response')
                    except (pika.exceptions.ConnectionClosed,
                            pika.exceptions.ChannelClosed):
                        connection = pika.BlockingConnection(conn_params)
                        channel = connection.channel()
                        nprops = pika.BasicProperties(
                                            correlation_id=props.correlation_id)
                        channel.basic_publish(exchange='',
                                              routing_key=self._hb_response_q,
                                              properties=nprops,
                                              body='response')

                    self._log.info('Sent heartbeat response')

                    channel.basic_ack(delivery_tag=method_frame.delivery_tag)

                except Exception as ex:
                    self._log.exception('Failed to respond to heartbeat, ' +
                                        'error: %s', ex)
                    raise EnTKError(ex) from ex
            # ------------------------------------------------------------------

            self._prof.prof('tmgr process started', uid=self._uid)
            self._log.info('Task Manager process started')

            # Acquire a connection+channel to the rmq server
            mq_connection = pika.BlockingConnection(rmq_conn_params)
            mq_channel = mq_connection.channel()

            # Make sure the heartbeat response queue is empty
            mq_channel.queue_delete(queue=self._hb_response_q)
            mq_channel.queue_declare(queue=self._hb_response_q)

            # Queue for communication between threads of this process
            task_queue = queue.Queue()

            # Pickle file for task id history.
            # TODO: How do you take care the first execution.
            if os.path.exists('.task_submitted.pkl'):
                with open('.task_submitted.pkl', 'rb') as f:
                    self._submitted_tasks = pickle.load(f)

            # Start second thread to receive tasks and push to RTS
            self._rts_runner = mt.Thread(target=self._process_tasks,
                                         args=(task_queue, rmgr,
                                               rmq_conn_params))
            self._rts_runner.start()

            self._prof.prof('tmgr infrastructure setup done', uid=uid)

            # While we are supposed to run and the thread that does the work is
            # alive go.
            while not self._tmgr_terminate.is_set() and \
                      self._rts_runner.is_alive():

                try:

                    # Get tasks from the pending queue
                    method_frame, _ , body = \
                                    mq_channel.basic_get(queue=pending_queue[0])

                    if body:

                        body = json.loads(body)
                        task_queue.put(body)

                        mq_channel.basic_ack(
                                delivery_tag=method_frame.delivery_tag)

                    heartbeat_response(mq_channel, rmq_conn_params)

                except Exception as ex:
                    self._log.exception('Error in task execution: %s', ex)
                    raise EnTKError(ex) from ex
            self._log.debug('Exited TMGR main loop')

        except KeyboardInterrupt as ex:

            self._log.exception('Execution interrupted (probably by Ctrl+C), '
                                'cancel tmgr process gracefully...')
            raise KeyboardInterrupt from ex

        except Exception as ex:

            self._log.exception('%s failed with %s', self._uid, ex)
            raise EnTKError(ex) from ex

        finally:

            self._prof.prof('tmgr_term', uid=uid)

            if self._rts_runner:
                self._rts_runner.join()

            self._log.debug('TMGR RTS Runner joined')

            mq_connection.close()
            self._log.debug('TMGR RMQ connection closed')
            self._prof.close()
            self._log.debug('TMGR profile closed')

    # --------------------------------------------------------------------------
    #
    def _process_tasks(self, task_queue, rmgr, rmq_conn_params):
        '''
        **Purpose**: The new thread that gets spawned by the main tmgr process
                     invokes this function. This function receives tasks from
                     'task_queue' and submits them to the RADICAL Pilot RTS.
        '''

        placeholders = dict()
        placeholder_lock = mt.Lock()

        # ----------------------------------------------------------------------
        def load_placeholder(task):
            with placeholder_lock:
                parent_pipeline = str(task.parent_pipeline['uid'])
                parent_stage = str(task.parent_stage['uid'])

                if parent_pipeline not in placeholders:
                    placeholders[parent_pipeline] = dict()

                if parent_stage not in placeholders[parent_pipeline]:
                    placeholders[parent_pipeline][parent_stage] = dict()

                if None not in [parent_pipeline, parent_stage, task.uid]:
                    placeholders[parent_pipeline][parent_stage][task.uid] = \
                                                            {'path': task.path,
                                                             'uid': task.uid}

        # ----------------------------------------------------------------------
        def task_state_cb(rp_task, state, cb_data):

            try:

                channel = cb_data['channel']
                conn_params = cb_data['params']
                self._log.debug('Task %s in state %s' % (rp_task.uid,
                                                         rp_task.state))

                if rp_task.state in rp.FINAL:

                    task = create_task_from_rp(rp_task, self._prof)

                    self._advance(task, 'Task', states.COMPLETED,
                                  channel, conn_params,
                                  '%s-cb-to-sync' % self._sid)

                    load_placeholder(task)

                    task_as_dict = json.dumps(task.to_dict())
                    try:
                        channel.basic_publish(exchange='',
                                              routing_key='%s-completedq-1' % self._sid,
                                              body=task_as_dict)
                    except (pika.exceptions.ConnectionClosed,
                            pika.exceptions.ChannelClosed):
                        connection = pika.BlockingConnection(conn_params)
                        channel = connection.channel()
                        channel.basic_publish(exchange='',
                                              routing_key='%s-completedq-1' % self._sid,
                                              body=task_as_dict)

                    self._log.info('Pushed task %s with state %s to completed '
                                   'queue %s-completedq-1',
                                   task.uid, task.state, self._sid)

            except KeyboardInterrupt as ex:
                self._log.exception('Execution interrupted (probably by Ctrl+C)'
                                    ' exit callback thread gracefully...')
                raise KeyboardInterrupt(ex) from ex

            except Exception as ex:
                self._log.exception('Error in RP callback thread: %s', ex)
                raise EnTKError(ex) from ex
        # ----------------------------------------------------------------------

        mq_connection = pika.BlockingConnection(rmq_conn_params)
        mq_channel = mq_connection.channel()

        rp_tmgr = rp.TaskManager(session=rmgr._session)
        rp_tmgr.add_pilots(rmgr.pilot)
        rp_tmgr.register_callback(task_state_cb,
                                  cb_data={'channel': mq_channel,
                                           'params' : rmq_conn_params})

        try:
            pickle_file = open('.task_submitted.pkl', 'wb')
            while not self._tmgr_terminate.is_set():

                body = None

                try:
                    body = task_queue.get(block=True, timeout=10)

                except queue.Empty:
                    # Ignore, we don't always have new tasks to run
                    pass

                if not body:
                    continue

                task_queue.task_done()

                bulk_tds   = list()

                for msg in body:

                    task = Task()
                    task.from_dict(msg)

                    load_placeholder(task)
                    bulk_tds.append(create_td_from_task(
                                            task, placeholders,
                                            self._submitted_tasks, pickle_file,
                                            self._prof))

                    self._advance(task, 'Task', states.SUBMITTING,
                                  mq_channel, rmq_conn_params,
                                  '%s-tmgr-to-sync' % self._sid)

                rp_tmgr.submit_tasks(bulk_tds)
            mq_connection.close()
            self._log.debug('Exited RTS main loop. TMGR terminating')
        except KeyboardInterrupt as ex:
            self._log.exception('Execution interrupted (probably by Ctrl+C), '
                                'cancel task processor gracefully...')
            raise KeyboardInterrupt from ex
        except Exception as ex:
            self._log.exception('%s failed with %s', self._uid, ex)
            raise EnTKError(ex) from ex

        finally:
            rp_tmgr.close()

    # --------------------------------------------------------------------------
    #
    def start_manager(self):
        """
        **Purpose**: Method to start the tmgr process. The tmgr function
                     is not to be accessed directly. The function is started
                     in a separate thread using this method.
        """
        # pylint: disable=attribute-defined-outside-init, access-member-before-definition
        if self._tmgr_process:
            self._log.warn('tmgr process already running!')
            return

        try:
            # Redeclare the heartbeat queues in case they got deleted because
            # of the task manager failure.
            # If the queues exist this has no effect.

            mq_connection = pika.BlockingConnection(self._rmq_conn_params)
            mq_channel = mq_connection.channel()

            mq_channel.queue_declare(queue=self._hb_response_q)
            mq_channel.queue_declare(queue=self._hb_request_q)

            self._prof.prof('creating tmgr process', uid=self._uid)
            self._tmgr_terminate = mp.Event()

            self._tmgr_process = mp.Process(target=self._tmgr,
                                            name='task-manager',
                                            args=(self._uid,
                                                  self._rmgr,
                                                  self._pending_queue,
                                                  self._completed_queue,
                                                  self._rmq_conn_params)
                                            )

            self._log.info('Starting task manager process')
            self._prof.prof('starting tmgr process', uid=self._uid)

            self._tmgr_process.start()
            self._log.debug('tmgr pid %s' % self._tmgr_process.pid)

            return True

        except Exception as ex:

            self._log.exception('Task manager not started, error: %s', ex)
            self.terminate_manager()
            raise EnTKError(ex) from ex
        # pylint: enable=attribute-defined-outside-init, access-member-before-definition

# ------------------------------------------------------------------------------

<commit_msg>ids saved.<commit_after>
__copyright__ = "Copyright 2017-2018, http://radical.rutgers.edu"
__author__    = "Vivek Balasubramanian <vivek.balasubramanian@rutgers.edu>"
__license__   = "MIT"

import json
import pika
import queue
import os
import pickle

import threading       as mt
import multiprocessing as mp

import radical.pilot   as rp

from ...exceptions     import EnTKError
from ...               import states, Task
from ..base            import Base_TaskManager
from .task_processor   import create_td_from_task, create_task_from_rp

# ------------------------------------------------------------------------------
#
class TaskManager(Base_TaskManager):
    """
    A Task Manager takes the responsibility of dispatching tasks it receives
    from a queue for execution on to the available resources using a runtime
    system. In this case, the runtime system being used RADICAL Pilot. Once
    the tasks have completed execution, they are pushed on to another queue for
    other components of EnTK to access.

    :arguments:
        :pending_queue:     (list) List of queue(s) with tasks ready to be
                            executed. Currently, only one queue.
        :completed_queue:   (list) List of queue(s) with tasks that have
                            finished execution. Currently, only one queue.
        :rmgr:              (ResourceManager) Object to be used to access the
                            Pilot where the tasks can be submitted
        :rmq_conn_params:   (pika.connection.ConnectionParameters) object of
                            parameters necessary to connect to RabbitMQ

    Currently, EnTK is configured to work with one pending queue and one
    completed queue. In the future, the number of queues can be varied for
    different throughput requirements at the cost of additional Memory and CPU
    consumption.
    """

    # --------------------------------------------------------------------------
    #
    def __init__(self, sid, pending_queue, completed_queue, rmgr,
                       rmq_conn_params):

        super(TaskManager, self).__init__(sid, pending_queue, completed_queue,
                                          rmgr, rmq_conn_params,
                                          rts='radical.pilot')
        self._rts_runner = None
        self._submitted_tasks = dict()
        self._log.info('Created task manager object: %s', self._uid)
        self._prof.prof('tmgr_create', uid=self._uid)

    # --------------------------------------------------------------------------
    #
    def _tmgr(self, uid, rmgr, pending_queue, completed_queue,
                    rmq_conn_params):
        """
        **Purpose**: This method has 3 purposes: Respond to a heartbeat thread
                     indicating the live-ness of the RTS, receive tasks from the
                     pending_queue, start a new thread that processes these
                     tasks and submits to the RTS.

                     It is important to separate the reception of the tasks from
                     their processing due to RMQ/AMQP design. The channel (i.e.,
                     the thread that holds the channel) that is receiving msgs
                     from the RMQ server needs to be non-blocking as it can
                     interfere with the heartbeat intervals of RMQ.  Processing
                     a large number of tasks can considerable time and can block
                     the communication channel. Hence, the two are separated.

                     The new thread is responsible for pushing completed tasks
                     (returned by the RTS) to the dequeueing queue. It also
                     converts Tasks into TDs and CUs into (partially described)
                     Tasks. This conversion is necessary since the current RTS
                     is RADICAL Pilot. Once Tasks are recovered from a CU, they
                     are then pushed to the completed_queue. At all state
                     transititons, they are synced (blocking) with the
                     AppManager in the master process.

                     In addition the tmgr also receives heartbeat 'request' msgs
                     from the heartbeat-request queue. It responds with
                     a 'response' message to the heartbeart-response queue.

        **Details**: The AppManager can re-invoke the tmgr process with this
                     function if the execution of the workflow is still
                     incomplete. There is also population of a dictionary,
                     `placeholders`, which stores the path of each of the
                     tasks on the remote machine.
        """

        try:

            # ------------------------------------------------------------------
            def heartbeat_response(mq_channel, conn_params):

                channel = mq_channel
                try:

                    # Get request from heartbeat-req for heartbeat response
                    method_frame, props, body = \
                                  channel.basic_get(queue=self._hb_request_q)

                    if not body:
                        return

                    self._log.info('Received heartbeat request')
                    try:
                        nprops = pika.BasicProperties(
                                            correlation_id=props.correlation_id)
                        channel.basic_publish(exchange='',
                                              routing_key=self._hb_response_q,
                                              properties=nprops,
                                              body='response')
                    except (pika.exceptions.ConnectionClosed,
                            pika.exceptions.ChannelClosed):
                        connection = pika.BlockingConnection(conn_params)
                        channel = connection.channel()
                        nprops = pika.BasicProperties(
                                            correlation_id=props.correlation_id)
                        channel.basic_publish(exchange='',
                                              routing_key=self._hb_response_q,
                                              properties=nprops,
                                              body='response')

                    self._log.info('Sent heartbeat response')

                    channel.basic_ack(delivery_tag=method_frame.delivery_tag)

                except Exception as ex:
                    self._log.exception('Failed to respond to heartbeat, ' +
                                        'error: %s', ex)
                    raise EnTKError(ex) from ex
            # ------------------------------------------------------------------

            self._prof.prof('tmgr process started', uid=self._uid)
            self._log.info('Task Manager process started')

            # Acquire a connection+channel to the rmq server
            mq_connection = pika.BlockingConnection(rmq_conn_params)
            mq_channel = mq_connection.channel()

            # Make sure the heartbeat response queue is empty
            mq_channel.queue_delete(queue=self._hb_response_q)
            mq_channel.queue_declare(queue=self._hb_response_q)

            # Queue for communication between threads of this process
            task_queue = queue.Queue()

            # Pickle file for task id history.
            # TODO: How do you take care the first execution.
            pkl_path = self._path + '/.task_submitted.pkl'
            if os.path.exists(pkl_path):
                with open('.task_submitted.pkl', 'rb') as f:
                    self._submitted_tasks = pickle.load(f)

            # Start second thread to receive tasks and push to RTS
            self._rts_runner = mt.Thread(target=self._process_tasks,
                                         args=(task_queue, rmgr,
                                               rmq_conn_params))
            self._rts_runner.start()

            self._prof.prof('tmgr infrastructure setup done', uid=uid)

            # While we are supposed to run and the thread that does the work is
            # alive go.
            while not self._tmgr_terminate.is_set() and \
                      self._rts_runner.is_alive():

                try:

                    # Get tasks from the pending queue
                    method_frame, _ , body = \
                                    mq_channel.basic_get(queue=pending_queue[0])

                    if body:

                        body = json.loads(body)
                        task_queue.put(body)

                        mq_channel.basic_ack(
                                delivery_tag=method_frame.delivery_tag)

                    heartbeat_response(mq_channel, rmq_conn_params)

                except Exception as ex:
                    self._log.exception('Error in task execution: %s', ex)
                    raise EnTKError(ex) from ex
            self._log.debug('Exited TMGR main loop')

        except KeyboardInterrupt as ex:

            self._log.exception('Execution interrupted (probably by Ctrl+C), '
                                'cancel tmgr process gracefully...')
            raise KeyboardInterrupt from ex

        except Exception as ex:

            self._log.exception('%s failed with %s', self._uid, ex)
            raise EnTKError(ex) from ex

        finally:

            self._prof.prof('tmgr_term', uid=uid)

            if self._rts_runner:
                self._rts_runner.join()

            self._log.debug('TMGR RTS Runner joined')

            mq_connection.close()
            self._log.debug('TMGR RMQ connection closed')
            self._prof.close()
            self._log.debug('TMGR profile closed')

    # --------------------------------------------------------------------------
    #
    def _process_tasks(self, task_queue, rmgr, rmq_conn_params):
        '''
        **Purpose**: The new thread that gets spawned by the main tmgr process
                     invokes this function. This function receives tasks from
                     'task_queue' and submits them to the RADICAL Pilot RTS.
        '''

        placeholders = dict()
        placeholder_lock = mt.Lock()

        # ----------------------------------------------------------------------
        def load_placeholder(task):
            with placeholder_lock:
                parent_pipeline = str(task.parent_pipeline['uid'])
                parent_stage = str(task.parent_stage['uid'])

                if parent_pipeline not in placeholders:
                    placeholders[parent_pipeline] = dict()

                if parent_stage not in placeholders[parent_pipeline]:
                    placeholders[parent_pipeline][parent_stage] = dict()

                if None not in [parent_pipeline, parent_stage, task.uid]:
                    placeholders[parent_pipeline][parent_stage][task.uid] = \
                                                            {'path': task.path,
                                                             'uid': task.uid}

        # ----------------------------------------------------------------------
        def task_state_cb(rp_task, state, cb_data):

            try:

                channel = cb_data['channel']
                conn_params = cb_data['params']
                self._log.debug('Task %s in state %s' % (rp_task.uid,
                                                         rp_task.state))

                if rp_task.state in rp.FINAL:

                    task = create_task_from_rp(rp_task, self._prof)

                    self._advance(task, 'Task', states.COMPLETED,
                                  channel, conn_params,
                                  '%s-cb-to-sync' % self._sid)

                    load_placeholder(task)

                    task_as_dict = json.dumps(task.to_dict())
                    try:
                        channel.basic_publish(exchange='',
                                              routing_key='%s-completedq-1' % self._sid,
                                              body=task_as_dict)
                    except (pika.exceptions.ConnectionClosed,
                            pika.exceptions.ChannelClosed):
                        connection = pika.BlockingConnection(conn_params)
                        channel = connection.channel()
                        channel.basic_publish(exchange='',
                                              routing_key='%s-completedq-1' % self._sid,
                                              body=task_as_dict)

                    self._log.info('Pushed task %s with state %s to completed '
                                   'queue %s-completedq-1',
                                   task.uid, task.state, self._sid)

            except KeyboardInterrupt as ex:
                self._log.exception('Execution interrupted (probably by Ctrl+C)'
                                    ' exit callback thread gracefully...')
                raise KeyboardInterrupt(ex) from ex

            except Exception as ex:
                self._log.exception('Error in RP callback thread: %s', ex)
                raise EnTKError(ex) from ex
        # ----------------------------------------------------------------------

        mq_connection = pika.BlockingConnection(rmq_conn_params)
        mq_channel = mq_connection.channel()

        rp_tmgr = rp.TaskManager(session=rmgr._session)
        rp_tmgr.add_pilots(rmgr.pilot)
        rp_tmgr.register_callback(task_state_cb,
                                  cb_data={'channel': mq_channel,
                                           'params' : rmq_conn_params})

        try:
            pkl_path = self._path + '/.task_submitted.pkl'
            pickle_file = open(pkl_path, 'wb')
            while not self._tmgr_terminate.is_set():

                body = None

                try:
                    body = task_queue.get(block=True, timeout=10)

                except queue.Empty:
                    # Ignore, we don't always have new tasks to run
                    pass

                if not body:
                    continue

                task_queue.task_done()

                bulk_tds   = list()

                for msg in body:

                    task = Task()
                    task.from_dict(msg)

                    load_placeholder(task)
                    bulk_tds.append(create_td_from_task(
                                            task, placeholders,
                                            self._submitted_tasks, pickle_file,
                                            self._prof))

                    self._advance(task, 'Task', states.SUBMITTING,
                                  mq_channel, rmq_conn_params,
                                  '%s-tmgr-to-sync' % self._sid)

                rp_tmgr.submit_tasks(bulk_tds)
            mq_connection.close()
            self._log.debug('Exited RTS main loop. TMGR terminating')
        except KeyboardInterrupt as ex:
            self._log.exception('Execution interrupted (probably by Ctrl+C), '
                                'cancel task processor gracefully...')
            raise KeyboardInterrupt from ex
        except Exception as ex:
            self._log.exception('%s failed with %s', self._uid, ex)
            raise EnTKError(ex) from ex

        finally:
            rp_tmgr.close()

    # --------------------------------------------------------------------------
    #
    def start_manager(self):
        """
        **Purpose**: Method to start the tmgr process. The tmgr function
                     is not to be accessed directly. The function is started
                     in a separate thread using this method.
        """
        # pylint: disable=attribute-defined-outside-init, access-member-before-definition
        if self._tmgr_process:
            self._log.warn('tmgr process already running!')
            return

        try:
            # Redeclare the heartbeat queues in case they got deleted because
            # of the task manager failure.
            # If the queues exist this has no effect.

            mq_connection = pika.BlockingConnection(self._rmq_conn_params)
            mq_channel = mq_connection.channel()

            mq_channel.queue_declare(queue=self._hb_response_q)
            mq_channel.queue_declare(queue=self._hb_request_q)

            self._prof.prof('creating tmgr process', uid=self._uid)
            self._tmgr_terminate = mp.Event()

            self._tmgr_process = mp.Process(target=self._tmgr,
                                            name='task-manager',
                                            args=(self._uid,
                                                  self._rmgr,
                                                  self._pending_queue,
                                                  self._completed_queue,
                                                  self._rmq_conn_params)
                                            )

            self._log.info('Starting task manager process')
            self._prof.prof('starting tmgr process', uid=self._uid)

            self._tmgr_process.start()
            self._log.debug('tmgr pid %s' % self._tmgr_process.pid)

            return True

        except Exception as ex:

            self._log.exception('Task manager not started, error: %s', ex)
            self.terminate_manager()
            raise EnTKError(ex) from ex
        # pylint: enable=attribute-defined-outside-init, access-member-before-definition

# ------------------------------------------------------------------------------

<|endoftext|>


<commit_before>#!/usr/bin/env python2

# Use regular sympy sparingly because it is slow
# Every time we explicitly use it, we should consider implementing such a line in C++
from csympy import *
from csympy.lib.csympy_wrapper import *
import sympy

z_norm = symbols('z_norm')
z_conj = symbols('z_conj')
delta  = symbols('delta')
delta_ext = symbols('delta_ext')

def delta_pole(nu, k, l, series):
    if series == 1:
        return 1 - l - 2 * k
    elif series == 2:
	return 1 + nu - k
    else:
	return 1 + l + 2 * nu - 2 * k

def delta_residue(nu, k, l, series):
    if series == 1:
        ret = - ((k * factorial(2 * k) ** 2) / (2 ** (4 * k - 1) * factorial(k) ** 4))
	if l == 0 and nu == 0:
	    # Take l to 0, then nu
	    return ret * 2
	else:
	    return ret * (sympy.rf(l + 2 * nu, 2 * k) / sympy.rf(l + nu, 2 * k))
    elif series == 2:
	return - sympy.rf(nu, k) * sympy.rf(1 - nu, k) * (sympy.rf((nu + l + 1 - k) / 2, k) ** 2 / sympy.rf((nu + l - k) / 2, k) ** 2) * (k / factorial(k) ** 2) * ((nu + l - k) / (nu + l + k))
    else:
	return - (sympy.rf(1 + l - 2 * k, 2 * k) / sympy.rf(1 + nu + l - 2 * k, 2 * k)) * ((k * factorial(2 * k) ** 2) / (2 ** (4 * k - 1) * factorial(k) ** 4))

def get_poles(dim, l, kept_pole_order):
    nu = sympy.Rational(dim, 2) - 1

    k = 1
    ret = []
    while (2 * k) <= kept_pole_order:
        if delta_residue(nu, k, l, 1) != 0:
	    ret.append(delta_pole(nu, k, l, 1))

	if k < nu + l or dim % 2 == 1:
	    if delta_residue(nu, k, l, 2) != 0:
	        ret.append(delta_pole(nu, k, l, 2))

	if k <= (l / 2):
	    if delta_residue(nu, k, l, 3) != 0:
	        ret.append(delta_pole(nu, k, l, 3))

	k += 1

    # This probably won't change anything
    if nu == 0 and l == 0:
        ret.append(-1)

    return ret

def omit_product(poles, special_pole):
    expression = 1
    omitted = False
    for p in poles:
        if p == special_pole and omitted == False:
	    omitted = True
	else:
	    expression *= (delta - p)
    return expression

def leading_block(nu, r, eta, l):
    if nu == 0:
        ret = sympy.chebyshevt(l, eta)
    else:
        ret = factorial(l) * sympy.gegenbauer(l, nu, eta) / sympy.rf(2 * nu, l)
    return ret / (((1 - r ** 2) ** nu) * sqrt((1 + r ** 2) ** 2 - 4 * (r * eta) ** 2))

def meromorphic_block(dim, r, eta, Delta, l_new, l, kept_pole_order, top):
    k = 1
    nu = sympy.Rational(dim, 2) - 1
    # When the recursion relation shifts l, this does not affect the appropriate poles and
    # residues to use which are still determined by the original spin.
    summation = leading_block(nu, r, eta, l_new)

    # Top says we have not recursed yet and our expression is still expected to have denominators
    # with the free variable delta. Cancelling them later is slow so we do it now.
    if top == True:
        poles = get_poles(dim, l, kept_pole_order)
	for p in poles:
	    summation *= (delta - p)

    while (2 * k) <= kept_pole_order:
        res = delta_residue(nu, k, l, 1)
	if res != 0:
	    pole = delta_pole(nu, k, l, 1)
	    if top == True:
	        new_term = res * (r ** (2 * k)) * omit_product(poles, pole)
	    else:
	        new_term = res * (r ** (2 * k)) / (Delta - pole)
	    summation += new_term * meromorphic_block(dim, r, eta, pole + 2 * k, l + 2 * k, l, kept_pole_order - 2 * k, False)

	if k < nu + l or dim % 2 == 1:
	    res = delta_residue(nu, k, l, 2)
	    if res != 0:
	        pole = delta_pole(nu, k, l, 2)
		if top == True:
	            new_term = res * (r ** (2 * k)) * omit_product(poles, pole)
	        else:
	            new_term = res * (r ** (2 * k)) / (Delta - pole)
	        summation += new_term * meromorphic_block(dim, r, eta, pole + 2 * k, l, l, kept_pole_order - 2 * k, False)

	if k <= (l / 2):
	    res = delta_residue(nu, k, l, 3)
	    if res != 0:
	        pole = delta_pole(nu, k, l, 3)
	        if top == True:
	            new_term = res * (r ** (2 * k)) * omit_product(poles, pole)
	        else:
	            new_term = res * (r ** (2 * k)) / (Delta - pole)
	        summation += new_term * meromorphic_block(dim, r, eta, pole + 2 * k, l - 2 * k, l, kept_pole_order - 2 * k, False)

	k += 1
    return summation

def conformal_block(dim, rho_norm, rho_conj, Delta, l, kept_pole_order):
    r = sqrt(rho_norm * rho_conj)
    eta = (rho_norm + rho_conj) / (2 * r)
    return (r ** Delta) * meromorphic_block(dim, r, eta, Delta, l, l, kept_pole_order, True)

class ConformalBlockTable:
    def __init__(self, dim, derivative_order, kept_pole_order, l_max, odd_spins = False):
	self.dim = dim
	self.derivative_order = derivative_order
	self.kept_pole_order = kept_pole_order
	self.l_max = l_max
	self.odd_spins = odd_spins
	self.m_order = []
	self.n_order = []
	self.table = []

	if odd_spins:
	    step = 1
	else:
	    step = 2

	z = symbols('z')
	rho_n = function_symbol('rho_n', z_norm)
	rho_c = function_symbol('rho_c', z_conj)
	rho = z / (1 + sqrt(1 - z)) ** 2

	rules = []
	for i in range(0, derivative_order + 1):
	    rules.append(rho.subs(z, 0.5))
	    rho = rho.diff(z)

	# We cache derivatives as we go
	# This is because csympy can only compute them one at a time, but it's faster anyway
	for l in range(0, l_max + 1, step):
	    derivatives = []
	    expressions = []

	    # Multiplying a list twice is weird in Python
	    # If this were a numpy array, it would have to store floats
	    for i in range(0, derivative_order + 1):
	        expressions.append([0] * (derivative_order + 1))

	    expressions[0][0] = conformal_block(dim, rho_n, rho_c, delta, l, kept_pole_order)
	    for m in range(0, derivative_order + 1):
		for n in range(0, min(m, derivative_order - m) + 1):
		    if n == 0 and m != 0:
		        expressions[m][0] = expressions[m - 1][0].diff(z_norm)
		    elif n > 0:
		        expressions[m][n] = expressions[m][n - 1].diff(z_conj)

		    # We need to expand before we evaluate to have the essential singularity cancel
		    deriv = expressions[m][n] * (sqrt(rho_n * rho_c) ** (-delta))
		    deriv = deriv.expand()

		    for i in range(m, 0, -1):
		        deriv = deriv.subs(Derivative(rho_n, [z_norm] * i), rules[i])
		    for j in range(n, 0, -1):
		        deriv = deriv.subs(Derivative(rho_c, [z_conj] * j), rules[j])
		    deriv = deriv.subs(rho_n, rules[0]).subs(rho_c, rules[0])

		    derivatives.append(deriv)
		    # For the 27th element of the list, say what m derivative and what n derivative it corresponds to
		    if l == 0:
		        self.m_order.append(m)
		        self.n_order.append(n)
	    self.table.append(derivatives)

class ConvolvedBlockTable:
    def __init__(self, block_table):
        # Copying everything but the unconvolved table is fine from a memory standpoint
        self.dim = block_table.dim
	self.derivative_order = block_table.derivative_order
	self.kept_pole_order = block_table.kept_pole_order
	self.l_max = block_table.l_max
	self.odd_spins = block_table.odd_spins
	self.table = []

	g = function_symbol('g', z_norm, z_conj)
	f = (((1 - z_norm) * (1 - z_conj)) ** delta_ext) * g

	# Same comments apply here
	for l in range(0, len(block_table.table)):
	    derivatives = []
	    expressions = []

	    for i in range(0, block_table.derivative_order + 1):
	        expressions.append([0] * (block_table.derivative_order + 1))

	    expressions[0][0] = f
	    for m in range(0, block_table.derivative_order + 1):
		for n in range(0, min(m, block_table.derivative_order - m) + 1):
		    if n == 0 and m != 0:
		        expressions[m][0] = expressions[m - 1][0].diff(z_norm)
		    elif n > 0:
		        expressions[m][n] = expressions[m][n - 1].diff(z_conj)

		    # Now we replace abstract derivatives with the expressions in the unconvolved table
		    deriv = expressions[m][n] / (factorial(m) * factorial(n))
		    for i in range(len(block_table.table[0]) - 1, 0, -1):
		        deriv = deriv.subs(Derivative(g, [z_norm] * block_table.m_order[i] + [z_conj] * block_table.n_order[i]), block_table.table[l][i])

		    deriv = deriv.subs(g, block_table.table[l][0])
		    derivatives.append(deriv.subs({z_norm : 0.5, z_conj : 0.5}))
	    self.table.append(derivatives)
<commit_msg>More help to performance<commit_after>#!/usr/bin/env python2

# Use regular sympy sparingly because it is slow
# Every time we explicitly use it, we should consider implementing such a line in C++
from csympy import *
from csympy.lib.csympy_wrapper import *
import sympy

z_norm = symbols('z_norm')
z_conj = symbols('z_conj')
delta  = symbols('delta')
delta_ext = symbols('delta_ext')

def delta_pole(nu, k, l, series):
    if series == 1:
        return 1 - l - 2 * k
    elif series == 2:
	return 1 + nu - k
    else:
	return 1 + l + 2 * nu - 2 * k

def delta_residue(nu, k, l, series):
    if series == 1:
        ret = - ((k * factorial(2 * k) ** 2) / (2 ** (4 * k - 1) * factorial(k) ** 4))
	if l == 0 and nu == 0:
	    # Take l to 0, then nu
	    return ret * 2
	else:
	    return ret * (sympy.rf(l + 2 * nu, 2 * k) / sympy.rf(l + nu, 2 * k))
    elif series == 2:
	return - sympy.rf(nu, k) * sympy.rf(1 - nu, k) * (sympy.rf((nu + l + 1 - k) / 2, k) ** 2 / sympy.rf((nu + l - k) / 2, k) ** 2) * (k / factorial(k) ** 2) * ((nu + l - k) / (nu + l + k))
    else:
	return - (sympy.rf(1 + l - 2 * k, 2 * k) / sympy.rf(1 + nu + l - 2 * k, 2 * k)) * ((k * factorial(2 * k) ** 2) / (2 ** (4 * k - 1) * factorial(k) ** 4))

def get_poles(dim, l, kept_pole_order):
    nu = sympy.Rational(dim, 2) - 1

    k = 1
    ret = []
    while (2 * k) <= kept_pole_order:
        if delta_residue(nu, k, l, 1) != 0:
	    ret.append(delta_pole(nu, k, l, 1))

	if k < nu + l or dim % 2 == 1:
	    if delta_residue(nu, k, l, 2) != 0:
	        ret.append(delta_pole(nu, k, l, 2))

	if k <= (l / 2):
	    if delta_residue(nu, k, l, 3) != 0:
	        ret.append(delta_pole(nu, k, l, 3))

	k += 1

    # This probably won't change anything
    if nu == 0 and l == 0:
        ret.append(-1)

    return ret

def omit_product(poles, special_pole):
    expression = 1
    omitted = False
    for p in poles:
        if p == special_pole and omitted == False:
	    omitted = True
	else:
	    expression *= (delta - p)
    return expression

def leading_block(nu, rho_norm, rho_conj, l):
    if nu == 0:
        ret = sympy.chebyshevt(l, (rho_norm + rho_conj) / (2 * sqrt(rho_norm * rho_conj)))
    else:
        ret = factorial(l) * sympy.gegenbauer(l, nu, (rho_norm + rho_conj) / (2 * sqrt(rho_norm * rho_conj))) / sympy.rf(2 * nu, l)
    return ret / (((1 - (rho_norm * rho_conj)) ** nu) * sqrt((1 + (rho_norm * rho_conj)) ** 2 - (rho_norm + rho_conj) ** 2))

def meromorphic_block(dim, rho_norm, rho_conj, Delta, l_new, l, kept_pole_order, top):
    k = 1
    nu = sympy.Rational(dim, 2) - 1
    # When the recursion relation shifts l, this does not affect the appropriate poles and
    # residues to use which are still determined by the original spin.
    summation = leading_block(nu, rho_norm, rho_conj, l_new)

    # Top says we have not recursed yet and our expression is still expected to have denominators
    # with the free variable delta. Cancelling them later is slow so we do it now.
    if top == True:
        poles = get_poles(dim, l, kept_pole_order)
	for p in poles:
	    summation *= (delta - p)

    while (2 * k) <= kept_pole_order:
        res = delta_residue(nu, k, l, 1)
	if res != 0:
	    pole = delta_pole(nu, k, l, 1)
	    if top == True:
	        new_term = res * ((rho_norm * rho_conj) ** k) * omit_product(poles, pole)
	    else:
	        new_term = res * ((rho_norm * rho_conj) ** k) / (Delta - pole)
	    summation += new_term * meromorphic_block(dim, rho_norm, rho_conj, pole + 2 * k, l + 2 * k, l, kept_pole_order - 2 * k, False)

	if k < nu + l or dim % 2 == 1:
	    res = delta_residue(nu, k, l, 2)
	    if res != 0:
	        pole = delta_pole(nu, k, l, 2)
		if top == True:
	            new_term = res * ((rho_norm * rho_conj) ** k) * omit_product(poles, pole)
	        else:
	            new_term = res * ((rho_norm * rho_conj) ** k) / (Delta - pole)
	        summation += new_term * meromorphic_block(dim, rho_norm, rho_conj, pole + 2 * k, l, l, kept_pole_order - 2 * k, False)

	if k <= (l / 2):
	    res = delta_residue(nu, k, l, 3)
	    if res != 0:
	        pole = delta_pole(nu, k, l, 3)
	        if top == True:
	            new_term = res * ((rho_norm * rho_conj) ** k) * omit_product(poles, pole)
	        else:
	            new_term = res * ((rho_norm * rho_conj) ** k) / (Delta - pole)
	        summation += new_term * meromorphic_block(dim, rho_norm, rho_conj, pole + 2 * k, l - 2 * k, l, kept_pole_order - 2 * k, False)

	k += 1
    return summation

def conformal_block(dim, rho_norm, rho_conj, Delta, l, kept_pole_order):
    return ((rho_norm * rho_conj) ** (Delta / 2)) * meromorphic_block(dim, rho_norm, rho_conj, Delta, l, l, kept_pole_order, True)

class ConformalBlockTable:
    def __init__(self, dim, derivative_order, kept_pole_order, l_max, odd_spins = False):
	self.dim = dim
	self.derivative_order = derivative_order
	self.kept_pole_order = kept_pole_order
	self.l_max = l_max
	self.odd_spins = odd_spins
	self.m_order = []
	self.n_order = []
	self.table = []

	if odd_spins:
	    step = 1
	else:
	    step = 2

	z = symbols('z')
	rho_n = function_symbol('rho_n', z_norm)
	rho_c = function_symbol('rho_c', z_conj)
	rho = z / (1 + sqrt(1 - z)) ** 2

	rules = []
	for i in range(0, derivative_order + 1):
	    rules.append(rho.subs(z, 0.5))
	    rho = rho.diff(z)

	# We cache derivatives as we go
	# This is because csympy can only compute them one at a time, but it's faster anyway
	for l in range(0, l_max + 1, step):
	    derivatives = []
	    old_expression = conformal_block(dim, rho_n, rho_c, delta, l, kept_pole_order)

	    for m in range(0, derivative_order + 1):
		for n in range(0, min(m, derivative_order - m) + 1):
		    # Each loop has one expansion that it can do without
		    if n == 0 and m == 0:
		        expression = old_expression
		    elif n == 0:
		        old_expression = old_expression.diff(z_norm).expand()
			expression = old_expression
		    else:
		        expression = expression.diff(z_conj).expand()

		    deriv = expression * ((rho_n * rho_c) ** (-delta / 2))
		    for i in range(m, 0, -1):
		        deriv = deriv.subs(Derivative(rho_n, [z_norm] * i), rules[i])
		    for j in range(n, 0, -1):
		        deriv = deriv.subs(Derivative(rho_c, [z_conj] * j), rules[j])
		    deriv = deriv.subs(rho_n, rules[0]).subs(rho_c, rules[0])

		    derivatives.append(deriv)
		    # For the 27th element of the list, say what m derivative and what n derivative it corresponds to
		    if l == 0:
		        self.m_order.append(m)
		        self.n_order.append(n)
	    self.table.append(derivatives)

class SDP:
    def __init__(self, block_table):
        # Copying everything but the unconvolved table is fine from a memory standpoint
        self.dim = block_table.dim
	self.derivative_order = block_table.derivative_order
	self.kept_pole_order = block_table.kept_pole_order
	self.l_max = block_table.l_max
	self.odd_spins = block_table.odd_spins
	self.table = []

	g = function_symbol('g', z_norm, z_conj)
	f = (((1 - z_norm) * (1 - z_conj)) ** delta_ext) * g

	# Same comments apply here
	for l in range(0, len(block_table.table)):
	    derivatives = []
	    old_expression = f

	    for m in range(0, block_table.derivative_order + 1):
		for n in range(0, min(m, block_table.derivative_order - m) + 1):
		    if n == 0 and m == 0:
		        expression = old_expression
		    elif n == 0:
		        old_expression = old_expression.diff(z_norm).expand()
			expression = old_expression
		    else:
		        expression = expression.diff(z_conj).expand()

		    # Now we replace abstract derivatives with the expressions in the unconvolved table
		    deriv = expression / (factorial(m) * factorial(n))
		    for i in range(len(block_table.table[0]) - 1, 0, -1):
		        deriv = deriv.subs(Derivative(g, [z_norm] * block_table.m_order[i] + [z_conj] * block_table.n_order[i]), block_table.table[l][i])

		    deriv = deriv.subs(g, block_table.table[l][0])
		    derivatives.append(deriv.subs({z_norm : 0.5, z_conj : 0.5}))
	    self.table.append(derivatives)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import os
import shutil
import unittest
import tempfile

from flask_reveal.app import FlaskReveal
from flask_reveal.blueprints.reveal import load_markdown_slides

class BlueprintTestCase(unittest.TestCase):
    slides = ['Slide1', 'Slide2', 'Slide3']

    def create_presentation_structure(self, slides=None):
        root = tempfile.mkdtemp()
        media = tempfile.mkdtemp()

        fd_cfg, config = tempfile.mkstemp('.py', dir=root)
        fd_img, image = tempfile.mkstemp('.jpg', dir=media)

        os.close(fd_cfg)
        os.close(fd_img)

        for index, slide in enumerate(slides):
            fd, _ = tempfile.mkstemp('.md', str(index), root)
            with os.fdopen(fd, 'w') as file:
                file.write(slide)

        return dict(root=root,
                    media=media,
                    config=config,
                    image=os.path.basename(image))

    def create_test_client(self, presentation_root, media_root, config):
        app = FlaskReveal('flask_reveal')

        app.load_user_config(presentation_root, media_root, config)
        app.config['TESTING'] = True

        return app.test_client()

    def setUp(self):
        self.presentation = self.create_presentation_structure(self.slides)

    def tearDown(self):
        shutil.rmtree(self.presentation['root'])

    def test_presentation_view_status(self):
        client = self.create_test_client(self.presentation['root'],
                                         self.presentation['media'],
                                         self.presentation['config'])

        with client.get('/') as response:
            self.assertEqual(response.status, '200 OK')

    def test_get_img_view_status(self):
        client = self.create_test_client(self.presentation['root'],
                                         self.presentation['media'],
                                         self.presentation['config'])
        url = '/img/{0}'.format(self.presentation['image'])

        with client.get(url) as response:
            self.assertEqual(response.status, '200 OK')

    def test_load_markdown_slides(self):
        slides = load_markdown_slides(self.presentation['root'])

        self.assertEqual(slides, self.slides)
<commit_msg>renaming make_test_client to make_client<commit_after># -*- coding: utf-8 -*-
import os
import shutil
import unittest
import tempfile

from flask_reveal.app import FlaskReveal
from flask_reveal.blueprints.reveal import load_markdown_slides

class BlueprintTestCase(unittest.TestCase):
    slides = ['Slide1', 'Slide2', 'Slide3']

    def create_presentation_structure(self, slides=None):
        root = tempfile.mkdtemp()
        media = tempfile.mkdtemp()

        fd_cfg, config = tempfile.mkstemp('.py', dir=root)
        fd_img, image = tempfile.mkstemp('.jpg', dir=media)

        os.close(fd_cfg)
        os.close(fd_img)

        for index, slide in enumerate(slides):
            fd, _ = tempfile.mkstemp('.md', str(index), root)
            with os.fdopen(fd, 'w') as file:
                file.write(slide)

        return dict(root=root,
                    media=media,
                    config=config,
                    image=os.path.basename(image))

    def create_client(self, presentation_root, media_root, config):
        app = FlaskReveal('flask_reveal')

        app.load_user_config(presentation_root, media_root, config)
        app.config['TESTING'] = True

        return app.test_client()

    def setUp(self):
        self.presentation = self.create_presentation_structure(self.slides)

    def tearDown(self):
        shutil.rmtree(self.presentation['root'])

    def test_presentation_view_status(self):
        client = self.create_client(self.presentation['root'],
                                    self.presentation['media'],
                                    self.presentation['config'])

        with client.get('/') as response:
            self.assertEqual(response.status, '200 OK')

    def test_get_img_view_status(self):
        client = self.create_client(self.presentation['root'],
                                    self.presentation['media'],
                                    self.presentation['config'])
        url = '/img/{0}'.format(self.presentation['image'])

        with client.get(url) as response:
            self.assertEqual(response.status, '200 OK')

    def test_load_markdown_slides(self):
        slides = load_markdown_slides(self.presentation['root'])

        self.assertEqual(slides, self.slides)
<|endoftext|>