<commit_before>import pandas as pd
import numpy as np
from util import cartesian
from itertools import product


class Variable(object):
    """A discrete variable, with a distribution over N states.

    The variable can be assigned a distribution if is being used as input, or
    as an intervention. Otherwise the distribution is calculated.
    """

    USED_NAMES = set()

    def __init__(self, name, n_states):
        """Name the variable, and say how many states it has. Variables start
        off as unassigned.
        """
        assert str(name) == name
        assert name.isalnum()
        # Don't recreate variables
        assert name not in Variable.USED_NAMES

        self.name = name
        Variable.USED_NAMES.add(name)

        # Generate the actual states; this makes it easy to work with.
        self.n_states = n_states
        self.states = range(n_states)

    def uniform(self):
        return np.ones(self.n_states) / float(self.n_states)

    def with_state(self, state):
        assert state in self.states
        dist = np.zeros(self.n_states)
        dist[state] = 1.0
        return dist

    def make_valid_distribution(self, distn):
        """Various checks to make sure nothing silly is happening"""
        np_distn = np.array(distn, dtype=float)
        assert np_distn.shape == (self.n_states,)
        assert np.isclose(np_distn.sum(), 1.0)
        return np_distn

    def __repr__(self):
        return '<{}>'.format(self.name)


def make_variables(strings, n_states):
    """Just a shortcut for making lots of variables"""
    varnames = strings.split()
    return [Variable(v, n_states) for v in varnames]


