OP_IN = 1 OP_OUT = 3 OP_ROOT = 2 OP_UP = 0 class TreeStructure(object):
    """Topological data representing a tree on features."""

    def __init__(self, num_vertices):
        logger.debug('TreeStructure with %d vertices', num_vertices)
        self._num_vertices = num_vertices
        self._num_edges = num_vertices - 1
        self.set_edges([(v, v + 1) for v in range(num_vertices - 1)])
        self._complete_grid = None  # Lazily constructed.
        self._vertices = np.arange(num_vertices, dtype=np.int32)

    def __eq__(self, other):
        return (self._num_vertices == other._num_vertices and  #
                (self._tree_grid == other._tree_grid).all())

    def set_edges(self, edges):
        """Sets the edges of this tree.

        Args:
            edges: A list of (vertex, vertex) pairs.
        """
        assert len(edges) == self._num_edges
        self._tree_grid = make_tree(edges)

    def get_edges(self):
        """Returns the edges of this tree.

        Returns:
            A list of (vertex, vertex) pairs.
        """
        return [(v1, v2) for e, v1, v2 in self._tree_grid.T]

    @property
    def num_vertices(self):
        return self._num_vertices

    @property
    def num_edges(self):
        return self._num_edges

    @property
    def tree_grid(self):
        """Array of (edge, vertex, vertex) triples defining the tree graph."""
        return self._tree_grid

    @property
    def complete_grid(self):
        """Array of (edge,vertex,vertex) triples defining a complete graph."""
        if self._complete_grid is None:
            self._complete_grid = make_complete_graph(self._num_vertices)
        return self._complete_grid

    @property
    def vertices(self):
        return self._vertices

    def gc(self):
        """Garbage collect temporary cached data structures."""
        self._complete_grid = None @profile
def estimate_tree(grid, edge_logits):
    """Compute a maximum likelihood spanning tree of a dense weighted graph.

    Args:
        grid: A 3 x K array as returned by make_complete_graph().
        edge_logits: A length-K array of nonnormalized log probabilities.

    Returns:
        A list of (vertex, vertex) pairs.
    """
    K = len(edge_logits)
    assert grid.shape == (3, K)
    weights = triangular_to_square(grid, edge_logits)
    weights *= -1
    weights -= weights.min()
    weights += 1.0
    csr = minimum_spanning_tree(weights, overwrite=True)
    coo = csr.tocoo()
    edges = zip(coo.row, coo.col)
    edges = sorted(tuple(sorted(pair)) for pair in edges)
    assert len(edges) == weights.shape[0] - 1
    return edges def make_propagation_program(grid, root=None):
    """Makes an efficient program for message passing on a tree.

    This creates a program of instructions to be intrepreted by various
    virtual machines. Each 8-byte instruction is broken into 4 2-byte chunks:

      [ op_code | vertex | relative_vertex | relative_edge ]

    where the relative_* operands are optional.
    The five instructions are (in order of occurrence at each vertex):

      OP_UP: Propagate upwards from observed state to latent state.
          vertex: The current vertex.
          relative_vertex: Not set.
          relative_edge: Not set.

      OP_IN: Propagate inwards from latent leaves towards latent root,
          assuming OP_UP has already been called on this vertex.
          vertex: The target parent vertex.
          relative_vertex: The source child vertex.
          relative_edge: The edge between parent and child.

      OP_ROOT: Process the root node after OP_UP and OP_IN have been called.
          vertex: The root vertex.
          relative_vertex: Not set.
          relative_edge: Not set.

      OP_OUT: Popagate outwards from the latent root towards latent leaves,
          assuming OP_UP and OP_IN have been called as needed.
          vertex: The target child vertex.
          relative_vertex: The source parent vertex.
          relative_edge: The edge between parent and child.

      OP_DOWN: Propagate downwards from latent state to observed state,
          assuming OP_UP, OP_IN, and OP_OUT have been called as needed.
          vertices and OP_IN has been called on all edges.
          vertex: The current vertex.
          relative_vertex: Not set.
          relative_edge: Not set.

    Args:
        grid: A tree graph as returned by make_tree().
        root: Optional root vertex, defaults to find_center_of_tree(grid).

    Returns:
        A [V+E+1+E+V, 4]-shaped numpy array whose rows are instructions.
    """
    if root is None:
        root = find_center_of_tree(grid)
    E = grid.shape[1]
    V = 1 + E
    neighbors = [set() for _ in range(V)]
    edge_dict = {}
    for e, v1, v2 in grid.T:
        neighbors[v1].add(v2)
        neighbors[v2].add(v1)
        edge_dict[v1, v2] = e
        edge_dict[v2, v1] = e

    # Construct a nested program.
    nested_program = []
    queue = deque()
    queue.append((root, None))
    while queue:
        v, parent = queue.popleft()
        nested_program.append((v, parent, []))
        for v2 in sorted(neighbors[v]):
            if v2 != parent:
                queue.append((v2, v))
    for v, parent, children in nested_program:
        for v2 in sorted(neighbors[v]):
            if v2 != parent:
                children.append(v2)

    # Construct a flattened program.
    program = np.zeros([V + E + 1 + E + V, 4], np.int16)
    pos = 0
    for v, parent, children in reversed(nested_program):
        program[pos, :] = [OP_UP, v, 0, 0]
        pos += 1
        for child in children:
            program[pos, :] = [OP_IN, v, child, edge_dict[v, child]]
            pos += 1
    program[pos, :] = [OP_ROOT, root, 0, 0]
    pos += 1
    for v, parent, children in nested_program:
        if parent is not None:
            program[pos, :] = [OP_OUT, v, parent, edge_dict[v, parent]]
            pos += 1
        program[pos, :] = [OP_DOWN, v, 0, 0]
        pos += 1
    assert pos == program.shape[0]

    return program @profile
