File size: 3,413 Bytes
ef53368 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """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 SI Table S8 values, for the exact-reproduction check: label -> (molecules, 1H, 13C).
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: # offsets are n_molecules + 1
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.
"""
# data_dir is the repo's data/ folder; its parent is the repo root the deposit fallback needs.
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") # MagNET-Zero initial round
pepper_ft = counts(pepper, "wp04_wb97xd_pcSseg2") # fine-tuning: sigma-pepper part
sigma_concentrate = counts(concentrate) # fine-tuning: sigma-concentrate part
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]
)
|