CKKSVector Polynomial Evaluation

tenseal/binding.cpp

@@ -114,6 +114,8 @@ PYBIND11_MODULE(_tenseal_cpp, m) {
              py::overload_cast<double>(&CKKSVector::mul_plain_inplace))
         .def("mul_plain_", py::overload_cast<const vector<double> &>(
                                &CKKSVector::mul_plain_inplace))
+        .def("polyval", &CKKSVector::polyval)
+        .def("polyval_", &CKKSVector::polyval_inplace)
         // because dot doesn't have a magic function like __add__
         // we prefer to overload it instead of having dot_plain functions
         .def("dot", &CKKSVector::dot_product)

tenseal/tensors/ckksvector.cpp

@@ -316,8 +316,11 @@ CKKSVector& CKKSVector::_mul_plain_inplace(const T& to_mul) {
         this->context->evaluator->multiply_plain_inplace(this->ciphertext,
                                                          plaintext);
     } catch (const std::logic_error& e) {  // result ciphertext is transparent
+        // TODO: chech if error e is exactly a "ciphertext is transparent" error
         // replace by encryption of zero
         this->context->encryptor->encrypt_zero(this->ciphertext);
+        this->ciphertext.scale() = this->init_scale;
+        return *this;
     }
 
     if (this->context->auto_rescale()) {
@@ -412,4 +415,59 @@ CKKSVector& CKKSVector::replicate_first_slot_inplace(size_t n) {
     return *this;
 }
 
+CKKSVector CKKSVector::polyval(const vector<double>& coefficients) {
+    CKKSVector new_vector = *this;
+    return new_vector.polyval_inplace(coefficients);
+}
+
+CKKSVector& CKKSVector::polyval_inplace(const vector<double>& coefficients) {
+    if (coefficients.size() == 0) {
+        throw invalid_argument(
+            "the coefficients vector need to have at least one element");
+    }
+
+    int degree = static_cast<int>(coefficients.size()) - 1;
+    while (degree >= 0) {
+        if (coefficients[degree] == 0.0)
+            degree--;
+        else
+            break;
+    }
+
+    // null polynomial: output should be an encrypted 0
+    // we can multiply by 0, or return the encryption of zero
+    if (degree == -1) {
+        // we set the vector to the encryption of zero
+        this->context->encryptor->encrypt_zero(this->ciphertext);
+        this->ciphertext.scale() = this->init_scale;
+        return *this;
+    }
+
+    // set result accumulator to the constant coefficient
+    vector<double> cst_coeff(this->size(), coefficients[0]);
+    CKKSVector result(this->context, cst_coeff, this->init_scale);
+
+    // pre-compute squares of x
+    CKKSVector x = *this;
+    int max_square = static_cast<int>(floor(log2(degree)));
+    vector<CKKSVector> x_squares;
+    x_squares.reserve(max_square + 1);
+    x_squares.push_back(x);  // x
+    for (int i = 1; i <= max_square; i++) {
+        // TODO: use square
+        x.mul_inplace(x);
+        x_squares.push_back(x);  // x^(2^i)
+    }
+
+    // coefficients[1] * x + ... + coefficients[degree] * x^(degree)
+    for (int i = 1; i <= degree; i++) {
+        if (coefficients[i] == 0.0) continue;
+        x = compute_polynomial_term(i, coefficients[i], x_squares);
+        result.add_inplace(x);
+    }
+
+    this->ciphertext = result.ciphertext;
+    return *this;
+}
+
 }  // namespace tenseal

tenseal/tensors/ckksvector.h

