<commit_before>
import scipy.special
from numpy import logical_and, asarray, pi, zeros_like, \
     piecewise, array, arctan2, tan, zeros, arange, floor
from numpy.core.umath import sqrt, exp, greater, less, cos, add, sin, \
     less_equal, greater_equal
from spline import *      # C-modules
from scipy.misc import comb

gamma = scipy.special.gamma
def factorial(n):
    return gamma(n+1)

def spline_filter(Iin, lmbda=5.0):
    """Smoothing spline (cubic) filtering of a rank-2 array.

    Filter an input data set, `Iin`, using a (cubic) smoothing spline of
    fall-off `lmbda`.
    """
    intype = Iin.dtype.char
    hcol = array([1.0,4.0,1.0],'f')/6.0
    if intype in ['F','D']:
        Iin = Iin.astype('F')
        ckr = cspline2d(Iin.real,lmbda)
        cki = cspline2d(Iin.imag,lmbda)
        outr = sepfir2d(ckr,hcol,hcol)
        outi = sepfir2d(cki,hcol,hcol)
        out = (outr + 1j*outi).astype(intype)
    elif intype in ['f','d']:
        ckr = cspline2d(Iin,lmbda)
        out = sepfir2d(ckr, hcol, hcol)
        out = out.astype(intype)
    else:
        raise TypeError;
    return out

_splinefunc_cache = {}

def _bspline_piecefunctions(order):
    """Returns the function defined over the left-side pieces for a bspline of
    a given order.

    The 0th piece is the first one less than 0.  The last piece is a function
    identical to 0 (returned as the constant 0).  (There are order//2 + 2 total
    pieces).

    Also returns the condition functions that when evaluated return boolean
    arrays for use with `numpy.piecewise`.
    """
    try:
        return _splinefunc_cache[order]
    except KeyError:
        pass

    def condfuncgen(num, val1, val2):
        if num == 0:
            return lambda x: logical_and(less_equal(x, val1),
                                         greater_equal(x, val2))
        elif num == 2:
            return lambda x: less_equal(x, val2)
        else:
            return lambda x: logical_and(less(x, val1),
                                         greater_equal(x, val2))

    last = order // 2 + 2
    if order % 2:
        startbound = -1.0
    else:
        startbound = -0.5
    condfuncs = [condfuncgen(0, 0, startbound)]
    bound = startbound
    for num in xrange(1,last-1):
        condfuncs.append(condfuncgen(1, bound, bound-1))
        bound = bound-1
    condfuncs.append(condfuncgen(2, 0, -(order+1)/2.0))

    # final value of bound is used in piecefuncgen below

    # the functions to evaluate are taken from the left-hand-side
    #  in the general expression derived from the central difference
    #  operator (because they involve fewer terms).

    fval = factorial(order)
    def piecefuncgen(num):
        Mk = order // 2 - num
        if (Mk < 0): return 0  # final function is 0
        coeffs = [(1-2*(k%2))*float(comb(order+1, k, exact=1))/fval for k in xrange(Mk+1)]
        shifts = [-bound - k for k in xrange(Mk+1)]
        #print "Adding piece number %d with coeffs %s and shifts %s" % (num, str(coeffs), str(shifts))
        def thefunc(x):
            res = 0.0
            for k in range(Mk+1):
                res += coeffs[k]*(x+shifts[k])**order
            return res
        return thefunc

    funclist = [piecefuncgen(k) for k in xrange(last)]

    _splinefunc_cache[order] = (funclist, condfuncs)

    return funclist, condfuncs

def bspline(x, n):
    """B-spline basis function of order n.

    Notes
    -----
    Uses numpy.piecewise and automatic function-generator.

    """
    ax = -abs(asarray(x))
    # number of pieces on the left-side is (n+1)/2
    funclist, condfuncs = _bspline_piecefunctions(n)
    condlist = [func(ax) for func in condfuncs]
    return piecewise(ax, condlist, funclist)

def gauss_spline(x, n):
    """Gaussian approximation to B-spline basis function of order n.
    """
    signsq = (n+1) / 12.0
    return 1/sqrt(2*pi*signsq) * exp(-x**2 / 2 / signsq)

def cubic(x):
    """A cubic B-spline.

    This is a special case of `bspline`, and equivalent to ``bspline(x, 3)``.
    """
    ax = abs(asarray(x))
    res = zeros_like(ax)
    cond1 = less(ax, 1)
    if cond1.any():
        ax1 = ax[cond1]
        res[cond1] = 2.0/3 - 1.0/2*ax1**2 * (2-ax1)
    cond2 = ~cond1 & less(ax, 2)
    if cond2.any():
        ax2 = ax[cond2]
        res[cond2] = 1.0/6*(2-ax2)**3
    return res

def quadratic(x):
    """A quadratic B-spline.

    This is a special case of `bspline`, and equivalent to ``bspline(x, 2)``.
    """
    ax = abs(asarray(x))
    res = zeros_like(ax)
    cond1 = less(ax, 0.5)
    if cond1.any():
        ax1 = ax[cond1]
        res[cond1] = 0.75-ax1**2
    cond2 = ~cond1 & less(ax, 1.5)
    if cond2.any():
        ax2 = ax[cond2]
        res[cond2] = (ax2-1.5)**2 / 2.0
    return res

def c0_P(order):
    # values taken from Unser, et.al. 1993 IEEE
    if order == 0:
        c0 = 1
        P = array([1])
    elif order == 1:
        c0 = 1
        P = array([0,1])
    elif order == 2:
        c0 = 8
        P = array([1,6,1])
    elif order == 3:
        c0 = 6
        P = array([1,4,1])
    elif order == 4:
        c0 = 384
        P = array([1,76,230,76,1])
    elif order == 5:
        c0 = 120
        P = array([1,26,66,26,1])
    elif order == 6:
        c0 = 46080
        P = array([1,722,10543,23548, 10543, 722, 1])
    elif order == 7:
        c0 = 5040
        P = array([1,120,1191,2416,1191, 120, 1])
    else:
        raise ValueError, "Unknown order."

def _coeff_smooth(lam):
    xi = 1 - 96*lam + 24*lam * sqrt(3 + 144*lam)
    omeg = arctan2(sqrt(144*lam-1),sqrt(xi))
    rho = (24*lam - 1 - sqrt(xi)) / (24*lam)
    rho = rho * sqrt((48*lam + 24*lam * sqrt(3+144*lam))/xi)
    return rho,omeg

def _hc(k, cs, rho, omega):
    return cs / sin(omega) * (rho**k)*sin(omega*(k+1))*(greater(k,-1))

def _hs(k, cs, rho, omega):
    c0 = cs*cs * (1 + rho*rho) / (1 - rho*rho) / (1-2*rho*rho*cos(2*omega) + rho**4)
    gamma = (1-rho*rho) / (1+rho*rho) / tan(omega)
    ak = abs(k)
    return c0 * rho**ak * (cos(omega*ak) + gamma*sin(omega*ak))

def _cubic_smooth_coeff(signal, lamb):
    rho, omega = _coeff_smooth(lamb)
    cs = 1-2*rho*cos(omega) + rho*rho
    K = len(signal)
    yp = zeros((K,),signal.dtype.char)
    k = arange(K)
    yp[0] = _hc(0,cs,rho,omega)*signal[0] + \
            add.reduce(_hc(k+1,cs,rho,omega)*signal)

    yp[1] = _hc(0,cs,rho,omega)*signal[0] + \
            _hc(1,cs,rho,omega)*signal[1] + \
            add.reduce(_hc(k+2,cs,rho,omega)*signal)

    for n in range(2,K):
        yp[n] = cs * signal[n] + 2*rho*cos(omega)*yp[n-1] - rho*rho*yp[n-2]

    y = zeros((K,),signal.dtype.char)

    y[K-1] = add.reduce((_hs(k,cs,rho,omega) + _hs(k+1,cs,rho,omega))*signal[::-1])
    y[K-2] = add.reduce((_hs(k-1,cs,rho,omega) + _hs(k+2,cs,rho,omega))*signal[::-1])

    for n in range(K-3,-1,-1):
        y[n] = cs*yp[n] + 2*rho*cos(omega)*y[n+1] - rho*rho*y[n+2]

    return y

def _cubic_coeff(signal):
    zi = -2 + sqrt(3)
    K = len(signal)
    yplus = zeros((K,),signal.dtype.char)
    powers = zi**arange(K)
    yplus[0] = signal[0] + zi*add.reduce(powers*signal)
    for k in range(1,K):
        yplus[k] = signal[k] + zi*yplus[k-1]
    output = zeros((K,),signal.dtype)
    output[K-1] = zi / (zi-1)*yplus[K-1]
    for k in range(K-2,-1,-1):
        output[k] = zi*(output[k+1]-yplus[k])
    return output*6.0

def _quadratic_coeff(signal):
    zi = -3 + 2*sqrt(2.0)
    K = len(signal)
    yplus = zeros((K,),signal.dtype.char)
    powers = zi**arange(K)
    yplus[0] = signal[0] + zi*add.reduce(powers*signal)
    for k in range(1,K):
        yplus[k] = signal[k] + zi*yplus[k-1]
    output = zeros((K,),signal.dtype.char)
    output[K-1] = zi / (zi-1)*yplus[K-1]
    for k in range(K-2,-1,-1):
        output[k] = zi*(output[k+1]-yplus[k])
    return output*8.0

def cspline1d(signal, lamb=0.0):
    """
    Compute cubic spline coefficients for rank-1 array.

    Find the cubic spline coefficients for a 1-D signal assuming
    mirror-symmetric boundary conditions.   To obtain the signal back from the
    spline representation mirror-symmetric-convolve these coefficients with a
    length 3 FIR window [1.0, 4.0, 1.0]/ 6.0 .

    Parameters
    ----------
    signal : ndarray
        A rank-1 array representing samples of a signal.
    lamb : float, optional
        Smoothing coefficient, default is 0.0.

    Returns
    -------
    c : ndarray
        Cubic spline coefficients.

    """
    if lamb != 0.0:
        return _cubic_smooth_coeff(signal,lamb)
    else:
        return _cubic_coeff(signal)

def qspline1d(signal, lamb=0.0):
    """Compute quadratic spline coefficients for rank-1 array.

    Find the quadratic spline coefficients for a 1-D signal assuming
    mirror-symmetric boundary conditions.   To obtain the signal back from the
    spline representation mirror-symmetric-convolve these coefficients with a
    length 3 FIR window [1.0, 6.0, 1.0]/ 8.0 .

    Parameters
    ----------
    signal : ndarray
        A rank-1 array representing samples of a signal.
    lamb : float, optional
        Smoothing coefficient (must be zero for now).

    Returns
    -------
    c : ndarray
        Cubic spline coefficients.
    """
    if lamb != 0.0:
        raise ValueError, "Smoothing quadratic splines not supported yet."
    else:
        return _quadratic_coeff(signal)

