1from ..core import Integer, diff, sympify
2from ..integrals import integrate
3from .coordsysrect import CoordSysCartesian
4from .dyadic import Dyadic
5from .scalar import BaseScalar
6from .vector import BaseVector, Vector
7
8
9def express(expr, system, system2=None, variables=False):
10    """
11    Global function for 'express' functionality.
12
13    Re-expresses a Vector, Dyadic or scalar(diofant-able) in the given
14    coordinate system.
15
16    If 'variables' is True, then the coordinate variables (base scalars)
17    of other coordinate systems present in the vector/scalar field or
18    dyadic are also substituted in terms of the base scalars of the
19    given system.
20
21    Parameters
22    ==========
23
24    expr : Vector/Dyadic/scalar(diofant-able)
25        The expression to re-express in CoordSysCartesian 'system'
26
27    system: CoordSysCartesian
28        The coordinate system the expr is to be expressed in
29
30    system2: CoordSysCartesian
31        The other coordinate system required for re-expression
32        (only for a Dyadic Expr)
33
34    variables : boolean
35        Specifies whether to substitute the coordinate variables present
36        in expr, in terms of those of parameter system
37
38    Examples
39    ========
40
41    >>> N = CoordSysCartesian('N')
42    >>> q = Symbol('q')
43    >>> B = N.orient_new_axis('B', q, N.k)
44    >>> express(B.i, N)
45    (cos(q))*N.i + (sin(q))*N.j
46    >>> express(N.x, B, variables=True)
47    -sin(q)*B.y + cos(q)*B.x
48    >>> d = N.i.outer(N.i)
49    >>> express(d, B, N)
50    (cos(q))*(B.i|N.i) + (-sin(q))*(B.j|N.i)
51
52    """
53    if expr in (0, Vector.zero):
54        return expr
55
56    if not isinstance(system, CoordSysCartesian):
57        raise TypeError('system should be a CoordSysCartesian instance')
58
59    if isinstance(expr, Vector):
60        if system2 is not None:
61            raise ValueError('system2 should not be provided for Vectors')
62        # Given expr is a Vector
63        if variables:
64            # If variables attribute is True, substitute
65            # the coordinate variables in the Vector
66            system_list = []
67            for x in expr.atoms(BaseScalar, BaseVector):
68                if x.system != system:
69                    system_list.append(x.system)
70            system_list = set(system_list)
71            subs_dict = {}
72            for f in system_list:
73                subs_dict.update(f.scalar_map(system))
74            expr = expr.subs(subs_dict)
75        # Re-express in this coordinate system
76        outvec = Vector.zero
77        parts = expr.separate()
78        for x in parts:
79            if x != system:
80                temp = system.rotation_matrix(x) * parts[x].to_matrix(x)
81                outvec += matrix_to_vector(temp, system)
82            else:
83                outvec += parts[x]
84        return outvec
85
86    elif isinstance(expr, Dyadic):
87        if system2 is None:
88            system2 = system
89        if not isinstance(system2, CoordSysCartesian):
90            raise TypeError('system2 should be a CoordSysCartesian instance')
91        outdyad = Dyadic.zero
92        var = variables
93        for k, v in expr.components.items():
94            outdyad += (express(v, system, variables=var) *
95                        (express(k.args[0], system, variables=var) |
96                         express(k.args[1], system2, variables=var)))
97
98        return outdyad
99
100    else:
101        if system2 is not None:
102            raise ValueError('system2 should not be provided for Vectors')
103        if variables:
104            # Given expr is a scalar field
105            system_set = set()
106            expr = sympify(expr)
107            # Subsitute all the coordinate variables
108            for x in expr.atoms(BaseScalar):
109                if x.system != system:
110                    system_set.add(x.system)
111            subs_dict = {}
112            for f in system_set:
113                subs_dict.update(f.scalar_map(system))
114            return expr.subs(subs_dict)
115        return expr
116
117
118def curl(vect, coord_sys):
119    """
120    Returns the curl of a vector field computed wrt the base scalars
121    of the given coordinate system.
122
123    Parameters
124    ==========
125
126    vect : Vector
127        The vector operand
128
129    coord_sys : CoordSysCartesian
130        The coordinate system to calculate the curl in
131
132    Examples
133    ========
134
135    >>> R = CoordSysCartesian('R')
136    >>> v1 = R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k
137    >>> curl(v1, R)
138    0
139    >>> v2 = R.x*R.y*R.z*R.i
140    >>> curl(v2, R)
141    R.x*R.y*R.j + (-R.x*R.z)*R.k
142
143    """
144    return coord_sys.delop.cross(vect).doit()
145
146
147def divergence(vect, coord_sys):
148    """
149    Returns the divergence of a vector field computed wrt the base
150    scalars of the given coordinate system.
151
152    Parameters
153    ==========
154
155    vect : Vector
156        The vector operand
157
158    coord_sys : CoordSysCartesian
159        The coordinate system to calculate the divergence in
160
161    Examples
162    ========
163
164    >>> R = CoordSysCartesian('R')
165    >>> v1 = R.x*R.y*R.z * (R.i+R.j+R.k)
166    >>> divergence(v1, R)
167    R.x*R.y + R.x*R.z + R.y*R.z
168    >>> v2 = 2*R.y*R.z*R.j
169    >>> divergence(v2, R)
170    2*R.z
171
172    """
173    return coord_sys.delop.dot(vect).doit()
174
175
176def gradient(scalar, coord_sys):
177    """
178    Returns the vector gradient of a scalar field computed wrt the
179    base scalars of the given coordinate system.
180
181    Parameters
182    ==========
183
184    scalar : Diofant Expr
185        The scalar field to compute the gradient of
186
187    coord_sys : CoordSysCartesian
188        The coordinate system to calculate the gradient in
189
190    Examples
191    ========
192
193    >>> R = CoordSysCartesian('R')
194    >>> s1 = R.x*R.y*R.z
195    >>> gradient(s1, R)
196    R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k
197    >>> s2 = 5*R.x**2*R.z
198    >>> gradient(s2, R)
199    10*R.x*R.z*R.i + 5*R.x**2*R.k
200
201    """
202    return coord_sys.delop(scalar).doit()
203
204
205def is_conservative(field):
206    """
207    Checks if a field is conservative.
208
209    Parameters
210    ==========
211
212    field : Vector
213        The field to check for conservative property
214
215    Examples
216    ========
217
218    >>> R = CoordSysCartesian('R')
219    >>> is_conservative(R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k)
220    True
221    >>> is_conservative(R.z*R.j)
222    False
223
224    """
225    # Field is conservative irrespective of system
226    # Take the first coordinate system in the result of the
227    # separate method of Vector
228    if not isinstance(field, Vector):
229        raise TypeError('field should be a Vector')
230    if field == Vector.zero:
231        return True
232    coord_sys = list(field.separate())[0]
233    return curl(field, coord_sys).simplify() == Vector.zero
234
235
236def is_solenoidal(field):
237    """
238    Checks if a field is solenoidal.
239
240    Parameters
241    ==========
242
243    field : Vector
244        The field to check for solenoidal property
245
246    Examples
247    ========
248
249    >>> R = CoordSysCartesian('R')
250    >>> is_solenoidal(R.y*R.z*R.i + R.x*R.z*R.j + R.x*R.y*R.k)
251    True
252    >>> is_solenoidal(R.y * R.j)
253    False
254
255    """
256    # Field is solenoidal irrespective of system
257    # Take the first coordinate system in the result of the
258    # separate method in Vector
259    if not isinstance(field, Vector):
260        raise TypeError('field should be a Vector')
261    if field == Vector.zero:
262        return True
263    coord_sys = list(field.separate())[0]
264    return divergence(field, coord_sys).simplify() == Integer(0)
265
266
267def scalar_potential(field, coord_sys):
268    """
269    Returns the scalar potential function of a field in a given
270    coordinate system (without the added integration constant).
271
272    Parameters
273    ==========
274
275    field : Vector
276        The vector field whose scalar potential function is to be
277        calculated
278
279    coord_sys : CoordSysCartesian
280        The coordinate system to do the calculation in
281
282    Examples
283    ========
284
285    >>> R = CoordSysCartesian('R')
286    >>> scalar_potential(R.k, R) == R.z
287    True
288    >>> scalar_field = 2*R.x**2*R.y*R.z
289    >>> grad_field = gradient(scalar_field, R)
290    >>> scalar_potential(grad_field, R)
291    2*R.x**2*R.y*R.z
292
293    """
294    # Check whether field is conservative
295    if not is_conservative(field):
296        raise ValueError('Field is not conservative')
297    if field == Vector.zero:
298        return Integer(0)
299    # Express the field exntirely in coord_sys
300    # Subsitute coordinate variables also
301    if not isinstance(coord_sys, CoordSysCartesian):
302        raise TypeError('coord_sys must be a CoordSysCartesian')
303    field = express(field, coord_sys, variables=True)
304    dimensions = coord_sys.base_vectors()
305    scalars = coord_sys.base_scalars()
306    # Calculate scalar potential function
307    temp_function = integrate(field.dot(dimensions[0]), scalars[0])
308    for i, dim in enumerate(dimensions[1:]):
309        partial_diff = diff(temp_function, scalars[i + 1])
310        partial_diff = field.dot(dim) - partial_diff
311        temp_function += integrate(partial_diff, scalars[i + 1])
312    return temp_function
313
314
315def scalar_potential_difference(field, coord_sys, point1, point2):
316    """
317    Returns the scalar potential difference between two points in a
318    certain coordinate system, wrt a given field.
319
320    If a scalar field is provided, its values at the two points are
321    considered. If a conservative vector field is provided, the values
322    of its scalar potential function at the two points are used.
323
324    Returns (potential at point2) - (potential at point1)
325
326    The position vectors of the two Points are calculated wrt the
327    origin of the coordinate system provided.
328
329    Parameters
330    ==========
331
332    field : Vector/Expr
333        The field to calculate wrt
334
335    coord_sys : CoordSysCartesian
336        The coordinate system to do the calculations in
337
338    point1 : Point
339        The initial Point in given coordinate system
340
341    position2 : Point
342        The second Point in the given coordinate system
343
344    Examples
345    ========
346
347    >>> R = CoordSysCartesian('R')
348    >>> P = R.origin.locate_new('P', R.x*R.i + R.y*R.j + R.z*R.k)
349    >>> vectfield = 4*R.x*R.y*R.i + 2*R.x**2*R.j
350    >>> scalar_potential_difference(vectfield, R, R.origin, P)
351    2*R.x**2*R.y
352    >>> Q = R.origin.locate_new('O', 3*R.i + R.j + 2*R.k)
353    >>> scalar_potential_difference(vectfield, R, P, Q)
354    -2*R.x**2*R.y + 18
355
356    """
357    if not isinstance(coord_sys, CoordSysCartesian):
358        raise TypeError('coord_sys must be a CoordSysCartesian')
359    if isinstance(field, Vector):
360        # Get the scalar potential function
361        scalar_fn = scalar_potential(field, coord_sys)
362    else:
363        # Field is a scalar
364        scalar_fn = field
365    # Express positions in required coordinate system
366    origin = coord_sys.origin
367    position1 = express(point1.position_wrt(origin), coord_sys,
368                        variables=True)
369    position2 = express(point2.position_wrt(origin), coord_sys,
370                        variables=True)
371    # Get the two positions as substitution dicts for coordinate variables
372    subs_dict1 = {}
373    subs_dict2 = {}
374    scalars = coord_sys.base_scalars()
375    for i, x in enumerate(coord_sys.base_vectors()):
376        subs_dict1[scalars[i]] = x.dot(position1)
377        subs_dict2[scalars[i]] = x.dot(position2)
378    return scalar_fn.subs(subs_dict2) - scalar_fn.subs(subs_dict1)
379
380
381def matrix_to_vector(matrix, system):
382    """
383    Converts a vector in matrix form to a Vector instance.
384
385    It is assumed that the elements of the Matrix represent the
386    measure numbers of the components of the vector along basis
387    vectors of 'system'.
388
389    Parameters
390    ==========
391
392    matrix : Diofant Matrix, Dimensions: (3, 1)
393        The matrix to be converted to a vector
394
395    system : CoordSysCartesian
396        The coordinate system the vector is to be defined in
397
398    Examples
399    ========
400
401    >>> m = Matrix([1, 2, 3])
402    >>> C = CoordSysCartesian('C')
403    >>> v = matrix_to_vector(m, C)
404    >>> v
405    C.i + 2*C.j + 3*C.k
406    >>> v.to_matrix(C) == m
407    True
408
409    """
410    outvec = Vector.zero
411    vects = system.base_vectors()
412    for i, x in enumerate(matrix):
413        outvec += x * vects[i]
414    return outvec
415
416
417def _path(from_object, to_object):
418    """
419    Calculates the 'path' of objects starting from 'from_object'
420    to 'to_object', along with the index of the first common
421    ancestor in the tree.
422
423    Returns (index, list) tuple.
424
425    """
426    if from_object._root != to_object._root:
427        raise ValueError('No connecting path found between ' +
428                         str(from_object) + ' and ' + str(to_object))
429
430    other_path = []
431    obj = to_object
432    while obj._parent is not None:
433        other_path.append(obj)
434        obj = obj._parent
435    other_path.append(obj)
436    object_set = set(other_path)
437    from_path = []
438    obj = from_object
439    while obj not in object_set:
440        from_path.append(obj)
441        obj = obj._parent
442    index = len(from_path)
443    i = other_path.index(obj)
444    while i >= 0:
445        from_path.append(other_path[i])
446        i -= 1
447    return index, from_path
448