feat: add custom correct column to database-grounded FDR control

tests/fdr/test_database_grounded.py

@@ -78,6 +78,21 @@ def test_fit_with_parameters(self, db_fdr_control):
         assert len(db_fdr_control.preds) == 1
         assert db_fdr_control.preds.iloc[0]["confidence"] == 0.9
 
+    def test_fit_uses_custom_correct_column(self, db_fdr_control):
+        """Test that fit() stores correctness under a custom column name."""
+        sample_df = pd.DataFrame(
+            {
+                "sequence": ["PEPTIDE", "PROTEIN"],
+                "prediction": [list("PEPTIDE"), list("PROTIEN")],
+                "confidence": [0.9, 0.8],
+            }
+        )
+
+        db_fdr_control.fit(sample_df, correct_column="proteome_hit")
+
+        assert list(db_fdr_control.preds.columns) == ["proteome_hit", "confidence"]
+        assert db_fdr_control.preds["proteome_hit"].tolist() == [True, False]
+
     def test_fit_with_empty_data(self, db_fdr_control):
         """Test that fit method handles empty data."""
         empty_data = pd.DataFrame()

winnow/fdr/database_grounded.py

@@ -34,6 +34,7 @@ def __init__(
     def fit(  # type: ignore
         self,
         dataset: pd.DataFrame,
+        correct_column: str = "correct",
     ) -> None:
         """Computes the precision-recall curve by comparing model predictions to database-grounded peptide sequences.
 
@@ -43,6 +44,8 @@ def fit(  # type: ignore
                 - 'sequence': Ground-truth peptide sequences.
                 - 'prediction': Model-predicted peptide sequences.
                 - Confidence column`confidence_feature` specified in the DatabaseGroundedFDRControl constructor.
+            correct_column (str):
+                Name of the column to store per-row correctness labels. Defaults to ``"correct"``.
         """
         assert len(dataset) > 0, "Fit method requires non-empty data"
 
@@ -57,18 +60,18 @@ def fit(  # type: ignore
             lambda row: self.metrics._novor_match(row["sequence"], row["prediction"]),
             axis=1,
         )
-        dataset["correct"] = dataset.apply(
+        dataset[correct_column] = dataset.apply(
             lambda row: (
                 row["num_matches"] == len(row["sequence"]) == len(row["prediction"])
             ),
             axis=1,
         )
-        self.preds = dataset[["correct", self.confidence_feature]]
+        self.preds = dataset[[correct_column, self.confidence_feature]]
 
         dataset = dataset.sort_values(
             by=self.confidence_feature, axis=0, ascending=False
         )
-        precision = np.cumsum(dataset["correct"]) / np.arange(1, len(dataset) + 1)
+        precision = np.cumsum(dataset[correct_column]) / np.arange(1, len(dataset) + 1)
         confidence = np.array(dataset[self.confidence_feature])
 
         self._fdr_values = np.array(1 - precision[self.drop :])