def cspline1d_eval(cj, newx, dx=1.0, x0=0):
    """Evaluate a spline at the new set of points.

    `dx` is the old sample-spacing while `x0` was the old origin.  In
    other-words the old-sample points (knot-points) for which the `cj`
    represent spline coefficients were at equally-spaced points of:

      oldx = x0 + j*dx  j=0...N-1, with N=len(cj)

    Edges are handled using mirror-symmetric boundary conditions.

    """
    newx = (asarray(newx)-x0)/float(dx)
    res = zeros_like(newx)
    if (res.size == 0):
        return res
    N = len(cj)
    cond1 = newx < 0
    cond2 = newx > (N-1)
    cond3 = ~(cond1 | cond2)
    # handle general mirror-symmetry
    res[cond1] = cspline1d_eval(cj, -newx[cond1])
    res[cond2] = cspline1d_eval(cj, 2*(N-1)-newx[cond2])
    newx = newx[cond3]
    if newx.size == 0:
        return res
    result = zeros_like(newx)
    jlower = floor(newx-2).astype(int)+1
    for i in range(4):
        thisj = jlower + i
        indj = thisj.clip(0,N-1) # handle edge cases
        result += cj[indj] * cubic(newx - thisj)
    res[cond3] = result
    return res

def qspline1d_eval(cj, newx, dx=1.0, x0=0):
    """Evaluate a quadratic spline at the new set of points.

    `dx` is the old sample-spacing while `x0` was the old origin.  In
    other-words the old-sample points (knot-points) for which the `cj`
    represent spline coefficients were at equally-spaced points of:

      oldx = x0 + j*dx  j=0...N-1, with N=len(cj)

    Edges are handled using mirror-symmetric boundary conditions.

    """
    newx = (asarray(newx)-x0)/dx
    res = zeros_like(newx)
    if (res.size == 0):
        return res
    N = len(cj)
    cond1 = newx < 0
    cond2 = newx > (N-1)
    cond3 = ~(cond1 | cond2)
    # handle general mirror-symmetry
    res[cond1] = qspline1d_eval(cj, -newx[cond1])
    res[cond2] = qspline1d_eval(cj, 2*(N-1)-newx[cond2])
    newx = newx[cond3]
    if newx.size == 0:
        return res
    result = zeros_like(newx)
    jlower = floor(newx-1.5).astype(int)+1
    for i in range(3):
        thisj = jlower + i
        indj = thisj.clip(0,N-1) # handle edge cases
        result += cj[indj] * quadratic(newx - thisj)
    res[cond3] = result
    return res
<commit_msg>ENH: signal: In bsplines.py, don't use 'import *', and update 'raise' statements.<commit_after>
import scipy.special
from numpy import logical_and, asarray, pi, zeros_like, \
     piecewise, array, arctan2, tan, zeros, arange, floor
from numpy.core.umath import sqrt, exp, greater, less, cos, add, sin, \
     less_equal, greater_equal

# From splinemodule.c
from spline import cspline2d, qspline2d, sepfir2d, symiirorder1, symiirorder2

from scipy.misc import comb

gamma = scipy.special.gamma

def factorial(n):
    return gamma(n+1)

def spline_filter(Iin, lmbda=5.0):
    """Smoothing spline (cubic) filtering of a rank-2 array.

    Filter an input data set, `Iin`, using a (cubic) smoothing spline of
    fall-off `lmbda`.
    """
    intype = Iin.dtype.char
    hcol = array([1.0,4.0,1.0],'f')/6.0
    if intype in ['F','D']:
        Iin = Iin.astype('F')
        ckr = cspline2d(Iin.real,lmbda)
        cki = cspline2d(Iin.imag,lmbda)
        outr = sepfir2d(ckr,hcol,hcol)
        outi = sepfir2d(cki,hcol,hcol)
        out = (outr + 1j*outi).astype(intype)
    elif intype in ['f','d']:
        ckr = cspline2d(Iin,lmbda)
        out = sepfir2d(ckr, hcol, hcol)
        out = out.astype(intype)
    else:
        raise TypeError("Invalid data type for Iin")
    return out

_splinefunc_cache = {}

def _bspline_piecefunctions(order):
    """Returns the function defined over the left-side pieces for a bspline of
    a given order.

    The 0th piece is the first one less than 0.  The last piece is a function
    identical to 0 (returned as the constant 0).  (There are order//2 + 2 total
    pieces).

    Also returns the condition functions that when evaluated return boolean
    arrays for use with `numpy.piecewise`.
    """
    try:
        return _splinefunc_cache[order]
    except KeyError:
        pass

    def condfuncgen(num, val1, val2):
        if num == 0:
            return lambda x: logical_and(less_equal(x, val1),
                                         greater_equal(x, val2))
        elif num == 2:
            return lambda x: less_equal(x, val2)
        else:
            return lambda x: logical_and(less(x, val1),
                                         greater_equal(x, val2))

    last = order // 2 + 2
    if order % 2:
        startbound = -1.0
    else:
        startbound = -0.5
    condfuncs = [condfuncgen(0, 0, startbound)]
    bound = startbound
    for num in xrange(1,last-1):
        condfuncs.append(condfuncgen(1, bound, bound-1))
        bound = bound-1
    condfuncs.append(condfuncgen(2, 0, -(order+1)/2.0))

    # final value of bound is used in piecefuncgen below

    # the functions to evaluate are taken from the left-hand-side
    #  in the general expression derived from the central difference
    #  operator (because they involve fewer terms).

    fval = factorial(order)
    def piecefuncgen(num):
        Mk = order // 2 - num
        if (Mk < 0): return 0  # final function is 0
        coeffs = [(1-2*(k%2))*float(comb(order+1, k, exact=1))/fval for k in xrange(Mk+1)]
        shifts = [-bound - k for k in xrange(Mk+1)]
        #print "Adding piece number %d with coeffs %s and shifts %s" % (num, str(coeffs), str(shifts))
        def thefunc(x):
            res = 0.0
            for k in range(Mk+1):
                res += coeffs[k]*(x+shifts[k])**order
            return res
        return thefunc

    funclist = [piecefuncgen(k) for k in xrange(last)]

    _splinefunc_cache[order] = (funclist, condfuncs)

    return funclist, condfuncs

def bspline(x, n):
    """B-spline basis function of order n.

    Notes
    -----
    Uses numpy.piecewise and automatic function-generator.

    """
    ax = -abs(asarray(x))
    # number of pieces on the left-side is (n+1)/2
    funclist, condfuncs = _bspline_piecefunctions(n)
    condlist = [func(ax) for func in condfuncs]
    return piecewise(ax, condlist, funclist)

def gauss_spline(x, n):
    """Gaussian approximation to B-spline basis function of order n.
    """
    signsq = (n+1) / 12.0
    return 1/sqrt(2*pi*signsq) * exp(-x**2 / 2 / signsq)

def cubic(x):
    """A cubic B-spline.

    This is a special case of `bspline`, and equivalent to ``bspline(x, 3)``.
    """
    ax = abs(asarray(x))
    res = zeros_like(ax)
    cond1 = less(ax, 1)
    if cond1.any():
        ax1 = ax[cond1]
        res[cond1] = 2.0/3 - 1.0/2*ax1**2 * (2-ax1)
    cond2 = ~cond1 & less(ax, 2)
    if cond2.any():
        ax2 = ax[cond2]
        res[cond2] = 1.0/6*(2-ax2)**3
    return res

def quadratic(x):
    """A quadratic B-spline.

    This is a special case of `bspline`, and equivalent to ``bspline(x, 2)``.
    """
    ax = abs(asarray(x))
    res = zeros_like(ax)
    cond1 = less(ax, 0.5)
    if cond1.any():
        ax1 = ax[cond1]
        res[cond1] = 0.75-ax1**2
    cond2 = ~cond1 & less(ax, 1.5)
    if cond2.any():
        ax2 = ax[cond2]
        res[cond2] = (ax2-1.5)**2 / 2.0
    return res

def c0_P(order):
    # values taken from Unser, et.al. 1993 IEEE
    if order == 0:
        c0 = 1
        P = array([1])
    elif order == 1:
        c0 = 1
        P = array([0,1])
    elif order == 2:
        c0 = 8
        P = array([1,6,1])
    elif order == 3:
        c0 = 6
        P = array([1,4,1])
    elif order == 4:
        c0 = 384
        P = array([1,76,230,76,1])
    elif order == 5:
        c0 = 120
        P = array([1,26,66,26,1])
    elif order == 6:
        c0 = 46080
        P = array([1,722,10543,23548, 10543, 722, 1])
    elif order == 7:
        c0 = 5040
        P = array([1,120,1191,2416,1191, 120, 1])
    else:
        raise ValueError("Unknown order %d" % order)

def _coeff_smooth(lam):
    xi = 1 - 96*lam + 24*lam * sqrt(3 + 144*lam)
    omeg = arctan2(sqrt(144*lam-1),sqrt(xi))
    rho = (24*lam - 1 - sqrt(xi)) / (24*lam)
    rho = rho * sqrt((48*lam + 24*lam * sqrt(3+144*lam))/xi)
    return rho,omeg

def _hc(k, cs, rho, omega):
    return cs / sin(omega) * (rho**k)*sin(omega*(k+1))*(greater(k,-1))

def _hs(k, cs, rho, omega):
    c0 = cs*cs * (1 + rho*rho) / (1 - rho*rho) / (1-2*rho*rho*cos(2*omega) + rho**4)
    gamma = (1-rho*rho) / (1+rho*rho) / tan(omega)
    ak = abs(k)
    return c0 * rho**ak * (cos(omega*ak) + gamma*sin(omega*ak))

def _cubic_smooth_coeff(signal, lamb):
    rho, omega = _coeff_smooth(lamb)
    cs = 1-2*rho*cos(omega) + rho*rho
    K = len(signal)
    yp = zeros((K,),signal.dtype.char)
    k = arange(K)
    yp[0] = _hc(0,cs,rho,omega)*signal[0] + \
            add.reduce(_hc(k+1,cs,rho,omega)*signal)

    yp[1] = _hc(0,cs,rho,omega)*signal[0] + \
            _hc(1,cs,rho,omega)*signal[1] + \
            add.reduce(_hc(k+2,cs,rho,omega)*signal)

    for n in range(2,K):
        yp[n] = cs * signal[n] + 2*rho*cos(omega)*yp[n-1] - rho*rho*yp[n-2]

    y = zeros((K,),signal.dtype.char)

    y[K-1] = add.reduce((_hs(k,cs,rho,omega) + _hs(k+1,cs,rho,omega))*signal[::-1])
    y[K-2] = add.reduce((_hs(k-1,cs,rho,omega) + _hs(k+2,cs,rho,omega))*signal[::-1])

    for n in range(K-3,-1,-1):
        y[n] = cs*yp[n] + 2*rho*cos(omega)*y[n+1] - rho*rho*y[n+2]

    return y

