add uncertainty band

galactic-forensics · Apr 15, 2024 · a89177c · a89177c
1 parent e6e9562
commit a89177c
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 1 deletion.
diff --git a/src/lib.rs b/src/lib.rs
@@ -31,6 +31,7 @@
 //! ```
 
 use std::fmt;
+use ndarray::Array1;
 
 mod regression;
 
@@ -84,6 +85,14 @@ impl fmt::Display for LinearFit {
     }
 }
 
+pub struct UncertaintyBand {
+    pub x: Array1<f64>,
+    pub y_ub_min: Array1<f64>,
+    pub y_ub_max: Array1<f64>,
+}
+
+// todo add docs, impl. Display for UncertaintyBand
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/src/regression.rs b/src/regression.rs
@@ -1,7 +1,7 @@
 /// Module for linear regression calculations.
 use ndarray::prelude::*;
 
-use crate::LinearFit;
+use crate::{LinearFit, UncertaintyBand};
 
 /// Data structure.
 ///
@@ -105,6 +105,57 @@ impl Data {
         Ok(result)
     }
 
+    pub fn uncertainty_band(
+        &mut self,
+        sigma: Option<f64>,
+        bins: Option<usize>,
+        xrange: Option<&Array1<f64>>,
+    ) -> UncertaintyBand {
+        let xrange_return = match xrange {
+            Some(x) => x.clone(),
+            None => {
+                let num_bins = bins.unwrap_or_else(|| 100);
+                let min_x = *self
+                    .xdat
+                    .iter()
+                    .min_by(|x, y| x.partial_cmp(y).unwrap())
+                    .unwrap();
+                let max_x = *self
+                    .xdat
+                    .iter()
+                    .max_by(|x, y| x.partial_cmp(y).unwrap())
+                    .unwrap();
+                Array1::linspace(min_x, max_x, num_bins)
+            }
+        };
+
+        let mut y_ub: Array1<f64> = Array1::zeros(xrange_return.len());
+
+        // save the current xdat
+        let xdat_save = self.xdat.clone();
+
+        for (it, deltax) in xrange_return.iter().enumerate() {
+            self.xdat = xdat_save.mapv(|x: f64| x - deltax);
+            let result = self.linear_fit().unwrap();
+            y_ub[it] = result.intercept[1];
+        }
+
+        y_ub *= sigma.unwrap_or(1.0);
+
+        // write the saved xdat back
+        self.xdat = xdat_save;
+        let result = self.linear_fit().unwrap();
+
+        let y_ub_min = &xrange_return * result.slope[0] + result.intercept[0] - &y_ub;
+        let y_ub_max = &xrange_return * result.slope[0] + result.intercept[0] - &y_ub;
+
+        UncertaintyBand {
+            x: xrange_return,
+            y_ub_min,
+            y_ub_max,
+        }
+    }
+
     /// Calculate MSWD of the linear regression.
     fn calculate_mswd(&self, slope: f64, intercept: f64, weights: &Array1<f64>) -> f64 {
         let chi_square = (weights