def sample_tree(grid, edge_logits, edges):
    """Sample a random spanning tree of a dense weighted graph using MCMC.

    This uses Gibbs sampling on edges. Consider E undirected edges that can
    move around a graph of V=1+E vertices. The edges are constrained so that no
    two edges can span the same pair of vertices and so that the edges must
    form a spanning tree. To Gibbs sample, chose one of the E edges at random
    and move it anywhere else in the graph. After we remove the edge, notice
    that the graph is split into two connected components. The constraints
    imply that the edge must be replaced so as to connect the two components.
    Hence to Gibbs sample, we collect all such bridging (vertex,vertex) pairs
    and sample from them in proportion to exp(edge_logits).

    Args:
        grid: A 3 x K array as returned by make_complete_graph().
        edge_logits: A length-K array of nonnormalized log probabilities.
        edges: A list of E initial edges in the form of (vertex,vertex) pairs.

    Returns:
        A list of (vertex, vertex) pairs.
    """
    logger.debug('sample_tree sampling a random spanning tree')
    COUNTERS.sample_tree_calls += 1
    if len(edges) <= 1:
        return edges
    E = len(edges)
    V = E + 1
    K = V * (V - 1) // 2
    e2k = np.zeros(E, np.int32)
    neighbors = np.zeros((V, V), np.int16)
    components = np.zeros(V, np.bool_)
    for e in range(E):
        v1, v2 = edges[e]
        e2k[e] = find_complete_edge(v1, v2)
        jit_set_add(neighbors[v1], v2)
        jit_set_add(neighbors[v2], v1)
    valid_edges = np.empty(K, np.int32)

    accepted = 0
    log2_choices = []
    for e in np.random.permutation(E):  # Sequential scanning doesn't work.
        k1 = jit_remove_edge(grid, e2k, neighbors, components, e)
        num_valid_edges = find_valid_edges(components, valid_edges)
        valid_probs = edge_logits[valid_edges[:num_valid_edges]]
        valid_probs -= valid_probs.max()
        valid_probs = np.exp(valid_probs)
        total_prob = valid_probs.sum()
        if total_prob > 0:
            k2 = valid_edges[sample_from_probs(valid_probs)]
        else:
            k2 = k1
            COUNTERS.sample_tree_infeasible += 1
        jit_add_edge(grid, e2k, neighbors, components, e, k2)

        accepted += (k1 != k2)
        log2_choices.append(num_valid_edges.bit_length())
    COUNTERS.sample_tree_accept += accepted
    COUNTERS.sample_tree_propose += E
    SERIES.sample_tree_accepted.append(accepted)
    HISTOGRAMS.sample_tree_log2_choices.update(log2_choices)

    edges = sorted((grid[1, k], grid[2, k]) for k in e2k)
    assert len(edges) == E
    return edges TY_MULTINOMIAL = 0 class Table(object):
    """Internal representation of read-only row-oriented heterogeneous data."""

    def __init__(
            self,
            feature_types,
            ragged_index,
            data,
            categorical_data=None, ):
        """Create a read-only table of heterogeneous data.

        Args:
            feature_types: A [V]-shaped numpy array of datatype ids.
            ragged_index: A [V+1]-shaped numpy array of indices into the ragged
                data array.
            data: An [N, _]-shaped numpy array of ragged data, where the vth
                column is stored in data[:, ragged_index[v]:ragged_index[v+1]].
            categorical_data: An [N, Vc]-shaped numpy array of category
                indices, where Vc is the number of categorical features and N
                is the number fo rows.
        """
        N = data.shape[0]
        V = len(feature_types)

        feature_types = np.asarray(feature_types, dtype=np.int8)
        feature_types.flags.writeable = False
        assert len(feature_types.shape) == 1

        # Multinomial features.
        ragged_index = np.asarray(ragged_index, np.int32)
        ragged_index.flags.writeable = False
        assert len(ragged_index.shape) == 1
        assert ragged_index.shape[0] == 1 + V
        data = np.asarray(data, np.int8)
        data.flags.writeable = False
        assert len(data.shape) == 2
        assert data.shape[0] == N
        assert data.shape[1] == ragged_index[-1]

        # Categorical features.
        if categorical_data is None:
            categorical_data = np.empty([N, 0], dtype=np.int16)
        categorical_data = np.asarray(categorical_data, dtype=np.int16)
        assert len(categorical_data.shape) == 2
        assert categorical_data.shape[0] == N

        self._feature_types = feature_types
        self._multinomial_ragged_index = ragged_index
        self._multinomial_data = data
        self._categorical_data = categorical_data

    @property
    def num_rows(self):
        return self._multinomial_data.shape[0]

    @property
    def num_cols(self):
        return self._feature_types.shape[0]

    @property
    def feature_types(self):
        return self._feature_types

    @property
    def ragged_index(self):
        return self._multinomial_ragged_index

    @property
    def data(self):
        return self._multinomial_data

    @property
    def categorical_data(self):
        return self._categorical_data SERIES = ProfilingSet(list) def TODO(message=''):
    raise NotImplementedError('TODO {}'.format(message)) TREECAT_JIT = int(os.environ.get('TREECAT_JIT', 1))
