def generate_dataset(num_rows, num_cols, num_cats=4, rate=1.0):
    """Generate a random dataset.

    Returns:
        A dataset dict with fields 'schema' and 'table'.
    """
    set_random_seed(0)
    N = num_rows
    V = num_cols
    K = V * (V - 1) // 2
    ragged_index = np.arange(0, num_cats * (V + 1), num_cats, np.int32)
    ragged_index.flags.writeable = False
    data = np.zeros((N, V * num_cats), np.int8)
    for v in range(V):
        beg, end = ragged_index[v:v + 2]
        column = data[:, beg:end]
        probs = np.random.dirichlet(np.zeros(num_cats) + 0.5)
        for n in range(N):
            count = np.random.poisson(rate)
            column[n, :] = np.random.multinomial(count, probs)
    data.flags.writeable = False
    feature_types = [TY_MULTINOMIAL] * V
    table = Table(feature_types, ragged_index, data)
    dataset = {
        'schema': {
            'ragged_index': ragged_index,
            'tree_prior': np.zeros(K, np.float32),
        },
        'table': table,
    }
    return dataset def generate_fake_ensemble(num_rows, num_cols, num_cats, num_components):
    dataset = generate_dataset(num_rows, num_cols, num_cats)
    ensemble = []
    config = make_config(model_num_clusters=num_components, seed=0)
    for sub_seed in range(3):
        sub_config = config.copy()
        sub_config['seed'] += sub_seed
        set_random_seed(sub_config['seed'])
        model = generate_fake_model(num_rows, num_cols, num_cats,
                                    num_components, dataset)
        model['config'] = sub_config
        ensemble.append(model)
    return ensemble def generate_fake_model(num_rows,
                        num_cols,
                        num_cats,
                        num_components,
                        dataset=None):
    tree = generate_tree(num_cols)
    assignments = np.random.choice(num_components, size=(num_rows, num_cols))
    assignments = assignments.astype(np.int32)
    if dataset is None:
        dataset = generate_dataset(num_rows, num_cols, num_cats)
    table = dataset['table']
    N = num_rows
    V = num_cols
    E = V - 1
    K = V * (V - 1) // 2
    C = num_cats
    M = num_components
    vert_ss = np.zeros((V, M), dtype=np.int32)
    edge_ss = np.zeros((E, M, M), dtype=np.int32)
    feat_ss = np.zeros((V * C, M), dtype=np.int32)
    meas_ss = np.zeros([V, M], np.int32)
    for v in range(V):
        vert_ss[v, :] = np.bincount(assignments[:, v], minlength=M)
    for e, v1, v2 in tree.tree_grid.T:
        pairs = assignments[:, v1].astype(np.int32) * M + assignments[:, v2]
        edge_ss[e, :, :] = np.bincount(pairs, minlength=M * M).reshape((M, M))
    for v in range(V):
        beg, end = table.ragged_index[v:v + 2]
        data_block = table.data[:, beg:end]
        feat_ss_block = feat_ss[beg:end, :]
        for n in range(N):
            feat_ss_block[:, assignments[n, v]] += data_block[n, :]
            meas_ss[v, assignments[n, v]] += data_block[n, :].sum()
    edge_logits = np.exp(-np.random.rand(K))
    model = {
        'tree': tree,
        'assignments': assignments,
        'edge_logits': edge_logits,
        'suffstats': {
            'ragged_index': table.ragged_index,
            'vert_ss': vert_ss,
            'edge_ss': edge_ss,
            'feat_ss': feat_ss,
            'meas_ss': meas_ss,
        },
    }
    return model class EnsembleServer(ServerBase):
    """Class for serving queries against a trained TreeCat ensemble."""

    def __init__(self, ensemble):
        logger.info('EnsembleServer of size %d', len(ensemble))
        assert ensemble
        ServerBase.__init__(self, ensemble[0]['suffstats']['ragged_index'])
        self._ensemble = [TreeCatServer(model) for model in ensemble]

        # These are used to inspect and visualize latent structure.
        self._edge_logits = self._ensemble[0].edge_logits.copy()
        for server in self._ensemble[1:]:
            self._edge_logits += server.edge_logits
        self._edge_logits /= len(self._ensemble)
        grid = self._ensemble[0]._tree.complete_grid
        self._estimated_tree = tuple(estimate_tree(grid, self._edge_logits))
        self._ensemble[0]._tree.gc()

    @property
    def edge_logits(self):
        return self._edge_logits

    def estimate_tree(self):
        return self._estimated_tree

    def sample_tree(self, num_samples):
        size = len(self._ensemble)
        pvals = np.ones(size, dtype=np.float32) / size
        sub_nums = np.random.multinomial(num_samples, pvals)
        samples = []
        for server, sub_num in zip(self._ensemble, sub_nums):
            samples += server.sample_tree(sub_num)
        np.random.shuffle(samples)
        assert len(samples) == num_samples
        return samples

    def sample(self, N, counts, data=None):
        size = len(self._ensemble)
        pvals = np.ones(size, dtype=np.float32) / size
        sub_Ns = np.random.multinomial(N, pvals)
        samples = np.concatenate([
            server.sample(sub_N, counts, data)
            for server, sub_N in zip(self._ensemble, sub_Ns)
        ])
        np.random.shuffle(samples)
        assert samples.shape[0] == N
        return samples

    def logprob(self, data):
        logprobs = np.stack(
            [server.logprob(data) for server in self._ensemble])
        logprobs = logsumexp(logprobs, axis=0)
        logprobs -= np.log(len(self._ensemble))
        assert logprobs.shape == (data.shape[0], )
        return logprobs

    def marginals(self, data):
        """Compute observed marginals conditioned on data."""
        TODO('https://github.com/posterior/treecat/issues/2')

    def mode(self, counts, data):
        """Compute a maximum a posteriori data value conditioned on data."""
        TODO('https://github.com/posterior/treecat/issues/3')

    def observed_perplexity(self, counts):
        """Compute perplexity = exp(entropy) of observed variables.

        Perplexity is an information theoretic measure of the number of
        clusters or observed classes. Perplexity is a real number in the range
        [1, dim[v]], where dim[v] is the number of categories in an observed
        categorical variable or 2 for an ordinal variable.

        Args:
            counts: A [V]-shaped array of multinomial counts.

        Returns:
            A [V]-shaped numpy array of perplexity.
        """
        result = self._ensemble[0].observed_perplexity(counts)
        for server in self._ensemble[1:]:
            result += server.observed_perplexity(counts)
        result /= len(self._ensemble)
        return result

    def latent_perplexity(self):
        """Compute perplexity = exp(entropy) of latent variables.

        Perplexity is an information theoretic measure of the number of
        clusters or latent classes. Perplexity is a real number in the range
        [1, M], where M is model_num_clusters.

        Returns:
            A [V]-shaped numpy array of perplexity.
        """
        result = self._ensemble[0].latent_perplexity()
        for server in self._ensemble[1:]:
            result += server.latent_perplexity()
        result /= len(self._ensemble)
        return result

    def latent_correlation(self):
        """Compute correlation matrix among latent features.

        This computes the generalization of Pearson's correlation to discrete
        data. Let I(X;Y) be the mutual information. Then define correlation as

          rho(X,Y) = sqrt(1 - exp(-2 I(X;Y)))

        Returns:
            A [V, V]-shaped numpy array of feature-feature correlations.
        """
        result = self._ensemble[0].latent_correlation()
        for server in self._ensemble[1:]:
            result += server.latent_correlation()
        result /= len(self._ensemble)
        return result class TreeCatServer(ServerBase):
    """Class for serving queries against a trained TreeCat model."""

    def __init__(self, model):
        """Create a TreeCat server.

        Args:
            model: A dict with fields:
                tree: A TreeStructure.
                suffstats: A dict of sufficient statistics.
                edge_logits: A K-sized array of nonnormalized edge
                    probabilities.
                config: A global config dict.
        """
        tree = model['tree']
        suffstats = model['suffstats']
        config = model['config']
        logger.info('TreeCatServer with %d features', tree.num_vertices)
        assert isinstance(tree, TreeStructure)
        ragged_index = suffstats['ragged_index']
        ServerBase.__init__(self, ragged_index)
        self._tree = tree
        self._config = config
        self._program = make_propagation_program(tree.tree_grid)

        # These are useful dimensions to import into locals().
        V = self._tree.num_vertices
        E = V - 1  # Number of edges in the tree.
        M = self._config['model_num_clusters']  # Number of latent clusters.
        R = ragged_index[-1]  # Size of ragged data.
        self._VEMR = (V, E, M, R)

        # Use Jeffreys priors.
        vert_prior = 0.5
        edge_prior = 0.5 / M
        feat_prior = 0.5 / M

        # These are posterior marginals for vertices and pairs of vertices.
        self._vert_probs = suffstats['vert_ss'].astype(np.float32) + vert_prior
        self._vert_probs /= self._vert_probs.sum(axis=1, keepdims=True)
        self._edge_probs = suffstats['edge_ss'].astype(np.float32) + edge_prior
        self._edge_probs /= self._edge_probs.sum(axis=(1, 2), keepdims=True)

        # This represents information in the pairwise joint posterior minus
        # information in the individual factors.
        self._edge_trans = self._edge_probs.copy()
        for e, v1, v2 in tree.tree_grid.T:
            self._edge_trans[e, :, :] /= self._vert_probs[v1, :, np.newaxis]
            self._edge_trans[e, :, :] /= self._vert_probs[v2, np.newaxis, :]

        # This is the conditional distribution of features given latent.
        self._feat_cond = suffstats['feat_ss'].astype(np.float32) + feat_prior
        for v in range(V):
            beg, end = ragged_index[v:v + 2]
            feat_block = self._feat_cond[beg:end, :]
            feat_block /= feat_block.sum(axis=0, keepdims=True)

        # These are used to inspect and visualize latent structure.
        self._edge_logits = model['edge_logits']
        self._estimated_tree = tuple(
            estimate_tree(self._tree.complete_grid, self._edge_logits))
        self._tree.gc()

    @property
    def edge_logits(self):
        return self._edge_logits

    def estimate_tree(self):
        return self._estimated_tree

    def sample_tree(self, num_samples):
        """Returns a num_samples-long list of trees, each a list of pairs."""
        samples = []
        edge_logits = self.edge_logits
        edges = self.estimate_tree()
        for _ in range(num_samples):
            edges = sample_tree(self._tree.complete_grid, edge_logits, edges)
            samples.append(edges)
        return samples

    @profile
    def sample(self, N, counts, data=None):
        """Draw N samples from the posterior distribution.

        Args:
            size: The number of samples to draw.
            counts: A [V]-shaped numpy array of requested counts of
                multinomials to sample.
            data: An optional single row of conditioning data, as a [R]-shaped
                ragged numpy array of multinomial counts,
                where R = server.ragged_size.

        Returns:
            An [N, R]-shaped numpy array of sampled multinomial data.
        """
        logger.debug('sampling data')
        V, E, M, R = self._VEMR
        if data is None:
            data = self._zero_row
        assert data.shape == (R, )
        assert data.dtype == np.int8
        assert counts.shape == (V, )
        assert counts.dtype == np.int8
        edge_trans = self._edge_trans
        feat_cond = self._feat_cond

        messages_in = self._vert_probs.copy()
        messages_out = np.tile(self._vert_probs[:, np.newaxis, :], (1, N, 1))
        vert_samples = np.zeros([V, N], np.int8)
        feat_samples = np.zeros([N, R], np.int8)
        range_N = np.arange(N, dtype=np.int32)

        for op, v, v2, e in self._program:
            if op == OP_UP:
                # Propagate upward from observed to latent.
                message = messages_in[v, :]
                beg, end = self._ragged_index[v:v + 2]
                for r in range(beg, end):
                    # This uses a with-replacement approximation that is exact
                    # for categorical data but approximate for multinomial.
                    if data[r]:
                        message *= feat_cond[r, :]**data[r]
            elif op == OP_IN:
                # Propagate latent state inward from children to v.
                message = messages_in[v, :]
                trans = edge_trans[e, :, :]
                if v > v2:
                    trans = trans.T
                message *= np.dot(trans, messages_in[v2, :])
                message /= message.max()  # Scale for numerical stability.
            elif op == OP_ROOT:
                # Process root node.
                messages_out[v, :, :] = messages_in[v, np.newaxis, :]
            elif op == OP_OUT:
                # Propagate latent state outward from parent to v.
                message = messages_out[v, :, :]
                message[...] = messages_in[v, np.newaxis, :]
                trans = edge_trans[e, :, :]
                if v2 > v:
                    trans = trans.T
                message *= trans[vert_samples[v2, :], :]
                # Scale for numerical stability.
                message /= message.max(axis=1, keepdims=True)
            elif op == OP_DOWN:
                # Sample latent and observed assignment.
                message = messages_out[v, :, :]
                vert_samples[v, :] = sample_from_probs2(message)
                beg, end = self._ragged_index[v:v + 2]
                feat_block = feat_cond[beg:end, :].T
                probs = feat_block[vert_samples[v, :], :]
                samples_block = feat_samples[:, beg:end]
                for _ in range(counts[v]):
                    samples_block[range_N, sample_from_probs2(probs)] += 1

        return feat_samples

    @profile
    def logprob(self, data):
        """Compute non-normalized log probabilies of many rows of data.

        To compute conditional probabilty, use the identity:

          log P(data|evidence) = server.logprob(data + evidence)
                               - server.logprob(evidence)

        Args:
            data: A [N,R]-shaped ragged nummpy array of multinomial count data,
                where N is the number of rows, and R = server.ragged_size.

        Returns:
            An [N]-shaped numpy array of log probabilities.
        """
        logger.debug('computing logprob')
        V, E, M, R = self._VEMR
        N = data.shape[0]
        assert data.shape == (N, R)
        assert data.dtype == np.int8
        edge_trans = self._edge_trans
        feat_cond = self._feat_cond

        messages = np.tile(self._vert_probs[:, :, np.newaxis], (1, 1, N))
        assert messages.shape == (V, M, N)
        logprob = np.zeros(N, np.float32)

        for op, v, v2, e in self._program:
            message = messages[v, :, :]
            if op == OP_UP:
                # Propagate upward from observed to latent.
                beg, end = self._ragged_index[v:v + 2]
                for r in range(beg, end):
                    # This uses a with-replacement approximation that is exact
                    # for categorical data but approximate for multinomial.
                    power = data[np.newaxis, :, r]
                    message *= feat_cond[r, :, np.newaxis]**power
            elif op == OP_IN:
                # Propagate latent state inward from children to v.
                trans = edge_trans[e, :, :]
                if v > v2:
                    trans = trans.T
                message *= np.dot(trans, messages[v2, :, :])
                message_sum = message.sum(axis=0, keepdims=True)
                message /= message_sum
                logprob += np.log(message_sum[0, :])
            elif op == OP_ROOT:
                # Aggregate total log probability at the root node.
                logprob += np.log(message.sum(axis=0))
                return logprob

    @profile
    def marginals(self, data):
        """Compute observed marginals conditioned on data.

        Args:
            data: An [N, R]-shaped numpy array of row of conditioning data, as
                a ragged nummpy array of multinomial counts,
                where R = server.ragged_size.

        Returns:
            An real-valued array of the same shape as data.
        """
        logger.debug('computing marginals')
        V, E, M, R = self._VEMR
        N = data.shape[0]
        assert data.shape == (N, R)
        assert data.dtype == np.int8
        edge_trans = self._edge_trans
        feat_cond = self._feat_cond

        messages_in = np.empty([V, M, N], dtype=np.float32)
        messages_in[...] = self._vert_probs[:, :, np.newaxis]
        messages_out = np.empty_like(messages_in.copy())
        result = np.zeros([N, R], np.float32)

        for op, v, v2, e in self._program:
            message = messages_in[v, :, :]
            if op == OP_UP:
                # Propagate upward from observed to latent.
                beg, end = self._ragged_index[v:v + 2]
                for r in range(beg, end):
                    # This uses a with-replacement approximation that is exact
                    # for categorical data but approximate for multinomial.
                    power = data[np.newaxis, :, r]
                    message *= feat_cond[r, :, np.newaxis]**power
                messages_out[v, :, :] = message
            elif op == OP_IN:
                # Propagate latent state inward from children to v.
                trans = edge_trans[e, :, :]
                if v > v2:
                    trans = trans.T
                message *= np.dot(trans, messages_in[v2, :, :])
                # Scale for numerical stability.
                message /= message.max(axis=0, keepdims=True)
            elif op == OP_OUT:
                # Propagate latent state outward from parent to v.
                trans = edge_trans[e, :, :]
                if v > v2:
                    trans = trans.T
                from_parent = np.dot(trans, messages_out[v2, :, :])
                messages_out[v, :, :] *= from_parent
                message *= from_parent
            elif op == OP_DOWN:
                # Propagate downward from latent state to observations.
                beg, end = self._ragged_index[v:v + 2]
                marginal = result[:, beg:end]
                marginal[...] = np.dot(feat_cond[beg:end, :], message).T
                marginal /= marginal.sum(axis=1, keepdims=True)

        return result

    def mode(self, counts, data):
        """Compute a maximum a posteriori data value conditioned on data."""
        TODO('https://github.com/posterior/treecat/issues/3')

    def observed_perplexity(self, counts):
        """Compute perplexity = exp(entropy) of observed variables.

        Perplexity is an information theoretic measure of the number of
        clusters or latent classes. Perplexity is a real number in the range
        [1, M], where M is model_num_clusters.

        Args:
            counts: A [V]-shaped array of multinomial counts.

        Returns:
            A [V]-shaped numpy array of perplexity.
        """
        V, E, M, R = self._VEMR
        if counts is not None:
            counts = np.ones(V, dtype=np.int8)
        assert counts.shape == (V, )
        assert counts.dtype == np.int8
        assert np.all(counts > 0)
        observed_entropy = np.empty(V, dtype=np.float32)
        for v in range(V):
            beg, end = self._ragged_index[v:v + 2]
            probs = np.dot(self._feat_cond[beg:end, :], self._vert_probs[v, :])
            observed_entropy[v] = multinomial_entropy(probs, counts[v])
        return np.exp(observed_entropy)

    def latent_perplexity(self):
        """Compute perplexity = exp(entropy) of latent variables.

        Perplexity is an information theoretic measure of the number of
        clusters or latent classes. Perplexity is a real number in the range
        [1, M], where M is model_num_clusters.

        Returns:
            A [V]-shaped numpy array of perplexity.
        """
        V, E, M, R = self._VEMR
        latent_entropy = np.array(
            [entropy(self._vert_probs[v, :]) for v in range(V)],
            dtype=np.float32)
        return np.exp(latent_entropy)

    @profile
    def latent_correlation(self):
        """Compute correlation matrix among latent features.

        This computes the generalization of Pearson's correlation to discrete
        data. Let I(X;Y) be the mutual information. Then define correlation as

          rho(X,Y) = sqrt(1 - exp(-2 I(X;Y)))

        Returns:
            A [V, V]-shaped numpy array of feature-feature correlations.
        """
        logger.debug('computing latent correlation')
        V, E, M, R = self._VEMR
        edge_probs = self._edge_probs
        vert_probs = self._vert_probs
        result = np.zeros([V, V], np.float32)
        for root in range(V):
            messages = np.empty([V, M, M])
            program = make_propagation_program(self._tree.tree_grid, root)
            for op, v, v2, e in program:
                if op == OP_ROOT:
                    # Initialize correlation at this node.
                    messages[v, :, :] = np.diagflat(vert_probs[v, :])
                elif op == OP_OUT:
                    # Propagate correlation outward from parent to v.
                    trans = edge_probs[e, :, :]
                    if v > v2:
                        trans = trans.T
                    messages[v, :, :] = np.dot(  #
                        trans / vert_probs[v2, np.newaxis, :],
                        messages[v2, :, :])
            for v in range(V):
                result[root, v] = correlation(messages[v, :, :])
        return result TINY_CONFIG = make_config(
    learning_init_epochs=2,  #
    model_num_clusters=7,  #
    model_ensemble_size=3) TINY_TABLE = Table(TINY_FEATURE_TYPES, TINY_RAGGED_INDEX, TINY_DATA) def make_seed(*args):
    """Create a random seed by fingerprinting args."""
    digest = fingerprint(args)
    bigseed = int('0x{}'.format(digest), 0)
    seed = bigseed % (2**32)
    return seed def numpy_seterr():
    np.seterr(divide='raise', invalid='raise') def xfail_param(*args, **kwargs):
    return pytest.param(*args, marks=pytest.mark.xfail(**kwargs)) def train_ensemble(table, tree_prior, config):
    """Train a TreeCat ensemble model using subsample-annealed MCMC.

    The ensemble size is controlled by config['model_ensemble_size'].
    Let N be the number of data rows and V be the number of features.

    Args:
        table: A Table instance holding N rows of V features of data.
        tree_prior: A [K]-shaped numpy array of prior edge log odds, where
            K is the number of edges in the complete graph on V vertices.
        config: A global config dict.

    Returns:
        A trained model as a dictionary with keys:
            tree: A TreeStructure instance with the learned latent structure.
            suffstats: Sufficient statistics of features, vertices, and edges.
            assignments: An [N, V] numpy array of latent cluster ids for each
                cell in the dataset.
    """
    tasks = []
    for sub_seed in range(config['model_ensemble_size']):
        sub_config = config.copy()
        sub_config['seed'] += sub_seed
        tasks.append((table, tree_prior, sub_config))
    return parallel_map(_train_model, tasks) def train_model(table, tree_prior, config):
    """Train a TreeCat model using subsample-annealed MCMC.

    Let N be the number of data rows and V be the number of features.

    Args:
        table: A Table instance holding N rows of V features of data.
        tree_prior: A [K]-shaped numpy array of prior edge log odds, where
            K is the number of edges in the complete graph on V vertices.
        config: A global config dict.

    Returns:
        A trained model as a dictionary with keys:
            config: A global config dict.
            tree: A TreeStructure instance with the learned latent structure.
            edge_logits: A [K]-shaped array of all edge logits.
            suffstats: Sufficient statistics of features, vertices, and edges.
            assignments: An [N, V] numpy array of latent cluster ids for each
                cell in the dataset.
    """
    assert isinstance(table, Table)
    M = config['model_num_clusters']
    D = config['model_latent_dim']
    assert M >= 1
    assert D >= 0
    if D == 0:
        Trainer = TreeCatTrainer
    elif M == 1:
        Trainer = TreeGaussTrainer
    else:
        Trainer = TreeMogTrainer
    return Trainer(table, tree_prior, config).train() def set_random_seed(seed):
    """Set random seeds for both numpy and numba."""
    np.random.seed(seed)
    jit_random_seed(seed)
