From 85a83acc1412e6473640a37905a038019f9f1c76 Mon Sep 17 00:00:00 2001
From: NCCU Schultz Lab <nccu-schultz-lab@users.noreply.github.com>
Date: Sun, 30 Aug 2026 15:48:20 +0000
Subject: [PATCH] Fix blank Mulliken 3D viewer with geometry persistence and
 lazy render

The Mulliken Populations panel now resolves geometry from result.json,
history context, and disk sidecars; defers py3Dmol rendering until the
accordion is visible; shows an inline message when coordinates are missing;
and refreshes on theme/backend re-render like other Analysis viewers.

Co-authored-by: Jonathan Schultz <jonathanschultzNU@users.noreply.github.com>
---
 quantui/app.py                 |  6 +++
 quantui/app_analysis.py        | 21 ++++++--
 quantui/app_history.py         | 13 ++++-
 quantui/populations_overlay.py | 36 +++++++++----
 quantui/results_storage.py     | 82 ++++++++++++++++++++++++++++
 tests/test_mulliken_panel.py   | 98 ++++++++++++++++++++++++++++++++++
 tests/test_results_storage.py  | 33 ++++++++++++
 7 files changed, 275 insertions(+), 14 deletions(-)

diff --git a/quantui/app.py b/quantui/app.py
index af0d22e..0c03e0f 100644
--- a/quantui/app.py
+++ b/quantui/app.py
@@ -1696,6 +1696,7 @@ def __init__(self) -> None:
         # Analysis-tab viewer when the toggle changes.
         self._analysis_displayed_molecule: Any = None
         self._mulliken_displayed_molecule: Any = None
+        self._mulliken_pending_molecule: Any = None
 
         # ── Build → wire → assemble ───────────────────────────────────────
         self._build_widgets()
@@ -3636,6 +3637,10 @@ def _rerender_3d_views(self) -> None:
                 self._set_html_output(self._analysis_mol_output, html)
                 self._update_analysis_backend_label(chosen)
 
+        # Mulliken Populations dedicated viewer (py3Dmol charge overlays).
+        if getattr(self, "_last_mulliken_charges", None):
+            self._show_mulliken_viewer()
+
     def _show_mulliken_viewer(self, molecule=None) -> None:
         from quantui.populations_overlay import show_mulliken_viewer
 
@@ -5991,6 +5996,7 @@ def _run_required_final_single_point(target_mol, reason: str):
                     pyscf_log=log.getvalue(),
                     calc_type=save_type,
                     spectra=save_spectra,
+                    molecule=calc_mol,
                 )
                 _run_saved = True
                 self._last_result_dir = _saved_dir
diff --git a/quantui/app_analysis.py b/quantui/app_analysis.py
index ee8e53c..354e0ab 100644
--- a/quantui/app_analysis.py
+++ b/quantui/app_analysis.py
@@ -255,6 +255,7 @@ def apply_analysis_context(app: Any, ctx: Any) -> None:
     if mulliken_out is not None:
         mulliken_out.clear_output()
     app._mulliken_displayed_molecule = None
+    app._mulliken_pending_molecule = None
     note = getattr(app, "_mulliken_overlay_note", None)
     if note is not None:
         note.value = ""
@@ -855,6 +856,18 @@ def _mulliken_molecule(app: Any, ctx: Any = None) -> Any:
         mol = getattr(ctx, "molecule", None)
         if mol is not None:
             return mol
+        result_dir = getattr(ctx, "result_dir", None)
+        if result_dir is not None:
+            try:
+                from quantui import load_result
+                from quantui.app_history import mol_from_result_dir
+
+                data = load_result(result_dir)
+                mol = mol_from_result_dir(result_dir, data)
+                if mol is not None:
+                    return mol
+            except Exception:
+                pass
     return getattr(app, "_analysis_displayed_molecule", None) or getattr(
         app, "_molecule", None
     )
@@ -958,10 +971,10 @@ def show_mulliken_populations(
     )
 
     update_mulliken_figure(app)
-    try:
-        app._show_mulliken_viewer(molecule)
-    except Exception:  # noqa: BLE001 — viewer must never block the panel
-        pass
+    # Cache geometry for lazy render when the accordion becomes visible.
+    # Voilà often skips <script> execution in hidden DOM; the authoritative
+    # draw happens in ``_on_mulliken_accordion_show``.
+    app._mulliken_pending_molecule = molecule
     return True
 
 