TREECAT_THREADS = int(
    os.environ.get('TREECAT_THREADS', multiprocessing.cpu_count()))
PROFILING = int(os.environ.get('TREECAT_PROFILE', 0))
LOG_LEVEL = int(os.environ.get('TREECAT_LOG_LEVEL', logging.CRITICAL))
LOG_FILENAME = os.environ.get('TREECAT_LOG_FILE')
LOG_FORMAT = '%(levelname).1s %(process)d %(name)s %(message)s'
    N, M = probs.shape
    V = ragged_index.shape[0] - 1
    R = ragged_index[-1]
    N, R = data.shape
    V = len(ragged_index) - 1
POOL = None
        POOL = multiprocessing.Pool(TREECAT_THREADS)
TIMERS = defaultdict(ProfileTimer)
COUNTERS = ProfilingSet(lambda: 0)
HISTOGRAMS = ProfilingSet(Counter)
SERIES = ProfilingSet(list)
def TODO(message=''):
def no_jit(*args, **kwargs):
def jit_random_seed(seed):
def set_random_seed(seed):
def np_printoptions(**kwargs):
def sample_from_probs(probs):
def sample_from_probs2(probs, out=None):
def quantize_from_probs2(probs, resolution):
def make_ragged_index(columns):
def make_ragged_mask(ragged_index, mask):
def count_observations(ragged_index, data):
def guess_counts(ragged_index, data):
def parallel_map(fun, args):
    def __init__(self):
    def __enter__(self):
    def __exit__(self, type, value, traceback):
def profile_timed(fun):
    def profiled_fun(*args, **kwargs):
