Implement evaluation

.github/workflows/ci.yml

@@ -25,6 +25,13 @@ jobs:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
+      - name: MOI
+        shell: julia --project=@. {0}
+        run: |
+          using Pkg
+          Pkg.add([
+              PackageSpec(name="MathOptInterface", rev="bl/arraydiff"),
+          ])
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         with:

README.md

@@ -9,6 +9,29 @@ Experimental addition of array support to `MOI.Nonlinear.ReverseAD`
 > [!WARNING]
 > This code is still very much experimental
 
+You need to use the following branch of [MathOptInterface](https://github.com/jump-dev/MathOptInterface.jl)
+```julia
+Pkg.add(PackageSpec(name="MathOptInterface", rev="bl/arraydiff"))
+```
+
+Supported operators:
+
+- [x] `vect`, e.g, `[1, 2]`.
+- [x] `dot`
+- [ ] `row`, e.g. `[1 2; 3 4]`
+- [ ] `hcat`
+- [ ] `vcat`
+- [ ] `norm`
+- [ ] Matrix-Vector product
+- [ ] Matrix-Matrix product
+- [ ] Broadcasting scalar operator
+
+Supported levels of AD:
+
+- [x] 0-th order
+- [ ] 1-st order
+- [ ] 2-nd order
+
 [build-img]: https://github.com/blegat/ArrayDiff.jl/actions/workflows/ci.yml/badge.svg?branch=main
 [build-url]: https://github.com/blegat/ArrayDiff.jl/actions?query=workflow%3ACI
 [codecov-img]: https://codecov.io/gh/blegat/ArrayDiff.jl/branch/main/graph/badge.svg

src/reverse_mode.jl

@@ -227,7 +227,27 @@ function _forward_eval(
                 @inbounds f.partials_storage[children_arr[idx1+2]] =
                     !(condition == 1)
                 f.forward_storage[k] = ifelse(condition == 1, lhs, rhs)
-            else # atan, min, max or vect
+            elseif node.index == 10 # vect
+                for j in _eachindex(f.sizes, k)
+                    ix = children_arr[children_indices[j]]
+                    @s f.partials_storage[ix] = one(T)
+                    val = @s f.forward_storage[ix]
+                    @j f.forward_storage[k] = val
+                end
+            elseif node.index == 11 # dot
+                idx1, idx2 = children_indices
+                ix1 = children_arr[idx1]
+                ix2 = children_arr[idx2]
+                tmp_dot = zero(T)
+                for j in _eachindex(f.sizes, ix1)
+                    v1 = @j f.forward_storage[ix1]
+                    v2 = @j f.forward_storage[ix2]
+                    @j f.partials_storage[ix1] = v2
+                    @j f.partials_storage[ix2] = v1
+                    tmp_dot += v1 * v2
+                end
+                @s f.forward_storage[k] = tmp_dot
+            else # atan, min, max
                 f_input = _UnsafeVectorView(d.jac_storage, N)
                 ∇f = _UnsafeVectorView(d.user_output_buffer, N)
                 for (r, i) in enumerate(children_indices)
@@ -251,13 +271,27 @@ function _forward_eval(
             end
         elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE
             child_idx = children_arr[f.adj.colptr[k]]
-            ret_f, ret_f′ = Nonlinear.eval_univariate_function_and_gradient(
-                operators,
-                node.index,
-                f.forward_storage[child_idx],
-            )
-            f.forward_storage[k] = ret_f
-            f.partials_storage[child_idx] = ret_f′
+            if node.index == 1 # :+
+                for j in _eachindex(f.sizes, k)
+                    @j f.partials_storage[child_idx] = one(T)
+                    val = @j f.forward_storage[child_idx]
+                    @j f.forward_storage[k] = val
+                end
+            elseif node.index == 2 # :-
+                for j in _eachindex(f.sizes, k)
+                    @j f.partials_storage[child_idx] = -one(T)
+                    val = @j f.forward_storage[child_idx]
+                    @j f.forward_storage[k] = -val
+                end
+            else
+                ret_f, ret_f′ = Nonlinear.eval_univariate_function_and_gradient(
+                    operators,
+                    node.index,
+                    f.forward_storage[child_idx],
+                )
+                f.forward_storage[k] = ret_f
+                f.partials_storage[child_idx] = ret_f′
+            end
         elseif node.type == Nonlinear.NODE_COMPARISON
             children_idx = SparseArrays.nzrange(f.adj, k)
             result = true

src/sizes.jl

@@ -41,6 +41,14 @@ function _storage_range(sizes::Sizes, k::Int)
     return sizes.storage_offset[k] .+ _eachindex(sizes, k)
 end
 
+function _getscalar(x, sizes::Sizes, k::Int)
+    return x[sizes.storage_offset[k]+1]
+end
+
+function _setscalar!(x, value, sizes::Sizes, k::Int)
+    return x[sizes.storage_offset[k]+1] = value
+end
+
 function _getindex(x, sizes::Sizes, k::Int, j)
     return x[sizes.storage_offset[k]+j]
 end
@@ -49,6 +57,34 @@ function _setindex!(x, value, sizes::Sizes, k::Int, j)
     return x[sizes.storage_offset[k]+j] = value
 end
 
+"""
+    @s(storage[node]) -> _getscalar(storage, f.sizes, node)
+    @s(storage[node] = value) -> _setscalar!(storage, value, f.sizes, node)
+
+This "at scalar" converts `getindex` and `setindex!` calls to access the
+scalar in a vector corresponding to a node.
+"""
+macro s(expr)
+    if Meta.isexpr(expr, :(=)) && length(expr.args) == 2
+        lhs, rhs = expr.args
+        @assert Meta.isexpr(lhs, :ref)
+        @assert length(expr.args) == 2
+        return Expr(
+            :call,
+            :_setscalar!,
+            esc(lhs.args[1]),
+            esc(rhs),
+            esc(:(f.sizes)),
+            esc(lhs.args[2]),
+        )
+    elseif Meta.isexpr(expr, :ref) && length(expr.args) == 2
+        arr, idx = expr.args
+        return Expr(:call, :_getscalar, esc(arr), esc(:(f.sizes)), esc(idx))
+    else
+        error("Unsupported expression `$expr`")
+    end
+end
+
 """
     @j(storage[node]) -> _getindex(storage, f.sizes, node, j)
     @j(storage[node] = value) -> _setindex!(storage, value, f.sizes, node, j)

test/ArrayDiff.jl

@@ -21,67 +21,42 @@ function runtests()
     return
 end
 
-function test_objective_quadratic_univariate()
+function test_objective_dot_univariate()
+    model = Nonlinear.Model()
     x = MOI.VariableIndex(1)
-    scalar = Nonlinear.Model()
-    Nonlinear.set_objective(model, :($x * $x))
-    vector = Nonlinear.Model()
-    Nonlinear.set_objective(vector, :([$x] * [$x]))
+    Nonlinear.set_objective(model, :(dot([$x], [$x])))
     evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x])
-    MOI.initialize(evaluator, [:Grad, :Jac, :Hess])
-    @test MOI.eval_objective(evaluator, [1.2]) == 1.2^2 + 1
-    g = [NaN]
-    MOI.eval_objective_gradient(evaluator, g, [1.2])
-    @test g == [2.4]
-    @test MOI.hessian_objective_structure(evaluator) == [(1, 1)]
-    H = [NaN]
-    MOI.eval_hessian_objective(evaluator, H, [1.2])
-    @test H == [2.0]
-    @test MOI.hessian_lagrangian_structure(evaluator) == [(1, 1)]
-    H = [NaN]
-    MOI.eval_hessian_lagrangian(evaluator, H, [1.2], 1.5, Float64[])
-    @test H == 1.5 .* [2.0]
-    MOI.eval_hessian_lagrangian_product(
-        evaluator,
-        H,
-        [1.2],
-        [1.2],
-        1.5,
-        Float64[],
+    MOI.initialize(evaluator, [:Grad])
+    sizes = evaluator.backend.objective.expr.sizes
+    @test sizes.ndims == [0, 1, 0, 1, 0]
+    @test sizes.size_offset == [0, 1, 0, 0, 0]
+    @test sizes.size == [1, 1]
+    @test sizes.storage_offset == [0, 1, 2, 3, 4, 5]
+    @test MOI.eval_objective(evaluator, [1.2]) == 1.2^2
+    return
+end
+
+function test_objective_dot_bivariate()
+    model = Nonlinear.Model()
+    x = MOI.VariableIndex(1)
+    y = MOI.VariableIndex(1)
+    Nonlinear.set_objective(
+        model,
+        :(dot([$x, $y] - [1, 2], -[1, 2] + [$x, $y])),
     )
-    @test H == [1.5 * 2.0 * 1.2]
+    evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x])
+    MOI.initialize(evaluator, [:Grad])
+    sizes = evaluator.backend.objective.expr.sizes
+    @test sizes.ndims == [0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0]
+    @test sizes.size_offset == [0, 6, 5, 0, 0, 4, 0, 0, 3, 2, 1, 0, 0, 0, 0, 0]
+    @test sizes.size == [2, 2, 2, 2, 2, 2, 2]
+    @test sizes.storage_offset ==
+          [0, 1, 3, 5, 6, 7, 9, 10, 11, 13, 15, 17, 18, 19, 21, 22, 23]
+    g = [NaN]
+    @test MOI.eval_objective(evaluator, [5, -1]) ≈ 25
     return
 end
 
 end  # module
 
 TestArrayDiff.runtests()
-
-import MathOptInterface as MOI
-import ArrayDiff
-using Test
-const Nonlinear = MOI.Nonlinear
-model = Nonlinear.Model()
-x = MOI.VariableIndex(1)
-Nonlinear.set_objective(model, :(dot([$x], [$x])))
-evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x])
-MOI.initialize(evaluator, [:Grad])
-sizes = evaluator.backend.objective.expr.sizes
-@test MOI.eval_objective(evaluator, [1.2]) == 1.2^2
-@test sizes.ndims == [0, 1, 0, 1, 0]
-@test sizes.size_offset == [0, 1, 0, 0, 0]
-@test sizes.size == [1, 1]
-@test sizes.storage_offset == [0, 1, 2, 3, 4, 5]
-
-y = MOI.VariableIndex(1)
-Nonlinear.set_objective(model, :(dot([$x, $y] - [1, 2], -[1, 2] + [$x, $y])))
-MOI.initialize(evaluator, [:Grad])
-sizes = evaluator.backend.objective.expr.sizes
-@test MOI.eval_objective(evaluator, [1.2]) == 1.2^2
-@test sizes.ndims == [0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0]
-@test sizes.size_offset == [0, 6, 5, 0, 0, 4, 0, 0, 3, 2, 1, 0, 0, 0, 0, 0]
-@test sizes.size == [2, 2, 2, 2, 2, 2, 2]
-@test sizes.storage_offset ==
-      [0, 1, 3, 5, 6, 7, 9, 10, 11, 13, 15, 17, 18, 19, 21, 22, 23]
-g = [NaN]
-MOI.eval_objective_gradient(evaluator, g, [1.2])

test/runtests.jl

@@ -1,3 +1,2 @@
 include("ReverseAD.jl")
-
-import SparseArrays
+include("ArrayDiff.jl")