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