Encapsulate Binary Classification parameters

dpkatz · Jun 12, 2018 · b9c613f · b9c613f
1 parent 3facc11
commit b9c613f
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 6 deletions.
diff --git a/examples/binary_classification/Main.hs b/examples/binary_classification/Main.hs
@@ -17,7 +17,7 @@ import           LightGBM.Utils.Test (fileDiff)
 
 trainParams :: [P.Param]
 trainParams =
-  [ P.Objective P.BinaryClassification
+  [ P.Objective $ P.BinaryClassification []
   , P.Metric [P.BinaryLogloss, P.AUC]
   , P.TrainingMetric True
   , P.LearningRate $$(refineTH 0.1)

diff --git a/examples/titanic/Main.hs b/examples/titanic/Main.hs
@@ -25,7 +25,7 @@ import           ConvertData (csvFilter, predsToKaggleFormat, testFilter)
 
 trainParams :: [P.Param]
 trainParams =
-  [ P.Objective P.BinaryClassification
+  [ P.Objective $ P.BinaryClassification []
   , P.Metric [P.BinaryLogloss, P.AUC]
   , P.TrainingMetric True
   , P.LearningRate $$(refineTH 0.1)

diff --git a/src/LightGBM/Internal/CommandLineWrapper.hs b/src/LightGBM/Internal/CommandLineWrapper.hs
@@ -81,12 +81,14 @@ applicationPMap =
     , (P.Regression P.Quantile, "quantile")
     , (P.Regression P.MAPE, "mape")
     , (P.Regression P.Gamma, "gamma")
-    , (P.BinaryClassification, "binary")
     , (P.CrossEntropy P.XEntropy, "xentropy")
     , (P.CrossEntropy P.XEntropyLambda, "xentlambda")
     , (P.LambdaRank, "lambdarank")
     ]
 
+mkBinaryClassString :: P.BinaryClassParam -> String
+mkBinaryClassString (P.IsUnbalance b) = "is_unbalance=" ++ fmap toLower (show b)
+
 mkTweedieString :: P.TweedieRegressionParam -> String
 mkTweedieString (P.TweedieVariancePower p) = "tweedie_variance_power=" ++ show p
 
@@ -114,6 +116,8 @@ colSelPrefix (P.ColName _) = "name:"
 
 -- | Construct the option string for the command.
 mkOptionString :: P.Param -> [String]
+mkOptionString (P.Objective (P.BinaryClassification bcParams)) =
+  ["application=binary"] ++ map mkBinaryClassString bcParams
 mkOptionString (P.Objective (P.MultiClass P.MultiClassSimple n)) =
   ["application=multiclass", "num_classes=" ++ show n]
 mkOptionString (P.Objective (P.MultiClass P.MultiClassOneVsAll n)) =
@@ -215,7 +219,6 @@ mkOptionString (P.Sigmoid d) = ["sigmoid=" ++ show (unrefine d)]
 mkOptionString (P.Alpha d) = ["alpha=" ++ show (unrefine d)]
 mkOptionString (P.ScalePosWeight d) = ["scale_pos_weight=" ++ show d]
 mkOptionString (P.BoostFromAverage b) = ["boost_from_average=" ++ show b]
-mkOptionString (P.IsUnbalance b) = ["is_unbalance=" ++ show b]
 mkOptionString (P.MaxPosition n) = ["max_position=" ++ show (unrefine n)]
 mkOptionString (P.LabelGain ds) =
   ["label_gain=" ++ intercalate "," (map show ds)]

diff --git a/src/LightGBM/Parameters.hs b/src/LightGBM/Parameters.hs
@@ -14,6 +14,7 @@ module LightGBM.Parameters
   ( -- * Parameters
     Param(..)
   , Application(..)
+  , BinaryClassParam(..)
   , Booster(..)
   , DARTParam(..)
   , Device(..)
@@ -123,7 +124,6 @@ data Param
   | Alpha OpenProperFraction -- ^ Used in Huber loss and Quantile regression
   | ScalePosWeight Double -- ^ Used in Binary classification
   | BoostFromAverage Bool -- ^ Used only in RegressionL2 task
-  | IsUnbalance Bool -- ^ Used in Binary classification (set to true if training data are unbalanced)
   | MaxPosition PositiveInt -- ^ Used in LambdaRank
   | LabelGain [Double] -- ^ Used in LambdaRank
   | RegSqrt Bool -- ^ Only used in RegressionL2
@@ -132,6 +132,11 @@ data Param
   | TrainingMetric Bool
   deriving (Eq, Show)
 
+data BinaryClassParam =
+  IsUnbalance Bool -- ^ Set to true if training data are unbalanced
+  deriving (Eq, Show, Generic)
+instance Hashable BinaryClassParam
+
 --- | Parameters for Fair loss regression
 data FairRegressionParam =
   FairC PositiveDouble
@@ -269,7 +274,7 @@ type NumClasses = Natural
 -- | LightGBM can be used for a variety of applications
 data Application
   = Regression RegressionApp
-  | BinaryClassification
+  | BinaryClassification [BinaryClassParam]
   | MultiClass MultiClassStyle NumClasses
   | CrossEntropy XEApp
   | LambdaRank                  -- ^ A ranking algorithm