<commit_before>#!/usr/bin/env python
"""
Data Augmentation wrappers and functions.
Functions to be used as "processing" arguments
are here.
"""

from __future__ import division

import neukrill_net.image_processing as image_processing
import numpy as np
import skimage.util
import itertools

def augmentation_wrapper(units='float64', **augment_settings):
    """
    Takes settings for augmentation as **kwargs
    and produces a "processing" function to make more
    training data.
    The returned function will return a list of images
    containing EVERY possible combinations of options
    as given in the augmentation settings.
    """
    # components of the processing pipeline
    components = []

    # Datatype unit conversion
    # We very probably want this to be float throughout
    if units == 'float64' or units == 'float':
        unitconvert = lambda images: [skimage.util.img_as_float(image)
                                        for image in images]

    elif units == 'uint8':
        unitconvert = lambda images: [skimage.util.img_as_ubyte(image)
                                        for image in images]

    elif units == None or units == 'auto':
        unitconvert = lambda images: images

    else:
        raise ValueError('Unrecognised output units: {}'.format(
                            augment_settings['units']))

    # Shape-fixing without resizing
    if 'shape' in augment_settings:
        #
        shapefix = lambda images: [image_processing.shape_fix(
                                    image,
                                    augment_settings['shape'])
                                    for image in images]
    else:
        shapefix = lambda images: images

    # Resize
    if 'resize' in augment_settings:
        # apply our resize
        resize = lambda images: [image_processing.resize_image(
                                    image,
                                    augment_settings['resize'])
                                    for image in images]
    else:
        resize = lambda images: images

    # Rescale
    if 'scale' in augment_settings and not augment_settings['scale']==None:
        scale = lambda images: [image_processing.scale_image(image, sf)
                                for sf in augment_settings['scale']
                                for image in images]
    else:
        scale = lambda images: images

    # Landscapise
    # Set to True if you want to ensure all the images are landscape
    if 'landscapise' in augment_settings and augment_settings['landscapise']:
        landscapise = lambda images: [image_processing.landscapise_image(image)
                                        for image in images]
    else:
        landscapise = lambda images: images

    # Rotate
    if 'rotate' in augment_settings and not augment_settings['rotate']==None:
        # apply rotate with options
        # will require a function here to call it a number of times
        rotate = lambda images: [rotatedImage for image in images
                                    for rotatedImage in
                                    rotations(image,
                                        augment_settings['rotate'],
                                        augment_settings['rotate_is_resizable']
                                    )
                                ]
    else:
        rotate = lambda images: images

    # Shear
    if 'shear' in augment_settings and not augment_settings['shear']==None:
        shear = lambda images: [image_processing.shear_image(image, this_shear)
                                for this_shear in augment_settings['shear']
                                for image in images]
    else:
        shear = lambda images: images

    # Flip (always horizontally)
    # All other relfections can be acheived by coupling with an appropriate reflection
    # Flip setting should either be True or False in settings
    if 'flip' in augment_settings and augment_settings['flip']:
        flip = lambda images: images + [image_processing.flip_image(image, True)
                                            for image in images]
    else:
        flip = lambda images: images

    # Crop (every side or not at all)
    if 'crop' in augment_settings and augment_settings['crop']:
        crop = lambda images: images + [croppedImage for image in images
                                            for croppedImage in
                                            allcrops(image)]
    else:
        crop = lambda images: images

    # Pad to translate within shape_fix window
    if 'translations' in augment_settings and not augment_settings['translations']==None:
        pad = lambda images: [paddedImage for image in images
                                            for paddedImage in
                                            allpads(image, augment_settings['translations'])]
    else:
        pad = lambda images: images

    # Add pixel noise
    if 'noise' in augment_settings and not augment_settings['noise']==None:
        noisify = lambda images: [image_processing.noisify_image(image, augment_settings['noise'])
                                    for image in images]
    else:
        noisify = lambda images: images

    # Stack all our functions together
    # Order matters here:
    # - Landscapise first because we don't want to undo rotations
    # - Rotate before anything else because then it has higher accuracy
    #   (might want to move in pixels which would otherwise be cropped
    #    don't want the loss of resolution caused by resize)
    # - Flip
    # - Crop
    # - Pad
    # - Align shape window
    # - Resize last because it is lossy
    # - Add noise independent for all output pixels
    processing=lambda image:noisify(
                                resize(
                                    shapefix(
                                        pad(
                                            crop(
                                                shear(
                                                    flip(
                                                        rotate(
                                                            landscapise(
                                                                unitconvert(
                                                                    [ image ]
                                                                )
                                                            )
                                                        )
                                                    )
                                                )
                                            )
                                        )
                                    )
                                )
                            )

    return processing

def rotations(image, num_rotations, resizable):
    """
    Returns a number of rotations depending on the settings.
    """
    return [image_processing.rotate_image(image, angle, resizable)
        for angle in np.linspace(0, 360, num_rotations, endpoint=False)]

def allcrops(image):
    """
    Returns a number of cropped copies of an image
    """
    return [image_processing.crop_image(image,side_id)
        for side_id in range(4)]

def allpads(image, pad_amounts):
    """
    Returns a list of padded images, with centre shift amounts
    specified by pad_amounts
    """
    # Can go in all four directions, so include +ve and -ve
    pad_amounts = np.array(pad_amounts)
    pad_amounts = np.union1d(-pad_amounts, pad_amounts)

    # Make permutations of all x and y shifts
    return [image_processing.padshift_image(image, shift)
        for shift in itertools.product(pad_amounts, repeat=2)]

class RandomAugment(object):
    """
    A class for performing random augmentations.
    When you call it with an image it gives
    you back a randomly manipulated version of it.
    """

    def __init__(self, random_seed=42, units='float64', rotate_is_resizable=False, **kwargs):
        """
        Inputs: random_seed
                keyword assignment of augmentation settings.
        """
        # Initialise the random number generator
        self.rng = np.random.RandomState(seed=random_seed)
        # Store settings as a dictionary
        self.settings = kwargs
        # Add kwargs with defaults into dictionary
        self.settings['units'] = units
        # Don't even need this anymore
        #self.settings['rotate_is_resizable'] = rotate_is_resizable

    def __call__(self, image):
        """
        Maps raw image to augmented image.
        """

        # Note down the original type
        original_dtype = image.dtype

        # Convert to float while we process it
        image = skimage.util.img_as_float(image)

        #####################################################
        # Pre-augmentation processing

        # Landscapise
        # Set to True if you want to ensure all the images are landscape
        if 'landscapise' in self.settings and self.settings['landscapise']:
            image = image_processing.landscapise_image(image)

        # Ensure image is square now if we are going to reshape
        if 'resize' in self.settings or 'shape' in self.settings:
            image = image_processing.pad_to_square(image)

        #####################################################
        # Random augmentation

        # Flip (always horizontally)
        # All other relfections can be acheived by coupling with an appropriate reflection
        # Flip setting should either be True or False in settings
        if 'flip' in self.settings and self.settings['flip']:
            # Flip if coin-toss says so
            if self.rng.binomial(1, 0.5):
                image = image_processing.flip_image(image, True)

        ### Affine Transformation
        # Rescale
        if ('scale' not in self.settings or
                self.settings['scale']==None or
                self.settings['scale']==0.0 ):
            scalefactor = 1.0
        else:
            if isinstance(self.settings['scale'],(list,tuple)):
                sf_index = self.rng.randint(0, len(self.settings['scale']))
                scalefactor = self.settings['scale'][sf_index]
            else:
                scalefactor = self.rng.normal(loc=1.0, scale=self.settings['scale'])
                # Ignore extreme events
                if (scalefactor-1.0) > 3*self.settings['scale']:
                    scalefactor = 1.0
                elif (scalefactor-1.0) < -3*self.settings['scale']:
                    scalefactor = 1.0

        # Asymmetric x and y scaling
        if ('scale_asym' not in self.settings
                or self.settings['scale_asym']==None
                or self.settings['scale_asym']==0.0 ):
            # x and y scales are the same
            scalefactor = (scalefactor,scalefactor)
        else:
            # Pick a new x and y scale near the overall scale
            x_sf = self.rng.normal(scale=self.settings['scale_asym'])
            y_sf = self.rng.normal(scale=self.settings['scale_asym'])
            # Ignore extreme events
            if x_sf > 3*self.settings['scale_asym'] or x_sf < -3*self.settings['scale_asym']:
                x_sf = 0.0
            if y_sf > 3*self.settings['scale_asym'] or y_sf < -3*self.settings['scale_asym']:
                y_sf = 0.0
            # Combine x and y scales as a tuple
            scalefactor = (scalefactor+x_sf, scalefactor+y_sf)

        # Define translation (shunt)
        if ('shunt' not in self.settings
                or self.settings['shunt']==None
                or self.settings['shunt']==0.0 ):
            shuntlist = (0,0)
        else:
            shuntlist = []
            for i in range(2):
                if isinstance(self.settings['shunt'],(list,tuple)):
                    # Pick from list
                    shunt_index = self.rng.randint(0, len(self.settings['shunt']))
                    shunt = self.settings['shunt'][shunt_index]
                    # Ensure symmetry
                    if self.rng.binomial(1, 0.5):
                        shunt = -shunt
                else:
                    # Chose from gaussian
                    shunt = self.rng.normal(loc=0.0, scale=self.settings['shunt'])
                    # Ignore extreme events
                    if shunt<-3*self.settings['shunt'] or shunt>3*self.settings['shunt']:
                        shunt = 0
                # Add to the list
                shuntlist += [shunt]
            # Convert list to tuple
            shuntlist = tuple(shuntlist)

        # Rotate
        if ('rotate' not in self.settings
                or self.settings['rotate']==None
                or self.settings['rotate']==0.0 ):
            rot_angle = 0.0
        else:
            if isinstance(self.settings['rotate'],(list,tuple)):
                # Pick from list of potential rotations
                rot_index = self.rng.randint(0, len(self.settings['rotate']))
                rot_angle = self.settings['rotate'][rot_index]

            elif self.settings['rotate']==-1:
                # Uniform distribution of rotations
                rot_angle = self.rng.uniform(low=0.0, high=360.0)

            else:
                # Select from implied list
                rot_index = self.rng.randint(0, self.settings['rotate'])
                rot_angle = 360.0 * rot_index / self.settings['rotate']

        # Shear
        if ('shear' not in self.settings
                or self.settings['shear']==None
                or self.settings['shear']==0.0 ):
            shear_angle = 0.0
        else:
            if isinstance(self.settings['shear'],(list,tuple)):
                shear_index = self.rng.randint(0, len(self.settings['shear']))
                shear_angle = self.settings['shear'][shear_index]
            else:
                # Chose from gaussian
                shear_angle = self.rng.normal(loc=0.0, scale=self.settings['shear'])

        # Order of affine transformation
        if 'transform_order' not in self.settings or self.settings['transform_order']==None:
            transform_order = 0.5
        else:
            transform_order = self.settings['transform_order']

        # Perform affine transformation
        image = image_processing.custom_transform_nice_units(image, scale=scalefactor,
                    rotation=rot_angle, shear=shear_angle, translation=shuntlist, order=transform_order)

        # Crop by a random amount from each side
        if 'crop' in self.settings and not self.settings['crop']==None:
            for side_id in range(4):
                crop_index = self.rng.randint(0, len(self.settings['crop']))
                image = image_processing.crop_image(image, side_id,
                        self.settings['crop'][crop_index])

        #####################################################
        # Post-augmentation processing

        # Shape-fixing without resizing
        if 'shape' in self.settings:
            image = image_processing.shape_fix(image, self.settings['shape'])

        # Resize
        if not 'resize_order' in self.settings or self.settings['resize_order']==None:
            resize_order = 0.75
        if 'resize' in self.settings:
            image = image_processing.resize_image(image, self.settings['resize'], resize_order)

        #####################################################

        # Add pixel noise
        if ('noise' in self.settings
                and not self.settings['noise']==None
                and not self.settings['noise']==0.0 ):
            noise_seed = self.rng.randint(2**32)
            image = image_processing.noisify_image(image, self.settings['noise'], noise_seed)

        #####################################################

        if 'normalise' in self.settings:
            if self.settings['normalise']['global_or_pixel'] == 'global':
                mu = self.settings['normalise']['mu']
                sigma = self.settings['normalise']['sigma']
                image = (image - mu)/sigma
            elif self.settings['normalise']['global_or_pixel'] == 'pixel':
                for i in range(image.shape[0]):
                    for j in range(image.shape[1]):
                        mu = self.settings['normalise']['mu'][str((i,j))]
                        sigma = self.settings['normalise']['sigma'][str((i,j))]
                        image[i,j] = (image[i,j] - mu)/sigma
            else:
                raise ValueError("Invalid option for global_or_pixel, should be "
                                     "one of global or pixel.")

        #####################################################

        # Datatype unit conversion
        if ('units' not in self.settings or
                self.settings['units'] == None or self.settings['units'] == 'auto'):
            # Preserve the datatype
            # Ensure output matches input
            image = image_processing.img_as_dtype(image, original_dtype)

        elif self.settings['units'] == 'float64' or self.settings['units'] == 'float':
            image = skimage.util.img_as_float(image)

        elif self.settings['units'] == 'uint8':
            image = skimage.util.img_as_ubyte(image)

        else:
            raise ValueError('Unrecognised output units: {}'.format(
                                self.settings['units']))

        return image

<commit_msg>Offline augmentation now requires resize_order as a parameter<commit_after>#!/usr/bin/env python
"""
Data Augmentation wrappers and functions.
Functions to be used as "processing" arguments
are here.
"""

from __future__ import division

import neukrill_net.image_processing as image_processing
import numpy as np
import skimage.util
import itertools

def augmentation_wrapper(units='float64', **augment_settings):
    """
    Takes settings for augmentation as **kwargs
    and produces a "processing" function to make more
    training data.
    The returned function will return a list of images
    containing EVERY possible combinations of options
    as given in the augmentation settings.
    """
    # components of the processing pipeline
    components = []

    # Datatype unit conversion
    # We very probably want this to be float throughout
    if units == 'float64' or units == 'float':
        unitconvert = lambda images: [skimage.util.img_as_float(image)
                                        for image in images]

    elif units == 'uint8':
        unitconvert = lambda images: [skimage.util.img_as_ubyte(image)
                                        for image in images]

    elif units == None or units == 'auto':
        unitconvert = lambda images: images

    else:
        raise ValueError('Unrecognised output units: {}'.format(
                            augment_settings['units']))

    # Shape-fixing without resizing
    if 'shape' in augment_settings:
        #
        shapefix = lambda images: [image_processing.shape_fix(
                                    image,
                                    augment_settings['shape'])
                                    for image in images]
    else:
        shapefix = lambda images: images

    # Resize
    if 'resize' in augment_settings:

        # apply our resize
        resize = lambda images: [image_processing.resize_image(
                                    image,
                                    augment_settings['resize'],
                                    augment_settings['resize_order'])
                                    for image in images]
    else:
        resize = lambda images: images

    # Rescale
    if 'scale' in augment_settings and not augment_settings['scale']==None:
        scale = lambda images: [image_processing.scale_image(image, sf)
                                for sf in augment_settings['scale']
                                for image in images]
    else:
        scale = lambda images: images

    # Landscapise
    # Set to True if you want to ensure all the images are landscape
    if 'landscapise' in augment_settings and augment_settings['landscapise']:
        landscapise = lambda images: [image_processing.landscapise_image(image)
                                        for image in images]
    else:
        landscapise = lambda images: images

    # Rotate
    if 'rotate' in augment_settings and not augment_settings['rotate']==None:
        # apply rotate with options
        # will require a function here to call it a number of times
        rotate = lambda images: [rotatedImage for image in images
                                    for rotatedImage in
                                    rotations(image,
                                        augment_settings['rotate'],
                                        augment_settings['rotate_is_resizable']
                                    )
                                ]
    else:
        rotate = lambda images: images

    # Shear
    if 'shear' in augment_settings and not augment_settings['shear']==None:
        shear = lambda images: [image_processing.shear_image(image, this_shear)
                                for this_shear in augment_settings['shear']
                                for image in images]
    else:
        shear = lambda images: images

    # Flip (always horizontally)
    # All other relfections can be acheived by coupling with an appropriate reflection
    # Flip setting should either be True or False in settings
    if 'flip' in augment_settings and augment_settings['flip']:
        flip = lambda images: images + [image_processing.flip_image(image, True)
                                            for image in images]
    else:
        flip = lambda images: images

    # Crop (every side or not at all)
    if 'crop' in augment_settings and augment_settings['crop']:
        crop = lambda images: images + [croppedImage for image in images
                                            for croppedImage in
                                            allcrops(image)]
    else:
        crop = lambda images: images

    # Pad to translate within shape_fix window
    if 'translations' in augment_settings and not augment_settings['translations']==None:
        pad = lambda images: [paddedImage for image in images
                                            for paddedImage in
                                            allpads(image, augment_settings['translations'])]
    else:
        pad = lambda images: images

    # Add pixel noise
    if 'noise' in augment_settings and not augment_settings['noise']==None:
        noisify = lambda images: [image_processing.noisify_image(image, augment_settings['noise'])
                                    for image in images]
    else:
        noisify = lambda images: images

    # Stack all our functions together
    # Order matters here:
    # - Landscapise first because we don't want to undo rotations
    # - Rotate before anything else because then it has higher accuracy
    #   (might want to move in pixels which would otherwise be cropped
    #    don't want the loss of resolution caused by resize)
    # - Flip
    # - Crop
    # - Pad
    # - Align shape window
    # - Resize last because it is lossy
    # - Add noise independent for all output pixels
    processing=lambda image:noisify(
                                resize(
                                    shapefix(
                                        pad(
                                            crop(
                                                shear(
                                                    flip(
                                                        rotate(
                                                            landscapise(
                                                                unitconvert(
                                                                    [ image ]
                                                                )
                                                            )
                                                        )
                                                    )
                                                )
                                            )
                                        )
                                    )
                                )
                            )

    return processing

def rotations(image, num_rotations, resizable):
    """
    Returns a number of rotations depending on the settings.
    """
    return [image_processing.rotate_image(image, angle, resizable)
        for angle in np.linspace(0, 360, num_rotations, endpoint=False)]

def allcrops(image):
    """
    Returns a number of cropped copies of an image
    """
    return [image_processing.crop_image(image,side_id)
        for side_id in range(4)]

def allpads(image, pad_amounts):
    """
    Returns a list of padded images, with centre shift amounts
    specified by pad_amounts
    """
    # Can go in all four directions, so include +ve and -ve
    pad_amounts = np.array(pad_amounts)
    pad_amounts = np.union1d(-pad_amounts, pad_amounts)

    # Make permutations of all x and y shifts
    return [image_processing.padshift_image(image, shift)
        for shift in itertools.product(pad_amounts, repeat=2)]

