<jupyter_start><jupyter_text>Clone repo and install requirements<jupyter_code># clone repo
!git clone https://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@github.com/nnuncert/nnuncert

# switch folder and install requirements
%cd nnuncert
!pip install -r requirements.txt<jupyter_output><empty_output><jupyter_text>Main imports<jupyter_code># general imports
import os
import numpy as np
import numexpr as ne
import tensorflow as tf
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.lines import Line2D

# thesis code
import nnuncert
from nnuncert.models import make_model, type2name
from nnuncert.models.dnnc import DNNCModel
from nnuncert.app.toy import make_toy_data, make_toy_plot, gen_2d_gaussian_samples, input2grid, contour_plot_2d
from nnuncert.utils.traintest import TrainTestSplit
from nnuncert.utils.indexing import index_to_rowcol
from scipy.io import loadmat

# some dummy class to load dnnc MATLAB results from
class LoadPred():
    def __init__(self, file):
        Ey = loadmat(file)["Ey"].ravel()
        Vary = loadmat(file)["Vary"].ravel()
        self.std_total = Vary**0.5
        self.pred_mean = Ey

dir1d = os.path.join(*["data", "toy", "1d"])
dir2d = os.path.join(*["data", "toy", "2d"])

legend_handler = {
    "GP-ReLU": "GP-ReLU",
    "MC dropout": "MC dropout",
    "MC dropout 400": "MC dropout",
    "PBP": "PBP",
    "PNN": "PNN",
    "PNN-E": "PNN-E",
    "NLM": "NLM",
    "NLM-E": "NLM-E",
    "DNNC-R": "DNNC-R",
    "DNNC-HS": "DNNC-HS",
}

# our implementation of the horseshoe variant is not 100 percent stable, so we
# ran the experiments with MATLAB and load the results.
MODELS_0 = ["GP-ReLU", "MC dropout 400", "PBP", "PNN", "PNN-E", "NLM", "NLM-E", "DNNC-R", "DNNC-HS"] #@param<jupyter_output><empty_output><jupyter_text>Make the dataGenerate input features on (-4, 4) uniformly and calculate noisy targets with true function 'x**3' and additive homoscedastic noise N(0, 3).<jupyter_code># set seed for reproducibility and make random number generator
seed = 21
rng = np.random.default_rng(seed)

# define function that generates true relationship, possible to use any kind of expression such as "sin(x)", "exp(x)", "3*x**2 - 8x + 14", ...
def reg_func(x):
    reg_func = "x**3"
    return ne.evaluate(reg_func)

# generate input data (x) uniformly from -4 to 4
ppc = 20  # points per cluster
clusters = [[-4, 4]]  # list of cluster bounds (left, right)
x = np.array([rng.choice(np.linspace(x1, x2, 1000), ppc) 
              for (x1, x2) in clusters]).ravel()

# generate responses with function 
noise_std = 3
data_1d = make_toy_data(x, reg_func, noise_std, seed=rng)
# have a look at the data
data_1d.head()
# plot ground truth
fig, ax = plt.subplots(figsize=(8, 4))

minx, maxx = [-6, 6]
x0 = np.linspace(minx, maxx, 80)
ax.scatter(data_1d["x1"], data_1d["y"])
ax.plot(x0, reg_func(x0), "--", color="black")<jupyter_output><empty_output><jupyter_text>Fit all models<jupyter_code># handle model settings
name2type_fitkw = {
    'GP-ReLU' : ('GP', {"sigma_w": 5, "sigma_b": 1, "noise_variance": noise_std**2}),
    'MC dropout' : ('MC dropout', {}),
    'MC dropout 400' : ('MC dropout', {"conv_factor" : 10}), 
    'PBP' : ('PBP', {}),
    'PNN' : ('PNN', {}),
    'PNN-E' : ('PNN-E', {}),
    'NLM' : ('NLM', {"tau2": 2}),
    'NLM-E' : ('NLM-E', {"tau2": 2}),
    'DNNC-R' : ('DNNC-R', {}),
    'DNNC-HS' : ('DNNC-HS', {}),
}
class TrainTestSplitToy1D(TrainTestSplit):
    def __init__(self, df, train_id = None, test_id = None, test_ratio=0.1, norm_x=False, rng=None):
        non_norm = []
        if norm_x is False:
            non_norm = ["x1"]
        super(TrainTestSplitToy1D, self).__init__(df, "y", non_norm=non_norm,train_id=train_id, test_id=test_id, ratio=test_ratio, rng=rng)

