"""
General tests for all estimators in sklearn.
"""

# Authors: Andreas Mueller <amueller@ais.uni-bonn.de>
#          Gael Varoquaux gael.varoquaux@normalesup.org
# License: BSD 3 clause
from __future__ import print_function

import os
import warnings
import sys
import pkgutil

from sklearn.externals.six import PY3
from sklearn.utils.testing import assert_false, clean_warning_registry
from sklearn.utils.testing import all_estimators
from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import assert_in
from sklearn.utils.testing import ignore_warnings

import sklearn
from sklearn.cluster.bicluster import BiclusterMixin

from sklearn.linear_model.base import LinearClassifierMixin
from sklearn.utils.estimator_checks import (
    _yield_all_checks,
    CROSS_DECOMPOSITION,
    check_parameters_default_constructible,
    check_class_weight_balanced_linear_classifier,
    check_transformer_n_iter,
    check_non_transformer_estimators_n_iter,
    check_get_params_invariance)

def test_all_estimator_no_base_class():
    # test that all_estimators doesn't find abstract classes.
    for name, Estimator in all_estimators():
        msg = ("Base estimators such as {0} should not be included"
               " in all_estimators").format(name)
        assert_false(name.lower().startswith('base'), msg=msg)

def test_all_estimators():
    # Test that estimators are default-constructible, clonable
    # and have working repr.
    estimators = all_estimators(include_meta_estimators=True)

    # Meta sanity-check to make sure that the estimator introspection runs
    # properly
    assert_greater(len(estimators), 0)

    for name, Estimator in estimators:
        # some can just not be sensibly default constructed
        yield check_parameters_default_constructible, name, Estimator

def test_non_meta_estimators():
    # input validation etc for non-meta estimators
    estimators = all_estimators()
    for name, Estimator in estimators:
        if issubclass(Estimator, BiclusterMixin):
            continue
        if name.startswith("_"):
            continue
        for check in _yield_all_checks(name, Estimator):
            yield check, name, Estimator

def test_configure():
    # Smoke test the 'configure' step of setup, this tests all the
    # 'configure' functions in the setup.pys in the scikit
    cwd = os.getcwd()
    setup_path = os.path.abspath(os.path.join(sklearn.__path__[0], '..'))
    setup_filename = os.path.join(setup_path, 'setup.py')
    if not os.path.exists(setup_filename):
        return
    try:
        os.chdir(setup_path)
        old_argv = sys.argv
        sys.argv = ['setup.py', 'config']
        clean_warning_registry()
        with warnings.catch_warnings():
            # The configuration spits out warnings when not finding
            # Blas/Atlas development headers
            warnings.simplefilter('ignore', UserWarning)
            if PY3:
                with open('setup.py') as f:
                    exec(f.read(), dict(__name__='__main__'))
            else:
                execfile('setup.py', dict(__name__='__main__'))
    finally:
        sys.argv = old_argv
        os.chdir(cwd)

def test_class_weight_balanced_linear_classifiers():
    classifiers = all_estimators(type_filter='classifier')

    clean_warning_registry()
    with warnings.catch_warnings(record=True):
        linear_classifiers = [
            (name, clazz)
            for name, clazz in classifiers
            if 'class_weight' in clazz().get_params().keys()
               and issubclass(clazz, LinearClassifierMixin)]

    for name, Classifier in linear_classifiers:
        if name == "LogisticRegressionCV":
            # Contrary to RidgeClassifierCV, LogisticRegressionCV use actual
            # CV folds and fit a model for each CV iteration before averaging
            # the coef. Therefore it is expected to not behave exactly as the
            # other linear model.
            continue
        yield check_class_weight_balanced_linear_classifier, name, Classifier

@ignore_warnings
def test_import_all_consistency():
    # Smoke test to check that any name in a __all__ list is actually defined
    # in the namespace of the module or package.
    pkgs = pkgutil.walk_packages(path=sklearn.__path__, prefix='sklearn.',
                                 onerror=lambda _: None)
    submods = [modname for _, modname, _ in pkgs]
    for modname in submods + ['sklearn']:
        if ".tests." in modname:
            continue
        package = __import__(modname, fromlist="dummy")
        for name in getattr(package, '__all__', ()):
            if getattr(package, name, None) is None:
                raise AttributeError(
                    "Module '{0}' has no attribute '{1}'".format(
                        modname, name))

def test_root_import_all_completeness():
    EXCEPTIONS = ('utils', 'tests', 'base', 'setup')
    for _, modname, _ in pkgutil.walk_packages(path=sklearn.__path__,
                                               onerror=lambda _: None):
        if '.' in modname or modname.startswith('_') or modname in EXCEPTIONS:
            continue
        assert_in(modname, sklearn.__all__)

def test_non_transformer_estimators_n_iter():
    # Test that all estimators of type which are non-transformer
    # and which have an attribute of max_iter, return the attribute
    # of n_iter atleast 1.
    for est_type in ['regressor', 'classifier', 'cluster']:
        regressors = all_estimators(type_filter=est_type)
        for name, Estimator in regressors:
            # LassoLars stops early for the default alpha=1.0 for
            # the iris dataset.
            if name == 'LassoLars':
                estimator = Estimator(alpha=0.)
            else:
                estimator = Estimator()
            if hasattr(estimator, "max_iter"):
                # These models are dependent on external solvers like
                # libsvm and accessing the iter parameter is non-trivial.
                if name in (['Ridge', 'SVR', 'NuSVR', 'NuSVC',
                             'RidgeClassifier', 'SVC', 'RandomizedLasso',
                             'LogisticRegressionCV']):
                    continue

                # Tested in test_transformer_n_iter below
                elif (name in CROSS_DECOMPOSITION or
                      name in ['LinearSVC', 'LogisticRegression']):
                    continue

                else:
                    # Multitask models related to ENet cannot handle
                    # if y is mono-output.
                    yield (check_non_transformer_estimators_n_iter,
                           name, estimator, 'Multi' in name)

def test_transformer_n_iter():
    transformers = all_estimators(type_filter='transformer')
    for name, Estimator in transformers:
        estimator = Estimator()
        # Dependent on external solvers and hence accessing the iter
        # param is non-trivial.
        external_solver = ['Isomap', 'KernelPCA', 'LocallyLinearEmbedding',
                           'RandomizedLasso', 'LogisticRegressionCV']

        if hasattr(estimator, "max_iter") and name not in external_solver:
            yield check_transformer_n_iter, name, estimator

def test_get_params_invariance():
    # Test for estimators that support get_params, that
    # get_params(deep=False) is a subset of get_params(deep=True)
    # Related to issue #4465

    estimators = all_estimators(include_meta_estimators=False, include_other=True)
    for name, Estimator in estimators:
        if hasattr(Estimator, 'get_params'):
            yield check_get_params_invariance, name, Estimator


# -*- coding: utf-8 -*-
"""
Name        : radcenter_distribution
Author      : Joan Juvert <trust.no.one.51@gmail.com>
Version     : 1.0
Description : This script calculates the influence of the distribution of
            : radiative centers in the active layer on the observed
            : spectrum.

Copyright 2012 Joan Juvert

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import multilayers as ml
import numpy as np
import bphysics as bp
import scipy.integrate as integ
import argparse as ap
import sys
import pdb

# Argument parsing
parser = ap.ArgumentParser(
        description = "This script calculates the effect of the " + \
        "distribution of radiative centers in the active layer on " + \
        "the modificator to the spectrum. The observation angle is " + \
        "a fixed parameter. Optionally, the output can be plotted " + \
        "and output to the standard output or to a file. The matrix " + \
        "containing the values of F(z, lambda) can be saved to a file " + \
        "and recovered in a following run of the program to avoid " + \
        "recalculating it in case we want to calculate the effect of " + \
        "different distributions on the same system.")
parser.add_argument(
        "--graph",
        help = "Plot the results",
        action = "store_true")
parser.add_argument(
        "-o",
        "--output",
        help = "Dump the results to a file")
parser.add_argument(
        "-s",
        "--savematrix",
        help = "Save the matrix with the F(z, lambda) values to a file")
parser.add_argument(
        "-l",
        "--loadmatrix",
        help = "Load the matrix with the F(z, lambda) values from a file")
args = parser.parse_args()

# Load the depth distribution of radiative centers. Note that the origin
# and units of z must be the same as in the multilayer.The distribution
# should be normalized to 1.
print("Loading the distribution...")
path = "/home/joan/Dropbox/CNM/projectes/simulations_report/figures/" + \
        "rcdistributions/"
distribution = bp.rdfile(path + "gaussian_m25_s07.dat", usecols = [0, 1])[1]
print("Done")

print("Checking the distribution...")
integral = integ.simps(distribution[:, 1], distribution[:, 0], 0)
np.testing.assert_almost_equal(integral, 1, 2)
print("Done")

# If we load the values of F(z, lambda) calculated in a previous
# execution we do not need to build the multilayer and repeat the
# calculation of the F function. Notice that the values of z at which
# the new distribution is sampled should be the same as the previous
# one.
if args.loadmatrix:
    print("Loading matrix...")
    fmatrix = np.load(args.loadmatrix)
    zlist = fmatrix['zlist']
    np.testing.assert_array_equal(zlist, distribution[:, 0])
    wlist = fmatrix['wlist']
    angle = fmatrix['angle']
    fte = fmatrix['fte']
    ftm = fmatrix['ftm']
    print("Done")
else:
    # Create the materials
    print("Loading materials... ")
    silicon = ml.Medium("silicon.dat")
    air = ml.Medium("air.dat")
    sio2 = ml.Medium("sio2.dat")
    poly = ml.Medium("polysilicon.dat")
    print("Done")

    # Set the fixed parameters.
    angle = np.deg2rad(0)

    # Create the multilayer
    print("Building multilayer and allocating memory... ")
    thicknesses = [300, 50]
    multilayer = ml.Multilayer([
            air,
            [poly, thicknesses[0]],
            [sio2, thicknesses[1]],
            silicon])

    # Define the wavelengths and z coordinates at which F will be calculated
    # and allocate memory for the results. We will use a structured array to
    # store the values of F(z, lambda).
    wstep = 1
    wmin = multilayer.getMinMaxWlength()[0]
    wmax = multilayer.getMinMaxWlength()[1]
    wlist = np.arange(wmin, wmax, wstep)
    zlist = distribution[:, 0]

    ftype = np.dtype([
            ('fx', np.complex128),
            ('fy', np.complex128),
            ('fz', np.complex128)])
    resmatrix = np.empty((zlist.size, wlist.size), dtype = ftype)
    print("Done")

    # I(wavelength, theta) = s(wavelength) * F'(wavelength, theta), where
    # F'(wav, theta) = integral[z](|F|^2 * rcdist(z). Therefore, we
    # calculate the new spectrum as a modification to the original spectrum.
    # The modification factor F'(wav, theta) is an integral over z.

    # First calculate |Fy|^2 for te and |Fx*cos^2 + Fz*sin^2|^2 for tm. We
    # do fx and fz in one loop and fy in another independent loop to avoid
    # recalculating the characteristic matrix at every iteration due to the
    # change of polarization.
    print("Calculating F...")
    for (widx, wlength) in enumerate(wlist):
        percent = (float(widx) / wlist.size) * 100
        print("%.2f%%" % percent)
        for (zidx, z) in enumerate(zlist):
            resmatrix[zidx][widx]['fx'] = multilayer.calculateFx(z, wlength, angle)
            resmatrix[zidx][widx]['fz'] = multilayer.calculateFz(z, wlength, angle)
        for (zidx, z) in enumerate(zlist):
            resmatrix[zidx][widx]['fy'] = multilayer.calculateFy(z, wlength, angle)

    # We are probably more interesed on the effect of the multilayer on the
    # energy rather than the electric field. What we want is |Fy(z)|^2 for
    # TE waves and |Fx(z) cosA^2 + Fz(z) sinA^2|^2 for TM waves.
    ftm = np.absolute(
            resmatrix['fx'] * np.cos(angle) ** 2 + \
            resmatrix['fz'] * np.sin(angle) ** 2) ** 2
    fte = np.absolute(resmatrix['fy']) ** 2
    print("Done")

    # Notice that until now we have not used the distribution of the
    # radiative ceneters, but the calculation of ftm and fte is costly.
    # If requested, we can save fte and ftm to a file. In a following
    # execution of the script, the matrix can be loaded from the file
    # instead of recalculated.
    if args.savematrix:
        print("Saving matrix...")
        np.savez(args.savematrix, fte = fte, ftm = ftm, zlist = zlist,
                wlist = wlist, angle = angle)
        print("Done")

# Build or load the original spectrum. It should be sampled at the same
# wavelengths defined in wlist. If we are interested only in the
# modificator to the spectrum, not in the modified spectrum, we can
# leave it at 1.
original_spec = 1

# Multiply each F(z, lambda) by the distribution.
print("Integrating...")
distval = distribution[:, 1].reshape(distribution[:, 1].size, 1)
fte_mplied = fte * distval
ftm_mplied = ftm * distval
fte_int = integ.simps(fte_mplied, zlist, axis = 0)
ftm_int = integ.simps(ftm_mplied, zlist, axis = 0)
spectrum_modte = original_spec * fte_int
spectrum_modtm = original_spec * ftm_int
print("Done")

# Dump data to file or stdout
comments = "# F_TE = |Fy^2|^2\n" + \
        "# F_TM = |Fx * cosA^2 + Fz * sinA^2|^2\n" + \
        "# Modified spectrum for TE and TM waves for a\n" + \
        "# distributions of the radiative centers.\n" + \
        "# wlength\tF_TE\tF_TM"

if args.output:
    bp.wdfile(args.output, comments,
            np.array([wlist, spectrum_modte, spectrum_modtm]).T, '%.6e')
else:
    print(comments)
    for i in xrange(wlist.size):
        print("%.6e\t%.6e\t%.6e" % (wlist[i], spectrum_modte[i],
                spectrum_modtm[i]))

# Plot data if requested
if args.graph:
    import matplotlib.pyplot as plt

    plt.plot(wlist, spectrum_modte, label='TE', color = 'r')
    plt.plot(wlist, spectrum_modtm, label='TM', color = 'b')
    plt.xlabel('Wavelength (nm)')
    plt.ylabel('Energy ratio')
    plt.grid()
    plt.legend(loc=2)
    plt.title('%.1f rad' % angle)
    plt.show()
    plt.close()


"""
===================================
Getting the location of the planets
===================================

