code: #
  #  Copyright 2019 The FATE Authors. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #  See the License for the specific language governing permissions and
  #  limitations under the License.
  #
  
  import unittest
  import functools
  import math
  from fate_arch.session import computing_session as session
  from federatedml.ensemble.basic_algorithms.decision_tree.tree_core.g_h_optim import PackedGHCompressor, GHPacker, fix_point_precision
  from federatedml.secureprotol.encrypt import PaillierEncrypt, FakeEncrypt
  from federatedml.ensemble.basic_algorithms.decision_tree.tree_core.splitter import SplitInfo
  from federatedml.secureprotol.encrypt_mode import EncryptModeCalculator
  from federatedml.util import consts
  import numpy as np
  
  
  np.random.seed(114514)
  
  
  def generate_bin_gh(num):
      # (-1, 1)
      g = np.random.random(num)
      h = np.random.random(num)
      g = g*2 - 1
      return g, h
  
  
  def generate_reg_gh(num, lower, upper):
      g = np.random.random(num)
      h = np.zeros(num) + 2
      g = g * (upper - lower) + lower
      return g, h
  
  
  def cmp(a, b):
      if a[0] > b[0]:
          return 1
      else:
          return -1
  
  
  def en_gh_list(g, h, en):
      en_g = [en.encrypt(i) for i in g]
      en_h = [en.encrypt(i) for i in h]
      return en_g, en_h
  
  
  def truncate(f, n=consts.TREE_DECIMAL_ROUND):
      return math.floor(f * 10 ** n) / 10 ** n
  
  
  def make_random_sum(collected_gh, g, h, en_g_l, en_h_l, max_sample_num):
  
      selected_sample_num = np.random.randint(max_sample_num) + 1# at least 1 sample
  
      idx = np.random.random(selected_sample_num)
      idx = np.unique((idx * max_sample_num).astype(int))
      print('randomly select {} samples'.format(len(idx)))
      selected_g = g[idx]
      selected_h = h[idx]
      g_sum = selected_g.sum()
      h_sum = selected_h.sum()
      g_h_list = sorted(collected_gh, key=functools.cmp_to_key(cmp))
      sum_gh = 0
      en_g_sum = 0
      en_h_sum = 0
      for i in idx:
          gh = g_h_list[i][1][0]
  
          sum_gh += gh
          en_g_sum += en_g_l[i]
          en_h_sum += en_h_l[i]
  
      return g_sum, h_sum, sum_gh, en_g_sum, en_h_sum, len(idx)
  
  
  class TestFeatureHistogram(unittest.TestCase):
  
      @staticmethod
      def prepare_testing_data(g, h, en, max_sample_num, sample_id, task_type, g_min=None, g_max=None):
  
          calculator = EncryptModeCalculator(encrypter=en)
          packer = GHPacker(max_sample_num, en_calculator=calculator, sync_para=False, task_type=task_type,
                            g_min=g_min, g_max=g_max)
          en_g_l, en_h_l = en_gh_list(g, h, en)
          data_list = [(id_, (g_, h_)) for id_, g_, h_ in zip(sample_id, g, h)]
          data_table = session.parallelize(data_list, 4, include_key=True)
          en_table = packer.pack_and_encrypt(data_table)
          collected_gh = list(en_table.collect())
  
          return packer, en_g_l, en_h_l, en_table, collected_gh
  
      @classmethod
      def setUpClass(cls):
  
          session.init("test_gh_packing")
  
          cls.max_sample_num = 1000
          cls.test_num = 10
          cls.split_info_test_num = 200
  
          key_length = 1024
          sample_id = [i for i in range(cls.max_sample_num)]
  
          # classification data
          cls.g, cls.h = generate_bin_gh(cls.max_sample_num)
  
          cls.p_en = PaillierEncrypt()
          cls.p_en.generate_key(key_length)
  
          cls.p_packer, cls.p_en_g_l, cls.p_en_h_l, cls.p_en_table, cls.p_collected_gh = \
              cls.prepare_testing_data(cls.g, cls.h, cls.p_en, cls.max_sample_num, sample_id, consts.CLASSIFICATION)
          cls.compressor = PackedGHCompressor(sync_para=False)
          cls.compressor.compressor._padding_length, cls.compressor.compressor._capacity = \
              cls.p_packer.packer.cipher_compress_suggest()
          print('paillier compress para {}'.format(cls.p_packer.packer.cipher_compress_suggest()))
  
          # regression data
          cls.g_reg, cls.h_reg = generate_reg_gh(cls.max_sample_num, -1000, 1000)
          cls.reg_p_packer, cls.reg_p_en_g_l, cls.reg_p_en_h_l, cls.reg_p_en_table, cls.reg_p_collected_gh = \
              cls.prepare_testing_data(cls.g_reg, cls.h_reg, cls.p_en, cls.max_sample_num, sample_id, consts.REGRESSION,
                                       g_min=-1000, g_max=1000)
          cls.reg_compressor = PackedGHCompressor(sync_para=False)
          cls.reg_compressor.compressor._padding_length, cls.reg_compressor.compressor._capacity = \
              cls.reg_p_packer.packer.cipher_compress_suggest()
          print('paillier compress para {}'.format(cls.p_packer.packer.cipher_compress_suggest()))
  
          print('initialization done')
  
      def run_gh_accumulate_test(self, test_num, collected_gh, en_g_l, en_h_l, packer, en, g, h, check=True):
  
          print('{} test to run'.format(test_num))
          for i in range(test_num):
              print('executing test {}'.format(i))
              g_sum, h_sum, en_sum, en_g_sum, en_h_sum, sample_num = make_random_sum(collected_gh, g, h,
                                                                                     en_g_l,
                                                                                     en_h_l,
                                                                                     self.max_sample_num)
              de_num = en.raw_decrypt(en_sum)
              unpack_num = packer.packer._unpack_an_int(de_num, packer.packer._bit_assignment[0])
  
              g_sum_ = unpack_num[0] / fix_point_precision - sample_num * packer.g_offset
              h_sum_ = unpack_num[1] / fix_point_precision
  
              g_sum_2 = en.decrypt(en_g_sum)
              h_sum_2 = en.decrypt(en_h_sum)
  
              print(g_sum, h_sum)
              print(g_sum_2, h_sum_2)
              print(g_sum_, h_sum_)
  
              g_sum, h_sum = truncate(g_sum), truncate(h_sum)
              g_sum_, h_sum_ = truncate(g_sum_), truncate(h_sum_)
              g_sum_2, h_sum_2 = truncate(g_sum_2), truncate(h_sum_2)
  
              print(g_sum, h_sum)
              print(g_sum_2, h_sum_2)
              print(g_sum_, h_sum_)
  
              if check:
                  # make sure packing result close to plaintext sum
                  self.assertTrue(g_sum_ == g_sum)
                  self.assertTrue(h_sum_ == h_sum)
  
              print('passed')
  
      def test_pack_gh_accumulate(self):
  
          # test the correctness of gh packing(in comparision to plaintext)
  
          # Paillier
          self.run_gh_accumulate_test(self.test_num, self.p_collected_gh, self.p_en_g_l, self.p_en_h_l, self.p_packer,
                                      self.p_en, self.g, self.h)
  
          print('*'*30)
          print('test paillier done')
          print('*'*30)
  
      def test_split_info_cipher_compress(self):
  
          # test the correctness of cipher compressing
          print('testing binary')
          collected_gh = self.p_collected_gh
          en_g_l = self.p_en_g_l
          en_h_l = self.p_en_h_l
          packer = self.p_packer
          en = self.p_en
  
          sp_list = []
          g_sum_list, h_sum_list = [], []
          pack_en_list = []
  
          for i in range(self.split_info_test_num):
              g_sum, h_sum, en_sum, en_g_sum, en_h_sum, sample_num = make_random_sum(collected_gh, self.g, self.h,
                                                                                     en_g_l,
                                                                                     en_h_l,
                                                                                     self.max_sample_num)
              sp = SplitInfo(sum_grad=en_sum, sum_hess=0, sample_count=sample_num)
              sp_list.append(sp)
              g_sum_list.append(g_sum)
              h_sum_list.append(h_sum)
              pack_en_list.append(en_sum)
  
          print('generating split-info done')
          packages = self.compressor.compress_split_info(sp_list[:-1], sp_list[-1])
          print('package length is {}'.format(len(packages)))
          unpack_rs = packer.decompress_and_unpack(packages)
          case_id = 0
          for s, g, h, en_gh in zip(unpack_rs, g_sum_list, h_sum_list, pack_en_list):
              print('*'*10)
              print(case_id)
              case_id += 1
              de_num = en.raw_decrypt(en_gh)
              unpack_num = packer.packer._unpack_an_int(de_num, packer.packer._bit_assignment[0])
              g_sum_ = unpack_num[0] / fix_point_precision - s.sample_count * packer.g_offset
              h_sum_ = unpack_num[1] / fix_point_precision
  
              print(s.sample_count)
              print(s.sum_grad, g_sum_, g)
              print(s.sum_hess, h_sum_, h)
  
              # make sure cipher compress is correct
              self.assertTrue(truncate(s.sum_grad) == truncate(g_sum_))
              self.assertTrue(truncate(s.sum_hess) == truncate(h_sum_))
          print('check passed')
  
      def test_regression_cipher_compress(self):
  
          # test the correctness of cipher compressing
          print('testing regression')
          collected_gh = self.reg_p_collected_gh
          en_g_l = self.reg_p_en_g_l
          en_h_l = self.reg_p_en_h_l
          packer = self.reg_p_packer
          en = self.p_en
  
          sp_list = []
          g_sum_list, h_sum_list = [], []
          pack_en_list = []
  
          for i in range(self.split_info_test_num):
              g_sum, h_sum, en_sum, en_g_sum, en_h_sum, sample_num = make_random_sum(collected_gh, self.g_reg, self.h_reg,
                                                                                     en_g_l,
                                                                                     en_h_l,
                                                                                     self.max_sample_num)
              sp = SplitInfo(sum_grad=en_sum, sum_hess=0, sample_count=sample_num)
              sp_list.append(sp)
              g_sum_list.append(g_sum)
              h_sum_list.append(h_sum)
              pack_en_list.append(en_sum)
  
          print('generating split-info done')
          packages = self.reg_compressor.compress_split_info(sp_list[:-1], sp_list[-1])
          print('package length is {}'.format(len(packages)))
          unpack_rs = packer.decompress_and_unpack(packages)
          case_id = 0
          for s, g, h, en_gh in zip(unpack_rs, g_sum_list, h_sum_list, pack_en_list):
              print('*' * 10)
              print(case_id)
              case_id += 1
              de_num = en.raw_decrypt(en_gh)            # make sure packing result close to plaintext sum
              unpack_num = packer.packer._unpack_an_int(de_num, packer.packer._bit_assignment[0])
              g_sum_ = unpack_num[0] / fix_point_precision - s.sample_count * packer.g_offset
              h_sum_ = unpack_num[1] / fix_point_precision
  
              print(s.sample_count)
              print(s.sum_grad, g_sum_, g)
              print(s.sum_hess, h_sum_, h)
  
              # make sure cipher compress is correct
              self.assertTrue(truncate(s.sum_grad) == truncate(g_sum_))
              self.assertTrue(truncate(s.sum_hess) == truncate(h_sum_))
          print('check passed')
  
      def test_regression_gh_packing(self):
  
          # Paillier
          self.run_gh_accumulate_test(self.test_num, self.reg_p_collected_gh, self.reg_p_en_g_l, self.reg_p_en_h_l, self.reg_p_packer,
                                      self.p_en, self.g_reg, self.h_reg, check=False)  # float error in regression is not controllable
  
      @classmethod
      def tearDownClass(self):
          session.stop()
  
  
  if __name__ == '__main__':
  
      unittest.main()
  
  
apis: numpy.random.random numpy.random.randint numpy.random.seed numpy.zeros