<commit_before>
import fnmatch
import getopt
import os
import os.path
import re
import sys
import unittest

def parse_cmdline():
    """ parses command line and returns test directories, verbosity, and test filter
        usage:
            runfiles.py  -v|--verbosity <level>  -f|--filter <regex>  dirs|files
    """
    verbosity = 2
    test_filter = None

    optlist, args = getopt.getopt(sys.argv[1:], "v:f:", ["verbosity=", "filter="])
    for opt,value in optlist:
        if opt in ("-v","--verbosity"):
            verbosity = value
        elif opt in ("-f","--filter"):
            if "," in value:
                test_filter = value.split(',')
            else:
                test_filter = [value]
    if type([]) != type(args):
        args = [args]
    return args, int(verbosity), test_filter

class PydevTestRunner:
    """ finds and runs a file or directory of files as a unit test """

    __py_extensions = ["*.py", "*.pyw"]
    __exclude_files = ["__init__.*"]

    def __init__(self, test_dir, test_filter=None, verbosity=2):
        self.test_dir = test_dir
        self.__adjust_path()
        self.test_filter = self.__setup_test_filter(test_filter)
        self.verbosity = verbosity

    def __adjust_path(self):
        """ add the current file or directory to the python path """
        path_to_append = None
        for n in xrange(len(self.test_dir)):
            orig = self.test_dir[n]
            dir_name = self.__unixify(self.test_dir[n])
            if os.path.isdir(dir_name):
                if not dir_name.endswith("/"):
                    self.test_dir[n] = dir_name + "/"
                path_to_append = os.path.normpath(dir_name)
            elif os.path.isfile(dir_name):
                path_to_append = os.path.dirname(dir_name)
            else:
                msg = ("unknown type. \n%s\nshould be file or a directory.\n" % (dir_name) )
                raise RuntimeError(msg)
        if path_to_append is not None:
            sys.path.insert(0, path_to_append)
        return

    def __setup_test_filter(self, test_filter):
        """ turn a filter string into a list of filter regexes """
        if test_filter is None or len(test_filter) == 0:
            return None
        return [re.compile("test%s"%f) for f in test_filter]

    def __is_valid_py_file(self, fname):
        """ tests that a particular file contains the proper file extension
            and is not in the list of files to exclude """
        is_valid_fname = 0
        for invalid_fname in self.__class__.__exclude_files:
            is_valid_fname += int(not fnmatch.fnmatch(fname, invalid_fname))
        if_valid_ext = 0
        for ext in self.__class__.__py_extensions:
            if_valid_ext += int(fnmatch.fnmatch(fname, ext))
        return is_valid_fname > 0 and if_valid_ext > 0

    def __unixify(self, s):
        """ stupid windows. converts the backslash to forwardslash for consistency """
        return os.path.normpath(s).replace( os.sep, "/" )

    def __importify(self, s):
        """ turns directory separators into dots and removes the ".py*" extension
            so the string can be used as import statement """
        dirname, fname = os.path.split(s)
        imp_stmt_pieces = [dirname.replace("\\","/").replace("/", "."),
                           os.path.splitext(fname)[0]]
        if len(imp_stmt_pieces[0]) == 0:
            imp_stmt_pieces = imp_stmt_pieces[1:]
        return ".".join(imp_stmt_pieces)

    def __add_files(self, pyfiles, root, files):
        """ if files match, appends them to pyfiles. used by os.path.walk fcn """
        for fname in files:
            if self.__is_valid_py_file(fname):
                name_without_base_dir = self.__unixify(os.path.join(root, fname)) #[len(base_dir):]
                pyfiles.append( name_without_base_dir )
        return

    def find_import_files(self):
        """ return a list of files to import """
        pyfiles = []
        for base_dir in self.test_dir:
            if os.path.isdir(base_dir):
                # argh, it would be nice to use os.walk, but jython2.1 is too old
                os.path.walk(base_dir, self.__add_files, pyfiles)
                for n in xrange(len(pyfiles)):
                    pyfiles[n] = pyfiles[n][len(base_dir):]
            elif os.path.isfile(base_dir):
                pyfiles.append( os.path.basename(base_dir) )
        return pyfiles

    def __get_module_from_str(self, modname):
        """ Import the module in the given import path.
            * Returns the "final" module, so importing "coilib40.subject.visu"
            return the "visu" module, not the "coilib40" as returned by __import__ """
        try:
            mod = __import__( modname )
            for part in modname.split('.')[1:]:
                mod = getattr(mod, part)
            return mod
        except ImportError:
            print >> sys.stderr, 'ERROR: Module :%s could not be imported (maybe the dir does not have __init__.py folders for all the packages?)' % (modname)
            return None

    def find_modules_from_files(self, pyfiles):
        """ returns a lisst of modules given a list of import strings """
        imports = [self.__importify(s) for s in pyfiles]
        return [self.__get_module_from_str(import_str) for import_str in imports if import_str is not None]

    def find_tests_from_modules(self, modules):
        """ returns the unittests given a list of modules """
        return [unittest.defaultTestLoader.loadTestsFromModule(m) for m in modules]

    def filter_tests(self, test_objs):
        """ based on a filter name, only return those tests that have
            the test case names that match """
        test_suite = []
        for test_obj in test_objs:
            if isinstance(test_obj, unittest.TestSuite):
                test_obj._tests = self.filter_tests(test_obj._tests)
                test_suite.append( test_obj )
            elif isinstance(test_obj, unittest.TestCase):
                test_cases = []
                for tc in test_objs:
                    if self.__match(self.test_filter, tc._TestCase__testMethodName):
                        test_cases.append( tc )
                return test_cases
        return test_suite

    def __match(self, filter_list, name):
        """ returns whether a test name matches the test filter """
        if filter_list is None:
            return 1
        for f in filter_list:
            if re.match(f, name):
                return 1
        return 0

    def run_tests(self):
        """ runs all tests """
        print "Finding files...",
        files = self.find_import_files()
        print self.test_dir, '... done'
        print "Importing test modules ...",
        modules = self.find_modules_from_files(files)
        print "done."
        all_tests = self.find_tests_from_modules(modules)
        if self.test_filter is not None:
            print 'Test Filter: %s' % [p.pattern for p in self.test_filter]
            all_tests = self.filter_tests(all_tests)
        print
        runner = unittest.TextTestRunner(stream=sys.stdout, descriptions=1, verbosity=verbosity)
        runner.run(unittest.TestSuite(all_tests))
        return

if __name__ == '__main__':
    dirs, verbosity, test_filter = parse_cmdline()
    PydevTestRunner(dirs, test_filter, verbosity).run_tests()
<commit_msg>*** empty log message ***<commit_after>
import fnmatch
import getopt
import os
import os.path
import re
import sys
import unittest

def parse_cmdline():
    """ parses command line and returns test directories, verbosity, and test filter
        usage:
            runfiles.py  -v|--verbosity <level>  -f|--filter <regex>  dirs|files
    """
    verbosity = 2
    test_filter = None

    optlist, args = getopt.getopt(sys.argv[1:], "v:f:", ["verbosity=", "filter="])
    for opt,value in optlist:
        if opt in ("-v","--verbosity"):
            verbosity = value
        elif opt in ("-f","--filter"):
            if "," in value:
                test_filter = value.split(',')
            else:
                test_filter = [value]
    if type([]) != type(args):
        args = [args]
    return args, int(verbosity), test_filter