def _cubic_coeff(signal):
    zi = -2 + sqrt(3)
    K = len(signal)
    yplus = zeros((K,),signal.dtype.char)
    powers = zi**arange(K)
    yplus[0] = signal[0] + zi*add.reduce(powers*signal)
    for k in range(1,K):
        yplus[k] = signal[k] + zi*yplus[k-1]
    output = zeros((K,),signal.dtype)
    output[K-1] = zi / (zi-1)*yplus[K-1]
    for k in range(K-2,-1,-1):
        output[k] = zi*(output[k+1]-yplus[k])
    return output*6.0

def _quadratic_coeff(signal):
    zi = -3 + 2*sqrt(2.0)
    K = len(signal)
    yplus = zeros((K,),signal.dtype.char)
    powers = zi**arange(K)
    yplus[0] = signal[0] + zi*add.reduce(powers*signal)
    for k in range(1,K):
        yplus[k] = signal[k] + zi*yplus[k-1]
    output = zeros((K,),signal.dtype.char)
    output[K-1] = zi / (zi-1)*yplus[K-1]
    for k in range(K-2,-1,-1):
        output[k] = zi*(output[k+1]-yplus[k])
    return output*8.0

def cspline1d(signal, lamb=0.0):
    """
    Compute cubic spline coefficients for rank-1 array.

    Find the cubic spline coefficients for a 1-D signal assuming
    mirror-symmetric boundary conditions.   To obtain the signal back from the
    spline representation mirror-symmetric-convolve these coefficients with a
    length 3 FIR window [1.0, 4.0, 1.0]/ 6.0 .

    Parameters
    ----------
    signal : ndarray
        A rank-1 array representing samples of a signal.
    lamb : float, optional
        Smoothing coefficient, default is 0.0.

    Returns
    -------
    c : ndarray
        Cubic spline coefficients.

    """
    if lamb != 0.0:
        return _cubic_smooth_coeff(signal,lamb)
    else:
        return _cubic_coeff(signal)

def qspline1d(signal, lamb=0.0):
    """Compute quadratic spline coefficients for rank-1 array.

    Find the quadratic spline coefficients for a 1-D signal assuming
    mirror-symmetric boundary conditions.   To obtain the signal back from the
    spline representation mirror-symmetric-convolve these coefficients with a
    length 3 FIR window [1.0, 6.0, 1.0]/ 8.0 .

    Parameters
    ----------
    signal : ndarray
        A rank-1 array representing samples of a signal.
    lamb : float, optional
        Smoothing coefficient (must be zero for now).

    Returns
    -------
    c : ndarray
        Cubic spline coefficients.
    """
    if lamb != 0.0:
        raise ValueError("Smoothing quadratic splines not supported yet.")
    else:
        return _quadratic_coeff(signal)

def cspline1d_eval(cj, newx, dx=1.0, x0=0):
    """Evaluate a spline at the new set of points.

    `dx` is the old sample-spacing while `x0` was the old origin.  In
    other-words the old-sample points (knot-points) for which the `cj`
    represent spline coefficients were at equally-spaced points of:

      oldx = x0 + j*dx  j=0...N-1, with N=len(cj)

    Edges are handled using mirror-symmetric boundary conditions.

    """
    newx = (asarray(newx)-x0)/float(dx)
    res = zeros_like(newx)
    if (res.size == 0):
        return res
    N = len(cj)
    cond1 = newx < 0
    cond2 = newx > (N-1)
    cond3 = ~(cond1 | cond2)
    # handle general mirror-symmetry
    res[cond1] = cspline1d_eval(cj, -newx[cond1])
    res[cond2] = cspline1d_eval(cj, 2*(N-1)-newx[cond2])
    newx = newx[cond3]
    if newx.size == 0:
        return res
    result = zeros_like(newx)
    jlower = floor(newx-2).astype(int)+1
    for i in range(4):
        thisj = jlower + i
        indj = thisj.clip(0,N-1) # handle edge cases
        result += cj[indj] * cubic(newx - thisj)
    res[cond3] = result
    return res

def qspline1d_eval(cj, newx, dx=1.0, x0=0):
    """Evaluate a quadratic spline at the new set of points.

    `dx` is the old sample-spacing while `x0` was the old origin.  In
    other-words the old-sample points (knot-points) for which the `cj`
    represent spline coefficients were at equally-spaced points of:

      oldx = x0 + j*dx  j=0...N-1, with N=len(cj)

    Edges are handled using mirror-symmetric boundary conditions.

    """
    newx = (asarray(newx)-x0)/dx
    res = zeros_like(newx)
    if (res.size == 0):
        return res
    N = len(cj)
    cond1 = newx < 0
    cond2 = newx > (N-1)
    cond3 = ~(cond1 | cond2)
    # handle general mirror-symmetry
    res[cond1] = qspline1d_eval(cj, -newx[cond1])
    res[cond2] = qspline1d_eval(cj, 2*(N-1)-newx[cond2])
    newx = newx[cond3]
    if newx.size == 0:
        return res
    result = zeros_like(newx)
    jlower = floor(newx-1.5).astype(int)+1
    for i in range(3):
        thisj = jlower + i
        indj = thisj.clip(0,N-1) # handle edge cases
        result += cj[indj] * quadratic(newx - thisj)
    res[cond3] = result
    return res
<|endoftext|>


<commit_before># Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.

from parlai.core.agents import Agent
from parlai.core.dict import DictionaryAgent

from torch.autograd import Variable
from torch import optim
import torch.nn as nn
import torch
import copy
import os
import random

