"""Creates RGB images from vector fields.
This module contains several routines for plotting colormaps from input
data consisting of 2D images of the 2D or 3D vector field. A variety of colormaps are
available using the Colorcet package.
https://colorcet.holoviz.org/user_guide/Continuous.html#cyclic-colormaps
Good Colour Maps: How to Design Them, Peter Kovesi (2015) https://arxiv.org/abs/1509.03700
"""
from __future__ import annotations
import colorsys
from typing import Optional, Union
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colors
from matplotlib.colors import Colormap
try:
import colorcet as cc
except ModuleNotFoundError:
cc = None
[docs]
def roll_cmap(
cmap: Union[Colormap, str],
frac: float,
invert: bool = False,
) -> Colormap:
"""
Shifts a matplotlib colormap by rolling.
Args:
cmap (mpl.colors.Colormap | str): The colormap to be shifted. Can be a colormap name or a Colormap object.
frac (float): The fraction of the colorbar by which to shift (must be between 0 and 1).
invert (bool, optional): Whether to invert the colormap. Defaults to False.
Returns:
mpl.colors.Colormap: The shifted colormap.
"""
N = 256
if isinstance(cmap, str):
cmap = get_cmap(cmap)
n = cmap.name
x = np.linspace(0, 1, N)
out = np.roll(x, int(N * frac))
if invert:
out = 1 - out
new_cmap = colors.LinearSegmentedColormap.from_list(f"{n}_s", cmap(out))
return new_cmap
[docs]
def shift_cmap_center(
cmap: Union[Colormap, str],
vmin: float = -1,
vmax: float = 1,
midpointval: float = 0,
invert: bool = False,
) -> Colormap:
"""
Shifts a matplotlib colormap such that the center is moved and the scaling is consistent.
Args:
cmap (mpl.colors.Colormap | str): The colormap to be shifted. Can be a colormap name or a Colormap object.
vmin (float, optional): Minimum value for scaling. Defaults to -1.
vmax (float, optional): Maximum value for scaling. Defaults to 1.
midpointval (float, optional): Value at the center of the colormap. Defaults to 0.
invert (bool, optional): Whether to invert the colormap. Defaults to False.
Returns:
mpl.colors.Colormap: The shifted colormap.
"""
assert vmax > vmin
midpoint_loc = (midpointval - vmin) / (vmax - vmin)
used_ratio = (vmax - vmin) / (2 * max(abs(vmax - midpointval), abs(vmin - midpointval)))
N = round(256 / used_ratio)
if isinstance(cmap, str):
cmap = get_cmap(cmap)
x = np.linspace(0, 1, N)
roll_frac = used_ratio if midpoint_loc < 0.5 else used_ratio * -1
roll_ind = round(N * roll_frac)
out = np.roll(x, roll_ind)
if midpoint_loc < 0.5:
out = out[:256]
elif midpoint_loc > 0.5:
out = out[-256:]
if invert:
out = 1 - out
new_cmap = colors.LinearSegmentedColormap.from_list(f"{cmap.name}_s", cmap(out))
return new_cmap
[docs]
def get_cmap(cmap: Optional[Union[str, Colormap]] = "linear", **kwargs) -> Colormap:
"""
Take a colormap or string input and return a Colormap object.
Args:
cmap (str | None, optional): String corresponding to a colorcet colormap name,
a mpl.colors.LinearSegmentedColormap object, or a
mpl.colors.ListedColormap object. Defaults to None -> matplotlib gray.
Cmap string options include:
- All matplotlib colormap names, e.g. "viridis"
You can see a full list of mpl colormaps by printing out plt.colormaps()
- Colorcet string names, e.g. "CET_L09" or "cet_CET_L09"
- "linear" -> mpl gray
- "diverging" -> mpl coolwarm
- "linear_cbl" -> cet CBL1 -- colorblind-safe linear
- "diverging_cbl" -> cet CBD1 -- colorblind-safe linear
- "cet_rainbow" -> cet R1 -- a good rainbow colormap
- "legacy4fold" -> cet C2 -- 4-fold colormap oriented as was the original PyLorentz default
- "purehsv" -> mpl hsv -- true hsv
- "6fold" or "hsv" -> cet C6 -- an improved hsv wheel
- "4fold" -> cet C7
- "isoluminant" -> cet C10
Keyword Args:
shift (float, optional): The amount to shift the colormap by in radians. Defaults to 0.
invert (bool, optional): Whether to invert the colormap. Defaults to False.
Raises:
TypeError: If the input type is not recognized.
Returns:
mpl.colors.Colormap: Matplotlib.colors.Colormap object.
"""
shift = kwargs.get("shift", 0)
invert = kwargs.get("invert", False)
if isinstance(cmap, colors.LinearSegmentedColormap) or isinstance(cmap, colors.ListedColormap):
return cmap
elif isinstance(cmap, str):
cmap_out = None
if cmap in plt.colormaps():
cmap_out = plt.get_cmap(cmap)
elif cmap.lower().startswith("cet") and cc is not None:
if "_" in cmap:
splits = cmap.split("_")
elif "-" in cmap:
splits = cmap.split("-")
else:
splits = ["CET", cmap[3:]]
# doesn't work for all, but parses many
cm2 = f"cet_CET_{splits[1].upper()}_{'_'.join(splits[2:])}".strip("_")
if cm2 in cc.colormaps():
cmap_out = plt.get_cmap(cm2)
elif "0" in cm2:
cm3 = cm2.replace("0", "")
if cm3 in cc.colormaps():
cmap_out = plt.get_cmap(cm3)
if cmap_out is None: # unable to find so far
try:
if cmap in ["linear", "lin", "", "default"]:
cmap_out = plt.get_cmap("gray")
elif cmap in ["diverging", "div"]:
cmap_out = plt.get_cmap("coolwarm")
elif cmap in ["linear_cbl", "cbl", "lin_cbl"] and cc is not None:
cmap_out = cc.cm["CET_CBL1"]
elif cmap in ["diverging_cbl", "div_cbl"] and cc is not None:
cmap_out = cc.cm["CET_CBD1"]
elif cmap in ["cet_rainbow", "cet_r1", "r1"] and cc is not None:
cmap_out = cc.cm.CET_R1
elif cmap in ["legacy4fold", "cet_c2", "c2", "cet_2"] and cc is not None:
cmap_out = cc.cm["CET_C2"]
shift += -np.pi / 2 # matching directions of legacy 4-fold
elif cmap in ["purehsv", "legacyhsv"]:
cmap_out = plt.get_cmap("hsv")
invert = not invert
shift += np.pi / 2
elif (
cmap
in ["cet_c6", "c6", "cet_6", "6fold", "6-fold", "sixfold", "hsv", "cyclic"]
and cc is not None
):
cmap_out = cc.cm["CET_C6"]
invert = not invert
shift += np.pi / 2
elif (
cmap in ["cet_c7", "c7", "cet_7", "4fold", "fourfold", "4-fold"]
and cc is not None
):
cmap_out = cc.cm["CET_C7"]
invert = not invert
elif cmap in ["cet_c8", "c8", "cet_8"] and cc is not None:
cmap_out = cc.cm["CET_C8"]
elif (
cmap in ["cet_c10", "c10", "cet_10", "isolum", "isoluminant", "iso"]
and cc is not None
):
cmap_out = cc.cm["CET_C10"]
elif cmap in ["cet_c11", "c11", "cet_11"] and cc is not None:
cmap_out = cc.cm["CET_C11"]
elif cmap in plt.colormaps():
cmap_out = plt.get_cmap(cmap)
else:
print(f"Unknown colormap input '{cmap}'.")
print("You can also pass a colormap object directly.")
print("Proceeding with default gray.")
cmap_out = plt.get_cmap("gray")
except NameError:
print("Colorcet not installed, proceeding with hsv from mpl")
cmap_out = plt.get_cmap("hsv")
invert = not invert
shift -= np.pi / 2
else:
raise TypeError(
f"Unknown input type {type(cmap)}, please input a matplotlib colormap or valid string"
)
if shift != 0: # given as radian convert to [0,1]
shift = shift % (2 * np.pi) / (2 * np.pi)
if shift != 0 or invert:
cmap_out = roll_cmap(cmap_out, shift, invert)
return cmap_out
[docs]
def color_im(
vx: np.ndarray,
vy: np.ndarray,
vz: Optional[np.ndarray] = None,
cmap: Optional[Union[str, Colormap]] = None,
rad: Optional[int] = None,
background: str = "black",
**kwargs,
) -> np.ndarray:
"""
Make the RGB image from vector maps. Takes 2D array inputs for x, y, (and
optionally z) vector components.
Unless otherwise specified, the color intensity corresponds to the magnitude of the
in-plane vector component normalized to the vector with the largest in-plane
magnitude. If a z-component is given, it will map from black (negative) to white
(positive).
Good colormaps are notoriously difficult to design [1], and cyclic colormaps
especially so. We recommend and use colormaps provided by the Colorcet package [2],
with the default being CET_C6, a nice 6-fold improved-hsv colorwheel. Other colormaps we suggest
include:
- C7: A nice 4-fold map
- C10: Isoluminescent 4-fold map
these can be specified with strings as detailed in get_cmap(). Additionally, any matplotlib or
colorcet Colormap object can be passed and will be used here.
[1] Kovesi, Peter. "Good colour maps: How to design them."
arXiv preprint arXiv:1509.03700 (2015).
[2] https://colorcet.holoviz.org
Args:
vx (np.ndarray): (M x N) array consisting of the x-component of the vector field.
vy (np.ndarray): (M x N) array consisting of the y-component of the vector field.
vz (np.ndarray | None, optional): (M x N) array consisting of the z-component of the vector field.
If vz is given, black corresponds to z<0 and white to z>0. Default is None.
cmap (str | mpl.colors.Colormap | None, optional): Specification for the colormap to be used.
This is passed to get_cmap() which returns the colormap object if a string is given.
Defaults to None -> colorcet.cm.CET_C7.
rad (int | None, optional): Radius of color-wheel in pixels. Set rad = 0 to remove color-wheel.
Default is None -> height/16.
background (str, optional):
- 'black' -- (default) magnetization magnitude corresponds to value.
- 'white' -- magnetization magnitude corresponds to saturation.
- if vz is given, this argument will not do anything. Default "black"
Keyword Args:
modulo (bool): Whether to map the direction or orientation of the vector field.
modulo = True will modulo the vectors by pi. Default False.
shift (float): Rotate the colorwheel and orientation map by the specified amount in radians. Default 0.
invert (bool): Whether or not to invert the directions of the orientation map. Default False.
mag_cutoff (bool): Normally the color intensity (saturation/value) corresponds to the magnitude of the vector,
scaled relative to the largest vector in the image. If mag_cutoff = True (specifying uniform_magnitude),
then all vectors larger with a magnitude larger than mag_cutoff_cutoff * max_magnitude will be displayed
while others will map to background colors or z-direction if vz is given. Default False.
Value [0,1], specifying the magnitude, as a fraction of the maximum vector length
in the image, above which a vector will be plotted. Default 0.5.
HSL (bool): When give a z-component, this function normally maps vector orientation to hue and vector magnitude
to saturation/value, with black/white corresponding to if the vector points in/out of the page.
This can cause problems as scaling RGB values can make it appear that there is strong in-plane signal
where there really isn't. An improvement is to use a HSL color space and map vector orientation to hue,
z-component to lightness, and in-plane magnitude to saturation. Setting HSL=True will use this type of mapping,
and set the colormap to a true hsv colormap, as more complex color maps get distorted when changing the
saturation and lightness independently. Default False.
Returns:
np.ndarray: Numpy array (M x N x 3) containing the RGB color image.
"""
vx, vy = np.squeeze(vx), np.squeeze(vy)
assert vx.ndim == vy.ndim == 2
assert np.shape(vy) == np.shape(vx)
if vz is not None:
HSL = kwargs.get("HSL", None)
if HSL is None:
HSL = kwargs.get("HLS", False) # i can never remember if it's HSL or HLS
else:
HSL = False
if HSL:
cmap = "purehsv"
elif cmap is None:
cmap = "hsv"
cmap = get_cmap(cmap, **kwargs)
if rad is None:
shape_frac = vx.shape[0] // 16
rad = max(shape_frac, 16)
raw_inp_mags = np.sqrt(vx**2 + vy**2)
if np.min(raw_inp_mags) == np.max(raw_inp_mags):
mags = np.ones_like(vx)
else:
mags = raw_inp_mags - np.min(raw_inp_mags)
mags = mags / np.max(mags) # normalize [0,1]
cutoff = kwargs.get("mag_cutoff")
if cutoff:
if not isinstance(cutoff, float):
cutoff = 0.4
mags = np.where(mags > cutoff, 1, 0)
if vz is None:
bkgs = np.where(mags == 0)
else:
bkgs = np.where(np.sqrt(vx**2 + vy**2 + vz**2) == 0)
if rad > 0:
pad = min(2 * (rad // 10), 40) # padding between edge of image and color-wheel
else:
pad = 0
rad = 0
dimy = np.shape(vy)[0]
if dimy < 2 * rad:
rad = dimy // 2
dimx = np.shape(vy)[1] + 2 * rad + pad
cimage = np.zeros((dimy, dimx, 3))
mod = kwargs.get("modulo", False)
if mod:
azimuth = np.mod((np.arctan2(vx, vy) + np.pi), mod) / mod
else:
azimuth = (np.arctan2(vx, vy) + np.pi) / (2 * np.pi)
# apply colormap to angle
imrgb = cmap(azimuth)[..., :3] # remove alpha channel
if vz is None:
if background.lower() == "black":
for i in range(3):
imrgb[:, :, i] *= mags
else:
for i in range(3):
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * mags
else: # mz > 1 -> white, mz < 1 -> black
theta = np.arctan2(vz, raw_inp_mags)
if HSL:
H = colors.rgb_to_hsv(imrgb)[:, :, 0]
# **3 is due to move more values closer to 0.5, max color intensity
L = (np.sin(theta) ** 3 + 1) / 2 # theta [-pi,pi] -> [0,1]
S = mags
S[bkgs] = 0
L[bkgs] = 0.5
for j in range(np.shape(vx)[0]):
for i in range(np.shape(vx)[1]):
imrgb[j, i, :] = colorsys.hls_to_rgb(H[j, i], L[j, i], S[j, i])
else:
if kwargs.get("mag_cutoff", False):
pos = np.where((np.sin(theta) > 0) & (mags == 0), 0, 1)
neg = np.where((np.sin(theta) < 0) & (mags == 0), 0, 1)
else:
neg = np.where(theta < 0, np.cos(theta) ** 2, 1)
pos = np.where(theta > 0, np.cos(theta) ** 2, 1)
for i in range(3):
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * pos
imrgb[:, :, i] *= neg
imrgb[:, :, i][bkgs] = 0.5
# colorwheel
if rad <= 0:
return imrgb
else:
cimage[:, : -2 * rad - pad, :] = imrgb
if vz is None:
wbkg = "black" if background == "white" else "white"
wheel = make_colorwheel(rad, cmap, background=wbkg, core=background, **kwargs)
if background == "black": # have white sidebar
cimage[:, dimx - 2 * rad - pad :] = 1
else:
wheel = make_colorwheelz(rad, cmap, **kwargs)
cimage[:, dimx - 2 * rad - pad :] = 0
cimage[:, dimx - 2 * rad - pad] = 1
cimage[
dimy // 2 - rad : dimy // 2 + rad,
dimx - 2 * rad - pad // 2 : -pad // 2,
:,
] = wheel
return cimage
[docs]
def make_colorwheel(
rad: int,
cmap: colors.Colormap,
background: str = "black",
core: Optional[Union[str, None]] = None,
**kwargs,
) -> np.ndarray:
"""
Makes an RGB image of a colorwheel for a given colormap.
Args:
rad (int): Radius of the colormap in pixels.
cmap (mpl.colors.Colormap): Matplotlib Colormap object.
background (str, optional): Background color. Defaults to "black".
core (str | None, optional): Core color. Defaults to background color.
Keyword Args:
modulo (bool): Whether to map the direction or orientation of the vector field.
modulo = True will modulo the vectors by pi. Default False.
mag_cutoff (bool): Cutoff value (0-1) for colormap showing orientation only, no magnitude.
Returns:
np.ndarray: Numpy array (rad*2 x rad*2 x 3) containing the RGB color image.
"""
cmap = get_cmap(cmap)
background = background.lower()
X, Y = np.mgrid[-rad:rad, -rad:rad]
mod = kwargs.get("modulo", False)
if mod:
azimuth = np.mod((np.arctan2(Y, X) + np.pi), mod) / mod
else:
azimuth = (np.arctan2(Y, X) + np.pi) / (2 * np.pi)
imrgb = cmap(azimuth)[..., :3]
rr = dist4(rad * 2)
mask = np.where(rr < rad, 1, 0)
rr *= mask
rr /= np.max(rr)
cutoff = kwargs.get("mag_cutoff")
if cutoff:
if not isinstance(cutoff, float):
cutoff = 0.4
rr = np.where(rr > cutoff, 1, 0)
if core is None:
core = background
if core == "black":
for i in range(3):
imrgb[:, :, i] *= rr * mask
else:
for i in range(3):
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * rr
if background == "black":
for i in range(3):
imrgb[:, :, i] = imrgb[:, :, i] * mask
else:
for i in range(3):
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * mask
return imrgb
[docs]
def make_colorwheelz(
rad: int,
cmap: Colormap,
outside: str = "black",
**kwargs,
) -> np.ndarray:
"""
Makes an RGB image of a colorwheel where z-direction corresponds to black/white.
Args:
rad (int): Radius of the colormap in pixels.
cmap (mpl.colors.Colormap): Matplotlib Colormap object.
outside (str, optional): Whether the outside portion of the colormap will be
black or white. Inside color will necessarily be the opposite. Defaults to "black".
Keyword Args:
modulo (bool): Whether to map the direction or orientation of the vector field.
modulo = True will modulo the vectors by pi. Default False.
mag_cutoff (bool): Ring colormap showing orientation only, no magnitude.
HSL (bool): Use HSL color space for mapping. Defaults to False.
Returns:
np.ndarray: Numpy array (rad*2 x rad*2 x 3) containing the RGB color image.
"""
cmap = get_cmap(cmap)
HSL = kwargs.get("HSL", None)
if HSL is None:
HSL = kwargs.get("HLS", False)
X, Y = np.mgrid[-rad:rad, -rad:rad]
mod = kwargs.get("modulo", False)
if mod:
azimuth = np.mod((np.arctan2(Y, X) + np.pi), mod) / mod
else:
azimuth = (np.arctan2(Y, X) + np.pi) / (2 * np.pi)
imrgb = cmap(azimuth)[..., :3]
rr = dist4(rad * 2)
mask = np.where(rr < rad, 1, 0)
rr *= mask
rr /= np.max(rr)
theta = (rr - 0.5) * np.pi
if kwargs.get("mag_cutoff", False):
cutoff = np.pi * 0.25
inner = np.where(theta + cutoff < 0, 0, 1)
outer = np.where(theta - cutoff > 0, 0, 1)
else:
if HSL:
H = colors.rgb_to_hsv(imrgb)[:, :, 0]
# **3 is due to move more values closer to 0.5, max color intensity
if outside == "black":
L = 1 - ((np.sin(theta) ** 3 + 1) / 2) # theta [-pi,pi] -> [0,1]
else:
L = (np.sin(theta) ** 3 + 1) / 2 # theta [-pi,pi] -> [0,1]
rr2 = np.abs(rr - 0.5) * -1 + 1 # rescale to maximize in middle
S = rr2**2
for j in range(2 * rad):
for i in range(2 * rad):
imrgb[j, i, :] = colorsys.hls_to_rgb(H[j, i], L[j, i], S[j, i])
if outside == "black":
for i in range(3):
imrgb[:, :, i] *= mask
else:
for i in range(3):
imrgb[:, :, i] += 1 - mask
inner = outer = 1 # for pylance
else:
inner = np.where(theta < 0, np.cos(theta) ** 2, 1)
outer = np.where(theta > 0, np.cos(theta) ** 2, 1)
if not HSL:
for i in range(3):
if outside == "black":
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * inner
imrgb[:, :, i] *= mask * outer
else:
imrgb[:, :, i] *= inner
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * outer
imrgb[:, :, i] = 1 - (1 - imrgb[:, :, i]) * mask
return imrgb
[docs]
def dist4(dim, shifted=True) -> np.ndarray:
"""
4-fold symmetric distance map (center is 0) even at small radii
centered in the middle (i.e. fft shifted) by default
"""
d = np.fft.fftfreq(dim, 1 / dim)
d = np.abs(d + 0.5) - 0.5 if dim % 2 == 0 else d
if shifted:
d = np.fft.fftshift(d)
rr = np.sqrt(d[None,] ** 2 + d[..., None] ** 2)
return rr
def _white_to_transparent_vec(image, magz=None):
rgba_image = np.ones((image.shape[0], image.shape[1], 4), dtype=float)
rgba_image[:, :, :3] = image[:, :, :3] # Convert RGB values to 0-255 range
if magz is not None:
alpha = np.where(magz > 0, 1 - magz**10, 1)
else:
print(rgba_image[:, :, :2].sum(axis=2).min(), rgba_image[:, :, :2].sum(axis=2).max())
alpha = np.where(rgba_image[:, :, :2].sum(axis=2) == 0, 0, 1).astype("float")
# alpha = rgba_image[:,:,:2].sum(axis=2)
# return alpha
rgba_image[:, :, 3] = alpha
return rgba_image