| 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 ValueError("Invalid argument for curl") |
|
|
|
|
| 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) |
| |
| 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 ValueError("Invalid argument for divergence") |
|
|
|
|
| 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 |
|
|
