<commit_before>from abc import abstractmethod, ABC

import numpy as np
from sklearn.base import BaseEstimator
from sklearn.utils.validation import check_is_fitted, check_X_y

from .visualization import plot_anomaly_score, plot_roc_curve

__all__ = ['is_outlier_detector', 'BaseOutlierDetector']

def is_outlier_detector(estimator):
    """Return True if the given estimator is (probably) an outlier detector.

    Parameters
    ----------
    estimator : object
        Estimator object to test.

    Returns
    -------
    out : bool
        True if estimator is an outlier detector and False otherwise.
    """

    return getattr(estimator, '_estimator_type', None) == 'outlier_detector'

class BaseOutlierDetector(BaseEstimator, ABC):
    """Base class for all outlier detectors in kenchi."""

    # TODO: Update anomaly_score method so that the normalized score can be
    # computed

    _estimator_type = 'outlier_detector'

    @abstractmethod
    def __init__(self, contamination=0.01, verbose=False):
        self.contamination = contamination
        self.verbose       = verbose

    @abstractmethod
    def fit(self, X, y=None, **fit_params):
        """Fit the model according to the given training data.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.

        y : ignored

        Returns
        -------
        self : object
            Return self.
        """

    @abstractmethod
    def anomaly_score(self, X):
        """Compute the anomaly score for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        Returns
        -------
        anomaly_score : array-like of shape (n_samples,)
            Anomaly score for each sample.
        """

    def _check_params(self):
        """Check validity of parameters and raise ValueError if not valid."""

        if not 0. <= self.contamination <= 0.5:
            raise ValueError(
                f'contamination must be between 0.0 and 0.5 inclusive '
                f'but was {self.contamination}'
            )

    def _get_threshold(self):
        """Define the threshold according to the given training data."""

        return np.percentile(
            self.anomaly_score_, 100. * (1. - self.contamination)
        )

    def predict(self, X, threshold=None):
        """Predict if a particular sample is an outlier or not.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        threshold : float, default None
            User-provided threshold.

        Returns
        -------
        y_pred : array-like of shape (n_samples,)
            Return 1 for inliers and -1 for outliers.
        """

        check_is_fitted(self, 'threshold_')

        if threshold is None:
            threshold = self.threshold_

        return np.where(self.anomaly_score(X) <= threshold, 1, -1)

    def fit_predict(self, X, y=None, threshold=None, **fit_params):
        """Fit the model according to the given training data and predict if a
        particular sample is an outlier or not.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training Data.

        y : ignored

        threshold : float, default None
            User-provided threshold.

        Returns
        -------
        y_pred : array-like of shape (n_samples,)
            Return 1 for inliers and -1 for outliers.
        """

        return self.fit(X, **fit_params).predict(X, threshold=threshold)

    def plot_anomaly_score(self, X, **kwargs):
        """Plot the anomaly score for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        ax : matplotlib Axes, default None
            Target axes instance.

        bins : int, str or array-like, default 'auto'
            Number of hist bins.

        figsize : tuple, default None
            Tuple denoting figure size of the plot.

        filename : str, default None
            If provided, save the current figure.

        grid : boolean, default True
            If True, turn the axes grids on.

        hist : bool, default True
            If True, plot a histogram of anomaly scores.

        title : string, default None
            Axes title. To disable, pass None.

        xlabel : string, default 'Samples'
            X axis title label. To disable, pass None.

        xlim : tuple, default None
            Tuple passed to `ax.xlim`.

        ylabel : string, default 'Anomaly score'
            Y axis title label. To disable, pass None.

        ylim : tuple, default None
            Tuple passed to `ax.ylim`.

        **kwargs : dict
            Other keywords passed to `ax.plot`.

        Returns
        -------
        ax : matplotlib Axes
            Axes on which the plot was drawn.
        """

        kwargs['threshold'] = self.threshold_

        return plot_anomaly_score(self.anomaly_score(X), **kwargs)

    def plot_roc_curve(self, X, y, **kwargs):
        """Plot the Receiver Operating Characteristic (ROC) curve.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        y : array-like of shape (n_samples,)
            Labels.

        ax : matplotlib Axes, default None
            Target axes instance.

        figsize: tuple, default None
            Tuple denoting figure size of the plot.

        filename : str, default None
            If provided, save the current figure.

        grid : boolean, default True
            If True, turn the axes grids on.

        label : str, default None
            Legend label.

        title : string, default None
            Axes title. To disable, pass None.

        xlabel : string, default 'FPR'
            X axis title label. To disable, pass None.

        ylabel : string, default 'TPR'
            Y axis title label. To disable, pass None.

        **kwargs : dict
            Other keywords passed to `ax.plot`.

        Returns
        -------
        ax : matplotlib Axes
            Axes on which the plot was drawn.
        """

        X, y = check_X_y(X, y)

        return plot_roc_curve(y, self.anomaly_score(X), **kwargs)
<commit_msg>add OutlierMixin<commit_after>from abc import abstractmethod, ABC

import numpy as np
from sklearn.base import BaseEstimator
from sklearn.utils.validation import check_is_fitted

from .visualization import plot_anomaly_score, plot_roc_curve

__all__ = ['is_outlier_detector', 'BaseOutlierDetector', 'OutlierMixin']

def is_outlier_detector(estimator):
    """Return True if the given estimator is (probably) an outlier detector.

    Parameters
    ----------
    estimator : object
        Estimator object to test.

    Returns
    -------
    out : bool
        True if estimator is an outlier detector and False otherwise.
    """

    return getattr(estimator, '_estimator_type', None) == 'outlier_detector'

class OutlierMixin:
    """Mixin class for all outlier detectors."""

    _estimator_type = 'outlier_detector'

    def fit_predict(self, X, y=None):
        """Fit the model according to the given training data and predict if a
        particular sample is an outlier or not.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training Data.

        y : ignored

        Returns
        -------
        y_pred : array-like of shape (n_samples,)
            Return 1 for inliers and -1 for outliers.
        """

        return self.fit(X).predict(X)

class BaseOutlierDetector(BaseEstimator, OutlierMixin, ABC):
    """Base class for all outlier detectors in kenchi."""

    # TODO: Update anomaly_score method so that the normalized score can be
    # computed

    @abstractmethod
    def __init__(self, contamination=0.01, verbose=False):
        self.contamination = contamination
        self.verbose       = verbose

    @abstractmethod
    def fit(self, X, y=None):
        """Fit the model according to the given training data.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.

        y : ignored

        Returns
        -------
        self : object
            Return self.
        """

    @abstractmethod
    def anomaly_score(self, X):
        """Compute the anomaly score for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        Returns
        -------
        anomaly_score : array-like of shape (n_samples,)
            Anomaly score for each sample.
        """

    def _check_params(self):
        """Check validity of parameters and raise ValueError if not valid."""

        if not 0. <= self.contamination <= 0.5:
            raise ValueError(
                f'contamination must be between 0.0 and 0.5 inclusive '
                f'but was {self.contamination}'
            )

    def _get_threshold(self):
        """Define the threshold according to the given training data."""

        return np.percentile(
            self.anomaly_score_, 100. * (1. - self.contamination)
        )

    def predict(self, X, threshold=None):
        """Predict if a particular sample is an outlier or not.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        threshold : float, default None
            User-provided threshold.

        Returns
        -------
        y_pred : array-like of shape (n_samples,)
            Return 1 for inliers and -1 for outliers.
        """

        check_is_fitted(self, 'threshold_')

        if threshold is None:
            threshold = self.threshold_

        return np.where(self.anomaly_score(X) <= threshold, 1, -1)

    def plot_anomaly_score(self, X, **kwargs):
        """Plot the anomaly score for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        ax : matplotlib Axes, default None
            Target axes instance.

        bins : int, str or array-like, default 'auto'
            Number of hist bins.

        figsize : tuple, default None
            Tuple denoting figure size of the plot.

        filename : str, default None
            If provided, save the current figure.

        grid : boolean, default True
            If True, turn the axes grids on.

        hist : bool, default True
            If True, plot a histogram of anomaly scores.

        title : string, default None
            Axes title. To disable, pass None.

        xlabel : string, default 'Samples'
            X axis title label. To disable, pass None.

        xlim : tuple, default None
            Tuple passed to `ax.xlim`.

        ylabel : string, default 'Anomaly score'
            Y axis title label. To disable, pass None.

        ylim : tuple, default None
            Tuple passed to `ax.ylim`.

        **kwargs : dict
            Other keywords passed to `ax.plot`.

        Returns
        -------
        ax : matplotlib Axes
            Axes on which the plot was drawn.
        """

        kwargs['threshold'] = self.threshold_

        return plot_anomaly_score(self.anomaly_score(X), **kwargs)

    def plot_roc_curve(self, X, y, **kwargs):
        """Plot the Receiver Operating Characteristic (ROC) curve.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data.

        y : array-like of shape (n_samples,)
            Labels.

        ax : matplotlib Axes, default None
            Target axes instance.

        figsize: tuple, default None
            Tuple denoting figure size of the plot.

        filename : str, default None
            If provided, save the current figure.

        grid : boolean, default True
            If True, turn the axes grids on.

        label : str, default None
            Legend label.

        title : string, default None
            Axes title. To disable, pass None.

        xlabel : string, default 'FPR'
            X axis title label. To disable, pass None.

        ylabel : string, default 'TPR'
            Y axis title label. To disable, pass None.

        **kwargs : dict
            Other keywords passed to `ax.plot`.

        Returns
        -------
        ax : matplotlib Axes
            Axes on which the plot was drawn.
        """

        return plot_roc_curve(y, self.anomaly_score(X), **kwargs)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# Copyright (c) 2014-2016 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.

