Repository: pybamm-team/PyBaMM (python)

Issue:
DiscreteTimeSum as child node
### Description

At the moment the DiscreteTimeSum is restricted in its use. You can only use it as the root node in an expression tree, because in Solution/ProcessedVariable only the root node is looked at when detecting this type of node (along with the integral node).

This issue proposes to modify the processing so that you can embed the DiscreteTimeSum as a child node, for example a sum like the Minkowski sum which raises the sum to the power of 1/p

$$
( \sum_i f(x_i) )^1/p
$$$

### Motivation

_No response_

### Possible Implementation

_No response_

### Additional context

_No response_

Pull request description:
feat: generalise pybamm.DiscreteTimeSum so that it can be inside an expression
# Description

Fixes #5039

## 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 #)

# Important checks:

Please confirm the following before marking the PR as ready for review:
- No style issues: `nox -s pre-commit`
- All tests pass: `nox -s tests`
- The documentation builds: `nox -s doctests`
- Code is commented for hard-to-understand areas
- Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ecf1816ef..2e299dbd5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## Features
 
+- Generalise `pybamm.DiscreteTimeSum` to allow it to be embedded in other expressions ([#5044](https://github.com/pybamm-team/PyBaMM/pull/5044))
+
 ## Optimizations
 
 ## Bug fixes
diff --git a/src/pybamm/expression_tree/scalar.py b/src/pybamm/expression_tree/scalar.py
index 269e9c067..e1d5c2bf6 100644
--- a/src/pybamm/expression_tree/scalar.py
+++ b/src/pybamm/expression_tree/scalar.py
@@ -86,9 +86,12 @@ class Scalar(pybamm.Symbol):
         perform_simplifications=True,
     ):
         """See :meth:`pybamm.Symbol.new_copy()`."""
-        if new_children is not None:
-            raise ValueError("Cannot create a copy of a scalar with new children")
-        return Scalar(self.value, self.name)
+        if new_children is None or len(new_children) == 0:
+            return Scalar(self.value, self.name)
+        else:
+            raise ValueError(
+                f"Cannot create a copy of a scalar with new children: {new_children}"
+            )
 
     def is_constant(self) -> Literal[True]:
         """See :meth:`pybamm.Symbol.is_constant()`."""
diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py
index af717d1ac..92177c641 100644
--- a/src/pybamm/solvers/processed_variable.py
+++ b/src/pybamm/solvers/processed_variable.py
@@ -5,7 +5,6 @@ import casadi
 import numpy as np
 import xarray as xr
 from pybammsolvers import idaklu
-from scipy.integrate import cumulative_trapezoid
 
 import pybamm
 
@@ -479,7 +478,28 @@ class ProcessedVariable:
                         f"ensure the correct times are used.\nSolution times: {ts}\nDiscrete Sum times: {self.time_integral.discrete_times}"
                     )
 
-            S_var = self._observe_postfix(S_var, ts)
+            if self.time_integral is not None:
+                post_sum_values = self.time_integral.postfix_sum(S_var, ts)
+                post_sum_node = self.time_integral.post_sum_node
+                if post_sum_node is None:
+                    S_var = post_sum_values
+                else:
+                    # If there is a post-sum expression, we need to compute the
+                    # sensitivity of the post-sum expression as well
+                    y_casadi = casadi.MX.sym("y", post_sum_values.shape[0])
+                    s_var_casadi = casadi.MX.sym("s_var", post_sum_values.shape)
+                    post_sum_casadi = post_sum_node.to_casadi(
+                        t_casadi, y_casadi, inputs=p_casadi
+                    )
+                    dpost_dy = casadi.jacobian(post_sum_casadi, y_casadi)
+                    dpost_dp = casadi.jacobian(post_sum_casadi, p_casadi_stacked)
+                    sens = dpost_dy @ s_var_casadi + dpost_dp
+                    sens_fun = casadi.Function(
+                        "sens_fun",
+                        [t_casadi, y_casadi, p_casadi_stacked, s_var_casadi],
+                        [sens],
+                    )
+                    S_var = sens_fun(0.0, self.data, inputs_stacked, post_sum_values)
 
             all_S_var.append(S_var)
 
