"""Exploration of Vectors and Frames.
Copyright 2012 Allen B. Downey
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html
"""
from __future__ import print_function, division
import sys
import numpy
import math
def println(s):
print(s, '\n')
class FrameError(ValueError):
"""Represents a problem with frame of reference."""
class Vector:
def __init__(self, array, frame=None):
"""A vector is an array of coordinates and a frame of reference.
array:
frame: Frame object
"""
self.array = array
self.frame = frame
def __str__(self):
if self.frame == None:
return '^{O}%s' % (str(self.array), )
else:
return '^{%s}%s' % (str(self.frame), str(self.array))
def __add__(self, other):
if self.frame != other.frame:
raise FrameError("Vectors must be relative to the same frame.")
return Vector(self.array + other.array, self.frame)
@staticmethod
def from_list(t, frame=None):
"""Makes a vector from a list.
t: list of coordinates
frame: reference Frame
"""
return Vector(numpy.array(t), frame)
class Rotation:
def __init__(self, array):
self.array = array
def __str__(self):
return 'Rotation\n%s' % str(self.array)
def __neg__(self):
return Rotation(-self.array)
def __mul__(self, other):
"""Apply the rotation to a Vector."""
return numpy.dot(self.array, other.array)
__call__ = __mul__
@staticmethod
def from_axis(axis, theta):
x, y, z = numpy.ravel(axis.array)
c = math.cos(theta)
u = 1.0-c
s = math.sqrt(1.0-c*c)
xu, yu, zu = x*u, y*u, z*u
v1 = [x*xu + c, x*yu - z*s, x*zu + y*s]
v2 = [x*yu + z*s, y*yu + c, y*zu - x*s]
v3 = [x*zu - y*s, y*zu + x*s, z*zu + c]
return Rotation(numpy.array([v1, v2, v3]))
def to_axis(self):
pass
def transpose(self):
return Rotation(numpy.transpose(self.array))
inverse = transpose
class Transform:
"""Represents a transform from one Frame to another."""
def __init__(self, rot, org, source=None):
"""Instantiates a Transform.
rot: Rotation object
org: origin Vector
source: source Frame
"""
self.rot = rot
self.org = org
self.dest = org.frame
self.source = source
self.source.add_transform(self)
def __str__(self):
"""Returns a string representation of the Transform."""
if self.dest == None:
return '%s' % self.source.name
return '_{%s}^{O}T' % self.source.name
else:
return '_{%s}^{%s}T' % (self.source.name, self.dest.name)
def __mul__(self, other):
"""Applies a Transform to a Vector or Transform."""
if isinstance(other, Vector):
return self.mul_vector(other)
if isinstance(other, Transform):
return self.mul_transform(other)
__call__ = __mul__
def mul_vector(self, p):
"""Applies a Transform to a Vector.
p: Vector
Returns: Vector
"""
if p.frame != self.source:
raise FrameError(
"The frame of the vector must be the source of the transform")
return Vector(self.rot * p, self.dest) + self.org
def mul_transform(self, other):
"""Applies a Transform to another Transform.
other: Transform
Returns Transform
"""
if other.dest != self.source:
raise FrameError(
"This frames source must be the other frame's destination.")
rot = Rotation(self.rot * other.rot)
t = Transform(rot, self * other.org, other.source)
return t
def inverse(self):
"""Computes the inverse transform.
Returns: Transform
"""
irot = self.rot.inverse()
iorg = Vector(-(irot * self.org), self.source)
t = Transform(irot, iorg, self.dest)
return t
class Frame:
"""Represents a frame of reference."""
roster = []
def __init__(self, name):
"""Instantiate a Frame.
name: string
"""
self.name = name
self.transforms = {}
Frame.roster.append(self)
def __str__(self):
return self.name
def add_transform(self, transform):
"""A frames is defined by a Transform relative to another Frame.
transform: Transform object
"""
if transform.source != self:
raise FrameError("Source of the transform must be this Frame.")
if transform.dest:
self.transforms[transform.dest] = transform
def dests(self):
"""Returns a list of the Frames we know how to Transform to."""
return self.transforms.keys()
class Vertex:
"""Represents a node in a graph."""
def __init__(self, frame):
self.frame = frame
self.dist = 1000000
self.out = []
def __str__(self):
return '%s %d' % (self.frame.name, self.dist)
def shortest_path(start, frames):
"""For a given list of frames and a starting frame,
find the shortest path of transforms from the
starting frame to all other frames.
The 'distance' is the number of inverse transformations
that must be calculated.
The result is a dictionary of vertices, where
each vertex is labeled with the frame it corresponds
to, the distance from the starting frame, and the prev
frame along the path from start. """
map = dict([(f, Vertex(f)) for f in frames])
length = {}
for v in map.values():
for dest in v.frame.transforms:
w = map[dest]
v.out.append(w)
length[(v, w)] = 0
w.out.append(v)
length[(w, v)] = 1
s = map[start]
s.dist = 0
queue = [s]
while queue:
v = queue.pop()
for w in v.out:
d = v.dist + length[(v,w)]
if d < w.dist:
w.dist = d
w.prev = v
if w not in queue: queue.append(w)
return map
def print_shortest_path(map):
for source, v in map.items():
try:
v.prev
print(source, v.dist, v.prev.frame)
except:
print(source, v.dist)
def print_length(length):
for v, w in length:
print(v.frame.name, w.frame.name, length[(v, w)])
print()
def main(name):
theta = math.pi/2
origin = Frame('O')
xhat = Vector.from_list([1, 0, 0], origin)
rx = Rotation.from_axis(xhat, theta)
a = Frame('A')
t_ao = Transform(rx, xhat, a)
yhat = Vector.from_list([0, 1, 0], a)
ry = Rotation.from_axis(yhat, theta)
b = Frame('B')
t_ba = Transform(ry, yhat, b)
zhat = Vector.from_list([0, 0, 1], b)
rz = Rotation.from_axis(zhat, theta)
c = Frame('C')
t_cb = Transform(rz, zhat, c)
p_c = Vector.from_list([1, 1, 1], c)
println(p_c)
p_b = t_cb(p_c)
println(p_b)
p_a = t_ba(p_b)
println(p_a)
p = t_ao(p_a)
println(p)
map = shortest_path(origin, Frame.roster)
print_shortest_path(map)
cbao = t_ao(t_ba(t_cb))
p = cbao(p_c)
println(p)
inv = cbao.inverse()
p_c = inv(p)
println(p_c)
map = shortest_path(origin, Frame.roster)
print_shortest_path(map)
if __name__ == '__main__':
main(*sys.argv)
