Sparsify dwdp to reduce computation time for adjoints (#858)

Fix super long computation times for adjoint sensitivities introduced in AMICI 0.10.0 (Fixes #849)

.github/workflows/test-large-model.yml

@@ -4,7 +4,8 @@ on:
     branches:
       - develop
       - master
-      - fix_862_powsimp
+      - feature_sparse_quadratures
+
   pull_request:
     branches:
       - master
@@ -73,3 +74,9 @@ jobs:
     - name: adjoint_sensitivities
       run: |
         check_time.sh adjoint_sensitivities tests/performance/test.py adjoint_sensitivities
+    - name: forward_simulation_non_optimal_parameters
+      run: |
+        check_time.sh forward_simulation_non_optimal_parameters tests/performance/test.py forward_simulation_non_optimal_parameters
+    - name: adjoint_sensitivities_non_optimal_parameters
+      run: |
+        check_time.sh adjoint_sensitivities_non_optimal_parameters tests/performance/test.py adjoint_sensitivities_non_optimal_parameters

include/amici/abstract_model.h

@@ -108,7 +108,8 @@ class AbstractModel {
                         const AmiVector &dx) = 0;
 
     /**
-     * @brief Parameter derivative of residual function
+     * @brief Model specific sparse implementation of
+              explicit parameter derivative of right hand side
      * @param t time
      * @param x state
      * @param dx time derivative of state (DAE only)
@@ -668,6 +669,18 @@ class AbstractModel {
                        const realtype *p, const realtype *k, const realtype *h,
                        const realtype *w, const realtype *tcl,
                        const realtype *stcl);
+
+    /**
+     * @brief Model specific implementation for dwdp, column pointers
+     * @param indexptrs column pointers
+     **/
+    virtual void fdwdp_colptrs(sunindextype *indexptrs);
+
+    /**
+     * @brief Model specific implementation for dwdp, row values
+     * @param indexvals row values
+     **/
+    virtual void fdwdp_rowvals(sunindextype *indexvals);
 
     /**
      * @brief Model specific sensitivity implementation of dwdp

include/amici/model.h

@@ -71,14 +71,18 @@ class Model : public AbstractModel {
      * @param plist indexes wrt to which sensitivities are to be computed
      * @param idlist indexes indicating algebraic components (DAE only)
      * @param z2event mapping of event outputs to events
+     * @param pythonGenerated flag indicating matlab or python wrapping
+     * @param ndxdotdp_explicit number of nonzero elements dxdotdp_explicit
+     * @param ndxdotdp_implicit number of nonzero elements dxdotdp_implicit
      */
     Model(int nx_rdata, int nxtrue_rdata, int nx_solver, int nxtrue_solver,
           int ny, int nytrue, int nz, int nztrue, int ne, int nJ, int nw,
-          int ndwdx, int ndwdp, int ndxdotdw, std::vector<int> ndJydy, int nnz,
-          int ubw, int lbw, amici::SecondOrderMode o2mode,
+          int ndwdx, int ndwdp, int ndxdotdw, std::vector<int> ndJydy,
+          int nnz, int ubw, int lbw, amici::SecondOrderMode o2mode,
           const std::vector<amici::realtype> &p, std::vector<amici::realtype> k,
           const std::vector<int> &plist, std::vector<amici::realtype> idlist,
-          std::vector<int> z2event);
+          std::vector<int> z2event, bool pythonGenerated=false,
+          int ndxdotdp_explicit=0, int ndxdotdp_implicit=0);
 
     /** destructor */
     ~Model() override = default;
@@ -130,6 +134,8 @@ class Model : public AbstractModel {
     using AbstractModel::fdsigmaydp;
     using AbstractModel::fdsigmazdp;
     using AbstractModel::fdwdp;
+    using AbstractModel::fdwdp_colptrs;
+    using AbstractModel::fdwdp_rowvals;
     using AbstractModel::fdwdx;
     using AbstractModel::fdwdx_colptrs;
     using AbstractModel::fdwdx_rowvals;
@@ -1033,12 +1039,6 @@ class Model : public AbstractModel {
      */
     bool getAlwaysCheckFinite() const;
 
-    /**
-     * @brief check whether the model was generated from python
-     * @return that
-     */
-    virtual bool wasPythonGenerated() const { return false; }
-
     /**
      * @brief Initial states
      * @param x pointer to state variables
@@ -1130,7 +1130,6 @@ class Model : public AbstractModel {
 
     /** number of nonzero entries in dxdotdw */
     int ndxdotdw{0};
-
     /** number of nonzero entries in dJydy */
     std::vector<int> ndJydy;
 
@@ -1145,18 +1144,34 @@ class Model : public AbstractModel {
 
     /** lower bandwith of the jacobian */
     int lbw{0};
-
+
+    /** flag indicating Matlab or python based model generation */
+    bool pythonGenerated;
+
+    /** number of nonzero entries in ndxdotdp_explicit */
+    int ndxdotdp_explicit{0};
+
+    /** number of nonzero entries in ndxdotdp_implicit */
+    int ndxdotdp_implicit{0};
+
     /** flag indicating whether for sensi == AMICI_SENSI_ORDER_SECOND
      * directional or full second order derivative will be computed */
     SecondOrderMode o2mode{SecondOrderMode::none};
 
     /** flag array for DAE equations */
     std::vector<realtype> idlist;
 
-    /** temporary storage of dxdotdp data across functions (dimension: nplist x
-     * nx_solver, row-major) */
+    /** temporary storage of dxdotdp data across functions, Python only
+     (dimension: nplist x nx_solver, nnz: ndxdotdp_explicit, type CSC_MAT) */
+    mutable SUNMatrixWrapper dxdotdp_explicit;
+
+    /** temporary storage of dxdotdp_implicit data across functions, Python only
+     (dimension: nplist x * nx_solver, nnz: ndxdotdp_implicit, type CSC_MAT) */
+    mutable SUNMatrixWrapper dxdotdp_implicit;
+
+    /** temporary storage of dxdotdp data across functions, Matlab only
+     (dimension: nplist x nx_solver, row-major) */
     AmiVectorArray dxdotdp;
-
     /** AMICI context */
     AmiciApplication *app = &defaultContext;
 
@@ -1258,7 +1273,7 @@ class Model : public AbstractModel {
     virtual void fdJydy_colptrs(sunindextype *indexptrs, int index);
 
     /**
-     * @brief Model specific implementation of fdxdotdw row vals
+     * @brief Model specific implementation of fdJydy row vals
      * @param indexptrs row val pointers
      * @param index ytrue index
      */
@@ -1560,6 +1575,9 @@ class Model : public AbstractModel {
     /** Sparse dxdotdw temporary storage (dimension: ndxdotdw) */
     mutable SUNMatrixWrapper dxdotdw;
 
+    /** Sparse dwdp temporary storage (dimension: ndwdp) */
+    mutable SUNMatrixWrapper dwdp;
+
     /** Sparse dwdx temporary storage (dimension: ndwdx) */
     mutable SUNMatrixWrapper dwdx;
 
@@ -1573,12 +1591,12 @@ class Model : public AbstractModel {
     mutable std::vector<realtype> mz;
 
     /** Sparse observable derivative of data likelihood,
-     * only used if wasPythonGenerated()==true
+     * only used if pythonGenerated==true
      * (dimension nytrue, nJ x ny, row-major) */
     mutable std::vector<SUNMatrixWrapper> dJydy;
 
     /** observable derivative of data likelihood,
-     * only used if wasPythonGenerated()==false
+     * only used if pythonGenerated==false
      * (dimension nJ x ny x nytrue, row-major)
      */
     mutable std::vector<realtype> dJydy_matlab;
@@ -1661,11 +1679,6 @@ class Model : public AbstractModel {
     /** tempory storage of w data across functions (dimension: nw) */
     mutable std::vector<realtype> w;
 
-    /** tempory storage of sparse/dense dwdp data across functions
-     * (dimension: ndwdp)
-     */
-    mutable std::vector<realtype> dwdp;
-
     /** tempory storage for flattened sx,
      * (dimension: nx_solver x nplist, row-major)
      */

include/amici/model_dae.h

@@ -58,21 +58,24 @@ class Model_DAE : public Model {
      * @param plist indexes wrt to which sensitivities are to be computed
      * @param idlist indexes indicating algebraic components (DAE only)
      * @param z2event mapping of event outputs to events
+     * @param pythonGenerated flag indicating matlab or python wrapping
+     * @param ndxdotdp_explicit number of nonzero elements dxdotdp_explicit
+     * @param ndxdotdp_implicit number of nonzero elements dxdotdp_implicit
      */
     Model_DAE(const int nx_rdata, const int nxtrue_rdata, const int nx_solver,
               const int nxtrue_solver, const int ny, const int nytrue,
               const int nz, const int nztrue, const int ne, const int nJ,
-              const int nw, const int ndwdx, const int ndwdp,
-              const int ndxdotdw, std::vector<int> ndJydy, const int nnz,
-              const int ubw, const int lbw, const SecondOrderMode o2mode,
+              const int nw, const int ndwdx, const int ndwdp, const int ndxdotdw,
+              std::vector<int> ndJydy, const int nnz, const int ubw,
+              const int lbw, const SecondOrderMode o2mode,
               std::vector<realtype> const &p, std::vector<realtype> const &k,
-              std::vector<int> const &plist,
-              std::vector<realtype> const &idlist,
-              std::vector<int> const &z2event)
+              std::vector<int> const &plist, std::vector<realtype> const &idlist,
+              std::vector<int> const &z2event, const bool pythonGenerated=false,
+              const int ndxdotdp_explicit=0, const int ndxdotdp_implicit=0)
         : Model(nx_rdata, nxtrue_rdata, nx_solver, nxtrue_solver, ny, nytrue,
-                nz, nztrue, ne, nJ, nw, ndwdx, ndwdp, ndxdotdw,
-                std::move(ndJydy), nnz, ubw, lbw, o2mode, p, k, plist, idlist,
-                z2event) {}
+                nz, nztrue, ne, nJ, nw, ndwdx, ndwdp, ndxdotdw, std::move(ndJydy),
+                nnz, ubw, lbw, o2mode, p, k, plist, idlist, z2event,
+                pythonGenerated, ndxdotdp_explicit, ndxdotdp_implicit) {}
 
     void fJ(realtype t, realtype cj, const AmiVector &x, const AmiVector &dx,
             const AmiVector &xdot, SUNMatrix J) override;

include/amici/model_ode.h

@@ -58,21 +58,25 @@ class Model_ODE : public Model {
      * @param plist indexes wrt to which sensitivities are to be computed
      * @param idlist indexes indicating algebraic components (DAE only)
      * @param z2event mapping of event outputs to events
+     * @param pythonGenerated flag indicating matlab or python wrapping
+     * @param ndxdotdp_explicit number of nonzero elements dxdotdp_explicit
+     * @param ndxdotdp_implicit number of nonzero elements dxdotdp_implicit
      */
     Model_ODE(const int nx_rdata, const int nxtrue_rdata, const int nx_solver,
               const int nxtrue_solver, const int ny, const int nytrue,
               const int nz, const int nztrue, const int ne, const int nJ,
               const int nw, const int ndwdx, const int ndwdp,
-              const int ndxdotdw, std::vector<int> ndJydy, const int nnz,
-              const int ubw, const int lbw, const SecondOrderMode o2mode,
-              std::vector<realtype> const &p, std::vector<realtype> const &k,
-              std::vector<int> const &plist,
+              const int ndxdotdw, std::vector<int> ndJydy,
+              const int nnz, const int ubw, const int lbw,
+              const SecondOrderMode o2mode, std::vector<realtype> const &p,
+              std::vector<realtype> const &k, std::vector<int> const &plist,
               std::vector<realtype> const &idlist,
-              std::vector<int> const &z2event)
+              std::vector<int> const &z2event, const bool pythonGenerated=false,
+              const int ndxdotdp_explicit=0, const int ndxdotdp_implicit=0)
         : Model(nx_rdata, nxtrue_rdata, nx_solver, nxtrue_solver, ny, nytrue,
-                nz, nztrue, ne, nJ, nw, ndwdx, ndwdp, ndxdotdw,
-                std::move(ndJydy), nnz, ubw, lbw, o2mode, p, k, plist, idlist,
-                z2event) {}
+                nz, nztrue, ne, nJ, nw, ndwdx, ndwdp, ndxdotdw, std::move(ndJydy),
+                nnz, ubw, lbw, o2mode, p, k, plist, idlist, z2event,
+                pythonGenerated, ndxdotdp_explicit, ndxdotdp_implicit) {}
 
     void fJ(realtype t, realtype cj, const AmiVector &x, const AmiVector &dx,
             const AmiVector &xdot, SUNMatrix J) override;
@@ -218,7 +222,7 @@ class Model_ODE : public Model {
      */
     void fdxdotdw(realtype t, const N_Vector x);
 
-    /** Sensitivity of dx/dt wrt model parameters p
+    /** Explicit sensitivity of dx/dt wrt model parameters p
      * @param t timepoint
      * @param x Vector with the states
      * @return status flag indicating successful execution
@@ -402,7 +406,7 @@ class Model_ODE : public Model {
                        const realtype *p, const realtype *k, const realtype *h,
                        const realtype *w) = 0;
 
-    /** model specific implementation of fdxdotdp, with w chainrule
+    /** model specific implementation of fdxdotdp, with w chainrule (Matlab)
      * @param dxdotdp partial derivative xdot wrt p
      * @param t timepoint
      * @param x Vector with the states
@@ -418,19 +422,39 @@ class Model_ODE : public Model {
                           const realtype *h, int ip, const realtype *w,
                           const realtype *dwdp);
 
-    /** model specific implementation of fdxdotdp, without w chainrule
-     * @param dxdotdp partial derivative xdot wrt p
+    /** model specific implementation of fdxdotdp_explicit, no w chainrule (Py)
+     * @param dxdotdp_explicit partial derivative xdot wrt p
      * @param t timepoint
      * @param x Vector with the states
      * @param p parameter vector
      * @param k constants vector
      * @param h heavyside vector
-     * @param ip parameter index
      * @param w vector with helper variables
      */
-    virtual void fdxdotdp(realtype *dxdotdp, realtype t, const realtype *x,
-                          const realtype *p, const realtype *k,
-                          const realtype *h, int ip, const realtype *w);
+    virtual void fdxdotdp_explicit(realtype *dxdotdp_explicit, realtype t,
+                                   const realtype *x, const realtype *p,
+                                   const realtype *k, const realtype *h,
+                                   const realtype *w);
+
+    /** model specific implementation of fdxdotdp_explicit, colptrs part
+     * @param indexptrs column pointers
+     */
+    virtual void fdxdotdp_explicit_colptrs(sunindextype *indexptrs);
+
+    /** model specific implementation of fdxdotdp_explicit, rowvals part
+     * @param indexvals row values
+     */
+    virtual void fdxdotdp_explicit_rowvals(sunindextype *indexvals);
+
+    /** model specific implementation of fdxdotdp_implicit, colptrs part
+     * @param indexptrs column pointers
+     */
+    virtual void fdxdotdp_implicit_colptrs(sunindextype *indexptrs);
+
+    /** model specific implementation of fdxdotdp_implicit, rowvals part
+     * @param indexvals row values
+     */
+    virtual void fdxdotdp_implicit_rowvals(sunindextype *indexvals);
 
     /** model specific implementation of fdxdotdw, data part
      * @param dxdotdw partial derivative xdot wrt w
@@ -450,12 +474,11 @@ class Model_ODE : public Model {
      */
     virtual void fdxdotdw_colptrs(sunindextype *indexptrs);
 
-    /** model specific implementation of fdxdotdw, colptrs part
+    /** model specific implementation of fdxdotdw, rowvals part
      * @param indexvals row values
      */
     virtual void fdxdotdw_rowvals(sunindextype *indexvals);
 };
-
 } // namespace amici
 
 #endif // MODEL_H

include/amici/serialization.h

@@ -111,6 +111,9 @@ void serialize(Archive &ar, amici::Model &u, const unsigned int version) {
     ar &u.pscale;
     ar &u.tstart;
     ar &u.stateIsNonNegative;
+    ar &u.pythonGenerated;
+    ar &u.ndxdotdp_explicit;
+    ar &u.ndxdotdp_implicit;
 }