Colormaps and Norms

Colorbars and scaling of 2D data make all the difference. Use the right combination of colorbar and matplotlib norms to communicate your data!

split data into 2 pieces along <w1>:
  0 : -inf to 1.60 eV (1, 6)
  1 : 1.60 to inf eV (1, 46)
/home/docs/checkouts/readthedocs.org/user_builds/wrighttools/envs/stable/lib/python3.12/site-packages/numpy/_core/numeric.py:2376: RuntimeWarning: One of rtol or atol is not valid, atol: inf, rtol: 0
  res = all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
data created at /tmp/_b194zxa.wt5::/
  axes: ('w2', 'w1')
  shape: (111, 111)
axis w2 converted from wn to eV
axis w1 converted from wn to eV
/home/docs/checkouts/readthedocs.org/user_builds/wrighttools/checkouts/stable/examples/norms.py:75: RuntimeWarning: divide by zero encountered in log10
  mantissa = lambda x: np.mod(np.log10(x), -1)
/home/docs/checkouts/readthedocs.org/user_builds/wrighttools/checkouts/stable/examples/norms.py:75: RuntimeWarning: invalid value encountered in log10
  mantissa = lambda x: np.mod(np.log10(x), -1)
/home/docs/checkouts/readthedocs.org/user_builds/wrighttools/checkouts/stable/examples/norms.py:75: RuntimeWarning: invalid value encountered in remainder
  mantissa = lambda x: np.mod(np.log10(x), -1)

import WrightTools as wt
from WrightTools import datasets

import matplotlib.colors as mpl_colors
from matplotlib import colormaps

import numpy as np

# --- colormaps -----------------------------------------------------------------------------------
unsigned_cmap = [
    wt.artists.colormaps["default"].copy(),
    colormaps["cubehelix_r"],
    colormaps["magma_r"],
    colormaps["viridis_r"],
][0]
unsigned_cmap.set_under([0.9, 0.9, 0.9, 1])

cyclic_cmap = "twilight"

signed_cmap = ["twilight_shifted", "RdBu", "coolwarm", "seismic"][1]


# --- data ----------------------------------------------------------------------------------------
p = datasets.wt5.v1p0p0_perovskite_TA
signed = wt.open(p).at(d2=[-15, "fs"]).split("w1", [1.6])[1]
signed.bring_to_front("signal_diff")

unsigned = wt.data.from_PyCMDS(r"https://osf.io/75vny/download")
unsigned.bring_to_front("signal_diff")
unsigned.convert("eV")
unsigned.signal_diff.normalize()


# --- plot ----------------------------------------------------------------------------------------
label_kwargs = dict(fontsize=14, corner="LR", background_alpha=0.6)
fig, gs = wt.artists.create_figure(width=8, cols=[1, 1], nrows=4, wspace=1.3)


# --- unsigned data -------------------------------------------------------------------------------
# --- --- linear ----------------------------------------------------------------------------------
ax00 = fig.add_subplot(gs[0, 0], label="00 - unsigned linear")
wt.artists.corner_text("linear", **label_kwargs)
art = ax00.pcolormesh(unsigned, cmap=unsigned_cmap, autolabel="y")
fig.colorbar(art, ax=ax00, extend="min")

ax00.set_title("unsigned data")

# --- --- sqrt ------------------------------------------------------------------------------------
ax10 = fig.add_subplot(gs[1, 0], label="10 - unsigned sqrt")
wt.artists.corner_text("sqrt", **label_kwargs)
norm = mpl_colors.PowerNorm(gamma=0.5, vmin=0)
art = ax10.pcolormesh(unsigned, norm=norm, cmap=unsigned_cmap, autolabel="y")
fig.colorbar(art, ax=ax10, extend="min")

# --- --- log10 -----------------------------------------------------------------------------------
ax20 = fig.add_subplot(gs[2, 0], label="20 - unsigned log")
wt.artists.corner_text("log10", **label_kwargs)
norm = mpl_colors.LogNorm(vmin=5e-4, clip=False)
art = ax20.pcolormesh(unsigned, norm=norm, cmap=unsigned_cmap, autolabel="y")
cblog = fig.colorbar(art, ax=ax20, extend="min")

# --- --- log10, decadic cycles -------------------------------------------------------------------
# currently a bit hackish, but works.  We could make tools for this
ax30 = fig.add_subplot(gs[3, 0], label="30 - unsigned log cyclic")
wt.artists.corner_text("log10, cyclic", **label_kwargs)
mantissa = lambda x: np.mod(np.log10(x), -1)
unsigned.create_channel("mantissa", values=10 ** mantissa(unsigned.signal_diff[:]))
norm = mpl_colors.LogNorm()
art = ax30.pcolormesh(unsigned, channel="mantissa", norm=norm, cmap=cyclic_cmap, autolabel="both")
cb = fig.colorbar(art, ax=ax30)
cb.ax.yaxis.set_minor_formatter("{x:0.1f}")
cb.ax.set_yticks([0.1, 0.2, 0.5], minor=True)  # , labels=["0.2", "0.5"], minor=True)

# --- plot signed data ---------------------------------------------------------------------------
# --- --- linear norm ----------------------------------------------------------------------------
ax01 = fig.add_subplot(gs[0, 1], label="01 - signed linear")
ax01.set_title("signed data")
wt.artists.corner_text("linear", **label_kwargs)
art = ax01.pcolormesh(signed, cmap=signed_cmap, autolabel="y")
fig.colorbar(art, ax=ax01)

# --- --- bilinear norm --------------------------------------------------------------------------
ax11 = fig.add_subplot(gs[1, 1], label="11 - signed bilinear")
wt.artists.corner_text("bilinear", **label_kwargs)
norm = mpl_colors.TwoSlopeNorm(vcenter=signed.signal_diff.null)
art = ax11.pcolormesh(signed, norm=norm, cmap=signed_cmap, autolabel="y")
cb = fig.colorbar(art, ax=ax11)
cb.ax.set_yscale("linear")

# --- --- asinh norm ------------------------------------------------------------------------------
ax21 = fig.add_subplot(gs[2, 1], label="21 - signed asinh")
wt.artists.corner_text("asinh", **label_kwargs)
norm = mpl_colors.AsinhNorm(
    linear_width=1e-3, vmin=-signed.signal_diff.mag(), vmax=signed.signal_diff.mag()
)
art = ax21.pcolormesh(signed, norm=norm, cmap=signed_cmap, autolabel="y")
fig.colorbar(art, ax=ax21)

# --- --- symlog norm -----------------------------------------------------------------------------
ax31 = fig.add_subplot(gs[3, 1], label="31 - signed symlog")
wt.artists.corner_text("symLog", **label_kwargs)
norm = mpl_colors.SymLogNorm(
    linthresh=1e-3, vmin=-signed.signal_diff.mag(), vmax=signed.signal_diff.mag()
)
art = ax31.pcolormesh(signed, norm=norm, cmap=signed_cmap, autolabel="both")
fig.colorbar(art, ax=ax31)


# --- final decorations ---------------------------------------------------------------------------

for ax in fig.axes:
    label = ax.get_label()
    if label == "<colorbar>":
        continue
    ax.grid(c="k", lw=1, ls="-.", alpha=0.5)
    if label[0] != "3":  # not bottom row
        ax.set_xticks(ax.get_xticks(), ax.get_xticklabels(), visible=False)
    ax.set_xticks(ax.get_yticks())
    if label[1] == "0":  # unsigned
        ax.set_xlim(unsigned.w2.min(), unsigned.w2.max())
    else:
        ax.set_xlim(signed.w1.min(), signed.w1.max())