Support MOI.UserDefinedFunction

[odow]

x-ref https://discourse.julialang.org/t/help-with-performance-in-large-mixed-complementarity-problem/106052

using JuMP
import BitOperations
import PATHSolver

function possible_shock_combinations(J)
    return [BitOperations.bget(j, i) for j in 0:(2^J-1), i in 0:(J-1)]
end

function shock_probabilities(sequences, probs)
    @assert size(sequences, 2) == length(probs)
    return map(1:size(sequences, 1)) do i
        return prod(
            s_i ? p_i : (1 - p_i) for (s_i, p_i) in zip(sequences[i, :], probs)
        )
    end
end

function solve_complementarity(param1, param2, shock, probs, w, i)
    model = Model(PATHSolver.Optimizer)
    set_silent(model)
    @variable(model, x[1:J] >= 0, start = 0.5)
    A = 1 .- shock .* (1 - param2)
    b = -(1 - param1) * w[i]
    @constraint(model, c[j in 1:J],
        b * sum(probs[s] * A[s, j] / (A[s, :]' * x) for s in 1:2^J) + w[j] ⟂ x[j]
    )
    optimize!(model)
    return value.(x)
end

function solve_complementarity(param1, param2, shock, probs, w, i)
    model = Model(PATHSolver.Optimizer)
    set_silent(model)
    @variable(model, x[1:J] >= 0, start = 0.5)
    A = 1 .- shock .* (1 - param2)
    b = -(1 - param1) * w[i]
    f(x...) = sum(probs[s] * A[s, j] / (A[s, :]' * x) for s in 1:2^J)
    function 𝝯f(g, x...)
        g .= 0.0
        for s in 1:2^J
            demoninator = A[s, :]' * x
            for j in 1:J
                g[j] += probs[s] * A[s, j]^2 / demoninator^2
            end
        end
        return
    end
    @operator(model, op_f, J, f, 𝝯f)
    @constraint(model, c[j in 1:J], b * op_f(x...) + w[j] ⟂ x[j])
    optimize!(model)
    return value.(x)
end

J = 10
param1 = 0.5
param2 = 0.1
prob = fill(0.5, J)
shock = possible_shock_combinations(J)
probs = shock_probabilities(shock, prob)
w = ones(J)
i = 1
x = solve_complementarity(param1, param2, shock, probs, w, i)