From b75f6632b5c37dd57e29f9ee55763c1f7b70df5a Mon Sep 17 00:00:00 2001 From: Kalen Josifovski Date: Wed, 8 Jul 2026 14:37:58 +0100 Subject: [PATCH 1/2] feat: add simple conformer sdf file caching. --- README.md | 3 + examples/basic_usage.py | 3 +- src/casper/config.py | 3 +- src/casper/featurize.py | 123 +++++++++++++++++++++++++++++++++++++--- tests/test_casper.py | 46 +++++++++++++++ 5 files changed, 167 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 75002ca..80216e9 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,8 @@ cfg = casper.CasperConfig( encoding=("hist", "autocorr"), n_bins=12, autocorr_bins=8, autocorr_max_dist=16.0, conf_pool=("mean", "max"), + cache_sdf=True, + cache_dir="conformer_cache", ) v = casper.featurize("CCO", cfg) @@ -97,6 +99,7 @@ Requires the viz extra: `pip install "casper-descriptor[viz]"`. |---|---|---| | `n_confs` | 10 | ETKDG conformers per molecule | | `optimize` | `"none"` | `"none"` (raw ETKDG, fast) / `"mmff"` / `"uff"` | +| `cache_sdf`, `cache_dir` | `False`, `None` | opt-in read-through SDF cache for embedded conformers | | `properties` | `("gasteiger","logp","mr")` | which atomic properties colour the surface | | `probe` | `0.0` | `0.0` = VdW surface; `1.4` = water-accessible | | `density` | `16` | surface dots per atom (knee of the accuracy/cost curve; cost is ~quadratic via autocorr) | diff --git a/examples/basic_usage.py b/examples/basic_usage.py index 17b9af9..da496b3 100644 --- a/examples/basic_usage.py +++ b/examples/basic_usage.py @@ -16,7 +16,8 @@ # 2. tuned config cfg = CasperConfig(n_confs=10, properties=("gasteiger", "logp", "mr", "tpsa"), - encoding=("hist", "autocorr"), density=16, conf_pool=("mean", "max")) + encoding=("hist", "autocorr"), density=16, conf_pool=("mean", "max"), + cache_sdf=True, cache_dir="conformer_cache") v, names = featurize(smiles[0], cfg, return_names=True) print("tuned feature dim:", len(v), "| example name:", names[0]) diff --git a/src/casper/config.py b/src/casper/config.py index 1756843..c77bbbc 100644 --- a/src/casper/config.py +++ b/src/casper/config.py @@ -15,6 +15,8 @@ class CasperConfig: prune_rms: float = 0.5 optimize: str = "none" # 'none' (default) | 'mmff' | 'uff' embed_threads: int = 1 # RDKit-internal threads per molecule (0=all) + cache_sdf: bool = False # enable read-through conformer SDF cache + cache_dir: str | None = None # explicit cache directory when cache_sdf=True # --- surface --- properties: tuple[str, ...] = ("gasteiger", "logp", "mr") probe: float = 0.0 # 0 = VDW surface, 1.4 = water SAS @@ -76,4 +78,3 @@ def pool(vectors: Sequence[np.ndarray], energies: np.ndarray | None, else: raise ValueError(f"unknown pooling '{m}'") return np.concatenate(parts) - diff --git a/src/casper/featurize.py b/src/casper/featurize.py index f78c6f3..779e9a3 100644 --- a/src/casper/featurize.py +++ b/src/casper/featurize.py @@ -1,4 +1,5 @@ """casper.featurize -- bag/pooled featurisers, batch helpers, sklearn transformer.""" +import hashlib import os from dataclasses import replace from typing import Sequence @@ -47,6 +48,107 @@ def _per_conf_dim(cfg, dist_edges): return d +_CACHE_CONF_PROP = "_casper_conf_id" +_CACHE_ENERGY_PROP = "_casper_energy" + + +def _cache_path(mol: Chem.Mol, cfg: CasperConfig) -> str | None: + if not (cfg.cache_sdf and cfg.cache_dir): + return None + mol_no_h = Chem.RemoveHs(Chem.Mol(mol)) + key = "|".join([ + Chem.MolToSmiles(mol_no_h, canonical=True), + str(cfg.n_confs), + str(cfg.seed), + str(cfg.prune_rms), + cfg.optimize, + ]) + digest = hashlib.sha256(key.encode("utf-8")).hexdigest() + return os.path.join(cfg.cache_dir, f"{digest}.sdf") + + +def _load_cached_embed(path: str, optimize: str) -> tuple[Chem.Mol, list[int], np.ndarray | None] | None: + try: + supplier = Chem.SDMolSupplier(path, removeHs=False) + except Exception: + return None + entries = [mol for mol in supplier if mol is not None] + if not entries: + return None + + ordered = [] + for idx, entry in enumerate(entries): + if not entry.HasProp(_CACHE_CONF_PROP): + return None + try: + conf_id = int(entry.GetProp(_CACHE_CONF_PROP)) + except ValueError: + return None + ordered.append((conf_id, idx, entry)) + ordered.sort() + + base = Chem.Mol(ordered[0][2]) + base.RemoveAllConformers() + energies = [] + saw_energy = False + for conf_id, _idx, entry in ordered: + conf = Chem.Conformer(entry.GetConformer()) + conf.SetId(conf_id) + base.AddConformer(conf, assignId=False) + if entry.HasProp(_CACHE_ENERGY_PROP): + try: + energies.append(float(entry.GetProp(_CACHE_ENERGY_PROP))) + saw_energy = True + except ValueError: + return None + else: + energies.append(np.nan) + + cids = [c.GetId() for c in base.GetConformers()] + if optimize == "none" or not saw_energy or len(energies) != len(cids): + out_energies = None + elif np.all(np.isfinite(energies)): + out_energies = np.asarray(energies, dtype=float) + else: + return None + return base, cids, out_energies + + +def _write_cached_embed(path: str, mol3d: Chem.Mol, cids: list[int], energies: np.ndarray | None) -> None: + os.makedirs(os.path.dirname(path), exist_ok=True) + writer = Chem.SDWriter(path) + try: + for i, cid in enumerate(cids): + cached = Chem.Mol(mol3d) + cached.RemoveAllConformers() + conf = Chem.Conformer(mol3d.GetConformer(cid)) + conf.SetId(cid) + cached.AddConformer(conf, assignId=False) + cached.SetProp(_CACHE_CONF_PROP, str(cid)) + if energies is not None and i < len(energies) and np.isfinite(energies[i]): + cached.SetProp(_CACHE_ENERGY_PROP, repr(float(energies[i]))) + writer.write(cached) + finally: + writer.close() + + +def _embed_with_cache(mol: Chem.Mol, cfg: CasperConfig) -> tuple[Chem.Mol, list[int], np.ndarray | None]: + path = _cache_path(mol, cfg) + if path and os.path.exists(path): + cached = _load_cached_embed(path, cfg.optimize) + if cached is not None: + return cached + + mol3d, cids, energies = embed(mol, cfg.n_confs, cfg.seed, cfg.prune_rms, + cfg.optimize, cfg.embed_threads) + if path and cids: + try: + _write_cached_embed(path, mol3d, cids, energies) + except Exception: + pass + return mol3d, cids, energies + + def bag_from_mol(mol3d: Chem.Mol, config: CasperConfig | None = None): """ Core primitive: encode each conformer of an already-embedded molecule WITHOUT @@ -121,8 +223,7 @@ def featurize(mol, config: CasperConfig | None = None, return_names: bool = Fals if mol is None: raise ValueError("could not parse molecule") - mol3d, _cids, energies = embed(mol, cfg.n_confs, cfg.seed, cfg.prune_rms, - cfg.optimize, cfg.embed_threads) + mol3d, _cids, energies = _embed_with_cache(mol, cfg) return featurize_from_mol(mol3d, energies, cfg, return_names) @@ -144,8 +245,7 @@ def featurize_bag(mol, config: CasperConfig | None = None): if mol is None: raise ValueError("could not parse molecule") - mol3d, _cids, _energies = embed(mol, cfg.n_confs, cfg.seed, cfg.prune_rms, - cfg.optimize, cfg.embed_threads) + mol3d, _cids, _energies = _embed_with_cache(mol, cfg) return bag_from_mol(mol3d, cfg) @@ -228,17 +328,18 @@ def _first_parsable(mols: Sequence): # Conformer cache: embed ONCE, reuse for many descriptor variants # # --------------------------------------------------------------------------- # def _embed_one(args): - smi_or_mol, n_confs, seed, prune_rms, optimize = args + smi_or_mol, cfg = args mol = Chem.MolFromSmiles(smi_or_mol) if isinstance(smi_or_mol, str) else smi_or_mol if mol is None: return (None, None) - mol3d, cids, energies = embed(mol, n_confs, seed, prune_rms, optimize, threads=1) + mol3d, cids, energies = _embed_with_cache(mol, replace(cfg, embed_threads=1)) return (mol3d.ToBinary() if cids else None, None if energies is None else np.asarray(energies)) def embed_many(mols: Sequence, n_confs: int = 10, seed: int = 0xF00D, - prune_rms: float = 0.5, optimize: str = "none", n_jobs: int = 1): + prune_rms: float = 0.5, optimize: str = "none", n_jobs: int = 1, + cache_sdf: bool = False, cache_dir: str | None = None): """ Embed many molecules ONCE (the expensive step). Returns a list of (mol_binary_or_None, energies_or_None) that you can pickle to disk and then @@ -246,7 +347,9 @@ def embed_many(mols: Sequence, n_confs: int = 10, seed: int = 0xF00D, Reconstruct a cached molecule with `Chem.Mol(binary)`. """ - args = [(m, n_confs, seed, prune_rms, optimize) for m in mols] + cfg = CasperConfig(n_confs=n_confs, seed=seed, prune_rms=prune_rms, + optimize=optimize, cache_sdf=cache_sdf, cache_dir=cache_dir) + args = [(m, cfg) for m in mols] if n_jobs == 1: return [_embed_one(a) for a in args] import os @@ -271,13 +374,14 @@ class CasperFeaturizer(BaseEstimator, TransformerMixin): density, n_bins, probe, properties, conf_pool, encoding. """ def __init__(self, n_confs=10, seed=0xF00D, prune_rms=0.5, optimize="none", - embed_threads=1, + embed_threads=1, cache_sdf=False, cache_dir=None, properties=("gasteiger", "logp", "mr"), probe=0.0, density=32, encoding=("hist",), n_bins=12, bin_range=None, autocorr_max_dist=12.0, autocorr_bins=8, autocorr_range=None, conf_pool=("mean",), boltzmann_T=300.0, n_jobs=1): self.n_confs = n_confs; self.seed = seed; self.prune_rms = prune_rms self.optimize = optimize; self.embed_threads = embed_threads + self.cache_sdf = cache_sdf; self.cache_dir = cache_dir self.properties = properties; self.probe = probe self.density = density; self.encoding = encoding; self.n_bins = n_bins self.bin_range = bin_range; self.autocorr_max_dist = autocorr_max_dist @@ -289,6 +393,7 @@ def _config(self) -> CasperConfig: return CasperConfig( n_confs=self.n_confs, seed=self.seed, prune_rms=self.prune_rms, optimize=self.optimize, embed_threads=self.embed_threads, + cache_sdf=self.cache_sdf, cache_dir=self.cache_dir, properties=tuple(self.properties), probe=self.probe, density=self.density, encoding=tuple(self.encoding), n_bins=self.n_bins, bin_range=self.bin_range, diff --git a/tests/test_casper.py b/tests/test_casper.py index 06fd84b..2f295e9 100644 --- a/tests/test_casper.py +++ b/tests/test_casper.py @@ -3,6 +3,7 @@ fixed-length output, the pool/bag relationship, density-invariance of the normalized autocorrelation, determinism, and graceful failure handling. """ +import importlib import numpy as np import pytest @@ -100,6 +101,51 @@ def test_featurize_many_failure_is_nan_row(): assert np.isnan(X[1]).all() +def test_sdf_cache_round_trip_and_hit(tmp_path, monkeypatch): + cfg = CasperConfig(n_confs=4, seed=7, cache_sdf=True, cache_dir=str(tmp_path)) + first = featurize("CCO", cfg) + files = list(tmp_path.glob("*.sdf")) + assert len(files) == 1 + + featurize_mod = importlib.import_module("casper.featurize") + monkeypatch.setattr(featurize_mod, "embed", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("embed called"))) + second = featurize("CCO", cfg) + assert np.allclose(first, second) + + +def test_sdf_cache_key_changes_with_embedding_inputs(tmp_path): + cfg_a = CasperConfig(n_confs=3, seed=7, cache_sdf=True, cache_dir=str(tmp_path)) + cfg_b = CasperConfig(n_confs=3, seed=8, cache_sdf=True, cache_dir=str(tmp_path)) + featurize("CCO", cfg_a) + featurize("CCO", cfg_b) + assert len(list(tmp_path.glob("*.sdf"))) == 2 + + +def test_sdf_cache_preserves_boltzmann_pooling(tmp_path, monkeypatch): + cached = CasperConfig(n_confs=4, seed=11, optimize="mmff", conf_pool=("boltzmann",), + cache_sdf=True, cache_dir=str(tmp_path)) + uncached = CasperConfig(n_confs=4, seed=11, optimize="mmff", conf_pool=("boltzmann",)) + expected = featurize("CCO", uncached) + first = featurize("CCO", cached) + assert np.allclose(first, expected) + + featurize_mod = importlib.import_module("casper.featurize") + monkeypatch.setattr(featurize_mod, "embed", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("embed called"))) + second = featurize("CCO", cached) + assert np.allclose(second, expected) + + +def test_corrupt_sdf_cache_falls_back_to_fresh_embed(tmp_path): + cfg = CasperConfig(n_confs=3, seed=5, cache_sdf=True, cache_dir=str(tmp_path)) + featurize("CCO", cfg) + [path] = list(tmp_path.glob("*.sdf")) + path.write_text("not an sdf") + v = featurize("CCO", cfg) + assert len(v) > 0 and np.isfinite(v).all() + + def test_register_property(): casper.register_property("const_one", lambda mol: np.ones(mol.GetNumAtoms()), (0.0, 1.0)) assert "const_one" in PROPERTIES From f8ee1dd45f771cdc1e3997e3b1f295814a887452 Mon Sep 17 00:00:00 2001 From: Kalen Josifovski Date: Thu, 9 Jul 2026 11:40:32 +0100 Subject: [PATCH 2/2] fix: import jazzy directly in the skip guard instead of import casper.jazzy_properties. --- tests/test_casper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_casper.py b/tests/test_casper.py index 2f295e9..845e094 100644 --- a/tests/test_casper.py +++ b/tests/test_casper.py @@ -199,7 +199,7 @@ def test_parse_feature_name(): # --- optional jazzy properties (skipped if jazzy not installed) ---------------- def test_jazzy_properties_if_available(): try: - import casper.jazzy_properties # noqa: registers eeq, alp, sa, sdc, sdx + import jazzy # noqa: required by casper.jazzy_properties except ImportError: pytest.skip("jazzy not installed") from casper import PROPERTIES, featurize