diff --git a/Cargo.lock b/Cargo.lock
index ff47160..556d6de 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1309,8 +1309,9 @@ dependencies = [
[[package]]
name = "rusty-fitpack"
-version = "0.1.0"
-source = "git+https://github.com/mitric-lab/Rusty-FITPACK#2275dfe6739b5e6c7cba50d4a023446d7e1345ce"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92efef01d576920212b373e3ba3a3938b81f8af25a79282eb42e9ada1e577c3b"
dependencies = [
"ndarray 0.12.1",
]
diff --git a/Cargo.toml b/Cargo.toml
index b3925bb..4389a08 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,7 +18,7 @@ levenberg-marquardt = "0.13.1"
nalgebra = "0.32.3"
num-complex = "0.4.3"
polyfit-rs = "0.2.1"
-rusty-fitpack = { git = "https://github.com/mitric-lab/Rusty-FITPACK", version = "0.1.0" }
+rusty-fitpack = "0.1.1"
rayon = "1.8.0"
@@ -42,4 +42,4 @@ name = "xas_group_benchmark_parallel"
harness = false
[profile.bench]
-debug = true
\ No newline at end of file
+debug = true
diff --git a/README.md b/README.md
index d9b55f9..7809159 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,7 @@ Additionally, this project seeks to leverage Rust's ecosystem to create a genera
- [x] Standard EXAFS analysis (find_e0, preedge postedge normalization, AUTOBK, FFT, IFFT)
- [x] Parallel processing using Rayon. (For example, M1 Macbook Pro with 10 cores can process 10000 spectra in 20 seconds, which is ~x10 enhancement without parallelization. Numpy + xraylarch takes 145 seconds.)
+- [x] Optimization on AUTOBK. The AUTOBK process were optimized with providing an analytical Jacobian to speed up the minimization process by Leverberg-Marquardt algorithm.
## Future Developments
diff --git a/benches/flamegraph_xraylarch.svg b/benches/flamegraph_xraylarch.svg
new file mode 100644
index 0000000..003d25d
--- /dev/null
+++ b/benches/flamegraph_xraylarch.svg
@@ -0,0 +1,14588 @@
+
+
+
diff --git a/benches/flamegraph_xraytsubaki_analytical_optimization.svg b/benches/flamegraph_xraytsubaki_analytical_optimization.svg
new file mode 100644
index 0000000..f7b0e0a
--- /dev/null
+++ b/benches/flamegraph_xraytsubaki_analytical_optimization.svg
@@ -0,0 +1,491 @@
+
\ No newline at end of file
diff --git a/benches/flamegraph_xraytsubaki_numerical_optimization.svg b/benches/flamegraph_xraytsubaki_numerical_optimization.svg
new file mode 100644
index 0000000..9033f8e
--- /dev/null
+++ b/benches/flamegraph_xraytsubaki_numerical_optimization.svg
@@ -0,0 +1,491 @@
+
\ No newline at end of file
diff --git a/benches/profiling.md b/benches/profiling.md
new file mode 100644
index 0000000..014bfb8
--- /dev/null
+++ b/benches/profiling.md
@@ -0,0 +1,22 @@
+# Profiling of the preedge, autobk, and fft process.
+
+Profinings were performed using the following command:
+
+```bash
+sudo cargo flamegraph --bench xas_group_benchmark_parallel
+```
+
+## Results
+
+In both cases (numerical and analytical Jacobian), AUTOBK algorithm takes most of the time, and the minimization process is the bottleneck of the entire process. The analytical Jacobian gives roughly x3-4 speedup compared to the numerical Jacobian, but we need to avoid minimization for further speedup.
+
+
+
+
+
+## Appectix: Profiling of xraylarch
+
+Profile of python + xraylarch (preedge, autobk, and fft) were also measured. In this case, AUTOBK is also the bottleneck of the entire process.
+xraytsubaki with single core and numerical Jacobian give similar performance to xraylarch.
+
+
diff --git a/benches/xas_group_benchmark_single.rs b/benches/xas_group_benchmark_single.rs
index 66b493e..9ab1e4b 100644
--- a/benches/xas_group_benchmark_single.rs
+++ b/benches/xas_group_benchmark_single.rs
@@ -22,7 +22,7 @@ fn criterion_benchmark(c: &mut Criterion) {
let mut group = XASGroup::new();
- for _ in 0..10_000 {
+ for _ in 0..100 {
group.add_spectrum(xafs_test_spectrum.clone());
}
diff --git a/src/xafs/background.rs b/src/xafs/background.rs
index b19b6b6..da49814 100644
--- a/src/xafs/background.rs
+++ b/src/xafs/background.rs
@@ -14,7 +14,7 @@ use rusty_fitpack;
// Import internal dependencies
use super::lmutils::LMParameters;
-use super::mathutils::{self, MathUtils};
+use super::mathutils::{self, splev_jacobian, MathUtils};
use super::normalization::{self, Normalization};
use super::nshare::{ToNalgebra, ToNdarray1};
use super::xafsutils::FTWindow;
@@ -375,6 +375,15 @@ impl AUTOBK {
None,
);
+ // Calculate the mu interpolated to the k grid
+ let mu_out = kout.to_vec().interpolate(
+ &kraw
+ .slice_axis(Axis(0), ndarray::Slice::from(0..iemax - iek0 + 1))
+ .to_vec(),
+ &mu.slice_axis(Axis(0), ndarray::Slice::from(iek0..iemax + 1))
+ .to_vec(),
+ )?;
+
let spline_opt = AUTOBKSpline {
coefs: DVector::from_vec(coefs),
knots: DVector::from_vec(knots),
@@ -386,11 +395,7 @@ impl AUTOBK {
.clone()
.to_owned()
.into_nalgebra(),
- mu: mu
- .slice_axis(Axis(0), ndarray::Slice::from(iek0..iemax + 1))
- .clone()
- .to_owned()
- .into_nalgebra(),
+ mu: DVector::from_vec(mu_out),
kout: kout.clone().into_nalgebra(),
ftwin: ftwin.into_nalgebra(),
kweight: self.kweight.unwrap(),
@@ -403,6 +408,7 @@ impl AUTOBK {
clamp_lo: self.clamp_lo.unwrap(),
clamp_hi: self.clamp_hi.unwrap(),
kstep: self.kstep.unwrap(),
+ ..Default::default()
};
let (fit_result, report) = LevenbergMarquardt::new()
@@ -577,14 +583,18 @@ fn spline_eval_nalgebra(
3,
));
- let chi: DVector = DVector::from_vec(
- kout.data
- .as_vec()
- .interpolate(kraw.data.as_vec(), &(mu - &bkg).data.as_vec())
- .unwrap(),
- );
+ // experimental
+ let bkg_out = DVector::from_vec(rusty_fitpack::splev(
+ knots.data.as_vec().clone(),
+ coefs.data.as_vec().clone(),
+ order,
+ kout.data.as_vec().clone(),
+ 3,
+ ));
+
+ let chi = mu - &bkg_out;
- (bkg, chi)
+ (bkg, chi.clone())
}
/// Struct for solving Levenberg-Marquardt optimization for AUTOBK
@@ -605,6 +615,7 @@ struct AUTOBKSpline {
pub clamp_lo: i32,
pub clamp_hi: i32,
pub kstep: f64,
+ pub scale: f64,
}
impl Default for AUTOBKSpline {
@@ -625,6 +636,7 @@ impl Default for AUTOBKSpline {
clamp_lo: 1,
clamp_hi: 1,
kstep: 0.05,
+ scale: 1.0,
}
}
}
@@ -647,15 +659,17 @@ impl AUTOBKSpline {
chi
};
- let chi = chi.component_mul(&self.ftwin);
-
- let mut out: DVector = chi.xftf_fast(self.nfft, self.kstep)[..self.irbkg].realimg();
+ let mut out: DVector = chi
+ .component_mul(&self.ftwin)
+ .xftf_fast(self.nfft, self.kstep)[..self.irbkg]
+ .realimg();
if self.nclamp == 0 {
return out;
}
let scale = 1.0 + 100.0 * out.dot(&out) / out.len() as f64;
+
let low_clamp = self.clamp_lo as f64 * scale * chi.view((0, 0), (self.nclamp as usize, 1));
let high_clamp = self.clamp_hi as f64
@@ -671,8 +685,83 @@ impl AUTOBKSpline {
out
}
+
+ pub fn residual_jacobian(&self, coefs: &DVector) -> DMatrix {
+ // just for calculating the scale
+
+ let scale = if self.nclamp != 0 {
+ let (_, chi) = spline_eval_nalgebra(
+ &self.kraw,
+ &self.mu,
+ &self.knots,
+ coefs,
+ self.order,
+ &self.kout,
+ );
+
+ let chi: DVector = if self.chi_std.is_some() {
+ chi - self.chi_std.as_ref().unwrap()
+ } else {
+ chi
+ };
+
+ let mut out: DVector = chi
+ .component_mul(&self.ftwin)
+ .xftf_fast(self.nfft, self.kstep)[..self.irbkg]
+ .realimg();
+
+ let scale = 1.0 + 100.0 * out.dot(&out) / out.len() as f64;
+
+ scale
+ } else {
+ 1.0
+ };
+
+ let spline_jacobian = -splev_jacobian(
+ self.knots.data.as_vec().clone(),
+ self.coefs.data.as_vec().clone(),
+ self.order,
+ self.kout.data.as_vec().clone(),
+ 3,
+ );
+ let num_cols = self.coefs.len();
+
+ let jacobian_columns = spline_jacobian
+ .column_iter()
+ .map(|chi_der| {
+ let mut out: DVector = chi_der
+ .component_mul(&self.ftwin)
+ .xftf_fast(self.nfft, self.kstep)[..self.irbkg]
+ .realimg();
+
+ if self.nclamp == 0 {
+ return out;
+ }
+
+ // let scale = 1.0 + 100.0 * out.dot(&out) / out.len() as f64;
+
+ let low_clamp =
+ self.clamp_lo as f64 * &scale * chi_der.view((0, 0), (self.nclamp as usize, 1));
+ let high_clamp = self.clamp_hi as f64
+ * &scale
+ * chi_der.view(
+ (chi_der.len() - self.nclamp as usize - 1, 0),
+ (self.nclamp as usize, 1),
+ );
+
+ out.extend(low_clamp.data.as_vec().to_owned());
+ out.extend(high_clamp.data.as_vec().to_owned());
+ out
+ })
+ .collect::>>();
+
+ DMatrix::from_columns(&jacobian_columns)
+ }
}
+use approx::assert_abs_diff_eq;
+use std::time::{Duration, Instant};
+
/// Implementation of LeastSquaresProblem trait for AUTOBK algorithm
impl LeastSquaresProblem for AUTOBKSpline {
type ParameterStorage = Owned;
@@ -694,8 +783,32 @@ impl LeastSquaresProblem for AUTOBKSpline {
/// Jacobian matrix for the Levenberg-Marquardt optimization
/// Jacobian matrix is calculated by numerical differentiation using foward difference
fn jacobian(&self) -> Option> {
- let residual_vec = |coefs: &DVector| AUTOBKSpline::residual_vec(&self, &coefs);
- Some(self.coefs.jacobian(&residual_vec))
+ // let residual_vec = |coefs: &DVector| AUTOBKSpline::residual_vec(&self, &coefs);
+ // Some(self.coefs.jacobian(&residual_vec))
+
+ // let start = Instant::now();
+
+ // let jac1 = self.coefs.jacobian(&residual_vec);
+ // let duration = start.elapsed();
+
+ // println!("jac1: {}", duration.as_secs_f64());
+
+ // let start = Instant::now();
+ // let jac2 = self.residual_jacobian(&self.coefs);
+
+ // let duration = start.elapsed();
+ // println!("jac2: {}", duration.as_secs_f64());
+
+ // println!("jac1: {:?}", jac1.shape());
+ // println!("jac2: {:?}", jac2.shape());
+
+ // jac1.iter().zip(jac2.iter()).for_each(|(x, y)| {
+ // println!("x: {}, y: {}", x, y);
+ // assert_abs_diff_eq!(x, y, epsilon = 1.0e-1);
+ // });
+
+ // Some(jac2)
+ Some(self.residual_jacobian(&self.coefs))
}
}
diff --git a/src/xafs/mathutils.rs b/src/xafs/mathutils.rs
index 15af359..609a942 100644
--- a/src/xafs/mathutils.rs
+++ b/src/xafs/mathutils.rs
@@ -3,9 +3,11 @@ use enterpolation::{
Generator,
};
use errorfunctions::ComplexErrorFunctions;
+use nalgebra::DMatrix;
use ndarray::{Array1, ArrayBase, Ix1, OwnedRepr};
use num_complex::Complex64;
use std::error::Error;
+
pub trait MathUtils {
fn interpolate(&self, x: &Vec, y: &Vec) -> Result
where
@@ -378,6 +380,60 @@ pub fn bessel_I0(x: f64) -> f64 {
sum
}
+/// Calculation jacobian of splev respect to c_i
+///
+///
+pub fn splev_jacobian(t: Vec, c: Vec, k: usize, x: Vec, e: usize) -> DMatrix {
+ let mut y: Vec = vec![0.0; x.len()];
+
+ let k1: usize = k + 1;
+ let k2: usize = k1 + 1;
+ let nk1: usize = t.len() - k1;
+ let tb: f64 = t[k1 - 1];
+ let te: f64 = t[nk1];
+ let mut l: usize = k1;
+ let mut l1: usize = l + 1;
+
+ let mut derivatives: Vec> = vec![vec![0.0; c.len()]; x.len()];
+
+ for (i, &arg) in x.iter().enumerate() {
+ let mut arg = arg;
+ if arg < tb && e == 3 {
+ arg = tb;
+ } else if arg > te && e == 3 {
+ arg = te;
+ }
+
+ let mut l = k1;
+ let mut l1 = l + 1;
+ while arg < t[l - 1] && l1 != k2 {
+ l1 = l;
+ l -= 1;
+ }
+ while arg >= t[l1 - 1] && l != nk1 {
+ l = l1;
+ l1 += 1;
+ }
+
+ let h = rusty_fitpack::fpbspl::fpbspl(arg, &t, k, l);
+
+ let mut ll = l - k1;
+ for j in 1..=k1 {
+ ll += 1;
+ if ll - 1 < c.len() {
+ derivatives[i][ll - 1] = h[j - 1];
+ }
+ }
+ }
+
+ DMatrix::from_vec(
+ derivatives[0].len(),
+ derivatives.len(),
+ derivatives.concat(),
+ )
+ .transpose()
+}
+
#[cfg(test)]
mod tests {
use std::vec;
@@ -387,6 +443,9 @@ mod tests {
use crate::xafs::tests::TEST_TOL;
use crate::xafs::tests::TOP_DIR;
use approx::assert_abs_diff_eq;
+ use nalgebra::DVector;
+
+ const NUMERICAL_TEST_TOL: f64 = 1e-6;
#[test]
fn test_argsort_float() {
@@ -599,4 +658,57 @@ mod tests {
let v = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
assert_eq!(v.diff(), Array1::from_vec(vec![1.0, 1.0, 1.0, 1.0]));
}
+
+ #[test]
+ fn test_splev_jacobian() {
+ let x = vec![
+ 0.0, 0.555, 1.111, 1.666, 2.222, 2.777, 3.333, 3.888, 4.444, 5.0,
+ ];
+ let y = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, -1.0, -2.0, -3.0, -4.0];
+ let order = 3;
+
+ let (t, c, k) = rusty_fitpack::splrep(
+ x.clone(),
+ y.clone(),
+ None,
+ None,
+ None,
+ Some(order),
+ None,
+ None,
+ None,
+ None,
+ None,
+ None,
+ );
+
+ let e = 3; // Example end condition
+
+ let analytical_jacobian = splev_jacobian(t.clone(), c.clone(), k, x.clone(), e);
+
+ use crate::xafs::lmutils::forward_jacobian_nalgebra_f64;
+
+ let spline = rusty_fitpack::splev(t.clone(), c.clone(), k, x.clone(), e);
+
+ let spline = DVector::from(spline);
+
+ let spline_function = |coef: &DVector| {
+ DVector::from(rusty_fitpack::splev(
+ t.clone(),
+ coef.data.as_vec().clone(),
+ k,
+ x.clone(),
+ e,
+ ))
+ };
+
+ let numerical_jacobian = forward_jacobian_nalgebra_f64(&DVector::from(c), &spline_function);
+
+ analytical_jacobian
+ .iter()
+ .zip(numerical_jacobian.iter())
+ .for_each(|(a, b)| {
+ assert_abs_diff_eq!(a, b, epsilon = NUMERICAL_TEST_TOL);
+ });
+ }
}