<commit_before>#!/usr/bin/env python

# L. Howard Copyright @2014
# Convert a CAD model (per the STEPTOOLS defined XML spec)
# into a JSON spec model
# Derived from Javascript version at
# https://github.com/ghemingway/cad.js/blob/master/scripts/xmlToJson.js

import argparse
import json
import math
from multiprocessing import cpu_count, Process, Queue
from operator import itemgetter
import os
from os.path import join
import sys
import time
import xml.etree.ElementTree as ET

import logging
logging.basicConfig(format='%(asctime)s %(levelname)s:%(message)s',
                    level=logging.INFO)
LOG = logging.getLogger(__name__)

# defaults and constants
DEFAULT_COLOR = "7d7d7d"
IDENTITY_TRANSFORM = "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"

#------------------------------------------------------------------------------

CONFIG = {
    'indexPoints': True,
    'indexNormals': True,
    'indexColors': True,
    'compressColors': True,
    'roundPrecision': 2
}

def roundFloat(val, precision):
    """floating point rounder"""
    if not precision:
        return val
    factor = math.pow(10, precision)
    return int(round(val * factor))

#------------------------------------------------------------------------------

def translateIndex(doc):
    """Returns the full JSON"""
    return {
        'root': doc.attrib['root'],
        'products': [translateProduct(x) for x in doc.iter('product')],
        'shapes': [translateShape(x) for x in doc.iter('shape')],
        'shells': [translateShell(x) for x in doc.iter('shell')],
        'annotations': [translateAnnotation(x) for x in doc.iter('annotation')]
    }

def translateProduct(product):
    """Translates a product"""
    data = {
        'id': product.attrib['id'],
        'step': product.attrib.get('step', ""),
        'name': product.attrib['name']
    }
    # Add children, if there are any
    if product.attrib.get('children'):
        data['children'] = product.attrib['children'].split(" ")
    # Add shapes, if there are any
    if product.attrib.get('shape'):
        data['shapes'] = product.attrib['shape'].split(" ")
    return data

def setTransform(transform):
    """Sets a transform"""
    return ("I" if transform == IDENTITY_TRANSFORM else
            [float(x) for x in transform.split(" ")])

def translateShape(shape):
    """Translates a shape"""
    data = {
        'id': shape.attrib['id'],
        # "unit": shape.attrib['unit'],
    }
    data.update({x: [] for x in ('shells', 'annotations', 'children')})
    for child in shape.iter('child'):
        # Add children, if there are any
        data['children'].append({
            'ref': child.attrib['ref'],
            'xform': setTransform(child.attrib['xform'])
        })
    # Add child annotations
    if shape.attrib.get('annotation'):
        data['annotations'] = shape.attrib['annotation'].split(" ")
    # Terminal Shape JSON
    if shape.attrib.get('shell'):
        data['shells'] = shape.attrib['shell'].split(" ")
    return data

def translateAnnotation(annotation):
    """Translates an annotation"""
    data = dict(id=annotation.attrib['id'])
    if 'href' in annotation.attrib:
        data['href'] = annotation.attrib['href'].replace("xml", "json")
    else:
        data['lines'] = []
        for polyline in annotation.iter('polyline'):
            points = []
            for p in polyline.iter('p'):
                points.extend([float(x) for x in p.attrib['l'].split(" ")])
            data['lines'].append(points)
    return data

#------------------------------------------------------------------------------

def make_index(data, ikey, ranger=None):
    """Create indexes, an abstraction of indexing functions in the original"""
    if ranger is None:
        ranger = xrange(len(data[ikey]))
    indexes = data[ikey + "Index"] = []
    values = data['values']
    for i in ranger:
        val = roundFloat(data[ikey][i], CONFIG['roundPrecision'])
        if val not in values:
            values[val] = len(values)
        indexes.append(values[val])
    del data[ikey]

# The following functions are named indexShellxxx in the original
# The renaming aligns them with settings in the indexing CONFIG
indexPoints = lambda d: make_index(d, 'points')
indexNormals = lambda d: make_index(d, 'normals')
indexColors = lambda d: make_index(d, 'colors')

def compressShellColors(data):
    """Color compression"""
    numTuples = len(data['colorsIndex']) / 3
    data['colorsData'] = []
    start = 0
    last = [data['colorsIndex'][x] for x in xrange(3)]
    # Short list comparison
    arraysIdentical = lambda a, b: all([a[x] == b[x] for x in xrange(3)])
    # Compress the rest
    for tupl in xrange(numTuples):
        index = tupl * 3
        tmp = [data['colorsIndex'][index + x] for x in xrange(3)]
        # Is this a new block?
        if not arraysIdentical(last, tmp):
            data['colorsData'].append(dict(data=last, duration=tupl - start))
            start = tupl
            last = tmp
    # append the final color block
    data['colorsData'].append(dict(data=last, duration=numTuples - start))
    # remove the colors index
    del data['colorsIndex']

#------------------------------------------------------------------------------

def translateShell(shell):
    """Translates a shell"""
    if 'href' in shell.attrib:
        # Do href here
        return {
            'id': shell.attrib['id'],
            'size': int(shell.attrib['size']),
            'bbox': [float(x) for x in shell.attrib['bbox'].split(" ")],
            'href': shell.attrib['href'].replace("xml", "json")
        }
    else:
        # Convert XML point/vert/color to new way
        points = loadPoints(shell.iter("verts"))
        defaultColor = parseColor(shell.attrib.get('color', DEFAULT_COLOR))
        data = dict(id=shell.attrib['id'], size=0)
        data.update({x: [] for x in ('points', 'normals', 'colors')})
        for facet in shell.iter('facets'):
            color = defaultColor
            if 'color' in facet.attrib:
                color = parseColor(facet.attrib['color'])
            for f in facet.iter('f'):
                # Get every vertex index and convert using points array
                indexVals = f.attrib['v'].split(" ")
                for i in range(3):
                    ival = int(indexVals[i]) * 3
                    data['points'].append(float(points[ival]))
                    data['points'].append(float(points[ival + 1]))
                    data['points'].append(float(points[ival + 2]))

                # Get the vertex normals
                norms = [x for x in f.iter('n')]
                for i in range(3):
                    normCoordinates = norms[i].attrib['d'].split(" ")
                    for j in range(3):
                        data['normals'].append(float(normCoordinates[j]))

                # Get the vertex colors
                for i in range(3):
                    for c in ('r', 'g', 'b'):
                        data['colors'].append(color[c])

        data['size'] = len(data['points']) / 9
        indexing = [x for x in CONFIG if x.startswith('index') and CONFIG[x]]
        for i in indexing:
            data['precision'] = CONFIG['roundPrecision']
            if 'values' not in data:
                data['values'] = {}
            globals()[i](data)
        if indexing:
            sorted_vals = sorted(data['values'].items(), key=itemgetter(1))
            data['values'] = map(itemgetter(0), sorted_vals)
        if CONFIG.get('indexColors') and CONFIG.get('compressColors'):
            compressShellColors(data)
        return data

def parseColor(hex):
    """Parse color values"""
    cval = int(hex, 16)
    x = lambda b: ((cval >> b) & 0xff) / 255.0
    return {k: x(v) for k, v in dict(r=16, g=8, b=0).iteritems()}

def loadPoints(verts):
    """Load all of the point information"""
    points = []
    for vert in verts:
        for v in vert:
            points.extend(v.attrib['p'].split(" "))
    return points

#------------------------------------------------------------------------------

class Worker(Process):
    """Worker process for parallelized translation"""

    def __init__(self, queue, exceptions):
        Process.__init__(self)
        self.queue = queue
        self.exceptions = exceptions

    def report_exception(self, job, reason):
        """Report a job exception"""
        info = dict(reason=reason)
        info.update(job)
        self.exceptions.put(info)

    def run(self):
        """Process jobs"""
        while True:
            job = self.queue.get()
            if job is None:
                break
            try:
                path = job['path']
                tree = ET.parse(path)
                root = tree.getroot()
            except Exception as e:
                reason = "Unable to parse XML file '{}'.".format(path)
                self.report_exception(job, reason)
                continue
            try:
                data = job['translator'](root)
            except Exception as e:
                reason = "Translation failure: '{}'.".format(e)
                self.report_exception(job, reason)
                continue
            out_path = os.path.splitext(path)[0] + ".json"
            try:
                with open(out_path, "w") as f:
                    json.dump(data, f)
            except Exception as e:
                reason = "Unable to output JSON '{}': {}.".format(out_path, e)
                self.report_exception(job, reason)
                continue

class XMLTranslator(object):
    """Translates STEP XML files to JSON"""

    def translate(self, xml_dir, xml_index):
        """Process index XML and enqueue jobs for workers"""
        if not os.path.isdir(xml_dir):
            LOG.error("'{}' is not a directory.".format(xml_dir))
            return True
        index_path = join(xml_dir, xml_index)
        if not os.path.isfile(index_path):
            LOG.error("Unable to locate index file '{}'.".format(index_path))
            return True
        try:
            tree = ET.parse(index_path)
            root = tree.getroot()
        except Exception as e:
            LOG.exception("Unable to parse '{}'.".format(index_path))
            return True
        try:
            data = translateIndex(root)
        except Exception as e:
            LOG.exception("Unable to translate index file.")
            return True

        pluck = lambda e, a: [x for x in data.get(e, []) if a in x]
        externalShells = pluck('shells', 'href')
        externalAnnotations = pluck('annotations', 'href')
        indexOut = join(xml_dir, os.path.splitext(xml_index)[0] + ".json")

        LOG.info("Writing new index file: " + indexOut)
        LOG.info("\tProducts: %s" % len(data.get('projects', [])))
        LOG.info("\tShapes: %s" % len(data.get('shapes', [])))
        LOG.info("\tAnnotations: %s" % len(data.get('annotations', [])))
        LOG.info("\tExternal Annotations: %s" % len(externalAnnotations))
        LOG.info("\tShells: %s" % len(data.get('shells', [])))
        LOG.info("\tExternal Shells: %s" % len(externalShells))

        try:
            with open(indexOut, "w") as f:
                json.dump(data, f)
        except Exception as e:
            LOG.exception("Unable to write JSON file '{}'.".format(indexOut))
            return True

        # start workers and queue jobs
        queue = Queue()
        exceptions = Queue()
        workers = [Worker(queue, exceptions) for w in xrange(cpu_count())]
        for w in workers:
            w.start()

        xml_path = lambda p: join(xml_dir, os.path.splitext(p)[0] + ".xml")
        for annotation in externalAnnotations:
            queue.put({
                'type': "annotation",
                'path': xml_path(annotation['href']),
                'translator': translateAnnotation
            })

        for shell in externalShells:
            queue.put({
                'type': "shell",
                'path': xml_path(shell['href']),
                'translator': translateShell
            })

        # add worker termination cues
        for w in workers:
            queue.put(None)

        # wait for completion
        while any([x.is_alive() for x in workers]):
            time.sleep(1)

        # report errors, if any
        has_errors = not exceptions.empty()
        while not exceptions.empty():
            info = exceptions.get()
            msg = "Error processing '{}': {}"
            LOG.error(msg.format(info['path'], info['reason']))

        # return success/failure result
        return has_errors

#------------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog='xmlToJson.py',
        description="Translates STEP XML to JSON")
    parser.add_argument("dir", help="directory containing STEP XML")
    parser.add_argument("index", help="index file")
    args = parser.parse_args()

    translator = XMLTranslator()
    has_errors = translator.translate(args.dir, args.index)
    sys.exit(1 if has_errors else 0)
