<jupyter_start><jupyter_text>1D CaseFirst let's generate a curve that has a discontinuity at some location<jupyter_code>x = np.linspace(0, 2, 100)
X = x[:, None]
ell_true = 0.2

kernel = RBFJump(length_scale=ell_true, jump=1.) + WhiteKernel(1e-5, noise_level_bounds='fixed')
K = kernel(X)
np.random.seed(1)
y = stats.multivariate_normal(cov=K).rvs(1)
x_train = x[::5]
X_train = X[::5]
y_train = y[::5]
plt.plot(x, y)
plt.plot(x_train, y_train, ls='', marker='o', c='C0', label='train')
plt.legend();<jupyter_output><empty_output><jupyter_text>Length scale estimation (Bad)Cool, now what if we found the posterior for the length scale assuming that there is no jump? I.e., it is just a vanilla draw from a GP with an RBF kernel.<jupyter_code>ell_vals = np.linspace(1e-5, 0.5, 200)
# ell_vals = np.linspace(0.5, 2, 200)
logpdf_rbf = []
ell_vals_rbf = np.linspace(1e-5, 0.5, 200)
kernel_rbf = RBF(length_scale=ell_true) + WhiteKernel(1e-5, noise_level_bounds='fixed')

for ell in ell_vals:
    kern_i = kernel_rbf.clone_with_theta(np.log([ell,]))
    K_i = kern_i(X_train)
    logpdf_rbf.append(stats.multivariate_normal(cov=K_i).logpdf(y_train))
    
logpdf_rbf = np.array(logpdf_rbf)
ell_pdf_rbf = np.exp(logpdf_rbf - np.max(logpdf_rbf))
plt.plot(ell_vals, ell_pdf_rbf)
plt.axvline(ell_true, 0, 1, color='k');
ell_map_bad = ell_vals[np.argmax(ell_pdf_rbf)]
gp_bad = GaussianProcessRegressor(
    kernel=kernel_rbf.clone_with_theta(np.log([ell_map_bad])), optimizer=None
)
gp_bad.fit(X_train, y_train)
y_pred_bad, std_pred_bad = gp_bad.predict(X, return_std=True)

plt.plot(x, y)
plt.plot(x, y_pred_bad, ls='--', lw=1)
plt.fill_between(x, y_pred_bad+2*std_pred_bad, y_pred_bad-2*std_pred_bad, facecolor='lightgray')
plt.plot(x_train, y_train, ls='', marker='o', c='C0', label='train');<jupyter_output><empty_output><jupyter_text>Not good. The length scale is drawn downwards by the sharp jump. Length scale estimation (Good)Now let's treat it as if the lower half is uncorrelated with the upper half (but they share the same length scale).<jupyter_code>logpdf = []

for ell in ell_vals:
    kern_i = kernel.clone_with_theta(np.log([ell,]))
    K_i = kern_i(X_train)
    logpdf.append(stats.multivariate_normal(cov=K_i).logpdf(y_train))
    
logpdf = np.array(logpdf)
ell_pdf = np.exp(logpdf - np.max(logpdf))
plt.plot(ell_vals, ell_pdf)
plt.axvline(ell_true, 0, 1, color='k');
ell_map_good = ell_vals[np.argmax(ell_pdf)]
gp_good = GaussianProcessRegressor(
    kernel=kernel.clone_with_theta(np.log([ell_map_good])), optimizer=None
)
gp_good.fit(X_train, y_train)
y_pred_good, std_pred_good = gp_good.predict(X, return_std=True)

plt.plot(x, y)
plt.plot(x, y_pred_good, ls='--', lw=1)
plt.fill_between(x, y_pred_good+2*std_pred_good, y_pred_good-2*std_pred_good, facecolor='lightgray')
plt.plot(x_train, y_train, ls='', marker='o', c='C0', label='train');<jupyter_output><empty_output><jupyter_text>Wow!  2D CaseWhy not?<jupyter_code>def cartesian(*arrays):
    """Makes the Cartesian product of arrays.

    Parameters
    ----------
    *arrays : array group, shapes = (N_1,), (N_2,), ..., (N_p,)
            1D arrays where earlier arrays loop more slowly than later ones
    Returns
    -------
    array, shape = (N_1 * N_2 * ... * N_p, p)
        The cartesian product
    """
    N = len(arrays)
    return np.stack(np.meshgrid(*arrays, indexing='ij'), -1).reshape(-1, N)
N1 = 40
N2 = 50
x_dim1 = np.linspace(0, 1, N1)
x_dim2 = np.linspace(0, 2, N2)
X2 = cartesian(x_dim1, x_dim2)
ell_dim1_true = 0.2
ell_dim2_true = 0.5
ell2_true = np.array([ell_dim1_true, ell_dim2_true])
jump2 = np.array([0.5, np.inf])

kernel = RBFJump(length_scale=ell2_true, jump=jump2) + WhiteKernel(1e-5, noise_level_bounds='fixed')
K2 = kernel(X2)
np.random.seed(2)
y2 = stats.multivariate_normal(cov=K2).rvs(1)
y2_mat = y2.reshape(N1, N2).T  # N2 x N1
plt.contourf(x_dim1, x_dim2, y2_mat)
plt.xlabel('x1')
plt.ylabel('x2');<jupyter_output><empty_output><jupyter_text>Space out the training points<jupyter_code>X2_train = X2[::10]
y2_train = y2[::10]
logpdf2 = []
ell_dim1_vals = np.linspace(1e-5, 0.5, 30)
ell_dim2_vals = np.linspace(0.2, 1, 30)
ell_vals2 = cartesian(ell_dim1_vals, ell_dim2_vals)

for ell in ell_vals2:
    kern_i = kernel.clone_with_theta(np.log(ell))
    K_i = kern_i(X2_train)
    logpdf2.append(stats.multivariate_normal(cov=K_i).logpdf(y2_train))
    
logpdf2 = np.array(logpdf2)
ell_pdf2 = np.exp(logpdf2 - np.max(logpdf2))
ell_pdf2_mat = ell_pdf2.reshape(len(ell_dim1_vals), len(ell_dim2_vals)).T
ell_logpdf2_mat = logpdf2.reshape(len(ell_dim1_vals), len(ell_dim2_vals)).T
plt.contourf(ell_dim1_vals, ell_dim2_vals, ell_pdf2_mat)
plt.axhline(ell_dim2_true, 0, 1, c='r')
plt.axvline(ell_dim1_true, 0, 1, c='r')
plt.xlabel('L1')
plt.ylabel('L2');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15457017360384856, 'fraction_numerical': 0.033884124660112945, 'mean_word_length': 5.623268698060942, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '512322', 'license': 'MIT', 'n_tokens_mistral': 2143, 'n_tokens_neox': 2043, 'n_words': 372, 'path': 'notebooks/test_cusp_kernel.ipynb', 'repo_name': 'buqeye/compton-scattering', 'size': '4781'}