@@ -528,16 +548,9 @@ class ProcessedVariable0D(ProcessedVariable):
     def _observe_postfix(self, entries, t):
         if self.time_integral is None:
             return entries
-        if self.time_integral.method == "discrete":
-            return np.sum(entries, axis=0, initial=self.time_integral.initial_condition)
-        elif self.time_integral.method == "continuous":
-            return cumulative_trapezoid(
-                entries, self.t_pts, initial=float(self.time_integral.initial_condition)
-            )
-        else:
-            raise ValueError(
-                "time_integral method must be 'discrete' or 'continuous'"
-            )  # pragma: no cover
+        return self.time_integral.postfix(
+            entries, self.t_pts, self.all_inputs_casadi[0]
+        )
 
     def _interp_setup(self, entries, t):
         # save attributes for interpolation
diff --git a/src/pybamm/solvers/processed_variable_time_integral.py b/src/pybamm/solvers/processed_variable_time_integral.py
index c18bbb079..77f223afe 100644
--- a/src/pybamm/solvers/processed_variable_time_integral.py
+++ b/src/pybamm/solvers/processed_variable_time_integral.py
@@ -3,8 +3,10 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Literal
 
+import casadi
 import numpy as np
 import numpy.typing as npt
+from scipy.integrate import cumulative_trapezoid
 
 import pybamm
 
@@ -12,21 +14,95 @@ import pybamm
 @dataclass
 class ProcessedVariableTimeIntegral:
     method: Literal["discrete", "continuous"]
-    initial_condition: npt.NDArray[np.float64] | None | float
+    sum_node: pybamm.Symbol
+    initial_condition: npt.NDArray[np.float64] | float
     discrete_times: npt.NDArray[np.float64] | None
+    post_sum_node: pybamm.Symbol | None = None
+    post_sum: casadi.Function | None = None
+
+    def postfix_sum(self, entries, t_pts) -> np.ndarray:
+        if self.method == "discrete":
+            return np.sum(
+                entries, axis=0, initial=self.initial_condition, keepdims=True
+            )
+        else:
+            return cumulative_trapezoid(
+                entries, t_pts, initial=float(self.initial_condition)
+            )
+
+    def postfix(self, entries, t_pts, inputs) -> np.ndarray:
+        """
+        Compute the postfix sum or integral of the entries over time points,
+        """
+        the_integral = self.postfix_sum(entries, t_pts)
+        if self.post_sum_node is None:
+            return the_integral
+        else:
+            return self.post_sum(0.0, the_integral, inputs)
+
+    @staticmethod
+    def to_post_sum_expr(
+        var: pybamm.Symbol,
+        integral_node: pybamm.ExplicitTimeIntegral | pybamm.DiscreteTimeSum,
+        nstates: int,
+    ) -> pybamm.Symbol:
+        if var == integral_node:
+            # If the variable is the integral node itself, return a state vector
+            return pybamm.StateVector(slice(0, nstates))
+
+        # Should not be any time or state vector nodes in the children
+        if any(isinstance(c, (pybamm.Time, pybamm.StateVector)) for c in var.children):
+            raise ValueError(
+                "For expressions containing time integrals, "
+                "time or state vector nodes should only appear within the "
+                "time integral node. "
+            )
+        new_children = [
+            ProcessedVariableTimeIntegral.to_post_sum_expr(c, integral_node, nstates)
+            for c in var.children
+        ]
+        return var.create_copy(new_children)
 
     @staticmethod
     def from_pybamm_var(
-        var: pybamm.DiscreteTimeSum | pybamm.ExplicitTimeIntegral,
-    ) -> ProcessedVariableTimeIntegral:
-        if isinstance(var, pybamm.DiscreteTimeSum):
+        var: pybamm.Symbol,
+        nstates: int,
+    ) -> ProcessedVariableTimeIntegral | None:
+        sum_node = None
+        for symbol in var.pre_order():
+            if isinstance(
+                symbol, (pybamm.ExplicitTimeIntegral, pybamm.DiscreteTimeSum)
+            ):
+                if sum_node is None:
+                    sum_node = symbol
+                else:
+                    raise ValueError(
+                        "More than one time integral node found in the variable"
+                        f" {var.name}. Only one time integral node is allowed."
+                    )
+        if sum_node is None:
+            return None
+        if sum_node == var:
+            post_sum_node = None
+        else:
+            sum_y_len = sum_node.evaluate_for_shape().shape[0]
+            post_sum_node = ProcessedVariableTimeIntegral.to_post_sum_expr(
+                var, sum_node, sum_y_len
+            )
+        if isinstance(sum_node, pybamm.DiscreteTimeSum):
             return ProcessedVariableTimeIntegral(
-                method="discrete", initial_condition=0.0, discrete_times=var.sum_times
+                method="discrete",
+                post_sum_node=post_sum_node,
+                sum_node=sum_node,
+                initial_condition=0.0,
+                discrete_times=sum_node.sum_times,
             )
