Vis: setup and comm_matrix

.github/workflows/unit-tests.yaml

@@ -12,9 +12,9 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         include:
-        - python-version: "3.6"
+        - python-version: "3.7"
           os: ubuntu-20.04
 
     steps:
@@ -26,7 +26,7 @@ jobs:
       - name: Install Python packages
         run: |
           pip install --upgrade pip
-          pip install --upgrade numpy pandas pytest otf2
+          pip install --upgrade numpy pandas pytest otf2 bokeh datashader
 
       - name: Lint and format check with flake8 and black
         if: ${{ matrix.python-version == 3.9 }}
@@ -54,7 +54,7 @@ jobs:
       - name: Install Python packages
         run: |
           pip install --upgrade pip
-          pip install --upgrade numpy pandas pytest otf2
+          pip install --upgrade numpy pandas pytest otf2 bokeh datashader
 
       - name: Basic test with pytest
         run: |

docs/examples/vis.ipynb

@@ -0,0 +1,80 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5841460",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "import sys\n",
+    "\n",
+    "sys.path.append(\"../../\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da5ebec2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pipit as pp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7070057d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ping_pong = pp.Trace.from_otf2(\"../../pipit/tests/data/ping-pong-otf2\")\n",
+    "ping_pong"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2476f1f3-a6cf-49f2-8e3e-0e004f088504",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ping_pong.plot_comm_matrix()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a4aff933",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ping_pong.plot_message_histogram()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pipit/trace.py

@@ -548,7 +548,7 @@ def flat_profile(
                 self.events.loc[self.events["Event Type"] == "Enter"]
                 .groupby([groupby_column, "Process"], observed=True)[metrics]
                 .sum()
-                .groupby(groupby_column)
+                .groupby(groupby_column, observed=True)
                 .mean()
             )
 
@@ -865,3 +865,21 @@ def detect_pattern(
             patterns.append(match_original)
 
         return patterns
+
+    def plot_comm_matrix(self, output="size", *args, **kwargs):
+        from .vis import plot_comm_matrix
+
+        # Generate the data
+        data = self.comm_matrix(output=output)
+
+        # Return the Bokeh plot
+        return plot_comm_matrix(data, output=output, *args, **kwargs)
+
+    def plot_message_histogram(self, bins=20, *args, **kwargs):
+        from .vis import plot_message_histogram
+
+        # Generate the data
+        data = self.message_histogram(bins=bins)
+
+        # Return the Bokeh plot
+        return plot_message_histogram(data, *args, **kwargs)

pipit/util/config.py

@@ -83,6 +83,23 @@ def url_validator(key, value):
         )
 
 
+# Validator to check if theme is valid YAML
+def theme_validator(key, value):
+    import yaml
+
+    try:
+        yaml.safe_load(value)
+    except yaml.YAMLError:
+        raise ValueError(
+            (
+                'Error loading configuration: The Value "{}" for Configuration "{}"'
+                + "must be a valid YAML"
+            ).format(value, key)
+        )
+    else:
+        return True
+
+
 registered_options = {
     "log_level": {
         "default": "INFO",
@@ -92,6 +109,34 @@ def url_validator(key, value):
         "default": "http://localhost:8888",
         "validator": url_validator,
     },
+    "theme": {
+        "default": """
+            attrs:
+                Plot:
+                    height: 350
+                    width: 700
+                    background_fill_color: "#fafafa"
+                Axis:
+                    axis_label_text_font_style: "bold"
+                    minor_tick_line_color: null
+                Toolbar:
+                    autohide: true
+                    logo: null
+                HoverTool:
+                    point_policy: "follow_mouse"
+                Legend:
+                    label_text_font_size: "8.5pt"
+                    spacing: 6
+                    border_line_color: null
+                    glyph_width: 16
+                    glyph_height: 16
+                Scatter:
+                    size: 9
+                DataRange1d:
+                    range_padding: 0.05
+        """,
+        "validator": theme_validator,
+    },
 }
 
 global_config = {key: registered_options[key]["default"] for key in registered_options}

pipit/vis/__init__.py

@@ -0,0 +1 @@
+from .core import plot_comm_matrix, plot_message_histogram  # noqa: F401

pipit/vis/core.py

