Simplify to HOMO orbital visualization using Hartree-Fock for faster computation

app.py

@@ -194,114 +194,120 @@ def name_to_3d_molecule(name: str, show_orbitals: bool = False) -> tuple:
         if show_orbitals:
             try:
                 import numpy as np
-                from pyscf import gto, lo, dft
+                from pyscf import gto, scf
                 
                 # Build PySCF molecule object
                 pyscf_elements = [atom.GetSymbol() for atom in mol.GetAtoms()]
                 pyscf_coords = []
                 for i in range(mol.GetNumAtoms()):
                     pos = conf.GetAtomPosition(i)
+                    # Convert to Angstrom (PySCF uses Angstrom)
                     pyscf_coords.append([pos.x, pos.y, pos.z])
                 
                 pyscf_atoms = [(elem, coord) for elem, coord in zip(pyscf_elements, pyscf_coords)]
                 
-                # Create PySCF molecule - use small basis for speed
-                pyscf_mole = gto.Mole(basis="sto-3g")
+                # Create PySCF molecule with minimal basis
+                pyscf_mole = gto.Mole()
                 pyscf_mole.atom = pyscf_atoms
+                pyscf_mole.basis = 'sto-3g'
+                pyscf_mole.unit = 'A'
+                pyscf_mole.verbose = 0
                 pyscf_mole.build()
                 
-                # Run fast DFT calculation
-                mf = dft.RKS(pyscf_mole)
-                mf.xc = 'b3lyp'
-                mf.verbose = 0  # Suppress output
-                mf.run()
+                # Run Hartree-Fock calculation (faster than DFT)
+                mf = scf.RHF(pyscf_mole)
+                mf.verbose = 0
+                mf.kernel()
                 
-                # Calculate Natural Atomic Orbitals (pre-NAOs for more localized orbitals)
-                dm = mf.make_rdm1()
-                naos = lo.nao.prenao(pyscf_mole, dm)
+                # Get HOMO and LUMO indices
+                n_electrons = pyscf_mole.nelectron
+                n_occ = n_electrons // 2
+                homo_idx = n_occ - 1
+                lumo_idx = n_occ
                 
-                # Create grid for orbital evaluation
-                grid_resolution = 25  # Lower for speed
-                margin = 3.0
+                # Create a coarser grid for visualization
+                margin = 2.5
+                grid_size = 30
                 
-                # Determine grid bounds
-                all_x = [coord[0] for coord in pyscf_coords]
-                all_y = [coord[1] for coord in pyscf_coords]
-                all_z = [coord[2] for coord in pyscf_coords]
+                x_coords_np = np.array(x_coords)
+                y_coords_np = np.array(y_coords)
+                z_coords_np = np.array(z_coords)
                 
-                x_min, x_max = min(all_x) - margin, max(all_x) + margin
-                y_min, y_max = min(all_y) - margin, max(all_y) + margin
-                z_min, z_max = min(all_z) - margin, max(all_z) + margin
+                x_min, x_max = x_coords_np.min() - margin, x_coords_np.max() + margin
+                y_min, y_max = y_coords_np.min() - margin, y_coords_np.max() + margin
+                z_min, z_max = z_coords_np.min() - margin, z_coords_np.max() + margin
                 
-                # Create meshgrid
-                xs = np.linspace(x_min, x_max, grid_resolution)
-                ys = np.linspace(y_min, y_max, grid_resolution)
-                zs = np.linspace(z_min, z_max, grid_resolution)
-                grid_x, grid_y, grid_z = np.meshgrid(xs, ys, zs, indexing='ij')
+                x_grid = np.linspace(x_min, x_max, grid_size)
+                y_grid = np.linspace(y_min, y_max, grid_size)
+                z_grid = np.linspace(z_min, z_max, grid_size)
                 
-                # Flatten grid for evaluation
-                grid_coords = np.column_stack([grid_x.ravel(), grid_y.ravel(), grid_z.ravel()])
+                coords_grid = np.array(np.meshgrid(x_grid, y_grid, z_grid, indexing='ij'))
+                grid_points = coords_grid.reshape(3, -1).T
                 
