diff --git a/ms2rescore/core.py b/ms2rescore/core.py
index 3f4d9cb..170f103 100644
--- a/ms2rescore/core.py
+++ b/ms2rescore/core.py
@@ -297,15 +297,14 @@ def _calculate_confidence(psm_list: PSMList) -> PSMList:
     )
     psm_df["is_target"] = ~psm_df["is_decoy"]
     lin_psm_data = LinearPsmDataset(
-        psms=psm_df[["index", "peptide", "score", "is_target"]],
+        psms=psm_df[["index", "peptide", "is_target"]],
         target_column="is_target",
         spectrum_columns="index",  # Use artificial index to allow multi-rank rescoring
         peptide_column="peptide",
-        feature_columns=["score"],
     )
 
     # Recalculate confidence
-    new_confidence = lin_psm_data.assign_confidence()
+    new_confidence = lin_psm_data.assign_confidence(scores=psm_list["score"])
 
     # Add new confidence estimations to PSMList
     add_psm_confidence(psm_list, new_confidence)
diff --git a/ms2rescore/report/utils.py b/ms2rescore/report/utils.py
index 272bdb8..069a641 100644
--- a/ms2rescore/report/utils.py
+++ b/ms2rescore/report/utils.py
@@ -51,35 +51,25 @@ def get_confidence_estimates(
             "was generated by MS²Rescore. Could not generate report."
         ) from e
 
+    score_after = psm_list["score"]
     peptide = (
         pd.Series(psm_list["peptidoform"]).astype(str).str.replace(r"(/\d+$)", "", n=1, regex=True)
     )
-    psms = pd.DataFrame(
-        {
-            "peptide": peptide,
-            "is_target": ~psm_list["is_decoy"],
-            "before": score_before,
-            "after": psm_list["score"],
-        }
-    ).reset_index()
-
+    psms = pd.DataFrame({"peptide": peptide, "is_target": ~psm_list["is_decoy"]}).reset_index()
+    lin_psm_dataset = LinearPsmDataset(
+        psms=psms,
+        target_column="is_target",
+        spectrum_columns="index",
+        peptide_column="peptide",
+    )
     if fasta_file:
         fasta = read_fasta(fasta_file)
+        lin_psm_dataset.add_proteins(fasta)
 
     confidence = dict()
-    for when in ["before", "after"]:
-        lin_psm_dataset = LinearPsmDataset(
-            psms=psms,
-            target_column="is_target",
-            spectrum_columns="index",
-            feature_columns=[when],
-            peptide_column="peptide",
-        )
-        if fasta_file:
-            lin_psm_dataset.add_proteins(fasta)
-
+    for when, scores in [("before", score_before), ("after", score_after)]:
         try:
-            confidence[when] = lin_psm_dataset.assign_confidence()
+            confidence[when] = lin_psm_dataset.assign_confidence(scores=scores)
         except RuntimeError:
             confidence[when] = None