336 lines
9.3 KiB
Python
336 lines
9.3 KiB
Python
import collections
|
|
from sympy.core.expr import Expr
|
|
from sympy.core import sympify, S, preorder_traversal
|
|
from sympy.vector.coordsysrect import CoordSys3D
|
|
from sympy.vector.vector import Vector, VectorMul, VectorAdd, Cross, Dot
|
|
from sympy.core.function import Derivative
|
|
from sympy.core.add import Add
|
|
from sympy.core.mul import Mul
|
|
|
|
|
|
def _get_coord_systems(expr):
|
|
g = preorder_traversal(expr)
|
|
ret = set()
|
|
for i in g:
|
|
if isinstance(i, CoordSys3D):
|
|
ret.add(i)
|
|
g.skip()
|
|
return frozenset(ret)
|
|
|
|
|
|
def _split_mul_args_wrt_coordsys(expr):
|
|
d = collections.defaultdict(lambda: S.One)
|
|
for i in expr.args:
|
|
d[_get_coord_systems(i)] *= i
|
|
return list(d.values())
|
|
|
|
|
|
class Gradient(Expr):
|
|
"""
|
|
Represents unevaluated Gradient.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, Gradient
|
|
>>> R = CoordSys3D('R')
|
|
>>> s = R.x*R.y*R.z
|
|
>>> Gradient(s)
|
|
Gradient(R.x*R.y*R.z)
|
|
|
|
"""
|
|
|
|
def __new__(cls, expr):
|
|
expr = sympify(expr)
|
|
obj = Expr.__new__(cls, expr)
|
|
obj._expr = expr
|
|
return obj
|
|
|
|
def doit(self, **hints):
|
|
return gradient(self._expr, doit=True)
|
|
|
|
|
|
class Divergence(Expr):
|
|
"""
|
|
Represents unevaluated Divergence.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, Divergence
|
|
>>> R = CoordSys3D('R')
|
|
>>> v = R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k
|
|
>>> Divergence(v)
|
|
Divergence(R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k)
|
|
|
|
"""
|
|
|
|
def __new__(cls, expr):
|
|
expr = sympify(expr)
|
|
obj = Expr.__new__(cls, expr)
|
|
obj._expr = expr
|
|
return obj
|
|
|
|
def doit(self, **hints):
|
|
return divergence(self._expr, doit=True)
|
|
|
|
|
|
class Curl(Expr):
|
|
"""
|
|
Represents unevaluated Curl.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, Curl
|
|
>>> R = CoordSys3D('R')
|
|
>>> v = R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k
|
|
>>> Curl(v)
|
|
Curl(R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k)
|
|
|
|
"""
|
|
|
|
def __new__(cls, expr):
|
|
expr = sympify(expr)
|
|
obj = Expr.__new__(cls, expr)
|
|
obj._expr = expr
|
|
return obj
|
|
|
|
def doit(self, **hints):
|
|
return curl(self._expr, doit=True)
|
|
|
|
|
|
def curl(vect, doit=True):
|
|
"""
|
|
Returns the curl of a vector field computed wrt the base scalars
|
|
of the given coordinate system.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
vect : Vector
|
|
The vector operand
|
|
|
|
doit : bool
|
|
If True, the result is returned after calling .doit() on
|
|
each component. Else, the returned expression contains
|
|
Derivative instances
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, curl
|
|
>>> R = CoordSys3D('R')
|
|
>>> v1 = R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k
|
|
>>> curl(v1)
|
|
0
|
|
>>> v2 = R.x*R.y*R.z*R.i
|
|
>>> curl(v2)
|
|
R.x*R.y*R.j + (-R.x*R.z)*R.k
|
|
|
|
"""
|
|
|
|
coord_sys = _get_coord_systems(vect)
|
|
|
|
if len(coord_sys) == 0:
|
|
return Vector.zero
|
|
elif len(coord_sys) == 1:
|
|
coord_sys = next(iter(coord_sys))
|
|
i, j, k = coord_sys.base_vectors()
|
|
x, y, z = coord_sys.base_scalars()
|
|
h1, h2, h3 = coord_sys.lame_coefficients()
|
|
vectx = vect.dot(i)
|
|
vecty = vect.dot(j)
|
|
vectz = vect.dot(k)
|
|
outvec = Vector.zero
|
|
outvec += (Derivative(vectz * h3, y) -
|
|
Derivative(vecty * h2, z)) * i / (h2 * h3)
|
|
outvec += (Derivative(vectx * h1, z) -
|
|
Derivative(vectz * h3, x)) * j / (h1 * h3)
|
|
outvec += (Derivative(vecty * h2, x) -
|
|
Derivative(vectx * h1, y)) * k / (h2 * h1)
|
|
|
|
if doit:
|
|
return outvec.doit()
|
|
return outvec
|
|
else:
|
|
if isinstance(vect, (Add, VectorAdd)):
|
|
from sympy.vector import express
|
|
try:
|
|
cs = next(iter(coord_sys))
|
|
args = [express(i, cs, variables=True) for i in vect.args]
|
|
except ValueError:
|
|
args = vect.args
|
|
return VectorAdd.fromiter(curl(i, doit=doit) for i in args)
|
|
elif isinstance(vect, (Mul, VectorMul)):
|
|
vector = [i for i in vect.args if isinstance(i, (Vector, Cross, Gradient))][0]
|
|
scalar = Mul.fromiter(i for i in vect.args if not isinstance(i, (Vector, Cross, Gradient)))
|
|
res = Cross(gradient(scalar), vector).doit() + scalar*curl(vector, doit=doit)
|
|
if doit:
|
|
return res.doit()
|
|
return res
|
|
elif isinstance(vect, (Cross, Curl, Gradient)):
|
|
return Curl(vect)
|
|
else:
|
|
raise Curl(vect)
|
|
|
|
|
|
def divergence(vect, doit=True):
|
|
"""
|
|
Returns the divergence of a vector field computed wrt the base
|
|
scalars of the given coordinate system.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
vector : Vector
|
|
The vector operand
|
|
|
|
doit : bool
|
|
If True, the result is returned after calling .doit() on
|
|
each component. Else, the returned expression contains
|
|
Derivative instances
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, divergence
|
|
>>> R = CoordSys3D('R')
|
|
>>> v1 = R.x*R.y*R.z * (R.i+R.j+R.k)
|
|
|
|
>>> divergence(v1)
|
|
R.x*R.y + R.x*R.z + R.y*R.z
|
|
>>> v2 = 2*R.y*R.z*R.j
|
|
>>> divergence(v2)
|
|
2*R.z
|
|
|
|
"""
|
|
coord_sys = _get_coord_systems(vect)
|
|
if len(coord_sys) == 0:
|
|
return S.Zero
|
|
elif len(coord_sys) == 1:
|
|
if isinstance(vect, (Cross, Curl, Gradient)):
|
|
return Divergence(vect)
|
|
# TODO: is case of many coord systems, this gets a random one:
|
|
coord_sys = next(iter(coord_sys))
|
|
i, j, k = coord_sys.base_vectors()
|
|
x, y, z = coord_sys.base_scalars()
|
|
h1, h2, h3 = coord_sys.lame_coefficients()
|
|
vx = _diff_conditional(vect.dot(i), x, h2, h3) \
|
|
/ (h1 * h2 * h3)
|
|
vy = _diff_conditional(vect.dot(j), y, h3, h1) \
|
|
/ (h1 * h2 * h3)
|
|
vz = _diff_conditional(vect.dot(k), z, h1, h2) \
|
|
/ (h1 * h2 * h3)
|
|
res = vx + vy + vz
|
|
if doit:
|
|
return res.doit()
|
|
return res
|
|
else:
|
|
if isinstance(vect, (Add, VectorAdd)):
|
|
return Add.fromiter(divergence(i, doit=doit) for i in vect.args)
|
|
elif isinstance(vect, (Mul, VectorMul)):
|
|
vector = [i for i in vect.args if isinstance(i, (Vector, Cross, Gradient))][0]
|
|
scalar = Mul.fromiter(i for i in vect.args if not isinstance(i, (Vector, Cross, Gradient)))
|
|
res = Dot(vector, gradient(scalar)) + scalar*divergence(vector, doit=doit)
|
|
if doit:
|
|
return res.doit()
|
|
return res
|
|
elif isinstance(vect, (Cross, Curl, Gradient)):
|
|
return Divergence(vect)
|
|
else:
|
|
raise Divergence(vect)
|
|
|
|
|
|
def gradient(scalar_field, doit=True):
|
|
"""
|
|
Returns the vector gradient of a scalar field computed wrt the
|
|
base scalars of the given coordinate system.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
scalar_field : SymPy Expr
|
|
The scalar field to compute the gradient of
|
|
|
|
doit : bool
|
|
If True, the result is returned after calling .doit() on
|
|
each component. Else, the returned expression contains
|
|
Derivative instances
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, gradient
|
|
>>> R = CoordSys3D('R')
|
|
>>> s1 = R.x*R.y*R.z
|
|
>>> gradient(s1)
|
|
R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k
|
|
>>> s2 = 5*R.x**2*R.z
|
|
>>> gradient(s2)
|
|
10*R.x*R.z*R.i + 5*R.x**2*R.k
|
|
|
|
"""
|
|
coord_sys = _get_coord_systems(scalar_field)
|
|
|
|
if len(coord_sys) == 0:
|
|
return Vector.zero
|
|
elif len(coord_sys) == 1:
|
|
coord_sys = next(iter(coord_sys))
|
|
h1, h2, h3 = coord_sys.lame_coefficients()
|
|
i, j, k = coord_sys.base_vectors()
|
|
x, y, z = coord_sys.base_scalars()
|
|
vx = Derivative(scalar_field, x) / h1
|
|
vy = Derivative(scalar_field, y) / h2
|
|
vz = Derivative(scalar_field, z) / h3
|
|
|
|
if doit:
|
|
return (vx * i + vy * j + vz * k).doit()
|
|
return vx * i + vy * j + vz * k
|
|
else:
|
|
if isinstance(scalar_field, (Add, VectorAdd)):
|
|
return VectorAdd.fromiter(gradient(i) for i in scalar_field.args)
|
|
if isinstance(scalar_field, (Mul, VectorMul)):
|
|
s = _split_mul_args_wrt_coordsys(scalar_field)
|
|
return VectorAdd.fromiter(scalar_field / i * gradient(i) for i in s)
|
|
return Gradient(scalar_field)
|
|
|
|
|
|
class Laplacian(Expr):
|
|
"""
|
|
Represents unevaluated Laplacian.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy.vector import CoordSys3D, Laplacian
|
|
>>> R = CoordSys3D('R')
|
|
>>> v = 3*R.x**3*R.y**2*R.z**3
|
|
>>> Laplacian(v)
|
|
Laplacian(3*R.x**3*R.y**2*R.z**3)
|
|
|
|
"""
|
|
|
|
def __new__(cls, expr):
|
|
expr = sympify(expr)
|
|
obj = Expr.__new__(cls, expr)
|
|
obj._expr = expr
|
|
return obj
|
|
|
|
def doit(self, **hints):
|
|
from sympy.vector.functions import laplacian
|
|
return laplacian(self._expr)
|
|
|
|
|
|
def _diff_conditional(expr, base_scalar, coeff_1, coeff_2):
|
|
"""
|
|
First re-expresses expr in the system that base_scalar belongs to.
|
|
If base_scalar appears in the re-expressed form, differentiates
|
|
it wrt base_scalar.
|
|
Else, returns 0
|
|
"""
|
|
from sympy.vector.functions import express
|
|
new_expr = express(expr, base_scalar.system, variables=True)
|
|
arg = coeff_1 * coeff_2 * new_expr
|
|
return Derivative(arg, base_scalar) if arg else S.Zero
|