simple_smooth_tree_predict.py

import jax.numpy as jnp

from smooth_binary_node import SmoothBinaryNode

node = SmoothBinaryNode()

# feature 1 will be used s criteria
# The threshold is 50
params = {'weights': jnp.array([0, 10, 0]),
          'leaves': jnp.array([-1, 1]),
          'bias': 50}

features = jnp.array([[1, 80, 7]])
pred = node.predict(params, features)
# 80 is greater than 50
# we expect a 1
print('Prediction', pred)
#> 1


features = jnp.array([[1, 20, 7]])
pred = node.predict(params, features)
# 20 is lesser than 50
# we expect a -1
print('Prediction', pred)
#> -1