Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3           1 : from __future__ import annotations
       4             : 
       5           1 : from typing import TYPE_CHECKING
       6             : 
       7           1 : import matplotlib as mpl
       8           1 : import matplotlib.pyplot as plt
       9           1 : import numpy as np
      10             : 
      11           1 : from pairinteraction.ket.ket_pair import KetPair, is_ket_atom_tuple
      12           1 : from pairinteraction.visualization.colormaps import alphamagma
      13             : 
      14             : if TYPE_CHECKING:
      15             :     from collections.abc import Sequence
      16             : 
      17             :     from matplotlib.axes import Axes
      18             :     from matplotlib.figure import Figure
      19             : 
      20             :     from pairinteraction.ket.ket_pair import KetPairLike
      21             :     from pairinteraction.system.system_pair import SystemPair
      22             :     from pairinteraction.units import NDArray
      23             : 
      24             : 
      25           1 : def plot_pair_potential(
      26             :     system_pairs: Sequence[SystemPair],
      27             :     ket_pair_of_interest: KetPairLike,
      28             :     ax: Axes | None = None,
      29             : ) -> tuple[Figure, Axes]:
      30             :     """Plot the pair potential of a sequence of SystemPair objects.
      31             : 
      32             :     Args:
      33             :         system_pairs: A sequence of SystemPair objects to plot.
      34             :         ket_pair_of_interest: The KetPair (or tuple of KetAtom) for whicht the overlaps will be plotted.
      35             :         ax: An optional matplotlib Axes object to plot on.
      36             :             If None, a new figure and axes will be created.
      37             : 
      38             :     """
      39           0 :     distances: NDArray = np.array([system.get_distance(unit="micrometer") for system in system_pairs])
      40             : 
      41           0 :     if is_ket_atom_tuple(ket_pair_of_interest):
      42           0 :         energy_of_interest = sum(ket.get_energy(unit="GHz") for ket in ket_pair_of_interest)
      43           0 :     elif isinstance(ket_pair_of_interest, KetPair):
      44           0 :         energy_of_interest = ket_pair_of_interest.get_energy(unit="GHz")
      45             :     else:
      46           0 :         raise TypeError("ket_pair_of_interest must be a KetPair or a tuple of KetAtom.")
      47             : 
      48           0 :     eigenenergies = [system.get_eigenenergies(unit="GHz") - energy_of_interest for system in system_pairs]
      49           0 :     overlaps = [system.get_eigenbasis().get_overlaps(ket_pair_of_interest) for system in system_pairs]
      50             : 
      51           0 :     if ax is None:
      52           0 :         owns_ax = True
      53           0 :         fig, ax = plt.subplots()
      54             :     else:
      55           0 :         owns_ax = False
      56           0 :         fig = ax.figure  # type: ignore[assignment]
      57             : 
      58           0 :     lw = mpl.rcParams["lines.linewidth"]
      59           0 :     ax.plot(distances, np.array(eigenenergies), c="k", lw=lw / 6, zorder=-10)
      60             : 
      61           0 :     x_repeated = np.hstack([val * np.ones_like(es) for val, es in zip(distances, eigenenergies, strict=True)])
      62           0 :     energies_flattend = np.hstack(eigenenergies)
      63           0 :     overlaps_flattend = np.hstack(overlaps)
      64           0 :     sorter = np.argsort(overlaps_flattend)
      65             : 
      66           0 :     ms = mpl.rcParams["lines.markersize"]
      67           0 :     scat = ax.scatter(
      68             :         x_repeated[sorter],
      69             :         energies_flattend[sorter],
      70             :         c=overlaps_flattend[sorter],
      71             :         s=ms * 2.5,
      72             :         vmin=0,
      73             :         vmax=1,
      74             :         cmap=alphamagma,
      75             :     )
      76             : 
      77           0 :     if owns_ax:
      78           0 :         ax.set_xlabel(r"Distance ($\mu$m)")
      79           0 :         ax.set_ylabel(r"Energy (h$\cdot$GHz)")
      80           0 :         fig.colorbar(scat, ax=ax, label="Overlap with state of interest")
      81           0 :         fig.tight_layout()
      82             : 
      83           0 :     return fig, ax