Skip to content

Commit

Permalink
Move number of split quantiles from hardwired to node class member
Browse files Browse the repository at this point in the history
  • Loading branch information
EgorKraevTransferwise committed May 14, 2024
1 parent 7e846f8 commit 98d59a8
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion wise_pizza/solve/tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ def __init__(
self.children = None
self.dim_split = dim_split or {}
self.model = None
# For dimension splitting candidates, hardwired for now
self.num_bins = 10

@property
def depth(self):
Expand Down Expand Up @@ -118,7 +120,7 @@ def error_improvement(self):
if np.any(np.isnan(self.df[dim + "_encoded"])): # pragma: no cover
raise ValueError("NaNs in encoded values")
# Get split candidates for brute force search
deciles = np.array([q / 10.0 for q in range(1, 10)])
deciles = np.array([q / self.num_bins for q in range(1, self.num_bins)])

splits = weighted_quantiles(
self.df[dim + "_encoded"], deciles, self.df["weights"]
Expand Down

0 comments on commit 98d59a8

Please sign in to comment.