Model import: parallelize computation of derivatives (#1740)

Allows to compute derivatives in parallel (using `multiprocessing`). Disabled by default. Enable by setting environment variable `AMICI_IMPORT_NPROCS` to the number of processes to use. For smaller models, the multiprocessing overhead dominates, therefore, this is only advisable to use for larger models (import time of ⪆ 5min).

Closes #1739 

Tested with https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab/tree/master/Benchmark-Models/Chen_MSB2009 on 8 core CPU, model import without compilation:

| n_procs | walltime m:ss |
|---------|---------------|
| 1       | 4:25.61       |
| 2       | 4:14.30       |
| 4       | 3:39.20       |
| 8       | 3:37.07       |

Tested with https://github.com/ICB-DCM/CS_Signalling_ERBB_RAS_AKT/tree/master/FroehlichKes2018/PEtab:

        n_procs=1 - walltime (m:ss): 25:05.64
        n_procs=2 - walltime (m:ss): 23:25.90
        n_procs=8 - walltime (m:ss): 17:29.66

.github/workflows/test_performance.yml

@@ -4,7 +4,7 @@ on:
     branches:
       - develop
       - master
-      - compile_without_optimization
+      - feature_1739_par_jac
 
   pull_request:
     branches:
@@ -63,7 +63,7 @@ jobs:
     # import test model
     - name: Import test model
       run: |
-        check_time.sh petab_import python tests/performance/test.py import
+        AMICI_IMPORT_NPROCS=2 check_time.sh petab_import python tests/performance/test.py import
 
     - name: "Upload artifact: CS_Signalling_ERBB_RAS_AKT_petab"
       uses: actions/upload-artifact@v1

python/amici/ode_export.py

@@ -418,15 +418,27 @@ def smart_jacobian(eq: sp.MutableDenseMatrix,
     :return:
         jacobian of eq wrt sym_var
     """
-    if min(eq.shape) and min(sym_var.shape) \
-            and not smart_is_zero_matrix(eq) \
-            and not smart_is_zero_matrix(sym_var):
+    if (
+        not min(eq.shape)
+        or not min(sym_var.shape)
+        or smart_is_zero_matrix(eq)
+        or smart_is_zero_matrix(sym_var)
+    ):
+        return sp.zeros(eq.shape[0], sym_var.shape[0])
+
+    if (n_procs := int(os.environ.get("AMICI_IMPORT_NPROCS", 1))) == 1:
+        # serial
         return sp.Matrix([
-            eq[i, :].jacobian(sym_var) if eq[i, :].has(*sym_var.flat())
-            else [0] * sym_var.shape[0]
+            _jacobian_row(eq[i, :], sym_var)
             for i in range(eq.shape[0])
         ])
-    return sp.zeros(eq.shape[0], sym_var.shape[0])
+
+    # parallel
+    from multiprocessing import Pool
+    with Pool(n_procs) as p:
+        mapped = p.starmap(_jacobian_row,
+                           ((eq[i, :], sym_var) for i in range(eq.shape[0])))
+    return sp.Matrix(mapped)
 
 
 @log_execution_time('running smart_multiply', logger)
@@ -3322,3 +3334,10 @@ def _custom_pow_eval_derivative(self, s):
         (self.base, sp.And(sp.Eq(self.base, 0), sp.Eq(dbase, 0))),
         (part2, True)
     )
+
+
+def _jacobian_row(eq_i, sym_var):
+    """Compute a row of a jacobian"""
+    if eq_i.has(*sym_var.flat()):
+        return eq_i.jacobian(sym_var)
+    return [0] * sym_var.shape[0]