From 7d3051a3d597dfb4cafaf9030fe9029b8b51b589 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 16 May 2024 17:25:55 +0100
Subject: [PATCH 01/25] Add a WeightedCost

---
 examples/scripts/spm_weighted_cost.py | 66 ++++++++++++++++++
 pybop/__init__.py                     |  2 +-
 pybop/costs/_likelihoods.py           | 38 ++++++-----
 pybop/costs/base_cost.py              | 96 +++++++++++++++++++++++++++
 pybop/costs/design_costs.py           |  2 +
 pybop/costs/fitting_costs.py          | 53 ++++++++++-----
 6 files changed, 223 insertions(+), 34 deletions(-)
 create mode 100644 examples/scripts/spm_weighted_cost.py

diff --git a/examples/scripts/spm_weighted_cost.py b/examples/scripts/spm_weighted_cost.py
new file mode 100644
index 000000000..4bd8380e9
--- /dev/null
+++ b/examples/scripts/spm_weighted_cost.py
@@ -0,0 +1,66 @@
+import numpy as np
+
+import pybop
+
+# Parameter set and model definition
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+model = pybop.lithium_ion.SPM(parameter_set=parameter_set)
+
+# Fitting parameters
+parameters = [
+    pybop.Parameter(
+        "Negative electrode active material volume fraction",
+        prior=pybop.Gaussian(0.68, 0.05),
+        bounds=[0.5, 0.8],
+    ),
+    pybop.Parameter(
+        "Positive electrode active material volume fraction",
+        prior=pybop.Gaussian(0.58, 0.05),
+        bounds=[0.4, 0.7],
+    ),
+]
+
+# Generate data
+sigma = 0.001
+t_eval = np.arange(0, 900, 3)
+values = model.predict(t_eval=t_eval)
+corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
+
+# Form dataset
+dataset = pybop.Dataset(
+    {
+        "Time [s]": t_eval,
+        "Current function [A]": values["Current [A]"].data,
+        "Voltage [V]": corrupt_values,
+    }
+)
+
+# Generate problem, cost function, and optimisation class
+problem = pybop.FittingProblem(model, parameters, dataset)
+cost1 = pybop.SumSquaredError(problem)
+cost2 = pybop.RootMeanSquaredError(problem)
+weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 100])
+
+for cost in [weighted_cost, cost1, cost2]:
+    optim = pybop.Optimisation(cost, optimiser=pybop.IRPropMin)
+    optim.set_max_iterations(60)
+
+    # Run the optimisation
+    x, final_cost = optim.run()
+    print(
+        "True parameters:",
+        [
+            parameters[0].true_value,
+            parameters[1].true_value,
+        ],
+    )
+    print("Estimated parameters:", x)
+
+    # Plot the timeseries output
+    pybop.quick_plot(problem, parameter_values=x, title="Optimised Comparison")
+
+    # Plot convergence
+    pybop.plot_convergence(optim)
+
+    # Plot the cost landscape with optimisation path
+    pybop.plot2d(optim, steps=15)
diff --git a/pybop/__init__.py b/pybop/__init__.py
index 034dcb4a4..6419f8f42 100644
--- a/pybop/__init__.py
+++ b/pybop/__init__.py
@@ -55,7 +55,7 @@
 #
 # Cost function class
 #
