text:
#!/usr/bin/env python
# Copyright (c) 2019 - The Procedural Generation for Gazebo authors
# For information on the respective copyright owner see the NOTICE file
#
# 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 os
import unittest
import numpy as np
from random import random, choice, randint
import shutil
import getpass
from pcg_gazebo.utils import generate_random_string
from pcg_gazebo.simulation import load_gazebo_models, get_gazebo_model_sdf, \
    get_gazebo_model_names, get_gazebo_model_path, \
    get_gazebo_model_sdf_filenames
from pcg_gazebo.simulation import Box, Cylinder, Sphere, Joint, \
    SimulationModel
from pcg_gazebo.simulation.properties import Pose
from pcg_gazebo.parsers import parse_sdf, parse_sdf_config
from pcg_gazebo.parsers.urdf import create_urdf_element

CUR_DIR = os.path.dirname(os.path.abspath(__file__))

INVALID_SIZES = [
    dict(),
    [2, 3],
    [1 for _ in range(8)],
    list(),
    'a',
    ['a' for _ in range(3)]
]

class TestSimulationObjects(unittest.TestCase):
    def check_joint_names(self, model, joint_names):
        for joint in joint_names:
            self.assertIsNotNone(
                model.get_joint_by_name(joint),
                'Cannot find joint of name <{}>'.format(joint))

    def check_link_names(self, model, link_names):
        for link in link_names:
            self.assertIsNotNone(model.get_link_by_name(link),
                                 'Cannot find link of name <{}>'.format(link))

    def check_number_of_links(self, model, n_links):
        self.assertEqual(len(model.links), n_links,
                         'Model should have {} links'.format(n_links))

    def check_number_of_joints(self, model, n_joints):
        self.assertEqual(len(model.joints), n_joints,
                         'Model should have {} joints'.format(n_joints))

    def check_attributes(self, obj, params):
        for tag in params:
            att = getattr(obj, tag)
            if hasattr(att, 'value'):
                value = att.value
            else:
                value = att
            self.assertEqual(
                value,
                params[tag],
                'Element {} should have value {}, retrieved={}'.format(
                    tag,
                    params[tag],
                    value))

    def check_pose(self, obj, pos, rpy):
        self.assertAlmostEqual(
            np.linalg.norm(obj.pose.position - pos),
            0,
            7,
            'Position was not parsed correctly, retrieved={}, '
            'expected={}'.format(
                obj.pose.position.tolist(),
                pos))
        diff = Pose.get_transform(obj.pose.quat, Pose.rpy2quat(*rpy))

        self.assertAlmostEqual(
            np.sum(diff[0:3]), 0, 7,
            'Orientation should be {}, retrieved={}'.format(
                Pose.rpy2quat(*rpy), obj.pose.quat))

    def create_random_box(self):
        box = Box()
        size = [random() for _ in range(3)]
        name = generate_random_string(5)
        # Setting random size
        box.size = size
        box.name = name
        box.add_inertial(random())
        return box

    def test_load_test_joint_fixed(self):
        model = SimulationModel.from_gazebo_model('test_joint_fixed')

        self.assertEqual(model.name, 'test_joint_fixed',
                         'Model has the wrong name')

        PARAMS = dict(
            allow_auto_disable=1,
            static=1,
            self_collide=0
        )

        for tag in PARAMS:
            self.assertEqual(
                getattr(
                    model,
                    tag),
                PARAMS[tag],
                'Element {} should have value {}, retrieved={}'.format(
                    tag,
                    PARAMS[tag],
                    getattr(
                        model,
                        tag)))

        self.assertEqual(
            model.pose.position.tolist(),
            [0, 0, 0],
            'Model position should be [0, 0, 0], '
            'retrieved={}'.format(
                model.pose.position))
        self.assertEqual(
            model.pose.rpy,
            [0, 0, 0],
            'Model orientation should be [0, 0, 0], '
            'retrieved={}'.format(
                model.pose.rpy))

        self.check_number_of_links(model, 2)
        self.check_number_of_joints(model, 1)

        self.check_link_names(model, ['link_1', 'link_2'])
        self.check_joint_names(model, ['fixed_joint'])

        link = model.get_link_by_name('link_1')
        self.check_pose(link, [0, 0, 0], [0, 0, 0])

        self.check_attributes(
            link.inertial,
            dict(mass=10, ixx=0.1, iyy=0.1, izz=0.1, ixy=0, ixz=0, iyz=0))

        link = model.get_link_by_name('link_2')
        self.check_pose(link, [10, 0, 0], [0.2, 0.4, 0])

        self.check_attributes(
            link.inertial,
            dict(mass=20, ixx=0.2, iyy=0.2, izz=0.2, ixy=0, ixz=0, iyz=0))

    def test_load_test_static_model_object(self):
        sdf = get_gazebo_model_sdf('test_static_model')

        self.assertEqual(
            sdf._NAME,
            'sdf',
            'Parsed SDF should be of type SDF, retrieved={}'.format(
                sdf._NAME))
        self.assertEqual(len(sdf.models), 1, 'SDF should have one model')

        sdf = sdf.models[0]

        model = SimulationModel.from_sdf(sdf)

        self.check_attributes(
            model,
            dict(
                name='test_static_model',
                allow_auto_disable=1,
                static=1,
                self_collide=0
            )
        )

        self.check_pose(model, [0, 0, 0], [0, 0, 0])

        self.check_number_of_links(model, 1)
        self.check_number_of_joints(model, 0)
        self.check_link_names(model, ['body'])

        link = model.get_link_by_name('body')

        self.check_attributes(
            link,
            dict(
                name='body',
                self_collide=0,
                kinematic=0
            )
        )

        self.assertEqual(len(link.visuals), 1,
                         'Link should have one visual element')
        self.assertIsNotNone(link.get_visual_by_name('body_visual'),
                             'Cannot find visual of name <body_visual>')

        visual = link.get_visual_by_name('body_visual')

        self.check_attributes(
            visual,
            dict(
                name='body_visual',
                cast_shadows=1,
                transparency=0
            )
        )

        self.assertEqual(visual.geometry.get_type(), 'box',
                         'Visual geometry should be a box')

        self.check_pose(visual, [0, 0.4, 0], [0, -0.7, 0])

        BOX_SIZE = [1, 2, 3]
        self.assertEqual(
            visual.geometry.get_param('size'),
            BOX_SIZE,
            'Visual geometry box should have size {}, retrieved={}'.format(
                BOX_SIZE,
                visual.geometry.get_param('size')))

        self.assertEqual(len(link.collisions), 1,
                         'Link should have one collision element')
        self.assertIsNotNone(link.get_collision_by_name('body_collision'),
                             'Cannot find collision of name <body_collision>')

        collision = link.get_collision_by_name('body_collision')

        self.check_attributes(
            collision,
            dict(
                name='body_collision',
                max_contacts=10
            )
        )

        self.assertEqual(collision.geometry.get_type(), 'cylinder',
                         'Collision geometry should be a cylinder')
        CYL_RADIUS = 2.4
        CYL_LENGTH = 5.6
        self.assertEqual(
            collision.geometry.get_param('radius'),
            CYL_RADIUS,
            'Collision cylinder radius is wrong, '
            'retrieved={}, expected={}'.format(
                collision.geometry.get_param('radius'),
                CYL_RADIUS))
        self.assertEqual(
            collision.geometry.get_param('length'),
            CYL_LENGTH,
            'Collision cylinder length is wrong, '
            'retrieved={}, expected={}'.format(
                collision.geometry.get_param('length'),
                CYL_LENGTH))

        self.check_pose(collision, [0.6, 1.2, 0.05], [0, 0, 0])

        PARAMS = dict(
            mu=1,
            mu2=0.5,
            slip1=0,
            slip2=0,
            fdir1=[0, 0, 0]
        )

        for tag in PARAMS:
            self.assertEqual(
                collision.get_ode_friction_param(tag),
                PARAMS[tag],
                'Element {} should have value {}, retrieved={}'.format(
                    tag,
                    PARAMS[tag],
                    collision.get_ode_friction_param(tag)))

        PARAMS = dict(
            soft_cfm=0,
            soft_erp=None,
            kp=None,
            kd=None,
            max_vel=None,
            min_depth=None
        )

        for tag in PARAMS:
            self.assertEqual(
                collision.get_ode_contact_param(tag),
                PARAMS[tag],
                'Element {} should have value {}, retrieved={}'.format(
                    tag,
                    PARAMS[tag],
                    collision.get_ode_contact_param(tag)))

    def test_load_test_static_model_sdf(self):
        sdf = get_gazebo_model_sdf('test_static_model')

        self.assertEqual(
            sdf._NAME,
            'sdf',
            'Parsed SDF should be of type SDF, retrieved={}'.format(
                sdf._NAME))
        self.assertEqual(len(sdf.models), 1, 'SDF should have one model')

        sdf = sdf.models[0]

        self.assertEqual(sdf.name, 'test_static_model',
                         'Model has the wrong name')

        self.check_attributes(
            sdf,
            dict(
                allow_auto_disable=1,
                static=1,
                self_collide=0,
                pose=[0 for _ in range(6)]
            )
        )

        # Parsing links
        self.assertEqual(len(sdf.links), 1, 'Model should have one link')

        sdf = sdf.links[0]
        self.assertEqual(sdf.name, 'body', 'Link should be named <body>')

        PARAMS = dict(
            name='body',
            pose=[1, 0, 3.5, 0, 0, 0],
            self_collide=False,
            kinematic=False
        )

        self.check_attributes(sdf, PARAMS)

        # Parsing visual
        self.assertEqual(len(sdf.visuals), 1,
                         'Link should have one visual element')
        self.assertEqual(
            sdf.visuals[0].name,
            'body_visual',
            'Visual must have name body_visual')
        self.assertIsNotNone(sdf.visuals[0].geometry.box,
                             'Visual has no box element')

        BOX_SIZE = [1, 2, 3]
        self.assertEqual(
            sdf.visuals[0].geometry.box.size.value,
            BOX_SIZE,
            'Box has the wrong size, retrieved={}, expected={}'.format(
                sdf.visuals[0].geometry.box.size.value,
                BOX_SIZE))

        POSE_VISUAL = [0, 0.4, 0, 0, -0.7, 0]
        self.assertEqual(
            sdf.visuals[0].pose.value,
            POSE_VISUAL,
            'Visual pose is wrong, retrieved={}, expected={}'.format(
                sdf.visuals[0].pose.value,
                POSE_VISUAL))

        self.assertTrue(sdf.visuals[0].cast_shadows.value,
                        'cast_shadows should be enabled')
        self.assertFalse(sdf.visuals[0].transparency.value,
                         'transparency should be disabled')

        # Parsing collision
        self.assertEqual(len(sdf.collisions), 1,
                         'Link should have one collision element')
        self.assertEqual(sdf.collisions[0].name, 'body_collision',
                         'Collision must have name body_collision')
        self.assertIsNotNone(sdf.collisions[0].geometry.cylinder,
                             'Collision has no cylinder element')

        CYL_RADIUS = 2.4
        CYL_LENGTH = 5.6
        self.assertEqual(
            sdf.collisions[0].geometry.cylinder.radius.value,
            CYL_RADIUS,
            'Collision cylinder radius is wrong, '
            'retrieved={}, expected={}'.format(
                sdf.collisions[0].geometry.cylinder.radius.value,
                CYL_RADIUS))
        self.assertEqual(
            sdf.collisions[0].geometry.cylinder.length.value,
            CYL_LENGTH,
            'Collision cylinder length is wrong, '
            'retrieved={}, expected={}'.format(
                sdf.collisions[0].geometry.cylinder.length.value,
                CYL_LENGTH))

        self.assertEqual(sdf.collisions[0].max_contacts.value, 10,
                         'Max. number of contacts is wrong')

        POSE_COLLISION = [0.6, 1.2, 0.05, 0, 0, 0]
        self.assertEqual(
            sdf.collisions[0].pose.value,
            POSE_COLLISION,
            'Collision pose is wrong, retrieved={}, expected={}'.format(
                sdf.collisions[0].pose.value,
                POSE_COLLISION))

        MU = 1
        MU2 = 0.5
        self.assertEqual(
            sdf.collisions[0].surface.friction.ode.mu.value,
            MU,
            'Collision mu is wrong, retrieved={}, expected={}'.format(
                sdf.collisions[0].surface.friction.ode.mu.value,
                MU))
        self.assertEqual(
            sdf.collisions[0].surface.friction.ode.mu2.value,
            MU2,
            'Collision mu2 is wrong, retrieved={}, expected={}'.format(
                sdf.collisions[0].surface.friction.ode.mu2.value,
                MU2))

        SLIP = 0.5
        self.assertEqual(
            sdf.collisions[0].surface.friction.torsional.ode.slip.value,
            SLIP,
            'Collision torsional slip is wrong, retrieved={}, '
            'expected={}'.format(
                sdf.collisions[0].surface.friction.torsional.ode.slip.value,
                SLIP))

    def test_load_gazebo_models(self):
        load_gazebo_models()

        self.assertGreater(len(get_gazebo_model_names()), 1,
                           'There should be at least one model in the list')

        root_test_model_path = os.path.join(
            CUR_DIR, 'gazebo_models')

        LIST_TEST_MODELS = dict(
            test_static_model=dict(
                path=os.path.join(root_test_model_path, 'test_static_model'),
                sdf=['model.sdf']))

        for test_model in LIST_TEST_MODELS:
            self.assertIn(test_model, get_gazebo_model_names(),
                          'Model {} not found in list of Gazebo models')

            self.assertEqual(
                get_gazebo_model_path(test_model),
                LIST_TEST_MODELS[test_model]['path'],
                'Invalid path returned for model {}'.format(
                    test_model))

            self.assertEqual(
                get_gazebo_model_sdf_filenames(test_model),
                LIST_TEST_MODELS[test_model]['sdf'],
                'Invalid list of SDF files returned '
                'for model {}'.format(test_model))

        INVALID_MODELS = [generate_random_string(4) for _ in range(10)]

        for model_name in INVALID_MODELS:
            self.assertIsNone(
                get_gazebo_model_path(model_name),
                'Function should return None for non-existent model')

    def test_model(self):
        pass

    def test_joint(self):
        # Test assertion error at joint constructor with missing arguments
        joint_args = [
            dict(
                name=None,
                parent=generate_random_string(3),
                child=generate_random_string(3)),
            dict(
                name=generate_random_string(3),
                parent=None,
                child=generate_random_string(3)),
            dict(
                name=generate_random_string(3),
                parent=generate_random_string(3),
                child=None),
            dict(
                name=generate_random_string(3),
                parent=generate_random_string(3),
                child=generate_random_string(3),
                joint_type='abc')]
        for item in joint_args:
            with self.assertRaises(AssertionError):
                joint = Joint(**item)

        joint_types = ['revolute', 'revolute2', 'gearbox', 'prismatic',
                       'ball', 'screw', 'universal', 'fixed']

        N_TESTS = 10
        valid_inputs = [dict(
            name=generate_random_string(3),
            parent=generate_random_string(3),
            child=generate_random_string(3),
            joint_type=choice(joint_types)) for _ in range(N_TESTS)]

        for joint_input in valid_inputs:
            joint = Joint(**joint_input)
            self.assertEqual(
                joint.name,
                joint_input['name'],
                'Joint name was not set, expected={}, '
                'retrieved={}'.format(
                    joint_input['name'],
                    joint.name))
            self.assertEqual(
                joint.parent,
                joint_input['parent'],
                'Joint parent name was not set, expected={}, '
                'retrieved={}'.format(
                    joint_input['parent'],
                    joint.parent))
            self.assertEqual(
                joint.child,
                joint_input['child'],
                'Joint child name was not set, expected={}, '
                'retrieved={}'.format(
                    joint_input['child'],
                    joint.child))
            self.assertEqual(
                joint.type,
                joint_input['joint_type'],
                'Joint type was not set, expected={}, '
                'retrieved={}'.format(
                    joint_input['joint_type'],
                    joint.type))

        # Testing parameter setting for fixed joints
        joint = Joint(
            name=generate_random_string(3),
            parent=generate_random_string(3),
            child=generate_random_string(3),
            joint_type='fixed')

        self.assertFalse(joint.set_axis_xyz([0, 0, 1]),
                         'Axis vector should be ignored for fixed joints')
        self.assertFalse(joint.set_axis_limits(
            lower=-1 * random(),
            upper=random()), 'Fixed joints have no limits')
        self.assertFalse(joint.set_axis_dynamics(
            damping=random(), friction=random()),
            'Fixed joints have no dynamic parameters')

        # TODO Add tests for universal and revolute2 joints
        # TODO Add tests for revolute and prismatic joints
        # TODO Add SDF parsing tests

    def test_sphere_object(self):
        sphere = Sphere()
        self.assertEqual(sphere.radius, 1, 'Default radius must be 1')

        N_TESTS = 10
        valid_radius = [random() for _ in range(10)]
        names = [generate_random_string(5) for _ in range(N_TESTS)]

        for s_radius, s_name in zip(valid_radius, names):
            sphere = Sphere(name=s_name, radius=s_radius)
            self.assertEqual(
                sphere.radius,
                s_radius,
                'Radius should be {}'.format(s_radius))
            self.assertEqual(
                sphere.name,
                s_name,
                'Name should be {}'.format(s_name))

            # Test setting attributes
            sphere = Sphere()
            sphere.name = s_name
            sphere.radius = s_radius
            self.assertEqual(
                sphere.radius,
                s_radius,
                'Radius should be {}'.format(s_radius))
            self.assertEqual(
                sphere.name,
                s_name,
                'Name should be {}'.format(s_name))

            sdf = sphere.to_sdf('sphere')
            self.assertEqual(
                sdf.radius.value,
                s_radius,
                'SDF converted radius is invalid')

        # Test setting different pose vectors (position + quaternion)
        valid_positions = [[random() for _ in range(3)]
                           for _ in range(N_TESTS)]

        for pos in valid_positions:
            sphere = Sphere()
            sphere.pose = pos + [0.0 for _ in range(3)]

            self.assertEqual(
                sphere.pose.position.tolist(),
                pos,
                'Sphere position should be {}, retrieved={}'.format(
                    pos,
                    sphere.pose.position))

        sdf_types = [
            'sphere',
            'geometry',
            'collision',
            'visual',
            'link',
            'model']
        for t in sdf_types:
            sphere = Sphere()
            radius = random()
            name = generate_random_string(5)
            # Setting random size
            sphere.radius = radius
            sphere.name = name
            sdf = sphere.to_sdf(t)

            self.assertEqual(
                sdf._NAME, t, 'Wrong SDF type in request for {}'.format(t))
            self.assertTrue(
                sdf.is_valid(),
                'Invalid SDF structure was returned')

            sdf_radius = None
            if sdf._NAME == 'sphere':
                sdf_radius = sdf.radius.value
            elif sdf._NAME == 'geometry':
                sdf_radius = sdf.sphere.radius.value
            elif sdf._NAME in ['collision', 'visual']:
                sdf_radius = sdf.geometry.sphere.radius.value

            if sdf_radius is not None:
                self.assertEqual(
                    sdf_radius, radius, 'SDF element has the wrong radius,'
                    ' expected={}, received={}'.format(
                        radius, sdf_radius))

            if sdf._NAME in ['link', 'model']:
                self.assertEqual(
                    sdf.name,
                    name,
                    'Name property for {} was not set, expected={}, '
                    'retrieved={}'.format(
                        sdf._NAME,
                        name,
                        sdf.name))

    def test_cylinder_object(self):
        cylinder = Cylinder()
        self.assertEqual(cylinder.radius, 1, 'Default radius must be 1')
        self.assertEqual(cylinder.length, 1, 'Default length must be 1')

        N_TESTS = 10
        valid_sizes = [[random(), random()] for _ in range(N_TESTS)]
        names = [generate_random_string(5) for _ in range(N_TESTS)]

        for cyl_size, cyl_name in zip(valid_sizes, names):
            # Testing setting cylinder size on constructor
            cylinder = Cylinder(
                name=cyl_name,
                radius=cyl_size[0],
                length=cyl_size[1])
            self.assertEqual(
                cylinder.radius,
                cyl_size[0],
                'Radius should be {}'.format(
                    cyl_size[0]))
            self.assertEqual(
                cylinder.length,
                cyl_size[1],
                'Length should be {}'.format(
                    cyl_size[1]))
            self.assertEqual(
                cylinder.name,
                cyl_name,
                'Name should be {}'.format(cyl_name))

            # Test setting attribute
            cylinder = Cylinder()
            cylinder.name = cyl_name
            cylinder.radius = cyl_size[0]
            cylinder.length = cyl_size[1]
            self.assertEqual(
                cylinder.radius,
                cyl_size[0],
                'Radius should be {}'.format(
                    cyl_size[0]))
            self.assertEqual(
                cylinder.length,
                cyl_size[1],
                'Length should be {}'.format(
                    cyl_size[1]))
            self.assertEqual(
                cylinder.name,
                cyl_name,
                'Name should be {}'.format(cyl_name))

            sdf = cylinder.to_sdf('cylinder')
            self.assertEqual(
                sdf.radius.value,
                cyl_size[0],
                'SDF converted radius is invalid')
            self.assertEqual(
                sdf.length.value,
                cyl_size[1],
                'SDF converted length is invalid')

        # Test setting different pose vectors (position + Euler angles)
        valid_poses = \
            [[random() for _ in range(3)] +
             [random() * 2 * np.pi for _ in range(3)] for _ in range(N_TESTS)]
        for pose in valid_poses:
            cylinder = Cylinder()
            cylinder.pose = pose
            self.assertEqual(
                cylinder.pose.position.tolist(),
                pose[0:3],
                'Box position should be {}, retrieved={}'.format(
                    pose[0:3], cylinder.pose.position))

            ref_q = Pose.rpy2quat(*pose[3::])
            q = cylinder.pose.quat
            diff = Pose.get_transform(ref_q, q)

            self.assertAlmostEqual(
                np.sum(diff[0:3]), 0, 7,
                'Box orientation element rpy should be {}, '
                'retrieved={}'.format(pose[3::], q))

        # Test setting different pose vectors (position + quaternion)
        valid_positions = [[random() for _ in range(3)]
                           for _ in range(N_TESTS)]
        valid_quat = [Pose.random_orientation().quat for _ in range(N_TESTS)]

        for pos, quat in zip(valid_positions, valid_quat):
            cylinder = Cylinder()
            cylinder.pose = pos + list(quat)

            self.assertEqual(
                cylinder.pose.position.tolist(),
                pos,
                'Box position should be {}, retrieved={}'.format(
                    pos,
                    cylinder.pose.position))

            ref_q = list(quat)
            q = list(cylinder.pose.quat)

            self.assertEqual(
                ref_q,
                q,
                'Quaternion was assigned incorrectly, '
                'expected={}, retrieved={}'.format(
                    ref_q,
                    q))

        sdf_types = [
            'cylinder',
            'geometry',
            'collision',
            'visual',
            'link',
            'model']
        for t in sdf_types:
            cylinder = Cylinder()
            size = [random(), random()]
            name = generate_random_string(5)
            # Setting random size
            cylinder.radius = size[0]
            cylinder.length = size[1]
            cylinder.name = name
            sdf = cylinder.to_sdf(t)

            self.assertEqual(
                sdf._NAME, t, 'Wrong SDF type in request for {}'.format(t))
            self.assertTrue(
                sdf.is_valid(),
                'Invalid SDF structure was returned')

            sdf_radius = None
            sdf_length = None
            if sdf._NAME == 'cylinder':
                sdf_radius = sdf.radius.value
                sdf_length = sdf.length.value
            elif sdf._NAME == 'geometry':
                sdf_radius = sdf.cylinder.radius.value
                sdf_length = sdf.cylinder.length.value
            elif sdf._NAME in ['collision', 'visual']:
                sdf_radius = sdf.geometry.cylinder.radius.value
                sdf_length = sdf.geometry.cylinder.length.value

            if sdf_radius is not None and sdf_length is not None:
                self.assertEqual(
                    sdf_radius, size[0], 'SDF element has the wrong radius,'
                    ' expected={}, received={}'.format(
                        size[0], sdf_radius))
                self.assertEqual(
                    sdf_length, size[1], 'SDF element has the wrong length,'
                    ' expected={}, received={}'.format(
                        size[1], sdf_length))

            if sdf._NAME in ['link', 'model']:
                self.assertEqual(
                    sdf.name,
                    name,
                    'Name property for {} was not set, expected={}, '
                    'retrieved={}'.format(
                        sdf._NAME,
                        name,
                        sdf.name))

    def test_box_object(self):
        box = Box()
        self.assertEqual(box.size, [1, 1, 1], 'Default size should be 1, 1, 1')

        # Test setting different box sizes
        N_TESTS = 10
        valid_sizes = [[random() for _ in range(3)] for _ in range(N_TESTS)]
        names = [generate_random_string(5) for _ in range(N_TESTS)]

        for box_size, box_name in zip(valid_sizes, names):
            # Test setting box size on constructor
            box = Box(name=box_name, size=box_size)
            self.assertEqual(
                box.size,
                box_size,
                'Size should be {}'.format(box_size))
            self.assertEqual(
                box.name,
                box_name,
                'Name should be {}'.format(box_name))

            # Test setting the attribute
            box = Box()
            box.name = box_name
            box.size = box_size
            self.assertEqual(
                box.size,
                box_size,
                'Size should be {}'.format(box_size))
            self.assertEqual(
                box.name,
                box_name,
                'Name should be {}'.format(box_name))

            sdf = box.to_sdf('box')
            self.assertEqual(
                sdf.size.value,
                box_size,
                'SDF converted size is invalid')

        # Test setting different pose vectors (position + Euler angles)
        valid_poses = \
            [[random() for _ in range(3)] +
             [random() * 2 * np.pi for _ in range(3)] for _ in range(N_TESTS)]
        for pose in valid_poses:
            box = Box()
            box.pose = pose
            self.assertEqual(
                box.pose.position.tolist(),
                pose[0:3],
                'Box position should be {}, retrieved={}'.format(
                    pose[0:3], box.pose.position))

            ref_q = Pose.rpy2quat(*pose[3::])
            q = box.pose.quat
            diff = Pose.get_transform(ref_q, q)

            self.assertAlmostEqual(np.sum(
                diff[0:3]), 0, 7, 'Box orientation element rpy '
                'should be {}, retrieved={}'.format(pose[3::], q))

        # Test setting different pose vectors (position + quaternion)
        valid_positions = [[random() for _ in range(3)]
                           for _ in range(N_TESTS)]
        valid_quat = [Pose.random_orientation().quat for _ in range(N_TESTS)]

        for pos, quat in zip(valid_positions, valid_quat):
            box = Box()
            box.pose = pos + list(quat)

            self.assertEqual(
                box.pose.position.tolist(),
                pos,
                'Box position should be {}, retrieved={}'.format(
                    pos,
                    box.pose.position))

            ref_q = list(quat)
            q = list(box.pose.quat)

            self.assertEqual(
                ref_q,
                q,
                'Quaternion was assigned incorrectly, '
                'expected={}, retrieved={}'.format(
                    ref_q,
                    q))

        for value in INVALID_SIZES:
            with self.assertRaises(AssertionError):
                box = Box(size=value)

        invalid_sdf_types = ['a', dict(), None, True, 1]
        for t in invalid_sdf_types:
            with self.assertRaises(AssertionError):
                box.to_sdf(t)

        sdf_types = ['box', 'geometry', 'collision', 'visual', 'link', 'model']
        for t in sdf_types:
            box = Box()
            size = [random() for _ in range(3)]
            name = generate_random_string(5)
            # Setting random size
            box.size = size
            box.name = name

            sdf = box.to_sdf(t)

            self.assertEqual(
                sdf._NAME, t, 'Wrong SDF type in request for {}'.format(t))
            self.assertTrue(
                sdf.is_valid(),
                'Invalid SDF structure was returned')

            sdf_size = None
            if sdf._NAME == 'box':
                sdf_size = sdf.size.value
            elif sdf._NAME == 'geometry':
                sdf_size = sdf.box.size.value
            elif sdf._NAME in ['collision', 'visual']:
                sdf_size = sdf.geometry.box.size.value

            if sdf_size is not None:
                self.assertEqual(
                    sdf_size, size, 'SDF element has the wrong size vector,'
                    ' expected={}, received={}'.format(
                        size, sdf_size))

            if sdf._NAME in ['link', 'model']:
                self.assertEqual(
                    sdf.name,
                    name,
                    'Name property for {} was not set, expected={}, '
                    'retrieved={}'.format(
                        sdf._NAME,
                        name,
                        sdf.name))

    def test_export_model_to_sdf(self):
        box = self.create_random_box()

        sdf = box.to_sdf('model')
        filename = '/tmp/box.sdf'
        sdf.export_xml(filename, version='1.6')

        self.assertTrue(os.path.isfile(filename))

        # Parse the exported file
        parsed_sdf = parse_sdf(filename)
        self.assertIsNotNone(parsed_sdf)
        self.assertEqual(sdf, parsed_sdf)

    def test_export_model_to_gazebo_model(self):
        box = self.create_random_box()
        model = SimulationModel.from_sdf(box.to_sdf())

        self.assertIsNotNone(model)

        # Export Gazebo model with default parameters
        model.to_gazebo_model()
        default_dir = os.path.join(
            os.path.expanduser('~'), '.gazebo', 'models')
        model_dir = os.path.join(default_dir, box.name)
        # Check if all model files were created
        self.assertTrue(os.path.isdir(model_dir))
        self.assertTrue(
            os.path.isfile(
                os.path.join(
                    model_dir,
                    'model.config')))
        self.assertTrue(os.path.isfile(os.path.join(model_dir, 'model.sdf')))

        # Parse model config file
        sdf_config = parse_sdf_config(os.path.join(model_dir, 'model.config'))
        self.assertIsNotNone(sdf_config)

        # Get name of current user, used in default model metadata input
        username = getpass.getuser()

        self.assertEqual(sdf_config.name.value, box.name)
        self.assertEqual(len(sdf_config.authors), 1)
        self.assertEqual(sdf_config.authors[0].name.value, username)
        self.assertEqual(
            sdf_config.authors[0].email.value,
            '{}@email.com'.format(username))
        self.assertEqual(sdf_config.description.value, '')
        self.assertEqual(sdf_config.version.value, '1.6')
        self.assertEqual(len(sdf_config.sdfs), 1)
        self.assertEqual(sdf_config.sdfs[0].value, 'model.sdf')
        self.assertEqual(sdf_config.sdfs[0].version, '1.6')

        sdf = parse_sdf(os.path.join(model_dir, 'model.sdf'))
        self.assertIsNotNone(sdf)
        self.assertEqual(sdf.xml_element_name, 'sdf')
        self.assertIsNotNone(sdf.models)
        self.assertEqual(len(sdf.models), 1)

        # Export Gazebo model with default parameters
        author = generate_random_string(5)
        email = generate_random_string(5)
        description = generate_random_string(10)
        model_metaname = generate_random_string(10)

        model.to_gazebo_model(
            author=author,
            email=email,
            description=description,
            model_metaname=model_metaname,
            overwrite=True
        )

        # Check if all model files were created
        self.assertTrue(os.path.isdir(model_dir))
        self.assertTrue(
            os.path.isfile(
                os.path.join(
                    model_dir,
                    'model.config')))
        self.assertTrue(os.path.isfile(os.path.join(model_dir, 'model.sdf')))

        # Parse model config file
        sdf_config = parse_sdf_config(os.path.join(model_dir, 'model.config'))
        self.assertIsNotNone(sdf_config)

        self.assertEqual(sdf_config.name.value, model_metaname)
        self.assertEqual(len(sdf_config.authors), 1)
        self.assertEqual(sdf_config.authors[0].name.value, author)
        self.assertEqual(sdf_config.authors[0].email.value, email)
        self.assertEqual(sdf_config.description.value, description)
        self.assertEqual(sdf_config.version.value, '1.6')
        self.assertEqual(len(sdf_config.sdfs), 1)
        self.assertEqual(sdf_config.sdfs[0].value, 'model.sdf')
        self.assertEqual(sdf_config.sdfs[0].version, '1.6')

        # Parse the SDF file
        sdf = parse_sdf(os.path.join(model_dir, 'model.sdf'))
        self.assertIsNotNone(sdf)
        self.assertEqual(sdf.xml_element_name, 'sdf')
        self.assertIsNotNone(sdf.models)
        self.assertEqual(len(sdf.models), 1)

        # Delete generated Gazebo model directory
        shutil.rmtree(model_dir)

    def test_convert_from_urdf(self):
        # Test correct conversion of joint relative
        # poses into absolute link poses

        # Test random sets of poses
        pose_sets = [
            [Pose.random() for _ in range(5)],
            [Pose.random_position() for _ in range(5)],
            [Pose.random_orientation() for _ in range(5)]
        ]

        for poses in pose_sets:
            urdf = create_urdf_element('robot')
            # Create links
            for i in range(len(poses) + 1):
                urdf.add_link('link_{}'.format(i))

            for i in range(len(poses)):
                joint_urdf = create_urdf_element('joint')
                joint_urdf.parent.link = 'link_{}'.format(i)
                joint_urdf.child.link = 'link_{}'.format(i + 1)
                joint_urdf.origin = poses[i].to_urdf()

                urdf.add_joint('joint_{}'.format(i), joint_urdf)

            model = SimulationModel.from_urdf(urdf)

            self.assertIsNotNone(model)

            cur_pose = None
            for i in range(len(poses) + 1):
                link = model.get_link_by_name('link_{}'.format(i))
                self.assertIsNotNone(link)

                if i == 0:
                    self.assertEqual(
                        link.pose, Pose(
                            pos=[
                                0, 0, 0], rot=[
                                0, 0, 0]))
                else:
                    if cur_pose is None:
                        cur_pose = poses[i - 1]
                    else:
                        cur_pose = cur_pose + poses[i - 1]
                    self.check_pose(link, cur_pose.position, cur_pose.rpy)

    def test_star_urdf_structure(self):
        urdf = create_urdf_element('robot')

        # Add origin link
        i = 0
        urdf.add_link('link_{}'.format(i))
        i += 1

        for _ in range(3):
            for j in range(randint(3, 5)):
                urdf.add_link('link_{}'.format(i))

                joint_urdf = create_urdf_element('joint')
                if j == 0:
                    joint_urdf.parent.link = 'link_{}'.format(0)
                else:
                    joint_urdf.parent.link = 'link_{}'.format(i - 1)
                joint_urdf.child.link = 'link_{}'.format(i)
                joint_urdf.origin = Pose.random().to_urdf()

                urdf.add_joint('joint_{}'.format(i), joint_urdf)
                i += 1

        model = SimulationModel.from_urdf(urdf)
        self.assertIsNotNone(model)

    def test_invalid_disconnected_robot_graph(self):
        urdf = create_urdf_element('robot')
        for i in range(3):
            urdf.add_link('link_{}'.format(i))

        with self.assertRaises(ValueError):
            SimulationModel.from_urdf(urdf)

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


