Fix Plot crashing when given a MultiFit inside a list (fixes #246)

kafe2/fit/_base/plot.py

@@ -675,29 +675,23 @@ class Plot:
     subplot and legend management.
     """
 
-    FIT_INFO_STRING_FORMAT_CHI2 = textwrap.dedent(
-        """\
+    FIT_INFO_STRING_FORMAT_CHI2 = textwrap.dedent("""\
         {model_function}
         {parameters}
             $\\hookrightarrow${fit_quality}
             $\\hookrightarrow \\chi^2 \\, \\mathrm{{probability =}}${chi2_probability}
-    """
-    )
-    FIT_INFO_STRING_FORMAT_SATURATED = textwrap.dedent(
-        """\
+    """)
+    FIT_INFO_STRING_FORMAT_SATURATED = textwrap.dedent("""\
         {model_function}
         {parameters}
             $\\hookrightarrow${fit_quality}
-    """
-    )
-    FIT_INFO_STRING_FORMAT_NOT_SATURATED = textwrap.dedent(
-        """\
+    """)
+    FIT_INFO_STRING_FORMAT_NOT_SATURATED = textwrap.dedent("""\
         {model_function}
         {parameters}
             $\\hookrightarrow${cost}
             $\\hookrightarrow${fit_quality}
-    """
-    )
+    """)
 
     def __init__(self, fit_objects, separate_figures=False):
         """
@@ -725,9 +719,24 @@ def __init__(self, fit_objects, separate_figures=False):
             else:
                 fit_objects = [_f["fit"] for _f in _fit_history[fit_objects:]]
         try:
-            iter(fit_objects)
+            fit_objects = list(fit_objects)
         except TypeError:
-            fit_objects = (fit_objects,)
+            fit_objects = [fit_objects]
+
+        if self._multifit is None:
+            # A MultiFit is not iterable, so passing one via e.g. a list (as the plot() wrapper
+            # function does) ends up here rather than being caught by the isinstance check above.
+            _multifits_in_seq = [_fo for _fo in fit_objects if isinstance(_fo, MultiFit)]
+            if _multifits_in_seq:
+                if len(fit_objects) != 1:
+                    raise NotImplementedError(
+                        "Plotting a MultiFit together with other fits in the same Plot is not supported. "
+                        "Pass the MultiFit object on its own instead, e.g. Plot(multi_fit) instead of "
+                        "Plot([multi_fit, other_fit])."
+                    )
+                self._multifit = _multifits_in_seq[0]
+                fit_objects = self._multifit.fits
+
         self._from_container = tuple(isinstance(_fo, DataContainerBase) for _fo in fit_objects)
         self._fits = tuple(Fit(_fo) if _fc else _fo for _fc, _fo in zip(self._from_container, fit_objects))
 

kafe2/test/fit/test_plot.py

@@ -5,7 +5,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from kafe2 import ContoursProfiler, Plot, XYFit
+from kafe2 import ContoursProfiler, MultiFit, Plot, XYFit
+from kafe2.fit.util.wrapper import plot as wrapper_plot
 
 
 class TestXYPlot(unittest.TestCase):
@@ -280,6 +281,36 @@ def test_save_raise(self):
             self.plot_sep.save(fname=["fit_0.png", "fit_1.png", "fit_2.png"])
 
 
+class TestMultiFitPlot(unittest.TestCase):
+    # Regression tests for https://github.com/PhiLFitters/kafe2/issues/246
+    def setUp(self):
+        self.fit1 = XYFit(xy_data=[[0, 1, 2], [0.2, 1.1, 1.2]])
+        self.fit1.add_error("y", 0.1)
+        self.fit2 = XYFit(xy_data=[[0, 1, 2], [0.3, 0.9, 1.1]])
+        self.fit2.add_error("y", 0.1)
+        self.multi_fit = MultiFit([self.fit1, self.fit2])
+        self.multi_fit.do_fit()
+
+    def tearDown(self):
+        plt.close("all")
+
+    def test_plot_multifit_directly(self):
+        Plot(self.multi_fit).plot()
+
+    def test_plot_multifit_in_list(self):
+        Plot([self.multi_fit]).plot()
+
+    def test_plot_wrapper_with_multifit(self):
+        wrapper_plot(self.multi_fit, show=False, save=False)
+
+    def test_plot_multifit_mixed_with_other_fit_raises(self):
+        _fit3 = XYFit(xy_data=[[0, 1, 2], [0.3, 0.9, 1.1]])
+        _fit3.add_error("y", 0.1)
+        _fit3.do_fit()
+        with self.assertRaises(NotImplementedError):
+            Plot([self.multi_fit, _fit3])
+
+
 class TestContoursProfiler(unittest.TestCase):
     def setUp(self):
         self._ref_data = [[1, 2, 3], [0.9, 2.1, 3.0]]

kafe2/test/fit/test_representers_constraint_yaml.py

@@ -31,12 +31,9 @@
 uncertainty: 1.2
 """
 
-TEST_SIMPLE_GAUSSIAN_CONSTRAINT_EXTRA_KEYWORD = (
-    TEST_SIMPLE_GAUSSIAN_CONSTRAINT_ABS
-    + """
+TEST_SIMPLE_GAUSSIAN_CONSTRAINT_EXTRA_KEYWORD = TEST_SIMPLE_GAUSSIAN_CONSTRAINT_ABS + """
 extra_keyword: 3.14
 """
-)
 
 
 class TestSimpleGaussianConstraintRepresenter(unittest.TestCase):