class RandomAugment(object):
    """
    A class for performing random augmentations.
    When you call it with an image it gives
    you back a randomly manipulated version of it.
    """

    def __init__(self, random_seed=42, units='float64', rotate_is_resizable=False, **kwargs):
        """
        Inputs: random_seed
                keyword assignment of augmentation settings.
        """
        # Initialise the random number generator
        self.rng = np.random.RandomState(seed=random_seed)
        # Store settings as a dictionary
        self.settings = kwargs
        # Add kwargs with defaults into dictionary
        self.settings['units'] = units
        # Don't even need this anymore
        #self.settings['rotate_is_resizable'] = rotate_is_resizable

    def __call__(self, image):
        """
        Maps raw image to augmented image.
        """

        # Note down the original type
        original_dtype = image.dtype

        # Convert to float while we process it
        image = skimage.util.img_as_float(image)

        #####################################################
        # Pre-augmentation processing

        # Landscapise
        # Set to True if you want to ensure all the images are landscape
        if 'landscapise' in self.settings and self.settings['landscapise']:
            image = image_processing.landscapise_image(image)

        # Ensure image is square now if we are going to reshape
        if 'resize' in self.settings or 'shape' in self.settings:
            image = image_processing.pad_to_square(image)

        #####################################################
        # Random augmentation

        # Flip (always horizontally)
        # All other relfections can be acheived by coupling with an appropriate reflection
        # Flip setting should either be True or False in settings
        if 'flip' in self.settings and self.settings['flip']:
            # Flip if coin-toss says so
            if self.rng.binomial(1, 0.5):
                image = image_processing.flip_image(image, True)

        ### Affine Transformation
        # Rescale
        if ('scale' not in self.settings or
                self.settings['scale']==None or
                self.settings['scale']==0.0 ):
            scalefactor = 1.0
        else:
            if isinstance(self.settings['scale'],(list,tuple)):
                sf_index = self.rng.randint(0, len(self.settings['scale']))
                scalefactor = self.settings['scale'][sf_index]
            else:
                scalefactor = self.rng.normal(loc=1.0, scale=self.settings['scale'])
                # Ignore extreme events
                if (scalefactor-1.0) > 3*self.settings['scale']:
                    scalefactor = 1.0
                elif (scalefactor-1.0) < -3*self.settings['scale']:
                    scalefactor = 1.0

        # Asymmetric x and y scaling
        if ('scale_asym' not in self.settings
                or self.settings['scale_asym']==None
                or self.settings['scale_asym']==0.0 ):
            # x and y scales are the same
            scalefactor = (scalefactor,scalefactor)
        else:
            # Pick a new x and y scale near the overall scale
            x_sf = self.rng.normal(scale=self.settings['scale_asym'])
            y_sf = self.rng.normal(scale=self.settings['scale_asym'])
            # Ignore extreme events
            if x_sf > 3*self.settings['scale_asym'] or x_sf < -3*self.settings['scale_asym']:
                x_sf = 0.0
            if y_sf > 3*self.settings['scale_asym'] or y_sf < -3*self.settings['scale_asym']:
                y_sf = 0.0
            # Combine x and y scales as a tuple
            scalefactor = (scalefactor+x_sf, scalefactor+y_sf)

        # Define translation (shunt)
        if ('shunt' not in self.settings
                or self.settings['shunt']==None
                or self.settings['shunt']==0.0 ):
            shuntlist = (0,0)
        else:
            shuntlist = []
            for i in range(2):
                if isinstance(self.settings['shunt'],(list,tuple)):
                    # Pick from list
                    shunt_index = self.rng.randint(0, len(self.settings['shunt']))
                    shunt = self.settings['shunt'][shunt_index]
                    # Ensure symmetry
                    if self.rng.binomial(1, 0.5):
                        shunt = -shunt
                else:
                    # Chose from gaussian
                    shunt = self.rng.normal(loc=0.0, scale=self.settings['shunt'])
                    # Ignore extreme events
                    if shunt<-3*self.settings['shunt'] or shunt>3*self.settings['shunt']:
                        shunt = 0
                # Add to the list
                shuntlist += [shunt]
            # Convert list to tuple
            shuntlist = tuple(shuntlist)

        # Rotate
        if ('rotate' not in self.settings
                or self.settings['rotate']==None
                or self.settings['rotate']==0.0 ):
            rot_angle = 0.0
        else:
            if isinstance(self.settings['rotate'],(list,tuple)):
                # Pick from list of potential rotations
                rot_index = self.rng.randint(0, len(self.settings['rotate']))
                rot_angle = self.settings['rotate'][rot_index]

            elif self.settings['rotate']==-1:
                # Uniform distribution of rotations
                rot_angle = self.rng.uniform(low=0.0, high=360.0)

            else:
                # Select from implied list
                rot_index = self.rng.randint(0, self.settings['rotate'])
                rot_angle = 360.0 * rot_index / self.settings['rotate']

        # Shear
        if ('shear' not in self.settings
                or self.settings['shear']==None
                or self.settings['shear']==0.0 ):
            shear_angle = 0.0
        else:
            if isinstance(self.settings['shear'],(list,tuple)):
                shear_index = self.rng.randint(0, len(self.settings['shear']))
                shear_angle = self.settings['shear'][shear_index]
            else:
                # Chose from gaussian
                shear_angle = self.rng.normal(loc=0.0, scale=self.settings['shear'])

        # Order of affine transformation
        if 'transform_order' not in self.settings or self.settings['transform_order']==None:
            transform_order = 0.5
        else:
            transform_order = self.settings['transform_order']

        # Perform affine transformation
        image = image_processing.custom_transform_nice_units(image, scale=scalefactor,
                    rotation=rot_angle, shear=shear_angle, translation=shuntlist, order=transform_order)

        # Crop by a random amount from each side
        if 'crop' in self.settings and not self.settings['crop']==None:
            for side_id in range(4):
                crop_index = self.rng.randint(0, len(self.settings['crop']))
                image = image_processing.crop_image(image, side_id,
                        self.settings['crop'][crop_index])

        #####################################################
        # Post-augmentation processing

        # Shape-fixing without resizing
        if 'shape' in self.settings:
            image = image_processing.shape_fix(image, self.settings['shape'])

        # Resize
        if not 'resize_order' in self.settings or self.settings['resize_order']==None:
            resize_order = 0.75
        if 'resize' in self.settings:
            image = image_processing.resize_image(image, self.settings['resize'], resize_order)

        #####################################################

        # Add pixel noise
        if ('noise' in self.settings
                and not self.settings['noise']==None
                and not self.settings['noise']==0.0 ):
            noise_seed = self.rng.randint(2**32)
            image = image_processing.noisify_image(image, self.settings['noise'], noise_seed)

        #####################################################

        if 'normalise' in self.settings:
            if self.settings['normalise']['global_or_pixel'] == 'global':
                mu = self.settings['normalise']['mu']
                sigma = self.settings['normalise']['sigma']
                image = (image - mu)/sigma
            elif self.settings['normalise']['global_or_pixel'] == 'pixel':
                for i in range(image.shape[0]):
                    for j in range(image.shape[1]):
                        mu = self.settings['normalise']['mu'][str((i,j))]
                        sigma = self.settings['normalise']['sigma'][str((i,j))]
                        image[i,j] = (image[i,j] - mu)/sigma
            else:
                raise ValueError("Invalid option for global_or_pixel, should be "
                                     "one of global or pixel.")

        #####################################################

        # Datatype unit conversion
        if ('units' not in self.settings or
                self.settings['units'] == None or self.settings['units'] == 'auto'):
            # Preserve the datatype
            # Ensure output matches input
            image = image_processing.img_as_dtype(image, original_dtype)

        elif self.settings['units'] == 'float64' or self.settings['units'] == 'float':
            image = skimage.util.img_as_float(image)

        elif self.settings['units'] == 'uint8':
            image = skimage.util.img_as_ubyte(image)

        else:
            raise ValueError('Unrecognised output units: {}'.format(
                                self.settings['units']))

        return image

<|endoftext|>


<commit_before>"""
Optimizer objects live here.
"""

import os
from inspyred import ec
from inspyred.ec import observers
from inspyred.ec import terminators
from inspyred.ec import selectors
from inspyred.ec import replacers
from inspyred.ec import variators
from random import Random
from time import time
import logging

class __Optimizer(object):
    """
    Base optimization class
    """

    def __init__(
        self,
        max_constraints,
        min_constraints,
        evaluator,
        mutation_rate,
        maximize,
        seeds,
        population_size,
    ):

        self.max_constraints = max_constraints
        self.min_constraints = min_constraints
        self.evaluator = evaluator
        self.population_size = population_size
        self.maximize = maximize
        self.mutation_rate = mutation_rate
        self.seeds = seeds

    def uniform_random_chromosome(self, random, args):
        chromosome = []
        for lo, hi in zip(self.max_constraints, self.min_constraints):
            chromosome.append(random.uniform(lo, hi))
        return chromosome

    def print_report(self, final_pop, do_plot, stat_file_name):
        print(max(final_pop))
        # Sort and print the fittest individual, which will be at index 0.
        final_pop.sort(reverse=True)
        print("\n  Fittest individual:\n  %s" % final_pop[0])

        if do_plot:
            from inspyred.ec import analysis

            try:
                stat_file = open(stat_file_name, "r")
                analysis.generation_plot(stat_file, errorbars=False)
            finally:
                stat_file.close()

class CustomOptimizerA(__Optimizer):
    def __init__(
        self,
        max_constraints,
        min_constraints,
        evaluator,
        mutation_rate=0.2,
        max_evaluations=100,
        population_size=10,
        num_selected=None,
        tourn_size=2,
        num_elites=1,
        maximize=False,
        num_offspring=None,
        seeds=[],
        verbose=False,
    ):

        super(CustomOptimizerA, self).__init__(
            max_constraints,
            min_constraints,
            evaluator,
            mutation_rate,
            maximize,
            seeds,
            population_size,
        )

        self.max_evaluations = max_evaluations
        self.tourn_size = tourn_size
        self.num_elites = num_elites
        self.mutation_rate = mutation_rate
        self.verbose = verbose

        if num_selected == None:
            self.num_selected = population_size
        else:
            self.num_selected = num_selected
        if num_offspring == None:
            self.num_offspring = population_size - self.num_selected
        else:
            self.num_offspring = num_offspring

    def optimize(self, do_plot=True, seed=int(time()), summary_dir=None):

        rand = Random()
        rand.seed(seed)

        if summary_dir is None:
            cwd = os.getcwd()
            summary_dir = os.path.dirname(cwd) + "/data/"

        if not os.path.exists(summary_dir):
            os.mkdir(summary_dir)

        stat_file_name = summary_dir + "/ga_statistics.csv"
        ind_file_name = summary_dir + "/ga_individuals.csv"

        stat_file = open(stat_file_name, "w")
        ind_file = open(ind_file_name, "w")
        print("Created files: %s and %s" % (stat_file_name, ind_file_name))

        if self.verbose:
            logger = logging.getLogger("inspyred.ec")
            logger.setLevel(logging.DEBUG)

            ch = logging.StreamHandler()
            # ch.setLevel(logging.DEBUG)
            formatter = logging.Formatter(">>> EC: - %(levelname)s - %(message)s")
            ch.setFormatter(formatter)
            logger.addHandler(ch)

        algorithm = ec.EvolutionaryComputation(rand)
        algorithm.observer = observers.file_observer
        algorithm.terminator = terminators.evaluation_termination
        algorithm.selector = selectors.tournament_selection
        algorithm.replacer = replacers.steady_state_replacement
        algorithm.variator = [variators.blend_crossover, variators.gaussian_mutation]

        final_pop = algorithm.evolve(
            generator=self.uniform_random_chromosome,
            evaluator=self.evaluator.evaluate,
            pop_size=self.population_size,
            maximize=self.maximize,
            bounder=ec.Bounder(
                lower_bound=self.min_constraints, upper_bound=self.max_constraints
            ),
            num_selected=self.num_selected,
            tourn_size=self.tourn_size,
            num_elites=self.num_elites,
            num_offspring=self.num_offspring,
            max_evaluations=self.max_evaluations,
            mutation_rate=self.mutation_rate,
            statistics_file=stat_file,
            seeds=self.seeds,
            individuals_file=ind_file,
        )

        stat_file.close()
        ind_file.close()

        self.print_report(final_pop, do_plot, stat_file_name)

        # return the parameter set for the best individual

        return final_pop[0].candidate, final_pop[0].fitness
<commit_msg>feat(optimizers): add basic sanity checks<commit_after>"""
Optimizer objects live here.
"""

import os
from inspyred import ec
from inspyred.ec import observers
from inspyred.ec import terminators
from inspyred.ec import selectors
from inspyred.ec import replacers
from inspyred.ec import variators
from random import Random
from time import time
import logging

class __Optimizer(object):
    """
    Base optimization class
    """

    def __init__(
        self,
        max_constraints,
        min_constraints,
        evaluator,
        mutation_rate,
        maximize,
        seeds,
        population_size,
    ):

        self.max_constraints = max_constraints
        self.min_constraints = min_constraints
        self.evaluator = evaluator
        self.population_size = population_size
        self.maximize = maximize
        self.mutation_rate = mutation_rate
        self.seeds = seeds

        # check that constraints are of equal lengths
        assert (len(self.max_constraints) == len(self.min_constraints)), "Min ({}) and max ({}) constraint lists are not of equal size.".format(len(self.min_constraints), len(self.max_constraints))

        # check that min constraints are always less than max constraints
        for i in range(0, len(self.max_constraints)):
            assert (self.min_constraints[i] <= self.max_constraints[i]), "Min constraint {} is larger than max constraint {}.".format(self.min_constraints[i], self.max_constraints[i])

    def uniform_random_chromosome(self, random, args):
        chromosome = []
        for lo, hi in zip(self.max_constraints, self.min_constraints):
            chromosome.append(random.uniform(lo, hi))
        return chromosome

    def print_report(self, final_pop, do_plot, stat_file_name):
        print(max(final_pop))
        # Sort and print the fittest individual, which will be at index 0.
        final_pop.sort(reverse=True)
        print("\n  Fittest individual:\n  %s" % final_pop[0])

        if do_plot:
            from inspyred.ec import analysis

            try:
                stat_file = open(stat_file_name, "r")
                analysis.generation_plot(stat_file, errorbars=False)
            finally:
                stat_file.close()