text:
import numpy as np
from time import time
import tqdm

from xbbo.search_space import problem_register
from xbbo.search_algorithm import alg_register
from xbbo.configspace import build_space
from xbbo.utils.constants import MAXINT
from xbbo.utils.record import Record

class BBO:

    def __init__(self, cfg, seed):
        # setup TestProblem
        self.cfg = cfg
        self.rng = np.random.RandomState(seed)
        self.function_instance = problem_register[cfg.TEST_PROBLEM.name](seed=self.rng.randint(MAXINT), **cfg.TEST_PROBLEM.kwargs)

        self.api_config = self.function_instance.get_api_config()  # 优化的hp
        self.config_spaces = build_space(self.api_config,seed=self.rng.randint(MAXINT))

        # Setup optimizer
        opt_class = alg_register[cfg.OPTM.name]
        self.optimizer_instance = opt_class(self.config_spaces,total_limit=cfg.OPTM.max_call,seed=self.rng.randint(MAXINT), **dict(cfg.OPTM.kwargs))

        self.n_suggestions = cfg.OPTM.n_suggestions
        self.n_obj = cfg.OPTM.n_obj

        assert self.n_suggestions >= 1, "batch size must be at least 1"
        assert self.n_obj >= 1, "Must be at least one objective"

        # self.suggest_time = np.zeros(n_calls)
        # self.observe_time = np.zeros(n_calls)
        # self.eval_time = np.zeros((n_calls, n_suggestions))
        # self.function_evals = np.zeros((n_calls, n_suggestions, self.n_obj))
        # self.suggest_log = [None] * n_calls
        self.n_calls = cfg.OPTM.max_call
        self.record = Record(self.cfg.GENERAL.exp_dir)

    def evaluate(self, param):
        return self.function_instance.evaluate(param)

    def run(self):
        pbar = tqdm.tqdm(range(self.n_calls))
        pbar.set_description(f"Optimizer {self.cfg.OPTM.name} is running:")

        for ii in pbar:

            tt = time()
            trial_list = self.optimizer_instance.suggest(self.n_suggestions)  # TODO 1

            # try:
            #     next_points, features = self.optimizer_instance.suggest(self.n_suggestions)  # TODO 1
            # except Exception as e:
            #     # logger.warning("Failure in optimizer suggest. Falling back to random search.")
            #     # logger.exception(e, exc_info=True)
            #     print(json.dumps({"optimizer_suggest_exception": {'iter': ii}}))
            #     # api_config = self.function_instance.get_api_config()
            #     # TODO 直接随机采样
            #     x_guess_configs = self.optimizer_instance.space.sample_configuration(size=self.n_suggestions) # a list
            #     next_points = [x_guess_config.get_dict_unwarped() for x_guess_config in x_guess_configs]
            #     features = [x_guess_config.get_array() for x_guess_config in x_guess_configs]

            suggest_time = time() - tt

            assert len(trial_list) == self.n_suggestions, "invalid number of suggestions provided by the optimizer"
            # eval_time = [None for _ in range(self.n_suggestions)]
            function_evals = []
            # losses = []
            for trial in (trial_list):
                # try:
                tt = time()

                f_current_eval = self.evaluate(trial.config_dict) # TODO 2
                eval_time = time() - tt

                # except Exception as e:
                #     f_current_eval = np.full((len(self.cfg.TEST_PROBLEM.func_evals),), np.inf, dtype=float)
                #     loss = np.full((len(self.cfg.TEST_PROBLEM.losses),), np.inf, dtype=float)

                trial.add_observe_value(observe_value=f_current_eval,time=eval_time)
                function_evals.append(f_current_eval)
            # if self.cfg.OPTM.n_obj == 1:
            #     eval_list = np.asarray(function_evals)[:, :self.cfg.OPTM.n_obj].ravel().tolist() # TODO
            # else:
            #     raise NotImplementedError()
                # eval_list = np.array(losses)[:, :self.cfg.OPTM.n_obj].tolist() # TODO
            # assert self.cfg.OPTM.n_obj == 1
            tt = time()
            self.optimizer_instance.observe(trial_list)  # TODO 3
            # try:
            #     self.optimizer_instance.observe(features, eval_list)  # TODO 3
            # except Exception as e:
                # logger.warning("Failure in optimizer observe. Ignoring these observations.")
                # logger.exception(e, exc_info=True)
                # print(json.dumps({"optimizer_observe_exception": {'iter': ii}}))
            observe_time = time() - tt
            timing = {
                         'suggest_time_per_suggest': suggest_time,
                         'observe_time_per_suggest': observe_time,
                         'eval_time_per_suggest': sum(trial.time for trial in trial_list)
            }
            self.record.append([trial.dense_array if trial.sparse_array is None else trial.sparse_array for trial in trial_list], function_evals, timing=timing, suggest_point=[trial.config_dict for trial in trial_list])
            # print(self.optimizer_instance.trials.best_observe_value)
            print(function_evals)

        print(self.optimizer_instance.trials.best_observe_value)

