"""Metrics to assess performance on regression task

Functions named as ``*_score`` return a scalar value to maximize: the higher
the better

Function named as ``*_error`` or ``*_loss`` return a scalar value to minimize:
the lower the better
"""

# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Mathieu Blondel <mathieu@mblondel.org>
#          Olivier Grisel <olivier.grisel@ensta.org>
#          Arnaud Joly <a.joly@ulg.ac.be>
#          Jochen Wersdorfer <jochen@wersdoerfer.de>
#          Lars Buitinck <L.J.Buitinck@uva.nl>
#          Joel Nothman <joel.nothman@gmail.com>
#          Noel Dawe <noel@dawe.me>
#          Manoj Kumar <manojkumarsivaraj334@gmail.com>
#          Michael Eickenberg <michael.eickenberg@gmail.com>
#          Konstantin Shmelkov <konstantin.shmelkov@polytechnique.edu>
# License: BSD 3 clause

from __future__ import division

import numpy as np

from ..utils.validation import check_array, check_consistent_length
from ..utils.validation import column_or_1d

import warnings

__ALL__ = [
    "mean_absolute_error",
    "mean_squared_error",
    "median_absolute_error",
    "r2_score",
    "explained_variance_score"
]

def _check_reg_targets(y_true, y_pred, multioutput):
    """Check that y_true and y_pred belong to the same regression task

    Parameters
    ----------
    y_true : array-like,

    y_pred : array-like,

    multioutput : array-like or string in ['raw_values', uniform_average',
        'variance_weighted'] or None
        None is accepted due to backward compatibility of r2_score().

    Returns
    -------
    type_true : one of {'continuous', continuous-multioutput'}
        The type of the true target data, as output by
        'utils.multiclass.type_of_target'

    y_true : array-like of shape = (n_samples, n_outputs)
        Ground truth (correct) target values.

    y_pred : array-like of shape = (n_samples, n_outputs)
        Estimated target values.

    multioutput : array-like of shape = (n_outputs) or string in ['raw_values',
        uniform_average', 'variance_weighted'] or None
        Custom output weights if ``multioutput`` is array-like or
        just the corresponding argument if ``multioutput`` is a
        correct keyword.

    """
    check_consistent_length(y_true, y_pred)
    y_true = check_array(y_true, ensure_2d=False)
    y_pred = check_array(y_pred, ensure_2d=False)

    if y_true.ndim == 1:
        y_true = y_true.reshape((-1, 1))

    if y_pred.ndim == 1:
        y_pred = y_pred.reshape((-1, 1))

    if y_true.shape[1] != y_pred.shape[1]:
        raise ValueError("y_true and y_pred have different number of output "
                         "({0}!={1})".format(y_true.shape[1], y_pred.shape[1]))

    n_outputs = y_true.shape[1]
    multioutput_options = (None, 'raw_values', 'uniform_average',
                           'variance_weighted')
    if multioutput not in multioutput_options:
        multioutput = check_array(multioutput, ensure_2d=False)
        if n_outputs == 1:
            raise ValueError("Custom weights are useful only in "
                             "multi-output cases.")
        elif n_outputs != len(multioutput):
            raise ValueError(("There must be equally many custom weights "
                              "(%d) as outputs (%d).") %
                             (len(multioutput), n_outputs))
    y_type = 'continuous' if n_outputs == 1 else 'continuous-multioutput'

    return y_type, y_true, y_pred, multioutput

def mean_absolute_error(y_true, y_pred,
                        sample_weight=None,
                        multioutput='uniform_average'):
    """Mean absolute error regression loss

    Read more in the :ref:`User Guide <mean_absolute_error>`.

    Parameters
    ----------
    y_true : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Ground truth (correct) target values.

    y_pred : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Estimated target values.

    sample_weight : array-like of shape = (n_samples), optional
        Sample weights.

    multioutput : string in ['raw_values', 'uniform_average']
        or array-like of shape (n_outputs)
        Defines aggregating of multiple output values.
        Array-like value defines weights used to average errors.

        'raw_values' :
            Returns a full set of errors in case of multioutput input.

        'uniform_average' :
            Errors of all outputs are averaged with uniform weight.

    Returns
    -------
    loss : float or ndarray of floats
        If multioutput is 'raw_values', then mean absolute error is returned
        for each output separately.
        If multioutput is 'uniform_average' or an ndarray of weights, then the
        weighted average of all output errors is returned.

        MAE output is non-negative floating point. The best value is 0.0.

    Examples
    --------
    >>> from sklearn.metrics import mean_absolute_error
    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]
    >>> mean_absolute_error(y_true, y_pred)
    0.5
    >>> y_true = [[0.5, 1], [-1, 1], [7, -6]]
    >>> y_pred = [[0, 2], [-1, 2], [8, -5]]
    >>> mean_absolute_error(y_true, y_pred)
    0.75
    >>> mean_absolute_error(y_true, y_pred, multioutput='raw_values')
    array([ 0.5,  1. ])
    >>> mean_absolute_error(y_true, y_pred, multioutput=[0.3, 0.7])
    ... # doctest: +ELLIPSIS
    0.849...
    """
    y_type, y_true, y_pred, multioutput = _check_reg_targets(
        y_true, y_pred, multioutput)
    output_errors = np.average(np.abs(y_pred - y_true),
                               weights=sample_weight, axis=0)
    if multioutput == 'raw_values':
        return output_errors
    elif multioutput == 'uniform_average':
        # pass None as weights to np.average: uniform mean
        multioutput = None

    return np.average(output_errors, weights=multioutput)

def mean_squared_error(y_true, y_pred,
                       sample_weight=None,
                       multioutput='uniform_average'):
    """Mean squared error regression loss

    Read more in the :ref:`User Guide <mean_squared_error>`.

    Parameters
    ----------
    y_true : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Ground truth (correct) target values.

    y_pred : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Estimated target values.

    sample_weight : array-like of shape = (n_samples), optional
        Sample weights.

    multioutput : string in ['raw_values', 'uniform_average']
        or array-like of shape (n_outputs)
        Defines aggregating of multiple output values.
        Array-like value defines weights used to average errors.

        'raw_values' :
            Returns a full set of errors in case of multioutput input.

        'uniform_average' :
            Errors of all outputs are averaged with uniform weight.

    Returns
    -------
    loss : float or ndarray of floats
        A non-negative floating point value (the best value is 0.0), or an
        array of floating point values, one for each individual target.

    Examples
    --------
    >>> from sklearn.metrics import mean_squared_error
    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]
    >>> mean_squared_error(y_true, y_pred)
    0.375
    >>> y_true = [[0.5, 1],[-1, 1],[7, -6]]
    >>> y_pred = [[0, 2],[-1, 2],[8, -5]]
    >>> mean_squared_error(y_true, y_pred)  # doctest: +ELLIPSIS
    0.708...
    >>> mean_squared_error(y_true, y_pred, multioutput='raw_values')
    ... # doctest: +ELLIPSIS
    array([ 0.416...,  1.        ])
    >>> mean_squared_error(y_true, y_pred, multioutput=[0.3, 0.7])
    ... # doctest: +ELLIPSIS
    0.824...

    """
    y_type, y_true, y_pred, multioutput = _check_reg_targets(
        y_true, y_pred, multioutput)
    output_errors = np.average((y_true - y_pred) ** 2, axis=0,
                               weights=sample_weight)
    if multioutput == 'raw_values':
        return output_errors
    elif multioutput == 'uniform_average':
        # pass None as weights to np.average: uniform mean
        multioutput = None

    return np.average(output_errors, weights=multioutput)

def median_absolute_error(y_true, y_pred):
    """Median absolute error regression loss

    Read more in the :ref:`User Guide <median_absolute_error>`.

    Parameters
    ----------
    y_true : array-like of shape = (n_samples)
        Ground truth (correct) target values.

    y_pred : array-like of shape = (n_samples)
        Estimated target values.

    Returns
    -------
    loss : float
        A positive floating point value (the best value is 0.0).

    Examples
    --------
    >>> from sklearn.metrics import median_absolute_error
    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]
    >>> median_absolute_error(y_true, y_pred)
    0.5

    """
    y_type, y_true, y_pred, _ = _check_reg_targets(y_true, y_pred,
                                                   'uniform_average')
    if y_type == 'continuous-multioutput':
        raise ValueError("Multioutput not supported in median_absolute_error")
    return np.median(np.abs(y_pred - y_true))

def explained_variance_score(y_true, y_pred,
                             sample_weight=None,
                             multioutput='uniform_average'):
    """Explained variance regression score function

    Best possible score is 1.0, lower values are worse.

    Read more in the :ref:`User Guide <explained_variance_score>`.

    Parameters
    ----------
    y_true : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Ground truth (correct) target values.

    y_pred : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Estimated target values.

    sample_weight : array-like of shape = (n_samples), optional
        Sample weights.

    multioutput : string in ['raw_values', 'uniform_average', \
                'variance_weighted'] or array-like of shape (n_outputs)
        Defines aggregating of multiple output scores.
        Array-like value defines weights used to average scores.

        'raw_values' :
            Returns a full set of scores in case of multioutput input.

        'uniform_average' :
            Scores of all outputs are averaged with uniform weight.

        'variance_weighted' :
            Scores of all outputs are averaged, weighted by the variances
            of each individual output.

    Returns
    -------
    score : float or ndarray of floats
        The explained variance or ndarray if 'multioutput' is 'raw_values'.

    Notes
    -----
    This is not a symmetric function.

    Examples
    --------
    >>> from sklearn.metrics import explained_variance_score
    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]
    >>> explained_variance_score(y_true, y_pred)  # doctest: +ELLIPSIS
    0.957...
    >>> y_true = [[0.5, 1], [-1, 1], [7, -6]]
    >>> y_pred = [[0, 2], [-1, 2], [8, -5]]
    >>> explained_variance_score(y_true, y_pred, multioutput='uniform_average')
    ... # doctest: +ELLIPSIS
    0.983...

    """
    y_type, y_true, y_pred, multioutput = _check_reg_targets(
        y_true, y_pred, multioutput)

    y_diff_avg = np.average(y_true - y_pred, weights=sample_weight, axis=0)
    numerator = np.average((y_true - y_pred - y_diff_avg) ** 2,
                           weights=sample_weight, axis=0)

    y_true_avg = np.average(y_true, weights=sample_weight, axis=0)
    denominator = np.average((y_true - y_true_avg) ** 2,
                             weights=sample_weight, axis=0)

    nonzero_numerator = numerator != 0
    nonzero_denominator = denominator != 0
    valid_score = nonzero_numerator & nonzero_denominator
    output_scores = np.ones(y_true.shape[1])

    output_scores[valid_score] = 1 - (numerator[valid_score] /
                                      denominator[valid_score])
    output_scores[nonzero_numerator & ~nonzero_denominator] = 0.
    if multioutput == 'raw_values':
        # return scores individually
        return output_scores
    elif multioutput == 'uniform_average':
        # passing to np.average() None as weights results is uniform mean
        avg_weights = None
    elif multioutput == 'variance_weighted':
        avg_weights = denominator
    else:
        avg_weights = multioutput

    return np.average(output_scores, weights=avg_weights)