class CustomOptimizerA(__Optimizer):
    def __init__(
        self,
        max_constraints,
        min_constraints,
        evaluator,
        mutation_rate=0.2,
        max_evaluations=100,
        population_size=10,
        num_selected=None,
        tourn_size=2,
        num_elites=1,
        maximize=False,
        num_offspring=None,
        seeds=[],
        verbose=False,
    ):

        super(CustomOptimizerA, self).__init__(
            max_constraints,
            min_constraints,
            evaluator,
            mutation_rate,
            maximize,
            seeds,
            population_size,
        )

        self.max_evaluations = max_evaluations
        self.tourn_size = tourn_size
        self.num_elites = num_elites
        self.mutation_rate = mutation_rate
        self.verbose = verbose

        if num_selected == None:
            self.num_selected = population_size
        else:
            self.num_selected = num_selected
        if num_offspring == None:
            self.num_offspring = population_size - self.num_selected
        else:
            self.num_offspring = num_offspring

    def optimize(self, do_plot=True, seed=int(time()), summary_dir=None):

        rand = Random()
        rand.seed(seed)

        if summary_dir is None:
            cwd = os.getcwd()
            summary_dir = os.path.dirname(cwd) + "/data/"

        if not os.path.exists(summary_dir):
            os.mkdir(summary_dir)

        stat_file_name = summary_dir + "/ga_statistics.csv"
        ind_file_name = summary_dir + "/ga_individuals.csv"

        stat_file = open(stat_file_name, "w")
        ind_file = open(ind_file_name, "w")
        print("Created files: %s and %s" % (stat_file_name, ind_file_name))

        if self.verbose:
            logger = logging.getLogger("inspyred.ec")
            logger.setLevel(logging.DEBUG)

            ch = logging.StreamHandler()
            # ch.setLevel(logging.DEBUG)
            formatter = logging.Formatter(">>> EC: - %(levelname)s - %(message)s")
            ch.setFormatter(formatter)
            logger.addHandler(ch)

        algorithm = ec.EvolutionaryComputation(rand)
        algorithm.observer = observers.file_observer
        algorithm.terminator = terminators.evaluation_termination
        algorithm.selector = selectors.tournament_selection
        algorithm.replacer = replacers.steady_state_replacement
        algorithm.variator = [variators.blend_crossover, variators.gaussian_mutation]

        final_pop = algorithm.evolve(
            generator=self.uniform_random_chromosome,
            evaluator=self.evaluator.evaluate,
            pop_size=self.population_size,
            maximize=self.maximize,
            bounder=ec.Bounder(
                lower_bound=self.min_constraints, upper_bound=self.max_constraints
            ),
            num_selected=self.num_selected,
            tourn_size=self.tourn_size,
            num_elites=self.num_elites,
            num_offspring=self.num_offspring,
            max_evaluations=self.max_evaluations,
            mutation_rate=self.mutation_rate,
            statistics_file=stat_file,
            seeds=self.seeds,
            individuals_file=ind_file,
        )

        stat_file.close()
        ind_file.close()

        self.print_report(final_pop, do_plot, stat_file_name)

        # return the parameter set for the best individual

        return final_pop[0].candidate, final_pop[0].fitness
<|endoftext|>


<commit_before>from cocos.actions import *
import cocos

import pyglet

class TextLayer(cocos.layer.Layer):
    """
        Layer that setups text labels and apply some actions on these labels
    """
    def __init__(self):
        super(TextLayer, self).__init__()

        window_width, window_height = cocos.director.director.get_window_size()

        label = cocos.text.Label("WoW Python Vigo!",
                                 font_name='Comic Sans MS',
                                 font_size=48,
                                 color=(255, 255, 255, 255),
                                 anchor_x='center',
                                 anchor_y='center')
        label.position = window_width*0.5, window_height*0.2
        label.do(Repeat(JumpTo(label.position, duration=2)))
        self.add(label, z=1)

        label2 = cocos.text.Label("Dat Scene",
                                  font_name='Comic Sans MS',
                                  font_size=52,
                                  color=(255, 0, 255, 255),
                                  anchor_x='center',
                                  anchor_y='center')
        label2.position = window_width*0.8, window_height*0.8
        label2.do(Repeat(ScaleBy(1.1, 0.2) + Reverse(ScaleBy(1.1, 0.2))))
        self.add(label2, z=1)

        label3 = cocos.text.Label("Such effects",
                                  font_name='Comic Sans MS',
                                  font_size=40,
                                  color=(0, 255, 255, 255),
                                  anchor_x='center',
                                  anchor_y='center')
        label3.position = window_width*0.25, window_height*0.6
        label3.do(Repeat(FadeTo(0, 1) + FadeTo(255, 1)))
        self.add(label3, z=1)

class DogeLayer(cocos.layer.Layer):
    """
        Layer that represents Doge and his jumping sunglasses!
    """
    def __init__(self):
        super(DogeLayer, self).__init__()

        window_width, window_height = cocos.director.director.get_window_size()
        self.position = 0, 100

        sprite = cocos.sprite.Sprite('sunglasses.png')
        sprite.scale = 2.5
        sprite.rotation = 6
        sprite.position = window_width*0.5, window_height
        sprite.do(JumpTo((window_width*0.4, window_height*0.6), duration=1) +
                  Repeat(RotateBy(10, 0.6) + Reverse(RotateBy(10, 0.6))))
        self.add(sprite, z=1)

        sprite = cocos.sprite.Sprite('doge.jpg')
        sprite.position = window_width*0.5, window_height*0.4
        self.add(sprite, z=0)

if __name__ == '__main__':
    # We configure the base assets dir
    pyglet.resource.path.append('../../assets/')
    pyglet.resource.reindex()

    # Scenes and layers configuration
    cocos.director.director.init(width=1024, height=768)
    main_scene = cocos.scene.Scene()
    main_scene.add(TextLayer(), z=1)
    main_scene.add(DogeLayer(), z=0)
    cocos.director.director.run(main_scene)

<commit_msg>Modifying the actions example<commit_after>from cocos.actions import *
import cocos

import pyglet

class TextLayer(cocos.layer.Layer):
    """
        Layer that setups text labels and apply some actions on these labels
    """
    def __init__(self):
        super(TextLayer, self).__init__()

        window_width, window_height = cocos.director.director.get_window_size()

        label = cocos.text.Label("WoW Python Vigo!",
                                 font_name='Comic Sans MS',
                                 font_size=48,
                                 color=(255, 255, 255, 255),
                                 anchor_x='center',
                                 anchor_y='center')
        label.position = window_width*0.5, window_height*0.2
        label.do(Repeat(JumpTo(label.position, height=100, duration=2)))
        self.add(label, z=1)

        label2 = cocos.text.Label("Dat Scene",
                                  font_name='Comic Sans MS',
                                  font_size=52,
                                  color=(255, 0, 255, 255),
                                  anchor_x='center',
                                  anchor_y='center')
        label2.position = window_width*0.8, window_height*0.8
        label2.do(Repeat(ScaleBy(1.1, 0.2) + Reverse(ScaleBy(1.1, 0.2))))
        self.add(label2, z=1)

        label3 = cocos.text.Label("Such effects",
                                  font_name='Comic Sans MS',
                                  font_size=40,
                                  color=(0, 255, 255, 255),
                                  anchor_x='center',
                                  anchor_y='center')
        label3.position = window_width*0.25, window_height*0.6
        label3.do(Repeat(FadeTo(0, 1) + FadeTo(255, 1)))
        self.add(label3, z=1)

class DogeLayer(cocos.layer.Layer):
    """
        Layer that represents Doge and his jumping sunglasses!
    """
    def __init__(self):
        super(DogeLayer, self).__init__()

        window_width, window_height = cocos.director.director.get_window_size()
        self.position = 0, 100

        sprite = cocos.sprite.Sprite('sunglasses.png')
        sprite.scale = 2.5
        sprite.rotation = 6
        sprite.position = window_width*0.5, window_height
        sprite.do(JumpTo((window_width*0.4, window_height*0.6), duration=1) +
                  Repeat(RotateBy(10, 0.6) + Reverse(RotateBy(10, 0.6))))
        self.add(sprite, z=1)

        sprite = cocos.sprite.Sprite('doge.jpg')
        sprite.position = window_width*0.5, window_height*0.4
        self.add(sprite, z=0)

if __name__ == '__main__':
    # We configure the base assets dir
    pyglet.resource.path.append('../../assets/')
    pyglet.resource.reindex()

    # Scenes and layers configuration
    cocos.director.director.init(width=1024, height=768)
    main_scene = cocos.scene.Scene()
    main_scene.add(TextLayer(), z=1)
    main_scene.add(DogeLayer(), z=0)
    cocos.director.director.run(main_scene)

<|endoftext|>


<commit_before>"""
Core geometric operations
=========================

Mechanisms for converting axis-aligned bounding boxes (:py:class:`Envelope`) and polygonal boundaries
(:py:class:`Boundary`) from one spatial reference system to another. It is used internally to convert the regions
represented by tiles between spatial references.

"""

import _gdal
import json
from TileStache.Geography import Point
from osgeo import osr, ogr, gdal, gdal_array
import numpy as np
import pyproj

def boundary_from_envelope(envelope):
    """Construct a :py:class:`Boundary` from an :py:class:`Envelope`. The boundary is the bounding box which encloses
    the envelope.

    :param envelope: the envelope to convert into a boundary
    :type envelope: :py:class:`Envelope`

    """

    wkt = 'POLYGON ((%s))' % (
            ','.join(['%f %f' % x for x in [
                (envelope.left,envelope.top),
                (envelope.right,envelope.top),
                (envelope.right,envelope.bottom),
                (envelope.left,envelope.bottom),
                (envelope.left,envelope.top)
            ]]),
    )
    geom = ogr.CreateGeometryFromWkt(wkt)
    geom.AssignSpatialReference(envelope.spatial_reference)
    return Boundary(geom)

class ProjectionError(RuntimeError):
    """Error raised when projection from one spatial reference to another has failed.

    :param message: a description of the error
    :type message: str

    """

    def __init__(self, message):
        super(ProjectionError, self).__init__(message)

