Coverage for cosmolayer/store/segments.py: 97%

1"""A segment-data store: on-disk arrays of COSMO segment coordinates, 

2charges, and areas, plus the per-molecule table describing them. 

3 

4Build one from ``.cosmo`` files with ``SegmentStore.from_cosmo_files``, or 

5load an existing one with ``SegmentStore.load``. Both use a flat 

6``storage_dir``:: 

7 

8 <storage_dir>/data.npy float32 (n_segs_total, 5): [x, y, z, charge, area] 

9 <storage_dir>/atom_indices.npy int64 (n_segs_total,): global atom index per 

10 segment 

11 <storage_dir>/molecules.parquet columns: smiles, filename, segment_offsets, 

12 atom_offsets, num_atoms, volume, cluster_id, 

13 cluster_distance, split (optional) 

14 <storage_dir>/atoms.parquet columns: id, element, x, y, z -- one row per atom, 

15 in global atom index order 

16 <storage_dir>/metadata.json {num_molecules, num_cosmo_parse_failures, schemes} 

17 <storage_dir>/<scheme>.npy float32 (n_segs_total,): one per averaging scheme 

18 

19``segment_offsets`` (the ``molecules_df`` column of that name) must describe 

20*exactly* the molecules present in the accompanying segment-level arrays: 

21the last molecule's segments run to the end of those arrays. 

22""" 

23 

24import json 

25import os 

26import pathlib 

27import tempfile 

28from collections.abc import Mapping, Sequence 

29from dataclasses import dataclass, field 

30from typing import Any, cast 

31 

32import numpy as np 

33import pandas as pd 

34from numpy.typing import NDArray 

35from rdkit import Chem 

36from tqdm.auto import tqdm 

37 

38from cosmolayer.parser import parse_cosmo_file 

39 

40from .averaging import ( 

41 AVERAGING_SCHEMES, 

42 AveragingScheme, 

43 average_sigmas_by_molecule, 

44) 

45from .clustering import ( 

46 ClusteringSpecs, 

47 FingerprintGenerator, 

48 butina_cluster, 

49 cluster_medoid_distances, 

50) 

51from .coarse_graining import compute_atom_remap 

52from .grid import DEFAULT_SIGMA_GRID, SigmaGrid 

53from .profiles import SigmaProfileTable 

54from .splitting import greedy_cluster_split 

55from .subsampling import apportion_counts, restrict_to_molecules 

56 

57DATA_FILE = pathlib.Path("data.npy") 

58ATOM_INDICES_FILE = pathlib.Path("atom_indices.npy") 

59MOLECULES_FILE = pathlib.Path("molecules.parquet") 

60ATOMS_FILE = pathlib.Path("atoms.parquet") 

61METADATA_FILE = pathlib.Path("metadata.json") 

62 

63_STORE_FILES = (DATA_FILE, ATOM_INDICES_FILE, MOLECULES_FILE, ATOMS_FILE, METADATA_FILE) 

64 

65# Averaged sigmas are written to "<name>.npy", so scheme names must not 

66# collide with the store's own .npy stems (data, atom_indices). 

67_RESERVED_SCHEME_NAMES = frozenset( 

68 f.stem for f in _STORE_FILES if f.suffix == DATA_FILE.suffix 

69) 

70 

71# COSMO atom tables include explicit hydrogens; disable RDKit's default 

72# folding of terminal [H] so SMILES atoms match the COSMO file. 

73_SMILES_PARSER_PARAMS = Chem.SmilesParserParams() 

74cast(Any, _SMILES_PARSER_PARAMS).removeHs = False 

75 

76 

77@dataclass 

78class StoreMetadata: 

79 """Typed view of a ``SegmentStore``'s ``metadata.json``. 

80 

81 Parameters 

82 ---------- 

83 num_molecules : int 

84 Number of molecules successfully stored. 

85 num_cosmo_parse_failures : int 

86 Number of molecules skipped because the COSMO file was missing, 

87 or because they failed to parse or validate (parse/validate 

88 skips require ``ignore_errors=True``, see 

89 ``SegmentStore.from_cosmo_files``). 

90 schemes : dict[str, AveragingScheme] 

91 Averaging schemes this store has computed sigmas for, keyed by 

92 scheme name. 

93 """ 

94 

95 num_molecules: int 

96 num_cosmo_parse_failures: int 

97 schemes: dict[str, AveragingScheme] = field(default_factory=dict) 

98 

99 def to_dict(self) -> dict[str, Any]: 

100 """Serialize to the JSON-compatible shape written to ``metadata.json``.""" 

101 return { 

102 "num_molecules": self.num_molecules, 

103 "num_cosmo_parse_failures": self.num_cosmo_parse_failures, 

104 "schemes": { 

105 name: { 

106 "averaging_radius": scheme.averaging_radius, 

107 "f_decay": scheme.f_decay, 

108 } 

109 for name, scheme in self.schemes.items() 

110 }, 

111 } 

