<commit_before>import os
import mimetypes

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.core.files.base import File
from django.core.files.storage import Storage
from django.core.exceptions import ImproperlyConfigured

try:
    from boto.s3.connection import S3Connection, S3ResponseError
    from boto.s3.key import Key
except ImportError:
    raise ImproperlyConfigured, "Could not load Boto's S3 bindings.\
    \nSee http://code.google.com/p/boto/"

ACCESS_KEY_NAME     = getattr(settings, 'AWS_ACCESS_KEY_ID', None)
SECRET_KEY_NAME     = getattr(settings, 'AWS_SECRET_ACCESS_KEY', None)
HEADERS             = getattr(settings, 'AWS_HEADERS', {})
STORAGE_BUCKET_NAME = getattr(settings, 'AWS_STORAGE_BUCKET_NAME', None)
AUTO_CREATE_BUCKET  = getattr(settings, 'AWS_AUTO_CREATE_BUCKET', True)
DEFAULT_ACL         = getattr(settings, 'AWS_DEFAULT_ACL', 'public-read')
QUERYSTRING_AUTH    = getattr(settings, 'AWS_QUERYSTRING_AUTH', True)
QUERYSTRING_EXPIRE  = getattr(settings, 'AWS_QUERYSTRING_EXPIRE', 3600)
LOCATION            = getattr(settings, 'AWS_LOCATION', '')
IS_GZIPPED          = getattr(settings, 'AWS_IS_GZIPPED', False)
GZIP_CONTENT_TYPES  = getattr(settings, 'GZIP_CONTENT_TYPES', (
    'text/css',
    'application/javascript',
    'application/x-javascript'
))

if IS_GZIPPED:
    from gzip import GzipFile

class S3BotoStorage(Storage):
    """Amazon Simple Storage Service using Boto"""

    def __init__(self, bucket=STORAGE_BUCKET_NAME, access_key=None,
                       secret_key=None, acl=DEFAULT_ACL, headers=HEADERS,
                       gzip=IS_GZIPPED, gzip_content_types=GZIP_CONTENT_TYPES):
        self.acl = acl
        self.headers = headers
        self.gzip = gzip
        self.gzip_content_types = gzip_content_types

        if not access_key and not secret_key:
             access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(access_key, secret_key)
        self.bucket = self._get_or_create_bucket(bucket)
        self.bucket.set_acl(self.acl)

    def _get_access_keys(self):
        access_key = ACCESS_KEY_NAME
        secret_key = SECRET_KEY_NAME
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_or_create_bucket(self, name):
        """Retrieves a bucket if it exists, otherwise creates it."""
        try:
            return self.connection.get_bucket(name)
        except S3ResponseError, e:
            if AUTO_CREATE_BUCKET:
                return self.connection.create_bucket(name)
            raise ImproperlyConfigured, ("Bucket specified by "
            "AWS_STORAGE_BUCKET_NAME does not exist. Buckets can be "
            "automatically created by setting AWS_AUTO_CREATE_BUCKET=True")

    def _clean_name(self, name):
        # Useful for windows' paths
        return os.path.normpath(name).replace('\\', '/')

    def _compress_content(self, content):
        """Gzip a given string."""
        zbuf = StringIO()
        zfile = GzipFile(mode='wb', compresslevel=6, fileobj=zbuf)
        zfile.write(content.read())
        zfile.close()
        content.file = zbuf
        return content

    def _open(self, name, mode='rb'):
        name = self._clean_name(name)
        return S3BotoStorageFile(name, mode, self)

    def _save(self, name, content):
        name = self._clean_name(name)
        headers = self.headers
        content_type = mimetypes.guess_type(name)[0] or Key.DefaultContentType

        if self.gzip and content_type in self.gzip_content_types:
            content = self._compress_content(content)
            headers.update({'Content-Encoding': 'gzip'})

        headers.update({
            'Content-Type': content_type,
        })

        content.name = name
        k = self.bucket.get_key(name)
        if not k:
            k = self.bucket.new_key(name)
        k.set_contents_from_file(content, headers=headers, policy=self.acl)
        return name

    def delete(self, name):
        name = self._clean_name(name)
        self.bucket.delete_key(name)

    def exists(self, name):
        name = self._clean_name(name)
        k = Key(self.bucket, name)
        return k.exists()

    def listdir(self, name):
        name = self._clean_name(name)
        return [l.name for l in self.bucket.list() if not len(name) or l.name[:len(name)] == name]

    def size(self, name):
        name = self._clean_name(name)
        return self.bucket.get_key(name).size

    def url(self, name):
        name = self._clean_name(name)
        return self.connection.generate_url(QUERYSTRING_EXPIRE, method='GET', \
                bucket=self.bucket.name, key=name, query_auth=QUERYSTRING_AUTH)

    def get_available_name(self, name):
        """ Overwrite existing file with the same name. """
        name = self._clean_name(name)
        return name

class S3BotoStorageFile(File):
    def __init__(self, name, mode, storage):
        self._storage = storage
        self.name = name
        self._mode = mode
        self.key = storage.bucket.get_key(name)
        self._is_dirty = False
        self.file = StringIO()

    @property
    def size(self):
        return self.key.size

    def read(self, *args, **kwargs):
        self.file = StringIO()
        self._is_dirty = False
        self.key.get_contents_to_file(self.file)
        return self.file.getvalue()

    def write(self, content):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self.file = StringIO(content)
        self._is_dirty = True

    def close(self):
        if self._is_dirty:
            self.key.set_contents_from_string(self.file.getvalue(), headers=self._storage.headers, acl=self._storage.acl)
        self.key.close()
<commit_msg>Use the new_key method as someone might have changed the key_class on the bucket.<commit_after>import os
import mimetypes

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.core.files.base import File
from django.core.files.storage import Storage
from django.core.exceptions import ImproperlyConfigured

try:
    from boto.s3.connection import S3Connection, S3ResponseError
    from boto.s3.key import Key
except ImportError:
    raise ImproperlyConfigured, "Could not load Boto's S3 bindings.\
    \nSee http://code.google.com/p/boto/"

ACCESS_KEY_NAME     = getattr(settings, 'AWS_ACCESS_KEY_ID', None)
SECRET_KEY_NAME     = getattr(settings, 'AWS_SECRET_ACCESS_KEY', None)
HEADERS             = getattr(settings, 'AWS_HEADERS', {})
STORAGE_BUCKET_NAME = getattr(settings, 'AWS_STORAGE_BUCKET_NAME', None)
AUTO_CREATE_BUCKET  = getattr(settings, 'AWS_AUTO_CREATE_BUCKET', True)
DEFAULT_ACL         = getattr(settings, 'AWS_DEFAULT_ACL', 'public-read')
QUERYSTRING_AUTH    = getattr(settings, 'AWS_QUERYSTRING_AUTH', True)
QUERYSTRING_EXPIRE  = getattr(settings, 'AWS_QUERYSTRING_EXPIRE', 3600)
LOCATION            = getattr(settings, 'AWS_LOCATION', '')
IS_GZIPPED          = getattr(settings, 'AWS_IS_GZIPPED', False)
GZIP_CONTENT_TYPES  = getattr(settings, 'GZIP_CONTENT_TYPES', (
    'text/css',
    'application/javascript',
    'application/x-javascript'
))

if IS_GZIPPED:
    from gzip import GzipFile

class S3BotoStorage(Storage):
    """Amazon Simple Storage Service using Boto"""

    def __init__(self, bucket=STORAGE_BUCKET_NAME, access_key=None,
                       secret_key=None, acl=DEFAULT_ACL, headers=HEADERS,
                       gzip=IS_GZIPPED, gzip_content_types=GZIP_CONTENT_TYPES):
        self.acl = acl
        self.headers = headers
        self.gzip = gzip
        self.gzip_content_types = gzip_content_types

        if not access_key and not secret_key:
             access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(access_key, secret_key)
        self.bucket = self._get_or_create_bucket(bucket)
        self.bucket.set_acl(self.acl)

    def _get_access_keys(self):
        access_key = ACCESS_KEY_NAME
        secret_key = SECRET_KEY_NAME
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_or_create_bucket(self, name):
        """Retrieves a bucket if it exists, otherwise creates it."""
        try:
            return self.connection.get_bucket(name)
        except S3ResponseError, e:
            if AUTO_CREATE_BUCKET:
                return self.connection.create_bucket(name)
            raise ImproperlyConfigured, ("Bucket specified by "
            "AWS_STORAGE_BUCKET_NAME does not exist. Buckets can be "
            "automatically created by setting AWS_AUTO_CREATE_BUCKET=True")

    def _clean_name(self, name):
        # Useful for windows' paths
        return os.path.normpath(name).replace('\\', '/')

    def _compress_content(self, content):
        """Gzip a given string."""
        zbuf = StringIO()
        zfile = GzipFile(mode='wb', compresslevel=6, fileobj=zbuf)
        zfile.write(content.read())
        zfile.close()
        content.file = zbuf
        return content

    def _open(self, name, mode='rb'):
        name = self._clean_name(name)
        return S3BotoStorageFile(name, mode, self)

    def _save(self, name, content):
        name = self._clean_name(name)
        headers = self.headers
        content_type = mimetypes.guess_type(name)[0] or Key.DefaultContentType

        if self.gzip and content_type in self.gzip_content_types:
            content = self._compress_content(content)
            headers.update({'Content-Encoding': 'gzip'})

        headers.update({
            'Content-Type': content_type,
        })

        content.name = name
        k = self.bucket.get_key(name)
        if not k:
            k = self.bucket.new_key(name)
        k.set_contents_from_file(content, headers=headers, policy=self.acl)
        return name

    def delete(self, name):
        name = self._clean_name(name)
        self.bucket.delete_key(name)

    def exists(self, name):
        name = self._clean_name(name)
        k = self.bucket.new_key(name)
        return k.exists()

    def listdir(self, name):
        dirlist = self._bucket.list(name)
        files = []
        dirs = set()
        base_parts = name.split("/") if name else []
        for item in dirlist:
            parts = item.name.split("/")
            parts = parts[len(base_parts):]
            if len(parts) == 1:
                # File
                files.append(parts[0])
            elif len(parts) > 1:
                # Directory
                dirs.add(parts[0])
        return list(dirs),files

    def size(self, name):
        name = self._clean_name(name)
        return self.bucket.get_key(name).size

    def url(self, name):
        name = self._clean_name(name)
        return self.connection.generate_url(QUERYSTRING_EXPIRE, method='GET', \
                bucket=self.bucket.name, key=name, query_auth=QUERYSTRING_AUTH)

    def get_available_name(self, name):
        """ Overwrite existing file with the same name. """
        name = self._clean_name(name)
        return name

