<jupyter_start><jupyter_text>RBF Calculation Speed-Up<jupyter_code>import numpy as np
from sklearn.model_selection import train_test_split
from krr import KRR, rbf_derivative
from sklearn.kernel_ridge import KernelRidge
from sklearn.model_selection import GridSearchCV
from sklearn.metrics.pairwise import rbf_kernel

import cProfile
from rbf_derivative_cy import rbf_derivative as rbf_derivative_cy<jupyter_output><empty_output><jupyter_text>Get Data<jupyter_code>random_state = 123
num_points = 1000
x_data = np.arange(0, num_points)
y_data = np.sin(x_data)

# split into training and testing
train_prnt = 0.7
x_train, x_test, y_train, y_test = \
    train_test_split(x_data, y_data,
                     train_size=train_prnt,
                     random_state=random_state)
# make a new axis D [N x D]
x_train, x_test = x_train[:, np.newaxis], x_test[:, np.newaxis]
y_train, y_test = y_train[:, np.newaxis], y_test[:, np.newaxis]

# remove the mean from y training
y_train = y_train - np.mean(y_train)

param_grid = {"alpha": [1e0, 1e-1, 1e-2, 1e-3],
              "gamma": np.logspace(-2, 2, 5)}

# initialize the kernel ridge regression model
KRR_model = GridSearchCV(KernelRidge(kernel='rbf'), 
                         cv=5, param_grid=param_grid)

# fit model to data
KRR_model.fit(x_train, y_train)

# predict using the krr model
y_pred = KRR_model.predict(x_test)<jupyter_output><empty_output><jupyter_text>RBF Derivative<jupyter_code>def rbf_derivative(x_train, x_function, weights, kernel_mat,
                   n_derivative=1, gamma=1.0):
    """This function calculates the rbf derivative
    Parameters
    ----------
    x_train : array, [N x D]
        The training data used to find the kernel model.

    x_function  : array, [M x D]
        The test points (or vector) to use.

    weights   : array, [N x D]
        The weights found from the kernel model
            y = K * weights

    kernel_mat: array, [N x M], default: None
        The rbf kernel matrix with the similarities between the test
        points and the training points.

    n_derivative : int, (default = 1) {1, 2}
        chooses which nth derivative to calculate

    gamma : float, default: None
        the parameter for the rbf_kernel matrix function

    Returns
    -------

    derivative : array, [M x D]
        returns the derivative with respect to training points used in
        the kernel model and the test points.

    Information
    -----------
    Author: <NAME>
    Email : <EMAIL>
            <EMAIL>
    """

    # initialize rbf kernel
    derivative = np.zeros(np.shape(x_function))

    # consolidate the parameters
    theta = 2 * gamma

    # 1st derivative
    if n_derivative == 1:

        # loop through dimensions
        for dim in np.arange(0, np.shape(x_function)[1]):

            # loop through the number of test points
            for iTest in np.arange(0, np.shape(x_function)[0]):

                # loop through the number of test points
                for iTrain in np.arange(0, np.shape(x_train)[0]):

                    # calculate the derivative for the test points
                    derivative[iTest, dim] += theta * weights[iTrain] * \
                                              (x_train[iTrain, dim] -
                                               x_function[iTest, dim]) * \
                                              kernel_mat[iTrain, iTest]

    # 2nd derivative
    elif n_derivative == 2:

        # loop through dimensions
        for dim in np.arange(0, np.shape(x_function)[1]):

            # loop through the number of test points
            for iTest in np.arange(0, np.shape(x_function)[0]):

                # loop through the number of test points
                for iTrain in np.arange(0, np.shape(x_train)[0]):
                    derivative[iTest, dim] += weights[iTrain] * \
                                              (theta ** 2 *
                                               (x_train[iTrain, dim] - x_function[iTest, dim]) ** 2
                                               - theta) * \
                                              kernel_mat[iTrain, iTest]

    return derivative

# extract necessary parameters
gamma = KRR_model.best_params_['gamma']
weights = KRR_model.best_estimator_.dual_coef_
kernel = rbf_kernel(x_train, gamma=gamma)
n_derivative = 1
derivative = rbf_derivative(x_train, x_test, weights=weights, kernel_mat=kernel, n_derivative=n_derivative, gamma=gamma)
derivative_cy = rbf_derivative_cy(np.float64(x_train), 
                                  np.float64(x_test), 
                                  weights=weights.squeeze(), 
                                  kernel_mat=kernel, 
                                  n_derivative=n_derivative, 
                                  gamma=gamma)
