Faster matrix computation

EmmanuelMess · Apr 1, 2024 · 9ebc544 · 9ebc544
1 parent 09b4950
commit 9ebc544
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/code/simulator/SimulationFunctions.py b/code/simulator/SimulationFunctions.py
@@ -17,6 +17,15 @@ def precompiledForceCalculation(J: np.ndarray, l: np.float64) -> np.ndarray:
         """
         return (J.T @ l).reshape((-1, 2))
 
+    @staticmethod
+    @numba.njit
+    def precompiledMatricesComputation(ks: np.float64, kd: np.float64, dq: np.ndarray, Q: np.ndarray, C: np.ndarray,
+                                       dC: np.ndarray, W: np.ndarray, J: np.ndarray, dJ: np.ndarray)\
+            -> Tuple[np.ndarray, np.ndarray]:
+        f = dJ @ dq + J @ W @ Q + ks * C + kd * dC
+        g = J @ W @ J.T
+        return f, g
+
     @staticmethod
     def matrices(particles: IndexerIterator[Particle], constraints: IndexerIterator[Constraint],
                  weight: np.float64 = np.float64(1))\
@@ -61,8 +70,7 @@ def matrices(particles: IndexerIterator[Particle], constraints: IndexerIterator[
         J = J.reshape((m, n * d))
         dJ = dJ.reshape((m, n * d))
 
-        f = dJ @ dq + J @ W @ Q + ks * C + kd * dC
-        g = J @ W @ J.T
+        f, g = SimulationFunctions.precompiledMatricesComputation(ks, kd, dq, Q, C, dC, W, J, dJ)
 
         return f, g, J