<commit_before># -*- coding: utf-8 -*-
#
# plotting.py
#
# purpose:  Plotting methods for ctd DataFrame.
# author:   Filipe P. A. Fernandes
# e-mail:   ocefpaf@gmail
# web:      http://ocefpaf.tiddlyspot.com/
# created:  23-Jul-2013
# modified: Tue 23 Jul 2013 01:31:06 PM BRT
#
# obs:
#

# Scientific stack.
import gsw
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
import mpl_toolkits.axisartist as AA

from pandas import Series
from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1 import host_subplot

from utilities import extrap1d

__all__ = ['get_maxdepth',
           'extrap_sec',
           'gen_topomask',
           'plot',
           'plot_vars',
           'plot_section']

def get_maxdepth(self):
    valid_last_depth = self.apply(Series.notnull).values.T
    return np.float_(self.index.values * valid_last_depth).max(axis=1)

def extrap_sec(data, dist, depth, w1=1., w2=0):
    """
    Extrapolates `data` to zones where the shallow stations are shadowed by
    the deep stations.  The shadow region usually cannot be extrapolates via
    linear interpolation.

    The extrapolation is applied using the gradients of the `data` at a certain
    level.

    Parameters
    ----------
    data : array_like
          Data to be extrapolated
    dist : array_like
           Stations distance
    fd : float
         Decay factor [0-1]

    Returns
    -------
    Sec_extrap : array_like
                 Extrapolated variable

    Examples
    --------
    >>> import data, dist, z
    >>> Sec_extrap = extrap_sec(data, dist, z, fd=1.)
    """
    new_data1 = []
    for row in data:
        mask = ~np.isnan(row)
        if mask.any():
            y = row[mask]
            if y.size == 1:
                row = np.repeat(y, len(mask))
            else:
                x = dist[mask]
                f_i = interp1d(x, y)
                f_x = extrap1d(f_i)
                row = f_x(dist)
        new_data1.append(row)

    new_data2 = []
    for col in data.T:
        mask = ~np.isnan(col)
        if mask.any():
            y = col[mask]
            if y.size == 1:
                col = np.repeat(y, len(mask))
            else:
                z = depth[mask]
                f_i = interp1d(z, y)
                f_z = extrap1d(f_i)
                col = f_z(depth)
        new_data2.append(col)

    new_data = np.array(new_data1) * w1 + np.array(new_data2).T * w2
    return new_data

def gen_topomask(h, lon, lat, dx=1., kind='linear', plot=False):
    """
    Generates a topography mask from an oceanographic transect taking the
    deepest CTD scan as the depth of each station.

    Inputs
    ------
    h : array
        Pressure of the deepest CTD scan for each station [dbar].
    lons : array
           Longitude of each station [decimal degrees east].
    lat : Latitude of each station. [decimal degrees north].
    dx : float
         Horizontal resolution of the output arrays [km].
    kind : string, optional
           Type of the interpolation to be performed.
           See scipy.interpolate.interp1d documentation for details.
    plot : bool
           Whether to plot mask for visualization.

    Outputs
    -------
    xm : array
         Horizontal distances [km].
    hm : array
         Local depth [m].

    Examples
    --------
    >>> import gsw
    >>> import df  # FIXME: Add a dataset.
    >>> h = df.get_maxdepth()
    >>> # TODO: method to output distance.
    >>> x = np.append(0, np.cumsum(gsw.distance(df.lon, df.lat)[0] / 1e3))
    >>> xm, hm = gen_topomask(h, df.lon, df.lat, dx=1., kind='linear')
    >>> fig, ax = plt.subplots()
    >>> ax.plot(xm, hm, 'k', linewidth=1.5)
    >>> ax.plot(x, h, 'ro')
    >>> ax.set_xlabel('Distance [km]')
    >>> ax.set_ylabel('Depth [m]')
    >>> ax.grid(True)
    >>> plt.show()

    Author
    ------
    André Palóczy Filho (paloczy@gmail.com) --  October/2012
    """

    h, lon, lat = map(np.asanyarray, (h, lon, lat))
    # Distance in km.
    x = np.append(0, np.cumsum(gsw.distance(lon, lat)[0] / 1e3))
    h = -gsw.z_from_p(h, lat.mean())
    Ih = interp1d(x, h, kind=kind, bounds_error=False, fill_value=h[-1])
    xm = np.arange(0, x.max() + dx, dx)
    hm = Ih(xm)

    return xm, hm

def plot(self, **kwds):
    """
    Plot a CTD variable against the index (pressure or depth).
    """
    fig, ax = plt.subplots()
    ax.plot(self.values, self.index, **kwds)
    ax.set_ylabel(self.index.name)
    ax.set_xlabel(self.name)
    ax.invert_yaxis()
    offset = 0.01
    x1, x2 = ax.get_xlim()[0] - offset, ax.get_xlim()[1] + offset
    ax.set_xlim(x1, x2)
    return fig, ax

def plot_vars(self, variables=None, **kwds):
    """
    Plot CTD temperature and salinity.
    """
    fig = plt.figure(figsize=(8, 10))
    ax0 = host_subplot(111, axes_class=AA.Axes)
    ax1 = ax0.twiny()

    # Axis location.
    host_new_axis = ax0.get_grid_helper().new_fixed_axis
    ax0.axis["bottom"] = host_new_axis(loc="top", axes=ax0, offset=(0, 0))
    par_new_axis = ax1.get_grid_helper().new_fixed_axis
    ax1.axis["top"] = par_new_axis(loc="bottom", axes=ax1, offset=(0, 0))

    ax0.plot(self[variables[0]], self.index, 'r.', label='Temperature')
    ax1.plot(self[variables[1]], self.index, 'b.', label='Salinity')

    ax0.set_ylabel("Pressure [dbar]")
    ax0.set_xlabel(u"Temperature [\u00b0C]")
    ax1.set_xlabel("Salinity [kg g$^{-1}$]")
    ax1.invert_yaxis()

    try:  # FIXME with metadata.
        fig.suptitle(r"Station %s profile" % self.name)
    except AttributeError:
        pass

    ax0.legend(shadow=True, fancybox=True,
               numpoints=1, loc='lower right')

    offset = 0.01
    x1, x2 = ax0.get_xlim()[0] - offset, ax0.get_xlim()[1] + offset
    ax0.set_xlim(x1, x2)

    offset = 0.01
    x1, x2 = ax1.get_xlim()[0] - offset, ax1.get_xlim()[1] + offset
    ax1.set_xlim(x1, x2)

    return fig, (ax0, ax1)

def plot_section(self, reverse=False, filled=False, **kw):
    lon, lat, data = map(np.asanyarray, (self.lon, self.lat, self.values))
    data = ma.masked_invalid(data)
    h = self.get_maxdepth()
    if reverse:
        lon = lon[::-1]
        lat = lat[::-1]
        data = data.T[::-1].T
        h = h[::-1]
    x = np.append(0, np.cumsum(gsw.distance(lon, lat)[0] / 1e3))
    z = self.index.values.astype(float)

    if filled:  # FIXME: Cause discontinuities.
        data = data.filled(fill_value=np.nan)
        data = extrap_sec(data, x, z, w1=0.97, w2=0.03)

    # Contour key words.
    fmt = kw.pop('fmt', '%1.0f')
    extend = kw.pop('extend', 'both')
    fontsize = kw.pop('fontsize', 12)
    labelsize = kw.pop('labelsize', 11)
    cmap = kw.pop('cmap', plt.cm.rainbow)
    levels = kw.pop('levels', np.arange(np.floor(data.min()),
                    np.ceil(data.max()) + 0.5, 0.5))

    # Colorbar key words.
    pad = kw.pop('pad', 0.04)
    aspect = kw.pop('aspect', 40)
    shrink = kw.pop('shrink', 0.9)
    fraction = kw.pop('fraction', 0.05)

    # Topography mask key words.
    dx = kw.pop('dx', 1.)
    kind = kw.pop('kind', 'linear')

    # Station symbols key words.
    color = kw.pop('color', 'k')
    offset = kw.pop('offset', -5)
    linewidth = kw.pop('linewidth', 1.5)

    # Figure.
    fig, ax = plt.subplots()
    xm, hm = gen_topomask(h, lon, lat, dx=dx, kind=kind)
    ax.plot(xm, hm, color='black', linewidth=linewidth, zorder=3)
    ax.fill_between(xm, hm, y2=hm.max(), color='0.9', zorder=3)

    ax.plot(x, [offset] * len(h), color=color, marker='v',
            alpha=0.5, zorder=5)
    ax.set_xlabel('Cross-shore distance [km]', fontsize=fontsize)
    ax.set_ylabel('Depth [m]', fontsize=fontsize)
    ax.set_ylim(offset, hm.max())
    ax.invert_yaxis()

    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')
    ax.yaxis.set_ticks_position('left')
    ax.yaxis.set_label_position('left')
    ax.xaxis.set_tick_params(tickdir='out', labelsize=labelsize, pad=1)
    ax.yaxis.set_tick_params(tickdir='out', labelsize=labelsize, pad=1)

    if False:  # TODO: +/- Black-and-White version.
        cs = ax.contour(x, z, data, colors='grey', levels=levels,
                        extend=extend, linewidths=1., alpha=1., zorder=2)
        ax.clabel(cs, fontsize=8, colors='grey', fmt=fmt, zorder=1)
        cb = None
    if True:  # Color version.
        cs = ax.contourf(x, z, data, cmap=cmap, levels=levels, alpha=1.,
                         extend=extend, zorder=2)  # manual=True
        # Colorbar.
        cb = fig.colorbar(mappable=cs, ax=ax, orientation='vertical',
                          aspect=aspect, shrink=shrink, fraction=fraction,
                          pad=pad)
    return fig, ax, cb

