Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[jvm-packages] Updates to Java Booster to support other feature importance measures (#3801)

Browse Source 
* Updates to Booster to support other feature importances

* Add returns for Java methods

* Pass Scala style checks

* Pass Java style checks

* Fix indents

* Use class instead of enum

* Return map string double

* A no longer broken build, thanks to mvn package local build

* Add a unit test to increase code coverage back

* Address code review on main code

* Add more unit tests for different feature importance scores

* Address more CR
This commit is contained in:
Shayak Banerjee
2019-01-02 04:13:14 -05:00
committed by Nan Zhu
parent 1f022929f4
commit 431c850c03
3 changed files with 195 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/Booster.scala
View File

@@ -204,7 +204,7 @@ class Booster private[xgboost4j](private[xgboost4j] var booster: JBooster)
/**
* Get importance of each feature
* Get importance of each feature based on weight only (number of splits)
*
* @return featureScoreMap key: feature index, value: feature importance score
*/
@@ -214,7 +214,8 @@ class Booster private[xgboost4j](private[xgboost4j] var booster: JBooster)
}
/**
* Get importance of each feature with specified feature names.
* Get importance of each feature based on weight only
* (number of splits), with specified feature names.
*
* @return featureScoreMap key: feature name, value: feature importance score
*/
@@ -223,6 +224,31 @@ class Booster private[xgboost4j](private[xgboost4j] var booster: JBooster)
booster.getFeatureScore(featureNames).asScala
}
/**
* Get importance of each feature based on information gain or cover
* Supported: ["gain, "cover", "total_gain", "total_cover"]
*
* @return featureScoreMap key: feature index, value: feature importance score
*/
@throws(classOf[XGBoostError])
def getScore(featureMap: String, importanceType: String): Map[String, Double] = {
Map(booster.getScore(featureMap, importanceType)
.asScala.mapValues(_.doubleValue).toSeq: _*)
}
/**
* Get importance of each feature based on information gain or cover
* , with specified feature names.
* Supported: ["gain, "cover", "total_gain", "total_cover"]
*
* @return featureScoreMap key: feature name, value: feature importance score
*/
@throws(classOf[XGBoostError])
def getScore(featureNames: Array[String], importanceType: String): Map[String, Double] = {
Map(booster.getScore(featureNames, importanceType)
.asScala.mapValues(_.doubleValue).toSeq: _*)
}
def getVersion: Int = booster.getVersion
def toByteArray: Array[Byte] = {