mlip-playground / src /streamlit_app.py
ManasSharma07's picture
Update src/streamlit_app.py
514c0ba verified
raw
history blame
26.5 kB
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
import pandas as pd
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")
login(token = hf_token)
except Exception as e:
print("hf secret not defined/assigned")
import os
os.environ["STREAMLIT_WATCHER_TYPE"] = "none"
# 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"
)
# Add CSS for better formatting
# st.markdown("""
# <style>
# .stApp {
# max-width: 1200px;
# margin: 0 auto;
# }
# .main-header {
# font-size: 2.5rem;
# font-weight: bold;
# margin-bottom: 1rem;
# }
# .section-header {
# font-size: 1.5rem;
# font-weight: bold;
# margin-top: 1.5rem;
# margin-bottom: 1rem;
# }
# .info-text {
# font-size: 1rem;
# color: #555;
# }
# </style>
# """, unsafe_allow_html=True)
# Title and description
st.markdown('## MLIP Playground', unsafe_allow_html=True)
st.write('#### Run, test and compare >17 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):
"""
Generate visualization of atomic structure with optional unit cell display
Parameters:
-----------
atoms_obj : ase.Atoms
ASE Atoms object containing the structure
style : str
Visualization style: 'ball_stick', 'stick', or 'ball'
show_unit_cell : bool
Whether to display unit cell for periodic systems
width, height : int
Dimensions of the visualization window
Returns:
--------
py3Dmol.view object
"""
# Convert atoms to XYZ format
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"
# Create a py3Dmol visualization
view = py3Dmol.view(width=width, height=height)
view.addModel(xyz_str, "xyz")
# Set molecular style based on input
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:
# Default to stick if unknown style
view.setStyle({'stick': {'radius': 0.15}})
# Add unit cell visualization for periodic systems
if show_unit_cell and any(atoms_obj.pbc):
cell = atoms_obj.get_cell()
# Define unit cell edges
origin = np.array([0.0, 0.0, 0.0])
edges = [
# Bottom face
(origin, cell[0]), # a
(origin, cell[1]), # b
(cell[0], cell[0] + cell[1]), # a+b from a
(cell[1], cell[0] + cell[1]), # a+b from b
# Top face
(cell[2], cell[2] + cell[0]), # a from c
(cell[2], cell[2] + cell[1]), # b from c
(cell[2] + cell[0], cell[2] + cell[0] + cell[1]), # a+b from c+a
(cell[2] + cell[1], cell[2] + cell[0] + cell[1]), # a+b from c+b
# Vertical edges
(origin, cell[2]), # c
(cell[0], cell[0] + cell[2]), # c from a
(cell[1], cell[1] + cell[2]), # c from b
(cell[0] + cell[1], cell[0] + cell[1] + cell[2]) # c from a+b
]
# Add unit cell lines
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
# Custom logger that updates the table
def streamlit_log(opt):
energy = opt.atoms.get_potential_energy()
forces = opt.atoms.get_forces()
fmax_step = np.max(np.linalg.norm(forces, axis=1))
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)
# Function to check atom count limits
def check_atom_limit(atoms_obj, selected_model):
if atoms_obj is None:
return True
num_atoms = len(atoms_obj)
if ('UMA' in selected_model or 'ESEN MD' in selected_model) and num_atoms > MAX_ATOMS_CLOUD_UMA:
st.error(f"⚠️ Error: Your structure contains {num_atoms} atoms, which exceeds the {MAX_ATOMS_CLOUD_UMA} atom limit for Streamlit Cloud deployments for large sized FairChem models. For larger systems, please download the repository from GitHub and run it locally on your machine where no atom limit applies.")
st.info("💡 Running locally allows you to process much larger structures and use your own computational resources more efficiently.")
return False
if num_atoms > MAX_ATOMS_CLOUD:
st.error(f"⚠️ Error: Your structure contains {num_atoms} atoms, which exceeds the {MAX_ATOMS_CLOUD} atom limit for Streamlit Cloud deployments. For larger systems, please download the repository from GitHub and run it locally on your machine where no atom limit applies.")
st.info("💡 Running locally allows you to process much larger structures and use your own computational resources more efficiently.")
return False
return True
# Define the available MACE models
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-large-density-agnesi-stress.model",
"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",
}
# Define the available FairChem models
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"
}
@st.cache_resource
def get_mace_model(model_path, device, selected_default_dtype):
# Create a model of the specified type.
return mace_mp(model=model_path, device=device, default_dtype=selected_default_dtype)
@st.cache_resource
def get_fairchem_model(selected_model, model_path, device, selected_task_type):
predictor = pretrained_mlip.get_predict_unit(model_path, device=device)
if selected_model == "UMA Small":
calc = FAIRChemCalculator(predictor, task_name=selected_task_type)
else:
calc = FAIRChemCalculator(predictor)
return calc
# Sidebar for file input and parameters
st.sidebar.markdown("## Input Options")
# Input method selection
input_method = st.sidebar.radio("Choose Input Method:", ["Select Example", "Upload File", "Paste Content"])
# Initialize atoms variable
atoms = None
# File upload option
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 is not None:
# Create a temporary file to save the uploaded content
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:
# Read the structure using ASE
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)}")
# Clean up the temporary file
os.unlink(tmp_filepath)
# Example structure selection
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)}")
# Paste content option
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 and st.sidebar.button("Parse Content"):
try:
# Create a temporary file with the pasted content
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
# Read the structure using ASE
atoms = read(tmp_filepath)
st.sidebar.success(f"Successfully parsed structure with {len(atoms)} atoms!")
# Clean up the temporary file
os.unlink(tmp_filepath)
except Exception as e:
st.sidebar.error(f"Error parsing content: {str(e)}")
# Model selection
st.sidebar.markdown("## Model Selection")
model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem"])
selected_task_type = None
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) 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', '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. This model was developed by the Fundamental AI Research (FAIR) team at Meta. By using it, you agree to their acceptable use policy, which prohibits using their models to violate the law or others' rights, plan or develop activities that present a risk of death or harm, and deceive or mislead others.")
selected_task_type = st.sidebar.selectbox("Select UMA Model Task Type:", ["omol", "omat", "omc", "odac", "oc20"])
# Check atom count limit
if atoms is not None:
check_atom_limit(atoms, selected_model)
#st.sidebar.success(f"Successfully parsed structure with {len(atoms)} atoms!")
# Device selection
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. Calculations will be slower.")
elif device == "cpu" and not torch.cuda.is_available():
st.sidebar.info("No GPU detected. Using CPU for calculations.")
# Task selection
st.sidebar.markdown("## Task Selection")
task = st.sidebar.selectbox("Select Calculation Task:",
["Energy Calculation",
"Energy + Forces Calculation",
"Geometry Optimization",
"Cell + Geometry Optimization"])
# Optimization parameters
if "Optimization" in task:
st.sidebar.markdown("### Optimization Parameters")
max_steps = st.sidebar.slider("Maximum Steps:", min_value=10, max_value=50, value=25, step=1)
fmax = st.sidebar.slider("Convergence Threshold (eV/Å):",
min_value=0.001, max_value=0.1, value=0.05, step=0.001, format="%.3f")
optimizer = st.sidebar.selectbox("Optimizer:", ["BFGS", "LBFGS", "FIRE"], index=1)
# Main content area
if atoms is not None:
col1, col2 = st.columns(2)
with col1:
st.markdown('### Structure Visualization', unsafe_allow_html=True)
# Generate visualization
def get_structure_viz(atoms_obj):
# Convert atoms to XYZ format
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"
# Create a py3Dmol visualization
view = py3Dmol.view(width=400, height=400)
view.addModel(xyz_str, "xyz")
view.setStyle({'stick': {}})
view.zoomTo()
view.setBackgroundColor('white')
return view
# Display the 3D structure
view = get_structure_viz2(atoms, style='stick', show_unit_cell=True, width=400, height=400)
# view = get_structure_viz(atoms)
html_str = view._make_html()
st.components.v1.html(html_str, width=400, height=400)
# Display structure information
st.markdown("### Structure Information")
atoms_info = {
"Number of Atoms": len(atoms),
"Chemical Formula": atoms.get_chemical_formula(),
"Cell Dimensions": atoms.cell.cellpar() if atoms.cell else "No cell defined",
"Atom Types": ", ".join(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)
# Display calculation details
st.markdown("### Selected Model")
st.write(f"**Model Type:** {model_type}")
st.write(f"**Model:** {selected_model}")
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}")
# Run calculation button
run_calculation = st.button("Run Calculation", type="primary")
if run_calculation:
try:
with st.spinner("Running calculation..."):
# Copy atoms to avoid modifying the original
calc_atoms = atoms.copy()
# Set up calculator based on selected model
if model_type == "MACE":
st.write("Setting up MACE calculator...")
calc = get_mace_model(model_path, device, selected_default_dtype)
else: # FairChem
st.write("Setting up FairChem calculator...")
# Seems like the FairChem models use float32 and when switching from MACE 64 model to FairChem float32 model we get an error
# probably due to both sharing the underlying torch implementation
# So just a dummy statement to swithc torch to 32 bit
calc = get_mace_model('https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model', 'cpu', 'float32')
calc = get_fairchem_model(selected_model, model_path, device, selected_task_type)
# Attach calculator to atoms
calc_atoms.calc = calc
# Perform the selected task
results = {}
if task == "Energy Calculation":
# Calculate energy
energy = calc_atoms.get_potential_energy()
results["Energy"] = f"{energy:.6f} eV"
elif task == "Energy + Forces Calculation":
# Calculate energy and forces
energy = calc_atoms.get_potential_energy()
forces = calc_atoms.get_forces()
max_force = np.max(np.sqrt(np.sum(forces**2, axis=1)))
results["Energy"] = f"{energy:.6f} eV"
results["Maximum Force"] = f"{max_force:.6f} eV/Å"
elif task == "Geometry Optimization":
# Set up optimizer
if optimizer == "BFGS":
opt = BFGS(calc_atoms)
elif optimizer == "LBFGS":
opt = LBFGS(calc_atoms)
else: # FIRE
opt = FIRE(calc_atoms)
# Streamlit placeholder for live-updating table
table_placeholder = st.empty()
# Container for log data
opt_log = []
# Attach the Streamlit logger to the optimizer
opt.attach(lambda: streamlit_log(opt), interval=1)
# Run optimization
st.write("Running geometry optimization...")
opt.run(fmax=fmax, steps=max_steps)
# Get results
energy = calc_atoms.get_potential_energy()
forces = calc_atoms.get_forces()
max_force = np.max(np.sqrt(np.sum(forces**2, axis=1)))
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"
elif task == "Cell + Geometry Optimization":
# Set up optimizer with FrechetCellFilter
fcf = FrechetCellFilter(calc_atoms)
if optimizer == "BFGS":
opt = BFGS(fcf)
elif optimizer == "LBFGS":
opt = LBFGS(fcf)
else: # FIRE
opt = FIRE(fcf)
# Streamlit placeholder for live-updating table
table_placeholder = st.empty()
# Container for log data
opt_log = []
# Attach the Streamlit logger to the optimizer
opt.attach(lambda: streamlit_log(opt), interval=1)
# Run optimization
st.write("Running cell + geometry optimization...")
opt.run(fmax=fmax, steps=max_steps)
# Get results
energy = calc_atoms.get_potential_energy()
forces = calc_atoms.get_forces()
max_force = np.max(np.sqrt(np.sum(forces**2, axis=1)))
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"
results["Final Cell Parameters"] = np.round(calc_atoms.cell.cellpar(), 4)
# Show results
st.success("Calculation completed successfully!")
st.markdown("### Results")
for key, value in results.items():
st.write(f"**{key}:** {value}")
# If we did an optimization, show the final structure
if "Optimization" in task:
st.markdown("### Optimized Structure")
view = get_structure_viz(calc_atoms)
html_str = view._make_html()
st.components.v1.html(html_str, width=400, height=400)
# Add download option for optimized structure
# First save the structure to a file
with tempfile.NamedTemporaryFile(delete=False, suffix=".xyz") as tmp_file:
write(tmp_file.name, calc_atoms)
tmp_filepath = tmp_file.name
# Read the content for downloading
with open(tmp_filepath, 'r') as file:
xyz_content = file.read()
st.download_button(
label="Download Optimized Structure (XYZ)",
data=xyz_content,
file_name="optimized_structure.xyz",
mime="chemical/x-xyz"
)
# Clean up the temp file
os.unlink(tmp_filepath)
except Exception as e:
st.error(f"Calculation error: {str(e)}")
st.error("Please make sure the structure is valid and compatible with the selected model.")
else:
# Display instructions if no structure is loaded
st.info("Please select a structure using the sidebar options to begin.")
# Footer
st.markdown("---")
with st.expander('## About This App'):
# Show some information about the app
st.write("""
Test, compare and benchmark universal machine learning interatomic potentials (MLIPs).
This app allows you to perform atomistic simulations using pre-trained foundational MLIPs such as those from the MACE and FairChem libraries.
### Features:
- Upload structure files (XYZ, CIF, POSCAR, etc.) or select from examples
- Choose between MACE and FairChem ML models (more models coming soon)
- Perform energy calculations, forces calculations, or geometry optimizations
- Visualize structures in 3D
- Download optimized structures
### Getting Started:
1. Select an input method in the sidebar
2. Choose a model and computational parameters
3. Select a calculation task
4. Run the calculation and analyze the results
""")
st.markdown("Universal MLIP Playground App | Created with Streamlit, ASE, MACE, FairChem and ❤️")
st.markdown("Made by [Manas Sharma](https://manas.bragitoff.com/) in the groups of [Prof. Ananth Govind Rajan](https://www.agrgroup.org/) and [Prof. Sudeep Punnathanam](https://chemeng.iisc.ac.in/sudeep/)")