Feature/sunadjoint

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

doc/arkode/guide/source/Constants.rst

@@ -540,6 +540,9 @@ contains the ARKODE output constants.
    +-------------------------------------+------+------------------------------------------------------------+
    | :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,112 @@ 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 = 0, \quad k = n - 1, \cdots 0, \\
+   \lambda_{\ell} &= \lambda_{\ell} + \left( \frac{\partial \varphi_k}{\partial y_{\ell}} (y_{k-1}) \right)^T \lambda_k,\quad \ell = k-1, \cdots, 0, \quad  k = n, \cdots 0. \\
+   \mu_n &= g_p^T(y_n, p), \quad \mu_k = 0, \quad k = n - 1, \cdots 0, \\
+   \mu_{\ell} &= \mu_{\ell} + \left( \frac{\partial \varphi_k}{\partial p} (y_{k-1}) \right)^T \mu_k,\quad \ell = k-1, \cdots, 0, \quad  k = n, \cdots 0.
+   :label: ARKODE_DISCRETE_ADJOINT
+
+The solution of the discrete adjoint equations :eq:`ARKODE_DISCRETE_ADJOINT` is the sensitivitrwes 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,57 @@
+.. _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 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 two subvectors.
+      The first subvector must hold :math:`\partial g/\partial y |_{t=t_f} \in \mathbb{R}^N` and the second subvector must hold  :math:`\partial g / \partial p |_{t=t_f} \in \mathbb{R}^d`.
+   :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/User_callable.rst

@@ -867,29 +867,30 @@ Optional inputs for ARKODE
 
 .. cssclass:: table-bordered
 
-================================================  =======================================  =======================
-Optional input                                    Function name                            Default
-================================================  =======================================  =======================
-Return ARKODE parameters to their defaults        :c:func:`ARKodeSetDefaults`              internal
-Set integrator method order                       :c:func:`ARKodeSetOrder`                 4
-Set dense output interpolation type (SPRKStep)    :c:func:`ARKodeSetInterpolantType`       ``ARK_INTERP_LAGRANGE``
-Set dense output interpolation type (others)      :c:func:`ARKodeSetInterpolantType`       ``ARK_INTERP_HERMITE``
-Set dense output polynomial degree                :c:func:`ARKodeSetInterpolantDegree`     5
-Disable time step adaptivity (fixed-step mode)    :c:func:`ARKodeSetFixedStep`             disabled
-Supply an initial step size to attempt            :c:func:`ARKodeSetInitStep`              estimated
-Maximum no. of warnings for :math:`t_n+h = t_n`   :c:func:`ARKodeSetMaxHnilWarns`          10
-Maximum no. of internal steps before *tout*       :c:func:`ARKodeSetMaxNumSteps`           500
-Maximum absolute step size                        :c:func:`ARKodeSetMaxStep`               :math:`\infty`
-Minimum absolute step size                        :c:func:`ARKodeSetMinStep`               0.0
-Set a value for :math:`t_{stop}`                  :c:func:`ARKodeSetStopTime`              undefined
-Interpolate at :math:`t_{stop}`                   :c:func:`ARKodeSetInterpolateStopTime`   ``SUNFALSE``
-Disable the stop time                             :c:func:`ARKodeClearStopTime`            N/A
-Supply a pointer for user data                    :c:func:`ARKodeSetUserData`              ``NULL``
-Maximum no. of ARKODE error test failures         :c:func:`ARKodeSetMaxErrTestFails`       7
-Set inequality constraints on solution            :c:func:`ARKodeSetConstraints`           ``NULL``
-Set max number of constraint failures             :c:func:`ARKodeSetMaxNumConstrFails`     10
-================================================  =======================================  =======================
-
+=================================================  ==========================================  =======================
+Optional input                                     Function name                               Default
+=================================================  ==========================================  =======================
+Return ARKODE parameters to their defaults         :c:func:`ARKodeSetDefaults`                 internal
+Set integrator method order                        :c:func:`ARKodeSetOrder`                    4
+Set dense output interpolation type (SPRKStep)     :c:func:`ARKodeSetInterpolantType`          ``ARK_INTERP_LAGRANGE``
+Set dense output interpolation type (others)       :c:func:`ARKodeSetInterpolantType`          ``ARK_INTERP_HERMITE``
+Set dense output polynomial degree                 :c:func:`ARKodeSetInterpolantDegree`        5
+Disable time step adaptivity (fixed-step mode)     :c:func:`ARKodeSetFixedStep`                disabled
+Supply an initial step size to attempt             :c:func:`ARKodeSetInitStep`                 estimated
+Maximum no. of warnings for :math:`t_n+h = t_n`    :c:func:`ARKodeSetMaxHnilWarns`             10
+Maximum no. of internal steps before *tout*        :c:func:`ARKodeSetMaxNumSteps`              500
+Maximum absolute step size                         :c:func:`ARKodeSetMaxStep`                  :math:`\infty`
+Minimum absolute step size                         :c:func:`ARKodeSetMinStep`                  0.0
+Set a value for :math:`t_{stop}`                   :c:func:`ARKodeSetStopTime`                 undefined
+Interpolate at :math:`t_{stop}`                    :c:func:`ARKodeSetInterpolateStopTime`      ``SUNFALSE``
+Disable the stop time                              :c:func:`ARKodeClearStopTime`               N/A
+Supply a pointer for user data                     :c:func:`ARKodeSetUserData`                 ``NULL``
+Maximum no. of ARKODE error test failures          :c:func:`ARKodeSetMaxErrTestFails`          7
+Set inequality constraints on solution             :c:func:`ARKodeSetConstraints`              ``NULL``
+Set max number of constraint failures              :c:func:`ARKodeSetMaxNumConstrFails`        10
+Set the checkpointing scheme to use (for adjoint)  :c:func:`ARKodeSetAdjointCheckpointScheme`  ``NULL``
+Set the checkpointing step index (for adjoint)     :c:func:`ARKodeSetAdjointCheckpointIndex`   0
+=================================================  ==========================================  =======================
 
 
 
