<commit_before>'''
Created on Mar 28, 2013

@author: Ash Booth

'''

import sys
import math
from collections import deque
from cStringIO import StringIO

from ordertree import OrderTree

class OrderBook(object):
    def __init__(self):
        self.tape = deque(maxlen=None) # Index [0] is most recent trade
        self.bids = OrderTree()
        self.asks = OrderTree()
        self.lastTick = None
        self.lastTimestamp = 0
        self.tickSize = 0.0001
        self.time = 0
        self.nextQuoteID = 0

    def clipPrice(self, price):
        """ Clips the price according to the ticksize """
        return round(price, int(math.log10(1 / self.tickSize)))

    def updateTime(self):
        self.time+=1

    def processOrder(self, quote, fromData, verbose):
        orderType = quote['type']
        orderInBook = None
        self.updateTime()
        if quote['qty'] <= 0:
            sys.exit('processLimitOrder() given order of qty <= 0')
        if not fromData: self.nextQuoteID += 1
        quote['timestamp'] = self.time
        if orderType=='market':
            trades = self.processMarketOrder(quote, verbose)
        elif orderType=='limit':
            quote['price'] = self.clipPrice(quote['price'])
            trades, orderInBook = self.processLimitOrder(quote, fromData, verbose)
        else:
            sys.exit("processOrder() given neither 'market' nor 'limit'")
        return trades, orderInBook

    def processOrderList(self, side, orderlist,
                         qtyStillToTrade, quote, verbose):
        '''
        Takes an order list (stack of orders at one price) and
        an incoming order and matches appropriate trades given
        the orders quantity.
        '''
        trades = []
        qtyToTrade = qtyStillToTrade
        while len(orderlist) > 0 and qtyToTrade > 0:
            headOrder = orderlist.getHeadOrder()
            tradedPrice = headOrder.price
            counterparty = headOrder.tid
            if qtyToTrade < headOrder.qty:
                tradedQty = qtyToTrade
                # Amend book order
                newBookQty = headOrder.qty - qtyToTrade
                headOrder.updateQty(newBookQty, headOrder.timestamp)
                # Incoming done with
                qtyToTrade = 0
            elif qtyToTrade == headOrder.qty:
                tradedQty = qtyToTrade
                if side=='bid':
                    # Hit the bid
                    self.bids.removeOrderById(headOrder.idNum)
                else:
                    # Lift the ask
                    self.asks.removeOrderById(headOrder.idNum)
                # Incoming done with
                qtyToTrade = 0
            else:
                tradedQty = headOrder.qty
                if side=='bid':
                    # Hit the bid
                    self.bids.removeOrderById(headOrder.idNum)
                else:
                    # Lift the ask
                    self.asks.removeOrderById(headOrder.idNum)
                # We need to keep eating into volume at this price
                qtyToTrade -= tradedQty
            if verbose: print('>>> TRADE \nt=%d $%f n=%d p1=%d p2=%d' %
                              (self.time, tradedPrice, tradedQty,
                               counterparty, quote['tid']))

            transactionRecord = {'timestamp': self.time,
                                 'price': tradedPrice,
                                 'qty': tradedQty}
            if side == 'bid':
                transactionRecord['party1'] = [counterparty,
                                               'bid',
                                               headOrder.idNum]
                transactionRecord['party2'] = [quote['tid'],
                                               'ask',
                                               None]
            else:
                transactionRecord['party1'] = [counterparty,
                                               'ask',
                                               headOrder.idNum]
                transactionRecord['party2'] = [quote['tid'],
                                               'bid',
                                               None]
            self.tape.append(transactionRecord)
            trades.append(transactionRecord)
        return qtyToTrade, trades

    def processMarketOrder(self, quote, verbose):
        trades = []
        qtyToTrade = quote['qty']
        side = quote['side']
        if side == 'bid':
            while qtyToTrade > 0 and self.asks:
                bestPriceAsks = self.asks.minPriceList()
                qtyToTrade, newTrades = self.processOrderList('ask',
                                                                 bestPriceAsks,
                                                                 qtyToTrade,
                                                                 quote, verbose)
                trades += newTrades
        elif side == 'ask':
            while qtyToTrade > 0 and self.bids:
                bestPriceBids = self.bids.maxPriceList()
                qtyToTrade, newTrades = self.processOrderList('bid',
                                                                 bestPriceBids,
                                                                 qtyToTrade,
                                                                 quote, verbose)
                trades += newTrades
        else:
            sys.exit('processMarketOrder() received neither "bid" nor "ask"')
        return trades

    def processLimitOrder(self, quote, fromData, verbose):
        orderInBook = None
        trades = []
        qtyToTrade = quote['qty']
        side = quote['side']
        price = quote['price']
        if side == 'bid':
            while (self.asks and
                   price > self.asks.minPrice() and
                   qtyToTrade > 0):
                bestPriceAsks = self.asks.minPriceList()
                qtyToTrade, newTrades = self.processOrderList('ask',
                                                              bestPriceAsks,
                                                              qtyToTrade,
                                                              quote, verbose)
                trades += newTrades
            # If volume remains, add to book
            if qtyToTrade > 0:
                if not fromData:
                    quote['idNum'] = self.nextQuoteID
                quote['qty'] = qtyToTrade
                self.bids.insertOrder(quote)
                orderInBook = quote
        elif side == 'ask':
            while (self.bids and
                   price < self.bids.maxPrice() and
                   qtyToTrade > 0):
                bestPriceBids = self.bids.maxPriceList()
                qtyToTrade, newTrades = self.processOrderList('bid',
                                                              bestPriceBids,
                                                              qtyToTrade,
                                                              quote, verbose)
                trades += newTrades
            # If volume remains, add to book
            if qtyToTrade > 0:
                if not fromData:
                    quote['idNum'] = self.nextQuoteID
                quote['qty'] = qtyToTrade
                self.asks.insertOrder(quote)
                orderInBook = quote
        else:
            sys.exit('processLimitOrder() given neither bid nor ask')
        return trades, orderInBook

    def cancelOrder(self, side, idNum):
        self.updateTime()
        if side == 'bid':
            if self.bids.orderExists(idNum):
                self.bids.removeOrderById(idNum)
        elif side == 'ask':
            if self.asks.orderExists(idNum):
                self.asks.removeOrderById(idNum)
        else:
            sys.exit('cancelOrder() given neither bid nor ask')

    def modifyOrder(self, idNum, orderUpdate):
        self.updateTime()
        side = orderUpdate['side']
        orderUpdate['idNum'] = idNum
        orderUpdate['timestamp'] = self.time
        if side == 'bid':
            if self.bids.orderExists(orderUpdate['idNum']):
                self.bids.updateOrder(orderUpdate)
        elif side == 'ask':
            if self.asks.orderExists(orderUpdate['idNum']):
                self.asks.updateOrder(orderUpdate)
        else:
            sys.exit('modifyOrder() given neither bid nor ask')

    def getVolumeAtPrice(self, side, price):
        price = self.clipPrice(price)
        if side =='bid':
            vol = 0
            if self.bids.priceExists(price):
                vol = self.bids.getPrice(price).volume
            return vol
        elif side == 'ask':
            vol = 0
            if self.asks.priceExists(price):
                vol = self.asks.getPrice(price).volume
            return vol
        else:
            sys.exit('getVolumeAtPrice() given neither bid nor ask')

    def getBestBid(self):
        return self.bids.maxPrice()
    def getWorstBid(self):
        return self.bids.minPrice()
    def getBestAsk(self):
        return self.asks.minPrice()
    def getWorstAsk(self):
        return self.asks.maxPrice()

    def tapeDump(self, fname, fmode, tmode):
            dumpfile = open(fname, fmode)
            for tapeitem in self.tape:
                dumpfile.write('%s, %s, %s\n' % (tapeitem['time'],
                                                 tapeitem['price'],
                                                 tapeitem['qty']))
            dumpfile.close()
            if tmode == 'wipe':
                    self.tape = []

    def __str__(self):
        fileStr = StringIO()
        fileStr.write("------ Bids -------\n")
        if self.bids != None and len(self.bids) > 0:
            for k, v in self.bids.priceTree.items(reverse=True):
                fileStr.write('%s' % v)
        fileStr.write("\n------ Asks -------\n")
        if self.asks != None and len(self.asks) > 0:
            for k, v in self.asks.priceTree.items():
                fileStr.write('%s' % v)
        fileStr.write("\n------ Trades ------\n")
        if self.tape != None and len(self.tape) > 0:
            num = 0
            for entry in self.tape:
                if num < 5:
                    fileStr.write(str(entry['qty']) + " @ " +
                                  str(entry['price']) +
                                  " (" + str(entry['timestamp']) + ")\n")
                    num += 1
                else:
                    break
        fileStr.write("\n")
        return fileStr.getvalue()

<commit_msg>added ability to specify tick size<commit_after>'''
Created on Mar 28, 2013

@author: Ash Booth

'''

import sys
import math
from collections import deque
from cStringIO import StringIO

from ordertree import OrderTree

