Supervised PCA model via manifold optimization

supervised_pca.py

"""
Supervised PCA model.

Ritchie, A., Balzano, L., Kessler, D., Sripada, C. S., & Scott, C.
(2020). Supervised PCA: A Multiobjective Approach. arXiv:2011.05309.
"""

import numpy as onp
import autograd.numpy as np
from pymanopt.manifolds import Grassmann, Euclidean, Product
from pymanopt import Problem
from pymanopt.solvers import SteepestDescent

m, n, p = 20, 10, 2
lam = 0.5
y = np.array(
	onp.random.randn(n).tolist()
)
X = np.array(
	onp.random.randn(n, m).tolist()
)

manifold = Product(
	(Grassmann(m, p), Euclidean(p))
)

def cost(params):
	L, w = params
	XL = X @ L
	lr_resid = y - (XL @ w)
	pca_resid = X - (XL @ L.T)
	lr_obj = np.mean(lr_resid ** 2)
	pca_obj = np.mean(pca_resid ** 2)
	return (lam * lr_obj) + ((1 - lam) * pca_obj)

problem = Problem(manifold=manifold, cost=cost)
solver = SteepestDescent(logverbosity=2)
solver.solve(problem)