generalizing imshow(labels, x, y)

Author: nicolaskruchten

Diff for: doc/python/imshow.md

@@ -6,7 +6,7 @@ jupyter:
       extension: .md
       format_name: markdown
       format_version: '1.2'
-      jupytext_version: 1.3.0
+      jupytext_version: 1.3.1
   kernelspec:
     display_name: Python 3
     language: python
@@ -20,7 +20,7 @@ jupyter:
     name: python
     nbconvert_exporter: python
     pygments_lexer: ipython3
-    version: 3.7.3
+    version: 3.6.8
   plotly:
     description: How to display image data in Python with Plotly.
     display_as: scientific
@@ -93,7 +93,11 @@ fig.show()
 import plotly.express as px
 import numpy as np
 img = np.arange(100).reshape((10, 10))
-fig = px.imshow(img, color_continuous_scale='gray')
+fig = px.imshow(img, color_continuous_scale='gray', labels=dict(x="yoo", y="yaa", color="hey"),
+               width=600, height=600, 
+                x=["a","b","c","d","e","f","g","h","i","j"],
+                y=["a","b","c","d","e","f","g","h","i","j"]
+               )
 fig.show()
 ```
 
@@ -120,7 +124,7 @@ import xarray as xr
 # Load xarray from dataset included in the xarray tutorial
 airtemps = xr.tutorial.open_dataset('air_temperature').air.sel(lon=250.0)
 fig = px.imshow(airtemps.T, color_continuous_scale='RdBu_r', origin='lower',
-                #labels={'colorbar':airtemps.attrs['var_desc']}
+                labels={'color':airtemps.attrs['var_desc']}
                )
 fig.show()
 ```
@@ -133,9 +137,8 @@ For xarrays, by default `px.imshow` does not constrain pixels to be square, sinc
 import plotly.express as px
 import xarray as xr
 airtemps = xr.tutorial.open_dataset('air_temperature').air.isel(time=500)
-colorbar_title = airtemps.attrs['var_desc'] + '<br> (%s)'%airtemps.attrs['units']
-fig = px.imshow(airtemps, color_continuous_scale='RdBu_r', aspect='equal',
-                labels={'colorbar':colorbar_title})
+colorbar_title = airtemps.attrs['var_desc'] + '<br>(%s)'%airtemps.attrs['units']
+fig = px.imshow(airtemps, color_continuous_scale='RdBu_r', aspect='equal')
 fig.show()
 ```
 

Diff for: packages/python/plotly/plotly/express/_imshow.py

@@ -69,6 +69,8 @@ def imshow(
     zmax=None,
     origin=None,
     labels={},
+    x=None,
+    y=None,
     color_continuous_scale=None,
     color_continuous_midpoint=None,
     range_color=None,
@@ -167,33 +169,33 @@ def imshow(
     """
     args = locals()
     apply_default_cascade(args)
-    img_is_xarray = False
-    x_label = "x"
-    y_label = "y"
-    z_label = ""
-    if xarray_imported:
-        if isinstance(img, xarray.DataArray):
-            y_label, x_label = img.dims[0], img.dims[1]
-            # np.datetime64 is not handled correctly by go.Heatmap
-            for ax in [x_label, y_label]:
-                if np.issubdtype(img.coords[ax].dtype, np.datetime64):
-                    img.coords[ax] = img.coords[ax].astype(str)
+    labels = labels.copy()
+    if xarray_imported and isinstance(img, xarray.DataArray):
+        y_label, x_label = img.dims[0], img.dims[1]
+        # np.datetime64 is not handled correctly by go.Heatmap
+        for ax in [x_label, y_label]:
+            if np.issubdtype(img.coords[ax].dtype, np.datetime64):
+                img.coords[ax] = img.coords[ax].astype(str)
+        if x is None:
             x = img.coords[x_label]
+        if y is None:
             y = img.coords[y_label]
-            img_is_xarray = True
-            if aspect is None:
-                aspect = "auto"
-            z_label = xarray.plot.utils.label_from_attrs(img).replace("\n", "<br>")
-
-    if labels is not None:
-        if "x" in labels:
-            x_label = labels["x"]
-        if "y" in labels:
-            y_label = labels["y"]
-        if "color" in labels:
-            z_label = labels["color"]
-
-    if not img_is_xarray:
+        if aspect is None:
+            aspect = "auto"
+        if labels.get("x", None) is None:
+            labels["x"] = x_label
+        if labels.get("y", None) is None:
+            labels["y"] = y_label
+        if labels.get("color", None) is None:
+            labels["color"] = xarray.plot.utils.label_from_attrs(img)
+            labels["color"] = labels["color"].replace("\n", "<br>")
+    else:
+        if labels.get("x", None) is None:
+            labels["x"] = ""
+        if labels.get("y", None) is None:
+            labels["y"] = ""
+        if labels.get("color", None) is None:
+            labels["color"] = ""
         if aspect is None:
             aspect = "equal"
 
@@ -205,7 +207,7 @@ def imshow(
 
     # For 2d data, use Heatmap trace
     if img.ndim == 2:
-        trace = go.Heatmap(z=img, coloraxis="coloraxis1")
+        trace = go.Heatmap(x=x, y=y, z=img, coloraxis="coloraxis1")
         autorange = True if origin == "lower" else "reversed"
         layout = dict(yaxis=dict(autorange=autorange))
         if aspect == "equal":
@@ -224,8 +226,9 @@ def imshow(
             cmid=color_continuous_midpoint,
             cmin=range_color[0],
             cmax=range_color[1],
-            colorbar=dict(title=z_label),
         )
+        if labels["color"]:
+            layout["coloraxis1"]["colorbar"] = dict(title=labels["color"])
 
     # For 2D+RGB data, use Image trace
     elif img.ndim == 3 and img.shape[-1] in [3, 4]:
@@ -250,19 +253,13 @@ def imshow(
         layout_patch["margin"] = {"t": 60}
     fig = go.Figure(data=trace, layout=layout)
     fig.update_layout(layout_patch)
-    if img.ndim <= 2:
-        hovertemplate = (
-            x_label
-            + ": %{x} <br>"
-            + y_label
-            + ": %{y} <br>"
-            + z_label
-            + " : %{z}<extra></extra>"
-        )
-        fig.update_traces(hovertemplate=hovertemplate)
-    if img_is_xarray:
-        fig.update_traces(x=x, y=y)
-    fig.update_xaxes(title_text=x_label)
-    fig.update_yaxes(title_text=y_label)
+    fig.update_traces(
+        hovertemplate="%s: %%{x}<br>%s: %%{y}<br>%s: %%{z}<extra></extra>"
+        % (labels["x"] or "x", labels["y"] or "y", labels["color"] or "color",)
+    )
+    if labels["x"]:
+        fig.update_xaxes(title_text=labels["x"])
+    if labels["y"]:
+        fig.update_yaxes(title_text=labels["y"])
     fig.update_layout(template=args["template"], overwrite=True)
     return fig