<commit_before># audit_payload.py
#
# Writes logs of jobs starting and stopping based on ClassAd updates.
#
# This works automatically in a condor-ce if worker nodes have the environment
#   variable CONDORCE_COLLECTOR_HOST set to point to their site's condor-ce
#   server, if the jobs are run with condor glideins such as with GlideinWMS.
#
# Otherwise a job should use
#     condor_advertise -pool $CONDORCE_COLLECTOR_HOST:9619 UPDATE_STARTD_AD
#   at the start of the job and
#     condor_advertise -pool $CONDORCE_COLLECTOR_HOST:9619 INVALIDATE_STARTD_ADS
#   at the end of the job, sending to both comannds' standard input at least
#   these variables with format var = value, string values double-quoted:
#     Name (string) - name identifying the worker node, typically in the
#	format user@fully.qualified.domain.name
#     SlotID (integer) - slot identifier number on the worker node
#     MyType (string) - required to be set to the value of "Machine"
#   The condor_advertise at the begining of the message must also contain
#     GlobalJobId (string) - a globally unique identifier of the job
#     RemoteOwner (string) - a string identifying the owner of the job
#   and if the beginning of the message contains any of these they will
#   also be logged:
#     ClientMachine (string)
#     ProjectName (string)
#     Group (string)
#     x509UserProxyVOName (string)
#     x509userproxysubject (string)
#     x509UserProxyFQAN (string)
#     x509UserProxyEmail (string)
#
# There is one condor-ce configuration variable AUDIT_PAYLOAD_MAX_HOURS,
#   which is optional and indicates the maximum number of hours any job
#   is expected to run, default 72 hours (3 days).  After that time the
#   jobs will stop being tracked, in case a stop message was missed.
#
# Written by Dave Dykstra, June 2017
#

import htcondor
import time
from collections import OrderedDict

runningjobs = OrderedDict()

if 'AUDIT_PAYLOAD_MAX_HOURS' in htcondor.param:
    maxjobhours = int(htcondor.param['AUDIT_PAYLOAD_MAX_HOURS'])
else:
    maxjobhours = 3 * 24
htcondor.log(htcondor.LogLevel.Audit,
    "Audit payload maximum job hours: %d" % maxjobhours)
maxjobsecs = maxjobhours * 60 * 60

# a job may be being stopped
def stopjob(info):
    global runningjobs
    if 'Name' not in info or 'SlotID' not in info:
	return
    idx = (info['Name'], info['SlotID'])
    if idx not in runningjobs:
	return
    loginfo = {}
    loginfo['Name'] = info['Name']
    loginfo['SlotID'] = info['SlotID']
    loginfo['GlobalJobId'] = runningjobs[idx]['globaljobid']
    htcondor.log(htcondor.LogLevel.Audit, "Job stop: %s" % loginfo)

    del runningjobs[idx]

# a job may be being started
def startjob(info):
    global maxjobsecs
    global runningjobs

    if 'Name' not in info or 'SlotID' not in info or 'GlobalJobId' not in info:
	return

    idx = (info['Name'], info['SlotID'])
    globaljobid = info['GlobalJobId']
    if idx in runningjobs:
	if globaljobid == runningjobs[idx]['globaljobid']:
	    # just an update to a running job, ignore
	    return
	# first stop the existing job, the slot is being reused
	stopjob(info)

    htcondor.log(htcondor.LogLevel.Audit, "Job start: %s" % info)
    now = time.time()
    thisjob = {}
    thisjob['globaljobid'] = globaljobid
    thisjob['starttime'] = now
    runningjobs[idx] = thisjob

    # also look for expired jobs at the beginning of the list and stop them
    for idx in runningjobs:
	thisjob = runningjobs[idx]
	deltasecs = int(now - thisjob['starttime'])
	if deltasecs <= maxjobsecs:
	    break
	loginfo = {}
	loginfo['Name'] = idx[0]
	loginfo['SlotID'] = idx[1]
	loginfo['GlobalJobId'] = thisjob['globaljobid']
	htcondor.log(htcondor.LogLevel.Audit,
	    "Cleaning up %d-second expired job: %s" % (deltasecs, loginfo))
	del runningjobs[idx]

# this is the primary entry point called by the API
def update(command, ad):
    if command != "UPDATE_STARTD_AD":
        return
    if ad.get('State') == 'Unclaimed':
	stopjob(ad)  # stop previous job on this slot if any
	return
    info = {}
    keys = ['Name', 'SlotID', 'GlobalJobId',
	    'RemoteOwner', 'ClientMachine', 'ProjectName', 'Group',
	    'x509UserProxyVOName', 'x509userproxysubject', 'x509UserProxyEmail']
    for key in keys:
        if key in ad:
            info[key] = ad[key]
    startjob(info)

# this can also be called from the API when a job or slot is deleted
def invalidate(command, ad):
    if command != "INVALIDATE_STARTD_ADS":
        return
    stopjob(ad)
<commit_msg>add support for glideinWMS partitionable slots<commit_after># audit_payload.py
#
# Writes logs of jobs starting and stopping based on ClassAd updates.
#
# This works automatically in a condor-ce if worker nodes have the environment
#   variable CONDORCE_COLLECTOR_HOST set to point to their site's condor-ce
#   server, if the jobs are run with condor glideins such as with GlideinWMS.
#
# Otherwise a job should use
#     condor_advertise -pool $CONDORCE_COLLECTOR_HOST:9619 UPDATE_STARTD_AD
#   at the start of the job and
#     condor_advertise -pool $CONDORCE_COLLECTOR_HOST:9619 INVALIDATE_STARTD_ADS
#   at the end of the job, sending to both comannds' standard input at least
#   these variables with format var = value, string values double-quoted:
#     Name (string) - name identifying the worker node, typically in the
#       format user@fully.qualified.domain.name
#     SlotID (integer) - slot identifier number on the worker node
#     MyType (string) - required to be set to the value of "Machine"
#   The condor_advertise at the begining of the message must also contain
#     GlobalJobId (string) - a globally unique identifier of the job
#     RemoteOwner (string) - a string identifying the owner of the job
#   and if the beginning of the message contains any of these they will
#   also be logged:
#     ClientMachine (string)
#     ProjectName (string)
#     Group (string)
#     x509UserProxyVOName (string)
#     x509userproxysubject (string)
#     x509UserProxyFQAN (string)
#     x509UserProxyEmail (string)
#
# There is one condor-ce configuration variable AUDIT_PAYLOAD_MAX_HOURS,
#   which is optional and indicates the maximum number of hours any job
#   is expected to run, default 72 hours (3 days).  After that time the
#   jobs will stop being tracked, in case a stop message was missed.
#
# Written by Dave Dykstra, June 2017
#

import htcondor
import time
import re
from collections import OrderedDict

# Dictionary containing all tracked running jobs.
# Each entry is for a 'master', which is either a pilot job/glidein or
#  individual job.
# The index of the dictionary is a tuple of (mastername, slotid).
# The contents of each entry is a tuple of (starttime, jobs), where
#  jobs is a dictionary of individual job names running in that master
#  and each entry has a value of the GlobalJobId of that job.
runningmasters = OrderedDict()

if 'AUDIT_PAYLOAD_MAX_HOURS' in htcondor.param:
    maxjobhours = int(htcondor.param['AUDIT_PAYLOAD_MAX_HOURS'])
else:
    maxjobhours = 3 * 24
htcondor.log(htcondor.LogLevel.Audit,
    "Audit payload maximum job hours: %d" % maxjobhours)
maxjobsecs = maxjobhours * 60 * 60

# a job may be being stopped
def stopjob(info):
    global runningmasters
    if 'Name' not in info or 'SlotID' not in info:
        return
    name = info['Name']
    matchre = ""
    if 'GLIDEIN_MASTER_NAME' in info:
        idxname = info['GLIDEIN_MASTER_NAME']
        if idxname == name:
            # stop all jobs under this master
            matchre = '.*'
        else:
            # names of form "slotN@" stop that name and all "slotN_M@" names
            slotn = re.sub('^(slot[0-9]*)@.*', r'\1', name)
            if slotn != name:
                # match any name starting with slotN@ or slotN_
                matchre = '^' + slotn + '[@_]'
            # else take the default of matching only one name
    else:
        idxname = name
    idx = (idxname, info['SlotID'])
    if idx not in runningmasters:
        return

    runningjobs = runningmasters[idx][1]
    if matchre == "":
        # no match expression, just stop one
        if name not in runningjobs:
            return
        stopjobnames = [name]
    else:
        # select all jobs in this master
        stopjobnames = runningjobs.keys()
        if matchre != '.*':
            # restrict to the matching regular expression
            regex = re.compile(matchre)
            stopjobnames = filter(regex.search, stopjobnames)

    for stopjobname in stopjobnames:
        loginfo = {}
        loginfo['Name'] = stopjobname
        loginfo['SlotID'] = info['SlotID']
        loginfo['GlobalJobId'] = runningjobs[stopjobname]
        htcondor.log(htcondor.LogLevel.Audit, "Job stop: %s" % loginfo)
        del runningjobs[stopjobname]

    if len(runningjobs) == 0:
        del runningmasters[idx]

