Initial ERK discrete adjoint support in ARKODE

doc/arkode/guide/source/Constants.rst

@@ -626,9 +626,10 @@ contains the ARKODE output constants.
    +-------------------------------------+------+------------------------------------------------------------+
    | :index:`ARK_STEP_DIRECTION_ERR`     | -52  | An error occurred changing the step direction.             |
    +-------------------------------------+------+------------------------------------------------------------+
-   | :index:`ARK_UNRECOGNIZED_ERROR`     | -99  | An unknown error was encountered.                          |
+   | :index:`ARK_ADJ_RECOMPUTE_FAIL`     | -53  | An occurred recomputing steps during the adjoint           |
+   |                                     |      | integration.                                               |
    +-------------------------------------+------+------------------------------------------------------------+
-   |                                                                                                         |
+   | :index:`ARK_UNRECOGNIZED_ERROR`     | -99  | An unknown error was encountered.                          |
    +-------------------------------------+------+------------------------------------------------------------+
    | **ARKLS linear solver module output constants**                                                         |
    +-------------------------------------+------+------------------------------------------------------------+

doc/arkode/guide/source/Mathematics.rst

@@ -978,7 +978,7 @@ arise from the **separable** Hamiltonian system
 where
 
 .. math::
-   f_1(t, q) \equiv -\frac{\partial V(t, q)}{\partial q}, \qquad
+   f_1(t, q) \equiv \frac{\partial V(t, q)}{\partial q}, \qquad
    f_2(t, p) \equiv \frac{\partial T(t, p)}{\partial p}.
 
 When *H* is autonomous, then *H* is a conserved quantity. Often this corresponds
@@ -2697,3 +2697,111 @@ by :math:`\eta_\text{rf}`.
 
 For more information on utilizing relaxation Runge--Kutta methods, see
 :numref:`ARKODE.Usage.Relaxation`.