class PydevTestRunner:
    """ finds and runs a file or directory of files as a unit test """

    __py_extensions = ["*.py", "*.pyw"]
    __exclude_files = ["__init__.*"]

    def __init__(self, test_dir, test_filter=None, verbosity=2):
        self.test_dir = test_dir
        self.__adjust_path()
        self.test_filter = self.__setup_test_filter(test_filter)
        self.verbosity = verbosity

    def __adjust_path(self):
        """ add the current file or directory to the python path """
        path_to_append = None
        for n in xrange(len(self.test_dir)):
            orig = self.test_dir[n]
            dir_name = self.__unixify(self.test_dir[n])
            if os.path.isdir(dir_name):
                if not dir_name.endswith("/"):
                    self.test_dir[n] = dir_name + "/"
                path_to_append = os.path.normpath(dir_name)
            elif os.path.isfile(dir_name):
                path_to_append = os.path.dirname(dir_name)
            else:
                msg = ("unknown type. \n%s\nshould be file or a directory.\n" % (dir_name) )
                raise RuntimeError(msg)
        if path_to_append is not None:
            sys.path.insert(0, path_to_append)
        return

    def __setup_test_filter(self, test_filter):
        """ turn a filter string into a list of filter regexes """
        if test_filter is None or len(test_filter) == 0:
            return None
        return [re.compile("test%s"%f) for f in test_filter]

    def __is_valid_py_file(self, fname):
        """ tests that a particular file contains the proper file extension
            and is not in the list of files to exclude """
        is_valid_fname = 0
        for invalid_fname in self.__class__.__exclude_files:
            is_valid_fname += int(not fnmatch.fnmatch(fname, invalid_fname))
        if_valid_ext = 0
        for ext in self.__class__.__py_extensions:
            if_valid_ext += int(fnmatch.fnmatch(fname, ext))
        return is_valid_fname > 0 and if_valid_ext > 0

    def __unixify(self, s):
        """ stupid windows. converts the backslash to forwardslash for consistency """
        return os.path.normpath(s).replace( os.sep, "/" )

    def __importify(self, s):
        """ turns directory separators into dots and removes the ".py*" extension
            so the string can be used as import statement """
        dirname, fname = os.path.split(s)
        imp_stmt_pieces = [dirname.replace("\\","/").replace("/", "."),
                           os.path.splitext(fname)[0]]
        if len(imp_stmt_pieces[0]) == 0:
            imp_stmt_pieces = imp_stmt_pieces[1:]
        return ".".join(imp_stmt_pieces)

    def __add_files(self, pyfiles, root, files):
        """ if files match, appends them to pyfiles. used by os.path.walk fcn """
        for fname in files:
            if self.__is_valid_py_file(fname):
                name_without_base_dir = self.__unixify(os.path.join(root, fname)) #[len(base_dir):]
                pyfiles.append( name_without_base_dir )
        return

    def find_import_files(self):
        """ return a list of files to import """
        pyfiles = []
        for base_dir in self.test_dir:
            if os.path.isdir(base_dir):
                # argh, it would be nice to use os.walk, but jython2.1 is too old
                os.path.walk(base_dir, self.__add_files, pyfiles)
                for n in xrange(len(pyfiles)):
                    pyfiles[n] = pyfiles[n][len(base_dir):]
            elif os.path.isfile(base_dir):
                pyfiles.append( os.path.basename(base_dir) )
        return pyfiles

    def __get_module_from_str(self, modname):
        """ Import the module in the given import path.
            * Returns the "final" module, so importing "coilib40.subject.visu"
            return the "visu" module, not the "coilib40" as returned by __import__ """
        try:
            mod = __import__( modname )
            for part in modname.split('.')[1:]:
                mod = getattr(mod, part)
            return mod
        except ImportError:
            import traceback;traceback.print_exc()
            print >> sys.stderr, 'ERROR: Module :%s could not be imported (alternative reason: the dir does not have __init__.py folders for all the packages?)' % (modname)
            return None

    def find_modules_from_files(self, pyfiles):
        """ returns a lisst of modules given a list of import strings """
        imports = [self.__importify(s) for s in pyfiles]
        return [self.__get_module_from_str(import_str) for import_str in imports if import_str is not None]

    def find_tests_from_modules(self, modules):
        """ returns the unittests given a list of modules """
        return [unittest.defaultTestLoader.loadTestsFromModule(m) for m in modules]

    def filter_tests(self, test_objs):
        """ based on a filter name, only return those tests that have
            the test case names that match """
        test_suite = []
        for test_obj in test_objs:
            if isinstance(test_obj, unittest.TestSuite):
                test_obj._tests = self.filter_tests(test_obj._tests)
                test_suite.append( test_obj )
            elif isinstance(test_obj, unittest.TestCase):
                test_cases = []
                for tc in test_objs:
                    if self.__match(self.test_filter, tc._TestCase__testMethodName):
                        test_cases.append( tc )
                return test_cases
        return test_suite

    def __match(self, filter_list, name):
        """ returns whether a test name matches the test filter """
        if filter_list is None:
            return 1
        for f in filter_list:
            if re.match(f, name):
                return 1
        return 0

    def run_tests(self):
        """ runs all tests """
        print "Finding files...",
        files = self.find_import_files()
        print self.test_dir, '... done'
        print "Importing test modules ...",
        modules = self.find_modules_from_files(files)
        print "done."
        all_tests = self.find_tests_from_modules(modules)
        if self.test_filter is not None:
            print 'Test Filter: %s' % [p.pattern for p in self.test_filter]
            all_tests = self.filter_tests(all_tests)
        print
        runner = unittest.TextTestRunner(stream=sys.stdout, descriptions=1, verbosity=verbosity)
        runner.run(unittest.TestSuite(all_tests))
        return

if __name__ == '__main__':
    dirs, verbosity, test_filter = parse_cmdline()
    PydevTestRunner(dirs, test_filter, verbosity).run_tests()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import logging
import operator
import time
import numpy

from openquake.baselib import hdf5, datastore, parallel, performance, general
from openquake.baselib.python3compat import zip, encode
from openquake.calculators import base, event_based, getters
from openquake.calculators.export.loss_curves import get_loss_builder

U8 = numpy.uint8
U16 = numpy.uint16
U32 = numpy.uint32
F32 = numpy.float32
F64 = numpy.float64
U64 = numpy.uint64

def get_assets_ratios(assets, riskmodel, gmvs, imts):
    """
    :param assets: a list of assets on the same site
    :param riskmodel: a CompositeRiskModel instance
    :params gmvs: hazard on the given site, shape (E, M)
    :param imts: intensity measure types
    :returns: a list of (assets, loss_ratios)
    """
    imti = {imt: i for i, imt in enumerate(imts)}
    tdict = riskmodel.get_taxonomy_dict()  # taxonomy -> taxonomy index
    assets_by_t = general.groupby(assets, operator.attrgetter('taxonomy'))
    assets_ratios = []
    for taxo, rm in riskmodel.items():
        t = tdict[taxo]
        try:
            assets = assets_by_t[t]
        except KeyError:  # there are no assets of taxonomy taxo
            continue
        assets_ratios.append((assets, rm.get_loss_ratios(gmvs, imti)))
    return assets_ratios

def start_ebrisk(rupgetter, srcfilter, param, monitor):
    """
    Launcher for ebrisk tasks
    """
    with monitor('weighting ruptures'):
        rupgetter.set_weights(srcfilter, param['num_taxonomies'])
    if rupgetter.weights.sum() <= param['maxweight']:
        yield ebrisk(rupgetter, srcfilter, param, monitor)
    else:
        for rgetter in rupgetter.split(param['maxweight']):
            yield ebrisk, rgetter, srcfilter, param

def ebrisk(rupgetter, srcfilter, param, monitor):
    """
    :param rupgetter:
        a RuptureGetter instance
    :param srcfilter:
        a SourceFilter instance
    :param param:
        a dictionary of parameters
    :param monitor:
        :class:`openquake.baselib.performance.Monitor` instance
    :returns:
        an ArrayWrapper with shape (E, L, T, ...)
    """
    riskmodel = param['riskmodel']
    L = len(riskmodel.lti)
    N = len(srcfilter.sitecol.complete)
    mon = monitor('getting assets', measuremem=False)
    with datastore.read(srcfilter.filename) as dstore:
        assgetter = getters.AssetGetter(dstore)
    getter = getters.GmfGetter(rupgetter, srcfilter, param['oqparam'])
    with monitor('getting hazard'):
        getter.init()  # instantiate the computers
        hazard = getter.get_hazard()  # sid -> (rlzi, sid, eid, gmv)
    mon_risk = monitor('computing losses', measuremem=False)
    mon_agg = monitor('aggregating losses', measuremem=False)
    imts = getter.imts
    events = rupgetter.get_eid_rlz()
    eid2idx = {eid: idx for idx, eid in enumerate(events['eid'])}
    tagnames = param['aggregate_by']
    shape = assgetter.tagcol.agg_shape((len(events), L), tagnames)
    elt_dt = [('eid', U64), ('rlzi', U16), ('loss', (F32, shape[1:]))]
    acc = numpy.zeros(shape, F32)  # shape (E, L, T...)
    if param['avg_losses']:
        losses_by_A = numpy.zeros((assgetter.num_assets, L), F32)
    else:
        losses_by_A = None
    times = numpy.zeros(N)  # risk time per site_id
    for sid, haz in hazard.items():
        t0 = time.time()
        weights = getter.weights[haz['rlzi']]
        assets_on_sid, tagidxs = assgetter.get(sid, tagnames)
        eidx = [eid2idx[eid] for eid in haz['eid']]
        mon.duration += time.time() - t0
        mon.counts += 1
        with mon_risk:
            assets_ratios = get_assets_ratios(
                assets_on_sid, riskmodel, haz['gmv'], imts)
        with mon_agg:
            for assets, triples in assets_ratios:
                for lti, (lt, imt, loss_ratios) in enumerate(triples):
                    w = weights[imt]
                    for asset in assets:
                        aid = asset.ordinal
                        losses = loss_ratios * asset.value(lt)
                        acc[(eidx, lti) + tagidxs[aid]] += losses
                        if param['avg_losses']:
                            losses_by_A[aid, lti] += losses @ w
            times[sid] = time.time() - t0
    with monitor('building event loss table'):
        elt = numpy.fromiter(
            ((event['eid'], event['rlz'], losses)
             for event, losses in zip(events, acc) if losses.sum()), elt_dt)
        agg = general.AccumDict(accum=numpy.zeros(shape[1:], F32))  # rlz->agg
        for rec in elt:
            agg[rec['rlzi']] += rec['loss'] * param['ses_ratio']
    return {'elt': elt, 'agg_losses': agg, 'losses_by_A': losses_by_A,
            'times': times}

