Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import os
+import io
 import tempfile
 import torch
 # FOR CPU only mode
@@ -24,7 +25,11 @@ import yaml # Added for FairChem reference energies
 import subprocess
 import sys
 import pkg_resources
+from ase.vibrations import Vibrations
+import matplotlib.pyplot as plt
 mattersim_available = False
+if mattersim_available:
+    from mattersim.forcefield import MatterSimCalculator
 # try:
 #     subprocess.check_call([sys.executable, "-m", "pip", "install", "mattersim"])
 # except Exception as e:
@@ -1028,14 +1033,14 @@ 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'])
-    selected_default_dtype = 'float64'
+    selected_default_dtype = st.sidebar.selectbox("Select Precision (default_dtype):", ['float32-high', 'float32-highest'])
 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()))
     if selected_model == '7net-mf-ompa':
-        selected_modal_7net = st.sidebar.selectbox("Select Modal (multi fidelity model):", ['omat', 'mpa'])
+        selected_modal_7net = st.sidebar.selectbox("Select Modal (multi fidelity model):", ['omat24', 'mpa'])
     model_path = SEVEN_NET_MODELS[selected_model]
 if atoms is not None:
     if not check_atom_limit(atoms, selected_model):
@@ -1055,7 +1060,8 @@ task = st.sidebar.selectbox("Select Calculation Task:",
                             "Energy + Forces Calculation", 
                             "Atomization/Cohesive Energy", # New Task Added
                             "Geometry Optimization", 
-                            "Cell + Geometry Optimization"])
+                            "Cell + Geometry Optimization",
+                            "Vibrational Mode Analysis"])
 
 if "Optimization" in task:
     st.sidebar.markdown("### Optimization Parameters")
@@ -1117,18 +1123,19 @@ if atoms is not None:
                  table_placeholder = st.empty() # Recreate placeholder for table
 
             try:
+                torch.set_default_dtype(torch.float32)
                 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)
+                        calc = get_mace_model(model_path, device, 'float32')
                     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
+                        # 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...")
@@ -1136,10 +1143,14 @@ if atoms is not None:
                         calc = ORBCalculator(orbff, device=device)
                     elif model_type == "MatterSim":
                         # st.write("Setting up MatterSim calculator...")
+                        # NOTE: Running mattersim on windows requires changing source code file
+                        # https://github.com/microsoft/mattersim/issues/112
+                        # mattersim/datasets/utils/convertor.py: 117
+                        # to pbc_ = np.array(structure.pbc, dtype=np.int64)
                         calc = MatterSimCalculator(load_path=model_path, device=device)
                     elif model_type == "SEVEN_NET":
                         # st.write("Setting up SEVENNET calculator...")
-                        if model_path=='7net-mf-omp':
+                        if model_path=='7net-mf-ompa':
                             calc = SevenNetCalculator(model=model_path, modal=selected_modal_7net, device=device)
                         else:
                             calc = SevenNetCalculator(model=model_path, device=device)
@@ -1234,6 +1245,7 @@ if atoms is not None:
                                 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
+                        is_periodic = any(calc_atoms.pbc)
                         opt_atoms_obj = FrechetCellFilter(calc_atoms) if task == "Cell + Geometry Optimization" else calc_atoms
                         # Create temporary trajectory file
                         traj_filename = tempfile.NamedTemporaryFile(delete=False, suffix=".traj").name
@@ -1268,30 +1280,15 @@ if atoms is not None:
                     
                     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=viz_style, show_unit_cell=True, width=400, height=400)
+                        opt_view = get_structure_viz2(opt_atoms_obj, style=viz_style, 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="extxyz")
+                            if is_periodic:
+                                write(tmp_file_opt.name, calc_atoms, format="extxyz")
+                            else:
+                                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:
@@ -1309,81 +1306,6 @@ if atoms is not None:
                         show_optimized_structure_download_button()
                         os.unlink(tmp_filepath_opt)
 
