From 34ab5b1994c6404f1995900e684f657438613919 Mon Sep 17 00:00:00 2001
From: "Matwey V. Kornilov" <matwey.kornilov@gmail.com>
Date: Wed, 29 Jan 2025 08:49:04 +0100
Subject: [PATCH] Use use_threads_if

---
 pyproject.toml                    | 2 +-
 src/coniferest/calc_paths_sum.pyx | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 83f03a4..7f97cfd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ requires = [
     "setuptools",
     "wheel",
     "numpy>=2.0",
-    "Cython",
+    "Cython==3.1.0a1",
 ]
 build-backend = "setuptools.build_meta"
 
diff --git a/src/coniferest/calc_paths_sum.pyx b/src/coniferest/calc_paths_sum.pyx
index a173335..a067cee 100644
--- a/src/coniferest/calc_paths_sum.pyx
+++ b/src/coniferest/calc_paths_sum.pyx
@@ -1,8 +1,11 @@
+# cython: profile=True
+
 import numpy as np
 from cython.parallel cimport prange, parallel
 
 cimport numpy as np
 cimport cython
+cimport openmp
 
 
 def calc_paths_sum(selector_t [::1] selectors,
@@ -21,6 +24,9 @@ def calc_paths_sum(selector_t [::1] selectors,
     if indices[-1] > sellen:
         raise ValueError('indices are out of range of the selectors')
 
+    if num_threads < 0:
+        num_threads = openmp.omp_get_max_threads()
+
     _paths_sum(selectors, indices, data, paths_view, weights, num_threads, chunksize)
     return paths
 
@@ -101,8 +107,9 @@ cdef void _paths_sum(selector_t [::1] selectors,
     cdef selector_t selector
     cdef Py_ssize_t tree_offset
     cdef np.int32_t feature, i
+    cdef int use_threads_if = (2 * num_threads < data.shape[0])
 
-    with nogil, parallel(num_threads=num_threads):
+    with nogil, parallel(num_threads=num_threads, use_threads_if=use_threads_if):
         trees = indices.shape[0] - 1
 
         for x_index in prange(data.shape[0], schedule='static', chunksize=chunksize):