Refactor e2e tests + Add more models (#51)

BayesWitnesses · Feb 4, 2019 · 854185e · 854185e
1 parent 818db42
commit 854185e
Show file tree

Hide file tree

Showing 5 changed files with 240 additions and 157 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -9,7 +9,7 @@ install:
   - pip install -r requirements-test.txt
 
 script:
-  - pytest -v tests/ --cov=m2cgen/
+  - pytest -v tests/ --cov=m2cgen/ --durations=0
   - flake8 .
 
 after_success:

diff --git a/README.md b/README.md
@@ -14,6 +14,10 @@ m2cgen (Model 2 Code Generator) is a lightweight library which provides an easy
 
 ## Supported models
 
+### SVM
+- LinearSVC
+- LinearSVR
+
 ### Linear Regressors
 - LinearRegression
 - HuberRegressor
@@ -26,15 +30,22 @@ m2cgen (Model 2 Code Generator) is a lightweight library which provides an easy
 - LassoCV
 - LassoLars
 - LassoLarsIC
-- MultiTaskElasticNetCV (?!?)
-- MultiTaskLassoCV (?!?)
-- MultiTaskElasticNet (?!?)
-- MultiTaskLasso (?!?)
-
+- OrthogonalMatchingPursuit
+- OrthogonalMatchingPursuitCV
+- Ridge
+- RidgeCV
+- BayesianRidge
+- ARDRegression
+- SGDRegressor
+- PassiveAggressiveRegressor
 
 ### Logistic Regressors
 - LogisticRegression
 - LogisticRegressionCV
+- RidgeClassifier
+- RidgeClassifierCV
+- SGDClassifier
+- PassiveAggressiveClassifier
 
 ### Decision trees
 - DecisionTreeClassifier

diff --git a/m2cgen/exporters.py b/m2cgen/exporters.py
@@ -3,12 +3,50 @@
 
 
 SUPPORTED_MODELS = {
+    # SVM
+    "LinearSVC": assemblers.LinearModelAssembler,
+    "LinearSVR": assemblers.LinearModelAssembler,
+
+    # Linear Regressors
     "LinearRegression": assemblers.LinearModelAssembler,
+    "HuberRegressor": assemblers.LinearModelAssembler,
+    "ElasticNet": assemblers.LinearModelAssembler,
+    "ElasticNetCV": assemblers.LinearModelAssembler,
+    "TheilSenRegressor": assemblers.LinearModelAssembler,
+    "Lars": assemblers.LinearModelAssembler,
+    "LarsCV": assemblers.LinearModelAssembler,
+    "Lasso": assemblers.LinearModelAssembler,
+    "LassoCV": assemblers.LinearModelAssembler,
+    "LassoLars": assemblers.LinearModelAssembler,
+    "LassoLarsIC": assemblers.LinearModelAssembler,
+    "OrthogonalMatchingPursuit": assemblers.LinearModelAssembler,
+    "OrthogonalMatchingPursuitCV": assemblers.LinearModelAssembler,
+    "Ridge": assemblers.LinearModelAssembler,
+    "RidgeCV": assemblers.LinearModelAssembler,
+    "BayesianRidge": assemblers.LinearModelAssembler,
+    "ARDRegression": assemblers.LinearModelAssembler,
+    "SGDRegressor": assemblers.LinearModelAssembler,
+    "PassiveAggressiveRegressor": assemblers.LinearModelAssembler,
+
+    # Logistic Regressors
     "LogisticRegression": assemblers.LinearModelAssembler,
+    "LogisticRegressionCV": assemblers.LinearModelAssembler,
+    "RidgeClassifier": assemblers.LinearModelAssembler,
+    "RidgeClassifierCV": assemblers.LinearModelAssembler,
+    "SGDClassifier": assemblers.LinearModelAssembler,
+    "PassiveAggressiveClassifier": assemblers.LinearModelAssembler,
+
+    # Decision trees
     "DecisionTreeRegressor": assemblers.TreeModelAssembler,
     "DecisionTreeClassifier": assemblers.TreeModelAssembler,
+    "ExtraTreeRegressor": assemblers.TreeModelAssembler,
+    "ExtraTreeClassifier": assemblers.TreeModelAssembler,
+
+    # Ensembles
     "RandomForestRegressor": assemblers.RandomForestModelAssembler,
     "RandomForestClassifier": assemblers.RandomForestModelAssembler,
+    "ExtraTreesRegressor": assemblers.RandomForestModelAssembler,
+    "ExtraTreesClassifier": assemblers.RandomForestModelAssembler,
 }