about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--pyhegp/pyhegp.py16
-rw-r--r--tests/test_pyhegp.py11
2 files changed, 19 insertions, 8 deletions
diff --git a/pyhegp/pyhegp.py b/pyhegp/pyhegp.py
index 2a82690..676e0b6 100644
--- a/pyhegp/pyhegp.py
+++ b/pyhegp/pyhegp.py
@@ -33,15 +33,19 @@ Stats = namedtuple("Stats", "n mean std")
 def random_key(rng, n):
     return special_ortho_group.rvs(n, random_state=rng)
 
-def standardize(matrix, mean, standard_deviation):
+def center(matrix, mean):
     m, _ = matrix.shape
-    return ((matrix - np.tile(mean, (m, 1)))
-            @ np.diag(1 / standard_deviation))
+    return matrix - np.tile(mean, (m, 1))
+
+def uncenter(matrix, mean):
+    return center(matrix, -mean)
+
+def standardize(matrix, mean, standard_deviation):
+    return center(matrix, mean) @ np.diag(1 / standard_deviation)
 
 def unstandardize(matrix, mean, standard_deviation):
-    m, _ = matrix.shape
-    return ((matrix @ np.diag(standard_deviation))
-            + np.tile(mean, (m, 1)))
+    return uncenter(matrix @ np.diag(standard_deviation),
+                    mean)
 
 def hegp_encrypt(plaintext, key):
     return key @ plaintext
diff --git a/tests/test_pyhegp.py b/tests/test_pyhegp.py
index cdf3a7f..c3cf47f 100644
--- a/tests/test_pyhegp.py
+++ b/tests/test_pyhegp.py
@@ -1,5 +1,5 @@
 ### pyhegp --- Homomorphic encryption of genotypes and phenotypes
-### Copyright © 2025 Arun Isaac <arunisaac@systemreboot.net>
+### Copyright © 2025–2026 Arun Isaac <arunisaac@systemreboot.net>
 ###
 ### This file is part of pyhegp.
 ###
@@ -29,7 +29,7 @@ import pandas as pd
 import pytest
 from pytest import approx
 
-from pyhegp.pyhegp import Stats, main, hegp_encrypt, hegp_decrypt, random_key, pool_stats, standardize, unstandardize, genotype_summary, encrypt_genotype, encrypt_phenotype, cat_genotype, cat_phenotype
+from pyhegp.pyhegp import Stats, main, hegp_encrypt, hegp_decrypt, random_key, pool_stats, center, uncenter, standardize, unstandardize, genotype_summary, encrypt_genotype, encrypt_phenotype, cat_genotype, cat_phenotype
 from pyhegp.serialization import Summary, read_summary, read_genotype, is_genotype_metadata_column, is_phenotype_metadata_column
 from pyhegp.utils import negate
 
@@ -95,6 +95,13 @@ def test_hegp_encryption_decryption_are_inverses(plaintext):
 
 @given(arrays("float64",
               array_shapes(min_dims=2, max_dims=2),
+              elements=st.floats(min_value=0, max_value=100)))
+def test_center_uncenter_are_inverses(matrix):
+    mean = np.mean(matrix, axis=0)
+    assert uncenter(center(matrix, mean), mean) == approx(matrix)
+
+@given(arrays("float64",
+              array_shapes(min_dims=2, max_dims=2),
               elements=st.floats(min_value=0, max_value=100))
        # Reject matrices with zero standard deviation columns since
        # they trigger a division by zero.