class OrderBook(object):
    def __init__(self, tick_size = 0.0001):
        self.tape = deque(maxlen=None) # Index [0] is most recent trade
        self.bids = OrderTree()
        self.asks = OrderTree()
        self.lastTick = None
        self.lastTimestamp = 0
        self.tickSize = tick_size
        self.time = 0
        self.nextQuoteID = 0

    def clipPrice(self, price):
        """ Clips the price according to the ticksize """
        return round(price, int(math.log10(1 / self.tickSize)))

    def updateTime(self):
        self.time+=1

    def processOrder(self, quote, fromData, verbose):
        orderType = quote['type']
        orderInBook = None
        self.updateTime()
        if quote['qty'] <= 0:
            sys.exit('processLimitOrder() given order of qty <= 0')
        if not fromData: self.nextQuoteID += 1
        quote['timestamp'] = self.time
        if orderType=='market':
            trades = self.processMarketOrder(quote, verbose)
        elif orderType=='limit':
            quote['price'] = self.clipPrice(quote['price'])
            trades, orderInBook = self.processLimitOrder(quote, fromData, verbose)
        else:
            sys.exit("processOrder() given neither 'market' nor 'limit'")
        return trades, orderInBook

    def processOrderList(self, side, orderlist,
                         qtyStillToTrade, quote, verbose):
        '''
        Takes an order list (stack of orders at one price) and
        an incoming order and matches appropriate trades given
        the orders quantity.
        '''
        trades = []
        qtyToTrade = qtyStillToTrade
        while len(orderlist) > 0 and qtyToTrade > 0:
            headOrder = orderlist.getHeadOrder()
            tradedPrice = headOrder.price
            counterparty = headOrder.tid
            if qtyToTrade < headOrder.qty:
                tradedQty = qtyToTrade
                # Amend book order
                newBookQty = headOrder.qty - qtyToTrade
                headOrder.updateQty(newBookQty, headOrder.timestamp)
                # Incoming done with
                qtyToTrade = 0
            elif qtyToTrade == headOrder.qty:
                tradedQty = qtyToTrade
                if side=='bid':
                    # Hit the bid
                    self.bids.removeOrderById(headOrder.idNum)
                else:
                    # Lift the ask
                    self.asks.removeOrderById(headOrder.idNum)
                # Incoming done with
                qtyToTrade = 0
            else:
                tradedQty = headOrder.qty
                if side=='bid':
                    # Hit the bid
                    self.bids.removeOrderById(headOrder.idNum)
                else:
                    # Lift the ask
                    self.asks.removeOrderById(headOrder.idNum)
                # We need to keep eating into volume at this price
                qtyToTrade -= tradedQty
            if verbose: print('>>> TRADE \nt=%d $%f n=%d p1=%d p2=%d' %
                              (self.time, tradedPrice, tradedQty,
                               counterparty, quote['tid']))

            transactionRecord = {'timestamp': self.time,
                                 'price': tradedPrice,
                                 'qty': tradedQty}
            if side == 'bid':
                transactionRecord['party1'] = [counterparty,
                                               'bid',
                                               headOrder.idNum]
                transactionRecord['party2'] = [quote['tid'],
                                               'ask',
                                               None]
            else:
                transactionRecord['party1'] = [counterparty,
                                               'ask',
                                               headOrder.idNum]
                transactionRecord['party2'] = [quote['tid'],
                                               'bid',
                                               None]
            self.tape.append(transactionRecord)
            trades.append(transactionRecord)
        return qtyToTrade, trades

    def processMarketOrder(self, quote, verbose):
        trades = []
        qtyToTrade = quote['qty']
        side = quote['side']
        if side == 'bid':
            while qtyToTrade > 0 and self.asks:
                bestPriceAsks = self.asks.minPriceList()
                qtyToTrade, newTrades = self.processOrderList('ask',
                                                                 bestPriceAsks,
                                                                 qtyToTrade,
                                                                 quote, verbose)
                trades += newTrades
        elif side == 'ask':
            while qtyToTrade > 0 and self.bids:
                bestPriceBids = self.bids.maxPriceList()
                qtyToTrade, newTrades = self.processOrderList('bid',
                                                                 bestPriceBids,
                                                                 qtyToTrade,
                                                                 quote, verbose)
                trades += newTrades
        else:
            sys.exit('processMarketOrder() received neither "bid" nor "ask"')
        return trades

    def processLimitOrder(self, quote, fromData, verbose):
        orderInBook = None
        trades = []
        qtyToTrade = quote['qty']
        side = quote['side']
        price = quote['price']
        if side == 'bid':
            while (self.asks and
                   price > self.asks.minPrice() and
                   qtyToTrade > 0):
                bestPriceAsks = self.asks.minPriceList()
                qtyToTrade, newTrades = self.processOrderList('ask',
                                                              bestPriceAsks,
                                                              qtyToTrade,
                                                              quote, verbose)
                trades += newTrades
            # If volume remains, add to book
            if qtyToTrade > 0:
                if not fromData:
                    quote['idNum'] = self.nextQuoteID
                quote['qty'] = qtyToTrade
                self.bids.insertOrder(quote)
                orderInBook = quote
        elif side == 'ask':
            while (self.bids and
                   price < self.bids.maxPrice() and
                   qtyToTrade > 0):
                bestPriceBids = self.bids.maxPriceList()
                qtyToTrade, newTrades = self.processOrderList('bid',
                                                              bestPriceBids,
                                                              qtyToTrade,
                                                              quote, verbose)
                trades += newTrades
            # If volume remains, add to book
            if qtyToTrade > 0:
                if not fromData:
                    quote['idNum'] = self.nextQuoteID
                quote['qty'] = qtyToTrade
                self.asks.insertOrder(quote)
                orderInBook = quote
        else:
            sys.exit('processLimitOrder() given neither bid nor ask')
        return trades, orderInBook

    def cancelOrder(self, side, idNum):
        self.updateTime()
        if side == 'bid':
            if self.bids.orderExists(idNum):
                self.bids.removeOrderById(idNum)
        elif side == 'ask':
            if self.asks.orderExists(idNum):
                self.asks.removeOrderById(idNum)
        else:
            sys.exit('cancelOrder() given neither bid nor ask')

    def modifyOrder(self, idNum, orderUpdate):
        self.updateTime()
        side = orderUpdate['side']
        orderUpdate['idNum'] = idNum
        orderUpdate['timestamp'] = self.time
        if side == 'bid':
            if self.bids.orderExists(orderUpdate['idNum']):
                self.bids.updateOrder(orderUpdate)
        elif side == 'ask':
            if self.asks.orderExists(orderUpdate['idNum']):
                self.asks.updateOrder(orderUpdate)
        else:
            sys.exit('modifyOrder() given neither bid nor ask')

    def getVolumeAtPrice(self, side, price):
        price = self.clipPrice(price)
        if side =='bid':
            vol = 0
            if self.bids.priceExists(price):
                vol = self.bids.getPrice(price).volume
            return vol
        elif side == 'ask':
            vol = 0
            if self.asks.priceExists(price):
                vol = self.asks.getPrice(price).volume
            return vol
        else:
            sys.exit('getVolumeAtPrice() given neither bid nor ask')

    def getBestBid(self):
        return self.bids.maxPrice()
    def getWorstBid(self):
        return self.bids.minPrice()
    def getBestAsk(self):
        return self.asks.minPrice()
    def getWorstAsk(self):
        return self.asks.maxPrice()

    def tapeDump(self, fname, fmode, tmode):
            dumpfile = open(fname, fmode)
            for tapeitem in self.tape:
                dumpfile.write('%s, %s, %s\n' % (tapeitem['time'],
                                                 tapeitem['price'],
                                                 tapeitem['qty']))
            dumpfile.close()
            if tmode == 'wipe':
                    self.tape = []

    def __str__(self):
        fileStr = StringIO()
        fileStr.write("------ Bids -------\n")
        if self.bids != None and len(self.bids) > 0:
            for k, v in self.bids.priceTree.items(reverse=True):
                fileStr.write('%s' % v)
        fileStr.write("\n------ Asks -------\n")
        if self.asks != None and len(self.asks) > 0:
            for k, v in self.asks.priceTree.items():
                fileStr.write('%s' % v)
        fileStr.write("\n------ Trades ------\n")
        if self.tape != None and len(self.tape) > 0:
            num = 0
            for entry in self.tape:
                if num < 5:
                    fileStr.write(str(entry['qty']) + " @ " +
                                  str(entry['price']) +
                                  " (" + str(entry['timestamp']) + ")\n")
                    num += 1
                else:
                    break
        fileStr.write("\n")
        return fileStr.getvalue()

<|endoftext|>


<commit_before>"""
Allows users to change the range of the current graph
"""
from PyQt5 import QtCore
from PyQt5 import QtGui
from PyQt5 import QtWidgets
import sas.qtgui.path_prepare

import matplotlib as mpl
import numpy

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from sas.qtgui.Plotting.PlotterData import Data2D
from sas.qtgui.Utilities.GuiUtils import formatNumber, DoubleValidator
from .rangeSlider import RangeSlider

DEFAULT_MAP = 'jet'

# Local UI
from sas.qtgui.UI import main_resources_rc
from sas.qtgui.Plotting.UI.ColorMapUI import Ui_ColorMapUI