-        elif isinstance(var, pybamm.ExplicitTimeIntegral):
+        elif isinstance(sum_node, pybamm.ExplicitTimeIntegral):
             return ProcessedVariableTimeIntegral(
                 method="continuous",
-                initial_condition=var.initial_condition.evaluate(),
+                post_sum_node=post_sum_node,
+                sum_node=sum_node,
+                initial_condition=sum_node.initial_condition.evaluate(),
                 discrete_times=None,
             )
         else:
diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py
index ea37ce0f7..cd38e24a3 100644
--- a/src/pybamm/solvers/solution.py
+++ b/src/pybamm/solvers/solution.py
@@ -442,15 +442,25 @@ class Solution:
                     "solve. Please re-run the solve with `output_variables` set to "
                     "include this variable."
                 )
-            elif isinstance(
-                var_pybamm, (pybamm.ExplicitTimeIntegral, pybamm.DiscreteTimeSum)
-            ):
+            elif variable in model._variables_casadi:
+                var_casadi = model._variables_casadi[variable]
+            else:
                 time_integral = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var(
-                    var_pybamm
+                    var_pybamm, self.all_ys[i].shape[0]
                 )
-                var_pybamm = var_pybamm.child
-                if variable in model._variables_casadi:
-                    var_casadi = model._variables_casadi[variable]
+                if time_integral is not None:
+                    vars_pybamm[i] = time_integral.sum_node.child
+                    var_casadi = self.process_casadi_var(
+                        time_integral.sum_node.child,
+                        inputs,
+                        ys.shape,
+                    )
+                    if time_integral.post_sum_node is not None:
+                        time_integral.post_sum = self.process_casadi_var(
+                            time_integral.post_sum_node,
+                            inputs,
+                            ys.shape,
+                        )
                 else:
                     var_casadi = self.process_casadi_var(
                         var_pybamm,
@@ -458,16 +468,6 @@ class Solution:
                         ys.shape,
                     )
                     model._variables_casadi[variable] = var_casadi