How to get the position of planetary bodies im the solar system using
`astropy's solar system ephemeris <http://docs.astropy.org/en/stable/coordinates/solarsystem.html#solar-system-ephemerides>`__ information and SunPy.
"""
import matplotlib.pyplot as plt

from astropy.time import Time

from sunpy.coordinates import get_body_heliographic_stonyhurst

##############################################################################
# Lets grab the positions of each of the planets in Heliographic Stonyhurst
# coordinates.
obstime = Time('2014-05-15T07:54:00.005')
planet_list = ['earth', 'venus', 'mars', 'mercury', 'jupiter', 'neptune', 'uranus', 'sun']
planet_coord = [get_body_heliographic_stonyhurst(this_planet, time=obstime) for this_planet in planet_list]

##############################################################################
# Let's plot the results. Remember the Sun is at the center of this coordinate
# system.
ax = plt.subplot(projection='polar')
for this_planet, this_coord in zip(planet_list, planet_coord):
    plt.polar(this_coord.lon.to('rad'), this_coord.radius, 'o', label=this_planet)
plt.legend()
plt.show()


# -*- coding: utf-8 -*-
"""
Created on Fri Oct  2 13:47:48 2015

@author: GILLES Armand
"""

import pandas as pd
import glob

df = pd.read_csv('source/PajeCom2009.csv', sep=";")

df.columns = ['Communes', 'Codes_Insee', 'NB_Allocataires_2009',
              'ALL_PAJE_2009', 'ALL_PRIM_2009', 'ALL_BASEP_2009',
              'ALL_ASMA_2009','ALL_Clca_Colca_2009']

files = glob.glob('source/PajeCom*')

for path_file in files:
    year = str(path_file[-8:-4])
    if (year != '2009'):
        df_temp = pd.read_csv(path_file, sep=';')

        # Rename Col with year
        year_col = ['Communes', 'Codes_Insee']
        features_col = []
        for col in df_temp.columns[2:]:
            year_col.append(col +"_"+ year)
            features_col.append(col +"_"+ year)

        # Adding key for mergeing
        features_col.append('Codes_Insee')
        df_temp.columns = year_col
        df = pd.merge(df, df_temp[features_col], how='inner', on='Codes_Insee')

# Rename col to have unique name in futur merge
list_col = []
for col in df.columns:
    if "nb_allocataires" in col.lower(): # NB_Allocataires (2009) != NB_allocataires (2010)
        list_col.append(col+"_PC") # PC = PageCom
    else:
        list_col.append(col)
df.columns = list_col

df.to_csv('data/full_PageCom.csv', encoding='utf-8', index=False)

## Features
#u'NB_Allocataires_2009_PC',
#       u'ALL_PAJE_2009', u'ALL_PRIM_2009', u'ALL_BASEP_2009', u'ALL_ASMA_2009',
#       u'ALL_Clca_Colca_2009', u'NB_Allocataires_2010_PC', u'ALL_PAJE_2010',
#       u'ALL_PRIM_2010', u'ALL_BASEP_2010', u'ALL_ASMA_2010',
#       u'ALL_Clca_Colca_2010', u'NB_Allocataires_2011_PC', u'ALL_PAJE_2011',
#       u'ALL_PRIM_2011', u'ALL_BASEP_2011', u'ALL_ASMA_2011',
#       u'ALL_Clca_Colca_2011', u'NB_Allocataires_2012_PC', u'ALL_PAJE_2012',
#       u'ALL_PRIM_2012', u'ALL_BASEP_2012', u'ALL_ASMA_2012',
#       u'ALL_Clca_Colca_2012', u'NB_Allocataires_2013_PC', u'ALL_PAJE_2013',
#       u'ALL_PRIM_2013', u'ALL_BASEP_2013', u'ALL_ASMA_2013',
#       u'ALL_Clca_Colca_2013', u'NB_Allocataires_2014_PC', u'ALL_PAJE_2014',
#       u'ALL_PRIM_2014', u'ALL_BASEP_2014', u'ALL_CMG_2014',
#       u'ALL_CMG_ASMA_2014', u'ALL_CMG_DOM_2014', u'ALL_CMG_A_2014',
#       u'ALL_Clca_Colca_2014', u'NB_Allocataires_2015_PC', u'ALL_PAJE_2015',
#       u'ALL_PRIM_2015', u'ALL_BASEP_2015', u'ALL_ASMA_2015',
#       u'ALL_Clca_Colca_2015'


# Author: Vlad Niculae
# License: BSD 3 clause

import sys

import numpy as np

from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import SkipTest
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_false
from sklearn.utils.testing import if_safe_multiprocessing_with_blas

from sklearn.decomposition import SparsePCA, MiniBatchSparsePCA
from sklearn.utils import check_random_state

def generate_toy_data(n_components, n_samples, image_size, random_state=None):
    n_features = image_size[0] * image_size[1]

    rng = check_random_state(random_state)
    U = rng.randn(n_samples, n_components)
    V = rng.randn(n_components, n_features)

    centers = [(3, 3), (6, 7), (8, 1)]
    sz = [1, 2, 1]
    for k in range(n_components):
        img = np.zeros(image_size)
        xmin, xmax = centers[k][0] - sz[k], centers[k][0] + sz[k]
        ymin, ymax = centers[k][1] - sz[k], centers[k][1] + sz[k]
        img[xmin:xmax][:, ymin:ymax] = 1.0
        V[k, :] = img.ravel()

    # Y is defined by : Y = UV + noise
    Y = np.dot(U, V)
    Y += 0.1 * rng.randn(Y.shape[0], Y.shape[1])  # Add noise
    return Y, U, V

# SparsePCA can be a bit slow. To avoid having test times go up, we
# test different aspects of the code in the same test

def test_correct_shapes():
    rng = np.random.RandomState(0)
    X = rng.randn(12, 10)
    spca = SparsePCA(n_components=8, random_state=rng)
    U = spca.fit_transform(X)
    assert_equal(spca.components_.shape, (8, 10))
    assert_equal(U.shape, (12, 8))
    # test overcomplete decomposition
    spca = SparsePCA(n_components=13, random_state=rng)
    U = spca.fit_transform(X)
    assert_equal(spca.components_.shape, (13, 10))
    assert_equal(U.shape, (12, 13))

def test_fit_transform():
    alpha = 1
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    spca_lars = SparsePCA(n_components=3, method='lars', alpha=alpha,
                          random_state=0)
    spca_lars.fit(Y)

    # Test that CD gives similar results
    spca_lasso = SparsePCA(n_components=3, method='cd', random_state=0,
                           alpha=alpha)
    spca_lasso.fit(Y)
    assert_array_almost_equal(spca_lasso.components_, spca_lars.components_)

@if_safe_multiprocessing_with_blas
def test_fit_transform_parallel():
    alpha = 1
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    spca_lars = SparsePCA(n_components=3, method='lars', alpha=alpha,
                          random_state=0)
    spca_lars.fit(Y)
    U1 = spca_lars.transform(Y)
    # Test multiple CPUs
    spca = SparsePCA(n_components=3, n_jobs=2, method='lars', alpha=alpha,
                     random_state=0).fit(Y)
    U2 = spca.transform(Y)
    assert_true(not np.all(spca_lars.components_ == 0))
    assert_array_almost_equal(U1, U2)

def test_transform_nan():
    # Test that SparsePCA won't return NaN when there is 0 feature in all
    # samples.
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    Y[:, 0] = 0
    estimator = SparsePCA(n_components=8)
    assert_false(np.any(np.isnan(estimator.fit_transform(Y))))

def test_fit_transform_tall():
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 65, (8, 8), random_state=rng)  # tall array
    spca_lars = SparsePCA(n_components=3, method='lars',
                          random_state=rng)
    U1 = spca_lars.fit_transform(Y)
    spca_lasso = SparsePCA(n_components=3, method='cd', random_state=rng)
    U2 = spca_lasso.fit(Y).transform(Y)
    assert_array_almost_equal(U1, U2)

def test_initialization():
    rng = np.random.RandomState(0)
    U_init = rng.randn(5, 3)
    V_init = rng.randn(3, 4)
    model = SparsePCA(n_components=3, U_init=U_init, V_init=V_init, max_iter=0,
                      random_state=rng)
    model.fit(rng.randn(5, 4))
    assert_array_equal(model.components_, V_init)

def test_mini_batch_correct_shapes():
    rng = np.random.RandomState(0)
    X = rng.randn(12, 10)
    pca = MiniBatchSparsePCA(n_components=8, random_state=rng)
    U = pca.fit_transform(X)
    assert_equal(pca.components_.shape, (8, 10))
    assert_equal(U.shape, (12, 8))
    # test overcomplete decomposition
    pca = MiniBatchSparsePCA(n_components=13, random_state=rng)
    U = pca.fit_transform(X)
    assert_equal(pca.components_.shape, (13, 10))
    assert_equal(U.shape, (12, 13))

def test_mini_batch_fit_transform():
    raise SkipTest("skipping mini_batch_fit_transform.")
    alpha = 1
    rng = np.random.RandomState(0)
    Y, _, _ = generate_toy_data(3, 10, (8, 8), random_state=rng)  # wide array
    spca_lars = MiniBatchSparsePCA(n_components=3, random_state=0,
                                   alpha=alpha).fit(Y)
    U1 = spca_lars.transform(Y)
    # Test multiple CPUs
    if sys.platform == 'win32':  # fake parallelism for win32
        import sklearn.externals.joblib.parallel as joblib_par
        _mp = joblib_par.multiprocessing
        joblib_par.multiprocessing = None
        try:
            U2 = MiniBatchSparsePCA(n_components=3, n_jobs=2, alpha=alpha,
                                    random_state=0).fit(Y).transform(Y)
        finally:
            joblib_par.multiprocessing = _mp
    else:  # we can efficiently use parallelism
        U2 = MiniBatchSparsePCA(n_components=3, n_jobs=2, alpha=alpha,
                                random_state=0).fit(Y).transform(Y)
    assert_true(not np.all(spca_lars.components_ == 0))
    assert_array_almost_equal(U1, U2)
    # Test that CD gives similar results
    spca_lasso = MiniBatchSparsePCA(n_components=3, method='cd', alpha=alpha,
                                    random_state=0).fit(Y)
    assert_array_almost_equal(spca_lasso.components_, spca_lars.components_)


# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# 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.

"""Tests for smu_utils_lib."""

import copy
import os
import tempfile

from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
import pandas as pd
from rdkit import Chem

from google.protobuf import text_format
from smu import dataset_pb2
from smu.parser import smu_parser_lib
from smu.parser import smu_utils_lib

MAIN_DAT_FILE = 'x07_sample.dat'
STAGE1_DAT_FILE = 'x07_stage1.dat'
TESTDATA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                             'testdata')

def str_to_bond_topology(s):
  bt = dataset_pb2.BondTopology()
  text_format.Parse(s, bt)
  return bt

def get_stage1_conformer():
  parser = smu_parser_lib.SmuParser(
      os.path.join(TESTDATA_PATH, STAGE1_DAT_FILE))
  conformer, _ = next(parser.process_stage1())
  return conformer

def get_stage2_conformer():
  parser = smu_parser_lib.SmuParser(os.path.join(TESTDATA_PATH, MAIN_DAT_FILE))
  conformer, _ = next(parser.process_stage2())
  return conformer

class SpecialIDTest(absltest.TestCase):

  def test_from_dat_id(self):
    self.assertIsNone(
        smu_utils_lib.special_case_bt_id_from_dat_id(123456, 'CC'))
    self.assertEqual(smu_utils_lib.special_case_bt_id_from_dat_id(999998, 'O'),
                     899650)
    self.assertEqual(smu_utils_lib.special_case_bt_id_from_dat_id(0, 'O'),
                     899650)
    with self.assertRaises(ValueError):
      smu_utils_lib.special_case_bt_id_from_dat_id(0, 'NotASpecialCaseSmiles')

  def test_from_bt_id(self):
    self.assertIsNone(smu_utils_lib.special_case_dat_id_from_bt_id(123456))
    self.assertEqual(
        smu_utils_lib.special_case_dat_id_from_bt_id(899651), 999997)

class GetCompositionTest(absltest.TestCase):

  def test_simple(self):
    bt = dataset_pb2.BondTopology()
    bt.atoms.extend([dataset_pb2.BondTopology.ATOM_C,
                     dataset_pb2.BondTopology.ATOM_C,
                     dataset_pb2.BondTopology.ATOM_N,
                     dataset_pb2.BondTopology.ATOM_H,
                     dataset_pb2.BondTopology.ATOM_H,
                     dataset_pb2.BondTopology.ATOM_H])
    self.assertEqual('x03_c2nh3', smu_utils_lib.get_composition(bt))

class GetCanonicalStoichiometryWithHydrogensTest(absltest.TestCase):

  def test_cyclobutane(self):
    bt = smu_utils_lib.create_bond_topology('CCCC', '110011', '2222')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(bt), '(ch2)4')

  def test_ethylene(self):
    bt = smu_utils_lib.create_bond_topology('CC', '2', '22')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(bt), '(ch2)2')

  def test_acrylic_acid(self):
    bt = smu_utils_lib.create_bond_topology('CCCOO', '2000100210', '21001')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(bt),
        '(c)(ch)(ch2)(o)(oh)')

  def test_fluorine(self):
    bt = smu_utils_lib.create_bond_topology('OFF', '110', '000')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(bt), '(o)(f)2')

  def test_fully_saturated(self):
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(
            smu_utils_lib.create_bond_topology('C', '', '4')), '(ch4)')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(
            smu_utils_lib.create_bond_topology('N', '', '3')), '(nh3)')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(
            smu_utils_lib.create_bond_topology('O', '', '2')), '(oh2)')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(
            smu_utils_lib.create_bond_topology('F', '', '1')), '(fh)')

  def test_nplus_oneg(self):
    bt = smu_utils_lib.create_bond_topology('NO', '1', '30')
    self.assertEqual(
        smu_utils_lib.get_canonical_stoichiometry_with_hydrogens(bt),
        '(nh3)(o)')

class ParseBondTopologyTest(absltest.TestCase):

  def test_4_heavy(self):
    num_atoms, atoms_str, matrix, hydrogens = smu_utils_lib.parse_bond_topology_line(
        ' 4  N+O O O-  010110  3000')
    self.assertEqual(num_atoms, 4)
    self.assertEqual(atoms_str, 'N+O O O-')
    self.assertEqual(matrix, '010110')
    self.assertEqual(hydrogens, '3000')

  def test_7_heavy(self):
    num_atoms, atoms_str, matrix, hydrogens = smu_utils_lib.parse_bond_topology_line(
        ' 7  N+O O O O-F F   001011101001000000000  1000000')
    self.assertEqual(num_atoms, 7)
    self.assertEqual(atoms_str, 'N+O O O O-F F ')  # Note the trailing space
    self.assertEqual(matrix, '001011101001000000000')
    self.assertEqual(hydrogens, '1000000')

class CreateBondTopologyTest(absltest.TestCase):

  def test_no_charged(self):
    got = smu_utils_lib.create_bond_topology('CNFF', '111000', '1200')
    expected_str = '''
atoms: ATOM_C
atoms: ATOM_N
atoms: ATOM_F
atoms: ATOM_F
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
bonds {
  atom_b: 1
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 2
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 3
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 4
  bond_type: BOND_SINGLE
}
bonds {
  atom_a: 1
  atom_b: 5
  bond_type: BOND_SINGLE
}
bonds {
  atom_a: 1
  atom_b: 6
  bond_type: BOND_SINGLE
}
'''
    expected = str_to_bond_topology(expected_str)
    self.assertEqual(str(expected), str(got))

  def test_charged(self):
    # This is actually C N N+O-
    got = smu_utils_lib.create_bond_topology('CNNO', '200101', '2020')
    expected_str = '''
atoms: ATOM_C
atoms: ATOM_N
atoms: ATOM_NPOS
atoms: ATOM_ONEG
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
bonds {
  atom_b: 1
  bond_type: BOND_DOUBLE
}
bonds {
  atom_a: 1
  atom_b: 2
  bond_type: BOND_SINGLE
}
bonds {
  atom_a: 2
  atom_b: 3
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 4
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 5
  bond_type: BOND_SINGLE
}
bonds {
  atom_a: 2
  atom_b: 6
  bond_type: BOND_SINGLE
}
bonds {
  atom_a: 2
  atom_b: 7
  bond_type: BOND_SINGLE
}
'''
    expected = str_to_bond_topology(expected_str)
    self.assertEqual(str(expected), str(got))

  def test_one_heavy(self):
    got = smu_utils_lib.create_bond_topology('C', '', '4')
    expected_str = '''
atoms: ATOM_C
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
bonds {
  atom_b: 1
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 2
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 3
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 4
  bond_type: BOND_SINGLE
}
'''
    expected = str_to_bond_topology(expected_str)
    self.assertEqual(str(expected), str(got))

class FromCSVTest(absltest.TestCase):

  def test_basic(self):
    infile = tempfile.NamedTemporaryFile(mode='w', delete=False)
    infile.write(
        'id,num_atoms,atoms_str,connectivity_matrix,hydrogens,smiles\n')
    infile.write('68,3,C N+O-,310,010,[NH+]#C[O-]\n')
    infile.write('134,4,N+O-F F ,111000,1000,[O-][NH+](F)F\n')
    infile.close()

    out = smu_utils_lib.generate_bond_topologies_from_csv(infile.name)

    bt = next(out)
    self.assertEqual(68, bt.bond_topology_id)
    self.assertLen(bt.atoms, 4)
    self.assertEqual(bt.smiles, '[NH+]#C[O-]')

    bt = next(out)
    self.assertEqual(134, bt.bond_topology_id)
    self.assertLen(bt.atoms, 5)
    self.assertEqual(bt.smiles, '[O-][NH+](F)F')

class ParseDuplicatesFileTest(absltest.TestCase):

  def test_basic(self):
    df = smu_utils_lib.parse_duplicates_file(
        os.path.join(TESTDATA_PATH, 'small.equivalent_isomers.dat'))
    pd.testing.assert_frame_equal(
        pd.DataFrame(
            columns=['name1', 'stoich1', 'btid1', 'shortconfid1', 'confid1',
                     'name2', 'stoich2', 'btid2', 'shortconfid2', 'confid2'],
            data=[
                ['x07_c2n2o2fh3.224227.004',
                 'c2n2o2fh3', 224227, 4, 224227004,
                 'x07_c2n2o2fh3.224176.005',
                 'c2n2o2fh3', 224176, 5, 224176005],
                ['x07_c2n2o2fh3.260543.005',
                 'c2n2o2fh3', 260543, 5, 260543005,
                 'x07_c2n2o2fh3.224050.001',
                 'c2n2o2fh3', 224050, 1, 224050001],
            ]),
        df,
        check_like=True)

class BondTopologyToMoleculeTest(absltest.TestCase):

  def test_o2(self):
    bond_topology = str_to_bond_topology('''
atoms: ATOM_O
atoms: ATOM_O
bonds {
  atom_b: 1
  bond_type: BOND_DOUBLE
}
''')
    got = smu_utils_lib.bond_topology_to_molecule(bond_topology)
    self.assertEqual('O=O', Chem.MolToSmiles(got))

  def test_methane(self):
    bond_topology = str_to_bond_topology('''
atoms: ATOM_C
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
atoms: ATOM_H
bonds {
  atom_b: 1
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 2
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 3
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 4
  bond_type: BOND_SINGLE
}
''')
    got = smu_utils_lib.bond_topology_to_molecule(bond_topology)
    self.assertEqual('[H]C([H])([H])[H]', Chem.MolToSmiles(got))

  # This molecule is an N+ central atom, bonded to C (triply), O-, and F
  def test_charged_molecule(self):
    bond_topology = str_to_bond_topology('''
atoms: ATOM_C
atoms: ATOM_NPOS
atoms: ATOM_ONEG
atoms: ATOM_F
bonds {
  atom_b: 1
  bond_type: BOND_TRIPLE
}
bonds {
  atom_a: 1
  atom_b: 2
  bond_type: BOND_SINGLE
}
bonds {
  atom_a: 1
  atom_b: 3
  bond_type: BOND_SINGLE
}
''')
    got = smu_utils_lib.bond_topology_to_molecule(bond_topology)
    self.assertEqual('C#[N+]([O-])F', Chem.MolToSmiles(got))

class ConformerToMoleculeTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.conformer = get_stage2_conformer()

    # We'll make a new initial_geometry which is just the current one with all
    # coordinates multiplied by 1000
    self.conformer.initial_geometries.append(
        self.conformer.initial_geometries[0])
    new_geom = self.conformer.initial_geometries[1]
    for atom_pos in new_geom.atom_positions:
      atom_pos.x = atom_pos.x * 1000
      atom_pos.y = atom_pos.y * 1000
      atom_pos.z = atom_pos.z * 1000

    # For the extra bond_topology, we'll just copy the existing one and change
    # the id. Through the dumb luck of the molecule we picked there's not a
    # simple way to make this a new bond topology and still have it look valid
    # to RDKit
    self.conformer.bond_topologies.append(self.conformer.bond_topologies[0])
    self.conformer.bond_topologies[1].bond_topology_id = 99999

  def test_all_outputs(self):
    mols = list(smu_utils_lib.conformer_to_molecules(self.conformer))
    self.assertLen(mols, 6)  # 2 bond topologies * (1 opt geom + 2 init_geom)
    self.assertEqual([m.GetProp('_Name') for m in mols], [
        'SMU 618451001 bt=618451(0/2) geom=init(0/2)',
        'SMU 618451001 bt=618451(0/2) geom=init(1/2)',
        'SMU 618451001 bt=618451(0/2) geom=opt',
        'SMU 618451001 bt=99999(1/2) geom=init(0/2)',
        'SMU 618451001 bt=99999(1/2) geom=init(1/2)',
        'SMU 618451001 bt=99999(1/2) geom=opt'
    ])
    self.assertEqual(
        '[H]C(F)=C(OC([H])([H])[H])OC([H])([H])[H]',
        Chem.MolToSmiles(mols[0], kekuleSmiles=True, isomericSmiles=False))
    self.assertEqual(
        '[H]C(F)=C(OC([H])([H])[H])OC([H])([H])[H]',
        Chem.MolToSmiles(mols[4], kekuleSmiles=True, isomericSmiles=False))

  def test_initial_only(self):
    mols = list(
        smu_utils_lib.conformer_to_molecules(
            self.conformer,
            include_initial_geometries=True,
            include_optimized_geometry=False,
            include_all_bond_topologies=False))
    self.assertLen(mols, 2)
    self.assertEqual([m.GetProp('_Name') for m in mols], [
        'SMU 618451001 bt=618451(0/2) geom=init(0/2)',
        'SMU 618451001 bt=618451(0/2) geom=init(1/2)',
    ])
    # This is just one random atom I picked from the .dat file and converted to
    # angstroms instead of bohr.
    self.assertEqual('C', mols[0].GetAtomWithIdx(1).GetSymbol())
    np.testing.assert_allclose([0.6643, -3.470301, 3.4766],
                               list(mols[0].GetConformer().GetAtomPosition(1)),
                               atol=1e-6)

    self.assertEqual('C', mols[1].GetAtomWithIdx(1).GetSymbol())
    np.testing.assert_allclose([664.299998, -3470.300473, 3476.600215],
                               list(mols[1].GetConformer().GetAtomPosition(1)),
                               atol=1e-6)

  def test_optimized_only(self):
    mols = list(
        smu_utils_lib.conformer_to_molecules(
            self.conformer,
            include_initial_geometries=False,
            include_optimized_geometry=True,
            include_all_bond_topologies=False))
    self.assertLen(mols, 1)
    self.assertEqual(
        mols[0].GetProp('_Name'),
        'SMU 618451001 bt=618451(0/2) geom=opt',
    )
    self.assertEqual(
        '[H]C(F)=C(OC([H])([H])[H])OC([H])([H])[H]',
        Chem.MolToSmiles(mols[0], kekuleSmiles=True, isomericSmiles=False))
    # This is just two random atoms I picked from the .dat file and converted to
    # angstroms instead of bohr.
    self.assertEqual('C', mols[0].GetAtomWithIdx(1).GetSymbol())
    np.testing.assert_allclose([0.540254, -3.465543, 3.456982],
                               list(mols[0].GetConformer().GetAtomPosition(1)),
                               atol=1e-6)
    self.assertEqual('H', mols[0].GetAtomWithIdx(13).GetSymbol())
    np.testing.assert_allclose([2.135153, -1.817366, 0.226376],
                               list(mols[0].GetConformer().GetAtomPosition(13)),
                               atol=1e-6)

class SmilesCompareTest(absltest.TestCase):

  def test_string_format(self):
    # for some simplicity later on, we use shorter names
    self.assertEqual('MISSING', str(smu_utils_lib.SmilesCompareResult.MISSING))
    self.assertEqual('MISMATCH',
                     str(smu_utils_lib.SmilesCompareResult.MISMATCH))
    self.assertEqual('MATCH', str(smu_utils_lib.SmilesCompareResult.MATCH))

  def test_missing(self):
    bond_topology = str_to_bond_topology('''
atoms: ATOM_O
atoms: ATOM_O
bonds {
  atom_b: 1
  bond_type: BOND_DOUBLE
}
''')
    result, with_h, without_h = smu_utils_lib.bond_topology_smiles_comparison(
        bond_topology)
    self.assertEqual(smu_utils_lib.SmilesCompareResult.MISSING, result)
    self.assertEqual('O=O', with_h)
    self.assertEqual('O=O', without_h)

    # Also directly test compute_smiles_for_bond_topology
    self.assertEqual(
        'O=O',
        smu_utils_lib.compute_smiles_for_bond_topology(
            bond_topology, include_hs=True))

  def test_mismatch(self):
    bond_topology = str_to_bond_topology('''
atoms: ATOM_O
atoms: ATOM_O
bonds {
  atom_b: 1
  bond_type: BOND_DOUBLE
}
smiles: "BlahBlahBlah"
''')
    result, with_h, without_h = smu_utils_lib.bond_topology_smiles_comparison(
        bond_topology)
    self.assertEqual(smu_utils_lib.SmilesCompareResult.MISMATCH, result)
    self.assertEqual('O=O', with_h)
    self.assertEqual('O=O', without_h)

  def test_matched_and_h_stripping(self):
    bond_topology = str_to_bond_topology('''
atoms: ATOM_O
atoms: ATOM_H
atoms: ATOM_H
bonds {
  atom_b: 1
  bond_type: BOND_SINGLE
}
bonds {
  atom_b: 2
  bond_type: BOND_SINGLE
}
smiles: "O"
''')
    result, with_h, without_h = smu_utils_lib.bond_topology_smiles_comparison(
        bond_topology)
    self.assertEqual(smu_utils_lib.SmilesCompareResult.MATCH, result)
    self.assertEqual('[H]O[H]', with_h)
    self.assertEqual('O', without_h)

    # Also directly test compute_smiles_for_bond_topology
    self.assertEqual(
        '[H]O[H]',
        smu_utils_lib.compute_smiles_for_bond_topology(
            bond_topology, include_hs=True))
    self.assertEqual(
        'O',
        smu_utils_lib.compute_smiles_for_bond_topology(
            bond_topology, include_hs=False))

  def test_compute_smiles_from_molecule_no_hs(self):
    mol = Chem.MolFromSmiles('FOC', sanitize=False)
    self.assertEqual(
        smu_utils_lib.compute_smiles_for_molecule(mol, include_hs=False), 'COF')
    # This is expected. Even with include_hs=True, if there were no Hs in the
    # molecule, they will not be in the smiles.
    self.assertEqual(
        smu_utils_lib.compute_smiles_for_molecule(mol, include_hs=True), 'COF')

  def test_compute_smiles_from_molecule_with_hs(self):
    mol = Chem.MolFromSmiles('FOC', sanitize=False)
    Chem.SanitizeMol(mol, Chem.rdmolops.SanitizeFlags.SANITIZE_ADJUSTHS)
    mol = Chem.AddHs(mol)
    self.assertEqual(
        smu_utils_lib.compute_smiles_for_molecule(mol, include_hs=False), 'COF')
    self.assertEqual(
        smu_utils_lib.compute_smiles_for_molecule(mol, include_hs=True),
        '[H]C([H])([H])OF')

  def test_compute_smiles_from_molecule_special_case(self):
    mol = Chem.MolFromSmiles('C12=C3C4=C1C4=C23', sanitize=False)
    # Double check that this really is the special case -- we get back the
    # SMILES we put in even though it's not the one we want.
    self.assertEqual('C12=C3C4=C1C4=C23',
                     Chem.MolToSmiles(mol, kekuleSmiles=True))
    self.assertEqual(
        smu_utils_lib.compute_smiles_for_molecule(mol, include_hs=False),
        'C12=C3C1=C1C2=C31')

  def test_compute_smiles_from_molecule_labeled_with_h(self):
    mol = Chem.MolFromSmiles(
        '[O-][N+]([H])([H])N([H])OC([H])([H])F', sanitize=False)
    self.assertIsNotNone(mol)
    self.assertEqual(
        '[O-][N+:1]([H:2])([H:3])[N:4]([H:5])[O:6][C:7]([H:8])([H:9])[F:10]',
        smu_utils_lib.compute_smiles_for_molecule(
            mol, include_hs=True, labeled_atoms=True))

  def test_compute_smiles_from_molecule_labeled_no_h(self):
    mol = Chem.MolFromSmiles(
        '[O-][N+]([H])([H])N([H])OC([H])([H])F', sanitize=False)
    self.assertIsNotNone(mol)
    self.assertEqual(
        '[O-][NH2+:1][NH:2][O:3][CH2:4][F:5]',
        smu_utils_lib.compute_smiles_for_molecule(
            mol, include_hs=False, labeled_atoms=True))

class MergeConformersTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    # We are relying on the fact that the first conformer in both x07_sample.dat
    # and x07_stage1.dat are the same.
    self.stage1_conformer = get_stage1_conformer()
    self.stage2_conformer = get_stage2_conformer()

    self.duplicate_conformer = dataset_pb2.Conformer()
    self.duplicate_conformer.conformer_id = self.stage1_conformer.conformer_id
    # A real duplicate conformer wouldn't have both of these fields filled in,
    # but it's fine for the test to make sure everything is copied.
    self.duplicate_conformer.duplicated_by = 123
    self.duplicate_conformer.duplicate_of.extend([111, 222])

  def test_two_stage2(self):
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(self.stage2_conformer,
                                    self.stage2_conformer)

  def test_two_stage1(self):
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(self.stage1_conformer,
                                    self.stage1_conformer)

  def test_two_duplicates(self):
    duplicate_conformer2 = copy.deepcopy(self.duplicate_conformer)
    duplicate_conformer2.duplicate_of[:] = [333, 444]

    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.duplicate_conformer, duplicate_conformer2)
    self.assertIsNone(got_conflict)
    self.assertEqual(123, got_conf.duplicated_by)
    self.assertCountEqual([111, 222, 333, 444], got_conf.duplicate_of)

  def test_stage2_stage1(self):
    # Add a duplicate to stage1 to make sure it is copied
    self.stage1_conformer.duplicate_of.append(999)
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.stage1_conformer)
    self.assertIsNone(got_conflict)
    self.assertEqual(got_conf.duplicate_of, [999])
    # Just check a random field that is in stage2 but not stage1
    self.assertNotEmpty(got_conf.properties.normal_modes)

  def test_stage2_stage1_conflict_energy(self):
    self.stage2_conformer.properties.initial_geometry_energy.value = -1.23
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.stage1_conformer)
    self.assertEqual(got_conflict, [
        618451001,
        1, 1, 1, 1, -406.51179, 0.052254, -406.522079, 2.5e-05, True, True,
        1, 1, 1, 1, -1.23, 0.052254, -406.522079, 2.5e-05, True, True
    ])
    # Just check a random field that is in stage2 but not stage1
    self.assertNotEmpty(got_conf.properties.normal_modes)
    # This stage2 values should be returned
    self.assertEqual(got_conf.properties.initial_geometry_energy.value, -1.23)

  def test_stage2_stage1_conflict_error_codes(self):
    self.stage2_conformer.properties.errors.error_nstat1 = 999
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.stage1_conformer)
    self.assertEqual(got_conflict, [
        618451001,
        1, 1, 1, 1, -406.51179, 0.052254, -406.522079, 2.5e-05, True, True,
        999, 1, 1, 1, -406.51179, 0.052254, -406.522079, 2.5e-05, True, True
    ])
    # Just check a random field that is in stage2 but not stage1
    self.assertNotEmpty(got_conf.properties.normal_modes)

  def test_stage2_stage1_conflict_missing_geometry(self):
    self.stage2_conformer.ClearField('optimized_geometry')
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.stage1_conformer)
    self.assertEqual(got_conflict, [
        618451001,
        1, 1, 1, 1, -406.51179, 0.052254, -406.522079, 2.5e-05, True, True,
        1, 1, 1, 1, -406.51179, 0.052254, -406.522079, 2.5e-05, True, False
    ])
    # Just check a random field that is in stage2 but not stage1
    self.assertNotEmpty(got_conf.properties.normal_modes)

  def test_stage2_stage1_no_conflict_minus1(self):
    # If stage2 contains a -1, we keep that (stricter error checking later on)
    self.stage2_conformer.properties.initial_geometry_energy.value = -1.0
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.stage1_conformer)
    self.assertIsNone(got_conflict)
    self.assertEqual(got_conf.properties.initial_geometry_energy.value, -1.0)

  def test_stage2_stage1_no_conflict_approx_equal(self):
    self.stage2_conformer.properties.initial_geometry_energy.value += 1e-7
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.stage1_conformer)
    self.assertIsNone(got_conflict)
    # Just check a random field from stage2
    self.assertNotEmpty(got_conf.properties.normal_modes)

  def test_stage2_duplicate(self):
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage2_conformer, self.duplicate_conformer)
    self.assertIsNone(got_conflict)
    self.assertEqual(got_conf.duplicate_of, [111, 222])
    self.assertEqual(got_conf.duplicated_by, 123)
    # Just check a random field from stage2
    self.assertNotEmpty(got_conf.properties.normal_modes)

  def test_stage1_duplicate(self):
    got_conf, got_conflict = smu_utils_lib.merge_conformer(
        self.stage1_conformer, self.duplicate_conformer)
    self.assertIsNone(got_conflict)
    self.assertEqual(got_conf.duplicate_of, [111, 222])
    self.assertEqual(got_conf.duplicated_by, 123)
    # Just check a random field from stage1
    self.assertTrue(got_conf.properties.HasField('initial_geometry_energy'))

  def test_multiple_initial_geometries(self):
    bad_conformer = copy.deepcopy(self.stage1_conformer)
    bad_conformer.initial_geometries.append(bad_conformer.initial_geometries[0])
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(bad_conformer, self.stage2_conformer)
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(self.stage2_conformer, bad_conformer)

  def test_multiple_bond_topologies(self):
    bad_conformer = copy.deepcopy(self.stage1_conformer)
    bad_conformer.bond_topologies.append(bad_conformer.bond_topologies[0])
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(bad_conformer, self.stage2_conformer)
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(self.stage2_conformer, bad_conformer)

  def test_different_bond_topologies(self):
    self.stage1_conformer.bond_topologies[0].atoms[0] = (
        dataset_pb2.BondTopology.ATOM_H)
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(self.stage1_conformer,
                                    self.stage2_conformer)
    with self.assertRaises(ValueError):
      smu_utils_lib.merge_conformer(self.stage2_conformer,
                                    self.stage1_conformer)

class ConformerErrorTest(absltest.TestCase):

  def test_stage1_no_error(self):
    conformer = get_stage1_conformer()
    self.assertFalse(smu_utils_lib.conformer_has_calculation_errors(conformer))

  def test_stage1_error(self):
    conformer = get_stage2_conformer()
    conformer.properties.errors.error_frequencies = 123
    self.assertTrue(smu_utils_lib.conformer_has_calculation_errors(conformer))

  def test_stage2_no_error(self):
    conformer = get_stage2_conformer()
    self.assertFalse(smu_utils_lib.conformer_has_calculation_errors(conformer))

  def test_stage2_error_in_1_expected_field(self):
    conformer = get_stage2_conformer()
    conformer.properties.errors.error_rotational_modes = 123
    self.assertTrue(smu_utils_lib.conformer_has_calculation_errors(conformer))

  def test_stage2_error_in_0_expected_field(self):
    conformer = get_stage2_conformer()
    # This field is 0 to indicate no error. Why the discrepancy? Who knows!
    conformer.properties.errors.error_nsvg09 = 1
    self.assertTrue(smu_utils_lib.conformer_has_calculation_errors(conformer))

  def test_stage2_nstat1_is_3(self):
    # This is the other bizaare case. nstat1 of 3 is still considered success.
    conformer = get_stage2_conformer()
    conformer.properties.errors.error_nstat1 = 3
    self.assertFalse(smu_utils_lib.conformer_has_calculation_errors(conformer))

class FilterConformerByAvailabilityTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.conformer = dataset_pb2.Conformer()
    properties = self.conformer.properties
    # A STANDARD field
    properties.single_point_energy_pbe0d3_6_311gd.value = 1.23
    # A COMPLETE field
    properties.homo_pbe0_aug_pc_1.value = 1.23
    # An INTERNAL_ONLY field
    properties.nuclear_repulsion_energy.value = 1.23

  def test_standard(self):
    smu_utils_lib.filter_conformer_by_availability(self.conformer,
                                                   [dataset_pb2.STANDARD])
    self.assertTrue(
        self.conformer.properties.HasField(
            'single_point_energy_pbe0d3_6_311gd'))
    self.assertFalse(self.conformer.properties.HasField('homo_pbe0_aug_pc_1'))
    self.assertFalse(
        self.conformer.properties.HasField('nuclear_repulsion_energy'))

  def test_complete_and_internal_only(self):
    smu_utils_lib.filter_conformer_by_availability(
        self.conformer, [dataset_pb2.COMPLETE, dataset_pb2.INTERNAL_ONLY])
    self.assertFalse(
        self.conformer.properties.HasField(
            'single_point_energy_pbe0d3_6_311gd'))
    self.assertTrue(self.conformer.properties.HasField('homo_pbe0_aug_pc_1'))
    self.assertTrue(
        self.conformer.properties.HasField('nuclear_repulsion_energy'))

class ConformerToStandardTest(absltest.TestCase):

  def setUp(self):
    super().setUp()

    self.conformer = get_stage2_conformer()

  def test_field_filtering(self):
    # Check that the field which should be filtered starts out set
    self.assertTrue(self.conformer.properties.HasField(
        'single_point_energy_hf_6_31gd'))

    got = smu_utils_lib.conformer_to_standard(self.conformer)
    # Check for a field that was originally in self.conformer and should be
    # filtered and a field which should still be present.
    self.assertTrue(got.properties.HasField(
        'single_point_energy_pbe0d3_6_311gd'))
    self.assertFalse(
        got.properties.HasField('single_point_energy_hf_6_31gd'))

  def test_remove_error_conformer(self):
    self.conformer.properties.errors.error_frequencies = 123

    self.assertIsNone(smu_utils_lib.conformer_to_standard(self.conformer))

  def test_remove_duplicate(self):
    self.conformer.duplicated_by = 123

    self.assertIsNone(smu_utils_lib.conformer_to_standard(self.conformer))

class DetermineFateTest(parameterized.TestCase):

  def test_duplicate_same_topology(self):
    conformer = get_stage1_conformer()
    # bond topology is conformer_id // 1000
    conformer.duplicated_by = conformer.conformer_id + 1
    self.assertEqual(dataset_pb2.Conformer.FATE_DUPLICATE_SAME_TOPOLOGY,
                     smu_utils_lib.determine_fate(conformer))

  def test_duplicate_different_topology(self):
    conformer = get_stage1_conformer()
    # bond topology is conformer_id // 1000
    conformer.duplicated_by = conformer.conformer_id + 1000
    self.assertEqual(dataset_pb2.Conformer.FATE_DUPLICATE_DIFFERENT_TOPOLOGY,
                     smu_utils_lib.determine_fate(conformer))

  @parameterized.parameters(
      (2, dataset_pb2.Conformer.FATE_GEOMETRY_OPTIMIZATION_PROBLEM),
      (5, dataset_pb2.Conformer.FATE_DISASSOCIATED),
      (4, dataset_pb2.Conformer.FATE_FORCE_CONSTANT_FAILURE),
      (6, dataset_pb2.Conformer.FATE_DISCARDED_OTHER))
  def test_geometry_failures(self, nstat1, expected_fate):
    conformer = get_stage1_conformer()
    conformer.properties.errors.error_nstat1 = nstat1
    self.assertEqual(expected_fate, smu_utils_lib.determine_fate(conformer))

  def test_no_result(self):
    conformer = get_stage1_conformer()
    self.assertEqual(dataset_pb2.Conformer.FATE_NO_CALCULATION_RESULTS,
                     smu_utils_lib.determine_fate(conformer))

  def test_calculation_errors(self):
    conformer = get_stage2_conformer()
    # This is a random choice of an error to set. I just need some error.
    conformer.properties.errors.error_atomic_analysis = 999
    self.assertEqual(dataset_pb2.Conformer.FATE_CALCULATION_WITH_ERROR,
                     smu_utils_lib.determine_fate(conformer))

  def test_success(self):
    conformer = get_stage2_conformer()
    self.assertEqual(dataset_pb2.Conformer.FATE_SUCCESS,
                     smu_utils_lib.determine_fate(conformer))

class ToBondTopologySummaryTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.conformer = get_stage2_conformer()

  def test_dup_same(self):
    self.conformer.fate = dataset_pb2.Conformer.FATE_DUPLICATE_SAME_TOPOLOGY
    got = list(
        smu_utils_lib.conformer_to_bond_topology_summaries(self.conformer))
    self.assertLen(got, 1)
    self.assertEqual(got[0].bond_topology.bond_topology_id,
                     self.conformer.bond_topologies[0].bond_topology_id)
    self.assertEqual(got[0].count_attempted_conformers, 1)
    self.assertEqual(got[0].count_duplicates_same_topology, 1)

  def test_dup_diff(self):
    self.conformer.fate = (
        dataset_pb2.Conformer.FATE_DUPLICATE_DIFFERENT_TOPOLOGY)
    got = list(
        smu_utils_lib.conformer_to_bond_topology_summaries(self.conformer))
    self.assertLen(got, 1)
    self.assertEqual(got[0].count_attempted_conformers, 1)
    self.assertEqual(got[0].count_duplicates_different_topology, 1)

  def test_geometry_failed(self):
    self.conformer.fate = (dataset_pb2.Conformer.FATE_DISCARDED_OTHER)
    got = list(
        smu_utils_lib.conformer_to_bond_topology_summaries(self.conformer))
    self.assertLen(got, 1)
    self.assertEqual(got[0].count_attempted_conformers, 1)
    self.assertEqual(got[0].count_failed_geometry_optimization, 1)

  def test_missing_calculation(self):
    self.conformer.fate = dataset_pb2.Conformer.FATE_NO_CALCULATION_RESULTS
    got = list(
        smu_utils_lib.conformer_to_bond_topology_summaries(self.conformer))
    self.assertLen(got, 1)
    self.assertEqual(got[0].count_attempted_conformers, 1)
    self.assertEqual(got[0].count_kept_geometry, 1)
    self.assertEqual(got[0].count_missing_calculation, 1)

  def test_calculation_with_error(self):
    self.conformer.fate = dataset_pb2.Conformer.FATE_CALCULATION_WITH_ERROR
    self.conformer.bond_topologies.append(self.conformer.bond_topologies[0])
    self.conformer.bond_topologies[-1].bond_topology_id = 123
    got = list(
        smu_utils_lib.conformer_to_bond_topology_summaries(self.conformer))
    self.assertLen(got, 2)
    # We don't actually care about the order, but this is what comes out right
    # now.
    self.assertEqual(got[0].bond_topology.bond_topology_id, 123)
    self.assertEqual(got[0].count_attempted_conformers, 0)
    self.assertEqual(got[0].count_kept_geometry, 0)
    self.assertEqual(got[0].count_calculation_with_error, 0)
    self.assertEqual(got[0].count_detected_match_with_error, 1)

    self.assertEqual(got[1].bond_topology.bond_topology_id,
                     self.conformer.bond_topologies[0].bond_topology_id)
    self.assertEqual(got[1].count_attempted_conformers, 1)
    self.assertEqual(got[1].count_kept_geometry, 1)
    self.assertEqual(got[1].count_calculation_with_error, 1)
    self.assertEqual(got[1].count_detected_match_with_error, 0)

  def test_calculation_success(self):
    self.conformer.fate = dataset_pb2.Conformer.FATE_SUCCESS
    self.conformer.bond_topologies.append(self.conformer.bond_topologies[0])
    self.conformer.bond_topologies[-1].bond_topology_id = 123
    got = list(
        smu_utils_lib.conformer_to_bond_topology_summaries(self.conformer))
    self.assertLen(got, 2)
    # We don't actually care about the order, but this is what comes out right
    # now.
    self.assertEqual(got[0].bond_topology.bond_topology_id, 123)
    self.assertEqual(got[0].count_attempted_conformers, 0)
    self.assertEqual(got[0].count_kept_geometry, 0)
    self.assertEqual(got[0].count_calculation_success, 0)
    self.assertEqual(got[0].count_detected_match_success, 1)

    self.assertEqual(got[1].bond_topology.bond_topology_id,
                     self.conformer.bond_topologies[0].bond_topology_id)
    self.assertEqual(got[1].count_attempted_conformers, 1)
    self.assertEqual(got[1].count_kept_geometry, 1)
    self.assertEqual(got[1].count_calculation_success, 1)
    self.assertEqual(got[1].count_detected_match_success, 0)

class LabeledSmilesTester(absltest.TestCase):

  def test_atom_labels(self):
    mol = Chem.MolFromSmiles('FCON[NH2+][O-]', sanitize=False)
    self.assertIsNotNone(mol)
    smiles_before = Chem.MolToSmiles(mol)
    self.assertEqual(
        smu_utils_lib.labeled_smiles(mol), 'F[CH2:1][O:2][NH:3][NH2+:4][O-:5]')
    # Testing both the atom numbers and the smiles is redundant,
    # but guards against possible future changes.
    for atom in mol.GetAtoms():
      self.assertEqual(atom.GetAtomMapNum(), 0)
    self.assertEqual(Chem.MolToSmiles(mol), smiles_before)

if __name__ == '__main__':
  absltest.main()


import os
import re
from collections import defaultdict
from operator import itemgetter
import logging

import pandas
from scipy.interpolate import InterpolatedUnivariateSpline as spline
from george import kernels
import matplotlib.pyplot as plt
import numpy as np
import george
import emcee
import StarData
import SpectralTypeRelations

def classify_filename(fname, type='bright'):
    """
    Given a CCF filename, it classifies the star combination, temperature, metallicity, and vsini
    :param fname:
    :return:
    """
    # First, remove any leading directories
    fname = fname.split('/')[-1]

    # Star combination
    m1 = re.search('\.[0-9]+kps', fname)
    stars = fname[:m1.start()]
    star1 = stars.split('+')[0].replace('_', ' ')
    star2 = stars.split('+')[1].split('_{}'.format(type))[0].replace('_', ' ')

    # secondary star vsini
    vsini = float(fname[m1.start() + 1:].split('kps')[0])

    # Temperature
    m2 = re.search('[0-9]+\.0K', fname)
    temp = float(m2.group()[:-1])

    # logg
    m3 = re.search('K\+[0-9]\.[0-9]', fname)
    logg = float(m3.group()[1:])

    # metallicity
    metal = float(fname.split(str(logg))[-1])

    return star1, star2, vsini, temp, logg, metal

def get_ccf_data(basedir, primary_name=None, secondary_name=None, vel_arr=np.arange(-900.0, 900.0, 0.1), type='bright'):
    """
    Searches the given directory for CCF files, and classifies
    by star, temperature, metallicity, and vsini
    :param basedir: The directory to search for CCF files
    :keyword primary_name: Optional keyword. If given, it will only get the requested primary star data
    :keyword secondary_name: Same as primary_name, but only reads ccfs for the given secondary
    :keyword vel_arr: The velocities to interpolate each ccf at
    :return: pandas DataFrame
    """
    if not basedir.endswith('/'):
        basedir += '/'
    all_files = ['{}{}'.format(basedir, f) for f in os.listdir(basedir) if type in f.lower()]
    primary = []
    secondary = []
    vsini_values = []
    temperature = []
    gravity = []
    metallicity = []
    ccf = []
    for fname in all_files:
        star1, star2, vsini, temp, logg, metal = classify_filename(fname, type=type)
        if primary_name is not None and star1.lower() != primary_name.lower():
            continue
        if secondary_name is not None and star2.lower() != secondary_name.lower():
            continue
        vel, corr = np.loadtxt(fname, unpack=True)
        fcn = spline(vel, corr)
        ccf.append(fcn(vel_arr))
        primary.append(star1)
        secondary.append(star2)
        vsini_values.append(vsini)
        temperature.append(temp)
        gravity.append(logg)
        metallicity.append(metal)

    # Make a pandas dataframe with all this data
    df = pandas.DataFrame(data={'Primary': primary, 'Secondary': secondary, 'Temperature': temperature,
                                'vsini': vsini_values, 'logg': gravity, '[Fe/H]': metallicity, 'CCF': ccf})
    return df

def get_ccf_summary(basedir, vel_arr=np.arange(-900.0, 900.0, 0.1), velocity='highest', type='bright'):
    """
    Very similar to get_ccf_data, but does it in a way that is more memory efficient
    :param basedir: The directory to search for CCF files
    :keyword velocity: The velocity to measure the CCF at. The default is 'highest', and uses the maximum of the ccf
    :keyword vel_arr: The velocities to interpolate each ccf at
    :return: pandas DataFrame
    """
    if not basedir.endswith('/'):
        basedir += '/'
    all_files = ['{}{}'.format(basedir, f) for f in os.listdir(basedir) if type in f.lower()]
    file_dict = defaultdict(lambda: defaultdict(list))
    for fname in all_files:
        star1, star2, vsini, temp, logg, metal = classify_filename(fname, type=type)
        file_dict[star1][star2].append(fname)

    # Now, read the ccfs for each primary/secondary combo, and find the best combination
    summary_dfs = []
    for primary in file_dict.keys():
        for secondary in file_dict[primary].keys():
            data = get_ccf_data(basedir, primary_name=primary, secondary_name=secondary,
                                vel_arr=vel_arr, type=type)
            summary_dfs.append(find_best_pars(data, velocity=velocity, vel_arr=vel_arr))

    return pandas.concat(summary_dfs, ignore_index=True)

def find_best_pars(df, velocity='highest', vel_arr=np.arange(-900.0, 900.0, 0.1)):
    """
    Find the 'best-fit' parameters for each combination of primary and secondary star
    :param df: the dataframe to search in
    :keyword velocity: The velocity to measure the CCF at. The default is 'highest', and uses the maximum of the ccf
    :keyword vel_arr: The velocities to interpolate each ccf at
    :return: a dataframe with keys of primary, secondary, and the parameters
    """
    # Get the names of the primary and secondary stars
    primary_names = pandas.unique(df.Primary)
    secondary_names = pandas.unique(df.Secondary)

    # Find the ccf value at the given velocity
    if velocity == 'highest':
        fcn = lambda row: (np.max(row), vel_arr[np.argmax(row)])
        vals = df['CCF'].map(fcn)
        df['ccf_max'] = vals.map(lambda l: l[0])
        df['rv'] = vals.map(lambda l: l[1])
        # df['ccf_max'] = df['CCF'].map(np.max)
    else:
        df['ccf_max'] = df['CCF'].map(lambda arr: arr[np.argmin(np.abs(vel_arr - velocity))])

    # Find the best parameter for each combination
    d = defaultdict(list)
    for primary in primary_names:
        for secondary in secondary_names:
            good = df.loc[(df.Primary == primary) & (df.Secondary == secondary)]
            best = good.loc[good.ccf_max == good.ccf_max.max()]
            d['Primary'].append(primary)
            d['Secondary'].append(secondary)
            d['Temperature'].append(best['Temperature'].item())
            d['vsini'].append(best['vsini'].item())
            d['logg'].append(best['logg'].item())
            d['[Fe/H]'].append(best['[Fe/H]'].item())
            d['rv'].append(best['rv'].item())

    return pandas.DataFrame(data=d)

def get_detected_objects(df, tol=1.0):
    """
    Takes a summary dataframe with RV information. Finds the median rv for each star,
      and removes objects that are 'tol' km/s from the median value
    :param df: A summary dataframe, such as created by find_best_pars
    :param tol: The tolerance, in km/s, to accept an observation as detected
    :return: a dataframe containing only detected companions
    """
    secondary_names = pandas.unique(df.Secondary)
    secondary_to_rv = defaultdict(float)
    for secondary in secondary_names:
        rv = df.loc[df.Secondary == secondary]['rv'].median()
        secondary_to_rv[secondary] = rv
        print secondary, rv

    keys = df.Secondary.values
    good = df.loc[abs(df.rv.values - np.array(itemgetter(*keys)(secondary_to_rv))) < tol]
    return good

def add_actual_temperature(df, method='spt'):
    """
    Add the actual temperature to a given summary dataframe
    :param df: The dataframe to which we will add the actual secondary star temperature
    :param method: How to get the actual temperature. Options are:
                   - 'spt': Use main-sequence relationships to go from spectral type --> temperature
                   - 'excel': Use tabulated data, available in the file 'SecondaryStar_Temperatures.xls'
    :return: copy of the original dataframe, with an extra column for the secondary star temperature
    """
    # First, get a list of the secondary stars in the data
    secondary_names = pandas.unique(df.Secondary)
    secondary_to_temperature = defaultdict(float)
    secondary_to_error = defaultdict(float)

    if method.lower() == 'spt':
        MS = SpectralTypeRelations.MainSequence()
        for secondary in secondary_names:
            star_data = StarData.GetData(secondary)
            spt = star_data.spectype[0] + re.search('[0-9]\.*[0-9]*', star_data.spectype).group()
            T_sec = MS.Interpolate(MS.Temperature, spt)
            secondary_to_temperature[secondary] = T_sec

    elif method.lower() == 'excel':
        table = pandas.read_excel('SecondaryStar_Temperatures.xls', 0)
        for secondary in secondary_names:
            T_sec = table.loc[table.Star.str.lower().str.contains(secondary.strip().lower())]['Literature_Temp'].item()
            T_error = table.loc[table.Star.str.lower().str.contains(secondary.strip().lower())][
                'Literature_error'].item()
            secondary_to_temperature[secondary] = T_sec
            secondary_to_error[secondary] = T_error

    df['Tactual'] = df['Secondary'].map(lambda s: secondary_to_temperature[s])
    df['Tact_err'] = df['Secondary'].map(lambda s: secondary_to_error[s])
    return

def make_gaussian_process_samples(df):
    """
    Make a gaussian process fitting the Tactual-Tmeasured relationship
    :param df: pandas DataFrame with columns 'Temperature' (with the measured temperature)
                 and 'Tactual' (for the actual temperature)
    :return: emcee sampler instance
    """
    # First, find the uncertainties at each actual temperature
    # Tactual = df['Tactual'].values
    #Tmeasured = df['Temperature'].values
    #error = df['Tact_err'].values
    temp = df.groupby('Temperature').mean()['Tactual']
    Tmeasured = temp.keys().values
    Tactual = temp.values
    error = np.nan_to_num(df.groupby('Temperature').std(ddof=1)['Tactual'].values)
    default = np.median(error[error > 1])
    error = np.maximum(error, np.ones(error.size) * default)
    for Tm, Ta, e in zip(Tmeasured, Tactual, error):
        print Tm, Ta, e
    plt.figure(1)
    plt.errorbar(Tmeasured, Tactual, yerr=error, fmt='.k', capsize=0)
    plt.plot(Tmeasured, Tmeasured, 'r--')
    plt.xlim((min(Tmeasured) - 100, max(Tmeasured) + 100))
    plt.xlabel('Measured Temperature')
    plt.ylabel('Actual Temperature')
    plt.show(block=False)

    # Define some functions to use in the GP fit
    def model(pars, T):
        #polypars = pars[2:]
        #return np.poly1d(polypars)(T)
        return T

    def lnlike(pars, Tact, Tmeas, Terr):
        a, tau = np.exp(pars[:2])
        gp = george.GP(a * kernels.ExpSquaredKernel(tau))
        gp.compute(Tmeas, Terr)
        return gp.lnlikelihood(Tact - model(pars, Tmeas))

    def lnprior(pars):
        lna, lntau = pars[:2]
        polypars = pars[2:]
        if -20 < lna < 20 and 4 < lntau < 20:
            return 0.0
        return -np.inf

    def lnprob(pars, x, y, yerr):
        lp = lnprior(pars)
        return lp + lnlike(pars, x, y, yerr) if np.isfinite(lp) else -np.inf

    # Set up the emcee fitter
    initial = np.array([0, 6])#, 1.0, 0.0])
    ndim = len(initial)
    nwalkers = 100
    p0 = [np.array(initial) + 1e-8 * np.random.randn(ndim) for i in xrange(nwalkers)]
    sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(Tactual, Tmeasured, error))

    print 'Running first burn-in'
    p1, lnp, _ = sampler.run_mcmc(p0, 500)
    sampler.reset()

    print "Running second burn-in..."
    p_best = p1[np.argmax(lnp)]
    p2 = [p_best + 1e-8 * np.random.randn(ndim) for i in xrange(nwalkers)]
    p3, _, _ = sampler.run_mcmc(p2, 250)
    sampler.reset()

    print "Running production..."
    sampler.run_mcmc(p3, 1000)

    # Plot a bunch of the fits
    print "Plotting..."
    N = 100
    Tvalues = np.arange(3300, 7000, 20)
    idx = np.argsort(-sampler.lnprobability.flatten())[:N]  # Get N 'best' curves
    par_vals = sampler.flatchain[idx]
    for i, pars in enumerate(par_vals):
        a, tau = np.exp(pars[:2])
        gp = george.GP(a * kernels.ExpSquaredKernel(tau))
        gp.compute(Tmeasured, error)
        s = gp.sample_conditional(Tactual - model(pars, Tmeasured), Tvalues) + model(pars, Tvalues)
        plt.plot(Tvalues, s, 'b-', alpha=0.1)
    plt.draw()

    # Finally, get posterior samples at all the possibly measured temperatures
    print 'Generating posterior samples at all temperatures...'
    N = 10000  # This is 1/10th of the total number of samples!
    idx = np.argsort(-sampler.lnprobability.flatten())[:N]  # Get N 'best' curves
    par_vals = sampler.flatchain[idx]
    Tvalues = np.arange(3000, 6900, 100)
    gp_posterior = []
    for pars in par_vals:
        a, tau = np.exp(pars[:2])
        gp = george.GP(a * kernels.ExpSquaredKernel(tau))
        gp.compute(Tmeasured, error)
        s = gp.sample_conditional(Tactual - model(pars, Tmeasured), Tvalues) + model(pars, Tvalues)
        gp_posterior.append(s)

    # Finally, make confidence intervals for the actual temperatures
    gp_posterior = np.array(gp_posterior)
    l, m, h = np.percentile(gp_posterior, [16.0, 50.0, 84.0], axis=0)
    conf = pandas.DataFrame(data={'Measured Temperature': Tvalues, 'Actual Temperature': m,
                                  'Lower Bound': l, 'Upper bound': h})
    conf.to_csv('Confidence_Intervals.csv', index=False)

    return sampler, np.array(gp_posterior)

def check_posterior(df, posterior, Tvalues):
    """
    Checks the posterior samples: Are 95% of the measurements within 2-sigma of the prediction?
    :param df: The summary dataframe
    :param posterior: The MCMC predicted values
    :param Tvalues: The measured temperatures the posterior was made with
    :return: boolean, as well as some warning messages if applicable
    """
    # First, make 2-sigma confidence intervals
    l, m, h = np.percentile(posterior, [5.0, 50.0, 95.0], axis=0)

    # Save the confidence intervals
    # conf = pandas.DataFrame(data={'Measured Temperature': Tvalues, 'Actual Temperature': m,
    #                              'Lower Bound': l, 'Upper bound': h})
    #conf.to_csv('Confidence_Intervals.csv', index=False)

    Ntot = []  # The total number of observations with the given measured temperature
    Nacc = []  # The number that have actual temperatures within the confidence interval

    g = df.groupby('Temperature')
    for i, T in enumerate(Tvalues):
        if T in g.groups.keys():
            Ta = g.get_group(T)['Tactual']
            low, high = l[i], h[i]
            Ntot.append(len(Ta))
            Nacc.append(len(Ta.loc[(Ta >= low) & (Ta <= high)]))
            p = float(Nacc[-1]) / float(Ntot[-1])
            if p < 0.95:
                logging.warn(
                    'Only {}/{} of the samples ({:.2f}%) were accepted for T = {} K'.format(Nacc[-1], Ntot[-1], p * 100,
                                                                                T))
                print low, high
                print sorted(Ta)
        else:
            Ntot.append(0)
            Nacc.append(0)

    p = float(sum(Nacc)) / float(sum(Ntot))
    if p < 0.95:
        logging.warn('Only {:.2f}% of the total samples were accepted!'.format(p * 100))
        return False
    return True

if __name__ == '__main__':
    pass


# Copyright (c) Charl P. Botha, TU Delft.
# All rights reserved.
# See COPYRIGHT for details.

import config
from install_package import InstallPackage
import os
import shutil
import sys
import utils
from distutils import sysconfig

MPL_VER = "1.1.0"

if os.name == 'posix':
    MPL_ARCHIVE = "matplotlib-%s.tar.gz" % (MPL_VER,)
    MPL_URL = "http://surfnet.dl.sourceforge.net/sourceforge/matplotlib/%s" % \
          (MPL_ARCHIVE,)

elif os.name == 'nt':
    if config.WINARCH_STR == 'x64':
        WINTHINGY = 'win-amd64'

    else:
        WINTHINGY = 'win32'

    MPL_ARCHIVE = "matplotlib-%s.%s-py2.7.exe" % (MPL_VER, WINTHINGY)
    MPL_URL = "http://graphics.tudelft.nl/~cpbotha/files/devide/johannes_support/gohlke/%s" % (MPL_ARCHIVE,)

MPL_DIRBASE = "matplotlib-%s" % (MPL_VER,)

# I prefer that this be built with numpy, but it is not a dependency
# per se
dependencies = []

class matplotlib(InstallPackage):

    def __init__(self):
        self.tbfilename = os.path.join(config.archive_dir, MPL_ARCHIVE)
        self.build_dir = os.path.join(config.build_dir, MPL_DIRBASE)
        self.inst_dir = os.path.join(config.inst_dir, 'matplotlib')

    def get(self):
        if os.path.exists(self.tbfilename):
            utils.output("%s already present, not downloading." %
                         (MPL_ARCHIVE,))
        else:
            utils.goto_archive()
            utils.urlget(MPL_URL)

    def unpack(self):
        if os.path.isdir(self.build_dir):
            utils.output("MATPLOTLIB source already unpacked, not redoing.")
        else:
            if os.name == 'posix':
                utils.output("Unpacking MATPLOTLIB source.")
                utils.unpack_build(self.tbfilename)
            else:
                utils.output("Unpacking MATPLOTLIB binaries.")
                os.mkdir(self.build_dir)
                os.chdir(self.build_dir)
                utils.unpack(self.tbfilename)

    def configure(self):
        if os.name == 'nt':
            utils.output("Skipping configure (WINDOWS).")
            return

        # pre-configure setup.py and setupext.py so that everything is
        # found and configured as we want it.
        os.chdir(self.build_dir)

        if os.path.exists('setup.py.new'):
            utils.output('matplotlib already configured.  Skipping step.')

        else:
            # pre-filter setup.py
            repls = [("(BUILD_GTKAGG\s*=\s*).*", "\\1 0"),
                     ("(BUILD_GTK\s*=\s*).*", "\\1 0"),
                     ("(BUILD_TKAGG\s*=\s*).*", "\\1 0"),
                     ("(BUILD_WXAGG\s*=\s*).*", "\\1 1"),
                     ("(rc\s*=\s*dict\().*",
                      "\\1 {'backend':'PS', 'numerix':'numpy'} )")]

            utils.re_sub_filter_file(repls, 'setup.py')

    def build(self):
        if os.name == 'nt':
            utils.output("Skipping build (WINDOWS).")
            return

        os.chdir(self.build_dir)

        # weak test... there are .so files deeper, but they're in platform
        # specific directories
        if os.path.exists('build'):
            utils.output('matplotlib already built.  Skipping step.')

        else:

            # add wx bin to path so that wx-config can be found
            os.environ['PATH'] = "%s%s%s" % (config.WX_BIN_PATH,
                                             os.pathsep, os.environ['PATH'])

            ret = os.system('%s setup.py build' % (sys.executable,))

            if ret != 0:
                utils.error('matplotlib build failed.  Please fix and try again.')

    def install(self):
        # to test for install, just do python -c "import matplotlib"
        # and test the result (we could just import directly, but that would
        # only work once our invoking python has been stopped and started
        # again)
        os.chdir(config.archive_dir) # we need to be elsewhere!
        ret = os.system('%s -c "import matplotlib"' % (sys.executable,))
        if ret == 0:
            utils.output('matplotlib already installed.  Skipping step.')

        else:
            utils.output('ImportError test shows that matplotlib is not '
                         'installed.  Installing...')

            if os.name == 'nt':
                self.install_nt()
            else:
                self.install_posix()

            # make sure the backend is set to WXAgg
            # and that interactive is set to True
            rcfn = os.path.join(
                    config.PYTHON_SITE_PACKAGES,
                    'matplotlib', 'mpl-data', 'matplotlibrc')
            utils.re_sub_filter_file(
                    [("(\s*backend\s*\:).*", "\\1 WXAgg"),
                     ("#*(\s*interactive\s:).*","\\1 True")], rcfn)

    def install_nt(self):
        sp_dir = sysconfig.get_python_lib()
        utils.copy_glob(os.path.join(self.build_dir, 'PLATLIB', '*'), sp_dir)

    def install_posix(self):
        os.chdir(self.build_dir)

        # add wx bin to path so that wx-config can be found
        os.environ['PATH'] = "%s%s%s" % (config.WX_BIN_PATH,
                                         os.pathsep, os.environ['PATH'])

        ret = os.system('%s setup.py install' % (sys.executable,))

        if ret != 0:
            utils.error(
                'matplotlib install failed.  Please fix and try again.')

    def clean_build(self):
        utils.output("Removing build and install directories.")
        if os.path.exists(self.build_dir):
            shutil.rmtree(self.build_dir)

        from distutils import sysconfig
        matplotlib_instdir = os.path.join(sysconfig.get_python_lib(),
                                          'matplotlib')

        if os.path.exists(matplotlib_instdir):
            shutil.rmtree(matplotlib_instdir)

    def get_installed_version(self):
        import matplotlib
        return matplotlib.__version__


"""
==============================================================
Restricted Boltzmann Machine features for digit classification
==============================================================

