Dear all,
Something I've been missing in sympy is support for simple vector calculus.
So I wrote up a little module (with support for curvilinear
coordinate systems), and I attach it to this e-mail in case anyone else
finds it useful.
Cheers,
Jorn
--~--~---------~--~----~------------~-------~--~----~
You received this message because you are subscribed to the Google Groups
"sympy" group.
To post to this group, send email to [email protected]
To unsubscribe from this group, send email to [email protected]
For more options, visit this group at http://groups.google.com/group/sympy?hl=en
-~----------~----~----~----~------~----~------~--~---
from sympy import Matrix, Symbol, diff, sin, cos
#
# This code follows the book
#
# P. C. Matthews, Vector Calculus. Springer Verlag, London, 1998.
#
class CoordinateSystem3D:
"""
General three-dimensioal coordinate system class.
"""
@classmethod
def grad(self, func):
"""
Calculate the gradient of 'func'.
"""
return Matrix([diff(func, self.xs[0]) / self.h[0], \
diff(func, self.xs[1]) / self.h[1], \
diff(func, self.xs[2]) / self.h[2]])
@classmethod
def div(self, func):
"""
Calculate the divergence of 'func'.
"""
return (diff(func[0] * self.h[1] * self.h[2], self.xs[0]) + \
diff(func[1] * self.h[2] * self.h[0], self.xs[1]) + \
diff(func[2] * self.h[0] * self.h[1], self.xs[2])) / \
(self.h[0] * self.h[1] * self.h[2])
@classmethod
def curl(self, func):
"""
Calculate the curl (rotor) of 'func'.
"""
curl0 = diff(func[2] * self.h[2], self.xs[1]) - \
diff(func[1] * self.h[1], self.xs[2])
curl1 = diff(func[0] * self.h[0], self.xs[2]) - \
diff(func[2] * self.h[2], self.xs[0])
curl2 = diff(func[1] * self.h[1], self.xs[0]) - \
diff(func[0] * self.h[0], self.xs[1])
mat = Matrix([self.h[0] * curl0, \
self.h[1] * curl1, \
self.h[2] * curl2])
return mat / (self.h[0] * self.h[1] * self.h[2])
@classmethod
def laplacian(self, func):
"""
Calculate the Laplacian of 'func'.
"""
return self.div(self.grad(func))
class CartesianCS(CoordinateSystem3D):
"""
Cartesian coordinate system.
"""
# Coordinate symbols.
xs = [Symbol('x', real=True), \
Symbol('y', real=True), \
Symbol('z', real=True)]
x, y, z = xs
# Basis vectors.
e = [Matrix([1, 0, 0]), \
Matrix([0, 1, 0]),
Matrix([0, 0, 1])]
# Scale factors.
h = [1, 1, 1]
class CylindricalCS(CoordinateSystem3D):
"""
Cylindrical coordinate system.
"""
# Coordinate symbols.
xs = [Symbol('r', real=True), \
Symbol('\\phi', real=True), \
Symbol('z', real=True)]
r, phi, z = xs
# Basis vectors.
e = [Matrix([cos(phi), \
sin(phi), \
0]), \
Matrix([-sin(phi), \
cos(phi), \
0]),
Matrix([0, 0, 1])]
# Scale factors.
h = [1, r, 1]
class SphericalCS(CoordinateSystem3D):
"""
Spherical coordinate system.
"""
# Coordinate symbols.
xs = [Symbol('r', real=True), \
Symbol('\\phi', real=True), \
Symbol('\\theta', real=True)]
r, phi, theta = xs
# Basis vectors.
e = [Matrix([sin(theta) * cos(phi), \
sin(theta) * sin(phi), \
cos(theta)]), \
Matrix([cos(theta) * cos(phi), \
cos(theta) * sin(phi), \
-sin(theta)]),
Matrix([-sin(phi), \
cos(phi), \
0])]
# Scale factors.
h = [1, r, r * sin(theta)]
class EarthCS(CoordinateSystem3D):
"""
Earth coordinate system.
"""
# Coordinate symbols.
xs = [Symbol('r', real=True), \
Symbol('\\phi', real=True), \
Symbol('\\theta', real=True)]
r, phi, theta = xs
# Basis vectors.
e = [Matrix([cos(theta) * cos(phi), \
cos(theta) * sin(phi), \
sin(theta)]), \
Matrix([sin(theta) * cos(phi), \
sin(theta) * sin(phi), \
-cos(theta)]),
Matrix([-sin(phi), \
cos(phi), \
0])]
# Scale factors.
h = [1, r, r * cos(theta)]
def vprod(x, y):
"""
Vector product of vectors x and y.
"""
prod1 = x[1] * y[2] - x[2] * y[1]
prod2 = x[2] * y[0] - x[0] * y[2]
prod3 = x[0] * y[1] - x[1] * y[0]
return Matrix([prod1, prod2, prod3])
from __future__ import division
from sympy import Matrix, pi, sqrt, trim, log
from sympy.mpmath import mp
from veccalc import *
# Cartesian.
cs = CartesianCS
q = 1.602176487e-19
Phi = 1 / (4 * pi) * q / sqrt(cs.x**2 + cs.y**2 + cs.z**2)
E = Matrix([1 / (4 * pi) * q * cs.x / (cs.x**2 + cs.y**2 + cs.z**2)**(3/2), \
1 / (4 * pi) * q * cs.y / (cs.x**2 + cs.y**2 + cs.z**2)**(3/2),
1 / (4 * pi) * q * cs.z / (cs.x**2 + cs.y**2 + cs.z**2)**(3/2)])
assert trim(E[0]) == trim(-cs.grad(Phi)[0])
assert cs.curl(E) == Matrix([0, 0, 0])
# FIXME this is really 0. How to simplify the expression using sympy?
assert cs.div(E).evalf(subs={cs.x: 1, cs.y: 2, cs.z: 3}) < mp.eps
assert cs.grad(cs.x ** 2 + cs.y ** 2) == Matrix([2 * cs.x, 2 * cs.y, 0])
assert cs.laplacian(cs.x ** 2 - cs.y ** 2) == 0
# Cylindrical.
cs = CylindricalCS
i = 2
B = Matrix([0, i / (2 * pi * cs.r), 0])
assert cs.curl(B) == Matrix([0, 0, 0])
assert cs.div(B) == 0
B = Matrix([0, i / (2 * pi) * cs.r, 0])
assert cs.curl(B) == Matrix([0, 0, i / pi])
assert cs.div(B) == 0
assert cs.laplacian(log(cs.r)) == 0
# Spherical.
for cs in [SphericalCS, EarthCS]:
q = 1.602176487e-19
Phi = 1 / (4 * pi) * q / cs.r
E = Matrix([1 / (4 * pi) * q / cs.r**2, 0, 0])
assert E == -cs.grad(Phi)
assert cs.curl(E) == Matrix([0, 0, 0])
assert cs.div(E) == 0
Phi = -1 / (4 * pi) * q * cs.r ** 2 / 2
E = Matrix([1 / (4 * pi) * q * cs.r, 0, 0])
assert E == -cs.grad(Phi)
assert cs.curl(E) == Matrix([0, 0, 0])
assert cs.div(E) == q / (4 / 3 * pi)
assert cs.laplacian(1 / cs.r) == 0
cs = CartesianCS
# Vector product.
assert vprod(cs.e[0], cs.e[1]) == cs.e[2]
assert vprod(cs.e[1], cs.e[0]) == -cs.e[2]
assert vprod(cs.e[2], cs.e[0]) == cs.e[1]