Add BroadcastMultiplyAtom (#654)

* Add DotMultiplyAtom

* Update

* update

* Update

* Update format

Author: odow

src/atoms/BroadcastMultiplyAtom.jl

@@ -0,0 +1,120 @@
+# Copyright (c) 2014: Madeleine Udell and contributors
+#
+# Use of this source code is governed by a BSD-style license that can be found
+# in the LICENSE file or at https://opensource.org/license/bsd-2-clause
+
+mutable struct BroadcastMultiplyAtom <: AbstractExpr
+    children::Tuple{AbstractExpr,AbstractExpr}
+    size::Tuple{Int,Int}
+
+    function BroadcastMultiplyAtom(x::AbstractExpr, y::AbstractExpr)
+        (x_r, x_c), (y_r, y_c) = size(x), size(y)
+        if (x_r, x_c) == (y_r, y_c)
+            # Broadcasting over equal sized matrices
+            return new((x, y), (x_r, x_c))
+        elseif x_r == y_r && (x_c == 1 || y_c == 1)
+            # Broadcasting over columns
+            return new((x, y), (x_r, max(x_c, y_c)))
+        elseif x_c == y_c && (x_r == 1 || y_r == 1)
+            # Broadcasting over rows
+            return new((x, y), (max(x_r, y_r), y_c))
+        elseif x_r == y_c && x_c == y_r == 1
+            # x is a column vector and y is a row vector
+            return new((x, y), (x_r, y_c))
+        elseif x_c == y_r && x_r == y_c == 1
+            # x is a row vector and y is a column vector
+            return new((x, y), (y_r, x_c))
+        end
+        return error(
+            "[BroadcastMultiplyAtom] cannot multiply two expressions of sizes $(x.size) and $(y.size)",
+        )
+    end
+end
+
+head(io::IO, ::BroadcastMultiplyAtom) = print(io, ".*")
+
+Base.sign(x::BroadcastMultiplyAtom) = sign(x.children[1]) * sign(x.children[2])
+
+function monotonicity(x::BroadcastMultiplyAtom)
+    return (
+        sign(x.children[2]) * Nondecreasing(),
+        sign(x.children[1]) * Nondecreasing(),
+    )
+end
+
+function curvature(x::BroadcastMultiplyAtom)
+    lhs, rhs = x.children
+    if vexity(lhs) != ConstVexity() && vexity(rhs) != ConstVexity()
+        return NotDcp()
+    end
+    return ConstVexity()
+end
+
+function evaluate(x::BroadcastMultiplyAtom)
+    return reshape(evaluate(x.children[1]) .* evaluate(x.children[2]), size(x))
+end
+
+function new_conic_form!(
+    context::Context{T},
+    x::BroadcastMultiplyAtom,
+) where {T}
+    lhs, rhs = x.children
+    if vexity(lhs) != ConstVexity()
+        if vexity(rhs) != ConstVexity()
+            error(
+                "[BroadcastMultiplyAtom] multiplication of two non-constant expressions is not DCP compliant",
+            )
+        end
+        # Switch arguments so that the left-hand side is constant
+        lhs, rhs = rhs, lhs
+    end
+    # Start by assuming that the constant lhs matrix is the smaller object that
+    # will be broadcast over the larger RHS object. Let Julia automatically
+    # resize it by .* by `ones`.
+    coef = evaluate(lhs) .* ones(T, size(rhs))
+    if size(coef) != size(rhs)
+        # If coef is not the same size as rhs, then we must be broadcasting the
+        # smaller rhs object over the larger coef. In this case, rhs must be a
+        # row or column vector.
+        if size(rhs, 1) == 1
+            # rhs is a row vector. Stretch it out to have the same number of
+            # rows as coef.
+            rhs = ones(T, size(coef, 1)) * rhs
+        else
+            @assert size(rhs, 2) == 1
+            # rhs is a col vector. Stretch it out to have the same number of
+            # columns as coef.
+            rhs = rhs * ones(T, 1, size(coef, 2))
+        end
+    end
+    # For sanity, check that these are the same size.
+    @assert size(coef) == size(rhs)
+    # Represent the array x .* y as D(x) * y
+    ret = SparseArrays.sparse(LinearAlgebra.Diagonal(vec(coef))) * vec(rhs)
+    return conic_form!(context, reshape(ret, size(rhs, 1), size(rhs, 2)))
+end
+
+function Base.Broadcast.broadcasted(
+    ::typeof(*),
+    x::AbstractExpr,
+    y::AbstractExpr,
+)
+    if isequal(x, y)
+        return square(x)
+    elseif x.size == (1, 1) || y.size == (1, 1)
+        return x * y
+    end
+    return BroadcastMultiplyAtom(x, y)
+end
+
+function Base.Broadcast.broadcasted(::typeof(*), x::Value, y::AbstractExpr)
+    return constant(x) .* y
+end
+
+function Base.Broadcast.broadcasted(::typeof(*), x::AbstractExpr, y::Value)
+    return x .* constant(y)
+end
+
+function Base.Broadcast.broadcasted(::typeof(/), x::AbstractExpr, y::Value)
+    return x .* constant(1 ./ y)
+end

src/atoms/MultiplyAtom.jl

@@ -123,58 +123,3 @@ Base.:*(x::Value, y::AbstractExpr) = MultiplyAtom(constant(x), y)
 Base.:*(x::AbstractExpr, y::Value) = MultiplyAtom(x, constant(y))
 
 Base.:/(x::AbstractExpr, y::Value) = MultiplyAtom(x, constant(1 ./ y))
-
-function _dot_multiply(x, y)
-    if size(x) == (1, 1) || size(y) == (1, 1)
-        return x * y
-    end
-    if vexity(x) != ConstVexity()
-        if vexity(y) != ConstVexity()
-            error(
-                "[MultiplyAtom] multiplication of two non-constant expressions is not DCP compliant",
-            )
-        end
-        x, y = y, x
-    end
-    # promote the size of the coefficient matrix, so e.g., 3 .* x works
-    # regardless of the size of x
-    coeff = evaluate(x) .* ones(size(y))
-    # Promote the size of the variable. We've previously ensured neither x nor y
-    # is 1x1 and that the sizes are compatible, so if the sizes aren't equal the
-    # smaller one is size 1.
-    if size(y, 1) < size(coeff, 1)
-        y = ones(size(coeff, 1)) * y
-    elseif size(y, 2) < size(coeff, 2)
-        y = y * ones(1, size(coeff, 1))
-    end
-    ret = SparseArrays.sparse(LinearAlgebra.Diagonal(vec(coeff))) * vec(y)
-    return reshape(ret, size(y, 1), size(y, 2))
-end
-
-# if neither is a constant it's not DCP, but might be nice to support anyway for
-# eg MultiConvex
-function Base.Broadcast.broadcasted(
-    ::typeof(*),
-    x::AbstractExpr,
-    y::AbstractExpr,
-)
-    if isequal(x, y)
-        return square(x)
-    end
-    return _dot_multiply(x, y)
-end
-
-function Base.Broadcast.broadcasted(::typeof(*), x::Value, y::AbstractExpr)
-    return _dot_multiply(constant(x), y)
-end
-
-function Base.Broadcast.broadcasted(::typeof(*), x::AbstractExpr, y::Value)
-    return _dot_multiply(constant(y), x)
-end
-
-function Base.Broadcast.broadcasted(::typeof(/), x::AbstractExpr, y::Value)
-    return _dot_multiply(constant(1 ./ y), x)
-end
-
-# x ./ y and x / y for x constant, y variable is defined in
-# second_order_cone/qol_elemwise.jl

src/atoms/QolElemAtom.jl

@@ -62,7 +62,7 @@ sumsquares(x::AbstractExpr) = square(norm2(x))
 invpos(x::AbstractExpr) = QolElemAtom(constant(ones(x.size)), x)
 
 function Base.Broadcast.broadcasted(::typeof(/), x::Value, y::AbstractExpr)
-    return _dot_multiply(constant(x), invpos(y))
+    return constant(x) .* invpos(y)
 end
 
 function Base.:/(x::Value, y::AbstractExpr)

test/test_atoms.jl

@@ -635,17 +635,112 @@ function test_MultiplyAtom()
         ErrorException(
             "[MultiplyAtom] multiplication of two non-constant expressions is not DCP compliant",
         ),
-        _test_atom(_ -> Variable(2) .* Variable(2), ""),
+        _test_atom(_ -> Variable() * Variable(), ""),
     )