<commit_msg>support for shell batching<commit_after>#!/usr/bin/env python

# L. Howard Copyright @2014
# Convert a CAD model (per the STEPTOOLS defined XML spec)
# into a JSON spec model
# Derived from Javascript version at
# https://github.com/ghemingway/cad.js/blob/master/scripts/xmlToJson.js

import argparse
from datetime import datetime
import json
import math
from multiprocessing import cpu_count, Process, Queue
from operator import itemgetter
import os
from os.path import join
import re
import sys
import time
import xml.etree.cElementTree as ET

import logging
logging.basicConfig(format='%(asctime)s %(levelname)s:%(message)s',
                    level=logging.DEBUG)
LOG = logging.getLogger(__name__)

# defaults and constants
DEFAULT_COLOR = "7d7d7d"
IDENTITY_TRANSFORM = "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"
SHELL_REGEX = re.compile("shell_(.*?).json")

#------------------------------------------------------------------------------

CONFIG = {
    'indexPoints': True,
    'indexNormals': True,
    'compressColors': True,
    'roundPrecision': 2
}

def roundFloat(val, precision):
    """floating point rounder"""
    if not precision:
        return val
    factor = math.pow(10, precision)
    return int(round(val * factor))

#------------------------------------------------------------------------------

def translateIndex(doc):
    """Returns the full JSON"""
    return {
        'root': doc.attrib['root'],
        'products': [translateProduct(x) for x in doc.iter('product')],
        'shapes': [translateShape(x) for x in doc.iter('shape')],
        'shells': [translateShell(x) for x in doc.iter('shell')],
        'annotations': [translateAnnotation(x) for x in doc.iter('annotation')]
    }

def translateProduct(product):
    """Translates a product"""
    data = {
        'id': product.attrib['id'],
        'step': product.attrib.get('step', ""),
        'name': product.attrib['name']
    }
    # Add children, if there are any
    if product.attrib.get('children'):
        data['children'] = product.attrib['children'].split(" ")
    # Add shapes, if there are any
    if product.attrib.get('shape'):
        data['shapes'] = product.attrib['shape'].split(" ")
    return data

def setTransform(transform):
    """Sets a transform"""
    return ("I" if transform == IDENTITY_TRANSFORM else
            [float(x) for x in transform.split(" ")])

def translateShape(shape):
    """Translates a shape"""
    data = {
        'id': shape.attrib['id'],
        # "unit": shape.attrib['unit'],
    }
    data.update({x: [] for x in ('shells', 'annotations', 'children')})
    for child in shape.iter('child'):
        # Add children, if there are any
        data['children'].append({
            'ref': child.attrib['ref'],
            'xform': setTransform(child.attrib['xform'])
        })
    # Add child annotations
    if shape.attrib.get('annotation'):
        data['annotations'] = shape.attrib['annotation'].split(" ")
    # Terminal Shape JSON
    if shape.attrib.get('shell'):
        data['shells'] = shape.attrib['shell'].split(" ")
    return data

def translateAnnotation(annotation):
    """Translates an annotation"""
    data = dict(id=annotation.attrib['id'])
    if 'href' in annotation.attrib:
        data['href'] = annotation.attrib['href'].replace("xml", "json")
    else:
        data['lines'] = []
        for polyline in annotation.iter('polyline'):
            points = []
            for p in polyline.iter('p'):
                points.extend([float(x) for x in p.attrib['l'].split(" ")])
            data['lines'].append(points)
    return data

#------------------------------------------------------------------------------

def make_index(data, ikey, ranger=None):
    """Create indexes, an abstraction of indexing functions in the original"""
    if ranger is None:
        ranger = xrange(len(data[ikey]))
    indexes = data[ikey + "Index"] = []
    values = data['values']
    for i in ranger:
        val = roundFloat(data[ikey][i], CONFIG['roundPrecision'])
        if val not in values:
            values[val] = len(values)
        indexes.append(values[val])
    del data[ikey]

# The following functions are named indexShellxxx in the original
# The renaming aligns them with settings in the indexing CONFIG
indexPoints = lambda d: make_index(d, 'points')
indexNormals = lambda d: make_index(d, 'normals')

def compressShellColors(data):
    """Color compression"""
    numTuples = len(data['colors']) / 3
    data['colorsData'] = []
    start = 0
    last = [data['colors'][x] for x in xrange(3)]
    # Short list comparison
    arraysIdentical = lambda a, b: all([a[x] == b[x] for x in xrange(3)])
    # Compress the rest
    for tupl in xrange(numTuples):
        index = tupl * 3
        tmp = [data['colors'][index + x] for x in xrange(3)]
        # Is this a new block?
        if not arraysIdentical(last, tmp):
            data['colorsData'].append(dict(data=last, duration=tupl - start))
            start = tupl
            last = tmp
    # append the final color block
    data['colorsData'].append(dict(data=last, duration=numTuples - start))
    # remove the colors index
    del data['colors']

#------------------------------------------------------------------------------

def translateShell(shell):
    """Translates a shell"""
    if 'href' in shell.attrib:
        # Do href here
        return {
            'id': shell.attrib['id'],
            'size': int(shell.attrib['size']),
            'bbox': [float(x) for x in shell.attrib['bbox'].split(" ")],
            'href': shell.attrib['href'].replace("xml", "json")
        }
    else:
        # Convert XML point/vert/color to new way
        points = loadPoints(shell.iter("verts"))
        defaultColor = parseColor(shell.attrib.get('color', DEFAULT_COLOR))
        data = dict(id=shell.attrib['id'], size=0)
        data.update({x: [] for x in ('points', 'normals', 'colors')})
        for facet in shell.iter('facets'):
            color = defaultColor
            if 'color' in facet.attrib:
                color = parseColor(facet.attrib['color'])
            for f in facet.iter('f'):
                # Get every vertex index and convert using points array
                indexVals = f.attrib['v'].split(" ")
                for i in range(3):
                    ival = int(indexVals[i]) * 3
                    data['points'].append(float(points[ival]))
                    data['points'].append(float(points[ival + 1]))
                    data['points'].append(float(points[ival + 2]))

                # Get the vertex normals
                norms = [x for x in f.iter('n')]
                for i in range(3):
                    normCoordinates = norms[i].attrib['d'].split(" ")
                    for j in range(3):
                        data['normals'].append(float(normCoordinates[j]))

                # Get the vertex colors
                for i in range(3):
                    for c in ('r', 'g', 'b'):
                        data['colors'].append(color[c])

        data['size'] = len(data['points']) / 9
        indexing = [x for x in CONFIG if x.startswith('index') and CONFIG[x]]
        for i in indexing:
            data['precision'] = CONFIG['roundPrecision']
            if 'values' not in data:
                data['values'] = {}
            globals()[i](data)
        if indexing:
            sorted_vals = sorted(data['values'].items(), key=itemgetter(1))
            data['values'] = map(itemgetter(0), sorted_vals)
        if CONFIG.get('compressColors'):
            compressShellColors(data)
        return data

def parseColor(hex):
    """Parse color values"""
    cval = int(hex, 16)
    x = lambda b: ((cval >> b) & 0xff) / 255.0
    return {k: x(v) for k, v in dict(r=16, g=8, b=0).iteritems()}

def loadPoints(verts):
    """Load all of the point information"""
    points = []
    for vert in verts:
        for v in vert:
            points.extend(v.attrib['p'].split(" "))
    return points

#------------------------------------------------------------------------------

class BaseWorker(Process):
    """Base class for Workers"""

    def __init__(self, queue, exceptions):
        Process.__init__(self)
        self.queue = queue
        self.exceptions = exceptions

    def report_exception(self, job, reason):
        """Report a job exception"""
        info = dict(reason=reason)
        info.update(job)
        self.exceptions.put(info)

    def run(self):
        raise NotImplementedError

class BatchWorker(BaseWorker):
    """Worker process for parallelized shell batching"""

    def run(self):
        """Process jobs"""
        while True:
            job = self.queue.get()
            if job is None:
                break
            # process shells
            batch = {'shells': []}
            indexed = CONFIG['indexPoints'] or CONFIG['indexNormals']
            reindex = job['reindex'] and indexed
            if reindex:
                batch['values'] = {}
            for s in job['shells']:
                try:
                    with open(join(job['path'], s)) as f:
                        shell = json.load(f)
                    sid = SHELL_REGEX.match(s).group(1)
                    shell['id'] = sid
                    if reindex:
                        imap = {}
                        for i, value in enumerate(shell['values']):
                            if value not in batch['values']:
                                batch['values'][value] = len(batch['values'])
                            imap[i] = batch['values'][value]
                        del shell['values']
                        for item in ('points', 'normals'):
                            idx = item + 'Index'
                            if idx in shell:
                                shell[idx] = [imap[x] for x in shell[idx]]
                    batch['shells'].append(shell)
                except Exception as e:
                    reason = "Error batching shell '{}': {}".format(s, e)
                    self.report_exception(job, reason)
                    continue
            # transform values to list
            if reindex:
                sorted_v = sorted(batch['values'].items(), key=itemgetter(1))
                batch['values'] = map(itemgetter(0), sorted_v)
            # write batch
            out_path = join(job['path'], job['name'] + ".json")
            try:
                with open(out_path, "w") as f:
                    json.dump(batch, f)
            except Exception as e:
                reason = "Unable to output JSON '{}': {}.".format(out_path, e)
                self.report_exception(job, reason)

class TranslationWorker(BaseWorker):
    """Worker process for parallelized translation"""

    def run(self):
        """Process jobs"""
        while True:
            job = self.queue.get()
            if job is None:
                break
            try:
                path = job['path']
                tree = ET.parse(path)
                root = tree.getroot()
            except Exception as e:
                reason = "Unable to parse XML file '{}'.".format(path)
                self.report_exception(job, reason)
                continue
            try:
                data = job['translator'](root)
            except Exception as e:
                reason = "Translation failure: '{}'.".format(e)
                self.report_exception(job, reason)
                continue
            out_path = os.path.splitext(path)[0] + ".json"
            try:
                with open(out_path, "w") as f:
                    json.dump(data, f)
            except Exception as e:
                reason = "Unable to output JSON '{}': {}.".format(out_path, e)
                self.report_exception(job, reason)

