manifolds.jl

using StableRNGs
@testset "Manifolds" begin
    rng = StableRNG(213)

    # Test case: find eigenbasis for first two eigenvalues of a symmetric matrix by minimizing the Rayleigh quotient under orthogonality constraints
    n = 4
    m = 2
    A = Diagonal(range(1, stop=2, length=n))
    fmanif(x) = real(dot(x,A*x)/2)
    gmanif(x) = A*x
    gmanif!(stor,x) = copyto!(stor,gmanif(x))
    # A[2,2] /= 10 #optional: reduce the gap to make the problem artificially harder
    x0 = randn(rng,n,m)+im*randn(rng,n,m)

    @testset "Stiefel, retraction = $retraction" for retraction in [:SVD,:CholQR]
        manif = Optim.Stiefel(retraction)

        # AcceleratedGradientDescent should be compatible also, but I haven't been able to make it converge
        for ls in (Optim.BackTracking,Optim.HagerZhang,Optim.StrongWolfe,Optim.MoreThuente)
            for method in (Optim.GradientDescent, Optim.ConjugateGradient, Optim.LBFGS, Optim.BFGS,
                           Optim.NGMRES, Optim.OACCEL)
                debug_printing && printstyled("Solver: $(summary(method())), linesearch: $(summary(ls()))\n", color=:green)
                res = Optim.optimize(fmanif, gmanif!, x0, method(manifold=manif,linesearch=ls()), Optim.Options(allow_f_increases=true,g_tol=1e-6))
                debug_printing && printstyled("Iter\tf-calls\tg-calls\n", color=:green)
                debug_printing && printstyled("$(Optim.iterations(res))\t$(Optim.f_calls(res))\t$(Optim.g_calls(res))\n", color=:red)
                @test Optim.converged(res)
            end
        end
        res = Optim.optimize(fmanif, gmanif!, x0, Optim.MomentumGradientDescent(mu=0.0, manifold=manif))
        @test Optim.converged(res)
    end

    @testset "Power and Product" begin
        # Power
        @views fprod(x) = fmanif(x[:,1]) + fmanif(x[:,2])
        @views gprod!(stor,x) = (gmanif!(stor[:, 1],x[:, 1]);gmanif!(stor[:, 2],x[:, 2]);stor)
        m1 = Optim.PowerManifold(Optim.Sphere(), (n,), (2,))
        rng = StableRNG(0)
        x0 = randn(rng, n, 2) + im*randn(rng, n, 2)
        res = Optim.optimize(fprod, gprod!, x0, Optim.ConjugateGradient(manifold=m1))
        @test Optim.converged(res)
        minpow = Optim.minimizer(res)

        # Product
        @views fprod(x) = fmanif(x[1:n]) + fmanif(x[n+1:2n])
        @views gprod!(stor,x) = (gmanif!(stor[1:n],x[1:n]);gmanif!(stor[n+1:2n],x[n+1:2n]);stor)
        m2 = Optim.ProductManifold(Optim.Sphere(), Optim.Sphere(), (n,), (n,))
        rng = StableRNG(0)
        x0 = randn(rng,2n) + im*randn(rng,2n)
        res = Optim.optimize(fprod, gprod!, x0, Optim.ConjugateGradient(manifold=m2))
        @test Optim.converged(res)
        minprod = Optim.minimizer(res)

        # results should be exactly equal: same initial guess, same sequence of operations
        @test minpow[:,1] == minprod[1:n]
        @test minpow[:,2] == minprod[n+1:2n]
    end
end