# a job may be being started
def startjob(info):
    global maxjobsecs
    global runningmasters

    if 'Name' not in info or 'SlotID' not in info or 'GlobalJobId' not in info:
        return

    name = info['Name']
    if 'GLIDEIN_MASTER_NAME' in info:
        # Glidein may be partitioned and sometimes tear down all contained
        #  slots at once, so need to track those slots together
        idxname = info['GLIDEIN_MASTER_NAME']
    else:
        idxname = name
    idx = (idxname, info['SlotID'])
    globaljobid = info['GlobalJobId']
    now = 0
    if idx in runningmasters:
        thismaster = runningmasters[idx]
        runningjobs = thismaster[1]
        if name in runningjobs:
            if globaljobid == runningjobs[name]:
                # just an update to a running job, ignore
                return
            # first stop the existing job, the slot is being reused
            stopjob(info)
    else:
        # new master
        now = time.time()
        thismaster = (now, {})
        runningmasters[idx] = thismaster
    # add job to this master
    thismaster[1][name] = globaljobid

    printinfo = {}
    keys = ['Name', 'SlotID', 'GlobalJobId',
            'RemoteOwner', 'ClientMachine', 'ProjectName', 'Group',
            'x509UserProxyVOName', 'x509userproxysubject', 'x509UserProxyEmail']
    for key in keys:
        if key in info:
            printinfo[key] = info[key]
    htcondor.log(htcondor.LogLevel.Audit, "Job start: %s" % printinfo)

    if now == 0:
        return

    # also look for expired jobs at the beginning of the list and stop them
    for idx in runningmasters:
        thismaster = runningmasters[idx]
        deltasecs = int(now - thismaster[0])
        if deltasecs <= maxjobsecs:
            break
        loginfo = {}
        loginfo['SlotID'] = idx[1]
        runningjobs = thismaster[1]
        for jobname in runningjobs:
            loginfo['Name'] = jobname
            loginfo['GlobalJobId'] = runningjobs[jobname]
            htcondor.log(htcondor.LogLevel.Audit,
                "Cleaning up %d-second expired job: %s" % (deltasecs, loginfo))
        del runningmasters[idx]

# this is the primary entry point called by the API
def update(command, ad):
    if command != "UPDATE_STARTD_AD":
        return
    if ad.get('State') == 'Unclaimed':
        stopjob(ad)  # stop previous job on this slot if any
        return
    startjob(ad)

# this can also be called from the API when a job or slot is deleted
def invalidate(command, ad):
    if command != "INVALIDATE_STARTD_ADS":
        return
    stopjob(ad)
<|endoftext|>


<commit_before>import asyncio as __asyncio

from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy

from . import includes as __includes  # NOQA
from . import _patch  # NOQA
from .loop import Loop as __BaseLoop  # NOQA

__version__ = '0.12.2.dev0'
__all__ = ('new_event_loop', 'install', 'EventLoopPolicy')

class Loop(__BaseLoop, __asyncio.AbstractEventLoop):
    pass

def new_event_loop():
    """Return a new event loop."""
    return Loop()

def install():
    """A helper function to install uvloop policy."""
    __asyncio.set_event_loop_policy(EventLoopPolicy())

class EventLoopPolicy(__BasePolicy):
    """Event loop policy.

    The preferred way to make your application use uvloop:

    >>> import asyncio
    >>> import uvloop
    >>> asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
    >>> asyncio.get_event_loop()
    <uvloop.Loop running=False closed=False debug=False>
    """

    def _loop_factory(self):
        return new_event_loop()
<commit_msg>Bump the version to 0.13.0.dev0<commit_after>import asyncio as __asyncio

from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy

from . import includes as __includes  # NOQA
from . import _patch  # NOQA
from .loop import Loop as __BaseLoop  # NOQA

__version__ = '0.13.0.dev0'
__all__ = ('new_event_loop', 'install', 'EventLoopPolicy')

class Loop(__BaseLoop, __asyncio.AbstractEventLoop):
    pass

def new_event_loop():
    """Return a new event loop."""
    return Loop()

def install():
    """A helper function to install uvloop policy."""
    __asyncio.set_event_loop_policy(EventLoopPolicy())

class EventLoopPolicy(__BasePolicy):
    """Event loop policy.

    The preferred way to make your application use uvloop:

    >>> import asyncio
    >>> import uvloop
    >>> asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
    >>> asyncio.get_event_loop()
    <uvloop.Loop running=False closed=False debug=False>
    """

    def _loop_factory(self):
        return new_event_loop()
<|endoftext|>


<commit_before>import asyncio as __asyncio

from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy

from . import includes as __includes  # NOQA
from . import _patch  # NOQA
from .loop import Loop as __BaseLoop  # NOQA

__version__ = '0.13.0.dev0'
__all__ = ('new_event_loop', 'install', 'EventLoopPolicy')

class Loop(__BaseLoop, __asyncio.AbstractEventLoop):
    pass

def new_event_loop():
    """Return a new event loop."""
    return Loop()

def install():
    """A helper function to install uvloop policy."""
    __asyncio.set_event_loop_policy(EventLoopPolicy())

class EventLoopPolicy(__BasePolicy):
    """Event loop policy.

    The preferred way to make your application use uvloop:

    >>> import asyncio
    >>> import uvloop
    >>> asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
    >>> asyncio.get_event_loop()
    <uvloop.Loop running=False closed=False debug=False>
    """

    def _loop_factory(self):
        return new_event_loop()
<commit_msg>Bump the version to 0.14.0.dev0<commit_after>import asyncio as __asyncio

from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy

from . import includes as __includes  # NOQA
from . import _patch  # NOQA
from .loop import Loop as __BaseLoop  # NOQA

__version__ = '0.14.0.dev0'
__all__ = ('new_event_loop', 'install', 'EventLoopPolicy')

class Loop(__BaseLoop, __asyncio.AbstractEventLoop):
    pass

def new_event_loop():
    """Return a new event loop."""
    return Loop()

def install():
    """A helper function to install uvloop policy."""
    __asyncio.set_event_loop_policy(EventLoopPolicy())

class EventLoopPolicy(__BasePolicy):
    """Event loop policy.

    The preferred way to make your application use uvloop:

    >>> import asyncio
    >>> import uvloop
    >>> asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
    >>> asyncio.get_event_loop()
    <uvloop.Loop running=False closed=False debug=False>
    """

    def _loop_factory(self):
        return new_event_loop()
<|endoftext|>


<commit_before># --------------------------------------------------------------------------------------
# Copyright 2017, Benedikt J. Daurer, Filipe R.N.C. Maia, Max F. Hantke, Carl Nettelblad
# Hummingbird is distributed under the terms of the Simplified BSD License.
# -------------------------------------------------------------------------
"""Online backend for reading EUxfel events from zmq."""
from __future__ import print_function # Compatibility with python 2 and 3
import os
import logging
from backend.event_translator import EventTranslator
from backend.record import Record, add_record
import numpy
import datetime
from pytz import timezone
from . import ureg
from backend import Worker
import ipc
import zmq
import msgpack
import msgpack_numpy
msgpack_numpy.patch()
import time as timemodule

from hummingbird import parse_cmdline_args

_argparser = None
def add_cmdline_args():
    global _argparser
    from utils.cmdline_args import argparser
    _argparser = argparser
    #group = _argparser.add_argument_group('EUxfel', 'Options for the EUxfel event translator')
    #group.add_argument('--euxfel-socket', metavar='euxfel_socket', default='tcp://127.0.0.1:4500',
    #                    help="EuXFEL socket address",
    #                    type=str)
    # TODO
    #group.add_argument('--euxfel-number-of-frames', metavar='euxfel_number_of_frames', nargs='?',
    #                    help="number of frames to be processed",
    #                    type=int)

#pulsecount = 'header.pulseCount'
pulsecount = 30

class EUxfelTranslator(object):
    """Translate between EUxfel events and Hummingbird ones"""
    """Note: Karabo provides full trains. We extract pulses from those."""
    def __init__(self, state):

        # Hack for timing
        self.t0 = timemodule.time()
        self.ntrains = 0

        self.timestamps = None
        cmdline_args = _argparser.parse_args()
        self._source = state['euxfel/agipd']

        # Reading data over ZMQ using socket adress (this is blocking, so this backend only works with one source at a time)
        self._zmq_context = zmq.Context()
        self._zmq_request = self._zmq_context.socket(zmq.REQ)
        self._zmq_request.connect(self._source['socket'])

        # Counters
        self._num_read_ahead = 1
        self._pos = 0
        self._data = None
        self._asked_data = False
        self.library = 'EUxfel'

        # Define how to translate between euxfel keys and Hummingbird ones
        # TODO: pulseEnergies, photonEnergies, train meta data, ..., ...
        # AGIPD
        self._n2c = {}
        self._mainsource = self._source['source']
        # Using the AGIPD metadata as our master source of metadata
        self._n2c[self._mainsource] = ['photonPixelDetectors', 'eventID']

        # Calculate the inverse mapping
        self._c2n = {}
        for k, v in self._n2c.items():
            if type(v) is not list:
                v = [v]
            for v2 in v:
                self._c2n[v2] = self._c2n.get(v2, [])
                self._c2n[v2].append(k)

        # Define how to translate between EuXFEL sources and Hummingbird ones
        self._s2c = {}
        # AGIPD
        self._s2c[self._mainsource] = self._mainsource

    def check_asked_data(self):
        """"Call for new data if needed."""
        if self._asked_data:
            return

        if self._data is None or self._pos >= pulsecount - self._num_read_ahead:
            self._zmq_request.send(b'next')
            self._asked_data = True

    def next_event(self):
        """Grabs the next event and returns the translated version"""
        # Old comment from Onda
        # FM: When running with vetoeing we get data on cells 2,4,6...,28
        # corresponding to indices 4,8,...,56
        if self._data is None or self._pos == pulsecount:
            self.check_asked_data()
            msg = self._zmq_request.recv()
            self._data = msgpack.loads(msg)
            self._asked_data = False
            self._pos = 0
            self.ntrains += 1.
            #print(self._data.keys())
            #print(self._data[self._mainsource]['image.pulseId'].shape)
            #print("Train count: ", self.ntrains)
            #print("Trains per sec: ",self.ntrains / (timemodule.time()-self.t0))

        self.check_asked_data()
        result = EventTranslator((self._pos, self._data), self)

        # Update pulse position counter
        self._pos = self._pos + 1

        return result

    def event_keys(self, evt):
        """Returns the translated keys available"""
        native_keys = evt[1].keys()
        common_keys = set()
        for k in native_keys:
            for c in self._native_to_common(k):
                common_keys.add(c)
        # analysis is for values added later on
        return list(common_keys)+['analysis']

    def _native_to_common(self, key):
        """Translates a native key to a hummingbird one"""
        if(key in self._n2c):
            val = self._n2c[key]
            if type(val) is not list:
                val = [val]
            return val
        else:
            return []

    def event_native_keys(self, evt):
        """Returns the native keys available"""
        return evt[1].keys()

    def translate(self, evt, key):
        """Returns a dict of Records that match a given hummingbird key"""
        values = {}
        if(key in self._c2n):
            return self.translate_core(evt, key)
        elif(key == 'analysis'):
            return {}
        elif(key == 'stream'):
            return {}
        else:
            # check if the key matches any of the existing keys in the event
            event_keys = evt[1].keys()
            values = {}
            found = False

            if key in event_keys:
                obj = evt[1][key]
                for subkey in obj.keys():
                    add_record(values, 'native', '%s[%s]' % (self._s2c[key], subkey),
                               obj[subkey], ureg.ADU)
                return values
            else:
                print('%s not found in event' % (key))

    def translate_core(self, evt, key):
        """Returns a dict of Records that matchs a core Hummingbird key.

        Core keys include  all except: parameters, any psana create key,
        any native key."""
        values = {}
        #print(evt,evt[0],evt[1])
        for k in self._c2n[key]:
            if k in evt[1]:
                if key == 'eventID':
                    self._tr_event_id(values, evt[1][k], evt[0])
                elif key == 'photonPixelDetectors':
                    self._tr_photon_detector(values, evt[1][k], k, evt[0])
                else:
                    raise RuntimeError('%s not yet supported with key %s' % (k, key))

        return values

    def event_id(self, evt):
        """Returns an id which should be unique for each
        shot and increase monotonically"""
        return self.translate(evt, 'eventID')['Timestamp'].timestamp

    def event_id2(self, evt):
        """Returns the int64 pulse ID"""
        return self.translate(evt, 'eventID')['Timestamp'].timestamp2

    def _tr_photon_detector(self, values, obj, evt_key, pos):
        """Translates pixel detector into Humminbird ADU array"""
        add_record(values, 'photonPixelDetectors', self._s2c[evt_key],
                   obj['image.data'][pos], ureg.ADU)

    def _tr_event_id(self, values, obj, pos):
        """Translates euxfel event ID from some source into a hummingbird one"""
        src_timestamp = obj['metadata']['timestamp']
        #print(src_timestamp, pos)
        timestamp = src_timestamp['sec'] + src_timestamp['frac'] * 1e-18 + pos * 1e-2
        #print(timestamp)
        #print(timemodule.time()-timestamp)

        time = datetime.datetime.fromtimestamp(timestamp, tz=timezone('utc'))
        time = time.astimezone(tz=timezone('CET'))
        rec = Record('Timestamp', time, ureg.s)
        #rec.fiducials = obj.fiducials()
        rec.pulseCount = pulsecount #obj[pulsecount]
        rec.pulseNo = pos
        #rec.timestamp2 = obj['trailer.trainId']
        #rec.timestamp = obj['image.pulseId'][rec.pulseNo]
        rec.timestamp = timestamp
        values[rec.name] = rec
<commit_msg>Improved euxfel backend<commit_after># --------------------------------------------------------------------------------------
# Copyright 2017, Benedikt J. Daurer, Filipe R.N.C. Maia, Max F. Hantke, Carl Nettelblad
# Hummingbird is distributed under the terms of the Simplified BSD License.
# -------------------------------------------------------------------------
"""Online backend for reading EUxfel events from zmq."""
from __future__ import print_function # Compatibility with python 2 and 3
import os
import logging
from backend.event_translator import EventTranslator
from backend.record import Record, add_record
import numpy
import datetime
from pytz import timezone
from . import ureg
from backend import Worker
import ipc
import zmq
import msgpack
import msgpack_numpy
msgpack_numpy.patch()
import time as timemodule

from hummingbird import parse_cmdline_args

_argparser = None
def add_cmdline_args():
    global _argparser
    from utils.cmdline_args import argparser
    _argparser = argparser
    #group = _argparser.add_argument_group('EUxfel', 'Options for the EUxfel event translator')
    #group.add_argument('--euxfel-socket', metavar='euxfel_socket', default='tcp://127.0.0.1:4500',
    #                    help="EuXFEL socket address",
    #                    type=str)
    # TODO
    #group.add_argument('--euxfel-number-of-frames', metavar='euxfel_number_of_frames', nargs='?',
    #                    help="number of frames to be processed",
    #                    type=int)

class EUxfelTranslator(object):
    """Translate between EUxfel events and Hummingbird ones"""
    """Note: Karabo provides full trains. We extract pulses from those."""
    def __init__(self, state):

        # Hack for timing
        self.t0 = timemodule.time()
        self.ntrains = 0

        self.timestamps = None
        cmdline_args = _argparser.parse_args()
        self._source = state['euxfel/agipd']

        # Reading data over ZMQ using socket adress (this is blocking, so this backend only works with one source at a time)
        self._zmq_context = zmq.Context()
        self._zmq_request = self._zmq_context.socket(zmq.REQ)
        self._zmq_request.connect(self._source['socket'])

        # Counters
        self._num_read_ahead = 1
        self._pos = 0
        self._pulsecount = None
        self._data = None
        self._asked_data = False
        self.library = 'EUxfel'

        # Define how to translate between euxfel keys and Hummingbird ones
        # TODO: pulseEnergies, photonEnergies, train meta data, ..., ...
        # AGIPD
        self._n2c = {}
        self._mainsource = self._source['source']
        # Using the AGIPD metadata as our master source of metadata
        self._n2c[self._mainsource] = ['photonPixelDetectors', 'eventID']

        # Calculate the inverse mapping
        self._c2n = {}
        for k, v in self._n2c.items():
            if type(v) is not list:
                v = [v]
            for v2 in v:
                self._c2n[v2] = self._c2n.get(v2, [])
                self._c2n[v2].append(k)

        # Define how to translate between EuXFEL sources and Hummingbird ones
        self._s2c = {}
        # AGIPD
        self._s2c[self._mainsource] = self._mainsource

    def check_asked_data(self):
        """"Call for new data if needed."""
        if self._asked_data:
            return

        if self._data is None or (self._pos >= self._pulsecount - self._num_read_ahead):
            self._zmq_request.send(b'next')
            self._asked_data = True

    def next_event(self):
        """Grabs the next event and returns the translated version"""
        # Old comment from Onda
        # FM: When running with vetoeing we get data on cells 2,4,6...,28
        # corresponding to indices 4,8,...,56
        if self._data is None or self._pos == self._pulsecount:
            self.check_asked_data()
            msg = self._zmq_request.recv()
            self._data = msgpack.loads(msg)
            self._asked_data = False
            self._pos = 0
            self._pulsecount = len(self._data[self._mainsource]['image.pulseId'].squeeze())
            self.ntrains += 1.
            #print("Train count: ", self.ntrains)
            #print("Trains per sec: ",self.ntrains / (timemodule.time()-self.t0))

        self.check_asked_data()
        result = EventTranslator((self._pos, self._data), self)

        # Update pulse position counter
        self._pos = self._pos + 1

        return result

    def event_keys(self, evt):
        """Returns the translated keys available"""
        native_keys = evt[1].keys()
        common_keys = set()
        for k in native_keys:
            for c in self._native_to_common(k):
                common_keys.add(c)
        # analysis is for values added later on
        return list(common_keys)+['analysis']

    def _native_to_common(self, key):
        """Translates a native key to a hummingbird one"""
        if(key in self._n2c):
            val = self._n2c[key]
            if type(val) is not list:
                val = [val]
            return val
        else:
            return []

    def event_native_keys(self, evt):
        """Returns the native keys available"""
        return evt[1].keys()

    def translate(self, evt, key):
        """Returns a dict of Records that match a given hummingbird key"""
        values = {}
        if(key in self._c2n):
            return self.translate_core(evt, key)
        elif(key == 'analysis'):
            return {}
        elif(key == 'stream'):
            return {}
        else:
            # check if the key matches any of the existing keys in the event
            event_keys = evt[1].keys()
            values = {}
            found = False

            if key in event_keys:
                obj = evt[1][key]
                for subkey in obj.keys():
                    add_record(values, 'native', '%s[%s]' % (self._s2c[key], subkey),
                               obj[subkey], ureg.ADU)
                return values
            else:
                print('%s not found in event' % (key))

    def translate_core(self, evt, key):
        """Returns a dict of Records that matchs a core Hummingbird key.

        Core keys include  all except: parameters, any psana create key,
        any native key."""
        values = {}
        #print(evt,evt[0],evt[1])
        for k in self._c2n[key]:
            if k in evt[1]:
                if key == 'eventID':
                    self._tr_event_id(values, evt[1][k], evt[0])
                elif key == 'photonPixelDetectors':
                    self._tr_photon_detector(values, evt[1][k], k, evt[0])
                else:
                    raise RuntimeError('%s not yet supported with key %s' % (k, key))

        return values

    def event_id(self, evt):
        """Returns an id which should be unique for each
        shot and increase monotonically"""
        return self.translate(evt, 'eventID')['Timestamp'].timestamp

    def event_id2(self, evt):
        """Returns the int64 pulse ID"""
        return self.translate(evt, 'eventID')['Timestamp'].timestamp2

    def _tr_photon_detector(self, values, obj, evt_key, pos):
        """Translates pixel detector into Humminbird ADU array"""
        if self._source['format'] == 'combined':
            img = obj['image.data'][pos]
        elif self._source['format'] == 'panel':
            img = obj['image.data'][:,:,0,pos]
        add_record(values, 'photonPixelDetectors', self._s2c[evt_key], img, ureg.ADU)

    def _tr_event_id(self, values, obj, pos):
        """Translates euxfel event ID from some source into a hummingbird one"""
        src_timestamp = obj['metadata']['timestamp']
        #print(src_timestamp, pos)
        timestamp = src_timestamp['sec'] + src_timestamp['frac'] * 1e-18 + pos * 1e-2
        #print(timestamp)
        #print(timemodule.time()-timestamp)
        time = datetime.datetime.fromtimestamp(timestamp, tz=timezone('utc'))
        time = time.astimezone(tz=timezone('CET'))
        rec = Record('Timestamp', time, ureg.s)
        rec.pulseCount = self._pulsecount
        rec.pulseNo = pos
        rec.timestamp = timestamp
        values[rec.name] = rec
