Use functions from solver module in Diffusion class

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."
    ]
   },

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,19 +9,26 @@ 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()
     >>> m.time
     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

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

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