| """Public molecular property data, scaffold splits, and CPU surrogate fits.""" |
|
|
| from pathlib import Path |
| import hashlib, json, time |
| import requests, numpy as np, pandas as pd, joblib |
| from rdkit import Chem, DataStructs |
| from rdkit.Chem import rdFingerprintGenerator |
| from rdkit.Chem.Scaffolds import MurckoScaffold |
| from sklearn.ensemble import ExtraTreesRegressor |
| from sklearn.metrics import mean_absolute_error, r2_score |
| from .chemistry import canonical |
|
|
| DATASETS = { |
| "caco2": { |
| "url": "https://dataverse.harvard.edu/api/access/datafile/4259569", |
| "sep": "\t", |
| "smiles": "Drug", |
| "target": "Y", |
| "units": "log10(cm/s)", |
| }, |
| "bace": { |
| "url": "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/bace.csv", |
| "sep": ",", |
| "smiles": "mol", |
| "target": "pIC50", |
| "units": "pIC50", |
| }, |
| } |
|
|
|
|
| def download(name, directory): |
| """Fetch a released dataset once and save its source URL and SHA-256 hash.""" |
| spec = DATASETS[name] |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| path = directory / (name + ".csv") |
| if not path.exists(): |
| r = requests.get(spec["url"], timeout=120) |
| r.raise_for_status() |
| path.write_bytes(r.content) |
| manifest = { |
| "dataset": name, |
| "source": spec["url"], |
| "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), |
| "bytes": path.stat().st_size, |
| "units": spec["units"], |
| } |
| (directory / (name + "_source.json")).write_text(json.dumps(manifest, indent=2)) |
| return path |
|
|
|
|
| def fingerprints(smiles): |
| """Return a float32 array of shape (molecules, 1024) with Morgan bits.""" |
| fp = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=1024) |
| out = np.zeros((len(smiles), 1024), dtype=np.float32) |
| for i, s in enumerate(smiles): |
| DataStructs.ConvertToNumpyArray( |
| fp.GetFingerprint(Chem.MolFromSmiles(s)), out[i] |
| ) |
| return out |
|
|
|
|
| def load_data(name, path): |
| """Canonicalize input molecules, average duplicate labels, and record exclusions.""" |
| spec = DATASETS[name] |
| df = pd.read_csv(path, sep=spec["sep"]) |
| rows = [] |
| excluded = 0 |
| for _, r in df.iterrows(): |
| try: |
| s = canonical(r[spec["smiles"]]) |
| y = float(r[spec["target"]]) |
| if not np.isfinite(y): |
| raise ValueError() |
| rows.append((s, y)) |
| except (ValueError, TypeError): |
| excluded += 1 |
| clean = ( |
| pd.DataFrame(rows, columns=["smiles", "value"]) |
| .groupby("smiles", as_index=False) |
| .value.mean() |
| ) |
| return clean, { |
| "raw_rows": len(df), |
| "excluded_rows": excluded, |
| "unique_molecules": len(clean), |
| } |
|
|
|
|
| def scaffold_split(smiles, seed=0, train_fraction=0.8): |
| """Partition entire Bemis-Murcko groups into train and test index arrays.""" |
| groups = {} |
| for i, s in enumerate(smiles): |
| scaffold = MurckoScaffold.MurckoScaffoldSmiles(smiles=s, includeChirality=False) |
| |
| groups.setdefault(scaffold, []).append(i) |
| rng = np.random.default_rng(seed) |
| items = list(groups.items()) |
| rng.shuffle(items) |
| items.sort(key=lambda x: -len(x[1])) |
| train = [] |
| test = [] |
| trsc = [] |
| tesc = [] |
| for key, inds in items: |
| if len(train) + len(inds) <= train_fraction * len(smiles): |
| train += inds |
| trsc.append(key) |
| else: |
| test += inds |
| tesc.append(key) |
| if not train or not test: |
| raise ValueError("Insufficient scaffold groups for a held-out split") |
| assert not set(trsc) & set(tesc) |
| return np.array(train), np.array(test) |
|
|
|
|
| def fit_property(name, data_directory, output, seed=0): |
| """Fit a 256-tree predictor and save weights, split membership, and test scores.""" |
| out = Path(output) |
| out.mkdir(parents=True, exist_ok=True) |
| path = download(name, data_directory) |
| df, counts = load_data(name, path) |
| smiles = df.smiles.tolist() |
| x = fingerprints(smiles) |
| y = df.value.to_numpy() |
| tr, te = scaffold_split(smiles, seed) |
| start = time.perf_counter() |
| model = ExtraTreesRegressor( |
| n_estimators=256, |
| min_samples_leaf=2, |
| max_features=0.5, |
| random_state=seed, |
| n_jobs=2, |
| ) |
| model.fit(x[tr], y[tr]) |
| pred = model.predict(x[te]) |
| result = { |
| **counts, |
| "dataset": name, |
| "seed": seed, |
| "split": "Bemis-Murcko scaffold, 80/20", |
| "train": len(tr), |
| "test": len(te), |
| "mae": mean_absolute_error(y[te], pred), |
| "r2": r2_score(y[te], pred), |
| "seconds": time.perf_counter() - start, |
| "units": DATASETS[name]["units"], |
| } |
| joblib.dump(model, out / (name + ".joblib")) |
| (out / (name + "_metrics.json")).write_text(json.dumps(result, indent=2)) |
| df.assign(split=np.where(np.isin(np.arange(len(df)), tr), "train", "test")).to_csv( |
| out / (name + "_split.csv"), index=False |
| ) |
| pd.DataFrame( |
| {"smiles": df.smiles.iloc[te], "observed": y[te], "predicted": pred} |
| ).to_csv(out / (name + "_test_predictions.csv"), index=False) |
| return result |
|
|
|
|
| def property_rewards(graph, model_directory, weights=(0.5, 0.5), concentration=5.0): |
| """Predicted BACE inhibition and Caco-2 transport, normalized to utilities.""" |
| if not np.isfinite(concentration) or concentration <= 0: |
| raise ValueError("Concentration must be finite and positive") |
| smiles = list(graph.terminals) |
| x = fingerprints(smiles) |
| directory = Path(model_directory) |
| bace = joblib.load(directory / "bace.joblib").predict(x) |
| caco = joblib.load(directory / "caco2.joblib").predict(x) |
| ub = np.clip((bace - 4) / 5, 0, 1) |
| uc = np.clip((caco + 7) / 3, 0, 1) |
| w = np.asarray(weights, dtype=float) |
| if w.shape != (2,) or np.any(w < 0) or not np.all(np.isfinite(w)) or w.sum() <= 0: |
| raise ValueError("Two nonnegative property weights with positive sum required") |
| w = w / w.sum() |
| rewards = { |
| s: float(concentration * (w[0] * a + w[1] * b)) |
| for s, a, b in zip(smiles, ub, uc) |
| } |
| scores = pd.DataFrame( |
| { |
| "smiles": smiles, |
| "predicted_bace_pIC50": bace, |
| "predicted_caco2_log10_cm_s": caco, |
| "bace_utility": ub, |
| "caco2_utility": uc, |
| } |
| ) |
| return rewards, scores |
|
|