class XMLTranslator(object):
    """Translates STEP XML files to JSON"""

    def __init__(self, batches=None, reindex=None):
        self.batches = batches
        self.reindex = reindex

    def assign(self, batches, shell):
        """simple bin packing"""
        name, size = shell
        blist = batches.values()
        best = min([x['total_size'] for x in blist])
        selected = [x for x in blist if x['total_size'] == best][0]
        selected['total_size'] += size
        selected['shells'].append(name)

    def get_batches(self, shells):
        """assign shells to batches, leveling by size"""
        batches = {'batch%s' % i:  {'total_size': 0, 'shells': []}
                   for i in xrange(self.batches)}
        for shell in shells:
            self.assign(batches, shell)
        return batches

    def batch(self, xml_dir):
        """Generates batched shell files"""
        is_shell = lambda x: SHELL_REGEX.match(x)
        size_of = lambda x: os.path.getsize(join(xml_dir, x))
        shells = [(x, size_of(x)) for x in os.listdir(xml_dir) if is_shell(x)]
        shells.sort(key=itemgetter(1), reverse=True)
        batches = self.get_batches(shells)

        # start workers and queue jobs
        queue = Queue()
        exceptions = Queue()
        count = min(cpu_count(), self.batches)
        workers = [BatchWorker(queue, exceptions) for w in xrange(count)]
        for w in workers:
            w.start()

        # enqueue jobs
        for batch, info in batches.items():
            job = {'path': xml_dir, 'name': batch, 'shells': info['shells'],
                   'reindex': self.reindex}
            queue.put(job)

        # add worker termination cues
        for w in workers:
            queue.put(None)

        # wait for completion
        while any([x.is_alive() for x in workers]):
            time.sleep(1)

        # report errors, if any
        has_errors = not exceptions.empty()
        while not exceptions.empty():
            info = exceptions.get()
            msg = "Error processing '{}': {}"
            LOG.error(msg.format(info['path'], info['reason']))

        if not has_errors:
            # report achieved compression
            shells_size = sum([size for name, size in shells])
            msg = "Shells.  Count: {} Total Size: {} bytes."
            LOG.debug(msg.format(len(shells), shells_size))
            regex = re.compile("batch[0-9]*.json")
            is_batch = lambda x: regex.match(x)
            sizes = [size_of(x) for x in os.listdir(xml_dir) if is_batch(x)]
            batches_size = sum(sizes)
            msg = "Batches.  Count: {} Total Size: {} bytes."
            LOG.debug(msg.format(len(sizes), batches_size))
            compression = float(batches_size) / float(shells_size)
            LOG.debug("Compression: {}".format(compression))

        return has_errors

    def translate(self, xml_dir, xml_index):
        """Process index XML and enqueue jobs for workers"""
        if not os.path.isdir(xml_dir):
            LOG.error("'{}' is not a directory.".format(xml_dir))
            return True
        index_path = join(xml_dir, xml_index)
        if not os.path.isfile(index_path):
            LOG.error("Unable to locate index file '{}'.".format(index_path))
            return True
        try:
            tree = ET.parse(index_path)
            root = tree.getroot()
        except Exception as e:
            LOG.exception("Unable to parse '{}'.".format(index_path))
            return True
        try:
            data = translateIndex(root)
        except Exception as e:
            LOG.exception("Unable to translate index file.")
            return True

        pluck = lambda e, a: [x for x in data.get(e, []) if a in x]
        externalShells = pluck('shells', 'href')
        externalAnnotations = pluck('annotations', 'href')
        indexOut = join(xml_dir, os.path.splitext(xml_index)[0] + ".json")

        LOG.debug("Writing new index file: " + indexOut)
        LOG.debug("\tProducts: %s" % len(data.get('projects', [])))
        LOG.debug("\tShapes: %s" % len(data.get('shapes', [])))
        LOG.debug("\tAnnotations: %s" % len(data.get('annotations', [])))
        LOG.debug("\tExternal Annotations: %s" % len(externalAnnotations))
        LOG.debug("\tShells: %s" % len(data.get('shells', [])))
        LOG.debug("\tExternal Shells: %s" % len(externalShells))
        if self.batches and len(externalShells):
            if len(externalShells) < self.batches:
                self.batches = 1
            LOG.debug("\tBatches: %s" % self.batches)
            data['batches'] = self.batches
        else:
            self.batches = 0

        try:
            with open(indexOut, "w") as f:
                json.dump(data, f)
        except Exception as e:
            LOG.exception("Unable to write JSON file '{}'.".format(indexOut))
            return True

        # start workers and queue jobs
        queue = Queue()
        exceptions = Queue()
        count = cpu_count()
        workers = [TranslationWorker(queue, exceptions) for w in xrange(count)]
        for w in workers:
            w.start()

        xml_path = lambda p: join(xml_dir, os.path.splitext(p)[0] + ".xml")
        for annotation in externalAnnotations:
            queue.put({
                'type': "annotation",
                'path': xml_path(annotation['href']),
                'translator': translateAnnotation
            })

        for shell in externalShells:
            queue.put({
                'type': "shell",
                'path': xml_path(shell['href']),
                'translator': translateShell
            })

        # add worker termination cues
        for w in workers:
            queue.put(None)

        # wait for completion
        while any([x.is_alive() for x in workers]):
            time.sleep(1)

        # report errors, if any
        has_errors = not exceptions.empty()
        while not exceptions.empty():
            info = exceptions.get()
            msg = "Error processing '{}': {}"
            LOG.error(msg.format(info['path'], info['reason']))

        if has_errors or not self.batches:
            return has_errors

        return self.batch(xml_dir)

#------------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog='xmlToJson.py',
        description="Translates STEP XML to JSON")
    parser.add_argument("dir", help="directory containing STEP XML")
    parser.add_argument("index", help="index file")
    h = "create batches of shells"
    parser.add_argument("-b", "--batches", type=int, default=0, help=h)
    h = "re-index when batching shells"
    parser.add_argument("-r", "--reindex", action="store_true", help=h)
    args = parser.parse_args()

    start = datetime.now()
    translator = XMLTranslator(args.batches, args.reindex)
    has_errors = translator.translate(args.dir, args.index)
    dt = datetime.now() - start
    LOG.info("xmlToJson Elapsed time: {} secs".format(dt.seconds))
    sys.exit(1 if has_errors else 0)
<|endoftext|>


<commit_before>import numpy as np
import unittest
import tensorflow as tf
import GPflow
import GPinv

class test_vgp(unittest.TestCase):
    def setUp(self):
        rng = np.random.RandomState(0)
        self.X = np.linspace(0.,1.,20)
        self.Y = np.sin(3.*self.X) + rng.randn(20)*0.1

    def test_build(self):
        m = GPinv.stvgp.StVGP(self.X.reshape(-1,1), self.Y.reshape(-1,1),
                    GPinv.kernels.RBF(1,output_dim=1),
                    GPinv.likelihoods.Gaussian())
        m._compile()

    def test_optimize(self):
        # reference GPR
        m_ref = GPflow.gpr.GPR(
                    self.X.reshape(-1,1),
                    self.Y.reshape(-1,1),
                    kern = GPflow.kernels.RBF(1))
        rslt_ref = m_ref.optimize()
        # tested StVGP
        tf.set_random_seed(1)
        m = GPinv.stvgp.StVGP(self.X.reshape(-1,1), self.Y.reshape(-1,1),
                    kern = GPinv.kernels.RBF(1,output_dim=1),
                    likelihood=GPinv.likelihoods.Gaussian())
        trainer = tf.train.AdamOptimizer(learning_rate=0.002)
        # first optimize by scipy
        m.optimize()
        # Stochastic optimization by tf.train
        rslt = m.optimize(trainer, maxiter=3000)
        print(rslt['fun'], rslt_ref['fun'])
        self.assertTrue(np.allclose(rslt['fun'], rslt_ref['fun'], atol=0.5))
        print(m.kern)
        print(m_ref.kern)
        print(m.likelihood)
        print(m_ref.likelihood)
        self.assertTrue(np.allclose(
            m.kern.lengthscales.value, m_ref.kern.lengthscales.value, rtol=0.2))
        self.assertTrue(np.allclose(
            m.kern.variance.value, m_ref.kern.variance.value, rtol=0.4))
        self.assertTrue(np.allclose(
            m.likelihood.variance.value, m_ref.likelihood.variance.value, rtol=0.2))
        # test prediction
        Xnew = np.linspace(0.,1.,22)
        mu, var = m.predict_f(Xnew.reshape(-1,1))
        mu_ref, var_ref = m_ref.predict_f(Xnew.reshape(-1,1))
        print(mu.flatten())
        print(mu_ref.flatten())
        print(var.flatten())
        print(var_ref.flatten())
        self.assertTrue(np.allclose(mu, mu_ref, atol=0.01))
        self.assertTrue(np.allclose(var, var_ref, atol=0.001))

if __name__ == '__main__':
    unittest.main()
<commit_msg>Test improved again.<commit_after>import numpy as np
import unittest
import tensorflow as tf
import GPflow
import GPinv

class test_vgp(unittest.TestCase):
    def setUp(self):
        rng = np.random.RandomState(0)
        self.X = np.linspace(0.,1.,20)
        self.Y = np.sin(3.*self.X) + rng.randn(20)*0.1

    def test_build(self):
        m = GPinv.stvgp.StVGP(self.X.reshape(-1,1), self.Y.reshape(-1,1),
                    GPinv.kernels.RBF(1,output_dim=1),
                    GPinv.likelihoods.Gaussian())
        m._compile()

    def test_optimize(self):
        # reference GPR
        m_ref = GPflow.gpr.GPR(
                    self.X.reshape(-1,1),
                    self.Y.reshape(-1,1),
                    kern = GPflow.kernels.RBF(1))
        rslt_ref = m_ref.optimize()
        # tested StVGP
        tf.set_random_seed(1)
        m = GPinv.stvgp.StVGP(self.X.reshape(-1,1), self.Y.reshape(-1,1),
                    kern = GPinv.kernels.RBF(1,output_dim=1),
                    likelihood=GPinv.likelihoods.Gaussian())
        trainer = tf.train.AdamOptimizer(learning_rate=0.002)
        # first optimize by scipy
        m.optimize()
        # Stochastic optimization by tf.train
        rslt = m.optimize(trainer, maxiter=3000)
        print(rslt['fun'], rslt_ref['fun'])
        self.assertTrue(np.allclose(rslt['fun'], rslt_ref['fun'], atol=0.5))
        print(m.kern)
        print(m_ref.kern)
        print(m.likelihood)
        print(m_ref.likelihood)
        self.assertTrue(np.allclose(
            m.kern.lengthscales.value, m_ref.kern.lengthscales.value, rtol=0.2))
        self.assertTrue(np.allclose(
            m.kern.variance.value, m_ref.kern.variance.value, rtol=0.4))
        self.assertTrue(np.allclose(
            m.likelihood.variance.value, m_ref.likelihood.variance.value, rtol=0.2))
        # test prediction
        Xnew = np.linspace(0.,1.,22)
        mu, var = m.predict_f(Xnew.reshape(-1,1))
        mu_ref, var_ref = m_ref.predict_f(Xnew.reshape(-1,1))
        print(mu.flatten())
        print(mu_ref.flatten())
        print(var.flatten())
        print(var_ref.flatten())
        self.assertTrue(np.allclose(mu, mu_ref, atol=0.03))
        self.assertTrue(np.allclose(var, var_ref, atol=0.003))

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


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Apache Beam SDK for Python setup file."""

from distutils.version import StrictVersion

import os
import platform
import warnings

import setuptools

from pkg_resources import get_distribution, DistributionNotFound

def get_version():
  global_names = {}
  execfile(os.path.normpath('./apache_beam/version.py'),
           global_names)
  return global_names['__version__']

PACKAGE_NAME = 'apache-beam-sdk'
PACKAGE_VERSION = get_version()
PACKAGE_DESCRIPTION = 'Apache Beam SDK for Python'
PACKAGE_URL = 'https://beam.incubator.apache.org'
PACKAGE_DOWNLOAD_URL = 'TBD'
PACKAGE_AUTHOR = 'Apache Software Foundation'
PACKAGE_EMAIL = 'dev@beam.incubator.apache.org'
PACKAGE_KEYWORDS = 'apache beam'
PACKAGE_LONG_DESCRIPTION = '''
TBD
'''

REQUIRED_PIP_VERSION = '7.0.0'
_PIP_VERSION = get_distribution('pip').version
if StrictVersion(_PIP_VERSION) < StrictVersion(REQUIRED_PIP_VERSION):
  warnings.warn(
      "You are using version {0} of pip. " \
      "However, version {1} is recommended.".format(
          _PIP_VERSION, REQUIRED_PIP_VERSION
      )
  )

REQUIRED_CYTHON_VERSION = '0.23.2'
try:
  _CYTHON_VERSION = get_distribution('cython').version
  if StrictVersion(_CYTHON_VERSION) < StrictVersion(REQUIRED_CYTHON_VERSION):
    warnings.warn(
        "You are using version {0} of cython. " \
        "However, version {1} is recommended.".format(
            _CYTHON_VERSION, REQUIRED_CYTHON_VERSION
        )
    )
except DistributionNotFound:
  # do nothing if Cython is not installed
  pass

# Currently all compiled modules are optional  (for performance only).
if platform.system() == 'Windows':
  # Windows doesn't always provide int64_t.
  cythonize = lambda *args, **kwargs: []
else:
  try:
    # pylint: disable=wrong-import-position
    from Cython.Build import cythonize
  except ImportError:
    cythonize = lambda *args, **kwargs: []

REQUIRED_PACKAGES = [
    'avro>=1.7.7,<2.0.0',
    'dill>=0.2.5,<0.3',
    'google-apitools>=0.5.2,<1.0.0',
    'googledatastore==6.4.0',
    'httplib2>=0.8,<0.10',
    'mock>=1.0.1,<3.0.0',
    'oauth2client>=2.0.1,<4.0.0',
    'protobuf==3.0.0',
    'protorpc>=0.9.1,<0.12',
    'python-gflags>=2.0,<4.0.0',
    'pyyaml>=3.10,<4.0.0',
    ]

REQUIRED_TEST_PACKAGES = [
    'pyhamcrest>=1.9,<2.0',
    ]

setuptools.setup(
    name=PACKAGE_NAME,
    version=PACKAGE_VERSION,
    description=PACKAGE_DESCRIPTION,
    long_description=PACKAGE_LONG_DESCRIPTION,
    url=PACKAGE_URL,
    download_url=PACKAGE_DOWNLOAD_URL,
    author=PACKAGE_AUTHOR,
    author_email=PACKAGE_EMAIL,
    packages=setuptools.find_packages(),
    package_data={'apache_beam': ['**/*.pyx', '**/*.pxd']},
    ext_modules=cythonize([
        '**/*.pyx',
        'apache_beam/coders/coder_impl.py',
        'apache_beam/runners/common.py',
        'apache_beam/transforms/cy_combiners.py',
        'apache_beam/utils/counters.py',
        'apache_beam/utils/windowed_value.py',
    ]),
    setup_requires=['nose>=1.0'],
    install_requires=REQUIRED_PACKAGES,
    test_suite='nose.collector',
    tests_require=REQUIRED_TEST_PACKAGES,
    zip_safe=False,
    # PyPI package information.
    classifiers=[
        'Intended Audience :: End Users/Desktop',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python :: 2.7',
        'Topic :: Software Development :: Libraries',
        'Topic :: Software Development :: Libraries :: Python Modules',
        ],
    license='Apache License, Version 2.0',
    keywords=PACKAGE_KEYWORDS,
    entry_points={
        'nose.plugins.0.10': [
            'beam_test_plugin = test_config:BeamTestPlugin'
            ]}
    )
<commit_msg>Closes #1453<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Apache Beam SDK for Python setup file."""

