"""
==================================================
Probability Calibration for 3-class classification
==================================================

This example illustrates how sigmoid calibration changes predicted
probabilities for a 3-class classification problem. Illustrated is the
standard 2-simplex, where the three corners correspond to the three classes.
Arrows point from the probability vectors predicted by an uncalibrated
classifier to the probability vectors predicted by the same classifier after
sigmoid calibration on a hold-out validation set. Colors indicate the true
class of an instance (red: class 1, green: class 2, blue: class 3).

The base classifier is a random forest classifier with 25 base estimators
(trees). If this classifier is trained on all 800 training datapoints, it is
overly confident in its predictions and thus incurs a large log-loss.
Calibrating an identical classifier, which was trained on 600 datapoints, with
method='sigmoid' on the remaining 200 datapoints reduces the confidence of the
predictions, i.e., moves the probability vectors from the edges of the simplex
towards the center. This calibration results in a lower log-loss. Note that an
alternative would have been to increase the number of base estimators which
would have resulted in a similar decrease in log-loss.
"""
print(__doc__)

# Author: Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
# License: BSD Style.

import matplotlib.pyplot as plt

import numpy as np

from sklearn.datasets import make_blobs
from sklearn.ensemble import RandomForestClassifier
from sklearn.calibration import CalibratedClassifierCV
from sklearn.metrics import log_loss

# nodebox section
if __name__ == '__builtin__':
    # were in nodebox
    import os
    import tempfile
    W = 800
    inset = 20
    size(W, 600)
    plt.cla()
    plt.clf()
    plt.close('all')
    def tempimage():
        fob = tempfile.NamedTemporaryFile(mode='w+b', suffix='.png', delete=False)
        fname = fob.name
        fob.close()
        return fname
    imgx = 20
    imgy = 0
    def pltshow(plt, dpi=150):
        global imgx, imgy
        temppath = tempimage()
        plt.savefig(temppath, dpi=dpi)
        dx,dy = imagesize(temppath)
        w = min(W,dx)
        image(temppath,imgx,imgy,width=w)
        imgy = imgy + dy + 20
        os.remove(temppath)
        size(W, HEIGHT+dy+40)
else:
    def pltshow(mplpyplot):
        mplpyplot.show()
# nodebox section end

np.random.seed(0)

# Generate data
X, y = make_blobs(n_samples=1000, n_features=2, random_state=42,
                  cluster_std=5.0)
X_train, y_train = X[:600], y[:600]
X_valid, y_valid = X[600:800], y[600:800]
X_train_valid, y_train_valid = X[:800], y[:800]
X_test, y_test = X[800:], y[800:]

# Train uncalibrated random forest classifier on whole train and validation
# data and evaluate on test data
clf = RandomForestClassifier(n_estimators=25)
clf.fit(X_train_valid, y_train_valid)
clf_probs = clf.predict_proba(X_test)
score = log_loss(y_test, clf_probs)

# Train random forest classifier, calibrate on validation data and evaluate
# on test data
clf = RandomForestClassifier(n_estimators=25)
clf.fit(X_train, y_train)
clf_probs = clf.predict_proba(X_test)
sig_clf = CalibratedClassifierCV(clf, method="sigmoid", cv="prefit")
sig_clf.fit(X_valid, y_valid)
sig_clf_probs = sig_clf.predict_proba(X_test)
sig_score = log_loss(y_test, sig_clf_probs)

# Plot changes in predicted probabilities via arrows
plt.figure(0)
colors = ["r", "g", "b"]
for i in range(clf_probs.shape[0]):
    plt.arrow(clf_probs[i, 0], clf_probs[i, 1],
              sig_clf_probs[i, 0] - clf_probs[i, 0],
              sig_clf_probs[i, 1] - clf_probs[i, 1],
              color=colors[y_test[i]], head_width=1e-2)

# Plot perfect predictions
plt.plot([1.0], [0.0], 'ro', ms=20, label="Class 1")
plt.plot([0.0], [1.0], 'go', ms=20, label="Class 2")
plt.plot([0.0], [0.0], 'bo', ms=20, label="Class 3")

# Plot boundaries of unit simplex
plt.plot([0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], 'k', label="Simplex")