import itertools
import numpy as np
import pytest
from goftests import multinomial_goodness_of_fit
from treecat.generate import generate_dataset
from treecat.generate import generate_fake_ensemble
from treecat.generate import generate_fake_model
from treecat.serving import EnsembleServer
from treecat.serving import TreeCatServer
from treecat.testutil import TINY_CONFIG
from treecat.testutil import TINY_TABLE
from treecat.testutil import make_seed
from treecat.testutil import numpy_seterr
from treecat.testutil import xfail_param
from treecat.training import train_ensemble
from treecat.training import train_model
from treecat.util import set_random_seed
@pytest.mark.parametrize('N,V,C,M', NVCM_EXAMPLES_FOR_GOF)
def test_ensemble_unconditional_gof(N, V, C, M):
    ensemble = generate_fake_ensemble(N, V, C, M)
    server = EnsembleServer(ensemble)
    validate_gof(N, V, C, M, server, conditional=False)


@pytest.mark.parametrize('N,V,C,M', NVCM_EXAMPLES_FOR_GOF)
def test_ensemble_conditional_gof(N, V, C, M):
    ensemble = generate_fake_ensemble(N, V, C, M)
    server = EnsembleServer(ensemble)
    validate_gof(N, V, C, M, server, conditional=True)


def validate_gof(N, V, C, M, server, conditional):
    # Generate samples.
    expected = C**V
    num_samples = 1000 * expected
    ones = np.ones(V, dtype=np.int8)
    if conditional:
        cond_data = server.sample(1, ones)[0, :]
    else:
        cond_data = server.make_zero_row()
    samples = server.sample(num_samples, ones, cond_data)
    logprobs = server.logprob(samples + cond_data[np.newaxis, :])
    counts = {}
    probs = {}
    for sample, logprob in zip(samples, logprobs):
        key = tuple(sample)
        if key in counts:
            counts[key] += 1
        else:
            counts[key] = 1
            probs[key] = np.exp(logprob)
    assert len(counts) == expected

    # Check accuracy using Pearson's chi-squared test.
    keys = sorted(counts.keys(), key=lambda key: -probs[key])
    counts = np.array([counts[k] for k in keys], dtype=np.int32)
    probs = np.array([probs[k] for k in keys])
    probs /= probs.sum()

    # Truncate to avoid low-precision.
    truncated = False
    valid = (probs * num_samples > 20)
    if not valid.all():
        T = valid.argmin()
        T = max(8, T)  # Avoid truncating too much
        probs = probs[:T]
        counts = counts[:T]
        truncated = True

    gof = multinomial_goodness_of_fit(
        probs, counts, num_samples, plot=True, truncated=truncated)
    assert 1e-2 < gof


