<commit_before>'''
Implementation of graph.

Author: mbforbes

Current approach (order matters):
-   (1) add RVs
-   (2) add factors to connect them
-   (3) set beliefs on factors
-   (4) run inference
-   (5) compute marginals

For some things below, we'll want to represent what's going on in mathematical
notation. Let's define some variables that we'll use throughout to help:

RV vars:
    X       the set of n random variables
    X_i     random variable i (1 <= i <= n)
    v_i     number of values that X_i can take (nonstandard but I wanted one)
    x_ij    a particular value for X_i (1 <= j <= v_i)
    x_i     a simpler (lazy) notation for x_ij (which j doesn't matter)
    x       a set of x_i for i = 1..n (imagine x_1j, x_2k, ..., x_nz)

Factor vars:
    F       the set of m factors
    f_a     factor a (1 <= a <= m) connecting a subset of X
    X_a     the subset of X (RVs) that f_a connects
    x_a     the subset of x (values for RVs) that f_a connects

Functions:
    p(x)    joint distribution for p(X = x)

Notes:
    f_a(x) = f_a(x_a)   Because f_a only touches (is only a function of) x_a,
                        it will "ignore" the other x_i in x that aren't in x_a.
                        Thus, we write f_a(x_a) for convenience to show exactly
                        what f_a operates on.
'''

# Imports
# -----------------------------------------------------------------------------

# 3rd party
import numpy as np

# Constants
# -----------------------------------------------------------------------------

# Settings

# Use this to turn all debugging on or off. Intended use: keep on when you're
# trying stuff out. Once you know stuff works, turn off for speed. Can also
# specify when creating each instance, but this global switch is provided for
# convenience.
DEBUG_DEFAULT = True

# Classes
# -----------------------------------------------------------------------------

class Graph(object):
    '''
    Graph right now has no point, really (except bookkeeping all the RVs and
    factors, assuming we remember to add them), so this might be removed or
    functionality might be stuffed in here later.
    '''

    def __init__(self, debug=DEBUG_DEFAULT):
        # add now
        self.debug = debug

        # added later
        self._rvs = {}
        self._factors = []  # TODO: Better data structure?

    def add_rv(self, rv):
        '''
        Node (RV|Factor)
        '''
        # Check RV with same name not already added.
        if self.debug:
            assert rv.name not in self._rvs
        # Add it.
        self._rvs[rv.name] = rv

    def add_factor(self, factor):
        if self.debug:
            # Check the same factor hasn't already been added.
            assert factor not in self._factors

            # Check factor connecting to exactly the same set of nodes doesn't
            # already exist. This isn't mandated by factor graphs by any means,
            # but it's a heuristic to prevent bugs; if you're adding factors
            # that connect the same set of ndoes, you're either doing something
            # weird (and can probably reformulate your graph structure to avoid
            # this duplication), or you have a bug.
            factor_rvs = sorted(factor._rvs)
            for f in self._factors:
                rvs = sorted(f._rvs)
                assert factor_rvs != rvs
        # Add it.
        self._factors += [factor]

    def joint(self, x):
        '''
        Joint is over the factors.

        For a probability, we use the normalization constant 1/Z

            p(x) = 1/Z \product_a^{1..m} f_a(x_a)

        If we don't care what the normalization is, we just write this without
        1/Z:

            p(x) = \product_a^{1..m} f_a(x_a)

        This is currently implemented without normalization. I might want to
        add it in the future. I don't know yet.

        Args:
            x ({str: str|int}) map of node names to assignments. The
                assignments can be labels or indexes
        '''
        # ensure the assignment x given is valid
        if self.debug:
            # check the length (that assignments to all RVs are provided)
            assert len(x) == len(self._rvs)

            # check that each assignment is valid (->)
            for name, label in x.iteritems():
                assert name in self._rvs
                assert self._rvs[name].has_label(label)

            # check that each RV has a valid assignment (<-)
            for name, rv in self._rvs.iteritems():
                assert name in x
                assert rv.has_label(x[name])

        # Do the actual computation.
        # NOTE: This could be sped up as all factors can be computed in
        # parallel.
        prod = 1.0
        for f in self._factors:
            prod *= f.eval(x)
        return prod

    def bf_best_joint(self):
        '''
        Brute-force algorithm to compute the best joint assignment to the
        factor graph given the current beliefs in the factors.

        This takes O(v^n) time, where
            v   is the # of possible assignments to each RV
            n   is the # of RVs

        This is bad. This function is given for debugging / proof of concept
        only.

        Returns:
            ({str: int}, float)
        '''
        return self._bf_bj_recurse({}, self._rvs.values())

    def _bf_bj_recurse(self, assigned, todo):
        '''
        Helper method for bf_best_joint.

        Args:
            assigned ({str: int})
            todo ([RV])
        '''
        # base case: just look up the current assignment
        if len(todo) == 0:
            return assigned, self.joint(assigned)

        # recursive case: pull off one RV and try all options.
        best_a, best_r = None, 0.0
        rv = todo[0]
        todo = todo[1:]
        for val in range(rv.n_opts):
            new_a = assigned.copy()
            new_a[rv.name] = val
            full_a, r = self._bf_bj_recurse(new_a, todo)
            if r > best_r:
                best_r = r
                best_a = full_a
        return best_a, best_r

    def print_stats(self):
        print 'Graph stats:'
        print '\t%d RVs' % (len(self._rvs))
        print '\t%d factors' % (len(self._factors))

class RV(object):

    def __init__(self, name, n_opts, labels=[], debug=DEBUG_DEFAULT):
        '''
        name (str)
        n_opts (int)              how many values it can take
        labels ([str], opt)
        debug (bool, opt)
        '''
        # validation
        if debug:
            # labels must be [str] if provided
            for l in labels:
                assert type(l) is str

            # must have n_opts labels if provided
            assert len(labels) == 0 or len(labels) == n_opts

        # vars set at construction time
        self.name = name
        self.n_opts = n_opts
        self.labels = labels
        self.debug = debug

        # vars added later
        self._factors = []

    def has_label(self, label):
        '''
        Returns whether label indicates a valid value for this RV.

        Args:
            label (int|str)

        returns
            bool
        '''
        # If int, make sure fits in n_opts. If str, make sure it's in the list.
        if len(self.labels) == 0:
            # Tracking ints only. Provided label must be int.
            if self.debug:
                assert type(label) is int
            return label < self.n_opts
        else:
            # Tracking strs only. Provided label can be int or str.
            if self.debug:
                assert type(label) is int or type(label) is str
            if type(label) is str:
                return label in self.labels
            # Default: int
            return label < self.n_opts

    def get_int_label(self, label):
        '''
        Returns the integer-valued label for this label. The provided label
        might be an integer (in which case it's already in the correct form and
        will be returned unchanged) or a string (in which case it will be
        turned into an int).

        This assumes the caller has already ensured this is a valid label with
        has_label.

        Args:
            label (int|str)

        returns
            int
        '''
        if type(label) is int:
            return label
        # assume string otherwise
        return self.labels.index(label)

    def attach(self, factor):
        '''
        Don't call this; automatically called by Factor's attach(...). This
        doesn't update the factor's attachment (which is why you shouldn't call
        it).

        factor (Factor)
        '''
        # check whether factor already added to rv; reach factor should be
        # added at most once to an rv.
        if self.debug:
            for f in self._factors:
                # We really only need to worry about the exact instance here,
                # so just using the builtin object (mem address) equals.
                assert f != factor, ('Can\'t re-add factor %r to rv %r' %
                                     (factor, self))

        # Do the registration
        self._factors += [factor]

class Factor(object):

    def __init__(self, rvs, name='', debug=DEBUG_DEFAULT):
        '''
        rvs ([RV])
        name (str, opt)
        debug (bool, opt)
        '''
        # at construction time
        self.name = name
        self.debug = debug

        # add later using methods
        self._rvs = []
        self._belief = None

        # set the rvs now
        for rv in rvs:
            self.attach(rv)

    def attach(self, rv):
        '''
        Call this to attach this factor to the RV rv. Clears any belief that
        has been set.

        rv (RV)
        '''
        # check whether rv already added to factor; reach rv should be added at
        # most once to a factor.
        if self.debug:
            for r in self._rvs:
                # We really only need to worry about the exact instance here,
                # so just using the builtin object (mem address) equals.
                assert r != rv, 'Can\'t re-add RV %r to factor %r' % (rv, self)

        # register with rv
        rv.attach(self)

        # register rv here
        self._rvs += [rv]

        # Clear belief as dimensions no longer match.
        self._belief = None

    def set_belief(self, b):
        '''
        Call this to set the belief for a factor. The passed belief b must
        dimensionally match all attached RVs.

        The dimensions can be a bit confusing. They iterate through the
        dimensions of the RVs in order.

        For example, say we have three RVs, which can each take on the
        following values:

            A {a, b, c}
            B {d, e}
            C {f, g}

        Now, say we have a facor which connects all of them (i.e. f(A,B,C)).
        The dimensions of the belief for this factor are 3 x 2 x 2. You can
        imagine a 3d table of numbers:

                        a b c
            a b c     +------
          + -----   d | g g g
        d | f f f / e | g g g
        e | f f f /

        This looks like you have two "sheets" of numbers. The lower sheet (on
        the left) contains the values for C = f, and the upper sheet (on the
        right) contains the values for C = g. A single cell contains the joint.
        For example, the top-left cell of the bottom sheet contains the value
        for f(A=a, B=d, C=f), and the middle-bottom cell of the top sheet
        contains the value for f(A=b, B=e, c=g).

        The confusing thing (for me) is that a single belief of shape (3, 2, 2)
        is represented in numpy as the following array:

           [[[n, n],
             [n, n]],

            [[n, n],
             [n, n]],

            [[n, n],
             [n, n]]]

        Though this still has twelve numbers, it wasn't how I was
        conceptualizing it. What gives? Well, what we're doing is indexing in
        the correct order. So, the first dimension, 3, indexes the value for
        the variable A. This visually splits our table into three areas, one
        each for A=a, A=b, and A=c. For each area, we have a 2 x 2 table. These
        would be represented in our 3d diagram above by three vertial sheets.
        Each 2 x 2 table has the values for B and C.

        So it really turns out I'd drawn my first table wrong for thinking
        about numpy arrays. You want to draw them by splitting up tables by the
        earlier Rvs. This would look like:

        A = a:
                d e
              +----
            d | n n
            e | n n

        A = b:
                f g
              +----
            d | n n
            e | n n

        A = c:
                f g
              +----
            d | n n
            e | n n

        Args:
            b (np.array)
        '''
        # check that the new belief has the correct shape
        if self.debug:
            # ensure overall dims match
            got = len(b.shape)
            want = len(self._rvs)
            assert got == want, ('Belief %r has %d dims but needs %d' %
                                 (b, got, want))

            # Ensure each dim matches.
            for i, d in enumerate(b.shape):
                got = d
                want = self._rvs[i].n_opts
                assert got == want, (
                    'Belief %r dim #%d has %d opts but rv has %d opts' %
                    (b, i+1, got, want))

        # Set it
        self._belief = b

    def eval(self, x):
        '''
        If this factor is f_a, x may have assignments to more than x_a; that
        is, it might be the full x.

        This checks (if debug is on) that all attached RVs have a valid
        assignment in x. Note that if this is begin called from Graph.joint(),
        this property is also checked there.
        '''
        if self.debug:
            # check that each RV has a valid assignment (<-)
            for rv in self._rvs:
                assert rv.name in x
                assert rv.has_label(x[rv.name])

        # slim down belief into desired value.
        ret = self._belief
        for r in self._rvs:
            ret = ret[r.get_int_label(x[r.name])]

        # should return a single number
        if self.debug:
            assert type(ret) is not np.ndarray

        return ret