class S3BotoStorageFile(File):
    def __init__(self, name, mode, storage):
        self._storage = storage
        self.name = name
        self._mode = mode
        self.key = storage.bucket.get_key(name)
        self._is_dirty = False
        self.file = StringIO()

    @property
    def size(self):
        return self.key.size

    def read(self, *args, **kwargs):
        self.file = StringIO()
        self._is_dirty = False
        self.key.get_contents_to_file(self.file)
        return self.file.getvalue()

    def write(self, content):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self.file = StringIO(content)
        self._is_dirty = True

    def close(self):
        if self._is_dirty:
            self.key.set_contents_from_string(self.file.getvalue(), headers=self._storage.headers, acl=self._storage.acl)
        self.key.close()
<|endoftext|>


<commit_before>'''
This module offers easy-to-understand methods that you can use to
interact with LEGO Mindstorms EV3 bricks.

Tested on:
- OS X 10.9.3 with Python 2.7.6

To do:
- all other EV3 commands! (only a couple are implemented here)
- USB capability
- WiFi capability
- error checking
- multiple layers support
- Linux support
- Windows support

Author: Thiago Marzagao (tmarzagao at gmail dot com)

See firmware source code (https://github.com/mindboards/ev3sources) for
info on the data types and opcodes used here.
'''

from dec_to_hex import h # decimal-to-hexadecimal dictionary

# globals
PRIMPAR_SHORT = 0x00
PRIMPAR_LONG = 0x80
PRIMPAR_CONST = 0x00
PRIMPAR_VALUE = 0x3F
PRIMPAR_GLOBAL = 0x20
PRIMPAR_INDEX = 0x1F
PRIMPAR_VARIABEL = 0x40
PRIMPAR_1_BYTE = 1
PRIMPAR_2_BYTES = 2
PRIMPAR_4_BYTES = 3

def LC0(v):
    '''
    format 1-byte local constant
    '''
    byte1 = ((v & PRIMPAR_VALUE) | PRIMPAR_SHORT | PRIMPAR_CONST)
    return h[byte1]

def LC1(v):
    '''
    format 2-byte local constant
    '''
    byte1 = (PRIMPAR_LONG | PRIMPAR_CONST | PRIMPAR_1_BYTE)
    byte2 = (v & 0xFF)
    return h[byte1] + h[byte2]

def LC2(v):
    '''
    format 3-byte local constant
    '''
    byte1 = (PRIMPAR_LONG | PRIMPAR_CONST | PRIMPAR_2_BYTES)
    byte2 = (v & 0xFF)
    byte3 = ((v >> 8) & 0xFF)
    return h[byte1] + h[byte2] + h[byte3]

def LC4(v):
    '''
    format 5-byte local constant
    '''
    byte1 = (PRIMPAR_LONG  | PRIMPAR_CONST | PRIMPAR_4_BYTES)
    byte2 = (v & 0xFF)
    byte3 = ((v >> 8) & 0xFF)
    byte4 = ((v >> 16) & 0xFF)
    byte5 = ((v >> 24) & 0xFF)
    return h[byte1] + h[byte2] + h[byte3] + h[byte4] + h[byte5]

def GV0(v):
    '''
    format 1-byte global variable
    '''
    return h[((v & PRIMPAR_INDEX) | PRIMPAR_SHORT | PRIMPAR_VARIABEL | PRIMPAR_GLOBAL)]

class ev3:

    '''
    this class lets you interact with EV3 bricks

    arguments used in several methods:

    ports: a, b, c, d, any combinations thereof
        type: str or iterable
        obs.: motor ports
        examples: 'a', 'bcd', ['d', 'c'], ['a'], ('c', 'b')
    port: 0, 1, 2, 3
        type: int
        obs.: sensor port
    power: -100...+100
        type: int
        obs.: -100 is full power backward, 100 is full power forward
    layer: 0, 1, 2, 3
        type: int
        obs.: use 0 unless you have multiple EV3 bricks
    '''

    def __init__(self):

        '''
        data used by several methods
        '''

        self.ports_to_int = {'a': 1, 'b': 2, 'c': 4, 'd': 8}
        self.stops = {'brake': 1, 'coast': 0}

    def connect(self, conn_type):

        '''
        connect to EV3

        conn: 'bt' (bluetooth), 'usb', 'wifi'
        '''

        conn_types = {'bt': '/dev/tty.EV3-SerialPort',
                      'usb': '', # must add usb support
                      'wifi': ''} # must add wifi support
        self.brick = open(conn_types[conn_type], mode = 'w+', buffering = 0)

    def start_motors(self, ports, power, layer=0):

        '''
        start motors at specified ports, power, and layer
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # opOUTPUT_POWER
        comm_1 = '\xA4' + LC0(layer) + LC0(ports) + LC1(power)

        # opOUTPUT_START
        comm_2 = '\xA6' + LC0(layer) + LC0(ports)

        # set message size, message counter, command type, vars
        msg_size = len(comm_1 + comm_2) + 5
        comm_0 = h[msg_size] + '\x00\x00\x00\x80\x00\x00'

        # assemble command and send to EV3
        command = comm_0 + comm_1 + comm_2
        self.brick.write(command)

    def stop_motors(self, ports, stop='coast', layer=0):

        '''
        stop motors at specified ports and layer

        stop: 'brake', 'coast'
            type: str
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # map mode: str->int
        stop = self.stops[stop]

        # set message size, message counter, command type, vars
        comm_0 = '\x09\x00\x01\x00\x80\x00\x00'

        # opOUTPUT_STOP
        comm_1 = '\xA3' + LC0(layer) + LC0(ports) + LC0(stop)

        # assemble command and send to EV3
        command = comm_0 + comm_1
        self.brick.write(command)

    def start_motors_degrees(self, ports, power, degrees, stop='brake',
                             ramp_up=0, ramp_down=0, layer=0):

        '''
        start motors at specified ports, power, and layer, and
        stop them after specified number of degrees (accurate to
        +/- 1 degree)

        degrees: 0...MAX
            type: int
            obs.: unit is degrees (360 degrees = 1 full turn)
        mode: 'brake', 'coast'
            type: str
        ramp_up: 0...MAX
            type: int
            obs.: makes acceleration constant; unit is degrees
        ramp_down: 0...MAX
            type: int
            obs.: makes deceleration constant; unit is degrees
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # map mode: str->int
        stop = self.stops[stop]

        # opOUTPUT_STEP_POWER
        comm_1 = '\xAC' + LC0(layer) + LC0(ports) + LC1(power) \
                 + LC4(ramp_up) + LC4(degrees) + LC4(ramp_down) \
                 + LC0(stop)

        # opOUTPUT_START
        comm_2 = '\xA6' + LC0(layer) + LC0(ports)

        msg_size = len(comm_1 + comm_2) + 5
        comm_0 = h[msg_size] + '\x00\x00\x00\x80\x00\x00'

        # assemble command and send to EV3
        command = comm_0 + comm_1 + comm_2
        self.brick.write(command)

    def start_motors_time(self, ports, power, time, stop='brake',
                             ramp_up=0, ramp_down=0, layer=0):

        '''
        start motors at specified ports, power, and layer, and
        stop them after specified number of milliseconds

        time: 0...MAX
            type: int
            obs.: unit is milliseconds
        mode: 'brake', 'coast'
            type: str
        ramp_up: 0...MAX
            type: int
            obs.: makes acceleration constant; unit is milliseconds
        ramp_down: 0...MAX
            type: int
            obs.: makes deceleration constant; unit is milliseconds
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # map mode: str->int
        stop = self.stops[stop]

        # opOUTPUT_STEP_POWER
        comm_1 = '\xAD' + LC0(layer) + LC0(ports) + LC1(power) \
                 + LC4(ramp_up) + LC4(time) + LC4(ramp_down) \
                 + LC0(stop)

        # opOUTPUT_START
        comm_2 = '\xA6' + LC0(layer) + LC0(ports)

        msg_size = len(comm_1 + comm_2) + 5
        comm_0 = h[msg_size] + '\x00\x00\x00\x80\x00\x00'

        # assemble command and send to EV3
        command = comm_0 + comm_1 + comm_2
        self.brick.write(command)

    def read_sensor(self, port, layer=0):

        '''
        read sensor from specified port (unit: percentage)
        '''

        # set message size, message counter, command type, vars
        comm_0 = '\x0B\x00\x00\x00\x00\x01\x00'

        # opINPUT_READ
        comm_1 = '\x9A' + LC0(layer) + LC0(port) + '\x00\x00\x60'

        # assemble command and send to EV3
        command = comm_0 + comm_1
        self.brick.write(command)

        # retrieve sensor value (5th byte) and convert to int
        sensor_data = self.brick.read(6)
        return int(hex(ord(sensor_data[5])), 16)

    def disconnect(self):
        self.brick.close()
<commit_msg>shortened method names<commit_after>'''
This module offers easy-to-understand methods that you can use to
interact with LEGO Mindstorms EV3 bricks.

Tested on:
- OS X 10.9.3 with Python 2.7.6

To do:
- all other EV3 commands! (only a couple are implemented here)
- USB capability
- WiFi capability
- error checking
- multiple layers support
- Linux support
- Windows support

Author: Thiago Marzagao (tmarzagao at gmail dot com)

See firmware source code (https://github.com/mindboards/ev3sources) for
info on the data types and opcodes used here.
'''

from dec_to_hex import h # decimal-to-hexadecimal dictionary

# globals
PRIMPAR_SHORT = 0x00
PRIMPAR_LONG = 0x80
PRIMPAR_CONST = 0x00
PRIMPAR_VALUE = 0x3F
PRIMPAR_GLOBAL = 0x20
PRIMPAR_INDEX = 0x1F
PRIMPAR_VARIABEL = 0x40
PRIMPAR_1_BYTE = 1
PRIMPAR_2_BYTES = 2
PRIMPAR_4_BYTES = 3

def LC0(v):
    '''
    format 1-byte local constant
    '''
    byte1 = ((v & PRIMPAR_VALUE) | PRIMPAR_SHORT | PRIMPAR_CONST)
    return h[byte1]

