CKKS square and power functions

tenseal/binding.cpp

@@ -87,6 +87,12 @@ PYBIND11_MODULE(_tenseal_cpp, m) {
         .def("neg", &CKKSVector::negate)
         .def("neg_", &CKKSVector::negate_inplace)
         .def("neg_inplace", &CKKSVector::negate_inplace)
+        .def("square", &CKKSVector::square)
+        .def("square_", &CKKSVector::square_inplace)
+        .def("square_inplace", &CKKSVector::square_inplace)
+        .def("pow", &CKKSVector::power)
+        .def("pow_", &CKKSVector::power_inplace)
+        .def("pow_inplace", &CKKSVector::power_inplace)
         .def("add", &CKKSVector::add)
         .def("add_", &CKKSVector::add_inplace)
         .def("add_plain", py::overload_cast<double>(&CKKSVector::add_plain))
@@ -130,6 +136,8 @@ PYBIND11_MODULE(_tenseal_cpp, m) {
         .def("mm_", &CKKSVector::matmul_plain_inplace)
         // python arithmetic
         .def("__neg__", &CKKSVector::negate)
+        .def("__pow__", &CKKSVector::power)
+        .def("__ipow__", &CKKSVector::power_inplace)
         .def("__add__", &CKKSVector::add)
         .def("__add__", py::overload_cast<double>(&CKKSVector::add_plain))
         .def("__add__",

tenseal/tensors/ckksvector.cpp

@@ -84,6 +84,63 @@ CKKSVector& CKKSVector::negate_inplace() {
     return *this;
 }
 
+CKKSVector CKKSVector::square() {
+    CKKSVector new_vector = *this;
+    new_vector.square_inplace();
+
+    return new_vector;
+}
+
+CKKSVector& CKKSVector::square_inplace() {
+    this->context->evaluator->square_inplace(this->ciphertext);
+
+    if (this->context->auto_relin()) {
+        this->context->evaluator->relinearize_inplace(
+            this->ciphertext, *this->context->relin_keys());
+    }
+
+    if (this->context->auto_rescale()) {
+        this->context->evaluator->rescale_to_next_inplace(this->ciphertext);
+        this->ciphertext.scale() = this->init_scale;
+    }
+
+    return *this;
+}
+
+CKKSVector CKKSVector::power(unsigned int power) {
+    CKKSVector new_vector = *this;
+    new_vector.power_inplace(power);
+
+    return new_vector;
+}
+
+CKKSVector& CKKSVector::power_inplace(unsigned int power) {
+    // if the power is zero, return a new encrypted vector of ones
+    if (power == 0) {
+        vector<double> ones(this->size(), 1);
+        *this = CKKSVector(this->context, ones, this->init_scale);
+        return *this;
+    }
+
+    if (power == 1) {
+        return *this;
+    }
+
+    if (power == 2) {
+        this->square_inplace();
+        return *this;
+    }
+
+    int closest_power_of_2 = 1 << static_cast<int>(floor(log2(power)));
+    power -= closest_power_of_2;
+    if (power == 0) {
+        this->power_inplace(closest_power_of_2 / 2).square_inplace();
+    } else {
+        this->power_inplace(power).mul_inplace(this->power(closest_power_of_2));
+    }
+    return *this;
+}
+
 CKKSVector CKKSVector::add(CKKSVector to_add) {
     CKKSVector new_vector = *this;
     new_vector.add_inplace(to_add);
@@ -454,8 +511,7 @@ CKKSVector& CKKSVector::polyval_inplace(const vector<double>& coefficients) {
     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.square_inplace();
         x_squares.push_back(x);  // x^(2^i)
     }
 

tenseal/tensors/ckksvector.h

@@ -57,6 +57,18 @@ class CKKSVector {
     CKKSVector negate();
     CKKSVector& negate_inplace();
 
+    /*
+    Compute the square of the CKKSVector.
+    */
+    CKKSVector square();
+    CKKSVector& square_inplace();
+
+    /*
+    Compute the power of the CKKSVector with minimal multiplication depth.
+    */
+    CKKSVector power(unsigned int power);
+    CKKSVector& power_inplace(unsigned int power);
+
     /*
     Encrypted evaluation function operates on two encrypted vectors and
     returns a new CKKSVector which is the result of either addition,

tests/tensors/test_ckks_vector.py

@@ -43,6 +43,100 @@ def test_negate_inplace(context, plain_vec):
     assert _almost_equal(decrypted_result, expected, 1), "Decryption of vector is incorrect"
 
 
+@pytest.mark.parametrize(
+    "plain_vec, power",
+    [
+        ([], 2),
+        ([0], 3),
+        ([0, 1, -1, 2, -2], 0),
+        ([1, -1, 2, -2], 1),
+        ([1, -1, 2, -2], 2),
+        ([1, -1, 2, -2], 3),
+        ([1, -2, 3, -4], 1),
+        ([1, -2, 3, -4], 2),
+        ([1, -2, 3, -4], 3),
+    ],
+)
+def test_power(context, plain_vec, power):
+    ckks_vec = ts.ckks_vector(context, plain_vec)
+    expected = [np.power(v, power) for v in plain_vec]
+    new_vec = ckks_vec ** power
+    decrypted_result = new_vec.decrypt()
+    assert _almost_equal(decrypted_result, expected, 1), "Decryption of vector is incorrect"
+    assert _almost_equal(ckks_vec.decrypt(), plain_vec, 1), "Something went wrong in memory."
+
+
+@pytest.mark.parametrize(
+    "plain_vec, power",
+    [
+        ([], 2),
+        ([0], 3),
+        ([0, 1, -1, 2, -2], 0),
+        ([1, -1, 2, -2], 1),
+        ([1, -1, 2, -2], 2),
+        ([1, -1, 2, -2], 3),
+        ([1, -2, 3, -4], 4),
+        ([1, -2, 3, -4], 1),
+        ([1, -2, 3, -4], 2),
+        ([1, -2, 3, -4], 3),
+        ([1, -2, 3, -4], 4),
+    ],
+)
+def test_power_inplace(context, plain_vec, power):
+    ckks_vec = ts.ckks_vector(context, plain_vec)
+    expected = [np.power(v, power) for v in plain_vec]
+    ckks_vec **= power
+    decrypted_result = ckks_vec.decrypt()
+    assert _almost_equal(decrypted_result, expected, 1), "Decryption of vector is incorrect"
+
+
+@pytest.mark.parametrize(
+    "plain_vec",
+    [
+        [],
+        [0],
+        [1],
+        [2],
+        [1, -1, 2, -2],
+        [1, -1, 6, -2],
+        [2, 1, -2, -2],
+        [1, -2, 3, -4],
+        [3, -2, 5, -4],
+        [1, -4, 3, 5],
+    ],
+)
+def test_square(context, plain_vec):
+    ckks_vec = ts.ckks_vector(context, plain_vec)
+    expected = [np.power(v, 2) for v in plain_vec]
+    new_vec = ckks_vec.square()
+    decrypted_result = new_vec.decrypt()
+    assert _almost_equal(decrypted_result, expected, 1), "Decryption of vector is incorrect"
+    assert _almost_equal(ckks_vec.decrypt(), plain_vec, 1), "Something went wrong in memory."
+
+
+@pytest.mark.parametrize(
+    "plain_vec",
+    [
+        [],
+        [0],
+        [1],
+        [2],
+        [1, -1, 2, -2],
+        [1, -1, 6, -2],
+        [2, 1, -2, -2],
+        [1, -2, 3, -4],
+        [3, -2, 5, -4],
+        [1, -4, 3, 5],
+    ],
+)
+def test_square_inplace(context, plain_vec):
+    ckks_vec = ts.ckks_vector(context, plain_vec)
+    expected = [np.power(v, 2) for v in plain_vec]
+    ckks_vec.square_()
+    decrypted_result = ckks_vec.decrypt()
+    assert _almost_equal(decrypted_result, expected, 1), "Decryption of vector is incorrect"
+
+
 @pytest.mark.parametrize(
     "vec1, vec2",
     [