# TODO(mbforbes): Remove all the main stuff once you have tests for this.
def main():
    # rvs
    r1 = RV('foo', 2)
    r2 = RV('bar', 3)

    # factors
    f = Factor([r1, r2])
    f_bar = Factor([r2])

    # beliefs
    b = np.array([[3, 2, 0.32453563], [2, 5, 5]])
    f.set_belief(b)
    b2 = np.array([8.0, 2.0, 3.0])
    f_bar.set_belief(b2)

    # add to a graph
    g = Graph()
    g.add_rv(r1)
    g.add_rv(r2)
    g.add_factor(f)
    g.add_factor(f_bar)

    # Print stuff
    g.print_stats()
    print 'Joint for foo=2, bar=1:', g.joint({'foo': 1, 'bar': 1})
    print 'Best joint:', g.bf_best_joint()

if __name__ == '__main__':
    main()
<commit_msg>Add sketch of lbp algo.<commit_after>'''
Implementation of factor graph.

Author: mbforbes

Current approach (order matters):
-   (1) add RVs
-   (2) add factors to connect them
-   (3) set beliefs on factors
-   (4) run inference
-   (5) compute marginals

For some things below, we'll want to represent what's going on in mathematical
notation. Let's define some variables that we'll use throughout to help:

RV vars:
    X       the set of n random variables
    X_i     random variable i (1 <= i <= n)
    v_i     number of values that X_i can take (nonstandard but I wanted one)
    x_ij    a particular value for X_i (1 <= j <= v_i)
    x_i     a simpler (lazy) notation for x_ij (which j doesn't matter)
    x       a set of x_i for i = 1..n (imagine x_1j, x_2k, ..., x_nz)

Factor vars:
    F       the set of m factors
    f_a     factor a (1 <= a <= m) connecting a subset of X
    X_a     the subset of X (RVs) that f_a connects
    x_a     the subset of x (values for RVs) that f_a connects

Functions:
    p(x)    joint distribution for p(X = x)

Notes:
    f_a(x) = f_a(x_a)   Because f_a only touches (is only a function of) x_a,
                        it will "ignore" the other x_i in x that aren't in x_a.
                        Thus, we write f_a(x_a) for convenience to show exactly
                        what f_a operates on.
'''

# Imports
# -----------------------------------------------------------------------------

# 3rd party
import numpy as np

# Constants
# -----------------------------------------------------------------------------

# Settings

# Use this to turn all debugging on or off. Intended use: keep on when you're
# trying stuff out. Once you know stuff works, turn off for speed. Can also
# specify when creating each instance, but this global switch is provided for
# convenience.
DEBUG_DEFAULT = True

# Classes
# -----------------------------------------------------------------------------

class Graph(object):
    '''
    Graph right now has no point, really (except bookkeeping all the RVs and
    factors, assuming we remember to add them), so this might be removed or
    functionality might be stuffed in here later.
    '''

    def __init__(self, debug=DEBUG_DEFAULT):
        # add now
        self.debug = debug

        # added later
        self._rvs = {}
        self._factors = []  # TODO: Better data structure?

    def add_rv(self, rv):
        '''
        Node (RV|Factor)
        '''
        # Check RV with same name not already added.
        if self.debug:
            assert rv.name not in self._rvs
        # Add it.
        self._rvs[rv.name] = rv

    def add_factor(self, factor):
        if self.debug:
            # Check the same factor hasn't already been added.
            assert factor not in self._factors

            # Check factor connecting to exactly the same set of nodes doesn't
            # already exist. This isn't mandated by factor graphs by any means,
            # but it's a heuristic to prevent bugs; if you're adding factors
            # that connect the same set of ndoes, you're either doing something
            # weird (and can probably reformulate your graph structure to avoid
            # this duplication), or you have a bug.
            factor_rvs = sorted(factor._rvs)
            for f in self._factors:
                rvs = sorted(f._rvs)
                assert factor_rvs != rvs
        # Add it.
        self._factors += [factor]

    def joint(self, x):
        '''
        Joint is over the factors.

        For a probability, we use the normalization constant 1/Z

            p(x) = 1/Z \product_a^{1..m} f_a(x_a)

        If we don't care what the normalization is, we just write this without
        1/Z:

            p(x) = \product_a^{1..m} f_a(x_a)

        This is currently implemented without normalization. I might want to
        add it in the future. I don't know yet.

        Args:
            x ({str: str|int}) map of node names to assignments. The
                assignments can be labels or indexes
        '''
        # ensure the assignment x given is valid
        if self.debug:
            # check the length (that assignments to all RVs are provided)
            assert len(x) == len(self._rvs)

            # check that each assignment is valid (->)
            for name, label in x.iteritems():
                assert name in self._rvs
                assert self._rvs[name].has_label(label)

            # check that each RV has a valid assignment (<-)
            for name, rv in self._rvs.iteritems():
                assert name in x
                assert rv.has_label(x[name])

        # Do the actual computation.
        # NOTE: This could be sped up as all factors can be computed in
        # parallel.
        prod = 1.0
        for f in self._factors:
            prod *= f.eval(x)
        return prod

    def bf_best_joint(self):
        '''
        Brute-force algorithm to compute the best joint assignment to the
        factor graph given the current beliefs in the factors.

        This takes O(v^n) time, where
            v   is the # of possible assignments to each RV
            n   is the # of RVs

        This is bad. This function is given for debugging / proof of concept
        only.

        Returns:
            ({str: int}, float)
        '''
        return self._bf_bj_recurse({}, self._rvs.values())

    def _bf_bj_recurse(self, assigned, todo):
        '''
        Helper method for bf_best_joint.

        Args:
            assigned ({str: int})
            todo ([RV])
        '''
        # base case: just look up the current assignment
        if len(todo) == 0:
            return assigned, self.joint(assigned)

        # recursive case: pull off one RV and try all options.
        best_a, best_r = None, 0.0
        rv = todo[0]
        todo = todo[1:]
        for val in range(rv.n_opts):
            new_a = assigned.copy()
            new_a[rv.name] = val
            full_a, r = self._bf_bj_recurse(new_a, todo)
            if r > best_r:
                best_r = r
                best_a = full_a
        return best_a, best_r

    def lbp(self):
        '''
        Loopy belief propagation.
        '''
        # Sketch of algorithm:
        # -------------------
        # preprocessing:
        # - sort nodes (rvs? factors?) by number of edges
        #
        # note:
        # - every time message sent, normalize if too large or small
        #
        # running (loop):
        # - initialize messages to 1
        # - for each node in sorted list (fewest edges to most):
        #     - send messages to neighbors (from which message wasn't received)
        # - for each node in *reverse* sorted list (most edges to fewest):
        #     - send messages to neighbors (from which message wasn't received)
        # - for each node:
        #     - ensure message was sent on every edge
        # - check convergence of messages
        #
        # after finished:
        # - compute marginals for rvs (I think?) and factors
        #
        # questions:
        # - do we have to do all RVs then all Factors? Can we mix up the order?
        #   *Should* we do one or the other?
        #
        # - do factors get updated? or is it just messages?
        pass

    def print_stats(self):
        print 'Graph stats:'
        print '\t%d RVs' % (len(self._rvs))
        print '\t%d factors' % (len(self._factors))

