CHGE: Change ADLift structure (#22)

Author: Axect

src/numerical/root.rs

@@ -84,9 +84,10 @@
 //!
 //! * Walter Gautschi, *Numerical Analysis*, Springer (2012)
 
-use crate::structure::ad::{AD, AD::AD1, ADLift};
+use crate::structure::ad::{AD, AD::AD1, ADFn};
 use std::fmt::Display;
 use RootState::{I, P};
+use crate::traits::stable::StableFn;
 //use std::collections::HashMap;
 //use std::marker::PhantomData;
 
@@ -264,11 +265,11 @@ impl<F: Fn(AD) -> AD> RootFinder<F> {
         match self.method {
             RootFind::Bisection => match self.curr {
                 I(a, b) => {
-                    let lift = ADLift::new(|x| self.f(x));
+                    let f_ad = ADFn::new(|x| self.f(x));
                     let x = 0.5 * (a + b);
-                    let fa = lift.f_0(a);
-                    let fx = lift.f_0(x);
-                    let fb = lift.f_0(b);
+                    let fa = f_ad.call_stable(a);
+                    let fx = f_ad.call_stable(x);
+                    let fb = f_ad.call_stable(b);
                     if (a - b).abs() <= self.tol {
                         self.find = RootBool::Find;
                         self.root = x;
@@ -289,11 +290,11 @@ impl<F: Fn(AD) -> AD> RootFinder<F> {
             },
             RootFind::FalsePosition => match self.curr {
                 I(a, b) => {
-                    let lift = ADLift::new(|x| self.f(x));
-                    let fa = lift.f_0(a);
-                    let fb = lift.f_0(b);
+                    let f_ad = ADFn::new(|x| self.f(x));
+                    let fa = f_ad.call_stable(a);
+                    let fb = f_ad.call_stable(b);
                     let x = (a * fb - b * fa) / (fb - fa);
-                    let fx = lift.f_0(x);
+                    let fx = f_ad.call_stable(x);
                     if (a - b).abs() <= self.tol || fx.abs() <= self.tol {
                         self.find = RootBool::Find;
                         self.root = x;
@@ -327,9 +328,9 @@ impl<F: Fn(AD) -> AD> RootFinder<F> {
             },
             RootFind::Secant => match self.curr {
                 I(xn_1, xn) => {
-                    let lift = ADLift::new(|x| self.f(x));
-                    let fxn_1 = lift.f_0(xn_1);
-                    let fxn = lift.f_0(xn);
+                    let f_ad = ADFn::new(|x| self.f(x));
+                    let fxn_1 = f_ad.call_stable(xn_1);
+                    let fxn = f_ad.call_stable(xn);
                     let x = xn - (xn - xn_1) / (fxn - fxn_1) * fxn;
                     if (x - xn).abs() <= self.tol {
                         self.find = RootBool::Find;

src/structure/ad.rs

@@ -1055,45 +1055,53 @@ impl Div<AD> for f64 {
 ///     let ad0 = 2f64;
 ///     let ad1 = AD1(2f64, 1f64);
 ///
-///     let lift = ADLift::new(f_ad);
+///     let f_ad = ADFn::new(f);
 ///
 ///     let ans_ad0 = ad0.powi(2);
-///     let ans_ad1 = ad1.powi(2);
+///     let ans_ad1 = ad1.powi(2).dx();
 ///
-///     // f_0: f64 -> f64
-///     // f:   AD -> AD
-///     assert_eq!(ans_ad0, lift.f_0(ad0));
-///     assert_eq!(ans_ad1, lift.f(ad1));
+///     assert_eq!(ans_ad0, f_ad.call_stable(ad0));
+///
+///     let f_grad = f_ad.grad();
+///     assert_eq!(ans_ad1, f_grad.call_stable(ad0));
 /// }
 ///
-/// fn f_ad(x: AD) -> AD {
+/// fn f(x: AD) -> AD {
 ///     x.powi(2)
 /// }
 /// ```
-pub struct ADLift<F> {
+pub struct ADFn<F> {
     f: Box<F>,
+    grad_level: usize,
 }
 
-impl<F: Fn(AD) -> AD> ADLift<F> {
+impl<F: Fn(AD) -> AD + Clone> ADFn<F> {
     pub fn new(f: F) -> Self {
         Self {
             f: Box::new(f),
+            grad_level: 0usize,
         }
     }
 
-    pub fn f(&self, t: AD) -> AD {
-        (self.f)(t)
-    }
-
-    pub fn f_0(&self, t: f64) -> f64 {
-        (self.f)(AD::from(t)).into()
+    /// Gradient
+    pub fn grad(&self) -> Self {
+        assert!(self.grad_level < 2, "Higher order AD is not allowed");
+        ADFn {
+            f: (self.f).clone(),
+            grad_level: self.grad_level + 1,
+        }
     }
 }
 
-impl<F: Fn(AD) -> AD> StableFn<f64> for ADLift<F> {
+impl<F: Fn(AD) -> AD> StableFn<f64> for ADFn<F> {
     type Output = f64;
     fn call_stable(&self, target: f64) -> Self::Output {
-        self.f(AD::from(target)).into()
+        match self.grad_level {
+            0 => (self.f)(AD::from(target)).into(),
+            1 => (self.f)(AD1(target, 1f64)).dx(),
+            2 => (self.f)(AD2(target, 1f64, 0f64)).ddx(),
+            _ => panic!("Higher order AD is not allowed"),
+        }
     }
 }
 

tests/ad_lift.rs

@@ -20,8 +20,8 @@ fn test_lift_ad2() {
 #[test]
 fn test_lift_f64() {
     let f = 2f64;
-    let lift = ADLift::new(f_ad);
-    let x = lift.f_0(f);
+    let lift = ADFn::new(f_ad);
+    let x = lift.call_stable(f);
     x.print();
     assert_eq!(x, f * 2f64);
 }