xad-rs 0.1.1

//! Full-Jacobian helpers for vector-valued functions.
//!
//! Computes the Jacobian matrix `J[i][j] = ∂output_i / ∂input_j` for a
//! function `f: Rⁿ → Rᵐ` given as a closure. Two mode-specific entry
//! points are provided and you should pick whichever is cheaper for your
//! problem shape:
//!
//! - [`compute_jacobian_rev`] — reverse-mode; one tape pass per output
//!   row. Efficient when `m ≪ n` (many inputs, few outputs).
//! - [`compute_jacobian_fwd`] — forward-mode; one forward pass per input
//!   column. Efficient when `n ≪ m` (few inputs, many outputs).
//!
//! Both helpers return a row-major `m × n` matrix as `Vec<Vec<T>>`. For
//! dense Jacobians with both `n` and `m` small, reverse mode usually wins
//! because the reverse sweep is heavily optimised in this crate.
//!
//! # Example
//!
//! ```
//! use xad_rs::jacobian::compute_jacobian_rev;
//!
//! // f(x, y) = [x · y, x + y]
//! //        J = [[y, x], [1, 1]]
//! let j = compute_jacobian_rev(&[3.0_f64, 5.0], |v| {
//!     vec![&v[0] * &v[1], &v[0] + &v[1]]
//! });
//! assert_eq!(j[0][0], 5.0); assert_eq!(j[0][1], 3.0);
//! assert_eq!(j[1][0], 1.0); assert_eq!(j[1][1], 1.0);
//! ```

use crate::areal::AReal;
use crate::freal::FReal;
use crate::traits::Scalar;
use crate::tape::{Tape, TapeStorage};

/// Compute the Jacobian using reverse (adjoint) mode.
///
/// Efficient when `num_outputs << num_inputs`. Requires one reverse pass per output.
///
/// `func` takes a slice of `AReal` inputs and returns a vector of `AReal` outputs.
/// Returns a row-major `num_outputs x num_inputs` matrix.
pub fn compute_jacobian_rev<T, F>(inputs: &[T], func: F) -> Vec<Vec<T>>
where
    T: TapeStorage,
    F: Fn(&[AReal<T>]) -> Vec<AReal<T>>,
{
    let mut tape = Tape::<T>::new(true);
    tape.activate();

    let mut ad_inputs: Vec<AReal<T>> = inputs.iter().map(|&v| AReal::new(v)).collect();
    AReal::register_input(&mut ad_inputs, &mut tape);

    let mut ad_outputs = func(&ad_inputs);
    AReal::register_output(&mut ad_outputs, &mut tape);

    let num_inputs = ad_inputs.len();
    let num_outputs = ad_outputs.len();
    let mut jacobian = vec![vec![T::zero(); num_inputs]; num_outputs];

    for i in 0..num_outputs {
        tape.clear_derivatives();
        ad_outputs[i].set_adjoint(&mut tape, T::one());
        tape.compute_adjoints();

        for j in 0..num_inputs {
            jacobian[i][j] = ad_inputs[j].adjoint(&tape);
        }
    }

    Tape::<T>::deactivate_all();
    jacobian
}

/// Compute the Jacobian using forward (tangent) mode.
///
/// Efficient when `num_inputs << num_outputs`. Requires one forward pass per input.
///
/// `func` takes a slice of `FReal` inputs and returns a vector of `FReal` outputs.
/// Returns a row-major `num_outputs x num_inputs` matrix.
pub fn compute_jacobian_fwd<T, F>(inputs: &[T], func: F) -> Vec<Vec<T>>
where
    T: Scalar,
    F: Fn(&[FReal<T>]) -> Vec<FReal<T>>,
{
    let zero_inputs: Vec<FReal<T>> = inputs
        .iter()
        .map(|&v| FReal::constant(v))
        .collect();
    let num_outputs = func(&zero_inputs).len();
    let num_inputs = inputs.len();
    let mut jacobian = vec![vec![T::zero(); num_inputs]; num_outputs];

    for j in 0..num_inputs {
        let mut fwd_inputs: Vec<FReal<T>> = inputs
            .iter()
            .map(|&v| FReal::constant(v))
            .collect();
        fwd_inputs[j].set_derivative(T::one());

        let fwd_outputs = func(&fwd_inputs);

        for i in 0..num_outputs {
            jacobian[i][j] = fwd_outputs[i].derivative();
        }
    }

    jacobian
}