# AUTOGENERATED! DO NOT EDIT! File to edit: 01_Pairwise_Distance.ipynb (unless otherwise specified).

__all__ = ['USE_64', 'gpu_dist_matrix', 'component_mixture_dist_matrix']

# Cell
import numpy as np
import numba as nb
from numba import cuda
from scipy.spatial.distance import pdist

# Cell

USE_64 = True

if USE_64:
    bits = 64
    np_type = np.float64
else:
    bits = 32
    np_type = np.float32

# Cell
@cuda.jit("void(float{}[:, :], float{}[:, :])".format(bits, bits))
def _euclidian_distance_matrix(mat, out):
    "CUDA kernel used to calculate the squared euclidian distance between all rows in a matrix"
    m = mat.shape[0]
    n = mat.shape[1]
    i, j = cuda.grid(2)
    d = 0
    if i < m and j > i and j < m:
        # calculate ||x - y||^2
        for k in range(n):
            tmp = mat[i, k] - mat[j, k]
            d += tmp * tmp
        out[i, j] = d

# Cell
def gpu_dist_matrix(mat):
    "calculate the squared euclidian distance between all pairs of rows in a given matrix"
    rows = mat.shape[0]

    block_dim = (16, 16)
    grid_dim = (int(rows/block_dim[0] + 1), int(rows/block_dim[1] + 1))

    stream = cuda.stream()
    mat2 = cuda.to_device(np.asarray(mat, dtype=np_type), stream=stream)
    out2 = cuda.device_array((rows, rows))
    _euclidian_distance_matrix[grid_dim, block_dim](mat2, out2)
    out = out2.copy_to_host(stream=stream)
    return out

# Cell

@cuda.jit("void(float{}[:, :], float{}[:, :], float{}[:, :])".format(bits, bits, bits))
def _pairwise_distance_matrix(compMat, mixMat, out):
    "CUDA kernel used to calcualte squared euclidian distance  between pairs of rows in two matrices"
    nC = compMat.shape[0]
    nM = mixMat.shape[0]
    dim = compMat.shape[1]
    i, j = cuda.grid(2)
    d = 0
    if i < nC and j < nM:
        # calculate ||c_i - m_j||^2
        for k in range(dim):
            tmp = compMat[i, k] - mixMat[j, k]
            d += tmp * tmp
        out[i, j] = d

# Cell

def component_mixture_dist_matrix(compMat, mixMat):
    "calculate the squared euclidian distance between pairs of rows in the two given matrices"
    compRows = compMat.shape[0]
    mixRows = mixMat.shape[0]

    block_dim = (16, 16)
    grid_dim = (int(compRows/block_dim[0] + 1), int(mixRows/block_dim[1] + 1))

    stream = cuda.stream()
    compMat2 = cuda.to_device(np.asarray(compMat, dtype=np_type), stream=stream)
    mixMat2 = cuda.to_device(np.asarray(mixMat, dtype=np_type), stream=stream)
    out2 = cuda.device_array((compRows, mixRows))
    _pairwise_distance_matrix[grid_dim, block_dim](compMat2, mixMat2, out2)
    out = out2.copy_to_host(stream=stream)
    return out
# AUTOGENERATED! DO NOT EDIT! File to edit: 01_Pairwise_Distance.ipynb (unless otherwise specified).

__all__ = ['USE_64', 'gpu_dist_matrix', 'component_mixture_dist_matrix']

# Cell
import numpy as np
import numba as nb
from numba import cuda
from scipy.spatial.distance import pdist

# Cell

USE_64 = True

if USE_64:
    bits = 64
    np_type = np.float64
else:
    bits = 32
    np_type = np.float32

# Cell
@cuda.jit("void(float{}[:, :], float{}[:, :])".format(bits, bits))
def _euclidian_distance_matrix(mat, out):
    "CUDA kernel used to calculate the squared euclidian distance between all rows in a matrix"
    m = mat.shape[0]
    n = mat.shape[1]
    i, j = cuda.grid(2)
    d = 0
    if i < m and j > i and j < m:
        # calculate ||x - y||^2
        for k in range(n):
            tmp = mat[i, k] - mat[j, k]
            d += tmp * tmp
        out[i, j] = d

# Cell
def gpu_dist_matrix(mat):
    "calculate the squared euclidian distance between all pairs of rows in a given matrix"
    rows = mat.shape[0]

    block_dim = (16, 16)
    grid_dim = (int(rows/block_dim[0] + 1), int(rows/block_dim[1] + 1))

    stream = cuda.stream()
    mat2 = cuda.to_device(np.asarray(mat, dtype=np_type), stream=stream)
    out2 = cuda.device_array((rows, rows))
    _euclidian_distance_matrix[grid_dim, block_dim](mat2, out2)
    out = out2.copy_to_host(stream=stream)
    return out

# Cell

@cuda.jit("void(float{}[:, :], float{}[:, :], float{}[:, :])".format(bits, bits, bits))
def _pairwise_distance_matrix(compMat, mixMat, out):
    "CUDA kernel used to calcualte squared euclidian distance  between pairs of rows in two matrices"
    nC = compMat.shape[0]
    nM = mixMat.shape[0]
    dim = compMat.shape[1]
    i, j = cuda.grid(2)
    d = 0
    if i < nC and j < nM:
        # calculate ||c_i - m_j||^2
        for k in range(dim):
            tmp = compMat[i, k] - mixMat[j, k]
            d += tmp * tmp
        out[i, j] = d

# Cell

def component_mixture_dist_matrix(compMat, mixMat):
    "calculate the squared euclidian distance between pairs of rows in the two given matrices"
    compRows = compMat.shape[0]
    mixRows = mixMat.shape[0]

    block_dim = (16, 16)
    grid_dim = (int(compRows/block_dim[0] + 1), int(mixRows/block_dim[1] + 1))

    stream = cuda.stream()
    compMat2 = cuda.to_device(np.asarray(compMat, dtype=np_type), stream=stream)
    mixMat2 = cuda.to_device(np.asarray(mixMat, dtype=np_type), stream=stream)
    out2 = cuda.device_array((compRows, mixRows))
    _pairwise_distance_matrix[grid_dim, block_dim](compMat2, mixMat2, out2)
    out = out2.copy_to_host(stream=stream)
    return out
# AUTOGENERATED! DO NOT EDIT! File to edit: 01_Pairwise_Distance.ipynb (unless otherwise specified). # Cell # Cell # Cell # calculate ||x - y||^2 # Cell # Cell # calculate ||c_i - m_j||^2 # Cell