<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
###
# Copyright (2016-2020) Hewlett Packard Enterprise Development LP
#
# Licensed under the Apache License, Version 2.0 (the "License");
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
###

import mock
import pytest
import yaml

from hpe_test_utils import OneViewBaseTest
from oneview_module_loader import NetworkSetModule

FAKE_MSG_ERROR = 'Fake message error'

NETWORK_SET = dict(
    name='OneViewSDK Test Network Set',
    networkUris=['/rest/ethernet-networks/aaa-bbb-ccc']
)

NETWORK_SET_WITH_NEW_NAME = dict(name='OneViewSDK Test Network Set - Renamed')

CONNECTION_TEMPLATE = dict(bandwidth=dict(maximumBandwidth=15000,
                           typicalBandwidth=3000),
                           name='name828625927-1592834188743',
                           type='connection-template',
                           uri='/rest/connection-templates/aaa-bbb-ccc')

PARAMS_FOR_PRESENT = dict(
    config='config.json',
    state='present',
    data=dict(name=NETWORK_SET['name'],
              networkUris=['/rest/ethernet-networks/aaa-bbb-ccc'])
)

PARAMS_WITH_CHANGES = dict(
    config='config.json',
    state='present',
    data=dict(name=NETWORK_SET['name'],
              newName=NETWORK_SET['name'] + " - Renamed",
              networkUris=['/rest/ethernet-networks/aaa-bbb-ccc', 'Name of a Network'])
)

YAML_PARAMS_WITH_CONNECTION_TEMPLATE = """
    config: "config.json"
    state: present
    data:
        name: 'name828625927-1592834188743'
        type: 'connection-template'
        uri: '/rest/connection-templates/aaa-bbb-ccc'
        bandwidth:
            maximumBandwidth: 15000
            typicalBandwidth: 3000
"""

YAML_PARAMS_WITH_CHANGES = """
    config: "config.json"
    state: present
    data:
      name: 'Test Network Set'
      purpose: Management
      connectionTemplateUri: ~
      bandwidth:
          maximumBandwidth: 3000
          typicalBandwidth: 2000
"""

PARAMS_FOR_ABSENT = dict(
    config='config.json',
    state='absent',
    data=dict(name=NETWORK_SET['name'])
)

DICT_PARAMS_WITH_CHANGES = yaml.load(YAML_PARAMS_WITH_CHANGES)["data"]
DICT_PARAMS_CONNECTION_TEMPLATE = yaml.load(YAML_PARAMS_WITH_CONNECTION_TEMPLATE)["data"]

@pytest.mark.resource(TestNetworkSetModule='network_sets')
class TestNetworkSetModule(OneViewBaseTest):
    """
    OneViewBaseTestCase has common tests for class constructor and main function,
    also provides the mocks used in this test case.
    """

    def test_should_create_new_network_set(self):
        self.resource.get_by_name.return_value = []
        self.resource.create.return_value = self.resource
        self.resource.data = NETWORK_SET

        self.mock_ansible_module.params = PARAMS_FOR_PRESENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_CREATED,
            ansible_facts=dict(network_set=NETWORK_SET)
        )

    def test_should_not_update_when_data_is_equals(self):
        self.resource.data = NETWORK_SET

        self.mock_ansible_module.params = PARAMS_FOR_PRESENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            msg=NetworkSetModule.MSG_ALREADY_PRESENT,
            ansible_facts=dict(network_set=NETWORK_SET)
        )

    def test_update_when_data_has_modified_attributes(self):
        data_merged = dict(name=NETWORK_SET['name'] + " - Renamed",
                           networkUris=['/rest/ethernet-networks/aaa-bbb-ccc',
                                        '/rest/ethernet-networks/ddd-eee-fff']
                           )

        self.resource.data = data_merged

        self.mock_ansible_module.params = PARAMS_WITH_CHANGES

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_UPDATED,
            ansible_facts=dict(network_set=data_merged)
        )

    def test_should_raise_exception_when_ethernet_network_not_found(self):
        self.resource.get_by.side_effect = [NETWORK_SET], []
        self.resource.get_by_name.return_value = None
        self.mock_ov_client.ethernet_networks.get_by.return_value = []
        self.mock_ansible_module.params = PARAMS_WITH_CHANGES.copy()
        self.mock_ansible_module.params['data']['networkUris'] = ['Name of a Network']

        NetworkSetModule().run()

        self.mock_ansible_module.fail_json.assert_called_once_with(
            exception=mock.ANY, msg=NetworkSetModule.MSG_ETHERNET_NETWORK_NOT_FOUND + "Name of a Network")

    def test_should_raise_exception_when_native_ethernet_network_not_found(self):
        self.resource.get_by.side_effect = [NETWORK_SET], []
        self.resource.get_by_name.return_value = None
        self.mock_ov_client.ethernet_networks.get_by.return_value = []
        self.mock_ansible_module.params = PARAMS_WITH_CHANGES.copy()
        self.mock_ansible_module.params['data']['networkUris'] = ['/rest/ethernet-networks/aaa-bbb-ccc']
        self.mock_ansible_module.params['data']['nativeNetworkUri'] = 'Name of a Native Network'

        NetworkSetModule().run()

        self.mock_ansible_module.fail_json.assert_called_once_with(
            exception=mock.ANY, msg=NetworkSetModule.MSG_ETHERNET_NETWORK_NOT_FOUND + "Name of a Native Network")

    def test_should_remove_network(self):
        self.resource.data = NETWORK_SET

        self.mock_ansible_module.params = PARAMS_FOR_ABSENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_DELETED
        )

    def test_update_when_only_bandwidth_has_modified_attributes(self):
        self.resource.data = DICT_PARAMS_WITH_CHANGES
        obj = mock.Mock()
        obj.data = CONNECTION_TEMPLATE
        self.mock_ov_client.connection_templates.get_by_uri.return_value = obj

        self.mock_ansible_module.params = yaml.load(YAML_PARAMS_WITH_CHANGES)

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_UPDATED,
            ansible_facts=dict(network_set=DICT_PARAMS_WITH_CHANGES)
        )

    def test_update_when_data_has_modified_attributes_but_bandwith_is_equal(self):
        self.resource.data = NETWORK_SET
        obj = mock.Mock()
        obj.data = {"bandwidth": DICT_PARAMS_WITH_CHANGES['bandwidth']}
        self.mock_ov_client.connection_templates.get_by_uri.return_value = obj

        self.mock_ansible_module.params = PARAMS_WITH_CHANGES

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_UPDATED,
            ansible_facts=dict(network_set=NETWORK_SET)
        )

    def test_should_not_update_when_connection_template_is_equal(self):
        self.resource.data = DICT_PARAMS_WITH_CHANGES
        obj = mock.Mock()
        obj.data = CONNECTION_TEMPLATE
        self.mock_ov_client.connection_templates.get_by_uri.return_value = obj

        self.mock_ansible_module.params = yaml.load(YAML_PARAMS_WITH_CONNECTION_TEMPLATE)

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            msg=NetworkSetModule.MSG_ALREADY_PRESENT,
            ansible_facts=dict(network_set=DICT_PARAMS_WITH_CHANGES)
        )

    def test_should_fail_when_connection_template_was_not_found(self):
        self.mock_ov_client.connection_templates.get_by_uri.return_value = []

        self.mock_ansible_module.params = yaml.load(YAML_PARAMS_WITH_CONNECTION_TEMPLATE)

        NetworkSetModule().run()

        self.mock_ansible_module.fail_json.assert_called_once_with(
            exception=mock.ANY,
            msg=NetworkSetModule.MSG_CONNECTION_TEMPLATE_NOT_FOUND
        )

    def test_should_do_nothing_when_network_set_not_exist(self):
        self.resource.get_by_name.return_value = None

        self.mock_ansible_module.params = PARAMS_FOR_ABSENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            msg=NetworkSetModule.MSG_ALREADY_ABSENT
        )

    def test_update_scopes_when_different(self):
        params_to_scope = PARAMS_FOR_PRESENT.copy()
        params_to_scope['data']['scopeUris'] = ['test']
        self.mock_ansible_module.params = params_to_scope

        resource_data = NETWORK_SET.copy()
        resource_data['scopeUris'] = ['fake']
        resource_data['uri'] = 'rest/network-sets/fake'
        self.resource.data = resource_data

        patch_return = resource_data.copy()
        patch_return['scopeUris'] = ['test']
        patch_return_obj = self.resource.copy()
        patch_return_obj.data = patch_return
        self.resource.patch.return_value = patch_return_obj
        NetworkSetModule().run()

        self.resource.patch.assert_called_once_with(operation='replace',
                                                    path='/scopeUris',
                                                    value=['test'])

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            ansible_facts=dict(network_set=patch_return),
            msg=NetworkSetModule.MSG_UPDATED
        )

    def test_should_do_nothing_when_scopes_are_the_same(self):
        params_to_scope = PARAMS_FOR_PRESENT.copy()
        params_to_scope['data']['scopeUris'] = ['test']
        self.mock_ansible_module.params = params_to_scope

        resource_data = NETWORK_SET.copy()
        resource_data['scopeUris'] = ['test']
        self.resource.data = resource_data

        NetworkSetModule().run()

        self.resource.patch.not_been_called()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            ansible_facts=dict(network_set=resource_data),
            msg=NetworkSetModule.MSG_ALREADY_PRESENT
        )

if __name__ == '__main__':
    pytest.main([__file__])
<commit_msg>fix typo<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
###
# Copyright (2016-2020) Hewlett Packard Enterprise Development LP
#
# Licensed under the Apache License, Version 2.0 (the "License");
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
###

import mock
import pytest
import yaml

from hpe_test_utils import OneViewBaseTest
from oneview_module_loader import NetworkSetModule

FAKE_MSG_ERROR = 'Fake message error'

NETWORK_SET = dict(
    name='OneViewSDK Test Network Set',
    networkUris=['/rest/ethernet-networks/aaa-bbb-ccc']
)

NETWORK_SET_WITH_NEW_NAME = dict(name='OneViewSDK Test Network Set - Renamed')

CONNECTION_TEMPLATE = dict(bandwidth=dict(maximumBandwidth=15000,
                           typicalBandwidth=3000),
                           name='name828625927-1592834188743',
                           type='connection-template',
                           uri='/rest/connection-templates/aaa-bbb-ccc')