class Seq2seqAgent(Agent):
    """Simple agent which uses an RNN to process incoming text observations.
    The RNN generates a vector which is used to represent the input text,
    conditioning on the context to generate an output token-by-token.

    For more information, see Sequence to Sequence Learning with Neural Networks
    `(Sutskever et al. 2014) <https://arxiv.org/abs/1409.3215>`_.
    """

    @staticmethod
    def add_cmdline_args(argparser):
        """Add command-line arguments specifically for this agent."""
        DictionaryAgent.add_cmdline_args(argparser)
        agent = argparser.add_argument_group('Seq2Seq Arguments')
        agent.add_argument('-hs', '--hiddensize', type=int, default=64,
            help='size of the hidden layers and embeddings')
        agent.add_argument('-nl', '--numlayers', type=int, default=2,
            help='number of hidden layers')
        agent.add_argument('-lr', '--learningrate', type=float, default=0.5,
            help='learning rate')
        agent.add_argument('-dr', '--dropout', type=float, default=0.1,
            help='dropout rate')
        agent.add_argument('--no-cuda', action='store_true', default=False,
            help='disable GPUs even if available')
        agent.add_argument('--gpu', type=int, default=-1,
            help='which GPU device to use')

    def __init__(self, opt, shared=None):
        # initialize defaults first
        super().__init__(opt, shared)
        if not shared:
            # this is not a shared instance of this class, so do full
            # initialization. if shared is set, only set up shared members.
            if opt.get('model_file') and os.path.isfile(opt['model_file']):
                # load model parameters if available
                print('Loading existing model params from ' + opt['model_file'])
                new_opt, self.states = self.load(opt['model_file'])
                # override options with stored ones
                opt = self.override_opt(new_opt)

            self.dict = DictionaryAgent(opt)
            self.id = 'Seq2Seq'
            # we use EOS markers to break input and output and end our output
            self.EOS = self.dict.eos_token
            self.observation = {'text': self.EOS, 'episode_done': True}
            self.EOS_TENSOR = torch.LongTensor(self.dict.parse(self.EOS))

            # store important params directly
            hsz = opt['hiddensize']
            self.hidden_size = hsz
            self.num_layers = opt['numlayers']
            self.learning_rate = opt['learningrate']
            self.longest_label = 1

            # set up modules
            self.criterion = nn.NLLLoss()
            # lookup table stores word embeddings
            self.lt = nn.Embedding(len(self.dict), hsz, padding_idx=0,
                                   scale_grad_by_freq=True)
            # encoder captures the input text
            self.encoder = nn.GRU(hsz, hsz, opt['numlayers'])
            # decoder produces our output states
            self.decoder = nn.GRU(hsz, hsz, opt['numlayers'])
            # linear layer helps us produce outputs from final decoder state
            self.h2o = nn.Linear(hsz, len(self.dict))
            # droput on the linear layer helps us generalize
            self.dropout = nn.Dropout(opt['dropout'])
            # softmax maps output scores to probabilities
            self.softmax = nn.LogSoftmax()

            # set up optims for each module
            lr = opt['learningrate']
            self.optims = {
                'lt': optim.SGD(self.lt.parameters(), lr=lr),
                'encoder': optim.SGD(self.encoder.parameters(), lr=lr),
                'decoder': optim.SGD(self.decoder.parameters(), lr=lr),
                'h2o': optim.SGD(self.h2o.parameters(), lr=lr),
            }

            if hasattr(self, 'states'):
                # set loaded states if applicable
                self.set_states(self.states)

            # check for cuda
            self.use_cuda = not opt.get('no_cuda') and torch.cuda.is_available()
            if self.use_cuda:
                print('[ Using CUDA ]')
                torch.cuda.set_device(opt['gpu'])
            if self.use_cuda:
                self.cuda()

        self.episode_done = True

    def override_opt(self, new_opt):
        """Print out each added key and each overriden key."""
        for k, v in new_opt.items():
            if k not in self.opt:
                print('Adding new option [ {k}: {v} ]'.format(k=k, v=v))
            elif self.opt[k] != v:
                print('Overriding option [ {k}: {old} => {v}]'.format(
                      k=k, old=self.opt[k], v=v))
            self.opt[k] = v
        return self.opt

    def parse(self, text):
        return torch.LongTensor(self.dict.txt2vec(text))

    def v2t(self, vec):
        return self.dict.vec2txt(vec)

    def cuda(self):
        self.EOS_TENSOR = self.EOS_TENSOR.cuda(async=True)
        self.criterion.cuda()
        self.lt.cuda()
        self.encoder.cuda()
        self.decoder.cuda()
        self.h2o.cuda()
        self.dropout.cuda()
        self.softmax.cuda()

    def hidden_to_idx(self, hidden, dropout=False):
        """Converts hidden state vectors into indices into the dictionary."""
        if hidden.size(0) > 1:
            raise RuntimeError('bad dimensions of tensor:', hidden)
        hidden = hidden.squeeze(0)
        scores = self.h2o(hidden)
        if dropout:
            scores = self.dropout(scores)
        scores = self.softmax(scores)
        _max_score, idx = scores.max(1)
        return idx, scores

    def zero_grad(self):
        for optimizer in self.optims.values():
            optimizer.zero_grad()

    def update_params(self):
        for optimizer in self.optims.values():
            optimizer.step()

    def init_zeros(self, bsz=1):
        t = torch.zeros(self.num_layers, bsz, self.hidden_size)
        if self.use_cuda:
            t = t.cuda(async=True)
        return Variable(t)

    def init_rand(self, bsz=1):
        t = torch.FloatTensor(self.num_layers, bsz, self.hidden_size)
        t.uniform_(0.05)
        if self.use_cuda:
            t = t.cuda(async=True)
        return Variable(t)

    def observe(self, observation):
        observation = copy.deepcopy(observation)
        if not self.episode_done:
            # if the last example wasn't the end of an episode, then we need to
            # recall what was said in that example
            prev_dialogue = self.observation['text']
            observation['text'] = prev_dialogue + '\n' + observation['text']
        self.observation = observation
        self.episode_done = observation['episode_done']
        return observation

    def predict(self, xs, ys=None):
        """Produce a prediction from our model. Update the model using the
        targets if available.
        """
        batchsize = len(xs)

        # first encode context
        xes = self.lt(xs).t()
        h0 = self.init_zeros(batchsize)
        _output, hn = self.encoder(xes, h0)

        # next we use EOS as an input to kick off our decoder
        x = Variable(self.EOS_TENSOR)
        xe = self.lt(x).unsqueeze(1)
        xes = xe.expand(xe.size(0), batchsize, xe.size(2))

        # list of output tokens for each example in the batch
        output_lines = [[] for _ in range(batchsize)]

        if ys is not None:
            # update the model based on the labels
            self.zero_grad()
            loss = 0
            # keep track of longest label we've ever seen
            self.longest_label = max(self.longest_label, ys.size(1))
            for i in range(ys.size(1)):
                output, hn = self.decoder(xes, hn)
                preds, scores = self.hidden_to_idx(output, dropout=True)
                y = ys.select(1, i)
                loss += self.criterion(scores, y)
                # use the true token as the next input instead of predicted
                # this produces a biased prediction but better training
                xes = self.lt(y).unsqueeze(0)
                for b in range(batchsize):
                    # convert the output scores to tokens
                    token = self.v2t([preds.data[b][0]])
                    output_lines[b].append(token)

            loss.backward()
            self.update_params()
        else:
            # just produce a prediction without training the model
            done = [False for _ in range(batchsize)]
            total_done = 0
            max_len = 0

            while(total_done < batchsize) and max_len < self.longest_label:
                # keep producing tokens until we hit EOS or max length for each
                # example in the batch
                output, hn = self.decoder(xes, hn)
                preds, scores = self.hidden_to_idx(output, dropout=False)
                xes = self.lt(preds.t())
                max_len += 1
                for b in range(batchsize):
                    if not done[b]:
                        # only add more tokens for examples that aren't done yet
                        token = self.v2t(preds.data[b])
                        if token == self.EOS:
                            # if we produced EOS, we're done
                            done[b] = True
                            total_done += 1
                        else:
                            output_lines[b].append(token)
            if random.random() < 0.1:
                print('prediction:', ' '.join(output_lines[0]))

        return output_lines

    def batchify(self, observations):
        """Convert a list of observations into input & target tensors."""
        # valid examples
        exs = [ex for ex in observations if 'text' in ex]
        # the indices of the valid (non-empty) tensors
        valid_inds = [i for i, ex in enumerate(observations) if 'text' in ex]

        # set up the input tensors
        batchsize = len(exs)
        # tokenize the text
        parsed = [self.parse(ex['text']) for ex in exs]
        max_x_len = max([len(x) for x in parsed])
        xs = torch.LongTensor(batchsize, max_x_len).fill_(0)
        # pack the data to the right side of the tensor for this model
        for i, x in enumerate(parsed):
            offset = max_x_len - len(x)
            for j, idx in enumerate(x):
                xs[i][j + offset] = idx
        if self.use_cuda:
            xs = xs.cuda(async=True)
        xs = Variable(xs)

        # set up the target tensors
        ys = None
        if 'labels' in exs[0]:
            # randomly select one of the labels to update on, if multiple
            # append EOS to each label
            labels = [random.choice(ex['labels']) + ' ' + self.EOS for ex in exs]
            parsed = [self.parse(y) for y in labels]
            max_y_len = max(len(y) for y in parsed)
            ys = torch.LongTensor(batchsize, max_y_len).fill_(0)
            for i, y in enumerate(parsed):
                for j, idx in enumerate(y):
                    ys[i][j] = idx
            if self.use_cuda:
                ys = ys.cuda(async=True)
            ys = Variable(ys)
        return xs, ys, valid_inds

    def batch_act(self, observations):
        batchsize = len(observations)
        # initialize a table of replies with this agent's id
        batch_reply = [{'id': self.getID()} for _ in range(batchsize)]

        # convert the observations into batches of inputs and targets
        # valid_inds tells us the indices of all valid examples
        # e.g. for input [{}, {'text': 'hello'}, {}, {}], valid_inds is [1]
        # since the other three elements had no 'text' field
        xs, ys, valid_inds = self.batchify(observations)

        if len(xs) == 0:
            # no valid examples, just return the empty responses we set up
            return batch_reply

        # produce predictions either way, but use the targets if available
        predictions = self.predict(xs, ys)

        for i in range(len(predictions)):
            # map the predictions back to non-empty examples in the batch
            # we join with spaces since we produce tokens one at a time
            batch_reply[valid_inds[i]]['text'] = ' '.join(
                c for c in predictions[i] if c != self.EOS)

        return batch_reply

    def act(self):
        # call batch_act with this batch of one
        return self.batch_act([self.observation])[0]

    def save(self, path=None):
        path = self.opt.get('model_file', None) if path is None else path

        if path:
            model = {}
            model['lt'] = self.lt.state_dict()
            model['encoder'] = self.encoder.state_dict()
            model['decoder'] = self.decoder.state_dict()
            model['h2o'] = self.h2o.state_dict()
            model['longest_label'] = self.longest_label
            model['opt'] = self.opt

            with open(path, 'wb') as write:
                torch.save(model, write)

    def load(self, path):
        """Return opt and model states."""
        with open(path, 'rb') as read:
            model = torch.load(read)

        return model['opt'], model

    def set_states(self, states):
        """Set the state dicts of the modules from saved states."""
        self.lt.load_state_dict(states['lt'])
        self.encoder.load_state_dict(states['encoder'])
        self.decoder.load_state_dict(states['decoder'])
        self.h2o.load_state_dict(states['h2o'])
        self.longest_label = states['longest_label']
<commit_msg>quick eos fix<commit_after># Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.

from parlai.core.agents import Agent
from parlai.core.dict import DictionaryAgent

from torch.autograd import Variable
from torch import optim
import torch.nn as nn
import torch
import copy
import os
import random