-                vars_pybamm[i] = var_pybamm
-            elif variable in model._variables_casadi:
-                var_casadi = model._variables_casadi[variable]
-            else:
-                var_casadi = self.process_casadi_var(
-                    var_pybamm,
-                    inputs,
-                    ys.shape,
-                )
-                model._variables_casadi[variable] = var_casadi
             vars_casadi.append(var_casadi)
         var = pybamm.process_variable(
             variable, vars_pybamm, vars_casadi, self, time_integral=time_integral


Test patch:
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index ec68ad210..64ebc4fc6 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -500,10 +500,45 @@ class TestSolution:
         time = sim.solution["Time [h]"](sim.solution.t)
         assert len(time) == 10
 
-    _solver_classes = [pybamm.CasadiSolver, pybamm.IDAKLUSolver]
+    def test_discrete_data_sum_errors(self):
+        data_times = np.array([0.0])
+        data_values = np.array([1.0])
+        data = pybamm.DiscreteTimeData(data_times, data_values, "test_data")
+        dts = pybamm.DiscreteTimeSum(data)
+
+        model = pybamm.BaseModel(name="test_model2")
+        c = pybamm.Variable("c")
+        model.rhs = {c: -c}
+        model.initial_conditions = {c: 1}
+        model.variables["dts"] = pybamm.t * dts
+        solver = pybamm.IDAKLUSolver()
+        with pytest.raises(
+            ValueError,
+            match="time or state vector nodes should only appear within the time integral node",
+        ):
+            solver.solve(model, t_eval=[0, 0.1])["dts"]
+
+        model = pybamm.BaseModel(name="test_model2")
+        c = pybamm.Variable("c")
+        model.rhs = {c: -c}
+        model.initial_conditions = {c: 1}
+        model.variables["dts"] = dts * dts
+        solver = pybamm.IDAKLUSolver()
+        with pytest.raises(
+            ValueError,
+            match="More than one time integral node found",
+        ):
+            solver.solve(model, t_eval=[0, 0.1])["dts"]
+
+    _solver_classes = [
+        (pybamm.CasadiSolver, False),
+        (pybamm.IDAKLUSolver, False),
+        (pybamm.CasadiSolver, True),
+        (pybamm.IDAKLUSolver, True),
+    ]
 
-    @pytest.mark.parametrize("solver_class", _solver_classes)
-    def test_discrete_data_sum(self, solver_class):
+    @pytest.mark.parametrize("solver_class,use_post_sum", _solver_classes)
+    def test_discrete_data_sum(self, solver_class, use_post_sum):
         model = pybamm.BaseModel(name="test_model")
         c = pybamm.Variable("c")
         model.rhs = {c: -2 * c}
@@ -521,13 +556,17 @@ class TestSolution:
         data_values = solver.solve(model, t_eval=t_eval, t_interp=t_interp)["c"].entries
 
         data = pybamm.DiscreteTimeData(data_times, data_values, "test_data")
-        data_comparison = pybamm.DiscreteTimeSum((c - data) ** 2)
+        if use_post_sum:
+            data_comparison = (pybamm.DiscreteTimeSum((c - data) ** 2)) ** 0.5
+        else:
+            data_comparison = pybamm.DiscreteTimeSum((c - data) ** 2)
 
         model = pybamm.BaseModel(name="test_model2")
         a = pybamm.InputParameter("a")
         b = pybamm.InputParameter("b")
-        model.rhs = {c: b * -a * c}
-        model.initial_conditions = {c: 1}
+        c2 = pybamm.Variable("c2")
+        model.rhs = {c: b * -a * c, c2: -2 * c2}
+        model.initial_conditions = {c: 1, c2: 1}
         model.variables["data_comparison"] = data_comparison
         model.variables["data"] = data
         model.variables["c"] = c
@@ -540,7 +579,10 @@ class TestSolution:
                 model, t_eval=t_eval, t_interp=t_interp, inputs={"a": a, "b": b}
             )
             y_sol = np.exp(b * -a * data_times)
-            expected = np.sum((y_sol - data_values) ** 2)
+            if use_post_sum:
+                expected = np.sqrt(np.sum((y_sol - data_values) ** 2))
+            else:
+                expected = np.sum((y_sol - data_values) ** 2)
             np.testing.assert_allclose(
                 sol["data_comparison"](), expected, rtol=1e-3, atol=1e-2
             )
@@ -556,6 +598,14 @@ class TestSolution:
                 )
                 y_sol = np.exp(b * -a * data_times)
                 dy_sol_da = -data_times * y_sol
+                if use_post_sum:
+                    expected_sens = (
+                        0.5
+                        * (expected ** (-0.5))
+                        * np.sum(2 * (y_sol - data_values) * dy_sol_da)
+                    )
+                else:
+                    expected_sens = np.sum(2 * (y_sol - data_values) * dy_sol_da)
 
                 np.testing.assert_allclose(
                     sol["data"].sensitivities["a"].flatten(),
@@ -577,7 +627,7 @@ class TestSolution:
                 )
                 np.testing.assert_allclose(
                     sol["data_comparison"].sensitivities["a"],
-                    np.sum(2 * (y_sol - data_values) * dy_sol_da),
+                    expected_sens,
                     rtol=1e-3,
                     atol=1e-2,
                 )