Vectorized 'rarefaction_wave' for Burgers.

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()'

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
+    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 uexact
+        return exact
 
     def implicit_burgers_solver(self, grid_values, mesh):
         """