VERBOSE = False
ENABLE_SLOW_TESTS = False
    _N, D = x_na.shape
    M = x_ma.shape[0]
    N, D = x_na.shape
    M = x_ma.shape[0]
        J = np.c_[Jl, Jt, Jw]
    M,D = x_ma.shape
    D = lin_ag.shape[0]
    D = lin_ag.shape[0]
    E,D = x_ea.shape
    N = y_ng.shape[0]
    M = nr_ma.shape[0]
    N,D = x_na.shape
    A = np.zeros((N+D+1, N+D+1))
    A[:N, :N] = K_nn
    A[:N, N:N+D] = x_na
    A[:N, N+D] = 1
    A[N:N+D,:N] = x_na.T
    A[N+D,:N] = 1
    A[N:N+D, N:N+D] = coef_ratio*np.eye(D)
    B = np.empty((N+D+1, D))
    B[:N] = y_ng
    B[N:N+D] = coef_ratio*np.eye(D)
    B[N+D] = 0
    X = np.linalg.solve(A, B)
    N = _u[:,n_cnts:]
    Q = np.c_[np.ones((n,1)), x_na, K_nn]
    WQ = wt_n[:,None] * Q
    QWQ = Q.T.dot(WQ)
    H = QWQ
    A = np.r_[np.zeros((d+1,d+1)), np.c_[np.ones((n,1)), x_na]].T
    N,D = x_na.shape 
    A = QQ_nn    
    B = Q_nn.T.dot(y_ng)
    X = np.linalg.solve(A,B)
    N,D = x_na.shape
                A = nr_coef*JJ_nn + QQ_nn + bend_coef*Bend_nn
                X0 = X_ng[:,g]
                B = nr_coef*np.dot(J_zn.T, np.dot(J_zn, X0) - nrerr_z) + Q_nn.T.dot(y_ng[:,g])
    N,D = x_na.shape
    M = nr_ma.shape[0]
    E = N + 4*M
    F = E - M
    Q = N + 3*M - 4
            X0 = X_fg[:,g]
    N,_D = x_na.shape
    A = QQ_nn
    B = Q_nn.T.dot(y_ng-x_na.dot(lin_ag))
    X = np.linalg.solve(A,B)
    N,_D = x_na.shape
def nan2zero(x):
def tps_apply_kernel(distmat, dim):
def tps_kernel_matrix(x_na):
def tps_kernel_matrix2(x_na, y_ma):
def tps_eval(x_ma, lin_ag, trans_g, w_ng, x_na):
def tps_grad(x_ma, lin_ag, _trans_g, w_ng, x_na):
def tps_nr_grad(x_ma, lin_ag, _trans_g, w_ng, x_na, return_tuple = False):
def tps_nr_err(x_ma, lin_ag, trans_g, w_ng, x_na):
def tps_cost(lin_ag, trans_g, w_ng, x_na, y_ng, bend_coef, K_nn=None, return_tuple=False, wt_n = None):
def tps_nr_cost_eval(lin_ag, trans_g, w_ng, x_na, y_ng, xnr_ma, bend_coef, nr_coef, K_nn = None, return_tuple=False):
def tps_nr_cost_eval_general(lin_ag, trans_g, w_eg, x_ea, y_ng, nr_ma, bend_coef, nr_coef, K_ee = None, return_tuple=True):
def tps_fit(x_na, y_ng, bend_coef, rot_coef, wt_n=None, K_nn = None):
def solve_eqp1(H, f, A):
def tps_fit3(x_na, y_ng, bend_coef, rot_coef, wt_n):
def tps_fit2(x_na, y_ng, bend_coef, rot_coef, wt_n=None):
def tps_nr_fit(x_na, y_ng, bend_coef, nr_ma, nr_coef, method="newton"):
def tps_nr_fit_enhanced(x_na, y_ng, bend_coef, nr_ma, nr_coef):
def tps_fit_fixedrot(x_na, y_ng, bend_coef, lin_ag, K_nn = None, wt_n=None):
def tps_fit_regrot(x_na, y_ng, bend_coef, rfunc, wt_n=None, max_iter = 1, inner_max_iter=100, rgrad=None, l_init=None):
        def f(x): 