if __name__ == '__main__':
    import doctest
    doctest.testmod()
<commit_msg>Improved kw for `plot`. Marker is now optional for `plot_section`.<commit_after># -*- coding: utf-8 -*-
#
# plotting.py
#
# purpose:  Plotting methods for ctd DataFrame.
# author:   Filipe P. A. Fernandes
# e-mail:   ocefpaf@gmail
# web:      http://ocefpaf.tiddlyspot.com/
# created:  23-Jul-2013
# modified: Tue 23 Jul 2013 01:31:06 PM BRT
#
# obs:
#

# Scientific stack.
import gsw
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
import mpl_toolkits.axisartist as AA

from pandas import Series
from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1 import host_subplot

from utilities import extrap1d

__all__ = ['get_maxdepth',
           'extrap_sec',
           'gen_topomask',
           'plot',
           'plot_vars',
           'plot_section']

def get_maxdepth(self):
    valid_last_depth = self.apply(Series.notnull).values.T
    return np.float_(self.index.values * valid_last_depth).max(axis=1)

def extrap_sec(data, dist, depth, w1=1., w2=0):
    """
    Extrapolates `data` to zones where the shallow stations are shadowed by
    the deep stations.  The shadow region usually cannot be extrapolates via
    linear interpolation.

    The extrapolation is applied using the gradients of the `data` at a certain
    level.

    Parameters
    ----------
    data : array_like
          Data to be extrapolated
    dist : array_like
           Stations distance
    fd : float
         Decay factor [0-1]

    Returns
    -------
    Sec_extrap : array_like
                 Extrapolated variable

    Examples
    --------
    >>> import data, dist, z
    >>> Sec_extrap = extrap_sec(data, dist, z, fd=1.)
    """
    new_data1 = []
    for row in data:
        mask = ~np.isnan(row)
        if mask.any():
            y = row[mask]
            if y.size == 1:
                row = np.repeat(y, len(mask))
            else:
                x = dist[mask]
                f_i = interp1d(x, y)
                f_x = extrap1d(f_i)
                row = f_x(dist)
        new_data1.append(row)

    new_data2 = []
    for col in data.T:
        mask = ~np.isnan(col)
        if mask.any():
            y = col[mask]
            if y.size == 1:
                col = np.repeat(y, len(mask))
            else:
                z = depth[mask]
                f_i = interp1d(z, y)
                f_z = extrap1d(f_i)
                col = f_z(depth)
        new_data2.append(col)

    new_data = np.array(new_data1) * w1 + np.array(new_data2).T * w2
    return new_data

def gen_topomask(h, lon, lat, dx=1., kind='linear', plot=False):
    """
    Generates a topography mask from an oceanographic transect taking the
    deepest CTD scan as the depth of each station.

    Inputs
    ------
    h : array
        Pressure of the deepest CTD scan for each station [dbar].
    lons : array
           Longitude of each station [decimal degrees east].
    lat : Latitude of each station. [decimal degrees north].
    dx : float
         Horizontal resolution of the output arrays [km].
    kind : string, optional
           Type of the interpolation to be performed.
           See scipy.interpolate.interp1d documentation for details.
    plot : bool
           Whether to plot mask for visualization.

    Outputs
    -------
    xm : array
         Horizontal distances [km].
    hm : array
         Local depth [m].

    Examples
    --------
    >>> import gsw
    >>> import df  # FIXME: Add a dataset.
    >>> h = df.get_maxdepth()
    >>> # TODO: method to output distance.
    >>> x = np.append(0, np.cumsum(gsw.distance(df.lon, df.lat)[0] / 1e3))
    >>> xm, hm = gen_topomask(h, df.lon, df.lat, dx=1., kind='linear')
    >>> fig, ax = plt.subplots()
    >>> ax.plot(xm, hm, 'k', linewidth=1.5)
    >>> ax.plot(x, h, 'ro')
    >>> ax.set_xlabel('Distance [km]')
    >>> ax.set_ylabel('Depth [m]')
    >>> ax.grid(True)
    >>> plt.show()

    Author
    ------
    André Palóczy Filho (paloczy@gmail.com) --  October/2012
    """

    h, lon, lat = map(np.asanyarray, (h, lon, lat))
    # Distance in km.
    x = np.append(0, np.cumsum(gsw.distance(lon, lat)[0] / 1e3))
    h = -gsw.z_from_p(h, lat.mean())
    Ih = interp1d(x, h, kind=kind, bounds_error=False, fill_value=h[-1])
    xm = np.arange(0, x.max() + dx, dx)
    hm = Ih(xm)

    return xm, hm

def plot(self, **kw):
    """
    Plot a CTD variable against the index (pressure or depth).
    """
    figsize = kw.pop('figsize', (5.5, 6))
    fig, ax = plt.subplots(figsize=figsize)

    ax.plot(self.values, self.index, **kw)
    ax.set_ylabel(self.index.name)
    ax.set_xlabel(self.name)
    ax.invert_yaxis()
    offset = 0.01
    x1, x2 = ax.get_xlim()[0] - offset, ax.get_xlim()[1] + offset
    ax.set_xlim(x1, x2)
    return fig, ax

def plot_vars(self, variables=None, **kwds):
    """
    Plot CTD temperature and salinity.
    """
    fig = plt.figure(figsize=(8, 10))
    ax0 = host_subplot(111, axes_class=AA.Axes)
    ax1 = ax0.twiny()

    # Axis location.
    host_new_axis = ax0.get_grid_helper().new_fixed_axis
    ax0.axis["bottom"] = host_new_axis(loc="top", axes=ax0, offset=(0, 0))
    par_new_axis = ax1.get_grid_helper().new_fixed_axis
    ax1.axis["top"] = par_new_axis(loc="bottom", axes=ax1, offset=(0, 0))

    ax0.plot(self[variables[0]], self.index, 'r.', label='Temperature')
    ax1.plot(self[variables[1]], self.index, 'b.', label='Salinity')

    ax0.set_ylabel("Pressure [dbar]")
    ax0.set_xlabel(u"Temperature [\u00b0C]")
    ax1.set_xlabel("Salinity [kg g$^{-1}$]")
    ax1.invert_yaxis()

    try:  # FIXME with metadata.
        fig.suptitle(r"Station %s profile" % self.name)
    except AttributeError:
        pass

    ax0.legend(shadow=True, fancybox=True,
               numpoints=1, loc='lower right')

    offset = 0.01
    x1, x2 = ax0.get_xlim()[0] - offset, ax0.get_xlim()[1] + offset
    ax0.set_xlim(x1, x2)

    offset = 0.01
    x1, x2 = ax1.get_xlim()[0] - offset, ax1.get_xlim()[1] + offset
    ax1.set_xlim(x1, x2)

    return fig, (ax0, ax1)

