<commit_before>import pandas as pd
import os.path as osp
import os
import numpy as np
from ogb.utils.url import decide_download, download_url, extract_zip
from tqdm import tqdm

### reading raw files from a directory.
def read_csv_graph_raw(raw_dir, add_inverse_edge = True, additional_node_files = [], additional_edge_files = []):
    '''
    raw_dir: path to the raw directory
    add_inverse_edge (bool): whether to add inverse edge or not

    return: graph_list, which is a list of graphs.
    Each graph is a dictionary, containing edge_index, edge_feat, node_feat, and num_nodes
    edge_feat and node_feat are optional: if a graph does not contain it, we will have None.

    additional_node_files and additional_edge_files must be in the raw directory.

    the name should be {additional_node_file, additional_edge_file}.csv.gz
    the length should be num_nodes or num_edges

    '''

    print('Loading necessary files...')
    # loading necessary files
    try:
        edge = pd.read_csv(osp.join(raw_dir, "edge.csv.gz"), compression="gzip", header = None).values.T.astype(np.int64) # (2, num_edge) numpy array
        num_node_list = pd.read_csv(osp.join(raw_dir, "num-node-list.csv.gz"), compression="gzip", header = None).values.T[0].astype(np.int64).tolist() # (num_graph, ) python list
        num_edge_list = pd.read_csv(osp.join(raw_dir, "num-edge-list.csv.gz"), compression="gzip", header = None).values.T[0].astype(np.int64).tolist() # (num_edge, ) python list
    except:
        raise RuntimeError("No necessary file")

    try:
        node_feat = pd.read_csv(osp.join(raw_dir, "node-feat.csv.gz"), compression="gzip", header = None).values
        if 'int' in str(node_feat.dtype):
            node_feat = node_feat.astype(np.int64)
        else:
            # float
            node_feat = node_feat.astype(np.float32)
    except:
        node_feat = None

    try:
        edge_feat = pd.read_csv(osp.join(raw_dir, "edge-feat.csv.gz"), compression="gzip", header = None).values
        if 'int' in str(edge_feat.dtype):
            edge_feat = edge_feat.astype(np.int64)
        else:
            #float
            edge_feat = edge_feat.astype(np.float32)

    except:
        edge_feat = None

    additional_node_info = {}
    for additional_file in additional_node_files:
        temp = pd.read_csv(osp.join(raw_dir, additional_file + ".csv.gz"), compression="gzip", header = None).values
        if 'int' in str(temp.dtype):
            additional_node_info[additional_file] = temp.astype(np.int64)
        else:
            # float
            additional_node_info[additional_file] = tempt.astype(np.float32)

    additional_edge_info = {}
    for additional_file in additional_edge_files:
        temp = pd.read_csv(osp.join(raw_dir, additional_file + ".csv.gz"), compression="gzip", header = None).values
        if 'int' in str(temp.dtype):
            additional_edge_info[additional_file] = temp.astype(np.int64)
        else:
            # float
            additional_edge_info[additional_file] = tempt.astype(np.float32)

    graph_list = []
    num_node_accum = 0
    num_edge_accum = 0

    print('Processing graphs...')
    for num_node, num_edge in tqdm(zip(num_node_list, num_edge_list), total=len(num_node_list)):

        graph = dict()

        ### handling edge
        if add_inverse_edge:
            ### duplicate edge
            duplicated_edge = np.repeat(edge[:, num_edge_accum:num_edge_accum+num_edge], 2, axis = 1)
            duplicated_edge[0, 1::2] = duplicated_edge[1,0::2]
            duplicated_edge[1, 1::2] = duplicated_edge[0,0::2]

            graph["edge_index"] = duplicated_edge

            if edge_feat is not None:
                graph["edge_feat"] = np.repeat(edge_feat[num_edge_accum:num_edge_accum+num_edge], 2, axis = 0)
            else:
                graph["edge_feat"] = None

            for key, value in additional_edge_info.items():
                graph[key] = np.repeat(value[num_edge_accum:num_edge_accum+num_edge], 2, axis = 0)

        else:
            graph["edge_index"] = edge[:, num_edge_accum:num_edge_accum+num_edge]

            if edge_feat is not None:
                graph["edge_feat"] = edge_feat[num_edge_accum:num_edge_accum+num_edge]
            else:
                graph["edge_feat"] = None

            for key, value in additional_edge_info.items():
                graph[key] = value[num_edge_accum:num_edge_accum+num_edge]

        num_edge_accum += num_edge

        ### handling node
        if node_feat is not None:
            graph["node_feat"] = node_feat[num_node_accum:num_node_accum+num_node]
        else:
            graph["node_feat"] = None

        for key, value in additional_node_info.items():
            graph[key] = value[num_node_accum:num_node_accum+num_node]

        graph["num_nodes"] = num_node
        num_node_accum += num_node

        graph_list.append(graph)

    return graph_list

if __name__ == "__main__":
    pass

<commit_msg>bug fix<commit_after>import pandas as pd
import os.path as osp
import os
import numpy as np
from ogb.utils.url import decide_download, download_url, extract_zip
from tqdm import tqdm

### reading raw files from a directory.
def read_csv_graph_raw(raw_dir, add_inverse_edge = True, additional_node_files = [], additional_edge_files = []):
    '''
    raw_dir: path to the raw directory
    add_inverse_edge (bool): whether to add inverse edge or not

    return: graph_list, which is a list of graphs.
    Each graph is a dictionary, containing edge_index, edge_feat, node_feat, and num_nodes
    edge_feat and node_feat are optional: if a graph does not contain it, we will have None.

    additional_node_files and additional_edge_files must be in the raw directory.

    the name should be {additional_node_file, additional_edge_file}.csv.gz
    the length should be num_nodes or num_edges

    '''

    print('Loading necessary files...')
    # loading necessary files
    try:
        edge = pd.read_csv(osp.join(raw_dir, "edge.csv.gz"), compression="gzip", header = None).values.T.astype(np.int64) # (2, num_edge) numpy array
        num_node_list = pd.read_csv(osp.join(raw_dir, "num-node-list.csv.gz"), compression="gzip", header = None).values.T[0].astype(np.int64).tolist() # (num_graph, ) python list
        num_edge_list = pd.read_csv(osp.join(raw_dir, "num-edge-list.csv.gz"), compression="gzip", header = None).values.T[0].astype(np.int64).tolist() # (num_edge, ) python list
    except:
        raise RuntimeError("No necessary file")

    try:
        node_feat = pd.read_csv(osp.join(raw_dir, "node-feat.csv.gz"), compression="gzip", header = None).values
        if 'int' in str(node_feat.dtype):
            node_feat = node_feat.astype(np.int64)
        else:
            # float
            node_feat = node_feat.astype(np.float32)
    except:
        node_feat = None

    try:
        edge_feat = pd.read_csv(osp.join(raw_dir, "edge-feat.csv.gz"), compression="gzip", header = None).values
        if 'int' in str(edge_feat.dtype):
            edge_feat = edge_feat.astype(np.int64)
        else:
            #float
            edge_feat = edge_feat.astype(np.float32)

    except:
        edge_feat = None

    additional_node_info = {}
    for additional_file in additional_node_files:
        temp = pd.read_csv(osp.join(raw_dir, additional_file + ".csv.gz"), compression="gzip", header = None).values
        if 'int' in str(temp.dtype):
            additional_node_info[additional_file] = temp.astype(np.int64)
        else:
            # float
            additional_node_info[additional_file] = temp.astype(np.float32)

    additional_edge_info = {}
    for additional_file in additional_edge_files:
        temp = pd.read_csv(osp.join(raw_dir, additional_file + ".csv.gz"), compression="gzip", header = None).values
        if 'int' in str(temp.dtype):
            additional_edge_info[additional_file] = temp.astype(np.int64)
        else:
            # float
            additional_edge_info[additional_file] = temp.astype(np.float32)

    graph_list = []
    num_node_accum = 0
    num_edge_accum = 0

    print('Processing graphs...')
    for num_node, num_edge in tqdm(zip(num_node_list, num_edge_list), total=len(num_node_list)):

        graph = dict()

        ### handling edge
        if add_inverse_edge:
            ### duplicate edge
            duplicated_edge = np.repeat(edge[:, num_edge_accum:num_edge_accum+num_edge], 2, axis = 1)
            duplicated_edge[0, 1::2] = duplicated_edge[1,0::2]
            duplicated_edge[1, 1::2] = duplicated_edge[0,0::2]

            graph["edge_index"] = duplicated_edge

            if edge_feat is not None:
                graph["edge_feat"] = np.repeat(edge_feat[num_edge_accum:num_edge_accum+num_edge], 2, axis = 0)
            else:
                graph["edge_feat"] = None

            for key, value in additional_edge_info.items():
                graph[key] = np.repeat(value[num_edge_accum:num_edge_accum+num_edge], 2, axis = 0)

        else:
            graph["edge_index"] = edge[:, num_edge_accum:num_edge_accum+num_edge]

            if edge_feat is not None:
                graph["edge_feat"] = edge_feat[num_edge_accum:num_edge_accum+num_edge]
            else:
                graph["edge_feat"] = None

            for key, value in additional_edge_info.items():
                graph[key] = value[num_edge_accum:num_edge_accum+num_edge]

        num_edge_accum += num_edge

        ### handling node
        if node_feat is not None:
            graph["node_feat"] = node_feat[num_node_accum:num_node_accum+num_node]
        else:
            graph["node_feat"] = None

        for key, value in additional_node_info.items():
            graph[key] = value[num_node_accum:num_node_accum+num_node]

        graph["num_nodes"] = num_node
        num_node_accum += num_node

        graph_list.append(graph)

    return graph_list

if __name__ == "__main__":
    pass

<|endoftext|>


<commit_before>import unittest

import numpy

import chainer
from chainer import cuda
from chainer import gradient_check
from chainer import links
from chainer import testing
from chainer.testing import attr
from chainer.testing import condition

def relu(x):
    return x * (x > 0)

@testing.parameterize(*testing.product({
    'use_cudnn': [True, False],
    'hidden_none': [True, False],
    'activation': ['tanh', 'relu'],
}))
class TestNStepRNN(unittest.TestCase):

    lengths = [3, 1, 2]
    n_layer = 2
    in_size = 3
    out_size = 2
    dropout = 0.0

    def setUp(self):
        shape = (self.n_layer, len(self.lengths), self.out_size)
        if self.hidden_none:
            self.h = numpy.zeros(shape, 'f')
        else:
            self.h = numpy.random.uniform(-1, 1, shape).astype('f')
        self.xs = [
            numpy.random.uniform(-1, 1, (l, self.in_size)).astype('f')
            for l in self.lengths]

        self.gh = numpy.random.uniform(-1, 1, shape).astype('f')
        self.gys = [
            numpy.random.uniform(-1, 1, (l, self.out_size)).astype('f')
            for l in self.lengths]
        if self.activation == 'tanh':
            rnn_link_class = links.NStepRNNTanh
        elif self.activation == 'relu':
            rnn_link_class = links.NStepRNNReLU
        self.rnn = rnn_link_class(
            self.n_layer, self.in_size, self.out_size, self.dropout,
            use_cudnn=self.use_cudnn)

        for layer in self.rnn:
            for p in layer.params():
                p.data[...] = numpy.random.uniform(-1, 1, p.data.shape)
        self.rnn.zerograds()

    def check_forward(self, h_data, xs_data):
        if self.hidden_none:
            h = None
        else:
            h = chainer.Variable(h_data)
        xs = [chainer.Variable(x) for x in xs_data]
        hy, ys = self.rnn(h, xs)

        self.assertEqual(hy.data.shape, h_data.shape)
        self.assertEqual(len(xs), len(ys))
        for x, y in zip(xs, ys):
            self.assertEqual(len(x.data), len(y.data))
            self.assertEqual(y.data.shape[1], self.out_size)

        self.rnn.to_cpu()

        for batch, seq in enumerate(self.xs):
            for layer in range(self.n_layer):
                p = self.rnn[layer]
                h_prev = self.h[layer, batch]
                hs = []
                for x in seq:
                    if self.activation == 'tanh':
                        activation_func = numpy.tanh
                    elif self.activation == 'relu':
                        activation_func = relu

                    h_prev = activation_func(x.dot(p.w0.data.T) +
                                             h_prev.dot(p.w1.data.T) +
                                             p.b0.data + p.b1.data)

                    hs.append(h_prev)

                seq = hs
                testing.assert_allclose(hy.data[layer, batch], h_prev)

            for y, ey in zip(ys[batch].data, seq):
                testing.assert_allclose(y, ey)

    def test_forward_cpu(self):
        self.check_forward(self.h, self.xs)

    @attr.gpu
    def test_forward_gpu(self):
        self.rnn.to_gpu()
        self.check_forward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs])

    def check_backward(
            self, h_data, xs_data, gh_data, gys_data):

        def fun(*args):
            if self.hidden_none:
                h = None
                xs = args
            else:
                h, = args[:1]
                xs = args[1:]
            hy, ys = self.rnn(h, xs)
            return tuple([hy, ] + list(ys))

        params = []
        for layer in self.rnn:
            for p in layer.params():
                params.append(p)

        if self.hidden_none:
            in_data = xs_data
        else:
            in_data = [h_data, ] + xs_data
        gradient_check.check_backward(
            fun, tuple(in_data),
            tuple([gh_data, ] + gys_data),
            tuple(params), eps=1e-2, rtol=1e-2, atol=5e-2)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(
            self.h, self.xs, self.gh, self.gys)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.rnn.to_gpu()
        self.check_backward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs],
            cuda.to_gpu(self.gh),
            [cuda.to_gpu(gy) for gy in self.gys])

@testing.parameterize(*testing.product({
    'use_cudnn': [True, False],
    'hidden_none': [True, False],
    'activation': ['tanh', 'relu'],
}))
class TestNStepBiRNN(unittest.TestCase):

    lengths = [3, 1, 2]
    n_layer = 2
    in_size = 3
    out_size = 2
    dropout = 0.0

    def setUp(self):
        shape = (self.n_layer * 2, len(self.lengths), self.out_size)
        if self.hidden_none:
            self.h = numpy.zeros(shape, 'f')
        else:
            self.h = numpy.random.uniform(-1, 1, shape).astype('f')
        self.xs = [
            numpy.random.uniform(-1, 1, (l, self.in_size)).astype('f')
            for l in self.lengths]

        self.gh = numpy.random.uniform(-1, 1, shape).astype('f')
        self.gys = [
            numpy.random.uniform(-1, 1, (l, self.out_size * 2)).astype('f')
            for l in self.lengths]
        if self.activation == 'tanh':
            rnn_link_class = links.NStepBiRNNTanh
        elif self.activation == 'relu':
            rnn_link_class = links.NStepBiRNNReLU
        self.rnn = rnn_link_class(
            self.n_layer, self.in_size, self.out_size, self.dropout,
            use_cudnn=self.use_cudnn)

        for layer in self.rnn:
            for p in layer.params():
                p.data[...] = numpy.random.uniform(-1, 1, p.data.shape)
        self.rnn.zerograds()

    def check_forward(self, h_data, xs_data):
        if self.hidden_none:
            h = None
        else:
            h = chainer.Variable(h_data)
        xs = [chainer.Variable(x) for x in xs_data]
        hy, ys = self.rnn(h, xs)

        self.assertEqual(hy.data.shape, h_data.shape)
        self.assertEqual(len(xs), len(ys))
        for x, y in zip(xs, ys):
            self.assertEqual(len(x.data), len(y.data))
            self.assertEqual(y.data.shape[1], self.out_size * 2)

        self.rnn.to_cpu()

        for batch, seq in enumerate(self.xs):
            for layer in range(self.n_layer):
                # forward
                di = 0
                layer_idx = layer * 2 + di
                p = self.rnn[layer_idx]
                h_prev = self.h[layer_idx, batch]
                hs_f = []
                for x in seq:
                    if self.activation == 'tanh':
                        activation_func = numpy.tanh
                    elif self.activation == 'relu':
                        activation_func = relu

                    h_prev = activation_func(x.dot(p.w0.data.T) +
                                             h_prev.dot(p.w1.data.T) +
                                             p.b0.data + p.b1.data)
                    hs_f.append(h_prev)

                testing.assert_allclose(hy.data[layer_idx, batch], h_prev)

                # backward
                di = 1
                layer_idx = layer * 2 + di
                p = self.rnn[layer_idx]
                h_prev = self.h[layer_idx, batch]
                hs_b = []
                for x in reversed(seq):
                    if self.activation == 'tanh':
                        activation_func = numpy.tanh
                    elif self.activation == 'relu':
                        activation_func = relu
                    h_prev = activation_func(x.dot(p.w0.data.T) +
                                             h_prev.dot(p.w1.data.T) +
                                             p.b0.data + p.b1.data)
                    hs_b.append(h_prev)
                testing.assert_allclose(hy.data[layer_idx, batch], h_prev)

                hs_b.reverse()
                seq = [numpy.concatenate([hfi, hbi], axis=0) for (hfi, hbi)
                       in zip(hs_f, hs_b)]

            for y, ey in zip(ys[batch].data, seq):
                testing.assert_allclose(y, ey)

    def test_forward_cpu(self):
        self.check_forward(self.h, self.xs)

    @attr.gpu
    def test_forward_gpu(self):
        self.rnn.to_gpu()
        self.check_forward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs])

    def check_backward(
            self, h_data, xs_data, gh_data, gys_data):

        def fun(*args):
            if self.hidden_none:
                h = None
                xs = args
            else:
                h, = args[:1]
                xs = args[1:]
            hy, ys = self.rnn(h, xs)
            return tuple([hy, ] + list(ys))

        params = []
        for layer in self.rnn:
            for p in layer.params():
                params.append(p)

        if self.hidden_none:
            in_data = xs_data
        else:
            in_data = [h_data, ] + xs_data
        gradient_check.check_backward(
            fun, tuple(in_data),
            tuple([gh_data, ] + gys_data),
            tuple(params), rtol=1e-2, atol=5e-2)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(
            self.h, self.xs, self.gh, self.gys)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.rnn.to_gpu()
        self.check_backward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs],
            cuda.to_gpu(self.gh),
            [cuda.to_gpu(gy) for gy in self.gys])

testing.run_module(__name__, __file__)
<commit_msg>modify eps value<commit_after>import unittest

import numpy

import chainer
from chainer import cuda
from chainer import gradient_check
from chainer import links
from chainer import testing
from chainer.testing import attr
from chainer.testing import condition

def relu(x):
    return x * (x > 0)

@testing.parameterize(*testing.product({
    'use_cudnn': [True, False],
    'hidden_none': [True, False],
    'activation': ['tanh', 'relu'],
}))
class TestNStepRNN(unittest.TestCase):

    lengths = [3, 1, 2]
    n_layer = 2
    in_size = 3
    out_size = 2
    dropout = 0.0

    def setUp(self):
        shape = (self.n_layer, len(self.lengths), self.out_size)
        if self.hidden_none:
            self.h = numpy.zeros(shape, 'f')
        else:
            self.h = numpy.random.uniform(-1, 1, shape).astype('f')
        self.xs = [
            numpy.random.uniform(-1, 1, (l, self.in_size)).astype('f')
            for l in self.lengths]

        self.gh = numpy.random.uniform(-1, 1, shape).astype('f')
        self.gys = [
            numpy.random.uniform(-1, 1, (l, self.out_size)).astype('f')
            for l in self.lengths]
        if self.activation == 'tanh':
            rnn_link_class = links.NStepRNNTanh
        elif self.activation == 'relu':
            rnn_link_class = links.NStepRNNReLU
        self.rnn = rnn_link_class(
            self.n_layer, self.in_size, self.out_size, self.dropout,
            use_cudnn=self.use_cudnn)

        for layer in self.rnn:
            for p in layer.params():
                p.data[...] = numpy.random.uniform(-1, 1, p.data.shape)
        self.rnn.zerograds()

    def check_forward(self, h_data, xs_data):
        if self.hidden_none:
            h = None
        else:
            h = chainer.Variable(h_data)
        xs = [chainer.Variable(x) for x in xs_data]
        hy, ys = self.rnn(h, xs)

        self.assertEqual(hy.data.shape, h_data.shape)
        self.assertEqual(len(xs), len(ys))
        for x, y in zip(xs, ys):
            self.assertEqual(len(x.data), len(y.data))
            self.assertEqual(y.data.shape[1], self.out_size)

        self.rnn.to_cpu()

        for batch, seq in enumerate(self.xs):
            for layer in range(self.n_layer):
                p = self.rnn[layer]
                h_prev = self.h[layer, batch]
                hs = []
                for x in seq:
                    if self.activation == 'tanh':
                        activation_func = numpy.tanh
                    elif self.activation == 'relu':
                        activation_func = relu

                    h_prev = activation_func(x.dot(p.w0.data.T) +
                                             h_prev.dot(p.w1.data.T) +
                                             p.b0.data + p.b1.data)

                    hs.append(h_prev)

                seq = hs
                testing.assert_allclose(hy.data[layer, batch], h_prev)

            for y, ey in zip(ys[batch].data, seq):
                testing.assert_allclose(y, ey)

    def test_forward_cpu(self):
        self.check_forward(self.h, self.xs)

    @attr.gpu
    def test_forward_gpu(self):
        self.rnn.to_gpu()
        self.check_forward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs])

    def check_backward(
            self, h_data, xs_data, gh_data, gys_data):

        def fun(*args):
            if self.hidden_none:
                h = None
                xs = args
            else:
                h, = args[:1]
                xs = args[1:]
            hy, ys = self.rnn(h, xs)
            return tuple([hy, ] + list(ys))

        params = []
        for layer in self.rnn:
            for p in layer.params():
                params.append(p)

        if self.hidden_none:
            in_data = xs_data
        else:
            in_data = [h_data, ] + xs_data
        gradient_check.check_backward(
            fun, tuple(in_data),
            tuple([gh_data, ] + gys_data),
            tuple(params), rtol=1e-2, atol=5e-2)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(
            self.h, self.xs, self.gh, self.gys)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.rnn.to_gpu()
        self.check_backward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs],
            cuda.to_gpu(self.gh),
            [cuda.to_gpu(gy) for gy in self.gys])