-                        # # Convert trajectory to XYZ for download
-                        # @st.fragment
-                        # def show_trajectory():
-                        #     if os.path.exists(traj_filename):
-                        #         try:
-                        #             from ase.io import read
-                        #             from ase.visualize import view
-                        #             import py3Dmol
-                            
-                        #             trajectory = read(traj_filename, index=':')
-                        #             st.markdown("### Optimization Trajectory")
-                        #             st.write(f"Captured {len(trajectory)} optimization steps")
-                            
-                        #             # Store the trajectory in session state
-                        #             if "traj_frames" not in st.session_state:
-                        #                 st.session_state.traj_frames = trajectory
-                        #                 st.session_state.traj_index = 0
-                            
-                        #             # Navigation Buttons
-                        #             col1, col2, col3, col4 = st.columns(4)
-                        #             with col1:
-                        #                 if st.button("⏮ First"):
-                        #                     st.session_state.traj_index = 0
-                        #             with col2:
-                        #                 if st.button("◀ Previous") and st.session_state.traj_index > 0:
-                        #                     st.session_state.traj_index -= 1
-                        #             with col3:
-                        #                 if st.button("Next ▶") and st.session_state.traj_index < len(st.session_state.traj_frames) - 1:
-                        #                     st.session_state.traj_index += 1
-                        #             with col4:
-                        #                 if st.button("Last ⏭"):
-                        #                     st.session_state.traj_index = len(st.session_state.traj_frames) - 1
-                            
-                        #             # Show current frame
-                        #             current_atoms = st.session_state.traj_frames[st.session_state.traj_index]
-                        #             st.write(f"Frame {st.session_state.traj_index + 1}/{len(st.session_state.traj_frames)}")
-                            
-                        #             # Convert to xyz string for py3Dmol
-                        #             def atoms_to_xyz_string(atoms):
-                        #                 xyz_str = f"{len(atoms)}\nStep {st.session_state.traj_index}, Energy = {atoms.get_potential_energy():.6f} eV\n"
-                        #                 for atom in atoms:
-                        #                     xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
-                        #                 return xyz_str
-                            
-                        #             xyz_str = atoms_to_xyz_string(current_atoms)
-                        #             view = py3Dmol.view(width=400, height=400)
-                        #             view.addModel(xyz_str, "xyz")
-                        #             view.setStyle({'stick': {}})
-                        #             view.zoomTo()
-                        #             view.setBackgroundColor("white")
-                        #             st.components.v1.html(view._make_html(), height=400, width=400)
-                            
-                        #             # Download entire trajectory
-                        #             @st.fragment
-                        #             def show_trajectory_download_button():
-                        #                 trajectory_xyz = ""
-                        #                 for i, atoms in enumerate(st.session_state.traj_frames):
-                        #                     trajectory_xyz += f"{len(atoms)}\nStep {i}, Energy = {atoms.get_potential_energy():.6f} eV\n"
-                        #                     for atom in atoms:
-                        #                         trajectory_xyz += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
-                        #                 st.download_button(
-                        #                     label="Download Optimization Trajectory (XYZ)",
-                        #                     data=trajectory_xyz,
-                        #                     file_name="optimization_trajectory.xyz",
-                        #                     mime="chemical/x-xyz"
-                        #                 )
-                        #             show_trajectory_download_button()
-                            
-                        #         except Exception as e:
-                        #             st.warning(f"Could not process trajectory: {e}")
-                            
-                        #         finally:
-                        #             os.unlink(traj_filename)
-
-                        # show_trajectory()
                         @st.fragment
                         def show_trajectory_and_controls():
                             from ase.io import read
@@ -1439,14 +1361,6 @@ if atoms is not None:
                                     xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
                                 return xyz_str
                         