+
+
+.. _ARKODE.Mathematics.ASA:
+
+Adjoint Sensitivity Analysis
+============================
+
+Consider :eq:`ARKODE_IVP_simple_explicit`, but where the ODE also depends on some parameters
+:math:`p` (that is, we have :math:`f(t,y,p)`). Now, suppose we have a functional :math:`g(y(t_f),p)`
+for which we would like to compute the gradients :math:`\partial g(y(t_f),p)/\partial y(t_0)`
+and/or :math:`\partial g(y(t_f),p)/\partial p`.  This most often arises in the form of an
+optimization problem such as
+
+.. math::
+   \min_{\xi} \bar{\Psi}(\xi) = g(y(t_f), p)
+   :label: ARKODE_OPTIMIZATION_PROBLEM
+
+where :math:`\xi \subset \{y(t_0), p\}`. The adjoint method is one approach to obtaining the
+gradients that is particularly efficient when there are relatively few functionals and a large
+number of parameters. While :ref:`CVODES <CVODES.Mathematics.ASA>` and
+:ref:`IDAS <IDAS.Mathematics.ASA>` *continuous* adjoint methods
+(differentiate-then-discretize), ARKODE provides *discrete* adjoint methods
+(discretize-then-differentiate). For the continuous approach, we derive and solve the adjoint ODE
+backwards in time
+
+.. math::
+   \lambda'(t) &= -f_y^T(t, y, p) \lambda,\quad \lambda(t_F) = g_y^T(y(t_f), p), \\
+   \mu'(t) &= -f_p^T(t, y, p) \mu,\quad \mu(t_F) = g_p^T(y(t_f), p), \quad t_f \geq t \geq t_0, \\
+   :label: ARKODE_CONTINUOUS_ADJOINT_ODE
+
+where :math:`\lambda(t) \in \mathbb{R}^N`, :math:`\mu(t) \in \mathbb{R}^{N_s}`
+:math:`f_y \equiv \partial f/\partial y \in \mathbb{R}^{N \times N}` is the Jacobian with respect to the dependent variable,
+and :math:`f_p \equiv \partial f/\partial p \in \mathbb{R}^{N \times N_s}` is the Jacobian with respect to the parameters
+(:math:`N` is the size of the original ODE, :math:`N_s` is the number of parameters).
+When solved with a numerical time integration scheme, the solution to the continuous adjoint ODE
+are numerical approximations of the continuous adjoint sensitivities
+
+.. math::
+   \lambda(t_0) \approx  g_y^T(y(t_0), p),\quad \mu(t_0) \approx g_p^T(y(t_0), p)
+   :label: ARKODE_CONTINUOUS_ADJOINT_SOLUTION
+
+For the discrete adjoint approach, we first numerically discretize the original ODE :eq:`ARKODE_IVP_simple_explicit`.
+In the context of ARKODE, this is done with a one-step time integration scheme :math:`\varphi` so that
+
+.. math::
+   y_0 = y(t_0),\quad y_n = \varphi(y_{n-1}).
+   :label: ARKODE_DISCRETE_ODE
+
+Reformulating the optimization problem for the discrete case, we have
+
+.. math::
+   \min_{\xi} \Psi(\xi) = g(y_n, p)
+   :label: ARKODE_DISCRETE_OPTIMIZATION_PROBLEM
+
+The gradients of :eq:`ARKODE_DISCRETE_OPTIMIZATION_PROBLEM` can be computed using the transposed chain
+rule backwards in time to obtain the discete adjoint variables :math:`\lambda_n, \lambda_{n-1}, \cdots, \lambda_0`
+and :math:`\mu_n, \mu_{n-1}, \cdots, \mu_0`,
+
+.. math::
+   \lambda_n &= g_y^T(y_n, p), \quad \lambda_k = \left(\frac{\partial \varphi}{\partial y_k}(y_k, p)\right)^T \lambda_{k+1} \\
+   \mu_n     &= g_p^T(y_n, p), \quad \mu_k     = \left(\frac{\partial \varphi}{\partial p}(y_k, p)\right)^T \lambda_{k+1},
+    \quad k = n - 1, \cdots, 0.
+   :label: ARKODE_DISCRETE_ADJOINT
+
+The solution of the discrete adjoint equations :eq:`ARKODE_DISCRETE_ADJOINT` is the sensitivities of the discrete cost function
+:eq:`ARKODE_DISCRETE_OPTIMIZATION_PROBLEM` with respect to changes in the discretized ODE :eq:`ARKODE_DISCRETE_ODE`.
+
+.. math::
+   \lambda_0 = g_y^T(y_0, p), \quad \mu_0 = g_p^T(y_0, p).
+   :label: ARKODE_DISCRETE_ADJOINT_SOLUTION
+
+Given an s-stage explicit Runge--Kutta method (as in :eq:`ARKODE_ERK`, but without the embedding), the discrete adjoint
+to compute :math:`\lambda_n` and :math:`\mu_n` starting from :math:`\lambda_{n+1}` and
+:math:`\mu_{n+1}` is given by
+
+.. math::
+   \Lambda_i &= h_n f_y^T(t_{n,i}, z_i) \left(b_i \lambda_{n+1} + \sum_{j=i+1}^s a_{j,i}
+   \Lambda_j \right), \quad \quad i = s, \dots, 1,\\
+   \nu_i     &= h_n f_p^T(t_{n,i}, z_i, p) \left(b_i \lambda_{n+1} + \sum_{j=i}^{s} a_{ji} \Lambda_j \right), \\
+   \lambda_n &= \lambda_{n+1} + \sum_{j=1}^{s} \Lambda_j, \\
+   \mu_n     &= \mu_{n+1} + \sum_{j=1}^{s} \nu_j.
+   :label: ARKODE_ERK_ADJOINT
+
+For more information on performing discrete adjoint sensitivity analysis using ARKODE see,
+:numref:`ARKODE.Usage.ARKStep.ASA`.
+
+For a detailed derivation of the discrete adjoint methods see :cite:p:`hager2000runge,sanduDiscrete2006`.
+For a detailed derivation of the continuous adjoint method see :ref:`CVODES <CVODES.Mathematics.ASA>`,
+or :cite:p:`CLPS:03`.
+
+
+Discrete vs. Continuous Adjoint Method
+--------------------------------------
+
+It is understood that the continuous adjoint method can be problematic in the context of
+optimization problems because the continuous adjoint method provides an approximation to the
+gradient of a continuous cost function while the optimizer is expecting the gradient of the discrete
+cost function. The discrepancy means that the optimizer can fail to converge further once it is near
+a local minimum :cite:p:`giles2000introduction`. On the other hand, the discrete adjoint method
+provides the exact gradient of the discrete cost function allowing the optimizer to fully converge.
+Consequently, the discrete adjoint method is often preferable in optimization despite its own
+drawbacks -- such as its (relatively) increased memory usage and the possible introduction of
+unphysical computational modes :cite:p:`sirkes1997finite`. This is not to say that the discrete
+adjoint approach is always the better choice over the continuous adjoint approach in optimization.
+Computational efficiency and stability of one approach over the other can be both problem and method
+dependent. Section 8 in the paper :cite:p:`rackauckas2020universal` discusses the tradeoffs further
+and provides numerous references that may help inform users in choosing between the discrete and
+continuous adjoint approaches.