from distutils.version import StrictVersion

import os
import platform
import warnings

import setuptools

from pkg_resources import get_distribution, DistributionNotFound

def get_version():
  global_names = {}
  execfile(os.path.normpath('./apache_beam/version.py'),
           global_names)
  return global_names['__version__']

PACKAGE_NAME = 'apache-beam-sdk'
PACKAGE_VERSION = get_version()
PACKAGE_DESCRIPTION = 'Apache Beam SDK for Python'
PACKAGE_URL = 'https://beam.incubator.apache.org'
PACKAGE_DOWNLOAD_URL = 'TBD'
PACKAGE_AUTHOR = 'Apache Software Foundation'
PACKAGE_EMAIL = 'dev@beam.incubator.apache.org'
PACKAGE_KEYWORDS = 'apache beam'
PACKAGE_LONG_DESCRIPTION = '''
TBD
'''

REQUIRED_PIP_VERSION = '7.0.0'
_PIP_VERSION = get_distribution('pip').version
if StrictVersion(_PIP_VERSION) < StrictVersion(REQUIRED_PIP_VERSION):
  warnings.warn(
      "You are using version {0} of pip. " \
      "However, version {1} is recommended.".format(
          _PIP_VERSION, REQUIRED_PIP_VERSION
      )
  )

REQUIRED_CYTHON_VERSION = '0.23.2'
try:
  _CYTHON_VERSION = get_distribution('cython').version
  if StrictVersion(_CYTHON_VERSION) < StrictVersion(REQUIRED_CYTHON_VERSION):
    warnings.warn(
        "You are using version {0} of cython. " \
        "However, version {1} is recommended.".format(
            _CYTHON_VERSION, REQUIRED_CYTHON_VERSION
        )
    )
except DistributionNotFound:
  # do nothing if Cython is not installed
  pass

# Currently all compiled modules are optional  (for performance only).
if platform.system() == 'Windows':
  # Windows doesn't always provide int64_t.
  cythonize = lambda *args, **kwargs: []
else:
  try:
    # pylint: disable=wrong-import-position
    from Cython.Build import cythonize
  except ImportError:
    cythonize = lambda *args, **kwargs: []

REQUIRED_PACKAGES = [
    'avro>=1.7.7,<2.0.0',
    'dill>=0.2.5,<0.3',
    'google-apitools>=0.5.2,<1.0.0',
    'googledatastore>=6.4.1,<7.0.0',
    'httplib2>=0.8,<0.10',
    'mock>=1.0.1,<3.0.0',
    'oauth2client>=2.0.1,<4.0.0',
    'protobuf==3.0.0',
    'protorpc>=0.9.1,<0.12',
    'python-gflags>=2.0,<4.0.0',
    'pyyaml>=3.10,<4.0.0',
    ]

REQUIRED_TEST_PACKAGES = [
    'pyhamcrest>=1.9,<2.0',
    ]

setuptools.setup(
    name=PACKAGE_NAME,
    version=PACKAGE_VERSION,
    description=PACKAGE_DESCRIPTION,
    long_description=PACKAGE_LONG_DESCRIPTION,
    url=PACKAGE_URL,
    download_url=PACKAGE_DOWNLOAD_URL,
    author=PACKAGE_AUTHOR,
    author_email=PACKAGE_EMAIL,
    packages=setuptools.find_packages(),
    package_data={'apache_beam': ['**/*.pyx', '**/*.pxd']},
    ext_modules=cythonize([
        '**/*.pyx',
        'apache_beam/coders/coder_impl.py',
        'apache_beam/runners/common.py',
        'apache_beam/transforms/cy_combiners.py',
        'apache_beam/utils/counters.py',
        'apache_beam/utils/windowed_value.py',
    ]),
    setup_requires=['nose>=1.0'],
    install_requires=REQUIRED_PACKAGES,
    test_suite='nose.collector',
    tests_require=REQUIRED_TEST_PACKAGES,
    zip_safe=False,
    # PyPI package information.
    classifiers=[
        'Intended Audience :: End Users/Desktop',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python :: 2.7',
        'Topic :: Software Development :: Libraries',
        'Topic :: Software Development :: Libraries :: Python Modules',
        ],
    license='Apache License, Version 2.0',
    keywords=PACKAGE_KEYWORDS,
    entry_points={
        'nose.plugins.0.10': [
            'beam_test_plugin = test_config:BeamTestPlugin'
            ]}
    )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Wed Dec  5 21:44:22 2012

@author: cryan
"""

import numpy as np
from numpy import sin, cos

from scipy.constants import pi
from scipy.linalg import expm, eigh

from PySim.SystemParams import SystemParams
from PySim.PulseSequence import PulseSequence
from PySim.Simulation import simulate_sequence_stack, simulate_sequence
from PySim.QuantumSystems import SCQubit, Hamiltonian, Dissipator

from numba import *

import matplotlib.pyplot as plt
from timeit import timeit

#Try to load the CPPBackEnd
try:
    import PySim.CySim
    CPPBackEnd = True
except ImportError:
    CPPBackEnd = False

#@jit(c16[:,:](c16[:,:], c16))
def expm_eigen(matIn, mult):
    '''
    Helper function to compute matrix exponential of Hermitian matrix
    '''
    dim = matIn.shape[0]
    D, V = eigh(matIn)
    return np.dot(V, np.exp(mult*D).repeat(dim).reshape((dim, dim))*V.conj().T), D, V

#@jit(c16[:,:](c16[:,:], c16[:,:,:], f8[:,:], f8[:]))
#@autojit
def evolution_numpy(Hnat, controlHams, controlFields, controlFreqs):

    timeStep = 0.01
    curTime = 0.0

    Uprop = np.eye(Hnat.shape[0])
    for timect in range(controlFields.shape[1]):
        tmpH = np.copy(Hnat)
        for controlct in range(controlFields.shape[0]):
            tmpH += controlFields[controlct, timect]*cos(2*pi*curTime*controlFreqs[controlct])*controlHams[controlct]

        Uprop = np.dot(expm_eigen(tmpH,-1j*2*pi*timeStep)[0],Uprop)
        curTime += timeStep

    return Uprop

@autojit()
def sum1d(my_double_array):
    sum = 0.0
    for i in range(my_double_array.shape[0]):
        sum += my_double_array[i]
    return sum

def sim_setup(dimension, numTimeSteps, numControls):
    #Create a random natural hamiltonian
    tmpMat = np.random.randn(dimension, dimension) + 1j*np.random.randn(dimension, dimension)
    Hnat = tmpMat+tmpMat.conj().T

    #Create random control Hamiltonians
    controlHams = np.zeros((numControls,dimension, dimension), dtype=np.complex128)
    for ct in range(numControls):
        tmpMat = np.random.randn(dimension, dimension) + 1j*np.random.randn(dimension, dimension)
        controlHams[ct] = tmpMat+tmpMat.conj().T

    #Create random controlfields
    controlFields = np.random.randn(numControls, numTimeSteps)

    #Control frequencies
    controlFreqs = np.random.randn(numControls)

    return Hnat, controlHams, controlFields, controlFreqs

def sim_setup_cython(Hnat, controlHams, controlFields, controlFreqs):
    systemParams = SystemParams()
    systemParams.Hnat = Hamiltonian(Hnat)
    pulseSeq = PulseSequence()
    pulseSeq.controlAmps = controlFields
    for ct in range(len(controlHams)):
        systemParams.add_control_ham(inphase=Hamiltonian(controlHams[ct]))
        pulseSeq.add_control_line(freq = controlFreqs[ct], phase=0, controlType='sinusoidal')
    for ct in range(np.int(np.log2(Hnat.shape[0]))):
        systemParams.add_sub_system(SCQubit(2,0e9, name='Q1', T1=1e-6))
    pulseSeq.timeSteps = 0.01*np.ones(controlFields.shape[1])
    pulseSeq.maxTimeStep = 1e6

    return systemParams, pulseSeq

if __name__ == '__main__':

    dims = 2**np.arange(1,6)
    cythonTimes = []
    numpyTimes = []
    for dim in dims:
        print(dim)
        Hnat, controlHams, controlFields, controlFreqs = sim_setup(dim, 2000, 4)
        systemParams, pulseSeq = sim_setup_cython(Hnat, controlHams, controlFields, controlFreqs)
        numpyTimes.append(timeit('evolution_numpy(Hnat, controlHams, controlFields, controlFreqs)',
                      setup='from __main__ import evolution_numpy, Hnat, controlHams, controlFields, controlFreqs', number=3)/3)
        cythonTimes.append(timeit('simulate_sequence(pulseSeq, systemParams)', setup='from __main__ import simulate_sequence, pulseSeq, systemParams', number=3)/3)

    plt.plot(dims, numpyTimes)
    plt.plot(dims, cythonTimes)
    plt.legend(('Numpy', 'Cython'))
    plt.xlabel('System Dimension')
    plt.show()

<commit_msg>A few percent faster with this expm_eigen.<commit_after># -*- coding: utf-8 -*-
"""
Created on Wed Dec  5 21:44:22 2012

