Refactor to speed up

[iory]

from datetime import datetime

import numpy as np
from skrobot.coordinates.math import matrix2quaternion


class RobotDataset(object):

    def __init__(self, robot_model):
        self.robot_model = robot_model

    def sample(self, n):
        robot_model = self.robot_model
        ndof = len(robot_model.joint_max_angles)

        j_max = robot_model.joint_max_angles.copy()
        j_min = robot_model.joint_min_angles.copy()
        j_max[j_max == np.inf] = 2 * np.pi
        j_min[j_min == -np.inf] = - 2 * np.pi
        return np.random.uniform(low=j_min,
                                 high=j_max,
                                 size=(n, ndof))

    def end_coords_from_samples(self, samples):
        robot_model = self.robot_model
        end_coords_list = []
        from tqdm import tqdm
        for s in tqdm(samples):
            robot_model.angle_vector(s)
            end_coords = robot_model.end_coords
            t_xyz = end_coords.worldpos()
            q_wxyz = matrix2quaternion(end_coords.worldrot())
            end_coords_list.append(
                np.concatenate(
                    [t_xyz, q_wxyz]))
        return np.array(end_coords_list)


if __name__ == '__main__':
    from skrobot.models import PR2
    robot_model = PR2()
    robot_dataset = RobotDataset(robot_model.rarm)

    n = 100000
    start = datetime.now()
    samples = robot_dataset.sample(n)
    end_coords_array = robot_dataset.end_coords_from_samples(samples)
    end = datetime.now()
    print(end - start)

0:00:10.482203 -> 0:00:06.810289