<commit_before># -*- coding: utf-8 -*-

# The MIT License (MIT)

# Copyright (c) 2014 <NAME>,
# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# Lead Developer: <NAME> <<EMAIL>>

# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:

# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.

import copy
import sys
import time
import cPickle as pickle
from math import log

import numpy as np
import matplotlib.pyplot as plt
from scipy.misc import logsumexp

import gpmcc.utils.general as gu
import gpmcc.utils.plots as pu

from gpmcc.view import View
from gpmcc.dim import Dim

_all_kernels = [
    'column_z',
    'state_alpha',
    'row_z',
    'column_hypers',
    'view_alphas'
    ]

class State(object):
    """State, the main crosscat object."""

    def __init__(self, X, cctypes, distargs=None, Zv=None, Zrcv=None, alpha=None,
            view_alphas=None, hypers=None, n_grid=30, seed=None):
        """Dim constructor provides a convenience method for bulk incorporate
        and unincorporate by specifying the data, and optinally view partition
        and row partition for each view.

        Parameters
        ----------
        X : np.ndarray
            A data matrix DxN, where D is the number of variabels and N is
            the number of observations.
        cctypes : list<str>
            Data type of each colum, see `utils.config` for valid cctypes.
        distargs : list<dict>, optional
            Distargs appropriate for each cctype in cctypes. For details on
            distargs see the documentation for each DistributionGpm. Empty
            distargs can be None or {}.
        n_grid : int, optional
            Number of bins for hyperparameter grids.
        Zv : list<int>, optional
            Assignmet of columns to views. If not specified a random
            partition is sampled.
        Zrcv : list(list<int>), optional
            Assignment of rows to clusters in each view, where Zrcv[k] is
            the Zr for View k. If not specified a random partition is
            sampled. If specified, then Zv must also be specified.
        seed : int
            Seed the random number generator.
        """
        # Seed.
        self.seed = 0 if seed is None else seed
        np.random.seed(self.seed)

        # Dataset.
        self.X = np.asarray(X)

        # Hyperparameters.
        self.n_grid = n_grid

        # Distargs.
        if distargs is None:
            distargs = [None] * len(cctypes)

        # Generate dimensions.
        self.dims = []
        for col in xrange(self.n_cols()):
            dim_hypers = None if hypers is None else hypers[col]
            self.dims.append(
                Dim(X[:,col], cctypes[col], col, n_grid=n_grid,
                hypers=dim_hypers, distargs=distargs[col]))

        # Generate CRP alpha.
        self.alpha_grid = gu.log_linspace(1./self.n_cols(), self.n_cols(),
            self.n_grid)
        if alpha is None:
            alpha = np.random.choice(self.alpha_grid)
        self.alpha = alpha

        # Generate view partition.
        if Zv is None:
            Zv = gu.simulate_crp(self.n_cols(), self.alpha)
        self.Zv = list(Zv)
        self.Nv = list(np.bincount(Zv))

        # Generate views.
        self.views = []
        for v in xrange(len(self.Nv)):
            dims = [self.dims[i] for i in xrange(self.n_cols()) if Zv[i] == v]
            Zr = None if Zrcv is None else np.asarray(Zrcv[v])
            alpha = None if view_alphas is None else view_alphas[v]
            V = View(self.X, dims, Zr=Zr, alpha=alpha, n_grid=n_grid)
            self.views.append(V)

        self._check_partitions()

    # --------------------------------------------------------------------------
    # Observe

    def incorporate_dim(self, X, cctype, distargs=None, v=None):
        """Incorporate a new Dim into this State.

        Parameters
        ----------
        X : np.array
            An array of data with length self.n_rows().
        cctype : str
            DistributionGpm name see `gpmcc.utils.config`.
        distargs : dict, optional.
            Distargs appropriate for the cctype. For details on
            distargs see the documentation for each DistributionGpm.
        v : int, optional
            Index of the view to assign the data. If unspecified, will be
            sampled. If 0 <= v < len(state.Nv) then will insert into an existing
            View. If v = len(state.Nv) a singleton view will be created with a
            partition from the prior.
        """
        assert len(X) == self.n_rows()
        self.X = np.column_stack((self.X, X))

        col = self.n_cols() - 1
        self.dims.append(Dim(X, cctype, col, n_grid=self.n_grid,
            distargs=distargs))

        for view in self.views:
            view.set_dataset(self.X)

        if 0 <= v < len(self.Nv):
            self.views[v].incorporate_dim(self.dims[-1])
            self.Zv.append(v)
            self.Nv[v] += 1
        elif v == len(self.Nv):
            self.views.append(
                View(self.X, [self.dims[-1]], n_grid=self.n_grid))
            self.Zv.append(v)
            self.Nv.append(1)
        else:
            self.views[0].incorporate_dim(self.dims[-1])
            self.Zv.append(0)
            self.Nv[0] += 1
            self.transition_columns(cols=[col])

        self.transition_column_hypers(cols=[col])
        self._check_partitions()

    def unincorporate_dim(self, col):
        """Unincorporate the existing dim with index col."""
        if self.n_cols() == 1:
            raise ValueError('State has only one dim, cannot unincorporate.')

        self.X = np.delete(self.X, col, 1)

        v = self.Zv[col]
        self.views[v].unincorporate_dim(self.dims[col])
        self.Nv[v] -= 1
        del self.Zv[col]

        if self.Nv[v] == 0:
            self.Zv = [i-1 if i>v else i for i in self.Zv]
            del self.Nv[v]
            del self.views[v]

        del self.dims[col]
        for i, dim in enumerate(self.dims):
            dim.index = i

        for view in self.views:
            view.set_dataset(self.X)
            view.reindex_dims()

        self._check_partitions()

    def incorporate_rows(self, X, k=None):
        """Incorporate list of new rows into global dataset X.

        Parameters
        ----------
        X : np.array
            A (r x self.n_cols) list of data, where r is the number of
            new rows to incorporate.
        k : list(list<int>), optional
            A (r x len(self.views)) list of integers, where r is the number of
            new rows to incorporate, and k[r][i] is the cluster to insert row r
            in view i. If k[r][i] is greater than the number of
            clusters in view[i] an error will be thrown. To specify cluster
            assignments for only some views, use None in all other locations
            i.e. k=[[None,2,None],[[0,None,1]]].
        """
        rowids = xrange(self.n_rows(), self.n_rows() + len(X))
        self.X = np.vstack((self.X, X))

        if k is None:
            k = [[None] * len(self.views)] * len(rowids)

        for v, view in enumerate(self.views):
            view.set_dataset(self.X)
            for r, rowid in enumerate(rowids):
                view.incorporate_row(rowid, k=k[r][v])

        self._check_partitions()

    def unincorporate_rows(self, rowids):
        """Unincorporate a list of rowids from dataset X. All r in rowids must
        be in range(0, State.n_rows())."""
        if self.n_rows() == 1:
            raise ValueError('State has only one row, cannot unincorporate.')

        self.X = np.delete(self.X, rowids, 0)

        for view in self.views:
            for rowid in rowids:
                view.unincorporate_row(rowid)
            view.set_dataset(self.X)
            view.reindex_rows()

        self._check_partitions()

    # --------------------------------------------------------------------------
    # logpdf

    def logpdf(self, rowid, query, evidence=None):
        """Compute density of query under the posterior predictive distirbution.

        Parameters
        ----------
        rowid : int
            The rowid of the member of the population to simulate from.
            If 0 <= rowid < state.n_rows then the latent variables of member
            rowid will be taken as conditioning variables.
            Otherwise logpdf for a hypothetical member is computed,
            marginalizing over latent variables.
        query : list(tuple<int>)
            A list of pairs (col, val) at which to query the logpdf.
        evidence : list(tuple<int>), optional
            A list of pairs (col, val) of observed values in the row to
            condition on

        Returns
        -------
        logpdf : float
            The logpdf(query|rowid, evidence).
        """
        if not 0 <= rowid < self.n_rows():
            return self.logpdf_unobserved(query, evidence=evidence)

        logpdf = 0
        for (col, val) in query:
            k = self.views[self.Zv[col]].Zr[rowid]
            logpdf += self.dims[col].logpdf(val, k)
        return logsumexp(logpdf)

    def logpdf_unobserved(self, query, evidence=None):
        """Simulates a hypothetical member, with no observed latents."""
        if evidence is None:
            evidence = []

        # Obtain all views of the query columns.
        query_views = set([self.Zv[col] for (col, _) in query])

        # Obtain the probability of hypothetical row belonging to each cluster.
        cluster_logps_for = dict()
        for v in query_views:
            # CRP densities.
            logp_crp = self._compute_cluster_crp_logps(v)
            # Evidence densities.
            logp_data = np.zeros(len(logp_crp))
            for (col, val) in evidence:
                if self.Zv[col] == v:
                    logp_data += self._compute_cluster_data_logps(col, val)
            cluster_logps_for[v] = gu.log_normalize(logp_crp+logp_data)

        logpdf = 0
        for (col, val) in query:
            # Query densities.
            logpdf += logsumexp(self._compute_cluster_data_logps(col, val)
                + cluster_logps_for[self.Zv[col]])

        return logpdf

    def logpdf_bulk(self, rowids, queries, evidences=None):
        """Evaluate multiple queries at once, used by Engine."""
        if evidences is None:
            evidences = [[] for _ in xrange(len(rowids))]
        assert len(rowids) == len(queries) == len(evidences)
        logpdfs = []
        for rowid, query, evidence in zip(rowids, queries, evidences):
            logpdfs.append(self.logpdf(rowid, query, evidence))
        return logpdfs

    def logpdf_marginal(self):
        return gu.logp_crp(len(self.Zv), self.Nv, self.alpha) + \
            sum(v.logpdf_marginal() for v in self.views)

    # --------------------------------------------------------------------------
    # Simulate

    def simulate(self, rowid, query, evidence=None, N=1):
        """Simulate from the posterior predictive distirbution.

        Parameters
        ----------
        rowid : int
            The rowid of the member of the population to simulate from.<commit_msg>Fix implementation of logpdf_unobserved, see Algorithm in comments.<commit_after># -*- coding: utf-8 -*-

# The MIT License (MIT)

# Copyright (c) 2014 <NAME>,
# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# Lead Developer: <NAME> <<EMAIL>>

# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:

# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.

import copy
import sys
import time
import cPickle as pickle
from math import log

import numpy as np
import matplotlib.pyplot as plt
from scipy.misc import logsumexp

import gpmcc.utils.general as gu
import gpmcc.utils.plots as pu

from gpmcc.view import View
from gpmcc.dim import Dim

_all_kernels = [
    'column_z',
    'state_alpha',
    'row_z',
    'column_hypers',
    'view_alphas'
    ]

class State(object):
    """State, the main crosscat object."""

    def __init__(self, X, cctypes, distargs=None, Zv=None, Zrcv=None, alpha=None,
            view_alphas=None, hypers=None, n_grid=30, seed=None):
        """Dim constructor provides a convenience method for bulk incorporate
        and unincorporate by specifying the data, and optinally view partition
        and row partition for each view.

        Parameters
        ----------
        X : np.ndarray
            A data matrix DxN, where D is the number of variabels and N is
            the number of observations.
        cctypes : list<str>
            Data type of each colum, see `utils.config` for valid cctypes.
        distargs : list<dict>, optional
            Distargs appropriate for each cctype in cctypes. For details on
            distargs see the documentation for each DistributionGpm. Empty
            distargs can be None or {}.
        n_grid : int, optional
            Number of bins for hyperparameter grids.
        Zv : list<int>, optional
            Assignmet of columns to views. If not specified a random
            partition is sampled.
        Zrcv : list(list<int>), optional
            Assignment of rows to clusters in each view, where Zrcv[k] is
            the Zr for View k. If not specified a random partition is
            sampled. If specified, then Zv must also be specified.
        seed : int
            Seed the random number generator.
        """
        # Seed.
        self.seed = 0 if seed is None else seed
        np.random.seed(self.seed)

        # Dataset.
        self.X = np.asarray(X)

        # Hyperparameters.
        self.n_grid = n_grid

        # Distargs.
        if distargs is None:
            distargs = [None] * len(cctypes)

        # Generate dimensions.
        self.dims = []
        for col in xrange(self.n_cols()):
            dim_hypers = None if hypers is None else hypers[col]
            self.dims.append(
                Dim(X[:,col], cctypes[col], col, n_grid=n_grid,
                hypers=dim_hypers, distargs=distargs[col]))

        # Generate CRP alpha.
        self.alpha_grid = gu.log_linspace(1./self.n_cols(), self.n_cols(),
            self.n_grid)
        if alpha is None:
            alpha = np.random.choice(self.alpha_grid)
        self.alpha = alpha

        # Generate view partition.
        if Zv is None:
            Zv = gu.simulate_crp(self.n_cols(), self.alpha)
        self.Zv = list(Zv)
        self.Nv = list(np.bincount(Zv))

        # Generate views.
        self.views = []
        for v in xrange(len(self.Nv)):
            dims = [self.dims[i] for i in xrange(self.n_cols()) if Zv[i] == v]
            Zr = None if Zrcv is None else np.asarray(Zrcv[v])
            alpha = None if view_alphas is None else view_alphas[v]
            V = View(self.X, dims, Zr=Zr, alpha=alpha, n_grid=n_grid)
            self.views.append(V)

        self._check_partitions()

    # --------------------------------------------------------------------------
    # Observe

    def incorporate_dim(self, X, cctype, distargs=None, v=None):
        """Incorporate a new Dim into this State.

        Parameters
        ----------
        X : np.array
            An array of data with length self.n_rows().
        cctype : str
            DistributionGpm name see `gpmcc.utils.config`.
        distargs : dict, optional.
            Distargs appropriate for the cctype. For details on
            distargs see the documentation for each DistributionGpm.
        v : int, optional
            Index of the view to assign the data. If unspecified, will be
            sampled. If 0 <= v < len(state.Nv) then will insert into an existing
            View. If v = len(state.Nv) a singleton view will be created with a
            partition from the prior.
        """
        assert len(X) == self.n_rows()
        self.X = np.column_stack((self.X, X))

        col = self.n_cols() - 1
        self.dims.append(Dim(X, cctype, col, n_grid=self.n_grid,
            distargs=distargs))

        for view in self.views:
            view.set_dataset(self.X)

        if 0 <= v < len(self.Nv):
            self.views[v].incorporate_dim(self.dims[-1])
            self.Zv.append(v)
            self.Nv[v] += 1
        elif v == len(self.Nv):
            self.views.append(
                View(self.X, [self.dims[-1]], n_grid=self.n_grid))
            self.Zv.append(v)
            self.Nv.append(1)
        else:
            self.views[0].incorporate_dim(self.dims[-1])
            self.Zv.append(0)
            self.Nv[0] += 1
            self.transition_columns(cols=[col])

        self.transition_column_hypers(cols=[col])
        self._check_partitions()

    def unincorporate_dim(self, col):
        """Unincorporate the existing dim with index col."""
        if self.n_cols() == 1:
            raise ValueError('State has only one dim, cannot unincorporate.')

        self.X = np.delete(self.X, col, 1)

        v = self.Zv[col]
        self.views[v].unincorporate_dim(self.dims[col])
        self.Nv[v] -= 1
        del self.Zv[col]

        if self.Nv[v] == 0:
            self.Zv = [i-1 if i>v else i for i in self.Zv]
            del self.Nv[v]
            del self.views[v]

        del self.dims[col]
        for i, dim in enumerate(self.dims):
            dim.index = i

        for view in self.views:
            view.set_dataset(self.X)
            view.reindex_dims()

        self._check_partitions()

    def incorporate_rows(self, X, k=None):
        """Incorporate list of new rows into global dataset X.

        Parameters
        ----------
        X : np.array
            A (r x self.n_cols) list of data, where r is the number of
            new rows to incorporate.
        k : list(list<int>), optional
            A (r x len(self.views)) list of integers, where r is the number of
            new rows to incorporate, and k[r][i] is the cluster to insert row r
            in view i. If k[r][i] is greater than the number of
            clusters in view[i] an error will be thrown. To specify cluster
            assignments for only some views, use None in all other locations
            i.e. k=[[None,2,None],[[0,None,1]]].
        """
        rowids = xrange(self.n_rows(), self.n_rows() + len(X))
        self.X = np.vstack((self.X, X))

        if k is None:
            k = [[None] * len(self.views)] * len(rowids)

        for v, view in enumerate(self.views):
            view.set_dataset(self.X)
            for r, rowid in enumerate(rowids):
                view.incorporate_row(rowid, k=k[r][v])

        self._check_partitions()

    def unincorporate_rows(self, rowids):
        """Unincorporate a list of rowids from dataset X. All r in rowids must
        be in range(0, State.n_rows())."""
        if self.n_rows() == 1:
            raise ValueError('State has only one row, cannot unincorporate.')

        self.X = np.delete(self.X, rowids, 0)

        for view in self.views:
            for rowid in rowids:
                view.unincorporate_row(rowid)
            view.set_dataset(self.X)
            view.reindex_rows()

        self._check_partitions()

    # --------------------------------------------------------------------------
    # logpdf

    def logpdf(self, rowid, query, evidence=None):
        """Compute density of query under the posterior predictive distirbution.

        Parameters
        ----------
        rowid : int
            The rowid of the member of the population to simulate from.
            If 0 <= rowid < state.n_rows then the latent variables of member
            rowid will be taken as conditioning variables.
            Otherwise logpdf for a hypothetical member is computed,
            marginalizing over latent variables.
        query : list(tuple<int>)
            A list of pairs (col, val) at which to query the logpdf.
        evidence : list(tuple<int>), optional
            A list of pairs (col, val) of observed values in the row to
            condition on

        Returns
        -------
        logpdf : float
            The logpdf(query|rowid, evidence).
        """
        if not 0 <= rowid < self.n_rows():
            return self.logpdf_unobserved(query, evidence=evidence)

        logpdf = 0
        for (col, val) in query:
            k = self.views[self.Zv[col]].Zr[rowid]
            logpdf += self.dims[col].logpdf(val, k)
        return logsumexp(logpdf)

    def logpdf_unobserved(self, query, evidence=None):
        """Simulates a hypothetical member, with no observed latents."""

        # Algorithm. Partition all columns in query and evidence by views.
        # P(x1,x2|x3,x4) where (x1...x4) in the same view.
        #   = \sum_z p(x1,x2|z,x3,x4)p(z|x3,x4)     marginalization
        #   = \sum_z p(x1,x2|z)p(z|x3,x4)           conditional independence
        #   = \sum_z p(x1|z)p(x2|z)p(z|x3,x4)       conditional independence
        # Now consider p(z|x3,x4)
        #   \propto p(z)p(x3|z)p(x4|z)              Bayes rule
        # [term]             [array]
        # p(z)             logp_crp
        # p(x3|z)p(x4|z)   logp_evidence
        # p(z|x3,x4)       logp_cluster
        # p(x1|z)p(x2|z)   logp_query

        if evidence is None:
            evidence = []

        logpdf = 0
        query_views = set([self.Zv[col] for (col, _) in query])

        for v in query_views:
            # Evidence densities.
            logp_crp = self._compute_cluster_crp_logps(v)
            logp_evidence = np.zeros(len(logp_crp))
            for (col, val) in evidence:
                if self.Zv[col] == v:
                    logp_evidence += self._compute_cluster_data_logps(col, val)
            logp_cluster = gu.log_normalize(logp_crp+logp_evidence)
            # Query densities.
            logp_query = np.zeros(len(logp_crp))
            for (col, val) in query:
                if self.Zv[col] == v:
                    logp_query += self._compute_cluster_data_logps(col, val)
            # Accumulate.
            logpdf += logsumexp(logp_query+logp_cluster)

        return logpdf

    def logpdf_bulk(self, rowids, queries, evidences=None):
        """Evaluate multiple queries at once, used by Engine."""
        if evidences is None:
            evidences = [[] for _ in xrange(len(rowids))]
        assert len(rowids) == len(queries) == len(evidences)
        logpdfs = []
        for rowid, query, evidence in zip(rowids, queries, evidences):
            logpdfs.append(self.logpdf(rowid, query, evidence))
        return logpdfs

    def logpdf_marginal(self):
        return gu.logp_crp(len(self.Zv), self.Nv, self.alpha) + \
            sum(v.logpdf_marginal() for v in self.views)

    # --------------------------------------------------------------------------
    # Simulate

    def simulate(self, rowid, query, evidence=None, N=1):
        """Simulate from the posterior predictive distirbution.

        Parameters
        ----------
        rowid : int
            The rowid of the member of the population to simulate from.
{'fraction_non_alphanumeric': 0.11370754035172248, 'fraction_numerical': 0.00481811611659841, 'mean_word_length': 1.897173432592765, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '597812', 'n_tokens_mistral': 7558, 'n_tokens_neox': 7251, 'n_words': 2373}