Source code for pybop.costs._weighted_cost

from typing import Optional, Union

import numpy as np

from pybop import BaseCost




class WeightedCost(BaseCost):
    """
    A subclass for constructing a linear combination of cost functions as
    a single weighted cost function.

    Parameters
    ----------
    costs : pybop.BaseCost
        The individual PyBOP cost objects.
    weights : list[float]
        A list of values with which to weight the cost values.
    has_identical_problems : bool
        If True, the shared problem will be evaluated once and saved before the
        self.compute() method of each cost is called (default: False).
    has_separable_problem: bool
        This attribute must be set to False for WeightedCost objects. If the
        corresponding attribute of an individual cost is True, the problem is
        separable from the cost function and will be evaluated before the
        individual cost evaluation is called.
    """

    def __init__(self, *costs, weights: Optional[list[float]] = None):
        if not all(isinstance(cost, BaseCost) for cost in costs):
            raise TypeError("All costs must be instances of BaseCost.")


        self.costs = [cost for cost in costs]

        if len(set(type(cost.problem) for cost in self.costs)) > 1:
            raise TypeError("All problems must be of the same class type.")

        # Check if weights are provided
        if weights is not None:
            try:
                self.weights = np.asarray(weights, dtype=float)
            except ValueError:
                raise ValueError("Weights must be numeric values.") from None

            if self.weights.size != len(self.costs):
                raise ValueError("Number of weights must match number of costs.")
        else:
            self.weights = np.ones(len(self.costs))

        # Check if all costs depend on the same problem


        self._has_identical_problems = all(
            cost.has_separable_problem and cost.problem is self.costs[0].problem
            for cost in self.costs
        )


        if self._has_identical_problems:
            super().__init__(self.costs[0].problem)
        else:
            super().__init__()

        for cost in self.costs:
            self.join_parameters(cost.parameters)

        # Apply the minimising property from each cost
        for i, cost in enumerate(self.costs):
            self.weights[i] = self.weights[i] * (1 if cost.minimising else -1)
        if all(not cost.minimising for cost in self.costs):
            # If all costs are maximising, convert the weighted cost to maximising
            self.weights = -self.weights
            self.minimising = False

        # Weighted costs do not use this functionality


        self._has_separable_problem = False




    def compute(
        self,
        y: dict,
        dy: Optional[np.ndarray] = None,
    ) -> Union[float, tuple[float, np.ndarray]]:
        """
        Computes the cost function for the given predictions.

        Parameters
        ----------
        y : dict
            The dictionary of predictions with keys designating the signals for fitting.
        dy : np.ndarray, optional
            The corresponding gradient with respect to the parameters for each signal.

        Returns
        -------
        float
            The weighted cost value.
        """
        if self._has_identical_problems:
            inputs = self.problem.parameters.as_dict()
            if dy is not None:
                y, dy = self.problem.evaluateS1(inputs)
            else:
                y = self.problem.evaluate(inputs)

        e = np.empty_like(self.costs)
        de = np.empty((len(self.parameters), len(self.costs)))

        for i, cost in enumerate(self.costs):
            if self._has_identical_problems:
                y, dy = (y, dy)
            elif cost.has_separable_problem:
                inputs = cost.parameters.as_dict()
                if dy is not None:
                    y, dy = cost.problem.evaluateS1(inputs)
                else:
                    y = cost.problem.evaluate(inputs)

            if dy is not None:
                e[i], de[:, i] = cost.compute(y, dy=dy)
            else:
                e[i] = cost.compute(y)

        e = np.dot(e, self.weights)
        if dy is not None:
            de = np.dot(de, self.weights)
            return e, de

        return e


    @property


    def has_identical_problems(self):
        return self._has_identical_problems