class ColorMap(QtWidgets.QDialog, Ui_ColorMapUI):
    apply_signal = QtCore.pyqtSignal(tuple, str)
    def __init__(self, parent=None, cmap=None, xmin=-1,xmax=10, ymin=-1.0, ymax=15.0, vmin=0.0, vmax=100.0, data=None):
        super(ColorMap, self).__init__()

        self.setupUi(self)
        assert(isinstance(data, Data2D))
        # disable the context help icon
        self.setWindowFlags(self.windowFlags() & ~QtCore.Qt.WindowContextHelpButtonHint)

        self.data = data
        self._cmap_orig = self._cmap = cmap if cmap is not None else DEFAULT_MAP
        self.all_maps = [m for m in mpl.cm.datad]
        self.maps = sorted(m for m in self.all_maps if not m.endswith("_r"))
        self.rmaps = sorted(set(self.all_maps) - set(self.maps))

        self.vmin = self.vmin_orig = vmin
        self.vmax = self.vmax_orig = vmax

        self.data.xmin = xmin
        self.data.xmax = xmax

        self.data.ymin = ymin
        self.data.ymax = ymax

        # Initialize detector labels
        self.initDetectorData()

        # Initialize the combo box
        self.initMapCombobox()

        self.initRangeSlider()

        # Add the color map component
        self.initColorMap()

        # Initialize validators on amplitude textboxes
        validator_min = DoubleValidator(self.txtMinAmplitude)
        validator_min.setNotation(0)
        self.txtMinAmplitude.setValidator(validator_min)
        validator_max = DoubleValidator(self.txtMaxAmplitude)
        validator_max.setNotation(0)
        self.txtMaxAmplitude.setValidator(validator_max)

        # Set the initial amplitudes
        self.txtMinAmplitude.setText(formatNumber(self.vmin))
        self.txtMaxAmplitude.setText(formatNumber(self.vmax))

        # Enforce constant size on the widget
        self.setFixedSize(self.minimumSizeHint())

        # Handle combobox changes
        self.cbColorMap.currentIndexChanged.connect(self.onMapIndexChange)

        # Handle checkbox changes
        self.chkReverse.stateChanged.connect(self.onColorMapReversed)

        # Handle the Reset button click
        self.buttonBox.button(QtWidgets.QDialogButtonBox.Reset).clicked.connect(self.onReset)

        # Handle the Apply button click
        self.buttonBox.button(QtWidgets.QDialogButtonBox.Apply).clicked.connect(self.onApply)

        # Handle the amplitude setup
        self.txtMinAmplitude.editingFinished.connect(self.onAmplitudeChange)
        self.txtMaxAmplitude.editingFinished.connect(self.onAmplitudeChange)

    def cmap(self):
        """
        Getter for the color map
        """
        return self._cmap

    def norm(self):
        """
        Getter for the color map norm
        """
        return (self._norm.vmin, self._norm.vmax)

    def onReset(self):
        """
        Respond to the Reset button click
        """
        # Go back to original settings
        self._cmap = self._cmap_orig
        self.vmin = self.vmin_orig
        self.vmax = self.vmax_orig
        self._norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
        self.txtMinAmplitude.setText(formatNumber(self.vmin))
        self.txtMaxAmplitude.setText(formatNumber(self.vmax))
        self.initMapCombobox()
        self.slider.setMinimum(self.vmin)
        self.slider.setMaximum(self.vmax)
        self.slider.setLowValue(self.vmin)
        self.slider.setHighValue(self.vmax)
        # Redraw the widget
        self.redrawColorBar()
        self.canvas.draw()

    def onApply(self):
        """
        Respond to the Apply button click.
        Send a signal to the plotter with vmin/vmax/cmap for chart update
        """
        self.apply_signal.emit(self.norm(), self.cmap())

    def initDetectorData(self):
        """
        Fill out the Detector labels
        """
        xnpts = len(self.data.x_bins)
        ynpts = len(self.data.y_bins)
        self.lblWidth.setText(formatNumber(xnpts))
        self.lblHeight.setText(formatNumber(ynpts))
        xmax = max(self.data.xmin, self.data.xmax)
        ymax = max(self.data.ymin, self.data.ymax)
        qmax = numpy.sqrt(numpy.power(xmax, 2) + numpy.power(ymax, 2))
        self.lblQmax.setText(formatNumber(qmax))
        self.lblStopRadius.setText(formatNumber(self.data.xmin))

    def initMapCombobox(self):
        """
        Fill out the combo box with all available color maps
        """
        if self._cmap in self.rmaps:
            maps = self.rmaps
            # Assure correct state of the checkbox
            self.chkReverse.setChecked(True)
        else:
            maps = self.maps
            # Assure correct state of the checkbox
            self.chkReverse.setChecked(False)

        self.cbColorMap.addItems(maps)
        # Set the default/passed map
        self.cbColorMap.setCurrentIndex(self.cbColorMap.findText(self._cmap))

    def initRangeSlider(self):
        """
        Create and display the double slider for data range mapping.
        """
        self.slider = RangeSlider()
        self.slider.setMinimum(self.vmin)
        self.slider.setMaximum(self.vmax)
        self.slider.setLowValue(self.vmin)
        self.slider.setHighValue(self.vmax)
        self.slider.setOrientation(QtCore.Qt.Horizontal)

        self.slider_label = QtWidgets.QLabel()
        self.slider_label.setText("Drag the sliders to adjust color range.")

        def set_vmin(value):
            self.vmin = value
            self.txtMinAmplitude.setText(str(value))
            self.updateMap()
        def set_vmax(value):
            self.vmax = value
            self.txtMaxAmplitude.setText(str(value))
            self.updateMap()

        self.slider.lowValueChanged.connect(set_vmin)
        self.slider.highValueChanged.connect(set_vmax)

    def updateMap(self):
        self._norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
        self.redrawColorBar()
        self.canvas.draw()

    def initColorMap(self):
        """
        Prepare and display the color map
        """
        self.fig = mpl.figure.Figure(figsize=(4, 1))
        self.ax1 = self.fig.add_axes([0.05, 0.65, 0.9, 0.15])

        self._norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
        self.redrawColorBar()
        self.canvas = FigureCanvas(self.fig)

        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.slider_label)
        layout.addWidget(self.slider)
        layout.addWidget(self.canvas)

        self.widget.setLayout(layout)

    def onMapIndexChange(self, index):
        """
        Respond to the color map change event
        """
        new_map = str(self.cbColorMap.itemText(index))
        self._cmap = new_map
        self.redrawColorBar()
        self.canvas.draw()

    def redrawColorBar(self):
        """
        Call ColorbarBase with current values, effectively redrawing the widget
        """
        self.cb = mpl.colorbar.ColorbarBase(self.ax1, cmap=mpl.pyplot.get_cmap(self._cmap),
                                            norm=self._norm,
                                            orientation='horizontal')
        self.cb.set_label('Color map range')

    def onColorMapReversed(self, isChecked):
        """
        Respond to ticking/unticking the color map reverse checkbox
        """
        current_map = str(self.cbColorMap.currentText())
        if isChecked:
            # Add "_r" to map name for the reversed version
            new_map = current_map + "_r"
            maps = self.rmaps
            # Assure the reversed map exists.
            if new_map not in maps:
                new_map = maps[0]
        else:
            new_map = current_map[:-2] # "_r" = last two chars
            maps = self.maps
            # Base map for the reversed map should ALWAYS exist,
            # but let's be paranoid here
            if new_map not in maps:
                new_map = maps[0]

        self._cmap = new_map
        # Clear the content of the combobox.
        # Needs signal blocking, or else onMapIndexChange() spoils it all
        self.cbColorMap.blockSignals(True)
        self.cbColorMap.clear()
        # Add the new set of maps
        self.cbColorMap.addItems(maps)
        # Remove the signal block before the proper index set
        self.cbColorMap.blockSignals(False)
        self.cbColorMap.setCurrentIndex(self.cbColorMap.findText(new_map))

    def onAmplitudeChange(self):
        """
        Respond to editing the amplitude fields
        """
        min_amp = self.vmin
        max_amp = self.vmax

        try:
            min_amp = float(self.txtMinAmplitude.text())
        except ValueError:
            pass
        try:
            max_amp = float(self.txtMaxAmplitude.text())
        except ValueError:
            pass

        self._norm = mpl.colors.Normalize(vmin=min_amp, vmax=max_amp)
        self.redrawColorBar()
        self.canvas.draw()
<commit_msg>revert changes for testing code<commit_after>"""
Allows users to change the range of the current graph
"""
from PyQt5 import QtCore
from PyQt5 import QtGui
from PyQt5 import QtWidgets
import sas.qtgui.path_prepare

import matplotlib as mpl
import numpy

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from sas.qtgui.Plotting.PlotterData import Data2D
from sas.qtgui.Utilities.GuiUtils import formatNumber, DoubleValidator
from .rangeSlider import RangeSlider

DEFAULT_MAP = 'jet'

# Local UI
from sas.qtgui.UI import main_resources_rc
from sas.qtgui.Plotting.UI.ColorMapUI import Ui_ColorMapUI

class ColorMap(QtWidgets.QDialog, Ui_ColorMapUI):
    apply_signal = QtCore.pyqtSignal(tuple, str)
    def __init__(self, parent=None, cmap=None, vmin=0.0, vmax=100.0, data=None):
        super(ColorMap, self).__init__()

        self.setupUi(self)
        assert(isinstance(data, Data2D))
        # disable the context help icon
        self.setWindowFlags(self.windowFlags() & ~QtCore.Qt.WindowContextHelpButtonHint)

        self.data = data
        self._cmap_orig = self._cmap = cmap if cmap is not None else DEFAULT_MAP
        self.all_maps = [m for m in mpl.cm.datad]
        self.maps = sorted(m for m in self.all_maps if not m.endswith("_r"))
        self.rmaps = sorted(set(self.all_maps) - set(self.maps))

        self.vmin = self.vmin_orig = vmin
        self.vmax = self.vmax_orig = vmax

        # Initialize detector labels
        self.initDetectorData()

        # Initialize the combo box
        self.initMapCombobox()

        self.initRangeSlider()

        # Add the color map component
        self.initColorMap()

        # Initialize validators on amplitude textboxes
        validator_min = DoubleValidator(self.txtMinAmplitude)
        validator_min.setNotation(0)
        self.txtMinAmplitude.setValidator(validator_min)
        validator_max = DoubleValidator(self.txtMaxAmplitude)
        validator_max.setNotation(0)
        self.txtMaxAmplitude.setValidator(validator_max)

        # Set the initial amplitudes
        self.txtMinAmplitude.setText(formatNumber(self.vmin))
        self.txtMaxAmplitude.setText(formatNumber(self.vmax))

        # Enforce constant size on the widget
        self.setFixedSize(self.minimumSizeHint())

        # Handle combobox changes
        self.cbColorMap.currentIndexChanged.connect(self.onMapIndexChange)

        # Handle checkbox changes
        self.chkReverse.stateChanged.connect(self.onColorMapReversed)

        # Handle the Reset button click
        self.buttonBox.button(QtWidgets.QDialogButtonBox.Reset).clicked.connect(self.onReset)

        # Handle the Apply button click
        self.buttonBox.button(QtWidgets.QDialogButtonBox.Apply).clicked.connect(self.onApply)

        # Handle the amplitude setup
        self.txtMinAmplitude.editingFinished.connect(self.onAmplitudeChange)
        self.txtMaxAmplitude.editingFinished.connect(self.onAmplitudeChange)

    def cmap(self):
        """
        Getter for the color map
        """
        return self._cmap

    def norm(self):
        """
        Getter for the color map norm
        """
        return (self._norm.vmin, self._norm.vmax)

    def onReset(self):
        """
        Respond to the Reset button click
        """
        # Go back to original settings
        self._cmap = self._cmap_orig
        self.vmin = self.vmin_orig
        self.vmax = self.vmax_orig
        self._norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
        self.txtMinAmplitude.setText(formatNumber(self.vmin))
        self.txtMaxAmplitude.setText(formatNumber(self.vmax))
        self.initMapCombobox()
        self.slider.setMinimum(self.vmin)
        self.slider.setMaximum(self.vmax)
        self.slider.setLowValue(self.vmin)
        self.slider.setHighValue(self.vmax)
        # Redraw the widget
        self.redrawColorBar()
        self.canvas.draw()

    def onApply(self):
        """
        Respond to the Apply button click.
        Send a signal to the plotter with vmin/vmax/cmap for chart update
        """
        self.apply_signal.emit(self.norm(), self.cmap())

    def initDetectorData(self):
        """
        Fill out the Detector labels
        """
        xnpts = len(self.data.x_bins)
        ynpts = len(self.data.y_bins)
        self.lblWidth.setText(formatNumber(xnpts))
        self.lblHeight.setText(formatNumber(ynpts))
        xmax = max(self.data.xmin, self.data.xmax)
        ymax = max(self.data.ymin, self.data.ymax)
        qmax = numpy.sqrt(numpy.power(xmax, 2) + numpy.power(ymax, 2))
        self.lblQmax.setText(formatNumber(qmax))
        self.lblStopRadius.setText(formatNumber(self.data.xmin))

    def initMapCombobox(self):
        """
        Fill out the combo box with all available color maps
        """
        if self._cmap in self.rmaps:
            maps = self.rmaps
            # Assure correct state of the checkbox
            self.chkReverse.setChecked(True)
        else:
            maps = self.maps
            # Assure correct state of the checkbox
            self.chkReverse.setChecked(False)

        self.cbColorMap.addItems(maps)
        # Set the default/passed map
        self.cbColorMap.setCurrentIndex(self.cbColorMap.findText(self._cmap))

    def initRangeSlider(self):
        """
        Create and display the double slider for data range mapping.
        """
        self.slider = RangeSlider()
        self.slider.setMinimum(self.vmin)
        self.slider.setMaximum(self.vmax)
        self.slider.setLowValue(self.vmin)
        self.slider.setHighValue(self.vmax)
        self.slider.setOrientation(QtCore.Qt.Horizontal)

        self.slider_label = QtWidgets.QLabel()
        self.slider_label.setText("Drag the sliders to adjust color range.")

        def set_vmin(value):
            self.vmin = value
            self.txtMinAmplitude.setText(str(value))
            self.updateMap()
        def set_vmax(value):
            self.vmax = value
            self.txtMaxAmplitude.setText(str(value))
            self.updateMap()

        self.slider.lowValueChanged.connect(set_vmin)
        self.slider.highValueChanged.connect(set_vmax)

    def updateMap(self):
        self._norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
        self.redrawColorBar()
        self.canvas.draw()

    def initColorMap(self):
        """
        Prepare and display the color map
        """
        self.fig = mpl.figure.Figure(figsize=(4, 1))
        self.ax1 = self.fig.add_axes([0.05, 0.65, 0.9, 0.15])

        self._norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
        self.redrawColorBar()
        self.canvas = FigureCanvas(self.fig)

        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.slider_label)
        layout.addWidget(self.slider)
        layout.addWidget(self.canvas)

        self.widget.setLayout(layout)

    def onMapIndexChange(self, index):
        """
        Respond to the color map change event
        """
        new_map = str(self.cbColorMap.itemText(index))
        self._cmap = new_map
        self.redrawColorBar()
        self.canvas.draw()

    def redrawColorBar(self):
        """
        Call ColorbarBase with current values, effectively redrawing the widget
        """
        self.cb = mpl.colorbar.ColorbarBase(self.ax1, cmap=mpl.pyplot.get_cmap(self._cmap),
                                            norm=self._norm,
                                            orientation='horizontal')
        self.cb.set_label('Color map range')

    def onColorMapReversed(self, isChecked):
        """
        Respond to ticking/unticking the color map reverse checkbox
        """
        current_map = str(self.cbColorMap.currentText())
        if isChecked:
            # Add "_r" to map name for the reversed version
            new_map = current_map + "_r"
            maps = self.rmaps
            # Assure the reversed map exists.
            if new_map not in maps:
                new_map = maps[0]
        else:
            new_map = current_map[:-2] # "_r" = last two chars
            maps = self.maps
            # Base map for the reversed map should ALWAYS exist,
            # but let's be paranoid here
            if new_map not in maps:
                new_map = maps[0]

        self._cmap = new_map
        # Clear the content of the combobox.
        # Needs signal blocking, or else onMapIndexChange() spoils it all
        self.cbColorMap.blockSignals(True)
        self.cbColorMap.clear()
        # Add the new set of maps
        self.cbColorMap.addItems(maps)
        # Remove the signal block before the proper index set
        self.cbColorMap.blockSignals(False)
        self.cbColorMap.setCurrentIndex(self.cbColorMap.findText(new_map))

    def onAmplitudeChange(self):
        """
        Respond to editing the amplitude fields
        """
        min_amp = self.vmin
        max_amp = self.vmax

        try:
            min_amp = float(self.txtMinAmplitude.text())
        except ValueError:
            pass
        try:
            max_amp = float(self.txtMaxAmplitude.text())
        except ValueError:
            pass

        self._norm = mpl.colors.Normalize(vmin=min_amp, vmax=max_amp)
        self.redrawColorBar()
        self.canvas.draw()
