[WIP] ensemble update

README.md

@@ -142,14 +142,25 @@ $$
 $$
 -->
 
-Here, confidence is the model's (self-reported) confidence in its prediction, calculated as
+Here, confidence is the model's (self-reported) confidence in its prediction. Each model has its own
+decision threshold $t_{m_i}$, calibrated on the validation set (see below), and confidence measures
+how far the prediction sits from that threshold — scaled separately on each side, so that a
+maximally confident negative ($p = 0$) and a maximally confident positive ($p = 1$) both count 1:
 $
-\text{confidence}_c^{m_i} = 2|p_c^{m_i} - 0.5|
+\text{confidence}_c^{m_i} = \begin{cases}
+(t_{m_i} - p_c^{m_i}) / t_{m_i} & \text{if } p_c^{m_i} < t_{m_i} \\
+(p_c^{m_i} - t_{m_i}) / (1 - t_{m_i}) & \text{otherwise}
+\end{cases}
 $
-For example, if a model makes a positive prediction with $p_c^{m_i} = 0.55$, the confidence is $2|0.55 - 0.5| = 0.1$.
-One could say that the model is not very confident in its prediction and very close to switching to a negative prediction.
-If another model is very sure about its negative prediction with $p_c^{m_j} = 0.1$, the confidence is $2|0.1 - 0.5| = 0.8$.
-Therefore, if in doubt, we are more confident in the negative prediction.
+For example, for a model with $t_{m_i} = 0.5$ and a positive prediction of $p_c^{m_i} = 0.55$, the
+confidence is $(0.55 - 0.5)/0.5 = 0.1$. One could say that the model is not very confident in its
+prediction and very close to switching to a negative prediction. If another model is very sure about
+its negative prediction with $p_c^{m_j} = 0.1$ (and $t_{m_j} = 0.5$), the confidence is
+$(0.5 - 0.1)/0.5 = 0.8$. Therefore, if in doubt, we are more confident in the negative prediction.
+
+The two-sided scaling matters whenever a model's threshold is not 0.5: with $t_{m_i} = 0.2$, a
+negative prediction only has a range of $0.2$ to move in and a positive one a range of $0.8$, so
+without rescaling the positive side would systematically outweigh the negative side.
 
 Confidence can be disabled by the `use_confidence` parameter of the predict method (default: True).
 
@@ -159,6 +170,37 @@ model independently of a given class.
 on a validation set for each class. If the `ensemble_type` is set to `wmv-f1`, the trust is calculated as F1-score $^{6.25}$.
 If the `ensemble_type` is set to `mv` (the default), the trust is set to 1 for all models.
 
