<commit_before>#    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 six
from tempest.api.volume import base
from tempest.common import waiters
from tempest import config
from tempest.lib.common.utils import data_utils
from tempest.lib import decorators

CONF = config.CONF

class VolumeMultiBackendTest(base.BaseVolumeAdminTest):

    @classmethod
    def skip_checks(cls):
        super(VolumeMultiBackendTest, cls).skip_checks()

        if not CONF.volume_feature_enabled.multi_backend:
            raise cls.skipException("Cinder multi-backend feature disabled")

    @classmethod
    def resource_setup(cls):
        super(VolumeMultiBackendTest, cls).resource_setup()

        # read backend name from a list .
        backend_names = set(CONF.volume.backend_names)

        cls.volume_id_list_with_prefix = []
        cls.volume_id_list_without_prefix = []

        # Volume/Type creation (uses volume_backend_name)
        # It is not allowed to create the same backend name twice
        if len(backend_names) < 2:
            raise cls.skipException("Requires at least two different "
                                    "backend names")
        for backend_name in backend_names:
            # Volume/Type creation (uses backend_name)
            cls._create_type_and_volume(backend_name, False)
            # Volume/Type creation (uses capabilities:volume_backend_name)
            cls._create_type_and_volume(backend_name, True)

    @classmethod
    def _create_type_and_volume(cls, backend_name_key, with_prefix):
        # Volume/Type creation
        type_name = data_utils.rand_name(cls.__name__ + '-Type')
        vol_name = data_utils.rand_name(cls.__name__ + '-Volume')
        spec_key_with_prefix = "capabilities:volume_backend_name"
        spec_key_without_prefix = "volume_backend_name"
        if with_prefix:
            extra_specs = {spec_key_with_prefix: backend_name_key}
        else:
            extra_specs = {spec_key_without_prefix: backend_name_key}
        cls.create_volume_type(name=type_name,
                               extra_specs=extra_specs)

        params = {'name': vol_name, 'volume_type': type_name,
                  'size': CONF.volume.volume_size}
        cls.volume = cls.create_volume(**params)
        if with_prefix:
            cls.volume_id_list_with_prefix.append(cls.volume['id'])
        else:
            cls.volume_id_list_without_prefix.append(
                cls.volume['id'])
        waiters.wait_for_volume_resource_status(cls.admin_volume_client,
                                                cls.volume['id'], 'available')

    @decorators.idempotent_id('c1a41f3f-9dad-493e-9f09-3ff197d477cc')
    def test_backend_name_reporting(self):
        # get volume id which created by type without prefix
        for volume_id in self.volume_id_list_without_prefix:
            self._test_backend_name_reporting_by_volume_id(volume_id)

    @decorators.idempotent_id('f38e647f-ab42-4a31-a2e7-ca86a6485215')
    def test_backend_name_reporting_with_prefix(self):
        # get volume id which created by type with prefix
        for volume_id in self.volume_id_list_with_prefix:
            self._test_backend_name_reporting_by_volume_id(volume_id)

    @decorators.idempotent_id('46435ab1-a0af-4401-8373-f14e66b0dd58')
    def test_backend_name_distinction(self):
        # get volume ids which created by type without prefix
        self._test_backend_name_distinction(self.volume_id_list_without_prefix)

    @decorators.idempotent_id('4236305b-b65a-4bfc-a9d2-69cb5b2bf2ed')
    def test_backend_name_distinction_with_prefix(self):
        # get volume ids which created by type without prefix
        self._test_backend_name_distinction(self.volume_id_list_with_prefix)

    def _get_volume_host(self, volume_id):
        return self.admin_volume_client.show_volume(
            volume_id)['volume']['os-vol-host-attr:host']

    def _test_backend_name_reporting_by_volume_id(self, volume_id):
        # this test checks if os-vol-attr:host is populated correctly after
        # the multi backend feature has been enabled
        # if multi-backend is enabled: os-vol-attr:host should be like:
        # host@backend_name
        volume = self.admin_volume_client.show_volume(volume_id)['volume']

        volume1_host = volume['os-vol-host-attr:host']
        msg = ("multi-backend reporting incorrect values for volume %s" %
               volume_id)
        self.assertGreater(len(volume1_host.split("@")), 1, msg)

    def _test_backend_name_distinction(self, volume_id_list):
        # this test checks that the volumes created at setUp don't
        # belong to the same backend (if they are, than the
        # volume backend distinction is not working properly)
        volume_hosts = [self._get_volume_host(volume) for volume in
                        volume_id_list]
        # assert that volumes are each created on separate hosts:
        msg = ("volumes %s were created in the same backend" % ", "
               .join(volume_hosts))
        six.assertCountEqual(self, volume_hosts, set(volume_hosts), msg)
<commit_msg>Move skip condition under skip_checks method<commit_after>#    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 six
from tempest.api.volume import base
from tempest.common import waiters
from tempest import config
from tempest.lib.common.utils import data_utils
from tempest.lib import decorators

CONF = config.CONF

class VolumeMultiBackendTest(base.BaseVolumeAdminTest):

    @classmethod
    def skip_checks(cls):
        super(VolumeMultiBackendTest, cls).skip_checks()

        if not CONF.volume_feature_enabled.multi_backend:
            raise cls.skipException("Cinder multi-backend feature disabled")

        if len(set(CONF.volume.backend_names)) < 2:
            raise cls.skipException("Requires at least two different "
                                    "backend names")

    @classmethod
    def resource_setup(cls):
        super(VolumeMultiBackendTest, cls).resource_setup()

        # read backend name from a list .
        backend_names = set(CONF.volume.backend_names)

        cls.volume_id_list_with_prefix = []
        cls.volume_id_list_without_prefix = []

        # Volume/Type creation (uses volume_backend_name)
        # It is not allowed to create the same backend name twice
        for backend_name in backend_names:
            # Volume/Type creation (uses backend_name)
            cls._create_type_and_volume(backend_name, False)
            # Volume/Type creation (uses capabilities:volume_backend_name)
            cls._create_type_and_volume(backend_name, True)

    @classmethod
    def _create_type_and_volume(cls, backend_name_key, with_prefix):
        # Volume/Type creation
        type_name = data_utils.rand_name(cls.__name__ + '-Type')
        vol_name = data_utils.rand_name(cls.__name__ + '-Volume')
        spec_key_with_prefix = "capabilities:volume_backend_name"
        spec_key_without_prefix = "volume_backend_name"
        if with_prefix:
            extra_specs = {spec_key_with_prefix: backend_name_key}
        else:
            extra_specs = {spec_key_without_prefix: backend_name_key}
        cls.create_volume_type(name=type_name,
                               extra_specs=extra_specs)

        params = {'name': vol_name, 'volume_type': type_name,
                  'size': CONF.volume.volume_size}
        cls.volume = cls.create_volume(**params)
        if with_prefix:
            cls.volume_id_list_with_prefix.append(cls.volume['id'])
        else:
            cls.volume_id_list_without_prefix.append(
                cls.volume['id'])
        waiters.wait_for_volume_resource_status(cls.admin_volume_client,
                                                cls.volume['id'], 'available')

    @decorators.idempotent_id('c1a41f3f-9dad-493e-9f09-3ff197d477cc')
    def test_backend_name_reporting(self):
        # get volume id which created by type without prefix
        for volume_id in self.volume_id_list_without_prefix:
            self._test_backend_name_reporting_by_volume_id(volume_id)

    @decorators.idempotent_id('f38e647f-ab42-4a31-a2e7-ca86a6485215')
    def test_backend_name_reporting_with_prefix(self):
        # get volume id which created by type with prefix
        for volume_id in self.volume_id_list_with_prefix:
            self._test_backend_name_reporting_by_volume_id(volume_id)

    @decorators.idempotent_id('46435ab1-a0af-4401-8373-f14e66b0dd58')
    def test_backend_name_distinction(self):
        # get volume ids which created by type without prefix
        self._test_backend_name_distinction(self.volume_id_list_without_prefix)

    @decorators.idempotent_id('4236305b-b65a-4bfc-a9d2-69cb5b2bf2ed')
    def test_backend_name_distinction_with_prefix(self):
        # get volume ids which created by type without prefix
        self._test_backend_name_distinction(self.volume_id_list_with_prefix)

    def _get_volume_host(self, volume_id):
        return self.admin_volume_client.show_volume(
            volume_id)['volume']['os-vol-host-attr:host']

    def _test_backend_name_reporting_by_volume_id(self, volume_id):
        # this test checks if os-vol-attr:host is populated correctly after
        # the multi backend feature has been enabled
        # if multi-backend is enabled: os-vol-attr:host should be like:
        # host@backend_name
        volume = self.admin_volume_client.show_volume(volume_id)['volume']

        volume1_host = volume['os-vol-host-attr:host']
        msg = ("multi-backend reporting incorrect values for volume %s" %
               volume_id)
        self.assertGreater(len(volume1_host.split("@")), 1, msg)

    def _test_backend_name_distinction(self, volume_id_list):
        # this test checks that the volumes created at setUp don't
        # belong to the same backend (if they are, than the
        # volume backend distinction is not working properly)
        volume_hosts = [self._get_volume_host(volume) for volume in
                        volume_id_list]
        # assert that volumes are each created on separate hosts:
        msg = ("volumes %s were created in the same backend" % ", "
               .join(volume_hosts))
        six.assertCountEqual(self, volume_hosts, set(volume_hosts), msg)
<|endoftext|>


<commit_before>import lmdb

import synapse.lib.lmdb as s_lmdb

from synapse.tests.common import *

class LmdbTest(SynTest):

    def test_lmdb_seqn(self):

        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)

            seqn = s_lmdb.Seqn(lenv, b'seqn:test')

            items = ('foo', 10, 20)

            with lenv.begin(write=True) as xact:
                seqn.save(xact, items)
                retn = tuple(seqn.iter(xact, 0))
                self.eq(retn, ((0, 'foo'), (1, 10), (2, 20)))

                self.raises(TypeError, seqn.save, xact, ({'set'},))
                retn = tuple(seqn.iter(xact, 0))
                self.eq(retn, ((0, 'foo'), (1, 10), (2, 20)))

            lenv.close()

            # Reopen the seqn and continue where we left off
            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)
            seqn = s_lmdb.Seqn(lenv, b'seqn:test')

            with lenv.begin(write=True) as xact:
                seqn.save(xact, items)

                retn = tuple(seqn.iter(xact, 0))
                self.eq(retn, ((0, 'foo'), (1, 10), (2, 20),
                               (3, 'foo'), (4, 10), (5, 20)))

            with lenv.begin() as xact:
                # We can also start in the middle of the sequence
                retn = tuple(seqn.iter(xact, 4))
                self.eq(retn, ((4, 10), (5, 20)))

            with lenv.begin() as xact:
                # iterating past the end yields nothing
                retn = tuple(seqn.iter(xact, 100))
                self.eq(retn, ())

            # A xact which is not a writer fails
            with lenv.begin() as xact:
                self.raises(lmdb.ReadonlyError, seqn.save, xact, items)

            # A subseqeunt write works.
            with lenv.begin(write=True) as xact:
                seqn.save(xact, items)
                retn = tuple(seqn.iter(xact, 0))
                self.len(9, retn)

            lenv.close()

    def test_lmdb_seqn_raceattempt(self):
        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)
            seqn = s_lmdb.Seqn(lenv, b'seqn:test')

            evt = threading.Event()
            evt1 = threading.Event()
            evt2 = threading.Event()
            evt3 = threading.Event()
            evt4 = threading.Event()

            @firethread
            def race(n, m, e):
                valus = [i for i in range(n, m)]
                e.set()
                evt.wait()
                with lenv.begin(write=True) as xact:
                    seqn.save(xact, valus)

            thr1 = race(10000, 20000, evt1)
            thr2 = race(20000, 30000, evt2)
            thr3 = race(30000, 40000, evt3)
            thr4 = race(40000, 50000, evt4)
            evt1.wait()
            evt2.wait()
            evt3.wait()
            evt4.wait()
            evt.set()

            thr1.join()
            thr2.join()
            thr3.join()
            thr4.join()

            with lenv.begin(write=False) as xact:
                retn = tuple(seqn.iter(xact, 0))
                offsets = [off for off, valu in retn]
                valus = [valu for off, valu in retn]
                self.len(40000, offsets)
                self.eq(offsets, [i for i in range(40000)])
                # Everything makes it in but order isn't guaranteed
                self.eq(sorted(valus), [i for i in range(10000, 50000)])

    def test_lmdb_metrics(self):

        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)
            metr = s_lmdb.Metrics(lenv)
            self.eq(metr.stat(), {})

            with lenv.begin(write=True) as xact:

                metr.inc(xact, 'woot', 20)
                metr.inc(xact, 'woot', 20)
                metr.record(xact, {'hehe': 10, 'haha': 20})
                metr.record(xact, {'hehe': 20, 'haha': 30})

            lenv.sync()
            lenv.close()

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)

            metr = s_lmdb.Metrics(lenv)
            with lenv.begin(write=True) as xact:

                self.eq(metr.info.get('woot'), 40)

                retn = list(metr.iter(xact, 1))
                self.len(1, retn)
                self.eq(retn[0][1].get('hehe'), 20)

                self.len(0, list(metr.iter(xact, 1234567890)))

            self.eq(metr.stat(), {'woot': 40})

    def test_lmdb_propstor(self):

        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)

            # this allows result sorting...
            buid0 = b'\x00' * 32
            buid1 = b'\x01' * 32

            psto = s_lmdb.PropStor(lenv)

            recs = (
                (buid0, (
                    (b'foo:bar', b'asdfqwer', 0),
                    (b'foo:baz', b'asdfasdf', 0),
                    (b'foo:intish', b'\x00\x80', 0),
                )),
                (buid1, (
                    (b'foo:bar', b'asdfzxcv', 0),
                    (b'foo:baz', b'qwer', 0),
                    (b'foo:intish', b'\x00\x01', 0),
                )),
            )

            with lenv.begin(write=True) as xact:
                setr = psto.getPropSetr(xact)
                retn = list(setr.put(recs))

            with lenv.begin(write=True) as xact:

                setr = psto.getPropSetr(xact=xact)
                self.true(setr.has(b'foo:baz', b'qwer'))
                self.false(setr.has(b'foo:lol', b'rofl'))
                self.true(setr.has(b'foo:baz', b'qwer'))

                self.true(setr.set(buid0, b'foo:lol', b'rofl'))
                self.true(setr.has(b'foo:lol', b'rofl'))
                self.true(setr.has(b'foo:baz', b'qwer'))

                self.false(setr.set(buid0, b'foo:lol', b'rofl'))
                self.false(setr.set(buid0, b'foo:lol', b'defv', flags=s_lmdb.STOR_FLAG_DEFVAL))
                self.false(setr.has(b'foo:lol', b'defv'))
                self.true(setr.has(b'foo:lol', b'rofl'))
                self.true(psto.has(xact, b'foo:lol', b'rofl'))

                self.true(setr.set(buid0, b'foo:lol', b'woo'))
                self.true(setr.has(b'foo:lol', b'woo'))
                self.true(psto.has(xact, b'foo:lol', b'woo'))
                self.false(setr.has(b'foo:lol', b'defv'))
                self.false(setr.has(b'foo:lol', b'rofl'))
                self.false(psto.has(xact, b'foo:lol', b'defv'))
                self.false(psto.has(xact, b'foo:lol', b'rofl'))

                self.true(setr.set(buid1, b'foo:multi', b'root', flags=s_lmdb.STOR_FLAG_MULTIVAL))
                self.true(setr.set(buid1, b'foo:multi', b'user', flags=s_lmdb.STOR_FLAG_MULTIVAL))
                self.true(setr.has(b'foo:multi', b'root'))
                self.true(setr.has(b'foo:multi', b'user'))
                self.false(setr.set(buid1, b'foo:multi', b'user', flags=s_lmdb.STOR_FLAG_MULTIVAL))

                retn = tuple(sorted(psto.pref(xact, b'foo:bar', b'asdf')))
                self.eq(retn, ((buid0, b'foo:bar', b'asdfqwer'), (buid1, b'foo:bar', b'asdfzxcv')))

                retn = tuple(sorted(psto.range(xact, b'foo:intish', b'\x00\x00', b'\x00\x80')))
                self.eq(retn, ((buid1, b'foo:intish', b'\x00\x01'),))

                retn = tuple(sorted(psto.eq(xact, b'foo:intish', b'\x00\x01')))
                self.eq(retn, ((buid1, b'foo:intish', b'\x00\x01'), ))
<commit_msg>Add additional LMDB unit test assertions<commit_after>import lmdb

import synapse.lib.lmdb as s_lmdb

from synapse.tests.common import *

class LmdbTest(SynTest):

    def test_lmdb_seqn(self):

        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)

            seqn = s_lmdb.Seqn(lenv, b'seqn:test')

            items = ('foo', 10, 20)

            with lenv.begin(write=True) as xact:
                seqn.save(xact, items)
                retn = tuple(seqn.iter(xact, 0))
                self.eq(retn, ((0, 'foo'), (1, 10), (2, 20)))

                self.raises(TypeError, seqn.save, xact, ({'set'},))
                retn = tuple(seqn.iter(xact, 0))
                self.eq(retn, ((0, 'foo'), (1, 10), (2, 20)))

            lenv.close()

            # Reopen the seqn and continue where we left off
            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)
            seqn = s_lmdb.Seqn(lenv, b'seqn:test')

            with lenv.begin(write=True) as xact:
                seqn.save(xact, items)

                retn = tuple(seqn.iter(xact, 0))
                self.eq(retn, ((0, 'foo'), (1, 10), (2, 20),
                               (3, 'foo'), (4, 10), (5, 20)))

            with lenv.begin() as xact:
                # We can also start in the middle of the sequence
                retn = tuple(seqn.iter(xact, 4))
                self.eq(retn, ((4, 10), (5, 20)))

            with lenv.begin() as xact:
                # iterating past the end yields nothing
                retn = tuple(seqn.iter(xact, 100))
                self.eq(retn, ())

            # A xact which is not a writer fails
            with lenv.begin() as xact:
                self.raises(lmdb.ReadonlyError, seqn.save, xact, items)

            # A subseqeunt write works.
            with lenv.begin(write=True) as xact:
                seqn.save(xact, items)
                retn = tuple(seqn.iter(xact, 0))
                self.len(9, retn)

            lenv.close()

    def test_lmdb_seqn_raceattempt(self):
        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)
            seqn = s_lmdb.Seqn(lenv, b'seqn:test')

            evt = threading.Event()
            evt1 = threading.Event()
            evt2 = threading.Event()
            evt3 = threading.Event()
            evt4 = threading.Event()

            @firethread
            def race(n, m, e):
                valus = [i for i in range(n, m)]
                e.set()
                evt.wait()
                with lenv.begin(write=True) as xact:
                    seqn.save(xact, valus)

            thr1 = race(10000, 20000, evt1)
            thr2 = race(20000, 30000, evt2)
            thr3 = race(30000, 40000, evt3)
            thr4 = race(40000, 50000, evt4)
            evt1.wait()
            evt2.wait()
            evt3.wait()
            evt4.wait()
            evt.set()

            thr1.join()
            thr2.join()
            thr3.join()
            thr4.join()

            with lenv.begin(write=False) as xact:
                retn = tuple(seqn.iter(xact, 0))
                offsets = [off for off, valu in retn]
                valus = [valu for off, valu in retn]
                self.len(40000, offsets)
                self.eq(offsets, [i for i in range(40000)])
                # Everything makes it in but order isn't guaranteed
                self.eq(sorted(valus), [i for i in range(10000, 50000)])

    def test_lmdb_metrics(self):

        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)
            metr = s_lmdb.Metrics(lenv)
            self.eq(metr.stat(), {})

            with lenv.begin(write=True) as xact:

                metr.inc(xact, 'woot', 20)
                metr.inc(xact, 'woot', 20)
                metr.record(xact, {'hehe': 10, 'haha': 20})
                metr.record(xact, {'hehe': 20, 'haha': 30})

            lenv.sync()
            lenv.close()

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)

            metr = s_lmdb.Metrics(lenv)
            with lenv.begin(write=True) as xact:

                self.eq(metr.info.get('woot'), 40)

                retn = list(metr.iter(xact, 1))
                self.len(1, retn)
                self.eq(retn[0][1].get('hehe'), 20)

                self.len(0, list(metr.iter(xact, 1234567890)))

            self.eq(metr.stat(), {'woot': 40})

    def test_lmdb_propstor(self):

        with self.getTestDir() as dirn:

            lenv = lmdb.open(dirn, writemap=True, max_dbs=128)

            # this allows result sorting...
            buid0 = b'\x00' * 32
            buid1 = b'\x01' * 32

            psto = s_lmdb.PropStor(lenv)

            recs = (
                (buid0, (
                    (b'foo:bar', b'asdfqwer', 0),
                    (b'foo:baz', b'asdfasdf', 0),
                    (b'foo:intish', b'\x00\x80', 0),
                )),
                (buid1, (
                    (b'foo:bar', b'asdfzxcv', 0),
                    (b'foo:baz', b'qwer', 0),
                    (b'foo:intish', b'\x00\x01', 0),
                )),
            )

            with lenv.begin(write=True) as xact:
                setr = psto.getPropSetr(xact)
                retn = list(setr.put(recs))

            with lenv.begin(write=True) as xact:

                setr = psto.getPropSetr(xact=xact)
                self.true(setr.has(b'foo:baz', b'qwer'))
                self.false(setr.has(b'foo:lol', b'rofl'))
                self.true(setr.has(b'foo:baz', b'qwer'))

                self.true(setr.set(buid0, b'foo:lol', b'rofl'))
                self.true(setr.has(b'foo:lol', b'rofl'))
                self.true(setr.has(b'foo:baz', b'qwer'))

                self.false(setr.set(buid0, b'foo:lol', b'rofl'))
                self.false(setr.set(buid0, b'foo:lol', b'defv', flags=s_lmdb.STOR_FLAG_DEFVAL))
                self.false(setr.has(b'foo:lol', b'defv'))
                self.true(setr.has(b'foo:lol', b'rofl'))
                self.true(psto.has(xact, b'foo:lol', b'rofl'))

                self.true(setr.set(buid0, b'foo:lol', b'woo'))
                self.true(setr.has(b'foo:lol', b'woo'))
                self.true(psto.has(xact, b'foo:lol', b'woo'))
                self.false(setr.has(b'foo:lol', b'defv'))
                self.false(setr.has(b'foo:lol', b'rofl'))
                self.false(psto.has(xact, b'foo:lol', b'defv'))
                self.false(psto.has(xact, b'foo:lol', b'rofl'))

                self.true(setr.set(buid1, b'foo:multi', b'root', flags=s_lmdb.STOR_FLAG_MULTIVAL))
                self.true(setr.set(buid1, b'foo:multi', b'user', flags=s_lmdb.STOR_FLAG_MULTIVAL))
                self.true(setr.has(b'foo:multi', b'root'))
                self.true(setr.has(b'foo:multi', b'user'))
                self.false(setr.set(buid1, b'foo:multi', b'user', flags=s_lmdb.STOR_FLAG_MULTIVAL))

                retn = tuple(sorted(psto.pref(xact, b'foo:bar', b'asdf')))
                self.eq(retn, ((buid0, b'foo:bar', b'asdfqwer'), (buid1, b'foo:bar', b'asdfzxcv')))
                retn = tuple(sorted(psto.pref(xact, b'foo:bar:nope', b'asdf')))
                self.eq(retn, ())

                retn = tuple(sorted(psto.range(xact, b'foo:intish', b'\x00\x00', b'\x00\x80')))
                self.eq(retn, ((buid1, b'foo:intish', b'\x00\x01'),))
                retn = tuple(sorted(psto.range(xact, b'foo:nothere', b'\x00\x00', b'\x00\x80')))
                self.eq(retn, ())

                retn = tuple(sorted(psto.eq(xact, b'foo:intish', b'\x00\x01')))
                self.eq(retn, ((buid1, b'foo:intish', b'\x00\x01'), ))
