<commit_before>import math
import numpy as np

from eigen3 import toNumpy, toEigen, Vector3d
import spacevecalg as sva

from mayavi import mlab
from tvtk.api import tvtk

import geometry



def setTransform(actor, transform):
  R = transform.rotation()
  T = transform.translation()

  actor.user_transform.set_matrix((R[0], R[1], R[2], T[0],
                                   R[3], R[4], R[5], T[1],
                                   R[6], R[7], R[8], T[2],
                                     0.,   0.,   0.,   1.))


  #actor.user_transform.set_matrix((R[0], R[3], R[6], T[0],
  #                                 R[1], R[4], R[7], T[1],
  #                                 R[2], R[5], R[8], T[2],
  #                                   0.,   0.,   0.,   1.))




class GraphicMultiBody:
  def __init__(self, mb, geom=None, staticTransform=sva.PTransform.Identity()):
    self.viewer = mlab.gcf()

    if geom == None:
      self.geom = geometry.DefaultGeometry(mb)
    else:
      self.geom = geom

    self.staticTransform = staticTransform

    self.transformVel = [(0., 0., 0., 0.)]*mb.nrJoints()

    actors = []
    for i in range(mb.nrBodies()):
      actors.append(self.geom.body(i)[0])
      if mb.parent(i) != -1:
        actors.append(self.geom.link(i)[0])
        actors.append(self.geom.joint(i)[0])

    self.viewer.scene.add_actors(actors)


    # todo : Replace by fixture
    # create velocity arrows
    self.velA = []
    for i in range(mb.nrBodies()):
      arrs = tvtk.ArrowSource()
      arrm = tvtk.PolyDataMapper(input=arrs.output)
      arra = tvtk.Actor(mapper=arrm)
      arra.visibility = 0
      self.IMat = arra.matrix
      self.viewer.scene.add_actors(arra)

      self.velA.append(arra)


  def render(self):
    self.viewer.scene.render()



  def draw(self, mb, mbc):
    bG = list(mbc.bodyPosW)

    for i in range(mb.nrBodies()):
      Xi = bG[i]

      #com = toNumpy(mb.body(i).inertia().momentum())/mb.body(i).inertia().mass()
      #com = toEigen(com)
      com = Vector3d.Zero()

      comPos = sva.PTransform(com)*Xi*self.staticTransform

      comA, comS = self.geom.body(i)
      setTransform(comA, comPos)

      if mb.parent(i) != -1:
        Xp = bG[mb.parent(i)]

        # show link from body base to joint
        linkA, linkS = self.geom.link(i)
        setTransform(linkA, linkS*Xp*self.staticTransform)

        # show joint
        jointA, jointS = self.geom.joint(i)
        setTransform(jointA, jointS*Xi*self.staticTransform)



  def drawVel(self, mb, mbc):
    bG = list(mbc.bodyPosW)
    bV = list(mbc.bodyVelW)

    for i in range(mb.nrBodies()):
      a = self.velA[i]

      XiT = bG[i].translation()
      ViL = bV[i].linear()

      arrAxe = np.mat([1., 0., 0.]).T
      nViL = toNumpy(ViL)
      speed = np.linalg.norm(ViL)
<commit_msg>Display MultiBody CoM and Body CoM.<commit_after>import math
import numpy as np

from eigen3 import toNumpy, toEigen, Vector3d
import spacevecalg as sva
import rbdyn as rbd

from mayavi import mlab
from tvtk.api import tvtk

import geometry



def setTransform(actor, transform):
  R = transform.rotation()
  T = transform.translation()

  actor.user_transform.set_matrix((R[0], R[1], R[2], T[0],
                                   R[3], R[4], R[5], T[1],
                                   R[6], R[7], R[8], T[2],
                                     0.,   0.,   0.,   1.))


  #actor.user_transform.set_matrix((R[0], R[3], R[6], T[0],
  #                                 R[1], R[4], R[7], T[1],
  #                                 R[2], R[5], R[8], T[2],
  #                                   0.,   0.,   0.,   1.))




class GraphicMultiBody:
  def __init__(self, mb, geom=None, staticTransform=sva.PTransform.Identity()):
    self.viewer = mlab.gcf()

    if geom == None:
      self.geom = geometry.DefaultGeometry(mb)
    else:
      self.geom = geom

    self.staticTransform = staticTransform

    self.transformVel = [(0., 0., 0., 0.)]*mb.nrJoints()

    ss = tvtk.SphereSource(radius=0.05)
    sm = tvtk.PolyDataMapper(input=ss.output)
    sa = tvtk.Actor(mapper=sm)
    sa.property.color = (1., 0., 0.)
    sa.user_transform = tvtk.Transform()
    sa.visibility = 1
    self.globalCoM = sa

    self.com = []
    actors = [sa]
    for i in range(mb.nrBodies()):
      ss = tvtk.SphereSource(radius=0.02*mb.body(i).inertia().mass())
      sm = tvtk.PolyDataMapper(input=ss.output)
      sa = tvtk.Actor(mapper=sm)
      sa.user_transform = tvtk.Transform()
      sa.property.color = (0., 1., 0.)
      sa.visibility = 1

      self.com.append(sa)
      actors.append(sa)

      actors.append(self.geom.body(i)[0])
      if mb.parent(i) != -1:
        actors.append(self.geom.link(i)[0])
        actors.append(self.geom.joint(i)[0])

    self.viewer.scene.add_actors(actors)


    # todo : Replace by fixture
    # create velocity arrows
    self.velA = []
    for i in range(mb.nrBodies()):
      arrs = tvtk.ArrowSource()
      arrm = tvtk.PolyDataMapper(input=arrs.output)
      arra = tvtk.Actor(mapper=arrm)
      arra.visibility = 0
      self.IMat = arra.matrix
      self.viewer.scene.add_actors(arra)

      self.velA.append(arra)


  def render(self):
    self.viewer.scene.render()



  def draw(self, mb, mbc):
    bG = list(mbc.bodyPosW)

    tCoM = sva.PTransform(rbd.computeCoM(mb, mbc))
    setTransform(self.globalCoM, tCoM)

    for i in range(mb.nrBodies()):
      Xi = bG[i]

      com = toNumpy(mb.body(i).inertia().momentum())/mb.body(i).inertia().mass()
      comPos = sva.PTransform(toEigen(com))*Xi*self.staticTransform
      setTransform(self.com[i], comPos)

      bodyA, bodyS = self.geom.body(i)
      bodyPos = Xi*self.staticTransform
      setTransform(bodyA, bodyPos)


      if mb.parent(i) != -1:
        Xp = bG[mb.parent(i)]

        # show link from body base to joint
        linkA, linkS = self.geom.link(i)
        setTransform(linkA, linkS*Xp*self.staticTransform)

        # show joint
        jointA, jointS = self.geom.joint(i)
        setTransform(jointA, jointS*Xi*self.staticTransform)



  def drawVel(self, mb, mbc):
    bG = list(mbc.bodyPosW)
    bV = list(mbc.bodyVelW)

    for i in range(mb.nrBodies()):
      a = self.velA[i]

      XiT = bG[i].translation()
      ViL = bV[i].linear()

      arrAxe = np.mat([1., 0., 0.]).T
      nViL = toNumpy(ViL)
      speed = np.linalg.norm(ViL)

{'fraction_non_alphanumeric': 0.15040906230333542, 'fraction_numerical': 0.017149150409062303, 'mean_word_length': 2.1346153846153846, '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': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6314007', 'n_tokens_mistral': 2451, 'n_tokens_neox': 2308, 'n_words': 454}