def tps_cost_regrot(lin_ag, trans_g, w_ng, x_na, y_ng, bend_coef, rfunc, K_nn = None, wt_n=None): def tps_kernel_matrix(x_na):
    dim = x_na.shape[1]
    distmat = ssd.squareform(ssd.pdist(x_na))
    return tps_apply_kernel(distmat,dim) def tps_eval(x_ma, lin_ag, trans_g, w_ng, x_na):
    K_mn = tps_kernel_matrix2(x_ma, x_na).astype(np.float32)
    return np.dot(K_mn, w_ng) + np.dot(x_ma, lin_ag) + trans_g[None,:] def dot_batch_nocheck(a_arr_gpu, b_arr_gpu, c_arr_gpu, a_ptr_gpu, b_ptr_gpu, c_ptr_gpu,
                      transa = 'N', transb = 'N', a = 1, b = 1, handle = None):
    """
    Implementation of batched dot products using cuda.    

    Parameters
    ----------
    a_arr_gpu : list of pycuda.gpuarray.GPUArray
        Input array.
    b_arr_gpu : list of pycuda.gpuarray.GPUArray
        Input array.
    c_arr_gpu : list of pycuda.gpuarray.GPUArray
        Input/Output array.
    a_ptr_gpu : pycuda.gpuarray.GPUArray
        Array of pointers to arrays in a_arr_gpu
    b_ptr_gpu : pycuda.gpuarray.GPUArray
        Array of pointers to arrays in b_arr_gpu
    c_ptr_gpu : pycuda.gpuarray.GPUArray
        Array of pointers to arrays in c_arr_gpu
    transa : char
        If 'T', compute the product of the transpose of `a_arr_gpu[i]`.
        If 'C', compute the product of the Hermitian of `a_arr_gpu[i]`.
    transb : char
        If 'T', compute the product of the transpose of `b_arr_gpu[i]`.
        If 'C', compute the product of the Hermitian of `b_arr_gpu[i]`.
    handle : int
        CUBLAS context. If no context is specified, the default handle from
        `scikits.cuda.misc._global_cublas_handle` is used.

    Returns
    -------
    None. Output is stored into the arrays in c_arr_gpu

    Notes
    -----
    The input matrices must all contain elements of the same data type.
    All matrics in a list must have same size.

    Examples
    --------
    >>> import pycuda.driver
    >>> import pycuda.autoinit
    >>> from pycuda import gpuarray
    >>> import numpy as np
    >>> from scikits.cuda import linalg
    >>> linalg.init()
    >>> a_arr = [np.asarray(np.random.rand(4, 2), np.float32) for i in range(10)]
    >>> b_arr = [np.asarray(np.random.rand(2, 2), np.float32) for i in range(10)]
    >>> c_arr = [np.asarray(np.random.rand(4, 2), np.float32) for i in range(10)]
    >>> a_arr_gpu = [gpuarray.to_gpu(a_gpu) for a_gpu in a_arr]
    >>> b_arr_gpu = [gpuarray.to_gpu(b_gpu) for b_gpu in b_arr]
    >>> c_arr_gpu = [gpuarray.to_gpu(c_gpu) for c_gpu in c_arr]
    >>> a_ptr_gpu = gpuarray.to_gpu(np.asarray([int(a_gpu.gpudata) for a_gpu in a_arr_gpu]))    
    >>> b_ptr_gpu = gpuarray.to_gpu(np.asarray([int(b_gpu.gpudata) for b_gpu in b_arr_gpu]))
    >>> c_ptr_gpu = gpuarray.to_gpu(np.asarray([int(c_gpu.gpudata) for c_gpu in c_arr_gpu]))
    >>> linalg.dot_batch_nocheck(a_arr_gpu, b_arr_gpu, c_arr_gpu, a_ptr_gpu, b_ptr_gpu, c_ptr_gpu)
    >>> for i in range(10):
    ...   print np.allclose(np.dot(a_arr[i], b_arr[i]) + c_arr[i], c_arr_gpu[i].get())
    ... 
    True
    True
    True
    True
    True
    True
    True
    True
    True
    True
    >>>
    """
    if handle is None:
        handle = misc._global_cublas_handle
    N = len(a_arr_gpu)
    a_shape = a_arr_gpu[0].shape
    a_dtype = a_arr_gpu[0].dtype
    b_shape = b_arr_gpu[0].shape
    b_dtype = b_arr_gpu[0].dtype
    c_shape = c_arr_gpu[0].shape
    c_dtype = c_arr_gpu[0].dtype
    
    for i in range(N):
        assert a_arr_gpu[i].shape == a_shape
        assert a_arr_gpu[i].dtype == a_dtype
        assert b_arr_gpu[i].shape == b_shape
        assert b_arr_gpu[i].dtype == b_dtype
        assert c_arr_gpu[i].shape == c_shape
        assert c_arr_gpu[i].dtype == c_dtype

        assert a_arr_gpu[i].flags.c_contiguous
        assert b_arr_gpu[i].flags.c_contiguous
        assert c_arr_gpu[i].flags.c_contiguous

    if len(a_shape) == 1:        
        a_shape = (1, a_shape[0])
    if len(b_shape) == 1:
        b_shape = (1, b_shape[0])
    if len(c_shape) == 1:
        c_shape = (1, c_shape[0])

    transa = lower(transa)
    transb = lower(transb)

    if transb in ['t', 'c']:
        m, k = b_shape
    elif transb in ['n']:
        k, m = b_shape
    else:
        raise ValueError('invalid value for transb')
    
    if transa in ['t', 'c']:
        l, n = a_shape
    elif transa in ['n']:
        n, l = a_shape
    else:
        raise ValueError('invalid value for transa')

    i, j = c_shape
    
    if l != k:
        raise ValueError('objects are not aligned')    
    if i != n:
        raise ValueError('objects are not aligned')
    if j != m:
        raise ValueError('objects are not aligned')
    
    if transb == 'n':
        lda = max(1, m)
    else:
        lda = max(1, k)

    if transa == 'n':
        ldb = max(1, k)
    else:
        ldb = max(1, n)

    ldc = max(1, m)
    
    if (a_dtype == np.complex64 and b_dtype == np.complex64 \
            and c_dtype == np.complex64):
        cublas_func = cublas.cublasCgemmBatched
        alpha = np.complex64(a)
        beta = np.complex64(b)
    elif (a_dtype == np.float32 and b_dtype == np.float32\
            and c_dtype == np.float32):
        cublas_func = cublas.cublasSgemmBatched
        alpha = np.float32(a)
        beta = np.float32(b)
    elif (a_dtype == np.complex128 and b_dtype == np.complex128\
            and c_dtype == np.complex128):
        cublas_func = cublas.cublasZgemmBatched
        alpha = np.complex128(a)
        beta = np.complex128(b)
    elif (a_dtype == np.float64 and b_dtype == np.float64\
            and c_dtype == np.float64):
        cublas_func = cublas.cublasDgemmBatched
        alpha = np.float64(a)
        beta = np.float64(b)
    else:
        raise ValueError('unsupported combination of input types')

    cublas_func(handle, transb, transa, m, n, k, alpha, b_ptr_gpu.gpudata, lda, 
                a_ptr_gpu.gpudata, ldb, beta, c_ptr_gpu.gpudata, ldc, N) def get_gpu_ptrs(arr, (m, n) = (0, 0)):
    shape = arr[0].shape
    if len(shape) == 1:        
        ptrs = [int(x[m:].gpudata) for x in arr]
    else:
        ptrs = [int(x[m:, n:].gpudata) for x in arr]
    return gpuarray.to_gpu(np.array(ptrs)) def downsample_cloud(cloud_xyz):
    return clouds.downsample(cloud_xyz, DS_SIZE) def batch_get_sol_params(x_nd, K_nn, bend_coefs, rot_coef):
    n, d = x_nd.shape

    x_gpu = gpuarray.to_gpu(x_nd)

    H_arr_gpu = []
    for b in bend_coefs:
        cur_offset = np.zeros((1 + d + n, 1 + d + n), np.float64)
        cur_offset[d+1:, d+1:] = b * K_nn
        cur_offset[1:d+1, 1:d+1] = np.diag(rot_coef)
        H_arr_gpu.append(gpuarray.to_gpu(cur_offset))
    H_ptr_gpu = get_gpu_ptrs(H_arr_gpu)

    A = np.r_[np.zeros((d+1,d+1)), np.c_[np.ones((n,1)), x_nd]].T
    n_cnts = A.shape[0]
    _u,_s,_vh = np.linalg.svd(A.T)
    N = _u[:,n_cnts:]
    F = np.zeros((n + d + 1, d), np.float64)
    F[1:d+1, :d] += np.diag(rot_coef)
    
    Q = np.c_[np.ones((n,1)), x_nd, K_nn].astype(np.float64)
    F = F.astype(np.float64)
    N = N.astype(np.float64)

    Q_gpu     = gpuarray.to_gpu(Q)
    Q_arr_gpu = [Q_gpu for _ in range(len(bend_coefs))]
    Q_ptr_gpu = get_gpu_ptrs(Q_arr_gpu)

    F_gpu     = gpuarray.to_gpu(F)
    F_arr_gpu = [F_gpu for _ in range(len(bend_coefs))]
    F_ptr_gpu = get_gpu_ptrs(F_arr_gpu)

    N_gpu = gpuarray.to_gpu(N)
    N_arr_gpu = [N_gpu for _ in range(len(bend_coefs))]
    N_ptr_gpu = get_gpu_ptrs(N_arr_gpu)
    
    dot_batch_nocheck(Q_arr_gpu, Q_arr_gpu, H_arr_gpu,
                      Q_ptr_gpu, Q_ptr_gpu, H_ptr_gpu,
                      transa = 'T')
    # N'HN
    NHN_arr_gpu, NHN_ptr_gpu = m_dot_batch((N_arr_gpu, N_ptr_gpu, 'T'),
                                           (H_arr_gpu, H_ptr_gpu, 'N'),
                                           (N_arr_gpu, N_ptr_gpu, 'N'))
    iH_arr = []
    for NHN in NHN_arr_gpu:
        iH_arr.append(scipy.linalg.inv(NHN.get()).copy())
    iH_arr_gpu = [gpuarray.to_gpu_async(iH) for iH in iH_arr]
    iH_ptr_gpu = get_gpu_ptrs(iH_arr_gpu)

    proj_mats   = m_dot_batch((N_arr_gpu,  N_ptr_gpu,   'N'),
                              (iH_arr_gpu, iH_ptr_gpu, 'N'),
                              (N_arr_gpu,  N_ptr_gpu,   'T'),
                              (Q_arr_gpu,  Q_ptr_gpu,   'T'))

    offset_mats = m_dot_batch((N_arr_gpu,  N_ptr_gpu,   'N'),
                              (iH_arr_gpu, iH_ptr_gpu, 'N'),
                              (N_arr_gpu,  N_ptr_gpu,   'T'),
                              (F_arr_gpu,  F_ptr_gpu,   'N'))

    return proj_mats, offset_mats def unit_boxify(x_na):    
    ranges = x_na.ptp(axis=0)
    dlarge = ranges.argmax()
    unscaled_translation = - (x_na.min(axis=0) + x_na.max(axis=0))/2
    scaling = 1./ranges[dlarge]
    scaled_translation = unscaled_translation * scaling
    return x_na*scaling + scaled_translation, (scaling, scaled_translation)     def loglinspace(a,b,n):
    "n numbers between a to b (inclusive) with constant ratio between consecutive numbers"
    return np.exp(np.linspace(np.log(a),np.log(b),n))     N_ITER_CHEAP       = 10 EM_ITER_CHEAP      = 1 DEFAULT_LAMBDA     = (.1, .001) MAX_CLD_SIZE       = 150 DATA_DIM           = 3 DS_SIZE            = 0.025 # for overhand N_STREAMS          = 10 DEFAULT_NORM_ITERS = 10 BEND_COEF_DIGITS   = 6 MAX_TRAJ_LEN       = 100
