From 2e51a5b077d5ac16375437045bf58c8b8fa7191d Mon Sep 17 00:00:00 2001 From: LSchueler Date: Sat, 30 Dec 2023 17:37:08 +0100 Subject: [PATCH] [WIP] Add global var. to set # of OpenMP threads --- docs/source/index.rst | 2 ++ src/gstools/config.py | 2 ++ src/gstools/field/generator.py | 10 ++++++++-- src/gstools/field/summator.pyx | 8 +++++--- src/gstools/krige/krigesum.pyx | 10 ++++++---- src/gstools/variogram/estimator.pyx | 17 ++++++++++++----- src/gstools/variogram/variogram.py | 8 ++++++-- 7 files changed, 41 insertions(+), 16 deletions(-) diff --git a/docs/source/index.rst b/docs/source/index.rst index df583778..7435d4bf 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -98,6 +98,8 @@ For the development version, you can do almost the same: export GSTOOLS_BUILD_PARALLEL=1 pip install git+git://github.com/GeoStat-Framework/GSTools.git@main +The number of parallel threads can be set with the global variable `config.NUM_THREADS`. + **Using experimental GSTools-Core for even more speed** You can install the optional dependency `GSTools-Core `_, diff --git a/src/gstools/config.py b/src/gstools/config.py index 30ac4554..a8362949 100644 --- a/src/gstools/config.py +++ b/src/gstools/config.py @@ -4,6 +4,8 @@ .. currentmodule:: gstools.config """ +NUM_THREADS = 1 + # pylint: disable=W0611 try: # pragma: no cover import gstools_core diff --git a/src/gstools/field/generator.py b/src/gstools/field/generator.py index 7780ec72..571c9733 100644 --- a/src/gstools/field/generator.py +++ b/src/gstools/field/generator.py @@ -209,7 +209,9 @@ def __call__(self, pos, add_nugget=True): the random modes """ pos = np.asarray(pos, dtype=np.double) - summed_modes = summate(self._cov_sample, self._z_1, self._z_2, pos) + summed_modes = summate( + self._cov_sample, self._z_1, self._z_2, pos, config.NUM_THREADS + ) nugget = self.get_nugget(summed_modes.shape) if add_nugget else 0.0 return np.sqrt(self.model.var / self._mode_no) * summed_modes + nugget @@ -489,7 +491,11 @@ def __call__(self, pos, add_nugget=True): """ pos = np.asarray(pos, dtype=np.double) summed_modes = summate_incompr( - self._cov_sample, self._z_1, self._z_2, pos + self._cov_sample, + self._z_1, + self._z_2, + pos, + config.NUM_THREADS, ) nugget = self.get_nugget(summed_modes.shape) if add_nugget else 0.0 e1 = self._create_unit_vector(summed_modes.shape) diff --git a/src/gstools/field/summator.pyx b/src/gstools/field/summator.pyx index 29239c07..a37a90a5 100644 --- a/src/gstools/field/summator.pyx +++ b/src/gstools/field/summator.pyx @@ -14,7 +14,8 @@ def summate( const double[:, :] cov_samples, const double[:] z_1, const double[:] z_2, - const double[:, :] pos + const double[:, :] pos, + const int num_threads=1, ): cdef int i, j, d cdef double phase @@ -25,7 +26,7 @@ def summate( cdef double[:] summed_modes = np.zeros(X_len, dtype=float) - for i in prange(X_len, nogil=True): + for i in prange(X_len, nogil=True, num_threads=num_threads): for j in range(N): phase = 0. for d in range(dim): @@ -49,7 +50,8 @@ def summate_incompr( const double[:, :] cov_samples, const double[:] z_1, const double[:] z_2, - const double[:, :] pos + const double[:, :] pos, + const int num_threads=1, ): cdef int i, j, d cdef double phase diff --git a/src/gstools/krige/krigesum.pyx b/src/gstools/krige/krigesum.pyx index 2f79d3ad..1dba1399 100644 --- a/src/gstools/krige/krigesum.pyx +++ b/src/gstools/krige/krigesum.pyx @@ -12,7 +12,8 @@ cimport numpy as np def calc_field_krige_and_variance( const double[:, :] krig_mat, const double[:, :] krig_vecs, - const double[:] cond + const double[:] cond, + const int num_threads=1, ): cdef int mat_i = krig_mat.shape[0] @@ -26,7 +27,7 @@ def calc_field_krige_and_variance( # error = krig_vecs * krig_mat * krig_vecs # field = cond * krig_mat * krig_vecs - for k in prange(res_i, nogil=True): + for k in prange(res_i, nogil=True, num_threads=num_threads): for i in range(mat_i): krig_fac = 0.0 for j in range(mat_i): @@ -40,7 +41,8 @@ def calc_field_krige_and_variance( def calc_field_krige( const double[:, :] krig_mat, const double[:, :] krig_vecs, - const double[:] cond + const double[:] cond, + const int num_threads=1, ): cdef int mat_i = krig_mat.shape[0] @@ -52,7 +54,7 @@ def calc_field_krige( cdef int i, j, k # field = cond * krig_mat * krig_vecs - for k in prange(res_i, nogil=True): + for k in prange(res_i, nogil=True, num_threads=num_threads): for i in range(mat_i): krig_fac = 0.0 for j in range(mat_i): diff --git a/src/gstools/variogram/estimator.pyx b/src/gstools/variogram/estimator.pyx index 611f5efb..ae464bce 100644 --- a/src/gstools/variogram/estimator.pyx +++ b/src/gstools/variogram/estimator.pyx @@ -181,6 +181,7 @@ def directional( const double bandwidth=-1.0, # negative values to turn of bandwidth search const bint separate_dirs=False, # whether the direction bands don't overlap str estimator_type='m', + const int num_threads=1, ): if pos.shape[1] != f.shape[1]: raise ValueError(f'len(pos) = {pos.shape[1]} != len(f) = {f.shape[1])}') @@ -208,7 +209,7 @@ def directional( cdef int i, j, k, m, d cdef double dist - for i in prange(i_max, nogil=True): + for i in prange(i_max, nogil=True, num_threads=num_threads): for j in range(j_max): for k in range(j+1, k_max): dist = dist_euclid(dim, pos, j, k) @@ -239,6 +240,7 @@ def unstructured( const double[:, :] pos, str estimator_type='m', str distance_type='e', + const int num_threads=1, ): cdef int dim = pos.shape[0] cdef _dist_func distance @@ -271,7 +273,7 @@ def unstructured( cdef int i, j, k, m cdef double dist - for i in prange(i_max, nogil=True): + for i in prange(i_max, nogil=True, num_threads=num_threads): for j in range(j_max): for k in range(j+1, k_max): dist = distance(dim, pos, j, k) @@ -287,7 +289,11 @@ def unstructured( return np.asarray(variogram), np.asarray(counts) -def structured(const double[:, :] f, str estimator_type='m'): +def structured( + const double[:, :] f, + str estimator_type='m', + const int num_threads=1, +): cdef _estimator_func estimator_func = choose_estimator_func(estimator_type) cdef _normalization_func normalization_func = ( choose_estimator_normalization(estimator_type) @@ -301,7 +307,7 @@ def structured(const double[:, :] f, str estimator_type='m'): cdef long[:] counts = np.zeros(k_max, dtype=long) cdef int i, j, k - with nogil, parallel(): + with nogil, parallel(num_threads=num_threads): for i in range(i_max): for j in range(j_max): for k in prange(1, k_max-i): @@ -316,6 +322,7 @@ def ma_structured( const double[:, :] f, const bint[:, :] mask, str estimator_type='m', + const int num_threads=1, ): cdef _estimator_func estimator_func = choose_estimator_func(estimator_type) cdef _normalization_func normalization_func = ( @@ -330,7 +337,7 @@ def ma_structured( cdef long[:] counts = np.zeros(k_max, dtype=long) cdef int i, j, k - with nogil, parallel(): + with nogil, parallel(num_threads=num_threads): for i in range(i_max): for j in range(j_max): for k in prange(1, k_max-i): diff --git a/src/gstools/variogram/variogram.py b/src/gstools/variogram/variogram.py index 69746658..5cdd78f0 100644 --- a/src/gstools/variogram/variogram.py +++ b/src/gstools/variogram/variogram.py @@ -374,6 +374,7 @@ def vario_estimate( pos, estimator_type=cython_estimator, distance_type=distance_type, + num_threads=config.NUM_THREADS, ) else: estimates, counts = directional( @@ -385,6 +386,7 @@ def vario_estimate( bandwidth, separate_dirs=_separate_dirs_test(direction, angles_tol), estimator_type=cython_estimator, + num_threads=config.num_threads, ) if dir_no == 1: estimates, counts = estimates[0], counts[0] @@ -485,8 +487,10 @@ def vario_estimate_axis( cython_estimator = _set_estimator(estimator) if masked: - return ma_structured(field, mask, cython_estimator) - return structured(field, cython_estimator) + return ma_structured( + field, mask, cython_estimator, num_threads=config.NUM_THREADS + ) + return structured(field, cython_estimator, num_threads=config.NUM_THREADS) # for backward compatibility