text:
'''
Text mode diagrams using UTF-8 characters and fancy colors.
'''

__author__ = 'Wijnand Modderman-Lenstra <maze@pyth0n.org>'
__copyright__ = 'Copyright 2014, maze.io labs'
__credits__ = ['Adam Tauber', 'Erik Rose', 'Jeff Quast', 'John-Paul Verkamp']
__license__ = 'MIT'

from collections import defaultdict
import curses
import locale
import math
import os
import sys
import warnings

# Optionally import numpy for faster arithmetic and curve filtering functions
try:
    import numpy as np
except ImportError:
    np = None

# Setup locale
if sys.platform == 'darwin':
    locale.setlocale(locale.LC_CTYPE, 'UTF-8')
else:
    locale.setlocale(locale.LC_ALL, '')

class Terminal(object):
    '''
    Terminal manipulation.
    '''

    def __init__(self):
        curses.setupterm()

    @property
    def colors(self):
        '''
        Get the number of colors supported by this terminal.
        '''
        number = curses.tigetnum('colors') or 0
        return 16 if number == 8 else number

    @property
    def encoding(self):
        '''
        Get the current terminal encoding.
        '''
        _, encoding = locale.getdefaultlocale()
        return encoding

    @property
    def height(self):
        '''
        Get the current terminal height.
        '''
        return self.size[1]

    @property
    def size(self):
        '''
        Get the current terminal size.
        '''
        for fd in range(3):
            cr = self._ioctl_GWINSZ(fd)
            if cr:
                break
        if not cr:
            try:
                fd = os.open(os.ctermid(), os.O_RDONLY)
                cr = self._ioctl_GWINSZ(fd)
                os.close(fd)
            except:
                pass

        if not cr:
            env = os.environ
            cr = (env.get('LINES', 25), env.get('COLUMNS', 80))

        return int(cr[1]), int(cr[0])

    @property
    def width(self):
        '''
        Get the current terminal width.
        '''
        return self.size[0]

    def _ioctl_GWINSZ(self, fd):
        '''
        Internal function that will try to request the ``TIOCGWINSZ`` against
        the selected file descriptor ``fd``.
        '''
        try:
            import fcntl, termios, struct, os
            return struct.unpack(
                'hh',
                fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234')
            )
        except:
            return None

    def color(self, index):
        '''
        Get the escape sequence for indexed color ``index``. The ``index`` is
        a color index in the 256 color space. The color space consists of:

        * 0x00-0x0f: default EGA colors
        * 0x10-0xe7: 6x6x6 RGB cubes
        * 0xe8-0xff: gray scale ramp
        '''
        if self.colors == 16:
            if index >= 8:
                return self.csi('bold') + self.csi('setaf', index - 8)
            else:
                return self.csi('sgr0') + self.csi('setaf', index)
        else:
            return self.csi('setaf', index)

    def csi(self, capname, *args):
        '''
        Returns the escape sequence for the selected Control Sequence
        Introducer.
        '''
        value = curses.tigetstr(capname)
        if value is None:
            return ''
        else:
            return curses.tparm(value, *args)

    def csi_wrap(self, value, capname, *args):
        '''
        Returns a value wrapped in the selected Control Sequence Introducer
        and resets the attributes afterwards.
        '''
        return u''.join([
            self.csi(capname, *args),
            value,
            self.csi('sgr0'),
        ])

# Bar characters (8 per map)
H_BAR = [(0x258f, 0x258e, 0x258d, 0x258c, 0x258b, 0x258a, 0x2589, 0x2588),
         (0x00a0, 0x2589, 0x258a, 0x258b, 0x258c, 0x258d, 0x258e, 0x258f)]
V_BAR = [(0x2581, 0x2582, 0x2583, 0x2584, 0x2585, 0x2586, 0x2587, 0x2588),
         (0x2588, 0x2587, 0x2586, 0x2585, 0x2584, 0x2583, 0x2582, 0x00a0)]

def filter_symlog(y, base=10.0):
    '''
    Symmetrical logarithmic scale.

    Optional arguments:

    *base*:
        The base of the logarithm.
    '''
    log_base = np.log(base)
    sign = np.sign(y)
    logs = np.log(np.abs(y) / log_base)
    return sign * logs

def filter_savitzky_golay(y, window_size=5, order=2, deriv=0, rate=1):
    '''
    Smooth (and optionally differentiate) data with a Savitzky-Golay filter.
    '''
    try:
        window_size = np.abs(np.int(window_size))
        order = np.abs(np.int(order))
    except ValueError:
        raise ValueError('window_size and order must be integers')

    if window_size % 2 != 1 or window_size < 1:
        raise ValueError('window_size size must be a positive odd number')
    if window_size < order + 2:
        raise ValueError('window_size is too small for the polynomials order')

    order_range = range(order + 1)
    half_window = (window_size - 1) // 2

    # precompute limits
    minimum = np.min(y)
    maximum = np.max(y)
    # precompute coefficients
    b = np.mat([
        [k ** i for i in order_range]
        for k in range(-half_window, half_window + 1)
    ])
    m = np.linalg.pinv(b).A[deriv] * rate ** deriv * math.factorial(deriv)
    # pad the signal at the extremes with values taken from the original signal
    firstvals = y[0] - np.abs(y[1:half_window+1][::-1] - y[0])
    lastvals = y[-1] + np.abs(y[-half_window-1:-1][::-1] - y[-1])
    y = np.concatenate((firstvals, y, lastvals))
    return np.clip(
        np.convolve(m[::-1], y, mode='valid'),
        minimum,
        maximum,
    )

# Graph filter functions
FUNCTION = dict(
    log=filter_symlog,
    smooth=filter_savitzky_golay,
)
FUNCTION_CONSTANT = {
    'e': math.e,
    '-e': -math.e,
    'pi': math.pi,
    '-pi': -math.pi,
}

# Indexed palettes
PALETTE = dict(
    red={
        0x010: (1, 9),
        0x100: [(16 + (o * 36)) for o in range(1, 6)] + [9],
    },
    green={
        0x010: (2, 10),
        0x100: [(16 + (o * 6)) for o in range(1, 6)] + [10],
    },
    blue={
        0x010: (4, 12),
        0x100: range(17, 22) + [12],
    },
    grey={
        0x010: (8, 7, 15),
        0x100: range(232, 257) + [15],
    },
    spectrum={
        0x010: [14, 6, 2, 10, 11, 3, 9],
        0x100: [
            51 - x for x in range(6)                # blue -> green
        ] + [
            46 + x * 36 for x in range(6)          # green -> yellow
        ] + [
            226 - x * 6 for x in range(6)          # yellow -> red
        ]
    },
)
PALETTE.update(dict(
    gray=PALETTE['grey'],
    default=PALETTE['spectrum'],
))

class Point(object):
    '''
    Holds a single, two-dimensional point.
    '''

    def __init__(self, coordinates):
        self.x = coordinates[0]
        self.y = coordinates[1]

    def __repr__(self):
        return 'Point(%r, %r)' % (self.x, self.y)

    def copy(self):
        '''
        Returns a fresh copy of the current point.
        '''
        return Point((self.x, self.y))

class Screen(object):
    '''
    Off-screen render buffer.
    '''
    def __init__(self, size, encoding=None, extend_x=False, extend_y=False):
        if isinstance(size, Point):
            self.size = size.copy()
        else:
            self.size = Point(size)

        self.term = Terminal()
        self.encoding = encoding or self.term.encoding
        self.extend_x = extend_x
        self.extend_y = extend_y

        defaultdictint = lambda: defaultdict(int)
        self.canvas = defaultdict(defaultdictint)

    @property
    def width(self):
        '''
        Get the buffer width.
        '''
        return self.size.x

    @property
    def height(self):
        '''
        Get the buffer height.
        '''
        return self.size.y

    def __contains__(self, point):
        '''
        Check if a point has a value.
        '''
        if not isinstance(point, Point):
            point = Point(point)

        if not point.y in self.canvas:
            return False
        else:
            return point.x in self.canvas[point.y]

    def __repr__(self):
        return '%s(%d, %d)' % (self.__class__.__name__, self.width, self.height)

    def __setitem__(self, point, value):
        '''
        Set a point value.
        '''
        if not isinstance(point, Point):
            point = Point(point)

        if point.y > self.size.y:
            if self.extend_y:
                self.size.y = point.y
            else:
                raise OverflowError('%r overflow y = %d > %d' % (
                    self,
                    point.y,
                    self.size.y,
                ))

        if point.x > self.size.x:
            if self.extend_x:
                self.size.x = point.x
            else:
                raise OverflowError('%r overflow x = %d > %d' % (
                    self,
                    point.x,
                    self.size.x,
                ))

        self.canvas[point.y][point.x] = value

    def __getitem__(self, point):
        '''
        Get a point value or None.
        '''
        if not isinstance(point, Point):
            point = Point(point)
        return self.canvas[point.y][point.x]

class Graph(object):
    '''
    Base class for graphs.
    '''

    def __init__(self, size, option):
        self.size = size
        self.option = option

        # Internal cycle duty counter
        self.cycle = 0L
        self.lines = 0L
        self.term = Terminal()

    def update(self, points):
        if np is None:
            if self.option.function:
                warnings.warn('numpy not available, function ignored')
            self.points = points
            self.minimum = min(self.points)
            self.maximum = max(self.points)
            self.extents = (self.maximum - self.minimum)

        else:
            self.points = self.apply_function(points)
            self.minimum = np.min(self.points)
            self.maximum = np.max(self.points)
            self.extents = (self.maximum - self.minimum)

    def color_ramp(self, size):
        '''
        Generate a color ramp for the current screen height.
        '''
        color = PALETTE.get(self.option.palette, {}).get(self.term.colors, None)
        color_ramp = []
        if color is not None:
            ratio = len(color) / float(size)
            for i in range(int(size)):
                color_ramp.append(self.term.color(color[int(ratio * i)]))

        return color_ramp

    def human(self, size, base=1000, units=' kMGTZ'):
        '''
        Converts the input ``size`` to human readable, short form.
        '''
        sign = '+' if size >= 0 else '-'
        size = abs(size)
        if size < 1000:
            return '%s%d' % (sign, size)
        for i, suffix in enumerate(units):
            unit = 1000 ** (i + 1)
            if size < unit:
                return ('%s%.01f%s' % (
                    sign,
                    size / float(unit) * base,
                    suffix,
                )).strip()
        raise OverflowError

    def apply_function(self, points):
        '''
        Run the filter function on the provided points.
        '''
        if not self.option.function:
            return points

        if np is None:
            print 'no numpy'
            raise ImportError('numpy is not available')

        if ':' in self.option.function:
            function, arguments = self.option.function.split(':', 1)
            arguments = arguments.split(',')
        else:
            function = self.option.function
            arguments = []

        # Resolve arguments
        arguments = map(self._function_argument, arguments)

        # Resolve function
        filter_function = FUNCTION.get(function)

        if filter_function is None:
            raise TypeError('Invalid function "%s"' % (function,))

        else:
            # We wrap in ``list()`` to consume generators and iterators, as
            # ``np.array`` doesn't do this for us.
            return filter_function(np.array(list(points)), *arguments)

    def _function_argument(self, value):
        if value in FUNCTION_CONSTANT:
            return FUNCTION_CONSTANT[value]
        else:
            return float(value)

    def line(self, p1, p2, resolution=1):
        '''
        Resolve the points to make a line between two points.
        '''
        xdiff = max(p1.x, p2.x) - min(p1.x, p2.x)
        ydiff = max(p1.y, p2.y) - min(p1.y, p2.y)
        xdir = [-1, 1][int(p1.x <= p2.x)]
        ydir = [-1, 1][int(p1.y <= p2.y)]
        r = int(round(max(xdiff, ydiff)))
        if r == 0:
            return

        for i in range((r + 1) * resolution):
            x = p1.x
            y = p1.y

            if xdiff:
                x += (float(i) * xdiff) / r * xdir / resolution
            if ydiff:
                y += (float(i) * ydiff) / r * ydir / resolution

            yield Point((x, y))

    @property
    def maximum_points(self):
        raise NotImplementedError()

    def render(self, stream):
        '''
        Render the graph to the selected output stream.
        '''
        raise NotImplementedError()

    def round(self, value):
        '''
        Get an integer value for the input value.
        '''
        if isinstance(value, (int, long)):
            return value

        else:
            return long(round(value))

    def set_text(self, point, text):
        '''
        Set a text value in the screen canvas.
        '''
        if not self.option.legend:
            return

        if not isinstance(point, Point):
            point = Point(point)

        for offset, char in enumerate(str(text)):
            self.screen.canvas[point.y][point.x + offset] = char

class AxisGraphScreen(Screen):
    '''
    Base class for axial graph buffers. The buffer size is larger than the
    actual screen size, because we have 8 pixels per character. The number
    of pixels per character is 2 horizontal and 4 vertical pixels.
    '''

    @property
    def width(self):
        return self.size.x * 2

    @property
    def height(self):
        return self.size.y * 4

class AxisGraph(Graph):
    '''
    Base class for axial graphs.
    '''

    # Graph characters, using braille characters
    base = 0x2800
    pixels = ((0x01, 0x08),
              (0x02, 0x10),
              (0x04, 0x20),
              (0x40, 0x80))

    def __init__(self, size, option):
        super(AxisGraph, self).__init__(size, option)

        self.size = Point((
            size.x or self.term.width,
            size.y or 10
        ))

    def render(self, stream):
        encoding = self.option.encoding or self.term.encoding

        if self.option.color:
            ramp = self.color_ramp(self.size.y)[::-1]
        else:
            ramp = None

        if self.cycle >= 1 and self.lines:
            stream.write(self.term.csi('cuu', self.lines))

        zero = int(self.null / 4)  # Zero crossing
        lines = 0
        for y in range(self.screen.size.y):
            if y == zero and self.size.y > 1:
                stream.write(self.term.csi('smul'))
            if ramp:
                stream.write(ramp[y])

            for x in range(self.screen.size.x):
                point = Point((x, y))
                if point in self.screen:
                    value = self.screen[point]
                    if isinstance(value, int):
                        stream.write(unichr(self.base + value).encode(encoding))
                    else:
                        stream.write(self.term.csi('sgr0'))
                        stream.write(self.term.csi_wrap(
                            unicode(value).encode(encoding),
                            'bold'
                        ))
                        if y == zero and self.size.y > 1:
                            stream.write(self.term.csi('smul'))
                        if ramp:
                            stream.write(ramp[y])
                else:
                    stream.write(' ')

            if y == zero and self.size.y > 1:
                stream.write(self.term.csi('rmul'))
            if ramp:
                stream.write(self.term.csi('sgr0'))

            stream.write('\n')
            lines += 1
        stream.flush()

        self.cycle = self.cycle + 1
        self.lines = lines

    @property
    def normalised(self):
        if np is None:
            return self._normalised_python()
        else:
            return self._normalised_numpy()

    def _normalised_numpy(self):
        dx = (self.screen.width / float(len(self.points)))
        oy = (self.screen.height)
        points = np.array(self.points) - self.minimum
        points = points * 4.0 / self.extents * self.size.y
        for x, y in enumerate(points):
            yield Point((
                dx * x,
                min(oy, oy - y),
            ))

    def _normalised_python(self):
        dx = (self.screen.width / float(len(self.points)))
        oy = (self.screen.height)
        for x, point in enumerate(self.points):
            y = (point - self.minimum) * 4.0 / self.extents * self.size.y
            yield Point((
                dx * x,
                min(oy, oy - y),
            ))

    @property
    def maximum_points(self):
        if hasattr(self, 'screen'):
            return self.screen.width
        else:
            return self.size.x

    @property
    def null(self):
        if not self.option.axis:
            return -1
        else:
            return self.screen.height - (
                -self.minimum * 4.0 / self.extents * self.size.y
            )

    def set(self, point):
        if not isinstance(point, Point):
            point = Point(point)

        rx = self.round(point.x)
        ry = self.round(point.y)

        item = Point((rx >> 1, min(ry >> 2, self.size.y)))
        self.screen[item] |= self.pixels[ry & 3][rx & 1]

    def unset(self, point):
        if not isinstance(point, Point):
            point = Point(point)

        x, y = self.round(point.x) >> 1, self.round(point.y) >> 2

        if (x, y) not in self.screen:
            return

        if isinstance(self.screen[y][x], (int, long)):
            self.screen[(x, y)] &= ~self.pixels[y & 3][x & 1]

        else:
            del self.screen[(x, y)]

        if not self.canvas.get(y):
            del self.screen[y]

    def update(self, points):
        super(AxisGraph, self).update(points)

        self.screen = AxisGraphScreen(self.size)

        # Plot lines between the points
        prev = Point((0, self.null))
        for curr in self.normalised:
            for point in self.line(prev, curr):
                self.set(point)
            prev = curr

        zero = int(self.null / 4)  # Zero crossing
        if self.size.y > 1:
            self.set_text(Point((0, zero)), '0')
            self.set_text(Point((0, 0)), self.human(self.maximum))
            self.set_text(Point((0, self.size.y - 1)), self.human(self.minimum))

class BarGraph(Graph):
    '''
    Base class for bar graphs.
    '''

    @property
    def normalised(self):
        for point in self.points:
            yield (point - self.minimum) / self.extents * self.scale

class HorizontalBarGraph(BarGraph):
    '''
    Horizontal bar graph.
    '''

    def __init__(self, size, option):
        super(BarGraph, self).__init__(size, option)

        if size.y:
            warnings.warn('Ignoring height on horizontal bar graph')

        self.size = Point((
            size.x or self.term.width,
            1,
        ))

        self.screen = Screen(
            Point((self.size.x, 1)),
            extend_y=True,
        )

        # Select block characters
        if self.option.reverse:
            self.blocks = H_BAR[1]
        else:
            self.blocks = H_BAR[0]

    def bar(self, size, y):
        full, frac = divmod(self.round(size * 8), 8)

        x = 0
        if self.option.reverse:
            for x in range(full):
                xr = self.screen.size.x - x
                self.screen[(xr, y)] = self.blocks[-1]
            if frac:
                x = x + 1 if x else x
                xr = self.screen.size.x - x
                self.screen[(xr, y)] = self.blocks[frac]
        else:
            for x in range(full):
                self.screen[(x, y)] = self.blocks[-1]
            if frac:
                x = x + 1 if x else x
                self.screen[(x, y)] = self.blocks[frac]

    @property
    def maximum_points(self):
        if self.option.height:
            return self.option.height
        else:
            return 10

    @property
    def scale(self):
        return float(self.screen.width)

    def render(self, stream):
        encoding = self.option.encoding or self.term.encoding()

        if self.option.color:
            ramp = self.color_ramp(self.screen.size.x)
            if self.option.reverse:
                ramp = ramp[::-1]
        else:
            ramp = None

        if self.cycle >= 1:
            stream.write(self.term.csi('cuu', self.lines))

        lines = 0
        stream.write(self.term.csi('el'))
        if self.option.legend:
            minimum_text = self.human(self.minimum)
            maximum_text = self.human(self.maximum)
            minimum_text = minimum_text.ljust(
                self.screen.width - len(maximum_text)
            )
            stream.write(self.term.csi_wrap(
                ''.join([minimum_text, maximum_text]),
                'bold',
            ))
            stream.write('\n')
            lines += 1

        for y in range(self.screen.size.y):
            prev_color = ''
            stream.write(self.term.csi('el'))
            for x in range(self.screen.size.x):
                if ramp:
                    curr_color = ramp[x]
                    if not self.option.reverse and curr_color != prev_color:
                        stream.write(curr_color)
                        prev_color = curr_color

                point = Point((x, y))
                if point in self.screen:
                    value = self.screen[point]
                    if isinstance(value, int):
                        if self.option.reverse:
                            stream.write(curr_color)
                            stream.write(self.term.csi('rev'))
                        stream.write(unichr(value).encode(encoding))
                    else:
                        stream.write(self.term.csi('sgr0'))
                        stream.write(self.term.csi_wrap(
                            unicode(value).encode(encoding),
                            'bold'
                        ))
                        if ramp:
                            stream.write(curr_color)
                        if self.option.reverse:
                            stream.write(self.term.csi('rev'))
                else:
                    stream.write(' ')

            if ramp:
                stream.write(self.term.csi('sgr0'))
            stream.write('\n')
            lines += 1

        self.cycle = self.cycle + 1
        self.lines = lines

    def update(self, points):
        super(HorizontalBarGraph, self).update(points)

        # Clear screen
        self.screen = Screen(
            self.screen.size.copy(),
            extend_y=True,
        )

        # Plot bars for each line
        for y, size in enumerate(self.normalised):
            self.bar(size, y)

class VerticalBarGraph(BarGraph):
    '''
    Vertical bar graph.
    '''

    def __init__(self, size, option):
        super(BarGraph, self).__init__(size, option)

        if size.x:
            warnings.warn('Ignoring width on horizontal bar graph')

        # Select block characters
        if self.option.reverse:
            self.blocks = V_BAR[1]
        else:
            self.blocks = V_BAR[0]

    def bar(self, size, x):
        full, frac = divmod(self.round(size * 8), 8)

        y = 0
        if self.option.reverse:
            for y in range(full):
                self.screen[(x, y)] = self.blocks[-1]
            if frac:
                y = y + 1 if y else y
                self.screen[(x, y)] = self.blocks[-frac]
        else:
            for y in range(self.size.y, self.size.y - full - 1, -1):
                self.screen[(x, y)] = self.blocks[-1]
            if frac:
                y = self.size.y - full - 1
                self.screen[(x, y)] = self.blocks[frac]

    @property
    def maximum_points(self):
        return self.term.width

    @property
    def scale(self):
        return float(self.screen.height)

    def render(self, stream):
        encoding = self.option.encoding or self.term.encoding()

        if self.option.color:
            ramp = self.color_ramp(self.size.y)
            if not self.option.reverse:
                ramp = ramp[::-1]
        else:
            ramp = []

        if self.cycle >= 1:
            stream.write(self.term.csi('cuu', self.lines))

        lines = 0
        for y in range(self.size.y):
            if ramp:
                stream.write(ramp[y])
            for x in range(self.screen.size.x):
                point = Point((x, y))
                if point in self.screen:
                    value = self.screen[point]
                    if isinstance(value, int):
                        if self.option.reverse:
                            if ramp:
                                stream.write(ramp[y])
                            stream.write(self.term.csi_wrap(
                                unichr(value),
                                'rev'
                            ))
                        else:
                            stream.write(unichr(value).encode(encoding))
                    else:
                        stream.write(self.term.csi('sgr0'))
                        stream.write(self.term.csi_wrap(
                            unicode(value).encode(encoding),
                            'bold'
                        ))
                        if ramp:
                            stream.write(ramp[y])

                else:
                    stream.write(' ')

            if ramp:
                stream.write(self.term.csi('sgr0'))
            stream.write('\n')
            lines += 1

        self.cycle = self.cycle + 1
        self.lines = lines

    def update(self, points):
        if self.option.reverse:
            points = points[::-1]

        maximum_width = self.maximum_points
        if len(points) > maximum_width:
            spoints = points[-maximum_points:]

        # If the legend is enabled, and we have sufficient room to shift the
        # columns to the right, we do so.
        elif len(points) < (maximum_width - 6) and self.option.legend and self.cycle == 0:
            for x in range(6):
                points.insert(0, min(points))

        super(VerticalBarGraph, self).update(points)

        self.size = Point((
            len(self.points),
            self.option.size.y or 10,
        ))

        self.screen = Screen(
            Point((1, self.size.y)),
            extend_x=True,
            extend_y=True,
        )

        # Plot bars for each line
        for x, size in enumerate(self.normalised):
            self.bar(size, x)

        # Plot legend, if there is sufficient space
        if self.size.y > 1 and self.option.legend:
            self.set_text(Point((0, 0)), self.human(self.maximum))
            self.set_text(Point((0, self.size.y - 1)), self.human(self.minimum))

def usage_function(parser):
    '''
    Shows usage and available curve functions.
    '''
    parser.print_usage()
    print ''
    print 'available functions:'
    for function in sorted(FUNCTION):
        doc = FUNCTION[function].__doc__.strip().splitlines()[0]
        print '    %-12s %s' % (function + ':', doc)

    return 0

def usage_palette(parser):
    '''
    Shows usage and available palettes.
    '''
    parser.print_usage()
    print ''
    print 'available palettes:'
    for palette in sorted(PALETTE):
        print '    %-12s' % (palette,)

    return 0

def run():
    '''
    Main entrypoint if invoked via the command line.
    '''
    import argparse

    parser = argparse.ArgumentParser()

    group = parser.add_argument_group('optional drawing mode')
    group.add_argument(
        '-G', '--graph',
        dest='mode', action='store_const', const='g',
        help='axis drawing mode (default)',
    )
    group.add_argument(
        '-H', '--horizontal-bars',
        dest='mode', action='store_const', const='h',
        help='horizontal drawing mode',
    )
    group.add_argument(
        '-V', '--vertical-bars',
        dest='mode', action='store_const', const='v',
        help='vertical drawing mode',
    )

    group = parser.add_argument_group('optional drawing arguments')
    group.add_argument(
        '-a', '--axis',
        dest='axis', action='store_const', const=True, default=True,
        help='draw axis (default: yes)',
    )
    group.add_argument(
        '-A', '--no-axis',
        dest='axis', action='store_const', const=False,
        help="don't draw axis",
    )
    group.add_argument(
        '-c', '--color',
        dest='color', action='store_const', const=True, default=True,
        help='use colors (default: yes)',
    )
    group.add_argument(
        '-C', '--no-color',
        dest='color', action='store_const', const=False,
        help="don't use colors",
    )
    group.add_argument(
        '-l', '--legend',
        dest='legend', action='store_const', const=True, default=True,
        help='draw y-axis legend (default: yes)',
    )
    group.add_argument(
        '-L', '--no-legend',
        dest='legend', action='store_const', const=False,
        help="don't draw y-axis legend",
    )
    group.add_argument(
        '-f', '--function',
        default=None, metavar='function',
        help='curve manipulation function, use "help" for a list',
    )
    group.add_argument(
        '-p', '--palette',
        default='default', metavar='palette',
        help='palette name, use "help" for a list',
    )
    group.add_argument(
        '-x', '--width',
        default=0, type=int, metavar='characters',
        help='drawing width (default: auto)',
    )
    group.add_argument(
        '-y', '--height',
        default=0, type=int, metavar='characters',
        help='drawing height (default: auto)',
    )
    group.add_argument(
        '-r', '--reverse',
        default=False, action='store_true',
        help='reverse draw graph',
    )

    group = parser.add_argument_group('optional input and output arguments')
    group.add_argument(
        '-b', '--batch',
        default=False, action='store_true',
        help='batch mode (default: no)',
    )
    group.add_argument(
        '-s', '--sleep',
        default=0, type=float,
        help='batch poll sleep time (default: none)',
    )
    group.add_argument(
        '-i', '--input',
        default='-', metavar='file',
        help='input file (default: stdin)',
    )
    group.add_argument(
        '-o', '--output',
        default='-', metavar='file',
        help='output file (default: stdout)',
    )
    group.add_argument(
        '-e', '--encoding',
        dest='encoding', default='',
        help='output encoding (default: auto)',
    )

    option = parser.parse_args()

    if option.function == 'help':
        return usage_function(parser)

    if option.palette == 'help':
        return usage_palette(parser)

    option.mode = option.mode or 'g'
    option.size = Point((option.width, option.height))

    if option.input in ['-', 'stdin']:
        istream = sys.stdin
    else:
        istream = open(option.input, 'r')

    if option.output in ['-', 'stdout']:
        ostream = sys.stdout
    else:
        ostream = open(option.output, 'r')

    option.encoding = option.encoding or Terminal().encoding

    if option.mode == 'g':
        engine = AxisGraph(option.size, option)

    elif option.mode == 'h':
        engine = HorizontalBarGraph(option.size, option)

    elif option.mode == 'v':
        engine = VerticalBarGraph(option.size, option)

    else:
        parser.error('invalid mode')
        return 1

    if option.batch:
        import select
        import time

        option.sleep = max(0.01, option.sleep)
        points = []
        while True:
            try:
                if select.select([istream.fileno()], [], [], option.sleep):
                    for point in istream.readline().split():
                        try:
                            points.append(float(point))
                        except ValueError:
                            pass

                    # We need at least two data points to draw a graph
                    if len(points) > 1:
                        points = points[-engine.maximum_points:]
                        engine.update(points)
                        engine.render(ostream)
                        time.sleep(option.sleep)

            except KeyboardInterrupt:
                break

    else:
        points = []
        for line in istream:
            for point in line.split():
                try:
                    points.append(float(point))
                except ValueError:
                    pass

        engine.update(points)
        engine.render(ostream)

if __name__ == '__main__':
    sys.exit(run())

max_stars_repo_name: mrjink/diagram

max_issues_repo_name: mrjink/diagram

max_forks_repo_name: mrjink/diagram


text:
"""Here we mostly test if it works at all."""
from skultrafast.dataset import TimeResSpec, PolTRSpec
from skultrafast.data_io import load_example
import numpy as np
from numpy.testing import assert_almost_equal

wl, t, data = load_example()

def test_integrate():
    ds = TimeResSpec(wl, t, data)
    ds.wn_i(15000, 20000)

def test_methods():
    ds = TimeResSpec(wl, t, data)
    bds = ds.bin_freqs(300)
    ds2 = TimeResSpec(1e7/wl, t, data, freq_unit='cm', disp_freq_unit='cm')
    bds2 = ds2.bin_freqs(50)
    assert(np.all(np.isfinite(bds2.data)))

    assert (len(bds.wavelengths) == 300)
    nds = ds.cut_freq(400, 600)
    assert (np.all(nds.wavelengths > 600))
    nds = ds.cut_time(-100, 1)
    assert (np.all(nds.t > .99))
    nds = ds.bin_times(5)
    assert (nds.t.size == np.ceil(ds.t.size / 5))
    ds.mask_freqs([(400, 600)])
    assert (np.all(ds.data.mask[:, ds.wl_idx(550)]))
    ds2 = ds.scale_and_shift(2, t_shift=1, wl_shift=10)
    assert_almost_equal(2*ds.data, ds2.data)
    assert_almost_equal(ds.t+1, ds2.t)
    assert_almost_equal(ds.wavelengths+10, ds2.wavelengths)
    assert_almost_equal(1e7/ds2.wavelengths, ds2.wavenumbers)

def test_est_disp():
    ds = TimeResSpec(wl, t, data)
    ds.auto_plot = False
    for s in ['abs', 'diff', 'gauss_diff', 'max']:
        ds.estimate_dispersion(heuristic=s)

def test_fitter():
    ds = TimeResSpec(wl, t, data)
    x0 = [0.1, 0.1, 1, 1000]
    out = ds.fit_exp(x0)

def test_error_calc():
    ds = TimeResSpec(wl, t, data)
    x0 = [0.1, 0.1, 1, 1000]
    out = ds.fit_exp(x0)
    out.calculate_stats()

def test_das_plots():
    ds = TimeResSpec(wl, t, data)
    x0 = [0.1, 0.1, 1, 1000]
    out = ds.fit_exp(x0)
    ds.plot.das()
    ds.plot.edas()

def test_das_pol_plots():
    ds = TimeResSpec(wl, t, data)
    pds = PolTRSpec(ds, ds) # fake pol
    x0 = [0.1, 0.1, 1, 1000]
    out = pds.fit_exp(x0)

    pds.plot.das()
    pds.plot.edas()

def test_sas_pol_plots():
    from skultrafast.kinetic_model import Model
    ds = TimeResSpec(wl, t, data)
    pds = PolTRSpec(ds, ds) # fake pol
    x0 = [0.1, 0.1, 1, 1000]
    out = pds.fit_exp(x0)
    m = Model()
    m.add_transition('S1', 'S1*', 'k1')
    m.add_transition('S1', 'zero', 'k2')

    pds.plot.sas(m)

def test_sas():
    from skultrafast.kinetic_model import Model
    ds = TimeResSpec(wl, t, data)
    x0 = [0.1, 0.1, 1, 1000]
    out = ds.fit_exp(x0)
    m = Model()
    m.add_transition('S1', 'S1*', 'k1')
    m.add_transition('S1', 'zero', 'k2')
    out.make_sas(m, {})

    m2 = Model()
    m2.add_transition('S1', 'S1*', 'k1', 'qy1')
    m2.add_transition('S1', 'zero', 'k2')
    out.make_sas(m2, {'qy1': 0.5})

def test_merge():
    ds = TimeResSpec(wl, t, data)
    nds = ds.merge_nearby_channels(10)
    assert (nds.wavelengths.size < ds.wavelengths.size)

def test_pol_tr():
    ds = TimeResSpec(wl, t, data)
    ds2 = TimeResSpec(wl, t, data)
    ps = PolTRSpec(para=ds, perp=ds2)
    out = ps.bin_freqs(10)
    assert (out.para.wavenumbers.size == 10)
    assert (out.perp.wavenumbers.size == 10)
    assert_almost_equal(out.perp.data, out.para.data)
    ps.subtract_background()
    ps.mask_freqs([(400, 550)])
    print(ps.para.data.mask, ps.para.data.mask[1, ps.para.wl_idx(520)])

    assert (ps.para.data.mask[1, ps.para.wl_idx(520)])
    out = ps.cut_freq(400, 550)
    assert (np.all(out.para.wavelengths >= 550))
    assert (np.all(out.perp.wavelengths >= 550))
    ps.bin_times(6)
    ps.scale_and_shift(1, 0.5)
    ps.copy()

def test_plot():
    ds = TimeResSpec(wl, t, data)
    ds = ds.bin_freqs(50)
    ds.plot.trans([550])
    ds.plot.spec([2, 10])
    ds.plot.trans_integrals((1e7 / 550, 1e7 / 600))
    ds.plot.trans_integrals((1e7 / 600, 1e7 / 500))
    ds.plot.trans([550], norm=1)
    ds.plot.trans([550], norm=1, marker='o')
    ds.plot.map(plot_con=0)
    ds.plot.svd()

def test_concat():
    ds = TimeResSpec(wl, t, data)
    ds = ds.bin_freqs(50)
    n = ds.wavelengths.size //2
    ds1 = ds.cut_freq(ds.wavelengths[n], np.inf)
    ds2 = ds.cut_freq(-np.inf, ds.wavelengths[n])
    dsc = ds1.concat_datasets(ds2)
    assert(np.allclose(dsc.data, ds.data))
    pol_ds = PolTRSpec(ds, ds)
    a = PolTRSpec(ds1, ds1)
    b = PolTRSpec(ds2, ds2)
    pol_dsc = a.concat_datasets(b)
    for p in 'para', 'perp', 'iso':
        assert(np.allclose(getattr(pol_dsc, p).data, getattr(pol_ds, p).data))

def test_pol_plot():
    ds = TimeResSpec(wl, t, data)
    ds = ds.bin_freqs(50)
    ds = PolTRSpec(para=ds, perp=ds)
    ds = ds.bin_freqs(50)
    ds.plot.trans([550])
    ds.plot.spec([2, 10])
    ds.plot.trans([550], norm=1)
    ds.plot.trans([550], norm=1, marker='o')

ext: py

max_stars_repo_name: cZahn/skultrafast

max_issues_repo_name: cZahn/skultrafast

max_forks_repo_name: cZahn/skultrafast

reason: import numpy,from numpy


text:
# from 3ML

import matplotlib.pyplot as plt
import copy
from operator import itemgetter, attrgetter
import numpy as np

from cosmogrb.utils.plotting import step_plot

class IntervalsDoNotOverlap(RuntimeError):
    pass

class IntervalsNotContiguous(RuntimeError):
    pass

class TimeInterval(object):
    def __init__(self, start, stop, counts=None, dead_time=0, swap_if_inverted=False):

        self._start = float(start)
        self._stop = float(stop)

        self._counts = counts
        self._dead_time = dead_time

        self._exposure = (self._stop - self._start) - dead_time

        # Note that this allows to have intervals of zero duration

        if self._stop < self._start:

            if swap_if_inverted:

                self._start = stop
                self._stop = start

            else:

                raise RuntimeError(
                    "Invalid time interval! TSTART must be before TSTOP and TSTOP-TSTART >0. "
                    "Got tstart = %s and tstop = %s" % (start, stop)
                )

    def __add__(self, number):
        """
        Return a new time interval equal to the original time interval shifted to the right by number

        :param number: a float
        :return: a new TimeInterval instance
        """

        return TimeInterval(self._start + number, self._stop + number)

    def __sub__(self, number):
        """
        Return a new time interval equal to the original time interval shifted to the left by number

        :param number: a float
        :return: a new TimeInterval instance
        """

        return TimeInterval(self._start - number, self._stop - number)

    @property
    def start(self):
        return self._start

    @property
    def stop(self):
        return self._stop

    @property
    def counts(self):
        return self._counts

    @property
    def dead_time(self):
        return self._dead_time

    @property
    def rate(self):
        if self._counts is not None:
            return self._counts / self._exposure

    @property
    def width(self):

        return self._stop - self._start

    @property
    def exposure(self):
        return self._exposure

    @property
    def mid_point(self):

        return (self._start + self._stop) / 2.0

    def __repr__(self):

        return " interval %s - %s (width: %s)" % (self.start, self.stop, self.width,)

    def intersect(self, interval):
        """
        Returns a new time interval corresponding to the intersection between this interval and the provided one.

        :param interval: a TimeInterval instance
        :type interval: Interval
        :return: new interval covering the intersection
        :raise IntervalsDoNotOverlap : if the intervals do not overlap
        """

        if not self.overlaps_with(interval):
            raise IntervalsDoNotOverlap(
                "Current interval does not overlap with provided interval"
            )

        new_start = max(self._start, interval.start)
        new_stop = min(self._stop, interval.stop)

        return TimeInterval(new_start, new_stop)

    def merge(self, interval):
        """
        Returns a new interval corresponding to the merge of the current and the provided time interval. The intervals
        must overlap.

        :param interval: a TimeInterval instance
         :type interval : Interval
        :return: a new TimeInterval instance
        """

        if self.overlaps_with(interval):

            new_start = min(self._start, interval.start)
            new_stop = max(self._stop, interval.stop)

            if (self._counts is not None) and (interval.counts is not None):

                new_counts = self._counts + interval.counts

            else:

                new_counts = None

            new_dead_time = self._dead_time + interval.dead_time

            return TimeInterval(new_start, new_stop, new_counts, new_dead_time)

        else:

            raise IntervalsDoNotOverlap("Could not merge non-overlapping intervals!")

    def overlaps_with(self, interval):
        """
        Returns whether the current time interval and the provided one overlap or not

        :param interval: a TimeInterval instance
        :type interval: Interval
        :return: True or False
        """

        if interval.start == self._start or interval.stop == self._stop:

            return True

        elif interval.start > self._start and interval.start < self._stop:

            return True

        elif interval.stop > self._start and interval.stop < self._stop:

            return True

        elif interval.start < self._start and interval.stop > self._stop:

            return True

        else:

            return False

    def to_string(self):
        """
        returns a string representation of the time interval that is like the
        argument of many interval reading funcitons

        :return:
        """

        return "%f-%f" % (self.start, self.stop)

    def __eq__(self, other):

        if not isinstance(other, TimeInterval):

            # This is needed for things like comparisons to None or other objects.
            # Of course if the other object is not even a TimeInterval, the two things
            # cannot be equal

            return False

        else:

            return self.start == other.start and self.stop == other.stop

class TimeIntervalSet(object):
    """
    A set of intervals

    """

    def __init__(self, list_of_intervals=()):

        self._intervals = list(list_of_intervals)

    @classmethod
    def from_starts_and_stops(cls, starts, stops, counts=None, dead_time=None):
        """
        Builds a TimeIntervalSet from a list of start and stop times:

        start = [-1,0]  ->   [-1,0], [0,1]
        stop =  [0,1]

        :param starts:
        :param stops:
        :return:
        """

        assert len(starts) == len(stops), (
            "starts length: %d and stops length: %d must have same length"
            % (len(starts), len(stops))
        )

        if counts is not None:

            assert len(starts) == len(counts), (
                "starts length: %d and counts length: %d must have same length"
                % (len(starts), len(counts))
            )

        else:

            counts = [None] * len(starts)

        if dead_time is not None:

            assert len(starts) == len(dead_time), (
                "starts length: %d and dead_time length: %d must have same length"
                % (len(starts), len(dead_time))
            )

        else:

            dead_time = [0] * len(starts)

        list_of_intervals = []

        for imin, imax, ic, idt in zip(starts, stops, counts, dead_time):
            list_of_intervals.append(TimeInterval(imin, imax, ic, idt))

        return cls(list_of_intervals)

    @classmethod
    def from_list_of_edges(cls, edges):
        """
        Builds a IntervalSet from a list of time edges:

        edges = [-1,0,1] -> [-1,0], [0,1]

        :param edges:
        :return:
        """

        raise NotImplementedError()

        # # sort the time edges

        # edges.sort()

        # list_of_intervals = []

        # for imin, imax in zip(edges[:-1], edges[1:]):
        #     list_of_intervals.append(cls.new_interval(imin, imax))

        # return cls(list_of_intervals)

    def merge_intersecting_intervals(self, in_place=False):
        """

        merges intersecting intervals into a contiguous intervals

        :return:
        """

        # get a copy of the sorted intervals

        sorted_intervals = self.sort()

        new_intervals = []

        while len(sorted_intervals) > 1:

            # pop the first interval off the stack

            this_interval = sorted_intervals.pop(0)

            # see if that interval overlaps with the the next one

            if this_interval.overlaps_with(sorted_intervals[0]):

                # if so, pop the next one

                next_interval = sorted_intervals.pop(0)

                # and merge the two, appending them to the new intervals

                new_intervals.append(this_interval.merge(next_interval))

            else:

                # otherwise just append this interval

                new_intervals.append(this_interval)

                # now if there is only one interval left
                # it should not overlap with any other interval
                # and the loop will stop
                # otherwise, we continue

        # if there was only one interval
        # or a leftover from the merge
        # we append it
        if sorted_intervals:

            assert (
                len(sorted_intervals) == 1
            ), "there should only be one interval left over, this is a bug"  # pragma: no cover

            # we want to make sure that the last new interval did not
            # overlap with the final interval
            if new_intervals:

                if new_intervals[-1].overlaps_with(sorted_intervals[0]):

                    new_intervals[-1] = new_intervals[-1].merge(sorted_intervals[0])

                else:

                    new_intervals.append(sorted_intervals[0])

            else:

                new_intervals.append(sorted_intervals[0])

        if in_place:

            self.__init__(new_intervals)

        else:

            return TimeIntervalSet(new_intervals)

    def extend(self, list_of_intervals):

        self._intervals.extend(list_of_intervals)

    def __len__(self):

        return len(self._intervals)

    def __iter__(self):

        for interval in self._intervals:
            yield interval

    def __getitem__(self, item):

        return self._intervals[item]

    def __eq__(self, other):

        for interval_this, interval_other in zip(self.argsort(), other.argsort()):

            if not self[interval_this] == other[interval_other]:
                return False

        return True

    def pop(self, index):

        return self._intervals.pop(index)

    def sort(self):
        """
        Returns a sorted copy of the set (sorted according to the tstart of the time intervals)

        :return:
        """

        if self.is_sorted:

            return copy.deepcopy(self)

        else:

            return TimeIntervalSet(
                np.atleast_1d(itemgetter(*self.argsort())(self._intervals))
            )

    def argsort(self):
        """
        Returns the indices which order the set

        :return:
        """

        # Gather all tstarts
        tstarts = [x.start for x in self._intervals]

        return [x[0] for x in sorted(enumerate(tstarts), key=itemgetter(1))]

    def is_contiguous(self, relative_tolerance=1e-5):
        """
        Check whether the time intervals are all contiguous, i.e., the stop time of one interval is the start
        time of the next

        :return: True or False
        """

        starts = [attrgetter("start")(x) for x in self._intervals]
        stops = [attrgetter("stop")(x) for x in self._intervals]

        return np.allclose(starts[1:], stops[:-1], rtol=relative_tolerance)

    @property
    def is_sorted(self):
        """
        Check whether the time intervals are sorted
        :return: True or False
        """

        return np.all(self.argsort() == np.arange(len(self)))

    def containing_bin(self, value):
        """
        finds the index of the interval containing
        :param value:
        :return:
        """

        # Get the index of the first ebounds upper bound larger than energy
        # (but never go below zero or above the last channel)
        idx = min(max(0, np.searchsorted(self.edges, value) - 1), len(self))

        return idx

    def containing_interval(self, start, stop, inner=True, as_mask=False):
        """

        returns either a mask of the intervals contained in the selection
        or a new set of intervals within the selection. NOTE: no sort is performed

        :param start: start of interval
        :param stop: stop of interval
        :param inner: if True, returns only intervals strictly contained within bounds, if False, returns outer bounds as well
        :param as_mask: if you want a mask or the intervals
        :return:
        """

        # loop only once because every call unpacks the array
        # we need to round for the comparison because we may have read from
        # strings which are rounded to six decimals

        starts = np.round(self.starts, decimals=6)
        stops = np.round(self.stops, decimals=6)

        start = np.round(start, decimals=6)
        stop = np.round(stop, decimals=6)

        condition = (starts >= start) & (stop >= stops)

        if not inner:

            # now we get the end caps
            lower_condition = (starts <= start) & (start <= stops)

            upper_condition = (starts <= stop) & (stop <= stops)

            condition = condition | lower_condition | upper_condition

        # if we just want the mask

        if as_mask:

            return condition

        else:

            return TimeIntervalSet(np.asarray(self._intervals)[condition])

    @property
    def starts(self):
        """
        Return the starts fo the set

        :return: list of start times
        """

        return np.array([interval.start for interval in self._intervals])

    @property
    def stops(self):
        """
        Return the stops of the set

        :return:
        """

        return np.array([interval.stop for interval in self._intervals])

    @property
    def counts(self):

        return np.array([interval.counts for interval in self._intervals])

    @property
    def rates(self):

        return np.array([interval.rate for interval in self._intervals])

    @property
    def exposures(self):

        return np.array([interval.exposure for interval in self._intervals])

    @property
    def mid_points(self):

        return np.array([interval.mid_point for interval in self._intervals])

    @property
    def widths(self):

        return np.array([interval._get_width() for interval in self._intervals])

    @property
    def absolute_start(self):
        """
        the minimum of the start times
        :return:
        """

        return min(self.starts)

    @property
    def absolute_stop(self):
        """
        the maximum of the stop times
        :return:
        """

        return max(self.stops)

    @property
    def edges(self):
        """
        return an array of time edges if contiguous
        :return:
        """

        if self.is_contiguous() and self.is_sorted:

            edges = [
                interval.start
                for interval in itemgetter(*self.argsort())(self._intervals)
            ]
            edges.append(
                [
                    interval.stop
                    for interval in itemgetter(*self.argsort())(self._intervals)
                ][-1]
            )

        else:

            raise IntervalsNotContiguous(
                "Cannot return edges for non-contiguous intervals"
            )

        return edges

    def to_string(self):
        """

        returns a set of string representaitons of the intervals
        :return:
        """

        return ",".join([interval.to_string() for interval in self._intervals])

    @property
    def bin_stack(self):
        """

        get a stacked view of the bins [[start_1,stop_1 ],
                                        [start_2,stop_2 ]]

        :return:
        """

        return np.vstack((self.starts, self.stops)).T

    def __add__(self, number):
        """
        Shift all time intervals to the right by number

        :param number: a float
        :return: new TimeIntervalSet instance
        """

        new_set = TimeIntervalSet()
        new_set.extend([time_interval + number for time_interval in self._intervals])

        return new_set

    def __sub__(self, number):
        """
        Shift all time intervals to the left by number (in place)

        :param number: a float
        :return: new TimeIntervalSet instance
        """

        new_set = TimeIntervalSet(
            [time_interval - number for time_interval in self._intervals]
        )

        return new_set

    def plot_intervals(self, as_rates=True, ax=None, **kwargs):
        """
        plot the intervals as rates or counts

        :param as_rates:
        :param ax:
        :returns:
        :rtype:

        """

        if ax is None:

            fig, ax = plt.subplots()

        else:

            fig = ax.get_figure()

        xbins = self.bin_stack

        if as_rates:

            y = self.rates
            y_label = "rate (cnts/s)"

        else:

            y = self.counts
            y_label = "cnts"

        step_plot(xbins, y, ax=ax, **kwargs)

        ax.set_xlabel("time")
        ax.set_ylabel(y_label)

        return fig

max_stars_repo_name: wematthias/cosmogrb

max_issues_repo_name: wematthias/cosmogrb

max_forks_repo_name: wematthias/cosmogrb

reason: import numpy


text:
"""Figure based on Open3D's visualizer."""
import numpy as np
import open3d as o3d
from .. import rotations as pr
from .. import transformations as pt
from .. import trajectories as ptr
from ._artists import (Line3D, Frame, Trajectory, Camera, Box, Sphere,
                       Cylinder, Mesh, Ellipsoid, Capsule, Cone, Graph)

class Figure:
    """The top level container for all the plot elements.

    You can close the visualizer with the keys `escape` or `q`.

    Parameters
    ----------
    window_name : str, optional (default: Open3D)
        Window title name.

    width : int, optional (default: 1920)
        Width of the window.

    height : int, optional (default: 1080)
        Height of the window.

    with_key_callbacks : bool, optional (default: False)
        Creates a visualizer that allows to register callbacks
        for keys.
    """
    def __init__(self, window_name="Open3D", width=1920, height=1080,
                 with_key_callbacks=False):
        if with_key_callbacks:
            self.visualizer = o3d.visualization.VisualizerWithKeyCallback()
        else:
            self.visualizer = o3d.visualization.Visualizer()
        self.visualizer.create_window(
            window_name=window_name, width=width, height=height)

    def add_geometry(self, geometry):
        """Add geometry to visualizer.

        Parameters
        ----------
        geometry : Geometry
            Open3D geometry.
        """
        self.visualizer.add_geometry(geometry)

    def _remove_geometry(self, geometry):
        """Remove geometry to visualizer.

        .. warning::

            This function is not public because the interface of the
            underlying visualizer might change in the future causing the
            signature of this function to change as well.

        Parameters
        ----------
        geometry : Geometry
            Open3D geometry.
        """
        self.visualizer.remove_geometry(geometry)

    def update_geometry(self, geometry):
        """Indicate that geometry has been updated.

        Parameters
        ----------
        geometry : Geometry
            Open3D geometry.
        """
        self.visualizer.update_geometry(geometry)

    def remove_artist(self, artist):
        """Remove artist from visualizer.

        Parameters
        ----------
        artist : Artist
            Artist that should be removed from this figure.
        """
        for g in artist.geometries:
            self._remove_geometry(g)

    def set_line_width(self, line_width):
        """Set render option line width.

        Note: this feature does not work in Open3D's visualizer at the
        moment.

        Parameters
        ----------
        line_width : float
            Line width.
        """
        self.visualizer.get_render_option().line_width = line_width
        self.visualizer.update_renderer()

    def set_zoom(self, zoom):
        """Set zoom.

        Parameters
        ----------
        zoom : float
            Zoom of the visualizer.
        """
        self.visualizer.get_view_control().set_zoom(zoom)

    def animate(self, callback, n_frames, loop=False, fargs=()):
        """Make animation with callback.

        Parameters
        ----------
        callback : callable
            Callback that will be called in a loop to update geometries.
            The first input of the function will be the current frame
            index from [0, `n_frames`). Further arguments can be given as
            `fargs`. The function should return one artist object or a
            list of artists that have been updated.

        n_frames : int
            Total number of frames.

        loop : bool, optional (default: False)
            Run callback in an infinite loop.

        fargs : list, optional (default: [])
            Arguments that will be passed to the callback.

        Raises
        ------
        RuntimeError
            When callback does not return any artists
        """
        initialized = False
        window_open = True
        while window_open and (loop or not initialized):
            for i in range(n_frames):
                drawn_artists = callback(i, *fargs)

                if drawn_artists is None:
                    raise RuntimeError(
                        "The animation function must return a "
                        "sequence of Artist objects.")
                try:
                    drawn_artists = [a for a in drawn_artists]
                except TypeError:
                    drawn_artists = [drawn_artists]

                for a in drawn_artists:
                    for geometry in a.geometries:
                        self.update_geometry(geometry)

                window_open = self.visualizer.poll_events()
                if not window_open:
                    break
                self.visualizer.update_renderer()
            initialized = True

    def view_init(self, azim=-60, elev=30):
        """Set the elevation and azimuth of the axes.

        Parameters
        ----------
        azim : float, optional (default: -60)
            Azimuth angle in the x,y plane in degrees.

        elev : float, optional (default: 30)
            Elevation angle in the z plane.
        """
        vc = self.visualizer.get_view_control()
        pcp = vc.convert_to_pinhole_camera_parameters()
        distance = np.linalg.norm(pcp.extrinsic[:3, 3])
        R_azim_elev_0_world2camera = np.array([
            [0, 1, 0],
            [0, 0, -1],
            [-1, 0, 0]])
        R_azim_elev_0_camera2world = R_azim_elev_0_world2camera.T
        # azimuth and elevation are defined in world frame
        R_azim = pr.active_matrix_from_angle(2, np.deg2rad(azim))
        R_elev = pr.active_matrix_from_angle(1, np.deg2rad(-elev))
        R_elev_azim_camera2world = R_azim.dot(R_elev).dot(
            R_azim_elev_0_camera2world)
        pcp.extrinsic = pt.transform_from(  # world2camera
            R=R_elev_azim_camera2world.T,
            p=[0, 0, distance])
        vc.convert_from_pinhole_camera_parameters(pcp)

    def plot(self, P, c=(0, 0, 0)):
        """Plot line.

        Parameters
        ----------
        P : array-like, shape (n_points, 3)
            Points of which the line consists.

        c : array-like, shape (n_points - 1, 3) or (3,), optional (default: black)
            Color can be given as individual colors per line segment or
            as one color for each segment. A color is represented by 3
            values between 0 and 1 indicate representing red, green, and
            blue respectively.

        Returns
        -------
        Line3D : line
            New line.
        """
        line3d = Line3D(P, c)
        line3d.add_artist(self)
        return line3d

    def plot_basis(self, R=None, p=np.zeros(3), s=1.0, strict_check=True):
        """Plot basis.

        Parameters
        ----------
        R : array-like, shape (3, 3), optional (default: I)
            Rotation matrix, each column contains a basis vector

        p : array-like, shape (3,), optional (default: [0, 0, 0])
            Offset from the origin

        s : float, optional (default: 1)
            Scaling of the frame that will be drawn

        strict_check : bool, optional (default: True)
            Raise a ValueError if the rotation matrix is not numerically
            close enough to a real rotation matrix. Otherwise we print a
            warning.

        Returns
        -------
        Frame : frame
            New frame.
        """
        if R is None:
            R = np.eye(3)
        R = pr.check_matrix(R, strict_check=strict_check)

        frame = Frame(pt.transform_from(R=R, p=p), s=s)
        frame.add_artist(self)

        return frame

    def plot_transform(self, A2B=None, s=1.0, name=None, strict_check=True):
        """Plot coordinate frame.

        Parameters
        ----------
        A2B : array-like, shape (4, 4)
            Transform from frame A to frame B

        s : float, optional (default: 1)
            Length of basis vectors

        name : str, optional (default: None)
            Name of the frame

        strict_check : bool, optional (default: True)
            Raise a ValueError if the transformation matrix is not
            numerically close enough to a real transformation matrix.
            Otherwise we print a warning.

        Returns
        -------
        Frame : frame
            New frame.
        """
        if A2B is None:
            A2B = np.eye(4)
        A2B = pt.check_transform(A2B, strict_check=strict_check)

        frame = Frame(A2B, name, s)
        frame.add_artist(self)

        return frame

    def plot_trajectory(self, P, n_frames=10, s=1.0, c=(0, 0, 0)):
        """Trajectory of poses.

        Parameters
        ----------
        P : array-like, shape (n_steps, 7), optional (default: None)
            Sequence of poses represented by positions and quaternions in
            the order (x, y, z, w, vx, vy, vz) for each step

        n_frames : int, optional (default: 10)
            Number of frames that should be plotted to indicate the
            rotation

        s : float, optional (default: 1)
            Scaling of the frames that will be drawn

        c : array-like, shape (3,), optional (default: black)
            A color is represented by 3 values between 0 and 1 indicate
            representing red, green, and blue respectively.

        Returns
        -------
        trajectory : Trajectory
            New trajectory.
        """
        H = ptr.matrices_from_pos_quat(P)
        trajectory = Trajectory(H, n_frames, s, c)
        trajectory.add_artist(self)
        return trajectory

    def plot_sphere(self, radius=1.0, A2B=np.eye(4), resolution=20, c=None):
        """Plot sphere.

        Parameters
        ----------
        radius : float, optional (default: 1)
            Radius of the sphere

        A2B : array-like, shape (4, 4)
            Transform from frame A to frame B

        resolution : int, optional (default: 20)
            The resolution of the sphere. The longitudes will be split into
            resolution segments (i.e. there are resolution + 1 latitude
            lines including the north and south pole). The latitudes will
            be split into 2 * resolution segments (i.e. there are
            2 * resolution longitude lines.)

        c : array-like, shape (3,), optional (default: None)
            Color

        Returns
        -------
        sphere : Sphere
            New sphere.
        """
        sphere = Sphere(radius, A2B, resolution, c)
        sphere.add_artist(self)
        return sphere

    def plot_box(self, size=np.ones(3), A2B=np.eye(4), c=None):
        """Plot box.

        Parameters
        ----------
        size : array-like, shape (3,), optional (default: [1, 1, 1])
            Size of the box per dimension

        A2B : array-like, shape (4, 4), optional (default: I)
            Center of the box

        c : array-like, shape (3,), optional (default: None)
            Color

        Returns
        -------
        box : Box
            New box.
        """
        box = Box(size, A2B, c)
        box.add_artist(self)
        return box

    def plot_cylinder(self, length=2.0, radius=1.0, A2B=np.eye(4),
                      resolution=20, split=4, c=None):
        """Plot cylinder.

        Parameters
        ----------
        length : float, optional (default: 1)
            Length of the cylinder.

        radius : float, optional (default: 1)
            Radius of the cylinder.

        A2B : array-like, shape (4, 4)
            Pose of the cylinder. The position corresponds to the center of the
            line segment and the z-axis to the direction of the line segment.

        resolution : int, optional (default: 20)
            The circle will be split into resolution segments

        split : int, optional (default: 4)
            The height will be split into split segments

        c : array-like, shape (3,), optional (default: None)
            Color

        Returns
        -------
        cylinder : Cylinder
            New cylinder.
        """
        cylinder = Cylinder(length, radius, A2B, resolution, split, c)
        cylinder.add_artist(self)
        return cylinder

    def plot_mesh(self, filename, A2B=np.eye(4), s=np.ones(3), c=None):
        """Plot mesh.

        Parameters
        ----------
        filename : str
            Path to mesh file

        A2B : array-like, shape (4, 4)
            Center of the mesh

        s : array-like, shape (3,), optional (default: [1, 1, 1])
            Scaling of the mesh that will be drawn

        c : array-like, shape (n_vertices, 3) or (3,), optional (default: None)
            Color(s)

        Returns
        -------
        mesh : Mesh
            New mesh.
        """
        mesh = Mesh(filename, A2B, s, c)
        mesh.add_artist(self)
        return mesh

    def plot_ellipsoid(self, radii=np.ones(3), A2B=np.eye(4), resolution=20,
                       c=None):
        """Plot ellipsoid.

        Parameters
        ----------
        radii : array-like, shape (3,)
            Radii along the x-axis, y-axis, and z-axis of the ellipsoid.

        A2B : array-like, shape (4, 4)
            Transform from frame A to frame B

        resolution : int, optional (default: 20)
            The resolution of the ellipsoid. The longitudes will be split into
            resolution segments (i.e. there are resolution + 1 latitude
            lines including the north and south pole). The latitudes will
            be split into 2 * resolution segments (i.e. there are
            2 * resolution longitude lines.)

        c : array-like, shape (3,), optional (default: None)
            Color

        Returns
        -------
        ellipsoid : Ellipsoid
            New ellipsoid.
        """
        ellipsoid = Ellipsoid(radii, A2B, resolution, c)
        ellipsoid.add_artist(self)
        return ellipsoid

    def plot_capsule(self, height=1, radius=1, A2B=np.eye(4), resolution=20,
                     c=None):
        """Plot capsule.

        A capsule is the volume covered by a sphere moving along a line segment.

        Parameters
        ----------
        height : float, optional (default: 1)
            Height of the capsule along its z-axis.

        radius : float, optional (default: 1)
            Radius of the capsule.

        A2B : array-like, shape (4, 4)
            Pose of the capsule. The position corresponds to the center of the
            line segment and the z-axis to the direction of the line segment.

        resolution : int, optional (default: 20)
            The resolution of the capsule. The longitudes will be split into
            resolution segments (i.e. there are resolution + 1 latitude lines
            including the north and south pole). The latitudes will be split
            into 2 * resolution segments (i.e. there are 2 * resolution
            longitude lines.)

        c : array-like, shape (3,), optional (default: None)
            Color

        Returns
        -------
        capsule : Capsule
            New capsule.
        """
        capsule = Capsule(height, radius, A2B, resolution, c)
        capsule.add_artist(self)
        return capsule

    def plot_cone(self, height=1, radius=1, A2B=np.eye(4), resolution=20,
                  c=None):
        """Plot cone.

        Parameters
        ----------
        height : float, optional (default: 1)
            Height of the cone along its z-axis.

        radius : float, optional (default: 1)
            Radius of the cone.

        A2B : array-like, shape (4, 4)
            Pose of the cone.

        resolution : int, optional (default: 20)
            The circle will be split into resolution segments.

        c : array-like, shape (3,), optional (default: None)
            Color
        """
        cone = Cone(height, radius, A2B, resolution, c)
        cone.add_artist(self)
        return cone

    def plot_graph(
            self, tm, frame, show_frames=False, show_connections=False,
            show_visuals=False, show_collision_objects=False,
            show_name=False, whitelist=None, s=1.0):
        """Plot graph of connected frames.

        Parameters
        ----------
        tm : TransformManager
            Representation of the graph

        frame : str
            Name of the base frame in which the graph will be displayed

        show_frames : bool, optional (default: False)
            Show coordinate frames

        show_connections : bool, optional (default: False)
            Draw lines between frames of the graph

        show_visuals : bool, optional (default: False)
            Show visuals that are stored in the graph

        show_collision_objects : bool, optional (default: False)
            Show collision objects that are stored in the graph

        show_name : bool, optional (default: False)
            Show names of frames

        whitelist : list, optional (default: all)
            List of frames that should be displayed

        s : float, optional (default: 1)
            Scaling of the frames that will be drawn

        Returns
        -------
        graph : Graph
            New graph.
        """
        graph = Graph(tm, frame, show_frames, show_connections,
                      show_visuals, show_collision_objects, show_name,
                      whitelist, s)
        graph.add_artist(self)
        return graph

    def plot_camera(self, M, cam2world=None, virtual_image_distance=1,
                    sensor_size=(1920, 1080), strict_check=True):
        """Plot camera in world coordinates.

        This function is inspired by Blender's camera visualization. It will
        show the camera center, a virtual image plane, and the top of the
        virtual image plane.

        Parameters
        ----------
        M : array-like, shape (3, 3)
            Intrinsic camera matrix that contains the focal lengths on the
            diagonal and the center of the the image in the last column. It
            does not matter whether values are given in meters or pixels as
            long as the unit is the same as for the sensor size.

        cam2world : array-like, shape (4, 4), optional (default: I)
            Transformation matrix of camera in world frame. We assume that the
            position is given in meters.

        virtual_image_distance : float, optional (default: 1)
            Distance from pinhole to virtual image plane that will be
            displayed. We assume that this distance is given in meters. The
            unit has to be consistent with the unit of the position in
            cam2world.

        sensor_size : array-like, shape (2,), optional (default: [1920, 1080])
            Size of the image sensor: (width, height). It does not matter
            whether values are given in meters or pixels as long as the unit is
            the same as for the sensor size.

        strict_check : bool, optional (default: True)
            Raise a ValueError if the transformation matrix is not numerically
            close enough to a real transformation matrix. Otherwise we print a
            warning.

        Returns
        -------
        camera : Camera
            New camera.
        """
        camera = Camera(M, cam2world, virtual_image_distance, sensor_size,
                        strict_check)
        camera.add_artist(self)
        return camera

    def save_image(self, filename):
        """Save rendered image to file.

        Parameters
        ----------
        filename : str
            Path to file in which the rendered image should be stored
        """
        self.visualizer.capture_screen_image(filename, True)

    def show(self):
        """Display the figure window."""
        self.visualizer.run()
        self.visualizer.destroy_window()

def figure(window_name="Open3D", width=1920, height=1080,
           with_key_callbacks=False):
    """Create a new figure.

    Parameters
    ----------
    window_name : str, optional (default: Open3D)
        Window title name.

    width : int, optional (default: 1920)
        Width of the window.

    height : int, optional (default: 1080)
        Height of the window.

    with_key_callbacks : bool, optional (default: False)
        Creates a visualizer that allows to register callbacks
        for keys.

    Returns
    -------
    figure : Figure
        New figure.
    """
    return Figure(window_name, width, height, with_key_callbacks)


text:
import argparse

import numpy as np

import matplotlib.pyplot as plt

from scipy import stats

from utils import StateFactory

def argmax(array):
    rows = range(array.shape[0])
    cols = array.argmax(axis=1)

    return rows, cols

def compute_prob(n_block, sigma, tau, state_factory):
    state = state_factory()

    def gaussian():
        return state.randn(n_block, 2)

    theta = gaussian() * sigma

    # making one observation per entry
    mean = (theta + gaussian() * tau) / 2

    # choosing between e0 and e1 within each block
    sample = mean + gaussian() / 2 ** 0.5
    index = argmax(sample)

    # updating the mean given new rewards
    rew = theta + gaussian() * tau
    mean[index] = (2 * mean + rew)[index] / 3

    # computing the directional mean and variance
    dir_mean = mean.sum()
    dir_var = n_block * (1 / 2 + 1 / 3)

    ratio = (tau - sigma) * dir_mean / dir_var ** 0.5

    return theta.sum() * (tau - sigma) > 0, stats.norm.cdf(ratio)

def run_experiments(n_iter, sigma, tau, seed):
    state_factory = StateFactory(seed)

    n_blocks = 2 ** np.arange(18)
    n_failures = np.zeros((n_iter,) + n_blocks.shape, dtype=np.float64)

    for i, n_block in enumerate(n_blocks):
        results = [
            compute_prob(n_block, sigma, tau, state_factory) for _ in range(n_iter)
        ]

        n_failures[:, i] = np.array([1 / p for _, p in results])
        n_subopt = sum(i for i, _ in results) / len(results)

        print(
            f"Number of blocks = {n_block:7d} -- proportion of times 0 is suboptimal = {n_subopt:.2f} -- "
            f"min number of steps TS will fail = {n_failures[:, i].min()}"
        )

    # plotting the data
    plt.yscale('log')
    plt.boxplot(n_failures, showfliers=False)
    plt.xticks(range(1, 19), [f"$2^{{{i}}}$" for i in range(1, 19)])

    plt.xlabel("Number of blocks (half of the dimension)")
    plt.ylabel("Expected number of failures for LinTS")
    plt.tight_layout()

    plt.savefig("plots/example-1.pdf")

def __main__():
    parser = argparse.ArgumentParser(
        description="Simulate the first example for TS failure."
    )

    parser.add_argument("--n-iter", type=int, help="number of iterations", default=50)
    parser.add_argument("--sigma", type=float, help="prior sd", default=1.0)
    parser.add_argument("--tau", type=float, help="noise sd", default=0.0)
    parser.add_argument("--seed", type=int, help="initial random seed", default=1)

    args = parser.parse_args()

    run_experiments(n_iter=args.n_iter, sigma=args.sigma, tau=args.tau, seed=args.seed)

if __name__ == "__main__":
    __main__()

max_stars_repo_name: yuweiluo/linear-ts

max_issues_repo_name: yuweiluo/linear-ts

max_forks_repo_name: yuweiluo/linear-ts

reason: import numpy,from scipy


text:
import cv2
from openvino.inference_engine import IENetwork, IECore
import numpy as np
'''
This is a sample class for a model. You may choose to use it as-is or make any changes to it.
This has been provided just to give you an idea of how to structure your model class.
'''
class facedetection:
    '''
    Class for the Face Detection Model.
    '''
    def __init__(self, model_name, device, extensions=None, thershold=0.5):
        '''
        TODO: Use this to set your instance variables.
        '''
        self.model_weights = model_name+'.bin'
        self.model_structure = model_name+'.xml'
        self.device = device
        self.threshold = thershold

        try:
            self.model=IENetwork(self.model_structure, self.model_weights)
        except Exception as e:
            raise ValueError("Could not Initialise the network. Have you enterred the correct model path?")

        self.infer_request_handle = None
        self.input_name=next(iter(self.model.inputs))
        self.input_shape=self.model.inputs[self.input_name].shape
        self.output_name=next(iter(self.model.outputs))
        self.output_shape=self.model.outputs[self.output_name].shape
        #raise NotImplementedError

    def load_model(self):
        '''
        TODO: You will need to complete this method.
        This method is for loading the model to the device specified by the user.
        If your model requires any Plugins, this is where you can load them.
        '''
        self.core = IECore()
        self.net = self.core.load_network(network=self.model, device_name=self.device, num_requests=1)
        #raise NotImplementedError

    def predict(self, image):
        '''
        TODO: You will need to complete this method.
        This method is meant for running predictions on the input image.
        '''
        self.processed_img = self.preprocess_input(image)
        input_dict = {self.input_name:self.processed_img}
        infer_request_handle = self.net.start_async(request_id=0, inputs=input_dict)
        if self.net.requests[0].wait(-1)==0:
            return self.preprocess_output(infer_request_handle.outputs)
        #raise NotImplementedError

    def check_model(self):
        raise NotImplementedError

    def preprocess_input(self, image):
        '''
            Before feeding the data into the model for inference,
        you might have to preprocess it. This function is where you can do that.
        '''
        self.image = image
        input_image = cv2.resize(image, (self.input_shape[3], self.input_shape[2]))
        input_image = input_image.transpose((2,0,1))
        input_image = input_image.reshape((1, *input_image.shape))
        return input_image
        #raise NotImplementedError

    def preprocess_output(self, outputs):
        '''
        Before feeding the output of this model to the next model,
        you might have to preprocess the output. This function is where you can do that.
        '''
        #image = self.image
        cr_image = np.array([[],[],[]])
        for box in outputs[self.output_name][0][0]:
            conf = box[2]
            if conf >= self.threshold:
                x_0 = int(box[3] * self.image.shape[1])
                y_0 = int(box[4] * self.image.shape[0])
                x_1 = int(box[5] * self.image.shape[1])
                y_1 = int(box[6] * self.image.shape[0])
                cr_image = self.image[y_0:y_1, x_0:x_1]
                coords_image = [x_0, y_0, x_1, y_1]
        return cr_image, coords_image
        #raise NotImplementedError


text:
#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import sys

if __name__ == "__main__":
    filename = sys.argv[1]
    data = np.genfromtxt(filename, delimiter=",", names=True)
    begin = data[data.dtype.names[0]]
    end = data[data.dtype.names[1]]
    event = ["{}".format(i) for i in range(len(begin))]
    plt.barh(range(len(begin)),  end-begin, left=begin)
    plt.yticks(range(len(begin)), event)
    plt.show()

max_stars_repo_name: nevesnunes/aggregables

max_issues_repo_name: nevesnunes/aggregables

max_forks_repo_name: nevesnunes/aggregables


text:
"""
Tests for sph2grd.
"""
import os

import numpy.testing as npt
from pygmt import sph2grd
from pygmt.helpers import GMTTempFile

def test_sph2grd_outgrid():
    """
    Test sph2grd with a set outgrid.
    """
    with GMTTempFile(suffix=".nc") as tmpfile:
        result = sph2grd(
            data="@EGM96_to_36.txt", outgrid=tmpfile.name, spacing=1, region="g"
        )
        assert result is None  # return value is None
        assert os.path.exists(path=tmpfile.name)  # check that outgrid exists

def test_sph2grd_no_outgrid():
    """
    Test sph2grd with no set outgrid.
    """
    temp_grid = sph2grd(data="@EGM96_to_36.txt", spacing=1, region="g")
    assert temp_grid.dims == ("y", "x")
    assert temp_grid.gmt.gtype == 0  # Cartesian grid
    assert temp_grid.gmt.registration == 0  # Gridline registration
    npt.assert_allclose(temp_grid.max(), 0.00021961, rtol=1e-4)
    npt.assert_allclose(temp_grid.min(), -0.0004326, rtol=1e-4)
    npt.assert_allclose(temp_grid.median(), -0.00010894, rtol=1e-4)
    npt.assert_allclose(temp_grid.mean(), -0.00010968, rtol=1e-4)

max_stars_repo_name: jbusecke/pygmt

max_issues_repo_name: jbusecke/pygmt

max_forks_repo_name: jbusecke/pygmt


text:
## Copyright (C) 2019, Guangke Chen <gkchen.shanghaitech@gmail.com>.
## This program is licenced under the BSD 2-Clause licence
## contained in the LICENCE file in this directory.

import numpy as np
import os
from scipy.io.wavfile import write
import subprocess
import shlex

class gmm_ubm_kaldiHelper(object):

    def __init__(self, pre_model_dir="pre-models", audio_dir=None, mfcc_dir=None, log_dir=None, score_dir=None):

        self.pre_model_dir = os.path.abspath(pre_model_dir)
        self.conf_dir = os.path.join(self.pre_model_dir, "conf")
        audio_dir = audio_dir if audio_dir else "audio"
        self.audio_dir = os.path.abspath(audio_dir)

        mfcc_dir = mfcc_dir if mfcc_dir else "mfcc"
        self.mfcc_dir = os.path.abspath(mfcc_dir)

        log_dir = log_dir if log_dir else "log"
        self.log_dir = os.path.abspath(log_dir)

        score_dir = score_dir if score_dir else "score"
        self.score_dir = os.path.abspath(score_dir)

        ''' deal with the protential permission issue
        '''
        all_files = (self.get_all_files(self.pre_model_dir + "/utils") +
                     self.get_all_files(self.pre_model_dir + "/steps") +
                     self.get_all_files(self.pre_model_dir + "/sid"))

        for file in all_files:
            change_permission_command = "chmod 777 " + file
            args = shlex.split(change_permission_command)
            p = subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            p.wait()

    def get_all_files(self, root_dir):

        files = []
        for name in os.listdir(root_dir):
            path = os.path.join(root_dir, name)
            if os.path.isfile(path):
                files.append(path)
            else:
                files_sub = self.get_all_files(path)
                files += files_sub

        return files

    def write_audio(self, audio_list, audio_dir=None, fs=16000):

        audio_dir = os.path.abspath(audio_dir) if audio_dir else self.audio_dir

        n_audios = len(audio_list)
        audio_path_list = list(range(n_audios))

        for i, audio in enumerate(audio_list):
            path = audio_dir + "/{}.wav".format(i)
            write(path, fs, audio)
            audio_path_list[i] = path

        return audio_path_list

    def data_prepare(self, audio_path_list, utt_id_list=None, spk_id_list=None, audio_dir=None, debug=False):

        ''' generate wav.scp, utt2spk, spk2utt in audio_dir according to utt_id_list,
            spk_id_list, audio_path_list. utt_id_list, spk_id_list, audio_path_list are
            three list objects with the same len
            Note: if utt_id_list and spk_id_listis provided, both of them must be in sorted order,
            otherwise, the returning scores may disorder
        '''

        audio_dir = os.path.abspath(audio_dir) if audio_dir else self.audio_dir
        if not audio_dir:
            print("--- Error:audio_dir is None, quit data_prepare ---")
            return

        ''' if audio_path_list is in relative path, make it in absolute path, otherwise kaldi may cannot finds them.
        '''
        audio_path_list = [os.path.abspath(path) for path in audio_path_list]

        n_audios = len(audio_path_list)
        spk_id_list = spk_id_list if spk_id_list else [("0000" + str(i+1))[-5:] for i in range(n_audios)] # this may cause a bug if there are more than 100000 audios in audio_list
        utt_id_list = utt_id_list if utt_id_list else [spk_id + "-{}".format(1) for spk_id in spk_id_list]

        # wav.scp, utt2spk
        wav_file = self.audio_dir + "/wav.scp"
        utt2spk_file = self.audio_dir + "/utt2spk"
        wav_mat = np.concatenate((np.array(utt_id_list)[:, np.newaxis], np.array(audio_path_list)[:, np.newaxis]), axis=1)
        np.savetxt(wav_file, wav_mat, fmt="%s")
        utt2spk_mat = np.concatenate((np.array(utt_id_list)[:, np.newaxis], np.array(spk_id_list)[:, np.newaxis]), axis=1)
        np.savetxt(utt2spk_file, utt2spk_mat, fmt="%s")

        # spk2utt
        spk_id_vector = utt2spk_mat[:, 1]
        spk_id_unique_vector = np.unique(spk_id_vector)
        spk2utt_list = []
        for spk_id in spk_id_unique_vector:
            utt_index = np.argwhere(spk_id_vector == spk_id).flatten()
            utts = ""
            for index in utt_index:
                utts += (utt2spk_mat[index, 0] + " ")
            utts.rstrip()
            spk2utt_list.append(utts)

        spk2utt_mat = np.concatenate((spk_id_unique_vector[:, np.newaxis], np.array(spk2utt_list)[:, np.newaxis]), axis=1)
        spk2utt_file = self.audio_dir + "/spk2utt"
        np.savetxt(spk2utt_file, spk2utt_mat, fmt="%s")

        # fix dir
        ''' Note: every time we execuate the shell script in utils/, steps/ or sid/,
            we need to change the current directory to pre_model_dir.
            To avoid mistask, we should make all the dirs (e.g., audio_dir, log_dir) in absoulute path
        '''
        current_dir = os.path.abspath(os.curdir)
        os.chdir(self.pre_model_dir)

        fix_dir_command = self.pre_model_dir + "/utils/fix_data_dir.sh " + audio_dir
        args = shlex.split(fix_dir_command)
        p = subprocess.Popen(args) if debug else subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        p.wait()

        os.chdir(current_dir)

    def make_mfcc(self, n_jobs=10, mfcc_conf=None, audio_dir=None, mfcc_dir=None, log_dir=None, debug=False):

        mfcc_conf = os.path.abspath(mfcc_conf) if mfcc_conf else (self.conf_dir + "/mfcc.conf")
        audio_dir = os.path.abspath(audio_dir) if audio_dir else self.audio_dir
        mfcc_dir = os.path.abspath(mfcc_dir) if mfcc_dir else self.mfcc_dir
        log_dir = os.path.abspath(log_dir) if log_dir else self.log_dir

        extract_mfcc_command = (self.pre_model_dir + "/steps/make_mfcc.sh --write-utt2num-frames true --mfcc-config " +
                                mfcc_conf + " --nj " + str(n_jobs) + " --cmd '$train_cmd' " + audio_dir + " " +
                                log_dir + " " + mfcc_dir)

        current_dir = os.path.abspath(os.curdir)
        os.chdir(self.pre_model_dir)

        args = shlex.split(extract_mfcc_command)
        p = subprocess.Popen(args) if debug else subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        p.wait()

        os.chdir(current_dir)

    def compute_vad(self, n_jobs=10, vad_conf=None, audio_dir=None, vad_dir=None, log_dir=None, debug=False):

        vad_conf = os.path.abspath(vad_conf) if vad_conf else (self.conf_dir + "/vad.conf")
        audio_dir = os.path.abspath(audio_dir) if audio_dir else self.audio_dir
        vad_dir = os.path.abspath(vad_dir) if vad_dir else self.mfcc_dir
        log_dir = os.path.abspath(log_dir) if log_dir else self.log_dir

        vad_command = (self.pre_model_dir + "/sid/compute_vad_decision.sh --nj " + str(n_jobs) +
                        " --cmd '$train_cmd' --vad-config " + vad_conf + " " +
                        audio_dir + " " + log_dir + " " + vad_dir)

        current_dir = os.path.abspath(os.curdir)
        os.chdir(self.pre_model_dir)

        args = shlex.split(vad_command)
        p = subprocess.Popen(args) if debug else subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        p.wait()

        os.chdir(current_dir)

    def get_frames_likes(self, model_path_list, n_jobs=10, debug=False):

        cmd = self.pre_model_dir + "/utils/run.pl"
        n_threads = 1
        get_frame_likes_log = self.log_dir + "/get_frame_likes.JOB.log"
        adverage="--average=true"

        # split data for parallel running
        current_dir = os.path.abspath(os.curdir)
        os.chdir(self.pre_model_dir)

        split_data_command = self.pre_model_dir + "/utils/split_data.sh" + " " + self.audio_dir + " " + str(n_jobs)
        args = shlex.split(split_data_command)
        p = subprocess.Popen(args) if debug else subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        p.wait()

        os.chdir(current_dir)

        sdata = self.audio_dir + "/split" + str(n_jobs)

        delta_opts_file = os.path.join(self.pre_model_dir, "delta_opts")
        with open(delta_opts_file, "r") as reader:
            delta_opts = reader.read()[:-1]

        add_deltas = ("add-deltas " + delta_opts + " scp:" + sdata + "/JOB/feats.scp" + " ark:- |")
        apply_cmvn = "apply-cmvn-sliding --norm-vars=false --center=true --cmn-window=300 ark:- ark:- |"
        select_voiced_frame = ("select-voiced-frames ark:- scp,s,cs:"  + sdata + "/JOB/vad.scp" + " ark:- |")
        feats = ("ark,s,cs:" + add_deltas + " " + apply_cmvn + " " + select_voiced_frame)

        job_scores_file = self.score_dir + "/score.JOB"

        for i, model_path in enumerate(model_path_list):

            current_dir = os.path.abspath(os.curdir)
            os.chdir(self.pre_model_dir)
            get_frames_likes_command = ("gmm-global-get-frame-likes " +
                                        adverage + " " + model_path + " " +
                                        shlex.quote(feats) + " ark,t:" + job_scores_file)
            # cmd_command = (cmd + " --num-threads " + str(n_threads) +
            #                " JOB=1:" + str(n_jobs) + " " +
            #                get_frame_likes_log + " " +
            #                shlex.quote(get_frames_likes_command) + " || exit 1;")
            cmd_command = (cmd + " --num-threads " + str(n_threads) +
                           " JOB=1:" + str(n_jobs) + " " +
                           get_frame_likes_log + " " +
                           get_frames_likes_command + " || exit 1;")

            args = shlex.split(cmd_command)
            p = subprocess.Popen(args) if debug else subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            p.wait()
            os.chdir(current_dir)

            scores_file = self.score_dir + "/" + str(i+1) + ".score"
            content = []
            for job_id in range(1, n_jobs+1):

                job_score_file = self.score_dir + "/score." + str(job_id)
                with open(job_score_file, "r") as reader:
                    content += reader.readlines()

                os.remove(job_score_file)

            with open(scores_file, "w") as writer:
                writer.writelines(content)

    def resolce_scores(self, model_path_list):

        score_array = []
        for i in range(1, len(model_path_list) + 1):
            scores_file = self.score_dir + "/" + str(i) + ".score"
            score = np.loadtxt(scores_file, usecols=1) # (n_audios, ) or a scalar(if only one audio)
            score_array.append(score)
        score_array = np.array(score_array) # (n_models, n_audios) or (n_models, ) (if only one audio)
        if len(score_array.shape) == 1:
            score_array = score_array[:, np.newaxis]
        score_array = score_array.T # (n_audios, n_models)

        return score_array

    def score_existing(self, model_path_list, audio_path_list, spk_id_list=None, utt_id_list=None,
                       fs=16000, n_jobs=10, debug=False, bits_per_sample=16, mfcc_conf=None):

        ''' make model path being absolute path
        '''
        model_path_list = [os.path.abspath(path) for path in model_path_list]

        n_audios = len(audio_path_list)
        if n_jobs > n_audios:
            n_jobs = n_audios

        self.data_prepare(audio_path_list, utt_id_list=utt_id_list, spk_id_list=spk_id_list, debug=debug)
        self.make_mfcc(n_jobs=n_jobs, mfcc_conf=mfcc_conf, debug=debug)
        self.compute_vad(n_jobs=n_jobs, debug=debug)
        self.get_frames_likes(model_path_list, n_jobs=n_jobs, debug=debug)
        score_array = self.resolce_scores(model_path_list)

        return score_array

    def score(self, model_path_list, audio_list, fs=16000, n_jobs=10, debug=False, bits_per_sample=16, mfcc_conf=None):

        ''' model_path_list is a list object and each item is the location of a gmm identity (XXX-identity.gmm)
        '''

        ''' make model path being absolute path
        '''
        model_path_list = [os.path.abspath(path) for path in model_path_list]

        n_audios = len(audio_list)
        if n_jobs > n_audios:
            n_jobs = n_audios

        audio_path_list = self.write_audio(audio_list, fs=fs)

        self.data_prepare(audio_path_list, debug=debug)
        self.make_mfcc(n_jobs=n_jobs, mfcc_conf=mfcc_conf, debug=debug)
        self.compute_vad(n_jobs=n_jobs, debug=debug)
        self.get_frames_likes(model_path_list, n_jobs=n_jobs, debug=debug)
        score_array = self.resolce_scores(model_path_list)

        return score_array

max_issues_repo_name: tom-doerr/FAKEBOB

max_forks_repo_name: tom-doerr/FAKEBOB