<commit_before>'''
Implementation of factor graph.

Author: mbforbes

Current approach (order matters):
-   (1) add RVs
-   (2) add factors to connect them
-   (3) set potentials 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
# -----------------------------------------------------------------------------

# Builtins
import code  # code.interact(local=dict(globals(), **locals()))

# 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

# This is the maximum number of iterations that we let loopy belief propagation
# run before cutting it off.
LBP_MAX_ITERS = 50


# 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.

    NOTE: All RVs must have unique names.

    TODO: Consider making Node base class which RV and Factor extend.

    TODO: convenience functions or modifications to consider (not worth making
    unless I need them):
        - getters (and setters?) for RVs and Factors
    '''

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

        # added later
        self._rvs = {}
        # TODO: Consider making dict for speed.
        self._factors = []

    # TODO(mbforbes): Learn about *args or **args or whatever and see whether I
    #                 can use here to clean this up.
    def rv(self, name, n_opts, labels=[], debug=DEBUG_DEFAULT):
        '''
        Creates an RV, adds it to this graph, and returns it. Convenience
        function.

        Args:
            name (str)                must be globally unique w.r.t. other RVs
            n_opts (int)              how many values it can take
            labels ([str], opt)       opt names for each var. len == n_opts
            debug (bool, opt)         a gazillion asserts

        Returns:
            RV
        '''
        rv = RV(name, n_opts, labels, debug)
        self.add_rv(rv)
        return rv

    def has_rv(self, rv_s):
        '''
        Args:
            rv_s (str): Potential name of RV

        Returns:
            bool
        '''
        return rv_s in self._rvs

    def add_rv(self, rv):
        '''
        Args:
            rv (RV)
        '''
        # 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

    # TODO(mbforbes): Learn about *args or **args or whatever and see whether I
    #                 can use here to clean this up.
    def factor(self, rvs, name='', potential=None, debug=DEBUG_DEFAULT):
        '''
        Creates a Factor, adds it to this graph, and returns it. Convenience
        function.

        Note that you can provide the name of an RV instead of the RV if you
        like. And you can mix and match. Wow!

        Args:<commit_msg>add support for pruning loner RVs. helpful before LBP OMG.<commit_after>'''
Implementation of factor graph.

Author: mbforbes

Current approach (order matters):
-   (1) add RVs
-   (2) add factors to connect them
-   (3) set potentials 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
# -----------------------------------------------------------------------------

# Builtins
import code  # code.interact(local=dict(globals(), **locals()))

# 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

# This is the maximum number of iterations that we let loopy belief propagation
# run before cutting it off.
LBP_MAX_ITERS = 50


# 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.

    NOTE: All RVs must have unique names.

    TODO: Consider making Node base class which RV and Factor extend.

    TODO: convenience functions or modifications to consider (not worth making
    unless I need them):
        - getters (and setters?) for RVs and Factors
    '''

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

        # added later
        self._rvs = {}
        # TODO: Consider making dict for speed.
        self._factors = []

    # TODO(mbforbes): Learn about *args or **args or whatever and see whether I
    #                 can use here to clean this up.
    def rv(self, name, n_opts, labels=[], debug=DEBUG_DEFAULT):
        '''
        Creates an RV, adds it to this graph, and returns it. Convenience
        function.

        Args:
            name (str)                must be globally unique w.r.t. other RVs
            n_opts (int)              how many values it can take
            labels ([str], opt)       opt names for each var. len == n_opts
            debug (bool, opt)         a gazillion asserts

        Returns:
            RV
        '''
        rv = RV(name, n_opts, labels, debug)
        self.add_rv(rv)
        return rv

    def has_rv(self, rv_s):
        '''
        Args:
            rv_s (str): Potential name of RV

        Returns:
            bool
        '''
        return rv_s in self._rvs

    def add_rv(self, rv):
        '''
        Args:
            rv (RV)
        '''
        # 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 remove_loner_rvs(self):
        '''
        Removes RVs from the graph that have no factors attached to them.

        Returns:
            int number removed
        '''
        removed = 0
        names = self._rvs.keys()
        for name in names:
            if self._rvs[name].n_edges() == 0:
                del self._rvs[name]
                removed += 1
        return removed

    # TODO(mbforbes): Learn about *args or **args or whatever and see whether I
    #                 can use here to clean this up.
    def factor(self, rvs, name='', potential=None, debug=DEBUG_DEFAULT):
        '''
        Creates a Factor, adds it to this graph, and returns it. Convenience
        function.

        Note that you can provide the name of an RV instead of the RV if you
        like. And you can mix and match. Wow!

        Args:
{'fraction_non_alphanumeric': 0.12098660429513076, 'fraction_numerical': 0.0032957686583032107, 'mean_word_length': 1.7180005778676684, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '707823', 'n_tokens_mistral': 2693, 'n_tokens_neox': 2565, 'n_words': 1155}