class RV(object):

    def __init__(self, name, n_opts, labels=[], debug=DEBUG_DEFAULT):
        '''
        name (str)
        n_opts (int)              how many values it can take
        labels ([str], opt)
        debug (bool, opt)
        '''
        # validation
        if debug:
            # labels must be [str] if provided
            for l in labels:
                assert type(l) is str

            # must have n_opts labels if provided
            assert len(labels) == 0 or len(labels) == n_opts

        # vars set at construction time
        self.name = name
        self.n_opts = n_opts
        self.labels = labels
        self.debug = debug

        # vars added later
        self._factors = []

    def has_label(self, label):
        '''
        Returns whether label indicates a valid value for this RV.

        Args:
            label (int|str)

        returns
            bool
        '''
        # If int, make sure fits in n_opts. If str, make sure it's in the list.
        if len(self.labels) == 0:
            # Tracking ints only. Provided label must be int.
            if self.debug:
                assert type(label) is int
            return label < self.n_opts
        else:
            # Tracking strs only. Provided label can be int or str.
            if self.debug:
                assert type(label) is int or type(label) is str
            if type(label) is str:
                return label in self.labels
            # Default: int
            return label < self.n_opts

    def get_int_label(self, label):
        '''
        Returns the integer-valued label for this label. The provided label
        might be an integer (in which case it's already in the correct form and
        will be returned unchanged) or a string (in which case it will be
        turned into an int).

        This assumes the caller has already ensured this is a valid label with
        has_label.

        Args:
            label (int|str)

        returns
            int
        '''
        if type(label) is int:
            return label
        # assume string otherwise
        return self.labels.index(label)

    def attach(self, factor):
        '''
        Don't call this; automatically called by Factor's attach(...). This
        doesn't update the factor's attachment (which is why you shouldn't call
        it).

        factor (Factor)
        '''
        # check whether factor already added to rv; reach factor should be
        # added at most once to an rv.
        if self.debug:
            for f in self._factors:
                # We really only need to worry about the exact instance here,
                # so just using the builtin object (mem address) equals.
                assert f != factor, ('Can\'t re-add factor %r to rv %r' %
                                     (factor, self))

        # Do the registration
        self._factors += [factor]

class Factor(object):

    def __init__(self, rvs, name='', debug=DEBUG_DEFAULT):
        '''
        rvs ([RV])
        name (str, opt)
        debug (bool, opt)
        '''
        # at construction time
        self.name = name
        self.debug = debug

        # add later using methods
        self._rvs = []
        self._belief = None

        # set the rvs now
        for rv in rvs:
            self.attach(rv)

    def attach(self, rv):
        '''
        Call this to attach this factor to the RV rv. Clears any belief that
        has been set.

        rv (RV)
        '''
        # check whether rv already added to factor; reach rv should be added at
        # most once to a factor.
        if self.debug:
            for r in self._rvs:
                # We really only need to worry about the exact instance here,
                # so just using the builtin object (mem address) equals.
                assert r != rv, 'Can\'t re-add RV %r to factor %r' % (rv, self)

        # register with rv
        rv.attach(self)

        # register rv here
        self._rvs += [rv]

        # Clear belief as dimensions no longer match.
        self._belief = None

    def set_belief(self, b):
        '''
        Call this to set the belief for a factor. The passed belief b must
        dimensionally match all attached RVs.

        The dimensions can be a bit confusing. They iterate through the
        dimensions of the RVs in order.

        For example, say we have three RVs, which can each take on the
        following values:

            A {a, b, c}
            B {d, e}
            C {f, g}

        Now, say we have a facor which connects all of them (i.e. f(A,B,C)).
        The dimensions of the belief for this factor are 3 x 2 x 2. You can
        imagine a 3d table of numbers:

                        a b c
            a b c     +------
          + -----   d | g g g
        d | f f f / e | g g g
        e | f f f /

        This looks like you have two "sheets" of numbers. The lower sheet (on
        the left) contains the values for C = f, and the upper sheet (on the
        right) contains the values for C = g. A single cell contains the joint.
        For example, the top-left cell of the bottom sheet contains the value
        for f(A=a, B=d, C=f), and the middle-bottom cell of the top sheet
        contains the value for f(A=b, B=e, c=g).

        The confusing thing (for me) is that a single belief of shape (3, 2, 2)
        is represented in numpy as the following array:

           [[[n, n],
             [n, n]],

            [[n, n],
             [n, n]],

            [[n, n],
             [n, n]]]

        Though this still has twelve numbers, it wasn't how I was
        conceptualizing it. What gives? Well, what we're doing is indexing in
        the correct order. So, the first dimension, 3, indexes the value for
        the variable A. This visually splits our table into three areas, one
        each for A=a, A=b, and A=c. For each area, we have a 2 x 2 table. These
        would be represented in our 3d diagram above by three vertial sheets.
        Each 2 x 2 table has the values for B and C.

        So it really turns out I'd drawn my first table wrong for thinking
        about numpy arrays. You want to draw them by splitting up tables by the
        earlier Rvs. This would look like:

        A = a:
                d e
              +----
            d | n n
            e | n n

        A = b:
                f g
              +----
            d | n n
            e | n n

        A = c:
                f g
              +----
            d | n n
            e | n n

        Args:
            b (np.array)
        '''
        # check that the new belief has the correct shape
        if self.debug:
            # ensure overall dims match
            got = len(b.shape)
            want = len(self._rvs)
            assert got == want, ('Belief %r has %d dims but needs %d' %
                                 (b, got, want))

            # Ensure each dim matches.
            for i, d in enumerate(b.shape):
                got = d
                want = self._rvs[i].n_opts
                assert got == want, (
                    'Belief %r dim #%d has %d opts but rv has %d opts' %
                    (b, i+1, got, want))

        # Set it
        self._belief = b

    def eval(self, x):
        '''
        If this factor is f_a, x may have assignments to more than x_a; that
        is, it might be the full x.

        This checks (if debug is on) that all attached RVs have a valid
        assignment in x. Note that if this is begin called from Graph.joint(),
        this property is also checked there.
        '''
        if self.debug:
            # check that each RV has a valid assignment (<-)
            for rv in self._rvs:
                assert rv.name in x
                assert rv.has_label(x[rv.name])

        # slim down belief into desired value.
        ret = self._belief
        for r in self._rvs:
            ret = ret[r.get_int_label(x[r.name])]

        # should return a single number
        if self.debug:
            assert type(ret) is not np.ndarray

        return ret

# TODO(mbforbes): Remove all the main stuff once you have tests for this.
def main():
    # rvs
    r1 = RV('foo', 2)
    r2 = RV('bar', 3)

    # factors
    f = Factor([r1, r2])
    f_bar = Factor([r2])

    # beliefs
    b = np.array([[3, 2, 0.32453563], [2, 5, 5]])
    f.set_belief(b)
    b2 = np.array([8.0, 2.0, 3.0])
    f_bar.set_belief(b2)

    # add to a graph
    g = Graph()
    g.add_rv(r1)
    g.add_rv(r2)
    g.add_factor(f)
    g.add_factor(f_bar)

    # Print stuff
    g.print_stats()
    print 'Joint for foo=2, bar=1:', g.joint({'foo': 1, 'bar': 1})
    print 'Best joint:', g.bf_best_joint()

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


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
Author: Timo Denissen
E-Mail: timo@communtu.com
About this program:
Todo: - Give the possibility to change the key name
      - GUI
'''

import yaml
import os

def file_processing():
    '''Working with the file.'''
    mydict = yaml.safe_load(template_yml)
    template_yml.close()
    new_dict = switch_keyvalue(mydict['de'])
    check_dict(new_dict)

def check_dict(dict):
    '''Check if there are multiple values for one key.'''
    for key, value in dict.iteritems():
        if len(value) > 1:
            print 'Please solve this conflict:'
            multiple_string = value
            print value, ':', key ,'\n'
            ask_string(multiple_string)
        else:
            pass

def ask_string(string):
    '''Ask user which string to use in templaye.yml'''
    temp_dict = {}
    x = 1
    for item in string:
        temp_dict[x] = item
        print x, item
        x += 1
    ask_number = input('Enter the number of the correct key: ')
    newkey = temp_dict[ask_number]
    string.remove(newkey)
    oldkey_temp = string
    edit_files(oldkey_temp, newkey)
#    update_template()

def edit_files(oldkey, newkey):
    '''Using "sed" to remove oldkey from source files.'''
    for item in oldkey:
        oldkey = item
        sed = 'find app lib -name "*rb" -exec sed -i \'s/t(:' + oldkey + '\([^:alnum:_]\)/t(:' + newkey + '\\1/g\' {} \;'
        print 'These files will be changed:\n'
        os.system('grep --include=\'*rb\' "' + oldkey + '" ./app -r')
        if raw_input('\nIs this ok? y/n ') == 'n':
            break
        else:
            pass
        os.system(sed)
        edit_template(item)
    print '----------'

def edit_template(key):
    '''Editing the template.yml and removing multiple key:value pairs.'''
    template_yml = open('config/locales/template.yml')
    template_dict_full = yaml.safe_load(template_yml)
    template_dict = template_dict_full['de']
    template_yml.close()
    template_yml = open('config/locales/template.yml', 'w')
    del template_dict[key]
#    print yaml.dump_all([template_dict_full], default_flow_style=False, width=300, line_break=False, allow_unicode=True, explicit_start=True)
    yaml.dump_all([template_dict_full], stream=template_yml, default_flow_style=False, width=1024, line_break=False, allow_unicode=True, explicit_start=True)
    template_yml.close()

def switch_keyvalue(dict):
    '''Switching the key:value pairs in the template.yml file.'''
    new_dict = {}
    for key, value in dict.iteritems():
        if value in new_dict:
            new_dict[value].append(key)
        else:
            new_dict[value] = [key]
    return new_dict

def update_template():
    '''Clearing false encodings from template.yml'''
    print 'Updating template.yml'
    replace = {'\\xDC': 'Ü',
               '\\xFC': 'ü',
               '\\xC4': 'Ä',
               '\\xE4': 'ä',
               '\\xD6': 'Ö',
               '\\xF6': 'ö',
               '\\xDF': 'ß'}
    for key, value in replace.iteritems():
        wrong_key = key
        correct_key = value
        sed = "sed -i 's/\\" + wrong_key + "/" + correct_key + "/g' config/locales/template.yml"
        os.system(sed)

'''Choosing the file to work with.'''
print 'Using file "config/locales/template.yml"'
template_yml = open('config/locales/template.yml')
file_processing()
print 'Exit'
<commit_msg>deletion of "update_template()" function and some minor changes when saving the updated template.yml file<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
Author: Timo Denissen
E-Mail: timo@communtu.com
About this program:
Todo: - Give the possibility to change the key name
      - GUI
'''

