Fix (some) ASA event handling (#2797)

* Fix backward simulation: In the forward simulation, measurements are handled after events. So for the backward integration, events should be handled after data points.
* Make Heaviside updates more visible
* Store Heaviside values and additional other pre-event state during event handling for the backward simulation
* Remove the now obsolete `Model::updateHeavisideB`

Author: dweindl

include/amici/forwardproblem.h

@@ -25,28 +25,40 @@ struct Discontinuity {
      * @brief Constructor.
      * @param time
      * @param root_info
+     * @param x_pre
      * @param xdot_pre
      * @param x_post
      * @param xdot_post
+     * @param h_pre
+     * @param total_cl_pre
      */
     explicit Discontinuity(
         realtype time, std::vector<int> const& root_info = std::vector<int>(),
+        AmiVector const& x_pre = AmiVector(),
         AmiVector const& xdot_pre = AmiVector(),
         AmiVector const& x_post = AmiVector(),
-        AmiVector const& xdot_post = AmiVector()
+        AmiVector const& xdot_post = AmiVector(),
+        std::vector<realtype> const& h_pre = std::vector<realtype>(),
+        std::vector<realtype> const& total_cl_pre = std::vector<realtype>(0)
     )
         : time(time)
         , x_post(x_post)
+        , x_pre(x_pre)
         , xdot_post(xdot_post)
         , xdot_pre(xdot_pre)
-        , root_info(root_info) {}
+        , root_info(root_info)
+        , h_pre(h_pre)
+        , total_cl_pre(total_cl_pre) {}
 
     /** Time of discontinuity. */
     realtype time;
 
     /** Post-event state vector (dimension nx). */
     AmiVector x_post;
 
+    /** Pre-event state vector (dimension nx). */
+    AmiVector x_pre;
+
     /** Post-event differential state vectors (dimension nx). */
     AmiVector xdot_post;
 
@@ -61,6 +73,16 @@ struct Discontinuity {
      * if not. See CVodeGetRootInfo for details.
      */
     std::vector<int> root_info;
+
+    /**
+     * Flag indicating whether a certain Heaviside function should be active or
+     * not, pre-event value (dimension: `ne`)
+     */
+    std::vector<realtype> h_pre;
+
+    /** Total abundances for conservation laws
+     (dimension: `nx_rdata - nx_solver`) */
+    std::vector<realtype> total_cl_pre;
 };
 
 /**

include/amici/model.h

@@ -1359,15 +1359,6 @@ class Model : public AbstractModel, public ModelDimensions {
      */
     void updateHeaviside(std::vector<int> const& rootsfound);
 
-    /**
-     * @brief Updates the Heaviside variables `h` on event occurrences in the
-     * backward problem.
-     * @param rootsfound Provides the direction of the zero-crossing, so adding
-     * it will give the right update to the Heaviside variables (zero if no root
-     * was found)
-     */
-    void updateHeavisideB(int const* rootsfound);
-
     /**
      * @brief Check if the given array has only finite elements.
      *

src/backwardproblem.cpp

@@ -129,7 +129,11 @@ void EventHandlingBwdSimulator::handleEventB(
         ws_->nroots_[ie]--;
     }
 
-    model_->updateHeavisideB(disc.root_info.data());
+    // apply pre-event state
+    auto state = model_->getModelState();
+    state.h = disc.h_pre;
+    state.total_cl = disc.total_cl_pre;
+    model_->setModelState(state);
 }
 
 void EventHandlingBwdSimulator::handleDataPointB(
@@ -218,18 +222,18 @@ void EventHandlingBwdSimulator::run(
             );
         }
 
-        // handle discontinuity
-        if (!ws_->discs_.empty() && tnext == ws_->discs_.back().time) {
-            handleEventB(ws_->discs_.back(), dJzdx);
-            ws_->discs_.pop_back();
-        }
-
         // handle data-point
         if (it >= 0 && tnext == timepoints[it]) {
             handleDataPointB(it, dJydx);
             it--;
         }
 
+        // handle discontinuity
+        if (!ws_->discs_.empty() && tnext == ws_->discs_.back().time) {
+            handleEventB(ws_->discs_.back(), dJzdx);
+            ws_->discs_.pop_back();
+        }
+
         // reinitialize state
         solver_->reInitB(ws_->which, t_, ws_->xB_, ws_->dxB_);
         solver_->quadReInitB(ws_->which, ws_->xQB_);

src/forwardproblem.cpp

@@ -304,10 +304,6 @@ void EventHandlingSimulator::handle_event(
                   store_event(edata);
 
               store_pre_event_state(seflag, initial_event);
-
-              if (!initial_event) {
-                  model_->updateHeaviside(ws_->roots_found);
-              }
           };
 
     // store post-event information that is to be saved
@@ -325,6 +321,10 @@ void EventHandlingSimulator::handle_event(
 
     store_pre_event_info(false);
 
+    if (!initial_event) {
+        model_->updateHeaviside(ws_->roots_found);
+    }
+
     // Collect all triggered events waiting for execution
     for (int ie = 0; ie < model_->ne; ie++) {
         // only consider transitions false -> true
@@ -379,7 +379,12 @@ void EventHandlingSimulator::handle_event(
         // and add it to the list of pending events if necessary
         if (detect_secondary_events()) {
             store_post_event_info();
+
             store_pre_event_info(true);
+
+            if (!initial_event) {
+                model_->updateHeaviside(ws_->roots_found);
+            }
         }
     }
     store_post_event_info();
@@ -474,6 +479,9 @@ void EventHandlingSimulator::store_pre_event_state(
         }
     } else if (solver_->computingASA()) {
         result.discs.back().xdot_pre = ws_->xdot_old;
+        result.discs.back().x_pre = ws_->x_old;
+        result.discs.back().h_pre = model_->getModelState().h;
+        result.discs.back().total_cl_pre = model_->getModelState().total_cl;
     }
 }
 

src/model.cpp

@@ -1559,12 +1559,6 @@ void Model::updateHeaviside(std::vector<int> const& rootsfound) {
     }
 }
 
-void Model::updateHeavisideB(int const* rootsfound) {
-    for (int ie = 0; ie < ne; ie++) {
-        state_.h.at(ie) -= rootsfound[ie];
-    }
-}
-
 int Model::checkFinite(
     gsl::span<realtype const> array, ModelQuantity model_quantity, realtype t
 ) const {