@base.calculators.add('ebrisk')
class EbriskCalculator(event_based.EventBasedCalculator):
    """
    Event based PSHA calculator generating event loss tables
    """
    core_task = start_ebrisk
    is_stochastic = True
    precalc = 'event_based'
    accept_precalc = ['event_based', 'event_based_risk', 'ucerf_hazard']

    def pre_execute(self):
        self.oqparam.ground_motion_fields = False
        super().pre_execute()
        # save a copy of the assetcol in hdf5cache
        self.hdf5cache = self.datastore.hdf5cache()
        with hdf5.File(self.hdf5cache, 'w') as cache:
            cache['sitecol'] = self.sitecol.complete
            cache['assetcol'] = self.assetcol
        self.param['ses_ratio'] = self.oqparam.ses_ratio
        self.param['aggregate_by'] = self.oqparam.aggregate_by
        # initialize the riskmodel
        self.riskmodel.taxonomy = self.assetcol.tagcol.taxonomy
        self.param['riskmodel'] = self.riskmodel
        self.L = L = len(self.riskmodel.loss_types)
        A = len(self.assetcol)
        self.datastore.create_dset('avg_losses', F32, (A, L))
        shp = self.get_shape(L)  # shape L, T...
        elt_dt = [('eid', U64), ('rlzi', U16), ('loss', (F32, shp))]
        self.datastore.create_dset('losses_by_event', elt_dt)
        self.zerolosses = numpy.zeros(shp, F32)  # to get the multi-index
        shp = self.get_shape(self.L, self.R)  # shape L, R, T...
        self.datastore.create_dset('agg_losses-rlzs', F32, shp)

    def execute(self):
        oq = self.oqparam
        self.set_param(
            num_taxonomies=self.assetcol.num_taxonomies_by_site(),
            maxweight=oq.ebrisk_maxweight / (oq.concurrent_tasks or 1))
        parent = self.datastore.parent
        if parent:
            hdf5path = parent.filename
            grp_indices = parent['ruptures'].attrs['grp_indices']
            nruptures = len(parent['ruptures'])
        else:
            hdf5path = self.datastore.hdf5cache()
            grp_indices = self.datastore['ruptures'].attrs['grp_indices']
            nruptures = len(self.datastore['ruptures'])
            with hdf5.File(hdf5path, 'r+') as cache:
                self.datastore.hdf5.copy('weights', cache)
                self.datastore.hdf5.copy('ruptures', cache)
                self.datastore.hdf5.copy('rupgeoms', cache)
        self.init_logic_tree(self.csm_info)
        smap = parallel.Starmap(
            self.core_task.__func__, monitor=self.monitor())
        trt_by_grp = self.csm_info.grp_by("trt")
        samples = self.csm_info.get_samples_by_grp()
        rlzs_by_gsim_grp = self.csm_info.get_rlzs_by_gsim_grp()
        ruptures_per_block = numpy.ceil(nruptures / (oq.concurrent_tasks or 1))
        for grp_id, rlzs_by_gsim in rlzs_by_gsim_grp.items():
            start, stop = grp_indices[grp_id]
            for indices in general.block_splitter(
                    range(start, stop), ruptures_per_block):
                rgetter = getters.RuptureGetter(
                    hdf5path, list(indices), grp_id,
                    trt_by_grp[grp_id], samples[grp_id], rlzs_by_gsim)
                smap.submit(rgetter, self.src_filter, self.param)
        return smap.reduce(self.agg_dicts, numpy.zeros(self.N))

    def agg_dicts(self, acc, dic):
        """
        :param dummy: unused parameter
        :param dic: a dictionary with keys eids, losses, losses_by_N
        """
        self.oqparam.ground_motion_fields = False  # hack
        if len(dic['elt']):
            with self.monitor('saving losses_by_event', autoflush=True):
                self.datastore.extend('losses_by_event', dic['elt'])
        with self.monitor('saving agg_losses-rlzs', autoflush=True):
            for r, aggloss in dic['agg_losses'].items():
                self.datastore['agg_losses-rlzs'][:, r] += aggloss
        if self.oqparam.avg_losses:
            with self.monitor('saving avg_losses', autoflush=True):
                self.datastore['avg_losses'] += dic['losses_by_A']
        return acc + dic['times']

    def get_shape(self, *sizes):
        """
        :returns: a shape (S1, ... SN, T1 ... TN)
        """
        return self.assetcol.tagcol.agg_shape(sizes, self.oqparam.aggregate_by)

    def build_datasets(self, builder):
        oq = self.oqparam
        stats = oq.hazard_stats()
        S = len(stats)
        P = len(builder.return_periods)
        C = len(oq.conditional_loss_poes)
        loss_types = ' '.join(self.riskmodel.loss_types)
        if oq.individual_curves or self.R == 1:
            shp = self.get_shape(P, self.R, self.L)  # shape P, R, L, T...
            self.datastore.create_dset('agg_curves-rlzs', F32, shp)
            self.datastore.set_attrs(
                'agg_curves-rlzs', return_periods=builder.return_periods,
                loss_types=loss_types)
        if oq.conditional_loss_poes:
            shp = self.get_shape(C, self.R, self.L)  # shape C, R, L, T...
            self.datastore.create_dset('agg_maps-rlzs', F32, shp)
        if self.R > 1:
            shp = self.get_shape(P, S, self.L)  # shape P, S, L, T...
            self.datastore.create_dset('agg_curves-stats', F32, shp)
            self.datastore.set_attrs(
                'agg_curves-stats', return_periods=builder.return_periods,
                stats=[encode(name) for (name, func) in stats],
                loss_types=loss_types)
            if oq.conditional_loss_poes:
                shp = self.get_shape(C, S, self.L)  # shape C, S, L, T...
                self.datastore.create_dset('agg_maps-stats', F32, shp)
                self.datastore.set_attrs(
                    'agg_maps-stats',
                    stats=[encode(name) for (name, func) in stats],
                    loss_types=loss_types)

    def post_execute(self, times):
        """
        Compute and store average losses from the losses_by_event dataset,
        and then loss curves and maps.
        """
        self.datastore.set_attrs('task_info/start_ebrisk', times=times)
        logging.info('Building losses_by_rlz')
        elt_length = len(self.datastore['losses_by_event'])
        oq = self.oqparam
        builder = get_loss_builder(self.datastore)
        self.build_datasets(builder)
        mon = performance.Monitor(hdf5=hdf5.File(self.datastore.hdf5cache()))
        smap = parallel.Starmap(compute_loss_curves_maps, monitor=mon)
        self.datastore.close()
        acc = []
        ct = oq.concurrent_tasks or 1
        for elt_slice in general.split_in_slices(elt_length, ct):
            smap.submit(self.datastore.filename, elt_slice,
                        oq.conditional_loss_poes, oq.individual_curves)
        for res in smap:
            idx = res.pop('idx')
            for name, arr in res.items():
                if arr is not None:
                    acc.append((name, idx, arr))
        # copy performance information from the cache to the datastore
        pd = mon.hdf5['performance_data'].value
        hdf5.extend3(self.datastore.filename, 'performance_data', pd)
        self.datastore.open('r+')  # reopen
        self.datastore['task_info/compute_loss_curves_and_maps'] = (
            mon.hdf5['task_info/compute_loss_curves_maps'].value)
        with self.monitor('saving loss_curves and maps', autoflush=True):
            for name, idx, arr in acc:
                for ij, val in numpy.ndenumerate(arr):
                    self.datastore[name][ij + idx] = val

def compute_loss_curves_maps(filename, elt_slice, clp, individual_curves,
                             monitor):
    """
    :param filename: path to the datastore
    :param elt_slice: slice of the event loss table
    :param clp: conditional loss poes used to computed the maps
    :param individual_curves: if True, build the individual curves and maps
    :param monitor: a Monitor instance
    :yields:
        dictionaries with keys idx, agg_curves-rlzs, agg_curves-stats,
        agg_maps-rlzs, agg_maps-stats
    """
    with datastore.read(filename) as dstore:
        oq = dstore['oqparam']
        stats = oq.hazard_stats()
        builder = get_loss_builder(dstore)
        R = len(dstore['weights'])
        losses = [[] for _ in range(R)]
        elt = dstore['losses_by_event'][elt_slice]
        for rec in elt:
            losses[rec['rlzi']].append(rec['loss'])
    for multi_index, _ in numpy.ndenumerate(elt[0]['loss']):
        result = {'idx': multi_index}
        thelosses = [[ls[multi_index] for ls in loss] for loss in losses]
        result['agg_curves-rlzs'], result['agg_curves-stats'] = (
            builder.build_pair(thelosses, stats))
        if R > 1 and individual_curves is False:
            del result['agg_curves-rlzs']
        if clp:
            result['agg_maps-rlzs'], result['agg_maps-stats'] = (
                builder.build_loss_maps(thelosses, clp, stats))
            if R > 1 and individual_curves is False:
                del result['agg_maps-rlzs']
        yield result
<commit_msg>Remove logging [skip CI]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import logging
import operator
import time
import numpy

from openquake.baselib import hdf5, datastore, parallel, performance, general
from openquake.baselib.python3compat import zip, encode
from openquake.calculators import base, event_based, getters
from openquake.calculators.export.loss_curves import get_loss_builder

U8 = numpy.uint8
U16 = numpy.uint16
U32 = numpy.uint32
F32 = numpy.float32
F64 = numpy.float64
U64 = numpy.uint64

def get_assets_ratios(assets, riskmodel, gmvs, imts):
    """
    :param assets: a list of assets on the same site
    :param riskmodel: a CompositeRiskModel instance
    :params gmvs: hazard on the given site, shape (E, M)
    :param imts: intensity measure types
    :returns: a list of (assets, loss_ratios)
    """
    imti = {imt: i for i, imt in enumerate(imts)}
    tdict = riskmodel.get_taxonomy_dict()  # taxonomy -> taxonomy index
    assets_by_t = general.groupby(assets, operator.attrgetter('taxonomy'))
    assets_ratios = []
    for taxo, rm in riskmodel.items():
        t = tdict[taxo]
        try:
            assets = assets_by_t[t]
        except KeyError:  # there are no assets of taxonomy taxo
            continue
        assets_ratios.append((assets, rm.get_loss_ratios(gmvs, imti)))
    return assets_ratios

