diff --git a/floris/simulation/flow_field.py b/floris/simulation/flow_field.py
index aae7e028d..2e6011a12 100644
--- a/floris/simulation/flow_field.py
+++ b/floris/simulation/flow_field.py
@@ -183,10 +183,7 @@ def calculate_speed_ups(self, het_map, x, y, z=None):
             # Calculate the 3-dimensional speed ups; reshape is needed as the generator
             # adds an extra dimension
             speed_ups = np.reshape(
-                [
-                    het_map[0][i](x[i:i+1,:,:,:,:], y[i:i+1,:,:,:,:], z[i:i+1,:,:,:,:])
-                    for i in range(len(het_map[0]))
-                ],
+                [het_map[0][i](x[i:i+1], y[i:i+1], z[i:i+1]) for i in range( len(het_map[0]))],
                 np.shape(x)
             )
 
@@ -195,10 +192,7 @@ def calculate_speed_ups(self, het_map, x, y, z=None):
             if np.isnan(speed_ups).any():
                 idx_nan = np.where(np.isnan(speed_ups))
                 speed_ups_out_of_region = np.reshape(
-                    [
-                        het_map[1][i](x[i:i+1,:,:,:,:], y[i:i+1,:,:,:,:], z[i:i+1,:,:,:,:])
-                        for i in range(len(het_map[1]))
-                    ],
+                    [het_map[1][i](x[i:i+1], y[i:i+1], z[i:i+1]) for i in range(len(het_map[1]))],
                     np.shape(x)
                 )
 
@@ -208,10 +202,7 @@ def calculate_speed_ups(self, het_map, x, y, z=None):
             # Calculate the 2-dimensional speed ups; reshape is needed as the generator
             # adds an extra dimension
             speed_ups = np.reshape(
-                [
-                    het_map[0][i](x[i:i+1,:,:,:,:], y[i:i+1,:,:,:,:])
-                    for i in range(len(het_map[0]))
-                ],
+                [het_map[0][i](x[i:i+1], y[i:i+1]) for i in range(len(het_map[0]))],
                 np.shape(x)
             )
 
@@ -220,10 +211,7 @@ def calculate_speed_ups(self, het_map, x, y, z=None):
             if np.isnan(speed_ups).any():
                 idx_nan = np.where(np.isnan(speed_ups))
                 speed_ups_out_of_region = np.reshape(
-                    [
-                        het_map[1][i](x[i:i+1,:,:,:,:], y[i:i+1,:,:,:,:])
-                        for i in range(len(het_map[1]))
-                    ],
+                    [het_map[1][i](x[i:i+1], y[i:i+1]) for i in range(len(het_map[1]))],
                     np.shape(x)
                 )