3D vector class in Python

Ilan Schnell, May 2008

Based on numpy, I wrote yet another 3D vector class in Python. Here's the source vec3.txt (original filename vec.py), which is only 100 lines. The class has all common features, and also supports conversion to and from spherical and cylindrical coordinates. The reason why it's so small is due to the fact that the class derives from numpy.ndarray, just like numpy's matrix class.

If you want to run this test suite (which is also the markdown from which this page was generated), you can also download doc.txt.

Creating 3D vector objects:

>>> from vec3 import *
>>> a = vec3(3, 4.2, 5)       # x, y, z as arguments
>>> a
vec3(3.0, 4.2, 5.0)
>>> vec3([2, 8, -9])          # coordinates as list
vec3(2.0, 8.0, -9.0)

The coordinates arguments may in fact be passed as sequence, which includes tuples, numpy arrays, and vec3 objects themselves, as long as the sequence has 3 elements. Without arguments, a zero vector is created, a new vec3 object is always a copy:

>>> b = vec3(a)
>>> a.x = 1.2
>>> a
vec3(1.2, 4.2, 5.0)
>>> b
vec3(3.0, 4.2, 5.0)

Components may be accessed and manipulated using the x, y and z attributes, as well as by indexing 0, 1 and 2:

>>> a[0], a.y, a[2] = 5.5, -3.2, 7.8
>>> a
vec3(5.5, -3.2, 7.8)
>>> a.x, a[1], a.z
(5.5, -3.2, 7.8)

Algebra:

>>> a + b
vec3(8.5, 1.0, 12.8)
>>> a - b
vec3(2.5, -7.4, 2.8)
>>> a * b                   # dot product
42.06
>>> 3.1 * b
vec3(9.3, 13.02, 15.5)
>>> a/2
vec3(2.75, -1.6, 3.9)
>>> cross(a, b)
vec3(-48.76, -4.1, 32.7)
>>> abs(a)
10.06628034578811
>>> a**2                    # norm
101.33
>>> a**1                    # same as abs(a)
10.06628034578811

String representation:

>>> print 25.7*a
[ 141.35  -82.24  200.46]

Copying:

>>> from copy import copy
>>> c = copy(a)             # or simply:   c = vec3(a)
>>> c
vec3(5.5, -3.2, 7.8)

Testing equality (accounting for round-off errors):

>>> a == c
True
>>> a == b
False

Conversion to list object:

>>> b.tolist()
[3.0, 4.2000000000000002, 5.0]

Converting to and from spherical coordinates:

>>> rtp = a.get_spherical()            # rtp = (r, theta, phi)
>>> rtp
(10.06628034578811, 0.68429613206766526, -0.52694322718942965)
>>> b.set_spherical(rtp)
>>> assert a == b
>>> b.set_spherical(1, 0, 0)           # north pole
>>> b
vec3(0.0, 0.0, 1.0)
>>> from math import pi
>>> b.set_spherical(1, pi, 0)          # south pole
>>> b
vec3(1.22460635382e-16, 0.0, -1.0)
>>> b.set_spherical(1, pi/2, pi/4)     # a point on equator
>>> b
vec3(0.707106781187, 0.707106781187, 6.12303176911e-17)

Converting to and from cylindrical coordinates:

>>> rpz = a.get_cylindrical()          # rpz = (rho, phi, z)
>>> rpz
(6.3631753079732132, -0.52694322718942965, 7.8)
>>> b.set_cylindrical(rpz)
>>> assert a == b
>>> a.set_cylindrical(1, pi/2, 0)
>>> a
vec3(6.12303176911e-17, 1.0, 0.0)

Pickling and unpickling:

>>> from pickle import loads, dumps
>>> a = vec3(1, -2, 3.678)
>>> assert a == loads(dumps(a))

Data Persistence using the objects in shelve module:

>>> import shelve
>>> d = shelve.open('/tmp/data')
>>> for i in xrange(1000):
...     d[str(i)] = vec3(0.1*i, 0.2*i, 0.3*i)
... 
>>> d['350']
vec3(35.0, 70.0, 105.0)
>>> d.close()
>>> 
>>> d = shelve.open('/tmp/data')
>>> d['751']
vec3(75.1, 150.2, 225.3)
>>> d.close()

Converting from numpy arrays:

>>> from numpy import array
>>> vec3(array([1, 7.5, 3]))
vec3(1.0, 7.5, 3.0)

Multiplying with a numpy 3x3 matrix:

>>> from numpy import matrix
>>> A = matrix([[1, -2, 4], [3, 7, -2], [0, 8, -5]])
>>> A
matrix([[ 1, -2,  4],
        [ 3,  7, -2],
        [ 0,  8, -5]])
>>> v = vec3(3, 4, -6)
>>> vec3(A * v)
vec3(-29.0, 49.0, 62.0)
>>> vec3(v * A)
vec3(15.0, -26.0, 34.0)
>>> assert vec3(v * A) * v == v * vec3(A * v) == (v*A*v).mean()
>>> w = vec3(23, 14, -26)
>>> assert vec3(w * A) * v == w * vec3(A * v) == (w*A*v).mean()