def temporal_smooth_smpls(ref_smpls, pose_fc=300, cam_fc=100):
"""
Args:
ref_smpls (np.ndarray): (length, 72)
pose_fc:
cam_fc:
Returns:
ref_smpls (np.ndarray): (length, 72)
"""
ref_rotvec = ref_smpls[:, 3:-10]
n = ref_rotvec.shape[0]
ref_rotvec = ref_rotvec.reshape((-1, 3))
ref_rotmat = R.from_rotvec(ref_rotvec).as_matrix()
ref_rotmat = torch.from_numpy(ref_rotmat)
ref_rot6d = rotations.rotmat_to_rot6d(ref_rotmat)
ref_rot6d = ref_rot6d.numpy()
ref_rot6d = ref_rot6d.reshape((n, -1))
ref_rot6d = get_smooth_params(ref_rot6d, fc=pose_fc)
ref_rot6d = ref_rot6d.reshape((-1, 6))
ref_rotmat = rotations.rot6d_to_rotmat(torch.from_numpy(ref_rot6d)).numpy()
ref_rotvec = R.from_matrix(ref_rotmat).as_rotvec()
ref_smpls[:, 3:-10] = ref_rotvec.reshape((n, -1))
ref_smpls[:, 0:3] = get_smooth_params(ref_smpls[:, 0:3], fc=cam_fc)
return ref_smpls
def set_pos_on_sphere_after_screen_rotate(self, mousedx, mousedy, screenw,
screenh):
theta = (math.pi / 5.) * math.sqrt((float(4 * mousedx * mousedx)) /
(screenw * screenw) +
(float(4 * mousedy * mousedy)) /
(screenh * screenh))
rotate_phi = 0.
if mousedx == 0:
if mousedy > 0:
rotate_phi = math.pi / 2.
else:
rotate_phi = -math.pi / 2.
else:
rotate_phi = math.atan2(mousedy, mousedx)
rot_vec = -np.dot(
self.T[:3, :3],
np.dot(
Rotation.from_rotvec(rotate_phi * np.array(
[0., 0., 1.])).as_dcm(), np.array([0., 1., 0.])))
self.vec_invalidate(self.T)
self.transforming_T[:] = self.T[:]
translate_temp = np.eye(4)
translate_temp[:3, 3] = -self.obj
self.transforming_T = np.dot(translate_temp, self.transforming_T)
rotate_temp = np.eye(4)
rotate_temp[:3, :3] = Rotation.from_rotvec(theta * rot_vec).as_dcm()
self.transforming_T = np.dot(rotate_temp, self.transforming_T)
self.transforming_T = np.dot(-translate_temp, self.transforming_T)
def samples_axis_rotation(pts, phi_r, delta_rs, snap=True):
"""
Samples poses with rotational offset in [-phi_r, phi_r] at a rate of [delta_r] per given axis. Optionally, the
translation to snap the given object (represented by [pts]) to the ground plane is computed.
:param pts: numpy array, surface points of the object.
:param phi_r: Rotational offset range.
:param delta_rs: Per axis, the rotational offset sampling rate.
:param snap: If true, determine translation such that object is snapped to ground plane.
:return: Rotation and translation of sampled poses.
"""
Rs = []
if delta_rs[0] > 0:
offset = np.arange(-phi_r, phi_r+1, delta_rs[0])
Rs += [Rotation.from_euler('x', roff, degrees=True).as_dcm() for roff in offset]
if delta_rs[1] > 0:
offset = np.arange(-phi_r, phi_r + 1, delta_rs[1])
Rs += [Rotation.from_euler('y', roff, degrees=True).as_dcm() for roff in offset]
if delta_rs[2] > 0:
offset = np.arange(-phi_r, phi_r + 1, delta_rs[2])
Rs += [Rotation.from_euler('z', roff, degrees=True).as_dcm() for roff in offset]
Rs = np.array(Rs).reshape(-1, 3, 3)
ts = np.array([[0, 0, -(R @ pts).min(axis=1)[2]] for R in Rs]) if snap else np.zeros((Rs.shape[0], 3))
return Rs, ts
def apply_gripper_tcp_offset(self, pose):
pose_trans = pose[:3]
pose_rot = Rotation.from_rotvec(pose[3:])
pose_tmat = Utils.trans_and_rot_to_tmat(pose_trans, pose_rot)
tcp_tmat = Utils.trans_and_rot_to_tmat([0, 0, -0.201], Rotation.from_rotvec([0, 0, 0]))
new_pose_tmat = pose_tmat @ tcp_tmat
new_pose_trans, new_pose_rot = Utils.tmat_to_trans_and_rot(new_pose_tmat)
new_pose_rotvec = new_pose_rot.as_rotvec()
return np.concatenate((new_pose_trans, new_pose_rotvec))
def rotate(self, direction, angle):
# rotate view around axis
if direction in ["RIGHT", "LEFT"]:
sign = 1 if direction == "RIGHT" else -1
rot = Rotation.from_rotvec([0, sign * angle, 0])
self.camera_rot = rot * self.camera_rot
else:
sign = 1 if direction == "UP" else -1
rot = Rotation.from_rotvec([sign * angle, 0, 0])
self.camera_rot = rot * self.camera_rot
def step(self, tau=1.0, start=None, goal=None, q=None):
# this pretty much mirrors original DMP, except it replaces transformation system
# with equivalent for quaternions
if goal is None:
g = self.g
else:
g = goal
if start is None:
q0 = self.q0
else:
q0 = start
# do closed-loop orientation control if current orientation given
if q is not None:
# ensure that quats are within 90 degrees of each other
# (so error term doesn't suddenly reverse direction)
if np.sum(self.q * q) >= 0:
self.q = q
else:
self.q = -q
theta = self.cs.step(tau)
psi = self.psi(theta)
ws = np.array([self.wx_dmp.w, self.wy_dmp.w, self.wz_dmp.w])
# note: forcing term excludes error dependence in this dmp
f = np.sum(ws * psi, axis=1, keepdims=True) * theta / np.sum(
psi, axis=1, keepdims=True)
# transformation system with quaternion error
self.omegad = - np.dot(self.K, q_err(g, self.q)) \
- np.dot(self.D, self.omega) \
+ np.dot(self.K, q_err(g, q0)) * theta \
+ np.dot(self.K, f)
self.omegad /= tau
self.omega += self.omegad * self.dt
# quaternion integration rule
wd, xd, yd, zd = 0.5 * np.dot(
np.block([[0, -self.omega.T], [self.omega, -q_cross(self.omega)]]),
np.concatenate(([self.q[3]], self.q[:3])))
self.qd = np.array([xd, yd, zd, wd]).reshape(-1, 1) / tau
# integrate off of omega (easier integration formula, more accurate than linear approx)
rq = R.from_quat(self.q.flatten()) * R.from_rotvec(
(self.omega * self.dt).flatten())
self.q = R.as_quat(rq).reshape(-1, 1)
return self.q, self.qd, self.omega, self.omegad
def apply_transform_moveit_to_real(pose):
pose_local = pose.copy()
tvec = pose_local[:3]
orientation = Rotation.from_rotvec(pose_local[3:])
pose_tmat = Utils.trans_and_rot_to_tmat(tvec, orientation)
rot_z_1 = Rotation.from_euler("xyz", [0, 0, -np.pi])
rot_z_2 = Rotation.from_euler("xyz", [0, 0, -np.pi / 2])
rot_y = Rotation.from_euler("xyz", [0, np.pi / 2, 0])
tmat_z_1 = Utils.trans_and_rot_to_tmat([0, 0, 0], rot_z_1)
tmat_z_2 = Utils.trans_and_rot_to_tmat([0, 0, 0], rot_z_2)
tmat_y = Utils.trans_and_rot_to_tmat([0, 0, 0], rot_y)
applied_tmat = tmat_z_1 @ pose_tmat @ tmat_y @ tmat_z_2
applied_trans, applied_rot = Utils.tmat_to_trans_and_rot(applied_tmat)
applied_rvect = applied_rot.as_rotvec()
return np.concatenate((applied_trans, applied_rvect))
def convert_to_ply():
# Get the subdirectories
subdirs = sorted(os.listdir(os.path.join(YCB_MODELS_DIR, 'models')))
for s in subdirs:
obj_mesh = read_obj(os.path.join(YCB_MODELS_DIR, 'models', s, 'textured_simple.obj'))
# From OpenGL coordinates
from scipy.spatial.transform.rotation import Rotation
rotx = np.eye(4)
rotx[:3, :3] = Rotation.from_euler('x', 180, degrees=True).as_dcm()
# pose = rotx @ pose
pose = np.eye(4)
pose = np.matmul(rotx, pose)
o3d_mesh = o3d.geometry.TriangleMesh()
if hasattr(obj_mesh, 'r'):
o3d_mesh.vertices = o3d.utility.Vector3dVector(np.copy(obj_mesh.r))
elif hasattr(obj_mesh, 'v'):
o3d_mesh.vertices = o3d.utility.Vector3dVector(np.copy(obj_mesh.v))
else:
raise Exception('Unknown Mesh format')
o3d_mesh.triangles = o3d.utility.Vector3iVector(np.copy(obj_mesh.f))
# Save as .ply
filename = os.path.join(YCB_MODELS_DIR, 'models', s, 'mesh.ply')
o3d.io.write_triangle_mesh(filename, o3d_mesh)
# Export with trimesh
mesh = trimesh.load(filename)
mesh.export(filename)
def run(self, input_cloud: DopplerPointCloud,
imu_in: ImuVelocityData) -> DopplerPointCloud:
"""Given the current Doppler point cloud with IMU data, this function will return a time adjusted point cloud.
This point cloud will include previously entered point_clouds, which will only be as accurate as the data provided.
Accuracy can be improved by modifying the allowable persistable steps.
Args:
input_cloud (DopplerPointCloud): [description]
imu_in (ImuData): [description]
Returns:
DopplerPointCloud: [description]
"""
t_delta: float = self._get_time_delta()
mv = imu_in.get_dxdydz()
# Simple state estimation based on imu
meters: tuple[float, float, float] = (mv[0] * t_delta, mv[1] * t_delta,
mv[2] * t_delta)
rot: Rotation = Rotation.from_euler(
'zyx', [x * t_delta
for x in imu_in.get_dyawdpitchdroll()]) # type: ignore
ret = DopplerPointCloud(input_cloud.get().copy())
for i in self._pts:
i.translate_rotate(meters, rot)
ret.append(i)
if len(self._pts) > self._steps:
self._pts.popleft()
self._pts.append(input_cloud)
return ret
def movel(self, pose, acceleration=1.0, velocity=0.2, use_mm=False, max_retries=3):
pose_local = pose.copy()
if use_mm:
pose_local[0] *= 0.001
pose_local[1] *= 0.001
pose_local[2] *= 0.001
pose_local = self.apply_gripper_tcp_offset(pose_local)
pose_local = apply_transform_real_to_moveit(pose_local)
print("moving to " + str(pose_local))
quaternion = Rotation.from_rotvec(pose_local[3:]).as_quat()
goal = MoveRobotGoal()
goal.action = "cartesian"
goal.cartesian_goal.position.x = pose_local[0]
goal.cartesian_goal.position.y = pose_local[1]
goal.cartesian_goal.position.z = pose_local[2]
goal.cartesian_goal.orientation.x = quaternion[0]
goal.cartesian_goal.orientation.y = quaternion[1]
goal.cartesian_goal.orientation.z = quaternion[2]
goal.cartesian_goal.orientation.w = quaternion[3]
retry_amount = 0
success = False
while retry_amount < max_retries:
self.client.send_goal(goal)
self.client.wait_for_result()
result = self.client.get_result()
success = result.success
if success:
break
retry_amount += 1
return success
def transform_quaternion(
self,
quaternion: Quaternion,
from_: Literal["robot", "asset"],
to_: Literal["robot", "asset"],
) -> Quaternion:
quaternion_from: Rotation = Rotation.from_quat(
quaternion.as_np_array())
if not isinstance(quaternion, Quaternion):
raise ValueError("Incorrect input format. Must be Quaternion.")
if from_ == self.transform.to_ and to_ == self.transform.from_:
"""Using the inverse transform"""
quaternion_to = quaternion_from * self.transform.rotation_object.inv(
)
elif from_ == self.transform.from_ and to_ == self.transform.to_:
quaternion_to = quaternion_from * self.transform.rotation_object
else:
raise ValueError("Transform not specified")
result: Quaternion = Quaternion.from_array(quaternion_to.as_quat(),
frame=to_)
return result
def mat_to_quaternion(R):
q1, q2, q3, q0 = Rotation.from_matrix(R).as_quat()
# q0 = np.sqrt(np.trace(R) + 1) / 2
# q1 = (R[2, 1] - R[1, 2]) / q0 / 4
# q2 = (R[0, 2] - R[2, 0]) / q0 / 4
# q3 = (R[1, 0] - R[0, 1]) / q0 / 4
return Quarternion((q0, q1, q2, q3))
def test_rmsd_rotated_and_translated(self):
s1 = self.get_test_structure()
s2 = self.get_test_structure()
# rotate and translate s2
AXIS_DIRECTION = np.array([11, 2, 0])
AXIS_DIRECTION = AXIS_DIRECTION / np.linalg.norm(
AXIS_DIRECTION) # following code expects a unit vector
ANGLE = np.pi / 4
TRANSLATION = np.array([1, 200, 7])
atoms = list(s2.get_atoms())
coords = np.array([a.get_coord() for a in atoms])
rotation = Rotation.from_rotvec(ANGLE * AXIS_DIRECTION)
for atom, new_coord in zip(atoms,
rotation.apply(coords) + TRANSLATION):
atom.set_coord(new_coord)
# test rmsd is still 0
chain_a = ChainResidues.from_bio_chain(s1[0]['A'])
chain_a_rotated = ChainResidues.from_bio_chain(s2[0]['A'])
get_c_alpha_coords = GetCAlphaCoords()
get_centroid = GetCentroid((get_c_alpha_coords, ))
get_centered_c_alpha_coords = GetCenteredCAlphaCoords(
(get_c_alpha_coords, get_centroid))
get_rmsd = GetRMSD((get_centered_c_alpha_coords,
GetRotationMatrix(
(get_centered_c_alpha_coords, ))))
rmsd = get_rmsd(chain_a, chain_a_rotated)
self.assertAlmostEqual(0, rmsd, places=4)
def plot_box(pd, pos, quat, size):
d = -size
p = size
X = np.array([[d[0], d[0], p[0], p[0], d[0], d[0], p[0], p[0]],
[d[1], p[1], p[1], d[1], d[1], p[1], p[1], d[1]],
[d[2], d[2], d[2], d[2], p[2], p[2], p[2], p[2]]])
R = Rotation.from_quat(quat)
X = R.apply(X.T) + pos
pd.append(
go.Mesh3d(
# 8 vertices of a cube
x=X[:, 0],
y=X[:, 1],
z=X[:, 2],
flatshading=True,
lighting=dict(facenormalsepsilon=0),
lightposition=dict(x=2000, y=1000),
color='yellow',
opacity=0.2,
# i, j and k give the vertices of triangles
i=[7, 0, 0, 0, 4, 4, 6, 6, 4, 0, 3, 2],
j=[3, 4, 1, 2, 5, 6, 5, 2, 0, 1, 6, 3],
k=[0, 7, 2, 3, 6, 7, 1, 1, 5, 5, 7, 6],
name='y',
showscale=False))
def load_object_and_pose(base_dir, seq_name, frame_id):
# Read annotation file
anno = read_annotation(base_dir, seq_name, frame_id, data_split)
# Make the 4x4 transformation matrix
pose = np.eye(4)
pose[:3, :3] = cv2.Rodrigues(anno['objRot'])[0]
pose[:3, 3] = anno['objTrans']
# Load the object model
# load object model
obj_mesh = read_obj(os.path.join(YCB_MODELS_DIR, 'models', anno['objName'], 'textured_simple.obj'))
## apply current pose to the object model
#objMesh.v = np.matmul(objMesh.v, cv2.Rodrigues(anno['objRot'])[0].T) + anno['objTrans']
# From OpenGL coordinates
from scipy.spatial.transform.rotation import Rotation
rotx = np.eye(4)
rotx[:3, :3] = Rotation.from_euler('x', 180, degrees=True).as_dcm()
# pose = rotx @ pose
pose = np.matmul(rotx, pose)
o3d_mesh = o3d.geometry.TriangleMesh()
if hasattr(obj_mesh, 'r'):
o3d_mesh.vertices = o3d.utility.Vector3dVector(np.copy(obj_mesh.r))
elif hasattr(obj_mesh, 'v'):
o3d_mesh.vertices = o3d.utility.Vector3dVector(np.copy(obj_mesh.v))
else:
raise Exception('Unknown Mesh format')
o3d_mesh.triangles = o3d.utility.Vector3iVector(np.copy(obj_mesh.f))
return anno['objName'], o3d_mesh, pose, anno['camMat']
def _make_camera_transform(rot, trans):
rot = np.array(rot)[[1, 2, 3, 0]]
rotation_matrix = Rotation.from_quat(rot).as_matrix()
transformation_matrix = np.eye(4)
transformation_matrix[:3, :3] = rotation_matrix
transformation_matrix[:3, -1] = trans
return torch.from_numpy(transformation_matrix).float()
def __init__(self, net, fov=90, **kwargs):
self.net = net
self.terminate = False
self.marker = o3d.geometry.TriangleMesh.create_octahedron(
self.MARK_SIZE)
self.marker.paint_uniform_color(np.r_[0, 1, 0.1])
self.vis = o3d.visualization.VisualizerWithKeyCallback()
self.vis.create_window("Perceptron Error Space",
width=self.WIDTH,
height=self.HEIGHT)
self.vis.register_key_callback(ord(" "), self.esc_cb)
ctr = self.vis.get_view_control()
ctr.change_field_of_view(step=fov)
self.vis.add_geometry(self.marker)
self.vis.get_render_option().show_coordinate_frame = True
self.vis.get_render_option().line_width = 10.0
self.arrows = [
o3d.geometry.TriangleMesh.create_arrow(self.MARK_SIZE * 0.6,
self.MARK_SIZE * 1,
self.MARK_SIZE * 4,
self.MARK_SIZE * 2, 10, 2)
for i in range(6)
]
rotations = [
[0, 90, 0],
[-90, 0, 0],
[0, 0, 0],
[0, -90, 0],
[90, 0, 0],
[-180, 0, 0],
]
for i, arrow in enumerate(self.arrows):
c = [0] * 3
c[i % 3] = 1
arrow.paint_uniform_color(np.r_[c])
arrow.rotate(
Rotation.from_euler("xyz", rotations[i],
degrees=True).as_dcm())
self.vis.add_geometry(arrow)
points = np.r_["0,2,1", [0, 0, 0], [1, 0, 0], [0, 1, 0],
[0, 0, 1], ] * self.RANGE * 1.5
lines = [
[0, 1],
[0, 2],
[0, 3],
]
colors = np.diag(np.ones(3))
axis = o3d.geometry.LineSet(
points=o3d.utility.Vector3dVector(points),
lines=o3d.utility.Vector2iVector(lines),
)
axis.colors = o3d.utility.Vector3dVector(colors)
self.vis.add_geometry(axis)
self._init_error_space(**kwargs)
def _bullet_to_world(self, position, orientation, id):
in_world = np.eye(4)
in_world[0:3, 0:3] = Rotation.from_quat(orientation).as_dcm()
# position was offset by center of mass (origin of collider in physics world)
offset = in_world[0:3, 0:3] @ np.array(
self.dataset.obj_coms[self.dataset.obj_ids.index(id)])
in_world[0:3, 3] = position - offset
return in_world
def __init__(self, altitude: float, dxdydz: tuple[float, float, float],
drolldpitchdyaw: tuple[float, float,
float], heading: float) -> None:
self._altitude: float = altitude
self._dxdydz: tuple[float, float, float] = dxdydz
self._drolldpitchdyaw: tuple[float, float, float] = drolldpitchdyaw
self._heading: float = heading
self._rot: Rotation = Rotation.from_euler(
'zyx', [self._drolldpitchdyaw]) #type: ignore
def plot_cube_at_pos_orientation(
pos=(0, 0, 0), size=(1, 1, 1), orientation=(0, 0, 0), ax=None,
**kwargs):
"""
Plot a cube/cuboid at any place, with any size and with any orientation
:param pos:
:param size:
:param orientation: as rotation matrix or euler angles (XYZ) in degrees
:param ax:
:param kwargs:
:return:
"""
if ax != None:
X, Y, Z = cuboid_data(pos, size)
# rotate coordinates
if len(orientation) == 9:
R = Rotation.from_dcm(orientation)
elif len(orientation) == 3:
R = Rotation.from_euler('XYZ', orientation, degrees=True)
else:
raise ValueError
X = np.array(X)
Y = np.array(Y)
Z = np.array(Z)
points = np.hstack((X.reshape(-1, 1), Y.reshape(-1,
1), Z.reshape(-1, 1)))
# shift to origin
points -= np.array(pos)
# rotate
points = R.apply(points)
# shift back
points += np.array(pos)
X = points[:, 0].reshape(4, 5)
Y = points[:, 1].reshape(4, 5)
Z = points[:, 2].reshape(4, 5)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, **kwargs)
def __init__(self,
initial_position=None,
initial_rotation=None,
colors=None):
self.position = np.zeros(
3) if initial_position is None else initial_position
self.rotation: Rotation = Rotation.identity(
) if initial_rotation is None else initial_rotation
self.colors = self.COLORS if colors is None else colors
self.update_verticies()
def _world_to_bullet(self, in_world, id):
R = in_world[0:3, 0:3]
orientation = Rotation.from_dcm(R).as_quat()
# position is offset by center of mass (origin of collider in physics world)
offset = R @ np.array(
self.dataset.obj_coms[self.dataset.obj_ids.index(id)])
position = in_world[0:3, 3] + offset
return position, orientation
def get_joint_pos(puck_position, rotation, psi):
ee_rotation = R.from_euler('zyx', [rotation[2], rotation[1], rotation[0]],
degrees=True)
ee_position = puck_position + np.array([0, 0, 0.1915])
ee_pose = np.concatenate(
[ee_position,
ee_rotation.as_quat()[3:],
ee_rotation.as_quat()[:3]])
res, q = kinematics.inverse_kinematics(torch.from_numpy(ee_pose),
torch.tensor(psi).double(),
[1, -1, 1])
return res, q.detach().numpy()
def process_scene(scene_no, input_directory, output_folder):
image_filenames = sorted((input_directory / 'imgs' / "scene_" + scene_no).files("*color*.png"))
depth_filenames = sorted((input_directory / 'imgs' / "scene_" + scene_no).files("*depth*.png"))
extrinsics = np.loadtxt(input_directory / 'pc' / scene_no + ".pose")
K = np.array([[570.3, 0.0, 320.0],
[0.0, 570.3, 240.0],
[0.0, 0.0, 1.0]])
homogen = np.zeros((1, 4))
homogen[0, 3] = 1
poses = []
for extrinsic in extrinsics:
w = extrinsic[0]
xyz = extrinsic[1:4]
rot = Rotation.from_quat(np.hstack((xyz, w))).as_matrix()
tra = np.reshape(extrinsic[4:], (3, 1))
extrinsic = np.hstack((rot, tra))
extrinsic = np.vstack((extrinsic, homogen))
poses.append(extrinsic)
current_output_dir = output_folder / "scene_" + scene_no
if os.path.isdir(current_output_dir):
shutil.rmtree(current_output_dir)
os.mkdir(current_output_dir)
os.mkdir(os.path.join(current_output_dir, "images"))
os.mkdir(os.path.join(current_output_dir, "depth"))
output_poses = []
for current_index in range(len(image_filenames)):
image = cv2.imread(image_filenames[current_index])
depth = cv2.imread(depth_filenames[current_index], cv2.IMREAD_ANYDEPTH)
depth = np.float32(depth)
depth = depth / 10000.0
depth[depth > 50.0] = 0.0
depth[np.isnan(depth)] = 0.0
depth[np.isinf(depth)] = 0.0
depth = (depth * 1000.0).astype(np.uint16)
output_poses.append(poses[current_index].ravel().tolist())
cv2.imwrite("{}/images/{}.png".format(current_output_dir, str(current_index).zfill(6)), image, [cv2.IMWRITE_PNG_COMPRESSION, 3])
cv2.imwrite("{}/depth/{}.png".format(current_output_dir, str(current_index).zfill(6)), depth, [cv2.IMWRITE_PNG_COMPRESSION, 3])
output_poses = np.array(output_poses)
np.savetxt("{}/poses.txt".format(current_output_dir), output_poses)
np.savetxt("{}/K.txt".format(current_output_dir), K)
return scene_no
def test_get_hinge_angle(self):
s1 = self.get_test_structure()
s2 = self.get_test_structure()
s2.id = f'{s1.id}_with_rotated_chain'
s1d1 = ChainResidues.from_bio_chain(s1[0]['A'])
s1d2 = ChainResidues.from_bio_chain(s1[0]['B'])
s2d1 = ChainResidues.from_bio_chain(s2[0]['A'])
s2d2 = ChainResidues.from_bio_chain(s2[0]['B'])
# rotate second domain over a defined screw axis, then check if GetHingeAngle indeed computes the correct parameters (angle, translation)
# define the screw axis
AXIS_DIRECTION = np.array([1, 2, 0])
AXIS_DIRECTION = AXIS_DIRECTION / np.linalg.norm(
AXIS_DIRECTION) # following code expects a unit vector
AXIS_LOCATION = np.array([
52.71183395385742, 44.92530822753906, -11.425999641418457
]) # a random pivot point (the axis goes through it)
ANGLE = np.pi / 4
TRANSLATION_IN_AXIS = 3
# move along the screw axis using scipy
s2d2_atoms = [atom for residue in s2d2 for atom in residue]
s2d2_atom_coords = np.array([atom.coord for atom in s2d2_atoms])
rotation = Rotation.from_rotvec(ANGLE * AXIS_DIRECTION)
rotated_s2d2_atom_coords = rotation.apply(
s2d2_atom_coords - AXIS_LOCATION) + AXIS_LOCATION
for atom, new_coord in zip(
s2d2_atoms, rotated_s2d2_atom_coords +
AXIS_DIRECTION * TRANSLATION_IN_AXIS):
atom.set_coord(new_coord)
# compute the hinge angle with GetHingeAngle
get_c_alpha_coords = GetCAlphaCoords()
get_centroid = GetCentroid((get_c_alpha_coords, ))
get_centered_c_alpha_coords = GetCenteredCAlphaCoords(
(get_c_alpha_coords, get_centroid))
get_hinge_angle = GetHingeAngle((get_c_alpha_coords, get_centroid,
GetRotationMatrix(
(get_centered_c_alpha_coords, ))))
screw_motion = get_hinge_angle(s1d1, s1d2, s2d1, s2d2)
self.assertAlmostEqual(ANGLE, screw_motion.angle, places=3)
self.assertAlmostEqual(TRANSLATION_IN_AXIS,
screw_motion.translation_in_axis,
places=3)
def translate_rotate(self, location: tuple[float, float, float],
pitch_rads: Rotation):
"""Translates and rotates the underlying object. This is done in-place, no further verification is done.
Args:
location (Tuple[float, float, float]): Tuple of float values to shift the underlying data with: (x meters, y meters, z meters)
pitch_rads (Rotation): Rotation matrix object from scipy.spatial.transform.rotation.Rotation
"""
if self._data.shape[0]:
self._data[:, 0] += location[0]
self._data[:, 1] += location[1]
self._data[:, 2] += location[2]
self._data[:, :3] = pitch_rads.apply(
self._data[:, :3]) #type: ignore
def __init__(self):
self.T = np.eye(4)
self.T[:3, 3] = np.array([0., 3., 4.])
self.T[:3, :3] = Rotation.from_rotvec(-math.atan(0.75) *
np.array([1., 0., 0.])).as_dcm()
self.transforming_T = self.T.copy()
self.distance = 5.
self.transforming = False
self.eye = np.zeros(3)
self.view = np.zeros(3)
self.obj = np.zeros(3)
self.up = np.zeros(3)
def euler_to_quaternion(euler: Euler,
sequence: str = "ZYX",
degrees: bool = False) -> Quaternion:
"""
Transform a quaternion into Euler angles.
:param euler: An Euler object.
:param sequence: Rotation sequence for the Euler angles.
:param degrees: Set to true if the provided Euler angles are in degrees. Default is radians.
:return: Quaternion object.
"""
rotation_object: Rotation = Rotation.from_euler(seq=sequence,
angles=euler.as_np_array(),
degrees=degrees)
quaternion: Quaternion = Quaternion.from_array(rotation_object.as_quat(),
frame="robot")
return quaternion
def move(self, imu_vel: ImuVelocityData, time_passed: float):
"""Based on some IMU velocity information and the amount of time which has passed, modify the current pose.
Args:
imu_vel (ImuVelocityData): Input IMU data.
time_passed (float): Time in seconds.
"""
rot_vec: np.ndarray = (Rotation.from_euler(
'zyx', [self._roll, self._pitch, self._yaw])).apply(
imu_vel.get_dxdydz()) * time_passed
self._x += rot_vec[0]
self._y += rot_vec[1]
self._z += rot_vec[2]
self._yaw += imu_vel.get_drolldpitchdyaw()[0] * time_passed
self._pitch += imu_vel.get_drolldpitchdyaw()[1] * time_passed
self._roll += imu_vel.get_drolldpitchdyaw()[2] * time_passed
def __probability_ellipsoid(self, Uxx, Uyy, Uzz, Uyz, Uxz, Uxy, p=0.99):
U = np.array([[Uxx,Uxy,Uxz],
[Uxy,Uyy,Uyz],
[Uxz,Uyz,Uzz]])
w, v = np.linalg.eig(np.linalg.inv(U))
r_eff = chi2.ppf(1-p, 3)
radii = np.sqrt(r_eff/w)
rot = Rotation.from_matrix(v)
euler_angles = rot.as_euler('ZXY', degrees=True)
return radii, euler_angles
