Create youden index metric

youden_index.py

@@ -11,85 +11,113 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""TODO: Add a description here."""
-
-import evaluate
+"""Youden index metric."""
 import datasets
+import evaluate
+import numpy as np
+from sklearn.metrics import (
+    precision_score,
+    roc_auc_score,
+    roc_curve,
+)
 
-
-# TODO: Add BibTeX citation
-_CITATION = """\
-@InProceedings{huggingface:module,
-title = {A great new module},
-authors={huggingface, Inc.},
-year={2020}
-}
-"""
-
-# TODO: Add description of the module here
-_DESCRIPTION = """\
-This new module is designed to solve this great ML task and is crafted with a lot of care.
+_DESCRIPTION = """
+This metric computes the Youden index based on the area under the curve (AUC) for the Receiver Operating Characteristic Curve (ROC). 
+The return values represent the ideal point in the ROC curve, where max(TPR - FPR) holds true (across all points in the curve).
+This metric only works with binary labels: The case in which there are only two different label classes, and each example gets only one label.
 """
 
-
-# TODO: Add description of the arguments of the module here
 _KWARGS_DESCRIPTION = """
-Calculates how good are predictions given some references, using certain scores
 Args:
-    predictions: list of predictions to score. Each predictions
-        should be a string with tokens separated by spaces.
-    references: list of reference for each prediction. Each
-        reference should be a string with tokens separated by spaces.
+- references (array-like of shape (n_samples,)): Ground truth labels.
+    - binary: expects an array-like of shape (n_samples,)
+- prediction_scores (array-like of shape (n_samples,)): Model predictions, probailities of the positive class.
+    - binary: expects an array-like of shape (n_samples,)
 Returns:
-    accuracy: description of the first score,
-    another_score: description of the second score,
-Examples:
-    Examples should be written in doctest format, and should illustrate how
-    to use the function.
-
-    >>> my_new_module = evaluate.load("my_new_module")
-    >>> results = my_new_module.compute(references=[0, 1], predictions=[0, 1])
-    >>> print(results)
-    {'accuracy': 1.0}
+     Dict[str, float]: Returns threshold (for classification), sensitivity and specificity at the optimal Youden index.
 """
 
-# TODO: Define external resources urls if needed
-BAD_WORDS_URL = "http://url/to/external/resource/bad_words.txt"
+_CITATION = """\
+@article{youden1950index,
+  title={Index for rating diagnostic tests},
+  author={Youden, William J},
+  journal={Cancer},
+  volume={3},
+  number={1},
+  pages={32--35},
+  year={1950},
+  publisher={Wiley Online Library}
+}
+@article{fluss2005estimation,
+  title={Estimation of the Youden Index and its associated cutoff point},
+  author={Fluss, Ronen and Faraggi, David and Reiser, Benjamin},
+  journal={Biometrical Journal: Journal of Mathematical Methods in Biosciences},
+  volume={47},
+  number={4},
+  pages={458--472},
+  year={2005},
+  publisher={Wiley Online Library}
+}
+@article{scikit-learn,
+title={Scikit-learn: Machine Learning in {P}ython},
+author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V. and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P. and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
+journal={Journal of Machine Learning Research},
+volume={12},
+pages={2825--2830},
+year={2011}
+}
+"""
 
 
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
-class youden_index(evaluate.Metric):
-    """TODO: Short description of my evaluation module."""
-
+class YoudenIndex(evaluate.Metric):
     def _info(self):
-        # TODO: Specifies the evaluate.EvaluationModuleInfo object
         return evaluate.MetricInfo(
-            # This is the description that will appear on the modules page.
-            module_type="metric",
             description=_DESCRIPTION,
             citation=_CITATION,
             inputs_description=_KWARGS_DESCRIPTION,
-            # This defines the format of each prediction and reference
-            features=datasets.Features({
-                'predictions': datasets.Value('int64'),
-                'references': datasets.Value('int64'),
-            }),
-            # Homepage of the module for documentation
-            homepage="http://module.homepage",
-            # Additional links to the codebase or references
-            codebase_urls=["http://github.com/path/to/codebase/of/new_module"],
-            reference_urls=["http://path.to.reference.url/new_module"]
+            features=datasets.Features(
+                {
+                    "prediction_scores": datasets.Sequence(datasets.Value("float")),
+                    "references": datasets.Value("int32"),
+                }
+            ),
+            reference_urls=[
+                "https://en.wikipedia.org/wiki/Youden%27s_J_statistic",
+                "https://gist.github.com/twolodzko/4fae2980a1f15f8682d243808e5859bb",
+            ],
         )
 
-    def _download_and_prepare(self, dl_manager):
-        """Optional: download external resources useful to compute the scores"""
-        # TODO: Download external resources if needed
-        pass
+    def _compute(
+        self,
+        references,
+        prediction_scores,
+    ):
+        """Compute the Youden index to determine the optimal threshold for classification.
+
+        Args:
+            references (array-like of shape (n_samples,)): Binary ground truth labels.
+            prediction_scores (array-like of shape (n_samples,)): Model predictions, probailities of the positive class.
 
-    def _compute(self, predictions, references):
-        """Returns the scores"""
-        # TODO: Compute the different scores of the module
-        accuracy = sum(i == j for i, j in zip(predictions, references)) / len(predictions)
+        Returns:
+            Dict[str, float]: Returns threshold (for classification), sensitivity, specificity, NPV and PPV
+                at the optimal Youden index and the ROC AUC score.
+        """
+        # 1. Compute ROC AUC
+        roc_auc = roc_auc_score(references, prediction_scores)
+        # 2. Determine the optimal threshold based on the Youden index
+        fpr, tpr, thresholds = roc_curve(references, prediction_scores)
+        idx = np.argmax(tpr - fpr)
+        optimal_threshold = thresholds[idx]
+        # 3. Calculate PPV (precision) and NPV based on the optimal threshold
+        optimal_predictions = np.where(prediction_scores > optimal_threshold, 1, 0)
+        ppv = precision_score(references, optimal_predictions)
+        npv = precision_score(references, optimal_predictions, pos_label=0)
         return {
-            "accuracy": accuracy,
-        }
+            "youden_threshold": optimal_threshold,
+            "sensitivity": tpr[idx],
+            "specificity": 1 - fpr[idx],
+            "roc_auc": roc_auc,
+            "ppv": ppv,
+            "npv": npv,
+        }