 codex/elevator-mode-vertical-columns.py

"""Illustrate a height-independent velocity field; save PNG in cwd."""
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def main():
    fig = plt.figure(figsize=(7.6, 5.4), dpi=100, facecolor="white")
    fig.suptitle("Elevator mode: vertical motion, no vertical variation",
                 y=0.97, fontsize=15, weight="bold", color="#172330")
    ax = fig.add_axes([0.12, 0.42, 0.82, 0.43])
    positions = np.linspace(0, 1, 13)
    for x in positions:
        speed = np.cos(2 * np.pi * x)
        if abs(speed) < 1e-12:
            ax.plot([x] * 4, [0.4, 1.1, 1.8, 2.5], ".", color="#9ba7ad")
            continue
        color = "#246a91" if speed > 0 else "#b65826"
        for z in [0.4, 1.1, 1.8, 2.5]:
            ax.annotate("", xy=(x, z + 0.22 * speed),
                        xytext=(x, z - 0.22 * speed),
                        arrowprops={"arrowstyle": "-|>", "color": color,
                                    "linewidth": 1.5})
    ax.set_xlim(-0.04, 1.04)
    ax.set_ylim(0, 3)
    ax.set_ylabel("height z (arbitrary units)", fontsize=11)
    ax.set_xticks([])
    ax.set_yticks([0, 1, 2, 3])
    ax.spines[["top", "right", "bottom"]].set_visible(False)
    ax.text(0.5, 1.08, r"$\mathbf{u}=(0,0,W\cos(kx))$; every row is identical",
            transform=ax.transAxes, ha="center", fontsize=12)
    below = fig.add_axes([0.12, 0.15, 0.82, 0.2])
    x = np.linspace(0, 1, 500)
    w = np.cos(2 * np.pi * x)
    below.plot(x, w, color="#53636d", linewidth=1.5)
    below.fill_between(x, 0, w, where=w >= 0, color="#246a91", alpha=0.2)
    below.fill_between(x, 0, w, where=w <= 0, color="#b65826", alpha=0.2)
    below.axhline(0, color="#8c979e", linewidth=0.7)
    below.set_xlim(-0.04, 1.04)
    below.set_ylim(-1.2, 1.2)
    below.set_xticks([0, 0.25, 0.5, 0.75, 1])
    below.set_yticks([-1, 0, 1])
    below.set_ylabel(r"$w/W$", fontsize=12)
    below.set_xlabel(r"horizontal coordinate $x/(2\pi/k)$", fontsize=11)
    below.spines[["top", "right"]].set_visible(False)
    fig.text(0.5, 0.025,
             "One wavelength of an ideal local mode; the plotted height limits are not walls.",
             ha="center", fontsize=10, color="#45535d")
    fig.savefig(Path.cwd() / (Path(__file__).stem + ".png"), dpi=100,
                facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()