import h5py
import sys
import numpy as np
import scipy.spatial.distance as ssd
import pycuda.driver as drv
import IPython as ipy
import time
        import scipy.spatial.distance as ssd
        import sys
        import sys
    import scikits.cuda.linalg as la
    import argparse
from pycuda import gpuarray
from scikits.cuda import linalg
from lfd.tpsopt import tps
from lfd.tpsopt.tps import tps_kernel_matrix, tps_eval
from lfd.tpsopt.culinalg_exts import dot_batch_nocheck, get_gpu_ptrs
from lfd.tpsopt.precompute import downsample_cloud, batch_get_sol_params
from cuda_funcs import init_prob_nm, norm_prob_nm, get_targ_pts, check_cuda_err, fill_mat, reset_cuda, sq_diffs, \
    closest_point_cost, scale_points, gram_mat_dist
from lfd.tpsopt.registration import unit_boxify, loglinspace
from lfd.tpsopt.settings import N_ITER_CHEAP, EM_ITER_CHEAP, DEFAULT_LAMBDA, MAX_CLD_SIZE, DATA_DIM, DS_SIZE, N_STREAMS, \
    DEFAULT_NORM_ITERS, BEND_COEF_DIGITS, MAX_TRAJ_LEN
    from lfd.tpsopt.transformations import ThinPlateSpline, set_ThinPlateSpline
                     self.corr_cm_ptrs, self.corr_rm_ptrs)
        norm_prob_nm(self.corr_cm_ptrs, self.corr_rm_ptrs, 
                     self.dims_gpu, other.dims_gpu, self.N, outlierfrac, norm_iters,
                     self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs)
        get_targ_pts(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.corr_cm_ptrs, self.corr_rm_ptrs,
                     self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs,
                     self.dims_gpu, other.dims_gpu, 
                     outliercutoff, self.N,
                     self.pt_t_ptrs, other.pt_t_ptrs)
        n, m = self.dims[0], other.dims[0]
        x = self.pts[0].get()[:n]
        xw = self.pts_w[0].get()[:n]
        xt = self.pts_t[0].get()[:n]
        y = other.pts[0].get()[:m]
        yw = other.pts_w[0].get()[:m]
        yt = other.pts_t[0].get()[:m]

        init_corr = self.corr_rm[0].get()        
        init_corr = init_corr.flatten()[:(n + 1) * (m + 1)].reshape(n+1, m+1).astype(np.float32)
        gpu_c_cn_coefs = self.c_coefs_cn[0].get().flatten()[:m + 1].reshape(m+1)
        gpu_c_rn_coefs = self.c_coefs_rn[0].get().flatten()[:m + 1].reshape(m+1)
        gpu_r_coefs = self.r_coefs[0].get().flatten()[:n + 1].reshape(n+1)

        rn_corr = (init_corr * gpu_c_rn_coefs[None, :]) * gpu_r_coefs[:, None]
        cn_corr = (init_corr * gpu_c_cn_coefs[None, :]) * gpu_r_coefs[:, None]        
        rn_corr = rn_corr[:n, :m]
        cn_corr = cn_corr[:n, :m]
        
        wt_n = rn_corr.sum(axis=1)
        inlier = wt_n > outliercutoff
        xtarg = np.empty((n, DATA_DIM))
        xtarg[inlier, :] = rn_corr.dot(y)[inlier, :]
        xtarg[~inlier, :] = xw[~inlier, :]

        if not np.allclose(xtarg, xt):
            print "xt values differ"
            ipy.embed()
            sys.exit(1)

        wt_m = cn_corr.sum(axis=0)
        inliner = wt_m > outliercutoff
        ytarg = np.empty((m, DATA_DIM))
        ytarg[inlier, :] = cn_corr.T.dot(x)[inliner, :]
        ytarg[~inlier, :] = yw[~inlier, :]
        if not np.allclose(ytarg, yt):
            print "yt values differ"
            ipy.embed()
            sys.exit(1)                            

    def unit_test(self, other):
        print "running basic unit tests"
        self.test_init_corr(other)
        self.test_norm_corr(other)
        self.test_get_targ(other)
        self.test_mapping_cost(other)
        self.test_bending_cost(other)
        print "UNIT TESTS PASSED"

