Add UMA as a calculator in the ASE job adapter

Build UMA (Meta FAIR fairchem-core, uma-s-1p1, task omol) on top of the generic
ASE adapter (PR #836) instead of a standalone adapter:

- ase_script.py: add a uma/fairchem branch to get_calculator; set total charge and
  spin (=multiplicity) on atoms.info (omol conditioning); use Sella order=1 for TS
  saddle-point searches when is_ts; add an irc job type via Sella IRC.
- ase.py: derive the calculator from the level method (so method='uma' works with no
  args), resolve UMA defaults (latest model, omol, cpu) via determine_settings, pass
  is_ts and irc_direction to the script, and warn on a UMA single point for an isolated
  atom or triplet O2 (unreliable absolute energy).
- settings.py: UMA_PYTHON=find_executable('uma_env'), ASE_CALCULATORS_ENV['uma'], and
  UMA_LATEST_MODEL.
- level.py: route method 'uma'/'uma-s-1'/'uma-s-1p1' to the 'ase' software.
- yaml.py: implement parse_irc_traj and parse_1d_scan_coords so UMA IRC/scan outputs
  round-trip.

Rotor scans run through ARC's directed_scan (constrained opt), already supported by the
ASE adapter. fairchem/Sella-IRC API points only confirmable inside uma_env are marked
with # VERIFY. Adds env-independent unit tests (routing, calculator/settings resolution,
input writing, sp warning, output round-trip) plus skip-guarded model tests.

Author: alongd

arc/job/adapters/ase.py

@@ -12,7 +12,7 @@
 from arc.job.adapters.common import _initialize_adapter
 from arc.job.factory import register_job_adapter
 from arc.imports import settings
-from arc.settings.settings import ARC_PYTHON, find_executable
+from arc.settings.settings import ARC_PYTHON, UMA_LATEST_MODEL, find_executable
 
 if TYPE_CHECKING:
     from arc.level import Level
@@ -25,8 +25,12 @@
 DEFAULT_ASE_ENV = {
     'torchani': 'TANI_PYTHON',
     'xtb': 'XTB_PYTHON',
+    'uma': 'UMA_PYTHON',
 }
 