print(np.array_equal(derivative, derivative_cy))
print(derivative.dtype)
print(derivative_cy.dtype)<jupyter_output>True
float64
float64<jupyter_text>Timing the Current Function<jupyter_code>original_python = %timeit -o rbf_derivative(x_train, x_test, weights=weights, kernel_mat=kernel, n_derivative=n_derivative, gamma=gamma)
#cProfile.run('rbf_derivative(x_train, x_test, weights=weights, kernel_mat=kernel, n_derivative=n_derivative, gamma=gamma)')
%prun rbf_derivative(x_train, x_test, weights=weights, kernel_mat=kernel, n_derivative=n_derivative, gamma=gamma)
%load_ext line_profiler
%lprun -f rbf_derivative rbf_derivative(x_train, x_test, weights=weights, kernel_mat=kernel, gamma=gamma)
from numba import double
from numba import jit, autojit, njit
rbf_derivative_numba = jit()(rbf_derivative)
# compile it the first time
derivative_numba = rbf_derivative_numba(x_train, x_test, kernel_mat=kernel, weights=weights, gamma=gamma)
print('Numba w/ Jit:')
jit_speedup = %timeit -o rbf_derivative_numba(x_train, x_test, kernel_mat=kernel, weights=weights, gamma=gamma)
print(original_python.best / jit_speedup.best)
%load_ext cython<jupyter_output><empty_output><jupyter_text>With GIL Lock<jupyter_code>%%cython -a

cimport cython
import numpy as np
cimport numpy as np

@cython.boundscheck(False)
@cython.wraparound(False)
def rbf_derivative_cython(np.float64_t[:, :] x_train, 
                   np.float64_t[:, :] x_function,
                   np.float64_t[:] weights,
                   np.float64_t[:, :] kernel_mat,
                   np.int_t n_derivative,
                   np.float64_t gamma):
    """This function calculates the rbf derivative using
    Cython. It has been fairly optimized and provides x100
    speedup over the original python function.
    
    Parameters
    ----------
    x_train : array, [N x D], int64
        The training data used to find the kernel model.

    x_function  : array, [M x D], int64
        The test points (or vector) to use.

    weights   : array, [N x D], float64
        The weights found from the kernel model
            y = K * weights

    kernel_mat: array, [N x M], float64
        The rbf kernel matrix with the similarities between the test
        points and the training points.

    n_derivative : int, (default = 1) {1, 2}, int
        chooses which nth derivative to calculate

    gamma : float, default: None, float64
        the parameter for the rbf_kernel matrix function

    Returns
    -------

    derivative : array, [M x D]
        returns the derivative with respect to training points used in
        the kernel model and the test points.

    Information
    -----------
    Author: <NAME>
    Email : <EMAIL>
            <EMAIL>
    """
    cdef int d_dimensions = x_function.shape[1]
    cdef int n_test = x_function.shape[0]
    cdef int n_train = x_train.shape[0]
    cdef int idim, iTest, iTrain
    
    # initialize the derivative
    cdef np.float64_t[:,:] derivative = np.zeros((n_test, d_dimensions))

    # consolidate the parameters
    cdef np.float64_t theta = 2.0 * gamma


    if n_derivative == 1:
        
        # loop through dimensions
        for idim in range(d_dimensions):

            # loop through the number of test points
            for iTest in range(n_test):

                # loop through the number of test points
                for iTrain in range(n_train):

                    # calculate the derivative for the test points
                    derivative[iTest, idim] += theta * weights[iTrain] * \
                                              (x_train[iTrain, idim] -
                                               x_function[iTest, idim]) * \
                                              kernel_mat[iTrain, iTest]
                        
    # 2nd derivative
    elif n_derivative == 2:

        # loop through dimensions
        for dim in range(d_dimensions):

            # loop through the number of test points
            for iTest in range(n_test):

                # loop through the number of test points
                for iTrain in range(n_train):
                    derivative[iTest, dim] += weights[iTrain] * \
                                              (theta ** 2 *
                                               (x_train[iTrain, dim] - x_function[iTest, dim]) ** 2
                                               - theta) * \
                                              kernel_mat[iTrain, iTest] 
    else:
        raise ValueError('n_derivative should be equal to 1 or 2.')
                            
    return np.asarray(derivative)

derivative_cython = rbf_derivative_cython(np.float64(x_train),
                                          np.float64(x_test),
                                          kernel_mat=kernel, 
                                          weights=weights.squeeze(), 
                                          n_derivative=n_derivative, 
                                          gamma=np.float64(gamma))