import yaml
import os

def file_processing():
    '''Working with the file.'''
    mydict = yaml.safe_load(template_yml)
    template_yml.close()
    new_dict = switch_keyvalue(mydict['de'])
    check_dict(new_dict)

def check_dict(dict):
    '''Check if there are multiple values for one key.'''
    for key, value in dict.iteritems():
        if len(value) > 1:
            print 'Please solve this conflict:'
            multiple_string = value
            print value, ':', key ,'\n'
            ask_string(multiple_string)
        else:
            pass

def ask_string(string):
    '''Ask user which string to use in templaye.yml'''
    temp_dict = {}
    x = 1
    for item in string:
        temp_dict[x] = item
        print x, item
        x += 1
    ask_number = input('Enter the number of the correct key: ')
    newkey = temp_dict[ask_number]
    string.remove(newkey)
    oldkey_temp = string
    edit_files(oldkey_temp, newkey)

def edit_files(oldkey, newkey):
    '''Using "sed" to remove oldkey from source files.'''
    for item in oldkey:
        oldkey = item
        sed = 'find app lib -name "*rb" -exec sed -i \'s/t(:' + oldkey + '\([^:alnum:_]\)/t(:' + newkey + '\\1/g\' {} \;'
        print 'These files will be changed:\n'
        os.system('grep --include=\'*rb\' "' + oldkey + '" ./app -r')
        if raw_input('\nIs this ok? y/n ') == 'n':
            break
        else:
            pass
        os.system(sed)
        edit_template(item)
    print '----------'

def edit_template(key):
    '''Editing the template.yml and removing multiple key:value pairs.'''
    template_yml = open('config/locales/template.yml')
    template_dict_full = yaml.safe_load(template_yml)
    template_dict = template_dict_full['de']
    template_yml.close()
    template_yml = open('config/locales/template.yml', 'w')
    del template_dict[key]
    yaml.dump_all([template_dict_full], stream=template_yml, default_flow_style=False, width=2048, line_break=False, allow_unicode=True, explicit_start=True)
    template_yml.close()

def switch_keyvalue(dict):
    '''Switching the key:value pairs in the template.yml file.'''
    new_dict = {}
    for key, value in dict.iteritems():
        if value in new_dict:
            new_dict[value].append(key)
        else:
            new_dict[value] = [key]
    return new_dict

'''Choosing the file to work with.'''
print 'Using file "config/locales/template.yml"'
template_yml = open('config/locales/template.yml')
file_processing()
print 'Exit'
<|endoftext|>


<commit_before>from os import path, makedirs
from hitchtest.utils import get_hitch_directory
import json

def is_modified(includes, excludes=None):
    """Return True if matching file(s) are either modified or no record of them being modified is kept."""

    # If checking just for one file, make it a list
    if type(includes) is str:
        includes = [includes, ]

    if excludes is None:
        excludes = []

    modififications_filename = path.join(get_hitch_directory(), "modified.json")

    if path.exists(modififications_filename):
        with open(modififications_filename, 'r') as modifications_filename_handle:
            modifications = json.loads(modifications_filename_handle.read())
    else:
        modifications = {}

    at_least_one_modified = False

    filenames = list(set(includes) - set(excludes))

    for filename in filenames:
        modified_time = path.getmtime(filename)

        if filename in modifications:
            if modified_time > modifications[filename]:
                modifications[filename] = modified_time
                at_least_one_modified = True
        else:
            modifications[filename] = modified_time
            at_least_one_modified = True

    with open(modififications_filename, 'w') as modifications_file_handle:
        modifications_file_handle.write(json.dumps(modifications))

    return at_least_one_modified
<commit_msg>BUG : Fixed lack of absolute filenames in is_modified checker.<commit_after>from os import path, makedirs
from hitchtest.utils import get_hitch_directory
import json

def is_modified(includes, excludes=None):
    """Return True if matching file(s) are either modified or no record of them being modified is kept."""

    # If checking just for one file, make it a list
    if type(includes) is str:
        includes = [includes, ]

    if excludes is None:
        excludes = []

    modififications_filename = path.join(get_hitch_directory(), "modified.json")

    if path.exists(modififications_filename):
        with open(modififications_filename, 'r') as modifications_filename_handle:
            modifications = json.loads(modifications_filename_handle.read())
    else:
        modifications = {}

    at_least_one_modified = False

    filenames = list(set(includes) - set(excludes))

    for relfilename in filenames:
        absfilename = path.abspath(relfilename)
        modified_time = path.getmtime(absfilename)

        if absfilename in modifications:
            if modified_time > modifications[absfilename]:
                modifications[absfilename] = modified_time
                at_least_one_modified = True
        else:
            modifications[absfilename] = modified_time
            at_least_one_modified = True

    with open(modififications_filename, 'w') as modifications_file_handle:
        modifications_file_handle.write(json.dumps(modifications))

    return at_least_one_modified
<|endoftext|>


<commit_before>#!/usr/bin/python3

# determine the connected groups of images.  Images without
# connections to each other cannot be correctly placed.

import argparse
import pickle
import os.path
import sys

sys.path.append('../lib')
import Groups
import ProjectMgr

parser = argparse.ArgumentParser(description='Keypoint projection.')
parser.add_argument('--project', required=True, help='project directory')
parser.add_argument('--area', required=True, help='sub area directory')
# parser.add_argument('--original-pairs', action='store_true', help='use original pair-rwise matches')
args = parser.parse_args()

proj = ProjectMgr.ProjectMgr(args.project)
proj.load_area_info(args.area)

area_dir = os.path.join(args.project, args.area)
source = 'matches_grouped'
print("Loading source matches:", source)
matches = pickle.load( open( os.path.join(area_dir, source), 'rb' ) )

print("features:", len(matches))

# if not args.original_pairs:
#     # recreate the pair-wise match structure (so we can honor any
#     # features/pairs removed by outlier rejection strategies
#     for i1 in proj.image_list:
#         i1.match_list = {}
#     for match in matches:
#         for p1 in match[1:]:
#             image1 = proj.image_list[p1[0]]
#             for p2 in match[1:]:
#                 if p1 == p2:
#                     continue
#                 image2 = proj.image_list[p2[0]]
#                 if image2.name in image1.match_list:
#                     image1.match_list[image2.name].append( [p1[1], p2[1]] )
#                 else:
#                     image1.match_list[image2.name] = [ [p1[1], p2[1]] ]
# else:
#     proj.load_match_pairs(extra_verbose=False)

# compute the group connections within the image set.

groups, used_features = Groups.groupByFeatureConnections(proj.image_list, matches)
# groups = Groups.groupByConnectedArea(proj.image_list, matches)
# groups = Groups.groupByImageConnections(proj)

groups.sort(key=len, reverse=True)

Groups.save(area_dir, groups)

print('Total images:', len(proj.image_list))
print('Group sizes:', end=" ")
for g in groups:
    print(len(g), end=" ")
print()

# generate the matches_used file
matches_used = []
for i, match in enumerate(matches):
    if used_features[i]:
        matches_used.append(match)

print("Writing matches_used...")
print("Features: %d/%d" % (len(matches_used), len(matches)))
pickle.dump(matches_used, open(os.path.join(area_dir, "matches_used"), "wb"))

# this is extra (and I'll put it here for now for lack of a better
# place), but for visualization's sake, create a gnuplot data file
# that will show all the match connectivity in the set.
# file = os.path.join(area_dir, 'connections.gnuplot')
# f = open(file, 'w')
# pair_dict = {}
# for match in matches:
#     for m1 in match[1:]:
#         for m2 in match[1:]:
#             if m1[0] == m2[0]:
#                 # skip selfies
#                 continue
#             key = '%d %d' % (m1[0], m2[0])
#             pair_dict[key] = [m1[0], m2[0]]
# for pair in pair_dict:
#     #print 'pair:', pair, pair_dict[pair]
#     image1 = proj.image_list[pair_dict[pair][0]]
#     image2 = proj.image_list[pair_dict[pair][1]]
#     (ned1, ypr1, quat1) = image1.get_camera_pose()
#     (ned2, ypr2, quat2) = image2.get_camera_pose()
#     f.write("%.2f %.2f\n" % (ned1[1], ned1[0]))
#     f.write("%.2f %.2f\n" % (ned2[1], ned2[0]))
#     f.write("\n")
# f.close()

<commit_msg>Clean up cruft.<commit_after>#!/usr/bin/python3

# determine the connected groups of images.  Images without
# connections to each other cannot be correctly placed.

import argparse
import pickle
import os.path
import sys

sys.path.append('../lib')
import Groups
import ProjectMgr

parser = argparse.ArgumentParser(description='Keypoint projection.')
parser.add_argument('--project', required=True, help='project directory')
parser.add_argument('--area', required=True, help='sub area directory')
args = parser.parse_args()

proj = ProjectMgr.ProjectMgr(args.project)
proj.load_area_info(args.area)

area_dir = os.path.join(args.project, args.area)
source = 'matches_grouped'
print("Loading source matches:", source)
matches = pickle.load( open( os.path.join(area_dir, source), 'rb' ) )

print("features:", len(matches))

# compute the group connections within the image set.

groups, used_features = Groups.groupByFeatureConnections(proj.image_list, matches)
# groups = Groups.groupByConnectedArea(proj.image_list, matches)
# groups = Groups.groupByImageConnections(proj)

groups.sort(key=len, reverse=True)
Groups.save(area_dir, groups)

print('Total images:', len(proj.image_list))
print('Group sizes:', end=" ")
for g in groups:
    print(len(g), end=" ")
