Merge branch 'master' into #93

msimet · Jul 14, 2018 · 0519cdb · 0519cdb
2 parents 966f332 + 6c01f37
commit 0519cdb
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 2 deletions.
diff --git a/stile/sys_tests.py b/stile/sys_tests.py
@@ -589,9 +589,16 @@ def plot(self, data, colors=['r', 'b'], log_yscale=False,
         fields = data.dtype.names
         # Pick which radius measurement to use
         # TreeCorr changed the name of the output columns
+        # This catches the case where we added the units to the label
+        fields_no_units = [f.split(' [')[0] for f in fields]
         for t_r in ['meanR', 'R_nom', '<R>', 'R_nominal', 'R']:
             if t_r in fields:
-                r = t_r
+                # Protect underscores since they blow up the plotting routines
+                r = '\\_'.join(t_r.split('\\'))
+                break
+            elif t_r in fields_no_units:
+                t_i = fields_no_units.index(t_r)
+                r = '\\_'.join(fields[t_i].split('\\'))
                 break
         else:
             raise ValueError('No radius parameter found in data')

diff --git a/tests/test_correlation_functions.py b/tests/test_correlation_functions.py
@@ -5,6 +5,7 @@
 import helper
 import unittest
 import treecorr
+import matplotlib
 
 try:
     import stile
@@ -81,7 +82,35 @@ def test_getCF(self):
         results3 = realshear(lens_data, source_data, config=stile_args)
         numpy.testing.assert_equal(results, results3)
 
-
+    def test_plot(self):
+        """ Test that the plotting routines successfully generate a plot """
+        stile_args = {'ra_units': 'degrees', 'dec_units': 'degrees', 'min_sep': 0.05, 'max_sep': 1,
+                      'sep_units': 'degrees', 'nbins': 20}
+        cf = stile.sys_tests.CorrelationFunctionSysTest()
+        lens_data = stile.ReadASCIITable('../examples/example_lens_catalog.dat',
+                    fields={'id': 0, 'ra': 1, 'dec': 2, 'z': 3, 'g1': 4, 'g2': 5})
+        source_data = stile.ReadASCIITable('../examples/example_source_catalog.dat',
+                    fields={'id': 0, 'ra': 1, 'dec': 2, 'z': 3, 'g1': 4, 'g2': 5})
+        object_list = ['GalaxyShear', 'BrightStarShear', 'StarXGalaxyShear', 'StarXStarShear']
+        for object_type in object_list:
+            obj = stile.CorrelationFunctionSysTest(object_type)
+            results = obj(lens_data, source_data, **stile_args)
+            pl = obj.plot(results)
+            self.assertIsInstance(pl, matplotlib.figure.Figure)
+        # Test underscore protection. If there's an underscore in the radius label somewhere,
+        # get rid of the non-underscore versions to make sure we hit that branch of the code,
+        # then test the plotting again
+        names = list(results.dtype.names)
+        names_no_units = [n.split(' [')[0] for n in names]
+        if 'R_nom' in names_no_units or 'R_nominal' in names_no_units:
+            for rname in ['meanR', '<R>', 'R']:
+                if rname in names_no_units:
+                    index = names_no_units.index(rname)
+                    names[index] = 'old_'+names[index]
+            results.dtype.names = names
+            pl = obj.plot(results)
+            self.assertIsInstance(pl, matplotlib.figure.Figure)
+            pl.savefig('examine.png')            
     def test_generator(self):
         """Make sure the CorrelationFunctionSysTest() generator returns the right objects"""
         object_list = ['GalaxyShear', 'BrightStarShear', 'StarXGalaxyDensity',  'StarXGalaxyShear',