def r2_score(y_true, y_pred,
             sample_weight=None,
             multioutput=None):
    """R^2 (coefficient of determination) regression score function.

    Best possible score is 1.0 and it can be negative (because the
    model can be arbitrarily worse). A constant model that always
    predicts the expected value of y, disregarding the input features,
    would get a R^2 score of 0.0.

    Read more in the :ref:`User Guide <r2_score>`.

    Parameters
    ----------
    y_true : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Ground truth (correct) target values.

    y_pred : array-like of shape = (n_samples) or (n_samples, n_outputs)
        Estimated target values.

    sample_weight : array-like of shape = (n_samples), optional
        Sample weights.

    multioutput : string in ['raw_values', 'uniform_average',
                'variance_weighted'] or None or array-like of shape (n_outputs)
        Defines aggregating of multiple output scores.
        Array-like value defines weights used to average scores.
        Default value correponds to 'variance_weighted', but
        will be changed to 'uniform_average' in next versions.

        'raw_values' :
            Returns a full set of scores in case of multioutput input.

        'uniform_average' :
            Scores of all outputs are averaged with uniform weight.

        'variance_weighted' :
            Scores of all outputs are averaged, weighted by the variances
            of each individual output.

    Returns
    -------
    z : float or ndarray of floats
        The R^2 score or ndarray of scores if 'multioutput' is
        'raw_values'.

    Notes
    -----
    This is not a symmetric function.

    Unlike most other scores, R^2 score may be negative (it need not actually
    be the square of a quantity R).

    References
    ----------
    .. [1] `Wikipedia entry on the Coefficient of determination
            <http://en.wikipedia.org/wiki/Coefficient_of_determination>`_

    Examples
    --------
    >>> from sklearn.metrics import r2_score
    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]
    >>> r2_score(y_true, y_pred)  # doctest: +ELLIPSIS
    0.948...
    >>> y_true = [[0.5, 1], [-1, 1], [7, -6]]
    >>> y_pred = [[0, 2], [-1, 2], [8, -5]]
    >>> r2_score(y_true, y_pred, multioutput='variance_weighted')  # doctest: +ELLIPSIS
    0.938...

    """
    y_type, y_true, y_pred, multioutput = _check_reg_targets(
        y_true, y_pred, multioutput)

    if sample_weight is not None:
        sample_weight = column_or_1d(sample_weight)
        weight = sample_weight[:, np.newaxis]
    else:
        weight = 1.

    numerator = (weight * (y_true - y_pred) ** 2).sum(axis=0,
                                                      dtype=np.float64)
    denominator = (weight * (y_true - np.average(
        y_true, axis=0, weights=sample_weight)) ** 2).sum(axis=0,
                                                          dtype=np.float64)
    nonzero_denominator = denominator != 0
    nonzero_numerator = numerator != 0
    valid_score = nonzero_denominator & nonzero_numerator
    output_scores = np.ones([y_true.shape[1]])
    output_scores[valid_score] = 1 - (numerator[valid_score] /
                                      denominator[valid_score])
    # arbitrary set to zero to avoid -inf scores, having a constant
    # y_true is not interesting for scoring a regression anyway
    output_scores[nonzero_numerator & ~nonzero_denominator] = 0.
    if multioutput is None and y_true.shape[1] != 1:
        # @FIXME change in 0.18
        warnings.warn("Default 'multioutput' behavior now corresponds to "
                      "'variance_weighted' value, it will be changed "
                      "to 'uniform_average' in 0.18.",
                      DeprecationWarning)
        multioutput = 'variance_weighted'
    if multioutput == 'raw_values':
        # return scores individually
        return output_scores
    elif multioutput == 'uniform_average':
        # passing None as weights results is uniform mean
        avg_weights = None
    elif multioutput == 'variance_weighted':
        avg_weights = denominator
        # avoid fail on constant y or one-element arrays
        if not np.any(nonzero_denominator):
            if not np.any(nonzero_numerator):
                return 1.0
            else:
                return 0.0
    else:
        avg_weights = multioutput

    return np.average(output_scores, weights=avg_weights)


import os, sys
lib_path = os.path.abspath("/home/slundquist/workspace/PetaVision/plab/")
sys.path.append(lib_path)
from plotRecon import plotRecon
from plotReconError import plotReconError

#For plotting
#import matplotlib.pyplot as plt

outputDir = "/nh/compneuro/Data/momentLearn/output/simple_momentum_out/"
skipFrames = 10 #Only print every 20th frame
startFrames = 0
doPlotRecon = True
doPlotErr = False
errShowPlots = False
layers = [
   "a1_ImageRescale",
   "a4_Recon",
   ]

#Layers for constructing recon error
preErrLayers = [
   "a1_LeftDownsample",
   "a5_RightDownsample",
]

postErrLayers = [
   "a3_LeftRecon",
   "a7_RightRecon",
]

gtLayers = None
#gtLayers = [
#   #"a25_DepthRescale",
#   #"a25_DepthRescale",
#   #"a25_DepthRescale",
#   "a25_DepthRescale",
#   "a25_DepthRescale",
#   "a25_DepthRescale",
#]

preToPostScale = [
   .007,
   .007,
]

if(doPlotRecon):
   print("Plotting reconstructions")
   plotRecon(layers, outputDir, skipFrames)

if(doPlotErr):
   print("Plotting reconstruction error")
   plotReconError(preErrLayers, postErrLayers, preToPostScale, outputDir, errShowPlots, skipFrames, gtLayers)


"""
===========================================================
Hierarchical clustering: structured vs unstructured ward
===========================================================

Example builds a swiss roll dataset and runs
hierarchical clustering on their position.

For more information, see :ref:`hierarchical_clustering`.

In a first step, the hierarchical clustering is performed without connectivity
constraints on the structure and is solely based on distance, whereas in
a second step the clustering is restricted to the k-Nearest Neighbors
graph: it's a hierarchical clustering with structure prior.

Some of the clusters learned without connectivity constraints do not
respect the structure of the swiss roll and extend across different folds of
the manifolds. On the opposite, when opposing connectivity constraints,
the clusters form a nice parcellation of the swiss roll.
"""

# Authors : Vincent Michel, 2010
#           Alexandre Gramfort, 2010
#           Gael Varoquaux, 2010
# License: BSD 3 clause

print(__doc__)

import time as time
import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3
from sklearn.cluster import AgglomerativeClustering
from sklearn.datasets.samples_generator import make_swiss_roll

###############################################################################
# Generate data (swiss roll dataset)
n_samples = 1500
noise = 0.05
X, _ = make_swiss_roll(n_samples, noise)
# Make it thinner
X[:, 1] *= .5

###############################################################################
# Compute clustering
print("Compute unstructured hierarchical clustering...")
st = time.time()
ward = AgglomerativeClustering(n_clusters=6, linkage='ward').fit(X)
elapsed_time = time.time() - st
label = ward.labels_
print("Elapsed time: %.2fs" % elapsed_time)
print("Number of points: %i" % label.size)

###############################################################################
# Plot result
fig = plt.figure()
ax = p3.Axes3D(fig)
ax.view_init(7, -80)
for l in np.unique(label):
    ax.plot3D(X[label == l, 0], X[label == l, 1], X[label == l, 2],
              'o', color=plt.cm.jet(np.float(l) / np.max(label + 1)))
plt.title('Without connectivity constraints (time %.2fs)' % elapsed_time)

###############################################################################
# Define the structure A of the data. Here a 10 nearest neighbors
from sklearn.neighbors import kneighbors_graph
connectivity = kneighbors_graph(X, n_neighbors=10, include_self=False)

###############################################################################
# Compute clustering
print("Compute structured hierarchical clustering...")
st = time.time()
ward = AgglomerativeClustering(n_clusters=6, connectivity=connectivity,
                               linkage='ward').fit(X)
elapsed_time = time.time() - st
label = ward.labels_
print("Elapsed time: %.2fs" % elapsed_time)
print("Number of points: %i" % label.size)

###############################################################################
# Plot result
fig = plt.figure()
ax = p3.Axes3D(fig)
ax.view_init(7, -80)
for l in np.unique(label):
    ax.plot3D(X[label == l, 0], X[label == l, 1], X[label == l, 2],
              'o', color=plt.cm.jet(float(l) / np.max(label + 1)))
plt.title('With connectivity constraints (time %.2fs)' % elapsed_time)

plt.show()


"""
To run this, you'll need to have installed.

  * glmnet-python
  * scikit-learn (of course)

Does two benchmarks

First, we fix a training set and increase the number of
samples. Then we plot the computation time as function of
the number of samples.

In the second benchmark, we increase the number of dimensions of the
training set. Then we plot the computation time as function of
the number of dimensions.

In both cases, only 10% of the features are informative.
"""
import numpy as np
import gc
from time import time
from sklearn.datasets.samples_generator import make_regression

alpha = 0.1
# alpha = 0.01

def rmse(a, b):
    return np.sqrt(np.mean((a - b) ** 2))

def bench(factory, X, Y, X_test, Y_test, ref_coef):
    gc.collect()

    # start time
    tstart = time()
    clf = factory(alpha=alpha).fit(X, Y)
    delta = (time() - tstart)
    # stop time

    print("duration: %0.3fs" % delta)
    print("rmse: %f" % rmse(Y_test, clf.predict(X_test)))
    print("mean coef abs diff: %f" % abs(ref_coef - clf.coef_.ravel()).mean())
    return delta

if __name__ == '__main__':
    from glmnet.elastic_net import Lasso as GlmnetLasso
    from sklearn.linear_model import Lasso as ScikitLasso
    # Delayed import of pylab
    import pylab as pl

    scikit_results = []
    glmnet_results = []
    n = 20
    step = 500
    n_features = 1000
    n_informative = n_features / 10
    n_test_samples = 1000
    for i in range(1, n + 1):
        print('==================')
        print('Iteration %s of %s' % (i, n))
        print('==================')

        X, Y, coef_ = make_regression(
            n_samples=(i * step) + n_test_samples, n_features=n_features,
            noise=0.1, n_informative=n_informative, coef=True)

        X_test = X[-n_test_samples:]
        Y_test = Y[-n_test_samples:]
        X = X[:(i * step)]
        Y = Y[:(i * step)]

        print("benchmarking scikit-learn: ")
        scikit_results.append(bench(ScikitLasso, X, Y, X_test, Y_test, coef_))
        print("benchmarking glmnet: ")
        glmnet_results.append(bench(GlmnetLasso, X, Y, X_test, Y_test, coef_))

    pl.clf()
    xx = range(0, n * step, step)
    pl.title('Lasso regression on sample dataset (%d features)' % n_features)
    pl.plot(xx, scikit_results, 'b-', label='scikit-learn')
    pl.plot(xx, glmnet_results, 'r-', label='glmnet')
    pl.legend()
    pl.xlabel('number of samples to classify')
    pl.ylabel('Time (s)')
    pl.show()

    # now do a benchmark where the number of points is fixed
    # and the variable is the number of features

    scikit_results = []
    glmnet_results = []
    n = 20
    step = 100
    n_samples = 500

    for i in range(1, n + 1):
        print('==================')
        print('Iteration %02d of %02d' % (i, n))
        print('==================')
        n_features = i * step
        n_informative = n_features / 10

        X, Y, coef_ = make_regression(
            n_samples=(i * step) + n_test_samples, n_features=n_features,
            noise=0.1, n_informative=n_informative, coef=True)

        X_test = X[-n_test_samples:]
        Y_test = Y[-n_test_samples:]
        X = X[:n_samples]
        Y = Y[:n_samples]

        print("benchmarking scikit-learn: ")
        scikit_results.append(bench(ScikitLasso, X, Y, X_test, Y_test, coef_))
        print("benchmarking glmnet: ")
        glmnet_results.append(bench(GlmnetLasso, X, Y, X_test, Y_test, coef_))

    xx = np.arange(100, 100 + n * step, step)
    pl.figure('scikit-learn vs. glmnet benchmark results')
    pl.title('Regression in high dimensional spaces (%d samples)' % n_samples)
    pl.plot(xx, scikit_results, 'b-', label='scikit-learn')
    pl.plot(xx, glmnet_results, 'r-', label='glmnet')
    pl.legend()
    pl.xlabel('number of features')
    pl.ylabel('Time (s)')
    pl.axis('tight')
    pl.show()


#!/usr/bin/env python

