From 34563a7b37996736887c1e6c90b7f888a79dd0b2 Mon Sep 17 00:00:00 2001 From: peekxc Date: Tue, 12 Dec 2023 15:05:45 -0500 Subject: [PATCH] matrix fucntion changes + removing openmp from arm build --- .cirrus.yml | 6 ++--- meson.build | 6 ++--- pyproject.toml | 2 +- src/primate/_lanczos.cpp | 17 +++++++++---- src/primate/_vapprox.cpp | 35 -------------------------- src/primate/meson.build | 2 ++ src/primate/pylinop.h | 40 ++++++++++++++++++++++++++++++ src/primate/sparse.py | 53 ++++++++++++++++++++++++++++++++++------ src/primate/special.py | 4 +++ src/primate/trace.py | 4 +-- tests/test_vapprox.py | 10 ++++++++ 11 files changed, 121 insertions(+), 58 deletions(-) delete mode 100644 src/primate/_vapprox.cpp create mode 100644 src/primate/pylinop.h diff --git a/.cirrus.yml b/.cirrus.yml index 9773e09..91e5a5b 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -64,10 +64,10 @@ windows_task: windows_container: image: cirrusci/windowsservercore:2019 only_if: changesInclude('.cirrus.yml', '**.{h,cpp,py}') - env: - PATH: ${PATH}:"C:\python311" - setup_python_script: | + setup_script: | choco install -y python311 + env: + PATH: '%PATH%;C:\ProgramData\chocolatey\bin;C:\Python311' dependencies_script: | bash tools/cibw_windows.sh before_script: | diff --git a/meson.build b/meson.build index 7ce7fa4..ed5202f 100644 --- a/meson.build +++ b/meson.build @@ -14,9 +14,9 @@ project( ) OS_platform = host_machine.system() env = environment() +use_openmp = get_option('use_openmp') and OS_platform != 'windows' and host_machine.cpu() == 'x86_64' -use_openmp = get_option('use_openmp') -if use_openmp and OS_platform != 'windows' +if use_openmp add_global_arguments('-DOMP_MULTITHREADED=1', language : 'cpp') message('Compiling with OpenMP support') endif @@ -85,7 +85,7 @@ dependency_map = {} # subdir('blas') # Configure BLAS / LAPACK ## Include OpenMP (mandatory ; but exclude on windows because it's too difficult to link) -if get_option('use_openmp') and OS_platform != 'windows' +if use_openmp omp = dependency('openmp', required: false) openmp_flags = [] if omp.found() diff --git a/pyproject.toml b/pyproject.toml index 4762642..ab1171f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = ['meson-python', 'wheel', 'ninja', 'pybind11', 'numpy'] # 'pythran-op [project] name = "primate" -version = "0.2.6" +version = "0.2.7" readme = "README.md" classifiers = [ "Intended Audience :: Science/Research", diff --git a/src/primate/_lanczos.cpp b/src/primate/_lanczos.cpp index ef8645c..29765b9 100644 --- a/src/primate/_lanczos.cpp +++ b/src/primate/_lanczos.cpp @@ -169,17 +169,24 @@ void _lanczos_wrapper(py::module& m, const std::string suffix, WrapperFunc wrap return std::unique_ptr< MatrixFunction< F, WrapperType > >(new MatrixFunction(op, sf, deg, rtol, orth)); })) .def_property_readonly("shape", &MatrixFunction< F, WrapperType >::shape) + .def_property_readonly("dtype", [](const MatrixFunction< F, WrapperType >& M) -> py::dtype { + auto dtype = pybind11::dtype(pybind11::format_descriptor< F >::format()); + return dtype; + }) .def_readonly("deg", &MatrixFunction< F, WrapperType >::deg) .def_readwrite("rtol", &MatrixFunction< F, WrapperType >::rtol) .def_readwrite("orth", &MatrixFunction< F, WrapperType >::orth) - .def("matvec", [](const MatrixFunction< F, WrapperType >& m, const py_array< F >& x) -> py_array< F >{ + .def("matvec", [](const MatrixFunction< F, WrapperType >& M, const py_array< F >& x) -> py_array< F >{ using VectorF = Eigen::Matrix< F, Dynamic, 1 >; - auto output = static_cast< ArrayF >(VectorF::Zero(m.shape().first)); - m.matvec(x.data(), output.data()); + auto output = static_cast< ArrayF >(VectorF::Zero(M.shape().first)); + M.matvec(x.data(), output.data()); return py::cast(output); }) - .def("matvec", [](const MatrixFunction< F, WrapperType >& m, const py_array< F >& x, py_array< F >& y) -> void { - m.matvec(x.data(), y.mutable_data()); + .def("matvec", [](const MatrixFunction< F, WrapperType >& M, const py_array< F >& x, py_array< F >& y) -> void { + if (M.shape().second != x.size() || M.shape().first != y.size()){ + throw std::invalid_argument("Input/output dimension mismatch; vector inputs must match shape of the operator."); + } + M.matvec(x.data(), y.mutable_data()); }) // .def_method("__repr__", &MatrixFunction::eval) ; diff --git a/src/primate/_vapprox.cpp b/src/primate/_vapprox.cpp deleted file mode 100644 index 93a9924..0000000 --- a/src/primate/_vapprox.cpp +++ /dev/null @@ -1,35 +0,0 @@ -#include -#include -#include -#include -#include -#include "_random_generator/vector_generator.h" -#include "_lanczos/lanczos.h" -#include "eigen_operators.h" -#include // constants -#include -#include -#include - -namespace py = pybind11; - - -// template< std::floating_point F, LinearOperator Matrix > -// struct MatrixFunction { -// using VectorF = Eigen::Matrix< F, Dynamic, 1 >; -// using ArrayF = Eigen::Array< F, Dynamic, 1 >; -// using EigenSolver = Eigen::SelfAdjointEigenSolver< DenseMatrix< F > >; - -// // Parameters -// const Matrix& op; -// std::function< F(F) > f; -// const int deg; -// const F rtol; -// const int orth; - -// MatrixFunction(const Matrix& A, std::function< F(F) > fun, int lanczos_degree, F lanczos_rtol, int add_orth) - -// PYBIND11_MODULE(_vapprox, m) { -// _matrix_function_wrapper< float, Eigen::SparseMatrix< float > >(m); -// } - diff --git a/src/primate/meson.build b/src/primate/meson.build index 829396a..f142b14 100644 --- a/src/primate/meson.build +++ b/src/primate/meson.build @@ -23,7 +23,9 @@ python_sources = [ 'diagonalize.py', 'plotting.py', 'trace.py', + 'sparse.py', 'stats.py', + 'special.py', '__init__.py' ] diff --git a/src/primate/pylinop.h b/src/primate/pylinop.h new file mode 100644 index 0000000..7079725 --- /dev/null +++ b/src/primate/pylinop.h @@ -0,0 +1,40 @@ +#include +#include +namespace py = pybind11; + +template< typename F > +using py_array = py::array_t< F, py::array::f_style | py::array::forcecast >; + +template< std::floating_point F > +struct PyLinearOperator { + using value_type = F; + const py::object _op; + PyLinearOperator(const py::object& op) : _op(op) { + if (!py::hasattr(op, "matvec")) { throw std::invalid_argument("Supplied object is missing 'matvec' attribute."); } + if (!py::hasattr(op, "shape")) { throw std::invalid_argument("Supplied object is missing 'shape' attribute."); } + // if (!op.has_attr("dtype")) { throw std::invalid_argument("Supplied object is missing 'dtype' attribute."); } + } + + // Calls the matvec in python, casts the result to py::array_t, and copies through + auto matvec(const F* inp, F* out) const { + py_array< F > input({ static_cast(shape().second) }, inp); + py::object matvec_out = _op.attr("matvec")(input); + py::array_t< F > output = matvec_out.cast< py::array_t< F > >(); + std::copy(output.data(), output.data() + output.size(), out); + } + + auto matvec(const py_array< F >& input) const -> py_array< F > { + auto out = std::vector< F >(static_cast< size_t >(shape().first), 0); + this->matvec(input.data(), static_cast< F* >(&out[0])); + return py::cast(out); + } + + auto shape() const -> pair< size_t, size_t > { + return _op.attr("shape").cast< std::pair< size_t, size_t > >(); + } + + auto dtype() const -> py::dtype { + auto dtype = pybind11::dtype(pybind11::format_descriptor< F >::format()); + return dtype; + } +}; \ No newline at end of file diff --git a/src/primate/sparse.py b/src/primate/sparse.py index dc7b84f..ec9a49e 100644 --- a/src/primate/sparse.py +++ b/src/primate/sparse.py @@ -1,14 +1,18 @@ import numpy as np -from typing import Union +from typing import * from scipy.sparse.linalg import LinearOperator -import _lanczos +from scipy.sparse import issparse +from .special import _builtin_matrix_functions +import _lanczos def matrix_function( A: Union[LinearOperator, np.ndarray], fun: Union[str, Callable] = "identity", deg: int = 20, - orth: int = 0 + rtol: float = 1e-8, + orth: int = 0, + **kwargs ) -> LinearOperator: """Constructs an operator approximating the action v |-> f(A)v @@ -23,23 +27,56 @@ def matrix_function( real-valued function defined on the spectrum of `A`. deg : int, default = 20 Degree of the Krylov expansion. + rtol : float, default = 1e-8 + Relative tolerance to consider two Lanczos vectors are numerically orthogonal. orth: int, default = 0 Number of additional Lanczos vectors to orthogonalize against when building the Krylov basis. - + kwargs : dict, optional + additional key-values to parameterize the chosen function 'fun'. + + Returns: + -------- + operator : LinearOperator + Operator approximating the action of `fun` on the spectrum of `A` + """ + attr_checks = [hasattr(A, "__matmul__"), hasattr(A, "matmul"), hasattr(A, "dot"), hasattr(A, "matvec")] + assert any(attr_checks), "Invalid operator; must have an overloaded 'matvec' or 'matmul' method" + assert hasattr(A, "shape") and len(A.shape) >= 2, "Operator must be at least two dimensional." + assert A.shape[0] == A.shape[1], "This function only works with square, symmetric matrices!" + + ## Parameterize the type of matrix function if isinstance(A, np.ndarray): module_func = "MatrixFunction_dense" - elif issparray(A): + elif issparse(A): module_func = "MatrixFunction_sparse" elif isinstance(A, LinearOperator): module_func = "MatrixFunction_linop" else: raise ValueError(f"Invalid type '{type(A)}' supplied for operator A") - #_lanczos.MatrixFunction_sparse(A_sparse, deg, rtol, orth, **dict(function="log")) - #_lanczos.MatrixFunction_sparse(A_sparse, deg, rtol, orth, **dict(function="log")) + ## Get the dtype; infer it if it's not available + f_dtype = (A @ np.zeros(A.shape[1])).dtype if not hasattr(A, "dtype") else A.dtype + i_dtype = np.int32 + assert f_dtype.type == np.float32 or f_dtype.type == np.float64, "Only 32- or 64-bit floating point numbers are supported." + + ## Argument checking + lanczos_rtol = np.finfo(f_dtype).eps # if lanczos_rtol is None else f_dtype.type(lanczos_rtol) + orth = int(orth) # Number of additional vectors should be an integer + deg = max(deg, 2) # Should be at least two + + ## Parameterize the matrix function and trace call + if isinstance(fun, str): + assert fun in _builtin_matrix_functions, "If given as a string, matrix_function be one of the builtin functions." + kwargs["function"] = fun # _builtin_matrix_functions.index(matrix_function) + elif isinstance(fun, Callable): + kwargs["function"] = "generic" + kwargs["matrix_func"] = fun + else: + raise ValueError(f"Invalid matrix function type '{type(fun)}'") + ## Construct the instance - M = getattr(_lanczos, module_func)(A, deg, orth, **dict(function=fun)) + M = getattr(_lanczos, module_func)(A, deg, rtol, orth, **kwargs) return M diff --git a/src/primate/special.py b/src/primate/special.py index 333ffa3..4fa0e7a 100644 --- a/src/primate/special.py +++ b/src/primate/special.py @@ -1,4 +1,8 @@ from typing import * +import numpy as np + +## Natively support matrix functions +_builtin_matrix_functions = ["identity", "sqrt", "exp", "pow", "log", "numrank", "smoothstep", "gaussian"] def soft_sign(x: np.ndarray = None, q: int = 1): """Soft-sign function. diff --git a/src/primate/trace.py b/src/primate/trace.py index dda3f2f..c9809fd 100644 --- a/src/primate/trace.py +++ b/src/primate/trace.py @@ -5,11 +5,9 @@ ## Package imports from .random import _engine_prefixes, _engines +from .special import _builtin_matrix_functions import _lanczos -## Natively support matrix functions -_builtin_matrix_functions = ["identity", "sqrt", "exp", "pow", "log", "numrank", "smoothstep", "gaussian"] - def sl_trace ( A: Union[LinearOperator, np.ndarray], fun: Union[str, Callable] = "identity", diff --git a/tests/test_vapprox.py b/tests/test_vapprox.py index 08d69c2..acdd56b 100644 --- a/tests/test_vapprox.py +++ b/tests/test_vapprox.py @@ -54,5 +54,15 @@ def test_mf_approx(): assert np.all(y_log_test != y_test) def test_mf_api(): + from primate.sparse import matrix_function + np.random.seed(1234) + n = 10 + A_sparse = csc_array(symmetric(n, psd = True), dtype=np.float32) + M = matrix_function(A_sparse) + v0 = np.random.normal(size=n) + assert np.max(np.abs(M.matvec(v0) - A_sparse @ v0)) <= 1e-6 assert True + # from scipy.sparse.linalg import LinearOperator + # assert isinstance(M, LinearOperator) +