PARAMS_FOR_PRESENT = dict(
    config='config.json',
    state='present',
    data=dict(name=NETWORK_SET['name'],
              networkUris=['/rest/ethernet-networks/aaa-bbb-ccc'])
)

PARAMS_WITH_CHANGES = dict(
    config='config.json',
    state='present',
    data=dict(name=NETWORK_SET['name'],
              newName=NETWORK_SET['name'] + " - Renamed",
              networkUris=['/rest/ethernet-networks/aaa-bbb-ccc', 'Name of a Network'])
)

YAML_PARAMS_WITH_CONNECTION_TEMPLATE = """
    config: "config.json"
    state: present
    data:
        name: 'name828625927-1592834188743'
        type: 'connection-template'
        uri: '/rest/connection-templates/aaa-bbb-ccc'
        bandwidth:
            maximumBandwidth: 15000
            typicalBandwidth: 3000
"""

YAML_PARAMS_WITH_CHANGES = """
    config: "config.json"
    state: present
    data:
      name: 'Test Network Set'
      purpose: Management
      connectionTemplateUri: ~
      bandwidth:
          maximumBandwidth: 3000
          typicalBandwidth: 2000
"""

PARAMS_FOR_ABSENT = dict(
    config='config.json',
    state='absent',
    data=dict(name=NETWORK_SET['name'])
)

DICT_PARAMS_WITH_CHANGES = yaml.load(YAML_PARAMS_WITH_CHANGES)["data"]
DICT_PARAMS_CONNECTION_TEMPLATE = yaml.load(YAML_PARAMS_WITH_CONNECTION_TEMPLATE)["data"]

@pytest.mark.resource(TestNetworkSetModule='network_sets')
class TestNetworkSetModule(OneViewBaseTest):
    """
    OneViewBaseTestCase has common tests for class constructor and main function,
    also provides the mocks used in this test case.
    """

    def test_should_create_new_network_set(self):
        self.resource.get_by_name.return_value = []
        self.resource.create.return_value = self.resource
        self.resource.data = NETWORK_SET

        self.mock_ansible_module.params = PARAMS_FOR_PRESENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_CREATED,
            ansible_facts=dict(network_set=NETWORK_SET)
        )

    def test_should_not_update_when_data_is_equals(self):
        self.resource.data = NETWORK_SET

        self.mock_ansible_module.params = PARAMS_FOR_PRESENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            msg=NetworkSetModule.MSG_ALREADY_PRESENT,
            ansible_facts=dict(network_set=NETWORK_SET)
        )

    def test_update_when_data_has_modified_attributes(self):
        data_merged = dict(name=NETWORK_SET['name'] + " - Renamed",
                           networkUris=['/rest/ethernet-networks/aaa-bbb-ccc',
                                        '/rest/ethernet-networks/ddd-eee-fff']
                           )

        self.resource.data = data_merged

        self.mock_ansible_module.params = PARAMS_WITH_CHANGES

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_UPDATED,
            ansible_facts=dict(network_set=data_merged)
        )

    def test_should_raise_exception_when_ethernet_network_not_found(self):
        self.resource.get_by.side_effect = [NETWORK_SET], []
        self.resource.get_by_name.return_value = None
        self.mock_ov_client.ethernet_networks.get_by.return_value = []
        self.mock_ansible_module.params = PARAMS_WITH_CHANGES.copy()
        self.mock_ansible_module.params['data']['networkUris'] = ['Name of a Network']

        NetworkSetModule().run()

        self.mock_ansible_module.fail_json.assert_called_once_with(
            exception=mock.ANY, msg=NetworkSetModule.MSG_ETHERNET_NETWORK_NOT_FOUND + "Name of a Network")

    def test_should_raise_exception_when_native_ethernet_network_not_found(self):
        self.resource.get_by.side_effect = [NETWORK_SET], []
        self.resource.get_by_name.return_value = None
        self.mock_ov_client.ethernet_networks.get_by.return_value = []
        self.mock_ansible_module.params = PARAMS_WITH_CHANGES.copy()
        self.mock_ansible_module.params['data']['networkUris'] = ['/rest/ethernet-networks/aaa-bbb-ccc']
        self.mock_ansible_module.params['data']['nativeNetworkUri'] = 'Name of a Native Network'

        NetworkSetModule().run()

        self.mock_ansible_module.fail_json.assert_called_once_with(
            exception=mock.ANY, msg=NetworkSetModule.MSG_ETHERNET_NETWORK_NOT_FOUND + "Name of a Native Network")

    def test_should_remove_network(self):
        self.resource.data = NETWORK_SET

        self.mock_ansible_module.params = PARAMS_FOR_ABSENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_DELETED
        )

    def test_update_when_only_bandwidth_has_modified_attributes(self):
        self.resource.data = DICT_PARAMS_WITH_CHANGES
        obj = mock.Mock()
        obj.data = CONNECTION_TEMPLATE
        self.mock_ov_client.connection_templates.get_by_uri.return_value = obj

        self.mock_ansible_module.params = yaml.load(YAML_PARAMS_WITH_CHANGES)

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_UPDATED,
            ansible_facts=dict(network_set=DICT_PARAMS_WITH_CHANGES)
        )

    def test_update_when_data_has_modified_attributes_but_bandwidth_is_equal(self):
        self.resource.data = NETWORK_SET
        obj = mock.Mock()
        obj.data = {"bandwidth": DICT_PARAMS_WITH_CHANGES['bandwidth']}
        self.mock_ov_client.connection_templates.get_by_uri.return_value = obj

        self.mock_ansible_module.params = PARAMS_WITH_CHANGES

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            msg=NetworkSetModule.MSG_UPDATED,
            ansible_facts=dict(network_set=NETWORK_SET)
        )

    def test_should_not_update_when_connection_template_is_equal(self):
        self.resource.data = DICT_PARAMS_WITH_CHANGES
        obj = mock.Mock()
        obj.data = CONNECTION_TEMPLATE
        self.mock_ov_client.connection_templates.get_by_uri.return_value = obj

        self.mock_ansible_module.params = yaml.load(YAML_PARAMS_WITH_CONNECTION_TEMPLATE)

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            msg=NetworkSetModule.MSG_ALREADY_PRESENT,
            ansible_facts=dict(network_set=DICT_PARAMS_WITH_CHANGES)
        )

    def test_should_fail_when_connection_template_was_not_found(self):
        self.mock_ov_client.connection_templates.get_by_uri.return_value = []

        self.mock_ansible_module.params = yaml.load(YAML_PARAMS_WITH_CONNECTION_TEMPLATE)

        NetworkSetModule().run()

        self.mock_ansible_module.fail_json.assert_called_once_with(
            exception=mock.ANY,
            msg=NetworkSetModule.MSG_CONNECTION_TEMPLATE_NOT_FOUND
        )

    def test_should_do_nothing_when_network_set_not_exist(self):
        self.resource.get_by_name.return_value = None

        self.mock_ansible_module.params = PARAMS_FOR_ABSENT

        NetworkSetModule().run()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            msg=NetworkSetModule.MSG_ALREADY_ABSENT
        )

    def test_update_scopes_when_different(self):
        params_to_scope = PARAMS_FOR_PRESENT.copy()
        params_to_scope['data']['scopeUris'] = ['test']
        self.mock_ansible_module.params = params_to_scope

        resource_data = NETWORK_SET.copy()
        resource_data['scopeUris'] = ['fake']
        resource_data['uri'] = 'rest/network-sets/fake'
        self.resource.data = resource_data

        patch_return = resource_data.copy()
        patch_return['scopeUris'] = ['test']
        patch_return_obj = self.resource.copy()
        patch_return_obj.data = patch_return
        self.resource.patch.return_value = patch_return_obj
        NetworkSetModule().run()

        self.resource.patch.assert_called_once_with(operation='replace',
                                                    path='/scopeUris',
                                                    value=['test'])

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=True,
            ansible_facts=dict(network_set=patch_return),
            msg=NetworkSetModule.MSG_UPDATED
        )

    def test_should_do_nothing_when_scopes_are_the_same(self):
        params_to_scope = PARAMS_FOR_PRESENT.copy()
        params_to_scope['data']['scopeUris'] = ['test']
        self.mock_ansible_module.params = params_to_scope

        resource_data = NETWORK_SET.copy()
        resource_data['scopeUris'] = ['test']
        self.resource.data = resource_data

        NetworkSetModule().run()

        self.resource.patch.not_been_called()

        self.mock_ansible_module.exit_json.assert_called_once_with(
            changed=False,
            ansible_facts=dict(network_set=resource_data),
            msg=NetworkSetModule.MSG_ALREADY_PRESENT
        )

if __name__ == '__main__':
    pytest.main([__file__])
<|endoftext|>


<commit_before>from . import heisenberg_dynamics_context
import matplotlib.pyplot as plt
import numpy as np
import vorpy

# Short names identifying the different kinds of plots that can be created
valid_quantity_to_plot_v = [
    'x,y',
    't,z',
    'error(H)',
    'error(J)',
    'sqd',
    'objective',
]

valid_plot_to_include_title_d = {
    'x,y'           : 'curve position in (x,y) only',
    't,z'           : 'curve position in z, time-parameterized',
    'p_x,p_y'       : 'curve momentum in (x,y) only',
    't,p_z'         : 'curve momentum in z, time-parameterized',
    'error(H)'      : 'abs(H)',
    'error(J)'      : 'abs(J - J(t=0))',
    'sqd'           : 'squared distance from initial',
    'objective'     : 'objective function history',
}

