[R] Support booster slicing. (#9948)

This commit is contained in:
david-cortes
2024-01-20 22:11:26 +01:00
committed by GitHub
parent c5d0608057
commit 5062a3ab46
8 changed files with 237 additions and 2 deletions

View File

@@ -64,6 +64,7 @@ extern SEXP XGDMatrixSliceDMatrix_R(SEXP, SEXP);
extern SEXP XGBSetGlobalConfig_R(SEXP);
extern SEXP XGBGetGlobalConfig_R(void);
extern SEXP XGBoosterFeatureScore_R(SEXP, SEXP);
extern SEXP XGBoosterSlice_R(SEXP, SEXP, SEXP, SEXP);
static const R_CallMethodDef CallEntries[] = {
{"XGDuplicate_R", (DL_FUNC) &XGDuplicate_R, 1},
@@ -114,6 +115,7 @@ static const R_CallMethodDef CallEntries[] = {
{"XGBSetGlobalConfig_R", (DL_FUNC) &XGBSetGlobalConfig_R, 1},
{"XGBGetGlobalConfig_R", (DL_FUNC) &XGBGetGlobalConfig_R, 0},
{"XGBoosterFeatureScore_R", (DL_FUNC) &XGBoosterFeatureScore_R, 2},
{"XGBoosterSlice_R", (DL_FUNC) &XGBoosterSlice_R, 4},
{NULL, NULL, 0}
};

View File

@@ -1289,3 +1289,18 @@ XGB_DLL SEXP XGBoosterFeatureScore_R(SEXP handle, SEXP json_config) {
return r_out;
}
XGB_DLL SEXP XGBoosterSlice_R(SEXP handle, SEXP begin_layer, SEXP end_layer, SEXP step) {
SEXP out = Rf_protect(XGBMakeEmptyAltrep());
R_API_BEGIN();
BoosterHandle handle_out = nullptr;
CHECK_CALL(XGBoosterSlice(R_ExternalPtrAddr(handle),
Rf_asInteger(begin_layer),
Rf_asInteger(end_layer),
Rf_asInteger(step),
&handle_out));
XGBAltrepSetPointer(out, handle_out);
R_API_END();
Rf_unprotect(1);
return out;
}

View File

@@ -402,4 +402,14 @@ XGB_DLL SEXP XGBoosterGetAttrNames_R(SEXP handle);
*/
XGB_DLL SEXP XGBoosterFeatureScore_R(SEXP handle, SEXP json_config);
/*!
* \brief Slice a fitted booster model (by rounds)
* \param handle handle to the fitted booster
* \param begin_layer start of the slice
* \param end_later end of the slice; end_layer=0 is equivalent to end_layer=num_boost_round
* \param step step size of the slice
* \return The sliced booster with the requested rounds only
*/
XGB_DLL SEXP XGBoosterSlice_R(SEXP handle, SEXP begin_layer, SEXP end_layer, SEXP step);
#endif // XGBOOST_WRAPPER_R_H_ // NOLINT(*)