class BBO_REBAR(BBO):
    def __init__(self):
        BBO.__init__(self)

    def evaluate(self,params):
        return BBO.evaluate(self, params)

    def run(self):
        pass

ext: py

reason: import numpy


text:
'''
Author: your name
Date: 2020-12-21 14:27:25
LastEditTime: 2020-12-23 10:28:06
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: /Workflows/server/PycharmProjects/Pacific_AvatarGame_Host/utils/view_human_pose_3d/view_pose_pyplot.py
'''

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import time

class HumanPose3DViewer():
    def __init__(self, human_mode):
        # plt.ion()  # for interactive plot
        self.fig = plt.figure("3D Human Pose Estimation", figsize=(10, 10))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ax.view_init(20, 45)  # unity view for human pose

        if human_mode == "pacific" or human_mode == "PACIFIC":
            self.data_order = [[6,2,1,0], [6,3,4,5], [6,7,8,9,10], [8,13,12,11], [8,14,15,16]]
        elif human_mode == "mpii" or human_mode == "MPII":
            self.data_order = [[6,2,1,0], [6,3,4,5], [6,7,8,9], [8,12,11,10], [8,13,14,15]]
        self.data_color = ['red', 'blue', 'green', 'red', 'blue']

        self.x_data = []
        self.y_data = []
        self.z_data = []

    def configureViewAngle(self):
        self.elev = self.ax.elev
        self.azim = self.ax.azim

        self.ax.view_init(self.elev, self.azim)

    def configureAxis(self, axis_range, axis_scale):
        self.ax.set_title("3D Human Pose Estimation")

        self.ax.set_xlabel('X-axis')
        self.ax.set_ylabel('Y-axis')
        self.ax.set_zlabel('Z-axis')
        self.ax.xaxis.label.set_color('red')
        self.ax.yaxis.label.set_color('green')
        self.ax.zaxis.label.set_color('blue')

        # TODO: Set the axis arrow type and the axis color
        # self.ax.axis['x'].set_axisline_style('-|>', size=3)
        # self.ax.axis['y'].set_axisline_style('-|>', size=3)
        # self.ax.axis['z'].set_axisline_style('-|>', size=3)
        self.ax.tick_params(axis='x', colors='red')
        self.ax.tick_params(axis='y', colors='green')
        self.ax.tick_params(axis='z', colors='blue')

        self.ax.set_xlim3d([axis_range[0][0], axis_range[0][1]])
        self.ax.set_ylim3d([axis_range[1][0], axis_range[1][1]])
        self.ax.set_zlim3d([axis_range[2][0] + axis_range[2][1], axis_range[2][1] - axis_range[2][0]])
        self.ax.set_xticks(np.arange(axis_range[0][0], axis_range[0][1], axis_scale[0]))
        self.ax.set_yticks(np.arange(axis_range[1][0], axis_range[1][1], axis_scale[1]))
        self.ax.set_zticks(np.arange(axis_range[2][0] + axis_range[2][1], axis_range[2][1] - axis_range[2][0], axis_scale[2]))

    def viewPose(self, human_kps_3d, axis_range, axis_scale):
        data_shape = human_kps_3d.shape
        human_data = human_kps_3d.tolist()

        for human in range(data_shape[0]):
            for branch in range(len(self.data_order)):
                for idx in range(len(self.data_order[branch])):
                    self.x_data.append(human_data[human][self.data_order[branch][idx]][0])
                    self.y_data.append(human_data[human][self.data_order[branch][idx]][1])
                    self.z_data.append(human_data[human][self.data_order[branch][idx]][2])  # set the axises for current(0,1,2)/unity(0,2,1) order

                    self.ax.text(human_data[human][self.data_order[branch][idx]][0], human_data[human][self.data_order[branch][idx]][1], \
                        human_data[human][self.data_order[branch][idx]][2], self.data_order[branch][idx], size='medium')

                self.configureViewAngle()
                self.configureAxis(axis_range, axis_scale)
                self.ax.scatter(self.x_data, self.y_data, self.z_data, color=self.data_color[branch], marker='o', s=10)
                self.ax.plot(self.x_data, self.y_data, self.z_data, color=self.data_color[branch], linewidth=1.5)

                self.x_data.clear()
                self.y_data.clear()
                self.z_data.clear()

        plt.pause(0.0005)
        plt.cla()

    def closeAll(self):
        plt.clf()
        plt.close()