"""
Pull all the temperature and salinity data out of a STOQS database no
matter what platform and write it out in Parquet file format.

This is a companion to select_data_in_columns_for_data_science.ipynb
where we operationalize the explorations demonstrated in this Notebook:

https://nbviewer.jupyter.org/github/stoqs/stoqs/blob/master/stoqs/contrib/notebooks/select_data_in_columns_for_data_science.ipynb

Sample command line executions:

(base) ➜  docker git:(master) ✗ docker-compose exec stoqs stoqs/contrib/parquet/extract_columns.py --db stoqs_canon_october2020 --platforms dorado -o dorado.parquet -v
INFO 2021-02-24 21:35:53,588 extract_columns.py _estimate_memory():161 Estimated required_memory = 146584085.76
INFO 2021-02-24 21:35:53,588 extract_columns.py _sql_to_df():65 Reading from SQL query into DataFrame...
INFO 2021-02-24 21:36:11,430 extract_columns.py _sql_to_df():77 df.shape: (2245467, 8) - read_sql_query() in 17.8 sec
INFO 2021-02-24 21:36:11,433 extract_columns.py _sql_to_df():78 df.memory_usage().sum(): 143710016
INFO 2021-02-24 21:36:13,876 extract_columns.py pivot_table_to_parquet():172 Writing data to file dorado.parquet...
INFO 2021-02-24 21:36:14,378 extract_columns.py pivot_table_to_parquet():176 dfp.shape: (169159, 16) - to_parquet() in 0.5 sec
INFO 2021-02-24 21:36:14,378 extract_columns.py pivot_table_to_parquet():177 Done
stoqs container peak memory usage: 2.1 GB

(base) ➜  docker git:(master) ✗ docker-compose exec stoqs stoqs/contrib/parquet/extract_columns.py --db stoqs_canon_october2020 --platforms pontus makai -o lrauv.parquet -v
INFO 2021-02-24 21:38:42,623 extract_columns.py _estimate_memory():161 Estimated required_memory = 645795521.28
INFO 2021-02-24 21:38:42,624 extract_columns.py _sql_to_df():65 Reading from SQL query into DataFrame...
INFO 2021-02-24 21:40:13,936 extract_columns.py _sql_to_df():77 df.shape: (9892701, 8) - read_sql_query() in 91.3 sec
INFO 2021-02-24 21:40:13,938 extract_columns.py _sql_to_df():78 df.memory_usage().sum(): 633132992
INFO 2021-02-24 21:40:30,517 extract_columns.py pivot_table_to_parquet():172 Writing data to file lrauv.parquet...
INFO 2021-02-24 21:40:31,798 extract_columns.py pivot_table_to_parquet():176 dfp.shape: (744662, 25) - to_parquet() in 1.3 sec
INFO 2021-02-24 21:40:31,798 extract_columns.py pivot_table_to_parquet():177 Done
stoqs container peak memory usage: 7.9 GB

(base) ➜  docker git:(master) ✗ docker-compose exec stoqs stoqs/contrib/parquet/extract_columns.py --db stoqs_canon_october2020 -o all_plats.parquet -v
INFO 2021-02-24 21:41:53,151 extract_columns.py _estimate_memory():161 Estimated required_memory = 896823624.96
INFO 2021-02-24 21:41:53,153 extract_columns.py _sql_to_df():65 Reading from SQL query into DataFrame...
INFO 2021-02-24 21:43:56,723 extract_columns.py _sql_to_df():77 df.shape: (13738107, 8) - read_sql_query() in 123.6 sec
INFO 2021-02-24 21:43:56,725 extract_columns.py _sql_to_df():78 df.memory_usage().sum(): 879238976
INFO 2021-02-24 21:44:20,578 extract_columns.py pivot_table_to_parquet():172 Writing data to file all_plats.parquet...
INFO 2021-02-24 21:44:23,349 extract_columns.py pivot_table_to_parquet():176 dfp.shape: (1123909, 61) - to_parquet() in 2.8 sec
INFO 2021-02-24 21:44:23,349 extract_columns.py pivot_table_to_parquet():177 Done
stoqs container peak memory usage: 11 GB

A regression of estimated df size to container memory usage gives a factor of 12.3

Mike McCann
MBARI 29 January 2021
"""

import os
import sys

# Insert Django App directory (parent of config) into python path
sys.path.insert(0, os.path.abspath(os.path.join(
    os.path.dirname(__file__), "../../")))
if 'DJANGO_SETTINGS_MODULE' not in os.environ:
    os.environ['DJANGO_SETTINGS_MODULE'] = 'config.settings.local'
# django >=1.7
try:
    import django

    django.setup()
except AttributeError:
    pass

import argparse
import logging
import pandas as pd
from django.db import connections
from stoqs.models import Platform
from time import time

class Columnar():

    logger = logging.getLogger(__name__)
    _handler = logging.StreamHandler()
    _formatter = logging.Formatter('%(levelname)s %(asctime)s %(filename)s '
                                   '%(funcName)s():%(lineno)d %(message)s')
    _handler.setFormatter(_formatter)
    _log_levels = (logging.WARN, logging.INFO, logging.DEBUG)
    logger.addHandler(_handler)

    # Set to GB of RAM that have been resourced to the Docker engine
    MAX_CONTAINER_MEMORY = 16
    DF_TO_RAM_FACTOR = 12.3

    def _set_platforms(self):
        '''Set plats and plat_list member variables
        '''
        platforms = (self.args.platforms or
                     Platform.objects.using(self.args.db).all()
                     .values_list('name', flat=True).order_by('name'))
        self.logger.debug(platforms)
        self.plats = ''
        self.plat_list = []
        for platform in platforms:
            if platform in self.args.platforms_omit:
                # Omit some platforms for shorter execution times
                continue
            self.plats += f"'{platform}',"
            self.plat_list.append(platform)
        self.plats = self.plats[:-2] + "'"

    def _sql_to_df(self, sql, extract=False):
        if extract:
            self.logger.info('Reading from SQL query into DataFrame...')

        # More than 10 GB of RAM is needed in Docker Desktop for reading data
        # from stoqs_canon_october2020. The chunksize option in read_sql_query()
        # does not help reduce the server side memory usage.
        # See: https://stackoverflow.com/a/31843091/1281657
        #      https://github.com/pandas-dev/pandas/issues/12265#issuecomment-181809005
        #      https://github.com/pandas-dev/pandas/issues/35689
        stime = time()
        df = pd.read_sql_query(sql, connections[self.args.db])
        etime = time() - stime
        if extract:
            self.logger.info(f"df.shape: {df.shape} <- read_sql_query() in {etime:.1f} sec")
            self.logger.info(f"Actual df.memory_usage().sum():"
                             f" {(df.memory_usage().sum()/1.e9):.3f} GB")
            self.logger.debug(f"Head of original df:\n{df.head()}")

        return df

    def _build_sql(self, limit=None, order=True, count=False):
        self._set_platforms()

        # Base query that's similar to the one behind the api/measuredparameter.csv request
        sql = f'''\nFROM public.stoqs_measuredparameter
        INNER JOIN stoqs_measurement ON (stoqs_measuredparameter.measurement_id = stoqs_measurement.id)
        INNER JOIN stoqs_instantpoint ON (stoqs_measurement.instantpoint_id = stoqs_instantpoint.id)
        INNER JOIN stoqs_activity ON (stoqs_instantpoint.activity_id = stoqs_activity.id)
        INNER JOIN stoqs_platform ON (stoqs_activity.platform_id = stoqs_platform.id)
        INNER JOIN stoqs_parameter ON (stoqs_measuredparameter.parameter_id = stoqs_parameter.id)
        WHERE stoqs_platform.name IN ({self.plats})
          AND stoqs_parameter.{self.args.collect} is not null'''

        if count:
            sql = 'SELECT count(*) ' + sql
        else:
            sql = f'''SELECT stoqs_platform.name as platform,
                stoqs_instantpoint.timevalue, stoqs_measurement.depth,
                ST_X(stoqs_measurement.geom) as longitude,
                ST_Y(stoqs_measurement.geom) as latitude,
                stoqs_parameter.{self.args.collect},
                stoqs_measuredparameter.datavalue {sql}'''
            if order:
                sql += ('\nORDER BY stoqs_platform.name, stoqs_instantpoint.timevalue,'
                            ' stoqs_measurement.depth, stoqs_parameter.name')

        if limit:
            sql += f"\nLIMIT {limit}"
        self.logger.debug(f'sql = {sql}')

        return sql

    def _estimate_memory(self):
        '''Perform a small query on the selection and extrapolate
        to estimate the server-side memory required for the full extraction.
        '''
        SAMPLE_SIZE = 100
        sql = self._build_sql(limit=SAMPLE_SIZE, order=False)
        df = self._sql_to_df(sql)

        sample_memory = df.memory_usage().sum()
        self.logger.debug(f"{sample_memory} B for {SAMPLE_SIZE} records")

        total_recs = self._sql_to_df(self._build_sql(count=True))['count'][0]
        self.logger.debug(f"total_recs = {total_recs}")

        required_memory = total_recs * sample_memory / SAMPLE_SIZE / 1.e9
        container_memory = self.DF_TO_RAM_FACTOR * required_memory
        self.logger.info(f"Estimated required_memory:"
                         f" {required_memory:.3f} GB for DataFrame,"
                         f" {container_memory:.3f} GB for container RAM,")

        if container_memory > self.MAX_CONTAINER_MEMORY:
            self.logger.exception(f"Request of {container_memory:.3f} GB would"
                                  f" exceed {self.MAX_CONTAINER_MEMORY} GB"
                                  f" of RAM available")
            sys.exit(-1)

    def pivot_table_to_parquet(self):
        '''Approach 4. Use Pandas do a pivot on data read into a DataFrame
        '''
        self._estimate_memory()
        sql = self._build_sql()
        df = self._sql_to_df(sql, extract=True)
        context = ['platform', 'timevalue', 'depth', 'latitude', 'longitude']
        dfp = df.pivot_table(index=context, columns=self.args.collect, values='datavalue')
        self.logger.debug(dfp.shape)

        self.logger.info(f'Writing data to file {self.args.output}...')
        stime = time()
        dfp.to_parquet(self.args.output)
        etime = time() - stime
        self.logger.info(f"dfp.shape: {dfp.shape} -> to_parquet() in {etime:.1f} sec")
        self.logger.debug(f"Head of pivoted df:\n{dfp.head()}")
        self.logger.info('Done')

    def process_command_line(self):
        parser = argparse.ArgumentParser(description='Transform STOQS data into columnar Parquet file format')

        parser.add_argument('--platforms', action='store', nargs='*',
                            help='Restrict to just these platforms')
        parser.add_argument('--platforms_omit', action='store', nargs='*', default=[],
                            help='Restrict to all but these platforms')
        parser.add_argument('--collect', action='store', default='name',
                            choices=['name', 'standard_name'],
                            help='The column to collect: name or standard_name')
        parser.add_argument('--db', action='store', required=True,
                            help='Database alias, e.g. stoqs_canon_october2020')
        parser.add_argument('-o', '--output', action='store', required=True,
                            help='Output file name')
        parser.add_argument('--start', action='store', help='Start time in YYYYMMDDTHHMMSS format',
                            default='19000101T000000')
        parser.add_argument('--end', action='store', help='End time in YYYYMMDDTHHMMSS format',
                            default='22000101T000000')
        parser.add_argument('-v', '--verbose', type=int, choices=range(3),
                            action='store', default=0, const=1, nargs='?',
                            help="verbosity level: " + ', '.join(
                                [f"{i}: {v}" for i, v, in enumerate(('WARN', 'INFO', 'DEBUG'))]))

        self.args = parser.parse_args()
        self.commandline = ' '.join(sys.argv)
        self.logger.setLevel(self._log_levels[self.args.verbose])
        self.logger.debug(f"Using databases at DATABASE_URL ="
                          f" {os.environ['DATABASE_URL']}")

if __name__ == '__main__':

    c = Columnar()
    c.process_command_line()
    c.pivot_table_to_parquet()


# Stacked Ensemble RfAmp Prediction Model
# Multiple ML regressors are individually trained and then combined via meta-regressor.
# Hyperparameters are tuned via GridSearchCV

# coding: utf-8

from __future__ import division

import optparse

import numpy as np
import pandas as pd
import os

if not os.getenv("DISPLAY", None):
    import matplotlib
    matplotlib.use("agg", warn=False)

import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.cm as cm

from sklearn.model_selection import train_test_split
from sklearn import preprocessing
from sklearn.externals import joblib
from mlxtend.regressor import StackingCVRegressor
from sklearn.datasets import load_boston
from sklearn.linear_model import Lasso
from sklearn.linear_model import Ridge
from sklearn.ensemble import RandomForestRegressor
from sklearn.svm import SVR
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import GridSearchCV

from keras.models import Sequential, model_from_json
from keras.layers import Dense, Dropout, Activation
from keras.optimizers import SGD, Adagrad, Adadelta, RMSprop, Adam
from keras import losses
from keras import callbacks
from keras.utils import plot_model

import pickle

__author__ = "Nikhil Mukund <nikhil.mukund@ligo.org>, Michael Coughlin <michael.coughlin.ligo.org>"
__version__ = 1.0
__date__ = "11/26/2017"

def parse_commandline():
    """@parse the options given on the command-line.
    """
    parser = optparse.OptionParser(usage=__doc__,version=__version__)

    parser.add_option("-f", "--earthquakesFile", help="Seismon earthquakes file.",default ="/home/mcoughlin/Seismon/Predictions/L1O1O2_CMT_GPR/earthquakes.txt")
    parser.add_option("-o", "--outputDirectory", help="output folder.",default ="/home/mcoughlin/Seismon/MLA/L1O1O2/")
    parser.add_option("-r", "--runType", help="run type (original, lowlatency, cmt)", default ="lowlatency")
    parser.add_option("-m", "--minMagnitude", help="Minimum earthquake magnitude.", default=5.0,type=float)
    parser.add_option("-N", "--Nepoch", help="number of epochs", default =10, type=int)
    parser.add_option("--doPlots",  action="store_true", default=False)
    parser.add_option("-v", "--verbose", action="store_true", default=False,
                      help="Run verbosely. (Default: False)")

    opts, args = parser.parse_args()

    # show parameters
    if opts.verbose:
        print >> sys.stderr, ""
        print >> sys.stderr, "running network_eqmon..."
        print >> sys.stderr, "version: %s"%__version__
        print >> sys.stderr, ""
        print >> sys.stderr, "***************** PARAMETERS ********************"
        for o in opts.__dict__.items():
          print >> sys.stderr, o[0]+":"
          print >> sys.stderr, o[1]
        print >> sys.stderr, ""

    return opts

'''
0: earthquake gps time
1: earthquake mag
2: p gps time
3: s gps time
4: r (2 km/s)
5: r (3.5 km/s)
6: r (5 km/s)
7: predicted ground motion (m/s)
8: lower bounding time
9: upper bounding time
10: latitude
11: longitude
12: distance
13: depth (m)
14: azimuth (deg)
15: nodalPlane1_strike
16: nodalPlane1_rake
17: nodalPlane1_dip
18: momentTensor_Mrt
19: momentTensor_Mtp
20: momentTensor_Mrp
21: momentTensor_Mtt
22: momentTensor_Mrr
23: momentTensor_Mpp
24: peak ground velocity gps time
25: peak ground velocity (m/s)
26: peak ground acceleration gps time
27: peak ground acceleration (m/s^2)
28: peak ground displacement gps time
29: peak ground displacement (m)
30: Lockloss time
31: Detector Status
'''

# Parse command line
opts = parse_commandline()

outputDirectory = os.path.join(opts.outputDirectory,opts.runType)
if not os.path.isdir(outputDirectory):
    os.makedirs(outputDirectory)

data = pd.read_csv(opts.earthquakesFile,delimiter=' ',header=None)
neqs, ncols = data.shape
if ncols == 32:
    fileType = "seismon"
elif ncols == 27:
    fileType = "usarray"
    data = data.drop(data.columns[[24]], 1)
    data = data.rename(index=int, columns={25: 24, 26: 25})
else:
    print("I do not understand the file type...")
    exit(0)

# find magnitudes greater than minimum magnitude
index = data[1] > opts.minMagnitude
data = data[:][index]

# find depth = 0
index = np.where(data[[13]] == 0)[0]
data.iloc[index,13] = 1.0

# shuffle data
data = data.reindex(np.random.permutation(data.index))

Mag_idx = 1
Dist_idx = 12
Depth_idx = 13
Rf_Amp_idx = 25

# Mag threshold
Rf_Amp_thresh = 1e-8;
index = data[Rf_Amp_idx] > Rf_Amp_thresh
data = data[:][index]

if opts.runType == "cmt":
    # Select features
    FeatSet_index = [1,10,11,12,13,14,15,16,17,18,19,20,21,22,23]
elif opts.runType == "lowlatency":
    #FeatSet_index = [1,7,10,11,12,13,14,15,16,17] #  these lower set paramaters makes  sense
    FeatSet_index = [1,10,11,12,13,14] #  these lower set paramaters makes  sense
elif opts.runType == "original":
    FeatSet_index = [1,12,13] #  Just Mag, Dist, Depth
else:
    print("--runType must be original, lowlatency, and cmt")
    exit(0)

Target_index = [Rf_Amp_idx]

# Artificially increase samples
data_temp = data
copy_num =  6
noise_level = 1e-2 # 1e-2

Rfamp_orig = data_temp[Target_index];
data_orig =  data_temp

def boost_samples(x_samples,y_samples,copy_num=3,noise_level=1e-2):
    # Artificially increase samples
    data_x_temp = x_samples
    data_y_temp = y_samples

    for i in range(copy_num):
        data_x_temp = np.vstack((data_x_temp,data_x_temp))
        data_y_temp = np.vstack((data_y_temp,data_y_temp))

    data_x_orig =  data_x_temp
    data_y_orig =  data_y_temp

    x1 = data_x_temp
    x2 = np.random.randn(*data_x_temp.shape)*noise_level
    x_samples_boosted = x1 + np.multiply(x1,x2)

    y1 = data_y_temp
    y2 = np.random.randn(*data_y_temp.shape)*noise_level
    y_samples_boosted = y1 + np.multiply(y1,y2)

    # Shuffle samples
    #IDX = np.random.permutation(y_samples_boosted.index)
    IDX = np.random.permutation(np.arange(0,len(y_samples_boosted)))
    x_samples_boosted = x_samples_boosted[IDX,:]
    y_samples_boosted = y_samples_boosted[IDX,:]

    return x_samples_boosted, y_samples_boosted

data = data_temp

# Take Log10 of certain features (Mag, Dist, Depth)
data[[Dist_idx, Depth_idx]] = np.log10(data[[Dist_idx, Depth_idx]])
data[Target_index] = np.log10(data[Target_index])

X = np.asarray(data[FeatSet_index])
Y = np.asarray(data[Target_index])

# Normalize samples
x_scaler = preprocessing.MinMaxScaler()
#x_scaler = preprocessing.data.QuantileTransformer()
X = x_scaler.fit_transform(X)
y_scaler = preprocessing.MinMaxScaler()
#y_scaler = preprocessing.data.QuantileTransformer()
Y = y_scaler.fit_transform(Y)

x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size=0.3,random_state=42)
x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.3,random_state=42)

# boost_samples + normalize + shuffle them
TUPLE1 = boost_samples(x_train,y_train,copy_num,noise_level)
TUPLE2 = boost_samples(x_val,y_val,copy_num,noise_level)
TUPLE3 = boost_samples(x_test,y_test,copy_num,noise_level)

x_train = TUPLE1[0]
y_train = TUPLE1[1]

x_val   = TUPLE2[0]
y_val   = TUPLE2[1]

x_test  = TUPLE3[0]
y_test  = TUPLE3[1]

#############################################
# Construct Stacked Ensemble Model #
#############################################

RANDOM_SEED = 42

ridge = Ridge()
lasso = Lasso()
svr_lin = SVR(kernel='linear')
svr_rbf = SVR(kernel='rbf')
lr = LinearRegression()
rf = RandomForestRegressor(random_state=RANDOM_SEED)

np.random.seed(RANDOM_SEED)

regressors = [svr_lin, svr_rbf, lr, ridge, lasso]

stack = StackingCVRegressor(regressors=regressors,
                            meta_regressor=rf,
                            use_features_in_secondary=True)

'''params = {'lasso__alpha': [0.1, 1.0, 10.0],
          'ridge__alpha': [0.1, 1.0, 10.0],
          'svr__C': [0.1, 1.0, 10.0],
          'meta-svr__C': [0.1, 1.0, 10.0, 100.0],
          'meta-svr__gamma': [0.1, 1.0, 10.0]}

params = {'lasso__alpha': [0.1, 1.0, 10.0],
          'ridge__alpha': [0.1, 1.0, 10.0]}'''

model = GridSearchCV(
    estimator=stack,
    param_grid={
        'lasso__alpha': [x/5.0 for x in range(1, 10)],
        'ridge__alpha': [x/20.0 for x in range(1, 10)],
        'meta-randomforestregressor__n_estimators': [10, 100]
    },
    cv=5,
    refit=True,
    verbose=10,
    n_jobs=8,
)

###################################################

model.fit(x_train, y_train.ravel())

print("Best: %f using %s" % (grid.best_score_, grid.best_params_))

###################################################
y_pred = model.predict(x_test)
y_pred = np.expand_dims(y_pred,axis=1)

# Rescale Back
y_pred = 10**y_scaler.inverse_transform(y_pred)
y_test = 10**y_scaler.inverse_transform(y_test)

# Reject test samples below certain threshold
Rf_thresh = 0.5*1e-7 # 0.5*1e-6
ijk = y_test > Rf_thresh
y_test = y_test[ijk]
y_pred = y_pred[ijk]
x_test = x_test[ijk.flatten(),:]

# Add bias
#y_pred = y_pred + 0.1*y_pred

# sort results in Ascending order
y_test_sort = np.sort(y_test,axis=0)
y_pred_sort = y_pred[np.argsort(y_test,axis=0)]

## Percentage within the specified factor
Fac = 2
IDX = y_pred_sort/(y_test_sort+np.finfo(float).eps) >= 1
K = y_pred_sort[IDX]
Q = y_test_sort[IDX]
L = y_pred_sort[~IDX]
M = y_test_sort[~IDX]
Upper_indices = [i for i, x in enumerate(K <= Fac*Q) if x == True]
Lower_indices = [i for i, x in enumerate(L >= M/Fac) if x == True]
Percent_within_Fac = (len(Upper_indices) + len(Lower_indices))/len(y_pred)*100
print("Percentage captured within a factor of {} = {:.2f}".format(Fac,Percent_within_Fac))

Diff = abs(y_pred_sort - y_test_sort)

# Errorbar values
yerr_lower = y_test_sort - y_test_sort/Fac
yerr_upper = Fac*y_test_sort - y_test_sort

idx = np.arange(0,len(y_test_sort))

