Update src/streamlit_app.py

src/streamlit_app.py

@@ -14,6 +14,7 @@ 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 mattersim.forcefield import MatterSimCalculator
 import pandas as pd
 import yaml # Added for FairChem reference energies
 
@@ -686,6 +687,11 @@ ORB_MODELS = {
     "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"
+}
 @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)
@@ -780,6 +786,9 @@ if model_type == "ORB":
     # 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 atoms is not None:
     if not check_atom_limit(atoms, selected_model):
         st.stop() # Stop execution if limit exceeded
@@ -870,6 +879,9 @@ if atoms is not None:
                         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)
                     calc_atoms.calc = calc
                     
                     if task == "Energy Calculation":