[dask] Support more meta data on functional interface. (#6132)

* Add base_margin, label_(lower|upper)_bound.
* Test survival training with dask.
This commit is contained in:
Jiaming Yuan
2020-09-21 16:56:37 +08:00
committed by GitHub
parent 7065779afa
commit 33d80ffad0
4 changed files with 154 additions and 48 deletions

View File

@@ -3,9 +3,10 @@ import pytest
import numpy as np
import xgboost as xgb
import json
from pathlib import Path
import os
dpath = os.path.join(tm.PROJECT_ROOT, 'demo', 'data')
dpath = Path('demo/data')
def test_aft_survival_toy_data():
# See demo/aft_survival/aft_survival_viz_demo.py
@@ -51,10 +52,10 @@ def test_aft_survival_toy_data():
for tree in model_json:
assert gather_split_thresholds(tree).issubset({2.5, 3.5, 4.5})
@pytest.mark.skipif(**tm.no_pandas())
@pytest.mark.skipif(**tm.no_pandas())
def test_aft_survival_demo_data():
import pandas as pd
df = pd.read_csv(dpath / 'veterans_lung_cancer.csv')
df = pd.read_csv(os.path.join(dpath, 'veterans_lung_cancer.csv'))
y_lower_bound = df['Survival_label_lower_bound']
y_upper_bound = df['Survival_label_upper_bound']