Enhances Arkane QM correction matching (#858)

Author: calvinp0

arc/main.py

@@ -41,7 +41,7 @@
 from arc.species.converter import str_to_xyz
 from arc.species.species import ARCSpecies
 from arc.statmech.adapter import StatmechEnum
-from arc.statmech.arkane import check_arkane_bacs
+from arc.statmech.arkane import check_arkane_aec, check_arkane_bacs
 from arc.utils.scale import determine_scaling_factors
 
 
@@ -1201,13 +1201,26 @@ def check_arkane_level_of_theory(self):
         """
         Check that the level of theory has AEC in Arkane.
         """
+        explicitly_set = self.arkane_level_of_theory is not None
         if self.arkane_level_of_theory is None:
             self.arkane_level_of_theory = self.composite_method if self.composite_method is not None \
                 else self.sp_level if self.sp_level is not None else None
         if self.arkane_level_of_theory is None:
             logger.warning('Could not determine a level of theory to be used for Arkane!')
-        elif self.bac_type is not None:
-            check_arkane_bacs(sp_level=self.arkane_level_of_theory, bac_type=self.bac_type, raise_error=self.compute_thermo)
+        else:
+            if explicitly_set:
+                source = ''
+            elif self.composite_method is not None:
+                source = ' (from composite method)'
+            else:
+                source = ' (from sp level)'
+            logger.info(f'Arkane level of theory:{source} {self.arkane_level_of_theory}')
+            if self.bac_type is not None:
+                check_arkane_bacs(sp_level=self.arkane_level_of_theory, bac_type=self.bac_type,
+                                  raise_error=self.compute_thermo)
+            else:
+                check_arkane_aec(sp_level=self.arkane_level_of_theory,
+                                 raise_error=self.compute_thermo)
 
     def backup_restart(self):
         """

arc/output.py

@@ -20,7 +20,12 @@
 from arc.job.local import execute_command
 from arc.parser.parser import parse_1d_scan_energies, parse_e_elect, parse_ess_version, parse_opt_steps, parse_zpe_correction
 from arc.species.converter import xyz_to_str
-from arc.statmech.arkane import get_arkane_model_chemistry
+from arc.statmech.arkane import (
+    AEC_SECTION_START, AEC_SECTION_END,
+    MBAC_SECTION_START, MBAC_SECTION_END,
+    PBAC_SECTION_START, PBAC_SECTION_END,
+    find_best_across_files, get_qm_corrections_files,
+)
 
 
 logger = get_logger()
@@ -282,23 +287,31 @@ def _get_energy_corrections(arkane_level_of_theory, bac_type: Optional[str]) ->
     Look up the AEC (per-atom, Hartree) and BAC (per-bond, kJ/mol) values
     that Arkane used from the RMG database for the given level of theory.
 
-    Uses ``get_arkane_model_chemistry`` to find the matched key, then calls
-    ``arc/scripts/get_qm_corrections.py`` as a subprocess to parse the
-    actual correction dicts from the RMG database.
+    Finds the AEC and BAC keys independently via fuzzy matching in their
+    respective database sections, then calls ``arc/scripts/get_qm_corrections.py``
+    as a subprocess to extract the actual correction dicts.
 
     Returns:
         (aec_dict_or_None, bac_dict_or_None)
     """
     if arkane_level_of_theory is None:
         return None, None
     try:
-        matched_key = get_arkane_model_chemistry(
-            sp_level=arkane_level_of_theory,
-            freq_scale_factor=1.0,  # dummy — we only need the energy key
-        )
-        if matched_key is None:
+        qm_corr_files = get_qm_corrections_files()
+
+        aec_key = find_best_across_files(arkane_level_of_theory, qm_corr_files,
+                                          AEC_SECTION_START, AEC_SECTION_END)
+        if aec_key is None:
             return None, None
 
+        bac_key = None
+        if bac_type in ('p', 'm'):
+            if bac_type == 'm':
+                bac_start, bac_end = MBAC_SECTION_START, MBAC_SECTION_END
+            else:
+                bac_start, bac_end = PBAC_SECTION_START, PBAC_SECTION_END
+            bac_key = find_best_across_files(arkane_level_of_theory, qm_corr_files, bac_start, bac_end)
+
         script_path = os.path.join(ARC_PATH, 'arc', 'scripts', 'get_qm_corrections.py')
         rmg_env = settings.get('RMG_ENV_NAME', 'rmg_env')
 
@@ -308,7 +321,8 @@ def _get_energy_corrections(arkane_level_of_theory, bac_type: Optional[str]) ->
             os.close(fd_in)
             os.close(fd_out)
             save_yaml_file(path=tmp_in, content={
-                'matched_key': matched_key,
+                'aec_key': aec_key,
+                'bac_key': bac_key,
                 'bac_type': bac_type,
             })
 

arc/output_test.py

@@ -479,6 +479,40 @@ def test_no_bac_when_type_none(self):
         aec, bac = _get_energy_corrections(lot, None)
         self.assertIsNone(bac)
 
+    def test_independent_aec_and_bac_keys(self):
+        """AEC and BAC keys should be resolved independently, not reusing the AEC key for BAC."""
+        lot = Level(method='wb97xd', basis='def2tzvp', software='gaussian')
+        aec_key = "LevelOfTheory(method='wb97xd',basis='def2tzvp',software='gaussian')"
+        bac_key = "LevelOfTheory(method='wb97xd',basis='def2tzvp')"  # different key (no software)
+
+        calls = []
+        def mock_find_best(level, files, start, end):
+            calls.append(start)
+            if 'atom_energies' in start:
+                return aec_key
+            elif 'pbac' in start:
+                return bac_key
+            return None
+
+        with patch('arc.output.find_best_across_files', side_effect=mock_find_best), \
+             patch('arc.output.get_qm_corrections_files', return_value=['/fake/data.py']), \
+             patch('arc.output.execute_command', return_value=('', '')), \
+             patch('arc.output.read_yaml_file', return_value={'aec': {'H': -0.5}, 'bac': {'C-H': -0.06}}), \
+             patch('arc.output.save_yaml_file') as mock_save:
+            aec, bac = _get_energy_corrections(lot, 'p')
+
+        # Verify both sections were searched independently
+        self.assertTrue(any('atom_energies' in c for c in calls))
+        self.assertTrue(any('pbac' in c for c in calls))
+        # Verify the script received separate keys
+        save_call = mock_save.call_args
+        saved_content = save_call[1].get('content') or save_call[0][1]
+        self.assertEqual(saved_content['aec_key'], aec_key)
+        self.assertEqual(saved_content['bac_key'], bac_key)
+        # Verify results returned
+        self.assertIsNotNone(aec)
+        self.assertIsNotNone(bac)
+
 
 class TestGetTsImagFreqFromFreqs(unittest.TestCase):
     """Tests for _get_ts_imag_freq using spc.freqs as primary source."""

arc/scripts/get_qm_corrections.py

@@ -44,24 +44,24 @@ def _lot_from_string(lot_str):
 def main(input_path, output_path):
     """Look up AEC and BAC for the given level of theory key."""
     params = read_yaml_file(input_path) or {}
-    matched_key = params.get('matched_key')
     bac_type = params.get('bac_type')
 
     result = {'aec': None, 'bac': None}
 
-    if not matched_key:
-        save_yaml_file(output_path, result)
-        return
+    # Support both old format (single matched_key) and new format (separate aec_key/bac_key)
+    aec_key = params.get('aec_key') or params.get('matched_key')
+    bac_key = params.get('bac_key') or params.get('matched_key')
 
-    lot = _lot_from_string(matched_key)
+    if aec_key:
+        lot = _lot_from_string(aec_key)
+        aec = atom_energies.get(lot)
+        if aec is not None:
+            result['aec'] = {str(k): float(v) for k, v in aec.items()}
 
-    aec = atom_energies.get(lot)
-    if aec is not None:
-        result['aec'] = {str(k): float(v) for k, v in aec.items()}
-
-    if bac_type in ('p', 'm'):
+    if bac_key and bac_type in ('p', 'm'):
+        bac_lot = _lot_from_string(bac_key)
         bac_dict = pbac if bac_type == 'p' else mbac
-        bac = bac_dict.get(lot)
+        bac = bac_dict.get(bac_lot)
         if bac is not None:
             result['bac'] = {str(k): float(v) for k, v in bac.items()}
 

arc/statmech/arkane.py

@@ -28,6 +28,15 @@
 RMG_ENV_NAME = settings.get('RMG_ENV_NAME', 'rmg_env')
 logger = get_logger()
 
+# Section boundary markers in the RMG quantum_corrections/data.py file.
+AEC_SECTION_START = "atom_energies = {"
+AEC_SECTION_END = "pbac = {"
+PBAC_SECTION_START = "pbac = {"
+PBAC_SECTION_END = "mbac = {"
+MBAC_SECTION_START = "mbac = {"
+MBAC_SECTION_END = "freq_dict ="
+FREQ_SECTION_START = "freq_dict = {"
+
 
 main_input_template = """#!/usr/bin/env python
 # -*- coding: utf-8 -*-
@@ -699,7 +708,7 @@ def _extract_section(file_path: str, section_start: str, section_end: Optional[s
 
 
 
-def _get_qm_corrections_files() -> List[str]:
+def get_qm_corrections_files() -> List[str]:
     """
     Return quantum corrections data.py paths from the RMG database.
     """
@@ -834,7 +843,7 @@ def _format_years(years: List[Optional[int]]) -> str:
     return ", ".join("none" if y is None else str(y) for y in years)
 
 
-def _find_best_across_files(level: "Level",
+def find_best_across_files(level: "Level",
                             qm_corr_files: List[str],
                             section_start: str,
                             section_end: Optional[str],
@@ -886,6 +895,10 @@ def _warn_no_match(level: "Level",
             f"available years: {_format_years(years)}. "
             f"Specify a year to select a matching entry."
         )
+    else:
+        logger.warning(
+            f"No Arkane {label} entry found for {level.simple()} in the RMG database."
+        )
 
 
 def _find_best_level_key_for_sp_level(level: "Level",
@@ -1015,33 +1028,28 @@ def get_arkane_model_chemistry(sp_level: 'Level',
     Returns:
         Optional[str]: Arkane-compatible model chemistry string.
     """
-    qm_corr_files = _get_qm_corrections_files()
-
-    aec_start = "atom_energies = {"
-    aec_end = "pbac = {"
-    freq_start = "freq_dict = {"
-    freq_end = None  # freq_dict is the last section in data.py; read to EOF
+    qm_corr_files = get_qm_corrections_files()
 
     # Composite methods and user-supplied freq_scale_factor both only need an AEC entry.
     if sp_level.method_type == 'composite' or freq_scale_factor is not None:
-        best_energy = _find_best_across_files(sp_level, qm_corr_files, aec_start, aec_end)
+        best_energy = find_best_across_files(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END)
         if best_energy is None:
-            _warn_no_match(sp_level, qm_corr_files, aec_start, aec_end, label="AEC")
+            _warn_no_match(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END, label="AEC")
             return None
         return best_energy
 
     # CompositeLevelOfTheory: need both energy (AEC) and frequency entries.
     if freq_level is None:
         raise ValueError("freq_level required when freq_scale_factor isn't provided")
 
-    best_energy = _find_best_across_files(sp_level, qm_corr_files, aec_start, aec_end)
-    best_freq = _find_best_across_files(freq_level, qm_corr_files, freq_start, freq_end)
+    best_energy = find_best_across_files(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END)
+    best_freq = find_best_across_files(freq_level, qm_corr_files, FREQ_SECTION_START, None)
 
     if best_energy is None or best_freq is None:
         if best_energy is None:
-            _warn_no_match(sp_level, qm_corr_files, aec_start, aec_end, label="AEC")
+            _warn_no_match(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END, label="AEC")
         if best_freq is None:
-            _warn_no_match(freq_level, qm_corr_files, freq_start, freq_end, label="frequency correction")
+            _warn_no_match(freq_level, qm_corr_files, FREQ_SECTION_START, None, label="frequency correction")
         return None
 
     return (
@@ -1052,6 +1060,33 @@ def get_arkane_model_chemistry(sp_level: 'Level',
     )
 
 
+def check_arkane_aec(sp_level: 'Level', raise_error: bool = False) -> bool:
+    """
+    Check that Arkane has AEC for the given sp level of theory (no BAC check).
+    Used when bac_type is None but we still want to verify AEC availability.
+
+    Args:
+        sp_level (Level): Level of theory for energy.
+        raise_error (bool): Whether to raise an error if AEC is missing.
+
+    Returns:
+        bool: True if AEC is available, False otherwise.
+    """
+    try:
+        qm_corr_files = get_qm_corrections_files()
+    except InputError as e:
+        logger.warning(f'Could not load Arkane quantum corrections data: {e}')
+        return False
+    best_aec_key = find_best_across_files(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END)
+    if best_aec_key is not None:
+        logger.info(f'Arkane atom energy corrections (AEC) matched for {best_aec_key} (BAC disabled)')
+    else:
+        _warn_no_match(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END, label="AEC")
+        if raise_error:
+            raise ValueError(f'Arkane has no atom energy corrections (AEC) for {_level_to_str(sp_level)}.')
+    return best_aec_key is not None
+
+
 def check_arkane_bacs(sp_level: 'Level',
                       bac_type: str = 'p',
                       raise_error: bool = False,
@@ -1071,19 +1106,19 @@ def check_arkane_bacs(sp_level: 'Level',
     Returns:
         bool: True if both AECs and BACs are available, False otherwise.
     """
-    qm_corr_files = _get_qm_corrections_files()
+    try:
+        qm_corr_files = get_qm_corrections_files()
+    except InputError as e:
+        logger.warning(f'Could not load Arkane quantum corrections data: {e}')
+        return False
 
-    aec_start = "atom_energies = {"
-    aec_end = "pbac = {"
     if bac_type.lower() == 'm':
-        bac_start = "mbac = {"
-        bac_end = "freq_dict ="
+        bac_start, bac_end = MBAC_SECTION_START, MBAC_SECTION_END
     else:
-        bac_start = "pbac = {"
-        bac_end = "mbac = {"
+        bac_start, bac_end = PBAC_SECTION_START, PBAC_SECTION_END
 
-    best_aec_key = _find_best_across_files(sp_level, qm_corr_files, aec_start, aec_end)
-    best_bac_key = _find_best_across_files(sp_level, qm_corr_files, bac_start, bac_end)
+    best_aec_key = find_best_across_files(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END)
+    best_bac_key = find_best_across_files(sp_level, qm_corr_files, bac_start, bac_end)
 
     has_aec = best_aec_key is not None
     has_bac = best_bac_key is not None
@@ -1092,7 +1127,7 @@ def check_arkane_bacs(sp_level: 'Level',
     if not has_encorr:
         repr_level = best_aec_key if best_aec_key is not None else _level_to_str(sp_level)
         year_note = ""
-        aec_years = _all_available_years(sp_level, qm_corr_files, aec_start, aec_end)
+        aec_years = _all_available_years(sp_level, qm_corr_files, AEC_SECTION_START, AEC_SECTION_END)
         bac_years = _all_available_years(sp_level, qm_corr_files, bac_start, bac_end)
         if sp_level.year is not None:
             year_note = (
@@ -1105,14 +1140,29 @@ def check_arkane_bacs(sp_level: 'Level',
                 f"available BAC years: {_format_years(bac_years)}. "
                 f"Specify a year to select a matching entry."
             )
-        mssg = (
-            f"Arkane does not have the required energy corrections for {repr_level} "
-            f"(AEC: {has_aec}, BAC: {has_bac}).{year_note}"
-        )
+        if has_aec and not has_bac:
+            mssg = (
+                f"Arkane atom energy corrections (AEC) matched for {repr_level}, "
+                f"but bond additivity corrections (BAC) were NOT found in the RMG database. "
+                f"Thermo/kinetics results will use AEC but lack BAC.{year_note}"
+            )
+        elif has_bac and not has_aec:
+            mssg = (
+                f"Arkane {bac_type.upper()}BAC matched for {best_bac_key}, "
+                f"but atom energy corrections (AEC) were NOT found in the RMG database. "
+                f"Energy corrections will be disabled.{year_note}"
+            )
+        else:
+            mssg = (
+                f"Arkane does not have atom energy corrections (AEC) or bond additivity corrections (BAC) "
+                f"for {repr_level}.{year_note}"
+            )
         if raise_error:
             raise ValueError(mssg)
         else:
             logger.warning(mssg)
+    else:
+        logger.info(f'Arkane energy corrections matched for {best_aec_key} (AEC and {bac_type.upper()}BAC)')
     return has_encorr