xad-rs 0.2.0

Automatic differentiation library for Rust — forward/reverse mode AD, a Rust port of the C++ XAD library (https://github.com/auto-differentiation/xad)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

//! `LabeledJacobian` + `compute_labeled_jacobian` — string-keyed Jacobian
//! helper, gated on the `labeled-ndarray` feature.
//!
//! This module is a thin wrapper around [`crate::jacobian::compute_jacobian_rev`]
//! that flattens the row-major `Vec<Vec<f64>>` output into an
//! `ndarray::Array2<f64>` and decorates it with row + column labels.
//!
//! # Example
//!
//! ```
//! # #[cfg(feature = "labeled-ndarray")]
//! # {
//! use xad_rs::areal::AReal;
//! use xad_rs::labeled::jacobian::compute_labeled_jacobian;
//!
//! // f(x, y) = [x · y, x + y]
//! let inputs  = vec![("x".to_string(), 3.0), ("y".to_string(), 5.0)];
//! let outputs = vec!["prod".to_string(), "sum".to_string()];
//! let j = compute_labeled_jacobian(&inputs, &outputs, |v: &[AReal<f64>]| {
//!     vec![&v[0] * &v[1], &v[0] + &v[1]]
//! });
//! assert_eq!(j.matrix[[0, 0]], 5.0);
//! assert_eq!(j.matrix[[0, 1]], 3.0);
//! assert_eq!(j.matrix[[1, 0]], 1.0);
//! assert_eq!(j.matrix[[1, 1]], 1.0);
//! # }
//! ```

use std::sync::Arc;

use ndarray::Array2;

use crate::areal::AReal;
use crate::jacobian::compute_jacobian_rev;
use crate::labeled::VarRegistry;

/// Row-labeled, column-labeled Jacobian matrix.
///
/// `matrix[[i, j]] = ∂outputs[i] / ∂inputs[j]`. Row order matches the
/// `outputs` slice passed to [`compute_labeled_jacobian`]; column order
/// matches the `inputs` slice (which becomes the inner [`VarRegistry`]).
pub struct LabeledJacobian {
    pub rows: Vec<String>,
    pub cols: Arc<VarRegistry>,
    pub matrix: Array2<f64>,
}

/// Compute a labeled Jacobian via reverse-mode AD.
///
/// Delegates the heavy lifting to [`compute_jacobian_rev`]; this wrapper
/// only attaches row + column labels and reshapes the output into an
/// `ndarray::Array2<f64>`.
///
/// # Arguments
///
/// * `inputs`  — input `(name, value)` pairs in the order they should appear
///   as columns of the Jacobian.
/// * `outputs` — output labels in the order they should appear as rows.
/// * `f`       — forward closure operating on positional `&[AReal<f64>]`,
///   returning `Vec<AReal<f64>>` (matches `compute_jacobian_rev`).
///
/// # Panics
///
/// Panics if `f(&inputs)` returns a vector whose length differs from
/// `outputs.len()`.
pub fn compute_labeled_jacobian<F>(
    inputs: &[(String, f64)],
    outputs: &[String],
    f: F,
) -> LabeledJacobian
where
    F: Fn(&[AReal<f64>]) -> Vec<AReal<f64>>,
{
    let input_values: Vec<f64> = inputs.iter().map(|(_, v)| *v).collect();
    let positional = compute_jacobian_rev(&input_values, f);

    let m = outputs.len();
    let n = inputs.len();
    assert_eq!(
        positional.len(),
        m,
        "compute_labeled_jacobian: closure produced {} outputs but {} output labels were given",
        positional.len(),
        m
    );

    let mut matrix = Array2::<f64>::zeros((m, n));
    for (i, row) in positional.iter().enumerate() {
        debug_assert_eq!(row.len(), n);
        for (j, &v) in row.iter().enumerate() {
            matrix[[i, j]] = v;
        }
    }

    let input_names: Vec<String> = inputs.iter().map(|(name, _)| name.clone()).collect();
    LabeledJacobian {
        rows: outputs.to_vec(),
        cols: Arc::new(VarRegistry::from_names(input_names)),
        matrix,
    }
}