diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000000000000000000000000000000000000..c96a322a2cf2732b62ec29d3df69de9c8688c91d --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,44 @@ +name: tests + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: pip install -r requirements.txt "pytest>=7" + - name: Run tests + run: pytest -q + # Tests run against a synthetic fixture; tests requiring the large + # Hugging Face-hosted .hdf5 files skip automatically when those files are absent. + + # The job above never installs torch, so conftest.py skips the whole magnet/ package. This job + # installs the CPU inference stack and actually imports + tests magnet/, catching packaging and API + # regressions (e.g. a clean-install `import magnet` failure). The weights-gated inference tests still + # skip here (checkpoints are not in the repo); everything else runs. + model: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install the CPU inference stack + run: | + pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/cpu + pip install -r magnet/requirements.txt + pip install --no-deps ./magnet + pip install "pytest>=7" + - name: Run the magnet package tests + run: pytest magnet/ -q diff --git a/.github/workflows/publish-pypi.yml b/.github/workflows/publish-pypi.yml new file mode 100644 index 0000000000000000000000000000000000000000..5dbc8281a552ef172018e48ad8665d7101eb3152 --- /dev/null +++ b/.github/workflows/publish-pypi.yml @@ -0,0 +1,34 @@ +name: publish to PyPI + +# Builds the magnet-nmr package and publishes it to PyPI when you push a version tag (e.g. v0.1.0). +# +# Uses PyPI Trusted Publishing (OIDC), so there is no token to store: configure the trusted publisher +# once at https://pypi.org/manage/project/magnet-nmr/settings/publishing/ with +# owner=ekwan repo=MagNET workflow=publish-pypi.yml environment=pypi +# +# The package lives in the magnet/ subdirectory (its own pyproject.toml), so every step runs there. + +on: + push: + tags: ["v*"] + +jobs: + publish: + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write # required for Trusted Publishing; no API token needed + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Build sdist and wheel + working-directory: magnet + run: | + pip install build + python -m build + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: magnet/dist/ diff --git a/.github/workflows/sync-to-hf.yml b/.github/workflows/sync-to-hf.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce105aa85621da487b5792a97236fa61411ecabf --- /dev/null +++ b/.github/workflows/sync-to-hf.yml @@ -0,0 +1,51 @@ +name: sync code to hugging face + +# Mirrors the CODE to the Hugging Face model repo ekwan16/MagNET on every push to main. +# +# The large model weights and datasets live ONLY on Hugging Face (as Git LFS objects). They are +# gitignored here, so a fresh checkout never contains them, and they are pushed once by hand from a +# workstation (see publish/build_hf_repo.py). This job must therefore NEVER delete anything on the +# Hub. +# +# `hf upload` is additive/overwrite-only: it removes remote files ONLY when given an explicit +# --delete flag. +# ==> NEVER add a --delete flag below. It would wipe the ~42 GB of LFS weights and datasets. <== +# The --exclude patterns below are defense-in-depth for the big-file extensions (a fresh checkout +# has none of them anyway); the actual safety comes from omitting --delete. + +on: + push: + branches: [main] + +concurrency: + group: sync-to-hf + cancel-in-progress: true + +jobs: + sync: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install the Hugging Face CLI + # Pinned as defense-in-depth so a future CLI default cannot silently change behavior. + # Tighten to the exact version you tested with once you have run this successfully. + run: pip install "huggingface_hub[cli]>=0.34,<1.0" + - name: Ensure the target repo exists and is private + # `hf upload` auto-creates a MISSING repo as PUBLIC. If this Action ever fires before the + # repo has been seeded by hand, this guard creates it private first so nothing is exposed + # early. If the repo already exists this create call errors and `|| true` ignores it (it + # never changes an existing repo's visibility); a bad token still fails loudly at upload. + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} # must be a WRITE token for ekwan16/MagNET + run: hf repo create ekwan16/MagNET --repo-type=model --private || true + - name: Upload code to the Hub (no deletions) + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} # must be a WRITE token for ekwan16/MagNET + run: | + hf upload ekwan16/MagNET . . \ + --repo-type=model \ + --exclude ".git/**" ".github/**" "*.hdf5" "*.mnova" "*.ckpt" \ + --commit-message "sync code from github@${{ github.sha }}" diff --git a/analysis/code/build_nb_dataset_summary.py b/analysis/code/build_nb_dataset_summary.py new file mode 100644 index 0000000000000000000000000000000000000000..2f8a21cad52005813db171b67a79610b335e59eb --- /dev/null +++ b/analysis/code/build_nb_dataset_summary.py @@ -0,0 +1,90 @@ +"""Source of truth for the Table S8 notebook (Dataset Summary Statistics). Edit the cell sources +here, then regenerate: + + python3 build_nb_dataset_summary.py si_table_s08_summary + jupyter nbconvert --to notebook --execute --inplace ../si_tables/si_table_s08_summary.ipynb + +Regenerating overwrites the .ipynb (clearing its execution outputs). All real code lives in +dataset_summary.py (the counting logic and the PUBLISHED_S8 reference values); the notebook only +runs it over the released data/ files. The notebook lives in analysis/si_tables/ (grouped by role), not here. + +Run with no arguments to list the available notebook names. +""" +import os +import sys + +from nb_build import md, code, save_notebook as _save + + +_BOOTSTRAP = r""" +import os, sys + +# make the in-repo modules importable (not pip-installed) +REPO = os.path.abspath("../..") +for _p in ("analysis/code", "analysis/code/shared"): + sys.path.insert(0, os.path.join(REPO, _p)) +""" + +_IMPORTS = r""" +import pandas as pd +import dataset_summary +""" + +_SETUP = r""" +DATA_DIR = os.path.join(REPO, "data") + +def document_path(name): + os.makedirs("documents", exist_ok=True) + return os.path.join("documents", name) +""" + +si_table_s08_summary = [ + md(r""" +# Table S8: Dataset Summary Statistics + +Molecule and ¹H/¹³C site counts for each training dataset (site counts from each HDF5's +`atomic_numbers`, ¹H=1/¹³C=6; MagNET-Zero combines both sigma-pepper rounds with sigma-concentrate). +"""), + code(_BOOTSTRAP), + code(_IMPORTS), + code(_SETUP), + code(r""" +table_s8 = dataset_summary.summary_table(DATA_DIR) +display(table_s8) + +# write the table to this notebook's documents/ folder +out = document_path("si_table_s08_summary.xlsx") +with pd.ExcelWriter(out) as writer: + table_s8.to_excel(writer, sheet_name="Table S8", index=False) +print("wrote", os.path.relpath(out, REPO)) +"""), + md("## Exact-reproduction check"), + code(r""" +# every count should match the published SI Table S8 value exactly +for _, row in table_s8.iterrows(): + pub = dataset_summary.PUBLISHED_S8[row["dataset"]] + got = (row["molecules"], row["n_1H_sites"], row["n_13C_sites"]) + assert got == pub, f"{row['dataset']}: {got} != published {pub}" +print("all rows match the published SI Table S8 exactly") +"""), +] + + +# name -> (cells, path relative to repo root) +NOTEBOOKS = { + "si_table_s08_summary": (si_table_s08_summary, "analysis/si_tables/si_table_s08_summary.ipynb"), +} + +if __name__ == "__main__": + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.abspath(os.path.join(here, "..", "..")) + names = sys.argv[1:] + if not names: + print("available names:", ", ".join(NOTEBOOKS)) + sys.exit(0) + unknown = [n for n in names if n not in NOTEBOOKS] + if unknown: + raise SystemExit(f"unknown notebook name(s): {unknown}; available: {', '.join(NOTEBOOKS)}") + for name in names: + cells, relpath = NOTEBOOKS[name] + _save(cells, os.path.join(repo, relpath)) diff --git a/analysis/code/build_nb_leveling.py b/analysis/code/build_nb_leveling.py new file mode 100644 index 0000000000000000000000000000000000000000..7ba9892575d38a2aed917b737b34c126fe48a578 --- /dev/null +++ b/analysis/code/build_nb_leveling.py @@ -0,0 +1,295 @@ +# build_nb_leveling.py -- generator for the SI Figure S3 leveling notebook. +# +# Source of truth for the notebook. Edit the cell sources here, then regenerate and re-execute: +# +# python3 build_nb_leveling.py si_figure_s03_leveling +# jupyter nbconvert --to notebook --execute --inplace ../si_figures/si_figure_s03_leveling.ipynb +# +# Regenerating overwrites the .ipynb (clearing its execution outputs), so do not hand-edit it. +# All real code lives in leveling.py; the notebook only runs the protocol and shows results. +# +# Covers NS372 across all 8 nuclei, plus delta22 (the same leveling.py analysis run on both +# datasets). Main-text Figure 2B (the clean delta22 correlation matrix) is built separately in +# analysis/manuscript_figures/fig2b_correlation.ipynb; Figure 2C is a hand-drawn schematic, not +# reproduced in code. This notebook computes both res_ns372 and res_delta22 (via the shared +# setup/load cells below) so its combined cross-dataset summary table is self-contained. +from nb_build import md, code, save_notebook as _save + + +_TITLE_S3 = r""" +# SI Figure S3: inter-method correlation and PCA of NS372 and delta-22 shieldings + +Inter-method correlation matrices and PC1/PC2 loadings for NS372 (44 functionals x 8 nuclei) and +delta22 (18 functionals, ¹H/¹³C), plus a combined summary table. +""" + +_METHOD_INTRO = r""" +Two independent NMR shielding datasets, analyzed separately (never pooled): + +| Dataset | Source | Methods | Nuclei | Reference | +|---|---|---|---|---| +| **NS372** | Schattenberg & Kaupp, *JCTC* **17**, 7602 (2021) | 44 DFT/WFT functionals | ¹H ¹¹B ¹³C ¹⁵N ¹⁷O ¹⁹F ³¹P ³³S | CCSD(T)/pcSseg-3 | +| **delta22** | in-house `delta22.hdf5` | 18 gas-phase functionals, largest basis (pcSseg-3; mp2 → pcSseg-2), PBE0/cc-pVTZ geometry | ¹H ¹³C | DSD-PBEP86 (highest-rung in-set method, stands in for CCSD(T)) | +""" + +_CONFIG_MD = "## 1. Configuration" + +_BOOTSTRAP = r""" +import os, sys + +REPO = os.path.abspath("../..") +for _p in ("analysis/code", "analysis/code/shared"): + sys.path.insert(0, os.path.join(REPO, _p)) +""" + +_IMPORTS = r""" +import glob +import pandas as pd +import matplotlib.pyplot as plt + +import paths +import leveling +import leveling_plots +""" + +_PATH_SETUP = r""" +# inputs: +# - the Kaupp NS372 spreadsheet is small and ships in the repo's data/ns372/ folder +# - delta22.hdf5 is large: it resolves from the repo's data/delta22/ folder, carried by the +# Hugging Face checkout via Git LFS (see analysis/code/paths.py) +# the delta22 file is encoded as whole numbers (real value x 10,000); the loader +# in leveling.py decodes it on read (a plain-decimal copy also works, since the +# decode step is a no-op on floating-point data). +KAUPP_XLSX = os.path.join(REPO, "data", "ns372", "ct1c00919_si_002.xlsx") +DELTA22_H5 = paths.dataset_file("delta22", root=REPO) +SAVE_DPI = 200 # SI-quality raster output + +def figure_path(name): + os.makedirs("figures", exist_ok=True) + return os.path.join("figures", name) + +# self-clean: this notebook builds PNG names dynamically (one per nucleus), so drop any +# previously written figures before regenerating (glob on a missing folder returns []) +for _stale in glob.glob(os.path.join("figures", "*.png")): + os.remove(_stale) + +pd.set_option('display.width', 150) +pd.set_option('display.max_columns', 25) +""" + +_STYLE_CONFIG = r""" +# functional-family colours used for every figure (family ordering lives in leveling.py) +FAMILY_COLORS = {'LDA':'#777777','GGA':'#1f77b4','mGGA':'#17becf','GH':'#2ca02c', + 'RSH':'#9467bd','LH':'#8c564b','DH':'#ff7f0e','WFT':'#d62728', + 'ref':'#FF1493'} + +# proper Unicode-superscript labels for nuclei in figure titles +NUC_DISPLAY = {'1H':'¹H','11B':'¹¹B','13C':'¹³C', + '15N':'¹⁵N','17O':'¹⁷O','19F':'¹⁹F', + '31P':'³¹P','33S':'³³S'} + +# adjustText parameters tuned per nucleus -- 1H and 13C have very dense central +# clusters and need stronger expansion; the paramagnetic-shielding nuclei +# (15N/17O/19F) already spread methods along the parabola and need a gentler +# pass to avoid over-flinging labels. +PCA_ADJUST_DEFAULT = dict( + force_text=(0.35, 0.55), force_explode=(0.25, 0.40), + force_static=(0.10, 0.15), force_pull=(0.02, 0.02), + expand=(1.25, 1.35), time_lim=4, +) +PCA_ADJUST = { + '1H': {**PCA_ADJUST_DEFAULT, 'force_text':(0.75, 1.05), + 'force_explode':(0.65, 0.90), 'expand':(1.7, 1.9), 'time_lim':7}, + '11B': {**PCA_ADJUST_DEFAULT, 'force_text':(0.50, 0.75), + 'force_explode':(0.40, 0.60), 'expand':(1.4, 1.55), 'time_lim':5}, + '13C': {**PCA_ADJUST_DEFAULT, 'force_text':(0.70, 1.00), + 'force_explode':(0.60, 0.85), 'expand':(1.6, 1.8), 'time_lim':6}, + '15N': {**PCA_ADJUST_DEFAULT, 'force_text':(0.60, 0.85), + 'force_explode':(0.50, 0.70), 'expand':(1.5, 1.65), 'time_lim':6}, + '17O': {**PCA_ADJUST_DEFAULT, 'force_text':(0.60, 0.85), + 'force_explode':(0.50, 0.70), 'expand':(1.5, 1.65), 'time_lim':6}, + '31P': {**PCA_ADJUST_DEFAULT, 'force_text':(0.65, 0.90), + 'force_explode':(0.55, 0.75), 'expand':(1.55, 1.7), 'time_lim':6}, + '33S': {**PCA_ADJUST_DEFAULT, 'force_text':(0.65, 0.90), + 'force_explode':(0.55, 0.75), 'expand':(1.55, 1.7), 'time_lim':6}, +} +""" + +_NS372_DEF_MD = r""" +## 2. NS372 (Kaupp) - definitions + +Conventional GIAO shieldings for 44 functionals across 8 main-group nuclei, with a +CCSD(T)/pcSseg-3 reference. Kaupp Reduced-Set exclusions (F₃⁻, O₃, BH - multireference +outliers) are applied. Input: `ct1c00919_si_002.xlsx`. +""" + +_DELTA22_DEF_MD = r""" +## 3. delta22 - definitions + +Gas-phase conventional GIAO shieldings from `delta22.hdf5`: 18 functionals at their +largest available basis (pcSseg-3; plain `mp2` only to pcSseg-2), at the PBE0/cc-pVTZ +geometry. Observations are pooled ¹H / ¹³C atom sites across all 22 solutes. There is no +CCSD(T) reference in the file; **DSD-PBEP86 is used as the reference** for delta22 - it +is the highest-rung double-hybrid available in this method set and stands in for CCSD(T) +in the same role. The stored whole-number shieldings are decoded on read. +""" + +_LOAD_MD = r""" +## 4. Load datasets + global colour scale + +Run the loaders, analyse every nucleus, and compute the figure-wide +`-log10(1-|r|)` maximum (`GLOBAL_VMAX`) used as the colour scale on every correlation matrix +below, for NS372 and delta22 alike. +""" + +_LOAD = r""" +ns372 = leveling.load_ns372(KAUPP_XLSX) +delta22 = leveling.load_delta22(DELTA22_H5) + +res_ns372 = {nuc: leveling.analyze_nucleus(d['M'], d['methods'], d['ref']) + for nuc, d in ns372.items()} +res_delta22 = {nuc: leveling.analyze_nucleus(d['M'], d['methods'], d['ref']) + for nuc, d in delta22.items()} + +GLOBAL_VMAX = max(leveling.dataset_logr_max(res_ns372), leveling.dataset_logr_max(res_delta22)) +print(f'NS372 : {len(ns372)} nuclei') +print(f'delta22: {len(delta22)} nuclei (reference = {leveling.DELTA22_REF})') +print(f'global colour scale: 0 -> {GLOBAL_VMAX} on the -log10(1-|r|) axis') +for nuc, d in ns372.items(): + print(f' NS372 {nuc:4s}: {d["M"].shape[0]:4d} mols x {d["M"].shape[1]-1} methods + CCSD(T)') +for nuc, d in delta22.items(): + print(f' delta22 {nuc:4s}: {d["M"].shape[0]:4d} sites x {d["M"].shape[1]} methods') +""" + +SHARED = [ + md(_CONFIG_MD), + code(_BOOTSTRAP), + code(_IMPORTS), + code(_PATH_SETUP), + code(_STYLE_CONFIG), + md(_NS372_DEF_MD), + md(_DELTA22_DEF_MD), + md(_LOAD_MD), + code(_LOAD), +] + +# ---------------------------------------------------------------------------- +# SI Figure S3: NS372 results across all 8 nuclei + the combined cross-dataset summary +# ---------------------------------------------------------------------------- +si_figure_s03_leveling = [ + md(_TITLE_S3), + md(_METHOD_INTRO), + *SHARED, + md('## 5. NS372 results'), + md('### 5.1 Leveling diagnostics - NS372'), + code(r""" +sum_ns372 = leveling.summarize(ns372, res_ns372) +print('NS372 - leveling diagnostics (CCSD(T) included as a method column):') +sum_ns372.round(5) +"""), + md(r""" +### 5.2 Per-method scaled RMSE vs CCSD(T) + +`scaled_rmse` = RMSE of residuals after a per-method linear fit +`sigma_method ~ a*sigma_CCSD(T) + b` - the error that survives empirical linear scaling. +"""), + code(r""" +rmse_ns372 = pd.DataFrame({nuc: res_ns372[nuc]['scaled_rmse'] for nuc in res_ns372}) +rmse_ns372 = rmse_ns372.drop(index='CCSD(T)', errors='ignore') +rmse_ns372 = rmse_ns372.reindex(ns372['1H']['methods'][:-1]) # family order +print('NS372 - per-method scaled RMSE vs CCSD(T) (ppm):') +rmse_ns372.round(3) +"""), + md('### 5.3 Inter-method correlation matrices - NS372 (one nucleus per file)'), + code(r""" +NS372_NUCS = ['1H', '11B', '13C', '15N', '17O', '19F', '31P', '33S'] +for nuc in NS372_NUCS: + fams = {nuc: ns372[nuc]['families']} + fig = leveling_plots.plot_corr_matrix('NS372', [nuc], res_ns372, fams, + GLOBAL_VMAX, ref_name='CCSD(T)', + nuc_display=NUC_DISPLAY) + fig.savefig(figure_path(f'si_figure_s03_corr_{nuc}.png'), + dpi=SAVE_DPI, bbox_inches='tight') + plt.show() +"""), + md('### 5.4 PC1/PC2 structure - NS372 (one nucleus per file)'), + code(r""" +for nuc in NS372_NUCS: + fams = {nuc: ns372[nuc]['families']} + fig = leveling_plots.plot_pca_pair('NS372', [nuc], res_ns372, fams, + family_colors=FAMILY_COLORS, nuc_display=NUC_DISPLAY, + pca_adjust=PCA_ADJUST, pca_adjust_default=PCA_ADJUST_DEFAULT) + fig.savefig(figure_path(f'si_figure_s03_pca_{nuc}.png'), + dpi=SAVE_DPI, bbox_inches='tight') + plt.show() +"""), + md(r""" +### 5.5 delta22 panels (lead the published figure: corr ¹H/¹³C + PCA) + +The same correlation + PCA layout on the delta22 gas-phase set. DSD-PBEP86 is the reference (delta22 +has no CCSD(T)); it appears as an ordinary double-hybrid point in the PCA, with no CCSD(T) star. +"""), + code(r""" +# delta22 correlation matrices (canonical S3 A = 1H, B = 13C) +for nuc in ['1H', '13C']: + fams = {nuc: delta22[nuc]['families']} + fig = leveling_plots.plot_corr_matrix('delta22', [nuc], res_delta22, fams, + GLOBAL_VMAX, ref_name=leveling.DELTA22_REF, + nuc_display=NUC_DISPLAY) + fig.savefig(figure_path(f'si_figure_s03_delta22_corr_{nuc}.png'), + dpi=SAVE_DPI, bbox_inches='tight') + plt.show() + +# delta22 PCA projection: 1H and 13C stacked in one figure (canonical S3 C) +fams = {nuc: delta22[nuc]['families'] for nuc in ['1H', '13C']} +fig = leveling_plots.plot_pca_pair('delta22', ['1H', '13C'], res_delta22, fams, + family_colors=FAMILY_COLORS, nuc_display=NUC_DISPLAY, + pca_adjust=PCA_ADJUST, pca_adjust_default=PCA_ADJUST_DEFAULT) +fig.savefig(figure_path('si_figure_s03_delta22_pca.png'), + dpi=SAVE_DPI, bbox_inches='tight') +plt.show() +"""), + md('## 6. Combined summary (both datasets)'), + code(r""" +# delta22's per-nucleus diagnostics (same summarize() used for NS372 above), computed here +# too so this combined table is self-contained -- it reuses res_delta22 from the shared load +# cell above, so this is cheap (no new correlation/PCA computation). +sum_delta22 = leveling.summarize(delta22, res_delta22) + +combined = pd.concat([sum_ns372.assign(dataset='NS372'), + sum_delta22.assign(dataset='delta22')], ignore_index=True) +combined = combined[['dataset','nucleus','n_obs','n_methods','r_min','r_median', + 'PC1_pct','PC2_pct','PC3plus_pct','parabola_R2']] +print('Leveling effect - both datasets (analyzed separately):') +print(f' PC1 range : {combined.PC1_pct.min():.2f}% - {combined.PC1_pct.max():.2f}%') +print(f' min pairwise r : {combined.r_min.min():.5f} (worst case, all nuclei)') +print(f' parabola R2 range : {combined.parabola_R2.min():.3f} - {combined.parabola_R2.max():.3f}') +print(f' global colour scale : 0 -> {GLOBAL_VMAX} on -log10(1-|r|)') +combined.round(5) +"""), +] + +# name -> (cells, path relative to repo root). Notebooks are grouped by role (manuscript/si_figures/ +# si_tables), not by dataset. +NOTEBOOKS = { + "si_figure_s03_leveling": (si_figure_s03_leveling, "analysis/si_figures/si_figure_s03_leveling.ipynb"), +} + +if __name__ == "__main__": + import os + import sys + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.abspath(os.path.join(here, "..", "..")) + names = sys.argv[1:] + if not names: + print("usage: python3 build_nb_leveling.py [ ...]") + print("regenerates ONLY the named notebook(s) -- pick just the one(s) you edited,") + print("since regenerating clears a notebook's execution outputs.") + print("available names:", ", ".join(NOTEBOOKS)) + raise SystemExit(1) + unknown = [n for n in names if n not in NOTEBOOKS] + if unknown: + raise SystemExit(f"unknown notebook name(s): {unknown}; available: {', '.join(NOTEBOOKS)}") + for name in names: + cells, relpath = NOTEBOOKS[name] + _save(cells, os.path.join(repo, relpath)) diff --git a/analysis/code/carbon_monoxide.py b/analysis/code/carbon_monoxide.py new file mode 100644 index 0000000000000000000000000000000000000000..4684d8f99cc567aa8cd8d27ba2d99371d57963d4 --- /dev/null +++ b/analysis/code/carbon_monoxide.py @@ -0,0 +1,63 @@ +"""Figure S12: MagNET vs DFT carbon-13 shielding of carbon monoxide across bond lengths. + +The CO bond is stretched/compressed; at each length the 13C shielding is computed by DFT +(PBE1PBE/pcSseg-1, gas, Gaussian) and by the MagNET foundation model. This module holds only the +loading; the plot is drawn inline in the figure notebook. + +Two files ship alongside: +- `co_magnet_vs_gaussian.csv` - one row per bond length: DFT and MagNET shieldings and their + difference, for carbon and oxygen. +- `co_bond_lengths_histogram.csv` - carbon-oxygen bond-length distribution across real molecules, + as (bin left edge, count) rows, for the shaded histogram. +""" +import os + +import pandas as pd + +from stats import mae + +HERE = os.path.dirname(os.path.abspath(__file__)) +COMPARISON_CSV = os.path.join(HERE, "..", "..", "data", "carbon_monoxide", "co_magnet_vs_gaussian.csv") +HISTOGRAM_CSV = os.path.join(HERE, "..", "..", "data", "carbon_monoxide", "co_bond_lengths_histogram.csv") + + +def load_comparison(path=COMPARISON_CSV): + """The bond-length scan as a DataFrame, sorted by bond length. Columns: `bond_length` (angstrom), + and for each of carbon and oxygen the `gaussian`, `magnet`, and `delta` (MagNET minus DFT) + shieldings in ppm.""" + raw = pd.read_csv(path) + return pd.DataFrame( + { + "bond_length": raw["bond_length_angstrom"].astype(float), + "gaussian_c": raw["gaussian_shielding_c_ppm"].astype(float), + "magnet_c": raw["magnet_shielding_c_ppm"].astype(float), + "delta_c": raw["delta_c_ppm"].astype(float), + "gaussian_o": raw["gaussian_shielding_o_ppm"].astype(float), + "magnet_o": raw["magnet_shielding_o_ppm"].astype(float), + "delta_o": raw["delta_o_ppm"].astype(float), + } + ).sort_values("bond_length", ignore_index=True) + + +def load_bond_length_histogram(path=HISTOGRAM_CSV): + """The sampled carbon-oxygen bond-length distribution as `(centers, counts, bin_width)`: the bin + centers in angstrom, the count in each bin, and the (uniform) bin width. The stored table holds + left bin edges and counts.""" + table = pd.read_csv(path) + left_edges = table["bin_left_angstrom"].to_numpy(dtype=float) + counts = table["count"].to_numpy(dtype=float) + bin_width = float(left_edges[1] - left_edges[0]) if len(left_edges) > 1 else 0.01 + centers = left_edges + bin_width / 2.0 + return centers, counts, bin_width + + +def equilibrium_vs_extreme_error(comparison=None, window=(1.0, 1.3), extreme=(0.9, 1.4)): + """The figure's claim as two numbers: the mean absolute carbon error (ppm) for bond lengths + inside the near-equilibrium `window`, and for the non-equilibrium tails outside `extreme`. The + second is far larger, which is the whole point of the panel.""" + df = load_comparison() if comparison is None else comparison + near = df[(df["bond_length"] >= window[0]) & (df["bond_length"] <= window[1])] + tails = df[(df["bond_length"] < extreme[0]) | (df["bond_length"] > extreme[1])] + return mae(near["delta_c"], 0), mae(tails["delta_c"], 0) + + diff --git a/analysis/code/cp3.py b/analysis/code/cp3.py new file mode 100644 index 0000000000000000000000000000000000000000..11690ea2a7d783100c986ac474561ec4d573d614 --- /dev/null +++ b/analysis/code/cp3.py @@ -0,0 +1,55 @@ +"""Figure 1A data and analysis: informative 1H shift differences are too small for DFT to resolve. + +Loads per-site 1H shift spread across the Goodman CP3 stereoisomers and classifies each site against +the experimental noise floor (0.02 ppm) and DFT's resolving power (0.10 ppm); ~36% fall in between. +Data: `goodman2009_cp3.xlsx` (Goodman, J. Org. Chem. 2009, 74, 4597), shipped alongside. Plotting is +in the figure notebook. +""" +import os + +import numpy as np +import pandas as pd + +HERE = os.path.dirname(os.path.abspath(__file__)) +CP3_XLSX = os.path.join(HERE, "..", "..", "data", "cp3", "goodman2009_cp3.xlsx") + +# accuracy thresholds in ppm: the experimental noise floor, the DFT resolving limit, and the cutoff +# above which a variation counts as large. These define the zones a site's variation falls into. +EXPERIMENTAL_LIMIT = 0.02 +DFT_LIMIT = 0.10 +LARGE_VARIATION = 0.30 + + +def load_variations(path=CP3_XLSX, nucleus="H"): + """The per-site chemical-shift standard deviations across the CP3 stereoisomers for one nucleus + ("H" or "C"), as a 1D array of finite values (ppm).""" + df = pd.read_excel(path) + x = df[df["nucleus"] == nucleus]["stdev"].to_numpy(dtype=float) + return x[np.isfinite(x)] + + +def _zone(value): + """Classify one variation (ppm) into its accuracy zone: within experimental noise, informative + but below the DFT resolving limit, resolvable by DFT, or large.""" + if value < EXPERIMENTAL_LIMIT: + return "below_experimental" + if value < DFT_LIMIT: + return "below_dft" + if value < LARGE_VARIATION: + return "dft_zone" + return "large" + + +def fraction_below_dft(variations, lo=EXPERIMENTAL_LIMIT, hi=DFT_LIMIT, bins=30): + """The fraction of the variation histogram's area between the experimental floor and the DFT + limit: the sites whose shift variation is informative but too small for DFT to resolve (the + paper's 36%). Computed as histogram area in [lo, hi] over total area, splitting the bins that + straddle a boundary, exactly as the figure does.""" + edges = np.linspace(0.0, float(variations.max()), bins + 1) + counts, edge = np.histogram(variations, bins=edges) + total = window = 0.0 + for height, left, right in zip(counts, edge[:-1], edge[1:]): + total += height * (right - left) + overlap = max(0.0, min(right, hi) - max(left, lo)) + window += height * overlap + return window / total if total > 0 else float("nan") diff --git a/analysis/code/delta22_plots.py b/analysis/code/delta22_plots.py new file mode 100644 index 0000000000000000000000000000000000000000..d02c5438ce43948df6edbf984d46a97d6da7a049 --- /dev/null +++ b/analysis/code/delta22_plots.py @@ -0,0 +1,606 @@ +"""Plotting engines for the delta-22 SI figure notebooks (analysis/si_figures/*.ipynb). + +Each function is the drawing engine behind one delta-22 SI figure panel. Notebook globals an engine +needs (color maps, axis-label lookups, N_SPLITS, solvent lists, etc.) are explicit function +parameters, passed in by the notebook at the call site. The numbers come from delta22.py; this +module only draws, except `ss_fits`, a small data-prep helper used only by the si_figure_s14 +notebook. +""" +import numpy as np +import pandas as pd +import seaborn as sns +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +from matplotlib.colors import ListedColormap +from scipy.stats import pearsonr, linregress +from adjustText import adjust_text + +import delta22 + + +# ---------------------------------------------------------------------------- +# color helpers shared by the engine/source panels (si_figure_s08, si_figure_s13) +# ---------------------------------------------------------------------------- +def darken_color(hex_color, factor=0.7): + """Scale a "#rrggbb" color's channels by factor (< 1 darkens).""" + hex_color = hex_color.lstrip("#") + r, g, b = (int(hex_color[i:i + 2], 16) for i in (0, 2, 4)) + r, g, b = (int(c * factor) for c in (r, g, b)) + return f"#{r:02x}{g:02x}{b:02x}" + + +def lighten_color(hex_color, amount=0.45): + """Blend a "#rrggbb" color toward white by amount (0-1).""" + hex_color = hex_color.lstrip("#") + r, g, b = (int(hex_color[i:i + 2], 16) for i in (0, 2, 4)) + r, g, b = (int(c + (255 - c) * amount) for c in (r, g, b)) + return f"#{r:02x}{g:02x}{b:02x}" + + +def display_solvent_name(name): + """Human-readable solvent label ("Water (TIP4P)" for TIP4P, else capitalized).""" + return "Water (TIP4P)" if name == "TIP4P" else name.capitalize() + + +# ---------------------------------------------------------------------------- +# SI Figure S4: correlation of implicit (PCM) corrections across solvents and methods +# ---------------------------------------------------------------------------- +def plot_correlation_matrix(corr_matrix, title, caption, colormap="Reds", show_values=True, save_path=None): + """Lower-triangle heatmap of -log10(1 - r) with a Pearson-R colorbar; a value of 3 means r=0.999.""" + arr = corr_matrix.to_numpy(dtype=float, copy=True) + np.fill_diagonal(arr, np.nan) + transformed = pd.DataFrame(-np.log10(1 - arr), index=corr_matrix.index, columns=corr_matrix.columns) + mask = np.tril(np.ones_like(transformed, dtype=bool), k=0).T + vmin, vmax = -np.log10(1 - 0.9), -np.log10(1 - 0.9999) + fig = plt.figure(figsize=(8, 8)) + ax = sns.heatmap(transformed, mask=mask, annot=transformed if show_values else False, fmt=".1f", + cmap=colormap, cbar_kws={"shrink": 0.5, "pad": -0.13}, + vmin=vmin, vmax=vmax, square=True) + cbar = ax.collections[0].colorbar + cbar.set_ticks([-np.log10(1 - r) for r in (0.99, 0.999, 0.9999)]) + cbar.set_ticklabels(["0.99", "0.999", "0.9999"]) + cbar.ax.set_title("Pearson R", fontweight="bold", pad=25) + ax.set_xlabel(""); ax.set_ylabel("") + ax.set_xticklabels(ax.get_xticklabels(), fontweight="bold") + ax.set_yticklabels(ax.get_yticklabels(), fontweight="bold") + ax.set_yticks(ax.get_yticks()[1:]); ax.set_xticks(ax.get_xticks()[:-1]) + plt.title(title, fontweight="bold") + plt.figtext(0.59, 0.895, caption, wrap=True, horizontalalignment="center", fontsize=10) + for side in ("top", "right", "bottom", "left"): + plt.gca().spines[side].set_visible(True) + plt.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + plt.show() + + +def plot_pcm_scatter(x_vals, y_vals, solvent_x, solvent_y, nucleus, save_path=None): + """Scatter of PCM corrections in one solvent vs another, with a best-fit line, slope, and Pearson R.""" + pearson_r = x_vals.corr(y_vals) + slope, _ = np.polyfit(x_vals, y_vals, 1) + cap = lambda s: s[0].upper() + s[1:] + fig = plt.figure(figsize=(6.5, 6.5)) + sns.scatterplot(x=x_vals, y=y_vals, s=60, color="black", edgecolor="none") + sns.regplot(x=x_vals, y=y_vals, scatter=False, color="black", ci=None, + line_kws={"linewidth": 1, "linestyle": "--"}) + plt.xlabel(f"{cap(solvent_x)} PCM correction") + plt.ylabel(f"{cap(solvent_y)} PCM correction") + plt.title(f"PCM corrections: {cap(solvent_x)} vs {cap(solvent_y)} ({nucleus})", fontweight="bold") + plt.text(0.05, 0.96, f"Slope = {slope:.3f}", transform=plt.gca().transAxes, va="top") + plt.text(0.05, 0.92, f"Pearson R = {pearson_r:.3f}", transform=plt.gca().transAxes, va="top") + plt.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + plt.show() + + +# ---------------------------------------------------------------------------- +# SI Figure S5: per-solvent PCM benefit vs bulk dielectric constant and polarizability +# ---------------------------------------------------------------------------- +def plot_pcm_benefit_vs_properties(benefit, dielectric, polarizability, nucleus_label, + exclude=(), title_extra="", figsize=(14, 6), save_path=None): + """Per-solvent PCM benefit (percent test-RMSE reduction) vs dielectric constant and + polarizability, one subplot each, with a best-fit line and Pearson R. exclude drops solvents.""" + solvents = [s for s in benefit.index if s not in set(exclude) and s in dielectric] + y = np.array([benefit[s] for s in solvents], dtype=float) + fig, axes = plt.subplots(1, 2, figsize=figsize) + for ax, prop, xlabel in [(axes[0], dielectric, "Dielectric Constant"), + (axes[1], polarizability, "Polarizability")]: + x = np.array([prop[s] for s in solvents], dtype=float) + r, _ = pearsonr(x, y) + slope, intercept, _, _, _ = linregress(x, y) + ax.scatter(x, y, s=100, color="black", alpha=1, edgecolors="black", linewidth=1, zorder=3) + ax.axhline(0, color="gray", linestyle="--", alpha=0.4, linewidth=1, zorder=1) + xr = np.linspace(x.min(), x.max(), 100) + ax.plot(xr, slope * xr + intercept, color="gray", linestyle="--", linewidth=2, + label=f"R = {r:.3f}", zorder=2) + ax.set_xlabel(xlabel, fontsize=13, fontweight="bold") + ax.set_ylabel("Percent Reduction in Test RMSE (%)", fontsize=13, fontweight="bold") + title = f"{nucleus_label} Nucleus: PCM Benefit vs {xlabel}" + if title_extra: + title += f"\n({title_extra})" + ax.set_title(title, fontsize=14, fontweight="bold") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=12) + texts = [ax.text(sx, sy, name, fontsize=9, fontweight="bold") + for sx, sy, name in zip(x, y, solvents)] + adjust_text(texts, ax=ax, arrowprops=dict(arrowstyle="-", color="gray", lw=0.5)) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + plt.show() + return fig + + +# ---------------------------------------------------------------------------- +# SI Figure S6: per-solvent test RMSE of six solvent-correction models +# ---------------------------------------------------------------------------- +def plot_formula_ladder_boxplot(results_df, formula_labels, solvent_groups, nucleus="H", + colors=None, figsize=(14, 7), title="", save_path=None): + """Per-solvent box-and-whisker of the test RMSE for a solvent-correction formula ladder, with + solvents grouped by class (a dashed line separates classes).""" + formulas = list(formula_labels) + n = len(formulas) + if colors is None: + colors = sns.color_palette("colorblind", n) + # lay the solvents out class by class (only those present), tracking where classes end + flat, boundaries = [], [] + for group in solvent_groups: + present = [s for s in solvent_groups[group] if s in set(results_df["solvent"])] + flat.extend(present) + boundaries.append(len(flat)) + fig, ax = plt.subplots(figsize=figsize) + box_width = 0.8 / n + gap = 0.6 + centers, separators = [], [] + pos = 0.0 + for si, solvent in enumerate(flat): + centers.append(pos) + for fi, formula in enumerate(formulas): + vals = results_df[(results_df["formula"] == formula) + & (results_df["solvent"] == solvent)]["test_RMSE"].dropna().to_numpy() + bp = ax.boxplot([vals], positions=[pos + (fi - (n - 1) / 2) * box_width], + widths=box_width * 0.9, showfliers=False, patch_artist=True, + manage_ticks=False) + bp["boxes"][0].set(facecolor=colors[fi], alpha=0.9) + bp["medians"][0].set(color="black") + if si == 0: + bp["boxes"][0].set_label(formula_labels[formula]) + pos += 1.0 + if (si + 1) in boundaries[:-1]: # a class just ended (not the last) + separators.append(pos - 0.5 + gap / 2) + pos += gap + for x in separators: + ax.axvline(x, color="gray", linestyle="--", linewidth=1) + ax.set_xticks(centers) + ax.set_xticklabels(flat, rotation=60, ha="right", fontweight="bold") + ax.set_ylabel("Test RMSE (ppm)", fontweight="bold") + ax.set_ylim(bottom=0) + ax.set_title(title, fontweight="bold") + ax.legend(fontsize=9, loc="upper left") + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +# ---------------------------------------------------------------------------- +# SI Figure S7: DFT 13C explicit-solvent correction, Desmond vs OpenMM +# ---------------------------------------------------------------------------- +def plot_desmond_vs_openmm_grid(pairs_by_solvent, figsize=(10, 10), save_path=None): + """Per OpenMM solvent, a scatter of the Desmond vs OpenMM explicit correction at each site. + Tight clustering on the diagonal shows the correction is independent of the MD engine.""" + solvents = list(pairs_by_solvent) + fig, axes = plt.subplots(2, 2, figsize=figsize) + axes = np.atleast_1d(axes).flatten() + for ax, solvent in zip(axes, solvents): + p = pairs_by_solvent[solvent] + ax.scatter(p["openMM"], p["desmond"], c="k", s=5) + ax.set_xlabel("OpenMM solvent correction (ppm)", fontweight="bold") + ax.set_ylabel("Desmond solvent correction (ppm)", fontweight="bold") + title = "Water (TIP4P)" if solvent == "TIP4P" else solvent.capitalize() + ax.set_title(title, fontweight="bold") + for ax in axes[len(solvents):]: + ax.set_visible(False) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +# ---------------------------------------------------------------------------- +# SI Figure S8 (panels B-D): explicit-solvent correction vs MD-frame count +# ---------------------------------------------------------------------------- +def plot_frame_convergence(running_by_label, finals_by_label, label_colors, label_names, + title="", xlabel="Number of Frames", figsize=(6, 5), save_path=None): + """Running-average correction vs number of frames, one line per label, with each label's + converged value drawn as a dashed horizontal line.""" + fig, ax = plt.subplots(figsize=figsize) + for label, running in running_by_label.items(): + color = label_colors.get(label) + valid = running[~np.isnan(running)] + display = label_names.get(label, label) + ax.plot(np.arange(1, len(valid) + 1), valid, color=color, lw=1.1, label=display) + final = finals_by_label[label] + ax.axhline(final, color=color, ls="--", lw=1, label=f"{display} final: {final:.3f} ppm") + ax.set_xlabel(xlabel, fontweight="bold") + ax.set_ylabel("Running Average Correction (ppm)", fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.legend(fontsize=8) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_frame_correction_histogram(values_by_label, total_frames_by_label, label_colors, label_names, + bins=30, title="", figsize=(6, 4.5), save_path=None): + """Distribution of valid per-frame corrections, one overlaid histogram per label; each bin's + height is a fraction of that label's total frame count, with mean/std in the legend.""" + fig, ax = plt.subplots(figsize=figsize) + for label, values in values_by_label.items(): + valid = values[~np.isnan(values)] + weights = np.full(len(valid), 1.0 / total_frames_by_label[label]) + display = label_names.get(label, label) + ax.hist(valid, bins=bins, weights=weights, color=label_colors.get(label), alpha=0.55, + edgecolor="black", linewidth=0.3, + label=f"{display} ($\\mu$={valid.mean():.3f}, $\\sigma$={valid.std():.3f})") + ax.set_xlabel("Correction per Frame (ppm)", fontweight="bold") + ax.set_ylabel("Frequency / Frame Count", fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.legend(fontsize=8) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_frame_autocorrelation(autocorr_by_label, label_colors, label_names, title="", + figsize=(8, 4.5), save_path=None): + """Autocorrelation of the per-frame correction vs lag (frames), one line per label.""" + fig, ax = plt.subplots(figsize=figsize) + for label, autocorr in autocorr_by_label.items(): + ax.plot(np.arange(len(autocorr)), autocorr, color=label_colors.get(label), lw=1, + label=label_names.get(label, label)) + ax.axhline(0, color="0.8", lw=0.8) + ax.set_xlabel("Lag (frames)", fontweight="bold") + ax.set_ylabel("Autocorrelation", fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.legend() + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_frame_validity_heatmaps(grids_by_engine, solvents_by_engine, engine_labels, solute="AcOH", + figsize=(10, 8), save_path=None): + """One heatmap per MD engine: which trajectory frames have computed DFT shielding data (light + blue) vs not (dark gray), one row per solvent and one column per frame index.""" + engines = list(grids_by_engine) + fig, axes = plt.subplots(len(engines), 1, figsize=figsize, squeeze=False) + cmap = ListedColormap(["#1a1a1a", "#a8dadc"]) # False = dark gray, True = light blue + for ax, engine in zip(axes[:, 0], engines): + ax.imshow(grids_by_engine[engine], aspect="auto", cmap=cmap, vmin=0, vmax=1, + interpolation="nearest") + ax.set_yticks(range(len(solvents_by_engine[engine]))) + ax.set_yticklabels(solvents_by_engine[engine], fontsize=8) + ax.set_title(f"{solute} Frame Validities ({engine_labels.get(engine, engine)})", + fontweight="bold", fontsize=10) + axes[-1, 0].set_xlabel("Frame index") + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +# ---------------------------------------------------------------------------- +# SI Figure S11: MagNET vs DFT rovibrational (QCD) corrections +# ---------------------------------------------------------------------------- +def plot_qcd_scatter(qcd, save_path=None): + """NN vs DFT QCD correction, one subplot per nucleus, with a y=x guide and an R^2/RMSE/MAE + stats box measured against that line.""" + fig, axes = plt.subplots(1, 2, figsize=(12, 6)) + for ax, nucleus in zip(axes, ["H", "C"]): + sub = qcd[qcd["nucleus"] == nucleus] + x, y = sub["qcd_dft"].to_numpy(), sub["qcd_nn"].to_numpy() + ss_res, ss_tot = np.sum((y - x) ** 2), np.sum((y - y.mean()) ** 2) + r2 = 1 - ss_res / ss_tot if ss_tot != 0 else np.nan + rmse, mae = np.sqrt(np.mean((y - x) ** 2)), np.mean(np.abs(y - x)) + print(f"{nucleus}: R^2={r2:.4f}, RMSE={rmse:.4f}, MAE={mae:.4f}") + ax.scatter(x, y, color="black") + raw = [min(x.min(), y.min()), max(x.max(), y.max())] + buf = 0.05 * (raw[1] - raw[0]) + lims = [raw[0] - buf, raw[1] + buf] + ax.plot(lims, lims, linestyle="--", color="gray", label="NN matches DFT", zorder=-1) + ax.set_xlim(lims); ax.set_ylim(lims) + ax.set_xlabel("DFT QCD"); ax.set_ylabel("NN QCD") + ax.set_title(f"DFT vs NN QCD ({nucleus})", fontweight="bold") + ax.legend() + ax.text(0.02, 0.98, f"R^2 = {r2:.3f}\nRMSE = {rmse:.3f}\nMAE = {mae:.3f}", + transform=ax.transAxes, ha="left", va="top", fontsize=12, fontweight="bold", + bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8)) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_qcd_error_histogram(qcd, save_path=None): + """Distribution of the NN-minus-DFT QCD error, one subplot per nucleus.""" + fig, axes = plt.subplots(1, 2, figsize=(12, 4), sharey=True) + for ax, nucleus in zip(axes, ["H", "C"]): + data = qcd[qcd["nucleus"] == nucleus]["error"] + ax.hist(data, bins=30, color="gray", alpha=0.8, edgecolor="black") + ax.set_title(f"QCD Error Distribution (NN - DFT) for {nucleus}", fontweight="bold") + ax.set_xlabel("Error"); ax.set_ylabel("Count") + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_qcd_correction_by_site(by_site, nucleus, legend_loc, dft_color="#A0A0A0", nn_color="#707070", + save_path=None): + """The DFT and NN QCD correction for every site, one column of points per solute; sites within + a solute are jittered horizontally and joined by a faint vertical line.""" + fig, ax = plt.subplots(figsize=(12, 6)) + solute_order = sorted(by_site["solute"].unique()) + x_pos = {s: i for i, s in enumerate(solute_order)} + jitter = 0.02 + site_offset = {} + for solute, group in by_site.groupby("solute"): + sites = sorted(group["site"].unique()) + if len(sites) == 1: + site_offset[(solute, sites[0])] = 0.0 + else: + for site, off in zip(sites, np.linspace(-jitter, jitter, len(sites))): + site_offset[(solute, site)] = off + for _, row in by_site.iterrows(): + x = x_pos[row["solute"]] + site_offset[(row["solute"], row["site"])] + ax.vlines(x, row["qcd_dft"], row["qcd_nn"], color=dft_color, alpha=0.2, linewidth=1) + ax.scatter(x, row["qcd_dft"], color=dft_color, s=32, edgecolor="black", linewidth=0.3, alpha=0.85) + ax.scatter(x, row["qcd_nn"], color=nn_color, s=32, edgecolor="black", linewidth=0.3, alpha=0.85) + ax.set_xticks(range(len(solute_order))) + ax.set_xticklabels(solute_order, rotation=65, ha="right", fontweight="bold") + ax.set_ylabel("QCD correction (ppm)") + ax.set_title(f"QCD: DFT vs NN ({nucleus})", fontweight="bold") + ax.grid(axis="y", linestyle="--", alpha=0.3) + handles = [Line2D([0], [0], color=dft_color, marker="o", linestyle="-", label="DFT QCD"), + Line2D([0], [0], color=nn_color, marker="o", linestyle="-", label="NN QCD")] + ax.legend(handles=handles, ncol=1, fontsize=9, frameon=True, loc=legend_loc) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +# ---------------------------------------------------------------------------- +# SI Figure S13: MagNET-x vs DFT explicit-solvent corrections +# ---------------------------------------------------------------------------- +def plot_dft_nn_scatter_by_engine(compare_df, engine_colors, engine_labels, save_path=None): + """NN vs DFT explicit correction at each site, both nuclei, Desmond and OpenMM overlaid, with + a y=x guide.""" + fig, axes = plt.subplots(1, 2, figsize=(12, 6)) + for ax, nucleus in zip(axes, ["H", "C"]): + sub = compare_df[compare_df["nucleus"] == nucleus] + allv = [] + for engine in ["desmond", "openMM"]: + d = sub[sub["engine"] == engine] + ax.scatter(d["dft_value"], d["nn_value"], alpha=0.9, color=engine_colors[engine], + label=engine_labels[engine]) + allv += [d["dft_value"].to_numpy(), d["nn_value"].to_numpy()] + allv = np.concatenate(allv) + raw = [allv.min(), allv.max()] + buf = 0.05 * (raw[1] - raw[0]) + lims = [raw[0] - buf, raw[1] + buf] + ax.plot(lims, lims, linestyle="--", color="gray", label="NN matches DFT", zorder=-1) + ax.set_xlim(lims); ax.set_ylim(lims) + ax.set_xlabel("DFT Explicit Correction"); ax.set_ylabel("NN Explicit Correction") + ax.set_title(f"DFT vs NN Explicit Corrections ({nucleus})", fontweight="bold") + ax.legend() + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_dft_nn_error_histogram_by_engine(compare_df, engine_colors, engine_labels, save_path=None): + """Distribution of the NN-minus-DFT explicit-correction error, both nuclei, Desmond and OpenMM + overlaid.""" + fig, axes = plt.subplots(1, 2, figsize=(12, 4), sharey=True) + for ax, nucleus in zip(axes, ["H", "C"]): + sub = compare_df[compare_df["nucleus"] == nucleus] + for engine in ["desmond", "openMM"]: + data = sub[sub["engine"] == engine]["error"].to_numpy() + if len(data) == 0: + continue + ax.hist(data, bins=30, color=engine_colors[engine], alpha=0.6, edgecolor="black", + label=engine_labels[engine]) + ax.set_title(f"Explicit Correction Error (NN - DFT) for {nucleus}", fontweight="bold") + ax.set_xlabel("Error"); ax.set_ylabel("Count") + ax.legend() + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_explicit_correction_by_site(pairs_df, title, engine_colors, save_path=None): + """The Desmond and OpenMM explicit correction from both DFT (base hue) and NN (darkened), one + column per solute; the two engines are offset, sites are jittered, and each DFT/NN pair is + joined by a faint vertical line.""" + fig, ax = plt.subplots(figsize=(12, 6)) + solute_order = sorted(pairs_df["solute"].unique()) + x_pos = {s: i for i, s in enumerate(solute_order)} + engine_offset = {"desmond": -0.12, "openMM": 0.12} + jitter = 0.02 + site_offset = {} + for solute, group in pairs_df.groupby("solute"): + sites = sorted(group["site"].unique()) + if len(sites) == 1: + site_offset[(solute, sites[0])] = 0.0 + else: + for site, off in zip(sites, np.linspace(-jitter, jitter, len(sites))): + site_offset[(solute, site)] = off + dft_color = {e: engine_colors[e] for e in ("desmond", "openMM")} + nn_color = {e: darken_color(engine_colors[e]) for e in ("desmond", "openMM")} + for engine in ["desmond", "openMM"]: + for _, row in pairs_df.iterrows(): + x = x_pos[row["solute"]] + engine_offset[engine] + site_offset[(row["solute"], row["site"])] + dft_val, nn_val = row[f"{engine}_dft"], row[f"{engine}_nn"] + ax.vlines(x, dft_val, nn_val, color=dft_color[engine], alpha=0.18, linewidth=1) + ax.scatter(x, dft_val, color=dft_color[engine], s=32, edgecolor="black", linewidth=0.3, alpha=0.85) + ax.scatter(x, nn_val, color=nn_color[engine], s=32, edgecolor="black", linewidth=0.3, alpha=0.85) + ax.set_xticks(range(len(solute_order))) + ax.set_xticklabels(solute_order, rotation=65, ha="right", fontweight="bold") + ax.set_ylabel("Explicit Solvent Correction (ppm)") + ax.set_title(title, fontweight="bold") + ax.grid(axis="y", linestyle="--", alpha=0.3) + handles = [Line2D([0], [0], color=dft_color["desmond"], marker="o", linestyle="-", label="DFT Desmond"), + Line2D([0], [0], color=nn_color["desmond"], marker="o", linestyle="-", label="NN Desmond"), + Line2D([0], [0], color=dft_color["openMM"], marker="o", linestyle="-", label="DFT OpenMM"), + Line2D([0], [0], color=nn_color["openMM"], marker="o", linestyle="-", label="NN OpenMM")] + ax.legend(handles=handles, ncol=2, fontsize=9, frameon=True, loc="upper left") + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +def plot_fitting_accuracy_boxplot(results_df, solvents, nuc_label, engine_colors, save_path=None): + """The semi-parsimonious composite model's test RMSE vs experiment, four boxes per solvent + (Desmond DFT/NN, OpenMM DFT/NN) so DFT-based and MagNET-x-based explicit terms sit side by side.""" + groups = [("desmond", "DFT", "Desmond (DFT)"), ("desmond", "NN", "Desmond (NN)"), + ("openMM", "DFT", "OpenMM (DFT)"), ("openMM", "NN", "OpenMM (NN)")] + colors = [lighten_color(engine_colors["desmond"], 0.15), lighten_color(engine_colors["desmond"], 0.55), + lighten_color(engine_colors["openMM"], 0.15), lighten_color(engine_colors["openMM"], 0.55)] + n = len(groups) + width = 0.8 / n + x_base = np.arange(len(solvents)) + fig, ax = plt.subplots(figsize=(11, 5)) + for gi, (engine, source, label) in enumerate(groups): + data = [results_df[(results_df["solvent"] == sv) & (results_df["engine"] == engine) + & (results_df["source"] == source)]["test_RMSE"].dropna().to_numpy() + for sv in solvents] + positions = x_base + (gi - (n - 1) / 2) * width + bp = ax.boxplot(data, positions=positions, widths=width * 0.9, patch_artist=True, + showfliers=False, manage_ticks=False) + for box in bp["boxes"]: + box.set(facecolor=colors[gi], alpha=0.9) + for median in bp["medians"]: + median.set(color="black") + bp["boxes"][0].set_label(label) + ax.set_xticks(x_base) + ax.set_xticklabels(solvents) + ax.set_ylabel(f"Accuracy vs. Experiment (RMSE, {nuc_label} ppm)") + ax.set_title(f"DFT vs NN Explicit Corrections, Fitting Accuracy ({nuc_label})") + ax.legend(ncol=2, fontsize=9) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +# ---------------------------------------------------------------------------- +# SI Figure S14: test RMSE with DFT vs MagNET features, by solvent +# ---------------------------------------------------------------------------- +def ss_fits(query, nucleus, formulas, solvents, n_splits, dft=False): + """Solvent-specific test RMSEs over the explicit solvents (SI Figure S14), nitromethane dropped.""" + q = query[(query["nucleus"] == nucleus) & (query["solute"] != "nitromethane")] + if dft: + # DFT features need one method per nucleus (the MagNET-Zero training reference: WP04/pcSseg2 + # for 1H, wB97X-D/pcSseg2 for 13C, AIMNet2 geometries) or the fit pools every method together. + # The NN table has only MagNET, so it needs no such filter. + method = delta22.MAGNET_PCM_OUTPUT_METHODS[nucleus] + q = q[(q["sap_nmr_method"] == method) & (q["sap_basis"] == "pcSseg2") + & (q["sap_geometry_type"] == "aimnet2")] + solutes = sorted(q["solute"].unique()) + return delta22.run_fits(q, solvents, formulas, n_splits=n_splits, solutes=solutes) + + +def plot_ss_boxplot_dft_vs_nn(dft_df, nn_df, solvents, formulas, labels, nucleus, + solvent_labels=None, colors=("#A72608", "#5D737E", "#D9FFF5"), + figsize=(14, 8), box_span=0.7, group_gap=0.35, save_path=None): + """Per solvent, box plots of the solvent-specific test-RMSE distributions for each composite + formula, computed with DFT features (solid) and MagNET/NN features (lightened), side by side.""" + def series(df, formula, solvent): + return df[(df["formula"] == formula) & (df["solvent"] == solvent)]["test_RMSE"].dropna().values + + labels_x = [(solvent_labels or {}).get(s, s) for s in solvents] + n_formulas = len(formulas) + n_boxes = 2 * n_formulas # all DFT boxes, then all NN boxes, per solvent + box_w = box_span / n_boxes + centers = [i * (box_span + group_gap) for i in range(len(solvents))] + fig, ax = plt.subplots(figsize=figsize) + handles = [] + # draw the DFT boxes (base colors) first, then the NN boxes (lightened), matching the SI legend + for source_index, (df, tag, lighten) in enumerate([(dft_df, "DFT", False), (nn_df, "NN", True)]): + for fi, (formula, label, color) in enumerate(zip(formulas, labels, colors)): + rgb = plt.cm.colors.to_rgb(color) + face = tuple(min(1.0, c + 0.3) for c in rgb) if lighten else color + slot = source_index * n_formulas + fi + offset = (slot - (n_boxes - 1) / 2) * box_w + data = [series(df, formula, s) for s in solvents] + ax.boxplot(data, positions=[c + offset for c in centers], widths=box_w * 0.9, + patch_artist=True, + boxprops=dict(color="gray", facecolor=face, alpha=0.9), + medianprops=dict(color="black", linewidth=0.5), + whiskerprops=dict(linewidth=0.4), capprops=dict(linewidth=0.4), + flierprops=dict(marker="o", markersize=0)) + handles.append(plt.Rectangle((0, 0), 1, 1, fc=face, ec="gray", alpha=0.9, + label=f"{label} ({tag})")) + ax.set_xticks(centers) + ax.set_xticklabels(labels_x, rotation=45, ha="right", fontweight="bold") + ax.set_ylabel("Test RMSE (ppm)", fontweight="bold") + ax.set_title(f"DFT vs NN Solvent-Specific Test RMSEs\nModel Comparison ({nucleus} nucleus)", + fontweight="bold") + ax.set_ylim(bottom=0) + ax.legend(handles=handles, loc="upper left", fontsize=9) + fig.tight_layout() + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig + + +# ---------------------------------------------------------------------------- +# SI Figures S16-S18: predicted vs experimental solvent-induced shifts by reference solvent +# ---------------------------------------------------------------------------- +def plot_shift_prediction_scatter_grid(diff_df, solvents, reference_label, axis_label, title_label, + n_cols=4, save_path=None): + """Per solvent, the implicit (PCM) and explicit (Desmond) predicted solvent-induced shift + differences (y) against the measured ones (x), with a y=x guide.""" + n = len(solvents) + n_rows = (n + n_cols - 1) // n_cols + fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 5 * n_rows)) + axes = np.atleast_1d(axes).flatten() + series = [("implicit_diff", "#FF6B6B", "D"), ("explicit_diff", "#4ECDC4", "o")] + for ax, solvent in zip(axes, solvents): + d = diff_df[diff_df["solvent"] == solvent] + vals = np.concatenate([d["exp_diff"].to_numpy()] + [d[c].to_numpy() for c, *_ in series]) + lim = float(np.nanmax(np.abs(vals))) * 1.05 if len(vals) else 1.0 + ax.plot([-lim, lim], [-lim, lim], "k--", alpha=0.3, lw=1.5, zorder=0) + for col, color, marker in series: + ax.scatter(d["exp_diff"], d[col], s=40, color=color, marker=marker, zorder=2) + tok = axis_label[solvent] + ax.set_xlabel(f"Experimental delta ({tok} - {reference_label}) [ppm]") + ax.set_ylabel(f"Correction delta ({reference_label} - {tok}) [ppm]") + ax.set_title(title_label[solvent], fontweight="bold") + ax.set_xlim(-lim, lim); ax.set_ylim(-lim, lim); ax.set_aspect("equal") + legend_handles = [ + Line2D([0], [0], marker="D", color="w", markerfacecolor="#FF6B6B", markersize=9, label="Implicit"), + Line2D([0], [0], marker="o", color="w", markerfacecolor="#4ECDC4", markersize=9, label="Explicit"), + Line2D([0], [0], color="gray", linestyle="--", label="Ideal (y = x)")] + empty = list(axes[n:]) + for ax in empty: + ax.set_visible(False) + if empty: # put the legend in the first empty grid slot (S16/S17) + empty[0].set_visible(True); empty[0].axis("off") + empty[0].legend(handles=legend_handles, loc="center", frameon=False, fontsize=12) + fig.tight_layout() + else: # full 12-panel grid (S18): legend centred below the grid + fig.tight_layout(rect=[0, 0.045, 1, 1]) + fig.legend(handles=legend_handles, loc="lower center", ncol=3, frameon=False, fontsize=12) + if save_path: + fig.savefig(save_path, dpi=300, bbox_inches="tight") + return fig diff --git a/analysis/code/shared/fixed_point.py b/analysis/code/shared/fixed_point.py new file mode 100644 index 0000000000000000000000000000000000000000..9775e9908ac3c374565e75381ca8d109bfbba890 --- /dev/null +++ b/analysis/code/shared/fixed_point.py @@ -0,0 +1,21 @@ +"""Decodes the release's fixed-point integer encoding for shieldings and similar fields. + +Values are stored as whole numbers equal to the real value times `scale`, with a reserved marker for +values that were never computed. This turns them back into real values (NaN where missing); float +input passes through untouched, so the same decoder reads either the released int32 encoding or an +older plain-float copy. +""" +import numpy as np + +FIXED_POINT_SCALE = 1e4 +MISSING_MARKER = -2147483648 + + +def decode_fixed_point(values, scale=FIXED_POINT_SCALE, missing_marker=MISSING_MARKER): + """Convert fixed-point integers back to real-valued floats, mapping the missing marker to NaN.""" + values = np.asarray(values) + if np.issubdtype(values.dtype, np.integer): + out = values.astype(np.float64) / scale + out[values == missing_marker] = np.nan + return out + return np.asarray(values, dtype=np.float64) diff --git a/analysis/code/shared/spreadsheet.py b/analysis/code/shared/spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..124ecd99c7537d98a345efed77909f0de2d408e8 --- /dev/null +++ b/analysis/code/shared/spreadsheet.py @@ -0,0 +1,8 @@ +"""Helpers for parsing the experimental spreadsheets' conventions.""" + + +def site_atom_indices(atom_numbers): + """0-based atom indices for a site from a spreadsheet's 1-based comma-separated atom_numbers + (e.g. "17,19" or "21"), the convention used by both the delta-22 and applications experimental + spreadsheets.""" + return [int(x) - 1 for x in str(atom_numbers).split(",")] diff --git a/analysis/code/shared/stats.py b/analysis/code/shared/stats.py new file mode 100644 index 0000000000000000000000000000000000000000..f212cc99469f32c81017f3249fe10ffba7779e08 --- /dev/null +++ b/analysis/code/shared/stats.py @@ -0,0 +1,58 @@ +"""Error statistics and least-squares fitting shared across the analysis modules. + +numpy.linalg.lstsq, not statsmodels: fitting by formula string through statsmodels' patsy parser is +about 1000x slower in a loop that fits many (solvent, solute, formula) combinations, so every fitting +harness here builds a plain design matrix and calls ols_fit directly. Where a module still needs to +report a p-value or similar statsmodels-only diagnostic, keep a separate, explicitly-named +statsmodels code path as a test oracle, not the harness itself. +""" +import numpy as np + + +def rmse(predicted, actual): + predicted = np.asarray(predicted, dtype=float) + actual = np.asarray(actual, dtype=float) + return float(np.sqrt(np.mean(np.square(predicted - actual)))) + + +def mae(predicted, actual): + predicted = np.asarray(predicted, dtype=float) + actual = np.asarray(actual, dtype=float) + return float(np.mean(np.abs(predicted - actual))) + + +def median_ae(predicted, actual): + predicted = np.asarray(predicted, dtype=float) + actual = np.asarray(actual, dtype=float) + return float(np.median(np.abs(predicted - actual))) + + +def summarize_errors(predicted, actual): + """{"n", "rmse", "mae", "median_ae"} for one (predicted, actual) pair.""" + predicted = np.asarray(predicted, dtype=float) + actual = np.asarray(actual, dtype=float) + return { + "n": int(predicted.size), + "rmse": rmse(predicted, actual), + "mae": mae(predicted, actual), + "median_ae": median_ae(predicted, actual), + } + + +def ols_fit(design, response): + """Least-squares coefficient vector for `design @ beta ~ response`. Callers build `design` + (including any intercept column) and drop non-finite rows themselves -- this is a thin, + single-purpose wrapper, not a formula parser.""" + design = np.asarray(design, dtype=float) + response = np.asarray(response, dtype=float) + beta, _residuals, _rank, _sv = np.linalg.lstsq(design, response, rcond=None) + return beta + + +def linear_fit_1d(x, y): + """(intercept, slope) of the least-squares line through (x, y).""" + x = np.asarray(x, dtype=float) + y = np.asarray(y, dtype=float) + design = np.column_stack([np.ones(len(x)), x]) + intercept, slope = ols_fit(design, y) + return float(intercept), float(slope) diff --git a/analysis/code/shared/test_fixed_point.py b/analysis/code/shared/test_fixed_point.py new file mode 100644 index 0000000000000000000000000000000000000000..1241440c9ee67ae76767574b991b4144239e4626 --- /dev/null +++ b/analysis/code/shared/test_fixed_point.py @@ -0,0 +1,42 @@ +import os +import sys + +import numpy as np + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import fixed_point as F # noqa: E402 + + +def test_decode_int32_with_missing_marker(): + values = np.array([10000, -2147483648, 25000], dtype=np.int32) + out = F.decode_fixed_point(values) + assert out[0] == 1.0 + assert np.isnan(out[1]) + assert out[2] == 2.5 + + +def test_decode_int64_also_works(): + # leveling_effect's existing test feeds int64, not just int32 -- the shared decoder must accept + # any integer dtype, not just int32. + values = np.array([10000, -2147483648, 25000], dtype=np.int64) + out = F.decode_fixed_point(values) + assert out[0] == 1.0 + assert np.isnan(out[1]) + assert out[2] == 2.5 + + +def test_float_input_passes_through(): + values = np.array([1.0, 2.5, float("nan")]) + out = F.decode_fixed_point(values) + assert out[0] == 1.0 + assert out[1] == 2.5 + assert np.isnan(out[2]) + + +def test_custom_scale_and_marker(): + values = np.array([100, -1], dtype=np.int32) + out = F.decode_fixed_point(values, scale=100.0, missing_marker=-1) + assert out[0] == 1.0 + assert np.isnan(out[1]) diff --git a/analysis/code/shared/test_spreadsheet.py b/analysis/code/shared/test_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..29c0560c670573c3eb971b43c00b6c24e422d05b --- /dev/null +++ b/analysis/code/shared/test_spreadsheet.py @@ -0,0 +1,19 @@ +import os +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import spreadsheet as S # noqa: E402 + + +def test_site_atom_indices_single(): + assert S.site_atom_indices("21") == [20] + + +def test_site_atom_indices_multiple(): + assert S.site_atom_indices("17,19") == [16, 18] + + +def test_site_atom_indices_accepts_int_input(): + assert S.site_atom_indices(21) == [20] diff --git a/analysis/code/shared/test_stats.py b/analysis/code/shared/test_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..000d8566b807373329aca3d900ce1d6a591b85f2 --- /dev/null +++ b/analysis/code/shared/test_stats.py @@ -0,0 +1,46 @@ +import os +import sys + +import numpy as np +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import stats as S # noqa: E402 + + +def test_rmse_mae_median_ae_hand_computed(): + predicted = [1, 2, 3] + actual = [1, 2, 4] + # errors = [0, 0, -1] + assert S.rmse(predicted, actual) == np.sqrt(1 / 3) + assert S.mae(predicted, actual) == 1 / 3 + assert S.median_ae(predicted, actual) == 0.0 + + +def test_summarize_errors(): + out = S.summarize_errors([1, 2, 3], [1, 2, 4]) + assert out["n"] == 3 + assert out["rmse"] == pytest.approx(np.sqrt(1 / 3)) + assert out["mae"] == pytest.approx(1 / 3) + assert out["median_ae"] == 0.0 + + +def test_linear_fit_1d_exact_recovery_no_noise(): + x = np.linspace(0, 10, 20) + y = 2 * x + 1 + intercept, slope = S.linear_fit_1d(x, y) + assert abs(intercept - 1) < 1e-8 + assert abs(slope - 2) < 1e-8 + + +def test_ols_fit_matches_polyfit(): + rng = np.random.default_rng(0) + x = rng.normal(size=200) + y = 3 * x - 2 + rng.normal(scale=0.01, size=200) + design = np.column_stack([np.ones(len(x)), x]) + beta = S.ols_fit(design, y) + ref_slope, ref_intercept = np.polyfit(x, y, 1) + assert abs(beta[0] - ref_intercept) < 1e-6 + assert abs(beta[1] - ref_slope) < 1e-6 diff --git a/analysis/code/test_cp3.py b/analysis/code/test_cp3.py new file mode 100644 index 0000000000000000000000000000000000000000..6ca7dc56980995305d3de28a915dd57578272d96 --- /dev/null +++ b/analysis/code/test_cp3.py @@ -0,0 +1,39 @@ +"""Tests for analysis/code/cp3.py. The fraction-below-DFT statistic is a quantitative claim in the +paper (36% of proton sites), so it is checked against the shipped Goodman CP3 data, along with the +zone logic on a synthetic array.""" +import os +import sys + +import numpy as np +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import cp3 # noqa: E402 + + +def test_fraction_below_dft_synthetic(): + # 10 values: 2 below the experimental floor, 5 in [0.02, 0.10), 3 at/above 0.10. Equal-width + # bins make the area fraction in [0.02, 0.10] equal to the count fraction there. + x = np.array([0.005, 0.015, 0.03, 0.04, 0.05, 0.06, 0.07, 0.12, 0.2, 0.3]) + frac = cp3.fraction_below_dft(x, bins=100) + assert frac == pytest.approx(0.5, abs=0.02) # 5 of 10 in [0.02, 0.10) + + +def test_zone_boundaries(): + assert cp3._zone(0.01) == "below_experimental" + assert cp3._zone(0.05) == "below_dft" + assert cp3._zone(0.2) == "dft_zone" + assert cp3._zone(0.4) == "large" + + +REAL = os.path.join(HERE, "..", "..", "data", "cp3", "goodman2009_cp3.xlsx") + + +@pytest.mark.skipif(not os.path.exists(REAL), reason="goodman2009_cp3.xlsx not present") +def test_reproduces_published_36_percent(): + variations = cp3.load_variations(REAL, nucleus="H") + assert len(variations) == 168 # the 1H sites in the CP3 set + frac = cp3.fraction_below_dft(variations) + assert frac == pytest.approx(0.36, abs=0.01) # paper: 36% of proton sites below DFT diff --git a/analysis/code/test_magnet_benchmark.py b/analysis/code/test_magnet_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..ae3e6e1e0f2ed3f257f3ebe26377c37dd9569700 --- /dev/null +++ b/analysis/code/test_magnet_benchmark.py @@ -0,0 +1,125 @@ +"""Tests for analysis/code/magnet_benchmark.py. + +The error statistics and the published-table structure are checked synthetically, so they run in CI +without any model or large dataset. Two more tests are opt-in, gated on real data being present +locally: test_shipped_results_reproduce_published_median_and_mae (the sampled, checkpoint-based +reproduction; see magnet_benchmark_run.py) and test_exact_stats_table_reproduces_published_values_ +exactly (the full, unsampled reproduction from data/magnet_test_predictions/, which is what the +shipped notebook actually uses -- see the module docstring's two reproduction paths). +""" +import os +import sys + +import numpy as np +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +import paths + +import magnet_benchmark as M # noqa: E402 + + +def test_collect_abs_errors_filters_by_element_and_nan(): + an = [np.array([1, 6, 1, 8])] + pred = [np.array([10.0, 100.0, 11.0, np.nan])] + dft = [np.array([10.1, 102.0, 10.9, 50.0])] + h = M.collect_abs_errors(pred, dft, an, z=1) + c = M.collect_abs_errors(pred, dft, an, z=6) + assert np.allclose(np.sort(h), [0.1, 0.1]) # the two hydrogens + assert np.allclose(c, [2.0]) # the one carbon + # oxygen has a NaN prediction -> dropped, so no oxygen errors survive + assert M.collect_abs_errors(pred, dft, an, z=8).size == 0 + + +def test_summarize_matches_numpy(): + e = np.array([0.1, 0.2, 0.3, 0.4]) + s = M.summarize(e) + assert s["median_ae"] == pytest.approx(0.25) + assert s["mae"] == pytest.approx(0.25) + assert s["rmse"] == pytest.approx(np.sqrt(np.mean(e ** 2))) + assert s["n"] == 4 + empty = M.summarize(np.array([])) + assert empty["n"] == 0 and np.isnan(empty["mae"]) + + +def test_error_table_both_nuclei(): + an = [np.array([1, 6])] + pred = [np.array([10.0, 100.0])] + dft = [np.array([10.5, 98.0])] + t = M.error_table(pred, dft, an) + assert t["1H"]["mae"] == pytest.approx(0.5) + assert t["13C"]["mae"] == pytest.approx(2.0) + + +def test_filter_supported_drops_out_of_vocabulary(): + # second structure contains phosphorus (15), which MagNET was never trained on + ans = [np.array([6, 1, 1]), np.array([6, 1, 15]), np.array([8, 1, 17])] + geos = [np.zeros((3, 3))] * 3 + dfts = [np.zeros(3)] * 3 + a, g, d, dropped = M.filter_supported(ans, geos, dfts) + assert dropped == 1 + assert len(a) == len(g) == len(d) == 2 + assert all(set(np.unique(x).tolist()) <= M.SUPPORTED_ELEMENTS for x in a) + + +RESULTS = os.path.join(HERE, "..", "..", "data", "magnet_benchmark", "performance_results.csv") + + +@pytest.mark.skipif(not os.path.exists(RESULTS), reason="performance_results.csv not present") +def test_shipped_results_reproduce_published_median_and_mae(): + """The shipped reproduced numbers match the published median and mean absolute error (the robust + statistics). RMSE is intentionally not checked: it is outlier-dominated for the vibrated sets and + needs the full test set (see the module docstring).""" + import csv + # tolerance per nucleus: 1H errors are ~0.03 ppm, 13C ~0.4 ppm, so allow a sampling margin + TOL = {"1H": 0.015, "13C": 0.12} + rows = list(csv.DictReader(open(RESULTS))) + assert len(rows) == 2 * len(M.MODELS) * len(M.TEST_SETS) # 16 rows + for r in rows: + pub_median, pub_mae, _ = M.PUBLISHED[r["nucleus"]][(r["model"], r["test_set"])] + tol = TOL[r["nucleus"]] + tag = f"{r['model']}/{r['test_set']}/{r['nucleus']}" + assert abs(float(r["median_ae"]) - pub_median) < tol, f"{tag} median {r['median_ae']} vs {pub_median}" + assert abs(float(r["mae"]) - pub_mae) < tol, f"{tag} mae {r['mae']} vs {pub_mae}" + + +def test_predictions_group_matches_documented_naming_scheme(): + # spot-check against the exact group names defined in magnet_benchmark.test_predictions_group + assert M.test_predictions_group("MagNET", "vibrated_external", "13C") == "gasphasedft8k_C_pretrained_C_vib1" + assert M.test_predictions_group("MagNET", "isolated_chloroform", "1H") == "solutesmd500isolated_chloroform_H_pretrained_H" + assert M.test_predictions_group("MagNET-x", "stationary_internal", "1H") == "gasphaseinternal_H_chloroform_H_vib0" + with pytest.raises(ValueError): + M.test_predictions_group("MagNET", "not_a_real_test_set", "1H") + + +PREDICTIONS_H5 = paths.dataset_file("magnet_test_predictions", file=__file__) + + +@pytest.mark.skipif(not os.path.exists(PREDICTIONS_H5), reason="magnet_test_predictions.hdf5 not present") +def test_exact_stats_table_reproduces_published_values_exactly(): + """Unlike the sampled test above, this reads the FULL released test sets (no sampling), so median, + MAE, and RMSE should all match the published SI values to float rounding, not just a sampling + tolerance -- this is what the shipped si_table_s03_s04_performance.ipynb notebook actually asserts.""" + sys.path.insert(0, os.path.join(HERE, "..", "..", "data", "magnet_test_predictions")) + import magnet_test_predictions_reader as R + rows = M.exact_stats_table(PREDICTIONS_H5, R) + assert len(rows) == 2 * len(M.MODELS) * len(M.TEST_SETS) + for r in rows: + pub_median, pub_mae, pub_rmse = M.PUBLISHED[r["nucleus"]][(r["model"], r["test_set"])] + tag = f"{r['model']}/{r['test_set']}/{r['nucleus']}" + assert r["median_ae"] == pytest.approx(pub_median, abs=1e-3), tag + assert r["mae"] == pytest.approx(pub_mae, abs=1e-3), tag + assert r["rmse"] == pytest.approx(pub_rmse, abs=1e-3), tag + + +def test_published_table_structure(): + # both nuclei, both models, all four test sets, triples of finite numbers + for nucleus in ("1H", "13C"): + for model in M.MODELS: + for ts in M.TEST_SETS: + triple = M.PUBLISHED[nucleus][(model, ts)] + assert len(triple) == 3 + assert all(np.isfinite(v) and v > 0 for v in triple) + # the 13C vibrated-internal RMSE is the documented outlier-dominated value + assert M.PUBLISHED["13C"][("MagNET", "vibrated_internal")][2] == pytest.approx(7.1377, abs=1e-3) diff --git a/analysis/code/test_paths.py b/analysis/code/test_paths.py new file mode 100644 index 0000000000000000000000000000000000000000..a8f204d6b83acdacbf2749e7390b0f57fa3d08b7 --- /dev/null +++ b/analysis/code/test_paths.py @@ -0,0 +1,94 @@ +import os +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import paths as P # noqa: E402 + + +def test_repo_root_resolves_this_file_two_levels_up(): + # requirements.txt lives only at the repo root, so its presence is what proves repo_root landed + # in the right place; the checkout folder's name is not part of the contract (a CI checkout, for + # example, names it after the GitHub repo, not this developer's local folder name). + root = P.repo_root(__file__) + assert os.path.isfile(os.path.join(root, "requirements.txt")) + assert os.path.basename(root) != "code" # sanity: didn't stop one level too early + + +def test_repo_root_from_synthetic_tree(tmp_path): + fake_module = tmp_path / "analysis" / "code" / "x.py" + fake_module.parent.mkdir(parents=True) + fake_module.write_text("") + assert P.repo_root(str(fake_module)) == str(tmp_path) + + +def test_ensure_on_path_inserts_at_front_and_dedupes(): + # sys.path is a single shared, mutable, process-global list -- other test modules in the same + # pytest session insert their own entries into the REAL sys.path (some also at index 0), so + # asserting against the real list is flaky no matter how carefully this test brackets its own + # calls. Swap in a private list for the duration of the test instead, so ensure_on_path's + # "insert at index 0, don't duplicate" contract can be checked hermetically. + real_sys_path = sys.path + sys.path = ["/some/other/existing/entry"] + try: + target = os.path.join(P.repo_root(__file__), "data", "delta22") + P.ensure_on_path("data", "delta22", file=__file__) + assert sys.path[0] == target # inserts at the FRONT, so local modules can shadow the rest + assert sys.path.count(target) == 1 + # inserting the same target again must not duplicate it + P.ensure_on_path("data", "delta22", file=__file__) + assert sys.path.count(target) == 1 + assert sys.path[0] == target + finally: + sys.path = real_sys_path + + +def test_ensure_on_path_needs_file_or_root(): + try: + P.ensure_on_path("data") + assert False, "expected ValueError" + except ValueError: + pass + + +def test_dataset_file_resolves_in_repo(tmp_path): + got = P.dataset_file("delta22", root=str(tmp_path)) + assert got == os.path.join(str(tmp_path), "data", "delta22", "delta22.hdf5") + + +def test_dataset_file_derives_root_from_file(tmp_path): + # the form most rewired modules use: dataset_file("", file=__file__). A module at + # analysis/code/x.py must resolve the file two directories up, at /data/... + fake = tmp_path / "analysis" / "code" / "x.py" + fake.parent.mkdir(parents=True) + fake.write_text("") + got = P.dataset_file("delta22", file=str(fake)) + assert got == os.path.join(str(tmp_path), "data", "delta22", "delta22.hdf5") + + +def test_dataset_file_custom_filename_and_needs_a_root(): + got = P.dataset_file("applications", "applications_md_geometries.hdf5", root="/r") + assert got == os.path.join("/r", "data", "applications", "applications_md_geometries.hdf5") + try: + P.dataset_file("delta22") # no root, no file + assert False, "expected ValueError" + except ValueError: + pass + + +def test_checkpoints_root(monkeypatch, tmp_path): + # point the resolver at an empty tree, so its model_checkpoints/ is absent (as in CI) + empty = tmp_path / "repo" + empty.mkdir() + monkeypatch.setattr(P, "repo_root", lambda _file: str(empty)) + assert P.checkpoints_root() is None + try: + P.checkpoints_root(required=True) + assert False, "expected RuntimeError" + except RuntimeError: + pass + # in-repo location: model_checkpoints/ present at the repo root + (empty / "model_checkpoints").mkdir() + assert P.checkpoints_root() == os.path.join(str(empty), "model_checkpoints") + assert P.checkpoints_root(required=True) == os.path.join(str(empty), "model_checkpoints") diff --git a/analysis/code/test_scaling_factors.py b/analysis/code/test_scaling_factors.py new file mode 100644 index 0000000000000000000000000000000000000000..3d7bdd5ae24158be59c444bcf47738ba155dd7c7 --- /dev/null +++ b/analysis/code/test_scaling_factors.py @@ -0,0 +1,272 @@ +"""Tests for analysis/code/scaling_factors.py. + +Synthetic tests exercise the two table builders and the prediction equation without any large file, +so the core math runs in CI. The opt-in real-data test reproduces the published SI numbers (Tables +S10 and S11) from the released delta-22 data when it is present. +""" +import os +import sys + +import numpy as np +import pandas as pd +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import scaling_factors as S # noqa: E402 +import paths as P # noqa: E402 + +REPO = os.path.abspath(os.path.join(HERE, "..", "..")) +REAL_H5 = P.dataset_file("delta22", root=REPO) +REAL_XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx") + + +# --------------------------------------------------------------------------- synthetic builders + +def _proton_nn(coeffs, n=6, seed=0): + """Synthetic proton table where experimental = a + b*stationary + c*pcm exactly, per solvent. + coeffs maps solvent -> (a, b, c).""" + rng = np.random.default_rng(seed) + rows = [] + for solvent, (a, b, c) in coeffs.items(): + for i in range(n): + stat = float(rng.normal(30, 3)) + pcm = float(rng.normal(0, 0.5)) + rows.append(dict(nucleus="H", solute=f"m{i}", site=f"m{i}_0", solvent=solvent, + stationary=stat, pcm=pcm, experimental=a + b * stat + c * pcm)) + return pd.DataFrame(rows) + + +def _carbon_nn(coeffs, factors, n=6, seed=1): + """Synthetic carbon table where experimental = a + b*(stationary + factor*pcm) exactly, per + solvent. coeffs maps solvent -> (a, b); factors maps solvent -> conversion factor.""" + rng = np.random.default_rng(seed) + rows = [] + for solvent, (a, b) in coeffs.items(): + factor = factors[solvent] + for i in range(n): + stat = float(rng.normal(100, 20)) + pcm = float(rng.normal(0, 0.5)) + rows.append(dict(nucleus="C", solute=f"m{i}", site=f"m{i}_0", solvent=solvent, + stationary=stat, pcm=pcm, experimental=a + b * (stat + factor * pcm))) + return pd.DataFrame(rows) + + +def _nitromethane_rows(nucleus, solvents): + """One garbage nitromethane row per solvent; the fits must exclude it.""" + return pd.DataFrame([dict(nucleus=nucleus, solute="nitromethane", site="x", solvent=s, + stationary=0.0, pcm=0.0, experimental=999.0) for s in solvents]) + + +# --------------------------------------------------------------------------- synthetic tests + +def test_recommended_model_forms(): + # Proton is the three-parameter model, carbon the two-parameter model (John's fitting trials). + assert S.RECOMMENDED_MODEL == {"H": "three_parameter", "C": "two_parameter"} + + +def test_predict_shift_equation(): + table = pd.DataFrame( + {"intercept": [31.0], "stationary": [-0.95], "pcm": [-0.85]}, + index=pd.Index(["chloroform"], name="solvent"), + ) + got = S.predict_shift(table, "chloroform", magnet_zero_shielding=25.0, + magnet_pcm_chloroform_correction=-0.1) + assert got == pytest.approx(31.0 - 0.95 * 25.0 - 0.85 * -0.1) + got_vec = S.predict_shift(table, "chloroform", [25.0, 26.0], [-0.1, 0.2]) + assert np.allclose(got_vec, [31.0 - 0.95 * 25.0 - 0.85 * -0.1, + 31.0 - 0.95 * 26.0 - 0.85 * 0.2]) + + +def test_proton_table_recovers_known_line_and_drops_nitromethane(): + coeffs = {"chloroform": (31.0, -0.97, -0.85), "benzene": (32.0, -1.00, 2.20)} + nn = pd.concat([_proton_nn(coeffs), _nitromethane_rows("H", coeffs)], ignore_index=True) + table = S.proton_scaling_table(nn, solvents=list(coeffs)) + assert list(table.columns) == ["intercept", "stationary", "pcm"] + assert table.index.name == "solvent" + for solvent, (a, b, c) in coeffs.items(): + row = table.loc[solvent] + # exact recovery (and the nitromethane garbage row did not perturb it -> it was excluded) + assert row["intercept"] == pytest.approx(a, abs=1e-6) + assert row["stationary"] == pytest.approx(b, abs=1e-6) + assert row["pcm"] == pytest.approx(c, abs=1e-6) + + +def test_carbon_table_reconstruction_and_drops_nitromethane(): + coeffs = {"chloroform": (171.0, -0.92), "benzene": (172.0, -0.93)} + factors = {"chloroform": 1.0, "benzene": 0.63} + nn = pd.concat([_carbon_nn(coeffs, factors), _nitromethane_rows("C", coeffs)], ignore_index=True) + table = S.carbon_scaling_table(nn, solvents=list(coeffs), conversion_factors=factors) + assert list(table.columns) == ["intercept", "stationary", "pcm"] + for solvent, (a, b) in coeffs.items(): + row = table.loc[solvent] + assert row["intercept"] == pytest.approx(a, abs=1e-6) + assert row["stationary"] == pytest.approx(b, abs=1e-6) + # the reported pcm is the shared slope times the conversion factor + assert row["pcm"] == pytest.approx(b * factors[solvent], abs=1e-6) + assert row["pcm"] == pytest.approx(row["stationary"] * factors[solvent], abs=1e-12) + + +def test_carbon_empty_or_nan_factor_gives_nan_row(): + # chloroform has data and a finite factor; benzene has no rows and a NaN factor. + nn = _carbon_nn({"chloroform": (171.0, -0.92)}, {"chloroform": 1.0}) + table = S.carbon_scaling_table(nn, solvents=["chloroform", "benzene"], + conversion_factors={"chloroform": 1.0, "benzene": np.nan}) + assert np.isfinite(table.loc["chloroform", "intercept"]) + # a degenerate solvent must be NaN, not a fake (0, 0, 0) fit + assert table.loc["benzene"].isna().all() + + +def test_carbon_requires_factors_or_dft(): + with pytest.raises(ValueError): + S.carbon_scaling_table(_carbon_nn({"chloroform": (171.0, -0.92)}, {"chloroform": 1.0}), + solvents=["chloroform"]) + + +# --------------------------------------------------------------------------- opt-in real-data test + +# Published SI values: parameter tuples are (intercept, stationary, pcm) per solvent, all 12 solvents. +# These also pin the shipped published_scaling_tables() copy (no-data test below): shipped CSV == +# this dict == build_scaling_tables() from delta-22 (real-data test). The chain keeps all three in sync. +PUBLISHED_S10_H = { + "chloroform": (31.294997, -0.9757947, -0.8526936), + "tetrahydrofuran": (31.321805, -0.9801755, -0.786672), + "dichloromethane": (31.3773285, -0.979871, -0.8085722), + "acetone": (31.512167, -0.9872612, -1.2383371), + "acetonitrile": (31.4961183, -0.9857168, -0.9744366), + "dimethylsulfoxide": (31.5987046, -0.9911032, -1.3551614), + "trifluoroethanol": (30.8759924, -0.9599975, -0.9589483), + "methanol": (31.2560285, -0.9764463, -1.2500843), + "TIP4P": (31.3591125, -0.978808, -1.4782606), + "benzene": (31.9876742, -1.0052929, 2.23638523), + "toluene": (31.7170517, -0.9967364, 1.94472899), + "chlorobenzene": (31.6814988, -0.9938483, 0.83014615), +} +PUBLISHED_S11_C = { + "chloroform": (171.728792, -0.9242313, -0.9368043), + "tetrahydrofuran": (171.054483, -0.9190001, -1.069509), + "dichloromethane": (171.509383, -0.9211671, -1.1150154), + "acetone": (171.308017, -0.9196101, -1.2395026), + "acetonitrile": (171.879832, -0.9226291, -1.2876643), + "dimethylsulfoxide": (170.598418, -0.9186745, -1.2965007), + "trifluoroethanol": (174.237665, -0.9390125, -1.2900728), + "methanol": (172.426154, -0.9275656, -1.2888469), + "TIP4P": (173.696364, -0.937854, -1.3426678), + "benzene": (171.967174, -0.9270331, -0.5937164), + "toluene": (171.690904, -0.9249871, -0.6184099), + "chlorobenzene": (171.075905, -0.921238, -0.997641), +} + + +def test_shipped_published_tables_match_si_values(): + """The shipped published_scaling_tables() copy (no data download needed) equals the published SI + values for all 12 solvents and both nuclei. Runs in CI without delta-22; the real-data test below + ties those same SI values back to a fit on the raw data, so the shipped copy cannot drift.""" + tables = S.published_scaling_tables() + for nucleus, published in (("H", PUBLISHED_S10_H), ("C", PUBLISHED_S11_C)): + table = tables[nucleus] + assert list(table.columns) == ["intercept", "stationary", "pcm"] + assert sorted(table.index) == sorted(published) + for solvent, (intercept, stationary, pcm) in published.items(): + row = table.loc[solvent] + assert row["intercept"] == pytest.approx(intercept, abs=5e-4), f"{nucleus} {solvent} int" + assert row["stationary"] == pytest.approx(stationary, abs=5e-4), f"{nucleus} {solvent} stat" + assert row["pcm"] == pytest.approx(pcm, abs=5e-4), f"{nucleus} {solvent} pcm" + + +@pytest.mark.skipif(not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX)), + reason="real delta22.hdf5 / experimental xlsx not present") +def test_reproduces_published_si_tables(): + """Both recommended-scaling tables reproduce the published SI numbers for all 12 solvents. The + tolerance is set above the int32-encoding floor (worst case ~1.3e-4, the benzene proton pcm) but + tight enough to catch a real regression.""" + tables = S.build_scaling_tables(REAL_H5, REAL_XLSX) + for nucleus, published in (("H", PUBLISHED_S10_H), ("C", PUBLISHED_S11_C)): + table = tables[nucleus] + assert list(table.columns) == ["intercept", "stationary", "pcm"] + assert sorted(table.index) == sorted(published) + for solvent, (intercept, stationary, pcm) in published.items(): + row = table.loc[solvent] + assert row["intercept"] == pytest.approx(intercept, abs=5e-4), f"{nucleus} {solvent} int" + assert row["stationary"] == pytest.approx(stationary, abs=5e-4), f"{nucleus} {solvent} stat" + assert row["pcm"] == pytest.approx(pcm, abs=5e-4), f"{nucleus} {solvent} pcm" + + +# --------------------------------------------------------------------------- symmetrized-inference smoke test + +CKPT_ROOT = P.checkpoints_root() or "" # "" (deposit env unset) -> os.path.exists(...) below is False +CKPT_ZERO = os.path.join(CKPT_ROOT, "MagNET-Zero") +CKPT_PCM = os.path.join(CKPT_ROOT, "MagNET-PCM") +_HAS_CHECKPOINTS = os.path.exists(CKPT_ZERO) and os.path.exists(CKPT_PCM) + + +@pytest.mark.skipif(not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX) and _HAS_CHECKPOINTS), + reason="real delta22.hdf5 / experimental xlsx / model checkpoints not present") +def test_symmetrized_build_scaling_tables_runs_and_is_close_to_published(): + """symmetrized=True re-derives the tables from live, reflection-symmetrized inference instead + of the HDF5's stored (unsymmetrized) shieldings. This is a deployment-quality table (e.g. for + serving MagNET-Zero/PCM on new molecules), not a replacement for the published SI numbers -- + it should be close (delta-22 solutes are small, so the correction is modest) but need not match + to the same tight tolerance as test_reproduces_published_si_tables.""" + pytest.importorskip("torch") + tables = S.build_scaling_tables(REAL_H5, REAL_XLSX, symmetrized=True, n_passes=2) + max_abs_delta = 0.0 + for nucleus, published in (("H", PUBLISHED_S10_H), ("C", PUBLISHED_S11_C)): + table = tables[nucleus] + assert list(table.columns) == ["intercept", "stationary", "pcm"] + assert sorted(table.index) == sorted(published) + for solvent, (intercept, stationary, pcm) in published.items(): + row = table.loc[solvent] + assert np.isfinite(row["intercept"]) and np.isfinite(row["stationary"]) and np.isfinite(row["pcm"]) + # loose tolerance: this is a methodology check (does it run and land in the right + # ballpark), not a bit-for-bit reproduction -- see module docstring for measured deltas + assert row["intercept"] == pytest.approx(intercept, abs=0.05), f"{nucleus} {solvent} int" + assert row["stationary"] == pytest.approx(stationary, abs=0.01), f"{nucleus} {solvent} stat" + assert row["pcm"] == pytest.approx(pcm, abs=0.1), f"{nucleus} {solvent} pcm" + max_abs_delta = max(max_abs_delta, abs(row["intercept"] - intercept), + abs(row["stationary"] - stationary), abs(row["pcm"] - pcm)) + # symmetrized inference must actually change something -- if nn_shieldings_override_df were + # silently dropped and this fell back to the exact HDF5 path, every delta above would be 0.0 + # and still pass the loose tolerances; this catches that failure mode specifically + assert max_abs_delta > 1e-4, "symmetrized=True produced numbers identical to published -- the override path did not run" + + +@pytest.mark.skipif(not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX) and _HAS_CHECKPOINTS), + reason="real delta22.hdf5 / experimental xlsx / model checkpoints not present") +def test_shipped_symmetrized_csvs_reproduce_from_live_inference(): + """Value-traceability for the deployment CSVs in data/scaling_factors/: they must reproduce from + a live symmetrized build at the n_passes they were generated with (10), not merely parse. A + swapped column, wrong solvent order, or stale hand-edit of the shipped CSVs fails here. The + tolerance absorbs the model's pass-to-pass inference noise (a single forward pass is not + deterministic) but is far tighter than any structural error.""" + pytest.importorskip("torch") + sys.path.insert(0, os.path.join(REPO, "data", "scaling_factors")) + import scaling_factors_reader as R # noqa: E402 + shipped = R.load_symmetrized_tables() + fresh = S.build_scaling_tables(REAL_H5, REAL_XLSX, symmetrized=True, n_passes=10) + for nucleus in ("H", "C"): + s, f = shipped[nucleus], fresh[nucleus] + assert list(s.columns) == list(f.columns) == ["intercept", "stationary", "pcm"] + assert sorted(s.index) == sorted(f.index) + for solvent in s.index: + assert s.loc[solvent, "intercept"] == pytest.approx(f.loc[solvent, "intercept"], abs=0.03), f"{nucleus} {solvent} intercept" + assert s.loc[solvent, "stationary"] == pytest.approx(f.loc[solvent, "stationary"], abs=0.01), f"{nucleus} {solvent} stationary" + assert s.loc[solvent, "pcm"] == pytest.approx(f.loc[solvent, "pcm"], abs=0.02), f"{nucleus} {solvent} pcm" + + +def test_magnet_package_scaling_matches_canonical_tables(): + """The tiny copy shipped in the importable package (magnet.scaling, used internally by + magnet.predict_shifts) must equal the canonical SI tables here, so the two can never drift.""" + pytest.importorskip("torch") # importing the magnet package pulls in the torch stack + from magnet import scaling as MS + canonical = S.published_scaling_tables() # DataFrames indexed by solvent + tiny = MS.published_scaling_tables() # {solvent: {column: value}} dicts + for nucleus in ("H", "C"): + assert set(tiny[nucleus]) == set(canonical[nucleus].index) + for solvent, coeffs in tiny[nucleus].items(): + for column in ("intercept", "stationary", "pcm"): + assert coeffs[column] == pytest.approx(float(canonical[nucleus].loc[solvent, column])) + # the two predict_shift implementations give the same shift + assert float(MS.predict_shift(tiny["C"], "chloroform", 170.0, -0.3)) == pytest.approx( + float(S.predict_shift(canonical["C"], "chloroform", 170.0, -0.3))) diff --git a/analysis/manuscript_figures/fig1a_cp3.ipynb b/analysis/manuscript_figures/fig1a_cp3.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..0fee154d01f134191fe955dc9ad5a4e65403fd6b --- /dev/null +++ b/analysis/manuscript_figures/fig1a_cp3.ipynb @@ -0,0 +1,143 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "01850ff6", + "metadata": {}, + "source": [ + "# Figure 1A: CP3 per-site ¹H shift spread vs the DFT and experimental resolution limits\n", + "\n", + "Histogram of per-site ¹H shift spread across the Goodman CP3 stereoisomers (Goodman, JOC 2009, 74,\n", + "4597), against the experimental noise floor (0.02 ppm) and DFT's resolving power (0.10 ppm)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a460d80b", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"analysis/code\",):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e0adf094", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import seaborn as sns\n", + "from matplotlib.ticker import FixedLocator\n", + "\n", + "import cp3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b68e0918", + "metadata": {}, + "outputs": [], + "source": [ + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f599e59b", + "metadata": {}, + "outputs": [], + "source": [ + "variations = cp3.load_variations(nucleus=\"H\")\n", + "frac = cp3.fraction_below_dft(variations)\n", + "print(f\"{len(variations)} proton sites; {frac:.1%} fall between the experimental floor \"\n", + " f\"({cp3.EXPERIMENTAL_LIMIT} ppm) and the DFT limit ({cp3.DFT_LIMIT} ppm)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c88605cd", + "metadata": {}, + "outputs": [], + "source": [ + "# color per zone; keys match cp3._zone's return values\n", + "zone_colors = {\n", + " \"below_experimental\": \"#2c4553\", # < 0.02 ppm: within experimental noise\n", + " \"below_dft\": \"#a72608\", # 0.02 - 0.10 ppm: informative but DFT cannot resolve\n", + " \"dft_zone\": \"#dacd82\", # 0.10 - 0.30 ppm: resolvable by DFT\n", + " \"large\": \"#68a79a\", # >= 0.30 ppm\n", + "}\n", + "boundaries = (cp3.EXPERIMENTAL_LIMIT, cp3.DFT_LIMIT, cp3.LARGE_VARIATION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "03efe205", + "metadata": {}, + "outputs": [], + "source": [ + "sns.set_theme(context=\"paper\", style=\"white\")\n", + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "\n", + "bins = 30\n", + "edges = np.linspace(0.0, float(variations.max()), bins + 1)\n", + "counts, edge = np.histogram(variations, bins=edges)\n", + "\n", + "# Draw each bar, splitting any bin that straddles a zone boundary so each segment gets its color.\n", + "for height, left, right in zip(counts, edge[:-1], edge[1:]):\n", + " splits = sorted({left, right} | {b for b in boundaries if left < b < right})\n", + " for a, b in zip(splits[:-1], splits[1:]):\n", + " ax.bar(a, height, width=b - a, align=\"edge\",\n", + " color=zone_colors[cp3._zone(0.5 * (a + b))],\n", + " edgecolor=\"white\", linewidth=0.6, alpha=0.9)\n", + "\n", + "ax.set_xlim(left=0.0)\n", + "ax.set_ylim(0, ax.get_ylim()[1])\n", + "ax.margins(x=0)\n", + "ax.set_title(\"Chemical shift variation (¹H)\", pad=8, fontsize=16)\n", + "ax.set_xlabel(\"Variation at Site (standard deviation, ¹H ppm)\", fontsize=13)\n", + "ax.set_ylabel(\"Count\", fontsize=13)\n", + "sns.despine(ax=ax, top=True, right=True)\n", + "for side in (\"left\", \"bottom\"):\n", + " ax.spines[side].set_linewidth(0.6)\n", + "ax.xaxis.set_major_locator(FixedLocator([0.0, 0.1, 0.2, 0.3, 0.4, 0.5]))\n", + "ax.tick_params(axis=\"both\", which=\"major\", labelsize=10, size=1.5, width=0.6)\n", + "ax.minorticks_off()\n", + "fig.tight_layout()\n", + "\n", + "fig.savefig(figure_path(\"fig1a_cp3_1H.png\"), dpi=300, bbox_inches=\"tight\", pad_inches=0.02)\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig2a_pareto.ipynb b/analysis/manuscript_figures/fig2a_pareto.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..1b90dde63577bba54e638109e52839d16dbf90ab --- /dev/null +++ b/analysis/manuscript_figures/fig2a_pareto.ipynb @@ -0,0 +1,82 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "55d51b05", + "metadata": {}, + "source": "# Figure 2A: accuracy-vs-compute-cost Pareto frontier (delta-22 ¹H, CDCl3)\n\nFitting RMSE vs total compute time (geometry + shielding), colored by method type, shaped by geometry source, sized by basis." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a5fb8cc9", + "metadata": {}, + "outputs": [], + "source": "import os, sys\n\n# make the in-repo modules importable (not pip-installed)\nREPO = os.path.abspath(\"../..\")\nfor _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n sys.path.insert(0, os.path.join(REPO, _p))" + }, + { + "cell_type": "code", + "id": "24c54da5", + "source": "import delta22\nimport pareto_plot\nimport paths", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "c2b77287", + "source": "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\nXLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n\ndef figure_path(name):\n os.makedirs(\"figures\", exist_ok=True)\n return os.path.join(\"figures\", name)", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c8daed8f", + "metadata": {}, + "outputs": [], + "source": "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\nnn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)\ndft_gas_timings, nn_timings = delta22.load_pareto_timings(DELTA22_HDF5)\n\n# one point per method/basis/geometry/nucleus/solvent (plus solvent-averaged rows); MagNET is a single\n# method (aimnet2, basis \"N/A\"). 100 seeded splits (train on the first 10 shuffled solutes, test on the\n# rest) to match the published panel.\npoints = delta22.fig2a_pareto_points(dft, nn, dft_gas_timings, nn_timings, n_splits=100)\nprint(points[\"nmr_method\"].nunique(), \"methods;\",\n \"MagNET total time\", float(points.query(\"nmr_method=='MagNET'\")[\"total_time\"].iloc[0]), \"s\")" + }, + { + "cell_type": "markdown", + "id": "556ad0b0", + "metadata": {}, + "source": [ + "## Proton Pareto panel" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c9a93fd", + "metadata": {}, + "outputs": [], + "source": "method_legend_label_map = {\n \"MagNET-Zero\": \"MagNET-Zero\", \"ab initio\": r\"$\\it{ab\\ initio}$\", \"double hybrids\": \"double hybrid\",\n \"NMR-specific\": \"NMR-specific\", \"DFT\": \"DFT\",\n}\ngeometry_legend_label_map = {\"pbe0_tz\": \"DFT (PBE0/cc-pVTZ)\", \"aimnet2\": \"ML (AIMNet2)\"}\nmethod_color_map = {\"ab initio\": \"#53585f\", \"MagNET-Zero\": \"#000000\", \"double hybrids\": \"#b3a560\",\n \"NMR-specific\": \"#a72608\", \"DFT\": \"#4a8075\"}\n\n# which points to draw: MagNET-Zero, all PBE0-geometry DFT methods, plus the wp04 NMR-specific point\nquery_str = (\"nucleus=='{nucleus}' and solvent == 'chloroform' and (\"\n \" (geometry_type=='aimnet2' and nmr_method.str.startswith('MagNET')) or\"\n \" (geometry_type=='pbe0_tz' and not nmr_method.str.startswith('MagNET')) or\"\n \" (geometry_type=='aimnet2' and nmr_method=='wp04' and basis=='pcSseg2' and nucleus=='H'))\")\n\n# hand-tuned label offsets (points, matching the published panel)\nmanual_label_positions = {\n (\"MagNET-Zero\", \"N/A\", \"aimnet2\"): {\"xytext\": (-28, 14)},\n (\"wp04\", \"pcSseg2\", \"aimnet2\"): {\"xytext\": (6, 4)},\n (\"pbe0\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (-22, -2)},\n (\"b2gp_plyp\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (4, 6)},\n (\"revdsd_pbep86\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (6, -4)},\n (\"revdsd_pbep86\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, 6)},\n (\"tpsstpss\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, -2)},\n (\"b97d3\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (-28, -2)},\n (\"mpw2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 2)},\n (\"b2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 0)},\n (\"tpsstpss\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (6, -12)},\n}" + }, + { + "cell_type": "code", + "id": "8d87338a", + "source": "pareto_plot.plot_pareto_panel(\n points, query_str, nucleus=\"H\",\n method_color_map=method_color_map,\n method_legend_label_map=method_legend_label_map,\n geometry_legend_label_map=geometry_legend_label_map,\n geometry_marker_map={\"aimnet2\": \"x\", \"pbe0_tz\": \"o\"},\n xlim_left=(1.7, 1.8), xlim_right=(3.8, 5.5), ylim=(0.08, 0.3),\n figsize=(12, 6), marker_alpha=0.8,\n manual_label_positions=manual_label_positions,\n save_png=figure_path(\"fig2a_pareto_1H.png\"),\n)", + "metadata": {}, + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig2b_correlation.ipynb b/analysis/manuscript_figures/fig2b_correlation.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..9bade551712e4a9727135aa234f3d75f02336ec7 --- /dev/null +++ b/analysis/manuscript_figures/fig2b_correlation.ipynb @@ -0,0 +1,111 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e8c73ac7", + "metadata": {}, + "source": [ + "# Figure 2B: inter-method correlation of delta-22 ¹H stationary shieldings\n", + "\n", + "Lower-triangle Pearson correlation matrix of per-site ¹H delta-22 stationary shieldings across methods\n", + "(pcSseg-2 basis). The colorbar is Pearson r; the published panel mislabels it r², a typo, corrected here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ae047e23", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d783476", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "import delta22\n", + "import paths\n", + "import fig2b_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9f91feb0", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "569d9eef", + "metadata": {}, + "outputs": [], + "source": [ + "# Gas-phase (\"stationary\") shieldings, one value per (solute, geometry, basis, site, nucleus, method).\n", + "# The stationary shielding does not depend on solvent, so keep one solvent (TIP4P) to deduplicate.\n", + "NUCLEUS = \"H\"\n", + "BASIS = \"pcSseg2\"\n", + "\n", + "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "corr_df = dft[[\"solute\", \"sap_geometry_type\", \"sap_nmr_method\", \"sap_basis\",\n", + " \"nucleus\", \"site\", \"solvent\", \"stationary\"]].copy()\n", + "corr_df.columns = [c.replace(\"sap_\", \"\") for c in corr_df.columns]\n", + "corr_df = corr_df.query('solvent == \"TIP4P\"').drop(columns=[\"solvent\"])\n", + "\n", + "pivot = corr_df.pivot(index=[\"solute\", \"geometry_type\", \"basis\", \"site\", \"nucleus\"],\n", + " columns=\"nmr_method\", values=\"stationary\")\n", + "sub = pivot.query(\"nucleus == @NUCLEUS and basis == @BASIS\")\n", + "correlations = sub.corr()\n", + "off_diag = correlations.where(~np.eye(len(correlations), dtype=bool))\n", + "print(f\"{correlations.shape[0]} methods; worst pairwise Pearson r = {off_diag.min().min():.5f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f866aab1", + "metadata": {}, + "outputs": [], + "source": [ + "fig2b_plots.plot_correlation_matrix(correlations, save_png=figure_path(\"fig2b_correlation_1H.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/manuscript_figures/fig2d_convergence.ipynb b/analysis/manuscript_figures/fig2d_convergence.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..daca4abd20c4b74150bc0a135ed7d1db4c9a9ccb --- /dev/null +++ b/analysis/manuscript_figures/fig2d_convergence.ipynb @@ -0,0 +1,148 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8e86520a", + "metadata": {}, + "source": [ + "# Figure 2D: per-method scaled error and slope vs a CCSD(T) reference (NS372)\n", + "\n", + "Each conventional method placed by **scaled error** (RMSE after rescaling into CCSD(T) units via the\n", + "per-method fit) and **slope** of that fit, vs a CCSD(T)/pcSseg-3 reference on the NS372 benchmark\n", + "(proton shieldings); third axis is the method's introduction year, colored by Jacob's-ladder rung." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "37fd61c1", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3de7e8c2", + "metadata": {}, + "outputs": [], + "source": "import leveling\nimport convergence_plot" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de9503d1", + "metadata": {}, + "outputs": [], + "source": [ + "# the Kaupp NS372 spreadsheet is small and ships in the repo's data/ns372/ folder\n", + "KAUPP_XLSX = os.path.join(REPO, \"data\", \"ns372\", \"ct1c00919_si_002.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a3bba2c1", + "metadata": {}, + "outputs": [], + "source": [ + "# the 1H / pcSseg-3 slice as a raw-named table, then per-method statistics (slope and scaled RMSE vs CCSD(T)).\n", + "cc_df = leveling.load_ns372_1h_convergence(KAUPP_XLSX)\n", + "fig2c_results_df = leveling.convergence_statistics(cc_df)\n", + "print(f\"{len(fig2c_results_df)} methods vs CCSD(T)/pcSseg-3\")\n", + "fig2c_results_df.sort_values(\"RMSE (scaled)\").head(6)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "afb2a097", + "metadata": {}, + "outputs": [], + "source": [ + "# rung colors follow Jacob's ladder (DFT approximation hierarchy).\n", + "METHOD_COLOR_MAP = {\n", + " \"Ab Initio\": \"#2B2B2B\",\n", + " \"LDA\": \"#3B82F6\",\n", + " \"GGA\": \"#1A2A6B\",\n", + " \"Meta-GGA\": \"#1E7A8B\",\n", + " \"Global Hybrid\": \"#22C55E\",\n", + " \"Range-Separated Hybrid\": \"#F59E0B\",\n", + " \"Local Hybrid\": \"#EF4444\",\n", + " \"Double Hybrid\": \"#A855F7\",\n", + "}\n", + "\n", + "methods_categorization = {\n", + " \"HF\": \"Ab Initio\", \"MP2\": \"Ab Initio\",\n", + " \"SVWN\": \"LDA\",\n", + " \"BP86\": \"GGA\", \"BLYP\": \"GGA\", \"PBE\": \"GGA\", \"KT1\": \"GGA\", \"KT2\": \"GGA\",\n", + " \"KT3\": \"GGA\", \"HCTH\": \"GGA\", \"B97D\": \"GGA\",\n", + " \"TPSS\": \"Meta-GGA\", \"τ-HCTH\": \"Meta-GGA\", \"M06-L\": \"Meta-GGA\", \"VSXC\": \"Meta-GGA\",\n", + " \"MN15-L\": \"Meta-GGA\", \"B97M-V\": \"Meta-GGA\", \"SCAN\": \"Meta-GGA\", \"rSCAN\": \"Meta-GGA\",\n", + " \"r2SCAN\": \"Meta-GGA\",\n", + " \"TPSSh\": \"Global Hybrid\", \"B3LYP\": \"Global Hybrid\", \"B97-2\": \"Global Hybrid\",\n", + " \"PBE0\": \"Global Hybrid\", \"M06\": \"Global Hybrid\", \"PW6B95\": \"Global Hybrid\",\n", + " \"BHLYP\": \"Global Hybrid\", \"MN15\": \"Global Hybrid\", \"M06-2X\": \"Global Hybrid\",\n", + " \"CAM-B3LYP\": \"Range-Separated Hybrid\", \"ωB97X-D\": \"Range-Separated Hybrid\",\n", + " \"ωB97X-V\": \"Range-Separated Hybrid\", \"ωB97M-V\": \"Range-Separated Hybrid\",\n", + " \"LH07s-SVWN\": \"Local Hybrid\", \"MPSTS\": \"Local Hybrid\", \"LHJ14\": \"Local Hybrid\",\n", + " \"LH07t-SVWN\": \"Local Hybrid\", \"LH12ct-SsirPW92\": \"Local Hybrid\",\n", + " \"LH12ct-SsifPW92\": \"Local Hybrid\", \"LH14t-calPBE\": \"Local Hybrid\", \"LH20t\": \"Local Hybrid\",\n", + " \"B2PLYP\": \"Double Hybrid\", \"B2GP-PLYP\": \"Double Hybrid\", \"DSD-PBEP86\": \"Double Hybrid\",\n", + "}\n", + "\n", + "year_map = {\n", + " \"ωB97X-D\": 2008, \"ωB97X-V\": 2014, \"r2SCAN\": 2020, \"M06-L\": 2006, \"B97D\": 1997,\n", + " \"SVWN\": 1980, \"MP2\": 1934, \"HF\": 1930, \"DSD-PBEP86\": 2011, \"BP86\": 1988, \"BLYP\": 1988,\n", + " \"PBE\": 1996, \"KT1\": 2003, \"KT2\": 2003, \"HCTH\": 1998, \"TPSS\": 2003, \"τ-HCTH\": 2002,\n", + " \"VSXC\": 1998, \"MN15-L\": 2016, \"B97M-V\": 2015, \"KT3\": 2004, \"SCAN\": 2015, \"rSCAN\": 2019,\n", + " \"TPSSh\": 2003, \"B3LYP\": 1994, \"B97-2\": 1998, \"PBE0\": 1999, \"M06\": 2008, \"PW6B95\": 2005,\n", + " \"BHLYP\": 1993, \"MN15\": 2016, \"M06-2X\": 2008, \"CAM-B3LYP\": 2004, \"ωB97M-V\": 2016,\n", + " \"LH07s-SVWN\": 2007, \"LH07t-SVWN\": 2007, \"LHJ14\": 2014, \"LH12ct-SsirPW92\": 2017,\n", + " \"LH12ct-SsifPW92\": 2017, \"LH14t-calPBE\": 2014, \"LH20t\": 2020, \"B2PLYP\": 2006,\n", + " \"B2GP-PLYP\": 2006, \"MPSTS\": 2021,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "62e50faa", + "metadata": {}, + "outputs": [], + "source": "highlight_methods = [\n \"DSD-PBEP86\", \"LH07t-SVWN\", \"LH14t-calPBE\", \"ωB97X-D\", \"ωB97X-V\",\n \"BHLYP\", \"M06-L\", \"r2SCAN\", \"B97D\", \"SVWN\", \"MP2\", \"HF\",\n]\n\nfig, ax, saved = convergence_plot.plot_statistics_3d_emphasize_rungs(\n x_param=\"Slope\",\n y_param=\"RMSE (scaled)\",\n z_param=None,\n results_df=fig2c_results_df,\n methods_categorization=methods_categorization,\n method_color_map=METHOD_COLOR_MAP,\n year_map=year_map,\n save_path=figure_path(\"fig2d_convergence_1H.png\"),\n dpi=200,\n highlight_methods=highlight_methods,\n highlight_marker_size=80, # uniform with the non-highlighted points\n highlight_edgewidth=1.3,\n highlight_text=True,\n highlight_text_size=8,\n highlight_text_dx_frac=0.03,\n highlight_text_dy_frac=-0.05,\n highlight_text_dz=2.0,\n year_tick_step=5,\n z_jitter_by_category=0.03,\n z_lane_by_category=True,\n z_lane_width=2,\n elev=12,\n azim=25,\n xlim=(0.99, 1.105),\n ylim=(0, 0.40),\n z_break=(1935, 1980),\n z_break_gap=8,\n zlim=(1930, 2025),\n xy_lane_by_category=True,\n xy_lane_strength=0.1,\n xy_lane_style=\"circle\",\n lane_exclude_categories=(\"ab initio\",),\n group_by_decade=False,\n)" + } + ], + "metadata": { + "kernelspec": { + "display_name": "magnet-venv", + "language": "python", + "name": "magnet-venv" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig3a_pcm_benefit.ipynb b/analysis/manuscript_figures/fig3a_pcm_benefit.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..01c55c9469cb71da45a54a32914f0a7adc0f2dc2 --- /dev/null +++ b/analysis/manuscript_figures/fig3a_pcm_benefit.ipynb @@ -0,0 +1,151 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d13d8aa5", + "metadata": {}, + "source": [ + "# Figure 3A: implicit-solvent (PCM) benefit by method, chloroform vs benzene\n", + "\n", + "Per-split percent benefit of adding a PCM solvent correction, for each method, in chloroform and benzene." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "80ceb0d1", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "88b3d113", + "metadata": {}, + "outputs": [], + "source": [ + "import delta22\n", + "import paths\n", + "import fig3_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e830f642", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "df43c551", + "metadata": {}, + "outputs": [], + "source": [ + "# the published panels use 250 seeded train/test splits\n", + "N_SPLITS = 250" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19a780b7", + "metadata": {}, + "outputs": [], + "source": [ + "query = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n", + "solutes = sorted(query[\"solute\"].unique())\n", + "print(len(query), \"rows;\", len(solutes), \"solutes;\", query[\"sap_nmr_method\"].nunique(), \"methods\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c68c523", + "metadata": {}, + "outputs": [], + "source": [ + "fig3a_by_solvent = delta22.fig3a_pcm_benefit_by_solvent(query, [\"chloroform\", \"benzene\"],\n", + " n_splits=N_SPLITS, solutes=solutes)\n", + "\n", + "# per-method median percent benefit, then mean +/- std across methods\n", + "chloro = fig3a_by_solvent[fig3a_by_solvent[\"solvent\"] == \"chloroform\"]\n", + "benz = fig3a_by_solvent[fig3a_by_solvent[\"solvent\"] == \"benzene\"]\n", + "chloro_by_method = chloro.groupby(\"sap_nmr_method\")[\"percent_benefit\"].median()\n", + "benz_by_method = benz.groupby(\"sap_nmr_method\")[\"percent_benefit\"].median()\n", + "print(f\"PCM benefit in chloroform: {chloro_by_method.mean():+.1f} +/- {chloro_by_method.std():.1f}%\"\n", + " f\" (positive = PCM helps)\")\n", + "print(f\"PCM benefit in benzene: {benz_by_method.mean():+.1f} +/- {benz_by_method.std():.1f}%\"\n", + " f\" (negative = PCM hurts)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f282780f", + "metadata": {}, + "outputs": [], + "source": [ + "# Fixed method order, grouped by family (ab initio -> conventional DFT -> NMR-specific -> double\n", + "# hybrid), each as (method, basis, geometry). This ordering is what gives the panel its family\n", + "# structure; sorting by benefit value would scramble it.\n", + "FIG3A_COMBOS = [\n", + " (\"hf\", \"pcSseg3\", \"pbe0_tz\"), (\"mp2\", \"pcSseg2\", \"pbe0_tz\"), (\"dlpno_mp2\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"b3lyp_d3bj\", \"pcSseg3\", \"pbe0_tz\"), (\"pbe0_d3bj\", \"pcSseg2\", \"pbe0_tz\"),\n", + " (\"m062x_d3\", \"pcSseg3\", \"pbe0_tz\"), (\"b97d3_d3bj\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"wb97xd\", \"pcSseg3\", \"pbe0_tz\"), (\"tpsstpss_d3bj\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"bp86_d3bj\", \"pcSseg3\", \"pbe0_tz\"), (\"blyp_d3bj\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"wp04\", \"pcSseg3\", \"pbe0_tz\"), (\"wc04\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"), (\"B2GP_PLYP\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"B2PLYP\", \"pcSseg3\", \"pbe0_tz\"), (\"mPW2PLYP\", \"pcSseg3\", \"pbe0_tz\"),\n", + " (\"revdsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"),\n", + "]\n", + "FIG3A_COLORS = {\"chloroform\": \"#44AA99\", \"benzene\": \"#DDCC77\"} # teal / tan, as published" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2fa6355b", + "metadata": {}, + "outputs": [], + "source": [ + "fig3_plots.plot_pcm_benefit_with_arrows(fig3a_by_solvent, FIG3A_COMBOS, FIG3A_COLORS,\n", + " save_path=figure_path(\"fig3a_pcm_benefit_1H.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/manuscript_figures/fig3b_shifts.ipynb b/analysis/manuscript_figures/fig3b_shifts.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..25428a3a3840c3b6a5a262accc95af1f9e89d9b7 --- /dev/null +++ b/analysis/manuscript_figures/fig3b_shifts.ipynb @@ -0,0 +1,105 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "182eaf82", + "metadata": {}, + "source": [ + "# Figure 3B: solvent-induced shift differences (experiment vs PCM)\n", + "\n", + "Measured vs PCM-predicted differences between solvent-induced ¹H shifts (methanol/benzene vs chloroform reference)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89397a6b", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc6077cf", + "metadata": {}, + "outputs": [], + "source": [ + "import delta22\n", + "import paths\n", + "import fig3_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a093d2ad", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ec2ee33", + "metadata": {}, + "outputs": [], + "source": [ + "query = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n", + "print(len(query), \"rows;\", query[\"sap_nmr_method\"].nunique(), \"methods\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "141afd11", + "metadata": {}, + "outputs": [], + "source": [ + "shift = delta22.fig3b_shift_differences(query, \"wp04\", \"pcSseg2\", \"aimnet2\", nucleus=\"H\",\n", + " x_solvent=\"methanol\", y_solvent=\"benzene\", reference=\"chloroform\")\n", + "print(len(shift), \"sites\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7edca767", + "metadata": {}, + "outputs": [], + "source": [ + "fig3_plots.plot_shift_vs_pcm(shift, save_path=figure_path(\"fig3b_shifts_1H.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/manuscript_figures/fig3d_solvent_corrections.ipynb b/analysis/manuscript_figures/fig3d_solvent_corrections.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..fd95d39f73f98096d2c3b27d67d9f77119bb663b --- /dev/null +++ b/analysis/manuscript_figures/fig3d_solvent_corrections.ipynb @@ -0,0 +1,149 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0954a86e", + "metadata": {}, + "source": [ + "# Figure 3D: implicit vs explicit solvent corrections, one box per solvent\n", + "\n", + "All 12 solvents individually (polar aprotic -> polar protic -> aromatic), four schemes each: implicit (PCM), implicit + vibrations, explicit, explicit + vibrations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5172978e", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c7e70d68", + "metadata": {}, + "outputs": [], + "source": [ + "import delta22\n", + "import paths\n", + "import fig3_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55397f4b", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "06416ecb", + "metadata": {}, + "outputs": [], + "source": [ + "# the published panels use 250 seeded train/test splits\n", + "N_SPLITS = 250" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e892c346", + "metadata": {}, + "outputs": [], + "source": [ + "query = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n", + "solutes = sorted(query[\"solute\"].unique())\n", + "print(len(query), \"rows;\", len(solutes), \"solutes;\", query[\"sap_nmr_method\"].nunique(), \"methods\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32748b6d", + "metadata": {}, + "outputs": [], + "source": [ + "FIG3D_LABELS = {\n", + " \"stationary + pcm\": \"Implicit Solvent (PCM)\",\n", + " \"stationary_plus_qcd + pcm\": \"Implicit Solvent (PCM) + Vibrations\",\n", + " \"stationary + desmond\": \"Explicit Solvent\",\n", + " \"stationary_plus_qcd + desmond\": \"Explicit Solvent + Vibrations\",\n", + "}\n", + "fig3d_results = delta22.fig3d_formula_regressions(\n", + " query, \"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\", list(FIG3D_LABELS),\n", + " delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS, solutes=solutes)\n", + "\n", + "# report implicit vs explicit median test RMSE per solvent (not pooled by class)\n", + "for solvent in delta22.DESMOND_SOLVENTS:\n", + " sub = fig3d_results[fig3d_results[\"solvent\"] == solvent]\n", + " mi = sub[sub[\"formula\"] == \"stationary + pcm\"][\"test_RMSE\"].median()\n", + " me = sub[sub[\"formula\"] == \"stationary + desmond\"][\"test_RMSE\"].median()\n", + " print(f\"{solvent:20s} implicit={mi:.4f} explicit={me:.4f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e826d46", + "metadata": {}, + "outputs": [], + "source": [ + "FIG3D_COLORS = [\"#A72608\", \"#F4BAAD\", \"#5D737E\", \"#D9FFF5\"] # dark red, pink, gray, mint (as published)\n", + "FIG3D_SOLVENT_LABELS = {\n", + " \"chloroform\": r\"CDCl$_3$\", \"dichloromethane\": \"DCM\", \"tetrahydrofuran\": \"THF\",\n", + " \"acetonitrile\": \"MeCN\", \"dimethylsulfoxide\": \"DMSO\", \"methanol\": \"MeOD\",\n", + " \"trifluoroethanol\": \"TFE\", \"chlorobenzene\": \"PhCl\",\n", + "}\n", + "# solvent order: polar aprotic -> polar protic -> aromatic (group ordering from delta22.SOLVENT_GROUPS)\n", + "FIG3D_SOLVENT_ORDER = [s for group in delta22.SOLVENT_GROUPS.values() for s in group]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c0c66631", + "metadata": {}, + "outputs": [], + "source": [ + "fig3_plots.plot_solvent_correction_boxplot(fig3d_results, FIG3D_LABELS, FIG3D_SOLVENT_ORDER,\n", + " FIG3D_COLORS, FIG3D_SOLVENT_LABELS,\n", + " save_path=figure_path(\"fig3d_explicit_solvation_1H.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/manuscript_figures/fig4a_implicit.ipynb b/analysis/manuscript_figures/fig4a_implicit.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..d0eaa25a4199c0e8f8fe30cc14897a855d0a2eab --- /dev/null +++ b/analysis/manuscript_figures/fig4a_implicit.ipynb @@ -0,0 +1,113 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a1be4f68", + "metadata": {}, + "source": [ + "# Figure 4A: predicted vs measured solvent-induced shift, implicit (PCM)\n", + "\n", + "Implicit (PCM) predicted vs measured solvent-induced shifts (Δδ vs. CDCl3, ¹H) for one solvent per class." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27283d57", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "576ad573", + "metadata": {}, + "outputs": [], + "source": [ + "import delta22\n", + "import paths\n", + "import fig4_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fff8d490", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "236a76a8", + "metadata": {}, + "outputs": [], + "source": [ + "METHOD, BASIS, GEOM = \"b3lyp_d3bj\", \"pcSseg2\", \"aimnet2\"\n", + "q = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "one = q[(q[\"sap_nmr_method\"] == METHOD) & (q[\"sap_basis\"] == BASIS) & (q[\"sap_geometry_type\"] == GEOM)]\n", + "one = delta22.add_solvent_mean(one)\n", + "print(len(one), \"rows for\", METHOD, BASIS, GEOM)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f28bd2a4", + "metadata": {}, + "outputs": [], + "source": [ + "shifts_chcl3 = delta22.solvent_induced_shifts(one, \"chloroform\", delta22.DESMOND_SOLVENTS,\n", + " nucleus=\"H\", explicit=\"desmond\")\n", + "print(len(shifts_chcl3), \"site/solvent differences\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6f3c02f9", + "metadata": {}, + "outputs": [], + "source": [ + "# Water == TIP4P\n", + "FIG4_HIGHLIGHT = {\"dimethylsulfoxide\": \"DMSO\", \"TIP4P\": \"Water\", \"benzene\": \"Benzene\"}\n", + "FIG4_ORDER = [\"DMSO\", \"Water\", \"Benzene\"]\n", + "FIG4_PALETTE = {\"DMSO\": \"#61a89a\", \"Water\": \"#A72608\", \"Benzene\": \"#dacd82\"}\n", + "FIG4_MARKERS = {\"DMSO\": \"o\", \"Water\": \"^\", \"Benzene\": \"X\"}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59bb0a61", + "metadata": {}, + "outputs": [], + "source": [ + "fig4_plots.save_panel(\n", + " lambda ax: fig4_plots.draw_shift_scatter(ax, shifts_chcl3, \"implicit_diff\",\n", + " FIG4_HIGHLIGHT, FIG4_ORDER, FIG4_PALETTE, FIG4_MARKERS),\n", + " figure_path(\"fig4a_implicit_1H.png\"),\n", + ")" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig4b_explicit.ipynb b/analysis/manuscript_figures/fig4b_explicit.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..b033373dc1dcb8a4d5b22b7d588ea469cd36b920 --- /dev/null +++ b/analysis/manuscript_figures/fig4b_explicit.ipynb @@ -0,0 +1,113 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "915187f5", + "metadata": {}, + "source": [ + "# Figure 4B: predicted vs measured solvent-induced shift, explicit (Desmond)\n", + "\n", + "Explicit (Desmond) predicted vs measured solvent-induced shifts (Δδ vs. CDCl3, ¹H) for one solvent per class." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d55e803", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18a69a1a", + "metadata": {}, + "outputs": [], + "source": [ + "import delta22\n", + "import paths\n", + "import fig4_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24c48a7e", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a56abe4a", + "metadata": {}, + "outputs": [], + "source": [ + "METHOD, BASIS, GEOM = \"b3lyp_d3bj\", \"pcSseg2\", \"aimnet2\"\n", + "q = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "one = q[(q[\"sap_nmr_method\"] == METHOD) & (q[\"sap_basis\"] == BASIS) & (q[\"sap_geometry_type\"] == GEOM)]\n", + "one = delta22.add_solvent_mean(one)\n", + "print(len(one), \"rows for\", METHOD, BASIS, GEOM)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3b2dc95", + "metadata": {}, + "outputs": [], + "source": [ + "shifts_chcl3 = delta22.solvent_induced_shifts(one, \"chloroform\", delta22.DESMOND_SOLVENTS,\n", + " nucleus=\"H\", explicit=\"desmond\")\n", + "print(len(shifts_chcl3), \"site/solvent differences\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bd10f351", + "metadata": {}, + "outputs": [], + "source": [ + "# Water == TIP4P\n", + "FIG4_HIGHLIGHT = {\"dimethylsulfoxide\": \"DMSO\", \"TIP4P\": \"Water\", \"benzene\": \"Benzene\"}\n", + "FIG4_ORDER = [\"DMSO\", \"Water\", \"Benzene\"]\n", + "FIG4_PALETTE = {\"DMSO\": \"#61a89a\", \"Water\": \"#A72608\", \"Benzene\": \"#dacd82\"}\n", + "FIG4_MARKERS = {\"DMSO\": \"o\", \"Water\": \"^\", \"Benzene\": \"X\"}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a7b4432", + "metadata": {}, + "outputs": [], + "source": [ + "fig4_plots.save_panel(\n", + " lambda ax: fig4_plots.draw_shift_scatter(ax, shifts_chcl3, \"explicit_diff\",\n", + " FIG4_HIGHLIGHT, FIG4_ORDER, FIG4_PALETTE, FIG4_MARKERS),\n", + " figure_path(\"fig4b_explicit_1H.png\"),\n", + ")" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig4c_rmse_range.ipynb b/analysis/manuscript_figures/fig4c_rmse_range.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..e6b34e31463d9a80c5cf44f2edb30b264754d67a --- /dev/null +++ b/analysis/manuscript_figures/fig4c_rmse_range.ipynb @@ -0,0 +1,127 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "6d0c4a9d", + "metadata": {}, + "source": [ + "# Figure 4C: fit RMSE vs experimental range, implicit vs explicit\n", + "\n", + "Per-solvent fit RMSE vs experimental shift range (¹H), implicit (PCM) vs explicit (Desmond)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8a3ec572", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "79b3bd64", + "metadata": {}, + "outputs": [], + "source": [ + "import delta22\n", + "import paths\n", + "import fig4_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "79171fdb", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdaaf850", + "metadata": {}, + "outputs": [], + "source": [ + "# solvent_mean is a synthetic pseudo-solvent representing the average across solvents; it is used below as the reference baseline.\n", + "METHOD, BASIS, GEOM = \"b3lyp_d3bj\", \"pcSseg2\", \"aimnet2\"\n", + "q = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "one = q[(q[\"sap_nmr_method\"] == METHOD) & (q[\"sap_basis\"] == BASIS) & (q[\"sap_geometry_type\"] == GEOM)]\n", + "one = delta22.add_solvent_mean(one)\n", + "print(len(one), \"rows for\", METHOD, BASIS, GEOM)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e5527018", + "metadata": {}, + "outputs": [], + "source": [ + "FIG4C_LABELS = {\n", + " \"chloroform\": \"CDCl3\", \"dichloromethane\": \"DCM\", \"tetrahydrofuran\": \"THF\",\n", + " \"acetonitrile\": \"MeCN\", \"dimethylsulfoxide\": \"DMSO\", \"acetone\": \"acetone\",\n", + " \"methanol\": \"MeOD\", \"TIP4P\": \"TIP4P\", \"trifluoroethanol\": \"trifluoroethanol\",\n", + " \"benzene\": \"benzene\", \"toluene\": \"toluene\", \"chlorobenzene\": \"chlorobenzene\",\n", + "}\n", + "FIG4C_COLORS = {\"Implicit (PCM)\": \"#A72608\", \"Explicit (Desmond)\": \"#61a89a\"} # red diamond / teal circle" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3cdab0a7", + "metadata": {}, + "outputs": [], + "source": [ + "# differences are taken against the solvent-averaged pseudo-solvent so every real solvent, including chloroform, appears as a point.\n", + "rmse_range_rows = []\n", + "for solvent in delta22.DESMOND_SOLVENTS:\n", + " sp = delta22.solvent_pair_differences(one, solvent, \"solvent_mean\", nucleus=\"H\", explicit=\"desmond\")\n", + " if len(sp) == 0:\n", + " continue\n", + " rmse_range_rows.append({\n", + " \"solvent\": solvent,\n", + " \"range\": float(sp[\"exp_diff\"].max() - sp[\"exp_diff\"].min()),\n", + " \"implicit\": delta22.fit_differences_to_experimental(sp, \"implicit_diff\")[\"rmse\"],\n", + " \"explicit\": delta22.fit_differences_to_experimental(sp, \"explicit_diff\")[\"rmse\"],\n", + " })\n", + "for r in sorted(rmse_range_rows, key=lambda r: r[\"range\"]):\n", + " print(f\"{FIG4C_LABELS[r['solvent']]:16s} range={r['range']:.3f} \"\n", + " f\"implicit={r['implicit']:.3f} explicit={r['explicit']:.3f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4601c2c7", + "metadata": {}, + "outputs": [], + "source": [ + "fig4_plots.save_panel(\n", + " lambda ax: fig4_plots.draw_rmse_dumbbell(ax, rmse_range_rows, FIG4C_LABELS, FIG4C_COLORS),\n", + " figure_path(\"fig4c_rmse_range_1H.png\"),\n", + ")" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig5b_dft8k.ipynb b/analysis/manuscript_figures/fig5b_dft8k.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..b4a4044ff0e8ef069b8b2fe2b5641423601512bb --- /dev/null +++ b/analysis/manuscript_figures/fig5b_dft8k.ipynb @@ -0,0 +1,188 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "706410fd", + "metadata": {}, + "source": [ + "# Figure 5B: MagNET-Zero vs DFT shieldings on the DFT8K benchmark\n", + "\n", + "Reproduces Figure 5B: MagNET-Zero gas-phase shielding predictions vs DFT on the DFT8K benchmark (~7,000\n", + "organics, untrained). Protons vs WP04/pcSseg-2, carbons vs wB97X-D/pcSseg-2; residual = DFT minus\n", + "MagNET." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7e0a3b95", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/dft8k\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c9838352", + "metadata": {}, + "outputs": [], + "source": [ + "import dft8k_residuals\n", + "import paths\n", + "import fig5b_plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d2231e1e", + "metadata": {}, + "outputs": [], + "source": [ + "DFT8K_HDF5 = paths.dataset_file(\"dft8k\", root=REPO)\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b7d02a23", + "metadata": {}, + "outputs": [], + "source": [ + "shieldings = dft8k_residuals.load_shieldings(DFT8K_HDF5)\n", + "print(\"atoms in DFT8K:\", len(shieldings[\"atomic_numbers\"]))" + ] + }, + { + "cell_type": "markdown", + "id": "867ed2b4", + "metadata": {}, + "source": [ + "## Residual statistics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02ca1b09", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus in [\"H\", \"C\"]:\n", + " errors = dft8k_residuals.residuals(shieldings, nucleus)\n", + " stats = dft8k_residuals.residual_stats(errors)\n", + " print(f\"{nucleus}: n={stats['n']:,} RMSE={stats['rmse']:.4f} ppm MAE={stats['mae']:.4f} ppm \"\n", + " f\"95% abs={stats['abs_p95']:.4f} ppm fraction < 0.1 ppm={stats['frac_below']:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "02d95406", + "metadata": {}, + "source": [ + "## Residual histograms" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bf9148b3", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus in [\"H\", \"C\"]:\n", + " errors = dft8k_residuals.residuals(shieldings, nucleus)\n", + " stats = dft8k_residuals.residual_stats(errors)\n", + " print(f\"{nucleus}: RMSE={stats['rmse']:.3f} ppm frac within 0.1 ppm={stats['frac_below']:.3f}\")\n", + " if nucleus == \"H\":\n", + " fig5b_plots.plot_dft8k_residual_histogram(errors, figure_path(\"fig5b_residuals_1H.png\"),\n", + " bin_width=0.0025, xlim=0.23, grey_box=0.1, x_tick=0.05)\n", + " else: # 13C: wider window, no +/-0.1 box since that threshold is 1H-specific\n", + " fig5b_plots.plot_dft8k_residual_histogram(errors, None,\n", + " bin_width=0.05, xlim=5.0, grey_box=None, x_tick=1.0)" + ] + }, + { + "cell_type": "markdown", + "id": "5eab006a", + "metadata": {}, + "source": [ + "## Extreme-residual callouts\n", + "\n", + "The largest positive and negative ¹H residuals over the full set. The published negative callout (a\n", + "sulfonium zwitterion, -1.040 ppm) is not the true global minimum -- see the printed note." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b1f2315b", + "metadata": {}, + "outputs": [], + "source": [ + "extreme_max = dft8k_residuals.find_extreme_residual(DFT8K_HDF5, \"H\", sign=\"max\")\n", + "print(\"largest positive 1H residual:\", extreme_max)\n", + "fig5b_plots.show_dft8k_molecule(extreme_max[\"smiles\"])\n", + "\n", + "zwitterion = dft8k_residuals.molecule_by_id(DFT8K_HDF5, \"H\", molecule_id=88779)\n", + "print(\"published sulfonium-zwitterion callout (id 88779):\", zwitterion)\n", + "fig5b_plots.show_dft8k_molecule(zwitterion[\"smiles\"])" + ] + }, + { + "cell_type": "markdown", + "id": "5f4276e7", + "metadata": {}, + "source": [ + "## Functional-group error breakdown\n", + "\n", + "Mean absolute residual for six functional groups (Carbonyls, Amines, Sulfonyl, Pyridines, Furans,\n", + "Nitroso), via RDKit SMARTS matching over every molecule's SMILES\n", + "(`dft8k_residuals.FUNCTIONAL_GROUP_SMARTS`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f95f9d58", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus in [\"H\", \"C\"]:\n", + " group_errors = dft8k_residuals.functional_group_errors(DFT8K_HDF5, nucleus)\n", + " print(f\"--- {nucleus} ---\")\n", + " for name, v in group_errors.items():\n", + " print(f\"{name:12s} mean|error|={v['mean_abs_error']:.4f} ppm n_molecules={v['n_molecules']:,}\")\n", + " save = figure_path(\"fig5b_functional_groups_1H.png\") if nucleus == \"H\" else None\n", + " fig5b_plots.plot_dft8k_functional_group_errors(group_errors, nucleus, save)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig5c.ipynb b/analysis/manuscript_figures/fig5c.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..d2d1ff6b866b07d103a2cdd77b9f0ca6238a86d6 --- /dev/null +++ b/analysis/manuscript_figures/fig5c.ipynb @@ -0,0 +1,115 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9f479dac", + "metadata": {}, + "source": [ + "# Figure 5C: implicit vs explicit correction RMSE, delta-22 vs the natural-products test set\n", + "\n", + "Bootstrap RMSE of implicit (PCM) vs explicit (OpenMM + vibrations) corrections, delta-22 vs the pooled complex/natural-product test set, chloroform and benzene (bars = mean, error bars = 2.5-97.5 percentile). Produces the ¹H panel (saved) and a ¹³C companion (inline only)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c4fd38cc", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/applications\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aef1da3f", + "metadata": {}, + "outputs": [], + "source": [ + "from applications_reader import Applications\n", + "import applications\n", + "import applications_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "178e31b6", + "metadata": {}, + "outputs": [], + "source": [ + "APPLICATIONS_HDF5 = paths.dataset_file(\"applications\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"applications\", \"applications_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a013c941", + "metadata": {}, + "outputs": [], + "source": [ + "loader = Applications(APPLICATIONS_HDF5, XLSX)\n", + "query_df_nn = applications.build_query_df_nn(loader)\n", + "seed = applications.build_bootstrap_seed_coeffs(loader)\n", + "\n", + "# pooled \"Test Set\" bootstrap RMSE distributions (one bin over all 13 complex molecules)\n", + "grouped = {}\n", + "for nuc in [\"H\", \"C\"]:\n", + " preds = applications.apply_bootstrap_params_to_full_dataset(query_df_nn, seed[nuc], nucleus=nuc)\n", + " grouped[nuc] = applications.compute_grouped_rmse(preds, applications.ALL_IN_ONE_BIN)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d3cf53a", + "metadata": {}, + "outputs": [], + "source": [ + "applications_plots.plot_nps_benefit_barplot(\n", + " loader.rmse_distribution(\"H\"), grouped[\"H\"], nucleus=\"H\",\n", + " solvents=[[\"chloroform\"], [\"benzene\"]], np_solute_groups=applications.ALL_IN_ONE_BIN,\n", + " formulas=[\"stationary_plus_pcm\", \"stationary_plus_qcd + openMM\"],\n", + " labels=[\"Implicit Solvent (SotA)\", \"Explicit Solvent + Vibrations (OpenMM)\"],\n", + " colors=[\"#A72608\", \"#61a89a\"], formula_remap=applications.FORMULA_REMAP,\n", + " figsize=(6, 5), y_min=0.0, y_max=0.37, save_path=figure_path(\"fig5c_benefit_1H.png\"))\n", + "\n", + "applications_plots.plot_nps_benefit_barplot(\n", + " loader.rmse_distribution(\"C\"), grouped[\"C\"], nucleus=\"C\",\n", + " solvents=[[\"chloroform\"], [\"benzene\"]], np_solute_groups=applications.ALL_IN_ONE_BIN,\n", + " formulas=[\"stationary_plus_pcm\", \"stationary_plus_op_vib + openMM\"],\n", + " labels=[\"Implicit Solvent (SotA)\", \"Explicit Solvent + Vibrations (OpenMM)\"],\n", + " colors=[\"#A72608\", \"#61a89a\"], formula_remap=applications.FORMULA_REMAP,\n", + " figsize=(6, 5), y_min=0.0, save_path=None) # inline-only companion, not saved to disk" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/manuscript_figures/fig5d.ipynb b/analysis/manuscript_figures/fig5d.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..3ef1a3c2b04205a76f96a723893511c884d841fc --- /dev/null +++ b/analysis/manuscript_figures/fig5d.ipynb @@ -0,0 +1,145 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4b208bf1", + "metadata": {}, + "source": [ + "# Figure 5D: per-solute explicit-correction RMSE across the test set\n", + "\n", + "Per-solute bootstrap RMSE (¹H, solvent-averaged) for the explicit-solvent correction, ordered delta-22 -> olefin/pyridine isomers -> natural products, with a dashed \"scaled to solute\" baseline and a red \"scaled to test set\" line." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66e2a22e", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/applications\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9b70b4ca", + "metadata": {}, + "outputs": [], + "source": [ + "from applications_reader import Applications\n", + "import applications\n", + "import applications_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2da375fc", + "metadata": {}, + "outputs": [], + "source": [ + "APPLICATIONS_HDF5 = paths.dataset_file(\"applications\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"applications\", \"applications_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a6acf673", + "metadata": {}, + "outputs": [], + "source": [ + "loader = Applications(APPLICATIONS_HDF5, XLSX)\n", + "query_df_nn = applications.build_query_df_nn(loader)\n", + "seed = applications.build_bootstrap_seed_coeffs(loader)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "889345fc", + "metadata": {}, + "outputs": [], + "source": [ + "site_counts = (query_df_nn.drop_duplicates(subset=[\"solute\", \"nucleus\", \"site\"])\n", + " .groupby([\"solute\", \"nucleus\"]).size().unstack(fill_value=0))\n", + "per_solute = applications.per_solute_fits(query_df_nn)\n", + "all_solute = applications.per_solvent_fits(query_df_nn)\n", + "preds_h = applications.apply_bootstrap_params_to_full_dataset(query_df_nn, seed[\"H\"], nucleus=\"H\")\n", + "bootstrap_rmses_h = applications.compute_solute_rmses(preds_h)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08869e35", + "metadata": {}, + "outputs": [], + "source": [ + "FIG5D_LABELS = {\n", + " \"isomer_1E\": \"Olefin 1 (E)\", \"isomer_1Z\": \"Olefin 1 (Z)\",\n", + " \"isomer_2E\": \"Olefin 2 (E)\", \"isomer_2Z\": \"Olefin 2 (Z)\",\n", + " \"isomer_3E\": \"Olefin 3 (E)\", \"isomer_3Z\": \"Olefin 3 (Z)\",\n", + " \"isomer_4N\": \"Pyridone 4\", \"isomer_4O\": \"Pyridine 4\",\n", + " \"vomicine\": \"Vomicine\", \"prednisone\": \"Prednisone\",\n", + " \"peptide\": \"Acetyl-L-alanyl-L-\\nglutamine\", \"flavone\": \"Flavone\",\n", + " \"dihydrotanshinone_I\": \"Dihydrotanshinone I\",\n", + "}\n", + "FIG5D_ORDER = [\"isomer_1E\", \"isomer_1Z\", \"isomer_2E\", \"isomer_2Z\", \"isomer_3E\", \"isomer_3Z\",\n", + " \"isomer_4N\", \"isomer_4O\", \"vomicine\", \"prednisone\", \"peptide\", \"flavone\",\n", + " \"dihydrotanshinone_I\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1ad5844c", + "metadata": {}, + "outputs": [], + "source": [ + "applications_plots.plot_nps_on_boxplot_delta22_simplified(\n", + " loader.rmse_distribution(\"H\"), bootstrap_rmses_h, per_solute[\"H\"],\n", + " nucleus=\"H\", formulas=[\"stationary_plus_qcd + openMM\"], colors=[\"#61a89a\"],\n", + " site_counts=site_counts, formula_remap=applications.FORMULA_REMAP,\n", + " solute_remap=FIG5D_LABELS, solute_order=FIG5D_ORDER,\n", + " title=\"Complex Molecules are Dominated by Additional Physics\",\n", + " solute_color_remap=applications.PEPTIDE_HIGHLIGHT_H,\n", + " figsize=(14, 8), box_width=0.20, box_gap=0.1,\n", + " show_baseline=True, baseline_annotation_text=\"Scaled to solute\",\n", + " baseline_annotation_x=0.215, baseline_annotation_y=0.39,\n", + " show_full_fit_line=True, full_fit_line_label=\"Scaled to Test Set\", full_fit_line_label_x=0.835,\n", + " all_solute_fitting_results=all_solute,\n", + " max_bar_height=0.06, site_count_axis_mode=\"inset\", site_count_inset_area_frac=0.1,\n", + " site_count_inset_axis_offset=-0.06,\n", + " save_path=figure_path(\"fig5d_complex_1H.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_composite_models_ablations.ipynb b/analysis/si_figures/si_composite_models_ablations.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..8a564634e6f46eee26f2c357cfd5c7be7ab86ee7 --- /dev/null +++ b/analysis/si_figures/si_composite_models_ablations.ipynb @@ -0,0 +1,176 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bfdd4bf5", + "metadata": {}, + "source": [ + "# SI Figure S15: composite-formula ablation workbook (`ablations.xlsx`)\n", + "\n", + "Rebuilds the `ablations.xlsx` workbook behind ~32 of SI Figure S15's images: ~25 composite-formula\n", + "variants (stationary geometry plus some combination of implicit PCM, explicit Desmond, and\n", + "rovibrational QCD corrections) across all 12 delta-22 solvents, at two reference levels:\n", + "DSD-PBEP86/pcSseg-3 (geometry PBE0/tz) and the MagNET-Zero training reference (WP04/pcSseg-2 for ¹H,\n", + "wB97X-D/pcSseg-2 for ¹³C)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cb68d1d6", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20914013", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import composite_models\n", + "import composite_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26ab5657", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)\n", + "\n", + "def document_path(name):\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)\n", + "\n", + "# the shipped ablations.xlsx used 250 seeded train/test splits per (formula, solvent)\n", + "N_SPLITS = 250" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "76d31292", + "metadata": {}, + "outputs": [], + "source": [ + "query_df_dft = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n", + "solutes = delta22.delta22_solutes(DELTA22_HDF5)\n", + "print(len(query_df_dft), \"rows;\", len(solutes), \"solutes\")" + ] + }, + { + "cell_type": "markdown", + "id": "2c351f43", + "metadata": {}, + "source": [ + "## Preview: one formula's mean test RMSE at the DSD-PBEP86 reference level" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eb03ec52", + "metadata": {}, + "outputs": [], + "source": [ + "# preview before the full build (all ~25 formulas x 12 solvents x 250 splits x 2 nuclei x 2\n", + "# reference levels)\n", + "level = composite_models.REFERENCE_LEVELS[\"dsd\"]\n", + "preview = composite_models.ablation_rmse_table(query_df_dft, \"H\", level[\"method_h\"], level[\"basis_h\"],\n", + " level[\"geometry_h\"], solutes, n_splits=20)\n", + "preview[[\"chloroform\", \"benzene\", \"Mean Test RMSE\"]].round(3)" + ] + }, + { + "cell_type": "markdown", + "id": "ee9136d7", + "metadata": {}, + "source": [ + "## Build the full ablations workbook (both reference levels, both nuclei)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5c9fbbd3", + "metadata": {}, + "outputs": [], + "source": [ + "output_path = document_path(\"ablations.xlsx\")\n", + "composite_models.build_ablations_workbook(query_df_dft, solutes, output_path, n_splits=N_SPLITS)" + ] + }, + { + "cell_type": "markdown", + "id": "78a693c7", + "metadata": {}, + "source": [ + "## Correlations Between Features\n", + "\n", + "Solvent-averaged Pearson r correlation matrix between the five composite-model features (stationary\n", + "shielding, PCM, Desmond, its vibrational analogue, QCD), both nuclei, plus a per-solvent\n", + "PCM-vs-Desmond table." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9ef593a6", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus, label in [(\"H\", \"Proton\"), (\"C\", \"Carbon\")]:\n", + " corr = delta22.si_s15_feature_correlations(query_df_dft, nucleus)\n", + " nuc_label = \"1H\" if nucleus == \"H\" else \"13C\"\n", + " nuc_title = \"$^{1}$H\" if nucleus == \"H\" else \"$^{13}$C\"\n", + " composite_plots.plot_feature_correlation_heatmap(\n", + " corr[\"r\"], vmin=-1, vmax=1, cmap=\"RdBu\",\n", + " title=f\"Solvent-Averaged Pearson $r$ Correlation Matrix for {nuc_title}\",\n", + " save_path=figure_path(f\"si_figure_s15_feature_corr_r_{nuc_label}.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5bc382f2", + "metadata": {}, + "outputs": [], + "source": [ + "pcm_desmond_corr = delta22.pcm_desmond_correlation_by_solvent(query_df_dft)\n", + "print(\"PCM vs. Desmond Pearson R by nucleus and solvent:\")\n", + "display(pcm_desmond_corr.round(3))\n", + "composite_plots.plot_pcm_desmond_correlation_table(pcm_desmond_corr,\n", + " save_path=figure_path(\"si_figure_s15_pcm_desmond_corr_table.png\"))\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s01.ipynb b/analysis/si_figures/si_figure_s01.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..781c7ee81a15d16be93941bcc87afdc72bed1833 --- /dev/null +++ b/analysis/si_figures/si_figure_s01.ipynb @@ -0,0 +1,85 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "6086dfef", + "metadata": {}, + "source": "# SI Figure S1: accuracy-vs-cost Pareto frontier (proton and carbon)\n\nFigure 2A for both nuclei in CDCl3: fitting RMSE vs total compute time, colored by method type, shaped by geometry source, sized by basis (panel A proton, B carbon)." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66768963", + "metadata": {}, + "outputs": [], + "source": "import os, sys\n\n# make the in-repo modules importable (not pip-installed)\nREPO = os.path.abspath(\"../..\")\nfor _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n sys.path.insert(0, os.path.join(REPO, _p))" + }, + { + "cell_type": "code", + "id": "ded31d36", + "source": "import delta22\nimport pareto_plot\nimport paths", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "354b195e", + "source": "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\nXLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n\ndef figure_path(name):\n os.makedirs(\"figures\", exist_ok=True)\n return os.path.join(\"figures\", name)", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bd7563ca", + "metadata": {}, + "outputs": [], + "source": "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\nnn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)\ndft_gas_timings, nn_timings = delta22.load_pareto_timings(DELTA22_HDF5)\n\npoints = delta22.fig2a_pareto_points(dft, nn, dft_gas_timings, nn_timings, n_splits=100)\nprint(points[\"nmr_method\"].nunique(), \"methods;\",\n \"MagNET total time\", float(points.query(\"nmr_method=='MagNET'\")[\"total_time\"].iloc[0]), \"s\")" + }, + { + "cell_type": "markdown", + "id": "06d579ec", + "metadata": {}, + "source": [ + "## Proton and carbon Pareto panels\n", + "\n", + "Panel A is the proton frontier, B the carbon; they differ only in y-range, the NMR reference point\n", + "(wp04 for ¹H, wb97xd for ¹³C), and label offsets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c5a5ed42", + "metadata": {}, + "outputs": [], + "source": "method_legend_label_map = {\n \"MagNET-Zero\": \"MagNET-Zero\", \"ab initio\": r\"$\\it{ab\\ initio}$\", \"double hybrids\": \"double hybrid\",\n \"NMR-specific\": \"NMR-specific\", \"DFT\": \"DFT\",\n}\ngeometry_legend_label_map = {\"pbe0_tz\": \"DFT (PBE0/cc-pVTZ)\", \"aimnet2\": \"ML (AIMNet2)\"}\nmethod_color_map = {\"ab initio\": \"#53585f\", \"MagNET-Zero\": \"#000000\", \"double hybrids\": \"#b3a560\",\n \"NMR-specific\": \"#a72608\", \"DFT\": \"#4a8075\"}\ngeometry_marker_map = {\"aimnet2\": \"x\", \"pbe0_tz\": \"o\"}\n\n# which points to draw in each panel: MagNET-Zero, all PBE0-geometry DFT methods, plus one\n# AIMNet2-geometry NMR-specific reference (wp04 for proton, wb97xd for carbon)\nquery_str = (\"nucleus=='{nucleus}' and solvent == 'chloroform' and (\"\n \" (geometry_type=='aimnet2' and nmr_method.str.startswith('MagNET')) or\"\n \" (geometry_type=='pbe0_tz' and not nmr_method.str.startswith('MagNET')) or\"\n \" (geometry_type=='aimnet2' and nmr_method=='wp04' and basis=='pcSseg2' and nucleus=='H') or\"\n \" (geometry_type=='aimnet2' and nmr_method=='wb97xd' and basis=='pcSseg2' and nucleus=='C'))\")\n\nproton_label_positions = {\n (\"MagNET-Zero\", \"N/A\", \"aimnet2\"): {\"xytext\": (-28, 14)},\n (\"wp04\", \"pcSseg2\", \"aimnet2\"): {\"xytext\": (6, 4)},\n (\"pbe0\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (-22, -2)},\n (\"b2gp_plyp\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (4, 6)},\n (\"revdsd_pbep86\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (6, -4)},\n (\"revdsd_pbep86\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, 6)},\n (\"tpsstpss\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, -2)},\n (\"b97d3\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (-28, -2)},\n (\"mpw2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 2)},\n (\"b2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 0)},\n (\"tpsstpss\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (6, -12)},\n}\ncarbon_label_positions = {\n (\"MagNET-Zero\", \"N/A\", \"aimnet2\"): {\"xytext\": (-28, 14)},\n (\"wb97xd\", \"pcSseg2\", \"aimnet2\"): {\"xytext\": (6, -4)},\n (\"b2gp_plyp\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (-20, 8)},\n (\"revdsd_pbep86\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (0, -12)},\n (\"hf\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (6, -6)},\n (\"revdsd_pbep86\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, 0)},\n (\"tpsstpss\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, -4)},\n (\"mpw2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 2)},\n (\"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 2)},\n (\"revdsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, -6)},\n (\"m062x\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 0)},\n (\"wp04\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 0)},\n}" + }, + { + "cell_type": "code", + "id": "ba4a0799", + "source": "pareto_plot.plot_pareto_panel(\n points, query_str, nucleus=\"H\", suptitle=\"MagNET Extends a Flat Pareto Frontier\",\n method_color_map=method_color_map, method_legend_label_map=method_legend_label_map,\n geometry_legend_label_map=geometry_legend_label_map, geometry_marker_map=geometry_marker_map,\n xlim_left=(1.7, 1.8), xlim_right=(3.8, 5.5), ylim=(0.08, 0.3),\n figsize=(12, 6), marker_alpha=0.8, manual_label_positions=proton_label_positions,\n save_png=figure_path(\"si_figure_s01_pareto_1H.png\"),\n)\n\npareto_plot.plot_pareto_panel(\n points, query_str, nucleus=\"C\", suptitle=\"MagNET Extends a Flat Pareto Frontier\",\n method_color_map=method_color_map, method_legend_label_map=method_legend_label_map,\n geometry_legend_label_map=geometry_legend_label_map, geometry_marker_map=geometry_marker_map,\n xlim_left=(1.7, 1.8), xlim_right=(3.8, 5.5), ylim=(1.2, 4.2),\n figsize=(12, 6), marker_alpha=0.8, manual_label_positions=carbon_label_positions,\n save_png=figure_path(\"si_figure_s01_pareto_13C.png\"),\n)", + "metadata": {}, + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/analysis/si_figures/si_figure_s02_diels_alder.ipynb b/analysis/si_figures/si_figure_s02_diels_alder.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..1234d90fc06fe2dcabeb8ff01b77cc24a034f859 --- /dev/null +++ b/analysis/si_figures/si_figure_s02_diels_alder.ipynb @@ -0,0 +1,130 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "69a77df0", + "metadata": {}, + "source": [ + "# SI Figure S2: literature Diels-Alder barrier predictions for ethylene + butadiene\n", + "\n", + "Literature quantum-chemistry activation barriers for the ethylene + butadiene Diels-Alder reaction, spanning ~15-45 kcal/mol (experiment ~23.3 kcal/mol)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "876d6185", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "65812dad", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "from matplotlib.colors import TwoSlopeNorm\n", + "\n", + "import diels_alder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "09cc398b", + "metadata": {}, + "outputs": [], + "source": [ + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f4057ed3", + "metadata": {}, + "outputs": [], + "source": [ + "print(len(diels_alder.DATA), 'literature predictions; experimental barrier', diels_alder.EXPERIMENTAL_BARRIER, 'kcal/mol')\n", + "print('range:', min(e for _, e in diels_alder.DATA), 'to', max(e for _, e in diels_alder.DATA), 'kcal/mol')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9054cad6", + "metadata": {}, + "outputs": [], + "source": [ + "# One bar per literature prediction, grouped by method family (diels_alder.ordered_layout) and sorted by\n", + "# barrier within each group, colored by distance from the experimental value (red above, blue below),\n", + "# with a dashed line at the experimental barrier.\n", + "layout = diels_alder.ordered_layout(diels_alder.DATA)\n", + "methods, energies, positions = layout[\"methods\"], layout[\"energies\"], layout[\"positions\"]\n", + "\n", + "norm = TwoSlopeNorm(vmin=min(energies), vcenter=diels_alder.EXPERIMENTAL_BARRIER, vmax=max(energies))\n", + "colormap = plt.cm.RdBu_r\n", + "colors = [colormap(norm(e)) for e in energies]\n", + "\n", + "fig, ax = plt.subplots(figsize=(16, 8))\n", + "ax.bar(positions, energies, color=colors, alpha=0.8, edgecolor=\"black\", linewidth=0.5)\n", + "ax.axhline(y=diels_alder.EXPERIMENTAL_BARRIER, color=\"black\", linestyle=\"--\", linewidth=2)\n", + "\n", + "# callout on the B3LYP/6-31G* bar at 23.1 kcal/mol (the level used elsewhere in the paper)\n", + "b3lyp = next((i for i, (m, e) in enumerate(zip(methods, energies))\n", + " if \"B3LYP/6-31G*\" in m and abs(e - 23.1) < 0.1), None)\n", + "if b3lyp is not None:\n", + " ax.annotate(f\"Experimental: {diels_alder.EXPERIMENTAL_BARRIER} kcal/mol\",\n", + " xy=(positions[b3lyp], diels_alder.EXPERIMENTAL_BARRIER),\n", + " xytext=(positions[b3lyp], diels_alder.EXPERIMENTAL_BARRIER + 8),\n", + " arrowprops=dict(arrowstyle=\"->\", color=\"black\", lw=1),\n", + " fontsize=10, fontweight=\"bold\", ha=\"center\",\n", + " bbox=dict(boxstyle=\"round,pad=0.3\", facecolor=\"white\", edgecolor=\"black\", alpha=0.8))\n", + "\n", + "ax.set_ylabel(\"Predicted barrier (kcal/mol)\")\n", + "ax.set_xticks(positions)\n", + "ax.set_xticklabels(methods, rotation=45, ha=\"right\", fontsize=8)\n", + "if b3lyp is not None:\n", + " ax.get_xticklabels()[b3lyp].set_fontweight(\"bold\")\n", + "\n", + "y_range = ax.get_ylim()[1] - ax.get_ylim()[0]\n", + "label_y = ax.get_ylim()[0] - y_range * 0.25\n", + "for category, (start, end) in zip(layout[\"group_labels\"], layout[\"group_boundaries\"]):\n", + " ax.text((start + end) / 2, label_y, category, ha=\"center\", va=\"top\",\n", + " fontweight=\"bold\", fontsize=11)\n", + "\n", + "fig.tight_layout()\n", + "fig.savefig(figure_path(\"si_figure_s02_diels_alder.png\"), dpi=300, bbox_inches=\"tight\")\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s04.ipynb b/analysis/si_figures/si_figure_s04.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..109360d220315078f9d1fd880220f5896cbabfc1 --- /dev/null +++ b/analysis/si_figures/si_figure_s04.ipynb @@ -0,0 +1,167 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "007f5b29", + "metadata": {}, + "source": [ + "# SI Figure S4: correlation of implicit (PCM) corrections across solvents and methods\n", + "\n", + "**S4A** correlation across solvents (both nuclei), **S4B** one solvent vs another (¹H), **S4C**\n", + "correlation across methods (both nuclei). Cells show -log10(1-r), so 3 means r=0.999." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34163231", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4fced82", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import delta22_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2d316f79", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cb9e3ebb", + "metadata": {}, + "outputs": [], + "source": [ + "METHOD, BASIS, GEOM = \"b3lyp_d3bj\", \"pcSseg2\", \"pbe0_tz\"\n", + "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "base = dft[(dft[\"sap_nmr_method\"] == METHOD) & (dft[\"sap_basis\"] == BASIS)\n", + " & (dft[\"sap_geometry_type\"] == GEOM)]\n", + "one = base[base[\"nucleus\"] == \"H\"]\n", + "\n", + "# solvent order that groups polar-aprotic -> polar-protic -> aromatic (matches the published panel)\n", + "ORDERED_SOLVENTS = [\"chloroform\", \"tetrahydrofuran\", \"dichloromethane\", \"acetone\", \"acetonitrile\",\n", + " \"dimethylsulfoxide\", \"trifluoroethanol\", \"methanol\", \"TIP4P\",\n", + " \"benzene\", \"toluene\", \"chlorobenzene\"]" + ] + }, + { + "cell_type": "markdown", + "id": "ee2d20ba", + "metadata": {}, + "source": [ + "## S4A: solvent-vs-solvent PCM correlation (¹H and ¹³C)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c8664491", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n", + " solvent_corr = delta22.correlation_matrix(base[base[\"nucleus\"] == nucleus], \"pcm\", [\"solute\", \"site\"], \"solvent\")\n", + " solvent_corr = solvent_corr.reindex(index=ORDERED_SOLVENTS, columns=ORDERED_SOLVENTS)\n", + " vals = solvent_corr.values[np.triu_indices_from(solvent_corr.values, k=1)]\n", + " print(f\"solvent PCM correlation ({label}): mean r={np.nanmean(vals):.4f} min r={np.nanmin(vals):.4f}\")\n", + " caption = (\"Correlation coefficients between solvents across all 22 solutes.\\n\"\n", + " f\"PCM corrections computed with {METHOD} with the {BASIS} basis.\\n\"\n", + " \"A value of 3 means the coefficient is 0.999.\")\n", + " delta22_plots.plot_correlation_matrix(solvent_corr, f\"Solvents are Highly Correlated ({nucleus})\", caption,\n", + " colormap=\"Reds\", show_values=True,\n", + " save_path=figure_path(f\"si_figure_s04a_{label}.png\"))" + ] + }, + { + "cell_type": "markdown", + "id": "133ed792", + "metadata": {}, + "source": [ + "## S4B: PCM corrections, chloroform vs acetonitrile (1H) -- one square of S4A" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21db6810", + "metadata": {}, + "outputs": [], + "source": [ + "pair = one.pivot_table(index=[\"solute\", \"site\"], columns=\"solvent\", values=\"pcm\")[\n", + " [\"chloroform\", \"acetonitrile\"]].dropna()\n", + "delta22_plots.plot_pcm_scatter(pair[\"chloroform\"], pair[\"acetonitrile\"], \"chloroform\", \"acetonitrile\", \"H\",\n", + " save_path=figure_path(\"si_figure_s04b_1H.png\"))" + ] + }, + { + "cell_type": "markdown", + "id": "1f0e76f9", + "metadata": {}, + "source": [ + "## S4C: method-vs-method correlation (chloroform, double hybrids excluded, ¹H and ¹³C)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f75897e7", + "metadata": {}, + "outputs": [], + "source": [ + "# one solvent, exclude the double hybrids whose PCM is the substituted reference value\n", + "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n", + " chcl3 = dft[(dft[\"sap_basis\"] == BASIS) & (dft[\"sap_geometry_type\"] == GEOM)\n", + " & (dft[\"nucleus\"] == nucleus) & (dft[\"solvent\"] == \"chloroform\")\n", + " & (~dft[\"sap_nmr_method\"].isin(delta22.DOUBLE_HYBRID_METHODS))]\n", + " method_corr = delta22.correlation_matrix(chcl3, \"pcm\", [\"solute\", \"site\"], \"sap_nmr_method\")\n", + " method_order = sorted(method_corr.index) # alphabetical, matches the published panel\n", + " method_corr = method_corr.reindex(index=method_order, columns=method_order)\n", + " mvals = method_corr.values[np.triu_indices_from(method_corr.values, k=1)]\n", + " print(f\"method PCM correlation ({label}): mean r={np.nanmean(mvals):.4f} min r={np.nanmin(mvals):.4f}\")\n", + " caption = (\"Correlation coefficients between NMR methods across all 22 solutes.\\n\"\n", + " f\"PCM corrections for chloroform with the {BASIS} basis.\")\n", + " delta22_plots.plot_correlation_matrix(method_corr, f\"NMR Methods are Highly Correlated ({nucleus})\", caption,\n", + " colormap=\"Reds\", show_values=True,\n", + " save_path=figure_path(f\"si_figure_s04c_{label}.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s05.ipynb b/analysis/si_figures/si_figure_s05.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..14d88dcdf2b1653d23a499b83519b4470c2e729f --- /dev/null +++ b/analysis/si_figures/si_figure_s05.ipynb @@ -0,0 +1,124 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1344ef50", + "metadata": {}, + "source": [ + "# SI Figure S5: per-solvent PCM benefit vs bulk dielectric constant and polarizability\n", + "\n", + "Per-solvent PCM benefit (% test-RMSE reduction) vs bulk dielectric constant and polarizability, for\n", + "the MagNET-Zero reference method per nucleus (WP04 ¹H, wB97X-D ¹³C). Panels A (¹H) and B (¹³C)\n", + "show all solvents; C repeats ¹H excluding aromatics and trifluoroethanol." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dfffb6dc", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "74921dd6", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import delta22_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c3053b56", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "99b0ba9e", + "metadata": {}, + "outputs": [], + "source": [ + "# the published panels use 250 seeded train/test splits\n", + "N_SPLITS = 250" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9aa30554", + "metadata": {}, + "outputs": [], + "source": [ + "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "solutes = sorted(dft[\"solute\"].unique())\n", + "benefit = {}\n", + "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n", + " method = delta22.MAGNET_PCM_OUTPUT_METHODS[nucleus]\n", + " benefit[nucleus] = delta22.pcm_benefit_per_solvent(dft, method, \"pcSseg2\", \"aimnet2\",\n", + " delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS,\n", + " solutes=solutes, nucleus=nucleus)\n", + " print(f\"{label} ({method}):\"); print(benefit[nucleus].round(1).to_string())" + ] + }, + { + "cell_type": "markdown", + "id": "329088d2", + "metadata": {}, + "source": [ + "## S5A (1H, all solvents), S5B (13C, all solvents), and S5C (1H, excluding aromatics and trifluoroethanol)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2d519d66", + "metadata": {}, + "outputs": [], + "source": [ + "# panels A (1H) and B (13C): all solvents\n", + "for nucleus, letter, nuc_label in [(\"H\", \"A\", \"H\"), (\"C\", \"B\", \"C\")]:\n", + " delta22_plots.plot_pcm_benefit_vs_properties(\n", + " benefit[nucleus], delta22.SOLVENT_DIELECTRIC, delta22.SOLVENT_POLARIZABILITY, nuc_label,\n", + " save_path=figure_path(f\"si_figure_s05{letter}_all.png\"))\n", + "# panel C is 1H only, excluding the aromatic solvents and trifluoroethanol (matches the published SI)\n", + "delta22_plots.plot_pcm_benefit_vs_properties(\n", + " benefit[\"H\"], delta22.SOLVENT_DIELECTRIC, delta22.SOLVENT_POLARIZABILITY, \"H\",\n", + " exclude=[\"benzene\", \"toluene\", \"chlorobenzene\", \"trifluoroethanol\"],\n", + " title_extra=\"Excluding Aromatic Solvents and Trifluoroethanol\",\n", + " save_path=figure_path(\"si_figure_s05A_no_aromatics_tfe.png\"))" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s06.ipynb b/analysis/si_figures/si_figure_s06.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..001987ee9a0212ec7871d70bf5541a0a56dbab1d --- /dev/null +++ b/analysis/si_figures/si_figure_s06.ipynb @@ -0,0 +1,152 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "29f74212", + "metadata": {}, + "source": [ + "# SI Figure S6: per-solvent test RMSE of six solvent-correction models\n", + "\n", + "Per nucleus, the per-solvent test RMSE of six solvent-correction models: SotA implicit (single-slope\n", + "PCM), the paper's implicit fit (stationary + PCM as separate terms), and explicit (Desmond), each\n", + "with and without a vibrational correction (QCD for ¹H, Desmond vibration for ¹³C). Level of theory:\n", + "dsd_pbep86/pcSseg3 on pbe0_tz geometries, PCM substituted from b3lyp_d3bj/pcSseg3." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6871225d", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3bad055e", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import delta22_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a55272ae", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "221a152c", + "metadata": {}, + "outputs": [], + "source": [ + "query = delta22.add_composite_columns(delta22.load_query_df_dft(\n", + " DELTA22_HDF5, XLSX, pcm_reference_method=\"b3lyp_d3bj\", pcm_reference_basis=\"pcSseg3\", verbose=False))\n", + "solutes = sorted(query[\"solute\"].unique())\n", + "\n", + "METHOD, BASIS, GEOM = \"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"\n", + "\n", + "# same six categories for both nuclei; the \"+ Vibrations\" term differs (qcd for H, desmond_vib for C)\n", + "LADDER = {\n", + " \"H\": {\n", + " \"stationary_plus_pcm\": \"SotA Implicit Solvent\",\n", + " \"stationary_plus_pcm_plus_qcd\": \"SotA Implicit Solvent + Vibrations\",\n", + " \"stationary + pcm\": \"Implicit Solvent\",\n", + " \"stationary_plus_qcd + pcm\": \"Implicit Solvent + Vibrations\",\n", + " \"stationary + desmond\": \"Explicit Solvent\",\n", + " \"stationary_plus_qcd + desmond\": \"Explicit Solvent + Vibrations\"},\n", + " \"C\": {\n", + " \"stationary_plus_pcm\": \"SotA Implicit Solvent\",\n", + " \"stationary_plus_pcm_plus_des_vib\": \"SotA Implicit Solvent + Vibrations\",\n", + " \"stationary + pcm\": \"Implicit Solvent\",\n", + " \"stationary_plus_des_vib + pcm\": \"Implicit Solvent + Vibrations\",\n", + " \"stationary + desmond\": \"Explicit Solvent\",\n", + " \"stationary_plus_des_vib + desmond\": \"Explicit Solvent + Vibrations\"},\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d1b678b5", + "metadata": {}, + "outputs": [], + "source": [ + "# one dark/light pair per solvent-treatment family (SotA implicit / implicit / explicit), light = +Vibrations\n", + "S6_COLORS = [\"#C4B037\", \"#F5EDA0\", \"#A72608\", \"#E4A0A0\", \"#5F7C8A\", \"#B9E7DF\"]" + ] + }, + { + "cell_type": "markdown", + "id": "53817ec5", + "metadata": {}, + "source": [ + "## Formula ladder per nucleus" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6cd062c4", + "metadata": {}, + "outputs": [], + "source": [ + "# the published panels use 250 seeded train/test splits\n", + "N_SPLITS = 250" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d57a9788", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus, labels in LADDER.items():\n", + " results = delta22.fig3d_formula_regressions(query, METHOD, BASIS, GEOM, list(labels),\n", + " delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS,\n", + " solutes=solutes, nucleus=nucleus)\n", + " med = results.groupby(\"formula\")[\"test_RMSE\"].median()\n", + " print(f\"--- {nucleus} ({METHOD}) median test RMSE ---\")\n", + " for formula in labels:\n", + " print(f\" {labels[formula]:38s} {med[formula]:.4f}\")\n", + " delta22_plots.plot_formula_ladder_boxplot(\n", + " results, labels, delta22.SOLVENT_GROUPS, nucleus=nucleus, colors=S6_COLORS,\n", + " title=f\"Explicit Solvation + Vibrational Effects Improve Predictions Across Solvents ({nucleus} Nucleus)\",\n", + " save_path=figure_path(f\"si_figure_s06_ladder_{'1H' if nucleus == 'H' else '13C'}.png\"))\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s07.ipynb b/analysis/si_figures/si_figure_s07.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..1662dde9206119839523ea8523444a84d44043c4 --- /dev/null +++ b/analysis/si_figures/si_figure_s07.ipynb @@ -0,0 +1,100 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "99d967bc", + "metadata": {}, + "source": [ + "# SI Figure S7: DFT ¹³C explicit-solvent correction, Desmond vs OpenMM\n", + "\n", + "For the four OpenMM solvents, the DFT ¹³C explicit-solvent correction from Desmond MD vs from OpenMM\n", + "(y=x reference line)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b506b1dc", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "97266e80", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import delta22_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "01aca94d", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f78afb6a", + "metadata": {}, + "outputs": [], + "source": [ + "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)" + ] + }, + { + "cell_type": "markdown", + "id": "b83900b6", + "metadata": {}, + "source": [ + "## DFT explicit corrections: Desmond vs OpenMM" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23f9e6f4", + "metadata": {}, + "outputs": [], + "source": [ + "SOLVENT_ORDER = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]\n", + "dft_pairs = {sv: delta22.explicit_correction_pairs(dft, sv, \"C\") for sv in SOLVENT_ORDER}\n", + "for sv, p in dft_pairs.items():\n", + " print(f\"DFT {sv:12s} n={len(p):3d} r={np.corrcoef(p['desmond'], p['openMM'])[0,1]:.3f}\")\n", + "delta22_plots.plot_desmond_vs_openmm_grid(dft_pairs, save_path=figure_path(\"si_figure_s07.png\"))\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s08_panelA_vomicine.ipynb b/analysis/si_figures/si_figure_s08_panelA_vomicine.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..1c8acca09dc5f9623c64fed876daa9056edcae49 --- /dev/null +++ b/analysis/si_figures/si_figure_s08_panelA_vomicine.ipynb @@ -0,0 +1,121 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "eef9d8ce", + "metadata": {}, + "source": [ + "# SI Figure S8 (panel A): explicit-solvent correction vs solvent-shell size (vomicine)\n", + "\n", + "The MagNET-x explicit-solvent correction (solvated minus isolated, averaged over a proton site's\n", + "atoms and MD frames) vs solvent-shell size (heavy-atom count of the solute-plus-solvent cluster),\n", + "for vomicine, one panel per solvent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "797c7b29", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/applications\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f174bf0", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "from applications_reader import Applications, SHELL_SIZES\n", + "import applications\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d616ff60", + "metadata": {}, + "outputs": [], + "source": [ + "DATA = os.path.join(REPO, \"data\", \"applications\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eab9a523", + "metadata": {}, + "outputs": [], + "source": [ + "loader = Applications(paths.dataset_file(\"applications\", root=REPO),\n", + " os.path.join(DATA, \"applications_experimental.xlsx\"))\n", + "\n", + "# compute the per-site, per-solvent shell-convergence corrections for vomicine (proton sites)\n", + "SOLUTE, NUCLEUS = \"vomicine\", \"H\"\n", + "corrections = applications.shell_convergence_corrections(loader, SOLUTE, NUCLEUS)\n", + "corrections.head()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f7d00a0", + "metadata": {}, + "outputs": [], + "source": [ + "# panel layout and font styling for the 2x2 grid (one panel per solvent)\n", + "solvent_order = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]\n", + "shell_cols = [f\"shell_{s}\" for s in SHELL_SIZES]\n", + "plt.rcParams.update({\"font.family\": \"serif\", \"font.serif\": [\"Arial\", \"DejaVu Serif\"],\n", + " \"axes.labelsize\": 12, \"axes.titlesize\": 14})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0350cca5", + "metadata": {}, + "outputs": [], + "source": [ + "# 2x2 panel, one solvent each; one line per proton site\n", + "fig, axes = plt.subplots(2, 2, figsize=(11, 8), sharex=True, sharey=True)\n", + "for ax, solvent in zip(axes.ravel(), solvent_order):\n", + " sub = corrections.xs(solvent, level=\"solvent\")\n", + " sub = sub.dropna(how=\"all\")\n", + " ax.set_title(\"TIP4P\" if solvent == \"TIP4P\" else solvent.capitalize(), fontweight=\"bold\")\n", + " for site, row in sub.iterrows():\n", + " ax.plot(SHELL_SIZES, row[shell_cols].values.astype(float),\n", + " marker=\"o\", linewidth=1.5, alpha=0.75, label=site)\n", + " ax.grid(True, alpha=0.3)\n", + " ax.set_xlabel(\"Shell Size (Heavy Atom Count)\", fontweight=\"bold\")\n", + " ax.set_ylabel(\"OpenMM correction (ppm)\", fontweight=\"bold\")\n", + "fig.suptitle(\"Shell-size Convergence for Vomicine\", fontweight=\"bold\", fontsize=18)\n", + "fig.tight_layout()\n", + "fig.savefig(figure_path(\"si_figure_s08_vomicine.png\"), dpi=300, bbox_inches=\"tight\")\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s08_panelsBCD.ipynb b/analysis/si_figures/si_figure_s08_panelsBCD.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..030b0184c5a4ddefb75353136339b7e6f33b7710 --- /dev/null +++ b/analysis/si_figures/si_figure_s08_panelsBCD.ipynb @@ -0,0 +1,248 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d45e897a", + "metadata": {}, + "source": [ + "# SI Figure S8 (panels B-D): explicit-solvent correction vs MD-frame count\n", + "\n", + "For an example site (AcOH proton in chloroform), the explicit-solvent correction vs number of MD frames included (running average): OpenMM vs Desmond (B1) and DFT vs MagNET-x/\"NN\" at fixed OpenMM (B3), with per-frame distributions (B2/B4) and autocorrelation (C). Panel D shows which frames have computed DFT data per solvent (also the source of B's histogram normalization)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e0fcff5a", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "62fa04d3", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import delta22_reader\n", + "import delta22_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5ba77909", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a93dfea1", + "metadata": {}, + "outputs": [], + "source": [ + "site_atoms = delta22_reader.load_site_atom_data(XLSX, verbose=False)\n", + "idx = delta22.site_atom_indices(site_atoms.loc[(\"AcOH\", \"H\", \"H\"), \"atom_numbers\"])\n", + "SOLUTE, SOLVENT = \"AcOH\", \"chloroform\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f50407ff", + "metadata": {}, + "outputs": [], + "source": [ + "# OpenMM/Desmond engine hues; DFT vs NN reuse the OpenMM hue (DFT lightened, NN full strength)\n", + "ENGINE_COLORS = {\"openMM\": \"#2E86AB\", \"desmond\": \"#A23B72\"}\n", + "ENGINE_LABELS = {\"openMM\": \"OpenMM\", \"desmond\": \"Desmond\"}\n", + "\n", + "SOURCE_COLORS = {\"dft\": delta22_plots.lighten_color(ENGINE_COLORS[\"openMM\"]), \"nn\": ENGINE_COLORS[\"openMM\"]}\n", + "SOURCE_LABELS = {\"dft\": \"DFT\", \"nn\": \"NN\"}\n", + "LABEL_COLORS = {**ENGINE_COLORS, **SOURCE_COLORS}\n", + "LABEL_NAMES = {**ENGINE_LABELS, **SOURCE_LABELS}" + ] + }, + { + "cell_type": "markdown", + "id": "398912a2", + "metadata": {}, + "source": [ + "## Panel B1/B2: OpenMM vs Desmond\n", + "\n", + "Running-average convergence (B1) and per-frame distribution (B2), comparing the two MD engines." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32f05604", + "metadata": {}, + "outputs": [], + "source": [ + "running_by_engine, finals_by_engine, per_frame_by_engine, n_frames_by_engine = {}, {}, {}, {}\n", + "for engine in [\"openMM\", \"desmond\"]:\n", + " perturbed = delta22_reader.load_perturbed_shieldings(DELTA22_HDF5, SOLUTE, SOLVENT, engine, \"dft\")\n", + " per_frame = delta22.frame_corrections(perturbed, idx)\n", + " running = delta22.running_average(per_frame)\n", + " running_by_engine[engine] = running\n", + " finals_by_engine[engine] = running[~np.isnan(running)][-1]\n", + " per_frame_by_engine[engine] = per_frame\n", + " n_frames_by_engine[engine] = len(per_frame) # the total trajectory length, valid or not\n", + " print(f\"{engine}: {int(np.sum(~np.isnan(per_frame)))} / {len(per_frame)} valid frames, \"\n", + " f\"converged correction {finals_by_engine[engine]:.4f} ppm\")\n", + "\n", + "delta22_plots.plot_frame_convergence(\n", + " running_by_engine, finals_by_engine, LABEL_COLORS, LABEL_NAMES,\n", + " title=f\"Convergence of Explicit Solvent Corrections\\n({SOLUTE} in {SOLVENT})\",\n", + " save_path=figure_path(\"si_figure_s08_b1_convergence_engine.png\"))\n", + "plt.show()\n", + "\n", + "delta22_plots.plot_frame_correction_histogram(\n", + " per_frame_by_engine, n_frames_by_engine, LABEL_COLORS, LABEL_NAMES,\n", + " title=\"Distribution of Frame-wise Corrections\",\n", + " save_path=figure_path(\"si_figure_s08_b2_histogram_engine.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "be656fba", + "metadata": {}, + "source": [ + "## Panel B3/B4: DFT vs NN\n", + "\n", + "Same running average (B3) and per-frame distribution (B4), at fixed OpenMM engine, comparing DFT vs\n", + "MagNET-x (\"NN\") shieldings." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9f996246", + "metadata": {}, + "outputs": [], + "source": [ + "running_by_source, finals_by_source, per_frame_by_source, n_frames_by_source = {}, {}, {}, {}\n", + "for source in [\"dft\", \"nn\"]:\n", + " perturbed = delta22_reader.load_perturbed_shieldings(DELTA22_HDF5, SOLUTE, SOLVENT, \"openMM\", source)\n", + " per_frame = delta22.frame_corrections(perturbed, idx)\n", + " running = delta22.running_average(per_frame)\n", + " running_by_source[source] = running\n", + " finals_by_source[source] = running[~np.isnan(running)][-1]\n", + " per_frame_by_source[source] = per_frame\n", + " n_frames_by_source[source] = len(per_frame)\n", + " print(f\"{source}: {int(np.sum(~np.isnan(per_frame)))} / {len(per_frame)} valid frames, \"\n", + " f\"converged correction {finals_by_source[source]:.4f} ppm\")\n", + "\n", + "delta22_plots.plot_frame_convergence(\n", + " running_by_source, finals_by_source, LABEL_COLORS, LABEL_NAMES, xlabel=\"Number of MD Frames\",\n", + " title=f\"OpenMM DFT vs. NN Convergence Comparison\\n({SOLUTE} in {SOLVENT})\",\n", + " save_path=figure_path(\"si_figure_s08_b3_convergence_source.png\"))\n", + "plt.show()\n", + "\n", + "delta22_plots.plot_frame_correction_histogram(\n", + " per_frame_by_source, n_frames_by_source, LABEL_COLORS, LABEL_NAMES,\n", + " title=\"Distribution of Frame-wise Corrections\",\n", + " save_path=figure_path(\"si_figure_s08_b4_histogram_source.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "483a55bb", + "metadata": {}, + "source": [ + "## Panel C: autocorrelation, DFT vs NN\n", + "\n", + "Autocorrelation of the per-frame correction to lag 200 (~100 frames per trajectory repeat), DFT vs NN,\n", + "same OpenMM trajectory as B3/B4." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b365dfe", + "metadata": {}, + "outputs": [], + "source": [ + "autocorr_by_source = {source: delta22.autocorrelation(per_frame_by_source[source], max_lag=200)\n", + " for source in [\"dft\", \"nn\"]}\n", + "for source, autocorr in autocorr_by_source.items():\n", + " print(f\"{source}: lag-1 autocorrelation {autocorr[1]:.3f}\")\n", + "\n", + "delta22_plots.plot_frame_autocorrelation(\n", + " autocorr_by_source, LABEL_COLORS, LABEL_NAMES,\n", + " title=f\"Autocorrelation of Frame-wise Corrections\\n({SOLUTE} in {SOLVENT})\",\n", + " save_path=figure_path(\"si_figure_s08_c_autocorrelation.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "bc704f01", + "metadata": {}, + "source": [ + "## Panel D: frame data availability\n", + "\n", + "Which trajectory frames have computed DFT shielding for AcOH, across all 12 Desmond and 4 OpenMM\n", + "solvents. DFT was computed for a non-contiguous subset (compute-budget limits; jobs queued randomly)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ece733ef", + "metadata": {}, + "outputs": [], + "source": [ + "# row order matching the published panel: chloroform first, then by solvent class (aprotic ->\n", + "# protic -> aromatic); OpenMM only has the four explicit-solvent solvents\n", + "DESMOND_ORDER = [\"chloroform\", \"tetrahydrofuran\", \"dichloromethane\", \"acetone\", \"acetonitrile\",\n", + " \"dimethylsulfoxide\", \"trifluoroethanol\", \"methanol\", \"TIP4P\",\n", + " \"benzene\", \"toluene\", \"chlorobenzene\"]\n", + "OPENMM_ORDER = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]\n", + "\n", + "grids_by_engine, solvents_by_engine = {}, {}\n", + "for engine, solvents in [(\"desmond\", DESMOND_ORDER), (\"openMM\", OPENMM_ORDER)]:\n", + " grid, frame_counts = delta22.frame_validity_grid(DELTA22_HDF5, SOLUTE, solvents, engine, shield_type=\"dft\")\n", + " grids_by_engine[engine] = grid\n", + " solvents_by_engine[engine] = [delta22_plots.display_solvent_name(s) for s in solvents]\n", + " print(f\"{engine}: {len(solvents)} solvents, frame counts {dict(zip(solvents, frame_counts))}\")\n", + "\n", + "delta22_plots.plot_frame_validity_heatmaps(grids_by_engine, solvents_by_engine, ENGINE_LABELS, solute=SOLUTE,\n", + " save_path=figure_path(\"si_figure_s08_d_frame_validity.png\"))\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s15.ipynb b/analysis/si_figures/si_figure_s15.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..5df3df17a347e3bb2789ae4d31fb3c10b6d2ca3a --- /dev/null +++ b/analysis/si_figures/si_figure_s15.ipynb @@ -0,0 +1,286 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a8256e57", + "metadata": {}, + "source": [ + "# SI Figure S15: per-solute RMSE with delta-22 composite coefficients applied\n", + "\n", + "Per-solute ¹H and ¹³C RMSE for the natural-products / olefin-isomer test set with delta-22\n", + "coefficients applied (reproduces Figure 5C in ¹H, plus the ¹³C analogue), plus fitting-RMSE\n", + "comparisons across coefficient choices, feature-space coverage by solvent, delta-22-plane residuals,\n", + "and the RMSE distribution shift by solvent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64581ad1", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/applications\", \"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "206af146", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "from applications_reader import Applications\n", + "import applications\n", + "import applications_plots\n", + "import delta22\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69c8b5ee", + "metadata": {}, + "outputs": [], + "source": [ + "DATA = os.path.join(REPO, \"data\", \"applications\")\n", + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "DELTA22_XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9cde73dc", + "metadata": {}, + "outputs": [], + "source": [ + "loader = Applications(paths.dataset_file(\"applications\", root=REPO),\n", + " os.path.join(DATA, \"applications_experimental.xlsx\"))\n", + "\n", + "# assemble the feature table and run the composite-model fits + bootstrap\n", + "query_df_nn = applications.build_query_df_nn(loader)\n", + "site_counts = (query_df_nn.drop_duplicates(subset=[\"solute\", \"nucleus\", \"site\"])\n", + " .groupby([\"solute\", \"nucleus\"]).size().unstack(fill_value=0))\n", + "\n", + "per_solute = applications.per_solute_fits(query_df_nn) # \"scaled to solute\" baseline\n", + "all_solute = applications.per_solvent_fits(query_df_nn) # \"scaled to test set\" full fit\n", + "\n", + "seed = applications.build_bootstrap_seed_coeffs(loader)\n", + "bootstrap_rmses = {}\n", + "for nuc in [\"H\", \"C\"]:\n", + " preds = applications.apply_bootstrap_params_to_full_dataset(query_df_nn, seed[nuc], nucleus=nuc)\n", + " bootstrap_rmses[nuc] = applications.compute_solute_rmses(preds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d218de52", + "metadata": {}, + "outputs": [], + "source": [ + "# published solute order (isomers, then the two pyridines, then the natural products); matches\n", + "# main-text Figure 5D. delta-22 is drawn first by the box-plot engine, so it is not listed here.\n", + "S15_SOLUTE_ORDER = [\"isomer_1E\", \"isomer_1Z\", \"isomer_2E\", \"isomer_2Z\", \"isomer_3E\", \"isomer_3Z\",\n", + " \"isomer_4N\", \"isomer_4O\", \"vomicine\", \"prednisone\", \"peptide\", \"flavone\",\n", + " \"dihydrotanshinone_I\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "85097f2d", + "metadata": {}, + "outputs": [], + "source": [ + "# 1H panel (reproduces Figure 5C)\n", + "applications_plots.plot_nps_on_boxplot_delta22_simplified(\n", + " loader.rmse_distribution(\"H\"), bootstrap_rmses[\"H\"], per_solute[\"H\"],\n", + " nucleus=\"H\", formulas=[\"stationary_plus_qcd + openMM\"], colors=[\"#61a89a\"],\n", + " site_counts=site_counts, formula_remap=applications.FORMULA_REMAP,\n", + " solute_remap=applications.SOLUTE_DISPLAY, solute_order=S15_SOLUTE_ORDER,\n", + " solute_color_remap=applications.PEPTIDE_HIGHLIGHT_H,\n", + " figsize=(14, 8), box_width=0.20, box_gap=0.1,\n", + " show_baseline=True, baseline_annotation_text=\"Scaled to solute\",\n", + " baseline_annotation_x=0.215, baseline_annotation_y=0.39,\n", + " show_full_fit_line=True, full_fit_line_label=\"Scaled to Test Set\", full_fit_line_label_x=0.835,\n", + " all_solute_fitting_results=all_solute,\n", + " max_bar_height=0.06, site_count_axis_mode=\"inset\", site_count_inset_area_frac=0.1,\n", + " site_count_inset_axis_offset=-0.06, site_count_inset_axis_label_pad=0.045,\n", + " save_path=figure_path(\"si_figure_s15_1H.png\"),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "973af7b2", + "metadata": {}, + "outputs": [], + "source": [ + "# 13C panel\n", + "applications_plots.plot_nps_on_boxplot_delta22_simplified(\n", + " loader.rmse_distribution(\"C\"), bootstrap_rmses[\"C\"], per_solute[\"C\"],\n", + " nucleus=\"C\", formulas=[\"stationary_plus_op_vib + openMM\"], colors=[\"#61a89a\"],\n", + " site_counts=site_counts, formula_remap=applications.FORMULA_REMAP,\n", + " solute_remap=applications.SOLUTE_DISPLAY, solute_order=S15_SOLUTE_ORDER,\n", + " solute_color_remap=applications.PEPTIDE_HIGHLIGHT_C,\n", + " figsize=(14, 8), box_width=0.20, box_gap=0.1,\n", + " show_baseline=True, baseline_annotation_text=\"Scaled to solute\",\n", + " baseline_annotation_x=0.28, baseline_annotation_y=0.355,\n", + " show_full_fit_line=True, full_fit_line_label=\"Scaled to Test Set\", full_fit_line_label_x=0.98,\n", + " all_solute_fitting_results=all_solute,\n", + " max_bar_height=0.8, site_count_axis_mode=\"inset\", site_count_inset_area_frac=0.1,\n", + " site_count_inset_axis_offset=-0.06, site_count_inset_axis_label_pad=0.045,\n", + " save_path=figure_path(\"si_figure_s15_13C.png\"),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "f079c075", + "metadata": {}, + "source": [ + "## Fitting RMSE Comparisons (chloroform)\n", + "\n", + "Per-solute RMSE under three coefficient choices (Scaled to Solute / Scaled to Test Set / Extrapolated\n", + "from delta22)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11ea3200", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus, formula in [(\"H\", \"stationary_plus_qcd + openMM\"), (\"C\", \"stationary_plus_op_vib + openMM\")]:\n", + " table = applications.fitting_rmse_comparison_table(query_df_nn, per_solute, all_solute, bootstrap_rmses[nucleus],\n", + " nucleus, \"chloroform\", formula)\n", + " table = table.rename(index=applications.SOLUTE_DISPLAY)\n", + " order = [applications.SOLUTE_DISPLAY.get(k, k) for k in S15_SOLUTE_ORDER] # published order; peptide has no chloroform data\n", + " table = table.reindex([n for n in order if n in table.index])\n", + " print(f\"--- {nucleus} ---\")\n", + " display(table.round(3))\n", + " applications_plots.plot_fitting_rmse_comparison_bars(\n", + " table, nucleus, \"chloroform\",\n", + " save_path=figure_path(f\"si_figure_s15_fitting_rmse_comparisons_{'1H' if nucleus == 'H' else '13C'}.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "47773c94", + "metadata": {}, + "source": [ + "## Feature Space Coverage by Solvent\n", + "\n", + "Test-set vs delta-22 mean-centered feature values across the four explicit-solvent solvents." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa1de7fe", + "metadata": {}, + "outputs": [], + "source": [ + "delta22_query_df_nn = delta22.add_composite_columns(delta22.load_query_df_nn(\n", + " DELTA22_HDF5, DELTA22_XLSX, verbose=False))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a19b095b", + "metadata": {}, + "outputs": [], + "source": [ + "_FEATURE_X_LABELS = {\n", + " \"H\": \"Gas-Phase Shielding + QCD Correction (centered, ppm)\",\n", + " \"C\": \"Gas-Phase Shielding + OpenMM Vibrational Correction (centered, ppm)\",\n", + "}\n", + "for nucleus in [\"H\", \"C\"]:\n", + " coverage = applications.feature_space_coverage_table(query_df_nn, delta22_query_df_nn, nucleus)\n", + " applications_plots.plot_feature_space_coverage_grid(\n", + " coverage, nucleus, _FEATURE_X_LABELS[nucleus],\n", + " save_path=figure_path(f\"si_figure_s15_feature_space_coverage_{'1H' if nucleus == 'H' else '13C'}.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "5d7bc6b3", + "metadata": {}, + "source": [ + "## Residuals for Delta22 Fitting Coefficients\n", + "\n", + "Residuals of a delta-22-only 2-feature plane applied to both delta-22 and the test set, vs\n", + "experimental shielding." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8dea4de5", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus in [\"H\", \"C\"]:\n", + " residuals = applications.delta22_plane_residuals_table(query_df_nn, delta22_query_df_nn, nucleus)\n", + " applications_plots.plot_delta22_plane_residuals_grid(\n", + " residuals, nucleus,\n", + " save_path=figure_path(f\"si_figure_s15_delta22_plane_residuals_{'1H' if nucleus == 'H' else '13C'}.png\"))\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "3fd43b43", + "metadata": {}, + "source": [ + "## Distribution Shift by Solvent\n", + "\n", + "Per-solvent mean test-set RMSE under the three coefficient choices." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24d4b26e", + "metadata": {}, + "outputs": [], + "source": [ + "for nucleus, formula in [(\"H\", \"stationary_plus_qcd + openMM\"), (\"C\", \"stationary_plus_op_vib + openMM\")]:\n", + " shift = applications.distribution_shift_by_solvent_table(per_solute, all_solute, bootstrap_rmses[nucleus],\n", + " nucleus, formula)\n", + " print(f\"--- {nucleus} ---\")\n", + " display(shift.round(3))\n", + " applications_plots.plot_distribution_shift_by_solvent_bars(\n", + " shift, nucleus,\n", + " save_path=figure_path(f\"si_figure_s15_distribution_shift_by_solvent_{'1H' if nucleus == 'H' else '13C'}.png\"))\n", + "plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_figures/si_figure_s16_s18.ipynb b/analysis/si_figures/si_figure_s16_s18.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..ef01eb1a703ffdc7f8c93cd7f642aedce1843b0f --- /dev/null +++ b/analysis/si_figures/si_figure_s16_s18.ipynb @@ -0,0 +1,143 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "36541c20", + "metadata": {}, + "source": [ + "# SI Figures S16-S18: predicted vs experimental solvent-induced shifts by reference solvent\n", + "\n", + "Per solvent, implicit (PCM) and explicit (Desmond) predicted solvent-induced shift differences vs\n", + "measured (y=x ideal), for three reference solvents: **S16** chloroform, **S17** benzene, **S18**\n", + "solvent-averaged. One panel per solvent, each point a proton site." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8464e639", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "354b6ae2", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "import delta22\n", + "import delta22_plots\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1b63be9a", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def figure_path(name):\n", + " os.makedirs(\"figures\", exist_ok=True)\n", + " return os.path.join(\"figures\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "710b94a1", + "metadata": {}, + "outputs": [], + "source": [ + "# Figure 4 uses one DFT method for the proton sites; differences are taken against a reference\n", + "# solvent, and the solvent_mean pseudo-solvent is added for the solvent-averaged reference.\n", + "METHOD, BASIS, GEOM = \"b3lyp_d3bj\", \"pcSseg2\", \"aimnet2\"\n", + "q = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "one = q[(q[\"sap_nmr_method\"] == METHOD) & (q[\"sap_basis\"] == BASIS) & (q[\"sap_geometry_type\"] == GEOM)]\n", + "one = delta22.add_solvent_mean(one)\n", + "print(len(one), \"rows for\", METHOD, BASIS, GEOM)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "94ad7330", + "metadata": {}, + "outputs": [], + "source": [ + "FIG4C_LABELS = {\n", + " \"chloroform\": \"CDCl3\", \"dichloromethane\": \"DCM\", \"tetrahydrofuran\": \"THF\",\n", + " \"acetonitrile\": \"MeCN\", \"dimethylsulfoxide\": \"DMSO\", \"acetone\": \"acetone\",\n", + " \"methanol\": \"MeOH\", \"TIP4P\": \"TIP4P\", \"trifluoroethanol\": \"TFE\",\n", + " \"benzene\": \"benzene\", \"toluene\": \"toluene\", \"chlorobenzene\": \"chlorobenzene\",\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ecc3f9be", + "metadata": {}, + "outputs": [], + "source": [ + "# panel order matches the canonical figures; a different order from delta22.DESMOND_SOLVENTS\n", + "SI_S16_18_SOLVENT_ORDER = [\n", + " \"chloroform\", \"tetrahydrofuran\", \"dichloromethane\", \"acetone\", \"acetonitrile\",\n", + " \"dimethylsulfoxide\", \"trifluoroethanol\", \"methanol\", \"TIP4P\", \"benzene\", \"toluene\", \"chlorobenzene\",\n", + "]\n", + "\n", + "# panel titles use the full solvent name; axis labels use the compact NMR abbreviations (THF, DCM,\n", + "# DMSO, TFE, CDCl3); the reference is CDCl3 / Benzene / \"Avg Corr.\" for S16 / S17 / S18.\n", + "TITLE_LABEL = {\"chloroform\": \"Chloroform\", \"tetrahydrofuran\": \"Tetrahydrofuran\",\n", + " \"dichloromethane\": \"Dichloromethane\", \"acetone\": \"Acetone\", \"acetonitrile\": \"Acetonitrile\",\n", + " \"dimethylsulfoxide\": \"Dimethylsulfoxide\", \"trifluoroethanol\": \"Trifluoroethanol\",\n", + " \"methanol\": \"Methanol\", \"TIP4P\": \"TIP4P\", \"benzene\": \"Benzene\", \"toluene\": \"Toluene\",\n", + " \"chlorobenzene\": \"Chlorobenzene\"}\n", + "AXIS_LABEL = {**TITLE_LABEL, \"chloroform\": \"CDCl3\", \"tetrahydrofuran\": \"THF\",\n", + " \"dichloromethane\": \"DCM\", \"dimethylsulfoxide\": \"DMSO\", \"trifluoroethanol\": \"TFE\"}\n", + "REF_LABEL = {\"chloroform\": \"CDCl3\", \"benzene\": \"Benzene\", \"solvent_mean\": \"Avg Corr.\"}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dea8a27d", + "metadata": {}, + "outputs": [], + "source": [ + "for figure, reference in [(\"S16\", \"chloroform\"), (\"S17\", \"benzene\"), (\"S18\", \"solvent_mean\")]:\n", + " sh = delta22.solvent_induced_shifts(one, reference, SI_S16_18_SOLVENT_ORDER, nucleus=\"H\", explicit=\"desmond\")\n", + " implicit_fit = delta22.fit_differences_to_experimental(sh, \"implicit_diff\")\n", + " explicit_fit = delta22.fit_differences_to_experimental(sh, \"explicit_diff\")\n", + " print(f\"{figure}: reference={reference:12s} n={explicit_fit['n']:4d} \"\n", + " f\"implicit fit RMSE={implicit_fit['rmse']:.3f} explicit fit RMSE={explicit_fit['rmse']:.3f}\")\n", + " panel_solvents = [s for s in SI_S16_18_SOLVENT_ORDER if s != reference]\n", + " delta22_plots.plot_shift_prediction_scatter_grid(\n", + " sh, panel_solvents, REF_LABEL[reference], AXIS_LABEL, TITLE_LABEL,\n", + " save_path=figure_path(f\"si_figure_{figure.lower()}.png\"))\n", + " plt.show()" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_tables/si_table_s01_s02_pareto.ipynb b/analysis/si_tables/si_table_s01_s02_pareto.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..2eb957ce452f8debc90fbc37093d037527d19476 --- /dev/null +++ b/analysis/si_tables/si_table_s01_s02_pareto.ipynb @@ -0,0 +1,125 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0d263621", + "metadata": {}, + "source": [ + "# SI Tables S1 and S2: the Pareto plot data\n", + "\n", + "Each method's total compute time and CDCl3 test RMSE, for proton (S1) and carbon (S2): a curated\n", + "55-row subset (MagNET, one AIMNet2-geometry row, and the DFT grid on PBE0/cc-pVTZ geometries)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b29adee9", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dc1b81aa", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import delta22\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1045389a", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def document_path(name):\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8673cb6b", + "metadata": {}, + "outputs": [], + "source": [ + "# Load the flat DFT and MagNET tables and the timing tables.\n", + "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "nn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)\n", + "dft_gas_timings, nn_timings = delta22.load_pareto_timings(DELTA22_HDF5)\n", + "\n", + "# One point per method/basis/geometry/nucleus/solvent, plus the solvent-averaged rows, each with its\n", + "# mean test RMSE and total compute time. MagNET appears as a single method (aimnet2, basis \"N/A\").\n", + "# Figure 2A uses 100 seeded splits, not the 250 the other delta-22 panels use, and trains on the\n", + "# first 10 shuffled solutes / tests on the rest; fig2a_pareto_points handles that split convention.\n", + "PARETO_N_SPLITS = 100\n", + "points = delta22.fig2a_pareto_points(dft, nn, dft_gas_timings, nn_timings, n_splits=PARETO_N_SPLITS)\n", + "print(points[\"nmr_method\"].nunique(), \"methods;\",\n", + " \"MagNET total time\", float(points.query(\"nmr_method=='MagNET'\")[\"total_time\"].iloc[0]), \"s\")" + ] + }, + { + "cell_type": "markdown", + "id": "19d525a3", + "metadata": {}, + "source": [ + "## Tables S1 and S2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3cb7258", + "metadata": {}, + "outputs": [], + "source": [ + "# full column set matching the SI tables: identity + fitting RMSE + the three time components and\n", + "# their log10 (all solvents here are chloroform). geometry_time and nmr_time sum to total_time.\n", + "COLUMNS = [\"geometry_type\", \"nmr_method\", \"basis\", \"solvent\", \"fitting_RMSE\",\n", + " \"geometry_time\", \"nmr_time\", \"total_time\", \"log10_total_time\"]\n", + "\n", + "curated = {}\n", + "for nucleus, label in [(\"H\", \"S1\"), (\"C\", \"S2\")]:\n", + " sub = delta22.pareto_table_curated(points, nucleus).copy()\n", + " sub[\"log10_total_time\"] = np.log10(sub[\"total_time\"])\n", + " sub = sub[COLUMNS]\n", + " curated[label] = sub\n", + " print(f\"Table {label} ({nucleus}): {len(sub)} rows x {sub.shape[1]} columns\")\n", + " print(sub.to_string(index=False), \"\\n\")\n", + "\n", + "# write the two tables to this notebook's documents/ folder, one sheet per SI table\n", + "out = document_path(\"si_table_s01_s02_pareto.xlsx\")\n", + "with pd.ExcelWriter(out) as writer:\n", + " curated[\"S1\"].to_excel(writer, sheet_name=\"Table S1 (1H)\", index=False)\n", + " curated[\"S2\"].to_excel(writer, sheet_name=\"Table S2 (13C)\", index=False)\n", + "print(\"wrote\", os.path.relpath(out, REPO))" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_tables/si_table_s03_s04_performance.ipynb b/analysis/si_tables/si_table_s03_s04_performance.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..d68c4280642331a66da908bc9d79aa264a21d399 --- /dev/null +++ b/analysis/si_tables/si_table_s03_s04_performance.ipynb @@ -0,0 +1,122 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "de02451b", + "metadata": {}, + "source": [ + "# Tables S3 and S4: MagNET performance across geometries\n", + "\n", + "How accurately MagNET reproduces its DFT training reference (PBE0/pcSseg-1) as geometries move from\n", + "stationary to vibrated to solvated, for the foundation model and the chloroform MagNET-x." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b6531dbc", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/magnet_test_predictions\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ab667bf", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "import magnet_test_predictions_reader\n", + "import magnet_benchmark\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "979ab00c", + "metadata": {}, + "outputs": [], + "source": [ + "PREDICTIONS = paths.dataset_file(\"magnet_test_predictions\", root=REPO)\n", + "\n", + "def document_path(name):\n", + " # table/spreadsheet outputs go under this notebook's documents/ folder, created on first save\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "059bba13", + "metadata": {}, + "outputs": [], + "source": [ + "rows = magnet_benchmark.exact_stats_table(PREDICTIONS, magnet_test_predictions_reader)\n", + "res = pd.DataFrame(rows).set_index([\"nucleus\", \"model\", \"test_set\"])\n", + "\n", + "table_rows = []\n", + "for nucleus in (\"1H\", \"13C\"):\n", + " for model in magnet_benchmark.MODELS:\n", + " for ts in magnet_benchmark.TEST_SETS:\n", + " r = res.loc[(nucleus, model, ts)]\n", + " pmed, pmae, prmse = magnet_benchmark.PUBLISHED[nucleus][(model, ts)]\n", + " table_rows.append(dict(nucleus=nucleus, model=model, test_set=ts, n=int(r.n),\n", + " median_repro=round(r.median_ae, 6), median_SI=round(pmed, 6),\n", + " mae_repro=round(r.mae, 6), mae_SI=round(pmae, 6),\n", + " rmse_repro=round(r.rmse, 5), rmse_SI=round(prmse, 5)))\n", + "table = pd.DataFrame(table_rows)\n", + "table_s3 = table[table.nucleus == \"1H\"].drop(columns=\"nucleus\")\n", + "table_s4 = table[table.nucleus == \"13C\"].drop(columns=\"nucleus\")\n", + "print(\"Table S3 (1H):\"); display(table_s3)\n", + "print(\"Table S4 (13C):\"); display(table_s4)\n", + "\n", + "# write the reproduced tables to this notebook's documents/ folder, one sheet per SI table\n", + "out = document_path(\"si_table_s03_s04_performance.xlsx\")\n", + "with pd.ExcelWriter(out) as writer:\n", + " table_s3.to_excel(writer, sheet_name=\"Table S3 (1H)\", index=False)\n", + " table_s4.to_excel(writer, sheet_name=\"Table S4 (13C)\", index=False)\n", + "print(\"wrote\", os.path.relpath(out, REPO))" + ] + }, + { + "cell_type": "markdown", + "id": "e5d795f1", + "metadata": {}, + "source": [ + "## Exact-reproduction check" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdb61e94", + "metadata": {}, + "outputs": [], + "source": [ + "# every row's reproduced median/MAE/RMSE should match the published SI value to a few parts per million\n", + "max_dev = (table[[\"median_repro\", \"median_SI\"]].diff(axis=1).iloc[:, -1].abs().max(),\n", + " table[[\"mae_repro\", \"mae_SI\"]].diff(axis=1).iloc[:, -1].abs().max(),\n", + " table[[\"rmse_repro\", \"rmse_SI\"]].diff(axis=1).iloc[:, -1].abs().max())\n", + "print(\"largest reproduced-vs-published deviation (median, MAE, RMSE):\", max_dev)\n", + "assert max(max_dev) < 1e-3, \"a row diverged from the SI by more than float rounding\"" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_tables/si_table_s05_qcd.ipynb b/analysis/si_tables/si_table_s05_qcd.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..ad0760dde787387294abc9f3524516bfec1cdbbf --- /dev/null +++ b/analysis/si_tables/si_table_s05_qcd.ipynb @@ -0,0 +1,115 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8a468c14", + "metadata": {}, + "source": [ + "# Table S5: Performance statistics for predicting QCD corrections\n", + "\n", + "Foundation MagNET's accuracy at predicting the rovibrational (QCD) correction (stationary vs\n", + "trajectory-averaged shielding) over qcdtraj2500 (2500 molecules), ¹H and ¹³C, all shieldings at\n", + "PBE0/pcSseg-1." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aca3e24b", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/magnet_test_predictions\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bbdcbe1d", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "import magnet_test_predictions_reader\n", + "import magnet_benchmark\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c98484b1", + "metadata": {}, + "outputs": [], + "source": [ + "PREDICTIONS = paths.dataset_file(\"magnet_test_predictions\", root=REPO)\n", + "\n", + "def document_path(name):\n", + " # table/spreadsheet outputs go under this notebook's documents/ folder, created on first save\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "da4e04e1", + "metadata": {}, + "outputs": [], + "source": [ + "rows = magnet_benchmark.qcd_stats_table(PREDICTIONS, magnet_test_predictions_reader)\n", + "table_rows = []\n", + "for r in rows:\n", + " nucleus = r[\"model\"].split(\"(\")[1].rstrip(\")\") # \"1H\" / \"13C\"\n", + " pmed, pmae, prmse = magnet_benchmark.PUBLISHED_S5[nucleus]\n", + " table_rows.append(dict(model=r[\"model\"], n=int(r[\"n\"]),\n", + " median_repro=round(r[\"median_ae\"], 8), median_SI=round(pmed, 8),\n", + " mae_repro=round(r[\"mae\"], 8), mae_SI=round(pmae, 8),\n", + " rmse_repro=round(r[\"rmse\"], 8), rmse_SI=round(prmse, 8)))\n", + "table_s5 = pd.DataFrame(table_rows)\n", + "print(\"Table S5 (QCD corrections):\"); display(table_s5)\n", + "\n", + "# write the reproduced table to this notebook's documents/ folder\n", + "out = document_path(\"si_table_s05_qcd.xlsx\")\n", + "with pd.ExcelWriter(out) as writer:\n", + " table_s5.to_excel(writer, sheet_name=\"Table S5\", index=False)\n", + "print(\"wrote\", os.path.relpath(out, REPO))" + ] + }, + { + "cell_type": "markdown", + "id": "9971c179", + "metadata": {}, + "source": [ + "## Exact-reproduction check" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b2c792c0", + "metadata": {}, + "outputs": [], + "source": [ + "# every reproduced median/MAE/RMSE should match the published SI value to a few parts per million\n", + "dev = max((table_s5[[\"median_repro\", \"median_SI\"]].diff(axis=1).iloc[:, -1].abs().max(),\n", + " table_s5[[\"mae_repro\", \"mae_SI\"]].diff(axis=1).iloc[:, -1].abs().max(),\n", + " table_s5[[\"rmse_repro\", \"rmse_SI\"]].diff(axis=1).iloc[:, -1].abs().max()))\n", + "print(\"largest reproduced-vs-published deviation:\", dev)\n", + "assert dev < 1e-3, \"a row diverged from the SI by more than float rounding\"" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_tables/si_table_s06_solvent_corrections.ipynb b/analysis/si_tables/si_table_s06_solvent_corrections.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..f28cfba7fc6a42e77b584c5ad6cde4f60be21d06 --- /dev/null +++ b/analysis/si_tables/si_table_s06_solvent_corrections.ipynb @@ -0,0 +1,115 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "aa1b7ba9", + "metadata": {}, + "source": [ + "# SI Table S6: per-solvent implicit vs explicit fit RMSE (chloroform reference, ¹H)\n", + "\n", + "Per solvent (chloroform reference), fit RMSE of experimental solvent-induced ¹H shifts using implicit\n", + "(PCM) vs explicit (Desmond), and the percent improvement." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f812183a", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f643aecf", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "import delta22\n", + "import paths" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cfa3b9d1", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def document_path(name):\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "130a1bdc", + "metadata": {}, + "outputs": [], + "source": [ + "# Figure 4 uses one DFT method for the proton sites; differences are taken against a reference\n", + "# solvent, and the solvent_mean pseudo-solvent is added for the solvent-averaged reference.\n", + "METHOD, BASIS, GEOM = \"b3lyp_d3bj\", \"pcSseg2\", \"aimnet2\"\n", + "q = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n", + "one = q[(q[\"sap_nmr_method\"] == METHOD) & (q[\"sap_basis\"] == BASIS) & (q[\"sap_geometry_type\"] == GEOM)]\n", + "one = delta22.add_solvent_mean(one)\n", + "print(len(one), \"rows for\", METHOD, BASIS, GEOM)" + ] + }, + { + "cell_type": "markdown", + "id": "fbf8e51b", + "metadata": {}, + "source": [ + "## Table S6: per-solvent implicit vs explicit fit RMSE (chloroform reference, 1H)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d80ea495", + "metadata": {}, + "outputs": [], + "source": [ + "rows = []\n", + "for solvent in ['tetrahydrofuran', 'dichloromethane', 'acetone', 'acetonitrile', 'dimethylsulfoxide', 'trifluoroethanol', 'methanol', 'TIP4P', 'benzene', 'toluene', 'chlorobenzene']:\n", + " sp = delta22.solvent_pair_differences(one, solvent, \"chloroform\", nucleus=\"H\", explicit=\"desmond\")\n", + " implicit = delta22.fit_differences_to_experimental(sp, \"implicit_diff\")[\"rmse\"]\n", + " explicit = delta22.fit_differences_to_experimental(sp, \"explicit_diff\")[\"rmse\"]\n", + " rows.append({\"solvent\": solvent,\n", + " \"implicit_fitting_rmse\": round(implicit, 4),\n", + " \"explicit_fitting_rmse\": round(explicit, 4),\n", + " \"explicit_improvement_pct\": round(100.0 * (implicit - explicit) / implicit, 2)})\n", + "table_s6 = pd.DataFrame(rows)\n", + "print(table_s6.to_string(index=False))\n", + "\n", + "# write the table to this notebook's documents/ folder\n", + "out = document_path(\"si_table_s06_solvent_corrections.xlsx\")\n", + "with pd.ExcelWriter(out) as writer:\n", + " table_s6.to_excel(writer, sheet_name=\"Table S6\", index=False)\n", + "print(\"wrote\", os.path.relpath(out, REPO))" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_tables/si_table_s08_summary.ipynb b/analysis/si_tables/si_table_s08_summary.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..d784cc424ef1f5b955134b34b24986e0c89ab104 --- /dev/null +++ b/analysis/si_tables/si_table_s08_summary.ipynb @@ -0,0 +1,102 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "dce84895", + "metadata": {}, + "source": [ + "# Table S8: Dataset Summary Statistics\n", + "\n", + "Molecule and ¹H/¹³C site counts for each training dataset (site counts from each HDF5's\n", + "`atomic_numbers`, ¹H=1/¹³C=6; MagNET-Zero combines both sigma-pepper rounds with sigma-concentrate)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "df321bd6", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "88710ad5", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "import dataset_summary" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "84e5a441", + "metadata": {}, + "outputs": [], + "source": [ + "DATA_DIR = os.path.join(REPO, \"data\")\n", + "\n", + "def document_path(name):\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6a606891", + "metadata": {}, + "outputs": [], + "source": [ + "table_s8 = dataset_summary.summary_table(DATA_DIR)\n", + "display(table_s8)\n", + "\n", + "# write the table to this notebook's documents/ folder\n", + "out = document_path(\"si_table_s08_summary.xlsx\")\n", + "with pd.ExcelWriter(out) as writer:\n", + " table_s8.to_excel(writer, sheet_name=\"Table S8\", index=False)\n", + "print(\"wrote\", os.path.relpath(out, REPO))" + ] + }, + { + "cell_type": "markdown", + "id": "3bf67791", + "metadata": {}, + "source": [ + "## Exact-reproduction check" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4d5f4416", + "metadata": {}, + "outputs": [], + "source": [ + "# every count should match the published SI Table S8 value exactly\n", + "for _, row in table_s8.iterrows():\n", + " pub = dataset_summary.PUBLISHED_S8[row[\"dataset\"]]\n", + " got = (row[\"molecules\"], row[\"n_1H_sites\"], row[\"n_13C_sites\"])\n", + " assert got == pub, f\"{row['dataset']}: {got} != published {pub}\"\n", + "print(\"all rows match the published SI Table S8 exactly\")" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/analysis/si_tables/si_table_s10_s11_scaling.ipynb b/analysis/si_tables/si_table_s10_s11_scaling.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..daa781b54e3b2c23d61af5f866038588dc972246 --- /dev/null +++ b/analysis/si_tables/si_table_s10_s11_scaling.ipynb @@ -0,0 +1,161 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "badbc639", + "metadata": {}, + "source": [ + "# Tables S10 and S11: MagNET-Zero/MagNET-PCM scaling factors\n", + "\n", + "Per-solvent linear coefficients (intercept, stationary, pcm) mapping MagNET-Zero shieldings plus a\n", + "MagNET-PCM correction to predicted shifts: ¹H (S10) and ¹³C (S11)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "672ff802", + "metadata": {}, + "outputs": [], + "source": [ + "import os, sys\n", + "\n", + "# make the in-repo modules importable (not pip-installed)\n", + "REPO = os.path.abspath(\"../..\")\n", + "for _p in (\"data/delta22\", \"data/scaling_factors\", \"data/applications\",\n", + " \"analysis/code\", \"analysis/code/shared\"):\n", + " sys.path.insert(0, os.path.join(REPO, _p))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a1fa320", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import scaling_factors\n", + "import scaling_factors_reader\n", + "import build_composite_model\n", + "import paths\n", + "from applications_reader import Applications" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "deba4a1e", + "metadata": {}, + "outputs": [], + "source": [ + "DELTA22 = paths.dataset_file(\"delta22\", root=REPO)\n", + "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n", + "\n", + "def document_path(name):\n", + " os.makedirs(\"documents\", exist_ok=True)\n", + " return os.path.join(\"documents\", name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e408fac2", + "metadata": {}, + "outputs": [], + "source": [ + "tables = scaling_factors.published_scaling_tables() # {\"H\": Table S10, \"C\": Table S11}\n", + "print(\"Table S10 (1H):\"); display(tables[\"H\"])\n", + "print(\"Table S11 (13C):\"); display(tables[\"C\"])\n", + "\n", + "out = document_path(\"si_table_s10_s11_scaling.xlsx\")\n", + "with pd.ExcelWriter(out) as writer:\n", + " tables[\"H\"].reset_index().to_excel(writer, sheet_name=\"Table S10 (1H)\", index=False)\n", + " tables[\"C\"].reset_index().to_excel(writer, sheet_name=\"Table S11 (13C)\", index=False)\n", + "print(\"wrote\", os.path.relpath(out, REPO))" + ] + }, + { + "cell_type": "markdown", + "id": "b6094980", + "metadata": {}, + "source": [ + "## Reproducibility check: re-derive from delta-22" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19f08d14", + "metadata": {}, + "outputs": [], + "source": [ + "derived = scaling_factors.build_scaling_tables(DELTA22, XLSX)\n", + "for nucleus in (\"H\", \"C\"):\n", + " p = tables[nucleus]\n", + " d = derived[nucleus].reindex(p.index)[p.columns]\n", + " max_dev = float(np.abs(p.values - d.values).max())\n", + " print(f\"{nucleus}: largest published-vs-rederived deviation = {max_dev:.2e}\")\n", + " assert max_dev < 1e-5, f\"{nucleus} scaling table diverged from the published values\"\n", + "print(\"both tables reproduce from delta-22\")" + ] + }, + { + "cell_type": "markdown", + "id": "8edf48b3", + "metadata": {}, + "source": [ + "## Deployment tables with reflection symmetrization\n", + "\n", + "The tables above match the published SI exactly. For serving new molecules, the recommended tables\n", + "average each prediction with its mirror image, correcting a reflection-parity error. Reproducing\n", + "them needs the model checkpoints; the shipped tables are shown below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "448f2e9f", + "metadata": {}, + "outputs": [], + "source": [ + "symmetrized = scaling_factors_reader.load_symmetrized_tables()\n", + "print(\"Symmetrized deployment Table S10 (1H):\"); display(symmetrized[\"H\"])\n", + "print(\"Symmetrized deployment Table S11 (13C):\"); display(symmetrized[\"C\"])" + ] + }, + { + "cell_type": "markdown", + "id": "ca457bc6", + "metadata": {}, + "source": [ + "## Composite-model coefficients (Figure 5C/5D, SI S15)\n", + "\n", + "The natural-products figures use a larger family of per-solvent coefficients fit on delta-22 (the\n", + "`composite_model` group in `applications.hdf5`): OLS fits, 1000-seed bootstrap resamples, their RMSE\n", + "distributions, and the PCM conversion factors. The group regenerates from released inputs alone." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "47c3d1e7", + "metadata": {}, + "outputs": [], + "source": [ + "reader = Applications(paths.dataset_file(\"applications\", root=REPO))\n", + "regenerated = build_composite_model.regenerate(reader)\n", + "max_dev = build_composite_model.verify(reader, regenerated)\n", + "print(f\"composite_model largest regenerated-vs-stored deviation = {max_dev:.2e}\")" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/data/scaling_factors/scaling_factors_symmetrized_H.csv b/data/scaling_factors/scaling_factors_symmetrized_H.csv new file mode 100644 index 0000000000000000000000000000000000000000..4eaf3bf333c8752559ce43f176cb3311b1a7a20e --- /dev/null +++ b/data/scaling_factors/scaling_factors_symmetrized_H.csv @@ -0,0 +1,13 @@ +solvent,intercept,stationary,pcm +chloroform,31.291983,-0.975682,-0.851715 +tetrahydrofuran,31.319664,-0.980082,-0.781373 +dichloromethane,31.374485,-0.979760,-0.806215 +acetone,31.510823,-0.987188,-1.229850 +acetonitrile,31.494061,-0.985628,-0.969617 +dimethylsulfoxide,31.597261,-0.991031,-1.348009 +trifluoroethanol,30.873172,-0.959885,-0.955808 +methanol,31.253627,-0.976349,-1.246680 +TIP4P,31.356622,-0.978710,-1.475907 +benzene,31.984616,-1.005170,2.239317 +toluene,31.714067,-0.996613,1.948607 +chlorobenzene,31.677950,-0.993715,0.830908 diff --git a/data/scaling_factors/test_scaling_factors_export.py b/data/scaling_factors/test_scaling_factors_export.py new file mode 100644 index 0000000000000000000000000000000000000000..7de5a0ec12c26197f710e4a8e3503cf79f15bc48 --- /dev/null +++ b/data/scaling_factors/test_scaling_factors_export.py @@ -0,0 +1,52 @@ +"""Structural checks for the shipped symmetrized scaling tables (data/scaling_factors/). + +These run on the small CSVs that ship in git, so they need no download. The file is named +test_scaling_factors_export.py (not test_scaling_factors.py) on purpose: analysis/code/ already has a +test_scaling_factors.py, and pytest's default import mode cannot collect two test modules that share +a basename. +""" +import os +import sys + +import pandas as pd +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) + +import scaling_factors_reader as R # noqa: E402 + + +def test_both_tables_load_with_expected_shape_and_columns(): + tables = R.load_symmetrized_tables() + assert set(tables) == {"H", "C"} + for nucleus, df in tables.items(): + assert df.shape == (12, 3), nucleus + assert list(df.columns) == ["intercept", "stationary", "pcm"], nucleus + assert "chloroform" in df.index + assert df.notnull().all().all(), nucleus + + +def test_slopes_are_negative_and_intercepts_in_range(): + tables = R.load_symmetrized_tables() + # the gas-phase (stationary) slope is a shielding-to-shift conversion, always near -1 + for nucleus, df in tables.items(): + assert (df["stationary"] < 0).all(), nucleus + # proton reference intercept sits near a bare-proton shielding (~31 ppm), carbon much higher + assert 25 < tables["H"]["intercept"].mean() < 35 + assert 150 < tables["C"]["intercept"].mean() < 200 + + +def test_predict_shift_uses_the_documented_linear_formula(): + # a synthetic one-row table with round coefficients, so the expected value is hand-computed + # independently of predict_shift's own arithmetic (this catches a sign flip or a swap of the + # stationary/pcm coefficients, which recomputing the same formula inline would not): + # 10 + (-1)*5 + 2*3 = 11 + table = pd.DataFrame({"intercept": [10.0], "stationary": [-1.0], "pcm": [2.0]}, index=["x"]) + table.index.name = "solvent" + assert R.predict_shift(table, "x", shielding=5.0, correction=3.0) == pytest.approx(11.0) + + +def test_reader_rejects_unknown_nucleus(): + with pytest.raises(ValueError): + R._csv_path("X") diff --git a/data/sigma-fresh/decode_sigma_fresh.py b/data/sigma-fresh/decode_sigma_fresh.py new file mode 100644 index 0000000000000000000000000000000000000000..f70d2fe115d21a4cede961b816c458e94dc2cf68 --- /dev/null +++ b/data/sigma-fresh/decode_sigma_fresh.py @@ -0,0 +1,148 @@ +""" +Reader for the sigma-fresh explicit-solvent dataset (numpy + h5py only). + +sigma-fresh holds PBE0/pcSseg-1 (solute) / PBE0/MIDI! (solvent) NMR shieldings on +explicitly-solvated MD snapshots, for ~10,000 solutes in each of four solvents +(chloroform, benzene, methanol, TIP4P water). It trains the MagNET-x explicit-solvent +model, which predicts the per-atom correction `shielding_solvated - shielding_isolated`. + +Layout: one HDF5 with a group per solvent, and a subgroup `solute_NNNNN` per solute. +Each solute stores all 102 trajectory frames of the full solvent cluster. Only the +solute atoms carry shieldings; many frames are geometry-only (shieldings not computed), +marked by `status` and by a sentinel value in the shielding arrays. + + from decode_sigma_fresh import SigmaFresh + with SigmaFresh("sigma-fresh.hdf5") as ds: + print(ds.solvents) # ['chloroform','benzene','methanol','TIP4P'] + print(len(ds.solutes("chloroform"))) # ~10084 + p = ds.pose("chloroform", 3, 5) # solvent, solute_index (1-based), frame_index (1-based) + print(p["name"], p["type"], p["status"]) # e.g. GLY182705 train complete + corr = p["shielding_solvated"] - p["shielding_isolated"] # NaN where not computed + +Requires: numpy, h5py. +""" +import numpy as np +import h5py + +SCALE = 1e-4 +SENTINEL = np.int32(-2147483648) +_STATUS = {0: "incomplete", 1: "geometries_only", 2: "partial_shieldings", + 3: "complete", 4: "error"} + + +class SigmaFresh: + """Reader for a sigma-fresh HDF5 file: per-solvent, per-solute explicit-solvation trajectories.""" + + def __init__(self, path): + self.f = h5py.File(path, "r") + self.scale = float(self.f.attrs.get("scale", SCALE)) + self.sentinel = np.int32(self.f.attrs.get("sentinel", SENTINEL)) + self.solvents = [s for s in self.f.attrs.get("solvents", "").split(",") if s] \ + or sorted(self.f.keys()) + + # ---- catalogue ---- + def solutes(self, solvent): + """Sorted list of solute subgroup names present for a solvent.""" + return sorted(self.f[solvent].keys()) + + def n_solutes(self, solvent): + """Number of solutes present for a solvent.""" + return len(self.f[solvent].keys()) + + def _grp(self, solvent, solute_index): + if solvent not in self.f: + raise KeyError(f"no such solvent: {solvent!r}") + name = f"solute_{int(solute_index):05d}" + if name not in self.f[solvent]: + raise IndexError(f"{solvent}/{name} not in dataset") + return self.f[solvent][name] + + def _decode_shield(self, raw): + out = raw.astype(np.float64) * self.scale + out[raw == self.sentinel] = np.nan + return out + + # ---- access ---- + def solute(self, solvent, solute_index): + """All frames for one solute as a dict (coordinates (102,A,3); shieldings (102,n_solute)).""" + g = self._grp(solvent, solute_index) + Z = g["atomic_numbers"][()] + ns = int(g.attrs["n_solute_atoms"]) + return { + "solvent": solvent, + "solute_index": int(solute_index), + "name": g.attrs["name"], + "type": g.attrs["type"], + "solute_charge": int(g.attrs["solute_charge"]), + "n_solute_atoms": ns, + "n_solvent_atoms": int(self.f[solvent].attrs["n_solvent_atoms"]), + "atomic_numbers": Z, + "solute_mask": np.arange(Z.shape[0]) < ns, + "coordinates": g["coordinates"][()].astype(np.float64) * self.scale, + "shielding_isolated": self._decode_shield(g["shielding_isolated"][()]), + "shielding_solvated": self._decode_shield(g["shielding_solvated"][()]), + "status": np.array([_STATUS[s] for s in g["status"][()]], dtype=object), + "status_code": g["status"][()], + "n_solvents_partial": g["n_solvents_partial"][()], + "partial_radius": g["partial_radius"][()].astype(np.float64) * self.scale, + "full_radius": g["full_radius"][()].astype(np.float64) * self.scale, + } + + def pose(self, solvent, solute_index, frame_index): + """One (solute, frame) pose. frame_index is 1-based (1..102), matching the source/CSVs. + + coordinates: (A,3) full solvent cluster, raw (uncentered). + shielding_isolated / shielding_solvated: (n_solute,), NaN where not computed. + For type='train' solutes the solvated shielding was computed on the partial pose + (the first `n_solvents_partial` solvent molecules); for type='test' it used the + full cluster. + """ + g = self._grp(solvent, solute_index) + n_frames = g["coordinates"].shape[0] + if not (1 <= int(frame_index) <= n_frames): + raise IndexError(f"frame_index {frame_index} out of range 1..{n_frames}") + fi = int(frame_index) - 1 + ns = int(g.attrs["n_solute_atoms"]) + return { + "solvent": solvent, + "solute_index": int(solute_index), + "frame_index": int(frame_index), + "name": g.attrs["name"], + "type": g.attrs["type"], + "n_solute_atoms": ns, + "atomic_numbers": g["atomic_numbers"][()], + "coordinates": g["coordinates"][fi].astype(np.float64) * self.scale, + "shielding_isolated": self._decode_shield(g["shielding_isolated"][fi]), + "shielding_solvated": self._decode_shield(g["shielding_solvated"][fi]), + "status": _STATUS[int(g["status"][fi])], + "n_solvents_partial": int(g["n_solvents_partial"][fi]), + "partial_radius": float(g["partial_radius"][fi]) * self.scale, + "full_radius": float(g["full_radius"][fi]) * self.scale, + } + + # ---- lifecycle ---- + def close(self): + """Close the underlying HDF5 file.""" + self.f.close() + + def __enter__(self): + return self + + def __exit__(self, *a): + self.close() + + def __repr__(self): + return f"" + + +if __name__ == "__main__": + import sys + path = sys.argv[1] if len(sys.argv) > 1 else "sigma-fresh.hdf5" + with SigmaFresh(path) as ds: + print(ds) + for s in ds.solvents: + print(f" {s}: {ds.n_solutes(s)} solutes") + p = ds.pose(ds.solvents[0], 3, 1) + print("pose chloroform/solute_00003/frame_1:", p["name"], p["type"], p["status"]) + print(" isolated[:4]:", p["shielding_isolated"][:4]) + print(" solvated[:4]:", p["shielding_solvated"][:4]) diff --git a/data/sigma-fresh/test_sigma_fresh.py b/data/sigma-fresh/test_sigma_fresh.py new file mode 100644 index 0000000000000000000000000000000000000000..e5027090f0cc0ac36e1fad97e1ed87e648f22a96 --- /dev/null +++ b/data/sigma-fresh/test_sigma_fresh.py @@ -0,0 +1,182 @@ +""" +Tests for the sigma-fresh reader (decode_sigma_fresh.SigmaFresh). + +A tiny synthetic dataset is built in the *exact* on-disk format (so the fixture also +serves as an executable spec of the layout) and the decoder is exercised against it. +One opt-in test runs against the real `sigma-fresh.hdf5` if present next to this file. + +Run: pytest test_sigma_fresh.py -q +Requires: pytest, numpy, h5py. +""" +import os +import numpy as np +import h5py +import pytest + +from decode_sigma_fresh import SigmaFresh, SENTINEL + +SCALE = 1e-4 + + +def _qc(x): + return np.round(np.asarray(x, float) / SCALE).astype(np.int32) + + +def _qs(x): + x = np.asarray(x, float) + out = np.where(np.isnan(x), SENTINEL, np.round(x / SCALE)).astype(np.int32) + return out + + +def _solute(g, name, *, idx, n_solute, znums, coords, iso, solv, status, nsolv, prad, frad, charge=0, typ="train"): + s = g.create_group(f"solute_{idx:05d}") + s.attrs["name"] = name + s.attrs["type"] = typ + s.attrs["n_solute_atoms"] = n_solute + s.attrs["solute_charge"] = charge + s.create_dataset("atomic_numbers", data=np.asarray(znums, np.int8)) + s.create_dataset("coordinates", data=_qc(coords)) + s.create_dataset("shielding_isolated", data=_qs(iso)) + s.create_dataset("shielding_solvated", data=_qs(solv)) + s.create_dataset("status", data=np.asarray(status, np.uint8)) + s.create_dataset("n_solvents_partial", data=np.asarray(nsolv, np.int32)) + s.create_dataset("partial_radius", data=_qc(prad)) + s.create_dataset("full_radius", data=_qc(frad)) + + +@pytest.fixture(scope="module") +def mini_path(tmp_path_factory): + p = str(tmp_path_factory.mktemp("sf") / "mini.hdf5") + with h5py.File(p, "w") as f: + f.attrs["dataset"] = "sigma-fresh" + f.attrs["solvents"] = "chloroform,benzene" + f.attrs["sentinel"] = int(SENTINEL) + f.attrs["scale"] = SCALE + f.attrs["level_of_theory"] = "PBE0/pcSseg-1 (solute), PBE0/MIDI! (solvent)" + # chloroform: solute_00001 has 3 frames; solute (2 atoms) + 1 chloroform (5 atoms) = 7 + gc = f.create_group("chloroform"); gc.attrs["n_solvent_atoms"] = 5 + Z = [6, 1, 6, 1, 17, 17, 17] # 2 solute + 5 solvent + n_solute = 2 + # 4 frames: complete, geometry-only (all NaN), complete, partial (isolated finite, solvated NaN) + coords = np.zeros((4, 7, 3), float) + for fr in range(4): + coords[fr] = np.arange(7 * 3).reshape(7, 3) * 0.1 + fr # raw, frame-shifted (uncentered) + iso = np.array([[150.0, 30.0], [np.nan, np.nan], [151.0, 31.0], [152.0, 32.0]]) + solv = np.array([[149.5, 29.8], [np.nan, np.nan], [150.6, 30.7], [np.nan, np.nan]]) + _solute(gc, "GLY000001", idx=1, n_solute=n_solute, znums=Z, coords=coords, + iso=iso, solv=solv, status=[3, 1, 3, 2], nsolv=[1, 1, 1, 1], + prad=[3.1, 3.1, 3.2, 3.2], frad=[9.9, 9.9, 9.9, 9.9], typ="train") + # benzene: a test solute, 1 frame complete + gb = f.create_group("benzene"); gb.attrs["n_solvent_atoms"] = 12 + Zb = [8, 1] # 2-atom solute, 0 solvent (degenerate but valid) + _solute(gb, "MOL000009", idx=9, n_solute=2, znums=Zb, + coords=np.array([[[0, 0, 0], [1.0, 0, 0]]], float), + iso=np.array([[280.0, 32.0]]), solv=np.array([[279.0, 31.5]]), + status=[3], nsolv=[0], prad=[0.0], frad=[0.0], typ="test") + return p + + +@pytest.fixture +def ds(mini_path): + with SigmaFresh(mini_path) as d: + yield d + + +def test_catalogue(ds): + assert ds.solvents == ["chloroform", "benzene"] + assert ds.n_solutes("chloroform") == 1 and ds.n_solutes("benzene") == 1 + assert ds.solutes("chloroform") == ["solute_00001"] + + +def test_pose_values_roundtrip(ds): + p = ds.pose("chloroform", 1, 1) # 1-based frame + assert p["name"] == "GLY000001" and p["type"] == "train" and p["status"] == "complete" + np.testing.assert_array_equal(p["atomic_numbers"], [6, 1, 6, 1, 17, 17, 17]) + np.testing.assert_allclose(p["shielding_isolated"], [150.0, 30.0], atol=5e-5) + np.testing.assert_allclose(p["shielding_solvated"], [149.5, 29.8], atol=5e-5) + assert p["n_solvents_partial"] == 1 + np.testing.assert_allclose(p["partial_radius"], 3.1, atol=5e-5) + + +def test_shieldings_solute_only(ds): + p = ds.pose("chloroform", 1, 1) + # shieldings have one entry per solute atom, not per cluster atom + assert p["shielding_isolated"].shape == (2,) + assert p["coordinates"].shape == (7, 3) + + +def test_geometry_only_frame_is_nan(ds): + p = ds.pose("chloroform", 1, 2) # frame 2 = geometry-only + assert p["status"] == "geometries_only" + assert np.isnan(p["shielding_isolated"]).all() + assert np.isnan(p["shielding_solvated"]).all() + # geometry is still valid on a geometry-only frame + assert np.isfinite(p["coordinates"]).all() + + +def test_raw_coordinates_not_centered(ds): + # frames are stored raw: distinct frames have distinct centroids (no per-frame centering) + c1 = ds.pose("chloroform", 1, 1)["coordinates"].mean(0) + c3 = ds.pose("chloroform", 1, 3)["coordinates"].mean(0) + assert not np.allclose(c1, c3) + + +def test_solute_bundle_and_mask(ds): + s = ds.solute("chloroform", 1) + assert s["coordinates"].shape == (4, 7, 3) + assert s["shielding_isolated"].shape == (4, 2) + assert s["solute_mask"].tolist() == [True, True, False, False, False, False, False] + assert list(s["status"]) == ["complete", "geometries_only", "complete", "partial_shieldings"] + # correction is finite on complete frames, NaN on the geometry-only and partial frames + corr = s["shielding_solvated"] - s["shielding_isolated"] + assert np.isfinite(corr[0]).all() and np.isnan(corr[1]).all() and np.isnan(corr[3]).all() + + +def test_partial_status_frame(ds): + # status==2: isolated shielding present, solvated not computed (sentinel -> NaN) + p = ds.pose("chloroform", 1, 4) + assert p["status"] == "partial_shieldings" + assert np.isfinite(p["shielding_isolated"]).all() + assert np.isnan(p["shielding_solvated"]).all() + + +def test_test_solute_type(ds): + p = ds.pose("benzene", 9, 1) + assert p["type"] == "test" and p["name"] == "MOL000009" + + +def test_bad_lookups(ds): + with pytest.raises(IndexError): + ds.pose("chloroform", 1, 0) # frames are 1-based + with pytest.raises(IndexError): + ds.pose("chloroform", 1, 99) + with pytest.raises(IndexError): + ds.pose("chloroform", 999, 1) + with pytest.raises(KeyError): + ds.pose("acetone", 1, 1) + + +def test_context_manager_closes(mini_path): + d = SigmaFresh(mini_path) + d.close() + assert not d.f.id.valid + + +# ---------- opt-in: real file, if present ---------- +REAL = os.path.join(os.path.dirname(__file__), "sigma-fresh.hdf5") + + +@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-fresh.hdf5 not present") +def test_real_file(): + with SigmaFresh(REAL) as ds: + assert set(ds.solvents) == {"chloroform", "benzene", "methanol", "TIP4P"} + for s in ds.solvents: + assert ds.n_solutes(s) > 9000 + p = ds.pose("chloroform", 3, 1) + assert p["coordinates"].shape[0] == p["atomic_numbers"].shape[0] + assert p["shielding_isolated"].shape == (p["n_solute_atoms"],) + # at least one complete frame exists for this solute, with finite shieldings + s = ds.solute("chloroform", 3) + complete = s["status_code"] == 3 + assert complete.any() + assert np.isfinite(s["shielding_solvated"][complete]).any() diff --git a/magnet/eqV2/__init__.py b/magnet/eqV2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f278104a5584f22c5c3901a8b9c1e59086537f32 --- /dev/null +++ b/magnet/eqV2/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa +from .equiformer_v2_NMR import EquiformerV2_NMR diff --git a/magnet/eqV2/activation.py b/magnet/eqV2/activation.py new file mode 100644 index 0000000000000000000000000000000000000000..5ff7a7f7fb27471c02b07ad5c1ef08b4df01119f --- /dev/null +++ b/magnet/eqV2/activation.py @@ -0,0 +1,192 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +#from .linear import Linear_gaussian_init + + +class ScaledSiLU(nn.Module): + def __init__(self, inplace=False): + super(ScaledSiLU, self).__init__() + self.inplace = inplace + self.scale_factor = 1.6791767923989418 + + + def forward(self, inputs): + return F.silu(inputs, inplace=self.inplace) * self.scale_factor + + + def extra_repr(self): + str = 'scale_factor={}'.format(self.scale_factor) + if self.inplace: + str = str + ', inplace=True' + return str + + +# Reference: https://github.com/facebookresearch/llama/blob/main/llama/model.py#L175 +class ScaledSwiGLU(nn.Module): + def __init__(self, in_channels, out_channels, bias=True): + super(ScaledSwiGLU, self).__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w = torch.nn.Linear(in_channels, 2 * out_channels, bias=bias) + self.act = ScaledSiLU() + + + def forward(self, inputs): + w = self.w(inputs) + w_1 = w.narrow(-1, 0, self.out_channels) + w_1 = self.act(w_1) + w_2 = w.narrow(-1, self.out_channels, self.out_channels) + out = w_1 * w_2 + return out + + +# Reference: https://github.com/facebookresearch/llama/blob/main/llama/model.py#L175 +class SwiGLU(nn.Module): + def __init__(self, in_channels, out_channels, bias=True): + super(SwiGLU, self).__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w = torch.nn.Linear(in_channels, 2 * out_channels, bias=bias) + self.act = torch.nn.SiLU() + + + def forward(self, inputs): + w = self.w(inputs) + w_1 = w.narrow(-1, 0, self.out_channels) + w_1 = self.act(w_1) + w_2 = w.narrow(-1, self.out_channels, self.out_channels) + out = w_1 * w_2 + return out + + +class SmoothLeakyReLU(torch.nn.Module): + def __init__(self, negative_slope=0.2): + super().__init__() + self.alpha = negative_slope + + + def forward(self, x): + x1 = ((1 + self.alpha) / 2) * x + x2 = ((1 - self.alpha) / 2) * x * (2 * torch.sigmoid(x) - 1) + return x1 + x2 + + + def extra_repr(self): + return 'negative_slope={}'.format(self.alpha) + + +class ScaledSmoothLeakyReLU(torch.nn.Module): + def __init__(self): + super().__init__() + self.act = SmoothLeakyReLU(0.2) + self.scale_factor = 1.531320475574866 + + + def forward(self, x): + return self.act(x) * self.scale_factor + + + def extra_repr(self): + return 'negative_slope={}, scale_factor={}'.format(self.act.alpha, self.scale_factor) + + +class ScaledSigmoid(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale_factor = 1.8467055342154763 + + + def forward(self, x): + return torch.sigmoid(x) * self.scale_factor + + +class GateActivation(torch.nn.Module): + def __init__(self, lmax, mmax, num_channels): + super().__init__() + + self.lmax = lmax + self.mmax = mmax + self.num_channels = num_channels + + # compute `expand_index` based on `lmax` and `mmax` + num_components = 0 + for l in range(1, self.lmax + 1): + num_m_components = min((2 * l + 1), (2 * self.mmax + 1)) + num_components = num_components + num_m_components + expand_index = torch.zeros([num_components]).long() + start_idx = 0 + for l in range(1, self.lmax + 1): + length = min((2 * l + 1), (2 * self.mmax + 1)) + expand_index[start_idx : (start_idx + length)] = (l - 1) + start_idx = start_idx + length + self.register_buffer('expand_index', expand_index) + + self.scalar_act = torch.nn.SiLU() #SwiGLU(self.num_channels, self.num_channels) # # + self.gate_act = torch.nn.Sigmoid() #torch.nn.SiLU() # # + + + def forward(self, gating_scalars, input_tensors): + ''' + `gating_scalars`: shape [N, lmax * num_channels] + `input_tensors`: shape [N, (lmax + 1) ** 2, num_channels] + ''' + + gating_scalars = self.gate_act(gating_scalars) + gating_scalars = gating_scalars.reshape(gating_scalars.shape[0], self.lmax, self.num_channels) + gating_scalars = torch.index_select(gating_scalars, dim=1, index=self.expand_index) + + input_tensors_scalars = input_tensors.narrow(1, 0, 1) + input_tensors_scalars = self.scalar_act(input_tensors_scalars) + + input_tensors_vectors = input_tensors.narrow(1, 1, input_tensors.shape[1] - 1) + input_tensors_vectors = input_tensors_vectors * gating_scalars + + output_tensors = torch.cat((input_tensors_scalars, input_tensors_vectors), dim=1) + + return output_tensors + + +class S2Activation(torch.nn.Module): + ''' + Assume we only have one resolution + ''' + def __init__(self, lmax, mmax): + super().__init__() + self.lmax = lmax + self.mmax = mmax + self.act = torch.nn.SiLU() + + + def forward(self, inputs, SO3_grid): + to_grid_mat = SO3_grid[self.lmax][self.mmax].get_to_grid_mat(device=None) # `device` is not used + from_grid_mat = SO3_grid[self.lmax][self.mmax].get_from_grid_mat(device=None) + x_grid = torch.einsum("bai, zic -> zbac", to_grid_mat, inputs) + x_grid = self.act(x_grid) + outputs = torch.einsum("bai, zbac -> zic", from_grid_mat, x_grid) + return outputs + + +class SeparableS2Activation(torch.nn.Module): + def __init__(self, lmax, mmax): + super().__init__() + + self.lmax = lmax + self.mmax = mmax + + self.scalar_act = torch.nn.SiLU() + self.s2_act = S2Activation(self.lmax, self.mmax) + + + def forward(self, input_scalars, input_tensors, SO3_grid): + output_scalars = self.scalar_act(input_scalars) + output_scalars = output_scalars.reshape(output_scalars.shape[0], 1, output_scalars.shape[-1]) + output_tensors = self.s2_act(input_tensors, SO3_grid) + outputs = torch.cat( + (output_scalars, output_tensors.narrow(1, 1, output_tensors.shape[1] - 1)), + dim=1 + ) + return outputs + + diff --git a/magnet/eqV2/drop.py b/magnet/eqV2/drop.py new file mode 100644 index 0000000000000000000000000000000000000000..3180b0b96b5adeccbc59d6b721b08eab00120047 --- /dev/null +++ b/magnet/eqV2/drop.py @@ -0,0 +1,146 @@ +''' + Add `extra_repr` into DropPath implemented by timm + for displaying more info. +''' + + +import torch +import torch.nn as nn +from e3nn import o3 +import torch.nn.functional as F + + +def drop_path(x, drop_prob: float = 0., training: bool = False): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor.floor_() # binarize + output = x.div(keep_prob) * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + def extra_repr(self): + return 'drop_prob={}'.format(self.drop_prob) + + +class GraphDropPath(nn.Module): + ''' + Consider batch for graph data when dropping paths. + ''' + def __init__(self, drop_prob=None): + super(GraphDropPath, self).__init__() + self.drop_prob = drop_prob + + + def forward(self, x, batch): + batch_size = batch.max() + 1 + shape = (batch_size,) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + ones = torch.ones(shape, dtype=x.dtype, device=x.device) + drop = drop_path(ones, self.drop_prob, self.training) + out = x * drop[batch] + return out + + + def extra_repr(self): + return 'drop_prob={}'.format(self.drop_prob) + + + +class EquivariantDropout(nn.Module): + def __init__(self, irreps, drop_prob): + super(EquivariantDropout, self).__init__() + self.irreps = irreps + self.num_irreps = irreps.num_irreps + self.drop_prob = drop_prob + self.drop = torch.nn.Dropout(drop_prob, True) + self.mul = o3.ElementwiseTensorProduct(irreps, + o3.Irreps('{}x0e'.format(self.num_irreps))) + + + def forward(self, x): + if not self.training or self.drop_prob == 0.0: + return x + shape = (x.shape[0], self.num_irreps) + mask = torch.ones(shape, dtype=x.dtype, device=x.device) + mask = self.drop(mask) + out = self.mul(x, mask) + return out + + +class EquivariantScalarsDropout(nn.Module): + def __init__(self, irreps, drop_prob): + super(EquivariantScalarsDropout, self).__init__() + self.irreps = irreps + self.drop_prob = drop_prob + + + def forward(self, x): + if not self.training or self.drop_prob == 0.0: + return x + out = [] + start_idx = 0 + for mul, ir in self.irreps: + temp = x.narrow(-1, start_idx, mul * ir.dim) + start_idx += mul * ir.dim + if ir.is_scalar(): + temp = F.dropout(temp, p=self.drop_prob, training=self.training) + out.append(temp) + out = torch.cat(out, dim=-1) + return out + + + def extra_repr(self): + return 'irreps={}, drop_prob={}'.format(self.irreps, self.drop_prob) + + +class EquivariantDropoutArraySphericalHarmonics(nn.Module): + def __init__(self, drop_prob, drop_graph=False): + super(EquivariantDropoutArraySphericalHarmonics, self).__init__() + self.drop_prob = drop_prob + self.drop = torch.nn.Dropout(drop_prob, True) + self.drop_graph = drop_graph + + + def forward(self, x, batch=None): + if not self.training or self.drop_prob == 0.0: + return x + assert len(x.shape) == 3 + + if self.drop_graph: + assert batch is not None + batch_size = batch.max() + 1 + shape = (batch_size, 1, x.shape[2]) + mask = torch.ones(shape, dtype=x.dtype, device=x.device) + mask = self.drop(mask) + out = x * mask[batch] + else: + shape = (x.shape[0], 1, x.shape[2]) + mask = torch.ones(shape, dtype=x.dtype, device=x.device) + mask = self.drop(mask) + out = x * mask + + return out + + + def extra_repr(self): + return 'drop_prob={}, drop_graph={}'.format(self.drop_prob, self.drop_graph) + \ No newline at end of file diff --git a/magnet/eqV2/edge_rot_mat.py b/magnet/eqV2/edge_rot_mat.py new file mode 100644 index 0000000000000000000000000000000000000000..6420b58be1a950e158104103ca498e92608172e2 --- /dev/null +++ b/magnet/eqV2/edge_rot_mat.py @@ -0,0 +1,71 @@ +import torch + + +def init_edge_rot_mat(edge_distance_vec): + edge_vec_0 = edge_distance_vec + edge_vec_0_distance = torch.sqrt(torch.sum(edge_vec_0**2, dim=1)) + + # Make sure the atoms are far enough apart + #assert torch.min(edge_vec_0_distance) < 0.0001 + if torch.min(edge_vec_0_distance) < 0.0001: + print( + "Error edge_vec_0_distance: {}".format( + torch.min(edge_vec_0_distance) + ) + ) + + norm_x = edge_vec_0 / (edge_vec_0_distance.view(-1, 1)) + + edge_vec_2 = torch.rand_like(edge_vec_0) - 0.5 + edge_vec_2 = edge_vec_2 / ( + torch.sqrt(torch.sum(edge_vec_2**2, dim=1)).view(-1, 1) + ) + # Create two rotated copys of the random vectors in case the random vector is aligned with norm_x + # With two 90 degree rotated vectors, at least one should not be aligned with norm_x + edge_vec_2b = edge_vec_2.clone() + edge_vec_2b[:, 0] = -edge_vec_2[:, 1] + edge_vec_2b[:, 1] = edge_vec_2[:, 0] + edge_vec_2c = edge_vec_2.clone() + edge_vec_2c[:, 1] = -edge_vec_2[:, 2] + edge_vec_2c[:, 2] = edge_vec_2[:, 1] + vec_dot_b = torch.abs(torch.sum(edge_vec_2b * norm_x, dim=1)).view( + -1, 1 + ) + vec_dot_c = torch.abs(torch.sum(edge_vec_2c * norm_x, dim=1)).view( + -1, 1 + ) + + vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1) + edge_vec_2 = torch.where( + torch.gt(vec_dot, vec_dot_b), edge_vec_2b, edge_vec_2 + ) + vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1) + edge_vec_2 = torch.where( + torch.gt(vec_dot, vec_dot_c), edge_vec_2c, edge_vec_2 + ) + + vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)) + # Check the vectors aren't aligned + assert torch.max(vec_dot) < 0.99 + + norm_z = torch.cross(norm_x, edge_vec_2, dim=1) + norm_z = norm_z / ( + torch.sqrt(torch.sum(norm_z**2, dim=1, keepdim=True)) + ) + norm_z = norm_z / ( + torch.sqrt(torch.sum(norm_z**2, dim=1)).view(-1, 1) + ) + norm_y = torch.cross(norm_x, norm_z, dim=1) + norm_y = norm_y / ( + torch.sqrt(torch.sum(norm_y**2, dim=1, keepdim=True)) + ) + + # Construct the 3D rotation matrix + norm_x = norm_x.view(-1, 3, 1) + norm_y = -norm_y.view(-1, 3, 1) + norm_z = norm_z.view(-1, 3, 1) + + edge_rot_mat_inv = torch.cat([norm_z, norm_x, norm_y], dim=2) + edge_rot_mat = torch.transpose(edge_rot_mat_inv, 1, 2) + + return edge_rot_mat.detach() \ No newline at end of file diff --git a/magnet/eqV2/equiformer_v2_NMR.py b/magnet/eqV2/equiformer_v2_NMR.py new file mode 100644 index 0000000000000000000000000000000000000000..6faf58922c09be69dae1edbc49f749a68323e362 --- /dev/null +++ b/magnet/eqV2/equiformer_v2_NMR.py @@ -0,0 +1,598 @@ +import logging +import time +import math +import numpy as np +import torch +import torch.nn as nn +from pyexpat.model import XML_CQUANT_OPT + +from .ocpmodels.common.registry import registry +from .ocpmodels.common.utils import conditional_grad +from .ocpmodels.models.base import BaseModel +from .ocpmodels.models.scn.sampling import CalcSpherePoints +from .ocpmodels.models.scn.smearing import ( + GaussianSmearing, + LinearSigmoidSmearing, + SigmoidSmearing, + SiLUSmearing, +) + +try: + from e3nn import o3 +except ImportError: + pass + +from .gaussian_rbf import GaussianRadialBasisLayer +from torch.nn import Linear +from .edge_rot_mat import init_edge_rot_mat +from .so3 import ( + CoefficientMappingModule, + SO3_Embedding, + SO3_Grid, + SO3_Rotation, + SO3_LinearV2 +) +from .module_list import ModuleListInfo +from .so2_ops import SO2_Convolution +from .radial_function import RadialFunction +from .layer_norm import ( + EquivariantLayerNormArray, + EquivariantLayerNormArraySphericalHarmonics, + EquivariantRMSNormArraySphericalHarmonics, + EquivariantRMSNormArraySphericalHarmonicsV2, + get_normalization_layer +) +from .transformer_block import ( + SO2EquivariantGraphAttention, + FeedForwardNetwork, + TransBlockV2, +) +from .input_block import EdgeDegreeEmbedding + + +# Statistics of IS2RE 100K +#_AVG_NUM_NODES = 77.81317 +#_AVG_DEGREE = 23.395238876342773 # IS2RE: 100k, max_radius = 5, max_neighbors = 100 + + +@registry.register_model("equiformer_v2") +class EquiformerV2_NMR(BaseModel): + """ + Equiformer with graph attention built upon SO(2) convolution and feedforward network built upon S2 activation + + Args: + use_pbc (bool): Use periodic boundary conditions + regress_forces (bool): Compute forces + otf_graph (bool): Compute graph On The Fly (OTF) + max_neighbors (int): Maximum number of neighbors per atom + max_radius (float): Maximum distance between nieghboring atoms in Angstroms + max_num_elements (int): Maximum atomic number + + num_layers (int): Number of layers in the GNN + sphere_channels (int): Number of spherical channels (one set per resolution) + attn_hidden_channels (int): Number of hidden channels used during SO(2) graph attention + num_heads (int): Number of attention heads + attn_alpha_head (int): Number of channels for alpha vector in each attention head + attn_value_head (int): Number of channels for value vector in each attention head + ffn_hidden_channels (int): Number of hidden channels used during feedforward network + norm_type (str): Type of normalization layer (['layer_norm', 'layer_norm_sh', 'rms_norm_sh']) + + lmax_list (int): List of maximum degree of the spherical harmonics (1 to 10) + mmax_list (int): List of maximum order of the spherical harmonics (0 to lmax) + grid_resolution (int): Resolution of SO3_Grid + + num_sphere_samples (int): Number of samples used to approximate the integration of the sphere in the output blocks + + edge_channels (int): Number of channels for the edge invariant features + use_atom_edge_embedding (bool): Whether to use atomic embedding along with relative distance for edge scalar features + share_atom_edge_embedding (bool): Whether to share `atom_edge_embedding` across all blocks + use_m_share_rad (bool): Whether all m components within a type-L vector of one channel share radial function weights + distance_function ("gaussian", "sigmoid", "linearsigmoid", "silu"): Basis function used for distances + + attn_activation (str): Type of activation function for SO(2) graph attention + use_s2_act_attn (bool): Whether to use attention after S2 activation. Otherwise, use the same attention as Equiformer + use_attn_renorm (bool): Whether to re-normalize attention weights + ffn_activation (str): Type of activation function for feedforward network + use_gate_act (bool): If `True`, use gate activation. Otherwise, use S2 activation + use_grid_mlp (bool): If `True`, use projecting to grids and performing MLPs for FFNs. + use_sep_s2_act (bool): If `True`, use separable S2 activation when `use_gate_act` is False. + + alpha_drop (float): Dropout rate for attention weights + drop_path_rate (float): Drop path rate + proj_drop (float): Dropout rate for outputs of attention and FFN in Transformer blocks + + weight_init (str): ['normal', 'uniform'] initialization of weights of linear layers except those in radial functions + """ + def __init__( + self, + num_atoms, # not used + bond_feat_dim, # not used + num_targets, # not used + + use_pbc=False, + regress_forces=False, + + otf_graph=True, + max_neighbors=500, + max_radius=5.0, + max_num_elements=90, + + num_layers=12, + sphere_channels=128, + attn_hidden_channels=128, + num_heads=8, + attn_alpha_channels=32, + attn_value_channels=16, + ffn_hidden_channels=512, + + norm_type='rms_norm_sh', + + lmax_list=[6], + mmax_list=[2], + grid_resolution=None, + + num_sphere_samples=128, + + edge_channels=128, + use_atom_edge_embedding=True, + share_atom_edge_embedding=False, + use_m_share_rad=False, + distance_function="gaussian", + num_distance_basis=512, + + attn_activation='scaled_silu', + use_s2_act_attn=False, + use_attn_renorm=True, + ffn_activation='scaled_silu', + use_gate_act=False, + use_grid_mlp=False, + use_sep_s2_act=True, + + alpha_drop=0.0, + drop_path_rate=0.0, + proj_drop=0.0, + + weight_init='normal', + + evidential_regression = False, + filter_solvent_edges = False, + solvent_edge_radius = 6.0, + ): + super().__init__() + + # NMR + self._AVG_NUM_NODES = 18.03065905448718 + self._AVG_DEGREE = 15.57930850982666 + + self.filter_solvent_edges = filter_solvent_edges + self.solvent_edge_radius = solvent_edge_radius + + self.use_pbc = use_pbc + self.regress_forces = regress_forces + self.otf_graph = otf_graph + self.max_neighbors = max_neighbors + self.max_radius = max_radius + self.cutoff = max_radius + self.max_num_elements = max_num_elements + + self.num_layers = num_layers + self.sphere_channels = sphere_channels + self.attn_hidden_channels = attn_hidden_channels + self.num_heads = num_heads + self.attn_alpha_channels = attn_alpha_channels + self.attn_value_channels = attn_value_channels + self.ffn_hidden_channels = ffn_hidden_channels + self.norm_type = norm_type + + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.grid_resolution = grid_resolution + + self.num_sphere_samples = num_sphere_samples + + self.edge_channels = edge_channels + self.use_atom_edge_embedding = use_atom_edge_embedding + self.share_atom_edge_embedding = share_atom_edge_embedding + if self.share_atom_edge_embedding: + assert self.use_atom_edge_embedding + self.block_use_atom_edge_embedding = False + else: + self.block_use_atom_edge_embedding = self.use_atom_edge_embedding + self.use_m_share_rad = use_m_share_rad + self.distance_function = distance_function + self.num_distance_basis = num_distance_basis + + self.attn_activation = attn_activation + self.use_s2_act_attn = use_s2_act_attn + self.use_attn_renorm = use_attn_renorm + self.ffn_activation = ffn_activation + self.use_gate_act = use_gate_act + self.use_grid_mlp = use_grid_mlp + self.use_sep_s2_act = use_sep_s2_act + + self.alpha_drop = alpha_drop + self.drop_path_rate = drop_path_rate + self.proj_drop = proj_drop + + self.weight_init = weight_init + assert self.weight_init in ['normal', 'uniform'] + + self.device = 'cpu' #torch.cuda.current_device() + + self.grad_forces = False + self.num_resolutions = len(self.lmax_list) + self.sphere_channels_all = self.num_resolutions * self.sphere_channels + + # Weights for message initialization + self.sphere_embedding = nn.Embedding(self.max_num_elements, self.sphere_channels_all) + + # Initialize the function used to measure the distances between atoms + assert self.distance_function in [ + 'gaussian', + ] + if self.distance_function == 'gaussian': + self.distance_expansion = GaussianSmearing( + 0.0, + self.cutoff, + 600, + 2.0, + ) + #self.distance_expansion = GaussianRadialBasisLayer(num_basis=self.num_distance_basis, cutoff=self.max_radius) + else: + raise ValueError + + # Initialize the sizes of radial functions (input channels and 2 hidden channels) + self.edge_channels_list = [int(self.distance_expansion.num_output)] + [self.edge_channels] * 2 + + # Initialize atom edge embedding + if self.share_atom_edge_embedding and self.use_atom_edge_embedding: + self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels_list[-1]) + self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels_list[-1]) + self.edge_channels_list[0] = self.edge_channels_list[0] + 2 * self.edge_channels_list[-1] + else: + self.source_embedding, self.target_embedding = None, None + + # Initialize the module that compute WignerD matrices and other values for spherical harmonic calculations + self.SO3_rotation = nn.ModuleList() + for i in range(self.num_resolutions): + self.SO3_rotation.append(SO3_Rotation(self.lmax_list[i])) + + # Initialize conversion between degree l and order m layouts + self.mappingReduced = CoefficientMappingModule(self.lmax_list, self.mmax_list) + + # Initialize the transformations between spherical and grid representations + self.SO3_grid = ModuleListInfo('({}, {})'.format(max(self.lmax_list), max(self.lmax_list))) + for l in range(max(self.lmax_list) + 1): + SO3_m_grid = nn.ModuleList() + for m in range(max(self.lmax_list) + 1): + SO3_m_grid.append( + SO3_Grid( + l, + m, + resolution=self.grid_resolution, + normalization='component' + ) + ) + self.SO3_grid.append(SO3_m_grid) + + # Edge-degree embedding + self.edge_degree_embedding = EdgeDegreeEmbedding( + self.sphere_channels, + self.lmax_list, + self.mmax_list, + self.SO3_rotation, + self.mappingReduced, + self.max_num_elements, + self.edge_channels_list, + self.block_use_atom_edge_embedding, + rescale_factor=self._AVG_DEGREE + ) + + # Initialize the blocks for each layer of EquiformerV2 + self.blocks = nn.ModuleList() + for i in range(self.num_layers): + block = TransBlockV2( + self.sphere_channels, + self.attn_hidden_channels, + self.num_heads, + self.attn_alpha_channels, + self.attn_value_channels, + self.ffn_hidden_channels, + self.sphere_channels, + self.lmax_list, + self.mmax_list, + self.SO3_rotation, + self.mappingReduced, + self.SO3_grid, + self.max_num_elements, + self.edge_channels_list, + self.block_use_atom_edge_embedding, + self.use_m_share_rad, + self.attn_activation, + self.use_s2_act_attn, + self.use_attn_renorm, + self.ffn_activation, + self.use_gate_act, + self.use_grid_mlp, + self.use_sep_s2_act, + self.norm_type, + self.alpha_drop, + self.drop_path_rate, + self.proj_drop + ) + self.blocks.append(block) + + + # Output blocks for energy and forces + self.norm = get_normalization_layer(self.norm_type, lmax=max(self.lmax_list), num_channels=self.sphere_channels) + + # predicting scalars (e.g., NMR shieldings) for each node + self.output_block = FeedForwardNetwork( + self.sphere_channels, + self.ffn_hidden_channels, + 1 if not evidential_regression else 4, + self.lmax_list, + self.mmax_list, + self.SO3_grid, + self.ffn_activation, + self.use_gate_act, + self.use_grid_mlp, + self.use_sep_s2_act + ) + + + """ + self.energy_block = FeedForwardNetwork( + self.sphere_channels, + self.ffn_hidden_channels, + 1, + self.lmax_list, + self.mmax_list, + self.SO3_grid, + self.ffn_activation, + self.use_gate_act, + self.use_grid_mlp, + self.use_sep_s2_act + ) + + + if self.regress_forces: + self.force_block = SO2EquivariantGraphAttention( + self.sphere_channels, + self.attn_hidden_channels, + self.num_heads, + self.attn_alpha_channels, + self.attn_value_channels, + 1, + self.lmax_list, + self.mmax_list, + self.SO3_rotation, + self.mappingReduced, + self.SO3_grid, + self.max_num_elements, + self.edge_channels_list, + self.block_use_atom_edge_embedding, + self.use_m_share_rad, + self.attn_activation, + self.use_s2_act_attn, + self.use_attn_renorm, + self.use_gate_act, + self.use_sep_s2_act, + alpha_drop=0.0 + ) + """ + + self.apply(self._init_weights) + self.apply(self._uniform_init_rad_func_linear_weights) + + + @conditional_grad(torch.enable_grad()) + def forward(self, data): + self.batch_size = len(data.natoms) + self.dtype = data.pos.dtype + self.device = data.pos.device + + atomic_numbers = data.atomic_numbers.long() + num_atoms = len(atomic_numbers) + pos = data.pos + + + ( + edge_index, + edge_distance, + edge_distance_vec, + _, # cell_offsets + _, # cell offset distances + neighbors, + ) = self.generate_graph(data) + + + # filter out certain edges here + # `data.keys` is a method in torch_geometric >= 2.6, so test the attribute directly + # (version-independent); this branch only runs for the solvent-aware MagNET-x models. + if (self.filter_solvent_edges) and hasattr(data, "solute"): + contains_solvent_edge = data.solute[edge_index] == 0 + + #is_solvent_solvent_edge = contains_solvent_edge[0,:] & contains_solvent_edge[1,:] # solvent-solvent edge + #remove_solvent_edge = (edge_distance > 4.0) & is_solvent_solvent_edge + + is_solvent_edge = contains_solvent_edge[1,:] # edges incoming to a solvent atom (target atom) + remove_solvent_edge = (edge_distance > self.solvent_edge_radius) & is_solvent_edge + + edge_index = edge_index[:, ~remove_solvent_edge] + edge_distance = edge_distance[~remove_solvent_edge] + edge_distance_vec = edge_distance_vec[~remove_solvent_edge] + neighbors = torch.ones_like(neighbors) * edge_index.shape[1] + + #print(edge_index.shape) + + + ############################################################### + # Initialize data structures + ############################################################### + + # Compute 3x3 rotation matrix per edge + edge_rot_mat = self._init_edge_rot_mat( + data, edge_index, edge_distance_vec + ) + + # Initialize the WignerD matrices and other values for spherical harmonic calculations + for i in range(self.num_resolutions): + self.SO3_rotation[i].set_wigner(edge_rot_mat) + + ############################################################### + # Initialize node embeddings + ############################################################### + + # Init per node representations using an atomic number based embedding + offset = 0 + x = SO3_Embedding( + num_atoms, + self.lmax_list, + self.sphere_channels, + self.device, + self.dtype, + ) + + offset_res = 0 + offset = 0 + # Initialize the l = 0, m = 0 coefficients for each resolution + for i in range(self.num_resolutions): + if self.num_resolutions == 1: + x.embedding[:, offset_res, :] = self.sphere_embedding(atomic_numbers) + else: + x.embedding[:, offset_res, :] = self.sphere_embedding( + atomic_numbers + )[:, offset : offset + self.sphere_channels] + offset = offset + self.sphere_channels + offset_res = offset_res + int((self.lmax_list[i] + 1) ** 2) + + # Edge encoding (distance and atom edge) + edge_distance = self.distance_expansion(edge_distance) + if self.share_atom_edge_embedding and self.use_atom_edge_embedding: + source_element = atomic_numbers[edge_index[0]] # Source atom atomic number + target_element = atomic_numbers[edge_index[1]] # Target atom atomic number + source_embedding = self.source_embedding(source_element) + target_embedding = self.target_embedding(target_element) + edge_distance = torch.cat((edge_distance, source_embedding, target_embedding), dim=1) + + # if we want to add some learned featurization of solvent vs solute atoms (which use different basis sets), do that here + # (if self.share_atom_edge_embedding == False, apply this in transformer_block.py instead.) + + # Edge-degree embedding + edge_degree = self.edge_degree_embedding( + atomic_numbers, + edge_distance, + edge_index) + x.embedding = x.embedding + edge_degree.embedding + + ############################################################### + # Update spherical node embeddings + ############################################################### + + for i in range(self.num_layers): + x = self.blocks[i]( + x, # SO3_Embedding + atomic_numbers, + edge_distance, + edge_index, + batch=data.batch # for GraphDropPath + ) + + # Final layer norm + x.embedding = self.norm(x.embedding) + + node_out = self.output_block(x) + node_out = node_out.embedding.narrow(1, 0, 1) + return node_out + + """ + ############################################################### + # Energy estimation + ############################################################### + node_energy = self.energy_block(x) + node_energy = node_energy.embedding.narrow(1, 0, 1) + energy = torch.zeros(len(data.natoms), device=node_energy.device, dtype=node_energy.dtype) + energy.index_add_(0, data.batch, node_energy.view(-1)) + energy = energy / self._AVG_NUM_NODES + + + ############################################################### + # Force estimation + ############################################################### + if self.regress_forces: + forces = self.force_block(x, + atomic_numbers, + edge_distance, + edge_index) + forces = forces.embedding.narrow(1, 1, 3) + forces = forces.view(-1, 3) + + if not self.regress_forces: + return energy + else: + return energy, forces + """ + + + # Initialize the edge rotation matrics + def _init_edge_rot_mat(self, data, edge_index, edge_distance_vec): + return init_edge_rot_mat(edge_distance_vec) + + + @property + def num_params(self): + return sum(p.numel() for p in self.parameters()) + + + def _init_weights(self, m): + if (isinstance(m, torch.nn.Linear) + or isinstance(m, SO3_LinearV2) + ): + if m.bias is not None: + torch.nn.init.constant_(m.bias, 0) + if self.weight_init == 'normal': + std = 1 / math.sqrt(m.in_features) + torch.nn.init.normal_(m.weight, 0, std) + + elif isinstance(m, torch.nn.LayerNorm): + torch.nn.init.constant_(m.bias, 0) + torch.nn.init.constant_(m.weight, 1.0) + + + def _uniform_init_rad_func_linear_weights(self, m): + if (isinstance(m, RadialFunction)): + m.apply(self._uniform_init_linear_weights) + + + def _uniform_init_linear_weights(self, m): + if isinstance(m, torch.nn.Linear): + if m.bias is not None: + torch.nn.init.constant_(m.bias, 0) + std = 1 / math.sqrt(m.in_features) + torch.nn.init.uniform_(m.weight, -std, std) + + + @torch.jit.ignore + def no_weight_decay(self): + no_wd_list = [] + named_parameters_list = [name for name, _ in self.named_parameters()] + for module_name, module in self.named_modules(): + if (isinstance(module, torch.nn.Linear) + or isinstance(module, SO3_LinearV2) + or isinstance(module, torch.nn.LayerNorm) + or isinstance(module, EquivariantLayerNormArray) + or isinstance(module, EquivariantLayerNormArraySphericalHarmonics) + or isinstance(module, EquivariantRMSNormArraySphericalHarmonics) + or isinstance(module, EquivariantRMSNormArraySphericalHarmonicsV2) + or isinstance(module, GaussianRadialBasisLayer)): + for parameter_name, _ in module.named_parameters(): + if (isinstance(module, torch.nn.Linear) + or isinstance(module, SO3_LinearV2) + ): + if 'weight' in parameter_name: + continue + global_parameter_name = module_name + '.' + parameter_name + assert global_parameter_name in named_parameters_list + no_wd_list.append(global_parameter_name) + return set(no_wd_list) diff --git a/magnet/eqV2/gaussian_rbf.py b/magnet/eqV2/gaussian_rbf.py new file mode 100644 index 0000000000000000000000000000000000000000..9f6a4d0889a2b68588dc2b3042634a62389571e4 --- /dev/null +++ b/magnet/eqV2/gaussian_rbf.py @@ -0,0 +1,46 @@ +import torch + + +@torch.jit.script +def gaussian(x, mean, std): + pi = 3.14159 + a = (2*pi) ** 0.5 + return torch.exp(-0.5 * (((x - mean) / std) ** 2)) / (a * std) + + +# From Graphormer +class GaussianRadialBasisLayer(torch.nn.Module): + def __init__(self, num_basis, cutoff): + super().__init__() + self.num_basis = num_basis + self.cutoff = cutoff + 0.0 + self.mean = torch.nn.Parameter(torch.zeros(1, self.num_basis)) + self.std = torch.nn.Parameter(torch.zeros(1, self.num_basis)) + self.weight = torch.nn.Parameter(torch.ones(1, 1)) + self.bias = torch.nn.Parameter(torch.zeros(1, 1)) + + self.std_init_max = 1.0 + self.std_init_min = 1.0 / self.num_basis + self.mean_init_max = 1.0 + self.mean_init_min = 0 + torch.nn.init.uniform_(self.mean, self.mean_init_min, self.mean_init_max) + torch.nn.init.uniform_(self.std, self.std_init_min, self.std_init_max) + torch.nn.init.constant_(self.weight, 1) + torch.nn.init.constant_(self.bias, 0) + + + def forward(self, dist, node_atom=None, edge_src=None, edge_dst=None): + x = dist / self.cutoff + x = x.unsqueeze(-1) + x = self.weight * x + self.bias + x = x.expand(-1, self.num_basis) + mean = self.mean + std = self.std.abs() + 1e-5 + x = gaussian(x, mean, std) + return x + + + def extra_repr(self): + return 'mean_init_max={}, mean_init_min={}, std_init_max={}, std_init_min={}'.format( + self.mean_init_max, self.mean_init_min, self.std_init_max, self.std_init_min) + \ No newline at end of file diff --git a/magnet/eqV2/input_block.py b/magnet/eqV2/input_block.py new file mode 100644 index 0000000000000000000000000000000000000000..3649da8249ff296946168f96373c71f9f5134a10 --- /dev/null +++ b/magnet/eqV2/input_block.py @@ -0,0 +1,125 @@ +import torch +import torch.nn as nn +import copy + +from .so3 import SO3_Embedding +from .radial_function import RadialFunction + + +class EdgeDegreeEmbedding(torch.nn.Module): + """ + + Args: + sphere_channels (int): Number of spherical channels + + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + + SO3_rotation (list:SO3_Rotation): Class to calculate Wigner-D matrices and rotate embeddings + mappingReduced (CoefficientMappingModule): Class to convert l and m indices once node embedding is rotated + + max_num_elements (int): Maximum number of atomic numbers + edge_channels_list (list:int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels]. + The last one will be used as hidden size when `use_atom_edge_embedding` is `True`. + use_atom_edge_embedding (bool): Whether to use atomic embedding along with relative distance for edge scalar features + + rescale_factor (float): Rescale the sum aggregation + """ + + def __init__( + self, + sphere_channels, + + lmax_list, + mmax_list, + + SO3_rotation, + mappingReduced, + + max_num_elements, + edge_channels_list, + use_atom_edge_embedding, + + rescale_factor + ): + super(EdgeDegreeEmbedding, self).__init__() + self.sphere_channels = sphere_channels + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.num_resolutions = len(self.lmax_list) + self.SO3_rotation = SO3_rotation + self.mappingReduced = mappingReduced + + self.m_0_num_coefficients = self.mappingReduced.m_size[0] + self.m_all_num_coefficents = len(self.mappingReduced.l_harmonic) + + # Create edge scalar (invariant to rotations) features + # Embedding function of the atomic numbers + self.max_num_elements = max_num_elements + self.edge_channels_list = copy.deepcopy(edge_channels_list) + self.use_atom_edge_embedding = use_atom_edge_embedding + + if self.use_atom_edge_embedding: + self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels_list[-1]) + self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels_list[-1]) + nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001) + nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001) + self.edge_channels_list[0] = self.edge_channels_list[0] + 2 * self.edge_channels_list[-1] + else: + self.source_embedding, self.target_embedding = None, None + + # Embedding function of distance + self.edge_channels_list.append(self.m_0_num_coefficients * self.sphere_channels) + self.rad_func = RadialFunction(self.edge_channels_list) + + self.rescale_factor = rescale_factor + + + def forward( + self, + atomic_numbers, + edge_distance, + edge_index + ): + + if self.use_atom_edge_embedding: + source_element = atomic_numbers[edge_index[0]] # Source atom atomic number + target_element = atomic_numbers[edge_index[1]] # Target atom atomic number + source_embedding = self.source_embedding(source_element) + target_embedding = self.target_embedding(target_element) + x_edge = torch.cat((edge_distance, source_embedding, target_embedding), dim=1) + else: + x_edge = edge_distance + + x_edge_m_0 = self.rad_func(x_edge) + x_edge_m_0 = x_edge_m_0.reshape(-1, self.m_0_num_coefficients, self.sphere_channels) + x_edge_m_pad = torch.zeros(( + x_edge_m_0.shape[0], + (self.m_all_num_coefficents - self.m_0_num_coefficients), + self.sphere_channels), + device=x_edge_m_0.device) + x_edge_m_all = torch.cat((x_edge_m_0, x_edge_m_pad), dim=1) + + x_edge_embedding = SO3_Embedding( + 0, + self.lmax_list.copy(), + self.sphere_channels, + device=x_edge_m_all.device, + dtype=x_edge_m_all.dtype + ) + x_edge_embedding.set_embedding(x_edge_m_all) + x_edge_embedding.set_lmax_mmax(self.lmax_list.copy(), self.mmax_list.copy()) + + # Reshape the spherical harmonics based on l (degree) + x_edge_embedding._l_primary(self.mappingReduced) + + # Rotate back the irreps + x_edge_embedding._rotate_inv(self.SO3_rotation, self.mappingReduced) + + # Compute the sum of the incoming neighboring messages for each target node + x_edge_embedding._reduce_edge(edge_index[1], atomic_numbers.shape[0]) + x_edge_embedding.embedding = x_edge_embedding.embedding / self.rescale_factor + + return x_edge_embedding + + diff --git a/magnet/eqV2/layer_norm.py b/magnet/eqV2/layer_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc8b946ef1b563d02b58f6a60fdff71c55e87e6 --- /dev/null +++ b/magnet/eqV2/layer_norm.py @@ -0,0 +1,377 @@ +''' + 1. Normalize features of shape (N, sphere_basis, C), + with sphere_basis = (lmax + 1) ** 2. + + 2. The difference from `layer_norm.py` is that all type-L vectors have + the same number of channels and input features are of shape (N, sphere_basis, C). +''' + +import torch +import torch.nn as nn + + +def get_normalization_layer(norm_type, lmax, num_channels, eps=1e-5, affine=True, normalization='component'): + assert norm_type in ['layer_norm', 'layer_norm_sh', 'rms_norm_sh'] + if norm_type == 'layer_norm': + norm_class = EquivariantLayerNormArray + elif norm_type == 'layer_norm_sh': + norm_class = EquivariantLayerNormArraySphericalHarmonics + elif norm_type == 'rms_norm_sh': + norm_class = EquivariantRMSNormArraySphericalHarmonicsV2 + else: + raise ValueError + return norm_class(lmax, num_channels, eps, affine, normalization) + + +def get_l_to_all_m_expand_index(lmax): + expand_index = torch.zeros([(lmax + 1) ** 2]).long() + for l in range(lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + expand_index[start_idx : (start_idx + length)] = l + return expand_index + + +class EquivariantLayerNormArray(nn.Module): + + def __init__(self, lmax, num_channels, eps=1e-5, affine=True, normalization='component'): + super().__init__() + + self.lmax = lmax + self.num_channels = num_channels + self.eps = eps + self.affine = affine + + if affine: + self.affine_weight = nn.Parameter(torch.ones(lmax + 1, num_channels)) + self.affine_bias = nn.Parameter(torch.zeros(num_channels)) + else: + self.register_parameter('affine_weight', None) + self.register_parameter('affine_bias', None) + + assert normalization in ['norm', 'component'] + self.normalization = normalization + + + def __repr__(self): + return f"{self.__class__.__name__}(lmax={self.lmax}, num_channels={self.num_channels}, eps={self.eps})" + + + @torch.cuda.amp.autocast(enabled=False) + def forward(self, node_input): + ''' + Assume input is of shape [N, sphere_basis, C] + ''' + + out = [] + + for l in range(self.lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + + feature = node_input.narrow(1, start_idx, length) + + # For scalars, first compute and subtract the mean + if l == 0: + feature_mean = torch.mean(feature, dim=2, keepdim=True) + feature = feature - feature_mean + + # Then compute the rescaling factor (norm of each feature vector) + # Rescaling of the norms themselves based on the option "normalization" + if self.normalization == 'norm': + feature_norm = feature.pow(2).sum(dim=1, keepdim=True) # [N, 1, C] + elif self.normalization == 'component': + feature_norm = feature.pow(2).mean(dim=1, keepdim=True) # [N, 1, C] + + feature_norm = torch.mean(feature_norm, dim=2, keepdim=True) # [N, 1, 1] + feature_norm = (feature_norm + self.eps).pow(-0.5) + + if self.affine: + weight = self.affine_weight.narrow(0, l, 1) # [1, C] + weight = weight.view(1, 1, -1) # [1, 1, C] + feature_norm = feature_norm * weight # [N, 1, C] + + feature = feature * feature_norm + + if self.affine and l == 0: + bias = self.affine_bias + bias = bias.view(1, 1, -1) + feature = feature + bias + + out.append(feature) + + out = torch.cat(out, dim=1) + + return out + + + +class EquivariantLayerNormArraySphericalHarmonics(nn.Module): + ''' + 1. Normalize over L = 0. + 2. Normalize across all m components from degrees L > 0. + 3. Do not normalize separately for different L (L > 0). + ''' + def __init__(self, lmax, num_channels, eps=1e-5, affine=True, normalization='component', std_balance_degrees=True): + super().__init__() + + self.lmax = lmax + self.num_channels = num_channels + self.eps = eps + self.affine = affine + self.std_balance_degrees = std_balance_degrees + + # for L = 0 + self.norm_l0 = torch.nn.LayerNorm(self.num_channels, eps=self.eps, elementwise_affine=self.affine) + + # for L > 0 + if self.affine: + self.affine_weight = nn.Parameter(torch.ones(self.lmax, self.num_channels)) + else: + self.register_parameter('affine_weight', None) + + assert normalization in ['norm', 'component'] + self.normalization = normalization + + if self.std_balance_degrees: + balance_degree_weight = torch.zeros((self.lmax + 1) ** 2 - 1, 1) + for l in range(1, self.lmax + 1): + start_idx = l ** 2 - 1 + length = 2 * l + 1 + balance_degree_weight[start_idx : (start_idx + length), :] = (1.0 / length) + balance_degree_weight = balance_degree_weight / self.lmax + self.register_buffer('balance_degree_weight', balance_degree_weight) + else: + self.balance_degree_weight = None + + + def __repr__(self): + return f"{self.__class__.__name__}(lmax={self.lmax}, num_channels={self.num_channels}, eps={self.eps}, std_balance_degrees={self.std_balance_degrees})" + + + @torch.cuda.amp.autocast(enabled=False) + def forward(self, node_input): + ''' + Assume input is of shape [N, sphere_basis, C] + ''' + + out = [] + + # for L = 0 + feature = node_input.narrow(1, 0, 1) + feature = self.norm_l0(feature) + out.append(feature) + + # for L > 0 + if self.lmax > 0: + num_m_components = (self.lmax + 1) ** 2 + feature = node_input.narrow(1, 1, num_m_components - 1) + + # Then compute the rescaling factor (norm of each feature vector) + # Rescaling of the norms themselves based on the option "normalization" + if self.normalization == 'norm': + feature_norm = feature.pow(2).sum(dim=1, keepdim=True) # [N, 1, C] + elif self.normalization == 'component': + if self.std_balance_degrees: + feature_norm = feature.pow(2) # [N, (L_max + 1)**2 - 1, C], without L = 0 + feature_norm = torch.einsum('nic, ia -> nac', feature_norm, self.balance_degree_weight) # [N, 1, C] + else: + feature_norm = feature.pow(2).mean(dim=1, keepdim=True) # [N, 1, C] + + feature_norm = torch.mean(feature_norm, dim=2, keepdim=True) # [N, 1, 1] + feature_norm = (feature_norm + self.eps).pow(-0.5) + + for l in range(1, self.lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + feature = node_input.narrow(1, start_idx, length) # [N, (2L + 1), C] + if self.affine: + weight = self.affine_weight.narrow(0, (l - 1), 1) # [1, C] + weight = weight.view(1, 1, -1) # [1, 1, C] + feature_scale = feature_norm * weight # [N, 1, C] + else: + feature_scale = feature_norm + feature = feature * feature_scale + out.append(feature) + + out = torch.cat(out, dim=1) + return out + + +class EquivariantRMSNormArraySphericalHarmonics(nn.Module): + ''' + 1. Normalize across all m components from degrees L >= 0. + ''' + def __init__(self, lmax, num_channels, eps=1e-5, affine=True, normalization='component'): + super().__init__() + + self.lmax = lmax + self.num_channels = num_channels + self.eps = eps + self.affine = affine + + # for L >= 0 + if self.affine: + self.affine_weight = nn.Parameter(torch.ones((self.lmax + 1), self.num_channels)) + else: + self.register_parameter('affine_weight', None) + + assert normalization in ['norm', 'component'] + self.normalization = normalization + + + def __repr__(self): + return f"{self.__class__.__name__}(lmax={self.lmax}, num_channels={self.num_channels}, eps={self.eps})" + + + @torch.cuda.amp.autocast(enabled=False) + def forward(self, node_input): + ''' + Assume input is of shape [N, sphere_basis, C] + ''' + + out = [] + + # for L >= 0 + feature = node_input + if self.normalization == 'norm': + feature_norm = feature.pow(2).sum(dim=1, keepdim=True) # [N, 1, C] + elif self.normalization == 'component': + feature_norm = feature.pow(2).mean(dim=1, keepdim=True) # [N, 1, C] + + feature_norm = torch.mean(feature_norm, dim=2, keepdim=True) # [N, 1, 1] + feature_norm = (feature_norm + self.eps).pow(-0.5) + + for l in range(0, self.lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + feature = node_input.narrow(1, start_idx, length) # [N, (2L + 1), C] + if self.affine: + weight = self.affine_weight.narrow(0, l, 1) # [1, C] + weight = weight.view(1, 1, -1) # [1, 1, C] + feature_scale = feature_norm * weight # [N, 1, C] + else: + feature_scale = feature_norm + feature = feature * feature_scale + out.append(feature) + + out = torch.cat(out, dim=1) + return out + + +class EquivariantRMSNormArraySphericalHarmonicsV2(nn.Module): + ''' + 1. Normalize across all m components from degrees L >= 0. + 2. Expand weights and multiply with normalized feature to prevent slicing and concatenation. + ''' + def __init__(self, lmax, num_channels, eps=1e-5, affine=True, normalization='component', centering=True, std_balance_degrees=True): + super().__init__() + + self.lmax = lmax + self.num_channels = num_channels + self.eps = eps + self.affine = affine + self.centering = centering + self.std_balance_degrees = std_balance_degrees + + # for L >= 0 + if self.affine: + self.affine_weight = nn.Parameter(torch.ones((self.lmax + 1), self.num_channels)) + if self.centering: + self.affine_bias = nn.Parameter(torch.zeros(self.num_channels)) + else: + self.register_parameter('affine_bias', None) + else: + self.register_parameter('affine_weight', None) + self.register_parameter('affine_bias', None) + + assert normalization in ['norm', 'component'] + self.normalization = normalization + + expand_index = get_l_to_all_m_expand_index(self.lmax) + self.register_buffer('expand_index', expand_index) + + if self.std_balance_degrees: + balance_degree_weight = torch.zeros((self.lmax + 1) ** 2, 1) + for l in range(self.lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + balance_degree_weight[start_idx : (start_idx + length), :] = (1.0 / length) + balance_degree_weight = balance_degree_weight / (self.lmax + 1) + self.register_buffer('balance_degree_weight', balance_degree_weight) + else: + self.balance_degree_weight = None + + + def __repr__(self): + return f"{self.__class__.__name__}(lmax={self.lmax}, num_channels={self.num_channels}, eps={self.eps}, centering={self.centering}, std_balance_degrees={self.std_balance_degrees})" + + + @torch.cuda.amp.autocast(enabled=False) + def forward(self, node_input): + ''' + Assume input is of shape [N, sphere_basis, C] + ''' + + feature = node_input + + if self.centering: + feature_l0 = feature.narrow(1, 0, 1) + feature_l0_mean = feature_l0.mean(dim=2, keepdim=True) # [N, 1, 1] + feature_l0 = feature_l0 - feature_l0_mean + feature = torch.cat((feature_l0, feature.narrow(1, 1, feature.shape[1] - 1)), dim=1) + + # for L >= 0 + if self.normalization == 'norm': + assert not self.std_balance_degrees + feature_norm = feature.pow(2).sum(dim=1, keepdim=True) # [N, 1, C] + elif self.normalization == 'component': + if self.std_balance_degrees: + feature_norm = feature.pow(2) # [N, (L_max + 1)**2, C] + feature_norm = torch.einsum('nic, ia -> nac', feature_norm, self.balance_degree_weight) # [N, 1, C] + else: + feature_norm = feature.pow(2).mean(dim=1, keepdim=True) # [N, 1, C] + + feature_norm = torch.mean(feature_norm, dim=2, keepdim=True) # [N, 1, 1] + feature_norm = (feature_norm + self.eps).pow(-0.5) + + if self.affine: + weight = self.affine_weight.view(1, (self.lmax + 1), self.num_channels) # [1, L_max + 1, C] + weight = torch.index_select(weight, dim=1, index=self.expand_index) # [1, (L_max + 1)**2, C] + feature_norm = feature_norm * weight # [N, (L_max + 1)**2, C] + + out = feature * feature_norm + + if self.affine and self.centering: + out[:, 0:1, :] = out.narrow(1, 0, 1) + self.affine_bias.view(1, 1, self.num_channels) + + return out + + +class EquivariantDegreeLayerScale(nn.Module): + ''' + 1. Similar to Layer Scale used in CaiT (Going Deeper With Image Transformers (ICCV'21)), we scale the output of both attention and FFN. + 2. For degree L > 0, we scale down the square root of 2 * L, which is to emulate halving the number of channels when using higher L. + ''' + def __init__(self, lmax, num_channels, scale_factor=2.0): + super().__init__() + + self.lmax = lmax + self.num_channels = num_channels + self.scale_factor = scale_factor + + self.affine_weight = nn.Parameter(torch.ones(1, (self.lmax + 1), self.num_channels)) + for l in range(1, self.lmax + 1): + self.affine_weight.data[0, l, :].mul_(1.0 / math.sqrt(self.scale_factor * l)) + expand_index = get_l_to_all_m_expand_index(self.lmax) + self.register_buffer('expand_index', expand_index) + + + def __repr__(self): + return f"{self.__class__.__name__}(lmax={self.lmax}, num_channels={self.num_channels}, scale_factor={self.scale_factor})" + + + def forward(self, node_input): + weight = torch.index_select(self.affine_weight, dim=1, index=self.expand_index) # [1, (L_max + 1)**2, C] + node_input = node_input * weight # [N, (L_max + 1)**2, C] + return node_input diff --git a/magnet/eqV2/module_list.py b/magnet/eqV2/module_list.py new file mode 100644 index 0000000000000000000000000000000000000000..622c4400687161667a6a05d19283b9bcbc43b4a3 --- /dev/null +++ b/magnet/eqV2/module_list.py @@ -0,0 +1,11 @@ +import torch + + +class ModuleListInfo(torch.nn.ModuleList): + def __init__(self, info_str, modules=None): + super().__init__(modules) + self.info_str = str(info_str) + + + def __repr__(self): + return self.info_str \ No newline at end of file diff --git a/magnet/eqV2/ocpmodels/__init__.py b/magnet/eqV2/ocpmodels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/magnet/eqV2/ocpmodels/common/__init__.py b/magnet/eqV2/ocpmodels/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/magnet/eqV2/ocpmodels/common/registry.py b/magnet/eqV2/ocpmodels/common/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..d60715cfc658d2d9a24c452c3673225669f16515 --- /dev/null +++ b/magnet/eqV2/ocpmodels/common/registry.py @@ -0,0 +1,313 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. + +This source code is licensed under the MIT license found in the +LICENSE file in the root directory of this source tree. +""" + +# Copyright (c) Facebook, Inc. and its affiliates. +# Borrowed from https://github.com/facebookresearch/pythia/blob/master/pythia/common/registry.py. +""" +Registry is central source of truth. Inspired from Redux's concept of +global store, Registry maintains mappings of various information to unique +keys. Special functions in registry can be used as decorators to register +different kind of classes. + +Import the global registry object using + +``from ocpmodels.common.registry import registry`` + +Various decorators for registry different kind of classes with unique keys + +- Register a model: ``@registry.register_model`` +""" +import importlib + + +def _get_absolute_mapping(name: str): + # in this case, the `name` should be the fully qualified name of the class + # e.g., `ocpmodels.tasks.base_task.BaseTask` + # we can use importlib to get the module (e.g., `ocpmodels.tasks.base_task`) + # and then import the class (e.g., `BaseTask`) + + module_name = ".".join(name.split(".")[:-1]) + class_name = name.split(".")[-1] + + try: + module = importlib.import_module(module_name) + except (ModuleNotFoundError, ValueError) as e: + raise RuntimeError( + f"Could not import module `{module_name}` for import `{name}`" + ) from e + + try: + return getattr(module, class_name) + except AttributeError as e: + raise RuntimeError( + f"Could not import class `{class_name}` from module `{module_name}`" + ) from e + + +class Registry: + r"""Class for registry object which acts as central source of truth.""" + mapping = { + # Mappings to respective classes. + "task_name_mapping": {}, + "dataset_name_mapping": {}, + "model_name_mapping": {}, + "logger_name_mapping": {}, + "trainer_name_mapping": {}, + "state": {}, + } + + @classmethod + def register_task(cls, name): + r"""Register a new task to registry with key 'name' + Args: + name: Key with which the task will be registered. + Usage:: + from ocpmodels.common.registry import registry + from ocpmodels.tasks import BaseTask + @registry.register_task("train") + class TrainTask(BaseTask): + ... + """ + + def wrap(func): + cls.mapping["task_name_mapping"][name] = func + return func + + return wrap + + @classmethod + def register_dataset(cls, name): + r"""Register a dataset to registry with key 'name' + + Args: + name: Key with which the dataset will be registered. + + Usage:: + + from ocpmodels.common.registry import registry + from ocpmodels.datasets import BaseDataset + + @registry.register_dataset("qm9") + class QM9(BaseDataset): + ... + """ + + def wrap(func): + cls.mapping["dataset_name_mapping"][name] = func + return func + + return wrap + + @classmethod + def register_model(cls, name): + r"""Register a model to registry with key 'name' + + Args: + name: Key with which the model will be registered. + + Usage:: + + from ocpmodels.common.registry import registry + from ocpmodels.modules.layers import CGCNNConv + + @registry.register_model("cgcnn") + class CGCNN(): + ... + """ + + def wrap(func): + cls.mapping["model_name_mapping"][name] = func + return func + + return wrap + + @classmethod + def register_logger(cls, name): + r"""Register a logger to registry with key 'name' + + Args: + name: Key with which the logger will be registered. + + Usage:: + + from ocpmodels.common.registry import registry + + @registry.register_logger("tensorboard") + class WandB(): + ... + """ + + def wrap(func): + from ocpmodels.common.logger import Logger + + assert issubclass( + func, Logger + ), "All loggers must inherit Logger class" + cls.mapping["logger_name_mapping"][name] = func + return func + + return wrap + + @classmethod + def register_trainer(cls, name): + r"""Register a trainer to registry with key 'name' + + Args: + name: Key with which the trainer will be registered. + + Usage:: + + from ocpmodels.common.registry import registry + + @registry.register_trainer("active_discovery") + class ActiveDiscoveryTrainer(): + ... + """ + + def wrap(func): + cls.mapping["trainer_name_mapping"][name] = func + return func + + return wrap + + @classmethod + def register(cls, name, obj): + r"""Register an item to registry with key 'name' + + Args: + name: Key with which the item will be registered. + + Usage:: + + from ocpmodels.common.registry import registry + + registry.register("config", {}) + """ + path = name.split(".") + current = cls.mapping["state"] + + for part in path[:-1]: + if part not in current: + current[part] = {} + current = current[part] + + current[path[-1]] = obj + + @classmethod + def __import_error(cls, name: str, mapping_name: str): + kind = mapping_name[: -len("_name_mapping")] + mapping = cls.mapping.get(mapping_name, {}) + existing_keys = list(mapping.keys()) + + existing_cls_path = ( + mapping.get(existing_keys[-1], None) if existing_keys else None + ) + if existing_cls_path is not None: + existing_cls_path = f"{existing_cls_path.__module__}.{existing_cls_path.__qualname__}" + else: + existing_cls_path = "ocpmodels.trainers.ForcesTrainer" + + existing_keys = [f"'{name}'" for name in existing_keys] + existing_keys = ( + ", ".join(existing_keys[:-1]) + " or " + existing_keys[-1] + ) + existing_keys_str = ( + f" (one of {existing_keys})" if existing_keys else "" + ) + return RuntimeError( + f"Failed to find the {kind} '{name}'. " + f"You may either use a {kind} from the registry{existing_keys_str} " + f"or provide the full import path to the {kind} (e.g., '{existing_cls_path}')." + ) + + @classmethod + def get_class(cls, name: str, mapping_name: str): + existing_mapping = cls.mapping[mapping_name].get(name, None) + if existing_mapping is not None: + return existing_mapping + + # mapping be class path of type `{module_name}.{class_name}` (e.g., `ocpmodels.trainers.ForcesTrainer`) + if name.count(".") < 1: + raise cls.__import_error(name, mapping_name) + + try: + return _get_absolute_mapping(name) + except RuntimeError as e: + raise cls.__import_error(name, mapping_name) from e + + @classmethod + def get_task_class(cls, name): + return cls.get_class(name, "task_name_mapping") + + @classmethod + def get_dataset_class(cls, name): + return cls.get_class(name, "dataset_name_mapping") + + @classmethod + def get_model_class(cls, name): + return cls.get_class(name, "model_name_mapping") + + @classmethod + def get_logger_class(cls, name): + return cls.get_class(name, "logger_name_mapping") + + @classmethod + def get_trainer_class(cls, name): + return cls.get_class(name, "trainer_name_mapping") + + @classmethod + def get(cls, name, default=None, no_warning=False): + r"""Get an item from registry with key 'name' + + Args: + name (string): Key whose value needs to be retreived. + default: If passed and key is not in registry, default value will + be returned with a warning. Default: None + no_warning (bool): If passed as True, warning when key doesn't exist + will not be generated. Useful for cgcnn's + internal operations. Default: False + Usage:: + + from ocpmodels.common.registry import registry + + config = registry.get("config") + """ + original_name = name + name = name.split(".") + value = cls.mapping["state"] + for subname in name: + value = value.get(subname, default) + if value is default: + break + + if ( + "writer" in cls.mapping["state"] + and value == default + and no_warning is False + ): + cls.mapping["state"]["writer"].write( + "Key {} is not present in registry, returning default value " + "of {}".format(original_name, default) + ) + return value + + @classmethod + def unregister(cls, name): + r"""Remove an item from registry with key 'name' + + Args: + name: Key which needs to be removed. + Usage:: + + from ocpmodels.common.registry import registry + + config = registry.unregister("config") + """ + return cls.mapping["state"].pop(name, None) + + +registry = Registry() diff --git a/magnet/eqV2/ocpmodels/common/utils.py b/magnet/eqV2/ocpmodels/common/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6b09dc811f6a48de762ad7c5700bf7fd3ce3d6e8 --- /dev/null +++ b/magnet/eqV2/ocpmodels/common/utils.py @@ -0,0 +1,1080 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. + +This source code is licensed under the MIT license found in the +LICENSE file in the root directory of this source tree. +""" + +import ast +import collections +import copy +import importlib +import itertools +import json +import logging +import os +import sys +import time +from argparse import Namespace +from bisect import bisect +from contextlib import contextmanager +from dataclasses import dataclass +from functools import wraps +from itertools import product +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional + +import numpy as np +import torch +import torch.nn as nn +import torch_geometric +# yaml and matplotlib are imported lazily inside the few (unused-at-inference) helpers below that +# need them, so `import magnet` does not require either package. +from torch_geometric.data import Data +from torch_geometric.utils import remove_self_loops +from torch_scatter import scatter, segment_coo, segment_csr + +if TYPE_CHECKING: + from torch.nn.modules.module import _IncompatibleKeys + + +def pyg2_data_transform(data: Data): + """ + if we're on the new pyg (2.0 or later) and if the Data stored is in older format + we need to convert the data to the new format + """ + if torch_geometric.__version__ >= "2.0" and "_store" not in data.__dict__: + return Data( + **{k: v for k, v in data.__dict__.items() if v is not None} + ) + + return data + + +def save_checkpoint( + state, checkpoint_dir="checkpoints/", checkpoint_file="checkpoint.pt" +): + filename = os.path.join(checkpoint_dir, checkpoint_file) + torch.save(state, filename) + return filename + + +class Complete(object): + def __call__(self, data): + device = data.edge_index.device + + row = torch.arange(data.num_nodes, dtype=torch.long, device=device) + col = torch.arange(data.num_nodes, dtype=torch.long, device=device) + + row = row.view(-1, 1).repeat(1, data.num_nodes).view(-1) + col = col.repeat(data.num_nodes) + edge_index = torch.stack([row, col], dim=0) + + edge_attr = None + if data.edge_attr is not None: + idx = data.edge_index[0] * data.num_nodes + data.edge_index[1] + size = list(data.edge_attr.size()) + size[0] = data.num_nodes * data.num_nodes + edge_attr = data.edge_attr.new_zeros(size) + edge_attr[idx] = data.edge_attr + + edge_index, edge_attr = remove_self_loops(edge_index, edge_attr) + data.edge_attr = edge_attr + data.edge_index = edge_index + + return data + + +def warmup_lr_lambda(current_step, optim_config): + """Returns a learning rate multiplier. + Till `warmup_steps`, learning rate linearly increases to `initial_lr`, + and then gets multiplied by `lr_gamma` every time a milestone is crossed. + """ + + # keep this block for older configs that have warmup_epochs instead of warmup_steps + # and lr_milestones are defined in epochs + if ( + any(x < 100 for x in optim_config["lr_milestones"]) + or "warmup_epochs" in optim_config + ): + raise Exception( + "ConfigError: please define lr_milestones in steps not epochs and define warmup_steps instead of warmup_epochs" + ) + + if current_step <= optim_config["warmup_steps"]: + alpha = current_step / float(optim_config["warmup_steps"]) + return optim_config["warmup_factor"] * (1.0 - alpha) + alpha + else: + idx = bisect(optim_config["lr_milestones"], current_step) + return pow(optim_config["lr_gamma"], idx) + + +def print_cuda_usage(): + print("Memory Allocated:", torch.cuda.memory_allocated() / (1024 * 1024)) + print( + "Max Memory Allocated:", + torch.cuda.max_memory_allocated() / (1024 * 1024), + ) + print("Memory Cached:", torch.cuda.memory_cached() / (1024 * 1024)) + print("Max Memory Cached:", torch.cuda.max_memory_cached() / (1024 * 1024)) + + +def conditional_grad(dec): + "Decorator to enable/disable grad depending on whether force/energy predictions are being made" + # Adapted from https://stackoverflow.com/questions/60907323/accessing-class-property-as-decorator-argument + def decorator(func): + @wraps(func) + def cls_method(self, *args, **kwargs): + f = func + if self.regress_forces and not getattr(self, "direct_forces", 0): + f = dec(func) + return f(self, *args, **kwargs) + + return cls_method + + return decorator + + +def plot_histogram(data, xlabel="", ylabel="", title=""): + from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas + from matplotlib.figure import Figure + + assert isinstance(data, list) + + # Preset + fig = Figure(figsize=(5, 4), dpi=150) + canvas = FigureCanvas(fig) + ax = fig.gca() + + # Plot + ax.hist(data, bins=20, rwidth=0.9, zorder=3) + + # Axes + ax.grid(color="0.95", zorder=0) + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + fig.tight_layout(pad=2) + + # Return numpy array + canvas.draw() + image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) + image_from_plot = image_from_plot.reshape( + fig.canvas.get_width_height()[::-1] + (3,) + ) + + return image_from_plot + + +# Override the collation method in `pytorch_geometric.data.InMemoryDataset` +def collate(data_list): + keys = data_list[0].keys + data = data_list[0].__class__() + + for key in keys: + data[key] = [] + slices = {key: [0] for key in keys} + + for item, key in product(data_list, keys): + data[key].append(item[key]) + if torch.is_tensor(item[key]): + s = slices[key][-1] + item[key].size( + item.__cat_dim__(key, item[key]) + ) + elif isinstance(item[key], int) or isinstance(item[key], float): + s = slices[key][-1] + 1 + else: + raise ValueError("Unsupported attribute type") + slices[key].append(s) + + if hasattr(data_list[0], "__num_nodes__"): + data.__num_nodes__ = [] + for item in data_list: + data.__num_nodes__.append(item.num_nodes) + + for key in keys: + if torch.is_tensor(data_list[0][key]): + data[key] = torch.cat( + data[key], dim=data.__cat_dim__(key, data_list[0][key]) + ) + else: + data[key] = torch.tensor(data[key]) + slices[key] = torch.tensor(slices[key], dtype=torch.long) + + return data, slices + + +def add_edge_distance_to_graph( + batch, + device="cpu", + dmin=0.0, + dmax=6.0, + num_gaussians=50, +): + # Make sure x has positions. + if not all(batch.pos[0][:] == batch.x[0][-3:]): + batch.x = torch.cat([batch.x, batch.pos.float()], dim=1) + # First set computations to be tracked for positions. + batch.x = batch.x.requires_grad_(True) + # Then compute Euclidean distance between edge endpoints. + pdist = torch.nn.PairwiseDistance(p=2.0) + distances = pdist( + batch.x[batch.edge_index[0]][:, -3:], + batch.x[batch.edge_index[1]][:, -3:], + ) + # Expand it using a gaussian basis filter. + gdf_filter = torch.linspace(dmin, dmax, num_gaussians) + var = gdf_filter[1] - gdf_filter[0] + gdf_filter, var = gdf_filter.to(device), var.to(device) + gdf_distances = torch.exp( + -((distances.view(-1, 1) - gdf_filter) ** 2) / var**2 + ) + # Reassign edge attributes. + batch.edge_weight = distances + batch.edge_attr = gdf_distances.float() + return batch + + +def _import_local_file(path: Path, *, project_root: Path): + """ + Imports a Python file as a module + + :param path: The path to the file to import + :type path: Path + :param project_root: The root directory of the project (i.e., the "ocp" folder) + :type project_root: Path + """ + + path = path.resolve() + project_root = project_root.resolve() + + module_name = ".".join( + path.absolute() + .relative_to(project_root.absolute()) + .with_suffix("") + .parts + ) + logging.debug(f"Resolved module name of {path} to {module_name}") + importlib.import_module(module_name) + + +def setup_experimental_imports(project_root: Path): + experimental_folder = (project_root / "experimental").resolve() + if not experimental_folder.exists() or not experimental_folder.is_dir(): + return + + experimental_files = [ + f.resolve().absolute() for f in experimental_folder.rglob("*.py") + ] + # Ignore certain directories within experimental + ignore_file = experimental_folder / ".ignore" + if ignore_file.exists(): + with open(ignore_file, "r") as f: + for line in f.read().splitlines(): + for ignored_file in (experimental_folder / line).rglob("*.py"): + experimental_files.remove( + ignored_file.resolve().absolute() + ) + + for f in experimental_files: + _import_local_file(f, project_root=project_root) + + +def _get_project_root(): + """ + Gets the root folder of the project (the "ocp" folder) + :return: The absolute path to the project root. + """ + from ocpmodels.common.registry import registry + + # Automatically load all of the modules, so that + # they register with registry + root_folder = registry.get("ocpmodels_root", no_warning=True) + + if root_folder is not None: + assert isinstance(root_folder, str), "ocpmodels_root must be a string" + root_folder = Path(root_folder).resolve().absolute() + assert root_folder.exists(), f"{root_folder} does not exist" + assert root_folder.is_dir(), f"{root_folder} is not a directory" + else: + root_folder = Path(__file__).resolve().absolute().parent.parent + + # root_folder is the "ocpmodes" folder, so we need to go up one more level + return root_folder.parent + + +# Copied from https://github.com/facebookresearch/mmf/blob/master/mmf/utils/env.py#L89. +def setup_imports(config: Optional[dict] = None): + from ocpmodels.common.registry import registry + + skip_experimental_imports = (config or {}).get( + "skip_experimental_imports", None + ) + + # First, check if imports are already setup + has_already_setup = registry.get("imports_setup", no_warning=True) + if has_already_setup: + return + + try: + project_root = _get_project_root() + logging.info(f"Project root: {project_root}") + importlib.import_module("ocpmodels.common.logger") + + import_keys = ["trainers", "datasets", "models", "tasks"] + for key in import_keys: + for f in (project_root / "ocpmodels" / key).rglob("*.py"): + _import_local_file(f, project_root=project_root) + + if not skip_experimental_imports: + setup_experimental_imports(project_root) + finally: + import nets # for equiformer_v2 + import oc20.trainer # for equiformer_v2 + registry.register("imports_setup", True) + + +def dict_set_recursively(dictionary, key_sequence, val): + top_key = key_sequence.pop(0) + if len(key_sequence) == 0: + dictionary[top_key] = val + else: + if top_key not in dictionary: + dictionary[top_key] = {} + dict_set_recursively(dictionary[top_key], key_sequence, val) + + +def parse_value(value): + """ + Parse string as Python literal if possible and fallback to string. + """ + try: + return ast.literal_eval(value) + except (ValueError, SyntaxError): + # Use as string if nothing else worked + return value + + +def create_dict_from_args(args: list, sep: str = "."): + """ + Create a (nested) dictionary from console arguments. + Keys in different dictionary levels are separated by sep. + """ + return_dict = {} + for arg in args: + arg = arg.strip("--") + keys_concat, val = arg.split("=") + val = parse_value(val) + key_sequence = keys_concat.split(sep) + dict_set_recursively(return_dict, key_sequence, val) + return return_dict + + +def load_config(path: str, previous_includes: list = []): + path = Path(path) + if path in previous_includes: + raise ValueError( + f"Cyclic config include detected. {path} included in sequence {previous_includes}." + ) + previous_includes = previous_includes + [path] + + import yaml + + direct_config = yaml.safe_load(open(path, "r")) + + # Load config from included files. + if "includes" in direct_config: + includes = direct_config.pop("includes") + else: + includes = [] + if not isinstance(includes, list): + raise AttributeError( + "Includes must be a list, '{}' provided".format(type(includes)) + ) + + config = {} + duplicates_warning = [] + duplicates_error = [] + + for include in includes: + include_config, inc_dup_warning, inc_dup_error = load_config( + include, previous_includes + ) + duplicates_warning += inc_dup_warning + duplicates_error += inc_dup_error + + # Duplicates between includes causes an error + config, merge_dup_error = merge_dicts(config, include_config) + duplicates_error += merge_dup_error + + # Duplicates between included and main file causes warnings + config, merge_dup_warning = merge_dicts(config, direct_config) + duplicates_warning += merge_dup_warning + + return config, duplicates_warning, duplicates_error + + +def build_config(args, args_override): + config, duplicates_warning, duplicates_error = load_config(args.config_yml) + if len(duplicates_warning) > 0: + logging.warning( + f"Overwritten config parameters from included configs " + f"(non-included parameters take precedence): {duplicates_warning}" + ) + if len(duplicates_error) > 0: + raise ValueError( + f"Conflicting (duplicate) parameters in simultaneously " + f"included configs: {duplicates_error}" + ) + + # Check for overridden parameters. + if args_override != []: + overrides = create_dict_from_args(args_override) + config, _ = merge_dicts(config, overrides) + + # Some other flags. + config["mode"] = args.mode + config["identifier"] = args.identifier + config["timestamp_id"] = args.timestamp_id + config["seed"] = args.seed + config["is_debug"] = args.debug + config["run_dir"] = args.run_dir + config["print_every"] = args.print_every + config["amp"] = args.amp + config["checkpoint"] = args.checkpoint + config["cpu"] = args.cpu + # Submit + config["submit"] = args.submit + config["summit"] = args.summit + # Distributed + config["local_rank"] = args.local_rank + config["distributed_port"] = args.distributed_port + config["world_size"] = args.num_nodes * args.num_gpus + config["distributed_backend"] = args.distributed_backend + config["noddp"] = args.no_ddp + config["gp_gpus"] = args.gp_gpus + + return config + + +def create_grid(base_config, sweep_file): + def _flatten_sweeps(sweeps, root_key="", sep="."): + flat_sweeps = [] + for key, value in sweeps.items(): + new_key = root_key + sep + key if root_key else key + if isinstance(value, collections.MutableMapping): + flat_sweeps.extend(_flatten_sweeps(value, new_key).items()) + else: + flat_sweeps.append((new_key, value)) + return collections.OrderedDict(flat_sweeps) + + def _update_config(config, keys, override_vals, sep="."): + for key, value in zip(keys, override_vals): + key_path = key.split(sep) + child_config = config + for name in key_path[:-1]: + child_config = child_config[name] + child_config[key_path[-1]] = value + return config + + import yaml + + sweeps = yaml.safe_load(open(sweep_file, "r")) + flat_sweeps = _flatten_sweeps(sweeps) + keys = list(flat_sweeps.keys()) + values = list(itertools.product(*flat_sweeps.values())) + + configs = [] + for i, override_vals in enumerate(values): + config = copy.deepcopy(base_config) + config = _update_config(config, keys, override_vals) + config["identifier"] = config["identifier"] + f"_run{i}" + configs.append(config) + return configs + + +def save_experiment_log(args, jobs, configs): + log_file = args.logdir / "exp" / time.strftime("%Y-%m-%d-%I-%M-%S%p.log") + log_file.parent.mkdir(exist_ok=True, parents=True) + with open(log_file, "w") as f: + for job, config in zip(jobs, configs): + print( + json.dumps( + { + "config": config, + "slurm_id": job.job_id, + "timestamp": time.strftime("%I:%M:%S%p %Z %b %d, %Y"), + } + ), + file=f, + ) + return log_file + + +def get_pbc_distances( + pos, + edge_index, + cell, + cell_offsets, + neighbors, + return_offsets=False, + return_distance_vec=False, +): + row, col = edge_index + + distance_vectors = pos[row] - pos[col] + + # correct for pbc + neighbors = neighbors.to(cell.device) + cell = torch.repeat_interleave(cell, neighbors, dim=0) + offsets = cell_offsets.float().view(-1, 1, 3).bmm(cell.float()).view(-1, 3) + distance_vectors += offsets + + # compute distances + distances = distance_vectors.norm(dim=-1) + + # redundancy: remove zero distances + nonzero_idx = torch.arange(len(distances), device=distances.device)[ + distances != 0 + ] + edge_index = edge_index[:, nonzero_idx] + distances = distances[nonzero_idx] + + out = { + "edge_index": edge_index, + "distances": distances, + } + + if return_distance_vec: + out["distance_vec"] = distance_vectors[nonzero_idx] + + if return_offsets: + out["offsets"] = offsets[nonzero_idx] + + return out + + +def radius_graph_pbc( + data, radius, max_num_neighbors_threshold, pbc=[True, True, True] +): + device = data.pos.device + batch_size = len(data.natoms) + + if hasattr(data, "pbc"): + data.pbc = torch.atleast_2d(data.pbc) + for i in range(3): + if not torch.any(data.pbc[:, i]).item(): + pbc[i] = False + elif torch.all(data.pbc[:, i]).item(): + pbc[i] = True + else: + raise RuntimeError( + "Different structures in the batch have different PBC configurations. This is not currently supported." + ) + + # position of the atoms + atom_pos = data.pos + + # Before computing the pairwise distances between atoms, first create a list of atom indices to compare for the entire batch + num_atoms_per_image = data.natoms + num_atoms_per_image_sqr = (num_atoms_per_image**2).long() + + # index offset between images + index_offset = ( + torch.cumsum(num_atoms_per_image, dim=0) - num_atoms_per_image + ) + + index_offset_expand = torch.repeat_interleave( + index_offset, num_atoms_per_image_sqr + ) + num_atoms_per_image_expand = torch.repeat_interleave( + num_atoms_per_image, num_atoms_per_image_sqr + ) + + # Compute a tensor containing sequences of numbers that range from 0 to num_atoms_per_image_sqr for each image + # that is used to compute indices for the pairs of atoms. This is a very convoluted way to implement + # the following (but 10x faster since it removes the for loop) + # for batch_idx in range(batch_size): + # batch_count = torch.cat([batch_count, torch.arange(num_atoms_per_image_sqr[batch_idx], device=device)], dim=0) + num_atom_pairs = torch.sum(num_atoms_per_image_sqr) + index_sqr_offset = ( + torch.cumsum(num_atoms_per_image_sqr, dim=0) - num_atoms_per_image_sqr + ) + index_sqr_offset = torch.repeat_interleave( + index_sqr_offset, num_atoms_per_image_sqr + ) + atom_count_sqr = ( + torch.arange(num_atom_pairs, device=device) - index_sqr_offset + ) + + # Compute the indices for the pairs of atoms (using division and mod) + # If the systems get too large this apporach could run into numerical precision issues + index1 = ( + torch.div( + atom_count_sqr, num_atoms_per_image_expand, rounding_mode="floor" + ) + ) + index_offset_expand + index2 = ( + atom_count_sqr % num_atoms_per_image_expand + ) + index_offset_expand + # Get the positions for each atom + pos1 = torch.index_select(atom_pos, 0, index1) + pos2 = torch.index_select(atom_pos, 0, index2) + + # Calculate required number of unit cells in each direction. + # Smallest distance between planes separated by a1 is + # 1 / ||(a2 x a3) / V||_2, since a2 x a3 is the area of the plane. + # Note that the unit cell volume V = a1 * (a2 x a3) and that + # (a2 x a3) / V is also the reciprocal primitive vector + # (crystallographer's definition). + + cross_a2a3 = torch.cross(data.cell[:, 1], data.cell[:, 2], dim=-1) + cell_vol = torch.sum(data.cell[:, 0] * cross_a2a3, dim=-1, keepdim=True) + + if pbc[0]: + inv_min_dist_a1 = torch.norm(cross_a2a3 / cell_vol, p=2, dim=-1) + rep_a1 = torch.ceil(radius * inv_min_dist_a1) + else: + rep_a1 = data.cell.new_zeros(1) + + if pbc[1]: + cross_a3a1 = torch.cross(data.cell[:, 2], data.cell[:, 0], dim=-1) + inv_min_dist_a2 = torch.norm(cross_a3a1 / cell_vol, p=2, dim=-1) + rep_a2 = torch.ceil(radius * inv_min_dist_a2) + else: + rep_a2 = data.cell.new_zeros(1) + + if pbc[2]: + cross_a1a2 = torch.cross(data.cell[:, 0], data.cell[:, 1], dim=-1) + inv_min_dist_a3 = torch.norm(cross_a1a2 / cell_vol, p=2, dim=-1) + rep_a3 = torch.ceil(radius * inv_min_dist_a3) + else: + rep_a3 = data.cell.new_zeros(1) + + # Take the max over all images for uniformity. This is essentially padding. + # Note that this can significantly increase the number of computed distances + # if the required repetitions are very different between images + # (which they usually are). Changing this to sparse (scatter) operations + # might be worth the effort if this function becomes a bottleneck. + max_rep = [rep_a1.max(), rep_a2.max(), rep_a3.max()] + + # Tensor of unit cells + cells_per_dim = [ + torch.arange(-rep, rep + 1, device=device, dtype=torch.float) + for rep in max_rep + ] + unit_cell = torch.cartesian_prod(*cells_per_dim) + num_cells = len(unit_cell) + unit_cell_per_atom = unit_cell.view(1, num_cells, 3).repeat( + len(index2), 1, 1 + ) + unit_cell = torch.transpose(unit_cell, 0, 1) + unit_cell_batch = unit_cell.view(1, 3, num_cells).expand( + batch_size, -1, -1 + ) + + # Compute the x, y, z positional offsets for each cell in each image + data_cell = torch.transpose(data.cell, 1, 2) + pbc_offsets = torch.bmm(data_cell, unit_cell_batch) + pbc_offsets_per_atom = torch.repeat_interleave( + pbc_offsets, num_atoms_per_image_sqr, dim=0 + ) + + # Expand the positions and indices for the 9 cells + pos1 = pos1.view(-1, 3, 1).expand(-1, -1, num_cells) + pos2 = pos2.view(-1, 3, 1).expand(-1, -1, num_cells) + index1 = index1.view(-1, 1).repeat(1, num_cells).view(-1) + index2 = index2.view(-1, 1).repeat(1, num_cells).view(-1) + # Add the PBC offsets for the second atom + pos2 = pos2 + pbc_offsets_per_atom + + # Compute the squared distance between atoms + atom_distance_sqr = torch.sum((pos1 - pos2) ** 2, dim=1) + atom_distance_sqr = atom_distance_sqr.view(-1) + + # Remove pairs that are too far apart + mask_within_radius = torch.le(atom_distance_sqr, radius * radius) + # Remove pairs with the same atoms (distance = 0.0) + mask_not_same = torch.gt(atom_distance_sqr, 0.0001) + mask = torch.logical_and(mask_within_radius, mask_not_same) + index1 = torch.masked_select(index1, mask) + index2 = torch.masked_select(index2, mask) + unit_cell = torch.masked_select( + unit_cell_per_atom.view(-1, 3), mask.view(-1, 1).expand(-1, 3) + ) + unit_cell = unit_cell.view(-1, 3) + atom_distance_sqr = torch.masked_select(atom_distance_sqr, mask) + + mask_num_neighbors, num_neighbors_image = get_max_neighbors_mask( + natoms=data.natoms, + index=index1, + atom_distance=atom_distance_sqr, + max_num_neighbors_threshold=max_num_neighbors_threshold, + ) + + if not torch.all(mask_num_neighbors): + # Mask out the atoms to ensure each atom has at most max_num_neighbors_threshold neighbors + index1 = torch.masked_select(index1, mask_num_neighbors) + index2 = torch.masked_select(index2, mask_num_neighbors) + unit_cell = torch.masked_select( + unit_cell.view(-1, 3), mask_num_neighbors.view(-1, 1).expand(-1, 3) + ) + unit_cell = unit_cell.view(-1, 3) + + edge_index = torch.stack((index2, index1)) + + return edge_index, unit_cell, num_neighbors_image + + +def get_max_neighbors_mask( + natoms, index, atom_distance, max_num_neighbors_threshold +): + """ + Give a mask that filters out edges so that each atom has at most + `max_num_neighbors_threshold` neighbors. + Assumes that `index` is sorted. + """ + device = natoms.device + num_atoms = natoms.sum() + + # Get number of neighbors + # segment_coo assumes sorted index + ones = index.new_ones(1).expand_as(index) + num_neighbors = segment_coo(ones, index, dim_size=num_atoms) + max_num_neighbors = num_neighbors.max() + num_neighbors_thresholded = num_neighbors.clamp( + max=max_num_neighbors_threshold + ) + + # Get number of (thresholded) neighbors per image + image_indptr = torch.zeros( + natoms.shape[0] + 1, device=device, dtype=torch.long + ) + image_indptr[1:] = torch.cumsum(natoms, dim=0) + num_neighbors_image = segment_csr(num_neighbors_thresholded, image_indptr) + + # If max_num_neighbors is below the threshold, return early + if ( + max_num_neighbors <= max_num_neighbors_threshold + or max_num_neighbors_threshold <= 0 + ): + mask_num_neighbors = torch.tensor( + [True], dtype=bool, device=device + ).expand_as(index) + return mask_num_neighbors, num_neighbors_image + + # Create a tensor of size [num_atoms, max_num_neighbors] to sort the distances of the neighbors. + # Fill with infinity so we can easily remove unused distances later. + distance_sort = torch.full( + [num_atoms * max_num_neighbors], np.inf, device=device + ) + + # Create an index map to map distances from atom_distance to distance_sort + # index_sort_map assumes index to be sorted + index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors + index_neighbor_offset_expand = torch.repeat_interleave( + index_neighbor_offset, num_neighbors + ) + index_sort_map = ( + index * max_num_neighbors + + torch.arange(len(index), device=device) + - index_neighbor_offset_expand + ) + distance_sort.index_copy_(0, index_sort_map, atom_distance) + distance_sort = distance_sort.view(num_atoms, max_num_neighbors) + + # Sort neighboring atoms based on distance + distance_sort, index_sort = torch.sort(distance_sort, dim=1) + # Select the max_num_neighbors_threshold neighbors that are closest + distance_sort = distance_sort[:, :max_num_neighbors_threshold] + index_sort = index_sort[:, :max_num_neighbors_threshold] + + # Offset index_sort so that it indexes into index + index_sort = index_sort + index_neighbor_offset.view(-1, 1).expand( + -1, max_num_neighbors_threshold + ) + # Remove "unused pairs" with infinite distances + mask_finite = torch.isfinite(distance_sort) + index_sort = torch.masked_select(index_sort, mask_finite) + + # At this point index_sort contains the index into index of the + # closest max_num_neighbors_threshold neighbors per atom + # Create a mask to remove all pairs not in index_sort + mask_num_neighbors = torch.zeros(len(index), device=device, dtype=bool) + mask_num_neighbors.index_fill_(0, index_sort, True) + + return mask_num_neighbors, num_neighbors_image + + +def get_pruned_edge_idx(edge_index, num_atoms=None, max_neigh=1e9): + assert num_atoms is not None + + # removes neighbors > max_neigh + # assumes neighbors are sorted in increasing distance + _nonmax_idx = [] + for i in range(num_atoms): + idx_i = torch.arange(len(edge_index[1]))[(edge_index[1] == i)][ + :max_neigh + ] + _nonmax_idx.append(idx_i) + _nonmax_idx = torch.cat(_nonmax_idx) + + return _nonmax_idx + + +def merge_dicts(dict1: dict, dict2: dict): + """Recursively merge two dictionaries. + Values in dict2 override values in dict1. If dict1 and dict2 contain a dictionary as a + value, this will call itself recursively to merge these dictionaries. + This does not modify the input dictionaries (creates an internal copy). + Additionally returns a list of detected duplicates. + Adapted from https://github.com/TUM-DAML/seml/blob/master/seml/utils.py + + Parameters + ---------- + dict1: dict + First dict. + dict2: dict + Second dict. Values in dict2 will override values from dict1 in case they share the same key. + + Returns + ------- + return_dict: dict + Merged dictionaries. + """ + if not isinstance(dict1, dict): + raise ValueError(f"Expecting dict1 to be dict, found {type(dict1)}.") + if not isinstance(dict2, dict): + raise ValueError(f"Expecting dict2 to be dict, found {type(dict2)}.") + + return_dict = copy.deepcopy(dict1) + duplicates = [] + + for k, v in dict2.items(): + if k not in dict1: + return_dict[k] = v + else: + if isinstance(v, dict) and isinstance(dict1[k], dict): + return_dict[k], duplicates_k = merge_dicts(dict1[k], dict2[k]) + duplicates += [f"{k}.{dup}" for dup in duplicates_k] + else: + return_dict[k] = dict2[k] + duplicates.append(k) + + return return_dict, duplicates + + +class SeverityLevelBetween(logging.Filter): + def __init__(self, min_level, max_level): + super().__init__() + self.min_level = min_level + self.max_level = max_level + + def filter(self, record): + return self.min_level <= record.levelno < self.max_level + + +def setup_logging(): + root = logging.getLogger() + + # Perform setup only if logging has not been configured + if not root.hasHandlers(): + root.setLevel(logging.INFO) + + log_formatter = logging.Formatter( + "%(asctime)s (%(levelname)s): %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + # Send INFO to stdout + handler_out = logging.StreamHandler(sys.stdout) + handler_out.addFilter( + SeverityLevelBetween(logging.INFO, logging.WARNING) + ) + handler_out.setFormatter(log_formatter) + root.addHandler(handler_out) + + # Send WARNING (and higher) to stderr + handler_err = logging.StreamHandler(sys.stderr) + handler_err.setLevel(logging.WARNING) + handler_err.setFormatter(log_formatter) + root.addHandler(handler_err) + + +def compute_neighbors(data, edge_index): + # Get number of neighbors + # segment_coo assumes sorted index + ones = edge_index[1].new_ones(1).expand_as(edge_index[1]) + num_neighbors = segment_coo( + ones, edge_index[1], dim_size=data.natoms.sum() + ) + + # Get number of neighbors per image + image_indptr = torch.zeros( + data.natoms.shape[0] + 1, device=data.pos.device, dtype=torch.long + ) + image_indptr[1:] = torch.cumsum(data.natoms, dim=0) + neighbors = segment_csr(num_neighbors, image_indptr) + return neighbors + + +def check_traj_files(batch, traj_dir): + if traj_dir is None: + return False + traj_dir = Path(traj_dir) + traj_files = [traj_dir / f"{id}.traj" for id in batch[0].sid.tolist()] + return all(fl.exists() for fl in traj_files) + + +@contextmanager +def new_trainer_context(*, config: Dict[str, Any], args: Namespace): + from ocpmodels.common import distutils, gp_utils + from ocpmodels.common.registry import registry + + if TYPE_CHECKING: + from ocpmodels.tasks.task import BaseTask + from ocpmodels.trainers import BaseTrainer + + @dataclass + class _TrainingContext: + config: Dict[str, Any] + task: "BaseTask" + trainer: "BaseTrainer" + + setup_logging() + original_config = config + config = copy.deepcopy(original_config) + + if args.distributed: + distutils.setup(config) + if config["gp_gpus"] is not None: + gp_utils.setup_gp(config) + try: + setup_imports(config) + trainer_cls = registry.get_trainer_class( + config.get("trainer", "energy") + ) + assert trainer_cls is not None, "Trainer not found" + trainer = trainer_cls( + task=config["task"], + model=config["model"], + dataset=config["dataset"], + optimizer=config["optim"], + identifier=config["identifier"], + timestamp_id=config.get("timestamp_id", None), + run_dir=config.get("run_dir", "./"), + is_debug=config.get("is_debug", False), + print_every=config.get("print_every", 10), + seed=config.get("seed", 0), + logger=config.get("logger", "tensorboard"), + local_rank=config["local_rank"], + amp=config.get("amp", False), + cpu=config.get("cpu", False), + slurm=config.get("slurm", {}), + noddp=config.get("noddp", False), + ) + + task_cls = registry.get_task_class(config["mode"]) + assert task_cls is not None, "Task not found" + task = task_cls(config) + start_time = time.time() + ctx = _TrainingContext( + config=original_config, task=task, trainer=trainer + ) + yield ctx + distutils.synchronize() + if distutils.is_master(): + logging.info(f"Total time taken: {time.time() - start_time}") + finally: + if args.distributed: + distutils.cleanup() + + +def _resolve_scale_factor_submodule(model: nn.Module, name: str): + from ocpmodels.modules.scaling.scale_factor import ScaleFactor + + try: + scale = model.get_submodule(name) + if not isinstance(scale, ScaleFactor): + return None + return scale + except AttributeError: + return None + + +def _report_incompat_keys( + model: nn.Module, + keys: "_IncompatibleKeys", + strict: bool = False, +): + # filter out the missing scale factor keys for the new scaling factor module + missing_keys: List[str] = [] + for full_key_name in keys.missing_keys: + parent_module_name, _ = full_key_name.rsplit(".", 1) + scale_factor = _resolve_scale_factor_submodule( + model, parent_module_name + ) + if scale_factor is not None: + continue + missing_keys.append(full_key_name) + + # filter out unexpected scale factor keys that remain from the old scaling modules + unexpected_keys: List[str] = [] + for full_key_name in keys.unexpected_keys: + parent_module_name, _ = full_key_name.rsplit(".", 1) + scale_factor = _resolve_scale_factor_submodule( + model, parent_module_name + ) + if scale_factor is not None: + continue + unexpected_keys.append(full_key_name) + + error_msgs = [] + if len(unexpected_keys) > 0: + error_msgs.insert( + 0, + "Unexpected key(s) in state_dict: {}. ".format( + ", ".join('"{}"'.format(k) for k in unexpected_keys) + ), + ) + if len(missing_keys) > 0: + error_msgs.insert( + 0, + "Missing key(s) in state_dict: {}. ".format( + ", ".join('"{}"'.format(k) for k in missing_keys) + ), + ) + + if len(error_msgs) > 0: + error_msg = "Error(s) in loading state_dict for {}:\n\t{}".format( + model.__class__.__name__, "\n\t".join(error_msgs) + ) + if strict: + raise RuntimeError(error_msg) + else: + logging.warning(error_msg) + + return missing_keys, unexpected_keys + + +def load_state_dict( + module: nn.Module, + state_dict: Mapping[str, torch.Tensor], + strict: bool = True, +): + incompat_keys = module.load_state_dict(state_dict, strict=False) # type: ignore + return _report_incompat_keys(module, incompat_keys, strict=strict) + + +def scatter_det(*args, **kwargs): + from ocpmodels.common.registry import registry + + if registry.get("set_deterministic_scatter", no_warning=True): + torch.use_deterministic_algorithms(mode=True) + + out = scatter(*args, **kwargs) + + if registry.get("set_deterministic_scatter", no_warning=True): + torch.use_deterministic_algorithms(mode=False) + + return out diff --git a/magnet/eqV2/ocpmodels/models/__init__.py b/magnet/eqV2/ocpmodels/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/magnet/eqV2/ocpmodels/models/base.py b/magnet/eqV2/ocpmodels/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..e3a1014d6ea71f75b2dbcfe9856241ded8244c79 --- /dev/null +++ b/magnet/eqV2/ocpmodels/models/base.py @@ -0,0 +1,111 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. + +This source code is licensed under the MIT license found in the +LICENSE file in the root directory of this source tree. +""" + +import logging + +import torch +import torch.nn as nn +from torch_geometric.nn import radius_graph + +from ..common.utils import ( + compute_neighbors, + conditional_grad, + get_pbc_distances, + radius_graph_pbc, +) + + +class BaseModel(nn.Module): + def __init__(self, num_atoms=None, bond_feat_dim=None, num_targets=None): + super(BaseModel, self).__init__() + self.num_atoms = num_atoms + self.bond_feat_dim = bond_feat_dim + self.num_targets = num_targets + + def forward(self, data): + raise NotImplementedError + + def generate_graph( + self, + data, + cutoff=None, + max_neighbors=None, + use_pbc=None, + otf_graph=None, + ): + cutoff = cutoff or self.cutoff + max_neighbors = max_neighbors or self.max_neighbors + use_pbc = use_pbc or self.use_pbc + otf_graph = otf_graph or self.otf_graph + + if not otf_graph: + try: + edge_index = data.edge_index + + if use_pbc: + cell_offsets = data.cell_offsets + neighbors = data.neighbors + + except AttributeError: + logging.warning( + "Turning otf_graph=True as required attributes not present in data object" + ) + otf_graph = True + + if use_pbc: + if otf_graph: + edge_index, cell_offsets, neighbors = radius_graph_pbc( + data, cutoff, max_neighbors + ) + + out = get_pbc_distances( + data.pos, + edge_index, + data.cell, + cell_offsets, + neighbors, + return_offsets=True, + return_distance_vec=True, + ) + + edge_index = out["edge_index"] + edge_dist = out["distances"] + cell_offset_distances = out["offsets"] + distance_vec = out["distance_vec"] + else: + if otf_graph: + edge_index = radius_graph( + data.pos, + r=cutoff, + batch=data.batch, + max_num_neighbors=max_neighbors, + ) + + j, i = edge_index + distance_vec = data.pos[j] - data.pos[i] + + edge_dist = distance_vec.norm(dim=-1) + cell_offsets = torch.zeros( + edge_index.shape[1], 3, device=data.pos.device + ) + cell_offset_distances = torch.zeros_like( + cell_offsets, device=data.pos.device + ) + neighbors = compute_neighbors(data, edge_index) + + return ( + edge_index, + edge_dist, + distance_vec, + cell_offsets, + cell_offset_distances, + neighbors, + ) + + @property + def num_params(self): + return sum(p.numel() for p in self.parameters()) diff --git a/magnet/eqV2/ocpmodels/models/scn/__init__.py b/magnet/eqV2/ocpmodels/models/scn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/magnet/eqV2/ocpmodels/models/scn/sampling.py b/magnet/eqV2/ocpmodels/models/scn/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..452af82eb245a27d458d204145cd9f065c9654c5 --- /dev/null +++ b/magnet/eqV2/ocpmodels/models/scn/sampling.py @@ -0,0 +1,46 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. + +This source code is licensed under the MIT license found in the +LICENSE file in the root directory of this source tree. +""" +import math + +import torch + +### Methods for sample points on a sphere + + +def CalcSpherePoints(num_points, device): + goldenRatio = (1 + 5**0.5) / 2 + i = torch.arange(num_points, device=device).view(-1, 1) + theta = 2 * math.pi * i / goldenRatio + phi = torch.arccos(1 - 2 * (i + 0.5) / num_points) + points = torch.cat( + [ + torch.cos(theta) * torch.sin(phi), + torch.sin(theta) * torch.sin(phi), + torch.cos(phi), + ], + dim=1, + ) + + # weight the points by their density + pt_cross = points.view(1, -1, 3) - points.view(-1, 1, 3) + pt_cross = torch.sum(pt_cross**2, dim=2) + pt_cross = torch.exp(-pt_cross / (0.5 * 0.3)) + scalar = 1.0 / torch.sum(pt_cross, dim=1) + scalar = num_points * scalar / torch.sum(scalar) + return points * (scalar.view(-1, 1)) + + +def CalcSpherePointsRandom(num_points, device): + pts = 2.0 * (torch.rand(num_points, 3, device=device) - 0.5) + radius = torch.sum(pts**2, dim=1) + while torch.max(radius) > 1.0: + replace_pts = 2.0 * (torch.rand(num_points, 3, device=device) - 0.5) + replace_mask = radius.gt(0.99) + pts.masked_scatter_(replace_mask.view(-1, 1).repeat(1, 3), replace_pts) + radius = torch.sum(pts**2, dim=1) + + return pts / radius.view(-1, 1) diff --git a/magnet/eqV2/ocpmodels/models/scn/smearing.py b/magnet/eqV2/ocpmodels/models/scn/smearing.py new file mode 100644 index 0000000000000000000000000000000000000000..cd83b0fe384fc10fc07f4a962da016b7ce789486 --- /dev/null +++ b/magnet/eqV2/ocpmodels/models/scn/smearing.py @@ -0,0 +1,74 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. + +This source code is licensed under the MIT license found in the +LICENSE file in the root directory of this source tree. +""" + +import torch +import torch.nn as nn + + +# Different encodings for the atom distance embeddings +class GaussianSmearing(torch.nn.Module): + def __init__( + self, start=-5.0, stop=5.0, num_gaussians=50, basis_width_scalar=1.0 + ): + super(GaussianSmearing, self).__init__() + self.num_output = num_gaussians + offset = torch.linspace(start, stop, num_gaussians) + self.coeff = ( + -0.5 / (basis_width_scalar * (offset[1] - offset[0])).item() ** 2 + ) + self.register_buffer("offset", offset) + + def forward(self, dist): + dist = dist.view(-1, 1) - self.offset.view(1, -1) + return torch.exp(self.coeff * torch.pow(dist, 2)) + + +class SigmoidSmearing(torch.nn.Module): + def __init__( + self, start=-5.0, stop=5.0, num_sigmoid=50, basis_width_scalar=1.0 + ): + super(SigmoidSmearing, self).__init__() + self.num_output = num_sigmoid + offset = torch.linspace(start, stop, num_sigmoid) + self.coeff = (basis_width_scalar / (offset[1] - offset[0])).item() + self.register_buffer("offset", offset) + + def forward(self, dist): + exp_dist = self.coeff * (dist.view(-1, 1) - self.offset.view(1, -1)) + return torch.sigmoid(exp_dist) + + +class LinearSigmoidSmearing(torch.nn.Module): + def __init__( + self, start=-5.0, stop=5.0, num_sigmoid=50, basis_width_scalar=1.0 + ): + super(LinearSigmoidSmearing, self).__init__() + self.num_output = num_sigmoid + offset = torch.linspace(start, stop, num_sigmoid) + self.coeff = (basis_width_scalar / (offset[1] - offset[0])).item() + self.register_buffer("offset", offset) + + def forward(self, dist): + exp_dist = self.coeff * (dist.view(-1, 1) - self.offset.view(1, -1)) + x_dist = torch.sigmoid(exp_dist) + 0.001 * exp_dist + return x_dist + + +class SiLUSmearing(torch.nn.Module): + def __init__( + self, start=-5.0, stop=5.0, num_output=50, basis_width_scalar=1.0 + ): + super(SiLUSmearing, self).__init__() + self.num_output = num_output + self.fc1 = nn.Linear(2, num_output) + self.act = nn.SiLU() + + def forward(self, dist): + x_dist = dist.view(-1, 1) + x_dist = torch.cat([x_dist, torch.ones_like(x_dist)], dim=1) + x_dist = self.act(self.fc1(x_dist)) + return x_dist diff --git a/magnet/eqV2/radial_function.py b/magnet/eqV2/radial_function.py new file mode 100644 index 0000000000000000000000000000000000000000..8b943d381d76ea31843ba5713adce90cb4cfc032 --- /dev/null +++ b/magnet/eqV2/radial_function.py @@ -0,0 +1,31 @@ +import torch +import torch.nn as nn + + +class RadialFunction(nn.Module): + ''' + Contruct a radial function (linear layers + layer normalization + SiLU) given a list of channels + ''' + def __init__(self, channels_list): + super().__init__() + modules = [] + input_channels = channels_list[0] + for i in range(len(channels_list)): + if i == 0: + continue + + modules.append(nn.Linear(input_channels, channels_list[i], bias=True)) + input_channels = channels_list[i] + + if i == len(channels_list) - 1: + break + + modules.append(nn.LayerNorm(channels_list[i])) + modules.append(torch.nn.SiLU()) + + self.net = nn.Sequential(*modules) + + + def forward(self, inputs): + return self.net(inputs) + \ No newline at end of file diff --git a/magnet/eqV2/so2_ops.py b/magnet/eqV2/so2_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..f70e011e55bccb49b3fe59d4df31d5a601dbbee6 --- /dev/null +++ b/magnet/eqV2/so2_ops.py @@ -0,0 +1,333 @@ +import torch +import torch.nn as nn +import math +import copy + +from torch.nn import Linear +from .so3 import SO3_Embedding +from .radial_function import RadialFunction + + +class SO2_m_Convolution(torch.nn.Module): + """ + SO(2) Conv: Perform an SO(2) convolution on features corresponding to +- m + + Args: + m (int): Order of the spherical harmonic coefficients + sphere_channels (int): Number of spherical channels + m_output_channels (int): Number of output channels used during the SO(2) conv + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + """ + def __init__( + self, + m, + sphere_channels, + m_output_channels, + lmax_list, + mmax_list + ): + super(SO2_m_Convolution, self).__init__() + + self.m = m + self.sphere_channels = sphere_channels + self.m_output_channels = m_output_channels + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.num_resolutions = len(self.lmax_list) + + num_channels = 0 + for i in range(self.num_resolutions): + num_coefficents = 0 + if self.mmax_list[i] >= self.m: + num_coefficents = self.lmax_list[i] - self.m + 1 + num_channels = num_channels + num_coefficents * self.sphere_channels + assert num_channels > 0 + + self.fc = Linear(num_channels, + 2 * self.m_output_channels * (num_channels // self.sphere_channels), + bias=False) + self.fc.weight.data.mul_(1 / math.sqrt(2)) + + + def forward(self, x_m): + x_m = self.fc(x_m) + x_r = x_m.narrow(2, 0, self.fc.out_features // 2) + x_i = x_m.narrow(2, self.fc.out_features // 2, self.fc.out_features // 2) + x_m_r = x_r.narrow(1, 0, 1) - x_i.narrow(1, 1, 1) #x_r[:, 0] - x_i[:, 1] + x_m_i = x_r.narrow(1, 1, 1) + x_i.narrow(1, 0, 1) #x_r[:, 1] + x_i[:, 0] + x_out = torch.cat((x_m_r, x_m_i), dim=1) + + return x_out + + +class SO2_Convolution(torch.nn.Module): + """ + SO(2) Block: Perform SO(2) convolutions for all m (orders) + + Args: + sphere_channels (int): Number of spherical channels + m_output_channels (int): Number of output channels used during the SO(2) conv + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + mappingReduced (CoefficientMappingModule): Used to extract a subset of m components + internal_weights (bool): If True, not using radial function to multiply inputs features + edge_channels_list (list:int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels]. + extra_m0_output_channels (int): If not None, return `out_embedding` (SO3_Embedding) and `extra_m0_features` (Tensor). + """ + def __init__( + self, + sphere_channels, + m_output_channels, + lmax_list, + mmax_list, + mappingReduced, + internal_weights=True, + edge_channels_list=None, + extra_m0_output_channels=None + ): + super(SO2_Convolution, self).__init__() + self.sphere_channels = sphere_channels + self.m_output_channels = m_output_channels + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.mappingReduced = mappingReduced + self.num_resolutions = len(lmax_list) + self.internal_weights = internal_weights + self.edge_channels_list = copy.deepcopy(edge_channels_list) + self.extra_m0_output_channels = extra_m0_output_channels + + num_channels_rad = 0 # for radial function + + num_channels_m0 = 0 + for i in range(self.num_resolutions): + num_coefficients = self.lmax_list[i] + 1 + num_channels_m0 = num_channels_m0 + num_coefficients * self.sphere_channels + + # SO(2) convolution for m = 0 + m0_output_channels = self.m_output_channels * (num_channels_m0 // self.sphere_channels) + if self.extra_m0_output_channels is not None: + m0_output_channels = m0_output_channels + self.extra_m0_output_channels + self.fc_m0 = Linear(num_channels_m0, m0_output_channels) + num_channels_rad = num_channels_rad + self.fc_m0.in_features + + # SO(2) convolution for non-zero m + self.so2_m_conv = nn.ModuleList() + for m in range(1, max(self.mmax_list) + 1): + self.so2_m_conv.append( + SO2_m_Convolution( + m, + self.sphere_channels, + self.m_output_channels, + self.lmax_list, + self.mmax_list, + ) + ) + num_channels_rad = num_channels_rad + self.so2_m_conv[-1].fc.in_features + + # Embedding function of distance + self.rad_func = None + if not self.internal_weights: + assert self.edge_channels_list is not None + self.edge_channels_list.append(int(num_channels_rad)) + self.rad_func = RadialFunction(self.edge_channels_list) + + + def forward(self, x, x_edge): + + num_edges = len(x_edge) + out = [] + + # Reshape the spherical harmonics based on m (order) + x._m_primary(self.mappingReduced) + + # radial function + if self.rad_func is not None: + x_edge = self.rad_func(x_edge) + offset_rad = 0 + + # Compute m=0 coefficients separately since they only have real values (no imaginary) + x_0 = x.embedding.narrow(1, 0, self.mappingReduced.m_size[0]) + x_0 = x_0.reshape(num_edges, -1) + if self.rad_func is not None: + x_edge_0 = x_edge.narrow(1, 0, self.fc_m0.in_features) + x_0 = x_0 * x_edge_0 + x_0 = self.fc_m0(x_0) + + x_0_extra = None + # extract extra m0 features + if self.extra_m0_output_channels is not None: + x_0_extra = x_0.narrow(-1, 0, self.extra_m0_output_channels) + x_0 = x_0.narrow(-1, self.extra_m0_output_channels, (self.fc_m0.out_features - self.extra_m0_output_channels)) + + x_0 = x_0.view(num_edges, -1, self.m_output_channels) + #x.embedding[:, 0 : self.mappingReduced.m_size[0]] = x_0 + out.append(x_0) + offset_rad = offset_rad + self.fc_m0.in_features + + # Compute the values for the m > 0 coefficients + offset = self.mappingReduced.m_size[0] + for m in range(1, max(self.mmax_list) + 1): + # Get the m order coefficients + x_m = x.embedding.narrow(1, offset, 2 * self.mappingReduced.m_size[m]) + x_m = x_m.reshape(num_edges, 2, -1) + + # Perform SO(2) convolution + if self.rad_func is not None: + x_edge_m = x_edge.narrow(1, offset_rad, self.so2_m_conv[m - 1].fc.in_features) + x_edge_m = x_edge_m.reshape(num_edges, 1, self.so2_m_conv[m - 1].fc.in_features) + x_m = x_m * x_edge_m + x_m = self.so2_m_conv[m - 1](x_m) + x_m = x_m.view(num_edges, -1, self.m_output_channels) + #x.embedding[:, offset : offset + 2 * self.mappingReduced.m_size[m]] = x_m + out.append(x_m) + offset = offset + 2 * self.mappingReduced.m_size[m] + offset_rad = offset_rad + self.so2_m_conv[m - 1].fc.in_features + + out = torch.cat(out, dim=1) + out_embedding = SO3_Embedding( + 0, + x.lmax_list.copy(), + self.m_output_channels, + device=x.device, + dtype=x.dtype + ) + out_embedding.set_embedding(out) + out_embedding.set_lmax_mmax(self.lmax_list.copy(), self.mmax_list.copy()) + + # Reshape the spherical harmonics based on l (degree) + out_embedding._l_primary(self.mappingReduced) + + if self.extra_m0_output_channels is not None: + return out_embedding, x_0_extra + else: + return out_embedding + + +class SO2_Linear(torch.nn.Module): + """ + SO(2) Linear: Perform SO(2) linear for all m (orders). + + Args: + sphere_channels (int): Number of spherical channels + m_output_channels (int): Number of output channels used during the SO(2) conv + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + mappingReduced (CoefficientMappingModule): Used to extract a subset of m components + internal_weights (bool): If True, not using radial function to multiply inputs features + edge_channels_list (list:int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels]. + """ + def __init__( + self, + sphere_channels, + m_output_channels, + lmax_list, + mmax_list, + mappingReduced, + internal_weights=False, + edge_channels_list=None, + ): + super(SO2_Linear, self).__init__() + self.sphere_channels = sphere_channels + self.m_output_channels = m_output_channels + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.mappingReduced = mappingReduced + self.internal_weights = internal_weights + self.edge_channels_list = copy.deepcopy(edge_channels_list) + self.num_resolutions = len(lmax_list) + + num_channels_rad = 0 + + num_channels_m0 = 0 + for i in range(self.num_resolutions): + num_coefficients = self.lmax_list[i] + 1 + num_channels_m0 = num_channels_m0 + num_coefficients * self.sphere_channels + + # SO(2) linear for m = 0 + self.fc_m0 = Linear(num_channels_m0, + self.m_output_channels * (num_channels_m0 // self.sphere_channels)) + num_channels_rad = num_channels_rad + self.fc_m0.in_features + + # SO(2) linear for non-zero m + self.so2_m_fc = nn.ModuleList() + for m in range(1, max(self.mmax_list) + 1): + num_in_channels = 0 + for i in range(self.num_resolutions): + num_coefficents = 0 + if self.mmax_list[i] >= m: + num_coefficents = self.lmax_list[i] - m + 1 + num_in_channels = num_in_channels + num_coefficents * self.sphere_channels + assert num_in_channels > 0 + fc = Linear(num_in_channels, + self.m_output_channels * (num_in_channels // self.sphere_channels), + bias=False) + num_channels_rad = num_channels_rad + fc.in_features + self.so2_m_fc.append(fc) + + # Embedding function of distance + self.rad_func = None + if not self.internal_weights: + assert self.edge_channels_list is not None + self.edge_channels_list.append(int(num_channels_rad)) + self.rad_func = RadialFunction(self.edge_channels_list) + + + def forward(self, x, x_edge): + + batch_size = x.embedding.shape[0] + out = [] + + # Reshape the spherical harmonics based on m (order) + x._m_primary(self.mappingReduced) + + # radial function + if self.rad_func is not None: + x_edge = self.rad_func(x_edge) + offset_rad = 0 + + # Compute m=0 coefficients separately since they only have real values (no imaginary) + x_0 = x.embedding.narrow(1, 0, self.mappingReduced.m_size[0]) + x_0 = x_0.reshape(batch_size, -1) + if self.rad_func is not None: + x_edge_0 = x_edge.narrow(1, 0, self.fc_m0.in_features) + x_0 = x_0 * x_edge_0 + x_0 = self.fc_m0(x_0) + x_0 = x_0.view(batch_size, -1, self.m_output_channels) + out.append(x_0) + offset_rad = offset_rad + self.fc_m0.in_features + + # Compute the values for the m > 0 coefficients + offset = self.mappingReduced.m_size[0] + for m in range(1, max(self.mmax_list) + 1): + # Get the m order coefficients + x_m = x.embedding.narrow(1, offset, 2 * self.mappingReduced.m_size[m]) + x_m = x_m.reshape(batch_size, 2, -1) + if self.rad_func is not None: + x_edge_m = x_edge.narrow(1, offset_rad, self.so2_m_fc[m - 1].in_features) + x_edge_m = x_edge_m.reshape(batch_size, 1, self.so2_m_fc[m - 1].in_features) + x_m = x_m * x_edge_m + + # Perform SO(2) linear + x_m = self.so2_m_fc[m - 1](x_m) + x_m = x_m.view(batch_size, -1, self.m_output_channels) + out.append(x_m) + + offset = offset + 2 * self.mappingReduced.m_size[m] + offset_rad = offset_rad + self.so2_m_fc[m - 1].in_features + + out = torch.cat(out, dim=1) + out_embedding = SO3_Embedding( + 0, + x.lmax_list.copy(), + self.m_output_channels, + device=x.device, + dtype=x.dtype + ) + out_embedding.set_embedding(out) + out_embedding.set_lmax_mmax(self.lmax_list.copy(), self.mmax_list.copy()) + + # Reshape the spherical harmonics based on l (degree) + out_embedding._l_primary(self.mappingReduced) + + return out_embedding \ No newline at end of file diff --git a/magnet/eqV2/so3.py b/magnet/eqV2/so3.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0fd9015ab63f8b6278ff5cd0cf9a4d359649a2 --- /dev/null +++ b/magnet/eqV2/so3.py @@ -0,0 +1,670 @@ +""" +Copyright (c) Facebook, Inc. and its affiliates. + +This source code is licensed under the MIT license found in the +LICENSE file in the root directory of this source tree. + + +TODO: + 1. Simplify the case when `num_resolutions` == 1. + 2. Remove indexing when the shape is the same. + 3. Move some functions outside classes and to separate files. +""" + +import os +import math +import torch +import torch.nn as nn + +try: + from e3nn import o3 + from e3nn.o3 import FromS2Grid, ToS2Grid +except ImportError: + pass + +from .wigner import wigner_D +from torch.nn import Linear + + +class CoefficientMappingModule(torch.nn.Module): + """ + Helper module for coefficients used to reshape l <--> m and to get coefficients of specific degree or order + + Args: + lmax_list (list:int): List of maximum degree of the spherical harmonics + mmax_list (list:int): List of maximum order of the spherical harmonics + """ + + def __init__( + self, + lmax_list, + mmax_list, + ): + super().__init__() + + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.num_resolutions = len(lmax_list) + + # Temporarily use `cpu` as device and this will be overwritten. + self.device = 'cpu' + + # Compute the degree (l) and order (m) for each entry of the embedding + l_harmonic = torch.tensor([], device=self.device).long() + m_harmonic = torch.tensor([], device=self.device).long() + m_complex = torch.tensor([], device=self.device).long() + + res_size = torch.zeros([self.num_resolutions], device=self.device).long() + + offset = 0 + for i in range(self.num_resolutions): + for l in range(0, self.lmax_list[i] + 1): + mmax = min(self.mmax_list[i], l) + m = torch.arange(-mmax, mmax + 1, device=self.device).long() + m_complex = torch.cat([m_complex, m], dim=0) + m_harmonic = torch.cat( + [m_harmonic, torch.abs(m).long()], dim=0 + ) + l_harmonic = torch.cat( + [l_harmonic, m.fill_(l).long()], dim=0 + ) + res_size[i] = len(l_harmonic) - offset + offset = len(l_harmonic) + + num_coefficients = len(l_harmonic) + # `self.to_m` moves m components from different L to contiguous index + to_m = torch.zeros([num_coefficients, num_coefficients], device=self.device) + m_size = torch.zeros([max(self.mmax_list) + 1], device=self.device).long() + + # The following is implemented poorly - very slow. It only gets called + # a few times so haven't optimized. + offset = 0 + for m in range(max(self.mmax_list) + 1): + idx_r, idx_i = self.complex_idx(m, -1, m_complex, l_harmonic) + + for idx_out, idx_in in enumerate(idx_r): + to_m[idx_out + offset, idx_in] = 1.0 + offset = offset + len(idx_r) + + m_size[m] = int(len(idx_r)) + + for idx_out, idx_in in enumerate(idx_i): + to_m[idx_out + offset, idx_in] = 1.0 + offset = offset + len(idx_i) + + to_m = to_m.detach() + + # save tensors and they will be moved to GPU + self.register_buffer('l_harmonic', l_harmonic) + self.register_buffer('m_harmonic', m_harmonic) + self.register_buffer('m_complex', m_complex) + self.register_buffer('res_size', res_size) + self.register_buffer('to_m', to_m) + self.register_buffer('m_size', m_size) + + # for caching the output of `coefficient_idx` + self.lmax_cache, self.mmax_cache = None, None + self.mask_indices_cache = None + self.rotate_inv_rescale_cache = None + + + # Return mask containing coefficients of order m (real and imaginary parts) + def complex_idx(self, m, lmax, m_complex, l_harmonic): + ''' + Add `m_complex` and `l_harmonic` to the input arguments + since we cannot use `self.m_complex`. + ''' + if lmax == -1: + lmax = max(self.lmax_list) + + indices = torch.arange(len(l_harmonic), device=self.device) + # Real part + mask_r = torch.bitwise_and( + l_harmonic.le(lmax), m_complex.eq(m) + ) + mask_idx_r = torch.masked_select(indices, mask_r) + + mask_idx_i = torch.tensor([], device=self.device).long() + # Imaginary part + if m != 0: + mask_i = torch.bitwise_and( + l_harmonic.le(lmax), m_complex.eq(-m) + ) + mask_idx_i = torch.masked_select(indices, mask_i) + + return mask_idx_r, mask_idx_i + + + # Return mask containing coefficients less than or equal to degree (l) and order (m) + def coefficient_idx(self, lmax, mmax): + + if (self.lmax_cache is not None) and (self.mmax_cache is not None): + if (self.lmax_cache == lmax) and (self.mmax_cache == mmax): + if self.mask_indices_cache is not None: + return self.mask_indices_cache + + mask = torch.bitwise_and( + self.l_harmonic.le(lmax), self.m_harmonic.le(mmax) + ) + self.device = mask.device + indices = torch.arange(len(mask), device=self.device) + mask_indices = torch.masked_select(indices, mask) + self.lmax_cache, self.mmax_cache = lmax, mmax + self.mask_indices_cache = mask_indices + return self.mask_indices_cache + + + # Return the re-scaling for rotating back to original frame + # this is required since we only use a subset of m components for SO(2) convolution + def get_rotate_inv_rescale(self, lmax, mmax): + + if (self.lmax_cache is not None) and (self.mmax_cache is not None): + if (self.lmax_cache == lmax) and (self.mmax_cache == mmax): + if self.rotate_inv_rescale_cache is not None: + return self.rotate_inv_rescale_cache + + if self.mask_indices_cache is None: + self.coefficient_idx(lmax, mmax) + + rotate_inv_rescale = torch.ones((1, (lmax + 1)**2, (lmax + 1)**2), device=self.device) + for l in range(lmax + 1): + if l <= mmax: + continue + start_idx = l ** 2 + length = 2 * l + 1 + rescale_factor = math.sqrt(length / (2 * mmax + 1)) + rotate_inv_rescale[:, start_idx : (start_idx + length), start_idx : (start_idx + length)] = rescale_factor + rotate_inv_rescale = rotate_inv_rescale[:, :, self.mask_indices_cache] + self.rotate_inv_rescale_cache = rotate_inv_rescale + return self.rotate_inv_rescale_cache + + + def __repr__(self): + return f"{self.__class__.__name__}(lmax_list={self.lmax_list}, mmax_list={self.mmax_list})" + + +class SO3_Embedding(): + """ + Helper functions for performing operations on irreps embedding + + Args: + length (int): Batch size + lmax_list (list:int): List of maximum degree of the spherical harmonics + num_channels (int): Number of channels + device: Device of the output + dtype: type of the output tensors + """ + + def __init__( + self, + length, + lmax_list, + num_channels, + device, + dtype, + ): + super().__init__() + self.num_channels = num_channels + self.device = device + self.dtype = dtype + self.num_resolutions = len(lmax_list) + + self.num_coefficients = 0 + for i in range(self.num_resolutions): + self.num_coefficients = self.num_coefficients + int( + (lmax_list[i] + 1) ** 2 + ) + + embedding = torch.zeros( + length, + self.num_coefficients, + self.num_channels, + device=self.device, + dtype=self.dtype, + ) + + self.set_embedding(embedding) + self.set_lmax_mmax(lmax_list, lmax_list.copy()) + + + # Clone an embedding of irreps + def clone(self): + clone = SO3_Embedding( + 0, + self.lmax_list.copy(), + self.num_channels, + self.device, + self.dtype, + ) + clone.set_embedding(self.embedding.clone()) + return clone + + + # Initialize an embedding of irreps + def set_embedding(self, embedding): + self.length = len(embedding) + self.embedding = embedding + + + # Set the maximum order to be the maximum degree + def set_lmax_mmax(self, lmax_list, mmax_list): + self.lmax_list = lmax_list + self.mmax_list = mmax_list + + + # Expand the node embeddings to the number of edges + def _expand_edge(self, edge_index): + embedding = self.embedding[edge_index] + self.set_embedding(embedding) + + + # Initialize an embedding of irreps of a neighborhood + def expand_edge(self, edge_index): + x_expand = SO3_Embedding( + 0, + self.lmax_list.copy(), + self.num_channels, + self.device, + self.dtype, + ) + x_expand.set_embedding(self.embedding[edge_index]) + return x_expand + + + # Compute the sum of the embeddings of the neighborhood + def _reduce_edge(self, edge_index, num_nodes): + new_embedding = torch.zeros( + num_nodes, + self.num_coefficients, + self.num_channels, + device=self.embedding.device, + dtype=self.embedding.dtype, + ) + new_embedding.index_add_(0, edge_index, self.embedding) + self.set_embedding(new_embedding) + + + # Reshape the embedding l -> m + def _m_primary(self, mapping): + self.embedding = torch.einsum("nac, ba -> nbc", self.embedding, mapping.to_m) + + + # Reshape the embedding m -> l + def _l_primary(self, mapping): + self.embedding = torch.einsum("nac, ab -> nbc", self.embedding, mapping.to_m) + + + # Rotate the embedding + def _rotate(self, SO3_rotation, lmax_list, mmax_list): + + if self.num_resolutions == 1: + embedding_rotate = SO3_rotation[0].rotate(self.embedding, lmax_list[0], mmax_list[0]) + else: + offset = 0 + embedding_rotate = torch.tensor([], device=self.device, dtype=self.dtype) + for i in range(self.num_resolutions): + num_coefficients = int((self.lmax_list[i] + 1) ** 2) + embedding_i = self.embedding[:, offset : offset + num_coefficients] + embedding_rotate = torch.cat([ + embedding_rotate, + SO3_rotation[i].rotate(embedding_i, lmax_list[i], mmax_list[i])], + dim=1) + offset = offset + num_coefficients + + self.embedding = embedding_rotate + self.set_lmax_mmax(lmax_list.copy(), mmax_list.copy()) + + + # Rotate the embedding by the inverse of the rotation matrix + def _rotate_inv(self, SO3_rotation, mappingReduced): + + if self.num_resolutions == 1: + embedding_rotate = SO3_rotation[0].rotate_inv(self.embedding, self.lmax_list[0], self.mmax_list[0]) + else: + offset = 0 + embedding_rotate = torch.tensor([], device=self.device, dtype=self.dtype) + for i in range(self.num_resolutions): + num_coefficients = mappingReduced.res_size[i] + embedding_i = self.embedding[:, offset : offset + num_coefficients] + embedding_rotate = torch.cat([ + embedding_rotate, + SO3_rotation[i].rotate_inv(embedding_i, self.lmax_list[i], self.mmax_list[i])], + dim=1) + offset = offset + num_coefficients + self.embedding = embedding_rotate + + # Assume mmax = lmax when rotating back + for i in range(self.num_resolutions): + self.mmax_list[i] = int(self.lmax_list[i]) + self.set_lmax_mmax(self.lmax_list, self.mmax_list) + + + # Compute point-wise spherical non-linearity + def _grid_act(self, SO3_grid, act, mappingReduced): + offset = 0 + for i in range(self.num_resolutions): + + num_coefficients = mappingReduced.res_size[i] + + if self.num_resolutions == 1: + x_res = self.embedding + else: + x_res = self.embedding[:, offset : offset + num_coefficients].contiguous() + to_grid_mat = SO3_grid[self.lmax_list[i]][self.mmax_list[i]].get_to_grid_mat(self.device) + from_grid_mat = SO3_grid[self.lmax_list[i]][self.mmax_list[i]].get_from_grid_mat(self.device) + + x_grid = torch.einsum("bai, zic -> zbac", to_grid_mat, x_res) + x_grid = act(x_grid) + x_res = torch.einsum("bai, zbac -> zic", from_grid_mat, x_grid) + if self.num_resolutions == 1: + self.embedding = x_res + else: + self.embedding[:, offset : offset + num_coefficients] = x_res + offset = offset + num_coefficients + + + # Compute a sample of the grid + def to_grid(self, SO3_grid, lmax=-1): + if lmax == -1: + lmax = max(self.lmax_list) + + to_grid_mat_lmax = SO3_grid[lmax][lmax].get_to_grid_mat(self.device) + grid_mapping = SO3_grid[lmax][lmax].mapping + + offset = 0 + x_grid = torch.tensor([], device=self.device) + + for i in range(self.num_resolutions): + num_coefficients = int((self.lmax_list[i] + 1) ** 2) + if self.num_resolutions == 1: + x_res = self.embedding + else: + x_res = self.embedding[:, offset : offset + num_coefficients].contiguous() + to_grid_mat = to_grid_mat_lmax[:, :, grid_mapping.coefficient_idx(self.lmax_list[i], self.lmax_list[i])] + x_grid = torch.cat([x_grid, torch.einsum("bai, zic -> zbac", to_grid_mat, x_res)], dim=3) + offset = offset + num_coefficients + + return x_grid + + + # Compute irreps from grid representation + def _from_grid(self, x_grid, SO3_grid, lmax=-1): + if lmax == -1: + lmax = max(self.lmax_list) + + from_grid_mat_lmax = SO3_grid[lmax][lmax].get_from_grid_mat(self.device) + grid_mapping = SO3_grid[lmax][lmax].mapping + + offset = 0 + offset_channel = 0 + for i in range(self.num_resolutions): + from_grid_mat = from_grid_mat_lmax[:, :, grid_mapping.coefficient_idx(self.lmax_list[i], self.lmax_list[i])] + if self.num_resolutions == 1: + temp = x_grid + else: + temp = x_grid[:, :, :, offset_channel : offset_channel + self.num_channels] + x_res = torch.einsum("bai, zbac -> zic", from_grid_mat, temp) + num_coefficients = int((self.lmax_list[i] + 1) ** 2) + + if self.num_resolutions == 1: + self.embedding = x_res + else: + self.embedding[:, offset : offset + num_coefficients] = x_res + + offset = offset + num_coefficients + offset_channel = offset_channel + self.num_channels + + +class SO3_Rotation(torch.nn.Module): + """ + Helper functions for Wigner-D rotations + + Args: + lmax_list (list:int): List of maximum degree of the spherical harmonics + """ + + def __init__( + self, + lmax, + ): + super().__init__() + self.lmax = lmax + self.mapping = CoefficientMappingModule([self.lmax], [self.lmax]) + + + def set_wigner(self, rot_mat3x3): + self.device, self.dtype = rot_mat3x3.device, rot_mat3x3.dtype + length = len(rot_mat3x3) + self.wigner = self.RotationToWignerDMatrix(rot_mat3x3, 0, self.lmax) + self.wigner_inv = torch.transpose(self.wigner, 1, 2).contiguous() + self.wigner = self.wigner.detach() + self.wigner_inv = self.wigner_inv.detach() + + + # Rotate the embedding + def rotate(self, embedding, out_lmax, out_mmax): + out_mask = self.mapping.coefficient_idx(out_lmax, out_mmax) + wigner = self.wigner[:, out_mask, :] + return torch.bmm(wigner, embedding) + + + # Rotate the embedding by the inverse of the rotation matrix + def rotate_inv(self, embedding, in_lmax, in_mmax): + in_mask = self.mapping.coefficient_idx(in_lmax, in_mmax) + wigner_inv = self.wigner_inv[:, :, in_mask] + wigner_inv_rescale = self.mapping.get_rotate_inv_rescale(in_lmax, in_mmax) + wigner_inv = wigner_inv * wigner_inv_rescale + return torch.bmm(wigner_inv, embedding) + + + # Compute Wigner matrices from rotation matrix + def RotationToWignerDMatrix(self, edge_rot_mat, start_lmax, end_lmax): + x = edge_rot_mat @ edge_rot_mat.new_tensor([0.0, 1.0, 0.0]) + alpha, beta = o3.xyz_to_angles(x) + R = ( + o3.angles_to_matrix( + alpha, beta, torch.zeros_like(alpha) + ).transpose(-1, -2) + @ edge_rot_mat + ) + gamma = torch.atan2(R[..., 0, 2], R[..., 0, 0]) + + size = (end_lmax + 1) ** 2 - (start_lmax) ** 2 + wigner = torch.zeros(len(alpha), size, size, device=self.device) + start = 0 + for lmax in range(start_lmax, end_lmax + 1): + block = wigner_D(lmax, alpha, beta, gamma) + end = start + block.size()[1] + wigner[:, start:end, start:end] = block + start = end + + return wigner.detach() + + +class SO3_Grid(torch.nn.Module): + """ + Helper functions for grid representation of the irreps + + Args: + lmax (int): Maximum degree of the spherical harmonics + mmax (int): Maximum order of the spherical harmonics + """ + + def __init__( + self, + lmax, + mmax, + normalization='integral', + resolution=None, + ): + super().__init__() + self.lmax = lmax + self.mmax = mmax + self.lat_resolution = 2 * (self.lmax + 1) + if lmax == mmax: + self.long_resolution = 2 * (self.mmax + 1) + 1 + else: + self.long_resolution = 2 * (self.mmax) + 1 + if resolution is not None: + self.lat_resolution = resolution + self.long_resolution = resolution + + self.mapping = CoefficientMappingModule([self.lmax], [self.lmax]) + + device = 'cpu' + + to_grid = ToS2Grid( + self.lmax, + (self.lat_resolution, self.long_resolution), + normalization=normalization, #normalization="integral", + device=device, + ) + to_grid_mat = torch.einsum("mbi, am -> bai", to_grid.shb, to_grid.sha).detach() + # rescale based on mmax + if lmax != mmax: + for l in range(lmax + 1): + if l <= mmax: + continue + start_idx = l ** 2 + length = 2 * l + 1 + rescale_factor = math.sqrt(length / (2 * mmax + 1)) + to_grid_mat[:, :, start_idx : (start_idx + length)] = to_grid_mat[:, :, start_idx : (start_idx + length)] * rescale_factor + to_grid_mat = to_grid_mat[:, :, self.mapping.coefficient_idx(self.lmax, self.mmax)] + + from_grid = FromS2Grid( + (self.lat_resolution, self.long_resolution), + self.lmax, + normalization=normalization, #normalization="integral", + device=device, + ) + from_grid_mat = torch.einsum("am, mbi -> bai", from_grid.sha, from_grid.shb).detach() + # rescale based on mmax + if lmax != mmax: + for l in range(lmax + 1): + if l <= mmax: + continue + start_idx = l ** 2 + length = 2 * l + 1 + rescale_factor = math.sqrt(length / (2 * mmax + 1)) + from_grid_mat[:, :, start_idx : (start_idx + length)] = from_grid_mat[:, :, start_idx : (start_idx + length)] * rescale_factor + from_grid_mat = from_grid_mat[:, :, self.mapping.coefficient_idx(self.lmax, self.mmax)] + + # save tensors and they will be moved to GPU + self.register_buffer('to_grid_mat', to_grid_mat) + self.register_buffer('from_grid_mat', from_grid_mat) + + + # Compute matrices to transform irreps to grid + def get_to_grid_mat(self, device): + return self.to_grid_mat + + + # Compute matrices to transform grid to irreps + def get_from_grid_mat(self, device): + return self.from_grid_mat + + + # Compute grid from irreps representation + def to_grid(self, embedding, lmax, mmax): + to_grid_mat = self.to_grid_mat[:, :, self.mapping.coefficient_idx(lmax, mmax)] + grid = torch.einsum("bai, zic -> zbac", to_grid_mat, embedding) + return grid + + + # Compute irreps from grid representation + def from_grid(self, grid, lmax, mmax): + from_grid_mat = self.from_grid_mat[:, :, self.mapping.coefficient_idx(lmax, mmax)] + embedding = torch.einsum("bai, zbac -> zic", from_grid_mat, grid) + return embedding + + +class SO3_Linear(torch.nn.Module): + def __init__(self, in_features, out_features, lmax, bias=True): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.lmax = lmax + self.linear_list = torch.nn.ModuleList() + for l in range(lmax + 1): + if l == 0: + self.linear_list.append(Linear(in_features, out_features, bias=bias)) + else: + self.linear_list.append(Linear(in_features, out_features, bias=False)) + + + def forward(self, input_embedding, output_scale=None): + out = [] + for l in range(self.lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + features = input_embedding.embedding.narrow(1, start_idx, length) + features = self.linear_list[l](features) + if output_scale is not None: + scale = output_scale.narrow(1, l, 1) + features = features * scale + out.append(features) + out = torch.cat(out, dim=1) + + out_embedding = SO3_Embedding( + 0, + input_embedding.lmax_list.copy(), + self.out_features, + device=input_embedding.device, + dtype=input_embedding.dtype + ) + out_embedding.set_embedding(out) + out_embedding.set_lmax_mmax(input_embedding.lmax_list.copy(), input_embedding.lmax_list.copy()) + + return out_embedding + + + def __repr__(self): + return f"{self.__class__.__name__}(in_features={self.in_features}, out_features={self.out_features}, lmax={self.lmax})" + + +class SO3_LinearV2(torch.nn.Module): + def __init__(self, in_features, out_features, lmax, bias=True): + ''' + 1. Use `torch.einsum` to prevent slicing and concatenation + 2. Need to specify some behaviors in `no_weight_decay` and weight initialization. + ''' + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.lmax = lmax + + self.weight = torch.nn.Parameter(torch.randn((self.lmax + 1), out_features, in_features)) + bound = 1 / math.sqrt(self.in_features) + torch.nn.init.uniform_(self.weight, -bound, bound) + self.bias = torch.nn.Parameter(torch.zeros(out_features)) + + expand_index = torch.zeros([(lmax + 1) ** 2]).long() + for l in range(lmax + 1): + start_idx = l ** 2 + length = 2 * l + 1 + expand_index[start_idx : (start_idx + length)] = l + self.register_buffer('expand_index', expand_index) + + + def forward(self, input_embedding): + + weight = torch.index_select(self.weight, dim=0, index=self.expand_index) # [(L_max + 1) ** 2, C_out, C_in] + out = torch.einsum('bmi, moi -> bmo', input_embedding.embedding, weight) # [N, (L_max + 1) ** 2, C_out] + bias = self.bias.view(1, 1, self.out_features) + out[:, 0:1, :] = out.narrow(1, 0, 1) + bias + + out_embedding = SO3_Embedding( + 0, + input_embedding.lmax_list.copy(), + self.out_features, + device=input_embedding.device, + dtype=input_embedding.dtype + ) + out_embedding.set_embedding(out) + out_embedding.set_lmax_mmax(input_embedding.lmax_list.copy(), input_embedding.lmax_list.copy()) + + return out_embedding + + + def __repr__(self): + return f"{self.__class__.__name__}(in_features={self.in_features}, out_features={self.out_features}, lmax={self.lmax})" \ No newline at end of file diff --git a/magnet/eqV2/transformer_block.py b/magnet/eqV2/transformer_block.py new file mode 100644 index 0000000000000000000000000000000000000000..eee38ad422688e05b34578023f802a0b7f1d0b1b --- /dev/null +++ b/magnet/eqV2/transformer_block.py @@ -0,0 +1,635 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +import torch_geometric +import copy + +from .activation import ( + ScaledSiLU, + ScaledSwiGLU, + SwiGLU, + ScaledSmoothLeakyReLU, + SmoothLeakyReLU, + GateActivation, + SeparableS2Activation, + S2Activation +) +from .layer_norm import ( + EquivariantLayerNormArray, + EquivariantLayerNormArraySphericalHarmonics, + EquivariantRMSNormArraySphericalHarmonics, + get_normalization_layer +) +from .so2_ops import ( + SO2_Convolution, + SO2_Linear +) +from .so3 import ( + SO3_Embedding, + SO3_Linear, + SO3_LinearV2 +) +from .radial_function import RadialFunction +from .drop import ( + GraphDropPath, + EquivariantDropoutArraySphericalHarmonics +) + + +class SO2EquivariantGraphAttention(torch.nn.Module): + """ + SO2EquivariantGraphAttention: Perform MLP attention + non-linear message passing + SO(2) Convolution with radial function -> S2 Activation -> SO(2) Convolution -> attention weights and non-linear messages + attention weights * non-linear messages -> Linear + + Args: + sphere_channels (int): Number of spherical channels + hidden_channels (int): Number of hidden channels used during the SO(2) conv + num_heads (int): Number of attention heads + attn_alpha_head (int): Number of channels for alpha vector in each attention head + attn_value_head (int): Number of channels for value vector in each attention head + output_channels (int): Number of output channels + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + + SO3_rotation (list:SO3_Rotation): Class to calculate Wigner-D matrices and rotate embeddings + mappingReduced (CoefficientMappingModule): Class to convert l and m indices once node embedding is rotated + SO3_grid (SO3_grid): Class used to convert from grid the spherical harmonic representations + + max_num_elements (int): Maximum number of atomic numbers + edge_channels_list (list:int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels]. + The last one will be used as hidden size when `use_atom_edge_embedding` is `True`. + use_atom_edge_embedding (bool): Whether to use atomic embedding along with relative distance for edge scalar features + use_m_share_rad (bool): Whether all m components within a type-L vector of one channel share radial function weights + + activation (str): Type of activation function + use_s2_act_attn (bool): Whether to use attention after S2 activation. Otherwise, use the same attention as Equiformer + use_attn_renorm (bool): Whether to re-normalize attention weights + use_gate_act (bool): If `True`, use gate activation. Otherwise, use S2 activation. + use_sep_s2_act (bool): If `True`, use separable S2 activation when `use_gate_act` is False. + + alpha_drop (float): Dropout rate for attention weights + """ + + def __init__( + self, + sphere_channels, + hidden_channels, + num_heads, + attn_alpha_channels, + attn_value_channels, + output_channels, + lmax_list, + mmax_list, + SO3_rotation, + mappingReduced, + SO3_grid, + max_num_elements, + edge_channels_list, + use_atom_edge_embedding=True, + use_m_share_rad=False, + activation='scaled_silu', + use_s2_act_attn=False, + use_attn_renorm=True, + use_gate_act=False, + use_sep_s2_act=True, + alpha_drop=0.0, + ): + super(SO2EquivariantGraphAttention, self).__init__() + + self.sphere_channels = sphere_channels + self.hidden_channels = hidden_channels + self.num_heads = num_heads + self.attn_alpha_channels = attn_alpha_channels + self.attn_value_channels = attn_value_channels + self.output_channels = output_channels + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.num_resolutions = len(self.lmax_list) + + self.SO3_rotation = SO3_rotation + self.mappingReduced = mappingReduced + self.SO3_grid = SO3_grid + + # Create edge scalar (invariant to rotations) features + # Embedding function of the atomic numbers + self.max_num_elements = max_num_elements + self.edge_channels_list = copy.deepcopy(edge_channels_list) + self.use_atom_edge_embedding = use_atom_edge_embedding + self.use_m_share_rad = use_m_share_rad + + if self.use_atom_edge_embedding: + self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels_list[-1]) + self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels_list[-1]) + nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001) + nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001) + self.edge_channels_list[0] = self.edge_channels_list[0] + 2 * self.edge_channels_list[-1] + else: + self.source_embedding, self.target_embedding = None, None + + # if we want to add some learned featurization of solvent vs solute atoms (which use different basis sets), do that here... + + + self.use_s2_act_attn = use_s2_act_attn + self.use_attn_renorm = use_attn_renorm + self.use_gate_act = use_gate_act + self.use_sep_s2_act = use_sep_s2_act + + assert not self.use_s2_act_attn # since this is not used + + # Create SO(2) convolution blocks + extra_m0_output_channels = None + if not self.use_s2_act_attn: + extra_m0_output_channels = self.num_heads * self.attn_alpha_channels + if self.use_gate_act: + extra_m0_output_channels = extra_m0_output_channels + max(self.lmax_list) * self.hidden_channels + else: + if self.use_sep_s2_act: + extra_m0_output_channels = extra_m0_output_channels + self.hidden_channels + + if self.use_m_share_rad: + self.edge_channels_list = self.edge_channels_list + [2 * self.sphere_channels * (max(self.lmax_list) + 1)] + self.rad_func = RadialFunction(self.edge_channels_list) + expand_index = torch.zeros([(max(self.lmax_list) + 1) ** 2]).long() + for l in range(max(self.lmax_list) + 1): + start_idx = l ** 2 + length = 2 * l + 1 + expand_index[start_idx : (start_idx + length)] = l + self.register_buffer('expand_index', expand_index) + + self.so2_conv_1 = SO2_Convolution( + 2 * self.sphere_channels, + self.hidden_channels, + self.lmax_list, + self.mmax_list, + self.mappingReduced, + internal_weights=( + False if not self.use_m_share_rad + else True + ), + edge_channels_list=( + self.edge_channels_list if not self.use_m_share_rad + else None + ), + extra_m0_output_channels=extra_m0_output_channels # for attention weights and/or gate activation + ) + + if self.use_s2_act_attn: + self.alpha_norm = None + self.alpha_act = None + self.alpha_dot = None + else: + if self.use_attn_renorm: + self.alpha_norm = torch.nn.LayerNorm(self.attn_alpha_channels) + else: + self.alpha_norm = torch.nn.Identity() + self.alpha_act = SmoothLeakyReLU() + self.alpha_dot = torch.nn.Parameter(torch.randn(self.num_heads, self.attn_alpha_channels)) + #torch_geometric.nn.inits.glorot(self.alpha_dot) # Following GATv2 + std = 1.0 / math.sqrt(self.attn_alpha_channels) + torch.nn.init.uniform_(self.alpha_dot, -std, std) + + self.alpha_dropout = None + if alpha_drop != 0.0: + self.alpha_dropout = torch.nn.Dropout(alpha_drop) + + if self.use_gate_act: + self.gate_act = GateActivation( + lmax=max(self.lmax_list), + mmax=max(self.mmax_list), + num_channels=self.hidden_channels + ) + else: + if self.use_sep_s2_act: + # separable S2 activation + self.s2_act = SeparableS2Activation( + lmax=max(self.lmax_list), + mmax=max(self.mmax_list) + ) + else: + # S2 activation + self.s2_act = S2Activation( + lmax=max(self.lmax_list), + mmax=max(self.mmax_list) + ) + + self.so2_conv_2 = SO2_Convolution( + self.hidden_channels, + self.num_heads * self.attn_value_channels, + self.lmax_list, + self.mmax_list, + self.mappingReduced, + internal_weights=True, + edge_channels_list=None, + extra_m0_output_channels=( + self.num_heads if self.use_s2_act_attn + else None + ) # for attention weights + ) + + self.proj = SO3_LinearV2(self.num_heads * self.attn_value_channels, self.output_channels, lmax=self.lmax_list[0]) + + + def forward( + self, + x, + atomic_numbers, + edge_distance, + edge_index + ): + + # Compute edge scalar features (invariant to rotations) + # Uses atomic numbers and edge distance as inputs + if self.use_atom_edge_embedding: + source_element = atomic_numbers[edge_index[0]] # Source atom atomic number + target_element = atomic_numbers[edge_index[1]] # Target atom atomic number + source_embedding = self.source_embedding(source_element) + target_embedding = self.target_embedding(target_element) + x_edge = torch.cat((edge_distance, source_embedding, target_embedding), dim=1) + else: + x_edge = edge_distance + + x_source = x.clone() + x_target = x.clone() + x_source._expand_edge(edge_index[0, :]) + x_target._expand_edge(edge_index[1, :]) + + x_message_data = torch.cat((x_source.embedding, x_target.embedding), dim=2) + x_message = SO3_Embedding( + 0, + x_target.lmax_list.copy(), + x_target.num_channels * 2, + device=x_target.device, + dtype=x_target.dtype + ) + x_message.set_embedding(x_message_data) + x_message.set_lmax_mmax(self.lmax_list.copy(), self.mmax_list.copy()) + + # radial function (scale all m components within a type-L vector of one channel with the same weight) + if self.use_m_share_rad: + x_edge_weight = self.rad_func(x_edge) + x_edge_weight = x_edge_weight.reshape(-1, (max(self.lmax_list) + 1), 2 * self.sphere_channels) + x_edge_weight = torch.index_select(x_edge_weight, dim=1, index=self.expand_index) # [E, (L_max + 1) ** 2, C] + x_message.embedding = x_message.embedding * x_edge_weight + + # Rotate the irreps to align with the edge + x_message._rotate(self.SO3_rotation, self.lmax_list, self.mmax_list) + + # First SO(2)-convolution + if self.use_s2_act_attn: + x_message = self.so2_conv_1(x_message, x_edge) + else: + x_message, x_0_extra = self.so2_conv_1(x_message, x_edge) + + # Activation + x_alpha_num_channels = self.num_heads * self.attn_alpha_channels + if self.use_gate_act: + # Gate activation + x_0_gating = x_0_extra.narrow(1, x_alpha_num_channels, x_0_extra.shape[1] - x_alpha_num_channels) # for activation + x_0_alpha = x_0_extra.narrow(1, 0, x_alpha_num_channels) # for attention weights + x_message.embedding = self.gate_act(x_0_gating, x_message.embedding) + else: + if self.use_sep_s2_act: + x_0_gating = x_0_extra.narrow(1, x_alpha_num_channels, x_0_extra.shape[1] - x_alpha_num_channels) # for activation + x_0_alpha = x_0_extra.narrow(1, 0, x_alpha_num_channels) # for attention weights + x_message.embedding = self.s2_act(x_0_gating, x_message.embedding, self.SO3_grid) + else: + x_0_alpha = x_0_extra + x_message.embedding = self.s2_act(x_message.embedding, self.SO3_grid) + ##x_message._grid_act(self.SO3_grid, self.value_act, self.mappingReduced) + + # Second SO(2)-convolution + if self.use_s2_act_attn: + x_message, x_0_extra = self.so2_conv_2(x_message, x_edge) + else: + x_message = self.so2_conv_2(x_message, x_edge) + + # Attention weights + if self.use_s2_act_attn: + alpha = x_0_extra + else: + x_0_alpha = x_0_alpha.reshape(-1, self.num_heads, self.attn_alpha_channels) + x_0_alpha = self.alpha_norm(x_0_alpha) + x_0_alpha = self.alpha_act(x_0_alpha) + alpha = torch.einsum('bik, ik -> bi', x_0_alpha, self.alpha_dot) + alpha = torch_geometric.utils.softmax(alpha, edge_index[1]) + alpha = alpha.reshape(alpha.shape[0], 1, self.num_heads, 1) + if self.alpha_dropout is not None: + alpha = self.alpha_dropout(alpha) + + # Attention weights * non-linear messages + attn = x_message.embedding + attn = attn.reshape(attn.shape[0], attn.shape[1], self.num_heads, self.attn_value_channels) + attn = attn * alpha + attn = attn.reshape(attn.shape[0], attn.shape[1], self.num_heads * self.attn_value_channels) + x_message.embedding = attn + + # Rotate back the irreps + x_message._rotate_inv(self.SO3_rotation, self.mappingReduced) + + # Compute the sum of the incoming neighboring messages for each target node + x_message._reduce_edge(edge_index[1], len(x.embedding)) + + # Project + out_embedding = self.proj(x_message) + + return out_embedding + + +class FeedForwardNetwork(torch.nn.Module): + """ + FeedForwardNetwork: Perform feedforward network with S2 activation or gate activation + + Args: + sphere_channels (int): Number of spherical channels + hidden_channels (int): Number of hidden channels used during feedforward network + output_channels (int): Number of output channels + + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + + SO3_grid (SO3_grid): Class used to convert from grid the spherical harmonic representations + + activation (str): Type of activation function + use_gate_act (bool): If `True`, use gate activation. Otherwise, use S2 activation + use_grid_mlp (bool): If `True`, use projecting to grids and performing MLPs. + use_sep_s2_act (bool): If `True`, use separable grid MLP when `use_grid_mlp` is True. + """ + + def __init__( + self, + sphere_channels, + hidden_channels, + output_channels, + lmax_list, + mmax_list, + SO3_grid, + activation='scaled_silu', + use_gate_act=False, + use_grid_mlp=False, + use_sep_s2_act=True + ): + super(FeedForwardNetwork, self).__init__() + self.sphere_channels = sphere_channels + self.hidden_channels = hidden_channels + self.output_channels = output_channels + self.lmax_list = lmax_list + self.mmax_list = mmax_list + self.num_resolutions = len(lmax_list) + self.sphere_channels_all = self.num_resolutions * self.sphere_channels + self.SO3_grid = SO3_grid + self.use_gate_act = use_gate_act + self.use_grid_mlp = use_grid_mlp + self.use_sep_s2_act = use_sep_s2_act + + self.max_lmax = max(self.lmax_list) + + self.so3_linear_1 = SO3_LinearV2(self.sphere_channels_all, self.hidden_channels, lmax=self.max_lmax) + if self.use_grid_mlp: + if self.use_sep_s2_act: + self.scalar_mlp = nn.Sequential( + nn.Linear(self.sphere_channels_all, self.hidden_channels, bias=True), + nn.SiLU(), + ) + else: + self.scalar_mlp = None + self.grid_mlp = nn.Sequential( + nn.Linear(self.hidden_channels, self.hidden_channels, bias=False), + nn.SiLU(), + nn.Linear(self.hidden_channels, self.hidden_channels, bias=False), + nn.SiLU(), + nn.Linear(self.hidden_channels, self.hidden_channels, bias=False) + ) + else: + if self.use_gate_act: + self.gating_linear = torch.nn.Linear(self.sphere_channels_all, self.max_lmax * self.hidden_channels) + self.gate_act = GateActivation(self.max_lmax, self.max_lmax, self.hidden_channels) + else: + if self.use_sep_s2_act: + self.gating_linear = torch.nn.Linear(self.sphere_channels_all, self.hidden_channels) + self.s2_act = SeparableS2Activation(self.max_lmax, self.max_lmax) + else: + self.gating_linear = None + self.s2_act = S2Activation(self.max_lmax, self.max_lmax) + self.so3_linear_2 = SO3_LinearV2(self.hidden_channels, self.output_channels, lmax=self.max_lmax) + + + def forward(self, input_embedding): + + gating_scalars = None + if self.use_grid_mlp: + if self.use_sep_s2_act: + gating_scalars = self.scalar_mlp(input_embedding.embedding.narrow(1, 0, 1)) + else: + if self.gating_linear is not None: + gating_scalars = self.gating_linear(input_embedding.embedding.narrow(1, 0, 1)) + + input_embedding = self.so3_linear_1(input_embedding) + + if self.use_grid_mlp: + # Project to grid + input_embedding_grid = input_embedding.to_grid(self.SO3_grid, lmax=self.max_lmax) + # Perform point-wise operations + input_embedding_grid = self.grid_mlp(input_embedding_grid) + # Project back to spherical harmonic coefficients + input_embedding._from_grid(input_embedding_grid, self.SO3_grid, lmax=self.max_lmax) + + if self.use_sep_s2_act: + input_embedding.embedding = torch.cat( + (gating_scalars, input_embedding.embedding.narrow(1, 1, input_embedding.embedding.shape[1] - 1)), + dim=1 + ) + else: + if self.use_gate_act: + input_embedding.embedding = self.gate_act(gating_scalars, input_embedding.embedding) + else: + if self.use_sep_s2_act: + input_embedding.embedding = self.s2_act(gating_scalars, input_embedding.embedding, self.SO3_grid) + else: + input_embedding.embedding = self.s2_act(input_embedding.embedding, self.SO3_grid) + + input_embedding = self.so3_linear_2(input_embedding) + + return input_embedding + + +class TransBlockV2(torch.nn.Module): + """ + + Args: + sphere_channels (int): Number of spherical channels + attn_hidden_channels (int): Number of hidden channels used during SO(2) graph attention + num_heads (int): Number of attention heads + attn_alpha_head (int): Number of channels for alpha vector in each attention head + attn_value_head (int): Number of channels for value vector in each attention head + ffn_hidden_channels (int): Number of hidden channels used during feedforward network + output_channels (int): Number of output channels + + lmax_list (list:int): List of degrees (l) for each resolution + mmax_list (list:int): List of orders (m) for each resolution + + SO3_rotation (list:SO3_Rotation): Class to calculate Wigner-D matrices and rotate embeddings + mappingReduced (CoefficientMappingModule): Class to convert l and m indices once node embedding is rotated + SO3_grid (SO3_grid): Class used to convert from grid the spherical harmonic representations + + max_num_elements (int): Maximum number of atomic numbers + edge_channels_list (list:int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels]. + The last one will be used as hidden size when `use_atom_edge_embedding` is `True`. + use_atom_edge_embedding (bool): Whether to use atomic embedding along with relative distance for edge scalar features + use_m_share_rad (bool): Whether all m components within a type-L vector of one channel share radial function weights + + attn_activation (str): Type of activation function for SO(2) graph attention + use_s2_act_attn (bool): Whether to use attention after S2 activation. Otherwise, use the same attention as Equiformer + use_attn_renorm (bool): Whether to re-normalize attention weights + ffn_activation (str): Type of activation function for feedforward network + use_gate_act (bool): If `True`, use gate activation. Otherwise, use S2 activation + use_grid_mlp (bool): If `True`, use projecting to grids and performing MLPs for FFN. + use_sep_s2_act (bool): If `True`, use separable S2 activation when `use_gate_act` is False. + + norm_type (str): Type of normalization layer (['layer_norm', 'layer_norm_sh']) + + alpha_drop (float): Dropout rate for attention weights + drop_path_rate (float): Drop path rate + proj_drop (float): Dropout rate for outputs of attention and FFN + """ + + def __init__( + self, + sphere_channels, + attn_hidden_channels, + num_heads, + attn_alpha_channels, + attn_value_channels, + ffn_hidden_channels, + output_channels, + + lmax_list, + mmax_list, + + SO3_rotation, + mappingReduced, + SO3_grid, + + max_num_elements, + edge_channels_list, + use_atom_edge_embedding=True, + use_m_share_rad=False, + + attn_activation='silu', + use_s2_act_attn=False, + use_attn_renorm=True, + ffn_activation='silu', + use_gate_act=False, + use_grid_mlp=False, + use_sep_s2_act=True, + + norm_type='rms_norm_sh', + + alpha_drop=0.0, + drop_path_rate=0.0, + proj_drop=0.0 + ): + super(TransBlockV2, self).__init__() + + max_lmax = max(lmax_list) + self.norm_1 = get_normalization_layer(norm_type, lmax=max_lmax, num_channels=sphere_channels) + + self.ga = SO2EquivariantGraphAttention( + sphere_channels=sphere_channels, + hidden_channels=attn_hidden_channels, + num_heads=num_heads, + attn_alpha_channels=attn_alpha_channels, + attn_value_channels=attn_value_channels, + output_channels=sphere_channels, + lmax_list=lmax_list, + mmax_list=mmax_list, + SO3_rotation=SO3_rotation, + mappingReduced=mappingReduced, + SO3_grid=SO3_grid, + max_num_elements=max_num_elements, + edge_channels_list=edge_channels_list, + use_atom_edge_embedding=use_atom_edge_embedding, + use_m_share_rad=use_m_share_rad, + activation=attn_activation, + use_s2_act_attn=use_s2_act_attn, + use_attn_renorm=use_attn_renorm, + use_gate_act=use_gate_act, + use_sep_s2_act=use_sep_s2_act, + alpha_drop=alpha_drop, + ) + + self.drop_path = GraphDropPath(drop_path_rate) if drop_path_rate > 0. else None + self.proj_drop = EquivariantDropoutArraySphericalHarmonics(proj_drop, drop_graph=False) if proj_drop > 0.0 else None + + self.norm_2 = get_normalization_layer(norm_type, lmax=max_lmax, num_channels=sphere_channels) + + self.ffn = FeedForwardNetwork( + sphere_channels=sphere_channels, + hidden_channels=ffn_hidden_channels, + output_channels=output_channels, + lmax_list=lmax_list, + mmax_list=mmax_list, + SO3_grid=SO3_grid, + activation=ffn_activation, + use_gate_act=use_gate_act, + use_grid_mlp=use_grid_mlp, + use_sep_s2_act=use_sep_s2_act + ) + + if sphere_channels != output_channels: + self.ffn_shortcut = SO3_LinearV2(sphere_channels, output_channels, lmax=max_lmax) + else: + self.ffn_shortcut = None + + + def forward( + self, + x, # SO3_Embedding + atomic_numbers, + edge_distance, + edge_index, + batch # for GraphDropPath + ): + + output_embedding = x + + x_res = output_embedding.embedding + output_embedding.embedding = self.norm_1(output_embedding.embedding) + output_embedding = self.ga(output_embedding, + atomic_numbers, + edge_distance, + edge_index) + + if self.drop_path is not None: + output_embedding.embedding = self.drop_path(output_embedding.embedding, batch) + if self.proj_drop is not None: + output_embedding.embedding = self.proj_drop(output_embedding.embedding, batch) + + output_embedding.embedding = output_embedding.embedding + x_res + + x_res = output_embedding.embedding + output_embedding.embedding = self.norm_2(output_embedding.embedding) + output_embedding = self.ffn(output_embedding) + + if self.drop_path is not None: + output_embedding.embedding = self.drop_path(output_embedding.embedding, batch) + if self.proj_drop is not None: + output_embedding.embedding = self.proj_drop(output_embedding.embedding, batch) + + if self.ffn_shortcut is not None: + shortcut_embedding = SO3_Embedding( + 0, + output_embedding.lmax_list.copy(), + self.ffn_shortcut.in_features, + device=output_embedding.device, + dtype=output_embedding.dtype + ) + shortcut_embedding.set_embedding(x_res) + shortcut_embedding.set_lmax_mmax(output_embedding.lmax_list.copy(), output_embedding.lmax_list.copy()) + shortcut_embedding = self.ffn_shortcut(shortcut_embedding) + x_res = shortcut_embedding.embedding + + output_embedding.embedding = output_embedding.embedding + x_res + + return output_embedding \ No newline at end of file diff --git a/magnet/eqV2/wigner.py b/magnet/eqV2/wigner.py new file mode 100644 index 0000000000000000000000000000000000000000..9ff2ec92de44ed803185e5370a5b145ee419d711 --- /dev/null +++ b/magnet/eqV2/wigner.py @@ -0,0 +1,38 @@ +import os +import torch + + +# Borrowed from e3nn @ 0.4.0: +# https://github.com/e3nn/e3nn/blob/0.4.0/e3nn/o3/_wigner.py#L10 +# _Jd is a list of tensors of shape (2l+1, 2l+1) +_Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt")) + + +# Borrowed from e3nn @ 0.4.0: +# https://github.com/e3nn/e3nn/blob/0.4.0/e3nn/o3/_wigner.py#L37 +# +# In 0.5.0, e3nn shifted to torch.matrix_exp which is significantly slower: +# https://github.com/e3nn/e3nn/blob/0.5.0/e3nn/o3/_wigner.py#L92 +def wigner_D(l, alpha, beta, gamma): + if not l < len(_Jd): + raise NotImplementedError( + f"wigner D maximum l implemented is {len(_Jd) - 1}, send us an email to ask for more" + ) + + alpha, beta, gamma = torch.broadcast_tensors(alpha, beta, gamma) + J = _Jd[l].to(dtype=alpha.dtype, device=alpha.device) + Xa = _z_rot_mat(alpha, l) + Xb = _z_rot_mat(beta, l) + Xc = _z_rot_mat(gamma, l) + return Xa @ J @ Xb @ J @ Xc + + +def _z_rot_mat(angle, l): + shape, device, dtype = angle.shape, angle.device, angle.dtype + M = angle.new_zeros((*shape, 2 * l + 1, 2 * l + 1)) + inds = torch.arange(0, 2 * l + 1, 1, device=device) + reversed_inds = torch.arange(2 * l, -1, -1, device=device) + frequencies = torch.arange(l, -l - 1, -1, dtype=dtype, device=device) + M[..., inds, reversed_inds] = torch.sin(frequencies * angle[..., None]) + M[..., inds, inds] = torch.cos(frequencies * angle[..., None]) + return M \ No newline at end of file