<|endoftext|>


<commit_before>"""
.. module:: errorgeopy
   :platform: Unix
   :synopsis: Python geocoding in a disagreeing world.

.. moduleauthor:: Richard Law <richard.m.law@gmail.com>
"""

__version__ = '1.0.0'

def _get_default_pool_members():
    from geopy.geocoders import Nominatim, GoogleV3
    return (Nominatim(), GoogleV3())

DEFAULT_GEOCODER_POOL = _get_default_pool_members()
<commit_msg>increment version<commit_after>"""
.. module:: errorgeopy
   :platform: Unix
   :synopsis: Python geocoding in a disagreeing world.

.. moduleauthor:: Richard Law <richard.m.law@gmail.com>
"""

__version__ = '1.0.1'

def _get_default_pool_members():
    from geopy.geocoders import Nominatim, GoogleV3
    return (Nominatim(), GoogleV3())

DEFAULT_GEOCODER_POOL = _get_default_pool_members()
<|endoftext|>


<commit_before>__version__ = '0.1'
<commit_msg>Bump version for dev.<commit_after>__version__ = '0.2.dev1'
<|endoftext|>


<commit_before>from Configuration import Configuration
# import pUtil
import sys
import os

def set_environment():

    env = Configuration()

    # env['version'] = pUtil.getPilotVersion()
    env['pilot_version_tag'] = 'RC' #'PR' #'RC'
    env['force_devpilot'] = False

    # get the python executable path, to be used by child job
    env['pyexe'] = sys.executable

    # default job dispatcher web server
    env['pshttpurl'] = "pandaserver.cern.ch"   # Should be set with pilot option -w <url>

    # cache URL used by some experiment classes
    env['cache'] = ""

    # URL for the schedconfig
    env['schedconfigURL'] = "http://" + env['pshttpurl'] # Could be set with pilot option -I <url>

    # global variables with some default (test) values
    env['psport'] = 25443                      # PanDA server listening port
    env['uflag'] = None                        # Flag if this pilot is to only get a user analysis job or ordinary job
    env['abortDQSpace'] = 10                   # Unit is in GB
    env['warnDQSpace'] = 100                   # Unit is in GB
    env['localsizelimit_stdout'] = 2*1024**2   # Size limit of payload stdout size during running. unit is in kB
    env['localspacelimit'] = 2*1024**2         # Space limit of remaining local disk size during running. unit is in kB
    env['localspacelimit_user'] = 7*1024**2    # Maximum size of user work area. unit is in kB
    env['localspacelimit0'] = 5*1024**2        # Initial space limit before asking for a job. unit is in kB
    env['outputlimit'] = 500*1024**3           # Maximum allowed size of an output file. unit is B
    env['pilotId'] = 'xtestP001'               # Pilot id
    env['jobSchedulerId'] = 'xtestJS001'       # Scheduler id
    env['heartbeatPeriod'] = 30*60             # 30 minutes limit for checking job state file
    env['maxjobrecDefault'] = 20               # Default maximum number of job recoveries
    env['maxjobrec'] = env['maxjobrecDefault'] # Maximum number of job recoveries
    env['zombieJobList'] = []                  # Zombie list
    env['prodJobDone'] = False                 # Will be set to true when production job has finished
    env['rmwkdir'] = None                      # If set to true, workdir will always be deleted
    env['jobrec'] = True                       # Perform job recovery per default
    env['jobRequestFlag'] = True               # Ask server for initial job (read from file otherwise)
    env['debugLevel'] = 0                      # 0: debug info off, 1: display function name when called, 2: full debug info
    env['maxNumberOfRecoveryAttempts'] = 15    # As attempted by the job recovery
    env['stagein'] = False                     # Set to True during stagein phase
    env['stageout'] = False                    # Set to True during stageout phase
    env['queuename'] = ""                      # Name of queue used to download config info
    env['loopingLimitDefaultProd'] = 12*3600   # If job does not write anything in N hours, it is considered a looping job (production jobs)
    env['loopingLimitDefaultUser'] = 3*3600    # If job does not write anything in N hours, it is considered a looping job (user analysis jobs)
    env['loopingLimitMinDefault'] = 2*3600     # Minimum allow looping limit
    env['stageinretry'] = 2                    # Number of stage-in tries
    env['stageoutretry'] = 2                   # Number of stage-out tries
    env['logFileDir'] = ""                     # Log file directory
    env['testLevel'] = "0"                     # Test suite control variable (0: no test, 1: put error, 2: ...)
    env['loggingMode'] = None                  # True puts the pilot in logging mode (will not run a job), False is a normal pilot. None means run as before
    env['updateServerFlag'] = True             # Switch on/off Panda server updates (for testing only), default is True
    env['globalSite'] = None                   # Site object
    env['globalWorkNode'] = None               # Work node object
    env['globalJob'] = None                    # Job object
    env['memory'] = None                       # For memory restrictions when asking the dispatcher for a new job (MB)
    env['proxycheckFlag'] = True               # True (default): perform proxy validity checks, False: no check
    env['wrapperFlag'] = False                 # True for wrappers that expect an exit code via return, False (default) when exit() can be used
    env['cleanupLimit'] = 2                    # Cleanup time limit in hours, see Cleaner.py
    env['logTransferred'] = False              # Boolean to keep track of whether the log has been transferred or not
    env['errorLabel'] = "WARNING"              # Set to FAILED when job recovery is not used
    env['nSent'] = 0                           # Throttle variable
    env['proxyguard'] = None                   # Global proxyguard object needed for the signal handler to be able to restore the proxy if necessary
    env['jobRecoveryMode'] = None              # When true, the pilot only runs job recovery, no payload download
    env['pilotToken'] = None                   # Pilot authentication token
    env['countryGroup'] = ""                   # Country group selector for getJob request
    env['workingGroup'] = ""                   # Working group selector for getJob request
    env['allowOtherCountry'] = False           #
    env['allowSameUser'] = True                # Boolean, multi-jobs will only come from same taskID (and thus same user)
    env['taskID'] = ""                         # TaskID of the current job or the first downloaded job in case allowSameUser is set to True
    env['inputDir'] = ""                       # Location of input files (source for mv site mover)
    env['outputDir'] = ""                      # Location of output files (destination for mv site mover)
    env['jobIds'] = []                         # Global job id list
    env['stageoutStartTime'] = None            # Set when pilot receives "stageout" state message, used by looping job killer
    env['timefloor_default'] = None            # Time limit for multi-jobs in minutes (turned off by default)
    env['useCoPilot'] = False                  # CERNVM Co-Pilot framework (on: let Co-Pilot finish job, off: let pilot finish job (default))
    env['update_freq_server'] = 30*60          # Server update frequency, 30 minutes
    env['experiment'] = "ATLAS"                # Current experiment (can be set with pilot option -F <experiment>)
    env['getjobmaxtime'] = 3*60                # Maximum time the pilot will attempt to download a single job (seconds)
    env['pandaJobDataFileName'] = "pandaJobData.out" # Job definition file name
    env['verifySoftwareDirectory'] = True      # Normally the software directory should be verified, but potentially not on an HPC system
    env['workerNode'] = None
    env['hasQueueData'] = None
    env['stdout_tail'] = ""
    env['stdout_path'] = ""
    env['maxtime'] = 999999                    # The maximum time a job is allowed to run (always set in pilot using schedconfig.maxtime)
    # panda proxy variables
    env['panda_proxy_url'] = ""
    env['panda_proxy_port'] = ""
    env['resourceType'] = ""

    # to test site mover
    env['copytool'] = 'gfal-copy'
    env['copytoolin'] = 'gfal-copy'

    # Harvester
    env['harvester'] = False                    # Harvester mode, True if pilot is launched by Harvester

    # In case the PILOT_INITDIR file has not been created yet, which means that this module is being
    # used by pilot.py, it will be created here using the current directory as init dir
    print "Current directory:", os.getcwd()
    path = os.path.join(os.getcwd(), "PILOT_INITDIR")
    if not os.path.exists(path):
        writeToFile(path, os.getcwd())
    else:
        print "File exists:", path
    env['pilot_initdir'] = readStringFromFile("PILOT_INITDIR")

    # data structure for job information in main loop
    # jobDic = {'prod/analy':[pid,job,os.getpgrp]}
    # (analysis label currently not used)
    env['jobDic'] = {}

    # some default values
    env['sitename'] = "testsite"
    env['workdir'] = os.getcwd() # "/tmp"
    env['appdir'] = ""
    env['return_code'] = None
    env['return'] = 0

    return env

def readStringFromFile(filename):
    """ read exit code from file <workdir>/EXITCODE """

    s = ""
    if os.path.exists(filename):
        try:
            f = open(filename, "r")
        except Exception, e:
            print "Failed to open %s: %s" % (filename, str(e))
        else:
            s = f.read()
            print "Found string %s in file %s" % (s, filename)
            f.close()
    else:
        print "No string to report (file %s does not exist)" % (filename)

    return s