if opts.doPlots:
    font = {'weight' : 'bold',
        'size'   : 15}
    plt.rc('font', **font)
    plt.rc('legend',**{'fontsize':15})

    plt.figure(figsize=(10,8))
    plt.style.use('dark_background')
    #plt.style.use('ggplot')

    diff_plt = plt.scatter(idx,Diff,color='lightgreen',alpha=0.1)

    errorbar_plt = plt.errorbar(idx,y_test_sort,yerr=[yerr_lower,yerr_upper], alpha=0.05 ,color='lightgrey')
    actual_plt = plt.scatter(idx,y_test_sort,color='#1f77b4',alpha=0.9)

    idx2 = np.arange(0,len(y_pred_sort))
    pred_plt = plt.scatter(idx2,y_pred_sort,color='#d62728',alpha=0.2)

    plt.yscale('log')
    plt.grid()
    plt.ylim([1e-7, 1e-3])
    #plt.ylim([0, 1])
    #plt.ylabel('Rf Amplitude (m/s) \n (Normalized to 1)',fontsize=25)
    plt.ylabel('Rf Amplitude (m/s) ',fontsize=25)
    plt.xlabel('Samples',fontsize=25)
    plt.title("Percentage captured within a factor of {} = {:.2f}".format(Fac,Percent_within_Fac))
    legend_plt = plt.legend([pred_plt,actual_plt, diff_plt],['Prediction', 'Actual', 'Difference'],loc=2,markerscale=2., scatterpoints=100)

    plt.autoscale(enable=True, axis='x', tight=True)
    plt.grid(linestyle=':')

    plt.savefig(os.path.join(outputDirectory,'performance.pdf'),bbox_inches='tight')
    plt.close()

# Save Model
# serialize model & pickle
pickle.dump(model, open("%s/model.p"%outputDirectory, "wb"))
print("Saved model to disk")

'''
# Load Saved Model
# load pickle
pickle_file = open('%s/model.p'%outputDirectory, 'rb')
loaded_model_pickle = pickle.load(pickle_file)
print("Loaded model from disk")
'''


#!/usr/bin/env python

"""
Python class wrapper for data fitting.
Includes the following external methods:
getFunctions returns the list of function names (dictionary keys)
FitRegion performs the fitting
Note that FitRegion will plot on top of the current data using MPlots routines
if the current curve and the current plot instance are passed.

"""
# January, 2009
# Paul B. Manis, Ph.D.
# UNC Chapel Hill
# Department of Otolaryngology/Head and Neck Surgery
# Supported by NIH Grants DC000425-22 and DC004551-07 to PBM.
# Copyright Paul Manis, 2009
#
"""
    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/>.
"""
"""
    Additional Terms:
    The author(s) would appreciate that any modifications to this program, or
    corrections of erros, be reported to the principal author, Paul Manis, at
    pmanis@med.unc.edu, with the subject line "PySounds Modifications".

    Note: This program also relies on the TrollTech Qt libraries for the GUI.
    You must obtain these libraries from TrollTech directly, under their license
    to use the program.
"""

import sys
import numpy
import scipy

try:
    import openopt
    HAVE_OPENOPT = True
except ImportError:
    HAVE_OPENOPT = False
    print "There was an error importing openopt. Continuing...."

import ctypes
import numpy.random

#from numba import autojit

usingMPlot = False
if usingMPlot:
    import MPlot # we include plotting as part of the fitting

def debug_trace():
  '''Set a tracepoint in the Python debugger that works with Qt'''
  if pyqt:
      from PyQt4.QtCore import pyqtRemoveInputHook
  from pdb import set_trace
  if pyqt:
      pyqtRemoveInputHook()
  set_trace()

class Fitting():
    # dictionary contains:
    # name of function: function call, initial parameters, iterations, plot color, then x and y for testing
    # target valutes, names of parameters, contant values, and derivative function if needed.
    #
    def __init__(self):
        self.fitfuncmap = {
        'exp0'  : (self.exp0eval, [0.0, 20.0], 2000, 'k', [0, 100, 1.],
                   [1.0, 5.0], ['A0', 'tau'], None, None),
        'exp1'  : (self.expeval, [0.0, 0.0, 20.0], 2000, 'k', [0, 100, 1.],
                   [0.5, 1.0, 5.0], ['DC', 'A0', 'tau'], None, self.expevalprime),
        'expsum'  : (self.expsumeval,  [0.0, -0.5, 200.0, -0.25, 450.0], 500000, 'k',  [0, 1000, 1.],
                   [0.0, -1.0, 150.0, -0.25, 350.0], ['DC', 'A0', 'tau0', 'A1', 'tau1'], None, None),
        'expsum2'  : (self.expsumeval2,  [0., -0.5, -0.250], 50000, 'k',  [0, 1000, 1.],
                   [0., -0.5, -0.25], ['A0', 'A1'], [5., 20.], None),
        'exp2'  : (self.exp2eval,  [0.0, -0.5, 200.0, -0.25, 450.0], 500000, 'k',  [0, 1000, 1.],
                   [0.0, -1.0, 150.0, -0.25, 350.0], ['DC', 'A0', 'tau0', 'A1', 'tau1'], None, None),
        'exppow'  : (self.exppoweval,  [0.0, 1.0, 100, ], 2000, 'k',  [0, 100, 0.1],
                   [0.0, 1.0, 100.0], ['DC', 'A0', 'tau'], None, None),
        'exppulse'  : (self.expPulse,  [3.0, 2.5, 0.2, 2.5, 2.0, 0.5], 2000, 'k',  [0, 10, 0.3],
                  [0.0, 0., 0.75, 4., 1.5, 1.], ['DC', 't0', 'tau1', 'tau2', 'amp', 'width'], None, None),
        'boltz' : (self.boltzeval,  [0.0, 1.0, -50.0, -5.0], 5000, 'r', [-130., -30., 1.],
                   [0.00, 0.010, -100.0, 7.0],  ['DC', 'A0', 'x0', 'k'], None, None),

        'gauss' : (self.gausseval,  [1.0, 0.0, 0.5], 2000, 'y',  [-10., 10., 0.2],
                   [1.0, 1.0, 2.0], ['A', 'mu', 'sigma'], None, None),
        'line'  : (self.lineeval,  [1.0, 0.0], 500, 'r', [-10., 10., 0.5],
                   [0.0, 2.0], ['m', 'b'], None, None),
        'poly2' : (self.poly2eval, [1.0, 1.0, 0.0], 500, 'r', [0, 100, 1.],
                    [0.5, 1.0, 5.0], ['a', 'b', 'c'], None, None),
        'poly3' : (self.poly3eval, [1.0, 1.0, 1.0, 0.0], 1000, 'r', [0., 100., 1.],
                   [0.5, 1.0, 5.0, 2.0], ['a', 'b', 'c', 'd'], None, None),
        'poly4' : (self.poly4eval, [1.0, 1.0, 1.0, 1.0, 0.0], 1000, 'r', [0., 100., 1.],
                   [0.1, 0.5, 1.0, 5.0, 2.0], ['a', 'b', 'c', 'd', 'e'], None, None),
        'sin'   : (self.sineeval, [-1., 1.0, 4.0, 0.0], 1000, 'r', [0., 100., 0.2],
                   [0.0, 1.0, 9.0, 0.0], ['DC', 'A', 'f', 'phi'], None, None),
        'boltz2' : (self.boltzeval2,  [0.0, 0.5, -50.0, 5.0, 0.5, -20.0, 3.0], 1200, 'r',
                    [-100., 50., 1.], [0.0, 0.3, -45.0, 4.0, 0.7, 10.0, 12.0],
                    ['DC', 'A1', 'x1', 'k1', 'A2', 'x2', 'k2'], None, None),
        'taucurve' : (self.taucurve,  [50., 300.0, 60.0, 10.0, 8.0, 65.0, 10.0], 50000, 'r',
                    [-150., 50., 1.], [0.0, 237.0, 60.0, 12.0, 17.0, 60.0, 14.0],
                    ['DC', 'a1', 'v1', 'k1', 'a2', 'v2', 'k2'],  None, self.taucurveder),
        }
        self.fitSum2Err = 0

    def getFunctions(self):
        return(self.fitfuncmap.keys())

    def exp0eval(self, p, x, y=None, C = None, sumsq = False):
        """
        Exponential function with an amplitude and 0 offset
        """
        yd = p[0] * numpy.exp(-x/p[1])
        if y is None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def expsumeval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        """
        Sum of two exponentials with independent time constants and amplitudes,
        and a DC offset
        """
        yd = p[0] + (p[1]* numpy.exp(-x/p[2])) + (p[3]*numpy.exp(-x/p[4]))
        if y is None:
            return yd
        else:
            yerr = y - yd
            if weights is not None:
                yerr = yerr * weights
            if sumsq is True:
                return numpy.sum(yerr**2)
            else:
                return yerr

    def expsumeval2(self, p, x, y=None, C = None, sumsq = False, weights=None):
        """
        Sum of two exponentials, with predefined time constants , allowing
        only the amplitudes and DC offset to vary
        """
        yd = p[0] + (p[1]* numpy.exp(-x/C[0])) + (p[2]*numpy.exp(-x/C[1]))
        if y is None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def expeval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        """
        Exponential with offset
        """
        yd = p[0] + p[1] * numpy.exp(-x/p[2])
#        print yd.shape
#        print y.shape
        if y is None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def expevalprime(self, p, x, y=None, C = None, sumsq = False, weights=None):
        """
        Derivative for exponential with offset
        """
        ydp = p[1] * numpy.exp(-x/p[2])/(p[2]*p[2])
        yd = p[0] + p[1] * numpy.exp(-x/p[2])
        print y
        if y is None:
            return (yd, ydp)
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def exppoweval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        """
        Single exponential function, rising to a ppower
        """

        if C is None:
            cx = 1.0
        else:
            cx = C[0]
        yd = p[0] + p[1] * (1.0-numpy.exp(-x/p[2]))**cx
        if y is None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def exp2eval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        """
        For fit to activation currents...
        """
        yd = p[0] + (p[1] * (1.0 - numpy.exp(-x/p[2]))**2.0 ) + (p[3] * (1.0 - numpy.exp(-x/p[4])))
        if y == None:
            return yd
        else:
            if sumsq is True:
                ss = numpy.sqrt(numpy.sum((y - yd)**2.0))
#                if p[4] < 3.0*p[2]:
#                    ss = ss*1e6 # penalize them being too close
                return ss
            else:
                return y - yd

