Fix inaccuracies in rank estimation when using FFT histogram method

CHANGELOG.rst

@@ -12,6 +12,10 @@ Not released
 * Support numpy 2.
 * Raise minimum supported python version to 3.10.
 
+* Added a new histogram-based ranking method called 'scaledhist'. 
+  This method enables correct FFT convolutions using scaled 
+  histograms and improves precision by incorporating upper and lower histograms.
+
 v0.6.0 (2024/09/03)
 -------------------
 

setup.py

@@ -70,7 +70,7 @@ def env_true(env):
 
 print(f"Build config: {noflags=} {portable=} {x86_64_v3=} {rustflags=}.")
 
-scalib_features = ["pyo3/abi3"]
+scalib_features = ["pyo3/abi3", "ntl"]
 
 if env_true(os.environ.get("SCALIB_BLIS")):
     scalib_features.append("blis")

src/scalib/postprocessing/rankestimation.py

@@ -12,6 +12,14 @@
 to specify the number of bins in the histograms (whereas `rank_accuracy` tunes
 this parameter automatically).
 
+By default, the `rank_accuracy` function uses scaled histograms for convolution. 
+Unlike the standard histogram approach, this method scales down the bins before
+each convolution by a factor chosen to prevent exceeding a predefined limit.
+After the convolution, the histogram is scaled back to restore its original magnitude.
+This ensures correct functionality by avoiding negative values due to f64 overflows.
+Additionally, two histograms are used — one for the lower bound and one for the upper
+bound — to mitigate rounding errors and maintain precision.
+
 Examples
 --------
 
@@ -82,7 +90,8 @@ def rank_nbin(costs, key, nbins, method="hist"):
     method : string
         Method used to estimate the rank. Can be the following:
 
-        * "hist": using histograms (default).
+        * "scaledhist": using two scaled histograms (default).
+        * "hist": using histograms.
         * "naive": enumerate possible keys (very slow).
         * "histbignum": using NTL library, allows better precision.
 
@@ -100,7 +109,7 @@ def rank_nbin(costs, key, nbins, method="hist"):
         )
 
 