def LC1(v):
    '''
    format 2-byte local constant
    '''
    byte1 = (PRIMPAR_LONG | PRIMPAR_CONST | PRIMPAR_1_BYTE)
    byte2 = (v & 0xFF)
    return h[byte1] + h[byte2]

def LC2(v):
    '''
    format 3-byte local constant
    '''
    byte1 = (PRIMPAR_LONG | PRIMPAR_CONST | PRIMPAR_2_BYTES)
    byte2 = (v & 0xFF)
    byte3 = ((v >> 8) & 0xFF)
    return h[byte1] + h[byte2] + h[byte3]

def LC4(v):
    '''
    format 5-byte local constant
    '''
    byte1 = (PRIMPAR_LONG  | PRIMPAR_CONST | PRIMPAR_4_BYTES)
    byte2 = (v & 0xFF)
    byte3 = ((v >> 8) & 0xFF)
    byte4 = ((v >> 16) & 0xFF)
    byte5 = ((v >> 24) & 0xFF)
    return h[byte1] + h[byte2] + h[byte3] + h[byte4] + h[byte5]

def GV0(v):
    '''
    format 1-byte global variable
    '''
    return h[((v & PRIMPAR_INDEX) | PRIMPAR_SHORT | PRIMPAR_VARIABEL | PRIMPAR_GLOBAL)]