<|endoftext|>


<commit_before>__cmdname__ = "add_check"

import json
import os
import sys

import circonusapi
import log
import util

class Module(object):
    def __init__(self, api, account):
        self.api = api
        self.account = account

    def command(self, opts, template_name, target, agent, *params):
        """Adds a check based on a template

        Arguments:
            template_name   -- the name of the template file
            check_name      -- the name for the check
            target          -- the target of the check (can be a hostname)
            agent           -- the agent to run the check from
            params          -- other parameters (see below)

        Other parameters are specified as "param_name=value" and will
        be substituted in the template. Use {param_name} in the template.

        Some predefined parameters:

            {agent}         -- the agent provided on the command line
            {target}        -- the target provided on the command line
            {targetip}      -- the ip of the target resolved in dns
                               (may be same as target if target was provided
                                as an IP)
        """
        template = util.Template(template_name, "check")
        targetip = util.resolve_target(target)
        template_params = {
            'agent': agent,
            'target': target,
            'targetip': targetip}
        template_params.update(template.parse_nv_params(params))

        substituted = template.sub(template_params)
        # Add required parameters
        substituted['agent_id'] = util.get_agent(self.api, agent)
        substituted['target'] = targetip

        try:
            self.api.add_check_bundle(**substituted)
            log.msg("Check added")
        except circonusapi.CirconusAPIError, e:
            log.error("Failed to add check: %s" % e.error)
<commit_msg>Allow specifying metric by regex in check templates<commit_after>__cmdname__ = "add_check"

import json
import os
import re
import sys

import circonusapi
import log
import util

class Module(object):
    def __init__(self, api, account):
        self.api = api
        self.account = account

    def command(self, opts, template_name, target, agent, *params):
        """Adds a check based on a template

        Arguments:
            template_name   -- the name of the template file
            check_name      -- the name for the check
            target          -- the target of the check (can be a hostname)
            agent           -- the agent to run the check from
            params          -- other parameters (see below)

        Other parameters are specified as "param_name=value" and will
        be substituted in the template. Use {param_name} in the template.

        Some predefined parameters:

            {agent}         -- the agent provided on the command line
            {target}        -- the target provided on the command line
            {targetip}      -- the ip of the target resolved in dns
                               (may be same as target if target was provided
                                as an IP)
        """
        template = util.Template(template_name, "check")
        targetip = util.resolve_target(target)
        template_params = {
            'agent': agent,
            'target': target,
            'targetip': targetip}
        template_params.update(template.parse_nv_params(params))

        substituted = template.sub(template_params)
        # Add required parameters
        substituted['agent_id'] = util.get_agent(self.api, agent)
        substituted['target'] = targetip

        # Allow matching metrics by regex, find available metrics, and test
        # them against each regex in the metrics_regex key in the template to
        # see if they match.
        if 'metric_regex' in substituted:
            log.msg("Fetching available metrics for regex match")
            try:
                substituted['test_mode'] = 1
                rv = self.api.add_check_bundle(**substituted)
                del substituted['test_mode']
            except circonusapi.CirconusAPIError, e:
                log.error("Failed to fetch available metrics: %s" % e.error)
                sys.exit(1)
            available_metrics = rv['metrics']
            for metric_type in available_metrics:
                for m in available_metrics[metric_type]:
                    for regex in substituted['metric_regex']:
                        if re.match(regex, m):
                            substituted['metric_name'].append(m)
                            break
            log.msg("Metrics to include in the check:")
            for m in sorted(substituted['metric_name']):
                print "    %s" % m
            if not util.confirm():
                sys.exit(1)

        try:
            self.api.add_check_bundle(**substituted)
            log.msg("Check added")
        except circonusapi.CirconusAPIError, e:
            log.error("Failed to add check: %s" % e.error)
<|endoftext|>


<commit_before>import numpy
import numpy.linalg

# Disable divide-by-zero warnings so that 1.0 / (1.0, 0.0, 0.0) returns (1.0, Inf, Inf) without complaining
numpy.seterr(divide='ignore')

import math

from copy import deepcopy

##  Simple 3D-vector class based on numpy arrays.
#
#   This class represents a 3-dimensional vector.
class Vector(object):
    Unit_X = None
    Unit_Y = None
    Unit_Z = None

    def __init__(self,x = 0 ,y = 0,z = 0):
        self._data = numpy.array([x, y, z],dtype=numpy.float32)

    ##  Set the data of the vector
    #   \param x X coordinate of vector.
    #   \param y Y coordinate of vector.
    #   \param z Z coordinate of vector.
    def setData(self, x = 0,y = 0,z = 0):
        self._data = numpy.array([x,y,z],dtype=numpy.float32)

    ##  Get numpy array with the data
    #   \returns numpy array of length 3 holding xyz data.
    def getData(self):
        return self._data

    ##  Return the x component of this vector
    @property
    def x(self):
        return self._data[0]

    ##  Set the x component of this vector
    #   \param value The value for the x component
    def setX(self, value):
        self._data[0] = value

    ##  Return the y component of this vector
    @property
    def y(self):
        return self._data[1]

    ##  Set the y component of this vector
    #   \param value The value for the y component
    def setY(self, value):
        self._data[1] = value

    ## Return the z component of this vector
    @property
    def z(self):
        return self._data[2]

    ##  Set the z component of this vector
    #   \param value The value for the z component
    def setZ(self, value):
        self._data[2] = value

    ##  Get the angle from this vector to another
    def angleToVector(self, vector,):
        v0 = numpy.array(self._data, dtype=numpy.float32, copy=False)
        v1 = numpy.array(vector.getData(), dtype = numpy.float32, copy=False)
        dot = numpy.sum(v0 * v1)
        dot /= self._normalizeVector(v0) * self._normalizeVector(v1)
        return numpy.arccos(numpy.fabs(dot))

    def normalize(self):
        l = self.length()
        if l != 0:
            self._data /= l
        return self

    def getNormalized(self):
        other = deepcopy(self)
        return other.normalize()

    ##  Return length, i.e. Euclidean norm, of ndarray along axis.
    def _normalizeVector(self, data):
        data = numpy.array(data, dtype=numpy.float32, copy=True)
        if data.ndim == 1:
            return math.sqrt(numpy.dot(data, data))
        data *= data
        out = numpy.atleast_1d(numpy.sum(data))
        numpy.sqrt(out, out)
        return out

    def length(self):
        return numpy.linalg.norm(self._data)

    def dot(self, other):
        return numpy.dot(self._data, other._data)

    def cross(self, other):
        result = numpy.cross(self._data, other._data)
        return Vector(result[0], result[1], result[2])

    def rotated(self, matrix):
        result = Vector()
        result.setX(matrix.at(0, 0) * self._data[0] + matrix.at(1, 0) * self._data[1] + matrix.at(2, 0) * self._data[2])
        result.setY(matrix.at(0, 1) * self._data[0] + matrix.at(1, 1) * self._data[1] + matrix.at(2, 1) * self._data[2])
        result.setZ(matrix.at(0, 2) * self._data[0] + matrix.at(1, 2) * self._data[1] + matrix.at(2, 2) * self._data[2])
        return result

    def __eq__(self, other):
        return self._data.all() == other._data.all()

    def __add__(self, other):
        v = Vector(self._data[0], self._data[1], self._data[2])
        v += other
        return v

    def __iadd__(self, other):
        if type(other) is float:
            self._data[0] += other
            self._data[1] += other
            self._data[2] += other
        elif type(other) is Vector:
            self._data[0] += other._data[0]
            self._data[1] += other._data[1]
            self._data[2] += other._data[2]
        else:
            raise NotImplementedError()

        return self

    def __sub__(self, other):
        v = Vector(self._data[0], self._data[1], self._data[2])
        v -= other
        return v

    def __isub__(self, other):
        if type(other) is float:
            self._data[0] -= other
            self._data[1] -= other
            self._data[2] -= other
        elif type(other) is Vector:
            self._data[0] -= other._data[0]
            self._data[1] -= other._data[1]
            self._data[2] -= other._data[2]
        else:
            raise NotImplementedError()

        return self

    def __truediv__(self, other):
        v = Vector(self._data[0], self._data[1], self._data[2])
        v /= other
        return v

    def __itruediv__(self, other):
        if type(other) is float:
            self._data /= other
            return self
        else:
            raise NotImplementedError()

    def __rtruediv__(self, other):
        if type(other) is float:
            v = numpy.float32(other) / self._data
            return Vector(v[0], v[1], v[2])
        else:
            raise NotImplementedError()

    def __neg__(self):
        self._data = -self._data
        return self

    def __pos__(self):
        self._data = +self._data
        return self

    def __repr__(self):
        return "Vector({0}, {1}, {2})".format(self._data[0], self._data[1], self._data[2])

