Add proper orbital shapes: s (sphere), p (dumbbell), d (cloverleaf) with better checkbox visibility

app.py

@@ -167,35 +167,141 @@ def name_to_3d_molecule(name: str, show_orbitals: bool = False) -> tuple:
         
         if show_orbitals:
             import numpy as np
-            # Add semi-transparent spheres around atoms to represent electron orbitals
+            
+            # Orbital size reference
             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))
+            # Get electron configuration for each atom to determine orbital types
+            def get_orbitals(atomic_num):
+                """Return list of orbitals for an atom based on electron configuration"""
+                # Simplified orbital visualization - show valence orbitals
+                if atomic_num == 1:  # H: 1s
+                    return ['s']
+                elif atomic_num <= 2:  # He: 1s
+                    return ['s']
+                elif atomic_num <= 10:  # Li-Ne: has s and p
+                    return ['s', 'p']
+                elif atomic_num <= 18:  # Na-Ar: has s and p
+                    return ['s', 'p']
+                elif atomic_num <= 36:  # K-Kr: has s, p, d
+                    return ['s', 'p', 'd']
+                else:  # Rb onwards: has s, p, d, f
+                    return ['s', 'p', 'd', 'f']
+            
+            # Create orbital shapes around each atom
+            for idx, (x, y, z, elem) in enumerate(zip(x_coords, y_coords, z_coords, elements)):
+                atom = atoms[idx]
+                atomic_num = atom.GetAtomicNum()
+                orbitals = get_orbitals(atomic_num)
                 
+                # Base radius for orbitals
+                base_radius = orbital_radius.get(elem, 0.8)
                 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)
