4
\$\begingroup\$

I have completed my project which models the paths of particles in a plasma experiencing a force. The program creates phase diagrams of a particle's perpendicular velocity against its parallel velocity from 0 to 180 degrees.

I would like your thoughts and improvements on my implementation.

Here is the code:

import matplotlib.pyplot as plt
import numpy as np

alpha_0 = [5, 30, 60, 85] # initial pitch angles in degrees
degrees = np.linspace(0, 180, int(round(180/1 + 1))) # I choose to plot phase diagram from 0 to 180 deg 
v = 1

v_perp_5=v*np.sin(np.deg2rad(alpha_0[0] + degrees))
v_par_5=pow(v**2 - v_perp_5**2,0.5)

v_perp_30=v*np.sin(np.deg2rad(alpha_0[1] + degrees))
v_par_30=pow(v**2 - v_perp_30**2,0.5)

v_perp_60=v*np.sin(np.deg2rad(alpha_0[2] + degrees))
v_par_60=pow(v**2 - v_perp_60**2,0.5)

v_perp_85=v*np.sin(np.deg2rad(alpha_0[3] + degrees))
v_par_85=pow(v**2 - v_perp_85**2,0.5)

# plot lines

plt.rcParams["figure.figsize"] = [13,13]
fig, axs = plt.subplots(2, 2)
fig.suptitle('Phase Diagrams from ' + r'$0^\circ$' + ' to ' +  r'$180^\circ$', fontsize=21)

axs[0, 0].plot(v_par_5, v_perp_5)
axs[0, 0].set_title(r'$\alpha_o = 5^\circ$', fontsize=15)
axs[0, 0].set_xlim([-1,1])
axs[0, 0].set_ylim([-1,1])
axs[0, 0].set_xlabel(r'$v_{\parallel}$', fontsize=15)
axs[0, 0].set_ylabel(r'$v_{\perp}$', fontsize=15)
axs[0, 0].annotate(r'$\alpha_o$', (v_par_5[0], v_perp_5[0]), fontsize=14)
cir_init_5 = plt.Circle((v_par_5[0], v_perp_5[0]), 0.05, color='r',fill=False)
axs[0, 0].add_patch(cir_init_5)
axs[0, 0].arrow(x=0, y=0.85, dx=0, dy=0.08, width=.015)
axs[0, 0].text(-0.04, 0.79, r'$\alpha_m$', fontsize=14)

axs[0, 1].plot(v_par_30, v_perp_30, 'tab:orange')
axs[0, 1].set_title(r'$\alpha_o = 30^\circ$', fontsize=15)
axs[0, 1].set_xlim([-1,1])
axs[0, 1].set_ylim([-1,1])
axs[0, 1].set_xlabel(r'$v_{\parallel}$', fontsize=15)
axs[0, 1].set_ylabel(r'$v_{\perp}$', fontsize=15)
axs[0, 1].annotate(r'$\alpha_o$', (v_par_30[0], v_perp_30[0]), fontsize=14)
cir_init_30 = plt.Circle((v_par_30[0], v_perp_30[0]), 0.05, color='r',fill=False)
axs[0, 1].add_patch(cir_init_30)
axs[0, 1].arrow(x=0, y=0.85, dx=0, dy=0.08, width=.015)
axs[0, 1].text(-0.04, 0.79, r'$\alpha_m$', fontsize=14)

axs[1, 0].plot(v_par_60, v_perp_60, 'tab:green')
axs[1, 0].set_title(r'$\alpha_o = 60^\circ$', fontsize=15)
axs[1, 0].set_xlim([-1,1])
axs[1, 0].set_ylim([-1,1])
axs[1, 0].set_xlabel(r'$v_{\parallel}$', fontsize=15)
axs[1, 0].set_ylabel(r'$v_{\perp}$', fontsize=15)
axs[1, 0].annotate(r'$\alpha_o$', (v_par_60[0], v_perp_60[0]), fontsize=14)
cir_init_60 = plt.Circle((v_par_60[0], v_perp_60[0]), 0.05, color='r',fill=False)
axs[1, 0].add_patch(cir_init_60)
axs[1, 0].arrow(x=0, y=0.85, dx=0, dy=0.08, width=.015)
axs[1, 0].text(-0.04, 0.79, r'$\alpha_m$', fontsize=14)

