diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py index baca33d..8a1567f 100644 --- a/tests/tools/test_plotting.py +++ b/tests/tools/test_plotting.py @@ -13,13 +13,22 @@ def test_plot_amplitude(show_plot): plotting.plot_amplitude(im, basis=np.array([[0, -1], [-1, 0], [0, 0]]), title='Test Amplitude') if show_plot: plt.show() - - # Test with numpy array - im = scipy.datasets.ascent().astype(np.complex128) + plt.close('all') + + # Test with numpy array and an extra dimension + im = np.stack([scipy.datasets.ascent().astype(np.complex128)]*3, axis=0) plotting.plot_amplitude(im, title='Test Amplitude') if show_plot: plt.show() - + plt.close('all') + + # Test with pytorch tensor and two extra dimensions + im = t.as_tensor(np.stack([im]*5, axis=0)) + plotting.plot_amplitude(im, title='Test Amplitude', + additional_axis_labels=['Hi','There']) + if show_plot: + plt.show() + plt.close('all') def test_plot_phase(show_plot): # Test with tensor @@ -27,12 +36,14 @@ def test_plot_phase(show_plot): plotting.plot_phase(im, title='Test Phase') if show_plot: plt.show() - + plt.close('all') + # Test with numpy array im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1]).numpy() plotting.plot_phase(im, title='Test Phase', basis=np.array([[0, -1], [-1, 0], [0, 0]])) if show_plot: plt.show() + plt.close('all') def test_plot_colorized(show_plot): @@ -42,9 +53,83 @@ def test_plot_colorized(show_plot): plotting.plot_colorized(im, title='Test Colorize', basis=np.array([[0, -1], [-1, 0], [0, 0]])) if show_plot: plt.show() + plt.close('all') - # Test with numpy array + # Test with numpy array and hsv im = im.numpy() - plotting.plot_colorized(im, title='Test Colorize') + plotting.plot_colorized(im, title='Test Colorize', use_cmocean=False) if show_plot: plt.show() + plt.close('all') + + +def test_plot_translations(show_plot): + rng = np.random.default_rng(0) + trans_np = rng.uniform(-5e-6, 5e-6, (20, 2)) + trans_t = t.as_tensor(trans_np) + + # numpy, defaults + plotting.plot_translations(trans_np) + if show_plot: + plt.show() + plt.close('all') + + # torch tensor and reuse figure + fig = plotting.plot_translations(trans_t) + plotting.plot_translations(trans_np, lines=False, color='red', label='scan', fig=fig, clear_fig=False) + if show_plot: + plt.show() + plt.close('all') + +def test_plot_nanomap(show_plot): + rng = np.random.default_rng(0) + trans_np = rng.uniform(-5e-6, 5e-6, (20, 2)) + values_np = np.random.default_rng(1).uniform(0, 1, 20) + trans_t = t.as_tensor(trans_np) + values_t = t.as_tensor(values_np) + + # numpy, defaults + plotting.plot_nanomap(trans_np, values_np) + if show_plot: + plt.show() + plt.close('all') + + # torch tensors + plotting.plot_nanomap(trans_t, values_t, units='nm', cmap_label='Intensity', convention='sample') + if show_plot: + plt.show() + plt.close('all') + + +def test_plot_nanomap_with_images(show_plot): + rng = np.random.default_rng(0) + trans_np = rng.uniform(-5e-6, 5e-6, (20, 2)) + values_np = np.random.default_rng(1).uniform(0, 1, 20) + # plot_nanomap_with_images requires tensor translations + trans_t = t.as_tensor(trans_np) + values_t = t.as_tensor(values_np) + + def get_image_2d(i): + return np.random.default_rng(i).uniform(0, 1, (32, 32)) + + def get_image_3d(i): + return np.random.default_rng(i).uniform(0, 1, (4, 32, 32)) + + # basic call, no values + plotting.plot_nanomap_with_images(trans_np, get_image_2d) + if show_plot: + plt.show() + plt.close('all') + + # with explicit values + plotting.plot_nanomap_with_images(trans_t, get_image_2d, values=values_np) + if show_plot: + plt.show() + plt.close('all') + + # 3D image stack + fig = plt.figure(figsize=(11,7)) + plotting.plot_nanomap_with_images(trans_np, get_image_3d, values=values_t, fig=fig) + if show_plot: + plt.show() + plt.close('all')