def storePilotInitdir(targetdir, pilot_initdir):
    """ Store the pilot launch directory in a file used by environment.py """

    # This function is used to store the location of the init directory in the init directory itself as well as in the
    # site work directory. The location file is used by environment.py to set the global env['pilot_initdir'] used
    # by the pilot and the Monitor

    # This function must be called before the global env variable is instantiated in the pilot

    print "Storing PILOT_INITDIR2"
    path = os.path.join(targetdir, "PILOT_INITDIR2")
    writeToFile(path, pilot_initdir)

def writeToFile(filename, s):
    """ write string s to file """

    try:
        f = open(filename, "w")
    except Exception, e:
        print "!!WARNING!!2990!! Could not open: %s, %s" % (filename, str(e))
    else:
        f.write("%s" % (s))
        f.close()
        print 'Wrote string "%s" to file: %s' % (s, filename)
<commit_msg>Updated rc<commit_after>from Configuration import Configuration
# import pUtil
import sys
import os

def set_environment():

    env = Configuration()

    # env['version'] = pUtil.getPilotVersion()
    env['pilot_version_tag'] = 'PR' #'RC'
    env['force_devpilot'] = False

    # get the python executable path, to be used by child job
    env['pyexe'] = sys.executable

    # default job dispatcher web server
    env['pshttpurl'] = "pandaserver.cern.ch"   # Should be set with pilot option -w <url>

    # cache URL used by some experiment classes
    env['cache'] = ""

    # URL for the schedconfig
    env['schedconfigURL'] = "http://" + env['pshttpurl'] # Could be set with pilot option -I <url>

    # global variables with some default (test) values
    env['psport'] = 25443                      # PanDA server listening port
    env['uflag'] = None                        # Flag if this pilot is to only get a user analysis job or ordinary job
    env['abortDQSpace'] = 10                   # Unit is in GB
    env['warnDQSpace'] = 100                   # Unit is in GB
    env['localsizelimit_stdout'] = 2*1024**2   # Size limit of payload stdout size during running. unit is in kB
    env['localspacelimit'] = 2*1024**2         # Space limit of remaining local disk size during running. unit is in kB
    env['localspacelimit_user'] = 7*1024**2    # Maximum size of user work area. unit is in kB
    env['localspacelimit0'] = 5*1024**2        # Initial space limit before asking for a job. unit is in kB
    env['outputlimit'] = 500*1024**3           # Maximum allowed size of an output file. unit is B
    env['pilotId'] = 'xtestP001'               # Pilot id
    env['jobSchedulerId'] = 'xtestJS001'       # Scheduler id
    env['heartbeatPeriod'] = 30*60             # 30 minutes limit for checking job state file
    env['maxjobrecDefault'] = 20               # Default maximum number of job recoveries
    env['maxjobrec'] = env['maxjobrecDefault'] # Maximum number of job recoveries
    env['zombieJobList'] = []                  # Zombie list
    env['prodJobDone'] = False                 # Will be set to true when production job has finished
    env['rmwkdir'] = None                      # If set to true, workdir will always be deleted
    env['jobrec'] = True                       # Perform job recovery per default
    env['jobRequestFlag'] = True               # Ask server for initial job (read from file otherwise)
    env['debugLevel'] = 0                      # 0: debug info off, 1: display function name when called, 2: full debug info
    env['maxNumberOfRecoveryAttempts'] = 15    # As attempted by the job recovery
    env['stagein'] = False                     # Set to True during stagein phase
    env['stageout'] = False                    # Set to True during stageout phase
    env['queuename'] = ""                      # Name of queue used to download config info
    env['loopingLimitDefaultProd'] = 12*3600   # If job does not write anything in N hours, it is considered a looping job (production jobs)
    env['loopingLimitDefaultUser'] = 3*3600    # If job does not write anything in N hours, it is considered a looping job (user analysis jobs)
    env['loopingLimitMinDefault'] = 2*3600     # Minimum allow looping limit
    env['stageinretry'] = 2                    # Number of stage-in tries
    env['stageoutretry'] = 2                   # Number of stage-out tries
    env['logFileDir'] = ""                     # Log file directory
    env['testLevel'] = "0"                     # Test suite control variable (0: no test, 1: put error, 2: ...)
    env['loggingMode'] = None                  # True puts the pilot in logging mode (will not run a job), False is a normal pilot. None means run as before
    env['updateServerFlag'] = True             # Switch on/off Panda server updates (for testing only), default is True
    env['globalSite'] = None                   # Site object
    env['globalWorkNode'] = None               # Work node object
    env['globalJob'] = None                    # Job object
    env['memory'] = None                       # For memory restrictions when asking the dispatcher for a new job (MB)
    env['proxycheckFlag'] = True               # True (default): perform proxy validity checks, False: no check
    env['wrapperFlag'] = False                 # True for wrappers that expect an exit code via return, False (default) when exit() can be used
    env['cleanupLimit'] = 2                    # Cleanup time limit in hours, see Cleaner.py
    env['logTransferred'] = False              # Boolean to keep track of whether the log has been transferred or not
    env['errorLabel'] = "WARNING"              # Set to FAILED when job recovery is not used
    env['nSent'] = 0                           # Throttle variable
    env['proxyguard'] = None                   # Global proxyguard object needed for the signal handler to be able to restore the proxy if necessary
    env['jobRecoveryMode'] = None              # When true, the pilot only runs job recovery, no payload download
    env['pilotToken'] = None                   # Pilot authentication token
    env['countryGroup'] = ""                   # Country group selector for getJob request
    env['workingGroup'] = ""                   # Working group selector for getJob request
    env['allowOtherCountry'] = False           #
    env['allowSameUser'] = True                # Boolean, multi-jobs will only come from same taskID (and thus same user)
    env['taskID'] = ""                         # TaskID of the current job or the first downloaded job in case allowSameUser is set to True
    env['inputDir'] = ""                       # Location of input files (source for mv site mover)
    env['outputDir'] = ""                      # Location of output files (destination for mv site mover)
    env['jobIds'] = []                         # Global job id list
    env['stageoutStartTime'] = None            # Set when pilot receives "stageout" state message, used by looping job killer
    env['timefloor_default'] = None            # Time limit for multi-jobs in minutes (turned off by default)
    env['useCoPilot'] = False                  # CERNVM Co-Pilot framework (on: let Co-Pilot finish job, off: let pilot finish job (default))
    env['update_freq_server'] = 30*60          # Server update frequency, 30 minutes
    env['experiment'] = "ATLAS"                # Current experiment (can be set with pilot option -F <experiment>)
    env['getjobmaxtime'] = 3*60                # Maximum time the pilot will attempt to download a single job (seconds)
    env['pandaJobDataFileName'] = "pandaJobData.out" # Job definition file name
    env['verifySoftwareDirectory'] = True      # Normally the software directory should be verified, but potentially not on an HPC system
    env['workerNode'] = None
    env['hasQueueData'] = None
    env['stdout_tail'] = ""
    env['stdout_path'] = ""
    env['maxtime'] = 999999                    # The maximum time a job is allowed to run (always set in pilot using schedconfig.maxtime)
    # panda proxy variables
    env['panda_proxy_url'] = ""
    env['panda_proxy_port'] = ""
    env['resourceType'] = ""

    # to test site mover
    env['copytool'] = 'gfal-copy'
    env['copytoolin'] = 'gfal-copy'

    # Harvester
    env['harvester'] = False                    # Harvester mode, True if pilot is launched by Harvester

    # In case the PILOT_INITDIR file has not been created yet, which means that this module is being
    # used by pilot.py, it will be created here using the current directory as init dir
    print "Current directory:", os.getcwd()
    path = os.path.join(os.getcwd(), "PILOT_INITDIR")
    if not os.path.exists(path):
        writeToFile(path, os.getcwd())
    else:
        print "File exists:", path
    env['pilot_initdir'] = readStringFromFile("PILOT_INITDIR")

    # data structure for job information in main loop
    # jobDic = {'prod/analy':[pid,job,os.getpgrp]}
    # (analysis label currently not used)
    env['jobDic'] = {}

    # some default values
    env['sitename'] = "testsite"
    env['workdir'] = os.getcwd() # "/tmp"
    env['appdir'] = ""
    env['return_code'] = None
    env['return'] = 0

    return env

def readStringFromFile(filename):
    """ read exit code from file <workdir>/EXITCODE """

    s = ""
    if os.path.exists(filename):
        try:
            f = open(filename, "r")
        except Exception, e:
            print "Failed to open %s: %s" % (filename, str(e))
        else:
            s = f.read()
            print "Found string %s in file %s" % (s, filename)
            f.close()
    else:
        print "No string to report (file %s does not exist)" % (filename)

    return s

def storePilotInitdir(targetdir, pilot_initdir):
    """ Store the pilot launch directory in a file used by environment.py """

    # This function is used to store the location of the init directory in the init directory itself as well as in the
    # site work directory. The location file is used by environment.py to set the global env['pilot_initdir'] used
    # by the pilot and the Monitor

    # This function must be called before the global env variable is instantiated in the pilot

    print "Storing PILOT_INITDIR2"
    path = os.path.join(targetdir, "PILOT_INITDIR2")
    writeToFile(path, pilot_initdir)

def writeToFile(filename, s):
    """ write string s to file """

    try:
        f = open(filename, "w")
    except Exception, e:
        print "!!WARNING!!2990!! Could not open: %s, %s" % (filename, str(e))
    else:
        f.write("%s" % (s))
        f.close()
        print 'Wrote string "%s" to file: %s' % (s, filename)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

from calendar import monthrange
from decimal import Decimal
import json

import numpy as np

