alternating_minimization.jl

# https://discourse.julialang.org/t/convex-jl-objective-function-of-matrix-factorization-with-missing-data/34253
using Clarabel
using Convex
using MathOptInterface
const MOI = MathOptInterface
using JuMP
using Test

# Generate fake data matrix
function gen_data(m, n, k)
    return (10 * rand(m, k) * 2 * rand(k, n))
end

function gen_masks(A, holdout)
    training_mask = rand(size(A)...) .> 1 - holdout
    validation_mask = .!training_mask
    return training_mask, validation_mask
end

function alternating_minimization(f, A, M, Y_init, k, MAX_ITERS)
    m, n = size(A)

    X = Variable(m, k)
    Y = Variable(k, n)

    objective =
        (norm(M .* A - M .* (X * Y), 2) + γ1 * norm(X, 2) + γ2 * norm(Y, 1))

    constraints = [X * Y >= ϵ]

    problem = minimize(objective, constraints)

    Y.value = Y_init
    for i in 1:MAX_ITERS
        fix!(Y)
        res = f(problem)
        free!(Y)
        fix!(X)
        res = f(problem)
        free!(X)
    end

    return problem, X, Y
end

function JuMP_setup!(model, X, Y, A, M, k)
    m, n = size(A)

    @variable(model, t1 >= 0)
    @variable(model, t2 >= 0)
    @variable(model, Y_abs[1:n*k] >= 0)
    @constraint(model, vcat(t1, vec(M .* A - M .* (X * Y))) ∈ SecondOrderCone())
    @constraint(model, vcat(t2, vec(X)) ∈ SecondOrderCone())
    @constraint(model, vec(Y) .<= Y_abs)
    @constraint(model, -vec(Y) .<= Y_abs)
    @objective(model, Min, t1 + &gamma;1 * t2 + &gamma;2 * sum(Y_abs))
    @constraint(model, X * Y .>= &#1013;)
end

function alternating_minimization_JuMP(A, M, Y_init, k, MAX_ITERS)
    m, n = size(A)

    function Y_fix_model(Y)
        model = Model(() -> Clarabel.Optimizer(verbose = false))
        @variable(model, X[1:m, j = 1:k])
        JuMP_setup!(model, X, Y, A, M, k)
        return X, model
    end

    function X_fix_model(X)
        model = Model(() -> Clarabel.Optimizer(verbose = false))
        @variable(model, Y[1:k, j = 1:n])
        JuMP_setup!(model, X, Y, A, M, k)
        return Y, model
    end

    local Xval, model
    Yval = Y_init
    for i in 1:MAX_ITERS
        X, model = Y_fix_model(Yval)
        JuMP.optimize!(model)
        Xval = value.(X)

        Y, model = X_fix_model(Xval)
        JuMP.optimize!(model)
        Yval = value.(Y)
    end

    return model, Xval, Yval
end

const &gamma;1 = 1.0
const &gamma;2 = 1.0
const &#1013; = 0.0001
MAX_ITERS = 2

m, n, k = 125, 125, 3
# m, n, k = 500, 500, 5
holdout = 0.80

A = gen_data(m, n, k)
Mt, Mv = gen_masks(A, holdout)

Y_init = rand(k, n)
@info "Running with Convex.jl..." (m, n, k)
@time p1, X1, Y1 =
    alternating_minimization(A, Mt, Y_init, k, MAX_ITERS) do problem
        return solve!(problem, () -> Clarabel.Optimizer(verbose = false))
    end;

@info "Running with JuMP..." (m, n, k)
@time model, X3, Y3 = alternating_minimization_JuMP(A, Mt, Y_init, k, MAX_ITERS);

@testset "Same results" begin
    @test evaluate(X1) &asymp; X3 atol = 1e-2 rtol = 1e-2
    @test evaluate(Y1) &asymp; Y3 atol = 1e-2 rtol = 1e-2
end