From c8c6d26d7c70f3a4b2711f255b63be8420f4c2dc Mon Sep 17 00:00:00 2001 From: Jonathan Schultz Date: Sun, 16 Aug 2026 16:51:00 +0000 Subject: [PATCH] M-EXPORT2: reorg-geometry XYZ export + generalized PNG capture bridge (EXP2.1-2.3, EXP2.5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - EXP2.3: export_destination() — one place deciding where a new export lands and what it's called, scoped to new exporters only (the 3 existing structure exporters are left alone to avoid regressing working code). - EXP2.1: Export XYZ button on the reorg-geometry viewer, one file per distinct geometry (R_neutral + each R_ion), each carrying its own charge/multiplicity in the header comment. reorg_geometries() now threads charge/multiplicity through (the source payload already had it). - EXP2.2: generalized the isosurface's PNG capture bridge — _png_capture_controls now takes a capture_fn name instead of hardcoding the isosurface's bare window.__quantuiIsoCapture global (default preserved for exact backward compatibility). Wired a Save-PNG button into the reorg-geometry viewer with its own uid-scoped capture function and its own inbox, since every render gets a fresh uid and a shared global would risk one render's button capturing another's viewer. Molecule/trajectory/vibrational viewers are deliberately deferred — the roadmap's "cost of passing a capture_class" framing wasn't accurate until this generalization existed; this PR proves the reusable pattern on one viewer rather than retrofitting all of them. - EXP2.5: 20 new tests covering export_destination, the XYZ exporter, the PNG capture handler, and the capture-wiring HTML (isosurface default unchanged, reorg viewer gets a uid-scoped function, button omitted when capture_class is empty). Full suite: 2538 passed, 14 failed (pre-existing NMR failures, unrelated), 23 skipped. ruff + black + mypy (pre-push hook-stage) all clean. Contributions: - Claude (Opus 4.8): implementation, tests, verification - Jonathan Schultz: direction and review Co-authored-by: Jonathan Schultz Co-authored-by: Claude --- quantui/app.py | 21 ++ quantui/app_analysis.py | 11 +- quantui/app_builders.py | 28 +++ quantui/app_exports.py | 199 +++++++++++++++- quantui/app_visualization.py | 53 ++++- quantui/orbital_visualization.py | 23 +- quantui/reorganization_energy.py | 7 + tests/test_export2_reorg_and_destination.py | 250 ++++++++++++++++++++ 8 files changed, 584 insertions(+), 8 deletions(-) create mode 100644 tests/test_export2_reorg_and_destination.py diff --git a/quantui/app.py b/quantui/app.py index 043a21a..edd5c0b 100644 --- a/quantui/app.py +++ b/quantui/app.py @@ -147,6 +147,9 @@ from quantui.app_exports import ( on_export_pdb as _exp_on_export_pdb, ) +from quantui.app_exports import ( + on_export_reorg_geometries as _exp_on_export_reorg_geometries, +) from quantui.app_exports import ( on_export_xyz as _exp_on_export_xyz, ) @@ -156,6 +159,9 @@ from quantui.app_exports import ( on_orb_png_captured as _exp_on_orb_png_captured, ) +from quantui.app_exports import ( + on_reorg_png_captured as _exp_on_reorg_png_captured, +) from quantui.app_formatters import ( format_freq_result as _fmt_freq_result, ) @@ -1199,6 +1205,9 @@ class QuantUIApp: _reorg_overlay_pair: Any _reorg_exaggerate: Any _reorg_mode_dd: Any + _reorg_export_btn: Any + _reorg_export_status: Any + _reorg_png_inbox: Any _result_dir_label: Any _result_log_accordion: Any _result_log_output: Any @@ -2101,6 +2110,7 @@ def _wire_callbacks(self) -> None: self.export_xyz_btn.on_click(self._on_export_xyz) self.export_mol_btn.on_click(self._on_export_mol) self.export_pdb_btn.on_click(self._on_export_pdb) + self._reorg_export_btn.on_click(self._safe_cb(self._on_export_reorg_geometries)) # History self.past_dd.observe(self._safe_cb(self._on_past_dd_changed), names="value") self.past_refresh_btn.on_click(self._on_past_refresh) @@ -2202,6 +2212,11 @@ def _wire_callbacks(self) -> None: self._orb_png_inbox.observe( self._safe_cb(self._on_orb_png_captured), names="value" ) + # Same bridge, own inbox — the reorg-geometry viewer's Save-PNG button + # (M-EXPORT2 EXP2.2). + self._reorg_png_inbox.observe( + self._safe_cb(self._on_reorg_png_captured), names="value" + ) # Persist the grid choice so it survives a relaunch (ORBX.2). self._iso_resolution_dd.observe( self._safe_cb(self._on_iso_resolution_changed), names="value" @@ -3609,6 +3624,9 @@ def _on_export_mol(self, btn) -> None: def _on_export_pdb(self, btn) -> None: _exp_on_export_pdb(self, btn) + def _on_export_reorg_geometries(self, btn) -> None: + _exp_on_export_reorg_geometries(self, btn) + def _on_reorg_view_changed(self, change) -> None: _ana_on_reorg_view_changed(self, change) @@ -3621,6 +3639,9 @@ def _on_iso_appearance_changed(self, change) -> None: def _on_orb_png_captured(self, change) -> None: _exp_on_orb_png_captured(self, change) + def _on_reorg_png_captured(self, change) -> None: + _exp_on_reorg_png_captured(self, change) + def _on_iso_resolution_changed(self, change) -> None: """Persist the isosurface grid choice. diff --git a/quantui/app_analysis.py b/quantui/app_analysis.py index 7a0241d..1fca82e 100644 --- a/quantui/app_analysis.py +++ b/quantui/app_analysis.py @@ -339,7 +339,16 @@ def render_reorg_geometries(app: Any) -> None: exaggerate=float(getattr(app._reorg_exaggerate, "value", 1.0)), ) else: - html = build_reorg_geometry_viewer_html(geoms, bgcolor=bg) + from quantui.app_builders import _REORG_PNG_INBOX_CLASS + + capture_class = ( + _REORG_PNG_INBOX_CLASS + if getattr(app, "_reorg_png_inbox", None) is not None + else "" + ) + html = build_reorg_geometry_viewer_html( + geoms, bgcolor=bg, capture_class=capture_class + ) app._set_html_output(app._reorg_geom_output, html) except Exception as exc: # noqa: BLE001 app._set_html_output( diff --git a/quantui/app_builders.py b/quantui/app_builders.py index 5759d8c..9ed988d 100644 --- a/quantui/app_builders.py +++ b/quantui/app_builders.py @@ -29,6 +29,11 @@ # it is defined once here and passed down rather than spelled twice. _ORB_PNG_INBOX_CLASS = "quantui-orb-png-inbox" +# Same pattern, own inbox: the reorg-geometry viewer's Save-PNG button +# (M-EXPORT2 EXP2.2) writes here rather than sharing _ORB_PNG_INBOX_CLASS, so +# a capture from one viewer can never be mistaken for the other's. +_REORG_PNG_INBOX_CLASS = "quantui-reorg-png-inbox" + # Friendlier labels for the library category filter. _CATEGORY_LABELS = { "diatomic": "Diatomics", @@ -2229,6 +2234,24 @@ def _plot_export_row(prefix: str) -> widgets.HBox: style={"description_width": "70px"}, layout=layout_fn(width="380px", display="none"), ) + # Export the retained geometries as XYZ (M-EXPORT2 EXP2.1) — one file per + # distinct geometry (R_neutral + each R_ion), not per energy. + app._reorg_export_btn = widgets.Button( + description="Export XYZ", + icon="download", + layout=layout_fn(width="130px", margin="4px 0 0 0"), + tooltip="Save each geometry (R_neutral, R_ion, ...) as its own XYZ file", + ) + app._reorg_export_status = widgets.HTML( + value="", layout=layout_fn(margin="4px 0 0 8px") + ) + # Hidden inbox for the viewer's Save-PNG button (M-EXPORT2 EXP2.2) — same + # write-into-DOM-node-then-sync-to-kernel bridge as _orb_png_inbox, own + # class so the two capture flows can never cross-fire each other. + app._reorg_png_inbox = widgets.Textarea( + value="", layout=layout_fn(width="1px", height="1px", visibility="hidden") + ) + app._reorg_png_inbox.add_class(_REORG_PNG_INBOX_CLASS) app._reorg_geom_body = widgets.VBox( [ widgets.HTML( @@ -2242,6 +2265,11 @@ def _plot_export_row(prefix: str) -> widgets.HBox: app._reorg_overlay_pair, app._reorg_exaggerate, app._reorg_geom_output, + widgets.HBox( + [app._reorg_export_btn, app._reorg_export_status], + layout=layout_fn(align_items="center"), + ), + app._reorg_png_inbox, ], layout=layout_fn(padding="8px"), ) diff --git a/quantui/app_exports.py b/quantui/app_exports.py index 92e5cec..8d9826f 100644 --- a/quantui/app_exports.py +++ b/quantui/app_exports.py @@ -4,11 +4,67 @@ import logging from pathlib import Path -from typing import Any +from typing import Any, cast from .results_storage import _safe_name +def export_destination( + app: Any, + category: str, + *name_parts: str, + suffix: str, +) -> Path: + """The one place that decides where a new export lands and what it's + called (M-EXPORT2 EXP2.3). + + Before this, each exporter picked its own destination and built its own + filename inline (``on_export_xyz`` / ``on_export_mol`` / ``on_export_pdb`` + below all repeat the same three-line pattern), so "where did it save?" + had a different answer depending on which button was pressed. New + exporters should call this instead of repeating that pattern; the three + existing structure exporters are left as-is deliberately — they already + work, and retrofitting working export paths carries real regression risk + for no user-facing benefit. This is about not repeating the inconsistency + as the export surface grows (EXP2.1, and whatever comes after it). + + Every part of ``name_parts`` is sanitised (:func:`_safe_name`) and joined + with underscores, so a caller passes meaningful pieces (formula, a + geometry label, ...) instead of building a filename by hand. + + Files land next to the calculation's own results (``app._last_result_dir``) + so everything about one run stays in one folder. Unlike the existing + exporters (which fall back to the current working directory when no + result folder exists yet), this raises — silently writing outside the + result folder is a worse default for anything added from here on. + + Args: + app: the running QuantUIApp (only ``_last_result_dir`` is read). + category: a short, human-readable export kind, used only in the + error message (e.g. ``"reorg geometry"``). + *name_parts: filename-stem pieces, sanitised and joined with ``"_"``. + suffix: file extension including the leading dot (e.g. ``".xyz"``). + + Returns: + The full destination path (parent directory already exists, since it + is always an existing result directory). + + Raises: + ValueError: no result directory is available yet. + """ + result_dir = getattr(app, "_last_result_dir", None) + if result_dir is None or not isinstance(result_dir, Path): + raise ValueError( + f"No result folder yet — run a calculation before exporting {category}." + ) + stem = "_".join(_safe_name(str(p)) for p in name_parts if p) + # result_dir is narrowed to Path by the isinstance check above, but mypy's + # Path.__truediv__ overload resolution still infers Any from an + # originally-Any-typed (getattr on `app: Any`) operand — verified in + # isolation; the isinstance check is the real, working type guard. + return cast(Path, result_dir / f"{stem}{suffix}") + + def on_export(app: Any, btn: Any) -> None: """Export a standalone Python calculation script.""" if app._molecule is None: @@ -100,6 +156,73 @@ def on_export_pdb(app: Any, btn: Any) -> None: app.struct_export_status.value = f"Error: {exc}" +def on_export_reorg_geometries(app: Any, btn: Any) -> None: + """Export every retained reorg-energy geometry as its own XYZ file + (M-EXPORT2 EXP2.1). + + ``app._reorg_geometries`` (built by ``reorg_geometries()``) is already + deduplicated to the DISTINCT geometries behind a run — R_neutral once, + plus one R_ion per channel — so this writes exactly that many files, not + one per energy. Naming follows EXP2.1's request directly: + ``_R_neutral__.xyz`` / + ``_R_hole__.xyz`` rather than three files all + called ``geometry.xyz``. + + Provenance (EXP2.4): each file's comment line carries the charge, + multiplicity, method/basis, and which of the four λ energies were + evaluated on that geometry (``note`` — already computed by + ``reorg_geometries()`` for the viewer, reused here rather than + recomputed) — the whole point of a free-text XYZ comment line, and + otherwise unrecoverable from the file six months later. + """ + status = getattr(app, "_reorg_export_status", None) + + def _set_status(msg: str) -> None: + if status is not None: + status.value = msg + + geoms = getattr(app, "_reorg_geometries", None) + if not geoms: + _set_status("No geometries to export yet.") + return + + from quantui.molecule import Molecule + + method = app.method_dd.value + basis = app.basis_dd.value + saved: list[str] = [] + try: + for g in geoms: + # "R_neutral — optimized neutral" -> "R_neutral" + tag = g["label"].split(" — ")[0] + mol = Molecule( + atoms=list(g["atoms"]), + coordinates=[list(c) for c in g["coordinates"]], + charge=int(g.get("charge", 0)), + multiplicity=int(g.get("multiplicity", 1)), + ) + dest = export_destination( + app, + "reorg geometry", + mol.get_formula(), + tag, + method, + basis, + suffix=".xyz", + ) + comment = ( + f"{tag} charge={mol.charge} multiplicity={mol.multiplicity} " + f"{method}/{basis} {g.get('note', '')}" + ) + full_xyz = f"{len(mol.atoms)}\n{comment}\n{mol.to_xyz_string()}\n" + dest.write_text(full_xyz, encoding="utf-8") + saved.append(dest.name) + except Exception as exc: + _set_status(f"Error: {exc}") + return + _set_status(f"Saved {len(saved)} file(s): " + ", ".join(saved)) + + def export_molecule_and_label(app: Any) -> tuple[Any, str, str]: """Return (molecule, method, basis) for structure export. @@ -250,6 +373,80 @@ def _clear_inbox(a: Any) -> None: _clear_inbox(app) +def on_reorg_png_captured(app: Any, change: dict) -> None: + """Write a PNG captured from the live reorg-geometry viewer (M-EXPORT2 EXP2.2). + + Same capture/decode/save shape as ``on_orb_png_captured``, fed by its own + inbox (``_reorg_png_inbox`` / ``_REORG_PNG_INBOX_CLASS``) so a capture from + this viewer is never mistaken for an isosurface capture. Filename/location + goes through ``export_destination`` (EXP2.3) since this is a new exporter, + not a retrofit of existing behaviour. + + Deliberately skips the isosurface panel's DPI-stamping/custom-name extras + (``_iso_png_dpi`` / ``_iso_png_name``) — those are isosurface-panel + controls, and wiring them in here would make an unrelated panel's PNG + export depend on a setting the user set for a different viewer. + """ + import base64 + import binascii + + uri = (change or {}).get("new") or "" + if not uri: + return + + status = getattr(app, "_reorg_export_status", None) + + def _fail(msg: str) -> None: + if status is not None: + status.value = f'{msg}' + _clear_inbox(app) + + def _clear_inbox(a: Any) -> None: + box = getattr(a, "_reorg_png_inbox", None) + if box is not None and box.value: + box.value = "" + + if not uri.startswith(_PNG_URI_PREFIX): + logger.warning("reorg PNG capture: unexpected data URI prefix") + _fail("Capture failed (unexpected image format).") + return + if len(uri) > _MAX_PNG_BYTES: + _fail("Capture failed (image too large).") + return + + try: + raw = base64.b64decode(uri[len(_PNG_URI_PREFIX) :], validate=True) + except (binascii.Error, ValueError) as exc: + logger.warning("reorg PNG capture: could not decode payload: %s", exc) + _fail("Capture failed (corrupt image data).") + return + + mol = getattr(app, "_molecule", None) + formula = mol.get_formula() if mol is not None else "molecule" + method = str(getattr(getattr(app, "method_dd", None), "value", "") or "") + basis = str(getattr(getattr(app, "basis_dd", None), "value", "") or "") + + try: + dest = export_destination( + app, "reorg geometry PNG", formula, "geometry", method, basis, suffix=".png" + ) + except ValueError as exc: + _fail(str(exc)) + return + + try: + dest.write_bytes(raw) + except OSError as exc: + logger.warning("reorg PNG capture: could not write %s: %s", dest, exc) + _fail("Could not write the image (see log).") + return + + logger.info("Saved reorg geometry PNG: %s (%d bytes)", dest, len(raw)) + if status is not None: + status.value = f'Saved: {dest.name}' + _clear_inbox(app) + + def on_iso_export_cube(app: Any, btn: Any) -> None: """Copy the last-generated cube file to the result folder. diff --git a/quantui/app_visualization.py b/quantui/app_visualization.py index 046d18b..78aa914 100644 --- a/quantui/app_visualization.py +++ b/quantui/app_visualization.py @@ -14,6 +14,7 @@ from quantui import theme as _theme from quantui.app_builders import _ORB_PNG_INBOX_CLASS +from quantui.orbital_visualization import _png_capture_controls logger = logging.getLogger(__name__) @@ -2386,12 +2387,39 @@ def _frame_stepper_controls( return bar + f"" +# UID-scoped capture function for the reorg-geometry viewer (M-EXPORT2 EXP2.2). +# Every render of build_reorg_geometry_viewer_html gets a fresh uid (full +# atomic HTML swap — see app_analysis.render_reorg_geometries), so, unlike the +# isosurface viewer's single bare window.__quantuiIsoCapture global, this +# defines a per-render global named after the uid to avoid one render's button +# capturing a different render's (possibly already-detached) viewer. +_REORG_CAPTURE_JS = """ +(function(){ + var UID="__UID__"; + window["__CAPFN__"] = function(transparent){ + var vw = window["viewer_"+UID]; + if(!vw || !vw.pngURI){ return null; } + if(!transparent){ return vw.pngURI(); } + var uri=null; + try{ + vw.setBackgroundColor(__BG__, 0.0); vw.render(); + uri=vw.pngURI(); + } finally { + vw.setBackgroundColor(__BG__, 1.0); vw.render(); + } + return uri; + }; +})(); +""" + + def build_reorg_geometry_viewer_html( geometries: list[dict], *, bgcolor: str = "white", width: int = 560, height: int = 420, + capture_class: str = "", ) -> str: """Step through the DISTINCT geometries behind a Marcus 4-point run. @@ -2409,6 +2437,10 @@ def build_reorg_geometry_viewer_html( was never computed. Built on ``_frame_stepper_controls`` so the camera behaviour, offline loading and control styling match the trajectory viewer rather than being reinvented. + + ``capture_class`` wires a "Save PNG" button (M-EXPORT2 EXP2.2), mirroring + the isosurface viewer's capture bridge (ORBX.1) but with a uid-scoped + capture function — see ``_REORG_CAPTURE_JS``. Empty omits the button. """ import json import re @@ -2434,7 +2466,24 @@ def build_reorg_geometry_viewer_html( view_html = view._make_html() m = re.search(r"3dmolviewer_(\w+)", view_html) - if m is None or len(geometries) < 2: + if m is None: + return _theme.frame_viewer_html(view_html, width=width) + uid = m.group(1) + + if capture_class: + capture_fn = f"__quantuiReorgCapture_{uid}" + capture_js = ( + _REORG_CAPTURE_JS.replace("__UID__", uid) + .replace("__CAPFN__", capture_fn) + .replace("__BG__", json.dumps(bgcolor)) + ) + view_html = ( + view_html + + f"" + + _png_capture_controls(uid, capture_class, capture_fn=capture_fn) + ) + + if len(geometries) < 2: return _theme.frame_viewer_html(view_html, width=width) labels = json.dumps( @@ -2442,7 +2491,7 @@ def build_reorg_geometry_viewer_html( ) notes = json.dumps([g.get("note", "") for g in geometries]) controls = _frame_stepper_controls( - m.group(1), + uid, len(geometries), 1000, # unused: loop=False and Play is not the point here label_js=( diff --git a/quantui/orbital_visualization.py b/quantui/orbital_visualization.py index 5e560c4..4c42da9 100644 --- a/quantui/orbital_visualization.py +++ b/quantui/orbital_visualization.py @@ -1317,7 +1317,7 @@ def _build_iso_viewer( var btn=document.getElementById("orb_png_"+UID); if(!btn){ return; } btn.addEventListener("click", function(){ - var cap=window["__quantuiIsoCapture"]; + var cap=window["__CAPFN__"]; if(!cap){ btn.textContent="\\u26a0 viewer not ready"; return; } // Transparency is decided HERE, at capture, not by the live viewer — the // preview stays opaque while the exported file has no background. @@ -1339,9 +1339,24 @@ def _build_iso_viewer( """ -def _png_capture_controls(uid: str, capture_class: str) -> str: - """A 'Save PNG' button wired to the viewer identified by *uid*.""" - js = _PNG_CAPTURE_JS.replace("__UID__", uid).replace("__CLS__", capture_class) +def _png_capture_controls( + uid: str, capture_class: str, capture_fn: str = "__quantuiIsoCapture" +) -> str: + """A 'Save PNG' button wired to the viewer identified by *uid*. + + *capture_fn* is the name of the global JS function (already defined + elsewhere, e.g. ``window[capture_fn] = function(transparent){...}``) that + does the actual ``pngURI()`` capture. Defaults to the isosurface viewer's + bare, unscoped global for backward compatibility (ORBX.1). Callers with + multiple live viewers of the same kind on a page (e.g. a fresh uid per + render) must pass a uid-scoped name to avoid one viewer's button + capturing another viewer's frame. + """ + js = ( + _PNG_CAPTURE_JS.replace("__UID__", uid) + .replace("__CLS__", capture_class) + .replace("__CAPFN__", capture_fn) + ) return ( f'
' f'