added tests for calc.py

src/calcs.py

@@ -1,47 +1,68 @@
 import os
+import yaml
 import torch
 from mace.calculators import MACECalculator
 from newtonnet.utils.ase_interface import MLAseCalculator
-from ase import Atoms
+from typing import Any, Dict
 
 
-def calc_mace():
+def load_config(config_path: str = 'tests/config.yml') -> Dict[str, Any]:
+    """
+    Load the configuration file.
 
-    mlcalculator = MACECalculator(
-        model_paths='/global/home/users/kumaranu/Documents/gpu_jobs/MACE_model.model',
-        #model_paths='/global/home/users/kumaranu/Documents/mep_tests/src/MACE_model_cpu.model',
-        device='cuda',
-        #device='cpu',
+    Parameters
+    ----------
+    config_path : str, optional
+        Path to the configuration file, by default 'config.yml'
+
+    Returns
+    -------
+    Dict[str, Any]
+        Configuration dictionary.
+    """
+    with open(config_path, 'r') as file:
+        config = yaml.safe_load(file)
+    return config
+
+
+def calc_mace() -> MACECalculator:
+    """
+    Initialize and return a MACECalculator instance.
+
+    Returns
+    -------
+    MACECalculator
+        An instance of the MACECalculator initialized with the specified model path and device.
+    """
+    config = load_config()
+    model_path = config['mace']['model_path']
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    ml_calculator = MACECalculator(
+        model_paths=model_path,
+        device=device,
         default_dtype="float64",
     )
-    return mlcalculator
-
-
-def calc():
-    use_cuda = torch.cuda.is_available()
-    device = 'cuda' if use_cuda else 'cpu'
-    mlcalculator = MLAseCalculator(
-        model_path=[
-                    os.path.join(
-                        os.path.expanduser("~"),
-                        'Documents/NewtonNet/example/predict/training_52/models/best_model_state.tar'
-                    ),
-                    #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_53/models/best_model_state.tar',
-                    #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_54/models/best_model_state.tar',
-                    #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_55/models/best_model_state.tar',
-                   ],    # path to model file, str or list of str
-        settings_path=[
-                       os.path.join(
-                           os.path.expanduser("~"),
-                           'Documents/NewtonNet/example/predict/training_52/run_scripts/config0.yml'
-                       ),
-                       #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_53/run_scripts/config2.yml',
-                       #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_54/run_scripts/config1.yml',
-                       #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_55/run_scripts/config3.yml',
-                      ],    # path to configuration file, str or list of str
-        hess_method=None,    # method to calculate hessians. 'autograd', 'fwd_diff', 'cnt_diff', or None (default: 'autograd')
-        # hess_precision=1e-5,    # hessian gradient calculation precision for 'fwd_diff' and 'cnt_diff', ignored otherwise (default: None)
-        disagreement='std',    # method to calculate disagreement among models. 'std', 'std_outlierremoval', 'range':, 'values', or None (default: 'std')
+    return ml_calculator
+
+
+def calc() -> MLAseCalculator:
+    """
+    Initialize and return an MLAseCalculator instance.
+
+    Returns
+    -------
+    MLAseCalculator
+        An instance of the MLAseCalculator initialized with the specified model path, config path, and device.
+    """
+    config = load_config()
+    model_path = config['newtonnet']['model_path']
+    config_path = config['newtonnet']['config_path']
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    ml_calculator = MLAseCalculator(
+        model_path=[model_path],
+        settings_path=[config_path],
+        disagreement='std',
         device=device
     )
-    return mlcalculator
+    return ml_calculator

src/setup_images.py

@@ -1,41 +1,71 @@
 import os
 from calcs import calc
 from ase.io import read, write
-from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union
 from sella_wrapper import sella_wrapper
 
 
 def setup_images(
-                 logdir=None,
-                 xyz_r_p=None,
-                 xyz_ts=None,
-                 N_intermediate=40,
-                ):
-    reactant = read(xyz_r_p, index='0')
-    reactant.calc = calc()
-    sella_wrapper(reactant, traj_file=logdir + '/reactant_opt.traj', sella_order=0)
-
-    product = read(xyz_r_p, index='1')
-    product.calc = calc()
-    sella_wrapper(product, traj_file=logdir + '/product_opt.traj', sella_order=0)
-
-    os.makedirs(logdir, exist_ok=True)
-    os.chdir(logdir)
-    write("r_p.xyz", [reactant.copy(), product.copy()])
-    os.system('geodesic_interpolate r_p.xyz --output output.xyz --nimages ' + str(N_intermediate))
-    images = read('output.xyz', index=':')
-
-    for image in images:
-        image.calc = calc()
-        mlcalculator = calc()
-        mlcalculator.calculate(image)
+        logdir: str,
+        xyz_r_p: str,
+        n_intermediate: int = 40,
+):
+    """
+    Sets up intermediate images for NEB calculations between reactant and product states.
+
+    Parameters:
+    logdir (str): Directory to save the intermediate files.
+    xyz_r_p (str): Path to the XYZ file containing reactant and product structures.
+    n_intermediate (int): Number of intermediate images to generate.
+
+    Returns:
+    List: List of ASE Atoms objects with calculated energies and forces.
+    """
+    try:
+        # Ensure the log directory exists
+        os.makedirs(logdir, exist_ok=True)
+
+        # Read reactant and product structures
+        reactant = read(xyz_r_p, index='0')
+        product = read(xyz_r_p, index='1')
+
+        # Optimize reactant and product structures using sella
+        for atom, name in zip([reactant, product], ['reactant', 'product']):
+            atom.calc = calc()
+            traj_file = os.path.join(logdir, f'{name}_opt.traj')
+            sella_wrapper(atom, traj_file=traj_file, sella_order=0)
+
+        # Save optimized reactant and product structures
+        r_p_path = os.path.join(logdir, "r_p.xyz")
+        write(r_p_path, [reactant.copy(), product.copy()])
+
+        # Generate intermediate images using geodesic interpolation
+        output_xyz = os.path.join(logdir, 'output.xyz')
+        os.system(f'geodesic_interpolate {r_p_path} --output {output_xyz} --nimages {n_intermediate}')
+        images = read(output_xyz, index=':')
+
+        # Calculate energies and forces for each intermediate image
+        for image in images:
+            image.calc = calc()
+            ml_calculator = calc()
+            ml_calculator.calculate(image)
 
-        energy = mlcalculator.results['energy']
-        forces = mlcalculator.results['forces']
+            energy = ml_calculator.results['energy']
+            forces = ml_calculator.results['forces']
 
-        image.info['energy'] = energy
-        #image.arrays['forces'] = forces
+            image.info['energy'] = energy
+            image.arrays['forces'] = forces
 
-    write('geodesic_path.xyz', images)
+        # Save the geodesic path
+        geodesic_path = os.path.join(logdir, 'geodesic_path.xyz')
+        write(geodesic_path, images)
+
+        return images
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return []
+
+
+if __name__ == '__main__':
 
-    return images
+    setup_images()

tests/config.yml

@@ -0,0 +1,5 @@
+mace:
+  model_path: 'MACE_model.model'
+newtonnet:
+  model_path: 'tests/best_model_state.tar'
+  config_path: 'tests/config0.yml'