Coverage for cosmolayer/store/__main__.py: 96%
74 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""CLI: build a segment-data store (if missing) and report summary
2statistics for atom- and molecule-level charges, areas, and sigma
3profiles.
5Run as::
7 cosmostore --storage-dir DIR \\
8 --cosmo-files-dir DIR --filenames-to-smiles FILE.json
9"""
11import argparse
12import json
13import pathlib
14from collections.abc import Sequence
16import numpy as np
17from numpy.typing import NDArray
19from .binning import compute_per_atom_properties, compute_per_molecule_properties
20from .grid import SigmaGrid
21from .profiles import SigmaProfileTable
22from .reporting import print_stats
23from .segments import SegmentStore
26def get_parser() -> argparse.ArgumentParser:
27 """Return the argument parser for cosmostore (used by sphinx-argparse)."""
28 arg_parser = argparse.ArgumentParser(
29 prog="cosmostore",
30 description=(
31 "Build the segment data store (if missing) and print summary "
32 "statistics for atom- and molecule-level charges, areas, and "
33 "sigma profiles."
34 ),
35 )
36 arg_parser.add_argument(
37 "--storage-dir",
38 type=str,
39 required=True,
40 help="The directory to store (or load) the segment data.",
41 )
42 arg_parser.add_argument(
43 "--cosmo-files-dir",
44 type=str,
45 default=None,
46 help=(
47 "Directory containing the .cosmo files referenced by "
48 "--filenames-to-smiles. Required only if --storage-dir doesn't "
49 "already hold a built store."
50 ),
51 )
52 arg_parser.add_argument(
53 "--filenames-to-smiles",
54 type=str,
55 default=None,
56 help=(
57 "Path to a JSON file mapping each .cosmo filename (relative "
58 "to --cosmo-files-dir) to that file's atom-mapped SMILES. "
59 "Required only if --storage-dir doesn't already hold a built "
60 "store."
61 ),
62 )
63 arg_parser.add_argument(
64 "--num-threads",
65 type=int,
66 default=None,
67 help=(
68 "Number of threads to use for every threaded step (segment "
69 "averaging, sigma-profile binning). Defaults to every "
70 "available CPU core -- lower this on a shared machine to "
71 "leave headroom for other users."
72 ),
73 )
74 arg_parser.add_argument(
75 "--sigma-scheme",
76 type=str,
77 default=None,
78 help=(
79 "Averaging scheme for every statistic (e.g. 'cosmo-rs', "
80 "'cosmo-sac-2002', 'cosmo-sac-2010'). The store must already "
81 "contain that scheme. Default: raw sigma."
82 ),
83 )
84 arg_parser.add_argument(
85 "--no-progress",
86 action="store_true",
87 help=(
88 "Disable tqdm progress bars. By default, bars are shown for "
89 "COSMO parsing, averaging, clustering, and sigma-profile work."
90 ),
91 )
92 arg_parser.add_argument(
93 "--ignore-errors",
94 action="store_true",
95 help=(
96 "Skip .cosmo files that fail to parse or validate when "
97 "building a store. By default those failures abort the run."
98 ),
99 )
100 return arg_parser
103def _ensure_store_built(
104 args: argparse.Namespace,
105 arg_parser: argparse.ArgumentParser,
106 storage_dir: pathlib.Path,
107) -> None:
108 """Build the store at ``storage_dir`` if it doesn't already exist."""
109 if SegmentStore.exists(storage_dir):
110 print("Segment data already exists.")
111 return
112 if args.cosmo_files_dir is None or args.filenames_to_smiles is None:
113 arg_parser.error(
114 "--cosmo-files-dir and --filenames-to-smiles are required "
115 f"when --storage-dir ({storage_dir}) doesn't already hold "
116 "a built store."
117 )
118 print("Storing segment data and averaged sigmas...")
119 cosmo_files_dir = pathlib.Path(args.cosmo_files_dir)
120 with open(args.filenames_to_smiles) as f:
121 filename_to_smiles = json.load(f)
122 SegmentStore.from_cosmo_files(
123 cosmo_files_dir,
124 filename_to_smiles,
125 storage_dir,
126 ignore_errors=args.ignore_errors,
127 num_threads=args.num_threads,
128 progress=not args.no_progress,
129 )
132def _compute_atom_data(
133 store: SegmentStore,
134 sigma_scheme: str | None,
135 num_threads: int | None,
136 *,
137 progress: bool,
138) -> tuple[
139 SigmaProfileTable, NDArray[np.float32], NDArray[np.float32], NDArray[np.intp]
140]:
141 """Compute per-atom properties and centered sigma profiles."""
142 total_num_atoms = int(store.molecules_df["num_atoms"].sum())
143 atom_charges = compute_per_atom_properties(
144 np.asarray(store.charges), store.atom_indices, total_num_atoms
145 )
146 atom_areas = compute_per_atom_properties(
147 np.asarray(store.areas), store.atom_indices, total_num_atoms
148 )
149 atom_segment_counts = np.bincount(store.atom_indices, minlength=total_num_atoms)
150 atom_sigma_profiles = store.compute_atom_sigma_profiles(
151 scheme=sigma_scheme,
152 grid=SigmaGrid(),
153 num_threads=num_threads,
154 centered=True,
155 progress=progress,
156 )
157 return atom_sigma_profiles, atom_areas, atom_charges, atom_segment_counts
160def _compute_molecule_data(
161 store: SegmentStore,
162 atom_sigma_profiles: SigmaProfileTable,
163 atom_areas: NDArray[np.float32],
164 atom_charges: NDArray[np.float32],
165 num_threads: int | None,
166 *,
167 progress: bool,
168) -> tuple[NDArray[np.float32], NDArray[np.float32], SigmaProfileTable]:
169 """Compute per-molecule properties and aggregated sigma profiles."""
170 atom_offsets = store.molecules_df["atom_offsets"].values.astype("int64")
171 molecule_areas = compute_per_molecule_properties(atom_areas, atom_offsets)
172 molecule_charges = compute_per_molecule_properties(atom_charges, atom_offsets)
173 molecule_sigma_profiles = atom_sigma_profiles.aggregate(
174 num_threads=num_threads, progress=progress
175 )
176 return molecule_areas, molecule_charges, molecule_sigma_profiles
179def _print_atom_stats(
180 atom_sigma_profiles: SigmaProfileTable,
181 atom_areas: NDArray[np.float32],
182 atom_charges: NDArray[np.float32],
183 atom_segment_counts: NDArray[np.intp],
184) -> None:
185 """Print per-atom statistics from already-computed arrays."""
186 print_stats("Atom charges", atom_charges)
187 print_stats("Atom areas", atom_areas)
188 print_stats("Atom segment counts", atom_segment_counts)
189 has_area = atom_sigma_profiles.areas > 0
190 print(
191 f"Atoms with no surface segments: {(~has_area).sum()} of "
192 f"{len(atom_sigma_profiles.areas)}"
193 )
194 first_moments = (
195 atom_sigma_profiles.profiles[has_area].astype(np.float64)
196 @ atom_sigma_profiles.sigma_values
197 )
198 print_stats("Atom profile first moments", first_moments, value_format=".3e")
201def _print_molecule_stats(
202 molecule_areas: NDArray[np.float32],
203 molecule_charges: NDArray[np.float32],
204 molecule_sigma_profiles: SigmaProfileTable,
205) -> None:
206 """Print per-molecule statistics from already-computed arrays."""
207 print_stats("Molecule areas", molecule_areas)
208 print_stats("Molecule charges", molecule_charges)
209 mass_err = np.abs(molecule_sigma_profiles.profiles.sum(axis=1) / molecule_areas - 1)
210 print(
211 f"Molecule profile mass conservation, max relative error: {mass_err.max():.2e}"
212 )
215def main(argv: Sequence[str] | None = None) -> int:
216 """Entry point for ``cosmostore``.
218 Parameters
219 ----------
220 argv : Sequence[str] | None, optional
221 Arguments to parse, by default None, meaning ``sys.argv[1:]``.
223 Returns
224 -------
225 int
226 Process exit code (0 on success).
227 """
228 arg_parser = get_parser()
229 args = arg_parser.parse_args(argv)
230 storage_dir = pathlib.Path(args.storage_dir)
232 _ensure_store_built(args, arg_parser, storage_dir)
233 store = SegmentStore.load(storage_dir)
235 if args.sigma_scheme is not None and args.sigma_scheme not in store.averaged_sigmas:
236 arg_parser.error(
237 f"No averaged sigmas for scheme {args.sigma_scheme!r} in "
238 f"{storage_dir}. Known schemes: {sorted(store.averaged_sigmas)}."
239 )
241 progress = not args.no_progress
242 atom_sigma_profiles, atom_areas, atom_charges, atom_segment_counts = (
243 _compute_atom_data(
244 store, args.sigma_scheme, args.num_threads, progress=progress
245 )
246 )
247 molecule_areas, molecule_charges, molecule_sigma_profiles = _compute_molecule_data(
248 store,
249 atom_sigma_profiles,
250 atom_areas,
251 atom_charges,
252 args.num_threads,
253 progress=progress,
254 )
255 _print_atom_stats(
256 atom_sigma_profiles, atom_areas, atom_charges, atom_segment_counts
257 )
258 _print_molecule_stats(molecule_areas, molecule_charges, molecule_sigma_profiles)
259 return 0
262if __name__ == "__main__":
263 raise SystemExit(main())