112 

113 @classmethod 

114 def from_dict(cls, data: dict[str, Any]) -> "StoreMetadata": 

115 """Reconstruct from the dict produced by ``to_dict``. 

116 

117 Parameters 

118 ---------- 

119 data : dict 

120 As produced by ``to_dict``. 

121 """ 

122 schemes = { 

123 name: AveragingScheme(name, params["averaging_radius"], params["f_decay"]) 

124 for name, params in data.get("schemes", {}).items() 

125 } 

126 return cls( 

127 num_molecules=data["num_molecules"], 

128 num_cosmo_parse_failures=data["num_cosmo_parse_failures"], 

129 schemes=schemes, 

130 ) 

131 

132 

133class SegmentStore: 

134 """On-disk COSMO segment arrays, the per-molecule table describing 

135 them, and any averaged sigmas computed for the store. 

136 

137 Prefer ``load`` (read an existing store, memory-mapped) or 

138 ``from_cosmo_files`` (build a new one). Direct construction is for 

139 wrapping arrays already in hand. 

140 

141 Parameters 

142 ---------- 

143 storage_dir : pathlib.Path 

144 Directory this store's files live in (or will be written to). 

145 data : np.ndarray 

146 ``(n_segs_total, 5)`` array, columns ``[x, y, z, charge, area]``. 

147 atom_indices : np.ndarray 

148 ``(n_segs_total,)`` global atom index of each segment. 

149 molecules_df : pd.DataFrame 

150 One row per molecule, with columns ``smiles``, ``segment_offsets``, 

151 ``atom_offsets``, ``num_atoms``, ``volume``, ``cluster_id``, and 

152 ``cluster_distance``. 

153 atoms_df : pd.DataFrame 

154 One row per atom, columns ``id``, ``element``, ``x``, ``y``, ``z``, 

155 in global atom index order (the same index space ``atom_indices`` 

156 points into). 

157 metadata : StoreMetadata 

158 Molecule and parse-failure counts, plus averaging schemes already 

159 computed for this store. 

160 averaged_sigmas : dict[str, np.ndarray] 

161 Scheme name to ``(n_segs_total,)`` averaged charge density. Empty 

162 if none have been computed. 

163 

164 Attributes 

165 ---------- 

166 coords, charges, areas : np.ndarray 

167 Views into ``data``'s columns. 

168 """ 

169 

170 def __init__( # noqa: PLR0913, PLR0917 

171 self, 

172 storage_dir: pathlib.Path, 

173 data: NDArray[np.float32], 

174 atom_indices: NDArray[np.int64], 

175 molecules_df: pd.DataFrame, 

176 atoms_df: pd.DataFrame, 

177 metadata: StoreMetadata, 

178 averaged_sigmas: dict[str, NDArray[np.float32]], 

179 ) -> None: 

180 self.storage_dir = pathlib.Path(storage_dir) 

181 self.data = data 

182 self.atom_indices = atom_indices 

183 self.molecules_df = molecules_df 

184 self.atoms_df = atoms_df 

185 self.metadata = metadata 

186 self.averaged_sigmas = averaged_sigmas 

187 self.coords = data[:, :3] 

188 self.charges = data[:, 3] 

189 self.areas = data[:, 4] 

190 

191 @staticmethod 

192 def _reorder_molecule(mol: Chem.Mol) -> Chem.Mol: 

193 """Reorder atoms into ascending AtomMapNum order. 

194 

195 After reordering, atom ``i`` has ``AtomMapNum == i`` (0-based) or 

196 ``AtomMapNum == i + 1`` (1-based), matching the COSMO file's 

197 0-based atom indices used as global atom indices. Atom-map numbers 

198 are required: RDKit's canonical SMILES output doesn't preserve 

199 input atom order, so an unmapped SMILES can't be trusted to match 

200 the COSMO file's atom order. 

201 

202 Parameters 

203 ---------- 

204 mol : Chem.Mol 

205 Molecule to reorder. 

206 

207 Returns 

208 ------- 

209 Chem.Mol 

210 Reordered molecule. 

211 

212 Raises 

213 ------ 

214 ValueError 

215 If the atom map numbers are not a 0-based or 1-based 

216 permutation of the atom indices (including the unmapped case, 

217 where every map number is 0). 

218 """ 

219 num_atoms = mol.GetNumAtoms() 

220 map_nums = {atom.GetAtomMapNum() for atom in mol.GetAtoms()} 

221 if map_nums in (set(range(num_atoms)), set(range(1, num_atoms + 1))): 

222 new_order = sorted( 

223 range(num_atoms), key=lambda i: mol.GetAtomWithIdx(i).GetAtomMapNum() 

224 ) 

225 return Chem.RenumberAtoms(mol, new_order) 