axs[1, 1].plot(v_par_85, v_perp_85, 'tab:red')
axs[1, 1].set_title(r'$\alpha_o = 85^\circ$', fontsize=15)
axs[1, 1].set_xlim([-1,1])
axs[1, 1].set_ylim([-1,1])
axs[1, 1].set_xlabel(r'$v_{\parallel}$', fontsize=15)
axs[1, 1].set_ylabel(r'$v_{\perp}$', fontsize=15)
axs[1, 1].annotate(r'$\alpha_o$', (v_par_85[0], v_perp_85[0]), fontsize=14)
cir_init_85 = plt.Circle((v_par_85[0], v_perp_85[0]), 0.05, color='r',fill=False)
axs[1, 1].add_patch(cir_init_85)
axs[1, 1].arrow(x=0, y=0.85, dx=0, dy=0.08, width=.015)
axs[1, 1].text(-0.04, 0.79, r'$\alpha_m$', fontsize=14)

fig.tight_layout()
\$\endgroup\$

1 Answer 1

5
\$\begingroup\$
  • Functions and type hints are your friend - use them
  • Vectorise your separate alpha values and calculations to a single calculation for each of perpendicular and parallel arrays
  • np.linspace(0, 180, int(round(180/1 + 1))) is a slightly bizarre way of writing np.arange(181)
  • Do not pow(x, 0.5); use np.sqrt
  • "13" is a quantity in inches. If you're rendering for print this will be larger than the standard 8.5x11. If you're rendering for the screen, most laptop screens are smaller than this as well. Between this and your custom font, the rendering was very poor and produced overlaps. I decreased this to 9" and used default fonts and the problems went away.
  • There isn't a benefit to your string concatenation for suptitle; you can have a single raw string.
  • Factor out your copy-and-paste subplot code to one function.
  • There is no semantic difference between your four subplots that hasn't already been explained by set_title, so there isn't a benefit to separate colours. Separate colours - with a legend - would be useful if you did not include the alpha value in the titles.
  • Don't ^\circ; just use \degree.

Coordinate space

Your conversion from polar to rectilinear coordinates is probably incorrect. Your use of the Euclidean norm drops a negative sign, so you're only ever seeing half of the coordinate space. Instead of sqrt for a reverse Pythagorean, just use cos.

In the (maybe unlikely?) case that this was deliberate, you can just apply an np.abs to your parallel coordinate axis. However, the less surprising method to get a (I, IV) quadrant effect would be to construct the angles themselves prior to projection using np.min to make a piecewise angle function, and to then use a plain sin,cos projection with no abs.

Suggested

from numbers import Real
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np


def calculate_v(v: Real) -> Tuple[
    np.ndarray,  # alpha: 4
    np.ndarray,  # perpendicular: n_alpha * n_degrees
    np.ndarray,  # parallel:      n_alpha * n_degrees
]:
    alpha_0 = np.array((5, 30, 60, 85), ndmin=2).T  # initial pitch angles in degrees
    degrees = np.arange(181)[np.newaxis, :]  # Plot phase diagram from 0 to 180 deg
    radians = np.deg2rad(alpha_0 + degrees)
    v_perp = v*np.sin(radians)
    v_par = v*np.cos(radians)
    return alpha_0[:, 0], v_perp, v_par


def plot_lines(alpha_0: np.ndarray, v_perp: np.ndarray, v_par: np.ndarray) -> plt.Figure:
    plt.rcParams['figure.figsize'] = (9, 9)
    fig, axes_grid = plt.subplots(2, 2)
    fig.suptitle(
        r'Phase Diagrams from $0^\circ$ to $180^\circ$'
    )

    axes = [
        ax
        for row in axes_grid
        for ax in row
    ]
    for args in zip(axes, v_perp, v_par, alpha_0):
        subplot(*args)

    fig.tight_layout()
    return fig


