Add a helper class PlotW7 for comparison plots.

tardis/tests/tests_slow/plot_w7.py

@@ -0,0 +1,38 @@
+import os
+import matplotlib.pyplot as plt
+
+
+class PlotW7:
+    """
+    A helper class to TestW7 for creating comparison plots.
+    """
+
+    def __init__(self, obtained, reference):
+        self.obtained = obtained
+        self.reference = reference
+
+    def plot_luminosity_density_lambda(self, passed):
+        spectrum_fig = plt.figure()
+        spectrum_fig.suptitle("Deviation in spectrum quantities", fontweight="bold")
+
+        # `ldl_` prefixed variables associated with `luminosity_density_lambda`.
+        ldl_axes = spectrum_fig.add_subplot(111)
+        ldl_axes.set_title("Deviation in luminosity_density_lambda")
+        ldl_axes.set_xlabel("N-th packet")
+        ldl_axes.set_ylabel("Relative error (1 - obtained / baseline)")
+        deviation = 1 - (
+            self.obtained.spectrum.luminosity_density_lambda.value /
+            self.reference['luminosity_density_lambda'].value)
+
+        if passed:
+            ldl_axes.text(0.9, 0.9, 'PASSED', transform=ldl_axes.transAxes,
+                          bbox={'facecolor': 'green', 'alpha': 0.5, 'pad': 10})
+            ldl_axes.plot(deviation, "g+")
+        else:
+            ldl_axes.set_yscale("log")
+            ldl_axes.text(0.9, 0.9, 'FAILED', transform=ldl_axes.transAxes,
+                          bbox={'facecolor': 'red', 'alpha': 0.5, 'pad': 10})
+            ldl_axes.plot(deviation, "rx")
+
+        # Figure is saved in `tmp` directory of slow tests baseline data.
+        plt.savefig(os.path.join("/tmp", "spectrum_plot.png"))

tardis/tests/tests_slow/test_w7.py

@@ -1,14 +1,14 @@
 import os
 import yaml
 import pytest
-import matplotlib.pyplot as plt
 from numpy.testing import assert_allclose
 from astropy.tests.helper import assert_quantity_allclose
 
 from tardis.atomic import AtomData
 from tardis.simulation.base import Simulation
 from tardis.model import Radial1DModel
 from tardis.io.config_reader import Configuration
+from tardis.tests.tests_slow.plot_w7 import PlotW7
 
 
 class TestW7(object):
@@ -57,6 +57,9 @@ def setup(self, request, reference, data_path, atomic_data_fname,
         # Get the reference data through the fixture.
         self.reference = reference
 
+        # An instance of helper class for plotting
+        self.plotter = PlotW7(self.obtained_radial1d_model, self.reference)
+
     def test_j_estimators(self):
         assert_allclose(
                 self.reference['j_estimators'],
@@ -116,36 +119,15 @@ def test_spectrum(self):
                 self.reference['wavelength'],
                 self.obtained_radial1d_model.spectrum.wavelength)
 
-        spectrum_fig = plt.figure()
-        spectrum_fig.suptitle("Deviation in spectrum quantities", fontweight="bold")
-
-        # `ldl_` prefixed variables associated with `luminosity_density_lambda`.
-        ldl_axes = spectrum_fig.add_subplot(111)
-        ldl_axes.set_title("Deviation in luminosity_density_lambda")
-        ldl_axes.set_xlabel("N-th packet")
-        ldl_axes.set_ylabel("Relative error (1 - obtained / baseline)")
-        deviation = 1 - (
-            self.obtained_radial1d_model.spectrum.luminosity_density_lambda.value /
-            self.reference['luminosity_density_lambda'].value)
-
         try:
             assert_quantity_allclose(
                     self.reference['luminosity_density_lambda'],
                     self.obtained_radial1d_model.spectrum.luminosity_density_lambda)
-
-            ldl_axes.text(0.9, 0.9, 'PASSED', transform=ldl_axes.transAxes,
-                          bbox={'facecolor': 'green', 'alpha': 0.5, 'pad': 10})
-            ldl_axes.plot(deviation, "g+")
+            self.plotter.plot_luminosity_density_lambda(passed=True)
         except Exception as e:
-            ldl_axes.set_yscale("log")
-            ldl_axes.text(0.9, 0.9, 'FAILED', transform=ldl_axes.transAxes,
-                          bbox={'facecolor': 'red', 'alpha': 0.5, 'pad': 10})
-            ldl_axes.plot(deviation, "rx")
+            self.plotter.plot_luminosity_density_lambda(passed=False)
             raise e
 
-        # Figure is saved in `tmp` directory of slow tests baseline data.
-        plt.savefig(os.path.join("/tmp", "spectrum_plot.png"))
-
     def test_montecarlo_properties(self):
         assert_quantity_allclose(
                 self.reference['montecarlo_luminosity'],