@@ -101,6 +101,13 @@ class CKKSVector {
     CKKSVector matmul_plain(const vector<vector<double>>& matrix);
     CKKSVector& matmul_plain_inplace(const vector<vector<double>>& matrix);
 
+    /*
+    Polynomial evaluation with `this` as variable.
+    p(x) = coefficients[0] + coefficients[1] * x + ... + coefficients[i] * x^i
+    */
+    CKKSVector polyval(const vector<double>& coefficients);
+    CKKSVector& polyval_inplace(const vector<double>& coefficients);
+
    private:
     /*
     Private evaluation functions to process both scalar and vector arguments.

tenseal/utils/utils.h

@@ -70,6 +70,28 @@ IMPORTANT: Tested only with CKKS.
 Ciphertext& sum_vector(shared_ptr<TenSEALContext> tenseal_context,
                        Ciphertext& vector, size_t size);
 
+template <typename T>
+T compute_polynomial_term(int degree, double coeff,
+                          const vector<T>& x_squares) {
+    if (degree < 1) {
+        throw invalid_argument("degree must be greater or equal to 1");
+    }
+
+    int closest_power_of_2 = static_cast<int>(floor(log2(degree)));
+    int new_degree = degree - (1 << closest_power_of_2);
+    T x = x_squares[closest_power_of_2];  // x^(2^closest_power_of_2)
+
+    if (new_degree == 0 && coeff != 1.0) {
+        // x^(2^closest_power_of_2) * coeff
+        x.mul_plain_inplace(coeff);
+    } else if (new_degree != 0) {
+        // x^(2^closest_power_of_2) * x^(new_degree) * coeff
+        x.mul_inplace(compute_polynomial_term(new_degree, coeff, x_squares));
+    }
+
+    return x;
+}
+
 }  // namespace tenseal
 
 #endif

tests/tensors/test_ckks_vector.py

@@ -848,58 +848,112 @@ def test_vec_plain_matrix_mul_depth2(context, vec, matrix1, matrix2):
 @pytest.mark.parametrize(
     "data, polynom",
     [
-        ([0, 1, 2, 3, 4], lambda x: x * x + x),
-        ([0, 1, 2, 3, 4], lambda x: x * x - x),
-        ([0, 1, 2, 3, 4], lambda x: x * x * x),
-        ([0, 0, 0, 0, 0], lambda x: x * x * x),
+        # null polynom
+        ([0, 1, 2, 3, 4], [0]),
+        ([0, 1, 2, 3, 4], [0, 0]),
+        ([0, 1, 2, 3, 4], [0, 0, 0]),
+        # power of two coeff
+        ([0, 1, 2, 3, 4], [1, 1]),
+        ([0, 1, 2, 3, 4], [0, 1, 1]),
+        ([0, 1, 2, 3, 4], [0, 1, 1, 0, 1]),
+        # random coeff
+        ([0, 1, 2, 3, 4], [-4, -2, 5]),
+        ([0, 0, 0, 0, 0], [0, 0, 0, 1]),
+        ([0, 1, 2, 3, 4], [0, 0, 0, 1]),
+        ([0, 1, 2, 3, 4], [3, 2, 4, 5]),
+        ([0, -1, -2, -3, -4], [-3, -2, -4, -5, 1]),
     ],
 )
-def test_simple_polynomial(context, data, polynom):
+def test_polynomial(context, data, polynom):
     ct = ts.ckks_vector(context, data)
-    expected = [polynom(x) for x in data]
-    result = polynom(ct)
+    expected = [np.polyval(polynom[::-1], x) for x in data]
+    result = ct.polyval(polynom)
 
     decrypted_result = result.decrypt()
     assert _almost_equal(decrypted_result, expected, 1), "Polynomial evaluation is incorrect."
-    # adding plain vector at the end
-    result += data
-    expected = [expected[i] + data[i] for i in range(len(data))]
+
+
+@pytest.mark.parametrize(
+    "data, polynom",
+    [
+        ## high data may result in bigger error (2 is enough for 0.1 error)
+        ([2, 2, 2, 2, 2], [0, 1, 1, 0, 1, 0, 0, 0, 1]),
+        ([0, -1, 2, -3, 4], [5, -3, 4, 73, -3]),
+        ([0, 1, -2, 3, -4], [-3, 0, 5, 1, -2]),
+        ([0, -1, 1, -2, 2], [3, -7, 2, 0, 1, -1, 7, 2, -3]),
+        ([0, -1, 1, -2, 2], [3, -7, 2, 0, 1, -1, 7, 2, -3, -7, 2, 0, 1, -1, 7, 2, -2]),
+        ([0, -1, 1, -2, 2], [0] * 1 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 2 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 3 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 4 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 5 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 6 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 7 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 8 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 9 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 10 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 11 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 12 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 13 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 14 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 15 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 16 + [1]),
+        ([0, -1, 1, -2, 2], [0] * 16 + [2]),
+    ],
+)
+def test_high_degree_polynomial(data, polynom):
+    # special context for higher depth
+    context = ts.context(
+        ts.SCHEME_TYPE.CKKS, 16384, coeff_mod_bit_sizes=[60, 40, 40, 40, 40, 40, 60]
+    )
+    context.global_scale = pow(2, 40)
+    ct = ts.ckks_vector(context, data)
+    expected = [np.polyval(polynom[::-1], x) for x in data]
+    result = ct.polyval(polynom)
+
     decrypted_result = result.decrypt()
-    assert _almost_equal(decrypted_result, expected, 1)
+    if len(polynom) >= 13:
+        # we allow greater error since some polynomial has terms with a high exponent
+        error_tolerance = -1
+    else:
+        error_tolerance = 1
+    assert _almost_equal(
+        decrypted_result, expected, error_tolerance
+    ), "Polynomial evaluation is incorrect."
 
 
 @pytest.mark.parametrize(
     "data, polynom",
     [
-        ([0, 1, 2, 3, 4], lambda x: x * x + x),
-        ([0, 1, 2, 3, 4], lambda x: x * x - x),
-        ([0, 1, 2, 3, 4], lambda x: x * x * x),
-        ([0, 0, 0, 0, 0], lambda x: x * x * x),
+        ([0, 1, 2, 3, 4], [0, 1, 1]),
+        ([0, 1, 2, 3, 4], [0, -1, 1]),
+        ([0, 1, 2, 3, 4], [0, 1, 0, 1]),
+        ([0, 0, 0, 0, 0], [0, 1, 0, 1]),
     ],
 )
-def test_simple_polynomial_modswitch_off(context, data, polynom):
+def test_polynomial_modswitch_off(context, data, polynom):
     context = ts.context(ts.SCHEME_TYPE.CKKS, 8192, 0, [60, 40, 40, 60])
     context.global_scale = 2 ** 40
     context.auto_mod_switch = False
 
     ct = ts.ckks_vector(context, data)
     with pytest.raises(ValueError) as e:
-        result = polynom(ct)
-    assert str(e.value) == "encrypted1 and encrypted2 parameter mismatch"
+        result = ct.polyval(polynom)
+    # encrypted1 and encrypted2 parameter mismatch (or encrypted_ntt and plain_ntt)
+    assert "parameter mismatch" in str(e.value)
 
 
 @pytest.mark.parametrize(
-    "data, polynom",
-    [([0, 1, 2, 3, 4], lambda x: x * x + x), ([0, 1, 2, 3, 4], lambda x: x * x - x),],
+    "data, polynom", [([0, 1, 2, 3, 4], [0, 1, 1]), ([0, 1, 2, 3, 4], [0, -1, 1]),],
 )
-def test_simple_polynomial_rescale_off(context, data, polynom):
+def test_polynomial_rescale_off(context, data, polynom):
     context = ts.context(ts.SCHEME_TYPE.CKKS, 8192, 0, [60, 40, 40, 60])
     context.global_scale = 2 ** 40
     context.auto_rescale = False
 
     ct = ts.ckks_vector(context, data)
     with pytest.raises(ValueError) as e:
-        result = polynom(ct)
+        result = ct.polyval(polynom)
     assert str(e.value) == "scale mismatch"