Upload folder using huggingface_hub

README.md

@@ -17,5 +17,6 @@ tags:
 
 # AIMNet2 Interactive Demo
 
-Fast neural network interatomic potential for molecular property prediction.
-Supports SMILES, XYZ, and PDB input formats.
+Neural network potential for molecular property prediction.
+Supports energy, forces, charges, geometry optimization, and vibrational frequencies.
+3D visualization with charge coloring.

app.py

@@ -1,24 +1,33 @@
-"""AIMNet2 Gradio Demo Space.
+"""AIMNet2 Interactive Demo v2.
 
-Interactive neural network potential: energy, forces, charges, 3D visualization.
-Atoms colored by predicted partial charge (red=negative, blue=positive).
+3D visualization, geometry optimization, vibrational analysis, charge coloring.
+https://huggingface.co/spaces/isayevlab/aimnet2-demo
 """
 
-import io
-import os
+from __future__ import annotations
 
+import html
+import json
+import tempfile
+import time
+from pathlib import Path
+
+import gradio as gr
 import numpy as np
+import plotly.graph_objects as go
 import torch
+from plotly.subplots import make_subplots
 
 torch.set_num_threads(2)
 
-import gradio as gr
-
 # ---------------------------------------------------------------------------
 # Constants
 # ---------------------------------------------------------------------------
 MAX_ATOMS = 200
+MAX_ATOMS_OPT = 50
 MAX_ATOMS_HESSIAN = 50
+REQUEST_TIMEOUT = 90  # seconds cumulative per request
+OPT_TIMEOUT = 85  # leave margin for Hessian
 
 HARTREE_TO_EV = 27.211386024367243
 EV_TO_KCAL = 23.06054783
@@ -30,128 +39,77 @@ ELEMENT_SYMBOLS = {
 }
 SYMBOL_TO_NUM = {v: k for k, v in ELEMENT_SYMBOLS.items()}
 
-# Jmol-style CPK colors for default view
-CPK_COLORS = {
-    "H": "#FFFFFF", "B": "#FFB5B5", "C": "#909090", "N": "#3050F8",
-    "O": "#FF0D0D", "F": "#90E050", "Si": "#F0C8A0", "P": "#FF8000",
-    "S": "#FFFF30", "Cl": "#1FF01F", "As": "#BD80E3", "Se": "#FFA100",
-    "Br": "#A62929", "Pd": "#006985", "I": "#940094",
+# Atomic masses in amu (IUPAC 2021)
+ATOMIC_MASSES = {
+    1: 1.008, 5: 10.81, 6: 12.011, 7: 14.007, 8: 15.999, 9: 18.998,
+    14: 28.085, 15: 30.974, 16: 32.06, 17: 35.45, 33: 74.922,
+    34: 78.971, 35: 79.904, 46: 106.42, 53: 126.904,
 }
 
+# Unit conversion: eigenvalue (eV/A^2/amu) -> s^-2
+_EV_TO_J = 1.602176634e-19
+_AMU_TO_KG = 1.66053906660e-27
+_A_TO_M = 1e-10
+_C_CM = 2.99792458e10  # speed of light in cm/s
+_FREQ_CONV = _EV_TO_J / (_A_TO_M**2 * _AMU_TO_KG)  # eV/(A^2*amu) -> s^-2
+
+
 # ---------------------------------------------------------------------------
-# Lazy model loader
+# Model loader (eager, singleton)
 # ---------------------------------------------------------------------------
-CALC = None
+BASE_CALC = None
 
 
-def get_calc():
-    global CALC
-    if CALC is None:
+def get_base_calc():
+    """Return shared AIMNet2Calculator singleton (thread-safe for read-only use)."""
+    global BASE_CALC
+    if BASE_CALC is None:
         from aimnet.calculators import AIMNet2Calculator
-        from aimnet.calculators.aimnet2ase import AIMNet2ASE
+        BASE_CALC = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3", device="cpu")
+    return BASE_CALC
 
-        base_calc = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3", device="cpu")
-        CALC = AIMNet2ASE(base_calc)
-    return CALC
 
+def make_ase_calc(charge: int = 0):
+    """Create a fresh AIMNet2ASE wrapper per request (concurrency-safe)."""
+    from aimnet.calculators.aimnet2ase import AIMNet2ASE
+    return AIMNet2ASE(get_base_calc(), charge=charge)
 
-# ---------------------------------------------------------------------------
-# 3D Visualization
-# ---------------------------------------------------------------------------
 
-def charge_to_rgb(q, qmin=-0.8, qmax=0.8):
-    """Map charge to color: red (negative) -> white (neutral) -> blue (positive)."""
-    t = np.clip((q - qmin) / (qmax - qmin), 0, 1)
-    if t < 0.5:
-        # red -> white
-        s = t * 2
-        r, g, b = 1.0, s, s
-    else:
-        # white -> blue
-        s = (t - 0.5) * 2
-        r, g, b = 1.0 - s, 1.0 - s, 1.0
-    return f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}"
-
-
-def build_3dmol_html(coords, numbers, charges=None, width=500, height=400):
-    """Generate HTML with embedded 3Dmol.js viewer, atoms colored by charge."""
-    # Build XYZ string
-    n = len(numbers)
-    xyz_lines = [str(n), "AIMNet2 prediction"]
-    for i in range(n):
-        sym = ELEMENT_SYMBOLS.get(int(numbers[i]), "X")
-        x, y, z = coords[i]
-        xyz_lines.append(f"{sym} {x:.6f} {y:.6f} {z:.6f}")
-    xyz_str = "\\n".join(xyz_lines)
-
-    # Build per-atom color assignments
-    if charges is not None:
-        qmin = float(np.min(charges))
-        qmax = float(np.max(charges))
-        # Symmetric range for better visualization
-        qlim = max(abs(qmin), abs(qmax), 0.3)
-        color_js = ""
-        for i in range(n):
-            c = charge_to_rgb(float(charges[i]), -qlim, qlim)
-            color_js += f'viewer.getModel().setAtomStyle({{index:{i}}},{{stick:{{radius:0.15}},sphere:{{scale:0.25,color:"{c}"}}}});'
-        color_mode_label = "colored by charge (red=&minus;, blue=+)"
-    else:
-        color_js = ""
-        color_mode_label = "CPK colors"
-
-    html = f"""
-    <div id="viewer-container" style="position:relative;width:{width}px;height:{height}px;margin:0 auto;">
-        <div id="mol-viewer" style="width:{width}px;height:{height}px;"></div>
-        <div style="position:absolute;bottom:4px;right:8px;font-size:11px;color:#888;">{color_mode_label}</div>
-    </div>
-    <script src="https://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.4.2/3Dmol-min.js"></script>
-    <script>
-    (function() {{
-        let viewer = $3Dmol.createViewer(document.getElementById("mol-viewer"), {{
-            backgroundColor: "white"
-        }});
-        viewer.addModel("{xyz_str}", "xyz");
-        viewer.setStyle({{}}, {{stick:{{radius:0.15}}, sphere:{{scale:0.25,colorscheme:"Jmol"}}}});
-        {color_js}
-        viewer.zoomTo();
-        viewer.render();
-    }})();
-    </script>
-    """
-    return html
+# Eager-load model at import time (during Space startup, not first request)
+try:
+    get_base_calc()
+except Exception:
+    pass  # Will fail on first request with a clear error instead
 
 
 # ---------------------------------------------------------------------------
 # Parsers
 # ---------------------------------------------------------------------------
 
-def parse_smiles(smiles: str):
-    """Return (coords, numbers, formal_charge)."""
+def parse_smiles(smiles: str) -> tuple[np.ndarray, np.ndarray, int]:
+    """Parse SMILES -> (coords, numbers, formal_charge)."""
     from rdkit import Chem
     from rdkit.Chem import AllChem
 
     mol = Chem.MolFromSmiles(smiles.strip())
     if mol is None:
-        raise ValueError(f"RDKit could not parse SMILES: {smiles!r}")
+        raise ValueError(f"Invalid SMILES: {smiles!r}")
     formal_charge = Chem.GetFormalCharge(mol)
     mol = Chem.AddHs(mol)
-    result = AllChem.EmbedMolecule(mol, AllChem.ETKDGv3())
-    if result == -1:
-        raise ValueError("Failed to generate 3D coordinates.")
+    if AllChem.EmbedMolecule(mol, AllChem.ETKDGv3()) == -1:
+        raise ValueError("Failed to generate 3D coordinates. Try a different molecule.")
     AllChem.MMFFOptimizeMolecule(mol)
     conf = mol.GetConformer()
-    coords = np.array([conf.GetAtomPosition(i) for i in range(mol.GetNumAtoms())], dtype=np.float64)
-    numbers = np.array([atom.GetAtomicNum() for atom in mol.GetAtoms()], dtype=np.int64)
+    coords = np.array([conf.GetAtomPosition(i) for i in range(mol.GetNumAtoms())])
+    numbers = np.array([a.GetAtomicNum() for a in mol.GetAtoms()])
     return coords, numbers, formal_charge
 
 
-def parse_xyz(xyz_text: str):
-    """Return (coords, numbers)."""
-    lines = [l.strip() for l in xyz_text.strip().splitlines()]
-    start = 0
-    if lines and lines[0].isdigit():
-        start = 2
-    coords_list, numbers_list = [], []
+def parse_xyz(text: str) -> tuple[np.ndarray, np.ndarray]:
+    """Parse XYZ format text -> (coords, numbers)."""
+    lines = [l.strip() for l in text.strip().splitlines()]
+    start = 2 if lines and lines[0].isdigit() else 0
+    coords_list, nums_list = [], []
     for line in lines[start:]:
         if not line:
             continue
@@ -161,17 +119,17 @@ def parse_xyz(xyz_text: str):
         sym = parts[0].capitalize()
         if sym not in SYMBOL_TO_NUM:
             raise ValueError(f"Unknown element: {sym!r}")
-        numbers_list.append(SYMBOL_TO_NUM[sym])
+        nums_list.append(SYMBOL_TO_NUM[sym])
         coords_list.append([float(parts[1]), float(parts[2]), float(parts[3])])
     if not coords_list:
         raise ValueError("No atoms found in XYZ input.")
-    return np.array(coords_list, dtype=np.float64), np.array(numbers_list, dtype=np.int64)
+    return np.array(coords_list), np.array(nums_list)
 
 
-def parse_pdb(pdb_text: str):
-    """Return (coords, numbers)."""
-    coords_list, numbers_list = [], []
-    for line in pdb_text.splitlines():
+def parse_pdb(text: str) -> tuple[np.ndarray, np.ndarray]:
+    """Parse PDB format text -> (coords, numbers)."""
+    coords_list, nums_list = [], []
+    for line in text.splitlines():
         if not line.startswith(("ATOM", "HETATM")):
             continue
         try:
@@ -184,264 +142,857 @@ def parse_pdb(pdb_text: str):
         elem = elem.capitalize()
         if elem not in SYMBOL_TO_NUM:
             raise ValueError(f"Unknown element in PDB: {elem!r}")
-        numbers_list.append(SYMBOL_TO_NUM[elem])
+        nums_list.append(SYMBOL_TO_NUM[elem])
         coords_list.append([x, y, z])
     if not coords_list:
         raise ValueError("No ATOM/HETATM records found.")
-    return np.array(coords_list, dtype=np.float64), np.array(numbers_list, dtype=np.int64)
+    return np.array(coords_list), np.array(nums_list)
+
+
+def parse_input(text: str, fmt: str) -> tuple[np.ndarray, np.ndarray, str]:
+    """Parse molecule input. Returns (coords, numbers, warning_str)."""
+    warning = ""
+    if fmt == "SMILES":
+        coords, numbers, _fc = parse_smiles(text)
+    elif fmt == "XYZ":
+        coords, numbers = parse_xyz(text)
+    elif fmt == "PDB":
+        coords, numbers = parse_pdb(text)
+    else:
+        raise ValueError(f"Unknown format: {fmt}")
+    return coords, numbers, warning
+
+
+def handle_file_upload(file_obj) -> tuple[str, str]:
+    """Process uploaded file. Returns (text_content, format_name).
+
+    Populates the text input and sets format radio.
+    """
+    if file_obj is None:
+        return "", "SMILES"
+    path = Path(file_obj.name if hasattr(file_obj, "name") else file_obj)
+    suffix = path.suffix.lower()
+    text = path.read_text()
+
+    if suffix == ".xyz":
+        return text, "XYZ"
+    elif suffix == ".pdb":
+        return text, "PDB"
+    elif suffix in (".sdf", ".mol"):
+        from rdkit import Chem
+        from rdkit.Chem import AllChem
+        suppl = Chem.SDMolSupplier(str(path), removeHs=False)
+        mol = next(suppl, None)
+        if mol is None:
+            raise ValueError("Could not read SDF file.")
+        if mol.GetNumConformers() == 0:
+            mol = Chem.AddHs(mol)
+            AllChem.EmbedMolecule(mol, AllChem.ETKDGv3())
+        # Convert to XYZ text
+        conf = mol.GetConformer()
+        n = mol.GetNumAtoms()
+        xyz_lines = [str(n), f"Converted from {path.name}"]
+        for i in range(n):
+            pos = conf.GetAtomPosition(i)
+            sym = mol.GetAtomWithIdx(i).GetSymbol()
+            xyz_lines.append(f"{sym} {pos.x:.6f} {pos.y:.6f} {pos.z:.6f}")
+        return "\n".join(xyz_lines), "XYZ"
+    else:
+        raise ValueError(f"Unsupported file type: {suffix}")
+
+
+# ---------------------------------------------------------------------------
+# 3D Viewer (iframe + 3Dmol.js)
+# ---------------------------------------------------------------------------
+
+def _charge_to_hex(q: float, qlim: float) -> str:
+    """Map charge to color: red (negative) -> white (0) -> blue (positive)."""
+    t = np.clip((q + qlim) / (2 * qlim), 0, 1)
+    if t < 0.5:
+        s = t * 2
+        r, g, b = 1.0, s, s
+    else:
+        s = (t - 0.5) * 2
+        r, g, b = 1.0 - s, 1.0 - s, 1.0
+    return f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}"
+
+
+def build_viewer_html(
+    coords: np.ndarray,
+    numbers: np.ndarray,
+    charges: np.ndarray | None = None,
+    height: int = 420,
+) -> str:
+    """Build iframe HTML with 3Dmol.js viewer and CPK/charge toggle."""
+    n = len(numbers)
+
+    # Build XYZ string
+    xyz_lines = [str(n), "AIMNet2"]
+    for i in range(n):
+        sym = ELEMENT_SYMBOLS.get(int(numbers[i]), "X")
+        x, y, z = coords[i]
+        xyz_lines.append(f"{sym} {x:.6f} {y:.6f} {z:.6f}")
+    xyz_string = "\n".join(xyz_lines)
+
+    # Build per-atom charge color JS (only if charges and <= 100 atoms)
+    charge_js = ""
+    has_toggle = charges is not None and n <= 100
+    if has_toggle:
+        qlim = max(float(np.max(np.abs(charges))), 0.3)
+        charge_styles = []
+        for i in range(n):
+            c = _charge_to_hex(float(charges[i]), qlim)
+            charge_styles.append(
+                f'viewer.getModel().setAtomStyle({{index:{i}}},'
+                f'{{stick:{{radius:0.15}},sphere:{{scale:0.25,color:"{c}"}}}});'
+            )
+        charge_js = "\n".join(charge_styles)
+
+    toggle_btn = ""
+    toggle_fn = ""
+    if has_toggle:
+        toggle_btn = (
+            '<button id="toggle-btn" onclick="toggleColors()" '
+            'style="position:absolute;top:8px;right:8px;z-index:10;'
+            'padding:4px 10px;font-size:12px;cursor:pointer;'
+            'border:1px solid #ccc;border-radius:4px;background:#f8f8f8;">'
+            'Color by charge</button>'
+        )
+        toggle_fn = f"""
+var cpkMode = true;
+function setCPK() {{
+    viewer.setStyle({{}}, {{stick:{{radius:0.15}}, sphere:{{scale:0.25, colorscheme:"Jmol"}}}});
+    viewer.render();
+}}
+function setCharges() {{
+    {charge_js}
+    viewer.render();
+}}
+function toggleColors() {{
+    cpkMode = !cpkMode;
+    if (cpkMode) {{
+        setCPK();
+        document.getElementById("toggle-btn").textContent = "Color by charge";
+    }} else {{
+        setCharges();
+        document.getElementById("toggle-btn").textContent = "CPK colors";
+    }}
+}}
+"""
+
+    inner_html = f"""<!DOCTYPE html>
+<html><head>
+<meta charset="utf-8">
+<style>
+body {{ margin:0; overflow:hidden; font-family:sans-serif; }}
+#viewer {{ width:100%; height:{height}px; position:relative; }}
+#fallback {{ display:none; padding:20px; color:#888; text-align:center; }}
+</style>
+</head><body>
+<div id="viewer"></div>
+{toggle_btn}
+<div id="fallback">3D viewer unavailable. Results are shown below.</div>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.4.2/3Dmol-min.js"></script>
+<script>
+try {{
+    var xyz = {json.dumps(xyz_string)};
+    var viewer = $3Dmol.createViewer("viewer", {{backgroundColor:"white"}});
+    viewer.addModel(xyz, "xyz");
+    viewer.setStyle({{}}, {{stick:{{radius:0.15}}, sphere:{{scale:0.25, colorscheme:"Jmol"}}}});
+    viewer.zoomTo();
+    viewer.render();
+    {toggle_fn}
+}} catch(e) {{
+    document.getElementById("viewer").style.display = "none";
+    document.getElementById("fallback").style.display = "block";
+}}
+</script>
+</body></html>"""
+
+    escaped = html.escape(inner_html, quote=True)
+    return (
+        f'<iframe srcdoc="{escaped}" width="100%" height="{height + 30}" '
+        f'frameborder="0" sandbox="allow-scripts" '
+        f'style="border:1px solid #eee;border-radius:8px;"></iframe>'
+    )
 
 
 # ---------------------------------------------------------------------------
-# Main prediction function — returns (markdown, html_viewer)
+# Frequency computation
+# ---------------------------------------------------------------------------
+
+def is_linear(coords: np.ndarray, numbers: np.ndarray, tol: float = 1e-3) -> bool:
+    """Check if molecule is linear via moment of inertia tensor."""
+    masses = np.array([ATOMIC_MASSES.get(int(z), 1.0) for z in numbers])
+    com = np.average(coords, weights=masses, axis=0)
+    r = coords - com
+    I = np.zeros((3, 3))
+    for m, ri in zip(masses, r):
+        I += m * (np.dot(ri, ri) * np.eye(3) - np.outer(ri, ri))
+    eigvals = np.linalg.eigvalsh(I)
+    return eigvals[0] / max(eigvals[-1], 1e-30) < tol
+
+
+def compute_frequencies(
+    hessian: np.ndarray,
+    numbers: np.ndarray,
+    coords: np.ndarray,
+) -> tuple[np.ndarray, int]:
+    """Compute vibrational frequencies from Hessian.
+
+    Parameters
+    ----------
+    hessian : ndarray, shape (N,3,N,3) or (3N,3N)
+        Hessian in eV/A^2.
+    numbers : ndarray, shape (N,)
+        Atomic numbers.
+    coords : ndarray, shape (N,3)
+        Atomic positions (for linearity check).
+
+    Returns
+    -------
+    freqs_cm : ndarray
+        Vibrational frequencies in cm^-1. Negative = imaginary.
+    n_imag : int
+        Number of imaginary frequencies.
+    """
+    n = len(numbers)
+    H = hessian.reshape(3 * n, 3 * n)
+
+    # Mass-weight
+    masses = np.array([ATOMIC_MASSES.get(int(z), 1.0) for z in numbers])
+    masses_3n = np.repeat(masses, 3)
+    H_mw = H / np.sqrt(np.outer(masses_3n, masses_3n))
+    H_mw = 0.5 * (H_mw + H_mw.T)  # symmetrize
+
+    eigenvalues = np.linalg.eigvalsh(H_mw)
+
+    # Convert to cm^-1
+    freqs = (
+        np.sign(eigenvalues)
+        * np.sqrt(np.abs(eigenvalues) * _FREQ_CONV)
+        / (2 * np.pi * _C_CM)
+    )
+
+    # Remove translation/rotation modes (count-based)
+    n_tr = 5 if is_linear(coords, numbers) else 6
+    sorted_idx = np.argsort(np.abs(freqs))
+    vib_idx = sorted_idx[n_tr:]
+    freqs_vib = np.sort(freqs[vib_idx])
+
+    n_imag = int(np.sum(freqs_vib < -10))
+    return freqs_vib, n_imag
+
+
+# ---------------------------------------------------------------------------
+# Plotting
+# ---------------------------------------------------------------------------
+
+def make_frequency_plot(freqs: np.ndarray) -> go.Figure:
+    """Create Plotly stick spectrum of vibrational frequencies."""
+    real = freqs[freqs > 0]
+    fig = go.Figure()
+    if len(real) > 0:
+        fig.add_trace(go.Bar(
+            x=real, y=np.ones_like(real),
+            width=3, marker_color="steelblue",
+            hovertemplate="%{x:.1f} cm\u207b\u00b9<extra></extra>",
+        ))
+    fig.update_layout(
+        xaxis_title="Frequency (cm\u207b\u00b9)",
+        yaxis_visible=False,
+        height=200, margin=dict(l=40, r=20, t=30, b=40),
+        title="Vibrational Spectrum",
+        showlegend=False,
+    )
+    return fig
+
+
+def make_convergence_plot(trajectory: list[dict]) -> go.Figure:
+    """Create dual-axis convergence plot (energy + max force vs step)."""
+    steps = [t["step"] for t in trajectory]
+    energies = [t["energy"] for t in trajectory]
+    fmaxes = [t["fmax"] for t in trajectory]
+
+    fig = make_subplots(specs=[[{"secondary_y": True}]])
+    fig.add_trace(
+        go.Scatter(x=steps, y=energies, name="Energy (eV)", mode="lines+markers",
+                   marker=dict(size=4), line=dict(color="steelblue")),
+        secondary_y=False,
+    )
+    fig.add_trace(
+        go.Scatter(x=steps, y=fmaxes, name="Max |F| (eV/\u00c5)", mode="lines+markers",
+                   marker=dict(size=4), line=dict(color="firebrick")),
+        secondary_y=True,
+    )
+    fig.update_xaxes(title_text="Step")
+    fig.update_yaxes(title_text="Energy (eV)", secondary_y=False)
+    fig.update_yaxes(title_text="Max |F| (eV/\u00c5)", secondary_y=True)
+    fig.update_layout(
+        height=280, margin=dict(l=60, r=60, t=30, b=40),
+        legend=dict(x=0.5, y=1.15, xanchor="center", orientation="h"),
+    )
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Geometry optimization
+# ---------------------------------------------------------------------------
+
+def run_optimization(
+    atoms,
+    max_steps: int,
+    fmax_target: float,
+    timeout: float = OPT_TIMEOUT,
+) -> tuple[list[dict], bool, float]:
+    """Run LBFGS optimization with timeout.
+
+    Returns (trajectory, converged, wall_time).
+    Reads from ASE cache to avoid double-computing.
+    """
+    from ase.optimize import LBFGS
+
+    opt = LBFGS(atoms, logfile=None)
+    trajectory = []
+    t0 = time.time()
+    converged = False
+
+    for step in range(max_steps):
+        if time.time() - t0 > timeout:
+            break
+        opt.step()
+        e = float(atoms.calc.results["energy"])
+        f = atoms.calc.results["forces"]
+        fmax = float(np.max(np.linalg.norm(f, axis=1)))
+        trajectory.append({"step": step + 1, "energy": e, "fmax": fmax})
+        if fmax < fmax_target:
+            converged = True
+            break
+
+    return trajectory, converged, time.time() - t0
+
+
+# ---------------------------------------------------------------------------
+# Reproduction script generator
+# ---------------------------------------------------------------------------
+
+def _fmt_array(arr: np.ndarray, name: str) -> str:
+    """Format numpy array as valid Python code."""
+    if arr.ndim == 1:
+        return f"{name} = {arr.tolist()!r}"
+    # 2D
+    rows = []
+    for row in arr:
+        rows.append("    [" + ", ".join(f"{v:.6f}" for v in row) + "],")
+    return f"{name} = np.array([\n" + "\n".join(rows) + "\n])"
+
+
+def generate_script(
+    coords: np.ndarray,
+    numbers: np.ndarray,
+    charge: int,
+    task: str = "single_point",
+    max_steps: int = 30,
+    fmax: float = 0.05,
+    compute_hessian: bool = False,
+) -> str:
+    """Generate Python reproduction script."""
+    lines = [
+        "# AIMNet2 calculation",
+        "# Generated by https://huggingface.co/spaces/isayevlab/aimnet2-demo",
+        "from aimnet.calculators import AIMNet2Calculator",
+        "from aimnet.calculators.aimnet2ase import AIMNet2ASE",
+        "from ase import Atoms",
+        "import numpy as np",
+        "",
+        _fmt_array(coords, "coords"),
+        f"numbers = {numbers.tolist()!r}",
+        f"charge = {charge}",
+        "",
+        'calc = AIMNet2ASE(AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3"), charge=charge)',
+        "atoms = Atoms(numbers=numbers, positions=coords)",
+        "atoms.calc = calc",
+        "",
+    ]
+
+    if task == "optimize":
+        lines += [
+            "from ase.optimize import LBFGS",
+            f"opt = LBFGS(atoms, logfile='-')",
+            f"opt.run(fmax={fmax}, steps={max_steps})",
+            "",
+            "energy = atoms.get_potential_energy()",
+            'print(f"Optimized energy: {energy:.6f} eV")',
+            'print(f"Max force: {max(np.linalg.norm(atoms.get_forces(), axis=1)):.6f} eV/A")',
+        ]
+    else:
+        lines += [
+            "energy = atoms.get_potential_energy()",
+            "forces = atoms.get_forces()",
+            'charges = atoms.calc.results["charges"]',
+            'print(f"Energy: {energy:.6f} eV")',
+        ]
+
+    if compute_hessian:
+        lines += [
+            "",
+            "# Hessian & frequencies",
+            "base_calc = calc.base_calc",
+            'hess_result = base_calc({"coord": atoms.get_positions(), '
+            '"numbers": atoms.numbers, "charge": float(charge)}, hessian=True)',
+            'hessian = hess_result["hessian"].detach().cpu().numpy()',
+            "# Diagonalize mass-weighted Hessian for frequencies (see demo source for details)",
+        ]
+
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# XYZ download helper
+# ---------------------------------------------------------------------------
+
+def write_xyz_file(coords: np.ndarray, numbers: np.ndarray,
+                   charges: np.ndarray | None = None,
+                   comment: str = "AIMNet2") -> str:
+    """Write XYZ to a temp file and return the path."""
+    n = len(numbers)
+    lines = [str(n), comment]
+    for i in range(n):
+        sym = ELEMENT_SYMBOLS.get(int(numbers[i]), "X")
+        x, y, z = coords[i]
+        q_str = f"  {charges[i]:+.4f}" if charges is not None else ""
+        lines.append(f"{sym:2s} {x:12.6f} {y:12.6f} {z:12.6f}{q_str}")
+    tmp = tempfile.NamedTemporaryFile(suffix=".xyz", delete=False, mode="w")
+    tmp.write("\n".join(lines))
+    tmp.close()
+    return tmp.name
+
+
+# ---------------------------------------------------------------------------
+# Tab 1: Single-point calculation
 # ---------------------------------------------------------------------------
 
 def predict(input_text, input_format, charge, compute_forces, compute_hessian):
-    """Run AIMNet2 and return (results_markdown, 3d_viewer_html)."""
+    """Run single-point calculation. Returns (markdown, viewer_html, freq_plot, xyz_file, script)."""
     charge = int(charge)
+    empty = ("", "", None, None, "")
 
-    # --- Parse ---
-    smiles_warning = ""
+    # Parse
     try:
-        if input_format == "SMILES":
-            coords, numbers, fc = parse_smiles(input_text)
-            if fc != charge:
-                smiles_warning = (
-                    f"\n> **Warning:** SMILES formal charge ({fc:+d}) != "
-                    f"supplied charge ({charge:+d}). Using supplied charge.\n"
-                )
-        elif input_format == "XYZ":
-            coords, numbers = parse_xyz(input_text)
-        elif input_format == "PDB":
-            coords, numbers = parse_pdb(input_text)
-        else:
-            return f"**Error:** Unknown format: {input_format}", ""
+        coords, numbers, warning = parse_input(input_text, input_format)
     except Exception as e:
-        return f"**Parse error:** {e}", ""
+        return (f"**Parse error:** {e}", *empty[1:])
 
-    n_atoms = len(numbers)
-    if n_atoms > MAX_ATOMS:
-        return f"**Error:** {n_atoms} atoms exceeds limit of {MAX_ATOMS}.", ""
-    if compute_hessian and n_atoms > MAX_ATOMS_HESSIAN:
-        return f"**Error:** Hessian limited to {MAX_ATOMS_HESSIAN} atoms ({n_atoms} given).", ""
+    n = len(numbers)
+    if n > MAX_ATOMS:
+        return (f"**Error:** {n} atoms exceeds limit of {MAX_ATOMS}.", *empty[1:])
+    if compute_hessian and n > MAX_ATOMS_HESSIAN:
+        return (f"**Error:** Hessian limited to {MAX_ATOMS_HESSIAN} atoms ({n} given).", *empty[1:])
 
-    unsupported = sorted({int(z) for z in numbers} - set(ELEMENT_SYMBOLS.keys()))
+    # Validate elements
+    unsupported = sorted({int(z) for z in numbers} - set(ELEMENT_SYMBOLS))
     if unsupported:
-        return f"**Error:** Unsupported elements: {unsupported}", ""
+        return (f"**Error:** Unsupported elements: {unsupported}", *empty[1:])
+
+    # SMILES charge validation
+    smiles_warn = ""
+    if input_format == "SMILES":
+        _, _, fc = parse_smiles(input_text)  # already parsed, just get charge
+        if fc != charge:
+            smiles_warn = (
+                f"\n> **Warning:** SMILES formal charge ({fc:+d}) != "
+                f"supplied charge ({charge:+d}). Using supplied charge.\n"
+            )
 
-    # --- Calculate ---
+    # Calculate
     try:
-        calc = get_calc()
-        calc.set_charge(charge)
-
+        ase_calc = make_ase_calc(charge)
         from ase import Atoms
         symbols = [ELEMENT_SYMBOLS[int(z)] for z in numbers]
         atoms = Atoms(symbols=symbols, positions=coords)
-        atoms.calc = calc
+        atoms.calc = ase_calc
 
         atoms.get_potential_energy()
-        energy_ev = float(atoms.calc.results["energy"])
-        charges_arr = atoms.calc.results.get("charges")
+        energy_ev = float(ase_calc.results["energy"])
+        charges_arr = ase_calc.results.get("charges")
+
+        if not np.isfinite(energy_ev):
+            return ("**Error:** Model produced NaN/Inf. Molecule may be outside training domain.", *empty[1:])
 
         forces_arr = None
         if compute_forces:
             atoms.get_forces()
-            forces_arr = atoms.calc.results["forces"]
+            forces_arr = ase_calc.results["forces"]
 
         hessian_arr = None
+        freqs = None
+        n_imag = 0
         if compute_hessian:
             data = {"coord": coords, "numbers": numbers, "charge": float(charge)}
-            hess_results = calc.base_calc(data, hessian=True)
-            hessian_arr = hess_results["hessian"].detach().cpu().numpy()
+            hess_result = get_base_calc()(data, hessian=True)
+            hessian_arr = hess_result["hessian"].detach().cpu().numpy()
+            freqs, n_imag = compute_frequencies(hessian_arr, numbers, coords)
 
     except Exception as e:
         import traceback
-        return f"**Calculation error:** {e}\n```\n{traceback.format_exc()}\n```", ""
+        return (f"**Calculation error:** {e}\n```\n{traceback.format_exc()}\n```", *empty[1:])
 
-    # --- 3D Viewer (colored by charge) ---
-    viewer_html = build_3dmol_html(coords, numbers, charges_arr)
+    # Build outputs
+    viewer_html = build_viewer_html(coords, numbers, charges_arr)
 
-    # --- Format results ---
+    # Results markdown
     energy_kcal = energy_ev * EV_TO_KCAL
     energy_ha = energy_ev / HARTREE_TO_EV
+    md = []
+    md.append("## AIMNet2 Results\n")
+    if smiles_warn:
+        md.append(smiles_warn)
+    md.append(f"**Atoms:** {n} | **Charge:** {charge:+d}\n")
+
+    md.append("### Energy\n| Unit | Value |\n|------|------:|")
+    md.append(f"| eV | {energy_ev:.6f} |")
+    md.append(f"| kcal/mol | {energy_kcal:.4f} |")
+    md.append(f"| Hartree | {energy_ha:.8f} |\n")
 
-    lines = []
-    lines.append("## AIMNet2 Results\n")
-    if smiles_warning:
-        lines.append(smiles_warning)
-    lines.append(f"**Atoms:** {n_atoms} | **Charge:** {charge:+d}\n")
-
-    # Energy table
-    lines.append("### Energy\n")
-    lines.append("| Unit | Value |")
-    lines.append("|------|------:|")
-    lines.append(f"| eV | {energy_ev:.6f} |")
-    lines.append(f"| kcal/mol | {energy_kcal:.4f} |")
-    lines.append(f"| Hartree | {energy_ha:.8f} |")
-    lines.append("")
-
-    # Charges table
     if charges_arr is not None:
-        lines.append("### Partial Charges (e)\n")
-        lines.append("| # | Elem | Charge |")
-        lines.append("|--:|:----:|-------:|")
+        md.append("### Partial Charges (e)\n| # | Elem | Charge |\n|--:|:----:|-------:|")
         for i, (z, q) in enumerate(zip(numbers, charges_arr)):
-            sym = ELEMENT_SYMBOLS.get(int(z), f"Z{z}")
-            lines.append(f"| {i+1} | {sym} | {q:+.4f} |")
-        lines.append("")
-        lines.append(f"*Sum: {float(np.sum(charges_arr)):+.4f} e*\n")
+            sym = ELEMENT_SYMBOLS.get(int(z), "?")
+            md.append(f"| {i+1} | {sym} | {q:+.4f} |")
+        md.append(f"\n*Sum: {float(np.sum(charges_arr)):+.4f} e*\n")
 
-    # Forces
     if forces_arr is not None:
         max_f = float(np.max(np.linalg.norm(forces_arr, axis=1)))
-        rms_f = float(np.sqrt(np.mean(forces_arr ** 2)))
-        lines.append("### Forces (eV/A)\n")
-        lines.append("| Metric | Value |")
-        lines.append("|--------|------:|")
-        lines.append(f"| Max |F| | {max_f:.6f} |")
-        lines.append(f"| RMS | {rms_f:.6f} |")
-        lines.append("")
+        rms_f = float(np.sqrt(np.mean(forces_arr**2)))
+        md.append("### Forces (eV/A)\n| Metric | Value |\n|--------|------:|")
+        md.append(f"| Max |F| | {max_f:.6f} |")
+        md.append(f"| RMS | {rms_f:.6f} |")
         if input_format == "SMILES":
-            lines.append("> *Geometry from MMFF, not AIMNet2-optimized. Non-zero forces expected.*\n")
+            md.append("\n> *Geometry from MMFF, not AIMNet2-optimized.*\n")
+
+    freq_plot = None
+    if freqs is not None:
+        real_f = freqs[freqs > 0]
+        imag_f = freqs[freqs < 0]
+        md.append("### Vibrational Frequencies\n")
+        if max_f > 0.05 if forces_arr is not None else True:
+            md.append("> *Frequencies at non-stationary point. Low modes may be unreliable.*\n")
+        if n_imag > 0:
+            md.append(f"> **{n_imag} imaginary frequency(ies)** -- not a true minimum.\n")
+        if len(real_f) > 0:
+            md.append("```")
+            for j, f in enumerate(real_f):
+                md.append(f"  {j+1:3d}:  {f:10.2f} cm-1")
+            md.append("```")
+        if len(imag_f) > 0:
+            md.append("\nImaginary:\n```")
+            for j, f in enumerate(imag_f):
+                md.append(f"  {j+1:3d}:  {abs(f):10.2f}i cm-1")
+            md.append("```")
+        freq_plot = make_frequency_plot(freqs)
+
+    md.append("\n---")
+    md.append("*AIMNet2 wB97M-D3 | [Model](https://huggingface.co/isayevlab/aimnet2-wb97m-d3) | [Paper](https://doi.org/10.1039/D4SC08572H)*")
+
+    xyz_file = write_xyz_file(coords, numbers, charges_arr,
+                              comment=f"Energy: {energy_ev:.6f} eV")
+
+    script = generate_script(coords, numbers, charge, "single_point",
+                             compute_hessian=compute_hessian)
+
+    return "\n".join(md), viewer_html, freq_plot, xyz_file, script
 
-    # Hessian
-    if hessian_arr is not None:
-        # Compute vibrational frequencies from Hessian
-        lines.append("### Vibrational Analysis\n")
-        try:
-            freqs = _compute_frequencies(hessian_arr, numbers)
-            real_freqs = freqs[freqs > 0]
-            imag_freqs = freqs[freqs < 0]
-            if len(real_freqs) > 0:
-                lines.append(f"**Real frequencies:** {len(real_freqs)}\n")
-                lines.append("```")
-                for j, f in enumerate(real_freqs):
-                    lines.append(f"  {j+1:3d}:  {f:10.2f} cm-1")
-                lines.append("```\n")
-            if len(imag_freqs) > 0:
-                lines.append(f"**Imaginary frequencies:** {len(imag_freqs)}\n")
-                lines.append("```")
-                for j, f in enumerate(imag_freqs):
-                    lines.append(f"  {j+1:3d}:  {f:10.2f}i cm-1")
-                lines.append("```\n")
-        except Exception as e:
-            lines.append(f"Frequency analysis failed: {e}\n")
-            lines.append(f"Hessian shape: `{hessian_arr.shape}`, norm: `{float(np.linalg.norm(hessian_arr)):.4f}`\n")
 
-    lines.append("---")
-    lines.append("*AIMNet2 wB97M-D3 | [Model card](https://huggingface.co/isayevlab/aimnet2-wb97m-d3) | [Paper](https://doi.org/10.1039/D4SC08572H)*")
+# ---------------------------------------------------------------------------
+# Tab 2: Geometry optimization
+# ---------------------------------------------------------------------------
 
-    return "\n".join(lines), viewer_html
+def optimize(input_text, input_format, charge, max_steps, fmax_target,
+             compute_freqs):
+    """Run geometry optimization. Returns (md, viewer_html, conv_plot, freq_plot, xyz_file, script)."""
+    charge = int(charge)
+    max_steps = int(max_steps)
+    fmax_target = float(fmax_target)
+    empty = ("", "", None, None, None, "")
 
+    # Auto-tighten fmax when frequencies requested
+    if compute_freqs and fmax_target > 0.02:
+        fmax_target = 0.02
 
-def _compute_frequencies(hessian, numbers):
-    """Compute vibrational frequencies from Hessian matrix.
+    # Validate fmax
+    if not 0.01 <= fmax_target <= 1.0:
+        return ("**Error:** fmax must be between 0.01 and 1.0 eV/A.", *empty[1:])
+
+    # Parse
+    try:
+        coords, numbers, _ = parse_input(input_text, input_format)
+    except Exception as e:
+        return (f"**Parse error:** {e}", *empty[1:])
 
