text:
"""
DualEPS: A class to combine bitmap compression and vector graphics

"""
from __future__ import absolute_import, print_function

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------
import pyx
import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt

from yt.config import \
    ytcfg
from yt.utilities.logger import ytLogger as mylog
from .plot_window import PlotWindow
from .profile_plotter import PhasePlot, ProfilePlot
from yt.units.yt_array import YTQuantity
from yt.units.unit_object import Unit
from yt.funcs import issue_deprecation_warning

def convert_frac_to_tex(string):
    frac_pos = string.find(r'\frac')
    result = string[frac_pos+5:]
    level = [0]*len(result)
    clevel = 0
    for i in range(len(result)):
        if result[i] == "{":
            clevel += 1
        elif result[i] == "}":
            clevel -= 1
        level[i] = clevel
    div_pos = level.index(0)
    end_pos = level.index(0,div_pos+1)
    result = r'${' + result[:div_pos+1] + r'\over' + result[div_pos+1:end_pos] + \
             r'}$' + result[end_pos:]
    result = result.replace(r'\ ', r'\;')
    return result

def pyxize_label(string):
    frac_pos = string.find(r'\frac')
    if frac_pos >= 0:
        pre = string[:frac_pos]
        result = pre + convert_frac_to_tex(string)
    else:
        result = string
    result = result.replace('$', '')
    result = r'$' + result + r'$'
    return result

class DualEPS(object):
    def __init__(self, figsize=(12,12)):
        r"""Initializes the DualEPS class to which we can progressively add layers
        of vector graphics and compressed bitmaps.

        Parameters
        ----------
        figsize : tuple of floats
            The width and height of a single figure in centimeters.
        """
        pyx.unit.set(xscale=1.4)
        self.figsize = figsize
        self.canvas = None
        self.colormaps = None
        self.field = None
        self.axes_drawn = False

    def hello_world(self):
        r"""A simple test.
        """
        if self.canvas is None:
            self.canvas = pyx.canvas.canvas()
        p = pyx.path.line(0,0,1,1)
        self.canvas.stroke(p)
        self.canvas.text(0,0, "Hello world.")

#=============================================================================

    def return_field(self, plot):
        if isinstance(plot, (PlotWindow, PhasePlot)):
            return list(plot.plots.keys())[0]
        else:
            return None

#=============================================================================

    def axis_box(self, xrange=(0,1), yrange=(0,1), xlabel="", ylabel="",
                 xlog=False, ylog=False, xdata=None, ydata=None,
                 tickcolor=None, bare_axes=False,
                 pos=(0,0), xaxis_side=0, yaxis_side=0, size=None):
        r"""Draws an axis box in the figure.

        Parameters
        ----------
        xrange : tuple of floats
            The min and max of the x-axis
        yrange : tuple of floats
            The min and max of the y-axis
        xlabel : string
            Label for the x-axis
        ylabel : string
            Label for the y-axis
        xlog : boolean
            Flag to use a logarithmic x-axis
        ylog : boolean
            Flag to use a logarithmic y-axis
        tickcolor : `pyx.color.*.*`
            Color for the tickmarks.  Example: pyx.color.cmyk.black
        bare_axes : boolean
            Set to true to have no annotations or tick marks on all of the
            axes.
        pos : tuple of floats
            (x,y) position in centimeters of the origin in the figure
        xaxis_side : integer
            Set to 0 for the x-axis annotations to be on the left.  Set
            to 1 to print them on the right side.
        yaxis_side : integer
            Set to 0 for the y-axis annotations to be on the bottom.  Set
            to 1 to print them on the top.
        size : tuple of floats
            Size of axis box in units of figsize

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.save_fig()
        """

        if isinstance(xrange[0], YTQuantity):
            xrange = (xrange[0].value, xrange[1].value)
        if isinstance(yrange[0], YTQuantity):
            yrange = (yrange[0].value, yrange[1].value)
        if tickcolor is None:
            c1 = pyx.graph.axis.painter.regular\
                 (tickattrs=[pyx.color.cmyk.black])
            c2 = pyx.graph.axis.painter.regular\
                 (tickattrs=[pyx.color.cmyk.black], labelattrs=None)
        else:
            c1 = pyx.graph.axis.painter.regular(tickattrs=[tickcolor])
            c2 = pyx.graph.axis.painter.regular\
                 (tickattrs=[tickcolor], labelattrs=None)

        if size is None:
            psize = self.figsize
        else:
            psize = (size[0]*self.figsize[0], size[1]*self.figsize[1])

        xticklabels = True
        yticklabels = True
        if xaxis_side == 0:
            xleftlabel = xlabel
            xrightlabel = ""
            c1x = c1
            c2x = c2
        elif xaxis_side == 1:
            xleftlabel = ""
            xrightlabel = xlabel
            c1x = c2
            c2x = c1
        else:
            xticklabels = False
            xleftlabel = ""
            xrightlabel = ""
            c1x = c1
            c2x = c2
        if yaxis_side == 0:
            yleftlabel = ylabel
            yrightlabel = ""
            c1y = c1
            c2y = c2
        elif yaxis_side == 1:
            yleftlabel = ""
            yrightlabel = ylabel
            c1y = c2
            c2y = c1
        else:
            yticklabels = False
            yleftlabel = ""
            yrightlabel = ""
            c1y = c1
            c2y = c2

        if xlog:
            if xticklabels:
                xaxis = pyx.graph.axis.log(min=xrange[0],max=xrange[1],
                                           title=xleftlabel, painter=c1x)
                xaxis2 = pyx.graph.axis.log(min=xrange[0],max=xrange[1],
                                            title=xrightlabel, painter=c2x)
            else:
                xaxis = pyx.graph.axis.log(min=xrange[0],max=xrange[1],
                                           title=xleftlabel, painter=c1x,
                                           parter=None)
                xaxis2 = pyx.graph.axis.log(min=xrange[0],max=xrange[1],
                                            title=xrightlabel, painter=c2x,
                                            parter=None)
        else:
            if xticklabels:
                xaxis = pyx.graph.axis.lin(min=xrange[0],max=xrange[1],
                                           title=xleftlabel, painter=c1x)
                xaxis2 = pyx.graph.axis.lin(min=xrange[0],max=xrange[1],
                                            title=xrightlabel, painter=c2x)
            else:
                xaxis = pyx.graph.axis.lin(min=xrange[0],max=xrange[1],
                                           title=xleftlabel, painter=c1x,
                                           parter=None)
                xaxis2 = pyx.graph.axis.lin(min=xrange[0],max=xrange[1],
                                            title=xrightlabel, painter=c2x,
                                            parter=None)
        if ylog:
            if yticklabels:
                yaxis = pyx.graph.axis.log(min=yrange[0],max=yrange[1],
                                           title=yleftlabel, painter=c1y)
                yaxis2 = pyx.graph.axis.log(min=yrange[0],max=yrange[1],
                                            title=yrightlabel, painter=c2y)
            else:
                yaxis = pyx.graph.axis.log(min=yrange[0],max=yrange[1],
                                           title=yleftlabel, painter=c1y,
                                           parter=None)
                yaxis2 = pyx.graph.axis.log(min=yrange[0],max=yrange[1],
                                            title=yrightlabel, painter=c2y,
                                            parter=None)
        else:
            if yticklabels:
                yaxis = pyx.graph.axis.lin(min=yrange[0],max=yrange[1],
                                           title=yleftlabel, painter=c1y)
                yaxis2 = pyx.graph.axis.lin(min=yrange[0],max=yrange[1],
                                            title=yrightlabel, painter=c2y)
            else:
                yaxis = pyx.graph.axis.lin(min=yrange[0],max=yrange[1],
                                           title=yleftlabel, painter=c1y,
                                           parter=None)
                yaxis2 = pyx.graph.axis.lin(min=yrange[0],max=yrange[1],
                                            title=yrightlabel, painter=c2y,
                                            parter=None)

        if bare_axes:
            if ylog:
                yaxis = pyx.graph.axis.log(min=yrange[0], max=yrange[1],
                                           title=yleftlabel, parter=None)
                yaxis2 = pyx.graph.axis.log(min=yrange[0], max=yrange[1],
                                            title=yrightlabel, parter=None)
            else:
                yaxis = pyx.graph.axis.lin(min=yrange[0], max=yrange[1],
                                           title=yleftlabel, parter=None)
                yaxis2 = pyx.graph.axis.lin(min=yrange[0], max=yrange[1],
                                            title=yrightlabel, parter=None)
            if xlog:
                xaxis = pyx.graph.axis.log(min=xrange[0], max=xrange[1],
                                           title=xleftlabel, parter=None)
                xaxis2 = pyx.graph.axis.log(min=xrange[0], max=xrange[1],
                                            title=xrightlabel, parter=None)
            else:
                xaxis = pyx.graph.axis.lin(min=xrange[0], max=xrange[1],
                                           title=xleftlabel, parter=None)
                xaxis2 = pyx.graph.axis.lin(min=xrange[0], max=xrange[1],
                                            title=xrightlabel, parter=None)

        blank_data = pyx.graph.data.points([(-1e20,-1e20),(-1e19,-1e19)], x=1,y=2)
        if self.canvas is None:
            self.canvas = pyx.graph.graphxy \
                          (width=psize[0], height=psize[1],
                           x=xaxis, y=yaxis, x2=xaxis2, y2=yaxis2,
                           xpos=pos[0], ypos=pos[1])
            if xdata is None:
                self.canvas.plot(blank_data)
            else:
                data = pyx.graph.data.points(np.array([xdata, ydata]).T, x=1, y=2)
                self.canvas.plot(data, [pyx.graph.style.line([pyx.style.linewidth.Thick])])
        else:
            plot = pyx.graph.graphxy \
                   (width=psize[0], height=psize[1],
                    x=xaxis, y=yaxis, x2=xaxis2, y2=yaxis2,
                    xpos=pos[0], ypos=pos[1])
            if xdata is None:
                plot.plot(blank_data)
            else:
                data = pyx.graph.data.points(np.array([xdata, ydata]).T, x=1, y=2)
                plot.plot(data, [pyx.graph.style.line([pyx.style.linewidth.Thick])])
            self.canvas.insert(plot)
        self.axes_drawn = True

#=============================================================================

    def axis_box_yt(self, plot, units=None, bare_axes=False,
                    tickcolor=None, xlabel=None, ylabel=None, **kwargs):
        r"""Wrapper around DualEPS.axis_box to automatically fill in the
        axis ranges and labels from a yt plot.

        This also accepts any parameters that DualEPS.axis_box takes.

        Parameters
        ----------
        plot : `yt.visualization.plot_window.PlotWindow`
            yt plot on which the axes are based.
        units : string
            Unit description that overrides yt's unit description.  Only
            affects the axis label.
        bare_axes : boolean
            Set to true to have no annotations or tick marks on all of the
            axes.

        Examples
        --------
        >>> p = SlicePlot(ds, 0, 'density')
        >>> d = DualEPS()
        >>> d.axis_box_yt(p)
        >>> d.save_fig()
        """

        if isinstance(plot, (PlotWindow, PhasePlot)):
            plot.refresh()
        if isinstance(plot, PlotWindow):
            data = plot.frb
            width = plot.width[0]
            if units is None:
                units = plot.ds.get_smallest_appropriate_unit(width)
            width = width.in_units(str(units))
            xc = 0.5*(plot.xlim[0] + plot.xlim[1])
            yc = 0.5*(plot.ylim[0] + plot.ylim[1])
            _xrange = [(plot.xlim[i] - xc).in_units(units) for i in (0, 1)]
            _yrange = [(plot.ylim[i] - yc).in_units(units) for i in (0, 1)]
            _xlog = False
            _ylog = False
            if bare_axes:
                _xlabel = ""
                _ylabel = ""
            else:
                if xlabel is not None:
                    _xlabel = xlabel
                else:
                    if data.axis != 4:
                        xi = plot.ds.coordinates.x_axis[data.axis]
                        x_name = plot.ds.coordinates.axis_name[xi]
                        _xlabel = '%s (%s)' % (x_name, units)
                    else:
                        _xlabel = 'x (%s)' % (units)
                if ylabel is not None:
                    _ylabel = ylabel
                else:
                    if data.axis != 4:
                        yi = plot.ds.coordinates.y_axis[data.axis]
                        y_name = plot.ds.coordinates.axis_name[yi]
                        _ylabel = '%s (%s)' % (y_name, units)
                    else:
                        _ylabel = 'y (%s)' % (units)
            if tickcolor is None:
                _tickcolor = pyx.color.cmyk.white
        elif isinstance(plot, ProfilePlot):
            subplot = plot.axes.values()[0]
            # limits for axes
            xlimits = subplot.get_xlim()
            _xrange = (YTQuantity(xlimits[0], 'm'), YTQuantity(xlimits[1], 'm')) # unit hardcoded but afaik it is not used anywhere so it doesn't matter
            if list(plot.axes.ylim.viewvalues())[0][0] is None:
                ylimits = subplot.get_ylim()
            else:
                ylimits = list(plot.axes.ylim.viewvalues())[0]
            _yrange = (YTQuantity(ylimits[0], 'm'), YTQuantity(ylimits[1], 'm')) # unit hardcoded but afaik it is not used anywhere so it doesn't matter
            # axis labels
            xaxis = subplot.xaxis
            _xlabel = pyxize_label(xaxis.label.get_text())
            yaxis = subplot.yaxis
            _ylabel = pyxize_label(yaxis.label.get_text())
            # set log if necessary
            if subplot.get_xscale() == "log":
                _xlog = True
            else:
                _xlog = False
            if subplot.get_yscale() == "log":
                _ylog = True
            else:
                _ylog = False
            _tickcolor = None
        elif isinstance(plot, PhasePlot):
            k = list(plot.plots.keys())[0]
            _xrange = plot[k].axes.get_xlim()
            _yrange = plot[k].axes.get_ylim()
            _xlog = plot.profile.x_log
            _ylog = plot.profile.y_log
            if bare_axes:
                _xlabel = ""
                _ylabel = ""
            else:
                if xlabel is not None:
                    _xlabel = xlabel
                else:
                    _xlabel = plot[k].axes.get_xlabel()
                if ylabel is not None:
                    _ylabel = ylabel
                else:
                    _ylabel = plot[k].axes.get_ylabel()
                _xlabel = pyxize_label(_xlabel)
                _ylabel = pyxize_label(_ylabel)
            if tickcolor is None:
                _tickcolor = None
        elif isinstance(plot, np.ndarray):
            ax = plt.gca()
            _xrange = ax.get_xlim()
            _yrange = ax.get_ylim()
            _xlog=False
            _ylog=False
            if bare_axes:
                _xlabel = ""
                _ylabel = ""
            else:
                if xlabel is not None:
                    _xlabel = xlabel
                else:
                    _xlabel = ax.get_xlabel()
                if ylabel is not None:
                    _ylabel = ylabel
                else:
                    _ylabel = ax.get_ylabel()
            if tickcolor is None:
                _tickcolor = None
        else:
            _xrange = plot._axes.get_xlim()
            _yrange = plot._axes.get_ylim()
            _xlog = plot._log_x
            _ylog = plot._log_y
            if bare_axes:
                _xlabel = ""
                _ylabel = ""
            else:
                if xlabel is not None:
                    _xlabel = xlabel
                else:
                    _xlabel = plot._x_label
                if ylabel is not None:
                    _ylabel = ylabel
                else:
                    _ylabel = plot._y_label
            if tickcolor is None:
                _tickcolor = None
        if tickcolor is not None:
            _tickcolor = tickcolor
        self.axis_box(xrange=_xrange, yrange=_yrange, xlabel=_xlabel,
                      ylabel=_ylabel, tickcolor=_tickcolor, xlog=_xlog,
                      ylog=_ylog, bare_axes=bare_axes, **kwargs)

#=============================================================================

    def insert_image(self, filename, pos=(0,0), size=None):
        r"""Inserts a JPEG file in the figure.

        Parameters
        ----------
        filename : string
            Name of the JPEG file
        pos : tuple of floats
            Position of the origin of the image in centimeters
        size : tuple of flots
            Size of image in units of figsize

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.insert_image("image.jpg")
        >>> d.save_fig()
        """
        if size is not None:
            width = size[0]*self.figsize[0]
            height = size[1]*self.figsize[1]
        else:
            width = self.figsize[0]
            height = self.figsize[1]
        image = pyx.bitmap.jpegimage(filename)
        if self.canvas is None:
            self.canvas = pyx.canvas.canvas()
        self.canvas.insert(pyx.bitmap.bitmap(pos[0], pos[1], image,
                                             compressmode=None,
                                             width=width,
                                             height=height))

#=============================================================================

    def insert_image_yt(self, plot, field=None, pos=(0,0), scale=1.0):
        r"""Inserts a bitmap taken from a yt plot.

        Parameters
        ----------
        plot : `yt.visualization.plot_window.PlotWindow`
            yt plot that provides the image
        pos : tuple of floats
            Position of the origin of the image in centimeters.

        Examples
        --------
        >>> p = SlicePlot(ds, 0, 'density')
        >>> d = DualEPS()
        >>> d.axis_box_yt(p)
        >>> d.insert_image_yt(p)
        >>> d.save_fig()

        Notes
        -----
        For best results, set use_colorbar=False when creating the yt
        image.
        """
        from ._mpl_imports import FigureCanvasAgg

        # We need to remove the colorbar (if necessary), remove the
        # axes, and resize the figure to span the entire figure
        force_square = False
        if self.canvas is None:
            self.canvas = pyx.canvas.canvas()
        if isinstance(plot, (PlotWindow, PhasePlot)):
            if field is None:
                self.field = list(plot.plots.keys())[0]
                mylog.warning("No field specified.  Choosing first field (%s)" % \
                              str(self.field))
            else:
                self.field = plot.data_source._determine_fields(field)[0]
            if self.field not in plot.plots.keys():
                raise RuntimeError("Field '%s' does not exist!" % str(self.field))
            if isinstance(plot, PlotWindow):
                plot.hide_colorbar()
                plot.hide_axes()
            else:
                plot.plots[self.field]._toggle_axes(False)
                plot.plots[self.field]._toggle_colorbar(False)
            plot.refresh()
            _p1 = plot.plots[self.field].figure
            force_square = True
        elif isinstance(plot, ProfilePlot):
            _p1 = plot.figures.items()[0][1]
        elif isinstance(plot, np.ndarray):
            plt.figure()
            iplot = plt.figimage(plot)
            _p1 = iplot.figure
            _p1.set_size_inches(self.figsize[0], self.figsize[1])
            ax = plt.gca()
            _p1.add_axes(ax)
        else:
            raise RuntimeError("Unknown plot type")

        _p1.axes[0].set_position([0,0,1,1])  # rescale figure
        _p1.set_facecolor('w')  # set background color
        figure_canvas = FigureCanvasAgg(_p1)
        figure_canvas.draw()
        size = (_p1.get_size_inches() * _p1.dpi).astype('int')

        # Account for non-square images after removing the colorbar.
        scale *= 1.0 - 1.0 / (_p1.dpi * self.figsize[0])
        if force_square:
            yscale = scale * float(size[1]) / float(size[0])
        else:
            yscale = scale
        image = pyx.bitmap.image(size[0], size[1], "RGB",
                                 figure_canvas.tostring_rgb())
        self.canvas.insert(pyx.bitmap.bitmap(pos[0], pos[1], image,
                                             width=scale*self.figsize[0],
                                             height=yscale*self.figsize[1]))

#=============================================================================

    def colorbar(self, name, zrange=(0,1), label="", log=False, tickcolor=None,
                 orientation="right", pos=[0,0], shrink=1.0):
        r"""Places a colorbar adjacent to the current figure.

        Parameters
        ----------
        name : string
            name of the (matplotlib) colormap to use
        zrange : tuple of floats
            min and max of the colorbar's range
        label : string
            colorbar label
        log : boolean
            Flag to use a logarithmic scale
        tickcolor : `pyx.color.*.*`
            Color for the tickmarks.  Example: pyx.color.cmyk.black
        orientation : string
            Placement of the colorbar.  Can be "left", "right", "top",
            or "bottom".
        pos : list of floats
            (x,y) position of the origin of the colorbar in centimeters.
        shrink : float
            Factor to shrink the colorbar's size.  A value of 1 means the
            colorbar will have a height / width of the figure.

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.insert_image("image.jpg")
        >>> d.colorbar("hot", xrange=(1e-2, 1e-4), log=True,
                       label="Density [cm$^{-3}$]")
        >>> d.save_fig()
        """
        if orientation == "right":
            origin = (pos[0]+self.figsize[0]+0.5, pos[1])
            size = (0.1*self.figsize[0], self.figsize[1])
            imsize = (1,256)
        elif orientation == "left":
            origin = (pos[0]-0.5-0.1*self.figsize[0], pos[1])
            size = (0.1*self.figsize[0], self.figsize[1])
            imsize = (1,256)
        elif orientation == "top":
            origin = (pos[0], pos[1]+self.figsize[1]+0.5)
            imorigin = (pos[0]+self.figsize[0], pos[1]+self.figsize[1]+0.5)
            size = (self.figsize[0], 0.1*self.figsize[1])
            imsize = (256,1)
        elif orientation == "bottom":
            origin = (pos[0], pos[1]-0.5-0.1*self.figsize[1])
            imorigin = (pos[0]+self.figsize[0], pos[1]-0.5-0.1*self.figsize[1])
            size = (self.figsize[0], 0.1*self.figsize[1])
            imsize = (256,1)
        else:
            raise RuntimeError("orientation %s unknown" % orientation)
            return

        # If shrink is a scalar, then convert into tuple
        if not isinstance(shrink, (tuple,list)):
            shrink = (shrink, shrink)

        # Scale the colorbar
        shift = (0.5*(1.0-shrink[0])*size[0], 0.5*(1.0-shrink[1])*size[1])
        # To facilitate stretching rather than shrinking
        # If stretched in both directions (makes no sense?) then y dominates.
        if(shrink[0] > 1.0):
            shift = (0.05*self.figsize[0], 0.5*(1.0-shrink[1])*size[1])
        if(shrink[1] > 1.0):
            shift = (0.5*(1.0-shrink[0])*size[0], 0.05*self.figsize[1])
        size = (size[0] * shrink[0], size[1] * shrink[1])
        origin = (origin[0] + shift[0], origin[1] + shift[1])

        # Convert the colormap into a string
        x = np.linspace(1,0,256)
        cm_string = cm.cmap_d[name](x, bytes=True)[:,0:3].tostring()

        cmap_im = pyx.bitmap.image(imsize[0], imsize[1], "RGB", cm_string)
        if orientation == "top" or orientation == "bottom":
            imorigin = (imorigin[0] - shift[0], imorigin[1] + shift[1])
            self.canvas.insert(pyx.bitmap.bitmap(imorigin[0], imorigin[1], cmap_im,
                                                 width=-size[0], height=size[1]))
        else:
            self.canvas.insert(pyx.bitmap.bitmap(origin[0], origin[1], cmap_im,
                                                 width=size[0], height=size[1]))

        if tickcolor is None:
            c1 = pyx.graph.axis.painter.regular(tickattrs=[pyx.color.cmyk.black])
            pyx.graph.axis.painter.regular(tickattrs=[pyx.color.cmyk.black],
                                           labelattrs=None)
        else:
            c1 = pyx.graph.axis.painter.regular(tickattrs=[tickcolor])
            pyx.graph.axis.painter.regular(tickattrs=[tickcolor],
                                           labelattrs=None)
        if log:
            yaxis = pyx.graph.axis.log(min=zrange[0],max=zrange[1],
                                       title=label, painter=c1)
            yaxis2 = pyx.graph.axis.log(min=zrange[0],max=zrange[1],parter=None)
        else:
            yaxis = pyx.graph.axis.lin(min=zrange[0],max=zrange[1],
                                       title=label, painter=c1)
            yaxis2 = pyx.graph.axis.lin(min=zrange[0], max=zrange[1], parter=None)
        xaxis = pyx.graph.axis.lin(parter=None)

        if orientation == "right":
            _colorbar = pyx.graph.graphxy(width=size[0], height=size[1],
                                          xpos=origin[0], ypos=origin[1],
                                          x=xaxis, y=yaxis2, y2=yaxis)
        elif orientation == "left":
            _colorbar = pyx.graph.graphxy(width=size[0], height=size[1],
                                          xpos=origin[0], ypos=origin[1],
                                          x=xaxis, y2=yaxis2, y=yaxis)
        elif orientation == "top":
            _colorbar = pyx.graph.graphxy(width=size[0], height=size[1],
                                          xpos=origin[0], ypos=origin[1],
                                          y=xaxis, x=yaxis2, x2=yaxis)
        elif orientation == "bottom":
            _colorbar = pyx.graph.graphxy(width=size[0], height=size[1],
                                          xpos=origin[0], ypos=origin[1],
                                          y=xaxis, x2=yaxis2, x=yaxis)

        blank_data = pyx.graph.data.points([(-1e10,-1e10),(-9e10,-9e10)],
                                           x=1, y=2)
        _colorbar.plot(blank_data)
        self.canvas.insert(_colorbar)

#=============================================================================

    def colorbar_yt(self, plot, field=None, cb_labels = None, **kwargs):
        r"""Wrapper around DualEPS.colorbar to take information from a yt plot.

        Accepts all parameters that DualEPS.colorbar takes.

        Parameters
        ----------
        plot : A yt plot
            yt plot from which the information is taken.
        cb_labels : list of labels for the colorbars. List should be the same
                    size as the number of colorbars used. Should be passed
                    into this function by either the singleplot or multiplot api.

        Examples
        --------
        >>> p = SlicePlot(ds, 0, 'density')
        >>> p.hide_colorbar()
        >>> d = DualEPS()
        >>> d.axis_box_yt(p)
        >>> d.insert_image_yt(p)
        >>> d.colorbar_yt(p)
        >>> d.save_fig()
        """

        if isinstance(plot, ProfilePlot):
            raise RuntimeError("When using ProfilePlots you must either set yt_nocbar=True or provide colorbar flags so that the profiles don't have colorbars")
        _cmap = None
        if field is not None:
            self.field = plot.data_source._determine_fields(field)[0]
        if isinstance(plot, (PlotWindow, PhasePlot)):
            _cmap = plot._colormaps[self.field]
        else:
            if plot.cmap is not None:
                _cmap = plot.cmap.name
        if _cmap is None:
            _cmap = ytcfg.get("yt", "default_colormap")
        if isinstance(plot, (PlotWindow, PhasePlot)):
            if isinstance(plot, PlotWindow):
                try:
                    _zlabel = plot.frb[self.field].info["label"]
                    _unit = Unit(plot.frb[self.field].units,
                                 registry=plot.ds.unit_registry)
                    units = _unit.latex_representation()
                    # PyX does not support \frac because it's based on TeX.
                    units = pyxize_label(units)
                    _zlabel += r' (' + units + r')'
                except NotImplementedError:
                    print("Colorbar label not available")
                    _zlabel = ''
            else:
                _, _, z_title = plot._get_field_title(self.field, plot.profile)
                _zlabel = pyxize_label(z_title)
            _zlabel = _zlabel.replace("_","\;")
            _zlog = plot.get_log(self.field)[self.field]
            if plot.plots[self.field].zmin is None:
                zmin = plot.plots[self.field].image._A.min()
            else:
                zmin = plot.plots[self.field].zmin
            if plot.plots[self.field].zmax is None:
                zmax = plot.plots[self.field].image._A.max()
            else:
                zmax = plot.plots[self.field].zmax
            _zrange = (zmin, zmax)
        else:
            _zlabel = plot._z_label.replace("_","\;")
            _zlog = plot._log_z
            _zrange = (plot.norm.vmin, plot.norm.vmax)
        if cb_labels is not None:  #Overrides deduced labels
            _zlabel = cb_labels.pop()
        self.colorbar(_cmap, zrange=_zrange, label=_zlabel, log=_zlog, **kwargs)

#=============================================================================

    def circle(self, radius=0.2, loc=(0.5,0.5),
               color=pyx.color.cmyk.white,
               linewidth=pyx.style.linewidth.normal):
        r"""Draws a circle in the current figure.

        Parameters
        ----------
        radius : float
            Radius of the circle in units of figsize
        loc : tuple of floats
            Location of the circle's center in units of figsize
        color : `pyx.color.*.*`
            Color of the circle stroke.  Example: pyx.color.cmyk.white
        linewidth : `pyx.style.linewidth.*`
            Width of the circle stroke width. Example:
            pyx.style.linewidth.normal

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.insert_image("image.jpg")
        >>> d.circle(radius=0.1, color=pyx.color.cmyk.Red)
        >>> d.save_fig()
        """
        circle = pyx.path.circle(self.figsize[0]*loc[0],
                                 self.figsize[1]*loc[1],
                                 self.figsize[0]*radius)
        self.canvas.stroke(circle, [color, linewidth])

#=============================================================================

    def arrow(self, size=0.2, label="", loc=(0.05,0.08), labelloc="top",
              color=pyx.color.cmyk.white,
              linewidth=pyx.style.linewidth.normal,
              rotation=0.0):
        r"""Draws an arrow in the current figure

        Parameters
        ----------
        size : float
            Length of arrow (base to tip) in units of the figure size.
        label : string
            Annotation label of the arrow.
        loc : tuple of floats
            Location of the left hand side of the arrow in units of
            the figure size.
        rotation : float
            Orientation angle of the arrow in units of degrees
        labelloc : string
            Location of the label with respect to the line.  Can be
            "top" or "bottom"
        color : `pyx.color.*.*`
            Color of the arrow.  Example: pyx.color.cymk.white
        linewidth : `pyx.style.linewidth.*`
            Width of the arrow.  Example: pyx.style.linewidth.normal

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.insert_image("arrow_image.jpg")
        >>> d.arrow(size=0.2, label="Black Hole!", loc=(0.05, 0.1))
        >>> d.save_fig()
        """
        line = pyx.path.line(self.figsize[0]*loc[0],
                             self.figsize[1]*loc[1],
                             self.figsize[0]*(loc[0]+size*np.cos(np.pi*rotation/180)),
                             self.figsize[1]*(loc[1]+size*np.sin(np.pi*rotation/180)))
        self.canvas.stroke(line, [linewidth, color, pyx.deco.earrow()])

        if labelloc == "bottom":
            yoff = -0.1*size
            valign = pyx.text.valign.top
        else:
            yoff = +0.1*size
            valign = pyx.text.valign.bottom
        if label != "":
            self.canvas.text(self.figsize[0]*(loc[0]+0.5*size),
                             self.figsize[1]*(loc[1]+yoff), label,
                             [color, valign, pyx.text.halign.center])

#=============================================================================

    def scale_line(self, size=0.2, label="", loc=(0.05,0.08), labelloc="top",
                   color=pyx.color.cmyk.white,
                   linewidth=pyx.style.linewidth.normal):
        r"""Draws a scale line in the current figure.

        Parameters
        ----------
        size : float
            Length of the scale line in units of the figure size.
        label : string
            Annotation label of the scale line.
        loc : tuple of floats
            Location of the left hand side of the scale line in units of
            the figure size.
        labelloc : string
            Location of the label with respect to the line.  Can be
            "top" or "bottom"
        color : `pyx.color.*.*`
            Color of the scale line.  Example: pyx.color.cymk.white
        linewidth : `pyx.style.linewidth.*`
            Width of the scale line.  Example: pyx.style.linewidth.normal

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.insert_image("image.jpg")
        >>> d.scale_line(size=0.2, label="1 kpc", loc=(0.05, 0.1))
        >>> d.save_fig()
        """

        line = pyx.path.line(self.figsize[0]*loc[0],
                             self.figsize[1]*loc[1],
                             self.figsize[0]*(loc[0]+size),
                             self.figsize[1]*loc[1])
        self.canvas.stroke(line, [linewidth, color])
        line = pyx.path.line(self.figsize[0]*loc[0],
                             self.figsize[1]*(loc[1]-0.1*size),
                             self.figsize[0]*loc[0],
                             self.figsize[1]*(loc[1]+0.1*size))
        self.canvas.stroke(line, [linewidth, color])
        line = pyx.path.line(self.figsize[0]*(loc[0]+size),
                             self.figsize[1]*(loc[1]-0.1*size),
                             self.figsize[0]*(loc[0]+size),
                             self.figsize[1]*(loc[1]+0.1*size))
        self.canvas.stroke(line, [linewidth, color])

        if labelloc == "bottom":
            yoff = -0.1*size
            valign = pyx.text.valign.top
        else:
            yoff = +0.1*size
            valign = pyx.text.valign.bottom
        if label != "":
            self.canvas.text(self.figsize[0]*(loc[0]+0.5*size),
                             self.figsize[1]*(loc[1]+yoff), label,
                             [color, valign, pyx.text.halign.center])

