diff --git a/src/atomate2/openmm/utils.py b/src/atomate2/openmm/utils.py index 5939a51742..25aa0759fd 100644 --- a/src/atomate2/openmm/utils.py +++ b/src/atomate2/openmm/utils.py @@ -4,6 +4,7 @@ import io import re +import shutil import tempfile import time import warnings @@ -22,15 +23,18 @@ def download_opls_xml( - names_smiles: dict[str, str], output_dir: str | Path, overwrite_files: bool = False + names_params: dict[str, dict[str, str]], + output_dir: str | Path, + overwrite_files: bool = False, ) -> None: """Download an OPLS-AA/M XML file from the LigParGen website using Selenium.""" try: from selenium import webdriver from selenium.webdriver.chrome.service import Service as ChromeService from selenium.webdriver.common.by import By + from selenium.webdriver.support import expected_conditions as ec + from selenium.webdriver.support.ui import WebDriverWait from webdriver_manager.chrome import ChromeDriverManager - except ImportError: warnings.warn( "The `selenium` package is not installed. " @@ -41,8 +45,12 @@ def download_opls_xml( # Initialize the Chrome driver driver = webdriver.Chrome(service=ChromeService(ChromeDriverManager().install())) - for name, smiles in names_smiles.items(): + for name, params in names_params.items(): final_file = Path(output_dir) / f"{name}.xml" + smiles = params.get("smiles") + checkopt = params.get("checkopt", 3) + charge = params.get("charge", 0) + if final_file.exists() and not overwrite_files: continue try: @@ -62,9 +70,23 @@ def download_opls_xml( driver.get("https://zarbi.chem.yale.edu/ligpargen/") # Find the SMILES input box and enter the SMILES code - smiles_input = driver.find_element(By.ID, "smiles") + smiles_input = WebDriverWait(driver, 10).until( + ec.presence_of_element_located((By.ID, "smiles")) + ) smiles_input.send_keys(smiles) + # Find Molecule Optimization Iterations dropdown menu and select + checkopt_input = WebDriverWait(driver, 10).until( + ec.presence_of_element_located((By.NAME, "checkopt")) + ) + checkopt_input.send_keys(checkopt) + + # Find Charge dropdown menu and select + charge_input = WebDriverWait(driver, 10).until( + ec.presence_of_element_located((By.NAME, "dropcharge")) + ) + charge_input.send_keys(charge) + # Find and click the "Submit Molecule" button submit_button = driver.find_element( By.XPATH, @@ -73,7 +95,9 @@ def download_opls_xml( submit_button.click() # Wait for the second page to load - # time.sleep(2) # Adjust this delay as needed based on the loading time + time.sleep( + 2 + 0.5 * int(checkopt) + ) # Adjust based on loading time & optimization iterations # Find and click the "XML" button under Downloads and OpenMM xml_button = driver.find_element( @@ -86,7 +110,7 @@ def download_opls_xml( file = next(Path(tmpdir).iterdir()) # copy downloaded file to output_file using os - Path(file).rename(final_file) + shutil.copy(file, final_file) except Exception as e: # noqa: BLE001 warnings.warn(