+#### Learned aggregation (`ltr` and `des`)
+
+Two further `ensemble_type`s replace the fixed voting rule by a model that is fitted on the
+validation split. Both restrict themselves to a candidate set (per molecule, the union of each
+base learner's top-`candidate_k` classes) and both emit the same net score as the voting
+ensembles, so inconsistency resolution and the decision threshold apply unchanged.
+
+- `ltr` — **learning to rank**, an adaptation of
+  [GOLabeler](https://doi.org/10.1093/bioinformatics/bty130): the base learner scores for a
+  (molecule, class) pair become the feature vector of a LambdaMART ranker (LightGBM) that ranks
+  ChEBI classes per molecule. Features are the raw base learner scores plus the number of covering
+  models and the max/mean/std over them; a global cutoff on the ranker score is calibrated on a
+  held-out 20% of the validation split.
+- `des` — **dynamic ensemble selection**, an adaptation of
+  [META-DES.H](https://arxiv.org/pdf/1811.01742): a `GaussianNB` meta-classifier estimates, per
+  (molecule, class, base learner), how competent that base learner is *for this molecule*, and only
+  the competent ones vote, weighted by that competence. Competence is described by the paper's five
+  meta-feature sets over two neighbourhoods — the `region_size` nearest molecules by Tanimoto
+  similarity on ECFP4, and the `profile_size` nearest output profiles. Because the neighbourhoods
+  are looked up at prediction time, calibration stores the reference predictions, labels and
+  fingerprints in the ensemble directory (~1 GB for a 20-model ensemble on ChEBI50).
+
+Both calibrate their hyperparameters by 5-fold cross-validation on the validation split, scoring
+macro-F1 on each held-out fold (the cutoff is tuned on a fold-internal dev set, so the reported
+score is not tuned on the fold it is measured on). Only the parameters that moved the result in
+previous experiments are searched: `candidate_k` for `ltr`, and `region_size` / `profile_size` /
+`vote` for `des`. The ranker's own tree hyperparameters, and `des`'s consensus and competence
+thresholds, sit on a plateau and are left at their published values. Passing any searched parameter
+to the constructor skips the search for it. Results are written to `hyperparameter_search.csv` and
+`best_hyperparameters.csv` in the ensemble directory, as for `wmv-f1`.
+
 ### Inconsistency resolution
 After a decision has been made for each class independently, the consistency of the predictions with regard to the ChEBI hierarchy
 and disjointness axioms is checked. This is
@@ -173,3 +215,41 @@ both, we select one with the higher class score and set the other to 0.
 with a small change. For a pair of classes $A \subseteq B$ with predictions $1$ and $0$, instead of setting $B$ to $1$,
 we now set $A$ to $0$. This has the advantage that we cannot introduce new disjointness-inconsistencies and don't have
 to repeat step 2.
+
+#### Alternative methods
+
+The method above is `--inconsistency-resolution score-based` (`-ir`, the default). Two alternatives
+from the literature are available at the same point in the pipeline; all of them consume the net
+score and return a net score, so the decision threshold applies unchanged.
+
+- `ilr-godel`, `ilr-lukasiewicz` — **Iterative Local Refinement**
+  ([Daniele et al. 2023](https://doi.org/10.1007/s10994-023-06310-3)). Subsumption becomes the
+  implication $A \rightarrow B$ and disjointness the formula $\neg (A \wedge B)$, both as hard
+  constraints ($\hat t = 1$). Each constraint is repaired by its *minimal refinement function* —
+  the closest truth vector satisfying it — and the repairs are iterated to a fixpoint instead of
+  running the fixed 3-step schedule above. The two variants differ in how they split a violation:
+  Gödel is winner-take-all (it raises the parent to the child, and zeroes the weaker side of a
+  disjoint pair), whereas Łukasiewicz shares the correction — a disjointness violation with scores
+  $0.8$ and $0.7$ becomes $0.55$ and $0.45$ rather than $0.8$ and $0$.
+- `hex` — **HEX graphs**
+  ([Deng et al. 2014](https://doi.org/10.1007/978-3-319-10590-1_4)). A CRF over binary label
+  vectors in which hierarchy edges forbid $(B, A) = (0, 1)$ and exclusion edges forbid
+  $(1, 1)$. Illegal states have probability zero, so the marginals satisfy
+  $P(A) \le P(B)$ for $A \subseteq B$ and $P(A) + P(B) \le 1$ for disjoint $A, B$ by construction.
+
+Both convert the net score to a probability with $p = \sigma(k \cdot \mathrm{score})$ and back with
+$\mathrm{logit}(p) / k$, so the decision boundary stays at $0$. `k` (and `delta` for `hex`) can be
+passed with `-irp k=2.0` and tuned with `scripts/calibrate_resolution.py`. Note that `k` does not
+affect `ilr-godel`'s decisions: every Gödel operation is order-preserving, so a monotone
+reparametrisation cannot change the sign of any score.
+
+Applied as published, HEX inference is intractable here. Its cost is bounded by
+$O(\min(|V|2^w, |V|2^{\Omega}))$, and on a 2117-class ChEBI label set the maximum overlap is
+$\Omega = 2115$ and the junction tree width is $\le 62$, with over 5 million legal states in the
+largest cliques — the paper's efficiency argument assumes labels are mostly mutually exclusive,
+whereas ChEBI labels overwhelmingly overlap (~25 classes hold per molecule). This implementation
+therefore *clamps*: classes whose score is further than `delta` from the boundary, and which are
+not involved in a violation, are fixed to their sign; that assignment is propagated to a fixpoint;
+and exact inference runs only on the connected components of what remains (typically fewer than 30
+classes). Components above `max_component_size` fall back to `score-based`, counted in
+`n_fallbacks`. This is a deviation from the published method and should be reported as such.

chebifier/__init__.py

@@ -2,10 +2,10 @@
 # even if multiple subpackages are imported later.
 
 from ._custom_cache import PerSmilesPerModelLRUCache, modelwise_smiles_lru_cache
-from .ensemble.base_ensemble import BaseEnsemble
+from .ensemble.voting_ensemble import VotingEnsemble
 
 __all__ = [
-    "BaseEnsemble",
+    "VotingEnsemble",
     "PerSmilesPerModelLRUCache",
     "modelwise_smiles_lru_cache",
 ]

chebifier/build_ensemble.py

@@ -0,0 +1,85 @@
+import os
+
+import torch
+
+from chebifier.predict import (
+    base_learner_cache_path,
+    collect_base_learner_predictions,
+    load_dense_predictions,
+    save_dense_predictions,
+)
+
+
+class EnsembleBuilder:
+    """
+    A class to build an ensemble model from base learners and validation data.
+
+    Attributes:
+        base_learners (dict[str, BasePredictor]): A dictionary of base learner models.
+        ensemble_model (BaseEnsemble): An instance of a BaseEnsemble model.
+        validation_data (list[Chem.Mol]): Validation data for calibration.
+        validation_labels (pd.DataFrame): Validation labels for calibration, one column per class.
+            The column names define the label set the base learner predictions are mapped onto.
+        prediction_cache_dir (str): Directory to cache predictions.
+    """
+
+    def __init__(
+        self,
+        base_learners,
+        ensemble_model,
+        validation_data,
+        validation_labels,
+        prediction_cache_dir,
+    ):
+        self.base_learners = base_learners
+        self.ensemble_model = ensemble_model
+        self.validation_data = validation_data
+        self.validation_labels = validation_labels
+        self.prediction_cache_dir = prediction_cache_dir
+        os.makedirs(self.prediction_cache_dir, exist_ok=True)
+
+    def build_ensemble(self):
+        """
+        Build an ensemble model from base learners and validation data.
+
+        Base learner predictions are cached to avoid recomputation.
+        """
+
+        # Step 1: Get predictions from base learners on validation data
+        validation_predictions = {}
+        classes = {}
+        # get cached predictions if available, otherwise compute and cache them
+        for model_name, model in self.base_learners.items():
+            cache_path = base_learner_cache_path(
+                self.prediction_cache_dir, model_name, "validation"
+            )
+            if os.path.exists(cache_path):
+                print(f"{model_name} validation predictions found in cache, loading...")
+                validation_predictions[model_name] = load_dense_predictions(cache_path)
+            else:
+                print(f"Computing {model_name} validation predictions...")
+                validation_predictions[model_name] = model.predict_dense(
+                    self.validation_data
+                )
+                save_dense_predictions(cache_path, *validation_predictions[model_name])
+
+        # Base learners may be trained on different label sets (e.g. ChEBI25 vs. ChEBI25_3_STAR),
+        # so their union does not match the labels we calibrate against. Map every base learner
+        # onto the label set of the validation data instead.
+        label_classes = [str(cls) for cls in self.validation_labels.columns]
+        validation_predictions, classes = collect_base_learner_predictions(
+            validation_predictions, classes=label_classes
+        )
+        validation_labels = torch.from_numpy(
+            self.validation_labels.to_numpy(dtype=bool)
+        )
+
+        print(
+            f"Collected validation predictions from {len(validation_predictions)} base learners with {len(classes)} unique classes. Calibrating ensemble model..."
+        )
+        # Step 2: Calibrate the ensemble model using validation predictions
+        self.ensemble_model.calibrate(
+            validation_predictions, self.validation_data, validation_labels
+        )
+
+        return self.ensemble_model