# Annotate points on the simplex
plt.annotate(r'($\frac{1}{3}$, $\frac{1}{3}$, $\frac{1}{3}$)',
             xy=(1.0/3, 1.0/3), xytext=(1.0/3, .23), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.plot([1.0/3], [1.0/3], 'ko', ms=5)
plt.annotate(r'($\frac{1}{2}$, $0$, $\frac{1}{2}$)',
             xy=(.5, .0), xytext=(.5, .1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($0$, $\frac{1}{2}$, $\frac{1}{2}$)',
             xy=(.0, .5), xytext=(.1, .5), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($\frac{1}{2}$, $\frac{1}{2}$, $0$)',
             xy=(.5, .5), xytext=(.6, .6), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($0$, $0$, $1$)',
             xy=(0, 0), xytext=(.1, .1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($1$, $0$, $0$)',
             xy=(1, 0), xytext=(1, .1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($0$, $1$, $0$)',
             xy=(0, 1), xytext=(.1, 1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
# Add grid
plt.grid("off")
for x in [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]:
    plt.plot([0, x], [x, 0], 'k', alpha=0.2)
    plt.plot([0, 0 + (1-x)/2], [x, x + (1-x)/2], 'k', alpha=0.2)
    plt.plot([x, x + (1-x)/2], [0, 0 + (1-x)/2], 'k', alpha=0.2)

plt.title("Change of predicted probabilities after sigmoid calibration")
plt.xlabel("Probability class 1")
plt.ylabel("Probability class 2")
plt.xlim(-0.05, 1.05)
plt.ylim(-0.05, 1.05)
plt.legend(loc="best")

print("Log-loss of")
print(" * uncalibrated classifier trained on 800 datapoints: %.3f "
      % score)
print(" * classifier trained on 600 datapoints and calibrated on "
      "200 datapoint: %.3f" % sig_score)

# Illustrate calibrator
plt.figure(1)
# generate grid over 2-simplex
p1d = np.linspace(0, 1, 20)
p0, p1 = np.meshgrid(p1d, p1d)
p2 = 1 - p0 - p1
p = np.c_[p0.ravel(), p1.ravel(), p2.ravel()]
p = p[p[:, 2] >= 0]

calibrated_classifier = sig_clf.calibrated_classifiers_[0]
prediction = np.vstack([calibrator.predict(this_p)
                        for calibrator, this_p in
                        zip(calibrated_classifier.calibrators_, p.T)]).T
prediction /= prediction.sum(axis=1)[:, None]

# Plot modifications of calibrator
for i in range(prediction.shape[0]):
    plt.arrow(p[i, 0], p[i, 1],
              prediction[i, 0] - p[i, 0], prediction[i, 1] - p[i, 1],
              head_width=1e-2, color=colors[np.argmax(p[i])])
# Plot boundaries of unit simplex
plt.plot([0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], 'k', label="Simplex")

plt.grid("off")
for x in [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]:
    plt.plot([0, x], [x, 0], 'k', alpha=0.2)
    plt.plot([0, 0 + (1-x)/2], [x, x + (1-x)/2], 'k', alpha=0.2)
    plt.plot([x, x + (1-x)/2], [0, 0 + (1-x)/2], 'k', alpha=0.2)

plt.title("Illustration of sigmoid calibrator")
plt.xlabel("Probability class 1")
plt.ylabel("Probability class 2")
plt.xlim(-0.05, 1.05)
plt.ylim(-0.05, 1.05)

# plt.show()
pltshow(plt)


"""
Support for plotting vector fields.

Presently this contains Quiver and Barb. Quiver plots an arrow in the
direction of the vector, with the size of the arrow related to the
magnitude of the vector.

Barbs are like quiver in that they point along a vector, but
the magnitude of the vector is given schematically by the presence of barbs
or flags on the barb.

This will also become a home for things such as standard
deviation ellipses, which can and will be derived very easily from
the Quiver code.
"""

import numpy as np
from numpy import ma
import matplotlib.collections as collections
import matplotlib.transforms as transforms
import matplotlib.text as mtext
import matplotlib.artist as martist
import matplotlib.font_manager as font_manager
from matplotlib.cbook import delete_masked_points
from matplotlib.patches import CirclePolygon
import math

_quiver_doc = """
Plot a 2-D field of arrows.

call signatures::

  quiver(U, V, **kw)
  quiver(U, V, C, **kw)
  quiver(X, Y, U, V, **kw)
  quiver(X, Y, U, V, C, **kw)

Arguments:

  *X*, *Y*:

    The x and y coordinates of the arrow locations (default is tail of
    arrow; see *pivot* kwarg)

  *U*, *V*:

    give the *x* and *y* components of the arrow vectors

  *C*:
    an optional array used to map colors to the arrows

All arguments may be 1-D or 2-D arrays or sequences. If *X* and *Y*
are absent, they will be generated as a uniform grid.  If *U* and *V*
are 2-D arrays but *X* and *Y* are 1-D, and if len(*X*) and len(*Y*)
match the column and row dimensions of *U*, then *X* and *Y* will be
expanded with :func:`numpy.meshgrid`.

*U*, *V*, *C* may be masked arrays, but masked *X*, *Y* are not
supported at present.

Keyword arguments:

  *units*: ['width' | 'height' | 'dots' | 'inches' | 'x' | 'y' ]
    arrow units; the arrow dimensions *except for length* are in
    multiples of this unit.

    * 'width' or 'height': the width or height of the axes

    * 'dots' or 'inches': pixels or inches, based on the figure dpi

    * 'x' or 'y': *X* or *Y* data units

    The arrows scale differently depending on the units.  For
    'x' or 'y', the arrows get larger as one zooms in; for other
    units, the arrow size is independent of the zoom state.  For
    'width or 'height', the arrow size increases with the width and
    height of the axes, respectively, when the the window is resized;
    for 'dots' or 'inches', resizing does not change the arrows.

   *angles*: ['uv' | 'xy' | array]
    With the default 'uv', the arrow aspect ratio is 1, so that
    if *U*==*V* the angle of the arrow on the plot is 45 degrees
    CCW from the *x*-axis.
    With 'xy', the arrow points from (x,y) to (x+u, y+v).
    Alternatively, arbitrary angles may be specified as an array
    of values in degrees, CCW from the *x*-axis.

  *scale*: [ None | float ]
    data units per arrow unit, e.g. m/s per plot width; a smaller
    scale parameter makes the arrow longer.  If *None*, a simple
    autoscaling algorithm is used, based on the average vector length
    and the number of vectors.

  *width*:
    shaft width in arrow units; default depends on choice of units,
    above, and number of vectors; a typical starting value is about
    0.005 times the width of the plot.

  *headwidth*: scalar
    head width as multiple of shaft width, default is 3

  *headlength*: scalar
    head length as multiple of shaft width, default is 5

  *headaxislength*: scalar
    head length at shaft intersection, default is 4.5

  *minshaft*: scalar
    length below which arrow scales, in units of head length. Do not
    set this to less than 1, or small arrows will look terrible!
    Default is 1

  *minlength*: scalar
    minimum length as a multiple of shaft width; if an arrow length
    is less than this, plot a dot (hexagon) of this diameter instead.
    Default is 1.

  *pivot*: [ 'tail' | 'middle' | 'tip' ]
    The part of the arrow that is at the grid point; the arrow rotates
    about this point, hence the name *pivot*.

  *color*: [ color | color sequence ]
    This is a synonym for the
    :class:`~matplotlib.collections.PolyCollection` facecolor kwarg.
    If *C* has been set, *color* has no effect.

The defaults give a slightly swept-back arrow; to make the head a
triangle, make *headaxislength* the same as *headlength*. To make the
arrow more pointed, reduce *headwidth* or increase *headlength* and
*headaxislength*. To make the head smaller relative to the shaft,
scale down all the head parameters. You will probably do best to leave
minshaft alone.

linewidths and edgecolors can be used to customize the arrow
outlines. Additional :class:`~matplotlib.collections.PolyCollection`
keyword arguments:

%(PolyCollection)s
""" % martist.kwdocd

_quiverkey_doc = """
Add a key to a quiver plot.

call signature::

  quiverkey(Q, X, Y, U, label, **kw)

Arguments:

  *Q*:
    The Quiver instance returned by a call to quiver.

  *X*, *Y*:
    The location of the key; additional explanation follows.

  *U*:
    The length of the key

  *label*:
    a string with the length and units of the key

Keyword arguments:

  *coordinates* = [ 'axes' | 'figure' | 'data' | 'inches' ]
    Coordinate system and units for *X*, *Y*: 'axes' and 'figure' are
    normalized coordinate systems with 0,0 in the lower left and 1,1
    in the upper right; 'data' are the axes data coordinates (used for
    the locations of the vectors in the quiver plot itself); 'inches'
    is position in the figure in inches, with 0,0 at the lower left
    corner.

  *color*:
    overrides face and edge colors from *Q*.

  *labelpos* = [ 'N' | 'S' | 'E' | 'W' ]
    Position the label above, below, to the right, to the left of the
    arrow, respectively.

  *labelsep*:
    Distance in inches between the arrow and the label.  Default is
    0.1

  *labelcolor*:
    defaults to default :class:`~matplotlib.text.Text` color.

  *fontproperties*:
    A dictionary with keyword arguments accepted by the
    :class:`~matplotlib.font_manager.FontProperties` initializer:
    *family*, *style*, *variant*, *size*, *weight*

Any additional keyword arguments are used to override vector
properties taken from *Q*.

The positioning of the key depends on *X*, *Y*, *coordinates*, and
*labelpos*.  If *labelpos* is 'N' or 'S', *X*, *Y* give the position
of the middle of the key arrow.  If *labelpos* is 'E', *X*, *Y*
positions the head, and if *labelpos* is 'W', *X*, *Y* positions the
tail; in either of these two cases, *X*, *Y* is somewhere in the
middle of the arrow+label key object.
"""

class QuiverKey(martist.Artist):
    """ Labelled arrow for use as a quiver plot scale key.
    """
    halign = {'N': 'center', 'S': 'center', 'E': 'left',   'W': 'right'}
    valign = {'N': 'bottom', 'S': 'top',    'E': 'center', 'W': 'center'}
    pivot  = {'N': 'mid',    'S': 'mid',    'E': 'tip',    'W': 'tail'}

    def __init__(self, Q, X, Y, U, label, **kw):
        martist.Artist.__init__(self)
        self.Q = Q
        self.X = X
        self.Y = Y
        self.U = U
        self.coord = kw.pop('coordinates', 'axes')
        self.color = kw.pop('color', None)
        self.label = label
        self._labelsep_inches = kw.pop('labelsep', 0.1)
        self.labelsep = (self._labelsep_inches * Q.ax.figure.dpi)

        def on_dpi_change(fig):
            self.labelsep = (self._labelsep_inches * fig.dpi)
            self._initialized = False # simple brute force update
                                      # works because _init is called
                                      # at the start of draw.

        Q.ax.figure.callbacks.connect('dpi_changed', on_dpi_change)

        self.labelpos = kw.pop('labelpos', 'N')
        self.labelcolor = kw.pop('labelcolor', None)
        self.fontproperties = kw.pop('fontproperties', dict())
        self.kw = kw
        _fp = self.fontproperties
        #boxprops = dict(facecolor='red')
        self.text = mtext.Text(text=label,  #      bbox=boxprops,
                       horizontalalignment=self.halign[self.labelpos],
                       verticalalignment=self.valign[self.labelpos],
                       fontproperties=font_manager.FontProperties(**_fp))
        if self.labelcolor is not None:
            self.text.set_color(self.labelcolor)
        self._initialized = False
        self.zorder = Q.zorder + 0.1

    __init__.__doc__ = _quiverkey_doc

    def _init(self):
        if True: ##not self._initialized:
            self._set_transform()
            _pivot = self.Q.pivot
            self.Q.pivot = self.pivot[self.labelpos]
            self.verts = self.Q._make_verts(np.array([self.U]),
                                                        np.zeros((1,)))
            self.Q.pivot = _pivot
            kw = self.Q.polykw
            kw.update(self.kw)
            self.vector = collections.PolyCollection(self.verts,
                                         offsets=[(self.X,self.Y)],
                                         transOffset=self.get_transform(),
                                         **kw)
            if self.color is not None:
                self.vector.set_color(self.color)
            self.vector.set_transform(self.Q.get_transform())
            self._initialized = True

    def _text_x(self, x):
        if self.labelpos == 'E':
            return x + self.labelsep
        elif self.labelpos == 'W':
            return x - self.labelsep
        else:
            return x

    def _text_y(self, y):
        if self.labelpos == 'N':
            return y + self.labelsep
        elif self.labelpos == 'S':
            return y - self.labelsep
        else:
            return y

    def draw(self, renderer):
        self._init()
        self.vector.draw(renderer)
        x, y = self.get_transform().transform_point((self.X, self.Y))
        self.text.set_x(self._text_x(x))
        self.text.set_y(self._text_y(y))
        self.text.draw(renderer)

    def _set_transform(self):
        if self.coord == 'data':
            self.set_transform(self.Q.ax.transData)
        elif self.coord == 'axes':
            self.set_transform(self.Q.ax.transAxes)
        elif self.coord == 'figure':
            self.set_transform(self.Q.ax.figure.transFigure)
        elif self.coord == 'inches':
            self.set_transform(self.Q.ax.figure.dpi_scale_trans)
        else:
            raise ValueError('unrecognized coordinates')

    def set_figure(self, fig):
        martist.Artist.set_figure(self, fig)
        self.text.set_figure(fig)

    def contains(self, mouseevent):
        # Maybe the dictionary should allow one to
        # distinguish between a text hit and a vector hit.
        if (self.text.contains(mouseevent)[0]
                or self.vector.contains(mouseevent)[0]):
            return True, {}
        return False, {}

    quiverkey_doc = _quiverkey_doc

class Quiver(collections.PolyCollection):
    """
    Specialized PolyCollection for arrows.

    The only API method is set_UVC(), which can be used
    to change the size, orientation, and color of the
    arrows; their locations are fixed when the class is
    instantiated.  Possibly this method will be useful
    in animations.

    Much of the work in this class is done in the draw()
    method so that as much information as possible is available
    about the plot.  In subsequent draw() calls, recalculation
    is limited to things that might have changed, so there
    should be no performance penalty from putting the calculations
    in the draw() method.
    """
    def __init__(self, ax, *args, **kw):
        self.ax = ax
        X, Y, U, V, C = self._parse_args(*args)
        self.X = X
        self.Y = Y
        self.XY = np.hstack((X[:,np.newaxis], Y[:,np.newaxis]))
        self.N = len(X)
        self.scale = kw.pop('scale', None)
        self.headwidth = kw.pop('headwidth', 3)
        self.headlength = float(kw.pop('headlength', 5))
        self.headaxislength = kw.pop('headaxislength', 4.5)
        self.minshaft = kw.pop('minshaft', 1)
        self.minlength = kw.pop('minlength', 1)
        self.units = kw.pop('units', 'width')
        self.angles = kw.pop('angles', 'uv')
        self.width = kw.pop('width', None)
        self.color = kw.pop('color', 'k')
        self.pivot = kw.pop('pivot', 'tail')
        kw.setdefault('facecolors', self.color)
        kw.setdefault('linewidths', (0,))
        collections.PolyCollection.__init__(self, [], offsets=self.XY,
                                            transOffset=ax.transData,
                                            closed=False,
                                            **kw)
        self.polykw = kw
        self.set_UVC(U, V, C)
        self._initialized = False

        self.keyvec = None
        self.keytext = None

        def on_dpi_change(fig):
            self._new_UV = True # vertices depend on width, span
                                # which in turn depend on dpi
            self._initialized = False # simple brute force update
                                      # works because _init is called
                                      # at the start of draw.

        self.ax.figure.callbacks.connect('dpi_changed', on_dpi_change)

    __init__.__doc__ = """
        The constructor takes one required argument, an Axes
        instance, followed by the args and kwargs described
        by the following pylab interface documentation:
        %s""" % _quiver_doc

    def _parse_args(self, *args):
        X, Y, U, V, C = [None]*5
        args = list(args)
        if len(args) == 3 or len(args) == 5:
            C = ma.asarray(args.pop(-1)).ravel()
        V = ma.asarray(args.pop(-1))
        U = ma.asarray(args.pop(-1))
        nn = np.shape(U)
        nc = nn[0]
        nr = 1
        if len(nn) > 1:
            nr = nn[1]
        if len(args) == 2: # remaining after removing U,V,C
            X, Y = [np.array(a).ravel() for a in args]
            if len(X) == nc and len(Y) == nr:
                X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
        else:
            indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
            X, Y = [np.ravel(a) for a in indexgrid]
        return X, Y, U, V, C

    def _init(self):
        """initialization delayed until first draw;
        allow time for axes setup.
        """
        # It seems that there are not enough event notifications
        # available to have this work on an as-needed basis at present.
        if True: ##not self._initialized:
            trans = self._set_transform()
            ax = self.ax
            sx, sy = trans.inverted().transform_point(
                                            (ax.bbox.width, ax.bbox.height))
            self.span = sx
            sn = max(8, min(25, math.sqrt(self.N)))
            if self.width is None:
                self.width = 0.06 * self.span / sn

    def draw(self, renderer):
        self._init()
        if self._new_UV or self.angles == 'xy':
            verts = self._make_verts(self.U, self.V)
            self.set_verts(verts, closed=False)
            self._new_UV = False
        collections.PolyCollection.draw(self, renderer)

    def set_UVC(self, U, V, C=None):
        self.U = U.ravel()
        self.V = V.ravel()
        if C is not None:
            self.set_array(C.ravel())
        self._new_UV = True

    def _set_transform(self):
        ax = self.ax
        if self.units in ('x', 'y'):
            if self.units == 'x':
                dx0 = ax.viewLim.width
                dx1 = ax.bbox.width
            else:
                dx0 = ax.viewLim.height
                dx1 = ax.bbox.height
            dx = dx1/dx0
        else:
            if self.units == 'width':
                dx = ax.bbox.width
            elif self.units == 'height':
                dx = ax.bbox.height
            elif self.units == 'dots':
                dx = 1.0
            elif self.units == 'inches':
                dx = ax.figure.dpi
            else:
                raise ValueError('unrecognized units')
        trans = transforms.Affine2D().scale(dx)
        self.set_transform(trans)
        return trans

    def _angles(self, U, V, eps=0.001):
        xy = self.ax.transData.transform(self.XY)
        uv = ma.hstack((U[:,np.newaxis], V[:,np.newaxis])).filled(0)
        xyp = self.ax.transData.transform(self.XY + eps * uv)
        dxy = xyp - xy
        ang = ma.arctan2(dxy[:,1], dxy[:,0])
        return ang

    def _make_verts(self, U, V):
        uv = ma.asarray(U+V*1j)
        a = ma.absolute(uv)
        if self.scale is None:
            sn = max(10, math.sqrt(self.N))
            scale = 1.8 * a.mean() * sn / self.span # crude auto-scaling
            self.scale = scale
        length = a/(self.scale*self.width)
        X, Y = self._h_arrows(length)
        if self.angles == 'xy':
            theta = self._angles(U, V).filled(0)[:,np.newaxis]
        elif self.angles == 'uv':
            theta = np.angle(ma.asarray(uv[..., np.newaxis]).filled(0))
        else:
            theta = ma.asarray(self.angles*np.pi/180.0).filled(0)
        xy = (X+Y*1j) * np.exp(1j*theta)*self.width
        xy = xy[:,:,np.newaxis]
        XY = ma.concatenate((xy.real, xy.imag), axis=2)
        return XY

    def _h_arrows(self, length):
        """ length is in arrow width units """
        # It might be possible to streamline the code
        # and speed it up a bit by using complex (x,y)
        # instead of separate arrays; but any gain would be slight.
        minsh = self.minshaft * self.headlength
        N = len(length)
        length = length.reshape(N, 1)
        # x, y: normal horizontal arrow
        x = np.array([0, -self.headaxislength,
                        -self.headlength, 0], np.float64)
        x = x + np.array([0,1,1,1]) * length
        y = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
        y = np.repeat(y[np.newaxis,:], N, axis=0)
        # x0, y0: arrow without shaft, for short vectors
        x0 = np.array([0, minsh-self.headaxislength,
                        minsh-self.headlength, minsh], np.float64)
        y0 = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
        ii = [0,1,2,3,2,1,0]
        X = x.take(ii, 1)
        Y = y.take(ii, 1)
        Y[:, 3:] *= -1
        X0 = x0.take(ii)
        Y0 = y0.take(ii)
        Y0[3:] *= -1
        shrink = length/minsh
        X0 = shrink * X0[np.newaxis,:]
        Y0 = shrink * Y0[np.newaxis,:]
        short = np.repeat(length < minsh, 7, axis=1)
        #print 'short', length < minsh
        # Now select X0, Y0 if short, otherwise X, Y
        X = ma.where(short, X0, X)
        Y = ma.where(short, Y0, Y)
        if self.pivot[:3] == 'mid':
            X -= 0.5 * X[:,3, np.newaxis]
        elif self.pivot[:3] == 'tip':
            X = X - X[:,3, np.newaxis]   #numpy bug? using -= does not
                                         # work here unless we multiply
                                         # by a float first, as with 'mid'.
        tooshort = length < self.minlength
        if tooshort.any():
            # Use a heptagonal dot:
            th = np.arange(0,7,1, np.float64) * (np.pi/3.0)
            x1 = np.cos(th) * self.minlength * 0.5
            y1 = np.sin(th) * self.minlength * 0.5
            X1 = np.repeat(x1[np.newaxis, :], N, axis=0)
            Y1 = np.repeat(y1[np.newaxis, :], N, axis=0)
            tooshort = ma.repeat(tooshort, 7, 1)
            X = ma.where(tooshort, X1, X)
            Y = ma.where(tooshort, Y1, Y)
        return X, Y

    quiver_doc = _quiver_doc

_barbs_doc = """
Plot a 2-D field of barbs.

call signatures::

  barb(U, V, **kw)
  barb(U, V, C, **kw)
  barb(X, Y, U, V, **kw)
  barb(X, Y, U, V, C, **kw)

Arguments:

  *X*, *Y*:
    The x and y coordinates of the barb locations
    (default is head of barb; see *pivot* kwarg)

  *U*, *V*:
    give the *x* and *y* components of the barb shaft

  *C*:
    an optional array used to map colors to the barbs

All arguments may be 1-D or 2-D arrays or sequences. If *X* and *Y*
are absent, they will be generated as a uniform grid.  If *U* and *V*
are 2-D arrays but *X* and *Y* are 1-D, and if len(*X*) and len(*Y*)
match the column and row dimensions of *U*, then *X* and *Y* will be
expanded with :func:`numpy.meshgrid`.

*U*, *V*, *C* may be masked arrays, but masked *X*, *Y* are not
supported at present.

Keyword arguments:

  *length*:
    Length of the barb in points; the other parts of the barb
    are scaled against this.
    Default is 9

  *pivot*: [ 'tip' | 'middle' ]
    The part of the arrow that is at the grid point; the arrow rotates
    about this point, hence the name *pivot*.  Default is 'tip'

  *barbcolor*: [ color | color sequence ]
    Specifies the color all parts of the barb except any flags.  This
    parameter is analagous to the *edgecolor* parameter for polygons,
    which can be used instead. However this parameter will override
    facecolor.

  *flagcolor*: [ color | color sequence ]
    Specifies the color of any flags on the barb.  This parameter is
    analagous to the *facecolor* parameter for polygons, which can be
    used instead. However this parameter will override facecolor.  If
    this is not set (and *C* has not either) then *flagcolor* will be
    set to match *barbcolor* so that the barb has a uniform color. If
    *C* has been set, *flagcolor* has no effect.

  *sizes*:
    A dictionary of coefficients specifying the ratio of a given
    feature to the length of the barb. Only those values one wishes to
    override need to be included.  These features include:

        - 'spacing' - space between features (flags, full/half barbs)

        - 'height' - height (distance from shaft to top) of a flag or
          full barb

        - 'width' - width of a flag, twice the width of a full barb

        - 'emptybarb' - radius of the circle used for low magnitudes

  *fill_empty*:
    A flag on whether the empty barbs (circles) that are drawn should
    be filled with the flag color.  If they are not filled, they will
    be drawn such that no color is applied to the center.  Default is
    False

  *rounding*:
    A flag to indicate whether the vector magnitude should be rounded
    when allocating barb components.  If True, the magnitude is
    rounded to the nearest multiple of the half-barb increment.  If
    False, the magnitude is simply truncated to the next lowest
    multiple.  Default is True

  *barb_increments*:
    A dictionary of increments specifying values to associate with
    different parts of the barb. Only those values one wishes to
    override need to be included.

        - 'half' - half barbs (Default is 5)

        - 'full' - full barbs (Default is 10)

        - 'flag' - flags (default is 50)

  *flip_barb*:
    Either a single boolean flag or an array of booleans.  Single
    boolean indicates whether the lines and flags should point
    opposite to normal for all barbs.  An array (which should be the
    same size as the other data arrays) indicates whether to flip for
    each individual barb.  Normal behavior is for the barbs and lines
    to point right (comes from wind barbs having these features point
    towards low pressure in the Northern Hemisphere.)  Default is
    False

Barbs are traditionally used in meteorology as a way to plot the speed
and direction of wind observations, but can technically be used to
plot any two dimensional vector quantity.  As opposed to arrows, which
give vector magnitude by the length of the arrow, the barbs give more
quantitative information about the vector magnitude by putting slanted
lines or a triangle for various increments in magnitude, as show
schematically below::

 :     /\    \\
 :    /  \    \\
 :   /    \    \    \\
 :  /      \    \    \\
 : ------------------------------

.. note the double \\ at the end of each line to make the figure
.. render correctly

The largest increment is given by a triangle (or "flag"). After those
come full lines (barbs). The smallest increment is a half line.  There
is only, of course, ever at most 1 half line.  If the magnitude is
small and only needs a single half-line and no full lines or
triangles, the half-line is offset from the end of the barb so that it
can be easily distinguished from barbs with a single full line.  The
magnitude for the barb shown above would nominally be 65, using the
standard increments of 50, 10, and 5.

linewidths and edgecolors can be used to customize the barb.
Additional :class:`~matplotlib.collections.PolyCollection` keyword
arguments:

%(PolyCollection)s
""" % martist.kwdocd

class Barbs(collections.PolyCollection):
    '''
    Specialized PolyCollection for barbs.

    The only API method is :meth:`set_UVC`, which can be used to
    change the size, orientation, and color of the arrows.  Locations
    are changed using the :meth:`set_offsets` collection method.
    Possibly this method will be useful in animations.

    There is one internal function :meth:`_find_tails` which finds
    exactly what should be put on the barb given the vector magnitude.
    From there :meth:`_make_barbs` is used to find the vertices of the
    polygon to represent the barb based on this information.
    '''
    #This may be an abuse of polygons here to render what is essentially maybe
    #1 triangle and a series of lines.  It works fine as far as I can tell
    #however.
    def __init__(self, ax, *args, **kw):
        self._pivot = kw.pop('pivot', 'tip')
        self._length = kw.pop('length', 7)
        barbcolor = kw.pop('barbcolor', None)
        flagcolor = kw.pop('flagcolor', None)
        self.sizes = kw.pop('sizes', dict())
        self.fill_empty = kw.pop('fill_empty', False)
        self.barb_increments = kw.pop('barb_increments', dict())
        self.rounding = kw.pop('rounding', True)
        self.flip = kw.pop('flip_barb', False)

        #Flagcolor and and barbcolor provide convenience parameters for setting
        #the facecolor and edgecolor, respectively, of the barb polygon.  We
        #also work here to make the flag the same color as the rest of the barb
        #by default
        if None in (barbcolor, flagcolor):
            kw['edgecolors'] = 'face'
            if flagcolor:
                kw['facecolors'] = flagcolor
            elif barbcolor:
                kw['facecolors'] = barbcolor
            else:
                #Set to facecolor passed in or default to black
                kw.setdefault('facecolors', 'k')
        else:
            kw['edgecolors'] = barbcolor
            kw['facecolors'] = flagcolor

        #Parse out the data arrays from the various configurations supported
        x, y, u, v, c = self._parse_args(*args)
        self.x = x
        self.y = y
        xy = np.hstack((x[:,np.newaxis], y[:,np.newaxis]))

        #Make a collection
        barb_size = self._length**2 / 4 #Empirically determined
        collections.PolyCollection.__init__(self, [], (barb_size,), offsets=xy,
            transOffset=ax.transData, **kw)
        self.set_transform(transforms.IdentityTransform())

        self.set_UVC(u, v, c)

    __init__.__doc__ = """
        The constructor takes one required argument, an Axes
        instance, followed by the args and kwargs described
        by the following pylab interface documentation:
        %s""" % _barbs_doc

    def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
        '''
        Find how many of each of the tail pieces is necessary.  Flag
        specifies the increment for a flag, barb for a full barb, and half for
        half a barb. Mag should be the magnitude of a vector (ie. >= 0).

        This returns a tuple of:

            (*number of flags*, *number of barbs*, *half_flag*, *empty_flag*)

        *half_flag* is a boolean whether half of a barb is needed,
        since there should only ever be one half on a given
        barb. *empty_flag* flag is an array of flags to easily tell if
        a barb is empty (too low to plot any barbs/flags.
        '''

        #If rounding, round to the nearest multiple of half, the smallest
        #increment
        if rounding:
            mag = half * (mag / half + 0.5).astype(np.int)

        num_flags = np.floor(mag / flag).astype(np.int)
        mag = np.mod(mag, flag)

        num_barb = np.floor(mag / full).astype(np.int)
        mag = np.mod(mag, full)

        half_flag = mag >= half
        empty_flag = ~(half_flag | (num_flags > 0) | (num_barb > 0))

        return num_flags, num_barb, half_flag, empty_flag

    def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
        pivot, sizes, fill_empty, flip):
        '''
        This function actually creates the wind barbs.  *u* and *v*
        are components of the vector in the *x* and *y* directions,
        respectively.

        *nflags*, *nbarbs*, and *half_barb*, empty_flag* are,
        *respectively, the number of flags, number of barbs, flag for
        *half a barb, and flag for empty barb, ostensibly obtained
        *from :meth:`_find_tails`.

        *length* is the length of the barb staff in points.

        *pivot* specifies the point on the barb around which the
        entire barb should be rotated.  Right now, valid options are
        'head' and 'middle'.

        *sizes* is a dictionary of coefficients specifying the ratio
        of a given feature to the length of the barb. These features
        include:

            - *spacing*: space between features (flags, full/half
               barbs)

            - *height*: distance from shaft of top of a flag or full
               barb

            - *width* - width of a flag, twice the width of a full barb

            - *emptybarb* - radius of the circle used for low
               magnitudes

        *fill_empty* specifies whether the circle representing an
        empty barb should be filled or not (this changes the drawing
        of the polygon).

        *flip* is a flag indicating whether the features should be flipped to
        the other side of the barb (useful for winds in the southern
        hemisphere.

        This function returns list of arrays of vertices, defining a polygon for
        each of the wind barbs.  These polygons have been rotated to properly
        align with the vector direction.
        '''

        #These control the spacing and size of barb elements relative to the
        #length of the shaft
        spacing = length * sizes.get('spacing', 0.125)
        full_height = length * sizes.get('height', 0.4)
        full_width = length * sizes.get('width', 0.25)
        empty_rad = length * sizes.get('emptybarb', 0.15)

        #Controls y point where to pivot the barb.
        pivot_points = dict(tip=0.0, middle=-length/2.)

        #Check for flip
        if flip: full_height = -full_height

        endx = 0.0
        endy = pivot_points[pivot.lower()]

        #Get the appropriate angle for the vector components.  The offset is due
        #to the way the barb is initially drawn, going down the y-axis.  This
        #makes sense in a meteorological mode of thinking since there 0 degrees
        #corresponds to north (the y-axis traditionally)
        angles = -(ma.arctan2(v, u) + np.pi/2)

        #Used for low magnitude.  We just get the vertices, so if we make it
        #out here, it can be reused.  The center set here should put the
        #center of the circle at the location(offset), rather than at the
        #same point as the barb pivot; this seems more sensible.
        circ = CirclePolygon((0,0), radius=empty_rad).get_verts()
        if fill_empty:
            empty_barb = circ
        else:
            #If we don't want the empty one filled, we make a degenerate polygon
            #that wraps back over itself
            empty_barb = np.concatenate((circ, circ[::-1]))

        barb_list = []
        for index, angle in np.ndenumerate(angles):
            #If the vector magnitude is too weak to draw anything, plot an
            #empty circle instead
            if empty_flag[index]:
                #We can skip the transform since the circle has no preferred
                #orientation
                barb_list.append(empty_barb)
                continue

            poly_verts = [(endx, endy)]
            offset = length

            #Add vertices for each flag
            for i in range(nflags[index]):
                #The spacing that works for the barbs is a little to much for
                #the flags, but this only occurs when we have more than 1 flag.
                if offset != length: offset += spacing / 2.
                poly_verts.extend([[endx, endy + offset],
                    [endx + full_height, endy - full_width/2 + offset],
                    [endx, endy - full_width + offset]])

                offset -= full_width + spacing

            #Add vertices for each barb.  These really are lines, but works
            #great adding 3 vertices that basically pull the polygon out and
            #back down the line
            for i in range(nbarbs[index]):
                poly_verts.extend([(endx, endy + offset),
                    (endx + full_height, endy + offset + full_width/2),
                    (endx, endy + offset)])

                offset -= spacing

            #Add the vertices for half a barb, if needed
            if half_barb[index]:
                #If the half barb is the first on the staff, traditionally it is
                #offset from the end to make it easy to distinguish from a barb
                #with a full one
                if offset == length:
                    poly_verts.append((endx, endy + offset))
                    offset -= 1.5 * spacing
                poly_verts.extend([(endx, endy + offset),
                    (endx + full_height/2, endy + offset + full_width/4),
                    (endx, endy + offset)])

            #Rotate the barb according the angle. Making the barb first and then
            #rotating it made the math for drawing the barb really easy.  Also,
            #the transform framework makes doing the rotation simple.
            poly_verts = transforms.Affine2D().rotate(-angle).transform(
                poly_verts)
            barb_list.append(poly_verts)

        return barb_list

    #Taken shamelessly from Quiver
    def _parse_args(self, *args):
        X, Y, U, V, C = [None]*5
        args = list(args)
        if len(args) == 3 or len(args) == 5:
            C = ma.asarray(args.pop(-1)).ravel()
        V = ma.asarray(args.pop(-1))
        U = ma.asarray(args.pop(-1))
        nn = np.shape(U)
        nc = nn[0]
        nr = 1
        if len(nn) > 1:
            nr = nn[1]
        if len(args) == 2: # remaining after removing U,V,C
            X, Y = [np.array(a).ravel() for a in args]
            if len(X) == nc and len(Y) == nr:
                X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
        else:
            indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
            X, Y = [np.ravel(a) for a in indexgrid]
        return X, Y, U, V, C

    def set_UVC(self, U, V, C=None):
        self.u = ma.asarray(U).ravel()
        self.v = ma.asarray(V).ravel()
        if C is not None:
            c = ma.asarray(C).ravel()
            x,y,u,v,c = delete_masked_points(self.x.ravel(), self.y.ravel(),
                self.u, self.v, c)
        else:
            x,y,u,v = delete_masked_points(self.x.ravel(), self.y.ravel(),
                self.u, self.v)

        magnitude = np.sqrt(u*u + v*v)
        flags, barbs, halves, empty = self._find_tails(magnitude,
            self.rounding, **self.barb_increments)

        #Get the vertices for each of the barbs

        plot_barbs = self._make_barbs(u, v, flags, barbs, halves, empty,
            self._length, self._pivot, self.sizes, self.fill_empty, self.flip)
        self.set_verts(plot_barbs)

        #Set the color array
        if C is not None:
            self.set_array(c)

        #Update the offsets in case the masked data changed
        xy = np.hstack((x[:,np.newaxis], y[:,np.newaxis]))
        self._offsets = xy

    def set_offsets(self, xy):
        '''
        Set the offsets for the barb polygons.  This saves the offets passed in
        and actually sets version masked as appropriate for the existing U/V
        data. *offsets* should be a sequence.

        ACCEPTS: sequence of pairs of floats
        '''
        self.x = xy[:,0]
        self.y = xy[:,1]
        x,y,u,v = delete_masked_points(self.x.ravel(), self.y.ravel(), self.u,
            self.v)
        xy = np.hstack((x[:,np.newaxis], y[:,np.newaxis]))
        collections.PolyCollection.set_offsets(self, xy)
    set_offsets.__doc__ = collections.PolyCollection.set_offsets.__doc__

    barbs_doc = _barbs_doc


###
# Copyright (c) 2016, 2017 Diamond Light Source Ltd.
#
# All rights reserved. This program and the accompanying materials
# are made available under the terms of the Eclipse Public License v1.0
# which accompanies this distribution, and is available at
# http://www.eclipse.org/legal/epl-v10.html
#
# Contributors:
#    Tom Cobb - initial API and implementation and/or initial documentation
#    Gary Yendell - initial API and implementation and/or initial documentation
#    Charles Mita - initial API and implementation and/or initial documentation
#
###

from scanpointgenerator import CompoundGenerator, RectangularROI, CircularROI

MARKER_SIZE = 10

def plot_generator(gen, excluder=None, show_indexes=True):
    from matplotlib.patches import Rectangle, Circle
    import matplotlib.pyplot as plt
    import numpy as np
    from scipy import interpolate

    if excluder is not None:
        for roi in excluder.rois:
            overlay = plt.subplot(111, aspect='equal')
            if isinstance(roi, RectangularROI):
                overlay.add_patch(Rectangle(roi.start, roi.width, roi.height, fill=False))
            if isinstance(roi, CircularROI):
                overlay.add_patch(Circle(roi.centre, roi.radius, fill=False))

    if not isinstance(gen, CompoundGenerator):
        excluders = [] if excluder is None else [excluder]
        gen = CompoundGenerator([gen], excluders, [])
    gen.prepare()

    # points for spline generation
    x, y = [], []
    # capture points and indexes
    capx, capy, capi = [], [], []
    # segment start for colour changing
    starts = []
    for point in gen.iterator():
        # If lower is different from last then include it
        xlower = point.lower["x"]
        ylower = point.lower.get("y", 0)
        if len(x) == 0 or x[-1] != xlower or y[-1] != ylower:
            if len(x) != 0:
                # add in a tiny fractional distance
                xneg = x[-1] - xlower > 0
                yneg = y[-1] - ylower > 0
                xdiff = (x[-1] - x[-2]) * 0.01
                ydiff = (y[-1] - y[-2]) * 0.01
                for i in range(3):
                    x.append(x[-1] + xdiff)
                    y.append(y[-1] + ydiff)
                # add the padding on the input
                if xneg:
                    xdiff *= -1
                if yneg:
                    ydiff *= -1
                for i in reversed(range(3)):
                    x.append(xlower + xdiff * (i + 1))
                    y.append(ylower + ydiff * (i + 1))
            starts.append(len(x))
            x.append(xlower)
            y.append(ylower)
        # Add in capture points
        xpos = point.positions["x"]
        ypos = point.positions.get("y", 0)
        x.append(xpos)
        y.append(ypos)
        capx.append(xpos)
        capy.append(ypos)
        capi.append(point.indexes)
        # And upper point
        starts.append(len(x))
        x.append(point.upper["x"])
        y.append(point.upper.get("y", 0))

    # # Plot labels
    plt.xlabel("X (%s)" % gen.units["x"])
    if "y" in gen.units:
        plt.ylabel("Y (%s)" % gen.units["y"])
    else:
        plt.tick_params(left='off', labelleft='off')

    # Define curves parametrically
    x = np.array(x)
    y = np.array(y)
    t = np.zeros(len(x))
    t[1:] = np.sqrt((x[1:] - x[:-1])**2 + (y[1:] - y[:-1])**2)
    t = np.cumsum(t)
    t /= t[-1]
    tck, _ = interpolate.splprep([x, y], s=0)

    # Plot each line
    for i, start in enumerate(starts):
        if i + 1 < len(starts):
            end = starts[i+1]
        else:
            end = len(x) - 1
        tnew = np.linspace(t[start], t[end], num=1001, endpoint=True)
        sx, sy = interpolate.splev(tnew, tck)
        plt.plot(sx, sy, linewidth=2)

    # And the capture points
    plt.plot(capx, capy, linestyle="", marker="x", color="k",
             markersize=MARKER_SIZE)

    # And a start position
    plt.plot([x[0]], [y[0]], 'bo')
    plt.annotate("Start", (x[0], y[0]), xytext=(MARKER_SIZE/2, MARKER_SIZE/2),
                 textcoords='offset points')

    # And the indexes
    if show_indexes:
        for i, x, y in zip(capi, capx, capy):
            plt.annotate(i, (x, y), xytext=(MARKER_SIZE/2, MARKER_SIZE/2),
                         textcoords='offset points')
        #indexes = ["%s (size %d)" % z for z in zip(gen.index_names, gen.index_dims)]
        #plt.title("Dataset: [%s]" % (", ".join(indexes)))
    plt.show()


#!/usr/bin/env python
#
# Copyright 2013 Heiko Burau, Axel Huebl
#
# This file is part of PIConGPU.
#
# PIConGPU 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.
#
# PIConGPU 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 PIConGPU.
# If not, see <http://www.gnu.org/licenses/>.
#

#___________P A R A M E T E R S___________

omega_plasma = 6.718e13		# SI unit: 1/s
v_th = 1.0e8               	# SI unit: m/s
c = 2.9979e8			    # SI unit: m/s
delta_t = 2.5e-15           # SI unit: s
delta_z = c * delta_t       # SI unit: m

#_________________________________________

from numpy import *
from matplotlib import pyplot as plt
from matplotlib.ticker import FormatStrFormatter

data_trans = loadtxt("eField_zt_trans.dat")
data_long = loadtxt("eField_zt_long.dat")

N_z = len(data_trans[:,0])
N_t = len(data_trans[0,:])

omega_max = pi*(N_t-1)/(N_t*delta_t)/omega_plasma
k_max = pi * (N_z-1)/(N_z*delta_z)

# __________________transversal plot______________________

ax = plt.subplot(211, autoscale_on=False, xlim=(-k_max, k_max), ylim=(-1, 10))
ax.xaxis.set_major_formatter(FormatStrFormatter('%2.2e'))
ax.yaxis.set_major_formatter(FormatStrFormatter('%0.0f'))

plt.xlabel(r"$k [1/m]$")
plt.ylabel(r"$\omega / \omega_{pe} $")

data_trans = fft.fftshift(fft.fft2(data_trans))

plt.imshow(abs(data_trans), extent=(-k_max, k_max, -omega_max, omega_max), aspect='auto', interpolation='nearest')
plt.colorbar()

# plot analytical dispersion relation
x = linspace(-k_max, k_max, 200)
y = sqrt(c**2 * x**2 + omega_plasma**2)/omega_plasma
plt.plot(x, y, 'r--', linewidth=1)

# ___________________longitudinal plot_____________________

ax = plt.subplot(212, autoscale_on=False, xlim=(-k_max, k_max), ylim=(-1, 10))
ax.xaxis.set_major_formatter(FormatStrFormatter('%2.2e'))
ax.yaxis.set_major_formatter(FormatStrFormatter('%0.0f'))

plt.xlabel(r"$k [1/m]$")
plt.ylabel(r"$\omega / \omega_{pe} $")

data_long = fft.fftshift(fft.fft2(data_long))

plt.imshow(abs(data_long), extent=(-k_max, k_max, -omega_max, omega_max), aspect='auto', interpolation='nearest')
plt.colorbar()

# plot analytical dispersion relation
x = linspace(-k_max, k_max, 200)
y = sqrt(3 * v_th**2 * x**2 + omega_plasma**2)/omega_plasma
plt.plot(x, y, 'r--', linewidth=1)

plt.show()


"""
Linear Discriminant Analysis and Quadratic Discriminant Analysis
"""

# Authors: Clemens Brunner
#          Martin Billinger
#          Matthieu Perrot
#          Mathieu Blondel

# License: BSD 3-Clause

from __future__ import print_function
import warnings

import numpy as np
from scipy import linalg
from .externals.six import string_types
from .externals.six.moves import xrange

from .base import BaseEstimator, TransformerMixin, ClassifierMixin
from .linear_model.base import LinearClassifierMixin
from .covariance import ledoit_wolf, empirical_covariance, shrunk_covariance
from .utils.multiclass import unique_labels
from .utils import check_array, check_X_y
from .utils.validation import check_is_fitted
from .utils.fixes import bincount
from .preprocessing import StandardScaler

__all__ = ['LinearDiscriminantAnalysis', 'QuadraticDiscriminantAnalysis']

def _cov(X, shrinkage=None):
    """Estimate covariance matrix (using optional shrinkage).

    Parameters
    ----------
    X : array-like, shape (n_samples, n_features)
        Input data.

    shrinkage : string or float, optional
        Shrinkage parameter, possible values:
          - None or 'empirical': no shrinkage (default).
          - 'auto': automatic shrinkage using the Ledoit-Wolf lemma.
          - float between 0 and 1: fixed shrinkage parameter.

    Returns
    -------
    s : array, shape (n_features, n_features)
        Estimated covariance matrix.
    """
    shrinkage = "empirical" if shrinkage is None else shrinkage
    if isinstance(shrinkage, string_types):
        if shrinkage == 'auto':
            sc = StandardScaler()  # standardize features
            X = sc.fit_transform(X)
            s = sc.std_ * ledoit_wolf(X)[0] * sc.std_  # scale back
        elif shrinkage == 'empirical':
            s = empirical_covariance(X)
        else:
            raise ValueError('unknown shrinkage parameter')
    elif isinstance(shrinkage, float) or isinstance(shrinkage, int):
        if shrinkage < 0 or shrinkage > 1:
            raise ValueError('shrinkage parameter must be between 0 and 1')
        s = shrunk_covariance(empirical_covariance(X), shrinkage)
    else:
        raise TypeError('shrinkage must be of string or int type')
    return s

def _class_means(X, y):
    """Compute class means.

    Parameters
    ----------
    X : array-like, shape (n_samples, n_features)
        Input data.

    y : array-like, shape (n_samples,) or (n_samples, n_targets)
        Target values.

    Returns
    -------
    means : array-like, shape (n_features,)
        Class means.
    """
    means = []
    classes = np.unique(y)
    for group in classes:
        Xg = X[y == group, :]
        means.append(Xg.mean(0))
    return np.asarray(means)

def _class_cov(X, y, priors=None, shrinkage=None):
    """Compute class covariance matrix.

    Parameters
    ----------
    X : array-like, shape (n_samples, n_features)
        Input data.

    y : array-like, shape (n_samples,) or (n_samples, n_targets)
        Target values.

    priors : array-like, shape (n_classes,)
        Class priors.

    shrinkage : string or float, optional
        Shrinkage parameter, possible values:
          - None: no shrinkage (default).
          - 'auto': automatic shrinkage using the Ledoit-Wolf lemma.
          - float between 0 and 1: fixed shrinkage parameter.

    Returns
    -------
    cov : array-like, shape (n_features, n_features)
        Class covariance matrix.
    """
    classes = np.unique(y)
    covs = []
    for group in classes:
        Xg = X[y == group, :]
        covs.append(np.atleast_2d(_cov(Xg, shrinkage)))
    return np.average(covs, axis=0, weights=priors)

class LinearDiscriminantAnalysis(BaseEstimator, LinearClassifierMixin,
                                 TransformerMixin):
    """Linear Discriminant Analysis

    A classifier with a linear decision boundary, generated by fitting class
    conditional densities to the data and using Bayes' rule.

    The model fits a Gaussian density to each class, assuming that all classes
    share the same covariance matrix.

    The fitted model can also be used to reduce the dimensionality of the input
    by projecting it to the most discriminative directions.

    Parameters
    ----------
    solver : string, optional
        Solver to use, possible values:
          - 'svd': Singular value decomposition (default). Does not compute the
                covariance matrix, therefore this solver is recommended for
                data with a large number of features.
          - 'lsqr': Least squares solution, can be combined with shrinkage.
          - 'eigen': Eigenvalue decomposition, can be combined with shrinkage.

    shrinkage : string or float, optional
        Shrinkage parameter, possible values:
          - None: no shrinkage (default).
          - 'auto': automatic shrinkage using the Ledoit-Wolf lemma.
          - float between 0 and 1: fixed shrinkage parameter.

        Note that shrinkage works only with 'lsqr' and 'eigen' solvers.

    priors : array, optional, shape (n_classes,)
        Class priors.

    n_components : int, optional
        Number of components (< n_classes - 1) for dimensionality reduction.

    store_covariance : bool, optional
        Additionally compute class covariance matrix (default False).

    tol : float, optional
        Threshold used for rank estimation in SVD solver.

    Attributes
    ----------
    coef_ : array, shape (n_features,) or (n_classes, n_features)
        Weight vector(s).

    intercept_ : array, shape (n_features,)
        Intercept term.

    covariance_ : array-like, shape (n_features, n_features)
        Covariance matrix (shared by all classes).

    explained_variance_ratio_ : array, shape (n_components,)
        Percentage of variance explained by each of the selected components.
        If ``n_components`` is not set then all components are stored and the
        sum of explained variances is equal to 1.0. Only available when eigen
        solver is used.

    means_ : array-like, shape (n_classes, n_features)
        Class means.

    priors_ : array-like, shape (n_classes,)
        Class priors (sum to 1).

    scalings_ : array-like, shape (rank, n_classes - 1)
        Scaling of the features in the space spanned by the class centroids.

    xbar_ : array-like, shape (n_features,)
        Overall mean.

    classes_ : array-like, shape (n_classes,)
        Unique class labels.

    See also
    --------
    sklearn.discriminant_analysis.QuadraticDiscriminantAnalysis: Quadratic
        Discriminant Analysis

    Notes
    -----
    The default solver is 'svd'. It can perform both classification and
    transform, and it does not rely on the calculation of the covariance
    matrix. This can be an advantage in situations where the number of features
    is large. However, the 'svd' solver cannot be used with shrinkage.

    The 'lsqr' solver is an efficient algorithm that only works for
    classification. It supports shrinkage.

    The 'eigen' solver is based on the optimization of the between class
    scatter to within class scatter ratio. It can be used for both
    classification and transform, and it supports shrinkage. However, the
    'eigen' solver needs to compute the covariance matrix, so it might not be
    suitable for situations with a high number of features.

    Examples
    --------
    >>> import numpy as np
    >>> from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
    >>> X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
    >>> y = np.array([1, 1, 1, 2, 2, 2])
    >>> clf = LinearDiscriminantAnalysis()
    >>> clf.fit(X, y)
    LinearDiscriminantAnalysis(n_components=None, priors=None, shrinkage=None,
                  solver='svd', store_covariance=False, tol=0.0001)
    >>> print(clf.predict([[-0.8, -1]]))
    [1]
    """
    def __init__(self, solver='svd', shrinkage=None, priors=None,
                 n_components=None, store_covariance=False, tol=1e-4):
        self.solver = solver
        self.shrinkage = shrinkage
        self.priors = priors
        self.n_components = n_components
        self.store_covariance = store_covariance  # used only in svd solver
        self.tol = tol  # used only in svd solver

    def _solve_lsqr(self, X, y, shrinkage):
        """Least squares solver.

        The least squares solver computes a straightforward solution of the
        optimal decision rule based directly on the discriminant functions. It
        can only be used for classification (with optional shrinkage), because
        estimation of eigenvectors is not performed. Therefore, dimensionality
        reduction with the transform is not supported.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        y : array-like, shape (n_samples,) or (n_samples, n_classes)
            Target values.

        shrinkage : string or float, optional
            Shrinkage parameter, possible values:
              - None: no shrinkage (default).
              - 'auto': automatic shrinkage using the Ledoit-Wolf lemma.
              - float between 0 and 1: fixed shrinkage parameter.

        Notes
        -----
        This solver is based on [1]_, section 2.6.2, pp. 39-41.

        References
        ----------
        .. [1] R. O. Duda, P. E. Hart, D. G. Stork. Pattern Classification
           (Second Edition). John Wiley & Sons, Inc., New York, 2001. ISBN
           0-471-05669-3.
        """
        self.means_ = _class_means(X, y)
        self.covariance_ = _class_cov(X, y, self.priors_, shrinkage)
        self.coef_ = linalg.lstsq(self.covariance_, self.means_.T)[0].T
        self.intercept_ = (-0.5 * np.diag(np.dot(self.means_, self.coef_.T))
                           + np.log(self.priors_))

    def _solve_eigen(self, X, y, shrinkage):
        """Eigenvalue solver.

        The eigenvalue solver computes the optimal solution of the Rayleigh
        coefficient (basically the ratio of between class scatter to within
        class scatter). This solver supports both classification and
        dimensionality reduction (with optional shrinkage).

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        y : array-like, shape (n_samples,) or (n_samples, n_targets)
            Target values.

        shrinkage : string or float, optional
            Shrinkage parameter, possible values:
              - None: no shrinkage (default).
              - 'auto': automatic shrinkage using the Ledoit-Wolf lemma.
              - float between 0 and 1: fixed shrinkage constant.

        Notes
        -----
        This solver is based on [1]_, section 3.8.3, pp. 121-124.

        References
        ----------
        .. [1] R. O. Duda, P. E. Hart, D. G. Stork. Pattern Classification
           (Second Edition). John Wiley & Sons, Inc., New York, 2001. ISBN
           0-471-05669-3.
        """
        self.means_ = _class_means(X, y)
        self.covariance_ = _class_cov(X, y, self.priors_, shrinkage)

        Sw = self.covariance_  # within scatter
        St = _cov(X, shrinkage)  # total scatter
        Sb = St - Sw  # between scatter

        evals, evecs = linalg.eigh(Sb, Sw)
        self.explained_variance_ratio_ = np.sort(evals / np.sum(evals))[::-1]
        evecs = evecs[:, np.argsort(evals)[::-1]]  # sort eigenvectors
        # evecs /= np.linalg.norm(evecs, axis=0)  # doesn't work with numpy 1.6
        evecs /= np.apply_along_axis(np.linalg.norm, 0, evecs)

        self.scalings_ = evecs
        self.coef_ = np.dot(self.means_, evecs).dot(evecs.T)
        self.intercept_ = (-0.5 * np.diag(np.dot(self.means_, self.coef_.T))
                           + np.log(self.priors_))

    def _solve_svd(self, X, y):
        """SVD solver.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        y : array-like, shape (n_samples,) or (n_samples, n_targets)
            Target values.
        """
        n_samples, n_features = X.shape
        n_classes = len(self.classes_)

        self.means_ = _class_means(X, y)
        if self.store_covariance:
            self.covariance_ = _class_cov(X, y, self.priors_)

        Xc = []
        for idx, group in enumerate(self.classes_):
            Xg = X[y == group, :]
            Xc.append(Xg - self.means_[idx])

        self.xbar_ = np.dot(self.priors_, self.means_)

        Xc = np.concatenate(Xc, axis=0)

        # 1) within (univariate) scaling by with classes std-dev
        std = Xc.std(axis=0)
        # avoid division by zero in normalization
        std[std == 0] = 1.
        fac = 1. / (n_samples - n_classes)

        # 2) Within variance scaling
        X = np.sqrt(fac) * (Xc / std)
        # SVD of centered (within)scaled data
        U, S, V = linalg.svd(X, full_matrices=False)

        rank = np.sum(S > self.tol)
        if rank < n_features:
            warnings.warn("Variables are collinear.")
        # Scaling of within covariance is: V' 1/S
        scalings = (V[:rank] / std).T / S[:rank]

        # 3) Between variance scaling
        # Scale weighted centers
        X = np.dot(((np.sqrt((n_samples * self.priors_) * fac)) *
                    (self.means_ - self.xbar_).T).T, scalings)
        # Centers are living in a space with n_classes-1 dim (maximum)
        # Use SVD to find projection in the space spanned by the
        # (n_classes) centers
        _, S, V = linalg.svd(X, full_matrices=0)

        rank = np.sum(S > self.tol * S[0])
        self.scalings_ = np.dot(scalings, V.T[:, :rank])
        coef = np.dot(self.means_ - self.xbar_, self.scalings_)
        self.intercept_ = (-0.5 * np.sum(coef ** 2, axis=1)
                           + np.log(self.priors_))
        self.coef_ = np.dot(coef, self.scalings_.T)
        self.intercept_ -= np.dot(self.xbar_, self.coef_.T)

    def fit(self, X, y, store_covariance=None, tol=None):
        """Fit LinearDiscriminantAnalysis model according to the given
           training data and parameters.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        y : array, shape (n_samples,)
            Target values.
        """
        if store_covariance:
            warnings.warn("The parameter 'store_covariance' is deprecated as "
                          "of version 0.17 and will be removed in 0.19. The "
                          "parameter is no longer necessary because the value "
                          "is set via the estimator initialisation or "
                          "set_params method.", DeprecationWarning)
            self.store_covariance = store_covariance
        if tol:
            warnings.warn("The parameter 'tol' is deprecated as of version "
                          "0.17 and will be removed in 0.19. The parameter is "
                          "no longer necessary because the value is set via "
                          "the estimator initialisation or set_params method.",
                          DeprecationWarning)
            self.tol = tol
        X, y = check_X_y(X, y, ensure_min_samples=2, estimator=self)
        self.classes_ = unique_labels(y)

        if self.priors is None:  # estimate priors from sample
            _, y_t = np.unique(y, return_inverse=True)  # non-negative ints
            self.priors_ = bincount(y_t) / float(len(y))
        else:
            self.priors_ = np.asarray(self.priors)

        if (self.priors_ < 0).any():
            raise ValueError("priors must be non-negative")
        if self.priors_.sum() != 1:
            warnings.warn("The priors do not sum to 1. Renormalizing",
                          UserWarning)
            self.priors_ = self.priors_ / self.priors_.sum()

        if self.solver == 'svd':
            if self.shrinkage is not None:
                raise NotImplementedError('shrinkage not supported')
            self._solve_svd(X, y)
        elif self.solver == 'lsqr':
            self._solve_lsqr(X, y, shrinkage=self.shrinkage)
        elif self.solver == 'eigen':
            self._solve_eigen(X, y, shrinkage=self.shrinkage)
        else:
            raise ValueError("unknown solver {} (valid solvers are 'svd', "
                             "'lsqr', and 'eigen').".format(self.solver))
        if self.classes_.size == 2:  # treat binary case as a special case
            self.coef_ = np.array(self.coef_[1, :] - self.coef_[0, :], ndmin=2)
            self.intercept_ = np.array(self.intercept_[1] - self.intercept_[0],
                                       ndmin=1)
        return self

    def transform(self, X):
        """Project data to maximize class separation.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Input data.

        Returns
        -------
        X_new : array, shape (n_samples, n_components)
            Transformed data.
        """
        if self.solver == 'lsqr':
            raise NotImplementedError("transform not implemented for 'lsqr' "
                                      "solver (use 'svd' or 'eigen').")
        check_is_fitted(self, ['xbar_', 'scalings_'], all_or_any=any)

        X = check_array(X)
        if self.solver == 'svd':
            X_new = np.dot(X - self.xbar_, self.scalings_)
        elif self.solver == 'eigen':
            X_new = np.dot(X, self.scalings_)
        n_components = X.shape[1] if self.n_components is None \
            else self.n_components
        return X_new[:, :n_components]

    def predict_proba(self, X):
        """Estimate probability.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Input data.

        Returns
        -------
        C : array, shape (n_samples, n_classes)
            Estimated probabilities.
        """
        prob = self.decision_function(X)
        prob *= -1
        np.exp(prob, prob)
        prob += 1
        np.reciprocal(prob, prob)
        if len(self.classes_) == 2:  # binary case
            return np.column_stack([1 - prob, prob])
        else:
            # OvR normalization, like LibLinear's predict_probability
            prob /= prob.sum(axis=1).reshape((prob.shape[0], -1))
            return prob

    def predict_log_proba(self, X):
        """Estimate log probability.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Input data.

        Returns
        -------
        C : array, shape (n_samples, n_classes)
            Estimated log probabilities.
        """
        return np.log(self.predict_proba(X))

class QuadraticDiscriminantAnalysis(BaseEstimator, ClassifierMixin):
    """
    Quadratic Discriminant Analysis

    A classifier with a quadratic decision boundary, generated
    by fitting class conditional densities to the data
    and using Bayes' rule.

    The model fits a Gaussian density to each class.

    Parameters
    ----------
    priors : array, optional, shape = [n_classes]
        Priors on classes

    reg_param : float, optional
        Regularizes the covariance estimate as
        ``(1-reg_param)*Sigma + reg_param*np.eye(n_features)``

    Attributes
    ----------
    covariances_ : list of array-like, shape = [n_features, n_features]
        Covariance matrices of each class.

    means_ : array-like, shape = [n_classes, n_features]
        Class means.

    priors_ : array-like, shape = [n_classes]
        Class priors (sum to 1).

    rotations_ : list of arrays
        For each class k an array of shape [n_features, n_k], with
        ``n_k = min(n_features, number of elements in class k)``
        It is the rotation of the Gaussian distribution, i.e. its
        principal axis.

    scalings_ : list of arrays
        For each class k an array of shape [n_k]. It contains the scaling
        of the Gaussian distributions along its principal axes, i.e. the
        variance in the rotated coordinate system.

    store_covariances : boolean
        If True the covariance matrices are computed and stored in the
        `self.covariances_` attribute.

    tol : float, optional, default 1.0e-4
        Threshold used for rank estimation.

    Examples
    --------
    >>> from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
    >>> import numpy as np
    >>> X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
    >>> y = np.array([1, 1, 1, 2, 2, 2])
    >>> clf = QuadraticDiscriminantAnalysis()
    >>> clf.fit(X, y)
    ... # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    QuadraticDiscriminantAnalysis(priors=None, reg_param=0.0,
                                  store_covariances=False, tol=0.0001)
    >>> print(clf.predict([[-0.8, -1]]))
    [1]

    See also
    --------
    sklearn.discriminant_analysis.LinearDiscriminantAnalysis: Linear
        Discriminant Analysis
    """

    def __init__(self, priors=None, reg_param=0., store_covariances=False,
                 tol=1.0e-4):
        self.priors = np.asarray(priors) if priors is not None else None
        self.reg_param = reg_param
        self.store_covariances = store_covariances
        self.tol = tol

    def fit(self, X, y, store_covariances=None, tol=None):
        """Fit the model according to the given training data and parameters.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            Training vector, where n_samples in the number of samples and
            n_features is the number of features.

        y : array, shape = [n_samples]
            Target values (integers)
        """
        if store_covariances:
            warnings.warn("The parameter 'store_covariances' is deprecated as "
                          "of version 0.17 and will be removed in 0.19. The "
                          "parameter is no longer necessary because the value "
                          "is set via the estimator initialisation or "
                          "set_params method.", DeprecationWarning)
            self.store_covariances = store_covariances
        if tol:
            warnings.warn("The parameter 'tol' is deprecated as of version "
                          "0.17 and will be removed in 0.19. The parameter is "
                          "no longer necessary because the value is set via "
                          "the estimator initialisation or set_params method.",
                          DeprecationWarning)
            self.tol = tol
        X, y = check_X_y(X, y)
        self.classes_, y = np.unique(y, return_inverse=True)
        n_samples, n_features = X.shape
        n_classes = len(self.classes_)
        if n_classes < 2:
            raise ValueError('y has less than 2 classes')
        if self.priors is None:
            self.priors_ = bincount(y) / float(n_samples)
        else:
            self.priors_ = self.priors

        cov = None
        if self.store_covariances:
            cov = []
        means = []
        scalings = []
        rotations = []
        for ind in xrange(n_classes):
            Xg = X[y == ind, :]
            meang = Xg.mean(0)
            means.append(meang)
            if len(Xg) == 1:
                raise ValueError('y has only 1 sample in class %s, covariance '
                                 'is ill defined.' % str(self.classes_[ind]))
            Xgc = Xg - meang
            # Xgc = U * S * V.T
            U, S, Vt = np.linalg.svd(Xgc, full_matrices=False)
            rank = np.sum(S > self.tol)
            if rank < n_features:
                warnings.warn("Variables are collinear")
            S2 = (S ** 2) / (len(Xg) - 1)
            S2 = ((1 - self.reg_param) * S2) + self.reg_param
            if self.store_covariances:
                # cov = V * (S^2 / (n-1)) * V.T
                cov.append(np.dot(S2 * Vt.T, Vt))
            scalings.append(S2)
            rotations.append(Vt.T)
        if self.store_covariances:
            self.covariances_ = cov
        self.means_ = np.asarray(means)
        self.scalings_ = scalings
        self.rotations_ = rotations
        return self

    def _decision_function(self, X):
        check_is_fitted(self, 'classes_')

        X = check_array(X)
        norm2 = []
        for i in range(len(self.classes_)):
            R = self.rotations_[i]
            S = self.scalings_[i]
            Xm = X - self.means_[i]
            X2 = np.dot(Xm, R * (S ** (-0.5)))
            norm2.append(np.sum(X2 ** 2, 1))
        norm2 = np.array(norm2).T   # shape = [len(X), n_classes]
        u = np.asarray([np.sum(np.log(s)) for s in self.scalings_])
        return (-0.5 * (norm2 + u) + np.log(self.priors_))

    def decision_function(self, X):
        """Apply decision function to an array of samples.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            Array of samples (test vectors).

        Returns
        -------
        C : array, shape = [n_samples, n_classes] or [n_samples,]
            Decision function values related to each class, per sample.
            In the two-class case, the shape is [n_samples,], giving the
            log likelihood ratio of the positive class.
        """
        dec_func = self._decision_function(X)
        # handle special case of two classes
        if len(self.classes_) == 2:
            return dec_func[:, 1] - dec_func[:, 0]
        return dec_func

    def predict(self, X):
        """Perform classification on an array of test vectors X.

        The predicted class C for each sample in X is returned.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = [n_samples]
        """
        d = self._decision_function(X)
        y_pred = self.classes_.take(d.argmax(1))
        return y_pred

    def predict_proba(self, X):
        """Return posterior probabilities of classification.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            Array of samples/test vectors.

        Returns
        -------
        C : array, shape = [n_samples, n_classes]
            Posterior probabilities of classification per class.
        """
        values = self._decision_function(X)
        # compute the likelihood of the underlying gaussian models
        # up to a multiplicative constant.
        likelihood = np.exp(values - values.max(axis=1)[:, np.newaxis])
        # compute posterior probabilities
        return likelihood / likelihood.sum(axis=1)[:, np.newaxis]

    def predict_log_proba(self, X):
        """Return posterior probabilities of classification.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            Array of samples/test vectors.

        Returns
        -------
        C : array, shape = [n_samples, n_classes]
            Posterior log-probabilities of classification per class.
        """
        # XXX : can do better to avoid precision overflows
        probas_ = self.predict_proba(X)
        return np.log(probas_)


# Author: Gael Varoquaux
# License: BSD 3 clause

import numpy as np
import scipy.sparse as sp

from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_false
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_not_equal
from sklearn.utils.testing import assert_raises

from sklearn.base import BaseEstimator, clone, is_classifier
from sklearn.svm import SVC
from sklearn.pipeline import Pipeline
from sklearn.grid_search import GridSearchCV
from sklearn.utils import deprecated

#############################################################################
# A few test classes
class MyEstimator(BaseEstimator):

    def __init__(self, l1=0, empty=None):
        self.l1 = l1
        self.empty = empty

class K(BaseEstimator):
    def __init__(self, c=None, d=None):
        self.c = c
        self.d = d

class T(BaseEstimator):
    def __init__(self, a=None, b=None):
        self.a = a
        self.b = b

class DeprecatedAttributeEstimator(BaseEstimator):
    def __init__(self, a=None, b=None):
        self.a = a
        if b is not None:
            DeprecationWarning("b is deprecated and renamed 'a'")
            self.a = b

    @property
    @deprecated("Parameter 'b' is deprecated and renamed to 'a'")
    def b(self):
        return self._b

class Buggy(BaseEstimator):
    " A buggy estimator that does not set its parameters right. "

    def __init__(self, a=None):
        self.a = 1

class NoEstimator(object):
    def __init__(self):
        pass

    def fit(self, X=None, y=None):
        return self

    def predict(self, X=None):
        return None

class VargEstimator(BaseEstimator):
    """Sklearn estimators shouldn't have vargs."""
    def __init__(self, *vargs):
        pass

#############################################################################
# The tests

def test_clone():
    # Tests that clone creates a correct deep copy.
    # We create an estimator, make a copy of its original state
    # (which, in this case, is the current state of the estimator),
    # and check that the obtained copy is a correct deep copy.

    from sklearn.feature_selection import SelectFpr, f_classif

    selector = SelectFpr(f_classif, alpha=0.1)
    new_selector = clone(selector)
    assert_true(selector is not new_selector)
    assert_equal(selector.get_params(), new_selector.get_params())

    selector = SelectFpr(f_classif, alpha=np.zeros((10, 2)))
    new_selector = clone(selector)
    assert_true(selector is not new_selector)

def test_clone_2():
    # Tests that clone doesn't copy everything.
    # We first create an estimator, give it an own attribute, and
    # make a copy of its original state. Then we check that the copy doesn't
    # have the specific attribute we manually added to the initial estimator.

    from sklearn.feature_selection import SelectFpr, f_classif

    selector = SelectFpr(f_classif, alpha=0.1)
    selector.own_attribute = "test"
    new_selector = clone(selector)
    assert_false(hasattr(new_selector, "own_attribute"))

def test_clone_buggy():
    # Check that clone raises an error on buggy estimators.
    buggy = Buggy()
    buggy.a = 2
    assert_raises(RuntimeError, clone, buggy)

    no_estimator = NoEstimator()
    assert_raises(TypeError, clone, no_estimator)

    varg_est = VargEstimator()
    assert_raises(RuntimeError, clone, varg_est)

def test_clone_empty_array():
    # Regression test for cloning estimators with empty arrays
    clf = MyEstimator(empty=np.array([]))
    clf2 = clone(clf)
    assert_array_equal(clf.empty, clf2.empty)

    clf = MyEstimator(empty=sp.csr_matrix(np.array([[0]])))
    clf2 = clone(clf)
    assert_array_equal(clf.empty.data, clf2.empty.data)

def test_clone_nan():
    # Regression test for cloning estimators with default parameter as np.nan
    clf = MyEstimator(empty=np.nan)
    clf2 = clone(clf)

    assert_true(clf.empty is clf2.empty)

def test_repr():
    # Smoke test the repr of the base estimator.
    my_estimator = MyEstimator()
    repr(my_estimator)
    test = T(K(), K())
    assert_equal(
        repr(test),
        "T(a=K(c=None, d=None), b=K(c=None, d=None))"
    )

    some_est = T(a=["long_params"] * 1000)
    assert_equal(len(repr(some_est)), 415)

def test_str():
    # Smoke test the str of the base estimator
    my_estimator = MyEstimator()
    str(my_estimator)

def test_get_params():
    test = T(K(), K())

    assert_true('a__d' in test.get_params(deep=True))
    assert_true('a__d' not in test.get_params(deep=False))

    test.set_params(a__d=2)
    assert_true(test.a.d == 2)
    assert_raises(ValueError, test.set_params, a__a=2)

def test_get_params_deprecated():
    # deprecated attribute should not show up as params
    est = DeprecatedAttributeEstimator(a=1)

    assert_true('a' in est.get_params())
    assert_true('a' in est.get_params(deep=True))
    assert_true('a' in est.get_params(deep=False))

    assert_true('b' not in est.get_params())
    assert_true('b' not in est.get_params(deep=True))
    assert_true('b' not in est.get_params(deep=False))

def test_is_classifier():
    svc = SVC()
    assert_true(is_classifier(svc))
    assert_true(is_classifier(GridSearchCV(svc, {'C': [0.1, 1]})))
    assert_true(is_classifier(Pipeline([('svc', svc)])))
    assert_true(is_classifier(Pipeline([('svc_cv',
                              GridSearchCV(svc, {'C': [0.1, 1]}))])))

def test_set_params():
    # test nested estimator parameter setting
    clf = Pipeline([("svc", SVC())])
    # non-existing parameter in svc
    assert_raises(ValueError, clf.set_params, svc__stupid_param=True)
    # non-existing parameter of pipeline
    assert_raises(ValueError, clf.set_params, svm__stupid_param=True)
    # we don't currently catch if the things in pipeline are estimators
    # bad_pipeline = Pipeline([("bad", NoEstimator())])
    # assert_raises(AttributeError, bad_pipeline.set_params,
    #               bad__stupid_param=True)

def test_score_sample_weight():
    from sklearn.tree import DecisionTreeClassifier
    from sklearn.tree import DecisionTreeRegressor
    from sklearn import datasets

    rng = np.random.RandomState(0)

    # test both ClassifierMixin and RegressorMixin
    estimators = [DecisionTreeClassifier(max_depth=2),
                  DecisionTreeRegressor(max_depth=2)]
    sets = [datasets.load_iris(),
            datasets.load_boston()]

    for est, ds in zip(estimators, sets):
        est.fit(ds.data, ds.target)
        # generate random sample weights
        sample_weight = rng.randint(1, 10, size=len(ds.target))
        # check that the score with and without sample weights are different
        assert_not_equal(est.score(ds.data, ds.target),
                         est.score(ds.data, ds.target,
                                   sample_weight=sample_weight),
                         msg="Unweighted and weighted scores "
                             "are unexpectedly equal")


#!/usr/bin/env python

from __future__ import division, print_function, absolute_import
import sys
import os

PPRZ_SRC = os.getenv("PAPARAZZI_SRC", "../../../..")
sys.path.append(PPRZ_SRC + "/sw/lib/python")

from pprz_math.geodetic import *
from pprz_math.algebra import DoubleRMat, DoubleEulers, DoubleVect3
from math import radians, degrees, tan
import matplotlib.pyplot as plt
import numpy as np

# Origin at ENAC
UTM_EAST0 = 377349  # in m
UTM_NORTH0 = 4824583  # in m
UTM_ZONE0 = 31
ALT0 = 147.000  # in m

utm_origin = UtmCoor_d(north=UTM_NORTH0, east=UTM_EAST0, alt=ALT0, zone=UTM_ZONE0)
print("origin %s" % utm_origin)

lla_origin = utm_origin.to_lla()
ecef_origin = lla_origin.to_ecef()
ltp_origin = ecef_origin.to_ltp_def()
print(ltp_origin)

# convergence angle to "true north" is approx 1 deg here
earth_radius = 6378137.0
n = 0.9996 * earth_radius
UTM_DELTA_EAST = 500000.
dist_to_meridian = utm_origin.east - UTM_DELTA_EAST
conv = dist_to_meridian / n * tan(lla_origin.lat)
# or (middle meridian of UTM zone 31 is at 3deg)
#conv = atan(tan(lla_origin.lon - radians(3))*sin(lla_origin.lat))
print("approx. convergence angle (north error compared to meridian): %f deg" % degrees(conv))
# Rotation matrix to correct for "true north"
R = DoubleEulers(psi=-conv).to_rmat()

# calculate ENU coordinates for 100 points in 100m distance
nb_points = 100
dist_points = 100
enu_res = np.zeros((nb_points, 2))
enu_res_c = np.zeros((nb_points, 2))
utm_res = np.zeros((nb_points, 2))
for i in range(0, nb_points):
    utm = UtmCoor_d()
    utm.north = i * dist_points + utm_origin.north
    utm.east = i * dist_points+ utm_origin.east
    utm.alt = utm_origin.alt
    utm.zone = utm_origin.zone
    #print(utm)
    utm_res[i, 0] = utm.east - utm_origin.east
    utm_res[i, 1] = utm.north - utm_origin.north
    lla = utm.to_lla()
    #print(lla)
    ecef = lla.to_ecef()
    enu = ecef.to_enu(ltp_origin)
    enu_res[i, 0] = enu.x
    enu_res[i, 1] = enu.y
    enu_c = R * DoubleVect3(enu.x, enu.y, enu.z)
    enu_res_c[i, 0] = enu_c.x
    enu_res_c[i, 1] = enu_c.y
    #print(enu)

dist = np.linalg.norm(utm_res, axis=1)
error = np.linalg.norm(utm_res - enu_res, axis=1)
error_c = np.linalg.norm(utm_res - enu_res_c, axis=1)

plt.figure(1)
plt.subplot(311)
plt.title("utm vs. enu")
plt.plot(enu_res[:, 0], enu_res[:, 1], 'g', label="ENU")
plt.plot(utm_res[:, 0], utm_res[:, 1], 'r', label="UTM")
plt.ylabel("y/north [m]")
plt.xlabel("x/east [m]")
plt.legend(loc='upper left')
plt.subplot(312)
plt.plot(dist, error, 'r')
plt.xlabel("dist from origin [m]")
plt.ylabel("error [m]")
plt.subplot(313)
plt.plot(dist, error_c, 'r')
plt.xlabel("dist from origin [m]")
plt.ylabel("error with north fix [m]")

plt.show()


# -*- coding: utf-8 -*-
# <nbformat>3.0</nbformat>

# <headingcell level=1>

# Goals

# <markdowncell>

# **To practice using generators to yield geographical entities of various types.**
#
# Generators are a bit complicated, and I won't try to explain all the intricacies here.  I will show you how to use `yield` in a function definition to return a generator.  From [Definition of a generator](http://docs.python.org/2/glossary.html#term-generator):
#
#
A function which returns an iterator. It looks like a normal function except that it contains yield statements for producing a series a values usable in a for-loop or that can be retrieved one at a time with the next() function. Each yield temporarily suspends processing, remembering the location execution state (including local variables and pending try-statements). When the generator resumes, it picks-up where it left-off (in contrast to functions which start fresh on every invocation)

#
# For some background on Python generators:
#
# * [iterator - The Python yield keyword explained - Stack Overflow](http://stackoverflow.com/questions/231767/the-python-yield-keyword-explained/231855#231855)
# * [Improve Your Python: 'yield' and Generators Explained](http://www.jeffknupp.com/blog/2013/04/07/improve-your-python-yield-and-generators-explained/)
#
# Why use generators: http://stackoverflow.com/a/102632/7782
#
#
Generators are good for calculating large sets of results (in particular calculations involving loops themselves) where you don't know if you are going to need all results, or where you don't want to allocate the memory for all results at the same time.

#
# Also, let's also practice using [itertools.islice](http://www.python.org/doc//current/library/itertools.html#itertools.islice) and [enumerate](http://docs.python.org/2/library/functions.html#enumerate) -- two of my favorite constructions in Python

# <markdowncell>

# From http://api.census.gov/data/2010/sf1/geo.html, geographic entities we are specifically interested in this exercise:
#
# * state-county
# * state-county-tract
#
# * state-place
# * state-metropolitan statistical area/micropolitan statistical area
# * state-metropolitan statistical area/micropolitan statistical area-metropolitan division
# * state-combined statistical area

# <codecell>

# usual imports for numpy, pandas, matplotlib

import numpy as np
import matplotlib.pyplot as plt
from pandas import DataFrame, Series, Index
import pandas as pd

# <codecell>

# check that CENSUS_KEY is defined

import census
import us

import settings
assert settings.CENSUS_KEY is not None

# <codecell>

# instantiate our Census object

c = census.Census(key=settings.CENSUS_KEY)

# <headingcell level=1>

# A bit of warmup with Generators

# <codecell>

import string
print list(string.lowercase)

# <codecell>

def abcs():
    alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k',
                'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
                'w', 'x', 'y', 'z']
    """a generator that returns """
    for letter in alphabet:
        yield letter

# a generator that gives you the letters of the alphabet a letter at a time
say_abcs =  abcs()

# <codecell>

# run this line over and over again to see the letters one at a time
say_abcs.next()

# <codecell>

# you can use list to grab all the items in an iterator.  But be careful if the number
# of items is large or even infinite!  In this case, we're ok

list(abcs())

# <markdowncell>

# Demonstration of how to use [enumerate](http://docs.python.org/2/library/functions.html#enumerate):
#
#
Return an enumerate object. sequence must be a sequence, an iterator, or some other object which supports iteration. The next() method of the iterator returned by enumerate() returns a tuple containing a count (from start which defaults to 0) and the values obtained from iterating over sequence

#

# <codecell>

for (i, letter) in enumerate(abcs()):
    print i, letter

# <markdowncell>

# You can use itertools.islice [itertools.islice](http://www.python.org/doc//current/library/itertools.html#itertools.islice) to return parts of the iterator.
#
#
Make an iterator that returns selected elements from the iterable. If start is non-zero, then elements from the iterable are skipped until start is reached. Afterward, elements are returned consecutively unless step is set higher than one which results in items being skipped. If stop is None, then iteration continues until the iterator is exhausted, if at all; otherwise, it stops at the specified position. Unlike regular slicing, islice() does not support negative values for start, stop, or step. Can be used to extract related fields from data where the internal structure has been flattened (for example, a multi-line report may list a name field on every third line).

# <codecell>

# let's get the first 10 letters of the alphabet

from itertools import islice
list(islice(abcs(), 10))

# <codecell>

# you can use None to get all items in islice
# from docs: "If stop is None, then iteration continues until the iterator is exhausted,"

list(islice(abcs(), None))

# <codecell>

# itertools.count can in principle generate an infinite sequence
# http://www.python.org/doc//current/library/itertools.html#itertools.count

from itertools import count

# count starting zero
my_counter = count(0)

# <codecell>

# try it out
my_counter.next()

# <codecell>

# DON'T do list(count(0))  -> you'll be trying to generate an infinite list
# but use an upper limit

list(islice(count(0),10))

# <codecell>

# start, stop
list(islice(count(),1,3))

# <headingcell level=1>

# Generator for US Counties

# <codecell>

# get the syntax down for getting counties from CA -- so that then we can use it later

r = c.sf1.get('NAME,P0010001', geo={'for':'county:*',
                                'in':'state:{fips}'.format(fips=us.states.CA.fips)})
r[:5]

# <markdowncell>

# With the census API, you can get the counties with one single call to the census API or state-by-state.  The `counties` generator below takes the first approach while `counties2` takes the second approach. Although `counties` is more efficient in most cases I can think of, it will be useful to know how to do calls on a state-by-state basis.  For example, when we to query on a census tract level or below, we will need to work on a state-by-state basis.

# <codecell>

def counties(variables='NAME'):
    """ask for all the states"""

    # tabulate a set of fips codes for the states
    states_fips = set([s.fips for s in us.states.STATES])

    geo={'for':'county:*',
             'in':'state:*'}
    for county in c.sf1.get(variables, geo=geo):
        # eliminate counties whose states aren't in a state or DC
        if county['state'] in states_fips:
            yield county

def counties2(variables='NAME'):
    """generator for all counties"""

    # since we can get all the counties in one call,
    # this function is for demonstrating the use of walking through
    # the states to get at the counties

    for state in us.states.STATES:
        geo={'for':'county:*',
             'in':'state:{fips}'.format(fips=state.fips)}
        for county in c.sf1.get(variables, geo=geo):
            yield county

# <codecell>

counties_list = list(counties('NAME,P0010001'))

# <codecell>

# add up the population to make sure we have the total right
counties_df = DataFrame(counties_list)
counties_df.P0010001 = counties_df.P0010001.astype('int')
counties_df.P0010001.sum()

# <markdowncell>

# One reason for writing all the counties in the form of a Python generator is tha you can easily control the number of counties we work with at any given time -- and then easily scaling out to get all of them.

# <codecell>

# make a list of the first ten counties

from itertools import islice
list(islice(counties2(),10))

# <headingcell level=1>

# Generator for Census Tracts

# <markdowncell>

# The following generator loops through all the states to get at the individual counties to then get at the census tracts.

# <codecell>

def tracts(variables='NAME'):
    for state in us.states.STATES:

        # handy to print out state to monitor progress
        print state.fips, state
        counties_in_state={'for':'county:*',
             'in':'state:{fips}'.format(fips=state.fips)}

        for county in c.sf1.get('NAME', geo=counties_in_state):

            # print county['state'], county['NAME']
            tracts_in_county = {'for':'tract:*',
              'in': 'state:{s_fips} county:{c_fips}'.format(s_fips=state.fips,
                                                            c_fips=county['county'])}

            for tract in c.sf1.get(variables,geo=tracts_in_county):
                yield tract

# <codecell>

r = list(islice(tracts('NAME,P0010001'),10))
tracts_df = DataFrame(r)
tracts_df.P0010001 = tracts_df.P0010001.astype('int')
tracts_df['FIPS'] = tracts_df.apply(lambda s: s['state']+s['county']+s['tract'], axis=1)
print "number of tracts", len(tracts_df)
print "total pop", tracts_df.P0010001.sum()
tracts_df.head()

# <markdowncell>

# Good to save the DataFrame so we can load up the census tracts without having call the census api again.
#
# I/O: http://pandas.pydata.org/pandas-docs/dev/io.html
#
# Today, we'll use [pickle format](http://docs.python.org/2/library/pickle.html) and look at other formats.

# <codecell>

TRACT_FILE_PICKLE = "tracts.pickle"

# UNCOMMENT THIS LINE TO SAVE YOUR FILE
# tracts_df.to_pickle(TRACT_FILE_PICKLE)

# <markdowncell>

# Let's read the DataFrame from disk to confirm that we were able to save the file properly.

# <codecell>

df = pd.read_pickle(TRACT_FILE_PICKLE)
df.head()

# <codecell>

# UNCOMMENT TO DO COMPARISON
# you can compare the saved file to the file from disk
# np.all(tracts_df == df)