#=============================================================================

    def title_box(self, text, color=pyx.color.cmyk.black,
                  bgcolor=pyx.color.cmyk.white, loc=(0.02,0.98),
                  halign=pyx.text.halign.left,
                  valign=pyx.text.valign.top,
                  text_opts=[]):
        r"""Inserts a box with text in the current figure.

        Parameters
        ----------
        text : string
            String to insert in the textbox.
        color : `pyx.color.*.*`
            Color of the text.  Example: pyx.color.cmyk.black
        bgcolor : `pyx.color.*.*`
            Color of the textbox background.  Example: pyx.color.cmyk.white
        loc : tuple of floats
            Location of the textbox origin in units of the figure size.
        halign : `pyx.text.halign.*`
            Horizontal alignment of the text.  Example: pyx.text.halign.left
        valign : `pyx.text.valign.*`
            Vertical alignment of the text.  Example: pyx.text.valign.top

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.insert_image("image.jpg")
        >>> d.title_box("Halo 1", loc=(0.05,0.95))
        >>> d.save_fig()
        """
        tbox = self.canvas.text(self.figsize[0]*loc[0],
                                self.figsize[1]*loc[1],
                                text, [color, valign, halign] + text_opts)
        if bgcolor is not None:
            tpath = tbox.bbox().enlarged(2*pyx.unit.x_pt).path()
            self.canvas.draw(tpath, [pyx.deco.filled([bgcolor]),
                                     pyx.deco.stroked()])
        self.canvas.insert(tbox)

#=============================================================================

    def save_fig(self, filename="test", format="eps", resolution=250):
        r"""Saves current figure to a file.

        Parameters
        ----------
        filename : string
            Name of the saved file without the extension.
        format : string
            Format type.  Can be "eps" or "pdf"

        Examples
        --------
        >>> d = DualEPS()
        >>> d.axis_box(xrange=(0,100), yrange=(1e-3,1), ylog=True)
        >>> d.save_fig("image1", format="pdf")
        """
        if format =="eps":
            self.canvas.writeEPSfile(filename)
        elif format == "pdf":
            self.canvas.writePDFfile(filename)
        elif format == "png":
            self.canvas.writeGSfile(filename+".png", "png16m", resolution=resolution)
        elif format == "jpg":
            self.canvas.writeGSfile(filename+".jpeg", "jpeg", resolution=resolution)
        else:
            raise RuntimeError("format %s unknown." % (format))

#=============================================================================
#=============================================================================
#=============================================================================

def multiplot(ncol, nrow, yt_plots=None, fields=None, images=None,
              xranges=None, yranges=None, xlabels=None, ylabels=None,
              xdata=None, ydata=None, colorbars=None,
              shrink_cb=0.95, figsize=(8,8), margins=(0,0), titles=None,
              savefig=None, format="eps", yt_nocbar=False, bare_axes=False,
              xaxis_flags=None, yaxis_flags=None,
              cb_flags=None, cb_location=None, cb_labels=None):
    r"""Convenience routine to create a multi-panel figure from yt plots or
    JPEGs.  The images are first placed from the origin, and then
    bottom-to-top and left-to-right.

    Parameters
    ----------
    ncol : integer
        Number of columns in the figure.
    nrow : integer
        Number of rows in the figure.
    yt_plots : list of yt plot instances
        yt plots to include in the figure.
    images : list of strings
        JPEG filenames to include in the figure.
    xranges : list of tuples
        The min and max of the x-axes
    yranges : list of tuples
        The min and max of the y-axes
    xlabels : list of strings
        Labels for the x-axes
    ylabels : list of strings
        Labels for the y-axes
    colorbars : list of dicts
        Dicts that describe the type of colorbar to be used in each
        figure.  Use the function return_colormap() to create these dicts.
    shrink_cb : float
        Factor by which the colorbar is shrunk.
    figsize : tuple of floats
        The width and height of a single figure in centimeters.
    margins : tuple of floats
        The horizontal and vertical margins between panels in centimeters.
    titles : list of strings
        Titles that are placed in textboxes in each panel.
    savefig : string
        Name of the saved file without the extension.
    format : string
        File format of the figure. eps or pdf accepted.
    yt_nocbar : boolean
        Flag to indicate whether or not colorbars are created.
    bare_axes : boolean
        Set to true to have no annotations or tick marks on all of the
        axes.
    cb_flags : list of booleans
        Flags for each plot to have a colorbar or not.
    cb_location : list of strings
        Strings to control the location of the colorbar (left, right,
        top, bottom)
    cb_labels : list of labels for the colorbars. List should be the same
                size as the number of colorbars used.

    Examples
    --------
    >>> images = ["density.jpg", "hi_density.jpg", "entropy.jpg",
    >>>           "special.jpg"]
    >>> cbs=[]
    >>> cbs.append(return_colormap("arbre", "Density [cm$^{-3}$]", (0,10), False))
    >>> cbs.append(return_colormap("kelp", "HI Density", (0,5), False))
    >>> cbs.append(return_colormap("hot", r"Entropy [K cm$^2$]", (1e-2,1e6), True))
    >>> cbs.append(return_colormap("Spectral", "Stuff$_x$!", (1,300), True))
    >>>
    >>> mp = multiplot(2,2, images=images, margins=(0.1,0.1),
    >>>                titles=["1","2","3","4"],
    >>>                xlabels=["one","two"], ylabels=None, colorbars=cbs,
    >>>                shrink_cb=0.95)
    >>> mp.scale_line(label="$r_{vir}$", labelloc="top")
    >>> mp.save_fig("multiplot")

    Notes
    -----
    If given both yt_plots and images, this will get preference to the
    yt plots.
    """
    # Error check
    npanels = ncol*nrow
    if(cb_labels is not None):
        cb_labels.reverse()   #Because I pop the list

    if images is not None:
        if len(images) != npanels:
            raise RuntimeError("Number of images (%d) doesn't match nrow(%d)"\
                               " x ncol(%d)." % (len(images), nrow, ncol))
            return
    if yt_plots is None and images is None:
        raise RuntimeError("Must supply either yt_plots or image filenames.")
        return
    if yt_plots is not None and images is not None:
        mylog.warning("Given both images and yt plots.  Ignoring images.")
    if yt_plots is not None:
        _yt = True
    else:
        _yt = False
    if fields is None:
        fields = [None] * npanels

    # If no ranges or labels given and given only images, fill them in.
    if not _yt:
        if xranges is None:
            xranges = []
            for i in range(npanels): xranges.append((0,1))
        if yranges is None:
            yranges = []
            for i in range(npanels): yranges.append((0,1))
        if xlabels is None:
            xlabels = []
            for i in range(npanels): xlabels.append("")
        if ylabels is None:
            ylabels = []
            for i in range(npanels): ylabels.append("")

    d = DualEPS(figsize=figsize)
    for j in range(nrow):
        invj = nrow - j - 1
        ypos = invj*(figsize[1] + margins[1])
        for i in range(ncol):
            xpos = i*(figsize[0] + margins[0])
            index = j*ncol + i
            if isinstance(yt_plots, list):
                this_plot = yt_plots[index]
            else:
                this_plot = yt_plots
            if j == nrow-1:
                xaxis = 0
            elif j == 0:
                xaxis = 1
            else:
                xaxis = -1
            if i == 0:
                yaxis = 0
            elif i == ncol-1:
                yaxis = 1
            else:
                yaxis = -1
            if xdata is None:
                _xdata = None
            else:
                _xdata = xdata[index]
            if ydata is None:
                _ydata = None
            else:
                _ydata = ydata[index]
            if xaxis_flags is not None:
                if xaxis_flags[index] is not None:
                    xaxis = xaxis_flags[index]
            if yaxis_flags is not None:
                if yaxis_flags[index] is not None:
                    yaxis = yaxis_flags[index]
            if _yt:
                this_plot._setup_plots()
                if xlabels is not None:
                    xlabel = xlabels[i]
                else:
                    xlabel = None
                if ylabels is not None:
                    ylabel = ylabels[j]
                else:
                    ylabel = None
                d.insert_image_yt(this_plot, pos=(xpos, ypos),
                                  field=fields[index])
                d.axis_box_yt(this_plot, pos=(xpos, ypos),
                              bare_axes=bare_axes, xaxis_side=xaxis,
                              yaxis_side=yaxis,
                              xlabel=xlabel, ylabel=ylabel,
                              xdata=_xdata, ydata=_ydata)
            else:
                d.insert_image(images[index], pos=(xpos,ypos))
                d.axis_box(pos = (xpos, ypos),
                           xrange=xranges[index], yrange=yranges[index],
                           xlabel=xlabels[i], ylabel=ylabels[j],
                           bare_axes=bare_axes, xaxis_side=xaxis, yaxis_side=yaxis,
                           xdata=_xdata, ydata=_ydata)
            if titles is not None:
                if titles[index] is not None:
                    d.title_box(titles[index],
                                loc=(i+0.05+i*margins[0]/figsize[0],
                                     j+0.98+j*margins[1]/figsize[1]))

    # Insert colorbars after all axes are placed because we want to
    # put them on the edges of the bounding box.
    bbox = (100.0 * d.canvas.bbox().left().t,
            100.0 * d.canvas.bbox().right().t - d.figsize[0],
            100.0 * d.canvas.bbox().bottom().t,
            100.0 * d.canvas.bbox().top().t - d.figsize[1])
    for j in range(nrow):
        invj = nrow - j - 1
        ypos0 = invj*(figsize[1] + margins[1])
        for i in range(ncol):
            xpos0 = i*(figsize[0] + margins[0])
            index = j*ncol + i
            if isinstance(yt_plots, list):
                this_plot = yt_plots[index]
            else:
                this_plot = yt_plots
            if (not _yt and colorbars is not None) or (_yt and not yt_nocbar):
                if cb_flags is not None:
                    if not cb_flags[index]:
                        continue
                if cb_location is None:
                    if ncol == 1:
                        orientation = "right"
                    elif i == 0:
                        orientation = "left"
                    elif i+1 == ncol:
                        orientation = "right"
                    elif j == 0:
                        orientation = "bottom"
                    elif j+1 == nrow:
                        orientation = "top"
                    else:
                        orientation = None  # Marker for interior plot
                else:
                    if isinstance(cb_location, dict):
                        if fields[index] not in cb_location.keys():
                            raise RuntimeError("%s not found in cb_location dict" %
                                               fields[index])
                            return
                        orientation = cb_location[fields[index]]
                    elif isinstance(cb_location, list):
                        orientation = cb_location[index]
                    else:
                        raise RuntimeError("Bad format: cb_location")
                if orientation == "right":
                    xpos = bbox[1]
                    ypos = ypos0
                elif orientation == "left":
                    xpos = bbox[0]
                    ypos = ypos0
                elif orientation == "bottom":
                    ypos = bbox[2]
                    xpos = xpos0
                elif orientation == "top":
                    ypos = bbox[3]
                    xpos = xpos0
                else:
                    mylog.warning("Unknown colorbar location %s. "
                                  "No colorbar displayed." % orientation)
                    orientation = None  # Marker for interior plot

                if orientation is not None:
                    if _yt:
                        # Set field if undefined
                        if fields[index] is None:
                            fields[index] = d.return_field(yt_plots[index])

                        d.colorbar_yt(this_plot,
                                      field=fields[index],
                                      pos=[xpos,ypos],
                                      shrink=shrink_cb,
                                      orientation=orientation,
                                      cb_labels=cb_labels)
                    else:
                        d.colorbar(colorbars[index]["cmap"],
                                   zrange=colorbars[index]["range"],
                                   label=colorbars[index]["name"],
                                   log=colorbars[index]["log"],
                                   orientation=orientation,
                                   pos=[xpos,ypos],
                                   shrink=shrink_cb)

    if savefig is not None:
        d.save_fig(savefig, format=format)

    return d

#=============================================================================

def multiplot_yt(ncol, nrow, plots, fields=None, **kwargs):
    r"""Wrapper for multiplot that takes a yt PlotWindow

    Accepts all parameters used in multiplot.

    Parameters
    ----------
    ncol : integer
        Number of columns in the figure.
    nrow : integer
        Number of rows in the figure.
    plots : ``PlotWindow`` instance, ``PhasePlot`` instance, or list of plots
        yt plots to be used.

    Examples
    --------
    >>> p1 = SlicePlot(ds, 0, 'density')
    >>> p1.set_width(10, 'kpc')
    >>>
    >>> p2 = SlicePlot(ds, 0, 'temperature')
    >>> p2.set_width(10, 'kpc')
    >>> p2.set_colormap('temperature', 'hot')
    >>>
    >>> sph = ds.sphere(ds.domain_center, (10, 'kpc'))
    >>> p3 = PhasePlot(sph, 'radius', 'density', 'temperature',
    ...                weight_field='cell_mass')
    >>>
    >>> p4 = PhasePlot(sph, 'radius', 'density', 'pressure', 'cell_mass')
    >>>
    >>> mp = multiplot_yt(2, 2, [p1, p2, p3, p4], savefig="yt", shrink_cb=0.9,
    ...                   bare_axes=True, yt_nocbar=False, margins=(0.5,0.5))
    """
    # Determine whether the plots are organized in a PlotWindow, or list
    # of PlotWindows
    if isinstance(plots, (PlotWindow, PhasePlot)):
        if fields is None:
            fields = plots.fields
        if len(fields) < nrow*ncol:
            raise RuntimeError("Number of plots ({0}) is less "
                               "than nrow({1}) x ncol({2}).".format(
                               len(fields), nrow, ncol))
            return
        figure = multiplot(ncol, nrow, yt_plots=plots, fields=fields, **kwargs)
    elif isinstance(plots, list) and isinstance(plots[0], (PlotWindow, PhasePlot)):
        if len(plots) < nrow*ncol:
            raise RuntimeError("Number of plots ({0}) is less "
                               "than nrow({1}) x ncol({2}).".format(
                               len(fields), nrow, ncol))
            return
        figure = multiplot(ncol, nrow, yt_plots=plots, fields=fields, **kwargs)
    else:
        raise RuntimeError("Unknown plot type in multiplot_yt")
        return
    return figure

#=============================================================================

def single_plot(plot, field=None, figsize=(12,12), cb_orient="right",
                bare_axes=False, savefig=None, colorbar=True,
                file_format='eps', **kwargs):
    r"""Wrapper for DualEPS routines to create a figure directly from a yt
    plot.  Calls insert_image_yt, axis_box_yt, and colorbar_yt.

    Parameters
    ----------
    plot : `yt.visualization.plot_window.PlotWindow`
        yt plot that provides the image and metadata
    figsize : tuple of floats
        Size of the figure in centimeters.
    cb_orient : string
        Placement of the colorbar.  Can be "left", "right", "top", or
        "bottom".
    bare_axes : boolean
        Set to true to have no annotations or tick marks on all of the axes.
    savefig : string
        Name of the saved file without the extension.
    colorbar : boolean
        Set to true to include a colorbar
    file_format : string
        Format type.  Can be "eps" or "pdf"

    Examples
    --------
    >>> p = SlicePlot(ds, 0, 'density')
    >>> p.set_width(0.1,'kpc')
    >>> single_plot(p, savefig="figure1")
    """
    d = DualEPS(figsize=figsize)
    d.insert_image_yt(plot, field=field)
    d.axis_box_yt(plot, bare_axes=bare_axes, **kwargs)
    if colorbar and not isinstance(plot, ProfilePlot):
        d.colorbar_yt(plot, orientation=cb_orient)
    if savefig is not None:
        d.save_fig(savefig, format=file_format)
    return d

#=============================================================================
def return_cmap(cmap=None, label="", range=(0,1), log=False):
    issue_deprecation_warning("Deprecated alias. Use return_colormap instead.")
    return return_colormap(cmap=cmap, label=label, crange=range, log=log)

def return_colormap(cmap=None, label="", range=(0,1), log=False):
    r"""Returns a dict that describes a colorbar.  Exclusively for use with
    multiplot.

    Parameters
    ----------
    cmap : string
        name of the (matplotlib) colormap to use
    label : string
        colorbar label
    range : tuple of floats
        min and max of the colorbar's range
    log : boolean
        Flag to use a logarithmic scale

    Examples
    --------
    >>> cb = return_colormap("arbre", "Density [cm$^{-3}$]", (0,10), False)
    """
    if cmap is None:
        cmap = ytcfg.get("yt", "default_colormap")
    return {'cmap': cmap, 'name': label, 'range': range, 'log': log}

max_stars_repo_name: NeilZaim/yt

max_issues_repo_name: NeilZaim/yt

max_forks_repo_name: NeilZaim/yt


text:
#!/usr/bin/python3

import numpy as np

from .non_uniform_bspline import NonUniformBSpline

class CentralizedCardinalBSpline(NonUniformBSpline):
  def __init__(self, p, nu=0):
    super().__init__(p, np.linspace(-(p+1)/2, (p+1)/2, p+2), nu=nu)

ext: py

max_stars_repo_name: valentjn/thesis

max_issues_repo_name: valentjn/thesis

max_forks_repo_name: valentjn/thesis

reason: import numpy


text:
import os
from functools import partial
from typing import Callable, Dict, List, Optional, Tuple, Union

import attr
import numpy as np
import pandas as pd
import tabmat as tm
from dask_ml.preprocessing import DummyEncoder
from git_root import git_root
from joblib import Memory
from scipy.sparse import csc_matrix

from .data import (
    generate_housing_dataset,
    generate_intermediate_insurance_dataset,
    generate_narrow_insurance_dataset,
    generate_real_insurance_dataset,
    generate_wide_insurance_dataset,
)
from .util import cache_location, exposure_and_offset_to_weights, get_tweedie_p

joblib_memory = Memory(cache_location, verbose=0)

@attr.s
class Problem:
    """Store metadata about which problem we should run."""

    data_loader = attr.ib(type=Callable)
    distribution = attr.ib(type=str)
    regularization_strength = attr.ib(type=float)
    l1_ratio = attr.ib(type=float)

@joblib_memory.cache
def load_data(
    loader_func: Callable[
        [Optional[int], Optional[float], Optional[str]],
        Tuple[pd.DataFrame, np.ndarray, np.ndarray],
    ],
    num_rows: int = None,
    storage: str = "dense",
    single_precision: bool = False,
    noise: float = None,
    distribution: str = "poisson",
    data_setup: str = "weights",
) -> Dict[str, np.ndarray]:
    """
    Load the data.

    A note about weights and exposures: Due to the way we have set up this problem, by
    rescaling the target variable, it is appropriate to pass what is modeled as an
    'exposure' as a weight. Everywhere else, exposures will be referred to as weights.
    """
    # TODO: add a weights_and_offset option
    # Step 1) Load the data.
    if data_setup not in ["weights", "offset", "no-weights"]:
        raise NotImplementedError
    X_in, y, exposure = loader_func(num_rows, noise, distribution)

    # Step 2) Convert to needed precision level.
    if single_precision:
        X_in = X_in.astype(np.float32)
        y = y.astype(np.float32)
        if exposure is not None:
            exposure = exposure.astype(np.float32)

    # Step 3) One hot encode columns if we are not using CategoricalMatrix
    def transform_col(i: int, dtype) -> Union[pd.DataFrame, tm.CategoricalMatrix]:
        if dtype.name == "category":
            if storage == "cat":
                return tm.CategoricalMatrix(X_in.iloc[:, i])
            return DummyEncoder().fit_transform(X_in.iloc[:, [i]])
        return X_in.iloc[:, [i]]

    mat_parts = [transform_col(i, dtype) for i, dtype in enumerate(X_in.dtypes)]
    # TODO: add a threshold for the number of categories needed to make a categorical
    #  matrix

    # Step 4) Convert the matrix to the appopriate storage type.
    if storage == "auto":
        dtype = np.float32 if single_precision else np.float64
        X = tm.from_pandas(X_in, dtype, sparse_threshold=0.1, cat_threshold=3)
    elif storage == "cat":
        cat_indices_in_expanded_arr: List[np.ndarray] = []
        dense_indices_in_expanded_arr: List[int] = []
        i = 0
        for elt in mat_parts:
            assert elt.ndim == 2
            if isinstance(elt, tm.CategoricalMatrix):
                ncol = elt.shape[1]
                cat_indices_in_expanded_arr.append(np.arange(i, i + ncol))
                i += ncol
            else:
                dense_indices_in_expanded_arr.append(i)
                i += 1

        non_cat_part = tm.DenseMatrix(
            np.hstack(
                [
                    elt.values
                    for elt in mat_parts
                    if not isinstance(elt, tm.CategoricalMatrix)
                ]
            )
        )
        X = tm.SplitMatrix(
            matrices=[non_cat_part]
            + [elt for elt in mat_parts if isinstance(elt, tm.CategoricalMatrix)],
            indices=[np.array(dense_indices_in_expanded_arr)]
            + cat_indices_in_expanded_arr,
        )
    elif storage == "sparse":
        X = csc_matrix(pd.concat(mat_parts, axis=1))
    elif storage.startswith("split"):
        threshold = float(storage.split("split")[1])
        X = tm.from_csc(csc_matrix(pd.concat(mat_parts, axis=1)), threshold)
    else:  # Fall back to using a dense matrix.
        X = pd.concat(mat_parts, axis=1)

    # Step 5) Handle weights or offsets if needed.
    if data_setup == "weights":
        # The exposure correction doesn't make sense for these distributions since
        # they don't use a log link (plus binomial isn't in the tweedie family),
        # but this is what we were doing before.
        if distribution in ["gaussian", "binomial"]:
            return dict(X=X, y=y, sample_weight=exposure)
        # when poisson, should be y=y, sample_weight=exposure
        # instead have y = y / exposure, weight = exposure
        y, sample_weight = exposure_and_offset_to_weights(
            get_tweedie_p(distribution), y, exposure
        )
        return dict(X=X, y=y * exposure, sample_weight=sample_weight)

    if data_setup == "offset":
        log_exposure = np.log(exposure)
        assert np.all(np.isfinite(log_exposure))
        # y has already been divided by exposure loader_func, so undo it here
        return dict(X=X, y=y * exposure, offset=log_exposure)
    # data_setup = "no_weights"
    return dict(X=X, y=y)

def get_all_problems() -> Dict[str, Problem]:
    """
    Return the name of all possible problems.

    Returns
    -------
    Dict mapping problem names to Problem instances.

    """
    regularization_strength = 0.001

    housing_distributions = ["gaussian", "gamma", "binomial"]
    housing_load_funcs = {
        "intermediate-housing": generate_housing_dataset,
    }

    insurance_distributions = [
        "gaussian",
        "poisson",
        "gamma",
        "tweedie-p=1.5",
        "binomial",
    ]
    insurance_load_funcs = {
        "intermediate-insurance": generate_intermediate_insurance_dataset,
        "narrow-insurance": generate_narrow_insurance_dataset,
        "wide-insurance": generate_wide_insurance_dataset,
    }
    if os.path.isfile(git_root("data", "X.parquet")):
        insurance_load_funcs["real-insurance"] = generate_real_insurance_dataset

    problems = {}
    for penalty_str, l1_ratio in [("l2", 0.0), ("net", 0.5), ("lasso", 1.0)]:
        # Add housing problems
        for distribution in housing_distributions:
            suffix = penalty_str + "-" + distribution
            dist = distribution
            for problem_name, load_fn in housing_load_funcs.items():
                for data_setup in ["no-weights", "offset"]:
                    problems["-".join((problem_name, data_setup, suffix))] = Problem(
                        data_loader=partial(
                            load_data, load_fn, distribution=dist, data_setup=data_setup
                        ),
                        distribution=distribution,
                        regularization_strength=regularization_strength,
                        l1_ratio=l1_ratio,
                    )
        # Add insurance problems
        for distribution in insurance_distributions:
            suffix = penalty_str + "-" + distribution
            dist = distribution
            for problem_name, load_fn in insurance_load_funcs.items():
                for data_setup in ["weights", "no-weights", "offset"]:
                    problems["-".join((problem_name, data_setup, suffix))] = Problem(
                        data_loader=partial(
                            load_data, load_fn, distribution=dist, data_setup=data_setup
                        ),
                        distribution=distribution,
                        regularization_strength=regularization_strength,
                        l1_ratio=l1_ratio,
                    )

    return problems

max_stars_repo_name: readthedocs-assistant/glum

max_issues_repo_name: readthedocs-assistant/glum

max_forks_repo_name: readthedocs-assistant/glum

reason: import numpy,from scipy


text:
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   distribution.py
@Contact :   Ancientshi@gmail.com
@Modify Time      @Author    @Version
------------      -------    --------
2020/12/9 5:18 下午   Ferdinand      1.0
@Desciption
----------------

----------------
'''

import numpy as np
def trainDistribution():
    '''
    查看训练数据中，各个relation的情况
    :return:
    '''
    fileNameList = ['train.txt', 'valid.txt', 'test.txt']
    context = ''
    dic = {'train.txt':{},
           'valid.txt':{},
           'test.txt':{}}
    for fileName in fileNameList:
        with open(context + fileName) as f:
            for line in f.readlines():
                array = line.strip().split('\t')
                relationName = array[1]
                if relationName not in dic[fileName].keys():
                    dic[fileName][relationName] = 0
                else:
                    dic[fileName][relationName] += 1
    # for i in dic.keys():
    #     print(i)
    #     for j in dic[i].keys():
    #         print('%s:%s'%(j,dic[i][j]))
    dic=np.array(dic)
    np.save('relationDistribution.npy',dic)
if __name__ == '__main__':
    trainDistribution()

max_stars_repo_name: Ancientshi/convKB-KerasVersion

max_issues_repo_name: Ancientshi/convKB-kerasVersion

max_forks_repo_name: Ancientshi/convKB-kerasVersion


text:
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo
import torch.nn.functional as F
import torch
import numpy as np
import cv2
import pdb
def sigmoid(x):
    return 1 / (1 + math.exp(-x))

def norm_angle(angle):
    norm_angle = sigmoid(10 * (abs(angle) / 0.7853975 - 1))
    return norm_angle

def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU()
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU()
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out = out + residual
        out = self.relu(out)

        return out

###''' self-attention; relation-attention '''

class ResNet_AT(nn.Module):
    def __init__(self, block, layers, num_classes=1000, end2end=True, at_type=''):
        self.inplanes = 64
        self.end2end = end2end
        super(ResNet_AT, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(0.5)
        self.dropout2 = nn.Dropout(0.6)
        self.alpha = nn.Sequential(nn.Linear(512, 1),
                                   nn.Sigmoid())

        self.beta = nn.Sequential(nn.Linear(1024, 1),
                                  nn.Sigmoid())

        self.pred_fc1 = nn.Linear(512, 7)
        self.pred_fc2 = nn.Linear(1024, 7)
        self.at_type = at_type

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x='', phrase='train', AT_level='first_level',vectors='',vm='',alphas_from1='',index_matrix=''):

        vs = []
        alphas = []

        assert phrase == 'train' or phrase == 'eval'
        assert AT_level == 'first_level' or AT_level == 'second_level' or AT_level == 'pred'
        if phrase == 'train':
            num_pair = 3

            for i in range(num_pair):
                f = x[:, :, :, :, i]  # x[128,3,224,224]

                f = self.conv1(f)
                f = self.bn1(f)
                f = self.relu(f)
                f = self.maxpool(f)

                f = self.layer1(f)
                f = self.layer2(f)
                f = self.layer3(f)
                f = self.layer4(f)
                f = self.avgpool(f)

                f = f.squeeze(3).squeeze(2)  # f[1, 512, 1, 1] ---> f[1, 512]

                # MN_MODEL(first Level)
                vs.append(f)
                alphas.append(self.alpha(self.dropout(f)))

            vs_stack = torch.stack(vs, dim=2)
            alphas_stack = torch.stack(alphas, dim=2)

            if self.at_type == 'self-attention':
                vm1 = vs_stack.mul(alphas_stack).sum(2).div(alphas_stack.sum(2))
            if self.at_type == 'self_relation-attention':
                vm1 = vs_stack.mul(alphas_stack).sum(2).div(alphas_stack.sum(2))
                betas = []
                for i in range(len(vs)):
                    vs[i] = torch.cat([vs[i], vm1], dim=1)
                    betas.append(self.beta(self.dropout(vs[i])))

                cascadeVs_stack = torch.stack(vs, dim=2)
                betas_stack = torch.stack(betas, dim=2)
                output = cascadeVs_stack.mul(betas_stack * alphas_stack).sum(2).div((betas_stack * alphas_stack).sum(2))

            if self.at_type == 'self-attention':
                vm1 = self.dropout(vm1)
                pred_score = self.pred_fc1(vm1)

            if self.at_type == 'self_relation-attention':
                output = self.dropout2(output)
                pred_score = self.pred_fc2(output)

            return pred_score

        if phrase == 'eval':
            if AT_level == 'first_level':
                f = self.conv1(x)
                f = self.bn1(f)
                f = self.relu(f)
                f = self.maxpool(f)

                f = self.layer1(f)
                f = self.layer2(f)
                f = self.layer3(f)
                f = self.layer4(f)
                f = self.avgpool(f)

                f = f.squeeze(3).squeeze(2)  # f[1, 512, 1, 1] ---> f[1, 512]
                # MN_MODEL(first Level)
                alphas = self.alpha(self.dropout(f))

                return f, alphas

            if AT_level == 'second_level':
                assert self.at_type == 'self_relation-attention'
                vms = index_matrix.permute(1, 0).mm(vm)  # [381, 21783] -> [21783,381] * [381,512] --> [21783, 512]
                vs_cate = torch.cat([vectors, vms], dim=1)

                betas = self.beta(self.dropout(vs_cate))
                ''' keywords: mean_fc ; weight_sourcefc; sum_alpha; weightmean_sourcefc '''
                ''' alpha * beta '''
                weight_catefc = vs_cate.mul(alphas_from1)  # [21570,512] * [21570,1] --->[21570,512]
                alpha_beta = alphas_from1.mul(betas)
                sum_alphabetas = index_matrix.mm(alpha_beta)  # [380,21570] * [21570,1] -> [380,1]
                weightmean_catefc = index_matrix.mm(weight_catefc).div(sum_alphabetas)

                weightmean_catefc = self.dropout2(weightmean_catefc)
                pred_score = self.pred_fc2(weightmean_catefc)

                return pred_score

            if AT_level == 'pred':
                if self.at_type == 'self-attention':
                    pred_score = self.pred_fc1(self.dropout(vm))

                return pred_score

''' self-attention; relation-attention '''
def resnet18_at(pretrained=False, **kwargs):
    # Constructs base a ResNet-18 model.
    model = ResNet_AT(BasicBlock, [2, 2, 2, 2], **kwargs)
    return model

max_stars_repo_name: J-asy/Emotion-FAN

max_issues_repo_name: J-asy/Emotion-FAN

max_forks_repo_name: J-asy/Emotion-FAN


text:
import math
import matplotlib.pyplot as plt
from numpy import *

	#Drop based Learning rate
def step_decay(epoch,initial_lrate,lrate_drop,lrate_epochs_drop):
	initial_lrate = initial_lrate
	drop = lrate_drop
	epochs_drop = lrate_epochs_drop
	lrate = initial_lrate * math.pow(drop, math.floor((1+epoch)/epochs_drop))
	return lrate

epochs = 20
initial_lrate = [0.1,0.06,0.01,0.001]
lrate_drop = 0.5
lrate_epochs_drop = 5
lrate = {}
lrate['0.01'] = []
lrate['0.1'] = []
lrate['0.001'] = []
lrate['0.06'] = []
for l in range(0,4):
	for epoch in range(0,epochs+1):
		lrate[str(initial_lrate[l])].append(step_decay(epoch,initial_lrate[l],lrate_drop,lrate_epochs_drop))

x_ticks = [1,4,8,12,16,20]
y_ticks = [0.0,0.0005,0.001,0.002,0.004,0.006,0.008,0.009,0.01]
x_tick_labels = ['1','4','8','12','16','20']
#y_tick_labels = ['0.000','0.01','0.02','0.04','0.06','0.08','0.1','0.12']
y_tick_labels = ['0.0','0.0005','0.001','0.002','0.004','0.006','0.008','0.009','0.01']
fig = plt.figure(figsize=(25,15))
#plt.plot(lrate['0.1'],'r',linewidth=3)
#plt.plot(lrate['0.06'],'g',linewidth=3)
plt.plot(lrate['0.01'],'b',linewidth=3)
#plt.plot(lrate['0.001'],linewidth=3)

plt.ylabel('Learning Rate', fontsize=35)
plt.xlabel('Epochs', fontsize=35)
plt.xticks(x_ticks, x_tick_labels, fontsize=35)
plt.yticks(y_ticks, y_tick_labels, fontsize=35)
fig.tight_layout()
plt.savefig('lr_scheduler.eps', format='eps', dpi=6000)
plt.show()

reason: from numpy


text:
"""
Implements various flows.
Each flow is invertible so it can be forward()ed and backward()ed.
Notice that backward() is not backward as in backprop but simply inversion.
Each flow also outputs its log det J "regularization"

Reference:

NICE: Non-linear Independent Components Estimation, Dinh et al. 2014
https://arxiv.org/abs/1410.8516

Variational Inference with Normalizing Flows, Rezende and Mohamed 2015
https://arxiv.org/abs/1505.05770

Density estimation using Real NVP, Dinh et al. May 2016
https://arxiv.org/abs/1605.08803
(Laurent's extension of NICE)

Improved Variational Inference with Inverse Autoregressive Flow, Kingma et al June 2016
https://arxiv.org/abs/1606.04934
(IAF)

Masked Autoregressive Flow for Density Estimation, Papamakarios et al. May 2017
https://arxiv.org/abs/1705.07057
"The advantage of Real NVP compared to MAF and IAF is that it can both generate data and estimate densities with one forward pass only, whereas MAF would need D passes to generate data and IAF would need D passes to estimate densities."
(MAF)

Glow: Generative Flow with Invertible 1x1 Convolutions, Kingma and Dhariwal, Jul 2018
https://arxiv.org/abs/1807.03039

"Normalizing Flows for Probabilistic Modeling and Inference"
https://arxiv.org/abs/1912.02762
(review paper)
"""

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

from spinflow.nets import LeafParam, MLP, ARMLP
from torch.distributions.utils import _standard_normal, broadcast_all
from torch.distributions.exp_family import ExponentialFamily
from numbers import Number
from torch.distributions import constraints
import itertools
from .spline_flows import NSF_CL
from .spline_embeddings import SplineNet

import math

class MyNormal(ExponentialFamily):

    arg_constraints = {'loc': constraints.real, 'logscale': constraints.real}
    support = constraints.real
    has_rsample = True
    _mean_carrier_measure = 0

    @property
    def mean(self):
        return self.loc

    @property
    def stddev(self):
        return self.scale

    @property
    def kl_normal(self):
        return (self.variance + self.mean ** 2) / 2 - self.logscale - 0.5

    @property
    def variance(self):
        return torch.exp(self.logscale * 2)

    def __init__(self, loc, logscale, validate_args=None):

        self.loc, self.logscale = broadcast_all(loc, logscale)

        if isinstance(loc, Number) and isinstance(logscale, Number):
            batch_shape = torch.Size()
        else:
            batch_shape = self.loc.size()
        super(MyNormal, self).__init__(batch_shape, validate_args=validate_args)

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(MyNormal, _instance)
        batch_shape = torch.Size(batch_shape)
        new.loc = self.loc.expand(batch_shape)
        new.logscale = self.logscale.expand(batch_shape)
        super(MyNormal, new).__init__(batch_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    @property
    def scale(self):
        return self.logscale.exp()

    def sample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        with torch.no_grad():
            return torch.normal(self.loc.expand(shape), self.scale.expand(shape))

    def rsample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        eps = _standard_normal(shape, dtype=self.loc.dtype, device=self.loc.device)
        return self.loc + eps * self.scale

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        return -((value - self.loc) ** 2) / (2 * self.variance) - self.logscale - math.log(math.sqrt(2 * math.pi))

    def cdf(self, value):
        if self._validate_args:
            self._validate_sample(value)
        return 0.5 * (1 + torch.erf((value - self.loc) * self.scale.reciprocal() / math.sqrt(2)))

    def icdf(self, value):
        if self._validate_args:
            self._validate_sample(value)
        return self.loc + self.scale * torch.erfinv(2 * value - 1) * math.sqrt(2)

    def entropy(self):
        raise NotImplementedError

    @property
    def _natural_params(self):
        return (self.loc / self.scale.pow(2), -0.5 * self.scale.pow(2).reciprocal())

    def _log_normalizer(self, x, y):
        return -0.25 * x.pow(2) / y + 0.5 * torch.log(-math.pi / y)

class AffineHalfEmbeddingFlow(nn.Module):
    """
    As seen in RealNVP, affine autoregressive flow (z = x * exp(s) + t), where half of the
    dimensions in x are linearly scaled/transfromed as a function of the other half.
    Which half is which is determined by the parity bit.
    - RealNVP both scales and shifts (default)
    - NICE only shifts
    """

    def __init__(self, dim, parity, scale=True, shift=True, emb=16, layer=64, delta=10, n_res=2):
        super().__init__()
        self.dim = dim
        self.parity = parity
        self.s_cond = lambda x: x.new_zeros(x.size(0), self.dim // 2)
        self.t_cond = lambda x: x.new_zeros(x.size(0), self.dim // 2)
        if scale:
            self.s_cond = SplineNet(self.dim // 2, emb=emb, layer=layer, delta=delta, n=self.dim // 2, n_res=n_res)
        if shift:
            self.t_cond = SplineNet(self.dim // 2, emb=emb, layer=layer, delta=delta, n=self.dim // 2, n_res=n_res)

    def forward(self, x, c=None):
        x0, x1 = x[:, ::2], x[:, 1::2]
        if self.parity:
            x0, x1 = x1, x0
        s = self.s_cond(x0)
        t = self.t_cond(x0)
        z0 = x0  # untouched half
        z1 = torch.exp(s) * x1 + t  # transform this half as a function of the other
        if self.parity:
            z0, z1 = z1, z0
        z = torch.cat([z0, z1], dim=1)
        log_det = torch.sum(s, dim=1)
        return z, log_det

    def backward(self, z, c=None):
        z0, z1 = z[:, ::2], z[:, 1::2]
        if self.parity:
            z0, z1 = z1, z0
        s = self.s_cond(z0)
        t = self.t_cond(z0)
        x0 = z0  # this was the same
        x1 = (z1 - t) * torch.exp(-s)  # reverse the transform on this half
        if self.parity:
            x0, x1 = x1, x0
        x = torch.cat([x0, x1], dim=1)
        log_det = torch.sum(-s, dim=1)
        return x, log_det

class AutoEncoderFlow(nn.Module):

    def __init__(self, nin, nout, nh, w=1e-3):
        super().__init__()
        self.enc = MLP(nin, nout, nh)
        self.dec = MLP(nout, nin, nh)
        self.w = w

    def forward(self, x, recursion=True):

        z = self.enc(x)

        if recursion:
            x_hat, _ = self.backward(z, recursion=False)
            mse = torch.square(x_hat - x.detach()).sum(dim=1) + self.w * torch.square(z).sum(dim=1)
        else:
            mse = torch.zeros(len(x), device=x.device)

        return z.detach(), -mse

    def backward(self, z, recursion=True):

        x = self.dec(z)

        if recursion:

            z_hat, _ = self.forward(x, recursion=False)
            mse = torch.square(z_hat - z.detach()).sum(dim=1)
        else:
            mse = torch.zeros(len(z), device=z.device)

        return x, -mse

class AffineConstantFlow(nn.Module):
    """
    Scales + Shifts the flow by (learned) constants per dimension.
    In NICE paper there is a Scaling layer which is a special case of this where t is None
    """
    def __init__(self, dim, scale=True, shift=True):
        super().__init__()
        self.s = nn.Parameter(torch.randn(1, dim, requires_grad=True)) if scale else None
        self.t = nn.Parameter(torch.randn(1, dim, requires_grad=True)) if shift else None

    def forward(self, x, c=None):
        s = self.s if self.s is not None else x.new_zeros(x.size())
        t = self.t if self.t is not None else x.new_zeros(x.size())
        z = x * torch.exp(s) + t
        log_det = torch.sum(s, dim=1)
        return z, log_det

    def backward(self, z, c=None):
        s = self.s if self.s is not None else z.new_zeros(z.size())
        t = self.t if self.t is not None else z.new_zeros(z.size())
        x = (z - t) * torch.exp(-s)
        log_det = torch.sum(-s, dim=1)
        return x, log_det

class SqrtNonLinearFlow(nn.Module):
    """
    Scales + Shifts the flow by (learned) constants per dimension.
    In NICE paper there is a Scaling layer which is a special case of this where t is None
    """

    def __init__(self):
        super().__init__()

    def forward(self, x, c=None):
        xh = torch.clamp_min(x, min=0)
        xl = torch.clamp_max(x, max=0)
        z = torch.sqrt(xh) - torch.sqrt(-xl)
        log_det = torch.sum(0.5 * torch.sqrt(1 / xh) + 0.5 * torch.sqrt(- 1 / xl), dim=1)
        return z, log_det

    def backward(self, z, c=None):
        zh = torch.clamp_min(z, min=0)
        zl = torch.clamp_max(z, max=0)
        x = torch.square(zh) - torch.square(zl)
        log_det = torch.sum(2 * zh - 2 * zl, dim=1)
        return x, log_det

class AffineConditionalFlow(nn.Module):
    """
    Scales + Shifts the flow by (learned) constants per dimension.
    In NICE paper there is a Scaling layer which is a special case of this where t is None
    """

    def __init__(self, dim, context_dim, nh=128, scale=True, shift=True):
        super().__init__()

        self.s = None
        self.t = None

        if scale:
            self.s_cond = MLP(context_dim, dim, nh)
        if shift:
            self.t_cond = MLP(context_dim, dim, nh)

    def forward(self, x, c=None):
        s = self.s(c) if self.s is not None else x.new_zeros(x.size())
        t = self.t(c) if self.t is not None else x.new_zeros(x.size())
        z = x * torch.exp(s) + t
        log_det = torch.sum(s, dim=1)
        return z, log_det

    def backward(self, z, c=None):
        s = self.s(c) if self.s is not None else z.new_zeros(z.size())
        t = self.t(c) if self.t is not None else z.new_zeros(z.size())
        x = (z - t) * torch.exp(-s)
        log_det = torch.sum(-s, dim=1)
        return x, log_det

class ActNorm(AffineConstantFlow):
    """
    Really an AffineConstantFlow but with a data-dependent initialization,
    where on the very first batch we clever initialize the s,t so that the output
    is unit gaussian. As described in Glow paper.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.data_dep_init_done = False
        self.eps = 1e-3

    def forward(self, x, c=None):
        # first batch is used for init
        if not self.data_dep_init_done:
            assert self.s is not None and self.t is not None # for now
            self.s.data = (-torch.log(torch.clamp_min(x.std(dim=0, keepdim=True), min=self.eps))).detach()
            self.t.data = (-(x * torch.exp(self.s)).mean(dim=0, keepdim=True)).detach()
            self.data_dep_init_done = True
        return super().forward(x)

class AffineHalfFlow(nn.Module):
    """
    As seen in RealNVP, affine autoregressive flow (z = x * exp(s) + t), where half of the
    dimensions in x are linearly scaled/transfromed as a function of the other half.
    Which half is which is determined by the parity bit.
    - RealNVP both scales and shifts (default)
    - NICE only shifts
    """
    def __init__(self, dim, parity, net_class=MLP, nh=24, scale=True, shift=True):
        super().__init__()
        self.dim = dim
        self.parity = parity
        self.s_cond = lambda x: x.new_zeros(x.size(0), self.dim // 2)
        self.t_cond = lambda x: x.new_zeros(x.size(0), self.dim // 2)
        if scale:
            self.s_cond = net_class(self.dim // 2, self.dim // 2, nh)
        if shift:
            self.t_cond = net_class(self.dim // 2, self.dim // 2, nh)

    def forward(self, x, c=None):
        x0, x1 = x[:,::2], x[:,1::2]
        if self.parity:
            x0, x1 = x1, x0
        s = self.s_cond(x0)
        t = self.t_cond(x0)
        z0 = x0 # untouched half
        z1 = torch.exp(s) * x1 + t # transform this half as a function of the other
        if self.parity:
            z0, z1 = z1, z0
        z = torch.cat([z0, z1], dim=1)
        log_det = torch.sum(s, dim=1)
        return z, log_det

    def backward(self, z, c=None):
        z0, z1 = z[:,::2], z[:,1::2]
        if self.parity:
            z0, z1 = z1, z0
        s = self.s_cond(z0)
        t = self.t_cond(z0)
        x0 = z0 # this was the same
        x1 = (z1 - t) * torch.exp(-s) # reverse the transform on this half
        if self.parity:
            x0, x1 = x1, x0
        x = torch.cat([x0, x1], dim=1)
        log_det = torch.sum(-s, dim=1)
        return x, log_det

class SlowMAF(nn.Module):
    """
    Masked Autoregressive Flow, slow version with explicit networks per dim
    """
    def __init__(self, dim, parity, net_class=MLP, nh=24):
        super().__init__()
        self.dim = dim
        self.layers = nn.ModuleDict()
        self.layers[str(0)] = LeafParam(2)
        for i in range(1, dim):
            self.layers[str(i)] = net_class(i, 2, nh)
        self.order = list(range(dim)) if parity else list(range(dim))[::-1]

    def forward(self, x, c=None):
        z = torch.zeros_like(x)
        log_det = torch.zeros(x.size(0))
        for i in range(self.dim):
            st = self.layers[str(i)](x[:, :i])
            s, t = st[:, 0], st[:, 1]
            z[:, self.order[i]] = x[:, i] * torch.exp(s) + t
            log_det += s
        return z, log_det

    def backward(self, z, c=None):
        x = torch.zeros_like(z)
        log_det = torch.zeros(z.size(0))
        for i in range(self.dim):
            st = self.layers[str(i)](x[:, :i])
            s, t = st[:, 0], st[:, 1]
            x[:, i] = (z[:, self.order[i]] - t) * torch.exp(-s)
            log_det += -s
        return x, log_det

class MAF(nn.Module):
    """ Masked Autoregressive Flow that uses a MADE-style network for fast forward """

    def __init__(self, dim, parity, net_class=ARMLP, nh=24):
        super().__init__()
        self.dim = dim
        self.net = net_class(dim, dim*2, nh)
        self.parity = parity

    def forward(self, x, c=None):
        # here we see that we are evaluating all of z in parallel, so density estimation will be fast
        st = self.net(x)
        s, t = st.split(self.dim, dim=1)
        z = x * torch.exp(s) + t
        # reverse order, so if we stack MAFs correct things happen
        z = z.flip(dims=(1,)) if self.parity else z
        log_det = torch.sum(s, dim=1)
        return z, log_det

    def backward(self, z, c=None):
        # we have to decode the x one at a time, sequentially
        x = torch.zeros_like(z)
        log_det = torch.zeros(z.size(0))
        z = z.flip(dims=(1,)) if self.parity else z
        for i in range(self.dim):
            st = self.net(x.clone()) # clone to avoid in-place op errors if using IAF
            s, t = st.split(self.dim, dim=1)
            x[:, i] = (z[:, i] - t[:, i]) * torch.exp(-s[:, i])
            log_det += -s[:, i]
        return x, log_det

class IAF(MAF):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        """
        reverse the flow, giving an Inverse Autoregressive Flow (IAF) instead,
        where sampling will be fast but density estimation slow
        """
        self.forward, self.backward = self.backward, self.forward

class Invertible1x1Conv(nn.Module):
    """
    As introduced in Glow paper.
    """

    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        Q = torch.nn.init.orthogonal_(torch.randn(dim, dim))
        P, L, U = torch.lu_unpack(*Q.lu())
        # remains fixed during optimization
        self.register_buffer('P', P)
        self.L = nn.Parameter(L)  # lower triangular portion
        self.S = nn.Parameter(U.diag())  # "crop out" the diagonal to its own parameter
        self.U = nn.Parameter(torch.triu(U, diagonal=1))  # "crop out" diagonal, stored in S

        self.W_inv = None

    def _assemble_W(self):
        """ assemble W from its pieces (P, L, U, S) """
        L = torch.tril(self.L, diagonal=-1) + torch.diag(torch.ones(self.dim, device=self.L.device))
        U = torch.triu(self.U, diagonal=1)
        W = self.P @ L @ (U + torch.diag(self.S))
        return W

    def forward(self, x, c=None):
        W = self._assemble_W()
        z = x @ W
        log_det = torch.sum(torch.log(torch.abs(self.S)))
        self.W_inv = None
        return z, log_det

    def backward(self, z, c=None):

        if self.W_inv is None:
            W = self._assemble_W()
            self.W_inv = torch.inverse(W)

        x = z @ self.W_inv
        log_det = -torch.sum(torch.log(torch.abs(self.S)))
        return x, log_det

# ------------------------------------------------------------------------

class NormalizingFlow(nn.Module):
    """ A sequence of Normalizing Flows is a Normalizing Flow """

    def __init__(self, flows):
        super().__init__()
        self.flows = nn.ModuleList(flows)

    def forward(self, x):
        m, _ = x.shape
        log_det = torch.zeros(m, device=x.device)
        zs = [x]
        for flow in self.flows:
            x, ld = flow.forward(x)
            log_det += ld
            zs.append(x)
        return zs, log_det

    def backward(self, z):
        m, _ = z.shape
        log_det = torch.zeros(m, device=z.device)
        xs = [z]
        for flow in self.flows[::-1]:
            z, ld = flow.backward(z)
            log_det += ld
            xs.append(z)
        return xs, log_det

class NormalizingFlowModel(nn.Module):
    """ A Normalizing Flow Model is a (prior, flow) pair """

    def __init__(self, prior, flows):
        super().__init__()
        self.prior = prior
        self.flow = NormalizingFlow(flows)

    def forward(self, x):
        zs, log_det = self.flow.forward(x)
        prior_logprob = self.prior.log_prob(zs[-1]).view(x.size(0), -1).sum(1)
        return zs, prior_logprob, log_det

    def backward(self, z):
        xs, log_det = self.flow.backward(z)
        return xs, log_det

    def sample(self, num_samples):
        z = self.prior.sample((num_samples,))
        xs, _ = self.flow.backward(z)
        return xs

class ConditionalNormalizingFlow(nn.Module):
    """ A sequence of Normalizing Flows is a Normalizing Flow """

    def __init__(self, flows):
        super().__init__()
        self.flows = nn.ModuleList(flows)

    def forward(self, x, c=None):
        m, _ = x.shape
        log_det = torch.zeros(m, device=x.device)
        zs = [x]
        for flow in self.flows:
            x, ld = flow.forward(x, c)
            log_det += ld
            zs.append(x)
        return zs, log_det

    def backward(self, z, c=None):
        m, _ = z.shape
        log_det = torch.zeros(m, device=z.device)
        xs = [z]
        for flow in self.flows[::-1]:
            z, ld = flow.backward(z, c)
            log_det += ld
            xs.append(z)
        return xs, log_det

class ConditionalNormalizingFlowModel(nn.Module):
    """ A Normalizing Flow Model is a (prior, flow) pair """

    def __init__(self, x_dim, c_dim, B=20, K=20, nh=128):
        super().__init__()

        # Define the primery flow

        flows = [NSF_CL(dim=x_dim, K=K, B=B, hidden_dim=16, base_network=MLP) for _ in range(1)]
        convs = [Invertible1x1Conv(dim=x_dim) for _ in flows]
        norms = [ActNorm(dim=x_dim) for _ in flows]
        flows = list(itertools.chain(*zip(norms, convs, flows)))
        self.flow_x = ConditionalNormalizingFlow(flows)

        # Define the secondary flow

        flows = [NSF_CL(dim=x_dim, K=K, B=B, hidden_dim=16, base_network=MLP) for _ in range(1)]
        convs = [Invertible1x1Conv(dim=x_dim) for _ in flows]
        norms = [ActNorm(dim=x_dim) for _ in flows]
        flows = list(itertools.chain(*zip(norms, convs, flows)))
        self.flow_c = ConditionalNormalizingFlow(flows)

        self.x_loc = MLP(c_dim, x_dim, nh, leak=0)
        self.x_logscale = MLP(c_dim, x_dim, nh, leak=0)

        self.register_buffer('prior_loc', torch.zeros(x_dim))
        self.register_buffer('prior_logscale', torch.zeros(x_dim))
        self.conditional_prior = None

    def forward(self, x, c):

        prior = MyNormal(loc=self.prior_loc, logscale=self.prior_logscale)

        loc = self.x_loc(c)
        logscale = self.x_logscale(c)

        params = {'loc': loc, 'logscale': logscale}
        conditional_prior = MyNormal(**params)

        kl = conditional_prior.kl_normal.sum(dim=1)

        self.conditional_prior = conditional_prior

        zs, log_det = self.flow_x.forward(x)
        z = zs[-1]

        prior_logprob = prior.log_prob(z).sum(1)

        zs, log_det_conditional = self.flow_c.forward(z)
        z = zs[-1]

        conditional_log_det = log_det + log_det_conditional
        conditional_logprob = conditional_prior.log_prob(z).sum(1)

        return prior_logprob, conditional_logprob, log_det, conditional_log_det, kl

    def freeze(self, fr=True):
        for param in self.parameters():
            param.requires_grad = not fr

    def sample(self, c=None, num_samples=1):

        if c is not None:
            loc = self.x_loc(c)
            logscale = self.x_logscale(c)

            params = {'loc': loc, 'logscale': logscale}
            conditional_prior = MyNormal(**params)

            z = conditional_prior.sample((num_samples,))
            _, b, _ = z.shape
            z = z.view(num_samples * b, -1)

            zs, _ = self.flow_c.backward(z)
            z = zs[-1]

        else:

            prior = MyNormal(loc=self.prior_loc, logscale=self.prior_logscale)
            z = prior.sample((num_samples,))

        xs, _ = self.flow_x.backward(z)
        x = xs[-1]

        if c is not None:
            x = x.view(num_samples, b, -1)

        return x.squeeze(0)

max_stars_repo_name: eladsar/spinningup

max_issues_repo_name: eladsar/spinningup

max_forks_repo_name: eladsar/spinningup