<commit_before>from abc import abstractmethod

from kernel_exp_family.estimators.estimator_oop import EstimatorBase
from kernel_exp_family.tools.assertions import assert_array_shape

try:
    import autograd.numpy as np  # Thinly-wrapped numpy
except ImportError:
    import numpy as np

def SE(x, y, l=2):
    # ASSUMES COLUMN VECTORS
    diff = x - y;
    return np.squeeze(np.exp(-np.dot(diff.T, diff) / (2 * l ** 2)))

def SE_dx(x, y, l=2):
    return SE(x, y, l) * (y - x) / l ** 2

def SE_dx_dx(x, y, l=2):
    # Doing SE(x,y,l)*((y-x)**2/l**4 - 1/l**2) does not work!
    return SE(x, y, l) * (y - x) ** 2 / l ** 4 - SE(x, y, l) / l ** 2

def SE_dx_dy(x, y, l=2):
    SE_tmp = SE(x, y, l)
    term1 = SE_tmp * np.eye(x.size) / l ** 2
    term2 = SE_tmp * np.dot((x - y), (x - y).T) / l ** 4
    return term1 - term2

def SE_dx_dx_dy(x, y, l=2):
    SE_tmp = SE(x, y, l)
    term1 = SE_tmp * np.dot((x - y) ** 2, (x - y).T) / l ** 6
    term2 = SE_tmp * 2 * np.diag((x - y)[:, 0]) / l ** 4
    term3 = SE_tmp * np.repeat((x - y), x.size, 1).T / l ** 4
    return term1 - term2 - term3

def SE_dx_dx_dy_dy(x, y, l=2):
    SE_tmp = SE(x, y, l)
    term1 = SE_tmp * np.dot((x - y), (x - y).T) ** 2 / l ** 8
    term2 = SE_tmp * 6 * np.diagflat((x - y) ** 2) / l ** 6  # diagonal (x-y)
    term3 = (1 - np.eye(x.size)) * SE_tmp * np.repeat((x - y), x.size, 1) ** 2 / l ** 6  # (x_i-y_i)^2 off-diagonal 
    term4 = (1 - np.eye(x.size)) * SE_tmp * np.repeat((x - y).T, x.size, 0) ** 2 / l ** 6  # (x_j-y_j)^2 off-diagonal
    term5 = SE_tmp * (1 + 2 * np.eye(x.size)) / l ** 4
    
    return term1 - term2 - term3 - term4 + term5


def SE_dx_i_dx_j(x, y, l=2):
    """ Matrix of \frac{\partial k}{\partial x_i \partial x_j}"""
    pairwise_dist = (y-x).dot((y-x).T)

    term1 = SE(x, y, l)*pairwise_dist/l**4
    term2 = SE(x, y, l)*np.eye(pairwise_dist.shape[0])/l**2

    return term1 - term2


def SE_dx_i_dx_i_dx_j(x, y, l=2):
    """ Matrix of \frac{\partial k}{\partial x_i^2 \partial x_j}"""
    pairwise_dist_squared_i = ((y-x)**2).dot((y-x).T)
    row_repeated_distances = np.repeat((y-x).T,
                                       pairwise_dist_squared_i.shape[0],
                                       axis=0)

    term1 = SE(x, y, l)*pairwise_dist_squared_i/l**6
    term2 = SE(x, y, l)*row_repeated_distances/l**4
    term3 = term2*2*np.eye(pairwise_dist_squared_i.shape[0])

    return term1 - term2 - term3

def compute_h(kernel_dx_dx_dy, data):
    n, d = data.shape
    h = np.zeros((n, d))
    for b in range(n):
        for a in range(n):
            h[b, :] += np.sum(kernel_dx_dx_dy(data[a, :].reshape(-1, 1), data[b, :].reshape(-1, 1)), axis=0)
            
    return h / n

def compute_G(kernel_dx_dy, data):
    n, d = data.shape
    G = np.zeros((n, n, d, d))
    for a in range(n):
        for b in range(n):
            x = data[a, :].reshape(-1, 1)
            y = data[b, :].reshape(-1, 1)
            G[a, b, :, :] = kernel_dx_dy(x, y)
            
    return G

# compute_G(SE_dx_dy, test)

def compute_xi_norm_2(kernel_dx_dx_dy_dy, data):
    n, _ = data.shape
    norm_2 = 0
    for a in range(n):
        for b in range(n):
            x = data[a, :].reshape(-1, 1)
            y = data[b, :].reshape(-1, 1)
            norm_2 += np.sum(kernel_dx_dx_dy_dy(x, y))
            
    return norm_2 / n ** 2