def plot_section(self, reverse=False, filled=False, **kw):
    lon, lat, data = map(np.asanyarray, (self.lon, self.lat, self.values))
    data = ma.masked_invalid(data)
    h = self.get_maxdepth()
    if reverse:
        lon = lon[::-1]
        lat = lat[::-1]
        data = data.T[::-1].T
        h = h[::-1]
    x = np.append(0, np.cumsum(gsw.distance(lon, lat)[0] / 1e3))
    z = self.index.values.astype(float)

    if filled:  # FIXME: Cause discontinuities.
        data = data.filled(fill_value=np.nan)
        data = extrap_sec(data, x, z, w1=0.97, w2=0.03)

    # Contour key words.
    fmt = kw.pop('fmt', '%1.0f')
    extend = kw.pop('extend', 'both')
    fontsize = kw.pop('fontsize', 12)
    labelsize = kw.pop('labelsize', 11)
    cmap = kw.pop('cmap', plt.cm.rainbow)
    levels = kw.pop('levels', np.arange(np.floor(data.min()),
                    np.ceil(data.max()) + 0.5, 0.5))

    # Colorbar key words.
    pad = kw.pop('pad', 0.04)
    aspect = kw.pop('aspect', 40)
    shrink = kw.pop('shrink', 0.9)
    fraction = kw.pop('fraction', 0.05)

    # Topography mask key words.
    dx = kw.pop('dx', 1.)
    kind = kw.pop('kind', 'linear')
    linewidth = kw.pop('linewidth', 1.5)

    # Station symbols key words.
    color = kw.pop('color', 'k')
    offset = kw.pop('offset', -5)
    marker = kw.pop('marker', 'v')
    alpha = kw.pop('alpha', 0.5)

    # Figure.
    figsize = kw.pop('figsize', (12, 6))
    fig, ax = plt.subplots(figsize=figsize)
    xm, hm = gen_topomask(h, lon, lat, dx=dx, kind=kind)
    ax.plot(xm, hm, color='black', linewidth=linewidth, zorder=3)
    ax.fill_between(xm, hm, y2=hm.max(), color='0.9', zorder=3)

    if marker:
        ax.plot(x, [offset] * len(h), color=color, marker=marker, alpha=alpha,
                zorder=5)
    ax.set_xlabel('Cross-shore distance [km]', fontsize=fontsize)
    ax.set_ylabel('Depth [m]', fontsize=fontsize)
    ax.set_ylim(offset, hm.max())
    ax.invert_yaxis()

    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')
    ax.yaxis.set_ticks_position('left')
    ax.yaxis.set_label_position('left')
    ax.xaxis.set_tick_params(tickdir='out', labelsize=labelsize, pad=1)
    ax.yaxis.set_tick_params(tickdir='out', labelsize=labelsize, pad=1)

    if False:  # TODO: +/- Black-and-White version.
        cs = ax.contour(x, z, data, colors='grey', levels=levels,
                        extend=extend, linewidths=1., alpha=1., zorder=2)
        ax.clabel(cs, fontsize=8, colors='grey', fmt=fmt, zorder=1)
        cb = None
    if True:  # Color version.
        cs = ax.contourf(x, z, data, cmap=cmap, levels=levels, alpha=1.,
                         extend=extend, zorder=2)  # manual=True
        # Colorbar.
        cb = fig.colorbar(mappable=cs, ax=ax, orientation='vertical',
                          aspect=aspect, shrink=shrink, fraction=fraction,
                          pad=pad)
    return fig, ax, cb

if __name__ == '__main__':
    import doctest
    doctest.testmod()
<|endoftext|>


<commit_before>import socketio
import eventlet
import eventlet.wsgi
from flask import Flask, render_template

sio = socketio.Server()
app = Flask(__name__)

@app.route('/')
def index():
    """Serve the client-side application."""
    return render_template('index.html')

@sio.on('connect')
def connect(sid, environ):
    print("connected", sid)

@sio.on('swing')
def message(sid, data):
    print("got it", data)
    sio.emit('color', "blue")

@sio.on('disconnect', namespace='/chat')
def disconnect(sid):
    print('disconnect ', sid)

if __name__ == '__main__':
    # wrap Flask application with engineio's middleware
    app = socketio.Middleware(sio, app)

    # deploy as an eventlet WSGI server
    eventlet.wsgi.server(eventlet.listen(('localhost', 8000)), app)
    <commit_msg>Cleaned up whitespace<commit_after>import socketio
import eventlet
import eventlet.wsgi
from flask import Flask, render_template

sio = socketio.Server()
app = Flask(__name__)

@app.route('/')
def index():
    """Serve the client-side application."""
    return render_template('index.html')

@sio.on('connect')
def connect(sid, environ):
    print("connected", sid)

@sio.on('swing')
def message(sid, data):
    print("got it", data)
    sio.emit('color', "blue")

@sio.on('disconnect', namespace='/chat')
def disconnect(sid):
    print('disconnect ', sid)

if __name__ == '__main__':
    # wrap Flask application with engineio's middleware
    app = socketio.Middleware(sio, app)

    # deploy as an eventlet WSGI server
    eventlet.wsgi.server(eventlet.listen(('localhost', 8000)), app)
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
"""
.. module:: burpui.api
    :platform: Unix
    :synopsis: Burp-UI api module.

.. moduleauthor:: Ziirish <hi+burpui@ziirish.me>

"""
import os
import sys
import uuid
import hashlib
import logging

from flask import Blueprint, Response, request, current_app, session, abort
from flask_restplus import Api as ApiPlus
from flask_login import current_user
from importlib import import_module
from functools import wraps

from .custom.namespace import Namespace
from .._compat import to_bytes
from ..desc import __version__, __release__, __url__, __doc__
from ..server import BUIServer  # noqa
from ..exceptions import BUIserverException
from ..config import config

bui = current_app  # type: BUIServer
EXEMPT_METHODS = set(['OPTIONS'])

def cache_key():
    key = '{}-{}-{}-{}-{}'.format(
        session.get('login', uuid.uuid4()),
        request.path,
        request.values,
        request.headers.get('X-Session-Tag', ''),
        session.get('language', '')
    )
    key = hashlib.sha256(to_bytes(key)).hexdigest()
    return key

def api_login_required(func):
    """Custom login decorator that is able to parse Basic credentials as well as
    Cookies set with the traditional login.
    """
    @wraps(func)
    def decorated_view(*args, **kwargs):
        """decorator"""
        if request.method in EXEMPT_METHODS:  # pragma: no cover
            return func(*args, **kwargs)
        # 'func' is a Flask.view.MethodView so we have access to some special
        # params
        cls = func.view_class
        login_required = getattr(cls, 'login_required', True)
        if (bui.auth != 'none' and
                login_required and
                not bui.config.get('LOGIN_DISABLED', False)):
            if not current_user.is_authenticated:
                if request.headers.get('X-From-UI', False):
                    abort(403)
                return Response(
                    'Could not verify your access level for that URL.\n'
                    'You have to login with proper credentials', 401,
                    {'WWW-Authenticate': 'Basic realm="Login Required"'})
        return func(*args, **kwargs)
    return decorated_view

def check_acl(func):
    """Custom decorator to check if the ACL are in use or not"""
    @wraps(func)
    def decorated_view(*args, **kwargs):
        if request.method in EXEMPT_METHODS:  # pragma: no cover
            return func(*args, **kwargs)
        # 'func' is a Flask.view.MethodView so we have access to some special
        # params
        cls = func.view_class
        login_required = getattr(cls, 'login_required', True)
        if (bui.auth != 'none' and
                login_required and
                not bui.config.get('LOGIN_DISABLED', False)):
            if current_user.is_anonymous:
                abort(403)
        return func(*args, **kwargs)
    return decorated_view

class Api(ApiPlus):
    """Wrapper class around :class:`flask_restplus.Api`"""
    logger = logging.getLogger('burp-ui')
    # TODO: should use global object instead of reference
    loaded = False
    release = __release__
    __doc__ = __doc__
    __url__ = __url__
    CELERY_REQUIRED = ['async']

    def load_all(self):
        if config['WITH_LIMIT']:
            from ..ext.limit import limiter
            self.decorators.append(limiter.limit(config['BUI_RATIO']))
        """hack to automatically import api modules"""
        if not self.loaded:
            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            self.loaded = True
            for f in os.listdir(__path__[0]):
                name, ext = os.path.splitext(f)
                if (os.path.isfile(os.path.join(__path__[0], f)) and
                        ext == '.py' and
                        name not in ['__init__', '.', '..']):
                    mod = '.' + name
                    if name not in self.CELERY_REQUIRED or \
                            config['WITH_CELERY']:
                        self.logger.debug('Loading API module: {}'.format(mod))
                        try:
                            import_module(mod, __name__)
                        except:  # pragma: no cover
                            import traceback
                            self.logger.critical(
                                'Unable to load {}:\n{}'.format(
                                    mod,
                                    traceback.format_exc()
                                )
                            )
                    else:
                        self.logger.warning(
                            'Skipping API module: {}'.format(mod)
                        )

    def acl_admin_required(self, message='Access denied', code=403):
        def decorator(func):
            @wraps(func)
            def decorated(resource, *args, **kwargs):
                if not current_user.is_anonymous and \
                        not current_user.acl.is_admin():
                    resource.abort(code, message)
                return func(resource, *args, **kwargs)
            return decorated
        return decorator

    def acl_own_or_admin(self, key='name', message='Access denied', code=403):
        def decorator(func):
            @wraps(func)
            def decorated(resource, *args, **kwargs):
                if key not in kwargs:  # pragma: no cover
                    resource.abort(500, "key '{}' not found".format(key))
                if kwargs[key] != current_user.name and \
                        not current_user.is_anonymous and \
                        not current_user.acl.is_admin():
                    resource.abort(code, message)
                return func(resource, *args, **kwargs)
            return decorated
        return decorator

    def disabled_on_demo(self):
        def decorator(func):
            @wraps(func)
            def decorated(resource, *args, **kwargs):
                if config['BUI_DEMO']:
                    resource.abort(
                        405,
                        'Sorry, this feature is not available on the demo'
                    )
                return func(resource, *args, **kwargs)
            return decorated
        return decorator

    def namespace(self, *args, **kwargs):
        """A namespace factory

        :returns Namespace: a new namespace instance
        """
        ns = Namespace(*args, **kwargs)
        self.add_namespace(ns)
        return ns

