<commit_before>#!/usr/bin/env python

# Copyright (c) 2017, DIANA-HEP
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import numbers
import math

import numpy
try:
    import numba
    import llvmlite
except ImportError:
    numba = None

import uproot.rootio

class TH1Methods(object):
    # makes __doc__ attribute mutable before Python 3.3
    __metaclass__ = type.__new__(type, "type", (uproot.rootio.ROOTObject.__metaclass__,), {})

    def __repr__(self):
        if self.fName is None:
            return "<{0} at 0x{1:012x}>".format(self.classname, id(self))
        else:
            return "<{0} {1} 0x{2:012x}>".format(self.classname, repr(self.fName), id(self))

    @property
    def name(self):
        return self.fName

    @property
    def title(self):
        return self.fTitle

    @property
    def numbins(self):
        return self.fXaxis.fNbins

    @property
    def low(self):
        return self.fXaxis.fXmin

    @property
    def high(self):
        return self.fXaxis.fXmax

    @property
    def underflows(self):
        return self[0]

    @property
    def overflows(self):
        return self[-1]

    @property
    def values(self):
        return self[1:-1]

    @property
    def allvalues(self):
        return self[:]

    @property
    def numpy(self):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        norm = (high - low) / self.fXaxis.fNbins
        freq = numpy.array(self.values, dtype=self._dtype.newbyteorder("="))
        edges = numpy.array([i*norm + low for i in range(self.numbins + 1)])
        return freq, edges

    def interval(self, index):
        if index < 0:
            index += len(self)

        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if index == 0:
            return (float("-inf"), low)
        elif index == len(self) - 1:
            return (high, float("inf"))
        else:
            norm = (high - low) / self.fXaxis.fNbins
            return (index - 1)*norm + low, index*norm + low

    def index(self, data):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if data < low:
            return 0
        elif data >= high:
            return len(self) - 1
        else:
            return int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1

    def fill(self, data, weights=None):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax

        if isinstance(data, numbers.Real):
            if weights is None:
                weights = 1

            if data < low:
                self[0] += weight
            elif data >= high:
                self[-1] += weight
            else:
                self[int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1] += weights

        else:
            if not isinstance(data, numpy.ndarray):
                data = numpy.array(data)

            if isinstance(weights, numbers.Real):
                weights = numpy.empty_like(data)

            freq, edges = numpy.histogram(data,
                                          bins=numpy.linspace(low, high, self.fXaxis.fNbins + 1),
                                          weights=weights,
                                          density=False)
            for i, x in enumerate(freq):
                self[i + 1] += x

            underflows = (data < low)
            overflows = (data >= high)

            if isinstance(weights, numpy.ndarray):
                self[0] += weights[underflows].sum()
                self[-1] += weights[overflows].sum()
            else:
                self[0] += underflows.sum()
                self[-1] += overflows.sum()

    def show(self, width=80, minimum=None, maximum=None, stream=sys.stdout):
        if minimum is None:
            minimum = min(self)
            if minimum < 0:
                minimum *= 1.05
            else:
                minimum = 0

        if maximum is None:
            maximum = max(self) * 1.05

        if maximum <= minimum:
            average = (minimum + maximum) / 2.0
            minimum = average - 0.5
            maximum = average + 0.5

        intervals = ["[{0:<.4g}, {1:<.4g})".format(l, h) for l, h in [self.interval(i) for i in range(len(self))]]
        intervals[-1] = intervals[-1][:-1] + "]"   # last interval is closed on top edge
        intervalswidth = max(len(x) for x in intervals)

        values = ["{0:<.4g}".format(float(x)) for x in self]
        valueswidth = max(len(x) for x in values)

        minimumtext = "{0:<.4g}".format(minimum)
        maximumtext = "{0:<.4g}".format(maximum)

        plotwidth = max(len(minimumtext) + len(maximumtext), width - (intervalswidth + 1 + valueswidth + 1 + 2))
        scale = minimumtext + " "*(plotwidth + 2 - len(minimumtext) - len(maximumtext)) + maximumtext

        norm = float(plotwidth) / float(maximum - minimum)
        zero = int(round((0.0 - minimum)*norm))
        line = numpy.empty(plotwidth, dtype=numpy.uint8)

        formatter = "{0:<%s} {1:<%s} |{2}|" % (intervalswidth, valueswidth)
        line[:] = ord("-")
        if minimum != 0 and 0 <= zero < plotwidth:
            line[zero] = ord("+")
        capstone = " " * (intervalswidth + 1 + valueswidth + 1) + "+" + str(line.tostring().decode("ascii")) + "+"

        out = [" "*(intervalswidth + valueswidth + 2) + scale]
        out.append(capstone)
        for interval, value, x in zip(intervals, values, self):
            line[:] = ord(" ")

            pos = int(round((x - minimum)*norm))
            if x < 0:
                line[pos:zero] = ord("*")
            else:
                line[zero:pos] = ord("*")

            if minimum != 0 and 0 <= zero < plotwidth:
                line[zero] = ord("|")

            out.append(formatter.format(interval, value, str(line.tostring().decode("ascii"))))

        out.append(capstone)
        out = "\n".join(out)
        if stream is None:
            return out
        else:
            stream.write(out)
            stream.write("\n")

    @property
    def holoviews(self):
        import uproot._connect.to_holoviews
        return uproot._connect.to_holoviews.TH1Methods_holoviews(self)

    @property
    def bokeh(self):
        import uproot._connect.to_bokeh
        return uproot._connect.to_bokeh.TH1Methods_bokeh(self)

uproot.rootio.methods["TH1"] = TH1Methods

class TH1(TH1Methods, list):
    def _type(self):
        if all(isinstance(x, numbers.Integral) for x in self):
            return int
        elif all(isinstance(x, numbers.Real) for x in self):
            return float
        else:
            raise TypeError("histogram bin values must be integers or floats")

    @property
    def classname(self):
        if self._type() is int:
            return "TH1I"
        else:
            return "TH1D"

    @property
    def _dtype(self):
        if self._type() is int:
            return numpy.dtype(">i4")
        else:
            return numpy.dtype(">f8")

class TAxis(object):
    classname = "TAxis"

def hist(numbins, low, high, name=None, title=None):
    out = TH1()
    out.fXaxis = TAxis()
    out.fXaxis.fNbins = int(numbins)
    out.fXaxis.fXmin = float(low)
    out.fXaxis.fXmax = float(high)
    out.fName = name
    out.fTitle = title
    out.extend([0] * (numbins + 2))
    return out
<commit_msg>starting to Numbafy it<commit_after>#!/usr/bin/env python

# Copyright (c) 2017, DIANA-HEP
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import numbers
import math

import numpy
try:
    import numba
    import llvmlite
except ImportError:
    numba = None

import uproot.rootio

