xarray support in imshow

.circleci/create_conda_optional_env.sh

@@ -19,5 +19,5 @@ if [ ! -d $HOME/miniconda/envs/circle_optional ]; then
 requests nbformat six retrying psutil pandas decorator pytest mock nose poppler xarray scikit-image ipython jupyter ipykernel ipywidgets
 
     # Install orca into environment
-    $HOME/miniconda/bin/conda install --yes -n circle_optional -c plotly plotly-orca
+    $HOME/miniconda/bin/conda install --yes -n circle_optional -c plotly plotly-orca=1.2.1
 fi

doc/python/datashader.md

@@ -20,7 +20,7 @@ jupyter:
     name: python
     nbconvert_exporter: python
     pygments_lexer: ipython3
-    version: 3.6.8
+    version: 3.7.3
   plotly:
     description:
       How to use datashader to rasterize large datasets, and visualize
@@ -98,32 +98,18 @@ fig.show()
 ```
 
 ```python
-import plotly.graph_objects as go
+import plotly.express as px
 import pandas as pd
 import numpy as np
 import datashader as ds
 df = pd.read_parquet('https://raw.githubusercontent.com/plotly/datasets/master/2015_flights.parquet')
 
 cvs = ds.Canvas(plot_width=100, plot_height=100)
 agg = cvs.points(df, 'SCHEDULED_DEPARTURE', 'DEPARTURE_DELAY')
-x = np.array(agg.coords['SCHEDULED_DEPARTURE'])
-y = np.array(agg.coords['DEPARTURE_DELAY'])
-
-# Assign nan to zero values so that the corresponding pixels are transparent
-agg = np.array(agg.values, dtype=np.float)
-agg[agg<1] = np.nan
-
-fig = go.Figure(go.Heatmap(
-    z=np.log10(agg), x=x, y=y,
-    hoverongaps=False,
-    hovertemplate='Scheduled departure: %{x:.1f}h <br>Depature delay: %{y} <br>Log10(Count): %{z}',
-    colorbar=dict(title='Count (Log)', tickprefix='1.e')))
-fig.update_xaxes(title_text='Scheduled departure')
-fig.update_yaxes(title_text='Departure delay')
+agg.values = np.log10(agg.values)
+agg.attrs['long_name'] = 'Log10(count)'
+fig = px.imshow(agg, origin='lower')
+fig.update_traces(hoverongaps=False)
+fig.update_layout(coloraxis_colorbar=dict(title='Count (Log)', tickprefix='1.e'))
 fig.show()
-
-```
-
-```python
-
 ```

doc/python/imshow.md

@@ -110,6 +110,35 @@ fig.update_layout(coloraxis_showscale=False)
 fig.show()
 ```
 
+### Display an xarray image with px.imshow
+
+[xarrays](http://xarray.pydata.org/en/stable/) are labeled arrays (with labeled axes and coordinates). If you pass an xarray image to `px.imshow`, its axes labels and coordinates will be used for axis titles. If you don't want this behavior, you can pass `img.values` which is a NumPy array if `img` is an xarray. Alternatively, you can override axis titles hover labels and colorbar title using the `labels` attribute, as in the next example.
+
+```python
+import plotly.express as px
+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']}
+               )
+fig.show()
+```
+
+### Display an xarray image with square pixels
+
+For xarrays, by default `px.imshow` does not constrain pixels to be square, since axes often correspond to different physical quantities (e.g. time and space), contrary to a plain camera image where pixels are square (most of the time). If you want to impose square pixels, set the parameter `aspect` to "equal" as below.
+
+```python
+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})
+fig.show()
+```
+
 ### Display multichannel image data with go.Image
 
 It is also possible to use the `go.Image` trace from the low-level `graph_objects` API in order to display image data. Note that `go.Image` only accepts multichannel images. For single images, use [`go.Heatmap`](/python/heatmaps).

packages/python/plotly/plotly/express/_imshow.py

@@ -3,6 +3,13 @@
 from ._core import apply_default_cascade
 import numpy as np
 
+try:
+    import xarray
+
+    xarray_imported = True
+except ImportError:
+    xarray_imported = False
+
 _float_types = []
 
 # Adapted from skimage.util.dtype
@@ -61,21 +68,23 @@ def imshow(
     zmin=None,
     zmax=None,
     origin=None,
+    labels={},
     color_continuous_scale=None,
     color_continuous_midpoint=None,
     range_color=None,
     title=None,
     template=None,
     width=None,
     height=None,
+    aspect=None,
 ):
     """
     Display an image, i.e. data on a 2D regular raster.
 
     Parameters
     ----------
 
-    img: array-like image
+    img: array-like image, or xarray
         The image data. Supported array shapes are
 
         - (M, N): an image with scalar data. The data is visualized
@@ -95,6 +104,13 @@ def imshow(
         position of the [0, 0] pixel of the image array, in the upper left or lower left
         corner. The convention 'upper' is typically used for matrices and images.
 
+    labels : dict with str keys and str values (default `{}`)
+        Overrides names used in the figure for axis titles (keys ``x`` and ``y``),
+        colorbar title and hover (key ``color``). The values should correspond
+        to the desired label to be displayed. If ``img`` is an xarray, dimension
+        names are used for axis titles, and long name for the colorbar title
+        (unless overridden in ``labels``). Possible keys are: x, y, and color.
+
     color_continuous_scale : str or list of str
         colormap used to map scalar data to colors (for a 2D image). This parameter is
         not used for RGB or RGBA images. If a string is provided, it should be the name
@@ -122,6 +138,14 @@ def imshow(
     height: number
         The figure height in pixels, defaults to 600.
 
+    aspect: 'equal', 'auto', or None
+      - 'equal': Ensures an aspect ratio of 1 or pixels (square pixels)
+      - 'auto': The axes is kept fixed and the aspect ratio of pixels is
+        adjusted so that the data fit in the axes. In general, this will
+        result in non-square pixels.
+      - if None, 'equal' is used for numpy arrays and 'auto' for xarrays
+        (which have typically heterogeneous coordinates)
+
     Returns
     -------
     fig : graph_objects.Figure containing the displayed image
@@ -137,11 +161,41 @@ def imshow(
 
     In order to update and customize the returned figure, use
     `go.Figure.update_traces` or `go.Figure.update_layout`.
+
+    If an xarray is passed, dimensions names and coordinates are used for
+    axes labels and ticks.
     """
     args = locals()
     apply_default_cascade(args)
+    img_is_xarray = False
+    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)
+            x = img.coords[x_label]
+            y = img.coords[y_label]
+            img_is_xarray = True
+            if aspect is None:
+                aspect = "auto"
+            z_name = xarray.plot.utils.label_from_attrs(img).replace("\n", "<br>")
+
+    if labels is not None:
+        if "x" in labels:
+            y_label = labels["x"]
+        if "y" in labels:
+            y_label = labels["y"]
+        if "color" in labels:
+            z_name = labels["color"]
+
+    if not img_is_xarray:
+        if aspect is None:
+            aspect = "equal"
 
     img = np.asanyarray(img)