#    @autojit
    def expPulse(self, p, x, y=None, C=None, sumsq = False, weights = None):
        """Exponential pulse function (rising exponential with optional variable-length
        plateau followed by falling exponential)
        Parameter p is [yOffset, t0, tau1, tau2, amp, width]
        """
        yOffset, t0, tau1, tau2, amp, width = p
        yd = numpy.empty(x.shape)
        yd[x<t0] = yOffset
        m1 = (x>=t0)&(x<(t0+width))
        m2 = (x>=(t0+width))
        x1 = x[m1]
        x2 = x[m2]
        yd[m1] = amp*(1-numpy.exp(-(x1-t0)/tau1))+yOffset
        amp2 = amp*(1-numpy.exp(-width/tau1)) ## y-value at start of decay
        yd[m2] = ((amp2)*numpy.exp(-(x2-(width+t0))/tau2))+yOffset
        if y == None:
            return yd
        else:
            if sumsq is True:
                ss = numpy.sqrt(numpy.sum((y-yd)**2.0))
                return ss
            else:
                return y-yd

    def boltzeval(self,p, x, y=None, C = None, sumsq = False, weights=None):
        yd = p[0] + (p[1]-p[0])/(1.0 + numpy.exp((x-p[2])/p[3]))
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sqrt(numpy.sum((y - yd)**2.0))
            else:
                return y - yd

    def boltzeval2(self,p, x, y=None, C = None, sumsq = False, weights=None):
        yd = p[0] + p[1]/(1 + numpy.exp((x-p[2])/p[3])) + p[4]/(1 + numpy.exp((x-p[5])/p[6]))
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def gausseval(self,p, x, y=None, C = None, sumsq = False, weights=None):
        yd = (p[0]/(p[2]*numpy.sqrt(2.0*numpy.pi)))*numpy.exp(-((x - p[1])**2.0)/(2.0*(p[2]**2.0)))
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def lineeval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        yd = p[0]*x + p[1]
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def poly2eval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        yd = p[0]*x**2.0 + p[1]*x + p[2]
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def poly3eval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        yd = p[0]*x**3.0 + p[1]*x**2.0 + p[2]*x +p[3]
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def poly4eval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        yd = p[0]*x**4.0 + p[1]*x**3.0 + p[2]*x**2.0 + p[3]*x +p[4]
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def sineeval(self, p, x, y=None, C = None, sumsq = False, weights=None):
        yd =  p[0] + p[1]*numpy.sin((x*2.0*numpy.pi/p[2])+p[3])
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sum((y - yd)**2)
            else:
                return y - yd

    def taucurve(self, p, x, y=None, C = None, sumsq=True, weights=None):
        """
        HH-like description of activation/inactivation function
        'DC', 'a1', 'v1', 'k1', 'a2', 'v2', 'k2'
        """
        yd = p[0] + 1.0/(p[1]*numpy.exp((x+p[2])/p[3]) +p[4]*numpy.exp(-(x+p[5])/p[6]))
        if y == None:
            return yd
        else:
            if sumsq is True:
                return numpy.sqrt(numpy.sum((y - yd)**2))
            else:
                return y - yd

    def taucurveder(self, p, x):
        """
        Derivative for taucurve
        'DC', 'a1', 'v1', 'k1', 'a2', 'v2', 'k2'
        """
        y = -(p[1]*numpy.exp((p[2] + x)/p[3])/p[3] - p[4]*numpy.exp(-(p[5] + x)/p[6])/p[6])/(p[1]*numpy.exp((p[2] + x)/p[3]) +
p[4]*numpy.exp(-(p[5] + x)/p[6]))**2.0
      #  print 'dy: ', y
        return y

    def getClipData(self, x, y, t0, t1):
        """
        Return the values in y that match the x range in tx from
        t0 to t1. x must be monotonic increasing or decreasing.
        Allow for reverse ordering. """
        it0 = (numpy.abs(x-t0)).argmin()
        it1 = (numpy.abs(x-t1)).argmin()
        if it0 > it1:
            t = it1
            it1 = it0
            it0 = t
        return(x[it0:it1], y[it0:it1])

    def FitRegion(self, whichdata, thisaxis, tdat, ydat, t0 = None, t1 = None,
                  fitFunc = 'exp1', fitFuncDer = None, fitPars = None, fixedPars = None,
                  fitPlot = None, plotInstance = None, dataType= 'xy', method = None,
                  bounds=None, weights=None, constraints=()):
        """
        **Arguments**
        ============= ===================================================
        whichdata
        thisaxis
        tdat
        ydat
        t0            (optional) Minimum of time data - determined from tdat if left unspecified
        t1            (optional) Maximum of time data - determined from tdat if left unspecified
        fitFunc       (optional) The function to fit the data to (as defined in __init__). Default is 'exp1'.
        fitFuncDer    (optional) default=None
        fitPars       (optional) Initial fit parameters. Use the values defined in self.fitfuncmap if unspecified.
        fixedPars     (optional) Fixed parameters to pass to the function. Default=None
        fitPlot       (optional) default=None
        plotInstance  (optional) default=None
        dataType      (optional) Options are ['xy', 'blocks']. Default='xy'
        method        (optional) Options are ['curve_fit', 'fmin', 'simplex', 'Nelder-Mead', 'bfgs', 'TNC', 'SLSQP', 'COBYLA', 'L-BFGS-B', 'openopt']. Default='leastsq'
        bounds        (optional) default=None
        weights       (optional) default=None
        constraints   (optional) default=()
        ============= ===================================================

        To call with tdat and ydat as simple arrays:
        FitRegion(1, 0, tdat, ydat, FitFunc = 'exp1')
        e.g., the first argument should be 1, but this axis is ignored if datatype is 'xy'
        """
        self.fitSum2Err = 0.0
        if t0 == t1:
            if plotInstance is not None and usingMPlot:
                (x, y) = plotInstance.getCoordinates()
                t0 = x[0]
                t1 = x[1]
        if t1 is None:
            t1 = numpy.max(tdat)
        if t0 is None:
            t0 = numpy.min(tdat)
        func = self.fitfuncmap[fitFunc]
        if func is None:
            print "FitRegion: unknown function %s" % (fitFunc)
            return
        xp = []
        xf = []
        yf = []
        yn = []
        tx = []
        names = func[6]
        if fitPars is None:
            fpars = func[1]
        else:
            fpars = fitPars
        if method == 'simplex': # remap calls if needed for newer versions of scipy (>= 0.11)
            method = 'Nelder-Mead'
        if ydat.ndim == 1 or dataType == 'xy' or dataType == '2d': # check if 1-d, then "pretend" its only a 1-element block
            nblock = 1
        else:
            nblock = ydat.shape[0] # otherwise, this is the number of traces in the block
        # print 'datatype: ', dataType
        # print 'nblock: ', nblock
        # print 'whichdata: ', whichdata
        # for block in range(nblock):
            for record in whichdata:
                if dataType == 'blocks':
                    (tx, dy) = self.getClipData(tdat[block], ydat[block][record, thisaxis, :], t0, t1)
                else:
                    (tx, dy) = self.getClipData(tdat, ydat[record,:], t0, t1)
                # print 'Fitting.py: block, type, Fit data: ', block, dataType
                # print tx.shape
                # print dy.shape
                yn.append(names)
                if not any(tx):
                    continue # no data in the window...
                ier = 0
                #
                # Different optimization methods are included here. Not all have been tested fully with
                # this wrapper.
                #
                if method is None or method == 'leastsq': # use standard leastsq, no bounds
                        plsq, cov, infodict, mesg, ier = scipy.optimize.leastsq(func[0], fpars,
                                                args=(tx.astype('float64'), dy.astype('float64'), fixedPars),
                                                full_output = 1, maxfev = func[2])
                        if ier > 4:
                            print "optimize.leastsq error flag is: %d" % (ier)
                            print mesg
                elif method == 'curve_fit':
                    print fpars
                    print fixedPars
                    plsq, cov = scipy.optimize.curve_fit(func[0], tx.astype('float64'), dy.astype('float64'), p0=fpars)
                    ier = 0
                elif method in ['fmin', 'simplex', 'Nelder-Mead', 'bfgs', 'TNC', 'SLSQP', 'COBYLA', 'L-BFGS-B']: # use standard wrapper from scipy for those routintes
                    res = scipy.optimize.minimize(func[0], fpars, args=(tx.astype('float64'), dy.astype('float64'), fixedPars, True),
                     method=method, jac=None, hess=None, hessp=None, bounds=bounds, constraints=constraints, tol=None, callback=None,
                     options={'maxiter': func[2], 'disp': False })
                    plsq = res.x
                    #print "    method:", method
                    #print "    bounds:", bounds
                    #print "    result:", plsq

                # next section is replaced by the code above - kept here for reference if needed...
                # elif method == 'fmin' or method == 'simplex':
                #     plsq = scipy.optimize.fmin(func[0], fpars, args=(tx.astype('float64'), dy.astype('float64'), fixedPars, True),
                #                 maxfun = func[2]) # , iprint=0)
                #     ier = 0
                # elif method == 'bfgs':
                #     plsq, cov, infodict = scipy.optimize.fmin_l_bfgs_b(func[0], fpars, fprime=func[8],
                #                 args=(tx.astype('float64'), dy.astype('float64'), fixedPars, True, weights),
                #                 maxfun = func[2], bounds = bounds,
                #                 approx_grad = True) # , disp=0, iprint=-1)
                elif method == 'openopt': # use OpenOpt's routines - usually slower, but sometimes they converge better
                    if not HAVE_OPENOPT:
                        raise Exception("Requested openopt fitting method but openopt is not installed.")

                    if bounds is not None:
                        # unpack bounds
                        lb = [y[0] for y in bounds]
                        ub = [y[1] for y in bounds]
                        fopt = openopt.DFP(func[0], fpars, tx, dy, df = fitFuncDer, lb=lb, ub=ub)
                       # fopt.df = func[8]
                        r = fopt.solve('nlp:ralg', plot=0, iprint = 10)
                        plsq = r.xf
                        ier = 0
                    else:
                        fopt = openopt.DFP(func[0], fpars, tx, dy, df = fitFuncDer)
                        print func[8]
                      #  fopt.df = func[7]
                        fopt.checkdf()
                        r = fopt.solve('nlp:ralg', plot=0, iprint = 10)
                        plsq = r.xf
                        ier = 0
                else:
                    print 'method %s not recognized, please check Fitting.py' % (method)
                    return
                xfit = numpy.arange(min(tx), max(tx), (max(tx)-min(tx))/100.0)
                yfit = func[0](plsq, xfit, C=fixedPars)
                yy = func[0](plsq, tx, C=fixedPars) # calculate function
                self.fitSum2Err = numpy.sum((dy - yy)**2)
                if usingMPlot and FitPlot != None and plotInstance != None:
                    self.FitPlot(xFit = xfit, yFit = yfit, fitFunc = fund[0],
                            fitPars = plsq, plot = fitPlot, plotInstance = plotInstance)
                xp.append(plsq) # parameter list
                xf.append(xfit) # x plot point list
                yf.append(yfit) # y fit point list