text:
# -*- coding: utf-8 -*-
"""Testing machine type checkers for scitypes."""

__author__ = ["fkiraly"]

import numpy as np

from sktime.datatypes import MTYPE_REGISTER, SCITYPE_REGISTER
from sktime.datatypes._check import check_dict, check_is_mtype
from sktime.datatypes._check import mtype as infer_mtype
from sktime.datatypes._examples import get_examples

SCITYPES = [sci[0] for sci in SCITYPE_REGISTER]

# scitypes where mtype inference is not unique
# alignment is excluded since mtypes can be ambiguous
#   (indices could be both loc or iloc when integers)
SCITYPES_AMBIGUOUS_MTYPE = ["Alignment"]

def _get_all_mtypes_for_scitype(scitype):
    """Return list of all mtypes for scitype.

    Parameters
    ----------
    scitype : str - scitype

    Returns
    -------
    mtypes : list of str - list of mtypes for scitype
    """
    if scitype not in [s[0] for s in SCITYPE_REGISTER]:
        raise RuntimeError(scitype + " is not in the SCITYPE_REGISTER")
    mtypes = [key[0] for key in MTYPE_REGISTER if key[1] == scitype]

    if len(mtypes) == 0:
        # if there are no mtypes, this must have been reached by mistake/bug
        raise RuntimeError("no mtypes defined for scitype " + scitype)

    return mtypes

