Catch class balance errors and test L matrix edge cases (#1449)

Author: paroma

snorkel/labeling/model/label_model.py

@@ -539,12 +539,25 @@ def _set_class_balance(
         """
         if class_balance is not None:
             self.p = np.array(class_balance)
+            if len(self.p) != self.cardinality:
+                raise ValueError(
+                    f"class_balance has {len(self.p)} entries. Does not match LabelModel cardinality {self.cardinality}."
+                )
         elif Y_dev is not None:
             class_counts = Counter(Y_dev)
             sorted_counts = np.array([v for k, v in sorted(class_counts.items())])
             self.p = sorted_counts / sum(sorted_counts)
+            if len(self.p) != self.cardinality:
+                raise ValueError(
+                    f"Y_dev has {len(self.p)} class(es). Does not match LabelModel cardinality {self.cardinality}."
+                )
         else:
             self.p = (1 / self.cardinality) * np.ones(self.cardinality)
+
+        if np.any(self.p == 0):
+            raise ValueError(
+                f"Class balance prior is 0 for class(es) {np.where(self.p)[0]}."
+            )
         self.P = torch.diag(torch.from_numpy(self.p)).float()
 
     def _set_constants(self, L: np.ndarray) -> None:
@@ -742,8 +755,8 @@ def fit(
                 f"L_train has cardinality {L_shift.max()}, cardinality={self.cardinality} passed in."
             )
 
-        self._set_class_balance(class_balance, Y_dev)
         self._set_constants(L_shift)
+        self._set_class_balance(class_balance, Y_dev)
         self._create_tree()
         lf_analysis = LFAnalysis(L_train)
         self.coverage = lf_analysis.lf_coverages()

test/labeling/model/test_label_model.py

@@ -29,26 +29,56 @@ def _set_up_model(self, L: np.ndarray, class_balance: List[float] = [0.5, 0.5]):
 
     def test_L_form(self):
         label_model = LabelModel(cardinality=2, verbose=False)
-        L = np.array([[0, 1, 0], [0, 1, 0], [1, 0, 0], [0, 1, 0]])
+        L = np.array([[-1, 1, -1], [-1, 1, -1], [1, -1, -1], [-1, 1, -1]])
         label_model._set_constants(L)
         self.assertEqual(label_model.n, 4)
         self.assertEqual(label_model.m, 3)
 
-        L = np.array([[0, 1, 2], [0, 1, 2], [1, 0, 2], [0, 1, 0]])
+        L = np.array([[-1, 0, 1], [-1, 0, 2], [0, -1, 2], [-1, 0, -1]])
         with self.assertRaisesRegex(ValueError, "L_train has cardinality"):
             label_model.fit(L, n_epochs=1)
 
-        L = np.array([[0], [1], [-1]])
+        L = np.array([[0, 1], [1, 1], [0, 1]])
         with self.assertRaisesRegex(ValueError, "L_train should have at least 3"):
             label_model.fit(L, n_epochs=1)
 
+    def test_mv_default(self):
+        # less than 2 LFs have overlaps
+        label_model = LabelModel(cardinality=2, verbose=False)
+        L = np.array([[-1, -1, 1], [-1, 1, -1], [0, -1, -1]])
+        label_model.fit(L, n_epochs=100)
+        np.testing.assert_array_almost_equal(
+            label_model.predict(L), np.array([1, 1, 0])
+        )
+
+        # less than 2 LFs have conflicts
+        L = np.array([[-1, -1, 1], [-1, 1, 1], [1, 1, 1]])
+        label_model.fit(L, n_epochs=100)
+        np.testing.assert_array_almost_equal(
+            label_model.predict(L), np.array([1, 1, 1])
+        )
+
     def test_class_balance(self):
         label_model = LabelModel(cardinality=2, verbose=False)
         # Test class balance
         Y_dev = np.array([0, 0, 1, 1, 0, 0, 0, 0, 1, 1])
         label_model._set_class_balance(class_balance=None, Y_dev=Y_dev)
         np.testing.assert_array_almost_equal(label_model.p, np.array([0.6, 0.4]))
 
+        class_balance = np.array([0.0, 1.0])
+        with self.assertRaisesRegex(ValueError, "Class balance prior is 0"):
+            label_model._set_class_balance(class_balance=class_balance, Y_dev=Y_dev)
+
+        class_balance = np.array([0.0])
+        with self.assertRaisesRegex(ValueError, "class_balance has 1 entries."):
+            label_model._set_class_balance(class_balance=class_balance, Y_dev=Y_dev)
+
+        Y_dev_one_class = np.array([0, 0, 0])
+        with self.assertRaisesRegex(
+            ValueError, "Does not match LabelModel cardinality"
+        ):
+            label_model._set_class_balance(class_balance=None, Y_dev=Y_dev_one_class)
+
     def test_generate_O(self):
         L = np.array([[0, 1, 0], [0, 1, 0], [1, 0, 0], [0, 1, 1]])
         label_model = self._set_up_model(L)