class OrbitPlot:
    def __init__ (self, *, curve_description_v, quantity_to_plot_v):
        assert len(frozenset(curve_description_v)) == len(curve_description_v), 'must specify unique values in curve_description_v'
        assert frozenset(quantity_to_plot_v).issubset(frozenset(valid_quantity_to_plot_v)), 'specified invalid elements of quantity_to_plot_v: {0}'.format(','.join(frozenset(quantity_to_plot_v).difference(frozenset(valid_quantity_to_plot_v))))

        self.curve_description_v  = curve_description_v
        self.quantity_to_plot_v  = quantity_to_plot_v

        row_count               = len(curve_description_v)
        col_count               = len(quantity_to_plot_v)

        self.fig,self.axis_vv   = plt.subplots(row_count, col_count, squeeze=False, figsize=(15*col_count,15*row_count))

    def plot_curve_quantity (self, *, axis, curve_description, quantity_to_plot, smo, objective_history_v=None, disable_plot_decoration=False):
        """
        Plot a particular quantity from the curve.  This is called by plot_curve, and you shouldn't
        normally need to call this, unless you want very specific control over the plotting.
        """
        assert quantity_to_plot in valid_quantity_to_plot_v, 'quantity_to_plot (which is {0}) was not found in valid_quantity_to_plot_v (which is {1})'.format(quantity_to_plot, valid_quantity_to_plot_v)

        title = '{0} {1}'.format(curve_description, valid_plot_to_include_title_d[quantity_to_plot])

        if quantity_to_plot == 'x,y':
            flow_curve = smo.flow_curve()
            Q_global_min_index = smo.Q_global_min_index()

            axis.plot(0, 0, 'o', color='black')
            axis.plot(flow_curve[:,0,0], flow_curve[:,0,1])
            axis.plot(flow_curve[0,0,0], flow_curve[0,0,1], 'o', color='green', alpha=0.5)
            # TODO: Plot the interpolated position/momentum (based on smo.t_min() instead of Q_global_min_index)
            if Q_global_min_index is not None and not disable_plot_decoration:
                axis.plot(flow_curve[Q_global_min_index,0,0], flow_curve[Q_global_min_index,0,1], 'o', color='red', alpha=0.5)
            axis.set_aspect('equal')
            # TODO: option to disable axes
        elif quantity_to_plot == 't,z':
            title += ' (red line indicates t_min, which is {0:.10e})'.format(smo.t_min())
            axis.axhline(0, color='black')
            axis.plot(smo.t_v(), smo.flow_curve()[:,0,2])
            axis.axvline(smo.t_min(), color='red')
        elif quantity_to_plot == 'p_x,p_y':
            flow_curve = smo.flow_curve()
            Q_global_min_index = smo.Q_global_min_index()

            axis.plot(0, 0, 'o', color='black')
            axis.plot(flow_curve[:,1,0], flow_curve[:,1,1])
            axis.plot(flow_curve[0,1,0], flow_curve[0,1,1], 'o', color='green', alpha=0.5)
            # TODO: Plot the interpolated position/momentum (based on smo.t_min() instead of Q_global_min_index)
            if Q_global_min_index is not None and not disable_plot_decoration:
                axis.plot(flow_curve[Q_global_min_index,1,0], flow_curve[Q_global_min_index,1,1], 'o', color='red', alpha=0.5)
            axis.set_aspect('equal')
            # TODO: option to disable axes
        elif quantity_to_plot == 't,p_z':
            title += ' (red line indicates t_min, which is {0:.10e})'.format(smo.t_min())
            axis.axhline(0, color='black')
            axis.plot(smo.t_v(), smo.flow_curve()[:,1,2])
            axis.axvline(smo.t_min(), color='red')
        elif quantity_to_plot == 'error(H)':
            H_v = vorpy.apply_along_axes(heisenberg_dynamics_context.Numeric.H, (-2,-1), (smo.flow_curve(),), output_axis_v=(), func_output_shape=())
            abs_H_v = np.abs(H_v)
            title += ' (should stay close to 0)\nmax(abs(H)) = {0:.2e}, H(t=0) = {1:e}'.format(np.max(abs_H_v), H_v[0])
            axis.semilogy(smo.t_v(), abs_H_v)
        elif quantity_to_plot == 'error(J)':
            J_v = vorpy.apply_along_axes(heisenberg_dynamics_context.Numeric.J, (-2,-1), (smo.flow_curve(),), output_axis_v=(), func_output_shape=())
            J_0 = J_v[0]
            J_v -= J_0
            abs_J_minus_J_0 = np.abs(J_v)
            title += ' (should stay close to 0)\nJ(t=0) = {0}; max(abs(J - J(t=0))) = {1:.2e}'.format(J_0, np.max(abs_J_minus_J_0))
            axis.semilogy(smo.t_v(), abs_J_minus_J_0)
        elif quantity_to_plot == 'sqd':
            title += '\nt_min = {0:.10e}, min sqd = {1:.17e}'.format(smo.t_min(), smo.objective())
            axis.semilogy(smo.t_v(), smo.Q_v())
            axis.axvline(smo.t_min(), color='red')
        elif quantity_to_plot == 'objective':
            assert objective_history_v is not None, 'must specify objective_history_v in order to plot {0}'.format(quantity_to_plot)
            title += '\nminimum objective value = {0:.17e}'.format(np.min(objective_history_v))
            axis.semilogy(optimizer.obj_history_v)
        else:
            assert False, 'this should never happen'

        if disable_plot_decoration:
            axis.set_axis_off()
        else:
            axis.set_title(title)

    def plot_curve (self, *, curve_description, smo, objective_history_v=None, disable_plot_decoration=False):
        """
        curve_description specifies which row to plot to (indexed by curve_description_v specified in the constructor).
        The quantity_to_plot_v value specified in the constructor indicates which quantities from smo to plot.
        """

        assert curve_description in self.curve_description_v, 'curve_description (which is {0}) was not found in curve_description_v (which is {1})'.format(curve_description, curve_description_v)

        row_index = self.curve_description_v.index(curve_description)
        axis_v = self.axis_vv[row_index]

        print('plot_curve; curve_description = "{0}"'.format(curve_description))
        assert len(axis_v) == len(self.quantity_to_plot_v)
        for axis,quantity_to_plot in zip(axis_v, self.quantity_to_plot_v):
            plot_it = True
            if quantity_to_plot == 'objective' and objective_history_v is None:
                plot_it = False
            if plot_it:
                print('plotting quantity "{0}" for curve "{1}"'.format(quantity_to_plot, curve_description))
                self.plot_curve_quantity(axis=axis, curve_description=curve_description, quantity_to_plot=quantity_to_plot, smo=smo, objective_history_v=objective_history_v, disable_plot_decoration=disable_plot_decoration)
            else:
                print('NOT plotting quantity "{0}" for curve "{1}"'.format(quantity_to_plot, curve_description))

    def savefig_and_clear (self, *, filename):
        self.fig.tight_layout()
        plt.savefig(filename, bbox_inches='tight')
        print('wrote to file "{0}"'.format(filename))
        # VERY important to do this -- otherwise your memory will slowly fill up!
        # Not sure which one is actually sufficient -- apparently none of them are, YAY!
        plt.clf()
        plt.close(self.fig)
        plt.close('all')
        del self.fig
        del self.axis_vv

<commit_msg>Fixed a copy/paste bug.<commit_after>from . import heisenberg_dynamics_context
import matplotlib.pyplot as plt
import numpy as np
import vorpy

# Short names identifying the different kinds of plots that can be created
valid_quantity_to_plot_v = [
    'x,y',
    't,z',
    'error(H)',
    'error(J)',
    'sqd',
    'objective',
]

valid_plot_to_include_title_d = {
    'x,y'           : 'curve position in (x,y) only',
    't,z'           : 'curve position in z, time-parameterized',
    'p_x,p_y'       : 'curve momentum in (x,y) only',
    't,p_z'         : 'curve momentum in z, time-parameterized',
    'error(H)'      : 'abs(H)',
    'error(J)'      : 'abs(J - J(t=0))',
    'sqd'           : 'squared distance from initial',
    'objective'     : 'objective function history',
}