class Seq2seqAgent(Agent):
    """Simple agent which uses an RNN to process incoming text observations.
    The RNN generates a vector which is used to represent the input text,
    conditioning on the context to generate an output token-by-token.

    For more information, see Sequence to Sequence Learning with Neural Networks
    `(Sutskever et al. 2014) <https://arxiv.org/abs/1409.3215>`_.
    """

    @staticmethod
    def add_cmdline_args(argparser):
        """Add command-line arguments specifically for this agent."""
        DictionaryAgent.add_cmdline_args(argparser)
        agent = argparser.add_argument_group('Seq2Seq Arguments')
        agent.add_argument('-hs', '--hiddensize', type=int, default=64,
            help='size of the hidden layers and embeddings')
        agent.add_argument('-nl', '--numlayers', type=int, default=2,
            help='number of hidden layers')
        agent.add_argument('-lr', '--learningrate', type=float, default=0.5,
            help='learning rate')
        agent.add_argument('-dr', '--dropout', type=float, default=0.1,
            help='dropout rate')
        agent.add_argument('--no-cuda', action='store_true', default=False,
            help='disable GPUs even if available')
        agent.add_argument('--gpu', type=int, default=-1,
            help='which GPU device to use')

    def __init__(self, opt, shared=None):
        # initialize defaults first
        super().__init__(opt, shared)
        if not shared:
            # this is not a shared instance of this class, so do full
            # initialization. if shared is set, only set up shared members.
            if opt.get('model_file') and os.path.isfile(opt['model_file']):
                # load model parameters if available
                print('Loading existing model params from ' + opt['model_file'])
                new_opt, self.states = self.load(opt['model_file'])
                # override options with stored ones
                opt = self.override_opt(new_opt)

            self.dict = DictionaryAgent(opt)
            self.id = 'Seq2Seq'
            # we use EOS markers to break input and output and end our output
            self.EOS = self.dict.end_token
            self.observation = {'text': self.EOS, 'episode_done': True}
            self.EOS_TENSOR = torch.LongTensor(self.dict.parse(self.EOS))

            # store important params directly
            hsz = opt['hiddensize']
            self.hidden_size = hsz
            self.num_layers = opt['numlayers']
            self.learning_rate = opt['learningrate']
            self.longest_label = 1

            # set up modules
            self.criterion = nn.NLLLoss()
            # lookup table stores word embeddings
            self.lt = nn.Embedding(len(self.dict), hsz, padding_idx=0,
                                   scale_grad_by_freq=True)
            # encoder captures the input text
            self.encoder = nn.GRU(hsz, hsz, opt['numlayers'])
            # decoder produces our output states
            self.decoder = nn.GRU(hsz, hsz, opt['numlayers'])
            # linear layer helps us produce outputs from final decoder state
            self.h2o = nn.Linear(hsz, len(self.dict))
            # droput on the linear layer helps us generalize
            self.dropout = nn.Dropout(opt['dropout'])
            # softmax maps output scores to probabilities
            self.softmax = nn.LogSoftmax()

            # set up optims for each module
            lr = opt['learningrate']
            self.optims = {
                'lt': optim.SGD(self.lt.parameters(), lr=lr),
                'encoder': optim.SGD(self.encoder.parameters(), lr=lr),
                'decoder': optim.SGD(self.decoder.parameters(), lr=lr),
                'h2o': optim.SGD(self.h2o.parameters(), lr=lr),
            }

            if hasattr(self, 'states'):
                # set loaded states if applicable
                self.set_states(self.states)

            # check for cuda
            self.use_cuda = not opt.get('no_cuda') and torch.cuda.is_available()
            if self.use_cuda:
                print('[ Using CUDA ]')
                torch.cuda.set_device(opt['gpu'])
            if self.use_cuda:
                self.cuda()

        self.episode_done = True

    def override_opt(self, new_opt):
        """Print out each added key and each overriden key."""
        for k, v in new_opt.items():
            if k not in self.opt:
                print('Adding new option [ {k}: {v} ]'.format(k=k, v=v))
            elif self.opt[k] != v:
                print('Overriding option [ {k}: {old} => {v}]'.format(
                      k=k, old=self.opt[k], v=v))
            self.opt[k] = v
        return self.opt

    def parse(self, text):
        return torch.LongTensor(self.dict.txt2vec(text))

    def v2t(self, vec):
        return self.dict.vec2txt(vec)

    def cuda(self):
        self.EOS_TENSOR = self.EOS_TENSOR.cuda(async=True)
        self.criterion.cuda()
        self.lt.cuda()
        self.encoder.cuda()
        self.decoder.cuda()
        self.h2o.cuda()
        self.dropout.cuda()
        self.softmax.cuda()

    def hidden_to_idx(self, hidden, dropout=False):
        """Converts hidden state vectors into indices into the dictionary."""
        if hidden.size(0) > 1:
            raise RuntimeError('bad dimensions of tensor:', hidden)
        hidden = hidden.squeeze(0)
        scores = self.h2o(hidden)
        if dropout:
            scores = self.dropout(scores)
        scores = self.softmax(scores)
        _max_score, idx = scores.max(1)
        return idx, scores

    def zero_grad(self):
        for optimizer in self.optims.values():
            optimizer.zero_grad()

    def update_params(self):
        for optimizer in self.optims.values():
            optimizer.step()

    def init_zeros(self, bsz=1):
        t = torch.zeros(self.num_layers, bsz, self.hidden_size)
        if self.use_cuda:
            t = t.cuda(async=True)
        return Variable(t)

    def init_rand(self, bsz=1):
        t = torch.FloatTensor(self.num_layers, bsz, self.hidden_size)
        t.uniform_(0.05)
        if self.use_cuda:
            t = t.cuda(async=True)
        return Variable(t)

    def observe(self, observation):
        observation = copy.deepcopy(observation)
        if not self.episode_done:
            # if the last example wasn't the end of an episode, then we need to
            # recall what was said in that example
            prev_dialogue = self.observation['text']
            observation['text'] = prev_dialogue + '\n' + observation['text']
        self.observation = observation
        self.episode_done = observation['episode_done']
        return observation

    def predict(self, xs, ys=None):
        """Produce a prediction from our model. Update the model using the
        targets if available.
        """
        batchsize = len(xs)

        # first encode context
        xes = self.lt(xs).t()
        h0 = self.init_zeros(batchsize)
        _output, hn = self.encoder(xes, h0)

        # next we use EOS as an input to kick off our decoder
        x = Variable(self.EOS_TENSOR)
        xe = self.lt(x).unsqueeze(1)
        xes = xe.expand(xe.size(0), batchsize, xe.size(2))

        # list of output tokens for each example in the batch
        output_lines = [[] for _ in range(batchsize)]

        if ys is not None:
            # update the model based on the labels
            self.zero_grad()
            loss = 0
            # keep track of longest label we've ever seen
            self.longest_label = max(self.longest_label, ys.size(1))
            for i in range(ys.size(1)):
                output, hn = self.decoder(xes, hn)
                preds, scores = self.hidden_to_idx(output, dropout=True)
                y = ys.select(1, i)
                loss += self.criterion(scores, y)
                # use the true token as the next input instead of predicted
                # this produces a biased prediction but better training
                xes = self.lt(y).unsqueeze(0)
                for b in range(batchsize):
                    # convert the output scores to tokens
                    token = self.v2t([preds.data[b][0]])
                    output_lines[b].append(token)

            loss.backward()
            self.update_params()
        else:
            # just produce a prediction without training the model
            done = [False for _ in range(batchsize)]
            total_done = 0
            max_len = 0

            while(total_done < batchsize) and max_len < self.longest_label:
                # keep producing tokens until we hit EOS or max length for each
                # example in the batch
                output, hn = self.decoder(xes, hn)
                preds, scores = self.hidden_to_idx(output, dropout=False)
                xes = self.lt(preds.t())
                max_len += 1
                for b in range(batchsize):
                    if not done[b]:
                        # only add more tokens for examples that aren't done yet
                        token = self.v2t(preds.data[b])
                        if token == self.EOS:
                            # if we produced EOS, we're done
                            done[b] = True
                            total_done += 1
                        else:
                            output_lines[b].append(token)
            if random.random() < 0.1:
                print('prediction:', ' '.join(output_lines[0]))

        return output_lines

    def batchify(self, observations):
        """Convert a list of observations into input & target tensors."""
        # valid examples
        exs = [ex for ex in observations if 'text' in ex]
        # the indices of the valid (non-empty) tensors
        valid_inds = [i for i, ex in enumerate(observations) if 'text' in ex]

        # set up the input tensors
        batchsize = len(exs)
        # tokenize the text
        parsed = [self.parse(ex['text']) for ex in exs]
        max_x_len = max([len(x) for x in parsed])
        xs = torch.LongTensor(batchsize, max_x_len).fill_(0)
        # pack the data to the right side of the tensor for this model
        for i, x in enumerate(parsed):
            offset = max_x_len - len(x)
            for j, idx in enumerate(x):
                xs[i][j + offset] = idx
        if self.use_cuda:
            xs = xs.cuda(async=True)
        xs = Variable(xs)

        # set up the target tensors
        ys = None
        if 'labels' in exs[0]:
            # randomly select one of the labels to update on, if multiple
            # append EOS to each label
            labels = [random.choice(ex['labels']) + ' ' + self.EOS for ex in exs]
            parsed = [self.parse(y) for y in labels]
            max_y_len = max(len(y) for y in parsed)
            ys = torch.LongTensor(batchsize, max_y_len).fill_(0)
            for i, y in enumerate(parsed):
                for j, idx in enumerate(y):
                    ys[i][j] = idx
            if self.use_cuda:
                ys = ys.cuda(async=True)
            ys = Variable(ys)
        return xs, ys, valid_inds

    def batch_act(self, observations):
        batchsize = len(observations)
        # initialize a table of replies with this agent's id
        batch_reply = [{'id': self.getID()} for _ in range(batchsize)]

        # convert the observations into batches of inputs and targets
        # valid_inds tells us the indices of all valid examples
        # e.g. for input [{}, {'text': 'hello'}, {}, {}], valid_inds is [1]
        # since the other three elements had no 'text' field
        xs, ys, valid_inds = self.batchify(observations)

        if len(xs) == 0:
            # no valid examples, just return the empty responses we set up
            return batch_reply

        # produce predictions either way, but use the targets if available
        predictions = self.predict(xs, ys)

        for i in range(len(predictions)):
            # map the predictions back to non-empty examples in the batch
            # we join with spaces since we produce tokens one at a time
            batch_reply[valid_inds[i]]['text'] = ' '.join(
                c for c in predictions[i] if c != self.EOS)

        return batch_reply

    def act(self):
        # call batch_act with this batch of one
        return self.batch_act([self.observation])[0]

    def save(self, path=None):
        path = self.opt.get('model_file', None) if path is None else path

        if path:
            model = {}
            model['lt'] = self.lt.state_dict()
            model['encoder'] = self.encoder.state_dict()
            model['decoder'] = self.decoder.state_dict()
            model['h2o'] = self.h2o.state_dict()
            model['longest_label'] = self.longest_label
            model['opt'] = self.opt

            with open(path, 'wb') as write:
                torch.save(model, write)

    def load(self, path):
        """Return opt and model states."""
        with open(path, 'rb') as read:
            model = torch.load(read)

        return model['opt'], model

    def set_states(self, states):
        """Set the state dicts of the modules from saved states."""
        self.lt.load_state_dict(states['lt'])
        self.encoder.load_state_dict(states['encoder'])
        self.decoder.load_state_dict(states['decoder'])
        self.h2o.load_state_dict(states['h2o'])
        self.longest_label = states['longest_label']
<|endoftext|>


<commit_before>import numpy
import os

from pylearn2.datasets import dense_design_matrix
from pylearn2.datasets import retina

class NORBSmall(dense_design_matrix.DenseDesignMatrix):
    """
    A pylearn2 dataset object for the small NORB dataset (v1.0).
    """

    @classmethod
    def load(cls, which_set, desc):

        assert desc in ['dat','cat','info']

        base = '%s/norb_small/original_npy/smallnorb-' % os.getenv('PYLEARN2_DATA_PATH')
        if which_set == 'train':
            base += '5x46789x9x18x6x2x96x96-training'
        else:
            base += '5x01235x9x18x6x2x96x96-testing'

        fp = open(base + '-%s.npy' % desc, 'r')
        data = numpy.load(fp)
        fp.close()

        return data

    def __init__(self, which_set, center=False, multi_target = False):
        """
        :param which_set: one of ['train','test']
        :param center: data is in range [0,256], center=True subtracts 127.5.
        :param multi_target: load extra information as additional labels.
        """
        assert which_set in ['train','test']

        X = NORBSmall.load(which_set, 'dat')

        # put things in pylearn2's DenseDesignMatrix format
        X = numpy.cast['float32'](X)
        X = X.reshape(-1, 2*96*96)

        #this is uint8
        y = NORBSmall.load(which_set, 'cat')
        if multi_target:
            y_extra = NORBSmall.load(which_set, 'info')
            y = numpy.hstack((y[:,numpy.newaxis],y_extra))

        if center:
            X -= 127.5

        view_converter = dense_design_matrix.DefaultViewConverter((96,96,2))

        super(NORBSmall,self).__init__(X = X, y = y, view_converter = view_converter)

