Refactor storage of discontinuities (#2763)

Store together what belongs together.

Author: dweindl

include/amici/backwardproblem.h

@@ -2,6 +2,7 @@
 #define AMICI_BACKWARDPROBLEM_H
 
 #include "amici/defines.h"
+#include "amici/forwardproblem.h"
 #include "amici/vector.h"
 
 #include <vector>
@@ -80,8 +81,9 @@ class BackwardProblem {
     /**
      * @brief Execute everything necessary for the handling of events
      * for the backward problem
+     * @param disc The discontinuity to handle
      */
-    void handleEventB();
+    void handleEventB(Discontinuity const& disc);
 
     /**
      * @brief Execute everything necessary for the handling of data
@@ -114,22 +116,14 @@ class BackwardProblem {
     AmiVector dxB_;
     /** quadrature state vector */
     AmiVector xQB_;
-    /** array of state vectors at discontinuities*/
-    std::vector<AmiVector> x_disc_;
-    /** array of differential state vectors at discontinuities*/
-    std::vector<AmiVector> xdot_disc_;
-    /** array of old differential state vectors at discontinuities*/
-    std::vector<AmiVector> xdot_old_disc_;
     /** sensitivity state vector array */
     AmiVectorArray sx0_;
     /** array of number of found roots for a certain event type */
     std::vector<int> nroots_;
     /** array containing the time-points of discontinuities*/
-    std::vector<realtype> discs_;
+    std::vector<Discontinuity> discs_;
     /** index of the backward problem */
     int which = 0;
-    /** array of index which root has been found */
-    std::vector<std::vector<int>> root_idx_;
 
     /** state derivative of data likelihood */
     std::vector<realtype> dJydx_;

include/amici/forwardproblem.h

@@ -16,6 +16,52 @@ class Solver;
 class SteadystateProblem;
 class FinalStateStorer;
 
+/**
+ * @brief Data structure to store some state of a simulation at a discontinuity.
+ */
+struct Discontinuity {
+    /**
+     * @brief Constructor.
+     * @param time
+     * @param root_info
+     * @param xdot_pre
+     * @param x_post
+     * @param xdot_post
+     */
+    explicit Discontinuity(
+        realtype time, std::vector<int> const& root_info = std::vector<int>(),
+        AmiVector const& xdot_pre = AmiVector(),
+        AmiVector const& x_post = AmiVector(),
+        AmiVector const& xdot_post = AmiVector()
+    )
+        : time(time)
+        , x_post(x_post)
+        , xdot_post(xdot_post)
+        , xdot_pre(xdot_pre)
+        , root_info(root_info) {}
+
+    /** Time of discontinuity. */
+    realtype time;
+
+    /** Post-event state vector (dimension nx). */
+    AmiVector x_post;
+
+    /** Post-event differential state vectors (dimension nx). */
+    AmiVector xdot_post;
+
+    /** Pre-event differential state vectors (dimension nx). */
+    AmiVector xdot_pre;
+
+    /**
+     * @brief Array of flags indicating which root has been found.
+     *
+     * Array of length nr (ne) with the indices of the user functions gi found
+     * to have a root. For i = 0, . . . ,nr 1 or -1 if gi has a root, and = 0
+     * if not. See CVodeGetRootInfo for details.
+     */
+    std::vector<int> root_info;
+};
+
 /**
  * @brief The ForwardProblem class groups all functions for solving the
  * forward problem.
@@ -80,48 +126,18 @@ class ForwardProblem {
      */
     AmiVectorArray const& getStateSensitivity() const { return sx_; }
 
-    /**
-     * @brief Accessor for x_disc
-     * @return x_disc
-     */
-    std::vector<AmiVector> const& getStatesAtDiscontinuities() const {
-        return x_disc_;
-    }
-
-    /**
-     * @brief Accessor for xdot_disc
-     * @return xdot_disc
-     */
-    std::vector<AmiVector> const& getRHSAtDiscontinuities() const {
-        return xdot_disc_;
-    }
-
-    /**
-     * @brief Accessor for xdot_old_disc
-     * @return xdot_old_disc
-     */
-    std::vector<AmiVector> const& getRHSBeforeDiscontinuities() const {
-        return xdot_old_disc_;
-    }
-
     /**
      * @brief Accessor for nroots
      * @return nroots
      */
     std::vector<int> const& getNumberOfRoots() const { return nroots_; }
 
     /**
-     * @brief Accessor for discs
-     * @return discs
-     */
-    std::vector<realtype> const& getDiscontinuities() const { return discs_; }
-
-    /**
-     * @brief Accessor for rootidx
-     * @return rootidx
+     * @brief Get information on the discontinuities encountered so far.
+     * @return The vector of discontinuities.
      */
-    std::vector<std::vector<int>> const& getRootIndexes() const {
-        return root_idx_;
+    std::vector<Discontinuity> const& getDiscontinuities() const {
+        return discs_;
     }
 
     /**
@@ -294,7 +310,7 @@ class ForwardProblem {
      */
     void fillEvents(int nmaxevent) {
         if (checkEventsToFill(nmaxevent)) {
-            discs_.push_back(t_);
+            discs_.emplace_back(t_);
             storeEvent();
         }
     }
@@ -305,11 +321,6 @@ class ForwardProblem {
      */
     SimulationState getSimulationState();
 
-    /** array of index vectors (dimension ne) indicating whether the respective
-     * root has been detected for all so far encountered discontinuities,
-     * extended as needed (dimension: dynamic) */
-    std::vector<std::vector<int>> root_idx_;
-
     /** array of number of found roots for a certain event type
      * (dimension: ne) */
     std::vector<int> nroots_;
@@ -321,30 +332,8 @@ class ForwardProblem {
      * (dimension: ne) */
     std::vector<realtype> rval_tmp_;
 
-    /** array containing the time-points of discontinuities
-     * (dimension: dynamic) */
-    std::vector<realtype> discs_;
-
-    /**
-     * array of post-event state vectors (dimension nx) for all so far
-     * encountered discontinuities, (extended as needed;
-     * after event processing, same dimension as discs_)
-     */
-    std::vector<AmiVector> x_disc_;
-
-    /**
-     * array of post-event differential state vectors (dimension nx) for all so
-     * far encountered discontinuities, (extended as needed; after event
-     * processing, same dimension as discs_)
-     */
-    std::vector<AmiVector> xdot_disc_;
-
-    /**
-     * array of old (pre-event) differential state vectors (dimension nx) for
-     * all so far encountered discontinuities, (extended as needed; after event
-     * processing, same dimension as discs_)
-     */
-    std::vector<AmiVector> xdot_old_disc_;
+    /** Discontinuities encountered so far (dimension: dynamic) */
+    std::vector<Discontinuity> discs_;
 
     /** Events that are waiting to be handled at the current timepoint. */
     EventQueue pending_events_;

src/backwardproblem.cpp

@@ -20,13 +20,9 @@ BackwardProblem::BackwardProblem(
     , xB_(fwd.model->nx_solver, solver_->getSunContext())
     , dxB_(fwd.model->nx_solver, solver_->getSunContext())
     , xQB_(fwd.model->nJ * fwd.model->nplist(), solver_->getSunContext())
-    , x_disc_(fwd.getStatesAtDiscontinuities())
-    , xdot_disc_(fwd.getRHSAtDiscontinuities())
-    , xdot_old_disc_(fwd.getRHSBeforeDiscontinuities())
     , sx0_(fwd.getStateSensitivity())
     , nroots_(fwd.getNumberOfRoots())
     , discs_(fwd.getDiscontinuities())
-    , root_idx_(fwd.getRootIndexes())
     , dJydx_(fwd.getDJydx())
     , dJzdx_(fwd.getDJzdx()) {
     /* complement dJydx from postequilibration. This shouldn't overwrite
@@ -95,9 +91,9 @@ void BackwardProblem::workBackwardProblem() {
             }
 
             /* handle discontinuity */
-            if (!discs_.empty() && tnext == discs_.back()) {
+            if (!discs_.empty() && tnext == discs_.back().time) {
+                handleEventB(discs_.back());
                 discs_.pop_back();
-                handleEventB();
             }
 
             /* handle data-point */
@@ -128,30 +124,18 @@ void BackwardProblem::workBackwardProblem() {
     }
 }
 
-void BackwardProblem::handleEventB() {
-    auto rootidx = root_idx_.back();
-    this->root_idx_.pop_back();
-
-    auto x_disc = this->x_disc_.back();
-    this->x_disc_.pop_back();
-
-    auto xdot_disc = this->xdot_disc_.back();
-    this->xdot_disc_.pop_back();
-
-    auto xdot_old_disc = this->xdot_old_disc_.back();
-    this->xdot_old_disc_.pop_back();
-
+void BackwardProblem::handleEventB(Discontinuity const& disc) {
     for (int ie = 0; ie < model_->ne; ie++) {
 
-        if (rootidx[ie] == 0) {
+        if (disc.root_info[ie] == 0) {
             continue;
         }
 
         model_->addAdjointQuadratureEventUpdate(
-            xQB_, ie, t_, x_disc, xB_, xdot_disc, xdot_old_disc
+            xQB_, ie, t_, disc.x_post, xB_, disc.xdot_post, disc.xdot_pre
         );
         model_->addAdjointStateEventUpdate(
-            xB_, ie, t_, x_disc, xdot_disc, xdot_old_disc
+            xB_, ie, t_, disc.x_post, disc.xdot_post, disc.xdot_pre
         );
 
         if (model_->nz > 0) {
@@ -167,7 +151,7 @@ void BackwardProblem::handleEventB() {
         nroots_[ie]--;
     }
 
-    model_->updateHeavisideB(rootidx.data());
+    model_->updateHeavisideB(disc.root_info.data());
 }
 
 void BackwardProblem::handleDataPointB(int const it) {
@@ -196,8 +180,8 @@ realtype BackwardProblem::getTnext(int const it) {
     }
 
     if (!discs_.empty()
-        && (it < 0 || discs_.back() > model_->getTimepoint(it))) {
-        double tdisc = discs_.back();
+        && (it < 0 || discs_.back().time > model_->getTimepoint(it))) {
+        double tdisc = discs_.back().time;
         return tdisc;
     }
 

src/forwardproblem.cpp

@@ -242,8 +242,7 @@ void ForwardProblem::handleEvent(
 
         // store timepoint at which the event occurred, the root function
         // values, and the direction of any zero crossings of the root function
-        discs_.push_back(t_);
-        root_idx_.push_back(roots_found_);
+        discs_.emplace_back(t_, roots_found_);
         rval_tmp_ = rootvals_;
 
         if (model->nz > 0)
@@ -261,9 +260,9 @@ void ForwardProblem::handleEvent(
     //  that did not trigger a secondary event
     auto store_post_event_info = [this]() {
         if (solver->computingASA()) {
-            // store x to compute jump in discontinuity
-            x_disc_.push_back(x_);
-            xdot_disc_.push_back(xdot_);
+            // store updated x to compute jump in discontinuity
+            discs_.back().x_post = x_;
+            discs_.back().xdot_post = xdot_;
         }
     };
 
@@ -344,7 +343,7 @@ void ForwardProblem::storeEvent() {
         for (int ie = 0; ie < model->ne; ie++) {
             roots_found_.at(ie) = (nroots_.at(ie) < model->nMaxEvent()) ? 1 : 0;
         }
-        root_idx_.push_back(roots_found_);
+        discs_.back().root_info = roots_found_;
     }
 
     if (getRootCounter() < getEventCounter()) {
@@ -409,7 +408,7 @@ void ForwardProblem::store_pre_event_state(bool seflag, bool initial_event) {
             std::ranges::fill(stau_, 0.0);
         }
     } else if (solver->computingASA()) {
-        xdot_old_disc_.push_back(xdot_old_);
+        discs_.back().xdot_pre = xdot_old_;
     }
 }
 

src/rdata.cpp

@@ -291,13 +291,14 @@ void ReturnData::processForwardProblem(
 
     // process event data
     if (nz > 0) {
-        auto rootidx = fwd.getRootIndexes();
+        auto const& discontinuities = fwd.getDiscontinuities();
+        Expects(static_cast<int>(discontinuities.size()) == fwd.getEventCounter() + 1);
         for (int iroot = 0; iroot <= fwd.getEventCounter(); iroot++) {
             auto const simulation_state = fwd.getSimulationStateEvent(iroot);
             model.setModelState(simulation_state.state);
             getEventOutput(
-                simulation_state.t, rootidx.at(iroot), model, simulation_state,
-                edata
+                simulation_state.t, discontinuities.at(iroot).root_info, model,
+                simulation_state, edata
             );
         }
     }