226 

227 raise ValueError( 

228 "Bad atom map numbers: must be a 0-based or 1-based permutation " 

229 "of the atom indices" 

230 ) 

231 

232 @classmethod 

233 def _parse_molecule( 

234 cls, 

235 cosmo_files_dir: pathlib.Path, 

236 filename: str, 

237 smi: str, 

238 fingerprint_generator: FingerprintGenerator, 

239 ) -> tuple[Chem.Mol, pd.DataFrame, pd.DataFrame, float, NDArray[np.int8]]: 

240 """Parse one ``.cosmo`` file/SMILES pair and fingerprint the 

241 molecule. 

242 

243 Parameters 

244 ---------- 

245 cosmo_files_dir : pathlib.Path 

246 Directory holding ``filename``. 

247 filename : str 

248 ``.cosmo`` filename, relative to ``cosmo_files_dir``. 

249 smi : str 

250 SMILES string for this molecule. 

251 fingerprint_generator : FingerprintGenerator 

252 Generator used to fingerprint the parsed molecule. 

253 

254 Returns 

255 ------- 

256 mol, atom_df, segment_df, volume, fingerprint 

257 The reordered molecule, its atom and segment tables, its 

258 volume, and its fingerprint. 

259 

260 Raises 

261 ------ 

262 ValueError 

263 If the SMILES can't be parsed, or its atom count doesn't 

264 match the ``.cosmo`` file's. 

265 """ 

266 _, atom_df, segment_df, volume = parse_cosmo_file( 

267 (cosmo_files_dir / filename).read_text(encoding="utf-8", errors="replace") 

268 ) 

269 mol = Chem.MolFromSmiles(smi, _SMILES_PARSER_PARAMS) 

270 if mol is None: 

271 raise ValueError(f"RDKit could not parse SMILES {smi!r}") 

272 if mol.GetNumAtoms() != len(atom_df): 

273 raise ValueError( 

274 f"SMILES {smi!r} has {mol.GetNumAtoms()} atoms, but " 

275 f"{filename} has {len(atom_df)}" 

276 ) 

277 mol = cls._reorder_molecule(mol) 

278 for i, atom in enumerate(mol.GetAtoms()): 

279 expected_element = atom_df["element"].iat[i] 

280 if atom.GetSymbol() != expected_element: 

281 raise ValueError( 

282 f"SMILES {smi!r} has element {atom.GetSymbol()!r} at " 

283 f"local atom index {i}, but {filename} has element " 

284 f"{expected_element!r} at that same index." 

285 ) 

286 # Re-stamp 1-based, so mol's map numbers survive 

287 # Chem.MolToSmiles's re-canonicalization once stored (GH #43). 

288 atom.SetAtomMapNum(i + 1) 

289 fingerprint = fingerprint_generator.generate(mol) 

290 return mol, atom_df, segment_df, volume, fingerprint 

291 

292 @staticmethod 

293 def _build_molecules_df( # noqa: PLR0913 

294 successful_molecules: list[str], 

295 filenames: list[str], 

296 segment_offsets: list[int], 

297 atom_offsets: list[int], 

298 num_atoms: list[int], 

299 volumes: list[float], 

300 cluster_ids: NDArray[np.int64], 

301 cluster_distance: NDArray[np.float64], 

302 ) -> pd.DataFrame: 

303 """Assemble the per-molecule table written to ``molecules.parquet``.""" 

304 return pd.DataFrame( 

305 { 

306 "smiles": successful_molecules, 

307 "filename": filenames, 

308 "segment_offsets": np.array(segment_offsets, dtype="int64"), 

309 "atom_offsets": np.array(atom_offsets, dtype="int64"), 

310 "num_atoms": np.array(num_atoms, dtype="int64"), 

311 "volume": np.array(volumes, dtype="float64"), 

312 "cluster_id": np.array(cluster_ids, dtype="int64"), 

313 "cluster_distance": np.array(cluster_distance, dtype="float64"), 

314 } 

315 ) 

316 

317 def compute_averaged_sigmas( 

318 self, 

319 schemes: Sequence[AveragingScheme] | None = None, 

320 num_threads: int | None = None, 

321 *, 

322 progress: bool = False, 

323 ) -> dict[str, NDArray[np.float32]]: 

324 """Compute averaged charge densities under each scheme, without 

325 writing to disk or mutating this store. 

326 

327 Parameters 

328 ---------- 

329 schemes : Sequence[AveragingScheme] | None, optional 

330 Schemes to apply. ``None`` (default) uses ``AVERAGING_SCHEMES``. 

331 num_threads : int | None, optional 

332 Thread count. ``None`` (default) uses every CPU core. 

333 progress : bool, optional 

334 If True, show a tqdm bar while averaging. Default False. 

335 

336 Returns 

337 ------- 

338 dict[str, np.ndarray] 

339 Scheme name to ``(n_segs_total,)`` float32 averaged charge 

340 density. 

341 

342 Raises 

343 ------ 

344 ValueError 

345 If a scheme's name would overwrite ``data.npy`` or 

346 ``atom_indices.npy`` on ``save``. 

347 """ 

