Merge pull request #81 from bleykauf/feature/better-plotting

Add better plotting methods

Author: bleykauf

aisim/atoms.py

@@ -1,5 +1,8 @@
 """Classes and functions related to the atomic cloud."""
 
+from typing import Literal
+
+import matplotlib.pyplot as plt
 import numpy as np
 import scipy.linalg as splin
 
@@ -32,13 +35,6 @@ class AtomicEnsemble:
     phase_space_vectors : ndarray
         n x 6 dimensional array representing the phase space vectors
         (x0, y0, z0, vx, vy, vz) of the atoms in an atomic ensemble
-    position
-    velocity
-    state_kets
-    state_bras
-    density_matrices
-    density_matrix
-
     """
 
     def __init__(self, phase_space_vectors, state_kets=[1, 0], time=0):
@@ -53,7 +49,7 @@ def __init__(self, phase_space_vectors, state_kets=[1, 0], time=0):
     def __getitem__(self, key):
         """Select  certain atoms "from the ensemble.
 
-            Parameters
+        Parameters
         ----------
             key : int or slice or bool map
                 for example 2, 1:15 or a boolean map
@@ -210,6 +206,49 @@ def fidelity(self, rho_target):
         """
         return _fidelity(self.density_matrix, rho_target)
 
+    def plot(
+        self,
+        ax: plt.Axes | None = None,
+        view_from: Literal["x", "y", "z"] = "z",
+        bins: int = 50,
+        **kwargs,
+    ) -> tuple[plt.Figure, plt.Axes]:
+        """Plot the positions of the atoms in the ensemble.
+
+        ax : Axis , optional
+            If axis is provided, they will be used for the plot. if not provided, a new
+            plot will automatically be created.
+        view_from : str
+            View from which direction the plot is created. Options are "x", "y", "z".
+        bins : int
+            Number of bins for the histogram
+        **kwargs
+            Additional keyword arguments for the plot function
+
+        Returns
+        -------
+        fig, ax : tuple of plt.Figure and plt.Axes
+            The figure and axis of the plot
+        """
+        if ax is None:
+            fig, ax = plt.subplots(subplot_kw=dict(projection="polar"))
+        else:
+            fig = ax.figure
+
+        view = {
+            "x": (1, 2),
+            "y": (0, 2),
+            "z": (0, 1),
+        }
+        ax.hist2d(
+            self.position[:, view[view_from][0]],
+            self.position[:, view[view_from][1]],
+            bins=bins,
+            **kwargs,
+        )
+
+        return fig, ax
+
 
 def create_random_ensemble_from_gaussian_distribution(
     pos_params, vel_params, n_samples, seed=None, **kwargs

aisim/beam.py

@@ -2,6 +2,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib.colors import Colormap
 
 from . import convert
 from .zern import FIRST_INDEX_J, ZernikeNorm, ZernikeOrder, ZernikePolynomial
@@ -176,7 +177,13 @@ def get_value(self, pos: np.ndarray) -> np.ndarray:
             values[rho > self.r_beam / self.r_wf] = np.nan
         return values
 
-    def plot(self, ax: plt.Axes | None = None) -> tuple[plt.Figure, plt.Axes]:
+    def plot(
+        self,
+        ax: plt.Axes | None = None,
+        cmap: str | Colormap = "RdBu",
+        levels: int = 100,
+        **kwargs,
+    ) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the wavefront data.
 
@@ -185,7 +192,17 @@ def plot(self, ax: plt.Axes | None = None) -> tuple[plt.Figure, plt.Axes]:
         ax : Axis , optional
             If axis is provided, they will be used for the plot. if not provided, a new
             plot will automatically be created.
+        cmap : str or Colormap
+            Colormap for the plot
+        level : int
+            Number of levels for the contour plot
+        **kwargs
+            Additional keyword arguments for the plot function
 
+        Returns
+        -------
+        fig, ax : tuple of plt.Figure and plt.Axes
+            The figure and axis of the plot
         """
         azimuths = np.radians(np.linspace(0, 360, 180))
         zeniths = np.linspace(0, self.r_wf, 50)
@@ -204,14 +221,16 @@ def plot(self, ax: plt.Axes | None = None) -> tuple[plt.Figure, plt.Axes]:
         theta = theta.reshape(n_dim, m_dim)
         rho = rho.reshape(n_dim, m_dim)
         values = values.reshape(n_dim, m_dim)
-        contour = ax.contourf(theta, rho, values)
+        contour = ax.contourf(theta, rho, values, cmap=cmap, levels=levels, **kwargs)
         cbar = plt.colorbar(contour)
         cbar.set_label(r"Aberration / $\lambda$", rotation=90)
         plt.tight_layout()
 
         return fig, ax
 
-    def plot_coeff(self, ax: plt.Axes | None = None) -> tuple[plt.Figure, plt.Axes]:
+    def plot_coeff(
+        self, ax: plt.Axes | None = None, **kwargs
+    ) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the coefficients as a bar chart.
 
@@ -220,13 +239,18 @@ def plot_coeff(self, ax: plt.Axes | None = None) -> tuple[plt.Figure, plt.Axes]:
         ax : Axis , optional
             If axis is provided, they will be used for the plot. if not provided, a new
             plot will automatically be created.
+
+        Returns
+        -------
+        fig, ax : tuple of plt.Figure and plt.Axes
+            The figure and axis of the plot
         """
         if ax is None:
             fig, ax = plt.subplots()
         else:
             fig = ax.figure
 
-        ax.bar(list(self.coeff.keys()), list(self.coeff.values()))
+        ax.bar(list(self.coeff.keys()), list(self.coeff.values()), **kwargs)
         ax.set_xlabel("Zernike polynomial $j$")
         ax.set_ylabel(r"Zernike coefficient $Z_j$ / $\lambda$")
         ax.set_xlim(min(self.coeff.keys()) - 1, max(self.coeff.keys()) + 1)