print()

# generate the matches_used file
matches_used = []
for i, match in enumerate(matches):
    if used_features[i]:
        matches_used.append(match)

print("Writing matches_used...")
print("Features: %d/%d" % (len(matches_used), len(matches)))
pickle.dump(matches_used, open(os.path.join(area_dir, "matches_used"), "wb"))

# this is extra (and I'll put it here for now for lack of a better
# place), but for visualization's sake, create a gnuplot data file
# that will show all the match connectivity in the set.
# file = os.path.join(area_dir, 'connections.gnuplot')
# f = open(file, 'w')
# pair_dict = {}
# for match in matches:
#     for m1 in match[1:]:
#         for m2 in match[1:]:
#             if m1[0] == m2[0]:
#                 # skip selfies
#                 continue
#             key = '%d %d' % (m1[0], m2[0])
#             pair_dict[key] = [m1[0], m2[0]]
# for pair in pair_dict:
#     #print 'pair:', pair, pair_dict[pair]
#     image1 = proj.image_list[pair_dict[pair][0]]
#     image2 = proj.image_list[pair_dict[pair][1]]
#     (ned1, ypr1, quat1) = image1.get_camera_pose()
#     (ned2, ypr2, quat2) = image2.get_camera_pose()
#     f.write("%.2f %.2f\n" % (ned1[1], ned1[0]))
#     f.write("%.2f %.2f\n" % (ned2[1], ned2[0]))
#     f.write("\n")
# f.close()

<|endoftext|>


<commit_before>__version__ = '0.2'
<commit_msg>0.3alpha<commit_after>__version__ = '0.3a'
<|endoftext|>


<commit_before>#Some house_ids with data:
#6   15   83   360   362   365   368   380   382   385   386   390   404   405   409

import sys
import pandas as pd
import numpy as np

pd.set_option('max_columns', 50)

#allcsv = pd.read_csv('newcompact.csv')
#URS = pd.read_csv('newURS.csv')

def householdPrice(houseID, allcsv, URS):

    us_state_abbrev = {
        'Alabama': 'AL',
        'Alaska': 'AK',
        'Arizona': 'AZ',
        'Arkansas': 'AR',
        'California': 'CA',
        'Colorado': 'CO',
        'Connecticut': 'CT',
        'Delaware': 'DE',
        'Florida': 'FL',
        'Georgia': 'GA',
        'Hawaii': 'HI',
        'Idaho': 'ID',
        'Illinois': 'IL',
        'Indiana': 'IN',
        'Iowa': 'IA',
        'Kansas': 'KS',
        'Kentucky': 'KY',
        'Louisiana': 'LA',
        'Maine': 'ME',
        'Maryland': 'MD',
        'Massachusetts': 'MA',
        'Michigan': 'MI',
        'Minnesota': 'MN',
        'Mississippi': 'MS',
        'Missouri': 'MO',
        'Montana': 'MT',
        'Nebraska': 'NE',
        'Nevada': 'NV',
        'New Hampshire': 'NH',
        'New Jersey': 'NJ',
        'New Mexico': 'NM',
        'New York': 'NY',
        'North Carolina': 'NC',
        'North Dakota': 'ND',
        'Ohio': 'OH',
        'Oklahoma': 'OK',
        'Oregon': 'OR',
        'Pennsylvania': 'PA',
        'Rhode Island': 'RI',
        'South Carolina': 'SC',
        'South Dakota': 'SD',
        'Tennessee': 'TN',
        'Texas': 'TX',
        'Utah': 'UT',
        'Vermont': 'VT',
        'Virginia': 'VA',
        'Washington': 'WA',
        'West Virginia': 'WV',
        'Wisconsin': 'WI',
        'Wyoming': 'WY',
    }

    inv_map = {v: k for k, v in us_state_abbrev.items()}

    house=allcsv[allcsv.unit_id==houseID]

    houseState= inv_map[house.STATE.values[0]]
    houseISP = house.isp.values[0]
    print houseState
    print houseISP
    rates = URS[URS.State==houseState and URS.isp==houseISP]
    return rates

#print householdPrice(6)
#print householdPrice(15)

# to merge=[]
# for i in list(allcsv.unit_id):
# 	tomerge.append({'unit_id':i, 'price':pricefind(i)})
# def pricefind(unit):
# 	speedup=allcsv[allcsv.unit_id==unit]["SK UP"]
# 	speeddown=allcsv[allcsv.unit_id==unit]["SK down"]
# 	isp=allcsv[allcsv.unit_id==unit]["isp"]
# 	tech=allcsv[allcsv.unit_id==unit]["TECHNOLOGY"]
# 	state=allcsv[allcsv.unit_id==unit]["STATE"]
# 	URS=abs(URS[URS["Download Bandwidth Mbps"])
#profilemba = pd.read_csv('fixedishunitprofile.csv')
#print URS
'''URS = URS[URS.Technology != "Other"]
URS = URS.rename(columns={'Provider': 'isp', 'Technology': 'TECHNOLOGY'})
URS=URS.replace("Provider","isp")
URS=URS.replace("Technology","TECHNOLOGY")
URS=URS.replace("AT&T Services, Inc.", "AT&T")
URS=URS.replace("Charter Communications, Inc.","Charter")
URS=URS.replace("Charter Communications","Charter")
URS=URS.replace('Mediacom Minnesota LLC','Mediacom')
URS=URS.replace('Mediacom Illinois LLC','Mediacom')
URS=URS.replace('Mediacom Iowa LLC','Mediacom')
URS=URS.replace('Verizon New England Inc.','Verizon')
URS=URS.replace('Verizon California Inc.','Verizon')
URS=URS.replace('Verizon New Jersey Inc.','Verizon')
URS=URS.replace('Verizon Florida LLC','Verizon')
URS=URS.replace('Verizon Pennsylvania Inc.','Verizon')
URS=URS.replace('Verizon New York Inc.','Verizon')
URS=URS.replace('Verizon Virginia Inc.','Verizon')
#profilemba=profilemba.replace('Verizon DSL','Verizon')
URS=URS.replace('Frontier Communications Corporation','Frontier')
URS=URS.replace('Cox Communications, Inc','Cox')
URS=URS.replace('Time Warner Cable Enterprises LLC','Time Warner Cable')
URS=URS.replace('Time Warner Cable Inc.','Time Warner Cable')
URS=URS.replace('BENTON CABLEVISION, INC.','Cablevision')
URS=URS.replace('Windstream Corporation','Windstream')
URS=URS.replace('COMCAST CABLE COMMUNICATIONS, INC.','Comcast')
URS=URS.replace('CenturyLink, Inc.','CenturyLink')
URS=URS.replace('Fixed wireless','SATELLITE')
URS=URS.replace('FTTH','FIBER')
URS=URS.replace('Cable','CABLE')
URS.to_csv("testchange.csv")
#print "written"
URSlist= ['Frontier Communications','Liberty Cablevision of Puerto Rico LLC','Canby Telephone Association', 'COMCAST CABLE COMMUNICATIONS, INC.', 'Northern New England Telephone Operations LLC', 'RCN Telecom Services (Lehigh) LLC', 'SureWest Kansas Operations, LLC', 'Hopper Telecommunications LLC', 'Savage Communications Inc.', 'Simply Bits, LLC', 'Sonic Telecom, LLC', 'The Computer Works', 'DerbyNet, LLC', 'Cable One, Inc.', 'RCN Telecom Services of New York, LP', 'WireFree Communications, Inc.', 'Armstrong Utilities, Inc.', 'RCN Telecom Services of Philadelphia, LLC', 'Cincinnati Bell Extended Territories LLC', 'Ranch Wireless, Inc.', 'ImOn Communications, LLC', 'Freedom Wireless Broadband, LLC', 'TDS TELECOMMUNICATIONS CORPORATION', 'Bright House Networks, LLC', 'Excel.Net, Inc.', 'Shidler Telephone Company', 'Chesnee Cable, Inc.', 'Farmers Telephone Cooperative, Inc.', 'MTCO Communications, Inc.', 'Google Fiber Kansas, LLC', 'DigitalPath, Inc.', 'MCC Georgia LLC', 'Netwurx', 'Puerto Rico Cable Acquisition Company Inc.', 'Otelco Telephone', 'GoldStar Communications, LLC', 'WEHCO Video, Inc.', 'Puerto Rico Telephone Company, Inc.', 'Enterprises LLC' , 'Digis LLC', 'Hawaiian Telcom, Inc.', 'Mega Broadband, Inc', 'E-Vergent.Com LLC', 'CentraCom Interactive', 'CSC Holdings LLC', 'Community Cable & Broadband, Inc.', 'Moose Lake, Village Of/dba, Moose-Tec', 'Chesnee Telephone Company, Inc. d/b/a Chesnee Communications', 'SpeedConnect LLC', '3 Rooms Communications LLC', 'Texas Communications of Bryan Inc.', 'FTC Diversified Services, LLC', 'BlountBroadband LLC', 'Atlantic Telephone Membership Corporation', 'Knology of Alabama, Inc.', 'Miles Wireless', 'Blountsville Telephone LLC', 'Fire2Wire', 'Blue Ridge Cable Technologies, Inc.', 'VTX Communications, LLC', 'Surf Air Wireless', '4 SIWI LLC', 'Beehive Broadband LLC UTAH', 'Delta Communications', 'ALASKA COMMUNICATIONS SYSTEMS HOLDING, INC.', 'Suddenlink Communications', 'Cimarron Telephone Company', 'MetaLINK Technologies, Inc.', 'Google Fiber Missouri, LLC', 'MetroCast Communications of Connecticut, LLC', 'Buford Communications 1, L.P.', 'Atlantic Broadband (Penn), LLC', 'Ultra Communications Group, LLC', 'FBN Indiana Inc.', 'WaveDivision Holdings, LLC', 'Orlando Telephone Company', 'GCI Communication Corp.', 'VTX Telecom, LLC.', 'Amplex Electric, Inc.', 'SureWest TeleVideo', 'Service Electric Cable TV, Inc.', 'AlasConnect, Inc.', 'Knology of Florida, Inc.', 'TRANSWORLD NETWORK, CORP', 'Cequel Communications, LLC']
for i in URSlist:
	URS=URS[URS.isp != i]
mbalist = ['Wildblue/ViaSat', 'Brighthouse', 'Hughes', 'Insight',  'Qwest']
for j in mbalist:
	profilemba=profilemba[profilemba.isp != j]'''
