Maximize a function of an orthogonal matrix

Question

I'm trying to write up a small code that, given a set of normal vibrational modes for a molecule, will convert them to localized vibrational modes. To do this I'm following the procedure from J. Chem. Phys. 130, 084106 (2009) where they introduce $\tilde{Q}=QU$ as the orthogonal transformation ($U$) from normal modes ($Q$) to local modes ($\tilde{Q}$). $U$ is a $k\times k$ matrix and Q is a $3n\times k$ matrix where $n$ is the number of molecules and k is the number of normal modes. The $3n$ comes from the fact that each molecule has x,y,z coordinates (the $Q$ matrix is occasionally expressed as a three index quantity so the cartesian coordinates can be summed over).

To find the $U$ that makes this transformation, they define the function: $$\xi(U)=\sum_p^k\sum_i^n\bigg(\sum_{\alpha=1}^3\sum_q^k(Q_{i\alpha,q}U_{qp})^2\bigg)^2$$ and say that the correct $U$ for the transformation is the one that maximizes this expression.

I have written up this function and its gradient in Python and am trying to optimize it using scipy.optimize.minimize for a small case, but all methods fail with the message "Desired error not necessarily achieved due to precision loss.".

The key issue seems to be maintaining the orthogonality of $U$, as without this constraint I suspect the problem is underdetermined. I have attempted to include this constraint by including it in the gradient as recommended by a question on Cross Validated, but it doesn't seem to work.

My question basically boils down to this: is there something obvious I'm missing that could improve the convergence with scipy.optimize.minimize or is the problem just intractable for the methods scipy uses and I need to try a different approach? My hope is not to have to reinvent the wheel, so suggestions that already have an implementation are preferred.

If it helps, I have included the code I'm using below.

import numpy as np
from scipy.optimize import minimize

#Q is the matrix of normal modes (3N x k). 

#U is the desired unitary transform matrix (k x k). Optimize only accepts functions of vectors, so have to flatten U. 

#Pipek-Mezey function/criterion
def PM(U,Q,NAtoms,NModes):
    U=np.reshape(U,(NModes,NModes))
    Qtilde=np.einsum('ij,jk->ik',Q,U)

    #Make a 3-index array of cartesian, number of atoms, number of modes for ease of summing
    Qtilde=np.reshape(Qtilde,(3,NAtoms,NModes))
    ##A**2 squares all elements of array, not the square of a matrix
    ##i...->... sums over first index
    C=np.einsum('i...->...',Qtilde**2)

    result=-np.einsum('ij->',C**2)

    return result  

#Gradient of PM condition, shares args with PM. Gradient includes additional factors to 
#ensure U remains orthogonal 
def PM_grad(U,Q,NAtoms,NModes):
    gradient=np.zeros((NModes,NModes))
    U=np.reshape(U,(NModes,NModes))

    #Taking derivatives with respect to U_rs. Sum over s first to reuse intermediates
    for s in range(NModes):
        #Each term of the gradient only contains a column of Qtilde
        Qtildecol=np.einsum('ij,j->i',Q,U[:,s])
        Qtildecol=np.reshape(Qtildecol,(3,NAtoms))
        #Similarly, only a column of C
        Ccol= np.einsum('i...->...',Qtildecol**2)    
        for r in range(NModes):
            Qcol=np.reshape(Q[:,r],(3,NAtoms))
            Intermediate=np.einsum('ij,ij->j',Qcol,Qtildecol)
            gradient[r,s]=-4*np.einsum('i,i->',Intermediate,Ccol)    

    #Gradient correction to ensure orthogonality
    gradient=gradient-(U@np.transpose(gradient)@U) 

    return np.reshape(gradient,(NModes**2))

NModes=3
NAtoms=3
#U is the desired unitary transform matrix (k x k). Optimize only accepts functions of vectors, so have to flatten U. 
U=[0.0,-.800,-.6,.8,-.36,.48,.6,.48,-.64]
#Q is the matrix of normal modes (3N x k).
Q=[0.00000,0.00000,0.00000,0.00000,0.00000,-0.06812,-0.06867,0.05247,0.00000,0.00000,0.00000,0.00000,0.44800,
0.57030,0.54056,0.54492,-0.41639,-0.45330,0.00000,0.00000,0.00000,-0.44800,-0.57030,0.54056,0.54492,-0.41639,0.45330]
Q=np.transpose(np.reshape(Q,(3,9)))


result=minimize(PM,U,args=(Q,NAtoms,NModes),method="BFGS",jac=PM_grad,options={'disp': True})
print(result.x)

Answer 1

There are specialized methods for the minimization of a differentiable function $f(X)$ subject to the orthogonality constraint $X^{T}X=I$. See for example:

Lai, Rongjie, and Stanley Osher. “A Splitting Method for Orthogonality Constrained Problems.” Journal of Scientific Computing 58, no. 2 (2014): 431–449.

Wen, Zaiwen, and Wotao Yin. “A Feasible Method for Optimization with Orthogonality Constraints.” Mathematical Programming 142, no. 1–2 (2013): 397–434.

Answer 2

A simple way to introduce orthogonality constraints is to parametrize all orthogonal matrices using, either, the Cayley transform, $Q=(I-A)(I+A)^{-1}$, or the matrix exponential, $Q = \exp(A)$. In both cases, $Q$ will be orthogonal if $A$ is skew-symmetric. Searching over the space of skew-symmetric matrices is easy, as an element of the space can be represented by a linear combination of the skew-symmetric basis matrices. The Cayley transform is its own inverse, so any orthogonal $Q$ that does not have -1 as an eigenvalue can be represented in this way. Similarly, $A=\log(Q)$, so as long as $Q$ is invertible (which it must be if it is orthogonal), its logarithm exists, so any orthogonal $Q$ can be represented as $\exp{A}$.

Another good reference is: "The Geometry of Algorithms with Orthogonality Constraints", Alan Edelman, Tomás A. Arias, and Steven T. Smith

Answer 3

I wound up using a different approach to get the optimization working. In the paper, they use an approach based on Jacobi rotations where they rotate pairs of modes by a calculated optimal angle to increase the function value. This approach had previously been used to perform a similar orthogonal optimization to localize molecular orbitals (Rev. Mod. Phys. 35, 457). These optimal rotations never decrease the function value and the rotations maintain the orthogonality of an initial guess, so they eventually converge to an orthogonal solution. They perform these rotations, sweeping over all the pairs of modes, until all rotations cause less than some threshold of change to the function value.

I initially couldn't get this approach to work which led to the attempt in my original question, but I have managed to fix my own implementation of their Jacobi Sweep approach. If it is of interest, I can include the code or a link to it here.