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

# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from math import log
from math import pi
from math import sqrt

from collections import OrderedDict
from collections import namedtuple

import numpy as np

from numpy.linalg import det

from scipy.special import gammaln

from cgpm.cgpm import CGpm
from cgpm.mixtures.dim import Dim
from cgpm.utils import config as cu
from cgpm.utils import data as du
from cgpm.utils import general as gu


LOG2PI = log(2*pi)
Data = namedtuple('Data', ['x', 'Y'])


class LinearRegression(CGpm):
    """Bayesian linear model with normal prior on regression parameters and
    inverse-gamma prior on both observation and regression variance.

    Reference
    http://www.biostat.umn.edu/~ph7440/pubh7440/BayesianLinearModelGoryDetails.pdf


    Y_i = w' X_i + \sigma^2
        Response data                   Y_i \in R
        Covariate vector                X_i \in R^p
        Regression coefficients         w \in R^p
        Regression variance             \sigma^2 \in R

    Hyperparameters:                    a=1, b=1, V=I, mu=[0], dimension=p

    Parameters:                         \sigma2 ~ Inverse-Gamma(a, b)
                                        w ~ MVNormal(\mu, \sigma2*I)

    Data                                Y_i|x_i ~ Normal(w' x_i, \sigma2)
    """

    def __init__(self, outputs, inputs, hypers=None, params=None, distargs=None,
            rng=None):
        # io data.
        self.outputs = outputs
        self.inputs = inputs
        self.rng = gu.gen_rng() if rng is None else rng
        assert len(self.outputs) == 1
        assert len(self.inputs) >= 1
        assert self.outputs[0] not in self.inputs
        assert len(distargs['inputs']['stattypes']) == len(self.inputs)
        self.input_cctypes = distargs['inputs']['stattypes']
        self.input_ccargs = distargs['inputs']['statargs']
        # Add a heuristic imputer?
        self.impute = distargs.get('impute', None)
        # Determine number of covariates (with 1 bias term) and number of
        # categories for categorical covariates.
        p, counts = zip(*[
            self._predictor_count(cctype, ccarg) for cctype, ccarg
            in zip(self.input_cctypes, self.input_ccargs)])
        self.p = sum(p)+1
        self.inputs_discrete = {i:c for i, c in enumerate(counts) if c}
        # For numerical covariates, map index in inputs to index in code.
        self.lookup_numerical_index = self.input_to_code_index()
        # Dataset.
        self.N = 0
        self.data = Data(x=OrderedDict(), Y=OrderedDict())
        # Hyper parameters.
        if hypers is None: hypers = {}
        self.a = hypers.get('a', 1.)
        self.b = hypers.get('b', 1.)
        self.mu = hypers.get('mu', np.zeros(self.p))
        self.V = hypers.get('V', np.eye(self.p))

    def incorporate(self, rowid, query, evidence=None):
        assert rowid not in self.data.x
        assert rowid not in self.data.Y
        if self.outputs[0] not in query:
            raise ValueError('No query in incorporate: %s.' % query)
        x, y = self.preprocess(query, evidence)
        self.N += 1
        self.data.x[rowid] = x
        self.data.Y[rowid] = y

    def unincorporate(self, rowid):
        try:
            del self.data.x[rowid]
            del self.data.Y[rowid]
        except KeyError:
            raise ValueError('No such observation: %d' % rowid)
        self.N -= 1

    def logpdf(self, rowid, query, evidence=None):
        xt, yt = self.preprocess(query, evidence)
        assert rowid not in self.data.x
        return LinearRegression.calc_predictive_logp(
            xt, yt, self.N, self.data.Y.values(), self.data.x.values(), self.a,
            self.b, self.mu, self.V)

    def simulate(self, rowid, query, evidence=None, N=None):
        if N is not None:
            return [self.simulate(rowid, query, evidence) for i in xrange(N)]
        assert query == self.outputs
        if rowid in self.data.x:
            return {self.outputs[0]: self.data.x[rowid]}
        xt, yt = self.preprocess(None, evidence)
        an, bn, mun, Vn_inv = LinearRegression.posterior_hypers(
            self.N, self.data.Y.values(), self.data.x.values(), self.a, self.b,
            self.mu, self.V)
        sigma2, b = LinearRegression.sample_parameters(
            an, bn, mun, np.linalg.inv(Vn_inv), self.rng)
        x = self.rng.normal(np.dot(yt, b), np.sqrt(sigma2))
        return {self.outputs[0]: x}

    def logpdf_score(self):
        return LinearRegression.calc_logpdf_marginal(
            self.N, self.data.Y.values(), self.data.x.values(),
            self.a, self.b, self.mu, self.V)

    ##################
    # NON-GPM METHOD #
    ##################

    def transition(self, N=None):
        self.transition_hypers(N=N)

    def transition_hypers(self, N=None):
        if N is None:<commit_msg>Allow user to simulate_params from linreg.py.<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from math import log
from math import pi
from math import sqrt

from collections import OrderedDict
from collections import namedtuple

import numpy as np

