Refactor submodels to use a single build method with no inputs

docs/source/examples/notebooks/simulations_and_experiments/custom-experiments.ipynb

@@ -41,14 +41,32 @@
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "<pybamm.solvers.solution.Solution at 0x16b7c6310>"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
+     "ename": "IndexError",
+     "evalue": "list index out of range",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:615\u001b[0m, in \u001b[0;36mParameterValues.process_symbol\u001b[0;34m(self, symbol)\u001b[0m\n\u001b[1;32m    614\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 615\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_processed_symbols\u001b[49m\u001b[43m[\u001b[49m\u001b[43msymbol\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m    616\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n",
+      "\u001b[0;31mKeyError\u001b[0m: Addition(-0x18abcaa58950fbd0, +, children=['-0.02', 'Anode potential [V]'], domains={})",
+      "\nDuring handling of the above exception, another exception occurred:\n",
+      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:615\u001b[0m, in \u001b[0;36mParameterValues.process_symbol\u001b[0;34m(self, symbol)\u001b[0m\n\u001b[1;32m    614\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 615\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_processed_symbols\u001b[49m\u001b[43m[\u001b[49m\u001b[43msymbol\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m    616\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n",
+      "\u001b[0;31mKeyError\u001b[0m: CoupledVariable(0x576128e6b8f0fabb, Anode potential [V], children=[], domains={})",
+      "\nDuring handling of the above exception, another exception occurred:\n",
+      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[2], line 44\u001b[0m\n\u001b[1;32m     41\u001b[0m sim \u001b[38;5;241m=\u001b[39m pybamm\u001b[38;5;241m.\u001b[39mSimulation(model, parameter_values\u001b[38;5;241m=\u001b[39mparameter_values, experiment\u001b[38;5;241m=\u001b[39mexperiment)\n\u001b[1;32m     43\u001b[0m \u001b[38;5;66;03m# for a charge we start as SOC 0\u001b[39;00m\n\u001b[0;32m---> 44\u001b[0m \u001b[43msim\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msolve\u001b[49m\u001b[43m(\u001b[49m\u001b[43minitial_soc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/simulation.py:541\u001b[0m, in \u001b[0;36mSimulation.solve\u001b[0;34m(self, t_eval, solver, save_at_cycles, calc_esoh, starting_solution, initial_soc, callbacks, showprogress, inputs, t_interp, **kwargs)\u001b[0m\n\u001b[1;32m    539\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moperating_mode \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwith experiment\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m    540\u001b[0m     callbacks\u001b[38;5;241m.\u001b[39mon_experiment_start(logs)\n\u001b[0;32m--> 541\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbuild_for_experiment\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    542\u001b[0m \u001b[43m        \u001b[49m\u001b[43minitial_soc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minitial_soc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msolve_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    543\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    544\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m t_eval \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    545\u001b[0m         pybamm\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mwarning(\n\u001b[1;32m    546\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIgnoring t_eval as solution times are specified by the experiment\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    547\u001b[0m         )\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/simulation.py:358\u001b[0m, in \u001b[0;36mSimulation.build_for_experiment\u001b[0;34m(self, initial_soc, inputs, solve_kwargs)\u001b[0m\n\u001b[1;32m    356\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m    357\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 358\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mset_up_and_parameterise_experiment\u001b[49m\u001b[43m(\u001b[49m\u001b[43msolve_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    360\u001b[0m     \u001b[38;5;66;03m# Can process geometry with default parameter values (only electrical\u001b[39;00m\n\u001b[1;32m    361\u001b[0m     \u001b[38;5;66;03m# parameters change between parameter values)\u001b[39;00m\n\u001b[1;32m    362\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_parameter_values\u001b[38;5;241m.\u001b[39mprocess_geometry(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_geometry)\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/simulation.py:239\u001b[0m, in \u001b[0;36mSimulation.set_up_and_parameterise_experiment\u001b[0;34m(self, solve_kwargs)\u001b[0m\n\u001b[1;32m    237\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mexperiment_unique_steps_to_model \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m    238\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mexperiment\u001b[38;5;241m.\u001b[39munique_steps:\n\u001b[0;32m--> 239\u001b[0m     parameterised_model \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_model\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameter_values\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    240\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mexperiment_unique_steps_to_model[step\u001b[38;5;241m.\u001b[39mbasic_repr()] \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    241\u001b[0m         parameterised_model\n\u001b[1;32m    242\u001b[0m     )\n\u001b[1;32m    244\u001b[0m \u001b[38;5;66;03m# Set up rest model if experiment has start times\u001b[39;00m\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/experiment/step/base_step.py:358\u001b[0m, in \u001b[0;36mBaseStep.process_model\u001b[0;34m(self, model, parameter_values)\u001b[0m\n\u001b[1;32m    355\u001b[0m     new_parameter_values[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAmbient temperature [K]\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtemperature\n\u001b[1;32m    357\u001b[0m \u001b[38;5;66;03m# Parameterise the model\u001b[39;00m\n\u001b[0;32m--> 358\u001b[0m parameterised_model \u001b[38;5;241m=\u001b[39m \u001b[43mnew_parameter_values\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    359\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnew_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minplace\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\n\u001b[1;32m    360\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    362\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m parameterised_model\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:490\u001b[0m, in \u001b[0;36mParameterValues.process_model\u001b[0;34m(self, unprocessed_model, inplace)\u001b[0m\n\u001b[1;32m    486\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m event \u001b[38;5;129;01min\u001b[39;00m unprocessed_model\u001b[38;5;241m.\u001b[39mevents:\n\u001b[1;32m    487\u001b[0m     pybamm\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mverbose(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mProcessing parameters for event \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mevent\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    488\u001b[0m     new_events\u001b[38;5;241m.\u001b[39mappend(\n\u001b[1;32m    489\u001b[0m         pybamm\u001b[38;5;241m.\u001b[39mEvent(\n\u001b[0;32m--> 490\u001b[0m             event\u001b[38;5;241m.\u001b[39mname, \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_symbol\u001b[49m\u001b[43m(\u001b[49m\u001b[43mevent\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexpression\u001b[49m\u001b[43m)\u001b[49m, event\u001b[38;5;241m.\u001b[39mevent_type\n\u001b[1;32m    491\u001b[0m         )\n\u001b[1;32m    492\u001b[0m     )\n\u001b[1;32m    494\u001b[0m interpolant_events \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_interpolant_events(model)\n\u001b[1;32m    495\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m event \u001b[38;5;129;01min\u001b[39;00m interpolant_events:\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:617\u001b[0m, in \u001b[0;36mParameterValues.process_symbol\u001b[0;34m(self, symbol)\u001b[0m\n\u001b[1;32m    615\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_processed_symbols[symbol]\n\u001b[1;32m    616\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[0;32m--> 617\u001b[0m     processed_symbol \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_process_symbol\u001b[49m\u001b[43m(\u001b[49m\u001b[43msymbol\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    618\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_processed_symbols[symbol] \u001b[38;5;241m=\u001b[39m processed_symbol\n\u001b[1;32m    620\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m processed_symbol\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:756\u001b[0m, in \u001b[0;36mParameterValues._process_symbol\u001b[0;34m(self, symbol)\u001b[0m\n\u001b[1;32m    750\u001b[0m \u001b[38;5;66;03m# Functions, BinaryOperators & Concatenations\u001b[39;00m\n\u001b[1;32m    751\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m (\n\u001b[1;32m    752\u001b[0m     \u001b[38;5;28misinstance\u001b[39m(symbol, pybamm\u001b[38;5;241m.\u001b[39mFunction)\n\u001b[1;32m    753\u001b[0m     \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(symbol, pybamm\u001b[38;5;241m.\u001b[39mConcatenation)\n\u001b[1;32m    754\u001b[0m     \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(symbol, pybamm\u001b[38;5;241m.\u001b[39mBinaryOperator)\n\u001b[1;32m    755\u001b[0m ):\n\u001b[0;32m--> 756\u001b[0m     new_children \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_symbol\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchild\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m child \u001b[38;5;129;01min\u001b[39;00m symbol\u001b[38;5;241m.\u001b[39mchildren]\n\u001b[1;32m    757\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m symbol\u001b[38;5;241m.\u001b[39mcreate_copy(new_children)\n\u001b[1;32m    759\u001b[0m \u001b[38;5;66;03m# Variables: update scale\u001b[39;00m\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:617\u001b[0m, in \u001b[0;36mParameterValues.process_symbol\u001b[0;34m(self, symbol)\u001b[0m\n\u001b[1;32m    615\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_processed_symbols[symbol]\n\u001b[1;32m    616\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[0;32m--> 617\u001b[0m     processed_symbol \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_process_symbol\u001b[49m\u001b[43m(\u001b[49m\u001b[43msymbol\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    618\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_processed_symbols[symbol] \u001b[38;5;241m=\u001b[39m processed_symbol\n\u001b[1;32m    620\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m processed_symbol\n",
+      "File \u001b[0;32m~/pybamm/new-submodel-structure/pybamm/src/pybamm/parameters/parameter_values.py:776\u001b[0m, in \u001b[0;36mParameterValues._process_symbol\u001b[0;34m(self, symbol)\u001b[0m\n\u001b[1;32m    773\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m pybamm\u001b[38;5;241m.\u001b[39mScalar(symbol)\n\u001b[1;32m    775\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(symbol, pybamm\u001b[38;5;241m.\u001b[39mCoupledVariable):\n\u001b[0;32m--> 776\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_symbol(\u001b[43msymbol\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mchildren\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m)\n\u001b[1;32m    778\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    779\u001b[0m     \u001b[38;5;66;03m# Backup option: return the object\u001b[39;00m\n\u001b[1;32m    780\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m symbol\n",
+      "\u001b[0;31mIndexError\u001b[0m: list index out of range"
+     ]
     }
    ],
    "source": [
@@ -66,7 +84,7 @@
     "# Create a custom termination event for the anode potential cut-off at 0.02V\n",
     "# We use 0.02V instead of 0V to give a safety factor\n",
     "def anode_potential_cutoff(variables):\n",
-    "    return variables[\"Anode potential [V]\"] - 0.02\n",
+    "    return pybamm.CoupledVariable(\"Anode potential [V]\") - 0.02\n",
     "\n",
     "\n",
     "# The CustomTermination class takes a name and function\n",
@@ -472,7 +490,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

src/pybamm/discretisations/discretisation.py

@@ -688,6 +688,8 @@ def process_dict(self, var_eqn_dict, ics=False):
                     eqn = pybamm.FullBroadcast(eqn, broadcast_domains=eqn_key.domains)
 
             pybamm.logger.debug(f"Discretise {eqn_key!r}")
+            if eqn_key == "Electrolyte concentration concatenation [mol.m-3]":
+                print("hi")
             processed_eqn = self.process_symbol(eqn)
             # Calculate scale if the key has a scale
             scale = getattr(eqn_key, "scale", 1)

src/pybamm/experiment/step/base_step.py

@@ -4,6 +4,9 @@
 import pybamm
 import numpy as np
 from datetime import datetime
+
+import pybamm.expression_tree
+import pybamm.expression_tree.coupled_variable
 from .step_termination import _read_termination
 import numbers
 
@@ -349,6 +352,7 @@ def process_model(self, model, parameter_values):
         new_parameter_values = parameter_values.copy()
         new_model, new_parameter_values = self.set_up(new_model, new_parameter_values)
         self.update_model_events(new_model)
+        new_model.link_coupled_variables()
 
         # Update temperature
         if self.temperature is not None:
@@ -364,6 +368,11 @@ def process_model(self, model, parameter_values):
     def update_model_events(self, new_model):
         for term in self.termination:
             event = term.get_event(new_model.variables, self)
+            if event is not None:
+                coupled_variables_from_event = pybamm.expression_tree.coupled_variable.find_and_save_coupled_variables(
+                    event.expression
+                )
+                new_model.coupled_variables.update(coupled_variables_from_event)
             if event is not None:
                 new_model.events.append(event)
 
@@ -440,6 +449,7 @@ def set_up(self, new_model, new_parameter_values):
         new_parameter_values["Current function [A]"] = self.current_value(
             new_model.variables
         )
+        new_model.link_coupled_variables()
         return new_model, new_parameter_values
 
 
@@ -463,19 +473,17 @@ def set_up(self, new_model, new_parameter_values):
         # Build the new submodel and update the model with it
         submodel = self.get_submodel(new_model)
         variables = new_model.variables
-        submodel.variables = submodel.get_fundamental_variables()
+        coupled_variables = new_model.coupled_variables
+        submodel.build(new_model.submodels)
         variables.update(submodel.variables)
-        submodel.variables.update(submodel.get_coupled_variables(variables))
-        variables.update(submodel.variables)
-        submodel.set_rhs(variables)
-        submodel.set_algebraic(variables)
-        submodel.set_initial_conditions(variables)
+        coupled_variables.update(submodel.coupled_variables)
         new_model.rhs.update(submodel.rhs)
         new_model.algebraic.update(submodel.algebraic)
         new_model.initial_conditions.update(submodel.initial_conditions)
-
-        # Set the "current function" to be the variable defined in the submodel
+        new_model.boundary_conditions.update(submodel.boundary_conditions)
         new_parameter_values["Current function [A]"] = submodel.variables["Current [A]"]
+        new_model.link_coupled_variables()
+
         # Update any other parameters as necessary
         new_parameter_values.update(
             self.get_parameter_values(variables), check_already_exists=False

src/pybamm/experiment/step/steps.py

@@ -410,9 +410,14 @@ def __init__(self, current_rhs_function, control="algebraic", **kwargs):
         self.kwargs = kwargs
 
     def get_submodel(self, model):
-        return pybamm.external_circuit.FunctionControl(
+        cv = pybamm.expression_tree.coupled_variable.find_and_save_coupled_variables(
+            self.current_rhs_function(None)
+        )
+        submodel = pybamm.external_circuit.FunctionControl(
             model.param, self.current_rhs_function, model.options, control=self.control
         )
+        submodel.coupled_variables = cv
+        return submodel
 
     def copy(self):
         return CustomStepImplicit(