-                # Evaluate a few representative valence orbitals
-                # Select orbitals around HOMO (Highest Occupied Molecular Orbital)
-                n_orbitals_to_show = min(3, naos.shape[1])  # Show up to 3 orbitals
-                start_orbital = max(0, naos.shape[1] - 5)  # Start near valence orbitals
+                # Evaluate HOMO orbital on grid
+                ao_value = pyscf_mole.eval_gto('GTOval_sph', grid_points)
+                mo_coeff = mf.mo_coeff[:, homo_idx]
+                homo_values = np.dot(ao_value, mo_coeff)
+                homo_grid = homo_values.reshape(grid_size, grid_size, grid_size)
                 
-                for orbital_idx in range(start_orbital, min(start_orbital + n_orbitals_to_show, naos.shape[1])):
-                    # Evaluate orbital on grid
-                    ao_values = pyscf_mole.eval_gto('GTOval_sph', grid_coords)
-                    orbital_values = np.dot(ao_values, naos[:, orbital_idx])
-                    orbital_grid = orbital_values.reshape(grid_x.shape)
+                # Create isosurfaces using marching cubes
+                try:
+                    from skimage import measure
                     
-                    # Create isosurface for positive lobe
-                    isoval_positive = 0.02
-                    try:
-                        from skimage import measure
-                        verts_pos, faces_pos, _, _ = measure.marching_cubes(orbital_grid, level=isoval_positive)
+                    # HOMO positive isosurface
+                    iso_val = 0.03
+                    if np.abs(homo_grid).max() > iso_val:
+                        verts, faces, _, _ = measure.marching_cubes(homo_grid, level=iso_val)
                         
-                        # Transform vertices to real coordinates
-                        verts_pos[:, 0] = x_min + verts_pos[:, 0] * (x_max - x_min) / (grid_resolution - 1)
-                        verts_pos[:, 1] = y_min + verts_pos[:, 1] * (y_max - y_min) / (grid_resolution - 1)
-                        verts_pos[:, 2] = z_min + verts_pos[:, 2] * (z_max - z_min) / (grid_resolution - 1)
+                        # Scale vertices to actual coordinates
+                        verts_scaled = np.zeros_like(verts)
+                        verts_scaled[:, 0] = x_min + verts[:, 0] * (x_max - x_min) / (grid_size - 1)
+                        verts_scaled[:, 1] = y_min + verts[:, 1] * (y_max - y_min) / (grid_size - 1)
+                        verts_scaled[:, 2] = z_min + verts[:, 2] * (z_max - z_min) / (grid_size - 1)
                         