def _generate_scitype_mtype_combinations():
    """Return scitype/mtype tuples for pytest_generate_tests.

    Fixtures parameterized
    ----------------------
    scitype : str - scitype of fixture
    mtype : str - mtype of fixture
    """
    # collect fixture tuples here

    sci_mtype_tuples = []

    for scitype in SCITYPES:

        mtypes = _get_all_mtypes_for_scitype(scitype)

        for mtype in mtypes:
            sci_mtype_tuples += [(scitype, mtype)]

    return sci_mtype_tuples

def _generate_scitype_mtype_fixtureindex_combinations():
    """Return fixture tuples for pytest_generate_tests.

    Fixtures parameterized
    ----------------------
    scitype : str - scitype of fixture
    mtype : str - mtype of fixture
    fixture_index : int - index of fixture tuple with that scitype and mtype
    """
    # collect fixture tuples here

    sci_mtype_tuples = _generate_scitype_mtype_combinations()

    sci_mtype_index_tuples = []

    for tuple_j in sci_mtype_tuples:
        scitype = tuple_j[0]
        mtype = tuple_j[1]
        n_fixtures = len(get_examples(mtype=mtype, as_scitype=scitype))

        for i in range(n_fixtures):
            sci_mtype_index_tuples += [(scitype, mtype, i)]

    return sci_mtype_index_tuples

def pytest_generate_tests(metafunc):
    """Test parameterization routine for pytest.

    Fixtures parameterized
    ----------------------
    scitype : str - scitype of fixture
    mtype : str - mtype of fixture
    fixture_index : int - index of fixture tuple with that scitype and mtype
    """
    # we assume all four arguments are present in the test below

    fixturenames = set(metafunc.fixturenames)

    if set(["scitype", "mtype", "fixture_index"]).issubset(fixturenames):
        keys = _generate_scitype_mtype_fixtureindex_combinations()

        ids = []
        for tuple in keys:
            ids += [f"{tuple[0]}-{tuple[1]}-fixture:{tuple[2]}"]

        # parameterize test with from-mtpes
        metafunc.parametrize("scitype,mtype,fixture_index", keys, ids=ids)

    elif set(["scitype", "mtype"]).issubset(fixturenames):
        keys = _generate_scitype_mtype_combinations()

        ids = []
        for tuple in keys:
            ids += [f"{tuple[0]}-{tuple[1]}"]

        # parameterize test with from-mtpes
        metafunc.parametrize("scitype,mtype", keys, ids=ids)

def test_check_positive(scitype, mtype, fixture_index):
    """Tests that check_is_mtype correctly confirms the mtype of examples.

    Parameters
    ----------
    scitype : str - name of scitype for which mtype conversions are tested

    Raises
    ------
    RuntimeError if scitype is not defined or has no mtypes or examples
    AssertionError if examples are not correctly identified
    error if check itself raises an error
    """
    # retrieve fixture for checking
    fixture = get_examples(mtype=mtype, as_scitype=scitype).get(fixture_index)

    # todo: possibly remove this once all checks are defined
    check_is_defined = (mtype, scitype) in check_dict.keys()

    # check fixtures that exist against checks that exist
    if fixture is not None and check_is_defined:
        check_result = check_is_mtype(fixture, mtype, scitype, return_metadata=True)
        if not check_result[0]:
            msg = (
                f"check_is_mtype returns False on scitype {scitype}, mtype {mtype} "
                f"fixture {fixture_index}, message: "
            )
            msg = msg + check_result[1]
        assert check_result[0], msg

def test_check_metadata_inference(scitype, mtype, fixture_index):
    """Tests that check_is_mtype correctly infers metadata of examples.

    Parameters
    ----------
    scitype : str - name of scitype for which mtype conversions are tested

    Raises
    ------
    RuntimeError if scitype is not defined or has no mtypes or examples
    AssertionError if example metadata is not correctly inferred
    error if check itself raises an error
    """
    # retrieve fixture for checking
    fixture, _, expected_metadata = get_examples(
        mtype=mtype, as_scitype=scitype, return_metadata=True
    ).get(fixture_index)

    # todo: possibly remove this once all checks are defined
    check_is_defined = (mtype, scitype) in check_dict.keys()
    # if the examples have no metadata to them, don't test
    metadata_provided = expected_metadata is not None

    # check fixtures that exist against checks that exist
    if fixture is not None and check_is_defined and metadata_provided:
        check_result = check_is_mtype(fixture, mtype, scitype, return_metadata=True)
        metadata = check_result[2]

        # remove mtype & scitype key if exists, since comparison is on scitype level
        if "mtype" in metadata:
            del metadata["mtype"]
        if "scitype" in metadata:
            del metadata["scitype"]

        msg = (
            f"check_is_mtype returns wrong metadata on scitype {scitype}, "
            f"mtype {mtype}, fixture {fixture_index}. "
            f"returned: {metadata}; expected: {expected_metadata}"
        )

        assert metadata == expected_metadata, msg

def test_check_negative(scitype, mtype):
    """Tests that check_is_mtype correctly identifies wrong mtypes of examples.

    Parameters
    ----------
    scitype : str - name of scitype for which mtype conversions are tested

    Raises
    ------
    RuntimeError if scitype is not defined or has no mtypes or examples
    AssertionError if a examples are correctly identified as incompatible
    error if check itself raises an error
    """
    # if the scitype is ambiguous, we can't assume that other mtypes are negative
    if scitype in SCITYPES_AMBIGUOUS_MTYPE:
        return None

    mtypes = _get_all_mtypes_for_scitype(scitype)
    fixtures = dict()

    for other_mtype in mtypes:
        fixtures[other_mtype] = get_examples(mtype=other_mtype, as_scitype=scitype)

    n_fixtures = np.max([len(fixtures[mtype]) for mtype in mtypes])

    if n_fixtures == 0:
        raise RuntimeError("no fixtures defined for scitype " + scitype)

    for i in range(n_fixtures):
        # if mtype is not ambiguous, other mtypes are negative examples
        for wrong_mtype in list(set(mtypes).difference(set([mtype]))):

            # retrieve fixture for checking
            fixture_wrong_type = fixtures[wrong_mtype].get(i)

            # todo: possibly remove this once all checks are defined
            check_is_defined = (mtype, scitype) in check_dict.keys()

            # check fixtures that exist against checks that exist
            if fixture_wrong_type is not None and check_is_defined:
                assert not check_is_mtype(fixture_wrong_type, mtype, scitype), (
                    f"check_is_mtype {mtype} returns True "
                    f"on {wrong_mtype} fixture {i}"
                )

def test_mtype_infer(scitype, mtype, fixture_index):
    """Tests that mtype correctly infers the mtype of examples.

    Parameters
    ----------
    scitype : str - name of scitype for which mtype conversions are tested

    Raises
    ------
    RuntimeError if scitype is not defined or has no mtypes or examples
    AssertionError if mtype of examples is not correctly identified
    error if check itself raises an error
    """
    # if mtypes are ambiguous, then this test should be skipped
    if scitype in SCITYPES_AMBIGUOUS_MTYPE:
        return None

    # retrieve fixture for checking
    fixture = get_examples(mtype=mtype, as_scitype=scitype).get(fixture_index)

    # todo: possibly remove this once all checks are defined
    check_is_defined = (mtype, scitype) in check_dict.keys()

    # check fixtures that exist against checks that exist
    if fixture is not None and check_is_defined:
        assert mtype == infer_mtype(
            fixture, as_scitype=scitype
        ), f"mtype {mtype} not correctly identified for fixture {fixture_index}"


text:
"""Helper functions.

Helper functions to suport other modules.

"""
import numpy as np

from scipy.ndimage import label

import torch

__all__ = ['label_connected_components', 'compute_importance']

def label_connected_components(label_images, start_label=1, is_3d=False):
    """Label connected components in a label image.

    Create new label images where labels are changed so that each \
    connected componnent has a diffetent label. \
    To find the connected components, it is used a 8-neighborhood.

    Parameters
    ----------
    label_images : ndarray
        Label images with size :math:`(N, H, W)`.

    start_labeel: int
        First label, by default 1.

    Returns
    -------
    ndarray
        Label images with new labels.

    """
    # label_image = label_image.squeeze()

    if label_images.ndim == 2:
        label_images = np.expand_dims(label_images, 0)

    if is_3d and label_images.ndim == 3:
        label_images = np.expand_dims(label_images, 0)

    new_label_images = np.zeros_like(label_images).astype(np.int)

    _c = start_label

    for label_image, new_label_image in zip(label_images, new_label_images):
        num_labels = label_image.astype(np.int32).max()

        #new_label_image = np.zeros_like(label_image).astype(np.int)
        if is_3d:
            structure = np.ones((3, 3, 3), dtype=np.uint8)
        else:
            structure = np.ones((3, 3), dtype=np.uint8)

        for _l in range(1, num_labels + 1):
            _label_image = label(label_image == _l, structure=structure)[0]

            for new_l in range(1, _label_image.max() + 1):
                mask = _label_image == new_l
                new_label_image[mask] = _c
                _c += 1

    return new_label_images

def _normalize_image(image):
    f = image

    max = f.max()
    min = f.min()
    if max - min > 0:
        f = image
        f = 1*(f - min)/(max - min)
    elif max > 0:
        f = f/max

    return f

def _compute_channel_importance(image_features, image_markers, label):
    importance_to_label = {}
    num_channels = image_features[0].shape[0]

    pos_act_by_channel = np.zeros(num_channels)
    neg_act_by_channel = np.zeros(num_channels)

    importance_to_label = np.zeros(num_channels)
    num_label_pixels = 0
    num_other_label_pixels = 0

    for features, markers in zip(image_features, image_markers):
        other_label_mask = np.logical_and(markers != label, markers != 0)
        label_mask = markers == label
        num_label_pixels += label_mask.sum()
        num_other_label_pixels += other_label_mask.sum()

        for i, channel in enumerate(features):
            _channel = _normalize_image(channel)
            pos_act_by_channel[i] = (_channel[label_mask].sum())
            neg_act_by_channel[i] = 1 - _channel[other_label_mask].sum()

    # reg_factor_pos = 1 - num_label_pixels/(num_label_pixels + num_other_label_pixels)
    # reg_factor_neg = 1 - reg_factor_pos

    importance_to_label = pos_act_by_channel + neg_act_by_channel
    importance_to_label = importance_to_label - importance_to_label.mean()

    #importance_to_label = importance_to_label/np.linalg.norm(importance_to_label)

    return importance_to_label

def compute_importance(images, markers, n_classes):

    importance_by_channel = []

    images = images.transpose(0, 3, 1, 2)

    for label in range(n_classes):

        importance_by_channel.append(_compute_channel_importance(images, markers, label+1))

    return np.array(importance_by_channel)

max_stars_repo_name: LIDS-UNICAMP/FLIM

max_issues_repo_name: LIDS-UNICAMP/FLIM

max_forks_repo_name: LIDS-UNICAMP/FLIM

reason: import numpy,from scipy


text:
#!/usr/bin/env python
# -*- coding: utf8 -*-
#
# Copyright (c) 2013 Nautilabs
#
# Licensed under the MIT License,
# https://github.com/baptistelabat/robokite
# Authors: Baptiste LABAT,
#
# This script is used as a gateway between a mobile phone and an arduino commanding a motor
# The mobile phone sensors can then be used to control the motor
#
# The sensors data are streamed over wifi by sensorUDP application on android mobile phone.
# The UDP packet are read by the python script then forwarded
import os
import sys
import threading
import time
import numpy as np
import serial
sys.path.append(os.getcwd())
sys.path.append('../Sensors')
import mobileState