+
     # Cast bools to uint8 (also one byte)
     if img.dtype == np.bool:
         img = 255 * img.astype(np.uint8)
@@ -150,10 +204,10 @@ def imshow(
     if img.ndim == 2:
         trace = go.Heatmap(z=img, coloraxis="coloraxis1")
         autorange = True if origin == "lower" else "reversed"
-        layout = dict(
-            xaxis=dict(scaleanchor="y", constrain="domain"),
-            yaxis=dict(autorange=autorange, constrain="domain"),
-        )
+        layout = dict(yaxis=dict(autorange=autorange))
+        if aspect == "equal":
+            layout["xaxis"] = dict(scaleanchor="y", constrain="domain")
+            layout["yaxis"]["constrain"] = "domain"
         colorscale_validator = ColorscaleValidator("colorscale", "imshow")
         if zmin is not None and zmax is None:
             zmax = img.max()
@@ -167,6 +221,7 @@ def imshow(
             cmid=color_continuous_midpoint,
             cmin=range_color[0],
             cmax=range_color[1],
+            colorbar=dict(title=z_name),
         )
 
     # For 2D+RGB data, use Image trace
@@ -185,12 +240,25 @@ def imshow(
         )
 
     layout_patch = dict()
-    for v in ["title", "height", "width"]:
-        if args[v]:
-            layout_patch[v] = args[v]
+    for attr_name in ["title", "height", "width"]:
+        if args[attr_name]:
+            layout_patch[attr_name] = args[attr_name]
     if "title" not in layout_patch and args["template"].layout.margin.t is None:
         layout_patch["margin"] = {"t": 60}
     fig = go.Figure(data=trace, layout=layout)
     fig.update_layout(layout_patch)
+    if img_is_xarray:
+        if img.ndim <= 2:
+            hovertemplate = (
+                x_label
+                + ": %{x} <br>"
+                + y_label
+                + ": %{y} <br>"
+                + z_name
+                + " : %{z}<extra></extra>"
+            )
+            fig.update_traces(x=x, y=y, hovertemplate=hovertemplate)
+        fig.update_xaxes(title_text=x_label)
+        fig.update_yaxes(title_text=y_label)
     fig.update_layout(template=args["template"], overwrite=True)
     return fig

packages/python/plotly/plotly/tests/test_core/test_px/test_imshow.py

@@ -1,6 +1,7 @@
 import plotly.express as px
 import numpy as np
 import pytest
+import xarray as xr
 
 img_rgb = np.array([[[255, 0, 0], [0, 255, 0], [0, 0, 255]]], dtype=np.uint8)
 img_gray = np.arange(100).reshape((10, 10))
@@ -58,8 +59,9 @@ def test_colorscale():
 
 def test_wrong_dimensions():
     imgs = [1, np.ones((5,) * 3), np.ones((5,) * 4)]
+    msg = "px.imshow only accepts 2D single-channel, RGB or RGBA images."
     for img in imgs:
-        with pytest.raises(ValueError) as err_msg:
+        with pytest.raises(ValueError, match=msg):
             fig = px.imshow(img)
 
 
@@ -114,3 +116,13 @@ def test_zmin_zmax_range_color():
     fig = px.imshow(img, zmax=0.8)
     assert fig.layout.coloraxis.cmin == 0.0
     assert fig.layout.coloraxis.cmax == 0.8
+
+
+def test_imshow_xarray():
+    img = np.random.random((20, 30))
+    da = xr.DataArray(img, dims=["dim_rows", "dim_cols"])
+    fig = px.imshow(da)
+    # Dimensions are used for axis labels and coordinates
+    assert fig.layout.xaxis.title.text == "dim_cols"
+    assert fig.layout.yaxis.title.text == "dim_rows"
+    assert np.all(np.array(fig.data[0].x) == np.array(da.coords["dim_cols"]))

packages/python/plotly/tox.ini

@@ -58,6 +58,7 @@ deps=
     retrying==1.3.3
     pytest==3.5.1
     pandas==0.24.2
+    xarray==0.10.9
     backports.tempfile==1.0
     optional: --editable=file:///{toxinidir}/../plotly-geo
     optional: numpy==1.16.5
@@ -71,7 +72,6 @@ deps=
     optional: pyshp==1.2.10
     optional: pillow==5.2.0
     optional: matplotlib==2.2.3
-    optional: xarray==0.10.9
     optional: scikit-image==0.14.4
 
 ; CORE ENVIRONMENTS