def build_system(X, sigma, lmbda):
        # esben parametrised kernel in terms of l as exp(-||---|| / (2*l^2)
    # therefore sigma = 2*(l**2)
    l = np.sqrt(np.float(sigma) / 2)
    
    n, d = X.shape
    
    SE_dx_dx_dy_l = lambda x, y: SE_dx_dx_dy(x, y, l)
    SE_dx_dy_l = lambda x, y: SE_dx_dy(x, y, l)
    SE_dx_dx_dy_dy_l = lambda x, y: SE_dx_dx_dy_dy(x, y, l)
    
    h = compute_h(SE_dx_dx_dy_l, X)
    G = compute_G(SE_dx_dy_l, X)
    xi_norm_2 = compute_xi_norm_2(SE_dx_dx_dy_dy_l, X)
    
    A = np.zeros((n * d + 1, n * d + 1))
    
    # Top left element
    A[0, 0] = np.sum(h ** 2) / n + lmbda * xi_norm_2
    
    # First row and first column
    for b in range(n):
        for j in range(d):
            A[0, 1 + b * d + j] = np.sum(G[:, b, :, j] * h) / n + lmbda * h[b, j]
            A[1 + b * d + j, 0] = A[0, 1 + b * d + j]
    
    # All other elements - (n*d)x(n*d) lower right submatrix
    for a in range(n):
        for i in range(d):
            for b in range(n):
                for j in range(d):
                    A[1 + b * d + j, 1 + a * d + i] = np.sum(G[a, :, i, :] * G[:, b, :, j]) / n + lmbda * G[a, b, i, j]
    
    b = np.zeros((n * d + 1, 1))

    b[0] = -xi_norm_2
    for a in range(n):
        for i in range(d):
            b[1 + a * d + i] = -h[a, i]
            
    return A, b

def fit(X, sigma, lmbda):
    A, b = build_system(X, sigma, lmbda)
    x = np.linalg.solve(A, b)
    alpha = x[0]
    beta = x[1:].reshape(n, d)
    return alpha, beta<commit_msg>Fix minor bug in full estimator<commit_after>from abc import abstractmethod

from kernel_exp_family.estimators.estimator_oop import EstimatorBase
from kernel_exp_family.tools.assertions import assert_array_shape

try:
    import autograd.numpy as np  # Thinly-wrapped numpy
except ImportError:
    import numpy as np

def SE(x, y, l=2):
    # ASSUMES COLUMN VECTORS
    diff = x - y;
    return np.squeeze(np.exp(-np.dot(diff.T, diff) / (2 * l ** 2)))

def SE_dx(x, y, l=2):
    return SE(x, y, l) * (y - x) / l ** 2

def SE_dx_dx(x, y, l=2):
    # Doing SE(x,y,l)*((y-x)**2/l**4 - 1/l**2) does not work!
    return SE(x, y, l) * (y - x) ** 2 / l ** 4 - SE(x, y, l) / l ** 2

def SE_dx_dy(x, y, l=2):
    SE_tmp = SE(x, y, l)
    term1 = SE_tmp * np.eye(x.size) / l ** 2
    term2 = SE_tmp * np.dot((x - y), (x - y).T) / l ** 4
    return term1 - term2

def SE_dx_dx_dy(x, y, l=2):
    SE_tmp = SE(x, y, l)
    term1 = SE_tmp * np.dot((x - y) ** 2, (x - y).T) / l ** 6
    term2 = SE_tmp * 2 * np.diag((x - y)[:, 0]) / l ** 4
    term3 = SE_tmp * np.repeat((x - y), x.size, 1).T / l ** 4
    return term1 - term2 - term3

def SE_dx_dx_dy_dy(x, y, l=2):
    SE_tmp = SE(x, y, l)
    term1 = SE_tmp * np.dot((x - y), (x - y).T) ** 2 / l ** 8
    term2 = SE_tmp * 6 * np.diagflat((x - y) ** 2) / l ** 6  # diagonal (x-y)
    term3 = (1 - np.eye(x.size)) * SE_tmp * np.repeat((x - y), x.size, 1) ** 2 / l ** 6  # (x_i-y_i)^2 off-diagonal 
    term4 = (1 - np.eye(x.size)) * SE_tmp * np.repeat((x - y).T, x.size, 0) ** 2 / l ** 6  # (x_j-y_j)^2 off-diagonal
    term5 = SE_tmp * (1 + 2 * np.eye(x.size)) / l ** 4
    
    return term1 - term2 - term3 - term4 + term5