-                            # xyz_str = atoms_to_xyz_string(current_atoms, st.session_state.traj_index)
-                        
-                            # view = py3Dmol.view(width=400, height=400)
-                            # view.addModel(xyz_str, "xyz")
-                            # view.setStyle({'stick': {}})
-                            # view.zoomTo()
-                            # view.setBackgroundColor("white")
-                            # st.components.v1.html(view._make_html(), height=400, width=400)
                             traj_view = get_structure_viz2(current_atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
                             st.components.v1.html(traj_view._make_html(), width=400, height=400)
                         
@@ -1462,7 +1376,73 @@ if atoms is not None:
                             )
                         
                         show_trajectory_and_controls()
+                    elif task == "Vibrational Mode Analysis":
 
+                        st.write("Running vibrational mode analysis using finite differences...")
+
+                        natoms = len(calc_atoms)
+                        is_linear = False  # Set manually or auto-detect
+                        nmodes_expected = 3 * natoms - (5 if is_linear else 6)
+
+                        # Create temporary directory to store .vib files
+                        with tempfile.TemporaryDirectory() as tmpdir:
+                            vib = Vibrations(calc_atoms, name=os.path.join(tmpdir, 'vib'))
+
+                            with st.spinner("Calculating vibrational modes... This may take a few minutes."):
+                                vib.run()
+                                freqs = vib.get_frequencies()
+                            
+                            # Convert frequencies to cm⁻¹
+                            freqs_cm = freqs #/ cm
+
+                            # Classify frequencies
+                            mode_data = []
+                            for i, freq in enumerate(freqs_cm):
+                                if freq < 0:
+                                    label = "Imaginary"
+                                elif abs(freq) < 500:
+                                    label = "Low"
+                                else:
+                                    label = "Physical"
+                                mode_data.append({
+                                    "Mode": i + 1,
+                                    "Frequency (cm⁻¹)": round(freq, 2),
+                                    "Type": label
+                                })
+
+                            df_modes = pd.DataFrame(mode_data)
+
+                            # Display summary and mode count
+                            st.success("Vibrational analysis completed.")
+                            st.write(f"Number of atoms: {natoms}")
+                            st.write(f"Expected vibrational modes: {nmodes_expected}")
+                            st.write(f"Found {len(freqs_cm)} modes (including translational/rotational modes).")
+
+                            # Show table of modes
+                            st.write("### Vibrational Mode Summary")
+                            st.dataframe(df_modes, use_container_width=True)
+
+                            # Store in results dictionary
+                            results["Vibrational Modes"] = df_modes.to_dict(orient="records")
+
+                            # Histogram plot of vibrational frequencies
+                            st.write("### Frequency Distribution Histogram")
+                            fig, ax = plt.subplots()
+                            ax.hist(freqs_cm, bins=30, color='skyblue', edgecolor='black')
+                            ax.set_xlabel("Frequency (cm⁻¹)")
+                            ax.set_ylabel("Number of Modes")
+                            ax.set_title("Distribution of Vibrational Frequencies")
+                            st.pyplot(fig)
+
+                            # CSV download
+                            csv_buffer = io.StringIO()
+                            df_modes.to_csv(csv_buffer, index=False)
+                            st.download_button(
+                                label="Download Vibrational Frequencies (CSV)",
+                                data=csv_buffer.getvalue(),
+                                file_name="vibrational_modes.csv",
+                                mime="text/csv"
+                            )
             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).")
@@ -1499,4 +1479,4 @@ with st.expander('ℹ️ About This App & Foundational MLIPs'):
     - 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, SevenNet, ORB and ❤️")
-st.markdown("Developed by [Manas Sharma](https://manas.bragitoff.com/) in the groups of [Prof. Ananth Govind Rajan Group](https://www.agrgroup.org/) and [Prof. Sudeep Punnathanam](https://chemeng.iisc.ac.in/sudeep/) at  [IISc Bangalore](https://iisc.ac.in/)")
+st.markdown("Developed by [Manas Sharma](https://manas.bragitoff.com/) in the groups of [Prof. Ananth Govind Rajan Group](https://www.agrgroup.org/) and [Prof. Sudeep Punnathanam](https://chemeng.iisc.ac.in/sudeep/) at  [IISc Bangalore](https://iisc.ac.in/)")