adds elstruct parameters/readers/writers for ase calls to psi4 and

autorun ase runners for single point energies
elstruct now has MLIP models allowed in the par and
reader/writers for energy jobs with them in ASE. autorun can run these
for MACE, more to come

Author: snelliott

src/autorun/__init__.py

@@ -7,6 +7,8 @@
 from autorun._run import from_input_string
 from autorun._run import run_script
 from autorun._run import write_input
+from autorun._run import dump_input
+from autorun._run import dump_output
 from autorun._run import read_output
 from autorun._host import host_node
 from autorun._host import process_id
@@ -23,6 +25,7 @@
 from autorun import projrot
 from autorun import thermp
 from autorun import varecof
+from autorun import ase
 
 # MultiProgram Runners
 from autorun._multiprog import projected_frequencies
@@ -35,6 +38,8 @@
     'run_script',
     'from_input_string',
     'write_input',
+    'dump_input',
+    'dump_output',
     'read_output',
     'host_node',
     'process_id',
@@ -50,6 +55,7 @@
     'projrot',
     'thermp',
     'varecof',
+    'ase',
     # MultiProgram Runners
     'projected_frequencies',
     'thermo'

src/autorun/_run.py

@@ -5,6 +5,7 @@
 import subprocess
 import warnings
 import stat
+import yaml
 
 
 SCRIPT_NAME = 'run.sh'