class FoveatedNORB(dense_design_matrix.DenseDesignMatrix):

    @classmethod
    def load(cls, which_set):

        base = '%s/norb_small/foveated/smallnorb-' % os.getenv('PYLEARN2_DATA_PATH')
        if which_set == 'train':
            base += '5x46789x9x18x6x2x96x96-training-dat'
        else:
            base += '5x01235x9x18x6x2x96x96-testing-dat'

        data = numpy.load(base + '.npy', 'r')
        return data

    def __init__(self, which_set, center=False, multi_target = False, scale = False,
            start = None, stop = None, one_hot = False):
        """
        :param which_set: one of ['train','test']
        :param center: data is in range [0,256], center=True subtracts 127.5.
        :param multi_target: load extra information as additional labels.
        """
        if which_set not in ['train','test']:
            raise ValueError("Unrecognized which_set value: " + which_set)

        X = FoveatedNORB.load(which_set)

        # put things in pylearn2's DenseDesignMatrix format
        X = numpy.cast['float32'](X)

        #this is uint8
        y = NORBSmall.load(which_set, 'cat')
        if multi_target:
            y_extra = NORBSmall.load(which_set, 'info')
            y = numpy.hstack((y[:,numpy.newaxis],y_extra))

        if center:
            X -= 127.5
            if scale:
                X /= 127.5
        else:
            if scale:
                X /= 255.

        view_converter = retina.RetinaCodingViewConverter((96,96,2), (8,4,2,2))

        super(FoveatedNORB, self).__init__(X=X, y=y, view_converter=view_converter)

        if one_hot:
            self.convert_to_one_hot()

        self.restrict(start, stop)

        self.y = self.y.astype('float32')

<commit_msg>added restrict_instances parameter to FoveatedNORB<commit_after>import numpy
np = numpy
import os

from pylearn2.datasets import dense_design_matrix
from pylearn2.datasets import retina

class NORBSmall(dense_design_matrix.DenseDesignMatrix):
    """
    A pylearn2 dataset object for the small NORB dataset (v1.0).
    """

    @classmethod
    def load(cls, which_set, desc):

        assert desc in ['dat','cat','info']

        base = '%s/norb_small/original_npy/smallnorb-' % os.getenv('PYLEARN2_DATA_PATH')
        if which_set == 'train':
            base += '5x46789x9x18x6x2x96x96-training'
        else:
            base += '5x01235x9x18x6x2x96x96-testing'

        fp = open(base + '-%s.npy' % desc, 'r')
        data = numpy.load(fp)
        fp.close()

        return data

    def __init__(self, which_set, center=False, multi_target = False):
        """
        :param which_set: one of ['train','test']
        :param center: data is in range [0,256], center=True subtracts 127.5.
        :param multi_target: load extra information as additional labels.
        """
        assert which_set in ['train','test']

        X = NORBSmall.load(which_set, 'dat')

        # put things in pylearn2's DenseDesignMatrix format
        X = numpy.cast['float32'](X)
        X = X.reshape(-1, 2*96*96)

        #this is uint8
        y = NORBSmall.load(which_set, 'cat')
        if multi_target:
            y_extra = NORBSmall.load(which_set, 'info')
            y = numpy.hstack((y[:,numpy.newaxis],y_extra))

        if center:
            X -= 127.5

        view_converter = dense_design_matrix.DefaultViewConverter((96,96,2))

        super(NORBSmall,self).__init__(X = X, y = y, view_converter = view_converter)

class FoveatedNORB(dense_design_matrix.DenseDesignMatrix):

    @classmethod
    def load(cls, which_set):

        base = '%s/norb_small/foveated/smallnorb-' % os.getenv('PYLEARN2_DATA_PATH')
        if which_set == 'train':
            base += '5x46789x9x18x6x2x96x96-training-dat'
        else:
            base += '5x01235x9x18x6x2x96x96-testing-dat'

        data = numpy.load(base + '.npy', 'r')
        return data

    def __init__(self, which_set, center=False, scale = False,
            start = None, stop = None, one_hot = False, restrict_instances=None):
        """
        :param which_set: one of ['train','test']
        :param center: data is in range [0,256], center=True subtracts 127.5.
        :param multi_target: load extra information as additional labels.
        """
        if which_set not in ['train','test']:
            raise ValueError("Unrecognized which_set value: " + which_set)

        X = FoveatedNORB.load(which_set)

        # put things in pylearn2's DenseDesignMatrix format
        X = numpy.cast['float32'](X)

        #this is uint8
        y = NORBSmall.load(which_set, 'cat')
        y_extra = NORBSmall.load(which_set, 'info')

        assert y_extra.shape[0] == y.shape[0]
        instance = y_extra[:, 0]
        assert instance.min() >= 0
        assert instance.max() <= 9

        if restrict_instances is not None:
            mask = reduce(np.maximum, [instance == ins for ins in restrict_instances])
            mask = mask.astype('bool')
            X = X[mask,:]
            y = y[mask]
            assert X.shape[0] == y.shape[0]
            expected = sum([(instance == ins).sum() for ins in restrict_instances])
            assert X.shape[0] == expected

        if center:
            X -= 127.5
            if scale:
                X /= 127.5
        else:
            if scale:
                X /= 255.

        view_converter = retina.RetinaCodingViewConverter((96,96,2), (8,4,2,2))

        super(FoveatedNORB, self).__init__(X=X, y=y, view_converter=view_converter)

        if one_hot:
            self.convert_to_one_hot()

        self.restrict(start, stop)

        self.y = self.y.astype('float32')

<|endoftext|>


<commit_before>#! /usr/bin/env python
import os
import textwrap

import jinja2

from .bmi_metadata import load_bmi_metadata, ModelInfo
from ..babel import BabelConfigError

_DOCSTRING = u"""
Basic Model Interface for {{ name }}.

{{ desc|trim|wordwrap(70) if desc }}

Author: {{ author }}
Version: {{ version }}
License: {{ license }}
DOI: {{ doi }}
URL: {{ url }}

{% if cite_as -%}
Cite as:
{% for citation in cite_as %}
{{ citation|trim|indent(width=4, indentfirst=True) }}
{% endfor %}
{%- endif %}
{% if parameters %}
Parameters
----------
{% for param in parameters|sort(attribute='name') -%}
{{param.name}} : {{param.type}}, optional
    {{ "%s [default=%s %s]"|format(param.desc, param.value, param.units)|trim|wordwrap(70)|indent(4) }}
{% endfor %}
{% endif -%}

Examples
--------
>>> from pymt.components import {{name}}
>>> model = {{name}}()
>>> (fname, initdir) = model.setup()
>>> model.initialize(fname, dir=initdir)
>>> for _ in xrange(10):
...     model.update()
>>> model.finalize()
""".strip()

def bmi_docstring(name, author=None, version=None, license=None, doi=None,
                  url=None, parameters=None, summary=None, cite_as=None,
                  email=None):
    """Build the docstring for a BMI model.

    Parameters
    ----------
    name : str
        Name of a BMI component.
    author : str, optional
        Name of author or authors.
    version : str, optional
        Version string for the component.
    license : str, optional
        Name of the license of the component.
    doi : str, optional
        A DOI for the component.
    url : str, optional
        URL of the component's location on the internet.
    parameters : iterable, optional
        List of input parameters for the component. Each parameter object
        must have attributes for name, type, value, units, and desc.
    cite_as : iterable of str, optional
        List of citations for this component.
    email : str, optional
        Contact email address.

    Returns
    -------
    str
        The docstring.

    Examples
    --------
    >>> from pymt.framework.bmi_docstring import bmi_docstring
    >>> print bmi_docstring('Model', author='Walt Disney') #doctest: +ELLIPSIS
    Basic Model Interface for Model.
    ...
    """
    try:
        meta = load_bmi_metadata(name)
    except (IOError, BabelConfigError):
        info = ModelInfo(name=name, author=author, version=version,
                         license=license, doi=doi, url=url, summary=summary,
                         cite_as=cite_as, email=None)
        defaults = []
    else:
        info = meta['info']
        defaults = meta['defaults'].values()

    author = author or info.author
    email = email or '-'
    version = version or info.version
    license = license or info.license
    doi = doi or info.doi
    url = url or info.url
    summary = summary or info.summary
    cite_as = cite_as or info.cite_as
    parameters = parameters or defaults

    env = jinja2.Environment(loader=jinja2.DictLoader({'docstring': _DOCSTRING}))
    return env.get_template('docstring').render(
        desc=summary,
        name=name,
        parameters=parameters,
        author=author,
        version=version,
        license=license,
        doi=doi,
        url=url,
        cite_as=cite_as,
    )
<commit_msg>Render author list as one per line.<commit_after>#! /usr/bin/env python
import os
import textwrap

import jinja2

from .bmi_metadata import load_bmi_metadata, ModelInfo
from ..babel import BabelConfigError

_DOCSTRING = u"""
Basic Model Interface for {{ name }}.

{{ desc|trim|wordwrap(70) if desc }}

{% if author -%}
Author:
{%- for name in author %}
{{"- " + name|trim }}{% endfor %}
{%- endif %}
Version: {{ version }}
License: {{ license }}
DOI: {{ doi }}
URL: {{ url }}

{% if cite_as -%}
Cite as:
{% for citation in cite_as %}
{{ citation|trim|indent(width=4, indentfirst=True) }}
{% endfor %}
{%- endif %}
{% if parameters %}
Parameters
----------
{% for param in parameters|sort(attribute='name') -%}
{{param.name}} : {{param.type}}, optional
    {{ "%s [default=%s %s]"|format(param.desc, param.value, param.units)|trim|wordwrap(70)|indent(4) }}
{% endfor %}
{% endif -%}

Examples
--------
>>> from pymt.components import {{name}}
>>> model = {{name}}()
>>> (fname, initdir) = model.setup()
>>> model.initialize(fname, dir=initdir)
>>> for _ in xrange(10):
...     model.update()
>>> model.finalize()
""".strip()

def bmi_docstring(name, author=None, version=None, license=None, doi=None,
                  url=None, parameters=None, summary=None, cite_as=None,
                  email=None):
    """Build the docstring for a BMI model.

    Parameters
    ----------
    name : str
        Name of a BMI component.
    author : str, optional
        Name of author or authors.
    version : str, optional
        Version string for the component.
    license : str, optional
        Name of the license of the component.
    doi : str, optional
        A DOI for the component.
    url : str, optional
        URL of the component's location on the internet.
    parameters : iterable, optional
        List of input parameters for the component. Each parameter object
        must have attributes for name, type, value, units, and desc.
    cite_as : iterable of str, optional
        List of citations for this component.
    email : str, optional
        Contact email address.

    Returns
    -------
    str
        The docstring.

    Examples
    --------
    >>> from pymt.framework.bmi_docstring import bmi_docstring
    >>> print bmi_docstring('Model', author='Walt Disney') #doctest: +ELLIPSIS
    Basic Model Interface for Model.
    ...
    """
    try:
        meta = load_bmi_metadata(name)
    except (IOError, BabelConfigError):
        info = ModelInfo(name=name, author=author, version=version,
                         license=license, doi=doi, url=url, summary=summary,
                         cite_as=cite_as, email=None)
        defaults = []
    else:
        info = meta['info']
        defaults = meta['defaults'].values()

    author = author or info.author
    email = email or '-'
    version = version or info.version
    license = license or info.license
    doi = doi or info.doi
    url = url or info.url
    summary = summary or info.summary
    cite_as = cite_as or info.cite_as
    parameters = parameters or defaults

    env = jinja2.Environment(loader=jinja2.DictLoader({'docstring': _DOCSTRING}))
    return env.get_template('docstring').render(
        desc=summary,
        name=name,
        parameters=parameters,
        author=author,
        version=version,
        license=license,
        doi=doi,
        url=url,
        cite_as=cite_as,
    )
