Convert initial_state to dict

examples/notebooks/multi_model_identification.ipynb

@@ -124,7 +124,7 @@
     "current = np.concatenate(\n",
     "    [np.ones(400) * parameter_set[\"Nominal cell capacity [A.h]\"], np.zeros(250)]\n",
     ")\n",
-    "init_soc = 0.5"
+    "initial_state = {\"Initial SoC\": 0.5}"
    ]
   },
   {
@@ -147,7 +147,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "synth_model.build(dataset, initial_state=init_soc)\n",
+    "synth_model.build(dataset, initial_state=initial_state)\n",
     "synth_model.signal = [\"Voltage [V]\"]\n",
     "values = synth_model.simulate(t_eval=t_eval, inputs={})"
    ]
@@ -2514,7 +2514,7 @@
     "xs = []\n",
     "for model in models:\n",
     "    print(f\"Running {model.name}\")\n",
-    "    model.set_initial_state(init_soc)\n",
+    "    model.set_initial_state(initial_state)\n",
     "    problem = pybop.FittingProblem(model, parameters, dataset)\n",
     "    cost = pybop.SumSquaredError(problem)\n",
     "    optim = pybop.XNES(\n",

examples/notebooks/multi_optimiser_identification.ipynb

@@ -115,8 +115,8 @@
    "outputs": [],
    "source": [
     "t_eval = np.arange(0, 2000, 10)\n",
-    "init_soc = 1.0\n",
-    "values = synth_model.predict(t_eval=t_eval, initial_state=init_soc)"
+    "initial_state = {\"Initial SoC\": 1.0}\n",
+    "values = synth_model.predict(t_eval=t_eval, initial_state=initial_state)"
    ]
   },
   {
@@ -448,7 +448,7 @@
    "source": [
     "optims = []\n",
     "xs = []\n",
-    "model.set_initial_state(init_soc)\n",
+    "model.set_initial_state(initial_state)\n",
     "problem = pybop.FittingProblem(model, parameters, dataset)\n",
     "cost = pybop.SumSquaredError(problem)\n",
     "for optimiser in gradient_optimisers:\n",

examples/notebooks/optimiser_calibration.ipynb

@@ -100,7 +100,7 @@
     "    }\n",
     ")\n",
     "model = pybop.lithium_ion.SPM(parameter_set=parameter_set)\n",
-    "init_soc = 0.4\n",
+    "initial_state = {\"Initial SoC\": 0.4}\n",
     "experiment = pybop.Experiment(\n",
     "    [\n",
     "        (\n",
@@ -110,7 +110,7 @@
     "    ]\n",
     "    * 2\n",
     ")\n",
-    "values = model.predict(initial_state=init_soc, experiment=experiment)"
+    "values = model.predict(initial_state=initial_state, experiment=experiment)"
    ]
   },
   {

examples/notebooks/spm_electrode_design.ipynb

@@ -171,7 +171,9 @@
    },
    "outputs": [],
    "source": [
-    "problem = pybop.DesignProblem(model, parameters, experiment, init_soc=1.0)\n",
+    "problem = pybop.DesignProblem(\n",
+    "    model, parameters, experiment, initial_state={\"Initial SoC\": 0.7}\n",
+    ")\n",
     "cost = pybop.GravimetricEnergyDensity(problem)"
    ]
   },

examples/scripts/cuckoo.py

@@ -27,7 +27,6 @@
         initial_value=0.41,
     ),
 )
-init_soc = 0.7
 experiment = pybop.Experiment(
     [
         (
@@ -36,7 +35,7 @@
         ),
     ]
 )
-values = model.predict(initial_state=init_soc, experiment=experiment)
+values = model.predict(initial_state={"Initial SoC": 0.7}, experiment=experiment)
 
 sigma = 0.002
 corrupt_values = values["Voltage [V]"].data + np.random.normal(

examples/scripts/spm_AdamW.py

@@ -19,7 +19,6 @@
 )
 
 # Generate data
-init_soc = 0.5
 sigma = 0.003
 experiment = pybop.Experiment(
     [
@@ -30,7 +29,7 @@
     ]
     * 2
 )
-values = model.predict(initial_state=init_soc, experiment=experiment)
+values = model.predict(initial_state={"Initial SoC": 0.5}, experiment=experiment)
 
 
 def noise(sigma):

examples/scripts/spm_MAP.py

@@ -4,7 +4,6 @@
 
 # Set variables
 sigma = 0.002
-init_soc = 0.7
 
 # Construct and update initial parameter values
 parameter_set = pybop.ParameterSet.pybamm("Chen2020")
@@ -45,9 +44,7 @@
         ),
     ]
 )
-values = model.predict(
-    initial_state=init_soc, experiment=experiment, inputs=parameters.true_value()
-)
+values = model.predict(initial_state={"Initial SoC": 0.7}, experiment=experiment)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(
     0, sigma, len(values["Voltage [V]"].data)
 )

