# train and test variables beginning with an underscore are unprocessed.
_x_train, y_train = cos_linear_edge_data(n_data=n_train, sigma=0.1, seed=seed)
preprocessor = StandardPreprocessor()
x_train, _x_plot, x_plot = preprocessor.preprocess_create_x_train_x_plot(
_x_train, test_ds=0.8)
y_ground_truth = ground_truth_cos_function(_x_plot)
layer_units = [50, 20] + [2]
layer_activations = ["relu"] * (len(layer_units) - 1) + ["linear"]
# %% codecell
y_lim = [-5, 5]
figsize = (10, 6)
fig, ax = plt.subplots(figsize=figsize)
plot_training_data(_x_train, y_train, fig=fig, ax=ax, y_lim=y_lim)
plot_ground_truth(_x_plot, y_ground_truth, fig=fig, ax=ax, alpha=0.2)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_ylim(y_lim)
fig.suptitle("Data with approximately linear edges", fontsize=15)
ax.tick_params(bottom=False, left=False, labelbottom=False, labelleft=False)
ax.legend()
fig.savefig(figure_dir.joinpath(f"data_{experiment_name}.pdf"),
bbox_inches="tight")
# %%
# General training
epochs = 100
batch_size = n_train
# %% markdown
layer_activations = ["relu"] * (len(layer_units) - 1) + ["linear"]
weight_prior = tfd.Normal(0, 1)
bias_prior = weight_prior
a = 0.5
b = 0.01
_var_d = tfd.InverseGamma(a, b)
noise_scale_prior = tfd.TransformedDistribution(distribution=_var_d,
bijector=tfp.bijectors.Invert(
tfp.bijectors.Square()))
# %% codecell
y_lim = [-5, 7]
fig, ax = plt.subplots(figsize=(8, 8))
plot_training_data(x_train, y_train, fig=fig, ax=ax, y_lim=y_lim)
plot_ground_truth(x_plot, y_ground_truth, fig=fig, ax=ax)
ax.legend()
# %%
initial_learning_rate = 0.05
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate, decay_steps=20, decay_rate=0.9, staircase=True)
net = MapDensityNetwork(
input_shape=input_shape,
layer_units=layer_units,
layer_activations=layer_activations,
weight_prior=weight_prior,
bias_prior=bias_prior,
n_train=n_train,
learning_rate=0.01,