Vector.Unit_X = Vector(1, 0, 0)
Vector.Unit_Y = Vector(0, 1, 0)
Vector.Unit_Z = Vector(0, 0, 1)
<commit_msg>Replace Vector::rotate with Vector::multiply<commit_after>import numpy
import numpy.linalg

# Disable divide-by-zero warnings so that 1.0 / (1.0, 0.0, 0.0) returns (1.0, Inf, Inf) without complaining
numpy.seterr(divide='ignore')

import math

from copy import deepcopy

##  Simple 3D-vector class based on numpy arrays.
#
#   This class represents a 3-dimensional vector.
class Vector(object):
    Unit_X = None
    Unit_Y = None
    Unit_Z = None

    def __init__(self,x = 0 ,y = 0,z = 0):
        self._data = numpy.array([x, y, z],dtype=numpy.float32)

    ##  Set the data of the vector
    #   \param x X coordinate of vector.
    #   \param y Y coordinate of vector.
    #   \param z Z coordinate of vector.
    def setData(self, x = 0,y = 0,z = 0):
        self._data = numpy.array([x,y,z],dtype=numpy.float32)

    ##  Get numpy array with the data
    #   \returns numpy array of length 3 holding xyz data.
    def getData(self):
        return self._data

    ##  Return the x component of this vector
    @property
    def x(self):
        return self._data[0]

    ##  Set the x component of this vector
    #   \param value The value for the x component
    def setX(self, value):
        self._data[0] = value

    ##  Return the y component of this vector
    @property
    def y(self):
        return self._data[1]

    ##  Set the y component of this vector
    #   \param value The value for the y component
    def setY(self, value):
        self._data[1] = value

    ## Return the z component of this vector
    @property
    def z(self):
        return self._data[2]

    ##  Set the z component of this vector
    #   \param value The value for the z component
    def setZ(self, value):
        self._data[2] = value

    ##  Get the angle from this vector to another
    def angleToVector(self, vector,):
        v0 = numpy.array(self._data, dtype=numpy.float32, copy=False)
        v1 = numpy.array(vector.getData(), dtype = numpy.float32, copy=False)
        dot = numpy.sum(v0 * v1)
        dot /= self._normalizeVector(v0) * self._normalizeVector(v1)
        return numpy.arccos(numpy.fabs(dot))

    def normalize(self):
        l = self.length()
        if l != 0:
            self._data /= l
        return self

    def getNormalized(self):
        other = deepcopy(self)
        return other.normalize()

    ##  Return length, i.e. Euclidean norm, of ndarray along axis.
    def _normalizeVector(self, data):
        data = numpy.array(data, dtype=numpy.float32, copy=True)
        if data.ndim == 1:
            return math.sqrt(numpy.dot(data, data))
        data *= data
        out = numpy.atleast_1d(numpy.sum(data))
        numpy.sqrt(out, out)
        return out

    def length(self):
        return numpy.linalg.norm(self._data)

    def dot(self, other):
        return numpy.dot(self._data, other._data)

    def cross(self, other):
        result = numpy.cross(self._data, other._data)
        return Vector(result[0], result[1], result[2])

    def multiply(self, matrix):
        d = numpy.empty(4, dtype=numpy.float32)
        d[0] = self._data[0]
        d[1] = self._data[1]
        d[2] = self._data[2]
        d[3] = 1.0

        d = d.dot(matrix.getData())

        return Vector(d[0], d[1], d[2])

    def __eq__(self, other):
        return self._data.all() == other._data.all()

    def __add__(self, other):
        v = Vector(self._data[0], self._data[1], self._data[2])
        v += other
        return v

    def __iadd__(self, other):
        if type(other) is float:
            self._data[0] += other
            self._data[1] += other
            self._data[2] += other
        elif type(other) is Vector:
            self._data[0] += other._data[0]
            self._data[1] += other._data[1]
            self._data[2] += other._data[2]
        else:
            raise NotImplementedError()

        return self

    def __sub__(self, other):
        v = Vector(self._data[0], self._data[1], self._data[2])
        v -= other
        return v

    def __isub__(self, other):
        if type(other) is float:
            self._data[0] -= other
            self._data[1] -= other
            self._data[2] -= other
        elif type(other) is Vector:
            self._data[0] -= other._data[0]
            self._data[1] -= other._data[1]
            self._data[2] -= other._data[2]
        else:
            raise NotImplementedError()

        return self

    def __truediv__(self, other):
        v = Vector(self._data[0], self._data[1], self._data[2])
        v /= other
        return v

    def __itruediv__(self, other):
        if type(other) is float:
            self._data /= other
            return self
        else:
            raise NotImplementedError()

    def __rtruediv__(self, other):
        if type(other) is float:
            v = numpy.float32(other) / self._data
            return Vector(v[0], v[1], v[2])
        else:
            raise NotImplementedError()

    def __neg__(self):
        self._data = -self._data
        return self

    def __pos__(self):
        self._data = +self._data
        return self

    def __repr__(self):
        return "Vector({0}, {1}, {2})".format(self._data[0], self._data[1], self._data[2])

Vector.Unit_X = Vector(1, 0, 0)
Vector.Unit_Y = Vector(0, 1, 0)
Vector.Unit_Z = Vector(0, 0, 1)
<|endoftext|>


<commit_before>"""A collection of classes and methods to deal with collections of
rates that together make up a network."""

# Common Imports
from __future__ import print_function

import functools
import math
from operator import mul
import os
from collections import OrderedDict

from ipywidgets import interact

import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import networkx as nx

# Import Rate
from pynucastro.rates import Rate, Nucleus, Library

mpl.rcParams['figure.dpi'] = 100

class Composition(object):
    """a composition holds the mass fractions of the nuclei in a network
    -- useful for evaluating the rates

    """
    def __init__(self, nuclei, small=1.e-16):
        """nuclei is an iterable of the nuclei (Nucleus objects) in the network"""
        if not isinstance(nuclei[0], Nucleus):
            raise ValueError("must supply an iterable of Nucleus objects")
        else:
            self.X = {k: small for k in nuclei}

    def set_solar_like(self, Z=0.02):
        """ approximate a solar abundance, setting p to 0.7, He4 to 0.3 - Z and
        the remainder evenly distributed with Z """
        num = len(self.X)
        rem = Z/(num-2)
        for k in self.X:
            if k == Nucleus("p"):
                self.X[k] = 0.7
            elif k.raw == "he4":
                self.X[k] = 0.3 - Z
            else:
                self.X[k] = rem

        self.normalize()

    def set_all(self, xval):
        """ set all species to a particular value """
        for k in self.X:
            self.X[k] = xval

    def set_nuc(self, name, xval):
        """ set nuclei name to the mass fraction xval """
        for k in self.X:
            if k.raw == name:
                self.X[k] = xval
                break

    def normalize(self):
        """ normalize the mass fractions to sum to 1 """
        X_sum = sum([self.X[k] for k in self.X])

        for k in self.X:
            self.X[k] /= X_sum

    def get_molar(self):
        """ return a dictionary of molar fractions"""
        molar_frac = {k: v/k.A for k, v in self.X.items()}
        return molar_frac

    def __str__(self):
        ostr = ""
        for k in self.X:
            ostr += "  X({}) : {}\n".format(k, self.X[k])
        return ostr