<|endoftext|>


<commit_before>#!/usr/bin/env python

__author__ = 'Dave Foster <dfoster@asascience.com>'
__license__ = 'Apache 2.0'

from pyon.util.unit_test import PyonTestCase
from pyon.util.int_test import IonIntegrationTestCase
from pyon.net.transport import NameTrio, BaseTransport, AMQPTransport, TransportError, TopicTrie, ZeroMQRouter
from pyon.core.bootstrap import get_sys_name

from nose.plugins.attrib import attr
from mock import Mock, MagicMock, sentinel
from gevent_zeromq import zmq
from gevent.event import Event
import time

@attr('UNIT')
class TestTransport(PyonTestCase):
    pass

@attr('UNIT')
class TestAMQPTransport(PyonTestCase):

    def setUp(self):
        self.client = MagicMock()

    def test__sync_call_no_ret_value(self):

        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam()

        tp = AMQPTransport()
        rv = tp._sync_call(self.client, async_func, 'callback')
        self.assertIsNone(rv)

    def test__sync_call_with_ret_value(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sentinel.val)

        tp = AMQPTransport()
        rv = tp._sync_call(self.client, async_func, 'callback')
        self.assertEquals(rv, sentinel.val)

    def test__sync_call_with_mult_rets(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sentinel.val, sentinel.val2)

        tp = AMQPTransport()
        rv = tp._sync_call(self.client, async_func, 'callback')
        self.assertEquals(rv, (sentinel.val, sentinel.val2))

    def test__sync_call_with_kwarg_rets(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sup=sentinel.val, sup2=sentinel.val2)

        tp = AMQPTransport()
        rv = tp._sync_call(self.client, async_func, 'callback')
        self.assertEquals(rv, {'sup':sentinel.val, 'sup2':sentinel.val2})

    def test__sync_call_with_normal_and_kwarg_rets(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sentinel.arg, sup=sentinel.val, sup2=sentinel.val2)

        tp = AMQPTransport()
        rv = tp._sync_call(self.client, async_func, 'callback')
        self.assertEquals(rv, (sentinel.arg, {'sup':sentinel.val, 'sup2':sentinel.val2}))

    def test__sync_call_with_error(self):
        tp = AMQPTransport()

        def async_func(*args, **kwargs):
            raise TransportError('haha')

        self.assertRaises(TransportError, tp._sync_call, self.client, async_func, 'callback')

@attr('UNIT')
class TestTopicTrie(PyonTestCase):
    def setUp(self):
        self.tt = TopicTrie()

    def test_many_matches(self):
        self.tt.add_topic_tree('a.b.c', sentinel.p1)
        self.tt.add_topic_tree('a.*.c', sentinel.p2)
        self.tt.add_topic_tree('*.b.c', sentinel.p3)
        self.tt.add_topic_tree('a.b.*', sentinel.p4)
        self.tt.add_topic_tree('*.*.c', sentinel.p5)
        self.tt.add_topic_tree('a.#',   sentinel.wild)
        self.tt.add_topic_tree('a.#.c', sentinel.middle_wild)

        self.assertEquals({sentinel.p1, sentinel.p2, sentinel.p3, sentinel.p4, sentinel.p5, sentinel.wild, sentinel.middle_wild},
                          set(self.tt.get_all_matches('a.b.c')))

        self.assertEquals({sentinel.p2, sentinel.p5, sentinel.wild, sentinel.middle_wild},
                          set(self.tt.get_all_matches('a.d.c')))

        self.assertEquals({sentinel.wild, sentinel.middle_wild},
                          set(self.tt.get_all_matches('a.b.b.b.b.b.b.c')))

        self.assertEquals({sentinel.wild},
                          set(self.tt.get_all_matches('a.b.b.b.b.b.b')))

@attr('UNIT')
class TestZeroMQRouter(PyonTestCase):
    pass