class DataModel(object):
    def __init__(self, data=None):
        self.__dict__.update(self.defaults())
        self.__dict__.update(data or {})
        self.__dict__.update(self.date_guides())

    def defaults(self):
        NLU = 16
        NAnimals = 9

        return {
            'BasinId': 0,
            'AgSlope3to8': 0,
            'NumUAs': 0,
            'UABasinArea': 0,

            # Variables that are passed to PRedICT and are no longer used.
            'InName': '',
            'OutName': '',
            'UnitsFileFlag': 1,
            'AssessDate': '',
            'VersionNo': '',
            'ProjName': '',

            'NLU': NLU,
            'NAnimals': NAnimals,

            'Landuse': np.zeros(NLU, dtype=object),
            'Area': np.zeros(NLU),
            'CN': np.zeros(NLU),
            'KF': np.zeros(NLU),
            'LS': np.zeros(NLU),
            'C': np.zeros(NLU),
            'P': np.zeros(NLU),

            'NumNormalSys': np.zeros(12, dtype=int),
            'NumPondSys': np.zeros(12),
            'NumShortSys': np.zeros(12),
            'NumDischargeSys': np.zeros(12),
            'NumSewerSys': np.zeros(12),

            'SEDFEN': 0,
            'NFEN': 0,
            'PFEN': 0,

            'n86': 0,
            'n87': 0,
            'n88': 0,
            'n89': 0,
            'n90': 0,
            'n91': 0,
            'n92': 0,
            'n93': 0,
            'n94': 0,
            'n95': 0,
            'n95b': 0,
            'n95c': 0,
            'n95d': 0,
            'n95e': 0,

            'n96': 0,
            'n97': 0,
            'n98': 0,
            'n99': 0,
            'n99b': 0,
            'n99c': 0,
            'n99d': 0,
            'n99e': 0,
            'n100': 0,
            'n101': 0,
            'n101b': 0,
            'n101c': 0,
            'n101d': 0,
            'n101e': 0,
            'n102': 0,
            'n103a': 0,
            'n103b': 0,
            'n103c': 0,
            'n103d': 0,

            'n104': 0,
            'n105': 0,
            'n106': 0,
            'n106b': 0,
            'n106c': 0,
            'n106d': 0,
            'n107': 0,
            'n107b': 0,
            'n107c': 0,
            'n107d': 0,
            'n107e': 0,

            'Storm': 0,
            'CSNAreaSim': 0,
            'CSNDevType': 'None',
        }

    def date_guides(self):
        model = self.__dict__
        output = {}
        if 'WxYrBeg' in model and 'WxYrEnd' in model:
            begyear = model['WxYrBeg']
            endyear = model['WxYrEnd']
            year_range = endyear - begyear + 1
            month_abbr = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun',
                          'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']

            if 'DaysMonth' not in model:
                output['DaysMonth'] = np.zeros((year_range, 12),
                                               dtype=np.int)
                for y in range(year_range):
                    year = begyear + y
                    for m in range(12):
                        output['DaysMonth'][y][m] = monthrange(year, m + 1)[1]

            if 'WxMonth' not in model:
                output['WxMonth'] = [month_abbr
                                     for y in range(year_range)]

            if 'WxYear' not in model:
                output['WxYear'] = [[begyear + y] * 12
                                    for y in range(year_range)]
        return output

    def __str__(self):
        return '<GWLF-E DataModel>'

    def tojson(self):
        return json.dumps(self.__dict__, cls=NumpyAwareJSONEncoder)

class NumpyAwareJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, Decimal):
            return float(obj)
        return json.JSONEncoder.default(self, obj)
<commit_msg>Set default initial value for `AdjUrbanQTotal`<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

from calendar import monthrange
from decimal import Decimal
import json

import numpy as np

class DataModel(object):
    def __init__(self, data=None):
        self.__dict__.update(self.defaults())
        self.__dict__.update(data or {})
        self.__dict__.update(self.date_guides())

    def defaults(self):
        NLU = 16
        NAnimals = 9

        return {
            'BasinId': 0,
            'AgSlope3to8': 0,
            'NumUAs': 0,
            'UABasinArea': 0,

            # Variables that are passed to PRedICT and are no longer used.
            'InName': '',
            'OutName': '',
            'UnitsFileFlag': 1,
            'AssessDate': '',
            'VersionNo': '',
            'ProjName': '',

            'NLU': NLU,
            'NAnimals': NAnimals,

            'Landuse': np.zeros(NLU, dtype=object),
            'Area': np.zeros(NLU),
            'CN': np.zeros(NLU),
            'KF': np.zeros(NLU),
            'LS': np.zeros(NLU),
            'C': np.zeros(NLU),
            'P': np.zeros(NLU),

            'NumNormalSys': np.zeros(12, dtype=int),
            'NumPondSys': np.zeros(12),
            'NumShortSys': np.zeros(12),
            'NumDischargeSys': np.zeros(12),
            'NumSewerSys': np.zeros(12),

            'SEDFEN': 0,
            'NFEN': 0,
            'PFEN': 0,

            'n86': 0,
            'n87': 0,
            'n88': 0,
            'n89': 0,
            'n90': 0,
            'n91': 0,
            'n92': 0,
            'n93': 0,
            'n94': 0,
            'n95': 0,
            'n95b': 0,
            'n95c': 0,
            'n95d': 0,
            'n95e': 0,

            'n96': 0,
            'n97': 0,
            'n98': 0,
            'n99': 0,
            'n99b': 0,
            'n99c': 0,
            'n99d': 0,
            'n99e': 0,
            'n100': 0,
            'n101': 0,
            'n101b': 0,
            'n101c': 0,
            'n101d': 0,
            'n101e': 0,
            'n102': 0,
            'n103a': 0,
            'n103b': 0,
            'n103c': 0,
            'n103d': 0,

            'n104': 0,
            'n105': 0,
            'n106': 0,
            'n106b': 0,
            'n106c': 0,
            'n106d': 0,
            'n107': 0,
            'n107b': 0,
            'n107c': 0,
            'n107d': 0,
            'n107e': 0,

            'Storm': 0,
            'CSNAreaSim': 0,
            'CSNDevType': 'None',
            'AdjUrbanQTotal': 0,
        }

    def date_guides(self):
        model = self.__dict__
        output = {}
        if 'WxYrBeg' in model and 'WxYrEnd' in model:
            begyear = model['WxYrBeg']
            endyear = model['WxYrEnd']
            year_range = endyear - begyear + 1
            month_abbr = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun',
                          'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']

            if 'DaysMonth' not in model:
                output['DaysMonth'] = np.zeros((year_range, 12),
                                               dtype=np.int)
                for y in range(year_range):
                    year = begyear + y
                    for m in range(12):
                        output['DaysMonth'][y][m] = monthrange(year, m + 1)[1]

            if 'WxMonth' not in model:
                output['WxMonth'] = [month_abbr
                                     for y in range(year_range)]

            if 'WxYear' not in model:
                output['WxYear'] = [[begyear + y] * 12
                                    for y in range(year_range)]
        return output

    def __str__(self):
        return '<GWLF-E DataModel>'

    def tojson(self):
        return json.dumps(self.__dict__, cls=NumpyAwareJSONEncoder)

class NumpyAwareJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, Decimal):
            return float(obj)
        return json.JSONEncoder.default(self, obj)
<|endoftext|>


<commit_before>from math import log

from ..ch3.model import LeastSquareModel, LinearModel
import numpy as np
from numpy import linalg as LA
from ..utils import lazy_method

class LinearRegression(LinearModel):
    def __init__(self, *args, **kwargs):
        self.K = kwargs.pop('K', 1)
        super().__init__(*args, **kwargs)

    @property
    @lazy_method
    def error_rate(self):
        return 1 - np.sum((self._raw_train_y == self.y_hat))/ self.N

class LinearRegressionIndicatorMatrix(LeastSquareModel):
    def __init__(self, *args, **kwargs):
        self.K = kwargs.pop('K', 1)
        super().__init__(*args, **kwargs)

    def _pre_processing_y(self, y):
        iy = y.flatten()
        N = y.shape[0]
        if self.K > 1:
            Y = np.zeros((N, self.K))
            for i in range(N):
                k = iy[i]
                # k starts from 1
                Y[i, k-1] = 1
        else:
            return super()._pre_processing_y(y)
        return Y

    def predict(self, X):
        Y_hat = super().predict(X)
        y = (Y_hat.argmax(axis=1)).reshape((-1, 1)) + 1
        return y

    @property
    @lazy_method
    def error_rate(self):
        return (1 - np.sum((self._raw_train_y == self.y_hat)) / self.N)

class LDAModel(LinearRegression):
    """
    Linear Discriminant Analysis
    from page 106
    """

    def _pre_processing_x(self, X):
        X = self.standardize(X)
        return X

    def train(self):
        X = self.train_x
        y = self.train_y
        K = self.K
        p = self.p

        self.Mu = np.zeros((K, p))
        self.Pi = np.zeros((K, 1))
        self.Sigma_hat = np.zeros((p, p))

        for k in range(K):
            mask = (y == k+1)
            N_k = sum(mask)

            X_k = X[mask.flatten(), :]

            self.Pi[k] = N_k / self.N
            self.Mu[k] = np.sum(X_k, axis=0).reshape((1, -1)) / N_k
            self.Sigma_hat = self.Sigma_hat + ((X_k - self.Mu[k]).T @ (X_k - self.Mu[k])) / (self.N - K)

    def linear_discriminant_func(self, x, k):
        """
        linear discriminant function.
        Define by (4.10)
        :return: delta_k(x)
        """
        mu_k = self.Mu[k]
        pi_k = self.Pi[k]
        sigma_inv = self.math.pinv(self.Sigma_hat)
        result = mu_k @ sigma_inv @ x.T - (mu_k @ sigma_inv @ mu_k.T)/2 + log(pi_k)
        return result

    def predict(self, X):
        X = self._pre_processing_x(X)
        N = X.shape[0]
        Y = np.zeros((N, self.K))

        for k in range(self.K):
            # delta_k is (N x 1)
            delta_k = self.linear_discriminant_func(X, k)
            Y[:, k] = delta_k

        # make the k start from 1
        y_hat = Y.argmax(axis=1).reshape((-1, 1)) + 1
        return y_hat

class QDAModel(LinearRegression):
    """
    Quadratic Discriminant Analysis
    pp. 110

    Ref
    ---
    http://www.wikicoursenote.com/wiki/Stat841#In_practice
    """

    def train(self):
        X = self.train_x
        y = self.train_y
        K = self.K
        p = self.p

        self.Mu = np.zeros((K, p))
        self.Pi = np.zeros((K, 1))
        self.Sigma_hat = []

        for k in range(K):
            mask = (y == k+1)
            N_k = sum(mask)
            X_k = X[mask.flatten(), :]

            self.Pi[k] = N_k / self.N
            self.Mu[k] = np.sum(X_k, axis=0).reshape((1, -1)) / N_k
            # We div by N_k instead of (N-K)
            self.Sigma_hat.append(((X_k - self.Mu[k]).T @ (X_k - self.Mu[k])) / N_k)

    def quadratic_discriminant_func(self, x, k):
        mu_k = self.Mu[k]
        pi_k = self.Pi[k]
        sigma_k = self.Sigma_hat[k]
        pinv = self.math.pinv

        # assume that each row of x contain observation
        result = -(np.log(np.linalg.det(sigma_k)))/2 - \
                 ((x - mu_k) @ pinv(sigma_k, rcond=0) @ (x - mu_k).T)/2 + log(pi_k)
        return result

    def predict(self, X):
        X = self._pre_processing_x(X)
        N = X.shape[0]
        Y = np.zeros((N, self.K))

        for k in range(self.K):
            # the intuitive solution is use np.apply_along_axis, but is too slow
            # delta_k is (N x 1)
            # delta_k_func = partial(self.linear_discriminant_func, k=k)
            # delta_k = np.apply_along_axis(delta_k_func, 1, X)

            # d is NxN,
            # Let B = A@A.T, the diagonal of B is [A[i] @ A[i].T for i in range(A.shape(0)]
            d = self.quadratic_discriminant_func(X, k)
            Y[:, k] = d.diagonal()

        # make the k start from 1
        y_hat = Y.argmax(axis=1).reshape((-1, 1)) + 1
        return y_hat