class RateCollection(object):
    """ a collection of rates that together define a network """

    pynucastro_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

    def __init__(self, rate_files=None, libraries=None, rates=None):
        """
        rate_files are the files that together define the network.  This
        can be any iterable or single string.

        This can include Reaclib library files storing multiple rates.

        If libraries is supplied, initialize a RateCollection using the rates
        in the Library object(s) in list 'libraries'.

        If rates is supplied, initialize a RateCollection using the
        Rate objects in the list 'rates'.

        Any combination of these options may be combined.
        """

        self.files = []
        self.rates = []
        self.library = None

        if rate_files:
            if isinstance(rate_files, str):
                rate_files = [rate_files]
            self._read_rate_files(rate_files)

        if rates:
            if isinstance(rates, Rate):
                rates = [rates]
            try:
                for r in rates:
                    assert(isinstance(r, Rate))
            except:
                print('Expected Rate object or list of Rate objects passed as the rates argument.')
                raise
            else:
                rlib = Library(rates=rates)
                if not self.library:
                    self.library = rlib
                else:
                    self.library = self.library + rlib

        if libraries:
            if isinstance(libraries, Library):
                libraries = [libraries]
            try:
                for lib in libraries:
                    assert(isinstance(lib, Library))
            except:
                print('Expected Library object or list of Library objects passed as the libraries argument.')
                raise
            else:
                if not self.library:
                    self.library = libraries.pop(0)
                for lib in libraries:
                    self.library = self.library + lib

        if self.library:
            self.rates = self.rates + self.library.get_rates()

        # get the unique nuclei
        u = []
        for r in self.rates:
            t = set(r.reactants + r.products)
            u = set(list(u) + list(t))

        self.unique_nuclei = sorted(u)

        # now make a list of each rate that touches each nucleus
        # we'll store this in a dictionary keyed on the nucleus
        self.nuclei_consumed = OrderedDict()
        self.nuclei_produced = OrderedDict()

        for n in self.unique_nuclei:
            self.nuclei_consumed[n] = [r for r in self.rates if n in r.reactants]
            self.nuclei_produced[n] = [r for r in self.rates if n in r.products]

        # Re-order self.rates so Reaclib rates come first,
        # followed by Tabular rates. This is needed if
        # reaclib coefficients are targets of a pointer array
        # in the Fortran network.
        # It is desired to avoid wasting array size
        # storing meaningless Tabular coefficient pointers.
        self.rates = sorted(self.rates,
                            key=lambda r: r.chapter == 't')

        self.tabular_rates = []
        self.reaclib_rates = []
        for n, r in enumerate(self.rates):
            if r.chapter == 't':
                self.tabular_rates.append(n)
            elif isinstance(r.chapter, int):
                self.reaclib_rates.append(n)
            else:
                print('ERROR: Chapter type unknown for rate chapter {}'.format(
                    str(r.chapter)))
                exit()

    def _read_rate_files(self, rate_files):
        # get the rates
        self.files = rate_files
        for rf in self.files:
            try:
                rflib = Library(rf)
            except:
                print("Error reading library from file: {}".format(rf))
                raise
            else:
                if not self.library:
                    self.library = rflib
                else:
                    self.library = self.library + rflib

    def get_nuclei(self):
        """ get all the nuclei that are part of the network """
        return self.unique_nuclei

    def evaluate_rates(self, rho, T, composition):
        """evaluate the rates for a specific density, temperature, and
        composition"""
        rvals = OrderedDict()
        ys = composition.get_molar()

        for r in self.rates:
            val = r.prefactor * rho**r.dens_exp * r.eval(T)
            yfac = functools.reduce(mul, [ys[q] for q in r.reactants])
            rvals[r] = yfac * val

        return rvals

    def network_overview(self):
        """ return a verbose network overview """
        ostr = ""
        for n in self.unique_nuclei:
            ostr += "{}\n".format(n)
            ostr += "  consumed by:\n"
            for r in self.nuclei_consumed[n]:
                ostr += "     {}\n".format(r.string)

            ostr += "  produced by:\n"
            for r in self.nuclei_produced[n]:
                ostr += "     {}\n".format(r.string)

            ostr += "\n"
        return ostr

    def write_network(self, *args, **kwargs):
        """Before writing the network, check to make sure the rates
        are distinguishable by name."""
        assert self._distinguishable_rates(), "ERROR: Rates not uniquely identified by Rate.fname"
        self._write_network(*args, **kwargs)

    def _distinguishable_rates(self):
        """Every Rate in this RateCollection should have a unique Rate.fname,
        as the network writers distinguish the rates on this basis."""
        names = [r.fname for r in self.rates]
        return len(set(names)) == len(self.rates)

    def _write_network(self, *args, **kwargs):
        """A stub for function to output the network -- this is implementation
        dependent."""
        print('To create network integration source code, use a class that implements a specific network type.')
        return

    def plot(self, outfile=None, rho=None, T=None, comp=None, size=(800, 600), dpi=100):
        """Make a plot of the network structure showing the links between nuclei"""

        G = nx.MultiDiGraph()
        G.position = {}
        G.labels = {}

        fig, ax = plt.subplots()
        #divider = make_axes_locatable(ax)
        #cax = divider.append_axes('right', size='15%', pad=0.05)

        ax.plot([0, 0], [8, 8], 'b-')

        # nodes -- the node nuclei will be all of the heavies, but not
        # p, n, alpha, unless we have p + p, 3-a, etc.
        node_nuclei = []
        for n in self.unique_nuclei:
            if n.raw not in ["p", "n", "he4"]:
                node_nuclei.append(n)
            else:
                for r in self.rates:
                    if r.reactants.count(n) > 1:
                        node_nuclei.append(n)
                        break

        for n in node_nuclei:
            G.add_node(n)
            G.position[n] = (n.N, n.Z)
            G.labels[n] = r"${}$".format(n.pretty)

        if rho is not None and T is not None and comp is not None:
            ydots = self.evaluate_rates(rho, T, comp)
        else:
            ydots = None

        #for rr in ydots:
        #    print("{}: {}".format(rr, ydots[rr]))

        # edges
        for n in node_nuclei:
            for r in self.nuclei_consumed[n]:
                for p in r.products:
                    if p in node_nuclei:
                        # networkx doesn't seem to keep the edges in
                        # any particular order, so we associate data
                        # to the edges here directly, in this case,
                        # the reaction rate, which will be used to
                        # color it
                        if ydots is None:
                            G.add_edges_from([(n, p)], weight=0.5)
                        else:
                            try:
                                rate_weight = math.log10(ydots[r])
                            except ValueError:
                                # if ydots[r] is zero, then set the weight
                                # to roughly the minimum exponent possible
                                # for python floats
                                rate_weight = -308
                            except:
                                raise
                            G.add_edges_from([(n, p)], weight=rate_weight)

        nx.draw_networkx_nodes(G, G.position,
                               node_color="#A0CBE2", alpha=1.0,
                               node_shape="o", node_size=1000, linewidth=2.0, zorder=10, ax=ax)

        nx.draw_networkx_labels(G, G.position, G.labels,
                                font_size=13, font_color="w", zorder=100, ax=ax)

        # get the edges and weights coupled in the same order
        edges, weights = zip(*nx.get_edge_attributes(G, 'weight').items())

        edges_lc = nx.draw_networkx_edges(G, G.position, width=3,
                                          edgelist=edges, edge_color=weights,
                                          node_size=1000,
                                          edge_cmap=plt.cm.viridis, zorder=1, ax=ax)

        # for networkx <= 2.0 draw_networkx_edges returns a
        # LineCollection matplotlib type which we can use for the
        # colorbar directly.  For networkx >= 2.1, it is a collection
        # of FancyArrowPatch-s, which we need to run through a
        # PatchCollection.  See:
        # https://stackoverflow.com/questions/18658047/adding-a-matplotlib-colorbar-from-a-patchcollection

        if ydots is not None:
            pc = mpl.collections.PatchCollection(edges_lc, cmap=plt.cm.viridis)
            pc.set_array(weights)
            plt.colorbar(pc, label="log10(rate)")

        Ns = [n.N for n in node_nuclei]
        Zs = [n.Z for n in node_nuclei]

        plt.xlim(min(Ns)-1, max(Ns)+1)
        #plt.ylim(min(Zs)-1, max(Zs)+1)
        plt.xlabel(r"$N$", fontsize="large")
        plt.ylabel(r"$Z$", fontsize="large")

        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')

        ax.set_aspect("equal", "datalim")

        fig.set_size_inches(size[0]/dpi, size[1]/dpi)

        if outfile is None:
            plt.show()
        else:
            plt.tight_layout()
            plt.savefig(outfile, dpi=dpi)

    def __repr__(self):
        string = ""
        for r in self.rates:
            string += "{}\n".format(r.string)
        return string

class Explorer(object):
    """ interactively explore a rate collection """
    def __init__(self, rc, comp, size=(800, 600)):
        """ take a RateCollection and a composition """
        self.rc = rc
        self.comp = comp
        self.size = size

    def _make_plot(self, logrho, logT):
        self.rc.plot(rho=10.0**logrho, T=10.0**logT, comp=self.comp, size=self.size)

    def explore(self):
        """Perform interactive exploration of the network structure."""
        interact(self._make_plot, logrho=(2, 6, 0.1), logT=(7, 9, 0.1))
<commit_msg>make the explorer limits an optional parameter<commit_after>"""A collection of classes and methods to deal with collections of
rates that together make up a network."""

# Common Imports
from __future__ import print_function

import functools
import math
from operator import mul
import os
from collections import OrderedDict

from ipywidgets import interact

import matplotlib as mpl
import matplotlib.pyplot as plt
#from mpl_toolkits.axes_grid1 import make_axes_locatable
import networkx as nx

# Import Rate
from pynucastro.rates import Rate, Nucleus, Library

mpl.rcParams['figure.dpi'] = 100

class Composition(object):
    """a composition holds the mass fractions of the nuclei in a network
    -- useful for evaluating the rates

    """
    def __init__(self, nuclei, small=1.e-16):
        """nuclei is an iterable of the nuclei (Nucleus objects) in the network"""
        if not isinstance(nuclei[0], Nucleus):
            raise ValueError("must supply an iterable of Nucleus objects")
        else:
            self.X = {k: small for k in nuclei}

    def set_solar_like(self, Z=0.02):
        """ approximate a solar abundance, setting p to 0.7, He4 to 0.3 - Z and
        the remainder evenly distributed with Z """
        num = len(self.X)
        rem = Z/(num-2)
        for k in self.X:
            if k == Nucleus("p"):
                self.X[k] = 0.7
            elif k.raw == "he4":
                self.X[k] = 0.3 - Z
            else:
                self.X[k] = rem

        self.normalize()

    def set_all(self, xval):
        """ set all species to a particular value """
        for k in self.X:
            self.X[k] = xval

    def set_nuc(self, name, xval):
        """ set nuclei name to the mass fraction xval """
        for k in self.X:
            if k.raw == name:
                self.X[k] = xval
                break

    def normalize(self):
        """ normalize the mass fractions to sum to 1 """
        X_sum = sum([self.X[k] for k in self.X])

        for k in self.X:
            self.X[k] /= X_sum

    def get_molar(self):
        """ return a dictionary of molar fractions"""
        molar_frac = {k: v/k.A for k, v in self.X.items()}
        return molar_frac

    def __str__(self):
        ostr = ""
        for k in self.X:
            ostr += "  X({}) : {}\n".format(k, self.X[k])
        return ostr

