def deserialize_config(config_str):
    """Deserialize a config dict form a short string."""
    config = make_config()
    keys = sorted(config.keys())
    values = config_str.split('-')
    assert len(keys) == len(values)
    for key, value_str in zip(keys, values):
        config[key] = int(value_str)
    return config def make_config(**options):
    """Create a new global config dict with default values."""
    config = DEFAULT_CONFIG.copy()
    for key, value in options.items():
        if key not in config:
            raise ValueError('Unknown option: {}. Try one of:\n  {}'.format(
                key, '\n  '.join(sorted(config.keys()))))
        config[key] = int(value)
    return config def serialize_config(config):
    """Serialize a config dict to a short string for use in filenames."""
    keys = sorted(config.keys())
    assert keys == sorted(make_config().keys())
    return '-'.join(str(int(config[key])) for key in keys) @contextmanager
def csv_reader(filename, encoding='utf-8'):
    with io.open(filename, 'r', encoding=encoding) as f:
        if sys.version_info >= (3, 0):
            yield csv.reader(f)
        else:
            yield csv.reader(line.encode(encoding, 'ignore') for line in f) def pickle_dump(data, filename):
    """Serialize data to file using gzip compression."""
    if filename.endswith('.pkz'):
        with gzip.open(filename, 'wb') as f:
            pickle.dump(data, f, protocol=2)  # Try to support python 2.
    elif filename.endswith('.jz'):
        with gzip.open(filename, 'wt') as f:
            f.write(json_dumps(data))
    else:
        raise ValueError(
            'Cannot determine format: {}'.format(os.path.basename(filename))) def pickle_load(filename):
    """Deserialize data from file using gzip compression."""
    if filename.endswith('.pkz'):
        with gzip.open(filename, 'rb') as f:
            return pickle.load(f)
    elif filename.endswith('.jz'):
        with gzip.open(filename, 'rt') as f:
            return json_loads(f.read())
    else:
        raise ValueError(
            'Cannot determine format: {}'.format(os.path.basename(filename))) def check_call_env(cmd, env):
    ret = Popen(cmd, env=env).wait()
    if ret:
        raise CalledProcessError(returncode=ret, cmd=cmd) 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 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 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 count_observations(ragged_index, data):
    """Count the observations in each cell of a ragged data array.

    Args:
        ragged_index: A [V+1]-shaped numpy array as returned by
            make_ragged_index.
        data: A [N, R]-shaped ragged array of multinomial count data, where
            N is the number of rows and R = ragged_index[-1].

    Returns:
        A [N, V]-shaped array whose entries are the number of observations
        in each cell of data.
    """
    N, R = data.shape
    assert R == ragged_index[-1]
    V = len(ragged_index) - 1
    counts = np.zeros([N, V], np.int8)
    for v in range(V):
        beg, end = ragged_index[v:v + 2]
        counts[:, v] = data[:, beg:end].sum(axis=1)
    return counts def get_profiling_stats():
    return {
        'timers': {
            name: {
                'elapsed': timer.elapsed,
                'count': timer.count,
            }
            for name, timer in TIMERS.items()
        },
        'counters': dict(COUNTERS),
        'histograms':
        {name: dict(counter)
         for name, counter in HISTOGRAMS.items()},
        'series': dict(SERIES),
    } def make_ragged_mask(ragged_index, mask):
    """Convert a boolean mask from dense to ragged format.

    Args:
        ragged_index: A [V+1]-shaped numpy array as returned by
            make_ragged_index.
        mask: A [V,...]-shaped numpy array of booleans.

    Returns:
        A [R,...]-shaped numpy array, where R = ragged_index[-1].
    """
    V = ragged_index.shape[0] - 1
    R = ragged_index[-1]
    assert mask.shape[0] == V
    assert mask.dtype == np.bool_
    ragged_mask = np.empty((R, ) + mask.shape[1:], dtype=np.bool_)
    for v in range(V):
        beg, end = ragged_index[v:v + 2]
        ragged_mask[beg:end] = mask[v]
    return ragged_mask 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): def set_random_seed(seed):
    """Set random seeds for both numpy and numba."""
    np.random.seed(seed)
    jit_random_seed(seed)
