From 8e518048454091414b6985aef8f549c8721e5645 Mon Sep 17 00:00:00 2001 From: Daniel Weindl Date: Fri, 8 Apr 2022 22:54:39 +0200 Subject: [PATCH] sparse --- python/amici/ode_export.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/amici/ode_export.py b/python/amici/ode_export.py index 92c6eae0fd..f2456916a2 100644 --- a/python/amici/ode_export.py +++ b/python/amici/ode_export.py @@ -3405,9 +3405,18 @@ def _parallel_applyfunc( # parallel from pickle import PicklingError from multiprocessing import Pool + from sympy.matrices.dense import DenseMatrix with Pool(n_procs) as p: try: - return obj._new(obj.rows, obj.cols, p.map(func, obj)) + if isinstance(obj, DenseMatrix): + return obj._new(obj.rows, obj.cols, p.map(func, obj)) + elif isinstance(obj, sp.SparseMatrix): + dok = obj.todok() + mapped = p.map(func, dok.values()) + dok = {k: v for k, v in zip(dok.keys(), mapped) if v != 0} + return obj._new(obj.rows, obj.cols, dok) + else: + raise ValueError(f"Unsupported matrix type {type(obj)}") except PicklingError as e: raise ValueError(f"Couldn't pickle {func}. This is likely due " "that the argument was not a module-level "