diff --git a/Framework/PythonInterface/test/python/mantid/plots/plots__init__Test.py b/Framework/PythonInterface/test/python/mantid/plots/plots__init__Test.py index 018bbed0675f8e9cd45734a9c87ed72ae2c72e17..b44dd4815c5a30311122e2a875a692bd4eeb12d4 100644 --- a/Framework/PythonInterface/test/python/mantid/plots/plots__init__Test.py +++ b/Framework/PythonInterface/test/python/mantid/plots/plots__init__Test.py @@ -9,10 +9,11 @@ from __future__ import (absolute_import, division, print_function) import matplotlib matplotlib.use('AGG') import matplotlib.pyplot as plt +from matplotlib.container import ErrorbarContainer import numpy as np import unittest -from mantid.simpleapi import CreateWorkspace, DeleteWorkspace +from mantid.simpleapi import CreateWorkspace, CreateSampleWorkspace, DeleteWorkspace class Plots__init__Test(unittest.TestCase): @@ -35,11 +36,89 @@ class Plots__init__Test(unittest.TestCase): def tearDownClass(cls): DeleteWorkspace('ws2d_histo') - def test_1d_plots(self): - fig, ax = plt.subplots(subplot_kw={'projection': 'mantid'}) - # ax.plot(self.ws2d_histo, 'rs', specNum=1) - ax.plot(self.ws2d_histo, specNum=2, linewidth=6) - ax.plot(np.arange(10), np.arange(10), 'bo-') + def setUp(self): + self.fig, self.ax = plt.subplots(subplot_kw={'projection': 'mantid'}) + + def tearDown(self): + plt.close('all') + self.fig, self.ax = None, None + + def test_line2d_plots(self): + self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6) + self.ax.plot(np.arange(10), np.arange(10), 'bo-') + + def test_errorbar_plots(self): + self.ax.errorbar(self.ws2d_histo, specNum=2, linewidth=6) + self.ax.errorbar(np.arange(10), np.arange(10), 0.1*np.ones((10,)), fmt='bo-') + + def test_imshow(self): + self.ax.imshow(self.ws2d_histo) + + def test_pcolor(self): + self.ax.pcolor(self.ws2d_histo) + + def test_pcolorfast(self): + self.ax.pcolorfast(self.ws2d_histo) + + def test_pcolormesh(self): + self.ax.pcolormesh(self.ws2d_histo) + + def test_remove_workspace_artist_for_known_workspace_removes_plot(self): + self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6) + is_empty = self.ax.remove_workspace_artists(self.ws2d_histo) + self.assertEqual(True, is_empty) + self.assertEqual(0, len(self.ax.lines)) + + def test_remove_workspace_artist_for_unknown_workspace_does_nothing(self): + self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6) + unknown_ws = CreateSampleWorkspace() + self.ax.remove_workspace_artists(unknown_ws) + self.assertEqual(1, len(self.ax.lines)) + + def test_remove_workspace_artist_for_removes_only_specified_workspace(self): + second_ws = CreateSampleWorkspace() + line_ws2d_histo = self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6)[0] + line_second_ws = self.ax.plot(second_ws, specNum=5)[0] + self.assertEqual(2, len(self.ax.lines)) + + self.ax.remove_workspace_artists(self.ws2d_histo) + self.assertEqual(1, len(self.ax.lines)) + self.assertTrue(line_ws2d_histo not in self.ax.lines) + self.assertTrue(line_second_ws in self.ax.lines) + DeleteWorkspace(second_ws) + + def test_replace_workspace_data_plot(self): + plot_data = CreateWorkspace(DataX=[10, 20, 30, 10, 20, 30, 10, 20, 30], + DataY=[3, 4, 5, 3, 4, 5], + DataE=[1, 2, 3, 4, 1, 1], + NSpec=3) + line_ws2d_histo = self.ax.plot(plot_data, specNum=2, color='r')[0] + plot_data = CreateWorkspace(DataX=[20, 30, 40, 20, 30, 40, 20, 30, 40], + DataY=[3, 4, 5, 3, 4, 5], + DataE=[1, 2, 3, 4, 1, 1], + NSpec=3) + self.ax.replace_workspace_artists(plot_data) + self.assertAlmostEqual(25, line_ws2d_histo.get_xdata()[0]) + self.assertAlmostEqual(35, line_ws2d_histo.get_xdata()[-1]) + self.assertEquals('r', line_ws2d_histo.get_color()) + + def test_replace_workspace_data_errorbar(self): + eb_data = CreateWorkspace(DataX=[10, 20, 30, 10, 20, 30, 10, 20, 30], + DataY=[3, 4, 5, 3, 4, 5], + DataE=[1, 2, 3, 4, 1, 1], + NSpec=3) + self.ax.errorbar(eb_data, specNum=2, color='r') + eb_data = CreateWorkspace(DataX=[20, 30, 40, 20, 30, 40, 20, 30, 40], + DataY=[3, 4, 5, 3, 4, 5], + DataE=[.1, .2, .3, .4, .1, .1], + NSpec=3) + self.ax.replace_workspace_artists(eb_data) + self.assertEqual(1, len(self.ax.containers)) + eb_container = self.ax.containers[0] + self.assertTrue(isinstance(eb_container, ErrorbarContainer)) + self.assertAlmostEqual(25, eb_container[0].get_xdata()[0]) + self.assertAlmostEqual(35, eb_container[0].get_xdata()[-1]) + self.assertEquals('r', eb_container[0].get_color()) def test_3d_plots(self): fig = plt.figure()