class Boundary(object):
    """A polygon defining an enclosing boundary for a region. A boundary *always* has an associated spatial reference.

    Generally you'll want to use the :py:func:`boundary_from_envelope` function to construct a boundary unless you have
    specific needs.

    :param geometry: boundary geometry
    :type geometry: :py:class:`ogr.Geometry` with polygon type

    .. py:data:: geometry

        An underlying instance of ogr.Geometry which specifies the boundary polygon.

    """

    def __init__(self, geometry):
        if geometry is None:
            raise ValueError('geometry cannot be None')
        if geometry.GetSpatialReference() is None:
            raise ValueError('geometry must have an associated spatial reference')

        self.geometry = geometry

    def envelope(self):
        """Calculate the bounding axis-aligned envelope which entirely contains this boundary.

        Since a boundary is an arbitrary polygon and an envelope is only an axis-aligned bounding box, the returned
        envelope will also contain point not within this boundary.

        :rtype: :py:class:`Envelope`

        """

        l,r,b,t = self.geometry.GetEnvelope()
        return Envelope(l,r,t,b,self.geometry.GetSpatialReference())

    def contains_point(self, x, y):
        """Convenience function to test if a point is contained within this boundary.

        :param x: the point's x-co-ordinate
        :type x: float
        :param y: the point's y-co-ordinate
        :type y: float
        :rtype: bool

        """

        pt = ogr.Geometry(ogr.wkbPoint)
        pt.AddPoint_2D(x,y)
        return self.geometry.Contains(pt)

    def transform_to(self, other_spatial_reference, src_seg_len=None, dst_seg_len=None):
        """Transform this boundary into another spatial reference.

        :param other_spatial_reference: the destination spatial reference
        :type other_spatial_reference: :py:class:`osr.SpatialReference`
        :param src_seg_len: the maximum length of a boundary segment in the source spatial reference
        :type src_seg_len: float or None
        :param dst_seg_len: the maximum length of a boundary segment in the detination spatial reference
        :type dst_seg_len: float or None
        :rtype: :py:class:`Boundary`

        This is a subtle function. The boundary is defined as a polygon of linear segments but, of course, lines do not
        necessarily map to lines in a projection. This function may optionally segmentise the boundary before
        transforming and simplify it afterwards.

        If *src_seg_len* is not None, it specifies the maximum length in the source spatial reference for boundary
        segments. If known, set this to the approximate length of a pixel in the source spatial reference.

        If *dst_seg_len* is not None, it specifies the maximum length in the destination spatial reference used to
        simplify the transformed boundary. If known, set this to the approximate length of a pixel in the destination
        spatial reference.

        """

        geom = self.geometry.Clone()
        if src_seg_len is not None:
            geom.Segmentize(float(src_seg_len))

        old_opt = gdal.GetConfigOption('OGR_ENABLE_PARTIAL_REPROJECTION')
        gdal.SetConfigOption('OGR_ENABLE_PARTIAL_REPROJECTION', 'TRUE')
        err = geom.TransformTo(other_spatial_reference)
        gdal.SetConfigOption('OGR_ENABLE_PARTIAL_REPROJECTION', old_opt)
        if err != 0:
            raise ProjectionError('Unable to project boundary into target projection (%s).' % (err,))

        if dst_seg_len is not None:
            geom.Simplify(dst_seg_len)

        return Boundary(geom)

    def __str__(self):
        return 'Boundary(%s, %s)' % (self.geometry, self.geometry.GetSpatialReference().ExportToPrettyWkt())

    def __repr__(self):
        return 'Boundary(%s, %s)' % (self.geometry, self.geometry.GetSpatialReference().ExportToWkt())

class Envelope(object):
    """An axis-aligned bounding box in a particular spatial reference.

    :param left: the left-most x co-ordinate
    :param right: the right-most x co-ordinate
    :param top: the top-most y co-ordinate
    :param bottom: the bottom-most y co-ordinate

    """

    def __init__(self, left, right, top, bottom, spatial_reference):
        self.left = left
        self.right = right
        self.top = top
        self.bottom = bottom
        self.spatial_reference = spatial_reference

    def top_left(self):
        """Obtain the top-left point for this envelope.

        :rtype: :py:class:`TileStache.Geography.Point`

        """
        return Point(self.left, self.top)

    def bottom_right(self):
        """Obtain the bottom-right point for this envelope.

        :rtype: :py:class:`TileStache.Geography.Point`

        """
        return Point(self.right, self.botom)

    def offset(self):
        """Return a tuple giving the offset from the top-left corner to the bottom right."""
        return (self.right - self.left, self.bottom - self.top)

    def size(self):
        """Return a tuple giving the *absolute* offset from the top-left corner to the bottom right."""
        return map(abs, self.offset())

    def transform_to(self, other_spatial_reference, src_seg_len=None, dst_seg_len=None):
        """Return a envelope which contains this envelope in a target spatial reference.

        This is a convenience function which converts the envelope into a rectangular boundary, transforms it and
        returns the envelope of the resulting boundary.

        :param other_spatial_reference: the destination spatial reference
        :type other_spatial_reference: :py:class:`osr.SpatialReference`
        :param src_seg_len: the maximum length of a boundary segment in the source spatial reference
        :type src_seg_len: float or None
        :param dst_seg_len: the maximum length of a boundary segment in the detination spatial reference
        :type dst_seg_len: float or None
        :rtype: :py:class:`Envelope`

        """
        return boundary_from_envelope(self). \
            transform_to(other_spatial_reference, src_seg_len, dst_seg_len).envelope()

    def __repr__(self):
        return 'Envelope(%f,%f,%f,%f)' % (self.left, self.right, self.top, self.bottom)

    def __str__(self):
        return '(%f => %f, %f => %f)' % (self.left, self.right, self.top, self.bottom)

def to_rgba_unknown(array):
    array = np.atleast_2d(array)
    rgba = np.empty(array.shape[:2] + (4,))
    red = np.reshape(np.arange(array.shape[1], dtype=np.float32) / array.shape[1], (1, array.shape[1]))
    green = np.reshape(np.arange(array.shape[0], dtype=np.float32) / array.shape[0], (array.shape[0], 1))
    alpha = 1
    mask = np.ma.getmask(array)
    if mask is not np.ma.nomask:
        alpha = np.where(np.any(mask, 2), 0.0, 1.0)
    rgba[:,:,0] = np.repeat(red, array.shape[0], 0) * alpha
    rgba[:,:,1] = np.repeat(green, array.shape[1], 1) * alpha
    rgba[:,:,2] = 0
    rgba[:,:,3] = alpha
    return rgba

class RgbaFromBands(object):
    # Band interpretations
    RED         = 'RED'
    GREEN       = 'GREEN'
    BLUE        = 'BLUE'
    ALPHA       = 'ALPHA'
    GRAY        = 'GRAY'
    NONE        = 'NONE'

    def __init__(self, bands, is_premultiplied):
        self.bands = bands
        self.is_premultiplied = is_premultiplied

    def __call__(self, array):
        rgba = to_rgba_unknown(array)
        if np.any(rgba[:,:,3] != 1.0):
            mask_alpha = rgba[:,:,3]
        else:
            mask_alpha = 1.0

        for idx, band in enumerate(self.bands):
            scale = (band[1] if len(band) >= 2 else 1.0) * mask_alpha
            interp = band[0]

            if interp is RgbaFromBands.GRAY:
                rgba[:,:,:3] = np.repeat(array[:,:,idx], 3, 2) * scale
            elif interp is RgbaFromBands.RED:
                rgba[:,:,0] = array[:,:,idx]
                rgba[:,:,0] *= scale
            elif interp is RgbaFromBands.GREEN:
                rgba[:,:,1] = array[:,:,idx]
                rgba[:,:,1] *= scale
            elif interp is RgbaFromBands.BLUE:
                rgba[:,:,2] = array[:,:,idx]
                rgba[:,:,2] *= scale
            elif interp is RgbaFromBands.ALPHA:
                rgba[:,:,3] = array[:,:,idx]
                rgba[:,:,3] *= scale

        if not self.is_premultiplied:
            for i in xrange(3):
                rgba[:,:,i] *= rgba[:,:,3]

        return rgba

class Raster(object):
    def __init__(self, array, envelope, to_rgba=None, can_interpolate=True, prototype=None):
        self.array = np.atleast_3d(np.float32(array))
        self.envelope = envelope
        if to_rgba is None:
            to_rgba = to_rgba_unknown
        self.to_rgba_cb = to_rgba
        self.can_interpolate = can_interpolate

        if prototype is not None:
            self.to_rgba_cb = prototype.to_rgba_cb
            self.can_interpolate = prototype.can_interpolate

    def to_rgba(self):
        return self.to_rgba_cb(self.array)

    def as_rgba_dataset(self):
        arr = self.to_rgba()
        if len(arr.shape) > 2:
            arr = arr.transpose((2,0,1))
        ds = gdal_array.OpenArray(np.uint8(255.0*np.clip(arr,0,1)))
        ds.SetProjection(self.envelope.spatial_reference.ExportToWkt())
        size = [ds.RasterXSize, ds.RasterYSize]
        xscale, yscale = [float(x[0])/float(x[1]) for x in zip(self.envelope.offset(), size)]
        ds.SetGeoTransform((self.envelope.left, xscale, 0, self.envelope.top, 0, yscale))

        return ds

    def as_dataset(self):
        arr = self.array
        if len(arr.shape) > 2:
            arr = arr.transpose((2,0,1))

        ds = gdal_array.OpenArray(arr)

        ds.SetProjection(self.envelope.spatial_reference.ExportToWkt())
        size = [ds.RasterXSize, ds.RasterYSize]
        xscale, yscale = [float(x[0])/float(x[1]) for x in zip(self.envelope.offset(), size)]
        ds.SetGeoTransform((self.envelope.left, xscale, 0, self.envelope.top, 0, yscale))

        mask = np.ma.getmask(self.array)
        mask_ds = None
        if mask is not np.ma.nomask:
            # There is some mask we need to overlay onto the dataset
            for i in xrange(1, ds.RasterCount+1):
                band = ds.GetRasterBand(i)
                band.SetNoDataValue(float('nan'))
                band.WriteArray(np.where(mask[:,:,i-1], float('nan'), self.array[:,:,i-1]))

        return ds

    def write_tiff(self, filename):
        driver = gdal.GetDriverByName('GTiff')
        driver.CreateCopy(filename, self.as_rgba_dataset())

    @classmethod
    def from_dataset(cls, ds, mask_band=None, **kwargs):
        ds_array = ds.ReadAsArray()
        if len(ds_array.shape) > 2:
            ds_array = ds_array.transpose((1,2,0))
        else:
            ds_array = np.atleast_3d(ds_array)

        if mask_band is not None:
            mask = np.atleast_3d(mask_band.ReadAsArray()) == 0
            if np.any(mask):
                ds_array = np.ma.masked_where(np.repeat(mask, ds_array.shape[2], 2), ds_array)

        srs = osr.SpatialReference()
        srs.ImportFromWkt(ds.GetProjection())
        envelope = _gdal.dataset_envelope(ds, srs)

        can_interpolate = gdal.GCI_PaletteIndex not in [
            ds.GetRasterBand(i).GetColorInterpretation()
            for i in xrange(1, ds.RasterCount+1)
        ]
        if 'can_interpolate' in kwargs:
            can_interpolate = can_interpolate and kwargs['can_interpolate']
            del kwargs['can_interpolate']

        return Raster(ds_array, envelope, can_interpolate=can_interpolate, **kwargs)