class RateCollection(object):
    """ a collection of rates that together define a network """

    pynucastro_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

    def __init__(self, rate_files=None, libraries=None, rates=None):
        """
        rate_files are the files that together define the network.  This
        can be any iterable or single string.

        This can include Reaclib library files storing multiple rates.

        If libraries is supplied, initialize a RateCollection using the rates
        in the Library object(s) in list 'libraries'.

        If rates is supplied, initialize a RateCollection using the
        Rate objects in the list 'rates'.

        Any combination of these options may be combined.
        """

        self.files = []
        self.rates = []
        self.library = None

        if rate_files:
            if isinstance(rate_files, str):
                rate_files = [rate_files]
            self._read_rate_files(rate_files)

        if rates:
            if isinstance(rates, Rate):
                rates = [rates]
            try:
                for r in rates:
                    assert(isinstance(r, Rate))
            except:
                print('Expected Rate object or list of Rate objects passed as the rates argument.')
                raise
            else:
                rlib = Library(rates=rates)
                if not self.library:
                    self.library = rlib
                else:
                    self.library = self.library + rlib

        if libraries:
            if isinstance(libraries, Library):
                libraries = [libraries]
            try:
                for lib in libraries:
                    assert(isinstance(lib, Library))
            except:
                print('Expected Library object or list of Library objects passed as the libraries argument.')
                raise
            else:
                if not self.library:
                    self.library = libraries.pop(0)
                for lib in libraries:
                    self.library = self.library + lib

        if self.library:
            self.rates = self.rates + self.library.get_rates()

        # get the unique nuclei
        u = []
        for r in self.rates:
            t = set(r.reactants + r.products)
            u = set(list(u) + list(t))

        self.unique_nuclei = sorted(u)

        # now make a list of each rate that touches each nucleus
        # we'll store this in a dictionary keyed on the nucleus
        self.nuclei_consumed = OrderedDict()
        self.nuclei_produced = OrderedDict()

        for n in self.unique_nuclei:
            self.nuclei_consumed[n] = [r for r in self.rates if n in r.reactants]
            self.nuclei_produced[n] = [r for r in self.rates if n in r.products]

        # Re-order self.rates so Reaclib rates come first,
        # followed by Tabular rates. This is needed if
        # reaclib coefficients are targets of a pointer array
        # in the Fortran network.
        # It is desired to avoid wasting array size
        # storing meaningless Tabular coefficient pointers.
        self.rates = sorted(self.rates,
                            key=lambda r: r.chapter == 't')

        self.tabular_rates = []
        self.reaclib_rates = []
        for n, r in enumerate(self.rates):
            if r.chapter == 't':
                self.tabular_rates.append(n)
            elif isinstance(r.chapter, int):
                self.reaclib_rates.append(n)
            else:
                print('ERROR: Chapter type unknown for rate chapter {}'.format(
                    str(r.chapter)))
                exit()

    def _read_rate_files(self, rate_files):
        # get the rates
        self.files = rate_files
        for rf in self.files:
            try:
                rflib = Library(rf)
            except:
                print("Error reading library from file: {}".format(rf))
                raise
            else:
                if not self.library:
                    self.library = rflib
                else:
                    self.library = self.library + rflib

    def get_nuclei(self):
        """ get all the nuclei that are part of the network """
        return self.unique_nuclei

    def evaluate_rates(self, rho, T, composition):
        """evaluate the rates for a specific density, temperature, and
        composition"""
        rvals = OrderedDict()
        ys = composition.get_molar()

        for r in self.rates:
            val = r.prefactor * rho**r.dens_exp * r.eval(T)
            yfac = functools.reduce(mul, [ys[q] for q in r.reactants])
            rvals[r] = yfac * val

        return rvals

    def network_overview(self):
        """ return a verbose network overview """
        ostr = ""
        for n in self.unique_nuclei:
            ostr += "{}\n".format(n)
            ostr += "  consumed by:\n"
            for r in self.nuclei_consumed[n]:
                ostr += "     {}\n".format(r.string)

            ostr += "  produced by:\n"
            for r in self.nuclei_produced[n]:
                ostr += "     {}\n".format(r.string)

            ostr += "\n"
        return ostr

    def write_network(self, *args, **kwargs):
        """Before writing the network, check to make sure the rates
        are distinguishable by name."""
        assert self._distinguishable_rates(), "ERROR: Rates not uniquely identified by Rate.fname"
        self._write_network(*args, **kwargs)

    def _distinguishable_rates(self):
        """Every Rate in this RateCollection should have a unique Rate.fname,
        as the network writers distinguish the rates on this basis."""
        names = [r.fname for r in self.rates]
        return len(set(names)) == len(self.rates)

    def _write_network(self, *args, **kwargs):
        """A stub for function to output the network -- this is implementation
        dependent."""
        print('To create network integration source code, use a class that implements a specific network type.')
        return

    def plot(self, outfile=None, rho=None, T=None, comp=None, size=(800, 600), dpi=100):
        """Make a plot of the network structure showing the links between nuclei"""

        G = nx.MultiDiGraph()
        G.position = {}
        G.labels = {}

        fig, ax = plt.subplots()
        #divider = make_axes_locatable(ax)
        #cax = divider.append_axes('right', size='15%', pad=0.05)

        ax.plot([0, 0], [8, 8], 'b-')

        # nodes -- the node nuclei will be all of the heavies, but not
        # p, n, alpha, unless we have p + p, 3-a, etc.
        node_nuclei = []
        for n in self.unique_nuclei:
            if n.raw not in ["p", "n", "he4"]:
                node_nuclei.append(n)
            else:
                for r in self.rates:
                    if r.reactants.count(n) > 1:
                        node_nuclei.append(n)
                        break

        for n in node_nuclei:
            G.add_node(n)
            G.position[n] = (n.N, n.Z)
            G.labels[n] = r"${}$".format(n.pretty)

        if rho is not None and T is not None and comp is not None:
            ydots = self.evaluate_rates(rho, T, comp)
        else:
            ydots = None

        #for rr in ydots:
        #    print("{}: {}".format(rr, ydots[rr]))

        # edges
        for n in node_nuclei:
            for r in self.nuclei_consumed[n]:
                for p in r.products:
                    if p in node_nuclei:
                        # networkx doesn't seem to keep the edges in
                        # any particular order, so we associate data
                        # to the edges here directly, in this case,
                        # the reaction rate, which will be used to
                        # color it
                        if ydots is None:
                            G.add_edges_from([(n, p)], weight=0.5)
                        else:
                            try:
                                rate_weight = math.log10(ydots[r])
                            except ValueError:
                                # if ydots[r] is zero, then set the weight
                                # to roughly the minimum exponent possible
                                # for python floats
                                rate_weight = -308
                            except:
                                raise
                            G.add_edges_from([(n, p)], weight=rate_weight)

        nx.draw_networkx_nodes(G, G.position,
                               node_color="#A0CBE2", alpha=1.0,
                               node_shape="o", node_size=1000, linewidth=2.0, zorder=10, ax=ax)

        nx.draw_networkx_labels(G, G.position, G.labels,
                                font_size=13, font_color="w", zorder=100, ax=ax)

        # get the edges and weights coupled in the same order
        edges, weights = zip(*nx.get_edge_attributes(G, 'weight').items())

        edges_lc = nx.draw_networkx_edges(G, G.position, width=3,
                                          edgelist=edges, edge_color=weights,
                                          node_size=1000,
                                          edge_cmap=plt.cm.viridis, zorder=1, ax=ax)

        # for networkx <= 2.0 draw_networkx_edges returns a
        # LineCollection matplotlib type which we can use for the
        # colorbar directly.  For networkx >= 2.1, it is a collection
        # of FancyArrowPatch-s, which we need to run through a
        # PatchCollection.  See:
        # https://stackoverflow.com/questions/18658047/adding-a-matplotlib-colorbar-from-a-patchcollection

        if ydots is not None:
            pc = mpl.collections.PatchCollection(edges_lc, cmap=plt.cm.viridis)
            pc.set_array(weights)
            plt.colorbar(pc, label="log10(rate)")

        Ns = [n.N for n in node_nuclei]
        Zs = [n.Z for n in node_nuclei]

        plt.xlim(min(Ns)-1, max(Ns)+1)
        #plt.ylim(min(Zs)-1, max(Zs)+1)
        plt.xlabel(r"$N$", fontsize="large")
        plt.ylabel(r"$Z$", fontsize="large")

        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')

        ax.set_aspect("equal", "datalim")

        fig.set_size_inches(size[0]/dpi, size[1]/dpi)

        if outfile is None:
            plt.show()
        else:
            plt.tight_layout()
            plt.savefig(outfile, dpi=dpi)

    def __repr__(self):
        string = ""
        for r in self.rates:
            string += "{}\n".format(r.string)
        return string

class Explorer(object):
    """ interactively explore a rate collection """
    def __init__(self, rc, comp, size=(800, 600)):
        """ take a RateCollection and a composition """
        self.rc = rc
        self.comp = comp
        self.size = size

    def _make_plot(self, logrho, logT):
        self.rc.plot(rho=10.0**logrho, T=10.0**logT, comp=self.comp, size=self.size)

    def explore(self, logrho=(2, 6, 0.1), logT=(7, 9, 0.1)):
        """Perform interactive exploration of the network structure."""
        interact(self._make_plot, logrho=logrho, logT=logT)
<|endoftext|>


<commit_before>version_info = (1, 2, 1)
__version__ = '.'.join(map(str, version_info))
<commit_msg>Back to work<commit_after>version_info = (1, 3, 0, 'dev0')
__version__ = '.'.join(map(str, version_info))
<|endoftext|>


<commit_before>import numpy.linalg
import numpy.random
import scipy.stats
import scipy.io
import numpy
import math
import sys
import os

###############################################################################

def mat_to_stdout (mat):

    for i in range(mat.shape[0]):
       for j in range(mat.shape[1]):
           sys.stdout.write ("%f "%mat[i][j])
       sys.stdout.write ("\n")

###############################################################################

def progress_bar (count, total, status=''):
    bar_len = 60
    filled_len = int(round(bar_len * count / float(total)))

    percents = round(100.0 * count / float(total), 1)
    bar = '=' * filled_len + '-' * (bar_len - filled_len)

    sys.stdout.write('[%s] %s%s ...%s\r' % (bar, percents, '%', status))
    sys.stdout.flush()

###############################################################################

filename1 = "ms.mat"
filename2 = "bp.mat"
step = 0.25
run = 100000
tprev = 37 # mesi previsione
namems = 'sp'
namebp = 'i_r'

if len(sys.argv) != 6:
    print "usage: ", sys.argv[0], " msmatfilename bpmatfilename step tprev run"
    exit(1)
else:
    filename1 = sys.argv[1]
    filename2 = sys.argv[2]
    step = float(sys.argv[3])
    tprev = int(sys.argv[4])
    run = int(sys.argv[5])

#numpy.random.seed(9001)

msd = scipy.io.loadmat(filename1)
bpd = scipy.io.loadmat(filename2)

if not(namems in msd.keys()):
    print "Cannot find ", namems, " in ", filename1
    print msd.keys()
    exit(1)

if not(namebp in bpd.keys()):
    print "Cannot find ", namebp, " in ", filename2
    print bpd.keys()
    exit(1)

if msd[namems].shape[0] != bpd[namebp].shape[0]:
    print "wrong dim of the input matrix"
    exit(1)

countries = msd[namems].shape[0]
rating = numpy.max(msd[namems])

ms = msd[namems]
i_r = bpd[namebp]
time = len(ms[1,:])

Nk = numpy.zeros((rating,rating,countries), dtype='int64')
Num = numpy.zeros((rating,rating), dtype='int64')
Den = numpy.zeros(rating, dtype='int64')
Pr = numpy.zeros((rating,rating), dtype='float64')

verbose = False

for k in range(countries):
    for t in range(time-1):
        for i in range(rating):
            for j in range(rating):
                if (ms[k][t] == (i+1)) and (ms[k][t+1] == (j+1)):
                    Nk[i][j][k] = Nk[i][j][k] + 1

                Num[i][j] = sum(Nk[i][j])

            Den[i] = sum(Num[i])

    progress_bar(k+1, countries)

#print Num
#print Den

for i in range(rating):
    for j in range(rating):
        if Den[i] != 0:
            Pr[i][j] = float(Num[i][j])/float(Den[i])
        else:
            Pr[i][j] = 0.0

newPr = Pr - numpy.identity(rating, dtype='float64')

#print newPr

s, v, d = numpy.linalg.svd(newPr)
#print numpy.mean(v)

