Replace 3Dmol.js with Plotly for interactive 3D molecule viewer with atom labels

app.py

@@ -80,14 +80,13 @@ def name_to_3d_molecule(name: str) -> tuple:
         # Import needed modules
         import io
         import base64
-        import json
         from rdkit.Chem.Draw import rdMolDraw2D
         
         # Generate 2D image for preview
         mol_2d = Chem.MolFromSmiles(smiles)  # Get 2D version
         
         # Create drawer for 2D structure
-        drawer = rdMolDraw2D.MolDraw2DCairo(500, 500)
+        drawer = rdMolDraw2D.MolDraw2DCairo(600, 600)
         drawer.DrawMolecule(mol_2d)
         drawer.FinishDrawing()
         img_2d_data = drawer.GetDrawingText()
@@ -96,81 +95,99 @@ def name_to_3d_molecule(name: str) -> tuple:
         # Get SMILES string
         canonical_smiles = Chem.MolToSmiles(mol_2d)
         
-        # Escape SDF for JavaScript
-        sdf_json = json.dumps(sdf_content)
-        viewer_id = f"viewer_{abs(hash(name))}"
-
-        # Create HTML with 3Dmol.js viewer
-        html_content = f"""
-        <div style="padding: 20px;">
-            <h3 style="text-align: center;">{name}</h3>
-            <p style="text-align: center;"><strong>SMILES:</strong> {canonical_smiles}</p>
-            
-            <div style="display: flex; gap: 20px; justify-content: center; flex-wrap: wrap; margin-top: 20px;">
-                <div style="text-align: center;">
-                    <h4>2D Structure</h4>
-                    <img src="data:image/png;base64,{img_2d_str}" style="max-width: 500px; width: 100%; border: 1px solid #ddd; border-radius: 8px; padding: 10px; background: white;">
-                </div>
-                
-                <div style="text-align: center;">
-                    <h4>3D Interactive Viewer</h4>
-                    <div id="{viewer_id}" style="width: 500px; height: 500px; border: 1px solid #ddd; border-radius: 8px; background: white; position: relative;"></div>
-                </div>
-            </div>
-        </div>
+        # Create 3D visualization using Plotly
+        import plotly.graph_objects as go
+        
+        # Get atom positions and elements
+        conf = mol.GetConformer()
+        atoms = mol.GetAtoms()
+        
+        x_coords = []
+        y_coords = []
+        z_coords = []
+        elements = []
+        atom_labels = []
+        
+        for i, atom in enumerate(atoms):
+            pos = conf.GetAtomPosition(i)
+            x_coords.append(pos.x)
+            y_coords.append(pos.y)
+            z_coords.append(pos.z)
+            element = atom.GetSymbol()
+            elements.append(element)
+            atom_labels.append(f"{element}{i}")
+        
+        # Define color scheme for common elements (Jmol colors)
+        color_map = {
+            'H': '#FFFFFF', 'C': '#909090', 'N': '#3050F8', 'O': '#FF0D0D',
+            'F': '#90E050', 'Cl': '#1FF01F', 'Br': '#A62929', 'I': '#940094',
+            'P': '#FF8000', 'S': '#FFFF30', 'B': '#FFB5B5', 'Si': '#F0C8A0'
+        }
+        
+        colors = [color_map.get(e, '#FF1493') for e in elements]
+        
+        # Create 3D scatter plot for atoms with labels
+        atoms_trace = go.Scatter3d(
+            x=x_coords, y=y_coords, z=z_coords,
+            mode='markers+text',
+            marker=dict(size=10, color=colors, line=dict(color='black', width=1)),
+            text=atom_labels,
+            textposition='top center',
+            textfont=dict(size=10, color='black'),
+            name='Atoms',
+            hovertext=[f"{lbl}<br>Pos: ({x:.2f}, {y:.2f}, {z:.2f})" 
+                      for lbl, x, y, z in zip(atom_labels, x_coords, y_coords, z_coords)],
+            hoverinfo='text'
+        )
         
-        <script src="https://code.jquery.com/jquery-3.6.0.min.js"></script>
-        <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
-        <script>
-            (function() {{
-                var sdfData = {sdf_json};
-                var viewerId = "{viewer_id}";
-                
-                function init3DMol() {{
-                    try {{
-                        var element = document.getElementById(viewerId);
-                        if (!element) {{
-                            console.error('Viewer element not found:', viewerId);
-                            return;
-                        }}
-                        
-                        var config = {{ backgroundColor: 'white' }};
-                        var viewer = $3Dmol.createViewer(element, config);
-                        
-                        viewer.addModel(sdfData, "sdf");
-                        viewer.setStyle({{}}, {{
-                            stick: {{ radius: 0.15, colorscheme: 'Jmol' }},
-                            sphere: {{ radius: 0.4, colorscheme: 'Jmol' }}
-                        }});
-                        
-                        viewer.addLabels({{ elem: true }}, {{
-                            fontSize: 12,
-                            fontColor: 'black',
-                            backgroundColor: 'white',
-                            backgroundOpacity: 0.7
-                        }});
-                        
-                        viewer.zoomTo();
-                        viewer.render();
-                        viewer.zoom(1.2);
-                        
-                        console.log('3DMol viewer initialized successfully');
-                    }} catch (e) {{
-                        console.error('Error initializing 3DMol viewer:', e);
-                    }}
-                }}
-                
-                // Wait for both jQuery and 3Dmol to load
-                if (typeof $3Dmol !== 'undefined') {{
-                    init3DMol();
-                }} else {{
-                    setTimeout(init3DMol, 1000);
-                }}
-            }})();
-        </script>
+        # Create bonds
+        bond_x = []
+        bond_y = []
+        bond_z = []
+        
+        for bond in mol.GetBonds():
+            i = bond.GetBeginAtomIdx()
+            j = bond.GetEndAtomIdx()
+            pos_i = conf.GetAtomPosition(i)
+            pos_j = conf.GetAtomPosition(j)
+            bond_x.extend([pos_i.x, pos_j.x, None])
+            bond_y.extend([pos_i.y, pos_j.y, None])
+            bond_z.extend([pos_i.z, pos_j.z, None])
+        
+        bonds_trace = go.Scatter3d(
+            x=bond_x, y=bond_y, z=bond_z,
+            mode='lines',
+            line=dict(color='#888888', width=4),
+            name='Bonds',
+            hoverinfo='skip'
+        )
+        
+        # Create figure
+        fig = go.Figure(data=[bonds_trace, atoms_trace])
+        
+        fig.update_layout(
+            title=dict(text=f"{name}<br>SMILES: {canonical_smiles}", x=0.5, xanchor='center'),
+            scene=dict(
+                xaxis=dict(showbackground=False, showgrid=True, zeroline=False, showticklabels=False, title=''),
+                yaxis=dict(showbackground=False, showgrid=True, zeroline=False, showticklabels=False, title=''),
+                zaxis=dict(showbackground=False, showgrid=True, zeroline=False, showticklabels=False, title=''),
+                bgcolor='white'
+            ),
+            showlegend=False,
+            width=800,
+            height=700,
+            margin=dict(l=0, r=0, t=50, b=0)
+        )
+        
+        # Generate 2D structure HTML
+        html_2d = f"""
+        <div style="text-align: center; padding: 20px;">
+            <h4>2D Structure</h4>
+            <img src="data:image/png;base64,{img_2d_str}" style="max-width: 600px; width: 100%; border: 1px solid #ddd; border-radius: 8px; padding: 10px; background: white;">
+        </div>
         """
 
-        return html_content, sdf_content
+        return html_2d, fig, sdf_content
 
     except Exception as e:
         raise gr.Error(f"Error creating molecule: {str(e)}")
@@ -221,12 +238,14 @@ molecule_3d_interface = gr.Interface(
     fn=name_to_3d_molecule,
     inputs=gr.Textbox(label="Chemical Name", placeholder="e.g., benzene, aspirin, caffeine, glucose"),
     outputs=[
-        gr.HTML(label="2D and 3D Interactive Molecular Viewer"),
+        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 viewer with atom labels. Rotate and zoom the 3D structure!",
+    description="View 2D structure and interactive 3D molecule with atom labels. Click and drag to rotate, scroll to zoom!",
     cache_examples=False,
+    examples=[["benzene"], ["aspirin"], ["caffeine"], ["glucose"]],
 )
 
 

requirements.txt

@@ -2,4 +2,5 @@
 rdkit
 gradio==4.44.1
 huggingface_hub==0.19.4
-cirpy
+cirpy
+plotly