% Quadratic vs linear predictor for a curved relation
% Textbook: Part I, "Linear Prediction", non-linear prediction exercise
%
% X = 0.9*W1 + 0.1*W2 and Y = W1^2, with W1, W2 ~ U[0,1] independent.
% Quadratic family Yhat = a*x^2: a = E[X^2 Y]/E[X^4],
%   MSE = E[Y^2] - E[X^2 Y]^2/E[X^4] = 1.478e-3.
% Linear family: MSE = Var[Y](1 - rho^2) = 6.572e-3, about four times larger.
% Both are the best member of their own family; the gap is the price of a
% straight line through a curved cloud.

clearvars; close all
rng(0)

% Theoretical moments, E[W^n] = 1/(n+1)
E_X2Y = 0.9^2/5 + 2*0.9*0.1*(1/4)*(1/2) + 0.1^2*(1/3)*(1/3);
E_X4 = 0.9^4/5 + 4*0.9^3*0.1*(1/4)*(1/2) + 6*0.9^2*0.1^2*(1/3)*(1/3) ...
    + 4*0.9*0.1^3*(1/2)*(1/4) + 0.1^4/5;
a = E_X2Y/E_X4;
mse_quad = 1/5 - E_X2Y^2/E_X4;

E_X = 1/2; E_Y = 1/3;
var_X = (0.9^2 + 0.1^2)/12;
var_Y = 1/5 - 1/9;
cov_XY = 3/40;
rho = cov_XY/sqrt(var_X*var_Y);
slope = cov_XY/var_X;
mse_lin = var_Y*(1 - rho^2);

% Simulation
N = 1e6;
W1 = rand(1, N);
W2 = rand(1, N);
X = 0.9*W1 + 0.1*W2;
Y = W1.^2;

fprintf('quadratic: a = %.4f, MSE theory = %.4e, simulation = %.4e\n', ...
    a, mse_quad, mean((Y - a*X.^2).^2))
fprintf('linear:    slope = %.4f, MSE theory = %.4e, simulation = %.4e\n', ...
    slope, mse_lin, mean((Y - E_Y - slope*(X - E_X)).^2))

n = 1500;
figure
plot(X(1:n), Y(1:n), '.')
hold on
x = linspace(0, 1, 100);
plot(x, a*x.^2, 'b', x, E_Y + slope*(x - E_X), 'r--', 'LineWidth', 1.5)
hold off
xlim([0 1])
xlabel('X')
ylabel('Y')
legend('data', sprintf('%.2f x^2', a), 'linear', 'Location', 'northwest')
grid on