apibp = Blueprint('api', __name__, url_prefix='/api')
api = Api(
    apibp,
    title='Burp-UI API',
    description='Burp-UI API to interact with burp',
    version=__version__,
    doc='/doc',
    decorators=[api_login_required, check_acl]
)

@api.errorhandler(BUIserverException)
def handle_bui_server_exception(error):
    """Forward a BUIserverException to the final user

    :param error: Custom exception
    :type error: :class:`burpui.exceptions.BUIserverException`
    """
    return {'message': error.description}, error.code
<commit_msg>revert force_acl check<commit_after># -*- coding: utf8 -*-
"""
.. module:: burpui.api
    :platform: Unix
    :synopsis: Burp-UI api module.

.. moduleauthor:: Ziirish <hi+burpui@ziirish.me>

"""
import os
import sys
import uuid
import hashlib
import logging

from flask import Blueprint, Response, request, current_app, session, abort
from flask_restplus import Api as ApiPlus
from flask_login import current_user
from importlib import import_module
from functools import wraps

from .custom.namespace import Namespace
from .._compat import to_bytes
from ..desc import __version__, __release__, __url__, __doc__
from ..server import BUIServer  # noqa
from ..exceptions import BUIserverException
from ..config import config

bui = current_app  # type: BUIServer
EXEMPT_METHODS = set(['OPTIONS'])

def cache_key():
    key = '{}-{}-{}-{}-{}'.format(
        session.get('login', uuid.uuid4()),
        request.path,
        request.values,
        request.headers.get('X-Session-Tag', ''),
        session.get('language', '')
    )
    key = hashlib.sha256(to_bytes(key)).hexdigest()
    return key

def api_login_required(func):
    """Custom login decorator that is able to parse Basic credentials as well as
    Cookies set with the traditional login.
    """
    @wraps(func)
    def decorated_view(*args, **kwargs):
        """decorator"""
        if request.method in EXEMPT_METHODS:  # pragma: no cover
            return func(*args, **kwargs)
        # 'func' is a Flask.view.MethodView so we have access to some special
        # params
        cls = func.view_class
        login_required = getattr(cls, 'login_required', True)
        if (bui.auth != 'none' and
                login_required and
                not bui.config.get('LOGIN_DISABLED', False)):
            if not current_user.is_authenticated:
                if request.headers.get('X-From-UI', False):
                    abort(403)
                return Response(
                    'Could not verify your access level for that URL.\n'
                    'You have to login with proper credentials', 401,
                    {'WWW-Authenticate': 'Basic realm="Login Required"'})
        return func(*args, **kwargs)
    return decorated_view

def check_acl(func):
    """Custom decorator to check if the ACL are in use or not"""
    @wraps(func)
    def decorated_view(*args, **kwargs):
        if request.method in EXEMPT_METHODS:  # pragma: no cover
            return func(*args, **kwargs)
        # 'func' is a Flask.view.MethodView so we have access to some special
        # params
        cls = func.view_class
        login_required = getattr(cls, 'login_required', True)
        if (bui.auth != 'none' and
                login_required and
                not bui.config.get('LOGIN_DISABLED', False)):
            if current_user.is_anonymous:
                abort(403)
        return func(*args, **kwargs)
    return decorated_view

class Api(ApiPlus):
    """Wrapper class around :class:`flask_restplus.Api`"""
    logger = logging.getLogger('burp-ui')
    # TODO: should use global object instead of reference
    loaded = False
    release = __release__
    __doc__ = __doc__
    __url__ = __url__
    CELERY_REQUIRED = ['async']

    def load_all(self):
        if config['WITH_LIMIT']:
            from ..ext.limit import limiter
            self.decorators.append(limiter.limit(config['BUI_RATIO']))
        """hack to automatically import api modules"""
        if not self.loaded:
            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            self.loaded = True
            for f in os.listdir(__path__[0]):
                name, ext = os.path.splitext(f)
                if (os.path.isfile(os.path.join(__path__[0], f)) and
                        ext == '.py' and
                        name not in ['__init__', '.', '..']):
                    mod = '.' + name
                    if name not in self.CELERY_REQUIRED or \
                            config['WITH_CELERY']:
                        self.logger.debug('Loading API module: {}'.format(mod))
                        try:
                            import_module(mod, __name__)
                        except:  # pragma: no cover
                            import traceback
                            self.logger.critical(
                                'Unable to load {}:\n{}'.format(
                                    mod,
                                    traceback.format_exc()
                                )
                            )
                    else:
                        self.logger.warning(
                            'Skipping API module: {}'.format(mod)
                        )

    def acl_admin_required(self, message='Access denied', code=403):
        def decorator(func):
            @wraps(func)
            def decorated(resource, *args, **kwargs):
                if not current_user.is_anonymous and \
                        not current_user.acl.is_admin():
                    resource.abort(code, message)
                return func(resource, *args, **kwargs)
            return decorated
        return decorator

    def acl_own_or_admin(self, key='name', message='Access denied', code=403):
        def decorator(func):
            @wraps(func)
            def decorated(resource, *args, **kwargs):
                if key not in kwargs:  # pragma: no cover
                    resource.abort(500, "key '{}' not found".format(key))
                if kwargs[key] != current_user.name and \
                        not current_user.is_anonymous and \
                        not current_user.acl.is_admin():
                    resource.abort(code, message)
                return func(resource, *args, **kwargs)
            return decorated
        return decorator

    def disabled_on_demo(self):
        def decorator(func):
            @wraps(func)
            def decorated(resource, *args, **kwargs):
                if config['BUI_DEMO']:
                    resource.abort(
                        405,
                        'Sorry, this feature is not available on the demo'
                    )
                return func(resource, *args, **kwargs)
            return decorated
        return decorator

    def namespace(self, *args, **kwargs):
        """A namespace factory

        :returns Namespace: a new namespace instance
        """
        ns = Namespace(*args, **kwargs)
        self.add_namespace(ns)
        return ns

apibp = Blueprint('api', __name__, url_prefix='/api')
api = Api(
    apibp,
    title='Burp-UI API',
    description='Burp-UI API to interact with burp',
    version=__version__,
    doc='/doc',
    decorators=[api_login_required]
)

@api.errorhandler(BUIserverException)
def handle_bui_server_exception(error):
    """Forward a BUIserverException to the final user

    :param error: Custom exception
    :type error: :class:`burpui.exceptions.BUIserverException`
    """
    return {'message': error.description}, error.code
<|endoftext|>


<commit_before>import os
import re

import pytest

from pip import __version__
from pip._internal.commands.show import search_packages_info

def test_basic_show(script):
    """
    Test end to end test for show command.
    """
    result = script.pip('show', 'pip')
    lines = result.stdout.splitlines()
    assert len(lines) == 10
    assert 'Name: pip' in lines
    assert 'Version: %s' % __version__ in lines
    assert any(line.startswith('Location: ') for line in lines)
    assert 'Requires: ' in lines

def test_show_with_files_not_found(script, data):
    """
    Test for show command with installed files listing enabled and
    installed-files.txt not found.
    """
    editable = data.packages.joinpath('SetupPyUTF8')
    script.pip('install', '-e', editable)
    result = script.pip('show', '-f', 'SetupPyUTF8')
    lines = result.stdout.splitlines()
    assert len(lines) == 12
    assert 'Name: SetupPyUTF8' in lines
    assert 'Version: 0.0.0' in lines
    assert any(line.startswith('Location: ') for line in lines)
    assert 'Requires: ' in lines
    assert 'Files:' in lines
    assert 'Cannot locate installed-files.txt' in lines

def test_show_with_files_from_wheel(script, data):
    """
    Test that a wheel's files can be listed
    """
    wheel_file = data.packages.joinpath('simple.dist-0.1-py2.py3-none-any.whl')
    script.pip('install', '--no-index', wheel_file)
    result = script.pip('show', '-f', 'simple.dist')
    lines = result.stdout.splitlines()
    assert 'Name: simple.dist' in lines
    assert 'Cannot locate installed-files.txt' not in lines[6], lines[6]
    assert re.search(r"Files:\n(  .+\n)+", result.stdout)

