# Source: content/notes/ml/unsupervised-learning.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.decomposition import PCA
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

rng = np.random.default_rng(55)
latent = rng.normal(size=(240, 2))
X = latent @ rng.normal(size=(2, 6)) + 0.05 * rng.normal(size=(240, 6))
X_train, X_test = train_test_split(X, random_state=55)
scaler = StandardScaler().fit(X_train)
Z_train, Z_test = scaler.transform(X_train), scaler.transform(X_test)
model = PCA(n_components=2, svd_solver="full").fit(Z_train)
scores = model.transform(Z_test)
reconstruction = model.inverse_transform(scores)
singular = np.linalg.svd(Z_train - Z_train.mean(axis=0), compute_uv=False)
assert np.allclose(model.explained_variance_, singular[:2] ** 2 / (len(Z_train) - 1))
assert scores.shape == (len(X_test), 2)
assert mean_squared_error(Z_test, reconstruction) < mean_squared_error(Z_test, np.zeros_like(Z_test))
print("Variance fractions:", model.explained_variance_ratio_)
print("Held-out reconstruction MSE:", mean_squared_error(Z_test, reconstruction))