add orbmol model and fix issues with orb models

src/streamlit_app.py

@@ -30,8 +30,11 @@ import sys
 import pkg_resources
 from ase.vibrations import Vibrations
 from mp_api.client import MPRester
+import pubchempy as pcp
+from io import StringIO
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.core.structure import Molecule
 import matplotlib.pyplot as plt
 mattersim_available = True
 if mattersim_available:
@@ -926,14 +929,16 @@ FAIRCHEM_MODELS = {
 }
 # 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",
+    "V3 OMOL Conservative": pretrained.orb_v3_conservative_omol,
+    "V3 OMOL Direct": pretrained.orb_v3_direct_omol,
+    "V3 OMAT Conservative (inf)": pretrained.orb_v3_conservative_inf_omat,
+    "V3 OMAT Conservative (20)": pretrained.orb_v3_conservative_20_omat,
+    "V3 OMAT Direct (inf)": pretrained.orb_v3_direct_inf_omat,
+    "V3 OMAT Direct (20)": pretrained.orb_v3_direct_20_omat,
+    "V3 MPA Conservative (inf)": pretrained.orb_v3_conservative_inf_mpa,
+    "V3 MPA Conservative (20)": pretrained.orb_v3_conservative_20_mpa,
+    "V3 MPA Direct (inf)": pretrained.orb_v3_direct_inf_mpa,
+    "V3 MPA Direct (20)": pretrained.orb_v3_direct_20_mpa,
 }
 # Define the available MatterSim models
 MATTERSIM_MODELS = {
@@ -959,117 +964,11 @@ def get_fairchem_model(selected_model_name, model_path_or_name, device, selected
         calc = FAIRChemCalculator(predictor, task_name="omol")
     return calc
 
-# st.sidebar.markdown("## Input Options")
-# input_method = st.sidebar.radio("Choose Input Method:", 
-#                                 ["Select Example", "Upload File", "Paste Content", "Materials Project ID"])
-# 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)
-# elif input_method == "Materials Project ID":
-#     # st.sidebar.info("Requires an API Key from [Materials Project](https://next-gen.materialsproject.org/api)")
-    
-#     # 1. Get API Key (You can also set this as an env variable)
-#     # mp_api_key = st.sidebar.text_input("MP API Key", type="password")
-#     mp_api_key = os.getenv("MP_API_KEY")
-#     # 2. Get Material ID
-#     material_id = st.sidebar.text_input("Enter Material ID:", value="mp-149")
-    
-#     # 3. Choose Cell Type
-#     cell_type = st.sidebar.radio("Unit Cell Type:", ['Primitive Cell', 'Conventional Unit Cell'])
-    
-#     # 1. Initialize the session state variable if it doesn't exist
-#     if "atoms" not in st.session_state:
-#         st.session_state.atoms = None
-
-#     if st.sidebar.button("Fetch Structure"):
-#         if not mp_api_key:
-#             st.sidebar.error("Please enter your Materials Project API Key.")
-#         elif not material_id:
-#             st.sidebar.error("Please enter a Material ID.")
-#         else:
-#             try:
-#                 with st.spinner(f"Fetching {material_id}..."):
-#                     # Connect to MP
-#                     with MPRester(mp_api_key) as mpr:
-#                         # Get Pymatgen Structure
-#                         pmg_structure = mpr.get_structure_by_material_id(material_id)
-                        
-#                         # Handle Symmetry (Primitive vs Conventional)
-#                         analyzer = SpacegroupAnalyzer(pmg_structure)
-                        
-#                         if cell_type == 'Conventional Unit Cell':
-#                             final_structure = analyzer.get_conventional_standard_structure()
-#                         else:
-#                             final_structure = analyzer.get_primitive_standard_structure()
-                            
-#                         # Convert Pymatgen Structure -> ASE Atoms
-#                         # 2. Store the result in session_state instead of a local variable
-#                         st.session_state.atoms = AseAtomsAdaptor.get_atoms(final_structure)
-                        
-#                         st.sidebar.success(f"Loaded {material_id} ({cell_type}) with {len(st.session_state.atoms)} atoms.")
-                        
-#             except Exception as e:
-#                 st.sidebar.error(f"Error fetching data: {str(e)}")
-#                 st.session_state.atoms = None
-
-#     # 3. Retrieve the data from session_state for use in the rest of your app
-#     atoms = st.session_state.atoms
-
-#     # --- Rest of your processing code ---
-#     # if atoms is not None:
-#     #     st.write(f"System is ready. Formula: {atoms.get_chemical_formula()}")
-#     # else:
-#     #     st.info("Please fetch a structure to begin.")
-# 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", 1) # Default spin (usually 2S for ASE, model might want 2S+1)
-
 # --- INITIALIZATION (Must be run first) ---
 if "atoms" not in st.session_state:
     st.session_state.atoms = None
+if "atoms_list" not in st.session_state:
+    st.session_state.atoms_list = []
 
 # Reset atoms state if input method changes, to prevent using old data
 # Use a key to track the currently active input method
@@ -1078,7 +977,7 @@ if 'current_input_method' not in st.session_state:
 
 st.sidebar.markdown("## Input Options")
 input_method = st.sidebar.radio("Choose Input Method:", 
-                                ["Select Example", "Upload File", "Paste Content", "Materials Project ID"])
+                                ["Select Example", "Upload File", "Paste Content", "Materials Project ID", "PubChem", "Batch Upload"])
 
 # If the input method changes, clear the loaded structure
 if input_method != st.session_state.current_input_method:
@@ -1168,6 +1067,92 @@ elif input_method == "Paste Content":
         # Clear structure if text area is empty
         st.session_state.atoms = None
 
+# --- PUBCHEM SEARCH MODE ---
+elif input_method == "PubChem":
+    
+
+    st.sidebar.markdown("### Search PubChem")
+
+    query = st.sidebar.text_input("Enter name or formula (e.g., H2O, water, methane):", 
+                                  key="pubchem_query", value="water")
+
+    # Reset atoms if no query
+    if query.strip() == "":
+        st.session_state.atoms = None
+
+    # Step 1: Search PubChem
+    if query and query.strip():
+        # Avoid re-searching if query is unchanged
+        if "pubchem_last_query" not in st.session_state or st.session_state.pubchem_last_query != query:
+            try:
+                with st.spinner("Searching PubChem..."):
+                    results = pcp.get_compounds(query, "name")  # name OR formula works
+                st.session_state.pubchem_results = results
+                st.session_state.pubchem_last_query = query
+            except Exception as e:
+                st.sidebar.error(f"Error searching PubChem: {str(e)}")
+                st.session_state.pubchem_results = None
+
+        results = st.session_state.get("pubchem_results", [])
+        if results:
+            # Convert to displayable table
+            df = pd.DataFrame(
+                [(c.cid, c.iupac_name, c.molecular_formula, c.molecular_weight, c.isomeric_smiles)
+                 for c in results],
+                columns=["CID", "Name", "Formula", "Weight", "SMILES"]
+            )
+            st.sidebar.success(f"Found {len(df)} result(s).")
+            st.sidebar.dataframe(df)
+
+            # Choose a CID
+            cid = st.sidebar.selectbox("Select CID", df["CID"], key="pubchem_cid")
+
+            # Step 2: Retrieve 3D structure for selected CID 
+            if cid:
+                if "pubchem_last_cid" not in st.session_state or st.session_state.pubchem_last_cid != cid:
+                    try:
+                        with st.spinner("Fetching 3D coordinates..."):
+                            # Function to format floating-point numbers with alignment
+                            def format_number(num, width=10, precision=5):
+                                # Handles positive/negative numbers while maintaining alignment
+                                return f"{num: {width}.{precision}f}"
+                            # CID to XYZ
+                            def generate_xyz_coordinates(cid):
+                                compound = pcp.Compound.from_cid(cid, record_type='3d')
+                                atoms = compound.atoms
+                                coords = [(atom.x, atom.y, atom.z) for atom in atoms]
+
+                                num_atoms = len(atoms)
+                                xyz_text = f"{num_atoms}\n{compound.cid}\n"
+
+                                for atom, coord in zip(atoms, coords):
+                                    atom_symbol = atom.element
+                                    x, y, z = coord
+                                    xyz_text += f"{atom_symbol} {format_number(x, precision=8)} {format_number(y, precision=8)} {format_number(z, precision=8)}\n"
+
+                                return xyz_text
+                            def get_molecule(cid):
+                                xyz_str = generate_xyz_coordinates(cid)
+                                return Molecule.from_str(xyz_str, fmt='xyz'), xyz_str
+                            # Fetch SDF with 3D conformer
+                            # sdf_str = pcp.Compound.from_cid(int(cid)).to_sdf()
+                            selected_molecule, xyz_str = get_molecule(cid)
+
+                        # Convert SDF → ASE Atoms using temporary memory buffer
+                        atoms_to_store = read(StringIO(xyz_str), format="xyz")
+
+                        atoms_to_store.info["source_name"] = f"PubChem CID {cid}"
+                        st.session_state.atoms = atoms_to_store
+                        st.session_state.pubchem_last_cid = cid
+
+                        st.sidebar.success(f"Loaded PubChem structure with {len(atoms_to_store)} atoms!")
+
+                    except Exception as e:
+                        st.sidebar.error(f"Unable to retrieve 3D structure: {str(e)}")
+                        st.session_state.atoms = None
+                        st.session_state.pubchem_last_cid = None
+        else:
+            st.sidebar.info("No PubChem results found.")
 
 # --- MATERIALS PROJECT ID ---
 elif input_method == "Materials Project ID":
@@ -1210,6 +1195,52 @@ elif input_method == "Materials Project ID":
     elif not material_id:
         st.sidebar.error("Please enter a Material ID.")
 
+# --- BATCH UPLOAD MULTIPLE FILES ---
+elif input_method == "Batch Upload":
+
+    uploaded_files = st.sidebar.file_uploader(
+        "Upload multiple structure files",
+        type=["xyz", "cif", "POSCAR", "vasp", "CONTCAR", "mol", "sdf", "tmol", "extxyz"],
+        accept_multiple_files=True
+    )
+
+    # Clear state if no files present
+    if not uploaded_files:
+        st.session_state.atoms_list = []
+        st.session_state.atoms = None
+
+    else:
+        atoms_list = []
+        errors = []
+
+        for file in uploaded_files:
+            try:
+                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.name)[1]) as tmp:
+                    tmp.write(file.getvalue())
+                    tmp_path = tmp.name
+
+                atoms_obj = read(tmp_path)
+                atoms_obj.info["source_name"] = file.name
+                atoms_list.append(atoms_obj)
+
+            except Exception as e:
+                errors.append(f"{file.name}: {str(e)}")
+
+            finally:
+                if "tmp_path" in locals() and os.path.exists(tmp_path):
+                    os.unlink(tmp_path)
+
+        # Store everything only if at least one success
+        if atoms_list:
+            st.session_state.atoms_list = atoms_list
+            st.session_state.atoms = atoms_list[0]  # default: first item
+            st.sidebar.success(f"Loaded {len(atoms_list)} structures successfully!")
+
+            if len(atoms_list) > 1:
+                st.sidebar.info("You can now process them as a batch.")
+
+        if errors:
+            st.sidebar.error("Some files could not be loaded:\n" + "\n".join(errors))
 # ----------------------------------------------------
 # --- FINAL STRUCTURE RETRIEVAL (The persistent structure) ---
 # ----------------------------------------------------
@@ -1224,8 +1255,6 @@ if atoms is not None:
     
     # Display confirmation in the main area (optional, helps the user confirm what's loaded)
     st.markdown(f"**Loaded Structure:** {atoms.get_chemical_formula()} ({len(atoms)} atoms)")
-    
-    # You can now add your processing buttons here, which will consistently use the 'atoms' variable.
 
 st.sidebar.markdown("## Model Selection")
 if mattersim_available:
@@ -1409,7 +1438,8 @@ if atoms is not None:
                         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)
+                        # orbff = pretrained.orb_v3_conservative_inf_omat(device=device, precision=selected_default_dtype)
+                        orbff = model_path(device=device, precision=selected_default_dtype)
                         calc = ORBCalculator(orbff, device=device)
                     elif model_type == "MatterSim":
                         # st.write("Setting up MatterSim calculator...")