348 if schemes is None: 

349 schemes = AVERAGING_SCHEMES 

350 for scheme in schemes: 

351 if scheme.name in _RESERVED_SCHEME_NAMES: 

352 raise ValueError( 

353 f"Averaging scheme name {scheme.name!r} collides with a " 

354 f"reserved store filename ({sorted(_RESERVED_SCHEME_NAMES)})." 

355 ) 

356 segment_offsets = self.molecules_df["segment_offsets"].values.astype("int64") 

357 

358 averaged = average_sigmas_by_molecule( 

359 np.asarray(self.coords, dtype=np.float64), 

360 np.asarray(self.charges, dtype=np.float64), 

361 np.asarray(self.areas, dtype=np.float64), 

362 segment_offsets, 

363 schemes, 

364 num_threads=num_threads, 

365 progress=progress, 

366 ) 

367 return { 

368 scheme.name: arr.astype(np.float32) 

369 for scheme, arr in zip(schemes, averaged, strict=True) 

370 } 

371 

372 def assign_splits(self, fractions: Mapping[str, float]) -> NDArray[np.str_]: 

373 """Partition this store's molecules into named splits (e.g. 

374 train/val/test) that respect ``cluster_id`` boundaries, and record 

375 the result as ``molecules_df["split"]``. 

376 

377 Splitting is independent of ``ClusteringSpecs``: it's a cheap pass 

378 over the already-computed ``cluster_id`` column, so it can be 

379 called (and re-called, e.g. to try different fractions) on a 

380 freshly built or a loaded store, without recomputing fingerprints 

381 or clusters. Each call overwrites any previous ``split`` column. 

382 Call ``save`` afterward to persist the result. 

383 

384 Parameters 

385 ---------- 

386 fractions : Mapping[str, float] 

387 Target fraction per split name, e.g. ``{"train": 0.8, "val": 

388 0.1, "test": 0.1}`` or ``{"train": 0.8, "test": 0.2}``. Values 

389 must be positive and sum to 1.0. 

390 

391 Returns 

392 ------- 

393 np.ndarray, shape (n_molecules,) 

394 Split name assigned to each molecule, in ``molecules_df`` row 

395 order (the same array written to ``molecules_df["split"]``). 

396 

397 Raises 

398 ------ 

399 ValueError 

400 If ``fractions`` is empty, contains non-positive values, or 

401 doesn't sum to 1.0. 

402 """ 

403 cluster_ids = self.molecules_df["cluster_id"].values.astype("int64") 

404 labels = greedy_cluster_split(cluster_ids, fractions) 

405 self.molecules_df["split"] = labels 

406 return labels 

407 

408 def subsample( 

409 self, num_molecules: int, shuffle_seed: int | None = None 

410 ) -> "SegmentStore": 

411 """Return a new, unsaved store shrunk to ``num_molecules``, without 

412 moving any molecule across its existing ``split``. 

413 

414 Requires ``molecules_df["split"]`` to already exist (see 

415 ``assign_splits``) -- split membership must be fixed *before* 

416 subsampling, since ``num_molecules`` is apportioned across the 

417 existing splits' own sizes and molecules are only ever dropped from 

418 within a split, never moved to another one. This is what keeps a 

419 test molecule from ending up in train/val just because the store 

420 got smaller. 

421 

422 Parameters 

423 ---------- 

424 num_molecules : int 

425 Total number of molecules to keep, summed across all splits. 

426 Must be positive and not exceed the current molecule count. 

427 shuffle_seed : int | None, optional 

428 If given, each split's share is a uniform random sample without 

429 replacement, drawn with this seed. ``None`` (default) instead 

430 keeps the first molecules in row order within each split, a 

431 deterministic choice. 

432 

433 Returns 

434 ------- 

435 SegmentStore 

436 A new store (see ``subsampling.restrict_to_molecules``) with 

437 the same ``storage_dir`` as this one -- pass a real directory 

438 to ``save`` to persist it. 

439 

440 Raises 

441 ------ 

442 ValueError 

443 If ``molecules_df`` has no ``split`` column, or if 

444 ``num_molecules`` is not positive or exceeds the current 

445 molecule count. 

446 """ 

447 if "split" not in self.molecules_df.columns: 

448 raise ValueError( 

449 "molecules_df has no 'split' column; call assign_splits " 

450 "before subsample." 

451 ) 

452 n_total = len(self.molecules_df) 

453 if num_molecules <= 0 or num_molecules > n_total: 

