Merge pull request #1008 from pybamm-team/issue-992-plot-arrays

Issue 992 plot arrays

Author: rtimms

.flake8

@@ -12,7 +12,7 @@ exclude=
     share,
     pyvenv.cfg,
     third-party,
-    sundials-5.0.0,
+    KLU_module_deps,
 ignore=
     # False positive for white space before ':' on list slice
     # black should format these correctly

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Features
 
 -   Added `BackwardIndefiniteIntegral` symbol ([#1014](https://github.com/pybamm-team/PyBaMM/pull/1014))
+-   Added `plot` and `plot2D` to enable easy plotting of `pybamm.Array` objects ([#1008](https://github.com/pybamm-team/PyBaMM/pull/1008))
 -   Added SEI film resistance as an option ([#994](https://github.com/pybamm-team/PyBaMM/pull/994))
 -   Added tab, edge, and surface cooling ([#965](https://github.com/pybamm-team/PyBaMM/pull/965))
 -   Added functionality to solver to automatically discretise a 0D model ([#947](https://github.com/pybamm-team/PyBaMM/pull/947))

docs/index.rst

@@ -66,7 +66,7 @@ API documentation
    source/solvers/index
    source/experiments/index
    source/simulation
-   source/quick_plot
+   source/plotting/index
    source/util
    source/citations
    source/parameters_cli

docs/source/expression_tree/array.rst

@@ -0,0 +1,9 @@
+Array
+=====
+
+.. autoclass:: pybamm.Array
+  :members:
+
+.. autofunction:: pybamm.linspace
+
+.. autofunction:: pybamm.meshgrid 

docs/source/expression_tree/index.rst

@@ -8,6 +8,7 @@ Expression Tree
   variable
   independent_variable
   scalar
+  array
   matrix
   vector
   state_vector

docs/source/plotting/dynamic_plot.rst

@@ -0,0 +1,4 @@
+Dynamic Plot
+============
+
+.. autofunction:: pybamm.dynamic_plot

docs/source/plotting/index.rst

@@ -0,0 +1,9 @@
+Plotting
+========
+
+.. toctree::
+
+  quick_plot
+  dynamic_plot
+  plot
+  plot_2D

docs/source/plotting/plot.rst

@@ -0,0 +1,4 @@
+Plot
+====
+
+.. autofunction:: pybamm.plot

docs/source/plotting/plot_2D.rst

@@ -0,0 +1,4 @@
+Plot 2D
+=======
+
+.. autofunction:: pybamm.plot2D

docs/source/quick_plot.rst docs/source/plotting/quick_plot.rst

@@ -68,7 +68,7 @@ def version(formatted=False):
 from .expression_tree.symbol import *
 from .expression_tree.binary_operators import *
 from .expression_tree.concatenations import *
-from .expression_tree.array import Array
+from .expression_tree.array import Array, linspace, meshgrid
 from .expression_tree.matrix import Matrix
 from .expression_tree.unary_operators import *
 from .expression_tree.functions import *
@@ -221,10 +221,16 @@ def version(formatted=False):
 from . import experiments
 
 #
-# other
+# Plotting
 #
-from .quick_plot import QuickPlot, dynamic_plot, ax_min, ax_max
+from .plotting.quick_plot import QuickPlot
+from .plotting.plot import plot
+from .plotting.plot2D import plot2D
+from .plotting.dynamic_plot import dynamic_plot
 
+#
+# Simulation
+#
 from .simulation import Simulation, load_sim, is_notebook
 
 #

pybamm/__init__.py

@@ -101,3 +101,24 @@ def new_copy(self):
     def _base_evaluate(self, t=None, y=None, y_dot=None, inputs=None):
         """ See :meth:`pybamm.Symbol._base_evaluate()`. """
         return self._entries
+
+
+def linspace(start, stop, num=50, **kwargs):
+    """
+    Creates a linearly spaced array by calling `numpy.linspace` with keyword
+    arguments 'kwargs'. For a list of 'kwargs' see the
+    `numpy linspace documentation <https://tinyurl.com/yc4ne47x>`_
+    """
+    return pybamm.Array(np.linspace(start, stop, num, **kwargs))
+
+
+def meshgrid(x, y, **kwargs):
+    """
+    Return coordinate matrices as from coordinate vectors by calling
+    `numpy.meshgrid` with keyword arguments 'kwargs'. For a list of 'kwargs'
+    see the `numpy meshgrid documentation <https://tinyurl.com/y8azewrj>`_
+    """
+    [X, Y] = np.meshgrid(x.entries, y.entries)
+    X = pybamm.Array(X)
+    Y = pybamm.Array(Y)
+    return X, Y

pybamm/expression_tree/array.py

@@ -0,0 +1,22 @@
+#
+# Method for creating a dynamic plot
+#
+import pybamm
+
+
+def dynamic_plot(*args, **kwargs):
+    """
+    Creates a :class:`pybamm.QuickPlot` object (with arguments 'args' and keyword
+    arguments 'kwargs') and then calls :meth:`pybamm.QuickPlot.dynamic_plot`.
+    The key-word argument 'testing' is passed to the 'dynamic_plot' method, not the
+    `QuickPlot` class.
+
+    Returns
+    -------
+    plot : :class:`pybamm.QuickPlot`
+        The 'QuickPlot' object that was created
+    """
+    kwargs_for_class = {k: v for k, v in kwargs.items() if k != "testing"}
+    plot = pybamm.QuickPlot(*args, **kwargs_for_class)
+    plot.dynamic_plot(kwargs.get("testing", False))
+    return plot

pybamm/plotting/__init__.py

@@ -0,0 +1,44 @@
+#
+# Method for creating a 1D plot of pybamm arrays
+#
+import pybamm
+from .quick_plot import ax_min, ax_max
+
+
+def plot(x, y, xlabel=None, ylabel=None, title=None, testing=False, **kwargs):
+    """
+    Generate a simple 1D plot. Calls `matplotlib.pyplot.plot` with keyword
+    arguments 'kwargs'. For a list of 'kwargs' see the
+    `matplotlib plot documentation <https://tinyurl.com/ycblw9bx>`_
+
+    Parameters
+    ----------
+    x : :class:`pybamm.Array`
+        The array to plot on the x axis
+    y : :class:`pybamm.Array`
+        The array to plot on the y axis
+    xlabel : str, optional
+        The label for the x axis
+    ylabel : str, optional
+        The label for the y axis
+    testing : bool, optional
+        Whether to actually make the plot (turned off for unit tests)
+
+    """
+    import matplotlib.pyplot as plt
+
+    if not isinstance(x, pybamm.Array):
+        raise TypeError("x must be 'pybamm.Array'")
+    if not isinstance(y, pybamm.Array):
+        raise TypeError("y must be 'pybamm.Array'")
+
+    plt.plot(x.entries, y.entries, **kwargs)
+    plt.ylim([ax_min(y.entries), ax_max(y.entries)])
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.title(title)
+
+    if not testing:  # pragma: no cover
+        plt.show()
+
+    return

pybamm/plotting/dynamic_plot.py

@@ -0,0 +1,66 @@
+#
+# Method for creating a filled contour plot of pybamm arrays
+#
+import pybamm
+from .quick_plot import ax_min, ax_max
+
+
+def plot2D(x, y, z, xlabel=None, ylabel=None, title=None, testing=False, **kwargs):
+    """
+    Generate a simple 2D plot. Calls `matplotlib.pyplot.contourf` with keyword
+    arguments 'kwargs'.  For a list of 'kwargs' see the
+    `matplotlib contourf documentation <https://tinyurl.com/y8mnadtn>`_
+
+    Parameters
+    ----------
+    x : :class:`pybamm.Array`
+        The array to plot on the x axis. Can be of shape (M, N) or (N, 1)
+    y : :class:`pybamm.Array`
+        The array to plot on the y axis. Can be of shape (M, N)  or (M, 1)
+    z : :class:`pybamm.Array`
+        The array to plot on the z axis. Is of shape (M, N)
+    xlabel : str, optional
+        The label for the x axis
+    ylabel : str, optional
+        The label for the y axis
+    title : str, optional
+        The title for the plot
+    testing : bool, optional
+        Whether to actually make the plot (turned off for unit tests)
+
+    """
+    import matplotlib.pyplot as plt
+
+    if not isinstance(x, pybamm.Array):
+        raise TypeError("x must be 'pybamm.Array'")
+    if not isinstance(y, pybamm.Array):
+        raise TypeError("y must be 'pybamm.Array'")
+    if not isinstance(z, pybamm.Array):
+        raise TypeError("z must be 'pybamm.Array'")
+
+    # Get correct entries of x and y depending on shape
+    if x.shape == y.shape == z.shape:
+        x_entries = x.entries
+        y_entries = y.entries
+    else:
+        x_entries = x.entries[:, 0]
+        y_entries = y.entries[:, 0]
+
+    plt.contourf(
+        x_entries,
+        y_entries,
+        z.entries,
+        vmin=ax_min(z.entries),
+        vmax=ax_max(z.entries),
+        cmap="coolwarm",
+        **kwargs
+    )
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.title(title)
+    plt.colorbar()
+
+    if not testing:  # pragma: no cover
+        plt.show()
+
+    return

pybamm/plotting/plot.py

@@ -44,24 +44,6 @@ def split_long_string(title, max_words=4):
         return first_line + "\n" + second_line
 
 
-def dynamic_plot(*args, **kwargs):
-    """
-    Creates a :class:`pybamm.QuickPlot` object (with arguments 'args' and keyword
-    arguments 'kwargs') and then calls :meth:`pybamm.QuickPlot.dynamic_plot`.
-    The key-word argument 'testing' is passed to the 'dynamic_plot' method, not the
-    `QuickPlot' class.
-
-    Returns
-    -------
-    plot : :class:`pybamm.QuickPlot`
-        The 'QuickPlot' object that was created
-    """
-    kwargs_for_class = {k: v for k, v in kwargs.items() if k != "testing"}
-    plot = pybamm.QuickPlot(*args, **kwargs_for_class)
-    plot.dynamic_plot(kwargs.get("testing", False))
-    return plot
-
-
 class QuickPlot(object):
     """
     Generates a quick plot of a subset of key outputs of the model so that the model

pybamm/plotting/plot2D.py

@@ -99,7 +99,7 @@ def test_compare_outputs_thermal(self):
             solutions = []
             t_eval = np.linspace(0, 3600, 100)
             for model in models:
-                solution = pybamm.CasadiSolver(dt_max=0.01).solve(model, t_eval)
+                solution = pybamm.CasadiSolver().solve(model, t_eval)
                 solutions.append(solution)
 
             # compare outputs

pybamm/quick_plot.py pybamm/plotting/quick_plot.py

@@ -0,0 +1,38 @@
+#
+# Tests for the Array class
+#
+import pybamm
+import numpy as np
+
+import unittest
+
+
+class TestArray(unittest.TestCase):
+    def test_name(self):
+        arr = pybamm.Array(np.array([1, 2, 3]))
+        self.assertEqual(arr.name, "Array of shape (3, 1)")
+
+    def test_linspace(self):
+        x = np.linspace(0, 1, 100)[:, np.newaxis]
+        y = pybamm.linspace(0, 1, 100)
+        np.testing.assert_array_equal(x, y.entries)
+
+    def test_meshgrid(self):
+        a = np.linspace(0, 5)
+        b = np.linspace(0, 3)
+        A, B = np.meshgrid(a, b)
+        c = pybamm.linspace(0, 5)
+        d = pybamm.linspace(0, 3)
+        C, D = pybamm.meshgrid(c, d)
+        np.testing.assert_array_equal(A, C.entries)
+        np.testing.assert_array_equal(B, D.entries)
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    pybamm.settings.debug_mode = True
+    unittest.main()

tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_compare_outputs.py

@@ -32,12 +32,6 @@ def test_matrix_operations(self):
         )
 
 
-class TestArray(unittest.TestCase):
-    def test_name(self):
-        arr = pybamm.Array(np.array([1, 2, 3]))
-        self.assertEqual(arr.name, "Array of shape (3, 1)")
-
-
 if __name__ == "__main__":
     print("Add -v for more debug output")
     import sys

tests/unit/test_expression_tree/test_array.py

@@ -0,0 +1,47 @@
+import pybamm
+import unittest
+import numpy as np
+
+
+class TestPlot(unittest.TestCase):
+    def test_plot(self):
+        x = pybamm.Array(np.array([0, 3, 10]))
+        y = pybamm.Array(np.array([6, 16, 78]))
+        pybamm.plot(x, y, xlabel="x", ylabel="y", title="title", testing=True)
+
+    def test_plot_fail(self):
+        x = pybamm.Array(np.array([0]))
+        with self.assertRaisesRegex(TypeError, "x must be 'pybamm.Array'"):
+            pybamm.plot("bad", x)
+        with self.assertRaisesRegex(TypeError, "y must be 'pybamm.Array'"):
+            pybamm.plot(x, "bad")
+
+    def test_plot2D(self):
+        x = pybamm.Array(np.array([0, 3, 10]))
+        y = pybamm.Array(np.array([6, 16, 78]))
+        X, Y = pybamm.meshgrid(x, y)
+
+        # plot with array directly
+        pybamm.plot2D(x, y, Y, xlabel="x", ylabel="y", title="title", testing=True)
+
+        # plot with meshgrid
+        pybamm.plot2D(X, Y, Y, xlabel="x", ylabel="y", title="title", testing=True)
+
+    def test_plot2D_fail(self):
+        x = pybamm.Array(np.array([0]))
+        with self.assertRaisesRegex(TypeError, "x must be 'pybamm.Array'"):
+            pybamm.plot2D("bad", x, x)
+        with self.assertRaisesRegex(TypeError, "y must be 'pybamm.Array'"):
+            pybamm.plot2D(x, "bad", x)
+        with self.assertRaisesRegex(TypeError, "z must be 'pybamm.Array'"):
+            pybamm.plot2D(x, x, "bad")
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    pybamm.settings.debug_mode = True
+    unittest.main()