class OrbitPlot:
    def __init__ (self, *, curve_description_v, quantity_to_plot_v):
        assert len(frozenset(curve_description_v)) == len(curve_description_v), 'must specify unique values in curve_description_v'
        assert frozenset(quantity_to_plot_v).issubset(frozenset(valid_quantity_to_plot_v)), 'specified invalid elements of quantity_to_plot_v: {0}'.format(','.join(frozenset(quantity_to_plot_v).difference(frozenset(valid_quantity_to_plot_v))))

        self.curve_description_v  = curve_description_v
        self.quantity_to_plot_v  = quantity_to_plot_v

        row_count               = len(curve_description_v)
        col_count               = len(quantity_to_plot_v)

        self.fig,self.axis_vv   = plt.subplots(row_count, col_count, squeeze=False, figsize=(15*col_count,15*row_count))

    def plot_curve_quantity (self, *, axis, curve_description, quantity_to_plot, smo, objective_history_v=None, disable_plot_decoration=False):
        """
        Plot a particular quantity from the curve.  This is called by plot_curve, and you shouldn't
        normally need to call this, unless you want very specific control over the plotting.
        """
        assert quantity_to_plot in valid_quantity_to_plot_v, 'quantity_to_plot (which is {0}) was not found in valid_quantity_to_plot_v (which is {1})'.format(quantity_to_plot, valid_quantity_to_plot_v)

        title = '{0} {1}'.format(curve_description, valid_plot_to_include_title_d[quantity_to_plot])

        if quantity_to_plot == 'x,y':
            flow_curve = smo.flow_curve()
            Q_global_min_index = smo.Q_global_min_index()

            axis.plot(0, 0, 'o', color='black')
            axis.plot(flow_curve[:,0,0], flow_curve[:,0,1])
            axis.plot(flow_curve[0,0,0], flow_curve[0,0,1], 'o', color='green', alpha=0.5)
            # TODO: Plot the interpolated position/momentum (based on smo.t_min() instead of Q_global_min_index)
            if Q_global_min_index is not None and not disable_plot_decoration:
                axis.plot(flow_curve[Q_global_min_index,0,0], flow_curve[Q_global_min_index,0,1], 'o', color='red', alpha=0.5)
            axis.set_aspect('equal')
            # TODO: option to disable axes
        elif quantity_to_plot == 't,z':
            title += ' (red line indicates t_min, which is {0:.10e})'.format(smo.t_min())
            axis.axhline(0, color='black')
            axis.plot(smo.t_v(), smo.flow_curve()[:,0,2])
            axis.axvline(smo.t_min(), color='red')
        elif quantity_to_plot == 'p_x,p_y':
            flow_curve = smo.flow_curve()
            Q_global_min_index = smo.Q_global_min_index()

            axis.plot(0, 0, 'o', color='black')
            axis.plot(flow_curve[:,1,0], flow_curve[:,1,1])
            axis.plot(flow_curve[0,1,0], flow_curve[0,1,1], 'o', color='green', alpha=0.5)
            # TODO: Plot the interpolated position/momentum (based on smo.t_min() instead of Q_global_min_index)
            if Q_global_min_index is not None and not disable_plot_decoration:
                axis.plot(flow_curve[Q_global_min_index,1,0], flow_curve[Q_global_min_index,1,1], 'o', color='red', alpha=0.5)
            axis.set_aspect('equal')
            # TODO: option to disable axes
        elif quantity_to_plot == 't,p_z':
            title += ' (red line indicates t_min, which is {0:.10e})'.format(smo.t_min())
            axis.axhline(0, color='black')
            axis.plot(smo.t_v(), smo.flow_curve()[:,1,2])
            axis.axvline(smo.t_min(), color='red')
        elif quantity_to_plot == 'error(H)':
            H_v = vorpy.apply_along_axes(heisenberg_dynamics_context.Numeric.H, (-2,-1), (smo.flow_curve(),), output_axis_v=(), func_output_shape=())
            abs_H_v = np.abs(H_v)
            title += ' (should stay close to 0)\nmax(abs(H)) = {0:.2e}, H(t=0) = {1:e}'.format(np.max(abs_H_v), H_v[0])
            axis.semilogy(smo.t_v(), abs_H_v)
        elif quantity_to_plot == 'error(J)':
            J_v = vorpy.apply_along_axes(heisenberg_dynamics_context.Numeric.J, (-2,-1), (smo.flow_curve(),), output_axis_v=(), func_output_shape=())
            J_0 = J_v[0]
            J_v -= J_0
            abs_J_minus_J_0 = np.abs(J_v)
            title += ' (should stay close to 0)\nJ(t=0) = {0}; max(abs(J - J(t=0))) = {1:.2e}'.format(J_0, np.max(abs_J_minus_J_0))
            axis.semilogy(smo.t_v(), abs_J_minus_J_0)
        elif quantity_to_plot == 'sqd':
            title += '\nt_min = {0:.10e}, min sqd = {1:.17e}'.format(smo.t_min(), smo.objective())
            axis.semilogy(smo.t_v(), smo.Q_v())
            axis.axvline(smo.t_min(), color='red')
        elif quantity_to_plot == 'objective':
            assert objective_history_v is not None, 'must specify objective_history_v in order to plot {0}'.format(quantity_to_plot)
            title += '\nminimum objective value = {0:.17e}'.format(np.min(objective_history_v))
            axis.semilogy(objective_history_v)
        else:
            assert False, 'this should never happen'

        if disable_plot_decoration:
            axis.set_axis_off()
        else:
            axis.set_title(title)

    def plot_curve (self, *, curve_description, smo, objective_history_v=None, disable_plot_decoration=False):
        """
        curve_description specifies which row to plot to (indexed by curve_description_v specified in the constructor).
        The quantity_to_plot_v value specified in the constructor indicates which quantities from smo to plot.
        """

        assert curve_description in self.curve_description_v, 'curve_description (which is {0}) was not found in curve_description_v (which is {1})'.format(curve_description, curve_description_v)

        row_index = self.curve_description_v.index(curve_description)
        axis_v = self.axis_vv[row_index]

        print('plot_curve; curve_description = "{0}"'.format(curve_description))
        assert len(axis_v) == len(self.quantity_to_plot_v)
        for axis,quantity_to_plot in zip(axis_v, self.quantity_to_plot_v):
            plot_it = True
            if quantity_to_plot == 'objective' and objective_history_v is None:
                plot_it = False
            if plot_it:
                print('plotting quantity "{0}" for curve "{1}"'.format(quantity_to_plot, curve_description))
                self.plot_curve_quantity(axis=axis, curve_description=curve_description, quantity_to_plot=quantity_to_plot, smo=smo, objective_history_v=objective_history_v, disable_plot_decoration=disable_plot_decoration)
            else:
                print('NOT plotting quantity "{0}" for curve "{1}"'.format(quantity_to_plot, curve_description))

    def savefig_and_clear (self, *, filename):
        self.fig.tight_layout()
        plt.savefig(filename, bbox_inches='tight')
        print('wrote to file "{0}"'.format(filename))
        # VERY important to do this -- otherwise your memory will slowly fill up!
        # Not sure which one is actually sufficient -- apparently none of them are, YAY!
        plt.clf()
        plt.close(self.fig)
        plt.close('all')
        del self.fig
        del self.axis_vv

<|endoftext|>


<commit_before>import unittest

import six

from example_tests import example_test

class TestBlackScholes(unittest.TestCase):

    def test_black_scholes(self):
        output = example_test.run_example(
            'finance/black_scholes.py', '--n-options', '10')
        six.assertRegex(
            self, output.decode('utf-8'),
            r'initializing...\n' +
            r'start computation\n' +
            r' CPU \(NumPy, Naive implementation\):\t[0-9\.]+ sec\n' +
            r' GPU \(CuPy, Naive implementation\):\t[0-9\.]+ sec\n' +
            r' GPU \(CuPy, Elementwise kernel\):\t[0-9\.]+ sec')

class TestMonteCarlo(unittest.TestCase):

    def test_monte_carlo(self):
        output = example_test.run_example(
            'finance/monte_carlo.py', '--n-options', '10',
            '--n-samples-per-thread', '10',
            '--n-threads-per-option', '10')
        six.assertRegex(
            self, output.decode('utf-8'),
            r'initializing...\n' +
            r'start computation\n' +
            r'    # of options: 10\n' +
            r'    # of samples per option: 100\n' +
            r'GPU \(CuPy, Monte Carlo method\):\t[0-9\.]+ sec\n' +
            r'Error: [0-9\.]+')
<commit_msg>Add Tests<commit_after>import unittest

import six

from example_tests import example_test

class TestBlackScholes(unittest.TestCase):

    def test_black_scholes(self):
        output = example_test.run_example(
            'finance/black_scholes.py', '--n-options', '10')
        six.assertRegex(
            self, output.decode('utf-8'),
            r'initializing...\n' +
            r'start computation\n' +
            r' CPU \(NumPy, Naive implementation\):\t[0-9\.]+ sec\n' +
            r' GPU \(CuPy, Naive implementation\):\t[0-9\.]+ sec\n' +
            r' GPU \(CuPy, Elementwise kernel\):\t[0-9\.]+ sec')

class TestMonteCarlo(unittest.TestCase):

    def test_monte_carlo(self):
        output = example_test.run_example(
            'finance/monte_carlo.py', '--n-options', '10',
            '--n-samples-per-thread', '10',
            '--n-threads-per-option', '10')
        six.assertRegex(
            self, output.decode('utf-8'),
            r'initializing...\n' +
            r'start computation\n' +
            r'    # of options: 10\n' +
            r'    # of samples per option: 100\n' +
            r'GPU \(CuPy, Monte Carlo method\):\t[0-9\.]+ sec\n' +
            r'Error: [0-9\.]+')

class TestMonteCarloWithMultiGPU(unittest.TestCase):

    def test_monte_carlo(self):
        output = example_test.run_example(
            'finance/monte_carlo_multigpu.py', '--n-options', '10',
            '--n-samples-per-thread', '10',
            '--n-threads-per-option', '10')
        six.assertRegex(
            self, output.decode('utf-8'),
            r'initializing...\n' +
            r'start computation\n' +
            r'    # of gpus: 1\n' +
            r'    # of options: 10\n' +
            r'    # of samples per option: 100\n' +
            r'GPU \(CuPy, Monte Carlo method\):\t[0-9\.]+ sec\n' +
            r'Error: [0-9\.]+')
<|endoftext|>


<commit_before>import json
import uuid

from contextlib import contextmanager

from sqlalchemy import and_
from sqlalchemy import String
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy import ForeignKey
from sqlalchemy import create_engine
from sqlalchemy.orm import relationship
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.orm.exc import MultipleResultsFound
from sqlalchemy.ext.declarative import declarative_base

from indexd import index

from indexd.index.driver import IndexDriverABC

from indexd.index.errors import NoRecordFound
from indexd.index.errors import MultipleRecordsFound
from indexd.index.errors import RevisionMismatch

Base = declarative_base()

class IndexRecord(Base):
    '''
    Base index record representation.
    '''
    __tablename__ = 'index_record'

    did = Column(String, primary_key=True)
    rev = Column(String)
    form = Column(String)
    size = Column(Integer)

    urls = relationship('IndexRecordUrl',
        backref='index_record',
        cascade='all, delete-orphan',
    )

    hashes = relationship('IndexRecordHash',
        backref='index_record',
        cascade='all, delete-orphan',
    )

class IndexRecordUrl(Base):
    '''
    Base index record url representation.
    '''
    __tablename__ = 'index_record_url'

    did = Column(String, ForeignKey('index_record.did'), primary_key=True)
    url = Column(String, primary_key=True)

class IndexRecordHash(Base):
    '''
    Base index record hash representation.
    '''
    __tablename__ = 'index_record_hash'

    did = Column(String, ForeignKey('index_record.did'), primary_key=True)
    hash_type = Column(String, primary_key=True)
    hash_value = Column(String)

