| """SI Table S8 (Dataset Summary Statistics): molecule and 1H/13C site counts for the training datasets. |
| |
| Site counts from each file's flat ``atomic_numbers`` (1H = 1, 13C = 6); molecule counts from each |
| file's per-molecule index. Vectorized reads, so seconds even on the multi-million-molecule files. The |
| MagNET-Zero rows combine the two sigma-pepper groups (initial PBE0/pcSseg-1, WP04/wB97X-D/pcSseg-2 |
| fine-tuning) with sigma-concentrate. |
| """ |
| import os |
|
|
| import h5py |
| import pandas as pd |
|
|
| import paths |
|
|
| |
| PUBLISHED_S8 = { |
| "sigma-shake": (4_787_816, 72_175_813, 50_542_154), |
| "MagNET-Zero (initial round)": (2_167_210, 24_047_999, 15_330_361), |
| "MagNET-Zero (fine-tuning)": (370_105, 4_042_135, 2_557_211), |
| "MagNET-Zero (both rounds combined)": (2_537_315, 28_090_134, 17_887_572), |
| "sigma-concentrate": (50_000, 754_437, 528_569), |
| } |
|
|
|
|
| def _counts(node): |
| """(molecules, 1H sites, 13C sites) for an open h5py file or group.""" |
| z = node["atomic_numbers"][:] |
| n_1h = int((z == 1).sum()) |
| n_13c = int((z == 6).sum()) |
| if "molecule_id" in node: |
| n_mol = node["molecule_id"].shape[0] |
| elif "smiles_offsets" in node: |
| n_mol = node["smiles_offsets"].shape[0] - 1 |
| elif "sigma_shake_id" in node: |
| n_mol = node["sigma_shake_id"].shape[0] |
| else: |
| raise KeyError("node has no recognized per-molecule key " |
| "(molecule_id / smiles_offsets / sigma_shake_id)") |
| return n_mol, n_1h, n_13c |
|
|
|
|
| def counts(path, group=None): |
| """(molecules, 1H sites, 13C sites) for one HDF5 file, optionally within one group.""" |
| with h5py.File(path, "r") as f: |
| return _counts(f[group] if group else f) |
|
|
|
|
| def summary_table(data_dir): |
| """SI Table S8 as a DataFrame with columns dataset / molecules / n_1H_sites / n_13C_sites. |
| |
| data_dir is the release repo's ``data/`` directory; the three source files are read from their |
| dataset subfolders. The two MagNET-Zero training rounds and their combination are assembled from |
| the sigma-pepper groups plus sigma-concentrate. |
| """ |
| |
| root = os.path.dirname(os.path.normpath(data_dir)) |
| shake = paths.dataset_file("sigma-shake", root=root) |
| pepper = paths.dataset_file("sigma-pepper", root=root) |
| concentrate = paths.dataset_file("sigma-concentrate", root=root) |
|
|
| sigma_shake = counts(shake) |
| initial = counts(pepper, "pbe0_pcSseg1") |
| pepper_ft = counts(pepper, "wp04_wb97xd_pcSseg2") |
| sigma_concentrate = counts(concentrate) |
| fine_tuning = tuple(a + b for a, b in zip(pepper_ft, sigma_concentrate)) |
| combined = tuple(a + b for a, b in zip(initial, fine_tuning)) |
|
|
| rows = [ |
| ("sigma-shake", sigma_shake), |
| ("MagNET-Zero (initial round)", initial), |
| ("MagNET-Zero (fine-tuning)", fine_tuning), |
| ("MagNET-Zero (both rounds combined)", combined), |
| ("sigma-concentrate", sigma_concentrate), |
| ] |
| return pd.DataFrame( |
| [{"dataset": label, "molecules": m, "n_1H_sites": h, "n_13C_sites": c} |
| for label, (m, h, c) in rows] |
| ) |
|
|