<commit_msg>do some pylint cleanups<commit_after>"""
Core geometric operations
=========================

Mechanisms for converting axis-aligned bounding boxes (:py:class:`Envelope`) and polygonal boundaries
(:py:class:`Boundary`) from one spatial reference system to another. It is used internally to convert the regions
represented by tiles between spatial references.

"""

# Hss to be first FSR otherwise there is a segfault :(
from TileStache.Geography import Point

from ._gdal import dataset_envelope
import json
import numpy as np
from osgeo import osr, ogr, gdal, gdal_array
import pyproj

def boundary_from_envelope(envelope):
    """Construct a :py:class:`Boundary` from an :py:class:`Envelope`. The boundary is the bounding box which encloses
    the envelope.

    :param envelope: the envelope to convert into a boundary
    :type envelope: :py:class:`Envelope`

    """

    wkt = 'POLYGON ((%s))' % (
            ','.join(['%f %f' % x for x in [
                (envelope.left,envelope.top),
                (envelope.right,envelope.top),
                (envelope.right,envelope.bottom),
                (envelope.left,envelope.bottom),
                (envelope.left,envelope.top)
            ]]),
    )
    geom = ogr.CreateGeometryFromWkt(wkt)
    geom.AssignSpatialReference(envelope.spatial_reference)
    return Boundary(geom)

class ProjectionError(RuntimeError):
    """Error raised when projection from one spatial reference to another has failed.

    :param message: a description of the error
    :type message: str

    """

    def __init__(self, message):
        super(ProjectionError, self).__init__(message)

class Boundary(object):
    """A polygon defining an enclosing boundary for a region. A boundary *always* has an associated spatial reference.

    Generally you'll want to use the :py:func:`boundary_from_envelope` function to construct a boundary unless you have
    specific needs.

    :param geometry: boundary geometry
    :type geometry: :py:class:`ogr.Geometry` with polygon type

    .. py:data:: geometry

        An underlying instance of ogr.Geometry which specifies the boundary polygon.

    """

    def __init__(self, geometry):
        if geometry is None:
            raise ValueError('geometry cannot be None')
        if geometry.GetSpatialReference() is None:
            raise ValueError('geometry must have an associated spatial reference')

        self.geometry = geometry

    def envelope(self):
        """Calculate the bounding axis-aligned envelope which entirely contains this boundary.

        Since a boundary is an arbitrary polygon and an envelope is only an axis-aligned bounding box, the returned
        envelope will also contain point not within this boundary.

        :rtype: :py:class:`Envelope`

        """

        left, right, bottom, top = self.geometry.GetEnvelope()
        return Envelope(left, right, top, bottom, self.geometry.GetSpatialReference())

    def contains_point(self, x, y):
        """Convenience function to test if a point is contained within this boundary.

        :param x: the point's x-co-ordinate
        :type x: float
        :param y: the point's y-co-ordinate
        :type y: float
        :rtype: bool

        """

        point = ogr.Geometry(ogr.wkbPoint)
        point.AddPoint_2D(x, y)
        return self.geometry.Contains(point)

    def transform_to(self, other_spatial_reference, src_seg_len=None, dst_seg_len=None):
        """Transform this boundary into another spatial reference.

        :param other_spatial_reference: the destination spatial reference
        :type other_spatial_reference: :py:class:`osr.SpatialReference`
        :param src_seg_len: the maximum length of a boundary segment in the source spatial reference
        :type src_seg_len: float or None
        :param dst_seg_len: the maximum length of a boundary segment in the detination spatial reference
        :type dst_seg_len: float or None
        :rtype: :py:class:`Boundary`

        This is a subtle function. The boundary is defined as a polygon of linear segments but, of course, lines do not
        necessarily map to lines in a projection. This function may optionally segmentise the boundary before
        transforming and simplify it afterwards.

        If *src_seg_len* is not None, it specifies the maximum length in the source spatial reference for boundary
        segments. If known, set this to the approximate length of a pixel in the source spatial reference.

        If *dst_seg_len* is not None, it specifies the maximum length in the destination spatial reference used to
        simplify the transformed boundary. If known, set this to the approximate length of a pixel in the destination
        spatial reference.

        """

        geom = self.geometry.Clone()
        if src_seg_len is not None:
            geom.Segmentize(float(src_seg_len))

        old_opt = gdal.GetConfigOption('OGR_ENABLE_PARTIAL_REPROJECTION')
        gdal.SetConfigOption('OGR_ENABLE_PARTIAL_REPROJECTION', 'TRUE')
        err = geom.TransformTo(other_spatial_reference)
        gdal.SetConfigOption('OGR_ENABLE_PARTIAL_REPROJECTION', old_opt)
        if err != 0:
            raise ProjectionError('Unable to project boundary into target projection (%s).' % (err,))

        if dst_seg_len is not None:
            geom.Simplify(dst_seg_len)

        return Boundary(geom)

    def __str__(self):
        return 'Boundary(%s, %s)' % (self.geometry, self.geometry.GetSpatialReference().ExportToPrettyWkt())

    def __repr__(self):
        return 'Boundary(%s, %s)' % (self.geometry, self.geometry.GetSpatialReference().ExportToWkt())

class Envelope(object):
    """An axis-aligned bounding box in a particular spatial reference.

    :param left: the left-most x co-ordinate
    :param right: the right-most x co-ordinate
    :param top: the top-most y co-ordinate
    :param bottom: the bottom-most y co-ordinate

    """

    def __init__(self, left, right, top, bottom, spatial_reference):
        self.left = left
        self.right = right
        self.top = top
        self.bottom = bottom
        self.spatial_reference = spatial_reference

    def top_left(self):
        """Obtain the top-left point for this envelope.

        :rtype: :py:class:`TileStache.Geography.Point`

        """
        return Point(self.left, self.top)

    def bottom_right(self):
        """Obtain the bottom-right point for this envelope.

        :rtype: :py:class:`TileStache.Geography.Point`

        """
        return Point(self.right, self.bottom)

    def offset(self):
        """Return a tuple giving the offset from the top-left corner to the bottom right."""
        return (self.right - self.left, self.bottom - self.top)

    def size(self):
        """Return a tuple giving the *absolute* offset from the top-left corner to the bottom right."""
        return [abs(x) for x in self.offset()]

    def transform_to(self, other_spatial_reference, src_seg_len=None, dst_seg_len=None):
        """Return a envelope which contains this envelope in a target spatial reference.

        This is a convenience function which converts the envelope into a rectangular boundary, transforms it and
        returns the envelope of the resulting boundary.

        :param other_spatial_reference: the destination spatial reference
        :type other_spatial_reference: :py:class:`osr.SpatialReference`
        :param src_seg_len: the maximum length of a boundary segment in the source spatial reference
        :type src_seg_len: float or None
        :param dst_seg_len: the maximum length of a boundary segment in the detination spatial reference
        :type dst_seg_len: float or None
        :rtype: :py:class:`Envelope`

        """
        return boundary_from_envelope(self). \
            transform_to(other_spatial_reference, src_seg_len, dst_seg_len).envelope()

    def __repr__(self):
        return 'Envelope(%f,%f,%f,%f)' % (self.left, self.right, self.top, self.bottom)

    def __str__(self):
        return '(%f => %f, %f => %f)' % (self.left, self.right, self.top, self.bottom)

def to_rgba_unknown(array):
    """Default conversion from an array of unknown type to an RGBA array. This simply fills the array with a default
    pattern.

    """

    array = np.atleast_2d(array)
    rgba = np.empty(array.shape[:2] + (4,))
    red = np.reshape(np.arange(array.shape[1], dtype=np.float32) / array.shape[1], (1, array.shape[1]))
    green = np.reshape(np.arange(array.shape[0], dtype=np.float32) / array.shape[0], (array.shape[0], 1))
    alpha = 1
    mask = np.ma.getmask(array)
    if mask is not np.ma.nomask:
        alpha = np.where(np.any(mask, 2), 0.0, 1.0)
    rgba[:, :, 0] = np.repeat(red, array.shape[0], 0) * alpha
    rgba[:, :, 1] = np.repeat(green, array.shape[1], 1) * alpha
    rgba[:, :, 2] = 0
    rgba[:, :, 3] = alpha
    return rgba

class RgbaFromBands(object):
    # Band interpretations
    RED         = 'RED'
    GREEN       = 'GREEN'
    BLUE        = 'BLUE'
    ALPHA       = 'ALPHA'
    GRAY        = 'GRAY'
    NONE        = 'NONE'

    def __init__(self, bands, is_premultiplied):
        self.bands = bands
        self.is_premultiplied = is_premultiplied

    def __call__(self, array):
        rgba = to_rgba_unknown(array)
        if np.any(rgba[:,:,3] != 1.0):
            mask_alpha = rgba[:,:,3]
        else:
            mask_alpha = 1.0

        for idx, band in enumerate(self.bands):
            scale = (band[1] if len(band) >= 2 else 1.0) * mask_alpha
            interp = band[0]

            if interp is RgbaFromBands.GRAY:
                rgba[:, :, :3] = np.repeat(array[:,:,idx], 3, 2) * scale
            elif interp is RgbaFromBands.RED:
                rgba[:, :, 0] = array[:,:,idx]
                rgba[:, :, 0] *= scale
            elif interp is RgbaFromBands.GREEN:
                rgba[:, :, 1] = array[:,:,idx]
                rgba[:, :, 1] *= scale
            elif interp is RgbaFromBands.BLUE:
                rgba[:, :, 2] = array[:,:,idx]
                rgba[:, :, 2] *= scale
            elif interp is RgbaFromBands.ALPHA:
                rgba[:, :, 3] = array[:,:,idx]
                rgba[:, :, 3] *= scale

        if not self.is_premultiplied:
            for i in xrange(3):
                rgba[:, :, i] *= rgba[:,:,3]

        return rgba

