| |
| """Run Morgan/protein-descriptor LightGBM controls on every split manifest.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.metadata |
| import json |
| import math |
| import time |
| from collections import Counter |
| from pathlib import Path |
|
|
| import lightgbm as lgb |
| import numpy as np |
| from rdkit import Chem |
| from rdkit.Chem import Descriptors, Lipinski, rdFingerprintGenerator |
|
|
| from mitointeract_recovery.metrics import regression_metrics |
|
|
| AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY" |
| DIPEPTIDES = tuple(a + b for a in AMINO_ACIDS for b in AMINO_ACIDS) |
| MORGAN_BITS = 2048 |
|
|
|
|
| def read_jsonl(path: Path) -> list[dict]: |
| with path.open() as handle: |
| return [json.loads(line) for line in handle if line.strip()] |
|
|
|
|
| def ligand_features(smiles: str, generator) -> np.ndarray: |
| mol = Chem.MolFromSmiles(smiles) |
| if mol is None: |
| raise ValueError(f"invalid canonical SMILES: {smiles}") |
| fingerprint = generator.GetFingerprintAsNumPy(mol).astype(np.float32) |
| descriptors = np.asarray( |
| [ |
| Descriptors.MolWt(mol) / 1000, |
| Descriptors.MolLogP(mol) / 10, |
| Descriptors.TPSA(mol) / 200, |
| Lipinski.NumHDonors(mol) / 10, |
| Lipinski.NumHAcceptors(mol) / 20, |
| Lipinski.NumRotatableBonds(mol) / 30, |
| Lipinski.RingCount(mol) / 20, |
| Lipinski.FractionCSP3(mol), |
| ], |
| dtype=np.float32, |
| ) |
| return np.concatenate([fingerprint, descriptors]) |
|
|
|
|
| def protein_features(sequence: str) -> np.ndarray: |
| sequence = sequence.upper() |
| length = max(1, len(sequence)) |
| counts = Counter(sequence) |
| amino_acid_composition = np.asarray( |
| [counts[amino_acid] / length for amino_acid in AMINO_ACIDS], |
| dtype=np.float32, |
| ) |
| dipeptide_counts = Counter( |
| sequence[index : index + 2] for index in range(length - 1) |
| ) |
| denominator = max(1, length - 1) |
| dipeptide_composition = np.asarray( |
| [dipeptide_counts[pair] / denominator for pair in DIPEPTIDES], |
| dtype=np.float32, |
| ) |
| return np.concatenate( |
| [ |
| np.asarray([math.log1p(length) / 10], dtype=np.float32), |
| amino_acid_composition, |
| dipeptide_composition, |
| ] |
| ) |
|
|
|
|
| def build_features( |
| rows: list[dict], |
| ) -> tuple[np.ndarray, np.ndarray, list[str], list[str]]: |
| generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=MORGAN_BITS) |
| ligand_cache = { |
| row["ligand_id"]: ligand_features(row["smiles"], generator) for row in rows |
| } |
| protein_cache = { |
| row["protein_id"]: protein_features(row["sequence"]) for row in rows |
| } |
| ligand = np.stack([ligand_cache[row["ligand_id"]] for row in rows]) |
| protein = np.stack([protein_cache[row["protein_id"]] for row in rows]) |
| ligand_names = [f"morgan_{index}" for index in range(MORGAN_BITS)] + [ |
| "mol_weight", |
| "mol_logp", |
| "tpsa", |
| "h_donors", |
| "h_acceptors", |
| "rotatable_bonds", |
| "ring_count", |
| "fraction_csp3", |
| ] |
| protein_names = ( |
| ["log_protein_length"] |
| + [f"aac_{amino_acid}" for amino_acid in AMINO_ACIDS] |
| + [f"dipeptide_{pair}" for pair in DIPEPTIDES] |
| ) |
| return ligand, protein, ligand_names, protein_names |
|
|
|
|
| def read_manifest(path: Path) -> dict[str, str]: |
| return {row["pair_id"]: row["split"] for row in read_jsonl(path)} |
|
|
|
|
| def fit_model( |
| name: str, |
| features: np.ndarray, |
| feature_names: list[str], |
| targets: np.ndarray, |
| split_indices: dict[str, np.ndarray], |
| seed: int, |
| ) -> dict: |
| started = time.monotonic() |
| model = lgb.LGBMRegressor( |
| objective="regression_l2", |
| n_estimators=1000, |
| learning_rate=0.03, |
| num_leaves=31, |
| min_child_samples=20, |
| subsample=0.8, |
| colsample_bytree=0.8, |
| reg_lambda=1.0, |
| random_state=seed, |
| n_jobs=8, |
| deterministic=True, |
| force_col_wise=True, |
| verbosity=-1, |
| ) |
| train = split_indices["train"] |
| validation = split_indices["validation"] |
| test = split_indices["test"] |
| model.fit( |
| features[train], |
| targets[train], |
| eval_X=features[validation], |
| eval_y=targets[validation], |
| eval_metric="rmse", |
| callbacks=[lgb.early_stopping(50, verbose=False)], |
| ) |
| importances = sorted( |
| zip(feature_names, model.feature_importances_, strict=True), |
| key=lambda item: item[1], |
| reverse=True, |
| )[:20] |
| return { |
| "name": name, |
| "best_iteration": int(model.best_iteration_), |
| "validation": regression_metrics( |
| targets[validation], model.predict(features[validation]) |
| ), |
| "test": regression_metrics(targets[test], model.predict(features[test])), |
| "top_feature_importance": [ |
| {"feature": feature, "gain_proxy": int(importance)} |
| for feature, importance in importances |
| ], |
| "fit_and_eval_seconds": time.monotonic() - started, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--data-dir", type=Path, required=True) |
| parser.add_argument("--target-key", required=True) |
| parser.add_argument("--target-name", required=True) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--output", type=Path, required=True) |
| args = parser.parse_args() |
|
|
| rows = read_jsonl(args.data_dir / "sample.jsonl") |
| targets = np.asarray([row[args.target_key] for row in rows], dtype=np.float64) |
| ligand, protein, ligand_names, protein_names = build_features(rows) |
| combined = np.concatenate([protein, ligand], axis=1) |
| feature_sets = { |
| "ligand_morgan_descriptors_lightgbm": (ligand, ligand_names), |
| "protein_aac_dipeptide_lightgbm": (protein, protein_names), |
| "combined_morgan_protein_lightgbm": ( |
| combined, |
| protein_names + ligand_names, |
| ), |
| } |
|
|
| report = { |
| "sample_rows": len(rows), |
| "target": args.target_name, |
| "seed": args.seed, |
| "packages": { |
| package: importlib.metadata.version(package) |
| for package in ("lightgbm", "numpy", "rdkit") |
| }, |
| "feature_dimensions": { |
| "ligand": int(ligand.shape[1]), |
| "protein": int(protein.shape[1]), |
| "combined": int(combined.shape[1]), |
| }, |
| "splits": {}, |
| } |
| pair_ids = [row["pair_id"] for row in rows] |
| for manifest_path in sorted(args.data_dir.glob("split-*.jsonl")): |
| manifest = read_manifest(manifest_path) |
| split_indices = { |
| split: np.asarray( |
| [ |
| index |
| for index, pair_id in enumerate(pair_ids) |
| if manifest[pair_id] == split |
| ] |
| ) |
| for split in ("train", "validation", "test") |
| } |
| mean = float(targets[split_indices["train"]].mean()) |
| report["splits"][manifest_path.stem.removeprefix("split-")] = { |
| "rows": { |
| split: int(len(indices)) for split, indices in split_indices.items() |
| }, |
| "mean_baseline": { |
| "prediction": mean, |
| "validation": regression_metrics( |
| targets[split_indices["validation"]], |
| np.full(len(split_indices["validation"]), mean), |
| ), |
| "test": regression_metrics( |
| targets[split_indices["test"]], |
| np.full(len(split_indices["test"]), mean), |
| ), |
| }, |
| "models": [ |
| fit_model( |
| name, |
| features, |
| feature_names, |
| targets, |
| split_indices, |
| args.seed, |
| ) |
| for name, (features, feature_names) in feature_sets.items() |
| ], |
| } |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