class TH1Methods(object):
    # makes __doc__ attribute mutable before Python 3.3
    __metaclass__ = type.__new__(type, "type", (uproot.rootio.ROOTObject.__metaclass__,), {})

    def __repr__(self):
        if self.fName is None:
            return "<{0} at 0x{1:012x}>".format(self.classname, id(self))
        else:
            return "<{0} {1} 0x{2:012x}>".format(self.classname, repr(self.fName), id(self))

    @property
    def name(self):
        return self.fName

    @property
    def title(self):
        return self.fTitle

    @property
    def numbins(self):
        return self.fXaxis.fNbins

    @property
    def low(self):
        return self.fXaxis.fXmin

    @property
    def high(self):
        return self.fXaxis.fXmax

    @property
    def underflows(self):
        return self[0]

    @property
    def overflows(self):
        return self[-1]

    @property
    def values(self):
        return self[1:-1]

    @property
    def allvalues(self):
        return self[:]

    @property
    def numpy(self):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        norm = (high - low) / self.fXaxis.fNbins
        freq = numpy.array(self.values, dtype=self._dtype.newbyteorder("="))
        edges = numpy.array([i*norm + low for i in range(self.numbins + 1)])
        return freq, edges

    # def _numba(self):
    #     return self, int(self.fXaxis.fNbins), float(self.fXaxis.fXmin), float(self.fXaxis.fXmax), numpy.array(self, dtype=numpy.float64)

    def interval(self, index):
        if index < 0:
            index += len(self)

        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if index == 0:
            return (float("-inf"), low)
        elif index == len(self) - 1:
            return (high, float("inf"))
        else:
            norm = (high - low) / self.fXaxis.fNbins
            return (index - 1)*norm + low, index*norm + low

    def index(self, data):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if data < low:
            return 0
        elif data >= high:
            return len(self) - 1
        else:
            return int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1

    def fill(self, data, weights=None):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax

        if isinstance(data, numbers.Real):
            if weights is None:
                weights = 1

            if data < low:
                self[0] += weight
            elif data >= high:
                self[-1] += weight
            else:
                self[int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1] += weights

        else:
            if not isinstance(data, numpy.ndarray):
                data = numpy.array(data)

            if isinstance(weights, numbers.Real):
                weights = numpy.empty_like(data)

            freq, edges = numpy.histogram(data,
                                          bins=numpy.linspace(low, high, self.fXaxis.fNbins + 1),
                                          weights=weights,
                                          density=False)
            for i, x in enumerate(freq):
                self[i + 1] += x

            underflows = (data < low)
            overflows = (data >= high)

            if isinstance(weights, numpy.ndarray):
                self[0] += weights[underflows].sum()
                self[-1] += weights[overflows].sum()
            else:
                self[0] += underflows.sum()
                self[-1] += overflows.sum()

    def show(self, width=80, minimum=None, maximum=None, stream=sys.stdout):
        if minimum is None:
            minimum = min(self)
            if minimum < 0:
                minimum *= 1.05
            else:
                minimum = 0

        if maximum is None:
            maximum = max(self) * 1.05

        if maximum <= minimum:
            average = (minimum + maximum) / 2.0
            minimum = average - 0.5
            maximum = average + 0.5

        intervals = ["[{0:<.4g}, {1:<.4g})".format(l, h) for l, h in [self.interval(i) for i in range(len(self))]]
        intervals[-1] = intervals[-1][:-1] + "]"   # last interval is closed on top edge
        intervalswidth = max(len(x) for x in intervals)

        values = ["{0:<.4g}".format(float(x)) for x in self]
        valueswidth = max(len(x) for x in values)

        minimumtext = "{0:<.4g}".format(minimum)
        maximumtext = "{0:<.4g}".format(maximum)

        plotwidth = max(len(minimumtext) + len(maximumtext), width - (intervalswidth + 1 + valueswidth + 1 + 2))
        scale = minimumtext + " "*(plotwidth + 2 - len(minimumtext) - len(maximumtext)) + maximumtext

        norm = float(plotwidth) / float(maximum - minimum)
        zero = int(round((0.0 - minimum)*norm))
        line = numpy.empty(plotwidth, dtype=numpy.uint8)

        formatter = "{0:<%s} {1:<%s} |{2}|" % (intervalswidth, valueswidth)
        line[:] = ord("-")
        if minimum != 0 and 0 <= zero < plotwidth:
            line[zero] = ord("+")
        capstone = " " * (intervalswidth + 1 + valueswidth + 1) + "+" + str(line.tostring().decode("ascii")) + "+"

        out = [" "*(intervalswidth + valueswidth + 2) + scale]
        out.append(capstone)
        for interval, value, x in zip(intervals, values, self):
            line[:] = ord(" ")

            pos = int(round((x - minimum)*norm))
            if x < 0:
                line[pos:zero] = ord("*")
            else:
                line[zero:pos] = ord("*")

            if minimum != 0 and 0 <= zero < plotwidth:
                line[zero] = ord("|")

            out.append(formatter.format(interval, value, str(line.tostring().decode("ascii"))))

        out.append(capstone)
        out = "\n".join(out)
        if stream is None:
            return out
        else:
            stream.write(out)
            stream.write("\n")

    @property
    def holoviews(self):
        import uproot._connect.to_holoviews
        return uproot._connect.to_holoviews.TH1Methods_holoviews(self)

    @property
    def bokeh(self):
        import uproot._connect.to_bokeh
        return uproot._connect.to_bokeh.TH1Methods_bokeh(self)

uproot.rootio.methods["TH1"] = TH1Methods

class TH1(TH1Methods, list):
    def _type(self):
        if all(isinstance(x, numbers.Integral) for x in self):
            return int
        elif all(isinstance(x, numbers.Real) for x in self):
            return float
        else:
            raise TypeError("histogram bin values must be integers or floats")

    @property
    def classname(self):
        if self._type() is int:
            return "TH1I"
        else:
            return "TH1D"

    @property
    def _dtype(self):
        if self._type() is int:
            return numpy.dtype(">i4")
        else:
            return numpy.dtype(">f8")

class TAxis(object):
    classname = "TAxis"

def hist(numbins, low, high, name=None, title=None):
    out = TH1()
    out.fXaxis = TAxis()
    out.fXaxis.fNbins = int(numbins)
    out.fXaxis.fXmin = float(low)
    out.fXaxis.fXmax = float(high)
    out.fName = name
    out.fTitle = title
    out.extend([0] * (numbins + 2))
    return out

if numba is not None:
    class Regular1dType(numba.types.Type):
        def __init__(self):
            super(Regular1dType, self).__init__(name="Regular1dType")

    regular1dType = Regular1dType()

    @numba.extending.typeof_impl.register(TH1Methods)
    def th1_typeof(val, c):
        assert isinstance(val, TH1Methods)
        return regular1dType

    @numba.extending.register_model(Regular1dType)
    class Regular1dModel(numba.datamodel.models.TupleModel):
        def __init__(self, dmm, fe_type):
            super(Regular1dModel, self).__init__(dmm, numba.types.Tuple(()))   # numba.types.Opaque, numba.types.int64, numba.types.float64, numba.types.float64, numba.types.float64[:])))

    # @numba.extending.infer_getattr
    # class StructAttribute(numba.typing.templates.AttributeTemplate):
    #     key = Regular1dType
    #     def generic_resolve(self, typ, attr):
    #         if attr == "numbins":
    #             return numba.types.int64
    #         elif attr == "low":
    #             return numba.types.float64
    #         elif attr == "high":
    #             return numba.types.float64
    #         elif attr == "data":
    #             return numba.types.float64[:]

    # @numba.extending.lower_getattr(Regular1dType, "numbins")
    # def th1_getattr_numbins_impl(context, builder, typ, val):
    #     res = builder.extract_value(val, 1)
    #     return numba.targets.imputils.impl_ret_borrowed(context, builder, typ.contents, res)

    # @numba.extending.lower_getattr(Regular1dType, "low")
    # def th1_getattr_low_impl(context, builder, typ, val):
    #     res = builder.extract_value(val, 2)
    #     return numba.targets.imputils.impl_ret_borrowed(context, builder, typ.contents, res)

    # @numba.extending.lower_getattr(Regular1dType, "high")
    # def th1_getattr_high_impl(context, builder, typ, val):
    #     res = builder.extract_value(val, 3)
    #     return numba.targets.imputils.impl_ret_borrowed(context, builder, typ.contents, res)

    # @numba.extending.lower_getattr(Regular1dType, "data")
    # def th1_getattr_data_impl(context, builder, typ, val):
    #     res = builder.extract_value(val, 4)
    #     return numba.targets.imputils.impl_ret_borrowed(context, builder, typ.contents, res)

    @numba.extending.unbox(Regular1dType)
    def th1_unbox(typ, obj, c):
        tuple_obj = c.pyapi.tuple_new(0)
        # c.pyapi.tuple_setitem(tuple_obj, 0, obj)
        # out = c.unbox(numba.types.Tuple((numba.types.Opaque,)), tuple_obj)
        out = c.unbox(numba.types.Tuple(()), tuple_obj)
        c.pyapi.decref(tuple_obj)
        return out

    # @numba.extending.box(Regular1dType)
    # def th1_box(typ, val, c):
    #     return c.box(numba.types.Opaque, c.builder.extract_value(val, 0))

@numba.njit
def testy(x):
    return 3.14

h = hist(10, -3.0, 3.0)
print h
print testy(h)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8
''' This is part of the MSS Python's module.
    Source: https://github.com/BoboTiG/python-mss
'''

from ctypes import (
    POINTER, Structure, byref, c_char_p, c_int, c_int32, c_long, c_uint,
    c_uint32, c_ulong, c_ushort, c_void_p, cast, cdll, create_string_buffer)
from ctypes.util import find_library
from os import environ
from os.path import abspath, dirname, isfile, realpath
from struct import pack
from sys import maxsize, version

from .base import MSSBase
from .exception import ScreenshotError

__all__ = ['MSS']

class Display(Structure):
    ''' Structure that serves as the connection to the X server
        and that contains all the information about that X server.
    '''

class XWindowAttributes(Structure):
    ''' Attributes for the specified window. '''

    _fields_ = [('x', c_int32), ('y', c_int32), ('width', c_int32),
                ('height', c_int32), ('border_width', c_int32),
                ('depth', c_int32), ('visual', c_ulong), ('root', c_ulong),
                ('class', c_int32), ('bit_gravity', c_int32),
                ('win_gravity', c_int32), ('backing_store', c_int32),
                ('backing_planes', c_ulong), ('backing_pixel', c_ulong),
                ('save_under', c_int32), ('colourmap', c_ulong),
                ('mapinstalled', c_uint32), ('map_state', c_uint32),
                ('all_event_masks', c_ulong), ('your_event_mask', c_ulong),
                ('do_not_propagate_mask', c_ulong),
                ('override_redirect', c_int32), ('screen', c_ulong)]

class XImage(Structure):
    ''' Description of an image as it exists in the client's memory.
        https://tronche.com/gui/x/xlib/graphics/images.html
    '''

    _fields_ = [('width', c_int), ('height', c_int), ('xoffset', c_int),
                ('format', c_int), ('data', c_void_p),
                ('byte_order', c_int), ('bitmap_unit', c_int),
                ('bitmap_bit_order', c_int), ('bitmap_pad', c_int),
                ('depth', c_int), ('bytes_per_line', c_int),
                ('bits_per_pixel', c_int), ('red_mask', c_ulong),
                ('green_mask', c_ulong), ('blue_mask', c_ulong)]

