Made y, z optional arguments for FunctionParameters that depend on space

CHANGELOG.md

@@ -13,6 +13,7 @@
 
 - Added `skip_ok` option to `step` to allow for steps to be skipped if they are infeasible at initial conditions. ([#4839](https://github.com/pybamm-team/PyBaMM/pull/4839))
 - Deprecated `CrateTermination` and renamed it to `CRateTermination`. ([#4834](https://github.com/pybamm-team/PyBaMM/pull/4834))
+- Made y, z optional arguments for FunctionParameters that depend on space ([#4883](https://github.com/pybamm-team/PyBaMM/pull/4883))
 
 ## Bug fixes
 

examples/scripts/DFN_ambient_temperature.py

@@ -18,7 +18,7 @@
 # load parameter values and process model and geometry
 
 
-def ambient_temperature(y, z, t):
+def ambient_temperature(t, y=None, z=None):
     return 300 + t * 100 / 3600
 
 

examples/scripts/pouch_cell_cooling.py

@@ -20,7 +20,7 @@
 parameter_values = model.default_parameter_values
 
 
-def T_amb(y, z, t):
+def T_amb(t, y=None, z=None):
     return 300 + 20 * pybamm.sin(np.pi * y / L_y) * pybamm.sin(np.pi * z / L_z)
 
 

src/pybamm/models/submodels/thermal/base_thermal.py

@@ -45,7 +45,7 @@ def _get_standard_fundamental_variables(self, T_dict):
         # (y, z) only and time
         y = pybamm.standard_spatial_vars.y
         z = pybamm.standard_spatial_vars.z
-        T_amb = self.param.T_amb(y, z, pybamm.t)
+        T_amb = self.param.T_amb(pybamm.t, y, z)
         T_amb_av = self._yz_average(T_amb)
 
         variables = {

src/pybamm/models/submodels/thermal/isothermal.py

@@ -29,7 +29,7 @@ def get_fundamental_variables(self):
         # Broadcast t to be the same size as y and z (to catch cases where the ambient
         # temperature is a function of time only)
         t_broadcast = pybamm.PrimaryBroadcast(pybamm.t, "current collector")
-        T_x_av = self.param.T_amb(y, z, t_broadcast)
+        T_x_av = self.param.T_amb(t_broadcast, y, z)
         T_vol_av = self._yz_average(T_x_av)
 
         T_dict = {

src/pybamm/parameters/thermal_parameters.py

@@ -40,22 +40,20 @@ def _set_parameters(self):
         # Initial temperature
         self.T_init = pybamm.Parameter("Initial temperature [K]")
 
-    def T_amb(self, y, z, t):
+    def T_amb(self, t, y=None, z=None):
         """Ambient temperature [K]"""
-        return pybamm.FunctionParameter(
-            "Ambient temperature [K]",
-            {
-                "Distance across electrode width [m]": y,
-                "Distance across electrode height [m]": z,
-                "Time [s]": t,
-            },
-        )
+        inputs = {"Time [s]": t}
+        if y is not None:
+            inputs["Distance across electrode width [m]"] = y
+        if z is not None:
+            inputs["Distance across electrode height [m]"] = z
+        return pybamm.FunctionParameter("Ambient temperature [K]", inputs)
 
     def T_amb_av(self, t):
         """YZ-averaged ambient temperature [K]"""
         y = pybamm.standard_spatial_vars.y
         z = pybamm.standard_spatial_vars.z
-        return pybamm.yz_average(self.T_amb(y, z, t))
+        return pybamm.yz_average(self.T_amb(t, y, z))
 
     def h_edge(self, y, z):
         """Cell edge heat transfer coefficient [W.m-2.K-1]"""

tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py

@@ -376,7 +376,7 @@ def test_basic_processing_temperature_interpolant(self):
         times = np.arange(0, 4000, 10)
         tmax = max(times)
 
-        def temp_drive_cycle(y, z, t):
+        def temp_drive_cycle(t, y=None, z=None):
             return pybamm.Interpolant(
                 times,
                 298.15 + 20 * (times / tmax),

tests/unit/test_experiments/test_experiment.py

@@ -355,3 +355,60 @@ def test_pchip_extrapolation(self):
         result = np.array(f(test_points)).flatten()
 
         np.testing.assert_allclose(result, expected, rtol=1e-7, atol=1e-6)
+
+    def test_ambient_temperature_interpolant_in_lumped_model(self):
+        options = {"thermal": "lumped"}
+        model = pybamm.lithium_ion.DFN(options)
+
+        geometry = model.default_geometry
+
+        times = np.arange(0, 1810, 10)
+        tmax = times[-1]
+
+        def ambient_temperature(t, y=None, z=None):
+            return pybamm.Interpolant(times, 298.15 + 20 * (times / tmax), pybamm.t)
+
+        param = model.default_parameter_values
+        param.update(
+            {"Ambient temperature [K]": ambient_temperature}, check_already_exists=False
+        )
+
+        param.process_model(model)
+        param.process_geometry(geometry)
+
+        var_pts = {"x_n": 30, "x_s": 30, "x_p": 30, "r_n": 10, "r_p": 10}
+        mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
+        disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+        disc.process_model(model)
+
+        t_eval = np.linspace(0, 1800, 100)
+        solver = pybamm.CasadiSolver(mode="fast", atol=1e-6, rtol=1e-3)
+        solution = solver.solve(model, t_eval)
+
+        amb_temp = solution["Ambient temperature [K]"]
+
+        t_sol = solution.t
+        computed = amb_temp(t_sol)
+
+        expected = 298.15 + 20 * (t_sol / tmax)
+
+        np.testing.assert_allclose(computed, expected, rtol=1e-2)
+
+    def test_ambient_function_required_spatial_params(self):
+        options = {"particle": "quadratic profile", "thermal": "lumped"}
+        model = pybamm.lithium_ion.SPMe(options)
+
+        def ambient_temperature(t, y=None, z=None):
+            return 300 + t * 100 / 3600 + 2 * y * z
+
+        param = pybamm.ParameterValues("Chen2020")
+        param.update(
+            {"Ambient temperature [K]": ambient_temperature}, check_already_exists=False
+        )
+
+        with pytest.raises(
+            ValueError, match=r"Dimensions .* do not exist\. Expected .*"
+        ):
+            sim = pybamm.Simulation(model, parameter_values=param)
+            solution = sim.solve([0, 3600])
+            solution["Ambient temperature [K]"](5, 2, 4)