#992 add simple 1D and 2D plots

pybamm/__init__.py

@@ -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 *
@@ -223,7 +223,7 @@ def version(formatted=False):
 #
 # other
 #
-from .quick_plot import QuickPlot, dynamic_plot, ax_min, ax_max
+from .quick_plot import QuickPlot, plot, plot2D, dynamic_plot, ax_min, ax_max
 
 from .simulation import Simulation, load_sim, is_notebook
 

pybamm/expression_tree/array.py

@@ -101,3 +101,16 @@ 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"
+    return pybamm.Array(np.linspace(start, stop, num, **kwargs))
+
+
+def meshgrid(x, y, **kwargs):
+    "Return coordinate matrices as from coordinate vectors"
+    [X, Y] = np.meshgrid(x.entries, y.entries)
+    X = pybamm.Array(X)
+    Y = pybamm.Array(Y)
+    return X, Y

pybamm/quick_plot.py

@@ -62,6 +62,96 @@ def dynamic_plot(*args, **kwargs):
     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

test.py

@@ -0,0 +1,13 @@
+import pybamm
+
+param = pybamm.ParameterValues(chemistry=pybamm.parameter_sets.Ecker2015)
+D_fun = param["Negative electrode diffusivity [m2.s-1]"]
+x = pybamm.linspace(0, 1, 100)
+y = pybamm.linspace(-10, 20, 100) + 300
+[X, Y] = pybamm.meshgrid(x, y)
+
+D = D_fun(x, 300)
+pybamm.plot(x, D, xlabel="x", ylabel="y")
+
+D = D_fun(X, Y)
+pybamm.plot2D(X, Y, D, xlabel="x", ylabel="y")