@pytest.mark.network
def test_show_with_all_files(script):
    """
    Test listing all files in the show command.
    """
    script.pip('install', 'initools==0.2')
    result = script.pip('show', '--files', 'initools')
    lines = result.stdout.splitlines()
    assert 'Cannot locate installed-files.txt' not in lines[6], lines[6]
    assert re.search(r"Files:\n(  .+\n)+", result.stdout)

def test_missing_argument(script):
    """
    Test show command with no arguments.
    """
    result = script.pip('show', expect_error=True)
    assert 'ERROR: Please provide a package name or names.' in result.stderr

def test_find_package_not_found():
    """
    Test trying to get info about a nonexistent package.
    """
    result = search_packages_info(['abcd3'])
    assert len(list(result)) == 0

def test_report_single_not_found(script):
    """
    Test passing one name and that isn't found.
    """
    # We choose a non-canonicalized name to test that the non-canonical
    # form is logged.
    # Also, the following should report an error as there are no results
    # to print. Consequently, there is no need to pass
    # allow_stderr_warning=True since this is implied by expect_error=True.
    result = script.pip('show', 'Abcd-3', expect_error=True)
    assert 'WARNING: Package(s) not found: Abcd-3' in result.stderr
    assert not result.stdout.splitlines()

def test_report_mixed_not_found(script):
    """
    Test passing a mixture of found and not-found names.
    """
    # We test passing non-canonicalized names.
    result = script.pip(
        'show', 'Abcd3', 'A-B-C', 'pip', allow_stderr_warning=True
    )
    assert 'WARNING: Package(s) not found: A-B-C, Abcd3' in result.stderr
    lines = result.stdout.splitlines()
    assert len(lines) == 10
    assert 'Name: pip' in lines

def test_search_any_case():
    """
    Search for a package in any case.

    """
    result = list(search_packages_info(['PIP']))
    assert len(result) == 1
    assert result[0]['name'] == 'pip'

def test_more_than_one_package():
    """
    Search for more than one package.

    """
    result = list(search_packages_info(['Pip', 'pytest', 'Virtualenv']))
    assert len(result) == 3

def test_show_verbose_with_classifiers(script):
    """
    Test that classifiers can be listed
    """
    result = script.pip('show', 'pip', '--verbose')
    lines = result.stdout.splitlines()
    assert 'Name: pip' in lines
    assert re.search(r"Classifiers:\n(  .+\n)+", result.stdout)
    assert "Intended Audience :: Developers" in result.stdout

def test_show_verbose_installer(script, data):
    """
    Test that the installer is shown (this currently needs a wheel install)
    """
    wheel_file = data.packages.joinpath('simple.dist-0.1-py2.py3-none-any.whl')
    script.pip('install', '--no-index', wheel_file)
    result = script.pip('show', '--verbose', 'simple.dist')
    lines = result.stdout.splitlines()
    assert 'Name: simple.dist' in lines
    assert 'Installer: pip' in lines

def test_show_verbose(script):
    """
    Test end to end test for verbose show command.
    """
    result = script.pip('show', '--verbose', 'pip')
    lines = result.stdout.splitlines()
    assert any(line.startswith('Metadata-Version: ') for line in lines)
    assert any(line.startswith('Installer: ') for line in lines)
    assert 'Entry-points:' in lines
    assert 'Classifiers:' in lines

def test_all_fields(script):
    """
    Test that all the fields are present
    """
    result = script.pip('show', 'pip')
    lines = result.stdout.splitlines()
    expected = {'Name', 'Version', 'Summary', 'Home-page', 'Author',
                'Author-email', 'License', 'Location', 'Requires',
                'Required-by'}
    actual = {re.sub(':.*$', '', line) for line in lines}
    assert actual == expected

def test_pip_show_is_short(script):
    """
    Test that pip show stays short
    """
    result = script.pip('show', 'pip')
    lines = result.stdout.splitlines()
    assert len(lines) <= 10

def test_pip_show_divider(script, data):
    """
    Expect a divider between packages
    """
    script.pip('install', 'pip-test-package', '--no-index',
               '-f', data.packages)
    result = script.pip('show', 'pip', 'pip-test-package')
    lines = result.stdout.splitlines()
    assert "---" in lines

def test_package_name_is_canonicalized(script, data):
    script.pip('install', 'pip-test-package', '--no-index', '-f',
               data.packages)

    dash_show_result = script.pip('show', 'pip-test-package')
    underscore_upper_show_result = script.pip('show', 'pip-test_Package')

    assert underscore_upper_show_result.returncode == 0
    assert underscore_upper_show_result.stdout == dash_show_result.stdout

def test_show_required_by_packages_simple(script, data):
    """
    Test that installed packages that depend on this package are shown
    """
    editable_path = os.path.join(data.src, 'requires_simple')
    script.pip(
        'install', '--no-index', '-f', data.find_links, editable_path
    )

    result = script.pip('show', 'simple')
    lines = result.stdout.splitlines()

    assert 'Name: simple' in lines
    assert 'Required-by: requires-simple' in lines

def test_show_required_by_packages_capitalized(script, data):
    """
    Test that installed packages that depend on this package are shown
    where the package has a capital letter
    """
    editable_path = os.path.join(data.src, 'requires_capitalized')
    script.pip(
        'install', '--no-index', '-f', data.find_links, editable_path
    )

    result = script.pip('show', 'simple')
    lines = result.stdout.splitlines()

    assert 'Name: simple' in lines
    assert 'Required-by: Requires-Capitalized' in lines

def test_show_required_by_with_mixed_capitalization(script, data):
    """
    Test that installed packages that depend on this package are shown
    where the package requires a name with a mix of
    lower and upper case letters
    """
    editable_path = os.path.join(data.src, 'required_by_mixed_capitalization')
    script.pip(
        'install', '--no-index', '-f', data.find_links, editable_path
    )

    result = script.pip('show', 'Requires_CapitalizeD')
    lines = result.stdout.splitlines()

    assert 'Name: Requires-Capitalized' in lines
    assert 'Required-by: simple' in lines
<commit_msg>Improve the documentation of the added tests<commit_after>import os
import re

import pytest

from pip import __version__
from pip._internal.commands.show import search_packages_info

def test_basic_show(script):
    """
    Test end to end test for show command.
    """
    result = script.pip('show', 'pip')
    lines = result.stdout.splitlines()
    assert len(lines) == 10
    assert 'Name: pip' in lines
    assert 'Version: %s' % __version__ in lines
    assert any(line.startswith('Location: ') for line in lines)
    assert 'Requires: ' in lines

def test_show_with_files_not_found(script, data):
    """
    Test for show command with installed files listing enabled and
    installed-files.txt not found.
    """
    editable = data.packages.joinpath('SetupPyUTF8')
    script.pip('install', '-e', editable)
    result = script.pip('show', '-f', 'SetupPyUTF8')
    lines = result.stdout.splitlines()
    assert len(lines) == 12
    assert 'Name: SetupPyUTF8' in lines
    assert 'Version: 0.0.0' in lines
    assert any(line.startswith('Location: ') for line in lines)
    assert 'Requires: ' in lines
    assert 'Files:' in lines
    assert 'Cannot locate installed-files.txt' in lines

def test_show_with_files_from_wheel(script, data):
    """
    Test that a wheel's files can be listed
    """
    wheel_file = data.packages.joinpath('simple.dist-0.1-py2.py3-none-any.whl')
    script.pip('install', '--no-index', wheel_file)
    result = script.pip('show', '-f', 'simple.dist')
    lines = result.stdout.splitlines()
    assert 'Name: simple.dist' in lines
    assert 'Cannot locate installed-files.txt' not in lines[6], lines[6]
    assert re.search(r"Files:\n(  .+\n)+", result.stdout)

@pytest.mark.network
def test_show_with_all_files(script):
    """
    Test listing all files in the show command.
    """
    script.pip('install', 'initools==0.2')
    result = script.pip('show', '--files', 'initools')
    lines = result.stdout.splitlines()
    assert 'Cannot locate installed-files.txt' not in lines[6], lines[6]
    assert re.search(r"Files:\n(  .+\n)+", result.stdout)

def test_missing_argument(script):
    """
    Test show command with no arguments.
    """
    result = script.pip('show', expect_error=True)
    assert 'ERROR: Please provide a package name or names.' in result.stderr

def test_find_package_not_found():
    """
    Test trying to get info about a nonexistent package.
    """
    result = search_packages_info(['abcd3'])
    assert len(list(result)) == 0

def test_report_single_not_found(script):
    """
    Test passing one name and that isn't found.
    """
    # We choose a non-canonicalized name to test that the non-canonical
    # form is logged.
    # Also, the following should report an error as there are no results
    # to print. Consequently, there is no need to pass
    # allow_stderr_warning=True since this is implied by expect_error=True.
    result = script.pip('show', 'Abcd-3', expect_error=True)
    assert 'WARNING: Package(s) not found: Abcd-3' in result.stderr
    assert not result.stdout.splitlines()