454 raise ValueError( 

455 f"num_molecules ({num_molecules}) must be positive and not " 

456 f"exceed the current molecule count ({n_total})." 

457 ) 

458 

459 split_labels = self.molecules_df["split"].to_numpy() 

460 split_names = list(dict.fromkeys(split_labels)) 

461 members_by_split = [ 

462 np.flatnonzero(split_labels == name) for name in split_names 

463 ] 

464 sizes = np.array([len(m) for m in members_by_split]) 

465 target_counts = apportion_counts(sizes, num_molecules) 

466 

467 rng = np.random.default_rng(shuffle_seed) if shuffle_seed is not None else None 

468 chosen_parts = [] 

469 for members, raw_target in zip(members_by_split, target_counts, strict=True): 

470 target = int(raw_target) 

471 chosen_parts.append( 

472 members[:target] 

473 if rng is None 

474 else rng.choice(members, size=target, replace=False) 

475 ) 

476 selected = np.sort(np.concatenate(chosen_parts)) 

477 return restrict_to_molecules(self, selected) 

478 

479 def coarse_grain(self) -> "SegmentStore": 

480 """Return a new, unsaved united-atom store: every hydrogen 

481 ``Chem.RemoveHs`` would actually remove is merged into the heavy 

482 atom it's bonded to, and every segment stays attributed to 

483 wherever its atom ended up. 

484 

485 Requires every molecule's ``molecules_df["smiles"]`` to carry 

486 atom-map numbers reflecting local (COSMO) atom index (guaranteed 

487 for stores built after GH issue #43's fix). No molecule or 

488 segment is ever dropped -- only atom attribution shrinks -- so 

489 ``segment_offsets``, ``data``, and every ``averaged_sigmas`` array 

490 carry through unchanged; ``atom_indices``, ``atoms_df`` (rows for 

491 merged hydrogens dropped), and ``molecules_df``'s 

492 ``smiles``/``num_atoms``/``atom_offsets`` columns change. 

493 

494 Returns 

495 ------- 

496 SegmentStore 

497 A new store (see ``coarse_graining.compute_atom_remap``) with 

498 the same ``storage_dir`` as this one -- pass a real directory 

499 to ``save`` to persist it. 

500 

501 Raises 

502 ------ 

503 ValueError 

504 If any molecule's ``smiles`` has no usable atom-map numbers. 

505 """ 

506 molecules_df = self.molecules_df 

507 n_mols = len(molecules_df) 

508 n_segs_total = len(self.data) 

509 old_atom_offsets = molecules_df["atom_offsets"].to_numpy().astype("int64") 

510 segment_offsets = molecules_df["segment_offsets"].to_numpy().astype("int64") 

511 segment_counts = np.diff(np.append(segment_offsets, n_segs_total)) 

512 segment_molecule = np.repeat(np.arange(n_mols), segment_counts) 

513 

514 new_num_atoms = np.empty(n_mols, dtype=np.int64) 

515 new_smiles: list[str] = [""] * n_mols 

516 remaps: list[dict[int, int]] = [{}] * n_mols 

517 survivor_masks: list[NDArray[np.bool_]] = [np.empty(0, dtype=np.bool_)] * n_mols 

518 for m, smi in enumerate(molecules_df["smiles"]): 

519 remap, survivors, mapped_smiles = compute_atom_remap(smi) 

520 remaps[m] = remap 

521 new_num_atoms[m] = len(set(remap.values())) 

522 new_smiles[m] = mapped_smiles 

523 num_atoms_m = int(molecules_df["num_atoms"].iat[m]) 

524 survivor_masks[m] = np.array( 

525 [j in survivors for j in range(num_atoms_m)], dtype=np.bool_ 

526 ) 

527 

528 new_atom_offsets = np.zeros(n_mols, dtype=np.int64) 

529 new_atom_offsets[1:] = np.cumsum(new_num_atoms)[:-1] 

530 

531 old_local = np.asarray(self.atom_indices) - old_atom_offsets[segment_molecule] 

532 new_atom_indices = np.empty(n_segs_total, dtype=np.int64) 

533 # Slice each molecule's own contiguous segment range directly from 

534 # segment_offsets, rather than recomputing `segment_molecule == m` 

535 # (an O(n_segs_total) scan of the FULL segment array) once per 

536 # molecule. segment_molecule is built by construction as contiguous 

537 # per-molecule blocks (np.repeat(np.arange(n_mols), segment_counts) 

538 # above), so `segment_bounds[m]:segment_bounds[m + 1]` is exactly 

539 # that boolean mask's nonzero range -- this makes the loop 

540 # O(n_segs_total) total instead of O(n_mols * n_segs_total), which on 

541 # a large store (e.g. ~53k molecules / ~100M segments) is the 

542 # difference between seconds and hours. 

543 segment_bounds = np.append(segment_offsets, n_segs_total) 

544 for m in range(n_mols): 

