diff --git a/lessons/python/functions.ipynb b/lessons/python/functions.ipynb
index b23d99c..7cf58de 100644
--- a/lessons/python/functions.ipynb
+++ b/lessons/python/functions.ipynb
@@ -72,7 +72,7 @@
    "outputs": [],
    "source": [
     "def calculate_time_step(grid_spacing, diffusivity):\n",
-    "    return grid_spacing**2 / diffusivity / 2.1"
+    "    return grid_spacing**2 / diffusivity / 4.0"
    ]
   },
   {
@@ -95,7 +95,7 @@
     "the grid spacing of the model and the diffusivity.\n",
     "The variables `grid_spacing` and `diffusivity` are *local* to the function--they don't exist outside of the body of the function.\n",
     "In the body of the function,\n",
-    "the time step is calculated from the stability criterion\n",
+    "the time step is calculated from a stability criterion\n",
     "and returned to the caller."
    ]
   },
diff --git a/lessons/python/ivy-diffusion/ivy_diffusion/__init__.py b/lessons/python/ivy-diffusion/ivy_diffusion/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/lessons/python/ivy-diffusion/ivy_diffusion/diffusion.py b/lessons/python/ivy-diffusion/ivy_diffusion/diffusion.py
index ea4383a..980d52f 100644
--- a/lessons/python/ivy-diffusion/ivy_diffusion/diffusion.py
+++ b/lessons/python/ivy-diffusion/ivy_diffusion/diffusion.py
@@ -1,5 +1,5 @@
 """Modeling the one-dimensional diffusion equation."""
-import numpy as np
+from .solver import calculate_time_step, set_initial_profile, solve1d
 
 
 class Diffusion:
@@ -9,12 +9,14 @@ class Diffusion:
     Examples
     --------
     >>> import numpy as np
-    >>> from diffusion import Diffusion
+    >>> from .diffusion import Diffusion
     >>> m = Diffusion(diffusivity=0.25)
     >>> m.concentration = np.zeros(m.shape)
-    >>> m.concentration[int(m.shape/2)] = 5
+    >>> m.concentration[m.shape//2] = 5
     >>> m.concentration
     array([0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 0.0])
+    >>> m.time_step
+    1.0
     >>> m.time
     0.0
     >>> m.update()
@@ -22,6 +24,11 @@ class Diffusion:
     1.0
     >>> m.concentration
     array([0.0, 0.0, 0.0, 0.0, 1.2, 2.5, 1.2, 0.0, 0.0, 0.0])
+    >>> m.update()
+    >>> m.time
+    2.0
+    >>> m.concentration
+    array([0.0, 0.0, 0.0, 0.3, 1.2, 1.9, 1.2, 0.3, 0.0, 0.0])
     """
 
     def __init__(self, shape=10, spacing=1.0, diffusivity=1.0):
@@ -40,16 +47,10 @@ def __init__(self, shape=10, spacing=1.0, diffusivity=1.0):
         self.spacing = spacing
         self.diffusivity = diffusivity
         self.time = 0.0
-        self.time_step = self.spacing**2 / (4.0 * self.diffusivity)
-
-        self.concentration = np.random.random(self.shape)
-
-    def solve(self):
-        """Solve the 1D diffusion equation."""
-        flux = -self.diffusivity * np.diff(self.concentration) / self.spacing
-        self.concentration[1:-1] -= self.time_step * np.diff(flux) / self.spacing
+        self.time_step = calculate_time_step(self.spacing, self.diffusivity)
+        self.concentration = set_initial_profile(self.shape)
 
     def update(self):
         """Calculate concentration at the next time step."""
-        self.solve()
+        solve1d(self.concentration, self.spacing, self.time_step, self.diffusivity)
         self.time += self.time_step
diff --git a/lessons/python/ivy-diffusion/ivy_diffusion/solver.py b/lessons/python/ivy-diffusion/ivy_diffusion/solver.py
index a657dd3..e470ec2 100644
--- a/lessons/python/ivy-diffusion/ivy_diffusion/solver.py
+++ b/lessons/python/ivy-diffusion/ivy_diffusion/solver.py
@@ -5,7 +5,7 @@
 
 
 def calculate_time_step(grid_spacing, diffusivity):
-    return grid_spacing**2 / diffusivity / 2.1
+    return grid_spacing**2 / diffusivity / 4.0
 
 
 def set_initial_profile(domain_size=100, boundary_left=500, boundary_right=0):
@@ -37,7 +37,7 @@ def solve1d(concentration, grid_spacing=1.0, time_step=1.0, diffusivity=1.0):
     Examples
     --------
     >>> import numpy as np
-    >>> from solver import solve1d
+    >>> from .solver import solve1d
     >>> z = np.zeros(5)
     >>> z[2] = 5
     >>> z
diff --git a/lessons/python/ivy-diffusion/tests/test_solver.py b/lessons/python/ivy-diffusion/tests/test_solver.py
index 7a48c8b..9c0535b 100644
--- a/lessons/python/ivy-diffusion/tests/test_solver.py
+++ b/lessons/python/ivy-diffusion/tests/test_solver.py
@@ -7,7 +7,7 @@
 DOMAIN_SIZE = 100
 GRID_SPACING = 1.0
 DIFFUSIVITY = 1.0
-TIME_STEP = 0.475
+TIME_STEP = 0.25
 TOLERANCE = 0.01
 ZMAX = 500.0