def computeXORChecksum(chksumdata):
	# Inspired from http://doschman.blogspot.fr/2013/01/calculating-nmea-sentence-checksums.html
    # Initializing XOR counter
    csum = 0

    # For each char in chksumdata, XOR against the previous
    # XOR char.  The final XOR of the last char will be the
    # checksum
    for c in chksumdata:
        # Makes XOR value of counter with the next char in line
        # and stores the new XOR value in csum
        csum ^= ord(c)
    h = hex(csum)
    return h[2:]#get hex data without 0x prefix

try:
  # Get the mobile orientation
  mobile = mobileState.mobileState()
  a = threading.Thread(None, mobileState.mobileState.checkUpdate, None, (mobile,))
  a.start()

  # Open the serial port
  locations=['/dev/ttyACM0','/dev/ttyACM1','/dev/ttyACM2','/dev/ttyUSB0','/dev/ttyUSB1','/dev/ttyUSB2','/dev/ttyUSB3','/dev/ttyS0','/  dev/ttyS1','/dev/ttyS2','/dev/ttyS3']
  for device in locations:
    try:
      print "Trying...",device
      ser = serial.Serial(
	port = device,
	baudrate = 9600
	)
      break
    except:
      print "Failed to connect on",device
  time.sleep(1.5) # Arduino is reset when opening port so wait before communicating
  ser.flush()
  ser.write('i1')
  dt = 0.1
  while True:
      mobile.computeRPY()
      ser.flush()
      alpha = np.max([-0.99, np.min([0.99, np.floor(100*mobile.pitch/2.0)/100.0])]
                          )
      msg = "ORPWM"+","+str(alpha)
      msg = "$"+msg +"*"+ computeXORChecksum(msg) + chr(13).encode('ascii')
      ser.write(msg
        )
      print msg
      time.sleep(dt)
except KeyboardInterrupt:
 mobile.stop_requested = True

max_stars_repo_name: alsuren/robokite

max_issues_repo_name: alsuren/robokite

max_forks_repo_name: alsuren/robokite


text:
"""
Customized data utils
The codes were refactored from [tensornets](https://github.com/taehoonlee/tensornets/blob/master/tensornets/datasets/voc.py)
"""
from __future__ import division

import os
import numpy as np
import xml.etree.ElementTree as ET

try:
    import cv2
except ImportError:
    cv2 = None

try:
    xrange          # Python 2
except NameError:
    xrange = range  # Python 3

try:
    reduce
except NameError:
    from functools import reduce

with open(os.path.join(os.path.dirname(__file__), 'yolf.names'), 'r') as f:
    classnames = [line.rstrip() for line in f.readlines()]

classes_num = 20
used_classes = len(classnames)

def classidx(classname):
    return dict((k, i) for (i, k) in enumerate(classnames))[classname]

def area(box):
    if box.ndim == 1:
        return (box[2] - box[0] + 1.) * (box[3] - box[1] + 1.)
    else:
        return (box[:, 2] - box[:, 0] + 1.) * (box[:, 3] - box[:, 1] + 1.)

def get_files(data_dir, data_name, total_num=None):
    with open("%s/ImageSets/Main/%s.txt" % (data_dir, data_name)) as f:
        files = [x.strip() for x in f.readlines()]
    if total_num is not None:
        files = files[:total_num]
    return files

def get_annotations(data_dir, files):
    annotations = {}
    for filename in files:
        tree = ET.parse("%s/Annotations/%s.xml" % (data_dir, filename))
        annotations[filename] = [[] for _ in range(classes_num)]
        for obj in tree.findall('object'):
            obj_struct = {}
            obj_struct['name'] = obj.find('name').text
            obj_struct['pose'] = obj.find('pose').text
            obj_struct['truncated'] = int(obj.find('truncated').text)
            obj_struct['difficult'] = int(obj.find('difficult').text)
            bbox = obj.find('bndbox')
            obj_struct['bbox'] = [int(bbox.find('xmin').text),
                                  int(bbox.find('ymin').text),
                                  int(bbox.find('xmax').text),
                                  int(bbox.find('ymax').text)]
            cidx = classidx(obj_struct['name'])
            annotations[filename][cidx].append(obj_struct)
    return annotations

def load(data_dir, data_name, min_shorter_side=None, max_longer_side=1000,
         batch_size=1, total_num=None):
    assert cv2 is not None, '`load` requires `cv2`.'
    files = get_files(data_dir, data_name, total_num)
    total_num = len(files)

    for batch_start in range(0, total_num, batch_size):
        x = cv2.imread("%s/JPEGImages/%s.jpg" % (data_dir, files[batch_start]))
        if min_shorter_side is not None:
            scale = float(min_shorter_side) / np.min(x.shape[:2])
        else:
            scale = 1.0
        if round(scale * np.max(x.shape[:2])) > max_longer_side:
            scale = float(max_longer_side) / np.max(x.shape[:2])
        x = cv2.resize(x, None, None, fx=scale, fy=scale,
                       interpolation=cv2.INTER_LINEAR)
        x = np.array([x], dtype=np.float32)
        scale = np.array([scale], dtype=np.float32)
        yield x, scale
        del x

def evaluate_class(ids, scores, boxes, annotations, files, ovthresh):
    if scores.shape[0] == 0:
        return 0.0, np.zeros(len(ids)), np.zeros(len(ids))

    # extract gt objects for this class
    diff = [np.array([obj['difficult'] for obj in annotations[filename]])
            for filename in files]
    total = sum([sum(x == 0) for x in diff])
    detected = dict(zip(files, [[False] * len(x) for x in diff]))

    # sort by confidence
    sorted_ind = np.argsort(-scores)
    ids = ids[sorted_ind]
    boxes = boxes[sorted_ind, :]

    # go down dets and mark TPs and FPs
    tp_list = []
    fp_list = []
    for d in range(len(ids)):
        actual = np.array([x['bbox'] for x in annotations[ids[d]]])
        difficult = np.array([x['difficult'] for x in annotations[ids[d]]])

        if actual.size > 0:
            iw = np.maximum(np.minimum(actual[:, 2], boxes[d, 2]) -
                            np.maximum(actual[:, 0], boxes[d, 0]) + 1, 0)
            ih = np.maximum(np.minimum(actual[:, 3], boxes[d, 3]) -
                            np.maximum(actual[:, 1], boxes[d, 1]) + 1, 0)
            inters = iw * ih
            overlaps = inters / (area(actual) + area(boxes[d, :]) - inters)
            jmax = np.argmax(overlaps)
            ovmax = overlaps[jmax]
        else:
            ovmax = -np.inf

        tp = 0.
        fp = 0.
        if ovmax > ovthresh:
            if difficult[jmax] == 0:
                if not detected[ids[d]][jmax]:
                    tp = 1.
                    detected[ids[d]][jmax] = True
                else:
                    fp = 1.
        else:
            fp = 1.
        tp_list.append(tp)
        fp_list.append(fp)

    tp = np.cumsum(tp_list)
    fp = np.cumsum(fp_list)
    recall = tp / float(total)
    precision = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
    ap = np.mean([0 if np.sum(recall >= t) == 0
                  else np.max(precision[recall >= t])
                  for t in np.linspace(0, 1, 11)])

    return ap, precision, recall

def evaluate(results, data_dir, data_name, ovthresh=0.5, verbose=True):
    files = get_files(data_dir, data_name)
    files = files[:len(results)]
    annotations = get_annotations(data_dir, files)
    aps = []

    for c in range(used_classes):
        ids = []
        scores = []
        boxes = []
        for (i, filename) in enumerate(files):
            pred = results[i][c]
            if pred.shape[0] > 0:
                for k in xrange(pred.shape[0]):
                    ids.append(filename)
                    scores.append(pred[k, -1])
                    boxes.append(pred[k, :4] + 1)
        ids = np.array(ids)
        scores = np.array(scores)
        boxes = np.array(boxes)
        _annotations = dict((k, v[c]) for (k, v) in annotations.items())
        ap, _, _ = evaluate_class(ids, scores, boxes, _annotations,
                                  files, ovthresh)
        aps += [ap]

    strs = ''
    for c in range(used_classes):
        strs += "| %6s " % classnames[c][:6]
    strs += '|\n'

    for ap in aps:
        strs += '|--------'
    strs += '|\n'

    for ap in aps:
        strs += "| %.4f " % ap
    strs += '|\n'

    strs += "Mean = %.4f" % np.mean(aps)
    return strs,np.mean(aps)

