#!/usr/bin/env python
#*****************
#Vector Operations
#*****************
#
#*Author:* Dustin Kleckner (2014)
import numpy as np
#------------------------------------------------------------------------------
# Basic Operations
#------------------------------------------------------------------------------
[docs]def mag(X):
'''Calculate the length of an array of vectors.'''
return np.sqrt((np.asarray(X)**2).sum(-1))
[docs]def mag1(X):
'''Calculate the length of an array of vectors, keeping the last dimension
index.'''
return np.sqrt((np.asarray(X)**2).sum(-1))[..., np.newaxis]
[docs]def dot(X, Y):
'''Calculate the dot product of two arrays of vectors.'''
return (np.asarray(X)*Y).sum(-1)
[docs]def dot1(X, Y):
'''Calculate the dot product of two arrays of vectors, keeping the last
dimension index'''
return (np.asarray(X)*Y).sum(-1)[..., np.newaxis]
[docs]def norm(X):
'''Computes a normalized version of an array of vectors.'''
return X / mag1(X)
[docs]def plus(X):
'''Return a shifted version of an array of vectors.'''
return np.roll(X, -1, 0)
[docs]def minus(X):
'''Return a shifted version of an array of vectors.'''
return np.roll(X, +1, 0)
[docs]def cross(X, Y):
'''Return the cross-product of two vectors.'''
return np.cross(X, Y)
[docs]def proj(X, Y):
r'''Return the projection of one vector onto another.
Parameters
----------
X, Y : vector array
Returns
-------
Z : vector array
:math:`\vec{Z} = \frac{\vec{X} \cdot \vec{Y}}{|Y|^2} \vec{Y}`
'''
Yp = norm(Y)
return dot1(Yp, X) * Yp
[docs]def midpoint_delta(X):
'''Returns center point and vector of each edge of the polygon defined by the points.'''
Xp = plus(X)
return (Xp + X) / 2., (Xp - X)
[docs]def arb_perp(V):
'''For each vector, return an arbitrary vector that is perpendicular.
**Note: arbitrary does not mean random!**'''
p = eye(3, dtype=V.dtype)[argmin(V, -1)]
return norm(p - proj(p, V))
[docs]def apply_basis(X, B):
'''Transform each vector into the specified basis/bases.
Parameters
----------
X : vector array, shape [..., 3]
B : orthonormal basis array, shape [..., 3, 3]
Returns
-------
Y : vector array, shape [..., 3]
X transformed into the basis given by B
'''
return (np.asarray(X)[..., np.newaxis, :] * B).sum(-1)
#------------------------------------------------------------------------------
# Building vectors intelligently
#------------------------------------------------------------------------------
[docs]def vec(x=[0], y=[0], z=[0]):
'''Generate a [..., 3] vector from seperate x, y, z.
Parameters
----------
x, y, z: array
coordinates; default to 0, may have any shape
Returns
-------
X : [..., 3] array'''
x, y, z = map(np.asarray, [x, y, z])
s = [1]
for a in (x, y, z):
while a.ndim > len(s): s.prepend(1)
s = [max(ss, n) for ss, n in zip(s, a.shape)]
v = np.empty(s + [3], 'd')
v[..., 0] = x
v[..., 1] = y
v[..., 2] = z
return v
#------------------------------------------------------------------------------
# Rotations and basis operations
#------------------------------------------------------------------------------
[docs]def rot(a, X=None, cutoff=1E-10):
'''Rotate points around an arbitrary axis.
Parameters
----------
a : [..., 3] array
Rotation vector, will rotate counter-clockwise around axis by an amount
given be the length of the vector (in radians). May be a single vector
or an array of vectors if each point is rotated separately.
X : [..., 3] array
Vectors to rotate; if not specified generates a rotation basis instead.
cutoff : float
If length of vector less than this value (1E-10 by default), no rotation
is performed. (Used to avoid basis errors)
Returns
-------
Y : [..., 3] array
Rotated vectors or rotation basis.
'''
#B = np.eye(3, dtype='d' if X is None else X.dtype)
a = np.asarray(a)
if X is None: X = np.eye(3).astype(a.dtype)
phi = mag(a)
if phi.max() < 1E-10: return X
#http://en.wikipedia.org/w/index.php?title=Rotation_matrix#Axis_and_angle
n = norm(a)
n[np.where(np.isnan(n).any(-1))] = (1, 0, 0)
B = np.zeros(a.shape[:-1] + (3, 3), dtype=a.dtype)
c = np.cos(phi)
s = np.sin(phi)
C = 1 - c
for i in range(3):
for j in range(3):
if i == j:
extra = c
else:
if (j - i)%3 == 2:
extra = +s * n[..., (j-1)%3]
else:
extra = -s * n[..., (j+1)%3]
B[..., i, j] = n[..., i]*n[..., j]*C + extra
##Create a new basis where the rotation is simply in x
#Ba = normalize_basis(a[..., np.newaxis, :])
#
###B = apply_basis(B, Ba) #This was pointless, B was 1
##y, z = B[..., 1, :].copy(), B[..., 2, :].copy()
##
##c, s = np.cos(phi), np.sin(phi)
##
###Rotate in new basis
##B[..., 1, :] = y*c - z*s
##B[..., 2, :] = y*s + z*c
##
##B = apply_basis(B, Ba.T).T
#
#B = np.zeros_like(Ba)
#B[..., 0, 0] = 1
#B[..., 1, 1] = +cos(phi)
#B[..., 1, 2] = -sin(phi)
#B[..., 2, 1] = +sin(phi)
#B[..., 2, 2] = +cos(phi)
#
#B = apply_basis(B, Ba)
if X is not None: return apply_basis(X, B)
else: return B
[docs]def normalize_basis(B):
'''Create right-handed orthonormal basis/bases from input basis.
Parameters
----------
B : [..., 1-3, 3] array
input basis, should be at least 2d. If the second to last axis has
1 vectors, it will automatically create an arbitrary orthonormal basis
with the specified first vector.
(note: even if three bases are specified, the last is always ignored,
and is generated by a cross product of the first two.)
Returns
-------
NB : [..., 3, 3] array
orthonormal basis
'''
B = np.asarray(B)
NB = np.empty(B.shape[:-2] + (3, 3), dtype='d')
v1 = norm(B[..., 0, :])
v1[np.where(np.isnan(v1).any(-1))] = (1, 0, 0)
v2 = B[..., 1, :] if B.shape[-2] >= 2 else np.eye(3)[np.argmin(abs(v1), axis=-1)]
v2 = norm(v2 - v1 * dot1(v1, v2))
v3 = cross(v1, v2)
for i, v in enumerate([v1, v2, v3]): NB[..., i, :] = v
return NB
if __name__ == '__main__':
TEST = 2
if TEST == 1:
x = np.eye(3)[0]
a = vec(z = np.linspace(0, np.pi, 5))
print a
y = rot(a, x)
print y
if TEST == 2:
from pylab import *
NP = 10
x = arange(NP) * 2 * pi / NP
yd = [sin(x), cos(x), -sin(x), -cos(x)]
for i in range(4):
y = yd[i]
yy = yd[0]
yi = d5p(yy, h=x[1]-x[0], d=i, closed=True)
plot(x, y+2*i)
plot(x, yi+2*i, '.')
print i, '%.12f' % mean(abs(y - yi))
show()