Repository: StingraySoftware/stingray
Pull request #338 (2018-03-01 15:03:36)

Created new plot function in cross spectrum
This PR contains:
1. `plot` in `crossspectrum.py`
2. a fix in `crosscorrelation.py` (taken help from my previous fix(#300) in `lightcurve.py`

Credits go to the person who wrote `plot` function in other classes.

Is there any other class that is missing `plot`?

@matteobachetti @abigailStev 


Patch:
diff --git a/stingray/crosscorrelation.py b/stingray/crosscorrelation.py
index 1882aa01..7e085e56 100644
--- a/stingray/crosscorrelation.py
+++ b/stingray/crosscorrelation.py
@@ -246,6 +246,8 @@ class CrossCorrelation(object):
                 plt.savefig('corr.pdf', format="pdf")
             else:
                 plt.savefig(filename)
+        else:
+            plt.show(block=False)
 
         return ax
 
diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 9f09d22e..ec804745 100644
--- a/stingray/crossspectrum.py
+++ b/stingray/crossspectrum.py
@@ -540,6 +540,73 @@ class Crossspectrum(object):
         else:
             raise AttributeError("Object has no attribute named 'time_lag' !")
 
+    def plot(self, labels=None, axis=None, title=None, marker='-', save=False,
+             filename=None):
+        """
+        Plot the amplitude of the cross spectrum vs. the frequency using ``matplotlib``.
+
+        Parameters
+        ----------
+        labels : iterable, default ``None``
+            A list of tuple with ``xlabel`` and ``ylabel`` as strings.
+
+        axis : list, tuple, string, default ``None``
+            Parameter to set axis properties of the ``matplotlib`` figure. For example
+            it can be a list like ``[xmin, xmax, ymin, ymax]`` or any other
+            acceptable argument for the``matplotlib.pyplot.axis()`` method.
+
+        title : str, default ``None``
+            The title of the plot.
+
+        marker : str, default '-'
+            Line style and color of the plot. Line styles and colors are
+            combined in a single format string, as in ``'bo'`` for blue
+            circles. See ``matplotlib.pyplot.plot`` for more options.
+
+        save : boolean, optional, default ``False``
+            If ``True``, save the figure with specified filename.
+
+        filename : str
+            File name of the image to save. Depends on the boolean ``save``.
+        """
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise ImportError("Matplotlib required for plot()")
+
+        fig = plt.figure('crossspectrum')
+        fig = plt.plot(self.freq, np.abs(self.power), marker, color='b', label='Amplitude')
+        fig = plt.plot(self.freq, np.abs(self.power.real), marker, color='r', alpha=0.5, label='Real Part')
+        fig = plt.plot(self.freq, np.abs(self.power.imag), marker, color='g', alpha=0.5, label='Imaginary Part')
+
+        if labels is not None:
+            try:
+                plt.xlabel(labels[0])
+                plt.ylabel(labels[1])
+            except TypeError:
+                utils.simon("``labels`` must be either a list or tuple with "
+                            "x and y labels.")
+                raise
+            except IndexError:
+                utils.simon("``labels`` must have two labels for x and y "
+                            "axes.")
+                # Not raising here because in case of len(labels)==1, only
+                # x-axis will be labelled.
+        plt.legend(loc='best')
+        if axis is not None:
+            plt.axis(axis)
+
+        if title is not None:
+            plt.title(title)
+
+        if save:
+            if filename is None:
+                plt.savefig('spec.png')
+            else:
+                plt.savefig(filename)
+        else:
+            plt.show(block=False)
+
 
 class AveragedCrossspectrum(Crossspectrum):
     """


Test patch:
diff --git a/stingray/tests/test_crossspectrum.py b/stingray/tests/test_crossspectrum.py
index fb0ffdb4..b94639a0 100644
--- a/stingray/tests/test_crossspectrum.py
+++ b/stingray/tests/test_crossspectrum.py
@@ -2,6 +2,7 @@ from __future__ import division
 import numpy as np
 import pytest
 import warnings
+import matplotlib.pyplot as plt
 from stingray import Lightcurve, AveragedPowerspectrum
 from stingray import Crossspectrum, AveragedCrossspectrum, coherence, time_lag
 from stingray import StingrayError
@@ -241,6 +242,9 @@ class TestCrossspectrum(object):
         with pytest.raises(AttributeError):
             lag = obj.time_lag()
 
+    def test_plot_simple(self):
+        self.cs.plot()
+        assert plt.fignum_exists('crossspectrum')
 
 class TestAveragedCrossspectrum(object):