@@ -133,12 +130,9 @@ def test_round_trip_with_stringstream(self):
 matrix: [[0.1, 0.1, 2.0], [0.1, 10.0, 30.0], [2.0, 30.0, 1000.0]]
 """
 
-TEST_MATRIX_GAUSSIAN_CONSTRAINT_EXTRA_KEYWORD = (
-    TEST_MATRIX_GAUSSIAN_CONSTRAINT_COV_ABS
-    + """
+TEST_MATRIX_GAUSSIAN_CONSTRAINT_EXTRA_KEYWORD = TEST_MATRIX_GAUSSIAN_CONSTRAINT_COV_ABS + """
 extra_keyword: 3.14
 """
-)
 
 
 class TestMatrixGaussianConstraintRepresenter(unittest.TestCase):

kafe2/test/fit/test_representers_container_yaml.py

@@ -81,12 +81,9 @@
   type: matrix
 """
 
-TEST_DATASET_INDEXED_EXTRA_KEYWORD = (
-    TEST_DATASET_INDEXED
-    + """
+TEST_DATASET_INDEXED_EXTRA_KEYWORD = TEST_DATASET_INDEXED + """
 extra_keyword: 3.14
 """
-)
 
 
 class TestIndexedContainerYamlRepresentation(unittest.TestCase):
@@ -252,12 +249,9 @@ def test_round_trip_with_stringstream(self):
   type: matrix
 """
 
-TEST_DATASET_XY_EXTRA_KEYWORD = (
-    TEST_DATASET_XY
-    + """
+TEST_DATASET_XY_EXTRA_KEYWORD = TEST_DATASET_XY + """
 extra_keyword: 3.14
 """
-)
 
 
 class TestXYContainerYamlRepresentation(unittest.TestCase):
@@ -421,12 +415,9 @@ def test_round_trip_with_stringstream(self):
 type: histogram
 """
 
-TEST_DATASET_HIST_EXTRA_KEYWORD = (
-    TEST_DATASET_HIST
-    + """
+TEST_DATASET_HIST_EXTRA_KEYWORD = TEST_DATASET_HIST + """
 extra_keyword: 3.14
 """
-)
 
 TEST_DATASET_HIST_MANUAL_HEIGHTS = """
 bin_edges:

kafe2/test/fit/test_representers_fit_yaml.py

@@ -126,12 +126,9 @@ def chi2(a=1.5, b=-0.5):
 type: custom
 """
 
