update app to include latest features available locally

src/streamlit_app.py

@@ -905,10 +905,14 @@ MACE_MODELS = {
     "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",
+    "MACE ANI-CC Large (500k)": "https://github.com/ACEsuit/mace/raw/main/mace/calculators/foundations_models/ani500k_large_CC.model",
+    "MACE OMOL-0 XL 4M": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_omol_0/mace-omol-0-extra-large-4M.model",
+    "MACE OMOL-0 XL 1024": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_omol_0/MACE-omol-0-extra-large-1024.model"
 }
 
 FAIRCHEM_MODELS = {
-    "UMA Small": "uma-s-1",
+    "UMA Small 1": "uma-s-1",
+    "UMA Small 1.1": "uma-s-1p1",
     "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"
@@ -936,8 +940,8 @@ SEVEN_NET_MODELS = {
     "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)
+def get_mace_model(model_path, dispersion, device, selected_default_dtype):
+    return mace_mp(model=model_path, dispersion=True, 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
@@ -999,7 +1003,7 @@ 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)
+    atoms.info["spin"] = atoms.info.get("spin", 1) # Default spin (usually 2S for ASE, model might want 2S+1)
 
 
 st.sidebar.markdown("## Model Selection")
@@ -1016,6 +1020,16 @@ if model_type == "MACE":
         st.sidebar.warning("Using model under Academic Software License (ASL).")
     # selected_default_dtype = st.sidebar.selectbox("Select Precision (default_dtype):", ['float32', 'float64'])
     selected_default_dtype = 'float64'
+    dispersion = st.sidebar.checkbox("Dispersion correction?", value=False)
+    if selected_model == "MACE OMOL-0 XL 4M" or selected_model == "MACE OMOL-0 XL 1024":
+        
+        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=1, value=int(atoms.info.get("spin",0) 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
+    # else:
+    #     atoms.info["charge"] = 0
+    #     atoms.info["spin"] = 1 # FairChem expects multiplicity
 if model_type == "FairChem":
     selected_model = st.sidebar.selectbox("Select FairChem Model:", list(FAIRCHEM_MODELS.keys()))
     model_path = FAIRCHEM_MODELS[selected_model]
@@ -1024,7 +1038,7 @@ if model_type == "FairChem":
         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
+            spin_multiplicity = st.sidebar.number_input("Spin Multiplicity (2S + 1)", min_value=1, max_value=11, step=1, value=int(atoms.info.get("spin",0) 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":
@@ -1061,7 +1075,9 @@ task = st.sidebar.selectbox("Select Calculation Task:",
                             "Atomization/Cohesive Energy", # New Task Added
                             "Geometry Optimization", 
                             "Cell + Geometry Optimization",
-                            "Vibrational Mode Analysis"])
+                            "Vibrational Mode Analysis",
+                            #"Phonons"
+                           ])
 
 if "Optimization" in task:
     st.sidebar.markdown("### Optimization Parameters")
@@ -1100,6 +1116,8 @@ if atoms is not None:
         st.write(f"**Model:** {selected_model}")
         if model_type == "FairChem" and selected_model == "UMA Small":
             st.write(f"**UMA Task Type:** {selected_task_type}")
+        if model_type == "MACE":
+            st.write(f"**Dispersion:** {dispersion}")
         st.write(f"**Device:** {device}")
         
         st.markdown("### Selected Task")
@@ -1129,7 +1147,7 @@ if atoms is not None:
                     
                     if model_type == "MACE":
                         # st.write("Setting up MACE calculator...")
-                        calc = get_mace_model(model_path, device, 'float32')
+                        calc = get_mace_model(model_path, dispersion, device, 'float32')
                     elif model_type == "FairChem":  # FairChem
                         # st.write("Setting up FairChem calculator...")
                         # Workaround for potential dtype issues when switching models
@@ -1377,6 +1395,15 @@ if atoms is not None:
                         
                         show_trajectory_and_controls()
                     elif task == "Vibrational Mode Analysis":
+                        # Conversion factors
+
+                        from ase.units import kB as kB_eVK, _Nav, J  # ASE's constants
+
+                        from scipy.constants import physical_constants
+
+                        kB_JK = physical_constants["Boltzmann constant"][0]  # J/K
+
+                        is_periodic = any(calc_atoms.pbc)
 
                         st.write("Running vibrational mode analysis using finite differences...")
 
@@ -1391,9 +1418,21 @@ if atoms is not None:
                             with st.spinner("Calculating vibrational modes... This may take a few minutes."):
                                 vib.run()
                                 freqs = vib.get_frequencies()
+                                energies = vib.get_energies()
+
+                                
+
+                                print('\n\n\n\n\n\n\n\n')
+                                # vib.get_hessian_2d()
+                                # st.write(vib.summary())
+                                # print('\n')
+                                # vib.tabulate()
+
+                            
+
+                            freqs_cm = freqs
+                            freqs_eV = energies
                             
-                            # Convert frequencies to cm⁻¹
-                            freqs_cm = freqs #/ cm
 
                             # Classify frequencies
                             mode_data = []
@@ -1443,6 +1482,86 @@ if atoms is not None:
                                 file_name="vibrational_modes.csv",
                                 mime="text/csv"
                             )
+                            # -------- Thermodynamic Analysis for Molecules --------
+                        if not is_periodic:
+
+                            # Filter physical frequencies > 1 cm⁻¹ (to avoid numerical issues)
+                            physical_freqs_eV = np.array([f for f in freqs_eV if f > 1e-5])
+
+                            # Zero-point vibrational energy (ZPE)
+                            ZPE = 0.5 * np.sum(physical_freqs_eV)  # in eV
+
+                            # Vibrational entropy (in eV/K)
+                            vib_entropy = 0.0
+                            for f in physical_freqs_eV:
+                                x = f / (kB_eVK * T)
+                                vib_entropy += (x / (np.exp(x) - 1) - np.log(1 - np.exp(-x)))
+
+                            S_vib_eVK = kB_eVK * vib_entropy  # eV/K
+                            S_vib_JmolK = S_vib_eVK * J * _Nav  # J/mol·K
+
+                            results["ZPE (eV)"] = ZPE.real
+                            results["Vibrational Entropy (eV/K)"] = S_vib_eVK
+                            results["Vibrational Entropy (J/mol·K)"] = S_vib_JmolK
+
+                            st.write(f"**Zero-point vibrational energy (ZPE)**: {ZPE.real:.6f} eV")
+                            st.write(f"**Vibrational entropy**: {S_vib_eVK:.6f} eV/K")
+
+                        else:
+                            st.info("Thermodynamic properties like ZPE and entropy are currently only meaningful for isolated molecules (non-periodic systems).")
+                    elif task == "Phonons":
+                        from ase.phonons import Phonons
+
+                        st.write("### Phonon Band Structure and Density of States")
+                        is_periodic = any(calc_atoms.pbc)
+
+                        if not is_periodic:
+                            st.error("Phonon calculations require a periodic structure. Please use a periodic system.")
+                        else:
+                            with tempfile.TemporaryDirectory() as tmpdir:
+                                st.info("Running phonon calculation using finite displacements...")
+                                sc = (7, 7, 7)
+
+                                # Create phonon object
+                                ph = Phonons(calc_atoms, calc_atoms.calc, supercell=sc, delta=0.001, name=os.path.join(tmpdir, 'phonon'))
+
+                                with st.spinner("Displacing atoms and computing forces..."):
+                                    ph.run()
+
+                                # Build dynamical matrix
+                                ph.read(acoustic=True)
+                                ph.clean()
+
+                                # Band path and DOS
+                                # path = calc_atoms.cell.bandpath('GXULGK', npoints=100)
+                                path = calc_atoms.cell.bandpath('GXKGL', npoints=100)
+                                # path = calc_atoms.cell.bandpath(eps=0.00001)
+                                bs = ph.get_band_structure(path)
+                                dos = ph.get_dos(kpts=(20, 20, 20)).sample_grid(npts=100, width=1e-3)
+
+                                # Plotting
+                                fig = plt.figure(figsize=(7, 4))
+                                ax = fig.add_axes([0.12, 0.07, 0.67, 0.85])
+
+                                emax = 0.075
+                                bs.plot(ax=ax, emin=0.0, emax=emax)
+
+                                dosax = fig.add_axes([0.8, 0.07, 0.17, 0.85])
+                                dosax.fill_between(
+                                    dos.get_weights(),
+                                    dos.get_energies(),
+                                    y2=0,
+                                    color='grey',
+                                    edgecolor='k',
+                                    lw=1,
+                                )
+                                dosax.set_ylim(0, emax)
+                                dosax.set_yticks([])
+                                dosax.set_xticks([])
+                                dosax.set_xlabel('DOS', fontsize=14)
+
+                                st.pyplot(fig)
+                                st.success("Phonon band structure and DOS successfully plotted.")
             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).")