Add center slicing in plt_slices
Also fix bug of hard coded `size`
This commit is contained in:
parent
9c8b331bf5
commit
61ca400942
@ -26,7 +26,6 @@ def plt_slices(*fields, size=64, title=None, cmap=None, norm=None):
|
|||||||
|
|
||||||
nc = max(field.shape[0] for field in fields)
|
nc = max(field.shape[0] for field in fields)
|
||||||
nf = len(fields)
|
nf = len(fields)
|
||||||
nd = fields[0].ndim - 1
|
|
||||||
|
|
||||||
if title is not None:
|
if title is not None:
|
||||||
assert len(title) == nf
|
assert len(title) == nf
|
||||||
@ -77,8 +76,19 @@ def plt_slices(*fields, size=64, title=None, cmap=None, norm=None):
|
|||||||
norm_col = SymLogNorm(linthresh=linthresh, 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
|
s0 = (c,) + tuple(d // 2 for d in field.shape[1:-2])
|
||||||
axes[c, f].pcolormesh(field[s], cmap=cmap_col, norm=norm_col)
|
s1 = (
|
||||||
|
slice(
|
||||||
|
(field.shape[-2] - size) // 2,
|
||||||
|
(field.shape[-2] + size) // 2,
|
||||||
|
),
|
||||||
|
slice(
|
||||||
|
(field.shape[-1] - size) // 2,
|
||||||
|
(field.shape[-1] + size) // 2,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
axes[c, f].pcolormesh(field[s0 + s1], cmap=cmap_col, norm=norm_col)
|
||||||
|
|
||||||
axes[c, f].set_aspect('equal')
|
axes[c, f].set_aspect('equal')
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user