From b0b013dd0a699ea57c0c424c5b85721081a5b00a Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Mon, 25 Jan 2021 15:19:58 -0600
Subject: [PATCH] Fix all tests

---
 miplearn/classifiers/adaptive.py           |  7 ++++
 miplearn/components/lazy_dynamic.py        |  6 ++--
 miplearn/components/steps/convert_tight.py |  6 ++--
 tests/components/test_lazy_dynamic.py      | 40 +++++++++++++++++-----
 tests/solvers/test_learning_solver.py      |  3 +-
 5 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/miplearn/classifiers/adaptive.py b/miplearn/classifiers/adaptive.py
index 8445d1f..94d753d 100644
--- a/miplearn/classifiers/adaptive.py
+++ b/miplearn/classifiers/adaptive.py
@@ -90,6 +90,13 @@ class AdaptiveClassifier(Classifier):
         n_samples = x_train.shape[0]
         assert y_train.shape == (n_samples, 2)
 
+        # If almost all samples belong to the same class, return a fixed prediction and
+        # skip all the other steps.
+        if y_train[:, 0].mean() > 0.999 or y_train[:, 1].mean() > 0.999:
+            self.classifier = CountingClassifier()
+            self.classifier.fit(x_train, y_train)
+            return
+
         best_name, best_clf, best_score = None, None, -float("inf")
         for (name, specs) in self.candidates.items():
             if n_samples < specs.min_samples:
diff --git a/miplearn/components/lazy_dynamic.py b/miplearn/components/lazy_dynamic.py
index be823a8..3236e04 100644
--- a/miplearn/components/lazy_dynamic.py
+++ b/miplearn/components/lazy_dynamic.py
@@ -85,8 +85,10 @@ class DynamicLazyConstraintsComponent(Component):
             disable=not sys.stdout.isatty(),
         ):
             logger.debug("Training: %s" % (str(v)))
-            label = np.zeros(len(training_instances))
-            label[violation_to_instance_idx[v]] = 1.0
+            label = [[True, False] for i in training_instances]
+            for idx in violation_to_instance_idx[v]:
+                label[idx] = [False, True]
+            label = np.array(label, dtype=np.bool8)
             classifier.fit(features, label)
 
     def predict(self, instance):
diff --git a/miplearn/components/steps/convert_tight.py b/miplearn/components/steps/convert_tight.py
index 7b25b4a..ee585db 100644
--- a/miplearn/components/steps/convert_tight.py
+++ b/miplearn/components/steps/convert_tight.py
@@ -116,9 +116,11 @@ class ConvertTightIneqsIntoEqsStep(Component):
                 if category not in y:
                     y[category] = []
                 if 0 <= slack <= self.slack_tolerance:
-                    y[category] += [[1]]
+                    y[category] += [[False, True]]
                 else:
-                    y[category] += [[0]]
+                    y[category] += [[True, False]]
+            for category in y.keys():
+                y[category] = np.array(y[category], dtype=np.bool8)
         return y
 
     def predict(self, x):
diff --git a/tests/components/test_lazy_dynamic.py b/tests/components/test_lazy_dynamic.py
index c7bf3ad..b2caf1d 100644
--- a/tests/components/test_lazy_dynamic.py
+++ b/tests/components/test_lazy_dynamic.py
@@ -6,6 +6,7 @@ from unittest.mock import Mock
 
 import numpy as np
 from numpy.linalg import norm
+from numpy.testing import assert_array_equal
 
 from miplearn.classifiers import Classifier
 from miplearn.components.lazy_dynamic import DynamicLazyConstraintsComponent
@@ -42,15 +43,36 @@ def test_lazy_fit():
     assert norm(expected_x_train_c - actual_x_train_c) < E
 
     # Should provide correct y_train to each classifier
-    expected_y_train_a = np.array([1.0, 0.0])
-    expected_y_train_b = np.array([1.0, 1.0])
-    expected_y_train_c = np.array([0.0, 1.0])
-    actual_y_train_a = component.classifiers["a"].fit.call_args[0][1]
-    actual_y_train_b = component.classifiers["b"].fit.call_args[0][1]
-    actual_y_train_c = component.classifiers["c"].fit.call_args[0][1]
-    assert norm(expected_y_train_a - actual_y_train_a) < E
-    assert norm(expected_y_train_b - actual_y_train_b) < E
-    assert norm(expected_y_train_c - actual_y_train_c) < E
+    expected_y_train_a = np.array(
+        [
+            [False, True],
+            [True, False],
+        ]
+    )
+    expected_y_train_b = np.array(
+        [
+            [False, True],
+            [False, True],
+        ]
+    )
+    expected_y_train_c = np.array(
+        [
+            [True, False],
+            [False, True],
+        ]
+    )
+    assert_array_equal(
+        component.classifiers["a"].fit.call_args[0][1],
+        expected_y_train_a,
+    )
+    assert_array_equal(
+        component.classifiers["b"].fit.call_args[0][1],
+        expected_y_train_b,
+    )
+    assert_array_equal(
+        component.classifiers["c"].fit.call_args[0][1],
+        expected_y_train_c,
+    )
 
 
 def test_lazy_before():
diff --git a/tests/solvers/test_learning_solver.py b/tests/solvers/test_learning_solver.py
index 983c1d1..fbb82da 100644
--- a/tests/solvers/test_learning_solver.py
+++ b/tests/solvers/test_learning_solver.py
@@ -3,6 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import logging
+import dill
 import pickle
 import tempfile
 import os
@@ -44,7 +45,7 @@ def test_learning_solver():
 
             # Assert solver is picklable
             with tempfile.TemporaryFile() as file:
-                pickle.dump(solver, file)
+                dill.dump(solver, file)
 
 
 def test_solve_without_lp():