Fix Rule-handling in PEtab import (#1753)

RateRules-targets in the condition table were previously ignored.

Closes #1750

Also fixes a bug where InitialAssignments were not honored for elements that occurred as column in the condition table and were set to `NaN` (i.e. "use model value").

.github/workflows/test_petab_test_suite.yml

@@ -63,9 +63,10 @@ jobs:
       # retrieve test models
       - name: Download and install PEtab test suite
         run: |
-          git clone --depth 1  --branch develop https://github.com/PEtab-dev/petab_test_suite \
+          git clone --depth 1 --branch develop \
+            https://github.com/PEtab-dev/petab_test_suite \
             && source ./build/venv/bin/activate \
-            && cd petab_test_suite && pip3 install -e . && cd ..
+            && cd petab_test_suite && pip3 install -e .
 
       - name: Run PEtab-related unit tests
         run: |

python/amici/parameter_mapping.py

@@ -191,14 +191,18 @@ def fill_in_parameters_for_condition(
     # Parameter mapping may contain parameter_ids as values, these *must*
     # be replaced
 
-    def _get_par(model_par, value):
+    def _get_par(model_par, value, mapping):
         """Replace parameter IDs in mapping dicts by values from
         problem_parameters where necessary"""
         if isinstance(value, str):
-            # estimated parameter
-            # (condition table overrides must have been handled already,
-            # e.g. by the PEtab parameter mapping)
-            return problem_parameters[value]
+            try:
+                # estimated parameter
+                return problem_parameters[value]
+            except KeyError:
+                # condition table overrides must have been handled already,
+                # e.g. by the PEtab parameter mapping, but parameters from
+                # InitialAssignments may still be present.
+                return _get_par(value, mapping[value], mapping)
         if model_par in problem_parameters:
             # user-provided
             return problem_parameters[model_par]
@@ -208,11 +212,11 @@ def _get_par(model_par, value):
         # constant value
         return value
 
-    map_preeq_fix = {key: _get_par(key, val)
+    map_preeq_fix = {key: _get_par(key, val, map_preeq_fix)
                      for key, val in map_preeq_fix.items()}
-    map_sim_fix = {key: _get_par(key, val)
+    map_sim_fix = {key: _get_par(key, val, map_sim_fix)
                    for key, val in map_sim_fix.items()}
-    map_sim_var = {key: _get_par(key, val)
+    map_sim_var = {key: _get_par(key, val, map_sim_var)
                    for key, val in map_sim_var.items()}
 
     # If necessary, (un)scale parameters

python/amici/petab_import.py

@@ -86,7 +86,8 @@ def get_fixed_parameters(
         # remove overridden parameters (`object`-type columns)
         fixed_parameters = [p for p in fixed_parameters
                             if condition_df[p].dtype != 'O'
-                            and sbml_model.getParameter(p) is not None]
+                            and sbml_model.getParameter(p) is not None
+                            and sbml_model.getRuleByVariable(p) is None]
         # must be unique
         if len(fixed_parameters) != len(set(fixed_parameters)):
             raise AssertionError(
@@ -512,12 +513,11 @@ def import_model_sbml(
     #  create a new parameter initial_${startOrCompartmentID}.
     #  feels dirty and should be changed (see also #924)
     # <BeginWorkAround>
+
     initial_states = [col for col in condition_df
-                      if sbml_model.getSpecies(col) is not None]
-    initial_sizes = [col for col in condition_df
-                     if sbml_model.getCompartment(col) is not None]
+                      if element_is_state(sbml_model, col)]
     fixed_parameters = []
-    if len(initial_states) or len(initial_sizes):
+    if initial_states:
         # add preequilibration indicator variable
         # NOTE: would only be required if we actually have preequilibration
         #  adding it anyways. can be optimized-out later
@@ -534,8 +534,8 @@ def import_model_sbml(
         logger.debug("Adding preequilibration indicator "
                      f"constant {PREEQ_INDICATOR_ID}")
     logger.debug("Adding initial assignments for "
-                 f"{initial_sizes + initial_states}")
-    for assignee_id in chain(initial_sizes, initial_states):
+                 f"{initial_states}")
+    for assignee_id in initial_states:
         init_par_id_preeq = f"initial_{assignee_id}_preeq"
         init_par_id_sim = f"initial_{assignee_id}_sim"
         for init_par_id in [init_par_id_preeq, init_par_id_sim]:
@@ -689,6 +689,20 @@ def show_model_info(sbml_model: 'libsbml.Model'):
     logger.info(f'Reactions: {len(sbml_model.getListOfReactions())}')
 
 
+def element_is_state(sbml_model: libsbml.Model, sbml_id: str) -> bool:
+    """Does the element with ID `sbml_id` correspond to a state variable?
+    """
+    if sbml_model.getCompartment(sbml_id) is not None:
+        return True
+    if sbml_model.getSpecies(sbml_id) is not None:
+        return True
+    if (rule := sbml_model.getRuleByVariable(sbml_id)) is not None \
+            and rule.getTypeCode() == libsbml.SBML_RATE_RULE:
+        return True
+
+    return False
+
+
 def _parse_cli_args():
     """
     Parse command line arguments

python/amici/petab_objective.py

@@ -22,7 +22,7 @@
 
 from . import AmiciModel, AmiciExpData
 from .logging import get_logger, log_execution_time
-from .petab_import import PREEQ_INDICATOR_ID
+from .petab_import import PREEQ_INDICATOR_ID, element_is_state
 from .parameter_mapping import (
     fill_in_parameters, ParameterMappingForCondition, ParameterMapping)
 
@@ -337,11 +337,11 @@ def create_parameter_mapping_for_condition(
     # ExpData.x0, but in the case of preequilibration this would not allow for
     # resetting initial states.
 
-    species_in_condition_table = [
+    states_in_condition_table = [
         col for col in petab_problem.condition_df
-        if petab_problem.sbml_model.getSpecies(col) is not None]
-
-    if species_in_condition_table:
+        if element_is_state(petab_problem.sbml_model, col)
+    ]
+    if states_in_condition_table:
         # set indicator fixed parameter for preeq
         # (we expect here, that this parameter was added during import and
         # that it was not added by the user with a different meaning...)
@@ -352,14 +352,33 @@ def create_parameter_mapping_for_condition(
         condition_map_sim[PREEQ_INDICATOR_ID] = 0.0
         condition_scale_map_sim[PREEQ_INDICATOR_ID] = LIN
 
-        def _set_initial_concentration(condition_id, species_id, init_par_id,
-                                       par_map, scale_map):
+        def _set_initial_state(condition_id, element_id, init_par_id,
+                               par_map, scale_map):
             value = petab.to_float_if_float(
-                petab_problem.condition_df.loc[condition_id, species_id])
+                petab_problem.condition_df.loc[condition_id, element_id])
             if pd.isna(value):
-                value = get_species_initial(
-                    petab_problem.sbml_model.getSpecies(species_id)
-                )
+                element = petab_problem.sbml_model.getElementBySId(element_id)
+                type_code = element.getTypeCode()
+                initial_assignment = petab_problem.sbml_model\
+                    .getInitialAssignmentBySymbol(element_id)
+                if initial_assignment:
+                    initial_assignment = sp.sympify(
+                        libsbml.formulaToL3String(initial_assignment.getMath())
+                    )
+                if type_code == libsbml.SBML_SPECIES:
+                    value = get_species_initial(element) \
+                        if initial_assignment is None else initial_assignment
+                elif type_code == libsbml.SBML_PARAMETER:
+                    value = element.getValue()\
+                        if initial_assignment is None else initial_assignment
+                elif type_code == libsbml.SBML_COMPARTMENT:
+                    value = element.getSize()\
+                        if initial_assignment is None else initial_assignment
+                else:
+                    raise NotImplementedError(
+                        f"Don't know what how to handle {element_id} in "
+                        "condition table.")
+
                 try:
                     value = float(value)
                 except (ValueError, TypeError):
@@ -368,11 +387,11 @@ def _set_initial_concentration(condition_id, species_id, init_par_id,
                         value = sp.nsimplify(value)
                     else:
                         raise NotImplementedError(
-                            "Cannot handle non-trivial expressions for "
-                            f"species initial for {species_id}: {value}")
+                            "Cannot handle non-trivial initial state "
+                            f"expression for {element_id}: {value}")
                     # this should be a parameter ID
                     value = str(value)
-                logger.debug(f'The species {species_id} has no initial value '
+                logger.debug(f'The species {element_id} has no initial value '
                              f'defined for the condition {condition_id} in '
                              'the PEtab conditions table. The initial value is '
                              f'now set to {value}, which is the initial value '
@@ -383,16 +402,17 @@ def _set_initial_concentration(condition_id, species_id, init_par_id,
                 scale_map[init_par_id] = petab.LIN
             else:
                 # parametric initial state
-                scale_map[init_par_id] = petab_problem.parameter_df.loc[
-                    value, PARAMETER_SCALE]
+                scale_map[init_par_id] = \
+                    petab_problem.parameter_df[PARAMETER_SCALE]\
+                        .get(value, petab.LIN)
 
-        for species_id in species_in_condition_table:
+        for element_id in states_in_condition_table:
             # for preequilibration
-            init_par_id = f'initial_{species_id}_preeq'
+            init_par_id = f'initial_{element_id}_preeq'
             if condition.get(PREEQUILIBRATION_CONDITION_ID):
                 condition_id = condition[PREEQUILIBRATION_CONDITION_ID]
-                _set_initial_concentration(
-                    condition_id, species_id, init_par_id, condition_map_preeq,
+                _set_initial_state(
+                    condition_id, element_id, init_par_id, condition_map_preeq,
                     condition_scale_map_preeq)
             else:
                 # need to set dummy value for preeq parameter anyways, as it
@@ -403,9 +423,9 @@ def _set_initial_concentration(condition_id, species_id, init_par_id,
 
             # for simulation
             condition_id = condition[SIMULATION_CONDITION_ID]
-            init_par_id = f'initial_{species_id}_sim'
-            _set_initial_concentration(
-                condition_id, species_id, init_par_id, condition_map_sim,
+            init_par_id = f'initial_{element_id}_sim'
+            _set_initial_state(
+                condition_id, element_id, init_par_id, condition_map_sim,
                 condition_scale_map_sim)
 
     ##########################################################################
@@ -547,22 +567,22 @@ def create_edata_for_condition(
         edata.id += "+" + condition.get(PREEQUILIBRATION_CONDITION_ID)
     ##########################################################################
     # enable initial parameters reinitialization
-    species_in_condition_table = [
+    states_in_condition_table = [
         col for col in petab_problem.condition_df
         if not pd.isna(petab_problem.condition_df.loc[
                            condition[SIMULATION_CONDITION_ID], col])
-           and petab_problem.sbml_model.getSpecies(col) is not None
+           and element_is_state(petab_problem.sbml_model, col)
     ]
     if condition.get(PREEQUILIBRATION_CONDITION_ID) \
-            and species_in_condition_table:
+            and states_in_condition_table:
         state_ids = amici_model.getStateIds()
         state_idx_reinitalization = [state_ids.index(s)
-                                     for s in species_in_condition_table]
+                                     for s in states_in_condition_table]
         edata.reinitialization_state_idxs_sim = state_idx_reinitalization
         logger.debug("Enabling state reinitialization for condition "
                      f"{condition.get(PREEQUILIBRATION_CONDITION_ID, '')} - "
                      f"{condition.get(SIMULATION_CONDITION_ID)} "
-                     f"{species_in_condition_table}")
+                     f"{states_in_condition_table}")
 
     ##########################################################################
     # timepoints

tests/petab_test_suite/test_petab_suite.py

@@ -1,15 +1,16 @@
 #!/usr/bin/env python3
-
 """Run PEtab test suite (https://github.com/PEtab-dev/petab_test_suite)"""
 
 import logging
 import sys
 
-import amici
+import pandas as pd
 import petab
 import petabtests
 import pytest
 from _pytest.outcomes import Skipped
+
+import amici
 from amici import SteadyStateSensitivityMode
 from amici.gradient_check import check_derivatives as amici_check_derivatives
 from amici.logging import get_logger, set_log_level
@@ -105,16 +106,26 @@ def _test_case(case, model_type):
     simulations_match = petabtests.evaluate_simulations(
         [simulation_df], gt_simulation_dfs, tol_simulations)
 
+    logger.log(logging.DEBUG if simulations_match else logging.ERROR,
+               f"Simulations: match = {simulations_match}")
+    if not simulations_match:
+        with pd.option_context('display.max_rows', None,
+                               'display.max_columns', None,
+                               'display.width', 200):
+            logger.log(logging.DEBUG, f"x_ss: {model.getStateIds()} "
+                                      f"{[rdata.x_ss for rdata in rdatas]}")
+            logger.log(logging.ERROR,
+                       f"Expected simulations:\n{gt_simulation_dfs}")
+            logger.log(logging.ERROR,
+                       f"Actual simulations:\n{simulation_df}")
     logger.log(logging.DEBUG if chi2s_match else logging.ERROR,
                f"CHI2: simulated: {chi2}, expected: {gt_chi2},"
                f" match = {chi2s_match}")
     logger.log(logging.DEBUG if simulations_match else logging.ERROR,
                f"LLH: simulated: {llh}, expected: {gt_llh}, "
                f"match = {llhs_match}")
-    logger.log(logging.DEBUG if simulations_match else logging.ERROR,
-               f"Simulations: match = {simulations_match}")
 
-    check_derivatives(problem, model)
+    check_derivatives(problem, model, solver)
 
     if not all([llhs_match, simulations_match]) or not chi2s_match:
         logger.error(f"Case {case} failed.")
@@ -124,19 +135,23 @@ def _test_case(case, model_type):
     logger.info(f"Case {case} passed.")
 
 
-def check_derivatives(problem: petab.Problem, model: amici.Model) -> None:
+def check_derivatives(
+        problem: petab.Problem,
+        model: amici.Model,
+        solver: amici.Solver
+) -> None:
     """Check derivatives using finite differences for all experimental
     conditions
 
     Arguments:
         problem: PEtab problem
         model: AMICI model matching ``problem``
+        solver: AMICI solver
     """
     problem_parameters = {t.Index: getattr(t, petab.NOMINAL_VALUE) for t in
                           problem.parameter_df.itertuples()}
-    solver = model.getSolver()
-    solver.setSensitivityMethod(amici.SensitivityMethod_forward)
-    solver.setSensitivityOrder(amici.SensitivityOrder_first)
+    solver.setSensitivityMethod(amici.SensitivityMethod.forward)
+    solver.setSensitivityOrder(amici.SensitivityOrder.first)
     # Required for case 9 to not fail in
     #  amici::NewtonSolver::computeNewtonSensis
     model.setSteadyStateSensitivityMode(