diff --git a/quantui/app_history.py b/quantui/app_history.py
index f0d7630..4e2a05d 100644
--- a/quantui/app_history.py
+++ b/quantui/app_history.py
@@ -387,12 +387,21 @@ def mol_from_result_dir(result_dir: Path, data: dict[str, Any]) -> Any:
     """Try to reconstruct a displayable Molecule from a saved result directory.
 
     Returns a Molecule or None if geometry data is not available.
-    Tries sources in order: frequency spectra -> orbitals_meta -> trajectory.
+    Tries sources in order: result.json geometry -> frequency spectra ->
+    orbitals_meta -> trajectory.
     """
     from quantui.molecule import Molecule
+    from quantui.results_storage import molecule_from_geometry_payload
 
     calc_type = data.get("calc_type", "")
 
+    geom = data.get("geometry")
+    if isinstance(geom, dict) and geom.get("atoms") and geom.get("coordinates"):
+        try:
+            return molecule_from_geometry_payload(geom)
+        except Exception:
+            pass
+
     # Frequency: geometry stored inside spectra.molecule
     if calc_type == "frequency":
         mol_data = data.get("spectra", {}).get("molecule", {})
@@ -614,6 +623,7 @@ def build_history_context(result_dir: Path, *, context_cls: Any) -> Optional[Any
         data = load_result(result_dir)
     except Exception:
         return None
+    molecule = mol_from_result_dir(result_dir, data)
     return context_cls(
         calc_type=data.get("calc_type", ""),
         formula=data.get("formula", result_dir.name),
@@ -623,4 +633,5 @@ def build_history_context(result_dir: Path, *, context_cls: Any) -> Optional[Any
         spectra_data=data.get("spectra", {}),
         timestamp=data.get("timestamp", ""),
         source="history",
+        molecule=molecule,
     )
diff --git a/quantui/populations_overlay.py b/quantui/populations_overlay.py
index 9a18863..fbb17f9 100644
--- a/quantui/populations_overlay.py
+++ b/quantui/populations_overlay.py
@@ -225,6 +225,27 @@ def _mulliken_vividness(app: Any) -> float:
     return float(getattr(slider, "value", 1.0))
 
 
+def _mulliken_viewer_unavailable_html() -> str:
+    style = "padding:12px 16px;color:#6b7280;font-size:13px;font-style:italic"
+    return (
+        f'<div style="{style}">'
+        "3D structure unavailable: this result has no saved coordinates. "
+        "Re-run the calculation to enable the charge-coloured viewer "
+        "(table and chart above are still valid).</div>"
+    )
+
+
+def _resolve_mulliken_molecule(app: Any, molecule: Any = None) -> Any:
+    """Return the best molecule for the Mulliken viewer slot."""
+    return (
+        molecule
+        or getattr(app, "_mulliken_pending_molecule", None)
+        or getattr(app, "_mulliken_displayed_molecule", None)
+        or getattr(app, "_analysis_displayed_molecule", None)
+        or getattr(app, "_molecule", None)
+    )
+
+
 def render_mulliken_viewer_html(app: Any, molecule: Any, *, render_html_fn: Any) -> str:
     """Return self-contained HTML for the Mulliken panel's py3Dmol viewer."""
     if render_html_fn is None or molecule is None:
@@ -239,6 +260,7 @@ def render_mulliken_viewer_html(app: Any, molecule: Any, *, render_html_fn: Any)
             style=app._viz_style,
             lighting=app._viz_lighting,
             bgcolor=bgcolor,
+            show_info=False,
         )
         return inject_populations_js(html)
     html = render_html_fn(
@@ -247,6 +269,7 @@ def render_mulliken_viewer_html(app: Any, molecule: Any, *, render_html_fn: Any)
         style=app._viz_style,
         lighting=app._viz_lighting,
         bgcolor=bgcolor,
+        show_info=False,
     )
     return cast(str, html)
 
@@ -258,27 +281,22 @@ def show_mulliken_viewer(
     out = getattr(app, "_mulliken_mol_output", None)
     if out is None:
         return
-    mol = (
-        molecule
-        or getattr(app, "_mulliken_displayed_molecule", None)
-        or getattr(app, "_analysis_displayed_molecule", None)
-        or getattr(app, "_molecule", None)
-    )
+    mol = _resolve_mulliken_molecule(app, molecule)
     if mol is None:
-        out.clear_output()
+        app._set_html_output(out, _mulliken_viewer_unavailable_html())
         app._mulliken_displayed_molecule = None
         return
     try:
         html = render_mulliken_viewer_html(app, mol, render_html_fn=render_html_fn)
         if not html:
-            out.clear_output()
+            app._set_html_output(out, _mulliken_viewer_unavailable_html())
             return
         app._set_html_output(out, html)
         app._mulliken_displayed_molecule = mol
         push_populations_overlay(app)
     except Exception as exc:  # noqa: BLE001 — viewer must never block the panel
         logger.debug("mulliken viewer render failed: %s", exc)
-        out.clear_output()
+        app._set_html_output(out, _mulliken_viewer_unavailable_html())
 
 
 def push_populations_overlay(app: Any) -> None:
diff --git a/quantui/results_storage.py b/quantui/results_storage.py
index 7f30113..8f11cc2 100644
--- a/quantui/results_storage.py
+++ b/quantui/results_storage.py
@@ -48,6 +48,79 @@
 _ANGSTROM_TO_BOHR = 1.0 / _BOHR_TO_ANGSTROM
 
 
+def _geometry_payload_from_molecule(molecule: Any) -> dict:
+    """Return a JSON-safe geometry dict from a :class:`~quantui.molecule.Molecule`."""
+    return {
+        "atoms": list(molecule.atoms),
+        "coordinates": [list(map(float, row)) for row in molecule.coordinates],
+        "charge": int(getattr(molecule, "charge", 0) or 0),
+        "multiplicity": int(getattr(molecule, "multiplicity", 1) or 1),
+    }
+
+
+def geometry_payload_for_result(result: object, molecule: Any = None) -> Optional[dict]:
+    """Best-effort geometry extraction for ``result.json`` persistence.
+
+    Prefers an explicit *molecule* (the geometry the calc actually used),
+    then duck-typed attributes on *result* (``molecule``, ``pyscf_mol_atom``,
+    final trajectory frame, …). Returns ``None`` when no coordinates are found.
+    """
+    if molecule is not None:
+        try:
+            return _geometry_payload_from_molecule(molecule)
+        except Exception:
+            return None
+
+    mol = getattr(result, "molecule", None)
+    if mol is not None:
+        try:
+            return _geometry_payload_from_molecule(mol)
+        except Exception:
+            pass
+
+    pyscf_mol_atom = getattr(result, "pyscf_mol_atom", None)
+    if pyscf_mol_atom:
+        try:
+            atoms = [str(sym) for sym, _ in pyscf_mol_atom]
+            coordinates = [list(map(float, coords)) for _, coords in pyscf_mol_atom]
+            if atoms and coordinates and len(atoms) == len(coordinates):
+                return {
+                    "atoms": atoms,
+                    "coordinates": coordinates,
+                    "charge": int(getattr(result, "charge", 0) or 0),
+                    "multiplicity": int(getattr(result, "multiplicity", 1) or 1),
+                }
+        except Exception:
+            pass
+
+    for traj_attr in ("trajectory", "coordinates_list"):
+        traj = getattr(result, traj_attr, None)
+        if traj:
+            try:
+                last = traj[-1]
+                return _geometry_payload_from_molecule(last)
+            except Exception:
+                pass
+
+    return None
+
+
+def molecule_from_geometry_payload(payload: dict) -> Any:
+    """Reconstruct a :class:`~quantui.molecule.Molecule` from a saved geometry dict."""
+    from quantui.molecule import Molecule
+
+    atoms = payload.get("atoms") or payload.get("atom_symbols")
+    coordinates = payload.get("coordinates") or payload.get("coords")
+    if not atoms or not coordinates:
+        raise ValueError("geometry payload missing atoms or coordinates")
+    return Molecule(
+        atoms=list(atoms),
+        coordinates=[list(map(float, row)) for row in coordinates],
+        charge=int(payload.get("charge", 0) or 0),
+        multiplicity=int(payload.get("multiplicity", 1) or 1),
+    )
+
+
 def _default_results_dir() -> Path:
     env = os.environ.get("QUANTUI_RESULTS_DIR")
     return Path(env) if env else Path("results")
@@ -168,6 +241,7 @@ def save_result(
     calc_type: str = "single_point",
     spectra: Optional[dict] = None,
     extras: Optional[dict] = None,
+    molecule: Any = None,
 ) -> Path:
     """Write *result* to a new timestamped subdirectory of *results_dir*.
 
@@ -201,6 +275,11 @@ def save_result(
         Keys clash with built-in result.json fields (``timestamp``,
         ``formula``, etc.) overwrite them — by design, since the
         caller is asserting they want to override.
+    molecule:
+        Optional :class:`~quantui.molecule.Molecule` for the geometry used
+        in this calculation. When provided, written to the top-level
+        ``geometry`` field in ``result.json`` so History replay and the
+        Mulliken panel viewer do not depend on ``orbitals_meta.json``.
 
     Returns
     -------
@@ -289,6 +368,9 @@ def save_result(
         "atom_symbols": _opt_str_list(getattr(result, "atom_symbols", None)),
         "spectra": spectra if spectra is not None else {},
     }
+    _geom = geometry_payload_for_result(result, molecule=molecule)
+    if _geom is not None:
+        data["geometry"] = _geom
     if extras:
         data.update(extras)
     (dest / "result.json").write_text(json.dumps(data, indent=2))
diff --git a/tests/test_mulliken_panel.py b/tests/test_mulliken_panel.py
index a262719..c1a9e0b 100644
--- a/tests/test_mulliken_panel.py
+++ b/tests/test_mulliken_panel.py
@@ -200,6 +200,104 @@ def test_writes_table_and_caches_state(self, app):
         assert "H2" in app._mulliken_table.value
         assert app._last_mulliken_dipole == pytest.approx(1.2)
 
+    def test_lazy_render_defers_viewer_until_accordion_expand(self, app):
+        mol = _water()
+        show_mulliken_populations(
+            app, ["O", "H", "H"], [-0.5, 0.25, 0.25], molecule=mol
+        )
+        assert app._mulliken_pending_molecule is mol
+        assert len(app._mulliken_mol_output.outputs) == 0
+        app._on_mulliken_accordion_show({"new": 0})
+        assert len(app._mulliken_mol_output.outputs) == 1
+        html = app._mulliken_mol_output.outputs[0]["data"]["text/html"]
+        assert html.strip() != ""
+        assert app._mulliken_displayed_molecule is mol
+
+
+class TestMullikenViewerRender:
+    def test_live_populate_renders_on_accordion_expand(self, app):
+        from quantui.app import _render_molecule_html
+
+        if _render_molecule_html is None:
+            pytest.skip("No 3D visualization backend installed")
+        ctx = SimpleNamespace(
+            calc_type="single_point",
+            live_result=_sp_result_with_charges(),
+            result_dir=None,
+            spectra_data={},
+            source="live",
+            formula="H2O",
+            method="RHF",
+            basis="STO-3G",
+            label="H2O",
+            timestamp="",
+            molecule=_water(),
+        )
+        assert pop_mulliken(app, ctx) is True
+        app._on_mulliken_accordion_show({"new": 0})
+        assert len(app._mulliken_mol_output.outputs) == 1
+        html = app._mulliken_mol_output.outputs[0]["data"]["text/html"]
+        assert "3D structure unavailable" not in html
+
+    def test_history_geometry_in_result_json_renders_without_orbitals(
+        self, tmp_path, app
+    ):
+        from quantui.app import _render_molecule_html
+
+        if _render_molecule_html is None:
+            pytest.skip("No 3D visualization backend installed")
+        mol = _water()
+        saved = save_result(
+            _sp_result_with_charges(),
+            results_dir=tmp_path,
+            calc_type="single_point",
+            spectra={},
+            molecule=mol,
+        )
+        for p in saved.glob("orbitals*"):
+            p.unlink()
+        ctx = app._build_history_context(saved)
+        assert ctx.molecule is not None
+        app._apply_analysis_context(ctx)
+        app._on_mulliken_accordion_show({"new": 0})
+        assert len(app._mulliken_mol_output.outputs) == 1
+        html = app._mulliken_mol_output.outputs[0]["data"]["text/html"]
+        assert "3D structure unavailable" not in html
+
+    def test_missing_geometry_shows_unavailable_message(self, tmp_path, app):
+        app._molecule = None
+        app._analysis_displayed_molecule = None
+        saved = save_result(
+            _sp_result_with_charges(),
+            results_dir=tmp_path,
+            calc_type="single_point",
+            spectra={},
+        )
+        for p in saved.glob("orbitals*"):
+            p.unlink()
+        ctx = app._build_history_context(saved)
+        app._apply_analysis_context(ctx)
+        app._on_mulliken_accordion_show({"new": 0})
+        assert len(app._mulliken_mol_output.outputs) == 1
+        html = app._mulliken_mol_output.outputs[0]["data"]["text/html"]
+        assert "3D structure unavailable" in html
+        assert "O1" in app._mulliken_table.value
+
+    def test_theme_rerender_refreshes_mulliken_viewer(self, app):
+        from quantui.app import _render_molecule_html
+
+        if _render_molecule_html is None:
+            pytest.skip("No 3D visualization backend installed")
+        show_mulliken_populations(
+            app, ["O", "H", "H"], [-0.5, 0.25, 0.25], molecule=_water()
+        )
+        app._on_mulliken_accordion_show({"new": 0})
+        before = app._mulliken_mol_output.outputs[0]["data"]["text/html"]
+        app._rerender_3d_views()
+        after = app._mulliken_mol_output.outputs[0]["data"]["text/html"]
+        assert before.strip() != ""
+        assert after.strip() != ""
+
 
 class TestHelpTopic:
     def test_mulliken_topic_exists(self):
diff --git a/tests/test_results_storage.py b/tests/test_results_storage.py
index 411365c..b53654f 100644
--- a/tests/test_results_storage.py
+++ b/tests/test_results_storage.py
@@ -16,10 +16,12 @@
 
 from quantui.molecule import Molecule
 from quantui.results_storage import (
+    geometry_payload_for_result,
     list_results,
     load_orbitals,
     load_result,
     load_trajectory,
+    molecule_from_geometry_payload,
     save_orbitals,
     save_result,
     save_thumbnail,
@@ -78,6 +80,37 @@ def test_result_json_fields(self, tmp_path):
         assert data["calc_type"] == "single_point"
         assert data["_schema_version"] == 2
 
+    def test_geometry_persisted_when_molecule_provided(self, tmp_path):
+        mol = Molecule(
+            ["O", "H", "H"],
+            [[0.0, 0.0, 0.0], [0.757, 0.586, 0.0], [-0.757, 0.586, 0.0]],
+        )
+        saved = save_result(
+            _make_result(), results_dir=tmp_path, calc_type="single_point", molecule=mol
+        )
+        data = json.loads((saved / "result.json").read_text())
+        assert data["geometry"]["atoms"] == ["O", "H", "H"]
+        assert len(data["geometry"]["coordinates"]) == 3
+        restored = molecule_from_geometry_payload(data["geometry"])
+        assert restored.atoms == mol.atoms
+        for row, expected in zip(restored.coordinates, mol.coordinates):
+            assert row == pytest.approx(expected)
+
+    def test_geometry_payload_from_pyscf_mol_atom(self, tmp_path):
+        result = _make_result(
+            pyscf_mol_atom=[
+                ("O", [0.0, 0.0, 0.0]),
+                ("H", [0.757, 0.586, 0.0]),
+                ("H", [-0.757, 0.586, 0.0]),
+            ]
+        )
+        payload = geometry_payload_for_result(result)
+        assert payload is not None
+        assert payload["atoms"] == ["O", "H", "H"]
+        saved = save_result(result, results_dir=tmp_path)
+        data = json.loads((saved / "result.json").read_text())
+        assert "geometry" in data
+
     def test_spectra_stored_in_json(self, tmp_path):
         spectra = {
             "ir": {"frequencies_cm1": [1000.0, 2000.0], "ir_intensities": [1.0, 2.0]}