added xgb classifier

alphamind/api.py

@@ -31,6 +31,7 @@ from alphamind.model import ConstLinearModel
 from alphamind.model import LogisticRegression
 from alphamind.model import RandomForestRegressor
 from alphamind.model import XGBRegressor
+from alphamind.model import XGBClassifier
 from alphamind.model import load_model
 from alphamind.model.data_preparing import fetch_data_package
 from alphamind.model.data_preparing import fetch_train_phase
@@ -70,6 +71,7 @@ __all__ = [
     'LogisticRegression',
     'RandomForestRegressor',
     'XGBRegressor',
+    'XGBClassifier',
     'load_model',
     'NaiveExecutor',
     'ThresholdExecutor',

alphamind/model/__init__.py

@@ -12,6 +12,7 @@ from alphamind.model.linearmodel import LogisticRegression
 
 from alphamind.model.treemodel import RandomForestRegressor
 from alphamind.model.treemodel import XGBRegressor
+from alphamind.model.treemodel import XGBClassifier
 
 from alphamind.model.loader import load_model
 
@@ -22,4 +23,5 @@ __all__ = ['LinearRegression',
            'LogisticRegression',
            'RandomForestRegressor',
            'XGBRegressor',
+           'XGBClassifier',
            'load_model']
\ No newline at end of file

alphamind/model/loader.py

@@ -12,6 +12,7 @@ from alphamind.model.linearmodel import LassoRegression
 from alphamind.model.linearmodel import LogisticRegression
 from alphamind.model.treemodel import RandomForestRegressor
 from alphamind.model.treemodel import XGBRegressor
+from alphamind.model.treemodel import XGBClassifier
 
 
 def load_model(model_desc: dict) -> ModelBase:
@@ -31,5 +32,7 @@ def load_model(model_desc: dict) -> ModelBase:
         return RandomForestRegressor.load(model_desc)
     elif 'XGBRegressor' in model_name_parts:
         return XGBRegressor.load(model_desc)
+    elif 'XGBClassifier' in model_name_parts:
+        return XGBClassifier.load(model_desc)
     else:
         raise ValueError('{0} is not currently supported in model loader.'.format(model_name))

alphamind/model/treemodel.py

@@ -11,6 +11,7 @@ from sklearn import __version__ as sklearn_version
 from sklearn.ensemble import RandomForestRegressor as RandomForestRegressorImpl
 from xgboost import __version__ as xgbboot_version
 from xgboost import XGBRegressor as XGBRegressorImpl
+from xgboost import XGBClassifier as XGBClassifierImpl
 from alphamind.model.modelbase import ModelBase
 from alphamind.utilities import alpha_logger
 
@@ -49,7 +50,8 @@ class XGBRegressor(ModelBase):
                  n_estimators: int=100,
                  learning_rate: float=0.1,
                  max_depth: int=3,
-                 features: List=None, **kwargs):
+                 features: List=None,
+                 **kwargs):
         super().__init__(features)
         self.impl = XGBRegressorImpl(n_estimators=n_estimators,
                                      learning_rate=learning_rate,
@@ -77,5 +79,39 @@ class XGBRegressor(ModelBase):
         return self.impl.feature_importances_.tolist()
 
 
+class XGBClassifier(ModelBase):
+
+    def __init__(self,
+                 n_estimators: int=100,
+                 learning_rate: float=0.1,
+                 max_depth: int=3,
+                 features: List = None,
+                 **kwargs):
+        super().__init__(features)
+        self.impl = XGBClassifierImpl(n_estimators=n_estimators,
+                                     learning_rate=learning_rate,
+                                     max_depth=max_depth,
+                                     **kwargs)
+
+    def save(self) -> dict:
+        model_desc = super().save()
+        model_desc['xgbboot_version'] = xgbboot_version
+        model_desc['importances'] = self.importances
+        return model_desc
+
+    @classmethod
+    def load(cls, model_desc: dict):
+        obj_layout = super().load(model_desc)
+
+        if LooseVersion(sklearn_version) < LooseVersion(model_desc['xgbboot_version']):
+            alpha_logger.warning('Current xgboost version {0} is lower than the model version {1}. '
+                                 'Loaded model may work incorrectly.'.format(xgbboot_version,
+                                                                             model_desc['xgbboot_version']))
+        return obj_layout
+
+    @property
+    def importances(self):
+        return self.impl.feature_importances_.tolist()
+
 
 

alphamind/tests/model/test_treemodel.py

@@ -10,6 +10,7 @@ import numpy as np
 from alphamind.model.loader import load_model
 from alphamind.model.treemodel import RandomForestRegressor
 from alphamind.model.treemodel import XGBRegressor
+from alphamind.model.treemodel import XGBClassifier
 
 
 class TestTreeModel(unittest.TestCase):
@@ -28,7 +29,7 @@ class TestTreeModel(unittest.TestCase):
         sample_x = np.random.randn(100, 10)
         np.testing.assert_array_almost_equal(model.predict(sample_x), new_model.predict(sample_x))
 
-    def tes_xgb_regress(self):
+    def test_xgb_regress_persistence(self):
         model = XGBRegressor(features=list(range(10)))
         x = np.random.randn(1000, 10)
         y = np.random.randn(1000)
@@ -40,4 +41,19 @@ class TestTreeModel(unittest.TestCase):
         self.assertEqual(model.features, new_model.features)
 
         sample_x = np.random.randn(100, 10)
-        np.testing.assert_array_almost_equal(model.predict(sample_x), new_model.predict(sample_x))
+        np.testing.assert_array_almost_equal(model.predict(sample_x), new_model.predict(sample_x))
\ No newline at end of file
+
+    def test_xgb_classify_persistence(self):
+        model = XGBClassifier(features=list(range(10)))
+        x = np.random.randn(1000, 10)
+        y = np.random.randn(1000)
+        y = np.where(y > 0, 1, 0)
+
+        model.fit(x, y)
+
+        desc = model.save()
+        new_model = load_model(desc)
+        self.assertEqual(model.features, new_model.features)
+
+        sample_x = np.random.randn(100, 10)
+        np.testing.assert_array_almost_equal(model.predict(sample_x), new_model.predict(sample_x))