optimization.py

import os
import sys
import itertools
from typing import Any, Callable, Iterable
from dataclasses import dataclass

import numpy as np
from scipy.optimize import fmin_slsqp
from scipy.optimize import least_squares
from typing import Union, Optional
from cyipopt import minimize_ipopt

import casadi as cas

current = os.path.dirname(os.path.realpath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)

from estimation.estimation import Estimator
from derivative_factory import GradientMethod, DerivativeEvaluator


UnionSXMX = Union[cas.SX, cas.MX]


@dataclass
class Constraints:
    eq: Optional[Callable] = None
    ieq: Optional[Callable] = None
    fprime_eq: Optional[Callable] = None
    fprime_ieq: Optional[Callable] = None
    hessvp_eq: Optional[Callable] = None   # hessian vector product
    hessvp_ieq: Optional[Callable] = None  # hessian vector product


class Optimizer:
    default_options_scipy = {'disp': 5,
                             'full_output': True,
                             'max_iter': 1000}

    default_options_ipopt = {'hessian_approximation': 'exact',
                             'limited_memory_max_history': 1000,
                              'print_level': 5}

    default_options_least_squares = {'x_scale': 1.,
                                     'verbose': 0}

    def __init__(self,
                 fun: Any,
                 x0: Any,
                 args: tuple = (),
                 kwargs: dict = {},
                 solver: str = 'slsqp',
                 jac: Optional[Callable] = None,
                 hess: Optional[Callable] = None,
                 hessp: Optional[Callable] = None,
                 bounds: Optional[tuple[Iterable, Iterable]] = None,
                 constraints: Optional[Constraints] = Constraints(),
                 tol: Any = None,
                 callback: Any = None,
                 options: Optional[dict] = None):
        """bounds is a tuple of lower bounds and upper bounds
        """

        self.fun = fun
        self.x0 = x0
        self.args = args      # Additional arguments sent to functions and jacobian
        self.kwargs = kwargs  # Send some options to solver
        self.solver = solver

        self.jac = jac
        self.hess = hess
        self.hessp = hessp
        self.bounds = bounds
        self.constraints = constraints

        self.tol = tol
        self.callback = callback
        self.options = options

    def solve(self, u0) -> Optional[object]:

        if self.solver == 'slsqp':
            sol = self.solve_with_slsqp(u0)
        elif self.solver == 'ipopt':
            sol = self.solve_with_ipopt(u0)
        elif self.solver == 'lsq':
            sol = self.solve_with_lsq(u0)
        elif self.solver == 'casadi':
            sol = self.solve_with_casadi(u0)
        else:
            sol = None

        self.p_opt = sol['x']

        return sol

    def solve_with_lsq(self, u0):
        """Solve with the specific least squares solver in SciPy"""
        options = self.default_options_least_squares | self.options

        sol = least_squares(self.fun,
                            x0=u0,
                            bounds=self.bounds,
                            jac=self.jac,
                            **options)
        return sol

    def solve_with_slsqp(self, u0):
        """Conduct parameter estimation using the SLSQP solver,
           h == 0,
           g >= 0
        """

        if self.bounds is not None:
            bounds = tuple(zip(self.bounds[0], self.bounds[1]))
        else:
            bounds = self.bounds

        # Solve with SLSQP
        options = self.default_options_scipy | self.options

        output = fmin_slsqp(self.fun, u0,
                            bounds=bounds,
                            args=self.args,
                            fprime=self.jac,
                            f_eqcons=self.constraints.eq,
                            f_ieqcons=self.constraints.ieq,
                            fprime_eqcons=self.constraints.fprime_eq,
                            fprime_ieqcons=self.constraints.fprime_ieq,
                            disp=options['disp'],
                            full_output=options['full_output'],
                            iter=options['max_iter'])
        sol = dict(zip(['x', 'f', 'iter', 'mode', 'smode'], output))
        return sol

    def solve_with_ipopt(self, u0):
        """Conduct parameter estimation using the IPOPT solver in the Scipy interface
           h == 0,
           g >= 0,
        """

        options = self.default_options_ipopt | self.options
        if 'disp' in options:
            options['print_level'] = options['disp']
            del options['disp']

        cons = [
            {'type': 'eq', 'fun': self.constraints.eq,
             'jac': self.constraints.fprime_eq, 'hess': self.constraints.hessvp_eq},
            {'type': 'ieq', 'fun': self.constraints.ieq,
             'jac': self.constraints.fprime_ieq, 'hess': self.constraints.hessvp_ieq}]

        if self.bounds is not None:
            bounds = tuple(zip(self.bounds[0], self.bounds[1]))
        else:
            bounds = self.bounds

        output = minimize_ipopt(self.fun, u0,
                                bounds=bounds,
                                args=self.args,
                                jac=self.jac,
                                hess=self.hess,
                                constraints=cons,
                                options=options)
        sol = dict(zip(['x', 'f', 'iter', 'mode', 'smode'], output))
        return sol

    def solve_with_casadi(self, u0):
        """Solve the optimization problem with casadi,
           This function doesn't work well yet. Please don't use.
        """
        # Define optimization problem
        prob = {'f': self.fun, 'x': self.x0}
        solver = cas.nlpsol('solver', 'ipopt', prob, {'ipopt': self.options})

        self.solver = solver
        sol = self.solver(u0)
        return sol


