Make PrimalSolutionComponent accept n_jobs argument

src/python/miplearn/components/primal.py

@@ -7,6 +7,7 @@ from copy import deepcopy
 from miplearn.classifiers.adaptive import AdaptiveClassifier
 from miplearn.components import classifier_evaluation_dict
 from sklearn.metrics import roc_curve
+from p_tqdm import p_map
 
 from .component import Component
 from ..extractors import *
@@ -45,54 +46,63 @@ class PrimalSolutionComponent(Component):
     def after_solve(self, solver, instance, model, results):
         pass
 
-    def fit(self, training_instances):
+    def fit(self, training_instances, n_jobs=1):
         logger.debug("Extracting features...")
         features = VariableFeaturesExtractor().extract(training_instances)
         solutions = SolutionExtractor().extract(training_instances)
 
-        for category in tqdm(features.keys(), desc="Fit (primal)"):
+        def _fit(args):
+            category, label = args[0], args[1]
             x_train = features[category]
             y_train = solutions[category]
-            for label in [0, 1]:
+            y = y_train[:, label].astype(int)
-                y = y_train[:, label].astype(int)
+
-
+            if isinstance(self.classifier_prototype, list):
-                logger.debug("Fitting predictors[%s, %s]:" % (category, label))
+                clf = deepcopy(self.classifier_prototype[label])
-                if isinstance(self.classifier_prototype, list):
+            else:
-                    pred = deepcopy(self.classifier_prototype[label])
+                clf = deepcopy(self.classifier_prototype)
-                else:
+            clf.fit(x_train, y)
-                    pred = deepcopy(self.classifier_prototype)
+
-                pred.fit(x_train, y)
+            y_avg = np.average(y)
-                self.classifiers[category, label] = pred
+            if (not self.dynamic_thresholds) or y_avg <= 0.001 or y_avg >= 0.999:
-
+                return {"classifier": clf,
-                # If y is either always one or always zero, set fixed threshold
+                        "threshold": self.min_threshold[label]}
-                y_avg = np.average(y)
+
-                if (not self.dynamic_thresholds) or y_avg <= 0.001 or y_avg >= 0.999:
+            proba = clf.predict_proba(x_train)
-                    self.thresholds[category, label] = self.min_threshold[label]
+            assert isinstance(proba, np.ndarray), \
-                    logger.debug("    Setting threshold to %.4f" % self.min_threshold[label])
+                "classifier should return numpy array"
-                    continue
+            assert proba.shape == (x_train.shape[0], 2), \
-
+                "classifier should return (%d,%d)-shaped array, not %s" % (
-                proba = pred.predict_proba(x_train)
+                    x_train.shape[0], 2, str(proba.shape))
-                assert isinstance(proba, np.ndarray), \
+
-                    "classifier should return numpy array"
+            y_scores = proba[:, 1]
-                assert proba.shape == (x_train.shape[0], 2), \
+            fpr, tpr, thresholds = roc_curve(y, y_scores)
-                    "classifier should return (%d,%d)-shaped array, not %s" % (
+            k = 0
-                        x_train.shape[0], 2, str(proba.shape))
+            while True:
-
+                if (k + 1) > len(fpr):
-                # Calculate threshold dynamically using ROC curve
+                    break
-                y_scores = proba[:, 1]
+                if fpr[k + 1] > self.max_fpr[label]:
-                fpr, tpr, thresholds = roc_curve(y, y_scores)
+                    break
-                k = 0
+                if thresholds[k + 1] < self.min_threshold[label]:
-                while True:
+                    break
-                    if (k + 1) > len(fpr):
+                k = k + 1
-                        break
+            self.thresholds[category, label] = thresholds[k]
-                    if fpr[k + 1] > self.max_fpr[label]:
+
-                        break
+            return {"classifier": clf,
-                    if thresholds[k + 1] < self.min_threshold[label]:
+                    "threshold": thresholds[k]}
-                        break
+
-                    k = k + 1
+        items = [(category, label)
-                logger.debug("    Setting threshold to %.4f (fpr=%.4f, tpr=%.4f)" %
+                 for category in features.keys()
-                             (thresholds[k], fpr[k], tpr[k]))
+                 for label in [0, 1]]
-                self.thresholds[category, label] = thresholds[k]
+
+        if n_jobs == 1:
+            results = list(map(_fit, tqdm(items, desc="Fit (primal)")))
+        else:
+            results = p_map(_fit, items, num_cpus=n_jobs)
+
+        for (idx, (category, label)) in enumerate(items):
+            self.thresholds[category, label] = results[idx]["threshold"]
+            self.classifiers[category, label] = results[idx]["classifier"]
 
     def predict(self, instance):
         x_test = VariableFeaturesExtractor().extract([instance])

src/python/miplearn/components/tests/test_primal.py

@@ -90,3 +90,11 @@ def test_evaluate():
                                    'True negative (%)': 50.0,
                                    'True positive': 1,
                                    'True positive (%)': 25.0}}}
+
+
+def test_primal_parallel_fit():
+    instances, models = get_training_instances_and_models()
+    comp = PrimalSolutionComponent()
+    comp.fit(instances, n_jobs=2)
+    assert len(comp.classifiers) == 2
+    assert len(comp.thresholds) == 2