#        print xp
#        print len(xp)
        return(xp, xf, yf, yn) # includes names with yn and range of tx

    def FitPlot(self, xFit = None, yFit = None, fitFunc = 'exp1',
                fitPars = None, fixedPars = None, fitPlot=None, plotInstance = None,
                color=None):
        """ Plot the fit data onto the fitPlot with the specified "plot Instance".
             if there is no xFit, or some parameters are missing, we just return.
             if there is xFit, but no yFit, then we try to compute the fit with
             what we have. The plot is superimposed on the specified "fitPlot" and
             the color is specified by the function color in the fitPars list.
             """
        if xFit is None or fitPars is None:
            return
        func = self.fitfuncmap[fitFunc]
        if color is None:
            fcolor = func[3]
        else:
            fcolor = color
        if yFit is None:
            yFit = numpy.array([])
            for k in range(0, len(fitPars)):
                yFit[k] = func[0](fitPars[k], xFit[k], C=fixedPars)
        if plotInstance is None or fitPlot is None:
            return(yfit)
        for k in range(0, len(fitPars)):
            plotInstance.PlotLine(fitPlot, xFit[k], yFit[k], color = fcolor)
        return(yfit)

    def getFitErr(self):
        """ Return the fit error for the most recent fit
             """
        return(self.fitSum2Err)

    def expfit(self, x, y):
        """ find best fit of a single exponential function to x and y
        using the chebyshev polynomial approximation.
        returns (DC, A, tau) for fit.

        Perform a single exponential fit to data using Chebyshev polynomial method.
        Equation fit: y = a1 * exp(-x/tau) + a0
        Call: [a0 a1 tau] = expfit(x,y);
        Calling parameter x is the time base, y is the data to be fit.
        Returned values: a0 is the offset, a1 is the amplitude, tau is the time
        constant (scaled in units of x).
        Relies on routines chebftd to generate polynomial coeffs, and chebint to compute the
        coefficients for the integral of the data. These are now included in this
        .py file source.
        This version is based on the one in the pClamp manual: HOWEVER, since
        I use the bounded [-1 1] form for the Chebyshev polynomials, the coefficients are different,
        and the resulting equation for tau is different. I manually optimized the tau
        estimate based on fits to some simulated noisy data. (Its ok to use the whole range of d1 and d0
        when the data is clean, but only the first few coeffs really hold the info when
        the data is noisy.)
        NOTE: The user is responsible for making sure that the passed data is appropriate,
        e.g., no large noise or electronic transients, and that the time constants in the
        data are adequately sampled.
        To do a double exp fit with this method is possible, but more complex.
        It would be computationally simpler to try breaking the data into two regions where
        the fast and slow components are dominant, and fit each separately; then use that to
        seed a non-linear fit (e.g., L-M) algorithm.
        Final working version 4/13/99 Paul B. Manis
        converted to Python 7/9/2009 Paul B. Manis. Seems functional.
        """
        n = 30; # default number of polynomials coeffs to use in fit
        a = numpy.amin(x)
        b = numpy.amax(x)
        d0 = self.chebftd(a, b, n, x, y) # coeffs for data trace...
        d1 = self.chebint(a, b, d0, n) # coeffs of integral...
        tau = -numpy.mean(d1[2:3]/d0[2:3])
        try:
            g = numpy.exp(-x/tau)
        except:
            g = 0.0
        dg = self.chebftd(a, b, n, x, g) # generate chebyshev polynomial for unit exponential function
        # now estimate the amplitude from the ratios of the coeffs.
        a1 = self.estimate(d0, dg, 1)
        a0 = (d0[0]-a1*dg[0])/2.0         # get the offset here
        return(a0, a1, tau)#

    def estimate(self, c, d, m):
        """ compute optimal estimate of parameter from arrays of data """
        n = len(c)
        a = sum(c[m:n]*d[m:n])/sum(d[m:n]**2.0)
        return(a)

    # note : the following routine is a bottleneck. It should be coded in C.

    def chebftd(self, a, b, n, t, d):
        """ Chebyshev fit; from Press et al, p 192.
            matlab code P. Manis 21 Mar 1999
            "Given a function func, lower and upper limits of the interval [a,b], and
            a maximum degree, n, this routine computes the n coefficients c[1..n] such that
            func(x) sum(k=1, n) of ck*Tk(y) - c0/2, where y = (x -0.5*(b+a))/(0.5*(b-a))
            This routine is to be used with moderately large n (30-50) the array of c's is
            subsequently truncated at the smaller value m such that cm and subsequent
            terms are negligible."
            This routine is modified so that we find close points in x (data array) - i.e., we find
            the best Chebyshev terms to describe the data as if it is an arbitrary function.
            t is the x data, d is the y data...
            """
        bma = 0.5*(b-a)
        bpa = 0.5*(b+a)
        inc = t[1]-t[0]
        f = numpy.zeros(n)
        for k in range(0, n):
            y = numpy.cos(numpy.pi*(k+0.5)/n)
            pos = int(0.5+(y*bma+bpa)/inc)
            if  pos < 0:
                pos = 0
            if pos >= len(d)-2:
                pos = len(d)-2
            try:
                f[k]= d[pos+1]
            except:
                print "error in chebftd: k = %d (len f = %d)  pos = %d, len(d) = %d\n" % (k, len(f), pos, len(d))
                print "you should probably make sure this doesn't happen"
        fac = 2.0/n
        c=numpy.zeros(n)
        for j in range(0, n):
           sum=0.0
           for k in range(0, n):
              sum = sum + f[k]*numpy.cos(numpy.pi*j*(k+0.5)/n)
           c[j]=fac*sum
        return(c)

    def chebint(self, a, b, c, n):
        """ Given a, b, and c[1..n] as output from chebft or chebftd, and given n,
        the desired degree of approximation (length of c to be used),
        this routine computes cint, the Chebyshev coefficients of the
        integral of the function whose coeffs are in c. The constant of
        integration is set so that the integral vanishes at a.
        Coded from Press et al, 3/21/99 P. Manis (Matlab)
        Python translation 7/8/2009 P. Manis
        """
        sum = 0.0
        fac = 1.0
        con = 0.25*(b-a) # factor that normalizes the interval
        cint = numpy.zeros(n)
        for j in range(1,n-2):
            cint[j]=con*(c[j-1]-c[j+1])/j
            sum = sum + fac * cint[j]
            fac = - fac
        cint[n-1] = con*c[n-2]/(n-1)
        sum = sum + fac*cint[n-1]
        cint[0] = 2.0*sum # set constant of integration.
        return(cint)

# routine to flatten an array/list.
#
    def flatten(self, l, ltypes=(list, tuple)):
        i = 0
        while i < len(l):
            while isinstance(l[i], ltypes):
                if not l[i]:
                    l.pop(i)
                    if not len(l):
                        break
                else:
                   l[i:i+1] = list(l[i])
            i += 1
        return l
    # flatten()

# run tests if we are "main"

if __name__ == "__main__":
#    import matplotlib.pyplot as pyplot
    import timeit
    import Fitting
    import matplotlib as MP
    MP.use('Qt4Agg')
    ################## Do not modify the following code
    # sets up matplotlib with sans-serif plotting...
    import matplotlib.gridspec as GS
    # import mpl_toolkits.axes_grid1.inset_locator as INSETS
    # #import inset_axes, zoomed_inset_axes
    # import mpl_toolkits.axes_grid1.anchored_artists as ANCHOR
    # # import AnchoredSizeBar

    stdFont = 'Arial'

    import  matplotlib.pyplot as pylab
    pylab.rcParams['text.usetex'] = True
    pylab.rcParams['interactive'] = False
    pylab.rcParams['font.family'] = 'sans-serif'
    pylab.rcParams['font.sans-serif'] = 'Arial'
    pylab.rcParams['mathtext.default'] = 'sf'
    pylab.rcParams['figure.facecolor'] = 'white'
    # next setting allows pdf font to be readable in Adobe Illustrator
    pylab.rcParams['pdf.fonttype'] = 42
    pylab.rcParams['text.dvipnghack'] = True
    ##################### to here (matplotlib stuff - touchy!

    Fits = Fitting.Fitting()
#    x = numpy.arange(0, 100.0, 0.1)
#    y = 5.0-2.5*numpy.exp(-x/5.0)+0.5*numpy.random.randn(len(x))
#    (dc, aFit,tauFit) = Fits.expfit(x,y)
#    yf = dc + aFit*numpy.exp(-x/tauFit)
 #   pyplot.figure(1)
  #  pyplot.plot(x,y,'k')
  #  pyplot.hold(True)
  #  pyplot.plot(x, yf, 'r')
  #  pyplot.show()
    exploreError = False

    if exploreError is True:
        # explore the error surface for a function:

        func = 'exppulse'
        f = Fits.fitfuncmap[func]
        p1range = numpy.arange(0.1, 5.0, 0.1)
        p2range = numpy.arange(0.1, 5.0, 0.1)

        err = numpy.zeros((len(p1range), len(p2range)))
        x = numpy.array(numpy.arange(f[4][0], f[4][1], f[4][2]))
        C = None
        if func == 'expsum2':
          C = f[7]

        # check exchange of tau1 ([1]) and width[4]
        C = None
        yOffset, t0, tau1, tau2, amp, width = f[1] # get inital parameters
        y0 = f[0](f[1], x, C=C)
        noise = numpy.random.random(y0.shape) - 0.5
        y0 += 0.0* noise
        sh = err.shape
        yp = numpy.zeros((sh[0], sh[1], len(y0)))
        for i, p1 in enumerate(p1range):
            tau1t = tau1*p1
            for j, p2 in enumerate(p2range):
                ampt = amp*p2
                pars = (yOffset, t0, tau1t, tau2, ampt, width) # repackage
                err[i,j] = f[0](pars, x, y0, C=C, sumsq = True)
                yp[i,j] = f[0](pars, x, C=C, sumsq = False)

        pylab.figure()
        CS=pylab.contour(p1range*tau1, p2range*width, err, 25)
        CB = pylab.colorbar(CS, shrink=0.8, extend='both')
        pylab.figure()
        for i, p1 in enumerate(p1range):
            for j, p2 in enumerate(p2range):
                pylab.plot(x, yp[i,j])
        pylab.plot(x, y0, 'r-', linewidth=2.0)

    # run tests for each type of fit, return results to compare parameters

    cons = None
    bnds = None

    signal_to_noise = 100000.
    for func in Fits.fitfuncmap:
        if func != 'exppulse':
            continue
        print "\nFunction: %s\nTarget: " % (func),
        f = Fits.fitfuncmap[func]
        for k in range(0,len(f[1])):
            print "%f " % (f[1][k]),
        print "\nStarting:     ",
        for k in range(0,len(f[5])):
            print "%f " % (f[5][k]),

#        nstep = 500.0
#        if func == 'sin':
#            nstep = 100.0
        x = numpy.array(numpy.arange(f[4][0], f[4][1], f[4][2]))
        C = None
        if func == 'expsum2':
            C = f[7]

        if func == 'exppulse':
            C = f[7]

        y = f[0](f[1], x, C=C)
        yd = numpy.array(y)
        noise = numpy.random.normal(0, 0.1, yd.shape)
        my = numpy.amax(yd)
        #yd = yd + sigmax*0.05*my*(numpy.random.random_sample(shape(yd))-0.5)
        yd += noise*my/signal_to_noise
        testMethod = 'SLSQP'
        if func == 'taucurve':
            continue
            bounds=[(0., 100.), (0., 1000.), (0.0, 500.0), (0.1, 50.0),
                (0., 1000), (0.0, 500.0), (0.1, 50.0)]
            (fpar, xf, yf, names) = Fits.FitRegion(numpy.array([1]), 0, x, yd, fitFunc = func, bounds=bounds, method=testMethod)
        elif func == 'boltz':
            continue
            bounds = [(-0.5,0.5), (0.0, 20.0), (-120., 0.), (-20., 0.)]
            (fpar, xf, yf, names) = Fits.FitRegion(numpy.array([1]), 0, x, yd, fitFunc = func, bounds=bounds, method=testMethod)

        elif func == 'exp2':
            bounds=[(-0.001, 0.001), (-5.0, 0.), (1.0, 500.0), (-5.0, 0.0),
                (1., 10000.)]
            (fpar, xf, yf, names) = Fits.FitRegion(numpy.array([1]), 0, x, yd, fitFunc = func, bounds=bounds, method=testMethod)

        elif func == 'exppulse':
            # set some constraints to the fitting
            # yOffset, tau1, tau2, amp, width = f[1]  # order of constraings
            dt = numpy.mean(numpy.diff(x))
            bounds = [(-5, 5), (-15., 15.), (-2, 2.0), (2-10, 10.), (-5, 5.), (0., 5.)]
            # cxample for constraints:
            # cons = ({'type': 'ineq', 'fun': lambda x:   x[4] - 3.0*x[2]},
            #         {'type': 'ineq', 'fun': lambda x:   - x[4] + 12*x[2]},
            #         {'type': 'ineq', 'fun': lambda x:   x[2]},
            #         {'type': 'ineq', 'fun': lambda x:  - x[4] + 2000},
            #         )
            cons = ({'type': 'ineq', 'fun': lambda x:   x[3] - x[2] }, # tau1 < tau2
                )
            C = None

            tv = f[5]
            initialgr = f[0](f[5], x, None )
            (fpar, xf, yf, names) = Fits.FitRegion(
                numpy.array([1]), 0, x, yd, fitFunc = func, fixedPars = C, constraints = cons, bounds = bounds, method=testMethod)
            # print xf
            # print yf
            # print fpar
            # print names

        else:
            (fpar, xf, yf, names) = Fits.FitRegion(
                numpy.array([1]), 0, x, yd, fitFunc = func, fixedPars = C, constraints = cons, bounds = bnds, method=testMethod)
        #print fpar
        s = numpy.shape(fpar)
        j = 0
        outstr = ""
        initstr = ""
        truestr = ""
        for i in range(0, len(names[j])):
#            print "%f " % fpar[j][i],
            outstr = outstr + ('%s = %f, ' % (names[j][i], fpar[j][i]))
            initstr = initstr + '%s = %f, ' % (names[j][i], tv[i])
            truestr = truestr + '%s = %f, ' % (names[j][i], f[1][i])
        print( "\nTrue(%d) : %s" % (j, truestr) )
        print( "FIT(%d)   : %s" % (j, outstr) )
        print( "init(%d) : %s" % (j, initstr) )
        print( "Error:   : %f" % (Fits.fitSum2Err))
        if func is 'exppulse':
            pylab.figure()
            pylab.plot(numpy.array(x), yd, 'ro-')
            pylab.hold(True)
            pylab.plot(numpy.array(x), initialgr, 'k--')
            pylab.plot(xf[0], yf[0], 'b-') # fit
            pylab.show()


