diff --git a/Snakefile b/Snakefile
index 582a83bbc1248f93655e2cf213aab6b4eb5cc635..ea3fb9b83c650d2bb010cde3e0976ee6a6e090f7 100644
--- a/Snakefile
+++ b/Snakefile
@@ -34,7 +34,7 @@ TODO: Move height adjustment and stretch factor into implicit solver
 TODO: Enable choice of equation -> Done
 TODO: Ensure termination of Burgers with DiscontinuousConstant function -> Done
 TODO: Ensure termination of Burgers with Sine function -> Done
-TODO: Vectorize 'rarefaction_wave()'
+TODO: Vectorize 'rarefaction_wave()' -> Done
 TODO: Vectorize 'burgers_volume_integral()'
 TODO: Vectorize 'burgers_local_Lax_Friedrich()'
 TODO: Vectorize 'burgers_boundary_matrix()'
diff --git a/scripts/tcd/Equation.py b/scripts/tcd/Equation.py
index 7817d12b1eb57800b14e41fd5bffb2cf806cf645..783dd22e8bad4c1b05c6dee9cbec5978f89461e4 100644
--- a/scripts/tcd/Equation.py
+++ b/scripts/tcd/Equation.py
@@ -495,19 +495,32 @@ class Burgers(Equation):
 
         return grid, exact
 
-    def rarefaction_wave(self, grid_values):
-        uexact = 0 * grid_values
-        N = np.size(grid_values)
-
-        for i in range(N):
-            if grid_values[0, i] <= - self._final_time:
-                uexact[0, i] = -1
-            elif - self._final_time < grid_values[0, i] < self._final_time:
-                uexact[0, i] = grid_values[0, i] / self._final_time
-            else:
-                uexact[0, i] = 1
-
-        return uexact
+    def rarefaction_wave(self, grid):
+        """Calculate the exact solution for the rarefaction wave.
+
+        Parameters
+        ----------
+        grid : ndarray
+            Array containing evaluation grid for a function.
+
+        Returns
+        -------
+        exact : ndarray
+            Array containing exact evaluation of a function.
+
+        """
+        exact = np.ones_like(grid)
+
+        # Set all values for the left boundary
+        mask = grid[0] <= -self._final_time
+        exact[0, mask] = -1
+
+        # Set all values between the boundaries
+        mask = np.logical_and(-self._final_time < grid[0],
+                              grid[0] < self._final_time)
+        exact[0, mask] = grid[0, mask]/self.final_time
+
+        return exact
 
     def implicit_burgers_solver(self, grid_values, mesh):
         """