def start_ebrisk(rupgetter, srcfilter, param, monitor):
    """
    Launcher for ebrisk tasks
    """
    with monitor('weighting ruptures'):
        rupgetter.set_weights(srcfilter, param['num_taxonomies'])
    if rupgetter.weights.sum() <= param['maxweight']:
        yield ebrisk(rupgetter, srcfilter, param, monitor)
    else:
        for rgetter in rupgetter.split(param['maxweight']):
            yield ebrisk, rgetter, srcfilter, param

def ebrisk(rupgetter, srcfilter, param, monitor):
    """
    :param rupgetter:
        a RuptureGetter instance
    :param srcfilter:
        a SourceFilter instance
    :param param:
        a dictionary of parameters
    :param monitor:
        :class:`openquake.baselib.performance.Monitor` instance
    :returns:
        an ArrayWrapper with shape (E, L, T, ...)
    """
    riskmodel = param['riskmodel']
    L = len(riskmodel.lti)
    N = len(srcfilter.sitecol.complete)
    mon = monitor('getting assets', measuremem=False)
    with datastore.read(srcfilter.filename) as dstore:
        assgetter = getters.AssetGetter(dstore)
    getter = getters.GmfGetter(rupgetter, srcfilter, param['oqparam'])
    with monitor('getting hazard'):
        getter.init()  # instantiate the computers
        hazard = getter.get_hazard()  # sid -> (rlzi, sid, eid, gmv)
    mon_risk = monitor('computing losses', measuremem=False)
    mon_agg = monitor('aggregating losses', measuremem=False)
    imts = getter.imts
    events = rupgetter.get_eid_rlz()
    eid2idx = {eid: idx for idx, eid in enumerate(events['eid'])}
    tagnames = param['aggregate_by']
    shape = assgetter.tagcol.agg_shape((len(events), L), tagnames)
    elt_dt = [('eid', U64), ('rlzi', U16), ('loss', (F32, shape[1:]))]
    acc = numpy.zeros(shape, F32)  # shape (E, L, T...)
    if param['avg_losses']:
        losses_by_A = numpy.zeros((assgetter.num_assets, L), F32)
    else:
        losses_by_A = None
    times = numpy.zeros(N)  # risk time per site_id
    for sid, haz in hazard.items():
        t0 = time.time()
        weights = getter.weights[haz['rlzi']]
        assets_on_sid, tagidxs = assgetter.get(sid, tagnames)
        eidx = [eid2idx[eid] for eid in haz['eid']]
        mon.duration += time.time() - t0
        mon.counts += 1
        with mon_risk:
            assets_ratios = get_assets_ratios(
                assets_on_sid, riskmodel, haz['gmv'], imts)
        with mon_agg:
            for assets, triples in assets_ratios:
                for lti, (lt, imt, loss_ratios) in enumerate(triples):
                    w = weights[imt]
                    for asset in assets:
                        aid = asset.ordinal
                        losses = loss_ratios * asset.value(lt)
                        acc[(eidx, lti) + tagidxs[aid]] += losses
                        if param['avg_losses']:
                            losses_by_A[aid, lti] += losses @ w
            times[sid] = time.time() - t0
    with monitor('building event loss table'):
        elt = numpy.fromiter(
            ((event['eid'], event['rlz'], losses)
             for event, losses in zip(events, acc) if losses.sum()), elt_dt)
        agg = general.AccumDict(accum=numpy.zeros(shape[1:], F32))  # rlz->agg
        for rec in elt:
            agg[rec['rlzi']] += rec['loss'] * param['ses_ratio']
    return {'elt': elt, 'agg_losses': agg, 'losses_by_A': losses_by_A,
            'times': times}

@base.calculators.add('ebrisk')
class EbriskCalculator(event_based.EventBasedCalculator):
    """
    Event based PSHA calculator generating event loss tables
    """
    core_task = start_ebrisk
    is_stochastic = True
    precalc = 'event_based'
    accept_precalc = ['event_based', 'event_based_risk', 'ucerf_hazard']

    def pre_execute(self):
        self.oqparam.ground_motion_fields = False
        super().pre_execute()
        # save a copy of the assetcol in hdf5cache
        self.hdf5cache = self.datastore.hdf5cache()
        with hdf5.File(self.hdf5cache, 'w') as cache:
            cache['sitecol'] = self.sitecol.complete
            cache['assetcol'] = self.assetcol
        self.param['ses_ratio'] = self.oqparam.ses_ratio
        self.param['aggregate_by'] = self.oqparam.aggregate_by
        # initialize the riskmodel
        self.riskmodel.taxonomy = self.assetcol.tagcol.taxonomy
        self.param['riskmodel'] = self.riskmodel
        self.L = L = len(self.riskmodel.loss_types)
        A = len(self.assetcol)
        self.datastore.create_dset('avg_losses', F32, (A, L))
        shp = self.get_shape(L)  # shape L, T...
        elt_dt = [('eid', U64), ('rlzi', U16), ('loss', (F32, shp))]
        self.datastore.create_dset('losses_by_event', elt_dt)
        self.zerolosses = numpy.zeros(shp, F32)  # to get the multi-index
        shp = self.get_shape(self.L, self.R)  # shape L, R, T...
        self.datastore.create_dset('agg_losses-rlzs', F32, shp)

    def execute(self):
        oq = self.oqparam
        self.set_param(
            num_taxonomies=self.assetcol.num_taxonomies_by_site(),
            maxweight=oq.ebrisk_maxweight / (oq.concurrent_tasks or 1))
        parent = self.datastore.parent
        if parent:
            hdf5path = parent.filename
            grp_indices = parent['ruptures'].attrs['grp_indices']
            nruptures = len(parent['ruptures'])
        else:
            hdf5path = self.datastore.hdf5cache()
            grp_indices = self.datastore['ruptures'].attrs['grp_indices']
            nruptures = len(self.datastore['ruptures'])
            with hdf5.File(hdf5path, 'r+') as cache:
                self.datastore.hdf5.copy('weights', cache)
                self.datastore.hdf5.copy('ruptures', cache)
                self.datastore.hdf5.copy('rupgeoms', cache)
        self.init_logic_tree(self.csm_info)
        smap = parallel.Starmap(
            self.core_task.__func__, monitor=self.monitor())
        trt_by_grp = self.csm_info.grp_by("trt")
        samples = self.csm_info.get_samples_by_grp()
        rlzs_by_gsim_grp = self.csm_info.get_rlzs_by_gsim_grp()
        ruptures_per_block = numpy.ceil(nruptures / (oq.concurrent_tasks or 1))
        for grp_id, rlzs_by_gsim in rlzs_by_gsim_grp.items():
            start, stop = grp_indices[grp_id]
            for indices in general.block_splitter(
                    range(start, stop), ruptures_per_block):
                rgetter = getters.RuptureGetter(
                    hdf5path, list(indices), grp_id,
                    trt_by_grp[grp_id], samples[grp_id], rlzs_by_gsim)
                smap.submit(rgetter, self.src_filter, self.param)
        return smap.reduce(self.agg_dicts, numpy.zeros(self.N))

    def agg_dicts(self, acc, dic):
        """
        :param dummy: unused parameter
        :param dic: a dictionary with keys eids, losses, losses_by_N
        """
        self.oqparam.ground_motion_fields = False  # hack
        if len(dic['elt']):
            with self.monitor('saving losses_by_event', autoflush=True):
                self.datastore.extend('losses_by_event', dic['elt'])
        with self.monitor('saving agg_losses-rlzs', autoflush=True):
            for r, aggloss in dic['agg_losses'].items():
                self.datastore['agg_losses-rlzs'][:, r] += aggloss
        if self.oqparam.avg_losses:
            with self.monitor('saving avg_losses', autoflush=True):
                self.datastore['avg_losses'] += dic['losses_by_A']
        return acc + dic['times']

    def get_shape(self, *sizes):
        """
        :returns: a shape (S1, ... SN, T1 ... TN)
        """
        return self.assetcol.tagcol.agg_shape(sizes, self.oqparam.aggregate_by)

    def build_datasets(self, builder):
        oq = self.oqparam
        stats = oq.hazard_stats()
        S = len(stats)
        P = len(builder.return_periods)
        C = len(oq.conditional_loss_poes)
        loss_types = ' '.join(self.riskmodel.loss_types)
        if oq.individual_curves or self.R == 1:
            shp = self.get_shape(P, self.R, self.L)  # shape P, R, L, T...
            self.datastore.create_dset('agg_curves-rlzs', F32, shp)
            self.datastore.set_attrs(
                'agg_curves-rlzs', return_periods=builder.return_periods,
                loss_types=loss_types)
        if oq.conditional_loss_poes:
            shp = self.get_shape(C, self.R, self.L)  # shape C, R, L, T...
            self.datastore.create_dset('agg_maps-rlzs', F32, shp)
        if self.R > 1:
            shp = self.get_shape(P, S, self.L)  # shape P, S, L, T...
            self.datastore.create_dset('agg_curves-stats', F32, shp)
            self.datastore.set_attrs(
                'agg_curves-stats', return_periods=builder.return_periods,
                stats=[encode(name) for (name, func) in stats],
                loss_types=loss_types)
            if oq.conditional_loss_poes:
                shp = self.get_shape(C, S, self.L)  # shape C, S, L, T...
                self.datastore.create_dset('agg_maps-stats', F32, shp)
                self.datastore.set_attrs(
                    'agg_maps-stats',
                    stats=[encode(name) for (name, func) in stats],
                    loss_types=loss_types)

    def post_execute(self, times):
        """
        Compute and store average losses from the losses_by_event dataset,
        and then loss curves and maps.
        """
        self.datastore.set_attrs('task_info/start_ebrisk', times=times)
        oq = self.oqparam
        elt_length = len(self.datastore['losses_by_event'])
        builder = get_loss_builder(self.datastore)
        self.build_datasets(builder)
        mon = performance.Monitor(hdf5=hdf5.File(self.datastore.hdf5cache()))
        smap = parallel.Starmap(compute_loss_curves_maps, monitor=mon)
        self.datastore.close()
        acc = []
        ct = oq.concurrent_tasks or 1
        for elt_slice in general.split_in_slices(elt_length, ct):
            smap.submit(self.datastore.filename, elt_slice,
                        oq.conditional_loss_poes, oq.individual_curves)
        for res in smap:
            idx = res.pop('idx')
            for name, arr in res.items():
                if arr is not None:
                    acc.append((name, idx, arr))
        # copy performance information from the cache to the datastore
        pd = mon.hdf5['performance_data'].value
        hdf5.extend3(self.datastore.filename, 'performance_data', pd)
        self.datastore.open('r+')  # reopen
        self.datastore['task_info/compute_loss_curves_and_maps'] = (
            mon.hdf5['task_info/compute_loss_curves_maps'].value)
        with self.monitor('saving loss_curves and maps', autoflush=True):
            for name, idx, arr in acc:
                for ij, val in numpy.ndenumerate(arr):
                    self.datastore[name][ij + idx] = val