<|endoftext|>


<commit_before>"""
Analysis tools for Topographica, other than plotting tools.

Configures the interface to the featuremapper and holoviews projects
and sets the appropriate Topographica-specific hooks.
"""

import numpy as np

from holoviews.interface.collector import Reference
from holoviews.core.options import Compositor
from holoviews.ipython import IPTestCase
from holoviews.operation.rgb import colormap
from holoviews.operation import chain, toRGB, matrix_overlay
import imagen.colorspaces
from featuremapper.command import Collector, measure_response

import topo
from topo.analysis.featureresponses import FeatureResponses, FeatureCurves,\
    FeatureMaps, ReverseCorrelation, MeasureResponseCommand, pattern_response,\
    topo_metadata_fn, StorageHook, get_feature_preference
from topo.base.projection import Projection
from topo.base.sheet import Sheet
from topo.base.sheetview import CFView
from topo.misc.ipython import RunProgress
from topo.misc import color

from command import measure_cog

CoG_spec = "Matrix.X CoG * Matrix.Y CoG * Matrix.BlueChannel"
XYCoG = chain.instance(value='XYCoG', name='XYCoG',
                       operations = [matrix_overlay.instance(spec=CoG_spec), toRGB.instance()])
Compositor.register(Compositor("Matrix.X CoG * Matrix.Y CoG", XYCoG, 'XYCoG', 'display'))

class TopoIPTestCase(IPTestCase):

    def __init__(self, *args, **kwargs):
        super(TopoIPTestCase, self).__init__(*args, **kwargs)
        self.addTypeEqualityFunc(CFView,   self.compare_cfviews)

    def compare_cfviews(self, view1, view2, msg):
        self.compare_matrix(view1, view2, msg)

    def compare_animation(self, view1, view2, msg):
        self.compare_viewmap(view1, view2, msg)

class SimRef(Reference):
    """
    A SimRef instance is installed on Collector to allow Topographica
    model elements to be referenced for collection.

    This is important to allow pickling and unpickling of Collectors
    that work correctly with Topographica in different execution
    environments (e.g. nodes of a cluster) and across different models
    without directly pickling the components (e.g. Sheets and
    Projections) themselves.

    More information about references can be found in the docstring of
    the holoviews.collector.Reference.
    """
    @property
    def resolved_type(self):
        if self.array_ref:
            return np.ndarray
        elif isinstance(self.obj, tuple):
            return Projection
        else:
            return Sheet

    def __init__(self, obj=None, array_ref=None):

        if topo.sim.model is not None:
            print "DEPRECATION WARNING: use topo.submodel.specifications instead of SimRef."

        if [obj, array_ref] == [None,None]:
            raise Exception("Please specify an object, a path string or an array_ref.")

        self.array_ref = None
        if obj is None:
            self.obj = None
            self.array_ref = array_ref
        elif isinstance(obj, str):
            self.obj = tuple(obj.split('.')) if '.' in obj else obj
        elif isinstance(obj, Projection):
            self.obj = (obj.dest.name, obj.name)
        else:
            self.obj = obj.name

    def resolve(self):
        from topo import sim
        if isinstance(self.obj, tuple):
            (sheet, proj) = self.obj
            return sim[sheet].projections()[proj]
        elif self.obj:
            return sim[self.obj]
        else:
            return eval('topo.sim.'+self.array_ref)

    def __repr__(self):
        if isinstance(self.obj, tuple):
            return "SimRef(%r)" % '.'.join(el for el in self.obj)
        elif self.obj is None:
            return "SimRef(array_ref=%r)" % self.array_ref
        else:
            return "SimRef(%r)" % self.obj

    def __str__(self):
        if isinstance(self.obj, tuple):
            return "topo.sim."+'.'.join(el for el in self.obj)
        elif self.obj is None:
            return "topo.sim." + self.array_ref
        else:
            return "topo.sim."+ self.obj

### Collection hooks

Collector.time_fn = topo.sim.time
Collector.interval_hook = RunProgress

def sheet_hook(obj, *args, **kwargs):
    """
    Return a Matrix of the Sheet activity.
    """
    return obj[:]

def projection_hook(obj, *args, **kwargs):
    """
    Return a Matrix of the projection activity, otherwise if
    grid=True, return a Grid of the CFs.
    """
    if kwargs.pop('grid', False):
        return obj.grid(**kwargs)
    else:
        return obj.projection_view()

def measurement_hook(obj, *args, **kwargs):
    return obj(*args, **kwargs)

# Configure Collector with appropriate hooks
Collector.sim = SimRef
Collector.for_type(Sheet, sheet_hook, referencer=SimRef)
Collector.for_type(Projection, projection_hook, referencer=SimRef)
Collector.for_type(measure_cog,  measurement_hook, mode='merge')

# Setting default channel operation for ON-OFF visualization
#Compositor.definitions.append(Compositor('Matrix.CFs * Matrix.CFs',
#                                   chain.instance(chain[colormap, operator.instance(operator=np.subtract)]), 'OnOff CFs', cmap='jet'))

# Featuremapper hooks

def empty_storage_hook(arg):
    """Use this to unset storage hook because lambda will not work
    with snapshots.

    This function is used in notebook_setup.py of the topographica
    IPython profile.
    """
    pass

FeatureResponses.metadata_fns = [topo_metadata_fn]
FeatureResponses.pattern_response_fn = pattern_response.instance()
FeatureMaps.measurement_storage_hook = StorageHook.instance(sublabel='Maps')
FeatureCurves.measurement_storage_hook = StorageHook.instance(sublabel='Curves')
ReverseCorrelation.measurement_storage_hook = StorageHook.instance(sublabel='RFs')
measure_response.measurement_storage_hook = StorageHook.instance(sublabel=None)
measure_cog.measurement_storage_hook = StorageHook.instance(sublabel='CoG')

MeasureResponseCommand.preference_lookup_fn = get_feature_preference
MeasureResponseCommand.pattern_response_fn = pattern_response.instance()

## Set optimized versions of color conversion functions
imagen.colorspaces.rgb_to_hsv = color._rgb_to_hsv_array_opt
imagen.colorspaces.hsv_to_rgb = color._hsv_to_rgb_array_opt

# Automatically discover all .py files in this directory.
import os,fnmatch
__all__ = [f.split('.py')[0] for f in os.listdir(__path__[0]) if fnmatch.fnmatch(f,'[!._]*.py')]
del f,os,fnmatch

<commit_msg>Updated CFView comparison<commit_after>"""
Analysis tools for Topographica, other than plotting tools.

Configures the interface to the featuremapper and holoviews projects
and sets the appropriate Topographica-specific hooks.
"""

import numpy as np

from holoviews.interface.collector import Reference
from holoviews.core.options import Compositor
from holoviews.ipython import IPTestCase
from holoviews.operation.rgb import colormap
from holoviews.operation import chain, toRGB, matrix_overlay
import imagen.colorspaces
from featuremapper.command import Collector, measure_response

import topo
from topo.analysis.featureresponses import FeatureResponses, FeatureCurves,\
    FeatureMaps, ReverseCorrelation, MeasureResponseCommand, pattern_response,\
    topo_metadata_fn, StorageHook, get_feature_preference
from topo.base.projection import Projection
from topo.base.sheet import Sheet
from topo.base.sheetview import CFView
from topo.misc.ipython import RunProgress
from topo.misc import color

from command import measure_cog

CoG_spec = "Matrix.X CoG * Matrix.Y CoG * Matrix.BlueChannel"
XYCoG = chain.instance(value='XYCoG', name='XYCoG',
                       operations = [matrix_overlay.instance(spec=CoG_spec), toRGB.instance()])
Compositor.register(Compositor("Matrix.X CoG * Matrix.Y CoG", XYCoG, 'XYCoG', 'display'))

class TopoIPTestCase(IPTestCase):

    def __init__(self, *args, **kwargs):
        super(TopoIPTestCase, self).__init__(*args, **kwargs)

    @classmethod
    def register(cls):
        super(TopoIPTestCase, cls).register()
        cls.equality_type_funcs[CFView] = cls.compare_cfview
        return cls.equality_type_funcs

    @classmethod
    def compare_cfview(cls, el1, el2, msg='CFView data'):
        cls.compare_matrix(el1, el2, msg=msg)

class SimRef(Reference):
    """
    A SimRef instance is installed on Collector to allow Topographica
    model elements to be referenced for collection.

    This is important to allow pickling and unpickling of Collectors
    that work correctly with Topographica in different execution
    environments (e.g. nodes of a cluster) and across different models
    without directly pickling the components (e.g. Sheets and
    Projections) themselves.

    More information about references can be found in the docstring of
    the holoviews.collector.Reference.
    """
    @property
    def resolved_type(self):
        if self.array_ref:
            return np.ndarray
        elif isinstance(self.obj, tuple):
            return Projection
        else:
            return Sheet

    def __init__(self, obj=None, array_ref=None):

        if topo.sim.model is not None:
            print "DEPRECATION WARNING: use topo.submodel.specifications instead of SimRef."

        if [obj, array_ref] == [None,None]:
            raise Exception("Please specify an object, a path string or an array_ref.")

        self.array_ref = None
        if obj is None:
            self.obj = None
            self.array_ref = array_ref
        elif isinstance(obj, str):
            self.obj = tuple(obj.split('.')) if '.' in obj else obj
        elif isinstance(obj, Projection):
            self.obj = (obj.dest.name, obj.name)
        else:
            self.obj = obj.name

    def resolve(self):
        from topo import sim
        if isinstance(self.obj, tuple):
            (sheet, proj) = self.obj
            return sim[sheet].projections()[proj]
        elif self.obj:
            return sim[self.obj]
        else:
            return eval('topo.sim.'+self.array_ref)

    def __repr__(self):
        if isinstance(self.obj, tuple):
            return "SimRef(%r)" % '.'.join(el for el in self.obj)
        elif self.obj is None:
            return "SimRef(array_ref=%r)" % self.array_ref
        else:
            return "SimRef(%r)" % self.obj

    def __str__(self):
        if isinstance(self.obj, tuple):
            return "topo.sim."+'.'.join(el for el in self.obj)
        elif self.obj is None:
            return "topo.sim." + self.array_ref
        else:
            return "topo.sim."+ self.obj

