Skip to main content

feos_core/ad/
mod.rs

1//! Automatic differentiation with respect to model parameters.
2use crate::Residual;
3use nalgebra::{Const, DefaultAllocator, Dim, U1, allocator::Allocator};
4use num_dual::{Derivative, DualNum, DualSVec};
5
6#[cfg(feature = "ndarray")]
7mod dataset;
8mod properties;
9#[cfg(feature = "ndarray")]
10pub use dataset::*;
11pub use properties::*;
12
13pub(crate) type Gradient<const P: usize> = DualSVec<f64, f64, P>;
14
15/// A model that can be evaluated with derivatives of its parameters.
16pub trait ParametersAD<N: Dim>: Residual<N>
17where
18    DefaultAllocator: Allocator<N>,
19{
20    /// Build the model by requesting each parameter by name.
21    ///
22    /// Call `f(name, differentiable)` for each parameter. The order of calls
23    /// defines the canonical parameter order.
24    ///
25    /// Set `differentiable` to `false` for fixed parameters.
26    fn build<D: DualNum<f64, Inner = f64> + Copy>(
27        f: impl FnMut(&'static str, bool) -> D,
28    ) -> Self::Lifted<D>;
29
30    /// Canonical parameter names in the order defined by [`build`](Self::build).
31    fn parameter_names() -> Vec<&'static str> {
32        let mut names = Vec::new();
33        let _ = Self::build(|name, _| {
34            names.push(name);
35            0.0
36        });
37        names
38    }
39
40    /// Parameter names that can be differentiated, in canonical order.
41    fn differentiable_parameters() -> Vec<&'static str> {
42        let mut names = Vec::new();
43        let _ = Self::build(|name, differentiable| {
44            if differentiable {
45                names.push(name);
46            }
47            0.0
48        });
49        names
50    }
51
52    /// Construct the model with derivative seeds for the `P` named parameters.
53    ///
54    /// - `parameter_values`: all parameter values in the canonical order
55    ///   defined by [`build`](Self::build).
56    /// - `derivative_names`: names of the parameters to differentiate with
57    ///   respect to. Gradient component `i` corresponds to
58    ///   `derivative_names[i]`.
59    fn seed_derivatives<const P: usize>(
60        parameter_values: &[f64],
61        derivative_names: [&str; P],
62    ) -> Self::Lifted<Gradient<P>> {
63        let mut idx = 0;
64        Self::build(|name, _differentiable| {
65            let i = idx;
66            idx += 1;
67            let mut d = Gradient::<P>::from(parameter_values[i]);
68            if let Some(seed_idx) = derivative_names.iter().position(|&n| n == name) {
69                d.eps =
70                    Derivative::<_, _, Const<P>, _>::derivative_generic(Const::<P>, U1, seed_idx);
71            }
72            d
73        })
74    }
75}