Precompute A2_minus.dot(A1_minus) and so on in make_update_func about butterfly-algorithm

[dhermes]

As of 76ba251

--- butterfly_algorithm.py  2015-02-17 22:48:16.833886489 -0800
+++ butterfly_algorithm.py  2015-02-17 22:51:38.929889758 -0800
@@ -176,6 +176,11 @@

 def make_update_func(A1_minus, A1_plus, A2_minus, A2_plus, delta_T):

+    top_left = A2_minus.dot(A1_minus)
+    top_right = A2_plus.dot(A1_minus)
+    bottom_left = A2_minus.dot(A1_plus)
+    bottom_right = A2_plus.dot(A1_plus)
+
     def update_func(left_val, right_val):
         tau, sigma, alpha_vals_left = left_val
         tau_same, sigma_prime, alpha_vals_right = right_val
@@ -194,12 +199,12 @@
         k_plus_sigma_prime = 1.0 / k_minus_sigma_prime

         new_alpha_vals_left = (
-            k_minus_sigma * A2_minus.dot(A1_minus.dot(alpha_vals_left)) +
-            k_minus_sigma_prime * A2_plus.dot(A1_minus.dot(alpha_vals_right))
+            k_minus_sigma * top_left.dot(alpha_vals_left) +
+            k_minus_sigma_prime * top_right.dot(alpha_vals_right)
         )
         new_alpha_vals_right = (
-            k_plus_sigma * A2_minus.dot(A1_plus.dot(alpha_vals_left)) +
-            k_plus_sigma_prime * A2_plus.dot(A1_plus.dot(alpha_vals_right))
+            k_plus_sigma * bottom_left.dot(alpha_vals_left) +
+            k_plus_sigma_prime * bottom_right.dot(alpha_vals_right)
         )

         new_left_val = (tau_left, sigma_minus, new_alpha_vals_left)

[dhermes]

An even larger speedup can be achieved by considering the special form of these products and noting that the ++ and -- combinations are paired as well as the +- and -+ combinations:

--- butterfly_algorithm.py  2015-02-17 22:48:16.833886489 -0800
+++ butterfly_algorithm.py  2015-02-17 22:52:42.885890793 -0800
@@ -174,7 +174,7 @@
     return result


-def make_update_func(A1_minus, A1_plus, A2_minus, A2_plus, delta_T):
+def make_update_func(top_left, top_right, bottom_left, bottom_right, delta_T):

     def update_func(left_val, right_val):
         tau, sigma, alpha_vals_left = left_val
@@ -194,12 +194,12 @@
         k_plus_sigma_prime = 1.0 / k_minus_sigma_prime

         new_alpha_vals_left = (
-            k_minus_sigma * A2_minus.dot(A1_minus.dot(alpha_vals_left)) +
-            k_minus_sigma_prime * A2_plus.dot(A1_minus.dot(alpha_vals_right))
+            k_minus_sigma * top_left.dot(alpha_vals_left) +
+            k_minus_sigma_prime * top_right.dot(alpha_vals_right)
         )
         new_alpha_vals_right = (
-            k_plus_sigma * A2_minus.dot(A1_plus.dot(alpha_vals_left)) +
-            k_plus_sigma_prime * A2_plus.dot(A1_plus.dot(alpha_vals_right))
+            k_plus_sigma * bottom_left.dot(alpha_vals_left) +
+            k_plus_sigma_prime * bottom_right.dot(alpha_vals_right)
         )

         new_left_val = (tau_left, sigma_minus, new_alpha_vals_left)
@@ -225,9 +225,14 @@
     A2_minus = A2(M, -delta_S)
     A2_plus = A2(M, delta_S)

+    top_left = A2_minus.dot(A1_minus)
+    top_right = A2_plus.dot(A1_minus)
+    bottom_left = A2_minus.dot(A1_plus)
+    bottom_right = A2_plus.dot(A1_plus)
+
     for ell in xrange(1, L + 1):
-        update_func = make_update_func(A1_minus, A1_plus, A2_minus,
-                                       A2_plus, delta_T)
+        update_func = make_update_func(top_left, top_right, bottom_left,
+                                       bottom_right, delta_T)
         coeff_vals = increase_tau_refinement(coeff_vals, num_tau,
                                              num_sigma, update_func)

@@ -235,9 +240,14 @@
         num_sigma = num_sigma / 2
         delta_T *= 0.5

-        A_update(A1_plus, scale_multiplier=0.5, upper_diags=True)
-        A_update(A1_minus, scale_multiplier=0.5, upper_diags=True)
-        A_update(A2_plus, scale_multiplier=2.0, upper_diags=False)
-        A_update(A2_minus, scale_multiplier=2.0, upper_diags=False)
+        A_update(top_left, scale_multiplier=0.5, upper_diags=True)
+        A_update(top_left, scale_multiplier=2.0, upper_diags=False)
+        A_update(top_right, scale_multiplier=0.5, upper_diags=True)
+        A_update(top_right, scale_multiplier=2.0, upper_diags=False)
+
+        A_update(bottom_left, scale_multiplier=0.5, upper_diags=True)
+        A_update(bottom_left, scale_multiplier=2.0, upper_diags=False)
+        A_update(bottom_right, scale_multiplier=0.5, upper_diags=True)
+        A_update(bottom_right, scale_multiplier=2.0, upper_diags=False)

     return compute_t_by_bins(t, coeff_vals, L)

[dhermes]

Speedups:

Running ./butterfly_on_whale.py --M=50 --L=10, total computation time is:

• Current code: 1.0735
• Optimization in make_update_func: 0.832383
• Optimization involving scaling the precomputed produces: 0.851653

So the "simple" optimization wins out