add metadata

Author: GraDje13

src/lib.rs

@@ -6,6 +6,7 @@ use log::Logger;
 use f128::f128;
 use float::FloatLike;
 use itertools::Itertools;
+use matrix::{DecompositionResult, SquareMatrix};
 use preprocessing::{TropicalGraph, TropicalSubgraphTable};
 use rand::Rng;
 use sampling::{sample, SamplingError};
@@ -30,6 +31,7 @@ pub struct TropicalSamplingSettings {
     pub upcast_on_failure: bool,
     pub matrix_stability_test: Option<f64>,
     pub print_debug_info: bool,
+    pub return_metadata: bool,
 }
 
 impl Default for TropicalSamplingSettings {
@@ -38,6 +40,7 @@ impl Default for TropicalSamplingSettings {
             upcast_on_failure: true,
             matrix_stability_test: None,
             print_debug_info: false,
+            return_metadata: false,
         }
     }
 }
@@ -72,6 +75,30 @@ pub struct TropicalSampleResult<T: FloatLike, const D: usize> {
     pub u: T,
     pub v: T,
     pub jacobian: T,
+    pub metadata: Option<Metadata<T, D>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct Metadata<T: FloatLike, const D: usize> {
+    pub q_vectors: Vec<Vector<T, D>>,
+    pub lambda: T,
+    pub l_matrix: SquareMatrix<T>,
+    pub decompoisiton_result: DecompositionResult<T>,
+    pub u_vectors: Vec<Vector<T, D>>,
+    pub shift: Vec<Vector<T, D>>,
+}
+
+impl<const D: usize> Metadata<f128, D> {
+    pub fn downcast(&self) -> Metadata<f64, D> {
+        Metadata {
+            q_vectors: self.q_vectors.iter().map(|v| v.downcast()).collect_vec(),
+            lambda: self.lambda.into(),
+            l_matrix: self.l_matrix.downcast(),
+            decompoisiton_result: self.decompoisiton_result.downcast(),
+            u_vectors: self.u_vectors.iter().map(|v| v.downcast()).collect_vec(),
+            shift: self.shift.iter().map(|v| v.downcast()).collect(),
+        }
+    }
 }
 
 impl<const D: usize> TropicalSampleResult<f128, D> {
@@ -83,6 +110,7 @@ impl<const D: usize> TropicalSampleResult<f128, D> {
             u: self.u.into(),
             v: self.v.into(),
             jacobian: self.jacobian.into(),
+            metadata: self.metadata.as_ref().map(|metadata| metadata.downcast()),
         }
     }
 }

src/matrix.rs

@@ -2,6 +2,7 @@ use smallvec::SmallVec;
 use std::ops::{Add, Index, IndexMut, Mul, Sub};
 
 use crate::{float::FloatLike, TropicalSamplingSettings};
+use f128::f128;
 
 #[derive(Debug, Clone, Default)]
 /// square symmetric matrix for use in the tropical sampling algorithm
@@ -264,18 +265,36 @@ impl<T: FloatLike> SquareMatrix<T> {
     }
 }
 
-#[derive(Debug)]
+impl SquareMatrix<f128> {
+    pub fn downcast(&self) -> SquareMatrix<f64> {
+        SquareMatrix {
+            data: self.data.iter().map(|&x| x.into()).collect(),
+            dim: self.dim,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
 pub struct DecompositionResult<T> {
     pub determinant: T,
     pub inverse: SquareMatrix<T>,
     pub q_transposed_inverse: SquareMatrix<T>,
 }
 
+impl DecompositionResult<f128> {
+    pub fn downcast(&self) -> DecompositionResult<f64> {
+        DecompositionResult {
+            determinant: self.determinant.into(),
+            inverse: self.inverse.downcast(),
+            q_transposed_inverse: self.q_transposed_inverse.downcast(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::assert_approx_eq;
-    use f128::f128;
 
     const EPSILON: f64 = 1e-12;
 

src/sampling.rs

@@ -7,7 +7,7 @@ use crate::matrix::{MatrixError, SquareMatrix};
 use crate::mimic_rng::MimicRng;
 use crate::vector::Vector;
 use crate::{float::FloatLike, gamma::inverse_gamma_lr};
-use crate::{TropicalSampleResult, TropicalSamplingSettings};
+use crate::{Metadata, TropicalSampleResult, TropicalSamplingSettings};
 
 fn box_muller<T: FloatLike>(x1: T, x2: T) -> (T, T) {
     let r = (-Into::<T>::into(2.) * x1.ln()).sqrt();
@@ -116,13 +116,27 @@ pub fn sample<T: FloatLike + Into<f64>, const D: usize, #[cfg(feature = "log")]
         .powf(Into::<T>::into(tropical_subgraph_table.tropical_graph.dod))
         * Into::<T>::into(tropical_subgraph_table.cached_factor);
 
+    let metadata = if settings.return_metadata {
+        Some(Metadata {
+            l_matrix,
+            q_vectors,
+            lambda,
+            shift: compute_only_shift(&decomposed_l_matrix.inverse, &u_vectors),
+            decompoisiton_result: decomposed_l_matrix,
+            u_vectors,
+        })
+    } else {
+        None
+    };
+
     Ok(TropicalSampleResult {
         loop_momenta,
         u_trop,
         v_trop,
         u,
         v,
         jacobian,
+        metadata,
     })
 }
 
@@ -378,3 +392,21 @@ fn compute_loop_momenta<T: FloatLike, const D: usize>(
         })
         .collect_vec()
 }
+
+fn compute_only_shift<T: FloatLike, const D: usize>(
+    l_inverse: &SquareMatrix<T>,
+    u_vectors: &[Vector<T, D>],
+) -> Vec<Vector<T, D>> {
+    let num_loops = l_inverse.get_dim();
+    (0..num_loops)
+        .map(|l| {
+            u_vectors
+                .iter()
+                .enumerate()
+                .fold(Vector::new(), |acc, (l_prime, u)| {
+                    let u_part: Vector<T, D> = u * l_inverse[(l, l_prime)];
+                    &acc + &u_part
+                })
+        })
+        .collect_vec()
+}