Issue 1087 endpoints interpolate 2D

CHANGELOG.md

@@ -14,6 +14,7 @@
 
 ## Bug fixes
 
+-   2D processed variables can now be evaluated at the domain boundaries ([#1088](https://github.com/pybamm-team/PyBaMM/pull/1088))
 -   Update the default variable points to better capture behaviour in the solid particles in li-ion models ([#1081](https://github.com/pybamm-team/PyBaMM/pull/1081))
 -   Fix `QuickPlot` to display variables discretised by FEM (in y-z) properly ([#1078](https://github.com/pybamm-team/PyBaMM/pull/1078))
 -   Add length scales to `EffectiveResistance` models ([#1071](https://github.com/pybamm-team/PyBaMM/pull/1071))

pybamm/solvers/processed_variable.py

@@ -261,12 +261,87 @@ def initialise_2D(self):
         """
         first_dim_nodes = self.mesh.nodes
         first_dim_edges = self.mesh.edges
-        second_dim_pts = self.base_variable.secondary_mesh.nodes
-        if self.base_eval.size // len(second_dim_pts) == len(first_dim_nodes):
+        second_dim_nodes = self.base_variable.secondary_mesh.nodes
+        second_dim_edges = self.base_variable.secondary_mesh.edges
+        if self.base_eval.size // len(second_dim_nodes) == len(first_dim_nodes):
             first_dim_pts = first_dim_nodes
-        elif self.base_eval.size // len(second_dim_pts) == len(first_dim_edges):
+        elif self.base_eval.size // len(second_dim_nodes) == len(first_dim_edges):
             first_dim_pts = first_dim_edges
 
+        second_dim_pts = second_dim_nodes
+        first_dim_size = len(first_dim_pts)
+        second_dim_size = len(second_dim_pts)
+        entries = np.empty((first_dim_size, second_dim_size, len(self.t_sol)))
+
+        # Evaluate the base_variable index-by-index
+        for idx in range(len(self.t_sol)):
+            t = self.t_sol[idx]
+            u = self.u_sol[:, idx]
+            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
+            if self.known_evals:
+                eval_and_known_evals = self.base_variable.evaluate(
+                    t, u, inputs=inputs, known_evals=self.known_evals[t]
+                )
+                entries[:, :, idx] = np.reshape(
+                    eval_and_known_evals[0],
+                    [first_dim_size, second_dim_size],
+                    order="F",
+                )
+                self.known_evals[t] = eval_and_known_evals[1]
+            else:
+                entries[:, :, idx] = np.reshape(
+                    self.base_variable.evaluate(t, u, inputs=inputs),
+                    [first_dim_size, second_dim_size],
+                    order="F",
+                )
+
+        # add points outside first dimension domain for extrapolation to
+        # boundaries
+        extrap_space_first_dim_left = np.array(
+            [2 * first_dim_pts[0] - first_dim_pts[1]]
+        )
+        extrap_space_first_dim_right = np.array(
+            [2 * first_dim_pts[-1] - first_dim_pts[-2]]
+        )
+        first_dim_pts = np.concatenate(
+            [extrap_space_first_dim_left, first_dim_pts, extrap_space_first_dim_right]
+        )
+        extrap_entries_left = np.expand_dims(2 * entries[0] - entries[1], axis=0)
+        extrap_entries_right = np.expand_dims(2 * entries[-1] - entries[-2], axis=0)
+        entries_for_interp = np.concatenate(
+            [extrap_entries_left, entries, extrap_entries_right], axis=0
+        )
+
+        # add points outside second dimension domain for extrapolation to
+        # boundaries
+        extrap_space_second_dim_left = np.array(
+            [2 * second_dim_pts[0] - second_dim_pts[1]]
+        )
+        extrap_space_second_dim_right = np.array(
+            [2 * second_dim_pts[-1] - second_dim_pts[-2]]
+        )
+        second_dim_pts = np.concatenate(
+            [
+                extrap_space_second_dim_left,
+                second_dim_pts,
+                extrap_space_second_dim_right,
+            ]
+        )
+        extrap_entries_second_dim_left = np.expand_dims(
+            2 * entries_for_interp[:, 0, :] - entries_for_interp[:, 1, :], axis=1
+        )
+        extrap_entries_second_dim_right = np.expand_dims(
+            2 * entries_for_interp[:, -1, :] - entries_for_interp[:, -2, :], axis=1
+        )
+        entries_for_interp = np.concatenate(
+            [
+                extrap_entries_second_dim_left,
+                entries_for_interp,
+                extrap_entries_second_dim_right,
+            ],
+            axis=1,
+        )
+
         # Process r-x or x-z
         if self.domain[0] in [
             "negative particle",
@@ -294,53 +369,31 @@ def initialise_2D(self):
                 "and auxiliary_domains '{}'".format(self.domain, self.auxiliary_domains)
             )
 
-        first_dim_size = len(first_dim_pts)
-        second_dim_size = len(second_dim_pts)
-        entries = np.empty((first_dim_size, second_dim_size, len(self.t_sol)))
-
-        # Evaluate the base_variable index-by-index
-        for idx in range(len(self.t_sol)):
-            t = self.t_sol[idx]
-            u = self.u_sol[:, idx]
-            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
-            if self.known_evals:
-                eval_and_known_evals = self.base_variable.evaluate(
-                    t, u, inputs=inputs, known_evals=self.known_evals[t]
-                )
-                entries[:, :, idx] = np.reshape(
-                    eval_and_known_evals[0],
-                    [first_dim_size, second_dim_size],
-                    order="F",
-                )
-                self.known_evals[t] = eval_and_known_evals[1]
-            else:
-                entries[:, :, idx] = np.reshape(
-                    self.base_variable.evaluate(t, u, inputs=inputs),
-                    [first_dim_size, second_dim_size],
-                    order="F",
-                )
-
         # assign attributes for reference
         self.entries = entries
         self.dimensions = 2
         first_length_scale = self.get_spatial_scale(
             self.first_dimension, self.domain[0]
         )
-        self.first_dim_pts = first_dim_pts * first_length_scale
+        first_dim_pts_for_interp = first_dim_pts * first_length_scale
 
         second_length_scale = self.get_spatial_scale(
             self.second_dimension, self.auxiliary_domains["secondary"][0]
         )
-        self.second_dim_pts = second_dim_pts * second_length_scale
+        second_dim_pts_for_interp = second_dim_pts * second_length_scale
+
+        # Set pts to edges for nicer plotting
+        self.first_dim_pts = first_dim_edges * first_length_scale
+        self.second_dim_pts = second_dim_edges * second_length_scale
 
         # set up interpolation
         if len(self.t_sol) == 1:
             # function of space only. Note the order of the points is the reverse
             # of what you'd expect
             interpolant = interp.interp2d(
-                self.second_dim_pts,
-                self.first_dim_pts,
-                entries[:, :, 0],
+                second_dim_pts_for_interp,
+                first_dim_pts_for_interp,
+                entries_for_interp[:, :, 0],
                 kind="linear",
                 fill_value=np.nan,
                 bounds_error=False,
@@ -353,8 +406,8 @@ def interp_fun(input):
         else:
             # function of space and time.
             self._interpolation_function = interp.RegularGridInterpolator(
-                (self.first_dim_pts, self.second_dim_pts, self.t_pts),
-                entries,
+                (first_dim_pts_for_interp, second_dim_pts_for_interp, self.t_pts),
+                entries_for_interp,
                 method="linear",
                 fill_value=np.nan,
                 bounds_error=False,

tests/unit/test_quick_plot.py

@@ -330,7 +330,7 @@ def test_loqs_spme(self):
                 pybamm.QuickPlot(solution, output_variables)
 
                 # check 2D (space) variables update properly for different time units
-                c_n = solution["Negative particle concentration [mol.m-3]"].entries
+                c_n = solution["Negative particle concentration [mol.m-3]"]
 
                 for unit, scale in zip(["seconds", "minutes", "hours"], [1, 60, 3600]):
                     quick_plot = pybamm.QuickPlot(
@@ -342,12 +342,16 @@ def test_loqs_spme(self):
                     qp_data = quick_plot.plots[
                         ("Negative particle concentration [mol.m-3]",)
                     ][0][1]
-                    np.testing.assert_array_almost_equal(qp_data, c_n[:, :, 0])
+                    c_n_eval = c_n(t_eval[0], r=c_n.first_dim_pts, x=c_n.second_dim_pts)
+                    np.testing.assert_array_almost_equal(qp_data, c_n_eval)
                     quick_plot.slider_update(t_eval[-1] / scale)
                     qp_data = quick_plot.plots[
                         ("Negative particle concentration [mol.m-3]",)
                     ][0][1]
-                    np.testing.assert_array_almost_equal(qp_data, c_n[:, :, 1])
+                    c_n_eval = c_n(
+                        t_eval[-1], r=c_n.first_dim_pts, x=c_n.second_dim_pts
+                    )
+                    np.testing.assert_array_almost_equal(qp_data, c_n_eval)
 
         pybamm.close_plots()
 
@@ -369,18 +373,20 @@ def test_plot_1plus1D_spme(self):
 
         # check 2D (x,z space) variables update properly for different time units
         # Note: these should be the transpose of the entries in the processed variable
-        c_e = solution["Electrolyte concentration [mol.m-3]"].entries
+        c_e = solution["Electrolyte concentration [mol.m-3]"]
 
         for unit, scale in zip(["seconds", "minutes", "hours"], [1, 60, 3600]):
             quick_plot = pybamm.QuickPlot(
                 solution, ["Electrolyte concentration [mol.m-3]"], time_unit=unit
             )
             quick_plot.plot(0)
             qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][1]
-            np.testing.assert_array_almost_equal(qp_data.T, c_e[:, :, 0])
+            c_e_eval = c_e(t_eval[0], x=c_e.first_dim_pts, z=c_e.second_dim_pts)
+            np.testing.assert_array_almost_equal(qp_data.T, c_e_eval)
             quick_plot.slider_update(t_eval[-1] / scale)
             qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][1]
-            np.testing.assert_array_almost_equal(qp_data.T, c_e[:, :, -1])
+            c_e_eval = c_e(t_eval[-1], x=c_e.first_dim_pts, z=c_e.second_dim_pts)
+            np.testing.assert_array_almost_equal(qp_data.T, c_e_eval)
 
         pybamm.close_plots()