<commit_before>import os
import json
from collections import OrderedDict
import numpy as np
import scipy.interpolate

from director import ikconstraints
from director import ikconstraintencoder
from director import ikparameters
from director import plannerPublisher
from director import transformUtils
from director import roboturdf
from director.fieldcontainer import FieldContainer
from director.simpletimer import FPSCounter
from director import vtkAll as vtk
from director import transformUtils

from director import lcmUtils

import pydrake
import pydrake.solvers.ik as pydrakeik

try:
    from robotlocomotion import robot_plan_t
except ImportError:
    from drc import robot_plan_t


class PyDrakePlannerPublisher(plannerPublisher.PlannerPublisher):

    def _setup(self):

        self.counter = FPSCounter()
        self.counter.printToConsole = True

        # disabled for now
        #self._setupLocalServer()

    def _setupLocalServer(self):

        initArgs = FieldContainer(
            urdfFile=self.ikPlanner.robotModel.getProperty('Filename'),
            packagePaths=roboturdf.getPackagePaths()
            )

        self.ikServer = PyDrakeIkServer()
        self.ikServer.initInstance(initArgs)

    def testEncodeDecode(self, fields):

        encoded = json.dumps(fields, cls=ikconstraintencoder.ConstraintEncoder)
        decoded = json.loads(encoded, object_hook=ikconstraintencoder.ConstraintDecoder)

        del decoded['class']
        fields = FieldContainer(**decoded)

        del fields.options['class']
        fields.options = ikparameters.IkParameters(**fields.options)

        constraints = []

        for c in fields.constraints:
            objClass = getattr(ikconstraints, c['class'])
            del c['class']
            obj = objClass()
            constraints.append(obj)

            for attr, value in c.iteritems():
                if isinstance(value, dict) and 'position' in value and 'quaternion' in value:
                    value = transformUtils.transformFromPose(value['position'], value['quaternion'])
                setattr(obj, attr, value)

        fields.constraints = constraints

        return fields

    def makePlanMessage(self, poses, poseTimes, info, fields):

        from director import robotstate

        # rescale poseTimes
        poseTimes = np.array(poseTimes)*2.0

        states = [robotstate.drakePoseToRobotState(pose) for pose in poses]
        for i, state in enumerate(states):
            state.utime = poseTimes[i]*1e6

        msg = robot_plan_t()
        msg.utime = fields.utime
        msg.robot_name = 'robot'
        msg.num_states = len(states)
        msg.plan = states
        msg.plan_info = [info]*len(states)
        msg.num_bytes = 0
        msg.matlab_data = []
        msg.num_grasp_transitions = 0
        msg.left_arm_control_type = msg.NONE
        msg.right_arm_control_type = msg.NONE
        msg.left_leg_control_type = msg.NONE
        msg.right_leg_control_type = msg.NONE

        return msg

    def processIK(self, constraints, ikParameters, positionCosts, nominalPoseName="", seedPoseName=""):

        fields = self.setupFields(constraints, ikParameters, positionCosts, nominalPoseName, seedPoseName)
        fields = self.testEncodeDecode(fields)
        endPose, info = self.ikServer.runIk(fields)
        self.counter.tick()
        return endPose, info

    def processTraj(self, constraints, ikParameters, positionCosts, nominalPoseName="", seedPoseName="", endPoseName=""):

        fields = self.setupFields(constraints, ikParameters, positionCosts, nominalPoseName, seedPoseName, endPoseName)
        fields = self.testEncodeDecode(fields)
        poses, poseTimes, info = self.ikServer.runIkTraj(fields)
        plan = self.makePlanMessage(poses, poseTimes, info, fields)
        lcmUtils.publish('CANDIDATE_MANIP_PLAN', plan)
        return plan, info


class RigidBodyTreeCompatNew(object):

    @staticmethod
    def get_body(rbt, i):
        return rbt.get_body(i)

    @staticmethod
    def get_num_bodies(rbt):
        return rbt.get_num_bodies()

    @staticmethod
    def get_body_name(rbt, i):
        return rbt.get_body(i).get_name()

    @staticmethod
    def get_position_name(rbt, i):
        return rbt.get_position_name(i)

    @staticmethod
    def get_num_positions(rbt):
        return rbt.get_num_positions()

    @staticmethod
    def get_num_velocities(rbt):
        return rbt.get_num_velocities()


class RigidBodyTreeCompatOld(object):

    @staticmethod
    def get_body(rbt, i):
        return rbt.bodies[i]

    @staticmethod
    def get_num_bodies(rbt):
        return len(rbt.bodies)

    @staticmethod
    def get_body_name(rbt, i):
        return rbt.bodies[i].linkname

    @staticmethod
    def get_position_name(rbt, i):
        return rbt.getPositionName(i)

    @staticmethod
    def get_num_positions(rbt):
        return rbt.num_positions

    @staticmethod
    def get_num_velocities(rbt):
        return rbt.num_velocities


if hasattr(pydrake.rbtree.RigidBodyTree, 'get_num_bodies'):
    rbt = RigidBodyTreeCompatNew
