From ce757d5a0f4a5c2381b91792d28e5a75ab5040d0 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 28 Jul 2026 18:09:49 +0200 Subject: [PATCH 01/15] separate ensemble from base learner predictions, generalise both ensemble and base learners for inclusion of new ensemble strategies that are not voting-based --- chebifier/__init__.py | 4 +- chebifier/build_ensemble.py | 61 ++++ chebifier/ensemble/base_ensemble.py | 315 ++---------------- chebifier/ensemble/voting_ensemble.py | 304 +++++++++++++++++ .../ensemble/weighted_majority_ensemble.py | 41 ++- chebifier/model_registry.py | 4 +- chebifier/predict.py | 1 + chebifier/prediction_models/base_predictor.py | 16 +- chebifier/prediction_models/c3p_predictor.py | 2 +- chebifier/prediction_models/chebi_lookup.py | 8 +- .../prediction_models/chemlog_predictor.py | 10 +- chebifier/prediction_models/nn_predictor.py | 2 +- pyproject.toml | 1 + 13 files changed, 447 insertions(+), 322 deletions(-) create mode 100644 chebifier/build_ensemble.py create mode 100644 chebifier/ensemble/voting_ensemble.py create mode 100644 chebifier/predict.py diff --git a/chebifier/__init__.py b/chebifier/__init__.py index a4f770c..34967ee 100644 --- a/chebifier/__init__.py +++ b/chebifier/__init__.py @@ -2,10 +2,10 @@ # even if multiple subpackages are imported later. from ._custom_cache import PerSmilesPerModelLRUCache, modelwise_smiles_lru_cache -from .ensemble.base_ensemble import BaseEnsemble +from .ensemble.voting_ensemble import VotingEnsemble __all__ = [ - "BaseEnsemble", + "VotingEnsemble", "PerSmilesPerModelLRUCache", "modelwise_smiles_lru_cache", ] diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py new file mode 100644 index 0000000..eaf2cba --- /dev/null +++ b/chebifier/build_ensemble.py @@ -0,0 +1,61 @@ +import os + +import torch + + +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 (torch.Tensor): Validation labels for calibration. + 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 + + 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 = {} + # get cached predictions if available, otherwise compute and cache them + for model_name, model in self.base_learners.items(): + cache_path = os.path.join( + self.prediction_cache_dir, f"{model_name}_validation_predictions.pt" + ) + if os.path.exists(cache_path): + validation_predictions[model_name] = torch.load( + cache_path, weights_only=False + ) + else: + validation_predictions[model_name] = model.predict_list( + self.validation_data + ) + torch.save(validation_predictions[model_name], cache_path) + + # Step 2: Calibrate the ensemble model using validation predictions + self.ensemble_model.calibrate( + validation_predictions, self.validation_data, self.validation_labels + ) + + return self.ensemble_model diff --git a/chebifier/ensemble/base_ensemble.py b/chebifier/ensemble/base_ensemble.py index 59601f4..67f9fa3 100644 --- a/chebifier/ensemble/base_ensemble.py +++ b/chebifier/ensemble/base_ensemble.py @@ -1,303 +1,40 @@ -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) - - with ( - importlib.resources.files("chebifier") - .joinpath("model_registry.yml") - .open("r") as f - ): - model_registry = yaml.safe_load(f) - - processed_configs = process_config(config, model_registry) - self.verbose_output = verbose_output - self.use_confidence = use_confidence - - self.chebi_graph = load_chebi_graph() - self.disjoint_files = get_disjoint_files() + """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) - 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"]) + Each ensemble gets a directory where it can store its calibration results (e.g. weights for WMV or meta-model parameters). - 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) + Not part of the ensemble are + - getting predictions from base learners + - resolving inconsistencies in the aggregated predictions + """ - 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 __init__(self, ensemble_dir: str): + os.makedirs(ensemble_dir, exist_ok=True) + self.ensemble_dir = ensemble_dir - 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] - ) + @property + def ensemble_name(self): + return self.__class__.__name__ - return ordered_logits, predicted_classes - - 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) - ) - - # 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 - ): - # 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 - + pass -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]) + def predict(self, test_predictions: dict[str, torch.Tensor]): + raise NotImplementedError() diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py new file mode 100644 index 0000000..38aa749 --- /dev/null +++ b/chebifier/ensemble/voting_ensemble.py @@ -0,0 +1,304 @@ +import importlib +import time +from pathlib import Path +from typing import Union + +import torch +import tqdm +import yaml + +from chebifier.check_env import check_package_installed +from chebifier.ensemble.base_ensemble import BaseEnsemble +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, +) + + +class VotingEnsemble(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) + + with ( + importlib.resources.files("chebifier") + .joinpath("model_registry.yml") + .open("r") as f + ): + model_registry = yaml.safe_load(f) + + processed_configs = process_config(config, model_registry) + self.verbose_output = verbose_output + self.use_confidence = use_confidence + + self.chebi_graph = load_chebi_graph() + self.disjoint_files = get_disjoint_files() + + 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"]) + + 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) + + 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 + + def consolidate_predictions( + self, + predictions, + classwise_weights, + return_intermediate_results=False, + **kwargs, + ): + """ + 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) + ) + + # 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 + ): + # 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 = VotingEnsemble( + { + "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]) diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 1353325..59d0c32 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,9 +1,12 @@ +import json +import os + import torch -from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.voting_ensemble import VotingEnsemble -class WMVwithPPVNPVEnsemble(BaseEnsemble): +class WMVwithPPVNPVEnsemble(VotingEnsemble): def __init__( self, config_path=None, weighting_strength=1, weighting_exponent=1, **kwargs @@ -19,6 +22,18 @@ def __init__( self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_classwise_weights = dict() + for model in self.models: + classwise_weights_path = os.path.join( + self.ensemble_dir, f"{model.model_name}_classwise_weights.json" + ) + if os.path.exists(classwise_weights_path): + self.model_classwise_weights[model.model_name] = json.load( + open(classwise_weights_path, encoding="utf-8") + ) + else: + self.model_classwise_weights[model.model_name] = None + def calculate_classwise_weights(self, predicted_classes): """ Given the positions of predicted classes in the predictions tensor, assign weights to each class. The @@ -30,9 +45,9 @@ def calculate_classwise_weights(self, predicted_classes): 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: + if self.model_classwise_weights[model.model_name] is None: continue - for cls, weights in model.classwise_weights.items(): + for cls, weights in self.model_classwise_weights[model.model_name].items(): if cls not in predicted_classes: continue ppv = ( @@ -64,7 +79,7 @@ def calculate_classwise_weights(self, predicted_classes): return positive_weights, negative_weights -class WMVwithF1Ensemble(BaseEnsemble): +class WMVwithF1Ensemble(VotingEnsemble): def __init__( self, config_path=None, weighting_strength=1, weighting_exponent=6.25, **kwargs @@ -77,6 +92,18 @@ def __init__( self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_classwise_weights = dict() + for model in self.models: + classwise_weights_path = os.path.join( + self.ensemble_dir, f"{model.model_name}_classwise_weights.json" + ) + if os.path.exists(classwise_weights_path): + self.model_classwise_weights[model.model_name] = json.load( + open(classwise_weights_path, encoding="utf-8") + ) + else: + self.model_classwise_weights[model.model_name] = None + def calculate_classwise_weights(self, predicted_classes): """ Given the positions of predicted classes in the predictions tensor, assign weights to each class. The @@ -86,9 +113,9 @@ def calculate_classwise_weights(self, predicted_classes): 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: + if self.model_classwise_weights[model.model_name] is None: continue - for cls, weights in model.classwise_weights.items(): + for cls, weights in self.model_classwise_weights[model.model_name].items(): if cls in predicted_classes: if (2 * weights["TP"] + weights["FP"] + weights["FN"]) > 0: f1 = ( diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index a8287bb..ef13e7e 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,4 +1,4 @@ -from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.voting_ensemble import VotingEnsemble from chebifier.ensemble.weighted_majority_ensemble import ( WMVwithF1Ensemble, WMVwithPPVNPVEnsemble, @@ -19,7 +19,7 @@ ) ENSEMBLES = { - "mv": BaseEnsemble, + "mv": VotingEnsemble, "wmv-ppvnpv": WMVwithPPVNPVEnsemble, "wmv-f1": WMVwithF1Ensemble, } diff --git a/chebifier/predict.py b/chebifier/predict.py new file mode 100644 index 0000000..9ff2530 --- /dev/null +++ b/chebifier/predict.py @@ -0,0 +1 @@ +# Get end-to-end predictions (from SMILES via base learners + ensemble + inconsistency resolution to ChEBI classes) diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 37851b2..1082b8e 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -1,6 +1,7 @@ -import json from abc import ABC +from rdkit import Chem + from .._custom_cache import modelwise_smiles_lru_cache @@ -9,27 +10,20 @@ 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]) -> dict: raise NotImplementedError() - def predict_smiles(self, smiles: str) -> dict: + def predict(self, molecule: str | Chem.Mol) -> dict: # by default, use list-based prediction - return self.predict_smiles_list([smiles])[0] + return self.predict_list([molecule])[0] @property def info_text(self): diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index bf6c39b..f0a1cfd 100644 --- a/chebifier/prediction_models/c3p_predictor.py +++ b/chebifier/prediction_models/c3p_predictor.py @@ -25,7 +25,7 @@ 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 result_list = [] diff --git a/chebifier/prediction_models/chebi_lookup.py b/chebifier/prediction_models/chebi_lookup.py index d68c9cd..006af48 100644 --- a/chebifier/prediction_models/chebi_lookup.py +++ b/chebifier/prediction_models/chebi_lookup.py @@ -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..f637e98 100644 --- a/chebifier/prediction_models/chemlog_predictor.py +++ b/chebifier/prediction_models/chemlog_predictor.py @@ -43,7 +43,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,7 +66,7 @@ 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: @@ -141,7 +141,7 @@ 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]: + def predict(self, smiles: str) -> Optional[dict]: from chemlog.cli import _smiles_to_mol, strategy_call mol = _smiles_to_mol(smiles) @@ -168,13 +168,13 @@ 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() diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 971a42d..2bfd389 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -29,7 +29,7 @@ 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. diff --git a/pyproject.toml b/pyproject.toml index 60cd4fb..d73db76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "pyyaml", "tqdm", "rdkit", + "chebi-utils>=0.3", # Package to install manually if required #"chebai>=1.0.1", #"chemlog>=1.0.4", From 4a167f96ff5aa77e7cfbf6aa7926813adfc86362 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Wed, 29 Jul 2026 10:43:28 +0200 Subject: [PATCH 02/15] remove chebi_graph builder (now done via chebi_utils library) --- chebifier/utils.py | 88 ---------------------------------------------- 1 file changed, 88 deletions(-) diff --git a/chebifier/utils.py b/chebifier/utils.py index 96f4a80..8bfa6f3 100644 --- a/chebifier/utils.py +++ b/chebifier/utils.py @@ -3,9 +3,6 @@ import os import pickle -import fastobo -import networkx as nx -import requests import yaml from rdkit import Chem @@ -29,82 +26,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 = [ @@ -168,12 +89,3 @@ def _smiles_to_mol(smiles: str): except Chem.KekulizeException as e: print(f"Failed to Kekulize {smiles}: {e}") 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"), - ) From b26ba7c96c1afb74ccd10d1fe2395afed41f2694 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Wed, 29 Jul 2026 12:56:55 +0200 Subject: [PATCH 03/15] update CLI and integrate MV ensemble into new workflow --- chebifier/build_ensemble.py | 8 + chebifier/cli.py | 341 +++++++++++++---- chebifier/ensemble/voting_ensemble.py | 347 +++++------------- chebifier/predict.py | 146 +++++++- chebifier/prediction_models/base_predictor.py | 4 +- chebifier/prediction_models/nn_predictor.py | 2 +- 6 files changed, 506 insertions(+), 342 deletions(-) diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index eaf2cba..5a9df68 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -2,6 +2,8 @@ import torch +from chebifier.predict import collect_base_learner_predictions + class EnsembleBuilder: """ @@ -28,6 +30,7 @@ def __init__( 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): """ @@ -38,6 +41,7 @@ def build_ensemble(self): # 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 = os.path.join( @@ -53,6 +57,10 @@ def build_ensemble(self): ) torch.save(validation_predictions[model_name], cache_path) + validation_predictions, classes = collect_base_learner_predictions( + validation_predictions + ) + # Step 2: Calibrate the ensemble model using validation predictions self.ensemble_model.calibrate( validation_predictions, self.validation_data, self.validation_labels diff --git a/chebifier/cli.py b/chebifier/cli.py index 2f0468c..be9bbd7 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -1,6 +1,160 @@ +import importlib.resources +import os +from typing import Literal + import click +import pandas as pd +import yaml +from rdkit import Chem + +from chebifier.build_ensemble import EnsembleBuilder +from chebifier.check_env import check_package_installed +from chebifier.hugging_face import download_model_files +from chebifier.model_registry import ENSEMBLES, MODEL_TYPES +from chebifier.predict import predict as predict_molecules +from chebifier.utils import get_default_configs, 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()]) + + mol_list = [] + for raw_input in raw_inputs: + try: + if raw_input.startswith("InChI="): + mol = Chem.MolFromInchi(raw_input, sanitize=False) + if mol is None: + click.echo(f"Failed to parse InChI: {raw_input}") + mol_list.append(None) + mol_list.append(mol) + elif Chem.MolFromSmiles(raw_input, sanitize=False) is None: + click.echo(f"Failed to parse SMILES: {raw_input}") + mol_list.append(None) + else: + mol_list.append(Chem.MolFromSmiles(raw_input, sanitize=False)) + except Exception as e: + click.echo(f"Error parsing molecule '{raw_input}': {e}.") + mol_list.append(None) + return mol_list + + +def build_base_learners(ensemble_config): + """Instantiate the base learners described by an ensemble configuration file.""" + 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 = load_chebi_graph() + base_learners = {} + for model_name, model_config in process_config(config, model_registry).items(): + 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 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 -from chebifier.model_registry import ENSEMBLES + +def ensemble_options(command): + """Options shared by all commands that use an ensemble.""" + 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( + "--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", + ), + click.option( + "--prediction-cache-dir", + type=click.Path(), + default=None, + help="Directory for caching base learner predictions", + ), + ] + ): + command = option(command) + return command + + +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,107 +164,138 @@ def cli(): @cli.command() +@ensemble_options +@data_options +def build( + ensemble_config, ensemble_type, ensemble_dir, prediction_cache_dir, data_path +): + """Build (calibrate) an ensemble on the ChEBI validation set.""" + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + + # 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() +@ensemble_options +@data_options @click.option( - "--ensemble-config", - "-e", - type=click.Path(exists=True), - default=None, - help="Configuration file for ensemble models", -) -@click.option("--smiles", "-s", multiple=True, help="SMILES strings to predict") -@click.option( - "--smiles-file", - "-f", - type=click.Path(exists=True), - help="File containing SMILES strings (one per line)", + "--resolve-inconsistencies/--no-resolve-inconsistencies", + default=True, + help="Resolve inconsistencies in the aggregated predictions (default: True)", ) @click.option( "--output", "-o", type=click.Path(), - help="Output file to save predictions (optional)", + default=None, + help="Output file to save the evaluation results (optional)", ) +def evaluate( + ensemble_config, + ensemble_type, + ensemble_dir, + prediction_cache_dir, + data_path, + resolve_inconsistencies, + output, +): + """Evaluate an ensemble on the ChEBI test set.""" + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + + # TODO: Hugging Face support + test_data, test_labels = load_dataset(data_path, split="test") + + predictions = predict_molecules( + base_learners, + ensemble_model, + test_data, + prediction_cache_dir=prediction_cache_dir, + resolve_inconsistencies=resolve_inconsistencies, + ) + + print(f"Predictions: {predictions}") + + # TODO: compare predictions to test_labels, report metrics and save them to output + + +@cli.command() +@ensemble_options @click.option( - "--ensemble-type", - "-t", - type=click.Choice(ENSEMBLES.keys()), - default="wmv-f1", - help="Type of ensemble to use (default: Weighted Majority Voting)", + "--molecules", "-m", multiple=True, help="SMILES or InChI strings to predict" ) @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)", + "--molecule-file", + "-f", + type=click.Path(exists=True), + default=None, + help="File containing SMILES or InChI strings (one per line)", ) @click.option( - "--resolve-inconsistencies", - "-r", - is_flag=True, + "--resolve-inconsistencies/--no-resolve-inconsistencies", default=True, - help="Resolve inconsistencies in predictions automatically (default: True)", + help="Resolve inconsistencies in the aggregated predictions (default: True)", +) +@click.option( + "--output", + "-o", + type=click.Path(), + default=None, + help="Output file to save the predictions (optional)", ) @click.option( - "--verbose", - "-v", - is_flag=True, - default=False, - help="Enable verbose output", + "--decision-threshold", + "-dt", + type=float, + default=0, + help="Threshold for classifying predictions (default: 0)", ) def predict( ensemble_config, - smiles, - smiles_file, - output, ensemble_type, - use_confidence, - resolve_inconsistencies=True, - verbose=False, + ensemble_dir, + prediction_cache_dir, + molecules, + molecule_file, + resolve_inconsistencies, + decision_threshold, + 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, - ) - - # 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()]) - - if not smiles_list: - click.echo("No SMILES strings provided. Use --smiles or --smiles-file options.") + """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 - # Make predictions - predictions = ensemble.predict_smiles_list(smiles_list) + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) - if output: - # save as json - import json - - with open(output, "w") as f: - json.dump( - {smiles: pred for smiles, pred in zip(smiles_list, predictions)}, - f, - indent=2, - ) + predictions = predict_molecules( + base_learners, + ensemble_model, + molecules_list, + prediction_cache_dir=prediction_cache_dir, + resolve_inconsistencies=resolve_inconsistencies, + decision_threshold=decision_threshold, + ) - 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") + print(f"Predictions: {predictions}") + # TODO: turn the aggregated predictions into ChEBI classes per molecule, print them / save to output if __name__ == "__main__": - cli() + # cli() + load_dataset( + os.path.join("data", "chebi_v252", "ChEBI25", "processed"), split="validation" + ) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 38aa749..1424d71 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -1,304 +1,131 @@ -import importlib -import time from pathlib import Path -from typing import Union import torch -import tqdm import yaml +from torchmetrics import F1Score -from chebifier.check_env import check_package_installed from chebifier.ensemble.base_ensemble import BaseEnsemble -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, -) class VotingEnsemble(BaseEnsemble): def __init__( self, - model_configs: Union[str, Path, dict, None] = None, - resolve_inconsistencies: bool = True, - verbose_output: bool = False, + ensemble_dir: str, 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) - - with ( - importlib.resources.files("chebifier") - .joinpath("model_registry.yml") - .open("r") as f - ): - model_registry = yaml.safe_load(f) - - processed_configs = process_config(config, model_registry) - self.verbose_output = verbose_output + super().__init__(ensemble_dir) self.use_confidence = use_confidence + self.macro_f1 = None - self.chebi_graph = load_chebi_graph() - self.disjoint_files = get_disjoint_files() - - 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"]) + def find_best_threshold(self, predictions, val_labels_tensor): - model_instance = model_cls( - model_name, - **model_config, - **hugging_face_kwargs, - chebi_graph=self.chebi_graph, + best_threshold = 0.5 + best_f1 = 0.0 + for threshold in range(0, 100): + threshold_value = threshold / 100 + macro_f1_score = self.macro_f1( + predictions > threshold_value, val_labels_tensor ) - assert isinstance(model_instance, BasePredictor) - self.models.append(model_instance) - - 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 + 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): + self.macro_f1 = F1Score( + task="multilabel", num_labels=len(validation_labels), average="macro" ) - 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 + prediction_thresholds = { + model_name: self.find_best_threshold(predictions, validation_labels) + for model_name, predictions in validation_predictions.items() + } + self._save_prediction_thresholds(prediction_thresholds) + + 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) + + def _load_prediction_thresholds(self) -> dict[str, float]: + 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 consolidate_predictions( - self, - predictions, - classwise_weights, - return_intermediate_results=False, - **kwargs, - ): + def predict(self, test_predictions: dict[str, torch.Tensor]): """ Aggregates predictions from multiple models using weighted majority voting. - Optimized version using tensor operations instead of for loops. + weights are only the self-reported confidence (=difference between prediction and threshold). If set to false, all models are weighted equally. """ - num_smiles, num_classes, num_models = predictions.shape - + 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) + 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 > self.positive_prediction_threshold + predictions_tensor > threshold_mask.unsqueeze(0).unsqueeze(0) ) & valid_predictions negative_mask = ( - predictions < self.positive_prediction_threshold + predictions_tensor < threshold_mask.unsqueeze(0).unsqueeze(0) ) & 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: + if self.use_confidence: confidence = 2 * torch.abs( - predictions.nan_to_num() - self.positive_prediction_threshold + predictions_tensor.nan_to_num() + - threshold_mask.unsqueeze(0).unsqueeze(0) ) 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) + confidence = torch.ones_like(predictions_tensor) # Calculate weighted predictions using broadcasting - # predictions shape: (num_smiles, num_classes, num_models) + # predictions shape: (num_molecules, 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) - ) + positive_weighted = positive_mask.float() * confidence + negative_weighted = negative_mask.float() * confidence # 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 - ): - # 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 = VotingEnsemble( - { - "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", - }, + 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 = positive_sum - negative_sum # 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, } - ) - 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]) diff --git a/chebifier/predict.py b/chebifier/predict.py index 9ff2530..c4c6811 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -1 +1,145 @@ -# Get end-to-end predictions (from SMILES via base learners + ensemble + inconsistency resolution to ChEBI classes) +# 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 torch +from rdkit import Chem + +from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother +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): + smoother.set_label_names(class_names) + smooth_net_score = smoother(aggregated_predictions["net_score"]) + aggregated_predictions["net_score"] = smooth_net_score + return aggregated_predictions + + +def collect_base_learner_predictions( + predictions: dict[str, list[dict | 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 lists of predictions. + Assumes those lists have the same length and each entry in the list is either None or a dict + mapping class labels to predicted values. + + 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 None, it will be replaced with NaN. + ensemble_classes (list): A list of class labels that are present in the predictions. + """ + collected_predictions = {} + ensemble_classes = set() + n_samples = -1 + # step 1: collect classes + for model_name, model_predictions in predictions.items(): + for pred in model_predictions: + if pred is not None: + ensemble_classes.update(pred.keys()) + if n_samples == -1: + n_samples = len(model_predictions) + else: + assert n_samples == len( + model_predictions + ), f"All prediction lists must have the same length. Model {model_name} has {len(model_predictions)} predictions, expected {n_samples}." + ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} + # step 2: map predictions to tensors + for model_name, model_predictions in predictions.items(): + # Replace None values with NaN + model_predictions = torch.zeros((n_samples, len(ensemble_classes))) * float( + "nan" + ) + for i, pred in enumerate(model_predictions): + if pred is not None: + for cls, value in pred.items(): + model_predictions[i, cls_to_idx[cls]] = value + collected_predictions[model_name] = model_predictions + + return collected_predictions, list(ensemble_classes) + + +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, + decision_threshold: float = 0, +) -> 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. + decision_threshold (float): Threshold for class decisions based on net score. Default is 0. + + Returns: + dict: A dictionary containing the final predictions and optionally the smoothed predictions. + """ + + # Step 1: Get predictions from base learners on test data + test_predictions = {} + for model_name, model in base_learners.items(): + if prediction_cache_dir is None: + test_predictions[model_name] = model.predict_list(molecules) + else: + cache_path = os.path.join( + prediction_cache_dir, f"{model_name}_test_predictions.pt" + ) + if os.path.exists(cache_path): + test_predictions[model_name] = torch.load( + cache_path, weights_only=False + ) + else: + test_predictions[model_name] = model.predict_list(molecules) + torch.save(test_predictions[model_name], cache_path) + + test_predictions, predicted_classes = collect_base_learner_predictions( + test_predictions + ) + + # Step 2: Get aggregated predictions from the ensemble model + aggregated_predictions = ensemble_model.predict(test_predictions) + # net_score, has_valid_predictions, intermediate_results_dict + + # Step 3: Optionally resolve inconsistencies in the aggregated predictions + if resolve_inconsistencies: + chebi_graph = load_chebi_graph() + disjoint_files = get_disjoint_files() + smoother = ScoreBasedPredictionSmoother( + chebi_graph=chebi_graph, label_names=None, disjoint_files=disjoint_files + ) + 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 diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 1082b8e..806cd6f 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -18,10 +18,10 @@ def __init__( self._description = kwargs.get("description", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_list(self, molecule_list: list[str | Chem.Mol]) -> dict: + def predict_list(self, molecule_list: list[str | Chem.Mol]) -> list[dict | None]: raise NotImplementedError() - def predict(self, molecule: str | Chem.Mol) -> dict: + def predict(self, molecule: str | Chem.Mol) -> dict | None: # by default, use list-based prediction return self.predict_list([molecule])[0] diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 2bfd389..cf48658 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -21,7 +21,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( From 8b48598e304e65e2c10f1d398521001de2f3cc32 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 15:09:52 +0200 Subject: [PATCH 04/15] bug fixes and optimized performance for aggregation --- chebifier/build_ensemble.py | 19 +++++- chebifier/cli.py | 5 +- chebifier/ensemble/voting_ensemble.py | 14 +++- chebifier/predict.py | 70 +++++++++++++++++--- chebifier/prediction_models/c3p_predictor.py | 26 ++++++-- 5 files changed, 106 insertions(+), 28 deletions(-) diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index 5a9df68..23de772 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -13,7 +13,8 @@ class EnsembleBuilder: 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 (torch.Tensor): Validation labels 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. """ @@ -48,22 +49,34 @@ def build_ensemble(self): self.prediction_cache_dir, f"{model_name}_validation_predictions.pt" ) if os.path.exists(cache_path): + print(f"{model_name} validation predictions found in cache, loading...") validation_predictions[model_name] = torch.load( cache_path, weights_only=False ) else: + print(f"Computing {model_name} validation predictions...") validation_predictions[model_name] = model.predict_list( self.validation_data ) torch.save(validation_predictions[model_name], cache_path) + # 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 + 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, self.validation_labels + validation_predictions, self.validation_data, validation_labels ) return self.ensemble_model diff --git a/chebifier/cli.py b/chebifier/cli.py index be9bbd7..49d345b 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -295,7 +295,4 @@ def predict( if __name__ == "__main__": - # cli() - load_dataset( - os.path.join("data", "chebi_v252", "ChEBI25", "processed"), split="validation" - ) + cli() diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 1424d71..8a224ee 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -18,7 +18,9 @@ def __init__( self.macro_f1 = None def find_best_threshold(self, predictions, val_labels_tensor): - + print( + f"Finding best threshold for predictions with shape {predictions.shape} and validation labels with shape {val_labels_tensor.shape}" + ) best_threshold = 0.5 best_f1 = 0.0 for threshold in range(0, 100): @@ -32,10 +34,15 @@ def find_best_threshold(self, predictions, val_labels_tensor): 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.macro_f1 = F1Score( - task="multilabel", num_labels=len(validation_labels), average="macro" + task="multilabel", num_labels=validation_labels.shape[1], average="macro" + ) + print( + f"Validation labels: {validation_labels.shape}, Validation predictions: {validation_predictions[list(validation_predictions.keys())[0]].shape}" ) - prediction_thresholds = { model_name: self.find_best_threshold(predictions, validation_labels) for model_name, predictions in validation_predictions.items() @@ -46,6 +53,7 @@ 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]: thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" diff --git a/chebifier/predict.py b/chebifier/predict.py index c4c6811..fb0726f 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -4,6 +4,7 @@ import os from typing import Optional +import numpy as np import torch from rdkit import Chem @@ -22,6 +23,7 @@ def apply_inconsistency_resolution(smoother, class_names, aggregated_predictions def collect_base_learner_predictions( predictions: dict[str, list[dict | None]], + classes: Optional[list[str]] = None, ) -> (dict[str, torch.Tensor], list[str]): """ Collect predictions from base learners into a single dictionary. @@ -30,6 +32,12 @@ def collect_base_learner_predictions( predictions (dict): A dictionary where keys are model names and values are lists of predictions. Assumes those lists have the same length and each entry in the list is either None or a dict mapping class labels to predicted values. + 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). @@ -39,30 +47,70 @@ def collect_base_learner_predictions( collected_predictions = {} ensemble_classes = set() n_samples = -1 + print(f"Collecting base learner predictions from {len(predictions)} models...") # step 1: collect classes + # Base learners typically return the same label set for every sample, so we only + # touch the class set when the label set actually changes from one sample to the next. for model_name, model_predictions in predictions.items(): - for pred in model_predictions: - if pred is not None: - ensemble_classes.update(pred.keys()) + if classes is None: + previous_keys = None + for pred in model_predictions: + if pred: + keys = tuple(pred) + if keys != previous_keys: + ensemble_classes.update(keys) + previous_keys = keys if n_samples == -1: n_samples = len(model_predictions) else: assert n_samples == len( model_predictions ), f"All prediction lists must have the same length. Model {model_name} has {len(model_predictions)} predictions, expected {n_samples}." - ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + if classes is None: + ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + else: + ensemble_classes = list(classes) cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} # step 2: map predictions to tensors + # Filled row-wise via numpy fancy indexing: one vectorised write per sample instead + # of one Python-level tensor assignment per predicted class. The column indices are + # reused as long as consecutive samples share the same label set (see step 1). for model_name, model_predictions in predictions.items(): - # Replace None values with NaN - model_predictions = torch.zeros((n_samples, len(ensemble_classes))) * float( - "nan" + # Samples without a prediction (and classes the model does not cover) stay NaN + predictions_array = np.full( + (n_samples, len(ensemble_classes)), np.nan, dtype=np.float32 ) + previous_keys = None + columns = None + # positions of the kept values within pred.values(), None if nothing is dropped + kept_positions = None for i, pred in enumerate(model_predictions): - if pred is not None: - for cls, value in pred.items(): - model_predictions[i, cls_to_idx[cls]] = value - collected_predictions[model_name] = model_predictions + if pred: + keys = tuple(pred) + if keys != previous_keys: + kept = [ + (position, cls_to_idx[cls]) + for position, cls in enumerate(keys) + if cls in cls_to_idx + ] + columns = np.fromiter( + (column for _, column in kept), dtype=np.intp, count=len(kept) + ) + kept_positions = ( + None + if len(kept) == len(keys) + else np.fromiter( + (position for position, _ in kept), + dtype=np.intp, + count=len(kept), + ) + ) + previous_keys = keys + values = np.fromiter(pred.values(), dtype=np.float32, count=len(pred)) + predictions_array[i, columns] = ( + values if kept_positions is None else values[kept_positions] + ) + collected_predictions[model_name] = torch.from_numpy(predictions_array) return collected_predictions, list(ensemble_classes) diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index f0a1cfd..ec59748 100644 --- a/chebifier/prediction_models/c3p_predictor.py +++ b/chebifier/prediction_models/c3p_predictor.py @@ -42,17 +42,29 @@ def predict_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 = [ + str(parent) for parent in self.chebi_graph.predecessors(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): From deb681ac5b0dba1f38a5e7d7d6a8005757f448db Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 15:54:00 +0200 Subject: [PATCH 05/15] integrate wmv-f1 ensemble into new workflow --- chebifier/ensemble/voting_ensemble.py | 25 ++- .../ensemble/weighted_majority_ensemble.py | 168 ++++++------------ chebifier/model_registry.py | 6 +- 3 files changed, 67 insertions(+), 132 deletions(-) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 8a224ee..6b23d5f 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -15,19 +15,16 @@ def __init__( ): super().__init__(ensemble_dir) self.use_confidence = use_confidence - self.macro_f1 = None + self.classwise_f1 = None def find_best_threshold(self, predictions, val_labels_tensor): - print( - f"Finding best threshold for predictions with shape {predictions.shape} and validation labels with shape {val_labels_tensor.shape}" - ) best_threshold = 0.5 best_f1 = 0.0 for threshold in range(0, 100): threshold_value = threshold / 100 - macro_f1_score = self.macro_f1( + 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 @@ -37,11 +34,8 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): print( f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." ) - self.macro_f1 = F1Score( - task="multilabel", num_labels=validation_labels.shape[1], average="macro" - ) - print( - f"Validation labels: {validation_labels.shape}, Validation predictions: {validation_predictions[list(validation_predictions.keys())[0]].shape}" + self.classwise_f1 = F1Score( + task="multilabel", num_labels=validation_labels.shape[1], average=None ) prediction_thresholds = { model_name: self.find_best_threshold(predictions, validation_labels) @@ -65,6 +59,10 @@ def _load_prediction_thresholds(self) -> dict[str, float]: f"Prediction thresholds file not found in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." ) + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: + # No trust for MV, only used in WMV + return 1 + def predict(self, test_predictions: dict[str, torch.Tensor]): """ Aggregates predictions from multiple models using weighted majority voting. @@ -112,11 +110,12 @@ def predict(self, test_predictions: dict[str, torch.Tensor]): else: confidence = torch.ones_like(predictions_tensor) + 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 - negative_weighted = negative_mask.float() * confidence + positive_weighted = positive_mask.float() * confidence * trust + negative_weighted = negative_mask.float() * confidence * trust # Sum over models dimension positive_sum = positive_weighted.sum( diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 59d0c32..495fa18 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,134 +1,74 @@ -import json -import os +from pathlib import Path import torch from chebifier.ensemble.voting_ensemble import VotingEnsemble -class WMVwithPPVNPVEnsemble(VotingEnsemble): - - def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=1, **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. - """ - super().__init__(config_path, **kwargs) - self.weighting_strength = weighting_strength - self.weighting_exponent = weighting_exponent - - self.model_classwise_weights = dict() - for model in self.models: - classwise_weights_path = os.path.join( - self.ensemble_dir, f"{model.model_name}_classwise_weights.json" - ) - if os.path.exists(classwise_weights_path): - self.model_classwise_weights[model.model_name] = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.model_classwise_weights[model.model_name] = 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 self.model_classwise_weights[model.model_name] is None: - continue - for cls, weights in self.model_classwise_weights[model.model_name].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: - print( - "Calculated model weightings. The averages for positive / negative weights are:" - ) - 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 - - class WMVwithF1Ensemble(VotingEnsemble): def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=6.25, **kwargs + self, + ensemble_dir: str, + use_confidence: bool = True, + weighting_strength=1, + weighting_exponent=1, + **kwargs, ): """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) + super().__init__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent - self.model_classwise_weights = dict() - for model in self.models: - classwise_weights_path = os.path.join( - self.ensemble_dir, f"{model.model_name}_classwise_weights.json" + 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) + + def _save_classwise_f1(self, validation_predictions, validation_labels): + thresholds = self._load_prediction_thresholds() + for model_name, predictions in validation_predictions.items(): + f1 = self.classwise_f1( + predictions > thresholds[model_name], validation_labels + ) + f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + with open(f1_path, "w+") as f: + f.write("\n".join(f1.tolist())) + print( + f"Saved class-wise F1 scores to {f1_path}: {len(f1.tolist())} classes (macro-f1: {f1.mean().item():.4f})." ) - if os.path.exists(classwise_weights_path): - self.model_classwise_weights[model.model_name] = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.model_classwise_weights[model.model_name] = 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 (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 self.model_classwise_weights[model.model_name] is None: - continue - for cls, weights in self.model_classwise_weights[model.model_name].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}") + def _load_classwise_f1(self, model_name: str) -> torch.Tensor: + 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." + ) - return weights_by_cls, weights_by_cls + 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) + 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] = ( + self.weighting_strength * classwise_f1 + (1 - self.weighting_strength) + ) ** self.weighting_exponent + return trust_tensor diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index ef13e7e..8d737de 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,8 +1,5 @@ from chebifier.ensemble.voting_ensemble import VotingEnsemble -from chebifier.ensemble.weighted_majority_ensemble import ( - WMVwithF1Ensemble, - WMVwithPPVNPVEnsemble, -) +from chebifier.ensemble.weighted_majority_ensemble import WMVwithF1Ensemble from chebifier.prediction_models import ( ChEBILookupPredictor, ChemlogPeptidesPredictor, @@ -20,7 +17,6 @@ ENSEMBLES = { "mv": VotingEnsemble, - "wmv-ppvnpv": WMVwithPPVNPVEnsemble, "wmv-f1": WMVwithF1Ensemble, } From fbafc8e8ae46a6965511ab408fd6c0524c9dd514 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 16:42:48 +0200 Subject: [PATCH 06/15] optimize prediction storage --- chebifier/build_ensemble.py | 16 ++- .../ensemble/weighted_majority_ensemble.py | 2 +- chebifier/predict.py | 128 ++++++++---------- chebifier/prediction_models/base_predictor.py | 37 +++++ chebifier/prediction_models/nn_predictor.py | 14 +- 5 files changed, 114 insertions(+), 83 deletions(-) diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index 23de772..64c6b0c 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -2,7 +2,11 @@ import torch -from chebifier.predict import collect_base_learner_predictions +from chebifier.predict import ( + collect_base_learner_predictions, + load_dense_predictions, + save_dense_predictions, +) class EnsembleBuilder: @@ -46,19 +50,17 @@ def build_ensemble(self): # get cached predictions if available, otherwise compute and cache them for model_name, model in self.base_learners.items(): cache_path = os.path.join( - self.prediction_cache_dir, f"{model_name}_validation_predictions.pt" + self.prediction_cache_dir, f"{model_name}_validation_predictions.npz" ) if os.path.exists(cache_path): print(f"{model_name} validation predictions found in cache, loading...") - validation_predictions[model_name] = torch.load( - cache_path, weights_only=False - ) + validation_predictions[model_name] = load_dense_predictions(cache_path) else: print(f"Computing {model_name} validation predictions...") - validation_predictions[model_name] = model.predict_list( + validation_predictions[model_name] = model.predict_dense( self.validation_data ) - torch.save(validation_predictions[model_name], cache_path) + 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 diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 495fa18..ff51020 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -35,7 +35,7 @@ def _save_classwise_f1(self, validation_predictions, validation_labels): ) f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" with open(f1_path, "w+") as f: - f.write("\n".join(f1.tolist())) + f.writelines(f"{x}\n" for x in f1.tolist()) print( f"Saved class-wise F1 scores to {f1_path}: {len(f1.tolist())} classes (macro-f1: {f1.mean().item():.4f})." ) diff --git a/chebifier/predict.py b/chebifier/predict.py index fb0726f..18a6fa8 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -21,17 +21,26 @@ def apply_inconsistency_resolution(smoother, class_names, aggregated_predictions return aggregated_predictions +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, list[dict | None]], + 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 lists of predictions. - Assumes those lists have the same length and each entry in the list is either None or a dict - mapping class labels to predicted values. + 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 @@ -41,78 +50,51 @@ def collect_base_learner_predictions( 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 None, it will be replaced with NaN. + If a prediction is missing, it will be NaN. ensemble_classes (list): A list of class labels that are present in the predictions. """ - collected_predictions = {} - ensemble_classes = set() - n_samples = -1 print(f"Collecting base learner predictions from {len(predictions)} models...") - # step 1: collect classes - # Base learners typically return the same label set for every sample, so we only - # touch the class set when the label set actually changes from one sample to the next. - for model_name, model_predictions in predictions.items(): - if classes is None: - previous_keys = None - for pred in model_predictions: - if pred: - keys = tuple(pred) - if keys != previous_keys: - ensemble_classes.update(keys) - previous_keys = keys + n_samples = -1 + for model_name, (_, scores) in predictions.items(): if n_samples == -1: - n_samples = len(model_predictions) + n_samples = scores.shape[0] else: - assert n_samples == len( - model_predictions - ), f"All prediction lists must have the same length. Model {model_name} has {len(model_predictions)} predictions, expected {n_samples}." + 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(ensemble_classes) # Sort for consistent ordering + 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)} - # step 2: map predictions to tensors - # Filled row-wise via numpy fancy indexing: one vectorised write per sample instead - # of one Python-level tensor assignment per predicted class. The column indices are - # reused as long as consecutive samples share the same label set (see step 1). - for model_name, model_predictions in predictions.items(): - # Samples without a prediction (and classes the model does not cover) stay NaN - predictions_array = np.full( - (n_samples, len(ensemble_classes)), np.nan, dtype=np.float32 - ) - previous_keys = None - columns = None - # positions of the kept values within pred.values(), None if nothing is dropped - kept_positions = None - for i, pred in enumerate(model_predictions): - if pred: - keys = tuple(pred) - if keys != previous_keys: - kept = [ - (position, cls_to_idx[cls]) - for position, cls in enumerate(keys) - if cls in cls_to_idx - ] - columns = np.fromiter( - (column for _, column in kept), dtype=np.intp, count=len(kept) - ) - kept_positions = ( - None - if len(kept) == len(keys) - else np.fromiter( - (position for position, _ in kept), - dtype=np.intp, - count=len(kept), - ) - ) - previous_keys = keys - values = np.fromiter(pred.values(), dtype=np.float32, count=len(pred)) - predictions_array[i, columns] = ( - values if kept_positions is None else values[kept_positions] - ) - collected_predictions[model_name] = torch.from_numpy(predictions_array) - - return collected_predictions, list(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 predict( @@ -144,18 +126,16 @@ def predict( test_predictions = {} for model_name, model in base_learners.items(): if prediction_cache_dir is None: - test_predictions[model_name] = model.predict_list(molecules) + test_predictions[model_name] = model.predict_dense(molecules) else: cache_path = os.path.join( - prediction_cache_dir, f"{model_name}_test_predictions.pt" + prediction_cache_dir, f"{model_name}_test_predictions.npz" ) if os.path.exists(cache_path): - test_predictions[model_name] = torch.load( - cache_path, weights_only=False - ) + test_predictions[model_name] = load_dense_predictions(cache_path) else: - test_predictions[model_name] = model.predict_list(molecules) - torch.save(test_predictions[model_name], cache_path) + test_predictions[model_name] = model.predict_dense(molecules) + save_dense_predictions(cache_path, *test_predictions[model_name]) test_predictions, predicted_classes = collect_base_learner_predictions( test_predictions diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 806cd6f..6caf8c7 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -1,9 +1,41 @@ 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__( @@ -25,6 +57,11 @@ def predict(self, molecule: str | Chem.Mol) -> dict | None: # by default, use list-based prediction 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): if self._description is None: diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index cf48658..7e4e1b4 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -1,11 +1,13 @@ 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 @@ -51,6 +53,16 @@ def predict_list(self, smiles_list: list[str]) -> list: else: return [None for _ in smiles_list] + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + raw_preds: Tensor = self.predictor.predict_smiles(molecule_list) + if raw_preds is None: + return [], np.full((len(molecule_list), 0), np.nan, dtype=SCORE_DTYPE) + classes = [str(label) for label in self.predictor._classification_labels] + scores = raw_preds.detach().cpu().numpy().astype(SCORE_DTYPE) + return classes, scores + def calculate_results(self, batch): collator = self.predictor._dm.reader.COLLATOR() dat = self.predictor._model._process_batch( From ea8177b1f07366b5695970ab43766ccdbe465906 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 17:18:19 +0200 Subject: [PATCH 07/15] add hyperparameter optimization for WMV-F1 --- chebifier/ensemble/voting_ensemble.py | 15 +- .../ensemble/weighted_majority_ensemble.py | 155 +++++++++++++++++- chebifier/prediction_models/nn_predictor.py | 26 +-- 3 files changed, 175 insertions(+), 21 deletions(-) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 6b23d5f..28e4e9c 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -16,6 +16,7 @@ def __init__( super().__init__(ensemble_dir) self.use_confidence = use_confidence self.classwise_f1 = None + self.prediction_thresholds = None def find_best_threshold(self, predictions, val_labels_tensor): best_threshold = 0.5 @@ -37,11 +38,15 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): self.classwise_f1 = F1Score( task="multilabel", num_labels=validation_labels.shape[1], average=None ) - prediction_thresholds = { - model_name: self.find_best_threshold(predictions, validation_labels) - for model_name, predictions in validation_predictions.items() + 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() } - self._save_prediction_thresholds(prediction_thresholds) def _save_prediction_thresholds(self, thresholds: dict[str, float]): thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" @@ -50,6 +55,8 @@ def _save_prediction_thresholds(self, thresholds: dict[str, float]): 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: diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index ff51020..eca7b41 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,9 +1,13 @@ from pathlib import Path +import pandas as pd import torch from chebifier.ensemble.voting_ensemble import VotingEnsemble +N_FOLDS = 5 +WEIGHTING_STRENGTH_GRID = [0, 0.25, 0.5, 0.75, 1] + class WMVwithF1Ensemble(VotingEnsemble): @@ -22,17 +26,27 @@ def __init__( super().__init__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_f1_scores = None 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() - for model_name, predictions in validation_predictions.items(): - f1 = self.classwise_f1( - predictions > thresholds[model_name], validation_labels - ) + 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()) @@ -41,6 +55,8 @@ def _save_classwise_f1(self, validation_predictions, validation_labels): ) 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: @@ -72,3 +88,134 @@ def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: self.weighting_strength * classwise_f1 + (1 - self.weighting_strength) ) ** self.weighting_exponent return trust_tensor + + 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, + ): + """Macro F1 of the aggregated predictions on each held-out fold.""" + 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" + ] + scores.append( + self.classwise_f1(decisions, validation_labels[test_idx]).mean().item() + ) + self.prediction_thresholds = None + self.model_f1_scores = None + return scores + + 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/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 7e4e1b4..53ab46b 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -10,7 +10,7 @@ from .base_predictor import SCORE_DTYPE, BasePredictor if TYPE_CHECKING: - from torch import Tensor + pass class NNPredictor(BasePredictor, ABC): @@ -36,18 +36,16 @@ 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) + raw_preds = 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)) + { + label: pred + for label, pred in zip( + self.predictor._classification_labels, pred_tensor.tolist() + ) + } + for pred_tensor in raw_preds ] return preds else: @@ -56,11 +54,13 @@ def predict_list(self, smiles_list: list[str]) -> list: def predict_dense( self, molecule_list: list[str | Chem.Mol] ) -> tuple[list[str], np.ndarray]: - raw_preds: Tensor = self.predictor.predict_smiles(molecule_list) + raw_preds = self.predictor.predict_smiles(molecule_list) if raw_preds is None: return [], np.full((len(molecule_list), 0), np.nan, dtype=SCORE_DTYPE) classes = [str(label) for label in self.predictor._classification_labels] - scores = raw_preds.detach().cpu().numpy().astype(SCORE_DTYPE) + scores = np.stack([pred.detach().cpu().numpy() for pred in raw_preds]).astype( + SCORE_DTYPE + ) return classes, scores def calculate_results(self, batch): From 844768a6ba2a18b75b4f1050c1daad0fab5f7239 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 17:24:30 +0200 Subject: [PATCH 08/15] apply optimized hyperparameters if available --- .../ensemble/weighted_majority_ensemble.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index eca7b41..56d08da 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -15,13 +15,17 @@ def __init__( self, ensemble_dir: str, use_confidence: bool = True, - weighting_strength=1, - weighting_exponent=1, + weighting_strength=None, + weighting_exponent=None, **kwargs, ): """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. + + 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__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength @@ -66,9 +70,26 @@ def _load_classwise_f1(self, model_name: str) -> torch.Tensor: 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 + 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] @@ -85,8 +106,8 @@ def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: # 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] = ( - self.weighting_strength * classwise_f1 + (1 - self.weighting_strength) - ) ** self.weighting_exponent + weighting_strength * classwise_f1 + (1 - weighting_strength) + ) ** weighting_exponent return trust_tensor def _build_folds(self, validation_predictions, validation_labels): From 4ea20f64ad0c13a821d979309fac9cfeb71e91f3 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 14:51:26 +0200 Subject: [PATCH 09/15] fix handling of missing predictions --- chebifier/prediction_models/nn_predictor.py | 35 ++++++++++++--------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 53ab46b..a39eff6 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -36,31 +36,36 @@ 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 = self.predictor.predict_smiles(smiles_list) - if raw_preds is not None: - preds = [ - { + 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 - ] - return preds - else: - return [None for _ in smiles_list] + ) + 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_smiles(molecule_list) - if raw_preds is None: - return [], np.full((len(molecule_list), 0), np.nan, dtype=SCORE_DTYPE) + raw_preds = self.predictor.predict_molecules(molecule_list) classes = [str(label) for label in self.predictor._classification_labels] - scores = np.stack([pred.detach().cpu().numpy() for pred in raw_preds]).astype( - SCORE_DTYPE - ) + # 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): From 32e14a154ae8306673b08ff81a129f414f2d3061 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 14:53:27 +0200 Subject: [PATCH 10/15] use chebi_utils SMILES / InChI parsing --- chebifier/cli.py | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/chebifier/cli.py b/chebifier/cli.py index 49d345b..a9f9901 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -5,7 +5,7 @@ import click import pandas as pd import yaml -from rdkit import Chem +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 @@ -22,24 +22,7 @@ def read_molecules(molecules, molecule_file): with open(molecule_file, "r", encoding="utf-8") as f: raw_inputs.extend([line.strip() for line in f if line.strip()]) - mol_list = [] - for raw_input in raw_inputs: - try: - if raw_input.startswith("InChI="): - mol = Chem.MolFromInchi(raw_input, sanitize=False) - if mol is None: - click.echo(f"Failed to parse InChI: {raw_input}") - mol_list.append(None) - mol_list.append(mol) - elif Chem.MolFromSmiles(raw_input, sanitize=False) is None: - click.echo(f"Failed to parse SMILES: {raw_input}") - mol_list.append(None) - else: - mol_list.append(Chem.MolFromSmiles(raw_input, sanitize=False)) - except Exception as e: - click.echo(f"Error parsing molecule '{raw_input}': {e}.") - mol_list.append(None) - return mol_list + return [smiles_or_inchi_to_mol(raw_input) for raw_input in raw_inputs] def build_base_learners(ensemble_config): From 4334525a5921d89672ebba69f41d04db1d580411 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 6 Aug 2026 15:01:25 +0200 Subject: [PATCH 11/15] add DES and LTR ensembles, major fixes to voting ensemble (rescaled confidence) and inconsistency resolution (transitive relations and is-a subgraph) --- README.md | 54 +- chebifier/cli.py | 36 +- chebifier/ensemble/base_ensemble.py | 9 +- .../ensemble/dynamic_selection_ensemble.py | 680 ++++++++++++++++++ .../ensemble/learning_to_rank_ensemble.py | 251 +++++++ chebifier/ensemble/level1.py | 150 ++++ chebifier/ensemble/voting_ensemble.py | 13 +- .../ensemble/weighted_majority_ensemble.py | 10 +- chebifier/inconsistency_resolution.py | 64 +- chebifier/model_registry.py | 4 + chebifier/predict.py | 28 +- chebifier/utils.py | 2 +- pyproject.toml | 3 + 13 files changed, 1247 insertions(+), 57 deletions(-) create mode 100644 chebifier/ensemble/dynamic_selection_ensemble.py create mode 100644 chebifier/ensemble/learning_to_rank_ensemble.py create mode 100644 chebifier/ensemble/level1.py diff --git a/README.md b/README.md index 8d75fcd..34c7f29 100644 --- a/README.md +++ b/README.md @@ -142,14 +142,25 @@ $$ $$ --> -Here, confidence is the model's (self-reported) confidence in its prediction, calculated as +Here, confidence is the model's (self-reported) confidence in its prediction. Each model has its own +decision threshold $t_{m_i}$, calibrated on the validation set (see below), and confidence measures +how far the prediction sits from that threshold — scaled separately on each side, so that a +maximally confident negative ($p = 0$) and a maximally confident positive ($p = 1$) both count 1: $ -\text{confidence}_c^{m_i} = 2|p_c^{m_i} - 0.5| +\text{confidence}_c^{m_i} = \begin{cases} +(t_{m_i} - p_c^{m_i}) / t_{m_i} & \text{if } p_c^{m_i} < t_{m_i} \\ +(p_c^{m_i} - t_{m_i}) / (1 - t_{m_i}) & \text{otherwise} +\end{cases} $ -For example, if a model makes a positive prediction with $p_c^{m_i} = 0.55$, the confidence is $2|0.55 - 0.5| = 0.1$. -One could say that the model is not very confident in its prediction and very close to switching to a negative prediction. -If another model is very sure about its negative prediction with $p_c^{m_j} = 0.1$, the confidence is $2|0.1 - 0.5| = 0.8$. -Therefore, if in doubt, we are more confident in the negative prediction. +For example, for a model with $t_{m_i} = 0.5$ and a positive prediction of $p_c^{m_i} = 0.55$, the +confidence is $(0.55 - 0.5)/0.5 = 0.1$. One could say that the model is not very confident in its +prediction and very close to switching to a negative prediction. If another model is very sure about +its negative prediction with $p_c^{m_j} = 0.1$ (and $t_{m_j} = 0.5$), the confidence is +$(0.5 - 0.1)/0.5 = 0.8$. Therefore, if in doubt, we are more confident in the negative prediction. + +The two-sided scaling matters whenever a model's threshold is not 0.5: with $t_{m_i} = 0.2$, a +negative prediction only has a range of $0.2$ to move in and a positive one a range of $0.8$, so +without rescaling the positive side would systematically outweigh the negative side. Confidence can be disabled by the `use_confidence` parameter of the predict method (default: True). @@ -159,6 +170,37 @@ model independently of a given class. on a validation set for each class. If the `ensemble_type` is set to `wmv-f1`, the trust is calculated as F1-score $^{6.25}$. If the `ensemble_type` is set to `mv` (the default), the trust is set to 1 for all models. +#### Learned aggregation (`ltr` and `des`) + +Two further `ensemble_type`s replace the fixed voting rule by a model that is fitted on the +validation split. Both restrict themselves to a candidate set (per molecule, the union of each +base learner's top-`candidate_k` classes) and both emit the same net score as the voting +ensembles, so inconsistency resolution and the decision threshold apply unchanged. + +- `ltr` — **learning to rank**, an adaptation of + [GOLabeler](https://doi.org/10.1093/bioinformatics/bty130): the base learner scores for a + (molecule, class) pair become the feature vector of a LambdaMART ranker (LightGBM) that ranks + ChEBI classes per molecule. Features are the raw base learner scores plus the number of covering + models and the max/mean/std over them; a global cutoff on the ranker score is calibrated on a + held-out 20% of the validation split. +- `des` — **dynamic ensemble selection**, an adaptation of + [META-DES.H](https://arxiv.org/pdf/1811.01742): a `GaussianNB` meta-classifier estimates, per + (molecule, class, base learner), how competent that base learner is *for this molecule*, and only + the competent ones vote, weighted by that competence. Competence is described by the paper's five + meta-feature sets over two neighbourhoods — the `region_size` nearest molecules by Tanimoto + similarity on ECFP4, and the `profile_size` nearest output profiles. Because the neighbourhoods + are looked up at prediction time, calibration stores the reference predictions, labels and + fingerprints in the ensemble directory (~1 GB for a 20-model ensemble on ChEBI50). + +Both calibrate their hyperparameters by 5-fold cross-validation on the validation split, scoring +macro-F1 on each held-out fold (the cutoff is tuned on a fold-internal dev set, so the reported +score is not tuned on the fold it is measured on). Only the parameters that moved the result in +previous experiments are searched: `candidate_k` for `ltr`, and `region_size` / `profile_size` / +`vote` for `des`. The ranker's own tree hyperparameters, and `des`'s consensus and competence +thresholds, sit on a plateau and are left at their published values. Passing any searched parameter +to the constructor skips the search for it. Results are written to `hyperparameter_search.csv` and +`best_hyperparameters.csv` in the ensemble directory, as for `wmv-f1`. + ### Inconsistency resolution After a decision has been made for each class independently, the consistency of the predictions with regard to the ChEBI hierarchy and disjointness axioms is checked. This is diff --git a/chebifier/cli.py b/chebifier/cli.py index a9f9901..8c72b03 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -3,6 +3,7 @@ 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 @@ -177,12 +178,18 @@ def build( default=True, help="Resolve inconsistencies in the aggregated predictions (default: True)", ) +@click.option( + "--split", + type=click.Choice(["validation", "test"]), + default="test", + help="Dataset split to evaluate on (default: test)", +) @click.option( "--output", "-o", type=click.Path(), default=None, - help="Output file to save the evaluation results (optional)", + help="Output file for the ensemble predictions (default: /_predictions.npz)", ) def evaluate( ensemble_config, @@ -191,26 +198,41 @@ def evaluate( prediction_cache_dir, data_path, resolve_inconsistencies, + split, output, ): - """Evaluate an ensemble on the ChEBI test set.""" + """Store the predictions of an ensemble on a ChEBI dataset split.""" base_learners = build_base_learners(ensemble_config) ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) # TODO: Hugging Face support - test_data, test_labels = load_dataset(data_path, split="test") + eval_data, eval_labels = load_dataset(data_path, split=split) predictions = predict_molecules( base_learners, ensemble_model, - test_data, + eval_data, prediction_cache_dir=prediction_cache_dir, resolve_inconsistencies=resolve_inconsistencies, + classes=[str(cls) for cls in eval_labels.columns], + split=split, ) - print(f"Predictions: {predictions}") - - # TODO: compare predictions to test_labels, report metrics and save them to output + if output is None: + output = os.path.join(ensemble_dir, f"{split}_predictions.npz") + os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True) + np.savez_compressed( + output, + classes=np.array(predictions["predicted_classes"]), + scores=predictions["net_score"].numpy(), + decisions=predictions["class_decisions"].numpy(), + has_valid_predictions=predictions["has_valid_predictions"].numpy(), + ) + print( + f"Saved {ensemble_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 {output}." + ) @cli.command() diff --git a/chebifier/ensemble/base_ensemble.py b/chebifier/ensemble/base_ensemble.py index 67f9fa3..a81d0b3 100644 --- a/chebifier/ensemble/base_ensemble.py +++ b/chebifier/ensemble/base_ensemble.py @@ -36,5 +36,12 @@ def calibrate( """ pass - def predict(self, test_predictions: dict[str, torch.Tensor]): + def predict( + self, + test_predictions: dict[str, torch.Tensor], + molecules: list[Chem.Mol] | None = None, + ): + """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..90000c0 --- /dev/null +++ b/chebifier/ensemble/dynamic_selection_ensemble.py @@ -0,0 +1,680 @@ +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, + 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 +REGION_SIZE_GRID = (1, 7) +PROFILE_SIZE_GRID = (5, 9, 15) +VOTE_GRID = ("plain", "confidence") +CHUNK_SIZE = 512 + + +def fingerprints(molecules, radius=MORGAN_RADIUS, n_bits=MORGAN_BITS): + from rdkit import Chem, RDLogger + from rdkit.Chem import rdFingerprintGenerator + + RDLogger.DisableLog("rdApp.*") + generator = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits) + 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 +): + 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) + + X = np.concatenate( + [ + 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], + ], + 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) + return (weight * direction).sum(axis=1).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, + **kwargs, + ): + super().__init__(ensemble_dir) + 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) + 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]) + + self._save( + classifier, + scores, + labels, + molecule_fingerprints, + dsel_mask, + coverage, + { + "model_names": model_names, + "candidate_k": int(self.candidate_k), + "region_size": int(region_size), + "profile_size": int(profile_size), + "vote": vote, + "consensus_threshold": float(self.consensus_threshold), + "competence_threshold": float(self.competence_threshold), + "tau": float(tau), + "n_classes": int(scores.shape[1]), + "n_dsel": int(dsel_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}, 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 _fit_meta_classifier( + self, candidates, region, profiles, dsel, coverage, molecule_mask, labels + ): + from sklearn.naive_bayes import GaussianNB + + classifier = GaussianNB() + meta_classes = np.array([0, 1]) + fitted = False + 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 + ) + 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 + classifier.partial_fit( + chunk.X[mask], chunk.alpha[mask], classes=meta_classes + ) + fitted = True + if not fitted: + raise RuntimeError( + "No meta-training samples survived the consensus filter. Increase " + "consensus_threshold or check the base learner predictions." + ) + 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 + ) + 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) + 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"] + ) + 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=MORGAN_BITS + ) + self._coverage = data["coverage"] + self.consensus_threshold = self._metadata["consensus_threshold"] + self.competence_threshold = self._metadata["competence_threshold"] + + 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) + 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 - self._metadata["tau"], + scores.shape[:2], + ) + return { + "net_score": torch.from_numpy(dense), + "has_valid_predictions": torch.from_numpy((~np.isnan(scores)).any(axis=2)), + } diff --git a/chebifier/ensemble/learning_to_rank_ensemble.py b/chebifier/ensemble/learning_to_rank_ensemble.py new file mode 100644 index 0000000..82afeef --- /dev/null +++ b/chebifier/ensemble/learning_to_rank_ensemble.py @@ -0,0 +1,251 @@ +import json +from pathlib import Path + +import numpy as np +import torch + +from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.level1 import ( + N_FOLDS, + RANDOM_SEED, + candidate_pairs, + cv_folds, + dense_from_pairs, + holdout_split, + pair_scorer, + save_hyperparameter_results, + select_candidates, + stack_predictions, +) + +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 + + +class LearningToRankEnsemble(BaseEnsemble): + + def __init__( + self, + ensemble_dir: str, + candidate_k=None, + candidate_k_grid=CANDIDATE_K_GRID, + n_estimators: int = 500, + **kwargs, + ): + super().__init__(ensemble_dir) + self.candidate_k = candidate_k + self.candidate_k_grid = tuple(candidate_k_grid) + self.n_estimators = n_estimators + self._booster = None + self._metadata = 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" + + def calibrate(self, validation_predictions, validation_data, validation_labels): + 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) + + candidate_k = self.candidate_k + best = None + if candidate_k is None: + candidate_k, best = self._optimize_candidate_k(scores, labels) + + 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 = self._fit(rows, labels, train_mask, dev_mask) + + booster.save_model(str(self._model_path), num_iteration=booster.best_iteration) + metadata = { + "model_names": model_names, + "candidate_k": int(candidate_k), + "tau": float(tau), + "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})." + ) + + def _fit(self, rows, labels, train_mask, dev_mask): + import lightgbm as lgb + + train_rows = train_mask[rows["molecule"]] + dev_rows = dev_mask[rows["molecule"]] + train_set = lgb.Dataset( + rows["X"][train_rows], + label=rows["y"][train_rows], + group=rows["group"][train_mask], + free_raw_data=False, + ) + dev_set = lgb.Dataset( + rows["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( + rows["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 + + def _score_candidate_k(self, rows, labels, folds): + 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, _ = self._fit(rows, labels, train_mask, inner_dev_mask) + test_rows = test_mask[rows["molecule"]] + net = booster.predict( + rows["X"][test_rows], num_iteration=booster.best_iteration + ).astype(np.float32) + scorer, _ = pair_scorer( + labels, rows["molecule"], rows["class_index"], test_mask + ) + scores.append(scorer.macro_f1(net, tau)) + return scores + + def _optimize_candidate_k(self, scores, labels): + 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) + 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)) + + 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"]) + net = ( + self._booster.predict(rows["X"]).astype(np.float32) - self._metadata["tau"] + ) + dense = dense_from_pairs( + rows["molecule"], rows["class_index"], net, scores.shape[:2] + ) + 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..8647e4d --- /dev/null +++ b/chebifier/ensemble/level1.py @@ -0,0 +1,150 @@ +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 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 = 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 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 index 28e4e9c..d56f882 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -21,7 +21,7 @@ def __init__( def find_best_threshold(self, predictions, val_labels_tensor): best_threshold = 0.5 best_f1 = 0.0 - for threshold in range(0, 100): + for threshold in range(1, 100): threshold_value = threshold / 100 macro_f1_score = self.classwise_f1( predictions > threshold_value, val_labels_tensor @@ -70,7 +70,7 @@ def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: # No trust for MV, only used in WMV return 1 - def predict(self, test_predictions: dict[str, torch.Tensor]): + def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): """ Aggregates predictions from multiple models using weighted majority voting. weights are only the self-reported confidence (=difference between prediction and threshold). If set to false, all models are weighted equally. @@ -110,9 +110,12 @@ def predict(self, test_predictions: dict[str, torch.Tensor]): ) & valid_predictions if self.use_confidence: - confidence = 2 * torch.abs( - predictions_tensor.nan_to_num() - - threshold_mask.unsqueeze(0).unsqueeze(0) + threshold = threshold_mask.unsqueeze(0).unsqueeze(0) + scores = predictions_tensor.nan_to_num() + confidence = torch.where( + scores < threshold, + (threshold - scores) / threshold, + (scores - threshold) / (1 - threshold), ) else: confidence = torch.ones_like(predictions_tensor) diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 56d08da..ed48c03 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -146,7 +146,9 @@ def _score_hyperparameters( weighting_strength, weighting_exponent, ): - """Macro F1 of the aggregated predictions on each held-out fold.""" + """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 = [] @@ -162,9 +164,9 @@ def _score_hyperparameters( decisions = (aggregated["net_score"] > 0) & aggregated[ "has_valid_predictions" ] - scores.append( - self.classwise_f1(decisions, validation_labels[test_idx]).mean().item() - ) + 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 diff --git a/chebifier/inconsistency_resolution.py b/chebifier/inconsistency_resolution.py index f442de8..fc613a3 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,8 +47,8 @@ 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 @@ -67,15 +67,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,20 +99,21 @@ 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=0.0) + ) 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 @@ -121,7 +129,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 @@ -145,12 +153,14 @@ class ScoreBasedPredictionSmoother(PredictionSmoother): 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 + 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 the one with the higher absolute value - preds_direction = preds_optimistic - preds_pessimistic > 0 + preds_direction = preds_optimistic.abs() > preds_pessimistic.abs() return torch.where(preds_direction, preds_optimistic, preds_pessimistic) diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index 8d737de..787df47 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,3 +1,5 @@ +from chebifier.ensemble.dynamic_selection_ensemble import DynamicSelectionEnsemble +from chebifier.ensemble.learning_to_rank_ensemble import LearningToRankEnsemble from chebifier.ensemble.voting_ensemble import VotingEnsemble from chebifier.ensemble.weighted_majority_ensemble import WMVwithF1Ensemble from chebifier.prediction_models import ( @@ -18,6 +20,8 @@ ENSEMBLES = { "mv": VotingEnsemble, "wmv-f1": WMVwithF1Ensemble, + "ltr": LearningToRankEnsemble, + "des": DynamicSelectionEnsemble, } diff --git a/chebifier/predict.py b/chebifier/predict.py index 18a6fa8..b2c545a 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -14,10 +14,20 @@ from chebifier.utils import get_disjoint_files, load_chebi_graph -def apply_inconsistency_resolution(smoother, class_names, aggregated_predictions): +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) - smooth_net_score = smoother(aggregated_predictions["net_score"]) - aggregated_predictions["net_score"] = smooth_net_score + net_score = aggregated_predictions["net_score"] + aggregated_predictions["net_score"] = torch.cat( + [ + smoother(net_score[start : start + batch_size]) + for start in range(0, net_score.shape[0], batch_size) + ] + ) return aggregated_predictions @@ -104,6 +114,8 @@ def predict( prediction_cache_dir: Optional[str] = None, resolve_inconsistencies: bool = True, decision_threshold: float = 0, + classes: Optional[list[str]] = None, + split: str = "test", ) -> dict: """ Get end-to-end predictions from base learners and an ensemble model. @@ -117,6 +129,10 @@ def predict( -> 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. decision_threshold (float): Threshold for class decisions based on net score. Default is 0. + 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. @@ -129,7 +145,7 @@ def predict( test_predictions[model_name] = model.predict_dense(molecules) else: cache_path = os.path.join( - prediction_cache_dir, f"{model_name}_test_predictions.npz" + prediction_cache_dir, f"{model_name}_{split}_predictions.npz" ) if os.path.exists(cache_path): test_predictions[model_name] = load_dense_predictions(cache_path) @@ -138,11 +154,11 @@ def predict( save_dense_predictions(cache_path, *test_predictions[model_name]) test_predictions, predicted_classes = collect_base_learner_predictions( - test_predictions + test_predictions, classes=classes ) # Step 2: Get aggregated predictions from the ensemble model - aggregated_predictions = ensemble_model.predict(test_predictions) + aggregated_predictions = ensemble_model.predict(test_predictions, molecules) # net_score, has_valid_predictions, intermediate_results_dict # Step 3: Optionally resolve inconsistencies in the aggregated predictions diff --git a/chebifier/utils.py b/chebifier/utils.py index 8bfa6f3..91b0450 100644 --- a/chebifier/utils.py +++ b/chebifier/utils.py @@ -17,7 +17,7 @@ def load_chebi_graph(filename=None): { "repo_id": "chebai/chebifier", "repo_type": "dataset", - "files": {"f": "chebi_graph.pkl"}, + "files": {"f": "chebi_graph_v252.pkl"}, } )["f"] else: diff --git a/pyproject.toml b/pyproject.toml index d73db76..df2f773 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,9 @@ dependencies = [ "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", From ea3afb15c28b7bf69c0522bd475867ecdc6a1d05 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Fri, 7 Aug 2026 14:18:39 +0200 Subject: [PATCH 12/15] add HEX and ILR inconsistency resolution methods, evaluate multiple ensembles at once --- README.md | 38 ++++ chebifier/build_ensemble.py | 5 +- chebifier/cli.py | 244 +++++++++++++++++++++----- chebifier/hex_graph.py | 210 ++++++++++++++++++++++ chebifier/ilr.py | 142 +++++++++++++++ chebifier/inconsistency_resolution.py | 54 +++++- chebifier/predict.py | 180 ++++++++++++++----- 7 files changed, 778 insertions(+), 95 deletions(-) create mode 100644 chebifier/hex_graph.py create mode 100644 chebifier/ilr.py diff --git a/README.md b/README.md index 34c7f29..94814c1 100644 --- a/README.md +++ b/README.md @@ -215,3 +215,41 @@ both, we select one with the higher class score and set the other to 0. with a small change. For a pair of classes $A \subseteq B$ with predictions $1$ and $0$, instead of setting $B$ to $1$, we now set $A$ to $0$. This has the advantage that we cannot introduce new disjointness-inconsistencies and don't have to repeat step 2. + +#### Alternative methods + +The method above is `--inconsistency-resolution score-based` (`-ir`, the default). Two alternatives +from the literature are available at the same point in the pipeline; all of them consume the net +score and return a net score, so the decision threshold applies unchanged. + +- `ilr-godel`, `ilr-lukasiewicz` — **Iterative Local Refinement** + ([Daniele et al. 2023](https://doi.org/10.1007/s10994-023-06310-3)). Subsumption becomes the + implication $A \rightarrow B$ and disjointness the formula $\neg (A \wedge B)$, both as hard + constraints ($\hat t = 1$). Each constraint is repaired by its *minimal refinement function* — + the closest truth vector satisfying it — and the repairs are iterated to a fixpoint instead of + running the fixed 3-step schedule above. The two variants differ in how they split a violation: + Gödel is winner-take-all (it raises the parent to the child, and zeroes the weaker side of a + disjoint pair), whereas Łukasiewicz shares the correction — a disjointness violation with scores + $0.8$ and $0.7$ becomes $0.55$ and $0.45$ rather than $0.8$ and $0$. +- `hex` — **HEX graphs** + ([Deng et al. 2014](https://doi.org/10.1007/978-3-319-10590-1_4)). A CRF over binary label + vectors in which hierarchy edges forbid $(B, A) = (0, 1)$ and exclusion edges forbid + $(1, 1)$. Illegal states have probability zero, so the marginals satisfy + $P(A) \le P(B)$ for $A \subseteq B$ and $P(A) + P(B) \le 1$ for disjoint $A, B$ by construction. + +Both convert the net score to a probability with $p = \sigma(k \cdot \mathrm{score})$ and back with +$\mathrm{logit}(p) / k$, so the decision boundary stays at $0$. `k` (and `delta` for `hex`) can be +passed with `-irp k=2.0` and tuned with `scripts/calibrate_resolution.py`. Note that `k` does not +affect `ilr-godel`'s decisions: every Gödel operation is order-preserving, so a monotone +reparametrisation cannot change the sign of any score. + +Applied as published, HEX inference is intractable here. Its cost is bounded by +$O(\min(|V|2^w, |V|2^{\Omega}))$, and on a 2117-class ChEBI label set the maximum overlap is +$\Omega = 2115$ and the junction tree width is $\le 62$, with over 5 million legal states in the +largest cliques — the paper's efficiency argument assumes labels are mostly mutually exclusive, +whereas ChEBI labels overwhelmingly overlap (~25 classes hold per molecule). This implementation +therefore *clamps*: classes whose score is further than `delta` from the boundary, and which are +not involved in a violation, are fixed to their sign; that assignment is propagated to a fixpoint; +and exact inference runs only on the connected components of what remains (typically fewer than 30 +classes). Components above `max_component_size` fall back to `score-based`, counted in +`n_fallbacks`. This is a deviation from the published method and should be reported as such. diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index 64c6b0c..06b5c42 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -3,6 +3,7 @@ import torch from chebifier.predict import ( + base_learner_cache_path, collect_base_learner_predictions, load_dense_predictions, save_dense_predictions, @@ -49,8 +50,8 @@ def build_ensemble(self): classes = {} # get cached predictions if available, otherwise compute and cache them for model_name, model in self.base_learners.items(): - cache_path = os.path.join( - self.prediction_cache_dir, f"{model_name}_validation_predictions.npz" + 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...") diff --git a/chebifier/cli.py b/chebifier/cli.py index 8c72b03..0f3cb3c 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -11,9 +11,17 @@ 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.utils import get_default_configs, load_chebi_graph, process_config +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): @@ -26,8 +34,13 @@ def read_molecules(molecules, molecule_file): return [smiles_or_inchi_to_mol(raw_input) for raw_input in raw_inputs] -def build_base_learners(ensemble_config): - """Instantiate the base learners described by an ensemble configuration file.""" +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: @@ -42,9 +55,21 @@ def build_base_learners(ensemble_config): ): model_registry = yaml.safe_load(f) - chebi_graph = load_chebi_graph() + 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: @@ -60,6 +85,19 @@ def build_base_learners(ensemble_config): 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 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") @@ -88,8 +126,8 @@ def load_dataset(data_path, split: Literal["train", "validation", "test"]): return mol_list, labels_df -def ensemble_options(command): - """Options shared by all commands that use an ensemble.""" +def base_learner_options(command): + """Options shared by all commands that use base learners.""" for option in reversed( [ click.option( @@ -99,6 +137,22 @@ def ensemble_options(command): 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", @@ -113,16 +167,10 @@ def ensemble_options(command): required=True, help="Directory where the calibration results of the ensemble are stored", ), - click.option( - "--prediction-cache-dir", - type=click.Path(), - default=None, - help="Directory for caching base learner predictions", - ), ] ): command = option(command) - return command + return base_learner_options(command) def data_options(command): @@ -171,12 +219,44 @@ def build( @cli.command() -@ensemble_options +@base_learner_options @data_options +@click.option( + "--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( + "--ensemble-dir", + "-d", + type=click.Path(), + multiple=True, + required=True, + help="Directory where the calibration results of the ensemble are stored (one per --ensemble-type)", +) @click.option( "--resolve-inconsistencies/--no-resolve-inconsistencies", default=True, - help="Resolve inconsistencies in the aggregated predictions (default: True)", + help="Resolve inconsistencies in the aggregated predictions (default: True). " + "--no-resolve-inconsistencies is equivalent to '-ir none'.", +) +@click.option( + "--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( + "--ir-param", + "-irp", + multiple=True, + help="Extra key=value parameter for the resolution method, e.g. -irp k=2.0 (repeatable)", ) @click.option( "--split", @@ -184,12 +264,20 @@ def build( default="test", help="Dataset split to evaluate on (default: test)", ) +@click.option( + "--skip-existing", + is_flag=True, + default=False, + help="Skip ensemble / resolution combinations whose output file already exists", +) @click.option( "--output", "-o", type=click.Path(), default=None, - help="Output file for the ensemble predictions (default: /_predictions.npz)", + 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, @@ -198,41 +286,92 @@ def evaluate( prediction_cache_dir, data_path, resolve_inconsistencies, + inconsistency_resolution, + ir_param, split, + skip_existing, output, ): - """Store the predictions of an ensemble on a ChEBI dataset split.""" - base_learners = build_base_learners(ensemble_config) - ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + """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." + ) - # TODO: Hugging Face support - eval_data, eval_labels = load_dataset(data_path, split=split) + 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") + + 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 - predictions = predict_molecules( - base_learners, - ensemble_model, - eval_data, - prediction_cache_dir=prediction_cache_dir, - resolve_inconsistencies=resolve_inconsistencies, - classes=[str(cls) for cls in eval_labels.columns], - split=split, + base_learners = build_base_learners( + ensemble_config, prediction_cache_dir=prediction_cache_dir, split=split ) - if output is None: - output = os.path.join(ensemble_dir, f"{split}_predictions.npz") - os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True) - np.savez_compressed( - output, - classes=np.array(predictions["predicted_classes"]), - scores=predictions["net_score"].numpy(), - decisions=predictions["class_decisions"].numpy(), - has_valid_predictions=predictions["has_valid_predictions"].numpy(), - ) - print( - f"Saved {ensemble_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 {output}." - ) + # TODO: Hugging Face support + eval_data, eval_labels = load_dataset(data_path, split=split) + + chebi_graph, disjoint_files = None, None + if any(variant != "none" for variant in variants): + chebi_graph = load_chebi_graph() + disjoint_files = get_disjoint_files() + + 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, + ) + 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(), + ) + 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}." + ) @cli.command() @@ -252,6 +391,13 @@ def evaluate( 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", @@ -259,6 +405,12 @@ def evaluate( 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", @@ -274,6 +426,8 @@ def predict( molecules, molecule_file, resolve_inconsistencies, + inconsistency_resolution, + ir_param, decision_threshold, output, ): @@ -292,6 +446,8 @@ def predict( 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, ) diff --git a/chebifier/hex_graph.py b/chebifier/hex_graph.py new file mode 100644 index 0000000..188abe8 --- /dev/null +++ b/chebifier/hex_graph.py @@ -0,0 +1,210 @@ +import numpy as np +import torch +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import connected_components + +from chebifier.inconsistency_resolution import ( + ScoreBasedPredictionSmoother, + densified_exclusion_matrix, + from_prob, +) + + +class HexSmoother(ScoreBasedPredictionSmoother): + def __init__( + self, + chebi_graph, + label_names=None, + disjoint_files=None, + verbose=False, + k=1.0, + delta=0.0, + max_states=2**20, + max_component_size=40, + ): + self.k = k + 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 > 0 + known = torch.ones_like(pos) if valid is None else valid + if valid is not None: + pos = pos & valid + active = ((scores.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]] = from_prob(torch.ones(1), self.k).item() + out[idx[forced_zero]] = from_prob(torch.zeros(1), self.k).item() + 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)] = from_prob(weights @ st, self.k) + return out + + def __call__(self, preds, valid_mask=None): + if preds.shape[1] == 0: + return preds + out = preds.clone() + f = self.k * 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..ecc46e3 --- /dev/null +++ b/chebifier/ilr.py @@ -0,0 +1,142 @@ +import torch + +from chebifier.inconsistency_resolution import ( + PredictionSmoother, + densified_exclusion_pairs, + from_prob, + to_prob, +) + + +class ILRSmoother(PredictionSmoother): + def __init__( + self, + chebi_graph, + label_names=None, + disjoint_files=None, + verbose=False, + k=1.0, + alpha=1.0, + max_iter=10, + tol=1e-4, + ): + self.k = k + 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 = to_prob(preds, self.k) + 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 from_prob(p, self.k) + + +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 fc613a3..fc33c73 100644 --- a/chebifier/inconsistency_resolution.py +++ b/chebifier/inconsistency_resolution.py @@ -53,6 +53,58 @@ def get_disjoint_groups(disjoint_files): return disjoint_all +def to_prob(scores, k): + return torch.sigmoid(k * scores) + + +def from_prob(p, k): + p = p.clamp(1e-6, 1 - 1e-6) + return (torch.log(p) - torch.log1p(-p)) / k + + +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""" @@ -118,7 +170,7 @@ def resolve_disjointness_violations(self, preds): 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 diff --git a/chebifier/predict.py b/chebifier/predict.py index b2c545a..16bafa7 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -9,7 +9,7 @@ from rdkit import Chem from chebifier.ensemble.base_ensemble import BaseEnsemble -from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother +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 @@ -22,15 +22,25 @@ def apply_inconsistency_resolution( """ 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]) + 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) @@ -107,66 +117,88 @@ def collect_base_learner_predictions( return collected_predictions, ensemble_classes -def predict( - base_learners: dict[str, BasePredictor], - ensemble_model: BaseEnsemble, +def get_base_learner_predictions( + base_learners: dict[str, Optional[BasePredictor]], molecules: list[str | Chem.Mol], prediction_cache_dir: Optional[str] = None, - resolve_inconsistencies: bool = True, - decision_threshold: float = 0, - classes: Optional[list[str]] = None, split: str = "test", -) -> dict: - """ - Get end-to-end predictions from base learners and an ensemble model. +) -> dict[str, tuple[list[str], np.ndarray]]: + """Get dense predictions from the base learners, using the cache where available. - 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. - decision_threshold (float): Threshold for class decisions based on net score. Default is 0. - 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. + 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. """ - - # Step 1: Get predictions from base learners on test data - test_predictions = {} + predictions = {} for model_name, model in base_learners.items(): - if prediction_cache_dir is None: - test_predictions[model_name] = model.predict_dense(molecules) - else: - cache_path = os.path.join( - prediction_cache_dir, f"{model_name}_{split}_predictions.npz" + 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})." ) - if os.path.exists(cache_path): - test_predictions[model_name] = load_dense_predictions(cache_path) - else: - test_predictions[model_name] = model.predict_dense(molecules) - save_dense_predictions(cache_path, *test_predictions[model_name]) + 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 ) - - # Step 2: Get aggregated predictions from the ensemble model 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, + chebi_graph=None, + disjoint_files=None, +) -> dict: + """Resolve inconsistencies in aggregated predictions and turn them into class decisions. - # Step 3: Optionally resolve inconsistencies in the aggregated predictions - if resolve_inconsistencies: - chebi_graph = load_chebi_graph() - disjoint_files = get_disjoint_files() - smoother = ScoreBasedPredictionSmoother( - chebi_graph=chebi_graph, label_names=None, disjoint_files=disjoint_files + `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 @@ -187,3 +219,55 @@ def predict( 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: float = 0, + 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 (float): Threshold for class decisions based on net score. Default is 0. + 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=decision_threshold, + ) From 8ec5a1cf9ec87f49cdde039030249c367ba98f7f Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 11 Aug 2026 09:41:58 +0200 Subject: [PATCH 13/15] more options for LTR / DES, normalize aggregation outputs --- README.md | 41 +++- chebifier/cli.py | 40 +++- chebifier/ensemble/base_ensemble.py | 9 + .../ensemble/dynamic_selection_ensemble.py | 206 +++++++++++++++--- .../ensemble/learning_to_rank_ensemble.py | 127 +++++++++-- chebifier/ensemble/level1.py | 66 +++++- chebifier/ensemble/voting_ensemble.py | 9 +- chebifier/predict.py | 11 +- 8 files changed, 441 insertions(+), 68 deletions(-) diff --git a/README.md b/README.md index 94814c1..6376a1c 100644 --- a/README.md +++ b/README.md @@ -182,7 +182,14 @@ ensembles, so inconsistency resolution and the decision threshold apply unchange (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. + 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 @@ -191,6 +198,31 @@ ensembles, so inconsistency resolution and the decision threshold apply unchange 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 @@ -198,7 +230,12 @@ score is not tuned on the fold it is measured on). Only the parameters that move previous experiments are searched: `candidate_k` for `ltr`, and `region_size` / `profile_size` / `vote` for `des`. The ranker's own tree hyperparameters, and `des`'s consensus and competence thresholds, sit on a plateau and are left at their published values. Passing any searched parameter -to the constructor skips the search for it. Results are written to `hyperparameter_search.csv` and +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 diff --git a/chebifier/cli.py b/chebifier/cli.py index 0f3cb3c..1f558e7 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -98,6 +98,25 @@ def parse_ir_params(ir_param): 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") @@ -198,12 +217,27 @@ def cli(): @cli.command() @ensemble_options @data_options +@click.option( + "--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.", +) def build( - ensemble_config, ensemble_type, ensemble_dir, prediction_cache_dir, data_path + 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) - ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + 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") @@ -357,6 +391,7 @@ def output_path(dir_, 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) @@ -366,6 +401,7 @@ def output_path(dir_, variant): 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}' " diff --git a/chebifier/ensemble/base_ensemble.py b/chebifier/ensemble/base_ensemble.py index a81d0b3..c01ee64 100644 --- a/chebifier/ensemble/base_ensemble.py +++ b/chebifier/ensemble/base_ensemble.py @@ -25,6 +25,15 @@ def __init__(self, ensemble_dir: str): def ensemble_name(self): return self.__class__.__name__ + @property + def decision_threshold(self): + """Value a net score has to exceed for the class to be predicted. + + Zero for every ensemble whose scores are already centred on their decision boundary; only + ensembles that emit a calibrated, unshifted scale need to override this. + """ + return 0.0 + def calibrate( self, validation_predictions: dict[str, torch.Tensor], diff --git a/chebifier/ensemble/dynamic_selection_ensemble.py b/chebifier/ensemble/dynamic_selection_ensemble.py index 90000c0..b653f4f 100644 --- a/chebifier/ensemble/dynamic_selection_ensemble.py +++ b/chebifier/ensemble/dynamic_selection_ensemble.py @@ -25,18 +25,27 @@ 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): +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) + 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): @@ -204,7 +213,14 @@ def __init__(self, X, covered, profiles, consensus, alpha, molecule): def build_chunk( - candidates, pair_slice, region, output_profiles, dsel, coverage, labels=None + candidates, + pair_slice, + region, + output_profiles, + dsel, + coverage, + labels=None, + model_id=False, ): first, last = pair_slice molecule = candidates.molecule[first:last] @@ -231,16 +247,21 @@ def build_chunk( filled_query = np.nan_to_num(query_scores, nan=POSITIVE_THRESHOLD) f5 = 2.0 * np.abs(filled_query - POSITIVE_THRESHOLD) - X = np.concatenate( - [ - 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], - ], - axis=2, - ) + 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: @@ -267,7 +288,10 @@ def aggregate(delta, chunk, competence_threshold, vote): direction = np.where(chunk.profiles > POSITIVE_THRESHOLD, 1.0, -1.0) if vote == "confidence": direction = direction * 2.0 * np.abs(chunk.profiles - POSITIVE_THRESHOLD) - return (weight * direction).sum(axis=1).astype(np.float32) + total = weight.sum(axis=1) + return ((weight * direction).sum(axis=1) / np.maximum(total, 1e-6)).astype( + np.float32 + ) class DynamicSelectionEnsemble(VotingEnsemble): @@ -285,9 +309,27 @@ def __init__( 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 @@ -325,7 +367,9 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): 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) + 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 " @@ -364,12 +408,15 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): ) tau, dev_macro_f1 = scorer.tune(net[keep]) + # 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, - dsel_mask, + reference_mask, coverage, { "model_names": model_names, @@ -377,17 +424,25 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): "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), "n_classes": int(scores.shape[1]), - "n_dsel": int(dsel_mask.sum()), + "n_dsel": int(reference_mask.sum()), "dev_macro_f1": float(dev_macro_f1), + "normalized": True, }, ) print( f"Saved meta-classifier to {self._classifier_path} (region_size={region_size}, " - f"profile_size={profile_size}, vote={vote}, tau={tau:.4f}, " + 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})." ) @@ -417,20 +472,24 @@ def _neighbourhoods( ) return region, profiles - def _fit_meta_classifier( + def _meta_samples( self, candidates, region, profiles, dsel, coverage, molecule_mask, labels ): - from sklearn.naive_bayes import GaussianNB - - classifier = GaussianNB() - meta_classes = np.array([0, 1]) - fitted = False + """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 + 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 @@ -440,17 +499,72 @@ def _fit_meta_classifier( mask = chunk.covered & keep[:, None] if not mask.any(): continue - classifier.partial_fit( - chunk.X[mask], chunk.alpha[mask], classes=meta_classes - ) - fitted = True - if not fitted: + 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( - "No meta-training samples survived the consensus filter. Increase " + 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, @@ -470,7 +584,13 @@ def _score( if molecule_mask is not None and not molecule_mask[start:end].any(): continue chunk = build_chunk( - candidates, pair_slice, region, profiles, dsel, coverage + 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(): @@ -612,19 +732,35 @@ def _load(self): ) with open(self._metadata_path, "r", encoding="utf-8") as f: self._metadata = json.load(f) + if not self._metadata.get("normalized"): + raise ValueError( + f"The ensemble in {self.ensemble_dir} was calibrated before net scores were " + "normalised by the selection weight. Its stored tau is on the old (unnormalised) " + "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=MORGAN_BITS + 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: @@ -636,7 +772,9 @@ def predict(self, test_predictions, molecules=None): scores, _ = stack_predictions( test_predictions, self._metadata["model_names"], dtype=np.float16 ) - query_fingerprints, parsed = fingerprints(molecules) + 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 diff --git a/chebifier/ensemble/learning_to_rank_ensemble.py b/chebifier/ensemble/learning_to_rank_ensemble.py index 82afeef..567e36f 100644 --- a/chebifier/ensemble/learning_to_rank_ensemble.py +++ b/chebifier/ensemble/learning_to_rank_ensemble.py @@ -4,19 +4,23 @@ import numpy as np import torch -from chebifier.ensemble.base_ensemble import BaseEnsemble from chebifier.ensemble.level1 import ( N_FOLDS, RANDOM_SEED, candidate_pairs, + class_statistics, cv_folds, dense_from_pairs, + fit_platt, holdout_split, + noncandidate_log_odds, 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 @@ -71,7 +75,13 @@ def build_rows(scores, k, labels=None): return rows -class LearningToRankEnsemble(BaseEnsemble): +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, @@ -79,14 +89,17 @@ def __init__( 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): @@ -96,27 +109,60 @@ def _model_path(self): 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): - print( - f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." - ) + 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) + 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 = self._fit(rows, labels, train_mask, dev_mask) + 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_logodds = noncandidate_log_odds( + 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_logodds": platt_a * float(tau) + platt_b, + "floor_logodds": floor_logodds, "n_classes": int(scores.shape[1]), "best_iteration": int(booster.best_iteration), "dev_macro_f1": float(dev_macro_f1), @@ -129,22 +175,30 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): 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})." + f"held-out macro-f1: {dev_macro_f1:.4f}). Calibrated to log-odds with " + f"a={platt_a:.4f}, b={platt_b:.4f} (decision threshold {metadata['tau_logodds']:.4f}); " + f"non-candidate pairs scored at {floor_logodds:.2f}." ) - def _fit(self, rows, labels, train_mask, dev_mask): + 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( - rows["X"][train_rows], + X[train_rows], label=rows["y"][train_rows], group=rows["group"][train_mask], free_raw_data=False, ) dev_set = lgb.Dataset( - rows["X"][dev_rows], + X[dev_rows], label=rows["y"][dev_rows], group=rows["group"][dev_mask], reference=train_set, @@ -159,32 +213,35 @@ def _fit(self, rows, labels, train_mask, dev_mask): callbacks=[lgb.early_stopping(EARLY_STOPPING_ROUNDS, verbose=False)], ) dev_scores = booster.predict( - rows["X"][dev_rows], num_iteration=booster.best_iteration + 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 + return booster, tau, dev_macro_f1, class_stats, dev_scores - def _score_candidate_k(self, rows, labels, folds): - 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, _ = self._fit(rows, labels, train_mask, inner_dev_mask) + booster, tau, _, class_stats, _ = self._fit( + rows, labels, train_mask, inner_dev_mask, scores, thresholds + ) test_rows = test_mask[rows["molecule"]] net = booster.predict( - rows["X"][test_rows], num_iteration=booster.best_iteration + 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 ) - scores.append(scorer.macro_f1(net, tau)) - return scores + fold_scores.append(scorer.macro_f1(net, tau)) + return fold_scores - def _optimize_candidate_k(self, scores, labels): + def _optimize_candidate_k(self, scores, labels, thresholds=None): print( f"Optimizing candidate_k with {N_FOLDS}-fold cross-validation on the validation set..." ) @@ -196,7 +253,9 @@ def _optimize_candidate_k(self, scores, labels): labels[rows["molecule"], rows["class_index"]].sum() / max(int(labels.sum()), 1) ) - fold_scores = self._score_candidate_k(rows, labels, folds) + fold_scores = self._score_candidate_k( + rows, labels, folds, scores, thresholds + ) mean_score = float(np.mean(fold_scores)) results.append( { @@ -234,16 +293,36 @@ def _load(self): 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): + """Decision threshold on the scale `predict` returns. + + Zero for rankers calibrated before Platt scaling was introduced - those still subtract tau + from the score themselves. + """ + self._load() + return float(self._metadata.get("tau_logodds", 0.0)) 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"]) - net = ( - self._booster.predict(rows["X"]).astype(np.float32) - self._metadata["tau"] + raw = self._booster.predict(design_matrix(rows, self._class_stats)).astype( + np.float32 ) + if "platt_a" in self._metadata: + # calibrated log-odds, with the decision threshold left to the caller so that the + # inconsistency resolution sees an unshifted scale + net = self._metadata["platt_a"] * raw + self._metadata["platt_b"] + floor = self._metadata["floor_logodds"] + else: + net, floor = raw - self._metadata["tau"], None dense = dense_from_pairs( - rows["molecule"], rows["class_index"], net, scores.shape[:2] + rows["molecule"], rows["class_index"], net, scores.shape[:2], floor=floor ) return { "net_score": torch.from_numpy(dense), diff --git a/chebifier/ensemble/level1.py b/chebifier/ensemble/level1.py index 8647e4d..55dde6e 100644 --- a/chebifier/ensemble/level1.py +++ b/chebifier/ensemble/level1.py @@ -46,6 +46,31 @@ 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) @@ -131,13 +156,50 @@ def pair_scorer(labels, molecule, class_index, molecule_mask): return PairScorer(labels[rows], remap[molecule[keep]], class_index[keep]), keep -def dense_from_pairs(molecule, class_index, net, shape): - floor = min(float(net.min()) - 1.0, -1.0) if net.size else -1.0 +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)` so that `a * score + b` is calibrated log-odds. + + 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 and + maps them through `sigmoid(k * score)`. 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. + """ + 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_log_odds(n_noncandidate, n_missed_positives): + """Log-odds 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 score that can be compared against the candidates. The Jeffreys-style pseudo-count + keeps the result finite when no positive is missed at all. + """ + p = (n_missed_positives + 0.5) / (n_noncandidate + 1.0) + return float(np.log(p / (1.0 - p))) + + 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) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index d56f882..d0fc4e8 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -74,6 +74,11 @@ def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): """ Aggregates predictions from multiple models using weighted majority voting. weights are only the self-reported confidence (=difference between prediction and threshold). If set to false, all models are weighted equally. + + The net score is normalised by the weight mass that was cast, so it is a signed agreement + fraction in [-1, 1] rather than a sum over models. This keeps classes covered by different + numbers of base learners comparable, which matters downstream: inconsistency resolution + compares scores across classes. The sign is unaffected, so class decisions do not change. """ predictions_tensor = torch.stack( list(test_predictions.values()), dim=2 @@ -136,7 +141,9 @@ def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): ) # Shape: (num_molecules, num_classes) # Determine which classes to include for each molecule - net_score = positive_sum - negative_sum # Shape: (num_molecules, num_classes) + net_score = (positive_sum - negative_sum) / (positive_sum + negative_sum).clamp( + min=1e-6 + ) # Shape: (num_molecules, num_classes) return { "net_score": net_score, "has_valid_predictions": has_valid_predictions, diff --git a/chebifier/predict.py b/chebifier/predict.py index 16bafa7..41ead5c 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -229,7 +229,7 @@ def predict( resolve_inconsistencies: bool = True, inconsistency_resolution: str = "score-based", inconsistency_resolution_params: Optional[dict] = None, - decision_threshold: float = 0, + decision_threshold: Optional[float] = None, classes: Optional[list[str]] = None, split: str = "test", ) -> dict: @@ -245,7 +245,8 @@ def predict( -> 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 (float): Threshold for class decisions based on net score. Default is 0. + 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 @@ -269,5 +270,9 @@ def predict( inconsistency_resolution if resolve_inconsistencies else "none" ), inconsistency_resolution_params=inconsistency_resolution_params, - decision_threshold=decision_threshold, + decision_threshold=( + ensemble_model.decision_threshold + if decision_threshold is None + else decision_threshold + ), ) From 89118306315574b3a2089e4d3077385fa6991d07 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 11 Aug 2026 11:56:22 +0200 Subject: [PATCH 14/15] ensemble output is now set to [0,1] range, same for inconsistency resolution input. mv and wmv-conf are now separate models --- README.md | 6 +- chebifier/__init__.py | 8 +- chebifier/ensemble/base_ensemble.py | 9 ++- .../ensemble/dynamic_selection_ensemble.py | 31 +++++--- .../ensemble/learning_to_rank_ensemble.py | 45 ++++++----- chebifier/ensemble/level1.py | 23 +++--- chebifier/ensemble/voting_ensemble.py | 76 +++++++++++++------ .../ensemble/weighted_majority_ensemble.py | 10 +-- chebifier/hex_graph.py | 21 ++--- chebifier/ilr.py | 8 +- chebifier/inconsistency_resolution.py | 32 +++++--- chebifier/model_registry.py | 8 +- chebifier/predict.py | 5 +- 13 files changed, 178 insertions(+), 104 deletions(-) diff --git a/README.md b/README.md index 6376a1c..525912b 100644 --- a/README.md +++ b/README.md @@ -162,13 +162,15 @@ The two-sided scaling matters whenever a model's threshold is not 0.5: with $t_{ 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`) diff --git a/chebifier/__init__.py b/chebifier/__init__.py index 34967ee..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.voting_ensemble import VotingEnsemble +from .ensemble.voting_ensemble import ( + MajorityVotingEnsemble, + VotingEnsemble, + WMVwithConfidenceEnsemble, +) __all__ = [ "VotingEnsemble", + "MajorityVotingEnsemble", + "WMVwithConfidenceEnsemble", "PerSmilesPerModelLRUCache", "modelwise_smiles_lru_cache", ] diff --git a/chebifier/ensemble/base_ensemble.py b/chebifier/ensemble/base_ensemble.py index c01ee64..6dc5bf4 100644 --- a/chebifier/ensemble/base_ensemble.py +++ b/chebifier/ensemble/base_ensemble.py @@ -27,12 +27,13 @@ def ensemble_name(self): @property def decision_threshold(self): - """Value a net score has to exceed for the class to be predicted. + """Probability a class has to exceed to be predicted. - Zero for every ensemble whose scores are already centred on their decision boundary; only - ensembles that emit a calibrated, unshifted scale need to override this. + 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.0 + return 0.5 def calibrate( self, diff --git a/chebifier/ensemble/dynamic_selection_ensemble.py b/chebifier/ensemble/dynamic_selection_ensemble.py index b653f4f..e49d14d 100644 --- a/chebifier/ensemble/dynamic_selection_ensemble.py +++ b/chebifier/ensemble/dynamic_selection_ensemble.py @@ -14,6 +14,7 @@ cv_folds, dense_from_pairs, holdout_split, + noncandidate_probability, pair_scorer, rescale_to_threshold, save_hyperparameter_results, @@ -289,9 +290,8 @@ def aggregate(delta, chunk, competence_threshold, vote): if vote == "confidence": direction = direction * 2.0 * np.abs(chunk.profiles - POSITIVE_THRESHOLD) total = weight.sum(axis=1) - return ((weight * direction).sum(axis=1) / np.maximum(total, 1e-6)).astype( - np.float32 - ) + 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): @@ -407,6 +407,11 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): 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 @@ -433,10 +438,10 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): "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), - "normalized": True, }, ) print( @@ -732,12 +737,11 @@ def _load(self): ) with open(self._metadata_path, "r", encoding="utf-8") as f: self._metadata = json.load(f) - if not self._metadata.get("normalized"): + if "floor" not in self._metadata: raise ValueError( - f"The ensemble in {self.ensemble_dir} was calibrated before net scores were " - "normalised by the selection weight. Its stored tau is on the old (unnormalised) " - "scale and would reject every prediction. Please re-run `chebifier build` for this " - "ensemble." + 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) @@ -809,10 +813,17 @@ def predict(self, test_predictions, molecules=None): dense = dense_from_pairs( candidates.molecule, candidates.class_index, - net - self._metadata["tau"], + 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 index 567e36f..17ae8e0 100644 --- a/chebifier/ensemble/learning_to_rank_ensemble.py +++ b/chebifier/ensemble/learning_to_rank_ensemble.py @@ -13,7 +13,7 @@ dense_from_pairs, fit_platt, holdout_split, - noncandidate_log_odds, + noncandidate_probability, pair_scorer, save_hyperparameter_results, select_candidates, @@ -75,6 +75,10 @@ def build_rows(scores, k, labels=None): 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"] @@ -142,7 +146,7 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): rows, labels, train_mask, dev_mask, scores, thresholds ) platt_a, platt_b = fit_platt(dev_scores, rows["y"][dev_mask[rows["molecule"]]]) - floor_logodds = noncandidate_log_odds( + floor = noncandidate_probability( scores.shape[0] * scores.shape[1] - rows["molecule"].size, int(labels.sum()) - int(rows["y"].sum()), ) @@ -161,8 +165,8 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): "tau": float(tau), "platt_a": platt_a, "platt_b": platt_b, - "tau_logodds": platt_a * float(tau) + platt_b, - "floor_logodds": floor_logodds, + "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), @@ -175,9 +179,9 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): 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 log-odds with " - f"a={platt_a:.4f}, b={platt_b:.4f} (decision threshold {metadata['tau_logodds']:.4f}); " - f"non-candidate pairs scored at {floor_logodds:.2f}." + 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): @@ -299,13 +303,9 @@ def _load(self): @property def decision_threshold(self): - """Decision threshold on the scale `predict` returns. - - Zero for rankers calibrated before Platt scaling was introduced - those still subtract tau - from the score themselves. - """ + """Probability a candidate has to exceed, the operating point macro-F1 peaks at.""" self._load() - return float(self._metadata.get("tau_logodds", 0.0)) + return float(self._metadata["tau_probability"]) def predict(self, test_predictions, molecules=None): self._load() @@ -314,15 +314,18 @@ def predict(self, test_predictions, molecules=None): raw = self._booster.predict(design_matrix(rows, self._class_stats)).astype( np.float32 ) - if "platt_a" in self._metadata: - # calibrated log-odds, with the decision threshold left to the caller so that the - # inconsistency resolution sees an unshifted scale - net = self._metadata["platt_a"] * raw + self._metadata["platt_b"] - floor = self._metadata["floor_logodds"] - else: - net, floor = raw - self._metadata["tau"], None + # 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, scores.shape[:2], floor=floor + rows["molecule"], + rows["class_index"], + net.astype(np.float32), + scores.shape[:2], + floor=self._metadata["floor"], ) return { "net_score": torch.from_numpy(dense), diff --git a/chebifier/ensemble/level1.py b/chebifier/ensemble/level1.py index 55dde6e..61b91ac 100644 --- a/chebifier/ensemble/level1.py +++ b/chebifier/ensemble/level1.py @@ -165,13 +165,15 @@ def dense_from_pairs(molecule, class_index, net, shape, floor=None): def fit_platt(scores, labels): - """Fit `P(positive) = sigmoid(a * score + b)` so that `a * score + b` is calibrated log-odds. + """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 and - maps them through `sigmoid(k * score)`. 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. + 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 @@ -187,17 +189,16 @@ def fit_platt(scores, labels): return slope, intercept -def noncandidate_log_odds(n_noncandidate, n_missed_positives): - """Log-odds that a pair outside the candidate set is nevertheless a positive. +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 score that can be compared against the candidates. The Jeffreys-style pseudo-count - keeps the result finite when no positive is missed at all. + 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. """ - p = (n_missed_positives + 0.5) / (n_noncandidate + 1.0) - return float(np.log(p / (1.0 - p))) + return float((n_missed_positives + 0.5) / (n_noncandidate + 1.0)) def save_hyperparameter_results(ensemble_dir, results, best): diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index d0fc4e8..cca29a5 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -8,13 +8,15 @@ 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, - use_confidence: bool = True, ): super().__init__(ensemble_dir) - self.use_confidence = use_confidence self.classwise_f1 = None self.prediction_thresholds = None @@ -66,19 +68,26 @@ def _load_prediction_thresholds(self) -> dict[str, float]: 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 for MV, only used in WMV + # 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 using weighted majority voting. - weights are only the self-reported confidence (=difference between prediction and threshold). If set to false, all models are weighted equally. - - The net score is normalised by the weight mass that was cast, so it is a signed agreement - fraction in [-1, 1] rather than a sum over models. This keeps classes covered by different - numbers of base learners comparable, which matters downstream: inconsistency resolution - compares scores across classes. The sign is unaffected, so class decisions do not change. + 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 @@ -114,16 +123,9 @@ def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): predictions_tensor < threshold_mask.unsqueeze(0).unsqueeze(0) ) & valid_predictions - if self.use_confidence: - threshold = threshold_mask.unsqueeze(0).unsqueeze(0) - scores = predictions_tensor.nan_to_num() - confidence = torch.where( - scores < threshold, - (threshold - scores) / threshold, - (scores - threshold) / (1 - threshold), - ) - else: - confidence = torch.ones_like(predictions_tensor) + confidence = self.calculate_confidence( + predictions_tensor, threshold_mask.unsqueeze(0).unsqueeze(0) + ) trust = self.calculate_trust(test_predictions) # Calculate weighted predictions using broadcasting @@ -141,8 +143,13 @@ def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): ) # Shape: (num_molecules, num_classes) # Determine which classes to include for each molecule - net_score = (positive_sum - negative_sum) / (positive_sum + negative_sum).clamp( - min=1e-6 + 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, @@ -153,3 +160,28 @@ def predict(self, test_predictions: dict[str, torch.Tensor], molecules=None): "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 ed48c03..edd7b00 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -3,23 +3,23 @@ import pandas as pd import torch -from chebifier.ensemble.voting_ensemble import VotingEnsemble +from chebifier.ensemble.voting_ensemble import WMVwithConfidenceEnsemble N_FOLDS = 5 WEIGHTING_STRENGTH_GRID = [0, 0.25, 0.5, 0.75, 1] -class WMVwithF1Ensemble(VotingEnsemble): +class WMVwithF1Ensemble(WMVwithConfidenceEnsemble): def __init__( self, ensemble_dir: str, - use_confidence: bool = True, weighting_strength=None, weighting_exponent=None, **kwargs, ): - """WMV ensemble that weights models based on their class-wise F1 scores. For each class, the weight is calculated as: + """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. @@ -27,7 +27,7 @@ def __init__( 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__(ensemble_dir, use_confidence, **kwargs) + super().__init__(ensemble_dir, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent self.model_f1_scores = None diff --git a/chebifier/hex_graph.py b/chebifier/hex_graph.py index 188abe8..1dbebcf 100644 --- a/chebifier/hex_graph.py +++ b/chebifier/hex_graph.py @@ -4,9 +4,10 @@ from scipy.sparse.csgraph import connected_components from chebifier.inconsistency_resolution import ( + NEUTRAL, ScoreBasedPredictionSmoother, densified_exclusion_matrix, - from_prob, + to_logit, ) @@ -17,12 +18,10 @@ def __init__( label_names=None, disjoint_files=None, verbose=False, - k=1.0, delta=0.0, max_states=2**20, max_component_size=40, ): - self.k = k self.delta = delta self.max_states = max_states self.max_component_size = max_component_size @@ -134,11 +133,13 @@ def _violating(self, pos): return viol def _resolve_row(self, f, scores, valid): - pos = scores > 0 + pos = scores > NEUTRAL known = torch.ones_like(pos) if valid is None else valid if valid is not None: pos = pos & valid - active = ((scores.abs() < self.delta) | self._violating(pos)) & known + active = ( + ((scores - NEUTRAL).abs() < self.delta) | self._violating(pos) + ) & known idx = torch.nonzero(active).flatten() if idx.numel() == 0: return scores @@ -165,8 +166,8 @@ def _resolve_row(self, f, scores, valid): break value[idx[forced_one]] = True value[idx[forced_zero]] = False - out[idx[forced_one]] = from_prob(torch.ones(1), self.k).item() - out[idx[forced_zero]] = from_prob(torch.zeros(1), self.k).item() + out[idx[forced_one]] = 1.0 + out[idx[forced_zero]] = 0.0 free = free & ~newly active = active.clone() active[idx[newly]] = False @@ -191,14 +192,16 @@ def _resolve_row(self, f, scores, valid): 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)] = from_prob(weights @ st, self.k) + 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() - f = self.k * preds + # 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) diff --git a/chebifier/ilr.py b/chebifier/ilr.py index ecc46e3..84b951e 100644 --- a/chebifier/ilr.py +++ b/chebifier/ilr.py @@ -3,8 +3,6 @@ from chebifier.inconsistency_resolution import ( PredictionSmoother, densified_exclusion_pairs, - from_prob, - to_prob, ) @@ -15,12 +13,10 @@ def __init__( label_names=None, disjoint_files=None, verbose=False, - k=1.0, alpha=1.0, max_iter=10, tol=1e-4, ): - self.k = k self.alpha = alpha self.max_iter = max_iter self.tol = tol @@ -79,7 +75,7 @@ def __call__(self, preds, valid_mask=None): if preds.shape[1] == 0: return preds self._valid = valid_mask - p = to_prob(preds, self.k) + p = preds.clamp(0.0, 1.0) original = p self.last_iterations = 0 for _ in range(self.max_iter): @@ -93,7 +89,7 @@ def __call__(self, preds, valid_mask=None): p = p_new self.max_iterations = max(self.max_iterations, self.last_iterations) self._valid = None - return from_prob(p, self.k) + return p class GodelILRSmoother(ILRSmoother): diff --git a/chebifier/inconsistency_resolution.py b/chebifier/inconsistency_resolution.py index fc33c73..b9ded0e 100644 --- a/chebifier/inconsistency_resolution.py +++ b/chebifier/inconsistency_resolution.py @@ -53,13 +53,18 @@ def get_disjoint_groups(disjoint_files): return disjoint_all -def to_prob(scores, k): - return torch.sigmoid(k * scores) +NEUTRAL = 0.5 -def from_prob(p, k): +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)) / k + return torch.log(p) - torch.log1p(-p) def densified_exclusion_matrix(label_names, label_successors, disjoint_groups): @@ -106,7 +111,11 @@ def get_smoother_class(name): 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 @@ -157,7 +166,7 @@ def resolve_disjointness_violations(self, preds): True ) preds[:, disj_group] = torch.where( - keep, group_preds, group_preds.clamp(max=0.0) + 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}") @@ -199,8 +208,9 @@ 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): @@ -213,6 +223,8 @@ def resolve_subsumption_violations(self, preds): 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 the one with the higher absolute value - preds_direction = preds_optimistic.abs() > preds_pessimistic.abs() + # 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 787df47..8fbcb2d 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,6 +1,9 @@ from chebifier.ensemble.dynamic_selection_ensemble import DynamicSelectionEnsemble from chebifier.ensemble.learning_to_rank_ensemble import LearningToRankEnsemble -from chebifier.ensemble.voting_ensemble import VotingEnsemble +from chebifier.ensemble.voting_ensemble import ( + MajorityVotingEnsemble, + WMVwithConfidenceEnsemble, +) from chebifier.ensemble.weighted_majority_ensemble import WMVwithF1Ensemble from chebifier.prediction_models import ( ChEBILookupPredictor, @@ -18,7 +21,8 @@ ) ENSEMBLES = { - "mv": VotingEnsemble, + "mv": MajorityVotingEnsemble, + "wmv-conf": WMVwithConfidenceEnsemble, "wmv-f1": WMVwithF1Ensemble, "ltr": LearningToRankEnsemble, "des": DynamicSelectionEnsemble, diff --git a/chebifier/predict.py b/chebifier/predict.py index 41ead5c..49f4982 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -178,12 +178,15 @@ def resolve_and_decide( predicted_classes: list[str], inconsistency_resolution: Optional[str] = "score-based", inconsistency_resolution_params: Optional[dict] = None, - decision_threshold: float = 0, + 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. From c43cb83c72180bb86fb20c138dd36b8068abcfe9 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 11 Aug 2026 13:23:05 +0200 Subject: [PATCH 15/15] update symbolic classifiers --- README.md | 5 ++- chebifier/cli.py | 4 +- chebifier/prediction_models/c3p_predictor.py | 35 +++++++++++++-- chebifier/prediction_models/chebi_lookup.py | 4 +- .../prediction_models/chemlog_predictor.py | 38 ++++++++-------- chebifier/utils.py | 43 ++++++++++++++++++- 6 files changed, 101 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 525912b..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: diff --git a/chebifier/cli.py b/chebifier/cli.py index 1f558e7..e4a8288 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -234,7 +234,9 @@ def build( ensemble_param, ): """Build (calibrate) an ensemble on the ChEBI validation set.""" - base_learners = build_base_learners(ensemble_config) + 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) ) diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index ec59748..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): @@ -28,6 +55,9 @@ def __init__( 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" @@ -55,9 +85,7 @@ def predict_list(self, smiles_list: list[str]) -> list: for result in tqdm.tqdm(result_list, desc="Reformatting C3P results"): chebi_id = result.class_id.split(":")[1] if result.is_match and self.chebi_graph is not None: - parents = [ - str(parent) for parent in self.chebi_graph.predecessors(chebi_id) - ] + parents = get_superclasses(self.chebi_graph, chebi_id) else: parents = [] for idx in indices_by_smiles[result.input_smiles]: @@ -74,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 006af48..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( diff --git a/chebifier/prediction_models/chemlog_predictor.py b/chebifier/prediction_models/chemlog_predictor.py index f637e98..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", @@ -70,41 +71,41 @@ 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): @@ -142,9 +143,9 @@ def __init__(self, model_name: str, **kwargs): print(f"Initialised ChemLog model {self.model_name}") def predict(self, smiles: str) -> Optional[dict]: - from chemlog.cli import _smiles_to_mol, strategy_call + 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(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 { @@ -181,7 +182,7 @@ def _predict_smiles_list(self, smiles_list: list[str]) -> list: 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/utils.py b/chebifier/utils.py index 91b0450..8c62db0 100644 --- a/chebifier/utils.py +++ b/chebifier/utils.py @@ -3,11 +3,15 @@ import os import pickle +import networkx as nx 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""" @@ -17,7 +21,7 @@ def load_chebi_graph(filename=None): { "repo_id": "chebai/chebifier", "repo_type": "dataset", - "files": {"f": "chebi_graph_v252.pkl"}, + "files": {"f": f"chebi_graph_v{CHEBI_VERSION}.pkl"}, } )["f"] else: @@ -89,3 +93,40 @@ def _smiles_to_mol(smiles: str): except Chem.KekulizeException as e: print(f"Failed to Kekulize {smiles}: {e}") return mol + + +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))