"""
Massively univariate analysis of a computation task from the Localizer dataset
==============================================================================

A permuted Ordinary Least Squares algorithm is run at each voxel in
order to determine which voxels are specifically active when a healthy subject
performs a computation task as opposed to a sentence reading task.

Randomized Parcellation Based Inference [1] is also used so as to illustrate
that it conveys more sensitivity.

"""
# Author: Virgile Fritsch, <virgile.fritsch@inria.fr>, Mar. 2014
import numpy as np
from nilearn import datasets
from scipy import linalg
from nilearn.input_data import NiftiMasker
from nilearn.mass_univariate import permuted_ols
from nilearn_sandbox.mass_univariate.rpbi import randomized_parcellation_based_inference

### Load Localizer motor contrast #############################################
n_samples = 20
# localizer_dataset = datasets.fetch_localizer_calculation_task(
#     n_subjects=n_samples)
localizer_dataset = datasets.fetch_localizer_contrasts(
    ["calculation vs sentences"],
    n_subjects=n_samples)

### Mask data #################################################################
nifti_masker = NiftiMasker(
    memory='nilearn_cache', memory_level=1)  # cache options
fmri_masked = nifti_masker.fit_transform(localizer_dataset.cmaps)

### Perform massively univariate analysis with permuted OLS ###################
tested_var = np.ones((n_samples, 1), dtype=float)  # intercept
neg_log_pvals, all_scores, h0 = permuted_ols(
    tested_var, fmri_masked, model_intercept=False,
    n_perm=5000,  # 5,000 for the sake of time. 10,000 is recommended
    two_sided_test=False,  # RPBI does not perform a two-sided test
    n_jobs=1)  # can be changed to use more CPUs
neg_log_pvals_unmasked = nifti_masker.inverse_transform(
    np.ravel(neg_log_pvals))

### Randomized Parcellation Based Inference ###################################
neg_log_pvals_rpbi, _, _ = randomized_parcellation_based_inference(
    tested_var, fmri_masked,
    np.asarray(nifti_masker.mask_img_.get_data()).astype(bool),
    n_parcellations=30,  # 30 for the sake of time, 100 is recommended
    n_parcels=1000,
    threshold='auto',
    n_perm=5000,  # 5,000 for the sake of time. 10,000 is recommended
    random_state=0, memory='nilearn_cache', n_jobs=1, verbose=True)
neg_log_pvals_rpbi_unmasked = nifti_masker.inverse_transform(
    neg_log_pvals_rpbi)

### Visualization #############################################################
import matplotlib.pyplot as plt
from nilearn.plotting import plot_stat_map

# Here, we should use a structural image as a background, when available.

# Various plotting parameters
z_slice = 39  # plotted slice
from nilearn.image.resampling import coord_transform
affine = neg_log_pvals_unmasked.get_affine()
_, _, k_slice = coord_transform(0, 0, z_slice,
                                linalg.inv(affine))
k_slice = round(k_slice)

threshold = - np.log10(0.1)  # 10% corrected
vmax = min(np.amax(neg_log_pvals),
           np.amax(neg_log_pvals_rpbi))

# Plot permutation p-values map
fig = plt.figure(figsize=(5, 7), facecolor='k')

display = plot_stat_map(neg_log_pvals_unmasked,
                        threshold=threshold, cmap=plt.cm.autumn,
                        display_mode='z', cut_coords=[z_slice],
                        figure=fig, vmax=vmax, black_bg=True)

neg_log_pvals_data = neg_log_pvals_unmasked.get_data()
neg_log_pvals_slice_data = \
    neg_log_pvals_data[..., k_slice]
n_detections = (neg_log_pvals_slice_data > threshold).sum()
title = ('Negative $\log_{10}$ p-values'
         '\n(Non-parametric + '
         '\nmax-type correction)'
         '\n%d detections') % n_detections

display.title(title, y=1.2)

# Plot RPBI p-values map
fig = plt.figure(figsize=(5, 7), facecolor='k')

display = plot_stat_map(neg_log_pvals_rpbi_unmasked,
                        threshold=threshold, cmap=plt.cm.autumn,
                        display_mode='z', cut_coords=[z_slice],
                        figure=fig, vmax=vmax, black_bg=True)

neg_log_pvals_rpbi_data = \
    neg_log_pvals_rpbi_unmasked.get_data()
neg_log_pvals_rpbi_slice_data = \
    neg_log_pvals_rpbi_data[..., k_slice]
n_detections = (neg_log_pvals_rpbi_slice_data > threshold).sum()
title = ('Negative $\log_{10}$ p-values' + '\n(RPBI)'
         '\n%d detections') % n_detections

display.title(title, y=1.2)

plt.show()


# Must satisfy the signature
# [t,X,D,P] = sim_function(T,X0,D0,P0,I0);

import numpy as np
from scipy.integrate import ode

import matplotlib.pyplot as PLT

PLOT = True

class SIM(object):
    def __init__(self, plt, pvt_init_data):
        #print I
        # atol = 1e-10
        rtol = 1e-5
        # tt,YY,dummy_D,dummy_P
        self.solver = ode(dyn).set_integrator('dopri5', rtol=rtol)
        return

    def sim(self, TT, X0, D, P, U, W, property_checker):

        if PLOT:
            num_dim_x = len(X0)
            plot_data = [np.empty(0, dtype=float), np.empty((0, num_dim_x), dtype=float)]
        else:
            plot_data = None

        Ti = TT[0]
        Tf = TT[1]
        T = Tf - Ti

        self.solver.set_solout(solout_fun(property_checker, plot_data))  # (2)
        self.solver.set_initial_value(X0, t=0.0)
        self.solver.set_f_params(W)
        X_ = self.solver.integrate(T)

        pvf = property_checker.check(Tf, X_)

        dummy_D = np.zeros(D.shape)
        dummy_P = np.zeros(P.shape)
        ret_t = Tf
        ret_X = X_
        ret_D = dummy_D
        ret_P = dummy_P

        if PLOT:
            PLT.figure(5)
            PLT.plot(plot_data[0], plot_data[1][:, 0])

        return (ret_t, ret_X, ret_D, ret_P), pvf

# State Space Modeling Template
# dx/dt = Ax + Bu
# y = Cx + Du
def dyn(t, X, w):

    if w > 0:
        u = 1.0
    elif w < 0:
        u = -1.0
    else:
        u = 0.0

    x2 = u
    X_ = np.array([x2])
    return X_

def solout_fun(property_checker, plot_data):

    def solout(t, Y):
        if PLOT:
            plot_data[0] = np.concatenate((plot_data[0], np.array([t])))
            plot_data[1] = np.concatenate((plot_data[1], np.array([Y])))

        if property_checker.check(t, Y):
            #violating_state[0] = (np.copy(t), np.copy(Y))
            # print 'violation found:', violating_state[0]
            # return -1 to stop integration
            return -1
        else:
            return 0

        return 0

    return solout


'''
Bayesian Online Compressed Sensing (2016)
Paulo V. Rossi & Yoshiyuki Kabashima
'''

from collections import namedtuple

import matplotlib.pyplot as plt
import numpy as np
from numpy.linalg import norm
from numpy.random import normal

from utils import DlnH, DDlnH, G, H, moments

def simulation(method='standard'):
    signal_length = 2000
    alpha_max = 20
    sigma_n_2 = 1e-1

    phi = prior()
    P = posterior(signal_length, phi)
    x0 = generate_signal(signal_length, phi)

    print('Simulation parameters:')
    print('N='+str(signal_length)+', sparsity='+str(phi.rho)+
          ', noise='+str(sigma_n_2)+', alpha_max='+str(alpha_max))
    print('Measurement model: '+method+'\n')

    number_of_measurements = alpha_max*signal_length
    mean_square_error = np.zeros(number_of_measurements)
    for measurement in range(number_of_measurements):
        P = update_posterior(P, phi, x0, signal_length, sigma_n_2, method)
        mean_square_error[measurement] = reconstruction_error(P, x0)
    plot_results(P, x0, mean_square_error, phi)

def prior():
    phi = namedtuple('prior_distribution', ['rho', 'sigma_x_2', 'bar_x'])
    phi.rho = 0.1
    phi.sigma_x_2 = 1.
    phi.bar_x = 0.
    return phi

def posterior(signal_length, phi):
    P = namedtuple('posterior_distribution', ['m', 'v', 'a', 'h'])
    P.m = np.zeros(signal_length)
    P.v = phi.rho * phi.sigma_x_2 * np.ones(signal_length)
    P.a = np.zeros(signal_length)
    P.h = np.zeros(signal_length)
    return P

def generate_signal (signal_length, phi):
    x0 = np.zeros(signal_length)
    number_of_non_zero_components = int(np.ceil(signal_length*phi.rho))
    x0[:number_of_non_zero_components] = normal(loc=phi.bar_x,
                                                scale=np.sqrt(phi.sigma_x_2),
                                                size=number_of_non_zero_components)
    return x0

def update_posterior(P, phi, x0, signal_length, sigma_n_2, method):
    A_t = measurement_vector(signal_length)
    P.a, P.h = update_and_project(method, A_t, x0, sigma_n_2, P)
    P.m, P.v = moments(P, phi)
    return P

def measurement_vector(signal_length):
    A_t = normal(size=signal_length)
    return A_t/norm(A_t)

def update_and_project(method, A_t, x0, sigma_n_2, P):
    m, v, a, h = P.m, P.v, P.a, P.h
    u0 = np.dot(A_t, x0)
    if sigma_n_2 > 0:
        noise = normal(scale=np.sqrt(sigma_n_2))
    else:
        noise = 0
    y = u0 + noise
    Delta = np.dot(A_t, m)
    chi = np.dot(A_t**2, v)
    if method == 'standard':
        da, dh = update_and_project_std(y, Delta, chi, sigma_n_2, A_t, m)
    elif method == '1bit':
        da, dh = update_and_project_1bit(y, Delta, chi, sigma_n_2, A_t, m)
    else:
        raise ValueError('Measurement model not recognized. Please use "standard" or "1bit".')
    return a+da, h+dh

def update_and_project_std(y, Delta, chi, sigma_n_2, A_t, m):
    da = A_t**2 / (sigma_n_2 + chi)
    dh = (y-Delta)*A_t / (sigma_n_2 + chi) + da*m
    return da, dh

def update_and_project_1bit(y, Delta, chi, sigma_n_2, A_t, m):
    y = np.sign(y)
    u = y * np.dot(A_t, m)
    chi_prime = chi + sigma_n_2
    z = -u/np.sqrt(chi_prime)
    da = -A_t**2/chi_prime * DDlnH(z)
    dh = -y*A_t/np.sqrt(chi_prime) * DlnH(z) + da*m
    return da, dh

def reconstruction_error(P, x0):
    return norm(x0 - P.m)**2 / norm(x0)**2

def plot_results(P, x0, mse_t, phi):
    plt.subplots(figsize=(10,20))
    plt.subplot(211)
    plt.plot(np.arange(len(mse_t))/float(len(P.m)), 10*np.log10(mse_t), color='k')
    plt.xlabel(r'$\alpha$')
    plt.ylabel(r'mse (dB)')
    plt.subplot(212)
    plt.plot(P.m, color='k', lw = 0.7, label=r'$m$')
    plt.scatter(range(int(len(x0)*phi.rho)), x0[:int(len(x0)*phi.rho)], \
        marker='o', facecolors='none', edgecolors='r', lw=1.5, label=r'$x^0$')
    plt.xlim([0,len(P.m)])
    plt.xlabel(r'Vector Component')
    plt.legend()
    plt.show()

if __name__ == '__main__':
    simulation(method='1bit')
    #simulation(method='standard')