class SrcContext(GPUContext):
    """
    specialized class to handle source clouds
    includes support for warped trajectories as well
    """
    def __init__(self, bend_coefs=None):
        GPUContext.__init__(self, bend_coefs)
        """
        items for the trajectory and warping thereof
        """
        self.l_traj          = []
        self.l_traj_ptrs     = None
        self.l_traj_K        = []
        self.l_traj_K_ptrs   = None
        self.l_traj_w        = []
        self.l_traj_w_ptrs   = None
        self.l_traj_dims     = []
        self.l_traj_dims_gpu = None

        self.r_traj          = []
        self.r_traj_ptrs     = None
        self.r_traj_K        = []
        self.r_traj_K_ptrs   = None
        self.r_traj_w        = []
        self.r_traj_w_ptrs   = None
        self.r_traj_dims     = []
        self.r_traj_dims_gpu = None

        self.traj_costs      = None
        self.tgt_traj_ptrs   = None
        self.tgt_dim_gpu     = None
        

    def update_ptrs(self):
        self.l_traj_ptrs   = get_gpu_ptrs(self.l_traj)
        self.l_traj_K_ptrs = get_gpu_ptrs(self.l_traj_K)
        self.l_traj_w_ptrs = get_gpu_ptrs(self.l_traj_w)
        self.r_traj_ptrs   = get_gpu_ptrs(self.r_traj)
        self.r_traj_K_ptrs = get_gpu_ptrs(self.r_traj_K)
        self.r_traj_w_ptrs = get_gpu_ptrs(self.r_traj_w)

        self.l_traj_dims_gpu = gpuarray.to_gpu(np.array(self.l_traj_dims, dtype=np.int32))
        self.r_traj_dims_gpu = gpuarray.to_gpu(np.array(self.r_traj_dims, dtype=np.int32))

        self.traj_costs = gpuarray.empty(self.N, np.float32)
        self.tgt_traj_ptrs = gpuarray.empty(self.N, np.int64)
        self.tgt_dim_gpu = gpuarray.empty(self.N, np.int32)

        GPUContext.update_ptrs(self)

    def add_cld(self, name, proj_mats, offset_mats, cloud_xyz, kernel, scale_params,
                r_traj, r_traj_K, l_traj, l_traj_K, update_ptrs = False):
        """
        does the normal add, but also adds the trajectories
        """
        # don't update ptrs there, do it after this
        GPUContext.add_cld(self, name, proj_mats, offset_mats, cloud_xyz, kernel, scale_params,
                           update_ptrs=False)
        self.r_traj.append(gpu_pad(r_traj, (MAX_TRAJ_LEN, DATA_DIM)))
        self.r_traj_K.append(gpu_pad(r_traj_K, (MAX_TRAJ_LEN, MAX_CLD_SIZE)))

        self.l_traj.append(gpu_pad(l_traj, (MAX_TRAJ_LEN, DATA_DIM)))