class ev3:

    '''
    this class lets you interact with EV3 bricks

    arguments used in several methods:

    ports: a, b, c, d, any combinations thereof
        type: str or iterable
        obs.: motor ports
        examples: 'a', 'bcd', ['d', 'c'], ['a'], ('c', 'b')
    port: 0, 1, 2, 3
        type: int
        obs.: sensor port
    power: -100...+100
        type: int
        obs.: -100 is full power backward, 100 is full power forward
    layer: 0, 1, 2, 3
        type: int
        obs.: use 0 unless you have multiple EV3 bricks
    '''

    def __init__(self):

        '''
        data used by several methods
        '''

        self.ports_to_int = {'a': 1, 'b': 2, 'c': 4, 'd': 8}
        self.stops = {'brake': 1, 'coast': 0}

    def connect(self, conn_type):

        '''
        connect to EV3

        conn: 'bt' (bluetooth), 'usb', 'wifi'
        '''

        conn_types = {'bt': '/dev/tty.EV3-SerialPort',
                      'usb': '', # must add usb support
                      'wifi': ''} # must add wifi support
        self.brick = open(conn_types[conn_type], mode = 'w+', buffering = 0)

    def motor_start(self, ports, power, layer=0):

        '''
        start motors at specified ports, power, and layer
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # opOUTPUT_POWER
        comm_1 = '\xA4' + LC0(layer) + LC0(ports) + LC1(power)

        # opOUTPUT_START
        comm_2 = '\xA6' + LC0(layer) + LC0(ports)

        # set message size, message counter, command type, vars
        msg_size = len(comm_1 + comm_2) + 5
        comm_0 = h[msg_size] + '\x00\x00\x00\x80\x00\x00'

        # assemble command and send to EV3
        command = comm_0 + comm_1 + comm_2
        self.brick.write(command)

    def motor_stop(self, ports, stop='coast', layer=0):

        '''
        stop motors at specified ports and layer

        stop: 'brake', 'coast'
            type: str
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # map mode: str->int
        stop = self.stops[stop]

        # set message size, message counter, command type, vars
        comm_0 = '\x09\x00\x01\x00\x80\x00\x00'

        # opOUTPUT_STOP
        comm_1 = '\xA3' + LC0(layer) + LC0(ports) + LC0(stop)

        # assemble command and send to EV3
        command = comm_0 + comm_1
        self.brick.write(command)

    def motor_degrees(self, ports, power, degrees, stop='brake',
                      ramp_up=0, ramp_down=0, layer=0):

        '''
        start motors at specified ports, power, and layer, and
        stop them after specified number of degrees (accurate to
        +/- 1 degree)

        degrees: 0...MAX
            type: int
            obs.: unit is degrees (360 degrees = 1 full turn)
        mode: 'brake', 'coast'
            type: str
        ramp_up: 0...MAX
            type: int
            obs.: makes acceleration constant; unit is degrees
        ramp_down: 0...MAX
            type: int
            obs.: makes deceleration constant; unit is degrees
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # map mode: str->int
        stop = self.stops[stop]

        # opOUTPUT_STEP_POWER
        comm_1 = '\xAC' + LC0(layer) + LC0(ports) + LC1(power) \
                 + LC4(ramp_up) + LC4(degrees) + LC4(ramp_down) \
                 + LC0(stop)

        # opOUTPUT_START
        comm_2 = '\xA6' + LC0(layer) + LC0(ports)

        msg_size = len(comm_1 + comm_2) + 5
        comm_0 = h[msg_size] + '\x00\x00\x00\x80\x00\x00'

        # assemble command and send to EV3
        command = comm_0 + comm_1 + comm_2
        self.brick.write(command)

    def motor_time(self, ports, power, time, stop='brake',
                   ramp_up=0, ramp_down=0, layer=0):

        '''
        start motors at specified ports, power, and layer, and
        stop them after specified number of milliseconds

        time: 0...MAX
            type: int
            obs.: unit is milliseconds
        mode: 'brake', 'coast'
            type: str
        ramp_up: 0...MAX
            type: int
            obs.: makes acceleration constant; unit is milliseconds
        ramp_down: 0...MAX
            type: int
            obs.: makes deceleration constant; unit is milliseconds
        '''

        # map ports: str->int
        ports = sum([self.ports_to_int[port] for port in ports])

        # map mode: str->int
        stop = self.stops[stop]

        # opOUTPUT_STEP_POWER
        comm_1 = '\xAD' + LC0(layer) + LC0(ports) + LC1(power) \
                 + LC4(ramp_up) + LC4(time) + LC4(ramp_down) \
                 + LC0(stop)

        # opOUTPUT_START
        comm_2 = '\xA6' + LC0(layer) + LC0(ports)

        msg_size = len(comm_1 + comm_2) + 5
        comm_0 = h[msg_size] + '\x00\x00\x00\x80\x00\x00'

        # assemble command and send to EV3
        command = comm_0 + comm_1 + comm_2
        self.brick.write(command)

    def read_sensor(self, port, layer=0):

        '''
        read sensor from specified port (unit: percentage)
        '''

        # set message size, message counter, command type, vars
        comm_0 = '\x0B\x00\x00\x00\x00\x01\x00'

        # opINPUT_READ
        comm_1 = '\x9A' + LC0(layer) + LC0(port) + '\x00\x00\x60'

        # assemble command and send to EV3
        command = comm_0 + comm_1
        self.brick.write(command)

        # retrieve sensor value (5th byte) and convert to int
        sensor_data = self.brick.read(6)
        return int(hex(ord(sensor_data[5])), 16)

    def disconnect(self):
        self.brick.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# 6/27/2019
# Author: Kevin McJunkin
# Use away if this is somehow relevant to ya

import os
import subprocess
import shutil

# Install required packages via yum
def install_packages():
    package_list = ['python-pip',
                    'gcc',
                    'python-devel',
                    'libyaml-devel',
                    'openssl-devel',
                    'libffi-devel',
                    'libxml2-devel',
                    'libxslt-devel',
                    'puppet']
    print('Installing packages')
    try:
        for package in package_list:
            print('Installing ' + package)
            os.system('yum install -y ' + package + '>/dev/null')
            print('Successful\n')
    except:
        print('Unsuccessful')

# Install required packages via pip
def pip_down():
    print('\nInstalling OpenStack HEAT requirements via pip')
    os_list = ['setuptools==36.2.7',
               '-U pip==20.1.1',
               '--ignore-installed os-apply-config==11.2.0',
               '--ignore-installed os-collect-config==5.0.0',
               'os-refresh-config==10.4.0',
               'dib-utils==0.0.11',
               '-U decorator',
               'ansible==2.4.3.0']
    try:
        for package in os_list:
            print('Installing ' + package)
            os.system('pip install ' + package + '>/dev/null')
            print('Successful')
    except:
        print('Unsuccessful')

# Remove git repo if it exist (should never come up but might as well)
# Clone git repo that has all our configuration files
def download_git():
    endpoint = 'https://github.com/kmcjunk/hotstrapper/archive/master.zip'
    os.system('wget {}'.format(endpoint))
    os.system('unzip master.zip')

# Move configuration files to the proper location on the OS
# ...and use a really ghetto create directory for the move
# chmod files properly

def configurate():
    file_list = ['opt/stack/os-config-refresh/configure.d/20-os-apply-config',
                 'opt/stack/os-config-refresh/configure.d/55-heat-config',
                 'usr/bin/heat-config-notify',
                 'var/lib/heat-config/hooks/ansible',
                 'var/lib/heat-config/hooks/script',
                 'var/lib/heat-config/hooks/puppet',
                 'etc/os-collect-config.conf',
                 'usr/libexec/os-apply-config/templates/var/run/heat-config/heat-config',
                 'usr/libexec/os-apply-config/templates/etc/os-collect-config.conf']
    print('Moving configuration files to the proper locations\n\n')
    for file in file_list:
        directory = os.path.dirname('/' + file)
        if not os.path.exists(directory):
            os.makedirs(directory)
        print('hotstrapper-master/rhel/7/' + file + '\t->\t' + '/' + file)
        shutil.move('hotstrapper-master/rhel/7/' + file, '/' + file)
    for i in range(3):
        os.chmod('/' + file_list[i], 0700)
    for i in range(3, 6):
        os.chmod('/' + file_list[i], 0755)
        # os.chmod('/' + file, 0700)

# Run os-collect to propagate the config & run it again
# Then run start_config to create/enable the os-collect service
# Also clean up the git repo cause it is dead to us
def jiggle_some_things():
    print('\nRunning os-collect-config & ensuring os-collect-config-exist')
    os.system('os-collect-config --one-time --debug')
    os.system('cat /etc/os-collect-config.conf')
    os.system('os-collect-config --one-time --debug')
    print('\nEnsuring everything is running & enabled on boot')
    subprocess.call('hotstrapper-master/rhel/7/start_config_agent.sh')
    print('\nCleaning up git folder')
    shutil.rmtree('hotstrapper-master/')
    os.system('rm -f master.zip')

# Ensure we don't get rekt by cloud-init next boot
def delete_some_other_things():
    print('Ensuring no cloud-init references exist')
    os.system('rm -rf /var/lib/cloud/instance')
    os.system('rm -rf /var/lib/cloud/instances/*')
    os.system('rm -rf /var/lib/cloud/data/*')
    os.system('rm -rf /var/lib/cloud/sem/config_scripts_per_once.once')
    os.system('rm -rf /var/log/cloud-init.log')
    os.system('rm -rf /var/log/cloud-init-output.log')
    print('\n\n\nDone!')

install_packages()
pip_down()
download_git()
configurate()
jiggle_some_things()
delete_some_other_things()
<commit_msg>Changing name schema<commit_after>#!/usr/bin/env python
# 6/27/2019
# Author: Kevin McJunkin
# Use away if this is somehow relevant to ya

import os
import subprocess
import shutil

# Install required packages via yum
def install_packages():
    package_list = ['python-pip',
                    'gcc',
                    'python-devel',
                    'libyaml-devel',
                    'openssl-devel',
                    'libffi-devel',
                    'libxml2-devel',
                    'libxslt-devel',
                    'puppet']
    print('Installing packages')
    try:
        for package in package_list:
            print('Installing ' + package)
            os.system('yum install -y ' + package + '>/dev/null')
            print('Successful\n')
    except:
        print('Unsuccessful')

# Install required packages via pip
def pip_down():
    print('\nInstalling OpenStack HEAT requirements via pip')
    os_list = ['setuptools==36.2.7',
               '-U pip==20.1.1',
               '--ignore-installed os-apply-config==11.2.0',
               '--ignore-installed os-collect-config==5.0.0',
               'os-refresh-config==10.4.0',
               'dib-utils==0.0.11',
               '-U decorator',
               'ansible==2.4.3.0']
    try:
        for package in os_list:
            print('Installing ' + package)
            os.system('pip install ' + package + '>/dev/null')
            print('Successful')
    except:
        print('Unsuccessful')

# Remove git repo if it exist (should never come up but might as well)
# Clone git repo that has all our configuration files
def download_git():
    endpoint = 'https://github.com/kmcjunk/hotstrapper/archive/master.zip'
    os.system('wget {}'.format(endpoint))
    os.system('unzip master.zip')

# Move configuration files to the proper location on the OS
# ...and use a really ghetto create directory for the move
# chmod files properly
def configurate():
    file_list = ['opt/stack/os-config-refresh/configure.d/20-os-apply-config',
                 'opt/stack/os-config-refresh/configure.d/55-heat-config',
                 'usr/bin/heat-config-notify',
                 'var/lib/heat-config/hooks/ansible',
                 'var/lib/heat-config/hooks/script',
                 'var/lib/heat-config/hooks/puppet',
                 'etc/os-collect-config.conf',
                 'usr/libexec/os-apply-config/templates/var/run/heat-config/heat-config',
                 'usr/libexec/os-apply-config/templates/etc/os-collect-config.conf']
    print('Moving configuration files to the proper locations\n\n')
    for file in file_list:
        directory = os.path.dirname('/' + file)
        if not os.path.exists(directory):
            os.makedirs(directory)
        print('hotstrapper-master/bootstrap/centos/7/' + file + '\t->\t' + '/' + file)
        shutil.move('hotstrapper-master/bootstrap/centos/7/' + file, '/' + file)
    for i in range(3):
        os.chmod('/' + file_list[i], 0700)
    for i in range(3, 6):
        os.chmod('/' + file_list[i], 0755)
        # os.chmod('/' + file, 0700)

# Run os-collect to propagate the config & run it again
# Then run start_config to create/enable the os-collect service
# Also clean up the git repo cause it is dead to us
def jiggle_some_things():
    print('\nRunning os-collect-config & ensuring os-collect-config-exist')
    os.system('os-collect-config --one-time --debug')
    os.system('cat /etc/os-collect-config.conf')
    os.system('os-collect-config --one-time --debug')
    print('\nEnsuring everything is running & enabled on boot')
    subprocess.call('hotstrapper-master/bootstrap/centos/7/start_config_agent.sh')
    print('\nCleaning up git folder')
    shutil.rmtree('hotstrapper-master/')
    os.system('rm -f master.zip')

# Ensure we don't get rekt by cloud-init next boot
def delete_some_other_things():
    print('Ensuring no cloud-init references exist')
    os.system('rm -rf /var/lib/cloud/instance')
    os.system('rm -rf /var/lib/cloud/instances/*')
    os.system('rm -rf /var/lib/cloud/data/*')
    os.system('rm -rf /var/lib/cloud/sem/config_scripts_per_once.once')
    os.system('rm -rf /var/log/cloud-init.log')
    os.system('rm -rf /var/log/cloud-init-output.log')
    os.system('rm -rf ')
    print('\n\n\nDone!')

install_packages()
pip_down()
download_git()
configurate()
jiggle_some_things()
delete_some_other_things()
exit(0)
<|endoftext|>


<commit_before># cerbero - a multi-platform build system for Open Source software
# Copyright (C) 2012 Andoni Morales Alastruey <ylatuya@gmail.com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Library General Public
# License as published by the Free Software Foundation; either
# version 2 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Library General Public License for more details.
#
# You should have received a copy of the GNU Library General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., 59 Temple Place - Suite 330,
# Boston, MA 02111-1307, USA.

import itertools

from cerbero.build.filesprovider import FilesProvider

class PackageBase(object):
    '''
    Base class for packages with the common field to describe a package

    @cvar name: name of the package
    @type name: str
    @cvar shortdesc: Short description of the package
    @type shortdesc: str
    @cvar longdesc: Long description of the package
    @type longdesc: str
    @cvar version: version of the package
    @type version: str
    @cvar uuid: unique id for this package
    @type uuid: str
    @cvar license:  package license
    @type license: str
    @cvar vendor: vendor for this package
    @type vendor: str
    @cvar org: organization for this package (eg: net.foo.bar)
    @type org: str
    @cvar url: url for this pacakge
    @type url: str
    '''
    name = 'default'
    shortdesc = 'default'
    longdesc = 'default'
    version = 'default'
    org = 'default'
    uuid = None
    licenses = 'GPL'
    vendor = 'default'
    url = 'default'

    def __init__(self, config):
        self.config = config

    def files_list(self):
        raise NotImplemented("'files_list' must be implemented by subclasses")

    def devel_files_list(self):
        raise NotImplemented("'devel_files_list' must be implemented by "
                             "subclasses")

    def all_files_list(self):
        raise NotImplemented("'all_files_list' must be implemented by "
                             "subclasses")

    def get_install_dir(self):
        try:
            return self.install_dir[self.config.target_platform]
        except:
            return self.config.install_dir

class Package(PackageBase):
    '''
    Describes a set of files to produce disctribution packages for the
    different target platforms

    @cvar deps: list of packages dependencies
    @type deps: list
    @cvar files: list of files included in this package
    @type files: list
    @cvar platform_files: list of platform files included in this package
    @type platform_files: list
    '''

    deps = list()
    files = list()
    platform_files = dict()

    def __init__(self, config, cookbook):
        PackageBase.__init__(self, config)
        self.cookbook = cookbook
        self._files = self.files + \
                self.platform_files.get(config.target_platform, [])
        self._parse_files()

    def recipes_dependencies(self):
        return [x.split(':')[0] for x in self._files]

    def files_list(self):
        files = []
        for recipe_name, categories in self._recipes_files.iteritems():
            recipe = self.cookbook.get_recipe(recipe_name)
            if len(categories) == 0:
                rfiles = recipe.dist_files_list()
            else:
                rfiles = recipe.files_list_by_categories(categories)
            files.extend(rfiles)
        return sorted(files)

    def devel_files_list(self):
        files = []
        for recipe, categories in self._recipes_files.iteritems():
            # only add development files for recipe from which used the 'libs'
            # category
            if len(categories) == 0 or FilesProvider.LIBS_CAT in categories:
                rfiles = self.cookbook.get_recipe(recipe).devel_files_list()
                files.extend(rfiles)
        return sorted(files)

    def all_files_list(self):
        files = self.files_list()
        files.extend(self.devel_files_list())
        return sorted(files)

    def _parse_files(self):
        self._recipes_files = {}
        for r in self._files:
            l = r.split(':')
            self._recipes_files[l[0]] = l[1:]

class MetaPackage(PackageBase):
    '''
    Group of packages used to build an installer package

    @cvar packages: list of packages grouped in this meta package
    @type packages: list
    @cvar icon: filename of the package icon
    @type icon: str
    '''

    icon = None
    packages = []

    def __init__(self, config, store):
        PackageBase.__init__(self, config)
        self.store = store

    def list_packages(self):
        return [p[0] for p in self.packages]

    def files_list(self):
        return self._list_files(Package.files_list)

    def devel_files_list(self):
        return self._list_files(Package.devel_files_list)

    def all_files_list(self):
        return self._list_files(Package.all_files_list)

    def _list_files(self, func):
        # for each package, call the function that list files
        files = []
        for name, _, _ in self.packages:
            files.extend(func(self.store.get_package(name)))
        return sorted(files)
<commit_msg>Include also dependencies in the metapackage files list<commit_after># cerbero - a multi-platform build system for Open Source software
# Copyright (C) 2012 Andoni Morales Alastruey <ylatuya@gmail.com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Library General Public
# License as published by the Free Software Foundation; either
# version 2 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Library General Public License for more details.
#
# You should have received a copy of the GNU Library General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., 59 Temple Place - Suite 330,
# Boston, MA 02111-1307, USA.

import itertools

from cerbero.build.filesprovider import FilesProvider

class PackageBase(object):
    '''
    Base class for packages with the common field to describe a package

    @cvar name: name of the package
    @type name: str
    @cvar shortdesc: Short description of the package
    @type shortdesc: str
    @cvar longdesc: Long description of the package
    @type longdesc: str
    @cvar version: version of the package
    @type version: str
    @cvar uuid: unique id for this package
    @type uuid: str
    @cvar license:  package license
    @type license: str
    @cvar vendor: vendor for this package
    @type vendor: str
    @cvar org: organization for this package (eg: net.foo.bar)
    @type org: str
    @cvar url: url for this pacakge
    @type url: str
    '''
    name = 'default'
    shortdesc = 'default'
    longdesc = 'default'
    version = 'default'
    org = 'default'
    uuid = None
    licenses = 'GPL'
    vendor = 'default'
    url = 'default'

    def __init__(self, config):
        self.config = config

    def files_list(self):
        raise NotImplemented("'files_list' must be implemented by subclasses")

    def devel_files_list(self):
        raise NotImplemented("'devel_files_list' must be implemented by "
                             "subclasses")

    def all_files_list(self):
        raise NotImplemented("'all_files_list' must be implemented by "
                             "subclasses")

    def get_install_dir(self):
        try:
            return self.install_dir[self.config.target_platform]
        except:
            return self.config.install_dir

class Package(PackageBase):
    '''
    Describes a set of files to produce disctribution packages for the
    different target platforms

    @cvar deps: list of packages dependencies
    @type deps: list
    @cvar files: list of files included in this package
    @type files: list
    @cvar platform_files: list of platform files included in this package
    @type platform_files: list
    '''

    deps = list()
    files = list()
    platform_files = dict()

    def __init__(self, config, cookbook):
        PackageBase.__init__(self, config)
        self.cookbook = cookbook
        self._files = self.files + \
                self.platform_files.get(config.target_platform, [])
        self._parse_files()

    def recipes_dependencies(self):
        return [x.split(':')[0] for x in self._files]

    def files_list(self):
        files = []
        for recipe_name, categories in self._recipes_files.iteritems():
            recipe = self.cookbook.get_recipe(recipe_name)
            if len(categories) == 0:
                rfiles = recipe.dist_files_list()
            else:
                rfiles = recipe.files_list_by_categories(categories)
            files.extend(rfiles)
        return sorted(files)

    def devel_files_list(self):
        files = []
        for recipe, categories in self._recipes_files.iteritems():
            # only add development files for recipe from which used the 'libs'
            # category
            if len(categories) == 0 or FilesProvider.LIBS_CAT in categories:
                rfiles = self.cookbook.get_recipe(recipe).devel_files_list()
                files.extend(rfiles)
        return sorted(files)

    def all_files_list(self):
        files = self.files_list()
        files.extend(self.devel_files_list())
        return sorted(files)

    def _parse_files(self):
        self._recipes_files = {}
        for r in self._files:
            l = r.split(':')
            self._recipes_files[l[0]] = l[1:]

class MetaPackage(PackageBase):
    '''
    Group of packages used to build an installer package

    @cvar packages: list of packages grouped in this meta package
    @type packages: list
    @cvar icon: filename of the package icon
    @type icon: str
    '''

    icon = None
    packages = []

    def __init__(self, config, store):
        PackageBase.__init__(self, config)
        self.store = store

    def list_packages(self):
        return [p[0] for p in self.packages]

    def files_list(self):
        return self._list_files(Package.files_list)

    def devel_files_list(self):
        return self._list_files(Package.devel_files_list)

    def all_files_list(self):
        return self._list_files(Package.all_files_list)

    def _list_files(self, func):
        # for each package, call the function that list files
        files = []
        for name in self.store.get_package_deps(self.name):
            files.extend(func(self.store.get_package(name)))
        return sorted(files)
<|endoftext|>


<commit_before>from django.contrib.auth.models import User
from django.core.management.base import BaseCommand, CommandError
from django.utils import timezone

from bulk_update.helper import bulk_update

from pontoon.administration.files import update_from_repository
from pontoon.administration.vcs import commit_to_vcs
from pontoon.base.models import Entity, Locale, Project, Resource, Translation
from pontoon.base.utils import match_attr
from pontoon.base.vcs_models import VCSProject

class Command(BaseCommand):
    args = '<project_slug project_slug ...>'
    help = 'Synchronize database and remote repositories.'

    def log(self, msg, *args, **kwargs):
        """Log a message to the console."""
        self.stdout.write(msg.format(*args, **kwargs))

    def info(self, msg, *args, **kwargs):
        """Log a message to the console if --verbosity=1 or more."""
        if self.verbosity >= 1:
            self.log(msg, *args, **kwargs)

    def debug(self, msg, *args, **kwargs):
        """Log a message to the console if --verbosity=2."""
        if self.verbosity == 2:
            self.log(msg, *args, **kwargs)

    def handle(self, *args, **options):
        self.verbosity = options['verbosity']

        self.log('SYNC PROJECTS: start')
        projects = Project.objects.filter(disabled=False)
        if args:
            projects = projects.filter(slug__in=args)

        if len(projects) < 1:
            raise CommandError('No matching projects found.')

        for project in projects:
            self.handle_project(project)
        self.log('SYNC PROJECTS: done')

        # Once we've synced, we can delete all translations scheduled
        # for deletion.
        Translation.deleted_objects.all().delete()

    def handle_project(self, db_project):
        if not db_project.can_commit:
            self.log('Skipping project {0}, cannot commit to repository.'
                     .format(db_project.name))
            return

        # Pull changes from VCS and update what we know about the files.
        update_from_repository(db_project)
        vcs_project = VCSProject(db_project)
        self.update_resources(db_project, vcs_project)

        # Collect all entities across VCS and the database and get their
        # keys so we can match up matching entities.
        vcs_entities = self.get_vcs_entities(vcs_project)
        db_entities = self.get_db_entities(db_project)
        entity_keys = set().union(db_entities.keys(), vcs_entities.keys())

        changeset = ChangeSet(db_project, vcs_project)
        for key in entity_keys:
            db_entity = db_entities.get(key, None)
            vcs_entity = vcs_entities.get(key, None)
            self.handle_entity(changeset, db_project, key, db_entity, vcs_entity)

        # Apply the changeset to the files and then commit them.
        changeset.execute()
        self.commit_changes(db_project)

        # Once we've successfully committed, update the last_synced date.
        db_project.last_synced = timezone.now()
        db_project.save()

        self.log('Synced project {0}', db_project.slug)

    def handle_entity(self, changeset, db_project, key, db_entity, vcs_entity):
        """
        Determine what needs to be synced between the database and VCS versions
        of a single entity and log what needs to be changed in the changeset.
        """
        if vcs_entity is None:
            if db_entity is None:
                # This should never happen. What? Hard abort.
                raise CommandError('No entities found for key {0}'.format(key))
            else:
                # VCS no longer has the entity, remove it from Pontoon.
                changeset.obsolete_db_entity(db_entity)
        elif db_entity is None:
            # New VCS entities are added to Pontoon.
            changeset.create_db_entity(vcs_entity)
        else:
            for locale in db_project.locales.all():
                if db_entity.has_changed(locale.code):
                    # Pontoon changes overwrite whatever VCS has.
                    changeset.update_vcs_entity(locale.code, db_entity, vcs_entity)
                elif vcs_entity.has_translation_for(locale.code):
                    # If Pontoon has nothing or has not changed, and the VCS
                    # still has the entity, update Pontoon with whatever may
                    # have changed.
                    changeset.update_db_entity(locale.code, db_entity, vcs_entity)

    def update_resources(self, db_project, vcs_project):
        """Update the database on what resource files exist in VCS."""
        relative_paths = vcs_project.resources.keys()
        db_project.resource_set.exclude(path__in=relative_paths).delete()

        for relative_path, vcs_resource in vcs_project.resources.items():
            resource, created = db_project.resource_set.get_or_create(path=relative_path)
            resource.format = Resource.get_path_format(relative_path)
            resource.entity_count = len(vcs_resource.entities)
            resource.save()

    def get_vcs_entities(self, vcs_project):
        return {self.entity_key(entity): entity for entity in vcs_project.entities}

    def get_db_entities(self, db_project):
        entities = (Entity.objects
                    .select_related('resource')
                    .filter(resource__project=db_project, obsolete=False))
        return {self.entity_key(entity): entity for entity in entities}

    def entity_key(self, entity):
        """
        Generate a key for the given entity that is unique within the
        project.
        """
        return ':'.join([entity.resource.path, entity.string])

    def commit_changes(self, db_project):
        """Commit the changes we've made back to the VCS."""
        # TODO: Set translation author properly.
        user = User.objects.filter(is_superuser=True).first()

        for locale in db_project.locales.all():
            commit_message = (
                'Pontoon: Updated {locale.name} ({locale.code}) localization of '
                '{project.name}'
            ).format(locale=locale, project=db_project)

            try:
                result = commit_to_vcs(
                    db_project.repository_type,
                    db_project.locale_directory_path(locale.code),
                    commit_message,
                    user,
                    db_project.repository_url
                )
            except Exception as err:
                result = {'message': unicode(err)}

            if result is not None:
                self.log(
                    'Committing project {project.name} for {locale.name} '
                    '({locale.code}) failed: {reason}',
                    project=db_project,
                    locale=locale,
                    reason=result['message']
                )

class ChangeSet(object):
    """
    Stores a set of changes to be made to the database and the
    translations stored in VCS. Once all the necessary changes have been
    stored, execute all the changes at once efficiently.
    """
    def __init__(self, db_project, vcs_project):
        self.db_project = db_project
        self.vcs_project = vcs_project

        self.executed = False
        self.changes = {
            'update_vcs': [],
            'update_db': [],
            'obsolete_db': [],
            'create_db': []
        }

        self.entities_to_update = []
        self.translations_to_update = []
        self.translations_to_create = []

    def update_vcs_entity(self, locale_code, db_entity, vcs_entity):
        """
        Replace the translations in VCS with the translations from the
        database.
        """
        self.changes['update_vcs'].append((locale_code, db_entity, vcs_entity))

    def create_db_entity(self, vcs_entity):
        """Create a new entity in the database."""
        self.changes['create_db'].append(vcs_entity)

    def update_db_entity(self, locale_code, db_entity, vcs_entity):
        """Update the database with translations from VCS."""
        self.changes['update_db'].append((locale_code, db_entity, vcs_entity))

    def obsolete_db_entity(self, db_entity):
        """Mark the given entity as obsolete."""
        self.changes['obsolete_db'].append(db_entity.pk)

    def execute(self):
        """
        Execute the changes stored in this changeset. Execute can only
        be called once per changeset; subsequent calls raise a
        RuntimeError, even if the changes failed.
        """
        if self.executed:
            raise RuntimeError('execute() can only be called once per changeset.')
        else:
            self.executed = True

        # Store locales and resources for FK relationships.
        self.locales = {l.code: l for l in Locale.objects.all()}
        self.resources = {r.path: r for r in self.db_project.resource_set.all()}

        # Perform the changes and fill the lists for bulk creation and
        # updating.
        self.execute_update_vcs()
        self.execute_create_db()
        self.execute_update_db()
        self.execute_obsolete_db()

        # Apply the built-up changes to the DB
        if len(self.entities_to_update) > 0:
            bulk_update(self.entities_to_update, update_fields=[
                'resource',
                'string',
                'string_plural',
                'key',
                'comment',
                'order',
                'source'
            ])

        Translation.objects.bulk_create(self.translations_to_create)
        if len(self.translations_to_update) > 0:
            bulk_update(self.translations_to_update, update_fields=[
                'entity',
                'locale',
                'string',
                'plural_form',
                'approved',
                'approved_user',
                'approved_date',
                'fuzzy',
                'extra'
            ])

    def execute_update_vcs(self):
        resources = self.vcs_project.resources
        changed_resources = set()

        for locale_code, db_entity, vcs_entity in self.changes['update_vcs']:
            changed_resources.add(resources[db_entity.resource.path])

            vcs_translation = vcs_entity.translations[locale_code]
            db_translations = (db_entity.translation_set
                               .filter(approved=True, locale__code=locale_code))

            # If no DB translations are fuzzy, set fuzzy to False.
            # Otherwise, it's true.
            vcs_translation.fuzzy = any(t for t in db_translations if t.fuzzy)

            # Replace existing translations with ones from the database.
            vcs_translation.strings = {
                db.plural_form: db.string for db in db_translations
            }

        for resource in changed_resources:
            resource.save()

    def get_entity_updates(self, vcs_entity):
        """
        Return a dict of the properties and values necessary to create
        or update a database entity from a VCS entity.
        """
        return {
            'resource': self.resources[vcs_entity.resource.path],
            'string': vcs_entity.string,
            'string_plural': '',  # TODO: Support plural source.
            'key': '',  # TODO: Support keys separate from source.
            'comment': '\n'.join(vcs_entity.comments),
            'order': vcs_entity.order,
            'source': ''  # TODO: Support source
        }

    def execute_create_db(self):
        for vcs_entity in self.changes['create_db']:
            entity = Entity(**self.get_entity_updates(vcs_entity))
            entity.save()  # We can't use bulk_create since we need a PK

            for locale_code, vcs_translation in vcs_entity.translations.items():
                for plural_form, string in vcs_translation.strings.items():
                    self.translations_to_create.append(Translation(
                        entity=entity,
                        locale=self.locales[locale_code],
                        string=string,
                        plural_form=plural_form,
                        approved=True,
                        approved_date=timezone.now(),
                        fuzzy=vcs_translation.fuzzy
                    ))

    def execute_update_db(self):
        # TODO: Optimize to only update translations if necessary.
        for locale_code, db_entity, vcs_entity in self.changes['update_db']:
            for field, value in self.get_entity_updates(vcs_entity).items():
                setattr(db_entity, field, value)
            self.entities_to_update.append(db_entity)

            # Update translations for the entity.
            vcs_translation = vcs_entity.translations[locale_code]
            db_translations = db_entity.translation_set.filter(locale__code=locale_code)
            approved_translations = []

            for plural_form, string in vcs_translation.strings.items():
                # Check if we need to modify an existing translation or
                # create a new one.
                db_translation = match_attr(db_translations,
                                            plural_form=plural_form,
                                            string=string)
                if db_translation:
                    if not db_translation.approved:
                        db_translation.approved = True
                        db_translation.approved_date = timezone.now()
                    db_translation.fuzzy = vcs_translation.fuzzy
                    db_translation.extra = vcs_translation.extra

                    self.translations_to_update.append(db_translation)
                    approved_translations.append(db_translation)
                else:
                    self.translations_to_create.append(Translation(
                        entity=db_entity,
                        locale=self.locales[locale_code],
                        string=string,
                        plural_form=plural_form,
                        approved=True,
                        approved_date=timezone.now(),
                        fuzzy=vcs_translation.fuzzy,
                        extra=vcs_translation.extra
                    ))

            # Any existing translations that were not approved get unapproved.
            for translation in db_translations:
                if translation not in approved_translations:
                    translation.approved = False
                    translation.approved_user = None
                    translation.approved_date = None
                    self.translations_to_update.append(translation)

    def execute_obsolete_db(self):
        (Entity.objects
            .filter(pk__in=self.changes['obsolete_db'])
            .update(obsolete=True))
<commit_msg>Add --no-commit option to sync_projects.<commit_after>from django.contrib.auth.models import User
from django.core.management.base import BaseCommand, CommandError
from django.utils import timezone

from bulk_update.helper import bulk_update

from pontoon.administration.files import update_from_repository
from pontoon.administration.vcs import commit_to_vcs
from pontoon.base.models import Entity, Locale, Project, Resource, Translation
from pontoon.base.utils import match_attr
from pontoon.base.vcs_models import VCSProject

class Command(BaseCommand):
    args = '<project_slug project_slug ...>'
    help = 'Synchronize database and remote repositories.'

    def add_arguments(self, parser):
        parser.add_argument(
            '--no-commit',
            action='store_true',
            dest='no_commit',
            default=False,
            help='Do not commit changes to VCS'
        )
    def log(self, msg, *args, **kwargs):
        """Log a message to the console."""
        self.stdout.write(msg.format(*args, **kwargs))

    def info(self, msg, *args, **kwargs):
        """Log a message to the console if --verbosity=1 or more."""
        if self.verbosity >= 1:
            self.log(msg, *args, **kwargs)

    def debug(self, msg, *args, **kwargs):
        """Log a message to the console if --verbosity=2."""
        if self.verbosity == 2:
            self.log(msg, *args, **kwargs)

    def handle(self, *args, **options):
        self.verbosity = options['verbosity']
        self.no_commit = options['no_commit']

        self.log('SYNC PROJECTS: start')
        projects = Project.objects.filter(disabled=False)
        if args:
            projects = projects.filter(slug__in=args)

        if len(projects) < 1:
            raise CommandError('No matching projects found.')

        for project in projects:
            self.handle_project(project)
        self.log('SYNC PROJECTS: done')

        # Once we've synced, we can delete all translations scheduled
        # for deletion.
        Translation.deleted_objects.all().delete()

    def handle_project(self, db_project):
        if not db_project.can_commit:
            self.log('Skipping project {0}, cannot commit to repository.'
                     .format(db_project.name))
            return

        # Pull changes from VCS and update what we know about the files.
        update_from_repository(db_project)
        vcs_project = VCSProject(db_project)
        self.update_resources(db_project, vcs_project)

        # Collect all entities across VCS and the database and get their
        # keys so we can match up matching entities.
        vcs_entities = self.get_vcs_entities(vcs_project)
        db_entities = self.get_db_entities(db_project)
        entity_keys = set().union(db_entities.keys(), vcs_entities.keys())

        changeset = ChangeSet(db_project, vcs_project)
        for key in entity_keys:
            db_entity = db_entities.get(key, None)
            vcs_entity = vcs_entities.get(key, None)
            self.handle_entity(changeset, db_project, key, db_entity, vcs_entity)

        # Apply the changeset to the files and then commit them.
        changeset.execute()
        self.commit_changes(db_project)

        # Once we've successfully committed, update the last_synced date.
        db_project.last_synced = timezone.now()
        db_project.save()

        self.log('Synced project {0}', db_project.slug)

    def handle_entity(self, changeset, db_project, key, db_entity, vcs_entity):
        """
        Determine what needs to be synced between the database and VCS versions
        of a single entity and log what needs to be changed in the changeset.
        """
        if vcs_entity is None:
            if db_entity is None:
                # This should never happen. What? Hard abort.
                raise CommandError('No entities found for key {0}'.format(key))
            else:
                # VCS no longer has the entity, remove it from Pontoon.
                changeset.obsolete_db_entity(db_entity)
        elif db_entity is None:
            # New VCS entities are added to Pontoon.
            changeset.create_db_entity(vcs_entity)
        else:
            for locale in db_project.locales.all():
                if db_entity.has_changed(locale.code):
                    # Pontoon changes overwrite whatever VCS has.
                    changeset.update_vcs_entity(locale.code, db_entity, vcs_entity)
                elif vcs_entity.has_translation_for(locale.code):
                    # If Pontoon has nothing or has not changed, and the VCS
                    # still has the entity, update Pontoon with whatever may
                    # have changed.
                    changeset.update_db_entity(locale.code, db_entity, vcs_entity)

    def update_resources(self, db_project, vcs_project):
        """Update the database on what resource files exist in VCS."""
        relative_paths = vcs_project.resources.keys()
        db_project.resource_set.exclude(path__in=relative_paths).delete()

        for relative_path, vcs_resource in vcs_project.resources.items():
            resource, created = db_project.resource_set.get_or_create(path=relative_path)
            resource.format = Resource.get_path_format(relative_path)
            resource.entity_count = len(vcs_resource.entities)
            resource.save()

    def get_vcs_entities(self, vcs_project):
        return {self.entity_key(entity): entity for entity in vcs_project.entities}

    def get_db_entities(self, db_project):
        entities = (Entity.objects
                    .select_related('resource')
                    .filter(resource__project=db_project, obsolete=False))
        return {self.entity_key(entity): entity for entity in entities}

    def entity_key(self, entity):
        """
        Generate a key for the given entity that is unique within the
        project.
        """
        return ':'.join([entity.resource.path, entity.string])

    def commit_changes(self, db_project):
        """Commit the changes we've made back to the VCS."""
        if self.no_commit:
            return

        # TODO: Set translation author properly.
        user = User.objects.filter(is_superuser=True).first()

        for locale in db_project.locales.all():
            commit_message = (
                'Pontoon: Updated {locale.name} ({locale.code}) localization of '
                '{project.name}'
            ).format(locale=locale, project=db_project)

            try:
                result = commit_to_vcs(
                    db_project.repository_type,
                    db_project.locale_directory_path(locale.code),
                    commit_message,
                    user,
                    db_project.repository_url
                )
            except Exception as err:
                result = {'message': unicode(err)}

            if result is not None:
                self.log(
                    'Committing project {project.name} for {locale.name} '
                    '({locale.code}) failed: {reason}',
                    project=db_project,
                    locale=locale,
                    reason=result['message']
                )

class ChangeSet(object):
    """
    Stores a set of changes to be made to the database and the
    translations stored in VCS. Once all the necessary changes have been
    stored, execute all the changes at once efficiently.
    """
    def __init__(self, db_project, vcs_project):
        self.db_project = db_project
        self.vcs_project = vcs_project

        self.executed = False
        self.changes = {
            'update_vcs': [],
            'update_db': [],
            'obsolete_db': [],
            'create_db': []
        }

        self.entities_to_update = []
        self.translations_to_update = []
        self.translations_to_create = []

    def update_vcs_entity(self, locale_code, db_entity, vcs_entity):
        """
        Replace the translations in VCS with the translations from the
        database.
        """
        self.changes['update_vcs'].append((locale_code, db_entity, vcs_entity))

    def create_db_entity(self, vcs_entity):
        """Create a new entity in the database."""
        self.changes['create_db'].append(vcs_entity)

    def update_db_entity(self, locale_code, db_entity, vcs_entity):
        """Update the database with translations from VCS."""
        self.changes['update_db'].append((locale_code, db_entity, vcs_entity))

    def obsolete_db_entity(self, db_entity):
        """Mark the given entity as obsolete."""
        self.changes['obsolete_db'].append(db_entity.pk)

    def execute(self):
        """
        Execute the changes stored in this changeset. Execute can only
        be called once per changeset; subsequent calls raise a
        RuntimeError, even if the changes failed.
        """
        if self.executed:
            raise RuntimeError('execute() can only be called once per changeset.')
        else:
            self.executed = True

        # Store locales and resources for FK relationships.
        self.locales = {l.code: l for l in Locale.objects.all()}
        self.resources = {r.path: r for r in self.db_project.resource_set.all()}

        # Perform the changes and fill the lists for bulk creation and
        # updating.
        self.execute_update_vcs()
        self.execute_create_db()
        self.execute_update_db()
        self.execute_obsolete_db()

        # Apply the built-up changes to the DB
        if len(self.entities_to_update) > 0:
            bulk_update(self.entities_to_update, update_fields=[
                'resource',
                'string',
                'string_plural',
                'key',
                'comment',
                'order',
                'source'
            ])

        Translation.objects.bulk_create(self.translations_to_create)
        if len(self.translations_to_update) > 0:
            bulk_update(self.translations_to_update, update_fields=[
                'entity',
                'locale',
                'string',
                'plural_form',
                'approved',
                'approved_user',
                'approved_date',
                'fuzzy',
                'extra'
            ])

    def execute_update_vcs(self):
        resources = self.vcs_project.resources
        changed_resources = set()

        for locale_code, db_entity, vcs_entity in self.changes['update_vcs']:
            changed_resources.add(resources[db_entity.resource.path])

            vcs_translation = vcs_entity.translations[locale_code]
            db_translations = (db_entity.translation_set
                               .filter(approved=True, locale__code=locale_code))

            # If no DB translations are fuzzy, set fuzzy to False.
            # Otherwise, it's true.
            vcs_translation.fuzzy = any(t for t in db_translations if t.fuzzy)

            # Replace existing translations with ones from the database.
            vcs_translation.strings = {
                db.plural_form: db.string for db in db_translations
            }

        for resource in changed_resources:
            resource.save()

    def get_entity_updates(self, vcs_entity):
        """
        Return a dict of the properties and values necessary to create
        or update a database entity from a VCS entity.
        """
        return {
            'resource': self.resources[vcs_entity.resource.path],
            'string': vcs_entity.string,
            'string_plural': '',  # TODO: Support plural source.
            'key': '',  # TODO: Support keys separate from source.
            'comment': '\n'.join(vcs_entity.comments),
            'order': vcs_entity.order,
            'source': ''  # TODO: Support source
        }

    def execute_create_db(self):
        for vcs_entity in self.changes['create_db']:
            entity = Entity(**self.get_entity_updates(vcs_entity))
            entity.save()  # We can't use bulk_create since we need a PK

            for locale_code, vcs_translation in vcs_entity.translations.items():
                for plural_form, string in vcs_translation.strings.items():
                    self.translations_to_create.append(Translation(
                        entity=entity,
                        locale=self.locales[locale_code],
                        string=string,
                        plural_form=plural_form,
                        approved=True,
                        approved_date=timezone.now(),
                        fuzzy=vcs_translation.fuzzy
                    ))

    def execute_update_db(self):
        # TODO: Optimize to only update translations if necessary.
        for locale_code, db_entity, vcs_entity in self.changes['update_db']:
            for field, value in self.get_entity_updates(vcs_entity).items():
                setattr(db_entity, field, value)
            self.entities_to_update.append(db_entity)

            # Update translations for the entity.
            vcs_translation = vcs_entity.translations[locale_code]
            db_translations = db_entity.translation_set.filter(locale__code=locale_code)
            approved_translations = []

            for plural_form, string in vcs_translation.strings.items():
                # Check if we need to modify an existing translation or
                # create a new one.
                db_translation = match_attr(db_translations,
                                            plural_form=plural_form,
                                            string=string)
                if db_translation:
                    if not db_translation.approved:
                        db_translation.approved = True
                        db_translation.approved_date = timezone.now()
                    db_translation.fuzzy = vcs_translation.fuzzy
                    db_translation.extra = vcs_translation.extra

                    self.translations_to_update.append(db_translation)
                    approved_translations.append(db_translation)
                else:
                    self.translations_to_create.append(Translation(
                        entity=db_entity,
                        locale=self.locales[locale_code],
                        string=string,
                        plural_form=plural_form,
                        approved=True,
                        approved_date=timezone.now(),
                        fuzzy=vcs_translation.fuzzy,
                        extra=vcs_translation.extra
                    ))

            # Any existing translations that were not approved get unapproved.
            for translation in db_translations:
                if translation not in approved_translations:
                    translation.approved = False
                    translation.approved_user = None
                    translation.approved_date = None
                    self.translations_to_update.append(translation)

    def execute_obsolete_db(self):
        (Entity.objects
            .filter(pk__in=self.changes['obsolete_db'])
            .update(obsolete=True))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import division, print_function

__all__ = []

import os
import pickle
import numpy as np
import matplotlib.pyplot as pl
from scipy.linalg import cho_factor, cho_solve
from itertools import product
import scipy.optimize as op

def kernel(theta, l_i, l_j, white_noise=True):
    a, tau0, tau1, tau2, s = np.exp(theta)
    K = a**2 * np.exp(-0.5*(((l_i-l_j)[:,:,0])**2/tau0**2 +
                           ((l_i-l_j)[:,:,1])**2/tau1**2 +
                           ((l_i-l_j)[:,:,2])**2/tau2**2))
    if white_noise:
        K[np.diag_indices_from(K)] += s**2
    return K

# Calculate alpha
def compute_l_and_alpha(theta, f, var, l_all):
    # compute the kernel
    K = kernel(theta, l_all[:, None], l_all[None, :])
    # add in the uncertainties in the fluxes to the kernel
    K[np.diag_indices_from(K)] += var
    # factorize the kernel
    L, flag = cho_factor(K, overwrite_a=True)
    return (L, flag), cho_solve((L, flag), f)

# Likelihood function
def ln_likelihood(theta, f, var, l_all):
    (L, flag), alpha = compute_l_and_alpha(theta, f, var, l_all)
    # compute half of the log determinant
    lndet = np.sum(np.log(np.diag(L)))
    return -0.5*np.dot(f, alpha)-lndet

# Evaluate the predictive mean
def predictive_mean(theta, test_labels):
    K_star = kernel(theta, test_labels[:, None],
                    l_all[None, :], white_noise=False)
    alpha = compute_l_and_alpha(theta)[1]
    f_star_bar = np.dot(K_star, alpha)
    return f_star_bar

# Evaluate the likelihood of a data point given model
def test_ln_likelihood_single_pix(label, flux, flux_var, theta, (L,flag), alpha, l_all):
    # Calculate the mean prediction.
    K_star = kernel(theta, label[None,None,:], l_all[None,:,:],
                    white_noise=False)
    f_star = np.dot(K_star, alpha)

    # Calculate the variance on the prediction.
    label_array = np.array(label[None,None,:])  # Hack so kernel works w scalar.
    pred_var = flux_var + float(kernel(theta, label_array, label_array))
    pred_var -= np.dot(K_star, cho_solve((L,flag), K_star[0]))

    return -0.5 * ((flux - f_star) ** 2 / pred_var + np.log(pred_var))

def test_ln_likelihood(label, f, var, theta_all, f_all, var_all, l_all):
    print(label)
    npix = len(f)
    npix = 150
    ln_likelihoods = np.zeros(npix)
    for pix in range(npix):
        (L,flag), alpha = compute_l_and_alpha(
                theta_all[pix,:], f_all[:,pix], var_all[:,pix], l_all)
        ln_likelihoods[pix] = test_ln_likelihood_single_pix(
                label, f[pix], var[pix], theta_all[pix], (L,flag), alpha, l_all)
    return sum(ln_likelihoods)

def pick_good_obj(pix, f_all, ivar_all, l_all):
    f = f_all[:, pix]
    ivar = ivar_all[:, pix]

    # this pixel is bad for the following objects: flux == 0 or flux == 1
    m = np.logical_and(f != 0.0, f != 1.0)
    f = f[m]
    f -= np.mean(f)
    l = l_all[m,:]
    var = 1./ivar[m]
    return f,var,l

def train_single_pix(pix, f_all, ivar_all, l_all):
    f, var, l = pick_good_obj(pix, f_all, ivar_all, l_all)
    nll = lambda theta: -ln_likelihood(theta, f, var, l)
    bounds = np.log([(0.001, 0.1), (100, 10000), (0.01, 10), (0.01, 10),
                     (0.0001, 0.1)])
    output = op.minimize(nll, np.log([0.01, 100, 0.5, 0.05, 0.01]),
                         method="L-BFGS-B", bounds=bounds)
    best_hyperparams = output.x
    print("best hyperparams: ")
    print(ln_likelihoods[pix])
    print(np.exp(best_hyperparams))
    return best_hyperparams

def infer_labels_single_pix(pix, best_hyperparams, f_all, var_all, l_all):
    f, var, l = pick_good_obj(pix, f_all, ivar_all, l_all)
    hyperparams = best_hyperparams[pix,:]
    (L,flag), alpha = compute_l_and_alpha(hyperparams, f, var, l)
    tll = lambda labels: -test_ln_likelihood(labels, f[obj], var[obj],
                                             hyperparams, (L,flag), alpha, l)
    p0 = np.array([np.mean(l[:,i]) for i in range(l.shape[1])])
    bounds = np.array([(3500,5500), (0,5), (-2.5,0.5)]) # training set dist

    print("fitting for labels")
    # want to minimize the negative log likelihood (maximize lnp)
    output = op.minimize(tll, p0, method="L-BFGS-B", bounds=bounds)
    return output.x

def infer_labels_single_obj(obj, hyperparams_all, f_all, var_all, l_all):
    f = f_all[obj,:]
    var = var_all[obj,:]
    # l = l_all[obj,:]
    tll = lambda labels: -test_ln_likelihood(
            labels, f, var, hyperparams_all, f_all, var_all, l_all)
    p0 = np.array([np.mean(l_all[:,i]) for i in range(l_all.shape[1])])
    bounds = np.array([(3500,5500), (0,5), (-2.5,0.5)]) # training set dist
    print("fitting for labels")
    output = op.minimize(tll, p0, method="L-BFGS-B", bounds=bounds)
    return output

f_all = pickle.load(open("data/norm_tr_fluxes.p", "r"))
l_all = pickle.load(open("data/tr_label_vals.p", "r"))
ivar_all = pickle.load(open("data/norm_tr_ivars.p", "r"))

# "bad" pixels are those with var == 0
keep = np.var(f_all, axis=0) > 0
f_all = f_all[:,keep]
ivar_all = ivar_all[:,keep]

npix = f_all.shape[1]

# Covariance function

hyperfn = "best_hyperparams.p"
if not os.path.exists(hyperfn):
    print("fitting for hyperparams")
    # Optimize for each pixel independently
    npix = f_all.shape[1]
    best_hyperparams_all = np.zeros((npix, 5))
    for pix in range(npix):
        print("training on pixel " + str(pix))
        best_hyperparams_all[pix,:] = train_single_pix(pix, f_all, ivar_all, l_all)
    pickle.dump(best_hyperparams_all, open(hyperfn, "wb"), -1)

else:
    print("loading")
    best_hyperparams = pickle.load(open(hyperfn, "r"))

##### TEST STEP #####

# Now optimize over the test label space
obj = 0
print(infer_labels_single_obj(obj, best_hyperparams, f_all, ivar_all, l_all))

# axis = 2
# test_labels[:,axis] = np.linspace(np.min(l_all[:,axis]),
#                                   np.max(l_all[:,axis]),
#                                   len(test_labels[:,axis]))
# mu = predictive_mean(best_hyperparams, test_labels)
# pl.scatter(l_all[:,axis], f, c="k")
# pl.plot(test_labels[:,axis], mu, "g", lw=2)
# pl.show()

# (L,flag), alpha = compute_l_and_alpha(best_hyperparams)
# pl.figure()
# for obj in range(5,9):
#     lls = [test_ln_likelihood(test_labels[i,:], f[obj], var[obj],
#            best_hyperparams, (L,flag), alpha) for i in range(0,len(test_labels))]
#     pl.plot(test_labels[:,axis], lls, label="%s"%np.round(l_all[obj,axis],2))
#     pl.legend()
#     pl.gca().axvline(l_all[obj,axis])
# pl.show()

<commit_msg>change training step to map()<commit_after># -*- coding: utf-8 -*-

from __future__ import division, print_function

__all__ = []

import os
import pickle
import numpy as np
import matplotlib.pyplot as pl
from scipy.linalg import cho_factor, cho_solve
from itertools import product
import scipy.optimize as op

def kernel(theta, l_i, l_j, white_noise=True):
    a, tau0, tau1, tau2, s = np.exp(theta)
    K = a**2 * np.exp(-0.5*(((l_i-l_j)[:,:,0])**2/tau0**2 +
                           ((l_i-l_j)[:,:,1])**2/tau1**2 +
                           ((l_i-l_j)[:,:,2])**2/tau2**2))
    if white_noise:
        K[np.diag_indices_from(K)] += s**2
    return K

# Calculate alpha
def compute_l_and_alpha(theta, f, var, l_all):
    # compute the kernel
    K = kernel(theta, l_all[:, None], l_all[None, :])
    # add in the uncertainties in the fluxes to the kernel
    K[np.diag_indices_from(K)] += var
    # factorize the kernel
    L, flag = cho_factor(K, overwrite_a=True)
    return (L, flag), cho_solve((L, flag), f)

# Likelihood function
def ln_likelihood(theta, f, var, l_all):
    (L, flag), alpha = compute_l_and_alpha(theta, f, var, l_all)
    # compute half of the log determinant
    lndet = np.sum(np.log(np.diag(L)))
    return -0.5*np.dot(f, alpha)-lndet

# Evaluate the predictive mean
def predictive_mean(theta, test_labels):
    K_star = kernel(theta, test_labels[:, None],
                    l_all[None, :], white_noise=False)
    alpha = compute_l_and_alpha(theta)[1]
    f_star_bar = np.dot(K_star, alpha)
    return f_star_bar

# Evaluate the likelihood of a data point given model
def test_ln_likelihood_single_pix(label, flux, flux_var, theta, (L,flag), alpha, l_all):
    # Calculate the mean prediction.
    K_star = kernel(theta, label[None,None,:], l_all[None,:,:],
                    white_noise=False)
    f_star = np.dot(K_star, alpha)

    # Calculate the variance on the prediction.
    label_array = np.array(label[None,None,:])  # Hack so kernel works w scalar.
    pred_var = flux_var + float(kernel(theta, label_array, label_array))
    pred_var -= np.dot(K_star, cho_solve((L,flag), K_star[0]))

    return -0.5 * ((flux - f_star) ** 2 / pred_var + np.log(pred_var))

def test_ln_likelihood(label, f, var, theta_all, f_all, var_all, l_all):
    #print(label)
    npix = len(f)
    npix = 10
    ln_likelihoods = np.zeros(npix)
    for pix in range(100,npix):
        (L,flag), alpha = compute_l_and_alpha(
                theta_all[pix,:], f_all[:,pix], var_all[:,pix], l_all)
        ln_likelihoods[pix] = test_ln_likelihood_single_pix(
                label, f[pix], var[pix], theta_all[pix], (L,flag), alpha, l_all)
    return sum(ln_likelihoods)

def pick_good_obj(pix, f_all, ivar_all, l_all):
    f = f_all[:, pix]
    ivar = ivar_all[:, pix]

    # this pixel is bad for the following objects: flux == 0 or flux == 1
    m = np.logical_and(f != 0.0, f != 1.0)
    f = f[m]
    f -= np.mean(f)
    l = l_all[m,:]
    var = 1./ivar[m]
    return f,var,l

def train_single_pix(pix, f_all, ivar_all, l_all):
    f, var, l = pick_good_obj(pix, f_all, ivar_all, l_all)
    nll = lambda theta: -ln_likelihood(theta, f, var, l)
    bounds = np.log([(0.001, 0.1), (100, 10000), (0.01, 10), (0.01, 10),
                     (0.0001, 0.1)])
    output = op.minimize(nll, np.log([0.01, 100, 0.5, 0.05, 0.01]),
                         method="L-BFGS-B", bounds=bounds)
    best_hyperparams = output.x
    print("best hyperparams: ")
    print(ln_likelihoods[pix])
    print(np.exp(best_hyperparams))
    return best_hyperparams

def infer_labels_single_pix(pix, best_hyperparams, f_all, var_all, l_all):
    f, var, l = pick_good_obj(pix, f_all, ivar_all, l_all)
    hyperparams = best_hyperparams[pix,:]
    (L,flag), alpha = compute_l_and_alpha(hyperparams, f, var, l)
    tll = lambda labels: -test_ln_likelihood(labels, f[obj], var[obj],
                                             hyperparams, (L,flag), alpha, l)
    p0 = np.array([np.mean(l[:,i]) for i in range(l.shape[1])])
    bounds = np.array([(3500,5500), (0,5), (-2.5,0.5)]) # training set dist

    print("fitting for labels")
    # want to minimize the negative log likelihood (maximize lnp)
    output = op.minimize(tll, p0, method="L-BFGS-B", bounds=bounds)
    return output.x

def infer_labels_single_obj(obj, hyperparams_all, f_all, var_all, l_all):
    f = f_all[obj,:]
    var = var_all[obj,:]
    # l = l_all[obj,:]
    tll = lambda labels: -test_ln_likelihood(
            labels, f, var, hyperparams_all, f_all, var_all, l_all)
    p0 = np.array([np.mean(l_all[:,i]) for i in range(l_all.shape[1])])
    #p0 = l_all[obj,:]
    bounds = np.array([(3500,5500), (0,5), (-2.5,0.5)]) # training set dist
    print("fitting for labels")
    output = op.minimize(tll, p0, method="L-BFGS-B", bounds=bounds)
    return output

f_all = pickle.load(open("data/norm_tr_fluxes.p", "r"))
l_all = pickle.load(open("data/tr_label_vals.p", "r"))
ivar_all = pickle.load(open("data/norm_tr_ivars.p", "r"))

# "bad" pixels are those with var == 0
keep = np.var(f_all, axis=0) > 0
f_all = f_all[:,keep]
ivar_all = ivar_all[:,keep]

npix = f_all.shape[1]

# Covariance function

hyperfn = "best_hyperparams.p"
if not os.path.exists(hyperfn):
    print("fitting for hyperparams")
    # Optimize for each pixel independently
    npix = f_all.shape[1]
    best_hyperparams_all = map(
            train_single_pix(pix, f_all, ivar_all, l_all), np.linspace(0,npix,npix+1))
    pickle.dump(best_hyperparams_all, open(hyperfn, "wb"), -1)

else:
    print("loading")
    best_hyperparams = pickle.load(open(hyperfn, "r"))

##### TEST STEP #####

# Now optimize over the test label space
obj = 0
print(infer_labels_single_obj(obj, best_hyperparams, f_all, ivar_all, l_all))

# axis = 2
# test_labels[:,axis] = np.linspace(np.min(l_all[:,axis]),
#                                   np.max(l_all[:,axis]),
#                                   len(test_labels[:,axis]))
# mu = predictive_mean(best_hyperparams, test_labels)
# pl.scatter(l_all[:,axis], f, c="k")
# pl.plot(test_labels[:,axis], mu, "g", lw=2)
# pl.show()

# (L,flag), alpha = compute_l_and_alpha(best_hyperparams)
# pl.figure()
# for obj in range(5,9):
#     lls = [test_ln_likelihood(test_labels[i,:], f[obj], var[obj],
#            best_hyperparams, (L,flag), alpha) for i in range(0,len(test_labels))]
#     pl.plot(test_labels[:,axis], lls, label="%s"%np.round(l_all[obj,axis],2))
#     pl.legend()
#     pl.gca().axvline(l_all[obj,axis])
# pl.show()

<|endoftext|>