mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 21:42:39 +02:00
111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
|
|
from CDTools.tools import cmath
|
|
import torch as t
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.colors import hsv_to_rgb
|
|
|
|
|
|
def colorize(z):
|
|
""" Returns RGB values for a complex color plot given a complex array
|
|
This function returns a set of RGB values that can be used directly
|
|
in a call to imshow based on an input complex numpy array (not a
|
|
torch tensor representing a complex field)
|
|
|
|
Args:
|
|
z (array_like) : A complex-valued array
|
|
Returns:
|
|
list : A list of arrays for R,G, and B channels of an image.
|
|
|
|
"""
|
|
|
|
amp = np.abs(z)
|
|
rmin = 0
|
|
rmax = np.max(amp)
|
|
amp = np.where(amp < rmin, rmin, amp)
|
|
amp = np.where(amp > rmax, rmax, amp)
|
|
ph = np.angle(z, deg=1) + 90
|
|
# HSV are values in range [0,1]
|
|
h = (ph % 360) / 360
|
|
s = 0.85 * np.ones_like(h)
|
|
v = (amp - rmin) / (rmax - rmin)
|
|
|
|
return hsv_to_rgb(np.dstack((h,s,v)))
|
|
|
|
|
|
def plot_1d(im, **kwargs):
|
|
pass
|
|
|
|
def plot_amplitude(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs):
|
|
""" Plots the amplitude of a complex Tensor or numpy array with dimensions NxMx2.
|
|
Args:
|
|
im (t.Tensor) : An image with dimensions NxMx2.
|
|
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
|
|
a new figure is created with an Axes subplot at 111.
|
|
basis (array-like) : The probe basis, used to put the axis labels in real space units.
|
|
Should have dimensions 3x2
|
|
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
|
|
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
|
|
"""
|
|
if fig is None:
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(111, **kwargs)
|
|
basis_norm = np.linalg.norm(basis, axis = -1)
|
|
if isinstance(im, t.Tensor):
|
|
absolute = cmath.cabs(im).detach().cpu().numpy()
|
|
else:
|
|
absolute = np.absolute(im)
|
|
plt.imshow(absolute, cmap = 'viridis', extent = [0, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]])
|
|
plt.colorbar()
|
|
return fig
|
|
|
|
def plot_phase(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs):
|
|
""" Plots the phase of a complex Tensor or numpy array with dimensions NxMx2.
|
|
Args:
|
|
im (t.Tensor) : An image with dimensions NxMx2.
|
|
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
|
|
a new figure is created with an Axes subplot at 111.
|
|
basis (array-like) : The probe basis, used to put the axis labels in real space units.
|
|
Should have dimensions 3x2
|
|
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
|
|
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
|
|
"""
|
|
if fig is None:
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(111, **kwargs)
|
|
# If the user has matplotlib >=3.0, use the preferred colormap
|
|
if isinstance(im, t.Tensor):
|
|
phase = cmath.cphase(im).detach().cpu().numpy()
|
|
else:
|
|
phase = np.angle(im)
|
|
basis_norm = np.linalg.norm(basis, axis = -1)
|
|
try: plt.imshow(phase, cmap = 'twilight', extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]])
|
|
except: plt.imshow(phase, cmap = 'hsv', extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]])
|
|
plt.colorbar()
|
|
return fig
|
|
|
|
def plot_colorized(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs):
|
|
""" Plots the colorized version of a complex Tensor or numpy array with dimensions NxMx2.
|
|
The darkness corresponds to the intensity of the image, and the color corresponds
|
|
to the phase.
|
|
|
|
Args:
|
|
im (t.Tensor) : An image with dimensions NxMx2.
|
|
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
|
|
a new figure is created with an Axes subplot at 111.
|
|
basis (array-like) : The probe basis, used to put the axis labels in real space units.
|
|
Should have dimensions 3x2
|
|
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
|
|
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
|
|
"""
|
|
if fig is None:
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(111, **kwargs)
|
|
if isinstance(im, t.Tensor):
|
|
im = cmath.torch_to_complex(im.detach().cpu())
|
|
basis_norm = np.linalg.norm(basis, axis = -1)
|
|
colorized = colorize(im)
|
|
plt.imshow(colorized, extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]])
|
|
return fig
|