Ticket #4529: logplots.py

from sage.all import *
import os, types

ALLOWED_EXTENSIONS = ['.eps', '.pdf', '.png', '.ps', '.svg']

class LogGraphics(SageObject):
    r"""
    class LogGraphics is use to create a graphics object which is one of
    the 'loglog', 'semilogx' or 'semilogy' plot type graphic. If both of
    your axes are linear then use the Graphics() class or plot() function
    of Sage directly.

    Examples::

    # This declares that p is an instance of type 'loglog' plot
        sage: p = LogGraphics('loglog')
        sage: pplot = p(range(10), map(lambda i:e^i, range(10)))
        sage: q = LogGraphics('loglog')
        sage: qplot = q(range(10), map(lambda i:5^i, range(10)))
        sage: r = pplot + qplot # This works

    # This declares that q is an instance of type 'semilogx' plot
        sage: q = LogGraphics('semilogx')
        sage: qplot = q(range(10), map(lambda i:5^i, range(10)))
        sage: r = pplot + qplot # Error. Can not combine different scalings
    """

    def __add__(self, g):
        if not isinstance(g, LogGraphics):
            raise TypeError("You can not add a LogGraphics object to "
            "a non LogGraphics object")
        if self._plottype != g._plottype:
            raise TypeError("Both the graphics must have the same scale "
            "along both the x-axis and y-axis")
        G = LogGraphics(self._plottype)
        G._xdata = self._xdata + g._xdata
        G._ydata = self._ydata + g._ydata
        G._options = self._options + g._options
        return G

    def __call__(self, xdata, ydata, **options):
        self._xdata = [xdata]
        self._ydata = [ydata]
        self._options[0].update(**options)
        return self

    def __init__(self, plottype):
        if plottype not in ['loglog', 'semilogx', 'semilogy']:
            raise ValueError("plottype must be one of 'loglog', 'semilogx'"
            " or 'semilogy'")
        self._options   = [dict(plotjoined=True, color='blue')]
        self._plottype  = plottype
        self._xdata     = []
        self._ydata     = []

    def _repr_(self):
        if sage.plot.plot.SHOW_DEFAULT:
            self.show()
            return ''
        else:
            return 'LogGraphics object of type %s with %d plots'%(
                self._plottype, len(self._xdata))

    def matplotlib(self, **options):
        r"""
        This function calls the actual matplotlib commands. It uses pylab
        """
        axes_labels = []        # x and y axis labels
        gopts       = dict()    # grid options
        has_legend  = False     # Whether there is a legend
        lopts       = dict()    # legend options
        title       = ''        # Title of the plot

        # Fix the xmin and xmax of the data
        [xmin, xmax] = self.get_min_max(self._xdata)
        [ymin, ymax] = self.get_min_max(self._ydata, flag=False)


        import matplotlib.pyplot as Pyl
        Pyl.hold(False)
        # Ensure that matplotlib does not flip the figure along x-axis
        Pyl.axis([xmin, xmax, ymin, ymax])

        for i in xrange(len(self._xdata)):
            o = self._options[i].copy()
            o.update(**options)
            if 'plotjoined' in o:
                if o['plotjoined']:
                    if 'linestyle' not in o:
                        o.update(linestyle='-')
                else:
                    if 'marker' not in o:
                        o.update(marker='o')
                    o.update(linestyle='')
                del o['plotjoined']
            
            for _o in o.copy():
                if _o == 'axes_labels':
                    axes_labels = o.pop(_o)
                elif _o == 'color':
                    from sage.plot.colors import rgbcolor, Color, to_mpl_color
                    o.update(color=to_mpl_color(o[_o]))
                elif _o == 'fancybox':
                    lopts.update(fancybox=o[_o])
                    del o[_o]
                elif _o == 'fontsize':
                    has_legend = True
                    from matplotlib.font_properties import FontProperties
                    f = FontProperties(size=o[_o])
                    lopts.update(prop=f)
                    del o[_o]
                elif _o == 'gridlines':
                    gopts.update(b=o[_o])
                    gopts.update(which='both')
                    del o[_o]
                elif _o == 'legend_label' or _o == 'legend':
                    has_legend = True
                    o.update(label=o[_o])
                    del o[_o]
                elif _o == 'linewidth':
                    o.update(linewidth=float(o[_o]))
                elif _o == 'loc':
                    lopts[_o] = o.pop(_o)
                elif _o == 'title':
                    title = o.pop(_o)
                elif _o == 'markersize':
                    o.update(markersize=float(o[_o]))
                elif _o == 'rgbcolor':
                    from sage.plot.colors import rgbcolor, Color, to_mpl_color
                    o.update(color=to_mpl_color(o[_o]))
                    del o[_o]
                elif _o == 'size':
                    o.update(markersize=float(o[_o]))
                    del o[_o]
                elif _o == 'thickness':
                    o.update(linewidth=float(o[_o]))
                    del o[_o]
                    
            if self._plottype == 'loglog':
                Pyl.loglog(self._xdata[i], self._ydata[i], **o)
            elif self._plottype == 'semilogx':
                Pyl.semilogx(self._xdata[i], self._ydata[i], **o)
            elif self._plottype == 'semilogy':
                Pyl.semilogy(self._xdata[i], self._ydata[i], **o)

            # Now set the hold
            Pyl.hold(True)
            
        if len(gopts) > 0:
            Pyl.grid(**gopts)
        if has_legend or len(lopts) > 0:
            Pyl.legend(**lopts)
        if title != '':
            Pyl.title(title)
        if axes_labels != []:
            Pyl.xlabel(axes_labels[0])
            Pyl.ylabel(axes_labels[1])

        return Pyl

    def get_min_max(self, list_of_lists, flag=None):
        """
        def get_min_max(list_of_lists)
        Returns the minimum of all possible values among all the lists in
        list_of_lists.
        
        Each list in list_of_lists is assumed to be sorted in either
        increasing or decreasing order. If any list is sorted in decreasing
        order, then the final output has lmin > lmax.

        Input:  a list of lists
        Output: [lmin, lmax], which corresponds to the min and max points
        of an axis.

        Examples::

            sage: r=[10^-i for i in range(5)]; s=[1/_ for _ in r]
            sage: LogGraphics('loglog').get_min_max([r, s])
            [10000, 1/10000]
            sage: t=[2*_ for _ in s]
            sage: LogGraphics('loglog').get_min_max([t, s])
            [1, 20000]
        """
        # When calling with .ydata, just get the min and max.
        # Don't reorder y-axis. To do this, just set flag=False while
        # calling this function.
        if flag is None:
            flag = False
            for l in list_of_lists:
                if l[0] > l[-1]:
                    flag = True
                    break

        if flag:
            # The plot will be decreasing along the axis. So, it must start
            # from largest and go to smallest.
            lmin = max([l[0]  if l[0] > l[-1] else l[-1] for l in
                        list_of_lists])
            lmax = min([l[-1] if l[0] > l[-1] else l[0]  for l in
                        list_of_lists])
        else:
            # This is the usual stuff. values increase along the axis
            lmin = min([l[0]  if l[0] < l[-1] else l[-1] for l in
                        list_of_lists])
            lmax = max([l[-1] if l[0] < l[-1] else l[0]  for l in
                        list_of_lists])

        return [lmin, lmax]


    def save(self, **options):
        r"""
        Save the graphics to an image file.
        """
        filename = options.pop('filename')
        if filename is None:
            filename = sage.misc.misc.graphics_filename()
        ext = os.path.splitext(filename)[1].lower()
        
        if ext not in ALLOWED_EXTENSIONS:
            raise ValueError("allowed file extensions for images are '"
                             + "', '".join(ALLOWED_EXTENSIONS) + "'!")
        else:
            figure = self.matplotlib(**options)
            figure.savefig(filename) # matplotlib handles the file ext.

    def show(self, **options):
        r"""
        Show the image on the sage notebook or on the system image viewer
        """
        # This option should not be passed on to save().
        linkmode = options.pop('linkmode', False)
        if sage.misc.latex.EMBEDDED_MODE:
            options.setdefault('filename', sage.misc.misc.graphics_filename())
            self.save(**options)
            if linkmode == True:
                return "<img src='cell://%s'>" % options['filename']
            else:
                html("<img src='cell://%s'>" % options['filename'])
        else:
            options.setdefault('filename', sage.misc.misc.tmp_filename()
                    + '.png')
            self.save(**options)
            os.system('%s %s 2>/dev/null 1>/dev/null &'
                      % (sage.misc.viewer.browser(), options['filename']))

#------------------------------------------------------------------------#
#                                                                        #
#   User facing functions are now given below. Main functions are        #
#   log_list_plot() and log_plot(). log_plot() calls log_list_plot()     #
#   eventually, but it first generates xdata and ydata.                  #
#                                                                        #
#------------------------------------------------------------------------#


def log_list_plot(xdata, ydata, scale, **options):
    r"""
    log_list_plot(xdata, ydata, scale, **options):
    Plot a loglog, semilogx, or semilogy plot of the points taken as
    pairwise tuples from xdata and ydata. xdata is a list of points on the
    horizontal axis and ydata is the corresponding list on the vertical
    axis.
    The scale has to be one of three strings: 'loglog', 'semilogx',
    'semilogy'.
    The additional relevant options are:
        - axes_labels: list of two strings. The first for the horizontal
          axis and the second for the vertical axis. Default None.
        - color: set the color of the line or marker. Default 'blue'
        - fontsize: the size of the fonts used in legend. Default ??
        - gridlines: whether to draw grid. Default is False
        - legend_label: the label for the legend in the plot. Default None
        - legend: (See legend_label)
        - loc: The location of the legend. See http://goo.gl/G5BBR
        - marker:   If you want to place a marker at every point then
          provide it as a string. Examples are '*', 'o', '^', etc.
          See http://goo.gl/e27PT for all the options. Default 'o'
        - markersize: the size of the marker. Default ??
        - size: (See markersize)
        - title: The title of the plot. Default is None
        - linestyle: The style of the line. Examples are '-', '--', '-.',
          etc. Default is '-' if plotjoined is True and '' if False.
          See http://goo.gl/wn1Cy for other options
        - plot_points: The number of points to evaluate the function at.
          Default 200
        - plotjoined:  Whether to draw a curve through the points at which
          the function is evaluated. Default True
        - rgbcolor: tuple of three numbers, each between 0 and 1, which
          represents the color of the line or marker as an RGB tuple
        - thickness: the thickness of the line drawn in the plot. Default 1
        - linewidth: (See thickness)

    Examples::
        
        sage: xdata = range(1,10); ydata = map(lambda i: 10^i, xdata)
        sage: log_list_plot(xdata, ydata, 'semilogy', 
                color='red', plotjoined=False)
        sage: semilogy_list_plot(xdata, ydata, 'semilogx',
                rgbcolor=(1,0,0))
        sage: semilogy_list_plot(xdata, ydata, 'loglog',
                legend_label='a plot')
    """
    if scale != 'loglog' and scale != 'semilogx' and scale != 'semilogy':
        raise ValueError("scale must be one of 'loglog', 'semilogx' or "
            "'semilogy'" )
    return LogGraphics(scale)(xdata, ydata, **options)

def log_plot(f, scale, *args, **options):
    r"""
    log_plot(f, scale, *args, **options):
    Plot a loglog, semilogx, or semilogy plot of the function f. The scale
    argument must be provided and it has to one of the strings 'loglog',
    'semilogx' or 'semilogy' to indicate which kind of plot is desired.

    The relevant options are:
        - axes_labels: list of two strings. The first for the horizontal
          axis and the second for the vertical axis. Default None.
        - color: set the color of the line or marker. Default 'blue'
        - fontsize: the size of the fonts used in legend. Default ??
        - gridlines: whether to draw grid. Default is False
        - legend_label: the label for the legend in the plot. Default None
        - legend: (See legend_label)
        - loc: The location of the legend. See http://goo.gl/G5BBR
        - marker:   If you want to place a marker at every point then
          provide it as a string. Examples are '*', 'o', '^', etc.
          See http://goo.gl/e27PT for all the options. Default 'o'
        - markersize: the size of the marker. Default ??
        - size: (See markersize)
        - title: The title of the plot. Default is None
        - linestyle: The style of the line. Examples are '-', '--', '-.',
          etc. Default is '-' if plotjoined is True and '' if False.
          See http://goo.gl/wn1Cy for other options
        - plot_points: The number of points to evaluate the function at.
          Default 200
        - plotjoined:  Whether to draw a curve through the points at which
          the function is evaluated. Default True
        - rgbcolor: tuple of three numbers, each between 0 and 1, which
          represents the color of the line or marker as an RGB tuple
        - thickness: the thickness of the line drawn in the plot. Default 1
        - linewidth: (See thickness)

    Examples::

        sage: log_plot(e^x, 'loglog', (x, 1, 10), color='red',
                plotjoined=False)
        sage: log_plot(e^x, 'semilogx', 1, 10, rgbcolor=(1,0,0),
                plotjoined=False)
        sage: log_plot(e^x, 'semilogy', legend_label='a plot', thickness=3)
        sage: log_plot(e^x, 'semilogy', marker='*', size=20)

    """
    if scale != 'loglog' and scale != 'semilogx' and scale != 'semilogy':
        raise ValueError("scale must be one of 'loglog', 'semilogx' or "
            "'semilogy'" )

    # Don't want to expose this function to outside
    def _generate_data(f, *args, **options):
        r"""
        This function is used internally by log_plot
        """
        n = len(args)
        if n == 0:
            ranges = (0, 1)
        elif n == 1:
            if len(args[0]) == 2:
                ranges = args[0]
            elif len(args[0]) == 3:
                ranges = args[0]
            else:
                raise ValueError("You must pass either a tuple of 2 entries "
                    "[for ex: (1, 2)] or of 3 entries [for ex: (x, 1, 2)] "
                    "after the function.")
        elif n == 2 or n == 3:
            ranges = tuple(args)
        else:
            raise KeyError("Too many arguments passed to function. Any extra"
                " arguments must be passed as keywords [ex: use color='red' "
                " instead of 'red'].")
        from sage.plot.misc import setup_for_eval_on_grid
        from sage.plot.plot import generate_plot_points
        plot_points = options.pop('plot_points', 200)
        f, ranges = setup_for_eval_on_grid(f, [ranges], plot_points=plot_points)
        data = generate_plot_points(f, ranges[0][:2], plot_points)
        return zip(*data)

    xdata, ydata = _generate_data(f, *args, **options)
    if 'plot_points' in options:
        del options['plot_points']

    return log_list_plot(xdata, ydata, scale, **options)

