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]

Reply via email to