interaction-eda-step 2026.11.7__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. interaction_eda_step/__init__.py +26 -0
  2. interaction_eda_step/__main__.py +20 -0
  3. interaction_eda_step/_version.py +21 -0
  4. interaction_eda_step/backends/__init__.py +9 -0
  5. interaction_eda_step/backends/base.py +124 -0
  6. interaction_eda_step/backends/gpu4pyscf.py +183 -0
  7. interaction_eda_step/data/configuration.txt +3 -0
  8. interaction_eda_step/data/gpu4pyscf.ini +33 -0
  9. interaction_eda_step/data/gpu4pyscf_eda.py +256 -0
  10. interaction_eda_step/data/properties.csv +17 -0
  11. interaction_eda_step/data/references.bib +40 -0
  12. interaction_eda_step/data/seamm-gpu4pyscf.yml +12 -0
  13. interaction_eda_step/fragments.py +189 -0
  14. interaction_eda_step/installer.py +423 -0
  15. interaction_eda_step/interaction_eda.py +634 -0
  16. interaction_eda_step/interaction_eda_parameters.py +291 -0
  17. interaction_eda_step/interaction_eda_step.py +95 -0
  18. interaction_eda_step/manybody.py +172 -0
  19. interaction_eda_step/metadata.py +67 -0
  20. interaction_eda_step/substep.py +99 -0
  21. interaction_eda_step/terms.py +172 -0
  22. interaction_eda_step/tk_interaction_eda.py +230 -0
  23. interaction_eda_step-2026.11.7.dist-info/METADATA +136 -0
  24. interaction_eda_step-2026.11.7.dist-info/RECORD +30 -0
  25. interaction_eda_step-2026.11.7.dist-info/WHEEL +6 -0
  26. interaction_eda_step-2026.11.7.dist-info/entry_points.txt +8 -0
  27. interaction_eda_step-2026.11.7.dist-info/licenses/AUTHORS.rst +5 -0
  28. interaction_eda_step-2026.11.7.dist-info/licenses/LICENSE +31 -0
  29. interaction_eda_step-2026.11.7.dist-info/top_level.txt +1 -0
  30. interaction_eda_step-2026.11.7.dist-info/zip-safe +1 -0
