Initial discrete adjoint capability with ERK

.gitignore

@@ -19,6 +19,7 @@
 /.pydevproject
 .vscode
 compile_commands.json
+.clangd
 
 # custom test environment setup script
 /test/env/env.sh
@@ -59,6 +60,7 @@ compile_commands.json
 /doc/*/*.log
 
 # directories created by/for Sphinx documentation
+/doc/.venv
 /doc/.pyvenv
 /doc/*/guide/build/
 /doc/*/examples/build/

CHANGELOG.md

@@ -6,6 +6,12 @@
 
 ### New Features and Enhancements
 
+Added the `SUNStepper` base class to represent a generic solution procedure for
+IVPs. This is used by the SplittingStep and ForcingStep modules of ARKODE. A
+SUNStepper can be created from an ARKODE memory block with the new function
+`ARKodeCreateSUNStepper`. To enable interoperability with `MRIStepInnerStepper`,
+the function `MRIStepInnerStepper_CreateFromSUNStepper` was added.
+
 The following DIRK schemes now have coefficients accurate to quad precision:
 * `ARKODE_BILLINGTON_3_3_2`
 * `ARKODE_KVAERNO_4_2_3`

doc/arkode/guide/source/Constants.rst

@@ -538,6 +538,11 @@ contains the ARKODE output constants.
    | :index:`ARK_STEPPER_UNSUPPORTED`    | -48  | An operation was not supported by the current              |
    |                                     |      | time-stepping module.                                      |
    +-------------------------------------+------+------------------------------------------------------------+
+   | :index:`ARK_SUNSTEPPER_ERR`         | -49  | An error occurred in the SUNStepper module.                |
+   +-------------------------------------+------+------------------------------------------------------------+
+   | :index:`ARK_ADJ_RECOMPUTE_FAIL`     | -50  | An occurred recomputing steps during the adjoint           |
+   |                                     |      | integration.                                               |
+   +-------------------------------------+------+------------------------------------------------------------+
    | :index:`ARK_UNRECOGNIZED_ERROR`     | -99  | An unknown error was encountered.                          |
    +-------------------------------------+------+------------------------------------------------------------+
    |                                                                                                         |

doc/arkode/guide/source/Mathematics.rst

@@ -2126,3 +2126,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 integation 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^n(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(\varphi^k_y(y_k, p)\right)^T \lambda_{k+1} \\
+   \mu_n     &= g_p^T(y_n, p), \quad \mu_k     = \mu_{k+1} + \left(\varphi^k_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/User_callable.rst

@@ -35,6 +35,11 @@ ARKStep supports *all categories*:
 * non-identity mass matrices
 * relaxation Runge--Kutta methods
 
+ARKStep also has forcing function support when converted to a
+:c:type:`SUNStepper` or :c:type:`MRIStepInnerStepper`. See
+:c:func:`ARKodeCreateSUNStepper` and :c:func:`ARKStepCreateMRIStepInnerStepper`
+for additional details.
+
 
 .. _ARKODE.Usage.ARKStep.Initialization:
 

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/User_callable.rst

@@ -33,6 +33,10 @@ ERKStep supports the following categories:
 * temporal adaptivity
 * relaxation Runge--Kutta methods
 
+ERKStep does not have forcing function support when converted to a
+:c:type:`SUNStepper` or :c:type:`MRIStepInnerStepper`. See
+:c:func:`ARKodeCreateSUNStepper` and :c:func:`ARKStepCreateMRIStepInnerStepper`
+for additional details.
 
 
 .. _ARKODE.Usage.ERKStep.Initialization:

doc/arkode/guide/source/Usage/MRIStep/Custom_Inner_Stepper/Description.rst

@@ -79,6 +79,32 @@ Creating and Destroying an Object
       for details on how to attach member data and method function pointers.
 
 
+.. c:function:: int MRIStepInnerStepper_CreateFromSUNStepper(SUNStepper sunstepper, MRIStepInnerStepper* stepper)
+
+   This utility function wraps a :c:type:`SUNStepper` as an
+   :c:type:`MRIStepInnerStepper`.
+
+   :param sunctx: the SUNDIALS simulation context.
+   :param sunstepper: the c:type:`SUNStepper` to wrap.
+   :param stepper: a pointer to an MRI inner stepper object.
+
+   :retval ARK_SUCCESS: if successful
+   :retval ARK_MEM_FAIL: if a memory allocation error occurs
+
+   **Example usage:**
+
+   .. code-block:: C
+
+      SUNStepper sunstepper = NULL;
+      SUNStepper_Create(ctx, &sunstepper);
+      /* Attach content and functions to the SUNStepper... */
+
+      MRIStepInnerStepper inner_stepper = NULL;
+      flag = MRIStepInnerStepper_CreateFromSUNStepper(sunstepper, &inner_stepper);
+
+   .. versionadded:: x.y.z
+
+
 .. c:function:: int MRIStepInnerStepper_Free(MRIStepInnerStepper *stepper)
 
    This function destroys an :c:type:`MRIStepInnerStepper` object.

doc/arkode/guide/source/Usage/MRIStep/User_callable.rst

@@ -33,6 +33,10 @@ MRIStep supports the following categories:
 
 * implicit nonlinear and/or linear solvers
 
+MRIStep does not have forcing function support when converted to a
+:c:type:`SUNStepper` or :c:type:`MRIStepInnerStepper`. See
+:c:func:`ARKodeCreateSUNStepper` and :c:func:`ARKStepCreateMRIStepInnerStepper`
+for additional details.
 
 
 .. _ARKODE.Usage.MRIStep.Initialization:

doc/arkode/guide/source/Usage/SPRKStep/User_callable.rst

@@ -30,6 +30,11 @@ SPRKStep supports only the basic set of user-callable functions, and
 does not support any of the restricted groups (time adaptivity, implicit
 solvers, etc.).
 
+SPRKStep does not have forcing function support when converted to a
+:c:type:`SUNStepper` or :c:type:`MRIStepInnerStepper`. See
+:c:func:`ARKodeCreateSUNStepper` and :c:func:`ARKStepCreateMRIStepInnerStepper`
+for additional details.
+
 
 .. _ARKODE.Usage.SPRKStep.Initialization: