simplex-tree-classifier 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Yasmin
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,157 @@
1
+ Metadata-Version: 2.4
2
+ Name: simplex-tree-classifier
3
+ Version: 0.1.0
4
+ Summary: Hierarchical simplex-tree classifier with barycentric embedding, GPU-accelerated transform, and non-convex boundary pruning.
5
+ Author: Yasmin
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/yasmin/SimplexTreeClassifier
8
+ Project-URL: Repository, https://github.com/yasmin/SimplexTreeClassifier
9
+ Keywords: machine-learning,classifier,simplex,barycentric,interpretability,surrogate-model
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Operating System :: OS Independent
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
14
+ Requires-Python: >=3.9
15
+ Description-Content-Type: text/markdown
16
+ License-File: LICENSE
17
+ Requires-Dist: numpy>=1.21
18
+ Requires-Dist: scipy>=1.7
19
+ Requires-Dist: scikit-learn>=1.0
20
+ Requires-Dist: torch>=1.12
21
+ Provides-Extra: viz
22
+ Requires-Dist: matplotlib>=3.4; extra == "viz"
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest>=7.0; extra == "dev"
25
+ Dynamic: license-file
26
+
27
+ # simplex-tree-classifier
28
+
29
+ A hierarchical **simplex-tree classifier**. It embeds points into the barycentric
30
+ coordinates of a subdivided simplex, then trains any scikit-learn estimator on
31
+ that sparse, geometry-aware feature space. The barycentric `transform` is
32
+ accelerated with **PyTorch** and runs on the GPU when one is available.
33
+
34
+ It works in two ways:
35
+
36
+ - **Dataset mode** - train directly on your data `(X, y)`.
37
+ - **Surrogate mode** - imitate an existing `model` (e.g. a neural network) to
38
+ produce an interpretable, piecewise-linear geometric surrogate.
39
+
40
+ Everything is **N-dimensional**.
41
+
42
+ ## Installation
43
+
44
+ ```bash
45
+ pip install simplex-tree-classifier # core (numpy, scipy, scikit-learn, torch)
46
+ pip install simplex-tree-classifier[viz] # + matplotlib for 2-D plots
47
+ ```
48
+
49
+ From source:
50
+
51
+ ```bash
52
+ cd simplex_tree_package
53
+ pip install -e .
54
+ ```
55
+
56
+
57
+
58
+ ## Quickstart
59
+
60
+
61
+
62
+ ### Dataset mode
63
+
64
+ ```python
65
+ import numpy as np
66
+ from simplex_tree_classifier import SimplexTreeClassifier
67
+
68
+ X = np.random.uniform(0, 1, size=(500, 4))
69
+ y = (X.sum(axis=1) > 2.0).astype(int)
70
+
71
+ clf = SimplexTreeClassifier(subdivision_levels=2) # default classifier: LinearSVC
72
+ clf.fit(X, y)
73
+ print(clf.predict(X[:5]))
74
+ ```
75
+
76
+
77
+
78
+ ### Surrogate mode (imitate a model)
79
+
80
+ Pass any object with `.predict(X)` or any callable `X -> labels`, plus the input
81
+ dimensionality. Fill points are sampled in `[0, 1]^d` and labeled by the model.
82
+
83
+ ```python
84
+ def black_box(X): # e.g. a wrapped PyTorch model
85
+ X = np.atleast_2d(X)
86
+ return (X[:, 0] + X[:, 1] > 1.0).astype(int)
87
+
88
+ surrogate = SimplexTreeClassifier(
89
+ model=black_box,
90
+ n_features=2,
91
+ subdivision_levels=4,
92
+ n_fill=5000,
93
+ )
94
+ surrogate.fit() # samples + labels internally
95
+ ```
96
+
97
+
98
+
99
+ ### Data-driven subdivision
100
+
101
+ Instead of uniform (barycentric) subdivision, split only where the classifier is
102
+ still wrong:
103
+
104
+ ```python
105
+ clf = SimplexTreeClassifier(
106
+ subdivision_strategy="data_driven",
107
+ subdivision_levels=6, # here: maximum tree depth
108
+ max_leaves=500,
109
+ )
110
+ clf.fit(X, y)
111
+ ```
112
+
113
+
114
+
115
+ ## API overview
116
+
117
+ Constructor: `SimplexTreeClassifier(classifier=None, model=None, n_features=None, subdivision_levels=1, subdivision_strategy="barycentric", n_fill=5000, max_leaves=None, normalize=None, margin=0.05, device=None, tolerance=1e-10, random_state=None, data_driven_max_iter=50)`
118
+
119
+ Core methods:
120
+
121
+ - `transform(X)` - sparse barycentric embedding (GPU-accelerated); rows sum to 1.
122
+ - `fit(X=None, y=None)` - trains the classifier. In surrogate mode `X`/`y` are
123
+ optional. Automatically records same-side simplices afterward (linear only).
124
+ - `predict(X)` - predicted labels.
125
+
126
+ Geometry queries:
127
+
128
+ - `find_containing_simplex(point)` - the leaf simplex a point falls in.
129
+ - `find_adjacent_simplexes(simplex)` - leaves sharing a `(d-1)`-face.
130
+ - `get_simplex_vertices()` - vertices of every leaf simplex (renamed from
131
+ `get_simplex_boundaries`).
132
+ - `is_in_simplex(point, simplex)` - boolean containment test.
133
+
134
+ Decision-boundary analysis (linear classifiers, multiclass-aware):
135
+
136
+ - `identify_crossing_simplices()` - leaves crossed by the decision boundary.
137
+ - `remove_nonconvex_leaves(criterion="distance"|"convexity_sign", removal_factor=0.15, epsilon=None, keep_frac=None, min_depth=0.0, max_remove_frac=0.25, max_iter=10, refit=True)` - iteratively prune leaves whose
138
+ boundary bends non-convexly.
139
+ - `find_same_side_simplices()` - subdivisions that don't touch the boundary (also
140
+ run automatically at the end of `fit`, stored in `same_side_keys_`).
141
+ - `compute_plane_equations()` - the boundary hyperplane within each crossing leaf.
142
+
143
+ Helpers: `SimplexTree`, `Simplex`, `VertexRegistry`, `make_enclosing_simplex(d)`,
144
+ `get_device()`.
145
+
146
+ ## GPU acceleration
147
+
148
+ `transform` stacks each leaf's inverse edge matrix into tensors and embeds an
149
+ entire batch of points against all leaves at once. The device is auto-detected
150
+ (CUDA if present); override with `device="cpu"` / `device="cuda"`. Points outside
151
+ every leaf (or inside a degenerate one) fall back to an exact per-point search.
152
+
153
+ ## TODO
154
+
155
+ - [ ] **Overleaf cross-check.** Reconcile this package's operations against the
156
+ list of supported operations in the paper (Overleaf) and add anything missing.
157
+
@@ -0,0 +1,131 @@
1
+ # simplex-tree-classifier
2
+
3
+ A hierarchical **simplex-tree classifier**. It embeds points into the barycentric
4
+ coordinates of a subdivided simplex, then trains any scikit-learn estimator on
5
+ that sparse, geometry-aware feature space. The barycentric `transform` is
6
+ accelerated with **PyTorch** and runs on the GPU when one is available.
7
+
8
+ It works in two ways:
9
+
10
+ - **Dataset mode** - train directly on your data `(X, y)`.
11
+ - **Surrogate mode** - imitate an existing `model` (e.g. a neural network) to
12
+ produce an interpretable, piecewise-linear geometric surrogate.
13
+
14
+ Everything is **N-dimensional**.
15
+
16
+ ## Installation
17
+
18
+ ```bash
19
+ pip install simplex-tree-classifier # core (numpy, scipy, scikit-learn, torch)
20
+ pip install simplex-tree-classifier[viz] # + matplotlib for 2-D plots
21
+ ```
22
+
23
+ From source:
24
+
25
+ ```bash
26
+ cd simplex_tree_package
27
+ pip install -e .
28
+ ```
29
+
30
+
31
+
32
+ ## Quickstart
33
+
34
+
35
+
36
+ ### Dataset mode
37
+
38
+ ```python
39
+ import numpy as np
40
+ from simplex_tree_classifier import SimplexTreeClassifier
41
+
42
+ X = np.random.uniform(0, 1, size=(500, 4))
43
+ y = (X.sum(axis=1) > 2.0).astype(int)
44
+
45
+ clf = SimplexTreeClassifier(subdivision_levels=2) # default classifier: LinearSVC
46
+ clf.fit(X, y)
47
+ print(clf.predict(X[:5]))
48
+ ```
49
+
50
+
51
+
52
+ ### Surrogate mode (imitate a model)
53
+
54
+ Pass any object with `.predict(X)` or any callable `X -> labels`, plus the input
55
+ dimensionality. Fill points are sampled in `[0, 1]^d` and labeled by the model.
56
+
57
+ ```python
58
+ def black_box(X): # e.g. a wrapped PyTorch model
59
+ X = np.atleast_2d(X)
60
+ return (X[:, 0] + X[:, 1] > 1.0).astype(int)
61
+
62
+ surrogate = SimplexTreeClassifier(
63
+ model=black_box,
64
+ n_features=2,
65
+ subdivision_levels=4,
66
+ n_fill=5000,
67
+ )
68
+ surrogate.fit() # samples + labels internally
69
+ ```
70
+
71
+
72
+
73
+ ### Data-driven subdivision
74
+
75
+ Instead of uniform (barycentric) subdivision, split only where the classifier is
76
+ still wrong:
77
+
78
+ ```python
79
+ clf = SimplexTreeClassifier(
80
+ subdivision_strategy="data_driven",
81
+ subdivision_levels=6, # here: maximum tree depth
82
+ max_leaves=500,
83
+ )
84
+ clf.fit(X, y)
85
+ ```
86
+
87
+
88
+
89
+ ## API overview
90
+
91
+ Constructor: `SimplexTreeClassifier(classifier=None, model=None, n_features=None, subdivision_levels=1, subdivision_strategy="barycentric", n_fill=5000, max_leaves=None, normalize=None, margin=0.05, device=None, tolerance=1e-10, random_state=None, data_driven_max_iter=50)`
92
+
93
+ Core methods:
94
+
95
+ - `transform(X)` - sparse barycentric embedding (GPU-accelerated); rows sum to 1.
96
+ - `fit(X=None, y=None)` - trains the classifier. In surrogate mode `X`/`y` are
97
+ optional. Automatically records same-side simplices afterward (linear only).
98
+ - `predict(X)` - predicted labels.
99
+
100
+ Geometry queries:
101
+
102
+ - `find_containing_simplex(point)` - the leaf simplex a point falls in.
103
+ - `find_adjacent_simplexes(simplex)` - leaves sharing a `(d-1)`-face.
104
+ - `get_simplex_vertices()` - vertices of every leaf simplex (renamed from
105
+ `get_simplex_boundaries`).
106
+ - `is_in_simplex(point, simplex)` - boolean containment test.
107
+
108
+ Decision-boundary analysis (linear classifiers, multiclass-aware):
109
+
110
+ - `identify_crossing_simplices()` - leaves crossed by the decision boundary.
111
+ - `remove_nonconvex_leaves(criterion="distance"|"convexity_sign", removal_factor=0.15, epsilon=None, keep_frac=None, min_depth=0.0, max_remove_frac=0.25, max_iter=10, refit=True)` - iteratively prune leaves whose
112
+ boundary bends non-convexly.
113
+ - `find_same_side_simplices()` - subdivisions that don't touch the boundary (also
114
+ run automatically at the end of `fit`, stored in `same_side_keys_`).
115
+ - `compute_plane_equations()` - the boundary hyperplane within each crossing leaf.
116
+
117
+ Helpers: `SimplexTree`, `Simplex`, `VertexRegistry`, `make_enclosing_simplex(d)`,
118
+ `get_device()`.
119
+
120
+ ## GPU acceleration
121
+
122
+ `transform` stacks each leaf's inverse edge matrix into tensors and embeds an
123
+ entire batch of points against all leaves at once. The device is auto-detected
124
+ (CUDA if present); override with `device="cpu"` / `device="cuda"`. Points outside
125
+ every leaf (or inside a degenerate one) fall back to an exact per-point search.
126
+
127
+ ## TODO
128
+
129
+ - [ ] **Overleaf cross-check.** Reconcile this package's operations against the
130
+ list of supported operations in the paper (Overleaf) and add anything missing.
131
+
@@ -0,0 +1,47 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "simplex-tree-classifier"
7
+ version = "0.1.0"
8
+ description = "Hierarchical simplex-tree classifier with barycentric embedding, GPU-accelerated transform, and non-convex boundary pruning."
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ authors = [{ name = "Yasmin" }]
14
+ keywords = [
15
+ "machine-learning",
16
+ "classifier",
17
+ "simplex",
18
+ "barycentric",
19
+ "interpretability",
20
+ "surrogate-model",
21
+ ]
22
+ classifiers = [
23
+ "Programming Language :: Python :: 3",
24
+ "Operating System :: OS Independent",
25
+ "Intended Audience :: Science/Research",
26
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
27
+ ]
28
+ dependencies = [
29
+ "numpy>=1.21",
30
+ "scipy>=1.7",
31
+ "scikit-learn>=1.0",
32
+ "torch>=1.12",
33
+ ]
34
+
35
+ [project.optional-dependencies]
36
+ viz = ["matplotlib>=3.4"]
37
+ dev = ["pytest>=7.0"]
38
+
39
+ [project.urls]
40
+ Homepage = "https://github.com/yasmin/SimplexTreeClassifier"
41
+ Repository = "https://github.com/yasmin/SimplexTreeClassifier"
42
+
43
+ [tool.setuptools.packages.find]
44
+ where = ["src"]
45
+
46
+ [tool.setuptools.package-dir]
47
+ "" = "src"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,28 @@
1
+ """simplex_tree_classifier - hierarchical simplex-tree classifier.
2
+
3
+ Public API:
4
+ - ``SimplexTreeClassifier``: the main estimator (dataset or surrogate mode).
5
+ - ``SimplexTree`` / ``Simplex`` / ``VertexRegistry``: the geometric core.
6
+ - ``make_enclosing_simplex``: build an N-D simplex enclosing ``[0, 1]^d``.
7
+ - ``get_device``: resolve the torch device used for the transform.
8
+ """
9
+
10
+ from .vertex_registry import VertexRegistry
11
+ from .simplex import Simplex
12
+ from .simplex_tree import SimplexTree, make_enclosing_simplex
13
+ from .plane_equation import PlaneEquation
14
+ from .classifier import SimplexTreeClassifier
15
+ from .backend import get_device
16
+
17
+ __version__ = "0.1.0"
18
+
19
+ __all__ = [
20
+ "SimplexTreeClassifier",
21
+ "SimplexTree",
22
+ "Simplex",
23
+ "VertexRegistry",
24
+ "PlaneEquation",
25
+ "make_enclosing_simplex",
26
+ "get_device",
27
+ "__version__",
28
+ ]
@@ -0,0 +1,181 @@
1
+ """PyTorch backend that accelerates the barycentric ``transform``.
2
+
3
+ Leaves produced by barycentric subdivision are full ``d``-simplices (``d + 1``
4
+ vertices), so each non-degenerate leaf has an invertible edge matrix ``A`` with a
5
+ precomputed inverse. This backend stacks those per-leaf inverses into tensors and
6
+ embeds an entire batch of points against every leaf at once, on the GPU when one
7
+ is available. For each point it selects the leaf whose barycentric coordinates
8
+ are all non-negative and most interior.
9
+
10
+ Points that are not resolved on the GPU (outside every leaf, or inside a
11
+ degenerate leaf excluded from the tensor stack) are reported back so the caller
12
+ can fall back to the exact per-point CPU search.
13
+ """
14
+
15
+ from typing import List, Optional
16
+
17
+ import numpy as np
18
+
19
+ try:
20
+ import torch
21
+ _TORCH_AVAILABLE = True
22
+ except ImportError: # pragma: no cover - torch is a hard dependency, guard anyway
23
+ torch = None
24
+ _TORCH_AVAILABLE = False
25
+
26
+
27
+ def get_device(device=None):
28
+ """Resolve a torch device, defaulting to CUDA when available, else CPU.
29
+
30
+ Args:
31
+ device: An explicit ``torch.device``/str, or ``None`` to auto-detect.
32
+
33
+ Returns:
34
+ A ``torch.device``.
35
+ """
36
+ if not _TORCH_AVAILABLE:
37
+ raise ImportError(
38
+ "PyTorch is required for simplex_tree_classifier. Install it with "
39
+ "`pip install torch`."
40
+ )
41
+ if device is not None:
42
+ return torch.device(device)
43
+ if torch.cuda.is_available():
44
+ return torch.device("cuda")
45
+ return torch.device("cpu")
46
+
47
+
48
+ class TransformBackend:
49
+ """Precomputes per-leaf tensors and batch-embeds points into barycentric coords."""
50
+
51
+ def __init__(self, device=None, tolerance: float = 1e-10,
52
+ element_budget: int = 40_000_000):
53
+ self.device = get_device(device)
54
+ self.tolerance = tolerance
55
+ # Caps the number of (chunk x leaves x dim) elements held at once.
56
+ self.element_budget = int(element_budget)
57
+ self.dtype = torch.float64
58
+ self._built = False
59
+
60
+ self.leaves: List = []
61
+ self.n_leaves = 0
62
+ self.dimension = 0
63
+ # Global vertex indices per leaf, in the order alphas are produced.
64
+ self.leaf_vertex_indices: Optional[np.ndarray] = None
65
+ self._V0 = None # (L, d)
66
+ self._A_inv = None # (L, d, d)
67
+
68
+ @property
69
+ def is_built(self) -> bool:
70
+ return self._built and self.n_leaves > 0
71
+
72
+ def build(self, leaves: List) -> None:
73
+ """Stack tensors for all full, non-degenerate leaves.
74
+
75
+ Args:
76
+ leaves: Iterable of leaf ``SimplexTree`` nodes.
77
+ """
78
+ self._built = False
79
+ self.leaves = []
80
+ v0_list = []
81
+ a_inv_list = []
82
+ vidx_list = []
83
+
84
+ dimension = None
85
+ for leaf in leaves:
86
+ d = leaf.dimension
87
+ # Only stack full simplices with a usable inverse.
88
+ if leaf.n_vertices != d + 1:
89
+ continue
90
+ if getattr(leaf, "is_degenerate", False) or getattr(leaf, "A_inv", None) is None:
91
+ continue
92
+ if dimension is None:
93
+ dimension = d
94
+ elif d != dimension:
95
+ # Mixed dimensions should not happen within one tree; skip oddities.
96
+ continue
97
+ self.leaves.append(leaf)
98
+ v0_list.append(np.asarray(leaf.vertices[0], dtype=np.float64))
99
+ a_inv_list.append(np.asarray(leaf.A_inv, dtype=np.float64))
100
+ vidx_list.append(np.asarray(leaf.vertex_indices, dtype=np.int64))
101
+
102
+ self.n_leaves = len(self.leaves)
103
+ self.dimension = dimension or 0
104
+
105
+ if self.n_leaves == 0:
106
+ self.leaf_vertex_indices = None
107
+ self._V0 = None
108
+ self._A_inv = None
109
+ self._built = True
110
+ return
111
+
112
+ self.leaf_vertex_indices = np.stack(vidx_list, axis=0) # (L, d+1)
113
+ self._V0 = torch.as_tensor(np.stack(v0_list, axis=0),
114
+ dtype=self.dtype, device=self.device) # (L, d)
115
+ self._A_inv = torch.as_tensor(np.stack(a_inv_list, axis=0),
116
+ dtype=self.dtype, device=self.device) # (L, d, d)
117
+ self._built = True
118
+
119
+ def _chunk_size(self) -> int:
120
+ per_point = max(self.n_leaves * max(self.dimension, 1), 1)
121
+ return max(1, self.element_budget // per_point)
122
+
123
+ def embed(self, points: np.ndarray):
124
+ """Embed a batch of points against every stacked leaf.
125
+
126
+ Args:
127
+ points: Array of shape ``(m, d)``.
128
+
129
+ Returns:
130
+ Tuple ``(leaf_index, found, alphas)`` where:
131
+ * ``leaf_index`` (m,) int array indexes into ``self.leaves`` / rows of
132
+ ``self.leaf_vertex_indices`` (``-1`` when unresolved),
133
+ * ``found`` (m,) bool array marks GPU-resolved points,
134
+ * ``alphas`` (m, d+1) float array holds the barycentric coordinates
135
+ for the chosen leaf (rows for unresolved points are meaningless).
136
+ """
137
+ points = np.asarray(points, dtype=np.float64)
138
+ if points.ndim == 1:
139
+ points = points.reshape(1, -1)
140
+ m = points.shape[0]
141
+
142
+ leaf_index = np.full(m, -1, dtype=np.int64)
143
+ found = np.zeros(m, dtype=bool)
144
+ alphas_out = np.zeros((m, self.dimension + 1), dtype=np.float64)
145
+
146
+ if not self.is_built:
147
+ return leaf_index, found, alphas_out
148
+
149
+ tol = self.tolerance
150
+ chunk = self._chunk_size()
151
+ P_all = torch.as_tensor(points, dtype=self.dtype, device=self.device)
152
+
153
+ for start in range(0, m, chunk):
154
+ end = min(start + chunk, m)
155
+ P = P_all[start:end] # (c, d)
156
+ b = P[:, None, :] - self._V0[None, :, :] # (c, L, d)
157
+ alpha_r = torch.einsum("lij,clj->cli", self._A_inv, b) # (c, L, d)
158
+ alpha0 = 1.0 - alpha_r.sum(dim=-1) # (c, L)
159
+ alphas = torch.cat([alpha0[..., None], alpha_r], dim=-1) # (c, L, d+1)
160
+
161
+ inside = (alphas >= -tol).all(dim=-1) # (c, L)
162
+ min_alpha = alphas.amin(dim=-1) # (c, L)
163
+ neg_inf = torch.full_like(min_alpha, float("-inf"))
164
+ scored = torch.where(inside, min_alpha, neg_inf) # (c, L)
165
+
166
+ best_val, best_leaf = scored.max(dim=1) # (c,), (c,)
167
+ chunk_found = torch.isfinite(best_val) # (c,)
168
+
169
+ rows = torch.arange(end - start, device=self.device)
170
+ best_alphas = alphas[rows, best_leaf] # (c, d+1)
171
+
172
+ cf = chunk_found.cpu().numpy()
173
+ bl = best_leaf.cpu().numpy()
174
+ ba = best_alphas.cpu().numpy()
175
+
176
+ sl = slice(start, end)
177
+ found[sl] = cf
178
+ leaf_index[sl] = np.where(cf, bl, -1)
179
+ alphas_out[sl] = ba
180
+
181
+ return leaf_index, found, alphas_out