Extract common subexpressions

.github/workflows/test_benchmark_collection_models.yml

@@ -20,8 +20,12 @@ jobs:
     runs-on: ubuntu-22.04
 
     strategy:
+      fail-fast: false
       matrix:
-        python-version: [ 3.8 ]
+        python-version: [ "3.8" ]
+        extract_subexpressions: ["true", "false"]
+    env:
+      AMICI_EXTRACT_CSE: ${{ matrix.extract_subexpressions }}
 
     steps:
     - name: Set up Python ${{ matrix.python-version }}

documentation/python_installation.rst

@@ -318,6 +318,8 @@ dependencies, you can run
    applied to *all* installed packages, including dependencies.)
 
 
+.. _amici_python_install_env_vars:
+
 Custom installation
 -------------------
 

documentation/python_interface.rst

@@ -146,6 +146,35 @@ Examples
    ExampleExperimentalConditions.ipynb
    ExampleEquilibrationLogic.ipynb
 
+Environment variables affecting model import
+============================================
+
+In addition to the environment variables listed
+:ref:`here <amici_python_install_env_vars>`, the following environment
+variables control various behaviours during model import and compilation:
+
+.. list-table:: Environment variables affecting model import
+   :widths: 25 50 25
+   :header-rows: 1
+
+   * - Variable
+     - Purpose
+     - Example
+   * - ``AMICI_EXTRACT_CSE``
+     - Extract common subexpressions. May significantly reduce file size and
+       compile time for large models, but makes the generated code less
+       readable. Disabled by default.
+     - ``AMICI_EXTRACT_CSE=1``
+   * - ``AMICI_IMPORT_NPROCS``
+     - Number of processes to be used for model import. Defaults to 1.
+       Speeds up import of large models. Will slow down import of small models,
+       benchmarking recommended.
+     - ``AMICI_IMPORT_NPROCS=4``
+   * - ``AMICI_EXPERIMENTAL_SBML_NONCONST_CLS``
+     - Compute conservation laws for non-constant species. SBML-import only.
+       See :py:func:`amici.sbml_import.SbmlImporter.sbml2amici`.
+     -
+
 
 Miscellaneous
 =============

python/amici/cxxcodeprinter.py

@@ -1,9 +1,11 @@
 """C++ code generation"""
+import os
 import re
-from typing import List, Optional, Tuple, Dict
+from typing import Dict, List, Optional, Tuple
 
 import sympy as sp
 from sympy.printing.cxx import CXX11CodePrinter
+from sympy.utilities.iterables import numbered_symbols
 
 
 class AmiciCxxCodePrinter(CXX11CodePrinter):
@@ -12,6 +14,10 @@ class AmiciCxxCodePrinter(CXX11CodePrinter):
     def __init__(self):
         super().__init__()
 
+        # extract common subexpressions in matrix functions?
+        self.extract_cse = (os.getenv("AMICI_EXTRACT_CSE", "0").lower()
+                            in ('1', 'on', 'true'))
+
     def doprint(self, expr: sp.Expr, assign_to: Optional[str] = None) -> str:
         try:
             code = super().doprint(expr, assign_to)
@@ -95,13 +101,38 @@ def _get_sym_lines_symbols(
         :return:
             C++ code as list of lines
         """
-        return [
-            f'{" " * indent_level}{sym} = {self.doprint(math)};'
-            f'  // {variable}[{index}]'.replace('\n',
-                                                '\n' + ' ' * indent_level)
+        indent = " " * indent_level
+        lines = []
+
+        if self.extract_cse:
+            # Extract common subexpressions
+            symbol_generator = numbered_symbols(
+                cls=sp.Symbol, prefix="__amici_")
+            replacements, reduced_exprs = sp.cse(
+                equations,
+                symbols=symbol_generator,
+                # subexpressions must not include symbols computed in this
+                # function, as they will be computed only after the extracted
+                # subexpressions
+                ignore=symbols,
+            )
+            assert len(reduced_exprs) == 1, "Unexpected input"
+            reduced_exprs = reduced_exprs[0]
+
+            if replacements:
+                lines = [
+                    f'{indent}const realtype {sym} = {self.doprint(math)};'
+                    for (sym, math) in replacements
+                ]
+                equations = reduced_exprs
+
+        lines.extend([
+            f'{indent}{sym} = {self.doprint(math)};'
+            f'  // {variable}[{index}]'.replace('\n', '\n' + indent)
             for index, (sym, math) in enumerate(zip(symbols, equations))
             if math not in [0, 0.0]
-        ]
+        ])
+        return lines
 
     def csc_matrix(
             self,
@@ -139,7 +170,6 @@ def csc_matrix(
         :return:
             symbol_col_ptrs, symbol_row_vals, sparse_list, symbol_list,
             sparse_matrix
-
         """
         idx = 0
 
@@ -229,4 +259,3 @@ def get_switch_statement(condition: str, cases: Dict[int, List[str]],
         lines.append(indent0 + '}')
 
     return lines
-

python/amici/petab_objective.py

@@ -762,9 +762,9 @@ def rdatas_to_measurement_df(
     # iterate over conditions
     for (_, condition), rdata in zip(simulation_conditions.iterrows(), rdatas):
         # current simulation matrix
-        y = rdata['y']
+        y = rdata.y
         # time array used in rdata
-        t = list(rdata['t'])
+        t = list(rdata.ts)
 
         # extract rows for condition
         cur_measurement_df = petab.get_rows_for_condition(

tests/benchmark-models/test_petab_model.py

@@ -12,6 +12,8 @@
 
 import petab
 import yaml
+
+import amici
 from amici.logging import get_logger
 from amici.petab_objective import (simulate_petab, rdatas_to_measurement_df,
                                    LLH, RDATAS)
@@ -92,6 +94,10 @@ def main():
     rdatas = res[RDATAS]
     llh = res[LLH]
 
+    for rdata in rdatas:
+        assert rdata.status == amici.AMICI_SUCCESS, \
+            f"Simulation failed for {rdata.id}"
+
     # create simulation PEtab table
     sim_df = rdatas_to_measurement_df(rdatas=rdatas, model=amici_model,
                                       measurement_df=problem.measurement_df)