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..64c6b0c --- /dev/null +++ b/chebifier/build_ensemble.py @@ -0,0 +1,84 @@ +import os + +import torch + +from chebifier.predict import ( + collect_base_learner_predictions, + load_dense_predictions, + save_dense_predictions, +) + + +class EnsembleBuilder: + """ + A class to build an ensemble model from base learners and validation data. + + Attributes: + base_learners (dict[str, BasePredictor]): A dictionary of base learner models. + ensemble_model (BaseEnsemble): An instance of a BaseEnsemble model. + validation_data (list[Chem.Mol]): Validation data for calibration. + validation_labels (pd.DataFrame): Validation labels for calibration, one column per class. + The column names define the label set the base learner predictions are mapped onto. + prediction_cache_dir (str): Directory to cache predictions. + """ + + def __init__( + self, + base_learners, + ensemble_model, + validation_data, + validation_labels, + prediction_cache_dir, + ): + self.base_learners = base_learners + self.ensemble_model = ensemble_model + self.validation_data = validation_data + self.validation_labels = validation_labels + self.prediction_cache_dir = prediction_cache_dir + os.makedirs(self.prediction_cache_dir, exist_ok=True) + + def build_ensemble(self): + """ + Build an ensemble model from base learners and validation data. + + Base learner predictions are cached to avoid recomputation. + """ + + # Step 1: Get predictions from base learners on validation data + validation_predictions = {} + classes = {} + # get cached predictions if available, otherwise compute and cache them + for model_name, model in self.base_learners.items(): + cache_path = os.path.join( + 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] = load_dense_predictions(cache_path) + else: + print(f"Computing {model_name} validation predictions...") + validation_predictions[model_name] = model.predict_dense( + self.validation_data + ) + save_dense_predictions(cache_path, *validation_predictions[model_name]) + + # Base learners may be trained on different label sets (e.g. ChEBI25 vs. ChEBI25_3_STAR), + # so their union does not match the labels we calibrate against. Map every base learner + # onto the label set of the validation data instead. + label_classes = [str(cls) for cls in self.validation_labels.columns] + validation_predictions, classes = collect_base_learner_predictions( + validation_predictions, classes=label_classes + ) + validation_labels = torch.from_numpy( + self.validation_labels.to_numpy(dtype=bool) + ) + + print( + f"Collected validation predictions from {len(validation_predictions)} base learners with {len(classes)} unique classes. Calibrating ensemble model..." + ) + # Step 2: Calibrate the ensemble model using validation predictions + self.ensemble_model.calibrate( + validation_predictions, self.validation_data, validation_labels + ) + + return self.ensemble_model diff --git a/chebifier/cli.py b/chebifier/cli.py index 2f0468c..a9f9901 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -1,6 +1,143 @@ +import importlib.resources +import os +from typing import Literal + import click +import pandas as pd +import yaml +from chebi_utils.read_molecule import smiles_or_inchi_to_mol + +from chebifier.build_ensemble import EnsembleBuilder +from chebifier.check_env import check_package_installed +from chebifier.hugging_face import download_model_files +from chebifier.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()]) + + 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.""" + 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,106 +147,134 @@ 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( - "--verbose", - "-v", - is_flag=True, - default=False, - help="Enable verbose output", + "--output", + "-o", + type=click.Path(), + default=None, + help="Output file to save the predictions (optional)", +) +@click.option( + "--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__": 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..28e4e9c --- /dev/null +++ b/chebifier/ensemble/voting_ensemble.py @@ -0,0 +1,145 @@ +from pathlib import Path + +import torch +import yaml +from torchmetrics import F1Score + +from chebifier.ensemble.base_ensemble import BaseEnsemble + + +class VotingEnsemble(BaseEnsemble): + 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 + + def find_best_threshold(self, predictions, val_labels_tensor): + best_threshold = 0.5 + best_f1 = 0.0 + for threshold in range(0, 100): + threshold_value = threshold / 100 + macro_f1_score = self.classwise_f1( + predictions > threshold_value, val_labels_tensor + ).mean() + if macro_f1_score > best_f1: + best_f1 = macro_f1_score.item() + best_threshold = threshold_value + return best_threshold + + def calibrate(self, validation_predictions, validation_data, validation_labels): + print( + f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." + ) + self.classwise_f1 = F1Score( + task="multilabel", num_labels=validation_labels.shape[1], average=None + ) + self._save_prediction_thresholds( + self._fit_prediction_thresholds(validation_predictions, validation_labels) + ) + + def _fit_prediction_thresholds(self, predictions, labels) -> dict[str, float]: + return { + model_name: self.find_best_threshold(model_predictions, labels) + for model_name, model_predictions in predictions.items() + } + + def _save_prediction_thresholds(self, thresholds: dict[str, float]): + thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" + with open(thresholds_path, "w+", encoding="utf-8") as f: + yaml.dump(thresholds, f) + print(f"Saved prediction thresholds to {thresholds_path}: {thresholds}") + + def _load_prediction_thresholds(self) -> dict[str, float]: + if self.prediction_thresholds is not None: + return self.prediction_thresholds + thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" + if thresholds_path.exists(): + with open(thresholds_path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + else: + raise FileNotFoundError( + f"Prediction thresholds file not found in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." + ) + + def calculate_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. + weights are only the self-reported confidence (=difference between prediction and threshold). If set to false, all models are weighted equally. + """ + predictions_tensor = torch.stack( + list(test_predictions.values()), dim=2 + ) # Shape: (num_molecules, num_classes, num_models) + # Get predictions for all classes + valid_predictions = ~torch.isnan(predictions_tensor) + valid_counts = valid_predictions.sum(dim=2) # Sum over models dimension + + thresholds = self._load_prediction_thresholds() + if any(model_name not in thresholds for model_name in test_predictions.keys()): + raise ValueError( + "Prediction thresholds not found for all models. Please calibrate the ensemble first. Models missing thresholds: " + + ", ".join( + model_name + for model_name in test_predictions.keys() + if model_name not in thresholds + ) + ) + threshold_mask = torch.tensor( + [thresholds[model_name] for model_name in test_predictions.keys()], + dtype=predictions_tensor.dtype, + device=predictions_tensor.device, + ) + + # Skip classes with no valid predictions + has_valid_predictions = valid_counts > 0 + + # Calculate positive and negative predictions for all classes at once + positive_mask = ( + predictions_tensor > threshold_mask.unsqueeze(0).unsqueeze(0) + ) & valid_predictions + negative_mask = ( + predictions_tensor < threshold_mask.unsqueeze(0).unsqueeze(0) + ) & valid_predictions + + if self.use_confidence: + confidence = 2 * torch.abs( + predictions_tensor.nan_to_num() + - threshold_mask.unsqueeze(0).unsqueeze(0) + ) + 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 * trust + negative_weighted = negative_mask.float() * confidence * trust + + # Sum over models dimension + positive_sum = positive_weighted.sum( + dim=2 + ) # Shape: (num_molecules, num_classes) + negative_sum = negative_weighted.sum( + dim=2 + ) # Shape: (num_molecules, num_classes) + + # Determine which classes to include for each molecule + net_score = 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, + } diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 1353325..56d08da 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,107 +1,242 @@ +from pathlib import Path + +import pandas as pd import torch -from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.voting_ensemble import VotingEnsemble + +N_FOLDS = 5 +WEIGHTING_STRENGTH_GRID = [0, 0.25, 0.5, 0.75, 1] -class WMVwithPPVNPVEnsemble(BaseEnsemble): +class WMVwithF1Ensemble(VotingEnsemble): def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=1, **kwargs + 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 positive / negative predictive values. For each class, the weight is calculated as: - weight = (weighting_strength * PPV + (1 - weighting_strength)) ** weighting_exponent - where PPV is the class-specific positive predictive value of the model on the validation set - or (if the prediction is negative): - weight = (weighting_strength * NPV + (1 - weighting_strength)) ** weighting_exponent - where NPV is the class-specific negative predictive value of the model on the validation set. + """WMV ensemble that weights models based on their class-wise F1 scores. For each class, the weight is calculated as: + weight = model_weight * (weighting_strength * F1 + (1 - weighting_strength)) ** weighting_exponent + where F1 is the class-specific F1 score ("trust") of the model on the validation set. + + weighting_strength and weighting_exponent default to the optimal values determined during + calibration (best_hyperparameters.csv in the ensemble directory), falling back to 1 if the + ensemble has not been calibrated. Values passed here take precedence over both. """ - super().__init__(config_path, **kwargs) + super().__init__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_f1_scores = None - def calculate_classwise_weights(self, predicted_classes): - """ - Given the positions of predicted classes in the predictions tensor, assign weights to each class. The - result is two tensors of shape (num_predicted_classes, num_models). The weight for each class is the model_weight - (default: 1) multiplied by the class-specific positive / negative weight (default 1). - """ - positive_weights = torch.ones(len(predicted_classes), len(self.models)) - negative_weights = torch.ones(len(predicted_classes), len(self.models)) - for j, model in enumerate(self.models): - positive_weights[:, j] *= model.model_weight - negative_weights[:, j] *= model.model_weight - if model.classwise_weights is None: - continue - for cls, weights in model.classwise_weights.items(): - if cls not in predicted_classes: - continue - ppv = ( - weights["TP"] / (weights["TP"] + weights["FP"]) - if (weights["TP"] + weights["FP"]) > 0 - else 1.0 - ) - npv = ( - weights["TN"] / (weights["TN"] + weights["FN"]) - if (weights["TN"] + weights["FN"]) > 0 - else 1.0 - ) - positive_weights[predicted_classes[cls], j] *= ( - ppv * self.weighting_strength + (1 - self.weighting_strength) - ) ** self.weighting_exponent - negative_weights[predicted_classes[cls], j] *= ( - npv * self.weighting_strength + (1 - self.weighting_strength) - ) ** self.weighting_exponent - - if self.verbose_output: + def calibrate(self, validation_predictions, validation_data, validation_labels): + super().calibrate(validation_predictions, validation_data, validation_labels) + self._save_classwise_f1(validation_predictions, validation_labels) + self._optimize_hyperparameters(validation_predictions, validation_labels) + + def _fit_classwise_f1(self, predictions, labels, thresholds): + return { + model_name: self.classwise_f1( + model_predictions > thresholds[model_name], labels + ) + for model_name, model_predictions in predictions.items() + } + + def _save_classwise_f1(self, validation_predictions, validation_labels): + thresholds = self._load_prediction_thresholds() + classwise_f1 = self._fit_classwise_f1( + validation_predictions, validation_labels, thresholds + ) + for model_name, f1 in classwise_f1.items(): + f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + with open(f1_path, "w+") as f: + f.writelines(f"{x}\n" for x in f1.tolist()) print( - "Calculated model weightings. The averages for positive / negative weights are:" + f"Saved class-wise F1 scores to {f1_path}: {len(f1.tolist())} classes (macro-f1: {f1.mean().item():.4f})." ) - for i, model in enumerate(self.models): - print( - f"{model.model_name}: {positive_weights[:, i].mean().item():.3f} / {negative_weights[:, i].mean().item():.3f}" - ) - return positive_weights, negative_weights + def _load_classwise_f1(self, model_name: str) -> torch.Tensor: + if self.model_f1_scores is not None: + return self.model_f1_scores[model_name] + classwise_f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + if classwise_f1_path.exists(): + with open(classwise_f1_path, "r", encoding="utf-8") as f: + return torch.tensor([float(x) for x in f.read().splitlines()]) + else: + raise FileNotFoundError( + f"Class-wise F1 scores file not found for model {model_name} in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." + ) + def _load_hyperparameters(self) -> tuple[float, int]: + """Hyperparameters set explicitly take precedence, otherwise the optimal values found during + calibration are used (falling back to 1 if the ensemble has not been calibrated). + """ + best = {} + best_path = Path(self.ensemble_dir) / "best_hyperparameters.csv" + if best_path.exists(): + best = pd.read_csv(best_path).iloc[0].to_dict() + strength = self.weighting_strength + if strength is None: + strength = float(best.get("weighting_strength", 1)) + exponent = self.weighting_exponent + if exponent is None: + exponent = int(best.get("weighting_exponent", 1)) + return strength, exponent -class WMVwithF1Ensemble(BaseEnsemble): + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: + # Calculate trust based on class-wise F1 scores for each model + # target shape: (num_molecules, num_classes, num_models) + weighting_strength, weighting_exponent = self._load_hyperparameters() + num_models = len(predictions) + num_molecules = list(predictions.values())[0].shape[0] + num_classes = list(predictions.values())[0].shape[1] + trust_tensor = torch.ones( + (num_molecules, num_classes, num_models), dtype=torch.float32 + ) + for model_idx, (model_name, prediction_tensor) in enumerate( + predictions.items() + ): + classwise_f1 = self._load_classwise_f1(model_name) + assert ( + classwise_f1.shape[0] == num_classes + ), f"Class-wise F1 scores for model {model_name} do not match number of classes in predictions." + # Expand classwise_f1 to match the shape of trust_tensor for broadcasting + classwise_f1 = classwise_f1.unsqueeze(0).expand(num_molecules, -1) + trust_tensor[:, :, model_idx] = ( + weighting_strength * classwise_f1 + (1 - weighting_strength) + ) ** weighting_exponent + return trust_tensor - def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=6.25, **kwargs + def _build_folds(self, validation_predictions, validation_labels): + """Split the validation set into N_FOLDS folds and calibrate prediction thresholds and + class-wise F1 scores on the training part of each fold. Returns one + (thresholds, class-wise F1 scores, held-out indices) tuple per fold.""" + permutation = torch.randperm( + validation_labels.shape[0], generator=torch.Generator().manual_seed(0) + ) + fold_indices = [permutation[i::N_FOLDS] for i in range(N_FOLDS)] + folds = [] + for fold, test_idx in enumerate(fold_indices): + print(f"Calibrating fold {fold + 1}/{N_FOLDS}...") + train_idx = torch.cat( + [idx for other, idx in enumerate(fold_indices) if other != fold] + ) + train_predictions = { + model_name: model_predictions[train_idx] + for model_name, model_predictions in validation_predictions.items() + } + train_labels = validation_labels[train_idx] + thresholds = self._fit_prediction_thresholds( + train_predictions, train_labels + ) + classwise_f1 = self._fit_classwise_f1( + train_predictions, train_labels, thresholds + ) + folds.append((thresholds, classwise_f1, test_idx)) + return folds + + def _score_hyperparameters( + self, + folds, + validation_predictions, + validation_labels, + weighting_strength, + weighting_exponent, ): - """WMV ensemble that weights models based on their class-wise F1 scores. For each class, the weight is calculated as: - weight = model_weight * (weighting_strength * F1 + (1 - weighting_strength)) ** weighting_exponent - where F1 is the class-specific F1 score ("trust") of the model on the validation set. - """ - super().__init__(config_path, **kwargs) + """Macro F1 of the aggregated predictions on each held-out fold.""" 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 calculate_classwise_weights(self, predicted_classes): - """ - Given the positions of predicted classes in the predictions tensor, assign weights to each class. The - result is two tensors of shape (num_predicted_classes, num_models). The weight for each class is the model_weight - (default: 1) multiplied by (1 + the class-specific validation-f1 (default 1)). - """ - weights_by_cls = torch.ones(len(predicted_classes), len(self.models)) - for j, model in enumerate(self.models): - weights_by_cls[:, j] *= model.model_weight - if model.classwise_weights is None: - continue - for cls, weights in model.classwise_weights.items(): - if cls in predicted_classes: - if (2 * weights["TP"] + weights["FP"] + weights["FN"]) > 0: - f1 = ( - 2 - * weights["TP"] - / (2 * weights["TP"] + weights["FP"] + weights["FN"]) - ) - weights_by_cls[predicted_classes[cls], j] *= ( - self.weighting_strength * f1 + 1 - self.weighting_strength - ) ** self.weighting_exponent - if self.verbose_output: - print("Calculated model weightings. The average weights are:") - for i, model in enumerate(self.models): - print(f"{model.model_name}: {weights_by_cls[:, i].mean().item():.3f}") - - return weights_by_cls, weights_by_cls + def _optimize_hyperparameters(self, validation_predictions, validation_labels): + print( + f"Optimizing hyperparameters with {N_FOLDS}-fold cross-validation on the validation set..." + ) + weighting_strength = self.weighting_strength + weighting_exponent = self.weighting_exponent + folds = self._build_folds(validation_predictions, validation_labels) + results = [] + + def score(stage, strength, exponent): + scores = self._score_hyperparameters( + folds, validation_predictions, validation_labels, strength, exponent + ) + mean_score = sum(scores) / len(scores) + results.append( + { + "stage": stage, + "weighting_strength": strength, + "weighting_exponent": exponent, + "mean_macro_f1": mean_score, + "std_macro_f1": torch.tensor(scores).std().item(), + **{f"fold_{i}_macro_f1": s for i, s in enumerate(scores)}, + } + ) + print( + f"weighting_strength={strength}, weighting_exponent={exponent}: macro-f1 {mean_score:.4f}" + ) + return mean_score + + strength_scores = { + strength: score("weighting_strength", strength, 1) + for strength in WEIGHTING_STRENGTH_GRID + } + best_strength = max(strength_scores, key=strength_scores.get) + + best_exponent = 1 + best_score = strength_scores[best_strength] + exponent = 2 + while True: + exponent_score = score("weighting_exponent", best_strength, exponent) + if exponent_score <= best_score: + break + best_score = exponent_score + best_exponent = exponent + exponent += 1 + + self.weighting_strength = weighting_strength + self.weighting_exponent = weighting_exponent + self._save_hyperparameter_results( + results, best_strength, best_exponent, best_score + ) + + def _save_hyperparameter_results( + self, results, best_strength, best_exponent, best_score + ): + results_path = Path(self.ensemble_dir) / "hyperparameter_search.csv" + pd.DataFrame(results).to_csv(results_path, index=False) + best_path = Path(self.ensemble_dir) / "best_hyperparameters.csv" + pd.DataFrame( + [ + { + "weighting_strength": best_strength, + "weighting_exponent": best_exponent, + "mean_macro_f1": best_score, + } + ] + ).to_csv(best_path, index=False) + print( + f"Saved hyperparameter search results to {results_path}. Recommended parameters (saved to {best_path}): " + f"weighting_strength={best_strength}, weighting_exponent={best_exponent} (cross-validated macro-f1: {best_score:.4f})." + ) diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index a8287bb..8d737de 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,8 +1,5 @@ -from chebifier.ensemble.base_ensemble import BaseEnsemble -from chebifier.ensemble.weighted_majority_ensemble import ( - WMVwithF1Ensemble, - WMVwithPPVNPVEnsemble, -) +from chebifier.ensemble.voting_ensemble import VotingEnsemble +from chebifier.ensemble.weighted_majority_ensemble import WMVwithF1Ensemble from chebifier.prediction_models import ( ChEBILookupPredictor, ChemlogPeptidesPredictor, @@ -19,8 +16,7 @@ ) ENSEMBLES = { - "mv": BaseEnsemble, - "wmv-ppvnpv": WMVwithPPVNPVEnsemble, + "mv": VotingEnsemble, "wmv-f1": WMVwithF1Ensemble, } diff --git a/chebifier/predict.py b/chebifier/predict.py new file mode 100644 index 0000000..18a6fa8 --- /dev/null +++ b/chebifier/predict.py @@ -0,0 +1,173 @@ +# Get end-to-end predictions (from SMILES / molecule list via base learners + ensemble + inconsistency resolution to ChEBI classes) + + +import os +from typing import Optional + +import numpy as np +import torch +from rdkit import Chem + +from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.inconsistency_resolution import 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 save_dense_predictions(path: str, classes: list[str], scores: np.ndarray) -> None: + np.savez_compressed(path, classes=np.array(classes), scores=scores) + + +def load_dense_predictions(path: str) -> tuple[list[str], np.ndarray]: + with np.load(path) as data: + return [str(cls) for cls in data["classes"]], data["scores"] + + +def collect_base_learner_predictions( + predictions: dict[str, tuple[list[str], np.ndarray]], + classes: Optional[list[str]] = None, +) -> (dict[str, torch.Tensor], list[str]): + """ + Collect predictions from base learners into a single dictionary. + + Args: + predictions (dict): A dictionary where keys are model names and values are + (class labels, score matrix) pairs as returned by BasePredictor.predict_dense. + Assumes all score matrices have the same number of rows. + classes (Optional[list]): Column space to map the predictions onto. If None (the default), the + union of all classes reported by the base learners is used. Pass an explicit list to align + the predictions with a fixed label set (e.g. the labels of an evaluation dataset) - base + learners may be trained on different label sets, so the union does not necessarily match the + labels you want to compare against. Classes predicted by a base learner but missing from + `classes` are dropped, classes that no base learner covers stay NaN. + + Returns: + dict: A dictionary where keys are model names and values are tensors of predictions with shape (num_samples, num_classes). + If a prediction is missing, it will be NaN. + ensemble_classes (list): A list of class labels that are present in the predictions. + """ + print(f"Collecting base learner predictions from {len(predictions)} models...") + n_samples = -1 + for model_name, (_, scores) in predictions.items(): + if n_samples == -1: + n_samples = scores.shape[0] + else: + assert ( + n_samples == scores.shape[0] + ), f"All prediction matrices must have the same length. Model {model_name} has {scores.shape[0]} predictions, expected {n_samples}." + + if classes is None: + ensemble_classes = sorted( + {cls for model_classes, _ in predictions.values() for cls in model_classes} + ) + else: + ensemble_classes = list(classes) + cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} + + collected_predictions = {} + for model_name, (model_classes, scores) in predictions.items(): + if model_classes == ensemble_classes: + collected_predictions[model_name] = torch.from_numpy( + scores.astype(np.float32) + ) + continue + shared = [ + (source, cls_to_idx[cls]) + for source, cls in enumerate(model_classes) + if cls in cls_to_idx + ] + mapped = np.full((n_samples, len(ensemble_classes)), np.nan, dtype=np.float32) + if shared: + source_idx = np.fromiter( + (source for source, _ in shared), dtype=np.intp, count=len(shared) + ) + target_idx = np.fromiter( + (target for _, target in shared), dtype=np.intp, count=len(shared) + ) + mapped[:, target_idx] = scores[:, source_idx] + collected_predictions[model_name] = torch.from_numpy(mapped) + + return collected_predictions, ensemble_classes + + +def 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_dense(molecules) + else: + cache_path = os.path.join( + prediction_cache_dir, f"{model_name}_test_predictions.npz" + ) + 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]) + + 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 37851b2..6caf8c7 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -1,35 +1,66 @@ -import json from abc import ABC +import numpy as np +from rdkit import Chem + from .._custom_cache import modelwise_smiles_lru_cache +SCORE_DTYPE = np.float16 + + +def dicts_to_dense(predictions: list[dict | None]) -> tuple[list[str], np.ndarray]: + class_set = set() + previous_keys = None + for pred in predictions: + if pred: + keys = tuple(pred) + if keys != previous_keys: + class_set.update(keys) + previous_keys = keys + classes = sorted(class_set) + cls_to_idx = {cls: idx for idx, cls in enumerate(classes)} + + scores = np.full((len(predictions), len(classes)), np.nan, dtype=SCORE_DTYPE) + previous_keys = None + columns = None + for i, pred in enumerate(predictions): + if pred: + keys = tuple(pred) + if keys != previous_keys: + columns = np.fromiter( + (cls_to_idx[cls] for cls in keys), dtype=np.intp, count=len(keys) + ) + previous_keys = keys + scores[i, columns] = np.fromiter( + pred.values(), dtype=SCORE_DTYPE, count=len(pred) + ) + return classes, scores + class BasePredictor(ABC): def __init__( self, model_name: str, model_weight: int = 1, - classwise_weights_path: str = None, **kwargs, ): self.model_name = model_name self.model_weight = model_weight - if classwise_weights_path is not None: - self.classwise_weights = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.classwise_weights = None self._description = kwargs.get("description", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> dict: + def predict_list(self, molecule_list: list[str | Chem.Mol]) -> list[dict | None]: raise NotImplementedError() - def predict_smiles(self, smiles: str) -> dict: + def predict(self, molecule: str | Chem.Mol) -> dict | None: # by default, use list-based prediction - return self.predict_smiles_list([smiles])[0] + return self.predict_list([molecule])[0] + + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + return dicts_to_dense(self.predict_list(molecule_list)) @property def info_text(self): diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index bf6c39b..ec59748 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 = [] @@ -42,17 +42,29 @@ def predict_smiles_list(self, smiles_list: list[str]) -> list: ) ) + # Look up the position of each SMILES via a dict instead of scanning smiles_list + # for every result (C3P returns one result per class and molecule, so the scan + # made reformatting quadratic in the number of molecules). Repeated SMILES map to + # all of their positions, which list.index could not do (it always returned the + # first one, leaving the later rows without any predictions). + indices_by_smiles: dict[str, list[int]] = {} + for idx, smiles in enumerate(smiles_list): + indices_by_smiles.setdefault(smiles, []).append(idx) + result_reformatted = [dict() for _ in range(len(smiles_list))] for result in tqdm.tqdm(result_list, desc="Reformatting C3P results"): chebi_id = result.class_id.split(":")[1] - result_reformatted[smiles_list.index(result.input_smiles)][ - chebi_id - ] = result.is_match if result.is_match and self.chebi_graph is not None: - for parent in list(self.chebi_graph.predecessors(chebi_id)): - result_reformatted[smiles_list.index(result.input_smiles)][ - str(parent) - ] = 1 + parents = [ + 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): 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..a39eff6 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -1,14 +1,16 @@ from abc import ABC from typing import TYPE_CHECKING +import numpy as np from chebai.result.prediction import Predictor +from rdkit import Chem from chebifier import modelwise_smiles_lru_cache -from .base_predictor import BasePredictor +from .base_predictor import SCORE_DTYPE, BasePredictor if TYPE_CHECKING: - from torch import Tensor + pass class NNPredictor(BasePredictor, ABC): @@ -21,7 +23,7 @@ def __init__( super().__init__(model_name, **kwargs) self.batch_size = kwargs.get("batch_size", None) # compile_model will run the model in eager mode, which gives better performance, but does not return intermediate states - # such as attention weights. Therfore, ELECTRA attention graphs will only work with compile_model=False. + # such as attention weights. Therefore, ELECTRA attention graphs will only work with compile_model=False. compile_model = kwargs.get("compile_model", True) # If batch_size is not provided, it will be set to default batch size used during training in Predictor self.predictor: Predictor = Predictor( @@ -29,27 +31,42 @@ def __init__( ) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: """ Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary of classes and predicted values. """ - raw_preds: Tensor = self.predictor.predict_smiles(smiles_list) - if raw_preds is not None: - preds = [ - ( - { - label: pred - for label, pred in zip( - self.predictor._classification_labels, raw_preds[i].tolist() - ) - } - ) - for i in range(len(smiles_list)) - ] - return preds - else: + raw_preds = self.predictor.predict_molecules(smiles_list) + if raw_preds is None: return [None for _ in smiles_list] + return [ + ( + None + if pred_tensor is None + else { + label: pred + for label, pred in zip( + self.predictor._classification_labels, pred_tensor.tolist() + ) + } + ) + for pred_tensor in raw_preds + ] + + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + raw_preds = self.predictor.predict_molecules(molecule_list) + classes = [str(label) for label in self.predictor._classification_labels] + # molecules the model could not process stay NaN, so the ensemble skips this + # model for those rows only (see dicts_to_dense in base_predictor.py) + scores = np.full((len(molecule_list), len(classes)), np.nan, dtype=SCORE_DTYPE) + if raw_preds is None: + return classes, scores + for idx, pred in enumerate(raw_preds): + if pred is not None: + scores[idx] = pred.detach().cpu().numpy() + return classes, scores def calculate_results(self, batch): collator = self.predictor._dm.reader.COLLATOR() diff --git a/chebifier/utils.py b/chebifier/utils.py index 96f4a80..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"), - ) 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",