#"FTTH"="FIBER"
#"Fixed wireless"="SATELLITE"
#"Cable"="CABLE"
#"DSL"="DSL"
#print "profile:"
#print list(set(profilemba["isp"]))
#print list(set(profilemba["TECHNOLOGY"]))

#print "URS:"
#print list(set(URS["isp"]))
#print list(set(URS["TECHNOLOGY"]))
#new=pd.merge(allcsv,profilemba, on='unit_id', how='inner')
#new.to_csv('newcompact.csv')
#print 'merged'
<commit_msg>updated with output<commit_after>#Some house_ids with data:
#6   15   83   360   362   365   368   380   382   385   386   390   404   405   409

import sys
import pandas as pd
import numpy as np

pd.set_option('max_columns', 50)

#allcsv = pd.read_csv('newcompact.csv')
#URS = pd.read_csv('newURS.csv')

def householdPrice(houseID, allcsv, URS):

    us_state_abbrev = {
        'Alabama': 'AL',
        'Alaska': 'AK',
        'Arizona': 'AZ',
        'Arkansas': 'AR',
        'California': 'CA',
        'Colorado': 'CO',
        'Connecticut': 'CT',
        'Delaware': 'DE',
        'Florida': 'FL',
        'Georgia': 'GA',
        'Hawaii': 'HI',
        'Idaho': 'ID',
        'Illinois': 'IL',
        'Indiana': 'IN',
        'Iowa': 'IA',
        'Kansas': 'KS',
        'Kentucky': 'KY',
        'Louisiana': 'LA',
        'Maine': 'ME',
        'Maryland': 'MD',
        'Massachusetts': 'MA',
        'Michigan': 'MI',
        'Minnesota': 'MN',
        'Mississippi': 'MS',
        'Missouri': 'MO',
        'Montana': 'MT',
        'Nebraska': 'NE',
        'Nevada': 'NV',
        'New Hampshire': 'NH',
        'New Jersey': 'NJ',
        'New Mexico': 'NM',
        'New York': 'NY',
        'North Carolina': 'NC',
        'North Dakota': 'ND',
        'Ohio': 'OH',
        'Oklahoma': 'OK',
        'Oregon': 'OR',
        'Pennsylvania': 'PA',
        'Rhode Island': 'RI',
        'South Carolina': 'SC',
        'South Dakota': 'SD',
        'Tennessee': 'TN',
        'Texas': 'TX',
        'Utah': 'UT',
        'Vermont': 'VT',
        'Virginia': 'VA',
        'Washington': 'WA',
        'West Virginia': 'WV',
        'Wisconsin': 'WI',
        'Wyoming': 'WY',
    }

    inv_map = {v: k for k, v in us_state_abbrev.items()}

    house=allcsv[allcsv.unit_id==houseID]

    houseState= inv_map[house.STATE.values[0]]
    houseISP = house.isp.values[0]
    print houseState
    print houseISP
    rates = URS[URS.State==houseState]
    rates = rates[rates.isp==houseISP]
    if rates.empty
        return float('NaN')
    return rates

#print householdPrice(6)
#print householdPrice(15)

# to merge=[]
# for i in list(allcsv.unit_id):
# 	tomerge.append({'unit_id':i, 'price':pricefind(i)})
# def pricefind(unit):
# 	speedup=allcsv[allcsv.unit_id==unit]["SK UP"]
# 	speeddown=allcsv[allcsv.unit_id==unit]["SK down"]
# 	isp=allcsv[allcsv.unit_id==unit]["isp"]
# 	tech=allcsv[allcsv.unit_id==unit]["TECHNOLOGY"]
# 	state=allcsv[allcsv.unit_id==unit]["STATE"]
# 	URS=abs(URS[URS["Download Bandwidth Mbps"])
#profilemba = pd.read_csv('fixedishunitprofile.csv')
#print URS
'''URS = URS[URS.Technology != "Other"]
URS = URS.rename(columns={'Provider': 'isp', 'Technology': 'TECHNOLOGY'})
URS=URS.replace("Provider","isp")
URS=URS.replace("Technology","TECHNOLOGY")
URS=URS.replace("AT&T Services, Inc.", "AT&T")
URS=URS.replace("Charter Communications, Inc.","Charter")
URS=URS.replace("Charter Communications","Charter")
URS=URS.replace('Mediacom Minnesota LLC','Mediacom')
URS=URS.replace('Mediacom Illinois LLC','Mediacom')
URS=URS.replace('Mediacom Iowa LLC','Mediacom')
URS=URS.replace('Verizon New England Inc.','Verizon')
URS=URS.replace('Verizon California Inc.','Verizon')
URS=URS.replace('Verizon New Jersey Inc.','Verizon')
URS=URS.replace('Verizon Florida LLC','Verizon')
URS=URS.replace('Verizon Pennsylvania Inc.','Verizon')
URS=URS.replace('Verizon New York Inc.','Verizon')
URS=URS.replace('Verizon Virginia Inc.','Verizon')
#profilemba=profilemba.replace('Verizon DSL','Verizon')
URS=URS.replace('Frontier Communications Corporation','Frontier')
URS=URS.replace('Cox Communications, Inc','Cox')
URS=URS.replace('Time Warner Cable Enterprises LLC','Time Warner Cable')
URS=URS.replace('Time Warner Cable Inc.','Time Warner Cable')
URS=URS.replace('BENTON CABLEVISION, INC.','Cablevision')
URS=URS.replace('Windstream Corporation','Windstream')
URS=URS.replace('COMCAST CABLE COMMUNICATIONS, INC.','Comcast')
URS=URS.replace('CenturyLink, Inc.','CenturyLink')
URS=URS.replace('Fixed wireless','SATELLITE')
URS=URS.replace('FTTH','FIBER')
URS=URS.replace('Cable','CABLE')
URS.to_csv("testchange.csv")
#print "written"
URSlist= ['Frontier Communications','Liberty Cablevision of Puerto Rico LLC','Canby Telephone Association', 'COMCAST CABLE COMMUNICATIONS, INC.', 'Northern New England Telephone Operations LLC', 'RCN Telecom Services (Lehigh) LLC', 'SureWest Kansas Operations, LLC', 'Hopper Telecommunications LLC', 'Savage Communications Inc.', 'Simply Bits, LLC', 'Sonic Telecom, LLC', 'The Computer Works', 'DerbyNet, LLC', 'Cable One, Inc.', 'RCN Telecom Services of New York, LP', 'WireFree Communications, Inc.', 'Armstrong Utilities, Inc.', 'RCN Telecom Services of Philadelphia, LLC', 'Cincinnati Bell Extended Territories LLC', 'Ranch Wireless, Inc.', 'ImOn Communications, LLC', 'Freedom Wireless Broadband, LLC', 'TDS TELECOMMUNICATIONS CORPORATION', 'Bright House Networks, LLC', 'Excel.Net, Inc.', 'Shidler Telephone Company', 'Chesnee Cable, Inc.', 'Farmers Telephone Cooperative, Inc.', 'MTCO Communications, Inc.', 'Google Fiber Kansas, LLC', 'DigitalPath, Inc.', 'MCC Georgia LLC', 'Netwurx', 'Puerto Rico Cable Acquisition Company Inc.', 'Otelco Telephone', 'GoldStar Communications, LLC', 'WEHCO Video, Inc.', 'Puerto Rico Telephone Company, Inc.', 'Enterprises LLC' , 'Digis LLC', 'Hawaiian Telcom, Inc.', 'Mega Broadband, Inc', 'E-Vergent.Com LLC', 'CentraCom Interactive', 'CSC Holdings LLC', 'Community Cable & Broadband, Inc.', 'Moose Lake, Village Of/dba, Moose-Tec', 'Chesnee Telephone Company, Inc. d/b/a Chesnee Communications', 'SpeedConnect LLC', '3 Rooms Communications LLC', 'Texas Communications of Bryan Inc.', 'FTC Diversified Services, LLC', 'BlountBroadband LLC', 'Atlantic Telephone Membership Corporation', 'Knology of Alabama, Inc.', 'Miles Wireless', 'Blountsville Telephone LLC', 'Fire2Wire', 'Blue Ridge Cable Technologies, Inc.', 'VTX Communications, LLC', 'Surf Air Wireless', '4 SIWI LLC', 'Beehive Broadband LLC UTAH', 'Delta Communications', 'ALASKA COMMUNICATIONS SYSTEMS HOLDING, INC.', 'Suddenlink Communications', 'Cimarron Telephone Company', 'MetaLINK Technologies, Inc.', 'Google Fiber Missouri, LLC', 'MetroCast Communications of Connecticut, LLC', 'Buford Communications 1, L.P.', 'Atlantic Broadband (Penn), LLC', 'Ultra Communications Group, LLC', 'FBN Indiana Inc.', 'WaveDivision Holdings, LLC', 'Orlando Telephone Company', 'GCI Communication Corp.', 'VTX Telecom, LLC.', 'Amplex Electric, Inc.', 'SureWest TeleVideo', 'Service Electric Cable TV, Inc.', 'AlasConnect, Inc.', 'Knology of Florida, Inc.', 'TRANSWORLD NETWORK, CORP', 'Cequel Communications, LLC']
for i in URSlist:
	URS=URS[URS.isp != i]
mbalist = ['Wildblue/ViaSat', 'Brighthouse', 'Hughes', 'Insight',  'Qwest']
for j in mbalist:
	profilemba=profilemba[profilemba.isp != j]'''