### Collection hooks

Collector.time_fn = topo.sim.time
Collector.interval_hook = RunProgress

def sheet_hook(obj, *args, **kwargs):
    """
    Return a Matrix of the Sheet activity.
    """
    return obj[:]

def projection_hook(obj, *args, **kwargs):
    """
    Return a Matrix of the projection activity, otherwise if
    grid=True, return a Grid of the CFs.
    """
    if kwargs.pop('grid', False):
        return obj.grid(**kwargs)
    else:
        return obj.projection_view()

def measurement_hook(obj, *args, **kwargs):
    return obj(*args, **kwargs)

# Configure Collector with appropriate hooks
Collector.sim = SimRef
Collector.for_type(Sheet, sheet_hook, referencer=SimRef)
Collector.for_type(Projection, projection_hook, referencer=SimRef)
Collector.for_type(measure_cog,  measurement_hook, mode='merge')

# Setting default channel operation for ON-OFF visualization
#Compositor.definitions.append(Compositor('Matrix.CFs * Matrix.CFs',
#                                   chain.instance(chain[colormap, operator.instance(operator=np.subtract)]), 'OnOff CFs', cmap='jet'))

# Featuremapper hooks

def empty_storage_hook(arg):
    """Use this to unset storage hook because lambda will not work
    with snapshots.

    This function is used in notebook_setup.py of the topographica
    IPython profile.
    """
    pass

FeatureResponses.metadata_fns = [topo_metadata_fn]
FeatureResponses.pattern_response_fn = pattern_response.instance()
FeatureMaps.measurement_storage_hook = StorageHook.instance(sublabel='Maps')
FeatureCurves.measurement_storage_hook = StorageHook.instance(sublabel='Curves')
ReverseCorrelation.measurement_storage_hook = StorageHook.instance(sublabel='RFs')
measure_response.measurement_storage_hook = StorageHook.instance(sublabel=None)
measure_cog.measurement_storage_hook = StorageHook.instance(sublabel='CoG')

MeasureResponseCommand.preference_lookup_fn = get_feature_preference
MeasureResponseCommand.pattern_response_fn = pattern_response.instance()

## Set optimized versions of color conversion functions
imagen.colorspaces.rgb_to_hsv = color._rgb_to_hsv_array_opt
imagen.colorspaces.hsv_to_rgb = color._hsv_to_rgb_array_opt

# Automatically discover all .py files in this directory.
import os,fnmatch
__all__ = [f.split('.py')[0] for f in os.listdir(__path__[0]) if fnmatch.fnmatch(f,'[!._]*.py')]
del f,os,fnmatch

<|endoftext|>


<commit_before>"""
Distribution class

$Id$
"""
__version__='$Revision$'
# To do:
#
# - wrap bins for cyclic histograms
# - check use of float() in count_mag() etc
# - clarify comment about negative selectivity
#
# - function to return value in a range (like a real histogram)
# - cache values
# - assumes cyclic axes start at 0: include a shift based on range
#
# - is there a way to make this work for arrays without mentioning
#   "array" anywhere in here?
# - should this be two classes: one for the core (which would be
#   small though) and another for statistics?

# The basic functions do not have any dependencies, but these imports
# are needed for some of the statistical functions (e.g. vector sum).
from math import pi
from Numeric import innerproduct, array, exp, argmax
from utils import arg

from utils import wrap

class Distribution(object):
    """
    Holds a distribution of the values f(x) associated with a variable x.

    A Distribution is a histogram-like object that is a dictionary of
    samples.  Each sample is an x:f(x) pair, where x is called the bin
    and f(x) is called the value(). Each bin's value is typically
    maintained as the sum of all the values that have been placed into
    it.

    The bin axis is continuous, and can represent a continuous
    quantity without discretization.  Alternatively, this class can be
    used as a traditional histogram by either discretizing the bin
    number before adding each sample, or by binning the values in the
    final Distribution.

    Distributions are bounded by the specified axis_bounds, and can
    either be cyclic (like directions or hues) or non-cyclic.  For
    cyclic distributions, samples provided outside the axis_bounds
    will be wrapped back into the bound range, as is appropriate for
    quantities like directions.  For non-cyclic distributions,
    providing samples outside the axis_bounds will result in a
    ValueError.

    In addition to the values, can also return the counts, i.e., the
    number of times that a sample has been added with the given bin.

    Not all instances of this class will be a true distribution in the
    mathematical sense; e.g. the values will have to be normalized
    before they can be considered a probability distribution.

    If keep_peak=True, the value stored in each bin will be the
    maximum of all values ever added, instead of the sum.  The
    distribution will thus be a record of the maximum value
    seen at each bin, also known as an envelope.
    """

    # Holds the number of times that undefined values have been
    # returned from calculations for any instance of this class,
    # e.g. calls to vector_direction() or vector_selectivity() when no
    # value is non-zero.  Useful for warning users when the values are
    # not meaningful.
    undefined_vals = 0

    def __init__(self, axis_bounds=(0.0, 2*pi), cyclic=False, keep_peak=False):
        self._data = {}
        self._counts = {}

        # total_count and total_value hold the total number and sum
        # (respectively) of values that have ever been provided for
        # each bin.  For a simple distribution these will be the same as
        # sum_counts() and sum_values().
        self.total_count = 0
        self.total_value = 0.0

        self.axis_bounds = axis_bounds
        self.axis_range = axis_bounds[1] - axis_bounds[0]
        self.cyclic = cyclic
        self.keep_peak = keep_peak

    ### JABHACKALERT!  The semantics of this operation are incorrect, because
    ### an expression like x+y should not modify x, while this does.  It could
    ### be renamed __iadd__, to implement += (which has the correct semantics),
    ### but it's not yet clear how to do that.
    def __add__(self,a):
        """
        Allows add() method to be used via the '+' operator; i.e.,
        Distribution + dictionary does Distribution.add(dictionary).
        """
        self.add(a)
        return None

    def get_value(self, bin):
        """
        Return the value of the specified bin.

        (Return None if there is no such bin.)
        """
        return self._data.get(bin)

    def get_count(self, bin):
        """
        Return the count from the specified bin.

        (Return None if there is no such bin.)
        """
        return self._counts.get(bin)

    def values(self):
        """
        Return a list of values.

        Various statistics can then be calculated if desired:

          sum(vals)  (total of all values)
          max(vals)  (highest value in any bin)

        Note that the bin-order of values returned does not necessarily
        match that returned by counts().
        """
        return self._data.values()

    def counts(self):
        """
        Return a list of values.

        Various statistics can then be calculated if desired:

          sum(counts)  (total of all counts)
          max(counts)  (highest count in any bin)

        Note that the bin-order of values returned does not necessarily
        match that returned by values().
        """
        return self._counts.values()

    def add(self, new_data):
        """
        Add a set of new data in the form of a dictionary of (bin,
        value) pairs.  If the bin already exists, the value is added
        to the current value.  If the bin doesn't exist, one is created
        with that value.

        Bin numbers outside axis_bounds are allowed for cyclic=True,
        but otherwise a ValueError is raised.

        If keep_peak=True, the value of the bin is the maximum of the
        current value and the supplied value.  That is, the bin stores
        the peak value seen so far.  Note that each call will increase
        the total_value and total_count (and thus decrease the
        value_mag() and count_mag()) even if the value doesn't happen
        to be the maximum seen so far, since each data point still
        helps improve the sampling and thus the confidence.
        """
        for bin in new_data.keys():

            if self.cyclic==False:
                if not (self.axis_bounds[0] <= bin <= self.axis_bounds[1]):
                    raise ValueError("Bin outside bounds.")
            # CEB (still to do):
            # new_bin = wrap(self.axis_bounds[0], self.axis_bounds[1], bin)
            new_bin = bin

            if new_bin not in self._data:
                self._data[new_bin] = 0.0
                self._counts[new_bin] = 0

            new_value = new_data[bin]
            self.total_value += new_value
            self._counts[new_bin] += 1
            self.total_count += 1

            if self.keep_peak == True:
                if new_value > self._data[new_bin]: self._data[new_bin] = new_value
            else:
                self._data[new_bin] += new_value

    def max_value_bin(self):
        """Return the bin with the largest value."""
        return self._data.keys()[argmax(self._data.values())]

    def weighted_average(self):
        """
        Return a continuous, interpolated equivalent of the max_value_bin().

        For a cyclic distribution, this is the direction of the vector
        sum (see vector_sum()).

        For a non-cyclic distribution, this is the arithmetic average
        of the data on the bin_axis, where each bin is weighted by its
        value.
        """
        if self.cyclic == True:
            return self.vector_sum()[1]
        else:
            return self._weighted_average()

    def vector_sum(self):
        """
        Return the vector sum of the data as a tuple (magnitude, avgbinnum).

        Each bin contributes a vector of length equal to its value, at
        a direction corresponding to the bin number.  Specifically,
        the total bin number range is mapped into a direction range
        [0,2pi].

        For a cyclic distribution, the avgbinnum will be a continuous
        measure analogous to the max_value_bin() of the distribution.
        But this quantity has more precision than max_value_bin()
        because it is computed from the entire distribution instead of
        just the peak bin.  However, it is likely to be useful only
        for uniform or very dense sampling; with sparse, non-uniform
        sampling the estimates will be biased significantly by the
        particular samples chosen.

        The avgbinnum is not meaningful when the magnitude is 0,
        because a zero-length vector has no direction.  To find out
        whether such cases occurred, you can compare the value of
        undefined_vals before and after a series of calls to this
        function.

        """
        # vectors are represented in polar form as complex numbers
        r = self._data.values()
        theta = self._bins_to_radians(array(self._data.keys()))
        v_sum = innerproduct(r, exp(theta*1j))

        magnitude = abs(v_sum)
        direction = arg(v_sum)

        if v_sum == 0:
            self.undefined_vals += 1

        direction_radians = self._radians_to_bins(direction)

        # wrap the direction because arctan2 returns principal values
        wrapped_direction = wrap(self.axis_bounds[0], self.axis_bounds[1], direction_radians)

        return (magnitude, wrapped_direction)

    def weighted_sum(self):
        """Return the sum of each value times its bin."""
        return innerproduct(self._data.keys(), self._data.values())

    def _weighted_average(self):
        """
        Return the weighted_sum divided by the sum of the values
        """
        return self._safe_divide(self.weighted_sum(),sum(self._data.values()))

    def selectivity(self):
        """
        Return a measure of the peakedness of the distribution.  The
        calculation differs depending on whether this is a cyclic
        variable.  For a cyclic variable, returns the magnitude of the
        vector_sum() divided by the sum_value() (see
        _vector_selectivity for more details).  For a non-cyclic
        variable, returns the max_value_bin()) as a proportion of the
        sum_value() (see _relative_selectivity for more details).
        """
        if self.cyclic == True:
            return self._vector_selectivity()
        else:
            return self._relative_selectivity()

    # not tested
    def _relative_selectivity(self):
        """
        Return max_value_bin()) as a proportion of the sum_value().

        This quantity is a measure of how strongly the distribution is
        biased towards the max_value_bin().  For a smooth,
        single-lobed distribution with an inclusive, non-cyclic range,
        this quantity is an analog to vector_selectivity.  To be a
        precise analog for arbitrary distributions, it would need to
        compute some measure of the selectivity that works like the
        weighted_average() instead of the max_value_bin().  The result
        is scaled such that if all bins are identical, the selectivity
        is 0.0, and if all bins but one are zero, the selectivity is
        1.0.
        """
        # A single bin is considered fully selective (but could also
        # arguably be considered fully unselective)
        if len(self._data) <= 1:
            return 1.0

        proportion = self._safe_divide( max(self._data.values()),
                                        sum(self._data.values()) )
        offset = 1.0/len(self._data)
        scaled = (proportion-offset)/(1.0-offset)

        # negative scaled is possible
        # e.g. 2 bins, with values that sum to less than 0.5
        # this probably isn't what should be done in those cases
        if scaled >= 0.0:
            return scaled
        else:
            return 0.0

    # not tested
    def _vector_selectivity(self):
        """
        Return the magnitude of the vector_sum() divided by the sum_value().

        This quantity is a vector-based measure of the peakedness of
        the distribution.  If only a single bin has a non-zero value(),
        the selectivity will be 1.0, and if all bins have the same
        value() then the selectivity will be 0.0.  Other distributions
        will result in intermediate values.

        For a distribution with a sum_value() of zero (i.e. all bins
        empty), the selectivity is undefined.  Assuming that one will
        usually be looking for high selectivity, we return zero in such
        a case so that high selectivity will not mistakenly be claimed.
        To find out whether such cases occurred, you can compare the
        value of undefined_values() before and after a series of
        calls to this function.
        """
        return self._safe_divide(self.vector_sum()[0], sum(self._data.values()))

    def value_mag(self, bin):
        """Return the value of a single bin as a proportion of total_value."""
        return self._safe_divide(self._data.get(bin), self.total_value)

    def count_mag(self, bin):
        """Return the count of a single bin as a proportion of total_count."""
        return self._safe_divide(float(self._counts.get(bin)), float(self.total_count))
        # use of float()

    def _bins_to_radians(self, bin):
        """
        Convert a bin number to a direction in radians.

        Works for Numeric arrays of bin numbers, returning
        an array of directions.
        """
        return (2*pi)*bin/self.axis_range

    def _radians_to_bins(self, direction):
        """
        Convert a direction in radians into a bin number.

        Works for Numeric arrays of direction, returning
        an array of bin numbers.
        """
        return direction*self.axis_range/(2*pi)

    def _safe_divide(self,numerator,denominator):
        """
        Division routine that avoids division-by-zero errors
        (returning zero in such cases) but keeps track of them
        for undefined_values().
        """
        if denominator==0:
            self.undefined_vals += 1
            return 0
        else:
            return numerator/denominator

