# 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 FactorAnalysis, FastICA, NMF
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

rng = np.random.default_rng(66)
positive = rng.uniform(0, 2, (180, 3)) @ rng.uniform(0, 2, (3, 8))
train, test = train_test_split(positive, random_state=66)
nmf = NMF(n_components=3, init="nndsvda", max_iter=1200, tol=1e-3, random_state=66)
nmf.fit(train)
codes = nmf.transform(test)
reconstruction = nmf.inverse_transform(codes)
assert np.all(codes >= 0)
print("NMF reconstruction MSE:", mean_squared_error(test, reconstruction))
sources = np.column_stack([rng.laplace(size=300), rng.uniform(-2, 2, 300)])
mixed = sources @ np.array([[1.0, 0.5], [0.2, 1.0]])
ica = FastICA(n_components=2, whiten="unit-variance", random_state=66, max_iter=600)
recovered = ica.fit_transform(mixed)
assert np.allclose(ica.inverse_transform(recovered), mixed, atol=1e-6)
correlations = np.abs(np.corrcoef(recovered.T, sources.T)[:2, 2:])
print("ICA source correlations, allowing sign/order changes:\n", correlations)
fa = FactorAnalysis(n_components=2, random_state=66).fit(train)
factors = fa.transform(test)
assert factors.shape == (len(test), 2) and np.isfinite(fa.score(test))
print("Factor-analysis held-out mean log likelihood:", fa.score(test))