class SQLAlchemyIndexDriver(IndexDriverABC):
    '''
    SQLAlchemy implementation of index driver.
    '''

    def __init__(self, conn, **config):
        '''
        Initialize the SQLAlchemy database driver.
        '''
        self.engine = create_engine(conn, **config)

        Base.metadata.bind = self.engine
        Base.metadata.create_all()

        self.Session = sessionmaker(bind=self.engine)

    @property
    @contextmanager
    def session(self):
        '''
        Provide a transactional scope around a series of operations.
        '''
        session = self.Session()

        yield session

        try: session.commit()
        except:
            session.rollback()
            raise
        finally:
            session.close()

    def ids(self, limit=100, start=None, size=None, urls=None, hashes=None):
        '''
        Returns list of records stored by the backend.
        '''
        # TODO add dids to filter on
        with self.session as session:
            query = session.query(IndexRecord)

            if start is not None:
                query = query.filter(IndexRecord.did > start)

            if size is not None:
                query = query.filter(IndexRecord.size == size)

            if urls is not None:
                query = query.join(IndexRecord.urls)
                for u in urls:
                    query = query.filter(IndexRecordUrl.url == u)

            if hashes is not None:
                query = query.join(IndexRecord.hashes)
                for h,v in hashes.items():
                    query = query.filter(and_(
                        IndexRecordHash.hash_type == h,
                        IndexRecordHash.hash_value == v,
                    ))

            query = query.limit(limit)

            return [i.did for i in query]

    def add(self, form, size=None, urls=[], hashes={}):
        '''
        Creates a new record given urls and hashes.
        '''
        if form not in index.FORMS:
            raise ValueError('form must be one of: %s' % index.FORMS)

        if size is not None and size < 0:
            raise ValueError('size must be non-negative')

        with self.session as session:
            record = IndexRecord()

            record.did = str(uuid.uuid4())
            record.rev = str(uuid.uuid4())[:8]

            record.form = form
            record.size = size

            record.urls = [IndexRecordUrl(
                did=record,
                url=url,
            ) for url in urls]

            record.hashes = [IndexRecordHash(
                did=record,
                hash_type=h,
                hash_value=v,
            ) for h,v in hashes.items()]

            session.add(record)

            return record.did, record.rev

    def get(self, did):
        '''
        Gets a record given the record id.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == did)

            try: record = query.one()
            except NoResultFound as err:
                raise NoRecordFound('no record found')
            except MultipleResultsFound as err:
                raise MultipleRecordsFound('multiple records found')

            rev = record.rev

            form = record.form
            size = record.size

            urls = [u.url for u in record.urls]
            hashes = {h.hash_type: h.hash_value for h in record.hashes}

        ret = {
            'did': did,
            'rev': rev,
            'size': size,
            'urls': urls,
            'hashes': hashes,
            'form': form,
        }

        return ret

    def update(self, did, rev, size=None, urls=None, hashes=None):
        '''
        Updates an existing record with new values.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == did)

            try: record = query.one()
            except NoResultFound as err:
                raise NoRecordFound('no record found')
            except MultipleResultsFound as err:
                raise MultipleRecordsFound('multiple records found')

            if rev != record.rev:
                raise RevisionMismatch('revision mismatch')

            if size is not None:
                record.size = size

            if urls is not None:
                record.urls = [IndexRecordUrl(
                    did=record,
                    url=url
                ) for url in urls]

            if hashes is not None:
                record.hashes = [IndexRecordHash(
                    did=record,
                    hash_type=h,
                    hash_value=v,
                ) for h,v in hashes.items()]

            record.rev = str(uuid.uuid4())[:8]

            session.add(record)

            return record.did, record.rev

    def delete(self, did, rev):
        '''
        Removes record if stored by backend.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == did)

            try: record = query.one()
            except NoResultFound as err:
                raise NoRecordFound('no record found')
            except MultipleResultsFound as err:
                raise MultipleRecordsFound('multiple records found')

            if rev != record.rev:
                raise RevisionMismatch('revision mismatch')

            session.delete(record)

    def __contains__(self, record):
        '''
        Returns True if record is stored by backend.
        Returns False otherwise.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == record)

            return query.exists()

    def __iter__(self):
        '''
        Iterator over unique records stored by backend.
        '''
        with self.session as session:
            for i in session.query(IndexRecord):
                yield i.did

    def __len__(self):
        '''
        Number of unique records stored by backend.
        '''
        with self.session as session:
            return session.query(IndexRecord).count()
<commit_msg>fixes conditional urls / hashes parsing<commit_after>import json
import uuid

from contextlib import contextmanager

from sqlalchemy import and_
from sqlalchemy import String
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy import ForeignKey
from sqlalchemy import create_engine
from sqlalchemy.orm import relationship
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.orm.exc import MultipleResultsFound
from sqlalchemy.ext.declarative import declarative_base

from indexd import index

from indexd.index.driver import IndexDriverABC

from indexd.index.errors import NoRecordFound
from indexd.index.errors import MultipleRecordsFound
from indexd.index.errors import RevisionMismatch

Base = declarative_base()

class IndexRecord(Base):
    '''
    Base index record representation.
    '''
    __tablename__ = 'index_record'

    did = Column(String, primary_key=True)
    rev = Column(String)
    form = Column(String)
    size = Column(Integer)

    urls = relationship('IndexRecordUrl',
        backref='index_record',
        cascade='all, delete-orphan',
    )

    hashes = relationship('IndexRecordHash',
        backref='index_record',
        cascade='all, delete-orphan',
    )

class IndexRecordUrl(Base):
    '''
    Base index record url representation.
    '''
    __tablename__ = 'index_record_url'

    did = Column(String, ForeignKey('index_record.did'), primary_key=True)
    url = Column(String, primary_key=True)

class IndexRecordHash(Base):
    '''
    Base index record hash representation.
    '''
    __tablename__ = 'index_record_hash'

    did = Column(String, ForeignKey('index_record.did'), primary_key=True)
    hash_type = Column(String, primary_key=True)
    hash_value = Column(String)