class RDAModel(QDAModel):
    def __init__(self, *args, alpha=1, **kwargs):
        self.alpha = alpha
        super().__init__(*args, **kwargs)

    def train(self):
        X = self.train_x
        y = self.train_y
        K = self.K
        p = self.p

        self.Mu = np.zeros((K, p))
        self.Pi = np.zeros((K, 1))
        # list of sigma_k
        self.Sigma_hat = []
        # the sum of sigma_k
        self.Sigma_tot = np.zeros((1, p))

        for k in range(K):
            mask = (y == k+1)
            N_k = sum(mask)

            X_k = X[mask.flatten(), :]

            self.Pi[k] = N_k / self.N
            self.Mu[k] = np.sum(X_k, axis=0).reshape((1, -1)) / N_k
            # We div by N_k instead of (N-K)
            self.Sigma_hat.append(((X_k - self.Mu[k]).T @ (X_k - self.Mu[k])) / N_k)
            self.Sigma_tot = self.Sigma_tot + (X_k - self.Mu[k]).T @ (X_k - self.Mu[k])

        self.Sigma_tot = self.Sigma_tot / (self.N - K)

        for k in range(K):
            self.Sigma_hat[k] = (self.Sigma_hat[k] * self.alpha) + self.Sigma_tot * (1 - self.alpha)

class LDAForComputation(LDAModel):

    def train(self):
        super().train()
        sigma = self.Sigma_hat
        D_, U = LA.eigh(sigma)
        D = np.diagflat(D_)
        self.A = np.power(LA.pinv(D), 0.5) @ U.T

    def predict(self, X):
        X = self._pre_processing_x(X)
        Y = np.zeros((X.shape[0], self.K))
        A = self.A

        # because X is (N x p), A is (K x p), we can to get the X_star (NxK)
        X_star = X @ A.T

        for k in range(self.K):
            # mu_s_star shape is (p,)
            mu_k_star = A @ self.Mu[k]

            # Ref: http://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.norm.html
            # Ref: http://stackoverflow.com/questions/1401712/how-can-the-euclidean-distance-be-calculated-with-numpy
            Y[:, k] = LA.norm(X_star - mu_k_star, axis=1) * 0.5 - log(self.Pi[k])

        # Python index start from 0, transform to start with 1
        y_hat = Y.argmin(axis=1).reshape((-1, 1)) + 1
        return y_hat

class ReducedRankLDAModel(LDAForComputation):
    """
    page 113, 4.3.3

    I also write a blog describe how to write RRLDA:
    http://littlezz.github.io/how-to-write-reduced-rank-linear-discriminant-analysis-with-python.html

    ref: http://sites.stat.psu.edu/~jiali/course/stat597e/notes2/lda2.pdf
    """

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

    def train(self):
        super().train()
        W = self.Sigma_hat
        # prior probabilities (K, 1)
        Pi = self.Pi
        # class centroids (K, p)
        Mu = self.Mu
        p = self.p
        # the number of class
        K = self.K
        # the dimension you want
        L = self.L

        # Mu is (K, p) matrix, Pi is (K, 1)
        mu = np.sum(Pi * Mu, axis=0)
        B = np.zeros((p, p))

        for k in range(K):
            # vector @ vector equal scalar, use vector[:, None] to transform to matrix
            # vec[:, None] equal to vec.reshape((1, vec.shape[0]))
            B = B + Pi[k]*((Mu[k] - mu)[:, None] @ ((Mu[k] - mu)[None, :]))

        # Be careful, the `eigh` method get the eigenvalues in ascending , which is opposite to R.
        Dw, Uw = LA.eigh(W)
        # reverse the Dw_ and Uw
        Dw = Dw[::-1]
        Uw = np.fliplr(Uw)

        W_half = self.math.pinv(np.diagflat(Dw**0.5) @ Uw.T)
        B_star = W_half.T @ B @ W_half
        D_, V = LA.eigh(B_star)

        # reverse V
        V = np.fliplr(V)

        # overwrite `self.A` so that we can reuse `predict` method define in parent class
        self.A = np.zeros((L, p))
        for l in range(L):
            self.A[l, :] = W_half @ V[:, l]

class BinaryLogisticRegression(LinearRegression):
    """
    page 119.
    two class case

    note that self.W is the second partial derivative.

    ref: http://sites.stat.psu.edu/~jiali/course/stat597e/notes2/logit.pdf
    """
    def __init__(self, *args, max_iter=500, **kwargs):
        self.max_iter = max_iter
        super().__init__(*args, **kwargs)

    def _pre_processing_x(self, X):
        # TODO: refactor that allow choose whether use standardizing
        # X = super()._pre_processing_x(X)
        X = np.insert(X, 0, [1], axis=1)
        return X

    def train(self):
        X = self.train_x
        y = self.train_y
        # include intercept
        beta = np.zeros((self.p+1, 1))

        iter_times = 0
        while True:
            e_X = np.exp(X @ beta)

            # N x 1
            self.P = e_X / (1 + e_X)

            # W is a vector
            self.W = (self.P * (1 - self.P)).flatten()

            beta = beta + self.math.pinv((X.T * self.W) @ X) @ X.T @ (y - self.P)

            iter_times += 1
            if iter_times > self.max_iter:
                break

        self.beta_hat = beta

    def predict(self, X):
        X = self._pre_processing_x(X)
        y = X @ self.beta_hat
        y[y>=0] = 1
        y[y<0] = 0
        return y

    @property
    def std_err(self):
        """
        ref: https://groups.google.com/d/msg/comp.soft-sys.stat.spss/Fv7Goxs_Bwk/ff0jCesG8REJ
        """
        return self.math.pinv(self.train_x.T * self.W @ self.train_x).diagonal() ** 0.5
<commit_msg>update comment<commit_after>from math import log

from ..ch3.model import LeastSquareModel, LinearModel
import numpy as np
from numpy import linalg as LA
from ..utils import lazy_method

class LinearRegression(LinearModel):
    def __init__(self, *args, **kwargs):
        self.K = kwargs.pop('K', 1)
        super().__init__(*args, **kwargs)

    @property
    @lazy_method
    def error_rate(self):
        return 1 - np.sum((self._raw_train_y == self.y_hat))/ self.N

class LinearRegressionIndicatorMatrix(LeastSquareModel):
    def __init__(self, *args, **kwargs):
        self.K = kwargs.pop('K', 1)
        super().__init__(*args, **kwargs)

    def _pre_processing_y(self, y):
        iy = y.flatten()
        N = y.shape[0]
        if self.K > 1:
            Y = np.zeros((N, self.K))
            for i in range(N):
                k = iy[i]
                # k starts from 1
                Y[i, k-1] = 1
        else:
            return super()._pre_processing_y(y)
        return Y

    def predict(self, X):
        Y_hat = super().predict(X)
        y = (Y_hat.argmax(axis=1)).reshape((-1, 1)) + 1
        return y

    @property
    @lazy_method
    def error_rate(self):
        return (1 - np.sum((self._raw_train_y == self.y_hat)) / self.N)

class LDAModel(LinearRegression):
    """
    Linear Discriminant Analysis
    from page 106
    """

    def _pre_processing_x(self, X):
        X = self.standardize(X)
        return X

    def train(self):
        X = self.train_x
        y = self.train_y
        K = self.K
        p = self.p

        self.Mu = np.zeros((K, p))
        self.Pi = np.zeros((K, 1))
        self.Sigma_hat = np.zeros((p, p))

        for k in range(K):
            mask = (y == k+1)
            N_k = sum(mask)

            X_k = X[mask.flatten(), :]

            self.Pi[k] = N_k / self.N
            self.Mu[k] = np.sum(X_k, axis=0).reshape((1, -1)) / N_k
            self.Sigma_hat = self.Sigma_hat + ((X_k - self.Mu[k]).T @ (X_k - self.Mu[k])) / (self.N - K)

    def linear_discriminant_func(self, x, k):
        """
        linear discriminant function.
        Define by (4.10)
        :return: delta_k(x)
        """
        mu_k = self.Mu[k]
        pi_k = self.Pi[k]
        sigma_inv = self.math.pinv(self.Sigma_hat)
        result = mu_k @ sigma_inv @ x.T - (mu_k @ sigma_inv @ mu_k.T)/2 + log(pi_k)
        return result

    def predict(self, X):
        X = self._pre_processing_x(X)
        N = X.shape[0]
        Y = np.zeros((N, self.K))

        for k in range(self.K):
            # delta_k is (N x 1)
            delta_k = self.linear_discriminant_func(X, k)
            Y[:, k] = delta_k

        # make the k start from 1
        y_hat = Y.argmax(axis=1).reshape((-1, 1)) + 1
        return y_hat

class QDAModel(LinearRegression):
    """
    Quadratic Discriminant Analysis
    pp. 110

    Ref
    ---
    http://www.wikicoursenote.com/wiki/Stat841#In_practice
    """

    def train(self):
        X = self.train_x
        y = self.train_y
        K = self.K
        p = self.p

        self.Mu = np.zeros((K, p))
        self.Pi = np.zeros((K, 1))
        self.Sigma_hat = []

        for k in range(K):
            mask = (y == k+1)
            N_k = sum(mask)
            X_k = X[mask.flatten(), :]

            self.Pi[k] = N_k / self.N
            self.Mu[k] = np.sum(X_k, axis=0).reshape((1, -1)) / N_k
            # We div by N_k instead of (N-K)
            self.Sigma_hat.append(((X_k - self.Mu[k]).T @ (X_k - self.Mu[k])) / N_k)

    def quadratic_discriminant_func(self, x, k):
        mu_k = self.Mu[k]
        pi_k = self.Pi[k]
        sigma_k = self.Sigma_hat[k]
        pinv = self.math.pinv

        # assume that each row of x contain observation
        result = -(np.log(np.linalg.det(sigma_k)))/2 - \
                 ((x - mu_k) @ pinv(sigma_k, rcond=0) @ (x - mu_k).T)/2 + log(pi_k)
        return result

    def predict(self, X):
        X = self._pre_processing_x(X)
        N = X.shape[0]
        Y = np.zeros((N, self.K))

        for k in range(self.K):
            # the intuitive solution is use np.apply_along_axis, but is too slow
            # delta_k is (N x 1)
            # delta_k_func = partial(self.linear_discriminant_func, k=k)
            # delta_k = np.apply_along_axis(delta_k_func, 1, X)

            # d is NxN,
            # Let B = A@A.T, the diagonal of B is [A[i] @ A[i].T for i in range(A.shape(0)]
            d = self.quadratic_discriminant_func(X, k)
            Y[:, k] = d.diagonal()

        # make the k start from 1
        y_hat = Y.argmax(axis=1).reshape((-1, 1)) + 1
        return y_hat

class RDAModel(QDAModel):
    def __init__(self, *args, alpha=1, **kwargs):
        self.alpha = alpha
        super().__init__(*args, **kwargs)

    def train(self):
        X = self.train_x
        y = self.train_y
        K = self.K
        p = self.p

        self.Mu = np.zeros((K, p))
        self.Pi = np.zeros((K, 1))
        # list of sigma_k
        self.Sigma_hat = []
        # the sum of sigma_k
        self.Sigma_tot = np.zeros((1, p))

        for k in range(K):
            mask = (y == k+1)
            N_k = sum(mask)

            X_k = X[mask.flatten(), :]

            self.Pi[k] = N_k / self.N
            self.Mu[k] = np.sum(X_k, axis=0).reshape((1, -1)) / N_k
            # We div by N_k instead of (N-K)
            self.Sigma_hat.append(((X_k - self.Mu[k]).T @ (X_k - self.Mu[k])) / N_k)
            self.Sigma_tot = self.Sigma_tot + (X_k - self.Mu[k]).T @ (X_k - self.Mu[k])

        self.Sigma_tot = self.Sigma_tot / (self.N - K)

        for k in range(K):
            self.Sigma_hat[k] = (self.Sigma_hat[k] * self.alpha) + self.Sigma_tot * (1 - self.alpha)

