<jupyter_start><jupyter_text>Test for OptimizationDefine usual functions:<jupyter_code>import numpy as np
import scipy.linalg as la
def data_batch(data, batch_size, seed = 123):
    n = data.shape[0]
    p = data.shape[1]
    if n % batch_size !=0:
        print('%d data dropped during batching' % (n%batch_size))
    sample_size = (n // batch_size)*batch_size
        
    #shuffle
    np.random.seed(seed)
    idx = np.arange(n)
    np.random.shuffle(idx)
    n_batch = n//batch_size
    data = data[idx]
    data = data[:sample_size].reshape(batch_size, p, n_batch)
    return(data, n_batch)

def is_pos_def(A):
    '''function to check if matrix is positive definite'''
    return np.all(np.linalg.eigvals(A) > 0)

def sghmc(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
    
    np.random.seed(seed)
    
    p = theta_0.shape[0]
    n = data.shape[0]
    
    theta_samp = np.zeros((p, epochs))
    theta_samp[:,0] = theta_0
    
    B_hat = 0.5*eps*V_hat
    
    if not is_pos_def(2*(C-B_hat)*eps):
        print("error: noise term is not positive definite")
        return
        
    sqrt_noise = np.sqrt(2*(C-B_hat)*eps)
    
    sqrtM = np.sqrt(la.inv(Minv))
    r = np.random.multivariate_normal(np.zeros(p), sqrtM).reshape(p, -1)
    
    dat_batch, nbatches = data_batch(data, batch_size)
    for i in range(epochs-1):
        
        theta = theta_samp[:,i]
        r = np.random.multivariate_normal(np.zeros(p), sqrtM).reshape(p, -1)
        
        for batch in range(nbatches):
            theta = theta + (eps*Minv@r).ravel()
            gradU_batch = gradU(theta, dat_batch[:,:, batch], n, batch_size).reshape(p, -1)
            r = r-eps*gradU_batch - eps*C@Minv@r \
                + np.random.multivariate_normal(np.zeros(p), sqrt_noise).reshape(p, -1)
            
        theta_samp[:,i+1] = theta
            
    return theta_samp[:, burns:]<jupyter_output><empty_output><jupyter_text>clean algorithm: use cholesky decomposition, and use cholesky based sampling<jupyter_code>def sghmc_clean(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
    
    
    np.random.seed(seed)
    
    p = theta_0.shape[0]
    n = data.shape[0]
    
    theta_samp = np.zeros((p, epochs))
    theta_samp[:,0] = theta_0
    
    B_hat = 0.5*eps*V_hat
    
    if not is_pos_def(2*(C-B_hat)*eps):
        print("error: noise term is not positive definite")
        return
    
    sqrt_noise = np.linalg.cholesky(2*(C-B_hat)*eps)
    
    sqrtM = np.linalg.cholesky(np.linalg.inv(Minv))
    r = sqrtM@np.random.normal(size = p).reshape(p, -1)
    
    dat_batch, nbatches = data_batch(data, batch_size)
    for i in range(epochs-1):
        
        theta = theta_samp[:,i]
        r = sqrtM@np.random.normal(size = p).reshape(p, -1)
        
        for batch in range(nbatches):
            theta = theta + (eps*Minv@r).ravel()
            gradU_batch = gradU(theta, dat_batch[:,:, batch], n, batch_size).reshape(p, -1)
            r = r-eps*gradU_batch - eps*C@Minv@r + sqrt_noise@np.random.normal(size = p).reshape(p, -1)
            
        theta_samp[:,i+1] = theta
            
    return theta_samp[:, burns:]<jupyter_output><empty_output><jupyter_text>Optimization<jupyter_code>from numba import jit, float64, int64
import time
import matplotlib.pyplot as plt
import autograd.numpy as np
import seaborn as sns
from autograd import jacobian<jupyter_output><empty_output><jupyter_text>approach 1: JIT to dat_batch and sghmc<jupyter_code>@jit
def dat_batch_numba(data, batch_size, seed = 123):
    n = data.shape[0]
    p = data.shape[1]
    if n % batch_size !=0:
        print('%d data dropped during batching' % (n%batch_size))
    sample_size = (n // batch_size)*batch_size
        
    #shuffle
    np.random.seed(seed)
    idx = np.arange(n)
    np.random.shuffle(idx)
    n_batch = n//batch_size
    data = data[idx]
    data = data[:sample_size].reshape(batch_size, p, n_batch)
    return(data, n_batch)
@jit([float64[:,:](float64[:], float64, float64[:,:], float64[:,:], float64[:], float64[:,:], int64, int64, float64[:,:], int64, int64)], cache = True)
def sghmc_numba(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
        
    np.random.seed(seed)
    
    p = theta_0.shape[0]
    n = data.shape[0]
    
    theta_samp = np.zeros((p, epochs))
    theta_samp[:,0] = theta_0
    
    B_hat = 0.5*eps*V_hat
    
    if not is_pos_def(2*(C-B_hat)*eps):
        print("error: noise term is not positive definite")
        return
    
    sqrt_noise = np.linalg.cholesky(2*(C-B_hat)*eps)
    
    sqrtM = np.linalg.cholesky(np.linalg.inv(Minv))
    r = sqrtM@np.random.normal(size = p).reshape(p, -1)
    
    dat_batch, nbatches = dat_batch_numba(data, batch_size)
    for i in range(epochs-1):
        
        theta = theta_samp[:,i]
        r = sqrtM@np.random.normal(size = p).reshape(p, -1)
        
        for batch in range(nbatches):
            theta = theta + (eps*Minv@r).ravel()
            gradU_batch = gradU(theta, dat_batch[:,:, batch], n, batch_size).reshape(p, -1)
            r = r-eps*gradU_batch - eps*C@Minv@r + sqrt_noise@np.random.normal(size = p).reshape(p, -1)
            
        theta_samp[:,i+1] = theta
            
    return theta_samp[:, burns:]<jupyter_output><ipython-input-6-cb78ff7807cb>:1: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "sghmc_numba" failed type inference due to: Untyped global name 'is_pos_def': cannot determine Numba type of <class 'function'>

File "<ipython-input-6-cb78ff7807cb>", line 14:
def sghmc_numba(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
    <source elided>
    
    if not is_pos_def(2*(C-B_hat)*eps):
    ^

  @jit([float64[:,:](float64[:], float64, float64[:,:], float64[:,:], float64[:], float64[:,:], int64, int64, float64[:,:], int64, int64)], cache = True)
<ipython-input-6-cb78ff7807cb>:1: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "sghmc_numba" failed type inference due to: Untyped global name 'is_pos_def': cannot determine Numba type of <class 'function'>

File "<ipython-input-6-cb78ff7807cb>", line 14:
def sghmc_numba(gradU, eps, C, Minv, theta_0, V_ha[...]<jupyter_text>approach 2: add JIT only to sghmc<jupyter_code>@jit([float64[:,:](float64[:], float64, float64[:,:], float64[:,:], float64[:], float64[:,:], int64, int64, float64[:,:], int64, int64)], cache = True)
def sghmc_numba2(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
    
    np.random.seed(seed)
    
    p = theta_0.shape[0]
    n = data.shape[0]
    
    theta_samp = np.zeros((p, epochs))
    theta_samp[:,0] = theta_0
    
    B_hat = 0.5*eps*V_hat
    
    if not is_pos_def(2*(C-B_hat)*eps):
        print("error: noise term is not positive definite")
        return
    
    sqrt_noise = np.linalg.cholesky(2*(C-B_hat)*eps)
    
    sqrtM = np.linalg.cholesky(np.linalg.inv(Minv))
    r = sqrtM@np.random.normal(size = p).reshape(p, -1)
    
    dat_batch, nbatches = data_batch(data, batch_size)
    for i in range(epochs-1):
        
        theta = theta_samp[:,i]
        r = sqrtM@np.random.normal(size = p).reshape(p, -1)
        
        for batch in range(nbatches):
            theta = theta + (eps*Minv@r).ravel()
            gradU_batch = gradU(theta, dat_batch[:,:, batch], n, batch_size).reshape(p, -1)
            r = r-eps*gradU_batch - eps*C@Minv@r +sqrt_noise@np.random.normal(size = p).reshape(p, -1)
            
        theta_samp[:,i+1] = theta
            
    return theta_samp[:, burns:]<jupyter_output><ipython-input-7-5b232f9a54e4>:1: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "sghmc_numba2" failed type inference due to: Untyped global name 'is_pos_def': cannot determine Numba type of <class 'function'>

File "<ipython-input-7-5b232f9a54e4>", line 14:
def sghmc_numba2(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
    <source elided>
    
    if not is_pos_def(2*(C-B_hat)*eps):
    ^

  @jit([float64[:,:](float64[:], float64, float64[:,:], float64[:,:], float64[:], float64[:,:], int64, int64, float64[:,:], int64, int64)], cache = True)
<ipython-input-7-5b232f9a54e4>:1: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "sghmc_numba2" failed type inference due to: Untyped global name 'is_pos_def': cannot determine Numba type of <class 'function'>

File "<ipython-input-7-5b232f9a54e4>", line 14:
def sghmc_numba2(gradU, eps, C, Minv, theta_0, [...]<jupyter_text>Test with Mixture normals<jupyter_code>mu = np.array([-3,3]).reshape(2,1)

def lprior(theta):
    return (-1/(2*10))*theta.T@theta

def ldatap(theta, x):
    return np.log(0.5 * np.exp(-0.5*(theta[0]-x)**2) + 0.5* np.exp(-0.5*(theta[1]-x)**2))

def U(theta, x, n, batch_size):
    return -lprior(theta) - (n/batch_size)*sum(ldatap(theta, x))

gradU = jacobian(U, argnum = 0)

@jit
def lprior_numba(theta):
    return (-1/(2*10))*theta.T@theta

@jit
def ldatap_numba(theta, x):
    return np.log(0.5 * np.exp(-0.5*(theta[0]-x)**2) + 0.5* np.exp(-0.5*(theta[1]-x)**2))

@jit
def U_numba(theta, x, n, batch_size):
    return -lprior_numba(theta) - (n/batch_size)*sum(ldatap_numba(theta, x))

gradU_numba = jacobian(U_numba, argnum = 0)

#Set up data and parameters
np.random.seed(123)
n = 100
x = np.r_[
    np.random.normal(mu[0], 1, n),
    np.random.normal(mu[1], 1, n)].reshape(-1,1)

theta_0 = np.array([-3, 3]) #start at true value
eps = 0.01
V_hat = np.eye(2)
C = np.eye(2)
epochs = 500
burns = 200
batch_size = 50
%timeit sghmc(gradU, eps, C, np.eye(2), theta_0, V_hat, epochs, burns, x, batch_size)
%timeit sghmc_clean(gradU, eps, C, np.eye(2), theta_0, V_hat, epochs, burns, x, batch_size)
%timeit sghmc_numba(gradU, eps, C, np.eye(2), theta_0, V_hat, epochs, burns, x, batch_size)
%timeit sghmc_numba2(gradU, eps, C, np.eye(2), theta_0, V_hat, epochs, burns, x, batch_size)<jupyter_output><ipython-input-101-5b232f9a54e4>:1: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "sghmc_numba2" failed type inference due to: non-precise type pyobject
[1] During: typing of argument at <ipython-input-101-5b232f9a54e4> (24)

File "<ipython-input-101-5b232f9a54e4>", line 24:
def sghmc_numba2(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, seed = 123):
    <source elided>
    dat_batch, nbatches = data_batch(data, batch_size)
    for i in range(epochs-1):
    ^

  @jit([float64[:,:](float64[:], float64, float64[:,:], float64[:,:], float64[:], float64[:,:], int64, int64, float64[:,:], int64, int64)], cache = True)
/opt/conda/lib/python3.6/site-packages/numba/object_mode_passes.py:178: NumbaWarning: Function "sghmc_numba2" was compiled in object mode without forceobj=True.

File "<ipython-input-101-5b232f9a54e4>", line 24:
def sghmc_numba2(gradU, eps, C, Minv, theta_0, V_hat, epochs, burns, data, batch_size, se[...]<jupyter_text>Test for another example<jupyter_code>def gradU(theta, data, n, batch_size):
    return (-n*np.sum(data-theta, axis = 0)/batch_size - np.sum(theta))

@jit
def gradU_numba(theta, data, n, batch_size):
    return (-n*np.sum(data-theta, axis = 0)/batch_size - np.sum(theta))
p = 100
x = np.random.normal(size = (10000, p))
V = np.eye(p)
eps = 0.01
theta_0 = np.zeros(p)
C = np.eye(p)
burn = 100
epochs = 200
batch_size = 500
%timeit sghmc(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_clean(gradU_numba, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_clean(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_numba(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_numba2(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
p = 10
x = np.random.normal(size = (10000, p))
V = np.eye(p)
eps = 0.01
theta_0 = np.zeros(p)
C = np.eye(p)
burn = 100
epochs = 200
batch_size = 500
%timeit sghmc(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_clean(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_numba(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)
%timeit sghmc_numba2(gradU, eps, C, np.eye(p), theta_0, V, epochs, burn, x, batch_size)<jupyter_output>1.22 s ± 13.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
440 ms ± 13.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
462 ms ± 15.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
457 ms ± 16.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
{'fraction_non_alphanumeric': 0.15381607757272442, 'fraction_numerical': 0.041914294651235535, 'mean_word_length': 3.52067868504772, 'pattern_counts': {'":': 0, '<': 44, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 9, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '278411', 'license': 'MIT', 'n_tokens_mistral': 5530, 'n_tokens_neox': 5176, 'n_words': 1179, 'path': 'opt:comparison/Optimization.ipynb', 'repo_name': 'xiangwenmo/SGHMC-final-project', 'size': '12788'}