def test_report_mixed_not_found(script):
    """
    Test passing a mixture of found and not-found names.
    """
    # We test passing non-canonicalized names.
    result = script.pip(
        'show', 'Abcd3', 'A-B-C', 'pip', allow_stderr_warning=True
    )
    assert 'WARNING: Package(s) not found: A-B-C, Abcd3' in result.stderr
    lines = result.stdout.splitlines()
    assert len(lines) == 10
    assert 'Name: pip' in lines

def test_search_any_case():
    """
    Search for a package in any case.

    """
    result = list(search_packages_info(['PIP']))
    assert len(result) == 1
    assert result[0]['name'] == 'pip'

def test_more_than_one_package():
    """
    Search for more than one package.

    """
    result = list(search_packages_info(['Pip', 'pytest', 'Virtualenv']))
    assert len(result) == 3

def test_show_verbose_with_classifiers(script):
    """
    Test that classifiers can be listed
    """
    result = script.pip('show', 'pip', '--verbose')
    lines = result.stdout.splitlines()
    assert 'Name: pip' in lines
    assert re.search(r"Classifiers:\n(  .+\n)+", result.stdout)
    assert "Intended Audience :: Developers" in result.stdout

def test_show_verbose_installer(script, data):
    """
    Test that the installer is shown (this currently needs a wheel install)
    """
    wheel_file = data.packages.joinpath('simple.dist-0.1-py2.py3-none-any.whl')
    script.pip('install', '--no-index', wheel_file)
    result = script.pip('show', '--verbose', 'simple.dist')
    lines = result.stdout.splitlines()
    assert 'Name: simple.dist' in lines
    assert 'Installer: pip' in lines

def test_show_verbose(script):
    """
    Test end to end test for verbose show command.
    """
    result = script.pip('show', '--verbose', 'pip')
    lines = result.stdout.splitlines()
    assert any(line.startswith('Metadata-Version: ') for line in lines)
    assert any(line.startswith('Installer: ') for line in lines)
    assert 'Entry-points:' in lines
    assert 'Classifiers:' in lines

def test_all_fields(script):
    """
    Test that all the fields are present
    """
    result = script.pip('show', 'pip')
    lines = result.stdout.splitlines()
    expected = {'Name', 'Version', 'Summary', 'Home-page', 'Author',
                'Author-email', 'License', 'Location', 'Requires',
                'Required-by'}
    actual = {re.sub(':.*$', '', line) for line in lines}
    assert actual == expected

def test_pip_show_is_short(script):
    """
    Test that pip show stays short
    """
    result = script.pip('show', 'pip')
    lines = result.stdout.splitlines()
    assert len(lines) <= 10

def test_pip_show_divider(script, data):
    """
    Expect a divider between packages
    """
    script.pip('install', 'pip-test-package', '--no-index',
               '-f', data.packages)
    result = script.pip('show', 'pip', 'pip-test-package')
    lines = result.stdout.splitlines()
    assert "---" in lines

def test_package_name_is_canonicalized(script, data):
    script.pip('install', 'pip-test-package', '--no-index', '-f',
               data.packages)

    dash_show_result = script.pip('show', 'pip-test-package')
    underscore_upper_show_result = script.pip('show', 'pip-test_Package')

    assert underscore_upper_show_result.returncode == 0
    assert underscore_upper_show_result.stdout == dash_show_result.stdout

def test_show_required_by_packages_simple(script, data):
    """
    Test that installed packages that depend on this package are shown
    """
    editable_path = os.path.join(data.src, 'requires_simple')
    script.pip(
        'install', '--no-index', '-f', data.find_links, editable_path
    )

    result = script.pip('show', 'simple')
    lines = result.stdout.splitlines()

    assert 'Name: simple' in lines
    assert 'Required-by: requires-simple' in lines

def test_show_required_by_packages_capitalized(script, data):
    """
    Test that the installed packages which depend on a package are shown
    where the package has a capital letter
    """
    editable_path = os.path.join(data.src, 'requires_capitalized')
    script.pip(
        'install', '--no-index', '-f', data.find_links, editable_path
    )

    result = script.pip('show', 'simple')
    lines = result.stdout.splitlines()

    assert 'Name: simple' in lines
    assert 'Required-by: Requires-Capitalized' in lines

def test_show_required_by_with_mixed_capitalization(script, data):
    """
    Test that the installed packages which depend on a package are shown
    where the package has a name with a mix of
    lower and upper case letters
    """
    editable_path = os.path.join(data.src, 'required_by_mixed_capitalization')
    script.pip(
        'install', '--no-index', '-f', data.find_links, editable_path
    )

    result = script.pip('show', 'Requires_CapitalizeD')
    lines = result.stdout.splitlines()

    assert 'Name: Requires-Capitalized' in lines
    assert 'Required-by: simple' in lines
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function

import logging
import os

import pytest
import torch

import pyro
import pyro.distributions as dist
from pyro.infer import EmpiricalMarginal
from pyro.infer.mcmc.mcmc import MCMC
from pyro.infer.mcmc.nuts import NUTS
from tests.common import assert_equal

from .test_hmc import GaussianChain, T, rmse

logger = logging.getLogger(__name__)

TEST_CASES = [
    T(
        GaussianChain(dim=10, chain_len=3, num_obs=1),
        num_samples=800,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.25, 0.50, 0.75],
        expected_precs=[1.33, 1, 1.33],
        mean_tol=0.06,
        std_tol=0.06,
    ),
    T(
        GaussianChain(dim=10, chain_len=4, num_obs=1),
        num_samples=800,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.20, 0.40, 0.60, 0.80],
        expected_precs=[1.25, 0.83, 0.83, 1.25],
        mean_tol=0.06,
        std_tol=0.06,
    ),
    pytest.param(*T(
        GaussianChain(dim=5, chain_len=2, num_obs=10000),
        num_samples=800,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.5, 1.0],
        expected_precs=[2.0, 10000],
        mean_tol=0.04,
        std_tol=0.04,
    ), marks=[pytest.mark.skipif('CI' in os.environ or 'CUDA_TEST' in os.environ,
                                 reason='Slow test - skip on CI/CUDA')]),
    pytest.param(*T(
        GaussianChain(dim=5, chain_len=9, num_obs=1),
        num_samples=1200,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90],
        expected_precs=[1.11, 0.63, 0.48, 0.42, 0.4, 0.42, 0.48, 0.63, 1.11],
        mean_tol=0.07,
        std_tol=0.07,
    ), marks=[pytest.mark.skipif('CI' in os.environ or 'CUDA_TEST' in os.environ,
                                 reason='Slow test - skip on CI/CUDA')])
]

TEST_IDS = [t[0].id_fn() if type(t).__name__ == 'TestExample'
            else t[0][0].id_fn() for t in TEST_CASES]

@pytest.mark.parametrize(
    'fixture, num_samples, warmup_steps, hmc_params, expected_means, expected_precs, mean_tol, std_tol',
    TEST_CASES,
    ids=TEST_IDS)
@pytest.mark.init(rng_seed=34)
@pytest.mark.disable_validation()
def test_nuts_conjugate_gaussian(fixture,
                                 num_samples,
                                 warmup_steps,
                                 hmc_params,
                                 expected_means,
                                 expected_precs,
                                 mean_tol,
                                 std_tol):
    pyro.get_param_store().clear()
    nuts_kernel = NUTS(fixture.model)
    mcmc_run = MCMC(nuts_kernel, num_samples, warmup_steps).run(fixture.data)
    for i in range(1, fixture.chain_len + 1):
        param_name = 'loc_' + str(i)
        marginal = EmpiricalMarginal(mcmc_run, sites=param_name)
        latent_loc = marginal.mean
        latent_std = marginal.variance.sqrt()
        expected_mean = torch.ones(fixture.dim) * expected_means[i - 1]
        expected_std = 1 / torch.sqrt(torch.ones(fixture.dim) * expected_precs[i - 1])

        # Actual vs expected posterior means for the latents
        logger.info('Posterior mean (actual) - {}'.format(param_name))
        logger.info(latent_loc)
        logger.info('Posterior mean (expected) - {}'.format(param_name))
        logger.info(expected_mean)
        assert_equal(rmse(latent_loc, expected_mean).item(), 0.0, prec=mean_tol)

        # Actual vs expected posterior precisions for the latents
        logger.info('Posterior std (actual) - {}'.format(param_name))
        logger.info(latent_std)
        logger.info('Posterior std (expected) - {}'.format(param_name))
        logger.info(expected_std)
        assert_equal(rmse(latent_std, expected_std).item(), 0.0, prec=std_tol)