-def rank_accuracy(costs, key, acc_bit=1.0, method="hist", max_nb_bin=2**26):
+def rank_accuracy(costs, key, acc_bit=1.0, method="scaledhist", max_nb_bin=2**26):
     r"""Estimate the rank of the full keys based on scores based on histograms.
 
     Parameters
@@ -121,7 +130,8 @@ def rank_accuracy(costs, key, acc_bit=1.0, method="hist", max_nb_bin=2**26):
     method : string
         Method used to estimate the rank. Can be the following:
 
-        * "hist": using histograms (default).
+        * "scaledhist": using two scaled histograms (default).
+        * "hist": using histograms.
         * "ntl": using NTL library, allows better precision.
     max_nb_bin : int, default: 2**26
         Maximum number of bins to use.

src/scalib_ext/ranklib/src/histogram.rs

@@ -6,12 +6,24 @@ mod bnp {
     #![allow(non_snake_case)]
     include!(concat!(env!("OUT_DIR"), "/binding_bnp.rs"));
 }
+//f64 has a 53 bit mantissa, to allow multiplications in the convolution use roughly the sqrt of the mantissa as limit
+const CONV_LIMIT: f64 = 67108864.0 as f64; //limit at 2^26
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum HistogramType {
+    F64Hist,
+    ScaledF64Hist,
+    BigNumHist,
+}
 
 pub trait Histogram {
     fn new(size: usize) -> Self;
-    fn convolve(&self, other: &Self) -> Self;
+    fn convolve(&mut self, other: &mut Self) -> Self;
     fn coefs_f64(&self) -> Vec<f64>;
+    fn coefs_f64_upper(&self) -> Vec<f64>;
     fn from_elems(size: usize, iter: impl Iterator<Item = usize>) -> Self;
+    fn scale_back(self) -> Self;
+    fn histogram_type(&self) -> HistogramType;
 }
 
 type FFT = std::sync::Arc<dyn realfft::RealToComplex<f64>>;
@@ -24,6 +36,15 @@ pub struct F64Hist {
     ifft: IFFT,
 }
 
+#[derive(Clone)]
+pub struct ScaledF64Hist {
+    state_lower: Vec<f64>,
+    state_upper: Vec<f64>,
+    fft: FFT,
+    ifft: IFFT,
+    scale: f64,
+}
+
 impl Histogram for F64Hist {
     fn new(size: usize) -> Self {
         let mut planner = realfft::RealFftPlanner::<f64>::new();
@@ -33,7 +54,7 @@ impl Histogram for F64Hist {
             ifft: planner.plan_fft_inverse(2 * size),
         }
     }
-    fn convolve(&self, other: &Self) -> Self {
+    fn convolve(&mut self, other: &mut Self) -> Self {
         assert_eq!(self.state.len(), other.state.len());
         let mut self_tr = {
             let mut tr = self.fft.make_output_vec();
@@ -70,6 +91,9 @@ impl Histogram for F64Hist {
     fn coefs_f64(&self) -> Vec<f64> {
         self.state.clone()
     }
+    fn coefs_f64_upper(&self) -> Vec<f64> {
+        self.state.clone()
+    }
     fn from_elems(size: usize, iter: impl Iterator<Item = usize>) -> Self {
         let mut res = Self::new(size);
         for elem in iter {
@@ -79,6 +103,156 @@ impl Histogram for F64Hist {
         }
         return res;
     }
+    fn scale_back(self) -> Self {
+        self
+    }
+    fn histogram_type(&self) -> HistogramType {
+        HistogramType::F64Hist
+    }
+}
+impl Histogram for ScaledF64Hist {
+    fn new(size: usize) -> Self {
+        let mut planner = realfft::RealFftPlanner::<f64>::new();
+        Self {
+            state_lower: vec![0.0; size],
+            state_upper: vec![0.0; size],
+            fft: planner.plan_fft_forward(2 * size),
+            ifft: planner.plan_fft_inverse(2 * size),
+            scale: 1.0,
+        }
+    }
+    /// Convolve two histogram objects
+    /// Each histogram object includes a compressed _lower and _upper histogram
+    /// The _lower and _upper histograms are convolved separatley
+    /// Multiply the scales of the histogram objects to keep track of the compression ratio
+    /// other: the histogram we want to convolve with
+    fn convolve(&mut self, other: &mut Self) -> Self {
+        assert_eq!(self.state_lower.len(), other.state_lower.len());
+        assert_eq!(self.state_upper.len(), other.state_upper.len());
+        self.rescale();
+        other.rescale();
+
+        let mut self_tr = {
+            let mut tr = self.fft.make_output_vec();
+            let mut input = vec![0.0; 2 * self.state_lower.len()];
+            input
+                .iter_mut()
+                .zip(self.state_lower.iter())
+                .for_each(|(i, s)| *i = *s);
+            self.fft.process(&mut input, &mut tr).unwrap();
+            tr
+        };
+        let mut self_tr_upper = {
+            let mut tr = self.fft.make_output_vec();
+            let mut input = vec![0.0; 2 * self.state_upper.len()];
+            input
+                .iter_mut()
+                .zip(self.state_upper.iter())
+                .for_each(|(i, s)| *i = *s);
+            self.fft.process(&mut input, &mut tr).unwrap();
+            tr
+        };
+        let other_tr = {
+            let mut tr = self.fft.make_output_vec();
+            let mut input = vec![0.0; 2 * self.state_lower.len()];
+            input
+                .iter_mut()
+                .zip(other.state_lower.iter())
+                .for_each(|(i, s)| *i = *s / (self.state_lower.len() as f64 * 2.0));
+            self.fft.process(&mut input, &mut tr).unwrap();
+            tr
+        };
+        let other_tr_upper = {
+            let mut tr = self.fft.make_output_vec();
+            let mut input = vec![0.0; 2 * self.state_upper.len()];
+            input
+                .iter_mut()
+                .zip(other.state_upper.iter())
+                .for_each(|(i, s)| *i = *s / (self.state_upper.len() as f64 * 2.0));
+            self.fft.process(&mut input, &mut tr).unwrap();
+            tr
+        };
+        self_tr
+            .iter_mut()
+            .zip(other_tr.iter())
+            .for_each(|(s, o)| *s *= *o);
+        self_tr_upper
+            .iter_mut()
+            .zip(other_tr_upper.iter())
+            .for_each(|(s, o)| *s *= *o);
+
+        let mut res = vec![0.0; 2 * self.state_lower.len()];
+        let mut res_upper = vec![0.0; 2 * self.state_upper.len()];
+        self.ifft.process(&mut self_tr, &mut res).unwrap();
+        self.ifft
+            .process(&mut self_tr_upper, &mut res_upper)
+            .unwrap();
+        return Self {
+            state_lower: res[..self.state_lower.len()]
+                .iter()
+                .map(|x| x.round())
+                .collect(),
+            fft: self.fft.clone(),
+            ifft: self.ifft.clone(),
+            state_upper: res_upper[..self.state_upper.len()]
+                .iter()
+                .map(|x| x.round())
+                .collect(),
+            scale: self.scale * other.scale,
+        };
+    }
+    fn coefs_f64(&self) -> Vec<f64> {
+        self.state_lower.clone()
+    }
+    fn coefs_f64_upper(&self) -> Vec<f64> {
+        self.state_upper.clone()
+    }
+    fn from_elems(size: usize, iter: impl Iterator<Item = usize>) -> Self {
+        let mut res = Self::new(size);
+        for elem in iter {
+            if elem < size {
+                res.state_lower[elem] += 1.0;
+                res.state_upper[elem] += 1.0;
+            }
+        }
+        return res;
+    }
+    /// Multiply the compressed histograms by the scaling factor to restore the original bin counts
+    fn scale_back(self) -> Self {
+        let new_scale = 1.0;
+        let new_state: Vec<f64> = self.state_lower.iter().map(|s| (s * self.scale)).collect();
+        let new_upper_state: Vec<f64> = self.state_upper.iter().map(|s| (s * self.scale)).collect();
+        return Self {
+            state_lower: new_state,
+            fft: self.fft.clone(),
+            ifft: self.ifft.clone(),
+            state_upper: new_upper_state,
+            scale: new_scale,
+        };
+    }
+    fn histogram_type(&self) -> HistogramType {
+        HistogramType::ScaledF64Hist
+    }
+}
+impl ScaledF64Hist {
+    /// Check whether the sum of the bin counts exceeds the predefined limit for safe convolution
+    /// If exceeds, adjust the state values by a scaling factor to ensure they remain within the convolution limit
+    /// Multiply the original scale by the new scaling factor to accurately track the total compression of the histogram
+    fn rescale(&mut self) {
+        let sum: f64 = self.state_upper.iter().sum();
+        if sum > CONV_LIMIT {
+            let scaler: f64 = sum / CONV_LIMIT;
+            let temp_scaler: f64 = 1.0 / scaler;
+            self.scale *= scaler;
+            self.state_lower
+                .iter_mut()
+                .zip(self.state_upper.iter_mut())
+                .for_each(|(low, up)| {
+                    *low = (*low * temp_scaler).floor();
+                    *up = (*up * temp_scaler).ceil();
+                });
+        }
+    }
 }
 
 #[cfg(feature = "ntl")]
@@ -89,7 +263,7 @@ impl Histogram for BigNumHist {
     fn new(nb_bins: usize) -> Self {
         return Self(unsafe { bnp::bnp_new_ZZX() }, nb_bins);
     }
-    fn convolve(&self, other: &Self) -> Self {
+    fn convolve(&mut self, other: &mut Self) -> Self {
         assert_eq!(self.1, other.1);
         let res = Self::new(self.1);
         unsafe {
@@ -101,6 +275,9 @@ impl Histogram for BigNumHist {
     fn coefs_f64(&self) -> Vec<f64> {
         self.coefs().collect()
     }
+    fn coefs_f64_upper(&self) -> Vec<f64> {
+        self.coefs().collect()
+    }
     fn from_elems(nb_bins: usize, iter: impl Iterator<Item = usize>) -> Self {
         unsafe {
             let hist = bnp::bnp_new_ZZX();
@@ -114,6 +291,12 @@ impl Histogram for BigNumHist {
             return Self(hist, nb_bins);
         }
     }
+    fn scale_back(self) -> Self {
+        self
+    }
+    fn histogram_type(&self) -> HistogramType {
+        HistogramType::BigNumHist
+    }
 }
 
 #[cfg(feature = "ntl")]

src/scalib_ext/ranklib/src/lib.rs

@@ -49,6 +49,7 @@ pub enum RankingMethod {
     #[cfg(feature = "hellib")]
     Hellib,
     Hist,
+    ScaledHist,
     #[cfg(feature = "ntl")]
     HistBigNum,
 }
@@ -78,6 +79,9 @@ impl RankingMethod {
                 rank_hellib(&problem.costs, &problem.real_key, nb_bin, merge.unwrap())
             }
             RankingMethod::Hist => merged_problem.rank_hist::<histogram::F64Hist>(nb_bin),
+            RankingMethod::ScaledHist => {
+                merged_problem.rank_hist::<histogram::ScaledF64Hist>(nb_bin)
+            }
             #[cfg(feature = "ntl")]
             RankingMethod::HistBigNum => merged_problem.rank_hist::<histogram::BigNumHist>(nb_bin),
         }