@author: cryan
"""

import numpy as np
from numpy import sin, cos

from scipy.constants import pi
from scipy.linalg import expm, eigh

from PySim.SystemParams import SystemParams
from PySim.PulseSequence import PulseSequence
from PySim.Simulation import simulate_sequence_stack, simulate_sequence
from PySim.QuantumSystems import SCQubit, Hamiltonian, Dissipator

from numba import *

import matplotlib.pyplot as plt
from timeit import timeit

#Try to load the CPPBackEnd
try:
    import PySim.CySim
    CPPBackEnd = True
except ImportError:
    CPPBackEnd = False

#@jit(c16[:,:](c16[:,:], c16))
def expm_eigen(matIn, mult):
    '''
    Helper function to compute matrix exponential of Hermitian matrix
    '''
    dim = matIn.shape[0]
    D, V = eigh(matIn)
    return V.dot(np.diag(np.exp(mult*D))).dot(V.conj().T)

#@jit(c16[:,:](c16[:,:], c16[:,:,:], f8[:,:], f8[:]))
#@autojit
def evolution_numpy(Hnat, controlHams, controlFields, controlFreqs):

    timeStep = 0.01
    curTime = 0.0

    Uprop = np.eye(Hnat.shape[0])
    for timect in range(controlFields.shape[1]):
        tmpH = np.copy(Hnat)
        for controlct in range(controlFields.shape[0]):
            tmpH += controlFields[controlct, timect]*cos(2*pi*curTime*controlFreqs[controlct])*controlHams[controlct]

        Uprop = np.dot(expm_eigen(tmpH,-1j*2*pi*timeStep)[0],Uprop)
        curTime += timeStep

    return Uprop

@autojit()
def sum1d(my_double_array):
    sum = 0.0
    for i in range(my_double_array.shape[0]):
        sum += my_double_array[i]
    return sum

def sim_setup(dimension, numTimeSteps, numControls):
    #Create a random natural hamiltonian
    tmpMat = np.random.randn(dimension, dimension) + 1j*np.random.randn(dimension, dimension)
    Hnat = tmpMat+tmpMat.conj().T

    #Create random control Hamiltonians
    controlHams = np.zeros((numControls,dimension, dimension), dtype=np.complex128)
    for ct in range(numControls):
        tmpMat = np.random.randn(dimension, dimension) + 1j*np.random.randn(dimension, dimension)
        controlHams[ct] = tmpMat+tmpMat.conj().T

    #Create random controlfields
    controlFields = np.random.randn(numControls, numTimeSteps)

    #Control frequencies
    controlFreqs = np.random.randn(numControls)

    return Hnat, controlHams, controlFields, controlFreqs

def sim_setup_cython(Hnat, controlHams, controlFields, controlFreqs):
    systemParams = SystemParams()
    systemParams.Hnat = Hamiltonian(Hnat)
    pulseSeq = PulseSequence()
    pulseSeq.controlAmps = controlFields
    for ct in range(len(controlHams)):
        systemParams.add_control_ham(inphase=Hamiltonian(controlHams[ct]))
        pulseSeq.add_control_line(freq = controlFreqs[ct], phase=0, controlType='sinusoidal')
    for ct in range(np.int(np.log2(Hnat.shape[0]))):
        systemParams.add_sub_system(SCQubit(2,0e9, name='Q1', T1=1e-6))
    pulseSeq.timeSteps = 0.01*np.ones(controlFields.shape[1])
    pulseSeq.maxTimeStep = 1e6

    return systemParams, pulseSeq

if __name__ == '__main__':

    dims = 2**np.arange(1,6)
    cythonTimes = []
    numpyTimes = []
    for dim in dims:
        print(dim)
        Hnat, controlHams, controlFields, controlFreqs = sim_setup(dim, 2000, 4)
        systemParams, pulseSeq = sim_setup_cython(Hnat, controlHams, controlFields, controlFreqs)
        numpyTimes.append(timeit('evolution_numpy(Hnat, controlHams, controlFields, controlFreqs)',
                      setup='from __main__ import evolution_numpy, Hnat, controlHams, controlFields, controlFreqs', number=3)/3)
        cythonTimes.append(timeit('simulate_sequence(pulseSeq, systemParams)', setup='from __main__ import simulate_sequence, pulseSeq, systemParams', number=3)/3)

    plt.plot(dims, numpyTimes)
    plt.plot(dims, cythonTimes)
    plt.legend(('Numpy', 'Cython'))
    plt.xlabel('System Dimension')
    plt.show()

<|endoftext|>


<commit_before># calculate
print(1-2)
print(4*5)
print(7/5)
print(3**2)

# check type
print(type(10))
print(type(2.78))
print(type("hello"))

# use variable
x = 10
print(x)
x = 100
print(x)<commit_msg>Add python list code<commit_after># calculate
print(1-2)
print(4*5)
print(7/5)
print(3**2)

# check type
print(type(10))
print(type(2.78))
print(type("hello"))

# use variable
x = 10
print(x)
x = 100
print(x)
y = 3.14
print(x*y)
print(type(x*y))

# list
a = [1, 2, 3, 4, 5]
print(a)
print(len(a))
print(a[0])
print(a[4])
a[4] = 99
print(a)
print(a[0:2])
print(a[1:])
print(a[:3])
print(a[:-1])
print(a[:-2])
<|endoftext|>


<commit_before>from base64 import urlsafe_b64encode, urlsafe_b64decode
from zlib import compress, decompress
from sys import version

from searx import settings, autocomplete
from searx.languages import language_codes as languages
from searx.url_utils import parse_qs, urlencode

if version[0] == '3':
    unicode = str

COOKIE_MAX_AGE = 60 * 60 * 24 * 365 * 5  # 5 years
LANGUAGE_CODES = [l[0] for l in languages]
LANGUAGE_CODES.append('all')
DISABLED = 0
ENABLED = 1

class MissingArgumentException(Exception):
    pass

class ValidationException(Exception):
    pass

class Setting(object):
    """Base class of user settings"""

    def __init__(self, default_value, **kwargs):
        super(Setting, self).__init__()
        self.value = default_value
        for key, value in kwargs.items():
            setattr(self, key, value)

        self._post_init()

    def _post_init(self):
        pass

    def parse(self, data):
        self.value = data

    def get_value(self):
        return self.value

    def save(self, name, resp):
        resp.set_cookie(name, self.value, max_age=COOKIE_MAX_AGE)

class StringSetting(Setting):
    """Setting of plain string values"""
    pass

class EnumStringSetting(Setting):
    """Setting of a value which can only come from the given choices"""

    def _validate_selection(self, selection):
        if selection not in self.choices:
            raise ValidationException('Invalid value: "{0}"'.format(selection))

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        self._validate_selection(self.value)

    def parse(self, data):
        self._validate_selection(data)
        self.value = data

class MultipleChoiceSetting(EnumStringSetting):
    """Setting of values which can only come from the given choices"""

    def _validate_selections(self, selections):
        for item in selections:
            if item not in self.choices:
                raise ValidationException('Invalid value: "{0}"'.format(selections))

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        self._validate_selections(self.value)

    def parse(self, data):
        if data == '':
            self.value = []
            return

        elements = data.split(',')
        self._validate_selections(elements)
        self.value = elements

    def parse_form(self, data):
        self.value = []
        for choice in data:
            if choice in self.choices and choice not in self.value:
                self.value.append(choice)

    def save(self, name, resp):
        resp.set_cookie(name, ','.join(self.value), max_age=COOKIE_MAX_AGE)

class SearchLanguageSetting(EnumStringSetting):
    """Available choices may change, so user's value may not be in choices anymore"""

    def parse(self, data):
        if data not in self.choices and data != self.value:
            # hack to give some backwards compatibility with old language cookies
            data = str(data).replace('_', '-')
            lang = data.split('-')[0]
            if data in self.choices:
                pass
            elif lang in self.choices:
                data = lang
            elif data == 'nb-NO':
                data = 'no-NO'
            elif data == 'ar-XA':
                data = 'ar-SA'
            else:
                data = self.value
        self.value = data

class MapSetting(Setting):
    """Setting of a value that has to be translated in order to be storable"""

    def _post_init(self):
        if not hasattr(self, 'map'):
            raise MissingArgumentException('missing argument: map')
        if self.value not in self.map.values():
            raise ValidationException('Invalid default value')

    def parse(self, data):
        if data not in self.map:
            raise ValidationException('Invalid choice: {0}'.format(data))
        self.value = self.map[data]
        self.key = data

    def save(self, name, resp):
        if hasattr(self, 'key'):
            resp.set_cookie(name, self.key, max_age=COOKIE_MAX_AGE)

class SwitchableSetting(Setting):
    """ Base class for settings that can be turned on && off"""

    def _post_init(self):
        self.disabled = set()
        self.enabled = set()
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('missing argument: choices')

    def transform_form_items(self, items):
        return items

    def transform_values(self, values):
        return values

    def parse_cookie(self, data):
        if data[DISABLED] != '':
            self.disabled = set(data[DISABLED].split(','))
        if data[ENABLED] != '':
            self.enabled = set(data[ENABLED].split(','))

    def parse_form(self, items):
        items = self.transform_form_items(items)

        self.disabled = set()
        self.enabled = set()
        for choice in self.choices:
            if choice['default_on']:
                if choice['id'] in items:
                    self.disabled.add(choice['id'])
            else:
                if choice['id'] not in items:
                    self.enabled.add(choice['id'])

    def save(self, resp):
        resp.set_cookie('disabled_{0}'.format(self.value), ','.join(self.disabled), max_age=COOKIE_MAX_AGE)
        resp.set_cookie('enabled_{0}'.format(self.value), ','.join(self.enabled), max_age=COOKIE_MAX_AGE)

    def get_disabled(self):
        disabled = self.disabled
        for choice in self.choices:
            if not choice['default_on'] and choice['id'] not in self.enabled:
                disabled.add(choice['id'])
        return self.transform_values(disabled)

    def get_enabled(self):
        enabled = self.enabled
        for choice in self.choices:
            if choice['default_on'] and choice['id'] not in self.disabled:
                enabled.add(choice['id'])
        return self.transform_values(enabled)

class EnginesSetting(SwitchableSetting):

    def _post_init(self):
        super(EnginesSetting, self)._post_init()
        transformed_choices = []
        for engine_name, engine in self.choices.items():
            for category in engine.categories:
                transformed_choice = dict()
                transformed_choice['default_on'] = not engine.disabled
                transformed_choice['id'] = '{}__{}'.format(engine_name, category)
                transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('engine_'):].replace('_', ' ').replace('  ', '__') for item in items]

    def transform_values(self, values):
        if len(values) == 1 and next(iter(values)) == '':
            return list()
        transformed_values = []
        for value in values:
            engine, category = value.split('__')
            transformed_values.append((engine, category))
        return transformed_values