-from .costs.base_cost import BaseCost
+from .costs.base_cost import BaseCost, WeightedCost
 from .costs.fitting_costs import (
     RootMeanSquaredError,
     SumSquaredError,
diff --git a/pybop/costs/_likelihoods.py b/pybop/costs/_likelihoods.py
index 91374cc07..848c34c2d 100644
--- a/pybop/costs/_likelihoods.py
+++ b/pybop/costs/_likelihoods.py
@@ -62,20 +62,22 @@ def __init__(self, problem, sigma):
 
     def _evaluate(self, x, grad=None):
         """
-        Calls the problem.evaluate method and calculates
-        the log-likelihood
+        Calculates the log-likelihood.
         """
-        y = self.problem.evaluate(x)
-
         for key in self.signal:
-            if len(y.get(key, [])) != len(self._target.get(key, [])):
+            if len(self._current_prediction.get(key, [])) != len(
+                self._target.get(key, [])
+            ):
                 return -np.float64(np.inf)  # prediction doesn't match target
 
         e = np.array(
             [
                 np.sum(
                     self._offset
-                    + self._multip * np.sum((self._target[signal] - y[signal]) ** 2)
+                    + self._multip
+                    * np.sum(
+                        (self._target[signal] - self._current_prediction[signal]) ** 2
+                    )
                 )
                 for signal in self.signal
             ]
@@ -88,20 +90,26 @@ def _evaluate(self, x, grad=None):
 
     def _evaluateS1(self, x, grad=None):
         """
-        Calls the problem.evaluateS1 method and calculates
-        the log-likelihood
+        Calculates the log-likelihood and sensitivities.
         """
-        y, dy = self.problem.evaluateS1(x)
-
         for key in self.signal:
-            if len(y.get(key, [])) != len(self._target.get(key, [])):
+            if len(self._current_prediction.get(key, [])) != len(
+                self._target.get(key, [])
+            ):
                 likelihood = np.float64(np.inf)
                 dl = self._dl * np.ones(self.n_parameters)
                 return -likelihood, -dl
 
-        r = np.array([self._target[signal] - y[signal] for signal in self.signal])
+        r = np.array(
+            [
+                self._target[signal] - self._current_prediction[signal]
+                for signal in self.signal
+            ]
+        )
         likelihood = self._evaluate(x)
-        dl = np.sum((self.sigma2 * np.sum((r * dy.T), axis=2)), axis=1)
+        dl = np.sum(
+            (self.sigma2 * np.sum((r * self._current_sensitivities.T), axis=2)), axis=1
+        )
         return likelihood, dl
 
 
@@ -119,6 +127,7 @@ def __init__(self, problem):
         super(GaussianLogLikelihood, self).__init__(problem)
         self._logpi = -0.5 * self.n_time_data * np.log(2 * np.pi)
         self._dl = np.ones(self._n_parameters + self.n_outputs)
+        self._fixed_problem = False  # keep problem evaluation within _evaluate
 
     def _evaluate(self, x, grad=None):
         """
@@ -161,8 +170,7 @@ def _evaluate(self, x, grad=None):
 
     def _evaluateS1(self, x, grad=None):
         """
-        Calls the problem.evaluateS1 method and calculates
-        the log-likelihood
+        Calculates the log-likelihood and sensitivities.
         """
         sigma = np.asarray(x[-self.n_outputs :])
 
diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index f3ac95170..28fdff7f7 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -39,6 +39,7 @@ def __init__(self, problem=None, sigma=None):
         self.bounds = None
         self.sigma0 = sigma
         self._minimising = True
+        self._fixed_problem = True
         if isinstance(self.problem, BaseProblem):
             self._target = problem._target
             self.parameters = problem.parameters
@@ -82,6 +83,9 @@ def evaluate(self, x, grad=None):
             If an error occurs during the calculation of the cost.
         """
         try:
+            if self._fixed_problem:
+                self._current_prediction = self.problem.evaluate(x)
+
             if self._minimising:
                 return self._evaluate(x, grad)
             else:  # minimise the negative cost
@@ -140,6 +144,11 @@ def evaluateS1(self, x):
             If an error occurs during the calculation of the cost or gradient.
         """
         try:
+            if self._fixed_problem:
+                self._current_prediction, self._current_sensitivities = (
+                    self.problem.evaluateS1(x)
+                )
+
             if self._minimising:
                 return self._evaluateS1(x)
             else:  # minimise the negative cost
@@ -173,3 +182,90 @@ def _evaluateS1(self, x):
             If the method has not been implemented by the subclass.
         """
         raise NotImplementedError
+
+
+class WeightedCost(BaseCost):
+    """
+    A subclass for constructing a linear combination of cost functions as
+    a single weighted cost function.
+
+    Inherits all parameters and attributes from ``BaseCost``.
+    """
+
+    def __init__(self, cost_list, weights=None):
+        self.cost_list = cost_list
+        self.weights = weights
+        self._different_problems = False
+
+        if not isinstance(self.cost_list, list):
+            raise TypeError("Expected a list of costs.")
+        if self.weights is None:
+            self.weights = np.ones(len(cost_list))
+        elif isinstance(self.weights, list):
+            self.weights = np.array(self.weights)
+        if not isinstance(self.weights, np.ndarray):
+            raise TypeError(
+                "Expected a list or array of weights the same length as cost_list."
+            )
+        if not len(self.weights) == len(self.cost_list):
+            raise ValueError(
+                "Expected a list or array of weights the same length as cost_list."
+            )
+
+        # Check if all costs depend on the same problem
+        for cost in self.cost_list:
+            if hasattr(cost, "problem") and (
+                not cost._fixed_problem or cost.problem is not self.cost_list[0].problem
+            ):
+                self._different_problems = True
+
+        if not self._different_problems:
+            super(WeightedCost, self).__init__(self.cost_list[0].problem)
+            self._fixed_problem = False
+        else:
+            super(WeightedCost, self).__init__()
+
+    def _evaluate(self, x, grad=None):
+        """
+        Calculate the weighted cost for a given set of parameters.
+        """
+        e = np.empty_like(self.cost_list)
+
+        if not self._different_problems:
+            current_prediction = self.problem.evaluate(x)
+
+        for i, cost in enumerate(self.cost_list):
+            if self._different_problems:
+                cost._current_prediction = cost.problem.evaluate(x)
+            else:
+                cost._current_prediction = current_prediction
+            e[i] = cost._evaluate(x, grad)
+
+        return np.dot(e, self.weights)
+
+    def _evaluateS1(self, x):
+        """
+        Compute the cost and its gradient with respect to the parameters.
+        """
+        e = np.empty_like(self.cost_list)
+        de = np.empty((len(self.parameters), len(self.cost_list)))
+
+        if not self._different_problems:
+            current_prediction, current_sensitivities = self.problem.evaluateS1(x)
+
+        for i, cost in enumerate(self.cost_list):
+            if self._different_problems:
+                cost._current_prediction, cost._current_sensitivities = (
+                    cost.problem.evaluateS1(x)
+                )
+            else:
+                cost._current_prediction, cost._current_sensitivities = (
+                    current_prediction,
+                    current_sensitivities,
+                )
+            e[i], de[:, i] = cost._evaluateS1(x)
+
+        e = np.dot(e, self.weights)
+        de = np.dot(de, self.weights)
+
+        return e, de
diff --git a/pybop/costs/design_costs.py b/pybop/costs/design_costs.py
index b02940bf0..aff621910 100644
--- a/pybop/costs/design_costs.py
+++ b/pybop/costs/design_costs.py
@@ -98,6 +98,7 @@ class GravimetricEnergyDensity(DesignCost):
 
     def __init__(self, problem, update_capacity=False):
         super(GravimetricEnergyDensity, self).__init__(problem, update_capacity)
+        self._fixed_problem = False  # keep problem evaluation within _evaluate
 
     def _evaluate(self, x, grad=None):
         """
@@ -157,6 +158,7 @@ class VolumetricEnergyDensity(DesignCost):
 
     def __init__(self, problem, update_capacity=False):
         super(VolumetricEnergyDensity, self).__init__(problem, update_capacity)
+        self._fixed_problem = False  # keep problem evaluation within _evaluate
 
     def _evaluate(self, x, grad=None):
         """
diff --git a/pybop/costs/fitting_costs.py b/pybop/costs/fitting_costs.py
index 0665454b0..24c78fff9 100644
--- a/pybop/costs/fitting_costs.py
+++ b/pybop/costs/fitting_costs.py
@@ -41,15 +41,19 @@ def _evaluate(self, x, grad=None):
             The root mean square error.
 
         """
-        prediction = self.problem.evaluate(x)
-
         for key in self.signal:
-            if len(prediction.get(key, [])) != len(self._target.get(key, [])):
+            if len(self._current_prediction.get(key, [])) != len(
+                self._target.get(key, [])
+            ):
                 return np.float64(np.inf)  # prediction doesn't match target
 
         e = np.array(
             [
-                np.sqrt(np.mean((prediction[signal] - self._target[signal]) ** 2))
+                np.sqrt(
+                    np.mean(
+                        (self._current_prediction[signal] - self._target[signal]) ** 2
+                    )
+                )
                 for signal in self.signal
             ]
         )
@@ -79,18 +83,24 @@ def _evaluateS1(self, x):
         ValueError
             If an error occurs during the calculation of the cost or gradient.
         """
-        y, dy = self.problem.evaluateS1(x)
-
         for key in self.signal:
-            if len(y.get(key, [])) != len(self._target.get(key, [])):
+            if len(self._current_prediction.get(key, [])) != len(
+                self._target.get(key, [])
+            ):
                 e = np.float64(np.inf)
                 de = self._de * np.ones(self.n_parameters)
                 return e, de
 
-        r = np.array([y[signal] - self._target[signal] for signal in self.signal])
+        r = np.array(
+            [
+                self._current_prediction[signal] - self._target[signal]
+                for signal in self.signal
+            ]
+        )
         e = np.sqrt(np.mean(r**2, axis=1))
-        de = np.mean((r * dy.T), axis=2) / (
-            np.sqrt(np.mean((r * dy.T) ** 2, axis=2)) + np.finfo(float).eps
+        de = np.mean((r * self._current_sensitivities.T), axis=2) / (
+            np.sqrt(np.mean((r * self._current_sensitivities.T) ** 2, axis=2))
+            + np.finfo(float).eps
         )
 
         if self.n_outputs == 1:
@@ -155,15 +165,15 @@ def _evaluate(self, x, grad=None):
         float
             The sum of squared errors.
         """
-        prediction = self.problem.evaluate(x)
-
         for key in self.signal:
-            if len(prediction.get(key, [])) != len(self._target.get(key, [])):
+            if len(self._current_prediction.get(key, [])) != len(
+                self._target.get(key, [])
+            ):
                 return np.float64(np.inf)  # prediction doesn't match target
 
         e = np.array(
             [
-                np.sum(((prediction[signal] - self._target[signal]) ** 2))
+                np.sum(((self._current_prediction[signal] - self._target[signal]) ** 2))
                 for signal in self.signal
             ]
         )
@@ -192,16 +202,22 @@ def _evaluateS1(self, x):
         ValueError
             If an error occurs during the calculation of the cost or gradient.
         """
-        y, dy = self.problem.evaluateS1(x)
         for key in self.signal:
-            if len(y.get(key, [])) != len(self._target.get(key, [])):
+            if len(self._current_prediction.get(key, [])) != len(
+                self._target.get(key, [])
+            ):
                 e = np.float64(np.inf)
                 de = self._de * np.ones(self.n_parameters)
                 return e, de
 
-        r = np.array([y[signal] - self._target[signal] for signal in self.signal])
+        r = np.array(
+            [
+                self._current_prediction[signal] - self._target[signal]
+                for signal in self.signal
+            ]
+        )
         e = np.sum(np.sum(r**2, axis=0), axis=0)
-        de = 2 * np.sum(np.sum((r * dy.T), axis=2), axis=1)
+        de = 2 * np.sum(np.sum((r * self._current_sensitivities.T), axis=2), axis=1)
 
         return e, de
 
@@ -235,6 +251,7 @@ class ObserverCost(BaseCost):
     def __init__(self, observer: Observer):
         super().__init__(problem=observer)
         self._observer = observer
+        self._fixed_problem = False  # keep problem evaluation within _evaluate
 
     def _evaluate(self, x, grad=None):
         """

From 7898e8868ae55e2a853728998fd3c1465b7c52e6 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 16 May 2024 19:06:48 +0100
Subject: [PATCH 02/25] Fix setting

---
 pybop/costs/base_cost.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 28fdff7f7..b6fa4885d 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -221,9 +221,9 @@ def __init__(self, cost_list, weights=None):
 
         if not self._different_problems:
             super(WeightedCost, self).__init__(self.cost_list[0].problem)
-            self._fixed_problem = False
         else:
             super(WeightedCost, self).__init__()
+            self._fixed_problem = False
 
     def _evaluate(self, x, grad=None):
         """

From 7a9bcc4f7d59be349c12630826bdaed56b62133f Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 16 May 2024 19:08:14 +0100
Subject: [PATCH 03/25] Add tests

---
 tests/unit/test_cost.py | 50 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 6ffeb58ea..3ed5245e6 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -1,3 +1,5 @@
+from copy import copy
+
 import numpy as np
 import pytest
 
@@ -227,3 +229,51 @@ def test_energy_density_costs(
         # Compute after updating nominal capacity
         cost = cost_class(problem, update_capacity=True)
         cost([0.4])
+
+    @pytest.mark.unit
+    def test_weighted_cost(self, problem, x0):
+        cost1 = pybop.SumSquaredError(problem)
+        cost2 = pybop.RootMeanSquaredError(problem)
+
+        # Test with and without weights
+        weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2])
+        np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
+        weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 1])
+        np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
+        weighted_cost = pybop.WeightedCost(
+            cost_list=[cost1, cost2], weights=np.array([1, 1])
+        )
+        np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
+        with pytest.raises(
+            TypeError,
+            match="Expected a list or array of weights the same length as cost_list.",
+        ):
+            weighted_cost = pybop.WeightedCost(
+                cost_list=[cost1, cost2], weights="Invalid string"
+            )
+        with pytest.raises(
+            ValueError,
+            match="Expected a list or array of weights the same length as cost_list.",
+        ):
+            weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1])
+
+        # Test with and without different problems
+        weighted_cost_2 = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 100])
+        assert weighted_cost_2._different_problems is False
+        assert weighted_cost_2.problem is problem
+        assert weighted_cost_2(x0) >= 0
+
+        cost3 = pybop.RootMeanSquaredError(copy(problem))
+        weighted_cost_3 = pybop.WeightedCost(cost_list=[cost1, cost3], weights=[1, 100])
+        assert weighted_cost_3._different_problems is True
+        assert weighted_cost_3.problem is None
+        assert weighted_cost_3(x0) >= 0
+
+        np.testing.assert_allclose(
+            weighted_cost_2._evaluate(x0), weighted_cost_3._evaluate(x0), atol=1e-5
+        )
+        weighted_cost_3.parameters = problem.parameters
+        errors_2, sensitivities_2 = weighted_cost_2._evaluateS1(x0)
+        errors_3, sensitivities_3 = weighted_cost_3._evaluateS1(x0)
+        np.testing.assert_allclose(errors_2, errors_3, atol=1e-5)
+        np.testing.assert_allclose(sensitivities_2, sensitivities_3, atol=1e-5)

