Remove atom numbers, show only element symbols and add orbital visualization toggle

app.py

@@ -50,7 +50,7 @@ def name_to_smiles(name: str) -> str:
         raise gr.Error(f"Error converting name to SMILES: {str(e)}")
 
 
-def name_to_3d_molecule(name: str) -> tuple:
+def name_to_3d_molecule(name: str, show_orbitals: bool = False) -> tuple:
     """Convert chemical name to 3D molecule SDF and 2D visualization"""
     try:
         # Convert name to SMILES
@@ -115,7 +115,7 @@ def name_to_3d_molecule(name: str) -> tuple:
             z_coords.append(pos.z)
             element = atom.GetSymbol()
             elements.append(element)
-            atom_labels.append(f"{element}{i}")
+            atom_labels.append(element)  # Only element symbol, no numbers
         
         # Define color scheme for common elements (Jmol colors)
         color_map = {
@@ -162,8 +162,43 @@ def name_to_3d_molecule(name: str) -> tuple:
             hoverinfo='skip'
         )
         
+        # Create orbital clouds if requested
+        data_traces = [bonds_trace, atoms_trace]
+        
+        if show_orbitals:
+            import numpy as np
+            # Add semi-transparent spheres around atoms to represent electron orbitals
+            orbital_radius = {
+                'H': 0.6, 'C': 0.9, 'N': 0.8, 'O': 0.75,
+                'F': 0.7, 'Cl': 1.0, 'Br': 1.15, 'I': 1.4,
+                'P': 1.1, 'S': 1.05, 'B': 0.95, 'Si': 1.2
+            }
+            
+            # Create mesh for each atom's orbital
+            for i, (x, y, z, elem) in enumerate(zip(x_coords, y_coords, z_coords, elements)):
+                radius = orbital_radius.get(elem, 0.8)
+                
+                # Create sphere mesh
+                u = np.linspace(0, 2 * np.pi, 20)
+                v = np.linspace(0, np.pi, 15)
+                sphere_x = x + radius * np.outer(np.cos(u), np.sin(v))
+                sphere_y = y + radius * np.outer(np.sin(u), np.sin(v))
+                sphere_z = z + radius * np.outer(np.ones(np.size(u)), np.cos(v))
+                
+                color = color_map.get(elem, '#FF1493')
+                
+                orbital_trace = go.Surface(
+                    x=sphere_x, y=sphere_y, z=sphere_z,
+                    colorscale=[[0, color], [1, color]],
+                    showscale=False,
+                    opacity=0.15,
+                    name=f'{elem} orbital',
+                    hoverinfo='skip'
+                )
+                data_traces.append(orbital_trace)
+        
         # Create figure
-        fig = go.Figure(data=[bonds_trace, atoms_trace])
+        fig = go.Figure(data=data_traces)
         
         fig.update_layout(
             title=dict(text=f"{name}<br>SMILES: {canonical_smiles}", x=0.5, xanchor='center'),
@@ -236,16 +271,19 @@ tpsa_interface = gr.Interface(
 
 molecule_3d_interface = gr.Interface(
     fn=name_to_3d_molecule,
-    inputs=gr.Textbox(label="Chemical Name", placeholder="e.g., benzene, aspirin, caffeine, glucose"),
+    inputs=[
+        gr.Textbox(label="Chemical Name", placeholder="e.g., benzene, aspirin, caffeine, glucose"),
+        gr.Checkbox(label="Show Electron Orbitals", value=False)
+    ],
     outputs=[
         gr.HTML(label="2D Structure"),
         gr.Plot(label="3D Interactive Molecule Viewer - Rotate and Zoom!"),
         gr.Textbox(label="3D SDF Content (Optional - for external viewers)", lines=10, max_lines=20, visible=False)
     ],
     api_name="name_to_molecule",
-    description="View 2D structure and interactive 3D molecule with atom labels. Click and drag to rotate, scroll to zoom!",
+    description="View 2D structure and interactive 3D molecule with atom labels. Click and drag to rotate, scroll to zoom! Toggle 'Show Electron Orbitals' to visualize electron clouds around atoms.",
     cache_examples=False,
-    examples=[["benzene"], ["aspirin"], ["caffeine"], ["glucose"]],
+    examples=[["benzene", False], ["aspirin", False], ["caffeine", True], ["glucose", False]],
 )