def test_predict(self, lung_X, lung_y):
tree = RangerTreeSurvival()
tree.fit(lung_X, lung_y)
pred = tree.predict(lung_X)
assert len(pred) == lung_X.shape[0]
# test with single record
lung_X_record = lung_X.values[0:1, :]
pred = tree.predict(lung_X_record)
assert len(pred) == 1
def test_sample_weight(self, lung_X, lung_y):
forest_w = RangerTreeSurvival()
forest_w.fit(lung_X, lung_y, sample_weight=[1] * len(lung_y))
tree = RangerTreeSurvival()
tree.fit(lung_X, lung_y)
pred_w = forest_w.predict(lung_X)
pred = tree.predict(lung_X)
np.testing.assert_array_equal(pred.reshape(-1, 1),
pred_w.reshape(-1, 1))
def test_sample_fraction(self, lung_X, lung_y):
tree = RangerTreeSurvival(sample_fraction=0.69)
tree.fit(lung_X, lung_y)
assert tree.sample_fraction_ == [0.69]
# test with single record
lung_X_record = lung_X.values[0:1, :]
pred = tree.predict(lung_X_record)
assert len(pred) == 1
def test_categorical_features(self, lung_X, lung_y,
respect_categorical_features):
# add a categorical feature
categorical_col = np.atleast_2d(
np.array([random.choice([0, 1]) for _ in range(lung_X.shape[0])]))
lung_X_c = np.hstack((lung_X, categorical_col.transpose()))
categorical_features = [lung_X.shape[1]]
tree = RangerTreeSurvival(
respect_categorical_features=respect_categorical_features, )
if respect_categorical_features not in [
"partition", "ignore", "order"
]:
with pytest.raises(ValueError):
tree.fit(lung_X_c,
lung_y,
categorical_features=categorical_features)
return
tree.fit(lung_X_c, lung_y, categorical_features=categorical_features)
tree.predict(lung_X_c)