upoqa.utils.model — UPOQA v1.0.1 documentation

# Source code for upoqa.utils.model

    # Copyright (c) 2025, Yichuan Liu and Yingzhou Li
    #
    # This program is free software: you can redistribute it and/or modify
    # it under the terms of the GNU General Public License as published by
    # the Free Software Foundation, either version 3 of the License, or
    # (at your option) any later version.
    #
    # This program is distributed in the hope that it will be useful,
    # but WITHOUT ANY WARRANTY; without even the implied warranty of
    # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    # GNU General Public License for more details.
    #
    # You should have received a copy of the GNU General Public License
    # along with this program. If not, see <http://www.gnu.org/licenses/>.

    """
    Interpolation Surrogate Model
    =============================

    Maintain two core classes:

    1. :class:`~upoqa.utils.model.QuadSurrogate`: Quadratic interpolation surrogate
       model using the derivative-free symmetric Broyden update [1]_. Serves as the
       element model in UPOQA.

    2. :class:`~upoqa.utils.model.OverallSurrogate`: Overall surrogate model formed by
       assembling multiple element models.

    References
    ----------
    .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic
        models that satisfy interpolation conditions. *Math. Program.* 100, 1
        (May 2004), 183-215.
    """

    import numpy as np
    import numpy.linalg as LA
    from copy import deepcopy
    from .interp_set import InterpSet, OverallInterpSet
    from typing import Any, Callable, Optional, Union, Tuple, List, Dict
    from .params import UPOQAParameterList

    __all__ = ["SurrogateLinAlgError", "QuadSurrogate", "OverallSurrogate"]

    class SurrogateLinAlgError(Exception):
        """
        Exception raised for linear algebra errors in the surrogate model

        Attributes
        ----------
        ele_idx : int or None
            Index of the element function that caused the error.
            If the error is not related to a specific element, this is None.
        ele_name : Any or None
            Name of the element function that caused the error.
            If the error is not related to a specific element, this is None.
        """

        def __init__(
            self,
            message: str,
            ele_idx: Optional[int] = None,
            ele_name: Optional[Any] = None,
        ) -> None:
            super().__init__(message)
            self.ele_idx = ele_idx
            self.ele_name = ele_name

    class QuadSurrogate:
        """
        Quadratic interpolation surrogate model using the *derivative-free symmetric
        Broyden update* [1]_. Serves as the elemental surrogate model in UPOQA.

        Parameters
        ----------
        interp_set: :class:`~upoqa.utils.interp_set.InterpSet`
            The interpolation set. It must be provided in order for the interpolation
            surrogate model to be built based on the interpolation points.
        center: ndarray, optional
            The initial center of the surrogate model. If provided, it overrides the ``n``
            argument. Otherwise, the center will be set as a zero vector.
        n : int, optional
            Problem dimension. Required if ``center`` is not provided.
            Note that either ``n`` or ``center`` must be provided.
        ref_surrogate : :class:`~upoqa.utils.model.QuadSurrogate`, optional
            A reference surrogate model to help initialize the KKT coefficients
            and other cached information.

        Attributes
        ----------
        n : int
            Problem dimension
        interp_set : :class:`~upoqa.utils.interp_set.InterpSet`
            The interpolation set used to build the surrogate model
        npt : int
            Number of interpolation points in the interpolation set
        model_center : ndarray, shape (n,)
            Center of the surrogate model
        model_cons : float
            Constant term of the model at the center
        model_grad : ndarray, shape (n,)
            Gradient of the model at the center
        model_hess_explicit : ndarray, shape (n, n)
            Explicit Hessian of the model
        model_hess_implicit : ndarray, shape (npt,)
            Implicit Hessian of the model

        References
        ----------
        .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic models that
            satisfy interpolation conditions. *Math. Program.* 100, 1 (May 2004), 183-215.
        """

        def __init__(
            self,
            interp_set: InterpSet,
            center: Optional[np.ndarray] = None,
            n: Optional[int] = None,
            ref_surrogate: Optional["QuadSurrogate"] = None,
        ) -> None:
            if center is None:
                assert n is not None
                center = np.zeros(n)
            n = center.size
            assert (
                interp_set.npt <= (n + 1) * (n + 2) // 2
            ), "QuadSurrogate requires npt to be <= (n + 1) * (n + 2) / 2."
            assert interp_set.npt >= n + 1, "QuadSurrogate requires npt to be >= n + 1."
            self.reset(interp_set, center, ref_surrogate)

        def reset(
            self,
            interp_set: InterpSet,
            center: np.ndarray,
            ref_surrogate: Optional["QuadSurrogate"] = None,
        ) -> None:
            """
            Reset the model with the given interpolation set, center point and dimension,
            and optionally a reference model.

            Parameters
            ----------
            interp_set: :class:`~upoqa.utils.interp_set.InterpSet`
                The interpolation set. It must be provided in order for the interpolation
                surrogate model to be built based on the interpolation points.
            center: ndarray
                The initial center of the surrogate model.
            ref_surrogate : :class:`~upoqa.utils.model.QuadSurrogate`, optional
                A reference surrogate model to help initialize the KKT coefficients
                and other cached information.
            """
            n = center.size
            self.n = n
            self.interp_set = interp_set
            self.npt = self.interp_set.npt

            # surrogate model coeff
            self.model_hess_explicit = np.zeros((n, n), dtype=np.float64)
            self.model_hess_implicit = np.zeros(interp_set.npt, dtype=np.float64)
            self.model_grad = np.zeros(n, dtype=np.float64)
            self.model_cons = 0.0
            self.model_center = (
                center.reshape((n,))
                if center is not None
                else np.zeros(n, dtype=np.float64)
            )

            if ref_surrogate is not None:
                # if ref_surrogate is provided, initialize KKT coefficients to match ref_surrogate
                self._KKT_R: np.ndarray = ref_surrogate._KKT_R.copy()
                self._KKT_B: np.ndarray = ref_surrogate._KKT_B.copy()
                self.cached_Y_shift: np.ndarray = (
                    ref_surrogate.cached_Y_shift.copy()
                )  # interp set points - center
                self.cached_x_anchor: np.ndarray = ref_surrogate.cached_x_anchor.copy()
                self.cached_x_anchor_idx: int = ref_surrogate.cached_x_anchor_idx
                self._negative_s_idx: int = ref_surrogate._negative_s_idx
                self._update_surrogate_coeff()
            else:
                # otherwise, initialize all KKT coefficients to zero and update them later
                self._KKT_R = np.zeros((self.npt - self.n - 1, self.npt), dtype=np.float64)
                self._KKT_B = np.zeros((self.npt + self.n, self.n), dtype=np.float64)
                self.cached_Y_shift = (
                    self.interp_set.get_interp_set()[0] - self.model_center
                )
                self.cached_x_anchor = self.interp_set.get_anchor()[0].copy()
                self.cached_x_anchor_idx = self.interp_set.x_anchor_idx
                self._negative_s_idx = 0
                self._init_KKT_and_model_coeff(self.interp_set.last_init_step_size)

        def _init_KKT_and_model_coeff(self, step_size: Optional[float] = None) -> None:
            """
            Initialize the KKT and model coefficients based on the interpolation set and
            the given step size. Detailed explanations can be found around formula (2.10)
            of [1]_.

            Parameters
            ----------
            step_size : float, optional
                The step size used for initializing the KKT and model coefficients. If not
                provided, the initialization step size from the interpolation set is used.

            References
            ----------
            .. [1] Michael J. D. Powell. 2009. The BOBYQA algorithm for bound constrained
                optimization without derivatives. *Cambridge NA Report NA2009/06,
                University of Cambridge, Cambridge* 26 (2009): 26-46.
            """
            init_step_size = step_size or self.interp_set.init_step_size

            self._negative_s_idx == 0
            xpt = np.zeros((self.npt, self.n))
            for k in range(self.npt):
                if 1 <= k <= self.n:
                    xpt[k, k - 1] = init_step_size
                    if self.npt <= k + self.n:
                        self._KKT_B[0, k - 1] = -1.0 / init_step_size
                        self._KKT_B[k, k - 1] = 1.0 / init_step_size
                        self._KKT_B[self.npt + k - 1, k - 1] = -0.5 * init_step_size**2.0
                elif self.n < k <= 2 * self.n:
                    xpt[k, k - self.n - 1] = -init_step_size
                    self._KKT_B[k, k - self.n - 1] = -0.5 / xpt[k - self.n, k - self.n - 1]
                    self._KKT_B[k - self.n, k - self.n - 1] = (
                        -self._KKT_B[0, k - self.n - 1] - self._KKT_B[k, k - self.n - 1]
                    )
                    self._KKT_R[k - self.n - 1, 0] = -np.sqrt(2.0) / (init_step_size**2.0)
                    self._KKT_R[k - self.n - 1, k] = np.sqrt(0.5) / init_step_size**2.0
                    self._KKT_R[k - self.n - 1, k - self.n] = (
                        -self._KKT_R[k - self.n - 1, 0] - self._KKT_R[k - self.n - 1, k]
                    )
                elif k > 2 * self.n:
                    shift = (k - self.n - 1) // self.n
                    i = k - (1 + shift) * self.n - 1
                    j = (i + shift) % self.n
                    self._KKT_R[k - self.n - 1, 0] = 1.0 / init_step_size**2.0
                    self._KKT_R[k - self.n - 1, k] = 1.0 / init_step_size**2.0
                    self._KKT_R[k - self.n - 1, i + 1] = -1.0 / init_step_size**2.0
                    self._KKT_R[k - self.n - 1, j + 1] = -1.0 / init_step_size**2.0

            self._update_surrogate_coeff()

        def inv_kkt_matrix_dot(
            self, x: np.ndarray, w_star: Optional[np.ndarray] = None
        ) -> np.ndarray:
            """
            Compute inv(W) @ x, where W is the coefficient matrix of the KKT equation generated
            by the derivative-free symmetric Broyden update [1]_, and return the vector formed
            by the first ``npt`` and last ``n`` dimensions of the result.

            In the multiplication, we can safely ignore the ``(npt + 1)``-th dimension of ``x``.
            Therefore, the input vector ``x`` only needs to include the first ``npt`` and the last
            ``n`` elements, instead of the full ``(npt + 1 + n)``-dimensional vector.

            Parameters
            ----------
            x : ndarray, shape (npt + n,)
                The vector to be dotted with the inverse of the KKT matrix.
            w_star : ndarray, shape (npt + n,), optional
                The auxiliary vector to assist in the calculation. It is formed by

                    ``[0.5 * ((yi - center) @ (Y - center).T) ** 2, yi - center]``,

                where ``Y`` is of shape ``(npt, n)`` which represents the interpolation points.
                ``yi`` can be any point in the interpolation set, and is ``x_anchor`` by default.

            Returns
            -------
            ndarray, shape (npt + n,)
                The vector formed by the first ``npt`` and last ``n`` dimensions of the result.

            References
            ----------
            .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic
                models that satisfy interpolation conditions. *Math. Program.* 100, 1
                (May 2004), 183-215.
            """
            # assert x.size == self.npt + self.n
            w_star = self._get_w() if w_star is None else w_star

            x = x.reshape((-1,))
            Zx = np.zeros((self.npt + self.n,), dtype=np.float64)
            tmp = x[: self.npt] @ self._KKT_R.T
            tmp[: self._negative_s_idx] = -tmp[: self._negative_s_idx]
            Zx[: self.npt] = tmp @ self._KKT_R
            Zx[: self.npt] += x[self.npt :] @ (self._KKT_B[: self.npt]).T
            Zx[: self.npt] -= w_star[self.npt :] @ (self._KKT_B[: self.npt]).T
            tmp = w_star[: self.npt] @ self._KKT_R.T
            tmp[: self._negative_s_idx] = -tmp[: self._negative_s_idx]
            Zx[: self.npt] -= tmp @ self._KKT_R
            Zx[self.npt :] = (x - w_star) @ self._KKT_B
            Zx[self.cached_x_anchor_idx] += 1
            return Zx

        def inv_kkt_matrix_partial_dot(self, x: np.ndarray) -> np.ndarray:
            """
            Compute inv(W) @ x, where W is the coefficient matrix of the KKT equation
            generated by the derivative-free symmetric Broyden update [1]_, and return the
            vector formed by the first ``npt`` and last ``n`` dimensions of the result.

            Unlike ``self.inv_kkt_matrix_dot(x)``, this method assumes that ``x`` is zero for
            indices ``[npt:]``, allowing for a simplified calculation.

            Parameters
            ----------
            x : ndarray, shape (npt + n,)
                The vector to be dotted with the inverse of the KKT matrix.

            Returns
            -------
            ndarray, shape (npt + n,)
                The vector formed by the first ``npt`` and last ``n`` dimensions of the result.

            References
            ----------
            .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic
                models that satisfy interpolation conditions. *Math. Program.* 100, 1
                (May 2004), 183-215.
            """
            x = x.reshape((-1,))
            tmp = x[: self.npt] @ self._KKT_R.T
            tmp[: self._negative_s_idx] = -tmp[: self._negative_s_idx]
            return np.hstack((tmp @ self._KKT_R, x[: self.npt] @ self._KKT_B[: self.npt]))

        def _get_alpha(self, k: Optional[int] = None) -> Union[np.ndarray, float]:
            if k is None:
                return (
                    LA.norm(self._KKT_R[self._negative_s_idx :], axis=0) ** 2
                    - LA.norm(self._KKT_R[: self._negative_s_idx], axis=0) ** 2
                )
            else:
                tmp = np.copy(self._KKT_R[:, k])
                tmp[: self._negative_s_idx] = -tmp[: self._negative_s_idx]
                return tmp.dot(self._KKT_R[:, k])

        def _get_beta(
            self,
            w: np.ndarray,
            Zw: np.ndarray,
            x_shift: np.ndarray,
            w_star: Optional[np.ndarray] = None,
        ) -> float:
            w_star = self._get_w() if w_star is None else w_star
            return (
                0.5 * (x_shift.dot(x_shift)) ** 2
                - w[self.cached_x_anchor_idx]
                - w.dot(Zw)
                + w_star.dot(Zw)
            )

        def _get_w(self, x_shift: Optional[np.ndarray] = None) -> np.ndarray:
            x_shift = (
                self.cached_x_anchor - self.model_center
                if x_shift is None
                else x_shift.copy()
            )
            w = np.zeros(self.n + self.npt, dtype=np.float64)
            w[: self.npt] = 0.5 * ((x_shift @ self.cached_Y_shift.T) ** 2)
            w[self.npt :] = x_shift
            return w.copy()

        def get_determinant_ratio(
            self, x_new: np.ndarray, idx: Optional[int] = None
        ) -> Tuple[
            Union[np.ndarray, float],
            float,
            Union[np.ndarray, float],
            Union[np.ndarray, float],
            np.ndarray,
        ]:
            """
            Calculate key intermediate quantities (``alpha``, ``beta``, ``tau``, ``sigma``) for
            updating the inverse KKT matrix in the derivative-free symmetric Broyden
            update [1]_. These quantities are defined in [2]_, where ``sigma`` (the denominator
            in the update equation) represents the ratio of the determinant of the inverse
            KKT matrix before and after update.

            Parameters
            ----------
            x_new : ndarray, shape (n,)
                New point to be added to the interpolation set (not yet added).
            idx : int, optional
                Index of the interpolation point to be replaced by ``x_new``.
                If provided, returns scalar values at this index; otherwise returns full arrays.

            Returns
            -------
            alpha : ndarray or float
                Quantity ``alpha``. Returns ``alpha[idx]`` if ``idx`` is provided.
            beta : float
                Quantity ``beta``.
            tau : ndarray or float
                Quantity ``tau``. Returns ``tau[idx]`` if ``idx`` is provided.
            sigma : ndarray or float
                Quantity ``sigma``. Returns ``sigma[idx]`` if ``idx`` is provided.
            ndarray
                Result of ``self.inv_kkt_matrix_dot(w)``, where:
                ``w = [0.5 * ((x_new - center) @ (Y - center).T) ** 2, x_new - center]``,
                where ``Y`` is of shape ``(npt, n)`` which represents the interpolation points.

            References
            ----------
            .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic
                models that satisfy interpolation conditions. *Math. Program.* 100, 1
                (May 2004), 183-215.

            .. [2] Michael J. D. Powell. 2004. On updating the inverse of a KKT matrix.
                *Numerical Linear Algebra and Optimization, ed. Ya-xiang Yuan, Science Press
                (Beijing)* (2004): 56-78.
            """
            x_new_shift = x_new - self.model_center
            w = self._get_w(x_new_shift)
            w_star = self._get_w()

            Zw = self.inv_kkt_matrix_dot(w, w_star)
            beta = self._get_beta(w, Zw, x_new_shift, w_star=w_star)
            alpha = self._get_alpha(idx)
            tau = Zw[idx] if idx is not None else Zw[: self.npt]
            sigma = alpha * beta + tau**2
            return alpha, beta, tau, sigma, Zw

        def update(
            self,
            cached_kkt_info: Optional[
                Tuple[np.ndarray, float, np.ndarray, np.ndarray, np.ndarray]
            ] = None,
        ) -> None:
            """
            Update the KKT and model coefficients using the method detailed in [1]_, [2]_ and [3]_.

            Parameters
            ----------
            cached_kkt_info : tuple, optional
                If not None, it should be the cached output of ``self.get_determinant_ratio(x_new)``,
                where ``x_new`` denotes the new point just added into the interpolation set.

            References
            ----------
            .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic models
                that satisfy interpolation conditions. *Math. Program.* 100, 1 (May 2004), 183-215.

            .. [2] Michael J. D. Powell. 2004. On updating the inverse of a KKT matrix. *Numerical
                Linear Algebra and Optimization, ed. Ya-xiang Yuan, Science Press (Beijing)*
                (2004): 56-78.

            .. [3] Michael J. D. Powell. 2006. *The NEWUOA Software for Unconstrained Optimization
                without Derivatives*. Springer US, Boston, MA, 255-297.
            """
            last_update_idx = self.interp_set.last_update_idx
            x_new = self.interp_set.get_interp_set(last_update_idx)[0]
            deleted_x = self.interp_set.just_deleted_point.copy()
            self._update_KKT_coeff(x_new, last_update_idx, cached_kkt_info)
            self._update_surrogate_coeff(deleted_x, last_update_idx)
            self.cached_x_anchor = self.interp_set.get_anchor()[0].copy()
            self.cached_x_anchor_idx = self.interp_set.x_anchor_idx

        def _update_KKT_coeff(
            self,
            x_new: np.ndarray,
            idx: int,
            cached_kkt_info: Optional[
                Tuple[np.ndarray, float, np.ndarray, np.ndarray, np.ndarray]
            ] = None,
        ) -> None:
            """
            Update the KKT coefficients using the method detailed in [1]_, [2]_ and [3]_.

            Parameters
            ----------
            x_new : ndarray, shape (n,)
                The new point which was just added into the interpolation set.
            idx : int
                The index in the interpolation set where ``x_new`` was added.
            cached_kkt_info : tuple, optional
                If not None, it should be the cached output of ``self.get_determinant_ratio(x_new)``,
                where ``x_new`` denotes the new point just added into the interpolation set.

            Notes
            -----
            The implementation here largely references that in COBYQA [4]_.

            References
            ----------
            .. [1] Michael J. D. Powell. 2004. Least Frobenius norm updating of quadratic
                models that satisfy interpolation conditions. *Math. Program.* 100, 1
                (May 2004), 183-215.

            .. [2] Michael J. D. Powell. 2004. On updating the inverse of a KKT matrix.
                *Numerical Linear Algebra and Optimization, ed. Ya-xiang Yuan, Science Press
                (Beijing)* (2004): 56-78.

            .. [3] Michael J. D. Powell. 2006. *The NEWUOA Software for Unconstrained Optimization
                without Derivatives*. Springer US, Boston, MA, 255-297.

            .. [4] Tom M. Ragonneau. *Model-Based Derivative-Free Optimization Methods
                and Software*. PhD thesis, Department of Applied Mathematics, The Hong
                Kong Polytechnic University, Hong Kong, China, 2022. URL:
                https://theses.lib.polyu.edu.hk/handle/200/12294.
            """

            def householder_transform(mat: np.ndarray, idx: int) -> None:
                """
                Perform an in-place Householder transformation on ``mat`` to make the ``idx``-th
                column align with the first standard basis vector ``e_1``.
                """
                vec = mat[:, idx].reshape((-1,)).copy()
                vec_norm = LA.norm(vec)
                if np.abs(vec_norm) > 0.0:
                    vec[0] += vec_norm if vec[0] >= 0 else -vec_norm
                    mat -= 2 * np.outer(vec, vec @ mat) / vec.dot(vec)

            # Perform Householder transformation on _KKT_R
            if self._negative_s_idx > 1:
                householder_transform(self._KKT_R[: self._negative_s_idx], idx)
            if self._negative_s_idx < self._KKT_R.shape[0] - 1:
                householder_transform(self._KKT_R[self._negative_s_idx :], idx)

            # Prepare alpha, beta, tau, sigma, and Zw
            if cached_kkt_info is None:
                alpha, beta, tau, sigma, Zw = self.get_determinant_ratio(x_new, idx)
            else:
                alpha, beta, tau, sigma, Zw = cached_kkt_info
                alpha, tau, sigma = alpha[idx], tau[idx], sigma[idx]
            b_max = np.max(np.abs(self._KKT_B), initial=1.0)
            z_max = np.max(np.abs(self._KKT_R), initial=1.0)
            if abs(sigma) < np.finfo(float).tiny * max(b_max, z_max):
                # The denominator of the updating formula is too small to safely
                # divide the coefficients of the KKT matrix of interpolation.
                # Theoretically, the value of abs(sigma) is always positive, and
                # becomes small only for ill-conditioned problems.
                raise ZeroDivisionError("The denominator of the updating formula is zero")

            # Update _KKT_B first
            _KKT_B_on_idx = self._KKT_B[idx].copy()
            Ze_ell = np.hstack((self._KKT_R[0] * self._KKT_R[0, idx], _KKT_B_on_idx))
            Zw_minus_e_ell = Zw.copy()
            Zw_minus_e_ell[idx] -= 1
            temp_vec1 = (alpha * Zw[self.npt :] - tau * _KKT_B_on_idx) / sigma
            temp_vec2 = (tau * Zw[self.npt :] + beta * _KKT_B_on_idx) / sigma
            self._KKT_B[: self.npt] += np.outer(
                Zw_minus_e_ell[: self.npt], temp_vec1
            ) - np.outer(Ze_ell[: self.npt], temp_vec2)
            self._KKT_B[self.npt :] += np.outer(Zw_minus_e_ell[self.npt :], temp_vec1)
            self._KKT_B[self.npt :] -= np.outer(Ze_ell[self.npt :], temp_vec2)

            # Then update _KKT_R
            jdz = (
                self._negative_s_idx if self._negative_s_idx < self.npt - self.n - 1 else 0
            )
            if jdz == 0:
                self._KKT_R[0] = (
                    tau * self._KKT_R[0] - self._KKT_R[0, idx] * Zw_minus_e_ell[: self.npt]
                ) / (np.abs(sigma) ** 0.5)
                if sigma < 0.0:
                    self._negative_s_idx = 1
            else:
                scala = self._KKT_R[0, idx] if jdz == 0 else -self._KKT_R[0, idx]
                scalb = 0.0 if jdz == 0 else self._KKT_R[jdz, idx]
                kdz = jdz if beta >= 0.0 else 0
                jdz -= kdz
                tempb = self._KKT_R[jdz, idx] * tau / sigma
                tmp1 = 1.0 / np.sqrt(abs(beta) * self._KKT_R[jdz, idx] ** 2.0 + tau**2.0)
                self._KKT_R[kdz] = (
                    tau * self._KKT_R[kdz]
                    - self._KKT_R[jdz, idx] * Zw_minus_e_ell[: self.npt]
                )
                self._KKT_R[kdz] *= tmp1
                self._KKT_R[jdz] -= (self._KKT_R[jdz, idx] * beta / sigma) * (
                    scala * self._KKT_R[0] + scalb * self._KKT_R[self._negative_s_idx]
                ) + tempb * Zw_minus_e_ell[: self.npt]
                self._KKT_R[jdz] *= tmp1 * np.sqrt(abs(sigma))
                if sigma <= 0.0:
                    if beta < 0.0:
                        # If sigma <= 0 and beta < 0, the positive-definite property is broken.
                        self._negative_s_idx += 1
                    else:
                        self._negative_s_idx -= 1
                        self._KKT_R[[0, self._negative_s_idx]] = self._KKT_R[
                            [self._negative_s_idx, 0]
                        ]

        def shift_center_to(self, x: np.ndarray) -> None:
            """
            Shift the model center to ``x``, then update the KKT and model coefficients.

            Parameters
            ----------
            x : ndarray, shape (n,)
                The new model center.

            Notes
            -----
            The computational cost is O(npt^3) with BLAS-3 operations.
            """
            x = x.copy()
            step = x - self.model_center

            # Prepare intermediate variables
            Phi_T = self.cached_Y_shift - 0.5 * step
            Psi_T = (step @ Phi_T.T).reshape((-1, 1)) * Phi_T + 0.25 * step.dot(step) * step
            R_T_Psi_T = self._KKT_R @ Psi_T
            Xi_Psi_T = self._KKT_B[: self.npt].T @ Psi_T
            S_R_T_Psi_T = np.copy(R_T_Psi_T)
            S_R_T_Psi_T[: self._negative_s_idx] = -S_R_T_Psi_T[: self._negative_s_idx]
            Omega_Psi_T = self._KKT_R.T @ S_R_T_Psi_T
            Psi_Omega_Psi_T = R_T_Psi_T.T @ S_R_T_Psi_T

            # Update _KKT_B
            # Note: _KKT_R remains invariant when shifting the model center.
            self._KKT_B[: self.npt] += Omega_Psi_T
            self._KKT_B[self.npt :] += Xi_Psi_T + Xi_Psi_T.T + Psi_Omega_Psi_T

            # Update surrogate coefficients
            center_shift = x - self.model_center
            Hess_center_shift = self.hess_operator(center_shift)
            self.model_cons = (
                self.model_cons
                + self.model_grad.dot(center_shift)
                + 0.5 * center_shift @ Hess_center_shift
            )
            self.model_grad = Hess_center_shift + self.model_grad
            self.model_center = x
            temp = (Phi_T.T * self.model_hess_implicit) @ np.tile(step, (self.npt, 1))
            self.model_hess_explicit += temp + temp.T
            self.cached_Y_shift = self.cached_Y_shift - step

        def _update_surrogate_coeff(
            self, deleted_x: Optional[np.ndarray] = None, idx: Optional[int] = None
        ) -> None:
            """
            Assuming the KKT coefficients have been updated, this function updates the model
            coefficients based on the updated KKT coefficients.

            Parameters
            ----------
            deleted_x : ndarray, shape (n,), optional
                The old interpolation point that was replaced with a new point at the index
                ``idx``. Both ``deleted_x`` and ``idx`` should be None if the model is being
                updated for the first time after being established.
            idx : int, optional
                The index where the old interpolation point ``deleted_x`` was replaced with a
                new point. Both ``deleted_x`` and ``idx`` should be None if the model is being
                updated for the first time after being established.
            """
            Y_new, f_Y_new = self.interp_set.get_interp_set()
            f_Y_shift = f_Y_new - self.fun_eval(Y_new)
            model_lambda, model_g = self.solve_interpolation_KKT(f_Y_shift=f_Y_shift)
            model_hess_explicit_p1 = self.model_hess_explicit.copy()
            model_hess_implicit_p1 = self.model_hess_implicit + model_lambda
            if idx is not None:
                # assert deleted_x is not None
                deleted_x = deleted_x - self.model_center
                model_hess_explicit_p1 += self.model_hess_implicit[idx] * np.outer(
                    deleted_x, deleted_x
                )
                model_hess_implicit_p1[idx] = model_lambda[idx]
            model_grad_p1 = self.grad_eval(self.model_center) + model_g
            self.model_hess_explicit = model_hess_explicit_p1
            self.model_hess_implicit = model_hess_implicit_p1
            self.model_grad = model_grad_p1
            self.cached_Y_shift = Y_new - self.model_center
            self.model_cons += f_Y_new[self.interp_set.x_anchor_idx] - self.fun_eval(
                self.interp_set.x_anchor
            )

        def solve_interpolation_KKT(
            self, f_Y_shift=np.ndarray
        ) -> Tuple[np.ndarray, np.ndarray]:
            """
            Given the interpolation set, solves the first-order KKT optimality conditions
            for the derivative-free symmetric Broyden update subproblem and returns the solution:
            the Lagrange multipliers and the gradient update term.

            Parameters
            ----------
            f_Y_shift : ndarray, shape (npt,)
                Represents ``fun(Y) - surrogate(Y)``, where ``Y`` denotes the interpolation points.

            Returns
            -------
            lambda : ndarray, shape (npt,)
                The Lagrange multipliers.
            g : ndarray, shape (n,)
                The gradient update term.
            """
            coeff = self.inv_kkt_matrix_partial_dot(f_Y_shift)
            return coeff[0 : self.npt], coeff[self.npt :]

        def reset_surrogate_coeff(self) -> None:
            """
            Reset all surrogate model coefficients to zero values.
            """
            self.model_hess_explicit = np.zeros((self.n, self.n), dtype=np.float64)
            self.model_hess_implicit = np.zeros(self.interp_set.npt, dtype=np.float64)
            self.model_grad = np.zeros(self.n, dtype=np.float64)
            self.model_cons = 0.0

        def reinit(self) -> None:
            """
            Reinitialize the surrogate model by first resetting and then updating the coefficients.
            """
            self.reset_surrogate_coeff()
            self._update_surrogate_coeff()

        def alt_grad_eval(self, x: np.ndarray) -> np.ndarray:
            """
            Evaluate gradient of a newly-instantiated surrogate model at point ``x``.

            Conceptually equivalent to:

            1. Creating a new surrogate model with current interpolation set
            2. Updating its coefficients from zero
            3. Returning its gradient at ``x``

            Avoids actual model reconstruction by leveraging KKT solution.

            Parameters
            ----------
            x : ndarray, shape (n,)
                Evaluation point for gradient calculation.

            Returns
            -------
            g : ndarray
                Gradient vector of the conceptual surrogate model at ``x``.
            """
            _, f_Y = self.interp_set.get_interp_set()
            model_lambda, model_g = self.solve_interpolation_KKT(f_Y_shift=f_Y)
            return (
                ((x - self.model_center) @ self.cached_Y_shift.T) * model_lambda
            ) @ self.cached_Y_shift + model_g

        def fun_eval(self, X: np.ndarray) -> Union[np.ndarray, float]:
            """
            Evaluate the surrogate model at given point(s).

            Parameters
            ----------
            x : ndarray, shape (n,) or ndarray, shape (n_samples, n)
                Evaluation point(s). Can be a single point (1D) or batch of points (2D).

            Returns
            -------
            fval : float or ndarray, shape (n_samples,)
                Model value(s). Scalar for single point input, array of shape (n_samples,)
                for multiple points.
            """
            X_shift = X - self.model_center
            if X_shift.ndim == 1:
                return (
                    X_shift.dot(0.5 * self.hess_operator(X_shift) + self.model_grad)
                    + self.model_cons
                )
            else:
                value = (
                    0.5 * np.diag(self.hess_operator(X_shift) @ X_shift.T)
                    + X_shift @ self.model_grad
                    + self.model_cons
                )
                return (
                    value.item()
                    if isinstance(value, np.ndarray) and value.size == 1
                    else value
                )

        def grad_eval(self, X: np.ndarray) -> np.ndarray:
            """
            Evaluate the gradient vector(s) of the surrogate model at given point(s).

            Parameters
            ----------
            x : ndarray, shape (n,) or ndarray, shape (n_samples, n)
                Evaluation point(s). Can be a single point (1D) or batch of points (2D).

            Returns
            -------
            g : ndarray, shape (n, ) or ndarray, shape (n_samples, n)
                Gradient vector(s). Shape (n,) for single point, (n_samples, n) for
                multiple points.
            """
            return self.hess_operator(X - self.model_center) + self.model_grad

        @property
        def model_hess(self) -> np.ndarray:
            """Full Hessian matrix of the surrogate model"""
            return (
                self.cached_Y_shift.T * self.model_hess_implicit
            ) @ self.cached_Y_shift + self.model_hess_explicit

        def hess_eval(self, *arg, **kwargs) -> np.ndarray:
            """
            Return the full Hessian matrix of the surrogate model.

            Returns
            -------
            H : ndarray, shape (n, n)
                Hessian matrix (same as ``model_hess`` property).
            """
            return self.model_hess

        def hess_operator(self, V: np.ndarray) -> np.ndarray:
            """
            Compute Hessian-vector product(s) for the surrogate model.

            Parameters
            ----------
            V : ndarray, shape (n, ) or ndarray, shape (n_samples, n)
                Input vector(s). Single vector (1D) or batch of vectors (2D).

            Returns
            -------
            HV : ndarray, shape (n, ) or ndarray, shape (n_samples, n)
                Hessian-vector product(s). Shape (n,) for single input vector,
                (n_samples, n) for multiple vectors.
            """
            if V.ndim == 1:
                V = V[np.newaxis, :]
            result = (
                (V @ self.cached_Y_shift.T) * self.model_hess_implicit
            ) @ self.cached_Y_shift + V @ self.model_hess_explicit
            return result if V.shape[0] > 1 else result[0]

    class OverallSurrogate:
        r"""
        Overall surrogate model formed by assembling multiple elemental surrogate models.

        Represent a full-system surrogate as a sum of component models, each defined on
        a subspace of the full design space.

        Attributes
        ----------
        n : int
            Problem dimension
        interp_set : :class:`~upoqa.utils.interp_set.OverallInterpSet`
            Interpolation set for the overall surrogate model
        proj_onto_ele : callable
            Projection function mapping full-space points to element subspaces:

            .. code-block:: python

                (full_point: ndarray, element_name: Any) -> element_point: ndarray

            For example, ``proj_onto_ele([1.0, 2.0], "f_1")`` yields ``[1.0,]``, where
            ``"f_1"`` denotes a function that depends only on the first variable of ``x``.
        coords : list of 1D arrays
            Variable indices for each element's subspace (parallel to ``ele_names``).
        ele_models : list
            List of elemental surrogate models
        xforms : list
            Transformations :math:`h_1, \ldots, h_q` applied to element function outputs
        weights : list
            Weights :math:`w_1, \ldots, w_q` for each element in the model
        extra_fun : list
            The white-box component :math:`f_0` of the objective function, which is a list of
            length 3 containing callables to evaluate the function value, gradient, and Hessian.
        params : :class:`~upoqa.utils.params.UPOQAParameterList`
            Parameter list used for the algorithm
        ele_names : list
            List containing the names of all elements in order
        model_center : ndarray, shape (n,)
            Center of the surrogate model
        """

        def __init__(
            self,
            n: int,
            interp_set: OverallInterpSet,
            proj_onto_ele: Callable[[np.ndarray, Any], np.ndarray],
            coords: List[Union[List, np.ndarray]],
            ele_models: List[QuadSurrogate],
            extra_fun: List[Callable[[np.ndarray], float]],
            ele_names: List[Any],
            params: UPOQAParameterList,
        ) -> None:
            self.n = n
            self._zero = np.zeros((n,))
            self.interp_set = interp_set
            self.ele_models = ele_models
            self.ele_names = ele_names
            self.ele_idxs = list(range(len(ele_names)))

            assert callable(proj_onto_ele)
            self.proj_onto_ele = proj_onto_ele

            self.coords = coords
            self.model_center = self.interp_set.x_opt

            self.extra_fun = extra_fun
            self.params = params
            self.xforms = None
            self.weights: Optional[np.ndarray] = None

        def __getitem__(self, ele_idx: int) -> QuadSurrogate:
            return self.ele_models[ele_idx]

        def set_weights(self, weights: np.ndarray) -> None:
            """
            Set the weights for the elements.

            Parameters
            ----------
            - weights : ndarray
                The weights to assign to each element.
            """
            self.weights = deepcopy(weights)

        def set_xforms(
            self,
            xforms: List[Optional[List[Callable[[np.ndarray], Union[np.ndarray, float]]]]],
        ) -> None:
            """
            Set transformation functions for the elements.

            Parameters
            ----------
            xforms : list
                Transformation functions to apply to each element.
            """
            self.xforms = deepcopy(xforms)

        def update_weights(
            self,
            weights: Union[
                Dict[Any, Union[int, float]], List[Union[int, float]], np.ndarray
            ],
        ) -> None:
            """
            Update the weights for the elements, handling both list and dictionary inputs.

            If ``weights`` is a dictionary, it maps element names to weight values.
            Missing elements retain their current weights.

            This method should only be called by
            :meth:`~upoqa.utils.manager.UPOQAManager.update_weights()`.

            Parameters
            ----------
            weights : list or dict or ndarray
                New weights for model elements. Can be a 1D array, or a list of weights
                for all elements or a dictionary mapping element names to weights.
            """
            if isinstance(weights, dict):
                weights_list = []
                for ele_idx in self.ele_idxs:
                    ele_name = self.ele_names[ele_idx]
                    if ele_name in weights:
                        weights_list.append(float(weights[ele_name]))
                    else:
                        weights_list.append(float(self.weights[ele_idx]))
                self.set_weights(weights_list)
            else:
                self.set_weights(weights)

        def update_xforms(
            self,
            xforms: List[List[Callable[[np.ndarray], Union[np.ndarray, float]]]],
        ) -> None:
            """
            Update transformation functions with input validation.
            This method should only be called by
            :meth:`~upoqa.utils.manager.UPOQAManager.update_xforms()`.

            Parameters
            ----------
            xforms : list
                List of updated transformation lists for all elements. each transformation
                list should contain exactly 3 callables: ``[function, gradient, hessian]``.
            """
            if isinstance(xforms, dict):
                xforms_list = []
                for ele_idx in self.ele_idxs:
                    ele_name = self.ele_names[ele_idx]
                    if ele_name in xforms:
                        if xforms[ele_name] is not None:
                            assert isinstance(
                                xforms[ele_name], list
                            ), f"xforms[{ele_name}] should be a list, but got {type(xforms[ele_name])}"
                            assert (
                                len(xforms[ele_name]) == 3
                            ), f"xforms[{ele_name}] should be a list of length 3, but got length {len(xforms[ele_name])}"
                            assert all(
                                [isinstance(op, Callable) for op in xforms[ele_name]]
                            ), f"xforms[{ele_name}] should contain only callable functions."
                    else:
                        xforms[ele_name] = self.xforms[ele_idx]
                    xforms_list.append(xforms[ele_name])
                self.set_xforms(xforms_list)
            else:
                self.set_xforms(xforms)

        def shift_center_to(self, x: np.ndarray) -> None:
            """
            Shift the model center to a new point ``x``.
            Note that this does not change the centers of the element models.

            Parameters
            ----------
            x : ndarray
                The new center point for the overall surrogate model.
            """
            self.model_center = x

        def update(
            self,
            want_update: List[bool],
            cached_KKT_info_eles: List[
                Optional[Tuple[np.ndarray, float, np.ndarray, np.ndarray, np.ndarray]]
            ] = None,
        ) -> None:
            """
            Update the element models based on the provided update flags and cached KKT
            information.

            Parameters
            ----------
            want_update : list of bool
                A list of flags indicating whether each element model should be updated.
            cached_KKT_info_eles : list of tuples, optional
                List of cached outputs of ``ele_surrogate.get_determinant_ratio(x_ele_new)``
                for each element model ``ele_surrogate``, where ``x_ele_new`` denotes the new
                point just added into the corresponding elemental interpolation set.
            """
            for ele_idx in self.ele_idxs:
                ele_surrogate = self.ele_models[ele_idx]
                if want_update[ele_idx]:
                    try:
                        ele_surrogate.update(
                            cached_kkt_info=(
                                cached_KKT_info_eles[ele_idx]
                                if cached_KKT_info_eles
                                else None
                            )
                        )
                    except (LA.LinAlgError, ZeroDivisionError) as e:
                        raise SurrogateLinAlgError(
                            str(e), ele_name=self.ele_names[ele_idx], ele_idx=ele_idx
                        )

        def fun_eval(self, x: np.ndarray) -> float:
            """
            Evaluate the overall surrogate model at given point.

            Parameters
            ----------
            x : ndarray, shape (n,)
                Evaluation point.

            Returns
            -------
            fval : float
                The model value of the overall surrogate model at the point ``x``.
            """
            value = float(self.extra_fun[0](x))
            for ele_idx in self.ele_idxs:
                ele_surrogate = self.ele_models[ele_idx]
                x_ele = self.proj_onto_ele(x, ele_idx)
                if self.xforms[ele_idx] is not None:
                    value_ele = self.weights[ele_idx] * self.xforms[ele_idx][0](
                        ele_surrogate.fun_eval(x_ele)
                    )
                else:
                    value_ele = self.weights[ele_idx] * ele_surrogate.fun_eval(x_ele)
                if self.params("debug.check_nan_fval") and np.any(np.isnan(value_ele)):
                    raise ValueError(
                        f"(element {self.ele_names[ele_idx]}) NaN encountered in the return value"
                        " of the surrogate function. "
                        "This may be caused by some numerical error or matrix singularity. "
                    )
                value += value_ele
            return value

        def grad_eval(self, x: np.ndarray) -> np.ndarray:
            """
            Evaluate the gradient vector of the overall surrogate model at given point.

            Parameters
            ----------
            x : ndarray, shape (n,)
                Evaluation point.

            Returns
            -------
            g : ndarray, shape (n, )
                The gradient vector of the overall surrogate model at the point ``x``.
            """
            grad = np.array(self.extra_fun[1](x), dtype=np.float64)
            for ele_idx in self.ele_idxs:
                ele_surrogate = self.ele_models[ele_idx]
                x_ele = self.proj_onto_ele(x, ele_idx)
                if self.xforms[ele_idx] is not None:
                    grad_ele = (
                        self.weights[ele_idx]
                        * self.xforms[ele_idx][1](ele_surrogate.fun_eval(x_ele))
                        * ele_surrogate.grad_eval(x_ele)
                    )
                else:
                    grad_ele = self.weights[ele_idx] * ele_surrogate.grad_eval(x_ele)
                if self.params("debug.check_nan_fval") and np.any(np.isnan(grad_ele)):
                    raise ValueError(
                        f"(element {self.ele_names[ele_idx]}) NaN encountered in the return value"
                        " of the surrogate gradient. "
                        "This may be caused by some numerical error or matrix singularity. "
                    )
                grad[self.coords[ele_idx]] += grad_ele
            return grad

        def hess_eval(self, x: np.ndarray) -> np.ndarray:
            """
            Evaluate the full Hessian matrix of the overall surrogate model at given point.

            Parameters
            ----------
            x : ndarray, shape (n,)
                Evaluation point.

            Returns
            -------
            H : ndarray, shape (n, n)
                The Hessian matrix of the overall surrogate model at the point ``x``.
            """
            hess = np.array(self.extra_fun[2](x), dtype=np.float64)
            for ele_idx in self.ele_idxs:
                ele_surrogate = self.ele_models[ele_idx]
                model_dim_idx = np.asarray(self.coords[ele_idx]).reshape((1, -1))
                x_ele = self.proj_onto_ele(x, ele_idx)
                if self.xforms[ele_idx] is not None:
                    ele_model_fval = ele_surrogate.fun_eval(x_ele)
                    ele_model_grad = ele_surrogate.grad_eval(x_ele)
                    xform_grad = self.xforms[ele_idx][1](ele_model_fval)
                    xform_hess = self.xforms[ele_idx][2](ele_model_fval)
                    hess_ele = self.weights[ele_idx] * (
                        xform_hess * np.outer(ele_model_grad, ele_model_grad)
                        + xform_grad * ele_surrogate.hess_eval(x_ele)
                    )
                else:
                    hess_ele = self.weights[ele_idx] * ele_surrogate.hess_eval(x_ele)
                if self.params("debug.check_nan_fval") and np.any(np.isnan(hess_ele)):
                    raise ValueError(
                        f"(element {self.ele_names[ele_idx]}) NaN encountered in the return value"
                        " of the surrogate hessian. "
                        "This may be caused by some numerical error or matrix singularity. "
                    )
                hess[model_dim_idx.T, model_dim_idx] += hess_ele
            return hess

        def hess_operator(self, x: np.ndarray, v: np.ndarray) -> np.ndarray:
            """
            Compute Hessian-vector product for the overall surrogate model at given point.

            Parameters
            ----------
            x : ndarray, shape (n,)
                The point at which to evaluate the Hessian-vector product.
            v : ndarray, shape (n, )
                The vector to multiply with the Hessian matrix.

            Returns
            -------
            Hv : ndarray, shape (n, )
                The result of the Hessian-vector product.
            """
            result = v @ np.array(self.extra_fun[2](x), dtype=np.float64)
            for ele_idx in self.ele_idxs:
                ele_surrogate = self.ele_models[ele_idx]
                v_ele = self.proj_onto_ele(v, ele_idx)
                if self.xforms[ele_idx] is not None:
                    x_ele = self.proj_onto_ele(x, ele_idx)
                    ele_model_fval = ele_surrogate.fun_eval(x_ele)
                    ele_model_grad = ele_surrogate.grad_eval(x_ele)
                    xform_grad = self.xforms[ele_idx][1](ele_model_fval)
                    xform_hess = self.xforms[ele_idx][2](ele_model_fval)
                    hess_operator_ele = self.weights[ele_idx] * (
                        xform_hess
                        * np.outer(ele_model_grad, v_ele.dot(ele_model_grad)).squeeze()
                        + xform_grad * ele_surrogate.hess_operator(v_ele)
                    )
                else:
                    hess_operator_ele = self.weights[ele_idx] * ele_surrogate.hess_operator(
                        v_ele
                    )
                if self.params("debug.check_nan_fval") and np.any(
                    np.isnan(hess_operator_ele)
                ):
                    raise ValueError(
                        f"(element {self.ele_names[ele_idx]}) NaN encountered in the return "
                        "value of the surrogate hessian operator. "
                        "This may be caused by some numerical error or matrix singularity. "
                    )
                result[self.coords[ele_idx]] += hess_operator_ele
            return result