doc/arkode/guide/source/Usage/ARKStep/ASA.rst

@@ -0,0 +1,58 @@
+.. _ARKODE.Usage.ARKStep.ASA:
+
+Adjoint Sensitivity Analysis
+============================
+
+The previous sections discuss using ARKStep for the integration of forward ODE models.
+This section discusses how to use ARKStep for adjoint sensitivity analysis as introduced
+in :numref:`ARKODE.Mathematics.ASA`. To use ARKStep for adjoint sensitivity analysis (ASA), users simply setup the forward
+integration as usual (following :numref:`ARKODE.Usage.Skeleton`) with one exception:
+a :c:type:`SUNAdjointCheckpointScheme` object must be created and passed to
+:c:func:`ARKodeSetAdjointCheckpointScheme` before the call to the :c:func:`ARKodeEvolve`
+function. After the forward model integration code, a :c:type:`SUNAdjointStepper` object
+can be created for the adjoint model integration by calling :c:func:`ARKStepCreateAdjointStepper`.
+The code snippet below demonstrates these steps in brief and the example code
+``examples/arkode/C_serial/ark_lotka_volterra_asa.c`` demonstrates these steps in detail.
+
+.. code-block:: C
+
+   // 1. Create a SUNAdjointCheckpointScheme object
+
+   // 2. Setup ARKStep for forward integration
+
+   // 3. Attach the SUNAdjointCheckpointScheme
+
+   // 4. Evolve the forward model
+
+   // 5. Create the SUNAdjointStepper
+
+   // 6. Setup the adjoint model
+
+   // 7. Evolve the adjoint model
+
+   // 8. Cleanup
+
+
+
+User Callable Functions
+-----------------------
+
+This section describes ARKStep-specific user-callable functions for performing
+adjoint sensitivity analysis with methods with ARKStep.
+
+.. c:function:: int ARKStepCreateAdjointStepper(void* arkode_mem, N_Vector sf, SUNAdjointStepper* adj_stepper_ptr)
+
+   Creates a :c:type:`SUNAdjointStepper` object compatible with the provided ARKStep instance for
+   integrating the adjoint sensitivity system :eq:`ARKODE_DISCRETE_ADJOINT`.
+
+   :param arkode_mem: a pointer to the ARKStep memory block.
+   :param sf: the sensitivity vector holding the adjoint system terminal condition.
+      This must be an instance of the ManyVector ``N_Vector`` implementation with at
+      least one subvector (depending on if sensitivities to parameters should be computed).
+      The first subvector must be :math:`\partial g_y(y(t_f)) \in \mathbb{R}^N`. If sensitivities to parameters should be computed, then the second subvector must be :math:`g_p(y(t_f), p) \in \mathbb{R}^{N_s}`.
+   :param adj_stepper_ptr: the newly created :c:type:`SUNAdjointStepper` object.
+
+   :return:
+      * ``ARK_SUCCESS`` if successful
+      * ``ARK_MEM_FAIL`` if a memory allocation failed
+      * ``ARK_ILL_INPUT`` if an argument has an illegal value.

