feat: add simple conformer sdf file caching.

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) |

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])
 

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)
-

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,25 +328,28 @@ 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
     feed to `featurize_from_mol` with as many different configs as you like.
 
     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,

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
@@ -153,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