<commit_before>
import logging
from collections import defaultdict
from itertools import count

from claripy.utils.orderedset import OrderedSet
from simuvex.s_variable import SimStackVariable, SimMemoryVariable, SimRegisterVariable, SimMemoryVariablePhi, \
    SimStackVariablePhi, SimRegisterVariablePhi

from .keyed_region import KeyedRegion
from .variable_access import VariableAccess

l = logging.getLogger('angr.knowledge.variable_manager')

class VariableType(object):
    REGISTER = 0
    MEMORY = 1

class LiveVariables(object):
    """
    A collection of live variables at a program point.
    """
    def __init__(self, register_region, stack_region):
        self.register_region = register_region
        self.stack_region = stack_region

class VariableManagerInternal(object):
    """
    Manage variables for a function. It is meant to be used internally by VariableManager.
    """
    def __init__(self, manager, func_addr=None):
        self.manager = manager

        self.func_addr = func_addr

        self._variables = OrderedSet()  # all variables that are added to any region
        self._stack_region = KeyedRegion()
        self._register_region = KeyedRegion()
        self._live_variables = { }  # a mapping between addresses of program points and live variable collections

        self._variable_accesses = defaultdict(set)
        self._insn_to_variable = defaultdict(set)
        self._block_to_variable = defaultdict(set)
        self._stmt_to_variable = defaultdict(set)
        self._variable_counters = {
            'register': count(),
            'stack': count(),
            'argument': count(),
            'phi': count(),
        }

    #
    # Public methods
    #

    def next_variable_ident(self, sort):
        if sort not in self._variable_counters:
            raise ValueError('Unsupported variable sort %s' % sort)

        if sort == 'register':
            prefix = "r"
        elif sort == 'stack':
            prefix = "s"
        elif sort == 'argument':
            prefix = 'arg'
        else:
            prefix = "m"

        return "i%s_%d" % (prefix, self._variable_counters[sort].next())

    def add_variable(self, sort, start, variable):
        if sort == 'stack':
            self._stack_region.add_variable(start, variable)
        elif sort == 'register':
            self._register_region.add_variable(start, variable)
        else:
            raise ValueError('Unsupported sort %s in add_variable().' % sort)

    def set_variable(self, sort, start, variable):
        if sort == 'stack':
            self._stack_region.set_variable(start, variable)
        elif sort == 'register':
            self._register_region.set_variable(start, variable)
        else:
            raise ValueError('Unsupported sort %s in add_variable().' % sort)

    def write_to(self, variable, offset, location, overwrite=False):
        self._record_variable_access('write', variable, offset, location, overwrite=overwrite)

    def read_from(self, variable, offset, location, overwrite=False):
        self._record_variable_access('read', variable, offset, location, overwrite=overwrite)

    def reference_at(self, variable, offset, location, overwrite=False):
        self._record_variable_access('reference', variable, offset, location, overwrite=overwrite)

    def _record_variable_access(self, sort, variable, offset, location, overwrite=False):
        self._variables.add(variable)
        if overwrite:
            self._variable_accesses[variable] = {VariableAccess(variable, sort, location)}
            self._insn_to_variable[location.ins_addr] = {(variable, offset)}
            self._block_to_variable[location.block_addr] = {(variable, offset)}
            self._stmt_to_variable[(location.block_addr, location.stmt_idx)] = {(variable, offset)}
        else:
            self._variable_accesses[variable].add(VariableAccess(variable, sort, location))
            self._insn_to_variable[location.ins_addr].add((variable, offset))
            self._block_to_variable[location.block_addr].add((variable, offset))
            self._stmt_to_variable[(location.block_addr, location.stmt_idx)].add((variable, offset))

    def make_phi_node(self, *variables):

        # unpack phi nodes
        existing_phi = [ ]
        unpacked = set()
        for var in variables:
            if isinstance(var, (SimRegisterVariablePhi, SimStackVariablePhi, SimMemoryVariablePhi)):
                unpacked |= var.variables
                existing_phi.append(var)
            else:
                unpacked.add(var)

        # optimization: if a phi node already contains all of the unpacked variables, just return that phi node
        for phi_node in existing_phi:
            if phi_node.variables.issuperset(unpacked):
                return phi_node

        variables = unpacked

        repre = next(iter(variables))
        repre_type = type(repre)
        if repre_type is SimRegisterVariable:
            cls = SimRegisterVariablePhi
            ident_sort = 'register'
        elif repre_type is SimMemoryVariable:
            cls = SimMemoryVariablePhi
            ident_sort = 'memory'
        elif repre_type is SimStackVariable:
            cls = SimStackVariablePhi
            ident_sort = 'stack'
        else:
            raise TypeError('make_phi_node(): Unsupported variable type "%s".' % type(repre))
        a = cls(ident=self.next_variable_ident(ident_sort),
                   region=self.func_addr,
                   variables=variables,
                   )
        return a

    def set_live_variables(self, addr, register_region, stack_region):
        lv = LiveVariables(register_region, stack_region)
        self._live_variables[addr] = lv

    def find_variables_by_insn(self, ins_addr, sort):
        if ins_addr not in self._insn_to_variable:
            return None

        if sort == VariableType.MEMORY or sort == 'memory':
            vars_and_offset = [(var, offset) for var, offset in self._insn_to_variable[ins_addr]
                        if isinstance(var, (SimStackVariable, SimMemoryVariable))]
        elif sort == VariableType.REGISTER or sort == 'register':
            vars_and_offset = [(var, offset) for var, offset in self._insn_to_variable[ins_addr]
                        if isinstance(var, SimRegisterVariable)]
        else:
            l.error('find_variable_by_insn(): Unsupported variable sort "%s".', sort)
            return [ ]

        return vars_and_offset

    def find_variable_by_stmt(self, block_addr, stmt_idx, sort):
        return next(iter(self.find_variables_by_stmt(block_addr, stmt_idx, sort)), None)

    def find_variables_by_stmt(self, block_addr, stmt_idx, sort):

        key = block_addr, stmt_idx

        if key not in self._stmt_to_variable:
            return [ ]

        variables = self._stmt_to_variable[key]
        if not variables:
            return [ ]

        if sort == 'memory':
            var_and_offsets = list((var, offset) for var, offset in self._stmt_to_variable[key]
                                   if isinstance(var, (SimStackVariable, SimMemoryVariable)))
        elif sort == 'register':
            var_and_offsets = list((var, offset) for var, offset in self._stmt_to_variable[key]
                                   if isinstance(var, SimRegisterVariable))
        else:
            l.error('find_variables_by_stmt(): Unsupported variable sort "%s".', sort)
            return [ ]

        return var_and_offsets

    def get_variable_accesses(self, variable, same_name=False):

        if not same_name:
            if variable in self._variable_accesses:
                return self._variable_accesses[variable]

            return [ ]

        # find all variables with the same variable name

        vars_list = [ ]

        for var in self._variable_accesses.keys():
            if variable.name == var.name:
                vars_list.append(var)

        accesses = [ ]
        for var in vars_list:
            accesses.extend(self.get_variable_accesses(var))

        return accesses

    def input_variables(self):
        """
        Get all variables that have never been written to.

        :return: A list of variables that are never written to.
        """

        def has_write_access(accesses):
            return any(acc for acc in accesses if acc.access_type == 'write')

        def has_read_access(accesses):
            return any(acc for acc in accesses if acc.access_type == 'read')

        input_variables = [ ]

        for variable, accesses in self._variable_accesses.iteritems():
            if not has_write_access(accesses) and has_read_access(accesses):
                input_variables.append(variable)

        return input_variables

    def assign_variable_names(self):
        """
        Assign default names to all variables.

        :return: None
        """

        for var in self._variables:
            if isinstance(var, SimStackVariable):
                if var.name is not None:
                    continue
                if var.ident.startswith('iarg'):
                    var.name = 'arg_%x' % var.offset
                else:
                    var.name = 's_%x' % (-var.offset)
                    # var.name = var.ident
            elif isinstance(var, SimRegisterVariable):
                if var.name is not None:
                    continue
                var.name = var.ident

class VariableManager(object):
    """
    Manage variables.
    """
    def __init__(self, kb=None):
        self.kb = kb
        self.global_manager = VariableManagerInternal(self)
        self.function_managers = { }

    def __getitem__(self, key):
        """
        Get the VariableManagerInternal object for a function or a region.

        :param str or int key: Key of the region. "global" for the global region, or a function address for the
                               function.
        :return:               The VariableManagerInternal object.
        :rtype:                VariableManagerInternal
        """

        if key == 'global':  # pylint:disable=no-else-return
            return self.global_manager

        else:
            # key refers to a function address
            return self.get_function_manager(key)

    def get_function_manager(self, func_addr):
        if not isinstance(func_addr, (int, long)):
            raise TypeError('Argument "func_addr" must be an int.')

        if func_addr not in self.function_managers:
            self.function_managers[func_addr] = VariableManagerInternal(self, func_addr=func_addr)

        return self.function_managers[func_addr]

    def initialize_variable_names(self):
        self.global_manager.assign_variable_names()
        for manager in self.function_managers.itervalues():
            manager.assign_variable_names()

    def get_variable_accesses(self, variable, same_name=False):
        """
        Get a list of all references to the given variable.

        :param simuvex.SimVariable variable: The variable.
        :param bool same_name:               Whether to include all variables with the same variable name, or just
                                             based on the variable identifier.
        :return:                All references to the variable.
        :rtype:                 list
        """

        if variable.region == 'global':
            return self.global_manager.get_variable_accesses(variable, same_name=same_name)

        elif variable.region in self.function_managers:
            return self.function_managers[variable.region].get_variable_accesses(variable, same_name=same_name)

        l.warning('get_variable_accesses(): Region %s is not found.', variable.region)
        return [ ]
<commit_msg>VariableManagerInternal: get_variables().<commit_after>
import logging
from collections import defaultdict
from itertools import count

from claripy.utils.orderedset import OrderedSet
from simuvex.s_variable import SimStackVariable, SimMemoryVariable, SimRegisterVariable, SimMemoryVariablePhi, \
    SimStackVariablePhi, SimRegisterVariablePhi

from .keyed_region import KeyedRegion
from .variable_access import VariableAccess

l = logging.getLogger('angr.knowledge.variable_manager')

class VariableType(object):
    REGISTER = 0
    MEMORY = 1

class LiveVariables(object):
    """
    A collection of live variables at a program point.
    """
    def __init__(self, register_region, stack_region):
        self.register_region = register_region
        self.stack_region = stack_region

