CafChem / app.py
cafierom's picture
Update app.py
68c78d9 verified
import gradio as gr
import shutil
from rdkit import Chem
from rdkit.Chem import AllChem
import pandas as pd
from ase import Atoms
from ase.optimize import BFGS
from ase.vibrations import Vibrations
from ase.thermochemistry import IdealGasThermo
from fairchem.core import FAIRChemCalculator, pretrained_mlip
predictor = pretrained_mlip.get_predict_unit("uma-s-1", device="cpu")
calculator = FAIRChemCalculator(predictor, task_name="omol")
model = "UMA-OMOL"
def calc_vibrations(smile: str, struct_type: str) -> str:
"""Calculate the vibrational frequencies of a molecule
using the META UMA MLIP, starting from a SMILES string.
Args:
smile: The SMILES string of the molecule
struct_type: The type of the molecule (linear or non-linear)
Returns:
The vibrational frequencies of the molecule in cm-1, the low
frequency motions (typically corresponding to translation or
rotation), and the number of imaginary frequencies in a text string
"""
mol = Chem.MolFromSmiles(smile)
molH = Chem.AddHs(mol)
AllChem.EmbedMolecule(molH)
AllChem.MMFFOptimizeMolecule(molH)
xyz_list = []
atoms_list = ""
for atom in molH.GetAtoms():
atoms_list += atom.GetSymbol()
pos = molH.GetConformer().GetAtomPosition(atom.GetIdx())
temp_tuple = (pos[0], pos[1], pos[2])
xyz_list.append(temp_tuple)
atoms = Atoms(atoms_list,xyz_list)
atoms.calc = calculator
energy = atoms.get_potential_energy()
dyn = BFGS(atoms)
dyn.run(fmax=0.05)
final_energy = atoms.get_potential_energy()
vib = Vibrations(atoms)
vib.run()
vib_energies = vib.get_energies()
real_vibs = []
number_imaginary = 0
for vib in vib_energies:
if vib.imag < 0.00001:
real_vibs.append(8065.56*vib.real.item())
else:
number_imaginary += 1
num_atoms = molH.GetNumAtoms()
if struct_type == "non-linear" or struct_type == "nonlinear":
non_vib_dof = 6 - number_imaginary
else:
non_vib_dof = 5 - number_imaginary
vib_dof = 3*num_atoms - non_vib_dof
out_text = ""
for i,vib in enumerate(real_vibs[3*num_atoms - vib_dof:]):
out_text += f"Vibrational frequency {i+1}: {vib:.3f} cm-1\n"
out_text += f"Also calculated the following low frequency motions:\n"
for i,vib in enumerate(real_vibs[:3*num_atoms - vib_dof]):
out_text += f"Frequency {i+1}: {vib:.3f} cm-1\n"
out_text += f"Also found the following number of imaginary frequencies: {number_imaginary}"
shutil.rmtree('vib')
return out_text
gradio_app = gr.Interface(
calc_vibrations,
inputs=[gr.Textbox(label="SMILES to calculate"),gr.Radio(label="Type of molecule", choices=["linear", "non-linear"])],
outputs=[gr.Textbox(label="Frequencies: ")],
title="Get Vibrational frequencies of a molecule from SMILES using the META UMA MLIP",
)
if __name__ == "__main__":
gradio_app.launch(mcp_server=True)