545 lo, hi = segment_bounds[m], segment_bounds[m + 1] 

546 lookup = np.array( 

547 [remaps[m][j] for j in range(int(molecules_df["num_atoms"].iat[m]))], 

548 dtype=np.int64, 

549 ) 

550 new_atom_indices[lo:hi] = lookup[old_local[lo:hi]] + new_atom_offsets[m] 

551 

552 new_molecules_df = molecules_df.copy() 

553 new_molecules_df["smiles"] = new_smiles 

554 new_molecules_df["num_atoms"] = new_num_atoms 

555 new_molecules_df["atom_offsets"] = new_atom_offsets 

556 

557 atom_survives = np.concatenate(survivor_masks) 

558 new_atoms_df = self.atoms_df.iloc[atom_survives].reset_index(drop=True) 

559 

560 metadata = StoreMetadata( 

561 num_molecules=n_mols, 

562 num_cosmo_parse_failures=self.metadata.num_cosmo_parse_failures, 

563 schemes=dict(self.metadata.schemes), 

564 ) 

565 return SegmentStore( 

566 self.storage_dir, 

567 self.data, 

568 new_atom_indices, 

569 new_molecules_df, 

570 new_atoms_df, 

571 metadata, 

572 self.averaged_sigmas, 

573 ) 

574 

575 def save(self, storage_dir: pathlib.Path | str | None = None) -> None: 

576 """Write this store's arrays, table, metadata, and averaged 

577 sigmas to disk. 

578 

579 ``metadata.json`` is written atomically (temp file + rename), so an 

580 interrupted save is never reported as complete by ``exists``. 

581 

582 Parameters 

583 ---------- 

584 storage_dir : pathlib.Path | str | None, optional 

585 Destination directory. ``None`` (default) uses 

586 ``self.storage_dir``. Created if missing. 

587 """ 

588 storage_dir = ( 

589 self.storage_dir if storage_dir is None else pathlib.Path(storage_dir) 

590 ) 

591 storage_dir.mkdir(parents=True, exist_ok=True) 

592 

593 np.save(storage_dir / DATA_FILE, self.data) 

594 np.save(storage_dir / ATOM_INDICES_FILE, self.atom_indices) 

595 self.molecules_df.to_parquet(storage_dir / MOLECULES_FILE, index=False) 

596 self.atoms_df.to_parquet(storage_dir / ATOMS_FILE, index=False) 

597 for name, arr in self.averaged_sigmas.items(): 

598 np.save(storage_dir / f"{name}.npy", np.asarray(arr, dtype=np.float32)) 

599 

600 fd, tmp_path = tempfile.mkstemp( 

601 dir=storage_dir, prefix=".metadata-", suffix=".json" 

602 ) 

603 try: 

604 with os.fdopen(fd, "w") as f: 

605 json.dump(self.metadata.to_dict(), f, indent=2) 

606 os.replace(tmp_path, storage_dir / METADATA_FILE) 

607 except BaseException: 

608 os.unlink(tmp_path) 

609 raise 

610 

611 @classmethod 

612 def exists(cls, storage_dir: pathlib.Path | str) -> bool: 

613 """Return whether ``storage_dir`` holds a complete store. 

614 

615 Requires the four fixed files and every scheme ``.npy`` listed in 

616 ``metadata.json``. 

617 

618 Parameters 

619 ---------- 

620 storage_dir : pathlib.Path | str 

621 Directory to check. 

622 

623 Returns 

624 ------- 

625 bool 

626 True if the store is complete. 

627 """ 

628 storage_dir = pathlib.Path(storage_dir) 

629 if not all((storage_dir / f).exists() for f in _STORE_FILES): 

630 return False 

631 try: 

632 with open(storage_dir / METADATA_FILE) as f: 

633 metadata = json.load(f) 

634 except (json.JSONDecodeError, OSError): 

635 return False 

636 scheme_names = metadata.get("schemes", {}) 

637 return all((storage_dir / f"{name}.npy").exists() for name in scheme_names) 

638 

639 @classmethod 

640 def load(cls, storage_dir: pathlib.Path | str) -> "SegmentStore": 

641 """Load an existing segment-data store from disk, memory-mapped. 

642 

643 Parameters 

644 ---------- 

645 storage_dir : pathlib.Path | str 

646 Directory holding a store built by ``from_cosmo_files`` (i.e. 

647 for which ``exists`` is True). 

648 

649 Returns 

650 ------- 

651 SegmentStore 

652 ``data``, ``atom_indices``, and any ``averaged_sigmas`` arrays 

653 are memory-mapped (``mmap_mode="r"``). 

654 

655 Raises 

656 ------ 

657 FileNotFoundError 

658 If ``storage_dir`` doesn't hold a complete store. 

659 """ 

660 storage_dir = pathlib.Path(storage_dir) 

