Exploit stoichiometric matrix in pysb import

include/amici/defines.h

@@ -180,7 +180,8 @@ enum class NonlinearSolverIteration {
 /** Sensitivity computation mode in steadyStateProblem */
 enum class SteadyStateSensitivityMode {
     newtonOnly,
-    simulationFSA
+    integrationOnly,
+    integrateIfNewtonFails
 };
 
 /** State in which the steady state computation finished */

python/amici/ode_export.py

@@ -1952,6 +1952,22 @@ def state_has_conservation_law(self, ix: int) -> bool:
         """
         return self._states[ix]._conservation_law is not None
 
+    def get_solver_indices(self) -> Dict[int, int]:
+        """
+        Returns a mapping that maps rdata species indices to solver indices
+
+        :return:
+            dictionary mapping rdata species indices to solver indices
+        """
+        solver_index = {}
+        ix_solver = 0
+        for ix in range(len(self._states)):
+            if self.state_has_conservation_law(ix):
+                continue
+            solver_index[ix] = ix_solver
+            ix_solver += 1
+        return solver_index
+
     def state_is_constant(self, ix: int) -> bool:
         """
         Checks whether the temporal derivative of the state is zero

python/amici/pysb_import.py

@@ -240,11 +240,110 @@ def ode_model_from_pysb_importer(
         for noise_distr in noise_distributions.values()
     )
 
+    _process_stoichiometric_matrix(model, ode, constant_parameters)
+
     ode.generate_basic_variables()
 
     return ode
 
 
+@log_execution_time('processing PySB stoich. matrix', logger)
+def _process_stoichiometric_matrix(pysb_model: pysb.Model,
+                                   ode_model: ODEModel,
+                                   constant_parameters: List[str]) -> None:
+
+    """
+    Exploits the PySB stoichiometric matrix to generate xdot derivatives
+
+    :param pysb_model:
+        pysb model instance
+
+    :param ode_model:
+        ODEModel instance
+
+    :param constant_parameters:
+        list of constant parameters
+    """
+
+    x = ode_model.sym('x')
+    w = list(ode_model.sym('w'))
+    p = list(ode_model.sym('p'))
+    x_rdata = list(ode_model.sym('x_rdata'))
+
+    n_x = len(x)
+    n_w = len(w)
+    n_p = len(p)
+    n_r = len(pysb_model.reactions)
+
+    solver_index = ode_model.get_solver_indices()
+    dflux_dx_dict = {}
+    dflux_dw_dict = {}
+    dflux_dp_dict = {}
+
+    w_idx = dict()
+    p_idx = dict()
+    wx_idx = dict()
+
+    def get_cached_index(symbol, sarray, index_cache):
+        idx = index_cache.get(symbol, None)
+        if idx is not None:
+            return idx
+        idx = sarray.index(symbol)
+        index_cache[symbol] = idx
+        return idx
+
+    for ir, rxn in enumerate(pysb_model.reactions):
+        for ix in np.unique(rxn['reactants']):
+            idx = solver_index.get(ix, None)
+            if idx is not None:
+                # species
+                values = dflux_dx_dict
+            else:
+                # conservation law
+                idx = get_cached_index(x_rdata[ix], w, wx_idx)
+                values = dflux_dw_dict
+
+            values[(ir, idx)] = sp.diff(rxn['rate'], x_rdata[ix])
+
+        # typically <= 3 free symbols in rate, we already account for
+        # species above so we only need to account for propensity, which
+        # can only be a parameter or expression
+        for fs in rxn['rate'].free_symbols:
+            # dw
+            if isinstance(fs, pysb.Expression):
+                var = w
+                idx_cache = w_idx
+                values = dflux_dw_dict
+            # dp
+            elif isinstance(fs, pysb.Parameter):
+                if fs.name in constant_parameters:
+                    continue
+                var = p
+                idx_cache = p_idx
+                values = dflux_dp_dict
+            else:
+                continue
+
+            idx = get_cached_index(fs, var, idx_cache)
+            values[(ir, idx)] = sp.diff(rxn['rate'], fs)
+
+    dflux_dx = sp.ImmutableSparseMatrix(n_r, n_x, dflux_dx_dict)
+    dflux_dw = sp.ImmutableSparseMatrix(n_r, n_w, dflux_dw_dict)
+    dflux_dp = sp.ImmutableSparseMatrix(n_r, n_p, dflux_dp_dict)
+
+    # use dok format to convert numeric csc to sparse symbolic
+    S = sp.ImmutableSparseMatrix(
+        n_x, n_r, # don't use shape here as we are eliminating rows
+        pysb_model.stoichiometry_matrix[
+            np.asarray(list(solver_index.keys())),:
+        ].todok()
+    )
+    # don't use `.dot` since it's awfully slow
+    ode_model._eqs['dxdotdx_explicit'] = S*dflux_dx
+    ode_model._eqs['dxdotdw'] = S*dflux_dw
+    ode_model._eqs['dxdotdp_explicit'] = S*dflux_dp
+
+
 @log_execution_time('processing PySB species', logger)
 def _process_pysb_species(pysb_model: pysb.Model,
                           ode_model: ODEModel) -> None:

python/examples/example_constant_species/ExampleEquilibrationLogic.ipynb

@@ -496,9 +496,9 @@
     }
    ],
    "source": [
-    "# Call simulation with singular Jacobian and simulationFSA mode\n",
+    "# Call simulation with singular Jacobian and integrateIfNewtonFails mode\n",
     "model.setTimepoints(np.full(1, np.inf))\n",
-    "model.setSteadyStateSensitivityMode(amici.SteadyStateSensitivityMode.simulationFSA)\n",
+    "model.setSteadyStateSensitivityMode(amici.SteadyStateSensitivityMode.integrateIfNewtonFails)\n",
     "solver = model.getSolver()\n",
     "solver.setNewtonMaxSteps(10)\n",
     "solver.setSensitivityMethod(amici.SensitivityMethod.forward)\n",
@@ -883,7 +883,7 @@
    ],
    "source": [
     "# Singluar Jacobian, use simulation\n",
-    "model.setSteadyStateSensitivityMode(amici.SteadyStateSensitivityMode.simulationFSA)\n",
+    "model.setSteadyStateSensitivityMode(amici.SteadyStateSensitivityMode.integrateIfNewtonFails)\n",
     "solver = model.getSolver()\n",
     "solver.setNewtonMaxSteps(10)\n",
     "solver.setSensitivityMethod(amici.SensitivityMethod.forward)\n",
@@ -1135,7 +1135,7 @@
    ],
    "source": [
     "# Non-singular Jacobian, use simulaiton\n",
-    "model_reduced.setSteadyStateSensitivityMode(amici.SteadyStateSensitivityMode.simulationFSA)\n",
+    "model_reduced.setSteadyStateSensitivityMode(amici.SteadyStateSensitivityMode.integrateIfNewtonFails)\n",
     "solver_reduced = model_reduced.getSolver()\n",
     "solver_reduced.setNewtonMaxSteps(0)\n",
     "solver_reduced.setSensitivityMethod(amici.SensitivityMethod.forward)\n",
@@ -1186,7 +1186,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.2"
+   "version": "3.8.10"
   },
   "toc": {
    "base_numbering": 1,

python/tests/test_compare_conservation_laws_sbml.py

@@ -10,21 +10,23 @@
 def edata_fixture():
     """edata is generated to test pre- and postequilibration"""
     edata_pre = amici.ExpData(2, 0, 0,
-                          np.array([0., 0.1, 0.2, 0.5, 1., 2., 5., 10.]))
+                              np.array([0., 0.1, 0.2, 0.5, 1., 2., 5., 10.]))
     edata_pre.setObservedData([1.5] * 16)
     edata_pre.fixedParameters = np.array([5., 20.])
     edata_pre.fixedParametersPreequilibration = np.array([0., 10.])
     edata_pre.reinitializeFixedParameterInitialStates = True
 
     # edata for postequilibration
     edata_post = amici.ExpData(2, 0, 0,
-                          np.array([float('inf')] * 3))
+                               np.array([float('inf')] * 3))
     edata_post.setObservedData([0.75] * 6)
     edata_post.fixedParameters = np.array([7.5, 30.])
 
     # edata with both equilibrations
     edata_full = amici.ExpData(2, 0, 0,
-        np.array([0., 0., 0., 1., 2., 2., 4., float('inf'), float('inf')]))
+                               np.array(
+                                   [0., 0., 0., 1., 2., 2., 4., float('inf'),
+                                    float('inf')]))
     edata_full.setObservedData([3.14] * 18)
     edata_full.fixedParameters = np.array([1., 2.])
     edata_full.fixedParametersPreequilibration = np.array([3., 4.])
