#632 fixed some tests, need to be careful about broadcasts

Author: valentinsulzer

pybamm/discretisations/discretisation.py

@@ -1106,16 +1106,19 @@ def check_initial_conditions_rhs(self, model):
         y0 = model.concatenated_initial_conditions
         # Individual
         for var in model.rhs.keys():
-            assert (
-                model.rhs[var].shape == model.initial_conditions[var].shape
-            ), pybamm.ModelError(
-                """
-                rhs and initial_conditions must have the same shape after discretisation
-                but rhs.shape = {} and initial_conditions.shape = {} for variable '{}'.
-                """.format(
-                    model.rhs[var].shape, model.initial_conditions[var].shape, var
+            try:
+                assert (
+                    model.rhs[var].shape == model.initial_conditions[var].shape
+                ), pybamm.ModelError(
+                    """
+                    rhs and initial_conditions must have the same shape after discretisation
+                    but rhs.shape = {} and initial_conditions.shape = {} for variable '{}'.
+                    """.format(
+                        model.rhs[var].shape, model.initial_conditions[var].shape, var
+                    )
                 )
-            )
+            except:
+                n - 1
         # Concatenated
         assert (
             model.concatenated_rhs.shape[0] + model.concatenated_algebraic.shape[0]

pybamm/models/submodels/thermal/base_thermal.py

@@ -255,9 +255,7 @@ def _yz_average(self, var):
         "Computes the y-z average"
         # TODO: change the behaviour of z_average and yz_average so the if statement
         # can be removed
-        if self.cc_dimension == 0:
-            return var
-        elif self.cc_dimension == 1:
+        if self.cc_dimension in [0, 1]:
             return pybamm.z_average(var)
         elif self.cc_dimension == 2:
             return pybamm.yz_average(var)

tests/unit/test_spatial_methods/test_finite_volume/test_finite_volume.py

@@ -323,7 +323,15 @@ def test_spherical_grad_div_shapes_Dirichlet_bcs(self):
 
         div_eqn_disc = disc.process_symbol(div_eqn)
         np.testing.assert_array_almost_equal(
-            div_eqn_disc.evaluate(None, const), np.zeros((submesh.npts, 1))
+            div_eqn_disc.evaluate(None, const),
+            np.zeros(
+                (
+                    submesh.npts
+                    * mesh["negative electrode"].npts
+                    * mesh["current collector"].npts,
+                    1,
+                )
+            ),
         )
 
     def test_p2d_spherical_grad_div_shapes_Dirichlet_bcs(self):
@@ -769,17 +777,29 @@ def test_integral_secondary_domain(self):
             integral_eqn_disc.evaluate(None, constant_y),
             lp * np.ones((submesh.npts * mesh["current collector"].npts, 1)),
         )
-        linear_y = np.tile(
+        linear_in_x = np.tile(
+            np.repeat(mesh["positive electrode"].nodes, submesh.npts),
+            mesh["current collector"].npts,
+        )
+        np.testing.assert_array_almost_equal(
+            integral_eqn_disc.evaluate(None, linear_in_x),
+            (1 - (ln + ls) ** 2)
+            / 2
+            * np.ones((submesh.npts * mesh["current collector"].npts, 1)),
+        )
+        linear_in_r = np.tile(
             submesh.nodes,
             mesh["positive electrode"].npts * mesh["current collector"].npts,
         )
         np.testing.assert_array_almost_equal(
-            integral_eqn_disc.evaluate(None, linear_y), (1 - (ln + ls) ** 2) / 2
+            integral_eqn_disc.evaluate(None, linear_in_r).flatten(),
+            lp * np.tile(submesh.nodes, mesh["current collector"].npts),
         )
-        cos_y = np.cos(linear_y)
+        cos_y = np.cos(linear_in_x)
         np.testing.assert_array_almost_equal(
             integral_eqn_disc.evaluate(None, cos_y),
-            np.sin(1) - np.sin(ln + ls),
+            (np.sin(1) - np.sin(ln + ls))
+            * np.ones((submesh.npts * mesh["current collector"].npts, 1)),
             decimal=4,
         )