else:
    rbt = RigidBodyTreeCompatOld


class PyDrakeIkServer(object):

    def __init__(self):
        self.rigidBodyTree = None
        self.trajInterpolationMode = 'cubic'

    def initInstance(self, fields):

        self.rigidBodyTree = self.loadRigidBodyTree(fields.urdfFile, fields.packagePaths)

        self.bodyNames = [rbt.get_body_name(self.rigidBodyTree, i) for i in xrange(rbt.get_num_bodies(self.rigidBodyTree))]
        self.bodyNameToId = {}
        for i, name in enumerate(self.bodyNames):
            self.bodyNameToId[name] = i

        self.positionNames = [rbt.get_position_name(self.rigidBodyTree, i) for i in xrange(rbt.get_num_positions(self.rigidBodyTree))]
        self.positionNameToId = {}
        for i, name in enumerate(self.positionNames):
            self.positionNameToId[name] = i

    def loadRigidBodyTree(self, urdfFile, packagePaths):

        assert os.path.isfile(urdfFile)

        packageMap = pydrake.rbtree.mapStringString()
        for path in packagePaths:
            packageMap[os.path.basename(path)] = path

        urdfString = open(urdfFile, 'r').read()
        baseDir = str(os.path.dirname(urdfFile))
        floatingBaseType = pydrake.rbtree.kRollPitchYaw
        weldFrame = None

        rigidBodyTree = pydrake.rbtree.RigidBodyTree()

        pydrake.rbtree.AddModelInstanceFromUrdfStringSearchingInRosPackages(
          urdfString,
          packageMap,
          baseDir,
          floatingBaseType,
          weldFrame,
          rigidBodyTree)

        return rigidBodyTree

    def makeIkOptions(self, fields):

        options = pydrakeik.IKoptions(self.rigidBodyTree)
        options.setQ(np.diag(fields.positionCosts))
        #options.setQv()
        #options.setQa()
        #options.setDebug(True)

        options.setMajorOptimalityTolerance(fields.options.majorOptimalityTolerance)
        options.setMajorFeasibilityTolerance(fields.options.majorFeasibilityTolerance)
        options.setMajorIterationsLimit(fields.options.majorIterationsLimit)
        #options.setIterationsLimit(limit)
        #options.setSuperbasicsLimit(limit)<commit_msg>Updated pydrakeik.py to use new PackageMap class.<commit_after>import os
import json
from collections import OrderedDict
import numpy as np
import scipy.interpolate

from director import ikconstraints
from director import ikconstraintencoder
from director import ikparameters
from director import plannerPublisher
from director import transformUtils
from director import roboturdf
from director.fieldcontainer import FieldContainer
from director.simpletimer import FPSCounter
from director import vtkAll as vtk
from director import transformUtils

from director import lcmUtils

import pydrake
import pydrake.solvers.ik as pydrakeik

try:
    from robotlocomotion import robot_plan_t
except ImportError:
    from drc import robot_plan_t


class PyDrakePlannerPublisher(plannerPublisher.PlannerPublisher):

    def _setup(self):

        self.counter = FPSCounter()
        self.counter.printToConsole = True

        # disabled for now
        #self._setupLocalServer()

    def _setupLocalServer(self):

        initArgs = FieldContainer(
            urdfFile=self.ikPlanner.robotModel.getProperty('Filename'),
            packagePaths=roboturdf.getPackagePaths()
            )

        self.ikServer = PyDrakeIkServer()
        self.ikServer.initInstance(initArgs)

    def testEncodeDecode(self, fields):

        encoded = json.dumps(fields, cls=ikconstraintencoder.ConstraintEncoder)
        decoded = json.loads(encoded, object_hook=ikconstraintencoder.ConstraintDecoder)

        del decoded['class']
        fields = FieldContainer(**decoded)

        del fields.options['class']
        fields.options = ikparameters.IkParameters(**fields.options)

        constraints = []

        for c in fields.constraints:
            objClass = getattr(ikconstraints, c['class'])
            del c['class']
            obj = objClass()
            constraints.append(obj)

            for attr, value in c.iteritems():
                if isinstance(value, dict) and 'position' in value and 'quaternion' in value:
                    value = transformUtils.transformFromPose(value['position'], value['quaternion'])
                setattr(obj, attr, value)

        fields.constraints = constraints

        return fields

    def makePlanMessage(self, poses, poseTimes, info, fields):

        from director import robotstate

        # rescale poseTimes
        poseTimes = np.array(poseTimes)*2.0

        states = [robotstate.drakePoseToRobotState(pose) for pose in poses]
        for i, state in enumerate(states):
            state.utime = poseTimes[i]*1e6

        msg = robot_plan_t()
        msg.utime = fields.utime
        msg.robot_name = 'robot'
        msg.num_states = len(states)
        msg.plan = states
        msg.plan_info = [info]*len(states)
        msg.num_bytes = 0
        msg.matlab_data = []
        msg.num_grasp_transitions = 0
        msg.left_arm_control_type = msg.NONE
        msg.right_arm_control_type = msg.NONE
        msg.left_leg_control_type = msg.NONE
        msg.right_leg_control_type = msg.NONE

        return msg

    def processIK(self, constraints, ikParameters, positionCosts, nominalPoseName="", seedPoseName=""):

        fields = self.setupFields(constraints, ikParameters, positionCosts, nominalPoseName, seedPoseName)
        fields = self.testEncodeDecode(fields)
        endPose, info = self.ikServer.runIk(fields)
        self.counter.tick()
        return endPose, info

    def processTraj(self, constraints, ikParameters, positionCosts, nominalPoseName="", seedPoseName="", endPoseName=""):

        fields = self.setupFields(constraints, ikParameters, positionCosts, nominalPoseName, seedPoseName, endPoseName)
        fields = self.testEncodeDecode(fields)
        poses, poseTimes, info = self.ikServer.runIkTraj(fields)
        plan = self.makePlanMessage(poses, poseTimes, info, fields)
        lcmUtils.publish('CANDIDATE_MANIP_PLAN', plan)
        return plan, info


