diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml
index a59b3664..d395f0ca 100644
--- a/.github/workflows/unit-tests.yaml
+++ b/.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: |
diff --git a/docs/examples/vis.ipynb b/docs/examples/vis.ipynb
new file mode 100644
index 00000000..2109ff99
--- /dev/null
+++ b/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
+}
diff --git a/pipit/trace.py b/pipit/trace.py
index 75946931..a84bec86 100644
--- a/pipit/trace.py
+++ b/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)
diff --git a/pipit/util/config.py b/pipit/util/config.py
index f71bb53f..c3f77882 100644
--- a/pipit/util/config.py
+++ b/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}
diff --git a/pipit/vis/__init__.py b/pipit/vis/__init__.py
new file mode 100644
index 00000000..12279d52
--- /dev/null
+++ b/pipit/vis/__init__.py
@@ -0,0 +1 @@
+from .core import plot_comm_matrix, plot_message_histogram  # noqa: F401
diff --git a/pipit/vis/core.py b/pipit/vis/core.py
new file mode 100644
index 00000000..7b413a12
--- /dev/null
+++ b/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)
diff --git a/pipit/vis/util.py b/pipit/vis/util.py
new file mode 100644
index 00000000..647b60d6
--- /dev/null
+++ b/pipit/vis/util.py
@@ -0,0 +1,116 @@
+import yaml
+from bokeh.models import BasicTicker, CustomJSHover, CustomJSTickFormatter
+from bokeh.plotting import output_notebook
+from bokeh.plotting import show as bk_show
+from bokeh.themes import Theme
+import pipit as pp
+
+# Custom tickers and formatters
+
+# JS expression to convert bytes to human-readable string
+# "x" is the value (in bytes) being compared
+# "y" is the value (in bytes) being formatted
+JS_FORMAT_SIZE = """
+    if(x < 1e3)
+        return (y).toFixed(2) + " B";
+    if(x < 1e6)
+        return (y / 1e3).toFixed(2) + " kB";
+    if(x < 1e9)
+        return (y / 1e6).toFixed(2) + " MB";
+    if(x < 1e12)
+        return (y / 1e9).toFixed(2) + " GB";
+    if(x < 1e15)
+        return (y / 1e12).toFixed(2) + " TB";
+    else
+        return (y / 1e15).toFixed(2) + " PB";
+"""
+
+
+def get_process_ticker(nranks):
+    return BasicTicker(
+        base=2, desired_num_ticks=min(nranks, 16), min_interval=1, num_minor_ticks=0
+    )
+
+
+def get_size_hover_formatter():
+    return CustomJSHover(
+        code=f"""
+            let x = value;
+            let y = value;
+            {JS_FORMAT_SIZE}
+        """
+    )
+
+
+def get_size_tick_formatter(ignore_range=False):
+    x = "tick" if ignore_range else "Math.max(...ticks) - Math.min(...ticks);"
+    return CustomJSTickFormatter(
+        code=f"""
+            let x = {x}
+            let y = tick;
+            {JS_FORMAT_SIZE}
+        """
+    )
+
+
+# Helper functions
+def in_notebook():
+    """Returns True if we are in notebook environment, False otherwise"""
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True
+
+
+def show(p, return_fig=False):
+    """Used to wrap return values of plotting functions.
+
+    If return_figure is True, then just returns the figure object, otherwise starts a
+    Bokeh server containing the figure. If we are in a notebook, displays the
+    figure in the output cell, otherwise shows figure in new browser tab.
+
+    See https://docs.bokeh.org/en/latest/docs/user_guide/output/jupyter.html#bokeh-server-applications,   # noqa E501
+    https://docs.bokeh.org/en/latest/docs/user_guide/server/library.html.
+    """
+    if return_fig:
+        return p
+
+    # Create a Bokeh app containing the figure
+    def bkapp(doc):
+        doc.clear()
+        doc.add_root(p)
+        doc.theme = Theme(
+            json=yaml.load(
+                pp.get_option("theme"),
+                Loader=yaml.FullLoader,
+            )
+        )
+
+    if in_notebook():
+        # If notebook, show it in output cell
+        output_notebook(hide_banner=True)
+        bk_show(bkapp, notebook_url=pp.get_option("notebook_url"))
+    else:
+        # If standalone, start HTTP server and show in browser
+        from bokeh.server.server import Server
+
+        server = Server({"/": bkapp}, port=0, allow_websocket_origin=["*"])
+        server.start()
+        server.io_loop.add_callback(server.show, "/")
+        server.io_loop.start()
+
+
+def clamp(value, min_val, max_val):
+    """Clamps value to min and max bounds"""
+
+    if value < min_val:
+        return min_val
+    if value > max_val:
+        return max_val
+    return value
diff --git a/requirements.txt b/requirements.txt
index 1f317793..7ad07110 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 numpy
 otf2
 pandas
+bokeh
\ No newline at end of file