class PluginsSetting(SwitchableSetting):

    def _post_init(self):
        super(PluginsSetting, self)._post_init()
        transformed_choices = []
        for plugin in self.choices:
            transformed_choice = dict()
            transformed_choice['default_on'] = plugin.default_on
            transformed_choice['id'] = plugin.id
            transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('plugin_'):] for item in items]

class Preferences(object):
    """Validates and saves preferences to cookies"""

    def __init__(self, themes, categories, engines, plugins):
        super(Preferences, self).__init__()

        self.key_value_settings = {'categories': MultipleChoiceSetting(['general'], choices=categories),
                                   'language': SearchLanguageSetting(settings['search']['language'],
                                                                     choices=LANGUAGE_CODES),
                                   'locale': EnumStringSetting(settings['ui']['default_locale'],
                                                               choices=list(settings['locales'].keys()) + ['']),
                                   'autocomplete': EnumStringSetting(settings['search']['autocomplete'],
                                                                     choices=list(autocomplete.backends.keys()) + ['']),
                                   'image_proxy': MapSetting(settings['server']['image_proxy'],
                                                             map={'': settings['server']['image_proxy'],
                                                                  '0': False,
                                                                  '1': True,
                                                                  'True': True,
                                                                  'False': False}),
                                   'method': EnumStringSetting('POST', choices=('GET', 'POST')),
                                   'safesearch': MapSetting(settings['search']['safe_search'], map={'0': 0,
                                                                                '1': 1,
                                                                                '2': 2}),
                                   'theme': EnumStringSetting(settings['ui']['default_theme'], choices=themes),
                                   'results_on_new_tab': MapSetting(False, map={'0': False,
                                                                                '1': True,
                                                                                'False': False,
                                                                                'True': True})}

        self.engines = EnginesSetting('engines', choices=engines)
        self.plugins = PluginsSetting('plugins', choices=plugins)
        self.unknown_params = {}

    def get_as_url_params(self):
        settings_kv = {}
        for k, v in self.key_value_settings.items():
            if isinstance(v, MultipleChoiceSetting):
                settings_kv[k] = ','.join(v.get_value())
            else:
                settings_kv[k] = v.get_value()

        settings_kv['disabled_engines'] = ','.join(self.engines.disabled)
        settings_kv['enabled_engines'] = ','.join(self.engines.enabled)

        settings_kv['disabled_plugins'] = ','.join(self.plugins.disabled)
        settings_kv['enabled_plugins'] = ','.join(self.plugins.enabled)

        return urlsafe_b64encode(compress(urlencode(settings_kv).encode('utf-8'))).decode('utf-8')

    def parse_encoded_data(self, input_data):
        decoded_data = decompress(urlsafe_b64decode(input_data.encode('utf-8')))
        self.parse_dict({x: y[0] for x,y in parse_qs(unicode(decoded_data)).items()})

    def parse_dict(self, input_data):
        for user_setting_name, user_setting in input_data.items():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name == 'disabled_engines':
                self.engines.parse_cookie((input_data.get('disabled_engines', ''),
                                           input_data.get('enabled_engines', '')))
            elif user_setting_name == 'disabled_plugins':
                self.plugins.parse_cookie((input_data.get('disabled_plugins', ''),
                                           input_data.get('enabled_plugins', '')))

    def parse_form(self, input_data):
        disabled_engines = []
        enabled_categories = []
        disabled_plugins = []
        for user_setting_name, user_setting in input_data.items():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name.startswith('engine_'):
                disabled_engines.append(user_setting_name)
            elif user_setting_name.startswith('category_'):
                enabled_categories.append(user_setting_name[len('category_'):])
            elif user_setting_name.startswith('plugin_'):
                disabled_plugins.append(user_setting_name)
            else:
                self.unknown_params[user_setting_name] = user_setting
        self.key_value_settings['categories'].parse_form(enabled_categories)
        self.engines.parse_form(disabled_engines)
        self.plugins.parse_form(disabled_plugins)

    # cannot be used in case of engines or plugins
    def get_value(self, user_setting_name):
        if user_setting_name in self.key_value_settings:
            return self.key_value_settings[user_setting_name].get_value()

    def save(self, resp):
        for user_setting_name, user_setting in self.key_value_settings.items():
            user_setting.save(user_setting_name, resp)
        self.engines.save(resp)
        self.plugins.save(resp)
        for k, v in self.unknown_params.items():
            resp.set_cookie(k, v, max_age=COOKIE_MAX_AGE)
        return resp
<commit_msg>[fix] pep8<commit_after>from base64 import urlsafe_b64encode, urlsafe_b64decode
from zlib import compress, decompress
from sys import version

from searx import settings, autocomplete
from searx.languages import language_codes as languages
from searx.url_utils import parse_qs, urlencode

if version[0] == '3':
    unicode = str

COOKIE_MAX_AGE = 60 * 60 * 24 * 365 * 5  # 5 years
LANGUAGE_CODES = [l[0] for l in languages]
LANGUAGE_CODES.append('all')
DISABLED = 0
ENABLED = 1

class MissingArgumentException(Exception):
    pass

class ValidationException(Exception):
    pass

class Setting(object):
    """Base class of user settings"""

    def __init__(self, default_value, **kwargs):
        super(Setting, self).__init__()
        self.value = default_value
        for key, value in kwargs.items():
            setattr(self, key, value)

        self._post_init()

    def _post_init(self):
        pass

    def parse(self, data):
        self.value = data

    def get_value(self):
        return self.value

    def save(self, name, resp):
        resp.set_cookie(name, self.value, max_age=COOKIE_MAX_AGE)

class StringSetting(Setting):
    """Setting of plain string values"""
    pass

class EnumStringSetting(Setting):
    """Setting of a value which can only come from the given choices"""

    def _validate_selection(self, selection):
        if selection not in self.choices:
            raise ValidationException('Invalid value: "{0}"'.format(selection))

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        self._validate_selection(self.value)

    def parse(self, data):
        self._validate_selection(data)
        self.value = data

class MultipleChoiceSetting(EnumStringSetting):
    """Setting of values which can only come from the given choices"""

    def _validate_selections(self, selections):
        for item in selections:
            if item not in self.choices:
                raise ValidationException('Invalid value: "{0}"'.format(selections))

    def _post_init(self):
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('Missing argument: choices')
        self._validate_selections(self.value)

    def parse(self, data):
        if data == '':
            self.value = []
            return

        elements = data.split(',')
        self._validate_selections(elements)
        self.value = elements

    def parse_form(self, data):
        self.value = []
        for choice in data:
            if choice in self.choices and choice not in self.value:
                self.value.append(choice)

    def save(self, name, resp):
        resp.set_cookie(name, ','.join(self.value), max_age=COOKIE_MAX_AGE)

class SearchLanguageSetting(EnumStringSetting):
    """Available choices may change, so user's value may not be in choices anymore"""

    def parse(self, data):
        if data not in self.choices and data != self.value:
            # hack to give some backwards compatibility with old language cookies
            data = str(data).replace('_', '-')
            lang = data.split('-')[0]
            if data in self.choices:
                pass
            elif lang in self.choices:
                data = lang
            elif data == 'nb-NO':
                data = 'no-NO'
            elif data == 'ar-XA':
                data = 'ar-SA'
            else:
                data = self.value
        self.value = data

class MapSetting(Setting):
    """Setting of a value that has to be translated in order to be storable"""

    def _post_init(self):
        if not hasattr(self, 'map'):
            raise MissingArgumentException('missing argument: map')
        if self.value not in self.map.values():
            raise ValidationException('Invalid default value')

    def parse(self, data):
        if data not in self.map:
            raise ValidationException('Invalid choice: {0}'.format(data))
        self.value = self.map[data]
        self.key = data

    def save(self, name, resp):
        if hasattr(self, 'key'):
            resp.set_cookie(name, self.key, max_age=COOKIE_MAX_AGE)

class SwitchableSetting(Setting):
    """ Base class for settings that can be turned on && off"""

    def _post_init(self):
        self.disabled = set()
        self.enabled = set()
        if not hasattr(self, 'choices'):
            raise MissingArgumentException('missing argument: choices')

    def transform_form_items(self, items):
        return items

    def transform_values(self, values):
        return values

    def parse_cookie(self, data):
        if data[DISABLED] != '':
            self.disabled = set(data[DISABLED].split(','))
        if data[ENABLED] != '':
            self.enabled = set(data[ENABLED].split(','))

    def parse_form(self, items):
        items = self.transform_form_items(items)

        self.disabled = set()
        self.enabled = set()
        for choice in self.choices:
            if choice['default_on']:
                if choice['id'] in items:
                    self.disabled.add(choice['id'])
            else:
                if choice['id'] not in items:
                    self.enabled.add(choice['id'])

    def save(self, resp):
        resp.set_cookie('disabled_{0}'.format(self.value), ','.join(self.disabled), max_age=COOKIE_MAX_AGE)
        resp.set_cookie('enabled_{0}'.format(self.value), ','.join(self.enabled), max_age=COOKIE_MAX_AGE)

    def get_disabled(self):
        disabled = self.disabled
        for choice in self.choices:
            if not choice['default_on'] and choice['id'] not in self.enabled:
                disabled.add(choice['id'])
        return self.transform_values(disabled)

    def get_enabled(self):
        enabled = self.enabled
        for choice in self.choices:
            if choice['default_on'] and choice['id'] not in self.disabled:
                enabled.add(choice['id'])
        return self.transform_values(enabled)

class EnginesSetting(SwitchableSetting):

    def _post_init(self):
        super(EnginesSetting, self)._post_init()
        transformed_choices = []
        for engine_name, engine in self.choices.items():
            for category in engine.categories:
                transformed_choice = dict()
                transformed_choice['default_on'] = not engine.disabled
                transformed_choice['id'] = '{}__{}'.format(engine_name, category)
                transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('engine_'):].replace('_', ' ').replace('  ', '__') for item in items]

    def transform_values(self, values):
        if len(values) == 1 and next(iter(values)) == '':
            return list()
        transformed_values = []
        for value in values:
            engine, category = value.split('__')
            transformed_values.append((engine, category))
        return transformed_values

class PluginsSetting(SwitchableSetting):

    def _post_init(self):
        super(PluginsSetting, self)._post_init()
        transformed_choices = []
        for plugin in self.choices:
            transformed_choice = dict()
            transformed_choice['default_on'] = plugin.default_on
            transformed_choice['id'] = plugin.id
            transformed_choices.append(transformed_choice)
        self.choices = transformed_choices

    def transform_form_items(self, items):
        return [item[len('plugin_'):] for item in items]