def test_logistic_regression():
    dim = 3
    data = torch.randn(2000, dim)
    true_coefs = torch.arange(1., dim + 1.)
    labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()

    def model(data):
        coefs_mean = torch.zeros(dim)
        coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(dim)))
        y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)
        return y

    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='beta')
    assert_equal(rmse(true_coefs, posterior.mean).item(), 0.0, prec=0.1)

@pytest.mark.parametrize(
    "step_size, adapt_step_size, adapt_mass_matrix, full_mass",
    [
        (0.1, False, False, False),
        (0.1, False, True, False),
        (None, True, False, False),
        (None, True, True, False),
        (None, True, True, True),
    ]
)
def test_beta_bernoulli(step_size, adapt_step_size, adapt_mass_matrix, full_mass):
    def model(data):
        alpha = torch.tensor([1.1, 1.1])
        beta = torch.tensor([1.1, 1.1])
        p_latent = pyro.sample("p_latent", dist.Beta(alpha, beta))
        pyro.sample("obs", dist.Bernoulli(p_latent), obs=data)
        return p_latent

    true_probs = torch.tensor([0.9, 0.1])
    data = dist.Bernoulli(true_probs).sample(sample_shape=(torch.Size((1000,))))
    nuts_kernel = NUTS(model, step_size, adapt_step_size, adapt_mass_matrix, full_mass)
    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
    assert_equal(posterior.mean, true_probs, prec=0.02)

def test_gamma_normal():
    def model(data):
        rate = torch.tensor([1.0, 1.0])
        concentration = torch.tensor([1.0, 1.0])
        p_latent = pyro.sample('p_latent', dist.Gamma(rate, concentration))
        pyro.sample("obs", dist.Normal(3, p_latent), obs=data)
        return p_latent

    true_std = torch.tensor([0.5, 2])
    data = dist.Normal(3, true_std).sample(sample_shape=(torch.Size((2000,))))
    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
    assert_equal(posterior.mean, true_std, prec=0.05)

def test_dirichlet_categorical():
    def model(data):
        concentration = torch.tensor([1.0, 1.0, 1.0])
        p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
        pyro.sample("obs", dist.Categorical(p_latent), obs=data)
        return p_latent

    true_probs = torch.tensor([0.1, 0.6, 0.3])
    data = dist.Categorical(true_probs).sample(sample_shape=(torch.Size((2000,))))
    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
    assert_equal(posterior.mean, true_probs, prec=0.02)

def test_gamma_beta():
    def model(data):
        alpha_prior = pyro.sample('alpha', dist.Gamma(concentration=1., rate=1.))
        beta_prior = pyro.sample('beta', dist.Gamma(concentration=1., rate=1.))
        pyro.sample('x', dist.Beta(concentration1=alpha_prior, concentration0=beta_prior), obs=data)

    true_alpha = torch.tensor(5.)
    true_beta = torch.tensor(1.)
    data = dist.Beta(concentration1=true_alpha, concentration0=true_beta).sample(torch.Size((5000,)))
    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=200).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites=['alpha', 'beta'])
    assert_equal(posterior.mean, torch.stack([true_alpha, true_beta]), prec=0.05)

def test_gaussian_mixture_model():
    K, N = 3, 1000

    def gmm(data):
        mix_proportions = pyro.sample("phi", dist.Dirichlet(torch.ones(K)))
        with pyro.plate("num_clusters", K):
            cluster_means = pyro.sample("cluster_means", dist.Normal(torch.arange(float(K)), 1.))
        with pyro.plate("data", data.shape[0]):
            assignments = pyro.sample("assignments", dist.Categorical(mix_proportions))
            pyro.sample("obs", dist.Normal(cluster_means[assignments], 1.), obs=data)
        return cluster_means

    true_cluster_means = torch.tensor([1., 5., 10.])
    true_mix_proportions = torch.tensor([0.1, 0.3, 0.6])
    cluster_assignments = dist.Categorical(true_mix_proportions).sample(torch.Size((N,)))
    data = dist.Normal(true_cluster_means[cluster_assignments], 1.0).sample()
    nuts_kernel = NUTS(gmm, max_plate_nesting=1)
    mcmc_run = MCMC(nuts_kernel, num_samples=300, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites=["phi", "cluster_means"]).mean.sort()[0]
    assert_equal(posterior[0], true_mix_proportions, prec=0.05)
    assert_equal(posterior[1], true_cluster_means, prec=0.2)

def test_bernoulli_latent_model():
    def model(data):
        y_prob = pyro.sample("y_prob", dist.Beta(1., 1.))
        with pyro.plate("data", data.shape[0]):
            y = pyro.sample("y", dist.Bernoulli(y_prob))
            z = pyro.sample("z", dist.Bernoulli(0.65 * y + 0.1))
            pyro.sample("obs", dist.Normal(2. * z, 1.), obs=data)

    N = 2000
    y_prob = torch.tensor(0.3)
    y = dist.Bernoulli(y_prob).sample(torch.Size((N,)))
    z = dist.Bernoulli(0.65 * y + 0.1).sample()
    data = dist.Normal(2. * z, 1.0).sample()
    nuts_kernel = NUTS(model, max_plate_nesting=1)
    mcmc_run = MCMC(nuts_kernel, num_samples=600, warmup_steps=200).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites="y_prob").mean
    assert_equal(posterior, y_prob, prec=0.05)

@pytest.mark.parametrize("num_steps,use_einsum", [
    (2, False),
    (3, False),
    (3, True),
    # This will crash without the einsum backend
    pytest.param(30, True,
                 marks=pytest.mark.skip(reason="https://github.com/pytorch/pytorch/issues/10661")),
])
def test_gaussian_hmm_enum_shape(num_steps, use_einsum):
    dim = 4

    def model(data):
        initialize = pyro.sample("initialize", dist.Dirichlet(torch.ones(dim)))
        transition = pyro.sample("transition", dist.Dirichlet(torch.ones(dim, dim)))
        emission_loc = pyro.sample("emission_loc", dist.Normal(torch.zeros(dim), torch.ones(dim)))
        emission_scale = pyro.sample("emission_scale", dist.LogNormal(torch.zeros(dim), torch.ones(dim)))
        x = None
        for t, y in enumerate(data):
            x = pyro.sample("x_{}".format(t), dist.Categorical(initialize if x is None else transition[x]))
            pyro.sample("y_{}".format(t), dist.Normal(emission_loc[x], emission_scale[x]), obs=y)
            # check shape
            effective_dim = sum(1 for size in x.shape if size > 1)
            assert effective_dim == 1

    data = torch.ones(num_steps)
    nuts_kernel = NUTS(model, max_plate_nesting=0, experimental_use_einsum=use_einsum)
    MCMC(nuts_kernel, num_samples=5, warmup_steps=5).run(data)
<commit_msg>Fix failing CUDA test in test_nuts (#1477)<commit_after>from __future__ import absolute_import, division, print_function

import logging
import os

import pytest
import torch

import pyro
import pyro.distributions as dist
from pyro.infer import EmpiricalMarginal
from pyro.infer.mcmc.mcmc import MCMC
from pyro.infer.mcmc.nuts import NUTS
from tests.common import assert_equal

from .test_hmc import GaussianChain, T, rmse

logger = logging.getLogger(__name__)

TEST_CASES = [
    T(
        GaussianChain(dim=10, chain_len=3, num_obs=1),
        num_samples=800,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.25, 0.50, 0.75],
        expected_precs=[1.33, 1, 1.33],
        mean_tol=0.06,
        std_tol=0.06,
    ),
    T(
        GaussianChain(dim=10, chain_len=4, num_obs=1),
        num_samples=800,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.20, 0.40, 0.60, 0.80],
        expected_precs=[1.25, 0.83, 0.83, 1.25],
        mean_tol=0.07,
        std_tol=0.06,
    ),
    pytest.param(*T(
        GaussianChain(dim=5, chain_len=2, num_obs=10000),
        num_samples=800,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.5, 1.0],
        expected_precs=[2.0, 10000],
        mean_tol=0.04,
        std_tol=0.04,
    ), marks=[pytest.mark.skipif('CI' in os.environ or 'CUDA_TEST' in os.environ,
                                 reason='Slow test - skip on CI/CUDA')]),
    pytest.param(*T(
        GaussianChain(dim=5, chain_len=9, num_obs=1),
        num_samples=1200,
        warmup_steps=200,
        hmc_params=None,
        expected_means=[0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90],
        expected_precs=[1.11, 0.63, 0.48, 0.42, 0.4, 0.42, 0.48, 0.63, 1.11],
        mean_tol=0.07,
        std_tol=0.07,
    ), marks=[pytest.mark.skipif('CI' in os.environ or 'CUDA_TEST' in os.environ,
                                 reason='Slow test - skip on CI/CUDA')])
]

TEST_IDS = [t[0].id_fn() if type(t).__name__ == 'TestExample'
            else t[0][0].id_fn() for t in TEST_CASES]