From a1c2ec6c15754e06a962864b234bcd83184fa399 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 23 May 2024 17:26:26 +0100
Subject: [PATCH 04/25] Update base_cost.py

---
 pybop/costs/base_cost.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 9c2ec6926..a734ccbd5 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -38,7 +38,7 @@ def __init__(self, problem=None, sigma=None):
         self.x0 = None
         self.bounds = None
         self.sigma0 = sigma
-        self._fixed_problem = True
+        self._fixed_problem = False
         if isinstance(self.problem, BaseProblem):
             self._target = problem._target
             self.parameters = problem.parameters
@@ -48,6 +48,7 @@ def __init__(self, problem=None, sigma=None):
             self.signal = problem.signal
             self._n_parameters = problem.n_parameters
             self.sigma0 = sigma or problem.sigma0 or np.zeros(self._n_parameters)
+            self._fixed_problem = True
 
     @property
     def n_parameters(self):

From 501064d37bd5b7c2a4c3444d78277707cf82a41d Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 24 May 2024 17:10:01 +0100
Subject: [PATCH 05/25] Update CHANGELOG.md

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9e26270dc..27e3b8ce5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#327](https://github.com/pybop-team/PyBOP/issues/327) - Adds the `WeightedCost` subclass, defines when to evaluate a problem and adds the `spm_weighted_cost` example script.
 - [#236](https://github.com/pybop-team/PyBOP/issues/236) - Restructures the optimiser classes, adds a new optimisation API through direct construction and keyword arguments, and fixes the setting of `max_iterations`, and `_minimising`. Introduces `pybop.BaseOptimiser`, `pybop.BasePintsOptimiser`, and `pybop.BaseSciPyOptimiser` classes.
 - [#321](https://github.com/pybop-team/PyBOP/pull/321) - Updates Prior classes with BaseClass, adds a `problem.sample_initial_conditions` method to improve stability of SciPy.Minimize optimiser.
 - [#249](https://github.com/pybop-team/PyBOP/pull/249) - Add WeppnerHuggins model and GITT example.

From 5fa3db4e84b65641edc48cb1bdbc582331b106d5 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Mon, 10 Jun 2024 10:41:30 +0100
Subject: [PATCH 06/25] Update imports

---
 pybop/costs/base_cost.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index e6c2f5408..558a96359 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from pybop import BaseProblem
 
 

From 23b10992a87894b3d99bd1078d6280d2e707807c Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Mon, 10 Jun 2024 11:19:42 +0100
Subject: [PATCH 07/25] Update x0 to [0.5]

---
 tests/unit/test_cost.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index fc643ef44..3f230fdab 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -234,7 +234,7 @@ def test_design_costs(
             cost([0.4])
 
     @pytest.mark.unit
-    def test_weighted_cost(self, problem, x0):
+    def test_weighted_cost(self, problem):
         cost1 = pybop.SumSquaredError(problem)
         cost2 = pybop.RootMeanSquaredError(problem)
 
@@ -264,19 +264,19 @@ def test_weighted_cost(self, problem, x0):
         weighted_cost_2 = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 100])
         assert weighted_cost_2._different_problems is False
         assert weighted_cost_2.problem is problem
-        assert weighted_cost_2(x0) >= 0
+        assert weighted_cost_2([0.5]) >= 0
 
         cost3 = pybop.RootMeanSquaredError(copy(problem))
         weighted_cost_3 = pybop.WeightedCost(cost_list=[cost1, cost3], weights=[1, 100])
         assert weighted_cost_3._different_problems is True
         assert weighted_cost_3.problem is None
-        assert weighted_cost_3(x0) >= 0
+        assert weighted_cost_3([0.5]) >= 0
 
         np.testing.assert_allclose(
-            weighted_cost_2._evaluate(x0), weighted_cost_3._evaluate(x0), atol=1e-5
+            weighted_cost_2._evaluate([0.5]), weighted_cost_3._evaluate([0.5]), atol=1e-5
         )
         weighted_cost_3.parameters = problem.parameters
-        errors_2, sensitivities_2 = weighted_cost_2._evaluateS1(x0)
-        errors_3, sensitivities_3 = weighted_cost_3._evaluateS1(x0)
+        errors_2, sensitivities_2 = weighted_cost_2._evaluateS1([0.5])
+        errors_3, sensitivities_3 = weighted_cost_3._evaluateS1([0.5])
         np.testing.assert_allclose(errors_2, errors_3, atol=1e-5)
         np.testing.assert_allclose(sensitivities_2, sensitivities_3, atol=1e-5)

From d6708d6e90a1e07998e257849c5ff481eb208b20 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 10 Jun 2024 10:20:01 +0000
Subject: [PATCH 08/25] style: pre-commit fixes

---
 tests/unit/test_cost.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 3f230fdab..f1e1fc4b9 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -273,7 +273,9 @@ def test_weighted_cost(self, problem):
         assert weighted_cost_3([0.5]) >= 0
 
         np.testing.assert_allclose(
-            weighted_cost_2._evaluate([0.5]), weighted_cost_3._evaluate([0.5]), atol=1e-5
+            weighted_cost_2._evaluate([0.5]),
+            weighted_cost_3._evaluate([0.5]),
+            atol=1e-5,
         )
         weighted_cost_3.parameters = problem.parameters
         errors_2, sensitivities_2 = weighted_cost_2._evaluateS1([0.5])

From 37ac6e2e0898b14470c020c30a89e7574c65c54e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 10 Jun 2024 19:03:53 +0000
Subject: [PATCH 09/25] style: pre-commit fixes

---
 pybop/costs/fitting_costs.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pybop/costs/fitting_costs.py b/pybop/costs/fitting_costs.py
index 069875c21..c63784e60 100644
--- a/pybop/costs/fitting_costs.py
+++ b/pybop/costs/fitting_costs.py
@@ -98,7 +98,9 @@ def _evaluateS1(self, x):
             ]
         )
         e = np.sqrt(np.mean(r**2, axis=1))
-        de = np.mean((r * self._current_sensitivities.T), axis=2) / (e + np.finfo(float).eps)
+        de = np.mean((r * self._current_sensitivities.T), axis=2) / (
+            e + np.finfo(float).eps
+        )
 
         if self.n_outputs == 1:
             return e.item(), de.flatten()

From 1c540d95514c5469e2c1bb5c5098461259ca4c87 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 11 Jun 2024 14:42:14 +0100
Subject: [PATCH 10/25] Update spm_weighted_cost.py

---
 examples/scripts/spm_weighted_cost.py | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/examples/scripts/spm_weighted_cost.py b/examples/scripts/spm_weighted_cost.py
index 4bd8380e9..14ba213d8 100644
--- a/examples/scripts/spm_weighted_cost.py
+++ b/examples/scripts/spm_weighted_cost.py
@@ -7,18 +7,20 @@
 model = pybop.lithium_ion.SPM(parameter_set=parameter_set)
 
 # Fitting parameters
-parameters = [
+parameters = pybop.Parameters(
     pybop.Parameter(
         "Negative electrode active material volume fraction",
         prior=pybop.Gaussian(0.68, 0.05),
         bounds=[0.5, 0.8],
+        true_value=parameter_set["Negative electrode active material volume fraction"],
     ),
     pybop.Parameter(
         "Positive electrode active material volume fraction",
         prior=pybop.Gaussian(0.58, 0.05),
         bounds=[0.4, 0.7],
+        true_value=parameter_set["Positive electrode active material volume fraction"],
     ),
-]
+)
 
 # Generate data
 sigma = 0.001
@@ -42,18 +44,11 @@
 weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 100])
 
 for cost in [weighted_cost, cost1, cost2]:
-    optim = pybop.Optimisation(cost, optimiser=pybop.IRPropMin)
-    optim.set_max_iterations(60)
+    optim = pybop.IRPropMin(cost, max_iterations=60)
 
     # Run the optimisation
     x, final_cost = optim.run()
-    print(
-        "True parameters:",
-        [
-            parameters[0].true_value,
-            parameters[1].true_value,
-        ],
-    )
+    print("True parameters:", parameters.true_value())
     print("Estimated parameters:", x)
 
     # Plot the timeseries output

From 3a9e10a08700b3fe4f3778a4b8d15b80443204a9 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 11 Jun 2024 15:21:58 +0100
Subject: [PATCH 11/25] Update TypeError with test

---
 pybop/costs/base_cost.py | 4 +++-
 tests/unit/test_cost.py  | 5 +++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 558a96359..c385a51f1 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -179,7 +179,9 @@ def __init__(self, cost_list, weights=None):
         self._different_problems = False
 
         if not isinstance(self.cost_list, list):