class Raster(object):
    def __init__(self, array, envelope, to_rgba=None, can_interpolate=True, prototype=None):
        self.array = np.atleast_3d(np.float32(array))
        self.envelope = envelope
        if to_rgba is None:
            to_rgba = to_rgba_unknown
        self.to_rgba_cb = to_rgba
        self.can_interpolate = can_interpolate

        if prototype is not None:
            self.to_rgba_cb = prototype.to_rgba_cb
            self.can_interpolate = prototype.can_interpolate

    def to_rgba(self):
        return self.to_rgba_cb(self.array)

    def as_rgba_dataset(self):
        arr = self.to_rgba()
        if len(arr.shape) > 2:
            arr = arr.transpose((2,0,1))
        ds = gdal_array.OpenArray(np.uint8(255.0*np.clip(arr,0,1)))
        ds.SetProjection(self.envelope.spatial_reference.ExportToWkt())
        size = [ds.RasterXSize, ds.RasterYSize]
        xscale, yscale = [float(x[0])/float(x[1]) for x in zip(self.envelope.offset(), size)]
        ds.SetGeoTransform((self.envelope.left, xscale, 0, self.envelope.top, 0, yscale))

        return ds

    def as_dataset(self):
        arr = self.array
        if len(arr.shape) > 2:
            arr = arr.transpose((2,0,1))

        ds = gdal_array.OpenArray(arr)

        ds.SetProjection(self.envelope.spatial_reference.ExportToWkt())
        size = [ds.RasterXSize, ds.RasterYSize]
        xscale, yscale = [float(x[0])/float(x[1]) for x in zip(self.envelope.offset(), size)]
        ds.SetGeoTransform((self.envelope.left, xscale, 0, self.envelope.top, 0, yscale))

        mask = np.ma.getmask(self.array)
        mask_ds = None
        if mask is not np.ma.nomask:
            # There is some mask we need to overlay onto the dataset
            for i in xrange(1, ds.RasterCount+1):
                band = ds.GetRasterBand(i)
                band.SetNoDataValue(float('nan'))
                band.WriteArray(np.where(mask[:,:,i-1], float('nan'), self.array[:,:,i-1]))

        return ds

    def write_tiff(self, filename):
        driver = gdal.GetDriverByName('GTiff')
        driver.CreateCopy(filename, self.as_rgba_dataset())

    @classmethod
    def from_dataset(cls, ds, mask_band=None, **kwargs):
        ds_array = ds.ReadAsArray()
        if len(ds_array.shape) > 2:
            ds_array = ds_array.transpose((1,2,0))
        else:
            ds_array = np.atleast_3d(ds_array)

        if mask_band is not None:
            mask = np.atleast_3d(mask_band.ReadAsArray()) == 0
            if np.any(mask):
                ds_array = np.ma.masked_where(np.repeat(mask, ds_array.shape[2], 2), ds_array)

        srs = osr.SpatialReference()
        srs.ImportFromWkt(ds.GetProjection())
        envelope = dataset_envelope(ds, srs)

        can_interpolate = gdal.GCI_PaletteIndex not in [
            ds.GetRasterBand(i).GetColorInterpretation()
            for i in xrange(1, ds.RasterCount+1)
        ]
        if 'can_interpolate' in kwargs:
            can_interpolate = can_interpolate and kwargs['can_interpolate']
            del kwargs['can_interpolate']

        return Raster(ds_array, envelope, can_interpolate=can_interpolate, **kwargs)
<|endoftext|>


<commit_before>#   Copyright 2020 Google LLC

#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

""" util unit tests
"""

import unittest

import numpy

from fqe import util

class UnitTest(unittest.TestCase):
    """unit tests
    """

    def test_alpha_beta_electrons(self):
        """Check to make sure that the correct number of alpha and beta
        electrons are parsed from the number and multiplicity
        """
        self.assertTupleEqual((1, 1), util.alpha_beta_electrons(2, 0))
        self.assertTupleEqual((4, 1), util.alpha_beta_electrons(5, 3))
        self.assertTupleEqual((0, 5), util.alpha_beta_electrons(5, -5))

    def test_alpha_beta_error(self):
        """Check to make sure that the correct number of alpha and beta
        electrons are parsed from the number and multiplicity
        """
        self.assertRaises(ValueError, util.alpha_beta_electrons, -1, 0)
        self.assertRaises(ValueError, util.alpha_beta_electrons, 1, 2)

    def test_bubblesort_order(self):
        """Check that bubble sort works.
        """
        length = 5
        test_list = [(length - 1 - i) for i in range(length)]
        ordered_list = [i for i in range(length)]
        util.bubblesort(test_list)
        self.assertListEqual(ordered_list, test_list)

    def test_bubblesort_permutation_count(self):
        """ Make sure that we are counting the correct number of permutations
        to sort the list
        """
        length = 2
        test_list = [(length - 1 - i) for i in range(length)]
        self.assertEqual(1, util.bubblesort(test_list))
        length = 3
        test_list = [(length - 1 - i) for i in range(length)]
        self.assertEqual(3, util.bubblesort(test_list))
        test_list = [2, 0, 1]
        self.assertEqual(2, util.bubblesort(test_list))

    def test_reverse_bubblesort_permutation_count(self):
        """ Make sure that we are counting the correct number of permutations
        to sort the list
        """
        test_list = [[0, 0], [1, 0]]
        self.assertEqual(1, util.reverse_bubble_list(test_list))
        test_list = [[0, 0], [1, 0], [2, 0]]
        self.assertEqual(3, util.reverse_bubble_list(test_list))
        test_list = [[0, 0], [2, 0], [1, 0]]
        self.assertEqual(2, util.reverse_bubble_list(test_list))

    def test_configuration_key_union_empty(self):
        """The union of no configuration keys should be an empty list
        """
        self.assertListEqual([], util.configuration_key_union())

    def test_configuration_key_union(self):
        """The union of no configuration keys should be an empty list
        """
        configs0 = [(2, 0), (3, 1)]
        configs1 = [(2, 0), (5, 1), (6, -2)]
        testset = set([(2, 0), (3, 1), (5, 1), (6, -2)])
        self.assertSetEqual(testset, set(util.configuration_key_union(configs0, configs1)))

    def test_configuration_key_union_many(self):
        """The union of many different keys should be all of them
        """
        configs0 = [(2, 0)]
        configs1 = [(5, 1)]
        configs2 = [(6, -2)]
        configs3 = [(3, -3)]
        refset = set([(2, 0), (5, 1), (6, -2), (3, -3)])
        testset = set(util.configuration_key_union(configs0, configs1,
                                                   configs2, configs3))
        self.assertSetEqual(testset, refset)

    def test_configuration_key_intersection_none(self):
        """If there are no keys in common the intersection should be zero
        """
        self.assertListEqual([], util.configuration_key_intersection([(2, 0)], [(2, 2)]))

    def test_configuration_key_intersection(self):
        """Check that the intersection returns the intersection
        """
        configs0 = [(10, 0), (3, 1), (5, -1)]
        configs1 = [(2, 0), (3, 1), (3, -1)]
        self.assertListEqual([(3, 1)], util.configuration_key_intersection(configs0, configs1))

    def test_invert_bitstring_with_mask(self):
        """When inverting the occupation we want to maintain the number of orbitals
        """
        ref = 8
        self.assertEqual(ref, util.invert_bitstring_with_mask(7, 4))
        ref = 8 + 16 + 32 + 64 + 128
        self.assertEqual(ref, util.invert_bitstring_with_mask(7, 8))

    def test_ltlt_index_min(self):
        """If we have a zero dimesnion tensor there should be no pointers to it
        """
        _gtest = util.ltlt_index_generator(0)
        _test = [i for i in _gtest]
        self.assertListEqual(_test, [])

    def test_ltlt_index(self):
        """Access unique elements of a lower triangular lower triangular
        matrix
        """
        index_list = [
            (0, 0, 0, 0),
            (1, 0, 0, 0),
            (1, 0, 1, 0),
            (1, 1, 0, 0),
            (1, 1, 1, 0),
            (1, 1, 1, 1)
            ]
        _gtest = util.ltlt_index_generator(2)
        _test = [i for i in _gtest]
        self.assertListEqual(_test, index_list)

    def test_bitstring_groundstate(self):
        """The ground state bitstring has the n lowest bits flipped
        """
        self.assertEqual(15, util.init_bitstring_groundstate(4))

    def test_qubit_particle_number_sector(self):
        """Find the vectors which are the basis for a particular particle
        number.
        """
        zero = 0
        one = 1
        ref = [
            numpy.array([zero, zero, one, zero], dtype=numpy.int),
            numpy.array([zero, one, zero, zero], dtype=numpy.int)
            ]
        test = util.qubit_particle_number_sector(2, 1)
        for i, j in zip(test, ref):
            self.assertEqual(i.all(), j.all())

    def test_qubit_particle_number_index_spin(self):
        """Find the indexes which point to the correct coefficients in a qubit
        particle number sector and return the total spin.
        """
        ref = [(3, 0), (5, 2), (6, 0), (9, 0), (10, -2), (12, 0)]
        test = util.qubit_particle_number_index_spin(4, 2)
        self.assertListEqual(ref, test)

    def test_qubit_config_sector(self):
        """Find the basis vectors for a particular particle number and spin
        configuration
        """
        zero = 0
        one = 1
        lowstate = [zero, zero, one, zero, zero, zero, zero, zero, zero, zero,
                    zero, zero, zero, zero, zero, zero]
        highstate = [zero, zero, zero, zero, zero, zero, zero, zero, one, zero,
                     zero, zero, zero, zero, zero, zero]
        ref = [numpy.array(lowstate, dtype=numpy.int),
               numpy.array(highstate, dtype=numpy.int)]
        test = util.qubit_config_sector(4, 1, 1)
        for i, j in zip(test, ref):
            self.assertEqual(i.all(), j.all())
        ref = [numpy.array([zero, zero, zero, one], dtype=numpy.int)]
        test = util.qubit_config_sector(2, 2, 0)
        for i, j in zip(test, ref):
            self.assertEqual(i.all(), j.all())

    def test_qubit_particle_number_index(self):
        """Find the indexes which point to the correct coefficients in a qubit
        particle number sector and return the total spin.
        """
        ref = [1, 2, 4, 8]
        test = util.qubit_particle_number_index(4, 1)
        self.assertListEqual(ref, test)

    def test_qubit_vacuum(self):
        """The qubit vacuum is the first vector in the qubit basis.
        """
        _gs = numpy.array([1. + .0j, 0. + .0j, 0. + .0j, 0. + .0j], dtype=numpy.complex64)
        self.assertListEqual(list(_gs), list(util.init_qubit_vacuum(2)))

    def test_sort_config_keys(self):
        """Keys are sorted by particle number and then by m_s
        """
        ref = [(0, 0), (1, -1), (3, -3), (3, 1), (5, -2), (5, 1)]
        keys = [(5, 1), (5, -2), (0, 0), (1, -1), (3, -3), (3, 1)]
        test = util.sort_configuration_keys(keys)
        self.assertListEqual(test, ref)

    def test_validate_config(self):
        """Make sure that the configuration validation routine identifies
        problematic values
        """
        self.assertRaises(ValueError, util.validate_config, 0, 0, -1)
        self.assertRaises(ValueError, util.validate_config, 3, 0, 2)
        self.assertRaises(ValueError, util.validate_config, 0, 3, 2)
        self.assertRaises(ValueError, util.validate_config, -1, 1, 2)
        self.assertRaises(ValueError, util.validate_config, 1, -1, 2)
        self.assertIsNone(util.validate_config(0, 0, 0))
        self.assertIsNone(util.validate_config(0, 0, 1))

    def test_zero_transform(self):
        """Ensure that things that should transform do and those that shouldn't
        dont
        """
        self.assertFalse(util.zero_transform(1 + 2 + 4, 8, 3, 6))
        self.assertTrue(util.zero_transform(2 + 4, 8, 1, 6))
        self.assertTrue(util.zero_transform(2 + 4, 4, 2, 6))

    def test_parity_sort_list(self):
        """Sort a list of lists according to the parity of the index in the
        0th element.
        """
        unchanged = [
            [6, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]],
            [7, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]],
            [3, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]],
            [15, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]]
        ]
        nswap, test = util.paritysort_list(unchanged)
        self.assertEqual(nswap, 0)
        self.assertListEqual(unchanged, test)