def compute_loss_curves_maps(filename, elt_slice, clp, individual_curves,
                             monitor):
    """
    :param filename: path to the datastore
    :param elt_slice: slice of the event loss table
    :param clp: conditional loss poes used to computed the maps
    :param individual_curves: if True, build the individual curves and maps
    :param monitor: a Monitor instance
    :yields:
        dictionaries with keys idx, agg_curves-rlzs, agg_curves-stats,
        agg_maps-rlzs, agg_maps-stats
    """
    with datastore.read(filename) as dstore:
        oq = dstore['oqparam']
        stats = oq.hazard_stats()
        builder = get_loss_builder(dstore)
        R = len(dstore['weights'])
        losses = [[] for _ in range(R)]
        elt = dstore['losses_by_event'][elt_slice]
        for rec in elt:
            losses[rec['rlzi']].append(rec['loss'])
    for multi_index, _ in numpy.ndenumerate(elt[0]['loss']):
        result = {'idx': multi_index}
        thelosses = [[ls[multi_index] for ls in loss] for loss in losses]
        result['agg_curves-rlzs'], result['agg_curves-stats'] = (
            builder.build_pair(thelosses, stats))
        if R > 1 and individual_curves is False:
            del result['agg_curves-rlzs']
        if clp:
            result['agg_maps-rlzs'], result['agg_maps-stats'] = (
                builder.build_loss_maps(thelosses, clp, stats))
            if R > 1 and individual_curves is False:
                del result['agg_maps-rlzs']
        yield result
<|endoftext|>


<commit_before>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2017, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
from __future__ import print_function
import sys
import csv
from openquake.baselib import sap
from openquake.hazardlib import valid

@sap.Script
def normalize(file_csv, sites_csv):
    """
    Produce a good `sites_csv` file from a generic CSV with fields longitude
    and latitude, by truncating to 5 digits and discarding duplicate sites.
    """
    nrows = 0
    with open(file_csv) as f:
        reader = csv.reader(f)
        header = next(reader)
        for i, field in enumerate(header):
            col = field.lower()
            if col in ('lon', 'longitude'):
                lon_idx = i
            if col in ('lat', 'latitude'):
                lat_idx = i
        try:
            lon_idx
        except NameError:
            sys.exit('No longitude field in the header of %s!' % file_csv)
        try:
            lat_idx
        except NameError:
            sys.exit('No latitude field in the header of %s!' % file_csv)
        points = set()
        for row in reader:
            point = valid.lon_lat('%s %s' % (row[lon_idx], row[lat_idx]))
            points.add(point)
            nrows += 1
    with open(sites_csv, 'w') as f:
        for point in sorted(points):
            f.write('%s,%s\n' % point)
    print('Written %d sites on %s' % (len(points), sites_csv))
    if len(points) < nrows:
        print('Found %d duplicated sites!' % (nrows - len(points)))

normalize.arg('file_csv', 'CSV file with lon and lat in the header')
normalize.arg('sites_csv', 'name of the normalized sites.csv file')
<commit_msg>Renamed the command oq normalize -> oq to_sites<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# =============================================================================
# @file ostap/frames/tests/test_frames_frames.py
# Test module for ostap/frames/frames.py.
# Copyright (c) Ostap developpers.
# =============================================================================
""" Test module for ostap/frames/frames.py.
"""
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' ==  __name__ : logger = getLogger ( 'test_frames_frames' )
else                       : logger = getLogger ( __name__            )
# =============================================================================
import ROOT, os
from ostap.core.pyrouts    import hID
from ostap.frames.frames   import DataFrame, frame_project, frame_statVar, frame_statCov
from ostap.utils.cleanup   import CleanUp
from ostap.trees.trees     import Tree
from ostap.plotting.canvas import use_canvas
from ostap.utils.utils     import wait
from ostap.core.meta_info  import root_version

tmpdir = CleanUp().tmpdir
fname  = os.path.join ( tmpdir , 'test_frame.root' )

# A simple helper function to fill a test tree
def fill_tree ( tname , fname ):

    tdf = DataFrame        ( 1000 )
    a   = tdf.ProgressBar  ( 1000 )
    tdf.Define   ("one", "1.0"                    )\
       .Define   ("b1" , "(double) tdfentry_ + 1 ")\
       .Define   ("b2" , "(1.0+b1)*(1.0+b1)"      )\
       .Snapshot ( tname, fname )

# We prepare an input tree to run on
tname = "myTree"
fill_tree ( tname, fname )

frame = DataFrame ( tname        , fname        )
tree  = Tree      ( name = tname , file = fname ).chain

from ostap.utils.utils  import implicitMT
from ostap.utils.timing import timing

