Add colormap() and contourf() functions

Author: spmccormack

examples/wavelet/plot_wavelet.py

@@ -18,6 +18,8 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
+import pymodalib
+
 os.chdir(os.path.abspath(os.path.dirname(__file__)))
 
 print("Plotting...")
@@ -39,7 +41,7 @@
 
 fig, (ax1, ax2) = plt.subplots(1, 2, gridspec_kw={"width_ratios": [3, 1,]}, sharey=True)
 
-ax1.contourf(mesh1, mesh2, wt_power)
+pymodalib.contourf(ax1, mesh1, mesh2, wt_power)
 ax1.set_title(f"WT power ({impl} implementation)")
 ax1.set_xlabel("Time (s)")
 

pymodalib/__init__.py

@@ -22,6 +22,7 @@
 from pymodalib.algorithms.signals import resampl_flow as downsample
 from pymodalib.algorithms.wavelet import wavelet_transform
 from pymodalib.algorithms.windowed_fourier import windowed_fourier_transform
+from pymodalib.utils.plotting import colormap, contourf
 from pymodalib.utils.cache import cachedarray, cleanup
 
 # PyMODAlib version.
@@ -42,6 +43,8 @@
     ridge_extraction,
     ecurve,
     rectfr,
+    colormap,
+    contourf,
 )
 
 # Useful aliases for functions.

pymodalib/utils/colormap.mat

@@ -0,0 +1,90 @@
+#  PyMODAlib, a Python implementation of the algorithms from MODA (Multiscale Oscillatory Dynamics Analysis).
+#  Copyright (C) 2020 Lancaster University
+#
+#  This program is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program. If not, see <https://www.gnu.org/licenses/>.
+import warnings
+
+import numpy as np
+from numpy import ndarray
+
+
+def contourf(
+    axes, x: ndarray, y: ndarray, z: ndarray, levels: int = 256, *args, **kwargs
+) -> None:
+    """
+    Plots a contour plot in PyMODA style. Useful for easily plotting a wavelet transform, etc.
+
+    This function is a Wrapper around matplotlib's 'contourf'.
+
+    Parameters
+    ----------
+    axes
+        The Axes object to plot on.
+    x, y : ndarray
+        The coordinates of the values in Z.
+
+        X and Y must both be 2-D with the same shape as Z (e.g. created via numpy.meshgrid),
+        or they must both be 1-D such that len(X) == M is the number of columns in Z and
+        len(Y) == N is the number of rows in Z.
+    z : ndarray
+        The height values over which the contour is drawn.
+    levels : int
+        Determines the number and positions of the contour lines / regions.
+
+        If an int n, use n data intervals; i.e. draw n+1 contour lines. The level heights are automatically chosen.
+
+        If array-like, draw contour lines at the specified levels. The values must be in increasing order.
+
+    .. note::
+        Documentation copied from the relevant matplotlib function, `matplotlib.pyplot.contourf`.
+    """
+    axes.contourf(
+        x,
+        y,
+        z,
+        levels,
+        vmin=np.nanmin(z),
+        vmax=np.nanmax(z),
+        cmap=colormap(),
+        *args,
+        **kwargs
+    )
+
+
+def colormap() -> "LinearSegmentedColormap":
+    """
+    Loads the colormap used by PyMODA.
+
+    Returns
+    -------
+    LinearSegmentedColormap
+        The colormap, as an object which can be passed to matplotlib functions.
+    """
+    from matplotlib.colors import LinearSegmentedColormap
+    from scipy.io import loadmat
+    import os
+
+    here = os.path.abspath(os.path.dirname(__file__))
+    filename = os.path.join(here, "colormap.mat")
+    cmap = loadmat(filename).get("cmap")
+
+    if cmap is None:
+        warnings.warn(
+            "Could not load colormap. The colormap data is not supplied with "
+            "source-distributions of PyMODAlib.",
+            RuntimeWarning,
+        )
+        return "jet"
+
+    return LinearSegmentedColormap.from_list("colours", cmap, N=len(cmap), gamma=1.0)

pymodalib/utils/plotting.py

@@ -1,3 +1,4 @@
 AsyncProcessScheduler>=0.5.0b1
 numpy>=1.18.1
 scipy>=1.4.1
+matplotlib>=3.1.1

requirements.txt

@@ -55,6 +55,8 @@ def find_packages() -> List[str]:
     name="PyMODAlib",
     version=version,
     packages=find_packages(),
+    include_package_data=True,
+    package_data={"": ["*.mat"]},
     python_requires=">=3.6",
     install_requires=requirements,
     description="Library providing Python implementations of MODA's algorithms.",