#"FTTH"="FIBER"
#"Fixed wireless"="SATELLITE"
#"Cable"="CABLE"
#"DSL"="DSL"
#print "profile:"
#print list(set(profilemba["isp"]))
#print list(set(profilemba["TECHNOLOGY"]))

#print "URS:"
#print list(set(URS["isp"]))
#print list(set(URS["TECHNOLOGY"]))
#new=pd.merge(allcsv,profilemba, on='unit_id', how='inner')
#new.to_csv('newcompact.csv')
#print 'merged'
<|endoftext|>


<commit_before># coding: utf-8

# Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

# Maintainer: joaander / All Developers are free to add commands for new
# features

import itertools
from functools import partial

from hoomd.md import _md
from hoomd.data.parameterdicts import ParameterDict
from hoomd.data.typeconverter import OnlyFrom
from hoomd.integrate import BaseIntegrator
from hoomd.data import syncedlist
from hoomd.md.methods import _Method
from hoomd.md.force import Force
from hoomd.md.constrain import ConstraintForce

def _preprocess_aniso(value):
    if value is True:
        return "true"
    elif value is False:
        return "false"
    else:
        return value

def _set_synced_list(old_list, new_list):
    old_list.clear()
    old_list.extend(new_list)

class _DynamicIntegrator(BaseIntegrator):
    def __init__(self, forces, constraints, methods):
        forces = [] if forces is None else forces
        constraints = [] if constraints is None else constraints
        methods = [] if methods is None else methods
        self._forces = syncedlist.SyncedList(
            Force, syncedlist._PartialGetAttr('_cpp_obj'), iterable=forces)

        self._constraints = syncedlist.SyncedList(
            ConstraintForce,
            syncedlist._PartialGetAttr('_cpp_obj'),
            iterable=constraints)

        self._methods = syncedlist.SyncedList(
            _Method, syncedlist._PartialGetAttr('_cpp_obj'), iterable=methods)

    def _attach(self):
        self.forces._sync(self._simulation, self._cpp_obj.forces)
        self.constraints._sync(self._simulation, self._cpp_obj.constraints)
        self.methods._sync(self._simulation, self._cpp_obj.methods)
        super()._attach()

    @property
    def forces(self):
        return self._forces

    @forces.setter
    def forces(self, value):
        _set_synced_list(self._forces, value)

    @property
    def constraints(self):
        return self._constraints

    @constraints.setter
    def constraints(self, value):
        _set_synced_list(self._constraints, value)

    @property
    def methods(self):
        return self._methods

    @methods.setter
    def methods(self, value):
        _set_synced_list(self._methods, value)

    @property
    def _children(self):
        children = list(self.forces)
        children.extend(self.constraints)
        children.extend(self.methods)

        for child in itertools.chain(self.forces, self.constraints,
                                     self.methods):
            children.extend(child._children)

        return children

class Integrator(_DynamicIntegrator):
    R""" Enables a variety of standard integration methods.

    Args:
        dt (float): Integrator time step size (in time units).

        methods (Sequence[hoomd.md.methods._Method]): Sequence of integration
            methods. Each integration method can be applied to only a specific
            subset of particles. The intersection of the subsets must be null.
            The default value of ``None`` initializes an empty list.

        forces (Sequence[hoomd.md.force.Force]): Sequence of forces applied to
            the particles in the system. All the forces are summed together.
            The default value of ``None`` initializes an empty list.

        aniso (str or bool): Whether to integrate rotational degrees of freedom
            (bool), default 'auto' (autodetect if there is anisotropic factor
            from any defined active or constraint forces).

        constraints (Sequence[hoomd.md.constrain.ConstraintForce]): Sequence of
            constraint forces applied to the particles in the system.
            The default value of ``None`` initializes an empty list.

    The following classes can be used as elements in `methods`

    - `hoomd.md.methods.Brownian`
    - `hoomd.md.methods.Langevin`
    - `hoomd.md.methods.NVE`
    - `hoomd.md.methods.NVT`
    - `hoomd.md.methods.NPT`

    The classes of following modules can be used as elements in `forces`

    - `hoomd.md.angle`
    - `hoomd.md.bond`
    - `hoomd.md.charge`
    - `hoomd.md.dihedral`
    - `hoomd.md.external`
    - `hoomd.md.force`
    - `hoomd.md.improper`
    - `hoomd.md.pair`
    - `hoomd.md.wall`
    - `hoomd.md.special_pair`

    The classes of the following module can be used as elements in `constraints`

    - `hoomd.md.constrain`

    Examples::

        nlist = hoomd.md.nlist.Cell()
        lj = hoomd.md.pair.LJ(nlist=nlist)
        lj.params.default = dict(epsilon=1.0, sigma=1.0)
        lj.r_cut[('A', 'A')] = 2**(1/6)
        nve = hoomd.md.methods.NVE(filter=hoomd.filter.All())
        integrator = hoomd.md.Integrator(dt=0.001, methods=[nve], forces=[lj])
        sim.operations.integrator = integrator

    Attributes:
        dt (float): Integrator time step size (in time units).

        methods (List[hoomd.md.methods._Method]): List of integration methods.
            Each integration method can be applied to only a specific subset of
            particles.

        forces (List[hoomd.md.force.Force]): List of forces applied to
            the particles in the system. All the forces are summed together.

        aniso (str): Whether rotational degrees of freedom are integrated.

        constraints (List[hoomd.md.constrain.ConstraintForce]): List of
            constraint forces applied to the particles in the system.
    """

    def __init__(self, dt, aniso='auto', forces=None, constraints=None,
                 methods=None):

        super().__init__(forces, constraints, methods)

        self._param_dict = ParameterDict(
            dt=float(dt),
            aniso=OnlyFrom(['true', 'false', 'auto'],
                           preprocess=_preprocess_aniso),
            _defaults=dict(aniso="auto")
            )
        if aniso is not None:
            self.aniso = aniso

    def _attach(self):
        # initialize the reflected c++ class
        self._cpp_obj = _md.IntegratorTwoStep(
            self._simulation.state._cpp_sys_def, self.dt)
        # Call attach from DynamicIntegrator which attaches forces,
        # constraint_forces, and methods, and calls super()._attach() itself.
        super()._attach()
<commit_msg>Remove unused import<commit_after># coding: utf-8

# Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

# Maintainer: joaander / All Developers are free to add commands for new
# features

import itertools

from hoomd.md import _md
from hoomd.data.parameterdicts import ParameterDict
from hoomd.data.typeconverter import OnlyFrom
from hoomd.integrate import BaseIntegrator
from hoomd.data import syncedlist
from hoomd.md.methods import _Method
from hoomd.md.force import Force
from hoomd.md.constrain import ConstraintForce

def _preprocess_aniso(value):
    if value is True:
        return "true"
    elif value is False:
        return "false"
    else:
        return value

def _set_synced_list(old_list, new_list):
    old_list.clear()
    old_list.extend(new_list)

class _DynamicIntegrator(BaseIntegrator):
    def __init__(self, forces, constraints, methods):
        forces = [] if forces is None else forces
        constraints = [] if constraints is None else constraints
        methods = [] if methods is None else methods
        self._forces = syncedlist.SyncedList(
            Force, syncedlist._PartialGetAttr('_cpp_obj'), iterable=forces)

        self._constraints = syncedlist.SyncedList(
            ConstraintForce,
            syncedlist._PartialGetAttr('_cpp_obj'),
            iterable=constraints)

        self._methods = syncedlist.SyncedList(
            _Method, syncedlist._PartialGetAttr('_cpp_obj'), iterable=methods)

    def _attach(self):
        self.forces._sync(self._simulation, self._cpp_obj.forces)
        self.constraints._sync(self._simulation, self._cpp_obj.constraints)
        self.methods._sync(self._simulation, self._cpp_obj.methods)
        super()._attach()

    @property
    def forces(self):
        return self._forces

    @forces.setter
    def forces(self, value):
        _set_synced_list(self._forces, value)

    @property
    def constraints(self):
        return self._constraints

    @constraints.setter
    def constraints(self, value):
        _set_synced_list(self._constraints, value)

    @property
    def methods(self):
        return self._methods

    @methods.setter
    def methods(self, value):
        _set_synced_list(self._methods, value)

    @property
    def _children(self):
        children = list(self.forces)
        children.extend(self.constraints)
        children.extend(self.methods)

        for child in itertools.chain(self.forces, self.constraints,
                                     self.methods):
            children.extend(child._children)

        return children

