Repository: dask/dask

Issue:
Categorise high level layers for display
There are currently a concrete number of subclasses of the base highlevelgraph.Layer. Some of these have specific contexts or collection linkage (array Vs dataframe), others do not. For the sake of the work being done by @freyam , it would be nice to create some categories for the purpose of being shown in `.visualize()`.

Layers allow for attaching attributes at instantiation. I suggests there might also be class attributes giving information about the layer type, which will be true for all instances.

Example: `DataFrameIOLayer` is IO by operation type and dataframe by collection.  It would be reasonable for these to be among the default annotations of all instances.



Discussion:
cc @GenevieveBuckley 
This is amazing! I will get started on finding more about this and share my findings here! :rocket: 
I suggest starting with a concrete list of the current known implementations of Layer.
# List of Layers

## General
```py
# High level graph layer
class Layer(collections.abc.Mapping)

# Fully materialized layer of `Layer`
class MaterializedLayer(Layer)

# Tensor Operation
class Blockwise(Layer)
```

## Array Layers

```py
# Specialized Blockwise Layer for array creation routines
class BlockwiseCreateArray(Blockwise)

# Simple HighLevelGraph array overlap layer
class ArrayOverlapLayer(Layer)
```

## Dataframe Layers

```py
# DataFrame-based HighLevelGraph Layer
class DataFrameLayer(Layer)

# High-level graph layer for a simple shuffle operation in which each output partition depends on all input partitions
class SimpleShuffleLayer(DataFrameLayer)

# High-level graph layer corresponding to a single stage of a multi-stage inter-partition shuffle operation.
class ShuffleLayer(SimpleShuffleLayer)

# High-level graph layer for a join operation requiring the smaller collection to be broadcasted to every partition of the larger collection.
class BroadcastJoinLayer(DataFrameLayer)

# DataFrame-based Blockwise Layer with IO
class DataFrameIOLayer(Blockwise, DataFrameLayer)
```

Copying over my comment from the Slack thread earlier today:

> I think figuring out what categories should be included will be the tricky thing, and trying to make those categories be something users will care about. It seems like it'd be more useful if we can use them to highlight bottlenecks. I assume that's why you suggest IO and shuffle. We seem a bit vague on things beyond that though, I don't know that "all other CPU computations" is likely to be very useful to display (and will probably be most cases)
So I think I'm more advocating that the Layer classes have some default annotations - and that's all. That way, we only need to look at these annotations, we don't need to do any `isinstance` checks. Of course, we still *might* want to look into the layer contents, to make the difference between GPU/CPU.

Totally agree that "CPU" should not have any associated decoration, as it will be the default, most common box.
@martindurant 
What would the default annotations be for all the layers? I can think of `layer_id`.

Meanwhile, I have come up with the visualization for the HLGs!

## Design
Group the listed layer_types into 4 categories

### IO (DataFrame)
1. `DataFrameIOLayer`

### Shuffle (DataFrame)
1. `SimpleShuffleLayer`
2. `ShuffleLayer`

### Blockwise (Array mostly + DataFrame)
1. `Blockwise`
2. `BlockwiseCreateArray`

### Materialized (Array  + DataFrame)
1. `MaterializedLayer`

The remaining layers:
1. `ArrayOverlapLayer`
2. `DataFrameLayer`
3. `BroadcastJoinLayer`
can be just left as it is.

### Visualization
 
#### Graphviz (`Layer.dask.visualize()`)
We can show the different groups with different colors.