@testing.parameterize(*testing.product({
    'use_cudnn': [True, False],
    'hidden_none': [True, False],
    'activation': ['tanh', 'relu'],
}))
class TestNStepBiRNN(unittest.TestCase):

    lengths = [3, 1, 2]
    n_layer = 2
    in_size = 3
    out_size = 2
    dropout = 0.0

    def setUp(self):
        shape = (self.n_layer * 2, len(self.lengths), self.out_size)
        if self.hidden_none:
            self.h = numpy.zeros(shape, 'f')
        else:
            self.h = numpy.random.uniform(-1, 1, shape).astype('f')
        self.xs = [
            numpy.random.uniform(-1, 1, (l, self.in_size)).astype('f')
            for l in self.lengths]

        self.gh = numpy.random.uniform(-1, 1, shape).astype('f')
        self.gys = [
            numpy.random.uniform(-1, 1, (l, self.out_size * 2)).astype('f')
            for l in self.lengths]
        if self.activation == 'tanh':
            rnn_link_class = links.NStepBiRNNTanh
        elif self.activation == 'relu':
            rnn_link_class = links.NStepBiRNNReLU
        self.rnn = rnn_link_class(
            self.n_layer, self.in_size, self.out_size, self.dropout,
            use_cudnn=self.use_cudnn)

        for layer in self.rnn:
            for p in layer.params():
                p.data[...] = numpy.random.uniform(-1, 1, p.data.shape)
        self.rnn.zerograds()

    def check_forward(self, h_data, xs_data):
        if self.hidden_none:
            h = None
        else:
            h = chainer.Variable(h_data)
        xs = [chainer.Variable(x) for x in xs_data]
        hy, ys = self.rnn(h, xs)

        self.assertEqual(hy.data.shape, h_data.shape)
        self.assertEqual(len(xs), len(ys))
        for x, y in zip(xs, ys):
            self.assertEqual(len(x.data), len(y.data))
            self.assertEqual(y.data.shape[1], self.out_size * 2)

        self.rnn.to_cpu()

        for batch, seq in enumerate(self.xs):
            for layer in range(self.n_layer):
                # forward
                di = 0
                layer_idx = layer * 2 + di
                p = self.rnn[layer_idx]
                h_prev = self.h[layer_idx, batch]
                hs_f = []
                for x in seq:
                    if self.activation == 'tanh':
                        activation_func = numpy.tanh
                    elif self.activation == 'relu':
                        activation_func = relu

                    h_prev = activation_func(x.dot(p.w0.data.T) +
                                             h_prev.dot(p.w1.data.T) +
                                             p.b0.data + p.b1.data)
                    hs_f.append(h_prev)

                testing.assert_allclose(hy.data[layer_idx, batch], h_prev)

                # backward
                di = 1
                layer_idx = layer * 2 + di
                p = self.rnn[layer_idx]
                h_prev = self.h[layer_idx, batch]
                hs_b = []
                for x in reversed(seq):
                    if self.activation == 'tanh':
                        activation_func = numpy.tanh
                    elif self.activation == 'relu':
                        activation_func = relu
                    h_prev = activation_func(x.dot(p.w0.data.T) +
                                             h_prev.dot(p.w1.data.T) +
                                             p.b0.data + p.b1.data)
                    hs_b.append(h_prev)
                testing.assert_allclose(hy.data[layer_idx, batch], h_prev)

                hs_b.reverse()
                seq = [numpy.concatenate([hfi, hbi], axis=0) for (hfi, hbi)
                       in zip(hs_f, hs_b)]

            for y, ey in zip(ys[batch].data, seq):
                testing.assert_allclose(y, ey)

    def test_forward_cpu(self):
        self.check_forward(self.h, self.xs)

    @attr.gpu
    def test_forward_gpu(self):
        self.rnn.to_gpu()
        self.check_forward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs])

    def check_backward(
            self, h_data, xs_data, gh_data, gys_data):

        def fun(*args):
            if self.hidden_none:
                h = None
                xs = args
            else:
                h, = args[:1]
                xs = args[1:]
            hy, ys = self.rnn(h, xs)
            return tuple([hy, ] + list(ys))

        params = []
        for layer in self.rnn:
            for p in layer.params():
                params.append(p)

        if self.hidden_none:
            in_data = xs_data
        else:
            in_data = [h_data, ] + xs_data
        gradient_check.check_backward(
            fun, tuple(in_data),
            tuple([gh_data, ] + gys_data),
            tuple(params), rtol=1e-2, atol=5e-2)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(
            self.h, self.xs, self.gh, self.gys)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.rnn.to_gpu()
        self.check_backward(
            cuda.to_gpu(self.h),
            [cuda.to_gpu(x) for x in self.xs],
            cuda.to_gpu(self.gh),
            [cuda.to_gpu(gy) for gy in self.gys])

testing.run_module(__name__, __file__)
<|endoftext|>


<commit_before>import types, copy, datetime, time
import logging

import asyncio
from pymongo.errors import DuplicateKeyError, BulkWriteError

from biothings.utils.common import timesofar, iter_n
from biothings.utils.dataload import merge_struct, merge_root_keys
from biothings.utils.mongo import get_src_db, check_document_size

class StorageException(Exception):
    pass

class BaseStorage(object):

    def __init__(self,db,dest_col_name,logger=logging):
        db = db or get_src_db()
        self.temp_collection = db[dest_col_name]
        self.logger = logger

    def process(self,iterable,*args,**kwargs):
        """
        Process iterable to store data. Must return the number
        of inserted records (even 0 if none)
        """
        raise NotImplementedError("implement-me in subclass")

    def check_doc_func(self,doc):
        return True

class CheckSizeStorage(BaseStorage):

    def check_doc_func(self,doc):
        ok = check_document_size(doc)
        # this is typically used to skip LFQSCWFLJHTTHZ-UHFFFAOYSA-N (Ethanol)
        # because there are too many elements in "ndc" list
        if not ok:
            self.logger.warning("Skip document '%s' because too large" % doc.get("_id"))
        return ok

class BasicStorage(BaseStorage):

    def doc_iterator(self, doc_d, batch=True, batch_size=10000):
        if isinstance(doc_d, types.GeneratorType) and batch:
            for doc_li in iter_n(doc_d, n=batch_size):
                doc_li = [d for d in doc_li if self.check_doc_func(d)]
                yield doc_li
        else:
            if batch:
                doc_li = []
                i = 0
            for _id, doc in doc_d.items():
                doc['_id'] = _id
                _doc = {}
                _doc.update(doc)
                if batch:
                    doc_li.append(_doc)
                    i += 1
                    if i % batch_size == 0:
                        doc_li = [d for d in doc_li if self.check_doc_func(d)]
                        yield doc_li
                        doc_li = []
                else:
                    yield self.check_doc_func(_doc)

            if batch:
                doc_li = [d for d in doc_li if self.check_doc_func(d)]
                yield doc_li

    def process(self, doc_d, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        total = 0
        for doc_li in self.doc_iterator(doc_d, batch=True, batch_size=batch_size):
            try:
                self.temp_collection.insert_many(doc_li, ordered=False)
            except BulkWriteError as bwe:
                # TODO:  Write logic for reciprocal key matching
                # - currently these should be considered serious errors
                self.logger.error("BulkWriteError - likely duplicate key without reciprocal matching on insert.")
                for we in bwe.details['writeErrors']:
                    self.logger.error("BulkWriteError Duplicate Key:  {}".format(we['op']['_id']))
                    # print the keylookup debug information if it is available
                    if 'dt_debug' in we['op'].keys():
                        self.logger.error("BulkWriteError KeyLookup Information:  {}".format(we['op']['dt_debug']))
            total += len(doc_li)
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class MergerStorage(BasicStorage):
    """
    This storage will try to merge documents when finding duplicated errors.
    It's useful when data is parsed using iterator. A record can be stored in database,
    then later, another record with the same ID is sent to the db, raising a duplicated error.
    These two documents would have been merged before using a 'put all in memory' parser.
    Since data is here read line by line, the merge is done while storing
    """

    merge_func = merge_struct

    def process(self, doc_d, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        aslistofdict = None
        total = 0
        for doc_li in self.doc_iterator(doc_d, batch=True, batch_size=batch_size):
            toinsert = len(doc_li)
            nbinsert = 0
            self.logger.info("Inserting %s records ... " % toinsert)
            try:
                bob = self.temp_collection.initialize_unordered_bulk_op()
                for d in doc_li:
                    aslistofdict = d.pop("__aslistofdict__",None)
                    bob.insert(d)
                res = bob.execute()
                nbinsert += res["nInserted"]
                self.logger.info("OK [%s]" % timesofar(tinner))
            except BulkWriteError as e:
                inserted = e.details["nInserted"]
                nbinsert += inserted
                self.logger.info("Fixing %d records " % len(e.details["writeErrors"]))
                ids = [d["op"]["_id"] for d in e.details["writeErrors"]]
                # build hash of existing docs
                docs = self.temp_collection.find({"_id" : {"$in" : ids}})
                hdocs = {}
                for doc in docs:
                    hdocs[doc["_id"]] = doc
                bob2 = self.temp_collection.initialize_unordered_bulk_op()
                for err in e.details["writeErrors"]:
                    errdoc = err["op"]
                    existing = hdocs[errdoc["_id"]]
                    assert "_id" in existing
                    _id = errdoc.pop("_id")
                    merged = self.__class__.merge_func(errdoc, existing, aslistofdict=aslistofdict)
                    bob2.find({"_id" : _id}).update_one({"$set" : merged})
                    # update previously fetched doc. if several errors are about the same doc id,
                    # we would't merged things properly without an updated document
                    assert "_id" in merged
                    hdocs[_id] = merged
                    nbinsert += 1

                res = bob2.execute()
                self.logger.info("OK [%s]" % timesofar(tinner))
            assert nbinsert == toinsert, "nb %s to %s" % (nbinsert,toinsert)
            # end of loop so it counts the time spent in doc_iterator
            tinner = time.time()
            total += nbinsert

        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class RootKeyMergerStorage(MergerStorage):
    """
    Just like MergerStorage, this storage deals with duplicated error
    by appending key's content to existing document. Keys in existing
    document will be converted to a list as needed.

    Note:
      - root keys must have the same type in each documents
      - inner structures aren't merged together, the merge happend
        at root key level
    """

    @classmethod
    def merge_func(klass, doc1, doc2, **kwargs):
        # caller popped it from doc1, let's take from doc2
        _id = doc2["_id"]
        # exclude_id will remove _if from doc2, that's why we kept it from before
        # also, copy doc2 ref as the merged doc will be stored in
        # a bulk op object, since doc2 is modified in place, this could impact
        # the bulk op and procude empty $set error from mongo
        doc = merge_root_keys(doc1,copy.copy(doc2),exclude=["_id"])
        doc["_id"] = _id
        return doc

class IgnoreDuplicatedStorage(BasicStorage):

    def process(self, iterable, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        total = 0
        for doc_li in self.doc_iterator(iterable, batch=True, batch_size=batch_size):
            try:
                bob = self.temp_collection.initialize_unordered_bulk_op()
                for d in doc_li:
                    bob.insert(d)
                res = bob.execute()
                total += res['nInserted']
                self.logger.info("Inserted %s records [%s]" % (res['nInserted'], timesofar(tinner)))
            except BulkWriteError as e:
                self.logger.info("Inserted %s records, ignoring %d [%s]" % (e.details['nInserted'],len(e.details["writeErrors"]),timesofar(tinner)))
            except Exception as e:
                raise
            tinner = time.time()
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class NoBatchIgnoreDuplicatedStorage(BasicStorage):
    """
    You should use IgnoreDuplicatedStorag, which works using batch
    and is thus way faster...
    """

    def process(self, doc_d, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        # force step = 1
        cnt = 0
        total = 0
        dups = 0
        for doc_li in self.doc_iterator(doc_d, batch=True, batch_size=1):
            try:
                res = self.temp_collection.insert(doc_li, manipulate=False, check_keys=False)
                cnt += 1
                total += 1
                if (cnt + dups) % batch_size == 0:
                    # we insert one by one but display progress on a "batch_size" base
                    self.logger.info("Inserted %s records, ignoring %s [%s]" % (cnt,dups,timesofar(tinner)))
                    cnt = 0
                    dups = 0
                    tinner = time.time()
            except DuplicateKeyError:
                dups += 1
                pass
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class UpsertStorage(BasicStorage):
    """Insert or update documents, based on _id"""

    def process(self, iterable, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        total = 0
        for doc_li in self.doc_iterator(iterable, batch=True, batch_size=batch_size):
            try:
                bob = self.temp_collection.initialize_unordered_bulk_op()
                for d in doc_li:
                    bob.find({"_id" : d["_id"]}).upsert().replace_one(d)
                res = bob.execute()
                nb = res["nUpserted"] + res["nModified"]
                total += nb
                self.logger.info("Upserted %s records [%s]" % (nb,timesofar(tinner)))
            except Exception as e:
                raise
            tinner = time.time()
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class NoStorage(object):
    """
    This a kind of a place-holder, this storage will just store nothing...
    (but it will respect storage interface)
    """

    def __init__(self,db_info,dest_col_name,logger):
        db = get_src_db()
        self.temp_collection = db[dest_col_name]
        self.logger = logger

    def process(self,iterable,*args,**kwargs):
        self.logger.info("NoStorage stores nothing, skip...")
        return 0

<commit_msg>storage.py - raise StorageException on BulkWriteError with duplicate keys<commit_after>import types, copy, datetime, time
import logging

import asyncio
from pymongo.errors import DuplicateKeyError, BulkWriteError

from biothings.utils.common import timesofar, iter_n
from biothings.utils.dataload import merge_struct, merge_root_keys
from biothings.utils.mongo import get_src_db, check_document_size

class StorageException(Exception):
    pass

class BaseStorage(object):

    def __init__(self,db,dest_col_name,logger=logging):
        db = db or get_src_db()
        self.temp_collection = db[dest_col_name]
        self.logger = logger

    def process(self,iterable,*args,**kwargs):
        """
        Process iterable to store data. Must return the number
        of inserted records (even 0 if none)
        """
        raise NotImplementedError("implement-me in subclass")

    def check_doc_func(self,doc):
        return True

class CheckSizeStorage(BaseStorage):

    def check_doc_func(self,doc):
        ok = check_document_size(doc)
        # this is typically used to skip LFQSCWFLJHTTHZ-UHFFFAOYSA-N (Ethanol)
        # because there are too many elements in "ndc" list
        if not ok:
            self.logger.warning("Skip document '%s' because too large" % doc.get("_id"))
        return ok

class BasicStorage(BaseStorage):

    def doc_iterator(self, doc_d, batch=True, batch_size=10000):
        if isinstance(doc_d, types.GeneratorType) and batch:
            for doc_li in iter_n(doc_d, n=batch_size):
                doc_li = [d for d in doc_li if self.check_doc_func(d)]
                yield doc_li
        else:
            if batch:
                doc_li = []
                i = 0
            for _id, doc in doc_d.items():
                doc['_id'] = _id
                _doc = {}
                _doc.update(doc)
                if batch:
                    doc_li.append(_doc)
                    i += 1
                    if i % batch_size == 0:
                        doc_li = [d for d in doc_li if self.check_doc_func(d)]
                        yield doc_li
                        doc_li = []
                else:
                    yield self.check_doc_func(_doc)

            if batch:
                doc_li = [d for d in doc_li if self.check_doc_func(d)]
                yield doc_li

    def process(self, doc_d, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        total = 0
        for doc_li in self.doc_iterator(doc_d, batch=True, batch_size=batch_size):
            try:
                self.temp_collection.insert_many(doc_li, ordered=False)
            except BulkWriteError as bwe:
                # TODO:  Write logic for reciprocal key matching
                # - currently these should be considered serious errors
                self.logger.error("BulkWriteError - likely duplicate key without reciprocal matching on insert.")
                for we in bwe.details['writeErrors']:
                    self.logger.error("BulkWriteError Duplicate Key:  {}".format(we['op']['_id']))
                    # print the keylookup debug information if it is available
                    if 'dt_debug' in we['op'].keys():
                        self.logger.error("BulkWriteError KeyLookup Information:  {}".format(we['op']['dt_debug']))
                # this is a serious error, so we raise
                raise StorageException("BulkWriteError - likely duplicate key without reciprocal matching on insert.")
            total += len(doc_li)
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class MergerStorage(BasicStorage):
    """
    This storage will try to merge documents when finding duplicated errors.
    It's useful when data is parsed using iterator. A record can be stored in database,
    then later, another record with the same ID is sent to the db, raising a duplicated error.
    These two documents would have been merged before using a 'put all in memory' parser.
    Since data is here read line by line, the merge is done while storing
    """

    merge_func = merge_struct

    def process(self, doc_d, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        aslistofdict = None
        total = 0
        for doc_li in self.doc_iterator(doc_d, batch=True, batch_size=batch_size):
            toinsert = len(doc_li)
            nbinsert = 0
            self.logger.info("Inserting %s records ... " % toinsert)
            try:
                bob = self.temp_collection.initialize_unordered_bulk_op()
                for d in doc_li:
                    aslistofdict = d.pop("__aslistofdict__",None)
                    bob.insert(d)
                res = bob.execute()
                nbinsert += res["nInserted"]
                self.logger.info("OK [%s]" % timesofar(tinner))
            except BulkWriteError as e:
                inserted = e.details["nInserted"]
                nbinsert += inserted
                self.logger.info("Fixing %d records " % len(e.details["writeErrors"]))
                ids = [d["op"]["_id"] for d in e.details["writeErrors"]]
                # build hash of existing docs
                docs = self.temp_collection.find({"_id" : {"$in" : ids}})
                hdocs = {}
                for doc in docs:
                    hdocs[doc["_id"]] = doc
                bob2 = self.temp_collection.initialize_unordered_bulk_op()
                for err in e.details["writeErrors"]:
                    errdoc = err["op"]
                    existing = hdocs[errdoc["_id"]]
                    assert "_id" in existing
                    _id = errdoc.pop("_id")
                    merged = self.__class__.merge_func(errdoc, existing, aslistofdict=aslistofdict)
                    bob2.find({"_id" : _id}).update_one({"$set" : merged})
                    # update previously fetched doc. if several errors are about the same doc id,
                    # we would't merged things properly without an updated document
                    assert "_id" in merged
                    hdocs[_id] = merged
                    nbinsert += 1

                res = bob2.execute()
                self.logger.info("OK [%s]" % timesofar(tinner))
            assert nbinsert == toinsert, "nb %s to %s" % (nbinsert,toinsert)
            # end of loop so it counts the time spent in doc_iterator
            tinner = time.time()
            total += nbinsert

        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class RootKeyMergerStorage(MergerStorage):
    """
    Just like MergerStorage, this storage deals with duplicated error
    by appending key's content to existing document. Keys in existing
    document will be converted to a list as needed.

    Note:
      - root keys must have the same type in each documents
      - inner structures aren't merged together, the merge happend
        at root key level
    """

    @classmethod
    def merge_func(klass, doc1, doc2, **kwargs):
        # caller popped it from doc1, let's take from doc2
        _id = doc2["_id"]
        # exclude_id will remove _if from doc2, that's why we kept it from before
        # also, copy doc2 ref as the merged doc will be stored in
        # a bulk op object, since doc2 is modified in place, this could impact
        # the bulk op and procude empty $set error from mongo
        doc = merge_root_keys(doc1,copy.copy(doc2),exclude=["_id"])
        doc["_id"] = _id
        return doc

class IgnoreDuplicatedStorage(BasicStorage):

    def process(self, iterable, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        total = 0
        for doc_li in self.doc_iterator(iterable, batch=True, batch_size=batch_size):
            try:
                bob = self.temp_collection.initialize_unordered_bulk_op()
                for d in doc_li:
                    bob.insert(d)
                res = bob.execute()
                total += res['nInserted']
                self.logger.info("Inserted %s records [%s]" % (res['nInserted'], timesofar(tinner)))
            except BulkWriteError as e:
                self.logger.info("Inserted %s records, ignoring %d [%s]" % (e.details['nInserted'],len(e.details["writeErrors"]),timesofar(tinner)))
            except Exception as e:
                raise
            tinner = time.time()
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class NoBatchIgnoreDuplicatedStorage(BasicStorage):
    """
    You should use IgnoreDuplicatedStorag, which works using batch
    and is thus way faster...
    """

    def process(self, doc_d, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        # force step = 1
        cnt = 0
        total = 0
        dups = 0
        for doc_li in self.doc_iterator(doc_d, batch=True, batch_size=1):
            try:
                res = self.temp_collection.insert(doc_li, manipulate=False, check_keys=False)
                cnt += 1
                total += 1
                if (cnt + dups) % batch_size == 0:
                    # we insert one by one but display progress on a "batch_size" base
                    self.logger.info("Inserted %s records, ignoring %s [%s]" % (cnt,dups,timesofar(tinner)))
                    cnt = 0
                    dups = 0
                    tinner = time.time()
            except DuplicateKeyError:
                dups += 1
                pass
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class UpsertStorage(BasicStorage):
    """Insert or update documents, based on _id"""

    def process(self, iterable, batch_size):
        self.logger.info("Uploading to the DB...")
        t0 = time.time()
        tinner = time.time()
        total = 0
        for doc_li in self.doc_iterator(iterable, batch=True, batch_size=batch_size):
            try:
                bob = self.temp_collection.initialize_unordered_bulk_op()
                for d in doc_li:
                    bob.find({"_id" : d["_id"]}).upsert().replace_one(d)
                res = bob.execute()
                nb = res["nUpserted"] + res["nModified"]
                total += nb
                self.logger.info("Upserted %s records [%s]" % (nb,timesofar(tinner)))
            except Exception as e:
                raise
            tinner = time.time()
        self.logger.info('Done[%s]' % timesofar(t0))

        return total

class NoStorage(object):
    """
    This a kind of a place-holder, this storage will just store nothing...
    (but it will respect storage interface)
    """

    def __init__(self,db_info,dest_col_name,logger):
        db = get_src_db()
        self.temp_collection = db[dest_col_name]
        self.logger = logger

    def process(self,iterable,*args,**kwargs):
        self.logger.info("NoStorage stores nothing, skip...")
        return 0

<|endoftext|>


<commit_before>import pytest
from ogusa import firm
import numpy as np

def test_get_r():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    Z = 0.5
    gamma = 0.5
    delta = 0.25
    tau_b = 0.5
    delta_tau = 1.0

    Y = np.array([2.0, 2.0])
    K = np.array([1.0, 1.0])

    epsilon = 0.0
    r = firm.get_r(Y, K, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(r, np.array([0.0, 0.0])))

    epsilon = 0.5
    delta = 0.5
    r = firm.get_r(Y, K, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(r, np.array([1.0, 1.0])))

def test_get_w():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    Z = 0.5
    gamma = 0.5
    Y = np.array([2.0, 2.0])
    L = np.array([1.0, 1.0])

    epsilon = 0.0
    w = firm.get_w(Y, L, (Z, gamma, epsilon))
    assert (np.allclose(w, np.array([0.5, 0.5])))

    epsilon = 0.5
    w = firm.get_w(Y, L, (Z, gamma, epsilon))
    assert (np.allclose(w, np.array([2.0, 2.0])))

def test_get_Y():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    Z = 2.0
    gamma = 0.5
    L = np.array([4.0, 4.0])
    K = np.array([9.0, 9.0])

    epsilon = 1
    Y = firm.get_Y(K, L, (Z, gamma, epsilon))
    assert (np.allclose(Y, np.array([12.0, 12.0])))

    epsilon = 0
    Y = firm.get_Y(K, L, (Z, gamma, epsilon))
    assert (np.allclose(Y, np.array([13.0, 13.0])))

    epsilon = 0.5
    Z = 2
    L = np.array([1/12.0, 1/12.0])
    K = np.array([1/4.0, 1/4.0])

    Y = firm.get_Y(K, L, (Z, gamma, epsilon))
    assert (np.allclose(Y, np.array([0.5, 0.5])))

def test_get_L():
    """
        Simulate data similar to observed and carry out get_L
        in simplest way possible.
    """
    T = 160
    s, j = 40, 2
    omega = np.random.rand(s).reshape(s, 1)
    lambdas = np.random.rand(j).reshape(1, j)
    n = np.random.rand(T * s * j).reshape(T, s, j)
    e = np.tile(np.random.rand(s, j), (T, 1, 1))

    L_loop = np.ones(T * s * j).reshape(T, s, j)
    for t in range(T):
        for i in range(s):
            for k in range(j):
                L_loop[t, i, k] *= (omega[i, 0] * lambdas[0, k] *
                                    n[t, i, k] * e[t, i, k])

    L_matrix = e * omega * lambdas * n
    assert (np.allclose(L_loop, L_matrix))

    method = 'SS'
    L = firm.get_L(n[0], (e[0], omega, lambdas, method))
    assert (np.allclose(L, L_loop[0].sum()))

    method = 'TPI'
    L = firm.get_L(n, (e, omega, lambdas, method))
    assert (np.allclose(L, L_loop.sum(1).sum(1)))

def test_get_I():
    """
        Simulate data similar to observed and carry out get_I
        in simplest way possible.
    """
    T = 160
    s, j = 40, 2

    b_splus1 = 10 * np.random.rand(T * s * j).reshape(T, s, j)
    K_p1 = 0.9 + np.random.rand(T)
    K = 0.9 + np.random.rand(T)

    delta = np.random.rand()
    g_y = np.random.rand()

    def shifted_arr():
        arr_t = []
        arr_shift_t = []
        for t in range(T):
            arr = np.random.rand(s).reshape(s, 1)
            arr_shift = np.append(arr[1:], [0.0])

            arr_t.append(arr)
            arr_shift_t.append(arr_shift)

        return (np.array(arr_t).reshape(T, s, 1),
                np.array(arr_shift_t).reshape(T, s, 1))

    lambdas = np.random.rand(2)
    imm_rates, imm_shift = shifted_arr()
    imm_rates = imm_rates - 0.5
    imm_shift = imm_shift - 0.5
    omega, omega_shift = shifted_arr()

    g_n = np.random.rand(T)

    res_loop = np.ones(T * s * j).reshape(T, s, j)
    for t in range(T):
        for i in range(s):
            for k in range(j):
                res_loop[t, i, k] *= (omega_shift[t, i, 0] * imm_shift[t, i, 0] *
                                      lambdas[k] * b_splus1[t, i, k])

    res_matrix = (b_splus1 * (imm_shift * omega_shift) * lambdas)

    assert (np.allclose(res_loop, res_matrix))

    aggI_SS_test = ((1 + g_n[0]) * np.exp(g_y) *
                    (K_p1[0] - res_matrix[0].sum() / (1 + g_n[0])) -
                    (1.0 - delta) * K[0])
    aggI_SS = firm.get_I(b_splus1[0], K_p1[0], K[0],
                         (delta, g_y, omega[0], lambdas, imm_rates[0], g_n[0], 'SS'))
    assert (np.allclose(aggI_SS, aggI_SS_test))

    aggI_TPI_test = ((1 + g_n) * np.exp(g_y) *
                     (K_p1 - res_matrix.sum(1).sum(1) / (1 + g_n)) -
                     (1.0 - delta) * K)
    aggI_TPI = firm.get_I(b_splus1, K_p1, K,
                          (delta, g_y, omega, lambdas, imm_rates, g_n, 'TPI'))
    assert (np.allclose(aggI_TPI, aggI_TPI_test))

def test_get_K():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    L = np.array([2.0, 2.0])
    r = np.array([1.0, 1.0])
    gamma = 0.5
    tau_b = 0.75
    delta = 4.0
    delta_tau = 4.0/3.0
    Z = 2.0

    epsilon = 1
    K = firm.get_K(L, r, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(K, np.array([2./(16**2), 2./(16**2)])))

    epsilon = 0
    K = firm.get_K(L, r, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(K, np.array([0.0, 0.0])))

    epsilon = 0.5
    Z = 4.0
    tau_b = 0.0
    delta_tau = 2.0
    delta = 3.0
    K = firm.get_K(L, r, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(K, np.array([2.0, 2.0])))
<commit_msg>update test parameter values<commit_after>import pytest
from ogusa import firm
import numpy as np

def test_get_r():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    Z = 0.5
    gamma = 0.5
    delta = 0.25
    tau_b = 0.5
    delta_tau = 0.35

    Y = np.array([2.0, 2.0])
    K = np.array([1.0, 1.0])

    epsilon = 0.0
    r = firm.get_r(Y, K, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(r, np.array([0.0, 0.0])))

    epsilon = 0.5
    delta = 0.5
    r = firm.get_r(Y, K, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(r, np.array([0.675, 0.675])))

def test_get_w():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    Z = 0.5
    gamma = 0.5
    Y = np.array([2.0, 2.0])
    L = np.array([1.0, 1.0])

    epsilon = 0.0
    w = firm.get_w(Y, L, (Z, gamma, epsilon))
    assert (np.allclose(w, np.array([0.5, 0.5])))

    epsilon = 0.5
    w = firm.get_w(Y, L, (Z, gamma, epsilon))
    assert (np.allclose(w, np.array([2.0, 2.0])))

def test_get_Y():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    Z = 2.0
    gamma = 0.5
    L = np.array([4.0, 4.0])
    K = np.array([9.0, 9.0])

    epsilon = 1
    Y = firm.get_Y(K, L, (Z, gamma, epsilon))
    assert (np.allclose(Y, np.array([12.0, 12.0])))

    epsilon = 0
    Y = firm.get_Y(K, L, (Z, gamma, epsilon))
    assert (np.allclose(Y, np.array([13.0, 13.0])))

    epsilon = 0.5
    Z = 2
    L = np.array([1/12.0, 1/12.0])
    K = np.array([1/4.0, 1/4.0])

    Y = firm.get_Y(K, L, (Z, gamma, epsilon))
    assert (np.allclose(Y, np.array([0.5, 0.5])))

def test_get_L():
    """
        Simulate data similar to observed and carry out get_L
        in simplest way possible.
    """
    T = 160
    s, j = 40, 2
    omega = np.random.rand(s).reshape(s, 1)
    lambdas = np.random.rand(j).reshape(1, j)
    n = np.random.rand(T * s * j).reshape(T, s, j)
    e = np.tile(np.random.rand(s, j), (T, 1, 1))

    L_loop = np.ones(T * s * j).reshape(T, s, j)
    for t in range(T):
        for i in range(s):
            for k in range(j):
                L_loop[t, i, k] *= (omega[i, 0] * lambdas[0, k] *
                                    n[t, i, k] * e[t, i, k])

    L_matrix = e * omega * lambdas * n
    assert (np.allclose(L_loop, L_matrix))

    method = 'SS'
    L = firm.get_L(n[0], (e[0], omega, lambdas, method))
    assert (np.allclose(L, L_loop[0].sum()))

    method = 'TPI'
    L = firm.get_L(n, (e, omega, lambdas, method))
    assert (np.allclose(L, L_loop.sum(1).sum(1)))

def test_get_I():
    """
        Simulate data similar to observed and carry out get_I
        in simplest way possible.
    """
    T = 160
    s, j = 40, 2

    b_splus1 = 10 * np.random.rand(T * s * j).reshape(T, s, j)
    K_p1 = 0.9 + np.random.rand(T)
    K = 0.9 + np.random.rand(T)

    delta = np.random.rand()
    g_y = np.random.rand()

    def shifted_arr():
        arr_t = []
        arr_shift_t = []
        for t in range(T):
            arr = np.random.rand(s).reshape(s, 1)
            arr_shift = np.append(arr[1:], [0.0])

            arr_t.append(arr)
            arr_shift_t.append(arr_shift)

        return (np.array(arr_t).reshape(T, s, 1),
                np.array(arr_shift_t).reshape(T, s, 1))

    lambdas = np.random.rand(2)
    imm_rates, imm_shift = shifted_arr()
    imm_rates = imm_rates - 0.5
    imm_shift = imm_shift - 0.5
    omega, omega_shift = shifted_arr()

    g_n = np.random.rand(T)

    res_loop = np.ones(T * s * j).reshape(T, s, j)
    for t in range(T):
        for i in range(s):
            for k in range(j):
                res_loop[t, i, k] *= (omega_shift[t, i, 0] * imm_shift[t, i, 0] *
                                      lambdas[k] * b_splus1[t, i, k])

    res_matrix = (b_splus1 * (imm_shift * omega_shift) * lambdas)

    assert (np.allclose(res_loop, res_matrix))

    aggI_SS_test = ((1 + g_n[0]) * np.exp(g_y) *
                    (K_p1[0] - res_matrix[0].sum() / (1 + g_n[0])) -
                    (1.0 - delta) * K[0])
    aggI_SS = firm.get_I(b_splus1[0], K_p1[0], K[0],
                         (delta, g_y, omega[0], lambdas, imm_rates[0], g_n[0], 'SS'))
    assert (np.allclose(aggI_SS, aggI_SS_test))

    aggI_TPI_test = ((1 + g_n) * np.exp(g_y) *
                     (K_p1 - res_matrix.sum(1).sum(1) / (1 + g_n)) -
                     (1.0 - delta) * K)
    aggI_TPI = firm.get_I(b_splus1, K_p1, K,
                          (delta, g_y, omega, lambdas, imm_rates, g_n, 'TPI'))
    assert (np.allclose(aggI_TPI, aggI_TPI_test))

def test_get_K():
    """
        choose values that simplify the calculations and are similar to
        observed values
    """
    L = np.array([2.0, 2.0])
    r = np.array([1.0, 1.0])
    gamma = 0.5
    tau_b = 0.75
    delta = 0.15
    delta_tau = 0.2 * delta
    Z = 2.0

    epsilon = 1
    K = firm.get_K(L, r, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(K, np.array([0.09832793, 0.09832793])))

    epsilon = 0
    K = firm.get_K(L, r, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(K, np.array([0.0, 0.0])))

    epsilon = 0.5
    Z = 4.0
    tau_b = 0.5
    delta_tau = 2.0
    delta = 3.0
    K = firm.get_K(L, r, (Z, gamma, epsilon, delta, tau_b, delta_tau))
    assert (np.allclose(K, np.array([1.26598632, 1.26598632])))
<|endoftext|>


<commit_before>import json
import copy

from functools import partial

from twisted.web import resource, http, util
from twisted.web.server import NOT_DONE_YET
from twisted.web.guard import HTTPAuthSessionWrapper
from twisted.web.guard import BasicCredentialFactory
from twisted.cred import portal
from twisted.internet.error import ConnectionDone
from twisted.internet.defer import Deferred, inlineCallbacks, returnValue

from vumi import errors
from vumi.blinkenlights import metrics
from vumi.message import TransportUserMessage, TransportEvent
from vumi.errors import InvalidMessage
from vumi import log

from go.apps.http_api.auth import ConversationRealm, ConversationAccessChecker

class BaseResource(resource.Resource):

    def __init__(self, worker, conversation_key):
        resource.Resource.__init__(self)
        self.worker = worker
        self.conversation_key = conversation_key
        self.vumi_api = self.worker.vumi_api
        self.user_apis = {}

    def get_user_api(self, user_account):
        if user_account in self.user_apis:
            return self.user_apis[user_account]

        user_api = self.vumi_api.get_user_api(user_account)
        self.user_apis[user_account] = user_api
        return user_api

    def get_conversation(self, user_account, conversation_key=None):
        conversation_key = conversation_key or self.conversation_key
        user_api = self.get_user_api(user_account)
        return user_api.get_wrapped_conversation(conversation_key)

class StreamResource(BaseResource):

    message_class = None

    def __init__(self, worker, conversation_key):
        BaseResource.__init__(self, worker, conversation_key)
        self.stream_ready = Deferred()
        self.stream_ready.addCallback(self.setup_stream)
        self._callback = None
        self._rk = self.routing_key % {
            'transport_name': self.worker.transport_name,
            'conversation_key': self.conversation_key,
        }

    def render_GET(self, request):
        # Twisted's Agent has trouble closing a connection when the server has
        # sent the HTTP headers but not the body, but sometimes we need to
        # close a connection when only the headers have been received.
        # Sending an empty string as a workaround gets the body consumer
        # stuff started anyway and then we have the ability to close the
        # connection.
        request.write('')
        done = request.notifyFinish()
        done.addBoth(self.teardown_stream)
        self._callback = partial(self.publish, request)
        self.stream_ready.callback(request)
        return NOT_DONE_YET

    def setup_stream(self, request):
        return self.worker.register_client(self._rk, self.message_class,
            self._callback)

    def teardown_stream(self, err):
        if not (err is None or err.trap(ConnectionDone)):
            log.error(err)
        return self.worker.unregister_client(self._rk, self._callback)

    def publish(self, request, message):
        line = u'%s\n' % (message.to_json(),)
        request.write(line.encode('utf-8'))

class InvalidAggregate(errors.VumiError):
    pass

class EventStream(StreamResource):

    message_class = TransportEvent
    routing_key = '%(transport_name)s.stream.event.%(conversation_key)s'

class MessageStream(StreamResource):

    message_class = TransportUserMessage
    routing_key = '%(transport_name)s.stream.message.%(conversation_key)s'

    def render_PUT(self, request):
        d = Deferred()
        d.addCallback(self.handle_PUT)
        d.callback(request)
        return NOT_DONE_YET

    def get_msg_options(self, payload, white_list=[]):
        raw_payload = copy.deepcopy(payload.copy())
        msg_options = dict((key, value)
                           for key, value in raw_payload.items()
                           if key in white_list)
        return msg_options

    def get_conversation_tag(self, conversation):
        return (conversation.delivery_tag_pool, conversation.delivery_tag)

    @inlineCallbacks
    def handle_PUT(self, request):
        try:
            payload = json.loads(request.content.read())
        except ValueError:
            request.setResponseCode(http.BAD_REQUEST, 'Invalid Message')
            request.finish()
            return

        in_reply_to = payload.get('in_reply_to')
        if in_reply_to:
            yield self.handle_PUT_in_reply_to(request, payload, in_reply_to)
        else:
            yield self.handle_PUT_send_to(request, payload)

    @inlineCallbacks
    def handle_PUT_in_reply_to(self, request, payload, in_reply_to):
        user_account = request.getUser()
        user_api = yield self.get_user_api(user_account)
        conversation = yield self.get_conversation(user_account)

        # Using the proxy's load() directly instead of
        # `mdb.get_inbound_message(msg_id)` because that gives us the
        # actual message, not the OutboundMessage. We need the
        # OutboundMessage to get the batch and verify the `user_account`
        reply_to = yield self.vumi_api.mdb.inbound_messages.load(in_reply_to)
        if reply_to is None:
            request.setResponseCode(http.BAD_REQUEST)
            request.write('Invalid in_reply_to value')
            request.finish()
            return

        batch = yield reply_to.batch.get()
        if batch is None or (batch.metadata['user_account']
                             != user_account):
            request.setResponseCode(http.BAD_REQUEST)
            request.write('Invalid in_reply_to value')
            request.finish()
            return

        msg_options = self.get_msg_options(payload,
                                           ['session_event', 'content'])

        # FIXME:    At some point this needs to be done better as it makes some
        #           assumption about how messages are routed which won't be
        #           true for very much longer.
        tag = self.get_conversation_tag(conversation)
        msg_options.update((yield user_api.msg_options(tag)))

        content = msg_options.pop('content')
        continue_session = (msg_options.pop('session_event', None)
                            != TransportUserMessage.SESSION_CLOSE)
        helper_metadata = msg_options.pop('helper_metadata')
        transport_type = msg_options.pop('transport_type')
        from_addr = msg_options.pop('from_addr')

        # NOTE:
        #
        # Not able to use `worker.reply_to()` because TransportUserMessage's
        # reply() method sets the `helper_metadata` which it isn't really
        # supposed to do. Because of this we're doing the manual call
        # to `worker._publish_message()` instead.
        msg = reply_to.msg.reply(content, continue_session, **msg_options)

        # These need to be set manually to ensure that stuff cannot be forged.
        msg['helper_metadata'].update(helper_metadata)
        msg['from_addr'] = from_addr
        msg['transport_type'] = transport_type
        msg['transport_name'] = self.worker.transport_name

        yield self.worker._publish_message(msg)

        request.setResponseCode(http.OK)
        request.write(msg.to_json())
        request.finish()

    @inlineCallbacks
    def handle_PUT_send_to(self, request, payload):
        user_account = request.getUser()
        user_api = yield self.get_user_api(user_account)
        conversation = yield self.get_conversation(user_account)

        msg_options = self.get_msg_options(payload, ['content', 'to_addr'])
        tag = self.get_conversation_tag(conversation)
        msg_options.update((yield user_api.msg_options(tag)))
        to_addr = msg_options.pop('to_addr')
        content = msg_options.pop('content')

        msg = yield self.worker.send_to(to_addr, content, **msg_options)

        request.setResponseCode(http.OK)
        request.write(msg.to_json())
        request.finish()

class MetricResource(BaseResource):

    DEFAULT_STORE_NAME = 'default'

    def render_PUT(self, request):
        d = Deferred()
        d.addCallback(self.handle_PUT)
        d.callback(request)
        return NOT_DONE_YET

    def find_aggregate(self, name):
        agg_class = getattr(metrics, name, None)
        if agg_class is None:
            raise InvalidAggregate('%s is not a valid aggregate.' % (name,))
        return agg_class

    def parse_metrics(self, data):
        metrics = []
        for name, value, aggregate in data:
            value = float(value)
            agg_class = self.find_aggregate(aggregate)
            metrics.append((name, value, agg_class))
        return metrics

    @inlineCallbacks
    def handle_PUT(self, request):
        data = json.loads(request.content.read())
        user_account = request.getUser()

        try:
            metrics = self.parse_metrics(data)
        except (ValueError, InvalidAggregate, InvalidMessage):
            request.setResponseCode(http.BAD_REQUEST, 'Invalid Message')
            request.finish()
            return

        conversation = yield self.get_conversation(user_account)
        metadata = conversation.get_metadata(default={})
        http_api_metadata = metadata.get('http_api', {})
        store = http_api_metadata.get('metrics_store', self.DEFAULT_STORE_NAME)
        for name, value, agg_class in metrics:
            self.worker.publish_account_metric(user_account, store, name,
                                                value, agg_class)

        request.finish()

class ConversationResource(resource.Resource):

    CONCURRENCY_LIMIT = 10

    def __init__(self, worker, conversation_key):
        resource.Resource.__init__(self)
        self.worker = worker
        self.redis = worker.redis
        self.conversation_key = conversation_key

    def key(self, *args):
        return ':'.join(['concurrency'] + map(unicode, args))

    @inlineCallbacks
    def is_allowed(self, user_id):
        count = int((yield self.redis.get(self.key(user_id))) or 0)
        returnValue(count < self.CONCURRENCY_LIMIT)

    def track_request(self, user_id):
        return self.redis.incr(self.key(user_id))

    def release_request(self, err, user_id):
        return self.redis.decr(self.key(user_id))

    def getChild(self, path, request):
        return util.DeferredResource(self.getDeferredChild(path, request))

    @inlineCallbacks
    def getDeferredChild(self, path, request):

        class_map = {
            'events.json': EventStream,
            'messages.json': MessageStream,
            'metrics.json': MetricResource,
        }
        stream_class = class_map.get(path)

        if stream_class is None:
            returnValue(resource.NoResource())

        user_id = request.getUser()
        if (yield self.is_allowed(user_id)):

            # remove track when request is closed
            finished = request.notifyFinish()
            finished.addBoth(self.release_request, user_id)

            yield self.track_request(user_id)
            returnValue(stream_class(self.worker, self.conversation_key))
        returnValue(resource.ErrorPage(http.FORBIDDEN, 'Forbidden',
                                        'Too many concurrent connections'))

class StreamingResource(resource.Resource):

    def __init__(self, worker):
        resource.Resource.__init__(self)
        self.worker = worker

    def getChild(self, conversation_key, request):
        if conversation_key:

            resource = ConversationResource(self.worker, conversation_key)

            checker = ConversationAccessChecker(self.worker.vumi_api,
                                                conversation_key)
            realm = ConversationRealm(resource)
            p = portal.Portal(realm, [checker])

            factory = BasicCredentialFactory("Conversation Stream")
            protected_resource = HTTPAuthSessionWrapper(p, [factory])

            return protected_resource
<commit_msg>preventing transport_name from reaching the reply() call (thanks @jerith)<commit_after>import json
import copy

from functools import partial

from twisted.web import resource, http, util
from twisted.web.server import NOT_DONE_YET
from twisted.web.guard import HTTPAuthSessionWrapper
from twisted.web.guard import BasicCredentialFactory
from twisted.cred import portal
from twisted.internet.error import ConnectionDone
from twisted.internet.defer import Deferred, inlineCallbacks, returnValue

from vumi import errors
from vumi.blinkenlights import metrics
from vumi.message import TransportUserMessage, TransportEvent
from vumi.errors import InvalidMessage
from vumi import log

from go.apps.http_api.auth import ConversationRealm, ConversationAccessChecker

class BaseResource(resource.Resource):

    def __init__(self, worker, conversation_key):
        resource.Resource.__init__(self)
        self.worker = worker
        self.conversation_key = conversation_key
        self.vumi_api = self.worker.vumi_api
        self.user_apis = {}

    def get_user_api(self, user_account):
        if user_account in self.user_apis:
            return self.user_apis[user_account]

        user_api = self.vumi_api.get_user_api(user_account)
        self.user_apis[user_account] = user_api
        return user_api

    def get_conversation(self, user_account, conversation_key=None):
        conversation_key = conversation_key or self.conversation_key
        user_api = self.get_user_api(user_account)
        return user_api.get_wrapped_conversation(conversation_key)

class StreamResource(BaseResource):

    message_class = None

    def __init__(self, worker, conversation_key):
        BaseResource.__init__(self, worker, conversation_key)
        self.stream_ready = Deferred()
        self.stream_ready.addCallback(self.setup_stream)
        self._callback = None
        self._rk = self.routing_key % {
            'transport_name': self.worker.transport_name,
            'conversation_key': self.conversation_key,
        }

    def render_GET(self, request):
        # Twisted's Agent has trouble closing a connection when the server has
        # sent the HTTP headers but not the body, but sometimes we need to
        # close a connection when only the headers have been received.
        # Sending an empty string as a workaround gets the body consumer
        # stuff started anyway and then we have the ability to close the
        # connection.
        request.write('')
        done = request.notifyFinish()
        done.addBoth(self.teardown_stream)
        self._callback = partial(self.publish, request)
        self.stream_ready.callback(request)
        return NOT_DONE_YET

    def setup_stream(self, request):
        return self.worker.register_client(self._rk, self.message_class,
            self._callback)

    def teardown_stream(self, err):
        if not (err is None or err.trap(ConnectionDone)):
            log.error(err)
        return self.worker.unregister_client(self._rk, self._callback)

    def publish(self, request, message):
        line = u'%s\n' % (message.to_json(),)
        request.write(line.encode('utf-8'))

class InvalidAggregate(errors.VumiError):
    pass

class EventStream(StreamResource):

    message_class = TransportEvent
    routing_key = '%(transport_name)s.stream.event.%(conversation_key)s'

class MessageStream(StreamResource):

    message_class = TransportUserMessage
    routing_key = '%(transport_name)s.stream.message.%(conversation_key)s'

    def render_PUT(self, request):
        d = Deferred()
        d.addCallback(self.handle_PUT)
        d.callback(request)
        return NOT_DONE_YET

    def get_msg_options(self, payload, white_list=[]):
        raw_payload = copy.deepcopy(payload.copy())
        msg_options = dict((key, value)
                           for key, value in raw_payload.items()
                           if key in white_list)
        return msg_options

    def get_conversation_tag(self, conversation):
        return (conversation.delivery_tag_pool, conversation.delivery_tag)

    @inlineCallbacks
    def handle_PUT(self, request):
        try:
            payload = json.loads(request.content.read())
        except ValueError:
            request.setResponseCode(http.BAD_REQUEST, 'Invalid Message')
            request.finish()
            return

        in_reply_to = payload.get('in_reply_to')
        if in_reply_to:
            yield self.handle_PUT_in_reply_to(request, payload, in_reply_to)
        else:
            yield self.handle_PUT_send_to(request, payload)

    @inlineCallbacks
    def handle_PUT_in_reply_to(self, request, payload, in_reply_to):
        user_account = request.getUser()
        user_api = yield self.get_user_api(user_account)
        conversation = yield self.get_conversation(user_account)

        # Using the proxy's load() directly instead of
        # `mdb.get_inbound_message(msg_id)` because that gives us the
        # actual message, not the OutboundMessage. We need the
        # OutboundMessage to get the batch and verify the `user_account`
        reply_to = yield self.vumi_api.mdb.inbound_messages.load(in_reply_to)
        if reply_to is None:
            request.setResponseCode(http.BAD_REQUEST)
            request.write('Invalid in_reply_to value')
            request.finish()
            return

        batch = yield reply_to.batch.get()
        if batch is None or (batch.metadata['user_account']
                             != user_account):
            request.setResponseCode(http.BAD_REQUEST)
            request.write('Invalid in_reply_to value')
            request.finish()
            return

        msg_options = self.get_msg_options(payload,
                                           ['session_event', 'content'])

        # FIXME:    At some point this needs to be done better as it makes some
        #           assumption about how messages are routed which won't be
        #           true for very much longer.
        tag = self.get_conversation_tag(conversation)
        msg_options.update((yield user_api.msg_options(tag)))

        content = msg_options.pop('content')
        continue_session = (msg_options.pop('session_event', None)
                            != TransportUserMessage.SESSION_CLOSE)
        helper_metadata = msg_options.pop('helper_metadata')
        transport_type = msg_options.pop('transport_type')
        # We need to pop this out of the msg_options, should it exist
        # because otherwise `msg.reply(...)` is unhappy
        msg_options.pop('transport_name', None)
        from_addr = msg_options.pop('from_addr')

        # NOTE:
        #
        # Not able to use `worker.reply_to()` because TransportUserMessage's
        # reply() method sets the `helper_metadata` which it isn't really
        # supposed to do. Because of this we're doing the manual call
        # to `worker._publish_message()` instead.
        msg = reply_to.msg.reply(content, continue_session, **msg_options)

        # These need to be set manually to ensure that stuff cannot be forged.
        msg['helper_metadata'].update(helper_metadata)
        msg['from_addr'] = from_addr
        msg['transport_type'] = transport_type
        msg['transport_name'] = self.worker.transport_name

        yield self.worker._publish_message(msg)

        request.setResponseCode(http.OK)
        request.write(msg.to_json())
        request.finish()

    @inlineCallbacks
    def handle_PUT_send_to(self, request, payload):
        user_account = request.getUser()
        user_api = yield self.get_user_api(user_account)
        conversation = yield self.get_conversation(user_account)

        msg_options = self.get_msg_options(payload, ['content', 'to_addr'])
        tag = self.get_conversation_tag(conversation)
        msg_options.update((yield user_api.msg_options(tag)))
        to_addr = msg_options.pop('to_addr')
        content = msg_options.pop('content')

        msg = yield self.worker.send_to(to_addr, content, **msg_options)

        request.setResponseCode(http.OK)
        request.write(msg.to_json())
        request.finish()

class MetricResource(BaseResource):

    DEFAULT_STORE_NAME = 'default'

    def render_PUT(self, request):
        d = Deferred()
        d.addCallback(self.handle_PUT)
        d.callback(request)
        return NOT_DONE_YET

    def find_aggregate(self, name):
        agg_class = getattr(metrics, name, None)
        if agg_class is None:
            raise InvalidAggregate('%s is not a valid aggregate.' % (name,))
        return agg_class

    def parse_metrics(self, data):
        metrics = []
        for name, value, aggregate in data:
            value = float(value)
            agg_class = self.find_aggregate(aggregate)
            metrics.append((name, value, agg_class))
        return metrics

    @inlineCallbacks
    def handle_PUT(self, request):
        data = json.loads(request.content.read())
        user_account = request.getUser()

        try:
            metrics = self.parse_metrics(data)
        except (ValueError, InvalidAggregate, InvalidMessage):
            request.setResponseCode(http.BAD_REQUEST, 'Invalid Message')
            request.finish()
            return

        conversation = yield self.get_conversation(user_account)
        metadata = conversation.get_metadata(default={})
        http_api_metadata = metadata.get('http_api', {})
        store = http_api_metadata.get('metrics_store', self.DEFAULT_STORE_NAME)
        for name, value, agg_class in metrics:
            self.worker.publish_account_metric(user_account, store, name,
                                                value, agg_class)

        request.finish()

class ConversationResource(resource.Resource):

    CONCURRENCY_LIMIT = 10

    def __init__(self, worker, conversation_key):
        resource.Resource.__init__(self)
        self.worker = worker
        self.redis = worker.redis
        self.conversation_key = conversation_key

    def key(self, *args):
        return ':'.join(['concurrency'] + map(unicode, args))

    @inlineCallbacks
    def is_allowed(self, user_id):
        count = int((yield self.redis.get(self.key(user_id))) or 0)
        returnValue(count < self.CONCURRENCY_LIMIT)

    def track_request(self, user_id):
        return self.redis.incr(self.key(user_id))

    def release_request(self, err, user_id):
        return self.redis.decr(self.key(user_id))

    def getChild(self, path, request):
        return util.DeferredResource(self.getDeferredChild(path, request))

    @inlineCallbacks
    def getDeferredChild(self, path, request):

        class_map = {
            'events.json': EventStream,
            'messages.json': MessageStream,
            'metrics.json': MetricResource,
        }
        stream_class = class_map.get(path)

        if stream_class is None:
            returnValue(resource.NoResource())

        user_id = request.getUser()
        if (yield self.is_allowed(user_id)):

            # remove track when request is closed
            finished = request.notifyFinish()
            finished.addBoth(self.release_request, user_id)

            yield self.track_request(user_id)
            returnValue(stream_class(self.worker, self.conversation_key))
        returnValue(resource.ErrorPage(http.FORBIDDEN, 'Forbidden',
                                        'Too many concurrent connections'))

class StreamingResource(resource.Resource):

    def __init__(self, worker):
        resource.Resource.__init__(self)
        self.worker = worker

    def getChild(self, conversation_key, request):
        if conversation_key:

            resource = ConversationResource(self.worker, conversation_key)

            checker = ConversationAccessChecker(self.worker.vumi_api,
                                                conversation_key)
            realm = ConversationRealm(resource)
            p = portal.Portal(realm, [checker])

            factory = BasicCredentialFactory("Conversation Stream")
            protected_resource = HTTPAuthSessionWrapper(p, [factory])

            return protected_resource
<|endoftext|>


<commit_before><commit_msg>removing empty file :eyes:<commit_after><|endoftext|>


<commit_before>from rest_framework import serializers

from ..comments.serializers import CommentSerializer
from .models import Card

class CardSerializer(serializers.ModelSerializer):
    thumbnail_sm_path = serializers.Field(source='get_thumbnail_sm_path')
    thumbnail_md_path = serializers.Field(source='get_thumbnail_md_path')
    thumbnail_lg_path = serializers.Field(source='get_thumbnail_lg_path')
    created_by = serializers.PrimaryKeyRelatedField(read_only=True)
    modified_by = serializers.PrimaryKeyRelatedField(read_only=True)

    class Meta:
        model = Card
        read_only_fields = ('slug', 'stack', )

    def validate_content(self, attrs, source):
        content = attrs.get(source)
        card_type = attrs.get('type')

        if card_type != 'stack' and not content:
            raise serializers.ValidationError(self.error_messages['required'])

        return attrs

    def save_object(self, obj, **kwargs):
        created = bool(obj.pk)
        featured_diff = obj.get_field_diff('featured')
        user = self.context['request'].user

        if not created:
            obj.created_by = user

        obj.modified_by = user

        super(CardSerializer, self).save_object(obj, **kwargs)

        if featured_diff and featured_diff[1]:
            obj.notify_featured(user)

class StackSerializer(CardSerializer):
    class Meta:
        model = Card
        read_only_fields = ('slug', )
        exclude = ('origin_url', 'content', 'thumbnail_sm_path',
                   'thumbnail_md_path', 'thumbnail_lg_path', 'file_size',
                   'mime_type', 'stack')

    def validate_cards(self, attrs, source):
        cards = attrs[source]
        card_type = attrs.get('type')

        if card_type != 'stack':
            attrs[source] = []
            return attrs

        request = self.context['request']

        user_cards_ids = request.user.cards.values_list('id', flat=True)

        for card in cards:
            if card.type == 'stack' or card == self.object \
                    or card.id not in user_cards_ids:
                msg = 'Invalid value.'
                raise serializers.ValidationError(msg)

        return attrs

class CardCommentSerializer(CommentSerializer):
    def save_object(self, obj, **kwargs):
        created = bool(obj.pk)
        card = self.context['content_object']
        user = self.context['request'].user

        if not created:
            obj.created_by = user

        obj.modified_by = user
        obj.content_object = card

        super(CardCommentSerializer, self).save_object(obj, **kwargs)

        if not created:
            card.notify_comment_created(user, obj)<commit_msg>Temporarily exclude card extra data field<commit_after>from rest_framework import serializers

from ..comments.serializers import CommentSerializer
from .models import Card

class CardSerializer(serializers.ModelSerializer):
    thumbnail_sm_path = serializers.Field(source='get_thumbnail_sm_path')
    thumbnail_md_path = serializers.Field(source='get_thumbnail_md_path')
    thumbnail_lg_path = serializers.Field(source='get_thumbnail_lg_path')
    created_by = serializers.PrimaryKeyRelatedField(read_only=True)
    modified_by = serializers.PrimaryKeyRelatedField(read_only=True)

    class Meta:
        model = Card
        read_only_fields = ('slug', 'stack', )
        exclude = ('data', )

    def validate_content(self, attrs, source):
        content = attrs.get(source)
        card_type = attrs.get('type')

        if card_type != 'stack' and not content:
            raise serializers.ValidationError(self.error_messages['required'])

        return attrs

    def save_object(self, obj, **kwargs):
        created = bool(obj.pk)
        featured_diff = obj.get_field_diff('featured')
        user = self.context['request'].user

        if not created:
            obj.created_by = user

        obj.modified_by = user

        super(CardSerializer, self).save_object(obj, **kwargs)

        if featured_diff and featured_diff[1]:
            obj.notify_featured(user)

class StackSerializer(CardSerializer):
    class Meta:
        model = Card
        read_only_fields = ('slug', )
        exclude = ('origin_url', 'content', 'thumbnail_sm_path',
                   'thumbnail_md_path', 'thumbnail_lg_path', 'file_size',
                   'mime_type', 'stack')

    def validate_cards(self, attrs, source):
        cards = attrs[source]
        card_type = attrs.get('type')

        if card_type != 'stack':
            attrs[source] = []
            return attrs

        request = self.context['request']

        user_cards_ids = request.user.cards.values_list('id', flat=True)

        for card in cards:
            if card.type == 'stack' or card == self.object \
                    or card.id not in user_cards_ids:
                msg = 'Invalid value.'
                raise serializers.ValidationError(msg)

        return attrs

class CardCommentSerializer(CommentSerializer):
    def save_object(self, obj, **kwargs):
        created = bool(obj.pk)
        card = self.context['content_object']
        user = self.context['request'].user

        if not created:
            obj.created_by = user

        obj.modified_by = user
        obj.content_object = card

        super(CardCommentSerializer, self).save_object(obj, **kwargs)

        if not created:
            card.notify_comment_created(user, obj)<|endoftext|>


<commit_before># Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import collections

from telemetry.core import extension_page
from telemetry.core.backends.chrome import inspector_backend_list

class ExtensionBackendList(inspector_backend_list.InspectorBackendList):
  """A dynamic sequence of extension_page.ExtensionPages."""

  def __init__(self, browser_backend):
    super(ExtensionBackendList, self).__init__(
        browser_backend, backend_wrapper=extension_page.ExtensionPage)

  def ShouldIncludeContext(self, context):
    return context['url'].startswith('chrome-extension://')

class ExtensionBackendDict(collections.Mapping):
  """A dynamic mapping of extension_id to extension_page.ExtensionPages."""

  def __init__(self, browser_backend):
    self._extension_backend_list = ExtensionBackendList(browser_backend)

  def __getitem__(self, extension_id):
    extensions = []
    for i, context_id in enumerate(self._extension_backend_list):
      if self.ContextIdToExtensionId(context_id) == extension_id:
        extensions.append(self._extension_backend_list[i])
    if not extensions:
      raise KeyError('Cannot find an extension with id=%s' % extension_id)
    return extensions

  def __iter__(self):
    for context_id in self._extension_backend_list:
      yield self.ContextIdToExtensionId(context_id)

  def __len__(self):
    return len(self._extension_backend_list)

  def ContextIdToExtensionId(self, context_id):
    context = self._extension_backend_list.GetContextInfo(context_id)
    return extension_page.UrlToExtensionId(context['url'])
<commit_msg>Add __contains__ implementation to avoid a race in __get_item__.<commit_after># Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import collections

from telemetry.core import extension_page
from telemetry.core.backends.chrome import inspector_backend_list

class ExtensionBackendList(inspector_backend_list.InspectorBackendList):
  """A dynamic sequence of extension_page.ExtensionPages."""

  def __init__(self, browser_backend):
    super(ExtensionBackendList, self).__init__(
        browser_backend, backend_wrapper=extension_page.ExtensionPage)

  def ShouldIncludeContext(self, context):
    return context['url'].startswith('chrome-extension://')

class ExtensionBackendDict(collections.Mapping):
  """A dynamic mapping of extension_id to extension_page.ExtensionPages."""

  def __init__(self, browser_backend):
    self._extension_backend_list = ExtensionBackendList(browser_backend)

  def __contains__(self, extension_id):
    return (extension_id in
        (self.ContextIdToExtensionId(context_id)
         for context_id in self._extension_backend_list))

  def __getitem__(self, extension_id):
    extensions = []
    for i, context_id in enumerate(self._extension_backend_list):
      if self.ContextIdToExtensionId(context_id) == extension_id:
        extensions.append(self._extension_backend_list[i])
    if not extensions:
      raise KeyError('Cannot find an extension with id=%s' % extension_id)
    return extensions

  def __iter__(self):
    for context_id in self._extension_backend_list:
      yield self.ContextIdToExtensionId(context_id)

  def __len__(self):
    return len(self._extension_backend_list)

  def ContextIdToExtensionId(self, context_id):
    context = self._extension_backend_list.GetContextInfo(context_id)
    return extension_page.UrlToExtensionId(context['url'])
<|endoftext|>