Updated based on comments from jkbradley

Author: imatiach-msft

mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala

@@ -177,8 +177,6 @@ class DecisionTreeClassificationModel private[ml] (
   /**
    * Construct a decision tree classification model.
    * @param rootNode  Root node of tree, with other nodes attached.
-   * @param numFeatures The number of features.
-   * @param numClasses The number of classes to predict.
    */
   private[ml] def this(rootNode: Node, numFeatures: Int, numClasses: Int) =
     this(Identifiable.randomUID("dtc"), rootNode, numFeatures, numClasses)

mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala

@@ -170,7 +170,7 @@ class GBTClassifier @Since("1.4.0") (
       maxDepth, maxBins, maxIter, maxMemoryInMB, minInfoGain, minInstancesPerNode,
       seed, stepSize, subsamplingRate, cacheNodeIds, checkpointInterval)
     instr.logNumFeatures(numFeatures)
-    instr.logNumClasses(2)
+    instr.logNumClasses(numClasses)
 
     val (baseLearners, learnerWeights) = GradientBoostedTrees.run(oldDataset, boostingStrategy,
       $(seed))
@@ -209,8 +209,7 @@ class GBTClassificationModel private[ml](
     @Since("1.6.0") override val uid: String,
     private val _trees: Array[DecisionTreeRegressionModel],
     private val _treeWeights: Array[Double],
-    @Since("1.6.0") override val numFeatures: Int,
-    @Since("2.2.0") override val numClasses: Int)
+    @Since("1.6.0") override val numFeatures: Int)
   extends ProbabilisticClassificationModel[Vector, GBTClassificationModel]
   with GBTClassifierParams with TreeEnsembleModel[DecisionTreeRegressionModel]
   with MLWritable with Serializable {
@@ -219,20 +218,6 @@ class GBTClassificationModel private[ml](
   require(_trees.length == _treeWeights.length, "GBTClassificationModel given trees, treeWeights" +
     s" of non-matching lengths (${_trees.length}, ${_treeWeights.length}, respectively).")
 
-  /**
-   * Construct a GBTClassificationModel
-   *
-   * @param _trees  Decision trees in the ensemble.
-   * @param _treeWeights  Weights for the decision trees in the ensemble.
-   * @param numFeatures  The number of features.
-   */
-  private[ml] def this(
-      uid: String,
-      _trees: Array[DecisionTreeRegressionModel],
-      _treeWeights: Array[Double],
-      numFeatures: Int) =
-    this(uid, _trees, _treeWeights, numFeatures, 2)
-
   /**
    * Construct a GBTClassificationModel
    *
@@ -241,7 +226,7 @@ class GBTClassificationModel private[ml](
    */
   @Since("1.6.0")
   def this(uid: String, _trees: Array[DecisionTreeRegressionModel], _treeWeights: Array[Double]) =
-    this(uid, _trees, _treeWeights, -1, 2)
+    this(uid, _trees, _treeWeights, -1)
 
   @Since("1.4.0")
   override def trees: Array[DecisionTreeRegressionModel] = _trees
@@ -294,7 +279,7 @@ class GBTClassificationModel private[ml](
 
   @Since("1.4.0")
   override def copy(extra: ParamMap): GBTClassificationModel = {
-    copyValues(new GBTClassificationModel(uid, _trees, _treeWeights, numFeatures, numClasses),
+    copyValues(new GBTClassificationModel(uid, _trees, _treeWeights, numFeatures),
       extra).setParent(parent)
   }
 
@@ -339,7 +324,6 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
 
   private val numFeaturesKey: String = "numFeatures"
   private val numTreesKey: String = "numTrees"
-  private val numClassesKey: String = "numClasses"
 
   @Since("2.0.0")
   override def read: MLReader[GBTClassificationModel] = new GBTClassificationModelReader
@@ -354,8 +338,7 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
 
       val extraMetadata: JObject = Map(
         numFeaturesKey -> instance.numFeatures,
-        numTreesKey -> instance.getNumTrees,
-        numClassesKey -> instance.numClasses)
+        numTreesKey -> instance.getNumTrees)
       EnsembleModelReadWrite.saveImpl(instance, path, sparkSession, extraMetadata)
     }
   }
@@ -372,7 +355,6 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
         EnsembleModelReadWrite.loadImpl(path, sparkSession, className, treeClassName)
       val numFeatures = (metadata.metadata \ numFeaturesKey).extract[Int]
       val numTrees = (metadata.metadata \ numTreesKey).extract[Int]
-      val numClasses = (metadata.metadata \ numClassesKey).extract[Int]
 
       val trees: Array[DecisionTreeRegressionModel] = treesData.map {
         case (treeMetadata, root) =>
@@ -384,7 +366,7 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
       require(numTrees == trees.length, s"GBTClassificationModel.load expected $numTrees" +
         s" trees based on metadata but found ${trees.length} trees.")
       val model = new GBTClassificationModel(metadata.uid,
-        trees, treeWeights, numFeatures, numClasses)
+        trees, treeWeights, numFeatures)
       DefaultParamsReader.getAndSetParams(model, metadata)
       model
     }