661 if not cls.exists(storage_dir): 

662 raise FileNotFoundError( 

663 f"No segment-data store in {storage_dir} (missing one of " 

664 f"{_STORE_FILES}, or a scheme .npy listed in {METADATA_FILE}; " 

665 "see SegmentStore.from_cosmo_files)." 

666 ) 

667 with open(storage_dir / METADATA_FILE) as f: 

668 metadata = StoreMetadata.from_dict(json.load(f)) 

669 data = np.load(storage_dir / DATA_FILE, mmap_mode="r") 

670 atom_indices = np.load(storage_dir / ATOM_INDICES_FILE, mmap_mode="r") 

671 molecules_df = pd.read_parquet(storage_dir / MOLECULES_FILE) 

672 atoms_df = pd.read_parquet(storage_dir / ATOMS_FILE) 

673 

674 averaged_sigmas = { 

675 name: np.load(storage_dir / f"{name}.npy", mmap_mode="r") 

676 for name in sorted(metadata.schemes) 

677 } 

678 

679 return cls( 

680 storage_dir, 

681 data, 

682 atom_indices, 

683 molecules_df, 

684 atoms_df, 

685 metadata, 

686 averaged_sigmas, 

687 ) 

688 

689 @classmethod 

690 def from_cosmo_files( # noqa: PLR0913, PLR0915, PLR0917 

691 cls, 

692 cosmo_files_dir: pathlib.Path, 

693 filename_to_smiles: Mapping[str, str], 

694 storage_dir: pathlib.Path, 

695 ignore_errors: bool = False, 

696 schemes: Sequence[AveragingScheme] | None = None, 

697 clustering_specs: ClusteringSpecs | None = None, 

698 split_fractions: Mapping[str, float] | None = None, 

699 num_threads: int | None = None, 

700 progress: bool = False, 

701 ) -> "SegmentStore": 

702 """Parse COSMO files, build a store, and write it to ``storage_dir``. 

703 

704 Atoms are numbered in the COSMO file's 0-based order. Each SMILES 

705 must be atom-mapped onto that indexing (0-based or 1-based) and 

706 have the same atom count and per-atom elements as its COSMO file. 

707 Averaged sigmas are computed and written unless ``schemes`` is an 

708 empty sequence. 

709 

710 Parameters 

711 ---------- 

712 cosmo_files_dir : pathlib.Path 

713 Directory containing the ``.cosmo`` files named by 

714 ``filename_to_smiles``'s keys. 

715 filename_to_smiles : Mapping[str, str] 

716 ``.cosmo`` filename (relative to ``cosmo_files_dir``) to that 

717 file's atom-mapped SMILES. Two files of the same molecule may 

718 share a SMILES (same atom order) or carry different SMILES 

719 (different atom orders); each key yields one ``molecules_df`` 

720 row. Keys whose files are not present under 

721 ``cosmo_files_dir`` are skipped and counted in 

722 ``metadata.num_cosmo_parse_failures``. 

723 storage_dir : pathlib.Path 

724 Destination directory for the output files. Created if 

725 missing. 

726 ignore_errors : bool, optional 

727 If True, skip molecules that fail to parse or validate and 

728 count them in ``metadata.num_cosmo_parse_failures``. Missing 

729 files are always skipped, even when this is False. Default 

730 False. 

731 schemes : Sequence[AveragingScheme] | None, optional 

732 Averaging schemes to compute. ``None`` (default) uses 

733 ``AVERAGING_SCHEMES``. Pass ``()`` to skip averaging. 

734 clustering_specs : ClusteringSpecs | None, optional 

735 Fingerprinting and Butina-clustering parameters, used to 

736 assign each molecule a ``cluster_id`` and 

737 ``cluster_distance``. ``None`` (default) uses 

738 ``ClusteringSpecs()``. 

739 split_fractions : Mapping[str, float] | None, optional 

740 Target fraction per named split (e.g. ``{"train": 0.8, "val": 

741 0.1, "test": 0.1}``), assigned via ``assign_splits`` and 

742 written to ``molecules_df["split"]``. ``None`` (default) skips 

743 splitting; the ``split`` column is omitted. Can also be 

744 applied later, without rebuilding, via ``assign_splits`` on a 

745 loaded store. 

746 num_threads : int | None, optional 

747 Thread count for averaging. ``None`` (default) uses every CPU 

748 core. 

749 progress : bool, optional 

750 If True, show tqdm while parsing COSMO files, averaging 

751 sigmas, clustering, and computing cluster medoid distances. 

752 Default False, so a library call stays quiet. 

753 

754 Returns 

755 ------- 

756 SegmentStore 

757 The newly built and saved store. 

758 

759 Raises 

760 ------ 

761 ValueError 

762 If a mapping value is not a SMILES string, if a molecule 

763 cannot be parsed and ``ignore_errors`` is False, if no 

764 molecule could be stored, or if a scheme name collides with a 

765 reserved store filename. 

766 """ 