def SE_dx_i_dx_j(x, y, l=2):
    """ Matrix of \frac{\partial k}{\partial x_i \partial x_j}"""
    pairwise_dist = (y-x).dot((y-x).T)

    term1 = SE(x, y, l)*pairwise_dist/l**4
    term2 = SE(x, y, l)*np.eye(pairwise_dist.shape[0])/l**2

    return term1 - term2


def SE_dx_i_dx_i_dx_j(x, y, l=2):
    """ Matrix of \frac{\partial k}{\partial x_i^2 \partial x_j}"""
    pairwise_dist_squared_i = ((y-x)**2).dot((y-x).T)
    row_repeated_distances = np.repeat((y-x).T,
                                       pairwise_dist_squared_i.shape[0],
                                       axis=0)

    term1 = SE(x, y, l)*pairwise_dist_squared_i/l**6
    term2 = SE(x, y, l)*row_repeated_distances/l**4
    term3 = term2*2*np.eye(pairwise_dist_squared_i.shape[0])

    return term1 - term2 - term3

def compute_h(kernel_dx_dx_dy, data):
    n, d = data.shape
    h = np.zeros((n, d))
    for b in range(n):
        for a in range(n):
            h[b, :] += np.sum(kernel_dx_dx_dy(data[a, :].reshape(-1, 1), data[b, :].reshape(-1, 1)), axis=0)
            
    return h / n

def compute_G(kernel_dx_dy, data):
    n, d = data.shape
    G = np.zeros((n, n, d, d))
    for a in range(n):
        for b in range(n):
            x = data[a, :].reshape(-1, 1)
            y = data[b, :].reshape(-1, 1)
            G[a, b, :, :] = kernel_dx_dy(x, y)
            
    return G

# compute_G(SE_dx_dy, test)

def compute_xi_norm_2(kernel_dx_dx_dy_dy, data):
    n, _ = data.shape
    norm_2 = 0
    for a in range(n):
        for b in range(n):
            x = data[a, :].reshape(-1, 1)
            y = data[b, :].reshape(-1, 1)
            norm_2 += np.sum(kernel_dx_dx_dy_dy(x, y))
            
    return norm_2 / n ** 2

def build_system(X, sigma, lmbda):
        # esben parametrised kernel in terms of l as exp(-||---|| / (2*l^2)
    # therefore sigma = 2*(l**2)
    l = np.sqrt(np.float(sigma) / 2)
    
    n, d = X.shape
    
    SE_dx_dx_dy_l = lambda x, y: SE_dx_dx_dy(x, y, l)
    SE_dx_dy_l = lambda x, y: SE_dx_dy(x, y, l)
    SE_dx_dx_dy_dy_l = lambda x, y: SE_dx_dx_dy_dy(x, y, l)
    
    h = compute_h(SE_dx_dx_dy_l, X)
    G = compute_G(SE_dx_dy_l, X)
    xi_norm_2 = compute_xi_norm_2(SE_dx_dx_dy_dy_l, X)
    
    A = np.zeros((n * d + 1, n * d + 1))
    
    # Top left element
    A[0, 0] = np.sum(h ** 2) / n + lmbda * xi_norm_2
    
    # First row and first column
    for b in range(n):
        for j in range(d):
            A[0, 1 + b * d + j] = np.sum(G[:, b, :, j] * h) / n + lmbda * h[b, j]
            A[1 + b * d + j, 0] = A[0, 1 + b * d + j]
    
    # All other elements - (n*d)x(n*d) lower right submatrix
    for a in range(n):
        for i in range(d):
            for b in range(n):
                for j in range(d):
                    A[1 + b * d + j, 1 + a * d + i] = np.sum(G[a, :, i, :] * G[:, b, :, j]) / n + lmbda * G[a, b, i, j]
    
    b = np.zeros((n * d + 1, 1))

    b[0] = -xi_norm_2
    for a in range(n):
        for i in range(d):
            b[1 + a * d + i] = -h[a, i]
            
    return A, b

def fit(X, sigma, lmbda):
    n, d = X.shape
    A, b = build_system(X, sigma, lmbda)
    x = np.linalg.solve(A, b)
    alpha = x[0]
    beta = x[1:].reshape(n, d)
    return alpha, beta
{'fraction_non_alphanumeric': 0.1794030156939173, 'fraction_numerical': 0.02748999897425377, 'mean_word_length': 1.997233323086382, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3391533', 'n_tokens_mistral': 4462, 'n_tokens_neox': 4304, 'n_words': 1095}