From c095b0059ab239f56706271f915628989ddf5bad Mon Sep 17 00:00:00 2001 From: Daniel Weindl Date: Fri, 25 Mar 2022 16:02:54 +0100 Subject: [PATCH] Model import: parallelize computation of derivatives (#1740) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 ++-- python/amici/ode_export.py | 31 +++++++++++++++++++++----- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test_performance.yml b/.github/workflows/test_performance.yml index 3ee2aad159..1d012ce958 100644 --- a/.github/workflows/test_performance.yml +++ b/.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 diff --git a/python/amici/ode_export.py b/python/amici/ode_export.py index 8c9fa7e96a..dd96678e30 100644 --- a/python/amici/ode_export.py +++ b/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]