-            raise TypeError("Expected a list of costs.")
+            raise TypeError(
+                f"Expected a list of costs. Received {type(self.cost_list)}"
+            )
         if self.weights is None:
             self.weights = np.ones(len(cost_list))
         elif isinstance(self.weights, list):
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index f1e1fc4b9..972862552 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -247,6 +247,11 @@ def test_weighted_cost(self, problem):
             cost_list=[cost1, cost2], weights=np.array([1, 1])
         )
         np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
+        with pytest.raises(
+            TypeError,
+            match=r"Expected a list of costs.",
+        ):
+            weighted_cost = pybop.WeightedCost(cost_list="Invalid string")
         with pytest.raises(
             TypeError,
             match="Expected a list or array of weights the same length as cost_list.",

From be1c5662ec9d58d314ddcbfe21738d6ff654dc82 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 4 Jul 2024 21:42:35 +0100
Subject: [PATCH 12/25] Update spm_weighted_cost.py

---
 examples/scripts/spm_weighted_cost.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/scripts/spm_weighted_cost.py b/examples/scripts/spm_weighted_cost.py
index 14ba213d8..557434613 100644
--- a/examples/scripts/spm_weighted_cost.py
+++ b/examples/scripts/spm_weighted_cost.py
@@ -52,7 +52,7 @@
     print("Estimated parameters:", x)
 
     # Plot the timeseries output
-    pybop.quick_plot(problem, parameter_values=x, title="Optimised Comparison")
+    pybop.quick_plot(problem, problem_inputs=x, title="Optimised Comparison")
 
     # Plot convergence
     pybop.plot_convergence(optim)

From f608aa722b078ae8ae6ec3c930a96db2ab41a083 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 5 Jul 2024 11:06:05 +0100
Subject: [PATCH 13/25] Update evaluate and _evaluate

---
 pybop/costs/base_cost.py     | 66 +++++++++++++++++++++++++++---------
 pybop/costs/fitting_costs.py |  8 ++++-
 tests/unit/test_cost.py      | 27 ++++++++++-----
 3 files changed, 76 insertions(+), 25 deletions(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 4fb0fd010..7fae9ee9c 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -208,46 +208,80 @@ def __init__(self, cost_list, weights=None):
         else:
             super(WeightedCost, self).__init__()
             self._fixed_problem = False
+            for cost in self.cost_list:
+                self.parameters.join(cost.parameters)
 
-    def _evaluate(self, x, grad=None):
+    def _evaluate(self, inputs: Inputs, grad=None):
         """
         Calculate the weighted cost for a given set of parameters.
+
+        Parameters
+        ----------
+        inputs : Inputs
+            The parameters for which to compute the cost.
+        grad : array-like, optional
+            An array to store the gradient of the cost function with respect
+            to the parameters.
+
+        Returns
+        -------
+        float
+            The weighted cost value.
         """
         e = np.empty_like(self.cost_list)
 
-        if not self._different_problems:
-            current_prediction = self.problem.evaluate(x)
+        if not self._fixed_problem and self._different_problems:
+            self.parameters.update(values=list(inputs.values()))
+        elif not self._fixed_problem:
+            self._current_prediction = self.problem.evaluate(inputs)
 
         for i, cost in enumerate(self.cost_list):
-            if self._different_problems:
-                cost._current_prediction = cost.problem.evaluate(x)
+            if not self._fixed_problem and self._different_problems:
+                inputs = cost.parameters.as_dict()
+                cost._current_prediction = cost.problem.evaluate(inputs)
             else:
-                cost._current_prediction = current_prediction
-            e[i] = cost._evaluate(x, grad)
+                cost._current_prediction = self._current_prediction
+            e[i] = cost._evaluate(inputs, grad)
 
         return np.dot(e, self.weights)
 
-    def _evaluateS1(self, x):
+    def _evaluateS1(self, inputs: Inputs):
         """
-        Compute the cost and its gradient with respect to the parameters.
+        Compute the weighted cost and its gradient with respect to the parameters.
+
+        Parameters
+        ----------
+        inputs : Inputs
+            The parameters for which to compute the cost and gradient.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the cost and the gradient. The cost is a float,
+            and the gradient is an array-like of the same length as `x`.
         """
         e = np.empty_like(self.cost_list)
         de = np.empty((len(self.parameters), len(self.cost_list)))
 
-        if not self._different_problems:
-            current_prediction, current_sensitivities = self.problem.evaluateS1(x)
+        if not self._fixed_problem and self._different_problems:
+            self.parameters.update(values=list(inputs.values()))
+        elif not self._fixed_problem:
+            self._current_prediction, self._current_sensitivities = (
+                self.problem.evaluateS1(inputs)
+            )
 
         for i, cost in enumerate(self.cost_list):
-            if self._different_problems:
+            if not self._fixed_problem and self._different_problems:
+                inputs = cost.parameters.as_dict()
                 cost._current_prediction, cost._current_sensitivities = (
-                    cost.problem.evaluateS1(x)
+                    cost.problem.evaluateS1(inputs)
                 )
             else:
                 cost._current_prediction, cost._current_sensitivities = (
-                    current_prediction,
-                    current_sensitivities,
+                    self._current_prediction,
+                    self._current_sensitivities,
                 )
-            e[i], de[:, i] = cost._evaluateS1(x)
+            e[i], de[:, i] = cost._evaluateS1(inputs)
 
         e = np.dot(e, self.weights)
         de = np.dot(de, self.weights)
diff --git a/pybop/costs/fitting_costs.py b/pybop/costs/fitting_costs.py
index d907c06ff..08eef2390 100644
--- a/pybop/costs/fitting_costs.py
+++ b/pybop/costs/fitting_costs.py
@@ -275,7 +275,7 @@ def _evaluate(self, inputs: Inputs, grad=None):
         )
         return -log_likelihood
 
-    def evaluateS1(self, inputs: Inputs):
+    def _evaluateS1(self, inputs: Inputs):
         """
         Compute the cost and its gradient with respect to the parameters.
 
@@ -347,6 +347,9 @@ def _evaluate(self, inputs: Inputs, grad=None):
         float
             The maximum a posteriori cost.
         """
+        if self._fixed_problem:
+            self.likelihood._current_prediction = self._current_prediction
+
         log_likelihood = self.likelihood._evaluate(inputs)
         log_prior = sum(
             self.parameters[key].prior.logpdf(value) for key, value in inputs.items()
@@ -376,6 +379,9 @@ def _evaluateS1(self, inputs: Inputs):
         ValueError
             If an error occurs during the calculation of the cost or gradient.
         """
+        if self._fixed_problem:
+            self.likelihood._current_prediction = self._current_prediction
+
         log_likelihood, dl = self.likelihood._evaluateS1(inputs)
         log_prior = sum(
             self.parameters[key].prior.logpdf(inputs[key]) for key in inputs.keys()
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 990682eb9..cc45deed6 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -285,24 +285,35 @@ def test_weighted_cost(self, problem):
             weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1])
 
         # Test with and without different problems
-        weighted_cost_2 = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 100])
+        weight = 100
+        weighted_cost_2 = pybop.WeightedCost(
+            cost_list=[cost1, cost2], weights=[1, weight]
+        )
         assert weighted_cost_2._different_problems is False
+        assert weighted_cost_2._fixed_problem is True
         assert weighted_cost_2.problem is problem
         assert weighted_cost_2([0.5]) >= 0
+        np.testing.assert_allclose(
+            weighted_cost_2.evaluate([0.6]),
+            cost1([0.6]) + weight * cost2([0.6]),
+            atol=1e-5,
+        )
 
         cost3 = pybop.RootMeanSquaredError(copy(problem))
-        weighted_cost_3 = pybop.WeightedCost(cost_list=[cost1, cost3], weights=[1, 100])
+        weighted_cost_3 = pybop.WeightedCost(
+            cost_list=[cost1, cost3], weights=[1, weight]
+        )
         assert weighted_cost_3._different_problems is True
+        assert weighted_cost_3._fixed_problem is False
         assert weighted_cost_3.problem is None
         assert weighted_cost_3([0.5]) >= 0
-
         np.testing.assert_allclose(
-            weighted_cost_2._evaluate([0.5]),
-            weighted_cost_3._evaluate([0.5]),
+            weighted_cost_3.evaluate([0.6]),
+            cost1([0.6]) + weight * cost3([0.6]),
             atol=1e-5,
         )
-        weighted_cost_3.parameters = problem.parameters
-        errors_2, sensitivities_2 = weighted_cost_2._evaluateS1([0.5])
-        errors_3, sensitivities_3 = weighted_cost_3._evaluateS1([0.5])
+
+        errors_2, sensitivities_2 = weighted_cost_2.evaluateS1([0.5])
+        errors_3, sensitivities_3 = weighted_cost_3.evaluateS1([0.5])
         np.testing.assert_allclose(errors_2, errors_3, atol=1e-5)
         np.testing.assert_allclose(sensitivities_2, sensitivities_3, atol=1e-5)

From 68b763d2e92f8dd7d0269d811cd7309d72dc3912 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 5 Jul 2024 11:23:12 +0100
Subject: [PATCH 14/25] Pass current_sensitivities to MAP

---
 pybop/costs/fitting_costs.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pybop/costs/fitting_costs.py b/pybop/costs/fitting_costs.py
index 08eef2390..8634e16f3 100644
--- a/pybop/costs/fitting_costs.py
+++ b/pybop/costs/fitting_costs.py
@@ -380,7 +380,10 @@ def _evaluateS1(self, inputs: Inputs):
             If an error occurs during the calculation of the cost or gradient.
         """
         if self._fixed_problem:
-            self.likelihood._current_prediction = self._current_prediction
+            (
+                self.likelihood._current_prediction,
+                self.likelihood._current_sensitivities,
+            ) = self._current_prediction, self._current_sensitivities
 
         log_likelihood, dl = self.likelihood._evaluateS1(inputs)
         log_prior = sum(

From 8c2632dfd0b636488d2e6c8eac55002798914e08 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 5 Jul 2024 12:10:52 +0100
Subject: [PATCH 15/25] Add test_weighted_design_cost

---
 pybop/costs/base_cost.py |  6 ++++--
 tests/unit/test_cost.py  | 43 +++++++++++++++++++++++++---------------
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 7fae9ee9c..b285ef261 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -198,13 +198,15 @@ def __init__(self, cost_list, weights=None):
 
         # Check if all costs depend on the same problem
         for cost in self.cost_list:
-            if hasattr(cost, "problem") and (
-                not cost._fixed_problem or cost.problem is not self.cost_list[0].problem
+            if (
+                hasattr(cost, "problem")
+                and cost.problem is not self.cost_list[0].problem
             ):
                 self._different_problems = True
 
         if not self._different_problems:
             super(WeightedCost, self).__init__(self.cost_list[0].problem)
+            self._fixed_problem = self.cost_list[0]._fixed_problem
         else:
             super(WeightedCost, self).__init__()
             self._fixed_problem = False
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index cc45deed6..0db06b05e 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -200,6 +200,12 @@ def test_costs(self, cost):
         # Test treatment of simulations that terminated early
         # by variation of the cut-off voltage.
 
+    @pytest.fixture
+    def design_problem(self, model, parameters, experiment, signal):
+        return pybop.DesignProblem(
+            model, parameters, experiment, signal=signal, init_soc=0.5
+        )
+
     @pytest.mark.parametrize(
         "cost_class",
         [
@@ -209,21 +215,9 @@ def test_costs(self, cost):
         ],
     )
     @pytest.mark.unit
-    def test_design_costs(
-        self,
-        cost_class,
-        model,
-        parameters,
-        experiment,
-        signal,
-    ):
-        # Construct Problem
-        problem = pybop.DesignProblem(
-            model, parameters, experiment, signal=signal, init_soc=0.5
-        )
-
+    def test_design_costs(self, cost_class, design_problem):
         # Construct Cost
-        cost = cost_class(problem)
+        cost = cost_class(design_problem)
 
         if cost_class in [pybop.DesignCost]:
             with pytest.raises(NotImplementedError):
@@ -249,11 +243,11 @@ def test_design_costs(
                 cost(["StringInputShouldNotWork"])
 
             # Compute after updating nominal capacity
-            cost = cost_class(problem, update_capacity=True)
+            cost = cost_class(design_problem, update_capacity=True)
             cost([0.4])
 
     @pytest.mark.unit
-    def test_weighted_cost(self, problem):
+    def test_weighted_fitting_cost(self, problem):
         cost1 = pybop.SumSquaredError(problem)
         cost2 = pybop.RootMeanSquaredError(problem)
 
@@ -317,3 +311,20 @@ def test_weighted_cost(self, problem):
         errors_3, sensitivities_3 = weighted_cost_3.evaluateS1([0.5])
         np.testing.assert_allclose(errors_2, errors_3, atol=1e-5)
         np.testing.assert_allclose(sensitivities_2, sensitivities_3, atol=1e-5)
+
+    @pytest.mark.unit
+    def test_weighted_design_cost(self, design_problem):
+        cost1 = pybop.GravimetricEnergyDensity(design_problem)
+        cost2 = pybop.RootMeanSquaredError(design_problem)
+
+        # Test with and without weights
+        weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2])
+        assert weighted_cost._different_problems is False
+        assert weighted_cost._fixed_problem is False
+        assert weighted_cost.problem is design_problem
+        assert weighted_cost([0.5]) >= 0
+        np.testing.assert_allclose(
+            weighted_cost.evaluate([0.6]),
+            cost1([0.6]) + cost2([0.6]),
+            atol=1e-5,
+        )

From c189d7a2ccc3ae6ab3c847ecd6c209e797d2b963 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 5 Jul 2024 12:35:53 +0000
Subject: [PATCH 16/25] style: pre-commit fixes

---
 pybop/costs/_likelihoods.py | 43 ++++++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/pybop/costs/_likelihoods.py b/pybop/costs/_likelihoods.py
index d05e0e9d4..696ae01db 100644
--- a/pybop/costs/_likelihoods.py
+++ b/pybop/costs/_likelihoods.py
@@ -45,7 +45,8 @@ def _evaluate(self, inputs: Inputs, grad: Union[None, np.ndarray] = None) -> flo
         Evaluates the Gaussian log-likelihood for the given parameters with known sigma.
         """
         if any(
-            len(self._current_prediction.get(key, [])) != len(self._target.get(key, [])) for key in self.signal
+            len(self._current_prediction.get(key, [])) != len(self._target.get(key, []))
+            for key in self.signal
         ):
             return -np.inf  # prediction length doesn't match target
 
@@ -53,7 +54,10 @@ def _evaluate(self, inputs: Inputs, grad: Union[None, np.ndarray] = None) -> flo
             [
                 np.sum(
                     self._offset
-                    + self._multip * np.sum((self._target[signal] - self._current_prediction[signal]) ** 2.0)
+                    + self._multip
+                    * np.sum(
+                        (self._target[signal] - self._current_prediction[signal]) ** 2.0
+                    )
                 )
                 for signal in self.signal
             ]
@@ -66,14 +70,22 @@ def _evaluateS1(self, inputs: Inputs) -> Tuple[float, np.ndarray]:
         Calculates the log-likelihood and gradient.
         """
         if any(
-            len(self._current_prediction.get(key, [])) != len(self._target.get(key, [])) for key in self.signal
+            len(self._current_prediction.get(key, [])) != len(self._target.get(key, []))
+            for key in self.signal
         ):
             return -np.inf, -self._dl
 
         likelihood = self._evaluate(inputs)
 
-        r = np.asarray([self._target[signal] - self._current_prediction[signal] for signal in self.signal])
-        dl = np.sum((np.sum((r * self._current_sensitivities.T), axis=2) / self.sigma2), axis=1)
+        r = np.asarray(
+            [
+                self._target[signal] - self._current_prediction[signal]
+                for signal in self.signal
+            ]
+        )
+        dl = np.sum(
+            (np.sum((r * self._current_sensitivities.T), axis=2) / self.sigma2), axis=1
+        )
 
         return likelihood, dl
 
@@ -193,7 +205,8 @@ def _evaluate(self, inputs: Inputs, grad: Union[None, np.ndarray] = None) -> flo
             return -np.inf
 
         if any(
-            len(self._current_prediction.get(key, [])) != len(self._target.get(key, [])) for key in self.signal
+            len(self._current_prediction.get(key, [])) != len(self._target.get(key, []))
+            for key in self.signal
         ):
             return -np.inf  # prediction length doesn't match target
 
@@ -202,7 +215,9 @@ def _evaluate(self, inputs: Inputs, grad: Union[None, np.ndarray] = None) -> flo
                 np.sum(
                     self._logpi
                     - self.n_time_data * np.log(sigma)
-                    - np.sum((self._target[signal] - self._current_prediction[signal]) ** 2.0)
+                    - np.sum(
+                        (self._target[signal] - self._current_prediction[signal]) ** 2.0
+                    )
                     / (2.0 * sigma**2.0)
                 )
                 for signal in self.signal
@@ -232,14 +247,22 @@ def _evaluateS1(self, inputs: Inputs) -> Tuple[float, np.ndarray]:
             return -np.inf, -self._dl
 
         if any(
-            len(self._current_prediction.get(key, [])) != len(self._target.get(key, [])) for key in self.signal
+            len(self._current_prediction.get(key, [])) != len(self._target.get(key, []))
+            for key in self.signal
         ):
             return -np.inf, -self._dl
 
         likelihood = self._evaluate(inputs)
 
-        r = np.asarray([self._target[signal] - self._current_prediction[signal] for signal in self.signal])
-        dl = np.sum((np.sum((r * self._current_sensitivities.T), axis=2) / (sigma**2.0)), axis=1)
+        r = np.asarray(
+            [
+                self._target[signal] - self._current_prediction[signal]
+                for signal in self.signal
+            ]
+        )
+        dl = np.sum(
+            (np.sum((r * self._current_sensitivities.T), axis=2) / (sigma**2.0)), axis=1
+        )
         dsigma = (
             -self.n_time_data / sigma + np.sum(r**2.0, axis=1) / (sigma**3.0)
         ) / self._dsigma_scale

From 2b1ae4c6af54fc719a3c3dabea9d3594663bbb24 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 5 Jul 2024 14:08:51 +0100
Subject: [PATCH 17/25] Add evaluate back into GaussianLogLikelihood

---
 pybop/costs/_likelihoods.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pybop/costs/_likelihoods.py b/pybop/costs/_likelihoods.py
index 696ae01db..8b867d4eb 100644
--- a/pybop/costs/_likelihoods.py
+++ b/pybop/costs/_likelihoods.py
@@ -204,6 +204,9 @@ def _evaluate(self, inputs: Inputs, grad: Union[None, np.ndarray] = None) -> flo
         if np.any(sigma <= 0):
             return -np.inf
 
+        self._current_prediction = self.problem.evaluate(
+            self.problem.parameters.as_dict()
+        )
         if any(
             len(self._current_prediction.get(key, [])) != len(self._target.get(key, []))
             for key in self.signal
@@ -246,6 +249,9 @@ def _evaluateS1(self, inputs: Inputs) -> Tuple[float, np.ndarray]:
         if np.any(sigma <= 0):
             return -np.inf, -self._dl
 
+        self._current_prediction, self._current_sensitivities = self.problem.evaluateS1(
+            self.problem.parameters.as_dict()
+        )
         if any(
             len(self._current_prediction.get(key, [])) != len(self._target.get(key, []))
             for key in self.signal

From 8c41327a302e6196074be3519f6673edee34ad57 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 11 Jul 2024 11:44:28 +0100
Subject: [PATCH 18/25] Update to super()

---
 pybop/costs/base_cost.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 322d082fc..0bd38192d 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -205,10 +205,10 @@ def __init__(self, cost_list, weights=None):
                 self._different_problems = True
 
         if not self._different_problems:
-            super(WeightedCost, self).__init__(self.cost_list[0].problem)
+            super().__init__(self.cost_list[0].problem)
             self._fixed_problem = self.cost_list[0]._fixed_problem
         else:
-            super(WeightedCost, self).__init__()
+            super().__init__()
             self._fixed_problem = False
             for cost in self.cost_list:
                 self.parameters.join(cost.parameters)

From 80a803e4251a05a7d75239ea8d4f68e4814985f7 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 15 Jul 2024 16:42:55 +0000
Subject: [PATCH 19/25] style: pre-commit fixes

---
 pybop/costs/_likelihoods.py | 2 +-
 pybop/costs/base_cost.py    | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/pybop/costs/_likelihoods.py b/pybop/costs/_likelihoods.py
index 308704153..43dc326aa 100644
--- a/pybop/costs/_likelihoods.py
+++ b/pybop/costs/_likelihoods.py
@@ -324,7 +324,7 @@ def _evaluate(self, inputs: Inputs, grad=None) -> float:
         if self._fixed_problem:
             self.likelihood._current_prediction = self._current_prediction
         log_likelihood = self.likelihood._evaluate(inputs)
-        
+
         posterior = log_likelihood + log_prior
         return posterior
 
diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 4226ed78f..eb2038b66 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -1,6 +1,7 @@
-import numpy as np
 from typing import Union
 
+import numpy as np
+
 from pybop import BaseProblem
 from pybop.parameters.parameter import Inputs, Parameters
 

From f18552349c6b5447d48601bf1c64505ab157e450 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Mon, 15 Jul 2024 17:48:53 +0100
Subject: [PATCH 20/25] Update prediction to self._current_prediction

---
 pybop/costs/fitting_costs.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pybop/costs/fitting_costs.py b/pybop/costs/fitting_costs.py
index 2ff465868..93c50a75c 100644
--- a/pybop/costs/fitting_costs.py
+++ b/pybop/costs/fitting_costs.py
@@ -239,7 +239,10 @@ def _evaluate(self, inputs: Inputs, grad=None):
 
         e = np.asarray(
             [
-                np.sum(np.abs(prediction[signal] - self._target[signal]) ** self.p)
+                np.sum(
+                    np.abs(self._current_prediction[signal] - self._target[signal])
+                    ** self.p
+                )
                 ** (1 / self.p)
                 for signal in self.signal
             ]
@@ -344,7 +347,10 @@ def _evaluate(self, inputs: Inputs, grad=None):
 
         e = np.asarray(
             [
-                np.sum(np.abs(prediction[signal] - self._target[signal]) ** self.p)
+                np.sum(
+                    np.abs(self._current_prediction[signal] - self._target[signal])
+                    ** self.p
+                )
                 for signal in self.signal
             ]
         )

From 23b77e895d9ee7aa23413ac1bad457192b4450f5 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Mon, 15 Jul 2024 17:52:01 +0100
Subject: [PATCH 21/25] Update y to self._current_prediction

---
 pybop/costs/fitting_costs.py | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/pybop/costs/fitting_costs.py b/pybop/costs/fitting_costs.py
index 93c50a75c..18cc752d6 100644
--- a/pybop/costs/fitting_costs.py
+++ b/pybop/costs/fitting_costs.py
@@ -273,16 +273,24 @@ def _evaluateS1(self, inputs):
         if not self.verify_prediction(self._current_prediction):
             return np.inf, self._de * np.ones(self.n_parameters)
 
-        r = np.asarray([y[signal] - self._target[signal] for signal in self.signal])
+        r = np.asarray(
+            [
+                self._current_prediction[signal] - self._target[signal]
+                for signal in self.signal
+            ]
+        )
         e = np.asarray(
             [
-                np.sum(np.abs(y[signal] - self._target[signal]) ** self.p)
+                np.sum(
+                    np.abs(self._current_prediction[signal] - self._target[signal])
+                    ** self.p
+                )
                 ** (1 / self.p)
                 for signal in self.signal
             ]
         )
         de = np.sum(
-            np.sum(r ** (self.p - 1) * dy.T, axis=2)
+            np.sum(r ** (self.p - 1) * self._current_sensitivities.T, axis=2)
             / (e ** (self.p - 1) + np.finfo(float).eps),
             axis=1,
         )
@@ -380,9 +388,16 @@ def _evaluateS1(self, inputs):
         if not self.verify_prediction(self._current_prediction):
             return np.inf, self._de * np.ones(self.n_parameters)
 
-        r = np.asarray([y[signal] - self._target[signal] for signal in self.signal])
+        r = np.asarray(
+            [
+                self._current_prediction[signal] - self._target[signal]
+                for signal in self.signal
+            ]
+        )
         e = np.sum(np.sum(np.abs(r) ** self.p))
-        de = self.p * np.sum(np.sum(r ** (self.p - 1) * dy.T, axis=2), axis=1)
+        de = self.p * np.sum(
+            np.sum(r ** (self.p - 1) * self._current_sensitivities.T, axis=2), axis=1
+        )
 
         return e, de
 

From 5157a8d8c96390e202877bb81af984ac50dda78d Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 16 Jul 2024 17:10:20 +0100
Subject: [PATCH 22/25] Update cost_list to args

---
 examples/scripts/spm_weighted_cost.py |  2 +-
 pybop/costs/base_cost.py              | 48 ++++++++++++++-------------
 tests/unit/test_cost.py               | 32 +++++++-----------
 3 files changed, 38 insertions(+), 44 deletions(-)

diff --git a/examples/scripts/spm_weighted_cost.py b/examples/scripts/spm_weighted_cost.py
index 557434613..74c43a33c 100644
--- a/examples/scripts/spm_weighted_cost.py
+++ b/examples/scripts/spm_weighted_cost.py
@@ -41,7 +41,7 @@
 problem = pybop.FittingProblem(model, parameters, dataset)
 cost1 = pybop.SumSquaredError(problem)
 cost2 = pybop.RootMeanSquaredError(problem)
-weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 100])
+weighted_cost = pybop.WeightedCost(cost1, cost2, weights=[1, 100])
 
 for cost in [weighted_cost, cost1, cost2]:
     optim = pybop.IRPropMin(cost, max_iterations=60)
diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index eb2038b66..9a3590530 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -212,45 +212,47 @@ class WeightedCost(BaseCost):
     a single weighted cost function.
 
     Inherits all parameters and attributes from ``BaseCost``.
+
+    Additional Attributes
+    ---------------------
+    costs : List[pybop.BaseCost]
+        A list of PyBOP cost objects.
     """
 
-    def __init__(self, cost_list, weights=None):
-        self.cost_list = cost_list
+    def __init__(self, *args, weights=None):
+        self.costs = []
+        for cost in args:
+            if not isinstance(cost, BaseCost):
+                raise TypeError(f"Received {type(cost)} instead of cost object.")
+            self.costs.append(cost)
         self.weights = weights
         self._different_problems = False
 
-        if not isinstance(self.cost_list, list):
-            raise TypeError(
-                f"Expected a list of costs. Received {type(self.cost_list)}"
-            )
         if self.weights is None:
-            self.weights = np.ones(len(cost_list))
+            self.weights = np.ones(len(self.costs))
         elif isinstance(self.weights, list):
             self.weights = np.array(self.weights)
         if not isinstance(self.weights, np.ndarray):
             raise TypeError(
-                "Expected a list or array of weights the same length as cost_list."
+                "Expected a list or array of weights the same length as costs."
             )
-        if not len(self.weights) == len(self.cost_list):
+        if not len(self.weights) == len(self.costs):
             raise ValueError(
-                "Expected a list or array of weights the same length as cost_list."
+                "Expected a list or array of weights the same length as costs."
             )
 
         # Check if all costs depend on the same problem
-        for cost in self.cost_list:
-            if (
-                hasattr(cost, "problem")
-                and cost.problem is not self.cost_list[0].problem
-            ):
+        for cost in self.costs:
+            if hasattr(cost, "problem") and cost.problem is not self.costs[0].problem:
                 self._different_problems = True
 
         if not self._different_problems:
-            super().__init__(self.cost_list[0].problem)
-            self._fixed_problem = self.cost_list[0]._fixed_problem
+            super().__init__(self.costs[0].problem)
+            self._fixed_problem = self.costs[0]._fixed_problem
         else:
             super().__init__()
             self._fixed_problem = False
-            for cost in self.cost_list:
+            for cost in self.costs:
                 self.parameters.join(cost.parameters)
 
     def _evaluate(self, inputs: Inputs, grad=None):
@@ -270,14 +272,14 @@ def _evaluate(self, inputs: Inputs, grad=None):
         float
             The weighted cost value.
         """
-        e = np.empty_like(self.cost_list)
+        e = np.empty_like(self.costs)
 
         if not self._fixed_problem and self._different_problems:
             self.parameters.update(values=list(inputs.values()))
         elif not self._fixed_problem:
             self._current_prediction = self.problem.evaluate(inputs)
 
-        for i, cost in enumerate(self.cost_list):
+        for i, cost in enumerate(self.costs):
             if not self._fixed_problem and self._different_problems:
                 inputs = cost.parameters.as_dict()
                 cost._current_prediction = cost.problem.evaluate(inputs)
@@ -302,8 +304,8 @@ def _evaluateS1(self, inputs: Inputs):
             A tuple containing the cost and the gradient. The cost is a float,
             and the gradient is an array-like of the same length as `x`.
         """
-        e = np.empty_like(self.cost_list)
-        de = np.empty((len(self.parameters), len(self.cost_list)))
+        e = np.empty_like(self.costs)
+        de = np.empty((len(self.parameters), len(self.costs)))
 
         if not self._fixed_problem and self._different_problems:
             self.parameters.update(values=list(inputs.values()))
@@ -312,7 +314,7 @@ def _evaluateS1(self, inputs: Inputs):
                 self.problem.evaluateS1(inputs)
             )
 
-        for i, cost in enumerate(self.cost_list):
+        for i, cost in enumerate(self.costs):
             if not self._fixed_problem and self._different_problems:
                 inputs = cost.parameters.as_dict()
                 cost._current_prediction, cost._current_sensitivities = (
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 2c9c2d5e5..ac6e58772 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -304,37 +304,31 @@ def test_weighted_fitting_cost(self, problem):
         cost2 = pybop.RootMeanSquaredError(problem)
 
         # Test with and without weights
-        weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2])
+        weighted_cost = pybop.WeightedCost(cost1, cost2)
         np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
-        weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1, 1])
+        weighted_cost = pybop.WeightedCost(cost1, cost2, weights=[1, 1])
         np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
-        weighted_cost = pybop.WeightedCost(
-            cost_list=[cost1, cost2], weights=np.array([1, 1])
-        )
+        weighted_cost = pybop.WeightedCost(cost1, cost2, weights=np.array([1, 1]))
         np.testing.assert_array_equal(weighted_cost.weights, np.ones(2))
         with pytest.raises(
             TypeError,
-            match=r"Expected a list of costs.",
+            match=r"Received <class 'str'> instead of cost object.",
         ):
-            weighted_cost = pybop.WeightedCost(cost_list="Invalid string")
+            weighted_cost = pybop.WeightedCost("Invalid string")
         with pytest.raises(
             TypeError,
-            match="Expected a list or array of weights the same length as cost_list.",
+            match="Expected a list or array of weights the same length as costs.",
         ):
-            weighted_cost = pybop.WeightedCost(
-                cost_list=[cost1, cost2], weights="Invalid string"
-            )
+            weighted_cost = pybop.WeightedCost(cost1, cost2, weights="Invalid string")
         with pytest.raises(
             ValueError,
-            match="Expected a list or array of weights the same length as cost_list.",
+            match="Expected a list or array of weights the same length as costs.",
         ):
-            weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2], weights=[1])
+            weighted_cost = pybop.WeightedCost(cost1, cost2, weights=[1])
 
         # Test with and without different problems
         weight = 100
-        weighted_cost_2 = pybop.WeightedCost(
-            cost_list=[cost1, cost2], weights=[1, weight]
-        )
+        weighted_cost_2 = pybop.WeightedCost(cost1, cost2, weights=[1, weight])
         assert weighted_cost_2._different_problems is False
         assert weighted_cost_2._fixed_problem is True
         assert weighted_cost_2.problem is problem
@@ -346,9 +340,7 @@ def test_weighted_fitting_cost(self, problem):
         )
 
         cost3 = pybop.RootMeanSquaredError(copy(problem))
-        weighted_cost_3 = pybop.WeightedCost(
-            cost_list=[cost1, cost3], weights=[1, weight]
-        )
+        weighted_cost_3 = pybop.WeightedCost(cost1, cost3, weights=[1, weight])
         assert weighted_cost_3._different_problems is True
         assert weighted_cost_3._fixed_problem is False
         assert weighted_cost_3.problem is None
@@ -370,7 +362,7 @@ def test_weighted_design_cost(self, design_problem):
         cost2 = pybop.RootMeanSquaredError(design_problem)
 
         # Test with and without weights
-        weighted_cost = pybop.WeightedCost(cost_list=[cost1, cost2])
+        weighted_cost = pybop.WeightedCost(cost1, cost2)
         assert weighted_cost._different_problems is False
         assert weighted_cost._fixed_problem is False
         assert weighted_cost.problem is design_problem

From 3e220b16ccc670ce071f1f4019be2ab1de0a594e Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 16 Jul 2024 18:04:58 +0100
Subject: [PATCH 23/25] Remove repeated line

---
 pybop/costs/_likelihoods.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pybop/costs/_likelihoods.py b/pybop/costs/_likelihoods.py
index 43dc326aa..1f96e2fb5 100644
--- a/pybop/costs/_likelihoods.py
+++ b/pybop/costs/_likelihoods.py
@@ -241,7 +241,6 @@ def _evaluateS1(self, inputs: Inputs) -> tuple[float, np.ndarray]:
         self._current_prediction, self._current_sensitivities = self.problem.evaluateS1(
             self.problem.parameters.as_dict()
         )
-        y, dy = self.problem.evaluateS1(self.problem.parameters.as_dict())
         if not self.verify_prediction(self._current_prediction):
             return -np.inf, -self._de * np.ones(self.n_parameters)
 

From 95600be95bd5994630023efbb2d6581c9a3ea9d2 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 16 Jul 2024 18:05:15 +0100
Subject: [PATCH 24/25] Add descriptions

---
 pybop/costs/base_cost.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index 9a3590530..b3cf7ae1e 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
@@ -24,9 +24,15 @@ class BaseCost:
         An array containing the target data to fit.
     n_outputs : int
         The number of outputs in the model.
+
+    Additional Attributes
+    ---------------------
+    _fixed_problem : bool
+        If True, the problem does not need to be rebuilt before the cost is
+        calculated (default: False).
     """
 
-    def __init__(self, problem=None):
+    def __init__(self, problem: Optional[BaseProblem] = None):
         self.parameters = Parameters()
         self.problem = problem
         self._fixed_problem = False
@@ -215,11 +221,16 @@ class WeightedCost(BaseCost):
 
     Additional Attributes
     ---------------------
-    costs : List[pybop.BaseCost]
+    costs : list[pybop.BaseCost]
         A list of PyBOP cost objects.
+    weights : list[float]
+        A list of values with which to weight the cost values.
+    _different_problems : bool
+        If True, the problem for each cost is evaluated independently during
+        each evaluation of the cost (default: False).
     """
 
-    def __init__(self, *args, weights=None):
+    def __init__(self, *args, weights: Optional[list[float]] = None):
         self.costs = []
         for cost in args:
             if not isinstance(cost, BaseCost):

From d739642faee723cb872384753fa733877922fc32 Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Thu, 18 Jul 2024 13:43:11 +0100
Subject: [PATCH 25/25] refactor: move WeightedCost into seperate file

---
 pybop/__init__.py             |   3 +-
 pybop/costs/_weighted_cost.py | 138 ++++++++++++++++++++++++++++++++++
 pybop/costs/base_cost.py      | 134 ---------------------------------
 3 files changed, 140 insertions(+), 135 deletions(-)
 create mode 100644 pybop/costs/_weighted_cost.py

diff --git a/pybop/__init__.py b/pybop/__init__.py
index 4a99ebc48..922a64803 100644
--- a/pybop/__init__.py
+++ b/pybop/__init__.py
@@ -81,7 +81,7 @@
 #
 # Cost function class
 #
-from .costs.base_cost import BaseCost, WeightedCost
+from .costs.base_cost import BaseCost
 from .costs.fitting_costs import (
     RootMeanSquaredError,
     SumSquaredError,
@@ -100,6 +100,7 @@
     GaussianLogLikelihoodKnownSigma,
     MAP,
 )
+from .costs._weighted_cost import WeightedCost
 
 #
 # Optimiser class
diff --git a/pybop/costs/_weighted_cost.py b/pybop/costs/_weighted_cost.py
new file mode 100644
index 000000000..effa5a510
--- /dev/null
+++ b/pybop/costs/_weighted_cost.py
@@ -0,0 +1,138 @@
+from typing import Optional
+
+import numpy as np
+
+from pybop import BaseCost
+from pybop.parameters.parameter import Inputs
+
+
+class WeightedCost(BaseCost):
+    """
+    A subclass for constructing a linear combination of cost functions as
+    a single weighted cost function.
+
+    Inherits all parameters and attributes from ``BaseCost``.
+
+    Additional Attributes
+    ---------------------
+    costs : list[pybop.BaseCost]
+        A list of PyBOP cost objects.
+    weights : list[float]
+        A list of values with which to weight the cost values.
+    _different_problems : bool
+        If True, the problem for each cost is evaluated independently during
+        each evaluation of the cost (default: False).
+    """
+
+    def __init__(self, *args, weights: Optional[list[float]] = None):
+        self.costs = []
+        for cost in args:
+            if not isinstance(cost, BaseCost):
+                raise TypeError(f"Received {type(cost)} instead of cost object.")
+            self.costs.append(cost)
+        self.weights = weights
+        self._different_problems = False
+
+        if self.weights is None:
+            self.weights = np.ones(len(self.costs))
+        elif isinstance(self.weights, list):
+            self.weights = np.array(self.weights)
+        if not isinstance(self.weights, np.ndarray):
+            raise TypeError(
+                "Expected a list or array of weights the same length as costs."
+            )
+        if not len(self.weights) == len(self.costs):
+            raise ValueError(
+                "Expected a list or array of weights the same length as costs."
+            )
+
+        # Check if all costs depend on the same problem
+        for cost in self.costs:
+            if hasattr(cost, "problem") and cost.problem is not self.costs[0].problem:
+                self._different_problems = True
+
+        if not self._different_problems:
+            super().__init__(self.costs[0].problem)
+            self._fixed_problem = self.costs[0]._fixed_problem
+        else:
+            super().__init__()
+            self._fixed_problem = False
+            for cost in self.costs:
+                self.parameters.join(cost.parameters)
+
+    def _evaluate(self, inputs: Inputs, grad=None):
+        """
+        Calculate the weighted cost for a given set of parameters.
+
+        Parameters
+        ----------
+        inputs : Inputs
+            The parameters for which to compute the cost.
+        grad : array-like, optional
+            An array to store the gradient of the cost function with respect
+            to the parameters.
+
+        Returns
+        -------
+        float
+            The weighted cost value.
+        """
+        e = np.empty_like(self.costs)
+
+        if not self._fixed_problem and self._different_problems:
+            self.parameters.update(values=list(inputs.values()))
+        elif not self._fixed_problem:
+            self._current_prediction = self.problem.evaluate(inputs)
+
+        for i, cost in enumerate(self.costs):
+            if not self._fixed_problem and self._different_problems:
+                inputs = cost.parameters.as_dict()
+                cost._current_prediction = cost.problem.evaluate(inputs)
+            else:
+                cost._current_prediction = self._current_prediction
+            e[i] = cost._evaluate(inputs, grad)
+
+        return np.dot(e, self.weights)
+
+    def _evaluateS1(self, inputs: Inputs):
+        """
+        Compute the weighted cost and its gradient with respect to the parameters.
+
+        Parameters
+        ----------
+        inputs : Inputs
+            The parameters for which to compute the cost and gradient.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the cost and the gradient. The cost is a float,
+            and the gradient is an array-like of the same length as `x`.
+        """
+        e = np.empty_like(self.costs)
+        de = np.empty((len(self.parameters), len(self.costs)))
+
+        if not self._fixed_problem and self._different_problems:
+            self.parameters.update(values=list(inputs.values()))
+        elif not self._fixed_problem:
+            self._current_prediction, self._current_sensitivities = (
+                self.problem.evaluateS1(inputs)
+            )
+
+        for i, cost in enumerate(self.costs):
+            if not self._fixed_problem and self._different_problems:
+                inputs = cost.parameters.as_dict()
+                cost._current_prediction, cost._current_sensitivities = (
+                    cost.problem.evaluateS1(inputs)
+                )
+            else:
+                cost._current_prediction, cost._current_sensitivities = (
+                    self._current_prediction,
+                    self._current_sensitivities,
+                )
+            e[i], de[:, i] = cost._evaluateS1(inputs)
+
+        e = np.dot(e, self.weights)
+        de = np.dot(de, self.weights)
+
+        return e, de
diff --git a/pybop/costs/base_cost.py b/pybop/costs/base_cost.py
index b3cf7ae1e..eedbbc2c0 100644
--- a/pybop/costs/base_cost.py
+++ b/pybop/costs/base_cost.py
@@ -1,7 +1,5 @@
 from typing import Optional, Union
 
-import numpy as np
-
 from pybop import BaseProblem
 from pybop.parameters.parameter import Inputs, Parameters
 
@@ -210,135 +208,3 @@ def verify_prediction(self, y):
             return False
 
         return True
-
-
-class WeightedCost(BaseCost):
-    """
-    A subclass for constructing a linear combination of cost functions as
-    a single weighted cost function.
-
-    Inherits all parameters and attributes from ``BaseCost``.
-
-    Additional Attributes
-    ---------------------
-    costs : list[pybop.BaseCost]
-        A list of PyBOP cost objects.
-    weights : list[float]
-        A list of values with which to weight the cost values.
-    _different_problems : bool
-        If True, the problem for each cost is evaluated independently during
-        each evaluation of the cost (default: False).
-    """
-
-    def __init__(self, *args, weights: Optional[list[float]] = None):
-        self.costs = []
-        for cost in args:
-            if not isinstance(cost, BaseCost):
-                raise TypeError(f"Received {type(cost)} instead of cost object.")
-            self.costs.append(cost)
-        self.weights = weights
-        self._different_problems = False
-
-        if self.weights is None:
-            self.weights = np.ones(len(self.costs))
-        elif isinstance(self.weights, list):
-            self.weights = np.array(self.weights)
-        if not isinstance(self.weights, np.ndarray):
-            raise TypeError(
-                "Expected a list or array of weights the same length as costs."
-            )
-        if not len(self.weights) == len(self.costs):
-            raise ValueError(
-                "Expected a list or array of weights the same length as costs."
-            )
-
-        # Check if all costs depend on the same problem
-        for cost in self.costs:
-            if hasattr(cost, "problem") and cost.problem is not self.costs[0].problem:
-                self._different_problems = True
-
-        if not self._different_problems:
-            super().__init__(self.costs[0].problem)
-            self._fixed_problem = self.costs[0]._fixed_problem
-        else:
-            super().__init__()
-            self._fixed_problem = False
-            for cost in self.costs:
-                self.parameters.join(cost.parameters)
-
-    def _evaluate(self, inputs: Inputs, grad=None):
-        """
-        Calculate the weighted cost for a given set of parameters.
-
-        Parameters
-        ----------
-        inputs : Inputs
-            The parameters for which to compute the cost.
-        grad : array-like, optional
-            An array to store the gradient of the cost function with respect
-            to the parameters.
-
-        Returns
-        -------
-        float
-            The weighted cost value.
-        """
-        e = np.empty_like(self.costs)
-
-        if not self._fixed_problem and self._different_problems:
-            self.parameters.update(values=list(inputs.values()))
-        elif not self._fixed_problem:
-            self._current_prediction = self.problem.evaluate(inputs)
-
-        for i, cost in enumerate(self.costs):
-            if not self._fixed_problem and self._different_problems:
-                inputs = cost.parameters.as_dict()
-                cost._current_prediction = cost.problem.evaluate(inputs)
-            else:
-                cost._current_prediction = self._current_prediction
-            e[i] = cost._evaluate(inputs, grad)
-
-        return np.dot(e, self.weights)
-
-    def _evaluateS1(self, inputs: Inputs):
-        """
-        Compute the weighted cost and its gradient with respect to the parameters.
-
-        Parameters
-        ----------
-        inputs : Inputs
-            The parameters for which to compute the cost and gradient.
-
-        Returns
-        -------
-        tuple
-            A tuple containing the cost and the gradient. The cost is a float,
-            and the gradient is an array-like of the same length as `x`.
-        """
-        e = np.empty_like(self.costs)
-        de = np.empty((len(self.parameters), len(self.costs)))
-
-        if not self._fixed_problem and self._different_problems:
-            self.parameters.update(values=list(inputs.values()))
-        elif not self._fixed_problem:
-            self._current_prediction, self._current_sensitivities = (
-                self.problem.evaluateS1(inputs)
-            )
-
-        for i, cost in enumerate(self.costs):
-            if not self._fixed_problem and self._different_problems:
-                inputs = cost.parameters.as_dict()
-                cost._current_prediction, cost._current_sensitivities = (
-                    cost.problem.evaluateS1(inputs)
-                )
-            else:
-                cost._current_prediction, cost._current_sensitivities = (
-                    self._current_prediction,
-                    self._current_sensitivities,
-                )
-            e[i], de[:, i] = cost._evaluateS1(inputs)
-
-        e = np.dot(e, self.weights)
-        de = np.dot(de, self.weights)
-
-        return e, de