doc: improve API docs (#229)

* doc: add some api docs for problem struct

* finish adding docstring for problem

* finish more docstrings

* doc: fill out documentation for builder

* doc: checkpointing, state

* doc: improve OdeSolverMethod

* doc: improve Vector

* doc: improve Matrix

* fix combine

* doc: improve Scalar

* remove some commented out code

* cargo fmt

* doc: diffsl, ops and jacobian

* cargo fmt

* doc: line_search

* doc: modules

* fix some bad links

* free up some disk space on ubuntu

* free disk space

Author: martinjrobins

.github/workflows/rust.yml

@@ -148,6 +148,16 @@ jobs:
       if: matrix.os == 'windows-latest'
       run: |
         echo "ADDITIONAL_FEATURES_FLAGS=--features diffsl-cranelift" >> $GITHUB_ENV
+    - name: Free Disk Space (Ubuntu)
+      if: matrix.os == 'ubuntu-latest'
+      uses: jlumbroso/free-disk-space@main
+      with:
+        android: true
+        dotnet: true
+        haskell: true
+        large-packages: false
+        docker-images: true
+        swap-storage: true
     - name: Run tests - default features
       if: matrix.tests == true
       run: cargo test --verbose

diffsol/src/jacobian/mod.rs

@@ -119,6 +119,31 @@ gen_find_non_zeros_linear!(
 
 use std::cell::RefCell;
 
+/// A structure for efficiently computing sparse Jacobians using graph coloring.
+///
+/// This struct uses graph coloring to group matrix columns that can be computed simultaneously
+/// using finite differences. By identifying columns that don't share any non-zero rows (i.e.,
+/// are structurally orthogonal), multiple Jacobian columns can be computed in parallel with a
+/// single function evaluation per color.
+///
+/// # Algorithm
+///
+/// The algorithm works as follows:
+/// 1. Construct a graph where each column is a node, and edges connect columns that share at least one non-zero row
+/// 2. Apply greedy graph coloring to assign colors to columns such that no two adjacent columns share the same color
+/// 3. For each color, perturb all columns of that color simultaneously and compute the corresponding Jacobian entries
+///
+/// To enable 3., we pre-calculate for all colors:
+///    - all the indices in the peturbed vector to peterb (i.e. set to 1) which is stored in `input_indices_per_color`.
+///    - all the indices in the output vector to read the results from which is stored in `src_indices_per_color`.
+///    - all the indices in the Jacobian matrix data vector to write the results to which is stored in `dst_indices_per_color`.
+///
+/// This reduces the number of function evaluations from O(n) to O(χ), where χ is the chromatic number
+/// of the graph (typically much smaller than n for sparse matrices).
+///
+/// # Type Parameters
+///
+/// * `M` - The matrix type used to store the Jacobian
 pub struct JacobianColoring<M: Matrix> {
     dst_indices_per_color: Vec<<M::V as Vector>::Index>,
     src_indices_per_color: Vec<<M::V as Vector>::Index>,
@@ -140,6 +165,21 @@ impl<M: Matrix> Clone for JacobianColoring<M> {
 }
 
 impl<M: Matrix> JacobianColoring<M> {
+    /// Create a new Jacobian coloring from a sparsity pattern and list of non-zero entries.
+    ///
+    /// This method constructs the graph coloring and organizes the indices for efficient
+    /// Jacobian computation. The coloring is computed using a greedy algorithm that aims
+    /// to minimize the number of colors (and thus function evaluations) needed.
+    ///
+    /// # Arguments
+    ///
+    /// * `sparsity` - The sparsity pattern of the Jacobian matrix
+    /// * `non_zeros` - A list of (row, column) pairs indicating non-zero entries in the Jacobian
+    /// * `ctx` - The context for creating vectors and indices
+    ///
+    /// # Returns
+    ///
+    /// A new `JacobianColoring` instance ready to compute Jacobians efficiently.
     pub fn new(sparsity: &impl MatrixSparsity<M>, non_zeros: &[(usize, usize)], ctx: M::C) -> Self {
         let ncols = sparsity.ncols();
         let graph = nonzeros2graph(non_zeros, ncols);
@@ -189,6 +229,22 @@ impl<M: Matrix> JacobianColoring<M> {
     //    Self::new_from_non_zeros(op, non_zeros)
     //}
 
+    /// Compute the Jacobian matrix in-place using the coloring scheme.
+    ///
+    /// This method uses the pre-computed coloring to efficiently compute the Jacobian by
+    /// evaluating `jac_mul` once per color instead of once per column. For each color group,
+    /// it perturbs multiple columns simultaneously and extracts the corresponding Jacobian entries.
+    ///
+    /// # Arguments
+    ///
+    /// * `op` - The nonlinear operator whose Jacobian is being computed
+    /// * `x` - The state vector at which to evaluate the Jacobian
+    /// * `t` - The time at which to evaluate the Jacobian
+    /// * `y` - The matrix to store the computed Jacobian (modified in-place)
+    ///
+    /// # Note
+    ///
+    /// The sparsity pattern of `y` must match the sparsity pattern used to create this coloring.
     pub fn jacobian_inplace<F: NonLinearOpJacobian<M = M, V = M::V, T = M::T, C = M::C>>(
         &self,
         op: &F,
@@ -209,6 +265,18 @@ impl<M: Matrix> JacobianColoring<M> {
         }
     }
 
+    /// Compute the sensitivity matrix (∂F/∂p) in-place using the coloring scheme.
+    ///
+    /// Similar to `jacobian_inplace`, but computes the sensitivity of the right-hand side
+    /// with respect to parameters rather than the state variables. This is used for forward
+    /// sensitivity analysis.
+    ///
+    /// # Arguments
+    ///
+    /// * `op` - The nonlinear operator whose sensitivity matrix is being computed
+    /// * `x` - The state vector at which to evaluate the sensitivities
+    /// * `t` - The time at which to evaluate the sensitivities
+    /// * `y` - The matrix to store the computed sensitivity matrix (modified in-place)
     pub fn sens_inplace<F: NonLinearOpSens<M = M, V = M::V, T = M::T, C = M::C>>(
         &self,
         op: &F,
@@ -229,6 +297,18 @@ impl<M: Matrix> JacobianColoring<M> {
         }
     }
 
+    /// Compute the transposed Jacobian (adjoint) matrix in-place using the coloring scheme.
+    ///
+    /// This method computes the transpose of the Jacobian, which is used in adjoint sensitivity
+    /// analysis. The coloring is applied to the columns of the transposed matrix (which correspond
+    /// to the rows of the original Jacobian).
+    ///
+    /// # Arguments
+    ///
+    /// * `op` - The nonlinear operator whose transposed Jacobian is being computed
+    /// * `x` - The state vector at which to evaluate the transposed Jacobian
+    /// * `t` - The time at which to evaluate the transposed Jacobian
+    /// * `y` - The matrix to store the computed transposed Jacobian (modified in-place)
     pub fn adjoint_inplace<F: NonLinearOpAdjoint<M = M, V = M::V, T = M::T, C = M::C>>(
         &self,
         op: &F,

diffsol/src/lib.rs

@@ -142,18 +142,126 @@ pub use diffsl::CraneliftJitModule;
 #[cfg(feature = "diffsl-llvm")]
 pub use diffsl::LlvmModule;
 
+/// Context objects for managing device resources for vectors and matrices (e.g. device streams, threading pools, etc.).
+///
+/// This module provides context types that encapsulate information about where data is stored and computed
+/// (CPU, GPU, etc.). Different backends like nalgebra and faer may require different context implementations.
+/// The [Context] trait defines the interface that must be implemented.
 pub mod context;
+
+/// Jacobian computation and coloring algorithms for efficient Jacobian evaluation.
+///
+/// This module provides utilities for:
+/// - Detecting the sparsity pattern of Jacobian matrices using NaN propagation
+/// - Computing efficient graph colorings of the Jacobian sparsity pattern
+/// - Using the coloring to compute sparse Jacobians with fewer function evaluations
+///
+/// The [JacobianColoring] struct is the main entry point for computing Jacobians efficiently.
 pub mod jacobian;
+
+/// Linear solver implementations and traits.
+///
+/// This module defines the [LinearSolver] trait for solving linear systems and provides implementations:
+/// - Direct solvers: [NalgebraLU], [FaerLU], [FaerSparseLU]
+/// - Optional sparse solvers: `KLU` (requires `suitesparse` feature)
+/// - GPU solvers: `CudaLU` (requires `cuda` feature)
+///
+/// The linear solver is a critical component used internally by nonlinear solvers to solve Newton systems.
 pub mod linear_solver;
+
+/// Matrix types and operations.
+///
+/// This module defines the [Matrix] trait and related traits for matrix operations:
+/// - [DenseMatrix] for dense column-major matrices
+/// - [MatrixView] and [MatrixViewMut] for borrowed views
+/// - Sparsity detection and handling
+///
+/// Implementations are provided for:
+/// - Dense matrices: [NalgebraMat], [FaerMat]
+/// - Sparse matrices: [FaerSparseMat]
+/// - GPU matrices: `CudaMat` (requires `cuda` feature)
 pub mod matrix;
+
+/// Nonlinear solver implementations and traits.
+///
+/// This module defines the [NonLinearSolver] trait and provides the [NewtonNonlinearSolver] implementation.
+/// It also includes:
+/// - [LineSearch] implementations for globalization ([NoLineSearch], [BacktrackingLineSearch])
+/// - Root finding algorithms via [RootFinder]
+/// - Convergence testing via [Convergence]
+///
+/// Nonlinear solvers are used internally by ODE solvers to solve implicit equations.
 pub mod nonlinear_solver;
+
+/// ODE equations and traits.
+///
+/// This module defines the [OdeEquations] trait and specialized variants:
+/// - [OdeEquationsImplicit] for implicit ODEs with mass matrices
+/// - [OdeEquationsImplicitSens] for forward sensitivity equations
+/// - [OdeEquationsAdjoint] for adjoint sensitivity equations
+///
+/// It also provides implementations:
+/// - [DiffSl] for equations specified in the DiffSL domain-specific language
+/// - [SensEquations] and [AdjointEquations] for sensitivity computations
+///
+/// All the test equations used in Diffsol's test suite are also provided here.
 pub mod ode_equations;
+
+/// ODE solver implementations and traits.
+///
+/// This module provides the complete ODE solving interface including:
+/// - [OdeSolverMethod] trait with implementations: [Bdf], [Sdirk], [ExplicitRk]
+/// - [OdeSolverProblem] for problem setup (equations, parameters, tolerances, solver options etc.)
+/// - [OdeSolverState] for managing solution state (including state vector, sensitivities, time, step size etc.)
+/// - [OdeBuilder] for convenient problem construction (builds and configures [OdeSolverProblem])
+/// - [Checkpointing] and [HermiteInterpolator] for solution interpolation
 pub mod ode_solver;
+
+/// Operator types and traits (e.g. non-linear, linear and constant operators; as well as their jacobians).
+///
+/// This module defines fundamental operator traits, for example:
+/// - [NonLinearOp] and variants for Jacobians and sensitivities
+/// - [LinearOp] for linear operators
+/// - [ConstantOp] for constants
+///
+/// It also provides concrete implementations, for example:
+/// - [Closure] for wrapping user-provided closures
+/// - [LinearClosure] for linear operators
+/// - [ConstantClosure] for constants
+/// - [MatrixOp] for explicit matrix operators
 pub mod op;
+
+/// Scalar types and traits.
+///
+/// This module defines the [Scalar] trait that all floating-point types used in DiffSol must implement.
+/// It aggregates requirements from nalgebra, faer, and num_traits to ensure compatibility with linear algebra operations.
+///
+/// Implementations are provided for `f32` and `f64`.
+/// GPU scalar types are available via `ScalarCuda` (requires `cuda` feature).
 pub mod scalar;
-pub mod solver;
+
+/// Vector types and traits.
+///
+/// This module defines the [Vector] trait and related traits for vector operations:
+/// - [VectorView] and [VectorViewMut] for borrowed views
+/// - [VectorIndex] for index collections
+/// - [VectorHost] for CPU-resident vectors with direct access
+///
+/// Implementations are provided for:
+/// - [NalgebraVec] using nalgebra vectors
+/// - [FaerVec] using faer vectors
+/// - `CudaVec` for GPU computation (requires `cuda` feature)
 pub mod vector;
 
+/// Error types and handling.
+///
+/// This module defines the [DiffsolError] enum and specialized error variants
+/// for different failure modes in ODE solving, including parsing, compilation,
+/// and numerical errors.
+pub mod error;
+
+pub use error::DiffsolError;
+
 #[cfg(feature = "sundials")]
 pub mod sundials_sys;
 
@@ -258,5 +366,3 @@ pub use scalar::cuda::{CudaType, ScalarCuda};
 pub use vector::cuda::{CudaIndex, CudaVec, CudaVecMut, CudaVecRef};
 
 pub use scalar::scale;
-
-pub mod error;