Merge pull request #51 from ealcobaca/model-based-updates

Model based updates

pymfe/mfe.py

@@ -49,6 +49,7 @@ def __init__(self,
                  num_cv_folds: int = 10,
                  shuffle_cv_folds: bool = False,
                  lm_sample_frac: float = 1.0,
+                 hypparam_model_dt: t.Optional[t.Dict[str, t.Any]] = None,
                  suppress_warnings: bool = False,
                  random_state: t.Optional[int] = None) -> None:
         """Provides easy access for metafeature extraction from datasets.
@@ -171,6 +172,19 @@ def __init__(self,
             Sample proportion used to produce the ``landmarking`` metafeatures.
             This argument must be in 0.5 and 1.0 (both inclusive) interval.
 
+        hypparam_model_dt : :obj:`dict`, optional
+            Dictionary providing extra hyperparameters for the Decision Tree
+            algorithm for building the Decision Tree model, used to extract the
+            model-based metafeatures. The class used to fit the model is the
+            ``sklearn.tree.DecisionTreeClassifier`` (sklearn library). Using
+            this argument, it is possible to provide extra arguments in the
+            DecisionTreeClassifier class initialization (e.g., ``max_depth``
+            and ``min_samples_split``.) In order to use this argument, provide
+            the DecisionTreeClassifier init argument name as the dictionary
+            keys and the corresponding custom values, as the dictionary values.
+            Example:
+            {"min_samples_split": 10, "criterion": "entropy"}
+
         suppress_warnings : :obj:`bool`, optional
             If True, then ignore all warnings invoked at the instantiation
             time.
@@ -291,6 +305,8 @@ def __init__(self,
                              .format(random_state))
 
         self.score = _internal.check_score(score, self.groups)
+        self.hypparam_model_dt = (hypparam_model_dt.copy()
+                                  if hypparam_model_dt else None)
 
     def _call_summary_methods(
             self,
@@ -904,6 +920,7 @@ def fit(self,
             "score": self.score,
             "random_state": self.random_state,
             "cat_cols": self._attr_indexes_cat,
+            "hypparam_model_dt": self.hypparam_model_dt,
         }
 
         # Custom arguments from preprocessing methods