# 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 threadpoolctl import threadpool_limits

rng = np.random.default_rng(112)
dimensions = [5, 15, 30, 39, 40, 41, 60, 100]
errors = {p: [] for p in dimensions}
training_errors = {p: [] for p in dimensions}
beta = np.array([1.0, -0.8, 0.6, 0.4, -0.2])
with threadpool_limits(limits=1):
    for repetition in range(30):
        X = rng.normal(size=(40, 100))
        y = X[:, :5] @ beta + rng.normal(scale=0.3, size=40)
        test_X = rng.normal(size=(250, 100))
        test_mean = test_X[:, :5] @ beta
        for p in dimensions:
            coefficient = np.linalg.lstsq(X[:, :p], y, rcond=None)[0]
            training_mse = np.mean((X[:, :p] @ coefficient - y) ** 2)
            test_mse = np.mean((test_X[:, :p] @ coefficient - test_mean) ** 2)
            assert np.isfinite(test_mse)
            if p >= 40:
                assert training_mse < 1e-12
            training_errors[p].append(training_mse)
            errors[p].append(test_mse)
for p in dimensions:
    print("features", p, "median training MSE", np.median(training_errors[p]),
          "median noiseless-test MSE", np.median(errors[p]))
assert np.median(errors[40]) > np.median(errors[5])