doc/arkode/guide/source/Usage/ARKStep/index.rst

@@ -30,3 +30,4 @@ are specific to ARKStep.
    User_callable
    Relaxation
    XBraid
+   ASA

doc/arkode/guide/source/Usage/ERKStep/ASA.rst

@@ -0,0 +1,58 @@
+.. _ARKODE.Usage.ERKStep.ASA:
+
+Adjoint Sensitivity Analysis
+============================
+
+The previous sections discuss using ARKStep for the integration of forward ODE models.
+This section discusses how to use ARKStep for adjoint sensitivity analysis as introduced
+in :numref:`ARKODE.Mathematics.ASA`. To use ARKStep for adjoint sensitivity analysis (ASA), users simply setup the forward
+integration as usual (following :numref:`ARKODE.Usage.Skeleton`) with one exception:
+a :c:type:`SUNAdjointCheckpointScheme` object must be created and passed to
+:c:func:`ARKodeSetAdjointCheckpointScheme` before the call to the :c:func:`ARKodeEvolve`
+function. After the forward model integration code, a :c:type:`SUNAdjointStepper` object
+can be created for the adjoint model integration by calling :c:func:`ERKStepCreateAdjointStepper`.
+The code snippet below demonstrates these steps in brief and the example code
+``examples/arkode/C_serial/ark_lotka_volterra_asa.c`` demonstrates these steps in detail.
+
+.. code-block:: C
+
+   // 1. Create a SUNAdjointCheckpointScheme object
+
+   // 2. Setup ERKStep for forward integration
+
+   // 3. Attach the SUNAdjointCheckpointScheme
+
+   // 4. Evolve the forward model
+
+   // 5. Create the SUNAdjointStepper
+
+   // 6. Setup the adjoint model
+
+   // 7. Evolve the adjoint model
+
+   // 8. Cleanup
+
+
+
+User Callable Functions
+-----------------------
+
+This section describes ERKStep-specific user-callable functions for performing
+adjoint sensitivity analysis with methods with ERKStep.
+
+.. c:function:: int ERKStepCreateAdjointStepper(void* arkode_mem, N_Vector sf, SUNAdjointStepper* adj_stepper_ptr)
+
+   Creates a :c:type:`SUNAdjointStepper` object compatible with the provided ARKStep instance for
+   integrating the adjoint sensitivity system :eq:`ARKODE_DISCRETE_ADJOINT`.
+
+   :param arkode_mem: a pointer to the ARKStep memory block.
+   :param sf: the sensitivity vector holding the adjoint system terminal condition.
+      This must be an instance of the ManyVector ``N_Vector`` implementation with at
+      least one subvector (depending on if sensitivities to parameters should be computed).
+      The first subvector must be :math:`\partial g_y(y(t_f)) \in \mathbb{R}^N`. If sensitivities to parameters should be computed, then the second subvector must be :math:`g_p(y(t_f), p) \in \mathbb{R}^{N_s}`.
+   :param adj_stepper_ptr: the newly created :c:type:`SUNAdjointStepper` object.
+
+   :return:
+      * ``ARK_SUCCESS`` if successful
+      * ``ARK_MEM_FAIL`` if a memory allocation failed
+      * ``ARK_ILL_INPUT`` if an argument has an illegal value.

doc/arkode/guide/source/Usage/ERKStep/index.rst

@@ -29,3 +29,4 @@ are specific to ERKStep.
 
    User_callable
    Relaxation
+   ASA