SYMBOLIC ROTATIONS in 3D with sympy

In [1]:
import sympy as sp
vy,vy,vz,theta,c,s,V = sp.symbols('vy vy vz theta c s V')
phi = sp.symbols('phi')
In [2]:
def Rx(theta,V):
    """
    Rotation of a 3d vector V
    of an angle theta around
    the x-axis
    """
    c = sp.cos(theta)
    s = sp.sin(theta)
    R = sp.Matrix(
        [[1, 0, 0], 
         [0, c, -s], 
         [0, s, c]]
    )
    return R*V
In [3]:
V = sp.Matrix([[1],[2],[3]])
Rx(theta,V)
Out[3]:
Matrix([
[                           1],
[-3*sin(theta) + 2*cos(theta)],
[ 2*sin(theta) + 3*cos(theta)]])
In [4]:
def Ry(theta,V):
    """
    Rotation of a 3d vector V
    of an angle theta around
    the y-axis
    """
    c = sp.cos(theta)
    s = sp.sin(theta)
    R = sp.Matrix(
        [[c, 0, s],
         [0, 1, 0],
         [-s, 0, c]]
    )
    return R*V
In [5]:
def Rz(theta,V):
    """
    Rotation of a 3d vector V
    of an angle theta around
    the z-axis
    """
    c = sp.cos(theta)
    s = sp.sin(theta)
    R = sp.Matrix(
        [[c, -s, 0], 
         [s, c, 0],
         [0, 0, 1]]
    )
    return R*V
In [6]:
Rz(sp.pi/4,V)
Out[6]:
Matrix([
[ -sqrt(2)/2],
[3*sqrt(2)/2],
[          3]])
In [7]:
def R(theta,phi,V):
    """
    Rotation of a 3d vector V
    of an angle theta around
    the y-axis followed by a rotation
    of phi around the z-axis
    """
    return Rz(phi,Ry(theta,V))
    

Examples:

In [8]:
V = sp.Matrix([[1],[0],[0]])
R(theta,phi,V)
Out[8]:
Matrix([
[cos(phi)*cos(theta)],
[sin(phi)*cos(theta)],
[        -sin(theta)]])
In [9]:
V = sp.Matrix([[0],[0],[1]])
R(theta,phi,V)
Out[9]:
Matrix([
[sin(theta)*cos(phi)],
[sin(phi)*sin(theta)],
[         cos(theta)]])
In [10]:
V = sp.Matrix([[1],[-sp.I],[0]])
R(theta,phi,V)
Out[10]:
Matrix([
[I*sin(phi) + cos(phi)*cos(theta)],
[sin(phi)*cos(theta) - I*cos(phi)],
[                     -sin(theta)]])
In [11]:
V = sp.Matrix([[-1],[-sp.I],[0]])
R(theta,phi,V)
Out[11]:
Matrix([
[ I*sin(phi) - cos(phi)*cos(theta)],
[-sin(phi)*cos(theta) - I*cos(phi)],
[                       sin(theta)]])
In [12]:
V = sp.Matrix([[0],[0],[1]])
R(theta,phi,V)
Out[12]:
Matrix([
[sin(theta)*cos(phi)],
[sin(phi)*sin(theta)],
[         cos(theta)]])
In [ ]: