buckets_yo

buckets.py

from itertools import permutations
import networkx as nx
import numpy as np
import matplotlib.pyplot as plt
import pylab

class BucketGraph:
    def __init__(self, buckets):
        self.buckets = buckets
        self.g = nx.DiGraph()
        self.vertices = self._generate_vertices(buckets)
        self.edges = self._generate_edges(self.vertices)
        pos = nx.random_layout(self.g)
        # nx.draw_networkx_nodes(self.g, pos)
        nx.draw_networkx(self.g, pos, node_size=1500)
        pylab.show()

    def _generate_vertices(self, buckets):
        self.volume = sum([bucket[1] for bucket in buckets])
        num_buckets = len(buckets)
        vertices = []

        def vertex_filter(v):
            for (i, b) in enumerate(v):
                if b > buckets[i][0]:
                    return False
            return True


        for (i, s) in enumerate(self._partition_lt(self.volume, len(buckets))):
            vertices.extend(
                set(filter(
                    vertex_filter,
                    [p for p in
                     permutations((s + [0] *
                                   num_buckets)[:num_buckets])]
                ))
            )

        return vertices

    def _generate_edges(self, vertices):
        def valid_transition(start, finish):
            def pour(vertex, source, dest):
                new_dest = vertex[dest] + vertex[source]
                destination_bucket_size = self.buckets[dest][0]
                new_source = 0
                if new_dest > destination_bucket_size:
                    new_source = new_dest - destination_bucket_size
                    new_dest = destination_bucket_size
                temp = [i for i in vertex]
                temp[source] = new_source
                temp[dest] = new_dest
                return tuple(temp)

            def transitions(vertex):
                t = []
                for (i, bucket) in enumerate(vertex):
                    if bucket == 0: continue # can't pour an empty bucket
                    for (j, other) in enumerate(vertex):
                        if i == j: continue # can't pour into self
                        if other == self.buckets[i][0]: continue # already full
                        t.append(pour(vertex, i, j))
                return t

            t = transitions(start)
            return finish in t

        # Now we have all of the possible partitions of the total volume
        # of the system, which means we can form a graph by connecting
        # each vertex to every other one as long as its allowed by the
        # rules given

        edges = []
        for (i, u) in enumerate(vertices):
            for (j, v) in enumerate(vertices):
                if valid_transition(u, v):
                    edges.append((i, j)) # vertex i --> j is allowed
                    self.g.add_edge(str(u), str(v))


    def _partition_lt(self, n, k):
        """
        Partition `n` into k, k-1, ... 1 parts
        """
        def ascending_filter(l):
            if len(l) > k:
                return False
            return True

        return filter(ascending_filter, self.ruleAsc(n))

    def ruleAsc(self, n):
        a = [0 for i in range(n + 1)]
        k = 1
        a[1] = n
        while k != 0:
            x = a[k - 1] + 1
            y = a[k] - 1
            k -= 1
            while x <= y:
                a[k] = x
                y -= x
                k += 1
            a[k] = x + y
            yield a[:k + 1]

def main():
    g = BucketGraph(zip([10, 4, 7],
                        [0,  4, 7]))


if __name__ == "__main__":
    main()