-                        orbital_trace_pos = go.Mesh3d(
-                            x=verts_pos[:, 0], y=verts_pos[:, 1], z=verts_pos[:, 2],
-                            i=faces_pos[:, 0], j=faces_pos[:, 1], k=faces_pos[:, 2],
-                            color='blue',
-                            opacity=0.3,
-                            name=f'Orbital {orbital_idx+1} (+)',
-                            hoverinfo='skip'
+                        homo_pos_trace = go.Mesh3d(
+                            x=verts_scaled[:, 0],
+                            y=verts_scaled[:, 1],
+                            z=verts_scaled[:, 2],
+                            i=faces[:, 0],
+                            j=faces[:, 1],
+                            k=faces[:, 2],
+                            color='lightblue',
+                            opacity=0.4,
+                            name='HOMO (+)',
+                            hoverinfo='name'
                         )
-                        data_traces.append(orbital_trace_pos)
-                    except:
-                        pass
-                    
-                    # Create isosurface for negative lobe
-                    isoval_negative = -0.02
-                    try:
-                        verts_neg, faces_neg, _, _ = measure.marching_cubes(orbital_grid, level=isoval_negative)
+                        data_traces.append(homo_pos_trace)
+                        
+                        # HOMO negative isosurface
+                        verts_neg, faces_neg, _, _ = measure.marching_cubes(homo_grid, level=-iso_val)
                         
-                        # Transform vertices to real coordinates
-                        verts_neg[:, 0] = x_min + verts_neg[:, 0] * (x_max - x_min) / (grid_resolution - 1)
-                        verts_neg[:, 1] = y_min + verts_neg[:, 1] * (y_max - y_min) / (grid_resolution - 1)
-                        verts_neg[:, 2] = z_min + verts_neg[:, 2] * (z_max - z_min) / (grid_resolution - 1)
+                        verts_neg_scaled = np.zeros_like(verts_neg)
+                        verts_neg_scaled[:, 0] = x_min + verts_neg[:, 0] * (x_max - x_min) / (grid_size - 1)
+                        verts_neg_scaled[:, 1] = y_min + verts_neg[:, 1] * (y_max - y_min) / (grid_size - 1)
+                        verts_neg_scaled[:, 2] = z_min + verts_neg[:, 2] * (z_max - z_min) / (grid_size - 1)
                         
-                        orbital_trace_neg = go.Mesh3d(
-                            x=verts_neg[:, 0], y=verts_neg[:, 1], z=verts_neg[:, 2],
-                            i=faces_neg[:, 0], j=faces_neg[:, 1], k=faces_neg[:, 2],
-                            color='red',
-                            opacity=0.3,
-                            name=f'Orbital {orbital_idx+1} (-)',
-                            hoverinfo='skip'
+                        homo_neg_trace = go.Mesh3d(
+                            x=verts_neg_scaled[:, 0],
+                            y=verts_neg_scaled[:, 1],
+                            z=verts_neg_scaled[:, 2],
+                            i=faces_neg[:, 0],
+                            j=faces_neg[:, 1],
+                            k=faces_neg[:, 2],
+                            color='salmon',
+                            opacity=0.4,
+                            name='HOMO (-)',
+                            hoverinfo='name'
                         )
-                        data_traces.append(orbital_trace_neg)
-                    except:
-                        pass
+                        data_traces.append(homo_neg_trace)
                         
+                except Exception as e:
+                    print(f"Isosurface generation failed: {e}")
+                    
             except Exception as e:
-                # If orbital calculation fails, add a note
-                print(f"Orbital calculation failed: {str(e)}")
-                # Fall back to simple message - orbitals are computationally expensive
+                print(f"Orbital calculation failed: {e}")
+                import traceback
+                traceback.print_exc()
         
         # Create figure
         fig = go.Figure(data=data_traces)
@@ -357,8 +363,8 @@ name_interface = gr.Interface(
 
 # Create Blocks interface for molecule viewer with autocomplete
 with gr.Blocks() as molecule_3d_block:
-    gr.Markdown("## 🔬 3D Molecule Viewer with Quantum Orbitals")
-    gr.Markdown("Enter a chemical name to view its 2D and 3D structure. Toggle orbitals to see **real quantum mechanical Natural Atomic Orbitals (NAOs)** computed with PySCF!")
+    gr.Markdown("## 🔬 3D Molecule Viewer with HOMO/LUMO Orbitals")
+    gr.Markdown("Enter a chemical name to view its 2D and 3D structure. Toggle orbitals to see **HOMO (Highest Occupied Molecular Orbital)** computed with quantum chemistry!")
     
     with gr.Row():
         with gr.Column():
@@ -373,9 +379,9 @@ with gr.Blocks() as molecule_3d_block:
                 filterable=True
             )
             orbital_checkbox = gr.Checkbox(
-                label="Show Quantum Orbitals (NAOs)",
+                label="Show HOMO Orbital",
                 value=False,
-                info="Compute and display real Natural Atomic Orbitals using quantum chemistry (DFT/B3LYP)"
+                info="Calculate and display HOMO (Highest Occupied Molecular Orbital) using Hartree-Fock"
             )
             submit_btn = gr.Button("Generate 3D Molecule", variant="primary")