diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 26e67bc..9ae5e79 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -8,11 +8,11 @@ repos:
     - id: trailing-whitespace
 
   - repo: https://github.com/psf/black
-    rev: 21.6b0
+    rev: 21.7b0
     hooks:
     - id: black
 
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.4.0
+    rev: 0.5.0
     hooks:
     - id: nbstripout
diff --git a/docs/conf.py b/docs/conf.py
index 1e9f5a2..daa25cf 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -40,6 +40,13 @@
 
 autosummary_generate = True
 
+nbsphinx_prolog = """
+|colab-icon|_ `Interactive online version <https://colab.research.google.com/github/spatial-model-editor/sme_contrib/blob/master/docs/{{ env.doc2path(env.docname, base=None) }}>`_
+
+.. |colab-icon| image:: https://colab.research.google.com/assets/colab-badge.svg
+.. _colab-icon: https://colab.research.google.com/github/spatial-model-editor/sme_contrib/blob/master/docs/{{ env.doc2path(env.docname, base=None) }}
+"""
+
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]
 
diff --git a/docs/notebooks/plot.ipynb b/docs/notebooks/plot.ipynb
index e98b7b5..3bf880e 100644
--- a/docs/notebooks/plot.ipynb
+++ b/docs/notebooks/plot.ipynb
@@ -15,6 +15,7 @@
    },
    "outputs": [],
    "source": [
+    "!pip install -q sme_contrib\n",
     "import sme\n",
     "import sme_contrib.plot as smeplot\n",
     "from matplotlib import pyplot as plt"
diff --git a/sme_contrib/tests/test_optimize.py b/sme_contrib/tests/test_optimize.py
index 33f829e..124aa5b 100644
--- a/sme_contrib/tests/test_optimize.py
+++ b/sme_contrib/tests/test_optimize.py
@@ -1,4 +1,5 @@
 import sme_contrib.optimize as opt
+import matplotlib
 import numdifftools
 import numpy as np
 import os.path
@@ -93,3 +94,24 @@ def test_steady_state() -> None:
         ss.get_model().compartments["Nucleus"].species["A_nucl"].diffusion_constant
         == params[0]
     )
+
+    # just check plots have correct titles for now
+    ax = ss.plot_target_concentration()
+    assert type(ax) == matplotlib.axes._subplots.Subplot
+    assert ax.get_title() == "Target Concentration"
+
+    ax = ss.plot_model_concentration()
+    assert type(ax) == matplotlib.axes._subplots.Subplot
+    assert ax.get_title() == "Model Concentration"
+
+    ax = ss.plot_cost_history()
+    assert type(ax) == matplotlib.axes._subplots.Subplot
+    assert ax.get_title() == "Best cost history"
+
+    ax = ss.plot_cost_history_pbest()
+    assert type(ax) == matplotlib.axes._subplots.Subplot
+    assert ax.get_title() == "Mean particle best cost history"
+
+    ax = ss.plot_timeseries(1, 1)
+    assert type(ax) == matplotlib.axes._subplots.Subplot
+    assert ax.get_title() == "Concentration time series"