def test_frame0 () :

    logger = getLogger ( 'test_frame0' )
    if root_version < (6,16) :
        logger.warning ( "Test is disabled for this version of ROOT %s" % str ( root_version ) )
        return

    for mt in  ( False , True ) :

        with implicitMT ( mt ) :

            logger.info ( 80*'*' )
            logger.info ( 'MT enabled?  %s' % ROOT.ROOT.IsImplicitMTEnabled() )

            logger.info ( 'Tree  :\n%s' % tree  )
            logger.info ( 'Frame :\n%s' % frame )

            logger.info ( 'Len:                          %30s vs %-30s '  % ( len ( frame )         , len ( tree ) )        )
            logger.info ( 'Branches:                     %30s vs %-30s '  % ( frame.branches()      , tree.branches() )     )
            logger.info ( 'nEff:                         %30s vs %-30s '  % ( frame.nEff    ()      , tree.nEff    () )     )
            logger.info ( 'nEff(b1):                     %30s vs %-30s '  % ( frame.nEff    ('b1')  , tree.nEff    ('b1') ) )
            logger.info ( "m(5,50,'b1','b1/(b2+1'):      %30s vs %-30s "  % ( frame.get_moment ( 5 , 50 ,  'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .get_moment ( 5 , 50 ,  'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "m(5,'b1','b1/(b2+1'):         %30s vs %-30s "  % ( frame.moment ( 5 , 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .moment ( 5 , 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "cm(5,'b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.central_moment ( 2 , 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .central_moment ( 2 , 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "mean('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.mean ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .mean ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "rms ('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.rms  ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .rms  ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "skew('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.skewness ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .skewness ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "kurt('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.kurtosis ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .kurtosis ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "quantile(0.3, 'b1','b1<500'): %30s vs %-30s "  % ( frame.quantile ( 0.3 , 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .quantile ( 0.3 , 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "median(0.3, 'b1','b1<500'):   %30s vs %-30s "  % ( frame.median   (       'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .median   (       'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "terciles('b1','b1<500'):      %30s vs %-30s "  % ( frame.terciles (       'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .terciles (       'b1' , 'b1/(b2+1)' ) ) )

def test_frame1 ( ) :

    logger = getLogger ( 'test_frame1' )
    if root_version < (6,16) :
        logger.warning ( "Test is disabled for this version of ROOT %s" % str ( root_version ) )
        return

    h1 = tree .draw ( 'b1' , '1/b1' )
    h2 = frame.draw ( 'b1' , '1/b1' )

    h1 = ROOT.TH1D( hID() , '' , 100 , 0 , 1000 )
    h2 = ROOT.TH1D( hID() , '' , 100 , 0 , 1000 )

    tree.project  ( h1    , 'b1' , '1.0/b1' )
    frame_project ( tree  , h2 , 'b1' , '2.0/b1' )

    with wait ( 3 ), use_canvas ( 'test_frame1' ) :
        h1.red  ()
        h2.blue ()
        h1.draw ()
        h2.draw ( 'same hist' )

def test_frame2 ( ) :

    logger = getLogger ( 'test_frame2' )
    if root_version < (6,16) :
        logger.warning ( "Test is disabled for this version of ROOT %s" % str ( root_version ) )
        return

    s1 = tree.statVar  (        'b1' , '1/b1' )
    s2 = frame_statVar ( tree , 'b1' , '1/b1' )

    logger.info ('StatTree  :%s' % s1 )
    logger.info ('StatFrame :%s' % s2 )

# =============================================================================
if '__main__' == __name__ :

    test_frame0 ()
    test_frame1 ()
    test_frame2 ()

    pass

# =============================================================================
##                                                                      The END
# =============================================================================
<commit_msg>Update test_frames_frames.py<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# =============================================================================
# @file ostap/frames/tests/test_frames_frames.py
# Test module for ostap/frames/frames.py.
# Copyright (c) Ostap developpers.
# =============================================================================
""" Test module for ostap/frames/frames.py.
"""
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' ==  __name__ : logger = getLogger ( 'test_frames_frames' )
else                       : logger = getLogger ( __name__            )
# =============================================================================
import ROOT, os
from ostap.core.pyrouts    import hID

from ostap.utils.cleanup   import CleanUp
from ostap.trees.trees     import Tree
from ostap.plotting.canvas import use_canvas
from ostap.utils.utils     import wait
from ostap.core.meta_info  import root_version

tmpdir = CleanUp().tmpdir
fname  = os.path.join ( tmpdir , 'test_frame.root' )

# A simple helper function to fill a test tree
def fill_tree ( tname , fname ):

    tdf = DataFrame        ( 1000 )
    a   = tdf.ProgressBar  ( 1000 )
    tdf.Define   ("one", "1.0"                    )\
       .Define   ("b1" , "(double) tdfentry_ + 1 ")\
       .Define   ("b2" , "(1.0+b1)*(1.0+b1)"      )\
       .Snapshot ( tname, fname )

# We prepare an input tree to run on
tname = "myTree"
fill_tree ( tname, fname )

if (6,16) <= root_version :
    from ostap.frames.frames   import DataFrame, frame_project, frame_statVar, frame_statCov
    frame = DataFrame ( tname        , fname        )
    tree  = Tree      ( name = tname , file = fname ).chain

from ostap.utils.utils  import implicitMT
from ostap.utils.timing import timing

def test_frame0 () :

    logger = getLogger ( 'test_frame0' )
    if root_version < (6,16) :
        logger.warning ( "Test is disabled for this version of ROOT %s" % str ( root_version ) )
        return

    for mt in  ( False , True ) :

        with implicitMT ( mt ) :

            logger.info ( 80*'*' )
            logger.info ( 'MT enabled?  %s' % ROOT.ROOT.IsImplicitMTEnabled() )

            logger.info ( 'Tree  :\n%s' % tree  )
            logger.info ( 'Frame :\n%s' % frame )

            logger.info ( 'Len:                          %30s vs %-30s '  % ( len ( frame )         , len ( tree ) )        )
            logger.info ( 'Branches:                     %30s vs %-30s '  % ( frame.branches()      , tree.branches() )     )
            logger.info ( 'nEff:                         %30s vs %-30s '  % ( frame.nEff    ()      , tree.nEff    () )     )
            logger.info ( 'nEff(b1):                     %30s vs %-30s '  % ( frame.nEff    ('b1')  , tree.nEff    ('b1') ) )
            logger.info ( "m(5,50,'b1','b1/(b2+1'):      %30s vs %-30s "  % ( frame.get_moment ( 5 , 50 ,  'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .get_moment ( 5 , 50 ,  'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "m(5,'b1','b1/(b2+1'):         %30s vs %-30s "  % ( frame.moment ( 5 , 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .moment ( 5 , 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "cm(5,'b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.central_moment ( 2 , 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .central_moment ( 2 , 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "mean('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.mean ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .mean ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "rms ('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.rms  ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .rms  ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "skew('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.skewness ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .skewness ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "kurt('b1','b1/(b2+1'):        %30s vs %-30s "  % ( frame.kurtosis ( 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .kurtosis ( 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "quantile(0.3, 'b1','b1<500'): %30s vs %-30s "  % ( frame.quantile ( 0.3 , 'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .quantile ( 0.3 , 'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "median(0.3, 'b1','b1<500'):   %30s vs %-30s "  % ( frame.median   (       'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .median   (       'b1' , 'b1/(b2+1)' ) ) )
            logger.info ( "terciles('b1','b1<500'):      %30s vs %-30s "  % ( frame.terciles (       'b1' , 'b1/(b2+1)' ) ,
                                                                              tree .terciles (       'b1' , 'b1/(b2+1)' ) ) )

def test_frame1 ( ) :

    logger = getLogger ( 'test_frame1' )
    if root_version < ( 6 , 16 ) :
        logger.warning ( "Test is disabled for this version of ROOT %s" % str ( root_version ) )
        return

    h1 = tree .draw ( 'b1' , '1/b1' )
    h2 = frame.draw ( 'b1' , '1/b1' )

    h1 = ROOT.TH1D( hID() , '' , 100 , 0 , 1000 )
    h2 = ROOT.TH1D( hID() , '' , 100 , 0 , 1000 )

    tree.project  ( h1    , 'b1' , '1.0/b1' )
    frame_project ( tree  , h2 , 'b1' , '2.0/b1' )

    with wait ( 3 ), use_canvas ( 'test_frame1' ) :
        h1.red  ()
        h2.blue ()
        h1.draw ()
        h2.draw ( 'same hist' )

def test_frame2 ( ) :

    logger = getLogger ( 'test_frame2' )
    if root_version <  ( 6 , 16 ) :
        logger.warning ( "Test is disabled for this version of ROOT %s" % str ( root_version ) )
        return

    s1 = tree.statVar  (        'b1' , '1/b1' )
    s2 = frame_statVar ( tree , 'b1' , '1/b1' )

    logger.info ('StatTree  :%s' % s1 )
    logger.info ('StatFrame :%s' % s2 )

# =============================================================================
if '__main__' == __name__ :

    test_frame0 ()
    test_frame1 ()
    test_frame2 ()

    pass

# =============================================================================
##                                                                      The END
# =============================================================================
<|endoftext|>


<commit_before>#!/usr/bin/env python

import rospy
import roslib
import rosparam
import smach
import smach_ros
import std_msgs.msg
from dynamic_reconfigure.client import Client as DynClient
from std_srvs.srv import Empty, EmptyResponse
import strands_webserver
import strands_webserver.client_utils
from mongodb_store.message_store import MessageStoreProxy
from aaf_walking_group.entertain import Entertain
from aaf_walking_group.guide_interface import GuideInterface
from aaf_walking_group.guiding import Guiding
from aaf_walking_group.reached_resting_point import RestingPoint
from aaf_walking_group.resting_chair_selection import SelectRestingChair
from aaf_walking_group.final_approach import FinalApproach
from aaf_walking_group.msg import GuidingAction, EmptyAction, StateMachineAction
from aaf_walking_group.srv import GetMediaId
from aaf_waypoint_sounds.srv import WaypointSoundsService, WaypointSoundsServiceRequest
from aaf_walking_group.utils import PTU, Gaze, RecoveryReconfigure
from aaf_walking_group.waypoint_manager import WaypointManager
import aaf_walking_group.utils as utils
from music_player.srv import MusicPlayerService
from sound_player_server.srv import PlaySoundService
from walking_group_recovery.srv import ToggleWalkingGroupRecovery
import actionlib
import json
import pprint

class WalkingGroupStateMachine(object):
    def __init__(self, name):
        rospy.loginfo("Creating " + name + " server...")
        self._as = actionlib.SimpleActionServer(
            name,
            StateMachineAction,
            self.execute,
            auto_start=False
        )
        self._as.register_preempt_callback(self.preempt_callback)

        rospy.loginfo("Creating guiding client...")
        nav_client = actionlib.SimpleActionClient("guiding", GuidingAction)
        nav_client.wait_for_server()
        rospy.loginfo(" ... done")
        rospy.loginfo("Creating wait client...")
        wait_client = actionlib.SimpleActionClient("wait_for_participant", EmptyAction)
        wait_client.wait_for_server()
        rospy.loginfo(" ... done")
        rospy.loginfo("Waiting for recovery toggle services...")
        rospy.loginfo(" ... images")
        self.rec_srv = rospy.ServiceProxy('/toggle_walking_group_recovery', ToggleWalkingGroupRecovery)
        self.rec_srv.wait_for_service()
        rospy.loginfo(" ... done")
        rospy.loginfo("Waiting for media server services...")
        rospy.loginfo(" ... images")
        s = rospy.ServiceProxy('/aaf_walking_group/image_server/get_id', GetMediaId)
        s.wait_for_service()
        rospy.loginfo(" ... video")
        s = rospy.ServiceProxy('/aaf_walking_group/video_server/get_id', GetMediaId)
        s.wait_for_service()
        rospy.loginfo(" ... sound")
        rospy.loginfo("  ... music")
        s = rospy.ServiceProxy('/music_player_service', MusicPlayerService)
        s.wait_for_service()
        rospy.loginfo("  ... jingles")
        s = rospy.ServiceProxy('/sound_player_server/sound_player_service', PlaySoundService)
        s.wait_for_service()
        rospy.loginfo("  ... waypoints")
        s = rospy.ServiceProxy('/aaf_waypoint_sounds_service', WaypointSoundsService)
        s.wait_for_service()
        rospy.loginfo(" ... done")
        # Get parameters
        self.volume = utils.get_master_volume()
        self.display_no = rospy.get_param("~display_no", 0)
        self.waypointset_name = rospy.get_param("~mongodb_params/waypointset_name", "")
        self.waypointset_collection = rospy.get_param("~mongodb_params/waypointset_collection", "aaf_walking_group")
        self.waypointset_meta = rospy.get_param("~mongodb_params/waypointset_meta", "waypoint_set")
        if self.waypointset_name == "":
            rospy.logfatal("Missing parameters.")
            rospy.logfatal("Please run with _mongodb_params/waypointset_name:=<waypointset_name>")
            return

        self.preempt_srv = None

        self.ptu = PTU()
        self.gaze = Gaze()
        self.recovery = RecoveryReconfigure(
            name="/monitored_navigation/recover_states/",
            whitelist=rosparam.load_file(rospy.get_param("~recovery_whitelist"))[0][0]["recover_states"]
        )
        rospy.loginfo("Creating dynamic reconfigure clients")
        self.han_dyn_client = DynClient(
            "/human_aware_navigation"
        )
        self.mov_dyn_client = DynClient(
            "/human_aware_navigation/DWAPlannerROS"
        )
        rospy.loginfo(" ...done")
        self.get_current_dyn_settings()

        rospy.loginfo(" ... starting " + name)
        self._as.start()
        rospy.loginfo(" ... started " + name)

    def get_current_dyn_settings(self):
        gazing = rospy.get_param("/human_aware_navigation/gaze_type")
        angle = round(rospy.get_param("/human_aware_navigation/detection_angle"),2)
        self.han_param = {
            'gaze_type': gazing,
            'detection_angle': angle
        }
        rospy.loginfo("Found following default values for human_aware_navigation: %s", self.han_param)
        max_vel_x = round(rospy.get_param("/human_aware_navigation/DWAPlannerROS/max_vel_x"), 2)
        max_trans_vel = round(rospy.get_param("/human_aware_navigation/DWAPlannerROS/max_trans_vel"),2)
        self.mov_param = {
            "max_vel_x": max_vel_x,
            "max_trans_vel": max_trans_vel
        }
        rospy.loginfo("Found following default values for human_aware_navigation/DWAPlannerROS: %s", self.mov_param)

    def execute(self, goal):
        rospy.loginfo("Starting state machine")

        try:
            rospy.loginfo("Creating recovery toggle service proxy and waiting ...")
            self.rec_srv.wait_for_service()
            rospy.loginfo(" .. calling recovery toggle service")
            self.rec_srv(True)
            rospy.loginfo(" .. started walking group recovery")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("waypoint sound service call failed: %s" % e)
        #Save current recovery behaviour states
        self.recovery.save_current_configuration()
        #and reconfigure
        self.recovery.reconfigure(RecoveryReconfigure.SET)

        # Setting http root
        http_root = roslib.packages.get_pkg_dir('aaf_walking_group') + '/www'
        strands_webserver.client_utils.set_http_root(http_root)

        self.ptu.turnPTU(-180, 15)
        dyn_param = {
            'timeout': 0.0,
            'gaze_type': 1,
            'detection_angle': 20.0
        }
        try:
            self.han_dyn_client.update_configuration(dyn_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)

        self.preempt_srv = rospy.Service('/walking_group/cancel', Empty, self.preempt_srv_cb)
        config = self.loadConfig(self.waypointset_name, collection_name=self.waypointset_collection, meta_name=self.waypointset_meta)
        pprint.pprint(config)

        dyn_param = {
            'max_vel_x': config[goal.group]["speed"],
            'max_trans_vel': config[goal.group]["speed"]
        }
        try:
            self.mov_dyn_client.update_configuration(dyn_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)

        try:
            rospy.loginfo("Creating waypoint sound service proxy and waiting ...")
            s = rospy.ServiceProxy('aaf_waypoint_sounds_service', WaypointSoundsService)
            s.wait_for_service()
            rospy.loginfo(" .. calling waypoint sound service")
            s(WaypointSoundsServiceRequest.RESUME)
            rospy.loginfo(" .. started waypoint sound service")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("waypoint sound service call failed: %s" % e)

        # Create a SMACH state machine
        self.sm = smach.StateMachine(outcomes=['succeeded', 'aborted', 'preempted'])
        self.sm.userdata.waypoints = WaypointManager([i for k,i in sorted(config[goal.group]["waypoints"].items(),key=lambda to_int: int(to_int[0]))])
        self.sm.userdata.play_music = True
        sis = smach_ros.IntrospectionServer(
            'walking_group_state_machine',
            self.sm,
            '/walking_group_machine'
        )
        sis.start()
        # Open the container
        with self.sm:
            # Add states to the container
            smach.StateMachine.add(
                'ENTERTAIN',
                Entertain(self.display_no, self.gaze),
                transitions={
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'GUIDE_INTERFACE',
                GuideInterface(),
                transitions={
                    'move_to_point': 'GUIDING',
                    'aborted': 'ENTERTAIN',
                    'select_chair': 'RESTING_CHAIR',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'GUIDING',
                Guiding(),
                transitions={
                    'reached_point': 'RESTING_CONT',
                    'reached_final_point': 'succeeded',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'RESTING_CONT',
                RestingPoint(self.display_no),
                transitions={
                    'rest': 'RESTING_CHAIR',
                    'continue': 'GUIDING',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'RESTING_CHAIR',
                SelectRestingChair(self.display_no),
                transitions={
                    'continue': 'FINAL_APPROACH',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'FINAL_APPROACH',
                FinalApproach(),
                transitions={
                    'reached_point': 'ENTERTAIN',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )

        # Execute SMACH plan
        self.sm.execute()

        sis.stop()
        self.preempt_srv.shutdown()
        self.ptu.turnPTU(0, 0)
        self.recovery.reconfigure(RecoveryReconfigure.RESET)
        utils.set_master_volume(self.volume)
        try:
            self.han_dyn_client.update_configuration(self.han_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)
        try:
            self.mov_dyn_client.update_configuration(self.mov_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)
        try:
            rospy.loginfo("Creating waypoint sound service proxy and waiting ...")
            s = rospy.ServiceProxy('aaf_waypoint_sounds_service', WaypointSoundsService)
            s.wait_for_service()
            rospy.loginfo(" ... calling waypoint sound service")
            s(WaypointSoundsServiceRequest.PAUSE)
            rospy.loginfo(" ... stopped waypoint sound service")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("Service call failed: %s" % e)
        try:
            rospy.loginfo("Creating recovery toggle service proxy and waiting ...")
            self.rec_srv.wait_for_service()
            rospy.loginfo(" .. calling recovery toggle service")
            self.rec_srv(False)
            rospy.loginfo(" .. stopped walking group recovery")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("Service call failed: %s" % e)
        if not self._as.is_preempt_requested() and self._as.is_active():
            self._as.set_succeeded()
        else:
            self._as.set_preempted()

        strands_webserver.client_utils.display_relative_page(self.display_no, "ende.html")

    def preempt_callback(self):
        rospy.logwarn("Walking group preempt requested")
        self.sm.request_preempt()

    def preempt_srv_cb(self, req):
        self.preempt_callback()
        return EmptyResponse()

    def loadConfig(self, dataset_name, collection_name="aaf_walking_group", meta_name="waypoint_set"):
        msg_store = MessageStoreProxy(collection=collection_name)
        query_meta = {}
        query_meta[meta_name] = dataset_name
        if len(msg_store.query(std_msgs.msg.String._type, {}, query_meta)) == 0:
            rospy.logerr("Desired data set '"+meta_name+": "+dataset_name+"' not in datacentre.")
            raise Exception("Can't find data in datacentre.")
        else:
            message = msg_store.query(std_msgs.msg.String._type, {}, query_meta)
            return json.loads(message[0][0].data)

if __name__ == '__main__':
    rospy.init_node('walking_group_smach')
    w = WalkingGroupStateMachine(rospy.get_name())
    rospy.spin()
<commit_msg>Publishing `none` on localise by topic topic after group ends (#371)<commit_after>#!/usr/bin/env python

import rospy
import roslib
import rosparam
import smach
import smach_ros
import std_msgs.msg
from dynamic_reconfigure.client import Client as DynClient
from std_srvs.srv import Empty, EmptyResponse
import strands_webserver
import strands_webserver.client_utils
from mongodb_store.message_store import MessageStoreProxy
from aaf_walking_group.entertain import Entertain
from aaf_walking_group.guide_interface import GuideInterface
from aaf_walking_group.guiding import Guiding
from aaf_walking_group.reached_resting_point import RestingPoint
from aaf_walking_group.resting_chair_selection import SelectRestingChair
from aaf_walking_group.final_approach import FinalApproach
from aaf_walking_group.msg import GuidingAction, EmptyAction, StateMachineAction
from aaf_walking_group.srv import GetMediaId
from aaf_waypoint_sounds.srv import WaypointSoundsService, WaypointSoundsServiceRequest
from aaf_walking_group.utils import PTU, Gaze, RecoveryReconfigure
from aaf_walking_group.waypoint_manager import WaypointManager
import aaf_walking_group.utils as utils
from music_player.srv import MusicPlayerService
from sound_player_server.srv import PlaySoundService
from walking_group_recovery.srv import ToggleWalkingGroupRecovery
import actionlib
import json
import pprint

class WalkingGroupStateMachine(object):
    def __init__(self, name):
        rospy.loginfo("Creating " + name + " server...")
        self._as = actionlib.SimpleActionServer(
            name,
            StateMachineAction,
            self.execute,
            auto_start=False
        )
        self._as.register_preempt_callback(self.preempt_callback)

        rospy.loginfo("Creating guiding client...")
        nav_client = actionlib.SimpleActionClient("guiding", GuidingAction)
        nav_client.wait_for_server()
        rospy.loginfo(" ... done")
        rospy.loginfo("Creating wait client...")
        wait_client = actionlib.SimpleActionClient("wait_for_participant", EmptyAction)
        wait_client.wait_for_server()
        rospy.loginfo(" ... done")
        rospy.loginfo("Waiting for recovery toggle services...")
        rospy.loginfo(" ... images")
        self.rec_srv = rospy.ServiceProxy('/toggle_walking_group_recovery', ToggleWalkingGroupRecovery)
        self.rec_srv.wait_for_service()
        rospy.loginfo(" ... done")
        rospy.loginfo("Waiting for media server services...")
        rospy.loginfo(" ... images")
        s = rospy.ServiceProxy('/aaf_walking_group/image_server/get_id', GetMediaId)
        s.wait_for_service()
        rospy.loginfo(" ... video")
        s = rospy.ServiceProxy('/aaf_walking_group/video_server/get_id', GetMediaId)
        s.wait_for_service()
        rospy.loginfo(" ... sound")
        rospy.loginfo("  ... music")
        s = rospy.ServiceProxy('/music_player_service', MusicPlayerService)
        s.wait_for_service()
        rospy.loginfo("  ... jingles")
        s = rospy.ServiceProxy('/sound_player_server/sound_player_service', PlaySoundService)
        s.wait_for_service()
        rospy.loginfo("  ... waypoints")
        s = rospy.ServiceProxy('/aaf_waypoint_sounds_service', WaypointSoundsService)
        s.wait_for_service()
        rospy.loginfo(" ... done")
        # Get parameters
        self.volume = utils.get_master_volume()
        self.display_no = rospy.get_param("~display_no", 0)
        self.waypointset_name = rospy.get_param("~mongodb_params/waypointset_name", "")
        self.waypointset_collection = rospy.get_param("~mongodb_params/waypointset_collection", "aaf_walking_group")
        self.waypointset_meta = rospy.get_param("~mongodb_params/waypointset_meta", "waypoint_set")
        if self.waypointset_name == "":
            rospy.logfatal("Missing parameters.")
            rospy.logfatal("Please run with _mongodb_params/waypointset_name:=<waypointset_name>")
            return

        self.preempt_srv = None

        self.ptu = PTU()
        self.gaze = Gaze()
        self.recovery = RecoveryReconfigure(
            name="/monitored_navigation/recover_states/",
            whitelist=rosparam.load_file(rospy.get_param("~recovery_whitelist"))[0][0]["recover_states"]
        )
        rospy.loginfo("Creating dynamic reconfigure clients")
        self.han_dyn_client = DynClient(
            "/human_aware_navigation"
        )
        self.mov_dyn_client = DynClient(
            "/human_aware_navigation/DWAPlannerROS"
        )
        rospy.loginfo(" ...done")
        self.get_current_dyn_settings()

        rospy.loginfo(" ... starting " + name)
        self._as.start()
        rospy.loginfo(" ... started " + name)

    def get_current_dyn_settings(self):
        gazing = rospy.get_param("/human_aware_navigation/gaze_type")
        angle = round(rospy.get_param("/human_aware_navigation/detection_angle"),2)
        self.han_param = {
            'gaze_type': gazing,
            'detection_angle': angle
        }
        rospy.loginfo("Found following default values for human_aware_navigation: %s", self.han_param)
        max_vel_x = round(rospy.get_param("/human_aware_navigation/DWAPlannerROS/max_vel_x"), 2)
        max_trans_vel = round(rospy.get_param("/human_aware_navigation/DWAPlannerROS/max_trans_vel"),2)
        self.mov_param = {
            "max_vel_x": max_vel_x,
            "max_trans_vel": max_trans_vel
        }
        rospy.loginfo("Found following default values for human_aware_navigation/DWAPlannerROS: %s", self.mov_param)

    def execute(self, goal):
        rospy.loginfo("Starting state machine")

        try:
            rospy.loginfo("Creating recovery toggle service proxy and waiting ...")
            self.rec_srv.wait_for_service()
            rospy.loginfo(" .. calling recovery toggle service")
            self.rec_srv(True)
            rospy.loginfo(" .. started walking group recovery")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("waypoint sound service call failed: %s" % e)
        #Save current recovery behaviour states
        self.recovery.save_current_configuration()
        #and reconfigure
        self.recovery.reconfigure(RecoveryReconfigure.SET)

        # Setting http root
        http_root = roslib.packages.get_pkg_dir('aaf_walking_group') + '/www'
        strands_webserver.client_utils.set_http_root(http_root)

        self.ptu.turnPTU(-180, 15)
        dyn_param = {
            'timeout': 0.0,
            'gaze_type': 1,
            'detection_angle': 20.0
        }
        try:
            self.han_dyn_client.update_configuration(dyn_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)

        self.preempt_srv = rospy.Service('/walking_group/cancel', Empty, self.preempt_srv_cb)
        config = self.loadConfig(self.waypointset_name, collection_name=self.waypointset_collection, meta_name=self.waypointset_meta)
        pprint.pprint(config)

        dyn_param = {
            'max_vel_x': config[goal.group]["speed"],
            'max_trans_vel': config[goal.group]["speed"]
        }
        try:
            self.mov_dyn_client.update_configuration(dyn_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)

        try:
            rospy.loginfo("Creating waypoint sound service proxy and waiting ...")
            s = rospy.ServiceProxy('aaf_waypoint_sounds_service', WaypointSoundsService)
            s.wait_for_service()
            rospy.loginfo(" .. calling waypoint sound service")
            s(WaypointSoundsServiceRequest.RESUME)
            rospy.loginfo(" .. started waypoint sound service")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("waypoint sound service call failed: %s" % e)

        # Create a SMACH state machine
        self.sm = smach.StateMachine(outcomes=['succeeded', 'aborted', 'preempted'])
        self.sm.userdata.waypoints = WaypointManager([i for k,i in sorted(config[goal.group]["waypoints"].items(),key=lambda to_int: int(to_int[0]))])
        self.sm.userdata.play_music = True
        sis = smach_ros.IntrospectionServer(
            'walking_group_state_machine',
            self.sm,
            '/walking_group_machine'
        )
        sis.start()
        # Open the container
        with self.sm:
            # Add states to the container
            smach.StateMachine.add(
                'ENTERTAIN',
                Entertain(self.display_no, self.gaze),
                transitions={
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'GUIDE_INTERFACE',
                GuideInterface(),
                transitions={
                    'move_to_point': 'GUIDING',
                    'aborted': 'ENTERTAIN',
                    'select_chair': 'RESTING_CHAIR',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'GUIDING',
                Guiding(),
                transitions={
                    'reached_point': 'RESTING_CONT',
                    'reached_final_point': 'succeeded',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'RESTING_CONT',
                RestingPoint(self.display_no),
                transitions={
                    'rest': 'RESTING_CHAIR',
                    'continue': 'GUIDING',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'RESTING_CHAIR',
                SelectRestingChair(self.display_no),
                transitions={
                    'continue': 'FINAL_APPROACH',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )
            smach.StateMachine.add(
                'FINAL_APPROACH',
                FinalApproach(),
                transitions={
                    'reached_point': 'ENTERTAIN',
                    'key_card': 'GUIDE_INTERFACE',
                    'killall': 'preempted'
                },
                remapping={'waypoints' : 'waypoints', 'play_music' : 'play_music'}
            )

        # Execute SMACH plan
        self.sm.execute()

        sis.stop()
        self.preempt_srv.shutdown()
        self.ptu.turnPTU(0, 0)
        self.recovery.reconfigure(RecoveryReconfigure.RESET)
        pub = rospy.Publisher("/aaf_walking_group/resting_node", std_msgs.msg.String, queue_size=1, latch=True)
        pub.publish("none") # Prevent localise by topic after walking group is done.
        utils.set_master_volume(self.volume)
        try:
            self.han_dyn_client.update_configuration(self.han_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)
        try:
            self.mov_dyn_client.update_configuration(self.mov_param)
        except rospy.ServiceException as e:
            rospy.logerr("Caught service exception: %s", e)
        try:
            rospy.loginfo("Creating waypoint sound service proxy and waiting ...")
            s = rospy.ServiceProxy('aaf_waypoint_sounds_service', WaypointSoundsService)
            s.wait_for_service()
            rospy.loginfo(" ... calling waypoint sound service")
            s(WaypointSoundsServiceRequest.PAUSE)
            rospy.loginfo(" ... stopped waypoint sound service")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("Service call failed: %s" % e)
        try:
            rospy.loginfo("Creating recovery toggle service proxy and waiting ...")
            self.rec_srv.wait_for_service()
            rospy.loginfo(" .. calling recovery toggle service")
            self.rec_srv(False)
            rospy.loginfo(" .. stopped walking group recovery")
        except (rospy.ServiceException, rospy.ROSInterruptException) as e:
            rospy.logwarn("Service call failed: %s" % e)
        if not self._as.is_preempt_requested() and self._as.is_active():
            self._as.set_succeeded()
        else:
            self._as.set_preempted()

        strands_webserver.client_utils.display_relative_page(self.display_no, "ende.html")

    def preempt_callback(self):
        rospy.logwarn("Walking group preempt requested")
        self.sm.request_preempt()

    def preempt_srv_cb(self, req):
        self.preempt_callback()
        return EmptyResponse()

    def loadConfig(self, dataset_name, collection_name="aaf_walking_group", meta_name="waypoint_set"):
        msg_store = MessageStoreProxy(collection=collection_name)
        query_meta = {}
        query_meta[meta_name] = dataset_name
        if len(msg_store.query(std_msgs.msg.String._type, {}, query_meta)) == 0:
            rospy.logerr("Desired data set '"+meta_name+": "+dataset_name+"' not in datacentre.")
            raise Exception("Can't find data in datacentre.")
        else:
            message = msg_store.query(std_msgs.msg.String._type, {}, query_meta)
            return json.loads(message[0][0].data)

if __name__ == '__main__':
    rospy.init_node('walking_group_smach')
    w = WalkingGroupStateMachine(rospy.get_name())
    rospy.spin()
<|endoftext|>