# standardize features x
toy = TrainTestSplitToy1D(data_1d, test_ratio=0, norm_x=True, rng=rng)
input_shape = toy.x_train.shape[1]
# handle general settings
# list of hidden layer description (size, act. func, dropout rate)
arch = [[50, "relu", 0]]  #@param
epochs = 40 #@param
verbose = 0 #@param
learning_rate = 0.01 #@param
load_dnnc = True #@param {type:"boolean"}

models = MODELS_0 #@param

xlim, ylim = (-6, 6),  (-100, 100) #@param {type:"raw"}
# handle plot setup
colors = ["mediumblue", "tab:blue", "#b3cde0"]
fig, ax = plt.subplots(3, 3, figsize=(11, 9), sharex=True, sharey=True)
plt.setp(ax, xlim=xlim, ylim=ylim)

# run models and get predictions
for i, mod in enumerate(models):
    if load_dnnc is True and mod == "DNNC-R":
        # assert hash1d == txt2str("data/toy/1d/toy1d_hash.txt")
        pred = LoadPred("data/toy/1d/toy1d_dnnc_ridge.mat")
    elif load_dnnc is True and mod == "DNNC-HS":
        # assert hash1d == txt2str("data/toy/1d/toy1d_hash.txt")
        pred = LoadPred("data/toy/1d/toy1d_dnnc_hs.mat")
    else:
        mtype, fit_kw = name2type_fitkw[mod]
        model = make_model(mtype, input_shape, arch)

        # compile model (loss automatically set in model)
        model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate), metrics=["mae", "mse"])

        # fit NN
        if isinstance(model, DNNCModel):
            dist = Dist._from_values(toy.y_train, method="ssv")
            fit_kw["dist"] = dist      
        model.fit(toy.x_train, toy.y_train, epochs=epochs, verbose=verbose, **fit_kw)

        # make prediction
        pred = model.make_prediction(toy.scale_x(x0))
    
    # get axis and plot to it
    ax_ = ax[index_to_rowcol(i, 3)]
    make_toy_plot(pred, toy.x_train_us, toy.y_train_us, x0=x0, 
                  reg_func=reg_func, std_devs=2, colors=colors, ax=ax_, 
                  title=legend_handler[mod])

# handle axis labeling
for i in [0, 1, 2]:
    ax[i, 0].set_ylabel(r'$y$')
    ax[-1, i].set_xlabel(r'$x$')

# make legend
fig.subplots_adjust(bottom=0.1)
legend_elements = [Line2D([0], [0], color='black', linestyle="--", lw=2),
                    Line2D([0], [0], marker='o', color='w', markerfacecolor=colors[0], markersize=10),
                    Line2D([0], [0], color=colors[1], lw=2),
                    Patch(facecolor=colors[2], alpha=0.7)]
leg = fig.legend(legend_elements, ["True function", "Training data", "Mean prediction", "+- 2 std"], ncol=4, loc='lower center', framealpha=1)
frame = leg.get_frame()
frame.set_linewidth(1)
frame.set_edgecolor("black")<jupyter_output>WARNING:tensorflow:5 out of the last 5 calls to <function Model.make_predict_function.<locals>.predict_function at 0x7ff9d18027a0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:6 out of the last 6 calls to <function Model.make_predict_function.<locals>.predict_function at 0x7ff9d18027a0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function [...]
{'fraction_non_alphanumeric': 0.1453804347826087, 'fraction_numerical': 0.026494565217391304, 'mean_word_length': 4.101178101178101, 'pattern_counts': {'":': 17, '<': 20, '<?xml version=': 0, '>': 20, 'https://': 3, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 1, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '333737', 'license': 'MIT', 'n_tokens_mistral': 2816, 'n_tokens_neox': 2695, 'n_words': 700, 'path': 'notebooks/0_toy_1d.ipynb', 'repo_name': 'pjoachims/nnuncert', 'size': '7360'}