From 6c5d62162993eeff0c23424185064db52dfbdcea Mon Sep 17 00:00:00 2001
From: Dilan Pathirana <59329744+dilpath@users.noreply.github.com>
Date: Mon, 6 Jan 2025 22:46:54 +0100
Subject: [PATCH] Compute criterion weights (by @vwiela) (#136)

---------

Co-authored-by: Vincent Wieland <v_wiela27@web.de>
---
 doc/analysis.rst             |  2 +-
 petab_select/analyze.py      | 35 +++++++++++++++++++++++++++++++++++
 test/analyze/test_analyze.py | 15 +++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/doc/analysis.rst b/doc/analysis.rst
index ee8aa46..bf238c2 100644
--- a/doc/analysis.rst
+++ b/doc/analysis.rst
@@ -7,5 +7,5 @@ The PEtab Select Python library provides some methods to help with this. Please
 See the Python API docs for the :class:`petab_select.Models` class, which provides some methods. In particular, :attr:`petab_select.Models.df` can be used
 to get a quick overview over all models, as a pandas dataframe.
 
-Additionally, see the Python API docs for the ``petab_select.analysis`` module, which contains some methods to subset and group models,
+Additionally, see the Python API docs for the :mod:`petab_select.analyze` module, which contains some methods to subset and group models,
 or compute "weights" (e.g. Akaike weights).
diff --git a/petab_select/analyze.py b/petab_select/analyze.py
index 9dad49b..a4188fc 100644
--- a/petab_select/analyze.py
+++ b/petab_select/analyze.py
@@ -3,6 +3,8 @@
 import warnings
 from collections.abc import Callable
 
+import numpy as np
+
 from .constants import Criterion
 from .model import Model, ModelHash, default_compare
 from .models import Models
@@ -12,6 +14,7 @@
     "group_by_predecessor_model",
     "group_by_iteration",
     "get_best_by_iteration",
+    "compute_weights",
 ]
 
 
@@ -159,3 +162,35 @@ def get_best_by_iteration(
         for iteration, iteration_models in iterations_models.items()
     }
     return best_by_iteration
+
+
+def compute_weights(
+    models: Models,
+    criterion: Criterion,
+    as_dict: bool = False,
+) -> list[float] | dict[ModelHash, float]:
+    """Compute criterion weights.
+
+    N.B.: regardless of the criterion, the formula used is the Akaike weights
+    formula, but with ``criterion`` values instead of the AIC.
+
+    Args:
+        models:
+            The models.
+        criterion:
+            The criterion.
+        as_dict:
+            Whether to return a dictionary, with model hashes for keys.
+
+    Returns:
+        The criterion weights.
+    """
+    relative_criterion_values = np.array(
+        models.get_criterion(criterion=criterion, relative=True)
+    )
+    weights = np.exp(-0.5 * relative_criterion_values)
+    weights /= weights.sum()
+    weights = weights.tolist()
+    if as_dict:
+        weights = dict(zip(models.hashes, weights, strict=False))
+    return weights
diff --git a/test/analyze/test_analyze.py b/test/analyze/test_analyze.py
index f37e601..d465f71 100644
--- a/test/analyze/test_analyze.py
+++ b/test/analyze/test_analyze.py
@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import numpy as np
 import pytest
 
 from petab_select import (
@@ -77,3 +78,17 @@ def test_relative_criterion_values(models: Models) -> None:
         for criterion_value in criterion_values
     ]
     assert test_value == expected_value
+
+
+def test_compute_weights(models: Models) -> None:
+    """Test ``analyze.compute_weights``."""
+    criterion_values = np.array(
+        models.get_criterion(criterion=Criterion.AIC, relative=True)
+    )
+    expected_weights = (
+        np.exp(-0.5 * criterion_values) / np.exp(-0.5 * criterion_values).sum()
+    )
+    test_weights = analyze.compute_weights(
+        models=models, criterion=Criterion.AIC
+    )
+    np.testing.assert_allclose(test_weights, expected_weights)