+# Level methods that select the UMA calculator. 'uma' resolves to the latest model.
+UMA_METHODS = ('uma', 'uma-s-1', 'uma-s-1p1')
+
 class ASEAdapter(JobAdapter):
     """
     A generic adapter for ASE (Atomic Simulation Environment) jobs.
@@ -77,12 +81,13 @@ def __init__(self,
         self.job_adapter = 'ase'
         self.execution_type = execution_type or 'incore'
         self.incore_capacity = 100
-        
+
         self.sp = None
         self.opt_xyz = None
         self.freqs = None
 
         self.args = args or dict()
+        self.level = level  # also set by _initialize_adapter; needed early by get_python_executable
         self.python_executable = self.get_python_executable()
         self.script_path = os.path.join(os.path.dirname(__file__), 'scripts', 'ase_script.py')
 
@@ -128,11 +133,46 @@ def __init__(self,
                             xyz=xyz,
                             )
 
+    def determine_calculator_name(self) -> str:
+        """
+        Determine the ASE calculator name, from ``args['keyword']['calculator']`` if given,
+        otherwise inferred from the level method (e.g., a 'uma' method selects the UMA calculator).
+
+        Returns:
+            str: The lowercased calculator name (empty string if undetermined).
+        """
+        calc = (self.args or dict()).get('keyword', dict()).get('calculator', '')
+        if not calc and self.level is not None and getattr(self.level, 'method', None) \
+                and self.level.method.lower() in UMA_METHODS:
+            calc = 'uma'
+        return calc.lower()
+
+    def determine_settings(self) -> dict:
+        """
+        Build the ``settings`` block passed to ase_script.py: the user's ``args['keyword']`` plus
+        a resolved ``calculator`` and, for UMA, default ``model`` (the level method, with 'uma'
+        resolving to the latest model), ``task``, and ``device``.
+
+        Returns:
+            dict: The resolved ASE run settings.
+        """
+        settings_dict = dict((self.args or dict()).get('keyword', dict()))
+        calc = self.determine_calculator_name()
+        if calc:
+            settings_dict.setdefault('calculator', calc)
+        if calc == 'uma':
+            if 'model' not in settings_dict:
+                method = self.level.method.lower() if self.level is not None and self.level.method else 'uma'
+                settings_dict['model'] = UMA_LATEST_MODEL if method == 'uma' else method
+            settings_dict.setdefault('task', 'omol')
+            settings_dict.setdefault('device', 'cpu')
+        return settings_dict
+
     def get_python_executable(self) -> str:
         """
         Identify the correct Python executable based on the calculator.
         """
-        calc = self.args.get('keyword', {}).get('calculator', '').lower()
+        calc = self.determine_calculator_name()
         env_mapping = settings.get('ASE_CALCULATORS_ENV', DEFAULT_ASE_ENV)
         env_var_name = env_mapping.get(calc)
 
@@ -157,15 +197,35 @@ def write_input_file(self) -> None:
             'xyz': self.xyz,
             'charge': self.charge,
             'multiplicity': self.multiplicity,
+            'is_ts': self.species[0].is_ts if self.species else False,
             'constraints': self.constraints,
-            'settings': self.args.get('keyword', {}),
+            'irc_direction': self.irc_direction,
+            'settings': self.determine_settings(),
         }
         save_yaml_file(os.path.join(self.local_path, 'input.yml'), input_dict)
 
+    def warn_if_unreliable_uma_sp(self) -> None:
+        """
+        Warn if this is a UMA single point on a species whose absolute UMA energy is unreliable
+        (an isolated atom or triplet O2). UMA's geometries/frequencies are fine; only the absolute
+        energy of these under-represented species is off, so a DFT single point is preferable.
+        """
+        if self.job_type not in ['sp', 'conf_sp'] or self.determine_calculator_name() != 'uma':
+            return
+        symbols = self.xyz['symbols'] if self.xyz is not None else tuple()
+        is_atom = len(symbols) == 1
+        is_triplet_o2 = len(symbols) == 2 and all(s == 'O' for s in symbols) and self.multiplicity == 3
+        if is_atom or is_triplet_o2:
+            label = self.species[0].label if self.species else 'species'
+            logger.warning(f'Computing a UMA single point for {label} (an isolated atom or triplet O2). '
+                           f'UMA absolute energies are unreliable for these under-represented species; '
+                           f'consider using a DFT single point instead.')
+
     def execute_incore(self) -> None:
         """
         Execute the job incore.
         """
+        self.warn_if_unreliable_uma_sp()
         self.write_input_file()
         cmd = [self.python_executable, self.script_path, '--yml_path', self.local_path]
         process = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

arc/job/adapters/scripts/ase_script.py

@@ -88,7 +88,18 @@ def get_calculator(calc_config: dict, charge: int = 0, multiplicity: int = 1):
     elif name == 'mopac':
         from ase.calculators.mopac import MOPAC
         return MOPAC(**kwargs)
-        
+
+    elif name in ('uma', 'fairchem'):
+        # UMA (Meta FAIR fairchem-core). Total charge and spin (= multiplicity) are conditioned on
+        # the ase.Atoms via atoms.info in main(); they are not calculator kwargs.
+        from fairchem.core import FAIRChemCalculator, pretrained_mlip
+        model = calc_config.get('model', 'uma-s-1p1')
+        device = calc_config.get('device', 'cpu')
+        task = calc_config.get('task', 'omol')
+        predictor = pretrained_mlip.get_predict_unit(model, device=device)
+        return FAIRChemCalculator(predictor, task_name=task)
+
+
     from ase.calculators.calculator import get_calculator_class
     try:
         calc_class = get_calculator_class(name)
@@ -279,8 +290,10 @@ def main():
     settings = input_dict.get('settings', {})
     charge = input_dict.get('charge', 0)
     multiplicity = input_dict.get('multiplicity', 1)
-    
+    is_ts = input_dict.get('is_ts', False)
+
     atoms = Atoms(symbols=xyz['symbols'], positions=xyz['coords'])
+    atoms.info.update({'charge': charge, 'spin': multiplicity})  # UMA (omol) conditions on these
     calc = get_calculator(settings, charge, multiplicity)
     atoms.calc = calc
 
@@ -308,13 +321,14 @@ def save_current_geometry(out_dict, atoms_obj, input_xyz):
             'scipyfminbfgs': SciPyFminBFGS, 'scipyfmincg': SciPyFminCG,
             'sella': None,
         }
-        if engine_name == 'sella':
+        logfile = os.path.join(os.path.dirname(input_path), 'opt.log')
+        if is_ts or engine_name == 'sella':
+            # A TS search needs a saddle-point optimizer; UMA ships none, so use Sella.
             from sella import Sella
-            opt_class = Sella
+            opt = Sella(atoms, order=1 if is_ts else 0, logfile=logfile)
         else:
-            opt_class = engine_dict.get(engine_name, BFGS)
-        opt = opt_class(atoms, logfile=os.path.join(os.path.dirname(input_path), 'opt.log'))
-        
+            opt = engine_dict.get(engine_name, BFGS)(atoms, logfile=logfile)
+
         try:
             opt.run(fmax=fmax, steps=steps)
             save_current_geometry(output, atoms, xyz)
@@ -326,6 +340,26 @@ def save_current_geometry(out_dict, atoms_obj, input_xyz):
         # For non-optimization jobs, still save the geometry
         save_current_geometry(output, atoms, xyz)
 
+    if job_type == 'irc':
+        from sella import IRC  # VERIFY the Sella IRC API in the installed sella
+        from ase.io import read
+        fmax = float(settings.get('fmax', 0.001))
+        steps = int(settings.get('steps', 1000))
+        direction = input_dict.get('irc_direction', 'forward')
+        traj_path = os.path.join(os.path.dirname(input_path), 'irc.traj')
+        try:
+            irc = IRC(atoms, logfile=os.path.join(os.path.dirname(input_path), 'irc.log'),
+                      trajectory=traj_path)
+            irc.run(fmax=fmax, steps=steps, direction=direction)
+            images = read(traj_path, index=':')
+            output['irc_traj'] = [
+                {'coords': tuple(map(tuple, image.get_positions().tolist())),
+                 'symbols': xyz['symbols'],
+                 'isotopes': xyz.get('isotopes') or tuple([None] * len(xyz['symbols']))}
+                for image in images]
+        except Exception as exc:
+            output['error'] = f"IRC failed: {exc}"
+
     if job_type in ['freq', 'optfreq']:
         try:
             freq_results = run_vibrational_analysis(atoms, settings)