-    Returns frequencies in cm^-1 (negative = imaginary).
-    """
-    # Atomic masses in amu
-    MASSES = {
-        1: 1.008, 5: 10.81, 6: 12.011, 7: 14.007, 8: 15.999, 9: 18.998,
-        14: 28.085, 15: 30.974, 16: 32.06, 17: 35.45, 33: 74.922,
-        34: 78.971, 35: 79.904, 46: 106.42, 53: 126.904,
-    }
     n = len(numbers)
+    if n > MAX_ATOMS_OPT:
+        return (f"**Error:** Optimization limited to {MAX_ATOMS_OPT} atoms ({n} given).", *empty[1:])
+
+    unsupported = sorted({int(z) for z in numbers} - set(ELEMENT_SYMBOLS))
+    if unsupported:
+        return (f"**Error:** Unsupported elements: {unsupported}", *empty[1:])
 
-    # Reshape Hessian from (N,3,N,3) to (3N, 3N)
-    if hessian.ndim == 4:
-        h = hessian.reshape(3 * n, 3 * n)
+    # Optimize
+    try:
+        ase_calc = make_ase_calc(charge)
+        from ase import Atoms
+        symbols = [ELEMENT_SYMBOLS[int(z)] for z in numbers]
+        atoms = Atoms(symbols=symbols, positions=coords)
+        atoms.calc = ase_calc
+
+        # Initial energy/forces
+        atoms.get_potential_energy()
+        e0 = float(ase_calc.results["energy"])
+        f0 = ase_calc.results["forces"]
+        fmax0 = float(np.max(np.linalg.norm(f0, axis=1)))
+
+        trajectory, converged, wall_time = run_optimization(
+            atoms, max_steps, fmax_target
+        )
+
+        opt_coords = atoms.get_positions()
+        e_final = trajectory[-1]["energy"] if trajectory else e0
+        fmax_final = trajectory[-1]["fmax"] if trajectory else fmax0
+        charges_arr = ase_calc.results.get("charges")
+
+        if not np.isfinite(e_final):
+            return ("**Error:** Model produced NaN/Inf during optimization.", *empty[1:])
+
+        # Frequencies at optimized geometry
+        freqs = None
+        n_imag = 0
+        if compute_freqs:
+            data = {
+                "coord": opt_coords,
+                "numbers": atoms.numbers,
+                "charge": float(charge),
+            }
+            hess_result = get_base_calc()(data, hessian=True)
+            hessian = hess_result["hessian"].detach().cpu().numpy()
+            freqs, n_imag = compute_frequencies(hessian, atoms.numbers, opt_coords)
+
+    except Exception as e:
+        import traceback
+        return (f"**Calculation error:** {e}\n```\n{traceback.format_exc()}\n```", *empty[1:])
+
+    # Build outputs
+    viewer_html = build_viewer_html(opt_coords, numbers, charges_arr)
+    conv_plot = make_convergence_plot(trajectory) if trajectory else None
+
+    md = []
+    md.append("## Optimization Results\n")
+    status = "Converged" if converged else "Not converged"
+    if not converged:
+        md.append(f"> **{status}** after {len(trajectory)} steps / {wall_time:.1f}s "
+                  f"(final fmax: {fmax_final:.4f} eV/A)\n")
     else:
-        h = hessian
-
-    # Mass-weight the Hessian
-    masses = np.array([MASSES.get(int(z), 1.0) for z in numbers])
-    mass_vec = np.repeat(masses, 3)
-    sqrt_mass = np.sqrt(mass_vec)
-    h_mw = h / np.outer(sqrt_mass, sqrt_mass)
-
-    # Symmetrize
-    h_mw = 0.5 * (h_mw + h_mw.T)
-
-    # Diagonalize
-    eigenvalues = np.linalg.eigvalsh(h_mw)
-
-    # Convert: eigenvalue (eV/A^2/amu) -> frequency (cm^-1)
-    # 1 eV = 1.602176634e-19 J, 1 A = 1e-10 m, 1 amu = 1.66053906660e-27 kg
-    # freq = sqrt(eigenvalue * eV_to_J / (A_to_m^2 * amu_to_kg)) / (2*pi*c)
-    eV_to_J = 1.602176634e-19
-    amu_to_kg = 1.66053906660e-27
-    A_to_m = 1e-10
-    c_cm = 2.99792458e10  # speed of light in cm/s
-
-    conv = eV_to_J / (A_to_m**2 * amu_to_kg)
-    freqs = []
-    for ev in eigenvalues:
-        if ev >= 0:
-            f = np.sqrt(ev * conv) / (2 * np.pi * c_cm)
-        else:
-            f = -np.sqrt(-ev * conv) / (2 * np.pi * c_cm)
-        freqs.append(f)
-
-    freqs = np.array(freqs)
-    # Remove 6 (or 5 for linear) near-zero translational/rotational modes
-    freqs = freqs[np.abs(freqs) > 10.0]
-    return np.sort(freqs)
+        md.append(f"**{status}** in {len(trajectory)} steps ({wall_time:.1f}s)\n")
+
+    md.append("| Property | Initial | Final |")
+    md.append("|----------|--------:|------:|")
+    md.append(f"| Energy (eV) | {e0:.6f} | {e_final:.6f} |")
+    md.append(f"| Energy (kcal/mol) | {e0*EV_TO_KCAL:.4f} | {e_final*EV_TO_KCAL:.4f} |")
+    md.append(f"| Max |F| (eV/A) | {fmax0:.6f} | {fmax_final:.6f} |")
+    md.append(f"| dE (eV) | | {e_final - e0:.6f} |")
+    md.append("")
+
+    if charges_arr is not None:
+        md.append("### Partial Charges (e)\n| # | Elem | Charge |\n|--:|:----:|-------:|")
+        for i, (z, q) in enumerate(zip(numbers, charges_arr)):
+            sym = ELEMENT_SYMBOLS.get(int(z), "?")
+            md.append(f"| {i+1} | {sym} | {q:+.4f} |")
+        md.append(f"\n*Sum: {float(np.sum(charges_arr)):+.4f} e*\n")
+
+    freq_plot = None
+    if freqs is not None:
+        real_f = freqs[freqs > 0]
+        imag_f = freqs[freqs < 0]
+        md.append("### Vibrational Frequencies\n")
+        if n_imag > 0:
+            md.append(f"> **{n_imag} imaginary frequency(ies)** -- not a true minimum.\n")
+        if len(real_f) > 0:
+            md.append("```")
+            for j, f in enumerate(real_f):
+                md.append(f"  {j+1:3d}:  {f:10.2f} cm-1")
+            md.append("```")
+        if len(imag_f) > 0:
+            md.append("\nImaginary:\n```")
+            for j, f in enumerate(imag_f):
+                md.append(f"  {j+1:3d}:  {abs(f):10.2f}i cm-1")
+            md.append("```")
+        freq_plot = make_frequency_plot(freqs)
+
+    md.append("\n---")
+    md.append("*AIMNet2 wB97M-D3 | [Model](https://huggingface.co/isayevlab/aimnet2-wb97m-d3)*")
+
+    xyz_file = write_xyz_file(opt_coords, numbers, charges_arr,
+                              comment=f"Optimized, E={e_final:.6f} eV, fmax={fmax_final:.6f}")
+
+    script = generate_script(coords, numbers, charge, "optimize",
+                             max_steps=max_steps, fmax=fmax_target,
+                             compute_hessian=compute_freqs)
+
+    return "\n".join(md), viewer_html, conv_plot, freq_plot, xyz_file, script
 
 
 # ---------------------------------------------------------------------------
 # Gradio UI
 # ---------------------------------------------------------------------------
 
-EXAMPLES = [
-    ["CCO",           "SMILES",  0,  True, False],
-    ["c1ccccc1",      "SMILES",  0,  True, False],
-    ["CC(=O)O",       "SMILES",  0,  True, False],
-    ["[NH4+]",        "SMILES",  1,  True, False],
-    ["CC(=O)[O-]",    "SMILES", -1,  True, False],
-    ["O=C(O)c1ccccc1","SMILES",  0,  True, False],
-    ["C1CCCC1",       "SMILES",  0,  True,  True],
+CALC_EXAMPLES = [
+    ["CCO",            "SMILES",  0, True, False],
+    ["c1ccccc1",       "SMILES",  0, True, False],
+    ["CC(=O)O",        "SMILES",  0, True, False],
+    ["[NH4+]",         "SMILES",  1, True, False],
+    ["CC(=O)[O-]",     "SMILES", -1, True, False],
+    ["O=C(O)c1ccccc1", "SMILES",  0, True, False],
+    ["O",              "SMILES",  0, True,  True],
+]
+
+OPT_EXAMPLES = [
+    ["CCO", "SMILES", 0, 30, 0.05, False],
+    ["O",   "SMILES", 0, 30, 0.05,  True],
 ]
 
+VIEWER_PLACEHOLDER = (
+    '<div style="height:420px;display:flex;align-items:center;'
+    'justify-content:center;color:#aaa;border:1px solid #eee;'
+    'border-radius:8px;">Run a calculation to see the 3D structure</div>'
+)
+
 with gr.Blocks(title="AIMNet2 Demo", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         "# AIMNet2 Interactive Demo\n"
-        "Neural network potential for molecular property prediction: "
-        "**energy, forces, charges, frequencies**.  \n"
-        "Atoms are colored by predicted partial charge "
-        "(red = negative, blue = positive)."
+        "Neural network potential: **energy, forces, charges, optimization, frequencies**.  \n"
+        "3D viewer with charge coloring (red = negative, blue = positive)."
     )
 
+    # --- Shared input region ---
     with gr.Row():
         with gr.Column(scale=1):
             input_format = gr.Radio(
-                choices=["SMILES", "XYZ", "PDB"],
-                value="SMILES",
-                label="Input Format",
+                ["SMILES", "XYZ", "PDB"], value="SMILES", label="Input Format"
             )
-            input_text = gr.Textbox(
-                lines=6,
-                label="Molecule",
-                placeholder="Enter SMILES (e.g. CCO), XYZ block, or PDB block...",
+            input_text = gr.Textbox(lines=6, label="Molecule",
+                                    placeholder="SMILES, XYZ block, or PDB...")
+            file_upload = gr.File(
+                label="Or upload file",
+                file_types=[".xyz", ".pdb", ".sdf", ".mol"],
             )
-            charge = gr.Number(value=0, precision=0, label="Charge")
-            compute_forces = gr.Checkbox(value=True, label="Compute Forces")
-            compute_hessian = gr.Checkbox(value=False, label="Compute Hessian & Frequencies")
-            run_btn = gr.Button("Run AIMNet2", variant="primary")
+            charge_input = gr.Number(value=0, precision=0, label="Charge")
 
-        with gr.Column(scale=1):
-            viewer = gr.HTML(label="3D Structure", value="<div style='height:400px;display:flex;align-items:center;justify-content:center;color:#aaa;'>Run a calculation to see the 3D structure</div>")
-            output = gr.Markdown(label="Results")
+    # File upload handler
+    def on_file_upload(file_obj):
+        if file_obj is None:
+            return gr.update(), gr.update()
+        try:
+            text, fmt = handle_file_upload(file_obj)
+            gr.Info(f"Loaded file ({fmt} format)")
+            return gr.update(value=text), gr.update(value=fmt)
+        except Exception as e:
+            gr.Warning(f"File upload failed: {e}")
+            return gr.update(), gr.update()
 
-    gr.Examples(
-        examples=EXAMPLES,
-        inputs=[input_text, input_format, charge, compute_forces, compute_hessian],
-        label="Example Molecules",
+    file_upload.change(
+        on_file_upload, inputs=[file_upload], outputs=[input_text, input_format]
     )
 
-    run_btn.click(
-        fn=predict,
-        inputs=[input_text, input_format, charge, compute_forces, compute_hessian],
-        outputs=[output, viewer],
-    )
+    # --- Tabs ---
+    with gr.Tabs():
+        # ===== Tab 1: Calculate =====
+        with gr.TabItem("Calculate"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    calc_forces = gr.Checkbox(value=True, label="Compute Forces")
+                    calc_hessian = gr.Checkbox(value=False,
+                                               label="Compute Hessian & Frequencies")
+                    calc_btn = gr.Button("Calculate", variant="primary")
+
+                    gr.Examples(
+                        examples=CALC_EXAMPLES,
+                        inputs=[input_text, input_format, charge_input,
+                                calc_forces, calc_hessian],
+                        label="Examples",
+                    )
+
+                with gr.Column(scale=2):
+                    calc_viewer = gr.HTML(value=VIEWER_PLACEHOLDER, label="3D Structure")
+                    calc_results = gr.Markdown(label="Results")
+                    calc_freq_plot = gr.Plot(label="Frequency Spectrum", visible=False)
+                    with gr.Accordion("Download XYZ", open=False):
+                        calc_xyz = gr.File(label="XYZ file", interactive=False)
+                    with gr.Accordion("Python code to reproduce", open=False):
+                        calc_script = gr.Code(language="python", label="Script")
+
+            def calc_wrapper(text, fmt, charge, forces, hessian):
+                md, viewer, fplot, xyz, script = predict(text, fmt, charge, forces, hessian)
+                return (
+                    md,
+                    viewer or VIEWER_PLACEHOLDER,
+                    gr.update(value=fplot, visible=fplot is not None),
+                    xyz,
+                    script,
+                )
+
+            calc_btn.click(
+                calc_wrapper,
+                inputs=[input_text, input_format, charge_input, calc_forces, calc_hessian],
+                outputs=[calc_results, calc_viewer, calc_freq_plot, calc_xyz, calc_script],
+            )
+
+        # ===== Tab 2: Optimize =====
+        with gr.TabItem("Optimize"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    opt_steps = gr.Slider(10, 50, value=30, step=1, label="Max Steps")
+                    opt_fmax = gr.Number(value=0.05, label="Convergence fmax (eV/A)",
+                                         minimum=0.01, maximum=1.0)
+                    opt_freqs = gr.Checkbox(value=False,
+                                            label="Compute frequencies at minimum")
+                    opt_btn = gr.Button("Optimize", variant="primary")
+
+                    gr.Examples(
+                        examples=OPT_EXAMPLES,
+                        inputs=[input_text, input_format, charge_input,
+                                opt_steps, opt_fmax, opt_freqs],
+                        label="Examples",
+                    )
+
+                with gr.Column(scale=2):
+                    opt_viewer = gr.HTML(value=VIEWER_PLACEHOLDER, label="Optimized Structure")
+                    opt_conv_plot = gr.Plot(label="Convergence")
+                    opt_results = gr.Markdown(label="Results")
+                    opt_freq_plot = gr.Plot(label="Frequency Spectrum", visible=False)
+                    with gr.Accordion("Download optimized XYZ", open=False):
+                        opt_xyz = gr.File(label="XYZ file", interactive=False)
+                    with gr.Accordion("Python code to reproduce", open=False):
+                        opt_script = gr.Code(language="python", label="Script")
+
+            def opt_wrapper(text, fmt, charge, steps, fmax, freqs):
+                md, viewer, conv, fplot, xyz, script = optimize(
+                    text, fmt, charge, steps, fmax, freqs
+                )
+                return (
+                    md,
+                    viewer or VIEWER_PLACEHOLDER,
+                    conv,
+                    gr.update(value=fplot, visible=fplot is not None),
+                    xyz,
+                    script,
+                )
+
+            opt_btn.click(
+                opt_wrapper,
+                inputs=[input_text, input_format, charge_input,
+                        opt_steps, opt_fmax, opt_freqs],
+                outputs=[opt_results, opt_viewer, opt_conv_plot,
+                         opt_freq_plot, opt_xyz, opt_script],
+            )
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -6,3 +6,4 @@ ase==3.27.0
 rdkit
 numpy
 gradio>=5.0
+plotly>=5.0