@@ -1397,6 +1398,34 @@ Set max number of constraint failures             :c:func:`ARKodeSetMaxNumConstr
    .. versionadded:: 6.1.0
 
 
+.. c:function:: int ARKodeSetAdjointCheckpointScheme(void* arkode_mem, SUNAdjointCheckpointScheme checkpoint_scheme)
+
+   Specifies the :c:type:`SUNAdjointCheckpointScheme` to use for saving states
+   during the forward integration, and loading states during backward integration
+   of an adjoint system.
+
+   :param arkode_mem: pointer to the ARKODE memory block.
+   :param checkpoint_scheme: the checkpoint scheme to use.
+
+   :retval ARK_SUCCESS: the function exited successfully.
+   :retval ARK_MEM_NULL: ``arkode_mem`` was ``NULL``.
+
+   .. versionadded:: x.y.z
+
+.. c:function:: int ARKodeSetAdjointCheckpointIndex(void* arkode_mem, long int step_index)
+
+   Specifies the step index (that is step number) to insert the next checkpoint at.
+   This is incremented along with the step count, but it is useful to be able to reset
+   this index during recomputations of missing states during the backward adjoint integration.
+
+   :param arkode_mem: pointer to the ARKODE memory block.
+   :param checkpoint_scheme: the checkpoint scheme to use
+
+   :retval ARK_SUCCESS: the function exited successfully.
+   :retval ARK_MEM_NULL: ``arkode_mem`` was ``NULL``.
+
+   .. versionadded:: x.y.z
+
 
 .. _ARKODE.Usage.ARKodeAdaptivityInputTable:
 

doc/arkode/guide/source/index.rst

@@ -66,6 +66,7 @@ with support by the `US Department of Energy <http://www.doe.gov>`_,
    sunnonlinsol/index.rst
    sunadaptcontroller/index.rst
    sunstepper/index.rst
