ENH: replace einsum with matmul for speed

jeffgortmaker · Jan 1, 2022 · ddc696e · ddc696e
1 parent af1103d
commit ddc696e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/pyblp/markets/problem_market.py b/pyblp/markets/problem_market.py
@@ -75,9 +75,13 @@ def solve(
                 probabilities, conditionals, compute_conditionals_tensor=self.K3 > 0
             )
             for p in range(self.parameters.P):
-                probabilities_tangent_mapping[p] += np.einsum('jki,j->ki', probabilities_tensor, xi_jacobian[:, p])
+                probabilities_tangent_mapping[p] += np.squeeze(
+                    np.moveaxis(probabilities_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
+                )
                 if conditionals_tensor is not None:
-                    conditionals_tangent_mapping[p] += np.einsum('jki,j->ki', conditionals_tensor, xi_jacobian[:, p])
+                    conditionals_tangent_mapping[p] += np.squeeze(
+                        np.moveaxis(conditionals_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
+                    )
 
         # compute contributions to micro moments, their Jacobian, and their covariances
         if moments.MM == 0:

diff --git a/pyblp/markets/results_market.py b/pyblp/markets/results_market.py
@@ -73,9 +73,13 @@ def safely_compute_jacobian_realizations(self, tilde_costs: Array) -> Tuple[Arra
                 probabilities, conditionals
             )
             for p in range(self.parameters.P):
-                probabilities_tangent_mapping[p] += np.einsum('jki,j->ki', probabilities_tensor, xi_jacobian[:, p])
+                probabilities_tangent_mapping[p] += np.squeeze(
+                    np.moveaxis(probabilities_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
+                )
                 if conditionals_tensor is not None:
-                    conditionals_tangent_mapping[p] += np.einsum('jki,j->ki', conditionals_tensor, xi_jacobian[:, p])
+                    conditionals_tangent_mapping[p] += np.squeeze(
+                        np.moveaxis(conditionals_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
+                    )
 
             # compute the supply-side Jacobian
             eta, capital_delta_inverse, eta_errors = self.compute_eta(