Improve performance of ExpAtom, LogAtom, and EntropyAtom (#618)

Author: ericphanson

src/atoms/EntropyAtom.jl

@@ -36,13 +36,12 @@ entropy(x::AbstractExpr) = sum(EntropyAtom(x))
 entropy_elementwise(x::AbstractExpr) = EntropyAtom(x)
 
 function new_conic_form!(context::Context, e::EntropyAtom)
-    # -x log x >= t  <=>  x exp(t/x) <= 1  <==>  (t,x,1) in exp cone
+    # -x log(x) >= t  <=>  x exp(t/x) <= 1 <=> (t, x, 1) in ExponentialCone()
     x = e.children[1]
-    m, n = size(x)
-    t = Variable(m, n)
-    for i in 1:m, j in 1:n
-        f = vcat(t[i, j], x[i, j], 1)
-        add_constraint!(context, GenericConstraint{MOI.ExponentialCone}(f))
-    end
-    return conic_form!(context, t)
+    # To choose the permutation, we want the elements of the constraint to be
+    #   (t, x, 1) in ExponentialCone()
+    # but with the identity permutation, the default is:
+    #   (x, 1, t) in ExponentialCone()
+    # so [3, 1, 2] permutes it to the correct order.
+    return _add_vectorized_exp_cone(context, x, [3, 1, 2])
 end

src/atoms/ExpAtom.jl

@@ -30,13 +30,67 @@ evaluate(x::ExpAtom) = exp.(evaluate(x.children[1]))
 Base.exp(x::AbstractExpr) = ExpAtom(x)
 
 function new_conic_form!(context::Context{T}, e::ExpAtom) where {T}
-    # exp(x) \leq z  <=>  (x,1,z) \in ExpCone
+    # exp(x) <= t  <=>  (x, 1, t) in ExponentialCone()
     x = e.children[1]
+    x_tape = conic_form!(context, x)
+    if x_tape isa SPARSE_VECTOR
+        return exp.(x_tape)
+    end
+    return _add_vectorized_exp_cone(context, x, [1, 2, 3])
+end
+
+"""
+    _add_vectorized_exp_cone(
+        context::Context{T},
+        x,
+        permutation::Vector{Int},
+    ) where {T}
+
+Constrains `(x[i], 1, t[i]) ∈ ExponentialCone()` for each element of `x` and
+returns `t`.
+
+Permutation is a permuted vector of `[1, 2, 3]` to reorder the triple before it
+is constrained. This is helpful for `LogAtom` and `EntropyAtom`.
+
+## Motivation
+
+A naive implementation of this method is:
+```julia
     m, n = size(x)
-    z = Variable(m, n)
+    t = Variable(m, n)
     for i in 1:m, j in 1:n
-        f = vcat(x[i, j], 1, z[i, j])
+        f = vcat(x[i, j], 1, t[i, j])[collect(permutation)]
         add_constraint!(context, GenericConstraint{MOI.ExponentialCone}(f))
     end
-    return conic_form!(context, z)
+    return conic_form!(context, t)
+end
+```
+This is slow because we are indexing on the Convex side, and Convex is based
+around vector/matrix operations. We don't want to produce n*m IndexAtoms!
+
+Instead, we will drop to the MOI level to implement this in terms of scalar
+operations.
+"""
+function _add_vectorized_exp_cone(
+    context::Context{T},
+    x,
+    permutation::Vector{Int},
+) where {T}
+    @assert issetequal(permutation, (1, 2, 3))
+    @assert length(permutation) == 3
+    m, n = size(x)
+    t = Variable(m, n)
+    t_tape = conic_form!(context, t)
+    ts = t_tape.variables::Vector{MOI.VariableIndex}
+    x_tape = conic_form!(context, x)
+    xs = MOI.Utilities.scalarize(to_vaf(x_tape))
+    for (xi, ti) in zip(xs, ts)
+        args = (xi, T(1), ti)[permutation]
+        MOI.add_constraint(
+            context.model,
+            MOI.Utilities.operate(vcat, T, args...),
+            MOI.ExponentialCone(),
+        )
+    end
+    return t_tape
 end

src/atoms/LogAtom.jl

@@ -30,13 +30,12 @@ evaluate(x::LogAtom) = log.(evaluate(x.children[1]))
 Base.log(x::AbstractExpr) = LogAtom(x)
 
 function new_conic_form!(context::Context, e::LogAtom)
-    # log(z) \geq x  <=> (x,1,z) \in ExpCone
-    z = e.children[1]
-    m, n = size(z)
-    x = Variable(m, n)
-    for i in 1:m, j in 1:n
-        f = vcat(x[i, j], 1, z[i, j])
-        add_constraint!(context, GenericConstraint{MOI.ExponentialCone}(f))
-    end
-    return conic_form!(context, x)
+    # log(x) >= t  <=> (t, 1, x) in ExponentialCone()
+    x = e.children[1]
+    # To choose the permutation, we want the elements of the constraint to be:
+    #   (t, 1, x) in ExponentialCone()
+    # The default is:
+    #   (x, 1, t) in ExponentialCone()
+    # so [3, 2, 1] permutes it to the correct order.
+    return _add_vectorized_exp_cone(context, x, [3, 2, 1])
 end

src/problem_depot/problems/exp.jl

@@ -22,6 +22,19 @@
         @test evaluate(exp(y)) ≈ 1 atol = atol rtol = rtol
     end
 
+    # Test for constant `exp` (#613)
+    y = Variable()
+    x = constant([1, 2, 3])
+    p = minimize(sum(exp(x)) + y, y >= 0; numeric_type = T)
+    if test
+        @test problem_vexity(p) == ConvexVexity()
+    end
+    handle_problem!(p)
+    if test
+        @test p.optval ≈ sum(exp.([1, 2, 3])) atol = atol rtol = rtol
+        @test evaluate(y) ≈ 0 atol = atol
+    end
+
     y = Variable()
     p = minimize(exp(y), y >= 1; numeric_type = T)