class Preferences(object):
    """Validates and saves preferences to cookies"""

    def __init__(self, themes, categories, engines, plugins):
        super(Preferences, self).__init__()

        self.key_value_settings = {'categories': MultipleChoiceSetting(['general'], choices=categories),
                                   'language': SearchLanguageSetting(settings['search']['language'],
                                                                     choices=LANGUAGE_CODES),
                                   'locale': EnumStringSetting(settings['ui']['default_locale'],
                                                               choices=list(settings['locales'].keys()) + ['']),
                                   'autocomplete': EnumStringSetting(settings['search']['autocomplete'],
                                                                     choices=list(autocomplete.backends.keys()) + ['']),
                                   'image_proxy': MapSetting(settings['server']['image_proxy'],
                                                             map={'': settings['server']['image_proxy'],
                                                                  '0': False,
                                                                  '1': True,
                                                                  'True': True,
                                                                  'False': False}),
                                   'method': EnumStringSetting('POST', choices=('GET', 'POST')),
                                   'safesearch': MapSetting(settings['search']['safe_search'], map={'0': 0,
                                                                                '1': 1,
                                                                                '2': 2}),
                                   'theme': EnumStringSetting(settings['ui']['default_theme'], choices=themes),
                                   'results_on_new_tab': MapSetting(False, map={'0': False,
                                                                                '1': True,
                                                                                'False': False,
                                                                                'True': True})}

        self.engines = EnginesSetting('engines', choices=engines)
        self.plugins = PluginsSetting('plugins', choices=plugins)
        self.unknown_params = {}

    def get_as_url_params(self):
        settings_kv = {}
        for k, v in self.key_value_settings.items():
            if isinstance(v, MultipleChoiceSetting):
                settings_kv[k] = ','.join(v.get_value())
            else:
                settings_kv[k] = v.get_value()

        settings_kv['disabled_engines'] = ','.join(self.engines.disabled)
        settings_kv['enabled_engines'] = ','.join(self.engines.enabled)

        settings_kv['disabled_plugins'] = ','.join(self.plugins.disabled)
        settings_kv['enabled_plugins'] = ','.join(self.plugins.enabled)

        return urlsafe_b64encode(compress(urlencode(settings_kv).encode('utf-8'))).decode('utf-8')

    def parse_encoded_data(self, input_data):
        decoded_data = decompress(urlsafe_b64decode(input_data.encode('utf-8')))
        self.parse_dict({x: y[0] for x, y in parse_qs(unicode(decoded_data)).items()})

    def parse_dict(self, input_data):
        for user_setting_name, user_setting in input_data.items():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name == 'disabled_engines':
                self.engines.parse_cookie((input_data.get('disabled_engines', ''),
                                           input_data.get('enabled_engines', '')))
            elif user_setting_name == 'disabled_plugins':
                self.plugins.parse_cookie((input_data.get('disabled_plugins', ''),
                                           input_data.get('enabled_plugins', '')))

    def parse_form(self, input_data):
        disabled_engines = []
        enabled_categories = []
        disabled_plugins = []
        for user_setting_name, user_setting in input_data.items():
            if user_setting_name in self.key_value_settings:
                self.key_value_settings[user_setting_name].parse(user_setting)
            elif user_setting_name.startswith('engine_'):
                disabled_engines.append(user_setting_name)
            elif user_setting_name.startswith('category_'):
                enabled_categories.append(user_setting_name[len('category_'):])
            elif user_setting_name.startswith('plugin_'):
                disabled_plugins.append(user_setting_name)
            else:
                self.unknown_params[user_setting_name] = user_setting
        self.key_value_settings['categories'].parse_form(enabled_categories)
        self.engines.parse_form(disabled_engines)
        self.plugins.parse_form(disabled_plugins)

    # cannot be used in case of engines or plugins
    def get_value(self, user_setting_name):
        if user_setting_name in self.key_value_settings:
            return self.key_value_settings[user_setting_name].get_value()

    def save(self, resp):
        for user_setting_name, user_setting in self.key_value_settings.items():
            user_setting.save(user_setting_name, resp)
        self.engines.save(resp)
        self.plugins.save(resp)
        for k, v in self.unknown_params.items():
            resp.set_cookie(k, v, max_age=COOKIE_MAX_AGE)
        return resp
<|endoftext|>


<commit_before>
import cv2 as cv
import numpy as np
import time
import sys

lcdEnable = False

if lcdEnable == True:
	import RPi.GPIO as GPIO

# Define GPIO to LCD mapping
LCD_RS = 7
LCD_E  = 8
LCD_D4 = 25
LCD_D5 = 24
LCD_D6 = 23
LCD_D7 = 18

redThresholdArea = 10000
greenThresholdArea = 10000

# Class for LCD
class LCD:

	def __init__(self, lcd_rs,lcd_en,lcd_d4,lcd_d5,lcd_d6,lcd_d7):

		self.LCD_RS = lcd_rs
		self.LCD_EN = lcd_en
		self.LCD_D4 = lcd_d4
		self.LCD_D5 = lcd_d5
		self.LCD_D6 = lcd_d6
		self.LCD_D7 = lcd_d7

		self.LCD_WIDTH = 20
		self.LCD_CHR = True
		self.LCD_CMD = False

		self.LCD_LINE_1 = 0x80
		self.LCD_LINE_2 = 0xC0

		self.E_PULSE = 0.00005
		self.E_DELAY = 0.00005

		GPIO.setmode(GPIO.BCM)
		GPIO.setup(self.LCD_EN, GPIO.OUT)  # E
		GPIO.setup(self.LCD_RS, GPIO.OUT) # RS
		GPIO.setup(self.LCD_D4, GPIO.OUT) # DB4
		GPIO.setup(self.LCD_D5, GPIO.OUT) # DB5
		GPIO.setup(self.LCD_D6, GPIO.OUT) # DB6
		GPIO.setup(self.LCD_D7, GPIO.OUT) # DB7

		self.lcd_init()

		# Initialise display
		self.lcd_init()

		# Send some test
		self.lcd_byte(LCD_LINE_1, LCD_CMD)
		self.lcd_string("Rasbperry Pi")
		self.lcd_byte(LCD_LINE_2, LCD_CMD)
		self.lcd_string("Model B")

	def lcd_init(self):
		# Initialise display
		self.lcd_byte(0x33,LCD_CMD)
		self.lcd_byte(0x32,LCD_CMD)
		self.lcd_byte(0x28,LCD_CMD)
		self.lcd_byte(0x0C,LCD_CMD)
		self.lcd_byte(0x06,LCD_CMD)
		self.lcd_byte(0x01,LCD_CMD)

	def lcd_string(message):
		# Send string to display

		message = message.ljust(LCD_WIDTH," ")

		for i in range(LCD_WIDTH):
			self.lcd_byte(ord(message[i]),LCD_CHR)

	def lcd_byte(bits, mode):
		# Send byte to data pins
		# bits = data
		# mode = True  for character
		#        False for command

		GPIO.output(LCD_RS, mode) # RS

		# High bits
		GPIO.output(LCD_D4, False)
		GPIO.output(LCD_D5, False)
		GPIO.output(LCD_D6, False)
		GPIO.output(LCD_D7, False)
		if bits&0x10==0x10:
			GPIO.output(LCD_D4, True)
		if bits&0x20==0x20:
			GPIO.output(LCD_D5, True)
		if bits&0x40==0x40:
			GPIO.output(LCD_D6, True)
		if bits&0x80==0x80:
			GPIO.output(LCD_D7, True)

		# Toggle 'Enable' pin
		time.sleep(E_DELAY)
		GPIO.output(LCD_EN, True)
		time.sleep(E_PULSE)
		GPIO.output(LCD_EN, False)
		time.sleep(E_DELAY)

		# Low bits
		GPIO.output(LCD_D4, False)
		GPIO.output(LCD_D5, False)
		GPIO.output(LCD_D6, False)
		GPIO.output(LCD_D7, False)
		if bits&0x01==0x01:
			GPIO.output(LCD_D4, True)
		if bits&0x02==0x02:
			GPIO.output(LCD_D5, True)
		if bits&0x04==0x04:
			GPIO.output(LCD_D6, True)
		if bits&0x08==0x08:
			GPIO.output(LCD_D7, True)

		# Toggle 'Enable' pin
		time.sleep(E_DELAY)
		GPIO.output(LCD_EN, True)
		time.sleep(E_PULSE)
		GPIO.output(LCD_EN, False)
		time.sleep(E_DELAY)

# Class for Image Processing
class Image:

	def __init__(self, _camera, _interval, filt, rlowHSV, rhighHSV, glowHSV, ghighHSV):

		# store
		self.cameraIndex = _camera
		self.frameInterval = _interval
		self.filt = filt
		self.capture = cv.VideoCapture(self.cameraIndex)

		#self.fwindow = cv.namedWindow("Result")
		self.redlowHSV = rlowHSV
		self.redhighHSV = rhighHSV

		self.greenlowHSV = glowHSV
		self.greenhighHSV = ghighHSV

		self.greenlastX = 0
		self.greenlastY = 0
		self.redlastX 	= 0
		self.redlastY 	= 0

	def doCaptureAndProcess(self):
		_, frame = self.capture.read()

		if self.filt == 0:
			filteredImage = frame
		elif self.filt == 1:
			filteredImage = cv.blur(frame, (5,5))
		elif self.filt == 2:
			filteredImage = cv.GaussianBlur(frame, (5,5),0)
		elif self.filt == 3:
			filteredImage = cv.medianBlur(frame, 5)
		elif self.filt == 4:
			filteredImage = cv.bilateralFilter(frame,9,75,75)

		kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE,(5,5))

		hsvImage = cv.cvtColor(filteredImage, cv.COLOR_BGR2HSV)

		RedimgThresholded = cv.inRange(hsvImage, self.redlowHSV, self.redhighHSV)
		GreenimgThresholded = cv.inRange(hsvImage, self.greenlowHSV, self.greenhighHSV)
		erosion = cv.erode(RedimgThresholded,kernel,iterations = 1)
		dilation = cv.dilate(RedimgThresholded,kernel,iterations = 1)

		dilation = cv.dilate(RedimgThresholded,kernel,iterations = 1)
		erosion = cv.erode(RedimgThresholded,kernel,iterations = 1)

		erosion = cv.erode(GreenimgThresholded,kernel,iterations = 1)
		dilation = cv.dilate(GreenimgThresholded,kernel,iterations = 1)

		dilation = cv.dilate(GreenimgThresholded,kernel,iterations = 1)
		erosion = cv.erode(GreenimgThresholded,kernel,iterations = 1)

		Redmoments = cv.moments(RedimgThresholded)
		Greenmoments = cv.moments(GreenimgThresholded)

		Redm01  = Redmoments['m01']
		Redm10  = Redmoments['m10']
		Redarea = Redmoments['m00']

		Greenm01  = Greenmoments['m01']
		Greenm10  = Greenmoments['m10']
		Greenarea = Greenmoments['m00']

		# print "RED " + str(Redarea)
		# print "GREEN " + str(Greenarea)

		if Redarea > redThresholdArea:
			#print "Red"
			posX = Redm10 / Redarea
			posY = Redm01 / Redarea

			# print "X " + str(posX)
			# print "Y " + str(posY)

			if self.redlastX >= 0 and self.redlastY >= 0 and posX >= 0  and posY >= 0:
				cv.line(filteredImage, (int(posX), int(posY)), (int(self.redlastX),int(self.redlastY)), (0,0,255), 5)

			redlastX = posX
			redlastY = posY
			if Redarea > Greenarea:
				print "RED"
				# Turn Servo Left

		if Greenarea > greenThresholdArea:
			#print "Green"
			posX = Greenm10 / Greenarea
			posY = Greenm01 / Greenarea

			# print "X " + str(posX)
			# print "Y " + str(posY)

			if self.greenlastX >= 0 and self.greenlastY >= 0 and posX >= 0  and posY >= 0:
				cv.line(filteredImage, (int(posX), int(posY)), (int(self.greenlastX),int(self.greenlastY)), (0,0,255), 5)

			greenlastX = posX
			greenlastY = posY

			if Greenarea > Redarea:
				print "GREEN"
				# Turn Servo Right

		cv.imshow("Red Thresholded Image", RedimgThresholded)
		cv.imshow("Green Thresholded Image", GreenimgThresholded)
		cv.imshow("Original Image", filteredImage)