<commit_msg>test adjustment due to a bug fix<commit_after>#   Copyright 2020 Google LLC

#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

""" util unit tests
"""

import unittest

import numpy

from fqe import util

class UnitTest(unittest.TestCase):
    """unit tests
    """

    def test_alpha_beta_electrons(self):
        """Check to make sure that the correct number of alpha and beta
        electrons are parsed from the number and multiplicity
        """
        self.assertTupleEqual((1, 1), util.alpha_beta_electrons(2, 0))
        self.assertTupleEqual((4, 1), util.alpha_beta_electrons(5, 3))
        self.assertTupleEqual((0, 5), util.alpha_beta_electrons(5, -5))

    def test_alpha_beta_error(self):
        """Check to make sure that the correct number of alpha and beta
        electrons are parsed from the number and multiplicity
        """
        self.assertRaises(ValueError, util.alpha_beta_electrons, -1, 0)
        self.assertRaises(ValueError, util.alpha_beta_electrons, 1, 2)

    def test_bubblesort_order(self):
        """Check that bubble sort works.
        """
        length = 5
        test_list = [(length - 1 - i) for i in range(length)]
        ordered_list = [i for i in range(length)]
        util.bubblesort(test_list)
        self.assertListEqual(ordered_list, test_list)

    def test_bubblesort_permutation_count(self):
        """ Make sure that we are counting the correct number of permutations
        to sort the list
        """
        length = 2
        test_list = [(length - 1 - i) for i in range(length)]
        self.assertEqual(1, util.bubblesort(test_list))
        length = 3
        test_list = [(length - 1 - i) for i in range(length)]
        self.assertEqual(3, util.bubblesort(test_list))
        test_list = [2, 0, 1]
        self.assertEqual(2, util.bubblesort(test_list))

    def test_reverse_bubblesort_permutation_count(self):
        """ Make sure that we are counting the correct number of permutations
        to sort the list
        """
        test_list = [[0, 0], [1, 0]]
        self.assertEqual(1, util.reverse_bubble_list(test_list))
        test_list = [[0, 0], [1, 0], [2, 0]]
        self.assertEqual(3, util.reverse_bubble_list(test_list))
        test_list = [[0, 0], [2, 0], [1, 0]]
        self.assertEqual(2, util.reverse_bubble_list(test_list))

    def test_configuration_key_union_empty(self):
        """The union of no configuration keys should be an empty list
        """
        self.assertListEqual([], util.configuration_key_union())

    def test_configuration_key_union(self):
        """The union of no configuration keys should be an empty list
        """
        configs0 = [(2, 0), (3, 1)]
        configs1 = [(2, 0), (5, 1), (6, -2)]
        testset = set([(2, 0), (3, 1), (5, 1), (6, -2)])
        self.assertSetEqual(testset, set(util.configuration_key_union(configs0, configs1)))

    def test_configuration_key_union_many(self):
        """The union of many different keys should be all of them
        """
        configs0 = [(2, 0)]
        configs1 = [(5, 1)]
        configs2 = [(6, -2)]
        configs3 = [(3, -3)]
        refset = set([(2, 0), (5, 1), (6, -2), (3, -3)])
        testset = set(util.configuration_key_union(configs0, configs1,
                                                   configs2, configs3))
        self.assertSetEqual(testset, refset)

    def test_configuration_key_intersection_none(self):
        """If there are no keys in common the intersection should be zero
        """
        self.assertListEqual([], util.configuration_key_intersection([(2, 0)], [(2, 2)]))

    def test_configuration_key_intersection(self):
        """Check that the intersection returns the intersection
        """
        configs0 = [(10, 0), (3, 1), (5, -1)]
        configs1 = [(2, 0), (3, 1), (3, -1)]
        self.assertListEqual([(3, 1)], util.configuration_key_intersection(configs0, configs1))

    def test_invert_bitstring_with_mask(self):
        """When inverting the occupation we want to maintain the number of orbitals
        """
        ref = 8
        self.assertEqual(ref, util.invert_bitstring_with_mask(7, 4))
        ref = 8 + 16 + 32 + 64 + 128
        self.assertEqual(ref, util.invert_bitstring_with_mask(7, 8))

    def test_ltlt_index_min(self):
        """If we have a zero dimesnion tensor there should be no pointers to it
        """
        _gtest = util.ltlt_index_generator(0)
        _test = [i for i in _gtest]
        self.assertListEqual(_test, [])

    def test_ltlt_index(self):
        """Access unique elements of a lower triangular lower triangular
        matrix
        """
        index_list = [
            (0, 0, 0, 0),
            (1, 0, 0, 0),
            (1, 0, 1, 0),
            (1, 1, 0, 0),
            (1, 1, 1, 0),
            (1, 1, 1, 1)
            ]
        _gtest = util.ltlt_index_generator(2)
        _test = [i for i in _gtest]
        self.assertListEqual(_test, index_list)

    def test_bitstring_groundstate(self):
        """The ground state bitstring has the n lowest bits flipped
        """
        self.assertEqual(15, util.init_bitstring_groundstate(4))

    def test_qubit_particle_number_sector(self):
        """Find the vectors which are the basis for a particular particle
        number.
        """
        zero = 0
        one = 1
        ref = [
            numpy.array([zero, zero, one, zero], dtype=numpy.int),
            numpy.array([zero, one, zero, zero], dtype=numpy.int)
            ]
        test = util.qubit_particle_number_sector(2, 1)
        for i, j in zip(test, ref):
            self.assertEqual(i.all(), j.all())

    def test_qubit_particle_number_index_spin(self):
        """Find the indexes which point to the correct coefficients in a qubit
        particle number sector and return the total spin.
        """
        ref = [(3, 0), (5, -2), (6, 0), (9, 0), (10, 2), (12, 0)]
        test = util.qubit_particle_number_index_spin(4, 2)
        self.assertListEqual(ref, test)

    def test_qubit_config_sector(self):
        """Find the basis vectors for a particular particle number and spin
        configuration
        """
        zero = 0
        one = 1
        lowstate = [zero, zero, one, zero, zero, zero, zero, zero, zero, zero,
                    zero, zero, zero, zero, zero, zero]
        highstate = [zero, zero, zero, zero, zero, zero, zero, zero, one, zero,
                     zero, zero, zero, zero, zero, zero]
        ref = [numpy.array(lowstate, dtype=numpy.int),
               numpy.array(highstate, dtype=numpy.int)]
        test = util.qubit_config_sector(4, 1, 1)
        for i, j in zip(test, ref):
            self.assertEqual(i.all(), j.all())
        ref = [numpy.array([zero, zero, zero, one], dtype=numpy.int)]
        test = util.qubit_config_sector(2, 2, 0)
        for i, j in zip(test, ref):
            self.assertEqual(i.all(), j.all())

    def test_qubit_particle_number_index(self):
        """Find the indexes which point to the correct coefficients in a qubit
        particle number sector and return the total spin.
        """
        ref = [1, 2, 4, 8]
        test = util.qubit_particle_number_index(4, 1)
        self.assertListEqual(ref, test)

    def test_qubit_vacuum(self):
        """The qubit vacuum is the first vector in the qubit basis.
        """
        _gs = numpy.array([1. + .0j, 0. + .0j, 0. + .0j, 0. + .0j], dtype=numpy.complex64)
        self.assertListEqual(list(_gs), list(util.init_qubit_vacuum(2)))

    def test_sort_config_keys(self):
        """Keys are sorted by particle number and then by m_s
        """
        ref = [(0, 0), (1, -1), (3, -3), (3, 1), (5, -2), (5, 1)]
        keys = [(5, 1), (5, -2), (0, 0), (1, -1), (3, -3), (3, 1)]
        test = util.sort_configuration_keys(keys)
        self.assertListEqual(test, ref)

    def test_validate_config(self):
        """Make sure that the configuration validation routine identifies
        problematic values
        """
        self.assertRaises(ValueError, util.validate_config, 0, 0, -1)
        self.assertRaises(ValueError, util.validate_config, 3, 0, 2)
        self.assertRaises(ValueError, util.validate_config, 0, 3, 2)
        self.assertRaises(ValueError, util.validate_config, -1, 1, 2)
        self.assertRaises(ValueError, util.validate_config, 1, -1, 2)
        self.assertIsNone(util.validate_config(0, 0, 0))
        self.assertIsNone(util.validate_config(0, 0, 1))

    def test_zero_transform(self):
        """Ensure that things that should transform do and those that shouldn't
        dont
        """
        self.assertFalse(util.zero_transform(1 + 2 + 4, 8, 3, 6))
        self.assertTrue(util.zero_transform(2 + 4, 8, 1, 6))
        self.assertTrue(util.zero_transform(2 + 4, 4, 2, 6))

    def test_parity_sort_list(self):
        """Sort a list of lists according to the parity of the index in the
        0th element.
        """
        unchanged = [
            [6, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]],
            [7, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]],
            [3, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]],
            [15, ['these', 'values', {'dont' : 6724}, tuple(['mat', 'ter'])]]
        ]
        nswap, test = util.paritysort_list(unchanged)
        self.assertEqual(nswap, 0)
        self.assertListEqual(unchanged, test)
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function, unicode_literals

import markdown
from nib import Processor, markup

@markup(['.md', '.mkdn'])
class MarkdownProcessor(Processor):
    def __init__(self, options):
        self.options = options
        self.markdown = markdown.Markdown(output_format='html5')

    def document(self, document):
        document.short = self.markdown.reset().convert(document.short)
        document.content = self.markdown.reset().convert(document.content)
        document.extension = '.html'
        return document
<commit_msg>Feature: enable markdown footnotes extension<commit_after>from __future__ import absolute_import, division, print_function, unicode_literals

import markdown
from nib import Processor, markup

@markup(['.md', '.mkdn'])
class MarkdownProcessor(Processor):
    def __init__(self, options):
        self.options = options
        config = {
            'markdown.extensions.footnotes': {
                'UNIQUE_IDS': True,
            },
        }
        self.markdown = markdown.Markdown(output_format='html5',
                                          extensions=config.keys(),
                                          extension_configs=config,
                                          )

    def document(self, document):
        document.short = self.markdown.reset().convert(document.short)
        document.content = self.markdown.reset().convert(document.content)
        document.extension = '.html'
        return document
<|endoftext|>