for i in range(len(i_r)):
    for j in range(len(i_r[0])):
        if math.isnan(i_r[i][j]):
           i_r[i][j] = float('inf')

benchmark = numpy.amin(i_r, 0)

r = numpy.zeros((countries,time), dtype='float64')

for k in range(countries):
    for Time in range(time):
        r[k][Time] = i_r[k][Time] - benchmark[Time]

for i in range(len(r)):
    for j in range(len(r[0])):
        if (r[i][j] == float('Inf')):
           r[i][j] = float('nan')

ist = numpy.zeros((rating,time*countries), dtype='float64')
Nn = numpy.zeros((rating), dtype='int')

for i in range(rating):
    for k in range(countries):
        for Time in range(time):
            if ms[k][Time] == i+1:
                Nn[i] = Nn[i] + 1
                ist[i][Nn[i]-1] = r[k][Time]

#Hfile = scipy.io.loadmat("ist.mat")
#ist1 = Hfile['ist']

Mean = numpy.zeros((rating), dtype='float64')
y = numpy.zeros((ist.shape[0], Nn[0]), dtype='float64')
for i in range(len(ist)):
    y[i] = ist[i][0:Nn[0]]

#Yfile = scipy.io.loadmat("y.mat")
#y = Yfile['y']

aaa = y[0][:Nn[0]]
aaa = aaa[numpy.isfinite(aaa)]

aa = []
a = []
bbb = []
bb = []
b = []
cc = []
d = []

if rating > 1:
  aa = y[1][:Nn[1]]
  aa = aa[numpy.isfinite(aa)]

if rating > 2:
 a = y[2][:Nn[2]]
 a = a[numpy.isfinite(a)]

if rating > 3:
  bbb = y[3][:Nn[3]]
  bbb = bbb[numpy.isfinite(bbb)]

if rating > 4:
  bb = y[4][:Nn[4]]
  bb = bb[numpy.isfinite(bb)]

if rating > 5:
  b = y[5][:Nn[5]]
  b = b[numpy.isfinite(b)]

if rating > 6:
  cc = y[6][:Nn[6]]
  cc = cc[numpy.isfinite(cc)]

if rating > 7:
  d = y[rating-1][:Nn[7]]

#ABCDfile = scipy.io.loadmat("yabcd.mat")
#aaa1 = ABCDfile['aaa'][0]
#aa1 = ABCDfile['aa'][0]
#bbb1 = ABCDfile['bbb'][0]
#bb1 = ABCDfile['bb'][0]
#b1 = ABCDfile['b'][0]
#cc1 = ABCDfile['cc'][0]
#a1 = ABCDfile['a'][0]
#d1 = ABCDfile['d'][0]

minh = min(aaa)
maxh = max(aaa)
nbins = int ((maxh - minh) / step)
AAA, xh = numpy.histogram(aaa, bins=nbins, range=(minh, maxh))
Paaa = AAA/float(sum(AAA))

minh = min(aa)
maxh = max(aa)
nbins = int ((maxh - minh) / step)
AA, xh = numpy.histogram(aa, bins=nbins, range=(minh, maxh));
Paa = AA/float(sum(AA))

minh = min(a)
maxh = max(a)
nbins = int ((maxh - minh) / step)
A, xh = numpy.histogram(a, bins=nbins, range=(minh, maxh));
Pa = A/float(sum(A))

minh = min(bbb)
maxh = max(bbb)
nbins = int ((maxh - minh) / step)
BBB, xh = numpy.histogram(bbb, bins=nbins, range=(minh, maxh));
Pbbb = BBB/float(sum(BBB))

minh = min(bb)
maxh = max(bb)
nbins = int ((maxh - minh) / step)
BB, xh = numpy.histogram(bb, bins=nbins, range=(minh, maxh));
Pbb = BB/float(sum(BB))

minh = min(b)
maxh = max(b)
nbins = int ((maxh - minh) / step)
B, xh = numpy.histogram(b, bins=nbins, range=(minh, maxh));
Pb = B/float(sum(B))

minh = min(cc)
maxh = max(cc)
nbins = int ((maxh - minh) / step)
CC, xh = numpy.histogram(cc, bins=nbins, range=(minh, maxh));
Pcc = CC/float(sum(CC))

#ALLHfile = scipy.io.loadmat("allhist.mat")
#Paaa = ALLHfile['Paaa'][0]
#Paa = ALLHfile['Paa'][0]
#Pa = ALLHfile['Pa'][0]
#Pbbb = ALLHfile['Pbbb'][0]
#Pbb = ALLHfile['Pbb'][0]
#Pb = ALLHfile['Pb'][0]
#Pcc = ALLHfile['Pcc'][0]

Taaa = sum((Paaa*(math.log(len(Paaa)))*Paaa))
Taa = sum((Paa*(math.log(len(Paa)))*Paa))
Ta = sum((Pa*(math.log(len(Pa)))*Pa))
Tbbb = sum((Pbbb*(math.log(len(Pbbb)))*Pbbb))
Tbb = sum((Pbb*(math.log(len(Pbb)))*Pbb))
Tb = sum((Pb*(math.log(len(Pb)))*Pb))
Tcc = sum((Pcc*(math.log(len(Pcc)))*Pcc))

Mean = [numpy.mean(aaa),numpy.mean(aa), \
        numpy.mean(a),numpy.mean(bbb),numpy.mean(bb), \
        numpy.mean(b),numpy.mean(cc)]

fval, pval = scipy.stats.f_oneway (aaa, aa, a, bbb, bb, b, cc)

print " "

oufilename = "1wayanova_"+str(run)+".txt"

if os.path.exists(oufilename):
    os.remove(oufilename)

outf = open(oufilename, "w")

outf.write("F-value: %f\n"%fval)
outf.write("P value: %f\n"%pval)

outf.close()

Ti = [Taaa, Taa, Ta, Tbbb, Tbb, Tb, Tcc]

s_t = numpy.zeros((countries,time), dtype='float64')

for i in range(r.shape[0]):
    for j in range(r.shape[1]):
        if math.isnan(r[i][j]):
            r[i][j] = 0.0

R_t = numpy.sum(r, axis=0)
T_t = numpy.zeros(time, dtype='float64')

for t in range(time):
    for k in range(countries):
        s_t[k][t] = r[k][t] / R_t[t]
        if s_t[k][t] != 0:
            T_t[t] += s_t[k][t]*math.log(float(countries) * s_t[k][t])

X = numpy.random.rand(countries,tprev,run)
cdf = numpy.zeros((rating,rating), dtype='float64')

#Xfile = scipy.io.loadmat("XcdfPr.mat")
#X = Xfile['X']
#cdf = Xfile['cdf']
#Pr = Xfile['Pr']
#Mean = Xfile['Mean'][0]
#Ti = Xfile['Ti'][0]

x = numpy.zeros((countries,tprev,run), dtype='int')
bp = numpy.zeros((countries,tprev,run), dtype='float64')
xm = numpy.zeros((countries,tprev), dtype='float64')
r_prev = numpy.zeros((tprev,run), dtype='float64')
Var = numpy.zeros((tprev), dtype='float64')
tot = numpy.zeros((rating,tprev,run), dtype='float64')
cont = numpy.zeros((rating,tprev,run), dtype='int')
ac = numpy.zeros((rating,tprev,run), dtype='float64')
t1 = numpy.zeros((tprev,run), dtype='float64')
t2 = numpy.zeros((tprev,run), dtype='float64')
term = numpy.zeros((tprev,run), dtype='float64')
entr = numpy.zeros((tprev,run), dtype='float64')
entropia = numpy.zeros(tprev, dtype='float64')
R_prev = numpy.zeros(tprev, dtype='float64')

for j in range(run):

    # da controllare
    for i in range(countries):
        x[i][0][j] = ms[i][time-1]

    for i in range (rating):
        cdf[i][0] = Pr[i][0]

    for i in range(rating):
        for k in range(1,rating):
            cdf[i][k] = Pr[i][k] + cdf[i][k-1]

    for c in range(countries):
        if X[c][0][j] <= cdf[x[c][0][j]-1][0]:
            x[c][1][j] = 1

        for k in range(1,rating):
            if (cdf[x[c][0][j]-1][k-1] < X[c][0][j]) and \
                    (X[c][0][j] <= cdf[x[c][0][j]-1][k] ):
               x[c][1][j] = k + 1

        for t in range(2,tprev):
            if X[c][t-1][j] <= cdf[x[c][t-1][j]-1][0]:
                x[c][t][j] = 1

            for k in range(1,rating):
                if (cdf[x[c][t-1][j]-1][k-1] < X[c][t-1][j]) \
                        and (X[c][t-1][j] <= cdf[x[c][t-1][j]-1][k]):
                  x[c][t][j] = k + 1

    for t in range(tprev):
        for c in range(countries):
            for i in range(rating):
                if x[c][t][j] == i+1:
                    bp[c][t][j] = Mean[i]
                    cont[i][t][j] = cont[i][t][j] + 1
                    tot[i][t][j] = cont[i][t][j] * Mean[i]

        summa = 0.0
        for a in range(bp.shape[0]):
            summa += bp[a][t][j]
        r_prev[t][j] = summa

    for t in range(tprev):
        for i in range(rating):
             ac[i][t][j] = tot[i][t][j]/r_prev[t][j]
             if ac[i][t][j] != 0.0:
                 t1[t][j] += (ac[i][t][j]*Ti[i])
                 t2[t][j] += (ac[i][t][j]*math.log(float(rating)*ac[i][t][j]))
                 if cont[i][t][j] != 0:
                    term[t][j] += ac[i][t][j]* \
                            math.log(float(countries)/(float(rating)*cont[i][t][j]))

        entr[t][j] = t1[t][j] + t2[t][j] + term[t][j]
        R_prev[t] = numpy.mean(r_prev[t][j])

    progress_bar(j+1, run)

print " "

oufilename = "entropy_"+str(run)+".txt"

if os.path.exists(oufilename):
    os.remove(oufilename)

outf = open(oufilename, "w")

for t in range(tprev):
    entropia[t] =numpy.mean(entr[t])
    Var[t] = numpy.std(entr[t])

for t in range(tprev):
    outf.write("%d %f %f \n"%(t+1, entropia[t], Var[t]))

outf.close()
<commit_msg>adding some extra check and started to reorganize the rating stuff<commit_after>import numpy.linalg
import numpy.random
import scipy.stats
import scipy.io
import numpy
import math
import sys
import os

###############################################################################

def mat_to_stdout (mat):

    for i in range(mat.shape[0]):
       for j in range(mat.shape[1]):
           sys.stdout.write ("%f "%mat[i][j])
       sys.stdout.write ("\n")

###############################################################################