class XRRModeInfo(Structure):
    ''' Voilà, voilà. '''

class XRRScreenResources(Structure):
    ''' Structure that contains arrays of XIDs that point to the
        available outputs and associated CRTCs.
    '''

    _fields_ = [('timestamp', c_ulong), ('configTimestamp', c_ulong),
                ('ncrtc', c_int), ('crtcs', POINTER(c_long)),
                ('noutput', c_int), ('outputs', POINTER(c_long)),
                ('nmode', c_int), ('modes', POINTER(XRRModeInfo))]

class XRRCrtcInfo(Structure):
    ''' Structure that contains CRTC informations. '''

    _fields_ = [('timestamp', c_ulong), ('x', c_int), ('y', c_int),
                ('width', c_int), ('height', c_int), ('mode', c_long),
                ('rotation', c_int), ('noutput', c_int),
                ('outputs', POINTER(c_long)), ('rotations', c_ushort),
                ('npossible', c_int), ('possible', POINTER(c_long))]

class MSS(MSSBase):
    ''' Mutliple ScreenShots implementation for GNU/Linux.
        It uses intensively the Xlib and Xrandr extension.
    '''

    # pylint: disable=too-many-instance-attributes

    display = None
    use_mss = False
    mss = None
    xlib = None
    xrandr = None
    display = None
    root = None

    def __del__(self):
        ''' Disconnect from X server. '''

        if self.display:
            self.xlib.XCloseDisplay(self.display)
            self.display = None

    def __init__(self, display=None):
        ''' GNU/Linux initialisations. '''

        if not display:
            try:
                if version > '3':
                    display = bytes(environ['DISPLAY'], 'utf-8')
                else:
                    display = environ['DISPLAY']
            except KeyError:
                err = '$DISPLAY not set. Stopping to prevent segfault.'
                raise ScreenshotError(err)
        elif not isinstance(display, bytes):
            display = bytes(display, 'utf-8')

        x11 = find_library('X11')
        if not x11:
            raise ScreenshotError('No X11 library found.')
        self.xlib = cdll.LoadLibrary(x11)

        xrandr = find_library('Xrandr')
        if not xrandr:
            raise ScreenshotError('No Xrandr extension found.')
        self.xrandr = cdll.LoadLibrary(xrandr)

        # libmss = find_library('mss')
        libmss = '{0}/linux/{1}/libmss.so'.format(
            dirname(realpath(abspath(__file__))), arch())
        if isfile(libmss):
            self.mss = cdll.LoadLibrary(libmss)
            self.use_mss = True
        else:
            print('No MSS library found. Using slow native function.')

        self._set_argtypes()
        self._set_restypes()

        self.display = self.xlib.XOpenDisplay(display)
        try:
            assert self.display.contents
        except ValueError:
            raise ScreenshotError('Cannot open display "{0}".'.format(display))
        screen = self.xlib.XDefaultScreen(self.display)
        self.root = self.xlib.XDefaultRootWindow(self.display, screen)

    def _set_argtypes(self):
        ''' Functions arguments.

            Curiously, if we set up XGetPixel arguments type,
            the entire process takes twice more time.
            So, no need to waste this precious time :)
            Note: this issue does not occur when using libmss.
        '''

        self.xlib.XOpenDisplay.argtypes = [c_char_p]
        self.xlib.XDefaultScreen.argtypes = [POINTER(Display)]
        self.xlib.XDefaultRootWindow.argtypes = [POINTER(Display), c_int]
        self.xlib.XGetWindowAttributes.argtypes = [POINTER(Display),
                                                   POINTER(XWindowAttributes),
                                                   POINTER(XWindowAttributes)]
        self.xlib.XAllPlanes.argtypes = []
        self.xlib.XGetImage.argtypes = [POINTER(Display), POINTER(Display),
                                        c_int, c_int, c_uint, c_uint, c_ulong,
                                        c_int]
        # self.xlib.XGetPixel.argtypes = [POINTER(XImage), c_int, c_int]
        self.xlib.XDestroyImage.argtypes = [POINTER(XImage)]
        self.xlib.XCloseDisplay.argtypes = [POINTER(Display)]
        self.xrandr.XRRGetScreenResources.argtypes = [POINTER(Display),
                                                      POINTER(Display)]
        self.xrandr.XRRGetCrtcInfo.argtypes = [POINTER(Display),
                                               POINTER(XRRScreenResources),
                                               c_long]
        self.xrandr.XRRFreeScreenResources.argtypes = \
            [POINTER(XRRScreenResources)]
        self.xrandr.XRRFreeCrtcInfo.argtypes = [POINTER(XRRCrtcInfo)]
        if self.use_mss:
            self.mss.GetXImagePixels.argtypes = [POINTER(XImage), c_void_p]

    def _set_restypes(self):
        ''' Functions return type. '''

        def validate(value):
            ''' Validate the returned value of xrandr.XRRGetScreenResources().
                We can end on a segfault if not:
                    Xlib:  extension "RANDR" missing on display "...".
            '''

            if value == 0:
                err = 'xrandr.XRRGetScreenResources() failed.'
                err += ' NULL pointer received.'
                raise ScreenshotError(err)

            return cast(value, POINTER(XRRScreenResources))

        self.xlib.XOpenDisplay.restype = POINTER(Display)
        self.xlib.XDefaultScreen.restype = c_int
        self.xlib.XGetWindowAttributes.restype = c_int
        self.xlib.XAllPlanes.restype = c_ulong
        self.xlib.XGetImage.restype = POINTER(XImage)
        self.xlib.XGetPixel.restype = c_ulong
        self.xlib.XDestroyImage.restype = c_void_p
        self.xlib.XCloseDisplay.restype = c_void_p
        self.xlib.XDefaultRootWindow.restype = POINTER(XWindowAttributes)
        self.xrandr.XRRGetScreenResources.restype = validate
        self.xrandr.XRRGetCrtcInfo.restype = POINTER(XRRCrtcInfo)
        self.xrandr.XRRFreeScreenResources.restype = c_void_p
        self.xrandr.XRRFreeCrtcInfo.restype = c_void_p
        if self.use_mss:
            self.mss.GetXImagePixels.restype = c_int

    def enum_display_monitors(self, force=False):
        ''' Get positions of monitors (see parent class). '''

        if not self.monitors or force:
            self.monitors = []

            # All monitors
            gwa = XWindowAttributes()
            self.xlib.XGetWindowAttributes(self.display, self.root, byref(gwa))
            self.monitors.append({
                'left': int(gwa.x),
                'top': int(gwa.y),
                'width': int(gwa.width),
                'height': int(gwa.height),
                'monitor': 0
            })

            # Each monitors
            # Fix for XRRGetScreenResources:
            #     expected LP_Display instance instead of LP_XWindowAttributes
            root = cast(self.root, POINTER(Display))
            mon = self.xrandr.XRRGetScreenResources(self.display, root)
            for num in range(mon.contents.ncrtc):
                crtc = self.xrandr.XRRGetCrtcInfo(self.display, mon,
                                                  mon.contents.crtcs[num])
                self.monitors.append({
                    'left': int(crtc.contents.x),
                    'top': int(crtc.contents.y),
                    'width': int(crtc.contents.width),
                    'height': int(crtc.contents.height),
                    'monitor': num
                })
                self.xrandr.XRRFreeCrtcInfo(crtc)
            self.xrandr.XRRFreeScreenResources(mon)

        return self.monitors

    def get_pixels(self, monitor):
        ''' Retrieve all pixels from a monitor. Pixels have to be RGB. '''

        self.width = monitor['width']
        self.height = monitor['height']
        left, top = monitor['left'], monitor['top']
        zpixmap = 2
        allplanes = self.xlib.XAllPlanes()

        # Fix for XGetImage:
        #     expected LP_Display instance instead of LP_XWindowAttributes
        root = cast(self.root, POINTER(Display))
        ximage = self.xlib.XGetImage(self.display, root, left, top,
                                     self.width, self.height, allplanes,
                                     zpixmap)
        if not ximage:
            err = 'xlib.XGetImage() failed. Monitor informations: '
            for key, val in sorted(monitor.items()):
                err = '{0}{1}: {2}, '.format(err, key, val)
            err = err.strip(', ')
            raise ScreenshotError(err)

        if not self.use_mss:
            self.get_pixels_slow(ximage)
        else:
            self.image = create_string_buffer(self.height * self.width * 3)
            ret = self.mss.GetXImagePixels(ximage, self.image)
            if not ret:
                self.xlib.XDestroyImage(ximage)
                err = 'libmss.GetXImagePixels() failed (retcode={0}).'
                raise ScreenshotError(err.format(ret))
        self.xlib.XDestroyImage(ximage)
        return self.image

    def get_pixels_slow(self, ximage):
        ''' Retrieve all pixels from a monitor. Pixels have to be RGB.
            (!) Insanely slow version, see doc/linux-slow-version. (!)
        '''

        # @TODO: this part takes most of the time. Need a better solution.
        def pix(pixel, _resultats={}, p__=pack):
            ''' Apply shifts to a pixel to get the RGB values.
                This method uses of memoization.
            '''

            # pylint: disable=dangerous-default-value

            if pixel not in _resultats:
                _resultats[pixel] = p__('<B', (pixel & rmask) >> 16) + \
                    p__('<B', (pixel & gmask) >> 8) + \
                    p__('<B', pixel & bmask)
            return _resultats[pixel]

        self.width = ximage.contents.width
        self.height = ximage.contents.height
        rmask = ximage.contents.red_mask
        bmask = ximage.contents.blue_mask
        gmask = ximage.contents.green_mask
        get_pix = self.xlib.XGetPixel
        pixels = [pix(get_pix(ximage, x, y))
                  for y in range(self.height) for x in range(self.width)]
        self.image = b''.join(pixels)
        return self.image

def arch():
    ''' Detect OS architecture.
        Returns an int: 32 or 64
    '''

    return 64 if maxsize > 2 ** 32 else 32
<commit_msg>Linux: get_pixels_slow() now uses EAFP<commit_after>#!/usr/bin/env python
# coding: utf-8
''' This is part of the MSS Python's module.
    Source: https://github.com/BoboTiG/python-mss
'''

from ctypes import (
    POINTER, Structure, byref, c_char_p, c_int, c_int32, c_long, c_uint,
    c_uint32, c_ulong, c_ushort, c_void_p, cast, cdll, create_string_buffer)
from ctypes.util import find_library
from os import environ
from os.path import abspath, dirname, isfile, realpath
from struct import pack
from sys import maxsize, version

from .base import MSSBase
from .exception import ScreenshotError

__all__ = ['MSS']

class Display(Structure):
    ''' Structure that serves as the connection to the X server
        and that contains all the information about that X server.
    '''

class XWindowAttributes(Structure):
    ''' Attributes for the specified window. '''

    _fields_ = [('x', c_int32), ('y', c_int32), ('width', c_int32),
                ('height', c_int32), ('border_width', c_int32),
                ('depth', c_int32), ('visual', c_ulong), ('root', c_ulong),
                ('class', c_int32), ('bit_gravity', c_int32),
                ('win_gravity', c_int32), ('backing_store', c_int32),
                ('backing_planes', c_ulong), ('backing_pixel', c_ulong),
                ('save_under', c_int32), ('colourmap', c_ulong),
                ('mapinstalled', c_uint32), ('map_state', c_uint32),
                ('all_event_masks', c_ulong), ('your_event_mask', c_ulong),
                ('do_not_propagate_mask', c_ulong),
                ('override_redirect', c_int32), ('screen', c_ulong)]

class XImage(Structure):
    ''' Description of an image as it exists in the client's memory.
        https://tronche.com/gui/x/xlib/graphics/images.html
    '''

    _fields_ = [('width', c_int), ('height', c_int), ('xoffset', c_int),
                ('format', c_int), ('data', c_void_p),
                ('byte_order', c_int), ('bitmap_unit', c_int),
                ('bitmap_bit_order', c_int), ('bitmap_pad', c_int),
                ('depth', c_int), ('bytes_per_line', c_int),
                ('bits_per_pixel', c_int), ('red_mask', c_ulong),
                ('green_mask', c_ulong), ('blue_mask', c_ulong)]

class XRRModeInfo(Structure):
    ''' Voilà, voilà. '''

class XRRScreenResources(Structure):
    ''' Structure that contains arrays of XIDs that point to the
        available outputs and associated CRTCs.
    '''

    _fields_ = [('timestamp', c_ulong), ('configTimestamp', c_ulong),
                ('ncrtc', c_int), ('crtcs', POINTER(c_long)),
                ('noutput', c_int), ('outputs', POINTER(c_long)),
                ('nmode', c_int), ('modes', POINTER(XRRModeInfo))]

class XRRCrtcInfo(Structure):
    ''' Structure that contains CRTC informations. '''

    _fields_ = [('timestamp', c_ulong), ('x', c_int), ('y', c_int),
                ('width', c_int), ('height', c_int), ('mode', c_long),
                ('rotation', c_int), ('noutput', c_int),
                ('outputs', POINTER(c_long)), ('rotations', c_ushort),
                ('npossible', c_int), ('possible', POINTER(c_long))]

class MSS(MSSBase):
    ''' Mutliple ScreenShots implementation for GNU/Linux.
        It uses intensively the Xlib and Xrandr extension.
    '''

    # pylint: disable=too-many-instance-attributes

    display = None
    use_mss = False
    mss = None
    xlib = None
    xrandr = None
    display = None
    root = None

    def __del__(self):
        ''' Disconnect from X server. '''

        if self.display:
            self.xlib.XCloseDisplay(self.display)
            self.display = None

    def __init__(self, display=None):
        ''' GNU/Linux initialisations. '''

        if not display:
            try:
                if version > '3':
                    display = bytes(environ['DISPLAY'], 'utf-8')
                else:
                    display = environ['DISPLAY']
            except KeyError:
                err = '$DISPLAY not set. Stopping to prevent segfault.'
                raise ScreenshotError(err)
        elif not isinstance(display, bytes):
            display = bytes(display, 'utf-8')

        x11 = find_library('X11')
        if not x11:
            raise ScreenshotError('No X11 library found.')
        self.xlib = cdll.LoadLibrary(x11)

        xrandr = find_library('Xrandr')
        if not xrandr:
            raise ScreenshotError('No Xrandr extension found.')
        self.xrandr = cdll.LoadLibrary(xrandr)

        # libmss = find_library('mss')
        libmss = '{0}/linux/{1}/libmss.so'.format(
            dirname(realpath(abspath(__file__))), arch())
        if isfile(libmss):
            self.mss = cdll.LoadLibrary(libmss)
            self.use_mss = True
        # else:
        #    print('No MSS library found. Using slow native function.')

        self._set_argtypes()
        self._set_restypes()

        self.display = self.xlib.XOpenDisplay(display)
        try:
            assert self.display.contents
        except ValueError:
            raise ScreenshotError('Cannot open display "{0}".'.format(display))
        screen = self.xlib.XDefaultScreen(self.display)
        self.root = self.xlib.XDefaultRootWindow(self.display, screen)

    def _set_argtypes(self):
        ''' Functions arguments.

            Curiously, if we set up XGetPixel arguments type,
            the entire process takes twice more time.
            So, no need to waste this precious time :)
            Note: this issue does not occur when using libmss.
        '''

        self.xlib.XOpenDisplay.argtypes = [c_char_p]
        self.xlib.XDefaultScreen.argtypes = [POINTER(Display)]
        self.xlib.XDefaultRootWindow.argtypes = [POINTER(Display), c_int]
        self.xlib.XGetWindowAttributes.argtypes = [POINTER(Display),
                                                   POINTER(XWindowAttributes),
                                                   POINTER(XWindowAttributes)]
        self.xlib.XAllPlanes.argtypes = []
        self.xlib.XGetImage.argtypes = [POINTER(Display), POINTER(Display),
                                        c_int, c_int, c_uint, c_uint, c_ulong,
                                        c_int]
        # self.xlib.XGetPixel.argtypes = [POINTER(XImage), c_int, c_int]
        self.xlib.XDestroyImage.argtypes = [POINTER(XImage)]
        self.xlib.XCloseDisplay.argtypes = [POINTER(Display)]
        self.xrandr.XRRGetScreenResources.argtypes = [POINTER(Display),
                                                      POINTER(Display)]
        self.xrandr.XRRGetCrtcInfo.argtypes = [POINTER(Display),
                                               POINTER(XRRScreenResources),
                                               c_long]
        self.xrandr.XRRFreeScreenResources.argtypes = \
            [POINTER(XRRScreenResources)]
        self.xrandr.XRRFreeCrtcInfo.argtypes = [POINTER(XRRCrtcInfo)]
        if self.use_mss:
            self.mss.GetXImagePixels.argtypes = [POINTER(XImage), c_void_p]

    def _set_restypes(self):
        ''' Functions return type. '''

        def validate(value):
            ''' Validate the returned value of xrandr.XRRGetScreenResources().
                We can end on a segfault if not:
                    Xlib:  extension "RANDR" missing on display "...".
            '''

            if value == 0:
                err = 'xrandr.XRRGetScreenResources() failed.'
                err += ' NULL pointer received.'
                raise ScreenshotError(err)

            return cast(value, POINTER(XRRScreenResources))

        self.xlib.XOpenDisplay.restype = POINTER(Display)
        self.xlib.XDefaultScreen.restype = c_int
        self.xlib.XGetWindowAttributes.restype = c_int
        self.xlib.XAllPlanes.restype = c_ulong
        self.xlib.XGetImage.restype = POINTER(XImage)
        self.xlib.XGetPixel.restype = c_ulong
        self.xlib.XDestroyImage.restype = c_void_p
        self.xlib.XCloseDisplay.restype = c_void_p
        self.xlib.XDefaultRootWindow.restype = POINTER(XWindowAttributes)
        self.xrandr.XRRGetScreenResources.restype = validate
        self.xrandr.XRRGetCrtcInfo.restype = POINTER(XRRCrtcInfo)
        self.xrandr.XRRFreeScreenResources.restype = c_void_p
        self.xrandr.XRRFreeCrtcInfo.restype = c_void_p
        if self.use_mss:
            self.mss.GetXImagePixels.restype = c_int

    def enum_display_monitors(self, force=False):
        ''' Get positions of monitors (see parent class). '''

        if not self.monitors or force:
            self.monitors = []

            # All monitors
            gwa = XWindowAttributes()
            self.xlib.XGetWindowAttributes(self.display, self.root, byref(gwa))
            self.monitors.append({
                'left': int(gwa.x),
                'top': int(gwa.y),
                'width': int(gwa.width),
                'height': int(gwa.height),
                'monitor': 0
            })

            # Each monitors
            # Fix for XRRGetScreenResources:
            #     expected LP_Display instance instead of LP_XWindowAttributes
            root = cast(self.root, POINTER(Display))
            mon = self.xrandr.XRRGetScreenResources(self.display, root)
            for num in range(mon.contents.ncrtc):
                crtc = self.xrandr.XRRGetCrtcInfo(self.display, mon,
                                                  mon.contents.crtcs[num])
                self.monitors.append({
                    'left': int(crtc.contents.x),
                    'top': int(crtc.contents.y),
                    'width': int(crtc.contents.width),
                    'height': int(crtc.contents.height),
                    'monitor': num
                })
                self.xrandr.XRRFreeCrtcInfo(crtc)
            self.xrandr.XRRFreeScreenResources(mon)

        return self.monitors

    def get_pixels(self, monitor):
        ''' Retrieve all pixels from a monitor. Pixels have to be RGB. '''

        self.width = monitor['width']
        self.height = monitor['height']
        left, top = monitor['left'], monitor['top']
        zpixmap = 2
        allplanes = self.xlib.XAllPlanes()

        # Fix for XGetImage:
        #     expected LP_Display instance instead of LP_XWindowAttributes
        root = cast(self.root, POINTER(Display))
        ximage = self.xlib.XGetImage(self.display, root, left, top,
                                     self.width, self.height, allplanes,
                                     zpixmap)
        if not ximage:
            err = 'xlib.XGetImage() failed. Monitor informations: '
            for key, val in sorted(monitor.items()):
                err = '{0}{1}: {2}, '.format(err, key, val)
            err = err.strip(', ')
            raise ScreenshotError(err)

        if not self.use_mss:
            self.get_pixels_slow(ximage)
        else:
            self.image = create_string_buffer(self.height * self.width * 3)
            ret = self.mss.GetXImagePixels(ximage, self.image)
            if not ret:
                self.xlib.XDestroyImage(ximage)
                err = 'libmss.GetXImagePixels() failed (retcode={0}).'
                raise ScreenshotError(err.format(ret))
        self.xlib.XDestroyImage(ximage)
        return self.image

    def get_pixels_slow(self, ximage):
        ''' Retrieve all pixels from a monitor. Pixels have to be RGB.

            (!) This version is damn slow, but if you consider that ~1 second
            for a screenshot of decent monitor is not too long, you can
            continue. There should be few things to tweak here to gain speed.
        '''

        # @TODO: this part takes most of the time. Need a better solution.
        def pix(pixel, _resultats={}, p__=pack):
            ''' Apply shifts to a pixel to get the RGB values.
                This method uses of memoization.
            '''

            # pylint: disable=dangerous-default-value

            try:
                return _resultats[pixel]
            except KeyError:
                _resultats[pixel] = p__('<B', (pixel & rmask) >> 16) + \
                    p__('<B', (pixel & gmask) >> 8) + \
                    p__('<B', pixel & bmask)
                return _resultats[pixel]

        self.width = ximage.contents.width
        self.height = ximage.contents.height
        rmask = ximage.contents.red_mask
        bmask = ximage.contents.blue_mask
        gmask = ximage.contents.green_mask
        get_pix = self.xlib.XGetPixel
        pixels = [pix(get_pix(ximage, x, y))
                  for y in range(self.height) for x in range(self.width)]
        self.image = b''.join(pixels)
        return self.image

def arch():
    ''' Detect OS architecture.
        Returns an int: 32 or 64
    '''

    return 64 if maxsize > 2 ** 32 else 32
<|endoftext|>


<commit_before>#!/usr/bin/env python

# Copyright (c) 2017, DIANA-HEP
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import numbers
import math

import numpy
try:
    import numba
    import llvmlite
except ImportError:
    numba = None

import uproot.rootio

class TH1Methods(object):
    # makes __doc__ attribute mutable before Python 3.3
    __metaclass__ = type.__new__(type, "type", (uproot.rootio.ROOTObject.__metaclass__,), {})

    def __repr__(self):
        if self.fName is None:
            return "<{0} at 0x{1:012x}>".format(self.classname, id(self))
        else:
            return "<{0} {1} 0x{2:012x}>".format(self.classname, repr(self.fName), id(self))

    @property
    def name(self):
        return self.fName

    @property
    def title(self):
        return self.fTitle

    @property
    def numbins(self):
        return self.fXaxis.fNbins

    @property
    def low(self):
        return self.fXaxis.fXmin

    @property
    def high(self):
        return self.fXaxis.fXmax

    @property
    def underflows(self):
        return self[0]

    @property
    def overflows(self):
        return self[-1]

    @property
    def values(self):
        return self[1:-1]

    @property
    def allvalues(self):
        return self[:]

    @property
    def numpy(self):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        norm = (high - low) / self.fXaxis.fNbins
        freq = numpy.array(self.values, dtype=self._dtype.newbyteorder("="))
        edges = numpy.array([i*norm + low for i in range(self.numbins + 1)])
        return freq, edges

    def interval(self, index):
        if index < 0:
            index += len(self)

        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if index == 0:
            return (float("-inf"), low)
        elif index == len(self) - 1:
            return (high, float("inf"))
        else:
            norm = (high - low) / self.fXaxis.fNbins
            return (index - 1)*norm + low, index*norm + low

    def index(self, data):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if data < low:
            return 0
        elif data >= high:
            return len(self) - 1
        else:
            return int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1

    def fill(self, data, weights=None):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax

        if isinstance(data, numbers.Real):
            if weights is None:
                weights = 1

            if data < low:
                self[0] += weight
            elif data >= high:
                self[-1] += weight
            else:
                self[int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1] += weights

        else:
            if not isinstance(data, numpy.ndarray):
                data = numpy.array(data)

            if isinstance(weights, numbers.Real):
                weights = numpy.empty_like(data)

            freq, edges = numpy.histogram(data,
                                          bins=numpy.linspace(low, high, self.fXaxis.fNbins + 1),
                                          weights=weights,
                                          density=False)
            for i, x in enumerate(freq):
                self[i + 1] += x

            underflows = (data < low)
            overflows = (data >= high)

            if isinstance(weights, numpy.ndarray):
                self[0] += weights[underflows].sum()
                self[-1] += weights[overflows].sum()
            else:
                self[0] += underflows.sum()
                self[-1] += overflows.sum()

    def show(self, width=80, minimum=None, maximum=None, stream=sys.stdout):
        if minimum is None:
            minimum = min(self)
            if minimum < 0:
                minimum *= 1.05
            else:
                minimum = 0

        if maximum is None:
            maximum = max(self) * 1.05

        if maximum <= minimum:
            average = (minimum + maximum) / 2.0
            minimum = average - 0.5
            maximum = average + 0.5

        intervals = ["[{0:<.4g}, {1:<.4g})".format(l, h) for l, h in [self.interval(i) for i in range(len(self))]]
        intervals[-1] = intervals[-1][:-1] + "]"   # last interval is closed on top edge
        intervalswidth = max(len(x) for x in intervals)

        values = ["{0:<.4g}".format(float(x)) for x in self]
        valueswidth = max(len(x) for x in values)

        minimumtext = "{0:<.4g}".format(minimum)
        maximumtext = "{0:<.4g}".format(maximum)

        plotwidth = max(len(minimumtext) + len(maximumtext), width - (intervalswidth + 1 + valueswidth + 1 + 2))
        scale = minimumtext + " "*(plotwidth + 2 - len(minimumtext) - len(maximumtext)) + maximumtext

        norm = float(plotwidth) / float(maximum - minimum)
        zero = int(round((0.0 - minimum)*norm))
        line = numpy.empty(plotwidth, dtype=numpy.uint8)

        formatter = "{0:<%s} {1:<%s} |{2}|" % (intervalswidth, valueswidth)
        line[:] = ord("-")
        if minimum != 0 and 0 <= zero < plotwidth:
            line[zero] = ord("+")
        capstone = " " * (intervalswidth + 1 + valueswidth + 1) + "+" + str(line.tostring().decode("ascii")) + "+"

        out = [" "*(intervalswidth + valueswidth + 2) + scale]
        out.append(capstone)
        for interval, value, x in zip(intervals, values, self):
            line[:] = ord(" ")

            pos = int(round((x - minimum)*norm))
            if x < 0:
                line[pos:zero] = ord("*")
            else:
                line[zero:pos] = ord("*")

            if minimum != 0 and 0 <= zero < plotwidth:
                line[zero] = ord("|")

            out.append(formatter.format(interval, value, str(line.tostring().decode("ascii"))))

        out.append(capstone)
        out = "\n".join(out)
        if stream is None:
            return out
        else:
            stream.write(out)
            stream.write("\n")

    @property
    def holoviews(self):
        import uproot._connect.to_holoviews
        return uproot._connect.to_holoviews.TH1Methods_holoviews(self)

    @property
    def bokeh(self):
        import uproot._connect.to_bokeh
        return uproot._connect.to_bokeh.TH1Methods_bokeh(self)

uproot.rootio.methods["TH1"] = TH1Methods

class TH1(TH1Methods, list):
    def _type(self):
        if all(isinstance(x, numbers.Integral) for x in self):
            return int
        elif all(isinstance(x, numbers.Real) for x in self):
            return float
        else:
            raise TypeError("histogram bin values must be integers or floats")

    @property
    def classname(self):
        if self._type() is int:
            return "TH1I"
        else:
            return "TH1D"

    @property
    def _dtype(self):
        if self._type() is int:
            return numpy.dtype(">i4")
        else:
            return numpy.dtype(">f8")

class TAxis(object):
    classname = "TAxis"

def hist(numbins, low, high, name=None, title=None):
    out = TH1()
    out.fXaxis = TAxis()
    out.fXaxis.fNbins = int(numbins)
    out.fXaxis.fXmin = float(low)
    out.fXaxis.fXmax = float(high)
    out.fName = name
    out.fTitle = title
    out.extend([0] * (numbins + 2))
    return out

if numba is not None:
    class Regular1dType(numba.types.Type):
        def __init__(self, valuetype):
            super(Regular1dType, self).__init__(name="Regular1dType")
            self.valuetype = valuetype
            self.obj = numba.types.PyObject("obj")
            self.numbins = numba.types.int64
            self.low = numba.types.float64
            self.high = numba.types.float64
            if self.valuetype is int:
                self.data = numba.types.int64[:]
            else:
                self.data = numba.types.float64[:]

    @numba.extending.typeof_impl.register(TH1Methods)
    def th1_typeof(val, c):
        assert isinstance(val, TH1Methods)
        if val.classname == "TH1F" or val.classname == "TH1D":
            return Regular1dType(float)
        else:
            return Regular1dType(int)

    @numba.extending.register_model(Regular1dType)
    class Regular1dModel(numba.datamodel.models.StructModel):
        def __init__(self, dmm, fe_type):
            members = [(x, getattr(fe_type, x)) for x in "obj", "numbins", "low", "high", "data"]
            super(Regular1dModel, self).__init__(dmm, fe_type, members)

    @numba.extending.unbox(Regular1dType)
    def th1_unbox(typ, obj, c):
        struct = numba.cgutils.create_struct_proxy(typ)(c.context, c.builder)

        numbins_obj = c.pyapi.object_getattr_string(obj, "numbins")
        low_obj = c.pyapi.object_getattr_string(obj, "low")
        high_obj = c.pyapi.object_getattr_string(obj, "high")

        array_fcn = c.pyapi.unserialize(c.pyapi.serialize_object(numpy.array))
        valuetype_obj = c.pyapi.unserialize(c.pyapi.serialize_object(typ.valuetype))
        array_obj = c.pyapi.call_function_objargs(array_fcn, (obj, valuetype_obj))

        struct.obj = obj
        struct.numbins = c.pyapi.long_as_longlong(numbins_obj)
        struct.low = c.pyapi.float_as_double(low_obj)
        struct.high = c.pyapi.float_as_double(high_obj)
        struct.data = c.unbox(typ.data, array_obj).value

        c.pyapi.decref(numbins_obj)
        c.pyapi.decref(low_obj)
        c.pyapi.decref(high_obj)
        c.pyapi.decref(array_fcn)
        c.pyapi.decref(valuetype_obj)
        c.pyapi.decref(array_obj)

        is_error = numba.cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
        return numba.extending.NativeValue(struct._getvalue(), is_error=is_error)

    def th1_merge(obj, data):
        for i, x in enumerate(data):
            obj[i] = x

    @numba.extending.box(Regular1dType)
    def th1_box(typ, val, c):
        struct = numba.cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
        obj = struct.obj

        array_obj = c.box(typ.data, struct.data)
        merge_fcn = c.pyapi.unserialize(c.pyapi.serialize_object(th1_merge))

        c.pyapi.call_function_objargs(merge_fcn, (obj, array_obj))

        c.pyapi.decref(array_obj)

        return obj

def doit():
    @numba.njit
    def testy(x):
        return x

    h = hist(10, -3.0, 3.0)
    h[3] = 7
    print h, h.values
    out = testy(h)
    print out
    print out.name, out.low, out.high, out.values
    out2 = testy(out)
    print out2
    print out2.name, out2.low, out2.high, out2.values
    print out
    print out.name, out.low, out.high, out.values
    print h
    print h.name, h.low, h.high, h.values
<commit_msg>actually filling<commit_after>#!/usr/bin/env python

# Copyright (c) 2017, DIANA-HEP
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import numbers
import math

import numpy
try:
    import numba
    import llvmlite
except ImportError:
    numba = None

import uproot.rootio

class TH1Methods(object):
    # makes __doc__ attribute mutable before Python 3.3
    __metaclass__ = type.__new__(type, "type", (uproot.rootio.ROOTObject.__metaclass__,), {})

    def __repr__(self):
        if self.fName is None:
            return "<{0} at 0x{1:012x}>".format(self.classname, id(self))
        else:
            return "<{0} {1} 0x{2:012x}>".format(self.classname, repr(self.fName), id(self))

    @property
    def name(self):
        return self.fName

    @property
    def title(self):
        return self.fTitle

    @property
    def numbins(self):
        return self.fXaxis.fNbins

    @property
    def low(self):
        return self.fXaxis.fXmin

    @property
    def high(self):
        return self.fXaxis.fXmax

    @property
    def underflows(self):
        return self[0]

    @property
    def overflows(self):
        return self[-1]

    @property
    def values(self):
        return self[1:-1]

    @property
    def allvalues(self):
        return self[:]

    @property
    def numpy(self):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        norm = (high - low) / self.fXaxis.fNbins
        freq = numpy.array(self.values, dtype=self._dtype.newbyteorder("="))
        edges = numpy.array([i*norm + low for i in range(self.numbins + 1)])
        return freq, edges

    def interval(self, index):
        if index < 0:
            index += len(self)

        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if index == 0:
            return (float("-inf"), low)
        elif index == len(self) - 1:
            return (high, float("inf"))
        else:
            norm = (high - low) / self.fXaxis.fNbins
            return (index - 1)*norm + low, index*norm + low

    def index(self, data):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax
        if data < low:
            return 0
        elif data >= high:
            return len(self) - 1
        elif not math.isnan(data):
            return int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1

    def fill(self, data, weights=None):
        low = self.fXaxis.fXmin
        high = self.fXaxis.fXmax

        if isinstance(data, numbers.Real):
            if weights is None:
                weights = 1

            if data < low:
                self[0] += weight
            elif data >= high:
                self[-1] += weight
            else:
                self[int(math.floor(self.fXaxis.fNbins * (data - low) / (high - low))) + 1] += weights

        else:
            if not isinstance(data, numpy.ndarray):
                data = numpy.array(data)

            if isinstance(weights, numbers.Real):
                weights = numpy.empty_like(data)

            freq, edges = numpy.histogram(data, bins=self.fXaxis.fNbins, range=(low, high), weights=weights, density=False)
            for i, x in enumerate(freq):
                self[i + 1] += x

            underflows = (data < low)
            overflows = (data >= high)

            if isinstance(weights, numpy.ndarray):
                self[0] += weights[underflows].sum()
                self[-1] += weights[overflows].sum()
            else:
                self[0] += underflows.sum()
                self[-1] += overflows.sum()

    def show(self, width=80, minimum=None, maximum=None, stream=sys.stdout):
        if minimum is None:
            minimum = min(self)
            if minimum < 0:
                minimum *= 1.05
            else:
                minimum = 0

        if maximum is None:
            maximum = max(self) * 1.05

        if maximum <= minimum:
            average = (minimum + maximum) / 2.0
            minimum = average - 0.5
            maximum = average + 0.5

        intervals = ["[{0:<.4g}, {1:<.4g})".format(l, h) for l, h in [self.interval(i) for i in range(len(self))]]
        intervals[-1] = intervals[-1][:-1] + "]"   # last interval is closed on top edge
        intervalswidth = max(len(x) for x in intervals)

        values = ["{0:<.4g}".format(float(x)) for x in self]
        valueswidth = max(len(x) for x in values)

        minimumtext = "{0:<.4g}".format(minimum)
        maximumtext = "{0:<.4g}".format(maximum)

        plotwidth = max(len(minimumtext) + len(maximumtext), width - (intervalswidth + 1 + valueswidth + 1 + 2))
        scale = minimumtext + " "*(plotwidth + 2 - len(minimumtext) - len(maximumtext)) + maximumtext

        norm = float(plotwidth) / float(maximum - minimum)
        zero = int(round((0.0 - minimum)*norm))
        line = numpy.empty(plotwidth, dtype=numpy.uint8)

        formatter = "{0:<%s} {1:<%s} |{2}|" % (intervalswidth, valueswidth)
        line[:] = ord("-")
        if minimum != 0 and 0 <= zero < plotwidth:
            line[zero] = ord("+")
        capstone = " " * (intervalswidth + 1 + valueswidth + 1) + "+" + str(line.tostring().decode("ascii")) + "+"

        out = [" "*(intervalswidth + valueswidth + 2) + scale]
        out.append(capstone)
        for interval, value, x in zip(intervals, values, self):
            line[:] = ord(" ")

            pos = int(round((x - minimum)*norm))
            if x < 0:
                line[pos:zero] = ord("*")
            else:
                line[zero:pos] = ord("*")

            if minimum != 0 and 0 <= zero < plotwidth:
                line[zero] = ord("|")

            out.append(formatter.format(interval, value, str(line.tostring().decode("ascii"))))

        out.append(capstone)
        out = "\n".join(out)
        if stream is None:
            return out
        else:
            stream.write(out)
            stream.write("\n")

    @property
    def holoviews(self):
        import uproot._connect.to_holoviews
        return uproot._connect.to_holoviews.TH1Methods_holoviews(self)

    @property
    def bokeh(self):
        import uproot._connect.to_bokeh
        return uproot._connect.to_bokeh.TH1Methods_bokeh(self)

uproot.rootio.methods["TH1"] = TH1Methods

class TH1(TH1Methods, list):
    def _type(self):
        if all(isinstance(x, numbers.Integral) for x in self):
            return int
        elif all(isinstance(x, numbers.Real) for x in self):
            return float
        else:
            raise TypeError("histogram bin values must be integers or floats")

    @property
    def classname(self):
        if self._type() is int:
            return "TH1I"
        else:
            return "TH1D"

    @property
    def _dtype(self):
        if self._type() is int:
            return numpy.dtype(">i4")
        else:
            return numpy.dtype(">f8")

class TAxis(object):
    classname = "TAxis"

def hist(numbins, low, high, name=None, title=None):
    out = TH1()
    out.fXaxis = TAxis()
    out.fXaxis.fNbins = int(numbins)
    out.fXaxis.fXmin = float(low)
    out.fXaxis.fXmax = float(high)
    out.fName = name
    out.fTitle = title
    out.extend([0] * (numbins + 2))
    return out

if numba is not None:
    class Regular1dType(numba.types.Type):
        def __init__(self, valuetype):
            super(Regular1dType, self).__init__(name="Regular1dType")
            self.valuetype = valuetype
            self.obj = numba.types.PyObject("obj")
            self.numbins = numba.types.int64
            self.low = numba.types.float64
            self.high = numba.types.float64
            if self.valuetype is int:
                self.data = numba.types.int64[:]
            else:
                self.data = numba.types.float64[:]

    @numba.extending.typeof_impl.register(TH1Methods)
    def th1_typeof(val, c):
        assert isinstance(val, TH1Methods)
        if val.classname == "TH1F" or val.classname == "TH1D":
            return Regular1dType(float)
        else:
            return Regular1dType(int)

    @numba.extending.register_model(Regular1dType)
    class Regular1dModel(numba.datamodel.models.StructModel):
        def __init__(self, dmm, fe_type):
            members = [(x, getattr(fe_type, x)) for x in "obj", "numbins", "low", "high", "data"]
            super(Regular1dModel, self).__init__(dmm, fe_type, members)

    numba.extending.make_attribute_wrapper(Regular1dType, "numbins", "numbins")
    numba.extending.make_attribute_wrapper(Regular1dType, "low", "low")
    numba.extending.make_attribute_wrapper(Regular1dType, "high", "high")
    numba.extending.make_attribute_wrapper(Regular1dType, "data", "_data")

    @numba.extending.unbox(Regular1dType)
    def th1_unbox(typ, obj, c):
        struct = numba.cgutils.create_struct_proxy(typ)(c.context, c.builder)

        numbins_obj = c.pyapi.object_getattr_string(obj, "numbins")
        low_obj = c.pyapi.object_getattr_string(obj, "low")
        high_obj = c.pyapi.object_getattr_string(obj, "high")

        array_fcn = c.pyapi.unserialize(c.pyapi.serialize_object(numpy.array))
        valuetype_obj = c.pyapi.unserialize(c.pyapi.serialize_object(typ.valuetype))
        array_obj = c.pyapi.call_function_objargs(array_fcn, (obj, valuetype_obj))

        struct.obj = obj
        struct.numbins = c.pyapi.long_as_longlong(numbins_obj)
        struct.low = c.pyapi.float_as_double(low_obj)
        struct.high = c.pyapi.float_as_double(high_obj)
        struct.data = c.unbox(typ.data, array_obj).value

        c.pyapi.decref(numbins_obj)
        c.pyapi.decref(low_obj)
        c.pyapi.decref(high_obj)
        c.pyapi.decref(array_fcn)
        c.pyapi.decref(valuetype_obj)
        c.pyapi.decref(array_obj)

        is_error = numba.cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
        return numba.extending.NativeValue(struct._getvalue(), is_error=is_error)

    def th1_merge(obj, data):
        obj[:] = data

    @numba.extending.box(Regular1dType)
    def th1_box(typ, val, c):
        struct = numba.cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
        obj = struct.obj

        array_obj = c.box(typ.data, struct.data)
        merge_fcn = c.pyapi.unserialize(c.pyapi.serialize_object(th1_merge))

        c.pyapi.call_function_objargs(merge_fcn, (obj, array_obj))

        c.pyapi.decref(array_obj)

        return obj

    @numba.extending.overload_method(Regular1dType, "fill")
    def th1_fill(regular1dType, data):
        if isinstance(data, numba.types.Number):
            def fill_impl(th1, data):
                if data < th1.low:
                    th1._data[0] += 1
                elif data >= th1.high:
                    th1._data[-1] += 1
                elif not math.isnan(data):
                    th1._data[int(math.floor(th1.numbins * (data - th1.low) / (th1.high - th1.low))) + 1] += 1
            return fill_impl

def doit():
    @numba.njit
    def testy(x):
        x.fill(0)
        x.fill(0)
        x.fill(0)
        return x

    h = hist(10, -3.0, 3.0)
    h[3] = 7
    print testy(h).values
<|endoftext|>


<commit_before>#!/usr/bin/env python
'''Base module components.'''

import copy
from collections import OrderedDict

import inspect

from contextlib import contextmanager
from sklearn.base import _pprint

import six

class BaseTransformer(object):
    '''The base class for all transformation objects.
    This class implements a single transformation (history)
    and some various niceties.'''

    # This bit gleefully stolen from sklearn.base
    @classmethod
    def _get_param_names(cls):
        '''Get the list of parameter names for the object'''

        init = cls.__init__

        if init is object.__init__:
            return []

        args, varargs = inspect.getargspec(init)[:2]

        if varargs is not None:
            raise RuntimeError('varargs ist verboten')

        args.pop(0)
        args.sort()
        return args

    def get_params(self, deep=True):
        '''Get the parameters for this object.  Returns as a dict.

        Parameters
        ----------
        deep : bool
            Recurse on nested objects

        Returns
        -------
        params : dict
            A dictionary containing all parameters for this object
        '''

        out = dict(__class__=self.__class__,
                   params=dict())

        for key in self._get_param_names():
            value = getattr(self, key, None)

            if deep and hasattr(value, 'get_params'):
                deep_items = value.get_params().items()
                out['params'][key] = dict(__class__=value.__class__)
                out['params'][key].update((k, val) for k, val in deep_items)
            else:
                out['params'][key] = value

        return out

    def __repr__(self):
        '''Pretty-print this object'''

        class_name = self.__class__.__name__
        return '{:s}({:s})'.format(class_name,
                                   _pprint(self.get_params(deep=False)['params'],
                                           offset=len(class_name),),)

    def __init__(self):
        '''Base-class initialization'''
        self.dispatch = OrderedDict()
        self._state = dict()

        self.n_samples = 1

    def get_state(self, jam):
        '''Build the state object for a static transformer'''
        return dict()

    @contextmanager
    def _transform_state(self, jam):
        '''Initialize state for static transformers.'''
        self._state = self.get_state(jam)
        yield

    def _transform(self, jam):
        '''Apply the transformation to audio and annotations.

        The input jam is copied and modified, and returned
        contained in a list.

        Parameters
        ----------
        jam : jams.JAMS
            A single jam object to modify

        Returns
        -------
        jam_list : list
            A length-1 list containing `jam` after transformation

        See also
        --------
        core.load_jam_audio
        '''

        if not hasattr(jam.sandbox, 'muda'):
            raise RuntimeError('No muda state found in jams sandbox.')

        # We'll need a working copy of this object for modification purposes
        jam_working = copy.deepcopy(jam)

        with self._transform_state(jam_working):
            # Push our reconstructor onto the history stack
            jam_working.sandbox.muda['history'].append({'transformer': self.__serialize__,
                                                        'state': self._state})

            if hasattr(self, 'audio'):
                self.audio(jam_working.sandbox.muda)

            if hasattr(self, 'metadata'):
                self.metadata(jam_working.file_metadata)

            # Walk over the list of deformers
            for query, function in six.iteritems(self.dispatch):
                for matched_annotation in jam_working.search(namespace=query):
                    function(matched_annotation)

        return [jam_working]

    def transform(self, jam):
        '''Iterative transformation generator

        Generates up to a fixed number of transformed jams
        from a single input.

        Parameters
        ----------
        jam : jams.JAMS
            The jam to transform

        Generates
        ---------
        jam_out : jams.JAMS
            Iterator of transformed jams
        '''

        # Reset the state
        self._state = dict()

        i = 0
        while self.n_samples is None or i < self.n_samples:
            for jam_out in self._transform(jam):
                yield jam_out
                i = i + 1

    @property
    def __serialize__(self):
        '''Serializer'''

        data = self.get_params()
        data['__class__'] = data['__class__'].__name__
        return data

class Pipeline(object):
    '''Wrapper which allows multiple transformers to be chained together'''

    def __init__(self, *steps):
        '''Transformation pipeline.

        A given JAMS object will be transformed sequentially by
        each stage of the pipeline.

        Parameters
        ----------
        steps : argument array
            steps[i] is a tuple of `(name, Transformer)`

        Examples
        --------
        >>> P = muda.deformers.PitchShift(semitones=5)
        >>> T = muda.deformers.TimeStretch(speed=1.25)
        >>> Pipe = Pipeline( ('Pitch:maj3', P), ('Speed:1.25x', T) )
        >>> output = Pipe.transform(data)
        '''

        named_steps = dict(steps)
        names, transformers = zip(*steps)

        if len(named_steps) != len(steps):
            raise ValueError("Names provided are not unique: "
                             " {:s}".format(names,))

        # shallow copy of steps
        self.steps = list(zip(names, transformers))

        for t in transformers:
            if not isinstance(t, BaseTransformer):
                raise TypeError('{:s} is not a BaseTransformer'.format(t))

    def get_params(self):
        '''Get the parameters for this object.  Returns as a dict.'''

        out = {}
        out['__class__'] = self.__class__
        out['params'] = []

        for name, step in self.steps:
            out['params'].append([name, step.get_params(deep=True)])

        return out

    def __repr__(self):
        '''Pretty-print the object'''

        class_name = self.__class__.__name__
        return '{:s}({:s})'.format(class_name,
                                   _pprint(self.get_params(),
                                           offset=len(class_name),),)

    def __recursive_transform(self, jam, steps):
        '''A recursive transformation pipeline'''

        if len(steps) > 0:
            head_transformer = steps[0][1]
            for t_jam in head_transformer.transform(jam):
                for q in self.__recursive_transform(t_jam, steps[1:]):
                    yield q
        else:
            yield jam

    def transform(self, jam):
        '''Apply the sequence of transformations to a single jam object.

        Parameters
        ----------
        jam : jams.JAMS
            The jam object to transform

        Generates
        ---------
        jam_stream : iterable

        See also
        --------
        BaseTransformer.transform
        '''

        for output in self.__recursive_transform(jam, self.steps):
            yield output
<commit_msg>minor code formatting changes<commit_after>#!/usr/bin/env python
'''Base module components.'''

import copy
from collections import OrderedDict

import inspect

from contextlib import contextmanager
from sklearn.base import _pprint

import six

class BaseTransformer(object):
    '''The base class for all transformation objects.
    This class implements a single transformation (history)
    and some various niceties.'''

    # This bit gleefully stolen from sklearn.base
    @classmethod
    def _get_param_names(cls):
        '''Get the list of parameter names for the object'''

        init = cls.__init__

        if init is object.__init__:
            return []

        args, varargs = inspect.getargspec(init)[:2]

        if varargs is not None:
            raise RuntimeError('varargs ist verboten')

        args.pop(0)
        args.sort()
        return args

    def get_params(self, deep=True):
        '''Get the parameters for this object.  Returns as a dict.

        Parameters
        ----------
        deep : bool
            Recurse on nested objects

        Returns
        -------
        params : dict
            A dictionary containing all parameters for this object
        '''

        out = dict(__class__=self.__class__,
                   params=dict())

        for key in self._get_param_names():
            value = getattr(self, key, None)

            if deep and hasattr(value, 'get_params'):
                deep_items = value.get_params().items()
                out['params'][key] = dict(__class__=value.__class__)
                out['params'][key].update((k, val) for k, val in deep_items)
            else:
                out['params'][key] = value

        return out

    def __repr__(self):
        '''Pretty-print this object'''

        class_name = self.__class__.__name__
        return '{:s}({:s})'.format(class_name,
                                   _pprint(self.get_params(deep=False)['params'],
                                           offset=len(class_name),),)

    def __init__(self):
        '''Base-class initialization'''
        self.dispatch = OrderedDict()
        self._state = dict()

        self.n_samples = 1

    def get_state(self, jam):
        '''Build the state object for a static transformer'''
        return dict()

    @contextmanager
    def _transform_state(self, jam):
        '''Initialize state for static transformers.'''
        self._state = self.get_state(jam)
        yield

    def _transform(self, jam):
        '''Apply the transformation to audio and annotations.

        The input jam is copied and modified, and returned
        contained in a list.

        Parameters
        ----------
        jam : jams.JAMS
            A single jam object to modify

        Returns
        -------
        jam_list : list
            A length-1 list containing `jam` after transformation

        See also
        --------
        core.load_jam_audio
        '''

        if not hasattr(jam.sandbox, 'muda'):
            raise RuntimeError('No muda state found in jams sandbox.')

        # We'll need a working copy of this object for modification purposes
        jam_w = copy.deepcopy(jam)

        with self._transform_state(jam_w):
            # Push our reconstructor onto the history stack
            jam_w.sandbox.muda['history'].append({'transformer': self.__serialize__,
                                                  'state': self._state})

            if hasattr(self, 'audio'):
                self.audio(jam_w.sandbox.muda)

            if hasattr(self, 'metadata'):
                self.metadata(jam_w.file_metadata)

            # Walk over the list of deformers
            for query, function in six.iteritems(self.dispatch):
                for matched_annotation in jam_w.search(namespace=query):
                    function(matched_annotation)

        return [jam_w]

    def transform(self, jam):
        '''Iterative transformation generator

        Generates up to a fixed number of transformed jams
        from a single input.

        Parameters
        ----------
        jam : jams.JAMS
            The jam to transform

        Generates
        ---------
        jam_out : jams.JAMS
            Iterator of transformed jams
        '''

        # Reset the state
        # XXX:    2015-03-25 22:37:15 by Brian McFee <brian.mcfee@nyu.edu>
        # this should be a context for state construction
        #   within _transform, we should iterate over states in contexts
        self._state = dict()

        i = 0
        while self.n_samples is None or i < self.n_samples:
            for jam_out in self._transform(jam):
                yield jam_out
                i = i + 1

    @property
    def __serialize__(self):
        '''Serializer'''

        data = self.get_params()
        data['__class__'] = data['__class__'].__name__
        return data

class Pipeline(object):
    '''Wrapper which allows multiple transformers to be chained together'''

    def __init__(self, *steps):
        '''Transformation pipeline.

        A given JAMS object will be transformed sequentially by
        each stage of the pipeline.

        Parameters
        ----------
        steps : argument array
            steps[i] is a tuple of `(name, Transformer)`

        Examples
        --------
        >>> P = muda.deformers.PitchShift(semitones=5)
        >>> T = muda.deformers.TimeStretch(speed=1.25)
        >>> Pipe = Pipeline( ('Pitch:maj3', P), ('Speed:1.25x', T) )
        >>> output = Pipe.transform(data)
        '''

        named_steps = dict(steps)
        names, transformers = zip(*steps)

        if len(named_steps) != len(steps):
            raise ValueError("Names provided are not unique: "
                             " {:s}".format(names,))

        # shallow copy of steps
        self.steps = list(zip(names, transformers))

        for t in transformers:
            if not isinstance(t, BaseTransformer):
                raise TypeError('{:s} is not a BaseTransformer'.format(t))

    def get_params(self):
        '''Get the parameters for this object.  Returns as a dict.'''

        out = {}
        out['__class__'] = self.__class__
        out['params'] = []

        for name, step in self.steps:
            out['params'].append([name, step.get_params(deep=True)])

        return out

    def __repr__(self):
        '''Pretty-print the object'''

        class_name = self.__class__.__name__
        return '{:s}({:s})'.format(class_name,
                                   _pprint(self.get_params(),
                                           offset=len(class_name),),)

    def __recursive_transform(self, jam, steps):
        '''A recursive transformation pipeline'''

        if len(steps) > 0:
            head_transformer = steps[0][1]
            for t_jam in head_transformer.transform(jam):
                for q in self.__recursive_transform(t_jam, steps[1:]):
                    yield q
        else:
            yield jam

    def transform(self, jam):
        '''Apply the sequence of transformations to a single jam object.

        Parameters
        ----------
        jam : jams.JAMS
            The jam object to transform

        Generates
        ---------
        jam_stream : iterable

        See also
        --------
        BaseTransformer.transform
        '''

        for output in self.__recursive_transform(jam, self.steps):
            yield output
<|endoftext|>