diff --git a/README.md b/README.md index 8d75fcd..4c4b129 100644 --- a/README.md +++ b/README.md @@ -10,8 +10,9 @@ Not all models can be installed automatically at the moment: - `chebai-graph` and its dependencies. To install them, follow the instructions in the [chebai-graph repository](https://github.com/ChEB-AI/python-chebai-graph). - `chemlog-extra` can be installed with `pip install git+https://github.com/ChEB-AI/chemlog-extra.git` -- The automatically installed version of `c3p` may not work under Windows. If you want to run chebifier on Windows, we -recommend using this forked version: `pip install git+https://github.com/sfluegel05/c3p.git` +- `c3p` reads its generated programs assuming a UTF-8 locale and guards each of them with a +SIGALRM-based timeout, neither of which holds on Windows. The `c3p` predictor works around both +(see `_patch_c3p`), at the price of running the programs without a timeout there. You can get the package from PyPI: @@ -142,22 +143,103 @@ $$ $$ --> -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). +Confidence is used by the weighted voting ensembles (`wmv-conf` and `wmv-f1`). If the `ensemble_type` +is set to `mv`, all votes count the same (confidence is fixed to 1), which gives an unweighted +majority-voting baseline. The`model_weight` can be set for each model in the configuration file (default: 1). This is used to favor a certain model independently of a given class. `Trust` is based on the model's performance on a validation set. After training, we evaluate the Machine Learning models 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. +For `mv` and `wmv-conf`, 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. Feature column *j* is always base learner *j*, so the + ranker can learn which model to trust — but the raw scores say nothing about the class being + scored. `class_stats` (on by default) adds that: one column per base learner holding its + validation F1 *for this class* (the same quantity `wmv-f1` weights by), plus the class prevalence + and its number of positives. To keep the labels of the scored molecules out of the features, the + statistics used during training are estimated on the training molecules only, while prediction + uses the statistics of the whole validation split. Set `class_stats=False` for the plain + GOLabeler feature set; that also skips the per-model threshold calibration the F1 scores need. +- `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). + The meta-features are otherwise purely behavioural — one meta-classifier is fitted over all + (molecule, class, base learner) rows pooled, and the paper's input identifies neither the base + learner nor the class, so competence is a function of local track record alone. `use_model_id` + (on by default) appends a one-hot encoding of the base learner, which lets the meta-classifier + express "model A is the stronger one here" instead of only "whichever model this is, it behaves + like *this*"; `use_model_id=False` restores the published feature set. + `meta_classifier="mlp"` replaces `GaussianNB` with a standardised two-layer `MLPClassifier`, + which drops the feature-independence assumption — a poor fit for these meta-features, since the + `region_size` correctness flags are strongly correlated with each other and with their own mean. + The MLP is fitted in one pass over the meta-training set rather than chunk-wise, which the + consensus filter keeps small (~130k rows for 8 base learners on ChEBI25 3-STAR); + `max_meta_samples` caps it if a larger ensemble overflows memory. + Two further options control the reference set rather than the meta-classifier. + `morgan_radius` / `morgan_bits` / `morgan_chirality` set the fingerprint the region of competence + is measured on. Plain ECFP4 cannot separate stereoisomers, which are distinct ChEBI classes, so + 6.6% of ChEBI25 3-STAR validation molecules share a fingerprint with one carrying different + labels; `morgan_chirality` is therefore on by default, which halves that to 3.9%. Widening + `morgan_bits` changes nothing — the degeneracy is structural, not hash collisions. + `full_dsel=True` stores the whole validation split as the reference set instead of only the 80% + that the meta-classifier is fitted on, for denser neighbourhoods at prediction time. + + Note that the region of competence excludes the query molecule itself during calibration but + not during prediction, where the query is genuinely unseen. Predicting for the validation split + therefore lets ~80% of molecules retrieve themselves as their own nearest neighbour, which makes + any validation-split metric for `des` optimistic. Use the test split. + +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 — `chebifier build` takes constructor arguments as +`-ep key=value`, e.g. +`-ep candidate_k=70 -ep class_stats=1` or `-ep region_size=7 -ep meta_classifier=mlp`. Arguments +that change the stored model are recorded in the ensemble's metadata, so `chebifier evaluate` picks +them up on its own. `scripts/reproduce_ablation_3star.ps1` compares the optional features above +against their baselines this way. 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 @@ -173,3 +255,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. diff --git a/chebifier/__init__.py b/chebifier/__init__.py index a4f770c..a1c7314 100644 --- a/chebifier/__init__.py +++ b/chebifier/__init__.py @@ -2,10 +2,16 @@ # 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 ( + MajorityVotingEnsemble, + VotingEnsemble, + WMVwithConfidenceEnsemble, +) __all__ = [ - "BaseEnsemble", + "VotingEnsemble", + "MajorityVotingEnsemble", + "WMVwithConfidenceEnsemble", "PerSmilesPerModelLRUCache", "modelwise_smiles_lru_cache", ] diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py new file mode 100644 index 0000000..06b5c42 --- /dev/null +++ b/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 diff --git a/chebifier/cli.py b/chebifier/cli.py index 2f0468c..e4a8288 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -1,6 +1,211 @@ +import importlib.resources +import os +from typing import Literal + import click +import numpy as np +import pandas as pd +import yaml +from chebi_utils.read_molecule import smiles_or_inchi_to_mol + +from chebifier.build_ensemble import EnsembleBuilder +from chebifier.check_env import check_package_installed +from chebifier.hugging_face import download_model_files +from chebifier.inconsistency_resolution import SMOOTHER_NAMES +from chebifier.model_registry import ENSEMBLES, MODEL_TYPES +from chebifier.predict import aggregate_predictions, base_learner_cache_path +from chebifier.predict import predict as predict_molecules +from chebifier.predict import resolve_and_decide +from chebifier.utils import ( + get_default_configs, + get_disjoint_files, + load_chebi_graph, + process_config, +) + + +def read_molecules(molecules, molecule_file): + """Collect molecules from CLI arguments and/or a file (one molecule per line) and convert them to RDKit mol objects.""" + raw_inputs = list(molecules) + if molecule_file: + with open(molecule_file, "r", encoding="utf-8") as f: + raw_inputs.extend([line.strip() for line in f if line.strip()]) + + return [smiles_or_inchi_to_mol(raw_input) for raw_input in raw_inputs] + + +def build_base_learners(ensemble_config, prediction_cache_dir=None, split=None): + """Instantiate the base learners described by an ensemble configuration file. + + If prediction_cache_dir and split are given, models whose predictions for that split are + already cached are not instantiated (their entry is None) - loading their checkpoints would + be a waste of time since the cached predictions are used instead. + """ + if ensemble_config is None: + config = get_default_configs() + else: + print(f"Loading ensemble configuration from {ensemble_config}") + with open(ensemble_config, "r") as f: + config = yaml.safe_load(f) + + with ( + importlib.resources.files("chebifier") + .joinpath("model_registry.yml") + .open("r") as f + ): + model_registry = yaml.safe_load(f) + + chebi_graph = None + base_learners = {} + for model_name, model_config in process_config(config, model_registry).items(): + if ( + prediction_cache_dir is not None + and split is not None + and os.path.exists( + base_learner_cache_path(prediction_cache_dir, model_name, split) + ) + ): + print(f"{model_name} {split} predictions found in cache, skipping model.") + base_learners[model_name] = None + continue + if chebi_graph is None: + chebi_graph = load_chebi_graph() + if "hugging_face" in model_config: + hugging_face_kwargs = download_model_files(model_config["hugging_face"]) + else: + hugging_face_kwargs = {} + if "package_name" in model_config: + check_package_installed(model_config["package_name"]) + base_learners[model_name] = MODEL_TYPES[model_config["type"]]( + model_name, + **model_config, + **hugging_face_kwargs, + chebi_graph=chebi_graph, + ) + return base_learners + + +def parse_ir_params(ir_param): + params = {} + for entry in ir_param: + if "=" not in entry: + raise click.BadParameter(f"Expected key=value, got '{entry}'") + key, value = entry.split("=", 1) + try: + params[key.strip()] = float(value) + except ValueError: + params[key.strip()] = value + return params + + +def parse_ensemble_params(ensemble_param): + """Parse key=value arguments into keyword arguments for the ensemble constructor. Unlike the + inconsistency resolution parameters, these are coerced to int where possible - passing a float + where the ensemble expects a count (e.g. region_size) fails deep inside numpy.""" + params = {} + for entry in ensemble_param: + if "=" not in entry: + raise click.BadParameter(f"Expected key=value, got '{entry}'") + key, value = entry.split("=", 1) + for cast in (int, float): + try: + value = cast(value) + break + except ValueError: + continue + params[key.strip()] = value + return params + + +def load_dataset(data_path, split: Literal["train", "validation", "test"]): + data_file = os.path.join(data_path, "data.pkl") + splits_file = os.path.join(data_path, "splits.csv") + if not os.path.exists(data_file) or not os.path.exists(splits_file): + raise FileNotFoundError( + f"Required dataset files not found. Expected to find 'data.pkl' and 'splits.csv' in the provided data path ({data_path})." + ) + data_df = pd.read_pickle(data_file) + splits_df = pd.read_csv(splits_file) + # merge dataframe on id column and filter by the specified split + splits_df["id"] = splits_df["id"].astype(str) + merged_df = data_df.merge(splits_df, left_on="chebi_id", right_on="id", how="inner") + merged_df = merged_df[merged_df["split"] == split].reset_index(drop=True) + + mol_list = merged_df["mol"].tolist() + + # extract labels from data_df: every column other than chebi_id/mol/id/split is a ChEBI class label + label_columns = [c for c in data_df.columns if c not in ("chebi_id", "mol")] + labels_df = merged_df[label_columns].astype(bool) + labels_df.columns = [str(c) for c in label_columns] + + print( + f"Loaded {len(mol_list)} molecules and {len(labels_df.columns)} labels for split '{split}' from {data_path}." + ) + + return mol_list, labels_df + + +def base_learner_options(command): + """Options shared by all commands that use base learners.""" + for option in reversed( + [ + click.option( + "--ensemble-config", + "-e", + type=click.Path(exists=True), + default=None, + help="Configuration file listing the base learners of the ensemble", + ), + click.option( + "--prediction-cache-dir", + type=click.Path(), + default=None, + help="Directory for caching base learner predictions", + ), + ] + ): + command = option(command) + return command + + +def ensemble_options(command): + """Options shared by all commands that use a single ensemble.""" + for option in reversed( + [ + click.option( + "--ensemble-type", + "-t", + type=click.Choice(ENSEMBLES.keys()), + default="wmv-f1", + help="Type of ensemble to use (default: Weighted Majority Voting with F1 weights)", + ), + click.option( + "--ensemble-dir", + "-d", + type=click.Path(), + required=True, + help="Directory where the calibration results of the ensemble are stored", + ), + ] + ): + command = option(command) + return base_learner_options(command) -from chebifier.model_registry import ENSEMBLES + +def data_options(command): + """Options shared by the commands that work on a ChEBI dataset split.""" + for option in reversed( + [ + click.option( + "--data-path", + type=str, + required=True, + help="Data source: local dataset directory or Hugging Face repo id", + ), + ] + ): + command = option(command) + return command @click.group() @@ -10,106 +215,282 @@ def cli(): @cli.command() +@ensemble_options +@data_options @click.option( - "--ensemble-config", - "-e", - type=click.Path(exists=True), - default=None, - help="Configuration file for ensemble models", + "--ensemble-param", + "-ep", + multiple=True, + help="Extra key=value argument for the ensemble constructor, e.g. -ep candidate_k=70 " + "(repeatable). Parameters that change the stored model are written to the ensemble's " + "metadata, so 'evaluate' does not need to be given them again.", ) -@click.option("--smiles", "-s", multiple=True, help="SMILES strings to predict") +def build( + ensemble_config, + ensemble_type, + ensemble_dir, + prediction_cache_dir, + data_path, + ensemble_param, +): + """Build (calibrate) an ensemble on the ChEBI validation set.""" + base_learners = build_base_learners( + ensemble_config, prediction_cache_dir=prediction_cache_dir, split="validation" + ) + ensemble_model = ENSEMBLES[ensemble_type]( + ensemble_dir, **parse_ensemble_params(ensemble_param) + ) + + # TODO: Hugging Face support + validation_data, validation_labels = load_dataset(data_path, split="validation") + + builder = EnsembleBuilder( + base_learners, + ensemble_model, + validation_data, + validation_labels, + prediction_cache_dir, + ) + builder.build_ensemble() + + +@cli.command() +@base_learner_options +@data_options @click.option( - "--smiles-file", - "-f", - type=click.Path(exists=True), - help="File containing SMILES strings (one per line)", + "--ensemble-type", + "-t", + type=click.Choice(ENSEMBLES.keys()), + multiple=True, + default=("wmv-f1",), + help="Type of ensemble to evaluate (repeatable, paired with --ensemble-dir)", ) @click.option( - "--output", - "-o", + "--ensemble-dir", + "-d", type=click.Path(), - help="Output file to save predictions (optional)", + multiple=True, + required=True, + help="Directory where the calibration results of the ensemble are stored (one per --ensemble-type)", ) @click.option( - "--ensemble-type", - "-t", - type=click.Choice(ENSEMBLES.keys()), - default="wmv-f1", - help="Type of ensemble to use (default: Weighted Majority Voting)", + "--resolve-inconsistencies/--no-resolve-inconsistencies", + default=True, + help="Resolve inconsistencies in the aggregated predictions (default: True). " + "--no-resolve-inconsistencies is equivalent to '-ir none'.", ) @click.option( - "--use-confidence", - "-c", - is_flag=True, - default=True, - help="Weight predictions based on how 'confident' a model is in its prediction (default: True)", + "--inconsistency-resolution", + "-ir", + type=click.Choice(SMOOTHER_NAMES + ["none"]), + multiple=True, + default=("score-based",), + help="Method used to resolve inconsistencies (repeatable, default: score-based). " + "All methods share the base learner predictions and the ensemble aggregation.", ) @click.option( - "--resolve-inconsistencies", - "-r", - is_flag=True, - default=True, - help="Resolve inconsistencies in predictions automatically (default: True)", + "--ir-param", + "-irp", + multiple=True, + help="Extra key=value parameter for the resolution method, e.g. -irp k=2.0 (repeatable)", +) +@click.option( + "--split", + type=click.Choice(["validation", "test"]), + default="test", + help="Dataset split to evaluate on (default: test)", ) @click.option( - "--verbose", - "-v", + "--skip-existing", is_flag=True, default=False, - help="Enable verbose output", + help="Skip ensemble / resolution combinations whose output file already exists", ) -def predict( +@click.option( + "--output", + "-o", + type=click.Path(), + default=None, + help="Output file for the ensemble predictions, only allowed for a single ensemble and " + "resolution method (default: /_predictions_.npz, " + "with 'noir' as method for '-ir none')", +) +def evaluate( ensemble_config, - smiles, - smiles_file, - output, ensemble_type, - use_confidence, - resolve_inconsistencies=True, - verbose=False, + ensemble_dir, + prediction_cache_dir, + data_path, + resolve_inconsistencies, + inconsistency_resolution, + ir_param, + split, + skip_existing, + output, ): - """Predict ChEBI classes for SMILES strings using an ensemble model.""" - - # Instantiate ensemble model - ensemble = ENSEMBLES[ensemble_type]( - ensemble_config, - resolve_inconsistencies=resolve_inconsistencies, - verbose_output=verbose, - use_confidence=use_confidence, - ) + """Store the predictions of one or more ensembles on a ChEBI dataset split.""" + if len(ensemble_type) != len(ensemble_dir): + raise click.BadParameter( + f"Got {len(ensemble_type)} --ensemble-type and {len(ensemble_dir)} --ensemble-dir " + f"values, expected one directory per ensemble type." + ) + variants = list(inconsistency_resolution) if resolve_inconsistencies else ["none"] + if output is not None and (len(ensemble_type) > 1 or len(variants) > 1): + raise click.BadParameter( + "--output can only be used with a single --ensemble-type and a single " + "--inconsistency-resolution." + ) - # Collect SMILES strings from arguments and/or file - smiles_list = list(smiles) - if smiles_file: - with open(smiles_file, "r") as f: - smiles_list.extend([line.strip() for line in f if line.strip()]) + def output_path(dir_, variant): + if output is not None: + return output + suffix = "noir" if variant == "none" else variant + return os.path.join(dir_, f"{split}_predictions_{suffix}.npz") - if not smiles_list: - click.echo("No SMILES strings provided. Use --smiles or --smiles-file options.") + jobs = {} + for type_, dir_ in zip(ensemble_type, ensemble_dir): + todo = [ + variant + for variant in variants + if not (skip_existing and os.path.exists(output_path(dir_, variant))) + ] + if todo: + jobs[(type_, dir_)] = todo + else: + print(f"All outputs for {type_} in {dir_} exist, skipping.") + if not jobs: return - # Make predictions - predictions = ensemble.predict_smiles_list(smiles_list) + base_learners = build_base_learners( + ensemble_config, prediction_cache_dir=prediction_cache_dir, split=split + ) + + # TODO: Hugging Face support + eval_data, eval_labels = load_dataset(data_path, split=split) - if output: - # save as json - import json + chebi_graph, disjoint_files = None, None + if any(variant != "none" for variant in variants): + chebi_graph = load_chebi_graph() + disjoint_files = get_disjoint_files() - with open(output, "w") as f: - json.dump( - {smiles: pred for smiles, pred in zip(smiles_list, predictions)}, - f, - indent=2, + ir_params = parse_ir_params(ir_param) + for (type_, dir_), todo in jobs.items(): + ensemble_model = ENSEMBLES[type_](dir_) + aggregated, predicted_classes = aggregate_predictions( + base_learners, + ensemble_model, + eval_data, + prediction_cache_dir=prediction_cache_dir, + classes=[str(cls) for cls in eval_labels.columns], + split=split, + ) + for variant in todo: + print(f"Resolving inconsistencies for {type_} with '{variant}'...") + predictions = resolve_and_decide( + aggregated, + predicted_classes, + inconsistency_resolution=variant, + inconsistency_resolution_params=ir_params, + chebi_graph=chebi_graph, + disjoint_files=disjoint_files, + decision_threshold=ensemble_model.decision_threshold, + ) + target = output_path(dir_, variant) + os.makedirs(os.path.dirname(os.path.abspath(target)), exist_ok=True) + np.savez_compressed( + target, + classes=np.array(predictions["predicted_classes"]), + scores=predictions["net_score"].numpy(), + decisions=predictions["class_decisions"].numpy(), + has_valid_predictions=predictions["has_valid_predictions"].numpy(), + decision_threshold=np.array(ensemble_model.decision_threshold), + ) + print( + f"Saved {type_} predictions for split '{split}' " + f"({predictions['class_decisions'].shape[0]} molecules, {len(predictions['predicted_classes'])} classes, " + f"{int(predictions['complete_failure'].sum())} molecules without any valid prediction) to {target}." ) - else: - # Print results - for i, (smiles, prediction) in enumerate(zip(smiles_list, predictions)): - click.echo(f"Result for: {smiles}") - if prediction: - click.echo(f" Predicted classes: {', '.join(map(str, prediction))}") - else: - click.echo(" No predictions") + +@cli.command() +@ensemble_options +@click.option( + "--molecules", "-m", multiple=True, help="SMILES or InChI strings to predict" +) +@click.option( + "--molecule-file", + "-f", + type=click.Path(exists=True), + default=None, + help="File containing SMILES or InChI strings (one per line)", +) +@click.option( + "--resolve-inconsistencies/--no-resolve-inconsistencies", + default=True, + help="Resolve inconsistencies in the aggregated predictions (default: True)", +) +@click.option( + "--inconsistency-resolution", + "-ir", + type=click.Choice(SMOOTHER_NAMES), + default="score-based", + help="Method used to resolve inconsistencies (default: score-based)", +) +@click.option( + "--output", + "-o", + type=click.Path(), + default=None, + help="Output file to save the predictions (optional)", +) +@click.option( + "--ir-param", + "-irp", + multiple=True, + help="Extra key=value parameter for the resolution method, e.g. -irp k=2.0 (repeatable)", +) +@click.option( + "--decision-threshold", + "-dt", + type=float, + default=0, + help="Threshold for classifying predictions (default: 0)", +) +def predict( + ensemble_config, + ensemble_type, + ensemble_dir, + prediction_cache_dir, + molecules, + molecule_file, + resolve_inconsistencies, + inconsistency_resolution, + ir_param, + decision_threshold, + output, +): + """Predict ChEBI classes for a list of SMILES / InChI strings.""" + molecules_list = read_molecules(molecules, molecule_file) + if not molecules_list: + click.echo("No molecules provided. Use --molecules or --molecule-file.") + return + + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + + predictions = predict_molecules( + base_learners, + ensemble_model, + molecules_list, + prediction_cache_dir=prediction_cache_dir, + resolve_inconsistencies=resolve_inconsistencies, + inconsistency_resolution=inconsistency_resolution, + inconsistency_resolution_params=parse_ir_params(ir_param), + decision_threshold=decision_threshold, + ) + + print(f"Predictions: {predictions}") + # TODO: turn the aggregated predictions into ChEBI classes per molecule, print them / save to output if __name__ == "__main__": diff --git a/chebifier/ensemble/base_ensemble.py b/chebifier/ensemble/base_ensemble.py index 59601f4..6dc5bf4 100644 --- a/chebifier/ensemble/base_ensemble.py +++ b/chebifier/ensemble/base_ensemble.py @@ -1,303 +1,57 @@ -import importlib -import time -from pathlib import Path -from typing import Union +import os import torch -import tqdm -import yaml - -from chebifier.check_env import check_package_installed -from chebifier.hugging_face import download_model_files -from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother -from chebifier.prediction_models.base_predictor import BasePredictor -from chebifier.utils import ( - get_default_configs, - get_disjoint_files, - load_chebi_graph, - process_config, -) +from rdkit import Chem class BaseEnsemble: - def __init__( - self, - model_configs: Union[str, Path, dict, None] = None, - resolve_inconsistencies: bool = True, - verbose_output: bool = False, - use_confidence: bool = True, - ): - # Deferred Import: To avoid circular import error - from chebifier.model_registry import MODEL_TYPES - - # Load configuration from YAML file - if not model_configs: - config = get_default_configs() - elif isinstance(model_configs, dict): - config = model_configs - else: - print(f"Loading ensemble configuration from {model_configs}") - with open(model_configs, "r") as f: - config = yaml.safe_load(f) + """Base class for ensemble models. + Each ensemble has to perform the following tasks: + 1. Calibration (e.g. calculating weights for WMV or fitting a meta-model) on validation data + 2. Prediction on test data (i.e., turning base learner predictions into aggregated predictions) - with ( - importlib.resources.files("chebifier") - .joinpath("model_registry.yml") - .open("r") as f - ): - model_registry = yaml.safe_load(f) + Each ensemble gets a directory where it can store its calibration results (e.g. weights for WMV or meta-model parameters). - processed_configs = process_config(config, model_registry) - self.verbose_output = verbose_output - self.use_confidence = use_confidence + Not part of the ensemble are + - getting predictions from base learners + - resolving inconsistencies in the aggregated predictions + """ - self.chebi_graph = load_chebi_graph() - self.disjoint_files = get_disjoint_files() + def __init__(self, ensemble_dir: str): + os.makedirs(ensemble_dir, exist_ok=True) + self.ensemble_dir = ensemble_dir - self.models = [] - self.positive_prediction_threshold = 0.5 - for model_name, model_config in processed_configs.items(): - model_cls = MODEL_TYPES[model_config["type"]] - if "hugging_face" in model_config: - hugging_face_kwargs = download_model_files(model_config["hugging_face"]) - else: - hugging_face_kwargs = {} - if "package_name" in model_config: - check_package_installed(model_config["package_name"]) + @property + def ensemble_name(self): + return self.__class__.__name__ - model_instance = model_cls( - model_name, - **model_config, - **hugging_face_kwargs, - chebi_graph=self.chebi_graph, - ) - assert isinstance(model_instance, BasePredictor) - self.models.append(model_instance) + @property + def decision_threshold(self): + """Probability a class has to exceed to be predicted. - if resolve_inconsistencies: - self.smoother = ScoreBasedPredictionSmoother( - self.chebi_graph, - label_names=None, - disjoint_files=self.disjoint_files, - verbose=self.verbose_output, - ) - else: - self.smoother = None - - def gather_predictions(self, smiles_list): - # get predictions from all models for the SMILES list - # order them alphabetically by label class - model_predictions = [] - predicted_classes = set() - for model in self.models: - model_predictions.append(model.predict_smiles_list(smiles_list)) - for logits_for_smiles in model_predictions[-1]: - if logits_for_smiles is not None: - for cls in logits_for_smiles: - predicted_classes.add(cls) - if self.verbose_output: - print(f"Sorting predictions from {len(model_predictions)} models...") - predicted_classes = sorted(list(predicted_classes)) - predicted_classes_dict = {cls: i for i, cls in enumerate(predicted_classes)} - ordered_logits = ( - torch.zeros(len(smiles_list), len(predicted_classes), len(self.models)) - * torch.nan - ) - for i, model_prediction in enumerate(model_predictions): - for j, logits_for_smiles in tqdm.tqdm( - enumerate(model_prediction), - total=len(model_prediction), - desc=f"Sorting predictions for {self.models[i].model_name}", - ): - if logits_for_smiles is not None: - for cls in logits_for_smiles: - ordered_logits[j, predicted_classes_dict[cls], i] = ( - logits_for_smiles[cls] - ) - - return ordered_logits, predicted_classes + Ensembles report probabilities, so the default is the neutral point: predict a class when + the ensemble believes it more likely than not. Ensembles that tune their own operating + point (to maximise macro-F1, say) override this. + """ + return 0.5 - def consolidate_predictions( + def calibrate( self, - predictions, - classwise_weights, - return_intermediate_results=False, - **kwargs, + validation_predictions: dict[str, torch.Tensor], + validation_data: list[Chem.Mol], + validation_labels: torch.Tensor, ): + """Calibrate the ensemble model using validation predictions and labels. + At the end, save the calibration results (e.g. weights for WMV or meta-model parameters) to self.ensemble_dir. """ - Aggregates predictions from multiple models using weighted majority voting. - Optimized version using tensor operations instead of for loops. - """ - num_smiles, num_classes, num_models = predictions.shape - - # Get predictions for all classes - valid_predictions = ~torch.isnan(predictions) - valid_counts = valid_predictions.sum(dim=2) # Sum over models dimension - - # Skip classes with no valid predictions - has_valid_predictions = valid_counts > 0 - - # Calculate positive and negative predictions for all classes at once - positive_mask = ( - predictions > self.positive_prediction_threshold - ) & valid_predictions - negative_mask = ( - predictions < self.positive_prediction_threshold - ) & valid_predictions - - # if use_confidence is passed in kwargs, it overrides the ensemble setting - use_confidence = kwargs.get("use_confidence", self.use_confidence) - if use_confidence: - confidence = 2 * torch.abs( - predictions.nan_to_num() - self.positive_prediction_threshold - ) - else: - confidence = torch.ones_like(predictions) - - # Extract positive and negative weights - pos_weights = classwise_weights[0] # Shape: (num_classes, num_models) - neg_weights = classwise_weights[1] # Shape: (num_classes, num_models) - - # Calculate weighted predictions using broadcasting - # predictions shape: (num_smiles, num_classes, num_models) - # weights shape: (num_classes, num_models) - positive_weighted = ( - positive_mask.float() * confidence * pos_weights.unsqueeze(0) - ) - negative_weighted = ( - negative_mask.float() * confidence * neg_weights.unsqueeze(0) - ) + pass - # Sum over models dimension - positive_sum = positive_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) - negative_sum = negative_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) - - # Determine which classes to include for each SMILES - net_score = positive_sum - negative_sum # Shape: (num_smiles, num_classes) - if return_intermediate_results: - return ( - net_score, - has_valid_predictions, - { - "positive_mask": positive_mask, - "negative_mask": negative_mask, - "confidence": confidence, - "positive_sum": positive_sum, - "negative_sum": negative_sum, - }, - ) - - return net_score, has_valid_predictions - - def apply_inconsistency_resolution( - self, net_score, class_names, has_valid_predictions + def predict( + self, + test_predictions: dict[str, torch.Tensor], + molecules: list[Chem.Mol] | None = None, ): - # Smooth predictions - start_time = time.perf_counter() - if self.smoother is not None: - self.smoother.set_label_names(class_names) - smooth_net_score = self.smoother(net_score) - class_decisions = ( - smooth_net_score > 0 - ) & has_valid_predictions # Shape: (num_smiles, num_classes) - else: - class_decisions = ( - net_score > 0 - ) & has_valid_predictions # Shape: (num_smiles, num_classes) - end_time = time.perf_counter() - if self.verbose_output: - print(f"Prediction smoothing took {end_time - start_time:.2f} seconds") - - complete_failure = torch.all(~has_valid_predictions, dim=1) - return class_decisions, complete_failure - - def calculate_classwise_weights(self, predicted_classes): - """No weights, simple majority voting""" - positive_weights = torch.ones(len(predicted_classes), len(self.models)) - negative_weights = torch.ones(len(predicted_classes), len(self.models)) - - return positive_weights, negative_weights - - def predict_smiles_list( - self, smiles_list, return_intermediate_results=False, **kwargs - ) -> list: - ordered_predictions, predicted_classes = self.gather_predictions(smiles_list) - if len(predicted_classes) == 0: - print("Warning: No classes have been predicted for the given SMILES list.") - predicted_classes = {cls: i for i, cls in enumerate(predicted_classes)} - - classwise_weights = self.calculate_classwise_weights(predicted_classes) - if return_intermediate_results: - net_score, has_valid_predictions, intermediate_results_dict = ( - self.consolidate_predictions( - ordered_predictions, - classwise_weights, - return_intermediate_results=return_intermediate_results, - ) - ) - else: - net_score, has_valid_predictions = self.consolidate_predictions( - ordered_predictions, classwise_weights - ) - class_decisions, is_failure = self.apply_inconsistency_resolution( - net_score, list(predicted_classes.keys()), has_valid_predictions - ) - - class_names = list(predicted_classes.keys()) - class_indices = {predicted_classes[cls]: cls for cls in class_names} - result = [ - ( - [ - class_indices[idx.item()] - for idx in torch.nonzero(i, as_tuple=True)[0] - ] - if not failure - else None - ) - for i, failure in zip(class_decisions, is_failure) - ] - if return_intermediate_results: - intermediate_results_dict["predicted_classes"] = predicted_classes - intermediate_results_dict["classwise_weights"] = classwise_weights - intermediate_results_dict["net_score"] = net_score - return result, intermediate_results_dict - - return result - - -if __name__ == "__main__": - ensemble = BaseEnsemble( - { - "resgated_0ps1g189": { - "type": "resgated", - "ckpt_path": "data/0ps1g189/epoch=122.ckpt", - "molecular_properties": [ - "chebai_graph.preprocessing.properties.AtomType", - "chebai_graph.preprocessing.properties.NumAtomBonds", - "chebai_graph.preprocessing.properties.AtomCharge", - "chebai_graph.preprocessing.properties.AtomAromaticity", - "chebai_graph.preprocessing.properties.AtomHybridization", - "chebai_graph.preprocessing.properties.AtomNumHs", - "chebai_graph.preprocessing.properties.BondType", - "chebai_graph.preprocessing.properties.BondInRing", - "chebai_graph.preprocessing.properties.BondAromaticity", - "chebai_graph.preprocessing.properties.RDKit2DNormalized", - ], - # "classwise_weights_path" : "../python-chebai/metrics_0ps1g189_80-10-10.json" - }, - "electra_14ko0zcf": { - "type": "electra", - "ckpt_path": "data/14ko0zcf/epoch=193.ckpt", - # "classwise_weights_path": "../python-chebai/metrics_electra_14ko0zcf_80-10-10.json", - }, - } - ) - r = ensemble.predict_smiles_list( - [ - "[NH3+]CCCC[C@H](NC(=O)[C@@H]([NH3+])CC([O-])=O)C([O-])=O", - "C[C@H](N)C(=O)NCC(O)=O#", - "", - ], - load_preds_if_possible=False, - ) - print(len(r), r[0]) + """Aggregate base learner predictions. `molecules` is only required by ensembles that + depend on the molecules themselves (e.g. dynamic selection, which looks up a region of + competence for each of them).""" + raise NotImplementedError() diff --git a/chebifier/ensemble/dynamic_selection_ensemble.py b/chebifier/ensemble/dynamic_selection_ensemble.py new file mode 100644 index 0000000..e49d14d --- /dev/null +++ b/chebifier/ensemble/dynamic_selection_ensemble.py @@ -0,0 +1,829 @@ +import json +import pickle +from pathlib import Path + +import numpy as np +import torch + +from chebifier.ensemble.level1 import ( + N_FOLDS, + POSITIVE_THRESHOLD, + RANDOM_SEED, + candidate_pairs, + coverage_of, + cv_folds, + dense_from_pairs, + holdout_split, + noncandidate_probability, + pair_scorer, + rescale_to_threshold, + save_hyperparameter_results, + select_candidates, + stack_predictions, + threshold_array, +) +from chebifier.ensemble.voting_ensemble import VotingEnsemble + +MORGAN_RADIUS = 2 +MORGAN_BITS = 2048 +# without chirality, ECFP4 gives stereoisomers - which are distinct ChEBI classes - the same +# fingerprint, so the region of competence can be anchored on a molecule of a different class +MORGAN_CHIRALITY = True +REGION_SIZE_GRID = (1, 7) +PROFILE_SIZE_GRID = (5, 9, 15) +VOTE_GRID = ("plain", "confidence") +CHUNK_SIZE = 512 +MAX_META_SAMPLES = 2_000_000 +MLP_HIDDEN_LAYERS = (64, 32) + + +def fingerprints( + molecules, radius=MORGAN_RADIUS, n_bits=MORGAN_BITS, chirality=MORGAN_CHIRALITY +): + from rdkit import Chem, RDLogger + from rdkit.Chem import rdFingerprintGenerator + + RDLogger.DisableLog("rdApp.*") + generator = rdFingerprintGenerator.GetMorganGenerator( + radius=radius, fpSize=n_bits, includeChirality=chirality + ) + out = np.zeros((len(molecules), n_bits), dtype=np.uint8) + ok = np.zeros(len(molecules), dtype=bool) + for i, molecule in enumerate(molecules): + if isinstance(molecule, str): + molecule = Chem.MolFromSmiles(molecule) + if molecule is None: + continue + out[i] = generator.GetFingerprintAsNumPy(molecule) + ok[i] = True + return out, ok + + +def tanimoto_knn( + query, reference, reference_ids, k, self_columns=None, chunk_size=1024 +): + query = query.astype(np.float32, copy=False) + reference = reference.astype(np.float32, copy=False) + query_bits = query.sum(axis=1) + reference_bits = reference.sum(axis=1) + k = min(k, reference.shape[0]) + out = np.zeros((query.shape[0], k), dtype=np.int32) + for start in range(0, query.shape[0], chunk_size): + end = min(start + chunk_size, query.shape[0]) + intersection = query[start:end] @ reference.T + union = query_bits[start:end, None] + reference_bits[None, :] - intersection + block = intersection / np.maximum(union, 1e-6) + if self_columns is not None: + rows = np.arange(end - start) + columns = self_columns[start:end] + present = columns >= 0 + block[rows[present], columns[present]] = -1.0 + take = min(4 * k, block.shape[1] - 1) + near = np.argpartition(-block, take, axis=1)[:, : take + 1] + near_similarity = np.take_along_axis(block, near, axis=1) + near_ids = reference_ids[near] + order = np.lexsort((near_ids, -near_similarity), axis=1)[:, :k] + out[start:end] = np.take_along_axis(near_ids, order, axis=1) + return out + + +def region_of_competence( + query_fingerprints, + query_ok, + reference_fingerprints, + reference_ids, + k, + self_columns=None, +): + neighbours = tanimoto_knn( + query_fingerprints, reference_fingerprints, reference_ids, k, self_columns + ) + neighbours[~query_ok] = -1 + return neighbours + + +class Candidates: + def __init__(self, scores, k, thresholds): + mask = select_candidates(scores, k) + self.molecule, self.class_index, group = candidate_pairs(mask) + self.profiles = rescale_to_threshold( + scores[self.molecule, self.class_index], thresholds + ) + self.n_molecules = scores.shape[0] + self.pointer = np.zeros(self.n_molecules + 1, dtype=np.int64) + np.cumsum(group, out=self.pointer[1:]) + + def slice_of(self, start, end): + return int(self.pointer[start]), int(self.pointer[end]) + + def iter_chunks(self, chunk_size): + for start in range(0, self.n_molecules, chunk_size): + end = min(start + chunk_size, self.n_molecules) + yield start, end, self.slice_of(start, end) + + +class Dsel: + def __init__(self, scores, labels, thresholds): + self.scores = scores + self.labels = labels + self.thresholds = thresholds + + def gather(self, rows, class_index): + n_pairs, n_neighbours = rows.shape + flat_rows = np.maximum(rows, 0).ravel() + flat_classes = np.repeat(class_index, n_neighbours) + scores = rescale_to_threshold( + self.scores[flat_rows, flat_classes], self.thresholds + ) + labels = self.labels[flat_rows, flat_classes] + return ( + scores.reshape(n_pairs, n_neighbours, -1), + labels.reshape(n_pairs, n_neighbours), + ) + + +def output_profile_knn( + query, dsel, dsel_mask, coverage, kp, same_split, row_chunk=4096 +): + n_classes = coverage.shape[0] + out = np.full((len(query.class_index), kp), -1, dtype=np.int32) + query_order = np.argsort(query.class_index, kind="stable") + dsel_order = np.argsort(dsel.class_index, kind="stable") + query_pointer = np.searchsorted( + query.class_index[query_order], np.arange(n_classes + 1) + ) + dsel_pointer = np.searchsorted( + dsel.class_index[dsel_order], np.arange(n_classes + 1) + ) + + for class_idx in range(n_classes): + query_rows = query_order[ + query_pointer[class_idx] : query_pointer[class_idx + 1] + ] + if len(query_rows) == 0: + continue + dsel_rows = dsel_order[dsel_pointer[class_idx] : dsel_pointer[class_idx + 1]] + dsel_rows = dsel_rows[dsel_mask[dsel.molecule[dsel_rows]]] + if len(dsel_rows) == 0: + continue + columns = np.flatnonzero(coverage[class_idx]) + reference = np.nan_to_num( + dsel.profiles[dsel_rows][:, columns], nan=POSITIVE_THRESHOLD + ) + reference_ids = dsel.molecule[dsel_rows] + reference_square = (reference**2).sum(axis=1) + effective_kp = min(kp, len(dsel_rows) - 1 if same_split else len(dsel_rows)) + if effective_kp < 1: + continue + + for start in range(0, len(query_rows), row_chunk): + block = query_rows[start : start + row_chunk] + probe = np.nan_to_num( + query.profiles[block][:, columns], nan=POSITIVE_THRESHOLD + ) + distance = ( + (probe**2).sum(axis=1)[:, None] + + reference_square[None, :] + - 2 * probe @ reference.T + ) + if same_split: + distance[query.molecule[block][:, None] == reference_ids[None, :]] = ( + np.inf + ) + take = min(4 * effective_kp, distance.shape[1] - 1) + near = np.argpartition(distance, take, axis=1)[:, : take + 1] + near_distance = np.take_along_axis(distance, near, axis=1) + near_ids = reference_ids[near] + order = np.lexsort((near_ids, near_distance), axis=1)[:, :effective_kp] + out[block, :effective_kp] = np.take_along_axis(near_ids, order, axis=1) + if effective_kp < kp: + out[block, effective_kp:] = out[block, effective_kp - 1][:, None] + return out + + +class MetaChunk: + __slots__ = ("X", "covered", "profiles", "consensus", "alpha", "molecule") + + def __init__(self, X, covered, profiles, consensus, alpha, molecule): + self.X = X + self.covered = covered + self.profiles = profiles + self.consensus = consensus + self.alpha = alpha + self.molecule = molecule + + +def build_chunk( + candidates, + pair_slice, + region, + output_profiles, + dsel, + coverage, + labels=None, + model_id=False, +): + first, last = pair_slice + molecule = candidates.molecule[first:last] + class_index = candidates.class_index[first:last] + query_scores = candidates.profiles[first:last] + covered = coverage[class_index] + + neighbours = region[molecule] + neighbour_scores, neighbour_labels = dsel.gather(neighbours, class_index) + matches = (neighbour_scores > POSITIVE_THRESHOLD) == neighbour_labels[:, :, None] + f1 = matches & covered[:, None, :] + f2 = np.where( + neighbour_labels[:, :, None], neighbour_scores, 1.0 - neighbour_scores + ) + f3 = f1.sum(axis=1, dtype=np.float32) / region.shape[1] + + profile_scores, profile_labels = dsel.gather( + output_profiles[first:last], class_index + ) + f4 = ( + (profile_scores > POSITIVE_THRESHOLD) == profile_labels[:, :, None] + ) & covered[:, None, :] + + filled_query = np.nan_to_num(query_scores, nan=POSITIVE_THRESHOLD) + f5 = 2.0 * np.abs(filled_query - POSITIVE_THRESHOLD) + + blocks = [ + f1.transpose(0, 2, 1).astype(np.float32), + np.nan_to_num(f2, nan=POSITIVE_THRESHOLD).transpose(0, 2, 1), + f3[:, :, None], + f4.transpose(0, 2, 1).astype(np.float32), + f5[:, :, None], + ] + if model_id: + n_models = covered.shape[1] + blocks.append( + np.broadcast_to( + np.eye(n_models, dtype=np.float32), (len(molecule), n_models, n_models) + ) + ) + X = np.concatenate(blocks, axis=2) + + alpha = None + if labels is not None: + truth = labels[molecule, class_index] + alpha = ((query_scores > POSITIVE_THRESHOLD) == truth[:, None]) & covered + + positive = ((query_scores > POSITIVE_THRESHOLD) & covered).sum(axis=1) + n_covered = covered.sum(axis=1) + agreement = np.maximum(positive, n_covered - positive) / np.maximum(n_covered, 1) + + unusable = (neighbours < 0).any(axis=1) + if unusable.any(): + covered = covered.copy() + covered[unusable] = False + + return MetaChunk(X, covered, filled_query, agreement, alpha, molecule) + + +def aggregate(delta, chunk, competence_threshold, vote): + selected = (delta > competence_threshold) & chunk.covered + empty = ~selected.any(axis=1) + selected[empty] = chunk.covered[empty] + weight = delta * selected + direction = np.where(chunk.profiles > POSITIVE_THRESHOLD, 1.0, -1.0) + if vote == "confidence": + direction = direction * 2.0 * np.abs(chunk.profiles - POSITIVE_THRESHOLD) + total = weight.sum(axis=1) + agreement = (weight * direction).sum(axis=1) / np.maximum(total, 1e-6) + return np.clip(0.5 + agreement / 2, 0.0, 1.0).astype(np.float32) + + +class DynamicSelectionEnsemble(VotingEnsemble): + + def __init__( + self, + ensemble_dir: str, + candidate_k: int = 50, + region_size=None, + profile_size=None, + vote=None, + consensus_threshold: float = 0.7, + competence_threshold: float = 0.5, + region_size_grid=REGION_SIZE_GRID, + profile_size_grid=PROFILE_SIZE_GRID, + vote_grid=VOTE_GRID, + chunk_size: int = CHUNK_SIZE, + use_model_id: bool = True, + meta_classifier: str = "nb", + max_meta_samples: int = MAX_META_SAMPLES, + morgan_radius: int = MORGAN_RADIUS, + morgan_bits: int = MORGAN_BITS, + morgan_chirality: bool = MORGAN_CHIRALITY, + full_dsel: bool = False, + **kwargs, + ): + super().__init__(ensemble_dir) + self.morgan_radius = int(morgan_radius) + self.morgan_bits = int(morgan_bits) + self.morgan_chirality = bool(morgan_chirality) + self.full_dsel = bool(full_dsel) + if meta_classifier not in ("nb", "mlp"): + raise ValueError( + f"Unknown meta_classifier '{meta_classifier}', expected 'nb' or 'mlp'." + ) + self.use_model_id = bool(use_model_id) + self.meta_classifier = meta_classifier + self.max_meta_samples = int(max_meta_samples) + self.candidate_k = candidate_k + self.region_size = region_size + self.profile_size = profile_size + self.vote = vote + self.consensus_threshold = consensus_threshold + self.competence_threshold = competence_threshold + self.region_size_grid = tuple(region_size_grid) + self.profile_size_grid = tuple(profile_size_grid) + self.vote_grid = tuple(vote_grid) + self.chunk_size = chunk_size + self._classifier = None + self._metadata = None + self._dsel = None + self._dsel_fingerprints = None + self._coverage = None + self._thresholds = None + + @property + def _classifier_path(self): + return Path(self.ensemble_dir) / "des_meta_classifier.pkl" + + @property + def _dsel_path(self): + return Path(self.ensemble_dir) / "des_dsel.npz" + + @property + def _metadata_path(self): + return Path(self.ensemble_dir) / "des_metadata.json" + + def calibrate(self, validation_predictions, validation_data, validation_labels): + super().calibrate(validation_predictions, validation_data, validation_labels) + scores, model_names = stack_predictions( + validation_predictions, dtype=np.float16 + ) + thresholds = threshold_array(self._load_prediction_thresholds(), model_names) + labels = np.asarray(validation_labels, dtype=bool) + coverage = coverage_of(scores) + molecule_fingerprints, parsed = fingerprints( + validation_data, self.morgan_radius, self.morgan_bits, self.morgan_chirality + ) + if not parsed.all(): + print( + f"{int((~parsed).sum())} of {len(parsed)} validation molecules could not be " + "fingerprinted and are excluded from the dynamic selection reference set." + ) + candidates = Candidates(scores, self.candidate_k, thresholds) + dsel = Dsel(scores, labels, thresholds) + + best = None + if None in (self.region_size, self.profile_size, self.vote): + best = self._optimize_hyperparameters( + candidates, dsel, coverage, molecule_fingerprints, parsed, labels + ) + region_size = self.region_size or best["region_size"] + profile_size = self.profile_size or best["profile_size"] + vote = self.vote or best["vote"] + + dsel_mask, dev_mask = holdout_split(parsed) + region, profiles = self._neighbourhoods( + candidates, + molecule_fingerprints, + parsed, + dsel_mask, + coverage, + region_size, + profile_size, + ) + classifier = self._fit_meta_classifier( + candidates, region, profiles, dsel, coverage, dsel_mask, labels + ) + net = self._score( + candidates, region, profiles, dsel, coverage, classifier, dev_mask, [vote] + )[vote] + scorer, keep = pair_scorer( + labels, candidates.molecule, candidates.class_index, dev_mask + ) + tau, dev_macro_f1 = scorer.tune(net[keep]) + floor = noncandidate_probability( + scores.shape[0] * scores.shape[1] - candidates.molecule.size, + int(labels.sum()) + - int(labels[candidates.molecule, candidates.class_index].sum()), + ) + + # the meta-classifier and tau are fitted with the dev molecules held out of the reference + # set, but nothing stops prediction from looking neighbours up in all of them + reference_mask = parsed if self.full_dsel else dsel_mask + self._save( + classifier, + scores, + labels, + molecule_fingerprints, + reference_mask, + coverage, + { + "model_names": model_names, + "candidate_k": int(self.candidate_k), + "region_size": int(region_size), + "profile_size": int(profile_size), + "vote": vote, + "use_model_id": self.use_model_id, + "meta_classifier": self.meta_classifier, + "morgan_radius": self.morgan_radius, + "morgan_bits": self.morgan_bits, + "morgan_chirality": self.morgan_chirality, + "full_dsel": self.full_dsel, + "consensus_threshold": float(self.consensus_threshold), + "competence_threshold": float(self.competence_threshold), + "tau": float(tau), + "floor": floor, + "n_classes": int(scores.shape[1]), + "n_dsel": int(reference_mask.sum()), + "dev_macro_f1": float(dev_macro_f1), + }, + ) + print( + f"Saved meta-classifier to {self._classifier_path} (region_size={region_size}, " + f"profile_size={profile_size}, vote={vote}, meta_classifier={self.meta_classifier}, " + f"use_model_id={self.use_model_id}, tau={tau:.4f}, " + f"held-out macro-f1: {dev_macro_f1:.4f})." + ) + + def _neighbourhoods( + self, + candidates, + molecule_fingerprints, + parsed, + dsel_mask, + coverage, + region_size, + profile_size, + ): + reference_ids = np.flatnonzero(dsel_mask) + self_columns = np.full(len(parsed), -1, dtype=np.int64) + self_columns[reference_ids] = np.arange(len(reference_ids)) + region = region_of_competence( + molecule_fingerprints, + parsed, + molecule_fingerprints[reference_ids], + reference_ids, + region_size, + self_columns=self_columns, + ) + profiles = output_profile_knn( + candidates, candidates, dsel_mask, coverage, profile_size, same_split=True + ) + return region, profiles + + def _meta_samples( + self, candidates, region, profiles, dsel, coverage, molecule_mask, labels + ): + """The meta-training samples, one block of (features, is the base learner correct?) rows + per chunk of molecules. Only pairs the base learners disagree on are kept.""" + for start, end, pair_slice in candidates.iter_chunks(self.chunk_size): + block = molecule_mask[start:end] + if not block.any(): + continue + chunk = build_chunk( + candidates, + pair_slice, + region, + profiles, + dsel, + coverage, + labels=labels, + model_id=self.use_model_id, + ) + keep = block[chunk.molecule - start] & ( + chunk.consensus < self.consensus_threshold + ) + if not keep.any(): + continue + mask = chunk.covered & keep[:, None] + if not mask.any(): + continue + yield chunk.X[mask], chunk.alpha[mask] + + def _fit_meta_classifier( + self, candidates, region, profiles, dsel, coverage, molecule_mask, labels + ): + blocks = self._meta_samples( + candidates, region, profiles, dsel, coverage, molecule_mask, labels + ) + if self.meta_classifier == "nb": + from sklearn.naive_bayes import GaussianNB + + classifier = GaussianNB() + meta_classes = np.array([0, 1]) + seen = np.zeros(2, dtype=np.int64) + for X, alpha in blocks: + classifier.partial_fit(X, alpha, classes=meta_classes) + seen += np.bincount(alpha.astype(np.int64), minlength=2) + else: + X, alpha = [], [] + for block_X, block_alpha in blocks: + X.append(block_X) + alpha.append(block_alpha) + if not X: + X, alpha = [np.zeros((0, 0), dtype=np.float32)], [ + np.zeros(0, dtype=bool) + ] + X, alpha = np.concatenate(X), np.concatenate(alpha) + seen = np.bincount(alpha.astype(np.int64), minlength=2) + classifier = None + if seen.all(): + if len(X) > self.max_meta_samples: + print( + f"Subsampling {len(X)} meta-training samples to {self.max_meta_samples}." + ) + take = np.random.default_rng(RANDOM_SEED).choice( + len(X), size=self.max_meta_samples, replace=False + ) + X, alpha = X[take], alpha[take] + classifier = self._fit_mlp(X, alpha) + if not seen.all(): + raise RuntimeError( + f"Meta-training set is unusable ({seen[1]} correct / {seen[0]} incorrect base " + "learner predictions survived the consensus filter). Increase " + "consensus_threshold or check the base learner predictions." + ) + return classifier + + def _fit_mlp(self, X, alpha): + from sklearn.neural_network import MLPClassifier + from sklearn.pipeline import make_pipeline + from sklearn.preprocessing import StandardScaler + + print(f"Fitting MLP meta-classifier on {X.shape[0]} x {X.shape[1]} samples...") + classifier = make_pipeline( + StandardScaler(), + MLPClassifier( + hidden_layer_sizes=MLP_HIDDEN_LAYERS, + early_stopping=True, + n_iter_no_change=5, + max_iter=200, + random_state=RANDOM_SEED, + ), + ) + classifier.fit(X, alpha) + return classifier + + def _score( + self, + candidates, + region, + profiles, + dsel, + coverage, + classifier, + molecule_mask, + votes, + ): + nets = { + vote: np.zeros(len(candidates.class_index), dtype=np.float32) + for vote in votes + } + for start, end, pair_slice in candidates.iter_chunks(self.chunk_size): + if molecule_mask is not None and not molecule_mask[start:end].any(): + continue + chunk = build_chunk( + candidates, + pair_slice, + region, + profiles, + dsel, + coverage, + model_id=self.use_model_id, + ) + delta = np.zeros(chunk.covered.shape, dtype=np.float32) + if chunk.covered.any(): + delta[chunk.covered] = classifier.predict_proba(chunk.X[chunk.covered])[ + :, 1 + ].astype(np.float32) + first, last = pair_slice + for vote in votes: + nets[vote][first:last] = aggregate( + delta, chunk, self.competence_threshold, vote + ) + return nets + + def _optimize_hyperparameters( + self, candidates, dsel, coverage, molecule_fingerprints, parsed, labels + ): + print( + f"Optimizing region_size / profile_size / vote with {N_FOLDS}-fold " + "cross-validation on the validation set..." + ) + widest_region = max(self.region_size_grid) + widest_profile = max(self.profile_size_grid) + fold_scores = { + (region_size, profile_size, vote): [] + for region_size in self.region_size_grid + for profile_size in self.profile_size_grid + for vote in self.vote_grid + } + + for fold, test_idx in enumerate(cv_folds(candidates.n_molecules)): + print(f"Calibrating fold {fold + 1}/{N_FOLDS}...") + test_mask = np.zeros(candidates.n_molecules, dtype=bool) + test_mask[test_idx] = True + dsel_mask, dev_mask = holdout_split( + ~test_mask & parsed, seed=RANDOM_SEED + fold + ) + region, profiles = self._neighbourhoods( + candidates, + molecule_fingerprints, + parsed, + dsel_mask, + coverage, + widest_region, + widest_profile, + ) + dev_scorer, dev_keep = pair_scorer( + labels, candidates.molecule, candidates.class_index, dev_mask + ) + test_scorer, test_keep = pair_scorer( + labels, candidates.molecule, candidates.class_index, test_mask + ) + for region_size in self.region_size_grid: + for profile_size in self.profile_size_grid: + classifier = self._fit_meta_classifier( + candidates, + region[:, :region_size], + profiles[:, :profile_size], + dsel, + coverage, + dsel_mask, + labels, + ) + nets = self._score( + candidates, + region[:, :region_size], + profiles[:, :profile_size], + dsel, + coverage, + classifier, + dev_mask | test_mask, + self.vote_grid, + ) + for vote, net in nets.items(): + tau, _ = dev_scorer.tune(net[dev_keep]) + fold_scores[(region_size, profile_size, vote)].append( + test_scorer.macro_f1(net[test_keep], tau) + ) + + results = [] + for (region_size, profile_size, vote), scores in fold_scores.items(): + mean_score = float(np.mean(scores)) + results.append( + { + "region_size": region_size, + "profile_size": profile_size, + "vote": vote, + "mean_macro_f1": mean_score, + "std_macro_f1": float(np.std(scores)), + **{f"fold_{i}_macro_f1": s for i, s in enumerate(scores)}, + } + ) + print( + f"region_size={region_size}, profile_size={profile_size}, vote={vote}: " + f"macro-f1 {mean_score:.4f}" + ) + best = max(results, key=lambda result: result["mean_macro_f1"]) + save_hyperparameter_results( + self.ensemble_dir, + results, + { + "region_size": best["region_size"], + "profile_size": best["profile_size"], + "vote": best["vote"], + "mean_macro_f1": best["mean_macro_f1"], + }, + ) + return best + + def _save( + self, + classifier, + scores, + labels, + molecule_fingerprints, + dsel_mask, + coverage, + metadata, + ): + with open(self._classifier_path, "wb") as f: + pickle.dump(classifier, f) + rows = np.flatnonzero(dsel_mask) + np.savez_compressed( + self._dsel_path, + scores=scores[rows], + labels=labels[rows], + fingerprints=np.packbits(molecule_fingerprints[rows], axis=1), + coverage=coverage, + ) + with open(self._metadata_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=2) + + def _load(self): + if self._classifier is not None: + return + if not self._metadata_path.exists(): + raise FileNotFoundError( + f"No calibrated meta-classifier found in ensemble directory: {self.ensemble_dir}. " + "Please calibrate the ensemble first." + ) + with open(self._metadata_path, "r", encoding="utf-8") as f: + self._metadata = json.load(f) + if "floor" not in self._metadata: + raise ValueError( + f"The ensemble in {self.ensemble_dir} was calibrated before scores became " + "probabilities. Its stored tau is on the old scale and would reject every " + "prediction. Please re-run `chebifier build` for this ensemble." + ) + with open(self._classifier_path, "rb") as f: + self._classifier = pickle.load(f) + self._thresholds = threshold_array( + self._load_prediction_thresholds(), self._metadata["model_names"] + ) + # the stored reference fingerprints are of whatever kind calibration used, so the query + # fingerprints have to be built the same way rather than from today's defaults + self.morgan_radius = self._metadata["morgan_radius"] + self.morgan_bits = self._metadata["morgan_bits"] + self.morgan_chirality = self._metadata["morgan_chirality"] + with np.load(self._dsel_path) as data: + self._dsel = Dsel(data["scores"], data["labels"], self._thresholds) + self._dsel_fingerprints = np.unpackbits( + data["fingerprints"], axis=1, count=self.morgan_bits + ) + self._coverage = data["coverage"] + self.consensus_threshold = self._metadata["consensus_threshold"] + self.competence_threshold = self._metadata["competence_threshold"] + # the stored classifier was fit on a feature layout these two decide - restoring them is + # what lets `evaluate` load a variant without being told which one it is + self.use_model_id = self._metadata["use_model_id"] + self.meta_classifier = self._metadata["meta_classifier"] + + def predict(self, test_predictions, molecules=None): + if molecules is None: + raise ValueError( + f"{self.ensemble_name} needs the molecules it predicts for, to look up their " + "region of competence. Pass them to chebifier.predict.predict." + ) + self._load() + scores, _ = stack_predictions( + test_predictions, self._metadata["model_names"], dtype=np.float16 + ) + query_fingerprints, parsed = fingerprints( + molecules, self.morgan_radius, self.morgan_bits, self.morgan_chirality + ) + candidates = Candidates(scores, self._metadata["candidate_k"], self._thresholds) + dsel_candidates = Candidates( + self._dsel.scores, self._metadata["candidate_k"], self._thresholds + ) + reference_ids = np.arange(self._dsel.scores.shape[0]) + region = region_of_competence( + query_fingerprints, + parsed, + self._dsel_fingerprints, + reference_ids, + self._metadata["region_size"], + ) + profiles = output_profile_knn( + candidates, + dsel_candidates, + np.ones(dsel_candidates.n_molecules, dtype=bool), + self._coverage, + self._metadata["profile_size"], + same_split=False, + ) + vote = self._metadata["vote"] + net = self._score( + candidates, + region, + profiles, + self._dsel, + self._coverage, + self._classifier, + None, + [vote], + )[vote] + dense = dense_from_pairs( + candidates.molecule, + candidates.class_index, + net, + scores.shape[:2], + floor=self._metadata["floor"], + ) + return { + "net_score": torch.from_numpy(dense), + "has_valid_predictions": torch.from_numpy((~np.isnan(scores)).any(axis=2)), + } + + @property + def decision_threshold(self): + """Probability a candidate has to exceed, the operating point macro-F1 peaks at.""" + self._load() + return float(self._metadata["tau"]) diff --git a/chebifier/ensemble/learning_to_rank_ensemble.py b/chebifier/ensemble/learning_to_rank_ensemble.py new file mode 100644 index 0000000..17ae8e0 --- /dev/null +++ b/chebifier/ensemble/learning_to_rank_ensemble.py @@ -0,0 +1,333 @@ +import json +from pathlib import Path + +import numpy as np +import torch + +from chebifier.ensemble.level1 import ( + N_FOLDS, + RANDOM_SEED, + candidate_pairs, + class_statistics, + cv_folds, + dense_from_pairs, + fit_platt, + holdout_split, + noncandidate_probability, + pair_scorer, + save_hyperparameter_results, + select_candidates, + stack_predictions, + threshold_array, +) +from chebifier.ensemble.voting_ensemble import VotingEnsemble + +CANDIDATE_K_GRID = (30, 50, 70) +EARLY_STOPPING_ROUNDS = 30 +LGB_PARAMS = { + "objective": "lambdarank", + "metric": "ndcg", + "ndcg_eval_at": [10, 30], + "label_gain": [0, 1], + "lambdarank_truncation_level": 60, + "max_depth": 4, + "num_leaves": 15, + "learning_rate": 0.05, + "min_data_in_leaf": 50, + "feature_fraction": 0.9, + "bagging_fraction": 0.9, + "bagging_freq": 1, + "verbosity": -1, + "seed": RANDOM_SEED, +} + + +def build_rows(scores, k, labels=None): + mask = select_candidates(scores, k) + molecule, class_index, group = candidate_pairs(mask) + block = scores[molecule, class_index] + valid = ~np.isnan(block) + n_valid = valid.sum(axis=1).astype(np.float32) + covered = n_valid > 0 + denominator = np.maximum(n_valid, 1) + mean = np.where(valid, block, 0.0).sum(axis=1) / denominator + variance = ( + np.where(valid, (block - mean[:, None]) ** 2, 0.0).sum(axis=1) / denominator + ) + maximum = np.where(valid, block, -np.inf).max(axis=1) + aggregates = np.column_stack( + [ + n_valid, + np.where(covered, maximum, np.nan), + np.where(covered, mean, np.nan), + np.where(covered, np.sqrt(variance), np.nan), + ] + ) + rows = { + "X": np.column_stack([block, aggregates]).astype(np.float32), + "molecule": molecule, + "class_index": class_index, + "group": group, + "n_molecules": scores.shape[0], + } + if labels is not None: + rows["y"] = labels[molecule, class_index].astype(np.int8) + return rows + + +def _sigmoid(x): + return float(1.0 / (1.0 + np.exp(-x))) + + +def design_matrix(rows, class_stats=None): + if class_stats is None: + return rows["X"] + return np.column_stack([rows["X"], class_stats[rows["class_index"]]]) + + +class LearningToRankEnsemble(VotingEnsemble): + + def __init__( + self, + ensemble_dir: str, + candidate_k=None, + candidate_k_grid=CANDIDATE_K_GRID, + n_estimators: int = 500, + class_stats: bool = True, + **kwargs, + ): + super().__init__(ensemble_dir) + self.candidate_k = candidate_k + self.candidate_k_grid = tuple(candidate_k_grid) + self.n_estimators = n_estimators + self.class_stats = bool(class_stats) + self._booster = None + self._metadata = None + self._class_stats = None + + @property + def _model_path(self): + return Path(self.ensemble_dir) / "ltr_ranker.txt" + + @property + def _metadata_path(self): + return Path(self.ensemble_dir) / "ltr_metadata.json" + + @property + def _class_stats_path(self): + return Path(self.ensemble_dir) / "ltr_class_stats.npy" + + def calibrate(self, validation_predictions, validation_data, validation_labels): + if self.class_stats: + # fits the per-model prediction thresholds the class-wise F1 scores are based on + super().calibrate( + validation_predictions, validation_data, validation_labels + ) + else: + print( + f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." + ) + scores, model_names = stack_predictions(validation_predictions) + labels = np.asarray(validation_labels, dtype=bool) + thresholds = ( + threshold_array(self._load_prediction_thresholds(), model_names) + if self.class_stats + else None + ) + + candidate_k = self.candidate_k + best = None + if candidate_k is None: + candidate_k, best = self._optimize_candidate_k(scores, labels, thresholds) + + rows = build_rows(scores, candidate_k, labels) + train_mask, dev_mask = holdout_split(np.ones(scores.shape[0], dtype=bool)) + booster, tau, dev_macro_f1, _, dev_scores = self._fit( + rows, labels, train_mask, dev_mask, scores, thresholds + ) + platt_a, platt_b = fit_platt(dev_scores, rows["y"][dev_mask[rows["molecule"]]]) + floor = noncandidate_probability( + scores.shape[0] * scores.shape[1] - rows["molecule"].size, + int(labels.sum()) - int(rows["y"].sum()), + ) + + booster.save_model(str(self._model_path), num_iteration=booster.best_iteration) + if self.class_stats: + # the ranker is trained on class statistics of the training molecules only, but + # predicts with the statistics of the whole validation split - the same out-of-fold + # arrangement the weighted majority vote ensemble uses for its class-wise F1 weights + self._class_stats = class_statistics(scores, labels, thresholds) + np.save(self._class_stats_path, self._class_stats) + metadata = { + "model_names": model_names, + "candidate_k": int(candidate_k), + "class_stats": self.class_stats, + "tau": float(tau), + "platt_a": platt_a, + "platt_b": platt_b, + "tau_probability": _sigmoid(platt_a * float(tau) + platt_b), + "floor": floor, + "n_classes": int(scores.shape[1]), + "best_iteration": int(booster.best_iteration), + "dev_macro_f1": float(dev_macro_f1), + "params": {key: value for key, value in LGB_PARAMS.items()}, + } + if best is not None: + metadata["cross_validated_macro_f1"] = best["mean_macro_f1"] + with open(self._metadata_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=2) + self._booster, self._metadata = booster, metadata + print( + f"Saved ranker to {self._model_path} (candidate_k={candidate_k}, tau={tau:.4f}, " + f"held-out macro-f1: {dev_macro_f1:.4f}). Calibrated to probabilities with " + f"a={platt_a:.4f}, b={platt_b:.4f} (decision threshold " + f"{metadata['tau_probability']:.4f}); non-candidate pairs scored at {floor:.2e}." + ) + + def _fit(self, rows, labels, train_mask, dev_mask, scores=None, thresholds=None): + import lightgbm as lgb + + class_stats = None + if self.class_stats: + # estimated on the training molecules only, otherwise the features would carry the + # labels of the molecules the ranker is scored on + class_stats = class_statistics(scores, labels, thresholds, train_mask) + X = design_matrix(rows, class_stats) + train_rows = train_mask[rows["molecule"]] + dev_rows = dev_mask[rows["molecule"]] + train_set = lgb.Dataset( + X[train_rows], + label=rows["y"][train_rows], + group=rows["group"][train_mask], + free_raw_data=False, + ) + dev_set = lgb.Dataset( + X[dev_rows], + label=rows["y"][dev_rows], + group=rows["group"][dev_mask], + reference=train_set, + free_raw_data=False, + ) + booster = lgb.train( + LGB_PARAMS, + train_set, + num_boost_round=self.n_estimators, + valid_sets=[dev_set], + valid_names=["dev"], + callbacks=[lgb.early_stopping(EARLY_STOPPING_ROUNDS, verbose=False)], + ) + dev_scores = booster.predict( + X[dev_rows], num_iteration=booster.best_iteration + ).astype(np.float32) + scorer, _ = pair_scorer(labels, rows["molecule"], rows["class_index"], dev_mask) + tau, dev_macro_f1 = scorer.tune(dev_scores) + return booster, tau, dev_macro_f1, class_stats, dev_scores + + def _score_candidate_k(self, rows, labels, folds, scores=None, thresholds=None): + fold_scores = [] + for fold, test_idx in enumerate(folds): + test_mask = np.zeros(rows["n_molecules"], dtype=bool) + test_mask[test_idx] = True + train_mask, inner_dev_mask = holdout_split( + ~test_mask, seed=RANDOM_SEED + fold + ) + booster, tau, _, class_stats, _ = self._fit( + rows, labels, train_mask, inner_dev_mask, scores, thresholds + ) + test_rows = test_mask[rows["molecule"]] + net = booster.predict( + design_matrix(rows, class_stats)[test_rows], + num_iteration=booster.best_iteration, + ).astype(np.float32) + scorer, _ = pair_scorer( + labels, rows["molecule"], rows["class_index"], test_mask + ) + fold_scores.append(scorer.macro_f1(net, tau)) + return fold_scores + + def _optimize_candidate_k(self, scores, labels, thresholds=None): + print( + f"Optimizing candidate_k with {N_FOLDS}-fold cross-validation on the validation set..." + ) + folds = cv_folds(scores.shape[0]) + results = [] + for candidate_k in self.candidate_k_grid: + rows = build_rows(scores, candidate_k, labels) + recall = float( + labels[rows["molecule"], rows["class_index"]].sum() + / max(int(labels.sum()), 1) + ) + fold_scores = self._score_candidate_k( + rows, labels, folds, scores, thresholds + ) + mean_score = float(np.mean(fold_scores)) + results.append( + { + "candidate_k": candidate_k, + "candidate_recall": recall, + "mean_macro_f1": mean_score, + "std_macro_f1": float(np.std(fold_scores)), + **{f"fold_{i}_macro_f1": s for i, s in enumerate(fold_scores)}, + } + ) + print( + f"candidate_k={candidate_k} (candidate recall {recall:.4f}): macro-f1 {mean_score:.4f}" + ) + best = max(results, key=lambda result: result["mean_macro_f1"]) + save_hyperparameter_results( + self.ensemble_dir, + results, + { + "candidate_k": best["candidate_k"], + "mean_macro_f1": best["mean_macro_f1"], + }, + ) + return best["candidate_k"], best + + def _load(self): + if self._booster is not None: + return + import lightgbm as lgb + + if not self._metadata_path.exists(): + raise FileNotFoundError( + f"No calibrated ranker found in ensemble directory: {self.ensemble_dir}. " + "Please calibrate the ensemble first." + ) + with open(self._metadata_path, "r", encoding="utf-8") as f: + self._metadata = json.load(f) + self._booster = lgb.Booster(model_file=str(self._model_path)) + self.class_stats = self._metadata.get("class_stats", False) + if self.class_stats: + self._class_stats = np.load(self._class_stats_path) + + @property + def decision_threshold(self): + """Probability a candidate has to exceed, the operating point macro-F1 peaks at.""" + self._load() + return float(self._metadata["tau_probability"]) + + def predict(self, test_predictions, molecules=None): + self._load() + scores, _ = stack_predictions(test_predictions, self._metadata["model_names"]) + rows = build_rows(scores, self._metadata["candidate_k"]) + raw = self._booster.predict(design_matrix(rows, self._class_stats)).astype( + np.float32 + ) + # the ranker score is not a probability, so it goes through the calibration fitted during + # `calibrate`; the decision threshold is left to the caller so that inconsistency + # resolution sees the probabilities themselves + net = 1.0 / ( + 1.0 + np.exp(-(self._metadata["platt_a"] * raw + self._metadata["platt_b"])) + ) + dense = dense_from_pairs( + rows["molecule"], + rows["class_index"], + net.astype(np.float32), + scores.shape[:2], + floor=self._metadata["floor"], + ) + return { + "net_score": torch.from_numpy(dense), + "has_valid_predictions": torch.from_numpy((~np.isnan(scores)).any(axis=2)), + } diff --git a/chebifier/ensemble/level1.py b/chebifier/ensemble/level1.py new file mode 100644 index 0000000..61b91ac --- /dev/null +++ b/chebifier/ensemble/level1.py @@ -0,0 +1,213 @@ +from pathlib import Path + +import numpy as np +import pandas as pd + +N_FOLDS = 5 +RANDOM_SEED = 42 +POSITIVE_THRESHOLD = 0.5 +THRESHOLD_STEPS = 80 + + +def stack_predictions(predictions, model_names=None, dtype=np.float32): + if model_names is None: + model_names = list(predictions) + missing = [name for name in model_names if name not in predictions] + if missing: + raise ValueError( + "Predictions are missing for models the ensemble was calibrated on: " + + ", ".join(missing) + ) + stacked = np.stack( + [np.asarray(predictions[name], dtype=dtype) for name in model_names], axis=2 + ) + return stacked, list(model_names) + + +def rescale_to_threshold(scores, thresholds): + scores = np.asarray(scores, dtype=np.float32) + thresholds = np.asarray(thresholds, dtype=np.float32) + return np.where( + scores < thresholds, + POSITIVE_THRESHOLD * scores / thresholds, + POSITIVE_THRESHOLD + + POSITIVE_THRESHOLD * (scores - thresholds) / (1 - thresholds), + ) + + +def threshold_array(thresholds, model_names): + missing = [name for name in model_names if name not in thresholds] + if missing: + raise ValueError("Prediction thresholds are missing for: " + ", ".join(missing)) + return np.array([thresholds[name] for name in model_names], dtype=np.float32) + + +def coverage_of(scores): + return ~np.isnan(scores).all(axis=0) + + +def class_statistics(scores, labels, thresholds, molecule_mask=None, chunk_size=512): + """Per-class statistics of the base learners: one F1 score per (class, model), followed by the + prevalence and the number of positives of the class. NaN scores count as negative predictions, + matching the class-wise F1 scores of the weighted majority vote ensemble.""" + rows = ( + np.arange(labels.shape[0]) + if molecule_mask is None + else np.flatnonzero(molecule_mask) + ) + counts = np.zeros((3, scores.shape[1], scores.shape[2]), dtype=np.int64) + for start in range(0, len(rows), chunk_size): + block = rows[start : start + chunk_size] + predicted = scores[block] > thresholds + truth = labels[block][:, :, None] + counts[0] += (predicted & truth).sum(axis=0) + counts[1] += (predicted & ~truth).sum(axis=0) + counts[2] += (~predicted & truth).sum(axis=0) + tp, fp, fn = counts + f1 = 2 * tp / np.maximum(2 * tp + fp + fn, 1) + positives = labels[rows].sum(axis=0) + return np.column_stack([f1, positives / max(len(rows), 1), positives]).astype( + np.float32 + ) + + +def select_candidates(scores, k): + n_molecules, n_classes, n_models = scores.shape + k = min(k, n_classes) + mask = np.zeros((n_molecules, n_classes), dtype=bool) + for model_idx in range(n_models): + column = scores[:, :, model_idx] + column = np.where(np.isnan(column), -np.inf, column) + top = np.argpartition(-column, k - 1, axis=1)[:, :k] + np.put_along_axis(mask, top, True, axis=1) + return mask + + +def candidate_pairs(mask): + molecule, class_index = np.nonzero(mask) + return ( + molecule.astype(np.int64), + class_index.astype(np.int64), + mask.sum(axis=1).astype(np.int64), + ) + + +def cv_folds(n_molecules, n_folds=N_FOLDS, seed=0): + permutation = np.random.default_rng(seed).permutation(n_molecules) + return [permutation[fold::n_folds] for fold in range(n_folds)] + + +def holdout_split(molecule_mask, dev_fraction=0.2, seed=RANDOM_SEED): + rows = np.flatnonzero(molecule_mask) + n_dev = max(int(len(rows) * dev_fraction), 1) + dev = np.zeros(len(molecule_mask), dtype=bool) + dev[np.random.default_rng(seed).choice(rows, size=n_dev, replace=False)] = True + return molecule_mask & ~dev, dev + + +class PairScorer: + def __init__(self, labels, molecule, class_index): + self.n_molecules, self.n_classes = labels.shape + self.molecule = molecule + self.class_index = class_index + self.y = labels[molecule, class_index] + self.class_positives = labels.sum(axis=0).astype(np.int64) + + def _class_counts(self, predicted): + tp = np.bincount( + self.class_index[predicted & self.y], minlength=self.n_classes + ).astype(np.int64) + positives = np.bincount( + self.class_index[predicted], minlength=self.n_classes + ).astype(np.int64) + return tp, positives - tp + + def macro_f1(self, net, tau): + tp, fp = self._class_counts(net > tau) + fn = self.class_positives - tp + denominator = 2 * tp + fp + fn + return float( + np.where(denominator > 0, 2 * tp / np.maximum(denominator, 1), 0.0).mean() + ) + + def micro_f1(self, net, tau): + tp, fp = self._class_counts(net > tau) + tp, fp = int(tp.sum()), int(fp.sum()) + fn = int(self.class_positives.sum()) - tp + denominator = 2 * tp + fp + fn + return 2 * tp / denominator if denominator > 0 else 0.0 + + def tune(self, net, metric=None, n_steps=THRESHOLD_STEPS): + metric = self.macro_f1 if metric is None else metric + grid = np.quantile(net, np.linspace(0.02, 0.9995, n_steps)) + best_tau, best_score = float(grid[0]), -1.0 + for tau in grid: + score = metric(net, tau) + if score > best_score: + best_tau, best_score = float(tau), score + return best_tau, best_score + + +def pair_scorer(labels, molecule, class_index, molecule_mask): + keep = molecule_mask[molecule] + rows = np.flatnonzero(molecule_mask) + remap = np.full(len(molecule_mask), -1, dtype=np.int64) + remap[rows] = np.arange(len(rows)) + return PairScorer(labels[rows], remap[molecule[keep]], class_index[keep]), keep + + +def dense_from_pairs(molecule, class_index, net, shape, floor=None): + if floor is None: + floor = min(float(net.min()) - 1.0, -1.0) if net.size else -1.0 + dense = np.full(shape, floor, dtype=np.float32) + dense[molecule, class_index] = net + return dense + + +def fit_platt(scores, labels): + """Fit `P(positive) = sigmoid(a * score + b)`, turning a raw score into a probability. + + A lambdarank score is only trained to order pairs within a group, so its scale carries no + probabilistic meaning. Inconsistency resolution needs one: it compares scores across classes. + Being a monotone map, this leaves the ensemble's own thresholded decisions untouched - what it + buys is a scale on which those downstream comparisons are meaningful. + + The voting ensembles need no equivalent: their weighted agreement fraction is already a well + calibrated probability, and a logistic fit measurably degrades it. + """ + from sklearn.linear_model import LogisticRegression + + scores = np.asarray(scores, dtype=np.float64).reshape(-1, 1) + platt = LogisticRegression(C=1e6).fit(scores, np.asarray(labels)) + slope, intercept = float(platt.coef_[0, 0]), float(platt.intercept_[0]) + if slope <= 0: + raise ValueError( + f"Platt calibration found a non-positive slope ({slope:.4g}), meaning the ranker scores " + "anti-correlate with the labels. Refusing to calibrate - check the base learner " + "predictions and the ranker training." + ) + return slope, intercept + + +def noncandidate_probability(n_noncandidate, n_missed_positives): + """Probability that a pair outside the candidate set is nevertheless a positive. + + Candidate selection keeps only the top-k classes per model, so the pairs it drops are not + "unknown" - they are overwhelmingly true negatives, and how overwhelmingly is measurable on the + validation set. This turns the fill value for those pairs from an arbitrary floor into a + calibrated probability that can be compared against the candidates. The Jeffreys-style + pseudo-count keeps the result non-zero when no positive is missed at all. + """ + return float((n_missed_positives + 0.5) / (n_noncandidate + 1.0)) + + +def save_hyperparameter_results(ensemble_dir, results, best): + results_path = Path(ensemble_dir) / "hyperparameter_search.csv" + pd.DataFrame(results).to_csv(results_path, index=False) + best_path = Path(ensemble_dir) / "best_hyperparameters.csv" + pd.DataFrame([best]).to_csv(best_path, index=False) + print( + f"Saved hyperparameter search results to {results_path}. Recommended parameters " + f"(saved to {best_path}): " + + ", ".join(f"{key}={value}" for key, value in best.items()) + ) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py new file mode 100644 index 0000000..cca29a5 --- /dev/null +++ b/chebifier/ensemble/voting_ensemble.py @@ -0,0 +1,187 @@ +from pathlib import Path + +import torch +import yaml +from torchmetrics import F1Score + +from chebifier.ensemble.base_ensemble import BaseEnsemble + + +class VotingEnsemble(BaseEnsemble): + """Base class for the voting ensembles. Subclasses define the vote weights by overriding + calculate_confidence (self-reported confidence of a model) and calculate_trust (measured + reliability of a model).""" + + def __init__( + self, + ensemble_dir: str, + ): + super().__init__(ensemble_dir) + self.classwise_f1 = None + self.prediction_thresholds = None + + def find_best_threshold(self, predictions, val_labels_tensor): + best_threshold = 0.5 + best_f1 = 0.0 + for threshold in range(1, 100): + threshold_value = threshold / 100 + macro_f1_score = self.classwise_f1( + predictions > threshold_value, val_labels_tensor + ).mean() + if macro_f1_score > best_f1: + best_f1 = macro_f1_score.item() + best_threshold = threshold_value + return best_threshold + + def calibrate(self, validation_predictions, validation_data, validation_labels): + print( + f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." + ) + self.classwise_f1 = F1Score( + task="multilabel", num_labels=validation_labels.shape[1], average=None + ) + self._save_prediction_thresholds( + self._fit_prediction_thresholds(validation_predictions, validation_labels) + ) + + def _fit_prediction_thresholds(self, predictions, labels) -> dict[str, float]: + return { + model_name: self.find_best_threshold(model_predictions, labels) + for model_name, model_predictions in predictions.items() + } + + def _save_prediction_thresholds(self, thresholds: dict[str, float]): + thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" + with open(thresholds_path, "w+", encoding="utf-8") as f: + yaml.dump(thresholds, f) + print(f"Saved prediction thresholds to {thresholds_path}: {thresholds}") + + def _load_prediction_thresholds(self) -> dict[str, float]: + if self.prediction_thresholds is not None: + return self.prediction_thresholds + thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" + if thresholds_path.exists(): + with open(thresholds_path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + else: + raise FileNotFoundError( + f"Prediction thresholds file not found in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." + ) + + def calculate_confidence( + self, predictions_tensor: torch.Tensor, thresholds: torch.Tensor + ) -> torch.Tensor: + raise NotImplementedError + + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: + # No trust unless a subclass measures it (e.g. WMVwithF1Ensemble) + return 1 + + def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): + """ + Aggregates predictions from multiple models by voting, each vote weighted by + calculate_confidence * calculate_trust. + + The net score is the weighted agreement among the models that voted, mapped onto [0, 1]: + 1 if they unanimously predict the class, 0 if they unanimously reject it, 0.5 if they are + evenly split. Normalising by the weight mass that was cast (rather than summing over models) + keeps classes covered by different numbers of base learners comparable, which is what + inconsistency resolution needs - it compares scores across classes. The agreement fraction + is a well calibrated probability on its own, so no further calibration is applied. + """ + predictions_tensor = torch.stack( + list(test_predictions.values()), dim=2 + ) # Shape: (num_molecules, num_classes, num_models) + # Get predictions for all classes + valid_predictions = ~torch.isnan(predictions_tensor) + valid_counts = valid_predictions.sum(dim=2) # Sum over models dimension + + thresholds = self._load_prediction_thresholds() + if any(model_name not in thresholds for model_name in test_predictions.keys()): + raise ValueError( + "Prediction thresholds not found for all models. Please calibrate the ensemble first. Models missing thresholds: " + + ", ".join( + model_name + for model_name in test_predictions.keys() + if model_name not in thresholds + ) + ) + threshold_mask = torch.tensor( + [thresholds[model_name] for model_name in test_predictions.keys()], + dtype=predictions_tensor.dtype, + device=predictions_tensor.device, + ) + + # Skip classes with no valid predictions + has_valid_predictions = valid_counts > 0 + + # Calculate positive and negative predictions for all classes at once + positive_mask = ( + predictions_tensor > threshold_mask.unsqueeze(0).unsqueeze(0) + ) & valid_predictions + negative_mask = ( + predictions_tensor < threshold_mask.unsqueeze(0).unsqueeze(0) + ) & valid_predictions + + confidence = self.calculate_confidence( + predictions_tensor, threshold_mask.unsqueeze(0).unsqueeze(0) + ) + + trust = self.calculate_trust(test_predictions) + # Calculate weighted predictions using broadcasting + # predictions shape: (num_molecules, num_classes, num_models) + # weights shape: (num_classes, num_models) + positive_weighted = positive_mask.float() * confidence * trust + negative_weighted = negative_mask.float() * confidence * trust + + # Sum over models dimension + positive_sum = positive_weighted.sum( + dim=2 + ) # Shape: (num_molecules, num_classes) + negative_sum = negative_weighted.sum( + dim=2 + ) # Shape: (num_molecules, num_classes) + + # Determine which classes to include for each molecule + net_score = ( + 0.5 + + (positive_sum - negative_sum) + / (positive_sum + negative_sum).clamp(min=1e-6) + / 2 + ).clamp( + 0.0, 1.0 + ) # Shape: (num_molecules, num_classes) + return { + "net_score": net_score, + "has_valid_predictions": has_valid_predictions, + "positive_sum": positive_sum, + "negative_sum": negative_sum, + "confidence": confidence, + "positive_mask": positive_mask, + "negative_mask": negative_mask, + } + + +class MajorityVotingEnsemble(VotingEnsemble): + """Plain majority voting: every model that votes counts the same, no weights at all.""" + + def calculate_confidence( + self, predictions_tensor: torch.Tensor, thresholds: torch.Tensor + ) -> torch.Tensor: + return torch.ones_like(predictions_tensor) + + +class WMVwithConfidenceEnsemble(VotingEnsemble): + """WMV ensemble that weights each vote by the model's self-reported confidence, i.e. how far + its prediction sits from its decision threshold, scaled separately on each side so that a + maximally confident negative and a maximally confident positive both count 1.""" + + def calculate_confidence( + self, predictions_tensor: torch.Tensor, thresholds: torch.Tensor + ) -> torch.Tensor: + scores = predictions_tensor.nan_to_num() + return torch.where( + scores < thresholds, + (thresholds - scores) / thresholds, + (scores - thresholds) / (1 - thresholds), + ) diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 1353325..edd7b00 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,107 +1,244 @@ +from pathlib import Path + +import pandas as pd import torch -from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.voting_ensemble import WMVwithConfidenceEnsemble + +N_FOLDS = 5 +WEIGHTING_STRENGTH_GRID = [0, 0.25, 0.5, 0.75, 1] -class WMVwithPPVNPVEnsemble(BaseEnsemble): +class WMVwithF1Ensemble(WMVwithConfidenceEnsemble): def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=1, **kwargs + self, + ensemble_dir: str, + weighting_strength=None, + weighting_exponent=None, + **kwargs, ): - """WMV ensemble that weights models based on their class-wise positive / negative predictive values. For each class, the weight is calculated as: - weight = (weighting_strength * PPV + (1 - weighting_strength)) ** weighting_exponent - where PPV is the class-specific positive predictive value of the model on the validation set - or (if the prediction is negative): - weight = (weighting_strength * NPV + (1 - weighting_strength)) ** weighting_exponent - where NPV is the class-specific negative predictive value of the model on the validation set. + """WMV ensemble that weights models based on their class-wise F1 scores, on top of the + confidence weighting of WMVwithConfidenceEnsemble. For each class, the weight is calculated as: + weight = model_weight * (weighting_strength * F1 + (1 - weighting_strength)) ** weighting_exponent + where F1 is the class-specific F1 score ("trust") of the model on the validation set. + + weighting_strength and weighting_exponent default to the optimal values determined during + calibration (best_hyperparameters.csv in the ensemble directory), falling back to 1 if the + ensemble has not been calibrated. Values passed here take precedence over both. """ - super().__init__(config_path, **kwargs) + super().__init__(ensemble_dir, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_f1_scores = None - def calculate_classwise_weights(self, predicted_classes): - """ - Given the positions of predicted classes in the predictions tensor, assign weights to each class. The - result is two tensors of shape (num_predicted_classes, num_models). The weight for each class is the model_weight - (default: 1) multiplied by the class-specific positive / negative weight (default 1). - """ - positive_weights = torch.ones(len(predicted_classes), len(self.models)) - negative_weights = torch.ones(len(predicted_classes), len(self.models)) - for j, model in enumerate(self.models): - positive_weights[:, j] *= model.model_weight - negative_weights[:, j] *= model.model_weight - if model.classwise_weights is None: - continue - for cls, weights in model.classwise_weights.items(): - if cls not in predicted_classes: - continue - ppv = ( - weights["TP"] / (weights["TP"] + weights["FP"]) - if (weights["TP"] + weights["FP"]) > 0 - else 1.0 - ) - npv = ( - weights["TN"] / (weights["TN"] + weights["FN"]) - if (weights["TN"] + weights["FN"]) > 0 - else 1.0 - ) - positive_weights[predicted_classes[cls], j] *= ( - ppv * self.weighting_strength + (1 - self.weighting_strength) - ) ** self.weighting_exponent - negative_weights[predicted_classes[cls], j] *= ( - npv * self.weighting_strength + (1 - self.weighting_strength) - ) ** self.weighting_exponent - - if self.verbose_output: + def calibrate(self, validation_predictions, validation_data, validation_labels): + super().calibrate(validation_predictions, validation_data, validation_labels) + self._save_classwise_f1(validation_predictions, validation_labels) + self._optimize_hyperparameters(validation_predictions, validation_labels) + + def _fit_classwise_f1(self, predictions, labels, thresholds): + return { + model_name: self.classwise_f1( + model_predictions > thresholds[model_name], labels + ) + for model_name, model_predictions in predictions.items() + } + + def _save_classwise_f1(self, validation_predictions, validation_labels): + thresholds = self._load_prediction_thresholds() + classwise_f1 = self._fit_classwise_f1( + validation_predictions, validation_labels, thresholds + ) + for model_name, f1 in classwise_f1.items(): + f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + with open(f1_path, "w+") as f: + f.writelines(f"{x}\n" for x in f1.tolist()) print( - "Calculated model weightings. The averages for positive / negative weights are:" + f"Saved class-wise F1 scores to {f1_path}: {len(f1.tolist())} classes (macro-f1: {f1.mean().item():.4f})." ) - for i, model in enumerate(self.models): - print( - f"{model.model_name}: {positive_weights[:, i].mean().item():.3f} / {negative_weights[:, i].mean().item():.3f}" - ) - return positive_weights, negative_weights + def _load_classwise_f1(self, model_name: str) -> torch.Tensor: + if self.model_f1_scores is not None: + return self.model_f1_scores[model_name] + classwise_f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + if classwise_f1_path.exists(): + with open(classwise_f1_path, "r", encoding="utf-8") as f: + return torch.tensor([float(x) for x in f.read().splitlines()]) + else: + raise FileNotFoundError( + f"Class-wise F1 scores file not found for model {model_name} in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." + ) + def _load_hyperparameters(self) -> tuple[float, int]: + """Hyperparameters set explicitly take precedence, otherwise the optimal values found during + calibration are used (falling back to 1 if the ensemble has not been calibrated). + """ + best = {} + best_path = Path(self.ensemble_dir) / "best_hyperparameters.csv" + if best_path.exists(): + best = pd.read_csv(best_path).iloc[0].to_dict() + strength = self.weighting_strength + if strength is None: + strength = float(best.get("weighting_strength", 1)) + exponent = self.weighting_exponent + if exponent is None: + exponent = int(best.get("weighting_exponent", 1)) + return strength, exponent -class WMVwithF1Ensemble(BaseEnsemble): + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: + # Calculate trust based on class-wise F1 scores for each model + # target shape: (num_molecules, num_classes, num_models) + weighting_strength, weighting_exponent = self._load_hyperparameters() + num_models = len(predictions) + num_molecules = list(predictions.values())[0].shape[0] + num_classes = list(predictions.values())[0].shape[1] + trust_tensor = torch.ones( + (num_molecules, num_classes, num_models), dtype=torch.float32 + ) + for model_idx, (model_name, prediction_tensor) in enumerate( + predictions.items() + ): + classwise_f1 = self._load_classwise_f1(model_name) + assert ( + classwise_f1.shape[0] == num_classes + ), f"Class-wise F1 scores for model {model_name} do not match number of classes in predictions." + # Expand classwise_f1 to match the shape of trust_tensor for broadcasting + classwise_f1 = classwise_f1.unsqueeze(0).expand(num_molecules, -1) + trust_tensor[:, :, model_idx] = ( + weighting_strength * classwise_f1 + (1 - weighting_strength) + ) ** weighting_exponent + return trust_tensor - def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=6.25, **kwargs + def _build_folds(self, validation_predictions, validation_labels): + """Split the validation set into N_FOLDS folds and calibrate prediction thresholds and + class-wise F1 scores on the training part of each fold. Returns one + (thresholds, class-wise F1 scores, held-out indices) tuple per fold.""" + permutation = torch.randperm( + validation_labels.shape[0], generator=torch.Generator().manual_seed(0) + ) + fold_indices = [permutation[i::N_FOLDS] for i in range(N_FOLDS)] + folds = [] + for fold, test_idx in enumerate(fold_indices): + print(f"Calibrating fold {fold + 1}/{N_FOLDS}...") + train_idx = torch.cat( + [idx for other, idx in enumerate(fold_indices) if other != fold] + ) + train_predictions = { + model_name: model_predictions[train_idx] + for model_name, model_predictions in validation_predictions.items() + } + train_labels = validation_labels[train_idx] + thresholds = self._fit_prediction_thresholds( + train_predictions, train_labels + ) + classwise_f1 = self._fit_classwise_f1( + train_predictions, train_labels, thresholds + ) + folds.append((thresholds, classwise_f1, test_idx)) + return folds + + def _score_hyperparameters( + self, + folds, + validation_predictions, + validation_labels, + weighting_strength, + weighting_exponent, ): - """WMV ensemble that weights models based on their class-wise F1 scores. For each class, the weight is calculated as: - weight = model_weight * (weighting_strength * F1 + (1 - weighting_strength)) ** weighting_exponent - where F1 is the class-specific F1 score ("trust") of the model on the validation set. - """ - super().__init__(config_path, **kwargs) + """Macro F1 of the aggregated predictions on each held-out fold, averaged only over + classes that have positive labels in that fold (a fold is too small for every class to + be present, and absent classes would otherwise contribute a hard 0).""" self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + scores = [] + for thresholds, classwise_f1, test_idx in folds: + self.prediction_thresholds = thresholds + self.model_f1_scores = classwise_f1 + aggregated = self.predict( + { + model_name: model_predictions[test_idx] + for model_name, model_predictions in validation_predictions.items() + } + ) + decisions = (aggregated["net_score"] > 0) & aggregated[ + "has_valid_predictions" + ] + fold_labels = validation_labels[test_idx] + fold_f1 = self.classwise_f1(decisions, fold_labels) + scores.append(fold_f1[fold_labels.sum(dim=0) > 0].mean().item()) + self.prediction_thresholds = None + self.model_f1_scores = None + return scores - def calculate_classwise_weights(self, predicted_classes): - """ - Given the positions of predicted classes in the predictions tensor, assign weights to each class. The - result is two tensors of shape (num_predicted_classes, num_models). The weight for each class is the model_weight - (default: 1) multiplied by (1 + the class-specific validation-f1 (default 1)). - """ - weights_by_cls = torch.ones(len(predicted_classes), len(self.models)) - for j, model in enumerate(self.models): - weights_by_cls[:, j] *= model.model_weight - if model.classwise_weights is None: - continue - for cls, weights in model.classwise_weights.items(): - if cls in predicted_classes: - if (2 * weights["TP"] + weights["FP"] + weights["FN"]) > 0: - f1 = ( - 2 - * weights["TP"] - / (2 * weights["TP"] + weights["FP"] + weights["FN"]) - ) - weights_by_cls[predicted_classes[cls], j] *= ( - self.weighting_strength * f1 + 1 - self.weighting_strength - ) ** self.weighting_exponent - if self.verbose_output: - print("Calculated model weightings. The average weights are:") - for i, model in enumerate(self.models): - print(f"{model.model_name}: {weights_by_cls[:, i].mean().item():.3f}") - - return weights_by_cls, weights_by_cls + def _optimize_hyperparameters(self, validation_predictions, validation_labels): + print( + f"Optimizing hyperparameters with {N_FOLDS}-fold cross-validation on the validation set..." + ) + weighting_strength = self.weighting_strength + weighting_exponent = self.weighting_exponent + folds = self._build_folds(validation_predictions, validation_labels) + results = [] + + def score(stage, strength, exponent): + scores = self._score_hyperparameters( + folds, validation_predictions, validation_labels, strength, exponent + ) + mean_score = sum(scores) / len(scores) + results.append( + { + "stage": stage, + "weighting_strength": strength, + "weighting_exponent": exponent, + "mean_macro_f1": mean_score, + "std_macro_f1": torch.tensor(scores).std().item(), + **{f"fold_{i}_macro_f1": s for i, s in enumerate(scores)}, + } + ) + print( + f"weighting_strength={strength}, weighting_exponent={exponent}: macro-f1 {mean_score:.4f}" + ) + return mean_score + + strength_scores = { + strength: score("weighting_strength", strength, 1) + for strength in WEIGHTING_STRENGTH_GRID + } + best_strength = max(strength_scores, key=strength_scores.get) + + best_exponent = 1 + best_score = strength_scores[best_strength] + exponent = 2 + while True: + exponent_score = score("weighting_exponent", best_strength, exponent) + if exponent_score <= best_score: + break + best_score = exponent_score + best_exponent = exponent + exponent += 1 + + self.weighting_strength = weighting_strength + self.weighting_exponent = weighting_exponent + self._save_hyperparameter_results( + results, best_strength, best_exponent, best_score + ) + + def _save_hyperparameter_results( + self, results, best_strength, best_exponent, best_score + ): + results_path = Path(self.ensemble_dir) / "hyperparameter_search.csv" + pd.DataFrame(results).to_csv(results_path, index=False) + best_path = Path(self.ensemble_dir) / "best_hyperparameters.csv" + pd.DataFrame( + [ + { + "weighting_strength": best_strength, + "weighting_exponent": best_exponent, + "mean_macro_f1": best_score, + } + ] + ).to_csv(best_path, index=False) + print( + f"Saved hyperparameter search results to {results_path}. Recommended parameters (saved to {best_path}): " + f"weighting_strength={best_strength}, weighting_exponent={best_exponent} (cross-validated macro-f1: {best_score:.4f})." + ) diff --git a/chebifier/hex_graph.py b/chebifier/hex_graph.py new file mode 100644 index 0000000..1dbebcf --- /dev/null +++ b/chebifier/hex_graph.py @@ -0,0 +1,213 @@ +import numpy as np +import torch +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import connected_components + +from chebifier.inconsistency_resolution import ( + NEUTRAL, + ScoreBasedPredictionSmoother, + densified_exclusion_matrix, + to_logit, +) + + +class HexSmoother(ScoreBasedPredictionSmoother): + def __init__( + self, + chebi_graph, + label_names=None, + disjoint_files=None, + verbose=False, + delta=0.0, + max_states=2**20, + max_component_size=40, + ): + self.delta = delta + self.max_states = max_states + self.max_component_size = max_component_size + self.n_fallbacks = 0 + self.max_component = 0 + self.dead_labels = [] + self._state_cache = {} + super().__init__(chebi_graph, label_names, disjoint_files, verbose) + self._build() + + def set_label_names(self, label_names): + super().set_label_names(label_names) + self._build() + + def _build(self): + if getattr(self, "label_names", None) is None: + return + if getattr(self, "disjoint_groups", None) is None: + return + self._state_cache = {} + succ = self.label_successors[0] + n = succ.shape[0] + excl = densified_exclusion_matrix( + self.label_names, self.label_successors, self.disjoint_groups + ) + self.excl_matrix = excl + self.excl_pairs = torch.nonzero(torch.triu(excl), as_tuple=False) + strict = succ & ~torch.eye(n, dtype=torch.bool) + self.sup_sets = [ + set(torch.nonzero(strict[i]).flatten().tolist()) for i in range(n) + ] + self.excl_sets = [ + set(torch.nonzero(excl[i]).flatten().tolist()) for i in range(n) + ] + self.dead_labels = self._find_dead(succ) + if self.verbose and self.dead_labels: + print(f"HEX: {len(self.dead_labels)} dead labels (inconsistent graph)") + + def _find_dead(self, succ): + index = {label: i for i, label in enumerate(self.label_names)} + dead = set() + for group in self.disjoint_groups: + members = [index[g] for g in group if g in index] + for i in range(len(members)): + for j in range(i + 1, len(members)): + both = succ[:, members[i]] & succ[:, members[j]] + dead.update(torch.nonzero(both).flatten().tolist()) + return sorted(dead) + + def _legal_states(self, comp): + key = frozenset(comp) + if key in self._state_cache: + return self._state_cache[key] + comp_set = set(comp) + order = sorted(comp, key=lambda c: len(self.sup_sets[c] & comp_set)) + local = {c: i for i, c in enumerate(order)} + sup_local = [ + [local[s] for s in self.sup_sets[c] & comp_set if local[s] < i] + for i, c in enumerate(order) + ] + excl_local = [ + [local[e] for e in self.excl_sets[c] & comp_set if local[e] < i] + for i, c in enumerate(order) + ] + states = [] + cur = [] + overflow = [False] + + def rec(i): + if overflow[0]: + return + if i == len(order): + if len(states) >= self.max_states: + overflow[0] = True + return + states.append(cur.copy()) + return + for bit in (0, 1): + if bit == 1: + if any(cur[j] == 0 for j in sup_local[i]): + continue + if any(cur[j] == 1 for j in excl_local[i]): + continue + cur.append(bit) + rec(i + 1) + cur.pop() + if overflow[0]: + return + + rec(0) + result = ( + None + if overflow[0] + else (np.array(order, dtype=np.int64), np.array(states, dtype=bool)) + ) + self._state_cache[key] = result + return result + + def _violating(self, pos): + bad = self.label_successors[0] & pos.unsqueeze(1) & ~pos.unsqueeze(0) + viol = bad.any(dim=1) | bad.any(dim=0) + if self.excl_pairs.shape[0]: + a, b = self.excl_pairs[:, 0], self.excl_pairs[:, 1] + both = pos[a] & pos[b] + if both.any(): + viol = viol.clone() + viol[a[both]] = True + viol[b[both]] = True + return viol + + def _resolve_row(self, f, scores, valid): + pos = scores > NEUTRAL + known = torch.ones_like(pos) if valid is None else valid + if valid is not None: + pos = pos & valid + active = ( + ((scores - NEUTRAL).abs() < self.delta) | self._violating(pos) + ) & known + idx = torch.nonzero(active).flatten() + if idx.numel() == 0: + return scores + succ = self.label_successors[0] + sup = succ[idx] + subs = succ[:, idx].T + exc = self.excl_matrix[idx] + value = pos.clone() + free = torch.ones(idx.numel(), dtype=torch.bool) + out = scores.clone() + for _ in range(20): + settled = (~active) & known + clamped_zero = settled & ~value + clamped_one = settled & value + forced_zero = ( + (sup & clamped_zero.unsqueeze(0)).any(dim=1) + | (exc & clamped_one.unsqueeze(0)).any(dim=1) + ) & free + forced_one = (subs & clamped_one.unsqueeze(0)).any(dim=1) & free + if bool((forced_zero & forced_one).any()): + return None + newly = forced_zero | forced_one + if not bool(newly.any()): + break + value[idx[forced_one]] = True + value[idx[forced_zero]] = False + out[idx[forced_one]] = 1.0 + out[idx[forced_zero]] = 0.0 + free = free & ~newly + active = active.clone() + active[idx[newly]] = False + idx = idx[free] + if idx.numel() == 0: + return out + sub = succ[idx][:, idx] + adj = (sub | sub.T | self.excl_matrix[idx][:, idx]).numpy() + n_comp, labels = connected_components( + csr_matrix(adj), directed=False, return_labels=True + ) + for c in range(n_comp): + members = np.nonzero(labels == c)[0] + if members.size < 2: + continue + self.max_component = max(self.max_component, int(members.size)) + if members.size > self.max_component_size: + return None + enumerated = self._legal_states(tuple(sorted(int(idx[m]) for m in members))) + if enumerated is None: + return None + order, states = enumerated + st = torch.from_numpy(states).to(f.dtype) + weights = torch.softmax(st @ f[torch.from_numpy(order)], dim=0) + out[torch.from_numpy(order)] = weights @ st + return out + + def __call__(self, preds, valid_mask=None): + if preds.shape[1] == 0: + return preds + out = preds.clone() + # the softmax over legal states is a log-linear model, P(state) proportional to + # exp(sum of the logits that are on), so this is the one step that needs log-odds + f = to_logit(preds) + for row in range(preds.shape[0]): + valid = valid_mask[row] if valid_mask is not None else None + resolved = self._resolve_row(f[row], preds[row], valid) + if resolved is None: + self.n_fallbacks += 1 + out[row] = super().__call__(preds[row : row + 1])[0] + else: + out[row] = resolved + return out diff --git a/chebifier/ilr.py b/chebifier/ilr.py new file mode 100644 index 0000000..84b951e --- /dev/null +++ b/chebifier/ilr.py @@ -0,0 +1,138 @@ +import torch + +from chebifier.inconsistency_resolution import ( + PredictionSmoother, + densified_exclusion_pairs, +) + + +class ILRSmoother(PredictionSmoother): + def __init__( + self, + chebi_graph, + label_names=None, + disjoint_files=None, + verbose=False, + alpha=1.0, + max_iter=10, + tol=1e-4, + ): + self.alpha = alpha + self.max_iter = max_iter + self.tol = tol + self.excl_pairs = torch.zeros((0, 2), dtype=torch.long) + self.last_iterations = 0 + self.max_iterations = 0 + self._valid = None + super().__init__(chebi_graph, label_names, disjoint_files, verbose) + self._build_exclusions() + + def set_label_names(self, label_names): + super().set_label_names(label_names) + self._build_exclusions() + + def _build_exclusions(self): + if getattr(self, "label_names", None) is None: + return + if getattr(self, "disjoint_groups", None) is None: + return + self.excl_pairs = densified_exclusion_pairs( + self.label_names, self.label_successors, self.disjoint_groups + ) + + def _source(self, p, fill): + if self._valid is None: + return p + return torch.where(self._valid, p, torch.full_like(p, fill)) + + def _up(self, p): + masked = torch.where( + torch.transpose(self.label_successors, 1, 2), + self._source(p, 0.0).unsqueeze(1), + -torch.inf, + ) + return masked.max(dim=2).values + + def _down(self, p): + masked = torch.where( + self.label_successors, self._source(p, 1.0).unsqueeze(1), torch.inf + ) + return masked.min(dim=2).values + + def _disjointness_deviation(self, p): + raise NotImplementedError + + def _subsumption_deviation(self, p): + raise NotImplementedError + + def _step(self, p): + sub_dev = self._subsumption_deviation(p) + disj_dev = self._disjointness_deviation(p) + dev = torch.where(sub_dev.abs() >= disj_dev.abs(), sub_dev, disj_dev) + return (p + self.alpha * dev).clamp(0.0, 1.0) + + def __call__(self, preds, valid_mask=None): + if preds.shape[1] == 0: + return preds + self._valid = valid_mask + p = preds.clamp(0.0, 1.0) + original = p + self.last_iterations = 0 + for _ in range(self.max_iter): + p_new = self._step(p) + if valid_mask is not None: + p_new = torch.where(valid_mask, p_new, original) + self.last_iterations += 1 + if torch.max((p_new - p).abs()) < self.tol: + p = p_new + break + p = p_new + self.max_iterations = max(self.max_iterations, self.last_iterations) + self._valid = None + return p + + +class GodelILRSmoother(ILRSmoother): + def _subsumption_deviation(self, p): + return self._up(p) - p + + def _disjointness_deviation(self, p): + if self.excl_pairs.shape[0] == 0: + return torch.zeros_like(p) + a, b = self.excl_pairs[:, 0], self.excl_pairs[:, 1] + src = self._source(p, 0.0) + pa, pb = src[:, a], src[:, b] + violated = torch.minimum(pa, pb) > 0 + a_is_min = pa <= pb + acc = torch.zeros_like(p) + n = p.shape[0] + acc.scatter_add_( + 1, a.unsqueeze(0).expand(n, -1), (violated & a_is_min).to(p.dtype) + ) + acc.scatter_add_( + 1, b.unsqueeze(0).expand(n, -1), (violated & ~a_is_min).to(p.dtype) + ) + return torch.where(acc > 0, -p, torch.zeros_like(p)) + + +class LukasiewiczILRSmoother(ILRSmoother): + def _subsumption_deviation(self, p): + raise_ = (self._up(p) - p).clamp(min=0.0) / 2 + lower = (p - self._down(p)).clamp(min=0.0) / 2 + return torch.where(raise_ >= lower, raise_, -lower) + + def _disjointness_deviation(self, p): + if self.excl_pairs.shape[0] == 0: + return torch.zeros_like(p) + a, b = self.excl_pairs[:, 0], self.excl_pairs[:, 1] + src = self._source(p, 0.0) + excess = (src[:, a] + src[:, b] - 1.0).clamp(min=0.0) / 2 + acc = torch.zeros_like(p) + n = p.shape[0] + acc.scatter_reduce_( + 1, a.unsqueeze(0).expand(n, -1), excess, reduce="amax", include_self=True + ) + acc.scatter_reduce_( + 1, b.unsqueeze(0).expand(n, -1), excess, reduce="amax", include_self=True + ) + return -acc diff --git a/chebifier/inconsistency_resolution.py b/chebifier/inconsistency_resolution.py index f442de8..b9ded0e 100644 --- a/chebifier/inconsistency_resolution.py +++ b/chebifier/inconsistency_resolution.py @@ -2,7 +2,9 @@ import os from pathlib import Path +import networkx as nx import torch +from chebi_utils.obo_extractor import get_hierarchy_subgraph def get_disjoint_groups(disjoint_files): @@ -26,18 +28,16 @@ def get_disjoint_groups(disjoint_files): if seg.startswith("rdf:Description ") or seg.startswith( "owl:Class" ): - left = int(seg.split('rdf:about="&obo;CHEBI_')[1].split('"')[0]) + left = seg.split('rdf:about="&obo;CHEBI_')[1].split('"')[0] elif seg.startswith("owl:disjointWith"): - right = int( - seg.split('rdf:resource="&obo;CHEBI_')[1].split('"')[0] - ) + right = seg.split('rdf:resource="&obo;CHEBI_')[1].split('"')[0] disjoint_pairs.append([left, right]) disjoint_groups = [] for seg in plaintext.split(""): if "owl;AllDisjointClasses" in seg: classes = seg.split('rdf:about="&obo;CHEBI_')[1:] - classes = [int(c.split('"')[0]) for c in classes] + classes = [c.split('"')[0] for c in classes] disjoint_groups.append(classes) else: raise NotImplementedError( @@ -47,14 +47,75 @@ def get_disjoint_groups(disjoint_files): disjoint_all = disjoint_pairs + disjoint_groups # one disjointness is commented out in the owl-file # (the correct way would be to parse the owl file and notice the comment symbols, but for this case, it should work) - if [22729, 51880] in disjoint_all: - disjoint_all.remove([22729, 51880]) + if ["22729", "51880"] in disjoint_all: + disjoint_all.remove(["22729", "51880"]) # print(f"Found {len(disjoint_all)} disjoint groups") return disjoint_all +NEUTRAL = 0.5 + + +def to_logit(p): + """Log-odds of a probability, for the one place that needs them: the HEX log-linear model. + + Ensembles report probabilities, which is what the resolution methods work in. Only the HEX + softmax over legal states needs logits, and its unanimous predictions sit exactly at 0 and 1, + so the clamp is what keeps them finite. + """ + p = p.clamp(1e-6, 1 - 1e-6) + return torch.log(p) - torch.log1p(-p) + + +def densified_exclusion_matrix(label_names, label_successors, disjoint_groups): + label_index = {label: i for i, label in enumerate(label_names)} + succ = label_successors[0] if label_successors.dim() == 3 else label_successors + n = succ.shape[0] + excl = torch.zeros((n, n), dtype=torch.bool) + for group in disjoint_groups: + members = [label_index[g] for g in group if g in label_index] + for gi in range(len(members)): + for gj in range(gi + 1, len(members)): + subs_a = succ[:, members[gi]] + subs_b = succ[:, members[gj]] + block = subs_a.unsqueeze(1) & subs_b.unsqueeze(0) + excl |= block | block.T + excl.fill_diagonal_(False) + return excl + + +def densified_exclusion_pairs(label_names, label_successors, disjoint_groups): + excl = densified_exclusion_matrix(label_names, label_successors, disjoint_groups) + return torch.nonzero(torch.triu(excl), as_tuple=False) + + +def get_smoother_class(name): + from chebifier.hex_graph import HexSmoother + from chebifier.ilr import GodelILRSmoother, LukasiewiczILRSmoother + + smoothers = { + "score-based": ScoreBasedPredictionSmoother, + "ilr-godel": GodelILRSmoother, + "ilr-lukasiewicz": LukasiewiczILRSmoother, + "hex": HexSmoother, + } + if name not in smoothers: + raise ValueError( + f"Unknown inconsistency resolution method '{name}'. " + f"Available: {', '.join(smoothers)}" + ) + return smoothers[name] + + +SMOOTHER_NAMES = ["score-based", "ilr-godel", "ilr-lukasiewicz", "hex"] + + class PredictionSmoother: - """Removes implication and disjointness violations from predictions""" + """Removes implication and disjointness violations from predictions. + + Predictions are probabilities in [0, 1]: NEUTRAL (0.5) means the ensemble is undecided, and a + class is predicted when its probability exceeds the ensemble's decision threshold. + """ def __init__( self, chebi_graph, label_names=None, disjoint_files=None, verbose=False @@ -67,15 +128,22 @@ def __init__( def set_label_names(self, label_names): if label_names is not None: self.label_names = label_names - chebi_subgraph = self.chebi_graph.subgraph(self.label_names) + # the ChEBI graph also contains non-subsumption relations (has role, conjugate + # acid/base, has functional parent, ...) which are not implications + isa_graph = get_hierarchy_subgraph(self.chebi_graph) + label_index = {label: i for i, label in enumerate(self.label_names)} self.label_successors = torch.zeros( (len(self.label_names), len(self.label_names)), dtype=torch.bool ) for i, label in enumerate(self.label_names): self.label_successors[i, i] = 1 - for p in chebi_subgraph.successors(label): - if p in self.label_names: - self.label_successors[i, self.label_names.index(p)] = 1 + if label not in isa_graph: + continue + # transitive closure: superclasses can be connected via intermediate + # classes that are not themselves labels + for p in nx.descendants(isa_graph, label): + if p in label_index: + self.label_successors[i, label_index[p]] = 1 self.label_successors = self.label_successors.unsqueeze(0) def resolve_subsumption_violations(self, preds): @@ -92,25 +160,26 @@ def resolve_disjointness_violations(self, preds): self.label_names.index(g) for g in disj_group if g in self.label_names ] if len(disj_group) > 1: - disj_max = torch.max(preds[:, disj_group], dim=1) - for i, row in enumerate(preds): - for l_ in range(len(preds[i])): - if l_ in disj_group and l_ != disj_group[disj_max.indices[i]]: - preds[i, l_] = 0 + group_preds = preds[:, disj_group] + keep = torch.zeros_like(group_preds, dtype=torch.bool) + keep[torch.arange(group_preds.shape[0]), group_preds.argmax(dim=1)] = ( + True + ) + preds[:, disj_group] = torch.where( + keep, group_preds, group_preds.clamp(max=NEUTRAL) + ) if self.verbose and torch.sum(preds) != preds_sum_orig: print(f"Preds change (step 2): {torch.sum(preds) - preds_sum_orig}") preds_sum_orig = torch.sum(preds) - # step 3: disjointness violation removal may have caused new implication inconsistencies -> set each prediction to min of predecessors + # step 3: disjointness violation removal may have caused new implication inconsistencies -> set each prediction to min of superclasses preds = preds.unsqueeze(1) - preds_masked_predec = torch.where( - torch.transpose(self.label_successors, 1, 2), preds, 1 - ) - preds = preds_masked_predec.min(dim=2).values + preds_masked_succ = torch.where(self.label_successors, preds, torch.inf) + preds = preds_masked_succ.min(dim=2).values if self.verbose and torch.sum(preds) != preds_sum_orig: print(f"Preds change (step 3): {torch.sum(preds) - preds_sum_orig}") return preds - def __call__(self, preds): + def __call__(self, preds, valid_mask=None): if preds.shape[1] == 0: # no labels predicted return preds @@ -121,7 +190,7 @@ def __call__(self, preds): if self.verbose and torch.sum(preds) != preds_sum_orig: print(f"Preds change (step 1): {torch.sum(preds) - preds_sum_orig}") - # step 2: eliminate disjointness violations: for group of disjoint classes, set all except max to 0.49 (if it is not already lower) + # step 2: eliminate disjointness violations: for group of disjoint classes, set all except max to 0 (if it is not already lower) preds = self.resolve_disjointness_violations(preds) return preds @@ -139,18 +208,23 @@ def resolve_subsumption_violations(self, preds): class ScoreBasedPredictionSmoother(PredictionSmoother): - """Removes implication violations from predictions based on net scores: for A subclassOf B where score(A) > score(B), either set score(B) = max(score(B), score(A)) - if abs(score(A)) > abs(score(B)) or set score(A) = min(score(A), score(B)) otherwise. + """Removes implication violations from predictions based on the predicted probabilities: for A + subclassOf B where score(A) > score(B), either set score(B) = max(score(B), score(A)) if A is + further from NEUTRAL than B, or set score(A) = min(score(A), score(B)) otherwise. """ def resolve_subsumption_violations(self, preds): preds = preds.unsqueeze(1) - preds_masked_succ = torch.where(self.label_successors, preds, 0) - preds_optimistic = preds_masked_succ.max(dim=2).values + # label_successors[i, j] means j is a superclass of i, so raising a class to the score of its + # subclasses means taking the max over its predecessors (and vice versa for lowering it). preds_masked_predec = torch.where( - torch.transpose(self.label_successors, 1, 2), preds, 1 + torch.transpose(self.label_successors, 1, 2), preds, -torch.inf ) - preds_pessimistic = preds_masked_predec.min(dim=2).values - # take the one with the higher absolute value - preds_direction = preds_optimistic - preds_pessimistic > 0 + preds_optimistic = preds_masked_predec.max(dim=2).values + preds_masked_succ = torch.where(self.label_successors, preds, torch.inf) + preds_pessimistic = preds_masked_succ.min(dim=2).values + # take whichever the ensemble is more confident about, i.e. further from NEUTRAL + preds_direction = (preds_optimistic - NEUTRAL).abs() > ( + preds_pessimistic - NEUTRAL + ).abs() return torch.where(preds_direction, preds_optimistic, preds_pessimistic) diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index a8287bb..8fbcb2d 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,8 +1,10 @@ -from chebifier.ensemble.base_ensemble import BaseEnsemble -from chebifier.ensemble.weighted_majority_ensemble import ( - WMVwithF1Ensemble, - WMVwithPPVNPVEnsemble, +from chebifier.ensemble.dynamic_selection_ensemble import DynamicSelectionEnsemble +from chebifier.ensemble.learning_to_rank_ensemble import LearningToRankEnsemble +from chebifier.ensemble.voting_ensemble import ( + MajorityVotingEnsemble, + WMVwithConfidenceEnsemble, ) +from chebifier.ensemble.weighted_majority_ensemble import WMVwithF1Ensemble from chebifier.prediction_models import ( ChEBILookupPredictor, ChemlogPeptidesPredictor, @@ -19,9 +21,11 @@ ) ENSEMBLES = { - "mv": BaseEnsemble, - "wmv-ppvnpv": WMVwithPPVNPVEnsemble, + "mv": MajorityVotingEnsemble, + "wmv-conf": WMVwithConfidenceEnsemble, "wmv-f1": WMVwithF1Ensemble, + "ltr": LearningToRankEnsemble, + "des": DynamicSelectionEnsemble, } diff --git a/chebifier/predict.py b/chebifier/predict.py new file mode 100644 index 0000000..49f4982 --- /dev/null +++ b/chebifier/predict.py @@ -0,0 +1,281 @@ +# Get end-to-end predictions (from SMILES / molecule list via base learners + ensemble + inconsistency resolution to ChEBI classes) + + +import os +from typing import Optional + +import numpy as np +import torch +from rdkit import Chem + +from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.inconsistency_resolution import get_smoother_class +from chebifier.prediction_models.base_predictor import BasePredictor +from chebifier.utils import get_disjoint_files, load_chebi_graph + + +def apply_inconsistency_resolution( + smoother, class_names, aggregated_predictions, batch_size: int = 16 +): + """Resolve inconsistencies in batches - the smoother materialises a + (batch_size, n_classes, n_classes) tensor, which does not fit into memory for a whole dataset split. + """ + smoother.set_label_names(class_names) + net_score = aggregated_predictions["net_score"] + valid = aggregated_predictions.get("has_valid_predictions") + aggregated_predictions["net_score"] = torch.cat( + [ + smoother( + net_score[start : start + batch_size], + None if valid is None else valid[start : start + batch_size], + ) + for start in range(0, net_score.shape[0], batch_size) + ] + ) + return aggregated_predictions + + +def base_learner_cache_path( + prediction_cache_dir: str, model_name: str, split: str +) -> str: + return os.path.join(prediction_cache_dir, f"{model_name}_{split}_predictions.npz") + + +def save_dense_predictions(path: str, classes: list[str], scores: np.ndarray) -> None: + np.savez_compressed(path, classes=np.array(classes), scores=scores) + + +def load_dense_predictions(path: str) -> tuple[list[str], np.ndarray]: + with np.load(path) as data: + return [str(cls) for cls in data["classes"]], data["scores"] + + +def collect_base_learner_predictions( + predictions: dict[str, tuple[list[str], np.ndarray]], + classes: Optional[list[str]] = None, +) -> (dict[str, torch.Tensor], list[str]): + """ + Collect predictions from base learners into a single dictionary. + + Args: + predictions (dict): A dictionary where keys are model names and values are + (class labels, score matrix) pairs as returned by BasePredictor.predict_dense. + Assumes all score matrices have the same number of rows. + classes (Optional[list]): Column space to map the predictions onto. If None (the default), the + union of all classes reported by the base learners is used. Pass an explicit list to align + the predictions with a fixed label set (e.g. the labels of an evaluation dataset) - base + learners may be trained on different label sets, so the union does not necessarily match the + labels you want to compare against. Classes predicted by a base learner but missing from + `classes` are dropped, classes that no base learner covers stay NaN. + + Returns: + dict: A dictionary where keys are model names and values are tensors of predictions with shape (num_samples, num_classes). + If a prediction is missing, it will be NaN. + ensemble_classes (list): A list of class labels that are present in the predictions. + """ + print(f"Collecting base learner predictions from {len(predictions)} models...") + n_samples = -1 + for model_name, (_, scores) in predictions.items(): + if n_samples == -1: + n_samples = scores.shape[0] + else: + assert ( + n_samples == scores.shape[0] + ), f"All prediction matrices must have the same length. Model {model_name} has {scores.shape[0]} predictions, expected {n_samples}." + + if classes is None: + ensemble_classes = sorted( + {cls for model_classes, _ in predictions.values() for cls in model_classes} + ) + else: + ensemble_classes = list(classes) + cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} + + collected_predictions = {} + for model_name, (model_classes, scores) in predictions.items(): + if model_classes == ensemble_classes: + collected_predictions[model_name] = torch.from_numpy( + scores.astype(np.float32) + ) + continue + shared = [ + (source, cls_to_idx[cls]) + for source, cls in enumerate(model_classes) + if cls in cls_to_idx + ] + mapped = np.full((n_samples, len(ensemble_classes)), np.nan, dtype=np.float32) + if shared: + source_idx = np.fromiter( + (source for source, _ in shared), dtype=np.intp, count=len(shared) + ) + target_idx = np.fromiter( + (target for _, target in shared), dtype=np.intp, count=len(shared) + ) + mapped[:, target_idx] = scores[:, source_idx] + collected_predictions[model_name] = torch.from_numpy(mapped) + + return collected_predictions, ensemble_classes + + +def get_base_learner_predictions( + base_learners: dict[str, Optional[BasePredictor]], + molecules: list[str | Chem.Mol], + prediction_cache_dir: Optional[str] = None, + split: str = "test", +) -> dict[str, tuple[list[str], np.ndarray]]: + """Get dense predictions from the base learners, using the cache where available. + + A base learner may be None if its predictions are known to be cached (see + cli.build_base_learners), which avoids loading model checkpoints that are never used. + """ + predictions = {} + for model_name, model in base_learners.items(): + cache_path = ( + None + if prediction_cache_dir is None + else base_learner_cache_path(prediction_cache_dir, model_name, split) + ) + if cache_path is not None and os.path.exists(cache_path): + predictions[model_name] = load_dense_predictions(cache_path) + continue + if model is None: + raise ValueError( + f"Base learner '{model_name}' was not instantiated, but its predictions are " + f"missing from the cache ({cache_path})." + ) + predictions[model_name] = model.predict_dense(molecules) + if cache_path is not None: + save_dense_predictions(cache_path, *predictions[model_name]) + return predictions + + +def aggregate_predictions( + base_learners: dict[str, Optional[BasePredictor]], + ensemble_model: BaseEnsemble, + molecules: list[str | Chem.Mol], + prediction_cache_dir: Optional[str] = None, + classes: Optional[list[str]] = None, + split: str = "test", +) -> tuple[dict, list[str]]: + """Get base learner predictions and aggregate them with the ensemble model. + + The result does not depend on the inconsistency resolution method, so it can be reused for + several resolution variants (see resolve_and_decide). + """ + test_predictions = get_base_learner_predictions( + base_learners, molecules, prediction_cache_dir=prediction_cache_dir, split=split + ) + test_predictions, predicted_classes = collect_base_learner_predictions( + test_predictions, classes=classes + ) + aggregated_predictions = ensemble_model.predict(test_predictions, molecules) + # net_score, has_valid_predictions, intermediate_results_dict + return aggregated_predictions, predicted_classes + + +def resolve_and_decide( + aggregated_predictions: dict, + predicted_classes: list[str], + inconsistency_resolution: Optional[str] = "score-based", + inconsistency_resolution_params: Optional[dict] = None, + decision_threshold: float = 0.5, + chebi_graph=None, + disjoint_files=None, +) -> dict: + """Resolve inconsistencies in aggregated predictions and turn them into class decisions. + + Net scores are probabilities, so `decision_threshold` defaults to the neutral point. Ensembles + that tune their own operating point report it as `BaseEnsemble.decision_threshold`. + + `aggregated_predictions` is not modified, so the same aggregation can be passed to several + resolution variants. Pass `inconsistency_resolution=None` or "none" to skip the resolution, + and chebi_graph / disjoint_files to avoid reloading them for every variant. + """ + aggregated_predictions = dict(aggregated_predictions) + if inconsistency_resolution not in (None, "none"): + if chebi_graph is None: + chebi_graph = load_chebi_graph() + if disjoint_files is None: + disjoint_files = get_disjoint_files() + smoother = get_smoother_class(inconsistency_resolution)( + chebi_graph=chebi_graph, + label_names=None, + disjoint_files=disjoint_files, + **(inconsistency_resolution_params or {}), + ) + aggregated_predictions = apply_inconsistency_resolution( + smoother, predicted_classes, aggregated_predictions + ) + + class_decisions = ( + aggregated_predictions["net_score"] > decision_threshold + ) & aggregated_predictions[ + "has_valid_predictions" + ] # Shape: (num_smiles, num_classes) + + complete_failure = torch.all( + ~aggregated_predictions["has_valid_predictions"], dim=1 + ) + aggregated_predictions["class_decisions"] = class_decisions + aggregated_predictions["complete_failure"] = complete_failure + + aggregated_predictions["predicted_classes"] = predicted_classes + + return aggregated_predictions + + +def predict( + base_learners: dict[str, BasePredictor], + ensemble_model: BaseEnsemble, + molecules: list[str | Chem.Mol], + prediction_cache_dir: Optional[str] = None, + resolve_inconsistencies: bool = True, + inconsistency_resolution: str = "score-based", + inconsistency_resolution_params: Optional[dict] = None, + decision_threshold: Optional[float] = None, + classes: Optional[list[str]] = None, + split: str = "test", +) -> dict: + """ + Get end-to-end predictions from base learners and an ensemble model. + + Args: + base_learners (dict[str, BasePredictor]): A dictionary of base learner models. + ensemble_model (BaseEnsemble): An instance of a BaseEnsemble model. + molecules (list[str | Chem.Mol]): List of molecules for prediction (either SMILES strings or molecule objects). + prediction_cache_dir (Optional[str]): Directory to cache predictions. If None, no caching is performed. If provided, + predictions from base learners will be cached to avoid recomputation (warning: not checked against the molecules provided + -> if the molecules change, you have to empty the cache or provide a new cache directory). + resolve_inconsistencies (bool): Whether to resolve inconsistencies in the aggregated predictions. + inconsistency_resolution (str): Which resolution method to use, see SMOOTHER_NAMES. + decision_threshold (Optional[float]): Threshold for class decisions based on net score. + If None (the default), the threshold reported by the ensemble is used. + classes (Optional[list[str]]): Column space to map the base learner predictions onto, see + collect_base_learner_predictions. If None (the default), the union of all classes is used. + split (str): Name of the dataset split, used to separate cached base learner predictions of + different splits within the same cache directory. + + Returns: + dict: A dictionary containing the final predictions and optionally the smoothed predictions. + """ + aggregated_predictions, predicted_classes = aggregate_predictions( + base_learners, + ensemble_model, + molecules, + prediction_cache_dir=prediction_cache_dir, + classes=classes, + split=split, + ) + return resolve_and_decide( + aggregated_predictions, + predicted_classes, + inconsistency_resolution=( + inconsistency_resolution if resolve_inconsistencies else "none" + ), + inconsistency_resolution_params=inconsistency_resolution_params, + decision_threshold=( + ensemble_model.decision_threshold + if decision_threshold is None + else decision_threshold + ), + ) diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 37851b2..6caf8c7 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -1,35 +1,66 @@ -import json from abc import ABC +import numpy as np +from rdkit import Chem + from .._custom_cache import modelwise_smiles_lru_cache +SCORE_DTYPE = np.float16 + + +def dicts_to_dense(predictions: list[dict | None]) -> tuple[list[str], np.ndarray]: + class_set = set() + previous_keys = None + for pred in predictions: + if pred: + keys = tuple(pred) + if keys != previous_keys: + class_set.update(keys) + previous_keys = keys + classes = sorted(class_set) + cls_to_idx = {cls: idx for idx, cls in enumerate(classes)} + + scores = np.full((len(predictions), len(classes)), np.nan, dtype=SCORE_DTYPE) + previous_keys = None + columns = None + for i, pred in enumerate(predictions): + if pred: + keys = tuple(pred) + if keys != previous_keys: + columns = np.fromiter( + (cls_to_idx[cls] for cls in keys), dtype=np.intp, count=len(keys) + ) + previous_keys = keys + scores[i, columns] = np.fromiter( + pred.values(), dtype=SCORE_DTYPE, count=len(pred) + ) + return classes, scores + class BasePredictor(ABC): def __init__( self, model_name: str, model_weight: int = 1, - classwise_weights_path: str = None, **kwargs, ): self.model_name = model_name self.model_weight = model_weight - if classwise_weights_path is not None: - self.classwise_weights = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.classwise_weights = None self._description = kwargs.get("description", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> dict: + def predict_list(self, molecule_list: list[str | Chem.Mol]) -> list[dict | None]: raise NotImplementedError() - def predict_smiles(self, smiles: str) -> dict: + def predict(self, molecule: str | Chem.Mol) -> dict | None: # by default, use list-based prediction - return self.predict_smiles_list([smiles])[0] + return self.predict_list([molecule])[0] + + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + return dicts_to_dense(self.predict_list(molecule_list)) @property def info_text(self): diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index bf6c39b..c005e09 100644 --- a/chebifier/prediction_models/c3p_predictor.py +++ b/chebifier/prediction_models/c3p_predictor.py @@ -1,3 +1,4 @@ +import functools from pathlib import Path from typing import List, Optional @@ -5,6 +6,32 @@ from chebifier import modelwise_smiles_lru_cache from chebifier.prediction_models import BasePredictor +from chebifier.utils import get_superclasses, to_smiles + + +def _patch_c3p(c3p_classifier): + """Two things C3P 0.5.0 assumes that do not hold on Windows, both of which make it return no + classification at all rather than a wrong one: + + - it reads its generated programs with `open(program, "r")`, i.e. in the platform default + encoding, while the programs are UTF-8. `open` is looked up in the module globals before the + builtins, so binding a UTF-8 `open` there fixes the reads without affecting any other module. + - it guards every program with timeout_decorator, which needs SIGALRM. Running the programs + without their 2s timeout is the only way to get predictions on a platform that has no + SIGALRM - timeout_decorator's signal-free mode forks a process per call, which is not + affordable for 300 programs per molecule. + """ + import signal + + if not hasattr(c3p_classifier, "open"): + c3p_classifier.open = functools.partial(open, encoding="utf-8") + if hasattr(signal, "SIGALRM"): + return + from c3p import learn + + if hasattr(learn.eval_with_timeout, "__wrapped__"): + print("No SIGALRM on this platform, running C3P programs without a timeout.") + learn.eval_with_timeout = learn.eval_with_timeout.__wrapped__ class C3PPredictor(BasePredictor): @@ -25,9 +52,12 @@ def __init__( self.chebi_graph = kwargs.get("chebi_graph", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: from c3p import classifier as c3p_classifier + _patch_c3p(c3p_classifier) + # C3P only takes SMILES, while the evaluation datasets hand out RDKit molecules + smiles_list = [to_smiles(molecule) for molecule in smiles_list] result_list = [] for batch_start in tqdm.tqdm( range(0, len(smiles_list), 32), desc="Classifying with C3P" @@ -42,17 +72,27 @@ def predict_smiles_list(self, smiles_list: list[str]) -> list: ) ) + # Look up the position of each SMILES via a dict instead of scanning smiles_list + # for every result (C3P returns one result per class and molecule, so the scan + # made reformatting quadratic in the number of molecules). Repeated SMILES map to + # all of their positions, which list.index could not do (it always returned the + # first one, leaving the later rows without any predictions). + indices_by_smiles: dict[str, list[int]] = {} + for idx, smiles in enumerate(smiles_list): + indices_by_smiles.setdefault(smiles, []).append(idx) + result_reformatted = [dict() for _ in range(len(smiles_list))] for result in tqdm.tqdm(result_list, desc="Reformatting C3P results"): chebi_id = result.class_id.split(":")[1] - result_reformatted[smiles_list.index(result.input_smiles)][ - chebi_id - ] = result.is_match if result.is_match and self.chebi_graph is not None: - for parent in list(self.chebi_graph.predecessors(chebi_id)): - result_reformatted[smiles_list.index(result.input_smiles)][ - str(parent) - ] = 1 + parents = get_superclasses(self.chebi_graph, chebi_id) + else: + parents = [] + for idx in indices_by_smiles[result.input_smiles]: + preds_i = result_reformatted[idx] + preds_i[chebi_id] = result.is_match + for parent in parents: + preds_i[parent] = 1 return result_reformatted def explain_smiles(self, smiles): @@ -62,6 +102,7 @@ def explain_smiles(self, smiles): """ from c3p import classifier as c3p_classifier + _patch_c3p(c3p_classifier) highlights = [] result_list = c3p_classifier.classify( [smiles], self.program_directory, self.chemical_classes, strict=False diff --git a/chebifier/prediction_models/chebi_lookup.py b/chebifier/prediction_models/chebi_lookup.py index d68c9cd..f01478e 100644 --- a/chebifier/prediction_models/chebi_lookup.py +++ b/chebifier/prediction_models/chebi_lookup.py @@ -6,7 +6,7 @@ from chebifier import modelwise_smiles_lru_cache from chebifier.prediction_models import BasePredictor -from chebifier.utils import _smiles_to_mol, load_chebi_graph +from chebifier.utils import _smiles_to_mol, get_superclasses, load_chebi_graph class ChEBILookupPredictor(BasePredictor): @@ -61,7 +61,7 @@ def build_smiles_lookup(self): smiles_lookup[canonical_smiles] = [] # if the canonical SMILES is already in the lookup, append "different interpretation of the SMILES" smiles_lookup[canonical_smiles].append( - (chebi_id, list(self.chebi_graph.predecessors(chebi_id))) + (chebi_id, list(get_superclasses(self.chebi_graph, chebi_id))) ) except Exception as e: print( @@ -69,7 +69,7 @@ def build_smiles_lookup(self): ) return smiles_lookup - def predict_smiles(self, smiles: str) -> Optional[dict]: + def predict(self, smiles: str) -> Optional[dict]: if not smiles: return None mol = _smiles_to_mol(smiles) @@ -96,10 +96,10 @@ def predict_smiles(self, smiles: str) -> Optional[dict]: return None @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: predictions = [] for smiles in smiles_list: - predictions.append(self.predict_smiles(smiles)) + predictions.append(self.predict(smiles)) return predictions @@ -150,5 +150,5 @@ def explain_smiles(self, smiles: str) -> dict: "C1=CC=CC=C1", "*C(=O)OC[C@H](COP(=O)([O-])OCC[N+](C)(C)C)OC(*)=O", ] # SMILES with 251 matches in ChEBI - predictions = predictor.predict_smiles_list(smiles_list) + predictions = predictor.predict_list(smiles_list) print(predictions) diff --git a/chebifier/prediction_models/chemlog_predictor.py b/chebifier/prediction_models/chemlog_predictor.py index 51fccd0..2fbd291 100644 --- a/chebifier/prediction_models/chemlog_predictor.py +++ b/chebifier/prediction_models/chemlog_predictor.py @@ -4,6 +4,7 @@ from chebifier import modelwise_smiles_lru_cache from chebifier.prediction_models.base_predictor import BasePredictor +from chebifier.utils import CHEBI_VERSION, get_superclasses, to_mol AA_DICT = { "A": "L-alanine", @@ -43,7 +44,7 @@ def __init__(self, model_name: str, **kwargs): ] @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: results = [] for predictor in self.predictors: predictor_results = predictor._predict_smiles_list(smiles_list) @@ -66,45 +67,45 @@ def __init__(self, model_name: str, **kwargs): self.classifier = None @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: return self._predict_smiles_list(smiles_list) def _predict_smiles_list(self, smiles_list: list[str]) -> list: - from chemlog.cli import _smiles_to_mol - - mol_list = [_smiles_to_mol(smiles) for smiles in smiles_list] + mol_list = [to_mol(molecule) for molecule in smiles_list] res = self.classifier.classify(mol_list) if self.chebi_graph is not None: for sample in res: sample_additions = dict() for cls in sample: if sample[cls] == 1: - successors = list(self.chebi_graph.predecessors(cls)) - if successors: - for succ in successors: - sample_additions[str(succ)] = 1 + for superclass in get_superclasses(self.chebi_graph, cls): + sample_additions[superclass] = 1 sample.update(sample_additions) return res class ChemlogXMolecularEntityPredictor(ChemlogExtraPredictor): - def __init__(self, model_name: str, **kwargs): + def __init__(self, model_name: str, chebi_version: int = CHEBI_VERSION, **kwargs): from chemlog_extra.alg_classification.by_element_classification import ( XMolecularEntityClassifier, ) super().__init__(model_name, **kwargs) - self.classifier = XMolecularEntityClassifier(chebi_graph=self.chebi_graph) + self.classifier = XMolecularEntityClassifier( + chebi_graph=self.chebi_graph, chebi_version=chebi_version + ) class ChemlogOrganoXCompoundPredictor(ChemlogExtraPredictor): - def __init__(self, model_name: str, **kwargs): + def __init__(self, model_name: str, chebi_version: int = CHEBI_VERSION, **kwargs): from chemlog_extra.alg_classification.by_element_classification import ( OrganoXCompoundClassifier, ) super().__init__(model_name, **kwargs) - self.classifier = OrganoXCompoundClassifier(chebi_graph=self.chebi_graph) + self.classifier = OrganoXCompoundClassifier( + chebi_graph=self.chebi_graph, chebi_version=chebi_version + ) class ChemlogLopsterPredictor(ChemlogExtraPredictor): @@ -141,10 +142,10 @@ def __init__(self, model_name: str, **kwargs): # fmt: on print(f"Initialised ChemLog model {self.model_name}") - def predict_smiles(self, smiles: str) -> Optional[dict]: - from chemlog.cli import _smiles_to_mol, strategy_call + def predict(self, smiles: str) -> Optional[dict]: + from chemlog.cli import strategy_call - mol = _smiles_to_mol(smiles) + mol = to_mol(smiles) if mol is None: return None pos_labels = [ @@ -157,9 +158,9 @@ def predict_smiles(self, smiles: str) -> Optional[dict]: ] if self.chebi_graph: indirect_pos_labels = [ - str(pr) + superclass for label in pos_labels - for pr in self.chebi_graph.predecessors(label) + for superclass in get_superclasses(self.chebi_graph, label) ] pos_labels = list(set(pos_labels + indirect_pos_labels)) return { @@ -168,20 +169,20 @@ def predict_smiles(self, smiles: str) -> Optional[dict]: } @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: return self._predict_smiles_list(smiles_list) def _predict_smiles_list(self, smiles_list: list[str]) -> list: results = [] for i, smiles in tqdm.tqdm(enumerate(smiles_list)): - results.append(self.predict_smiles(smiles)) + results.append(self.predict(smiles)) for classifier in self.classifier_instances.values(): classifier.on_finish() return results - def get_chemlog_result_info(self, smiles): + def get_chemlog_result_info(self, molecule): """Get classification for single molecule with additional information.""" from chemlog.alg_classification.charge_classifier import get_charge_category from chemlog.alg_classification.peptide_size_classifier import ( @@ -194,10 +195,9 @@ def get_chemlog_result_info(self, smiles): is_diketopiperazine, is_emericellamide, ) - from chemlog.cli import _smiles_to_mol - mol = _smiles_to_mol(smiles) - if mol is None or not smiles: + mol = to_mol(molecule) if molecule else None + if mol is None: return {"error": "Failed to parse SMILES"} charge_category = get_charge_category(mol) diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 971a42d..a39eff6 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -1,14 +1,16 @@ from abc import ABC from typing import TYPE_CHECKING +import numpy as np from chebai.result.prediction import Predictor +from rdkit import Chem from chebifier import modelwise_smiles_lru_cache -from .base_predictor import BasePredictor +from .base_predictor import SCORE_DTYPE, BasePredictor if TYPE_CHECKING: - from torch import Tensor + pass class NNPredictor(BasePredictor, ABC): @@ -21,7 +23,7 @@ def __init__( super().__init__(model_name, **kwargs) self.batch_size = kwargs.get("batch_size", None) # compile_model will run the model in eager mode, which gives better performance, but does not return intermediate states - # such as attention weights. Therfore, ELECTRA attention graphs will only work with compile_model=False. + # such as attention weights. Therefore, ELECTRA attention graphs will only work with compile_model=False. compile_model = kwargs.get("compile_model", True) # If batch_size is not provided, it will be set to default batch size used during training in Predictor self.predictor: Predictor = Predictor( @@ -29,27 +31,42 @@ def __init__( ) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: """ Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary of classes and predicted values. """ - raw_preds: Tensor = self.predictor.predict_smiles(smiles_list) - if raw_preds is not None: - preds = [ - ( - { - label: pred - for label, pred in zip( - self.predictor._classification_labels, raw_preds[i].tolist() - ) - } - ) - for i in range(len(smiles_list)) - ] - return preds - else: + raw_preds = self.predictor.predict_molecules(smiles_list) + if raw_preds is None: return [None for _ in smiles_list] + return [ + ( + None + if pred_tensor is None + else { + label: pred + for label, pred in zip( + self.predictor._classification_labels, pred_tensor.tolist() + ) + } + ) + for pred_tensor in raw_preds + ] + + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + raw_preds = self.predictor.predict_molecules(molecule_list) + classes = [str(label) for label in self.predictor._classification_labels] + # molecules the model could not process stay NaN, so the ensemble skips this + # model for those rows only (see dicts_to_dense in base_predictor.py) + scores = np.full((len(molecule_list), len(classes)), np.nan, dtype=SCORE_DTYPE) + if raw_preds is None: + return classes, scores + for idx, pred in enumerate(raw_preds): + if pred is not None: + scores[idx] = pred.detach().cpu().numpy() + return classes, scores def calculate_results(self, batch): collator = self.predictor._dm.reader.COLLATOR() diff --git a/chebifier/utils.py b/chebifier/utils.py index 96f4a80..8c62db0 100644 --- a/chebifier/utils.py +++ b/chebifier/utils.py @@ -3,14 +3,15 @@ import os import pickle -import fastobo import networkx as nx -import requests import yaml +from chebi_utils.obo_extractor import get_hierarchy_subgraph from rdkit import Chem from chebifier.hugging_face import download_model_files +CHEBI_VERSION = 252 + def load_chebi_graph(filename=None): """Load ChEBI graph from Hugging Face (if filename is None) or local file""" @@ -20,7 +21,7 @@ def load_chebi_graph(filename=None): { "repo_id": "chebai/chebifier", "repo_type": "dataset", - "files": {"f": "chebi_graph.pkl"}, + "files": {"f": f"chebi_graph_v{CHEBI_VERSION}.pkl"}, } )["f"] else: @@ -29,82 +30,6 @@ def load_chebi_graph(filename=None): return pickle.load(open(file, "rb")) -def term_callback(doc): - """Similar to the chebai function, but reduced to the necessary fields. Also, ChEBI IDs are strings""" - parents = [] - name = None - smiles = None - subset = None - for clause in doc: - if isinstance(clause, fastobo.term.PropertyValueClause): - t = clause.property_value - if str(t.relation) == "http://purl.obolibrary.org/obo/chebi/smiles": - assert smiles is None - smiles = t.value - # in older chebi versions, smiles strings are synonyms - # e.g. synonym: "[F-].[Na+]" RELATED SMILES [ChEBI] - elif isinstance(clause, fastobo.term.SynonymClause): - if "SMILES" in clause.raw_value(): - assert smiles is None - smiles = clause.raw_value().split('"')[1] - elif isinstance(clause, fastobo.term.IsAClause): - chebi_id = str(clause.term) - chebi_id = chebi_id[chebi_id.index(":") + 1 :] - parents.append(chebi_id) - elif isinstance(clause, fastobo.term.NameClause): - name = str(clause.name) - elif isinstance(clause, fastobo.term.SubsetClause): - subset = str(clause.subset) - if isinstance(clause, fastobo.term.IsObsoleteClause): - if clause.obsolete: - # if the term document contains clause as obsolete as true, skips this document. - return False - chebi_id = str(doc.id) - chebi_id = chebi_id[chebi_id.index(":") + 1 :] - return { - "id": chebi_id, - "parents": parents, - "name": name, - "smiles": smiles, - "subset": subset, - } - - -def build_chebi_graph(chebi_version=241): - """Creates a networkx graph for the ChEBI hierarchy. Usually, you don't want to call this function directly, but rather use the `load_chebi_graph` function.""" - chebi_path = os.path.join("data", f"chebi_v{chebi_version}", "chebi.obo") - os.makedirs(os.path.join("data", f"chebi_v{chebi_version}"), exist_ok=True) - if not os.path.exists(chebi_path): - url = f"http://purl.obolibrary.org/obo/chebi/{chebi_version}/chebi.obo" - r = requests.get(url, allow_redirects=True) - open(chebi_path, "wb").write(r.content) - with open(chebi_path, encoding="utf-8") as chebi: - chebi = "\n".join(line for line in chebi if not line.startswith("xref:")) - - elements = [] - for term_doc in fastobo.loads(chebi): - if ( - term_doc - and isinstance(term_doc.id, fastobo.id.PrefixedIdent) - and term_doc.id.prefix == "CHEBI" - ): - term_dict = term_callback(term_doc) - if term_dict: - elements.append(term_dict) - - g = nx.DiGraph() - for n in elements: - g.add_node(n["id"], **n) - - # Only take the edges which connect the existing nodes, to avoid internal creation of obsolete nodes - # https://github.com/ChEB-AI/python-chebai/pull/55#issuecomment-2386654142 - g.add_edges_from( - [(p, q["id"]) for q in elements for p in q["parents"] if g.has_node(p)], - label="direct_child", - ) - return nx.transitive_closure_dag(g) - - def get_disjoint_files(): """Gets local disjointness files if they are present in the right location, otherwise downloads them from Hugging Face.""" local_disjoint_files = [ @@ -170,10 +95,38 @@ def _smiles_to_mol(smiles: str): return mol -if __name__ == "__main__": - chebi_graph = build_chebi_graph(chebi_version=244) - os.makedirs(os.path.join("data", "chebi_v244"), exist_ok=True) - pickle.dump( - chebi_graph, - open(os.path.join("data", "chebi_v244", "chebi_graph.pkl"), "wb"), - ) +def to_mol(molecule: str | Chem.Mol): + """Molecules reach a predictor either as SMILES or as RDKit molecules (the evaluation datasets + store the latter). Rule-based classifiers expect kekulised molecules, and Kekulize works in + place, so a molecule that is not ours to modify is copied first.""" + if not isinstance(molecule, Chem.Mol): + return _smiles_to_mol(molecule) + molecule = Chem.Mol(molecule) + try: + Chem.Kekulize(molecule) + except Chem.KekulizeException as e: + print(f"Failed to Kekulize {Chem.MolToSmiles(molecule)}: {e}") + return molecule + + +def to_smiles(molecule: str | Chem.Mol) -> str: + return Chem.MolToSmiles(molecule) if isinstance(molecule, Chem.Mol) else molecule + + +@functools.lru_cache(maxsize=2) +def _isa_graph(chebi_graph): + return get_hierarchy_subgraph(chebi_graph) + + +@functools.lru_cache(maxsize=None) +def get_superclasses(chebi_graph, chebi_id: str) -> tuple[str, ...]: + """All transitive superclasses of a ChEBI class. + + is-a edges point from child to parent, and the graph also carries non-subsumption relations + (has role, conjugate acid/base, ...), so the superclasses of a node are neither its + predecessors nor all of its successors. + """ + isa_graph = _isa_graph(chebi_graph) + if chebi_id not in isa_graph: + return () + return tuple(str(cls) for cls in nx.descendants(isa_graph, chebi_id)) diff --git a/pyproject.toml b/pyproject.toml index 60cd4fb..df2f773 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,10 @@ dependencies = [ "pyyaml", "tqdm", "rdkit", + "chebi-utils>=0.3", + # level-1 ensembles: learning to rank (ltr) and dynamic ensemble selection (des) + "lightgbm", + "scikit-learn", # Package to install manually if required #"chebai>=1.0.1", #"chemlog>=1.0.4",