print(np.allclose(derivative, derivative_cython))
print(derivative.dtype)
print(derivative_cy.dtype)
print('Cython:')
cython_speedup = %timeit -o rbf_derivative_cython(np.float64(x_train), np.float64(x_test), kernel_mat=kernel, weights=weights.squeeze(), n_derivative=n_derivative, gamma=gamma)
print('Pure Python:')
print(original_python.best / original_python.best) 

print('Numba w/ Jit Speedup:')
print(original_python.best / jit_speedup.best)

print('Cython speedup:')
print(original_python.best / cython_speedup.best)<jupyter_output>Pure Python:
1.0
Numba w/ Jit Speedup:
1.0067572125699529
Cython speedup:
3772.616076209928<jupyter_text>Without GIL Lock<jupyter_code>%%cython

cimport cython
import numpy as np
cimport numpy as np
from cython.parallel import prange

@cython.boundscheck(False)
@cython.wraparound(False)
def rbf_derivative_cython(np.float64_t[:, :] x_train, 
                   np.float64_t[:, :] x_function,
                   np.float64_t[:] weights,
                   np.float64_t[:, :] kernel_mat,
                   np.int_t n_derivative,
                   np.float64_t gamma):
    """This function calculates the rbf derivative using
    Cython. It has been fairly optimized and provides x100
    speedup over the original python function.
    
    Parameters
    ----------
    x_train : array, [N x D], int64
        The training data used to find the kernel model.

    x_function  : array, [M x D], int64
        The test points (or vector) to use.

    weights   : array, [N x D], float64
        The weights found from the kernel model
            y = K * weights

    kernel_mat: array, [N x M], float64
        The rbf kernel matrix with the similarities between the test
        points and the training points.

    n_derivative : int, (default = 1) {1, 2}, int
        chooses which nth derivative to calculate

    gamma : float, default: None, float64
        the parameter for the rbf_kernel matrix function

    Returns
    -------

    derivative : array, [M x D]
        returns the derivative with respect to training points used in
        the kernel model and the test points.

    Information
    -----------
    Author: <NAME>
    Email : <EMAIL>
            <EMAIL>
    """
    cdef int d_dimensions = x_function.shape[1]
    cdef int n_test = x_function.shape[0]
    cdef int n_train = x_train.shape[0]
    cdef int idim, iTest, iTrain
    
    # initialize the derivative
    cdef np.float64_t[:,:] derivative = np.zeros((n_test, d_dimensions))

    # consolidate the parameters
    cdef np.float64_t theta = 2.0 * gamma


    if n_derivative == 1:
        
        # loop through dimensions
        for idim in prange(d_dimensions, nogil=True):

            # loop through the number of test points
            for iTest in prange(n_test, nogil=True):

                # loop through the number of test points
                for iTrain in prange(n_train, nogil=True):

                    # calculate the derivative for the test points
                    derivative[iTest, idim] += theta * weights[iTrain] * \
                                              (x_train[iTrain, idim] -
                                               x_function[iTest, idim]) * \
                                              kernel_mat[iTrain, iTest]
                        
    # 2nd derivative
    elif n_derivative == 2:

        # loop through dimensions
        for dim in prange(d_dimensions, nogil=True):

            # loop through the number of test points
            for iTest in prange(n_test, nogil=True):

                # loop through the number of test points
                for iTrain in prange(n_train, nogil=True):
                    derivative[iTest, dim] += weights[iTrain] * \
                                              (theta ** 2 *
                                               (x_train[iTrain, dim] - x_function[iTest, dim]) ** 2
                                               - theta) * \
                                              kernel_mat[iTrain, iTest] 
    else:
        raise ValueError('n_derivative should be equal to 1 or 2.')
                            
    return np.asarray(derivative)<jupyter_output>Error compiling Cython file:
------------------------------------------------------------
...
        
        # loop through dimensions
        for idim in prange(d_dimensions, nogil=True):

            # loop through the number of test points
            for iTest in prange(n_test, nogil=True):
                              ^
------------------------------------------------------------

/Users/eman/.ipython/cython/_cython_magic_029e0f0e560b830aefa06c993b2d5ad8.pyx:72:31: Trying to release the GIL while it was previously released.

Error compiling Cython file:
------------------------------------------------------------
...

            # loop through the number of test points
            for iTest in prange(n_test, nogil=True):

                # loop through the number of test points
                for iTrain in prange(n_train, nogil=True):
                                   ^
------------------------------------------------------------