@pytest.mark.parametrize(
    'fixture, num_samples, warmup_steps, hmc_params, expected_means, expected_precs, mean_tol, std_tol',
    TEST_CASES,
    ids=TEST_IDS)
@pytest.mark.init(rng_seed=34)
@pytest.mark.disable_validation()
def test_nuts_conjugate_gaussian(fixture,
                                 num_samples,
                                 warmup_steps,
                                 hmc_params,
                                 expected_means,
                                 expected_precs,
                                 mean_tol,
                                 std_tol):
    pyro.get_param_store().clear()
    nuts_kernel = NUTS(fixture.model)
    mcmc_run = MCMC(nuts_kernel, num_samples, warmup_steps).run(fixture.data)
    for i in range(1, fixture.chain_len + 1):
        param_name = 'loc_' + str(i)
        marginal = EmpiricalMarginal(mcmc_run, sites=param_name)
        latent_loc = marginal.mean
        latent_std = marginal.variance.sqrt()
        expected_mean = torch.ones(fixture.dim) * expected_means[i - 1]
        expected_std = 1 / torch.sqrt(torch.ones(fixture.dim) * expected_precs[i - 1])

        # Actual vs expected posterior means for the latents
        logger.info('Posterior mean (actual) - {}'.format(param_name))
        logger.info(latent_loc)
        logger.info('Posterior mean (expected) - {}'.format(param_name))
        logger.info(expected_mean)
        assert_equal(rmse(latent_loc, expected_mean).item(), 0.0, prec=mean_tol)

        # Actual vs expected posterior precisions for the latents
        logger.info('Posterior std (actual) - {}'.format(param_name))
        logger.info(latent_std)
        logger.info('Posterior std (expected) - {}'.format(param_name))
        logger.info(expected_std)
        assert_equal(rmse(latent_std, expected_std).item(), 0.0, prec=std_tol)

def test_logistic_regression():
    dim = 3
    data = torch.randn(2000, dim)
    true_coefs = torch.arange(1., dim + 1.)
    labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()

    def model(data):
        coefs_mean = torch.zeros(dim)
        coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(dim)))
        y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)
        return y

    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='beta')
    assert_equal(rmse(true_coefs, posterior.mean).item(), 0.0, prec=0.1)

@pytest.mark.parametrize(
    "step_size, adapt_step_size, adapt_mass_matrix, full_mass",
    [
        (0.1, False, False, False),
        (0.1, False, True, False),
        (None, True, False, False),
        (None, True, True, False),
        (None, True, True, True),
    ]
)
def test_beta_bernoulli(step_size, adapt_step_size, adapt_mass_matrix, full_mass):
    def model(data):
        alpha = torch.tensor([1.1, 1.1])
        beta = torch.tensor([1.1, 1.1])
        p_latent = pyro.sample("p_latent", dist.Beta(alpha, beta))
        pyro.sample("obs", dist.Bernoulli(p_latent), obs=data)
        return p_latent

    true_probs = torch.tensor([0.9, 0.1])
    data = dist.Bernoulli(true_probs).sample(sample_shape=(torch.Size((1000,))))
    nuts_kernel = NUTS(model, step_size, adapt_step_size, adapt_mass_matrix, full_mass)
    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
    assert_equal(posterior.mean, true_probs, prec=0.02)

def test_gamma_normal():
    def model(data):
        rate = torch.tensor([1.0, 1.0])
        concentration = torch.tensor([1.0, 1.0])
        p_latent = pyro.sample('p_latent', dist.Gamma(rate, concentration))
        pyro.sample("obs", dist.Normal(3, p_latent), obs=data)
        return p_latent

    true_std = torch.tensor([0.5, 2])
    data = dist.Normal(3, true_std).sample(sample_shape=(torch.Size((2000,))))
    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
    assert_equal(posterior.mean, true_std, prec=0.05)

def test_dirichlet_categorical():
    def model(data):
        concentration = torch.tensor([1.0, 1.0, 1.0])
        p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
        pyro.sample("obs", dist.Categorical(p_latent), obs=data)
        return p_latent

    true_probs = torch.tensor([0.1, 0.6, 0.3])
    data = dist.Categorical(true_probs).sample(sample_shape=(torch.Size((2000,))))
    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
    assert_equal(posterior.mean, true_probs, prec=0.02)

def test_gamma_beta():
    def model(data):
        alpha_prior = pyro.sample('alpha', dist.Gamma(concentration=1., rate=1.))
        beta_prior = pyro.sample('beta', dist.Gamma(concentration=1., rate=1.))
        pyro.sample('x', dist.Beta(concentration1=alpha_prior, concentration0=beta_prior), obs=data)

    true_alpha = torch.tensor(5.)
    true_beta = torch.tensor(1.)
    data = dist.Beta(concentration1=true_alpha, concentration0=true_beta).sample(torch.Size((5000,)))
    nuts_kernel = NUTS(model)
    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=200).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites=['alpha', 'beta'])
    assert_equal(posterior.mean, torch.stack([true_alpha, true_beta]), prec=0.05)

def test_gaussian_mixture_model():
    K, N = 3, 1000

    def gmm(data):
        mix_proportions = pyro.sample("phi", dist.Dirichlet(torch.ones(K)))
        with pyro.plate("num_clusters", K):
            cluster_means = pyro.sample("cluster_means", dist.Normal(torch.arange(float(K)), 1.))
        with pyro.plate("data", data.shape[0]):
            assignments = pyro.sample("assignments", dist.Categorical(mix_proportions))
            pyro.sample("obs", dist.Normal(cluster_means[assignments], 1.), obs=data)
        return cluster_means

    true_cluster_means = torch.tensor([1., 5., 10.])
    true_mix_proportions = torch.tensor([0.1, 0.3, 0.6])
    cluster_assignments = dist.Categorical(true_mix_proportions).sample(torch.Size((N,)))
    data = dist.Normal(true_cluster_means[cluster_assignments], 1.0).sample()
    nuts_kernel = NUTS(gmm, max_plate_nesting=1)
    mcmc_run = MCMC(nuts_kernel, num_samples=300, warmup_steps=100).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites=["phi", "cluster_means"]).mean.sort()[0]
    assert_equal(posterior[0], true_mix_proportions, prec=0.05)
    assert_equal(posterior[1], true_cluster_means, prec=0.2)

def test_bernoulli_latent_model():
    def model(data):
        y_prob = pyro.sample("y_prob", dist.Beta(1., 1.))
        with pyro.plate("data", data.shape[0]):
            y = pyro.sample("y", dist.Bernoulli(y_prob))
            z = pyro.sample("z", dist.Bernoulli(0.65 * y + 0.1))
            pyro.sample("obs", dist.Normal(2. * z, 1.), obs=data)

    N = 2000
    y_prob = torch.tensor(0.3)
    y = dist.Bernoulli(y_prob).sample(torch.Size((N,)))
    z = dist.Bernoulli(0.65 * y + 0.1).sample()
    data = dist.Normal(2. * z, 1.0).sample()
    nuts_kernel = NUTS(model, max_plate_nesting=1)
    mcmc_run = MCMC(nuts_kernel, num_samples=600, warmup_steps=200).run(data)
    posterior = EmpiricalMarginal(mcmc_run, sites="y_prob").mean
    assert_equal(posterior, y_prob, prec=0.05)

@pytest.mark.parametrize("num_steps,use_einsum", [
    (2, False),
    (3, False),
    (3, True),
    # This will crash without the einsum backend
    pytest.param(30, True,
                 marks=pytest.mark.skip(reason="https://github.com/pytorch/pytorch/issues/10661")),
])
def test_gaussian_hmm_enum_shape(num_steps, use_einsum):
    dim = 4

    def model(data):
        initialize = pyro.sample("initialize", dist.Dirichlet(torch.ones(dim)))
        transition = pyro.sample("transition", dist.Dirichlet(torch.ones(dim, dim)))
        emission_loc = pyro.sample("emission_loc", dist.Normal(torch.zeros(dim), torch.ones(dim)))
        emission_scale = pyro.sample("emission_scale", dist.LogNormal(torch.zeros(dim), torch.ones(dim)))
        x = None
        for t, y in enumerate(data):
            x = pyro.sample("x_{}".format(t), dist.Categorical(initialize if x is None else transition[x]))
            pyro.sample("y_{}".format(t), dist.Normal(emission_loc[x], emission_scale[x]), obs=y)
            # check shape
            effective_dim = sum(1 for size in x.shape if size > 1)
            assert effective_dim == 1

    data = torch.ones(num_steps)
    nuts_kernel = NUTS(model, max_plate_nesting=0, experimental_use_einsum=use_einsum)
    MCMC(nuts_kernel, num_samples=5, warmup_steps=5).run(data)
<|endoftext|>