def inv_kin(self, target=None):
if target is None:
target = self.target
a = target[0]**2 + target[1]**2 - self.l[0]**2 - self.l[1]**2
b = 2 * self.l[0] * self.l[1]
q2 = np.arccos(a / b)
c = np.arctan2(target[1], target[0])
q1 = c - np.arctan2(self.l[1] * np.sin(q2),
(self.l[0] + self.l[1] * np.cos(q2)))
return np.array([q1, q2])
def debug_random_angles():
from jax import random
import jax.numpy as jnp
import pylab as plt
S = 10000
p1 = random.uniform(random.PRNGKey(0),shape=(S,2))
p1 /= jnp.linalg.norm(p1, axis=-1, keepdims=True)
plt.hist(jnp.arctan2(p1[:,1],p1[:,0]), bins='auto', alpha=0.5)
p2 = random.normal(random.PRNGKey(0), shape=(S, 2))
p2 /= jnp.linalg.norm(p2, axis=-1, keepdims=True)
plt.hist(jnp.arctan2(p2[:,1],p2[:,0]), bins='auto', alpha=0.5)
plt.show()
def quaternion_to_angle_axis(quaternion: np.ndarray) -> np.ndarray:
"""Convert quaternion vector to angle axis of rotation.
The quaternion should be in (w, x, y, z) format.
Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
Args:
quaternion (torch.Tensor): tensor with quaternions.
Return:
torch.Tensor: tensor with angle axis of rotation.
Shape:
- Input: :math:`(*, 4)` where `*` means, any number of dimensions
- Output: :math:`(*, 3)`
Example:
>>> quaternion = torch.rand(2, 4) # Nx4
>>> angle_axis = kornia.quaternion_to_angle_axis(quaternion) # Nx3
"""
if not isinstance(quaternion, np.ndarray):
raise TypeError("Input type is not a np.ndarray. Got {}".format(
type(quaternion)))
if not quaternion.shape[-1] == 4:
raise ValueError(
"Input must be a tensor of shape Nx4 or 4. Got {}".format(
quaternion.shape))
# unpack input and compute conversion
q1: np.ndarray = quaternion[..., 1]
q2: np.ndarray = quaternion[..., 2]
q3: np.ndarray = quaternion[..., 3]
sin_squared_theta: np.ndarray = q1 * q1 + q2 * q2 + q3 * q3
sin_theta: np.ndarray = np.sqrt(sin_squared_theta)
cos_theta: np.ndarray = quaternion[..., 0]
two_theta: np.ndarray = 2.0 * np.where(
cos_theta < 0.0,
np.arctan2(-sin_theta, -cos_theta),
np.arctan2(sin_theta, cos_theta),
)
k_pos: np.ndarray = two_theta / sin_theta
k_neg: np.ndarray = 2.0 * np.ones_like(sin_theta)
k: np.ndarray = np.where(sin_squared_theta > 0.0, k_pos, k_neg)
angle_axis = np.concatenate((np.expand_dims(
q1 * k, -1), np.expand_dims(q2 * k, -1), np.expand_dims(q3 * k, -1)),
axis=-1)
return angle_axis
def sph2latlon(xsph: jnp.ndarray) -> Tuple[jnp.ndarray]:
"""Converts a point on the unit sphere into latitude and longitude measured in
radians.
Args:
xsph: Observations on the sphere.
Returns:
lat, lon: A tuple containing the latitude and longitude coordinates of
the input.
"""
x, y, z = xsph[..., 0], xsph[..., 1], xsph[..., 2]
lat = jnp.arctan2(z, jnp.sqrt(jnp.square(x) + jnp.square(y)))
lon = jnp.arctan2(y, x)
return lat, lon
def log(self: "SO3") -> jnp.ndarray:
# Reference:
# > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/so3.hpp#L247
w = self.wxyz[..., 0]
norm_sq = self.wxyz[..., 1:] @ self.wxyz[..., 1:]
use_taylor = norm_sq < get_epsilon(norm_sq.dtype)
# Shim to avoid NaNs in jnp.where branches, which cause failures for
# reverse-mode AD.
norm_safe = jnp.sqrt(
jnp.where(
use_taylor,
1.0, # Any non-zero value should do here.
norm_sq,
))
atan_n_over_w = jnp.arctan2(
jnp.where(w < 0, -norm_safe, norm_safe),
jnp.abs(w),
)
atan_factor = jnp.where(
use_taylor,
2.0 / w - 2.0 / 3.0 * norm_sq / w**3,
jnp.where(
jnp.abs(w) < get_epsilon(w.dtype),
jnp.where(w > 0, 1.0, -1.0) * jnp.pi / norm_safe,
2.0 * atan_n_over_w / norm_safe,
),
)
return atan_factor * self.wxyz[1:]
def get_roll_pitch_jax(y_in):
sq1, sq2, sq3, sq4, sq5, sq6, sq7, cq1, cq2, cq3, cq4, cq5, cq6, cq7 = get_sin_cos_jax(y_in)
r_32 = -cq7*(cq5*sq2*sq3 - sq5*(cq2*sq4 - cq3*cq4*sq2)) - sq7*(cq6*(cq5*(cq2*sq4 - cq3*cq4*sq2) + sq2*sq3*sq5) + sq6*(cq2*cq4 + cq3*sq2*sq4))
r_33 = -cq6*(cq2*cq4 + cq3*sq2*sq4) + sq6*(cq5*(cq2*sq4 - cq3*cq4*sq2) + sq2*sq3*sq5)
r_31 = cq7*(cq6*(cq5*(cq2*sq4 - cq3*cq4*sq2) + sq2*sq3*sq5) + sq6*(cq2*cq4 + cq3*sq2*sq4)) - sq7*(cq5*sq2*sq3 - sq5*(cq2*sq4 - cq3*cq4*sq2))
return jnp.arctan2(r_32, r_33), -jnp.arcsin(r_31)
def render(self, mode="human"):
if self.viewer is None:
from gym.envs.classic_control import rendering
self.viewer = rendering.Viewer(1000, 1000)
self.viewer.set_bounds(-2.5, 2.5, -2.5, 2.5)
fname = path.dirname(__file__) + "/classic/assets/lds_arrow.png"
self.img = rendering.Image(fname, 0.35, 0.35)
self.img.set_color(1.0, 1.0, 1.0)
self.imgtrans = rendering.Transform()
self.img.add_attr(self.imgtrans)
fnamewind = path.dirname(__file__) + "/classic/assets/lds_grid.png"
self.imgwind = rendering.Image(fnamewind, 5.0, 5.0)
self.imgwind.set_color(0.5, 0.5, 0.5)
self.imgtranswind = rendering.Transform()
self.imgwind.add_attr(self.imgtranswind)
self.viewer.add_onetime(self.imgwind)
self.viewer.add_onetime(self.img)
cur_x, cur_y = self.imgtrans.translation
# new_x, new_y = self.state[0], self.state[1]
new_x, new_y = jnp.dot(self.state.squeeze(),
self.proj_vector1), jnp.dot(
self.state.squeeze(), self.proj_vector2)
diff_x = new_x - cur_x
diff_y = new_y - cur_y
new_rotation = jnp.arctan2(diff_y, diff_x)
self.imgtrans.set_translation(new_x, new_y)
self.imgtrans.set_rotation(new_rotation)
return self.viewer.render(return_rgb_array=mode == "rgb_array")
def nngp_fn(cov12, var1, var2):
if 'Identity' in name:
res = cov12
elif 'Erf' in name:
prod = (1 + 2 * var1) * (1 + 2 * var2)
res = np.arcsin(2 * cov12 / np.sqrt(prod)) * 2 / np.pi
elif 'Sin' in name:
sum_ = (var1 + var2)
s1 = np.exp((-0.5 * sum_ + cov12))
s2 = np.exp((-0.5 * sum_ - cov12))
res = (s1 - s2) / 2
elif 'Relu' in name:
prod = var1 * var2
sqrt = np.sqrt(np.maximum(prod - cov12 ** 2, 1e-30))
angles = np.arctan2(sqrt, cov12)
dot_sigma = (1 - angles / np.pi) / 2
res = sqrt / (2 * np.pi) + dot_sigma * cov12
else:
raise NotImplementedError(name)
return res
def _dihedral(pos: jnp.ndarray, indices: jnp.ndarray,
tvecs: jnp.ndarray) -> float:
dx1 = pos[indices[1]] - pos[indices[0]] + tvecs[0]
dx2 = pos[indices[2]] - pos[indices[1]] + tvecs[1]
dx3 = pos[indices[3]] - pos[indices[2]] + tvecs[2]
numer = dx2 @ jnp.cross(jnp.cross(dx1, dx2), jnp.cross(dx2, dx3))
denom = jnp.linalg.norm(dx2) * jnp.cross(dx1, dx2) @ jnp.cross(dx2, dx3)
return jnp.arctan2(numer, denom)
def nngp_fn_diag(nngp: np.ndarray) -> np.ndarray:
square_root = np.sqrt(1. + 2. * nngp)
new_nngp = nngp / ((nngp + 1.) * np.sqrt(1. + 2. * nngp))
new_nngp += np.arctan2(nngp, square_root) / 2
new_nngp /= np.pi
new_nngp += 0.25
new_nngp *= nngp
return new_nngp
def calc_torsions(xyz: np.ndarray, struct_descr_dict):
b1 = xyz[struct_descr_dict['torsions'][:, 1]] - xyz[struct_descr_dict['torsions'][:, 0]]
b2 = xyz[struct_descr_dict['torsions'][:, 2]] - xyz[struct_descr_dict['torsions'][:, 1]]
b3 = xyz[struct_descr_dict['torsions'][:, 3]] - xyz[struct_descr_dict['torsions'][:, 2]]
b12 = np.cross(b1, b2)
b23 = np.cross(b2, b3)
return np.arctan2((np.cross(b12, b23) * b2).sum(axis=-1) / np.linalg.norm(b2, axis=-1), (b12 * b23).sum(axis=-1))
def angle(p1, p2, p3):
"""
Returns the angle defined by three points in space
(around the one in the middle).
"""
q = p1 - p2
r = p3 - p2
return np.arctan2(linalg.norm(np.cross(q, r)), np.dot(q, r))
def dihedral_angle(p1, p2, p3, p4):
"""
Returns the dihedral angle defined by four points in space
(around the line defined by the two central points).
"""
q = p3 - p2
r = np.cross(p2 - p1, q)
s = np.cross(q, p4 - p3)
return np.arctan2(np.dot(np.cross(r, s), q), np.dot(r, s) * linalg.norm(q))
def to_kappa_angle(self) -> Tuple[RealArray, RealArray]:
if self.mean_times_concentration.shape[-1] != 2:
raise ValueError
kappa = jnp.linalg.norm(self.mean_times_concentration, axis=-1)
angle = jnp.where(kappa == 0.0,
0.0,
jnp.arctan2(self.mean_times_concentration[..., 1],
self.mean_times_concentration[..., 0]))
return kappa, angle
def compute_yaw_radians(self) -> jnp.ndarray:
"""Compute yaw angle. Uses the ZYX mobile robot convention.
Returns:
Euler angle in radians.
"""
# https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Quaternion_to_Euler_angles_conversion
q0, q1, q2, q3 = self.wxyz
return jnp.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2**2 + q3**2))
def compute_roll_radians(self) -> hints.ScalarJax:
"""Compute roll angle.
Uses the ZYX mobile robot convention.
Returns:
Euler angle in radians.
"""
# https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Quaternion_to_Euler_angles_conversion
q0, q1, q2, q3 = self.wxyz
return jnp.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1**2 + q2**2))
def calc_tray_pose(q):
"""Calculate tray pose in terms of x1, y1, y2."""
x1, y1, y2 = q
Δy = y2 - y1
Δx = jnp.sqrt((2 * TRAY_W)**2 - Δy**2)
θ_tw = jnp.arctan2(Δy, Δx)
R_wt = jax_rotation_matrix(θ_tw)
r_c1w_w = jnp.array([x1, y1])
r_tw_w = r_c1w_w + R_wt @ r_tc1_t
return jnp.array([r_tw_w[0], r_tw_w[1], θ_tw])
def __call__(self, state, action):
"""__call__.
Args:
state:
action:
Returns:
"""
sin, cos, _ = state.arr
action = self.max_torque * jnp.tanh(action[0])
newthdot = jnp.arctan2(sin, cos) + (
-3.0 * self.g /
(2.0 * self.l) * jnp.sin(jnp.arctan2(sin, cos) + jnp.pi) + 3.0 /
(self.m * self.l**2) * action)
newth = jnp.arctan2(sin, cos) + newthdot * self.dt
newsin, newcos = jnp.sin(newth), jnp.cos(newth)
arr = jnp.array([newsin, newcos, newthdot])
return PendulumState(arr=arr, h=state.h + 1), arr
def _observation_function(x, s1, s2):
"""
Returns the observed angles as function of the state and the sensors locations
Parameters
----------
x: array_like
The current state
s1: array_like
The first sensor location
s2: array_like
The second sensor location
Returns
-------
y: array_like
The observed angles, the first component is the angle w.r.t. the first sensor, the second w.r.t the second.
"""
return jnp.array([jnp.arctan2(x[1] - s1[1], x[0] - s1[0]),
jnp.arctan2(x[1] - s2[1], x[0] - s2[0])])
def euclid2ang(x):
"""Converts points on the circle embedded in the plane into an angular
representation in polar coordinates. Note that the output is shifted so
that angles lie in the interval [0, 2pi).
Args:
x: Observations on the circle.
Returns:
out: Angular representation of the input.
"""
return jnp.arctan2(x[..., 1], x[..., 0]) + jnp.pi
def nngp_ntk_fn(
nngp: np.ndarray,
prod: np.ndarray,
ntk: Optional[np.ndarray] = None
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
square_root = _sqrt(prod - 4 * nngp**2)
nngp = factor * np.arctan2(2 * nngp, square_root)
if ntk is not None:
dot_sigma = 2 * factor / square_root
ntk *= dot_sigma
return nngp, ntk
def rect_to_polar(p_rect):
"""
Convert (x, y) position vector in rectangular coordinates to (rho, psi) polar coordinate vector
:param p_rect: (x, y) position vector in rectangular coordinates
:return: (rho, psi) polar coordinate vector [m], [rad]
"""
x, y = p_rect
rho = np.sqrt(x**2 + y**2)
psi = np.arctan2(y, x)
return np.array([rho, psi])
def sample(self, key, sample_shape=()):
"""
** References: **
1. A New Unified Approach for the Simulation of a Wide Class of Directional Distributions
John T. Kent, Asaad M. Ganeiber & Kanti V. Mardia (2018)
"""
assert is_prng_key(key)
phi_key, psi_key = random.split(key)
corr = self.correlation
conc = jnp.stack((self.phi_concentration, self.psi_concentration))
eig = 0.5 * (conc[0] - corr**2 / conc[1])
eig = jnp.stack((jnp.zeros_like(eig), eig))
eigmin = jnp.where(eig[1] < 0, eig[1],
jnp.zeros_like(eig[1], dtype=eig.dtype))
eig = eig - eigmin
b0 = self._bfind(eig)
total = _numel(sample_shape)
phi_den = log_I1(0, conc[1]).squeeze(0)
batch_size = _numel(self.batch_shape)
phi_shape = (total, 2, batch_size)
phi_state = SineBivariateVonMises._phi_marginal(
phi_shape,
phi_key,
jnp.reshape(conc, (2, batch_size)),
jnp.reshape(corr, (batch_size, )),
jnp.reshape(eig, (2, batch_size)),
jnp.reshape(b0, (batch_size, )),
jnp.reshape(eigmin, (batch_size, )),
jnp.reshape(phi_den, (batch_size, )),
)
phi = jnp.arctan2(phi_state.phi[:, 1:], phi_state.phi[:, :1])
alpha = jnp.sqrt(conc[1]**2 + (corr * jnp.sin(phi))**2)
beta = jnp.arctan(corr / conc[1] * jnp.sin(phi))
psi = VonMises(beta, alpha).sample(psi_key)
phi_psi = jnp.concatenate(
(
(phi + self.phi_loc + pi) % (2 * pi) - pi,
(psi + self.psi_loc + pi) % (2 * pi) - pi,
),
axis=1,
)
phi_psi = jnp.transpose(phi_psi, (0, 2, 1))
return phi_psi.reshape(*sample_shape, *self.batch_shape,
*self.event_shape)
def _dynamics_real(input_val): diff = (self.max_bounds - self.min_bounds) / 2.0 mean = (self.max_bounds + self.min_bounds) / 2.0 x, u = input_val u = diff * np.tanh(u) + mean sin_theta = x[0] cos_theta = x[1] theta_dot = x[2] torque = u[0] theta = np.arctan2(sin_theta, cos_theta) theta_dot_dot = -3.0*self.g/(2*self.l)*np.sin(theta+np.pi) theta_dot_dot += 3.0 / (1.1 * 1.1**2) * torque next_theta = theta + theta_dot * self.dt return np.array([np.sin(next_theta), np.cos(next_theta), theta_dot + theta_dot_dot * self.dt])
def xyz2equirect(xyz):
"""
Convert unit vector to equirectangular coordinate,
inverse of equirect2xyz
Args:
xyz: jnp.ndarray [..., 3] unit vectors
Returns:
uv: jnp.ndarray [...] coordinates (x, y) in image space in [-1.0, 1.0]
"""
lat = jnp.arcsin(jnp.clip(xyz[..., 1], -1.0, 1.0))
lon = jnp.arctan2(xyz[..., 0], xyz[..., 2])
x = lon / jnp.pi
y = 2.0 * lat / jnp.pi
return jnp.stack([x, y], axis=-1)
def signed_torsion_angle(ci, cj, ck, cl):
"""
Batch compute the signed angle of a torsion angle. The torsion angle
between two planes should be periodic but not necessarily symmetric.
Parameters
----------
ci: shape [num_torsions, 3] np.array
coordinates of the 1st atom in the 1-4 torsion angle
cj: shape [num_torsions, 3] np.array
coordinates of the 2nd atom in the 1-4 torsion angle
ck: shape [num_torsions, 3] np.array
coordinates of the 3rd atom in the 1-4 torsion angle
cl: shape [num_torsions, 3] np.array
coordinates of the 4th atom in the 1-4 torsion angle
Returns
-------
shape [num_torsions,] np.array
array of torsion angles.
"""
# Taken from the wikipedia arctan2 implementation:
# https://en.wikipedia.org/wiki/Dihedral_angle
# We use an identical but numerically stable arctan2
# implementation as opposed to the OpenMM energy function to
# avoid asingularity when the angle is zero.
rij = delta_r(cj, ci)
rkj = delta_r(cj, ck)
rkl = delta_r(cl, ck)
n1 = np.cross(rij, rkj)
n2 = np.cross(rkj, rkl)
lhs = np.linalg.norm(n1, axis=-1)
rhs = np.linalg.norm(n2, axis=-1)
bot = lhs * rhs
y = np.sum(np.multiply(np.cross(n1, n2),
rkj / np.linalg.norm(rkj, axis=-1, keepdims=True)),
axis=-1)
x = np.sum(np.multiply(n1, n2), -1)
return np.arctan2(y, x)
def reduce_state(self, state):
"""Reduces a non-angular state into an angular state by replacing
sin(theta) and cos(theta) with theta.
In this case, it converts:
[sin(theta), cos(theta), theta'] -> [theta, theta']
Args:
state: Augmented state vector [state_size].
Returns:
Reduced state size [reducted_state_size].
"""
sin_theta, cos_theta, theta_dot = state
theta = np.arctan2(sin_theta, cos_theta)
return np.hstack([theta, theta_dot])
def calc_ddR_wt(q, dq, ddq):
"""Calculate the second time derivative of tray's rotation matrix."""
Δy = q[2] - q[1]
dΔy = dq[2] - dq[1]
ddΔy = ddq[2] - ddq[1]
Δx = jnp.sqrt((2 * TRAY_W)**2 - Δy**2)
θ_tw = jnp.arctan2(Δy, Δx)
dθ_tw = dΔy / Δx # TODO possible division by zero
ddθ_tw = (ddΔy * Δx**2 + dΔy**2 * Δy) / Δx**3
S1 = skew1(1)
ddR_wt = (ddθ_tw * S1 - dθ_tw**2) @ jax_rotation_matrix(θ_tw)
return ddR_wt
def __quaternion_to_angle(quaternion: np.ndarray) -> np.ndarray:
"""Convert quaternion vector to angle of rotation.
The quaternion should be in (w, x, y, z) format.
Args:
quaternion (torch.Tensor): tensor with quaternions.
Return:
torch.Tensor: tensor with angle of rotation.
Shape:
- Input: :math:`(*, 4)` where `*` means, any number of dimensions
- Output: :math:`(*)`
"""
if not isinstance(quaternion, np.ndarray):
raise TypeError("Input type is not a torch.Tensor. Got {}".format(
type(quaternion)))
if not quaternion.shape[-1] == 4:
raise ValueError(
"Input must be a tensor of shape Nx4 or 4. Got {}".format(
quaternion.shape))
# unpack input and compute conversion
q1: np.ndarray = quaternion[..., 1]
q2: np.ndarray = quaternion[..., 2]
q3: np.ndarray = quaternion[..., 3]
sin_squared_theta: np.ndarray = q1 * q1 + q2 * q2 + q3 * q3
sin_theta: np.ndarray = np.sqrt(sin_squared_theta)
cos_theta: np.ndarray = quaternion[..., 0]
two_theta: np.ndarray = 2.0 * np.where(
cos_theta < 0.0,
np.arctan2(-sin_theta, -cos_theta),
np.arctan2(sin_theta, cos_theta),
)
return two_theta
def dynamics(self, x, u, t):
sin_theta = x[0]
cos_theta = x[1]
theta_dot = x[2]
torque = u[0]
theta = np.arctan2(sin_theta, cos_theta)
theta_dot_dot = -3.0 * self.g / (2 * self.l) * np.sin(theta + np.pi)
theta_dot_dot += 3.0 / (self.m * self.l**2) * torque
next_theta = theta + theta_dot * self.dt
x_new = np.array([
np.sin(next_theta),
np.cos(next_theta), theta_dot + theta_dot_dot * self.dt
])
return x_new
