# Source: content/notes/ml/bias-variance-and-generalization.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.linear_model import Ridge
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures
from threadpoolctl import threadpool_limits

rng = np.random.default_rng(111)
grid = np.linspace(-1, 1, 151)[:, None]
truth = np.sin(np.pi * grid[:, 0])
noise_std = 0.25
training_sets = []
for repetition in range(100):
    x = rng.uniform(-1, 1, size=(40, 1))
    y = np.sin(np.pi * x[:, 0]) + rng.normal(0, noise_std, size=40)
    training_sets.append((x, y))
with threadpool_limits(limits=1):
    for degree, penalty in ((1, 1e-8), (3, 1e-8), (12, 0.1)):
        predictions = []
        for x, y in training_sets:
            model = make_pipeline(PolynomialFeatures(degree, include_bias=False),
                                  Ridge(alpha=penalty, solver="svd"))
            model.fit(x, y)
            predictions.append(model.predict(grid))
        predictions = np.asarray(predictions)
        mean_prediction = predictions.mean(axis=0)
        bias_squared = np.mean((mean_prediction - truth) ** 2)
        variance = predictions.var(axis=0, ddof=0).mean()
        noiseless_error = np.mean((predictions - truth) ** 2)
        np.testing.assert_allclose(noiseless_error, bias_squared + variance, atol=1e-12)
        observed = truth + rng.normal(0, noise_std, size=predictions.shape)
        noisy_error = np.mean((predictions - observed) ** 2)
        predicted_risk = bias_squared + variance + noise_std ** 2
        assert abs(noisy_error - predicted_risk) < 0.015
        print("degree/penalty", degree, penalty, "bias^2", round(bias_squared, 4),
              "variance", round(variance, 4), "noise", noise_std ** 2,
              "measured risk", round(noisy_error, 4))