def get_profiling_stats():
def log_profiling_stats():
class ProfilingSet(defaultdict):
class ProfileTimer(object): def parallel_map(fun, args):
    global POOL
    args = list(args)
    if len(args) < 2 or TREECAT_THREADS == 1:
        return list(map(fun, args))
    if POOL is None:
        POOL = multiprocessing.Pool(TREECAT_THREADS)
    return POOL.map(fun, args) TREECAT_JIT = int(os.environ.get('TREECAT_JIT', 1))
TREECAT_THREADS = int(
    os.environ.get('TREECAT_THREADS', multiprocessing.cpu_count()))
PROFILING = int(os.environ.get('TREECAT_PROFILE', 0))
LOG_LEVEL = int(os.environ.get('TREECAT_LOG_LEVEL', logging.CRITICAL))
LOG_FILENAME = os.environ.get('TREECAT_LOG_FILE')
LOG_FORMAT = '%(levelname).1s %(process)d %(name)s %(message)s'
    N, M = probs.shape
    V = ragged_index.shape[0] - 1
    R = ragged_index[-1]
    N, R = data.shape
    V = len(ragged_index) - 1
POOL = None
        POOL = multiprocessing.Pool(TREECAT_THREADS)
TIMERS = defaultdict(ProfileTimer)
COUNTERS = ProfilingSet(lambda: 0)
HISTOGRAMS = ProfilingSet(Counter)
SERIES = ProfilingSet(list)
def TODO(message=''):
def no_jit(*args, **kwargs):
def jit_random_seed(seed):
def set_random_seed(seed):
def np_printoptions(**kwargs):
def sample_from_probs(probs):
def sample_from_probs2(probs, out=None):
def quantize_from_probs2(probs, resolution):
def make_ragged_index(columns):
def make_ragged_mask(ragged_index, mask):
def count_observations(ragged_index, data):
def guess_counts(ragged_index, data):
def parallel_map(fun, args):
    def __init__(self):
    def __enter__(self):
    def __exit__(self, type, value, traceback):
def profile_timed(fun):
    def profiled_fun(*args, **kwargs):
def get_profiling_stats():
def log_profiling_stats():
class ProfilingSet(defaultdict):
class ProfileTimer(object): TREECAT_JIT = int(os.environ.get('TREECAT_JIT', 1))
TREECAT_THREADS = int(
    os.environ.get('TREECAT_THREADS', multiprocessing.cpu_count()))
PROFILING = int(os.environ.get('TREECAT_PROFILE', 0))
LOG_LEVEL = int(os.environ.get('TREECAT_LOG_LEVEL', logging.CRITICAL))
LOG_FILENAME = os.environ.get('TREECAT_LOG_FILE')
LOG_FORMAT = '%(levelname).1s %(process)d %(name)s %(message)s'
    N, M = probs.shape
    V = ragged_index.shape[0] - 1
    R = ragged_index[-1]
    N, R = data.shape
    V = len(ragged_index) - 1
POOL = None
        POOL = multiprocessing.Pool(TREECAT_THREADS)
TIMERS = defaultdict(ProfileTimer)
COUNTERS = ProfilingSet(lambda: 0)
HISTOGRAMS = ProfilingSet(Counter)
SERIES = ProfilingSet(list)
def TODO(message=''):
def no_jit(*args, **kwargs):
def jit_random_seed(seed):
def set_random_seed(seed):
def np_printoptions(**kwargs):
def sample_from_probs(probs):
def sample_from_probs2(probs, out=None):
def quantize_from_probs2(probs, resolution):
def make_ragged_index(columns):
def make_ragged_mask(ragged_index, mask):
def count_observations(ragged_index, data):
def guess_counts(ragged_index, data):
def parallel_map(fun, args):
    def __init__(self):
    def __enter__(self):
    def __exit__(self, type, value, traceback):
def profile_timed(fun):
    def profiled_fun(*args, **kwargs):
def get_profiling_stats():
def log_profiling_stats():
class ProfilingSet(defaultdict):
class ProfileTimer(object): @jit(nopython=True, cache=True)
def sample_from_probs(probs):
    """Sample from a vector of non-normalized probabilitites.

    Args:
        probs: An [M]-shaped numpy array of non-normalized probabilities.

    Returns:
        An integer in range(M).
    """
    cdf = probs.cumsum()
    return (np.random.rand() * cdf[-1] < cdf).argmax() def set_random_seed(seed):
    """Set random seeds for both numpy and numba."""
    np.random.seed(seed)
    jit_random_seed(seed)