from numpy.linalg import det

from scipy.special import gammaln

from cgpm.cgpm import CGpm
from cgpm.mixtures.dim import Dim
from cgpm.utils import config as cu
from cgpm.utils import data as du
from cgpm.utils import general as gu


LOG2PI = log(2*pi)
Data = namedtuple('Data', ['x', 'Y'])


class LinearRegression(CGpm):
    """Bayesian linear model with normal prior on regression parameters and
    inverse-gamma prior on both observation and regression variance.

    Reference
    http://www.biostat.umn.edu/~ph7440/pubh7440/BayesianLinearModelGoryDetails.pdf


    Y_i = w' X_i + \sigma^2
        Response data                   Y_i \in R
        Covariate vector                X_i \in R^p
        Regression coefficients         w \in R^p
        Regression variance             \sigma^2 \in R

    Hyperparameters:                    a=1, b=1, V=I, mu=[0], dimension=p

    Parameters:                         \sigma2 ~ Inverse-Gamma(a, b)
                                        w ~ MVNormal(\mu, \sigma2*I)

    Data                                Y_i|x_i ~ Normal(w' x_i, \sigma2)
    """

    def __init__(self, outputs, inputs, hypers=None, params=None, distargs=None,
            rng=None):
        # io data.
        self.outputs = outputs
        self.inputs = inputs
        self.rng = gu.gen_rng() if rng is None else rng
        assert len(self.outputs) == 1
        assert len(self.inputs) >= 1
        assert self.outputs[0] not in self.inputs
        assert len(distargs['inputs']['stattypes']) == len(self.inputs)
        self.input_cctypes = distargs['inputs']['stattypes']
        self.input_ccargs = distargs['inputs']['statargs']
        # Add a heuristic imputer?
        self.impute = distargs.get('impute', None)
        # Determine number of covariates (with 1 bias term) and number of
        # categories for categorical covariates.
        p, counts = zip(*[
            self._predictor_count(cctype, ccarg) for cctype, ccarg
            in zip(self.input_cctypes, self.input_ccargs)])
        self.p = sum(p)+1
        self.inputs_discrete = {i:c for i, c in enumerate(counts) if c}
        # For numerical covariates, map index in inputs to index in code.
        self.lookup_numerical_index = self.input_to_code_index()
        # Dataset.
        self.N = 0
        self.data = Data(x=OrderedDict(), Y=OrderedDict())
        # Hyper parameters.
        if hypers is None: hypers = {}
        self.a = hypers.get('a', 1.)
        self.b = hypers.get('b', 1.)
        self.mu = hypers.get('mu', np.zeros(self.p))
        self.V = hypers.get('V', np.eye(self.p))

    def incorporate(self, rowid, query, evidence=None):
        assert rowid not in self.data.x
        assert rowid not in self.data.Y
        if self.outputs[0] not in query:
            raise ValueError('No query in incorporate: %s.' % query)
        x, y = self.preprocess(query, evidence)
        self.N += 1
        self.data.x[rowid] = x
        self.data.Y[rowid] = y

    def unincorporate(self, rowid):
        try:
            del self.data.x[rowid]
            del self.data.Y[rowid]
        except KeyError:
            raise ValueError('No such observation: %d' % rowid)
        self.N -= 1

    def logpdf(self, rowid, query, evidence=None):
        xt, yt = self.preprocess(query, evidence)
        assert rowid not in self.data.x
        return LinearRegression.calc_predictive_logp(
            xt, yt, self.N, self.data.Y.values(), self.data.x.values(), self.a,
            self.b, self.mu, self.V)

    def simulate(self, rowid, query, evidence=None, N=None):
        if N is not None:
            return [self.simulate(rowid, query, evidence) for i in xrange(N)]
        assert query == self.outputs
        if rowid in self.data.x:
            return {self.outputs[0]: self.data.x[rowid]}
        xt, yt = self.preprocess(None, evidence)
        sigma2, b = self.simulate_params()
        x = self.rng.normal(np.dot(yt, b), np.sqrt(sigma2))
        return {self.outputs[0]: x}

    def logpdf_score(self):
        return LinearRegression.calc_logpdf_marginal(
            self.N, self.data.Y.values(), self.data.x.values(),
            self.a, self.b, self.mu, self.V)

    def simulate_params(self):
        an, bn, mun, Vn_inv = LinearRegression.posterior_hypers(
            self.N, self.data.Y.values(), self.data.x.values(), self.a, self.b,
            self.mu, self.V)
        return LinearRegression.sample_parameters(
            an, bn, mun, np.linalg.inv(Vn_inv), self.rng)

    ##################
    # NON-GPM METHOD #
    ##################

    def transition(self, N=None):
        self.transition_hypers(N=N)

    def transition_hypers(self, N=None):
        if N is None:
{'fraction_non_alphanumeric': 0.11353315168029064, 'fraction_numerical': 0.008356039963669392, 'mean_word_length': 2.3105832832230906, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1972655', 'n_tokens_mistral': 3533, 'n_tokens_neox': 3351, 'n_words': 1012}