Skip to content

Commit 840543c

Browse files
alongdclaude
andcommitted
main: accept a YAML list form for adaptive_levels
Tuple keys (1, 6)/(7, 'inf') can't be produced by yaml.safe_load, so adaptive_levels was unusable from an input file. Accept a list of {atom_range, levels} entries, build the tuple-keyed structure internally, and serialize the same form on restart. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 3258243 commit 840543c

2 files changed

Lines changed: 118 additions & 73 deletions

File tree

arc/main.py

Lines changed: 61 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -98,20 +98,22 @@ class ARC(object):
9898
ess_settings (dict, optional): A dictionary of available ESS (keys) and a corresponding server list (values).
9999
bath_gas (str, optional): A bath gas. Currently used in OneDMin to calc L-J parameters.
100100
Allowed values are He, Ne, Ar, Kr, H2, N2, O2.
101-
adaptive_levels (dict, optional): A dictionary of levels of theory for ranges of the number of heavy atoms in
102-
the molecule. Keys are tuples of (min_num_atoms, max_num_atoms), values are dictionaries. Keys of the
103-
sub-dictionaries are tuples of job types, values are levels of theory (str, dict or Level).
101+
adaptive_levels (list, optional): A list of levels of theory for ranges of the number of heavy atoms in the
102+
molecule. Each entry is a dictionary with an ``atom_range`` 2-list (min_num_atoms, max_num_atoms; the upper
103+
bound may be ``'inf'``) and a ``levels`` mapping from job type to level of theory (str, dict or Level).
104+
Job types sharing a level may be given as a whitespace- or comma-separated key (e.g. ``'opt freq'``).
104105
Job types not defined in adaptive levels will have non-adaptive (regular) levels.
105106
Example::
106107
107-
adaptive_levels = {(1, 5): {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)',
108-
'sp': 'ccsd(t)-f12/aug-cc-pvtz-f12'},
109-
(6, 15): {('opt', 'freq'): 'b3lyp/cbsb7',
110-
'sp': 'dlpno-ccsd(t)/def2-tzvp'},
111-
(16, 30): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
112-
'sp': 'wb97xd/6-311+g(2d,2p)'},
113-
(31, 'inf'): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
114-
'sp': 'b3lyp/6-311+g(d,p)'}}
108+
adaptive_levels = [{'atom_range': [1, 5],
109+
'levels': {'opt freq': 'wb97xd/6-311+g(2d,2p)',
110+
'sp': 'ccsd(t)-f12/aug-cc-pvtz-f12'}},
111+
{'atom_range': [6, 15],
112+
'levels': {'opt freq': 'b3lyp/cbsb7',
113+
'sp': 'dlpno-ccsd(t)/def2-tzvp'}},
114+
{'atom_range': [16, 'inf'],
115+
'levels': {'opt freq': 'b3lyp/6-31g(d,p)',
116+
'sp': 'wb97xd/6-311+g(2d,2p)'}}]
115117
116118
freq_scale_factor (float, optional): The harmonic frequencies scaling factor. Could be automatically determined
117119
if not available in Arkane and not provided by the user.
@@ -163,9 +165,8 @@ class ARC(object):
163165
ts_guess_level (Level): Level of theory for comparisons of TS guesses between different methods.
164166
irc_level (Level): The level of theory to use for IRC calculations.
165167
orbitals_level (Level): Level of theory for molecular orbitals calculations.
166-
adaptive_levels (dict): A dictionary of levels of theory for ranges of the number of heavy atoms in
167-
the molecule. Keys are tuples of (min_num_atoms, max_num_atoms), values are dictionaries. Keys of the
168-
sub-dictionaries are tuples of job types, values are levels of theory (str, dict or Level).
168+
adaptive_levels (dict): The processed adaptive levels, keyed by (min_num_atoms, max_num_atoms) tuples, each
169+
mapping job-type tuples to ``Level`` objects (built from the user-facing ``adaptive_levels`` list).
169170
Job types not defined in adaptive levels will have non-adaptive (regular) levels.
170171
output (dict): Output dictionary with status and final QM file paths for all species. Only used for restarting,
171172
the actual object used is in the Scheduler class.
@@ -431,8 +432,10 @@ def as_dict(self) -> dict:
431432
"""
432433
restart_dict = dict()
433434
if self.adaptive_levels is not None:
434-
restart_dict['adaptive_levels'] = {atom_range: {job_type: level.as_dict() for job_type, level in levels_dict}
435-
for atom_range, levels_dict in self.adaptive_levels.items()}
435+
restart_dict['adaptive_levels'] = [
436+
{'atom_range': [atom_range[0], atom_range[1]],
437+
'levels': {' '.join(job_types): level.as_dict() for job_types, level in levels_dict.items()}}
438+
for atom_range, levels_dict in self.adaptive_levels.items()]
436439
if self.allow_nonisomorphic_2d:
437440
restart_dict['allow_nonisomorphic_2d'] = self.allow_nonisomorphic_2d
438441
if self.arkane_level_of_theory is not None:
@@ -1256,44 +1259,60 @@ def standardize_output_paths(self):
12561259
self.output = dict()
12571260

12581261

1259-
def process_adaptive_levels(adaptive_levels: dict | None) -> dict | None:
1262+
def process_adaptive_levels(adaptive_levels: list | None) -> dict | None:
12601263
"""
12611264
Process the ``adaptive_levels`` argument.
12621265
1266+
The user-facing form is a YAML-friendly list of entries, each a dictionary with an
1267+
``atom_range`` 2-list (the heavy-atom count range, the upper bound may be the string
1268+
``'inf'`` or ``float('inf')``) and a ``levels`` mapping of job types to levels of theory.
1269+
Job types that share a level may be given as a single whitespace- or comma-separated key
1270+
(e.g. ``'opt freq'``). A level value may be a string or a ``Level`` dictionary. For example::
1271+
1272+
adaptive_levels = [{'atom_range': [1, 6],
1273+
'levels': {'opt freq': 'wb97xd/def2tzvp',
1274+
'sp': 'ccsd(t)-f12/cc-pvtz-f12'}},
1275+
{'atom_range': [7, 'inf'],
1276+
'levels': {'opt freq': 'b3lyp/6-31g(d,p)',
1277+
'sp': 'dlpno-ccsd(t)/def2-tzvp'}}]
1278+
12631279
Args:
1264-
adaptive_levels (dict): The adaptive levels dictionary.
1280+
adaptive_levels (list): The adaptive levels specification (a list of entries).
12651281
1266-
Returns: dict
1267-
The processed adaptive levels dictionary.
1282+
Returns: dict | None
1283+
The processed adaptive levels keyed by ``(min_heavy_atoms, max_heavy_atoms)`` tuples,
1284+
each mapping job-type tuples to ``Level`` objects, or ``None`` if the input is ``None``.
12681285
"""
12691286
if adaptive_levels is None:
12701287
return None
1271-
processed = dict()
1272-
if not isinstance(adaptive_levels, dict):
1273-
raise InputError(f'The adaptive levels argument must be a dictionary, '
1288+
if not isinstance(adaptive_levels, list):
1289+
raise InputError(f'The adaptive levels argument must be a list of entries, '
12741290
f'got {adaptive_levels} which is a {type(adaptive_levels)}')
1275-
for atom_range, adaptive_level in adaptive_levels.items():
1276-
if not isinstance(atom_range, tuple) \
1277-
or not all([isinstance(a, int) or a == 'inf' for a in atom_range]) \
1278-
or len(atom_range) != 2:
1279-
raise InputError(f'Keys of the adaptive levels argument must be 2-length tuples of integers or an "inf" '
1280-
f'indicator, got {atom_range} which is a {type(atom_range)} in:\n{adaptive_levels}')
1281-
if not isinstance(adaptive_level, dict):
1282-
raise InputError(f'Each adaptive level in the adaptive levels argument must be a dictionary, '
1283-
f'got {adaptive_level} which is a {type(adaptive_level)} in:\n{adaptive_levels}')
1291+
processed = dict()
1292+
for entry in adaptive_levels:
1293+
if not isinstance(entry, dict) or 'atom_range' not in entry or 'levels' not in entry:
1294+
raise InputError(f'Each adaptive levels entry must be a dictionary with "atom_range" and "levels" '
1295+
f'keys, got {entry} which is a {type(entry)} in:\n{adaptive_levels}')
1296+
atom_range = entry['atom_range']
1297+
if not isinstance(atom_range, (list, tuple)) or len(atom_range) != 2 \
1298+
or not isinstance(atom_range[0], int) \
1299+
or not (isinstance(atom_range[1], int) or atom_range[1] in ('inf', float('inf'))):
1300+
raise InputError(f'The "atom_range" of each adaptive levels entry must be a 2-length list of an integer '
1301+
f'lower bound and an integer or "inf" upper bound, got {atom_range} '
1302+
f'in:\n{adaptive_levels}')
1303+
atom_range = (atom_range[0], 'inf' if atom_range[1] in ('inf', float('inf')) else atom_range[1])
1304+
levels = entry['levels']
1305+
if not isinstance(levels, dict):
1306+
raise InputError(f'The "levels" of each adaptive levels entry must be a dictionary, '
1307+
f'got {levels} which is a {type(levels)} in:\n{adaptive_levels}')
12841308
processed[atom_range] = dict()
1285-
for sub_key, level in adaptive_level.items():
1286-
new_sub_key = (sub_key,) if isinstance(sub_key, str) else sub_key
1287-
if not isinstance(new_sub_key, tuple):
1288-
raise InputError(f'Job types specifications in adaptive levels must be tuples, got {sub_key} '
1289-
f'which is a {type(sub_key)} in:\n{adaptive_levels}')
1290-
new_level = Level(repr=level)
1291-
processed[atom_range][new_sub_key] = new_level
1292-
atom_ranges = sorted(list(adaptive_levels.keys()), key=lambda x: x[0])
1309+
for job_types, level in levels.items():
1310+
job_type_tuple = tuple(job_types.replace(',', ' ').split())
1311+
processed[atom_range][job_type_tuple] = Level(repr=level)
1312+
atom_ranges = sorted(processed.keys(), key=lambda x: x[0])
12931313
for i, atom_range in enumerate(atom_ranges):
12941314
if i and atom_ranges[i-1][1] + 1 != atom_ranges[i][0]:
1295-
raise InputError(f'Atom ranges of adaptive levels must be consecutive. '
1296-
f'Got:\n{list(adaptive_levels.keys())}')
1315+
raise InputError(f'Atom ranges of adaptive levels must be consecutive. Got:\n{atom_ranges}')
12971316
if atom_ranges[-1][1] != 'inf':
12981317
raise InputError(f'The last atom range must be "inf", got {atom_ranges[-1][1]} in {atom_ranges}')
12991318
return processed

arc/main_test.py

Lines changed: 57 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -387,46 +387,72 @@ def test_add_hydrogen_for_bde(self):
387387
self.assertIn('H', [spc.label for spc in arc1.species])
388388

389389
def test_process_adaptive_levels(self):
390-
"""Test processing the adaptive levels"""
391-
adaptive_levels_1 = {(1, 5): {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)',
392-
('sp',): 'ccsd(t)-f12/aug-cc-pvtz-f12'},
393-
(6, 15): {('opt', 'freq'): 'b3lyp/cbsb7',
394-
'sp': 'dlpno-ccsd(t)/def2-tzvp'},
395-
(16, 30): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
396-
'sp': {'method': 'wb97xd', 'basis': '6-311+g(2d,2p)'}},
397-
(31, 'inf'): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
398-
'sp': 'b3lyp/6-311+g(d,p)'}}
399-
390+
"""Test processing the adaptive levels (YAML-friendly list-of-entries schema)"""
391+
# None passes through.
392+
self.assertIsNone(process_adaptive_levels(None))
393+
394+
# A normal, multi-range specification. Job types sharing a level are given as a
395+
# whitespace- or comma-separated key; a level may be a string or a Level dict.
396+
adaptive_levels_1 = [{'atom_range': [1, 5],
397+
'levels': {'opt freq': 'wb97xd/6-311+g(2d,2p)',
398+
'sp': 'ccsd(t)-f12/aug-cc-pvtz-f12'}},
399+
{'atom_range': [6, 15],
400+
'levels': {'opt, freq': 'b3lyp/cbsb7',
401+
'sp': 'dlpno-ccsd(t)/def2-tzvp'}},
402+
{'atom_range': [16, 30],
403+
'levels': {'opt freq': 'b3lyp/6-31g(d,p)',
404+
'sp': {'method': 'wb97xd', 'basis': '6-311+g(2d,2p)'}}},
405+
{'atom_range': [31, 'inf'],
406+
'levels': {'opt freq': 'b3lyp/6-31g(d,p)',
407+
'sp': 'b3lyp/6-311+g(d,p)'}}]
400408
processed_1 = process_adaptive_levels(adaptive_levels_1)
401409
self.assertEqual(processed_1[(6, 15)][('sp',)].simple(), 'dlpno-ccsd(t)/def2-tzvp')
402410
self.assertEqual(processed_1[(16, 30)][('sp',)].simple(), 'wb97xd/6-311+g(2d,2p)')
403-
404-
# test non dict
411+
self.assertEqual(processed_1[(1, 5)][('opt', 'freq')].simple(), 'wb97xd/6-311+g(2d,2p)')
412+
413+
# A single range covering everything, and a float 'inf' is accepted as the upper bound.
414+
processed_2 = process_adaptive_levels([{'atom_range': [1, float('inf')],
415+
'levels': {'opt freq': 'b3lyp/6-31g(d,p)',
416+
'sp': 'b3lyp/6-311+g(d,p)'}}])
417+
self.assertEqual(processed_2[(1, 'inf')][('sp',)].simple(), 'b3lyp/6-311+g(d,p)')
418+
419+
# Restart round-trip: as_dict() must emit the list form and reproduce the same structure.
420+
arc0 = ARC(project='adaptive_levels_test', adaptive_levels=adaptive_levels_1)
421+
restart_levels = arc0.as_dict()['adaptive_levels']
422+
self.assertIsInstance(restart_levels, list)
423+
reprocessed = process_adaptive_levels(restart_levels)
424+
self.assertEqual(reprocessed[(6, 15)][('sp',)].simple(), 'dlpno-ccsd(t)/def2-tzvp')
425+
self.assertEqual(set(reprocessed.keys()), set(processed_1.keys()))
426+
427+
# Not a list (the legacy tuple-dict form is no longer accepted).
405428
with self.assertRaises(InputError):
406429
process_adaptive_levels(4)
407-
# wrong atom range
408430
with self.assertRaises(InputError):
409-
process_adaptive_levels({5: {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)',
410-
('sp',): 'ccsd(t)-f12/aug-cc-pvtz-f12'},
411-
(6, 'inf'): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
412-
'sp': 'b3lyp/6-311+g(d,p)'}})
413-
# no 'inf
431+
process_adaptive_levels({(1, 5): {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)'},
432+
(6, 'inf'): {'sp': 'b3lyp/6-311+g(d,p)'}})
433+
# atom_range is not a 2-length list.
434+
with self.assertRaises(InputError):
435+
process_adaptive_levels([{'atom_range': [5], 'levels': {'sp': 'b3lyp/6-311+g(d,p)'}}])
436+
# 'inf' is only allowed as the upper bound.
437+
with self.assertRaises(InputError):
438+
process_adaptive_levels([{'atom_range': [float('inf'), 'inf'], 'levels': {'sp': 'b3lyp/6-311+g(d,p)'}}])
439+
with self.assertRaises(InputError):
440+
process_adaptive_levels([{'atom_range': ['inf', 10], 'levels': {'sp': 'b3lyp/6-311+g(d,p)'}}])
441+
# The last range does not end with 'inf'.
442+
with self.assertRaises(InputError):
443+
process_adaptive_levels([{'atom_range': [1, 5], 'levels': {'sp': 'wb97xd/def2tzvp'}},
444+
{'atom_range': [6, 75], 'levels': {'sp': 'b3lyp/6-311+g(d,p)'}}])
445+
# 'levels' is not a dict.
414446
with self.assertRaises(InputError):
415-
process_adaptive_levels({(1, 5): {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)',
416-
('sp',): 'ccsd(t)-f12/aug-cc-pvtz-f12'},
417-
(6, 75): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
418-
'sp': 'b3lyp/6-311+g(d,p)'}})
419-
# adaptive level not a dict
447+
process_adaptive_levels([{'atom_range': [1, 5], 'levels': {'sp': 'wb97xd/def2tzvp'}},
448+
{'atom_range': [6, 'inf'], 'levels': 'b3lyp/6-31g(d,p)'}])
449+
# Non-consecutive atom ranges.
420450
with self.assertRaises(InputError):
421-
process_adaptive_levels({(1, 5): {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)',
422-
('sp',): 'ccsd(t)-f12/aug-cc-pvtz-f12'},
423-
(6, 'inf'): 'b3lyp/6-31g(d,p)'})
424-
# non-consecutive atom ranges
451+
process_adaptive_levels([{'atom_range': [1, 5], 'levels': {'sp': 'wb97xd/def2tzvp'}},
452+
{'atom_range': [15, 'inf'], 'levels': {'sp': 'b3lyp/6-311+g(d,p)'}}])
453+
# An entry missing required keys.
425454
with self.assertRaises(InputError):
426-
process_adaptive_levels({(1, 5): {('opt', 'freq'): 'wb97xd/6-311+g(2d,2p)',
427-
('sp',): 'ccsd(t)-f12/aug-cc-pvtz-f12'},
428-
(15, 'inf'): {('opt', 'freq'): 'b3lyp/6-31g(d,p)',
429-
'sp': 'b3lyp/6-311+g(d,p)'}})
455+
process_adaptive_levels([{'levels': {'sp': 'wb97xd/def2tzvp'}}])
430456

431457
def test_process_level_of_theory(self):
432458
"""

0 commit comments

Comments
 (0)