import streamlit as st import os import tempfile import torch import numpy as np from ase import Atoms from ase.io import read, write from ase.optimize import BFGS, LBFGS, FIRE from ase.constraints import FixAtoms from ase.filters import FrechetCellFilter from ase.visualize import view import py3Dmol from mace.calculators import mace_mp from fairchem.core import pretrained_mlip, FAIRChemCalculator from orb_models.forcefield import pretrained from orb_models.forcefield.calculator import ORBCalculator from sevenn.calculator import SevenNetCalculator import pandas as pd import yaml # Added for FairChem reference energies import subprocess import sys import pkg_resources mattersim_available = False # try: # subprocess.check_call([sys.executable, "-m", "pip", "install", "mattersim"]) # except Exception as e: # print(f"Error during installation of mattersim: {e}") # try: # from mattersim.forcefield import MatterSimCalculator # mattersim_available = True # print("\n\n\n\n\n\n\nSuccessfully imported MatterSimCalculator.\n\n\n\n\n\n\n\n\n\n") # except ImportError as e: # print(f"Failed to import MatterSimCalculator: {e} \n\n\n\n\n\n\n\n") # mattersim_available = False # # Define version threshold # required_version = "2.0.0" # try: # installed_version = pkg_resources.get_distribution("numpy").version # if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version(required_version): # print(f"numpy version {installed_version} >= {required_version}. Installing numpy<2.0.0...") # subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy<2.0.0"]) # else: # print(f"numpy version {installed_version} is already < {required_version}. No action needed.") # except pkg_resources.DistributionNotFound: # print("numpy is not installed. Installing numpy<2.0.0...") # subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy<2.0.0"]) from huggingface_hub import login # try: # hf_token = st.secrets["HF_TOKEN"]["token"] # os.environ["HF_TOKEN"] = hf_token # login(token=hf_token) # except Exception as e: # print("streamlit hf secret not defined/assigned") try: hf_token = os.getenv("YOUR SECRET KEY") # Replace with your actual Hugging Face token or manage secrets appropriately if hf_token: login(token = hf_token) else: print("Hugging Face token not found. Some models might not be accessible.") except Exception as e: print(f"hf login error: {e}") os.environ["STREAMLIT_WATCHER_TYPE"] = "none" # YAML data for FairChem reference energies ELEMENT_REF_ENERGIES_YAML = """ oc20_elem_refs: - 0.0 - -0.16141512 - 0.03262098 - -0.04787699 - -0.06299825 - -0.14979306 - -0.11657468 - -0.10862579 - -0.10298174 - -0.03420248 - 0.02673997 - -0.03729558 - 0.00515243 - -0.07535697 - -0.13663351 - -0.12922852 - -0.11796547 - -0.07802946 - -0.00672682 - -0.04089589 - -0.00024177 - -1.74545186 - -1.54220241 - -1.0934019 - -1.16168372 - -1.23073475 - -0.78852824 - -0.71851599 - -0.52465053 - -0.02692092 - -0.00317922 - -0.06266862 - -0.10835274 - -0.12394474 - -0.11351727 - -0.07455817 - -0.00258354 - -0.04111325 - -0.02090265 - -1.89306078 - -1.30591887 - -0.63320009 - -0.26230344 - -0.2633669 - -0.5160055 - -0.95950798 - -1.45589361 - -0.0429969 - -0.00026949 - -0.05925609 - -0.09734631 - -0.12406852 - -0.11427538 - -0.07021442 - 0.01091345 - -0.05305289 - -0.02427209 - -0.19975668 - -1.71692859 - -1.53677781 - -3.89987009 - -10.70940462 - -6.71693816 - -0.28102249 - -8.86944824 - -7.95762687 - -7.13041437 - -6.64620014 - -5.11482482 - -4.42548227 - 0.00848295 - -0.06956227 - -2.6748853 - -2.21153293 - -1.67367741 - -1.07636151 - -0.79009981 - -0.16387243 - -0.18164401 - -0.04122529 - -0.00041833 - -0.05259382 - -0.0934314 - -0.11023834 - -0.10039175 - -0.06069209 - 0.01790437 - -0.04694024 - 0.00334084 - -0.06030621 - -0.58793619 - -1.27821808 - -4.97483577 - -5.66985655 - -8.43154622 - -11.15001317 - -12.95770812 - 0.0 - -14.47602729 - 0.0 odac_elem_refs: - 0.0 - -1.11737936 - -0.00011835 - -0.2941727 - -0.03868426 - -0.34862832 - -1.31552566 - -3.12457285 - -1.6052078 - -0.49653389 - -0.01137327 - -0.21957281 - -0.0008343 - -0.2750172 - -0.88417265 - -1.887378 - -0.94903558 - -0.31628167 - -0.02014536 - -0.15901053 - -0.00731884 - -1.96521355 - -1.89045209 - -2.53057428 - -5.43600675 - -5.09739336 - -3.03088746 - -1.23786562 - -0.40650749 - -0.2416017 - -0.01139188 - -0.26282496 - -0.82446455 - -1.70237206 - -0.84245376 - -0.28544892 - -0.02239991 - -0.14115912 - -0.02840799 - -2.09540994 - -1.85863996 - -1.12257399 - -4.32965355 - -3.30670045 - -1.19460755 - -1.26257601 - -1.46832888 - -0.19779414 - -0.0144274 - -0.23668767 - -0.70836953 - -1.43186113 - -0.71701186 - -0.24883129 - -0.01118184 - -0.13173447 - -0.0318395 - -0.41195547 - -1.23134873 - -2.03082996 - 0.1375954 - -5.45866275 - -7.59139905 - -5.99965965 - -8.43495767 - -2.6578407 - -7.77349787 - -5.30762201 - -5.15109657 - -4.41466995 - -0.02995219 - -0.2544495 - -3.23821202 - -3.45887214 - -4.53635003 - -4.60979468 - -2.90707964 - -1.28286153 - -0.57716664 - -0.18337108 - -0.01135944 - -0.22045398 - -0.66150479 - -1.32506342 - -0.66500178 - -0.22643927 - -0.00728197 - -0.11208472 - -0.00757856 - -0.21798637 - -0.91078787 - -1.78187161 - -3.89912261 - -3.94192659 - -7.59026042 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 omat_elem_refs: - 0.0 - -1.11700253 - 0.00079886 - -0.29731164 - -0.04129868 - -0.29106192 - -1.27751531 - -3.12342715 - -1.54797136 - -0.43969356 - -0.01250908 - -0.22855413 - -0.00943179 - -0.21707638 - -0.82619133 - -1.88667434 - -0.89093583 - -0.25816211 - -0.02414768 - -0.17662425 - -0.02568319 - -2.13001165 - -2.38688845 - -3.55934233 - -5.44700879 - -5.14749562 - -3.30662847 - -1.42167737 - -0.63181379 - -0.23449167 - -0.01146636 - -0.21291259 - -0.77939897 - -1.70148487 - -0.78386705 - -0.22690657 - -0.02245409 - -0.16092396 - -0.02798717 - -2.25685695 - -2.23690495 - -2.15347771 - -4.60251809 - -3.36416792 - -2.23062607 - -1.15550917 - -1.47553527 - -0.19918102 - -0.01475888 - -0.19767692 - -0.68005773 - -1.43073368 - -0.65790462 - -0.18915279 - -0.01179476 - -0.13507902 - -0.03056979 - -0.36017439 - -0.86279246 - -0.20573327 - -0.2734463 - -0.20046965 - -0.25444338 - -8.37972664 - -9.58424928 - -0.19466184 - -0.24860115 - -0.19531288 - -0.15401392 - -0.14577898 - -0.19655747 - -0.15645898 - -3.49380556 - -3.5317097 - -4.57108006 - -4.63425205 - -2.88247063 - -1.45679675 - -0.50290184 - -0.18521704 - -0.01123956 - -0.17483649 - -0.63132037 - -1.3248562 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - -0.24135757 - -1.04601971 - -2.04574044 - -3.84544799 - -7.28626119 - -7.3136314 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 omol_elem_refs: - 0.0 - -13.44558 - -78.82027 - -203.32564 - -398.94742 - -670.75275 - -1029.85403 - -1485.54188 - -2042.97832 - -2714.24015 - -3508.74317 - -4415.24203 - -5443.89712 - -6594.61834 - -7873.6878 - -9285.6593 - -10832.62132 - -12520.66852 - -14354.278 - -16323.54671 - -18436.47845 - -20696.18244 - -23110.5386 - -25682.99429 - -28418.37804 - -31317.92317 - -34383.42519 - -37623.46835 - -41039.92413 - -44637.38634 - -48417.14864 - -52373.87849 - -56512.76952 - -60836.14871 - -65344.28833 - -70041.24251 - -74929.56277 - -653.64777 - -833.31922 - -1038.0281 - -1273.96788 - -1542.45481 - -1850.74158 - -2193.91654 - -2577.18734 - -3004.13604 - -3477.52796 - -3997.31825 - -4563.75804 - -5171.82293 - -5828.85334 - -6535.61529 - -7291.54792 - -8099.87914 - -8962.17916 - -546.03214 - -690.6089 - -854.11237 - -12923.04096 - -14064.26124 - -15272.68689 - -16550.20551 - -17900.36515 - -19323.23406 - -20829.08848 - -22428.73258 - -24078.68008 - -25794.42097 - -27616.6819 - -29523.5526 - -31526.68012 - -33615.37779 - -1300.17791 - -1544.40924 - -1818.62298 - -2123.14417 - -2461.76028 - -2833.76287 - -3242.79895 - -3690.363 - -4174.99772 - -4691.75674 - -5245.36013 - -5838.12005 - -6469.07296 - -7140.86455 - -7854.60638 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 - 0.0 omc_elem_refs: - 0.0 - -0.02831808 - 4.512e-05 - -0.03227157 - -0.03842519 - -0.05829283 - -0.0845041 - -0.08806738 - -0.09021346 - -0.06669846 - -0.01218631 - -0.03650269 - -0.00059093 - -0.05787736 - -0.08730952 - -0.0975534 - -0.09264199 - -0.07124762 - -0.02374602 - -0.05299112 - -0.02631476 - -1.7772147 - -1.25083444 - -0.79579447 - -0.49099317 - -0.31414986 - -0.20292182 - -0.14011632 - -0.09929659 - -0.03771207 - -0.01117902 - -0.06168715 - -0.08873364 - -0.09512942 - -0.09035978 - -0.06910849 - -0.02244872 - -0.05303651 - -0.02871903 - -1.94805417 - -1.33379896 - -0.69169331 - -0.26184306 - -0.20631599 - -0.48251608 - -0.96911893 - -1.47569462 - -0.03845194 - -0.0142445 - -0.07118991 - -0.09940292 - -0.09235056 - -0.08755943 - -0.06544925 - -0.01246646 - -0.04692937 - -0.03225123 - -0.26086039 - -27.20024339 - -0.08412926 - -0.08225924 - -0.07799715 - -0.07806185 - 0.00043759 - -0.07459766 - 0.0 - -0.06842841 - -0.07758266 - -0.07025152 - -0.08055003 - -0.07118177 - -0.07159568 - -2.69202862 - -2.21926765 - -1.679756 - -1.06135075 - -0.4554231 - -0.14488432 - -0.18377098 - -0.03603118 - -0.01076585 - -0.06381411 - -0.0905623 - -0.10095787 - -0.09501217 - -0.0574478 - -0.00599173 - -0.04134751 - -0.0082683 - -0.08704692 - -0.49656425 - -5.24233138 - -2.32542606 - -4.3376616 - -5.96430676 - 0.0 - 0.0 - -0.03842519 - 0.0 - 0.0 """ try: ELEMENT_REF_ENERGIES = yaml.safe_load(ELEMENT_REF_ENERGIES_YAML) except yaml.YAMLError as e: # st.error(f"Error parsing YAML reference energies: {e}") # st objects can only be used in main script flow print(f"Error parsing YAML reference energies: {e}") ELEMENT_REF_ENERGIES = {} # Fallback # Check if running on Streamlit Cloud vs locally is_streamlit_cloud = os.environ.get('STREAMLIT_RUNTIME_ENV') == 'cloud' MAX_ATOMS_CLOUD = 500 # Maximum atoms allowed on Streamlit Cloud MAX_ATOMS_CLOUD_UMA = 500 # Set page configuration st.set_page_config( page_title="Molecular Structure Analysis", page_icon="đŸ§Ē", layout="wide" ) # Title and description st.markdown('## MLIP Playground', unsafe_allow_html=True) st.write('#### Run, test and compare 22 state-of-the-art universal machine learning interatomic potentials (MLIPs) for atomistic simulations of molecules and materials') st.markdown('Upload molecular structure files or select from predefined examples, then compute energies and forces using foundation models such as those from MACE or FairChem (Meta).', unsafe_allow_html=True) # Create a directory for sample structures if it doesn't exist SAMPLE_DIR = "sample_structures" os.makedirs(SAMPLE_DIR, exist_ok=True) # Dictionary of sample structures SAMPLE_STRUCTURES = { "Water": "H2O.xyz", "Methane": "CH4.xyz", "Benzene": "C6H6.xyz", "Ethane": "C2H6.xyz", "Caffeine": "caffeine.xyz", "Ibuprofen": "ibuprofen.xyz" } def get_structure_viz2(atoms_obj, style='stick', show_unit_cell=True, width=400, height=400): xyz_str = "" xyz_str += f"{len(atoms_obj)}\n" xyz_str += "Structure\n" for atom in atoms_obj: xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n" view = py3Dmol.view(width=width, height=height) view.addModel(xyz_str, "xyz") if style.lower() == 'ball_stick': view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'scale': 0.3}}) elif style.lower() == 'stick': view.setStyle({'stick': {}}) elif style.lower() == 'ball': view.setStyle({'sphere': {'scale': 0.4}}) else: view.setStyle({'stick': {'radius': 0.15}}) if show_unit_cell and atoms_obj.pbc.any(): # Check pbc.any() cell = atoms_obj.get_cell() origin = np.array([0.0, 0.0, 0.0]) if cell is not None and cell.any(): # Ensure cell is not None and not all zeros edges = [ (origin, cell[0]), (origin, cell[1]), (cell[0], cell[0] + cell[1]), (cell[1], cell[0] + cell[1]), (cell[2], cell[2] + cell[0]), (cell[2], cell[2] + cell[1]), (cell[2] + cell[0], cell[2] + cell[0] + cell[1]), (cell[2] + cell[1], cell[2] + cell[0] + cell[1]), (origin, cell[2]), (cell[0], cell[0] + cell[2]), (cell[1], cell[1] + cell[2]), (cell[0] + cell[1], cell[0] + cell[1] + cell[2]) ] for start, end in edges: view.addCylinder({ 'start': {'x': start[0], 'y': start[1], 'z': start[2]}, 'end': {'x': end[0], 'y': end[1], 'z': end[2]}, 'radius': 0.05, 'color': 'black', 'alpha': 0.7 }) view.zoomTo() view.setBackgroundColor('white') return view opt_log = [] # Define globally or pass around if necessary table_placeholder = st.empty() # Define globally if updated from callback def streamlit_log(opt): global opt_log, table_placeholder try: energy = opt.atoms.get_potential_energy() forces = opt.atoms.get_forces() fmax_step = np.max(np.linalg.norm(forces, axis=1)) if forces.shape[0] > 0 else 0.0 opt_log.append({ "Step": opt.nsteps, "Energy (eV)": round(energy, 6), "Fmax (eV/Å)": round(fmax_step, 6) }) df = pd.DataFrame(opt_log) table_placeholder.dataframe(df) except Exception as e: st.warning(f"Error in optimization logger: {e}") def check_atom_limit(atoms_obj, selected_model): if atoms_obj is None: return True num_atoms = len(atoms_obj) limit = MAX_ATOMS_CLOUD_UMA if ('UMA' in selected_model or 'ESEN MD' in selected_model) else MAX_ATOMS_CLOUD if num_atoms > limit: st.error(f"⚠ī¸ Error: Your structure contains {num_atoms} atoms, exceeding the {limit} atom limit for this model on Streamlit Cloud. Please run locally for larger systems.") return False return True MACE_MODELS = { "MACE MPA Medium": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model", "MACE OMAT Medium": "https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-medium.model", "MACE MATPES r2SCAN Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/MACE-matpes-r2scan-omat-ft.model", "MACE MATPES PBE Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/MACE-matpes-pbe-omat-ft.model", "MACE MP 0a Small": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model", "MACE MP 0a Medium": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-03-mace-128-L1_epoch-199.model", "MACE MP 0a Large": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2024-01-07-mace-128-L2_epoch-199.model", "MACE MP 0b Small": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0b/mace_agnesi_small.model", "MACE MP 0b Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0b/mace_agnesi_medium.model", "MACE MP 0b2 Small": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0b2/mace-small-density-agnesi-stress.model", # Corrected name from original code "MACE MP 0b2 Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0b2/mace-medium-density-agnesi-stress.model", "MACE MP 0b2 Large": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0b2/mace-large-density-agnesi-stress.model", "MACE MP 0b3 Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0b3/mace-mp-0b3-medium.model", } FAIRCHEM_MODELS = { "UMA Small": "uma-sm", "ESEN MD Direct All OMOL": "esen-md-direct-all-omol", "ESEN SM Conserving All OMOL": "esen-sm-conserving-all-omol", "ESEN SM Direct All OMOL": "esen-sm-direct-all-omol" } # Define the available ORB models ORB_MODELS = { "V3 OMAT Conservative (inf)": "orb-v3-conservative-inf-omat", "V3 OMAT Conservative (20)": "orb-v3-conservative-20-omat", "V3 OMAT Direct (inf)": "orb-v3-direct-inf-omat", "V3 OMAT Direct (20)": "orb-v3-direct-20-omat", "V3 MPA Conservative (inf)": "orb-v3-conservative-inf-mpa", "V3 MPA Conservative (20)": "orb-v3-conservative-20-mpa", "V3 MPA Direct (inf)": "orb-v3-direct-inf-mpa", "V3 MPA Direct (20)": "orb-v3-direct-20-mpa", } # Define the available MatterSim models MATTERSIM_MODELS = { "V1 SMALL": "MatterSim-v1.0.0-1M.pth", "V1 LARGE": "MatterSim-v1.0.0-5M.pth" } SEVEN_NET_MODELS = { "7net-0": "7net-0", "7net-l3i5": "7net-l3i5", "7net-omat": "7net-omat", "7net-mf-ompa": "7net-mf-ompa" } @st.cache_resource def get_mace_model(model_path, device, selected_default_dtype): return mace_mp(model=model_path, device=device, default_dtype=selected_default_dtype) @st.cache_resource def get_fairchem_model(selected_model_name, model_path_or_name, device, selected_task_type_fc): # Renamed args to avoid conflict predictor = pretrained_mlip.get_predict_unit(model_path_or_name, device=device) if selected_model_name == "UMA Small": calc = FAIRChemCalculator(predictor, task_name=selected_task_type_fc) else: calc = FAIRChemCalculator(predictor) return calc st.sidebar.markdown("## Input Options") input_method = st.sidebar.radio("Choose Input Method:", ["Select Example", "Upload File", "Paste Content"]) atoms = None if input_method == "Upload File": uploaded_file = st.sidebar.file_uploader("Upload structure file", type=["xyz", "cif", "POSCAR", "mol", "tmol", "vasp", "sdf", "CONTCAR"]) if uploaded_file: with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file: tmp_file.write(uploaded_file.getvalue()) tmp_filepath = tmp_file.name try: atoms = read(tmp_filepath) st.sidebar.success(f"Successfully loaded structure with {len(atoms)} atoms!") except Exception as e: st.sidebar.error(f"Error loading file: {str(e)}") finally: if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath): os.unlink(tmp_filepath) elif input_method == "Select Example": example_name = st.sidebar.selectbox("Select Example Structure:", list(SAMPLE_STRUCTURES.keys())) if example_name: file_path = os.path.join(SAMPLE_DIR, SAMPLE_STRUCTURES[example_name]) try: atoms = read(file_path) st.sidebar.success(f"Loaded {example_name} with {len(atoms)} atoms!") except Exception as e: st.sidebar.error(f"Error loading example: {str(e)}") elif input_method == "Paste Content": file_format = st.sidebar.selectbox("File Format:", ["XYZ", "CIF", "extXYZ", "POSCAR (VASP)", "Turbomole", "MOL"]) content = st.sidebar.text_area("Paste file content here:", height=200) if content: try: suffix_map = {"XYZ": ".xyz", "CIF": ".cif", "extXYZ": ".extxyz", "POSCAR (VASP)": ".vasp", "Turbomole": ".tmol", "MOL": ".mol"} suffix = suffix_map.get(file_format, ".xyz") with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file: tmp_file.write(content.encode()) tmp_filepath = tmp_file.name atoms = read(tmp_filepath) st.sidebar.success(f"Successfully parsed structure with {len(atoms)} atoms!") except Exception as e: st.sidebar.error(f"Error parsing content: {str(e)}") finally: if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath): os.unlink(tmp_filepath) if atoms is not None: if not hasattr(atoms, 'info'): atoms.info = {} atoms.info["charge"] = atoms.info.get("charge", 0) # Default charge atoms.info["spin"] = atoms.info.get("spin", 0) # Default spin (usually 2S for ASE, model might want 2S+1) st.sidebar.markdown("## Model Selection") if mattersim_available: model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem", "ORB", "MatterSim"]) else: model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem", "ORB", "SEVEN_NET"]) selected_task_type = None # For FairChem UMA if model_type == "MACE": selected_model = st.sidebar.selectbox("Select MACE Model:", list(MACE_MODELS.keys())) model_path = MACE_MODELS[selected_model] if selected_model == "MACE OMAT Medium": st.sidebar.warning("Using model under Academic Software License (ASL).") selected_default_dtype = st.sidebar.selectbox("Select Precision (default_dtype):", ['float32', 'float64']) if model_type == "FairChem": selected_model = st.sidebar.selectbox("Select FairChem Model:", list(FAIRCHEM_MODELS.keys())) model_path = FAIRCHEM_MODELS[selected_model] if selected_model == "UMA Small": st.sidebar.warning("Meta FAIR Acceptable Use Policy applies.") selected_task_type = st.sidebar.selectbox("Select UMA Model Task Type:", ["omol", "omat", "omc", "odac", "oc20"]) if selected_task_type == "omol" and atoms is not None: charge = st.sidebar.number_input("Total Charge", min_value=-10, max_value=10, value=atoms.info.get("charge",0)) spin_multiplicity = st.sidebar.number_input("Spin Multiplicity (2S + 1)", min_value=1, max_value=11, step=2, value=int(atoms.info.get("spin",0)*2+1 if atoms.info.get("spin",0) is not None else 1)) # Assuming spin in atoms.info is S atoms.info["charge"] = charge atoms.info["spin"] = spin_multiplicity # FairChem expects multiplicity if model_type == "ORB": selected_model = st.sidebar.selectbox("Select ORB Model:", list(ORB_MODELS.keys())) model_path = ORB_MODELS[selected_model] # if "omat" in selected_model: # st.sidebar.warning("Using model under Academic Software License (ASL) license, see [https://github.com/gabor1/ASL](https://github.com/gabor1/ASL). To use this model you accept the terms of the license.") selected_default_dtype = st.sidebar.selectbox("Select Precision (default_dtype):", ['float32-high', 'float32-highest', 'float64']) if model_type == "MatterSim": selected_model = st.sidebar.selectbox("Select MatterSim Model:", list(MATTERSIM_MODELS.keys())) model_path = MATTERSIM_MODELS[selected_model] if model_type == "SEVEN_NET": selected_model = st.sidebar.selectbox("Select SEVENNET Model:", list(SEVEN_NET_MODELS.keys())) model_path = SEVEN_NET_MODELS[selected_model] if atoms is not None: if not check_atom_limit(atoms, selected_model): st.stop() # Stop execution if limit exceeded device = st.sidebar.radio("Computation Device:", ["CPU", "CUDA (GPU)"], index=0 if not torch.cuda.is_available() else 1) device = "cuda" if device == "CUDA (GPU)" and torch.cuda.is_available() else "cpu" if device == "cpu" and torch.cuda.is_available(): st.sidebar.info("GPU is available but CPU was selected.") elif device == "cpu" and not torch.cuda.is_available(): st.sidebar.info("No GPU detected. Using CPU.") st.sidebar.markdown("## Task Selection") task = st.sidebar.selectbox("Select Calculation Task:", ["Energy Calculation", "Energy + Forces Calculation", "Atomization/Cohesive Energy", # New Task Added "Geometry Optimization", "Cell + Geometry Optimization"]) if "Optimization" in task: st.sidebar.markdown("### Optimization Parameters") max_steps = st.sidebar.slider("Maximum Steps:", min_value=10, max_value=200, value=50, step=1) # Increased max_steps fmax = st.sidebar.slider("Convergence Threshold (eV/Å):", min_value=0.001, max_value=0.1, value=0.01, step=0.001, format="%.3f") # Adjusted default fmax optimizer_type = st.sidebar.selectbox("Optimizer:", ["BFGS", "LBFGS", "FIRE"], index=1) # Renamed to optimizer_type if atoms is not None: col1, col2 = st.columns(2) with col1: st.markdown('### Structure Visualization', unsafe_allow_html=True) view_3d = get_structure_viz2(atoms, style='stick', show_unit_cell=True, width=400, height=400) st.components.v1.html(view_3d._make_html(), width=400, height=400) st.markdown("### Structure Information") atoms_info = { "Number of Atoms": len(atoms), "Chemical Formula": atoms.get_chemical_formula(), "Periodic Boundary Conditions (PBC)": atoms.pbc.tolist(), "Cell Dimensions": np.round(atoms.cell.cellpar(),3).tolist() if atoms.pbc.any() and atoms.cell is not None and atoms.cell.any() else "No cell / Non-periodic", "Atom Types": ", ".join(sorted(list(set(atoms.get_chemical_symbols())))) } for key, value in atoms_info.items(): st.write(f"**{key}:** {value}") with col2: st.markdown('## Calculation Setup', unsafe_allow_html=True) st.markdown("### Selected Model") st.write(f"**Model Type:** {model_type}") st.write(f"**Model:** {selected_model}") if model_type == "FairChem" and selected_model == "UMA Small": st.write(f"**UMA Task Type:** {selected_task_type}") st.write(f"**Device:** {device}") st.markdown("### Selected Task") st.write(f"**Task:** {task}") if "Optimization" in task: st.write(f"**Max Steps:** {max_steps}") st.write(f"**Convergence Threshold:** {fmax} eV/Å") st.write(f"**Optimizer:** {optimizer_type}") run_calculation = st.button("Run Calculation", type="primary") if run_calculation: results = {} #global table_placeholder # Ensure they are accessible opt_log = [] # Reset log for each run if "Optimization" in task: table_placeholder = st.empty() # Recreate placeholder for table try: with st.spinner("Running calculation... Please wait."): calc_atoms = atoms.copy() if model_type == "MACE": # st.write("Setting up MACE calculator...") calc = get_mace_model(model_path, device, selected_default_dtype) elif model_type == "FairChem": # FairChem # st.write("Setting up FairChem calculator...") # Workaround for potential dtype issues when switching models if device == "cpu": # Ensure torch default dtype matches if needed torch.set_default_dtype(torch.float32) _ = get_mace_model(MACE_MODELS["MACE MP 0a Small"], 'cpu', 'float32') # Dummy call calc = get_fairchem_model(selected_model, model_path, device, selected_task_type) elif model_type == "ORB": st.write("Setting up ORB calculator...") orbff = pretrained.orb_v3_conservative_inf_omat(device=device, precision=selected_default_dtype) calc = ORBCalculator(orbff, device=device) elif model_type == "MatterSim": st.write("Setting up MatterSim calculator...") calc = MatterSimCalculator(load_path=model_path, device=device) elif model_type == "SEVEN_NET": st.write("Setting up SEVENNET calculator...") calc = SevenNetCalculator(model=model_path, device=device) calc_atoms.calc = calc if task == "Energy Calculation": energy = calc_atoms.get_potential_energy() results["Energy"] = f"{energy:.6f} eV" elif task == "Energy + Forces Calculation": energy = calc_atoms.get_potential_energy() forces = calc_atoms.get_forces() max_force = np.max(np.sqrt(np.sum(forces**2, axis=1))) if forces.shape[0] > 0 else 0.0 results["Energy"] = f"{energy:.6f} eV" results["Maximum Force"] = f"{max_force:.6f} eV/Å" elif task == "Atomization/Cohesive Energy": st.write("Calculating system energy...") E_system = calc_atoms.get_potential_energy() num_atoms = len(calc_atoms) if num_atoms == 0: st.error("Cannot calculate atomization/cohesive energy for a system with zero atoms.") results["Error"] = "System has no atoms." else: atomic_numbers = calc_atoms.get_atomic_numbers() E_isolated_atoms_total = 0.0 calculation_possible = True if model_type == "FairChem": st.write("Fetching FairChem reference energies for isolated atoms...") ref_key_suffix = "_elem_refs" chosen_ref_list_name = None if selected_model == "UMA Small": if selected_task_type: chosen_ref_list_name = selected_task_type + ref_key_suffix elif "ESEN" in selected_model: chosen_ref_list_name = "omol" + ref_key_suffix if chosen_ref_list_name and chosen_ref_list_name in ELEMENT_REF_ENERGIES: ref_energies = ELEMENT_REF_ENERGIES[chosen_ref_list_name] missing_Z_refs = [] for Z_val in atomic_numbers: if Z_val > 0 and Z_val < len(ref_energies): E_isolated_atoms_total += ref_energies[Z_val] else: if Z_val not in missing_Z_refs: missing_Z_refs.append(Z_val) if missing_Z_refs: st.warning(f"Reference energy for atomic number(s) {sorted(list(set(missing_Z_refs)))} " f"not found in '{chosen_ref_list_name}' list (max Z defined: {len(ref_energies)-1}). " "These atoms are treated as having 0 reference energy.") else: st.error(f"Could not find or determine reference energy list for FairChem model: '{selected_model}' " f"and UMA task type: '{selected_task_type}'. Cannot calculate atomization/cohesive energy.") results["Error"] = "Missing FairChem reference energies." calculation_possible = False else:# == "MACE": st.write("Calculating isolated atom energies with MACE...") unique_atomic_numbers = sorted(list(set(atomic_numbers))) atom_counts = {Z_unique: np.count_nonzero(atomic_numbers == Z_unique) for Z_unique in unique_atomic_numbers} progress_text = "Calculating isolated atom energies: 0% complete" mace_progress_bar = st.progress(0, text=progress_text) for i, Z_unique in enumerate(unique_atomic_numbers): isolated_atom = Atoms(numbers=[Z_unique], cell=[20, 20, 20], pbc=False) if not hasattr(isolated_atom, 'info'): isolated_atom.info = {} isolated_atom.info["charge"] = 0 isolated_atom.info["spin"] = 0 isolated_atom.calc = calc # Use the same MACE calculator E_isolated_atom_type = isolated_atom.get_potential_energy() E_isolated_atoms_total += E_isolated_atom_type * atom_counts[Z_unique] progress_val = (i + 1) / len(unique_atomic_numbers) mace_progress_bar.progress(progress_val, text=f"Calculating isolated atom energies for Z={Z_unique}: {int(progress_val*100)}% complete") mace_progress_bar.empty() if calculation_possible: is_periodic = any(calc_atoms.pbc) if is_periodic: cohesive_E = (E_isolated_atoms_total - E_system) / num_atoms results["Cohesive Energy"] = f"{cohesive_E:.6f} eV/atom" else: atomization_E = E_isolated_atoms_total - E_system results["Atomization Energy"] = f"{atomization_E:.6f} eV" results["System Energy ($E_{system}$)"] = f"{E_system:.6f} eV" results["Total Isolated Atom Energy ($\sum E_{atoms}$)"] = f"{E_isolated_atoms_total:.6f} eV" elif "Optimization" in task: # Handles both Geometry and Cell+Geometry Opt opt_atoms_obj = FrechetCellFilter(calc_atoms) if task == "Cell + Geometry Optimization" else calc_atoms if optimizer_type == "BFGS": opt = BFGS(opt_atoms_obj) elif optimizer_type == "LBFGS": opt = LBFGS(opt_atoms_obj) else: # FIRE opt = FIRE(opt_atoms_obj) # opt.attach(streamlit_log, interval=1) # Removed lambda for simplicity if streamlit_log is defined correctly opt.attach(lambda: streamlit_log(opt), interval=1) st.write(f"Running {task.lower()}...") opt.run(fmax=fmax, steps=max_steps) energy = calc_atoms.get_potential_energy() forces = calc_atoms.get_forces() max_force = np.max(np.sqrt(np.sum(forces**2, axis=1))) if forces.shape[0] > 0 else 0.0 results["Final Energy"] = f"{energy:.6f} eV" results["Final Maximum Force"] = f"{max_force:.6f} eV/Å" results["Steps Taken"] = opt.get_number_of_steps() results["Converged"] = "Yes" if opt.converged() else "No" if task == "Cell + Geometry Optimization": results["Final Cell Parameters"] = np.round(calc_atoms.cell.cellpar(), 4).tolist() st.success("Calculation completed successfully!") st.markdown("### Results") for key, value in results.items(): st.write(f"**{key}:** {value}") if "Optimization" in task and "Final Energy" in results: # Check if opt was successful st.markdown("### Optimized Structure") # Need get_structure_viz function that takes atoms obj def get_structure_viz_simple(atoms_obj_viz): xyz_str_viz = f"{len(atoms_obj_viz)}\nStructure\n" for atom_viz in atoms_obj_viz: xyz_str_viz += f"{atom_viz.symbol} {atom_viz.position[0]:.6f} {atom_viz.position[1]:.6f} {atom_viz.position[2]:.6f}\n" view_viz = py3Dmol.view(width=400, height=400) view_viz.addModel(xyz_str_viz, "xyz") view_viz.setStyle({'stick': {}}) if any(atoms_obj_viz.pbc): # Show cell for optimized periodic structures cell_viz = atoms_obj_viz.get_cell() if cell_viz is not None and cell_viz.any(): # Simplified cell drawing for brevity, use get_structure_viz2 if full cell needed view_viz.addUnitCell({'box': {'lx':cell_viz.lengths()[0],'ly':cell_viz.lengths()[1],'lz':cell_viz.lengths()[2], 'hx':cell_viz.cellpar()[3],'hy':cell_viz.cellpar()[4],'hz':cell_viz.cellpar()[5]}}) view_viz.zoomTo() view_viz.setBackgroundColor('white') return view_viz opt_view = get_structure_viz2(calc_atoms, style='stick', show_unit_cell=True, width=400, height=400) st.components.v1.html(opt_view._make_html(), width=400, height=400) with tempfile.NamedTemporaryFile(delete=False, suffix=".xyz", mode="w+") as tmp_file_opt: write(tmp_file_opt.name, calc_atoms, format="xyz") tmp_filepath_opt = tmp_file_opt.name with open(tmp_filepath_opt, 'r') as file_opt: xyz_content_opt = file_opt.read() st.download_button( label="Download Optimized Structure (XYZ)", data=xyz_content_opt, file_name="optimized_structure.xyz", mime="chemical/x-xyz" ) os.unlink(tmp_filepath_opt) except Exception as e: st.error(f"🔴 Calculation error: {str(e)}") st.error("Please check the structure, model compatibility, and parameters. For FairChem UMA, ensure the task type (omol, omat etc.) is appropriate for your system (e.g. omol for molecules, omat for materials).") import traceback st.error(f"Traceback: {traceback.format_exc()}") else: st.info("👋 Welcome! Please select or upload a structure using the sidebar options to begin.") st.markdown("---") with st.expander('ℹī¸ About This App & Foundational MLIPs'): st.write(""" **Test, compare, and benchmark universal machine learning interatomic potentials (MLIPs).** This application allows you to perform atomistic simulations using pre-trained foundational MLIPs from the MACE and FairChem (by Meta AI) libraries. **Features:** - Upload structure files (XYZ, CIF, POSCAR, etc.) or use built-in examples. - Select from various MACE and FairChem models. - Calculate energies, forces, and perform geometry/cell optimizations. - **New**: Calculate atomization energy (for molecules) or cohesive energy (for periodic systems). - Visualize atomic structures in 3D and download results. **Quick Start:** 1. **Input**: Choose an input method in the sidebar (e.g., "Select Example"). 2. **Model**: Pick a model type (MACE/FairChem) and specific model. For FairChem UMA, select the appropriate task type (e.g., `omol` for molecules, `omat` for materials). 3. **Task**: Select a calculation task (e.g., "Energy Calculation", "Atomization/Cohesive Energy", "Geometry Optimization"). 4. **Run**: Click "Run Calculation" and view the results. **Atomization/Cohesive Energy Notes:** - **Atomization Energy** ($E_{\text{atomization}} = \sum E_{\text{isolated atoms}} - E_{\text{molecule}}$) is typically for non-periodic systems (molecules). - **Cohesive Energy** ($E_{\text{cohesive}} = (\sum E_{\text{isolated atoms}} - E_{\text{bulk system}}) / N_{\text{atoms}}$) is for periodic systems. - For **MACE models**, isolated atom energies are computed on-the-fly. - For **FairChem models**, isolated atom energies are based on pre-tabulated reference values (provided in a YAML-like structure within the app). Ensure the selected FairChem task type (`omol`, `omat`, etc. for UMA models) or model type (ESEN models use `omol` references) aligns with the system and has the necessary elemental references. """) st.markdown("Universal MLIP Playground App | Created with Streamlit, ASE, MACE, FairChem and ❤ī¸") st.markdown("Developed by [Manas Sharma](https://manas.bragitoff.com/) ([Fundamental AI Research (FAIR) team, Meta AI](https://ai.meta.com/research/fair/) and [Ananth Govind Rajan Group, IISc Bangalore](https://www.agrgroup.org/))")