examples/scripts/spm_NelderMead.py

@@ -19,7 +19,6 @@
 )
 
 # Generate data
-init_soc = 0.5
 sigma = 0.003
 experiment = pybop.Experiment(
     [
@@ -30,7 +29,7 @@
     ]
     * 2
 )
-values = model.predict(initial_state=init_soc, experiment=experiment)
+values = model.predict(initial_state={"Initial SoC": 0.5}, experiment=experiment)
 
 
 def noise(sigma):

examples/scripts/spm_weighted_cost.py

@@ -24,7 +24,6 @@
 
 # Generate data
 sigma = 0.001
-init_soc = 0.5
 experiment = pybop.Experiment(
     [
         (
@@ -34,7 +33,7 @@
     ]
     * 2
 )
-values = model.predict(experiment=experiment, initial_state=init_soc)
+values = model.predict(experiment=experiment, initial_state={"Initial SoC": 0.5})
 
 
 def noise(sigma):

examples/scripts/spme_max_energy.py

@@ -31,12 +31,11 @@
 experiment = pybop.Experiment(
     ["Discharge at 1C until 2.5 V (5 seconds period)"],
 )
-init_soc = 1  # start from full charge
 signal = ["Voltage [V]", "Current [A]"]
 
 # Generate problem
 problem = pybop.DesignProblem(
-    model, parameters, experiment, signal=signal, init_soc=init_soc
+    model, parameters, experiment, signal=signal, initial_state={"Initial SoC": 1.0}
 )
 
 # Generate multiple cost functions and combine them.

pybop/models/base_model.py

@@ -117,7 +117,7 @@ def build(
         dataset: Optional[Dataset] = None,
         parameters: Union[Parameters, dict] = None,
         check_model: bool = True,
-        initial_state: Optional[float] = None,
+        initial_state: Optional[dict] = None,
     ) -> None:
         """
         Construct the PyBaMM model if not already built, and set parameters.
@@ -134,10 +134,10 @@ def build(
             A pybop Parameters class or dictionary containing parameter values to apply to the model.
         check_model : bool, optional
             If True, the model will be checked for correctness after construction.
-        initial_state : float or str, optional
-            If float, this value is used as the initial state of charge (as a decimal between 0
-            and 1). If str ending in "V", this value is used as the initial open-circuit voltage.
-            Defaults to None, indicating that the existing initial concentrations will be used.
+        initial_state : dict, optional
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+            Defaults to None, indicating that the existing initial state of charge (for an ECM)
+            or initial concentrations (for an EChem model) will be used.
         """
         if parameters is not None:
             self.parameters = parameters
@@ -173,20 +173,50 @@ def build(
 
         self.n_states = self._built_model.len_rhs_and_alg  # len_rhs + len_alg
 
-    def set_initial_state(self, initial_state: Optional[float] = None):
+    def convert_to_pybamm_initial_state(self, initial_state: dict = {}):
         """
-        Set the initial concentrations for the battery model.
+        Convert an initial state of charge into a float and an initial open-circuit
+        voltage into a string ending in "V".
 
         Parameters
         ----------
-        initial_state : float or str, optional
+        initial_state : dict
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+
+        Returns
+        -------
+        float or str
             If float, this value is used as the initial state of charge (as a decimal between 0
             and 1). If str ending in "V", this value is used as the initial open-circuit voltage.
-            Defaults to None, indicating that the existing initial concentrations will be used.
+
+        Raises
+        ------
+        ValueError
+            If the input is not a dictionary with a single, valid key.
+        """
+        if len(initial_state) > 1:
+            raise ValueError("Expecting only one initial state.")
+        elif "Initial SoC" in initial_state.keys():
+            return initial_state["Initial SoC"]
+        elif "Initial open-circuit voltage [V]" in initial_state.keys():
+            return str(initial_state["Initial open-circuit voltage [V]"]) + "V"
+        else:
+            raise ValueError("Unrecognised initial state.")
+
+    def set_initial_state(self, initial_state: dict = {}):
         """
+        Set the initial state of charge or concentrations for the battery model.
+
+        Parameters
+        ----------
+        initial_state : dict
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+        """
+        initial_state = self.convert_to_pybamm_initial_state(initial_state)
+
         if isinstance(self.pybamm_model, pybamm.equivalent_circuit.Thevenin):
-            initial_soc = self.get_initial_state(initial_state)
-            self._unprocessed_parameter_set.update({"Initial SoC": initial_soc})
+            initial_state = self.get_initial_state(initial_state)
+            self._unprocessed_parameter_set.update({"Initial SoC": initial_state})
 
         else:
             # Temporary construction of attribute for PyBaMM
@@ -246,7 +276,7 @@ def rebuild(
         dataset: Optional[Dataset] = None,
         parameters: Union[Parameters, dict] = None,
         check_model: bool = True,
-        initial_state: Optional[float] = None,
+        initial_state: Optional[dict] = None,
     ) -> None:
         """
         Rebuild the PyBaMM model for a given set of inputs.
@@ -264,10 +294,10 @@ def rebuild(
             A pybop Parameters class or dictionary containing parameter values to apply to the model.
         check_model : bool, optional
             If True, the model will be checked for correctness after construction.
-        initial_state : float or str, optional
-            If float, this value is used as the initial state of charge. If str ending in "V", this
-            value is used as the initial open-circuit voltage. Defaults to None, indicating that the
-            initial concentrations in the parameter set should be used.
+        initial_state : dict, optional
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+            Defaults to None, indicating that the existing initial state of charge (for an ECM)
+            or initial concentrations (for an EChem model) will be used.
         """
         if parameters is not None:
             self.classify_and_update_parameters(parameters)
@@ -376,7 +406,7 @@ def step(self, state: TimeSeriesState, time: np.ndarray) -> TimeSeriesState:
         return TimeSeriesState(sol=new_sol, inputs=state.inputs, t=time)
 
     def simulate(
-        self, inputs: Inputs, t_eval: np.array, initial_state: Optional[float] = None
+        self, inputs: Inputs, t_eval: np.array, initial_state: Optional[dict] = None
     ) -> Union[pybamm.Solution, list[np.float64]]:
         """
         Execute the forward model simulation and return the result.
@@ -387,10 +417,10 @@ def simulate(
             The input parameters for the simulation.
         t_eval : array-like
             An array of time points at which to evaluate the solution.
-        initial_state : float or str, optional
-            If float, this value is used as the initial state of charge (as a decimal between 0
-            and 1). If str ending in "V", this value is used as the initial open-circuit voltage.
-            Defaults to None, indicating that the existing initial concentrations will be used.
+        initial_state : dict, optional
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+            Defaults to None, indicating that the existing initial state of charge (for an ECM)
+            or initial concentrations (for an EChem model) will be used.
 
         Returns
         -------
@@ -437,7 +467,7 @@ def simulate(
             return [np.inf]
 
     def simulateS1(
-        self, inputs: Inputs, t_eval: np.array, initial_state: Optional[float] = None
+        self, inputs: Inputs, t_eval: np.array, initial_state: Optional[dict] = None
     ):
         """
         Perform the forward model simulation with sensitivities.
@@ -449,10 +479,10 @@ def simulateS1(
         t_eval : array-like
             An array of time points at which to evaluate the solution and its
             sensitivities.
-        initial_state : float or str, optional
-            If float, this value is used as the initial state of charge (as a decimal between 0
-            and 1). If str ending in "V", this value is used as the initial open-circuit voltage.
-            Defaults to None, indicating that the existing initial concentrations will be used.
+        initial_state : dict, optional
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+            Defaults to None, indicating that the existing initial state of charge (for an ECM)
+            or initial concentrations (for an EChem model) will be used.
 
         Returns
         -------
@@ -499,7 +529,7 @@ def predict(
         t_eval: Optional[np.array] = None,
         parameter_set: Optional[ParameterSet] = None,
         experiment: Optional[Experiment] = None,
-        initial_state: Optional[float] = None,
+        initial_state: Optional[dict] = None,
     ) -> dict[str, np.ndarray[np.float64]]:
         """
         Solve the model using PyBaMM's simulation framework and return the solution.
@@ -522,10 +552,10 @@ def predict(
         experiment : pybamm.Experiment, optional
             A PyBaMM Experiment object specifying the experimental conditions under which
             the simulation should be run. Defaults to None, indicating no experiment.
-        initial_state : float or str, optional
-            If float, this value is used as the initial state of charge (as a decimal between 0
-            and 1). If str ending in "V", this value is used as the initial open-circuit voltage.
-            Defaults to None, indicating that the existing initial concentrations will be used.
+        initial_state : dict, optional
+            A valid initial state, e.g. the initial state of charge or open-circuit voltage.
+            Defaults to None, indicating that the existing initial state of charge (for an ECM)
+            or initial concentrations (for an EChem model) will be used.
 
         Returns
         -------
@@ -553,12 +583,13 @@ def predict(
             parameter_set.update(inputs)
 
         if initial_state is not None:
+            # Update the default initial state for consistency
+            self.set_initial_state(initial_state)
+
+            initial_state = self.convert_to_pybamm_initial_state(initial_state)
             if isinstance(self.pybamm_model, pybamm.equivalent_circuit.Thevenin):
-                parameter_set["Initial SoC"] = initial_state
+                parameter_set["Initial SoC"] = self._parameter_set["Initial SoC"]
                 initial_state = None
-            else:
-                # Update the default initial state just for consistency
-                self.set_initial_state(initial_state)
 
         if self.check_params(
             parameter_set=parameter_set,