def progress_bar (count, total, status=''):
    bar_len = 60
    filled_len = int(round(bar_len * count / float(total)))

    percents = round(100.0 * count / float(total), 1)
    bar = '=' * filled_len + '-' * (bar_len - filled_len)

    sys.stdout.write('[%s] %s%s ...%s\r' % (bar, percents, '%', status))
    sys.stdout.flush()

###############################################################################

def get_histo (v, step):

    minh = min(v)
    maxh = max(v)
    nbins = int ((maxh - minh) / step)
    h, xh = numpy.histogram(v, bins=nbins, range=(minh, maxh))
    p = h/float(sum(h))

    t = sum((p*(math.log(len(p)))*p))

    return p, t

###############################################################################

filename1 = "ms.mat"
filename2 = "bp.mat"
step = 0.25
run = 100000
tprev = 37 # mesi previsione
namems = 'ms'
namebp = 'i_r'

if len(sys.argv) != 6:
    print "usage: ", sys.argv[0], " msmatfilename bpmatfilename step tprev run"
    exit(1)
else:
    filename1 = sys.argv[1]
    filename2 = sys.argv[2]
    step = float(sys.argv[3])
    tprev = int(sys.argv[4])
    run = int(sys.argv[5])

numpy.random.seed(9001)

msd = scipy.io.loadmat(filename1)
bpd = scipy.io.loadmat(filename2)

if not(namems in msd.keys()):
    print "Cannot find ", namems, " in ", filename1
    print msd.keys()
    exit(1)

if not(namebp in bpd.keys()):
    print "Cannot find ", namebp, " in ", filename2
    print bpd.keys()
    exit(1)

if msd[namems].shape[0] != bpd[namebp].shape[0]:
    print "wrong dim of the input matrix"
    exit(1)

countries = msd[namems].shape[0]
rating = numpy.max(msd[namems])

if (rating <= 0) or (rating > 8):
    print "rating ", rating, " is not a valid value"
    exit(1)

ms = msd[namems]
i_r = bpd[namebp]
time = len(ms[1,:])

Nk = numpy.zeros((rating,rating,countries), dtype='int64')
Num = numpy.zeros((rating,rating), dtype='int64')
Den = numpy.zeros(rating, dtype='int64')
Pr = numpy.zeros((rating,rating), dtype='float64')

verbose = False

for k in range(countries):
    for t in range(time-1):
        for i in range(rating):
            for j in range(rating):
                if (ms[k][t] == (i+1)) and (ms[k][t+1] == (j+1)):
                    Nk[i][j][k] = Nk[i][j][k] + 1

                Num[i][j] = sum(Nk[i][j])

            Den[i] = sum(Num[i])

    progress_bar(k+1, countries)

#print Num
#print Den

for i in range(rating):
    for j in range(rating):
        if Den[i] != 0:
            Pr[i][j] = float(Num[i][j])/float(Den[i])
        else:
            Pr[i][j] = 0.0

newPr = Pr - numpy.identity(rating, dtype='float64')

#print newPr

s, v, d = numpy.linalg.svd(newPr)
#print numpy.mean(v)

for i in range(len(i_r)):
    for j in range(len(i_r[0])):
        if math.isnan(i_r[i][j]):
           i_r[i][j] = float('inf')

benchmark = numpy.amin(i_r, 0)

r = numpy.zeros((countries,time), dtype='float64')

for k in range(countries):
    for Time in range(time):
        r[k][Time] = i_r[k][Time] - benchmark[Time]

for i in range(len(r)):
    for j in range(len(r[0])):
        if (r[i][j] == float('Inf')):
           r[i][j] = float('nan')

ist = numpy.zeros((rating,time*countries), dtype='float64')
Nn = numpy.zeros((rating), dtype='int')

for i in range(rating):
    for k in range(countries):
        for Time in range(time):
            if ms[k][Time] == i+1:
                Nn[i] = Nn[i] + 1
                ist[i][Nn[i]-1] = r[k][Time]

Mean = numpy.zeros((rating), dtype='float64')
y = numpy.zeros((ist.shape[0], Nn[0]), dtype='float64')
for i in range(len(ist)):
    y[i] = ist[i][0:Nn[0]]

allratings = []
Mean = []
Ti = []

if rating > 0:
    aaa = y[0][:Nn[0]]
    aaa = aaa[numpy.isfinite(aaa)]
    Paaa, Taaa = get_histo(aaa, step)

    allratings.append(aaa)
    Mean.append(numpy.mean(aaa))
    Ti.append(Taaa)

if rating > 1:
    aa = y[1][:Nn[1]]
    aa = aa[numpy.isfinite(aa)]
    Paa, Taa = get_histo(aa, step)

    allratings.append(aa)
    Mean.append(numpy.mean(aa))
    Ti.append(Taa)

if rating > 2:
    a = y[2][:Nn[2]]
    a = a[numpy.isfinite(a)]
    Pa, Ta = get_histo(a, step)

    allratings.append(a)
    Mean.append(numpy.mean(a))
    Ti.append(Ta)

if rating > 3:
    bbb = y[3][:Nn[3]]
    bbb = bbb[numpy.isfinite(bbb)]
    Pbbb, Tbbb = get_histo(bbb, step)

    allratings.append(bbb)
    Mean.append(numpy.mean(bbb))
    Ti.append(Tbbb)

if rating > 4:
    bb = y[4][:Nn[4]]
    bb = bb[numpy.isfinite(bb)]
    Pbb, Tbb = get_histo(bb, step)

    allratings.append(bb)
    Mean.append(numpy.mean(bb))
    Ti.append(Tbb)

if rating > 5:
    b = y[5][:Nn[5]]
    b = b[numpy.isfinite(b)]
    Pb, Tb = get_histo(b, step)

    allratings.append(b)
    Mean.append(numpy.mean(b))
    Ti.append(Tb)

if rating > 6:
    cc = y[6][:Nn[6]]
    cc = cc[numpy.isfinite(cc)]
    Pcc, Tcc = get_histo(cc, step)

    allratings.append(cc)
    Mean.append(numpy.mean(cc))
    Ti.append(Tcc)

if rating > 7:
    d = y[rating-1][:Nn[7]]
    allratings.append(d)
    Pd, Td = get_histo(d, step)

    allratings.append(d)
    Mean.append(numpy.mean(d))
    Ti.append(Td)

#fval, pval = scipy.stats.f_oneway (aaa, aa, a, bbb, bb, b, cc)

#print " "

#oufilename = "1wayanova_"+str(run)+".txt"

#if os.path.exists(oufilename):
#    os.remove(oufilename)

#outf = open(oufilename, "w")

#outf.write("F-value: %f\n"%fval)
#outf.write("P value: %f\n"%pval)

#outf.close()

s_t = numpy.zeros((countries,time), dtype='float64')

for i in range(r.shape[0]):
    for j in range(r.shape[1]):
        if math.isnan(r[i][j]):
            r[i][j] = 0.0

R_t = numpy.sum(r, axis=0)
T_t = numpy.zeros(time, dtype='float64')

for t in range(time):
    for k in range(countries):
        s_t[k][t] = r[k][t] / R_t[t]
        if s_t[k][t] != 0:
            T_t[t] += s_t[k][t]*math.log(float(countries) * s_t[k][t])

X = numpy.random.rand(countries,tprev,run)
cdf = numpy.zeros((rating,rating), dtype='float64')

x = numpy.zeros((countries,tprev,run), dtype='int')
bp = numpy.zeros((countries,tprev,run), dtype='float64')
xm = numpy.zeros((countries,tprev), dtype='float64')
r_prev = numpy.zeros((tprev,run), dtype='float64')
Var = numpy.zeros((tprev), dtype='float64')
tot = numpy.zeros((rating,tprev,run), dtype='float64')
cont = numpy.zeros((rating,tprev,run), dtype='int')
ac = numpy.zeros((rating,tprev,run), dtype='float64')
t1 = numpy.zeros((tprev,run), dtype='float64')
t2 = numpy.zeros((tprev,run), dtype='float64')
term = numpy.zeros((tprev,run), dtype='float64')
entr = numpy.zeros((tprev,run), dtype='float64')
entropia = numpy.zeros(tprev, dtype='float64')
R_prev = numpy.zeros(tprev, dtype='float64')

for j in range(run):

    # da controllare
    for i in range(countries):
        x[i][0][j] = ms[i][time-1]

    for i in range (rating):
        cdf[i][0] = Pr[i][0]

    for i in range(rating):
        for k in range(1,rating):
            cdf[i][k] = Pr[i][k] + cdf[i][k-1]

    for c in range(countries):
        if X[c][0][j] <= cdf[x[c][0][j]-1][0]:
            x[c][1][j] = 1

        for k in range(1,rating):
            if (cdf[x[c][0][j]-1][k-1] < X[c][0][j]) and \
                    (X[c][0][j] <= cdf[x[c][0][j]-1][k] ):
               x[c][1][j] = k + 1

        for t in range(2,tprev):
            if X[c][t-1][j] <= cdf[x[c][t-1][j]-1][0]:
                x[c][t][j] = 1

            for k in range(1,rating):
                if (cdf[x[c][t-1][j]-1][k-1] < X[c][t-1][j]) \
                        and (X[c][t-1][j] <= cdf[x[c][t-1][j]-1][k]):
                  x[c][t][j] = k + 1

    for t in range(tprev):
        for c in range(countries):
            for i in range(rating):
                if x[c][t][j] == i+1:
                    bp[c][t][j] = Mean[i]
                    cont[i][t][j] = cont[i][t][j] + 1
                    tot[i][t][j] = cont[i][t][j] * Mean[i]

        summa = 0.0
        for a in range(bp.shape[0]):
            summa += bp[a][t][j]
        r_prev[t][j] = summa

    for t in range(tprev):
        for i in range(rating):
             ac[i][t][j] = tot[i][t][j]/r_prev[t][j]
             if ac[i][t][j] != 0.0:
                 t1[t][j] += (ac[i][t][j]*Ti[i])
                 t2[t][j] += (ac[i][t][j]*math.log(float(rating)*ac[i][t][j]))
                 if cont[i][t][j] != 0:
                    term[t][j] += ac[i][t][j]* \
                            math.log(float(countries)/(float(rating)*cont[i][t][j]))

        entr[t][j] = t1[t][j] + t2[t][j] + term[t][j]
        R_prev[t] = numpy.mean(r_prev[t][j])

    progress_bar(j+1, run)

print " "

oufilename = "entropy_"+str(run)+".txt"

if os.path.exists(oufilename):
    os.remove(oufilename)

outf = open(oufilename, "w")

for t in range(tprev):
    entropia[t] =numpy.mean(entr[t])
    Var[t] = numpy.std(entr[t])

for t in range(tprev):
    outf.write("%d %f %f \n"%(t+1, entropia[t], Var[t]))

outf.close()
<|endoftext|>