Skip to content

Commit

Permalink
[MLLIB] Add fit intercept term
Browse files Browse the repository at this point in the history
Added unit tests and changed docs in line with PR comments
  • Loading branch information
Omede Firouz committed Apr 3, 2015
1 parent bd9663c commit 329c1e2
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,8 @@ private[ml] trait HasFitIntercept extends Params {
* param for fitting the intercept term
* @group param
*/
val fitIntercept: BooleanParam = new BooleanParam(this, "fitIntercept", "fits the intercept term or not")
val fitIntercept: BooleanParam =
new BooleanParam(this, "fitIntercept", "indicates whether to fit an intercept term")

/** @group getParam */
def getFitIntercept: Boolean = get(fitIntercept)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ class LogisticRegressionSuite extends FunSuite with MLlibTestSparkContext {
assert(lr.getPredictionCol == "prediction")
assert(lr.getRawPredictionCol == "rawPrediction")
assert(lr.getProbabilityCol == "probability")
assert(lr.getFitIntercept == true)
val model = lr.fit(dataset)
model.transform(dataset)
.select("label", "probability", "prediction", "rawPrediction")
Expand All @@ -55,6 +56,14 @@ class LogisticRegressionSuite extends FunSuite with MLlibTestSparkContext {
assert(model.getPredictionCol == "prediction")
assert(model.getRawPredictionCol == "rawPrediction")
assert(model.getProbabilityCol == "probability")
assert(model.intercept != 0.0)
}

test("logistic regression doesn't fit intercept when fitIntercept is off") {
val lr = new LogisticRegression
lr.setFitIntercept(false)
val model = lr.fit(dataset)
assert(model.intercept == 0.0)
}

test("logistic regression with setters") {
Expand Down

0 comments on commit 329c1e2

Please sign in to comment.