class SymbolOptimizer:
    """
    To do list:
    We need to consider whether we need to get_obj or get_residuals again!
    """
    scipy_like_solvers = ('slsqp', 'ipopt', 'lsq')

    def __init__(self,
                 estimator: Estimator,
                 solver_name='ipopt',
                 options: [Optional[dict]] = None,
                 evaluate_jacobian=False,
                 gradient_by_jacobian=True,
                 gradient_method=GradientMethod.hybrid_ad_fd
                 ):

        self.estimator = estimator

        if solver_name != 'lsq':
            expr = self.estimator.get_obj()
        else:
            residuals = self.estimator.problem.get_residuals()
            expr = cas.vcat(list(itertools.chain.from_iterable(residuals.values())))
        self.derivative_evaluator = DerivativeEvaluator(expr,
                                                        estimator.problem.p_est_sym,
                                                        evaluate_jacobian=evaluate_jacobian,
                                                        gradient_by_jacobian=gradient_by_jacobian,
                                                        gradient_method=gradient_method)

        self.solver_name = solver_name
        self.options = options
        self.solver = None

        self.lbp = np.array(self.estimator.problem.lbp)
        self.ubp = np.array(self.estimator.problem.ubp)

    def estimate(self, u0) -> Optional[object]:
        if self.estimator.problem.log_param:
            u0 = np.log(np.array(u0))

        if self.solver_name in self.scipy_like_solvers:
            sol = self.estimate_with_scipy_like_solver(u0)
        elif self.solver_name == 'casadi':
            sol = self.estimate_with_casadi(u0)
        else:
            sol = None

        self.p_opt = sol['x']
        self.estimator.p_opt = self.p_opt

        return sol

    def estimate_with_scipy_like_solver(self, u0):
        """Solve with the specific least squares solver in SciPy"""
        """Conduct parameter estimation using the SLSQP solver"""

        # Solve with SLSQP
        self.solver = Optimizer(self.derivative_evaluator.obj_fun_py,
                                x0=u0,
                                bounds=(self.lbp, self.ubp),
                                jac=self.derivative_evaluator.grad_fun_py,
                                hess=self.derivative_evaluator.hess_fun_py,
                                solver=self.solver_name,
                                options=self.options)
        sol = self.solver.solve(u0)
        return sol

    def estimate_with_casadi(self, u0):
        """ Estimate the problem """

        # Get all the residuals
        obj = self.estimator.get_obj()

        # Define optimization problem
        self.solver = Optimizer(obj, x0=self.estimator.problem.p_est_sym, solver=self.solver_name)
        sol = self.solver.solve(u0)

        return sol


class EstimatorOptimizer:
    """
    To do list:
    We need to consider whether we need to get_obj or get_residuals again!
    """
    scipy_like_solvers = ('slsqp', 'ipopt', 'lsq')

    def __init__(self,
                 estimator: Estimator,
                 solver_name='ipopt',
                 options: [Optional[dict]] = None,
                 evaluate_jacobian=False,
                 gradient_by_jacobian=True,
                 gradient_method=GradientMethod.hybrid_ad_fd
                 ):

        self.estimator = estimator

        if solver_name != 'lsq':
            expr = self.estimator.get_obj()
        else:
            residuals = self.estimator.problem.get_residuals()
            expr = cas.vcat(list(itertools.chain.from_iterable(residuals.values())))
        self.derivative_evaluator = DerivativeEvaluator(expr,
                                                        estimator.problem.p_est_sym,
                                                        evaluate_jacobian=evaluate_jacobian,
                                                        gradient_by_jacobian=gradient_by_jacobian,
                                                        gradient_method=gradient_method)

        self.solver_name = solver_name
        self.options = options
        self.solver = None

        self.lbp = np.array(self.estimator.problem.lbp)
        self.ubp = np.array(self.estimator.problem.ubp)

    def estimate(self, u0) -> Optional[object]:
        if self.estimator.problem.log_param:
            u0 = np.log(np.array(u0))

        if self.solver_name in self.scipy_like_solvers:
            sol = self.estimate_with_scipy_like_solver(u0)
        elif self.solver_name == 'casadi':
            sol = self.estimate_with_casadi(u0)
        else:
            sol = None

        self.p_opt = sol['x']
        self.estimator.p_opt = self.p_opt

        return sol

    def estimate_with_scipy_like_solver(self, u0):
        """Solve with the specific least squares solver in SciPy"""
        """Conduct parameter estimation using the SLSQP solver"""

        # Solve with SLSQP
        self.solver = Optimizer(self.derivative_evaluator.obj_fun_py,
                                x0=u0,
                                bounds=(self.lbp, self.ubp),
                                jac=self.derivative_evaluator.grad_fun_py,
                                hess=self.derivative_evaluator.hess_fun_py,
                                solver=self.solver_name,
                                options=self.options)
        sol = self.solver.solve(u0)
        return sol

    def estimate_with_casadi(self, u0):
        """ Estimate the problem """

        # Get all the residuals
        obj = self.estimator.get_obj()

        # Define optimization problem
        self.solver = Optimizer(obj, x0=self.estimator.problem.p_est_sym, solver=self.solver_name)
        sol = self.solver.solve(u0)

        return sol