testorb.py

from ase.build import molecule
from ase.optimize import FIRE
from ase.vibrations import Vibrations
import time


def run_vibrational_analysis(atoms, calculator, calc_name, box_size=20.0, fmax=0.01, steps=500):
    """
    Function to perform optimization, vibrational analysis, and print the summary.
    
    Parameters:
    atoms - The ASE Atoms object (e.g., a molecule)
    calculator - The calculator to be used for optimization and vibrations
    calc_name - Name of the calculator (used in the output files)
    box_size - Size of the simulation box (default 20.0 Å)
    fmax - Maximum force for optimization (default 0.01 eV/Å)
    steps - Maximum steps for optimization (default 500)
    """
    
    # Get timestamp for file names
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    
    # Center molecule inside the box and set cell dimensions
    #atoms = molecule('H2O')
    atoms.center(vacuum=box_size / 2)
    atoms.set_cell([box_size, box_size, box_size])
    atoms.set_pbc([True, True, True])  # No periodic boundary conditions
    
    # Step 2: Set the calculator
    atoms.set_calculator(calculator)
    
    # Step 3: Perform optimization
    opt = FIRE(atoms)
    opt.run(fmax=fmax, steps=steps)
    
    # Step 4: Perform vibrational analysis
    vib = Vibrations(atoms)
    vib.run()
    
    # Step 5: Save vibrational analysis summary to a log file
    summary_file = f'{timestamp}_{calc_name}_vib_summary.txt'
    vib.summary(log=summary_file)
    vib.clean()

    # Print vibrational summary
    with open(summary_file, 'r') as file:
        print(f"Vibrational analysis summary for {calc_name}:")
        print(file.read())
        

# Example usage with ORBCalculator and SevenNetCalculator
if __name__ == "__main__":
    from orb_models.forcefield import pretrained
    from orb_models.forcefield.calculator import ORBCalculator
    from sevenn.sevennet_calculator import SevenNetCalculator

    # Prepare ORB Forcefield for calculation
    orbff = pretrained.orb_d3_v1(device='cuda')
    orbcalc = ORBCalculator(orbff, device="cuda")
    atoms = molecule('H2O')
    # Run vibrational analysis for ORBCalculator
    run_vibrational_analysis(atoms, orbcalc, 'ORBCalculator')

    # Prepare SevenNetCalculator for calculation
    scalc = SevenNetCalculator(device="cuda")

    run_vibrational_analysis(atoms, scalc, 'SevenNetCalculator')