@@ -0,0 +1,26 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ """
4
+ interaction_eda_step
5
+ A SEAMM plug-in for Interaction EDA
6
+ """
7
+
8
+ # Bring up the classes so that they appear to be directly in
9
+ # the interaction_eda_step package.
10
+
11
+ from .interaction_eda import InteractionEda # noqa: F401, E501
12
+ from .interaction_eda_parameters import InteractionEdaParameters # noqa: F401, E501
13
+ from .interaction_eda_step import InteractionEdaStep # noqa: F401, E501
14
+ from .tk_interaction_eda import TkInteractionEda # noqa: F401, E501
15
+
16
+ from .metadata import metadata # noqa: F401
17
+
18
+ # Handle versioneer
19
+ from ._version import get_versions
20
+
21
+ __author__ = "Paul Saxe"
22
+ __email__ = "psaxe@molssi.org"
23
+ versions = get_versions()
24
+ __version__ = versions["version"]
25
+ __git_revision__ = versions["full-revisionid"]
26
+ del get_versions, versions
@@ -0,0 +1,20 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ """Handle the installation of the Interaction EDA step."""
4
+
5
+ from .installer import Installer
6
+
7
+
8
+ def run():
9
+ """Handle the extra installation needed.
10
+
11
+ * Create the conda environment `seamm-gpu4pyscf` with GPU4PySCF for this
12
+ machine's NVIDIA driver.
13
+ * Add or update information in the SEAMM gpu4pyscf.ini file.
14
+ """
15
+ installer = Installer()
16
+ installer.run()
17
+
18
+
19
+ if __name__ == "__main__":
20
+ run()
@@ -0,0 +1,21 @@
1
+
2
+ # This file was generated by 'versioneer.py' (0.18) from
3
+ # revision-control system data, or from the parent directory name of an
4
+ # unpacked source archive. Distribution tarballs contain a pre-generated copy
5
+ # of this file.
6
+
7
+ import json
8
+
9
+ version_json = '''
10
+ {
11
+ "date": "2026-11-07T13:51:33-0500",
12
+ "dirty": false,
13
+ "error": null,
14
+ "full-revisionid": "b97b2b9549560991799ed04f040db639cacc76d0",
15
+ "version": "2026.11.7"
16
+ }
17
+ ''' # END VERSION_JSON
18
+
19
+
20
+ def get_versions():
21
+ return json.loads(version_json)
@@ -0,0 +1,9 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ """The codes that do the energy decomposition (see :mod:`.base`)."""
4
+
5
+ from .base import Backend, Calculation, Result, read_jsonl # noqa: F401
6
+ from .gpu4pyscf import GPU4PySCF # noqa: F401
7
+
8
+ #: The backends by name.
9
+ BACKENDS = {b.name: b for b in (GPU4PySCF,)}
@@ -0,0 +1,124 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ """The interface between the Interaction EDA step and the codes that do the EDA.
4
+
5
+ A backend runs a batch of EDA calculations, one per set of fragments, in its own
6
+ environment through ``seamm_exec``. The step knows only this interface:
7
+
8
+ * what the backend can do (:attr:`Backend.scheme`, :attr:`Backend.device`,
9
+ :attr:`Backend.max_fragments`, :meth:`Backend.options`), so the GUI can offer
10
+ only valid choices;
11
+ * :meth:`Backend.files` -- the input files for a batch;
12
+ * :attr:`Backend.ini` and :meth:`Backend.command` -- how ``seamm_exec`` runs it;
13
+ * :meth:`Backend.read_results` -- the native terms of each calculation, kJ/mol.
14
+
15
+ Backends planned: ``gpu4pyscf`` (ALMO-EDA2 on a GPU), ``psi4_almo`` (ALMO-EDA2 on
16
+ CPUs) and ``psi4_sapt`` (SAPT0/F-SAPT, dimers only).
17
+ """
18
+
19
+ from dataclasses import dataclass, field
20
+ import json
21
+ from pathlib import Path
22
+
23
+
24
+ @dataclass
25
+ class Calculation:
26
+ """One EDA calculation: a set of fragments.
27
+
28
+ Attributes
29
+ ----------
30
+ key : tuple(int, ...)
31
+ Ascending fragment indices.
32
+ fragments : [interaction_eda_step.fragments.Fragment]
33
+ """
34
+
35
+ key: tuple
36
+ fragments: list
37
+
38
+ @property
39
+ def id(self):
40
+ return "+".join(str(k) for k in self.key)
41
+
42
+
43
+ @dataclass
44
+ class Result:
45
+ """The outcome of one calculation.
46
+
47
+ Attributes
48
+ ----------
49
+ key : tuple(int, ...)
50
+ energies : {str: float} or None
51
+ Native terms, kJ/mol; None if it failed.
52
+ error : str or None
53
+ data : dict
54
+ Everything else the backend reports (timing, memory, versions, ...).
55
+ """
56
+
57
+ key: tuple
58
+ energies: dict = None
59
+ error: str = None
60
+ data: dict = field(default_factory=dict)
61
+
62
+
63
+ class Backend:
64
+ """Base class for the EDA backends."""
65
+
66
+ #: Short name, used in the parameters and the ini file name.
67
+ name = None
68
+ #: The EDA scheme (a key of :data:`interaction_eda_step.terms.SCHEMES`).
69
+ scheme = None
70
+ #: "gpu" or "cpu".
71
+ device = "cpu"
72
+ #: The most fragments in one calculation; None for no limit.
73
+ max_fragments = None
74
+ #: The ini file (in the installation's root) that says how to run the code.
75
+ ini = None
76
+ #: Results file the code writes.
77
+ results_file = "results.jsonl"
78
+
79
+ def options(self):
80
+ """The settings this backend accepts: {name: default}."""
81
+ return {}
82
+
83
+ def files(self, calculations, settings):
84
+ """The input files for a batch of calculations: {name: text}."""
85
+ raise NotImplementedError
86
+
87
+ def command(self):
88
+ """The command line, with ``{code}`` from the ini file."""
89
+ raise NotImplementedError
90
+
91
+ def read_results(self, directory, calculations):
92
+ """{key: Result} for the calculations, from the results file.
93
+
94
+ A calculation with no result is reported with an error.
95
+ """
96
+ raise NotImplementedError
97
+
98
+ def method_string(self, settings):
99
+ """A one-line description of the method, for labels and tables."""
100
+ raise NotImplementedError
101
+
102
+ def citations(self, settings):
103
+ """[(alias, bibtex, level, note)] to cite."""
104
+ return []
105
+
106
+
107
+ def read_jsonl(path):
108
+ """The records of a JSON-lines file; later records replace earlier ones
109
+ with the same id (a rerun after a failure)."""
110
+ records = {}
111
+ path = Path(path)
112
+ if path.exists():
113
+ for line in path.read_text().splitlines():
114
+ line = line.strip()
115
+ if not line:
116
+ continue
117
+ try:
118
+ record = json.loads(line)
119
+ except ValueError:
120
+ continue
121
+ old = records.get(record.get("id"))
122
+ if old is None or "error" in old or "error" not in record:
123
+ records[record.get("id")] = record
124
+ return records
@@ -0,0 +1,183 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ """ALMO-EDA2 with GPU4PySCF on a GPU."""
4
+
5
+ import importlib.resources
6
+ import json
7
+ from pathlib import Path
8
+
9
+ from .base import Backend, Result, read_jsonl
10
+
11
+ DRIVER = "gpu4pyscf_eda.py"
12
+
13
+ # The settings validated against SAPT and CCSD(T) and on ChemAI and ARC
14
+ # (2026-11-06): wB97M-V/def2-TZVPD, no density fitting, FERF SVD threshold 1e-4,
15
+ # coupled FERF-DQ response.
16
+ DEFAULTS = {
17
+ "xc": "wb97m-v",
18
+ "basis": "def2-tzvpd",
19
+ "auxbasis": None,
20
+ "nlc": None,
21
+ "disp": None,
22
+ "svd_threshold": 1.0e-4,
23
+ "ferf": "coupled",
24
+ "response_nlc": True,
25
+ "allow_low_fp64": False,
26
+ "verbose": 3,
27
+ }
28
+
29
+ BIBTEX = {
30
+ "gpu4pyscf": (
31
+ "@article{gpu4pyscf,\n"
32
+ " title = {Introducing {GPU} Acceleration into the {Python}-Based"
33
+ " Simulations of Chemistry Framework},\n"
34
+ " author = {Li, Rui and Sun, Qiming and Zhang, Xing and Chan, Garnet"
35
+ " Kin-Lic},\n"
36
+ " journal = {J. Phys. Chem. A},\n"
37
+ " volume = {129},\n number = {5},\n pages = {1459--1468},\n"
38
+ " year = {2025},\n"
39
+ " doi = {10.1021/acs.jpca.4c05876}\n}"
40
+ ),
41
+ "pyscf": (
42
+ "@article{pyscf2020,\n"
43
+ " title = {Recent developments in the {PySCF} program package},\n"
44
+ " author = {Sun, Qiming and others},\n"
45
+ " journal = {J. Chem. Phys.},\n"
46
+ " volume = {153},\n pages = {024109},\n year = {2020},\n"
47
+ " doi = {10.1063/5.0006074}\n}"
48
+ ),
49
+ "almo_eda2": (
50
+ "@article{almo_eda2,\n"
51
+ " title = {Probing non-covalent interactions with a second generation"
52
+ " energy decomposition analysis using absolutely localized molecular"
53
+ " orbitals},\n"
54
+ " author = {Horn, Paul R. and Mao, Yuezhi and Head-Gordon, Martin},\n"
55
+ " journal = {Phys. Chem. Chem. Phys.},\n"
56
+ " volume = {18},\n pages = {23067--23079},\n year = {2016},\n"
57
+ " doi = {10.1039/C6CP03784D}\n}"
58
+ ),
59
+ "frozen": (
60
+ "@article{almo_eda_frozen,\n"
61
+ " title = {Defining the contributions of permanent electrostatics,"
62
+ " {Pauli} repulsion, and dispersion in density functional theory"
63
+ " calculations of intermolecular interaction energies},\n"
64
+ " author = {Horn, Paul R. and Mao, Yuezhi and Head-Gordon, Martin},\n"
65
+ " journal = {J. Chem. Phys.},\n"
66
+ " volume = {144},\n pages = {114107},\n year = {2016},\n"
67
+ " doi = {10.1063/1.4942921}\n}"
68
+ ),
69
+ "ferf": (
70
+ "@article{ferf,\n"
71
+ " title = {Polarization contributions to intermolecular interactions"
72
+ " revisited with fragment electric-field response functions},\n"
73
+ " author = {Horn, Paul R. and Head-Gordon, Martin},\n"
74
+ " journal = {J. Chem. Phys.},\n"
75
+ " volume = {143},\n pages = {114111},\n year = {2015},\n"
76
+ " doi = {10.1063/1.4930534}\n}"
77
+ ),
78
+ }
79
+
80
+
81
+ class GPU4PySCF(Backend):
82
+ """Second-generation ALMO-EDA in GPU4PySCF (1.9 or later)."""
83
+
84
+ name = "GPU4PySCF"
85
+ scheme = "ALMO-EDA2"
86
+ device = "gpu"
87
+ max_fragments = None
88
+ ini = "gpu4pyscf.ini"
89
+
90
+ def options(self):
91
+ return dict(DEFAULTS)
92
+
93
+ def settings(self, settings):
94
+ """The full settings: the defaults updated with those given."""
95
+ out = dict(DEFAULTS)
96
+ out.update({k: v for k, v in (settings or {}).items() if k in DEFAULTS})
97
+ return out
98
+
99
+ def files(self, calculations, settings):
100
+ data = {
101
+ "settings": self.settings(settings),
102
+ "calculations": [
103
+ {
104
+ "id": calc.id,
105
+ "fragments": [
106
+ {
107
+ "label": f.label(),
108
+ "symbols": list(f.symbols),
109
+ "coordinates": [
110
+ [float(x) for x in row] for row in f.coordinates
111
+ ],
112
+ "charge": int(f.charge),
113
+ "spin": int(f.multiplicity) - 1,
114
+ }
115
+ for f in calc.fragments
116
+ ],
117
+ }
118
+ for calc in calculations
119
+ ],
120
+ }
121
+ driver = (
122
+ importlib.resources.files("interaction_eda_step") / "data" / DRIVER
123
+ ).read_text()
124
+ return {
125
+ "input.json": json.dumps(data, indent=1),
126
+ DRIVER: driver,
127
+ }
128
+
129
+ def command(self):
130
+ # The logs go in the step's directory, not a subdirectory: the executor's
131
+ # in-situ clean-up (seamm_exec 2026.11.5) cannot remove a new directory
132
+ # holding files it returns.
133
+ return [
134
+ "{code}",
135
+ DRIVER,
136
+ "input.json",
137
+ "-o",
138
+ self.results_file,
139
+ "--logdir",
140
+ ".",
141
+ ]
142
+
143
+ def read_results(self, directory, calculations):
144
+ records = read_jsonl(Path(directory) / self.results_file)
145
+ out = {}
146
+ for calc in calculations:
147
+ record = records.get(calc.id)
148
+ if record is None:
149
+ out[calc.key] = Result(calc.key, error="no result (not run?)")
150
+ elif "error" in record and "eda_kJ_mol" not in record:
151
+ out[calc.key] = Result(calc.key, error=record["error"], data=record)
152
+ else:
153
+ out[calc.key] = Result(
154
+ calc.key,
155
+ energies={k: float(v) for k, v in record["eda_kJ_mol"].items()},
156
+ data=record,
157
+ )
158
+ return out
159
+
160
+ def method_string(self, settings):
161
+ s = self.settings(settings)
162
+ parts = [
163
+ f"scheme={self.scheme}",
164
+ f"backend={self.name}",
165
+ f"xc={s['xc']}",
166
+ f"basis={s['basis']}",
167
+ f"auxbasis={s['auxbasis'] or 'none'}",
168
+ f"nlc={'default' if s['nlc'] is None else s['nlc']}",
169
+ f"response_nlc={'on' if s['response_nlc'] else 'off'}",
170
+ f"disp={s['disp'] or 'none'}",
171
+ f"ferf={s['ferf']}",
172
+ f"svd_threshold={s['svd_threshold']:g}",
173
+ ]
174
+ return ";".join(parts)
175
+
176
+ def citations(self, settings):
177
+ return [
178
+ ("gpu4pyscf", BIBTEX["gpu4pyscf"], 1, "The GPU4PySCF code."),
179
+ ("pyscf2020", BIBTEX["pyscf"], 1, "PySCF, which GPU4PySCF builds on."),
180
+ ("almo_eda2", BIBTEX["almo_eda2"], 1, "Second-generation ALMO-EDA."),
181
+ ("almo_eda_frozen", BIBTEX["frozen"], 2, "The frozen-term decomposition."),
182
+ ("ferf", BIBTEX["ferf"], 2, "Polarization from FERFs."),
183
+ ]
@@ -0,0 +1,3 @@
1
+ [interaction-eda-step]
2
+ # The Interaction EDA step has no options here yet. How GPU4PySCF is run is in
3
+ # gpu4pyscf.ini in the SEAMM installation's directory.
@@ -0,0 +1,33 @@
1
+ # How to run GPU4PySCF for the Interaction EDA step.
2
+ #
3
+ # GPU4PySCF lives in its own Python environment (GPU4PySCF 1.9.1 or later, with
4
+ # CuPy, cuTENSOR and pyscf-dispersion) and needs Linux and an NVIDIA GPU with good
5
+ # double precision: A100, A30, H100/H200 or B200, not L40S or T4.
6
+ #
7
+ # Each section is for one executor ('local' runs on this machine).
8
+
9
+ [local]
10
+ # The type of installation. Options are:
11
+ # conda: the Conda environment below, which interaction-eda-step-installer
12
+ # creates and fills (the default)
13
+ # local: 'code' is the full path to the Python of an environment you made
14
+
15
+ installation = conda
16
+
17
+ # For a conda installation, the conda executable and the environment.
18
+
19
+ # conda = /path/to/conda
20
+ conda-environment = seamm-gpu4pyscf
21
+
22
+ # The Python that runs the step's driver script: 'python' in the Conda
23
+ # environment, or e.g. /home/me/envs/gpu4pyscf/bin/python for a local one.
24
+
25
+ code = python
26
+
27
+ # Which CUDA build of GPU4PySCF the installer puts in the environment. 'auto'
28
+ # reads the NVIDIA driver with nvidia-smi and uses the CUDA 12 build (the
29
+ # validated one, which also runs with CUDA 13 drivers); with no driver, as on a
30
+ # cluster login node, it also uses the CUDA 12 build. Set cuda12x or cuda13x to
31
+ # force one.
32
+
33
+ cuda-build = auto
@@ -0,0 +1,256 @@
1
+ #!/usr/bin/env python
2
+ """Second-generation ALMO-EDA with GPU4PySCF, for SEAMM's Interaction EDA step.
3
+
4
+ Runs in the GPU4PySCF environment, which needs nothing from SEAMM:
5
+
6
+ python gpu4pyscf_eda.py input.json [-o results.jsonl]
7
+
8
+ ``input.json``::
9
+
10
+ {"settings": {"xc": "wb97m-v", "basis": "def2-tzvpd", "auxbasis": null,
11
+ "nlc": null, "disp": null, "svd_threshold": null,
12
+ "ferf": "coupled", "response_nlc": true, "verbose": 3},
13
+ "calculations": [{"id": "0+1",
14
+ "fragments": [{"label": "0:Li+", "symbols": ["Li"],
15
+ "coordinates": [[0.0, 0.0, 0.0]],
16
+ "charge": 1, "spin": 0}, ...]}, ...]}
17
+
18
+ ``nlc: false`` removes VV10 from the whole functional; ``response_nlc: false``
19
+ removes it only from the FERF response kernel (the SCF and SCF-MI keep it).
20
+ Coordinates are in Å. Each calculation appends one JSON line to the results file:
21
+ the ALMO-EDA terms in kJ/mol, the energies in hartree, the wall time, the peak GPU
22
+ memory and the versions -- or the error. Calculations already in the results file
23
+ are skipped, so a job that ran out of time can simply be run again.
24
+
25
+ The GPU needs good double-precision throughput (A100, A30, H100/H200, B200).
26
+ """
27
+
28
+ import argparse
29
+ import functools
30
+ import json
31
+ import os
32
+ import re
33
+ import sys
34
+ import threading
35
+ import time
36
+ import traceback
37
+
38
+ # GPUs made for single precision: FP64 at 1/32-1/64 of FP32 (whole words, so
39
+ # that "A10" does not match "A100").
40
+ LOW_FP64 = re.compile(r"\b(T4|L4|L40S?|A10|A16|A40|RTX|GEFORCE)\b", re.IGNORECASE)
41
+
42
+
43
+ def low_fp64(gpu):
44
+ """Whether the GPU, by name, has little double-precision throughput."""
45
+ return LOW_FP64.search(gpu) is not None
46
+
47
+
48
+ def gpu_refusal(gpu, settings):
49
+ """Why the calculations must not run on this GPU, or None if they may.
50
+
51
+ A GPU made for single precision runs ALMO-EDA tens of times slower than one
52
+ with full double precision, so the step refuses it unless told to allow it.
53
+ """
54
+ if not low_fp64(gpu) or settings.get("allow_low_fp64", False):
55
+ return None
56
+ return (
57
+ f"The {gpu} has little double-precision (FP64) throughput, so the EDA "
58
+ "would be very slow: nothing was run. Use a queue with A100, A30, "
59
+ "H100/H200, B200 or V100 GPUs, or set 'GPUs with little FP64' to 'allow'."
60
+ )
61
+
62
+
63
+ class PeakGPUMemory:
64
+ """Sample the device's used memory (all processes on it) in a thread."""
65
+
66
+ def __init__(self, cp, interval=0.1):
67
+ self.cp = cp
68
+ self.interval = interval
69
+ self.peak = 0
70
+ self._stop = threading.Event()
71
+
72
+ def _run(self):
73
+ while not self._stop.is_set():
74
+ free, total = self.cp.cuda.runtime.memGetInfo()
75
+ self.peak = max(self.peak, total - free)
76
+ time.sleep(self.interval)
77
+
78
+ def __enter__(self):
79
+ self.base = self.cp.cuda.runtime.memGetInfo()
80
+ self.thread = threading.Thread(target=self._run, daemon=True)
81
+ self.thread.start()
82
+ return self
83
+
84
+ def __exit__(self, *args):
85
+ self._stop.set()
86
+ self.thread.join()
87
+
88
+
89
+ def done_ids(path):
90
+ ids = set()
91
+ if os.path.exists(path):
92
+ with open(path) as fd:
93
+ for line in fd:
94
+ try:
95
+ record = json.loads(line)
96
+ except ValueError:
97
+ continue
98
+ if "eda_kJ_mol" in record:
99
+ ids.add(record["id"])
100
+ return ids
101
+
102
+
103
+ def versions():
104
+ out = {"python": sys.version.split()[0]}
105
+ for name in ("gpu4pyscf", "pyscf", "cupy"):
106
+ try:
107
+ module = __import__(name)
108
+ out[name] = getattr(module, "__version__", "?")
109
+ except Exception: # pragma: no cover
110
+ out[name] = None
111
+ return out
112
+
113
+
114
+ def main(argv=None):
115
+ p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
116
+ p.add_argument("input")
117
+ p.add_argument("-o", "--output", default="results.jsonl")
118
+ p.add_argument("--logdir", default="logs")
119
+ args = p.parse_args(argv)
120
+
121
+ with open(args.input) as fd:
122
+ data = json.load(fd)
123
+ settings = data.get("settings", {})
124
+ calculations = data["calculations"]
125
+
126
+ import cupy as cp
127
+ import pyscf
128
+ import gpu4pyscf.properties.eda as eda_module
129
+
130
+ ferf_options = {}
131
+ if settings.get("svd_threshold") is not None:
132
+ ferf_options["virtual_singular_value_threshold"] = float(
133
+ settings["svd_threshold"]
134
+ )
135
+ if settings.get("ferf", "coupled") == "uncoupled":
136
+ ferf_options["uncoupled_ferf"] = True
137
+ if ferf_options:
138
+ eda_module._get_eda_polarization_FERF_subspace_projector = functools.partial(
139
+ eda_module._get_eda_polarization_FERF_subspace_projector, **ferf_options
140
+ )
141
+
142
+ if not settings.get("response_nlc", True):
143
+ # VV10 off in the FERF response kernel only; the SCF and SCF-MI keep the
144
+ # full functional (the Psi4-parity mode). eda.py looks gen_vind up at
145
+ # call time, so rebinding the module attribute is enough.
146
+ import gpu4pyscf.properties.polarizability as polarizability
147
+
148
+ _gen_vind = polarizability.gen_vind
149
+
150
+ def gen_vind_no_nlc(mf, mo_coeff, mo_occ, with_nlc=True):
151
+ return _gen_vind(mf, mo_coeff, mo_occ, with_nlc=False)
152
+
153
+ polarizability.gen_vind = gen_vind_no_nlc
154
+
155
+ nlc = settings.get("nlc")
156
+ if isinstance(nlc, str) and nlc.lower() in ("false", "no", "none", "off", "0", ""):
157
+ nlc = False
158
+ xc = settings.get("xc", "wb97m-v")
159
+ basis = settings.get("basis", "def2-tzvpd")
160
+ auxbasis = settings.get("auxbasis")
161
+ disp = settings.get("disp")
162
+ verbose = int(settings.get("verbose", 3))
163
+
164
+ gpu = cp.cuda.runtime.getDeviceProperties(0)["name"]
165
+ gpu = gpu.decode() if isinstance(gpu, bytes) else str(gpu)
166
+ info = versions()
167
+ print(f"GPU: {gpu}; {info}", flush=True)
168
+ refusal = gpu_refusal(gpu, settings)
169
+ if refusal is not None:
170
+ print(f"REFUSED: {refusal}", file=sys.stderr, flush=True)
171
+ return 3
172
+ if low_fp64(gpu):
173
+ print(
174
+ f"WARNING: the {gpu} has little double-precision (FP64) throughput, so "
175
+ "this will be very slow ('GPUs with little FP64' is 'allow').",
176
+ flush=True,
177
+ )
178
+ os.makedirs(args.logdir, exist_ok=True)
179
+ skip = done_ids(args.output)
180
+
181
+ n_failed = 0
182
+ with open(args.output, "a") as out:
183
+ for calc in calculations:
184
+ cid = calc["id"]
185
+ if cid in skip:
186
+ print(f"{cid}: already done", flush=True)
187
+ continue
188
+ log = os.path.join(args.logdir, f"{cid}.log")
189
+ with open(log, "w") as logfd:
190
+ record = {"id": cid, "gpu": gpu, "versions": info}
191
+ t0 = time.time()
192
+ try:
193
+ mols = []
194
+ for frag in calc["fragments"]:
195
+ geom = "\n".join(
196
+ "{} {:.10f} {:.10f} {:.10f}".format(s, *xyz)
197
+ for s, xyz in zip(frag["symbols"], frag["coordinates"])
198
+ )
199
+ m = pyscf.M(
200
+ atom=geom,
201
+ basis=basis,
202
+ charge=int(frag.get("charge", 0)),
203
+ spin=int(frag.get("spin", 0)),
204
+ verbose=0,
205
+ )
206
+ m.stdout = logfd # one log for all fragments
207
+ m.verbose = verbose
208
+ mols.append(m)
209
+ cp.get_default_memory_pool().free_all_blocks()
210
+ with PeakGPUMemory(cp) as mem:
211
+ eda, dft = eda_module.eval_ALMO_EDA_2_energies(
212
+ mols,
213
+ xc=xc,
214
+ auxbasis=auxbasis,
215
+ verbose=verbose,
216
+ nlc=nlc,
217
+ disp=disp,
218
+ )
219
+ cp.cuda.runtime.deviceSynchronize()
220
+ record["eda_kJ_mol"] = {
221
+ k: float(v) for k, v in eda.items() if k != "unit"
222
+ }
223
+ record["energies_hartree"] = _plain(dft.get("energy"))
224
+ record["nao"] = int(sum(m.nao for m in mols))
225
+ record["peak_gpu_mem_GB"] = round(mem.peak / 1e9, 3)
226
+ except Exception:
227
+ n_failed += 1
228
+ record["error"] = traceback.format_exc()
229
+ print(f"{cid}: FAILED\n{record['error']}", flush=True)
230
+ record["wall_s"] = round(time.time() - t0, 2)
231
+ out.write(json.dumps(record) + "\n")
232
+ out.flush()
233
+ if "eda_kJ_mol" in record:
234
+ e = record["eda_kJ_mol"]
235
+ print(
236
+ f"{cid}: total {e['total']:.3f} kJ/mol ({record['wall_s']} s, "
237
+ f"{record['peak_gpu_mem_GB']} GB)",
238
+ flush=True,
239
+ )
240
+ return 1 if n_failed else 0
241
+
242
+
243
+ def _plain(value):
244
+ """JSON-safe copy of numbers, lists and dicts (numpy/cupy scalars too)."""
245
+ if isinstance(value, dict):
246
+ return {str(k): _plain(v) for k, v in value.items()}
247
+ if isinstance(value, (list, tuple)):
248
+ return [_plain(v) for v in value]
249
+ try:
250
+ return float(value)
251
+ except (TypeError, ValueError):
252
+ return str(value)
253
+
254
+
255
+ if __name__ == "__main__":
256
+ sys.exit(main())