#992 reorganise plots

pybamm/__init__.py

@@ -221,10 +221,16 @@ def version(formatted=False):
 from . import experiments
 
 #
-# other
+# Plotting
 #
-from .quick_plot import QuickPlot, plot, plot2D, 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/expression_tree/array.py

@@ -104,12 +104,20 @@ def _base_evaluate(self, t=None, y=None, y_dot=None, inputs=None):
 
 
 def linspace(start, stop, num=50, **kwargs):
-    "Creates a linearly spaced array"
+    """
+    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"
+    """
+    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)

pybamm/plotting/dynamic_plot.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/plot.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/plot2D.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/quick_plot.py → pybamm/plotting/quick_plot.py

@@ -44,114 +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
-
-
-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'.
-
-    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
-
-
-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'.
-
-    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
-    z : :class:`pybamm.Array`
-        The array to plot on the z axis
-    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'")
-
-    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
-
-
 class QuickPlot(object):
     """
     Generates a quick plot of a subset of key outputs of the model so that the model

tests/unit/test_plot.py

@@ -21,6 +21,10 @@ def test_plot2D(self):
         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):