@@ -42,7 +44,7 @@ def models():
     sbml_importer = amici.SbmlImporter(sbml_file)
 
     # Name of the model that will also be the name of the python module
-    model_name =  model_output_dir ='model_constant_species'
+    model_name = model_output_dir = 'model_constant_species'
     model_name_cl = model_output_dir_cl = 'model_constant_species_cl'
 
     # Define constants, observables, sigmas
@@ -67,9 +69,11 @@ def models():
                              compute_conservation_laws=False)
 
     # load both models
-    model_without_cl_module = amici.import_model_module(model_name,
+    model_without_cl_module = amici.import_model_module(
+        model_name,
         module_path=os.path.abspath(model_name))
-    model_with_cl_module = amici.import_model_module(model_name_cl,
+    model_with_cl_module = amici.import_model_module(
+        model_name_cl,
         module_path=os.path.abspath(model_name_cl))
 
     # get the models and return
@@ -81,8 +85,8 @@ def models():
 def get_results(model, edata=None, sensi_order=0,
                 sensi_meth=amici.SensitivityMethod.forward,
                 sensi_meth_preeq=amici.SensitivityMethod.forward,
+                stst_sensi_mode=amici.SteadyStateSensitivityMode.newtonOnly,
                 reinitialize_states=False):
-
     # set model and data properties
     model.setReinitializeFixedParameterInitialStates(reinitialize_states)
 
@@ -92,6 +96,7 @@ def get_results(model, edata=None, sensi_order=0,
     solver.setSensitivityOrder(sensi_order)
     solver.setSensitivityMethodPreequilibration(sensi_meth_preeq)
     solver.setSensitivityMethod(sensi_meth)
+    model.setSteadyStateSensitivityMode(stst_sensi_mode)
     if edata is None:
         model.setTimepoints(np.linspace(0, 5, 101))
     else:
@@ -134,9 +139,6 @@ def test_compare_conservation_laws_sbml(models, edata_fixture):
         assert np.isclose(rdata[field], rdata_cl[field]).all(), field
 
     # ----- compare simulations wo edata, sensi = 0, states and sensis -------
-    model_without_cl.setSteadyStateSensitivityMode(
-        amici.SteadyStateSensitivityMode.simulationFSA
-    )
 
     # run simulations
     edata, _, _ = edata_fixture
@@ -154,7 +156,10 @@ def test_compare_conservation_laws_sbml(models, edata_fixture):
     # run simulations
     rdata_cl = get_results(model_with_cl, edata=edata, sensi_order=1)
     assert rdata_cl['status'] == amici.AMICI_SUCCESS
-    rdata = get_results(model_without_cl, edata=edata, sensi_order=1)
+    rdata = get_results(
+        model_without_cl, edata=edata, sensi_order=1,
+        stst_sensi_mode=amici.SteadyStateSensitivityMode.integrateIfNewtonFails
+    )
     assert rdata['status'] == amici.AMICI_SUCCESS
     # check that steady state computation succeeded only by sim in full model
     assert (rdata['preeq_status'] == np.array([-3, 1, 0])).all()
@@ -178,43 +183,50 @@ def test_compare_conservation_laws_sbml(models, edata_fixture):
 
 def test_adjoint_pre_and_post_equilibration(edata_fixture):
     # get both models
-    model_module = amici.import_model_module('model_constant_species',
+    model_module = amici.import_model_module(
+        'model_constant_species',
         module_path=os.path.abspath('model_constant_species'))
     model = model_module.getModel()
-    model_module_cl = amici.import_model_module('model_constant_species_cl',
+    model_module_cl = amici.import_model_module(
+        'model_constant_species_cl',
         module_path=os.path.abspath('model_constant_species_cl'))
     model_cl = model_module_cl.getModel()
 
     # check gradient with and without state reinitialization
     for edata in edata_fixture:
         for reinit in [False, True]:
-    # --- compare different ways of preequilibration, full rank Jacobian ---------
+            # --- compare different ways of preequilibration, full rank Jacobian ------
             # forward preequilibration, forward simulation
-            rff_cl = get_results(model_cl, edata=edata, sensi_order=1,
-                                 sensi_meth=amici.SensitivityMethod.forward,
-                                 sensi_meth_preeq=amici.SensitivityMethod.forward,
-                                 reinitialize_states=reinit)
+            rff_cl = get_results(
+                model_cl, edata=edata, sensi_order=1,
+                sensi_meth=amici.SensitivityMethod.forward,
+                sensi_meth_preeq=amici.SensitivityMethod.forward,
+                reinitialize_states=reinit)
             # forward preequilibration, adjoint simulation
-            rfa_cl = get_results(model_cl, edata=edata, sensi_order=1,
-                                 sensi_meth=amici.SensitivityMethod.adjoint,
-                                 sensi_meth_preeq=amici.SensitivityMethod.forward,
-                                 reinitialize_states=reinit)
+            rfa_cl = get_results(
+                model_cl, edata=edata, sensi_order=1,
+                sensi_meth=amici.SensitivityMethod.adjoint,
+                sensi_meth_preeq=amici.SensitivityMethod.forward,
+                reinitialize_states=reinit)
             # adjoint preequilibration, adjoint simulation
-            raa_cl = get_results(model_cl, edata=edata, sensi_order=1,
-                                 sensi_meth=amici.SensitivityMethod.adjoint,
-                                 sensi_meth_preeq=amici.SensitivityMethod.adjoint,
-                                 reinitialize_states=reinit)
+            raa_cl = get_results(
+                model_cl, edata=edata, sensi_order=1,
+                sensi_meth=amici.SensitivityMethod.adjoint,
+                sensi_meth_preeq=amici.SensitivityMethod.adjoint,
+                reinitialize_states=reinit)
 
             # assert all are close
             assert np.isclose(rff_cl['sllh'], rfa_cl['sllh']).all()
             assert np.isclose(rfa_cl['sllh'], raa_cl['sllh']).all()
             assert np.isclose(raa_cl['sllh'], rff_cl['sllh']).all()
 
-    # --- compare fully adjoint approach to simulation with singular Jacobian ----
-            raa = get_results(model, edata=edata, sensi_order=1,
-                              sensi_meth=amici.SensitivityMethod.adjoint,
-                              sensi_meth_preeq=amici.SensitivityMethod.adjoint,
-                              reinitialize_states=reinit)
+            # --- compare fully adjoint approach to simulation with singular Jacobian ----
+            raa = get_results(
+                model, edata=edata, sensi_order=1,
+                sensi_meth=amici.SensitivityMethod.adjoint,
+                sensi_meth_preeq=amici.SensitivityMethod.adjoint,
+                stst_sensi_mode=amici.SteadyStateSensitivityMode.integrateIfNewtonFails,
+                reinitialize_states=reinit)
 
             # assert gradients are close (quadrature tolerances are laxer)
             assert np.isclose(raa_cl['sllh'], raa['sllh'], 1e-5, 1e-5).all()

python/tests/test_preequilibration.py

@@ -286,7 +286,8 @@ def test_equilibration_methods_with_adjoints(preeq_fixture):
 
     rdatas = {}
     equil_meths = [amici.SteadyStateSensitivityMode.newtonOnly,
-                   amici.SteadyStateSensitivityMode.simulationFSA]
+                   amici.SteadyStateSensitivityMode.integrationOnly,
+                   amici.SteadyStateSensitivityMode.integrateIfNewtonFails]
     sensi_meths = [amici.SensitivityMethod.forward,
                    amici.SensitivityMethod.adjoint]
     settings = itertools.product(equil_meths, sensi_meths)
@@ -333,7 +334,7 @@ def test_newton_solver_equilibration(preeq_fixture):
     edata.setObservedDataStdDev(np.hstack([stdy, stdy[0]]))
 
     rdatas = {}
-    settings = [amici.SteadyStateSensitivityMode.simulationFSA,
+    settings = [amici.SteadyStateSensitivityMode.integrationOnly,
                 amici.SteadyStateSensitivityMode.newtonOnly]
 
     solver.setNewtonStepSteadyStateCheck(True)
@@ -345,9 +346,6 @@ def test_newton_solver_equilibration(preeq_fixture):
         model.setSteadyStateSensitivityMode(equil_meth)
         if equil_meth == amici.SteadyStateSensitivityMode.newtonOnly:
             solver.setNewtonMaxSteps(10)
-        else:
-            solver.setSensiSteadyStateCheck(False)
-            solver.setNewtonMaxSteps(0)
 
         # add rdatas
         rdatas[equil_meth] = amici.runAmiciSimulation(model, solver, edata)

python/tests/test_pysb.py

@@ -201,7 +201,7 @@ def get_data(model):
     solver = model.getSolver()
     model.setTimepoints(np.linspace(0, 60, 61))
     model.setSteadyStateSensitivityMode(
-        amici.SteadyStateSensitivityMode.simulationFSA
+        amici.SteadyStateSensitivityMode.integrateIfNewtonFails
     )
 
     rdata = amici.runAmiciSimulation(model, solver)
@@ -220,7 +220,7 @@ def get_results(model, edata):
     edata.reinitializeFixedParameterInitialStates = True
     model.setTimepoints(np.linspace(0, 60, 61))
     model.setSteadyStateSensitivityMode(
-        amici.SteadyStateSensitivityMode.simulationFSA
+        amici.SteadyStateSensitivityMode.integrateIfNewtonFails
     )
     return amici.runAmiciSimulation(model, solver, edata)