/Users/eman/.ipython/cython/_cython_magic_0[...]<jupyter_text>Making Numba Faster<jupyter_code>@jit
def rbf_derivative_numba(x_train, x_function, weights, kernel_mat,
                   n_derivative=1, gamma=1.0):
    """This function calculates the rbf derivative
    Parameters
    ----------
    x_train : array, [N x D]
        The training data used to find the kernel model.

    x_function  : array, [M x D]
        The test points (or vector) to use.

    weights   : array, [N x D]
        The weights found from the kernel model
            y = K * weights

    kernel_mat: array, [N x M], default: None
        The rbf kernel matrix with the similarities between the test
        points and the training points.

    n_derivative : int, (default = 1) {1, 2}
        chooses which nth derivative to calculate

    gamma : float, default: None
        the parameter for the rbf_kernel matrix function

    Returns
    -------

    derivative : array, [M x D]
        returns the derivative with respect to training points used in
        the kernel model and the test points.

    Information
    -----------
    Author: <NAME>
    Email : <EMAIL>
            <EMAIL>
    """

    # initialize the derivative
    d_dimensions = x_function.shape[1]
    n_test = x_function.shape[0]
    n_train = x_train.shape[0]
    
    derivative = np.zeros((n_test, d_dimensions))

    # consolidate the parameters
    theta = 2.0 * gamma


    if n_derivative == 1:
        
        # loop through dimensions
        for idim in np.arange(0, d_dimensions):

            # loop through the number of test points
            for iTest in np.arange(0, n_test):

                # loop through the number of test points
                for iTrain in np.arange(0, n_train):

                    # calculate the derivative for the test points
                    derivative[iTest, idim] += theta * weights[iTrain] * \
                                              (x_train[iTrain, idim] -
                                               x_function[iTest, idim]) * \
                                              kernel_mat[iTrain, iTest]
                        
    # 2nd derivative
    elif n_derivative == 2:

        # loop through dimensions
        for dim in np.arange(0, d_dimensions):

            # loop through the number of test points
            for iTest in np.arange(0, n_test):

                # loop through the number of test points
                for iTrain in np.arange(0, n_train):
                    derivative[iTest, dim] += weights[iTrain] * \
                                              (theta ** 2 *
                                               (x_train[iTrain, dim] - x_function[iTest, dim]) ** 2
                                               - theta) * \
                                              kernel_mat[iTrain, iTest]

    return derivative
temp_d = rbf_derivative_numba(np.int64(x_train), 
                              np.int64(x_test), 
                              kernel_mat=np.float64(kernel), 
                              weights=np.float64(weights),
                              n_derivative=n_derivative,
                              gamma=gamma)
%%timeit
rbf_derivative_numba(np.int64(x_train), 
                              np.int64(x_test), 
                              kernel_mat=np.float64(kernel), 
                              weights=np.float64(weights),
                              n_derivative=n_derivative,
                              gamma=gamma)
rbf_derivative_numba.inspect_types()
print('Numba w/ Jit:')
jitted = %timeit -o rbf_derivative_numba(x_train, x_test, kernel_mat=kernel, weights=weights, gamma=gamma)

fast_numba = jit()(rbf_derivative)

print('Numba w/ Jit:')
%timeit fast_numba(x_train, x_test, kernel_mat=kernel, weights=weights, gamma=gamma)
rbf_derivative_numba = jit(void(double[:], double[:,:], double[:,:], double[:,:], double[:]))(rbf_derivative_numba)
%timeit rbf_derivative_numba(x_train, x_test, kernel_mat=kernel, weights=weights, n_derivative=1, gamma=gamma)
print(x_train.shape, x_train.dtype)
print(x_test.shape, x_test.dtype)
print(kernel.shape, kernel.dtype)
print(weights.squeeze().shape, weights.squeeze().dtype)
print(gamma.shape, gamma.dtype, gamma.astype(np.float).dtype)<jupyter_output>(700, 1) int64
(300, 1) int64
(700, 700) float64
(700,) float64
() float64 float64
{'fraction_non_alphanumeric': 0.09808897405355016, 'fraction_numerical': 0.014345641928634316, 'mean_word_length': 1.8259915117810626, 'pattern_counts': {'":': 2, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 72, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '569449', 'license': 'MIT', 'n_tokens_mistral': 5874, 'n_tokens_neox': 5451, 'n_words': 1536, 'path': 'notebooks/uncategorized/scale/krr/kernel_speedup.ipynb', 'repo_name': 'J-Garcke-SCAI/jaxkern', 'size': '19445'}