#
# Test merkleblock fetch/validation
#

from test_framework.test_framework import NavCoinTestFramework
from test_framework.util import *

class MerkleBlockTest(NavCoinTestFramework):

    def __init__(self):
        super().__init__()
        self.setup_clean_chain = True
        self.num_nodes = 4

    def setup_network(self):
        self.nodes = []
        # Nodes 0/1 are "wallet" nodes
        self.nodes.append(start_node(0, self.options.tmpdir, ["-debug"]))
        self.nodes.append(start_node(1, self.options.tmpdir, ["-debug"]))
        # Nodes 2/3 are used for testing
        self.nodes.append(start_node(2, self.options.tmpdir, ["-debug"]))
        self.nodes.append(start_node(3, self.options.tmpdir, ["-debug", "-txindex"]))
        connect_nodes(self.nodes[0], 1)
        connect_nodes(self.nodes[0], 2)
        connect_nodes(self.nodes[0], 3)

        self.is_network_split = False
        self.sync_all()

    def run_test(self):
        print("Mining blocks...")
        slow_gen(self.nodes[0], 105)
        self.sync_all()
        chain_height = self.nodes[1].getblockcount()
        assert_equal(chain_height, 105)
        assert_equal(self.nodes[1].getbalance(), 0)
        assert_equal(self.nodes[2].getbalance(), 0)

        node0utxos = self.nodes[0].listunspent(1, 104) // We generated 105 blocks so we are excluding the first block which contains 59 mill~ NAV
        tx1 = self.nodes[0].createrawtransaction([node0utxos.pop()], {self.nodes[1].getnewaddress(): 49.99})
        txid1 = self.nodes[0].sendrawtransaction(self.nodes[0].signrawtransaction(tx1)["hex"])
        tx2 = self.nodes[0].createrawtransaction([node0utxos.pop()], {self.nodes[1].getnewaddress(): 49.99})
        txid2 = self.nodes[0].sendrawtransaction(self.nodes[0].signrawtransaction(tx2)["hex"])
        assert_raises(JSONRPCException, self.nodes[0].gettxoutproof, [txid1])

        slow_gen(self.nodes[0], 1)
        blockhash = self.nodes[0].getblockhash(chain_height + 1)
        self.sync_all()

        txlist = []
        blocktxn = self.nodes[0].getblock(blockhash, True)["tx"]
        txlist.append(blocktxn[1])
        txlist.append(blocktxn[2])

        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1])), [txid1])
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1, txid2])), txlist)
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1, txid2], blockhash)), txlist)

        txin_spent = self.nodes[1].listunspent(1).pop()
        tx3 = self.nodes[1].createrawtransaction([txin_spent], {self.nodes[0].getnewaddress(): 49.98})
        self.nodes[0].sendrawtransaction(self.nodes[1].signrawtransaction(tx3)["hex"])
        slow_gen(self.nodes[0], 1)
        self.sync_all()
        txid_spent = txin_spent["txid"]
        txid_unspent = txid1 if txin_spent["txid"] != txid1 else txid2

        # We can't find the block from a fully-spent tx
        assert_raises(JSONRPCException, self.nodes[2].gettxoutproof, [txid_spent])
        # ...but we can if we specify the block
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid_spent], blockhash)), [txid_spent])
        # ...or if the first tx is not fully-spent
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid_unspent])), [txid_unspent])
        try:
            assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1, txid2])), txlist)
        except JSONRPCException:
            assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid2, txid1])), txlist)
        # ...or if we have a -txindex
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[3].gettxoutproof([txid_spent])), [txid_spent])

if __name__ == '__main__':
    MerkleBlockTest().main()

<commit_msg>fix test comment<commit_after>#!/usr/bin/env python3
# Copyright (c) 2014-2016 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.

#
# Test merkleblock fetch/validation
#

from test_framework.test_framework import NavCoinTestFramework
from test_framework.util import *

class MerkleBlockTest(NavCoinTestFramework):

    def __init__(self):
        super().__init__()
        self.setup_clean_chain = True
        self.num_nodes = 4

    def setup_network(self):
        self.nodes = []
        # Nodes 0/1 are "wallet" nodes
        self.nodes.append(start_node(0, self.options.tmpdir, ["-debug"]))
        self.nodes.append(start_node(1, self.options.tmpdir, ["-debug"]))
        # Nodes 2/3 are used for testing
        self.nodes.append(start_node(2, self.options.tmpdir, ["-debug"]))
        self.nodes.append(start_node(3, self.options.tmpdir, ["-debug", "-txindex"]))
        connect_nodes(self.nodes[0], 1)
        connect_nodes(self.nodes[0], 2)
        connect_nodes(self.nodes[0], 3)

        self.is_network_split = False
        self.sync_all()

    def run_test(self):
        print("Mining blocks...")
        slow_gen(self.nodes[0], 105)
        self.sync_all()
        chain_height = self.nodes[1].getblockcount()
        assert_equal(chain_height, 105)
        assert_equal(self.nodes[1].getbalance(), 0)
        assert_equal(self.nodes[2].getbalance(), 0)

        node0utxos = self.nodes[0].listunspent(1, 104)
        # We generated 105 blocks so we are excluding the first block which contains 59 mill~ NAV

        tx1 = self.nodes[0].createrawtransaction([node0utxos.pop()], {self.nodes[1].getnewaddress(): 49.99})
        txid1 = self.nodes[0].sendrawtransaction(self.nodes[0].signrawtransaction(tx1)["hex"])
        tx2 = self.nodes[0].createrawtransaction([node0utxos.pop()], {self.nodes[1].getnewaddress(): 49.99})
        txid2 = self.nodes[0].sendrawtransaction(self.nodes[0].signrawtransaction(tx2)["hex"])
        assert_raises(JSONRPCException, self.nodes[0].gettxoutproof, [txid1])

        slow_gen(self.nodes[0], 1)
        blockhash = self.nodes[0].getblockhash(chain_height + 1)
        self.sync_all()

        txlist = []
        blocktxn = self.nodes[0].getblock(blockhash, True)["tx"]
        txlist.append(blocktxn[1])
        txlist.append(blocktxn[2])

        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1])), [txid1])
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1, txid2])), txlist)
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1, txid2], blockhash)), txlist)

        txin_spent = self.nodes[1].listunspent(1).pop()
        tx3 = self.nodes[1].createrawtransaction([txin_spent], {self.nodes[0].getnewaddress(): 49.98})
        self.nodes[0].sendrawtransaction(self.nodes[1].signrawtransaction(tx3)["hex"])
        slow_gen(self.nodes[0], 1)
        self.sync_all()
        txid_spent = txin_spent["txid"]
        txid_unspent = txid1 if txin_spent["txid"] != txid1 else txid2

        # We can't find the block from a fully-spent tx
        assert_raises(JSONRPCException, self.nodes[2].gettxoutproof, [txid_spent])
        # ...but we can if we specify the block
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid_spent], blockhash)), [txid_spent])
        # ...or if the first tx is not fully-spent
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid_unspent])), [txid_unspent])
        try:
            assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid1, txid2])), txlist)
        except JSONRPCException:
            assert_equal(self.nodes[2].verifytxoutproof(self.nodes[2].gettxoutproof([txid2, txid1])), txlist)
        # ...or if we have a -txindex
        assert_equal(self.nodes[2].verifytxoutproof(self.nodes[3].gettxoutproof([txid_spent])), [txid_spent])

if __name__ == '__main__':
    MerkleBlockTest().main()

<|endoftext|>


<commit_before>import time

import requests
from weavelib.http import AppHTTPServer
from weavelib.messaging import Receiver
from weavelib.rpc import RPCServer, ServerAPI, RPCClient
from weavelib.services import BaseService, BackgroundThreadServiceStart

from weaveserver.services.core import CoreService
from weaveserver.services.http import HTTPService

from weaveserver.core.logger import configure_logging

configure_logging()

AUTH = {
    "auth1": {
        "type": "SYSTEM",
        "package": "core",
        "appid": "auth1"
    },
    "auth2": {
        "appid": "auth2",
        "type": "SYSTEM",
        "package": "http"
    },
    "auth3": {
        "appid": "auth3",
        "package": "dummy"
    }
}

class ThreadedHTTPService(BackgroundThreadServiceStart, HTTPService):
    pass

class DummyService(BaseService):
    def __init__(self, token):
        super(DummyService, self).__init__(token)
        self.rpc_server = RPCServer("name", "desc", [
            ServerAPI("api1", "desc2", [], self.api1),
        ], self)
        self.http = AppHTTPServer(self)

    def api1(self):
        return "OK"

    def on_service_start(self):
        self.rpc_server.start()
        self.http.start()
        self.relative_url = self.http.register_folder("test_dir")

    def on_service_stop(self):
        self.http.stop()
        self.rpc_server.stop()

