Change color normalization in plt_slices

This commit is contained in:
Yin Li 2020-08-04 20:44:32 -04:00
parent d6b7a528df
commit 6b65b7888d

View File

@ -9,6 +9,10 @@ from matplotlib.colors import Normalize, LogNorm, SymLogNorm
from matplotlib.cm import ScalarMappable from matplotlib.cm import ScalarMappable
def quantize(x):
return 2 ** round(log2(x), ndigits=1)
def plt_slices(*fields, size=64, title=None, cmap=None, norm=None): def plt_slices(*fields, size=64, title=None, cmap=None, norm=None):
"""Plot slices of fields of more than 2 spatial dimensions. """Plot slices of fields of more than 2 spatial dimensions.
""" """
@ -39,9 +43,6 @@ def plt_slices(*fields, size=64, title=None, cmap=None, norm=None):
gridspec_kw={'height_ratios': nc * [im_size] + [cbar_height]} gridspec_kw={'height_ratios': nc * [im_size] + [cbar_height]}
) )
def quantize(x):
return 2 ** round(log2(x), ndigits=1)
for f, (field, cmap_col, norm_col) in enumerate(zip(fields, cmap, norm)): for f, (field, cmap_col, norm_col) in enumerate(zip(fields, cmap, norm)):
all_non_neg = (field >= 0).all() all_non_neg = (field >= 0).all()
all_non_pos = (field <= 0).all() all_non_pos = (field <= 0).all()
@ -61,19 +62,19 @@ def plt_slices(*fields, size=64, title=None, cmap=None, norm=None):
if all_non_neg: if all_non_neg:
if h1 > 0.1 * h2: if h1 > 0.1 * h2:
norm_col = Normalize(vmin=0, vmax=quantize(2 * h2)) norm_col = Normalize(vmin=0, vmax=quantize(h2))
else: else:
norm_col = LogNorm(vmin=quantize(0.5 * l2), vmax=quantize(2 * h2)) norm_col = LogNorm(vmin=quantize(l2), vmax=quantize(h2))
elif all_non_pos: elif all_non_pos:
warnings.warn('no implementation for all non-positive values') warnings.warn('no implementation for all non-positive values yet')
norm_col = None norm_col = None
else: else:
if w1 > 0.1 * w2: vlim = quantize(max(-l2, h2))
vlim = quantize(2.5 * w1) if w1 > 0.1 * w2 or l1 * h1 >= 0:
norm_col = Normalize(vmin=-vlim, vmax=vlim) norm_col = Normalize(vmin=-vlim, vmax=vlim)
else: else:
vlim = quantize(w2) linthresh = 0.1 * quantize(min(-l1, h1))
norm_col = SymLogNorm(linthresh=0.1 * w1, vmin=-vlim, vmax=vlim) norm_col = SymLogNorm(linthresh=linthresh, vmin=-vlim, vmax=vlim)
for c in range(field.shape[0]): for c in range(field.shape[0]):
s = (c,) + (0,) * (nd - 2) + (slice(64),) * 2 s = (c,) + (0,) * (nd - 2) + (slice(64),) * 2