@attr('UNIT')
class TestZeroMQRouter(PyonTestCase):

    def setUp(self):
        self.zc = zmq.Context(1)
        self.zr = ZeroMQRouter(self.zc, get_sys_name())
        self.zr.start()
        self.addCleanup(self.zr.stop)

        self.pub = self.zc.socket(zmq.PUB)
        self.pub.connect(self.zr._connect_addr)

        # make a hook so we can tell when a message gets routed
        self.ev = Event()
        def new_route(*args, **kwargs):
            self.ev.set()
            self.zr._oldroute(*args, **kwargs)
        self.zr._oldroute = self.zr._route
        self.zr._route = new_route

    def test_publish_to_unknown_exchange(self):
        self.assertEquals(len(self.zr.errors), 0)
        self.pub.send_multipart(['ex', 'rkey', 'body', 'props'])
        self.ev.wait(timeout=10)
        self.assertEquals(len(self.zr.errors), 1)

    def test_publish_to_known_exchange(self):
        # declare exchange
        self.zr.declare_exchange('known')

        # send message
        self.pub.send_multipart(['known', 'rkey', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # message is binned but no errors
        self.assertEquals(len(self.zr.errors), 0)

    def test_publish_to_queue(self):
        # declare exchange/queue/binding
        self.zr.declare_exchange('known')
        self.zr.declare_queue('iamqueue')
        self.zr.bind('known', 'iamqueue', 'binzim')
        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 0)

        # send message
        self.pub.send_multipart(['known', 'binzim', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # should be sitting in a queue waiting
        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 1)
        self.assertIn(('known', 'binzim', 'body', 'props'), self.zr._queues['iamqueue'])

    def test_publish_to_many_queues(self):
        # declare exchange/queue/binding
        self.zr.declare_exchange('known')
        self.zr.declare_queue('q1')
        self.zr.bind('known', 'q1', 'a.*')

        self.zr.declare_queue('q2')
        self.zr.bind('known', 'q2', 'a.b')

        self.zr.declare_queue('q3')
        self.zr.bind('known', 'q3', '*.b')

        # send message
        self.pub.send_multipart(['known', 'a.b', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # should be in 3 queues
        for q in ['q1','q2','q3']:
            self.assertEquals(self.zr._queues[q].qsize(), 1)
            self.assertIn(('known', 'a.b', 'body', 'props'), self.zr._queues[q])

    def test_publish_to_queue_with_multiple_matching_binds_only_makes_one_message(self):
        # exchange/queue/bindings
        self.zr.declare_exchange('known')
        self.zr.declare_queue('iamqueue')
        self.zr.bind('known', 'iamqueue', 'a.*')
        self.zr.bind('known', 'iamqueue', 'a.b')
        self.zr.bind('known', 'iamqueue', '*.b')

        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 0)

        # send message
        self.pub.send_multipart(['known', 'a.b', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # should be in the queue
        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 1)
        self.assertIn(('known', 'a.b', 'body', 'props'), self.zr._queues['iamqueue'])

    def test_publish_with_binds_and_unbinds(self):
        # declare exchange/queue
        self.zr.declare_exchange('known')
        self.zr.declare_queue('ein')

        self.assertEquals(self.zr._queues['ein'].qsize(), 0)

        # send message
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # not bound, so doesn't go into the queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 0)

        # bind now
        self.zr.bind('known', 'ein', 'ein')

        # send message
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # should be in queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 1)

        # send again
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # now 2 in queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 2)

        # unbind
        self.zr.unbind('known', 'ein', 'ein')

        # send again
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # still 2 in queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 2)

<commit_msg>Updated transport unit tests<commit_after>#!/usr/bin/env python

__author__ = 'Dave Foster <dfoster@asascience.com>'
__license__ = 'Apache 2.0'

from pyon.util.unit_test import PyonTestCase
from pyon.util.int_test import IonIntegrationTestCase
from pyon.net.transport import NameTrio, BaseTransport, AMQPTransport, TransportError, TopicTrie, ZeroMQRouter
from pyon.core.bootstrap import get_sys_name

from nose.plugins.attrib import attr
from mock import Mock, MagicMock, sentinel
from gevent_zeromq import zmq
from gevent.event import Event
import time

@attr('UNIT')
class TestTransport(PyonTestCase):
    pass

@attr('UNIT')
class TestAMQPTransport(PyonTestCase):

    def setUp(self):
        pass

    def test__sync_call_no_ret_value(self):

        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam()

        tp = AMQPTransport(Mock())
        rv = tp._sync_call(async_func, 'callback')
        self.assertIsNone(rv)

    def test__sync_call_with_ret_value(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sentinel.val)

        tp = AMQPTransport(Mock())
        rv = tp._sync_call(async_func, 'callback')
        self.assertEquals(rv, sentinel.val)

    def test__sync_call_with_mult_rets(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sentinel.val, sentinel.val2)

        tp = AMQPTransport(Mock())
        rv = tp._sync_call(async_func, 'callback')
        self.assertEquals(rv, (sentinel.val, sentinel.val2))

    def test__sync_call_with_kwarg_rets(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sup=sentinel.val, sup2=sentinel.val2)

        tp = AMQPTransport(Mock())
        rv = tp._sync_call(async_func, 'callback')
        self.assertEquals(rv, {'sup':sentinel.val, 'sup2':sentinel.val2})

    def test__sync_call_with_normal_and_kwarg_rets(self):
        def async_func(*args, **kwargs):
            cbparam = kwargs.get('callback')
            cbparam(sentinel.arg, sup=sentinel.val, sup2=sentinel.val2)

        tp = AMQPTransport(Mock())
        rv = tp._sync_call(async_func, 'callback')
        self.assertEquals(rv, (sentinel.arg, {'sup':sentinel.val, 'sup2':sentinel.val2}))

    def test__sync_call_with_error(self):
        tp = AMQPTransport(Mock())

        def async_func(*args, **kwargs):
            raise TransportError('haha')

        self.assertRaises(TransportError, tp._sync_call, async_func, 'callback')

@attr('UNIT')
class TestTopicTrie(PyonTestCase):
    def setUp(self):
        self.tt = TopicTrie()

    def test_many_matches(self):
        self.tt.add_topic_tree('a.b.c', sentinel.p1)
        self.tt.add_topic_tree('a.*.c', sentinel.p2)
        self.tt.add_topic_tree('*.b.c', sentinel.p3)
        self.tt.add_topic_tree('a.b.*', sentinel.p4)
        self.tt.add_topic_tree('*.*.c', sentinel.p5)
        self.tt.add_topic_tree('a.#',   sentinel.wild)
        self.tt.add_topic_tree('a.#.c', sentinel.middle_wild)

        self.assertEquals({sentinel.p1, sentinel.p2, sentinel.p3, sentinel.p4, sentinel.p5, sentinel.wild, sentinel.middle_wild},
                          set(self.tt.get_all_matches('a.b.c')))

        self.assertEquals({sentinel.p2, sentinel.p5, sentinel.wild, sentinel.middle_wild},
                          set(self.tt.get_all_matches('a.d.c')))

        self.assertEquals({sentinel.wild, sentinel.middle_wild},
                          set(self.tt.get_all_matches('a.b.b.b.b.b.b.c')))

        self.assertEquals({sentinel.wild},
                          set(self.tt.get_all_matches('a.b.b.b.b.b.b')))

@attr('UNIT')
class TestZeroMQRouter(PyonTestCase):
    pass

@attr('UNIT')
class TestZeroMQRouter(PyonTestCase):

    def setUp(self):
        self.zc = zmq.Context(1)
        self.zr = ZeroMQRouter(self.zc, get_sys_name())
        self.zr.start()
        self.addCleanup(self.zr.stop)

        self.pub = self.zc.socket(zmq.PUB)
        self.pub.connect(self.zr._connect_addr)

        # make a hook so we can tell when a message gets routed
        self.ev = Event()
        def new_route(*args, **kwargs):
            self.ev.set()
            self.zr._oldroute(*args, **kwargs)
        self.zr._oldroute = self.zr._route
        self.zr._route = new_route

    def test_publish_to_unknown_exchange(self):
        self.assertEquals(len(self.zr.errors), 0)
        self.pub.send_multipart(['ex', 'rkey', 'body', 'props'])
        self.ev.wait(timeout=10)
        self.assertEquals(len(self.zr.errors), 1)

    def test_publish_to_known_exchange(self):
        # declare exchange
        self.zr.declare_exchange('known')

        # send message
        self.pub.send_multipart(['known', 'rkey', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # message is binned but no errors
        self.assertEquals(len(self.zr.errors), 0)

    def test_publish_to_queue(self):
        # declare exchange/queue/binding
        self.zr.declare_exchange('known')
        self.zr.declare_queue('iamqueue')
        self.zr.bind('known', 'iamqueue', 'binzim')
        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 0)

        # send message
        self.pub.send_multipart(['known', 'binzim', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # should be sitting in a queue waiting
        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 1)
        #self.assertIn(('known', 'binzim', 'body', 'props'), self.zr._queues['iamqueue'])

    def test_publish_to_many_queues(self):
        # declare exchange/queue/binding
        self.zr.declare_exchange('known')
        self.zr.declare_queue('q1')
        self.zr.bind('known', 'q1', 'a.*')

        self.zr.declare_queue('q2')
        self.zr.bind('known', 'q2', 'a.b')

        self.zr.declare_queue('q3')
        self.zr.bind('known', 'q3', '*.b')

        # send message
        self.pub.send_multipart(['known', 'a.b', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # should be in 3 queues
        for q in ['q1','q2','q3']:
            self.assertEquals(self.zr._queues[q].qsize(), 1)
            #self.assertIn(('known', 'a.b', 'body', 'props'), self.zr._queues[q])

    def test_publish_to_queue_with_multiple_matching_binds_only_makes_one_message(self):
        # exchange/queue/bindings
        self.zr.declare_exchange('known')
        self.zr.declare_queue('iamqueue')
        self.zr.bind('known', 'iamqueue', 'a.*')
        self.zr.bind('known', 'iamqueue', 'a.b')
        self.zr.bind('known', 'iamqueue', '*.b')

        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 0)

        # send message
        self.pub.send_multipart(['known', 'a.b', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)

        # should be in the queue
        self.assertEquals(self.zr._queues['iamqueue'].qsize(), 1)
        #self.assertIn(('known', 'a.b', 'body', 'props'), self.zr._queues['iamqueue'])

    def test_publish_with_binds_and_unbinds(self):
        # declare exchange/queue
        self.zr.declare_exchange('known')
        self.zr.declare_queue('ein')

        self.assertEquals(self.zr._queues['ein'].qsize(), 0)

        # send message
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # not bound, so doesn't go into the queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 0)

        # bind now
        self.zr.bind('known', 'ein', 'ein')

        # send message
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # should be in queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 1)

        # send again
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # now 2 in queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 2)

        # unbind
        self.zr.unbind('known', 'ein', 'ein')

        # send again
        self.pub.send_multipart(['known', 'ein', 'body', 'props'])

        # wait for route
        self.ev.wait(timeout=10)
        self.ev.clear() # we need to use again

        # still 2 in queue
        self.assertEquals(self.zr._queues['ein'].qsize(), 2)

<|endoftext|>


<commit_before>import torch
import itertools
import functools
from torch.autograd import Variable
from pyro.distributions.distribution import Distribution

class Bernoulli(Distribution):
    """
    :param ps: probabilities *(vector [0, 1])*

    Distribution over a vector of independent Bernoulli variables. Each element
    of the vector takes on a value in ``{0, 1}``.
    """

    def _sanitize_input(self, ps):
        if ps is not None:
            # stateless distribution
            return ps
        elif self.ps is not None:
            # stateful distribution
            return self.ps
        else:
            raise ValueError("Parameter(s) were None")

    def __init__(self, ps=None, batch_size=1, *args, **kwargs):
        """
        Params:
          ps = tensor of probabilities
        """
        self.ps = ps
        if ps is not None:
            if ps.dim() == 1:
                self.ps = ps.expand(batch_size, ps.size(0))
        super(Bernoulli, self).__init__(*args, **kwargs)

    def sample(self, ps=None, *args, **kwargs):
        """
        Reparameterized Bernoulli sampler.
        """
        _ps = self._sanitize_input(ps)
        return torch.bernoulli(_ps).type_as(_ps)

    def log_pdf(self, x, ps=None, batch_size=1, *args, **kwargs):
        """
        Bernoulli log-likelihood
        """
        _ps = self._sanitize_input(ps)
        x_1 = x - 1
        ps_1 = _ps - 1
        xmul = torch.mul(x, _ps)
        xmul_1 = torch.mul(x_1, ps_1)
        logsum = torch.log(torch.add(xmul, xmul_1))
        return torch.sum(logsum)

    def batch_log_pdf(self, x, ps=None, batch_size=1, *args, **kwargs):
        _ps = self._sanitize_input(ps)
        if x.dim() == 1 and _ps.dim() == 1 and batch_size == 1:
            return self.log_pdf(x, _ps)
        elif x.dim() == 1:
            x = x.expand(batch_size, x.size(0))
        if _ps.size() != x.size():
            _ps = _ps.expand_as(x)
        x_1 = x - 1
        ps_1 = _ps - 1
        xmul = torch.mul(x, _ps)
        xmul_1 = torch.mul(x_1, ps_1)
        logsum = torch.log(torch.add(xmul, xmul_1))
        return torch.sum(logsum, 1)

    def support(self, ps=None, *args, **kwargs):
        _ps = self._sanitize_input(ps)
        if _ps.dim() == 1:
            return iter([Variable(torch.ones(1).type_as(_ps.data)), Variable(torch.zeros(1).type_as(_ps))])
        size = functools.reduce(lambda x, y: x * y, _ps.size())
        return (Variable(torch.Tensor(list(x)).view_as(_ps))
                for x in itertools.product(torch.Tensor([0, 1]).type_as(_ps.data), repeat=size))
<commit_msg>Block AD through Bernoulli sampler.<commit_after>import torch
import itertools
import functools
from torch.autograd import Variable
from pyro.distributions.distribution import Distribution

class Bernoulli(Distribution):
    """
    :param ps: probabilities *(vector [0, 1])*

    Distribution over a vector of independent Bernoulli variables. Each element
    of the vector takes on a value in ``{0, 1}``.
    """

    def _sanitize_input(self, ps):
        if ps is not None:
            # stateless distribution
            return ps
        elif self.ps is not None:
            # stateful distribution
            return self.ps
        else:
            raise ValueError("Parameter(s) were None")

    def __init__(self, ps=None, batch_size=1, *args, **kwargs):
        """
        Params:
          ps = tensor of probabilities
        """
        self.ps = ps
        if ps is not None:
            if ps.dim() == 1:
                self.ps = ps.expand(batch_size, ps.size(0))
        super(Bernoulli, self).__init__(*args, **kwargs)

    def sample(self, ps=None, *args, **kwargs):
        """
        Bernoulli sampler.
        """
        _ps = self._sanitize_input(ps)
        return Variable(torch.bernoulli(_ps.data).type_as(_ps.data))

    def log_pdf(self, x, ps=None, batch_size=1, *args, **kwargs):
        """
        Bernoulli log-likelihood
        """
        _ps = self._sanitize_input(ps)
        x_1 = x - 1
        ps_1 = _ps - 1
        xmul = torch.mul(x, _ps)
        xmul_1 = torch.mul(x_1, ps_1)
        logsum = torch.log(torch.add(xmul, xmul_1))
        return torch.sum(logsum)

    def batch_log_pdf(self, x, ps=None, batch_size=1, *args, **kwargs):
        _ps = self._sanitize_input(ps)
        if x.dim() == 1 and _ps.dim() == 1 and batch_size == 1:
            return self.log_pdf(x, _ps)
        elif x.dim() == 1:
            x = x.expand(batch_size, x.size(0))
        if _ps.size() != x.size():
            _ps = _ps.expand_as(x)
        x_1 = x - 1
        ps_1 = _ps - 1
        xmul = torch.mul(x, _ps)
        xmul_1 = torch.mul(x_1, ps_1)
        logsum = torch.log(torch.add(xmul, xmul_1))
        return torch.sum(logsum, 1)

    def support(self, ps=None, *args, **kwargs):
        _ps = self._sanitize_input(ps)
        if _ps.dim() == 1:
            return iter([Variable(torch.ones(1).type_as(_ps.data)), Variable(torch.zeros(1).type_as(_ps))])
        size = functools.reduce(lambda x, y: x * y, _ps.size())
        return (Variable(torch.Tensor(list(x)).view_as(_ps))
                for x in itertools.product(torch.Tensor([0, 1]).type_as(_ps.data), repeat=size))
<|endoftext|>