<commit_before># -*- coding: utf-8 -*-
"""
Created on Sun Jul 17 15:59:25 2016

@author: timothyb0912
"""

import unittest
import warnings
from collections import OrderedDict
from copy import deepcopy

import numpy as np
import pandas as pd
from scipy.sparse import csr_matrix
import numpy.testing as npt
import pylogit.mixed_logit_calcs as mlc
import pylogit.mixed_logit as mixed_logit


# Use the following to always show the warnings
np.seterr(all='warn')
warnings.simplefilter("always")


def temp_utility_transform(sys_utility_array, *args, **kwargs):
    """
    Parameters
    ----------
    sys_utility_array : numpy array.
        Should have 1D or 2D. Should have been created by the dot product of a
        design matrix and an array  of index coefficients.

    Returns
    -------
    2D numpy array.
        The returned array will contain a representation of the
        `sys_utility_array`. If `sys_utility_array` is 2D, then
        `sys_utility_array` will be returned unaltered. Else, the function will
        return `sys_utility_array[:, None]`.
    """
    # Return a 2D array of systematic utility values
    if len(sys_utility_array.shape) == 1:
        systematic_utilities = sys_utility_array[:, np.newaxis]
    else:
        systematic_utilities = sys_utility_array

    return systematic_utilities


class NormalDrawsTests(unittest.TestCase):

    def test_return_format(self):
        n_obs = 10
        n_draws = 5
        n_vars = 3
        random_draws = mlc.get_normal_draws(n_obs, n_draws, n_vars)

        self.assertIsInstance(random_draws, list)
        self.assertEqual(len(random_draws), n_vars)
        for draws in random_draws:
            self.assertIsInstance(draws, np.ndarray)
            self.assertAlmostEqual(draws.shape, (n_obs, n_draws))

        return None


class MixingNamesToPositions(unittest.TestCase):

    def test_convert_mixing_names_to_positions(self):
        fake_index_vars = ["foo", "bar", "cake", "cereal"]
        fake_mixing_vars = ["bar", "cereal"]
        args = (fake_mixing_vars, fake_index_vars)
        mix_pos = mlc.convert_mixing_names_to_positions(*args)

        self.assertIsInstance(mix_pos, list)
        self.assertEqual(len(mix_pos), len(fake_mixing_vars))
        for pos, idx_val in enumerate(mix_pos):
            current_var = fake_mixing_vars[pos]
            self.assertEqual(idx_val, fake_index_vars.index(current_var))
        return None