@@ -395,15 +377,14 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
       oldModel: OldGBTModel,
       parent: GBTClassifier,
       categoricalFeatures: Map[Int, Int],
-      numFeatures: Int = -1,
-      numClasses: Int = 2): GBTClassificationModel = {
+      numFeatures: Int = -1): GBTClassificationModel = {
     require(oldModel.algo == OldAlgo.Classification, "Cannot convert GradientBoostedTreesModel" +
       s" with algo=${oldModel.algo} (old API) to GBTClassificationModel (new API).")
     val newTrees = oldModel.trees.map { tree =>
       // parent for each tree is null since there is no good way to set this.
       DecisionTreeRegressionModel.fromOld(tree, null, categoricalFeatures)
     }
     val uid = if (parent != null) parent.uid else Identifiable.randomUID("gbtc")
-    new GBTClassificationModel(uid, newTrees, oldModel.treeWeights, numFeatures, numClasses)
+    new GBTClassificationModel(uid, newTrees, oldModel.treeWeights, numFeatures)
   }
 }

mllib/src/main/scala/org/apache/spark/mllib/tree/loss/LogLoss.scala

@@ -52,9 +52,10 @@ object LogLoss extends ClassificationLoss {
     2.0 * MLUtils.log1pExp(-margin)
   }
 
-  override private[spark] def computeProbability(prediction: Double): Double = {
-    // The probability can be calculated as:
-    // p+(x) = 1 / (1 + e^(-2 * F(x)))
-    1.0 / (1.0 + math.exp(-2.0 * prediction))
+  /**
+   * Returns the estimated probability of a label of 1.0.
+   */
+  override private[spark] def computeProbability(margin: Double): Double = {
+    1.0 / (1.0 + math.exp(-2.0 * margin))
   }
 }

mllib/src/main/scala/org/apache/spark/mllib/tree/loss/Loss.scala

@@ -70,8 +70,6 @@ trait Loss extends Serializable {
 private[spark] trait ClassificationLoss extends Loss {
   /**
    * Computes the class probability given the margin.
-   * @param prediction The margin.
-   * @return The class probability from the margin.
    */
-  private[spark] def computeProbability(prediction: Double): Double
+  private[spark] def computeProbability(margin: Double): Double
 }

mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala

@@ -115,7 +115,7 @@ class GBTClassifierSuite extends SparkFunSuite with MLlibTestSparkContext
     val gbt2 = new GBTClassifier
     val threshold = Array(0.3, 0.7)
     gbt2.setThresholds(threshold)
-    assert(gbt2.getThresholds.zip(threshold).forall { case(t1, t2) => t1 === t2 })
+    assert(gbt2.getThresholds === threshold)
   }
 
   test("thresholds prediction") {