Is it possible to give a function with 2 outputs as an objective function gradient and a constraint gradient in scipy.optimize.minimize?

Question

I have a very basic optimization problem, which is a 6-dimensional Rosenbrock function. I am using simplex gradient for both the objective function and the constraint, but I want to calculate gradients within a single function to make it computationally cheaper. This is because the gradient calculation for the objective function also contains the necessary data for the constraint gradient calculation. Although it doesn't make sense to calculate them within a single function for this problem, it serves as a toy problem for me. I cannot explain my real, complex problem here, but if you help me with this, I will apply the same logic to my problem.

This is what I tried before but it does not work. I also tried to dump it in a pickle file then try to call it back but in that case optimization is not working properly.

import numpy as np
from scipy.optimize import minimize, NonlinearConstraint


def rosenbrock_6d(x):
    print("obj is calculating!")
    result = 0
    for i in range(5):
        result += 100 * (x[i + 1] - x[i] ** 2) ** 2 + (1 - x[i]) ** 2
    
    return result

def nonlinear_constraint(x):
    print("Nonlinear constraint is calculating!")
    return 10*x[0]**2 + 11*x[1]**2 + 12*x[2]**2 + 13*x[3]**2 + 14*x[4]**2 + 15*x[5]**2


def StoSAG(x):
    print("OBj StoSAG is calculating!")
    CU = np.eye(6)
    n_pert=10
    u_norm_perturbed = np.random.multivariate_normal(x, CU, n_pert)
    u_norm_perturbed = np.transpose(u_norm_perturbed)
    J_ = np.zeros(n_pert)
    NL_array_pert=np.zeros((n_pert,1))
    for j in range(n_pert):
        J_[j] = rosenbrock_6d(u_norm_perturbed[:, j])
        NL_array_pert[j] = nonlinear_constraint(u_norm_perturbed[:, j])
    J_=np.reshape(J_, (n_pert, 1))
    j=J_- rosenbrock_6d(x)
    U= u_norm_perturbed-np.reshape(x, (6, 1))
    du=np.linalg.pinv(np.transpose(U))
    grad_f=np.matmul(du, j).ravel()
    c_nl=(NL_array_pert- nonlinear_constraint(x))
    g_nl=np.matmul(du, c_nl)
    grad_g=np.transpose(g_nl)
    return grad_f, grad_g

nlc = NonlinearConstraint(nonlinear_constraint, -np.inf,2,jac=StoSAG()[1])

# Example usage:
# Define a sample 6-dimensional point
x_values = (1, 2, 3, 4, 5, 6)
result = minimize(rosenbrock_6d, x_values, method='SLSQP', jac=StoSAG()[0],constraints=nlc,options={'maxiter':200,'disp':True,'iprint':2,'ftol':1e-2})

# Print the result
print("Optimal solution: x =", result.x)
print("Optimal objective function value:", result.fun)

Answer 1

There is a long list of Numpy usage issues that I will not cover here as I consider them out-of-scope, but broadly - vectorise more things.

Don't use pickle. You could use SQLite but I consider this overkill without knowing more about your use case. The only practical way I see this being done is an LRU cache:

import numpy as np
from scipy.optimize import minimize, NonlinearConstraint
from functools import lru_cache, wraps


# https://stackoverflow.com/a/52332109/313768
def np_cache(function):
    @lru_cache()
    def cached_wrapper(hashable_array):
        array = np.array(hashable_array)
        return function(array)

    @wraps(function)
    def wrapper(array):
        return cached_wrapper(tuple(array))

    # copy lru_cache attributes over too
    wrapper.cache_info = cached_wrapper.cache_info
    wrapper.cache_clear = cached_wrapper.cache_clear

    return wrapper


def rosenbrock_6d(x: np.ndarray) -> float:
    result = 0
    for i in range(5):
        result += 100*(x[i + 1] - x[i]**2)**2 + (1 - x[i])**2

    return result


def nonlinear_constraint(x: np.ndarray) -> float:
    return 10*x[0]**2 + 11*x[1]**2 + 12*x[2]**2 + 13*x[3]**2 + 14*x[4]**2 + 15*x[5]**2


@np_cache
def StoSAG(x: np.ndarray) -> tuple[
    np.ndarray,  # objective gradient
    np.ndarray,  # constraint gradient
]:
    print("OBj StoSAG is calculating!", flush=True)

    CU = np.eye(6)
    n_pert = 10
    u_norm_perturbed = np.random.multivariate_normal(x, CU, n_pert)
    u_norm_perturbed = np.transpose(u_norm_perturbed)
    J_ = np.zeros(n_pert)
    NL_array_pert = np.zeros((n_pert, 1))

    for j in range(n_pert):
        J_[j] = rosenbrock_6d(u_norm_perturbed[:, j])
        NL_array_pert[j] = nonlinear_constraint(u_norm_perturbed[:, j])

    J_ = np.reshape(J_, (n_pert, 1))
    j = J_ - rosenbrock_6d(x)
    U = u_norm_perturbed - np.reshape(x, (6, 1))
    du = np.linalg.pinv(np.transpose(U))
    grad_f = np.matmul(du, j).ravel()
    c_nl = (NL_array_pert - nonlinear_constraint(x))
    g_nl = np.matmul(du, c_nl)
    grad_g = np.transpose(g_nl)

    return grad_f, grad_g


def StoSAG_objective_grad(x: np.ndarray) -> np.ndarray:
    print('Objective gradient is calculating!', flush=True)
    return StoSAG(x)[0]


def StoSAG_constraint_grad(x: np.ndarray) -> np.ndarray:
    print('Constraint gradient is calculating!', flush=True)
    return StoSAG(x)[1]


def main() -> None:
    nlc = NonlinearConstraint(
        fun=nonlinear_constraint, lb=-np.inf, ub=2,
        jac=StoSAG_constraint_grad,
    )

    # Example usage:
    # Define a sample 6-dimensional point
    x_values = (1, 2, 3, 4, 5, 6)
    result = minimize(
        fun=rosenbrock_6d,
        x0=x_values,
        method='SLSQP',
        jac=StoSAG_objective_grad,
        constraints=nlc,
        options={
            'maxiter': 200,
            'ftol': 1e-2,
        },
    )

    print("Optimal solution: x =", result.x)
    print("Optimal objective function value:", result.fun)


if __name__ == '__main__':
    main()

Constraint gradient is calculating!
OBj StoSAG is calculating!
Objective gradient is calculating!
...