def subplot(ax: plt.Axes, v_perp: np.ndarray, v_par: np.ndarray, alpha: int) -> None:
    ax.plot(v_par, v_perp)
    ax.set_title(rf'$\alpha_o = {alpha}^\circ$')
    ax.set_xlim(left=-1, right=1)
    ax.set_ylim(bottom=-1, top=1)
    ax.set_xlabel(r'$v_{\parallel}$')
    ax.set_ylabel(r'$v_{\perp}$')
    ax.annotate(r'$\alpha_o$', (v_par[0], v_perp[0]))
    cir_init = plt.Circle((v_par[0], v_perp[0]), 0.05, color='r', fill=False)
    ax.add_patch(cir_init)
    ax.arrow(x=0, y=0.85, dx=0, dy=0.08, width=.015)
    ax.text(-0.04, 0.79, r'$\alpha_m$')


def main() -> None:
    alpha, v_perp, v_par = calculate_v(v=1)
    plot_lines(alpha, v_perp, v_par)
    plt.show()


if __name__ == '__main__':
    main()

separated plot

Superimposed plots

A more informative and condensed visualisation superimposes all of your subplots and allows for easier comparison, including a legend:

from numbers import Real
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np


def calculate_v(v: Real) -> Tuple[
    np.ndarray,  # alpha: 4
    np.ndarray,  # perpendicular: n_alpha * n_degrees
    np.ndarray,  # parallel:      n_alpha * n_degrees
]:
    alpha_0 = np.array((5, 30, 60, 85), ndmin=2).T  # initial pitch angles in degrees
    degrees = np.arange(181)[np.newaxis, :]  # Plot phase diagram from 0 to 180 deg
    radians = np.deg2rad(alpha_0 + degrees)
    v_perp = v*np.sin(radians)
    v_par = v*np.cos(radians)
    return alpha_0[:, 0], v_perp, v_par


def plot_lines(alpha_0: np.ndarray, v_perp: np.ndarray, v_par: np.ndarray) -> plt.Figure:
    plt.rcParams['figure.figsize'] = (9, 9)
    fig, ax = plt.subplots()
    fig.suptitle(
        r'Phase Diagrams from $0^\circ$ to $180^\circ$'
    )

    ax.set_xlim(left=-1, right=1)
    ax.set_ylim(bottom=-1, top=1)
    ax.set_xlabel(r'$v_{\parallel}$')
    ax.set_ylabel(r'$v_{\perp}$')
    ax.arrow(x=0, y=0.85, dx=0, dy=0.08, width=.015)
    ax.text(-0.04, 0.79, r'$\alpha_m$')

    for args in reversed(tuple(zip(v_perp, v_par, alpha_0))):
        subplot(ax, *args)

    ax.legend(title=r'$\alpha_o$')
    fig.tight_layout()
    return fig


def subplot(ax: plt.Axes, v_perp: np.ndarray, v_par: np.ndarray, alpha: int) -> None:
    line, = ax.plot(
        v_par, v_perp,
        label=alpha,
    )
    ax.annotate(
        text=r'$\alpha_o$', xy=(v_par[0], v_perp[0]),
    )
    cir_init = plt.Circle(
        xy=(v_par[0], v_perp[0]), radius=0.05,
        fill=False, color=line.get_color(),
    )
    ax.add_patch(cir_init)


def main() -> None:
    alpha, v_perp, v_par = calculate_v(v=1)
    plot_lines(alpha, v_perp, v_par)
    plt.show()


if __name__ == '__main__':
    main()

superimposed plots

First-class Polar

It doesn't make a tonne of sense to use a Cartesian plot for your data, or to represent them in rectilinear coordinates. Everything is made more simple if you represent your data in polar coordinates - every single radius is just equal to 1, your angles are linear, and there's no need for sin, cos or sqrt.

from numbers import Real
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np


def calculate_v(v: Real) -> Tuple[
    np.ndarray,  # alpha: 4
    np.ndarray,  # radians: n_alpha * n_degrees
    np.ndarray,  # radii:   n_alpha * n_degrees
]:
    alpha_0 = np.array((5, 30, 60, 85), ndmin=2).T  # initial pitch angles in degrees
    degrees = np.arange(181)[np.newaxis, :]  # Plot phase diagram from 0 to 180 deg
    radians = np.deg2rad(alpha_0 + degrees)
    radii = np.full_like(radians, fill_value=v)
    return alpha_0[:, 0], radians, radii