class TestHTTPService(object):
    def setup_class(cls):
        cls.core_service = CoreService("auth1", {"core_config": {}})
        cls.core_service.service_start()
        cls.core_service.wait_for_start(30)

        cls.core_service.message_server.register_application(AUTH["auth2"])
        cls.core_service.message_server.register_application(AUTH["auth3"])

        # Wait till it starts.
        receiver = Receiver("/_system/root_rpc/request")
        while True:
            try:
                receiver.start()
                break
            except:
                time.sleep(1)

        cls.service = ThreadedHTTPService("auth2", None)
        cls.service.service_start()
        cls.service.wait_for_start(30)

    def teardown_class(cls):
        cls.service.service_stop()
        cls.core_service.service_stop()

    def setup_method(self):
        self.dummy_service = DummyService("auth3")
        self.dummy_service.service_start()

    def teardown_method(self):
        self.dummy_service.service_stop()

    def test_rpc(self):
        rpc = RPCClient(self.dummy_service.rpc_server.info_message)
        rpc.start()
        assert "OK" == rpc["api1"](_block=True)
        rpc.stop()

    def disabled_test_http_simple_request(self):
        base_url = "http://localhost:5000" + self.dummy_service.relative_url

        url = base_url + "/index.json"
        resp = requests.get(url, headers={"host": "http://localhost:1234/"})
        assert resp.json()["hello"] == "world"

        url = base_url + "/test.csv"
        resp = requests.get(url)
        assert resp.text == "a,b,c\n"
        assert resp.headers["Content-Type"] == "text/csv; charset=UTF-8"

    def test_rpc_info(self):
        info = self.dummy_service.rpc_client["rpc_info"]("dummy", "name",
                                                         _block=True)
        actual_info = self.dummy_service.rpc_server.info_message
        assert info["request_queue"] == actual_info["request_queue"]
        assert info["response_queue"] == actual_info["response_queue"]
<commit_msg>Fix test case.<commit_after>import hashlib
import os
import time

import requests
from weavelib.http import AppHTTPServer
from weavelib.messaging import Receiver
from weavelib.rpc import RPCServer, ServerAPI, RPCClient
from weavelib.services import BaseService, BackgroundThreadServiceStart

from weaveserver.services.core import CoreService
from weaveserver.services.http import HTTPService

from weaveserver.core.logger import configure_logging

configure_logging()

AUTH = {
    "auth1": {
        "type": "SYSTEM",
        "package": "core",
        "appid": "auth1"
    },
    "auth2": {
        "appid": "auth2",
        "type": "SYSTEM",
        "package": "http"
    },
    "auth3": {
        "appid": "auth3",
        "package": "dummy"
    }
}

class ThreadedHTTPService(BackgroundThreadServiceStart, HTTPService):
    pass

class DummyService(BaseService):
    def __init__(self, token):
        super(DummyService, self).__init__(token)
        self.rpc_server = RPCServer("name", "desc", [
            ServerAPI("api1", "desc2", [], self.api1),
        ], self)
        self.http = AppHTTPServer(self)

    def api1(self):
        return "OK"

    def on_service_start(self):
        self.rpc_server.start()
        self.http.start()
        self.relative_url = self.http.register_folder("test_dir")

    def on_service_stop(self):
        self.http.stop()
        self.rpc_server.stop()

class TestHTTPService(object):
    def setup_class(cls):
        cls.core_service = CoreService("auth1",
                                       {"core_config": {}, "apps": AUTH})
        cls.core_service.service_start()
        cls.core_service.wait_for_start(30)

        cls.core_service.message_server.register_application(AUTH["auth2"])
        cls.core_service.message_server.register_application(AUTH["auth3"])

        # Wait till it starts.
        receiver = Receiver("/_system/root_rpc/request")
        while True:
            try:
                receiver.start()
                break
            except:
                time.sleep(1)

        cls.service = ThreadedHTTPService("auth2", None)
        cls.service.service_start()
        cls.service.wait_for_start(30)

    def teardown_class(cls):
        cls.service.service_stop()
        cls.core_service.service_stop()

    def setup_method(self):
        self.dummy_service = DummyService("auth3")
        self.dummy_service.service_start()

    def teardown_method(self):
        self.dummy_service.service_stop()

    def test_rpc(self):
        rpc = RPCClient(self.dummy_service.rpc_server.info_message)
        rpc.start()
        assert "OK" == rpc["api1"](_block=True)
        rpc.stop()

    def disabled_test_http_simple_request(self):
        base_url = "http://localhost:5000" + self.dummy_service.relative_url

        url = base_url + "/index.json"
        resp = requests.get(url, headers={"host": "http://localhost:1234/"})
        assert resp.json()["hello"] == "world"

        url = base_url + "/test.csv"
        resp = requests.get(url)
        assert resp.text == "a,b,c\n"
        assert resp.headers["Content-Type"] == "text/csv; charset=UTF-8"

    def test_resource_creation(self):
        src = os.path.join(os.path.dirname(__file__), "test_dir")
        dest = os.path.join(self.service.plugin_dir.name, "auth3")
        src_files = [x[1:] for x in os.walk(src)]
        dest_files = [x[1:] for x in os.walk(dest)]

        assert src_files == dest_files
        src_full_paths = [os.path.join(src, x)
                          for folder_info in src_files
                          for x in folder_info[1]]
        dest_full_paths = [os.path.join(dest, x)
                           for folder_info in dest_files
                           for x in folder_info[1]]

        for src_path, dest_path in zip(src_full_paths, dest_full_paths):
            with open(src_path, "rb") as src_inp:
                src_hash = hashlib.md5(src_inp.read()).hexdigest()
            with open(dest_path, "rb") as dest_inp:
                dest_hash = hashlib.md5(dest_inp.read()).hexdigest()

            assert src_hash == dest_hash
<|endoftext|>


<commit_before>"""
This is originally from:

https://github.com/mahendra/kafka-python/blob/zookeeper/kafka/zookeeper.py

It is modified in a few places to work with more recent KafkaClient.

Also, multiprocess is substituted for threading. Since threading
is gevent friendly, where multiprocess is not.
"""
import logging
import threading
import os
import random
import socket
import uuid
import json
from functools import partial
from kafka.client import KafkaClient
from kafka.producer import SimpleProducer, KeyedProducer
from kafka.consumer import SimpleConsumer
from kazoo.client import KazooClient
import time
import sys

if 'gevent' in sys.modules:
    from kazoo.handlers.gevent import SequentialGeventHandler as kazoo_handler
else:
    from kazoo.handlers.threading import SequentialThreadingHandler as kazoo_handler

BROKER_IDS_PATH = 'brokers/ids/'      # Path where kafka stores broker info
PARTITIONER_PATH = 'python/kafka/'    # Path to use for consumer co-ordination
DEFAULT_TIME_BOUNDARY = 5

# Allocation states
ALLOCATION_COMPLETED = -1
ALLOCATION_CHANGING = -2
ALLOCATION_FAILED = -3
ALLOCATION_MISSED = -4
ALLOCATION_INACTIVE = -5

log = logging.getLogger("kafka")
random.seed()

def _get_brokers(zkclient, chroot='/'):
    """
    Get the list of available brokers registered in zookeeper
    """
    brokers = []
    root = os.path.join(chroot, BROKER_IDS_PATH)

    for broker_id in zkclient.get_children(root):
        path = os.path.join(root, broker_id)
        info, _ = zkclient.get(path)
        info = json.loads(info)
        brokers.append((info['host'], info['port']))

    log.debug("List of brokers fetched" + str(brokers))

    random.shuffle(brokers)
    return brokers

def get_client(zkclient, chroot='/'):
    """
    Given a zookeeper client, return a KafkaClient instance for use
    """
    brokers = _get_brokers(zkclient, chroot=chroot)
    brokers = ["%s:%s"%(host, port) for (host, port) in brokers]
    return KafkaClient(brokers)

# TODO: Make this a subclass of Producer later
class ZProducer(object):
    """
    A base Zookeeper producer to be used by other producer classes

    Args
    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    topic - The kafka topic to send messages to
    chroot - The kafka subdirectory to search for brokers
    """
    producer_kls = None

    def __init__(self, hosts, topic, chroot='/', **kwargs):

        if self.producer_kls is None:
            raise NotImplemented("Producer class needs to be mentioned")

        self.zkclient = KazooClient(hosts=hosts)
        self.zkclient.start()

        # Start the producer instance
        self.client = get_client(self.zkclient, chroot=chroot)
        self.producer = self.producer_kls(self.client, topic, **kwargs)

        # Stop Zookeeper
        self.zkclient.stop()
        self.zkclient.close()
        self.zkclient = None

    def stop(self):
        self.producer.stop()
        self.client.close()

class ZSimpleProducer(ZProducer):
    """
    A simple, round-robbin producer. Each message goes to exactly one partition

    Args:
    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    topic - The kafka topic to send messages to
    """
    producer_kls = SimpleProducer

    def send_messages(self, *msg):
        self.producer.send_messages(*msg)

class ZKeyedProducer(ZProducer):
    """
    A producer which distributes messages to partitions based on a
    partitioner function (class) and the key

    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    topic - The kafka topic to send messages to
    partitioner - A partitioner class that will be used to get the partition
        to send the message to. Must be derived from Partitioner
    """
    producer_kls = KeyedProducer

    def send(self, key, msg):
        self.producer.send(key, msg)

