Steady state sensitivity modes (#1727)

plakrisenko · Fabian Fröhlich · dweindl · commit cf375b9c3f71 · 2022-04-05T15:30:56.000+02:00
Settings for specifying sensitivities computation method at steady state have been implemented/reworked.
Three `SteadyStateSensitivityMode` are available for both `SensitivityMethod`s (`forward` and `adjoint`):
- `newtonOnly`: only Newton's method is used for sensitivities computation
- `integrationOnly`: only integration is used for sensitivities computation, which in forward sensitivity analysis case means that numerical integration has to be used to find the steady state
- `integrateIfNewtonFails`: more flexible option, where simulation is used if Newton's method fails

Note: Previously available `SteadyStateSensitivityMode` `simulationFSA` has been removed.

Co-authored-by: Fabian Fröhlich &lt;fabian_froehlich@hms.harvard.edu&gt;
diff --git a/include/amici/defines.h b/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 */
diff --git a/python/examples/example_constant_species/ExampleEquilibrationLogic.ipynb b/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,
diff --git a/python/tests/test_compare_conservation_laws_sbml.py b/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()
diff --git a/python/tests/test_preequilibration.py b/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)
diff --git a/python/tests/test_pysb.py b/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)
 
diff --git a/python/tests/test_sbml_import.py b/python/tests/test_sbml_import.py
@@ -205,10 +205,9 @@ def test_presimulation(sbml_example_presimulation_module):
     """Test 'presimulation' test model"""
     model = sbml_example_presimulation_module.getModel()
     solver = model.getSolver()
-    solver.setNewtonMaxSteps(0)
     model.setTimepoints(np.linspace(0, 60, 61))
     model.setSteadyStateSensitivityMode(
-        amici.SteadyStateSensitivityMode.simulationFSA
+        amici.SteadyStateSensitivityMode.integrationOnly
     )
     solver.setSensitivityOrder(amici.SensitivityOrder.first)
     model.setReinitializeFixedParameterInitialStates(True)
diff --git a/src/steadystateproblem.cpp b/src/steadystateproblem.cpp
@@ -94,16 +94,22 @@ void SteadystateProblem::workSteadyStateBackwardProblem(
 void SteadystateProblem::findSteadyState(const Solver &solver, Model &model,
                                          int it) {
     steady_state_status_.resize(3, SteadyStateStatus::not_run);
+    bool turnOffNewton = model.getSteadyStateSensitivityMode() ==
+        SteadyStateSensitivityMode::integrationOnly && 
+        ((it == -1 && solver.getSensitivityMethodPreequilibration() ==
+         SensitivityMethod::forward) || solver.getSensitivityMethod() ==
+        SensitivityMethod::forward);
 
     /* First, try to run the Newton solver */
-    findSteadyStateByNewtonsMethod(model, false);
+    if (!turnOffNewton)
+        findSteadyStateByNewtonsMethod(model, false);
 
     /* Newton solver didn't work, so try to simulate to steady state */
     if (!checkSteadyStateSuccess())
         findSteadyStateBySimulation(solver, model, it);
 
     /* Simulation didn't work, retry the Newton solver from last sim state. */
-    if (!checkSteadyStateSuccess())
+    if (!turnOffNewton && !checkSteadyStateSuccess())
         findSteadyStateByNewtonsMethod(model, true);
 
     /* Nothing worked, throw an as informative error as possible */
@@ -243,11 +249,17 @@ void SteadystateProblem::computeSteadyStateQuadrature(const Solver &solver,
      We therefore compute the integral over xB first and then do a
      matrix-vector multiplication */
 
-    /* Try to compute the analytical solution for quadrature algebraically */
-    getQuadratureByLinSolve(model);
+    auto sensitivityMode = model.getSteadyStateSensitivityMode();
 
-    /* Analytical solution didn't work, perform simulation instead */
-    if (!hasQuadrature())
+    /* Try to compute the analytical solution for quadrature algebraically */
+    if (sensitivityMode == SteadyStateSensitivityMode::newtonOnly 
+        || sensitivityMode == SteadyStateSensitivityMode::integrateIfNewtonFails)
+        getQuadratureByLinSolve(model);
+
+    /* Perform simulation */
+    if (sensitivityMode == SteadyStateSensitivityMode::integrationOnly || 
+        (sensitivityMode == SteadyStateSensitivityMode::integrateIfNewtonFails
+         && !hasQuadrature()))
         getQuadratureBySimulation(solver, model);
 
     /* If analytic solution and integration did not work, throw an Exception */
@@ -366,8 +378,10 @@ bool SteadystateProblem::getSensitivityFlag(const Model &model,
     bool forwardSensisAlreadyComputed =
         solver.getSensitivityOrder() >= SensitivityOrder::first &&
         steady_state_status_[1] == SteadyStateStatus::success &&
-        model.getSteadyStateSensitivityMode() ==
-            SteadyStateSensitivityMode::simulationFSA;
+        (model.getSteadyStateSensitivityMode() ==
+         SteadyStateSensitivityMode::integrationOnly || 
+         model.getSteadyStateSensitivityMode() ==
+         SteadyStateSensitivityMode::integrateIfNewtonFails);
 
     bool simulationStartedInSteadystate =
         steady_state_status_[0] == SteadyStateStatus::success &&
@@ -392,9 +406,13 @@ bool SteadystateProblem::getSensitivityFlag(const Model &model,
         !simulationStartedInSteadystate;
 
     /* When we're creating a new solver object */
-    bool needForwardSensiAtCreation =
-        needForwardSensisPreeq && model.getSteadyStateSensitivityMode() ==
-                                      SteadyStateSensitivityMode::simulationFSA;
+    bool needForwardSensiAtCreation = 
+        needForwardSensisPreeq &&
+        (model.getSteadyStateSensitivityMode() ==
+         SteadyStateSensitivityMode::integrationOnly || 
+         model.getSteadyStateSensitivityMode() ==
+         SteadyStateSensitivityMode::integrateIfNewtonFails
+        );
 
     /* Check if we need to store sensis */
     switch (context) {
diff --git a/tests/petab_test_suite/test_petab_suite.py b/tests/petab_test_suite/test_petab_suite.py
@@ -16,7 +16,7 @@
 from amici.petab_import import import_petab_problem, PysbPetabProblem
 from amici.petab_objective import (
     simulate_petab, rdatas_to_measurement_df, create_parameterized_edatas)
-from amici import SteadyStateSensitivityMode_simulationFSA
+from amici import SteadyStateSensitivityMode
 
 logger = get_logger(__name__, logging.DEBUG)
 set_log_level(get_logger("amici.petab_import"), logging.DEBUG)
@@ -134,7 +134,7 @@ def check_derivatives(problem: petab.Problem, model: amici.Model) -> None:
     # Required for case 9 to not fail in
     #  amici::NewtonSolver::computeNewtonSensis
     model.setSteadyStateSensitivityMode(
-        SteadyStateSensitivityMode_simulationFSA)
+        SteadyStateSensitivityMode.integrateIfNewtonFails)
 
     for edata in create_parameterized_edatas(
             amici_model=model, petab_problem=problem,