class Distribution(object):
    """A distribution over one or more variables
    """
    P_LABEL = 'Pr'

    def __init__(self, variables, pr=None):
        # Make a copy
        for v in variables:
            assert isinstance(v, Variable)
        # No duplicates!
        assert len(set(variables)) == len(variables)

        self.variables = list(variables)
        self.names = [v.name for v in variables]
        self.probabilities = pr

    def joint(self, *variables):
        """Generate the (sub)joint distribution over N variables

        Note: This uses pandas' MultiIndex to generate the joint distribution
        any number of variables together from the existing `probabilities`.
        """
        for v in variables:
            assert isinstance(v, Variable)
            assert v in self.variables

        # This is the function that does all the work. It is amazingly
        # flexible.
        pr = pd.pivot_table(
            # We need to treat the variables as columns not indexes for this
            # to work!
            self.probabilities.reset_index(),
            values=[Distribution.P_LABEL],
            index=[v.name for v in variables],
            aggfunc=np.sum,
        )

        # Return another (smaller?) distribution.
        return Distribution(variables, pr)

    def prob_of_state(self, assignments):
        assert set(assignments.keys()) == set(self.variables)
        # Construct a tuple from the assignments
        by_name = dict([(v.name, val) for v, val in assignments.items()])
        state = tuple([by_name[nm] for nm in self.probabilities.index.names])
        return self.probabilities.loc[state].values[0]

    def iter_conditional(self, a, b):
        # Build a lookup for b
        jb = self.joint(b)
        jab = self.joint(a, b)
        for ass, p in jab.iter_assignments():
            yield ass, p / jb.prob_of_state({b: ass[b]})

    def iter_assignments(self):
        # Indexes are returned as single values when there is only one..
        is_single = len(self.variables) == 1
        for indexes, columns in self.probabilities.iterrows():
            if is_single:
                assignments = {self.variables[0]: indexes}
            else:
                assignments = dict([(v, val) for v, val in zip(self.variables, indexes)])
            pr = columns[Distribution.P_LABEL]
            yield assignments, pr

    def entropy(self, *variables):
        return self._calc_entropy(self.joint(*variables).probabilities)

    def mutual_info(self, v1, v2, v3=None):
        """calculate the mutual (or conditional mutual) information.

        the simplest way to do this is to calculate the entropy of various
        joint distributions and bung them together. see here:

        https://en.wikipedia.org/wiki/mutual_information
        https://en.wikipedia.org/wiki/conditional_mutual_information<commit_msg>Documentation updates<commit_after>import pandas as pd
import numpy as np
from util import cartesian
from itertools import product


class Variable(object):
    """A discrete variable, with a distribution over N states.
    """

    USED_NAMES = set()

    def __init__(self, name, n_states):
        """Name the variable, and say how many states it has. Variables start
        off as unassigned.
        """
        assert isinstance(name, str)
        assert name.isalnum()
        # Names must uniquely identify the variables (as they're used to
        # identify the columns in our tables).
        assert name not in Variable.USED_NAMES

        self.name = name
        Variable.USED_NAMES.add(name)

        self.n_states = n_states

        # Generate the actual states; this makes it easy to work with.
        self.states = range(n_states)

    def uniform(self):
        """Return a uniform distribution for this variable."""
        return np.ones(self.n_states) / float(self.n_states)

    def with_state(self, state):
        """Return a distribution with just this state set."""
        assert state in self.states
        dist = np.zeros(self.n_states)
        dist[state] = 1.0
        return dist

    def make_valid_distribution(self, distn):
        """Convert distribution and check it."""
        np_distn = np.array(distn, dtype=float)
        assert np_distn.shape == (self.n_states,)
        assert np.isclose(np_distn.sum(), 1.0)
        return np_distn

    def __repr__(self):
        return '<{}>'.format(self.name)


def make_variables(strings, n_states):
    """Just a shortcut for making lots of variables"""
    varnames = strings.split()
    return [Variable(v, n_states) for v in varnames]


class Distribution(object):
    """Base class for a distribution over one or more variables
    """
    P_LABEL = 'Pr'

    def __init__(self, variables, pr=None):
        # Everything should be a variable
        for v in variables:
            assert isinstance(v, Variable)
        # No duplicates!
        assert len(set(variables)) == len(variables)

        self.variables = list(variables)
        self.names = [v.name for v in variables]
        self.probabilities = pr

    def joint(self, *variables):
        """Generate the (sub)joint distribution over included variables

        Note: This uses pandas' MultiIndex to generate the joint distribution
        any number of variables together from the existing probabilities.
        """
        for v in variables:
            assert isinstance(v, Variable)
            assert v in self.variables

        # This is the function that does all the work. It is amazingly
        # flexible.
        pr = pd.pivot_table(
            # We need to treat the variables as columns not indexes for this
            # to work!
            self.probabilities.reset_index(),
            values=[Distribution.P_LABEL],
            index=[v.name for v in variables],
            aggfunc=np.sum,
        )

        # Return another (smaller?) distribution.
        return Distribution(variables, pr)

    def prob_of_state(self, assignments):
        assert set(assignments.keys()) == set(self.variables)
        # Construct a tuple from the assignments
        by_name = dict([(v.name, val) for v, val in assignments.items()])
        state = tuple([by_name[nm] for nm in self.probabilities.index.names])
        return self.probabilities.loc[state].values[0]

    def iter_conditional(self, a, b):
        # Build a lookup for b
        jb = self.joint(b)
        jab = self.joint(a, b)
        for ass, p in jab.iter_assignments():
            yield ass, p / jb.prob_of_state({b: ass[b]})

    def iter_assignments(self):
        # Indexes are returned as single values when there is only one..
        is_single = len(self.variables) == 1
        for indexes, columns in self.probabilities.iterrows():
            if is_single:
                assignments = {self.variables[0]: indexes}
            else:
                assignments = dict([(v, val) for v, val in zip(self.variables, indexes)])
            pr = columns[Distribution.P_LABEL]
            yield assignments, pr

    def entropy(self, *variables):
        """Calculate the entropy of one or more variables in this distribution."""
        return self._calc_entropy(self.joint(*variables).probabilities)

    def mutual_info(self, v1, v2, v3=None):
        """calculate the mutual (or conditional mutual) information.

        the simplest way to do this is to calculate the entropy of various
        joint distributions and bung them together. see here:

        https://en.wikipedia.org/wiki/mutual_information
        https://en.wikipedia.org/wiki/conditional_mutual_information
{'fraction_non_alphanumeric': 0.09125229209362529, 'fraction_numerical': 0.0021572645885017797, 'mean_word_length': 2.390127970749543, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 4, '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': '2450574', 'n_tokens_mistral': 2584, 'n_tokens_neox': 2472, 'n_words': 864}