EDIT: Updated Graph Reprentation
![image](https://user-images.githubusercontent.com/62539811/126983749-f07450fd-b36c-437f-a5a7-8256c68af0bf.png)


### HTML Repr (`Layer.dask`)
Using the same colors used in the Graphviz, we change the tiny circles next to each layer to reflect its `layer_type`.

EDIT: Updated HTML Reprentation
![image](https://user-images.githubusercontent.com/62539811/126983834-cbe441c2-100d-4f54-9e5d-e0fa2501c987.png)


(I have also modified the `MaterializedLayer`'s existing color to be a little lighter. I feel it looks nicer and easier on the eyes)
It might seem too colorful, but that's because I am showing all the different types of layers at the same time. This is highly unlikely to been observed in real Dask Graphs.


/cc @jacobtomlinson I would love to hear your opinion as well on the colors. You have done an amazing job adding all the existing colors :sparkles: 
> What would the default annotations be for all the layers

Empty? If we can't think of anything useful, it's better not to complicate the visuals.
/cc @mrocklin What do you think about the use of colors?
(colors are open for discussion)
I hadn't thought of changing the colours in the dots next to the layer titles in the HTML representation, that's a nice touch.

Personally, I'm not sure we're gaining much by having multiple shades of green/blue/etc. for different types of layers in the same larger category. I think that adds more confusion.

We'll need some effort spent on:
1. Explaining what information the user should gain from this (eg: shuffles are inefficient so you don't want to see too many of them, blockwise layers are likely to be nicely parallelized, etc.). You'll need to be able to explain this very clearly for all of them.
2. How we are going to communicate that information to the user (a legend for graphviz? added information in the docs? something in the HTML repr?)

> Personally, I'm not sure we're gaining much by having multiple shades of green/blue/etc. for different types of layers in the same larger category. I think that adds more confusion.

You are right! This looks even cleaner :sparkles: 

I have also added a Legend below so users have a good idea of what the colors mean.
![image](https://user-images.githubusercontent.com/62539811/126964030-42093e8b-5327-4aca-a01e-fb569861a8d3.png)

and 

![image](https://user-images.githubusercontent.com/62539811/126961117-18664988-844f-4024-9e48-d6ca9a6bb5d9.png)

Note: The colors used in Layer 2, 5, and 6 are supposed to be Green (Blockwise) as well, but I have used different ones to show all of the colors


-----

> Explaining what information the user should gain from this (eg: shuffles are inefficient so you don't want to see too many of them, blockwise layers are likely to be nicely parallelized, etc.). You'll need to be able to explain this very clearly for all of them.

I'm unfamiliar with this section of the code. So I'm not sure which layers are more difficult for users and which are simple (quick?).

If I have more context about the different layers and what makes them slower than the rest, I can perhaps do something more with those layers' nodes.

For example, I can add a little node form **note** to all of the Shuffle and DataFrameIO layers. This could imply that this is an area worth investigating.

![image](https://user-images.githubusercontent.com/62539811/126968583-9fead6b3-a968-4406-8a5e-f7fc5a33f999.png)

How does this sound?








> For example, I can add a little node form **note** to all of the Shuffle and DataFrameIO layers. 

I'm sorry, what does this mean?
I think it means that the outline box is not just a rectangle, but has a little fold in the corner. Pretty subtle.

On July 26, 2021 10:28:04 PM EDT, Genevieve Buckley ***@***.***> wrote:
>> For example, I can add a little node form **note** to all of the
>Shuffle and DataFrameIO layers. 
>
>I'm sorry, what does this mean?
>
>-- 
>You are receiving this because you were mentioned.
>Reply to this email directly or view it on GitHub:
>https://github.com/dask/dask/issues/7919#issuecomment-887159227

-- 
Sent from my Android device with K-9 Mail. Please excuse my brevity.
@GenevieveBuckley Yes, Martin is correct.

So, in essence, the current node shape is "box." Another node shape is "note," which resembles a paper fold from the corner.

Normally, people fold a page to bookmark it so that it can be easily found later. In this case, the fold would indicate that this area warrants further investigation. 
> /cc @jacobtomlinson I would love to hear your opinion as well on the colors. You have done an amazing job adding all the existing colors ✨

Thanks @freyam but all the colour design was done by an external designer. See dask/community#135. I've added another comment to that issue to see if we can get the designer involved in choosing the more colours.
> I've added another comment to that issue to see if we can get the designer involved in choosing the more colours.

That's amazing! :rocket: 
> So, in essence, the current node shape is "box." Another node shape is "note," which resembles a paper fold from the corner.

I think that the difference is too subtle to see in your examples, and the meaning of the different shapes will not be obvious to users.
Just wanna confirm:
Shuffle and DataFrameIO layers are the 2 specific layers you want to highlight differently to point at the potential source of inefficiency?

Resolving patch:
diff --git a/dask/base.py b/dask/base.py
--- a/dask/base.py
+++ b/dask/base.py
@@ -571,15 +571,15 @@ def compute(*args, **kwargs):
 
 def visualize(*args, **kwargs):
     """
-    Visualize several dask graphs at once.
+    Visualize several low level dask graphs at once.
 
     Requires ``graphviz`` to be installed. All options that are not the dask
     graph(s) should be passed as keyword arguments.
 
     Parameters
     ----------
-    dsk : dict(s) or collection(s)
-        The dask graph(s) to visualize.
+    args : dict(s) or collection(s)
+        The low level dask graph(s) to visualize.
     filename : str or None, optional
         The name of the file to write to disk. If the provided `filename`
         doesn't include an extension, '.png' will be used by default.
@@ -591,8 +591,10 @@ def visualize(*args, **kwargs):
         If True, the graph is optimized before rendering.  Otherwise,
         the graph is displayed as is. Default is False.
     color : {None, 'order'}, optional
-        Options to color nodes.  Provide ``cmap=`` keyword for additional
+        Options to color nodes.
         colormap
+        - None, the default, no colors.
+        - 'order', colors the nodes' border based on the order they appear in the graph
     collapse_outputs : bool, optional
         Whether to collapse output boxes, which often have empty labels.
         Default is False.
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -870,6 +870,43 @@ def merge(cls, *graphs):
         return cls(layers, dependencies)
 
     def visualize(self, filename="dask-hlg.svg", format=None, **kwargs):
+        """
+        Visualize this dask high level graph.
+
+        Requires ``graphviz`` to be installed.
+
+        Parameters
+        ----------
+        filename : str or None, optional
+            The name of the file to write to disk. If the provided `filename`
+            doesn't include an extension, '.png' will be used by default.
+            If `filename` is None, no file will be written, and the graph is
+            rendered in the Jupyter notebook only.
+        format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
+            Format in which to write output file. Default is 'svg'.
+        color : {None, 'layer_type'}, optional (default: None)
+            Options to color nodes.
+            - None, no colors.
+            - layer_type, color nodes based on the layer type.
+        **kwargs
+           Additional keyword arguments to forward to ``to_graphviz``.
+
+        Examples
+        --------
+        >>> x.dask.visualize(filename='dask.svg')  # doctest: +SKIP
+        >>> x.dask.visualize(filename='dask.svg', color='layer_type')  # doctest: +SKIP
+
+        Returns
+        -------
+        result : IPython.diplay.Image, IPython.display.SVG, or None
+            See dask.dot.dot_graph for more information.
+
+        See Also
+        --------
+        dask.dot.dot_graph
+        dask.base.visualize # low level variant
+        """
+
         from .dot import graphviz_to_file
 
         g = to_graphviz(self, **kwargs)
@@ -1228,18 +1265,34 @@ def to_graphviz(
     for layer in hg.dependencies:
         n_tasks[layer] = len(hg.layers[layer])
 
-    mn = min(n_tasks.values())
-    mx = max(n_tasks.values())
+    min_tasks = min(n_tasks.values())
+    max_tasks = max(n_tasks.values())
 
     cache = {}
 
+    color = kwargs.get("color")
+    if color == "layer_type":
+        layer_colors = {
+            "DataFrameIOLayer": ["#CCC7F9", False],  # purple
+            "ShuffleLayer": ["#F9CCC7", False],  # rose
+            "SimpleShuffleLayer": ["#F9CCC7", False],  # rose
+            "ArrayOverlayLayer": ["#FFD9F2", False],  # pink
+            "BroadcastJoinLayer": ["#D9F2FF", False],  # blue
+            "Blockwise": ["#D9FFE6", False],  # green
+            "BlockwiseLayer": ["#D9FFE6", False],  # green
+            "BlockwiseCreateArray": ["#D9FFE6", False],  # green
+            "MaterializedLayer": ["#DBDEE5", False],  # gray
+        }
+
     for layer in hg.dependencies:
         layer_name = name(layer)
         attrs = data_attributes.get(layer, {})
 
         node_label = label(layer, cache=cache)
         node_size = (
-            20 if mx == mn else int(20 + ((n_tasks[layer] - mn) / (mx - mn)) * 20)
+            20
+            if max_tasks == min_tasks
+            else int(20 + ((n_tasks[layer] - min_tasks) / (max_tasks - min_tasks)) * 20)
         )
 
         layer_type = str(type(hg.layers[layer]).__name__)
@@ -1271,6 +1324,13 @@ def to_graphviz(
         attrs.setdefault("fontsize", str(node_size))
         attrs.setdefault("tooltip", str(node_tooltips))
 
+        if color == "layer_type":
+            node_color = layer_colors.get(layer_type)[0]
+            layer_colors.get(layer_type)[1] = True
+
+            attrs.setdefault("fillcolor", str(node_color))
+            attrs.setdefault("style", "filled")
+
         g.node(layer_name, **attrs)
 
     for layer, deps in hg.dependencies.items():
@@ -1278,6 +1338,28 @@ def to_graphviz(
         for dep in deps:
             dep_name = name(dep)
             g.edge(dep_name, layer_name)
+
+    if color == "layer_type":
+        legend_title = "Key"
+
+        legend_label = (
+            '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5">'
+            "<TR><TD><B>Legend: Layer types</B></TD></TR>"
+        )
+
+        for layer_type, color in layer_colors.items():
+            if color[1]:
+                legend_label += f'<TR><TD BGCOLOR="{color[0]}">{layer_type}</TD></TR>'
+
+        legend_label += "</TABLE>>"
+
+        attrs = data_attributes.get(legend_title, {})
+        attrs.setdefault("label", str(legend_label))
+        attrs.setdefault("fontsize", "20")
+        attrs.setdefault("margin", "0")
+
+        g.node(legend_title, **attrs)
+
     return g
 
 


Test patch: