<commit_before>"""Fairly basic set of tools for real-time data augmentation on image data.
Can easily be extended to include new transformations,
new preprocessing methods, etc...
"""
from __future__ import absolute_import
from __future__ import print_function

from collections import deque
import warnings

import numpy as np
import scipy.ndimage as ndi

try:
    from PIL import Image as pil_image
except ImportError:
    pil_image = None

from . import Augmenter

def flip_axis(x, axis):
    x = np.asarray(x).swapaxes(axis, 0)
    x = x[::-1, ...]
    x = x.swapaxes(0, axis)
    return x

def random_channel_shift(x, intensity, channel_axis=0):
    x = np.rollaxis(x, channel_axis, 0)
    min_x, max_x = np.min(x), np.max(x)
    channel_images = [
        np.clip(x_channel + np.random.uniform(-intensity, intensity), min_x,
                max_x) for x_channel in x
    ]
    x = np.stack(channel_images, axis=0)
    x = np.rollaxis(x, 0, channel_axis + 1)
    return x

def transform_matrix_offset_center(matrix, x, y):
    o_x = float(x) / 2 + 0.5
    o_y = float(y) / 2 + 0.5
    offset_matrix = np.array([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])
    reset_matrix = np.array([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])
    transform_matrix = np.dot(np.dot(offset_matrix, matrix), reset_matrix)
    return transform_matrix

def apply_transform(x,
                    transform_matrix,
                    channel_axis=0,
                    fill_mode='nearest',
                    cval=0.):
    """Apply the image transformation specified by a matrix.
    # Arguments
        x: 2D numpy array, single image.
        transform_matrix: Numpy array specifying the geometric transformation.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
    # Returns
        The transformed version of the input.
    """
    x = np.rollaxis(x, channel_axis, 0)
    final_affine_matrix = transform_matrix[:2, :2]
    final_offset = transform_matrix[:2, 2]
    channel_images = [
        ndi.interpolation.affine_transform(
            x_channel,
            final_affine_matrix,
            final_offset,
            order=0,
            mode=fill_mode,
            cval=cval) for x_channel in x
    ]
    x = np.stack(channel_images, axis=0)
    x = np.rollaxis(x, 0, channel_axis + 1)
    return x

class ImageDataAugmenter(Augmenter):
    """Generate minibatches of image data with real-time data augmentation.
    # Arguments
        samplewise_center: set each sample mean to 0.
        samplewise_std_normalization: divide each input by its std.
        rotation_range: degrees (0 to 180).
        width_shift_range: fraction of total width.
        height_shift_range: fraction of total height.
        shear_range: shear intensity (shear angle in radians).
        zoom_range: amount of zoom. if scalar z, zoom will be randomly picked
            in the range [1-z, 1+z]. A sequence of two can be passed instead
            to select this range.
        channel_shift_range: shift range for each channels.
        fill_mode: points outside the boundaries are filled according to the
            given mode ('constant', 'nearest', 'reflect' or 'wrap'). Default
            is 'nearest'.
        cval: value used for points outside the boundaries when fill_mode is
            'constant'. Default is 0.
        horizontal_flip: whether to randomly flip images horizontally.
        vertical_flip: whether to randomly flip images vertically.
        rescale: rescaling factor. If None or 0, no rescaling is applied,
            otherwise we multiply the data by the value provided. This is
            applied after the `preprocessing_function` (if any provided)
            but before any other transformation.
        preprocessing_function: function that will be implied on each input.
            The function will run before any other modification on it.
            The function should take one argument:
            one image (Numpy tensor with rank 3),
            and should output a Numpy tensor with the same shape.
        data_format: 'channels_first' or 'channels_last'. In 'channels_first'
            mode, the channels dimension (the depth) is at index 1, in
            'channels_last' mode it is at index 3. It defaults to
            "channels_last".
    """

    def __init__(self,
                 labels=True,
                 augment_labels=False,
                 samplewise_center=False,
                 samplewise_std_normalization=False,
                 rotation_range=0.,
                 width_shift_range=0.,
                 height_shift_range=0.,
                 shear_range=0.,
                 zoom_range=0.,
                 channel_shift_range=0.,
                 fill_mode='nearest',
                 cval=0.,
                 horizontal_flip=False,
                 vertical_flip=False,
                 rescale=None,
                 preprocessing_function=None,
                 seed=None,
                 data_format='channels_last',
                 save_inverses=False):

        super(ImageDataAugmenter, self).__init__(labels, augment_labels)
        self.samplewise_center = samplewise_center
        self.samplewise_std_normalization = samplewise_std_normalization
        self.rotation_range = rotation_range
        self.width_shift_range = width_shift_range
        self.height_shift_range = height_shift_range
        self.shear_range = shear_range
        self.zoom_range = zoom_range
        self.channel_shift_range = channel_shift_range
        self.fill_mode = fill_mode
        self.cval = cval
        self.horizontal_flip = horizontal_flip
        self.vertical_flip = vertical_flip
        self.rescale = rescale
        self.preprocessing_function = preprocessing_function
        self.save_inverses = save_inverses
        self.inverse_transforms = deque()

        if self.channel_shift_range != 0 and self.save_inverses:
            warnings.warn("Image augmenter cannot invert channel shifting")
        if self.augment_masks and self.save_inverses:
            warnings.warn("Be careful how you invert images. \
                Masks are augmented in batch first before images!")
        if (self.samplewise_center or self.samplewise_std_normalization) \
                and self.save_inverses:
            warnings.warn("Image augmenter cannot invert samplewise mean"
                          " and std normalization!")

        if data_format not in {'channels_last', 'channels_first'}:
            raise ValueError("`data_format` should be `channels_last` (channel\
             after row and column) or `channels_first` (channel before row and\
             column). Received arg: {} data_format".format_map(data_format))
        self.data_format = data_format
        if data_format == 'channels_first':
            self.channel_axis = 1
            self.row_axis = 2
            self.col_axis = 3
        if data_format == 'channels_last':
            self.channel_axis = 3
            self.row_axis = 1
            self.col_axis = 2

        self.mean = None
        self.std = None
        self.principal_components = None

        if np.isscalar(zoom_range):
            self.zoom_range = [1 - zoom_range, 1 + zoom_range]
        elif len(zoom_range) == 2:
            self.zoom_range = [zoom_range[0], zoom_range[1]]
        else:
            raise ValueError(
                '`zoom_range` should be a float or '
                'a tuple or list of two floats. '
                'Received arg: ', zoom_range)
        if seed is not None:
            np.random.seed(seed)

    def augment(self, x):
        for img in x:
            assert img.ndim == 3, "Each image must be a 3-dimensional array."
            "Your input is {}".format_map(img.ndim)
            self.random_transform(x, np.random.randint(2**32))

    def standardize(self, x):
        """Apply the normalization configuration to a batch of inputs.
        # Arguments
            x: batch of inputs to be normalized.
        # Returns
            The inputs, normalized.
        """
        if self.preprocessing_function:
            x = self.preprocessing_function(x)
        if self.rescale:
            x *= self.rescale
        # x is a single image, so it doesn't have image number at index 0
        img_channel_axis = self.channel_axis - 1
        # Check if x has any channels and is the right dimensions
        remove_channel_axis = False
        if x.ndim == 2:
            remove_channel_axis = True
            x = np.expand_dims(x, axis=img_channel_axis)
        elif x.ndim != 3:
            raise ValueError("Dim of input image must be 2 or 3, given ",
                             x.ndim)
        if self.samplewise_center:
            x -= np.mean(x, axis=img_channel_axis, keepdims=True)
        if self.samplewise_std_normalization:
            x /= (np.std(x, axis=img_channel_axis, keepdims=True) + 1e-7)

        # If we added a dimension for the channel, remove it
        if remove_channel_axis:
            x = np.squeeze(x, axis=img_channel_axis)
        return x

    def random_transform(self, x, seed=None):
        """Randomly augment a single image tensor.
        # Arguments
            x: 3D tensor, single image.
            seed: random seed.
        # Returns
            A randomly transformed version of the input (same shape).
        """
        # x is a single image, so it doesn't have image number at index 0
        img_row_axis = self.row_axis - 1
        img_col_axis = self.col_axis - 1
        img_channel_axis = self.channel_axis - 1

        # Check if x has any channels and is the right dimensions
        remove_channel_axis = False
        if x.ndim == 2:
            remove_channel_axis = True
            x = np.expand_dims(x, axis=img_channel_axis)
        elif x.ndim != 3:
            raise ValueError("Dim of input image must be 2 or 3, given ",
                             x.ndim)

        if seed is not None:
            np.random.seed(seed)

        # use composition of homographies
        # to generate final transform that needs to be applied
        if self.rotation_range:
            theta = np.pi / 180 * \
                np.random.uniform(-self.rotation_range, self.rotation_range)
        else:
            theta = 0

        if self.height_shift_range:
            tx = np.random.uniform(-self.height_shift_range,
                                   self.height_shift_range) * \
                                   x.shape[img_row_axis]
        else:
            tx = 0

        if self.width_shift_range:
            ty = np.random.uniform(
                -self.width_shift_range,
                self.width_shift_range) * x.shape[img_col_axis]
        else:
            ty = 0

        if self.shear_range:
            shear = np.random.uniform(-self.shear_range, self.shear_range)
        else:
            shear = 0

        if self.zoom_range[0] == 1 and self.zoom_range[1] == 1:
            zx, zy = 1, 1
        else:
            zx, zy = np.random.uniform(self.zoom_range[0], self.zoom_range[1],
                                       2)

        transform_matrix = None
        if theta != 0:
            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
                                        [np.sin(theta),
                                         np.cos(theta), 0], [0, 0, 1]])
            transform_matrix = rotation_matrix

        if tx != 0 or ty != 0:
            shift_matrix = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]])
            transform_matrix = shift_matrix if transform_matrix is None else \
                np.dot(transform_matrix, shift_matrix)

        if shear != 0:
            shear_matrix = np.array([[1, -np.sin(shear), 0],
                                     [0, np.cos(shear), 0], [0, 0, 1]])
            transform_matrix = shear_matrix if transform_matrix is None else \
                np.dot(transform_matrix, shear_matrix)

        if zx != 1 or zy != 1:
            zoom_matrix = np.array([[zx, 0, 0], [0, zy, 0], [0, 0, 1]])
            transform_matrix = zoom_matrix if transform_matrix is None else \
                np.dot(transform_matrix, zoom_matrix)

        inverse_transform = None
        if transform_matrix is not None:
            h, w = x.shape[img_row_axis], x.shape[img_col_axis]
            transform_matrix = transform_matrix_offset_center(
                transform_matrix, h, w)
            x = apply_transform(
                x,
                transform_matrix,
                img_channel_axis,
                fill_mode=self.fill_mode,
                cval=self.cval)
            inverse_transform = np.linalg.inv(transform_matrix)

        if self.channel_shift_range != 0:
            x = random_channel_shift(x, self.channel_shift_range,
                                     img_channel_axis)

        horizontal_flipped = False
        if self.horizontal_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_col_axis)
                horizontal_flipped = True

        vertically_flipped = False
        if self.vertical_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_row_axis)
                vertically_flipped = True

        # If we added a dimension for the channel, remove it
        if remove_channel_axis:
            x = np.squeeze(x, axis=img_channel_axis)

        # If we want to save inverses, push onto the inverse queue
        if self.save_inverses:
            self.inverse_transforms.appendleft(
                (inverse_transform, horizontal_flipped, vertically_flipped))
        return x

    def apply_inversion_transform(self, x):
        """
        Invert a random augment of a single image tensor.
        # Arguments
            x: 3D tensor, single image.
        # Returns
            The orignal image before the random transformed version
            of the input (same shape).
        """
        # x is a single image, so it doesn't have image number at index 0
        img_row_axis = self.row_axis - 1
        img_col_axis = self.col_axis - 1
        img_channel_axis = self.channel_axis - 1

        # Check if x has any channels and is the right dimensions
        remove_channel_axis = False
        if x.ndim == 2:
            remove_channel_axis = True
            x = np.expand_dims(x, axis=img_channel_axis)
        elif x.ndim != 3:
            raise ValueError("Dim of input image must be 2 or 3, given ",
                             x.ndim)

        # Pop from the queue
        inverse_transform, horizontally_flipped, vertically_flipped = \
            self.inverse_transforms.pop()
        # Undo any flipping
        if vertically_flipped:
            x = flip_axis(x, img_row_axis)

        if horizontally_flipped:
            x = flip_axis(x, img_col_axis)

        # Apply the inverse transform
        x = apply_transform(
            x,
            inverse_transform,
            img_channel_axis,
            fill_mode=self.fill_mode,
            cval=self.cval)

        # Remove the channel axis if we added it
        if remove_channel_axis:
            x = np.squeeze(x, img_channel_axis)

        return x

    def invert_images(self, x_images):
        if len(x_images) != len(self.inverse_transforms):
            warnings.warn("Have %s images, but only %s inverse transforms. "
                          "Using first transforms and discarding rest" %
                          (len(x_images), len(self.inverse_transforms)))
        inverted = np.stack(
            [self.apply_inversion_transform(x) for x in x_images], axis=0)
        self.inverse_transforms = deque()
        return inverted
<commit_msg>bugfix augment_masks undefined<commit_after>"""Fairly basic set of tools for real-time data augmentation on image data.
Can easily be extended to include new transformations,
new preprocessing methods, etc...
"""
from __future__ import absolute_import
from __future__ import print_function

from collections import deque
import warnings

import numpy as np
import scipy.ndimage as ndi

try:
    from PIL import Image as pil_image
except ImportError:
    pil_image = None

from . import Augmenter

def flip_axis(x, axis):
    x = np.asarray(x).swapaxes(axis, 0)
    x = x[::-1, ...]
    x = x.swapaxes(0, axis)
    return x

def random_channel_shift(x, intensity, channel_axis=0):
    x = np.rollaxis(x, channel_axis, 0)
    min_x, max_x = np.min(x), np.max(x)
    channel_images = [
        np.clip(x_channel + np.random.uniform(-intensity, intensity), min_x,
                max_x) for x_channel in x
    ]
    x = np.stack(channel_images, axis=0)
    x = np.rollaxis(x, 0, channel_axis + 1)
    return x

def transform_matrix_offset_center(matrix, x, y):
    o_x = float(x) / 2 + 0.5
    o_y = float(y) / 2 + 0.5
    offset_matrix = np.array([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])
    reset_matrix = np.array([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])
    transform_matrix = np.dot(np.dot(offset_matrix, matrix), reset_matrix)
    return transform_matrix

def apply_transform(x,
                    transform_matrix,
                    channel_axis=0,
                    fill_mode='nearest',
                    cval=0.):
    """Apply the image transformation specified by a matrix.
    # Arguments
        x: 2D numpy array, single image.
        transform_matrix: Numpy array specifying the geometric transformation.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
    # Returns
        The transformed version of the input.
    """
    x = np.rollaxis(x, channel_axis, 0)
    final_affine_matrix = transform_matrix[:2, :2]
    final_offset = transform_matrix[:2, 2]
    channel_images = [
        ndi.interpolation.affine_transform(
            x_channel,
            final_affine_matrix,
            final_offset,
            order=0,
            mode=fill_mode,
            cval=cval) for x_channel in x
    ]
    x = np.stack(channel_images, axis=0)
    x = np.rollaxis(x, 0, channel_axis + 1)
    return x

class ImageDataAugmenter(Augmenter):
    """Generate minibatches of image data with real-time data augmentation.
    # Arguments
        samplewise_center: set each sample mean to 0.
        samplewise_std_normalization: divide each input by its std.
        rotation_range: degrees (0 to 180).
        width_shift_range: fraction of total width.
        height_shift_range: fraction of total height.
        shear_range: shear intensity (shear angle in radians).
        zoom_range: amount of zoom. if scalar z, zoom will be randomly picked
            in the range [1-z, 1+z]. A sequence of two can be passed instead
            to select this range.
        channel_shift_range: shift range for each channels.
        fill_mode: points outside the boundaries are filled according to the
            given mode ('constant', 'nearest', 'reflect' or 'wrap'). Default
            is 'nearest'.
        cval: value used for points outside the boundaries when fill_mode is
            'constant'. Default is 0.
        horizontal_flip: whether to randomly flip images horizontally.
        vertical_flip: whether to randomly flip images vertically.
        rescale: rescaling factor. If None or 0, no rescaling is applied,
            otherwise we multiply the data by the value provided. This is
            applied after the `preprocessing_function` (if any provided)
            but before any other transformation.
        preprocessing_function: function that will be implied on each input.
            The function will run before any other modification on it.
            The function should take one argument:
            one image (Numpy tensor with rank 3),
            and should output a Numpy tensor with the same shape.
        data_format: 'channels_first' or 'channels_last'. In 'channels_first'
            mode, the channels dimension (the depth) is at index 1, in
            'channels_last' mode it is at index 3. It defaults to
            "channels_last".
    """

    def __init__(self,
                 labels=True,
                 augment_labels=False,
                 samplewise_center=False,
                 samplewise_std_normalization=False,
                 rotation_range=0.,
                 width_shift_range=0.,
                 height_shift_range=0.,
                 shear_range=0.,
                 zoom_range=0.,
                 channel_shift_range=0.,
                 fill_mode='nearest',
                 cval=0.,
                 horizontal_flip=False,
                 vertical_flip=False,
                 rescale=None,
                 preprocessing_function=None,
                 seed=None,
                 data_format='channels_last',
                 save_inverses=False):

        super(ImageDataAugmenter, self).__init__(labels, augment_labels)
        self.samplewise_center = samplewise_center
        self.samplewise_std_normalization = samplewise_std_normalization
        self.rotation_range = rotation_range
        self.width_shift_range = width_shift_range
        self.height_shift_range = height_shift_range
        self.shear_range = shear_range
        self.zoom_range = zoom_range
        self.channel_shift_range = channel_shift_range
        self.fill_mode = fill_mode
        self.cval = cval
        self.horizontal_flip = horizontal_flip
        self.vertical_flip = vertical_flip
        self.rescale = rescale
        self.preprocessing_function = preprocessing_function
        self.save_inverses = save_inverses
        self.inverse_transforms = deque()

        if self.channel_shift_range != 0 and self.save_inverses:
            warnings.warn("Image augmenter cannot invert channel shifting")
        if self.augment_labels and self.save_inverses:
            warnings.warn("Be careful how you invert images. \
                Masks are augmented in batch first before images!")
        if (self.samplewise_center or self.samplewise_std_normalization) \
                and self.save_inverses:
            warnings.warn("Image augmenter cannot invert samplewise mean"
                          " and std normalization!")

        if data_format not in {'channels_last', 'channels_first'}:
            raise ValueError("`data_format` should be `channels_last` (channel\
             after row and column) or `channels_first` (channel before row and\
             column). Received arg: {} data_format".format_map(data_format))
        self.data_format = data_format
        if data_format == 'channels_first':
            self.channel_axis = 1
            self.row_axis = 2
            self.col_axis = 3
        if data_format == 'channels_last':
            self.channel_axis = 3
            self.row_axis = 1
            self.col_axis = 2

        self.mean = None
        self.std = None
        self.principal_components = None

        if np.isscalar(zoom_range):
            self.zoom_range = [1 - zoom_range, 1 + zoom_range]
        elif len(zoom_range) == 2:
            self.zoom_range = [zoom_range[0], zoom_range[1]]
        else:
            raise ValueError(
                '`zoom_range` should be a float or '
                'a tuple or list of two floats. '
                'Received arg: ', zoom_range)
        if seed is not None:
            np.random.seed(seed)

    def augment(self, x):
        for img in x:
            assert img.ndim == 3, "Each image must be a 3-dimensional array."
            "Your input is {}".format_map(img.ndim)
            self.random_transform(x, np.random.randint(2**32))

    def standardize(self, x):
        """Apply the normalization configuration to a batch of inputs.
        # Arguments
            x: batch of inputs to be normalized.
        # Returns
            The inputs, normalized.
        """
        if self.preprocessing_function:
            x = self.preprocessing_function(x)
        if self.rescale:
            x *= self.rescale
        # x is a single image, so it doesn't have image number at index 0
        img_channel_axis = self.channel_axis - 1
        # Check if x has any channels and is the right dimensions
        remove_channel_axis = False
        if x.ndim == 2:
            remove_channel_axis = True
            x = np.expand_dims(x, axis=img_channel_axis)
        elif x.ndim != 3:
            raise ValueError("Dim of input image must be 2 or 3, given ",
                             x.ndim)
        if self.samplewise_center:
            x -= np.mean(x, axis=img_channel_axis, keepdims=True)
        if self.samplewise_std_normalization:
            x /= (np.std(x, axis=img_channel_axis, keepdims=True) + 1e-7)

        # If we added a dimension for the channel, remove it
        if remove_channel_axis:
            x = np.squeeze(x, axis=img_channel_axis)
        return x

    def random_transform(self, x, seed=None):
        """Randomly augment a single image tensor.
        # Arguments
            x: 3D tensor, single image.
            seed: random seed.
        # Returns
            A randomly transformed version of the input (same shape).
        """
        # x is a single image, so it doesn't have image number at index 0
        img_row_axis = self.row_axis - 1
        img_col_axis = self.col_axis - 1
        img_channel_axis = self.channel_axis - 1

        # Check if x has any channels and is the right dimensions
        remove_channel_axis = False
        if x.ndim == 2:
            remove_channel_axis = True
            x = np.expand_dims(x, axis=img_channel_axis)
        elif x.ndim != 3:
            raise ValueError("Dim of input image must be 2 or 3, given ",
                             x.ndim)

        if seed is not None:
            np.random.seed(seed)

        # use composition of homographies
        # to generate final transform that needs to be applied
        if self.rotation_range:
            theta = np.pi / 180 * \
                np.random.uniform(-self.rotation_range, self.rotation_range)
        else:
            theta = 0

        if self.height_shift_range:
            tx = np.random.uniform(-self.height_shift_range,
                                   self.height_shift_range) * \
                                   x.shape[img_row_axis]
        else:
            tx = 0

        if self.width_shift_range:
            ty = np.random.uniform(
                -self.width_shift_range,
                self.width_shift_range) * x.shape[img_col_axis]
        else:
            ty = 0

        if self.shear_range:
            shear = np.random.uniform(-self.shear_range, self.shear_range)
        else:
            shear = 0

        if self.zoom_range[0] == 1 and self.zoom_range[1] == 1:
            zx, zy = 1, 1
        else:
            zx, zy = np.random.uniform(self.zoom_range[0], self.zoom_range[1],
                                       2)

        transform_matrix = None
        if theta != 0:
            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
                                        [np.sin(theta),
                                         np.cos(theta), 0], [0, 0, 1]])
            transform_matrix = rotation_matrix

        if tx != 0 or ty != 0:
            shift_matrix = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]])
            transform_matrix = shift_matrix if transform_matrix is None else \
                np.dot(transform_matrix, shift_matrix)

        if shear != 0:
            shear_matrix = np.array([[1, -np.sin(shear), 0],
                                     [0, np.cos(shear), 0], [0, 0, 1]])
            transform_matrix = shear_matrix if transform_matrix is None else \
                np.dot(transform_matrix, shear_matrix)

        if zx != 1 or zy != 1:
            zoom_matrix = np.array([[zx, 0, 0], [0, zy, 0], [0, 0, 1]])
            transform_matrix = zoom_matrix if transform_matrix is None else \
                np.dot(transform_matrix, zoom_matrix)

        inverse_transform = None
        if transform_matrix is not None:
            h, w = x.shape[img_row_axis], x.shape[img_col_axis]
            transform_matrix = transform_matrix_offset_center(
                transform_matrix, h, w)
            x = apply_transform(
                x,
                transform_matrix,
                img_channel_axis,
                fill_mode=self.fill_mode,
                cval=self.cval)
            inverse_transform = np.linalg.inv(transform_matrix)

        if self.channel_shift_range != 0:
            x = random_channel_shift(x, self.channel_shift_range,
                                     img_channel_axis)

        horizontal_flipped = False
        if self.horizontal_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_col_axis)
                horizontal_flipped = True

        vertically_flipped = False
        if self.vertical_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_row_axis)
                vertically_flipped = True

        # If we added a dimension for the channel, remove it
        if remove_channel_axis:
            x = np.squeeze(x, axis=img_channel_axis)

        # If we want to save inverses, push onto the inverse queue
        if self.save_inverses:
            self.inverse_transforms.appendleft(
                (inverse_transform, horizontal_flipped, vertically_flipped))
        return x

    def apply_inversion_transform(self, x):
        """
        Invert a random augment of a single image tensor.
        # Arguments
            x: 3D tensor, single image.
        # Returns
            The orignal image before the random transformed version
            of the input (same shape).
        """
        # x is a single image, so it doesn't have image number at index 0
        img_row_axis = self.row_axis - 1
        img_col_axis = self.col_axis - 1
        img_channel_axis = self.channel_axis - 1

        # Check if x has any channels and is the right dimensions
        remove_channel_axis = False
        if x.ndim == 2:
            remove_channel_axis = True
            x = np.expand_dims(x, axis=img_channel_axis)
        elif x.ndim != 3:
            raise ValueError("Dim of input image must be 2 or 3, given ",
                             x.ndim)

        # Pop from the queue
        inverse_transform, horizontally_flipped, vertically_flipped = \
            self.inverse_transforms.pop()
        # Undo any flipping
        if vertically_flipped:
            x = flip_axis(x, img_row_axis)

        if horizontally_flipped:
            x = flip_axis(x, img_col_axis)

        # Apply the inverse transform
        x = apply_transform(
            x,
            inverse_transform,
            img_channel_axis,
            fill_mode=self.fill_mode,
            cval=self.cval)

        # Remove the channel axis if we added it
        if remove_channel_axis:
            x = np.squeeze(x, img_channel_axis)

        return x

    def invert_images(self, x_images):
        if len(x_images) != len(self.inverse_transforms):
            warnings.warn("Have %s images, but only %s inverse transforms. "
                          "Using first transforms and discarding rest" %
                          (len(x_images), len(self.inverse_transforms)))
        inverted = np.stack(
            [self.apply_inversion_transform(x) for x in x_images], axis=0)
        self.inverse_transforms = deque()
        return inverted
<|endoftext|>


<commit_before># Check http://docs.python.org/library/time.html

from __pyjamas__ import JS
import math

timezone = JS("60 * (new Date(new Date().getFullYear(), 0, 1)).getTimezoneOffset()")
altzone = JS("60 * (new Date(new Date().getFullYear(), 6, 1)).getTimezoneOffset()")
if altzone > timezone:
    # Probably on southern parth of the earth...
    d = timezone
    timezone = altzone
    altzone = d
d = JS("(new Date(new Date().getFullYear(), 0, 1))")
d = d.toLocaleString().split()[-1]
if d[0] == '(':
    d = d[1:-1]
tzname = (d, None)
del d

__c__days = ["Sunday", "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"]
__c__months = ["January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December"]

def time():
    JS("return new Date().getTime() / 1000.0;")

class struct_time(object):
    n_fields = 9
    n_sequence_fields = 9
    n_unnamed_fields = 0
    tm_year = None
    tm_mon = None
    tm_mday = None
    tm_hour = None
    tm_min = None
    tm_sec = None
    tm_wday = None
    tm_yday = None
    tm_isdst = None

    def __init__(self, ttuple = None):
        if not ttuple is None:
            self.tm_year = ttuple[0]
            self.tm_mon = ttuple[1]
            self.tm_mday = ttuple[2]
            self.tm_hour = ttuple[3]
            self.tm_min = ttuple[4]
            self.tm_sec = ttuple[5]
            self.tm_wday = ttuple[6]
            self.tm_yday = ttuple[7]
            self.tm_isdst = ttuple[8]

    def __str__(self):
        t = (
            self.tm_year,
            self.tm_mon,
            self.tm_mday,
            self.tm_hour,
            self.tm_min,
            self.tm_sec,
            self.tm_wday,
            self.tm_yday,
            self.tm_isdst,
        )
        return t.__str__()

    def __getitem__(self, idx):
        return [self.tm_year, self.tm_mon, self.tm_mday,
                self.tm_hour, self.tm_min, self.tm_sec,
                self.tm_wday, self.tm_yday, self.tm_isdst][idx]

    def __getslice__(self, lower, upper):
        return [self.tm_year, self.tm_mon, self.tm_mday,
                self.tm_hour, self.tm_min, self.tm_sec,
                self.tm_wday, self.tm_yday, self.tm_isdst][lower:upper]

def gmtime(t = None):
    if t == None:
        t = time()
    date = JS("new Date(t*1000)")
    tm = struct_time()
    tm.tm_year = date.getUTCFullYear()
    tm.tm_mon = date.getUTCMonth() + 1
    tm.tm_mday = date.getUTCDate()
    tm.tm_hour = date.getUTCHours()
    tm.tm_min = date.getUTCMinutes()
    tm.tm_sec = date.getUTCSeconds()
    tm.tm_wday = (date.getUTCDay() + 6) % 7
    tm.tm_isdst = 0
    startOfYear = JS("new Date('Jan 1 '+tm.tm_year+' GMT+0000')")
    tm.tm_yday = 1 + int((t - startOfYear.getTime()/1000)/86400)
    return tm

def localtime(t = None):
    if t == None:
        t = time()
    date = JS("new Date(t*1000)")
    dateOffset = date.getTimezoneOffset()
    tm = struct_time()
    tm.tm_year = date.getFullYear()
    tm.tm_mon = date.getMonth() + 1
    tm.tm_mday = date.getDate()
    tm.tm_hour = date.getHours()
    tm.tm_min = date.getMinutes()
    tm.tm_sec = date.getSeconds()
    tm.tm_wday = (date.getDay() + 6) % 7
    tm.tm_isdst = 0 if timezone == 60*date.getTimezoneOffset() else 1
    startOfYear = JS("new Date(tm.tm_year,0,1)") # local time
    startOfYearOffset = startOfYear.getTimezoneOffset()
    startOfDay = JS("new Date(tm.tm_year,tm.tm_mon-1,tm.tm_mday)")
    dt = (startOfDay.getTime() - startOfYear.getTime())/1000
    dt = dt + 60 * (startOfYearOffset - dateOffset)
    tm.tm_yday = 1 + int(dt/86400.0)
    return tm
    if startOfYearOffset != dateOffset:
        # Changed from std to dst or the opposite
        #if startOfYearOffset > dateOffset:
        # Changed from std to dst
        tm.tm_yday += 1
        dt2 = dt + 60 * (startOfYearOffset - dateOffset)
        print dt, dt/86400.0, (startOfYearOffset, dateOffset), dt2, dt2/86400.0
        tm.tm_yday = 1 + int(dt2/86400.0)
    #if tm.tm_isdst and 60*startOfYearOffset == timezone:
    #    tm.tm_yday += 1
    return tm

def mktime(t):
    tm_year = t[0]
    tm_mon = t[1] - 1
    tm_mday = t[2]
    tm_hour = t[3]
    tm_min = t[4]
    tm_sec = t[5]
    date = JS("new Date(tm_year, tm_mon, tm_mday, tm_hour, tm_min, tm_sec)") # local time
    if t[8] == 0:
        return date.getTime()/1000 - 60 * date.getTimezoneOffset()
    #return date.getTime()/1000 - 60 * date.getTimezoneOffset() + (60 * date.getTimezoneOffset() - timezone)
    return date.getTime()/1000 - timezone

def strftime(fmt, t = None):
    if t is None:
        t = localtime()
    else:
        if not isinstance(t, struct_time) and len(t) != 9:
            raise TypeError('argument must be 9-item sequence, not float')
    tm_year = t[0]
    tm_mon = t[1]
    tm_mday = t[2]
    tm_hour = t[3]
    tm_min = t[4]
    tm_sec = t[5]
    tm_wday = t[6]
    tm_yday = t[7]
    date = JS("new Date(tm_year, tm_mon - 1, tm_mday, tm_hour, tm_min, tm_sec)")
    startOfYear = JS("new Date(tm_year,0,1)")
    firstMonday = 1 - ((startOfYear.getDay() + 6) % 7) + 7
    firstWeek = JS("new Date(tm_year,0,firstMonday)")
    weekNo = date.getTime() - firstWeek.getTime()
    if weekNo < 0:
        weekNo = 0
    else:
        weekNo = 1 + int(weekNo/604800000)

    def format(c):
        if c == '%':
            return '%'
        elif c == 'a':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'A':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'b':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'B':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'c':
            return date.toLocaleString()
        elif c == 'd':
            return "%02d" % tm_mday
        elif c == 'H':
            return "%02d" % tm_hour
        elif c == 'I':
            return "%02d" % (tm_hour % 12)
        elif c == 'j':
            return "%03d" % tm_yday
        elif c == 'm':
            return "%02d" % tm_mon
        elif c == 'M':
            return "%02d" % tm_min
        elif c == 'p': # FIXME: should be locale dependent
            if tm_hour < 12:
                return "AM"
            return "PM"
        elif c == 'S':
            return "%02d" % tm_sec
        elif c == 'U':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'w':
            return "%d" % ((tm_wday+1) % 7)
        elif c == 'W':
            return "%d" % weekNo
        elif c == 'x':
            return "%s" % date.toLocaleDateString()
        elif c == 'X':
            return "%s" % date.toLocaleTimeString()
        elif c == 'y':
            return "%02d" % (tm_year % 100)
        elif c == 'Y':
            return "%04d" % tm_year
        elif c == 'Z':
            raise NotImplementedError("strftime format character '%s'" % c)
        return "%" + c
    result = ''
    remainder = fmt
    re_pct = JS("/([^%]*)%(.)(.*)/")
    JS("var a, fmtChar;")
    while remainder:
        JS("""
        a = re_pct.exec(remainder);
        if (!a) {
            result += remainder;
            remainder = null;
        } else {
            result += a[1];
            fmtChar = a[2];
            remainder = a[3];
            if (typeof fmtChar != 'undefined') {
                result += format(fmtChar);
            }
        }
        """)
    return result

def asctime(t = None):
    if t == None:
        t = localtime()
    return "%s %s %02d %02d:%02d:%02d %04d" % (__c__days[(t[6]+1)%7][:3], __c__months[t[1]-1], t[2], t[3], t[4], t[5], t[0])

def ctime(t = None):
    t = localtime()
    return asctime(t)
<commit_msg>JS @ fixup<commit_after># Check http://docs.python.org/library/time.html

from __pyjamas__ import JS
import math

timezone = JS("60 * (new Date(new Date().getFullYear(), 0, 1)).getTimezoneOffset()")
altzone = JS("60 * (new Date(new Date().getFullYear(), 6, 1)).getTimezoneOffset()")
if altzone > timezone:
    # Probably on southern parth of the earth...
    d = timezone
    timezone = altzone
    altzone = d
d = JS("(new Date(new Date().getFullYear(), 0, 1))")
d = d.toLocaleString().split()[-1]
if d[0] == '(':
    d = d[1:-1]
tzname = (d, None)
del d

__c__days = ["Sunday", "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"]
__c__months = ["January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December"]

def time():
    JS("return new Date().getTime() / 1000.0;")

class struct_time(object):
    n_fields = 9
    n_sequence_fields = 9
    n_unnamed_fields = 0
    tm_year = None
    tm_mon = None
    tm_mday = None
    tm_hour = None
    tm_min = None
    tm_sec = None
    tm_wday = None
    tm_yday = None
    tm_isdst = None

    def __init__(self, ttuple = None):
        if not ttuple is None:
            self.tm_year = ttuple[0]
            self.tm_mon = ttuple[1]
            self.tm_mday = ttuple[2]
            self.tm_hour = ttuple[3]
            self.tm_min = ttuple[4]
            self.tm_sec = ttuple[5]
            self.tm_wday = ttuple[6]
            self.tm_yday = ttuple[7]
            self.tm_isdst = ttuple[8]

    def __str__(self):
        t = (
            self.tm_year,
            self.tm_mon,
            self.tm_mday,
            self.tm_hour,
            self.tm_min,
            self.tm_sec,
            self.tm_wday,
            self.tm_yday,
            self.tm_isdst,
        )
        return t.__str__()

    def __getitem__(self, idx):
        return [self.tm_year, self.tm_mon, self.tm_mday,
                self.tm_hour, self.tm_min, self.tm_sec,
                self.tm_wday, self.tm_yday, self.tm_isdst][idx]

    def __getslice__(self, lower, upper):
        return [self.tm_year, self.tm_mon, self.tm_mday,
                self.tm_hour, self.tm_min, self.tm_sec,
                self.tm_wday, self.tm_yday, self.tm_isdst][lower:upper]

def gmtime(t = None):
    if t == None:
        t = time()
    date = JS("new Date(@{{t}}*1000)")
    tm = struct_time()
    tm.tm_year = date.getUTCFullYear()
    tm.tm_mon = date.getUTCMonth() + 1
    tm.tm_mday = date.getUTCDate()
    tm.tm_hour = date.getUTCHours()
    tm.tm_min = date.getUTCMinutes()
    tm.tm_sec = date.getUTCSeconds()
    tm.tm_wday = (date.getUTCDay() + 6) % 7
    tm.tm_isdst = 0
    startOfYear = JS("new Date('Jan 1 '+ @{{tm.tm_year}} +' GMT+0000')")
    tm.tm_yday = 1 + int((t - startOfYear.getTime()/1000)/86400)
    return tm

def localtime(t = None):
    if t == None:
        t = time()
    date = JS("new Date(@{{t}}*1000)")
    dateOffset = date.getTimezoneOffset()
    tm = struct_time()
    tm.tm_year = date.getFullYear()
    tm.tm_mon = date.getMonth() + 1
    tm.tm_mday = date.getDate()
    tm.tm_hour = date.getHours()
    tm.tm_min = date.getMinutes()
    tm.tm_sec = date.getSeconds()
    tm.tm_wday = (date.getDay() + 6) % 7
    tm.tm_isdst = 0 if timezone == 60*date.getTimezoneOffset() else 1
    startOfYear = JS("new Date(@{{tm.tm_year}},0,1)") # local time
    startOfYearOffset = startOfYear.getTimezoneOffset()
    startOfDay = JS("new Date(@{{tm.tm_year}},@{{tm.tm_mon}}-1,@{{tm.tm_mday}})")
    dt = (startOfDay.getTime() - startOfYear.getTime())/1000
    dt = dt + 60 * (startOfYearOffset - dateOffset)
    tm.tm_yday = 1 + int(dt/86400.0)
    return tm
    if startOfYearOffset != dateOffset:
        # Changed from std to dst or the opposite
        #if startOfYearOffset > dateOffset:
        # Changed from std to dst
        tm.tm_yday += 1
        dt2 = dt + 60 * (startOfYearOffset - dateOffset)
        print dt, dt/86400.0, (startOfYearOffset, dateOffset), dt2, dt2/86400.0
        tm.tm_yday = 1 + int(dt2/86400.0)
    #if tm.tm_isdst and 60*startOfYearOffset == timezone:
    #    tm.tm_yday += 1
    return tm

def mktime(t):
    tm_year = t[0]
    tm_mon = t[1] - 1
    tm_mday = t[2]
    tm_hour = t[3]
    tm_min = t[4]
    tm_sec = t[5]
    date = JS("new Date(@{{tm_year}}, @{{tm_mon}}, @{{tm_mday}}, @{{tm_hour}}, @{{tm_min}}, @{{tm_sec}})") # local time
    if t[8] == 0:
        return date.getTime()/1000 - 60 * date.getTimezoneOffset()
    #return date.getTime()/1000 - 60 * date.getTimezoneOffset() + (60 * date.getTimezoneOffset() - timezone)
    return date.getTime()/1000 - timezone

def strftime(fmt, t = None):
    if t is None:
        t = localtime()
    else:
        if not isinstance(t, struct_time) and len(t) != 9:
            raise TypeError('argument must be 9-item sequence, not float')
    tm_year = t[0]
    tm_mon = t[1]
    tm_mday = t[2]
    tm_hour = t[3]
    tm_min = t[4]
    tm_sec = t[5]
    tm_wday = t[6]
    tm_yday = t[7]
    date = JS("new Date(@{{tm_year}}, @{{tm_mon}} - 1, @{{tm_mday}}, @{{tm_hour}}, @{{tm_min}}, @{{tm_sec}})")
    startOfYear = JS("new Date(@{{tm_year}},0,1)")
    firstMonday = 1 - ((startOfYear.getDay() + 6) % 7) + 7
    firstWeek = JS("new Date(@{{tm_year}},0,@{{firstMonday}})")
    weekNo = date.getTime() - firstWeek.getTime()
    if weekNo < 0:
        weekNo = 0
    else:
        weekNo = 1 + int(weekNo/604800000)

    def format(c):
        if c == '%':
            return '%'
        elif c == 'a':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'A':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'b':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'B':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'c':
            return date.toLocaleString()
        elif c == 'd':
            return "%02d" % tm_mday
        elif c == 'H':
            return "%02d" % tm_hour
        elif c == 'I':
            return "%02d" % (tm_hour % 12)
        elif c == 'j':
            return "%03d" % tm_yday
        elif c == 'm':
            return "%02d" % tm_mon
        elif c == 'M':
            return "%02d" % tm_min
        elif c == 'p': # FIXME: should be locale dependent
            if tm_hour < 12:
                return "AM"
            return "PM"
        elif c == 'S':
            return "%02d" % tm_sec
        elif c == 'U':
            raise NotImplementedError("strftime format character '%s'" % c)
        elif c == 'w':
            return "%d" % ((tm_wday+1) % 7)
        elif c == 'W':
            return "%d" % weekNo
        elif c == 'x':
            return "%s" % date.toLocaleDateString()
        elif c == 'X':
            return "%s" % date.toLocaleTimeString()
        elif c == 'y':
            return "%02d" % (tm_year % 100)
        elif c == 'Y':
            return "%04d" % tm_year
        elif c == 'Z':
            raise NotImplementedError("strftime format character '%s'" % c)
        return "%" + c
    result = ''
    remainder = fmt
    re_pct = JS("/([^%]*)%(.)(.*)/")
    JS("var a, fmtChar;")
    while remainder:
        JS("""
        a = @{{re_pct}}.exec(@{{remainder}});
        if (!a) {
            result += @{{remainder}};
            remainder = null;
        } else {
            result += a[1];
            fmtChar = a[2];
            remainder = a[3];
            if (typeof fmtChar != 'undefined') {
                result += @{{format}}(fmtChar);
            }
        }
        """)
    return result

def asctime(t = None):
    if t == None:
        t = localtime()
    return "%s %s %02d %02d:%02d:%02d %04d" % (__c__days[(t[6]+1)%7][:3], __c__months[t[1]-1], t[2], t[3], t[4], t[5], t[0])

def ctime(t = None):
    t = localtime()
    return asctime(t)
<|endoftext|>


<commit_before>"""Utilities related to the compilation of Theano functions."""
import functools

__author__ = "David Warde-Farley"
__copyright__ = "Copyright 2012, David Warde-Farley / Universite de Montreal"
__license__ = "3-clause BSD"
__maintainer__ = "David Warde-Farley"
__email__ = "wardefar@iro"
__all__ = ["compiled_theano_function", "HasCompiledFunctions"]

def compiled_theano_function(fn):
    """
    Method decorator that enables lazy on-demand compilation of
    Theano functions.

    Parameters
    ----------
    fn : bound method
        Method that takes exactly one parameter (i.e. `self`). This
        method should return a compiled Theano function when called.

    Notes
    -----
    This will return an object property (i.e. the `property()` construct)
    that returns a Theano function, and store it in a dictionary attribute
    called `_compiled_functions` on the object. If the function has
    already been accessed it is not compiled again.

    Objects wishing to take advantage of this decorator should inherit
    from `HasCompiledFunctions` in this module to have the object's
    `_compiled_functions` attribute removed upon pickling.

    Examples
    --------
    >>> from pylearn2.utils.compile import compiled_theano
    >>> import theano
    >>> class Foo(object):
    ...     @compiled_theano_function
    ...     def bar(self):
    ...         x = theano.tensor.vector()
    ...         y = theano.tensor.vector()
    ...         print "Compiling..."
    ...         return theano.function([x, y], theano.tensor.dot(x, y))
    ...
    >>> from numpy.random import randn, seed
    >>> o = Foo()
    >>> seed(0)
    >>> xx, yy = randn(50), randn(50)
    >>> o.bar(xx, yy)  # first call, method body will be run
    Compiling...
    array(-3.0349256483108418)
    >>> o.bar(xx, yy)  # function already compiled, no print.
    array(-3.0349256483108418)
    >>> o.bar(randn(5), randn(5))  # different args, still no print
    array(5.294487561729036)
    """
    @functools.wraps(fn)
    def wrapped(self):
        try:
            func = self._compiled_functions[fn.func_name]
        except (AttributeError, KeyError):
            if not hasattr(self, '_compiled_functions'):
                self._compiled_functions = {}
            self._compiled_functions[fn.func_name] = func = fn(self)
        return func
    return property(wrapped)

class HasCompiledFunctions(object):
    """
    Base class/mixin that will automatically strip a `_compiled_functions`
    attribute when pickling.
    """
    def __getstate__(self):
        state = self.__dict__.copy()
        if '_compiled_functions' in state:
            del state['_compiled_functions']
        return state
<commit_msg>Format docstrings in utils/compile.py<commit_after>"""Utilities related to the compilation of Theano functions."""
import functools

__author__ = "David Warde-Farley"
__copyright__ = "Copyright 2012, David Warde-Farley / Universite de Montreal"
__license__ = "3-clause BSD"
__maintainer__ = "David Warde-Farley"
__email__ = "wardefar@iro"
__all__ = ["compiled_theano_function", "HasCompiledFunctions"]

def compiled_theano_function(fn):
    """
    Method decorator that enables lazy on-demand compilation of Theano
    functions.

    Parameters
    ----------
    fn : bound method
        Method that takes exactly one parameter (i.e. `self`). This method \
        should return a compiled Theano function when called.

    Notes
    -----
    This will return an object property (i.e. the `property()` construct)
    that returns a Theano function, and store it in a dictionary attribute
    called `_compiled_functions` on the object. If the function has
    already been accessed it is not compiled again.

    Objects wishing to take advantage of this decorator should inherit
    from `HasCompiledFunctions` in this module to have the object's
    `_compiled_functions` attribute removed upon pickling.

    Examples
    --------
    >>> from pylearn2.utils.compile import compiled_theano
    >>> import theano
    >>> class Foo(object):
    ...     @compiled_theano_function
    ...     def bar(self):
    ...         x = theano.tensor.vector()
    ...         y = theano.tensor.vector()
    ...         print "Compiling..."
    ...         return theano.function([x, y], theano.tensor.dot(x, y))
    ...
    >>> from numpy.random import randn, seed
    >>> o = Foo()
    >>> seed(0)
    >>> xx, yy = randn(50), randn(50)
    >>> o.bar(xx, yy)  # first call, method body will be run
    Compiling...
    array(-3.0349256483108418)
    >>> o.bar(xx, yy)  # function already compiled, no print.
    array(-3.0349256483108418)
    >>> o.bar(randn(5), randn(5))  # different args, still no print
    array(5.294487561729036)
    """
    @functools.wraps(fn)
    def wrapped(self):
        try:
            func = self._compiled_functions[fn.func_name]
        except (AttributeError, KeyError):
            if not hasattr(self, '_compiled_functions'):
                self._compiled_functions = {}
            self._compiled_functions[fn.func_name] = func = fn(self)
        return func
    return property(wrapped)

class HasCompiledFunctions(object):
    """
    Base class/mixin that will automatically strip a `_compiled_functions`
    attribute when pickling.
    """
    def __getstate__(self):
        """
        .. todo::

            WRITEME
        """
        state = self.__dict__.copy()
        if '_compiled_functions' in state:
            del state['_compiled_functions']
        return state
<|endoftext|>


<commit_before><commit_msg>pymatgen/io/qchem/sets.py<commit_after><|endoftext|>


<commit_before>
import math
import logging

from pymoku._instrument import *
from pymoku._oscilloscope import _CoreOscilloscope, VoltsData
from . import _instrument
from ._pid import PID
from ._sweep import SweepGenerator
log = logging.getLogger(__name__)

REGBASE_LLB_DEMOD			= 23
REGBASE_LLB_SCAN			= 32
REGBASE_LLB_AUX_SINE		= 41

REG_LLB_RATE_SEL			= 76

REGBASE_LLB_PID1			= 106
REGBASE_LLB_PID2			= 117

_LLB_PHASESCALE				= 2**64 / 360.0
_LLB_FREQSCALE				= 2**64 / 1e9
_LLB_HIGH_RATE				= 0
_LLB_LOW_RATE				= 1

_LLB_COEFFICIENT_WIDTH		= 24

class LaserLockBox(_CoreOscilloscope):
	def __init__(self):
		super(LaserLockBox).__init__()
		self._register_accessors(_llb_reg_hdl)

		self.id = 16
		self.type = "laserlockbox"

		self.fast_fs = 31.25e6
		seld.slow_fs = 31.25e6

		self.fast_pid = PID(self, reg_base = REGBASE_LLB_PID1, self.fast_fs)
		self.slow_pid = PID(self, reg_base = REGBASE_LLB_PID2, self.slow_fs)

		self.demod_sweep = SweepGenerator(self, reg_base = REGBASE_LLB_DEMOD)
		self.scan_sweep = SweepGenerator(self, reg_base = REGBASE_LLB_SCAN)
		self.aux_sine_sweep = SweepGenerator(self, reg_base = REGBASE_LLB_AUX_SINE)
		SweepGenerator.step = 0
		SweepGenerator.stop = 2**64 -1
		SweepGenerator.duration = 0
		SweepGenerator.waveform = 2
		SweepGenerator.start = 0
		SweepGenerator.wait_for_trig = False
		SweepGenerator.hold_last = False

	@needs_commit
	def set_pid_by_gain(self, pid_block, g=1, kp=1, ki=0, kd=0, si=None, sd=None):
		"""

		Configure the selected PID controller using gain coefficients.

		:type ch: int; [1,2]
		:param ch: Channel of the PID controller to be configured.

		:type g: float; [0,2^16 - 1]
		:param g: Gain

		:type kp: float; [-1e3,1e3]
		:param kp: Proportional gain factor

		:type ki: float;
		:param ki: Integrator gain factor

		:type kd: float;
		:param kd: Differentiator gain factor

		:type si: float; float; [-1e3,1e3]
		:param si: Integrator gain saturation

		:type sd: float; [-1e3,1e3]
		:param sd: Differentiator gain saturation

		:raises InvalidConfigurationException: if the configuration of PID gains is not possible.
		"""
		pid_array = [self.fast_pid, self.slow_pid]
		pid_array[pid_block -1].set_reg_by_gain(g, kp, ki, kd, si, sd)

	@needs_commit
	def set_pid_by_freq(self, pid_block, kp=1, i_xover=None, d_xover=None, si=None, sd=None):
		"""

		Configure the selected PID controller using crossover frequencies.

		:type ch: int; [1,2]
		:param ch: PID controller to  configure

		:type kp: float; [-1e3,1e3]
		:param kp: Proportional gain factor

		:type i_xover: float; [1e-3,1e6] Hz
		:param i_xover: Integrator crossover frequency

		:type d_xover: float; [1,10e6] Hz
		:param d_xover: Differentiator crossover frequency

		:type si: float; float; [-1e3,1e3]
		:param si: Integrator gain saturation

		:type sd: float; [-1e3,1e3]
		:param sd: Differentiator gain saturation

		:raises InvalidConfigurationException: if the configuration of PID gains is not possible.
		"""
		pid_array = [self.fast_pid, self.slow_pid]
		pid_array[pid_block -1].set_reg_by_frequency(kp, i_xover, d_xover, si, sd)

	@needs_commit
	def set_demodulation(self, frequency, phase):
		"""
		Configure the demodulation stage.

		:type frequency : float; [0, 200e6] Hz
		:param frequency : Internal demodulation frequency

		:type phase : float; [0, 360] degrees
		:param phase : float; Internal demodulation phase

		"""
		self.demod_sweep.step = frequency * _LLB_FREQSCALE
		self.demod_phase.start = phase * _LLB_PHASESCALE

	@needs_commit
	def set_scan(self, frequency, phase):
		"""
		Configure the scan signal.

		:type frequency : float; [0, 200e6] Hz
		:param frequency : Internal demodulation frequency

		:type phase : float; [0, 360] degrees
		:param phase : float; Internal demodulation phase

		"""
		self.scan_sweep.step = frequency * _LLB_FREQSCALE
		seld.scan_sweep.phase = phase * _LLB_PHASESCALE

	@needs_commit
	def set_aux_sine(self, frequency, phase):
		"""
		Configure the aux sine signal.

		:type frequency : float; [0, 200e6] Hz
		:param frequency : Internal demodulation frequency

		:type phase : float; [0, 360] degrees
		:param phase : float; Internal demodulation phase

		"""
		self.aux_sine_sweep.step = frequency * _LLB_FREQSCALE
		seld.aux_sine_sweep.phase = phase * _LLB_PHASESCALE

	@needs_commit
	def set_sample_rate(self, rate):
		"""
		Configure the sample rate of the filters and pid controllers of the laser locker.

		selectable rates:
			-**high** : 62.5 MHz
			-**low**  : 31.25 MHz

		:type rate : string; {'high', 'low'}
		:param rate: sample rate

		"""
		_str_to_rate = {
			'high' 	: _LLB_HIGH_RATE,
			'low'	: _LLB_LOW_RATE
		}
		self.rate_sel = _utils.str_to_val(_str_to_rate, rate, 'sampling rate')

# NOTE: This function avoids @needs_commit because it calls _set_mmap_access which requires an immediate commit
	def set_filter(self, filter_coefficients):
		"""
		Set SOS filter sample rate and filter coefficients. This also enables the input and outputs of the specified Moku:Lab channel.

		:type ch: int; {1,2}
		:param ch: target channel

		:type filter_coefficients: array;
		:param filter_coefficients: array containing SOS filter coefficients. Format is described in class documentation above.
		"""
		_utils.check_parameter_valid('set', ch, [1, 2], 'filter channel')
		_utils.check_parameter_valid('set', sample_rate, ['high', 'low'], 'filter sample rate')

		# Conversion of input array (typically generated by Scipy/Matlab) to HDL memory map format
		if filter_coefficients != None:

			# Deep copy to avoid modifying user's original input array
			intermediate_filter = deepcopy(filter_coefficients)

			# Array dimension check
			if len(filter_coefficients) != 5:
				_utils.check_parameter_valid('set', len(filter_coefficients), [5],'number of coefficient array rows')
			for m in range(4):
				if m == 0:
					if len(filter_coefficients[0]) != 1:
						_utils.check_parameter_valid('set', len(filter_coefficients[0]), [1],'number of columns in coefficient array row 0')
				else:
					if len(filter_coefficients[m]) != 6:
						_utils.check_parameter_valid('set', len(filter_coefficients[m]), [6],("number of columns in coefficient array row %s"%(m)))

			# Array values check
			_utils.check_parameter_valid('range', filter_coefficients[0][0], [-8e6,8e6 - 2**(-24)],("coefficient array entry m = %s, n = %s"%(0,0)))
			for m in range(1, 5):
				for n in range(6):
					_utils.check_parameter_valid('range', filter_coefficients[m][n], [-4.0,4.0 - 2**(-45)],("coefficient array entry m = %s, n = %s"%(0,0)))

			# multiply S coefficients into B coefficients and replace all S coefficients with 1.0
			for n in range(1,5):
				intermediate_filter[n][1] *= intermediate_filter[n][0]
				intermediate_filter[n][2] *= intermediate_filter[n][0]
				intermediate_filter[n][3] *= intermediate_filter[n][0]
				intermediate_filter[n][0] = 1.0

			# place gain factor G into S coefficient position 4 to comply with HDL requirements:
			intermediate_filter[4][0] = intermediate_filter[0][0]
			intermediate_filter = intermediate_filter[1:5]

			if ch == 1:
				self.filter_ch1 = intermediate_filter
			else:
				self.filter_ch2 = intermediate_filter

		# combine both filter arrays:
		filter_coeffs = [[0.0]*6]*4
		coeff_list = [ [ [0 for k in range(2)] for x in range(6)] for y in range(8) ]
		for n in range(4):
		 	filter_coeffs[n] = self.filter_ch1[n] + self.filter_ch2[n]

		for stage in range(4):
			for coeff in range(6):
				if coeff == 0:
					coeff_list[stage][coeff] = int(round( 2**(_LLB_COEFFICIENT_WIDTH - 24) * filter_coeffs[x][y + k*6]))
				else:
					coeff_list[stage][coeff] = int(round( 2**(_LLB_COEFFICIENT_WIDTH - 3) * filter_coeffs[x][y + k*6]))

		with open('.data.dat', 'wb') as f:
			for coeff in range(6):
				for stage in range(4):
					f.write(struct.pack('<q', coeff_list[stage][coeff]))

		self._set_mmap_access(True)
		self._moku._send_file('j', '.data.dat')
		self._set_mmap_access(False)
		os.remove('.data.dat')

		# Enable the output and input of the set channel
		if ch==1:
			self.output_en1 = True
			self.input_en1 = True
		else:
			self.output_en2 = True
			self.input_en2 = True

		# Manually commit the above register settings as @needs_commit is not used in this function
		self.commit()

	def _write_coeffs(self, ch, coeffs):
		assert ch in [1,2], "Invalid channel"
		assert len(coeffs) <= _FIR_NUM_BLOCKS * _FIR_BLOCK_SIZE, "Invalid number of filter coefficients."

		coeffs = list(coeffs)

		# Create a list of coefficients in each FIR block
		n = int(math.ceil(len(coeffs)/float(_FIR_NUM_BLOCKS)))
		blocks = [coeffs[x:x+n] for x in range(0, len(coeffs), n)]
		blocks += [[]] * (_FIR_NUM_BLOCKS - len(blocks))

		# Construct a bytearray from the FIR block contents
		coeff_bytes = bytearray()
		for b in blocks:
			b.reverse()
			coeff_bytes += bytearray(struct.pack('<I', len(b)))
			coeff_bytes += bytearray(struct.pack('<' + 'i'*len(b), *[int(round((2.0**24-1) * c)) for c in b]))
			coeff_bytes += bytearray(b'\x00'*4*(_FIR_BLOCK_SIZE-len(b)))

		# Sanity check the coefficient byte array length
		assert len(coeff_bytes) == (_FIR_BLOCK_SIZE+1) * _FIR_NUM_BLOCKS * 4, "Invalid length for FIR coefficient memory map."

		# Write the coefficients to the FIR coefficient memory map
		self._set_mmap_access(True)
		self._moku._send_file_bytes('j', '', coeff_bytes, offset=_FIR_MMAP_BLOCK_SIZE*(ch-1))
		self._set_mmap_access(False)

		# Release the memory map "file" to other resources
		self._moku._fs_finalise('j', '', _FIR_MMAP_BLOCK_SIZE*2)

_llb_reg_hdl = {
	'rate_sel':		(REG_LLB_RATE_SEL,	to_reg_unsigned(0, 1),
										from_reg_unsigned(0, 1)),
}<commit_msg>HG-2779 : modifications to the _laser_lock_box.<commit_after>
import math
import logging

from pymoku._instrument import *
from pymoku._oscilloscope import _CoreOscilloscope, VoltsData
from . import _instrument
from ._pid import PID
from ._sweep import SweepGenerator
log = logging.getLogger(__name__)

REGBASE_LLB_DEMOD			= 23
REGBASE_LLB_SCAN			= 32
REGBASE_LLB_AUX_SINE		= 41

REG_LLB_RATE_SEL			= 76

REGBASE_LLB_PID1			= 106
REGBASE_LLB_PID2			= 117

_LLB_PHASESCALE				= 2**64 / 360.0
_LLB_FREQSCALE				= 2**64 / 1e9
_LLB_HIGH_RATE				= 0
_LLB_LOW_RATE				= 1

_LLB_COEFFICIENT_WIDTH		= 24

class LaserLockBox(_CoreOscilloscope):
	def __init__(self):
		super(LaserLockBox).__init__()
		self._register_accessors(_llb_reg_hdl)

		self.id = 16
		self.type = "laserlockbox"

		self.fast_fs = 31.25e6
		seld.slow_fs = 31.25e6

		self.fast_pid = PID(self, reg_base = REGBASE_LLB_PID1, self.fast_fs)
		self.slow_pid = PID(self, reg_base = REGBASE_LLB_PID2, self.slow_fs)

		self.demod_sweep = SweepGenerator(self, reg_base = REGBASE_LLB_DEMOD)
		self.scan_sweep = SweepGenerator(self, reg_base = REGBASE_LLB_SCAN)
		self.aux_sine_sweep = SweepGenerator(self, reg_base = REGBASE_LLB_AUX_SINE)

		self.demod_sweep.step = 0
		self.demod_sweep.stop = 2**64 -1
		self.demod_sweep.duration = 0
		self.demod_sweep.waveform = 2
		self.demod_sweep.start = 0
		self.demod_sweep.wait_for_trig = False
		self.demod_sweep.hold_last = False

		self.scan_sweep.step = 0
		self.scan_sweep.stop = 2**64 -1
		self.scan_sweep.duration = 0
		self.scan_sweep.waveform = 2
		self.scan_sweep.start = 0
		self.scan_sweep.wait_for_trig = False
		self.scan_sweep.hold_last = False

		self.aux_sine_sweep.step = 0
		self.aux_sine_sweep.stop = 2**64 -1
		self.aux_sine_sweep.duration = 0
		self.aux_sine_sweep.waveform = 2
		self.aux_sine_sweep.start = 0
		self.aux_sine_sweep.wait_for_trig = False
		self.aux_sine_sweep.hold_last = False

	@needs_commit
	def set_pid_by_gain(self, pid_block, g=1, kp=1, ki=0, kd=0, si=None, sd=None):
		"""

		Configure the selected PID controller using gain coefficients.

		:type ch: int; [1,2]
		:param ch: Channel of the PID controller to be configured.

		:type g: float; [0,2^16 - 1]
		:param g: Gain

		:type kp: float; [-1e3,1e3]
		:param kp: Proportional gain factor

		:type ki: float;
		:param ki: Integrator gain factor

		:type kd: float;
		:param kd: Differentiator gain factor

		:type si: float; float; [-1e3,1e3]
		:param si: Integrator gain saturation

		:type sd: float; [-1e3,1e3]
		:param sd: Differentiator gain saturation

		:raises InvalidConfigurationException: if the configuration of PID gains is not possible.
		"""
		pid_array = [self.fast_pid, self.slow_pid]
		pid_array[pid_block -1].set_reg_by_gain(g, kp, ki, kd, si, sd)

	@needs_commit
	def set_pid_by_freq(self, pid_block, kp=1, i_xover=None, d_xover=None, si=None, sd=None):
		"""

		Configure the selected PID controller using crossover frequencies.

		:type ch: int; [1,2]
		:param ch: PID controller to  configure

		:type kp: float; [-1e3,1e3]
		:param kp: Proportional gain factor

		:type i_xover: float; [1e-3,1e6] Hz
		:param i_xover: Integrator crossover frequency

		:type d_xover: float; [1,10e6] Hz
		:param d_xover: Differentiator crossover frequency

		:type si: float; float; [-1e3,1e3]
		:param si: Integrator gain saturation

		:type sd: float; [-1e3,1e3]
		:param sd: Differentiator gain saturation

		:raises InvalidConfigurationException: if the configuration of PID gains is not possible.
		"""
		pid_array = [self.fast_pid, self.slow_pid]
		pid_array[pid_block -1].set_reg_by_frequency(kp, i_xover, d_xover, si, sd)

	@needs_commit
	def set_demodulation(self, frequency, phase):
		"""
		Configure the demodulation stage.

		:type frequency : float; [0, 200e6] Hz
		:param frequency : Internal demodulation frequency

		:type phase : float; [0, 360] degrees
		:param phase : float; Internal demodulation phase

		"""
		self.demod_sweep.step = frequency * _LLB_FREQSCALE
		self.demod_phase.start = phase * _LLB_PHASESCALE

	@needs_commit
	def set_scan(self, frequency, phase):
		"""
		Configure the scan signal.

		:type frequency : float; [0, 200e6] Hz
		:param frequency : Internal demodulation frequency

		:type phase : float; [0, 360] degrees
		:param phase : float; Internal demodulation phase

		"""
		self.scan_sweep.step = frequency * _LLB_FREQSCALE
		seld.scan_sweep.phase = phase * _LLB_PHASESCALE

	@needs_commit
	def set_aux_sine(self, frequency, phase):
		"""
		Configure the aux sine signal.

		:type frequency : float; [0, 200e6] Hz
		:param frequency : Internal demodulation frequency

		:type phase : float; [0, 360] degrees
		:param phase : float; Internal demodulation phase

		"""
		self.aux_sine_sweep.step = frequency * _LLB_FREQSCALE
		seld.aux_sine_sweep.phase = phase * _LLB_PHASESCALE

	@needs_commit
	def set_sample_rate(self, rate):
		"""
		Configure the sample rate of the filters and pid controllers of the laser locker.

		selectable rates:
			-**high** : 62.5 MHz
			-**low**  : 31.25 MHz

		:type rate : string; {'high', 'low'}
		:param rate: sample rate

		"""
		_str_to_rate = {
			'high' 	: _LLB_HIGH_RATE,
			'low'	: _LLB_LOW_RATE
		}
		self.rate_sel = _utils.str_to_val(_str_to_rate, rate, 'sampling rate')

# NOTE: This function avoids @needs_commit because it calls _set_mmap_access which requires an immediate commit
	def set_filter(self, filter_coefficients):
		"""
		Set SOS filter sample rate and filter coefficients. This also enables the input and outputs of the specified Moku:Lab channel.

		:type ch: int; {1,2}
		:param ch: target channel

		:type filter_coefficients: array;
		:param filter_coefficients: array containing SOS filter coefficients. Format is described in class documentation above.
		"""
		_utils.check_parameter_valid('set', ch, [1, 2], 'filter channel')
		_utils.check_parameter_valid('set', sample_rate, ['high', 'low'], 'filter sample rate')

		# Conversion of input array (typically generated by Scipy/Matlab) to HDL memory map format
		if filter_coefficients != None:

			# Deep copy to avoid modifying user's original input array
			intermediate_filter = deepcopy(filter_coefficients)

			# Array dimension check
			if len(filter_coefficients) != 5:
				_utils.check_parameter_valid('set', len(filter_coefficients), [5],'number of coefficient array rows')
			for m in range(4):
				if m == 0:
					if len(filter_coefficients[0]) != 1:
						_utils.check_parameter_valid('set', len(filter_coefficients[0]), [1],'number of columns in coefficient array row 0')
				else:
					if len(filter_coefficients[m]) != 6:
						_utils.check_parameter_valid('set', len(filter_coefficients[m]), [6],("number of columns in coefficient array row %s"%(m)))

			# Array values check
			_utils.check_parameter_valid('range', filter_coefficients[0][0], [-8e6,8e6 - 2**(-24)],("coefficient array entry m = %s, n = %s"%(0,0)))
			for m in range(1, 5):
				for n in range(6):
					_utils.check_parameter_valid('range', filter_coefficients[m][n], [-4.0,4.0 - 2**(-45)],("coefficient array entry m = %s, n = %s"%(0,0)))

			# multiply S coefficients into B coefficients and replace all S coefficients with 1.0
			for n in range(1,5):
				intermediate_filter[n][1] *= intermediate_filter[n][0]
				intermediate_filter[n][2] *= intermediate_filter[n][0]
				intermediate_filter[n][3] *= intermediate_filter[n][0]
				intermediate_filter[n][0] = 1.0

			# place gain factor G into S coefficient position 4 to comply with HDL requirements:
			intermediate_filter[4][0] = intermediate_filter[0][0]
			intermediate_filter = intermediate_filter[1:5]

			if ch == 1:
				self.filter_ch1 = intermediate_filter
			else:
				self.filter_ch2 = intermediate_filter

		# combine both filter arrays:
		filter_coeffs = [[0.0]*6]*4
		coeff_list = [ [ [0 for k in range(2)] for x in range(6)] for y in range(8) ]
		for n in range(4):
		 	filter_coeffs[n] = self.filter_ch1[n] + self.filter_ch2[n]

		for stage in range(4):
			for coeff in range(6):
				if coeff == 0:
					coeff_list[stage][coeff] = int(round( 2**(_LLB_COEFFICIENT_WIDTH - 24) * filter_coeffs[x][y + k*6]))
				else:
					coeff_list[stage][coeff] = int(round( 2**(_LLB_COEFFICIENT_WIDTH - 3) * filter_coeffs[x][y + k*6]))

		with open('.data.dat', 'wb') as f:
			for coeff in range(6):
				for stage in range(4):
					f.write(struct.pack('<q', coeff_list[stage][coeff]))

		self._set_mmap_access(True)
		self._moku._send_file('j', '.data.dat')
		self._set_mmap_access(False)
		os.remove('.data.dat')

		# Enable the output and input of the set channel
		if ch==1:
			self.output_en1 = True
			self.input_en1 = True
		else:
			self.output_en2 = True
			self.input_en2 = True

		# Manually commit the above register settings as @needs_commit is not used in this function
		self.commit()

	def _write_coeffs(self, ch, coeffs):
		assert ch in [1,2], "Invalid channel"
		assert len(coeffs) <= _FIR_NUM_BLOCKS * _FIR_BLOCK_SIZE, "Invalid number of filter coefficients."

		coeffs = list(coeffs)

		# Create a list of coefficients in each FIR block
		n = int(math.ceil(len(coeffs)/float(_FIR_NUM_BLOCKS)))
		blocks = [coeffs[x:x+n] for x in range(0, len(coeffs), n)]
		blocks += [[]] * (_FIR_NUM_BLOCKS - len(blocks))

		# Construct a bytearray from the FIR block contents
		coeff_bytes = bytearray()
		for b in blocks:
			b.reverse()
			coeff_bytes += bytearray(struct.pack('<I', len(b)))
			coeff_bytes += bytearray(struct.pack('<' + 'i'*len(b), *[int(round((2.0**24-1) * c)) for c in b]))
			coeff_bytes += bytearray(b'\x00'*4*(_FIR_BLOCK_SIZE-len(b)))

		# Sanity check the coefficient byte array length
		assert len(coeff_bytes) == (_FIR_BLOCK_SIZE+1) * _FIR_NUM_BLOCKS * 4, "Invalid length for FIR coefficient memory map."

		# Write the coefficients to the FIR coefficient memory map
		self._set_mmap_access(True)
		self._moku._send_file_bytes('j', '', coeff_bytes, offset=_FIR_MMAP_BLOCK_SIZE*(ch-1))
		self._set_mmap_access(False)

		# Release the memory map "file" to other resources
		self._moku._fs_finalise('j', '', _FIR_MMAP_BLOCK_SIZE*2)

_llb_reg_hdl = {
	'rate_sel':		(REG_LLB_RATE_SEL,	to_reg_unsigned(0, 1),
										from_reg_unsigned(0, 1)),
}<|endoftext|>


<commit_before># The MIT License (MIT)
#
# Copyright (c) 2014 Johannes Schlatow
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Helper class for drawing stacked and grouped bar charts with matplotlib
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

from math import *
from collections import OrderedDict

import matplotlib.pyplot as pyplot

import numpy as np
import brewer2mpl

class BarChart(object):
    def __init__(self, title="Title", ylabel="Unknown", xlabel=None, xticks=None, noxticks=False, width=0.15, colorshift=0, rotation=70,
            xticksize=12, legend_loc="upper right"):
        """
        Arguments
        ----------
        title -- The title of plot.
        ylabel -- The label of the y-axis.
        xlabel -- The label ot the x-axis
        xticks -- Override ticks on the x-axis
        noxticks -- Omit ticks on the x-axis
        width -- The (relative) width of the bars.
        colorshift -- Shift the color set. (default: 0)
        rotation -- Rotate the x-tick labels. (default: 70)
        xticksize -- Font size of the x-ticks
        legend_loc -- Location of the legend
        """

        self.groups = OrderedDict()
        self.categories = OrderedDict()
        self.bars = list()
        self.ylabel = ylabel
        self.xlabel = xlabel
        self.title  = title
        self.width  = width
        self.colorshift = colorshift
        self.rotation = rotation
        self.errvalues = dict()
        self.xticksize = xticksize
        self.xticks = xticks
        self.noxticks = noxticks
        self.legend_loc = legend_loc
        if rotation == "vertical" or rotation == "horizontal":
            self.xtickalign = "center"
        elif rotation > 0 and rotation < 90:
            self.xtickalign = "right"
        elif rotation < 0 and rotation > -90:
            self.xtickalign = "left"
        else:
            self.xtickalign = "center"

        # Get "11-class Paired" from ColorBrewer, a nice print-friendly color set.
        # For more on ColorBrewer, see http://colorbrewer2.org/
        self.colors = brewer2mpl.get_map('Paired', 'qualitative', 11).mpl_colors

    def plot(self, axis, legend=True, sort=False, stacked=False):
        """ Create bars on the given axis.
        Arguments
        ----------
        axis -- A pyplot axis object
        legend -- Draw a legend (default: True)
        sort -- Sort categories (default: False)
        stacked -- Plot stacked bars instead of groups (default: False)
        """

        self.cat_list = list(self.categories.keys())
        if sort:
            self.cat_list = sorted(self.cat_list)

        num_colors = len(self.colors)

        # assign colors
        j = self.colorshift
        for name in self.cat_list:
            self.categories[name]["color"] = self.colors[j % num_colors]
            j = j + 1

        N = len(self.groups)
        ind = np.arange(N) * 1.3
        bottom = np.zeros(N)
        j = 0
        if stacked:
            offset = self.width/2
            tick_offset = self.width
        else:
            M = len(self.cat_list)
            ind = ind * M
            offset = self.width*(M/2)
            tick_offset = self.width*M

        for name in self.cat_list:
            cat_data = list()
            for group in self.groups:
                if name in self.groups[group]:
                    cat_data.append(self.groups[group][name])
                else:
                    cat_data.append(0)

            cat_err = None
            if name in self.errvalues:
                cat_err = list()
                for group in self.groups:
                    cat_err.append(self.errvalues[name][group])

            self.bars.append(axis.bar(ind+offset, cat_data, self.width, bottom=bottom, color=self.categories[name]["color"], yerr=cat_err))

            if stacked:
                j = j + 1
                for i in range(0, N):
                    bottom[i] = bottom[i] + cat_data[i]
            else:
                offset = offset + self.width

        axis.set_ylabel(self.ylabel)
        if self.xlabel is not None:
            axis.set_xlabel(self.xlabel)
        axis.set_title(self.title)

        if self.xticks is not None:
            labels = list()
            minorticks = list()
            majorticks = list([0+offset])
            last = 0
            for tick in self.xticks:
                cur  = ind[tick['index']]+offset
                majorticks.append(cur)
                minorticks.append((last+cur)/2)
                last = cur
                labels.append(tick['label'])

            axis.xaxis.set_tick_params(direction="out", top=False, labelbottom=False, size=20, width=2)
            axis.set_xticks(majorticks)
            axis.set_xticks(minorticks, minor=True)
            axis.set_xticklabels(labels, rotation=self.rotation, fontsize=self.xticksize,
                    horizontalalignment=self.xtickalign, minor=True)
        elif self.noxticks:
            axis.xaxis.set_tick_params(top=False, bottom=False, labelbottom=False)
        else:
            axis.set_xticks(ind+tick_offset)
            axis.set_xticklabels(list(self.groups.keys()), rotation=self.rotation,
                    fontsize=self.xticksize, ha=self.xtickalign)

        if legend:
            self._add_legend()

    def add_category(self, name, label=None):
        """ Manually pick a category from the added data (use add_group_data() first).
        Arguments
        ----------
        name -- The name of the category (as in the added data).
        label -- The label of the category as it will be shown in the plot. (default: name)
        """
        self.categories[name] = dict()
        if label is None:
            self.categories[name]['label'] = name
        else:
            self.categories[name]['label'] = label

    def auto_add_categories(self):
        """ Automatically pick all categories from the added data (use add_group_data() first).
        """
        for group in self.groups:
            for name in self.groups[group]:
                self.add_category(name)

    def add_group_data(self, label, tuples):
        """ Add a group and its data to the barchart.
        Arguments
        ----------
        label -- The label of the group as it will be shown in the plot.
        tuples -- A list of tuples of category names and the corresponding data
        """
        self.groups[label] = dict(tuples)

    def add_group_error(self, group, category, errval):
        """ Add an optional error (e.g. variance) to the datum of a specific category in a group.
        By setting an error value, error bars will automatically be added to the barchart.
        Arguments
        ----------
        group -- The group to which an error value shall be added.
        category -- The category of the group to which an error value shall be added.
        errval -- The error value (e.g. variance).
        """
        if category not in self.errvalues:
            self.errvalues[category] = dict()
        self.errvalues[category][group] = errval

    def _add_legend(self):
        revbars = list(self.bars)
        revbars.reverse()
        revnames = list()
        for cat in self.cat_list:
            revnames.append(self.categories[cat]['label'])
        revnames.reverse()
        pyplot.legend( (revbars), (revnames), loc=self.legend_loc )

    def _normalize(self, groupname):
        for group in self.groups:
            basevalue = self.groups[group][groupname]
            for name in self.groups[group]:
                self.groups[group][name] /= basevalue

# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
<commit_msg>add legendsize parameter<commit_after># The MIT License (MIT)
#
# Copyright (c) 2014 Johannes Schlatow
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Helper class for drawing stacked and grouped bar charts with matplotlib
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

from math import *
from collections import OrderedDict

import matplotlib.pyplot as pyplot

import numpy as np
import brewer2mpl

class BarChart(object):
    def __init__(self, title="Title", ylabel="Unknown", xlabel=None, xticks=None, noxticks=False, width=0.15, colorshift=0, rotation=70,
            xticksize=12, legend_loc="upper right", legendsize=None):
        """
        Arguments
        ----------
        title -- The title of plot.
        ylabel -- The label of the y-axis.
        xlabel -- The label ot the x-axis
        xticks -- Override ticks on the x-axis
        noxticks -- Omit ticks on the x-axis
        width -- The (relative) width of the bars.
        colorshift -- Shift the color set. (default: 0)
        rotation -- Rotate the x-tick labels. (default: 70)
        xticksize -- Font size of the x-ticks
        legend_loc -- Location of the legend
        legendsize -- Font size of the legend
        """

        self.groups = OrderedDict()
        self.categories = OrderedDict()
        self.bars = list()
        self.ylabel = ylabel
        self.xlabel = xlabel
        self.title  = title
        self.width  = width
        self.colorshift = colorshift
        self.rotation = rotation
        self.errvalues = dict()
        self.xticksize = xticksize
        self.xticks = xticks
        self.noxticks = noxticks
        self.legend_loc = legend_loc
        if rotation == "vertical" or rotation == "horizontal":
            self.xtickalign = "center"
        elif rotation > 0 and rotation < 90:
            self.xtickalign = "right"
        elif rotation < 0 and rotation > -90:
            self.xtickalign = "left"
        else:
            self.xtickalign = "center"

        if legendsize is not None:
            self.legendsize = legendsize
        else:
            self.legendsize = xticksize

        # Get "11-class Paired" from ColorBrewer, a nice print-friendly color set.
        # For more on ColorBrewer, see http://colorbrewer2.org/
        self.colors = brewer2mpl.get_map('Paired', 'qualitative', 11).mpl_colors

    def plot(self, axis, legend=True, sort=False, stacked=False):
        """ Create bars on the given axis.
        Arguments
        ----------
        axis -- A pyplot axis object
        legend -- Draw a legend (default: True)
        sort -- Sort categories (default: False)
        stacked -- Plot stacked bars instead of groups (default: False)
        """

        self.cat_list = list(self.categories.keys())
        if sort:
            self.cat_list = sorted(self.cat_list)

        num_colors = len(self.colors)

        # assign colors
        j = self.colorshift
        for name in self.cat_list:
            self.categories[name]["color"] = self.colors[j % num_colors]
            j = j + 1

        N = len(self.groups)
        ind = np.arange(N) * 1.3
        bottom = np.zeros(N)
        j = 0
        if stacked:
            offset = self.width/2
            tick_offset = self.width
        else:
            M = len(self.cat_list)
            ind = ind * M
            offset = self.width*(M/2)
            tick_offset = self.width*M

        for name in self.cat_list:
            cat_data = list()
            for group in self.groups:
                if name in self.groups[group]:
                    cat_data.append(self.groups[group][name])
                else:
                    cat_data.append(0)

            cat_err = None
            if name in self.errvalues:
                cat_err = list()
                for group in self.groups:
                    cat_err.append(self.errvalues[name][group])

            self.bars.append(axis.bar(ind+offset, cat_data, self.width, bottom=bottom, color=self.categories[name]["color"], yerr=cat_err))

            if stacked:
                j = j + 1
                for i in range(0, N):
                    bottom[i] = bottom[i] + cat_data[i]
            else:
                offset = offset + self.width

        axis.set_ylabel(self.ylabel)
        if self.xlabel is not None:
            axis.set_xlabel(self.xlabel)
        axis.set_title(self.title)

        if self.xticks is not None:
            labels = list()
            minorticks = list()
            majorticks = list([0+offset])
            last = 0
            for tick in self.xticks:
                cur  = ind[tick['index']]+offset
                majorticks.append(cur)
                minorticks.append((last+cur)/2)
                last = cur
                labels.append(tick['label'])

            axis.xaxis.set_tick_params(direction="out", top=False, labelbottom=False, size=20, width=2)
            axis.set_xticks(majorticks)
            axis.set_xticks(minorticks, minor=True)
            axis.set_xticklabels(labels, rotation=self.rotation, fontsize=self.xticksize,
                    horizontalalignment=self.xtickalign, minor=True)
        elif self.noxticks:
            axis.xaxis.set_tick_params(top=False, bottom=False, labelbottom=False)
        else:
            axis.set_xticks(ind+tick_offset)
            axis.set_xticklabels(list(self.groups.keys()), rotation=self.rotation,
                    fontsize=self.xticksize, ha=self.xtickalign)

        if legend:
            self._add_legend()

    def add_category(self, name, label=None):
        """ Manually pick a category from the added data (use add_group_data() first).
        Arguments
        ----------
        name -- The name of the category (as in the added data).
        label -- The label of the category as it will be shown in the plot. (default: name)
        """
        self.categories[name] = dict()
        if label is None:
            self.categories[name]['label'] = name
        else:
            self.categories[name]['label'] = label

    def auto_add_categories(self):
        """ Automatically pick all categories from the added data (use add_group_data() first).
        """
        for group in self.groups:
            for name in self.groups[group]:
                self.add_category(name)

    def add_group_data(self, label, tuples):
        """ Add a group and its data to the barchart.
        Arguments
        ----------
        label -- The label of the group as it will be shown in the plot.
        tuples -- A list of tuples of category names and the corresponding data
        """
        self.groups[label] = dict(tuples)

    def add_group_error(self, group, category, errval):
        """ Add an optional error (e.g. variance) to the datum of a specific category in a group.
        By setting an error value, error bars will automatically be added to the barchart.
        Arguments
        ----------
        group -- The group to which an error value shall be added.
        category -- The category of the group to which an error value shall be added.
        errval -- The error value (e.g. variance).
        """
        if category not in self.errvalues:
            self.errvalues[category] = dict()
        self.errvalues[category][group] = errval

    def _add_legend(self):
        revbars = list(self.bars)
        revbars.reverse()
        revnames = list()
        for cat in self.cat_list:
            revnames.append(self.categories[cat]['label'])
        revnames.reverse()
        pyplot.legend( (revbars), (revnames), loc=self.legend_loc, fontsize=self.legendsize,
                handlelength=1, labelspacing=0.2)

    def _normalize(self, groupname):
        for group in self.groups:
            basevalue = self.groups[group][groupname]
            for name in self.groups[group]:
                self.groups[group][name] /= basevalue

# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
<|endoftext|>


<commit_before># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import

from .error import ValidationError
from .error import InvalidDataError
from .error import OpenError
from .error import InvalidTableNameError
from .error import InvalidFilePathError

from .data import TableData

from .csv.core import CsvTableFileLoader
from .csv.core import CsvTableTextLoader
from .html.core import HtmlTableFileLoader
from .html.core import HtmlTableTextLoader
from .mediawiki.core import MediaWikiTableFileLoader
from .mediawiki.core import MediaWikiTableTextLoader
from .spreadsheet.excelloader import ExcelTableFileLoader
from .json.core import JsonTableFileLoader
from .json.core import JsonTableTextLoader
<commit_msg>Add imports<commit_after># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import

from .error import ValidationError
from .error import InvalidTableNameError
from .error import InvalidHeaderNameError
from .error import InvalidFilePathError
from .error import InvalidDataError
from .error import EmptyDataError
from .error import OpenError

from .data import TableData

from .csv.core import CsvTableFileLoader
from .csv.core import CsvTableTextLoader
from .html.core import HtmlTableFileLoader
from .html.core import HtmlTableTextLoader
from .mediawiki.core import MediaWikiTableFileLoader
from .mediawiki.core import MediaWikiTableTextLoader
from .spreadsheet.excelloader import ExcelTableFileLoader
from .json.core import JsonTableFileLoader
from .json.core import JsonTableTextLoader
<|endoftext|>