##    # not tested - will it be needed?
##    def relative_value(self, bin):
##        """
##        Return the value of the given bin as a proportion of the
##        sum_value().
##
##        This quantity is on a scale from 0.0 (the
##        specified bin is zero and others are nonzero) to 1.0 (the
##        specified bin is the only nonzero bin).  If the distribution is
##        empty the result is undefined; in such a case zero is returned
##        and the value returned by undefined_values() is incremented.
##        """
##        return self._safe_divide(self._data.get(bin), sum(self.values()))

<commit_msg>Added reminder that we planned to refine the definition of selectivity for non-cyclic distributions.<commit_after>"""
Distribution class

$Id$
"""
__version__='$Revision$'
# To do:
#
# - wrap bins for cyclic histograms
# - check use of float() in count_mag() etc
# - clarify comment about negative selectivity
#
# - function to return value in a range (like a real histogram)
# - cache values
# - assumes cyclic axes start at 0: include a shift based on range
#
# - is there a way to make this work for arrays without mentioning
#   "array" anywhere in here?
# - should this be two classes: one for the core (which would be
#   small though) and another for statistics?

# The basic functions do not have any dependencies, but these imports
# are needed for some of the statistical functions (e.g. vector sum).
from math import pi
from Numeric import innerproduct, array, exp, argmax
from utils import arg

from utils import wrap