import itertools
import logging
import numpy as np
from abc import ABCMeta
from abc import abstractmethod
from scipy.special import gammaln
from six import add_metaclass
from six.moves import range
from treecat.structure import OP_IN
from treecat.structure import OP_OUT
from treecat.structure import OP_ROOT
from treecat.structure import OP_UP
from treecat.structure import TreeStructure
from treecat.structure import estimate_tree
from treecat.structure import make_propagation_program
from treecat.structure import sample_tree
from treecat.tables import TY_MULTINOMIAL
from treecat.tables import Table
from treecat.util import SERIES
from treecat.util import TODO
from treecat.util import jit
from treecat.util import parallel_map
from treecat.util import prange
from treecat.util import profile
from treecat.util import sample_from_probs
from treecat.util import set_random_seed
        """
        logger.info('TreeTrainer.train')
        set_random_seed(self._config['seed'])
        init_epochs = self._config['learning_init_epochs']
        full_epochs = self._config['learning_full_epochs']
        sample_tree_rate = self._config['learning_sample_tree_rate']
        num_rows = self._num_rows

        # Initialize using subsample annealing.
        assert len(self._added_rows) == 0
        schedule = make_annealing_schedule(num_rows, init_epochs,
                                           sample_tree_rate)
        for action, row_id in schedule:
            if action == 'add_row':
                self.add_row(row_id)
            elif action == 'remove_row':
                self.remove_row(row_id)
            elif action == 'sample_tree':
                edges, edge_logits = self.sample_tree()
                self.set_edges(edges)
            else:
                raise ValueError(action)

        # Run full gibbs scans.
        assert len(self._added_rows) == num_rows
        for step in range(full_epochs):
            edges, edge_logits = self.sample_tree()
            self.set_edges(edges)
            for row_id in range(num_rows):
                self.remove_row(row_id)
                self.add_row(row_id)

        # Compute optimal tree.
        assert len(self._added_rows) == num_rows
        edges, edge_logits = self.estimate_tree()
        if self._config['learning_estimate_tree']:
            self.set_edges(edges)

        self._tree.gc()

        return {
            'config': self._config,
            'tree': self._tree,
            'edge_logits': edge_logits,
        }


@jit(nopython=True, cache=True)
def treecat_add_cell(
        feature_type,
        ragged_index,
        data_row,
        message,
        feat_probs,
        meas_probs,
        v, ):
    if feature_type == TY_MULTINOMIAL:
        beg, end = ragged_index[v:v + 2]
        feat_block = feat_probs[beg:end, :]
        meas_block = meas_probs[v, :]
        for c, count in enumerate(data_row[beg:end]):
            for _ in range(count):
                message *= feat_block[c, :]
                message /= meas_block
                feat_block[c, :] += np.float32(1)
                meas_block += np.float32(1)
    else:
        raise NotImplementedError


@profile
@jit(nopython=True, cache=True)
def treecat_add_row(
        feature_types,
        ragged_index,
        data_row,
        tree_grid,
        program,
        assignments,
        vert_ss,
        edge_ss,
        feat_ss,
        meas_ss,
        vert_probs,
        edge_probs,
        feat_probs,
        meas_probs, ):
    # Sample latent assignments using dynamic programming.
    messages = vert_probs.copy()
    for i in range(len(program)):
        op, v, v2, e = program[i]
        message = messages[v, :]
        if op == OP_UP:
            # Propagate upward from observed to latent.
            treecat_add_cell(
                feature_types[v],
                ragged_index,
                data_row,
                message,
                feat_probs,
                meas_probs,
                v, )
        elif op == OP_IN:
            # Propagate latent state inward from children to v.
            trans = edge_probs[e, :, :]
            if v > v2:
                trans = trans.T
            message *= np.dot(trans, messages[v2, :] / vert_probs[v2, :])
            message /= vert_probs[v, :]
            message /= message.max()  # Scale for numerical stability.
        elif op == OP_ROOT:
            # Process root node.
            assignments[v] = sample_from_probs(message)
        elif op == OP_OUT:
            # Propagate latent state outward from parent to v.
            trans = edge_probs[e, :, :]
            if v2 > v:
                trans = trans.T
            message *= trans[assignments[v2], :]
            message /= vert_probs[v, :]

            assignments[v] = sample_from_probs(message)