diff options
| -rw-r--r-- | pyhegp/linalg.py | 56 | ||||
| -rw-r--r-- | tests/test_linalg.py | 65 |
2 files changed, 121 insertions, 0 deletions
diff --git a/pyhegp/linalg.py b/pyhegp/linalg.py new file mode 100644 index 0000000..455e2d7 --- /dev/null +++ b/pyhegp/linalg.py @@ -0,0 +1,56 @@ +### pyhegp --- Homomorphic encryption of genotypes and phenotypes +### Copyright © 2026 Arun Isaac <arunisaac@systemreboot.net> +### +### This file is part of pyhegp. +### +### pyhegp is free software: you can redistribute it and/or modify it +### under the terms of the GNU General Public License as published by +### the Free Software Foundation, either version 3 of the License, or +### (at your option) any later version. +### +### pyhegp is distributed in the hope that it will be useful, but +### WITHOUT ANY WARRANTY; without even the implied warranty of +### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +### General Public License for more details. +### +### You should have received a copy of the GNU General Public License +### along with pyhegp. If not, see <https://www.gnu.org/licenses/>. + +import numpy as np + +from itertools import accumulate, pairwise +from scipy.linalg import block_diag + +class BlockDiagonalMatrix: + def __init__(self, _blocks): + self.blocks = _blocks + self.shape = (sum(len(block) for block in self.blocks),) * 2 + + def __len__(self): + return self.shape[0] + + def __repr__(self): + return f"BlockDiagonalMatrix{self.blocks}" + + def __array_function__(self, func, types, args, kwargs): + if ((func is np.transpose) + and (all(issubclass(type, BlockDiagonalMatrix) + for type in types))): + return BlockDiagonalMatrix([np.transpose(block) + for block in self.blocks]) + else: + return NotImplemented + + def __array__(self): + return block_diag(*self.blocks) + + def __matmul__(self, multiplier): + return np.concatenate( + [block @ multiplier[start:stop, ...] + for (start, stop), block + in zip(pairwise(accumulate((len(block) for block in self.blocks), + initial=0)), + self.blocks)]) + + def savetxt(self, file, *args, **kwargs): + return np.savetxt(file, self.to_ndarray(), *args, **kwargs) diff --git a/tests/test_linalg.py b/tests/test_linalg.py new file mode 100644 index 0000000..f826a0a --- /dev/null +++ b/tests/test_linalg.py @@ -0,0 +1,65 @@ +### pyhegp --- Homomorphic encryption of genotypes and phenotypes +### Copyright © 2026 Arun Isaac <arunisaac@systemreboot.net> +### +### This file is part of pyhegp. +### +### pyhegp is free software: you can redistribute it and/or modify it +### under the terms of the GNU General Public License as published by +### the Free Software Foundation, either version 3 of the License, or +### (at your option) any later version. +### +### pyhegp is distributed in the hope that it will be useful, but +### WITHOUT ANY WARRANTY; without even the implied warranty of +### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +### General Public License for more details. +### +### You should have received a copy of the GNU General Public License +### along with pyhegp. If not, see <https://www.gnu.org/licenses/>. + +from hypothesis import given, strategies as st +from hypothesis.extra.numpy import arrays +import numpy as np +from pytest import approx + +from pyhegp.linalg import BlockDiagonalMatrix + +@st.composite +def block_diagonal_matrices(draw, max_block_size=10, max_number_of_blocks=None): + return BlockDiagonalMatrix( + [draw(arrays("float64", (n, n), + elements=st.floats(min_value=-10, + max_value=10, + allow_nan=False, + allow_infinity=False))) + for n in draw(st.lists(st.integers(min_value=1, + max_value=max_block_size), + min_size=1, + max_size=max_number_of_blocks))]) + +@given(block_diagonal_matrices(max_number_of_blocks=10)) +def test_block_diagonal_matrix_transpose(block_diagonal_matrix): + assert (np.transpose(block_diagonal_matrix).__array__() + == approx(np.transpose(block_diagonal_matrix.__array__()))) + +@st.composite +def block_diagonal_matrix_product_multiplicands(draw): + block_diagonal_matrix = draw(block_diagonal_matrices()) + return (block_diagonal_matrix, + draw(arrays("float64", + # Either a vector or a matrix + (len(block_diagonal_matrix), + *draw(st.one_of( + st.just(()), + st.builds(lambda x: (x,), + st.integers(min_value=1, + max_value=100))))), + elements=st.floats(min_value=-10, + max_value=10, + allow_nan=False, + allow_infinity=False)))) + +@given(block_diagonal_matrix_product_multiplicands()) +def test_block_diagonal_matrix_product(multiplicands): + block_diagonal_matrix, multiplier = multiplicands + assert ((block_diagonal_matrix @ multiplier) + == approx(block_diagonal_matrix.__array__() @ multiplier)) |