Add support to pass and use gradient callback in the optim interface

lang_julia/OpenInterfaces/src/OpenInterfaces.jl

@@ -529,21 +529,36 @@ function _make_c_func_wrapper_over_jl_fn(
         # Call the Julia function with the converted arguments.
         result = fn(jl_args...)
 
-        if result == nothing
-            return Int32(0)
-        elseif typeof(result) == Int32
-            return result
-        elseif typeof(result) == Int64
-            try
-                return Int32(result)
-            catch InexactError
-                throw(
-                    "Return type of the callback is Int64 with value `result = $(result)` " *
-                    "and cannot be converted to 32-bit integer due to truncation",
+        if restype == OIF_TYPE_INT
+            if result == nothing
+                return Int32(0)
+            elseif typeof(result) == Int32
+                return result
+            elseif typeof(result) == Int64
+                try
+                    return Int32(result)
+                catch InexactError
+                    throw(
+                        "Return type of the callback is Int64 with value `result = $(result)` " *
+                        "and cannot be converted to 32-bit integer due to truncation",
+                    )
+                end
+            elseif typeof(result) == Float64
+                error(
+                    "[OpenInterfaces.jl] Expected callback return type is " *
+                    "OIF_TYPE_I32, but instead got Float64. " *
+                    "Did you forget to add `return 0` to the end of the callback?",
+                )
+            end
+        elseif restype == OIF_TYPE_F64
+            if typeof(result) == Float64
+                return result
+            else
+                error(
+                    "[OpenInterfaces.jl] Expected callback return type is " *
+                    "OIF_TYPE_F64, got `$(typeof(result))` instead",
                 )
             end
-        elseif typeof(result) == Float64
-            return result
         else
             error(
                 "Unsupported type `$(typeof(result))`, where `result = $(result)` " *

lang_julia/OpenInterfaces/src/interfaces/optim.jl

@@ -30,14 +30,24 @@ mutable struct Self
     x0::Vector{Float64}
     x::Vector{Float64}
     objective_fn_wrapper::Any
+    grad_fn_wrapper::Any
     user_data::Ref{Any}
     oif_user_data::Any
     function Self(impl::String)
         implh = load_impl("optim", impl, 1, 0)
         if implh < OIF_IMPL_STARTING_NUMBER
             error("Could not load implementation '`$impl`' of the `optim` interface")
         end
-        self = new(ImplHandle(implh), 0, Float64[], Float64[], Nothing, Nothing, Nothing)
+        self = new(
+            ImplHandle(implh),
+            0,
+            Float64[],
+            Float64[],
+            Nothing,
+            Nothing,
+            Nothing,
+            Nothing,
+        )
         finalizer(finalizing, self)
     end
 
@@ -67,6 +77,16 @@ function set_objective_fn(self::Self, objective_fn::Function)
     call_impl(self.implh, "set_objective_fn", (self.objective_fn_wrapper,), ())
 end
 
+function set_grad_fn(self::Self, grad_fn::Function)
+    """Specify right-hand side function f."""
+    self.grad_fn_wrapper = make_oif_callback(
+        grad_fn,
+        (OIF_TYPE_ARRAY_F64, OIF_TYPE_ARRAY_F64, OIF_TYPE_USER_DATA),
+        OIF_TYPE_I32,
+    )
+    call_impl(self.implh, "set_grad_fn", (self.grad_fn_wrapper,), ())
+end
+
 # function set_tolerances(self::Self, rtol::Float64, atol::Float64)
 #     """Specify relative and absolute tolerances, respectively."""
 #     call_impl(self.implh, "set_tolerances", (rtol, atol), ())

lang_julia/_impl/optim/optim_jl/optim_jl.jl

@@ -26,13 +26,15 @@ GENERAL_OPTIONS_NAMES = [
 mutable struct Self
     x0::Vector{Float64}
     objective_fn::Union{Function,Nothing}
+    grad_fn::Union{Function,Nothing}
     user_data::Any
     method_name::String
     method_params::Dict
     method::Any
     general_options::Dict
+    grad_values::Any
     function Self()
-        return new([], nothing, nothing, "BFGS", Dict(), BFGS(), Dict())
+        return new([], nothing, nothing, nothing, "BFGS", Dict(), BFGS(), Dict(), nothing)
     end
 end
 
@@ -50,6 +52,11 @@ function set_objective_fn(self::Self, objective_fn)
     self.objective_fn(self.x0, self.user_data)
 end
 
+function set_grad_fn(self::Self, grad_fn)
+    self.grad_fn = grad_fn
+    self.grad_values = zeros(length(self.x0))
+end
+
 function set_method(self, method_name, method_params)
     println(
         "[optim::optim_jl] To check available configuration options, " *
@@ -72,7 +79,7 @@ function set_method(self, method_name, method_params)
     end
     self.method_name = method_name
     self.method_params = merge(self.method_params, method_params)
-    self.method = getfield(Optim, method_symbol)(; method_options)
+    self.method = getfield(Optim, method_symbol)(; method_options...)
     self.general_options = general_options
 end
 
@@ -84,8 +91,20 @@ function minimize(self::Self, out_x)::Tuple{Int,String}
 
     wrapper(x) = self.objective_fn(x, self.user_data)
 
-    result::Optim.MultivariateOptimizationResults =
-        optimize(wrapper, self.x0, self.method, Optim.Options(; self.general_options...))
+    result::Optim.MultivariateOptimizationResults = begin
+        if isnothing(self.grad_fn)
+            optimize(wrapper, self.x0, self.method, Optim.Options(; self.general_options...))
+        else
+            g_wrapper(grad_values, x) = self.grad_fn(x, grad_values, self.user_data)
+            optimize(
+                wrapper,
+                g_wrapper,
+                self.x0,
+                self.method,
+                Optim.Options(; self.general_options...),
+            )
+        end
+    end
     out_x[:] = copy(Optim.minimizer(result))
 
     println("res = ", result)

lang_python/oif_impl/openinterfaces/_impl/optim/scipy_optimize/scipy_optimize.py

@@ -17,10 +17,12 @@ def __init__(self) -> None:
         self.N = 0  # Problem dimension.
         self.x0: np.ndarray | None = None
         self.objective_fn: Callable | None = None
+        self.grad_fn: Callable | None = None
         self.user_data: object | None = None
         self.with_user_data = False
         self.method_name: str | None = None
         self.method_params: str | None = None
+        self.grad_values: np.ndarray | None = None
 
     def set_initial_guess(self, x0: np.ndarray):
         self.x0 = x0
@@ -42,6 +44,10 @@ def set_objective_fn(self, objective_fn):
             np.float64,
         ], msg
 
+    def set_grad_fn(self, grad_fn):
+        self.grad_fn = grad_fn
+        self.grad_values = np.empty((self.N,))
+
     def set_method(self, method_name, method_params):
         available_options = optimize.show_options(
             "minimize", method=method_name, disp=False
@@ -75,10 +81,13 @@ def minimize(self, out_x):
                 "Shapes of the output array and the initial-guess array differ"
             )
 
+        grad_fn = None if self.grad_fn is None else self.grad_wrapper
+
         result = optimize.minimize(
             self.objective_fn,
             self.x0,
-            args=self.user_data,
+            args=(self.user_data,),
+            jac=grad_fn,  # I wonder, why the gradient in scipy.optimize is `jac`?
             method=self.method_name,
             options=self.method_params,
         )
@@ -87,10 +96,8 @@ def minimize(self, out_x):
         print(result)
         return (result.status, result.message)
 
-    # def print_stats(self):
-    #     print("WARNING: `scipy_ode` does not provide statistics")
-
-    # def _rhs_fn_wrapper(self, t, y):
-    #     """Callback that satisfies signature expected by Open Interfaces."""
-    #     self.rhs(t, y, self.ydot, self.user_data)
-    #     return self.ydot
+    def grad_wrapper(self, x, __):
+        grad_values = np.empty_like(x)
+        # grad_values = self.grad_values
+        self.grad_fn(x, grad_values, self.user_data)
+        return grad_values

lang_python/oif_interfaces/openinterfaces/interfaces/optim.py

@@ -114,6 +114,14 @@ def set_objective_fn(self, objective_fn: ObjectiveFn):
         )
         self._binding.call("set_objective_fn", (self.wrapper,), ())
 
+    def set_grad_fn(self, grad_fn):
+        self.grad_wrapper = make_oif_callback(
+            grad_fn,
+            (OIF_TYPE_ARRAY_F64, OIF_TYPE_ARRAY_F64, OIF_USER_DATA),
+            OIF_TYPE_INT,
+        )
+        self._binding.call("set_grad_fn", (self.grad_wrapper,), ())
+
     def set_user_data(self, user_data: object):
         """Specify additional data that will be used for right-hand side function."""
         self.user_data = user_data

tests/lang_julia/test_optim.jl

@@ -7,6 +7,7 @@ using OpenInterfaces.Interfaces.Optim
 # -----------------------------------------------------------------------------
 # Checking what implementations are available.
 POTENTIAL_IMPLEMENTATIONS = ["scipy_optimize", "ipopt_jl", "optim_jl"]
+POTENTIAL_IMPLEMENTATIONS = ["scipy_optimize", "optim_jl"]
 IMPLEMENTATIONS = []
 for impl in POTENTIAL_IMPLEMENTATIONS
     implh = load_impl("optim", impl, 1, 0)
@@ -44,6 +45,26 @@ function rosenbrock_objective_fn(x, __)
     a, b = (0.5, 1.0)
     return sum(a * (x[2:end] - x[1:(end-1)] .^ 2.0) .^ 2.0 + (1 .- x[1:(end-1)]) .^ 2.0) + b
 end
+
+function rosenbrock_objective_fn_with_user_data(x, user_data)
+    """The Rosenbrock function with additional arguments"""
+    a, b = user_data
+    return sum(a * (x[2:end] - x[1:(end-1)] .^ 2.0) .^ 2.0 + (1 .- x[1:(end-1)]) .^ 2.0) + b
+end
+
+function rosenbrock_grad_fn(x, grad_f, user_data)
+    """Gradient of the Rosenbrock function"""
+    a, _ = user_data
+
+    # `xi` is x[i], `xip1` is x[i+1]
+    xi = @view x[1:(end-1)]
+    xip1 = @view x[2:end]
+
+    grad_f[1:(end-1)] .= -4.0 .* a .* xi .* (xip1 .- xi .^ 2.0) .- 2.0 .* (1.0 .- xi)
+    + 2.0 .* a .* (xip1 .- xi .^ 2.0)
+    grad_f[end] = 2.0 * a * (x[end] - x[end-1]^2)
+    return 0
+end
 # -----------------------------------------------------------------------------
 
 
@@ -109,6 +130,26 @@ end
         end
     end
 
+    @testset "test__rosenbrok_with_grad__converges" begin
+        test() do self
+            x0 = [3.14, 2.72, 42.0, 9.81, 8.31]
+            user_data = (0.5, 1.0)
+
+            Optim.set_initial_guess(self, x0)
+            Optim.set_user_data(self, user_data)
+            Optim.set_objective_fn(self, rosenbrock_objective_fn_with_user_data)
+            Optim.set_grad_fn(self, rosenbrock_grad_fn)
+            Optim.set_method(self, "BFGS", Dict())
+
+            status, message = Optim.minimize(self)
+            x = self.x
+
+            @test status == 0
+            @test length(x) == length(x0)
+            @test all(abs.(x .- 1.0) .< 1e-5)  # The solution is [1, 1, ..., 1].
+        end
+    end
+
     @testset "test__set_method_with__wrong_method_params__are_not_accepted" begin
         test() do self
             @test_throws ErrorException Optim.set_method(
@@ -119,6 +160,8 @@ end
         end
     end
 
+
+    # END General tests
     # -------------------------------------------------------------------------
     if "scipy_optimize" in IMPLEMENTATIONS
         @testset "scipy_minimize__rosenbrock_fn__converges_better_with_tighter_tol" begin

tests/lang_python/test_optim.py

@@ -34,6 +34,26 @@ def rosenbrock_objective_fn(x, __):
     return sum(a * (x[1:] - x[:-1] ** 2.0) ** 2.0 + (1 - x[:-1]) ** 2.0) + b
 
 
+def rosenbrock_objective_fn_with_user_data(x, user_data):
+    """The Rosenbrock function with additional arguments"""
+    a, b = user_data
+    return sum(a * (x[1:] - x[:-1] ** 2.0) ** 2.0 + (1 - x[:-1]) ** 2.0) + b
+
+
+def rosenbrock_grad_fn(x, grad_f, user_data):
+    """Gradient of the Rosenbrock function"""
+    a, __ = user_data
+    grad_f[-1] = 0.0
+    grad_f[:-1] = -4.0 * a * x[:-1] * (x[1:] - x[:-1] ** 2.0) - 2.0 * (1.0 - x[:-1])
+    grad_f[1:] += 2.0 * a * (x[1:] - x[:-1] ** 2.0)
+
+
+def rosenbrock_grad_fn_wrapper(x, user_data):
+    grad_values = np.empty_like(x)
+    rosenbrock_grad_fn(x, grad_values, user_data)
+    return grad_values
+
+
 @pytest.fixture(params=IMPLEMENTATIONS)
 def s(request):
     return Optim(request.param)
@@ -86,6 +106,24 @@ def test__parameterized_convex_problem__converges(s):
     assert np.all(np.abs(x - user_data) < 1e-6)
 
 
+def test__rosenbrok_with_grad__converges(s):
+    x0 = np.array([3.14, 2.72, 42.0, 9.81, 8.31])
+    user_data = (0.5, 1.0)
+
+    s.set_initial_guess(x0)
+    s.set_user_data(user_data)
+    s.set_objective_fn(rosenbrock_objective_fn_with_user_data)
+    s.set_grad_fn(rosenbrock_grad_fn)
+    s.set_method("BFGS")
+
+    status, message = s.minimize()
+    x = s.x
+
+    assert status == 0
+    assert len(x) == len(x0)
+    assert np.all(np.abs(x - 1.0) < 1e-5)  # The solution is [1, 1, ..., 1].
+
+
 def test__set_method_with__wrong_method_params__are_not_accepted(s):
     with pytest.raises(RuntimeError):
         s.set_method("BFGS", {"wrong_param": 42})