diff --git a/ch05/hydrogen_wavefunctions.ipynb b/ch05/hydrogen_wavefunctions.ipynb
index ac6291f9..30545366 100644
--- a/ch05/hydrogen_wavefunctions.ipynb
+++ b/ch05/hydrogen_wavefunctions.ipynb
@@ -42093,15 +42093,18 @@
         "    r = np.sqrt((x ** 2 + y ** 2))\n",
         "\n",
         "    # Ψnlm(r,θ,φ) = Rnl(r).Ylm(θ,φ)\n",
-        "    psi = radial_function(r, n, l) * sph_harm(m, l, 0, np.arctan(x / (y + 1e-10)))\n",
+        "    psi = radial_function(r, n, l) * sph_harm(m, l, 0, np.arctan(x / (y + 1e-7)))\n",
+        "    psi_sq = np.abs(psi)**2\n",
         "\n",
         "    # Plot orbitals\n",
         "    fig, ax = plt.subplots()\n",
-        "    ax.contour(normalized_wavefunction(n, l, m), 16, cmap='rocket', extent=[-max_r, max_r,-max_r, max_r])\n",
+        "    ax.contour(psi_sq, 40, cmap='RdBu', extent=[-max_r, max_r,-max_r, max_r])\n",
         "    ax.set_title(r'$|\\psi_{{({0}, {1}, {2})}}|^2$'.format(n, l, m), fontsize=15)\n",
         "    ax.set_aspect('equal')\n",
-        "\n",
-        "    #return np.abs(psi) ** 2"
+        "    ax.set_xlabel(r'$x$')\n",
+        "    ax.set_ylabel(r'$y$')\n",
+        "    ax.set_xlim(-max_r, max_r)\n",
+        "    ax.set_ylim(-max_r, max_r)"
       ]
     },
     {