Source code for decent_bench.costs._base._quadratic_cost

from __future__ import annotations

from functools import cached_property

import numpy as np
from numpy import float64
from numpy.typing import NDArray

import decent_bench.utils.interoperability as iop
from decent_bench.costs._base._cost import Cost
from decent_bench.utils.array import Array
from decent_bench.utils.types import SupportedDevices, SupportedFrameworks




class QuadraticCost(Cost):
    r"""
    Quadratic cost function.

    .. math:: f(\mathbf{x}) = \frac{1}{2} \mathbf{x}^T \mathbf{Ax} + \mathbf{b}^T \mathbf{x} + c
    """



    def __init__(
        self,
        A: Array,  # noqa: N803
        b: Array,
        c: float = 0,
    ):
        self.A: NDArray[float64] = iop.to_numpy(A)
        self.b: NDArray[float64] = iop.to_numpy(b)

        if self.A.ndim != 2:
            raise ValueError("Matrix A must be 2D")
        if self.A.shape[0] != self.A.shape[1]:
            raise ValueError("Matrix A must be square")
        if self.b.ndim != 1:
            raise ValueError("Vector b must be 1D")
        if self.A.shape[0] != self.b.shape[0]:
            raise ValueError(f"Dimension mismatch: A has shape {self.A.shape} but b has length {self.b.shape[0]}")

        self.A_sym = 0.5 * (self.A + self.A.T)
        self.c = c


    @property
    def shape(self) -> tuple[int, ...]:
        return self.b.shape

    @property
    def framework(self) -> SupportedFrameworks:
        return SupportedFrameworks.NUMPY

    @property
    def device(self) -> SupportedDevices:
        return SupportedDevices.CPU



    @cached_property
    def m_smooth(self) -> float:  # pyright: ignore[reportIncompatibleMethodOverride]
        r"""
        The cost function's smoothness constant.

        .. math::
            \max_{i} \left| \lambda_i \right|

        where :math:`\lambda_i` are the eigenvalues of :math:`\frac{1}{2} (\mathbf{A}+\mathbf{A}^T)`.

        For the general definition, see
        :attr:`Cost.m_smooth <decent_bench.costs.Cost.m_smooth>`.
        """
        eigs = np.linalg.eigvalsh(self.A_sym)
        return float(np.max(np.abs(eigs)))




    @cached_property
    def m_cvx(self) -> float:  # pyright: ignore[reportIncompatibleMethodOverride]
        r"""
        The cost function's convexity constant.

        .. math::
            \begin{array}{ll}
                \min_i \lambda_i, & \text{if } \min_i \lambda_i > 0, \\
                0, & \text{if } \min_i \lambda_i = 0, \\
                \text{NaN}, & \text{if } \min_i \lambda_i < 0
            \end{array}

        where :math:`\lambda_i` are the eigenvalues of :math:`\frac{1}{2} (\mathbf{A}+\mathbf{A}^T)`.

        For the general definition, see
        :attr:`Cost.m_cvx <decent_bench.costs.Cost.m_cvx>`.
        """
        eigs = np.linalg.eigvalsh(self.A_sym)
        l_min = float(np.min(eigs))
        tol = 1e-12
        if l_min > tol:
            return l_min
        if abs(l_min) <= tol:
            return 0
        return np.nan




    @iop.autodecorate_cost_method(Cost.function)
    def function(self, x: NDArray[float64]) -> float:
        r"""
        Evaluate function at x.

        .. math:: \frac{1}{2} \mathbf{x}^T \mathbf{Ax} + \mathbf{b}^T \mathbf{x} + c
        """
        return float(0.5 * x.dot(self.A.dot(x)) + self.b.dot(x) + self.c)




    @iop.autodecorate_cost_method(Cost.gradient)
    def gradient(self, x: NDArray[float64]) -> NDArray[float64]:
        r"""
        Gradient at x.

        .. math:: \frac{1}{2} (\mathbf{A}+\mathbf{A}^T)\mathbf{x} + \mathbf{b}
        """
        return self.A_sym @ x + self.b




    @iop.autodecorate_cost_method(Cost.hessian)
    def hessian(self, x: NDArray[float64]) -> NDArray[float64]:  # noqa: ARG002
        r"""
        Hessian at x.

        .. math:: \frac{1}{2} (\mathbf{A}+\mathbf{A}^T)
        """
        ret: NDArray[float64] = self.A_sym.copy()
        return ret




    @iop.autodecorate_cost_method(Cost.proximal)
    def proximal(self, x: NDArray[float64], penalty: float) -> NDArray[float64]:
        r"""
        Proximal at x.

        .. math::
            (\frac{\rho}{2} (\mathbf{A} + \mathbf{A}^T) + \mathbf{I})^{-1} (\mathbf{x} - \rho \mathbf{b})

        where :math:`\rho > 0` is the penalty.

        This is a closed form solution, see
        :meth:`Cost.proximal() <decent_bench.costs.Cost.proximal>`
        for the general proximal definition.
        """
        lhs = penalty * self.A_sym + np.eye(self.A.shape[1])
        rhs = x - self.b * penalty

        return np.asarray(np.linalg.solve(lhs, rhs), dtype=float64)




    def __add__(self, other: Cost) -> Cost:
        """Add another cost function."""
        self._validate_cost_operation(other)
        if isinstance(other, QuadraticCost):
            return QuadraticCost(
                A=iop.to_array(self.A + other.A, self.framework, self.device),
                b=iop.to_array(self.b + other.b, self.framework, self.device),
                c=self.c + other.c,
            )

        return super().__add__(other)


    def __sub__(self, other: Cost) -> Cost:
        """
        Subtract another cost function.

        Preserves :class:`QuadraticCost` when subtracting another quadratic cost.
        """
        self._validate_cost_operation(other)
        if isinstance(other, QuadraticCost):
            return self.__add__(
                QuadraticCost(
                    A=iop.to_array(-other.A, self.framework, self.device),
                    b=iop.to_array(-other.b, self.framework, self.device),
                    c=-other.c,
                )
            )
        return super().__sub__(other)