@@ -0,0 +1,148 @@
+import numpy as np
+from bokeh.models import (
+    ColorBar,
+    HoverTool,
+    LinearColorMapper,
+    LogColorMapper,
+    NumeralTickFormatter,
+)
+from bokeh.plotting import figure
+
+from .util import (
+    clamp,
+    get_process_ticker,
+    get_size_hover_formatter,
+    get_size_tick_formatter,
+    show,
+)
+
+
+def plot_comm_matrix(
+    data, output="size", cmap="log", palette="Viridis256", return_fig=False
+):
+    """Plots the trace's communication matrix.
+
+    Args:
+        data (numpy.ndarray): a 2D numpy array of shape (N, N) containing the
+            communication matrix between N processes.
+        output (str, optional): Specifies whether the matrix contains "size"
+            or "count" values. Defaults to "size".
+        cmap (str, optional): Specifies the color mapping. Options are "log",
+            "linear", and "any". Defaults to "log".
+        palette (str, optional): Name of Bokeh color palette to use. Defaults to
+            "Viridis256".
+        return_fig (bool, optional): Specifies whether to return the Bokeh figure
+            object. Defaults to False, which displays the result and returns nothing.
+
+    Returns:
+        Bokeh figure object if return_fig, None otherwise
+    """
+    nranks = data.shape[0]
+
+    # Define color mapper
+    if cmap == "linear":
+        color_mapper = LinearColorMapper(palette=palette, low=0, high=np.amax(data))
+    elif cmap == "log":
+        color_mapper = LogColorMapper(
+            palette=palette, low=max(np.amin(data), 1), high=np.amax(data)
+        )
+    elif cmap == "any":
+        color_mapper = LinearColorMapper(palette=palette, low=1, high=1)
+
+    # Create bokeh plot
+    p = figure(
+        x_axis_label="Receiver",
+        y_axis_label="Sender",
+        x_range=(-0.5, nranks - 0.5),
+        y_range=(nranks - 0.5, -0.5),
+        x_axis_location="above",
+        tools="hover,pan,reset,wheel_zoom,save",
+        width=90 + clamp(nranks * 30, 200, 500),
+        height=10 + clamp(nranks * 30, 200, 500),
+        toolbar_location="below",
+    )
+
+    # Add glyphs and layouts
+    p.image(
+        image=[np.flipud(data)],
+        x=-0.5,
+        y=-0.5,
+        dw=nranks,
+        dh=nranks,
+        color_mapper=color_mapper,
+        origin="top_left",
+    )
+
+    color_bar = ColorBar(
+        color_mapper=color_mapper,
+        formatter=(
+            get_size_tick_formatter(ignore_range=cmap == "log")
+            if output == "size"
+            else NumeralTickFormatter()
+        ),
+        width=15,
+    )
+    p.add_layout(color_bar, "right")
+
+    # Customize plot
+    p.axis.ticker = get_process_ticker(nranks=nranks)
+    p.grid.visible = False
+
+    # Configure hover
+    hover = p.select(HoverTool)
+    hover.tooltips = [
+        ("Sender", "$y{0.}"),
+        ("Receiver", "$x{0.}"),
+        ("Count", "@image") if output == "count" else ("Volume", "@image{custom}"),
+    ]
+    hover.formatters = {"@image": get_size_hover_formatter()}
+
+    # Return plot
+    return show(p, return_fig=return_fig)
+
+
+def plot_message_histogram(
+    data,
+    return_fig=False,
+):
+    """Plots the trace's message size histogram.
+
+    Args:
+        data (hist, edges): Histogram and edges
+        return_fig (bool, optional): Specifies whether to return the Bokeh figure
+            object. Defaults to False, which displays the result and returns nothing.
+
+    Returns:
+        Bokeh figure object if return_fig, None otherwise
+    """
+    hist, edges = data
+
+    # Create bokeh plot
+    p = figure(
+        x_axis_label="Message size",
+        y_axis_label="Number of messages",
+        tools="hover,save",
+    )
+    p.y_range.start = 0
+
+    # Add glyphs and layouts
+    p.quad(top=hist, bottom=0, left=edges[:-1], right=edges[1:])
+
+    # Customize plot
+    p.xaxis.formatter = get_size_tick_formatter()
+    p.yaxis.formatter = NumeralTickFormatter()
+    p.xgrid.visible = False
+
+    # Configure hover
+    hover = p.select(HoverTool)
+    hover.tooltips = [
+        ("Bin", "@left{custom} - @right{custom}"),
+        ("Count", "@top"),
+    ]
+    hover.formatters = {
+        "@left": get_size_hover_formatter(),
+        "@right": get_size_hover_formatter(),
+    }
+
+    # Return plot
+    return show(p, return_fig=return_fig)