class RigidBodyTreeCompatNew(object):

    @staticmethod
    def get_body(rbt, i):
        return rbt.get_body(i)

    @staticmethod
    def get_num_bodies(rbt):
        return rbt.get_num_bodies()

    @staticmethod
    def get_body_name(rbt, i):
        return rbt.get_body(i).get_name()

    @staticmethod
    def get_position_name(rbt, i):
        return rbt.get_position_name(i)

    @staticmethod
    def get_num_positions(rbt):
        return rbt.get_num_positions()

    @staticmethod
    def get_num_velocities(rbt):
        return rbt.get_num_velocities()


class RigidBodyTreeCompatOld(object):

    @staticmethod
    def get_body(rbt, i):
        return rbt.bodies[i]

    @staticmethod
    def get_num_bodies(rbt):
        return len(rbt.bodies)

    @staticmethod
    def get_body_name(rbt, i):
        return rbt.bodies[i].linkname

    @staticmethod
    def get_position_name(rbt, i):
        return rbt.getPositionName(i)

    @staticmethod
    def get_num_positions(rbt):
        return rbt.num_positions

    @staticmethod
    def get_num_velocities(rbt):
        return rbt.num_velocities


if hasattr(pydrake.rbtree.RigidBodyTree, 'get_num_bodies'):
    rbt = RigidBodyTreeCompatNew
else:
    rbt = RigidBodyTreeCompatOld


class PyDrakeIkServer(object):

    def __init__(self):
        self.rigidBodyTree = None
        self.trajInterpolationMode = 'cubic'

    def initInstance(self, fields):

        self.rigidBodyTree = self.loadRigidBodyTree(fields.urdfFile, fields.packagePaths)

        self.bodyNames = [rbt.get_body_name(self.rigidBodyTree, i) for i in xrange(rbt.get_num_bodies(self.rigidBodyTree))]
        self.bodyNameToId = {}
        for i, name in enumerate(self.bodyNames):
            self.bodyNameToId[name] = i

        self.positionNames = [rbt.get_position_name(self.rigidBodyTree, i) for i in xrange(rbt.get_num_positions(self.rigidBodyTree))]
        self.positionNameToId = {}
        for i, name in enumerate(self.positionNames):
            self.positionNameToId[name] = i

    def loadRigidBodyTree(self, urdfFile, packagePaths):

        assert os.path.isfile(urdfFile)

        packageMap = pydrake.rbtree.PackageMap()
        for path in packagePaths:
            packageMap.Add(os.path.basename(path), path)

        urdfString = open(urdfFile, 'r').read()
        baseDir = str(os.path.dirname(urdfFile))
        floatingBaseType = pydrake.rbtree.kRollPitchYaw
        weldFrame = None

        rigidBodyTree = pydrake.rbtree.RigidBodyTree()

        pydrake.rbtree.AddModelInstanceFromUrdfStringSearchingInRosPackages(
          urdfString,
          packageMap,
          baseDir,
          floatingBaseType,
          weldFrame,
          rigidBodyTree)

        return rigidBodyTree

    def makeIkOptions(self, fields):

        options = pydrakeik.IKoptions(self.rigidBodyTree)
        options.setQ(np.diag(fields.positionCosts))
        #options.setQv()
        #options.setQa()
        #options.setDebug(True)

        options.setMajorOptimalityTolerance(fields.options.majorOptimalityTolerance)
        options.setMajorFeasibilityTolerance(fields.options.majorFeasibilityTolerance)
        options.setMajorIterationsLimit(fields.options.majorIterationsLimit)
        #options.setIterationsLimit(limit)
        #options.setSuperbasicsLimit(limit)
{'fraction_non_alphanumeric': 0.08405977584059776, 'fraction_numerical': 0.0008302200083022001, 'mean_word_length': 2.9941972920696327, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6866236', 'n_tokens_mistral': 4538, 'n_tokens_neox': 4318, 'n_words': 867}