class Distribution(object):
    """
    Holds a distribution of the values f(x) associated with a variable x.

    A Distribution is a histogram-like object that is a dictionary of
    samples.  Each sample is an x:f(x) pair, where x is called the bin
    and f(x) is called the value(). Each bin's value is typically
    maintained as the sum of all the values that have been placed into
    it.

    The bin axis is continuous, and can represent a continuous
    quantity without discretization.  Alternatively, this class can be
    used as a traditional histogram by either discretizing the bin
    number before adding each sample, or by binning the values in the
    final Distribution.

    Distributions are bounded by the specified axis_bounds, and can
    either be cyclic (like directions or hues) or non-cyclic.  For
    cyclic distributions, samples provided outside the axis_bounds
    will be wrapped back into the bound range, as is appropriate for
    quantities like directions.  For non-cyclic distributions,
    providing samples outside the axis_bounds will result in a
    ValueError.

    In addition to the values, can also return the counts, i.e., the
    number of times that a sample has been added with the given bin.

    Not all instances of this class will be a true distribution in the
    mathematical sense; e.g. the values will have to be normalized
    before they can be considered a probability distribution.

    If keep_peak=True, the value stored in each bin will be the
    maximum of all values ever added, instead of the sum.  The
    distribution will thus be a record of the maximum value
    seen at each bin, also known as an envelope.
    """

    # Holds the number of times that undefined values have been
    # returned from calculations for any instance of this class,
    # e.g. calls to vector_direction() or vector_selectivity() when no
    # value is non-zero.  Useful for warning users when the values are
    # not meaningful.
    undefined_vals = 0

    def __init__(self, axis_bounds=(0.0, 2*pi), cyclic=False, keep_peak=False):
        self._data = {}
        self._counts = {}

        # total_count and total_value hold the total number and sum
        # (respectively) of values that have ever been provided for
        # each bin.  For a simple distribution these will be the same as
        # sum_counts() and sum_values().
        self.total_count = 0
        self.total_value = 0.0

        self.axis_bounds = axis_bounds
        self.axis_range = axis_bounds[1] - axis_bounds[0]
        self.cyclic = cyclic
        self.keep_peak = keep_peak

    ### JABHACKALERT!  The semantics of this operation are incorrect, because
    ### an expression like x+y should not modify x, while this does.  It could
    ### be renamed __iadd__, to implement += (which has the correct semantics),
    ### but it's not yet clear how to do that.
    def __add__(self,a):
        """
        Allows add() method to be used via the '+' operator; i.e.,
        Distribution + dictionary does Distribution.add(dictionary).
        """
        self.add(a)
        return None

    def get_value(self, bin):
        """
        Return the value of the specified bin.

        (Return None if there is no such bin.)
        """
        return self._data.get(bin)

    def get_count(self, bin):
        """
        Return the count from the specified bin.

        (Return None if there is no such bin.)
        """
        return self._counts.get(bin)

    def values(self):
        """
        Return a list of values.

        Various statistics can then be calculated if desired:

          sum(vals)  (total of all values)
          max(vals)  (highest value in any bin)

        Note that the bin-order of values returned does not necessarily
        match that returned by counts().
        """
        return self._data.values()

    def counts(self):
        """
        Return a list of values.

        Various statistics can then be calculated if desired:

          sum(counts)  (total of all counts)
          max(counts)  (highest count in any bin)

        Note that the bin-order of values returned does not necessarily
        match that returned by values().
        """
        return self._counts.values()

    def add(self, new_data):
        """
        Add a set of new data in the form of a dictionary of (bin,
        value) pairs.  If the bin already exists, the value is added
        to the current value.  If the bin doesn't exist, one is created
        with that value.

        Bin numbers outside axis_bounds are allowed for cyclic=True,
        but otherwise a ValueError is raised.

        If keep_peak=True, the value of the bin is the maximum of the
        current value and the supplied value.  That is, the bin stores
        the peak value seen so far.  Note that each call will increase
        the total_value and total_count (and thus decrease the
        value_mag() and count_mag()) even if the value doesn't happen
        to be the maximum seen so far, since each data point still
        helps improve the sampling and thus the confidence.
        """
        for bin in new_data.keys():

            if self.cyclic==False:
                if not (self.axis_bounds[0] <= bin <= self.axis_bounds[1]):
                    raise ValueError("Bin outside bounds.")
            # CEB (still to do):
            # new_bin = wrap(self.axis_bounds[0], self.axis_bounds[1], bin)
            new_bin = bin

            if new_bin not in self._data:
                self._data[new_bin] = 0.0
                self._counts[new_bin] = 0

            new_value = new_data[bin]
            self.total_value += new_value
            self._counts[new_bin] += 1
            self.total_count += 1

            if self.keep_peak == True:
                if new_value > self._data[new_bin]: self._data[new_bin] = new_value
            else:
                self._data[new_bin] += new_value

    def max_value_bin(self):
        """Return the bin with the largest value."""
        return self._data.keys()[argmax(self._data.values())]

    def weighted_average(self):
        """
        Return a continuous, interpolated equivalent of the max_value_bin().

        For a cyclic distribution, this is the direction of the vector
        sum (see vector_sum()).

        For a non-cyclic distribution, this is the arithmetic average
        of the data on the bin_axis, where each bin is weighted by its
        value.
        """
        if self.cyclic == True:
            return self.vector_sum()[1]
        else:
            return self._weighted_average()

    def vector_sum(self):
        """
        Return the vector sum of the data as a tuple (magnitude, avgbinnum).

        Each bin contributes a vector of length equal to its value, at
        a direction corresponding to the bin number.  Specifically,
        the total bin number range is mapped into a direction range
        [0,2pi].

        For a cyclic distribution, the avgbinnum will be a continuous
        measure analogous to the max_value_bin() of the distribution.
        But this quantity has more precision than max_value_bin()
        because it is computed from the entire distribution instead of
        just the peak bin.  However, it is likely to be useful only
        for uniform or very dense sampling; with sparse, non-uniform
        sampling the estimates will be biased significantly by the
        particular samples chosen.

        The avgbinnum is not meaningful when the magnitude is 0,
        because a zero-length vector has no direction.  To find out
        whether such cases occurred, you can compare the value of
        undefined_vals before and after a series of calls to this
        function.

        """
        # vectors are represented in polar form as complex numbers
        r = self._data.values()
        theta = self._bins_to_radians(array(self._data.keys()))
        v_sum = innerproduct(r, exp(theta*1j))

        magnitude = abs(v_sum)
        direction = arg(v_sum)

        if v_sum == 0:
            self.undefined_vals += 1

        direction_radians = self._radians_to_bins(direction)

        # wrap the direction because arctan2 returns principal values
        wrapped_direction = wrap(self.axis_bounds[0], self.axis_bounds[1], direction_radians)

        return (magnitude, wrapped_direction)

    def weighted_sum(self):
        """Return the sum of each value times its bin."""
        return innerproduct(self._data.keys(), self._data.values())

    def _weighted_average(self):
        """
        Return the weighted_sum divided by the sum of the values
        """
        return self._safe_divide(self.weighted_sum(),sum(self._data.values()))

    def selectivity(self):
        """
        Return a measure of the peakedness of the distribution.  The
        calculation differs depending on whether this is a cyclic
        variable.  For a cyclic variable, returns the magnitude of the
        vector_sum() divided by the sum_value() (see
        _vector_selectivity for more details).  For a non-cyclic
        variable, returns the max_value_bin()) as a proportion of the
        sum_value() (see _relative_selectivity for more details).
        """
        if self.cyclic == True:
            return self._vector_selectivity()
        else:
            return self._relative_selectivity()

    # CEBHACKALERT: the definition of selectivity for non-cyclic
    # quantities probably needs some more thought.
    # Additionally, this fails the test in testfeaturemap
    # (see the comment there).
    def _relative_selectivity(self):
        """
        Return max_value_bin()) as a proportion of the sum_value().

        This quantity is a measure of how strongly the distribution is
        biased towards the max_value_bin().  For a smooth,
        single-lobed distribution with an inclusive, non-cyclic range,
        this quantity is an analog to vector_selectivity.  To be a
        precise analog for arbitrary distributions, it would need to
        compute some measure of the selectivity that works like the
        weighted_average() instead of the max_value_bin().  The result
        is scaled such that if all bins are identical, the selectivity
        is 0.0, and if all bins but one are zero, the selectivity is
        1.0.
        """
        # A single bin is considered fully selective (but could also
        # arguably be considered fully unselective)
        if len(self._data) <= 1:
            return 1.0

        proportion = self._safe_divide( max(self._data.values()),
                                        sum(self._data.values()) )
        offset = 1.0/len(self._data)
        scaled = (proportion-offset)/(1.0-offset)

        # negative scaled is possible
        # e.g. 2 bins, with values that sum to less than 0.5
        # this probably isn't what should be done in those cases
        if scaled >= 0.0:
            return scaled
        else:
            return 0.0

    def _vector_selectivity(self):
        """
        Return the magnitude of the vector_sum() divided by the sum_value().

        This quantity is a vector-based measure of the peakedness of
        the distribution.  If only a single bin has a non-zero value(),
        the selectivity will be 1.0, and if all bins have the same
        value() then the selectivity will be 0.0.  Other distributions
        will result in intermediate values.

        For a distribution with a sum_value() of zero (i.e. all bins
        empty), the selectivity is undefined.  Assuming that one will
        usually be looking for high selectivity, we return zero in such
        a case so that high selectivity will not mistakenly be claimed.
        To find out whether such cases occurred, you can compare the
        value of undefined_values() before and after a series of
        calls to this function.
        """
        return self._safe_divide(self.vector_sum()[0], sum(self._data.values()))

    def value_mag(self, bin):
        """Return the value of a single bin as a proportion of total_value."""
        return self._safe_divide(self._data.get(bin), self.total_value)

    def count_mag(self, bin):
        """Return the count of a single bin as a proportion of total_count."""
        return self._safe_divide(float(self._counts.get(bin)), float(self.total_count))
        # use of float()

    def _bins_to_radians(self, bin):
        """
        Convert a bin number to a direction in radians.

        Works for Numeric arrays of bin numbers, returning
        an array of directions.
        """
        return (2*pi)*bin/self.axis_range

    def _radians_to_bins(self, direction):
        """
        Convert a direction in radians into a bin number.

        Works for Numeric arrays of direction, returning
        an array of bin numbers.
        """
        return direction*self.axis_range/(2*pi)

    def _safe_divide(self,numerator,denominator):
        """
        Division routine that avoids division-by-zero errors
        (returning zero in such cases) but keeps track of them
        for undefined_values().
        """
        if denominator==0:
            self.undefined_vals += 1
            return 0
        else:
            return numerator/denominator

##    # not tested - will it be needed?
##    def relative_value(self, bin):
##        """
##        Return the value of the given bin as a proportion of the
##        sum_value().
##
##        This quantity is on a scale from 0.0 (the
##        specified bin is zero and others are nonzero) to 1.0 (the
##        specified bin is the only nonzero bin).  If the distribution is
##        empty the result is undefined; in such a case zero is returned
##        and the value returned by undefined_values() is incremented.
##        """
##        return self._safe_divide(self._data.get(bin), sum(self.values()))

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (C) 2005 Edgewall Software
# Copyright (C) 2005 Christopher Lenz <cmlenz@gmx.de>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.org/wiki/TracLicense.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://trac.edgewall.org/log/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>

import os
import re
import weakref

from trac.core import *
from trac.db.api import IDatabaseConnector
from trac.db.util import ConnectionWrapper

_like_escape_re = re.compile(r'([/_%])')

try:
    import pysqlite2.dbapi2 as sqlite
    have_pysqlite = 2
    _ver = sqlite.sqlite_version_info
    sqlite_version = _ver[0] * 10000 + _ver[1] * 100 + int(_ver[2])

    class PyFormatCursor(sqlite.Cursor):
        def _rollback_on_error(self, function, *args, **kwargs):
            try:
                return function(self, *args, **kwargs)
            except sqlite.OperationalError, e:
                self.cnx.rollback()
                raise
        def execute(self, sql, args=None):
            if args:
                sql = sql % (('?',) * len(args))
            return self._rollback_on_error(sqlite.Cursor.execute, sql,
                                           args or [])
        def executemany(self, sql, args=None):
            if args:
                sql = sql % (('?',) * len(args[0]))
            return self._rollback_on_error(sqlite.Cursor.executemany, sql,
                                           args or [])

except ImportError:
    try:
        import sqlite
        have_pysqlite = 1
        _ver = sqlite._sqlite.sqlite_version_info()
        sqlite_version = _ver[0] * 10000 + _ver[1] * 100 + _ver[2]

        class SQLiteUnicodeCursor(sqlite.Cursor):
            def _convert_row(self, row):
                return tuple([(isinstance(v, str) and [v.decode('utf-8')] or [v])[0]
                              for v in row])
            def fetchone(self):
                row = sqlite.Cursor.fetchone(self)
                return row and self._convert_row(row) or None
            def fetchmany(self, num):
                rows = sqlite.Cursor.fetchmany(self, num)
                return rows != None and [self._convert_row(row)
                                         for row in rows] or []
            def fetchall(self):
                rows = sqlite.Cursor.fetchall(self)
                return rows != None and [self._convert_row(row)
                                         for row in rows] or []
    except ImportError:
        have_pysqlite = 0

def _to_sql(table):
    sql = ["CREATE TABLE %s (" % table.name]
    coldefs = []
    for column in table.columns:
        ctype = column.type.lower()
        if column.auto_increment:
            ctype = "integer PRIMARY KEY"
        elif len(table.key) == 1 and column.name in table.key:
            ctype += " PRIMARY KEY"
        elif ctype == "int":
            ctype = "integer"
        coldefs.append("    %s %s" % (column.name, ctype))
    if len(table.key) > 1:
        coldefs.append("    UNIQUE (%s)" % ','.join(table.key))
    sql.append(',\n'.join(coldefs) + '\n);')
    yield '\n'.join(sql)
    for index in table.indices:
        yield "CREATE INDEX %s_%s_idx ON %s (%s);" % (table.name,
              '_'.join(index.columns), table.name, ','.join(index.columns))