def loglog_list_plot(xdata, ydata, **options):
    r"""
    loglog_list_plot(xdata, ydata, **options):
    Both the horizontal and vertical axes are plotted on a logarithmic
    scale.

    See the documentation of log_list_plot for all the options.

    Examples::
        
        sage: xdata = range(1,10); ydata = map(lambda i: 10^i, xdata)
        sage: loglog_list_plot(xdata, ydata, color='red', plotjoined=False)
        sage: loglog_list_plot(xdata, ydata, rgbcolor=(1,0,0))
        sage: loglog_list_plot(xdata, ydata, legend_label='a plot')
    """
    return log_list_plot(xdata, ydata, 'loglog', **options)

def loglog_plot(f, *args, **options):
    r"""
    loglog_plot(f, *args, **options):
    Plot a log-log plot of the function f.  Both the horizontal and
    vertical axes are on a logarithmic scale.

    See the documention of log_plot for other options.

    Examples::

        sage: loglog_plot(e^x, (x, 1, 10), color='red', plotjoined=False)
        sage: loglog_plot(e^x, 1, 10, rgbcolor=(1,0,0), plotjoined=False)
        sage: loglog_plot(e^x, legend_label='a plot', thickness=3)
        sage: loglog_plot(e^x, marker='*', size=20)
    """
    return log_plot(f, 'loglog', *args, **options)

def semilogx_list_plot(xdata, ydata, **options):
    r"""
    semilogx_list_plot(xdata, ydata, **options):
    Plot a semilogx function, with vertical axis on linear scale and
    horizontal axis on log scale.

    See the documentation of log_list_plot for all the options.

    Examples::
        
        sage: xdata = range(1,10); ydata = map(lambda i: 10^i, xdata)
        sage: semilogx_list_plot(xdata, ydata, color='red', plotjoined=False)
        sage: semilogx_list_plot(xdata, ydata, rgbcolor=(1,0,0))
        sage: semilogx_list_plot(xdata, ydata, legend_label='a plot')
    """
    return log_list_plot(xdata, ydata, 'semilogx', **options)

def semilogx_plot(f, *args, **options):
    r"""
    semilogx_plot(f, *args, **options):
    Plot a semilogx plot of the function f. The horizontal axis is
    on a logarithmic scale and the vertical axis is on a linear scale.

    See the documention of log_plot for other options.

    Examples::

        sage: semilogx_plot(e^x, (x, 1, 10), color='red', plotjoined=False)
        sage: semilogx_plot(e^x, 1, 10, rgbcolor=(1,0,0), plotjoined=False)
        sage: semilogx_plot(e^x, legend_label='a plot', thickness=3)
        sage: semilogx_plot(e^x, marker='*', size=20)
    """
    return log_plot(f, 'semilogx', *args, **options)

def semilogy_list_plot(xdata, ydata, **options):
    r"""
    semilogy_list_plot(xdata, ydata, **options):
    Plot a semilogy function, with vertical axis on log scale and
    horizontal axis on linear scale.

    See the documentation of log_list_plot for all the options.

    Examples::
        
        sage: xdata = range(1,10); ydata = map(lambda i: 10^i, xdata)
        sage: semilogy_list_plot(xdata, ydata, color='red', plotjoined=False)
        sage: semilogy_list_plot(xdata, ydata, rgbcolor=(1,0,0))
        sage: semilogy_list_plot(xdata, ydata, legend_label='a plot')
    """
    return log_list_plot(xdata, ydata, 'semilogy', **options)

def semilogy_plot(f, *args, **options):
    r"""
    semilogy_plot(f, *args, **options):
    Plot a semilogy plot of the function f. The horizontal axis is
    on a linear scale while the vertical axis is logarithmic scale.

    See the documention of log_plot for other options.

    Examples::

        sage: semilogy_plot(e^x, (x, 1, 10), color='red', plotjoined=False)
        sage: semilogy_plot(e^x, 1, 10, rgbcolor=(1,0,0), plotjoined=False)
        sage: semilogy_plot(e^x, legend_label='a plot', thickness=3)
        sage: semilogy_plot(e^x, marker='*', size=20)
    """
    return log_plot(f, 'semilogy', *args, **options)