-TEST_FIT_CUSTOM_EXTRA_KEYWORD = (
-    TEST_FIT_CUSTOM
-    + """
+TEST_FIT_CUSTOM_EXTRA_KEYWORD = TEST_FIT_CUSTOM + """
 extra_keyword: 3.14
 """
-)
 
 
 class TestCustomFitYamlRepresenter(unittest.TestCase, AbstractTestFitRepresenter):
@@ -257,12 +254,9 @@ def hist_model_density(x, mu, sigma):
     model_parameters: [0.1, 1.0]
 """
 
-TEST_FIT_HIST_EXTRA_KEYWORD = (
-    TEST_FIT_HIST
-    + """
+TEST_FIT_HIST_EXTRA_KEYWORD = TEST_FIT_HIST + """
 extra_keyword: 3.14
 """
-)
 
 TEST_FIT_HIST_SIMPLE = """
 type: histogram
@@ -440,12 +434,9 @@ def linear_model(a, b):
     model_parameters: [1.5, -0.5]
 """
 
-TEST_FIT_INDEXED_EXTRA_KEYWORD = (
-    TEST_FIT_INDEXED
-    + """
+TEST_FIT_INDEXED_EXTRA_KEYWORD = TEST_FIT_INDEXED + """
 extra_keyword: 3.14
 """
-)
 
 TEST_FIT_INDEXED_SIMPLE = """
 type: indexed
@@ -608,12 +599,9 @@ def linear_model(x, a, b):
     model_parameters: [1.5, -0.5]
 """
 
-TEST_FIT_XY_EXTRA_KEYWORD = (
-    TEST_FIT_XY
-    + """
+TEST_FIT_XY_EXTRA_KEYWORD = TEST_FIT_XY + """
 extra_keyword: 3.14
 """
-)
 
 TEST_FIT_XY_SIMPLE = """
 x_data: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
@@ -832,12 +820,9 @@ def normal_distribution_pdf(x, mu, sigma):
 type: unbinned
 """
 
-TEST_FIT_UNBINNED_EXTRA_KEYWORD = (
-    TEST_FIT_UNBINNED
-    + """
+TEST_FIT_UNBINNED_EXTRA_KEYWORD = TEST_FIT_UNBINNED + """
 extra_keyword: 3.14
 """
-)
 
 TEST_FIT_UNBINNED_SIMPLE = """
 type: unbinned

kafe2/test/fit/test_representers_format_yaml.py

@@ -129,12 +129,9 @@ def test_round_trip_with_stringstream(self):
 latex_expression_string: '{0}{x}^2 + {1}{x} + {2}'
 """
 
-TEST_MODEL_FUNCTION_FORMATTER_INDEXED_EXTRA_KEYWORD = (
-    TEST_MODEL_FUNCTION_FORMATTER_INDEXED
-    + """
+TEST_MODEL_FUNCTION_FORMATTER_INDEXED_EXTRA_KEYWORD = TEST_MODEL_FUNCTION_FORMATTER_INDEXED + """
 extra_keyword: 3.14
 """
-)
 
 
 class TestIndexedModelFunctionFormatterYamlRepresenter(unittest.TestCase):

kafe2/test/fit/test_representers_model_function_yaml.py

@@ -18,16 +18,11 @@
 type: histogram
 """
 
-TEST_MODEL_FUNCTION_HIST_EXTRA_KEYWORD = (
-    TEST_MODEL_FUNCTION_HIST
-    + """
+TEST_MODEL_FUNCTION_HIST_EXTRA_KEYWORD = TEST_MODEL_FUNCTION_HIST + """
 extra_keyword: 3.14
 """
-)
 
-TEST_MODEL_FUNCTION_HIST_WITH_FORMATTER = (
-    TEST_MODEL_FUNCTION_HIST
-    + r"""
+TEST_MODEL_FUNCTION_HIST_WITH_FORMATTER = TEST_MODEL_FUNCTION_HIST + r"""
 model_function_formatter:
     name: linear_model
     latex_name: linear model
@@ -42,7 +37,6 @@
     expression_string: '{0} * {x} + {1}'
     latex_expression_string: '{0}{x} + {1}'
 """
-)
 
 
 class TestHistModelFunctionYamlRepresenter(unittest.TestCase):
@@ -145,16 +139,11 @@ def linear_model(a, b):
 type: indexed
 """
 
-TEST_MODEL_FUNCTION_INDEXED_EXTRA_KEYWORD = (
-    TEST_MODEL_FUNCTION_INDEXED
-    + """
+TEST_MODEL_FUNCTION_INDEXED_EXTRA_KEYWORD = TEST_MODEL_FUNCTION_INDEXED + """
 extra_keyword: 3.14
 """
-)
 
-TEST_MODEL_FUNCTION_INDEXED_WITH_FORMATTER = (
-    TEST_MODEL_FUNCTION_INDEXED
-    + r"""
+TEST_MODEL_FUNCTION_INDEXED_WITH_FORMATTER = TEST_MODEL_FUNCTION_INDEXED + r"""
 model_function_formatter:
     name: linear_model
     latex_name: linear model
@@ -170,7 +159,6 @@ def linear_model(a, b):
     expression_string: '{0} * {r} + {1}'
     latex_expression_string: '{0}{r} + {1}'
 """
-)
 
 
 class TestIndexedModelFunctionYamlRepresenter(unittest.TestCase):
@@ -273,16 +261,11 @@ def test_round_trip_with_stringstream(self):
 type: base
 """
 
-TEST_MODEL_FUNCTION_XY_EXTRA_KEYWORD = (
-    TEST_MODEL_FUNCTION_BASE
-    + """
+TEST_MODEL_FUNCTION_XY_EXTRA_KEYWORD = TEST_MODEL_FUNCTION_BASE + """
 extra_keyword: 3.14
 """
-)
 
-TEST_MODEL_FUNCTION_XY_WITH_FORMATTER = (
-    TEST_MODEL_FUNCTION_BASE
-    + r"""
+TEST_MODEL_FUNCTION_XY_WITH_FORMATTER = TEST_MODEL_FUNCTION_BASE + r"""
 model_function_formatter:
     name: linear_model
     latex_name: linear model
@@ -297,7 +280,6 @@ def test_round_trip_with_stringstream(self):
     expression_string: '{0} * {x} + {1}'
     latex_expression_string: '{0}{x} + {1}'
 """
-)
 
 
 class TestModelFunctionBaseYamlRepresenter(unittest.TestCase):