class LDAForComputation(LDAModel):

    def train(self):
        super().train()
        sigma = self.Sigma_hat
        D_, U = LA.eigh(sigma)
        D = np.diagflat(D_)
        self.A = np.power(LA.pinv(D), 0.5) @ U.T

    def predict(self, X):
        X = self._pre_processing_x(X)
        Y = np.zeros((X.shape[0], self.K))
        A = self.A

        # because X is (N x p), A is (K x p), we can to get the X_star (NxK)
        X_star = X @ A.T

        for k in range(self.K):
            # mu_s_star shape is (p,)
            mu_k_star = A @ self.Mu[k]

            # Ref: http://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.norm.html
            # Ref: http://stackoverflow.com/questions/1401712/how-can-the-euclidean-distance-be-calculated-with-numpy
            Y[:, k] = LA.norm(X_star - mu_k_star, axis=1) * 0.5 - log(self.Pi[k])

        # Python index start from 0, transform to start with 1
        y_hat = Y.argmin(axis=1).reshape((-1, 1)) + 1
        return y_hat

class ReducedRankLDAModel(LDAForComputation):
    """
    page 113, 4.3.3

    I also write a blog describe how to write RRLDA:
    http://littlezz.github.io/how-to-write-reduced-rank-linear-discriminant-analysis-with-python.html

    ref: http://sites.stat.psu.edu/~jiali/course/stat597e/notes2/lda2.pdf
    """

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

    def train(self):
        super().train()
        W = self.Sigma_hat
        # prior probabilities (K, 1)
        Pi = self.Pi
        # class centroids (K, p)
        Mu = self.Mu
        p = self.p
        # the number of class
        K = self.K
        # the dimension you want
        L = self.L

        # Mu is (K, p) matrix, Pi is (K, 1)
        mu = np.sum(Pi * Mu, axis=0)
        B = np.zeros((p, p))

        for k in range(K):
            # vector @ vector equal scalar, use vector[:, None] to transform to matrix
            # vec[:, None] equal to vec.reshape((1, vec.shape[0]))
            B = B + Pi[k]*((Mu[k] - mu)[:, None] @ ((Mu[k] - mu)[None, :]))

        # Be careful, the `eigh` method get the eigenvalues in ascending , which is opposite to R.
        Dw, Uw = LA.eigh(W)
        # reverse the Dw_ and Uw
        Dw = Dw[::-1]
        Uw = np.fliplr(Uw)

        W_half = self.math.pinv(np.diagflat(Dw**0.5) @ Uw.T)
        B_star = W_half.T @ B @ W_half
        D_, V = LA.eigh(B_star)

        # reverse V
        V = np.fliplr(V)

        # overwrite `self.A` so that we can reuse `predict` method define in parent class
        self.A = np.zeros((L, p))
        for l in range(L):
            self.A[l, :] = W_half @ V[:, l]

class BinaryLogisticRegression(LinearRegression):
    """
    page 119.
    two class case

    note that self.W is the second partial derivative.

    ref: http://sites.stat.psu.edu/~jiali/course/stat597e/notes2/logit.pdf
    """
    def __init__(self, *args, max_iter=500, **kwargs):
        self.max_iter = max_iter
        super().__init__(*args, **kwargs)

    def _pre_processing_x(self, X):
        # TODO: refactor that allow choose whether use standardizing
        # X = super()._pre_processing_x(X)
        X = np.insert(X, 0, [1], axis=1)
        return X

    def train(self):
        X = self.train_x
        y = self.train_y
        # include intercept
        beta = np.zeros((self.p+1, 1))

        iter_times = 0
        while True:
            e_X = np.exp(X @ beta)

            # N x 1
            self.P = e_X / (1 + e_X)

            # W is a vector
            self.W = (self.P * (1 - self.P)).flatten()

            beta = beta + self.math.pinv((X.T * self.W) @ X) @ X.T @ (y - self.P)

            iter_times += 1
            if iter_times > self.max_iter:
                break

        self.beta_hat = beta

    def predict(self, X):
        X = self._pre_processing_x(X)
        y = X @ self.beta_hat
        y[y>=0] = 1
        y[y<0] = 0
        return y

    @property
    def std_err(self):
        """
        ref: https://groups.google.com/d/msg/comp.soft-sys.stat.spss/Fv7Goxs_Bwk/ff0jCesG8REJ

        intuitive formula find in : http://data.princeton.edu/wws509/notes/c3.pdf  page 10
        """
        return self.math.pinv(self.train_x.T * self.W @ self.train_x).diagonal() ** 0.5
<|endoftext|>


<commit_before>import time

import RPi.GPIO as GPIO
GPIO.setmode(GPIO.BCM)

def parse_edge_key(edge):
    if edge in [GPIO.FALLING, GPIO.RISING, GPIO.BOTH]:
        edge
    elif edge.lower() == 'falling':
        edge = GPIO.FALLING
    elif edge.lower() == 'raising':
        edge = GPIO.RISING
    elif edge.lower() == 'both':
        edge = GPIO.BOTH
    else:
        raise KeyError('Unknown Edge type {edge}'.format(edge=edge))
    return edge

# This function takes a holding time (fractional seconds), a channel, a GPIO state and an action reference (function).
# It checks if the GPIO is in the state since the function was called. If the state
# changes it return False. If the time is over the function returns True.
def checkGpioStaysInState(holdingTime, gpioChannel, gpioHoldingState):
    # Get a reference start time (https://docs.python.org/3/library/time.html#time.perf_counter)
    startTime = time.perf_counter()
    # Continously check if time is not over
    while True:
        currentState = GPIO.input(gpioChannel)
        if holdingTime < (time.perf_counter() - startTime):
            break
        # Return if state does not match holding state
        if (gpioHoldingState != currentState):
            return False
        # Else: Wait

    if (gpioHoldingState != currentState):
        return False
    return True

class SimpleButton:
    def callbackFunctionHandler(self, *args):
        if self.hold_repeat:
            return self.holdAndRepeatHandler(*args)
        return self.when_pressed(*args)

    def __init__(self, pin, action=lambda *args: None, name=None, bouncetime=500, edge=GPIO.FALLING,
                 hold_time=.1, hold_repeat=False):
        self.edge = parse_edge_key(edge)
        self.hold_time = hold_time
        self.hold_repeat = hold_repeat
        self.pull_up = True

        self.pin = pin
        self.name = name
        self.bouncetime = bouncetime
        GPIO.setup(self.pin, GPIO.IN, pull_up_down=GPIO.PUD_UP)
        self.when_pressed = action
        GPIO.add_event_detect(self.pin, edge=self.edge, callback=self.callbackFunctionHandler,
                              bouncetime=self.bouncetime)

    def set_callbackFunction(self, callbackFunction):
        self.when_pressed = callbackFunction

    def holdAndRepeatHandler(self, *args):
        # Rise volume as requested
        self.when_pressed(*args)
        # Detect holding of button
        while checkGpioStaysInState(self.hold_time, self.pin, GPIO.LOW):
            self.when_pressed(*args)

    def __del__(self):
        GPIO.remove_event_detect(self.pin)

    @property
    def is_pressed(self):
        if self.pull_up:
            return not GPIO.input(self.pin)
        return GPIO.input(self.pin)

    def __repr__(self):
        return '<SimpleButton-{}(pin {},hold_repeat={},hold_time={})>'.format(
            self.name, self.pin, self.hold_repeat, self.hold_time
        )

if __name__ == "__main__":
    print('please enter pin no to test')
    pin = int(input())
    func = lambda *args: print('FunctionCall with {}'.format(args))
    btn = SimpleButton(pin=pin, action=func, hold_repeat=True)
    while True:
        pass
<commit_msg>logging output<commit_after>import time

import RPi.GPIO as GPIO
GPIO.setmode(GPIO.BCM)

def parse_edge_key(edge):
    if edge in [GPIO.FALLING, GPIO.RISING, GPIO.BOTH]:
        edge
    elif edge.lower() == 'falling':
        edge = GPIO.FALLING
    elif edge.lower() == 'raising':
        edge = GPIO.RISING
    elif edge.lower() == 'both':
        edge = GPIO.BOTH
    else:
        raise KeyError('Unknown Edge type {edge}'.format(edge=edge))
    return edge

# This function takes a holding time (fractional seconds), a channel, a GPIO state and an action reference (function).
# It checks if the GPIO is in the state since the function was called. If the state
# changes it return False. If the time is over the function returns True.
def checkGpioStaysInState(holdingTime, gpioChannel, gpioHoldingState):
    # Get a reference start time (https://docs.python.org/3/library/time.html#time.perf_counter)
    startTime = time.perf_counter()
    # Continously check if time is not over
    while True:
        currentState = GPIO.input(gpioChannel)
        if holdingTime < (time.perf_counter() - startTime):
            break
        # Return if state does not match holding state
        if (gpioHoldingState != currentState):
            return False
        # Else: Wait

    if (gpioHoldingState != currentState):
        return False
    return True

class SimpleButton:
    def callbackFunctionHandler(self, *args):
        if self.hold_repeat:
            return self.holdAndRepeatHandler(*args)
        return self.when_pressed(*args)

    def __init__(self, pin, action=lambda *args: None, name=None, bouncetime=500, edge=GPIO.FALLING,
                 hold_time=.1, hold_repeat=False):
        self.edge = parse_edge_key(edge)
        self.hold_time = hold_time
        self.hold_repeat = hold_repeat
        self.pull_up = True

        self.pin = pin
        self.name = name
        self.bouncetime = bouncetime
        GPIO.setup(self.pin, GPIO.IN, pull_up_down=GPIO.PUD_UP)
        self.when_pressed = action
        GPIO.add_event_detect(self.pin, edge=self.edge, callback=self.callbackFunctionHandler,
                              bouncetime=self.bouncetime)

    def set_callbackFunction(self, callbackFunction):
        self.when_pressed = callbackFunction

    def holdAndRepeatHandler(self, *args):
        # Rise volume as requested
        self.when_pressed(*args)
        # Detect holding of button
        while checkGpioStaysInState(self.hold_time, self.pin, GPIO.LOW):
            self.when_pressed(*args)

    def __del__(self):
        logger.debug('remove event detection')
        GPIO.remove_event_detect(self.pin)

    @property
    def is_pressed(self):
        if self.pull_up:
            return not GPIO.input(self.pin)
        return GPIO.input(self.pin)

    def __repr__(self):
        return '<SimpleButton-{}(pin {},hold_repeat={},hold_time={})>'.format(
            self.name, self.pin, self.hold_repeat, self.hold_time
        )

if __name__ == "__main__":
    print('please enter pin no to test')
    pin = int(input())
    func = lambda *args: print('FunctionCall with {}'.format(args))
    btn = SimpleButton(pin=pin, action=func, hold_repeat=True)
    pause()
<|endoftext|>