Closed
Conversation
nabbelbabbel
requested review from
ConradJohnston,
sobolevnrm and
wavefunction91
as code owners
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new visualization utility for plotting orbital entanglement diagrams as chord diagrams. The implementation provides a comprehensive function plot_orbital_entanglement that visualizes single-orbital entropies and mutual information from wavefunction objects, using arc lengths to represent entropy magnitudes and chord lines to show pairwise mutual information between orbitals.
Changes:
- Adds
plot_orbital_entanglement()function with extensive customization options for creating circular chord diagrams - Implements label staggering for readability in systems with many orbitals
- Provides comprehensive test coverage with 19 test cases covering various scenarios and edge cases
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| python/src/qdk_chemistry/utils/visualization/orbital_entanglement.py | New module implementing the chord diagram visualization with helper functions for drawing arcs, outlines, and chord lines |
| python/src/qdk_chemistry/utils/visualization/init.py | New package initialization file exporting the main visualization function |
| python/src/qdk_chemistry/utils/init.py | Updated to expose the visualization function at the utils level |
| python/tests/test_orbital_entanglement_plot.py | Comprehensive test suite with mock wavefunctions and 19 test cases covering functionality and edge cases |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| The diagram encodes: | ||
| * **Arc length** - proportional to the single-orbital entropy | ||
| (sum of mutual information per orbital). |
There was a problem hiding this comment.
The parenthetical note "(sum of mutual information per orbital)" could be misleading. The arc length is directly proportional to the single-orbital entropy value from s1, not computed as a sum of MI values. Consider rephrasing to clarify the relationship, for example: "proportional to the single-orbital entropy" without the parenthetical, or "proportional to the single-orbital entropy (which relates to the orbital's entanglement with the rest of the system)".
Suggested change
| (sum of mutual information per orbital). | |
| (which reflects how strongly that orbital is entangled with the rest of the system). |
Comment on lines
+127
to
+315
| class TestPlotOrbitalEntanglement: | ||
| """Unit tests for ``plot_orbital_entanglement``.""" | ||
|
|
||
| def test_returns_figure_and_axes(self, small_wfn): | ||
| fig, ax = plot_orbital_entanglement(small_wfn) | ||
| assert fig is not None | ||
| assert ax is not None | ||
|
|
||
| assert isinstance(fig, matplotlib.figure.Figure) | ||
| assert isinstance(ax, matplotlib.axes.Axes) | ||
|
|
||
| def test_default_labels_are_indices(self, small_wfn): | ||
| """Default labels should be '0', '1', '2', '3' without active space.""" | ||
| _, ax = plot_orbital_entanglement(small_wfn) | ||
| texts = [t.get_text() for t in ax.texts] | ||
| assert texts == ["0", "1", "2", "3"] | ||
|
|
||
| def test_active_space_labels(self, small_wfn_active): | ||
| """When orbitals have an active space, labels should be those indices.""" | ||
| _, ax = plot_orbital_entanglement(small_wfn_active) | ||
| texts = [t.get_text() for t in ax.texts] | ||
| assert texts == ["3", "5", "7", "9"] | ||
|
|
||
| def test_custom_labels(self, small_wfn): | ||
| labels = ["\u03c3", "\u03c3*", "\u03c0", "\u03c0*"] | ||
| _, ax = plot_orbital_entanglement(small_wfn, labels=labels) | ||
| texts = [t.get_text() for t in ax.texts] | ||
| assert texts == labels | ||
|
|
||
| def test_wrong_label_count_raises(self, small_wfn): | ||
| with pytest.raises(ValueError, match="Number of labels"): | ||
| plot_orbital_entanglement(small_wfn, labels=["a", "b"]) | ||
|
|
||
| def test_missing_entropy_raises(self): | ||
| wfn = MockWavefunctionNoEntropy(np.zeros(2), np.zeros((2, 2))) | ||
| with pytest.raises(RuntimeError, match="single-orbital entropies"): | ||
| plot_orbital_entanglement(wfn) | ||
|
|
||
| def test_missing_mi_raises(self): | ||
| wfn = MockWavefunctionNoMI(np.zeros(2), np.zeros((2, 2))) | ||
| with pytest.raises(RuntimeError, match="mutual information"): | ||
| plot_orbital_entanglement(wfn) | ||
|
|
||
| def test_zero_entropy_does_not_crash(self, zero_entropy_wfn): | ||
| """All-zero entropies should produce equal arcs, not divide-by-zero.""" | ||
| fig, _ = plot_orbital_entanglement(zero_entropy_wfn) | ||
| assert fig is not None | ||
|
|
||
| def test_save_path_creates_file(self, small_wfn): | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| p = Path(tmpdir) / "test_out.png" | ||
| plot_orbital_entanglement(small_wfn, save_path=p) | ||
| assert p.exists() | ||
| assert p.stat().st_size > 0 | ||
|
|
||
| def test_save_svg(self, small_wfn): | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| p = Path(tmpdir) / "test_out.svg" | ||
| plot_orbital_entanglement(small_wfn, save_path=p) | ||
| assert p.exists() | ||
| assert p.stat().st_size > 0 | ||
|
|
||
| def test_existing_axes(self, small_wfn): | ||
| """Drawing into a user-supplied axes should work.""" | ||
| fig, ax = plt.subplots() | ||
| fig2, ax2 = plot_orbital_entanglement(small_wfn, ax=ax) | ||
| assert ax2 is ax | ||
| assert fig2 is fig | ||
|
|
||
| def test_mi_threshold_filters_chords(self): | ||
| """With a high threshold, weak chords should be omitted.""" | ||
| s1 = np.array([0.5, 0.5, 0.5]) | ||
| mi = np.array( | ||
| [ | ||
| [0.0, 0.01, 0.8], | ||
| [0.01, 0.0, 0.01], | ||
| [0.8, 0.01, 0.0], | ||
| ] | ||
| ) | ||
| wfn = MockWavefunction(s1, mi) | ||
| # With threshold = 0.5, only the (0,2) chord should survive. | ||
| _, ax = plot_orbital_entanglement(wfn, mi_threshold=0.5) | ||
| # Count PathPatch objects (chord lines) | ||
| chord_patches = [p for p in ax.patches if isinstance(p, mpatches.PathPatch)] | ||
| assert len(chord_patches) == 1 | ||
|
|
||
| def test_title_none_suppresses_title(self, small_wfn): | ||
| _, ax = plot_orbital_entanglement(small_wfn, title=None) | ||
| assert ax.get_title() == "" | ||
|
|
||
| def test_custom_title(self, small_wfn): | ||
| _, ax = plot_orbital_entanglement(small_wfn, title="My Plot") | ||
| assert ax.get_title() == "My Plot" | ||
|
|
||
| def test_s1_vmax_and_mi_vmax(self, small_wfn): | ||
| """Custom v-max values should not crash.""" | ||
| fig, _ = plot_orbital_entanglement( | ||
| small_wfn, | ||
| s1_vmax=2.0, | ||
| mi_vmax=3.0, | ||
| ) | ||
| assert fig is not None | ||
|
|
||
| def test_selected_indices_draws_outlines(self, small_wfn): | ||
| """Passing selected_indices should add outline patches.""" | ||
| _, ax1 = plot_orbital_entanglement(small_wfn) | ||
| n_patches_without = len(ax1.patches) | ||
|
|
||
| _, ax2 = plot_orbital_entanglement( | ||
| small_wfn, | ||
| selected_indices=[0, 2], | ||
| ) | ||
| n_patches_with = len(ax2.patches) | ||
| # Should have more patches when outlines are drawn | ||
| assert n_patches_with > n_patches_without | ||
|
|
||
| def test_selected_indices_with_active_space(self, small_wfn_active): | ||
| """selected_indices should match against label strings (active indices).""" | ||
| # Active indices are [3, 5, 7, 9]; select orbitals 5 and 9 | ||
| fig, _ = plot_orbital_entanglement( | ||
| small_wfn_active, | ||
| selected_indices=[5, 9], | ||
| ) | ||
| assert fig is not None | ||
|
|
||
| def test_selection_color_and_linewidth(self, small_wfn): | ||
| """Custom selection styling should not crash.""" | ||
| fig, _ = plot_orbital_entanglement( | ||
| small_wfn, | ||
| selected_indices=[1], | ||
| selection_color="green", | ||
| selection_linewidth=5.0, | ||
| ) | ||
| assert fig is not None | ||
|
|
||
| def test_large_system(self): | ||
| """Smoke test with a larger orbital count.""" | ||
| rng = np.random.default_rng(789) | ||
| n = 30 | ||
| s1 = rng.random(n) * np.log(4.0) | ||
| mi = _make_symmetric(n, rng) | ||
| wfn = MockWavefunction(s1, mi) | ||
| fig, _ = plot_orbital_entanglement( | ||
| wfn, | ||
| figsize=(12, 13), | ||
| gap_deg=1.5, | ||
| ) | ||
| assert fig is not None | ||
|
|
||
| def test_very_large_system_labels_staggered(self): | ||
| """With many orbitals, labels should be staggered, not dropped.""" | ||
| rng = np.random.default_rng(101) | ||
| n = 100 | ||
| # Make a few orbitals dominant, rest near-zero | ||
| s1 = np.full(n, 0.01) | ||
| s1[:5] = rng.random(5) * np.log(4.0) | ||
| mi = np.zeros((n, n)) | ||
| wfn = MockWavefunction(s1, mi) | ||
| _, ax = plot_orbital_entanglement(wfn, gap_deg=0.5) | ||
| # All labels should still be drawn (staggered, not skipped) | ||
| n_labels = len(ax.texts) | ||
| assert n_labels == n | ||
|
|
||
| def test_auto_line_scale(self): | ||
| """line_scale=None should auto-scale based on orbital count.""" | ||
| rng = np.random.default_rng(202) | ||
| for n_orbs in [4, 20, 100]: | ||
| s1 = rng.random(n_orbs) * np.log(4.0) | ||
| mi = _make_symmetric(n_orbs, rng) | ||
| wfn = MockWavefunction(s1, mi) | ||
| fig, ax = plot_orbital_entanglement(wfn) | ||
| assert fig is not None | ||
| plt.close(fig) | ||
|
|
||
| def test_explicit_line_scale_overrides_auto(self, small_wfn): | ||
| """Passing line_scale explicitly should be honoured.""" | ||
| fig, _ = plot_orbital_entanglement(small_wfn, line_scale=5.0) | ||
| assert fig is not None | ||
|
|
||
| def test_figsize_parameter(self, small_wfn): | ||
| fig, _ = plot_orbital_entanglement(small_wfn, figsize=(8, 9)) | ||
| w, h = fig.get_size_inches() | ||
| assert w == pytest.approx(8) | ||
| assert h == pytest.approx(9) | ||
|
|
||
| def test_gap_deg_zero(self, small_wfn): | ||
| """Zero gap should work (arcs touch).""" | ||
| fig, _ = plot_orbital_entanglement(small_wfn, gap_deg=0.0) | ||
| assert fig is not None |
There was a problem hiding this comment.
The tests should close created matplotlib figures to avoid resource leaks and memory accumulation during test runs. Only test_auto_line_scale closes figures (line 299). Consider adding plt.close(fig) or using a fixture with autouse=True that closes all figures after each test. This is especially important since the test suite creates many figures (including a 100-orbital one).
| edgecolor=selection_color, | ||
| linewidth=selection_linewidth, | ||
| zorder=3, | ||
| ) |
There was a problem hiding this comment.
The arc_mids list is computed twice identically (here and at line 344). Consider computing it once and reusing it throughout the function to avoid redundant calculations.
| ax.set_facecolor("none") | ||
|
|
||
| # 6. Draw outer arcs and labels | ||
| # Strategy: keep labels at a legible font size (≥7pt), and when |
There was a problem hiding this comment.
The function does not handle the edge case of an empty wavefunction (n=0). Line 338 would raise a ValueError: "max() arg is an empty sequence" if labels is empty. Consider either adding validation early to reject empty wavefunctions with a clear error message, or handle the edge case if it's a valid scenario. Note: This may be acceptable if empty wavefunctions are not physically meaningful.
Comment on lines
+267
to
+269
| # Both scales run from black (0) through a saturated colour at the | ||
| # low-to-mid range out to light grey at the theoretical maximum. | ||
| # This gives more colour resolution where values typically cluster. |
There was a problem hiding this comment.
The comment is inaccurate. The colormaps actually run from light grey at 0 (not black) through saturated color in the middle to dark/black at the maximum. The colors are: light grey (#d8d8d8) at 0, saturated red/blue at mid-range, and nearly black (#1a1a1a) at the maximum. Consider updating the comment to accurately describe the color progression.
Suggested change
| # Both scales run from black (0) through a saturated colour at the | |
| # low-to-mid range out to light grey at the theoretical maximum. | |
| # This gives more colour resolution where values typically cluster. | |
| # Both scales run from light grey (0) through a saturated red/blue | |
| # in the low-to-mid range out to a nearly black value at the | |
| # theoretical maximum. This gives more colour resolution where | |
| # values typically cluster. |
📊 Coverage Summary
Detailed Coverage ReportsC++ Coverage DetailsPython Coverage DetailsPybind11 Coverage Details |
Collaborator
Author
|
Closing in favor of: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a utility function to generate orbital entanglement diagrams such as: