diff --git a/docs/notebooks/dwi_simulated_gp.ipynb b/docs/notebooks/dwi_simulated_gp.ipynb
index 7ddfe8e8..e1a4ba76 100644
--- a/docs/notebooks/dwi_simulated_gp.ipynb
+++ b/docs/notebooks/dwi_simulated_gp.ipynb
@@ -65,9 +65,7 @@
     "        gtab, mevals, S0=S0, angles=angles, fractions=fractions, snr=snr\n",
     "    )\n",
     "\n",
-    "    grad = np.vstack([gtab.bvecs.T, gtab.bvals])\n",
-    "\n",
-    "    return signal, sticks, grad"
+    "    return signal, sticks, gtab"
    ],
    "id": "a0f5bab019855954",
    "outputs": [],
@@ -84,12 +82,96 @@
    "cell_type": "code",
    "source": [
     "hsph_dirs = 90\n",
-    "signal, sticks, grad = create_multitensor_dmri_signal(hsph_dirs)"
+    "signal, sticks, gtab = create_multitensor_dmri_signal(hsph_dirs)"
    ],
    "id": "e9545781fe5cf3b8",
    "outputs": [],
    "execution_count": null
   },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Define the method that computes the signal attenuation.",
+   "id": "40b8ba2f818a49f8"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "def diffusion_distance(bvecs, Q):\n",
+    "    # Borrowed from https://github.com/arokem/osmosis/blob/72d848ddc172ce6a72baba4c0d4cf3675db1a381/osmosis/tensor.py#L233-L248\n",
+    "    # sphADC = np.dot(np.dot(bvecs.T, np.matrix(Q).getI()), bvecs)\n",
+    "    sphADC = np.dot(np.dot(bvecs.T, np.linalg.inv(Q)), bvecs)\n",
+    "    dist = np.diag(1 / np.sqrt(sphADC))\n",
+    "\n",
+    "    return dist"
+   ],
+   "id": "d40db09fa860c25e",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Perform the DTI fit on the signal and compute the signal attenuation along each diffusion-encoding gradient direction.",
+   "id": "25f32e883f520b86"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import dipy.reconst.dti as dti\n",
+    "\n",
+    "dti_model = dti.TensorModel(gtab)\n",
+    "dti_fit = dti_model.fit(signal)\n",
+    "\n",
+    "distance = diffusion_distance(gtab.bvecs.T, dti_fit.quadratic_form)"
+   ],
+   "id": "ba6f132d088a1c0a",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "We now define a function to plot the DWI signal.",
+   "id": "e9c2295d90cfbfc1"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "from matplotlib import pyplot as plt \n",
+    "%matplotlib inline\n",
+    "\n",
+    "def plot_dwi(dwi_data, bvecs, b0s_mask, voxel_idx, b0_idx, distance):\n",
+    "\n",
+    "    s0_voxel = dwi_data[voxel_idx[0], voxel_idx[1], voxel_idx[2], b0_idx]\n",
+    "    s_voxel = dwi_data[voxel_idx[0], voxel_idx[1], voxel_idx[2], ~b0s_mask]\n",
+    "\n",
+    "    # Scale gradient vector values with the distance\n",
+    "    # s_normal = s_voxel/s0_voxel\n",
+    "    # scaled_vectors = bvecs[:, ~b0s_mask] / s_normal[np.newaxis, :]\n",
+    "    scaled_vectors = bvecs[:, ~b0s_mask] / distance[~b0s_mask]\n",
+    "\n",
+    "    # Plot the DWI signal as a 3D point cloud\n",
+    "    fig = plt.figure()\n",
+    "    ax = fig.add_subplot(111, projection=\"3d\")\n",
+    "\n",
+    "    _ = ax.scatter(scaled_vectors[0, :], scaled_vectors[1, :], scaled_vectors[2, :], c=\"blue\", marker=\"*\")\n",
+    "\n",
+    "    # Set labels\n",
+    "    ax.set_xlabel(\"X\")\n",
+    "    ax.set_ylabel(\"Y\")\n",
+    "    ax.set_zlabel(\"Z\")\n",
+    "    ax.set_title(\"Normalized dMRI signal\")\n",
+    "\n",
+    "    return fig"
+   ],
+   "id": "1ed5c04d4b49d8d9",
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "metadata": {},
    "cell_type": "markdown",
@@ -99,13 +181,17 @@
   {
    "metadata": {},
    "cell_type": "code",
-   "execution_count": null,
    "source": [
     "voxel_idx = [0, 0, 0]\n",
-    "dwi_data = signal[np.newaxis, np.newaxis, np.newaxis, :]"
+    "dwi_data = signal[np.newaxis, np.newaxis, np.newaxis, :]\n",
+    "# There is only one b0 volume in the simulated signal\n",
+    "b0_idx = 0\n",
+    "_ = plot_dwi(dwi_data, gtab.bvecs.T, gtab.b0s_mask, voxel_idx, b0_idx, distance)\n",
+    "plt.show()"
    ],
    "id": "c07f103d9bd347cc",
-   "outputs": []
+   "outputs": [],
+   "execution_count": null
   },
   {
    "metadata": {},