class VariableManagerInternal(object):
    """
    Manage variables for a function. It is meant to be used internally by VariableManager.
    """
    def __init__(self, manager, func_addr=None):
        self.manager = manager

        self.func_addr = func_addr

        self._variables = OrderedSet()  # all variables that are added to any region
        self._stack_region = KeyedRegion()
        self._register_region = KeyedRegion()
        self._live_variables = { }  # a mapping between addresses of program points and live variable collections

        self._variable_accesses = defaultdict(set)
        self._insn_to_variable = defaultdict(set)
        self._block_to_variable = defaultdict(set)
        self._stmt_to_variable = defaultdict(set)
        self._variable_counters = {
            'register': count(),
            'stack': count(),
            'argument': count(),
            'phi': count(),
        }

    #
    # Public methods
    #

    def next_variable_ident(self, sort):
        if sort not in self._variable_counters:
            raise ValueError('Unsupported variable sort %s' % sort)

        if sort == 'register':
            prefix = "r"
        elif sort == 'stack':
            prefix = "s"
        elif sort == 'argument':
            prefix = 'arg'
        else:
            prefix = "m"

        return "i%s_%d" % (prefix, self._variable_counters[sort].next())

    def add_variable(self, sort, start, variable):
        if sort == 'stack':
            self._stack_region.add_variable(start, variable)
        elif sort == 'register':
            self._register_region.add_variable(start, variable)
        else:
            raise ValueError('Unsupported sort %s in add_variable().' % sort)

    def set_variable(self, sort, start, variable):
        if sort == 'stack':
            self._stack_region.set_variable(start, variable)
        elif sort == 'register':
            self._register_region.set_variable(start, variable)
        else:
            raise ValueError('Unsupported sort %s in add_variable().' % sort)

    def write_to(self, variable, offset, location, overwrite=False):
        self._record_variable_access('write', variable, offset, location, overwrite=overwrite)

    def read_from(self, variable, offset, location, overwrite=False):
        self._record_variable_access('read', variable, offset, location, overwrite=overwrite)

    def reference_at(self, variable, offset, location, overwrite=False):
        self._record_variable_access('reference', variable, offset, location, overwrite=overwrite)

    def _record_variable_access(self, sort, variable, offset, location, overwrite=False):
        self._variables.add(variable)
        if overwrite:
            self._variable_accesses[variable] = {VariableAccess(variable, sort, location)}
            self._insn_to_variable[location.ins_addr] = {(variable, offset)}
            self._block_to_variable[location.block_addr] = {(variable, offset)}
            self._stmt_to_variable[(location.block_addr, location.stmt_idx)] = {(variable, offset)}
        else:
            self._variable_accesses[variable].add(VariableAccess(variable, sort, location))
            self._insn_to_variable[location.ins_addr].add((variable, offset))
            self._block_to_variable[location.block_addr].add((variable, offset))
            self._stmt_to_variable[(location.block_addr, location.stmt_idx)].add((variable, offset))

    def make_phi_node(self, *variables):

        # unpack phi nodes
        existing_phi = [ ]
        unpacked = set()
        for var in variables:
            if isinstance(var, (SimRegisterVariablePhi, SimStackVariablePhi, SimMemoryVariablePhi)):
                unpacked |= var.variables
                existing_phi.append(var)
            else:
                unpacked.add(var)

        # optimization: if a phi node already contains all of the unpacked variables, just return that phi node
        for phi_node in existing_phi:
            if phi_node.variables.issuperset(unpacked):
                return phi_node

        variables = unpacked

        repre = next(iter(variables))
        repre_type = type(repre)
        if repre_type is SimRegisterVariable:
            cls = SimRegisterVariablePhi
            ident_sort = 'register'
        elif repre_type is SimMemoryVariable:
            cls = SimMemoryVariablePhi
            ident_sort = 'memory'
        elif repre_type is SimStackVariable:
            cls = SimStackVariablePhi
            ident_sort = 'stack'
        else:
            raise TypeError('make_phi_node(): Unsupported variable type "%s".' % type(repre))
        a = cls(ident=self.next_variable_ident(ident_sort),
                   region=self.func_addr,
                   variables=variables,
                   )
        return a

    def set_live_variables(self, addr, register_region, stack_region):
        lv = LiveVariables(register_region, stack_region)
        self._live_variables[addr] = lv

    def find_variables_by_insn(self, ins_addr, sort):
        if ins_addr not in self._insn_to_variable:
            return None

        if sort == VariableType.MEMORY or sort == 'memory':
            vars_and_offset = [(var, offset) for var, offset in self._insn_to_variable[ins_addr]
                        if isinstance(var, (SimStackVariable, SimMemoryVariable))]
        elif sort == VariableType.REGISTER or sort == 'register':
            vars_and_offset = [(var, offset) for var, offset in self._insn_to_variable[ins_addr]
                        if isinstance(var, SimRegisterVariable)]
        else:
            l.error('find_variable_by_insn(): Unsupported variable sort "%s".', sort)
            return [ ]

        return vars_and_offset

    def find_variable_by_stmt(self, block_addr, stmt_idx, sort):
        return next(iter(self.find_variables_by_stmt(block_addr, stmt_idx, sort)), None)

    def find_variables_by_stmt(self, block_addr, stmt_idx, sort):

        key = block_addr, stmt_idx

        if key not in self._stmt_to_variable:
            return [ ]

        variables = self._stmt_to_variable[key]
        if not variables:
            return [ ]

        if sort == 'memory':
            var_and_offsets = list((var, offset) for var, offset in self._stmt_to_variable[key]
                                   if isinstance(var, (SimStackVariable, SimMemoryVariable)))
        elif sort == 'register':
            var_and_offsets = list((var, offset) for var, offset in self._stmt_to_variable[key]
                                   if isinstance(var, SimRegisterVariable))
        else:
            l.error('find_variables_by_stmt(): Unsupported variable sort "%s".', sort)
            return [ ]

        return var_and_offsets

    def get_variable_accesses(self, variable, same_name=False):

        if not same_name:
            if variable in self._variable_accesses:
                return self._variable_accesses[variable]

            return [ ]

        # find all variables with the same variable name

        vars_list = [ ]

        for var in self._variable_accesses.keys():
            if variable.name == var.name:
                vars_list.append(var)

        accesses = [ ]
        for var in vars_list:
            accesses.extend(self.get_variable_accesses(var))

        return accesses

    def get_variables(self, sort=None, collapse_same_ident=False):
        """
        Get a list of variables.

        :param str or None sort:    Sort of the variable to get.
        :param collapse_same_ident: Whether variables of the same identifier should be collapsed or not.
        :return:                    A list of variables.
        :rtype:                     list
        """

        variables = [ ]

        if collapse_same_ident:
            raise NotImplementedError()

        for var in self._variables:
            if sort == 'stack' and not isinstance(var, SimStackVariable):
                continue
            if sort == 'reg' and not isinstance(var, SimRegisterVariable):
                continue
            variables.append(var)

        return variables

    def input_variables(self):
        """
        Get all variables that have never been written to.

        :return: A list of variables that are never written to.
        """

        def has_write_access(accesses):
            return any(acc for acc in accesses if acc.access_type == 'write')

        def has_read_access(accesses):
            return any(acc for acc in accesses if acc.access_type == 'read')

        input_variables = [ ]

        for variable, accesses in self._variable_accesses.iteritems():
            if not has_write_access(accesses) and has_read_access(accesses):
                input_variables.append(variable)

        return input_variables

    def assign_variable_names(self):
        """
        Assign default names to all variables.

        :return: None
        """

        for var in self._variables:
            if isinstance(var, SimStackVariable):
                if var.name is not None:
                    continue
                if var.ident.startswith('iarg'):
                    var.name = 'arg_%x' % var.offset
                else:
                    var.name = 's_%x' % (-var.offset)
                    # var.name = var.ident
            elif isinstance(var, SimRegisterVariable):
                if var.name is not None:
                    continue
                var.name = var.ident

class VariableManager(object):
    """
    Manage variables.
    """
    def __init__(self, kb=None):
        self.kb = kb
        self.global_manager = VariableManagerInternal(self)
        self.function_managers = { }

    def __getitem__(self, key):
        """
        Get the VariableManagerInternal object for a function or a region.

        :param str or int key: Key of the region. "global" for the global region, or a function address for the
                               function.
        :return:               The VariableManagerInternal object.
        :rtype:                VariableManagerInternal
        """

        if key == 'global':  # pylint:disable=no-else-return
            return self.global_manager

        else:
            # key refers to a function address
            return self.get_function_manager(key)

    def get_function_manager(self, func_addr):
        if not isinstance(func_addr, (int, long)):
            raise TypeError('Argument "func_addr" must be an int.')

        if func_addr not in self.function_managers:
            self.function_managers[func_addr] = VariableManagerInternal(self, func_addr=func_addr)

        return self.function_managers[func_addr]

    def initialize_variable_names(self):
        self.global_manager.assign_variable_names()
        for manager in self.function_managers.itervalues():
            manager.assign_variable_names()

    def get_variable_accesses(self, variable, same_name=False):
        """
        Get a list of all references to the given variable.

        :param simuvex.SimVariable variable: The variable.
        :param bool same_name:               Whether to include all variables with the same variable name, or just
                                             based on the variable identifier.
        :return:                All references to the variable.
        :rtype:                 list
        """

        if variable.region == 'global':
            return self.global_manager.get_variable_accesses(variable, same_name=same_name)

        elif variable.region in self.function_managers:
            return self.function_managers[variable.region].get_variable_accesses(variable, same_name=same_name)

        l.warning('get_variable_accesses(): Region %s is not found.', variable.region)
        return [ ]
<|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
#
#    https://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.

"""Pre-Processes the MAESTRO Dataset.

Pre-proceses the raw MAESTRO dataset into a data set
of music chunks with a predefined lenth and sampling rate.
The music chunks are then saved as a .npz file
"""

import os
import sys
import glob
import librosa
import mido
import copy
import numpy as np

# Overall Config #
APPROX_TOTAL_HOURS = 6
SAMPLE_RATE = 16000 # 16kHz
DATA_POINT_LENGTH = 2**14
BLOCK_IN_SECONDS = DATA_POINT_LENGTH / SAMPLE_RATE
BLOCK_N_QUANT = 128
BLOCK_ZONE_LENGTH = BLOCK_IN_SECONDS / BLOCK_N_QUANT
RAW_DATA_PATH = './data/maestro/2017'
PROCESSED_DATA_PATH = '../data/'

def main():
    audio_paths = glob.glob(RAW_DATA_PATH + '/**/*.wav', recursive=True)

    # Load audio files until we reach our desired
    # data set size.
    data = []
    midi_data = []
    hours_loaded = 0
    for audio_file_path in audio_paths:
        print(audio_file_path)
        wav, _ = librosa.load(audio_file_path, sr=SAMPLE_RATE)
        print(len(wav) / 16000)

        hours_loaded += len(wav) / SAMPLE_RATE / 60 / 60
        if hours_loaded >= APPROX_TOTAL_HOURS:
            break

        # Pad the song to ensure it can be evenly divided into
        # chunks of length 'DATA_POINT_LENGTH'
        padded_wav_length = int(DATA_POINT_LENGTH * np.ceil(len(wav) / DATA_POINT_LENGTH))
        wav = np.pad(wav, [[0, padded_wav_length - len(wav)]])

        chunks = np.reshape(wav, (-1, DATA_POINT_LENGTH))
        data.extend(chunks)

        # Load and process MIDI file
        midi_file_path = audio_file_path[0:-3] + 'midi'#.replace('.wav', '.midi')
        print(midi_file_path)
        mid = mido.MidiFile(midi_file_path, clip=True)
        track = mid.tracks[1]

        # Convert Track to Absolute time
        total_time = 0
        for msg in track:
            total_time += msg.time
            msg.time = total_time

        # Account for padding?
        print(total_time)
        total_time = mido.second2tick(padded_wav_length // SAMPLE_RATE, mid.ticks_per_beat, 500000)
        print(total_time)

        # Split into seconds and convert into array
        # of (time_steps, 88)
        start_idx = 0
        seconds_elapsed = 0.0
        current_time = 0.0
        end_idx = 1
        state = np.zeros(88)

        # Compute time chunks in 'tic-time'
        BLOCK_IN_TICKS = mido.second2tick(BLOCK_IN_SECONDS, mid.ticks_per_beat, 500000)
        print(BLOCK_IN_TICKS)
        time_chunks = np.arange(0, total_time, BLOCK_IN_TICKS, dtype=np.float32)
        time_chunks = np.append(time_chunks, total_time)

        print(len(time_chunks[1:]))
        print(len(chunks))
        block_zones = np.arange(-BLOCK_IN_TICKS, 0, BLOCK_IN_TICKS / (BLOCK_N_QUANT))[1:]
        block_zones = np.append(block_zones, 0)
        block_length = len(block_zones)+1
        state = np.zeros((88))
        for end_point in time_chunks[1:]:
            block = []#np.zeros((len(block_zones)+1, 88))
            while end_idx < len(track) and track[end_idx].time < end_point:
                end_idx += 1
                continue

            for zone in block_zones:
                while start_idx < end_idx and track[start_idx].time < (end_point + zone):
                    if not track[start_idx].type == 'note_on':
                        start_idx += 1
                        continue

                    state[(track[start_idx].note - 21)] = (track[start_idx].velocity / 127)
                    start_idx += 1
                block.append(copy.deepcopy(state))

            midi_data.append(block)

    data = np.array(data)
    midi_data = np.array(midi_data)
    print(data.shape)
    print(midi_data.shape)

    print('Dataset Stats:')
    print('Total Hours: ', len(data) / 60 / 60)
    print('Dataset Size: ', sys.getsizeof(data) / 1e9, 'GB')
    print('Dataset Shape: ', data.shape)

    print("Saving Waveform Dataset")
    np.savez_compressed(os.path.join(PROCESSED_DATA_PATH, 'MAESTRO_{}h.npz'\
                        .format(APPROX_TOTAL_HOURS)),
                        np.array(data))

if __name__ == '__main__':
    main()
<commit_msg>Fix Saving Bug Add code to handle saving the loaded midi data<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
#
#    https://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.

"""Pre-Processes the MAESTRO Dataset.

Pre-proceses the raw MAESTRO dataset into a data set
of music chunks with a predefined lenth and sampling rate.
The music chunks are then saved as a .npz file
"""

import os
import sys
import glob
import librosa
import mido
import copy
import numpy as np

# Overall Config #
APPROX_TOTAL_HOURS = 6
SAMPLE_RATE = 16000 # 16kHz
DATA_POINT_LENGTH = 2**14
BLOCK_IN_SECONDS = DATA_POINT_LENGTH / SAMPLE_RATE
BLOCK_N_QUANT = 128
BLOCK_ZONE_LENGTH = BLOCK_IN_SECONDS / BLOCK_N_QUANT
RAW_DATA_PATH = './data/maestro/2017'
PROCESSED_DATA_PATH = './data/'

def main():
    audio_paths = glob.glob(RAW_DATA_PATH + '/**/*.wav', recursive=True)

    # Load audio files until we reach our desired
    # data set size.
    data = []
    midi_data = []
    hours_loaded = 0
    for audio_file_path in audio_paths:
        print(audio_file_path)
        wav, _ = librosa.load(audio_file_path, sr=SAMPLE_RATE)
        print(len(wav) / 16000)

        hours_loaded += len(wav) / SAMPLE_RATE / 60 / 60
        if hours_loaded >= APPROX_TOTAL_HOURS:
            break

        # Pad the song to ensure it can be evenly divided into
        # chunks of length 'DATA_POINT_LENGTH'
        padded_wav_length = int(DATA_POINT_LENGTH * np.ceil(len(wav) / DATA_POINT_LENGTH))
        wav = np.pad(wav, [[0, padded_wav_length - len(wav)]])

        chunks = np.reshape(wav, (-1, DATA_POINT_LENGTH))
        data.extend(chunks)

        # Load and process MIDI file
        midi_file_path = audio_file_path[0:-3] + 'midi'#.replace('.wav', '.midi')
        print(midi_file_path)
        mid = mido.MidiFile(midi_file_path, clip=True)
        track = mid.tracks[1]

        # Convert Track to Absolute time
        total_time = 0
        for msg in track:
            total_time += msg.time
            msg.time = total_time

        # Account for padding?
        print(total_time)
        total_time = mido.second2tick(padded_wav_length // SAMPLE_RATE, mid.ticks_per_beat, 500000)
        print(total_time)

        # Split into seconds and convert into array
        # of (time_steps, 88)
        start_idx = 0
        seconds_elapsed = 0.0
        current_time = 0.0
        end_idx = 1
        state = np.zeros(88)

        # Compute time chunks in 'tic-time'
        BLOCK_IN_TICKS = mido.second2tick(BLOCK_IN_SECONDS, mid.ticks_per_beat, 500000)
        print(BLOCK_IN_TICKS)
        time_chunks = np.arange(0, total_time, BLOCK_IN_TICKS, dtype=np.float32)
        time_chunks = np.append(time_chunks, total_time)

        print(len(time_chunks[1:]))
        print(len(chunks))
        block_zones = np.arange(-BLOCK_IN_TICKS, 0, BLOCK_IN_TICKS / (BLOCK_N_QUANT))[1:]
        block_zones = np.append(block_zones, 0)
        block_length = len(block_zones)+1
        state = np.zeros((88))
        for end_point in time_chunks[1:]:
            block = []#np.zeros((len(block_zones)+1, 88))
            while end_idx < len(track) and track[end_idx].time < end_point:
                end_idx += 1
                continue

            for zone in block_zones:
                while start_idx < end_idx and track[start_idx].time < (end_point + zone):
                    if not track[start_idx].type == 'note_on':
                        start_idx += 1
                        continue

                    state[(track[start_idx].note - 21)] = (track[start_idx].velocity / 127)
                    start_idx += 1
                block.append(copy.deepcopy(state))

            midi_data.append(block)

    data = np.array(data)
    midi_data = np.array(midi_data)
    print(data.shape)
    print(midi_data.shape)

    print('Dataset Stats:')
    print('Total Hours: ', len(data) / 60 / 60)
    print('Dataset Size: ', sys.getsizeof(data) / 1e9, 'GB')
    print('Dataset Shape: ', data.shape)

    print("Saving Waveform Dataset")
    np.savez_compressed(os.path.join(PROCESSED_DATA_PATH, 'MAESTRO_{}h.npz'\
                        .format(APPROX_TOTAL_HOURS)),
                        np.array(data))
    np.savez_compressed(os.path.join(PROCESSED_DATA_PATH, 'MAESTRO_midi_{}h.npz'\
                        .format(APPROX_TOTAL_HOURS)),
                        np.array(midi_data))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Mon Nov  3 13:28:07 2014
A script to re-format a couple of files that were given me by Margie into the adequate format.
These scripts are SWAG1 and SWAG2 gene expression profiles.
@author: davidaalbores
"""
from __future__ import division, print_function, absolute_import
import pandas as pd

#Load the gene name dictionary
names= pd.read_csv("../input/c_elegans.PRJNA13758.WS241.livegeneIDs.unmaprm.txt", sep= "\t", comment= "#")

#Load the control database
control= pd.read_csv("../input/AVAcontrol.txt", sep= "\t", comment= "#")
control.columns= ['gene', 'reads']
##Experimental
#experimental= pd.read_csv("../input/glp-4dissectedgut(SWAG1).txt", sep= "\t", comment= "#")

#Discard the first line of the dataframes; it's useless.
#control= control[1:]
#experimental= experimental[1:]

#Make a new dataframe that only contains the counts and the gene name, dropping any NAs
control= control.loc[:, ['gene', 'reads']].dropna()
#experimental= experimental.loc[:, ['SWAG1', 'gene']].dropna()

#change the names into WB IDs
for i in range(0, len(control)):
    #get the gene associated with locus i
    g= control.iloc[i].gene
    if names[names.GeneName == g].empty == False:
       #make a variable x that extracts the worm ID from the name df. this is a
    #dataframe itself!
        print('here')
        x= names[names.GeneName == g].WBID
        #Replace the gene name in the control array
        control= control.replace(g, x.iloc[0])
    else:
        if names[names.WBID == g].empty == True:
            #There's going to be a problem.
            print(names[names.WBID == g])
            control= control.replace(g, float('nan'))
        else:
            print(names[names.WBID == g])

##change the names into WB IDs
#for i in range(0, len(experimental)):
#    #get the gene associated with locus i
#    g= experimental.iloc[i].gene
#    if names[names.GeneName == g].empty == False:
#       #make a variable x that extracts the worm ID from the name df. this is a
#    #dataframe itself!
#        x= names[names.GeneName == g].WBID
#        #Replace the gene name in the experimental array
#        experimental= experimental.replace(g, x.iloc[0])
#    else:
#        #There's going to be a problem.
#        experimental= experimental.replace(g, float('nan'))

control= control.dropna()
#experimental= experimental.dropna()

control.to_csv("../input/controlAVA131014.txt", sep=',')
#experimental.to_csv("experimentalGE.txt", sep=',')

<commit_msg>3rd Commit, FileFormatScript.py improved and debugged<commit_after># -*- coding: utf-8 -*-
"""
Created on Mon Nov  3 13:28:07 2014
A script to re-format a file
@author: davidaalbores
"""
from __future__ import division, print_function, absolute_import
import pandas as pd

#Filenames!
fileInput= '../input/glp-4dissectedgut(SWAG1).txt'; sepInput= '\t'; comInput= '#';
newFile= '../input/experimentalGE.txt'
errorFile= '../output/ErrorFile.txt'
summaryFile= '../output/Summary.txt'

#Filename and location of the dataframe you want to translate,
#kind of separator used and comment marker
filename= '../input/c_elegans.PRJNA13758.WS241.livegeneIDs.unmaprm.txt'
separator= "\t"
comments= "#"

#Load the gene name dictionary
names= pd.read_csv(filename, sep= separator, comment= comments)

#Load the control database
df= pd.read_csv(fileInput, sep= sepInput, comment= comInput)
df= df.loc[:, ['gene', 'reads'] ]
#df.columns= ['gene', 'reads']

nInitial= len(df)
nulls= []

#change the names into WB IDs
for i in range(0, len(df)):
    #get the gene associated with locus i
    g= df.iloc[i].gene
    if names[names.GeneName == g].empty == False:
        #make a variable x that extracts the worm ID from the name df. this is a
        #dataframe itself!
        x= names[names.GeneName == g].WBID
        #Replace the gene name in the control array
        df= df.replace(g, x.iloc[0])
    else:
        if names[names.WBID == g].empty == True:
            if names[names.HumanReadable == g].empty == False and g.lower() != 'nan':
                x= names[names.HumanReadable == g].WBID
                df= df.replace(g, x)
            else:
                #There's going to be a problem.
                nulls.append([g, df.iloc[0].reads])
                df= df.replace(g, float('nan'))

df= df.dropna()

#Write to main file
df.to_csv(newFile, sep=',', index= False)

#Write error file
f= open(errorFile, 'w')
f.write("#This file stores any genes that could not be translated into a WBID\n")
f.write("gene_Name\,reads\n")
for element in nulls:
    string= str(element[0])+","+str(element[1])+"\n"
    f.write(string)
f.close()

#Write summary
f1= open(summaryFile, 'w')
f1.write("#This file gives a summary of the results of running FileFormatScript.py\n")
n= len(df)
f1.write("Original list contained {0} entries. List with WBIDs contains {1} entries".format(nInitial, n))
f1.close()

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Columns StreamField block

This module provides a columns block for use with the wagtail StreamField
class.

TODO:
* Limit the row to a maximum of 100% width
"""

from django import forms
from django.template.loader import render_to_string
from django.core.exceptions import NON_FIELD_ERRORS
from django.contrib.staticfiles.templatetags.staticfiles import static
from django.utils.html import format_html, format_html_join

from wagtail.wagtailcore.blocks.utils import js_dict
from wagtail.wagtailcore import blocks
from wagtail.wagtailembeds.blocks import EmbedBlock
from wagtail.wagtailimages.blocks import ImageChooserBlock
from wagtail.contrib.table_block.blocks import TableBlock

from common.blocks import COLOUR_CHOICES, COLUMN_CHOICES
from common.blocks.maps import GoogleMapBlock
from common.blocks.twitter import TwitterBlock
from common.blocks.images import ImageBlock
from common.blocks.centered_text import CenteredTextBlock
from common.blocks.images import COSPhotoStreamBlock

class GenericContentStreamBlock(blocks.StreamBlock):

    def __init__(self, local_blocks=None, **kwargs):

        default_blocks = [
            ('appeal', blocks.StructBlock([
                        ('icon', blocks.ChoiceBlock(required=True, choices=[
                            ('none', 'none'),
                            ('flask', 'flask'),
                            ('group', 'group'),
                            ('laptop', 'laptop'),
                            ('sitemap', 'sitemap'),
                            ('user', 'user'),
                            ('book', 'book'),
                            ('download', 'download'),
                        ])),
                ('topic', blocks.CharBlock(required=True, max_length=35)),
                ('content', blocks.TextBlock(required=True, max_length=255)),
            ], classname='appeal', icon='tick', template='common/blocks/appeal.html')),
            ('heading', blocks.CharBlock(classname="full title")),
            ('statement', blocks.CharBlock()),
            ('twitter', TwitterBlock()),
            ('photo_stream', COSPhotoStreamBlock()),
            ('rich_text', blocks.RichTextBlock()),
            ('paragraph', blocks.TextBlock()),
            ('map', GoogleMapBlock()),
            ('customizedimage', ImageBlock()),
            ('raw_html', blocks.RawHTMLBlock(help_text='With great power comes great responsibility. This HTML is unescaped. Be careful!')),
            ('centered_text', CenteredTextBlock()),
            ('embed_block', EmbedBlock()),
            ('table_block', TableBlock())
        ]

        if local_blocks:
            default_blocks = default_blocks + local_blocks

        super(GenericContentStreamBlock, self).__init__(local_blocks=default_blocks, **kwargs)

    def render_form(self, value, prefix='', errors=None):
        error_dict = {}
        if errors:
            if len(errors) > 1:
                # We rely on StreamBlock.clean throwing a single
                # StreamBlockValidationError with a specially crafted 'params'
                # attribute that we can pull apart and distribute to the child
                # blocks
                raise TypeError('StreamBlock.render_form unexpectedly received multiple errors')
            error_dict = errors.as_data()[0].params

        # value can be None when the StreamField is in a formset
        if value is None:
            value = self.get_default()
        # drop any child values that are an unrecognised block type
        valid_children = [child for child in value if child.block_type in self.child_blocks]

        list_members_html = [
            self.render_list_member(child.block_type, child.value, "%s-%d" % (prefix, i), i,
                                    errors=error_dict.get(i))
            for (i, child) in enumerate(valid_children)
        ]

        return render_to_string('common/block_forms/column_content.html', {
            'prefix': prefix,
            'list_members_html': list_members_html,
            'child_blocks': self.child_blocks.values(),
            'header_menu_prefix': '%s-before' % prefix,
            'block_errors': error_dict.get(NON_FIELD_ERRORS),
        })

    class Meta:
        icon = 'arrow-left'
        label = 'Content'
        classname = ''
        form_template = 'common/block_forms/column_content.html'

class ColumnBlock(blocks.StructBlock):

    column_size = blocks.ChoiceBlock(choices=COLUMN_CHOICES, default="12")
    content = GenericContentStreamBlock()

    class Meta:
        form_template = 'common/block_forms/column.html'
        template = 'common/blocks/column_block.html'
        icon = 'placeholder'
        label = 'Column'

class ColumnsBlock(blocks.ListBlock):

    class Meta:
        form_template = 'common/block_forms/columns.html'
        icon = 'placeholder'
        label = 'Columns'

    @property
    def media(self):
        return forms.Media(js=[static('wagtailadmin/js/blocks/sequence.js'), static('wagtailadmin/js/blocks/list.js'), static('common/js/blocks/columns.js')])

    def __init__(self, **kwargs):
        return super(ColumnsBlock, self).__init__(ColumnBlock(), **kwargs)

    def render_basic(self, value, context=None):
        children = format_html_join(
            '\n', '{0}',
            [
                (self.child_block._render_with_context(child_value, context=context),)
                for child_value in value
            ]
        )
        return format_html("
{0}
", children)

    def render_form(self, value, prefix='', errors=None):
        if errors:
            if len(errors) > 1:
                # We rely on ListBlock.clean throwing a single ValidationError with a specially crafted
                # 'params' attribute that we can pull apart and distribute to the child blocks
                raise TypeError('ListBlock.render_form unexpectedly received multiple errors')
            error_list = errors.as_data()[0].params
        else:
            error_list = None

        list_members_html = [
            self.render_list_member(child_val, "%s-%d" % (prefix, i), i,
                                    errors=error_list[i] if error_list else None)
            for (i, child_val) in enumerate(value)
        ]

        return render_to_string('common/block_forms/columns.html', {
            'help_text': getattr(self.meta, 'help_text', None),
            'prefix': prefix,
            'list_members_html': list_members_html,
        })

    def js_initializer(self):
        opts = {'definitionPrefix': "'%s'" % self.definition_prefix}
        if self.child_js_initializer:
            opts['childInitializer'] = self.child_js_initializer
        return "ColumnsBlock(%s)" % js_dict(opts)

<commit_msg>ad margin to bottom of rows<commit_after># -*- coding: utf-8 -*-
"""Columns StreamField block

This module provides a columns block for use with the wagtail StreamField
class.

TODO:
* Limit the row to a maximum of 100% width
"""

from django import forms
from django.template.loader import render_to_string
from django.core.exceptions import NON_FIELD_ERRORS
from django.contrib.staticfiles.templatetags.staticfiles import static
from django.utils.html import format_html, format_html_join

from wagtail.wagtailcore.blocks.utils import js_dict
from wagtail.wagtailcore import blocks
from wagtail.wagtailembeds.blocks import EmbedBlock
from wagtail.wagtailimages.blocks import ImageChooserBlock
from wagtail.contrib.table_block.blocks import TableBlock

from common.blocks import COLOUR_CHOICES, COLUMN_CHOICES
from common.blocks.maps import GoogleMapBlock
from common.blocks.twitter import TwitterBlock
from common.blocks.images import ImageBlock
from common.blocks.centered_text import CenteredTextBlock
from common.blocks.images import COSPhotoStreamBlock

class GenericContentStreamBlock(blocks.StreamBlock):

    def __init__(self, local_blocks=None, **kwargs):

        default_blocks = [
            ('appeal', blocks.StructBlock([
                        ('icon', blocks.ChoiceBlock(required=True, choices=[
                            ('none', 'none'),
                            ('flask', 'flask'),
                            ('group', 'group'),
                            ('laptop', 'laptop'),
                            ('sitemap', 'sitemap'),
                            ('user', 'user'),
                            ('book', 'book'),
                            ('download', 'download'),
                        ])),
                ('topic', blocks.CharBlock(required=True, max_length=35)),
                ('content', blocks.TextBlock(required=True, max_length=255)),
            ], classname='appeal', icon='tick', template='common/blocks/appeal.html')),
            ('heading', blocks.CharBlock(classname="full title")),
            ('statement', blocks.CharBlock()),
            ('twitter', TwitterBlock()),
            ('photo_stream', COSPhotoStreamBlock()),
            ('rich_text', blocks.RichTextBlock()),
            ('paragraph', blocks.TextBlock()),
            ('map', GoogleMapBlock()),
            ('customizedimage', ImageBlock()),
            ('raw_html', blocks.RawHTMLBlock(help_text='With great power comes great responsibility. This HTML is unescaped. Be careful!')),
            ('centered_text', CenteredTextBlock()),
            ('embed_block', EmbedBlock()),
            ('table_block', TableBlock())
        ]

        if local_blocks:
            default_blocks = default_blocks + local_blocks

        super(GenericContentStreamBlock, self).__init__(local_blocks=default_blocks, **kwargs)

    def render_form(self, value, prefix='', errors=None):
        error_dict = {}
        if errors:
            if len(errors) > 1:
                # We rely on StreamBlock.clean throwing a single
                # StreamBlockValidationError with a specially crafted 'params'
                # attribute that we can pull apart and distribute to the child
                # blocks
                raise TypeError('StreamBlock.render_form unexpectedly received multiple errors')
            error_dict = errors.as_data()[0].params

        # value can be None when the StreamField is in a formset
        if value is None:
            value = self.get_default()
        # drop any child values that are an unrecognised block type
        valid_children = [child for child in value if child.block_type in self.child_blocks]

        list_members_html = [
            self.render_list_member(child.block_type, child.value, "%s-%d" % (prefix, i), i,
                                    errors=error_dict.get(i))
            for (i, child) in enumerate(valid_children)
        ]

        return render_to_string('common/block_forms/column_content.html', {
            'prefix': prefix,
            'list_members_html': list_members_html,
            'child_blocks': self.child_blocks.values(),
            'header_menu_prefix': '%s-before' % prefix,
            'block_errors': error_dict.get(NON_FIELD_ERRORS),
        })

    class Meta:
        icon = 'arrow-left'
        label = 'Content'
        classname = ''
        form_template = 'common/block_forms/column_content.html'

class ColumnBlock(blocks.StructBlock):

    column_size = blocks.ChoiceBlock(choices=COLUMN_CHOICES, default="12")
    content = GenericContentStreamBlock()

    class Meta:
        form_template = 'common/block_forms/column.html'
        template = 'common/blocks/column_block.html'
        icon = 'placeholder'
        label = 'Column'

class ColumnsBlock(blocks.ListBlock):

    class Meta:
        form_template = 'common/block_forms/columns.html'
        icon = 'placeholder'
        label = 'Columns'

    @property
    def media(self):
        return forms.Media(js=[static('wagtailadmin/js/blocks/sequence.js'), static('wagtailadmin/js/blocks/list.js'), static('common/js/blocks/columns.js')])

    def __init__(self, **kwargs):
        return super(ColumnsBlock, self).__init__(ColumnBlock(), **kwargs)

    def render_basic(self, value, context=None):
        children = format_html_join(
            '\n', '{0}',
            [
                (self.child_block._render_with_context(child_value, context=context),)
                for child_value in value
            ]
        )
        return format_html("
{0}
", children)

    def render_form(self, value, prefix='', errors=None):
        if errors:
            if len(errors) > 1:
                # We rely on ListBlock.clean throwing a single ValidationError with a specially crafted
                # 'params' attribute that we can pull apart and distribute to the child blocks
                raise TypeError('ListBlock.render_form unexpectedly received multiple errors')
            error_list = errors.as_data()[0].params
        else:
            error_list = None

        list_members_html = [
            self.render_list_member(child_val, "%s-%d" % (prefix, i), i,
                                    errors=error_list[i] if error_list else None)
            for (i, child_val) in enumerate(value)
        ]

        return render_to_string('common/block_forms/columns.html', {
            'help_text': getattr(self.meta, 'help_text', None),
            'prefix': prefix,
            'list_members_html': list_members_html,
        })

    def js_initializer(self):
        opts = {'definitionPrefix': "'%s'" % self.definition_prefix}
        if self.child_js_initializer:
            opts['childInitializer'] = self.child_js_initializer
        return "ColumnsBlock(%s)" % js_dict(opts)

<|endoftext|>


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

from omnilib.frameworks.framework_base import Framework_Base
from omnilib.util.dossl import _do_ssl
import omnilib.util.credparsing as credutils

from geni.util.urn_util import is_valid_urn, URN, string_to_urn_format

import os
import string
import sys

class Framework(Framework_Base):
    def __init__(self, config, opts):
        Framework_Base.__init__(self,config)
        config['cert'] = os.path.expanduser(config['cert'])
        if not os.path.exists(config['cert']):
            sys.exit('CHAPI Framework certfile %s doesnt exist' % config['cert'])
        if not os.path.getsize(config['cert']) > 0:
            sys.exit('CHAPI Framework certfile %s is empty' % config['cert'])
        config['key'] = os.path.expanduser(config['key'])
        if not os.path.exists(config['key']):
            sys.exit('CHAPI Framework keyfile %s doesnt exist' % config['key'])
        if not os.path.getsize(config['key']) > 0:
            sys.exit('CHAPI Framework keyfile %s is empty' % config['key'])
        if not config.has_key('verbose'):
            config['verbose'] = False
        self.config = config

        self.ch = self.make_client(config['ch']+'/CH', self.key, self.cert,
                                   verbose=config['verbose'])
        self.ma = self.make_client(config['ch']+'/MA', self.key, self.cert,
                                   verbose=config['verbose'])
        self.sa = self.make_client(config['ch']+'/SA', self.key, self.cert,
                                   verbose=config['verbose'])

        self.cert_string = file(config['cert'],'r').read()
        self.user_cred = self.init_user_cred( opts )
        self.logger = config['logger']

    def get_user_cred(self):
        message = ""
        if self.user_cred == None:
            self.logger.debug("Getting user credential from CHAPI MA %s", self.config['ch'])
            # TODO: supply empty filter/query critera now. Should we do getversion so that we can filter on nothing and query on *?
            (self.user_cred, message) = _do_ssl(self, None, ("lookup public member info on CHAPI CH %s" % self.config['ch']),
                                                self.ma.lookup_public_member_info,
                                                self.cert_string,
                                                {'match':{'USER_URN':[]}})

        return self.user_cred, message

    def get_slice_cred(self, urn):
        scred = '' #TODO: where do we get the cred to pass in??
        (cred, message) = _do_ssl(self, None, ("Get slice credentials  %s on CHAPI SA %s" % (urn, self.config['ch'])),
                                  self.sa.lookup_slices,
                                  scred,
                                  {'match':{'SLICE_URN':[urn]}})
        # FIXME: use any message?
        _ = message #Appease eclipse
        return cred

    def create_slice(self, urn):
        (cred, message) = _do_ssl(self, None, ("Create slice %s on CHAPI SA %s" % (urn, self.config['ch'])), self.sa.create_slice, urn)
        # FIXME: use any message?
        _ = message #Appease eclipse
        return cred

    def delete_slice(self, urn):
        self.logger.error("CHAPI SA does not allow delete_slice")
        return False

    def list_aggregates(self):
        # TODO: list of field names from getVersion - should we get all or assume we have URN and URL
        (res, message) = _do_ssl(self, None, ("List Aggregates at CHAPI CH %s" % self.config['ch']),
                                 self.ch.get_aggregates, ['SERVICE_URN', 'SERVICE_URL'])

        aggs = dict()
        if res['value'] is not None:
            for (urn, url) in res['value']:
                aggs[urn] = url

        return aggs

    def list_my_slices(self, user):
        '''List slices owned by the user (name or URN) provided, returning a list of slice URNs.'''

        if user is None or user.strip() == '':
            raise Exception('Empty user name')

        # construct a urn from that user
        if is_valid_urn(user):
            # FIXME: Validate type, authority?
            userurn = user
        else:
            if not self.config.has_key('authority'):
                raise Exception("Invalid configuration: no authority defined")

            auth = self.config['authority']
            userurn = URN(auth, "user", user).urn_string()

        # TODO: sa.lookup_slices (username), then ma.lookup_public/private_member_info

        # invoke ListMySlices(urn)
        (slices, message) = _do_ssl(self, None, ("List Slices for %s at CHAPI SA %s" % (user, self.config['ch'])),
                                    self.sa.ListMySlices, userurn)
        # FIXME: use any message?
        _ = message #Appease eclipse

#        # Return is a urn. Strip out the name
#        slicenames = list()
#        if slices and isinstance(slices, list):
#            for slice in slices:
#                slicelower = string.lower(slice)
#                if not string.find(slicelower, "+slice+"):
#                    continue
#                slicename = slice[string.index(slicelower,"+slice+") + len("+slice+"):]
#                slicenames.append(slicename)
#        return slicenames

        return slices

    def slice_name_to_urn(self, name):
        """Convert a slice name to a slice urn."""

        if name is None or name.strip() == '':
            raise Exception('Empty slice name')

        if is_valid_urn(name):
            urn = URN(None, None, None, name)
            if not urn.getType() == "slice":
                raise Exception("Invalid Slice name: got a non Slice URN %s", name)
            # if config has an authority, make sure it matches
            if self.config.has_key('authority'):
                auth = self.config['authority']
                urn_fmt_auth = string_to_urn_format(urn.getAuthority())
                if urn_fmt_auth != auth:
                    self.logger.warn("CAREFUL: slice' authority (%s) doesn't match current configured authority (%s)" % (urn_fmt_auth, auth))
                    self.logger.info("This may be OK though if you are using delegated slice credentials...")
#                    raise Exception("Invalid slice name: slice' authority (%s) doesn't match current configured authority (%s)" % (urn_fmt_auth, auth))
            return name

        if not self.config.has_key('authority'):
            raise Exception("Invalid configuration: no authority defined")

        auth = self.config['authority']
        return URN(auth, "slice", name).urn_string()

    def renew_slice(self, urn, expiration_dt):
        """See framework_base for doc.
        """
        expiration = expiration_dt.isoformat()
        cred = self.get_slice_cred(urn)
        (bool, message) = _do_ssl(self, None, ("Renew slice %s on CHAPI SA %s until %s" % (urn, self.config['ch'], expiration_dt)),
                                  self.sa.update_slice, urn, cred, {'SLICE_EXPIRATION':expiration})
        if bool:
            return expiration_dt
        else:
            # FIXME: use any message?
            _ = message #Appease eclipse
            return None

    def get_user_cred_struct(self):
        """
        Returns a user credential from the control framework as a string in a struct. And an error message if any.
        Struct is as per AM API v3:
        {
           geni_type: <string>,
           geni_version: <string>,
           geni_value: <the credential as a string>
        }
        """
        cred, message = self.get_user_cred()
        if cred:
            cred = self.wrap_cred(cred)
        return cred, message

    def get_slice_cred_struct(self, urn):
        """
        Retrieve a slice with the given urn and returns the signed
        credential as a string in the AM API v3 struct:
        {
           geni_type: <string>,
           geni_version: <string>,
           geni_value: <the credential as a string>
        }
        """
        cred = self.get_slice_cred(urn)
        return self.wrap_cred(cred)

    def wrap_cred(self, cred):
        """
        Wrap the given cred in the appropriate struct for this framework.
        """
        if isinstance(cred, dict):
            self.logger.warn("Called wrap on a cred that's already a dict? %s", cred)
            return cred
        elif not isinstance(cred, str):
            self.logger.warn("Called wrap on non string cred? Stringify. %s", cred)
            cred = str(cred)
        ret = dict(geni_type="geni_sfa", geni_version="2", geni_value=cred)
        if credutils.is_valid_v3(self.logger, cred):
            ret["geni_version"] = "3"
        return ret

    def get_version(self):
        pl_response = dict()
        versionstruct = dict()
        (pl_response, message) = _do_ssl(self, None, ("GetVersion of CHAPI CH %s using cert %s" % (self.config['ch'], self.config['cert'])), self.ch.get_version)
        _ = message #Appease eclipse
        if pl_response is None:
            self.logger.error("Failed to get version of CHAPI CH: %s", message)
            # FIXME: Return error message?
            return None, message
        if isinstance(pl_response, dict) and pl_response.has_key('code'):
            code = pl_response['code']
            if code:
                self.logger.error("Failed to get version of CHAPI CH: Received error code: %d", code)
                output = pl_response['output']
                self.logger.error("Received error message: %s", output)
            else:
                versionstruct = pl_response['value']
        else:
            versionstruct = pl_response
        return versionstruct, message
<commit_msg>Add comment for listmyslices<commit_after>#----------------------------------------------------------------------
# Copyright (c) 2011-2013 Raytheon BBN Technologies
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and/or hardware specification (the "Work") to
# deal in the Work without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Work, and to permit persons to whom the Work
# is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Work.
#
# THE WORK IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE WORK OR THE USE OR OTHER DEALINGS
# IN THE WORK.
#----------------------------------------------------------------------

from omnilib.frameworks.framework_base import Framework_Base
from omnilib.util.dossl import _do_ssl
import omnilib.util.credparsing as credutils

from geni.util.urn_util import is_valid_urn, URN, string_to_urn_format

import os
import string
import sys

class Framework(Framework_Base):
    def __init__(self, config, opts):
        Framework_Base.__init__(self,config)
        config['cert'] = os.path.expanduser(config['cert'])
        if not os.path.exists(config['cert']):
            sys.exit('CHAPI Framework certfile %s doesnt exist' % config['cert'])
        if not os.path.getsize(config['cert']) > 0:
            sys.exit('CHAPI Framework certfile %s is empty' % config['cert'])
        config['key'] = os.path.expanduser(config['key'])
        if not os.path.exists(config['key']):
            sys.exit('CHAPI Framework keyfile %s doesnt exist' % config['key'])
        if not os.path.getsize(config['key']) > 0:
            sys.exit('CHAPI Framework keyfile %s is empty' % config['key'])
        if not config.has_key('verbose'):
            config['verbose'] = False
        self.config = config

        self.ch = self.make_client(config['ch']+'/CH', self.key, self.cert,
                                   verbose=config['verbose'])
        self.ma = self.make_client(config['ch']+'/MA', self.key, self.cert,
                                   verbose=config['verbose'])
        self.sa = self.make_client(config['ch']+'/SA', self.key, self.cert,
                                   verbose=config['verbose'])

        self.cert_string = file(config['cert'],'r').read()
        self.user_cred = self.init_user_cred( opts )
        self.logger = config['logger']

    def get_user_cred(self):
        message = ""
        if self.user_cred == None:
            self.logger.debug("Getting user credential from CHAPI MA %s", self.config['ch'])
            # TODO: supply empty filter/query critera now. Should we do getversion so that we can filter on nothing and query on *?
            (self.user_cred, message) = _do_ssl(self, None, ("lookup public member info on CHAPI CH %s" % self.config['ch']),
                                                self.ma.lookup_public_member_info,
                                                self.cert_string,
                                                {'match':{'USER_URN':[]}})

        return self.user_cred, message

    def get_slice_cred(self, urn):
        scred = '' #TODO: where do we get the cred to pass in??
        (cred, message) = _do_ssl(self, None, ("Get slice credentials  %s on CHAPI SA %s" % (urn, self.config['ch'])),
                                  self.sa.lookup_slices,
                                  scred,
                                  {'match':{'SLICE_URN':[urn]}})
        # FIXME: use any message?
        _ = message #Appease eclipse
        return cred

    def create_slice(self, urn):
        (cred, message) = _do_ssl(self, None, ("Create slice %s on CHAPI SA %s" % (urn, self.config['ch'])), self.sa.create_slice, urn)
        # FIXME: use any message?
        _ = message #Appease eclipse
        return cred

    def delete_slice(self, urn):
        self.logger.error("CHAPI SA does not allow delete_slice")
        return False

    def list_aggregates(self):
        # TODO: list of field names from getVersion - should we get all or assume we have URN and URL
        (res, message) = _do_ssl(self, None, ("List Aggregates at CHAPI CH %s" % self.config['ch']),
                                 self.ch.get_aggregates, ['SERVICE_URN', 'SERVICE_URL'])

        aggs = dict()
        if res['value'] is not None:
            for (urn, url) in res['value']:
                aggs[urn] = url

        return aggs

    def list_my_slices(self, user):
        '''List slices owned by the user (name or URN) provided, returning a list of slice URNs.'''

        if user is None or user.strip() == '':
            raise Exception('Empty user name')

        # construct a urn from that user
        if is_valid_urn(user):
            # FIXME: Validate type, authority?
            userurn = user
        else:
            if not self.config.has_key('authority'):
                raise Exception("Invalid configuration: no authority defined")

            auth = self.config['authority']
            userurn = URN(auth, "user", user).urn_string()

        # TODO: sa.lookup_slices (username), then ma.lookup_public/private_member_info
        # TODO: can we filter lookup_slices with username directly, or do we need to go back to MA?

        # invoke ListMySlices(urn)
        (slices, message) = _do_ssl(self, None, ("List Slices for %s at CHAPI SA %s" % (user, self.config['ch'])),
                                    self.sa.ListMySlices, userurn)
        # FIXME: use any message?
        _ = message #Appease eclipse

#        # Return is a urn. Strip out the name
#        slicenames = list()
#        if slices and isinstance(slices, list):
#            for slice in slices:
#                slicelower = string.lower(slice)
#                if not string.find(slicelower, "+slice+"):
#                    continue
#                slicename = slice[string.index(slicelower,"+slice+") + len("+slice+"):]
#                slicenames.append(slicename)
#        return slicenames

        return slices

    def slice_name_to_urn(self, name):
        """Convert a slice name to a slice urn."""

        if name is None or name.strip() == '':
            raise Exception('Empty slice name')

        if is_valid_urn(name):
            urn = URN(None, None, None, name)
            if not urn.getType() == "slice":
                raise Exception("Invalid Slice name: got a non Slice URN %s", name)
            # if config has an authority, make sure it matches
            if self.config.has_key('authority'):
                auth = self.config['authority']
                urn_fmt_auth = string_to_urn_format(urn.getAuthority())
                if urn_fmt_auth != auth:
                    self.logger.warn("CAREFUL: slice' authority (%s) doesn't match current configured authority (%s)" % (urn_fmt_auth, auth))
                    self.logger.info("This may be OK though if you are using delegated slice credentials...")
#                    raise Exception("Invalid slice name: slice' authority (%s) doesn't match current configured authority (%s)" % (urn_fmt_auth, auth))
            return name

        if not self.config.has_key('authority'):
            raise Exception("Invalid configuration: no authority defined")

        auth = self.config['authority']
        return URN(auth, "slice", name).urn_string()

    def renew_slice(self, urn, expiration_dt):
        """See framework_base for doc.
        """
        expiration = expiration_dt.isoformat()
        cred = self.get_slice_cred(urn)
        (bool, message) = _do_ssl(self, None, ("Renew slice %s on CHAPI SA %s until %s" % (urn, self.config['ch'], expiration_dt)),
                                  self.sa.update_slice, urn, cred, {'SLICE_EXPIRATION':expiration})
        if bool:
            return expiration_dt
        else:
            # FIXME: use any message?
            _ = message #Appease eclipse
            return None

    def get_user_cred_struct(self):
        """
        Returns a user credential from the control framework as a string in a struct. And an error message if any.
        Struct is as per AM API v3:
        {
           geni_type: <string>,
           geni_version: <string>,
           geni_value: <the credential as a string>
        }
        """
        cred, message = self.get_user_cred()
        if cred:
            cred = self.wrap_cred(cred)
        return cred, message

    def get_slice_cred_struct(self, urn):
        """
        Retrieve a slice with the given urn and returns the signed
        credential as a string in the AM API v3 struct:
        {
           geni_type: <string>,
           geni_version: <string>,
           geni_value: <the credential as a string>
        }
        """
        cred = self.get_slice_cred(urn)
        return self.wrap_cred(cred)

    def wrap_cred(self, cred):
        """
        Wrap the given cred in the appropriate struct for this framework.
        """
        if isinstance(cred, dict):
            self.logger.warn("Called wrap on a cred that's already a dict? %s", cred)
            return cred
        elif not isinstance(cred, str):
            self.logger.warn("Called wrap on non string cred? Stringify. %s", cred)
            cred = str(cred)
        ret = dict(geni_type="geni_sfa", geni_version="2", geni_value=cred)
        if credutils.is_valid_v3(self.logger, cred):
            ret["geni_version"] = "3"
        return ret

    def get_version(self):
        pl_response = dict()
        versionstruct = dict()
        (pl_response, message) = _do_ssl(self, None, ("GetVersion of CHAPI CH %s using cert %s" % (self.config['ch'], self.config['cert'])), self.ch.get_version)
        _ = message #Appease eclipse
        if pl_response is None:
            self.logger.error("Failed to get version of CHAPI CH: %s", message)
            # FIXME: Return error message?
            return None, message
        if isinstance(pl_response, dict) and pl_response.has_key('code'):
            code = pl_response['code']
            if code:
                self.logger.error("Failed to get version of CHAPI CH: Received error code: %d", code)
                output = pl_response['output']
                self.logger.error("Received error message: %s", output)
            else:
                versionstruct = pl_response['value']
        else:
            versionstruct = pl_response
        return versionstruct, message
<|endoftext|>


<commit_before>"""
Copyright 2013 by National University of Singapore

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.
"""
import numpy

class LDA:
    """Latent Dirichlet Allocation Topic Modelling library

    Args:
            Word vectors \w;
        TODO - hyperparameters \alpha, \beta; topic number

    Returns:
            Topic associations \z; multinominal parameters \Phi and \Theta;
            hyperparameter estimates \alpha, \beta

    1) Create a LDA, call LDA.init
    Necessary!

    2) Do inference per round, call LDA.inference
    Necessary!

    3) Calculate perplexity, call LDA.perplexity
    Optional - v2
    """

    def __init__(self, docs, V):
        self.K = 2             # number of topics
        self.alpha = 0.01    # should have been a vector
        self.beta = 0.01        # should have been a vector
        self.seed = 1
        self.docs = docs
        self.M = len(self.docs)

        # topic assignments for words in documents
        self.z_m_n = list()
        self.n_m_k = numpy.zeros(self.M, self.K)
        self.n_m = numpy.zeros(self.M)
        self.n_k_t = numpy.zeros((self.K, V))
        self.n_k = numpy.zeros(self.K)

        numpy.random.seed(self.seed)

        for m, doc in enumerate(docs):
            z_n = numpy.zeros(len(doc))
            for n, word in enumerate(doc):
                # Smart Init sample topic index z_m_n = K from Multinomial(1/K)
                z = numpy.random.randint(0, self.K)
                z_n[n] = z
                # increment document-topic count sum_m_z += 1
                self.n_m_k[m, z] += 1
                # increment document-topic sum sum_m += 1
                self.n_m[m] += 1
                # increment topic-term count sum_z_t += 1
                self.n_k_t[z, n] += 1
                # increment topic-term sum sum_z += 1
                self.n_k[z] += 1
            self.z_m_n.append(z_n)

    def inference(self):
        """
        Gibbs Sampling over burn-in period and sampling period

        This procedure should be run multiple times until reaching a
        fixed number of iterations or other terminating conditions
        (e.g., perplexity)
        """

        for m in range(10):
            for n in range(3):
                zmn = self.z_m_n[m, n]
                self.n_m_k[m, zmn] -= 1
                self.n_m[m] -= 1
                self.n_k[zmn] -= 1
            # end for all words
        # end for all documents
        # end of method inference

    def perplexity(self):
        1                       # placeholder
        # end of method perplexity

# end of class LDA
<commit_msg>Inference Complete<commit_after>"""
Copyright 2013 by National University of Singapore

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.
"""
import numpy

class LDA:
    """Latent Dirichlet Allocation Topic Modelling library

    Args:
            Word vectors \w;
        TODO - hyperparameters \alpha, \beta; topic number

    Returns:
            Topic associations \z; multinominal parameters \Phi and \Theta;
            hyperparameter estimates \alpha, \beta

    1) Create a LDA, call LDA.init
    Necessary!

    2) Do inference per round, call LDA.inference
    Necessary!

    3) Calculate perplexity, call LDA.perplexity
    Optional - v2
    """

    def __init__(self, docs, V):
        self.K = 2             # number of topics
        self.alpha = 0.01    # should have been a vector
        self.beta = 0.01        # should have been a vector
        self.docs = docs
        self.M = len(self.docs)

        # Topics of words of documents.
        self.z_m_n = list()
        # Word count of each document and topic.
        self.n_m_k = numpy.zeros(self.M, self.K)
        # Word count of each document
        self.n_m = numpy.zeros(self.M)
        # Word count of each topic and vocabulary.
        self.n_k_t = numpy.zeros((self.K, V))
        # Word count of each topic.
        self.n_k = numpy.zeros(self.K)

        for m, doc in enumerate(docs):
            # Word count for the current document
            self.n_m[m] = len(doc)
            # Sampling a Multinomial Distribution for the current document
            doc_topic_dist = numpy.random.Dirichlet([self.alpha] * self.K)
            z_n = numpy.zeros(len(doc))
            for n, word in enumerate(doc):
                # Sampling a topic from the sampled Multinomial Distribution
                z = numpy.random.multinomial(1, doc_topic_dist).argmax()
                z_n[n] = z
                # increment document-topic count sum_m_z += 1
                self.n_m_k[m, z] += 1
                # increment topic-term count sum_z_t += 1
                self.n_k_t[z, word] += 1
                # increment topic-term sum sum_z += 1
                self.n_k[z] += 1
            self.z_m_n.append(z_n)

    def inference(self):
        """
        Gibbs Sampling over burn-in period and sampling period

        This procedure should be run multiple times until reaching a
        fixed number of iterations or other terminating conditions
        (e.g., perplexity)
        """

        for m, doc in enumerate(self.docs):
            for n, t in enumerate(doc):
                # Topic assigned to the current word
                zmn = self.z_m_n[m, n]
                # Decrement count from the document-topic matrix
                self.n_m_k[m, zmn] -= 1
                # Decrement count for the topic vocab matrix
                self.n_k_t[zmn, t] -= 1
                # Decrement count for topic vector
                self.n_k[zmn] -= 1
                #self.n_m[m] -= 1
                # Calculating the probability ratios without the current word
                p_z = (self.n_k_t[:, t] + self.beta) * (self.n_m_k[m] +
                                                        self.alpha) / \
                      (self.n_k + self.beta)
                # Feeding these probabilities to the multinomial and sampling a
                # new topic (the prob. values are normalized as well)
                new_z = numpy.random.multinominal(1, p_z / p_z.sum()).argmax()
                # Updating all the variables
                self.z_m_n[m, n] = new_z
                self.n_m_k[m, new_z] += 1
                self.n_k_t[new_z, t] += 1
                self.n_k[new_z] += 1

    def perplexity(self):
        1                       # placeholder
        # end of method perplexity

# end of class LDA
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
This code uses a loop along with our set of coupled differential equations and
matrix math to create arrays of 4-vector quaternions.

The old plotting functions need to be updated and incorperated into the end of
this code or a better visualization solution needs to be found.
"""

#------------------------------------------------------------------------------
#               Importing modules and copying functions
#                   AKA "Setting stuff up"
#------------------------------------------------------------------------------

import numpy as np
from time import time as checktime

# a set of init quaternions and the identity matrix for building general q-matrices
rm = np.identity(2)
im = np.array([[-1j,0],[0,1j]])
jm = np.array([[0,1],[-1,0]])
km = np.array([[0,-1j],[-1j,0]])

def vec_mat(v):
    '''
    Converts a quaternion vector into the 2x2 imaginary matrix representation
    '''
    return v[0]*rm + v[1]*im + v[2]*jm + v[3]*km

def mat_vec(M):
    '''
    Converts a 2x2 imaginary matrix quaternion into its vector representation
    '''
    return np.array([ M[1,1].real , M[1,1].imag , M[0,1].real , -M[0,1].imag ])

def qvecmult(vec1,vec2):
    '''
    Multiplies two 4-vector quaternions via matrix math
    '''
    return mat_vec(np.dot(vec_mat(vec1),vec_mat(vec2)))

def qmatcon(M):
    '''
    conjugates a 2x2 imaginary matrix quaternion
    '''
    return vec_mat(mat_vec(M)*np.array([1,-1,-1,-1]))

def qveccon(vec):
    '''
    conjugates 4-vector quaternion
    '''
    return vec*np.array([1,-1,-1,-1])

def qvecnorm(vec):
    '''
    normalizes a 4-vector quaternion
    '''
    return vec/np.sqrt(qvecmult(qveccon(vec),vec)[0])

def qmatnorm(M):
    '''
    piggy-backs off the previous function to normalize 2x2 imaginary matrices
    '''
    return vec_mat(qvecnorm(mat_vec(M)))

def qvecmagsqr(vec):
    '''
    returns the magnitude squared of a 4-vector quaternion
    '''
    return qvecmult(qveccon(vec),vec)[0]

def qmatmagsqr(M):
    '''
    piggy-backs off the previous function to give the magnitude squared of 2x2 imaginary matrix
    quaternions
    '''
    return qvecmagsqr(mat_vec(M))

#------------------------------------------------------------------------------
#                   Defining the differential equations
#              AKA "Bringing (first) order to the universe"
#------------------------------------------------------------------------------

def q1_dot(q1,q2,p1,p2,a):
    '''
    takes the current value of things that we know and calculates derivatives
    Function assumes 2x2 complex matrices as inputs for q1,q2,p1,p2
    a is the coupling constant
    '''
    return (p1 - a*np.dot(q1,np.dot(qmatcon(q2),p2)))   \
           #/(1. - qmatmagsqr(q1)*qmatmagsqr(q2)*a**2)

def p1_dot(q1,q2,q1dot,q2dot,a,w):
    '''
    takes the current values of things we know and the hopefully recently
    calculated derivatives of q1,q2 and uses them to find other derivatives
    '''
    return a*np.dot(q1dot,np.dot(qmatcon(q2dot),q2)) - q1*w**2

#------------------------------------------------------------------------------
#          Defining necessary constants and initial conditions
#                   AKA "on the first day..."
#------------------------------------------------------------------------------

w = 1. # \omega_0 in our notation
a = 0.01 # coupling constant. \alpha in our notation
print 'alpha =',a

seed = 42
np.random.seed(seed)
print 'seed =',seed

q1 = vec_mat([1,0,0,0])
q2 = vec_mat([1,0,0,0])
p1 = np.random.rand(4)
p2 = np.random.rand(4)
p1[0] = 0
p2[0] = 0

p1 = vec_mat(p1)
p2 = vec_mat(p2)

q1 = qmatnorm(q1)
q2 = qmatnorm(q2)
p1 = qmatnorm(p1)
p2 = qmatnorm(p2)

#------------------------------------------------------------------------------
#                     Defining loop parameters
#            AKA "Configuring the space-time continuum"
#------------------------------------------------------------------------------

dt = 0.01 #time step
t  = 0
print 'dt = ',dt

q1a = []
q2a = []
p1a = []
p2a = []
s1a = []
s2a = []

t1a = []
t2a = []

q1m = []
q2m = []
p1m = []
p2m = []
S1 = [0,0]
S2 = [0,0]

swaptime = 0.88/a #determined 'experimentally'

#------------------------------------------------------------------------------
#                   Checking conserved quantity
#                     AKA "might as well..."
#------------------------------------------------------------------------------

con = [] #checking to see if our conserved quantity is actually conserved

def conserved(q1,q2,p1,p2):
    return np.dot(qmatcon(p1),q1) + np.dot(qmatcon(p2),q2)

#------------------------------------------------------------------------------
#                   Creating the time loop
#                     AKA "Let 'er rip"
#------------------------------------------------------------------------------

runtime = checktime()
while t<swaptime:
    '''
    This integrator works on an RK4 algorithm.
    For a good explaination, see wikipedia
    https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
    note that the algorithm is modified slightly to fit our function
    '''
    q1k1 = q1_dot(q1,q2,p1,p2,a)
    q2k1 = q1_dot(q2,q1,p2,p1,a)
    p1k1 = p1_dot(q1,q2,q1k1,q2k1,a,w)
    p2k1 = p1_dot(q2,q1,q2k1,q1k1,a,w)
    q1k2 = q1_dot(q1+q1k1*dt/2.,q2+q2k1*dt/2.,p1+p1k1*dt/2.,p2+p2k1*dt/2.,a)
    q2k2 = q1_dot(q2+q2k1*dt/2.,q1+q1k1*dt/2.,p2+p2k1*dt/2.,p1+p1k1*dt/2.,a)
    p1k2 = p1_dot(q1+q1k1*dt/2.,q2+q2k1*dt/2.,q1k1,q2k1,a,w)
    p2k2 = p1_dot(q2+q2k1*dt/2.,q1+q1k1*dt/2.,q2k1,q1k1,a,w)
    q1k3 = q1_dot(q1+q1k2*dt/2.,q2+q2k2*dt/2.,p1+p1k2*dt/2.,p2+p2k2*dt/2.,a)
    q2k3 = q1_dot(q2+q2k2*dt/2.,q1+q1k2*dt/2.,p2+p2k2*dt/2.,p1+p1k2*dt/2.,a)
    p1k3 = p1_dot(q1+q1k2*dt/2.,q2+q2k2*dt/2.,q1k2,q2k2,a,w)
    p2k3 = p1_dot(q2+q2k2*dt/2.,q1+q1k2*dt/2.,q2k2,q1k2,a,w)
    q1k4 = q1_dot(q1+q1k3*dt,q2+q2k3*dt,p1+p1k3*dt,p2+p2k3*dt,a)
    q2k4 = q1_dot(q2+q2k3*dt,q1+q1k3*dt,p2+p2k3*dt,p1+p1k3*dt,a)
    p1k4 = p1_dot(q1+q1k3*dt,q2+q2k3*dt,q1k3,q2k3,a,w)
    p2k4 = p1_dot(q2+q2k3*dt,q1+q1k3*dt,q2k3,q1k3,a,w)
    q1 += (q1k1 + 2*q1k2 + 2*q1k3 + q1k4)*dt/6.
    q2 += (q2k1 + 2*q2k2 + 2*q2k3 + q2k4)*dt/6.
    p1 += (p1k1 + 2*p1k2 + 2*p1k3 + p1k4)*dt/6.
    p2 += (p2k1 + 2*p2k2 + 2*p2k3 + p2k4)*dt/6.
    t  += dt
    q1a.append(mat_vec(q1))
    p1a.append(mat_vec(p1))
    q1m.append(qmatmagsqr(q1))
    p1m.append(qmatmagsqr(p1))
    s1a.append(mat_vec(np.dot(qmatcon(p1),q1)))
    t1a.append(t)
    q2a.append(mat_vec(q2))
    p2a.append(mat_vec(p2))
    q2m.append(qmatmagsqr(q2))
    p2m.append(qmatmagsqr(p2))
    s2a.append(mat_vec(np.dot(qmatcon(p2),q2)))
    t2a.append(t)

runtime = checktime() - runtime
q1a = np.array(q1a)
q2a = np.array(q2a)
p1a = np.array(p1a)
p2a = np.array(p2a)
t1a = np.array(t1a)
t2a = np.array(t2a)
s1a = np.array(s1a)
s2a = np.array(s2a)
#con = np.array(con)

#------------------------------------------------------------------------------
#                       Plotting things
#                   AKA "Can we see it now?"
#------------------------------------------------------------------------------

import matplotlib.pyplot as plt

def vecplot(thing,time,name):
    plt.clf()
    plt.title(name)
    plt.plot(time,thing[:,0],label='Real', color = 'black')
    plt.plot(time,thing[:,1],label='i', color = 'red')
    plt.plot(time,thing[:,2],label='j', color = 'green')
    plt.plot(time,thing[:,3],label='k', color = 'blue')
    plt.legend(loc='best')
    plt.xlim([time[0], time[-1]])
    plt.grid()
    plt.show()

def scalarplot(thing,time,name):
    plt.clf()
    plt.title(name)
    plt.plot(time,thing,color = 'black')
    plt.grid()
    plt.xlim([time[0], time[-1]])
    plt.show()

vecplot(q1a,t1a,'$q_1$')
vecplot(q2a,t2a,'$q_2$')
vecplot(p1a,t1a,'$p_1$')
vecplot(p2a,t2a,'$p_2$')
vecplot(s1a,t1a,'$p_1^{\dagger}q_1$')
vecplot(s2a,t2a,'$p_2^{\dagger}q_2$')

print 'Initial:'
print 'q1 = ', q1a[0]
print 'q2 = ', q2a[0]
print 'p1 = ', p1a[0]
print 'p2 = ', p2a[0]
print 's1 = ', s1a[0]
print 's2 = ', s2a[0]
print 'Final:'
print 'q1 = ', q1a[-1]
print 'q2 = ', q2a[-1]
print 'p1 = ', p1a[-1]
print 'p2 = ', p2a[-1]
print 's1 = ', s1a[-1]
print 's2 = ', s2a[-1]
print 'runtime is',runtime, 'seconds'<commit_msg>Added files via upload<commit_after># -*- coding: utf-8 -*-
"""
This code uses a loop along with our set of coupled differential equations and
matrix math to create arrays of 4-vector quaternions.

The old plotting functions need to be updated and incorperated into the end of
this code or a better visualization solution needs to be found.
"""

#------------------------------------------------------------------------------
#               Importing modules and copying functions
#                   AKA "Setting stuff up"
#------------------------------------------------------------------------------

import numpy as np
from time import time as checktime

# a set of init quaternions and the identity matrix for building general q-matrices
rm = np.identity(2)
im = np.array([[-1j,0],[0,1j]])
jm = np.array([[0,1],[-1,0]])
km = np.array([[0,-1j],[-1j,0]])

def vec_mat(v):
    '''
    Converts a quaternion vector into the 2x2 imaginary matrix representation
    '''
    return v[0]*rm + v[1]*im + v[2]*jm + v[3]*km

def mat_vec(M):
    '''
    Converts a 2x2 imaginary matrix quaternion into its vector representation
    '''
    return np.array([ M[1,1].real , M[1,1].imag , M[0,1].real , -M[0,1].imag ])

def qvecmult(vec1,vec2):
    '''
    Multiplies two 4-vector quaternions via matrix math
    '''
    return mat_vec(np.dot(vec_mat(vec1),vec_mat(vec2)))

def qmatcon(M):
    '''
    conjugates a 2x2 imaginary matrix quaternion
    '''
    return vec_mat(mat_vec(M)*np.array([1,-1,-1,-1]))

def qveccon(vec):
    '''
    conjugates 4-vector quaternion
    '''
    return vec*np.array([1,-1,-1,-1])

def qvecnorm(vec):
    '''
    normalizes a 4-vector quaternion
    '''
    return vec/np.sqrt(qvecmult(qveccon(vec),vec)[0])

def qmatnorm(M):
    '''
    piggy-backs off the previous function to normalize 2x2 imaginary matrices
    '''
    return vec_mat(qvecnorm(mat_vec(M)))

def qvecmagsqr(vec):
    '''
    returns the magnitude squared of a 4-vector quaternion
    '''
    return qvecmult(qveccon(vec),vec)[0]

def qmatmagsqr(M):
    '''
    piggy-backs off the previous function to give the magnitude squared of 2x2 imaginary matrix
    quaternions
    '''
    return qvecmagsqr(mat_vec(M))

#------------------------------------------------------------------------------
#                   Defining the differential equations
#              AKA "Bringing (first) order to the universe"
#------------------------------------------------------------------------------

def q1_dot(q1,q2,p1,p2,a):
    '''
    takes the current value of things that we know and calculates derivatives
    Function assumes 2x2 complex matrices as inputs for q1,q2,p1,p2
    a is the coupling constant
    '''
    return (p1 - a*np.dot(q1,np.dot(qmatcon(q2),p2)))   \
           #/(1. - qmatmagsqr(q1)*qmatmagsqr(q2)*a**2)

def p1_dot(q1,q2,q1dot,q2dot,a,w):
    '''
    takes the current values of things we know and the hopefully recently
    calculated derivatives of q1,q2 and uses them to find other derivatives
    '''
    return a*np.dot(q1dot,np.dot(qmatcon(q2dot),q2)) - q1*w**2

#------------------------------------------------------------------------------
#          Defining necessary constants and initial conditions
#                   AKA "on the first day..."
#------------------------------------------------------------------------------

w = 1. # \omega_0 in our notation
a = 0.01 # coupling constant. \alpha in our notation
print 'alpha =',a

seed = 42
np.random.seed(seed)
print 'seed =',seed

q1 = vec_mat([1,0,0,0])
q2 = vec_mat([1,0,0,0])
p1 = np.random.rand(4)
p2 = np.random.rand(4)
p1[0] = 0
p2[0] = 0

p1 = vec_mat(p1)
p2 = vec_mat(p2)

q1 = qmatnorm(q1)
q2 = qmatnorm(q2)
p1 = qmatnorm(p1)
p2 = qmatnorm(p2)

#------------------------------------------------------------------------------
#                     Defining loop parameters
#            AKA "Configuring the space-time continuum"
#------------------------------------------------------------------------------

dt = 0.01 #time step
t  = 0
print 'dt = ',dt

q1a = [mat_vec(q1)]
p1a = [mat_vec(p1)]
s1a = [mat_vec(np.dot(qmatcon(p1),q1))]
q2a = [mat_vec(q2)]
p2a = [mat_vec(p2)]
s2a = [mat_vec(np.dot(qmatcon(p2),q2))]
time = [t]

swaptime = 0.8785/a #determined 'experimentally'

#------------------------------------------------------------------------------
#                   Checking conserved quantity
#                     AKA "might as well..."
#------------------------------------------------------------------------------

con = [] #checking to see if our conserved quantity is actually conserved

def conserved(q1,q2,p1,p2):
    return np.dot(qmatcon(p1),q1) + np.dot(qmatcon(p2),q2)

#------------------------------------------------------------------------------
#                   Creating the time loop
#                     AKA "Let 'er rip"
#------------------------------------------------------------------------------

runtime = checktime()
while t<swaptime:
    '''
    This integrator works on an RK4 algorithm.
    For a good explaination, see wikipedia
    https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
    note that the algorithm is modified slightly to fit our function
    '''
    q1k1 = q1_dot(q1,q2,p1,p2,a)
    q2k1 = q1_dot(q2,q1,p2,p1,a)
    p1k1 = p1_dot(q1,q2,q1k1,q2k1,a,w)
    p2k1 = p1_dot(q2,q1,q2k1,q1k1,a,w)
    q1k2 = q1_dot(q1+q1k1*dt/2.,q2+q2k1*dt/2.,p1+p1k1*dt/2.,p2+p2k1*dt/2.,a)
    q2k2 = q1_dot(q2+q2k1*dt/2.,q1+q1k1*dt/2.,p2+p2k1*dt/2.,p1+p1k1*dt/2.,a)
    p1k2 = p1_dot(q1+q1k1*dt/2.,q2+q2k1*dt/2.,q1k1,q2k1,a,w)
    p2k2 = p1_dot(q2+q2k1*dt/2.,q1+q1k1*dt/2.,q2k1,q1k1,a,w)
    q1k3 = q1_dot(q1+q1k2*dt/2.,q2+q2k2*dt/2.,p1+p1k2*dt/2.,p2+p2k2*dt/2.,a)
    q2k3 = q1_dot(q2+q2k2*dt/2.,q1+q1k2*dt/2.,p2+p2k2*dt/2.,p1+p1k2*dt/2.,a)
    p1k3 = p1_dot(q1+q1k2*dt/2.,q2+q2k2*dt/2.,q1k1,q2k1,a,w)
    p2k3 = p1_dot(q2+q2k2*dt/2.,q1+q1k2*dt/2.,q2k1,q1k1,a,w)
    q1k4 = q1_dot(q1+q1k3*dt,q2+q2k3*dt,p1+p1k3*dt,p2+p2k3*dt,a)
    q2k4 = q1_dot(q2+q2k3*dt,q1+q1k3*dt,p2+p2k3*dt,p1+p1k3*dt,a)
    p1k4 = p1_dot(q1+q1k3*dt,q2+q2k3*dt,q1k1,q2k1,a,w)
    p2k4 = p1_dot(q2+q2k3*dt,q1+q1k3*dt,q2k1,q1k1,a,w)
    q1 += (q1k1 + 2*q1k2 + 2*q1k3 + q1k4)*dt/6.
    q2 += (q2k1 + 2*q2k2 + 2*q2k3 + q2k4)*dt/6.
    p1 += (p1k1 + 2*p1k2 + 2*p1k3 + p1k4)*dt/6.
    p2 += (p2k1 + 2*p2k2 + 2*p2k3 + p2k4)*dt/6.
    t  += dt
    q1a.append(mat_vec(q1))
    p1a.append(mat_vec(p1))
    s1a.append(mat_vec(np.dot(qmatcon(p1),q1)))
    q2a.append(mat_vec(q2))
    p2a.append(mat_vec(p2))
    s2a.append(mat_vec(np.dot(qmatcon(p2),q2)))
    time.append(t)

runtime = checktime() - runtime

q1a = np.array(q1a)
q2a = np.array(q2a)
p1a = np.array(p1a)
p2a = np.array(p2a)
s1a = np.array(s1a)
s2a = np.array(s2a)
time = np.array(time)

#------------------------------------------------------------------------------
#                       Plotting things
#                   AKA "Can we see it now?"
#------------------------------------------------------------------------------

import matplotlib.pyplot as plt

def vecplot(thing,time,name):
    plt.clf()
    plt.title(name)
    plt.plot(time,thing[:,0],label='Real', color = 'black')
    plt.plot(time,thing[:,1],label='i', color = 'red')
    plt.plot(time,thing[:,2],label='j', color = 'green')
    plt.plot(time,thing[:,3],label='k', color = 'blue')
    plt.legend(loc='best')
    plt.xlim([time[0], time[-1]])
    plt.grid()
    plt.show()

def scalarplot(thing,time,name):
    plt.clf()
    plt.title(name)
    plt.plot(time,thing,color = 'black')
    plt.grid()
    plt.xlim([time[0], time[-1]])
    plt.show()

vecplot(q1a,time,'$q_1$')
vecplot(q2a,time,'$q_2$')
vecplot(p1a,time,'$p_1$')
vecplot(p2a,time,'$p_2$')
vecplot(s1a,time,'$p_1^{\dagger}q_1$')
vecplot(s2a,time,'$p_2^{\dagger}q_2$')

print 'Initial:'
print 'q1 = ', q1a[0]
print 'q2 = ', q2a[0]
print 'p1 = ', p1a[0]
print 'p2 = ', p2a[0]
print 's1 = ', s1a[0]
print 's2 = ', s2a[0]
print 'Final:'
print 'q1 = ', q1a[-1]
print 'q2 = ', q2a[-1]
print 'p1 = ', p1a[-1]
print 'p2 = ', p2a[-1]
print 's1 = ', s1a[-1]
print 's2 = ', s2a[-1]
print 'runtime is',runtime, 'seconds'<|endoftext|>