@@ -120,6 +121,90 @@ def write_input(run_dir, input_str,
                         aux_obj.write(fstring)
 
 
+def _simplify_data(obj):
+    """Convert numpy arrays, numpy scalars, and tuples to simple Python lists and values.
+    
+    Args:
+        obj: The object to simplify
+    Returns:
+        Simplified version of the object with standard Python types
+    """
+    import numpy as np
+    
+    if isinstance(obj, dict):
+        return {k: _simplify_data(v) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [_simplify_data(x) for x in obj]
+    elif isinstance(obj, np.ndarray):
+        return _simplify_data(obj.tolist())
+    elif isinstance(obj, np.generic):
+        return obj.item()
+    return obj
+
+
+def dump_input(run_dir, input_dct,
+                aux_dct=None,
+                input_name=INPUT_NAME):
+    """ Write input from a dictionary to a file in YAML format
+
+    Args:
+        run_dir: Directory where input file will be written
+        input_dict: Dictionary containing input data to be written
+        aux_dct: Dictionary of auxiliary files to write {filename: content}
+        input_name: Name of the main input file
+    """
+    if not os.path.exists(run_dir):
+        os.makedirs(run_dir)
+
+    # Simplify the data structures before dumping
+    simplified_dct = _simplify_data(input_dct)
+
+    with EnterDirectory(run_dir):
+        # Write the main input file as YAML
+        with open(input_name, mode='w', encoding='utf-8') as input_obj:
+            input_obj.write('__yaml__\n') 
+            yaml.dump(
+                simplified_dct,
+                input_obj,
+                default_flow_style=False,
+                sort_keys=False,
+                allow_unicode=True)
+
+        # Write all auxiliary input files
+        if aux_dct is not None:
+            for fname, fstring in aux_dct.items():
+                if fstring:
+                    with open(fname, mode='w', encoding='utf-8') as aux_obj:
+                        aux_obj.write(fstring)
+
+
+
+def dump_output(run_dir, output_dct,
+                output_name=OUTPUT_NAME):
+    """ Write output from a dictionary to a file in YAML format
+    Args:
+        run_dir: Directory where output file will be written
+        output_dct: Dictionary containing output data to be written
+        output_name: Name of the main output file
+    """
+    if not os.path.exists(run_dir):
+        os.makedirs(run_dir)
+
+    # Simplify the data structures before dumping
+    simplified_dct = _simplify_data(output_dct)
+    
+    with EnterDirectory(run_dir):
+        # Write the main output file as YAML
+        with open(output_name, mode='w', encoding='utf-8') as output_obj:
+            output_obj.write('__yaml__\n') 
+            yaml.dump(
+                simplified_dct,
+                output_obj, 
+                default_flow_style=False, 
+                sort_keys=False,
+                allow_unicode=True)
+
+            
 def read_output(run_dir, output_names=(OUTPUT_NAME,)):
     """ Read the output string from the run directory
     """

src/autorun/ase.py

@@ -0,0 +1,152 @@
+"""Functions to run ASE calculators."""
+
+from ase import Atoms
+from ase import units
+from ase.optimize import BFGS
+# from sella import Sella
+
+def get_calculator(calculator, family=None):
+    """Get a new instance of the appropriate ASE calculator.
+    
+    Args:
+        program: The program name ('psi4', 'mace_mp', etc.)
+    Returns:
+        The calculator class
+    Raises:
+        ValueError: If the program is not supported
+    """
+    # Lazy import the appropriate calculator
+    if calculator == 'psi4':
+        import psi4
+        from ase.calculators.psi4 import Psi4
+        return Psi4
+    elif calculator == 'mace':
+        if family == 'mace_mp':
+            from mace.calculators import mace_mp
+            return mace_mp
+        elif family == 'mace_off':
+            from mace.calculators import mace_off
+            return mace_off
+        elif family == 'mace_anicc':
+            from mace.calculators import mace_anicc
+            return mace_anicc
+        else:
+            raise ValueError(f"Unsupported MACE family: {family}")
+    # elif calculator == 'uma':
+    # elif calculator == 'nwx':
+    #     from ase.calculators.nwchem import NWChem
+    #     return NWChem
+
+
+def from_calc_dictionary(input_dct, script_str):
+    """ Run an ASE calculation from an input dictionary and return results
+
+    Args:
+        input_dct: Dictionary containing:
+            - calculator: Name of ASE calculator to use ('psi4' or 'nwx')
+            - atom_parameters:
+                - symbols: List of atomic symbols
+                - positions: List of [x,y,z] coordinates
+            - basis: Basis set name
+            - method: Quantum chemistry method
+            - charge: Molecular charge
+            - multiplicity: Spin multiplicity
+            - reference: Reference type ('rhf', 'uhf', etc.)
+        script_str: String of script that contains ase_<calculator>
+    Returns:
+        dict: Results containing:
+            - energy: Total energy in eV
+            - forces: Forces on atoms in eV/Å
+            - dipole: Dipole moment in e⋅Å
+            - charges: Atomic charges
+            - parameters: Input parameters used
+            - version: Version string of the program used
+    """
+    # Set up Atoms object
+    atom_parameters = input_dct.get('atom_parameters', {})
+    job = input_dct.get('job', 'energy')
+    atoms = Atoms(
+        symbols=atom_parameters.get('symbols', []),
+        positions=atom_parameters.get('positions', []),
+    )
+    # Set up calculator
+    calculator = None
+    if len(script_str.split('ase_')) > 1:
+        calculator = script_str.split('ase_')[-1].strip()
+        if 'method' in input_dct.keys():
+            kwargs = {
+                'atoms': atoms,
+                'basis': input_dct.get('basis'),
+                'method': input_dct.get('method'),
+                'charge': input_dct.get('charge', 0),
+                'multiplicity': input_dct.get('multiplicity', 1),
+                'reference': input_dct.get('reference')
+            }
+            family = None
+        elif 'family' in input_dct.keys():
+            family = input_dct.get('family')
+            if 'mlip' in input_dct.keys():
+                kwargs = {
+                    'model': input_dct.get('mlip'),
+                    'device': 'cpu' 
+                }
+            else:
+                kwargs = {}
+    calc = get_calculator(calculator, family)(**kwargs)
+    
+    # Attach calculator to atoms and run
+    atoms.calc = calc
+    
+    # Run calculation and get basic properties that all calculators support
+    if job == 'energy':
+        energy = atoms.get_potential_energy() / units.Hartree
+        positions = atoms.get_positions()
+    elif job == 'optimize':
+        if not input_dct.get('saddle', False):
+            dyn = BFGS(atoms)
+        # else:
+        #     dyn = Sella(atoms)
+        dyn.run(
+            fmax=input_dct.get('gconv', 0.05),
+            steps=input_dct.get('maxcyc', 100),
+            opt_type=input_dct.get('opt_type', 'standard')
+            ) 
+        energy = atoms.get_potential_energy()
+        positions = atoms.get_positions()
+
+    # Initialize results with guaranteed properties
+    results = {
+        'energy': energy,
+        'positions': positions,
+        # 'forces': forces.tolist(),
+        'parameters': dict(calc.parameters)  # Convert to regular dict for JSON serialization
+    }
+
+    # Add calculator-specific properties based on calculator type
+    if calculator == 'psi4':
+        results['version'] = f'ase-psi4'
+        results['normal_termination'] = True
+        # Psi4 supports both dipole moments and charges
+        # if 'properties' not in calc.parameters or 'dipole' in calc.parameters.get('properties', []):
+        #     results['dipole'] = atoms.get_dipole_moment().tolist()
+        # if 'properties' not in calc.parameters or 'mulliken' in calc.parameters.get('properties', []):
+        #     results['charges'] = atoms.get_charges().tolist()
+    elif 'mace' in calculator:
+        # MACE specific properties
+        results['version'] = f'ase-mace'
+        results['normal_termination'] = True
+        # MACE does not provide dipole moments or charges
+        # if 'properties' not in calc.parameters or 'dipole' in calc.parameters.get('properties', []):
+        #     results['dipole'] = atoms.get_dipole_moment().tolist()
+        # if 'properties' not in calc.parameters or 'mulliken' in calc.parameters.get('properties', []):
+        #     results['charges'] = atoms.get_charges().tolist()
+    elif calculator == 'nwx':
+        # NWChem specific properties
+        pass
+        # if 'dipole_moment' in calc.get_implemented_properties():
+        #     results['dipole'] = atoms.get_dipole_moment().tolist()
+        # if 'charges' in calc.get_implemented_properties():
+        #     results['charges'] = atoms.get_charges().tolist()
+
+    return results
+    

src/elstruct/__init__.py

@@ -12,13 +12,17 @@
 from elstruct.par import Program
 from elstruct.par import Reference
 from elstruct.par import Method
+from elstruct.par import Model
 from elstruct.par import Basis
 from elstruct.par import programs
+from elstruct.par import program_models
 from elstruct.par import program_methods
 from elstruct.par import program_dft_methods
 from elstruct.par import program_nondft_methods
 from elstruct.par import program_method_orbital_types
 from elstruct.par import program_bases
+from elstruct.par import method_is_mlip
+from elstruct.par import mlip_from_method
 
 
 __all__ = [
@@ -34,11 +38,15 @@
     'Program',
     'Reference',
     'Method',
+    'Model',
     'Basis',
     'programs',
+    'program_models',
     'program_methods',
     'program_dft_methods',
     'program_nondft_methods',
     'program_method_orbital_types',
     'program_bases',
+    'method_is_mlip',
+    'mlip_from_method'
 ]