Source code for pypolar.visualization

"""
A set of basic routines for visualizing polarization.

Functions for drawing the polarization ellipse (sectional pattern)::

   * draw_jones_ellipse(J, simple=False, **kwargs)
   * draw_stokes_ellipse(S, **kwargs)

Functions for drawing 2D and 3D representations::

    * draw_jones_field(J, offset=0, **kwargs)
    * draw_stokes_field(S, offset=0, **kwargs)

Functions for drawing animated 2D and 3D representations::

   * draw_jones_animated(J, nframes=64, **kwargs)
   * draw_stokes_animated(S, **kwargs)

Functions for drawing Poincaré representations::
   * draw_empty_sphere(ax=None, **kwargs)
   * draw_jones_poincare(J, ax=None, label=None, normalize="s0", text_kwargs=None, **kwargs)
   * draw_stokes_poincare(S, ax=None, label=None, normalize="s0", text_kwargs=None, **kwargs)
   * join_jones_poincare(J1, J2, ax=None, normalize="s0", **kwargs)
   * join_stokes_poincare(S1, S2, ax=None, normalize="s0", **kwargs)

Poincaré coordinates use reduced Stokes values (S1/S0, S2/S0, S3/S0),
so partially polarized states lie inside the unit sphere.

Jones-vector plots follow the package-wide sign convention set by
`pypolar.jones.use_alternate_convention(...)`.

Set `normalize="unit"` to project states onto the unit sphere using
`(S1,S2,S3) / sqrt(S1^2+S2^2+S3^2)`.

Example: Poincaré sphere plot of a Jones vector::

    J = pypolar.jones.field_linear(np.pi / 6)
    pypolar.visualization.draw_jones_poincare(J)

Example: Poincaré sphere plot of two Stokes vectors::

    S1 = pypolar.mueller.stokes_left_circular()
    S2 = pypolar.mueller.stokes_linear(np.radians(15))

    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')
    pypolar.visualization.draw_empty_sphere(ax)
    pypolar.visualization.draw_stokes_poincare(S1, ax, label='  S1')
    pypolar.visualization.draw_stokes_poincare(S2, ax, label='  S2')
    pypolar.visualization.join_stokes_poincare(S1, S2, ax, lw=2, ls=':', color='orange')
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib import animation

import pypolar.fresnel
import pypolar.mueller
import pypolar.jones
import pypolar.poincare as _poincare

draw_empty_sphere = _poincare.draw_empty_sphere
draw_jones_poincare = _poincare.draw_jones_poincare
draw_stokes_poincare = _poincare.draw_stokes_poincare
join_jones_poincare = _poincare.join_jones_poincare
join_stokes_poincare = _poincare.join_stokes_poincare
great_circle_points = _poincare.great_circle_points
spherical_angles = _poincare.spherical_angles

__all__ = (
    "draw_jones_field",
    "draw_jones_animated",
    "draw_jones_ellipse",
    "draw_stokes_ellipse",
    "draw_stokes_field",
    "draw_stokes_animated",
    "draw_empty_sphere",
    "draw_jones_poincare",
    "draw_stokes_poincare",
    "join_jones_poincare",
    "join_stokes_poincare",
)


def _jones_for_visualization(J):
    """Return a Jones vector in the plotting convention used by this module."""
    if pypolar.jones.alternate_sign_convention:
        return np.conjugate(J)
    return J


def _draw_optical_axis_3d(J, ax, last=4 * np.pi, **kwargs):
    """
    Draw the optical axis in a 3D plot.

    Args:
        J:    Jones vector
        ax:   matplotlib axis to use
        last: length of optical axis
        **kwargs: style arguments passed to line artists.
    """
    h_amp, v_amp = abs(J)
    the_max = max(h_amp, v_amp) * 1.1

    ax.plot([0, last * 1.15], [0, 0], [0, 0], "k", **kwargs)
    ax.plot([0, 0], [-the_max, the_max], [0, 0], "g", **kwargs)
    ax.plot([0, 0], [0, 0], [-the_max, the_max], "b", **kwargs)
    ax.text(0, 0, the_max * 1.1, "y", ha="center", va="bottom")
    ax.text(0, the_max * 1.1, 0, "x", va="center")
    ax.text(last * 1.2, 0, 0, "z", va="center")


def _draw_h_field_3d(J, ax, offset, last=4 * np.pi, **kwargs):
    """
    Draw the horizontal electric field in a 3D plot.

    Args:
        J:      Jones vector
        ax:     matplotlib axis to use
        offset: phase offset in radians added to the cosine argument.
                This shifts the x-component waveform along the optical axis.
        last:   length of optical axis
        **kwargs: style arguments passed to line artists.
    """
    t = np.linspace(0, last, 100)
    x = t
    y = np.abs(J[0]) * np.cos(t + offset - np.angle(J[0]))
    z = 0
    ax.plot(x, y, z, ":g", **kwargs)


def _draw_v_field_3d(J, ax, offset, last=4 * np.pi, **kwargs):
    """
    Draw the vertical electric field in a 3D plot.

    Args:
        J:      Jones vector
        ax:     matplotlib axis to use
        offset: phase offset in radians added to the cosine argument.
                This shifts the y-component waveform along the optical axis.
        last:   length of optical axis
        **kwargs: style arguments passed to line artists.
    """
    t = np.linspace(0, last, 100)
    x = t
    y = 0 * t
    z = np.abs(J[1]) * np.cos(t + offset - np.angle(J[1]))
    ax.plot(x, y, z, ":b", **kwargs)


def _draw_total_field_3d(J, ax, offset, last=4 * np.pi, **kwargs):
    """
    Draw the total electric field in a 3D plot.

    Args:
        J:      Jones vector
        ax:     matplotlib axis to use
        offset: phase offset in radians added to both component cosines.
                This picks the instantaneous phase shown in the 3D trace.
        last:   length of optical axis
        **kwargs: style arguments passed to line artists.
    """
    t = np.linspace(0, last, 100)
    x = t
    y = np.abs(J[0]) * np.cos(t + offset - np.angle(J[0]))
    z = np.abs(J[1]) * np.cos(t + offset - np.angle(J[1]))
    ax.plot(x, y, z, "r", **kwargs)


def _draw_projected_vector_3d(J, ax, offset, **kwargs):
    """
    Draw the projection vector of the polarization field in 3D.

    Args:
        J:      Jones vector
        ax:     matplotlib axis to use
        offset: phase offset in radians used to compute the instantaneous
                tip location of the field vector.
        **kwargs: style arguments passed to line artists.
    """
    y = np.abs(J[0]) * np.cos(offset - np.angle(J[0]))
    z = np.abs(J[1]) * np.cos(offset - np.angle(J[1]))

    x1, y1, z1 = 0, y, 0
    x2, y2, z2 = 0, y, z
    ax.plot([x1, x2], [y1, y2], [z1, z2], "g--", **kwargs)

    x1, y1, z1 = 0, 0, z
    ax.plot([x1, x2], [y1, y2], [z1, z2], "b--", **kwargs)

    x1, y1, z1 = 0, 0, 0
    ax.plot([x1, x2], [y1, y2], [z1, z2], "r", **kwargs)
    ax.plot([0], [y], [z], "ro", **kwargs)


def _draw_3D_field(J, ax, offset, **kwargs):
    """
    Draw a representation of the polarization fields in 3D.

    Args:
        J:      Jones vector
        ax:     matplotlib axis to use
        offset: phase offset in radians that controls the instantaneous
                snapshot of all 3D field components.
        **kwargs: style arguments passed to line artists.
    """
    _draw_optical_axis_3d(J, ax, **kwargs)
    _draw_h_field_3d(J, ax, offset, **kwargs)
    _draw_v_field_3d(J, ax, offset, **kwargs)
    _draw_total_field_3d(J, ax, offset, **kwargs)
    _draw_projected_vector_3d(J, ax, offset, **kwargs)

    ax.grid(False)
    ax.axis("off")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    ax.set_box_aspect(None, zoom=1.3)


def _draw_2D_field(J, ax, offset, **kwargs):
    """
    Draw a simple 2D representation of the projected field.

    Also called a sectional pattern.

    Args:
        J:      Jones vector
        ax:     matplotlib axis to use
        offset: phase offset in radians that chooses which point on the
                polarization ellipse is highlighted.
        **kwargs: style arguments passed to line artists.
    """
    h_amp, v_amp = np.abs(J)
    h_phi, v_phi = np.angle(J)
    the_max = max(h_amp, v_amp) * 1.2

    ax.plot([-the_max, the_max], [0, 0], "g", **kwargs)
    ax.plot([0, 0], [-the_max, the_max], "b", **kwargs)

    t = np.linspace(0, 2 * np.pi, 100)
    x = h_amp * np.cos(t + offset - h_phi)
    y = v_amp * np.cos(t + offset - v_phi)
    ax.plot(x, y, "k", **kwargs)

    x = h_amp * np.cos(offset - h_phi)
    y = v_amp * np.cos(offset - v_phi)
    ax.plot(x, y, "ro", **kwargs)
    ax.plot([x, x], [0, y], "g--", **kwargs)
    ax.plot([0, x], [y, y], "b--", **kwargs)
    ax.plot([0, x], [0, y], "r", **kwargs)

    ax.set_xlim(-the_max, the_max)
    ax.set_ylim(-the_max, the_max)
    ax.set_aspect("equal")
    ax.grid(False)
    ax.axis("off")
    ax.text(0, the_max, "y", ha="center", va="bottom")
    ax.text(the_max, 0, "x", va="center")


def _animation_update(offset, J, ax1, ax2, plot_kwargs):
    """
    Draw the next animation frame.

    Args:
        offset: frame phase in radians. Each frame updates this value to move
                the instantaneous field point around the ellipse.
        J:      Jones vector
        ax1:    matplotlib axis for 3D plot
        ax2:    matplotlib axis for 2D plot
        plot_kwargs: style arguments passed to line artists.
    """
    ax1.clear()
    ax2.clear()
    _draw_3D_field(J, ax1, offset, **plot_kwargs)
    _draw_2D_field(J, ax2, offset, **plot_kwargs)
    return ax1, ax2


def draw_ellipse_axes(J, ax, **kwargs):
    """
    Draw the sectional pattern with ellipse labels.

    Args:
        J:  Jones vector
        ax: plot axis
        **kwargs: style arguments passed to line artists.
    """
    Ex0, Ey0 = np.abs(J)
    phix, phiy = np.angle(J)

    alpha = pypolar.jones.ellipse_azimuth(J)
    a, b = pypolar.jones.ellipse_axes(J)

    t = np.linspace(0, 2 * np.pi, 100)
    xx = Ex0 * np.cos(t + phix)
    yy = Ey0 * np.cos(t + phiy)

    the_max = max(Ex0, Ey0) * 1.2

    ax.set_aspect("equal")
    ax.plot(xx, yy, "b", **kwargs)

    # semi-major diameter
    dx = a * np.cos(alpha)
    dy = a * np.sin(alpha)
    ax.plot([0, dx], [0, dy], "r", **kwargs)
    ax.text(dx / 2, dy / 2, "  a", color="red")
    ax.text(dx / 5, dy / 10, r"$\alpha$", va="center", ha="center")
    s = r"a=%.2f, b=%.2f, $\alpha$=%.2f°" % (a, b, np.degrees(alpha))
    ax.text(0, -1.15 * the_max, s, ha="center")

    # semi-minor diameter
    alpha += np.pi / 2
    dx = b * np.cos(alpha)
    dy = b * np.sin(alpha)
    ax.plot([0, dx], [0, dy], "g", **kwargs)
    ax.text(dx / 2, dy / 2, "  b", color="green")
    s = r"b / a=%.2f, " % (b / a)
    s += r"$\tan^{-1}(b / a)$=%.2f°" % np.degrees(pypolar.jones.ellipticity_angle(J))
    ax.text(0, -1.30 * the_max, s, ha="center")

    # draw x and y axes
    ax.plot([0, 0], [-the_max, the_max], "k", **kwargs)
    ax.plot([-the_max, the_max], [0, 0], "k", **kwargs)
    ax.set_xlim(-the_max, the_max)
    ax.set_ylim(-the_max, the_max)
    ax.set_xticks([])
    ax.set_yticks([])


def draw_ellipse_Ex_Ey(J, ax, **kwargs):
    """
    Draw the sectional pattern with field labels.

    Args:
        J:  Jones vector
        ax: plot axis
        **kwargs: style arguments passed to line artists.
    """
    Ex0, Ey0 = np.abs(J)
    phix, phiy = np.angle(J)

    t = np.linspace(0, 2 * np.pi, 100)
    xx = Ex0 * np.cos(t + phix)
    yy = Ey0 * np.cos(t + phiy)

    the_max = max(Ex0, Ey0) * 1.2
    ax.set_aspect("equal")
    ax.plot(xx, yy, "b", **kwargs)
    ax.plot([-Ex0, -Ex0, Ex0, Ex0, -Ex0], [-Ey0, Ey0, Ey0, -Ey0, -Ey0], ":g", **kwargs)
    ax.plot([-Ex0, Ex0], [-Ey0, Ey0], ":r", **kwargs)
    ax.plot([0, 0], [-the_max, the_max], "k", **kwargs)
    ax.plot([-the_max, the_max], [0, 0], "k", **kwargs)
    ax.text(Ex0, 0, r" $E_{x0}$", va="bottom", ha="left")
    ax.text(-Ex0, 0, r"$-E_{x0} $", va="bottom", ha="right")
    ax.text(0, Ey0, r"$E_{y0}$", va="bottom", ha="left")
    ax.text(0, -Ey0, r"$-E_{y0}$", va="top", ha="left")
    ax.text(0, Ey0 / 5, r" $\psi$", va="bottom", ha="left")
    ax.set_xlim(-the_max, the_max)
    ax.set_ylim(-the_max, the_max)
    ax.set_xticks([])
    ax.set_yticks([])
    psi = np.degrees(np.arctan2(Ex0, Ey0))
    s = r"$E_{0x}$=%.2f, $E_{0y}$=%.2f, $\psi$=%.2f°" % (Ex0, Ey0, psi)
    ax.text(0, -1.15 * the_max, s, ha="center")
    s = r"$\phi_x$=%.2f°, " % np.degrees(phix)
    s += r"$\phi_y$=%.2f°, " % np.degrees(phiy)
    s += r"$\phi_y-\phi_x$=%.2f°" % np.degrees(phiy - phix)
    ax.text(0, -1.30 * the_max, s, ha="center")


[docs] def draw_jones_ellipse(J, simple=False, **kwargs): """ Draw a 2D sectional pattern for a Jones vector. Args: J: Jones vector simple: if True then just draw a simple ellipse plot **kwargs: style arguments passed to line artists. Returns: tuple: `(fig, ax_or_axes, artists)` where `ax_or_axes` is one axis for `simple=True` and `(ax1, ax2)` for `simple=False`. """ JJ = _jones_for_visualization(J) if simple: Ex0, Ey0 = np.abs(JJ) phix, phiy = np.angle(JJ) the_max = max(Ex0, Ey0) * 1.2 t = np.linspace(0, 2 * np.pi, 100) xx = Ex0 * np.cos(t + phix) yy = Ey0 * np.cos(t + phiy) ax = plt.gca() fig = ax.figure n_lines = len(ax.lines) n_texts = len(ax.texts) ax.set_xlim(-the_max, the_max) ax.set_ylim(-the_max, the_max) ax.set_aspect("equal") axis_kwargs = dict(kwargs) if "color" not in axis_kwargs and "c" not in axis_kwargs: axis_kwargs["color"] = "black" ax.axhline(0, **axis_kwargs) ax.axvline(0, **axis_kwargs) ax.plot(xx, yy, "b", **kwargs) ax.plot([-Ex0, Ex0], [-Ey0, Ey0], ":r", **kwargs) ax.axis("off") ax.text(0, Ey0 / 5, r" $\psi$", va="bottom", ha="left") artists = {"lines": list(ax.lines[n_lines:]), "texts": list(ax.texts[n_texts:])} return fig, ax, artists fig = plt.figure(figsize=(8, 4)) gs = gridspec.GridSpec(1, 2, width_ratios=[1, 1]) ax1 = plt.subplot(gs[0]) draw_ellipse_axes(JJ, ax1, **kwargs) ax2 = plt.subplot(gs[1]) draw_ellipse_Ex_Ey(JJ, ax2, **kwargs) artists = { "ax1_lines": list(ax1.lines), "ax1_texts": list(ax1.texts), "ax2_lines": list(ax2.lines), "ax2_texts": list(ax2.texts), } return fig, (ax1, ax2), artists
[docs] def draw_stokes_ellipse(S, **kwargs): """ Draw polarization ellipse panels from a Stokes vector. Args: S: Stokes vector **kwargs: style arguments passed to `draw_jones_ellipse`. Returns: tuple: `(fig, ax_or_axes, artists)` as returned by :func:`draw_jones_ellipse`. """ J = pypolar.mueller.stokes_to_jones(S) return draw_jones_ellipse(J, **kwargs)
[docs] def draw_jones_field(J, offset=0, **kwargs): """ Draw 3D and 2D representations of a Jones vector. Args: J: Jones vector offset: phase offset in radians for the plotted snapshot. `offset=0` uses the default phase origin; changing it rotates the highlighted instantaneous field point around the ellipse. **kwargs: style arguments passed to line artists. Returns: tuple: `(fig, (ax3d, ax2d), artists)` where `artists` includes line and text handles for each axis. Example: Plot a static field representation:: import matplotlib.pyplot as plt import pypolar.jones as jones import pypolar.visualization as vis J = jones.field_left_circular() vis.draw_jones_field(J) plt.show() """ JJ = _jones_for_visualization(J) fig = plt.figure(figsize=(8, 4)) gs = gridspec.GridSpec(1, 2, width_ratios=[5, 3]) ax1 = fig.add_subplot(gs[0], projection="3d") ax2 = fig.add_subplot(gs[1]) _draw_3D_field(JJ, ax1, offset, **kwargs) _draw_2D_field(JJ, ax2, offset, **kwargs) artists = { "ax3d_lines": list(ax1.lines), "ax3d_collections": list(ax1.collections), "ax3d_texts": list(ax1.texts), "ax2d_lines": list(ax2.lines), "ax2d_texts": list(ax2.texts), } return fig, (ax1, ax2), artists
[docs] def draw_stokes_field(S, offset=0, **kwargs): """ Draw 3D and 2D field representations of a Stokes vector. Args: S: Stokes vector offset: phase offset in radians for the plotted snapshot after converting `S` to a Jones vector. Meaning matches :func:`draw_jones_field`. **kwargs: style arguments passed to `draw_jones_field`. Returns: tuple: `(fig, (ax3d, ax2d), artists)` as returned by :func:`draw_jones_field`. Example: Plot a static field representation from a Stokes vector:: import matplotlib.pyplot as plt import pypolar.mueller as mueller import pypolar.visualization as vis S = mueller.stokes_linear(0) vis.draw_stokes_field(S) plt.show() For animated output in notebooks, use :func:`draw_stokes_animated`. """ J = pypolar.mueller.stokes_to_jones(S) return draw_jones_field(J, offset, **kwargs)
[docs] def draw_jones_animated(J, nframes=64, **kwargs): """ Animate 3D and 2D representations of the polarization field. Args: J: Jones vector nframes: number of frames to create **kwargs: style arguments passed to line artists in each frame. Returns: matplotlib.animation.FuncAnimation: animation handle. The associated figure and axes are available via `ani._fig` and `ani._args[1:]`. Example: Display an animation in Jupyter notebooks:: import matplotlib.pyplot as plt import pypolar.jones as jones import pypolar.visualization as vis plt.rcParams["animation.html"] = "jshtml" J = jones.field_linear(0) ani = vis.draw_jones_animated(J, nframes=32) ani In scripts, save the animation to a file:: ani = vis.draw_jones_animated(J, nframes=32) ani.save("jones_field.mp4") """ JJ = _jones_for_visualization(J) fig = plt.figure(figsize=(8, 4)) gs = gridspec.GridSpec(1, 2, width_ratios=[5, 3]) ax1 = fig.add_subplot(gs[0], projection="3d") ax2 = fig.add_subplot(gs[1]) ani = animation.FuncAnimation( fig, _animation_update, frames=np.linspace(0, -2 * np.pi, nframes), fargs=(JJ, ax1, ax2, kwargs) ) ani.axes = (ax1, ax2) plt.close() return ani
[docs] def draw_stokes_animated(S, **kwargs): """ Draw animated 3D and 2D field representations from a Stokes vector. Args: S: Stokes vector **kwargs: style arguments passed to `draw_jones_animated`. Returns: matplotlib.animation.FuncAnimation: animation handle as returned by :func:`draw_jones_animated`. Example: Display an animation in Jupyter notebooks:: import matplotlib.pyplot as plt import pypolar.mueller as mueller import pypolar.visualization as vis plt.rcParams["animation.html"] = "jshtml" S = mueller.stokes_right_circular() ani = vis.draw_stokes_animated(S, nframes=32) ani """ J = pypolar.mueller.stokes_to_jones(S) return draw_jones_animated(J, **kwargs)