import multiprocessing
import os
import sys
import numpy as np
from parsable import parsable
from six.moves import range
from six.moves import zip
from treecat.config import deserialize_config
from treecat.config import make_config
from treecat.config import serialize_config
from treecat.format import csv_reader
from treecat.format import pickle_dump
from treecat.format import pickle_load
from treecat.profile import check_call_env
from treecat.serving import TreeCatServer
from treecat.tables import Table
from treecat.training import train_model
from treecat.util import count_observations
from treecat.util import get_profiling_stats
from treecat.util import make_ragged_mask
from treecat.util import parallel_map
from treecat.util import profile
from treecat.util import set_random_seed

def split_data(ragged_index, num_rows, num_parts, partid):
    """Split a dataset into training + holdout for n-fold crossvalidation.

    This splits a dataset into num_parts disjoint parts by randomly holding out
    cells. Note that whereas supervised crossvalidation typically holds out
    entire rows, our unsupervised crossvalidation is intended to evaluate a
    model of the full joint distribution.

    Args:
        ragged_index: A [V+1]-shaped numpy array of indices into the ragged
            data array, where V is the number of features.
        num_rows: An integer, the number of rows in the dataset.
        num_parts: An integer, the number of folds in n-fold crossvalidation.
        partid: An integer in [0, num_parts).

    Returns:
        A [N,R]-shaped mask where True means held-out and False means training.
        Here N = num_rows and R = ragged_index[-1].
    """
    set_random_seed(0)
    assert 0 <= partid < num_parts
    N = num_rows
    V = ragged_index.shape[0] - 1
    R = ragged_index[-1]
    dense_mask = (partid == np.random.randint(num_parts, size=(N, V)))
    ragged_mask = make_ragged_mask(ragged_index, dense_mask.T).T
    assert ragged_mask.shape == (N, R)
    return ragged_mask


def read_param_csv(param_csv_path, **options):
    """Reads configs from a csv file.

    Args:
        param_csv_path: The path to a csv file with one line per config.
        options: A dict of extra config parameters.

    Returns:
        A pair (header, configs), where:
        header is a list of parameters, and
        configs is list of config dicts.
    """
    # This is a hack: to configure crossvalidation from a config, we alias:
    #   num_parts = model_ensemble_size
    #   partid = seed
    num_parts = options.get('model_ensemble_size', 4)
    assert 'seed' not in options
    configs = []
    with csv_reader(param_csv_path) as reader:
        header = next(reader)
        assert 'seed' not in header
        assert 'model_ensemble_size' not in header
        for row in reader:
            if len(row) != len(header) or row[0].startswith('#'):
                continue
            for key, value in zip(header, row):
                options[key] = int(value)
            for partid in range(num_parts):
                configs.append(make_config(seed=partid, **options))
    return header, configs


@parsable
def train_task(dataset_path, model_path, config_str):
    """INTERNAL Train a single model."""
    print('Train {}'.format(os.path.basename(model_path)))
    config = deserialize_config(config_str)
    dataset = pickle_load(dataset_path)
    tree_prior = dataset['schema']['tree_prior']

    # Split data for crossvalidation.
    num_parts = config['model_ensemble_size']
    partid = config['seed']
    assert 0 <= partid < num_parts
    table = dataset['table']
    ragged_index = table.ragged_index
    mask = split_data(ragged_index, table.num_rows, num_parts, partid)
    training_data = table.data.copy()
    training_data[mask] = 0
    training_table = Table(table.feature_types, ragged_index, training_data)

    # Train a model.
    model = train_model(training_table, tree_prior, config)
    model['profiling_stats'] = get_profiling_stats()
    pickle_dump(model, model_path)


def process_train_task(task):
    (dataset_path, models_dir, config) = task
    config_str = serialize_config(config)
    model_path = os.path.join(models_dir, 'model.{}.pkz'.format(config_str))
    env = os.environ.copy()
    env['TREECAT_PROFILE'] = '1'
    env['TREECAT_THREADS'] = '1'
    check_call_env([
        sys.executable,
        '-O',
        '-m',
        'treecat.validate',
        'train_task',
        dataset_path,
        model_path,
        config_str,
    ], env)


@parsable
def train(dataset_path, param_csv_path, models_dir, **options):
    """Tune parameters specified in a csv file."""
    options = {k: int(v) for k, v in options.items()}
    if not os.path.exists(models_dir):
        os.makedirs(models_dir)
    header, configs = read_param_csv(param_csv_path, **options)
    configs.sort(key=lambda c: c['learning_init_epochs'])
    tasks = [(dataset_path, models_dir, c) for c in configs]
    print('Scheduling {} tasks'.format(len(tasks)))
    parallel_map(process_train_task, tasks)



@profile