def load_train(data_dir, data_name,
               batch_size=48, shuffle=True,
               target_size=416, anchors=5, classes=classes_num,
               total_num=None, dtype=np.float32):

    files = get_files(data_dir, data_name, total_num)
    annotations = get_annotations(data_dir, files)
    dirs = np.zeros(len(files), dtype=np.int)

    total_num = len(files)
    for f in files:
        annotations[f] = reduce(lambda x, y: x + y, annotations[f])

    if isinstance(target_size, int):
        target_size = (target_size, target_size)
    feature_size = [x // 32 for x in target_size]
    cells = feature_size[0] * feature_size[1]

    b = 0
    while True:
        if b == 0:
            if shuffle is True:
                idx = np.random.permutation(total_num)
            else:
                idx = np.arange(total_num)
        if b + batch_size > total_num:
            b = 0
            yield None, None
        else:
            batch_num = batch_size

        imgs = np.zeros((batch_num,) + target_size + (3,), dtype=dtype)
        probs = np.zeros((batch_num, cells, anchors, classes), dtype=dtype)
        confs = np.zeros((batch_num, cells, anchors), dtype=dtype)
        coord = np.zeros((batch_num, cells, anchors, 4), dtype=dtype)
        proid = np.zeros((batch_num, cells, anchors, classes), dtype=dtype)
        prear = np.zeros((batch_num, cells, 4), dtype=dtype)
        areas = np.zeros((batch_num, cells, anchors), dtype=dtype)
        upleft = np.zeros((batch_num, cells, anchors, 2), dtype=dtype)
        botright = np.zeros((batch_num, cells, anchors, 2), dtype=dtype)

        for i in range(batch_num):
            f = files[idx[b + i]]
            x = cv2.imread("%s/JPEGImages/%s.jpg" % (data_dir, f))
            h, w = x.shape[:2]
            cellx = 1. * w / feature_size[1]
            celly = 1. * h / feature_size[0]

            processed_objs = []
            for obj in annotations[f]:
                bbox = obj['bbox']
                centerx = .5 * (bbox[0] + bbox[2])  # xmin, xmax
                centery = .5 * (bbox[1] + bbox[3])  # ymin, ymax
                cx = centerx / cellx
                cy = centery / celly
                if cx >= feature_size[1] or cy >= feature_size[0]:
                    continue
                processed_objs += [[
                    classidx(obj['name']),
                    cx - np.floor(cx),  # centerx
                    cy - np.floor(cy),  # centery
                    np.sqrt(float(bbox[2] - bbox[0]) / w),
                    np.sqrt(float(bbox[3] - bbox[1]) / h),
                    int(np.floor(cy) * feature_size[1] + np.floor(cx))
                ]]

            # Calculate placeholders' values
            for obj in processed_objs:
                probs[i, obj[5], :, :] = [[0.] * classes] * anchors
                probs[i, obj[5], :, obj[0]] = 1.
                proid[i, obj[5], :, :] = [[1.] * classes] * anchors
                coord[i, obj[5], :, :] = [obj[1:5]] * anchors
                prear[i, obj[5], 0] = obj[1] - obj[3]**2 * .5 * feature_size[1]
                prear[i, obj[5], 1] = obj[2] - obj[4]**2 * .5 * feature_size[0]
                prear[i, obj[5], 2] = obj[1] + obj[3]**2 * .5 * feature_size[1]
                prear[i, obj[5], 3] = obj[2] + obj[4]**2 * .5 * feature_size[0]
                confs[i, obj[5], :] = [1.] * anchors

            # Finalise the placeholders' values
            ul = np.expand_dims(prear[i, :, 0:2], 1)
            br = np.expand_dims(prear[i, :, 2:4], 1)
            wh = br - ul
            area = wh[:, :, 0] * wh[:, :, 1]
            upleft[i, :, :, :] = np.concatenate([ul] * anchors, 1)
            botright[i, :, :, :] = np.concatenate([br] * anchors, 1)
            areas[i, :, :] = np.concatenate([area] * anchors, 1)

            imgs[i] = cv2.resize(x, target_size,
                                 interpolation=cv2.INTER_LINEAR)
        yield imgs, [probs, confs, coord, proid, areas, upleft, botright]
        b += batch_size

max_stars_repo_path: For custom data/yolf_data.py

max_issues_repo_path: For custom data/yolf_data.py

max_forks_repo_path: For custom data/yolf_data.py


text:
#!/usr/bin/env python
""" Fast tree and support functions for fast implementation of UniFrac """

from numpy import (logical_and, logical_or, sum, take, nonzero, repeat,
    array, concatenate, zeros, put, transpose, flatnonzero, newaxis,
    logical_xor, logical_not)
from numpy.random import permutation
from cogent.core.tree import PhyloNode

__author__ = "Rob Knight and Micah Hamady"
__copyright__ = "Copyright 2007-2012, The Cogent Project"
__credits__ = ["Rob Knight", "Micah Hamady", "Daniel McDonald"]
__license__ = "GPL"
__version__ = "1.5.3"
__maintainer__ = "Rob Knight, Micah Hamady"
__email__ = "rob@spot.colorado.edu, hamady@colorado.edu"
__status__ = "Prototype"

#bind reduce to local variables for speed in inner loops
lar = logical_and.reduce
lor = logical_or.reduce

class UniFracTreeNode(PhyloNode):
    """Slightly extended PhyloNode treenode for use with UniFrac

    Can expect Length and Name (= Name) to be set by DndParser.
    """
    def __nonzero__(self):
        """Returns True if self.Children."""
        return bool(self.Children)

def index_tree(t):
    """Returns tuple containing {node_id:node}, [node_id,first_child,last_child]

    Indexes nodes in-place as n._leaf_index.

    Algorithm is as follows:
    for each node in post-order traversal over tree:
        if the node has children:
            set an index on each child
            for each child with children:
                add the child and its start and end tips to the result
    """
    id_index = {}    #needs to be dict, not list, b/c adding out of order
    child_index = []
    curr_index = 0
    for n in t.traverse(self_before=False, self_after=True):
        for c in n.Children:
            c._leaf_index = curr_index
            id_index[curr_index] = c
            curr_index += 1
            if c:    #c has children itself, so need to add to result
                child_index.append((c._leaf_index, c.Children[0]._leaf_index,\
                    c.Children[-1]._leaf_index))
    #handle root, which should be t itself
    t._leaf_index = curr_index
    id_index[curr_index] = t
    #only want to add to the child_index if t has children...
    if t.Children:
        child_index.append((t._leaf_index, t.Children[0]._leaf_index,\
            t.Children[-1]._leaf_index))
    return id_index, child_index

def count_envs(lines, ignore_chars=0):
    """Reads env counts from lines. Returns dict of {name:{env:count}}.

    Assumes all name-env mappings are unique -- will overwrite counts with the
    last observed count rather than adding.
    """
    result = {}
    for line in lines:
        fields = line.split()
        #skip if we don't have at least label and field
        if len(fields) < 2:
            continue
        name, env = fields[:2]
        if ignore_chars:
            env = env[ignore_chars:]
        if len(fields) > 2:
            count = int(fields[2])
        else:
            count = 1
        if name not in result:
            result[name] = {}
        result[name][env] = count
    return result

def sum_env_dict(envs):
    """Sums counts from the data structure produced by count_envs."""
    return sum([sum(env.values()) for env in envs.values()])

def get_unique_envs(envs):
    """extract all unique envs from envs dict"""
    result = set()
    for v in envs.values():
        result.update(v.keys())
    #sort envs for convenience in testing and display
    return sorted(result), len(result)

def index_envs(env_counts, tree_index, array_constructor=int):
    """Returns array of taxon x env with counts of the taxon in each env.

    env_counts should be the output of count_envs(lines).
    tree_index should be the id_index of index_tree(t).
    array_constructor is int by default (may need to change to float later
        to handle microarray data).
    """
    num_nodes = len(tree_index)
    unique_envs, num_envs = get_unique_envs(env_counts)
    env_to_index = dict([(e, i) for i, e in enumerate(unique_envs)])
    result = zeros((num_nodes, num_envs), array_constructor)
    #figure out taxon label to index map
    node_to_index = {}
    for i, node in tree_index.items():
        if node.Name is not None:
            node_to_index[node.Name] = i
    #walk over env_counts, adding correct slots in array
    for name in env_counts:
        curr_row_index = node_to_index[name]
        for env, count in env_counts[name].items():
            result[curr_row_index, env_to_index[env]] = count
    #return all the data structures we created; will be useful for other tasks
    return result, unique_envs, env_to_index, node_to_index

def get_branch_lengths(tree_index):
    """Returns array of branch lengths, in tree index order."""
    result = zeros(len(tree_index), float)
    for i, node in tree_index.items():
        try:
            if node.Length is not None:
                result[i] = node.Length
        except AttributeError:
            pass
    return result

def bind_to_array(tree_index, a):
    """Binds tree_index to array a, returning result in list.

    Takes as input list of (node, first_child, last_child)
    returns list of (node_row, child_rows) such that node_row points to the
    row of a that corresponds to the current node, and child_rows points to the
    row or rows of a that correspond to the direct children of the current
    node.

    Order is assumed to be traversal order, i.e. for the typical case of
    postorder traversal iterating over the items in the result and consolidating
    each time should give the same result as postorder traversal of the original
    tree. Should also be able to modify for preorder traversal.
    """
    #note: range ends with end+1, not end, b/c end is included
    return [(a[node], a[start:end+1]) for node, start, end in tree_index]

def bind_to_parent_array(t, a):
    """Binds tree to array a, returning result in list.

    Takes as input tree t with _leaf_index set.

    Returns list of (node_row, parent_row such that node_row points to the
    row of a that corresponds to the current row, and parent_row points to
    the row of the parent.

    Order will be preorder traversal, i.e. for propagating attributes from
    the root to the tip.

    Typical usage of this function is to set up an array structure for many
    preorder traversals on the same tree, especially where you plan to change
    the data between traversals.
    """
    result = []
    for n in t.traverse(self_before=True, self_after=False):
        if n is not t:
            result.append([a[n._leaf_index], a[n.Parent._leaf_index]])
    return result

def _is_parent_empty(parent_children):
    """Returns True if the first element in a (parent,children) tuple is empty.

    This is used by delete_empty_parents to figure out which elements to filter
    out.
    """
    return bool(parent_children[0].sum())

def delete_empty_parents(bound_indices):
    """Deletes from list of (parent, children) bound indices empty parents.

    Expects as input the output of bind_to_array. Returns copy rather than
    acting in-place because deletions from long lists are expensive.

    This has the effect of pruning the tree, but by just skipping over the
    parents who have no children rather than altering memory. This is expected
    to be faster than trimming the array because it avoids copy operations.

    For pairwise environment bootstrapping or jackknifing, run this after
    running bool_descendants (or similar) to delete parents that only have
    offspring that are in other environments. Note that this does _not_
    collapse nodes so that you might have long "stalks" with many serial nodes.
    It might be worth testing whether collapsing these stalks provides time
    savings.
    """
    return filter(_is_parent_empty, bound_indices)

def traverse_reduce(bound_indices, f):
    """Applies a[i] = f(a[j:k]) over list of [(a[i], a[j:k])].

    If list is in traversal order, has same effect as consolidating the
    function over the tree, only much faster.

    Note that f(a[j:k]) must return an object that can be broadcast to the
    same shape as a[i], e.g. summing a 2D array to get a vector.
    """
    for i, s in bound_indices:
        i[:] = f(s, 0)

def bool_descendants(bound_indices):
    """For each internal node, sets col to True if any descendant is True."""
    traverse_reduce(bound_indices, lor)

def zero_branches_past_roots(bound_indices, sums):
    """Zeroes out internal nodes that are roots of each subtree."""
    for i, ignore in bound_indices:
        i *= (i != sums)

def sum_descendants(bound_indices):
    """For each internal node, sets col to sum of values in descendants."""
    traverse_reduce(bound_indices, sum)

class FitchCounterDense(object):
    """Returns parsimony result for set of child states, counting changes.

    WARNING: this version assumes that all tips are assigned to at least
    one env, and produces incorrect parsimony counts if this is not the
    case.
    """
    def __init__(self):
        """Returns new FitchCounter, with Changes = 0."""
        self.Changes = 0

    def __call__(self, a, ignored):
        """Returns intersection(a), or, if zero, union(a)."""
        result = lar(a)
        if not result.any():
            result = lor(a)
            self.Changes += 1
        return result

class FitchCounter(object):
    """Returns parsimony result for set of child states, counting changes.

    This version is slower but is robust to the case where some tips are
    missing envs.

    WARNING: logical_and.reduce(), if called on an empty array, will return
    all True if there are no values (I can only assume that it returns
    True because it is reporting that there are no values that return False:
    this isn't what I expected). Hence the code to explicitly trap this
    case based on the shape parameter.
    """
    def __init__(self):
        """Returns new FitchCounter, with Changes = 0."""
        self.Changes = 0

    def __call__(self, a, ignored):
        """Returns intersection(a), or, if zero, union(a)."""
        nonzero_rows = a[a.sum(1).nonzero()]
        if len(nonzero_rows):
            result = lar(nonzero_rows)
        else:
            result = zeros(nonzero_rows.shape[-1], bool)
        if not result.any():
            if nonzero_rows.any():
                result = lor(nonzero_rows)
                self.Changes += 1
        return result

def fitch_descendants(bound_indices, counter=FitchCounter):
    """Sets each internal node to Fitch parsimony assignment, returns # changes."""
    f = counter()
    traverse_reduce(bound_indices, f.__call__)
    return f.Changes

def tip_distances(a, bound_indices, tip_indices):
    """Sets each tip to its distance from the root."""
    for i, s in bound_indices:
        i += s
    mask = zeros(len(a))
    put(mask, tip_indices, 1)
    a *= mask[:,newaxis]

def permute_selected_rows(rows, orig, new, permutation_f=permutation):
    """Takes selected rows from orig, inserts into new in permuted order.

    NOTE: the traditional UniFrac permutation test, the P test, etc. shuffle
    the envs, i.e. they preserve all the correlations of seqs between envs.
    This function can also be used to shuffle each env (i.e. column)
    individually by applying it to column slices of orig and new. This
    latter method provides a potentially less biased but less conservative
    test.
    """
    shuffled = take(rows, permutation_f(len(rows)))
    for r, s in zip(rows, shuffled):
        new[s] = orig[r]

def prep_items_for_jackknife(col):
    """Takes column of a, returns vector with multicopy states unpacked.

    e.g. if index 3 has value 4, there will be 4 copies of index 3 in result.
    """
    nz = flatnonzero(col)
    result = [repeat(array((i,)), col[i]) for i in nz]
    return concatenate(result)

def jackknife_bool(orig_items, n, length, permutation_f=permutation):
    """Jackknifes vector of items so that only n remain.

    orig = flatnonzero(vec)
    length = len(vec)

    Returns all items if requested sample is larger than number of items.
    """
    permuted = take(orig_items, permutation_f(len(orig_items))[:n])
    result = zeros(length)
    put(result, permuted, 1)
    return result

def jackknife_int(orig_items, n, length, permutation_f=permutation):
    """Jackknifes new vector from vector of orig items.

    Returns all items if reqested sample is larger than number of items.
    """
    result = zeros(length)
    permuted = take(orig_items, permutation_f(len(orig_items))[:n])
    for p in permuted:
        result[p] += 1
    return result

def jackknife_array(mat, num_keep, axis=1, jackknife_f=jackknife_int,
    permutation_f=permutation):
    """ Jackknife array along specified axis, keeping specified num_keep"""
    cur_mat = mat
    if axis:
        cur_mat = mat.T
    num_r, num_c = cur_mat.shape
    jack_mat = zeros((num_r, num_c))
    for row_ix in range(num_r):
        in_prepped_array = prep_items_for_jackknife(cur_mat[row_ix,:])
        jack_mat[row_ix,:] = jackknife_f(orig_items=in_prepped_array,
            n=num_keep, length=num_c, permutation_f=permutation_f)
    if axis:
        jack_mat = jack_mat.T
    return jack_mat

def unifrac(branch_lengths, i, j):
    """Calculates unifrac(i,j) from branch lengths and cols i and j of m.

    This is the original, unweighted UniFrac metric.

    branch_lengths should be row vector, same length as # nodes in tree.

    i and j should be slices of states from m, same length as # nodes in
    tree. Slicing m (e.g. m[:,i]) returns a vector in the right format; note
    that it should be a row vector (the default), not a column vector.
    """
    return 1 - ((branch_lengths*logical_and(i,j)).sum()/\
        (branch_lengths*logical_or(i,j)).sum())

def unnormalized_unifrac(branch_lengths, i, j):
    """UniFrac, but omits normalization for frac of tree covered."""
    return (branch_lengths*logical_xor(i,j)).sum()/branch_lengths.sum()

def G(branch_lengths, i, j):
    """Calculates G(i,j) from branch lengths and cols i,j of m.

    This calculates fraction gain in branch length in i with respect to i+j,
    i.e. normalized for the parts of the tree that i and j cover.

    Note: G is the metric that we also call "asymmetric unifrac".
    """
    return (branch_lengths*logical_and(i, logical_not(j))).sum()/\
        (branch_lengths*logical_or(i,j)).sum()

def PD(branch_lengths, i):
    """Calculate PD(i) from branch lengths and col i of m.

    Calculates raw amount of branch length leading to tips in i, including
    branch length from the root.
    """
    return (branch_lengths * i.astype(bool)).sum()

def unnormalized_G(branch_lengths, i, j):
    """Calculates G(i,j) from branch length and cols i,j of m.

    This calculates the fraction gain in branch length of i with respect to j,
    divided by all the branch length in the tree.
    """
    return (branch_lengths*logical_and(i, logical_not(j))).sum()/\
        branch_lengths.sum()

def unifrac_matrix(branch_lengths, m, metric=unifrac, is_symmetric=True):
    """Calculates unifrac(i,j) for all i,j in m.

    branch_lengths is the array of branch lengths.

    m is 2D array: rows are taxa, states are columns. Assumes that ancestral
    states have already been calculated (either by logical_or or Fitch).

    metric: metric to use for combining each pair of columns i and j. Default
    is unifrac.

    is_symmetric indicates whether the metric is symmetric. Default is True.
    """
    num_cols = m.shape[-1]
    cols = [m[:,i] for i in range(num_cols)]
    result = zeros((num_cols,num_cols), float)
    if is_symmetric:
        #only calc half matrix and transpose
        for i in range(1, num_cols):
            first_col = cols[i]
            row_result = []
            for j in range(i):
                second_col = cols[j]
                row_result.append(metric(branch_lengths, first_col, second_col))
            result[i,:j+1] = row_result
        #note: can't use += because shared memory between a and transpose(a)
        result = result + transpose(result)
    else:
        #calc full matrix, incl. diagonal (which is probably 0...)
        for i in range(num_cols):
            first_col = cols[i]
            result[i] = [metric(branch_lengths, first_col, cols[j]) for \
                j in range(num_cols)]
    return result

def unifrac_one_sample(one_sample_idx, branch_lengths, m, metric=unifrac):
    """Calculates unifrac(one_sample_idx,j) for all environments j in m.

    branch_lengths is the array of branch lengths.

    m is 2D count array: rows are taxa (corresponding to branch_lenths),
    samples/states/envs are columns.
    Assumes that ancestral
    states have already been calculated (either by logical_or or Fitch).

    metric: metric to use for when comparing two environments. Default
    is unifrac. must be called like:
    metric(branch_lengths, env1_counts, env2counts)

    returns a numpy 1d array
    if asymmetric metric, returns metric(one_sample, other), usually a row in
    the mtx returned by unifrac_matrix
    """
    num_cols = m.shape[-1]
    cols = [m[:,i] for i in range(num_cols)]
    # result = zeros((num_cols), float)

    first_col = cols[one_sample_idx]
    # better to do loop into preallocated numpy array here?
    result = array([metric(branch_lengths, first_col, cols[j]) for \
        j in range(num_cols)],'float')
    return result

def env_unique_fraction(branch_lengths, m):
    """ Calculates unique branch length for each env.

    Returns unique branch len and unique fraction
    """
    total_bl = branch_lengths.sum()
    if total_bl <= 0:
        raise ValueError, "total branch length in tree must be > 0"

    n_rows_nodes, n_col_envs = m.shape
    env_bl_sums = zeros(n_col_envs)

    cols = [m[:, i] for i in range(n_col_envs)]
    col_sum = m.sum(1)
    env_bl_sums = zeros(n_col_envs)
    for env_ix, f in enumerate(cols):
        sing = (f == col_sum)
        # have  to mask zeros
        put(sing, nonzero(f == 0), 0)
        env_bl_sums[env_ix] = (sing * branch_lengths).sum()

    return env_bl_sums, env_bl_sums/total_bl

def unifrac_vector(branch_lengths, m, metric=unifrac):
    """Calculates unifrac(i, others) for each column i of m.

    Parameters as for unifrac_matrix. Use this when you want to calculate
    UniFrac or G of each state against the rest of the states, rather than
    of each state against each other state.
    """
    num_cols = m.shape[-1]
    cols = [m[:, i] for i in range(num_cols)]
    col_sum = m.sum(1)
    return array([metric(branch_lengths, col, col_sum-col) for col in cols])

def PD_vector(branch_lengths, m, metric=PD):
    """Calculates metric(i) for each column i of m.

    Parameters as for unifrac_matrix. Use this when you want to calculate
    PD or some other alpha diversity metric that depends solely on the branches
    within each state, rather than calculations that compare each state against
    each other state.
    """
    return array([metric(branch_lengths, col) for col in m.T])

def _weighted_unifrac(branch_lengths, i, j, i_sum, j_sum):
    """Calculates weighted unifrac(i,j) from branch lengths and cols i,j of m.
    """
    return (branch_lengths * abs((i/float(i_sum))-(j/float(j_sum)))).sum()

def _branch_correct(tip_distances, i, j, i_sum, j_sum):
    """Calculates weighted unifrac branch length correction.

    tip_distances  must be 0 except for tips.
    """
    result = tip_distances.ravel()*((i/float(i_sum))+(j/float(j_sum)))
    return result.sum()

def weighted_unifrac(branch_lengths, i, j, tip_indices, \
    unifrac_f=_weighted_unifrac):
    """Returns weighted unifrac(i,j) from branch lengths and cols i,j of m.

    Must pass in tip indices to calculate sums.

    Note: this calculation is not used in practice because it has to recalc.
    the sum each time. More efficient to calculate the sum first and pass it
    into _weighted_unifrac directly, as weighted_unifrac_matrix does.
    """
    i_sum = (take(i, tip_indices)).sum()
    j_sum = (take(j, tip_indices)).sum()
    return unifrac_f(branch_lengths, i, j, i_sum, j_sum)

def weighted_unifrac_matrix(branch_lengths, m, tip_indices, bl_correct=False,
        tip_distances=None, unifrac_f=_weighted_unifrac):
    """Calculates weighted_unifrac(i,j) for all i,j in m.

    Requires tip_indices for calculating sums, etc.
    bl_correct: if True (default: False), applies branch length correction.
        tip_distances is required for normalization for weighted unifrac.
    """
    num_cols = m.shape[-1]
    cols = [m[:,i] for i in range(num_cols)] #note that these will be row vecs
    sums = [take(m[:,i], tip_indices).sum() for i in range(num_cols)]
    result = zeros((num_cols,num_cols),float)
    for i in range(1, num_cols):
        i_sum = sums[i]
        first_col = cols[i]
        row_result = []
        for j in range(i):
            second_col = cols[j]
            j_sum = sums[j]
            curr = unifrac_f(branch_lengths, first_col, \
                second_col, i_sum, j_sum)
            if bl_correct:
                curr /= _branch_correct(tip_distances, first_col, \
                    second_col, i_sum, j_sum)
            row_result.append(curr)
        result[i,:j+1] = row_result
    result = result + transpose(result)
    return result

def weighted_one_sample(one_sample_idx, branch_lengths, m, tip_indices,
    bl_correct=False, tip_distances=None, unifrac_f=_weighted_unifrac):
    """Calculates weighted_unifrac(one_sample_idx,j) for all environments j in m

    Requires tip_indices for calculating sums, etc.
    bl_correct: if True (default: False), applies branch length correction.
        tip_distances is required for normalization for weighted unifrac.
    """
    num_cols = m.shape[-1]
    cols = [m[:,i] for i in range(num_cols)] #note that these will be row vecs
    sums = [take(m[:,i], tip_indices).sum() for i in range(num_cols)]
    result = zeros((num_cols),float)

    i_sum = sums[one_sample_idx]
    first_col = cols[one_sample_idx]
    row_result = []
    for j in range(num_cols):
        second_col = cols[j]
        j_sum = sums[j]
        curr = unifrac_f(branch_lengths, first_col, \
            second_col, i_sum, j_sum)
        if bl_correct:
            curr /= _branch_correct(tip_distances, first_col, \
                second_col, i_sum, j_sum)
        result[j] = curr
    return result

def weighted_unifrac_vector(branch_lengths, m, tip_indices, bl_correct=False,
    tip_distances=None, unifrac_f=_weighted_unifrac):
    """Calculates weighted_unifrac(i,rest) for i in m.

    Requires tip_indices for calculating sums, etc.
    bl_correct: if True (default: False), applies branch length correction.
        tip_distances is required for normalization for weighted unifrac.
    """
    num_cols = m.shape[-1]
    cols = [m[:,i] for i in range(num_cols)]
    sums = [take(m[:,i], tip_indices).sum() for i in range(num_cols)]
    sum_of_cols = m.sum(1)
    sum_of_sums = sum(sums)
    result = []
    for i, col in enumerate(cols):
        i_sum = sums[i]
        i_col = cols[i]
        rest_col = sum_of_cols - i_col
        rest_sum = sum_of_sums - i_sum
        curr = unifrac_f(branch_lengths, i_col,
            rest_col, i_sum, rest_sum)
        if bl_correct:
            curr /= _branch_correct(tip_distances, i_col,
            rest_col, i_sum, rest_sum)
        result.append(curr)
    return array(result)

max_stars_repo_name: sauloal/cnidaria

max_issues_repo_name: sauloal/cnidaria

max_forks_repo_name: sauloal/cnidaria

reason: from numpy