class MixedLogitCalculations(unittest.TestCase):

    # Note that for this set up, we will consider a situation with the
    # following parameters:
    # 3 Alternatives per individual
    # 2 Individuals
    # Individual 1 has 2 observed choice situations
    # Individual 2 has 1 observed choice situation
    # The true systematic utility depends on ASC_1, ASC_2, and a single X
    # The X coefficient is randomly distributed
    def setUp(self):
        # Fake random draws where Row 1 is for observation 1 and row 2 is
        # for observation 2. Column 1 is for draw 1 and column 2 is for draw 2
        self.fake_draws = np.array([[0.4, 0.8], [0.6, 0.2]])
        # Create the betas to be used during the tests
        self.fake_betas = np.array([0.3, -0.6, 0.2])
        self.fake_std = 1
        self.fake_betas_ext = np.concatenate((self.fake_betas,
                                              np.array([self.fake_std])),
                                             axis=0)

        # Create the fake design matrix with columns denoting ASC_1, ASC_2, X
        self.fake_design = np.array([[1, 0, 1],
                                     [0, 1, 2],
                                     [0, 0, 3],
                                     [1, 0, 1.5],
                                     [0, 1, 2.5],
                                     [0, 0, 3.5],
                                     [1, 0, 0.5],
                                     [0, 1, 1.0],
                                     [0, 0, 1.5]])
        # Record what positions in the design matrix are being mixed over
        self.mixing_pos = [2]

        # Create the arrays that specify the choice situation, individual id
        # and alternative ids
        self.situation_ids = np.array([1, 1, 1, 2, 2, 2, 3, 3, 3])
        self.individual_ids = np.array([1, 1, 1, 1, 1, 1, 2, 2, 2])
        self.alternative_ids = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3])
        # Create a fake array of choices
        self.choice_array = np.array([0, 1, 0, 0, 0, 1, 1, 0, 0])

        # Create the 'rows_to_mixers' sparse array for this dataset
        # Denote the rows that correspond to observation 1 and observation 2
        self.obs_1_rows = np.ones(self.fake_design.shape[0])
        # Make sure the rows for observation 2 are given a zero in obs_1_rows
        self.obs_1_rows[-3:] = 0
        self.obs_2_rows = 1 - self.obs_1_rows
        # Create the row_to_mixers scipy.sparse matrix
        self.fake_rows_to_mixers = csr_matrix(self.obs_1_rows[:, None] ==
                                              np.array([1, 0])[None, :])
        # Create the rows_to_obs scipy.sparse matrix
        self.fake_rows_to_obs = csr_matrix(self.situation_ids[:, None] ==
                                           np.arange(1, 4)[None, :])
        # Create the rows_to_alts scipy.sparse matrix
        self.fake_rows_to_alts = csr_matrix(self.alternative_ids[:, None] ==
                                            np.arange(1, 4)[None, :])

        # Create the design matrix that we should see for draw 1 and draw 2
        arrays_to_join = (self.fake_design.copy(),
                          self.fake_design.copy()[:, -1][:, None])
        self.fake_design_draw_1 = np.concatenate(arrays_to_join, axis=1)
        self.fake_design_draw_2 = self.fake_design_draw_1.copy()

        # Multiply the 'random' coefficient draws by the corresponding variable
        self.fake_design_draw_1[:, -1] *= (self.obs_1_rows *
                                           self.fake_draws[0, 0] +
                                           self.obs_2_rows *
                                           self.fake_draws[1, 0])
        self.fake_design_draw_2[:, -1] *= (self.obs_1_rows *
                                           self.fake_draws[0, 1] +
                                           self.obs_2_rows *
                                           self.fake_draws[1, 1])
        extended_design_draw_1 = self.fake_design_draw_1[:, None, :]
        extended_design_draw_2 = self.fake_design_draw_2[:, None, :]
        self.fake_design_3d = np.concatenate((extended_design_draw_1,
                                              extended_design_draw_2),
                                             axis=1)

        # Create the fake systematic utility values
        self.sys_utilities_draw_1 = (self.fake_design_draw_1
                                         .dot(self.fake_betas_ext))
        self.sys_utilities_draw_2 = (self.fake_design_draw_2
                                         .dot(self.fake_betas_ext))

        #####
        # Calculate the probabilities of each alternatve in each choice
        # situation
        #####
        long_exp_draw_1 = np.exp(self.sys_utilities_draw_1)
        long_exp_draw_2 = np.exp(self.sys_utilities_draw_2)
        ind_exp_sums_draw_1 = self.fake_rows_to_obs.T.dot(long_exp_draw_1)
        ind_exp_sums_draw_2 = self.fake_rows_to_obs.T.dot(long_exp_draw_2)
        long_exp_sum_draw_1 = self.fake_rows_to_obs.dot(ind_exp_sums_draw_1)
        long_exp_sum_draw_2 = self.fake_rows_to_obs.dot(ind_exp_sums_draw_2)
        long_probs_draw_1 = long_exp_draw_1 / long_exp_sum_draw_1
        long_probs_draw_2 = long_exp_draw_2 / long_exp_sum_draw_2
        self.prob_array = np.concatenate((long_probs_draw_1[:, None],
                                          long_probs_draw_2[:, None]),
                                         axis=1)

        ###########
        # Create a mixed logit object for later use.
        ##########
        # Create a fake old long format dataframe for mixed logit model object
        self.alt_id_column = "alt_id"
        self.situation_id_column = "situation_id"
        self.obs_id_column = "observation_id"
        self.choice_column = "choice"

        self.fake_old_df = pd.DataFrame({"x": self.fake_design[:, 2],
                                    self.alt_id_column: self.alternative_ids,
                                    self.situation_id_column:
                                                            self.situation_ids,
                                    self.obs_id_column: self.individual_ids,
                                    self.choice_column: self.choice_array})
        self.fake_old_df["intercept"] = 1

        # Create a fake specification
        self.fake_spec = OrderedDict()
        self.fake_names = OrderedDict()

        self.fake_spec["intercept"] = [1, 2]
        self.fake_names["intercept"] = ["ASC 1", "ASC 2"]

        self.fake_spec["x"] = [[1, 2, 3]]
        self.fake_names["x"] = ["beta_x"]

        # Specify the mixing variable
        self.fake_mixing_vars = ["beta_x"]

        # Create a fake version of a mixed logit model object
        args = [self.fake_old_df,<commit_msg>Made a tiny change to comply with PEP8.<commit_after># -*- coding: utf-8 -*-
"""
Created on Sun Jul 17 15:59:25 2016

@author: timothyb0912
"""

import unittest
import warnings
from collections import OrderedDict
from copy import deepcopy

import numpy as np
import pandas as pd
from scipy.sparse import csr_matrix
import numpy.testing as npt
import pylogit.mixed_logit_calcs as mlc
import pylogit.mixed_logit as mixed_logit


# Use the following to always show the warnings
np.seterr(all='warn')
warnings.simplefilter("always")


def temp_utility_transform(sys_utility_array, *args, **kwargs):
    """
    Parameters
    ----------
    sys_utility_array : numpy array.
        Should have 1D or 2D. Should have been created by the dot product of a
        design matrix and an array  of index coefficients.

    Returns
    -------
    2D numpy array.
        The returned array will contain a representation of the
        `sys_utility_array`. If `sys_utility_array` is 2D, then
        `sys_utility_array` will be returned unaltered. Else, the function will
        return `sys_utility_array[:, None]`.
    """
    # Return a 2D array of systematic utility values
    if len(sys_utility_array.shape) == 1:
        systematic_utilities = sys_utility_array[:, np.newaxis]
    else:
        systematic_utilities = sys_utility_array

    return systematic_utilities


class NormalDrawsTests(unittest.TestCase):

    def test_return_format(self):
        n_obs = 10
        n_draws = 5
        n_vars = 3
        random_draws = mlc.get_normal_draws(n_obs, n_draws, n_vars)

        self.assertIsInstance(random_draws, list)
        self.assertEqual(len(random_draws), n_vars)
        for draws in random_draws:
            self.assertIsInstance(draws, np.ndarray)
            self.assertAlmostEqual(draws.shape, (n_obs, n_draws))

        return None


class MixingNamesToPositions(unittest.TestCase):

    def test_convert_mixing_names_to_positions(self):
        fake_index_vars = ["foo", "bar", "cake", "cereal"]
        fake_mixing_vars = ["bar", "cereal"]
        args = (fake_mixing_vars, fake_index_vars)
        mix_pos = mlc.convert_mixing_names_to_positions(*args)

        self.assertIsInstance(mix_pos, list)
        self.assertEqual(len(mix_pos), len(fake_mixing_vars))
        for pos, idx_val in enumerate(mix_pos):
            current_var = fake_mixing_vars[pos]
            self.assertEqual(idx_val, fake_index_vars.index(current_var))
        return None


class MixedLogitCalculations(unittest.TestCase):

    # Note that for this set up, we will consider a situation with the
    # following parameters:
    # 3 Alternatives per individual
    # 2 Individuals
    # Individual 1 has 2 observed choice situations
    # Individual 2 has 1 observed choice situation
    # The true systematic utility depends on ASC_1, ASC_2, and a single X
    # The X coefficient is randomly distributed
    def setUp(self):
        # Fake random draws where Row 1 is for observation 1 and row 2 is
        # for observation 2. Column 1 is for draw 1 and column 2 is for draw 2
        self.fake_draws = np.array([[0.4, 0.8], [0.6, 0.2]])
        # Create the betas to be used during the tests
        self.fake_betas = np.array([0.3, -0.6, 0.2])
        self.fake_std = 1
        self.fake_betas_ext = np.concatenate((self.fake_betas,
                                              np.array([self.fake_std])),
                                             axis=0)

        # Create the fake design matrix with columns denoting ASC_1, ASC_2, X
        self.fake_design = np.array([[1, 0, 1],
                                     [0, 1, 2],
                                     [0, 0, 3],
                                     [1, 0, 1.5],
                                     [0, 1, 2.5],
                                     [0, 0, 3.5],
                                     [1, 0, 0.5],
                                     [0, 1, 1.0],
                                     [0, 0, 1.5]])
        # Record what positions in the design matrix are being mixed over
        self.mixing_pos = [2]

        # Create the arrays that specify the choice situation, individual id
        # and alternative ids
        self.situation_ids = np.array([1, 1, 1, 2, 2, 2, 3, 3, 3])
        self.individual_ids = np.array([1, 1, 1, 1, 1, 1, 2, 2, 2])
        self.alternative_ids = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3])
        # Create a fake array of choices
        self.choice_array = np.array([0, 1, 0, 0, 0, 1, 1, 0, 0])

        # Create the 'rows_to_mixers' sparse array for this dataset
        # Denote the rows that correspond to observation 1 and observation 2
        self.obs_1_rows = np.ones(self.fake_design.shape[0])
        # Make sure the rows for observation 2 are given a zero in obs_1_rows
        self.obs_1_rows[-3:] = 0
        self.obs_2_rows = 1 - self.obs_1_rows
        # Create the row_to_mixers scipy.sparse matrix
        self.fake_rows_to_mixers = csr_matrix(self.obs_1_rows[:, None] ==
                                              np.array([1, 0])[None, :])
        # Create the rows_to_obs scipy.sparse matrix
        self.fake_rows_to_obs = csr_matrix(self.situation_ids[:, None] ==
                                           np.arange(1, 4)[None, :])
        # Create the rows_to_alts scipy.sparse matrix
        self.fake_rows_to_alts = csr_matrix(self.alternative_ids[:, None] ==
                                            np.arange(1, 4)[None, :])

        # Create the design matrix that we should see for draw 1 and draw 2
        arrays_to_join = (self.fake_design.copy(),
                          self.fake_design.copy()[:, -1][:, None])
        self.fake_design_draw_1 = np.concatenate(arrays_to_join, axis=1)
        self.fake_design_draw_2 = self.fake_design_draw_1.copy()

        # Multiply the 'random' coefficient draws by the corresponding variable
        self.fake_design_draw_1[:, -1] *= (self.obs_1_rows *
                                           self.fake_draws[0, 0] +
                                           self.obs_2_rows *
                                           self.fake_draws[1, 0])
        self.fake_design_draw_2[:, -1] *= (self.obs_1_rows *
                                           self.fake_draws[0, 1] +
                                           self.obs_2_rows *
                                           self.fake_draws[1, 1])
        extended_design_draw_1 = self.fake_design_draw_1[:, None, :]
        extended_design_draw_2 = self.fake_design_draw_2[:, None, :]
        self.fake_design_3d = np.concatenate((extended_design_draw_1,
                                              extended_design_draw_2),
                                             axis=1)

        # Create the fake systematic utility values
        self.sys_utilities_draw_1 = (self.fake_design_draw_1
                                         .dot(self.fake_betas_ext))
        self.sys_utilities_draw_2 = (self.fake_design_draw_2
                                         .dot(self.fake_betas_ext))

        #####
        # Calculate the probabilities of each alternatve in each choice
        # situation
        #####
        long_exp_draw_1 = np.exp(self.sys_utilities_draw_1)
        long_exp_draw_2 = np.exp(self.sys_utilities_draw_2)
        ind_exp_sums_draw_1 = self.fake_rows_to_obs.T.dot(long_exp_draw_1)
        ind_exp_sums_draw_2 = self.fake_rows_to_obs.T.dot(long_exp_draw_2)
        long_exp_sum_draw_1 = self.fake_rows_to_obs.dot(ind_exp_sums_draw_1)
        long_exp_sum_draw_2 = self.fake_rows_to_obs.dot(ind_exp_sums_draw_2)
        long_probs_draw_1 = long_exp_draw_1 / long_exp_sum_draw_1
        long_probs_draw_2 = long_exp_draw_2 / long_exp_sum_draw_2
        self.prob_array = np.concatenate((long_probs_draw_1[:, None],
                                          long_probs_draw_2[:, None]),
                                         axis=1)

        ###########
        # Create a mixed logit object for later use.
        ##########
        # Create a fake old long format dataframe for mixed logit model object
        self.alt_id_column = "alt_id"
        self.situation_id_column = "situation_id"
        self.obs_id_column = "observation_id"
        self.choice_column = "choice"

        data = {"x": self.fake_design[:, 2],
                self.alt_id_column: self.alternative_ids,
                self.situation_id_column: self.situation_ids,
                self.obs_id_column: self.individual_ids,
                self.choice_column: self.choice_array}
        self.fake_old_df = pd.DataFrame(data)
        self.fake_old_df["intercept"] = 1

        # Create a fake specification
        self.fake_spec = OrderedDict()
        self.fake_names = OrderedDict()

        self.fake_spec["intercept"] = [1, 2]
        self.fake_names["intercept"] = ["ASC 1", "ASC 2"]

        self.fake_spec["x"] = [[1, 2, 3]]
        self.fake_names["x"] = ["beta_x"]

        # Specify the mixing variable
        self.fake_mixing_vars = ["beta_x"]

        # Create a fake version of a mixed logit model object
        args = [self.fake_old_df,
{'fraction_non_alphanumeric': 0.09171403585483166, 'fraction_numerical': 0.023557061652820288, 'mean_word_length': 1.8805100755667505, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 106, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1328946', 'n_tokens_mistral': 5776, 'n_tokens_neox': 5436, 'n_words': 1364}