molbuilder 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

molbuilder/visualization/electron_density_viz.py

@@ -0,0 +1,246 @@
+"""Electron cloud rendering for bond events.
+
+Uses an LCAO (Linear Combination of Atomic Orbitals) approximation to
+visualize electron density during bond formation and breaking.  The
+bonding orbital is modelled as:
+
+    psi_bond = c_a * phi_a(r - R_a) + c_b * phi_b(r - R_b)
+
+where phi are hydrogen-like orbitals evaluated with Slater effective
+nuclear charges, and c_a, c_b are mixing coefficients determined by
+the bond order.
+
+Reuses radial_wavefunction, real_spherical_harmonic, and
+wavefunction_real from molbuilder.atomic.wavefunctions.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from molbuilder.core.constants import BOHR_RADIUS_M
+from molbuilder.core.elements import SYMBOL_TO_Z
+from molbuilder.visualization.theme import ELECTRON_CLOUD_COLOR
+
+
+# Bohr radius in Angstroms
+_A0_ANGSTROM = BOHR_RADIUS_M * 1e10  # ~0.529 A
+
+
+def _hex_to_rgba(hex_color: str, alpha: float = 0.3) -> tuple:
+    """Convert hex color string to RGBA tuple (0-1 scale)."""
+    h = hex_color.lstrip("#")
+    r, g, b = (int(h[i:i+2], 16) / 255.0 for i in (0, 2, 4))
+    return (r, g, b, alpha)
+
+
+def _slater_zeff_simple(Z: int) -> float:
+    """Simplified Slater Z_eff for outermost valence s/p orbital.
+
+    Uses a rough approximation: Z_eff ~ Z - S where S is the total
+    screening.  For quick rendering this is sufficient.
+    """
+    # Very simplified: use known values for common elements
+    _ZEFF = {
+        1: 1.0, 6: 3.25, 7: 3.9, 8: 4.55, 9: 5.2,
+        15: 4.8, 16: 5.45, 17: 6.1, 35: 9.0, 53: 12.0,
+    }
+    return _ZEFF.get(Z, max(1.0, Z * 0.3))
+
+
+def _valence_nl(Z: int) -> tuple[int, int]:
+    """Return (n, l) for the outermost valence orbital."""
+    if Z <= 2:
+        return (1, 0)
+    if Z <= 10:
+        if Z <= 4:
+            return (2, 0)
+        return (2, 1)
+    if Z <= 18:
+        if Z <= 12:
+            return (3, 0)
+        return (3, 1)
+    if Z <= 36:
+        return (4, 0)
+    return (5, 0)
+
+
+class ElectronDensityRenderer:
+    """Renders electron density clouds using Monte Carlo point sampling.
+
+    Generates scatter points weighted by the LCAO bonding orbital
+    probability density, producing a visual representation of the
+    electron cloud during bond events.
+
+    Parameters
+    ----------
+    n_points : int
+        Number of scatter points to generate per bond.  More points
+        give smoother clouds but slower rendering.
+    """
+
+    def __init__(self, n_points: int = 5000):
+        self.n_points = n_points
+        self._rng = np.random.default_rng(42)
+
+    def compute_bond_density(self, pos_a: np.ndarray, pos_b: np.ndarray,
+                              z_a: int, z_b: int,
+                              bond_order: float = 1.0,
+                              ) -> tuple[np.ndarray, np.ndarray]:
+        """Compute electron cloud scatter points for a bond.
+
+        Uses rejection sampling: generates random points in the bonding
+        region and accepts them with probability proportional to
+        |psi_bond|^2.
+
+        Parameters
+        ----------
+        pos_a, pos_b : ndarray of shape (3,)
+            Atomic positions in Angstroms.
+        z_a, z_b : int
+            Atomic numbers.
+        bond_order : float
+            Fractional bond order (0 to 3).  Controls cloud density.
+
+        Returns
+        -------
+        points : ndarray of shape (n_accepted, 3)
+            Accepted scatter point positions.
+        colors : ndarray of shape (n_accepted, 4)
+            RGBA colors for each point.
+        """
+        if bond_order < 0.05:
+            return np.empty((0, 3)), np.empty((0, 4))
+
+        midpoint = 0.5 * (pos_a + pos_b)
+        bond_vec = pos_b - pos_a
+        bond_len = np.linalg.norm(bond_vec)
+        if bond_len < 1e-6:
+            return np.empty((0, 3)), np.empty((0, 4))
+
+        # Effective nuclear charges
+        zeff_a = _slater_zeff_simple(z_a)
+        zeff_b = _slater_zeff_simple(z_b)
+
+        # Valence orbital quantum numbers
+        n_a, l_a = _valence_nl(z_a)
+        n_b, l_b = _valence_nl(z_b)
+
+        # Sampling region: ellipsoid around the bond
+        spread = max(bond_len * 0.8, 1.5)
+
+        # Generate candidate points
+        pts = self._rng.normal(0, spread * 0.4, (self.n_points, 3))
+        pts += midpoint
+
+        # Evaluate approximate orbital values at each point
+        # Use simplified 1s-like orbitals for speed
+        r_a = np.linalg.norm(pts - pos_a, axis=1)
+        r_b = np.linalg.norm(pts - pos_b, axis=1)
+
+        # Slater-type orbital: phi ~ r^(n-1) * exp(-zeff * r / (n * a0))
+        decay_a = zeff_a / (n_a * _A0_ANGSTROM)
+        decay_b = zeff_b / (n_b * _A0_ANGSTROM)
+
+        phi_a = np.exp(-decay_a * r_a)
+        phi_b = np.exp(-decay_b * r_b)
+
+        # LCAO bonding combination
+        c_a = 1.0 / math.sqrt(2.0)
+        c_b = 1.0 / math.sqrt(2.0)
+        psi_bond = c_a * phi_a + c_b * phi_b
+        density = psi_bond ** 2
+
+        # Normalize and apply bond order scaling
+        max_dens = np.max(density)
+        if max_dens < 1e-30:
+            return np.empty((0, 3)), np.empty((0, 4))
+        prob = density / max_dens * min(bond_order, 1.5)
+
+        # Rejection sampling
+        rng_vals = self._rng.uniform(0, 1, self.n_points)
+        accept = rng_vals < prob
+
+        accepted_pts = pts[accept]
+        n_accepted = len(accepted_pts)
+
+        if n_accepted == 0:
+            return np.empty((0, 3)), np.empty((0, 4))
+
+        # Colors: base color with alpha proportional to density
+        base_rgba = _hex_to_rgba(ELECTRON_CLOUD_COLOR, 0.4)
+        colors = np.tile(base_rgba, (n_accepted, 1))
+        # Modulate alpha by density
+        accepted_density = density[accept]
+        alpha_scale = accepted_density / max_dens
+        colors[:, 3] = 0.1 + 0.4 * alpha_scale * min(bond_order, 1.5)
+
+        return accepted_pts, colors
+
+    def render_on_axis(self, ax, pos_a: np.ndarray, pos_b: np.ndarray,
+                       z_a: int, z_b: int,
+                       bond_order: float = 1.0,
+                       point_size: float = 2.0):
+        """Draw electron density cloud on a matplotlib 3D axis.
+
+        Parameters
+        ----------
+        ax : Axes3D
+            Matplotlib 3D axis to draw on.
+        pos_a, pos_b : ndarray of shape (3,)
+            Atomic positions in Angstroms.
+        z_a, z_b : int
+            Atomic numbers.
+        bond_order : float
+            Fractional bond order.
+        point_size : float
+            Size of scatter points.
+        """
+        points, colors = self.compute_bond_density(
+            pos_a, pos_b, z_a, z_b, bond_order)
+        if len(points) == 0:
+            return
+
+        ax.scatter(
+            points[:, 0], points[:, 1], points[:, 2],
+            c=colors,
+            s=point_size,
+            alpha=0.3,
+            depthshade=True,
+            edgecolors="none",
+        )
+
+    def render_on_axis_2d(self, ax, pos_a: np.ndarray, pos_b: np.ndarray,
+                           z_a: int, z_b: int,
+                           bond_order: float = 1.0,
+                           point_size: float = 2.0):
+        """Draw electron density cloud on a 2D matplotlib axis.
+
+        Uses the first two coordinates (x, y) for 2D projection.
+
+        Parameters
+        ----------
+        ax : Axes
+            Matplotlib 2D axis.
+        pos_a, pos_b : ndarray
+            Atomic positions.
+        z_a, z_b : int
+            Atomic numbers.
+        bond_order : float
+            Fractional bond order.
+        point_size : float
+            Scatter point size.
+        """
+        points, colors = self.compute_bond_density(
+            pos_a, pos_b, z_a, z_b, bond_order)
+        if len(points) == 0:
+            return
+
+        ax.scatter(
+            points[:, 0], points[:, 1],
+            c=colors,
+            s=point_size,
+            edgecolors="none",
+        )

molbuilder/visualization/interaction_controls.py

@@ -0,0 +1,211 @@
+"""Playback controls for the interaction visualizer.
+
+Provides keyboard bindings for interactive animation control and an
+optional tkinter panel for GUI embedding.
+
+Keyboard bindings:
+    Space     : Play / Pause
+    Right     : Step forward
+    Left      : Step backward
+    Up        : Increase speed (2x)
+    Down      : Decrease speed (0.5x)
+    E         : Toggle electron density
+    L         : Toggle labels
+    R         : Reset to beginning
+    Q / Esc   : Close
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from molbuilder.visualization.interaction_viz import InteractionVisualizer
+
+
+class PlaybackController:
+    """Keyboard-driven playback controller for InteractionVisualizer.
+
+    Connects to matplotlib key_press_event to provide interactive
+    playback controls.
+
+    Parameters
+    ----------
+    visualizer : InteractionVisualizer
+        The visualizer to control.
+    """
+
+    def __init__(self, visualizer: InteractionVisualizer):
+        self.viz = visualizer
+        self._speed_multiplier = 1.0
+        self._current_frame = 0
+        self._connected = False
+
+    def connect(self):
+        """Connect keyboard event handlers to the visualizer's figure."""
+        if self.viz.fig is None:
+            return
+        self.viz.fig.canvas.mpl_connect("key_press_event", self._on_key)
+        self._connected = True
+
+    def _on_key(self, event):
+        """Handle a key press event."""
+        if event.key == " ":
+            self._toggle_pause()
+        elif event.key == "right":
+            self._step_forward()
+        elif event.key == "left":
+            self._step_backward()
+        elif event.key == "up":
+            self._speed_up()
+        elif event.key == "down":
+            self._slow_down()
+        elif event.key == "e":
+            self._toggle_electron_density()
+        elif event.key == "l":
+            self._toggle_labels()
+        elif event.key == "r":
+            self._reset()
+        elif event.key in ("q", "escape"):
+            self._close()
+
+    def _toggle_pause(self):
+        """Toggle play/pause."""
+        anim = self.viz._anim
+        if anim is None:
+            return
+        if self.viz._paused:
+            anim.resume()
+            self.viz._paused = False
+        else:
+            anim.pause()
+            self.viz._paused = True
+
+    def _step_forward(self):
+        """Step one frame forward (pauses first)."""
+        anim = self.viz._anim
+        if anim is None:
+            return
+        if not self.viz._paused:
+            anim.pause()
+            self.viz._paused = True
+        self._current_frame = min(
+            self._current_frame + 1, self.viz.n_frames - 1)
+        self.viz._render_frame(self._current_frame)
+        self.viz.fig.canvas.draw_idle()
+
+    def _step_backward(self):
+        """Step one frame backward."""
+        anim = self.viz._anim
+        if anim is None:
+            return
+        if not self.viz._paused:
+            anim.pause()
+            self.viz._paused = True
+        self._current_frame = max(self._current_frame - 1, 0)
+        self.viz._render_frame(self._current_frame)
+        self.viz.fig.canvas.draw_idle()
+
+    def _speed_up(self):
+        """Double the playback speed."""
+        self._speed_multiplier *= 2.0
+        anim = self.viz._anim
+        if anim is not None:
+            new_interval = max(
+                1, int(1000 / (self.viz.config.fps * self._speed_multiplier)))
+            anim.event_source.interval = new_interval
+
+    def _slow_down(self):
+        """Halve the playback speed."""
+        self._speed_multiplier *= 0.5
+        anim = self.viz._anim
+        if anim is not None:
+            new_interval = int(
+                1000 / (self.viz.config.fps * self._speed_multiplier))
+            anim.event_source.interval = new_interval
+
+    def _toggle_electron_density(self):
+        """Toggle electron density cloud rendering."""
+        self.viz.config.show_electron_density = (
+            not self.viz.config.show_electron_density)
+
+    def _toggle_labels(self):
+        """Toggle time and energy labels."""
+        self.viz.config.show_time_label = not self.viz.config.show_time_label
+        self.viz.config.show_energy_bar = not self.viz.config.show_energy_bar
+
+    def _reset(self):
+        """Reset to the first frame."""
+        self._current_frame = 0
+        if self.viz._paused:
+            self.viz._render_frame(0)
+            self.viz.fig.canvas.draw_idle()
+
+    def _close(self):
+        """Close the figure."""
+        import matplotlib.pyplot as plt
+        plt.close(self.viz.fig)
+
+
+def create_tkinter_panel(parent, visualizer: InteractionVisualizer):
+    """Create an optional tkinter control panel for GUI embedding.
+
+    Parameters
+    ----------
+    parent : tk.Widget
+        Parent tkinter widget.
+    visualizer : InteractionVisualizer
+        The visualizer to control.
+
+    Returns
+    -------
+    tk.Frame
+        The control panel frame.
+    """
+    import tkinter as tk
+    from tkinter import ttk
+
+    controller = PlaybackController(visualizer)
+
+    frame = ttk.Frame(parent)
+
+    play_btn = ttk.Button(
+        frame, text="Play/Pause",
+        command=controller._toggle_pause)
+    play_btn.pack(side="left", padx=2)
+
+    step_back_btn = ttk.Button(
+        frame, text="<<",
+        command=controller._step_backward)
+    step_back_btn.pack(side="left", padx=2)
+
+    step_fwd_btn = ttk.Button(
+        frame, text=">>",
+        command=controller._step_forward)
+    step_fwd_btn.pack(side="left", padx=2)
+
+    slower_btn = ttk.Button(
+        frame, text="Slower",
+        command=controller._slow_down)
+    slower_btn.pack(side="left", padx=2)
+
+    faster_btn = ttk.Button(
+        frame, text="Faster",
+        command=controller._speed_up)
+    faster_btn.pack(side="left", padx=2)
+
+    edensity_var = tk.BooleanVar(value=visualizer.config.show_electron_density)
+    edensity_cb = ttk.Checkbutton(
+        frame, text="e- Density",
+        variable=edensity_var,
+        command=controller._toggle_electron_density)
+    edensity_cb.pack(side="left", padx=2)
+
+    labels_var = tk.BooleanVar(value=visualizer.config.show_time_label)
+    labels_cb = ttk.Checkbutton(
+        frame, text="Labels",
+        variable=labels_var,
+        command=controller._toggle_labels)
+    labels_cb.pack(side="left", padx=2)
+
+    return frame