class SQLAlchemyIndexDriver(IndexDriverABC):
    '''
    SQLAlchemy implementation of index driver.
    '''

    def __init__(self, conn, **config):
        '''
        Initialize the SQLAlchemy database driver.
        '''
        self.engine = create_engine(conn, **config)

        Base.metadata.bind = self.engine
        Base.metadata.create_all()

        self.Session = sessionmaker(bind=self.engine)

    @property
    @contextmanager
    def session(self):
        '''
        Provide a transactional scope around a series of operations.
        '''
        session = self.Session()

        yield session

        try: session.commit()
        except:
            session.rollback()
            raise
        finally:
            session.close()

    def ids(self, limit=100, start=None, size=None, urls=None, hashes=None):
        '''
        Returns list of records stored by the backend.
        '''
        # TODO add dids to filter on
        with self.session as session:
            query = session.query(IndexRecord)

            if start is not None:
                query = query.filter(IndexRecord.did > start)

            if size is not None:
                query = query.filter(IndexRecord.size == size)

            if urls is not None and urls:
                query = query.join(IndexRecord.urls)
                for u in urls:
                    query = query.filter(IndexRecordUrl.url == u)

            if hashes is not None and hashes:
                query = query.join(IndexRecord.hashes)
                for h,v in hashes.items():
                    query = query.filter(and_(
                        IndexRecordHash.hash_type == h,
                        IndexRecordHash.hash_value == v,
                    ))

            query = query.limit(limit)

            return [i.did for i in query]

    def add(self, form, size=None, urls=[], hashes={}):
        '''
        Creates a new record given urls and hashes.
        '''
        if form not in index.FORMS:
            raise ValueError('form must be one of: %s' % index.FORMS)

        if size is not None and size < 0:
            raise ValueError('size must be non-negative')

        with self.session as session:
            record = IndexRecord()

            record.did = str(uuid.uuid4())
            record.rev = str(uuid.uuid4())[:8]

            record.form = form
            record.size = size

            record.urls = [IndexRecordUrl(
                did=record,
                url=url,
            ) for url in urls]

            record.hashes = [IndexRecordHash(
                did=record,
                hash_type=h,
                hash_value=v,
            ) for h,v in hashes.items()]

            session.add(record)

            return record.did, record.rev

    def get(self, did):
        '''
        Gets a record given the record id.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == did)

            try: record = query.one()
            except NoResultFound as err:
                raise NoRecordFound('no record found')
            except MultipleResultsFound as err:
                raise MultipleRecordsFound('multiple records found')

            rev = record.rev

            form = record.form
            size = record.size

            urls = [u.url for u in record.urls]
            hashes = {h.hash_type: h.hash_value for h in record.hashes}

        ret = {
            'did': did,
            'rev': rev,
            'size': size,
            'urls': urls,
            'hashes': hashes,
            'form': form,
        }

        return ret

    def update(self, did, rev, size=None, urls=None, hashes=None):
        '''
        Updates an existing record with new values.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == did)

            try: record = query.one()
            except NoResultFound as err:
                raise NoRecordFound('no record found')
            except MultipleResultsFound as err:
                raise MultipleRecordsFound('multiple records found')

            if rev != record.rev:
                raise RevisionMismatch('revision mismatch')

            if size is not None:
                record.size = size

            if urls is not None:
                record.urls = [IndexRecordUrl(
                    did=record,
                    url=url
                ) for url in urls]

            if hashes is not None:
                record.hashes = [IndexRecordHash(
                    did=record,
                    hash_type=h,
                    hash_value=v,
                ) for h,v in hashes.items()]

            record.rev = str(uuid.uuid4())[:8]

            session.add(record)

            return record.did, record.rev

    def delete(self, did, rev):
        '''
        Removes record if stored by backend.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == did)

            try: record = query.one()
            except NoResultFound as err:
                raise NoRecordFound('no record found')
            except MultipleResultsFound as err:
                raise MultipleRecordsFound('multiple records found')

            if rev != record.rev:
                raise RevisionMismatch('revision mismatch')

            session.delete(record)

    def __contains__(self, record):
        '''
        Returns True if record is stored by backend.
        Returns False otherwise.
        '''
        with self.session as session:
            query = session.query(IndexRecord)
            query = query.filter(IndexRecord.did == record)

            return query.exists()

    def __iter__(self):
        '''
        Iterator over unique records stored by backend.
        '''
        with self.session as session:
            for i in session.query(IndexRecord):
                yield i.did

    def __len__(self):
        '''
        Number of unique records stored by backend.
        '''
        with self.session as session:
            return session.query(IndexRecord).count()
<|endoftext|>


<commit_before>import json

from flask import Flask, Blueprint
from .view_classes import BasicView, IndexView, JSONifyTestView
from nose.tools import *

app = Flask("blueprints")
bp = Blueprint("bptest", "bptest")
BasicView.register(bp)
IndexView.register(bp)
JSONifyTestView.register(bp)
app.register_blueprint(bp)

client = app.test_client()

def test_bp_index():
    resp = client.get("/basic/")
    eq_(b"Index", resp.data)

def test_bp_get():
    resp = client.get("/basic/1234")
    eq_(b"Get 1234", resp.data)

def test_bp_put():
    resp = client.put("/basic/1234")
    eq_(b"Put 1234", resp.data)

def test_bp_patch():
    resp = client.patch("/basic/1234")
    eq_(b"Patch 1234", resp.data)

def test_bp_post():
    resp = client.post("/basic/")
    eq_(b"Post", resp.data)

def test_bp_delete():
    resp = client.delete("/basic/1234")
    eq_(b"Delete 1234", resp.data)

def test_bp_custom_method():
    resp = client.get("/basic/custom_method/")
    eq_(b"Custom Method", resp.data)

def test_bp_custom_method_with_params():
    resp = client.get("/basic/custom_method_with_params/1234/abcd")
    eq_(b"Custom Method 1234 abcd", resp.data)

def test_bp_routed_method():
    resp = client.get("/basic/routed/")
    eq_(b"Routed Method", resp.data)

def test_bp_multi_routed_method():
    resp = client.get("/basic/route1/")
    eq_(b"Multi Routed Method", resp.data)

    resp = client.get("/basic/route2/")
    eq_(b"Multi Routed Method", resp.data)

def test_bp_no_slash():
    resp = client.get("/basic/noslash")
    eq_(b"No Slash Method", resp.data)

def test_bp_index_view_index():
    resp = client.get("/")
    eq_(b"Index", resp.data)

def test_bp_custom_http_method():
    resp = client.post("/basic/route3/")
    eq_(b"Custom HTTP Method", resp.data)

def test_bp_url_prefix():
    foo = Blueprint('foo', __name__)
    BasicView.register(foo, route_base="/")
    app.register_blueprint(foo, url_prefix='/foo')

    resp = client.get('/foo/')
    eq_(b"Index", resp.data)

def test_jsonify_normal_index():
    resp = client.get('/jsonify')
    eq_(resp.status_code, 200)
    eq_(json.loads(resp.data), dict(success=True))

def test_jsonify_post_custom_status_code():
    resp = client.post('/jsonify')
    eq_(resp.status_code, 201)
    eq_(json.loads(resp.data), dict(success=True))

def test_jsonify_not_found():
    resp = client.get('/jsonify/not-found')
    eq_(resp.status_code, 404)
    eq_(json.loads(resp.data), dict(success=False))

def test_custom_header():
    resp = client.get('/jsonify/custom-header')
    eq_(resp.status_code, 418)
    eq_(resp.headers['X-TEAPOT'], '1')
    eq_(json.loads(resp.data), dict(success=True))

def test_normal_jsonify():
    resp = client.get('/jsonify/normal')
    eq_(resp.status_code, 200)
    eq_(resp.headers != None, True)
    eq_(json.loads(resp.data), dict(success=True))
<commit_msg>@ #19 | should allows flask.jsonify to return status codes and headers: fix failing tests on Python 3.x<commit_after>import json

from flask import Flask, Blueprint
from .view_classes import BasicView, IndexView, JSONifyTestView
from nose.tools import *

app = Flask("blueprints")
bp = Blueprint("bptest", "bptest")
BasicView.register(bp)
IndexView.register(bp)
JSONifyTestView.register(bp)
app.register_blueprint(bp)

client = app.test_client()

def test_bp_index():
    resp = client.get("/basic/")
    eq_(b"Index", resp.data)

def test_bp_get():
    resp = client.get("/basic/1234")
    eq_(b"Get 1234", resp.data)

def test_bp_put():
    resp = client.put("/basic/1234")
    eq_(b"Put 1234", resp.data)

def test_bp_patch():
    resp = client.patch("/basic/1234")
    eq_(b"Patch 1234", resp.data)

def test_bp_post():
    resp = client.post("/basic/")
    eq_(b"Post", resp.data)

def test_bp_delete():
    resp = client.delete("/basic/1234")
    eq_(b"Delete 1234", resp.data)

def test_bp_custom_method():
    resp = client.get("/basic/custom_method/")
    eq_(b"Custom Method", resp.data)

def test_bp_custom_method_with_params():
    resp = client.get("/basic/custom_method_with_params/1234/abcd")
    eq_(b"Custom Method 1234 abcd", resp.data)

def test_bp_routed_method():
    resp = client.get("/basic/routed/")
    eq_(b"Routed Method", resp.data)

def test_bp_multi_routed_method():
    resp = client.get("/basic/route1/")
    eq_(b"Multi Routed Method", resp.data)

    resp = client.get("/basic/route2/")
    eq_(b"Multi Routed Method", resp.data)

def test_bp_no_slash():
    resp = client.get("/basic/noslash")
    eq_(b"No Slash Method", resp.data)

def test_bp_index_view_index():
    resp = client.get("/")
    eq_(b"Index", resp.data)

def test_bp_custom_http_method():
    resp = client.post("/basic/route3/")
    eq_(b"Custom HTTP Method", resp.data)

def test_bp_url_prefix():
    foo = Blueprint('foo', __name__)
    BasicView.register(foo, route_base="/")
    app.register_blueprint(foo, url_prefix='/foo')

    resp = client.get('/foo/')
    eq_(b"Index", resp.data)

def test_jsonify_normal_index():
    resp = client.get('/jsonify')
    eq_(resp.status_code, 200)
    eq_(json.loads(resp.data.decode('utf-8')), dict(success=True))

def test_jsonify_post_custom_status_code():
    resp = client.post('/jsonify')
    eq_(resp.status_code, 201)
    eq_(json.loads(resp.data.decode('utf-8')), dict(success=True))

def test_jsonify_not_found():
    resp = client.get('/jsonify/not-found')
    eq_(resp.status_code, 404)
    eq_(json.loads(resp.data.decode('utf-8')), dict(success=False))

def test_custom_header():
    resp = client.get('/jsonify/custom-header')
    eq_(resp.status_code, 418)
    eq_(resp.headers['X-TEAPOT'], '1')
    eq_(json.loads(resp.data.decode('utf-8')), dict(success=True))

def test_normal_jsonify():
    resp = client.get('/jsonify/normal')
    eq_(resp.status_code, 200)
    eq_(resp.headers is not None, True)
    eq_(json.loads(resp.data.decode('utf-8')), dict(success=True))
<|endoftext|>


<commit_before># Licensed Materials - Property of IBM
# Copyright IBM Corp. 2015

from streamsx.topology import graph

class Topology(object):
    """Topology that contains graph + operators"""
    def __init__(self, name, files=None):
        self.name=name
        self.graph = graph.SPLGraph(name)
        self.files = []

    def source(self, func):
        """
        Fetches information from an external system and presents that information as a stream.
        Takes a zero-argument callable that returns an iterable of tuples.
        Each tuple that is not None from the iterator returned
        from iter(func()) is present on the returned stream.
        :param func: A zero-argument callable that returns an iterable of tuples.
        The callable must be either
        * the name of a function defined at the top level of a module that takes no arguments, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self)`.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        A tuple is represented as a Python object that must be picklable.
        :return: A stream whose tuples are the result of the output obtained by invoking the provided callable.
        """
        op = self.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionSource", func)
        oport = op.addOutputPort()
        return Stream(self, oport)

    def subscribe(self, topic, schema):
        op = self.graph.addOperator(kind=schema.subscribeOp())
        oport = op.addOutputPort(schema=schema)
        topicParam = {"topic": [topic]}
        op.setParameters(topicParam)
        return Stream(self, oport)

class Stream(object):
    """
    Definition of a data stream in python.
    """
    def __init__(self, topology, oport):
        self.topology = topology
        self.oport = oport

    def sink(self, func):
        """
        Sends information as a stream to an external system.
        Takes a user provided callable that does not return a value.
        :param func: A callable that takes a single parameter for the tuple and returns None.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)`.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: None
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionSink", func)
        op.addInputPort(outputPort=self.oport)

    def filter(self, func):
        """
        Filters tuples from a stream using the supplied callable `func`.
        For each tuple on the stream the callable is called passing
        the tuple, if the callable return evalulates to true the
        tuple will be present on the returned stream, otherwise
        the tuple is filtered out.
        :param func: A callable that takes a single parameter for the tuple, and returns True or False.
        If True, the tuple is included on the returned stream.  If False, the tuple is filtered out.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)`.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: A stream containing tuples that have not been filtered out.
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionFilter", func)
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def transform(self, func):
        """
        Transforms each tuple from this stream into 0 or 1 tuples using the supplied callable `func`.
        For each tuple on this stream, the returned stream will contain a tuple
        that is the result of the callable when the return is not None.
        If the callable returns None then no tuple is submitted to the returned
        stream.
        :param func: A callable that takes a single parameter for the tuple, and returns a tuple or None.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)`.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: A stream containing transformed tuples.
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionTransform", func)
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def multi_transform(self, func):
        """
        Transforms each tuple from this stream into 0 or more tuples using the supplied callable `func`.
        For each tuple on this stream, the returned stream will contain all non-None tuples from
        the iterable.
        Tuples will be added to the returned stream in the order the iterable
        returns them.
        If the return is None or an empty iterable then no tuples are added to
        the returned stream.
        :param func: A callable that takes a single parameter for the tuple, and returns an iterable of tuples or None.
        The callable must return an iterable or None, otherwise a TypeError is raised.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)`.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: A stream containing transformed tuples.
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionMultiTransform", func)
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def isolate(self):
        """
        Guarantees that the upstream operation will run in a separate process from the downstream operation
        :param: None
        :return: None
        """
        op = self.topology.graph.addOperator("$Isolate$")
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def low_latency(self):
        """
	      The function is guaranteed to run in the same process as the
        upstream Stream function. All streams that are created from the returned stream
        are also guaranteed to run in the same process until end_low_latency()
        is called.
        """
        op = self.topology.graph.addOperator("$LowLatency$")
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def end_low_latency(self):
        """
        Return a Stream that is no longer guaranteed to run in the same process
        as the calling stream.
        """
        op = self.topology.graph.addOperator("$EndLowLatency$")
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def parallel(self, width):
        """
        Marks operator as parallel output with width
        :param: width
        :returns: Stream
        """
        iop = self.isolate()

        op2 = self.topology.graph.addOperator("$Parallel$")
        op2.addInputPort(outputPort=iop.getOport())
        oport = op2.addOutputPort(width)
        return Stream(self.topology, oport)

    def end_parallel(self):
        """
        Marks end of operators as parallel output with width
        :returns: Stream
        """
        lastOp = self.topology.graph.getLastOperator()
        outport = self.oport
        if (isinstance(lastOp, graph.Marker)):
            if (lastOp.kind == "$Union$"):
                pto = self.topology.graph.addPassThruOperator()
                pto.addInputPort(outputPort=self.oport)
                outport = pto.addOutputPort()
        op = self.topology.graph.addOperator("$EndParallel$")
        op.addInputPort(outputPort=outport)
        oport = op.addOutputPort()
        endP = Stream(self.topology, oport)
        return endP.isolate()

    def union(self, streamSet):
        """
        The Union operator merges the outputs of the streams in the set
        into a single stream.
        :param set of streams outputs to merge
        :returns Stream
        """
        if(not isinstance(streamSet,set)) :
            raise TypeError("The union operator parameter must be a set object")
        if(len(streamSet) == 0):
            return self
        op = self.topology.graph.addOperator("$Union$")
        op.addInputPort(outputPort=self.getOport())
        for stream in streamSet:
            op.addInputPort(outputPort=stream.getOport())
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def getOport(self):
        return self.oport
<commit_msg>update comments for instances of callable classes<commit_after># Licensed Materials - Property of IBM
# Copyright IBM Corp. 2015

from streamsx.topology import graph

class Topology(object):
    """Topology that contains graph + operators"""
    def __init__(self, name, files=None):
        self.name=name
        self.graph = graph.SPLGraph(name)
        self.files = []

    def source(self, func):
        """
        Fetches information from an external system and presents that information as a stream.
        Takes a zero-argument callable that returns an iterable of tuples.
        Each tuple that is not None from the iterator returned
        from iter(func()) is present on the returned stream.
        :param func: A zero-argument callable that returns an iterable of tuples.
        The callable must be either
        * the name of a function defined at the top level of a module that takes no arguments, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self)` and is picklable.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        A tuple is represented as a Python object that must be picklable.
        :return: A stream whose tuples are the result of the output obtained by invoking the provided callable.
        """
        op = self.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionSource", func)
        oport = op.addOutputPort()
        return Stream(self, oport)

    def subscribe(self, topic, schema):
        op = self.graph.addOperator(kind=schema.subscribeOp())
        oport = op.addOutputPort(schema=schema)
        topicParam = {"topic": [topic]}
        op.setParameters(topicParam)
        return Stream(self, oport)

class Stream(object):
    """
    Definition of a data stream in python.
    """
    def __init__(self, topology, oport):
        self.topology = topology
        self.oport = oport

    def sink(self, func):
        """
        Sends information as a stream to an external system.
        Takes a user provided callable that does not return a value.
        :param func: A callable that takes a single parameter for the tuple and returns None.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)` and is picklable.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: None
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionSink", func)
        op.addInputPort(outputPort=self.oport)

    def filter(self, func):
        """
        Filters tuples from a stream using the supplied callable `func`.
        For each tuple on the stream the callable is called passing
        the tuple, if the callable return evalulates to true the
        tuple will be present on the returned stream, otherwise
        the tuple is filtered out.
        :param func: A callable that takes a single parameter for the tuple, and returns True or False.
        If True, the tuple is included on the returned stream.  If False, the tuple is filtered out.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)` and is picklable.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: A stream containing tuples that have not been filtered out.
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionFilter", func)
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def transform(self, func):
        """
        Transforms each tuple from this stream into 0 or 1 tuples using the supplied callable `func`.
        For each tuple on this stream, the returned stream will contain a tuple
        that is the result of the callable when the return is not None.
        If the callable returns None then no tuple is submitted to the returned
        stream.
        :param func: A callable that takes a single parameter for the tuple, and returns a tuple or None.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)` and is picklable.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: A stream containing transformed tuples.
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionTransform", func)
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def multi_transform(self, func):
        """
        Transforms each tuple from this stream into 0 or more tuples using the supplied callable `func`.
        For each tuple on this stream, the returned stream will contain all non-None tuples from
        the iterable.
        Tuples will be added to the returned stream in the order the iterable
        returns them.
        If the return is None or an empty iterable then no tuples are added to
        the returned stream.
        :param func: A callable that takes a single parameter for the tuple, and returns an iterable of tuples or None.
        The callable must return an iterable or None, otherwise a TypeError is raised.
        The callable must be either
        * the name of a function defined at the top level of a module that takes a single parameter for the tuple, or
        * an instance of a callable class defined at the top level of a module that implements the method `__call__(self, tuple)` and is picklable.
        Using a callable class allows state information such as user-defined parameters to be stored during class
        initialization and utilized when the instance is called.
        The callable is invoked for each incoming tuple.
        :return: A stream containing transformed tuples.
        """
        op = self.topology.graph.addOperator("com.ibm.streamsx.topology.functional.python::PyFunctionMultiTransform", func)
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def isolate(self):
        """
        Guarantees that the upstream operation will run in a separate process from the downstream operation
        :param: None
        :return: None
        """
        op = self.topology.graph.addOperator("$Isolate$")
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def low_latency(self):
        """
	      The function is guaranteed to run in the same process as the
        upstream Stream function. All streams that are created from the returned stream
        are also guaranteed to run in the same process until end_low_latency()
        is called.
        """
        op = self.topology.graph.addOperator("$LowLatency$")
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def end_low_latency(self):
        """
        Return a Stream that is no longer guaranteed to run in the same process
        as the calling stream.
        """
        op = self.topology.graph.addOperator("$EndLowLatency$")
        op.addInputPort(outputPort=self.oport)
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def parallel(self, width):
        """
        Marks operator as parallel output with width
        :param: width
        :returns: Stream
        """
        iop = self.isolate()

        op2 = self.topology.graph.addOperator("$Parallel$")
        op2.addInputPort(outputPort=iop.getOport())
        oport = op2.addOutputPort(width)
        return Stream(self.topology, oport)

    def end_parallel(self):
        """
        Marks end of operators as parallel output with width
        :returns: Stream
        """
        lastOp = self.topology.graph.getLastOperator()
        outport = self.oport
        if (isinstance(lastOp, graph.Marker)):
            if (lastOp.kind == "$Union$"):
                pto = self.topology.graph.addPassThruOperator()
                pto.addInputPort(outputPort=self.oport)
                outport = pto.addOutputPort()
        op = self.topology.graph.addOperator("$EndParallel$")
        op.addInputPort(outputPort=outport)
        oport = op.addOutputPort()
        endP = Stream(self.topology, oport)
        return endP.isolate()

    def union(self, streamSet):
        """
        The Union operator merges the outputs of the streams in the set
        into a single stream.
        :param set of streams outputs to merge
        :returns Stream
        """
        if(not isinstance(streamSet,set)) :
            raise TypeError("The union operator parameter must be a set object")
        if(len(streamSet) == 0):
            return self
        op = self.topology.graph.addOperator("$Union$")
        op.addInputPort(outputPort=self.getOport())
        for stream in streamSet:
            op.addInputPort(outputPort=stream.getOport())
        oport = op.addOutputPort()
        return Stream(self.topology, oport)

    def getOport(self):
        return self.oport
<|endoftext|>


<commit_before>#Copyright (C) 2016 Paolo Galeone <nessuno@nerdz.eu>
#
#This Source Code Form is subject to the terms of the Mozilla Public
#License, v. 2.0. If a copy of the MPL was not distributed with this
#file, you can obtain one at http://mozilla.org/MPL/2.0/.
#Exhibit B is not attached; this software is compatible with the
#licenses expressed under Section 1.12 of the MPL v2.
"""./test_classification_pascal_rf.py --test-ds ~/data/PASCAL_2012/test/VOCdevkit/VOC2012/"""

import argparse
import os
import sys
from collections import defaultdict
from PIL import Image
import tensorflow as tf
import numpy as np
import train
import utils
from pgnet import model
from inputs import pascal

def main(args):
    """ main """

    if not os.path.exists(args.test_ds):
        print("{} does not exists".format(args.test_ds))
        return 1

    # export model.pb from session dir. Skip if model.pb already exists
    model.export(train.NUM_CLASSES, train.SESSION_DIR, "model-0",
                 train.MODEL_PATH)

    results_dir = "{}/results".format(
        os.path.dirname(os.path.abspath(__file__)))
    files = {label: open(
        results_dir + "/VOC2012/Main/comp1_cls_test_{}.txt".format(label), "w")
             for label in pascal.CLASSES}

    graph = model.load(train.MODEL_PATH, args.device)
    with graph.as_default():
        # (?, n, n, NUM_CLASSES) tensor
        logits = graph.get_tensor_by_name(model.OUTPUT_TENSOR_NAME + ":0")
        images_ = graph.get_tensor_by_name(model.INPUT_TENSOR_NAME + ":0")
        # each cell in coords (batch_position, i, j) -> is a probability vector
        per_region_probabilities = tf.nn.softmax(
            tf.reshape(logits, [-1, train.NUM_CLASSES]))
        # [tested positions, train.NUM_CLASSES]

        # array[0]=values, [1]=indices
        # get every probabiliy, because we can use localization to do classification
        top_k = tf.nn.top_k(per_region_probabilities, k=train.NUM_CLASSES)
        # each with shape [tested_positions, k]

        k = 1
        input_side = model.INPUT_SIDE + model.DOWNSAMPLING_FACTOR * model.LAST_CONV_INPUT_STRIDE * k

        test_queue, test_filename_queue = pascal.test(
            args.test_ds, 29, input_side,
            args.test_ds + "/ImageSets/Main/test.txt")

        init_op = tf.group(tf.initialize_all_variables(),
                           tf.initialize_local_variables())

        with tf.Session(config=tf.ConfigProto(
                allow_soft_placement=True)) as sess:

            sess.run(init_op)
            coordinator = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(
                sess=sess, coord=coordinator)

            try:
                processed = 0
                while not coordinator.should_stop():
                    image_batch, filename_batch = sess.run(
                        [test_queue, test_filename_queue])

                    probability_map, top_values, top_indices = sess.run(
                        [logits, top_k[0], top_k[1]],
                        feed_dict={images_: image_batch})

                    # let's think to the net as a big net, with the last layer (before the FC
                    # layers for classification) with a receptive field of
                    # LAST_KERNEL_SIDE x LAST_KERNEL_SIDE. Lets approximate the net with this last kernel:
                    # If the image is scaled down to  LAST_KERNEL_SIDE x LAST_KERNEL_SIDE
                    # the output is a single point.
                    # if the image is scaled down to something bigger
                    # (that make the output side of contolution integer) the result is a spacial map
                    # of points. Every point has a depth of num classes.

                    # for every image in the input batch
                    for batch_elem_id in range(len(image_batch)):
                        # scaling factor between original image and resized image
                        decoded_filename = filename_batch[
                            batch_elem_id].decode("utf-8")
                        image = Image.open(args.test_ds + "/JPEGImages/" +
                                           decoded_filename + ".jpg")
                        full_image_scaling_factors = np.array(
                            [image.size[0] / input_side,
                             image.size[1] / input_side])

                        probability_coords = 0
                        glance = defaultdict(list)

                        group = defaultdict(lambda: defaultdict(float))
                        for pmap_y in range(probability_map.shape[1]):
                            # calculate position in the downsampled image ds
                            ds_y = pmap_y * model.LAST_CONV_OUTPUT_STRIDE
                            for pmap_x in range(probability_map.shape[2]):
                                ds_x = pmap_x * model.LAST_CONV_OUTPUT_STRIDE

                                if top_indices[probability_coords][
                                        0] != pascal.BACKGROUND_CLASS_ID:

                                    # create coordinates of rect in the downsampled image
                                    # convert to numpy array in order to use broadcast ops
                                    coord = [ds_x, ds_y,
                                             ds_x + model.LAST_KERNEL_SIDE,
                                             ds_y + model.LAST_KERNEL_SIDE]
                                    # if something is found, append rectagle to the
                                    # map of rectalges per class
                                    rect = utils.upsample_and_shift(
                                        coord, model.DOWNSAMPLING_FACTOR,
                                        [0, 0], full_image_scaling_factors)

                                    prob = top_values[probability_coords][0]
                                    label = pascal.CLASSES[top_indices[
                                        probability_coords][0]]

                                    rect_prob = [rect, prob]
                                    glance[label].append(rect_prob)
                                    group[label]["count"] += 1
                                    group[label]["prob"] += prob

                                # update probability coord value
                                probability_coords += 1

                        classes = group.keys()
                        print('Found {} classes: {}'.format(
                            len(classes), classes))

                        # find out the minimum amount of intersection among regions
                        # in the original image, that can be used to trigger a match
                        looked_pos_side = probability_map.shape[
                            1]  # or 2, is s square. 0 dim is batch
                        looked_pos = looked_pos_side**2

                        # Save the relative frequency for every class
                        # To trigger a match, at least a fraction of intersection should be present
                        for label in group:
                            group[label]["rf"] = group[label][
                                "count"] / looked_pos
                            files[label].write("{} {}\n".format(
                                decoded_filename, group[label]["rf"]))
                        for label in set(set(pascal.CLASSES) - group.keys()):
                            files[label].write("{} {}\n".format(
                                decoded_filename, 0))

                        processed += 1

            except tf.errors.OutOfRangeError:
                print("[I] Done. Test completed!")
                print("Processed {} images".format(processed))
            finally:
                coordinator.request_stop()

            coordinator.join(threads)

        for label in files:
            files[label].close()

if __name__ == "__main__":
    PARSER = argparse.ArgumentParser(description="Test the model")
    PARSER.add_argument("--device", default="/gpu:1")
    PARSER.add_argument("--test-ds")
    sys.exit(main(PARSER.parse_args()))
<commit_msg>test_classification_pascal_rf.py: load original image using Tensorflow<commit_after>#Copyright (C) 2016 Paolo Galeone <nessuno@nerdz.eu>
#
#This Source Code Form is subject to the terms of the Mozilla Public
#License, v. 2.0. If a copy of the MPL was not distributed with this
#file, you can obtain one at http://mozilla.org/MPL/2.0/.
#Exhibit B is not attached; this software is compatible with the
#licenses expressed under Section 1.12 of the MPL v2.
"""./test_classification_pascal_rf.py --test-ds ~/data/PASCAL_2012/test/VOCdevkit/VOC2012/"""

import argparse
import os
import sys
from collections import defaultdict
import tensorflow as tf
import numpy as np
import train
import utils
from pgnet import model
from inputs import pascal, image_processing

def main(args):
    """ main """

    if not os.path.exists(args.test_ds):
        print("{} does not exists".format(args.test_ds))
        return 1

    # export model.pb from session dir. Skip if model.pb already exists
    model.export(train.NUM_CLASSES, train.SESSION_DIR, "model-0",
                 train.MODEL_PATH)

    results_dir = "{}/results".format(
        os.path.dirname(os.path.abspath(__file__)))
    files = {label: open(
        results_dir + "/VOC2012/Main/comp1_cls_test_{}.txt".format(label), "w")
             for label in pascal.CLASSES}

    graph = model.load(train.MODEL_PATH, args.device)
    with graph.as_default():
        # (?, n, n, NUM_CLASSES) tensor
        logits = graph.get_tensor_by_name(model.OUTPUT_TENSOR_NAME + ":0")
        images_ = graph.get_tensor_by_name(model.INPUT_TENSOR_NAME + ":0")
        # each cell in coords (batch_position, i, j) -> is a probability vector
        per_region_probabilities = tf.nn.softmax(
            tf.reshape(logits, [-1, train.NUM_CLASSES]))
        # [tested positions, train.NUM_CLASSES]

        # array[0]=values, [1]=indices
        # get every probabiliy, because we can use localization to do classification
        top_k = tf.nn.top_k(per_region_probabilities, k=train.NUM_CLASSES)
        # each with shape [tested_positions, k]

        k = 1
        input_side = model.INPUT_SIDE + model.DOWNSAMPLING_FACTOR * model.LAST_CONV_INPUT_STRIDE * k

        test_queue, test_filename_queue = pascal.test(
            args.test_ds, 29, input_side,
            args.test_ds + "/ImageSets/Main/test.txt")

        init_op = tf.group(tf.initialize_all_variables(),
                           tf.initialize_local_variables())

        with tf.Session(config=tf.ConfigProto(
                allow_soft_placement=True)) as sess:

            sess.run(init_op)
            coordinator = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(
                sess=sess, coord=coordinator)

            try:
                processed = 0
                while not coordinator.should_stop():
                    image_batch, filename_batch = sess.run(
                        [test_queue, test_filename_queue])

                    probability_map, top_values, top_indices = sess.run(
                        [logits, top_k[0], top_k[1]],
                        feed_dict={images_: image_batch})

                    # let's think to the net as a big net, with the last layer (before the FC
                    # layers for classification) with a receptive field of
                    # LAST_KERNEL_SIDE x LAST_KERNEL_SIDE. Lets approximate the net with this last kernel:
                    # If the image is scaled down to  LAST_KERNEL_SIDE x LAST_KERNEL_SIDE
                    # the output is a single point.
                    # if the image is scaled down to something bigger
                    # (that make the output side of contolution integer) the result is a spacial map
                    # of points. Every point has a depth of num classes.

                    # for every image in the input batch
                    for batch_elem_id in range(len(image_batch)):
                        # scaling factor between original image and resized image
                        decoded_filename = filename_batch[
                            batch_elem_id].decode("utf-8")

                        image = sess.run(
                            image_processing.read_image_jpg(
                                args.test_ds + "/JPEGImages/" +
                                decoded_filename + ".jpg"))
                        full_image_scaling_factors = np.array(
                            [image.shape[1] / input_side,
                             image.shape[0] / input_side])

                        probability_coords = 0
                        glance = defaultdict(list)

                        group = defaultdict(lambda: defaultdict(float))
                        for pmap_y in range(probability_map.shape[1]):
                            # calculate position in the downsampled image ds
                            ds_y = pmap_y * model.LAST_CONV_OUTPUT_STRIDE
                            for pmap_x in range(probability_map.shape[2]):
                                ds_x = pmap_x * model.LAST_CONV_OUTPUT_STRIDE

                                if top_indices[probability_coords][
                                        0] != pascal.BACKGROUND_CLASS_ID:

                                    # create coordinates of rect in the downsampled image
                                    # convert to numpy array in order to use broadcast ops
                                    coord = [ds_x, ds_y,
                                             ds_x + model.LAST_KERNEL_SIDE,
                                             ds_y + model.LAST_KERNEL_SIDE]
                                    # if something is found, append rectagle to the
                                    # map of rectalges per class
                                    rect = utils.upsample_and_shift(
                                        coord, model.DOWNSAMPLING_FACTOR,
                                        [0, 0], full_image_scaling_factors)

                                    prob = top_values[probability_coords][0]
                                    label = pascal.CLASSES[top_indices[
                                        probability_coords][0]]

                                    rect_prob = [rect, prob]
                                    glance[label].append(rect_prob)
                                    group[label]["count"] += 1
                                    group[label]["prob"] += prob

                                # update probability coord value
                                probability_coords += 1

                        classes = group.keys()
                        print('Found {} classes: {}'.format(
                            len(classes), classes))

                        # find out the minimum amount of intersection among regions
                        # in the original image, that can be used to trigger a match
                        # position 1or 2, is a square. 0 dim is batch
                        looked_pos_side = probability_map.shape[1]
                        looked_pos = looked_pos_side**2

                        # Save the relative frequency for every class
                        # To trigger a match, at least a fraction of intersection should be present
                        for label in group:
                            group[label]["rf"] = group[label][
                                "count"] / looked_pos
                            files[label].write("{} {}\n".format(
                                decoded_filename, group[label]["rf"]))
                        for label in set(set(pascal.CLASSES) - group.keys()):
                            files[label].write("{} {}\n".format(
                                decoded_filename, 0))

                        processed += 1

            except tf.errors.OutOfRangeError:
                print("[I] Done. Test completed!")
                print("Processed {} images".format(processed))
            finally:
                coordinator.request_stop()

            coordinator.join(threads)

        for label in files:
            files[label].close()

if __name__ == "__main__":
    PARSER = argparse.ArgumentParser(description="Test the model")
    PARSER.add_argument("--device", default="/gpu:1")
    PARSER.add_argument("--test-ds")
    sys.exit(main(PARSER.parse_args()))
<|endoftext|>