Fix corr mode and add tests

jetplot/images.py

@@ -28,8 +28,9 @@ def img(
 
     Args:
       img: array_like, The array to visualize.
-      mode: string, Either 'div' for a diverging image or 'seq' for
-        sequential (default: 'div').
+      mode: string, One of 'div' for a diverging image, 'seq' for
+        sequential, 'cov' for covariance matrices, or 'corr' for
+        correlation matrices (default: 'div').
       cmap: string, Colormap to use.
       aspect: string, Either 'equal' or 'auto'
     """
@@ -57,7 +58,7 @@ def img(
             cmap = "viridis"
     elif mode == "cov":
         vmin, vmax, cmap, cbar = 0, 1, "viridis", True
-    elif mode == "cov":
+    elif mode == "corr":
         vmin, vmax, cmap, cbar = -1, 1, "seismic", True
     else:
         raise ValueError("Unrecognized mode: '" + mode + "'")

tests/test_images.py

@@ -0,0 +1,17 @@
+import numpy as np
+from matplotlib import pyplot as plt
+from jetplot import images
+
+
+def test_img_corr_mode():
+    data = np.eye(3)
+    fig, ax = plt.subplots()
+    im = images.img(data, mode="corr", fig=fig, ax=ax)
+
+    # Check defaults for correlation mode
+    assert im.get_cmap().name == "seismic"
+    assert im.get_clim() == (-1, 1)
+
+    # Colorbar should have been added
+    assert len(fig.axes) == 2
+    plt.close(fig)