allow custom initial state

src/pybamm/parameters/parameter_values.py

@@ -65,10 +65,15 @@ def __init__(self, values):
                     f"'{values}' is not a valid parameter set. Parameter set must be one of:\n{valid_sets}"
                 )
 
-        if chemistry == "ecm":
-            self._set_initial_state = pybamm.equivalent_circuit.set_initial_state
-        else:
-            self._set_initial_state = pybamm.lithium_ion.set_initial_state
+        if "Initial state function" not in self._dict_items:
+            if chemistry == "ecm":
+                self._dict_items["Initial state function"] = (
+                    pybamm.equivalent_circuit.set_initial_state
+                )
+            else:
+                self._dict_items["Initial state function"] = (
+                    pybamm.lithium_ion.set_initial_state
+                )
 
         # Initialise empty _processed_symbols dict (for caching)
         self._processed_symbols = {}
@@ -234,7 +239,6 @@ def copy(self):
         """Returns a copy of the parameter values. Makes sure to copy the internal
         dictionary."""
         new_copy = ParameterValues(self._dict_items.copy())
-        new_copy._set_initial_state = self._set_initial_state
         return new_copy
 
     def search(self, key, print_values=True):
@@ -992,3 +996,7 @@ def __contains__(self, key):
 
     def __iter__(self):
         return iter(self._dict_items)
+
+    @property
+    def _set_initial_state(self):
+        return self._dict_items["Initial state function"]

tests/unit/test_parameters/test_base_parameters.py

@@ -33,6 +33,27 @@ def test_getattr__(self):
         assert param.n.cap_init == param.n.Q_init
         assert param.p.prim.cap_init == param.p.prim.Q_init
 
+    def test_that_initial_state_function_is_assigned(self):
+        param_1 = pybamm.ParameterValues("Chen2020")
+        assert param_1._set_initial_state == pybamm.lithium_ion.set_initial_state
+
+        def my_func(x):
+            return x
+
+        param_1.update({"Initial state function": my_func})
+        assert param_1._set_initial_state == my_func
+        param_2 = pybamm.ParameterValues("ECM_Example")
+        assert param_2._set_initial_state == pybamm.equivalent_circuit.set_initial_state
+
+        def my_error_func(*args, **kwargs):
+            raise NotImplementedError("this function should error")
+
+        param_1.update({"Initial state function": my_error_func})
+        model = pybamm.lithium_ion.SPM()
+        sim = pybamm.Simulation(model, parameter_values=param_1)
+        with pytest.raises(NotImplementedError, match="this function should error"):
+            sim.solve([0, 1], initial_soc=0.5)
+
     def test__setattr__(self):
         # domain gets added as a subscript
         param = pybamm.GeometricParameters()

tests/unit/test_parameters/test_parameter_values.py

@@ -50,6 +50,7 @@ def test_repr(self):
             "Boltzmann constant [J.K-1]",
             "Electron charge [C]",
             "a",
+            "Initial state function",
         ]
 
     def test_eq(self):
@@ -1098,4 +1099,4 @@ def test_iter_method(self):
             values={"Negative particle radius [m]": 1e-6}
         )
         pv = [i for i in parameter_values]
-        assert len(pv) == 5, "Should have 5 keys"
+        assert len(pv) == 6, "Should have 6 keys"