767 if clustering_specs is None: 

768 clustering_specs = ClusteringSpecs() 

769 

770 data_chunks, atoms_chunks, atom_tables = [], [], [] 

771 segment_offsets, segment_offset = [], 0 

772 atom_offsets, atom_offset = [], 0 

773 fingerprints = [] 

774 num_atoms = [] 

775 volumes = [] 

776 successful_molecules = [] 

777 filenames = [] 

778 present_items = [ 

779 (filename, smi) 

780 for filename, smi in filename_to_smiles.items() 

781 if (cosmo_files_dir / filename).is_file() 

782 ] 

783 num_cosmo_parse_failures = len(filename_to_smiles) - len(present_items) 

784 if num_cosmo_parse_failures: 

785 tqdm.write( 

786 f"Skipping {num_cosmo_parse_failures} missing COSMO file" 

787 f"{'s' if num_cosmo_parse_failures != 1 else ''}." 

788 ) 

789 fingerprint_generator = FingerprintGenerator(clustering_specs) 

790 

791 for filename, smi in tqdm( 

792 present_items, 

793 desc="Processing COSMO files", 

794 disable=not progress, 

795 ): 

796 if not isinstance(smi, str): 

797 raise ValueError( 

798 f"filename_to_smiles[{filename!r}] must be a SMILES " 

799 f"string, got {smi!r}." 

800 ) 

801 try: 

802 mol, atom_df, segment_df, volume, fingerprint = cls._parse_molecule( 

803 cosmo_files_dir, filename, smi, fingerprint_generator 

804 ) 

805 except (ValueError, AssertionError) as e: 

806 if ignore_errors: 

807 tqdm.write(f"Error parsing {filename}: {e}") 

808 num_cosmo_parse_failures += 1 

809 continue 

810 else: 

811 raise e 

812 

813 data_chunks.append( 

814 segment_df[["x", "y", "z", "charge", "area"]].values.astype("float32") 

815 ) 

816 atoms_chunks.append(segment_df["atom"].values.astype("int64") + atom_offset) 

817 atom_tables.append(atom_df[["id", "element", "x", "y", "z"]]) 

818 fingerprints.append(fingerprint) 

819 segment_offsets.append(segment_offset) 

820 segment_offset += len(segment_df) 

821 atom_offsets.append(atom_offset) 

822 atom_offset += len(atom_df) 

823 num_atoms.append(len(atom_df)) 

824 volumes.append(volume) 

825 successful_molecules.append(Chem.MolToSmiles(mol)) 

826 filenames.append(filename) 

827 

828 if not successful_molecules: 

829 raise ValueError("No COSMO files could be parsed successfully.") 

830 

831 fingerprint_array = np.stack(fingerprints, axis=0) 

832 cluster_ids = butina_cluster( 

833 fingerprint_array, clustering_specs.cutoff, progress=progress 

834 ) 

835 cluster_distance = cluster_medoid_distances( 

836 fingerprint_array, cluster_ids, progress=progress 

837 ) 

838 

839 data = np.concatenate(data_chunks, axis=0) 

840 atom_indices = np.concatenate(atoms_chunks) 

841 atoms_df = pd.concat(atom_tables, ignore_index=True) 

842 molecules_df = cls._build_molecules_df( 

843 successful_molecules, 

844 filenames, 

845 segment_offsets, 

846 atom_offsets, 

847 num_atoms, 

848 volumes, 

849 cluster_ids, 

850 cluster_distance, 

851 ) 

852 metadata = StoreMetadata( 

853 num_molecules=len(successful_molecules), 

854 num_cosmo_parse_failures=num_cosmo_parse_failures, 

855 ) 

856 

857 store = cls( 

858 pathlib.Path(storage_dir), 

859 data, 

860 atom_indices, 

861 molecules_df, 

862 atoms_df, 

863 metadata, 

864 {}, 

865 ) 

866 if schemes is None or schemes: 

867 resolved_schemes = AVERAGING_SCHEMES if schemes is None else schemes 

868 store.averaged_sigmas = store.compute_averaged_sigmas( 

869 schemes=resolved_schemes, 

870 num_threads=num_threads, 

871 progress=progress, 

872 ) 

873 store.metadata.schemes.update({s.name: s for s in resolved_schemes}) 

874 if split_fractions is not None: 

875 store.assign_splits(split_fractions) 

876 store.save() 

877 return store 

878 

879 def sigmas(self, scheme: str | None = None) -> NDArray[np.float64]: 

880 """Resolve this store's segment charge density for a given 

881 scheme. 

882 

883 Parameters 

884 ---------- 

885 scheme : str | None, optional 

886 Which charge density to return: None (default) returns raw 

887 ``charges / areas``; a scheme name returns