class Integrator(_DynamicIntegrator):
    R""" Enables a variety of standard integration methods.

    Args:
        dt (float): Integrator time step size (in time units).

        methods (Sequence[hoomd.md.methods._Method]): Sequence of integration
            methods. Each integration method can be applied to only a specific
            subset of particles. The intersection of the subsets must be null.
            The default value of ``None`` initializes an empty list.

        forces (Sequence[hoomd.md.force.Force]): Sequence of forces applied to
            the particles in the system. All the forces are summed together.
            The default value of ``None`` initializes an empty list.

        aniso (str or bool): Whether to integrate rotational degrees of freedom
            (bool), default 'auto' (autodetect if there is anisotropic factor
            from any defined active or constraint forces).

        constraints (Sequence[hoomd.md.constrain.ConstraintForce]): Sequence of
            constraint forces applied to the particles in the system.
            The default value of ``None`` initializes an empty list.

    The following classes can be used as elements in `methods`

    - `hoomd.md.methods.Brownian`
    - `hoomd.md.methods.Langevin`
    - `hoomd.md.methods.NVE`
    - `hoomd.md.methods.NVT`
    - `hoomd.md.methods.NPT`

    The classes of following modules can be used as elements in `forces`

    - `hoomd.md.angle`
    - `hoomd.md.bond`
    - `hoomd.md.charge`
    - `hoomd.md.dihedral`
    - `hoomd.md.external`
    - `hoomd.md.force`
    - `hoomd.md.improper`
    - `hoomd.md.pair`
    - `hoomd.md.wall`
    - `hoomd.md.special_pair`

    The classes of the following module can be used as elements in `constraints`

    - `hoomd.md.constrain`

    Examples::

        nlist = hoomd.md.nlist.Cell()
        lj = hoomd.md.pair.LJ(nlist=nlist)
        lj.params.default = dict(epsilon=1.0, sigma=1.0)
        lj.r_cut[('A', 'A')] = 2**(1/6)
        nve = hoomd.md.methods.NVE(filter=hoomd.filter.All())
        integrator = hoomd.md.Integrator(dt=0.001, methods=[nve], forces=[lj])
        sim.operations.integrator = integrator

    Attributes:
        dt (float): Integrator time step size (in time units).

        methods (List[hoomd.md.methods._Method]): List of integration methods.
            Each integration method can be applied to only a specific subset of
            particles.

        forces (List[hoomd.md.force.Force]): List of forces applied to
            the particles in the system. All the forces are summed together.

        aniso (str): Whether rotational degrees of freedom are integrated.

        constraints (List[hoomd.md.constrain.ConstraintForce]): List of
            constraint forces applied to the particles in the system.
    """

    def __init__(self, dt, aniso='auto', forces=None, constraints=None,
                 methods=None):

        super().__init__(forces, constraints, methods)

        self._param_dict = ParameterDict(
            dt=float(dt),
            aniso=OnlyFrom(['true', 'false', 'auto'],
                           preprocess=_preprocess_aniso),
            _defaults=dict(aniso="auto")
            )
        if aniso is not None:
            self.aniso = aniso

    def _attach(self):
        # initialize the reflected c++ class
        self._cpp_obj = _md.IntegratorTwoStep(
            self._simulation.state._cpp_sys_def, self.dt)
        # Call attach from DynamicIntegrator which attaches forces,
        # constraint_forces, and methods, and calls super()._attach() itself.
        super()._attach()
<|endoftext|>


<commit_before>#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
#

'''
calltree.py - Call tree manipulator

This script manipulates call tree generated by JASMIN CallTreeManager. It
prints, folds, and limit depth of the input call tree.
'''

import argparse
import copy
import json

from collections import OrderedDict

def fold_tree(tree):

    def fold_tree_recursive(tree, compare_equal):
        if not tree:
            return None
        new_tree = OrderedDict()
        new_tree["id"] = tree["id"]
        new_tree["cycle"] = tree["cycle"]
        new_tree["children"] = [fold_tree_recursive(t, compare_equal) for t in
                                tree["children"]]
        if not new_tree["children"]:
            return new_tree
        new_children = [new_tree["children"][0]]
        for i in xrange(1, len(new_tree["children"])):
            if not compare_equal(new_tree["children"][i], new_children[-1]):
                new_children.append(new_tree["children"][i])
        new_tree["children"] = new_children
        return new_tree

    def tree_equal(src, dst):
        if src["id"] != dst["id"]:
            return False
        elif len(src["children"]) != len(dst["children"]):
            return False
        else:
            for x, y in zip(src["children"], dst["children"]):
                if not tree_equal(x, y):
                    return False
            return True

    return fold_tree_recursive(tree, tree_equal)

def print_tree(tree, max_level=None):
    def print_tree_recursive(tree, level, max_level):
        if max_level and level > max_level:
            return
        print "--" * level, tree["id"], tree["cycle"]
        for c in tree["children"]:
            print_tree_recursive(c, level + 1, max_level)
    print_tree_recursive(tree, 1, max_level)

def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("call_tree",
                        help="Call tree file (json) generated by JASMIN")
    parser.add_argument("--save", help="Output filename", default=None)
    parser.add_argument("--max-depth", default=None, type=int,
                        help="Max depth for output tree")
    parser.add_argument("--fold", action="store_true",
                        help="Try hard to fold tree")

    args = parser.parse_args()

    tree = json.load(file(args.call_tree), object_hook=OrderedDict)
    if tree["id"] == "ROOT":
        tree = tree["children"][0]
    if args.fold:
        tree = fold_tree(tree)
    if args.save:
        json.dump(tree, file(args.save, "w"), indent=2)
    print_tree(tree, args.max_depth)

if __name__ == "__main__":
    main()
<commit_msg>Refactor and improve bentoo-calltree<commit_after>#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
#

'''
calltree.py - Call tree manipulator

This script manipulates call tree generated by JASMIN CallTreeManager. It
prints, folds, and limit depth of the input call tree.
'''

import argparse
import copy
import json
import fnmatch
import hashlib

from collections import OrderedDict

class TreeNode(object):

    @classmethod
    def serialize(cls, node):
        if not node:
            return None
        data = OrderedDict()
        data["id"] = node.id
        data["cycle"] = node.cycle
        data["children"] = []
        for c in node.children:
            child = cls.serialize(c)
            data["children"].append(child)
        return data

    @classmethod
    def deserialize(cls, data):
        if not data:
            return None
        node = TreeNode(data["id"], data["cycle"])
        for c in data["children"]:
            child = cls.deserialize(c)
            node.append_child(child)
        return node

    def __init__(self, id, cycle=1):
        self.id = id
        self.cycle = cycle
        self.children = []
        self.update_digest()

    def append_child(self, child):
        assert isinstance(child, TreeNode)
        self.children.append(child)
        self.update_digest()

    def update_digest(self):
        m = hashlib.md5()
        m.update(str(self.id))
        for c in self.children:
            m.update(str(c.digest))
        self.digest = m.digest()

    def __repr__(self):
        result = []
        result.append(" %s %s %r" % (self.id, self.cycle, str(self.digest)))
        for c in self.children:
            child_repr = repr(c).split("\n")
            result.extend(map(lambda x: "-%s" % x, child_repr))
        return "\n".join(result)

    def __eq__(self, x):
        return self.digest == x.digest

    def __ne__(self, x):
        return self.digest != x.digest

    def __le__(self, x):
        if self.id != x.id:
            return False
        me = [c.digest for c in self.children]
        other = [c.digest for c in x.children]
        index = 0
        for item in me:
            if item not in other[index:]:
                return False
            index = other.index(item, index)
        return True

    def __lt__(self, x):
        if len(self.children) >= len(x.children):
            return False
        return self <= x

    def __ge__(self, x):
        if self.id != x.id:
            return False
        me = [c.digest for c in self.children]
        other = [c.digest for c in x.children]
        index = 0
        for item in other:
            if item not in me[index:]:
                return False
            index = me.index(item, index)
        return True

    def __gt__(self, x):
        if len(self.children) <= len(x.children):
            return False
        return self <= x

def fold_tree(tree, keep_level=None, remove_calls=None):

    def fold_tree_recursive(tree):
        if not tree:
            return None
        folded_children = []
        for c in tree.children:
            folded_children.append(fold_tree_recursive(c))
        new_children = []
        if folded_children:
            new_children = [folded_children[0]]
            for i in xrange(1, len(folded_children)):
                if folded_children[i] != new_children[-1]:
                    new_children.append(folded_children[i])
        new_tree = OrderedDict()
        new_tree = TreeNode(tree.id, tree.cycle)
        for c in new_children:
            new_tree.append_child(c)
        return new_tree

    def cut_tree_recursive(tree, curr_level, max_level):
        if not tree:
            return None
        if curr_level >= max_level:
            return None
        children = []
        for c in tree.children:
            child = cut_tree_recursive(c, curr_level + 1, max_level)
            if child:
                children.append(child)
        new_tree = TreeNode(tree.id, tree.cycle)
        for c in children:
            new_tree.append_child(c)
        return new_tree

    def remove_node_recursive(tree, patterns):
        if not tree:
            return None
        for c in patterns:
            if fnmatch.fnmatch(tree.id, c):
                return None
        children = []
        for c in tree.children:
            child = remove_node_recursive(c, patterns)
            if child:
                children.append(child)
        new_tree = TreeNode(tree.id, tree.cycle)
        for c in children:
            new_tree.append_child(c)
        return new_tree

    if remove_calls:
        tree = remove_node_recursive(tree, remove_calls)
    if keep_level:
        tree = cut_tree_recursive(tree, 0, keep_level)

    return fold_tree_recursive(tree)

def print_tree(tree, max_level=None):
    def print_tree_recursive(tree, level, max_level):
        if max_level and level > max_level:
            return
        print "--" * level, tree.id, tree.cycle
        for c in tree.children:
            print_tree_recursive(c, level + 1, max_level)
    print_tree_recursive(tree, 1, max_level)

def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("call_tree",
                        help="Call tree file (json) generated by JASMIN")
    parser.add_argument("--save", help="Output filename", default=None)
    parser.add_argument("--max-depth", default=None, type=int,
                        help="Max depth for output tree")
    parser.add_argument("--keep-level", default=None, type=int,
                        help="Max levels to keep when folding tree")
    parser.add_argument("--remove-calls", default=[], nargs="+",
                        help="Remove matched calls, supports shell wildcards")
    parser.add_argument("--fold", action="store_true",
                        help="Try hard to fold tree")

    args = parser.parse_args()

    data = json.load(file(args.call_tree), object_hook=OrderedDict)
    tree = TreeNode.deserialize(data)
    if tree.id == "ROOT":
        tree = tree.children[0]
    if args.fold:
        tree = fold_tree(tree, args.keep_level, args.remove_calls)
    if args.save:
        data = TreeNode.serialize(tree)
        json.dump(data, file(args.save, "w"), indent=2)
    print_tree(tree, args.max_depth)

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