redlowHSV = np.array([170,150,60], dtype=np.uint8)
redhighHSV = np.array([179,255,255], dtype=np.uint8)

greenlowHSV = np.array([50,150,60], dtype=np.uint8)
greenhighHSV = np.array([75,255,255], dtype=np.uint8)

obj = Image(1, 1, 1, redlowHSV, redhighHSV, greenlowHSV, greenhighHSV)

# If LCD is enabled enable lcd for display
if lcdEnable == True:
	lcd = LCD(LCD_RS, LCD_E, LCD_D4, LCD_D5,LCD_D6,LCD_D7)

while True:

	obj.doCaptureAndProcess()

	# Break Operation if key interrupt
	k = cv.waitKey(5) & 0xFF
	if k == 27:
		break<commit_msg>Code has to test in pi<commit_after>
import cv2 as cv
import numpy as np
import time
import sys

lcdEnable = False
gpioEnable = False

if lcdEnable == True:
	import RPi.GPIO as GPIO

# Define GPIO to LCD mapping
LCD_RS = 7
LCD_E  = 8
LCD_D4 = 25
LCD_D5 = 24
LCD_D6 = 23
LCD_D7 = 18

SWITCH_CAPTURE 	= 11
SWITCH_STOP		= 9

LED_GREEN		= 10
LED_RED			= 22

redThresholdArea = 10000
greenThresholdArea = 10000

def gpioInit():
	GPIO.setmode(GPIO.BCM)
	GPIO.setup(LED_GREEN, GPIO.OUT)
	GPIO.setup(LED_RED, GPIO.OUT)
	GPIO.setup(SWITCH_CAPTURE, GPIO.IN, pull_up_down = GPIO.PUD_UP)
	GPIO.setup(SWITCH_STOP, GPIO.IN, pull_up_down = GPIO.PUD_UP)
	GPIO.output(LED_RED, True)
	GPIO.output(LED_GREEN, True)

# Class for LCD
class LCD:

	def __init__(self, lcd_rs,lcd_en,lcd_d4,lcd_d5,lcd_d6,lcd_d7):

		self.LCD_RS = lcd_rs
		self.LCD_EN = lcd_en
		self.LCD_D4 = lcd_d4
		self.LCD_D5 = lcd_d5
		self.LCD_D6 = lcd_d6
		self.LCD_D7 = lcd_d7

		self.LCD_WIDTH = 20
		self.LCD_CHR = True
		self.LCD_CMD = False

		self.LCD_LINE_1 = 0x80
		self.LCD_LINE_2 = 0xC0

		self.E_PULSE = 0.00005
		self.E_DELAY = 0.00005

		GPIO.setmode(GPIO.BCM)
		GPIO.setup(self.LCD_EN, GPIO.OUT)  # E
		GPIO.setup(self.LCD_RS, GPIO.OUT) # RS
		GPIO.setup(self.LCD_D4, GPIO.OUT) # DB4
		GPIO.setup(self.LCD_D5, GPIO.OUT) # DB5
		GPIO.setup(self.LCD_D6, GPIO.OUT) # DB6
		GPIO.setup(self.LCD_D7, GPIO.OUT) # DB7

		self.lcd_init()

		# Initialise display
		self.lcd_init()

		# Send some test
		self.lcd_byte(LCD_LINE_1, LCD_CMD)
		self.lcd_string("Rasbperry Pi")
		self.lcd_byte(LCD_LINE_2, LCD_CMD)
		self.lcd_string("Model B")

	def lcd_init(self):
		# Initialise display
		self.lcd_byte(0x33,LCD_CMD)
		self.lcd_byte(0x32,LCD_CMD)
		self.lcd_byte(0x28,LCD_CMD)
		self.lcd_byte(0x0C,LCD_CMD)
		self.lcd_byte(0x06,LCD_CMD)
		self.lcd_byte(0x01,LCD_CMD)

	def lcd_string(message):
		# Send string to display

		message = message.ljust(LCD_WIDTH," ")

		for i in range(LCD_WIDTH):
			self.lcd_byte(ord(message[i]),LCD_CHR)

	def lcd_byte(bits, mode):
		# Send byte to data pins
		# bits = data
		# mode = True  for character
		#        False for command

		GPIO.output(LCD_RS, mode) # RS

		# High bits
		GPIO.output(LCD_D4, False)
		GPIO.output(LCD_D5, False)
		GPIO.output(LCD_D6, False)
		GPIO.output(LCD_D7, False)
		if bits&0x10==0x10:
			GPIO.output(LCD_D4, True)
		if bits&0x20==0x20:
			GPIO.output(LCD_D5, True)
		if bits&0x40==0x40:
			GPIO.output(LCD_D6, True)
		if bits&0x80==0x80:
			GPIO.output(LCD_D7, True)

		# Toggle 'Enable' pin
		time.sleep(E_DELAY)
		GPIO.output(LCD_EN, True)
		time.sleep(E_PULSE)
		GPIO.output(LCD_EN, False)
		time.sleep(E_DELAY)

		# Low bits
		GPIO.output(LCD_D4, False)
		GPIO.output(LCD_D5, False)
		GPIO.output(LCD_D6, False)
		GPIO.output(LCD_D7, False)
		if bits&0x01==0x01:
			GPIO.output(LCD_D4, True)
		if bits&0x02==0x02:
			GPIO.output(LCD_D5, True)
		if bits&0x04==0x04:
			GPIO.output(LCD_D6, True)
		if bits&0x08==0x08:
			GPIO.output(LCD_D7, True)

		# Toggle 'Enable' pin
		time.sleep(E_DELAY)
		GPIO.output(LCD_EN, True)
		time.sleep(E_PULSE)
		GPIO.output(LCD_EN, False)
		time.sleep(E_DELAY)

# Class for Image Processing
class Image:

	def __init__(self, _camera, _interval, filt, rlowHSV, rhighHSV, glowHSV, ghighHSV):

		# store
		self.cameraIndex = _camera
		self.frameInterval = _interval
		self.filt = filt
		self.capture = cv.VideoCapture(self.cameraIndex)

		#self.fwindow = cv.namedWindow("Result")
		self.redlowHSV = rlowHSV
		self.redhighHSV = rhighHSV

		self.greenlowHSV = glowHSV
		self.greenhighHSV = ghighHSV

		self.greenlastX = 0
		self.greenlastY = 0
		self.redlastX 	= 0
		self.redlastY 	= 0

	def doCaptureAndProcess(self):
		_, frame = self.capture.read()

		if self.filt == 0:
			filteredImage = frame
		elif self.filt == 1:
			filteredImage = cv.blur(frame, (5,5))
		elif self.filt == 2:
			filteredImage = cv.GaussianBlur(frame, (5,5),0)
		elif self.filt == 3:
			filteredImage = cv.medianBlur(frame, 5)
		elif self.filt == 4:
			filteredImage = cv.bilateralFilter(frame,9,75,75)

		kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE,(5,5))

		hsvImage = cv.cvtColor(filteredImage, cv.COLOR_BGR2HSV)

		RedimgThresholded = cv.inRange(hsvImage, self.redlowHSV, self.redhighHSV)
		GreenimgThresholded = cv.inRange(hsvImage, self.greenlowHSV, self.greenhighHSV)
		erosion = cv.erode(RedimgThresholded,kernel,iterations = 1)
		dilation = cv.dilate(RedimgThresholded,kernel,iterations = 1)

		dilation = cv.dilate(RedimgThresholded,kernel,iterations = 1)
		erosion = cv.erode(RedimgThresholded,kernel,iterations = 1)

		erosion = cv.erode(GreenimgThresholded,kernel,iterations = 1)
		dilation = cv.dilate(GreenimgThresholded,kernel,iterations = 1)

		dilation = cv.dilate(GreenimgThresholded,kernel,iterations = 1)
		erosion = cv.erode(GreenimgThresholded,kernel,iterations = 1)

		Redmoments = cv.moments(RedimgThresholded)
		Greenmoments = cv.moments(GreenimgThresholded)

		Redm01  = Redmoments['m01']
		Redm10  = Redmoments['m10']
		Redarea = Redmoments['m00']

		Greenm01  = Greenmoments['m01']
		Greenm10  = Greenmoments['m10']
		Greenarea = Greenmoments['m00']

		# print "RED " + str(Redarea)
		# print "GREEN " + str(Greenarea)

		if Redarea > redThresholdArea:
			#print "Red"
			posX = Redm10 / Redarea
			posY = Redm01 / Redarea

			# print "X " + str(posX)
			# print "Y " + str(posY)

			if self.redlastX >= 0 and self.redlastY >= 0 and posX >= 0  and posY >= 0:
				cv.line(filteredImage, (int(posX), int(posY)), (int(self.redlastX),int(self.redlastY)), (0,0,255), 5)

			redlastX = posX
			redlastY = posY
			if Redarea > Greenarea:
				print "RED"
				if gpioEnable == True:
					GPIO.output(LED_GREEN, False)
					GPIO.output(LED_RED, True)
				if lcdEnable == True:
					lcd.lcd_byte(0xC0, False)
					lcd.lcd_string("GREEN   ")
				# Turn Servo Left

		if Greenarea > greenThresholdArea:
			#print "Green"
			posX = Greenm10 / Greenarea
			posY = Greenm01 / Greenarea

			# print "X " + str(posX)
			# print "Y " + str(posY)

			if self.greenlastX >= 0 and self.greenlastY >= 0 and posX >= 0  and posY >= 0:
				cv.line(filteredImage, (int(posX), int(posY)), (int(self.greenlastX),int(self.greenlastY)), (0,0,255), 5)

			greenlastX = posX
			greenlastY = posY

			if Greenarea > Redarea:
				print "GREEN"
				if gpioEnable == True:
					GPIO.output(LED_GREEN, True)
					GPIO.output(LED_RED, False)
				if lcdEnable == True:
					lcd.lcd_byte(0xC0, False)
					lcd.lcd_string("GREEN   ")
				# Turn Servo

		if lcdEnable == True:
			lcd.lcd_byte(0x94, False)
			lcd.lcd_string("RED: " + str(Redarea / 1000000))
			lcd.lcd_byte(0xD4, False)
			lcd.lcd_string("GREEN: " + str(Greenarea / 1000000))
		cv.imshow("Red Thresholded Image", RedimgThresholded)
		cv.imshow("Green Thresholded Image", GreenimgThresholded)
		cv.imshow("Original Image", filteredImage)

redlowHSV = np.array([170,150,60], dtype=np.uint8)
redhighHSV = np.array([179,255,255], dtype=np.uint8)

greenlowHSV = np.array([50,150,60], dtype=np.uint8)
greenhighHSV = np.array([75,255,255], dtype=np.uint8)

obj = Image(0, 1, 1, redlowHSV, redhighHSV, greenlowHSV, greenhighHSV)
if gpioEnable == True:
	gpioInit()

# If LCD is enabled enable lcd for display
if lcdEnable == True:
	lcd = LCD(LCD_RS, LCD_E, LCD_D4, LCD_D5,LCD_D6,LCD_D7)
	lcd.lcd_init()
	lcd.lcd_byte(0x80, False)
	lcd.lcd_string("  Color Detection  ")

while True:

	obj.doCaptureAndProcess()

	# Break Operation if key interrupt
	k = cv.waitKey(5) & 0xFF
	if k == 27:
		break<|endoftext|>