@pytest.mark.parametrize('N,V,C,M', [
    (10, 1, 2, 2),
    (10, 2, 3, 3),
    (10, 3, 4, 4),
    (10, 4, 5, 5),
    (10, 4, 6, 6),
    (10, 4, 7, 7),
    (10, 4, 8, 8),
])
def test_observed_perplexity(N, V, C, M):
    set_random_seed(make_seed(N, V, C, M))
    model = generate_fake_model(N, V, C, M)
    config = TINY_CONFIG.copy()
    config['model_num_clusters'] = M
    model['config'] = config
    server = TreeCatServer(model)

    for count in [1, 2, 3]:
        if count > 1 and C > 2:
            continue  # NotImplementedError.
        counts = 1
        perplexity = server.observed_perplexity(counts)
        print(perplexity)
        assert perplexity.shape == (V, )
        assert np.all(1 <= perplexity)
        assert np.all(perplexity <= count * C)


@pytest.mark.parametrize('N,V,C,M', [
    (10, 1, 2, 2),
    (10, 2, 2, 3),
    (10, 3, 2, 4),
    (10, 4, 2, 5),
    (10, 5, 2, 6),
    (10, 6, 2, 7),
    (10, 7, 2, 8),
])
def test_latent_perplexity(N, V, C, M):
    set_random_seed(make_seed(N, V, C, M))
    model = generate_fake_model(N, V, C, M)
    config = TINY_CONFIG.copy()
    config['model_num_clusters'] = M
    model['config'] = config
    server = TreeCatServer(model)

    perplexity = server.latent_perplexity()
    print(perplexity)
    assert perplexity.shape == (V, )
    assert np.all(1 <= perplexity)
    assert np.all(perplexity <= M)


@pytest.mark.parametrize('N,V,C,M', [
    (10, 1, 2, 2),
    (10, 2, 2, 3),
    (10, 3, 2, 4),
    (10, 4, 2, 5),
    (10, 5, 2, 6),
    (10, 6, 2, 7),
    (10, 7, 2, 8),
])
def test_ensemble_latent_perplexity(N, V, C, M):
    set_random_seed(make_seed(N, V, C, M))

    ensemble = generate_fake_ensemble(N, V, C, M)