def plot_lines(alpha_0: np.ndarray, radians: np.ndarray, radii: np.ndarray) -> plt.Figure:
    plt.rcParams['figure.figsize'] = (9, 9)
    fig: plt.Figure
    ax: plt.PolarAxes
    fig, ax = plt.subplots(
        subplot_kw={'projection': 'polar'}
    )
    fig.suptitle(
        r'Phase Diagrams from $0^\circ$ to $180^\circ$'
    )

    ax.set_rlabel_position(-22.5)
    ax.arrow(x=np.pi/2, y=0.85, dx=0, dy=0.08, width=.015)
    ax.text(x=np.pi/2, y=0.79, s=r'$\alpha_m$')

    for args in reversed(tuple(zip(radians, radii, alpha_0))):
        subplot(ax, *args)

    ax.grid(True)
    ax.legend(title=r'$\alpha_o$')
    fig.tight_layout()
    return fig


def subplot(ax: plt.PolarAxes, radians: np.ndarray, radii: np.ndarray, alpha: int) -> None:
    line, = ax.plot(radians, radii, label=alpha)
    ax.annotate(
        text=r'$\alpha_o$', xy=(radians[0], radii[0]),
    )
    cir_init = plt.Circle(
        xy=(radians[0], radii[0]), radius=0.05,
        fill=False, color=line.get_color(),
    )
    ax.add_patch(cir_init)


def main() -> None:
    alpha, radians, radii = calculate_v(v=1)
    plot_lines(alpha, radians, radii)
    plt.show()


if __name__ == '__main__':
    main()

polar plot

Rays instead of markers

Your circles are basically non-standard markers. You could use standard markers and just have a markevery setting that only shows the first data point in each series. However, that doesn't help with highlighting the end of the series. One way to highlight both the beginning and end of each polar series is to draw a ray to the beginning and end point:

from numbers import Real
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np


def calculate_v(v: Real) -> Tuple[
    np.ndarray,  # alpha: 4
    np.ndarray,  # radians: n_alpha * n_degrees
    np.ndarray,  # radii:   n_alpha * n_degrees
]:
    alpha_0 = np.array((5, 30, 60, 85), ndmin=2).T  # initial pitch angles in degrees
    degrees = np.arange(181)[np.newaxis, :]  # Plot phase diagram from 0 to 180 deg
    radians = np.deg2rad(alpha_0 + degrees)
    radii = np.full_like(radians, fill_value=v)
    return alpha_0[:, 0], radians, radii


def plot_lines(alpha_0: np.ndarray, radians: np.ndarray, radii: np.ndarray) -> plt.Figure:
    plt.rcParams['figure.figsize'] = (9, 9)
    fig: plt.Figure
    ax: plt.PolarAxes
    fig, ax = plt.subplots(
        subplot_kw={'projection': 'polar'}
    )
    ax.set_title(
        r'Phase Diagrams from $0\degree$ to $180\degree$, with bound rays'
    )

    ax.set_rlabel_position(-22.5)
    ax.arrow(x=np.pi/2, y=0.85, dx=0, dy=0.08, width=.015)
    ax.text(x=np.pi/2, y=0.79, s=r'$\alpha_m$')

    for args in reversed(tuple(zip(radians, radii, alpha_0))):
        subplot(ax, *args)

    ax.grid(True)
    ax.legend(title=r'$\alpha_o$')
    fig.tight_layout()
    return fig


def subplot(ax: plt.PolarAxes, radians: np.ndarray, radii: np.ndarray, alpha: int) -> None:
    line, = ax.plot(
        radians, radii,
        label=rf'${alpha}\degree$',
    )

    ax.plot(
        [radians[0], 0, radians[-1]],
        [radii[0], 0, radii[-1]],
        color=line.get_color(),
        linestyle='--',
    )


def main() -> None:
    alpha, radians, radii = calculate_v(v=1)
    plot_lines(alpha, radians, radii)
    plt.show()


if __name__ == '__main__':
    main()

rays included

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.