class SQLiteConnector(Component):
    """SQLite database support."""
    implements(IDatabaseConnector)

    def get_supported_schemes(self):
        return [('sqlite', 1)]

    def get_connection(self, path, params={}):
        return SQLiteConnection(path, params)

    def init_db(cls, path, params={}):
        if path != ':memory:':
            # make the directory to hold the database
            if os.path.exists(path):
                raise TracError, 'Database already exists at %s' % path
            os.makedirs(os.path.split(path)[0])
        cnx = sqlite.connect(path, timeout=int(params.get('timeout', 10000)))
        cursor = cnx.cursor()
        from trac.db_default import schema
        for table in schema:
            for stmt in cls.to_sql(table):
                cursor.execute(stmt)
        cnx.commit()

    def to_sql(cls, table):
        return _to_sql(table)

class SQLiteConnection(ConnectionWrapper):
    """Connection wrapper for SQLite."""

    __slots__ = ['_active_cursors']
    poolable = False

    def __init__(self, path, params={}):
        assert have_pysqlite > 0
        self.cnx = None
        if path != ':memory:':
            if not os.access(path, os.F_OK):
                raise TracError, 'Database "%s" not found.' % path

            dbdir = os.path.dirname(path)
            if not os.access(path, os.R_OK + os.W_OK) or \
                   not os.access(dbdir, os.R_OK + os.W_OK):
                from getpass import getuser
                raise TracError, 'The user %s requires read _and_ write ' \
                                 'permission to the database file %s and the ' \
                                 'directory it is located in.' \
                                 % (getuser(), path)

        if have_pysqlite == 2:
            self._active_cursors = weakref.WeakKeyDictionary()
            timeout = int(params.get('timeout', 10.0))
            cnx = sqlite.connect(path, detect_types=sqlite.PARSE_DECLTYPES,
                                 timeout=timeout)
        else:
            timeout = int(params.get('timeout', 10000))
            cnx = sqlite.connect(path, timeout=timeout, encoding='utf-8')

        ConnectionWrapper.__init__(self, cnx)

    if have_pysqlite == 2:
        def cursor(self):
            cursor = self.cnx.cursor(PyFormatCursor)
            self._active_cursors[cursor] = True
            cursor.cnx = self
            return cursor

        def rollback(self):
            for cursor in self._active_cursors.keys():
                cursor.close()
            self.cnx.rollback()

    else:
        def cursor(self):
            self.cnx._checkNotClosed("cursor")
            return SQLiteUnicodeCursor(self.cnx, self.cnx.rowclass)

    def cast(self, column, type):
        return 'CAST(%s AS %s)' % (column, type)

    def like(self):
        if sqlite_version >= 30100:
            return "LIKE %s ESCAPE '/'"
        else:
            return 'LIKE %s'

    def like_escape(self, text):
        if sqlite_version >= 30100:
            return _like_escape_re.sub(r'/\1', text)
        else:
            return text

    if have_pysqlite == 2:
        def get_last_id(self, cursor, table, column='id'):
            return cursor.lastrowid
    else:
        def get_last_id(self, cursor, table, column='id'):
            return self.cnx.db.sqlite_last_insert_rowid()
<commit_msg>simulate `CAST(column as int)` on older SQLite versions with `1*column` (fixes #3669)<commit_after># -*- coding: utf-8 -*-
#
# Copyright (C) 2005 Edgewall Software
# Copyright (C) 2005 Christopher Lenz <cmlenz@gmx.de>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.org/wiki/TracLicense.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://trac.edgewall.org/log/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>

import os
import re
import weakref

from trac.core import *
from trac.db.api import IDatabaseConnector
from trac.db.util import ConnectionWrapper

_like_escape_re = re.compile(r'([/_%])')

try:
    import pysqlite2.dbapi2 as sqlite
    have_pysqlite = 2
    _ver = sqlite.sqlite_version_info
    sqlite_version = _ver[0] * 10000 + _ver[1] * 100 + int(_ver[2])

    class PyFormatCursor(sqlite.Cursor):
        def _rollback_on_error(self, function, *args, **kwargs):
            try:
                return function(self, *args, **kwargs)
            except sqlite.OperationalError, e:
                self.cnx.rollback()
                raise
        def execute(self, sql, args=None):
            if args:
                sql = sql % (('?',) * len(args))
            return self._rollback_on_error(sqlite.Cursor.execute, sql,
                                           args or [])
        def executemany(self, sql, args=None):
            if args:
                sql = sql % (('?',) * len(args[0]))
            return self._rollback_on_error(sqlite.Cursor.executemany, sql,
                                           args or [])

except ImportError:
    try:
        import sqlite
        have_pysqlite = 1
        _ver = sqlite._sqlite.sqlite_version_info()
        sqlite_version = _ver[0] * 10000 + _ver[1] * 100 + _ver[2]

        class SQLiteUnicodeCursor(sqlite.Cursor):
            def _convert_row(self, row):
                return tuple([(isinstance(v, str) and [v.decode('utf-8')] or [v])[0]
                              for v in row])
            def fetchone(self):
                row = sqlite.Cursor.fetchone(self)
                return row and self._convert_row(row) or None
            def fetchmany(self, num):
                rows = sqlite.Cursor.fetchmany(self, num)
                return rows != None and [self._convert_row(row)
                                         for row in rows] or []
            def fetchall(self):
                rows = sqlite.Cursor.fetchall(self)
                return rows != None and [self._convert_row(row)
                                         for row in rows] or []
    except ImportError:
        have_pysqlite = 0

def _to_sql(table):
    sql = ["CREATE TABLE %s (" % table.name]
    coldefs = []
    for column in table.columns:
        ctype = column.type.lower()
        if column.auto_increment:
            ctype = "integer PRIMARY KEY"
        elif len(table.key) == 1 and column.name in table.key:
            ctype += " PRIMARY KEY"
        elif ctype == "int":
            ctype = "integer"
        coldefs.append("    %s %s" % (column.name, ctype))
    if len(table.key) > 1:
        coldefs.append("    UNIQUE (%s)" % ','.join(table.key))
    sql.append(',\n'.join(coldefs) + '\n);')
    yield '\n'.join(sql)
    for index in table.indices:
        yield "CREATE INDEX %s_%s_idx ON %s (%s);" % (table.name,
              '_'.join(index.columns), table.name, ','.join(index.columns))

class SQLiteConnector(Component):
    """SQLite database support."""
    implements(IDatabaseConnector)

    def get_supported_schemes(self):
        return [('sqlite', 1)]

    def get_connection(self, path, params={}):
        return SQLiteConnection(path, params)

    def init_db(cls, path, params={}):
        if path != ':memory:':
            # make the directory to hold the database
            if os.path.exists(path):
                raise TracError, 'Database already exists at %s' % path
            os.makedirs(os.path.split(path)[0])
        cnx = sqlite.connect(path, timeout=int(params.get('timeout', 10000)))
        cursor = cnx.cursor()
        from trac.db_default import schema
        for table in schema:
            for stmt in cls.to_sql(table):
                cursor.execute(stmt)
        cnx.commit()

    def to_sql(cls, table):
        return _to_sql(table)

class SQLiteConnection(ConnectionWrapper):
    """Connection wrapper for SQLite."""

    __slots__ = ['_active_cursors']
    poolable = False

    def __init__(self, path, params={}):
        assert have_pysqlite > 0
        self.cnx = None
        if path != ':memory:':
            if not os.access(path, os.F_OK):
                raise TracError, 'Database "%s" not found.' % path

            dbdir = os.path.dirname(path)
            if not os.access(path, os.R_OK + os.W_OK) or \
                   not os.access(dbdir, os.R_OK + os.W_OK):
                from getpass import getuser
                raise TracError, 'The user %s requires read _and_ write ' \
                                 'permission to the database file %s and the ' \
                                 'directory it is located in.' \
                                 % (getuser(), path)

        if have_pysqlite == 2:
            self._active_cursors = weakref.WeakKeyDictionary()
            timeout = int(params.get('timeout', 10.0))
            cnx = sqlite.connect(path, detect_types=sqlite.PARSE_DECLTYPES,
                                 timeout=timeout)
        else:
            timeout = int(params.get('timeout', 10000))
            cnx = sqlite.connect(path, timeout=timeout, encoding='utf-8')

        ConnectionWrapper.__init__(self, cnx)

    if have_pysqlite == 2:
        def cursor(self):
            cursor = self.cnx.cursor(PyFormatCursor)
            self._active_cursors[cursor] = True
            cursor.cnx = self
            return cursor

        def rollback(self):
            for cursor in self._active_cursors.keys():
                cursor.close()
            self.cnx.rollback()

    else:
        def cursor(self):
            self.cnx._checkNotClosed("cursor")
            return SQLiteUnicodeCursor(self.cnx, self.cnx.rowclass)

    def cast(self, column, type):
        if sqlite_version >= 30203:
            return 'CAST(%s AS %s)' % (column, type)
        elif type == 'int':
            # hack to force older SQLite versions to convert column to an int
            return '1*' + column
        else:
            return column

    def like(self):
        if sqlite_version >= 30100:
            return "LIKE %s ESCAPE '/'"
        else:
            return 'LIKE %s'

    def like_escape(self, text):
        if sqlite_version >= 30100:
            return _like_escape_re.sub(r'/\1', text)
        else:
            return text

    if have_pysqlite == 2:
        def get_last_id(self, cursor, table, column='id'):
            return cursor.lastrowid
    else:
        def get_last_id(self, cursor, table, column='id'):
            return self.cnx.db.sqlite_last_insert_rowid()
<|endoftext|>