class ZSimpleConsumer(object):
    """
    A consumer that uses Zookeeper to co-ordinate and share the partitions
    of a topic with other consumers

    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    group: a name for this consumer, used for offset storage and must be unique
    topic: the topic to consume
    chroot - The kafka subdirectory to search for brokers
    driver_type: The driver type to use for the consumer
    block_init: If True, the init method will block till the allocation is
        completed. If not, it will return immediately and user can invoke
        consumer.status() to check the status. Default True.
    time_boundary: The time interval, in seconds, to wait out before deciding
        on consumer changes in zookeeper. A higher value will ensure that a
        consumer restart will not cause two re-balances.
        (Default 10s)
    ignore_non_allocation: If set to True, the consumer will ignore the
        case where no partitions were allocated to it.
        This can be used to keep consumers in stand-by. They will take over
        when another consumer fails. (Default False)

    auto_commit: default True. Whether or not to auto commit the offsets
    auto_commit_every_n: default 100. How many messages to consume
                         before a commit
    auto_commit_every_t: default 5000. How much time (in milliseconds) to
                         wait before commit

    Auto commit details:
    If both auto_commit_every_n and auto_commit_every_t are set, they will
    reset one another when one is triggered. These triggers simply call the
    commit method on this class. A manual call to commit will also reset
    these triggers

    Partition allocation details
    * When the consumer is initialized, it blocks till it gets an allocation
    * If ignore_non_allocation is False, the consumer will throw an error
      in init or during other operations
    * During re-balancing of partitions, the consumer will not return any
      messages (iteration or get_messages)
    * After re-balancing, if the consumer does not get any partitions,
      ignore_non_allocation will control it's behaviour
    """
    def __init__(self,
                 hosts,
                 group,
                 topic,
                 chroot='/',
                 block_init=True,
                 time_boundary=DEFAULT_TIME_BOUNDARY,
                 ignore_non_allocation=False,
                 **kwargs):

        # User is not allowed to specify partitions
        if 'partitions' in kwargs:
            raise ValueError("Partitions cannot be specified")

        self.ignore_non_allocation = ignore_non_allocation
        self.time_boundary = time_boundary

        zkclient = KazooClient(hosts, handler=kazoo_handler())
        zkclient.start()

        self.client = get_client(zkclient, chroot=chroot)
        self.client.load_metadata_for_topics(topic)
        partitions = set(self.client.topic_partitions[topic])

        # create consumer id
        hostname = socket.gethostname()
        self.identifier = "%s-%s-%s" % (topic, group, hostname)
        log.debug("Consumer id set to: %s" % self.identifier)

        path = os.path.join(chroot, PARTITIONER_PATH, topic, group)
        log.debug("Using path %s for co-ordination" % path)

        # Create a function which can be used for creating consumers
        self.consumer = []
        self.consumer_fact = partial(SimpleConsumer,
                                     self.client,
                                     group,
                                     topic,
                                     **kwargs)

        # Keep monitoring for changes

        # Design:
        # * We will have a worker which will keep monitoring for rebalance
        # * The worker and main consumer will share data via shared memory
        #   protected by a lock
        # * If the worker gets new allocations, it will SET an Event()
        # * The main consumer will check this event to change itself
        # * Main consumer will SET another Event() to indicate worker to exit

        # This event will notify the worker to exit
        self.exit = threading.Event()

        # Used by the worker to indicate that allocation has changed
        self.changed = threading.Event()

        # The shared memory and lock used for sharing allocation info
        self.lock = threading.Lock()
        self.allocated = [0] * len(partitions)

        # Initialize the array
        self._set_partitions(self.allocated, [], ALLOCATION_CHANGING)
        self.consumer_state = ALLOCATION_CHANGING

        # Start the worker
        self.proc = threading.Thread(target=self._check_and_allocate,
                                     args=(hosts,
                                           path,
                                           partitions,
                                           self.allocated))
        self.proc.daemon = True
        self.proc.start()

        # Stop the Zookeeper client (worker will create one itself)
        zkclient.stop()
        zkclient.close()

        # Do the setup once and block till we get an allocation
        self._set_consumer(block=block_init, timeout=None)

    def __repr__(self):
        """
        Give a string representation of the consumer
        """
        partitions = filter(lambda x: x>=0, self.allocated)
        message = ','.join([str(i) for i in partitions])

        if not message:
            message = self.status()

        return u'ZSimpleConsumer<%s>' % message

    def status(self):
        """
        Returns the status of the consumer
        """
        self._set_consumer(block=False)

        if self.consumer_state == ALLOCATION_COMPLETED:
            return 'ALLOCATED'
        elif self.consumer_state == ALLOCATION_CHANGING:
            return 'ALLOCATING'
        elif self.consumer_state == ALLOCATION_FAILED:
            return 'FAILED'
        elif self.consumer_state == ALLOCATION_MISSED:
            return 'MISSED'
        elif self.consumer_state == ALLOCATION_INACTIVE:
            return 'INACTIVE'

    def _set_partitions(self, array, partitions, filler):
        """
        Update partition info in the shared memory array
        """
        i = 0
        for partition in partitions:
            array[i] = partition
            i += 1

        while i < len(array):
            array[i] = filler
            i += 1

    def _set_consumer(self, block=False, timeout=None):
        """
        Check if a new consumer has to be created because of a re-balance
        """
        if not block:
            timeout = 0

        if not self.changed.wait(timeout=timeout):
            return

        # There is a change. Get our new partitions
        with self.lock:
            partitions = [p for p in self.allocated]
            self.changed.clear()

        # If we have a consumer clean it up
        if self.consumer:
            self.consumer.stop()

        self.consumer = None
        self.consumer_state = partitions[0]

        if self.consumer_state == ALLOCATION_MISSED:
            # Check if we must change the consumer
            if not self.ignore_non_allocation:
                raise RuntimeError("Did not get any partition allocation")
            else:
                log.info("No partitions allocated. Ignoring")
                self.consumer = []

        elif self.consumer_state == ALLOCATION_FAILED:
            # Allocation has failed. Nothing we can do about it
            raise RuntimeError("Error in partition allocation")

        elif self.consumer_state == ALLOCATION_CHANGING:
            # Allocation is changing because of consumer changes
            self.consumer = []
            log.info("Partitions are being reassigned")
            return

        elif self.consumer_state == ALLOCATION_INACTIVE:
            log.info("Consumer is inactive")
        else:
            # Create a new consumer
            partitions = filter(lambda x: x >= 0, partitions)
            self.consumer = self.consumer_fact(partitions=partitions)
            self.consumer_state = ALLOCATION_COMPLETED
            log.info("Reinitialized consumer with partitions: %s" % partitions)

    def _check_and_allocate(self, hosts, path, partitions, array):
        """
        Checks if a new allocation is needed for the partitions.
        If so, co-ordinates with Zookeeper to get a set of partitions
        allocated for the consumer
        """

        old = None

        # Start zookeeper connection again
        zkclient = KazooClient(hosts, handler=kazoo_handler())
        zkclient.start()

        identifier = '%s-%d-%s' % (self.identifier,
                                   os.getpid(),
                                   uuid.uuid4().hex)

        # Set up the partitioner
        partitioner = zkclient.SetPartitioner(path=path, set=partitions,
                                              identifier=identifier,
                                              time_boundary=self.time_boundary)

        # Once allocation is done, sleep for some time between each checks
        sleep_time = self.time_boundary / 2.0

        # Keep running the allocation logic till we are asked to exit
        while not self.exit.is_set():
            if partitioner.acquired:
                # A new set of partitions has been acquired
                new = list(partitioner)

                # If there is a change, notify for a consumer change
                if new != old:
                    log.info("Acquired partitions: %s" % str(new))
                    old = new
                    with self.lock:
                        self._set_partitions(array, new, ALLOCATION_MISSED)
                        self.changed.set()

                # Wait for a while before checking again. In the meantime
                # wake up if the user calls for exit
                self.exit.wait(sleep_time)

            elif partitioner.release:
                # We have been asked to release the partitions
                log.info("Releasing partitions for reallocation")
                old = None
                with self.lock:
                    self._set_partitions(array, [], ALLOCATION_CHANGING)
                    self.changed.set()

                partitioner.release_set()

            elif partitioner.failed:
                # Partition allocation failed
                old = []
                with self.lock:
                    self._set_partitions(array, old, ALLOCATION_FAILED)
                    self.changed.set()
                break

            elif partitioner.allocating:
                # We have to wait till the partition is allocated
                log.info("Waiting for partition allocation")
                partitioner.wait_for_acquire(timeout=1)

        # Clean up
        partitioner.finish()

        with self.lock:
            self._set_partitions(array, [], ALLOCATION_INACTIVE)
            self.changed.set()

        zkclient.stop()
        zkclient.close()

    def __iter__(self):
        """
        Iterate through data available in partitions allocated to this
        instance
        """
        self._set_consumer(block=False)

        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")

        for msg in self.consumer:
            yield msg
            self._set_consumer(block=False)

    def get_messages(self, count=1, block=True, timeout=0.1):
        """
        Fetch the specified number of messages

        count: Indicates the maximum number of messages to be fetched
        block: If True, the API will block till some messages are fetched.
        timeout: If None, and block=True, the API will block infinitely.
                 If >0, API will block for specified time (in seconds)
        """
        self._set_consumer(block=False, timeout=timeout)

        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")
        elif not self.consumer:
            # This is needed in cases where gevent is used with
            # a thread that does not have any calls that would yield.
            # If we do not sleep here a greenlet could spin indefinitely.
            time.sleep(0)
            return []

        return self.consumer.get_messages(count, block, timeout)

    def stop(self):
        self.exit.set()
        self.proc.join()
        self._set_consumer(block=True)
        self.client.close()

    def commit(self):
        if self.consumer:
            self.consumer.commit()
        self._set_consumer(block=False)

    def seek(self, *args, **kwargs):
        self._set_consumer()

        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")
        elif not self.consumer:
            raise RuntimeError("Waiting for partition allocation")

        return self.consumer.seek(*args, **kwargs)

    def pending(self):
        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")
        elif not self.consumer:
            # We are in a transition/suspended state
            return 0

        return self.consumer.pending()