+   sunadjoint/index.rst
    sunmemory/index.rst
    sundials/Install_link.rst
    Constants

doc/arkode/guide/source/sunadjoint/SUNAdjoint_links.rst

@@ -0,0 +1,14 @@
+.. ----------------------------------------------------------------
+   SUNDIALS Copyright Start
+   Copyright (c) 2002-2024, Lawrence Livermore National Security
+   and Southern Methodist University.
+   All rights reserved.
+
+   See the top-level LICENSE and NOTICE files for details.
+
+   SPDX-License-Identifier: BSD-3-Clause
+   SUNDIALS Copyright End
+   ----------------------------------------------------------------
+
+.. include:: ../../../../shared/sunadjoint/SUNAdjointCheckpointScheme.rst
+.. include:: ../../../../shared/sunadjoint/SUNAdjointStepper.rst

doc/arkode/guide/source/sunadjoint/index.rst

@@ -0,0 +1,32 @@
+..
+   ----------------------------------------------------------------
+   SUNDIALS Copyright Start
+   Copyright (c) 2002-2024, Lawrence Livermore National Security
+   and Southern Methodist University.
+   All rights reserved.
+
+   See the top-level LICENSE and NOTICE files for details.
+
+   SPDX-License-Identifier: BSD-3-Clause
+   SUNDIALS Copyright End
+   ----------------------------------------------------------------
+
+.. _SUNAdjoint:
+
+############################
+Adjoint Sensitivity Analysis
+############################
+
+.. versionadded:: x.y.z
+
+The ``SUNAdjoint`` API consists of a few customizable modules that provide a framework for adjoint
+sensitivity analysis (ASA). The API itself does not implement ASA, but it provides a common
+interface for ASA capabilities implemented in the SUNDIALS packages. Right now it supports :ref:`the
+ASA capabilities in ARKODE <ARKODE.Mathematics.ASA>`, while the ASA capabilities in :ref:`CVODES
+<CVODES.Mathematics.ASA>` and :ref:`IDAS <IDAS.Mathematics.ASA>` must be used directly.
+*Users should read the package specific sections on ASA capabilities before this section.*
+
+.. toctree::
+   :maxdepth: 1
+
+   SUNAdjoint_links.rst

doc/arkode/guide/source/sunmemory/SUNMemory_links.rst

@@ -11,6 +11,7 @@
    ----------------------------------------------------------------
 
 .. include:: ../../../../shared/sunmemory/SUNMemory_Description.rst
+.. include:: ../../../../shared/sunmemory/SUNMemory_System.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_CUDA.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_HIP.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_SYCL.rst

doc/cvode/guide/source/sunmemory/SUNMemory_links.rst

@@ -11,6 +11,7 @@
    ----------------------------------------------------------------
 
 .. include:: ../../../../shared/sunmemory/SUNMemory_Description.rst
+.. include:: ../../../../shared/sunmemory/SUNMemory_System.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_CUDA.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_HIP.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_SYCL.rst

doc/cvodes/guide/source/sunmemory/SUNMemory_links.rst

@@ -11,6 +11,7 @@
    ----------------------------------------------------------------
 
 .. include:: ../../../../shared/sunmemory/SUNMemory_Description.rst
+.. include:: ../../../../shared/sunmemory/SUNMemory_System.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_CUDA.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_HIP.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_SYCL.rst

doc/ida/guide/source/sunmemory/SUNMemory_links.rst

@@ -11,6 +11,7 @@
    ----------------------------------------------------------------
 
 .. include:: ../../../../shared/sunmemory/SUNMemory_Description.rst
+.. include:: ../../../../shared/sunmemory/SUNMemory_System.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_CUDA.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_HIP.rst
 .. include:: ../../../../shared/sunmemory/SUNMemory_SYCL.rst