scheduler: keep reaction levels consistent under adaptive_levels

Levels were chosen per species from its own heavy-atom count, so a
reaction's large TS and its small wells could fall on different grains,
mixing levels of theory across the barrier. Key the whole reaction by
its (conserved) heavy-atom count, and for any well on a coarser grain
make an autonomous relabeled copy that the reaction uses at the
reaction-wide level. The new per-species thermo_at_own_level flag
(default True) keeps each species' own granular level for thermo; set
it False to evaluate the species at the reaction-wide level with no copy.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: alongd

arc/scheduler.py

@@ -52,6 +52,7 @@
 from arc.level import Level
 from arc.species.species import (ARCSpecies,
                                  are_coords_compliant_with_graph,
+                                 check_label,
                                  determine_rotor_symmetry,
                                  TSGuess)
 from arc.species.converter import (check_isomorphism,
@@ -331,6 +332,8 @@ def __init__(self,
         self.save_restart = False
 
         if len(self.rxn_list):
+            if self.adaptive_levels is not None:
+                self._apply_adaptive_reaction_levels()
             rxn_info_path = self.make_reaction_labels_info_file()
             for rxn in self.rxn_list:
                 logger.info('\n\n')
@@ -920,8 +923,10 @@ def run_job(self,
             torsions = species.rotors_dict[rotor_index]['torsion']
             torsions = [torsions] if not isinstance(torsions[0], list) else torsions
         if self.adaptive_levels is not None and label is not None:
+            spc = self.species_dict[label]
+            heavy_atoms = spc.adaptive_lot_n_heavy if spc.adaptive_lot_n_heavy is not None else spc.number_of_heavy_atoms
             level_of_theory = self.determine_adaptive_level(original_level_of_theory=level_of_theory, job_type=job_type,
-                                                            heavy_atoms=self.species_dict[label].number_of_heavy_atoms)
+                                                            heavy_atoms=heavy_atoms)
         job_adapter = job_adapter.lower() if job_adapter is not None else \
             self.deduce_job_adapter(level=Level(repr=level_of_theory), job_type=job_type)
         args = {'keyword': {}, 'block': {}}
@@ -3936,6 +3941,78 @@ def determine_adaptive_level(self,
         # for any other job type use the original level of theory regardless of the number of heavy atoms
         return original_level_of_theory
 
+    def _adaptive_atom_range(self, heavy_atoms: int) -> tuple | None:
+        """
+        Determine the adaptive levels atom range (the LOT grain) that a heavy-atom count falls into.
+
+        Args:
+            heavy_atoms (int): The number of heavy atoms.
+
+        Returns:
+            tuple | None: The ``(min, max)`` atom range key from ``self.adaptive_levels``, or ``None`` if not found.
+        """
+        for atom_range in self.adaptive_levels.keys():
+            if (atom_range[1] == 'inf' and heavy_atoms >= atom_range[0]) \
+                    or (atom_range[1] != 'inf' and atom_range[0] <= heavy_atoms <= atom_range[1]):
+                return atom_range
+        return None
+
+    def _apply_adaptive_reaction_levels(self):
+        """
+        Make every reaction's energetics internally consistent under adaptive levels of theory.
+
+        Heavy atoms are conserved across a reaction, so the reaction-wide heavy-atom count (the sum over the reactants,
+        equal to the TS supermolecule count) keys a single adaptive level for the whole reaction. A reaction
+        participant whose own heavy-atom count lands it on a *different* (finer) adaptive grain than the reaction
+        would otherwise be evaluated at an inconsistent level, mixing levels of theory across the barrier. To avoid
+        this, each such participant whose ``thermo_at_own_level`` is ``False`` (the default) is itself evaluated at
+        the reaction-wide level (no copy). If ``thermo_at_own_level`` is ``True``, an autonomous relabeled copy of
+        the species is created from the outset and used by the reaction (evaluated at the reaction-wide level),
+        while the original species is left to compute its own thermochemistry at its own granular level. A copy is
+        also created for a no-copy participant that is shared across reactions landing on different grains, since a
+        single per-species override cannot keep both reactions internally consistent.
+
+        This runs once during setup, before the reactions are processed, so each reaction is defined with its copies
+        from the start (its label, reactants, and products are kept mutually consistent for ``check_attributes`` and
+        restart).
+        """
+        for rxn_i, rxn in enumerate(self.rxn_list):
+            rxn_index = rxn.index if rxn.index is not None else rxn_i
+            reactant_species = [self.species_dict[label] for label in rxn.reactants if label in self.species_dict]
+            if len(reactant_species) != len(rxn.reactants) \
+                    or any(spc.number_of_heavy_atoms is None for spc in reactant_species):
+                logger.warning(f'Could not determine reaction-wide adaptive levels for {rxn.label}, '
+                               f'using per-species levels.')
+                continue
+            reaction_n_heavy = sum(spc.number_of_heavy_atoms for spc in reactant_species)
+            reaction_range = self._adaptive_atom_range(reaction_n_heavy)
+            for participants in (rxn.reactants, rxn.products):
+                for pos, label in enumerate(participants):
+                    spc = self.species_dict.get(label)
+                    if spc is None or spc.number_of_heavy_atoms is None \
+                            or self._adaptive_atom_range(spc.number_of_heavy_atoms) == reaction_range:
+                        continue
+                    if not spc.thermo_at_own_level:
+                        if spc.adaptive_lot_n_heavy is None \
+                                or self._adaptive_atom_range(spc.adaptive_lot_n_heavy) == reaction_range:
+                            spc.adaptive_lot_n_heavy = reaction_n_heavy
+                            continue
+                        logger.warning(f'Species {label} participates in reactions on different adaptive level '
+                                       f'grains, creating a dedicated copy of it for reaction {rxn.label} to keep '
+                                       f'the reaction internally consistent.')
+                    copy_label = check_label(f'{label}_TS{rxn_index}')[0]
+                    if copy_label not in self.species_dict:
+                        copy_spc = spc.copy()
+                        copy_spc.label = copy_label
+                        copy_spc.adaptive_lot_n_heavy = reaction_n_heavy
+                        copy_spc.compute_thermo = False
+                        copy_spc.include_in_thermo_lib = False
+                        self.species_list.append(copy_spc)
+                        self.species_dict[copy_label] = copy_spc
+                        self.initialize_output_dict(copy_label)
+                    participants[pos] = copy_label
+            rxn.label = rxn.arrow.join([rxn.plus.join(rxn.reactants), rxn.plus.join(rxn.products)])
+
     def initialize_output_dict(self, label: str | None = None):
         """
         Initialize self.output.

arc/scheduler_test.py

@@ -1241,5 +1241,125 @@ def test_spawn_ts_jobs_unknown_family_admission_predicate(self):
             self.assertEqual(admit_unknown_family, expected_admission)
 
 
+class TestSchedulerAdaptiveReactionLevels(unittest.TestCase):
+    """
+    Contains unit tests for the reaction-wide adaptive levels of theory logic
+    (Scheduler._apply_adaptive_reaction_levels and the spawn_job override).
+    """
+    @classmethod
+    def setUpClass(cls):
+        """A method that is run before all unit tests in this class."""
+        cls.ess_settings = {'gaussian': ['server1'], 'molpro': ['server2', 'server1'], 'qchem': ['server1']}
+        # A grain at exactly 1 heavy atom, and another for everything larger, so single-heavy-atom wells
+        # (e.g. CH4, OH) land on a different grain than a 2-heavy-atom reaction.
+        cls.adaptive_levels = {(1, 1): {('opt', 'freq'): Level(repr='wb97xd/def2tzvp'),
+                                        ('sp',): Level(repr='ccsd(t)-f12/cc-pvtz-f12')},
+                               (2, 'inf'): {('opt', 'freq'): Level(repr='b3lyp/6-31g(d,p)'),
+                                            ('sp',): Level(repr='b3lyp/6-311+g(d,p)')}}
+
+    def build_scheduler(self, rxn, species_list, name):
+        """
+        Build a testing Scheduler for a single reaction under the class adaptive levels.
+
+        Args:
+            rxn (ARCReaction): The reaction.
+            species_list (list): The species list.
+            name (str): A unique project name.
+
+        Returns:
+            Scheduler: The constructed (testing) scheduler.
+        """
+        project_directory = os.path.join(ARC_PATH, 'Projects', f'{name}_delete')
+        self.addCleanup(shutil.rmtree, project_directory, ignore_errors=True)
+        return Scheduler(project=name,
+                         ess_settings=self.ess_settings,
+                         species_list=species_list,
+                         rxn_list=[rxn],
+                         opt_level=Level(repr='b3lyp/6-31g(d,p)'),
+                         sp_level=Level(repr='b3lyp/6-311+g(d,p)'),
+                         freq_level=Level(repr='b3lyp/6-31g(d,p)'),
+                         adaptive_levels=self.adaptive_levels,
+                         project_directory=project_directory,
+                         job_types=initialize_job_types(),
+                         testing=True)
+
+    def test_bimolecular_creates_copies(self):
+        """Test that a reaction with wells on a different grain than the reaction gets relabeled copies"""
+        r = [ARCSpecies(label='CH4', smiles='C'), ARCSpecies(label='OH', smiles='[OH]')]
+        p = [ARCSpecies(label='CH3', smiles='[CH3]'), ARCSpecies(label='H2O', smiles='O')]
+        rxn = ARCReaction(label='CH4 + OH <=> CH3 + H2O', r_species=r, p_species=p)
+        sched = self.build_scheduler(rxn, r + p, 'adaptive_bimol')
+
+        # The reaction is now defined with the copies, and stays self-consistent.
+        self.assertEqual(rxn.reactants, ['CH4_TS0', 'OH_TS0'])
+        self.assertEqual(rxn.products, ['CH3_TS0', 'H2O_TS0'])
+        self.assertEqual([s.label for s in rxn.r_species], ['CH4_TS0', 'OH_TS0'])
+        self.assertEqual(rxn.label, 'CH4_TS0 + OH_TS0 <=> CH3_TS0 + H2O_TS0')
+        rxn.check_attributes()  # Must not raise.
+
+        # The copies are autonomous, evaluated at the reaction-wide level, and not part of the thermo library.
+        for copy_label in ['CH4_TS0', 'OH_TS0', 'CH3_TS0', 'H2O_TS0']:
+            self.assertIn(copy_label, sched.species_dict)
+            self.assertIn(copy_label, sched.output)
+            copy_spc = sched.species_dict[copy_label]
+            self.assertEqual(copy_spc.adaptive_lot_n_heavy, 2)
+            self.assertFalse(copy_spc.compute_thermo)
+            self.assertFalse(copy_spc.include_in_thermo_lib)
+
+        # The originals are untouched - they keep their own granular level and their own thermo.
+        for original_label in ['CH4', 'OH', 'CH3', 'H2O']:
+            original = sched.species_dict[original_label]
+            self.assertIsNone(original.adaptive_lot_n_heavy)
+            self.assertTrue(original.compute_thermo)
+
+    def test_thermo_at_own_level_false_no_copy(self):
+        """Test that with thermo_at_own_level=False the species itself takes the reaction-wide level, with no copy"""
+        r = [ARCSpecies(label='CH4', smiles='C', thermo_at_own_level=False),
+             ARCSpecies(label='OH', smiles='[OH]', thermo_at_own_level=False)]
+        p = [ARCSpecies(label='CH3', smiles='[CH3]', thermo_at_own_level=False),
+             ARCSpecies(label='H2O', smiles='O', thermo_at_own_level=False)]
+        rxn = ARCReaction(label='CH4 + OH <=> CH3 + H2O', r_species=r, p_species=p)
+        sched = self.build_scheduler(rxn, r + p, 'adaptive_noflag')
+
+        self.assertEqual(rxn.label, 'CH4 + OH <=> CH3 + H2O')
+        self.assertFalse(any('_TS' in label for label in sched.species_dict))
+        self.assertEqual(sched.species_dict['CH4'].adaptive_lot_n_heavy, 2)
+
+    def test_unimolecular_no_copy(self):
+        """Test that a reaction whose well shares the reaction's grain gets no copies"""
+        r = [ARCSpecies(label='nC3H7', smiles='[CH2]CC')]
+        p = [ARCSpecies(label='iC3H7', smiles='C[CH]C')]
+        rxn = ARCReaction(label='nC3H7 <=> iC3H7', r_species=r, p_species=p)
+        sched = self.build_scheduler(rxn, r + p, 'adaptive_unimol')
+
+        self.assertEqual(rxn.label, 'nC3H7 <=> iC3H7')
+        self.assertFalse(any('_TS' in label for label in sched.species_dict))
+        self.assertIsNone(sched.species_dict['nC3H7'].adaptive_lot_n_heavy)
+
+    def test_spawn_job_uses_override(self):
+        """Test that determine_adaptive_level is driven by adaptive_lot_n_heavy when set"""
+        spc = ARCSpecies(label='CH4', smiles='C')
+        sched = Scheduler(project='adaptive_override',
+                          ess_settings=self.ess_settings,
+                          species_list=[spc],
+                          opt_level=Level(repr='b3lyp/6-31g(d,p)'),
+                          sp_level=Level(repr='b3lyp/6-311+g(d,p)'),
+                          freq_level=Level(repr='b3lyp/6-31g(d,p)'),
+                          adaptive_levels=self.adaptive_levels,
+                          project_directory=os.path.join(ARC_PATH, 'Projects', 'adaptive_override_delete'),
+                          job_types=initialize_job_types(),
+                          testing=True)
+        self.addCleanup(shutil.rmtree, os.path.join(ARC_PATH, 'Projects', 'adaptive_override_delete'),
+                        ignore_errors=True)
+        original = Level(method='CBS-QB3')
+        # 1 heavy atom (own count) -> the (1, 1) grain.
+        self.assertEqual(sched.determine_adaptive_level(original, 'sp', spc.number_of_heavy_atoms).simple(),
+                         'ccsd(t)-f12/cc-pvtz-f12')
+        # With the override set to a 2-heavy-atom reaction, the (2, 'inf') grain is used instead.
+        spc.adaptive_lot_n_heavy = 2
+        heavy = spc.adaptive_lot_n_heavy if spc.adaptive_lot_n_heavy is not None else spc.number_of_heavy_atoms
+        self.assertEqual(sched.determine_adaptive_level(original, 'sp', heavy).simple(), 'b3lyp/6-311+g(d,p)')
+
+
 if __name__ == '__main__':
     unittest.main(testRunner=unittest.TextTestRunner(verbosity=2))

arc/species/species.py

@@ -109,6 +109,11 @@ class ARCSpecies(object):
         bond_corrections (dict, optional): The bond additivity corrections (BAC) to be used. Determined from the
                                            structure if not directly given.
         compute_thermo (bool, optional): Whether to calculate thermodynamic properties for this species.
+        thermo_at_own_level (bool, optional): Relevant only when ``adaptive_levels`` are used. If ``True`` (default),
+            the species' thermochemistry is computed at its own size-appropriate (granular) adaptive level even when
+            it participates in a reaction; the reaction's energetics then use a separate relabeled copy of the species
+            evaluated at the reaction-consistent level. If ``False``, the species itself is evaluated at the
+            reaction-consistent (coarser) level and no copy is made.
         include_in_thermo_lib (bool, optional): Whether to include in the output RMG library.
         e0_only (bool, optional): Whether to only run statmech (w/o thermo) to compute E0.
         species_dict (dict, optional): A dictionary to create this object from (used when restarting ARC).
@@ -231,6 +236,11 @@ class ARCSpecies(object):
         t1 (float): The T1 diagnostic parameter from Molpro.
         neg_freqs_trshed (list): A list of negative frequencies this species was troubleshooted for.
         compute_thermo (bool): Whether to calculate thermodynamic properties for this species.
+        thermo_at_own_level (bool): Whether to compute this species' thermochemistry at its own granular adaptive
+            level even when it participates in a reaction (see the constructor argument; relevant for ``adaptive_levels``).
+        adaptive_lot_n_heavy (int): An internal override for adaptive-levels selection - the heavy-atom count to key the
+            level by instead of this species' own count (set to the reaction-wide count for reaction participants).
+            ``None`` means use the species' own heavy-atom count.
         include_in_thermo_lib (bool): Whether to include in the output RMG library.
         e0_only (bool): Whether to only run statmech (w/o thermo) to compute E0.
         thermo (ThermoData): The thermo data calculated by ARC with 'H298' in kJ/mol and 'S298' in J/mol*K.
@@ -307,6 +317,8 @@ def __init__(self,
                  charge: int | None = None,
                  checkfile: str | None = None,
                  compute_thermo: bool | None = None,
+                 thermo_at_own_level: bool = True,
+                 adaptive_lot_n_heavy: int | None = None,
                  include_in_thermo_lib: bool | None = True,
                  consider_all_diastereomers: bool = True,
                  directed_rotors: dict | None = None,
@@ -409,6 +421,8 @@ def __init__(self,
             self.chosen_ts_method = None
             self.chosen_ts_list = list()
             self.compute_thermo = compute_thermo if compute_thermo is not None else not self.is_ts
+            self.thermo_at_own_level = thermo_at_own_level
+            self.adaptive_lot_n_heavy = adaptive_lot_n_heavy
             self.include_in_thermo_lib = include_in_thermo_lib
             self.e0_only = e0_only
             self.long_thermo_description = ''
@@ -688,6 +702,10 @@ def as_dict(self,
             species_dict['charge'] = self.charge
         if not self.compute_thermo and not self.is_ts:
             species_dict['compute_thermo'] = self.compute_thermo
+        if not self.thermo_at_own_level:
+            species_dict['thermo_at_own_level'] = self.thermo_at_own_level
+        if self.adaptive_lot_n_heavy is not None:
+            species_dict['adaptive_lot_n_heavy'] = self.adaptive_lot_n_heavy
         if not self.include_in_thermo_lib:
             species_dict['include_in_thermo_lib'] = self.include_in_thermo_lib
         species_dict['number_of_rotors'] = self.number_of_rotors
@@ -884,6 +902,8 @@ def from_dict(self, species_dict):
         self.multi_species = species_dict['multi_species'] if 'multi_species' in species_dict else None
         self.charge = species_dict['charge'] if 'charge' in species_dict else 0
         self.compute_thermo = species_dict['compute_thermo'] if 'compute_thermo' in species_dict else not self.is_ts
+        self.thermo_at_own_level = species_dict['thermo_at_own_level'] if 'thermo_at_own_level' in species_dict else True
+        self.adaptive_lot_n_heavy = species_dict['adaptive_lot_n_heavy'] if 'adaptive_lot_n_heavy' in species_dict else None
         self.include_in_thermo_lib = species_dict['include_in_thermo_lib'] if 'include_in_thermo_lib' in species_dict else True
         self.e0_only = species_dict['e0_only'] if 'e0_only' in species_dict else False
         self.number_of_radicals = species_dict['number_of_radicals'] if 'number_of_radicals' in species_dict else None