+    return
+end
+
+### BroadcastMultiplyAtom
+
+function test_BroadcastMultiplyAtom()
+    target = """
+    variables: x1, x2
+    minobjective: [0.25 * x1, 0.25 * x2]
+    """
+    _test_atom(target) do context
+        return Variable(2) ./ 4
+    end
+    _test_atom(target) do context
+        return 0.25 .* Variable(2)
+    end
+    _test_atom(target) do context
+        return Variable(2) .* 0.25
+    end
+    target = """
+    variables: x1, x2, x3, x4, x5, x6
+    minobjective: [0.5 * x1, 2.0 * x2, 0.5 * x3, 2.0 * x4, 0.5 * x5, 2.0 * x6]
+    """
+    _test_atom(target) do context
+        x = Variable(2, 3)
+        return x .* [0.5, 2.0]
+    end
+    _test_atom(target) do context
+        x = Variable(2, 3)
+        return [0.5, 2.0] .* x
+    end
+    target = """
+    variables: x1, x2, x3, x4, x5, x6
+    minobjective: [0.5 * x1, 0.5 * x2, 2.0 * x3, 2.0 * x4, 4.0 * x5, 4.0 * x6]
+    """
+    _test_atom(target) do context
+        x = Variable(2, 3)
+        return x .* [0.5 2.0 4.0]
+    end
+    _test_atom(target) do context
+        x = Variable(2, 3)
+        return [0.5 2.0 4.0] .* x
+    end
+    _test_atom(target) do context
+        x = Variable(2, 3)
+        return x ./ [2.0 0.5 0.25]
+    end
+    _test_atom(target) do context
+        x = Variable(2, 3)
+        return x ./ [2.0 0.5 0.25]
+    end
+    target = """
+    variables: x1, x2, x3, x4
+    minobjective: [1.0 * x1, 3.0 * x2, 2.0 * x3, 4.0 * x4]
+    """
+    _test_atom(target) do context
+        x = Variable(2, 2)
+        return x .* [1 2; 3 4]
+    end
+    target = """
+    variables: x1, x2
+    minobjective: [0.5 * x1, 0.5 * x2, 2.0 * x1, 2.0 * x2]
+    """
+    _test_atom(target) do context
+        x = Variable(2, 1)
+        return x .* [0.5 2.0]
+    end
+    target = """
+    variables: x1, x2
+    minobjective: [0.5 * x1, 2.0 * x1, 0.5 * x2, 2.0 * x2]
+    """
+    _test_atom(target) do context
+        x = Variable(1, 2)
+        return x .* [0.5, 2.0]
+    end
+    target = """
+    variables: t1, t2, x1, x2
+    minobjective: [1.0 * t1, 1.0 * t2]
+    [t1, 0.5, x1] in RotatedSecondOrderCone(3)
+    [t2, 0.5, x2] in RotatedSecondOrderCone(3)
+    """
+    _test_atom(target) do context
+        x = Variable(2)
+        return x .* x
+    end
     @test_throws(
         ErrorException(
-            "[MultiplyAtom] multiplication of two non-constant expressions is not DCP compliant",
+            "[BroadcastMultiplyAtom] multiplication of two non-constant expressions is not DCP compliant",
         ),
-        _test_atom(_ -> Variable() * Variable(), ""),
+        _test_atom(_ -> Variable(2) .* Variable(2), ""),
     )
     return
 end
 
+function test_BroadcastMultiply_issue_653()
+    x = Variable(2)
+    fix!(x, [1.0, 2.0])
+    atom = dot(x, [2.0, 1.0])
+    @test evaluate(atom) ≈ 4
+    fix!(x, [2.0, 1.0])
+    @test evaluate(atom) ≈ 5
+    return
+end
+
 ### affine/NegateAtom
 
 function test_NegateAtom()