Repository: pybamm-team/PyBaMM (python)

Issue:
[Bug]: model.set_initial_conditions_from assumes state variable names consistent with output variables
### PyBaMM Version

develop

### Python Version

n/a

### Describe the bug

the `BaseModel.set_initial_conditions` function extracts the ic values from a `pybamm.Solution` object. It looks in the output variables of the model to find each variable in turn, then concatenates a new state vector from the evaluation of those output variables.

This works fine with our library models, where for each variable `a = pybamm.Variable('A')`, we have defined an output variable `model.variables = { 'A': a}`. But third party models do not neccessarily have this structure, so the   `BaseModel.set_initial_conditions` function can easily fail when this third party model is used in an experiment

Perhaps instead we could:
1. disallow any output variables of a model with the same name as a model variable
2. automatically create output variables for each variable when model is discretised

### Steps to Reproduce

1. create a model that has a variable with name `A`, and an output variable with the same name but a different shape
2. use this model in an experiment

### Relevant log output

```shell

```

Pull request description:
feat: auto-add state variables to output variables
# Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes #4662

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

- [x] New feature (non-breaking change which adds functionality)
- [x] Bug fix (non-breaking change which fixes an issue)

# Key checklist:

- [x] No style issues: `$ pre-commit run` (or `$ nox -s pre-commit`) (see [CONTRIBUTING.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CONTRIBUTING.md#installing-and-using-pre-commit) for how to set this up to run automatically when committing locally, in just two lines of code)
- [x] All tests pass: `$ python run-tests.py --all` (or `$ nox -s tests`)
- [x] The documentation builds: `$ python run-tests.py --doctest` (or `$ nox -s doctests`)

You can run integration tests, unit tests, and doctests together at once, using `$ python run-tests.py --quick` (or `$ nox -s quick`).

## Further checks:

- [x] Code is commented, particularly in hard-to-understand areas
- [x] Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7a9e25ff..f03ea5d4e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,7 @@
 - Made composite electrode model compatible with particle size distribution ([#4687](https://github.com/pybamm-team/PyBaMM/pull/4687))
 - Added `Symbol.post_order()` method to return an iterable that steps through the tree in post-order fashion. ([#4684](https://github.com/pybamm-team/PyBaMM/pull/4684))
 - Added two more submodels (options) for the SEI: Lars von Kolzenberg (2020) model and Tunneling Limit model ([#4394](https://github.com/pybamm-team/PyBaMM/pull/4394))
-
+- Automatically add state variables of the model to the output variables if they are not already present ([#4700](https://github.com/pybamm-team/PyBaMM/pull/4700))
 
 ## Breaking changes
 
diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py
index 3d9579ff9..1a70f1dd5 100644
--- a/src/pybamm/discretisations/discretisation.py
+++ b/src/pybamm/discretisations/discretisation.py
@@ -193,6 +193,11 @@ class Discretisation:
 
         model_disc.bcs = self.bcs
 
+        # pre-process variables so that all state variables are included
+        pre_processed_variables = self._pre_process_variables(
+            model.variables, model.initial_conditions
+        )
+
         pybamm.logger.verbose(f"Discretise initial conditions for {model.name}")
         ics, concat_ics = self.process_initial_conditions(model)
         model_disc.initial_conditions = ics
@@ -202,7 +207,8 @@ class Discretisation:
         # Note that we **do not** discretise the keys of model.rhs,
         # model.initial_conditions and model.boundary_conditions
         pybamm.logger.verbose(f"Discretise variables for {model.name}")
-        model_disc.variables = self.process_dict(model.variables)
+
+        model_disc.variables = self.process_dict(pre_processed_variables)
 
         # Process parabolic and elliptic equations
         pybamm.logger.verbose(f"Discretise model equations for {model.name}")
@@ -657,6 +663,46 @@ class Discretisation:
 
         return mass_matrix, mass_matrix_inv
 
+    def _pre_process_variables(
+        self,
+        variables: dict[str, pybamm.Symbol],
+        initial_conditions: dict[pybamm.Variable, pybamm.Symbol],
+    ):
+        """
+        Pre-process variables before discretisation. This involves:
+        - ensuring that all the state variables are included in the variables,
+          any missing are added
+
+        Parameters
+        ----------
+        variables : dict
+            Dictionary of variables to pre-process
+        initial_conditions : dict
+            Dictionary of initial conditions
+
+        Returns
+        -------
+        dict
+            Pre-processed variables (copy of input variables with any missing state)
+
+        Raises
+        ------
+        :class:`pybamm.ModelError`
+            If any state variable names are already included but with
+            incorrect expressions
+        """
+        new_variables = {k: v for k, v in variables.items()}
+        for var in initial_conditions.keys():
+            if var.name not in new_variables:
+                new_variables[var.name] = var
+            else:
+                if new_variables[var.name] != var:
+                    raise pybamm.ModelError(
+                        f"Variable '{var.name}' should have expression "
+                        f"'{var}', but has expression '{new_variables[var.name]}'"
+                    )
+        return new_variables
+
     def process_dict(self, var_eqn_dict, ics=False):
         """Discretise a dictionary of {variable: equation}, broadcasting if necessary
         (can be model.rhs, model.algebraic, model.initial_conditions or
@@ -1008,7 +1054,6 @@ class Discretisation:
     def check_model(self, model):
         """Perform some basic checks to make sure the discretised model makes sense."""
         self.check_initial_conditions(model)
-        self.check_variables(model)
 
     def check_initial_conditions(self, model):
         # Check initial conditions are a numpy array
@@ -1049,40 +1094,6 @@ class Discretisation:
                     f"{model.algebraic[var].shape} and initial_conditions.shape = {model.initial_conditions[var].shape} for variable '{var}'."
                 )
 
-    def check_variables(self, model):
-        """
-        Check variables in variable list against rhs.
-        Be lenient with size check if the variable in model.variables is broadcasted, or
-        a concatenation
-        (if broadcasted, variable is a multiplication with a vector of ones)
-        """
-        for rhs_var in model.rhs.keys():
-            if rhs_var.name in model.variables.keys():
-                var = model.variables[rhs_var.name]
-
-                different_shapes = not np.array_equal(
-                    model.rhs[rhs_var].shape, var.shape
-                )
-
-                not_concatenation = not isinstance(var, pybamm.Concatenation)
-
-                not_mult_by_one_vec = not (
-                    isinstance(
-                        var, (pybamm.Multiplication, pybamm.MatrixMultiplication)
-                    )
-                    and (
-                        pybamm.is_matrix_one(var.left)
-                        or pybamm.is_matrix_one(var.right)
-                    )
-                )
-
-                if different_shapes and not_concatenation and not_mult_by_one_vec:
-                    raise pybamm.ModelError(
-                        "variable and its eqn must have the same shape after "
-                        "discretisation but variable.shape = "
-                        f"{var.shape} and rhs.shape = {model.rhs[rhs_var].shape} for variable '{var}'. "
-                    )
-
     def is_variable_independent(self, var, all_vars_in_eqns):
         pybamm.logger.verbose("Removing independent blocks.")
         if not isinstance(var, pybamm.Variable):
diff --git a/src/pybamm/expression_tree/concatenations.py b/src/pybamm/expression_tree/concatenations.py
index a1905f908..62adac526 100644
--- a/src/pybamm/expression_tree/concatenations.py
+++ b/src/pybamm/expression_tree/concatenations.py
@@ -4,6 +4,7 @@
 from __future__ import annotations
 import copy
 from collections import defaultdict
+from typing import Optional
 
 import numpy as np
 import sympy
@@ -146,9 +147,9 @@ class Concatenation(pybamm.Symbol):
         children before creating the new copy.
         """
         if perform_simplifications:
-            return concatenation(*children)
+            return concatenation(*children, name=self.name)
         else:
-            return self.__class__(*children)
+            return self.__class__(*children, name=self.name)
 
     def _concatenation_jac(self, children_jacs):
         """Calculate the Jacobian of a concatenation."""
@@ -468,17 +469,18 @@ class SparseStack(Concatenation):
 class ConcatenationVariable(Concatenation):
     """A Variable representing a concatenation of variables."""
 
-    def __init__(self, *children):
-        # Name is the intersection of the children names (should usually make sense
-        # if the children have been named consistently)
-        name = intersect(children[0].name, children[1].name)
-        for child in children[2:]:
-            name = intersect(name, child.name)
-        if len(name) == 0:
-            name = None
-        # name is unchanged if its length is 1
-        elif len(name) > 1:
-            name = name[0].capitalize() + name[1:]
+    def __init__(self, *children, name: Optional[str] = None):
+        if name is None:
+            # Name is the intersection of the children names (should usually make sense
+            # if the children have been named consistently)
+            name = intersect(children[0].name, children[1].name)
+            for child in children[2:]:
+                name = intersect(name, child.name)
+            if len(name) == 0:
+                name = None
+            # name is unchanged if its length is 1
+            elif len(name) > 1:
+                name = name[0].capitalize() + name[1:]
 
         if len(children) > 0:
             if all(child.scale == children[0].scale for child in children):
@@ -523,7 +525,7 @@ def intersect(s1: str, s2: str):
     return intersect.lstrip().rstrip()
 
 
-def simplified_concatenation(*children):
+def simplified_concatenation(*children, name: Optional[str] = None):
     """Perform simplifications on a concatenation."""
     # remove children that are None
     children = list(filter(lambda x: x is not None, children))
@@ -534,29 +536,29 @@ def simplified_concatenation(*children):
     elif len(children) == 1:
         return children[0]
     elif all(isinstance(child, pybamm.Variable) for child in children):
-        return pybamm.ConcatenationVariable(*children)
+        return pybamm.ConcatenationVariable(*children, name=name)
     else:
         # Create Concatenation to easily read domains
-        concat = Concatenation(*children)
+        concat = Concatenation(*children, name=name)
         if all(
             isinstance(child, pybamm.Broadcast) and child.child == children[0].child
             for child in children
         ):
             unique_child = children[0].orphans[0]
             if isinstance(children[0], pybamm.PrimaryBroadcast):
-                return pybamm.PrimaryBroadcast(unique_child, concat.domain)
+                return pybamm.PrimaryBroadcast(unique_child, concat.domain, name=name)
             else:
                 return pybamm.FullBroadcast(
-                    unique_child, broadcast_domains=concat.domains
+                    unique_child, broadcast_domains=concat.domains, name=name
                 )
         else:
             return concat
 
 
-def concatenation(*children):
+def concatenation(*children, name: Optional[str] = None):
     """Helper function to create concatenations."""
     # TODO: add option to turn off simplifications
-    return simplified_concatenation(*children)
+    return simplified_concatenation(*children, name=name)
 
 
 def simplified_numpy_concatenation(*children):
diff --git a/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py b/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py
index 7ecad6fa4..4066dcfed 100644
--- a/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py
+++ b/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py
@@ -36,7 +36,11 @@ class Full(BaseModel):
             domain="positive electrode",
             auxiliary_domains={"secondary": "current collector"},
         )
-        c_ox_s_p = pybamm.concatenation(c_ox_s, c_ox_p)
+        c_ox_s_p = pybamm.concatenation(
+            c_ox_s,
+            c_ox_p,
+            name="Separator and positive electrode oxygen concentration [mol.m-3]",
+        )
         variables = {
             "Separator and positive electrode oxygen concentration [mol.m-3]": c_ox_s_p
         }


Test patch:
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index 55d6546bd..270a2eaa3 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -1036,9 +1036,9 @@ class TestDiscretise:
         assert expr.children[2].evaluate(0, y).shape == (105, 1)
 
     def test_exceptions(self):
-        c_n = pybamm.Variable("c", domain=["negative electrode"])
+        c_n = pybamm.Variable("c_n", domain=["negative electrode"])
         N_n = pybamm.grad(c_n)
-        c_s = pybamm.Variable("c", domain=["separator"])
+        c_s = pybamm.Variable("c_s", domain=["separator"])
         N_s = pybamm.grad(c_s)
         model = pybamm.BaseModel()
         model.rhs = {c_n: pybamm.div(N_n), c_s: pybamm.div(N_s)}
@@ -1049,22 +1049,6 @@ class TestDiscretise:
         }
 
         disc = get_discretisation_for_testing()
-
-        # check raises error if different sized key and output var
-        model.variables = {c_n.name: c_s}
-        with pytest.raises(pybamm.ModelError, match="variable and its eqn"):
-            disc.process_model(model)
-
-        # check doesn't raise if concatenation
-        model.variables = {c_n.name: pybamm.concatenation(2 * c_n, 3 * c_s)}
-        disc.process_model(model, inplace=False)
-
-        # check doesn't raise if broadcast
-        model.variables = {
-            c_n.name: pybamm.PrimaryBroadcast(
-                pybamm.InputParameter("a"), ["negative electrode"]
-            )
-        }
         disc.process_model(model)
 
         # Check setting up a 0D spatial method with 1D mesh raises error
@@ -1277,3 +1261,29 @@ class TestDiscretise:
         disc = pybamm.Discretisation(remove_independent_variables_from_rhs=True)
         disc.process_model(model)
         assert len(model.rhs) == 3
+
+    def test_pre_process_variables(self):
+        a = pybamm.Variable("a")
+        b = pybamm.Variable("b")
+        model = pybamm.BaseModel()
+        model.rhs = {a: b, b: a}
+        model.initial_conditions = {
+            a: pybamm.Scalar(0),
+            b: pybamm.Scalar(1),
+        }
+        model.variables = {
+            "a": a,  # correct
+            # b missing
+        }
+        disc = pybamm.Discretisation()
+        disc_model = disc.process_model(model, inplace=False)
+        assert list(disc_model.variables.keys()) == ["a", "b"]
+
+        model.variables = {
+            "a": a,
+            "b": 2 * a,
+        }
+        with pytest.raises(
+            pybamm.ModelError, match="Variable 'b' should have expression"
+        ):
+            disc.process_model(model, inplace=False)