Repository: pybamm-team/PyBaMM
Pull request #4190 (2024-06-17 17:13:48)

move secondary broadcast outside of gradient
# Description

This fixes the edge case for #2837 / removes the need to use x-average

## 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.

- [ ] New feature (non-breaking change which adds functionality)
- [ ] Optimization (back-end change that speeds up the code)
- [ ] Bug fix (non-breaking change which fixes an issue)

# Key checklist:

- [ ] 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)
- [ ] All tests pass: `$ python run-tests.py --all` (or `$ nox -s tests`)
- [ ] 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:

- [ ] Code is commented, particularly in hard-to-understand areas
- [ ] Tests added that prove fix is effective or that feature works



Patch:
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index 9a9ad4b44..ace1cd994 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -1269,6 +1269,11 @@ def grad(symbol):
         else:
             new_child = pybamm.PrimaryBroadcast(0, symbol.child.domain)
         return pybamm.PrimaryBroadcastToEdges(new_child, symbol.domain)
+    elif isinstance(symbol, pybamm.SecondaryBroadcast):
+        # Take gradient of the child
+        # then broadcast back to the originalsymbol's secondary domain
+        # We can do this because gradient only acts on the primary domain
+        return pybamm.SecondaryBroadcast(grad(symbol.child), symbol.secondary_domain)
     elif isinstance(symbol, pybamm.FullBroadcast):
         return pybamm.FullBroadcastToEdges(0, broadcast_domains=symbol.domains)
     else:


Test patch:
diff --git a/tests/unit/test_expression_tree/test_unary_operators.py b/tests/unit/test_expression_tree/test_unary_operators.py
index c61301055..06f2e1fa5 100644
--- a/tests/unit/test_expression_tree/test_unary_operators.py
+++ b/tests/unit/test_expression_tree/test_unary_operators.py
@@ -223,6 +223,14 @@ class TestUnaryOperators(TestCase):
         grad = pybamm.grad(a)
         self.assertEqual(grad, pybamm.PrimaryBroadcastToEdges(0, "test domain"))
 
+        # gradient of a secondary broadcast moves the secondary out of the gradient
+        a = pybamm.Symbol("a", domain="test domain")
+        a_broad = pybamm.SecondaryBroadcast(a, "another domain")
+        grad = pybamm.grad(a_broad)
+        self.assertEqual(
+            grad, pybamm.SecondaryBroadcast(pybamm.grad(a), "another domain")
+        )
+
         # otherwise gradient should work
         a = pybamm.Symbol("a", domain="test domain")
         grad = pybamm.Gradient(a)
diff --git a/tests/unit/test_spatial_methods/test_finite_volume/test_finite_volume.py b/tests/unit/test_spatial_methods/test_finite_volume/test_finite_volume.py
index 52fa379ef..1d5844d7b 100644
--- a/tests/unit/test_spatial_methods/test_finite_volume/test_finite_volume.py
+++ b/tests/unit/test_spatial_methods/test_finite_volume/test_finite_volume.py
@@ -575,6 +575,51 @@ class TestFiniteVolume(TestCase):
         y = np.arange(n)[:, np.newaxis]
         self.assertEqual(evaluate_at_disc.evaluate(y=y), y[idx])
 
+    def test_inner(self):
+        # standard
+        mesh = get_mesh_for_testing()
+        spatial_methods = {
+            "macroscale": pybamm.FiniteVolume(),
+            "negative particle": pybamm.FiniteVolume(),
+        }
+
+        var = pybamm.Variable("var", domain="negative particle")
+        grad_var = pybamm.grad(var)
+        inner = pybamm.inner(grad_var, grad_var)
+
+        disc = pybamm.Discretisation(mesh, spatial_methods)
+        disc.set_variable_slices([var])
+        boundary_conditions = {
+            var: {
+                "left": (pybamm.Scalar(0), "Neumann"),
+                "right": (pybamm.Scalar(0), "Neumann"),
+            }
+        }
+        disc.bcs = boundary_conditions
+        inner_disc = disc.process_symbol(inner)
+
+        self.assertIsInstance(inner_disc, pybamm.Inner)
+        self.assertIsInstance(inner_disc.left, pybamm.MatrixMultiplication)
+        self.assertIsInstance(inner_disc.right, pybamm.MatrixMultiplication)
+
+        n = mesh["negative particle"].npts
+        y = np.ones(n)[:, np.newaxis]
+        np.testing.assert_array_equal(inner_disc.evaluate(y=y), np.zeros((n, 1)))
+        mesh = get_mesh_for_testing()
+
+        # with secondary broadcast
+        grad_var = pybamm.grad(pybamm.SecondaryBroadcast(var, "negative electrode"))
+        inner = pybamm.inner(grad_var, grad_var)
+
+        inner_disc = disc.process_symbol(inner)
+
+        self.assertIsInstance(inner_disc, pybamm.Inner)
+        self.assertIsInstance(inner_disc.left, pybamm.MatrixMultiplication)
+        self.assertIsInstance(inner_disc.right, pybamm.MatrixMultiplication)
+
+        m = mesh["negative electrode"].npts
+        np.testing.assert_array_equal(inner_disc.evaluate(y=y), np.zeros((n * m, 1)))
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")