+                # S orbital (spherical)
+                if 's' in orbitals:
+                    u = np.linspace(0, 2 * np.pi, 25)
+                    v = np.linspace(0, np.pi, 20)
+                    s_radius = base_radius * 0.6
+                    sphere_x = x + s_radius * np.outer(np.cos(u), np.sin(v))
+                    sphere_y = y + s_radius * np.outer(np.sin(u), np.sin(v))
+                    sphere_z = z + s_radius * np.outer(np.ones(np.size(u)), np.cos(v))
+                    
+                    s_orbital = go.Surface(
+                        x=sphere_x, y=sphere_y, z=sphere_z,
+                        colorscale=[[0, color], [1, color]],
+                        showscale=False,
+                        opacity=0.2,
+                        name=f'{elem}-s',
+                        hoverinfo='skip'
+                    )
+                    data_traces.append(s_orbital)
+                
+                # P orbitals (dumbbell shaped - 3 orientations: px, py, pz)
+                if 'p' in orbitals:
+                    p_radius = base_radius * 1.0
+                    p_length = base_radius * 1.5
+                    
+                    # Create dumbbell shape for p orbitals
+                    u = np.linspace(0, 2 * np.pi, 20)
+                    v = np.linspace(-1, 1, 15)
+                    
+                    # Pz orbital (along z-axis)
+                    pz_r = p_radius * np.sqrt(1 - v**2)
+                    pz_z_vals = z + p_length * v
+                    pz_x = x + np.outer(pz_r, np.cos(u))
+                    pz_y = y + np.outer(pz_r, np.sin(u))
+                    pz_z = np.outer(pz_z_vals, np.ones(len(u)))
+                    
+                    pz_orbital = go.Surface(
+                        x=pz_x, y=pz_y, z=pz_z,
+                        colorscale=[[0, color], [0.5, color], [1, color]],
+                        showscale=False,
+                        opacity=0.15,
+                        name=f'{elem}-pz',
+                        hoverinfo='skip'
+                    )
+                    data_traces.append(pz_orbital)
+                    
+                    # Px orbital (along x-axis)
+                    px_r = p_radius * np.sqrt(1 - v**2)
+                    px_x_vals = x + p_length * v
+                    px_x = np.outer(px_x_vals, np.ones(len(u)))
+                    px_y = y + np.outer(px_r, np.cos(u))
+                    px_z = z + np.outer(px_r, np.sin(u))
+                    
+                    px_orbital = go.Surface(
+                        x=px_x, y=px_y, z=px_z,
+                        colorscale=[[0, color], [0.5, color], [1, color]],
+                        showscale=False,
+                        opacity=0.15,
+                        name=f'{elem}-px',
+                        hoverinfo='skip'
+                    )
+                    data_traces.append(px_orbital)
+                    
+                    # Py orbital (along y-axis)
+                    py_r = p_radius * np.sqrt(1 - v**2)
+                    py_y_vals = y + p_length * v
+                    py_x = x + np.outer(py_r, np.cos(u))
+                    py_y = np.outer(py_y_vals, np.ones(len(u)))
+                    py_z = z + np.outer(py_r, np.sin(u))
+                    
+                    py_orbital = go.Surface(
+                        x=py_x, y=py_y, z=py_z,
+                        colorscale=[[0, color], [0.5, color], [1, color]],
+                        showscale=False,
+                        opacity=0.15,
+                        name=f'{elem}-py',
+                        hoverinfo='skip'
+                    )
+                    data_traces.append(py_orbital)
+                
+                # D orbitals (for transition metals) - simplified cloverleaf shape
+                if 'd' in orbitals and elem not in ['H', 'C', 'N', 'O', 'F', 'S', 'P', 'Cl']:
+                    d_radius = base_radius * 1.2
+                    theta = np.linspace(0, 2 * np.pi, 25)
+                    phi = np.linspace(0, np.pi, 15)
+                    
+                    # Simplified d-orbital (four-lobed)
+                    r = d_radius * np.abs(np.outer(np.sin(2 * phi), np.cos(2 * theta)))
+                    d_x = x + r * np.outer(np.sin(phi), np.cos(theta))
+                    d_y = y + r * np.outer(np.sin(phi), np.sin(theta))
+                    d_z = z + r * np.outer(np.cos(phi), np.ones(len(theta)))
+                    
+                    d_orbital = go.Surface(
+                        x=d_x, y=d_y, z=d_z,
+                        colorscale=[[0, color], [1, color]],
+                        showscale=False,
+                        opacity=0.12,
+                        name=f'{elem}-d',
+                        hoverinfo='skip'
+                    )
+                    data_traces.append(d_orbital)
         
         # Create figure
         fig = go.Figure(data=data_traces)
@@ -273,7 +379,7 @@ molecule_3d_interface = gr.Interface(
     fn=name_to_3d_molecule,
     inputs=[
         gr.Textbox(label="Chemical Name", placeholder="e.g., benzene, aspirin, caffeine, glucose"),
-        gr.Checkbox(label="Show Electron Orbitals", value=False)
+        gr.Checkbox(label="Show Electron Orbitals (s, p, d shapes)", value=False, info="Toggle to see s-orbitals (spheres), p-orbitals (dumbbells), and d-orbitals (cloverleaf) around atoms")
     ],
     outputs=[
         gr.HTML(label="2D Structure"),
@@ -281,9 +387,16 @@ molecule_3d_interface = gr.Interface(
         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! Toggle 'Show Electron Orbitals' to visualize electron clouds around atoms.",
+    description="🔬 View 2D structure and interactive 3D molecule with atom labels. Click and drag to rotate, scroll to zoom! ✨ Enable 'Show Electron Orbitals' to see s-orbitals (spheres), p-orbitals (dumbbells), and d-orbitals (cloverleaf shapes).",
     cache_examples=False,
-    examples=[["benzene", False], ["aspirin", False], ["caffeine", True], ["glucose", False]],
+    examples=[
+        ["benzene", False], 
+        ["benzene", True],
+        ["water", True],
+        ["aspirin", False], 
+        ["caffeine", True], 
+        ["glucose", False]
+    ],
 )