I'm trying to maximize the trace of $X^TAX$ subject to the columns of $X$ being orthonormal, where $A$ is a diagonal matrix and X is not necessarily square, but does not have more columns than rows. I can represent this constraint succinctly as $X^TX = I$, but I'm not sure how I would use this with the method of Lagrange multipliers without explicitly writing out a condition for each row, column pair, as I can't simply add $\lambda(I - X^TX)$ to the derivative to form the Lagrangian. How does one go about doing this?
Your Lagrangian is
Note that $tr(X^TAX)=tr((XX^T)A)$ and $P=XX^T$ is an orthogonal projector.