diff options
| -rw-r--r-- | pyhegp/pyhegp.py | 15 | ||||
| -rw-r--r-- | tests/test_pyhegp.py | 10 |
2 files changed, 10 insertions, 15 deletions
diff --git a/pyhegp/pyhegp.py b/pyhegp/pyhegp.py index a35cfcb..3229724 100644 --- a/pyhegp/pyhegp.py +++ b/pyhegp/pyhegp.py @@ -39,12 +39,11 @@ def unstandardize(matrix, mean, standard_deviation): return ((matrix @ np.diag(standard_deviation)) + np.tile(mean, (m, 1))) -def hegp_encrypt(plaintext, mean, standard_deviation, key): - return key @ standardize(plaintext, mean, standard_deviation) +def hegp_encrypt(plaintext, key): + return key @ plaintext -def hegp_decrypt(ciphertext, mean, standard_deviation, key): - return unstandardize(np.transpose(key) @ ciphertext, - mean, standard_deviation) +def hegp_decrypt(ciphertext, key): + return np.transpose(key) @ ciphertext def pool_stats(list_of_stats): sums = [stats.n*stats.mean for stats in list_of_stats] @@ -103,8 +102,10 @@ def encrypt(genotype_file, summary_file, key_file, ciphertext_file): summary = read_summary(summary_file) rng = np.random.default_rng() key = random_key(rng, len(genotype)) - encrypted_genotype = hegp_encrypt(genotype, summary.mean, - summary.std, key) + encrypted_genotype = hegp_encrypt(standardize(genotype, + summary.mean, + summary.std), + key) if key_file: np.savetxt(key_file, key, delimiter=",", fmt="%f") np.savetxt(ciphertext_file, encrypted_genotype, delimiter=",", fmt="%f") diff --git a/tests/test_pyhegp.py b/tests/test_pyhegp.py index 2d2a258..c494d30 100644 --- a/tests/test_pyhegp.py +++ b/tests/test_pyhegp.py @@ -53,17 +53,11 @@ def no_column_zero_standard_deviation(matrix): elements=st.integers(min_value=0, max_value=100)), arrays("float64", array_shapes(min_dims=2, max_dims=2, min_side=2), - elements=st.floats(min_value=0, max_value=100))) - # Reject matrices with zero standard deviation columns since - # they trigger a division by zero. - .filter(no_column_zero_standard_deviation)) + elements=st.floats(min_value=0, max_value=100)))) def test_hegp_encryption_decryption_are_inverses(plaintext): - mean = np.mean(plaintext, axis=0) - standard_deviation = np.std(plaintext, axis=0) rng = np.random.default_rng() key = random_key(rng, len(plaintext)) - assert hegp_decrypt(hegp_encrypt(plaintext, mean, standard_deviation, key), - mean, standard_deviation, key) == approx(plaintext) + assert hegp_decrypt(hegp_encrypt(plaintext, key), key) == approx(plaintext) @given(arrays("float64", array_shapes(min_dims=2, max_dims=2), |