Added noise equation for keyswitch over GGSW & updated associated test

2026-02-10 05:06:44 +01:00 · 2025-05-20 14:36:26 +02:00
parent 7d84477e64
commit a803127424
5 changed files with 107 additions and 29 deletions
--- a/core/src/test_fft64/automorphism_key.rs
+++ b/core/src/test_fft64/automorphism_key.rs
@@ -7,7 +7,7 @@ use crate::{
    glwe_ciphertext_fourier::GLWECiphertextFourier,
    glwe_plaintext::GLWEPlaintext,
    keys::{SecretKey, SecretKeyFourier},
-    test_fft64::gglwe::noise_gglwe_product,
+    test_fft64::gglwe::log2_std_noise_gglwe_product,
 };

 #[test]
@@ -101,7 +101,7 @@ fn test_automorphism(p0: i64, p1: i64, log_n: usize, basek: usize, k_ksk: usize,
            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk, col_i);

            let noise_have: f64 = pt.data.std(0, basek).log2();
-            let noise_want: f64 = noise_gglwe_product(
+            let noise_want: f64 = log2_std_noise_gglwe_product(
                module.n() as f64,
                basek,
                0.5,
@@ -192,7 +192,7 @@ fn test_automorphism_inplace(p0: i64, p1: i64, log_n: usize, basek: usize, k_ksk
            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk, col_i);

            let noise_have: f64 = pt.data.std(0, basek).log2();
-            let noise_want: f64 = noise_gglwe_product(
+            let noise_want: f64 = log2_std_noise_gglwe_product(
                module.n() as f64,
                basek,
                0.5,
--- a/core/src/test_fft64/gglwe.rs
+++ b/core/src/test_fft64/gglwe.rs
@@ -212,7 +212,7 @@ fn test_key_switch(
            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk0, col_i);

            let noise_have: f64 = pt.data.std(0, basek).log2();
-            let noise_want: f64 = noise_gglwe_product(
+            let noise_want: f64 = log2_std_noise_gglwe_product(
                module.n() as f64,
                basek,
                0.5,
@@ -314,7 +314,7 @@ fn test_key_switch_inplace(log_n: usize, basek: usize, k_ksk: usize, sigma: f64,
            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk0, col_i);

            let noise_have: f64 = pt.data.std(0, basek).log2();
-            let noise_want: f64 = noise_gglwe_product(
+            let noise_want: f64 = log2_std_noise_gglwe_product(
                module.n() as f64,
                basek,
                0.5,
@@ -571,7 +571,7 @@ fn test_external_product_inplace(log_n: usize, basek: usize, k: usize, sigma: f6
    });
 }

-pub(crate) fn noise_gglwe_product(
+pub(crate) fn var_noise_gglwe_product(
    n: f64,
    basek: usize,
    var_xs: f64,
@@ -597,7 +597,34 @@ pub(crate) fn noise_gglwe_product(
    let mut noise: f64 = (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
    noise += var_msg * var_a_err * a_scale * a_scale * n;
    noise *= rank_in;
+    noise /= b_scale * b_scale;
+    noise
+}
+
+pub(crate) fn log2_std_noise_gglwe_product(
+    n: f64,
+    basek: usize,
+    var_xs: f64,
+    var_msg: f64,
+    var_a_err: f64,
+    var_gct_err_lhs: f64,
+    var_gct_err_rhs: f64,
+    rank_in: f64,
+    a_logq: usize,
+    b_logq: usize,
+) -> f64 {
+    let mut noise: f64 = var_noise_gglwe_product(
+        n,
+        basek,
+        var_xs,
+        var_msg,
+        var_a_err,
+        var_gct_err_lhs,
+        var_gct_err_rhs,
+        rank_in,
+        a_logq,
+        b_logq,
+    );
    noise = noise.sqrt();
-    noise /= b_scale;
    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
 }
--- a/core/src/test_fft64/ggsw.rs
+++ b/core/src/test_fft64/ggsw.rs
@@ -13,9 +13,11 @@ use crate::{
    keys::{SecretKey, SecretKeyFourier},
    keyswitch_key::GLWESwitchingKey,
    tensor_key::TensorKey,
-    test_fft64::gglwe::noise_gglwe_product,
+    test_fft64::gglwe::log2_std_noise_gglwe_product,
 };

+use super::gglwe::var_noise_gglwe_product;
+
 #[test]
 fn encrypt_sk() {
    (1..4).for_each(|rank| {
@@ -146,14 +148,16 @@ fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64)
            ),
    );

+    let var_xs: f64 = 0.5;
+
    let mut sk_in: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
-    sk_in.fill_ternary_prob(0.5, &mut source_xs);
+    sk_in.fill_ternary_prob(var_xs, &mut source_xs);

    let mut sk_in_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
    sk_in_dft.dft(&module, &sk_in);

    let mut sk_out: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
-    sk_out.fill_ternary_prob(0.5, &mut source_xs);
+    sk_out.fill_ternary_prob(var_xs, &mut source_xs);

    let mut sk_out_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
    sk_out_dft.dft(&module, &sk_out);
@@ -213,11 +217,11 @@ fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64)
            module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);

            let noise_have: f64 = pt_have.data.std(0, basek).log2();
-            let noise_want: f64 = noise_gglwe_product(
+            let noise_want: f64 = noise_ggsw_keyswitch(
                module.n() as f64,
                basek,
-                0.5,
-                0.5,
+                col_j,
+                var_xs,
                0f64,
                sigma * sigma,
                0f64,
@@ -226,20 +230,67 @@ fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64)
                k,
            );

-            println!("{} {}", noise_have, noise_want);
-
-            // assert!(
-            // (noise_have - noise_want).abs() <= 0.1,
-            // "{} {}",
-            // noise_have,
-            // noise_want
-            // );
+            assert!(
+                (noise_have - noise_want).abs() <= 0.1,
+                "{} {}",
+                noise_have,
+                noise_want
+            );

            pt_want.data.zero();
        });
    });
 }

+pub(crate) fn noise_ggsw_keyswitch(
+    n: f64,
+    basek: usize,
+    col: usize,
+    var_xs: f64,
+    var_a_err: f64,
+    var_gct_err_lhs: f64,
+    var_gct_err_rhs: f64,
+    rank: f64,
+    a_logq: usize,
+    b_logq: usize,
+) -> f64 {
+    let var_si_x_sj: f64 = n * var_xs * var_xs;
+
+    // Initial KS for col = 0
+    let mut noise: f64 = var_noise_gglwe_product(
+        n,
+        basek,
+        var_xs,
+        var_xs,
+        var_a_err,
+        var_gct_err_lhs,
+        var_gct_err_rhs,
+        rank,
+        a_logq,
+        b_logq,
+    );
+
+    // Other GGSW reconstruction for col > 0
+    if col > 0 {
+        noise += var_noise_gglwe_product(
+            n,
+            basek,
+            var_xs,
+            var_si_x_sj,
+            var_a_err + 1f64 / 12.0,
+            var_gct_err_lhs,
+            var_gct_err_rhs,
+            rank,
+            a_logq,
+            b_logq,
+        );
+        noise += n * noise * var_xs * 0.5;
+    }
+
+    noise = noise.sqrt();
+    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
+}
+
 // fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
 // let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 // let rows: usize = (k_ggsw + basek - 1) / basek;
--- a/core/src/test_fft64/glwe.rs
+++ b/core/src/test_fft64/glwe.rs
@@ -14,7 +14,7 @@ use crate::{
    glwe_plaintext::GLWEPlaintext,
    keys::{GLWEPublicKey, SecretKey, SecretKeyFourier},
    keyswitch_key::GLWESwitchingKey,
-    test_fft64::{gglwe::noise_gglwe_product, ggsw::noise_ggsw_product},
+    test_fft64::{gglwe::log2_std_noise_gglwe_product, ggsw::noise_ggsw_product},
 };

 #[test]
@@ -326,7 +326,7 @@ fn test_keyswitch(
    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);

    let noise_have: f64 = pt_have.data.std(0, basek).log2();
-    let noise_want: f64 = noise_gglwe_product(
+    let noise_want: f64 = log2_std_noise_gglwe_product(
        module.n() as f64,
        basek,
        0.5,
@@ -411,7 +411,7 @@ fn test_keyswitch_inplace(log_n: usize, basek: usize, k_ksk: usize, k_ct: usize,
    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);

    let noise_have: f64 = pt_have.data.std(0, basek).log2();
-    let noise_want: f64 = noise_gglwe_product(
+    let noise_want: f64 = log2_std_noise_gglwe_product(
        module.n() as f64,
        basek,
        0.5,
@@ -502,7 +502,7 @@ fn test_automorphism(

    println!("{}", noise_have);

-    let noise_want: f64 = noise_gglwe_product(
+    let noise_want: f64 = log2_std_noise_gglwe_product(
        module.n() as f64,
        basek,
        0.5,
@@ -581,7 +581,7 @@ fn test_automorphism_inplace(log_n: usize, basek: usize, p: i64, k_autokey: usiz
    module.vec_znx_normalize_inplace(basek, &mut pt_have, 0, scratch.borrow());

    let noise_have: f64 = pt_have.data.std(0, basek).log2();
-    let noise_want: f64 = noise_gglwe_product(
+    let noise_want: f64 = log2_std_noise_gglwe_product(
        module.n() as f64,
        basek,
        0.5,
--- a/core/src/test_fft64/glwe_fourier.rs
+++ b/core/src/test_fft64/glwe_fourier.rs
@@ -6,7 +6,7 @@ use crate::{
    glwe_plaintext::GLWEPlaintext,
    keys::{SecretKey, SecretKeyFourier},
    keyswitch_key::GLWESwitchingKey,
-    test_fft64::{gglwe::noise_gglwe_product, ggsw::noise_ggsw_product},
+    test_fft64::{gglwe::log2_std_noise_gglwe_product, ggsw::noise_ggsw_product},
 };
 use base2k::{FFT64, FillUniform, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, VecZnxToMut, ZnxViewMut};
 use sampling::source::Source;
@@ -132,7 +132,7 @@ fn test_keyswitch(
    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);

    let noise_have: f64 = pt_have.data.std(0, basek).log2();
-    let noise_want: f64 = noise_gglwe_product(
+    let noise_want: f64 = log2_std_noise_gglwe_product(
        module.n() as f64,
        basek,
        0.5,
@@ -220,7 +220,7 @@ fn test_keyswitch_inplace(log_n: usize, basek: usize, k_ksk: usize, k_ct: usize,
    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);

    let noise_have: f64 = pt_have.data.std(0, basek).log2();
-    let noise_want: f64 = noise_gglwe_product(
+    let noise_want: f64 = log2_std_noise_gglwe_product(
        module.n() as f64,
        basek,
        0.5,