<commit_msg>major refactor of zkconsumer<commit_after>"""
This is originally from:

https://github.com/mahendra/kafka-python/blob/zookeeper/kafka/zookeeper.py

It is modified in a few places to work with more recent KafkaClient.

Also, multiprocess is substituted for threading. Since threading
is gevent friendly, where multiprocess is not.
"""
import logging
import threading
import os
import random
import socket
import uuid
import json
from functools import partial
from kafka.client import KafkaClient
from kafka.producer import SimpleProducer, KeyedProducer
from kafka.consumer import SimpleConsumer
from kazoo.client import KazooClient
import time
import sys

if 'gevent' in sys.modules:
    from kazoo.handlers.gevent import SequentialGeventHandler as kazoo_handler
else:
    from kazoo.handlers.threading import SequentialThreadingHandler as kazoo_handler

BROKER_IDS_PATH = 'brokers/ids/'      # Path where kafka stores broker info
PARTITIONER_PATH = 'python/kafka/'    # Path to use for consumer co-ordination
DEFAULT_TIME_BOUNDARY = 10

# Allocation states
ALLOCATION_COMPLETED = -1
ALLOCATION_CHANGING = -2
ALLOCATION_FAILED = -3
ALLOCATION_MISSED = -4
ALLOCATION_INACTIVE = -5

log = logging.getLogger("kafka")
random.seed()

def _get_brokers(zkclient, chroot='/'):
    """
    Get the list of available brokers registered in zookeeper
    """
    brokers = []
    root = os.path.join(chroot, BROKER_IDS_PATH)

    for broker_id in zkclient.get_children(root):
        path = os.path.join(root, broker_id)
        info, _ = zkclient.get(path)
        info = json.loads(info)
        brokers.append((info['host'], info['port']))

    log.debug("List of brokers fetched" + str(brokers))

    random.shuffle(brokers)
    return brokers

def get_client(zkclient, chroot='/'):
    """
    Given a zookeeper client, return a KafkaClient instance for use
    """
    brokers = _get_brokers(zkclient, chroot=chroot)
    brokers = ["%s:%s"%(host, port) for (host, port) in brokers]
    return KafkaClient(brokers)

# TODO: Make this a subclass of Producer later
class ZProducer(object):
    """
    A base Zookeeper producer to be used by other producer classes

    Args
    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    topic - The kafka topic to send messages to
    chroot - The kafka subdirectory to search for brokers
    """
    producer_kls = None

    def __init__(self, hosts, topic, chroot='/', **kwargs):

        if self.producer_kls is None:
            raise NotImplemented("Producer class needs to be mentioned")

        self.zkclient = KazooClient(hosts=hosts)
        self.zkclient.start()

        # Start the producer instance
        self.client = get_client(self.zkclient, chroot=chroot)
        self.producer = self.producer_kls(self.client, topic, **kwargs)

        # Stop Zookeeper
        self.zkclient.stop()
        self.zkclient.close()
        self.zkclient = None

    def stop(self):
        self.producer.stop()
        self.client.close()

class ZSimpleProducer(ZProducer):
    """
    A simple, round-robbin producer. Each message goes to exactly one partition

    Args:
    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    topic - The kafka topic to send messages to
    """
    producer_kls = SimpleProducer

    def send_messages(self, *msg):
        self.producer.send_messages(*msg)

class ZKeyedProducer(ZProducer):
    """
    A producer which distributes messages to partitions based on a
    partitioner function (class) and the key

    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    topic - The kafka topic to send messages to
    partitioner - A partitioner class that will be used to get the partition
        to send the message to. Must be derived from Partitioner
    """
    producer_kls = KeyedProducer

    def send(self, key, msg):
        self.producer.send(key, msg)

class ZSimpleConsumer(object):
    """
    A consumer that uses Zookeeper to co-ordinate and share the partitions
    of a topic with other consumers

    hosts: Comma-separated list of hosts to connect to
           (e.g. 127.0.0.1:2181,127.0.0.1:2182)
    group: a name for this consumer, used for offset storage and must be unique
    topic: the topic to consume
    chroot - The kafka subdirectory to search for brokers
    driver_type: The driver type to use for the consumer
    block_init: If True, the init method will block till the allocation is
        completed. If not, it will return immediately and user can invoke
        consumer.status() to check the status. Default True.
    time_boundary: The time interval, in seconds, to wait out before deciding
        on consumer changes in zookeeper. A higher value will ensure that a
        consumer restart will not cause two re-balances.
        (Default 10s)
    ignore_non_allocation: If set to True, the consumer will ignore the
        case where no partitions were allocated to it.
        This can be used to keep consumers in stand-by. They will take over
        when another consumer fails. (Default False)

    auto_commit: default True. Whether or not to auto commit the offsets
    auto_commit_every_n: default 100. How many messages to consume
                         before a commit
    auto_commit_every_t: default 5000. How much time (in milliseconds) to
                         wait before commit

    Auto commit details:
    If both auto_commit_every_n and auto_commit_every_t are set, they will
    reset one another when one is triggered. These triggers simply call the
    commit method on this class. A manual call to commit will also reset
    these triggers

    Partition allocation details
    * When the consumer is initialized, it blocks till it gets an allocation
    * If ignore_non_allocation is False, the consumer will throw an error
      in init or during other operations
    * During re-balancing of partitions, the consumer will not return any
      messages (iteration or get_messages)
    * After re-balancing, if the consumer does not get any partitions,
      ignore_non_allocation will control it's behaviour
    """
    def __init__(self,
                 hosts,
                 group,
                 topic,
                 chroot='/',
                 block_init=True,
                 time_boundary=DEFAULT_TIME_BOUNDARY,
                 ignore_non_allocation=False,
                 **kwargs):

        # User is not allowed to specify partitions
        if 'partitions' in kwargs:
            raise ValueError("Partitions cannot be specified")

        self.ignore_non_allocation = ignore_non_allocation
        self.time_boundary = time_boundary

        self.zkclient = KazooClient(hosts, handler=kazoo_handler())
        self.zkclient.start()

        self.client = get_client(self.zkclient, chroot=chroot)
        self.client.load_metadata_for_topics(topic)
        self.partitions = set(self.client.topic_partitions[topic])

        #self.allocated = [ALLOCATION_CHANGING] * len(partitions)

        # create consumer id
        hostname = socket.gethostname()
        self.identifier = "%s-%s-%s" % (topic, group, hostname)
        log.debug("Consumer id set to: %s" % self.identifier)

        self.path = os.path.join(chroot, PARTITIONER_PATH, topic, group)
        log.debug("Using path %s for co-ordination" % self.path)

        # Create a function which can be used for creating consumers
        self.consumer = []
        self.consumer_fact = partial(SimpleConsumer,
                                     self.client,
                                     group,
                                     topic,
                                     **kwargs)

        # Keep monitoring for changes

        # Design:
        # * We will have a worker which will keep monitoring for rebalance
        # * The worker and main consumer will share data via shared memory
        #   protected by a lock
        # * If the worker gets new allocations, it will SET an Event()
        # * The main consumer will check this event to change itself
        # * Main consumer will SET another Event() to indicate worker to exit

        # This event will notify the worker to exit
        self.exit = threading.Event()

        # Used by the worker to indicate that allocation has changed
        self.changed = threading.Event()

        # The shared memory and lock used for sharing allocation info
        self.lock = threading.Lock()

        # Initialize the array
        #self._set_partitions(self.allocated, [], ALLOCATION_CHANGING)
        self.consumer_state = ALLOCATION_CHANGING

        self.identifier = '%s-%d-%s' % (self.identifier,
                                   os.getpid(),
                                   uuid.uuid4().hex)
        # Start the worker
        self.proc = threading.Thread(target=self._check_and_allocate)

        self.proc.daemon = True
        self.proc.start()

        # Stop the Zookeeper client (worker will create one itself)
        # zkclient.stop()
        # zkclient.close()

        # Do the setup once and block till we get an allocation
        #self._set_consumer(block=block_init, timeout=None)

    def __repr__(self):
        """
        Give a string representation of the consumer
        """
        partitions = filter(lambda x: x>=0, self.allocated)
        message = ','.join([str(i) for i in partitions])

        if not message:
            message = self.status()

        return u'ZSimpleConsumer<%s>' % message

    def status(self):
        """
        Returns the status of the consumer
        """
        self._set_consumer(block=False)

        if self.consumer_state == ALLOCATION_COMPLETED:
            return 'ALLOCATED'
        elif self.consumer_state == ALLOCATION_CHANGING:
            return 'ALLOCATING'
        elif self.consumer_state == ALLOCATION_FAILED:
            return 'FAILED'
        elif self.consumer_state == ALLOCATION_MISSED:
            return 'MISSED'
        elif self.consumer_state == ALLOCATION_INACTIVE:
            return 'INACTIVE'

    # def _set_partitions(self, array, partitions, filler):
    #     """
    #     Update partition info in the shared memory array
    #     """
    #     i = 0
    #     for partition in partitions:
    #         array[i] = partition
    #         i += 1
    #
    #     while i < len(array):
    #         array[i] = filler
    #         i += 1

    def _get_new_partioner(self):
        return self.zkclient.SetPartitioner(path=self.path,
                                                   set=self.partitions,
                                                   identifier=self.identifier,
                                                   time_boundary=self.time_boundary)

    def _check_and_allocate(self):
        """
        Checks if a new allocation is needed for the partitions.
        If so, co-ordinates with Zookeeper to get a set of partitions
        allocated for the consumer
        """

        old = None

        # Set up the partitioner
        partitioner = self._get_new_partioner()

        # Once allocation is done, sleep for some time between each checks
        sleep_time = self.time_boundary / 2.0

        # Keep running the allocation logic till we are asked to exit
        while not self.exit.is_set():

            log.info("ZK Partitoner state: %s"%partitioner.state)

            if partitioner.acquired:
                # A new set of partitions has been acquired

                new = list(partitioner)

                # If there is a change, notify for a consumer change
                if new != old:
                    log.info("Acquired partitions: %s" % str(new))
                    self.consumer = self.consumer_fact(partitions=new)
                    old = new

                # Wait for a while before checking again. In the meantime
                # wake up if the user calls for exit
                self.exit.wait(sleep_time)

            elif partitioner.release:
                # We have been asked to release the partitions

                log.info("Releasing partitions for reallocation")
                #old = None
                self.consumer.stop()
                partitioner.release_set()

            elif partitioner.failed:
                # Partition allocation failed

                # Failure means we need to create a new SetPartitioner:
                # see: http://kazoo.readthedocs.org/en/latest/api/recipe/partitioner.html

                log.error("Partitioner Failed. Creating new partitioner.")

                partitioner = self._get_new_partioner()

            elif partitioner.allocating:
                # We have to wait till the partition is allocated
                log.info("Waiting for partition allocation")
                partitioner.wait_for_acquire(timeout=1)

        # Clean up
        partitioner.finish()

    def __iter__(self):
        """
        Iterate through data available in partitions allocated to this
        instance
        """
        self._set_consumer(block=False)

        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")

        for msg in self.consumer:
            yield msg
            self._set_consumer(block=False)

    def get_messages(self, count=1, block=True, timeout=0.1):
        """
        Fetch the specified number of messages

        count: Indicates the maximum number of messages to be fetched
        block: If True, the API will block till some messages are fetched.
        timeout: If None, and block=True, the API will block infinitely.
                 If >0, API will block for specified time (in seconds)
        """
        #self._set_consumer(block=False, timeout=timeout)

        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")
        elif not self.consumer:
            # This is needed in cases where gevent is used with
            # a thread that does not have any calls that would yield.
            # If we do not sleep here a greenlet could spin indefinitely.
            time.sleep(0)
            return []

        return self.consumer.get_messages(count, block, timeout)

    def stop(self):
        self.exit.set()
        self.proc.join()
        self.zkclient.stop()
        self.zkclient.close()
        #self._set_consumer(block=True)
        self.client.close()

    def commit(self):
        if self.consumer:
            self.consumer.commit()
        #self._set_consumer(block=False)

    def seek(self, *args, **kwargs):
        self._set_consumer()

        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")
        elif not self.consumer:
            raise RuntimeError("Waiting for partition allocation")

        return self.consumer.seek(*args, **kwargs)

    def pending(self):
        if self.consumer is None:
            raise RuntimeError("Error in partition allocation")
        elif not self.consumer:
            # We are in a transition/suspended state
            return 0

        return self.consumer.pending()
<|endoftext|>


<commit_before>#!/usr/bin/python -u

import cStringIO as StringIO
#import Image
import cv2
import picamera
import numpy as np
import time
import datetime
import os
#import ipdb

with picamera.PiCamera() as camera:
	image_data = StringIO.StringIO()
	camera.resolution = (160, 120)
	camera.capture(image_data, format="bmp", use_video_port=True)
	data = np.fromstring(image_data.getvalue(), dtype=np.uint8)
	imarr1 = cv2.imdecode(data, 1)
	image_data.reset()

	counter = 0
	active = 0
	start_time = time.time()
	while 1:
		image_data = StringIO.StringIO()
		camera.capture(image_data, format="bmp", use_video_port=True)
		data = np.fromstring(image_data.getvalue(), dtype=np.uint8)
		imarr2 = cv2.imdecode(data, 1)
		difference = 100*float(((imarr2-imarr1)**2).sum())/(imarr1**2).sum()
		brightness = 100*float(imarr2.sum())/(255*imarr2.shape[0]*imarr2.shape[1]*3)
		print "Difference: {:6.2f}%".format(difference),
		print "Brightness: {:6.2f}%".format(brightness),
		print "{:6.2f} fps".format(float(counter)/(time.time()-start_time)), "\r",
		imarr1 = imarr2
		counter += 1
		if brightness > 50 or difference > 20:
			print ""
			active = 5
			while active > 0:
				date = datetime.datetime.now()
				directory = "{}_{:02d}_{:02d}".format(date.year, date.month, date.day)
				if not os.path.exists(directory):
					os.mkdir(directory)
				filename = "{:02d}{:02d}{:02d}.jpg".format(date.hour, date.minute, date.second)
				print "Active! Writing to {}/{}".format(directory, filename), "\r",
				camera.capture(os.path.join(directory, filename), format="jpeg", use_video_port=True)
				active -= 1
				time.sleep(1)

<commit_msg>sensitiviy increased<commit_after>#!/usr/bin/python -u

import cStringIO as StringIO
#import Image
import cv2
import picamera
import numpy as np
import time
import datetime
import os
#import ipdb

PIC_DIR = "/shared/"

def shoot_pics(number, resolution, delay_sec):
	resolution_init = camera.resolution
	camera.resolution = resolution
	print ""
	while number > 0:
		date = datetime.datetime.now()
		directory = "{}_{:02d}_{:02d}".format(date.year, date.month, date.day)
		if not os.path.exists(os.path.join(PIC_DIR, directory)):
			os.mkdir(os.path.join(PIC_DIR, directory))
		filename = "{:02d}{:02d}{:02d}.jpg".format(date.hour, date.minute, date.second)
		print "Active! Writing to {}/{}/{}".format(PIC_DIR, directory, filename), "\r",
		camera.capture(os.path.join(PIC_DIR, directory, filename), format="jpeg", use_video_port=True)
		number -= 1
		time.sleep(delay_sec)
	camera.resolution = resolution_init

with picamera.PiCamera() as camera:
	image_data = StringIO.StringIO()
	camera.resolution = (160, 120)
	camera.capture(image_data, format="bmp", use_video_port=True)
	data = np.fromstring(image_data.getvalue(), dtype=np.uint8)
	imarr1 = cv2.imdecode(data, 1)
	image_data.reset()

	counter = 0
	active = 0
	start_time = time.time()
	while 1:
		image_data = StringIO.StringIO()
		camera.capture(image_data, format="bmp", use_video_port=True)
		data = np.fromstring(image_data.getvalue(), dtype=np.uint8)
		imarr2 = cv2.imdecode(data, 1)
		difference = 100*float(((imarr2-imarr1)**2).sum())/(imarr1**2).sum()
		brightness = 100*float(imarr2.sum())/(255*imarr2.shape[0]*imarr2.shape[1]*3)
		print "Difference: {:6.2f}%".format(difference),
		print "Brightness: {:6.2f}%".format(brightness),
		print "{:6.2f} fps".format(float(counter)/(time.time()-start_time)), "\r",
		imarr1 = imarr2
		counter += 1

		if difference > 5 and counter > 20:
			shoot_pics(5, (640,480), 1)
			counter = 0
			start_time = time.time()

		if counter == 1000:
			counter = 0

<|endoftext|>


<commit_before>""" Linear Total Correlation Explanation

Code below written by:
Greg Ver Steeg (gregv@isi.edu), 2015.
"""

import numpy as np
from scipy.stats import norm, rankdata  # Used for Gaussianizing data

class Corex(object):
    """
    Linear Total Correlation Explanation

    Conventions
    ----------
    Code follows sklearn naming/style (e.g. fit(X) to train).

    Parameters
    ----------
    n_hidden : int, default = 2
        The number of latent factors to use.

    max_iter : int, default = 100
        The max. number of iterations to reach convergence.

    noise : float default = 0.01
        We imagine some small fundamental measurement noise on Y. The value is arbitrary, but it sets
        the scale of the results, Y.

    verbose : int, optional
        Print verbose outputs.

    seed : integer or numpy.RandomState, optional
        A random number generator instance to define the state of the
        random permutations generator. If an integer is given, it fixes the
        seed. Defaults to the global numpy random number generator.

    Attributes
    ----------

    References
    ----------
    [1] Greg Ver Steeg and Aram Galstyan. "The Information Sieve", 2015.
    [2] ?, Greg Ver Steeg, and Aram Galstyan. "The Information Sieve for Continuous Variables" [In progress]
    [3] Greg Ver Steeg, ?, and Aram Galstyan. "Linear Total Correlation Explanation" [In progress]
    """

    def __init__(self, n_hidden=2, max_iter=10000, noise=1., lam=0., mu=0.01, tol=0.0001, additive=True,
                 gaussianize_marginals=False, verbose=False, seed=None, copy=True, **kwargs):
        self.m = n_hidden  # Number of latent factors to learn
        self.max_iter = max_iter  # Number of iterations to try
        self.noise = noise  # Sets the scale of Y
        self.mu = mu  # Initial step size for lagrange multiplier update
        self.tol = tol  # Threshold for convergence
        self.gaussianize_marginals = gaussianize_marginals
        self.verbose = verbose
        self.copy = copy  # Copy the data before subtracting the mean
        np.random.seed(seed)  # Set seed for deterministic results
        self.kwargs = kwargs

        # Initialize these when we fit on data
        self.n_samples, self.nv = 0, 0  # Number of samples/variables in input data
        self.ws = np.zeros((self.m, self.nv))  # m by nv array of weights
        self.additive = additive  # Whether or not to constrain to additive solutions
        self.lam = lam
        self.moments = {}  # dictionary of moments
        self.mean_x = None  # Mean is subtracted out, save for prediction/inversion
        self.updates = np.ones(self.m)  # Keep track of number of updates for each latent factor
        self.history = np.zeros((max_iter, 3))  # Keep track of values for TC, additivity and objective
        if verbose > 1:
            np.set_printoptions(precision=3, suppress=True, linewidth=160)
            print 'Linear CorEx with %d latent factors' % n_hidden

    @staticmethod
    def calculate_mi(moments):
        return (- 0.5 * np.log(1 - moments["X_i Y_j"] ** 2 / (moments["Y_j^2"] * moments["X_i^2"][:, np.newaxis]))).T

    @property
    def tc(self):
        """Sum_j TC(X;Y_j)"""
        return self.tcs.sum()

    def fit_transform(self, x, **kwargs):
        self.fit(x)
        return self.transform(x, **kwargs)

    def fit(self, x):
        x = np.array(x, dtype=float, copy=self.copy)
        self.n_samples, self.nv = x.shape  # Number of samples, variables in input data
        if self.gaussianize_marginals:
            x = gaussianize(x)
        else:
            self.mean_x = x.mean(axis=0)
            x -= self.mean_x
        var_x = np.einsum('li,li->i', x, x).clip(1e-10) / (self.n_samples - 1)  # Variance of x
        if self.ws.size == 0:  # Randomly initialize weights if not set
            self.ws = np.random.randn(self.m, self.nv) * self.noise ** 2 / np.sqrt(var_x)
        self.lam = self.lam * np.ones(self.nv)  # Initialize lagrange multipliers
        delta = np.inf

        for i_loop in range(self.max_iter):
            old_w = self.ws.copy()
            self._update_moments(x)  # Update moments based on w and samples, x.
            self.history[i_loop] = (self.tc, self.additivity.sum(), self.objective)  # Record convergence stats
            if self.additive and i_loop > 0:  # Update lambda dynamically to get additive solutions
                self.lam = (self.lam - self.mu * self.additivity).clip(0, 1)
            self._update_ws(i_loop)  # Update weights

            if self.additive:  # Adapt mu, the lambda step size, if necessary
                if i_loop % 10 == 9 and i_loop < self.max_iter / 10:
                    if self.additivity.sum() < - 0.01 * self.objective:
                        if (self.history[i_loop, 1] - self.history[i_loop-5, 1]) / 5. <= - 0.1 * self.additivity.sum():
                            self.mu = min(1., self.mu * 1.5)

            delta = np.max(np.sqrt(((old_w - self.ws)**2).sum(axis=1) / (self.ws**2).sum(axis=1)))  # max relative change for each factor
            if self.verbose > 1:
                print 'TC = %0.3f\tadd = %0.3f\tobj = %0.3f\tdelta = %0.5f\t\tmu %0.4f' % \
                      (self.tc, self.additivity.sum(), self.objective, delta, self.mu)
            if delta < self.tol:  # Check for convergence
                if self.verbose:
                    print '%d iterations to tol: %f' % (i_loop, self.tol)
                break
        else:
            if self.verbose:
                print "Warning: Convergence not achieved in %d iterations. Final delta: %f" % (self.max_iter, delta)

        self.history = self.history[:i_loop + 1]
        order = np.argsort(-self.tcs)  # Largest TC components first.
        self.ws = self.ws[order]
        self._update_moments(x)  # Update moments based on w and samples, x.
        return self

    def transform(self, x, details=False):
        """Transform an array of inputs, x, into an array of k latent factors, Y.
            Optionally, you can get the remainder information and/or stop at a specified level."""
        x = np.array(x, dtype=float, copy=self.copy)
        ns, nv = x.shape
        assert self.nv == nv, "Incorrect number of variables in input, %d instead of %d" % (nv, self.nv)
        if self.gaussianize_marginals:
            x = gaussianize(x)  # Should gaussianize wrt to original data...
        else:
            x -= self.mean_x
        if details:
            moments = self._calculate_moments(x)
            return x.dot(self.ws.T), moments
        else:
            return x.dot(self.ws.T)

    def predict(self, y):
        return np.dot(self.moments["X_i Z_j"], y.T).T + self.mean_x

    def _update_moments(self, x):
        """Calculate moments based on the weights and samples. We also calculate and save MI, TC, additivity, and
        the value of the objective."""
        # TODO: I'd like to be able to calculate all these properties for transformed data.
        m = self.moments  # Dictionary of moments
        y = x.dot(self.ws.T)  # + self.noise * np.random.randn(len(x), self.m)  # Noise is included analytically
        if "X_i^2" in self.moments:
            m["X_i^2"] = self.moments["X_i^2"]
        else:
            m["X_i^2"] = np.einsum('li,li->i', x, x).clip(1e-10) / (len(x) - 1)  # Variance of x, unbiased estimate
        m["X_i Y_j"] = x.T.dot(y) / len(y)  # nv by m,  <X_i Y_j_j>
        m["cy"] = self.ws.dot(m["X_i Y_j"]) + self.noise ** 2 * np.eye(self.m)  # cov(y.T), m by m
        m["X_i Z_j"] = np.linalg.solve(m["cy"], m["X_i Y_j"].T).T
        m["X_i^2 | Y"] = (m["X_i^2"] - np.einsum('ij,ij->i', m["X_i Z_j"], m["X_i Y_j"]))
        assert np.all(m["X_i^2 | Y"] > 0), \
            "Negative expected covariance suggests numerical instability in inversion of covariance matrix for Y."
        m["Y_j^2"] = np.diag(m["cy"]).copy()
        self.mis = self.calculate_mi(m)
        mi_yj_x = 0.5 * np.log(m["Y_j^2"]) - 0.5 * np.log(self.noise ** 2)
        mi_xi_y = 0.5 * np.log(m["X_i^2"]) - 0.5 * np.log(m["X_i^2 | Y"])
        self.tcs = self.mis.sum(axis=1) - mi_yj_x
        self.additivity = self.mis.sum(axis=0) - mi_xi_y
        self.objective = self.tc - self.additivity.sum()

    def _update_ws(self, i_loop):
        """Update weights, and also the lagrange multipliers."""
        m = self.moments  # Shorthand for readability
        H = ((1. - self.lam) / m["X_i^2 | Y"] * m["X_i Z_j"].T).dot(m["X_i Z_j"])  # np.einsum('ir,i,is->rs', m["X_i Z_j"], (1. - self.lam) / m["X_i^2 | Y"], m["X_i Z_j"])
        np.fill_diagonal(H, 0)
        if self.verbose == 3:
            print H
        if np.max(np.abs(H)) > 1. / self.noise**2:
            update = dominant_subset(H, 1. / self.noise**2, self.updates)
        else:
            update = list(range(self.m))
        self.updates[update] += 1
        Q = m["X_i Y_j"].T[update] / (m["X_i^2"] * m["Y_j^2"][update, np.newaxis] - m["X_i Y_j"].T[update] ** 2)
        R = m["X_i Z_j"].T[update] / m["X_i^2 | Y"]
        S = np.dot(H[update], self.ws)
        if i_loop < 10:
            self.ws[update] = 0.5 * self.ws[update] + 0.5 * self.noise**2 * (self.lam * Q + (1 - self.lam) * R - S)
        else:
            self.ws[update] = self.noise**2 * (self.lam * Q + (1 - self.lam) * R - S)
        if self.verbose >= 2:
            print 'updating %d / %d' % (len(update), self.m), np.min(self.lam), np.max(self.lam), np.mean(self.lam)

def gaussianize(x):
    """Return an empirically gaussianized version of either 1-d or 2-d data(processed column-wise)"""
    if len(x.shape) == 1:
        return norm.ppf((rankdata(x) - 0.5) / len(x))
    return np.array([norm.ppf((rankdata(x_i) - 0.5) / len(x_i)) for x_i in x.T]).T

def dominant_subset(H, threshold, updates):
    """Pick a submatrix of off-diagonal matrix, H, that is diagonally dominant."""
    H = np.abs(H) / threshold
    scores = np.sum(H, axis=0)
    update = list(range(len(H)))
    while np.any(scores > 1):
        p = scores * updates[update]
        p /= p.sum()
        i = np.random.choice(len(scores), p=p)  # weighted random
        # i = np.random.choice(len(scores))  # random
        # i = np.argmax(p)  # Greedy
        update.pop(i)
        subH = H[update][:, update]
        scores = np.sum(subH, axis=0)
    # print update, inhibit
    return update
<commit_msg>testing if I messed up git.<commit_after>""" Linear Total Correlation Explanation

Code below written by:
Greg Ver Steeg (gregv@isi.edu), 2015.
"""

import numpy as np
from scipy.stats import norm, rankdata  # Used for Gaussianizing data

class Corex(object):
    """
    Linear Total Correlation Explanation

    Conventions
    ----------
    Code follows sklearn naming/style (e.g. fit(X) to train).

    Parameters
    ----------
    n_hidden : int, default = 2
        The number of latent factors to use.

    max_iter : int, default = 100
        The max. number of iterations to reach convergence.

    noise : float default = 0.01
        We imagine some small fundamental measurement noise on Y. The value is arbitrary, but it sets
        the scale of the results, Y.

    verbose : int, optional
        Print verbose outputs.

    seed : integer or numpy.RandomState, optional
        A random number generator instance to define the state of the
        random permutations generator. If an integer is given, it fixes the
        seed. Defaults to the global numpy random number generator.

    Attributes
    ----------

    References
    ----------
    [1] Greg Ver Steeg and Aram Galstyan. "The Information Sieve", 2015.
    [2] ?, Greg Ver Steeg, and Aram Galstyan. "The Information Sieve for Continuous Variables" [In progress]
    [3] Greg Ver Steeg, ?, and Aram Galstyan. "Linear Total Correlation Explanation" [In progress]
    """

    def __init__(self, n_hidden=2, max_iter=10000, noise=1., lam=0., mu=0.01, tol=0.0001, additive=True,
                 gaussianize_marginals=False, verbose=False, seed=None, copy=True, **kwargs):
        self.m = n_hidden  # Number of latent factors to learn
        self.max_iter = max_iter  # Number of iterations to try
        self.noise = noise  # Sets the scale of Y
        self.mu = mu  # Initial step size for lagrange multiplier update
        self.tol = tol  # Threshold for convergence
        self.gaussianize_marginals = gaussianize_marginals
        self.verbose = verbose
        self.copy = copy  # Copy the data before subtracting the mean
        np.random.seed(seed)  # Set seed for deterministic results
        self.kwargs = kwargs

        # Initialize these when we fit on data
        self.n_samples, self.nv = 0, 0  # Number of samples/variables in input data
        self.ws = np.zeros((self.m, self.nv))  # m by nv array of weights
        self.additive = additive  # Whether or not to constrain to additive solutions
        self.lam = lam
        self.moments = {}  # dictionary of moments
        self.mean_x = None  # Mean is subtracted out, save for prediction/inversion
        self.updates = np.ones(self.m)  # Keep track of number of updates for each latent factor
        self.history = np.zeros((max_iter, 3))  # Keep track of values for TC, additivity and objective
        if verbose > 1:
            np.set_printoptions(precision=3, suppress=True, linewidth=160)
            print 'Linear CorEx with %d latent factors' % n_hidden

    @staticmethod
    def calculate_mi(moments):
        return (- 0.5 * np.log(1 - moments["X_i Y_j"] ** 2 / (moments["Y_j^2"] * moments["X_i^2"][:, np.newaxis]))).T

    @property
    def tc(self):
        """Sum_j TC(X;Y_j)"""
        return self.tcs.sum()

    def fit_transform(self, x, **kwargs):
        self.fit(x)
        return self.transform(x, **kwargs)

    def fit(self, x):
        x = np.array(x, dtype=float, copy=self.copy)
        self.n_samples, self.nv = x.shape  # Number of samples, variables in input data
        if self.gaussianize_marginals:
            x = gaussianize(x)
        else:
            self.mean_x = x.mean(axis=0)
            x -= self.mean_x
        var_x = np.einsum('li,li->i', x, x).clip(1e-10) / (self.n_samples - 1)  # Variance of x
        if self.ws.size == 0:  # Randomly initialize weights if not set
            self.ws = np.random.randn(self.m, self.nv) * self.noise ** 2 / np.sqrt(var_x)
        self.lam = self.lam * np.ones(self.nv)  # Initialize lagrange multipliers
        delta = np.inf

        for i_loop in range(self.max_iter):
            old_w = self.ws.copy()
            self._update_moments(x)  # Update moments based on w and samples, x.
            self.history[i_loop] = (self.tc, self.additivity.sum(), self.objective)  # Record convergence stats
            if self.additive and i_loop > 0:  # Update lambda dynamically to get additive solutions
                self.lam = (self.lam - self.mu * self.additivity).clip(0, 1)
            self._update_ws(i_loop)  # Update weights

            if self.additive:  # Adapt mu, the lambda step size, if necessary
                if i_loop % 10 == 9 and i_loop < self.max_iter / 10:
                    if self.additivity.sum() < - 0.01 * self.objective:
                        if (self.history[i_loop, 1] - self.history[i_loop-5, 1]) / 5. <= - 0.1 * self.additivity.sum():
                            self.mu = min(1., self.mu * 1.5)

            delta = np.max(np.sqrt(((old_w - self.ws)**2).sum(axis=1) / (self.ws**2).sum(axis=1)))  # max relative change for each factor
            if self.verbose > 1:
                print 'TC = %0.3f\tadd = %0.3f\tobj = %0.3f\tdelta = %0.5f\t\tmu %0.4f' % \
                      (self.tc, self.additivity.sum(), self.objective, delta, self.mu)
            if delta < self.tol:  # Check for convergence
                if self.verbose:
                    print '%d iterations to tol: %f' % (i_loop, self.tol)
                break
        else:
            if self.verbose:
                print "Warning: Convergence not achieved in %d iterations. Final delta: %f" % (self.max_iter, delta)

        self.history = self.history[:i_loop + 1]
        order = np.argsort(-self.tcs)  # Largest TC components first.
        self.ws = self.ws[order]
        self._update_moments(x)  # Update moments based on w and samples, x.
        return self

    def transform(self, x, details=False):
        """Transform an array of inputs, x, into an array of k latent factors, Y.
            Optionally, you can get the remainder information and/or stop at a specified level."""
        x = np.array(x, dtype=float, copy=self.copy)
        ns, nv = x.shape
        assert self.nv == nv, "Incorrect number of variables in input, %d instead of %d" % (nv, self.nv)
        if self.gaussianize_marginals:
            x = gaussianize(x)  # Should gaussianize wrt to original data...
        else:
            x -= self.mean_x
        if details:
            moments = self._calculate_moments(x)
            return x.dot(self.ws.T), moments
        else:
            return x.dot(self.ws.T)

    def predict(self, y):
        return np.dot(self.moments["X_i Z_j"], y.T).T + self.mean_x

    def _update_moments(self, x):
        """Calculate moments based on the weights and samples. We also calculate and save MI, TC, additivity, and
        the value of the objective."""
        # TODO: I'd like to be able to calculate all these properties for transformed data.
        m = self.moments  # Dictionary of moments
        y = x.dot(self.ws.T)  # + self.noise * np.random.randn(len(x), self.m)  # Noise is included analytically
        if "X_i^2" in self.moments:
            m["X_i^2"] = self.moments["X_i^2"]
        else:
            m["X_i^2"] = np.einsum('li,li->i', x, x).clip(1e-10) / (len(x) - 1)  # Variance of x, unbiased estimate
        m["X_i Y_j"] = x.T.dot(y) / len(y)  # nv by m,  <X_i Y_j_j>
        m["cy"] = self.ws.dot(m["X_i Y_j"]) + self.noise ** 2 * np.eye(self.m)  # cov(y.T), m by m
        m["X_i Z_j"] = np.linalg.solve(m["cy"], m["X_i Y_j"].T).T
        m["X_i^2 | Y"] = (m["X_i^2"] - np.einsum('ij,ij->i', m["X_i Z_j"], m["X_i Y_j"]))
        assert np.all(m["X_i^2 | Y"] > 0), \
            "Negative expected covariance suggests numerical instability in inversion of covariance matrix for Y."
        m["Y_j^2"] = np.diag(m["cy"]).copy()
        self.mis = self.calculate_mi(m)
        mi_yj_x = 0.5 * np.log(m["Y_j^2"]) - 0.5 * np.log(self.noise ** 2)
        mi_xi_y = 0.5 * np.log(m["X_i^2"]) - 0.5 * np.log(m["X_i^2 | Y"])
        self.tcs = self.mis.sum(axis=1) - mi_yj_x
        self.additivity = self.mis.sum(axis=0) - mi_xi_y
        self.objective = self.tc - self.additivity.sum()

    def _update_ws(self, i_loop):
        """Update weights, and also the lagrange multipliers."""
        m = self.moments  # Shorthand for readability
        H = ((1. - self.lam) / m["X_i^2 | Y"] * m["X_i Z_j"].T).dot(m["X_i Z_j"])  # np.einsum('ir,i,is->rs', m["X_i Z_j"], (1. - self.lam) / m["X_i^2 | Y"], m["X_i Z_j"])
        np.fill_diagonal(H, 0)
        if self.verbose == 3:
            print H
        if np.max(np.abs(H)) > 1. / self.noise**2:
            update = dominant_subset(H, 1. / self.noise**2, self.updates)
        else:
            update = list(range(self.m))
        self.updates[update] += 1
        Q = m["X_i Y_j"].T[update] / (m["X_i^2"] * m["Y_j^2"][update, np.newaxis] - m["X_i Y_j"].T[update] ** 2)
        R = m["X_i Z_j"].T[update] / m["X_i^2 | Y"]
        S = np.dot(H[update], self.ws)
        if i_loop < 10:
            self.ws[update] = 0.5 * self.ws[update] + 0.5 * self.noise**2 * (self.lam * Q + (1 - self.lam) * R - S)
        else:
            self.ws[update] = self.noise**2 * (self.lam * Q + (1 - self.lam) * R - S)
        if self.verbose >= 2:
            print 'updating %d / %d' % (len(update), self.m), np.min(self.lam), np.max(self.lam), np.mean(self.lam)

        # TESTing if I messed up git

def gaussianize(x):
    """Return an empirically gaussianized version of either 1-d or 2-d data(processed column-wise)"""
    if len(x.shape) == 1:
        return norm.ppf((rankdata(x) - 0.5) / len(x))
    return np.array([norm.ppf((rankdata(x_i) - 0.5) / len(x_i)) for x_i in x.T]).T

def dominant_subset(H, threshold, updates):
    """Pick a submatrix of off-diagonal matrix, H, that is diagonally dominant."""
    H = np.abs(H) / threshold
    scores = np.sum(H, axis=0)
    update = list(range(len(H)))
    while np.any(scores > 1):
        p = scores * updates[update]
        p /= p.sum()
        i = np.random.choice(len(scores), p=p)  # weighted random
        # i = np.random.choice(len(scores))  # random
        # i = np.argmax(p)  # Greedy
        update.pop(i)
        subH = H[update][:, update]
        scores = np.sum(subH, axis=0)
    # print update, inhibit
    return update
<|endoftext|>