For greyscale image data where pixel values can be interpreted as degrees of
blackness on a white background, like handwritten digit recognition, the
Bernoulli Restricted Boltzmann machine model (:class:`BernoulliRBM
<sklearn.neural_network.BernoulliRBM>`) can perform effective non-linear
feature extraction.

In order to learn good latent representations from a small dataset, we
artificially generate more labeled data by perturbing the training data with
linear shifts of 1 pixel in each direction.

This example shows how to build a classification pipeline with a BernoulliRBM
feature extractor and a :class:`LogisticRegression
<sklearn.linear_model.LogisticRegression>` classifier. The hyperparameters
of the entire model (learning rate, hidden layer size, regularization)
were optimized by grid search, but the search is not reproduced here because
of runtime constraints.

Logistic regression on raw pixel values is presented for comparison. The
example shows that the features extracted by the BernoulliRBM help improve the
classification accuracy.
"""

from __future__ import print_function

print(__doc__)

# Authors: Yann N. Dauphin, Vlad Niculae, Gabriel Synnaeve
# License: BSD

import numpy as np
import matplotlib.pyplot as plt

from scipy.ndimage import convolve
from sklearn import linear_model, datasets, metrics
from sklearn.cross_validation import train_test_split
from sklearn.neural_network import BernoulliRBM
from sklearn.pipeline import Pipeline

###############################################################################
# Setting up

def nudge_dataset(X, Y):
    """
    This produces a dataset 5 times bigger than the original one,
    by moving the 8x8 images in X around by 1px to left, right, down, up
    """
    direction_vectors = [
        [[0, 1, 0],
         [0, 0, 0],
         [0, 0, 0]],

        [[0, 0, 0],
         [1, 0, 0],
         [0, 0, 0]],

        [[0, 0, 0],
         [0, 0, 1],
         [0, 0, 0]],

        [[0, 0, 0],
         [0, 0, 0],
         [0, 1, 0]]]

    shift = lambda x, w: convolve(x.reshape((8, 8)), mode='constant',
                                  weights=w).ravel()
    X = np.concatenate([X] +
                       [np.apply_along_axis(shift, 1, X, vector)
                        for vector in direction_vectors])
    Y = np.concatenate([Y for _ in range(5)], axis=0)
    return X, Y

# Load Data
digits = datasets.load_digits()
X = np.asarray(digits.data, 'float32')
X, Y = nudge_dataset(X, digits.target)
X = (X - np.min(X, 0)) / (np.max(X, 0) + 0.0001)  # 0-1 scaling

X_train, X_test, Y_train, Y_test = train_test_split(X, Y,
                                                    test_size=0.2,
                                                    random_state=0)

# Models we will use
logistic = linear_model.LogisticRegression()
rbm = BernoulliRBM(random_state=0, verbose=True)

classifier = Pipeline(steps=[('rbm', rbm), ('logistic', logistic)])

###############################################################################
# Training

# Hyper-parameters. These were set by cross-validation,
# using a GridSearchCV. Here we are not performing cross-validation to
# save time.
rbm.learning_rate = 0.06
rbm.n_iter = 20
# More components tend to give better prediction performance, but larger
# fitting time
rbm.n_components = 100
logistic.C = 6000.0

# Training RBM-Logistic Pipeline
classifier.fit(X_train, Y_train)

# Training Logistic regression
logistic_classifier = linear_model.LogisticRegression(C=100.0)
logistic_classifier.fit(X_train, Y_train)

###############################################################################
# Evaluation

print()
print("Logistic regression using RBM features:\n%s\n" % (
    metrics.classification_report(
        Y_test,
        classifier.predict(X_test))))

print("Logistic regression using raw pixel features:\n%s\n" % (
    metrics.classification_report(
        Y_test,
        logistic_classifier.predict(X_test))))

###############################################################################
# Plotting

plt.figure(figsize=(4.2, 4))
for i, comp in enumerate(rbm.components_):
    plt.subplot(10, 10, i + 1)
    plt.imshow(comp.reshape((8, 8)), cmap=plt.cm.gray_r,
               interpolation='nearest')
    plt.xticks(())
    plt.yticks(())
plt.suptitle('100 components extracted by RBM', fontsize=16)
plt.subplots_adjust(0.08, 0.02, 0.92, 0.85, 0.08, 0.23)

plt.show()


TalkingData AdTracking Fraud Detection/perform_ensembling.py

import os
import glob
import shutil
import datetime
import numpy as np
import pandas as pd

# Dataset
PROJECT_NAME = "TalkingData AdTracking Fraud Detection"
PROJECT_FOLDER_PATH = os.path.join(os.path.expanduser("~"), "Documents/Dataset",
                                   PROJECT_NAME)

# Submission
TEAM_NAME = "Aurora"
SUBMISSION_FOLDER_PATH = os.path.join(PROJECT_FOLDER_PATH, "submission")
os.makedirs(SUBMISSION_FOLDER_PATH, exist_ok=True)

# Ensembling
WORKSPACE_FOLDER_PATH = os.path.join(PROJECT_FOLDER_PATH, "script/Mar_25_3")
KEYWORD = "DL"

# Generate a zip archive for a file
create_zip_archive = lambda file_path: shutil.make_archive(
    file_path[:file_path.rindex(".")], "zip",
    os.path.abspath(os.path.join(file_path, "..")), os.path.basename(file_path))

def run():
    print("Searching for submissions with keyword {} at {} ...".format(
        KEYWORD, WORKSPACE_FOLDER_PATH))
    submission_file_path_list = sorted(
        glob.glob(os.path.join(WORKSPACE_FOLDER_PATH, "*{}*".format(KEYWORD))))
    assert len(submission_file_path_list) != 0

    ranking_array_list = []
    for submission_file_path in submission_file_path_list:
        print("Loading {} ...".format(submission_file_path))
        submission_df = pd.read_csv(submission_file_path)

        print("Ranking the entries ...")
        index_series = submission_df["is_attributed"].argsort()
        ranking_array = np.zeros(index_series.shape, dtype=np.uint32)
        ranking_array[index_series] = np.arange(len(index_series))
        ranking_array_list.append(ranking_array)

    ensemble_df = submission_df.copy()
    ensemble_prediction_array = np.mean(ranking_array_list, axis=0)
    apply_normalization = lambda data_array: 1.0 * (data_array - np.min(
        data_array)) / (np.max(data_array) - np.min(data_array))
    ensemble_df["is_attributed"] = apply_normalization(
        ensemble_prediction_array)
    ensemble_file_path = os.path.join(
        SUBMISSION_FOLDER_PATH, "{} {} {}.csv".format(
            TEAM_NAME, KEYWORD,
            str(datetime.datetime.now()).split(".")[0]).replace(" ", "_"))
    print("Saving submission to {} ...".format(ensemble_file_path))
    ensemble_df.to_csv(ensemble_file_path, float_format="%.6f", index=False)
    compressed_ensemble_file_path = create_zip_archive(ensemble_file_path)
    print("Saving compressed submission to {} ...".format(
        compressed_ensemble_file_path))

    print("All done!")

if __name__ == "__main__":
    run()