def test_doit(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
assert VecDot(vn1, vn2).doit() - vn1.dot(vn2) == 0
assert VecDot(nabla, vn2).doit() - divergence(vn2) == 0
assert isinstance(VecDot(v1, v1).doit(), VecPow)
assert isinstance(VecDot(vn1, vn1).doit(), Integer)
expr = VecDot(v1, v2.norm).doit()
assert tuple(type(a) for a in expr.args) in [(VectorSymbol, VecMul),
(VecMul, VectorSymbol)]
expr = VecDot(v1, v2.norm).doit(deep=False)
assert tuple(type(a) for a in expr.args) in [(VectorSymbol, Normalize),
(Normalize, VectorSymbol)]
expr = VecDot(v1, VecAdd(vn2, vn1)).doit(deep=False)
assert tuple(type(a) for a in expr.args) in [(VectorSymbol, VecAdd),
(VecAdd, VectorSymbol)]
expr = VecDot(v1, VecAdd(vn2, vn1)).doit()
assert tuple(type(a) for a in expr.args) in [(VectorSymbol, VectorAdd),
(VectorAdd, VectorSymbol)]
expr = VecDot(vn1, vn2 + vn1).doit()
assert expr == (x + 2 * y + 3 * z + 14)
def test_doit(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
assert isinstance(Advection(v1, v2).doit(), Advection)
# identity test
a = C.x * C.i + C.y * C.j + C.z * C.k
b = C.y * C.z * C.i + C.x * C.z * C.j + C.x * C.y * C.k
assert self._check_args((nabla ^ (a ^ b)).doit(),
(divergence(b) * a) + Advection(b, a).doit() -
(divergence(a) * b) - Advection(a, b).doit())
assert isinstance(
Advection(VecAdd(a, b), C.x).doit(deep=False), Advection)
assert isinstance(
Advection(VecAdd(a, b), VecAdd(a, b)).doit(deep=False), Advection)
def test_doit(self):
v1, v2, zero, one, nabla, C, vn1, vn2 = self._get_vars()
assert Magnitude(VecAdd(v1, zero, evaluate=False)).doit() == v1.mag
assert Magnitude(vn1).doit() == sqrt(14)
assert Magnitude(vn2).doit() == sqrt(x**2 + y**2 + z**2)
assert isinstance(
Magnitude(VecCross(vn1, vn2)).doit(deep=False), Magnitude)
assert isinstance(Magnitude(VecCross(vn1, vn2)).doit(), Pow)
def collect_cross_dot(expr, pattern=VecCross, first_arg=True):
""" Collect additive dot/cross products with common arguments.
Example: expr = (a ^ b) + (a ^ c) + (b ^ c)
collect_cross_dot(expr)
(a ^ (b + c)) + (b ^ c)
collect_cross_dot(expr, VecCross, False)
(a ^ b) + ((a + b) ^ c)
Parameters
----------
expr : the expression to process
pattern : the class to look for. Can be VecCross or VecDot.
first_arg : Boolean, default to True. If True, look for common first
arguments in the instances of the class `pattern`. Otherwise, look
for common second arguments.
"""
if not issubclass(pattern, DotCross):
return expr
copy = expr
subs_list = []
for arg in preorder_traversal(expr):
# we are interested to the args of instance pattern.func in the
# current tree level
found = [a for a in arg.args if isinstance(a, pattern)]
terms = _terms_with_commong_args(found, first_arg)
if first_arg:
get_v = lambda t: pattern(t[0].args[0],
VecAdd(*[a.args[1] for a in t]))
else:
get_v = lambda t: pattern(VecAdd(*[a.args[0]
for a in t]), t[0].args[1])
if terms:
for t in terms:
subs_list.append({VecAdd(*t): get_v(t)})
for s in subs_list:
expr = expr.subs(s)
if copy == expr:
return expr
return collect_cross_dot(expr, pattern, first_arg)
def test_doit(self):
v1, v2, zero, one, nabla, C, vn1, vn2 = self._get_vars()
# return a VecMul object (a fraction, vector/magnitude)
assert isinstance(Normalize(v1).doit(), VecMul)
assert isinstance(Normalize(one).doit(), VecMul)
assert isinstance(Normalize(zero).doit(), VectorZero)
# symbolic Vector
assert Normalize(vn1).doit() - vn1.normalize() == Vector.zero
assert Normalize(vn2).doit() - vn2.normalize() == Vector.zero
assert Normalize(VecAdd(vn1, vn2)).doit() - (vn1 + vn2).normalize() == Vector.zero
assert isinstance(Normalize(VecAdd(vn1, vn2)).doit(deep=False), VecMul)
# test to check if the result is a vector
assert Normalize(v1).doit().is_Vector
assert Normalize(v1 + v2).doit().is_Vector
assert Normalize(one).doit().is_Vector
assert Normalize(zero).doit().is_Vector
assert Normalize(vn1).doit().is_Vector
assert Normalize(VecAdd(vn1, vn2)).doit().is_Vector
assert Normalize(VecAdd(vn1, vn2)).doit(deep=False).is_Vector
def test_creation(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
assert isinstance(VectorSymbol("v1"), VectorSymbol)
assert isinstance(VectorSymbol(Symbol("x")), VectorSymbol)
assert VectorSymbol(v1) == v1
def func(s):
with self.assertRaises(TypeError):
VectorSymbol(s)
func(x + y)
func(VecAdd(v1, v2))
func(Symbol("x") + Symbol("y"))
def test_is_vector(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
assert VecAdd(v1, v2).is_Vector
# assert not VecAdd(v1, v2).is_Vector_Scalar
assert not VecAdd(v1.mag, v2.mag).is_Vector
# assert VecAdd(v1.mag, v2.mag).is_Vector_Scalar
# with dot product and nested mul/dot
assert not VecAdd(2, v2 & v1).is_Vector
# assert VecAdd(2, v2 & v1).is_Vector_Scalar
assert not VecAdd(2, VecMul(x, v2 & v1)).is_Vector
# assert VecAdd(2, VecMul(x, v2 & v1)).is_Vector_Scalar
# with cross product and nested mul/cross
assert VecAdd(v1, v2 ^ v1).is_Vector
# assert not VecAdd(v1, v2 ^ v1).is_Vector_Scalar
assert VecAdd(v1, VecMul(x, v2 ^ v1)).is_Vector
def test_doit(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
expr = VecAdd(v1, v2, v1, evaluate=False)
assert not self._check_args(expr, VecAdd(VecMul(2, v1), v2))
assert self._check_args(expr.doit(), VecAdd(VecMul(2, v1), v2))
expr = VecAdd(v1, v2, VecAdd(v1, v2, evaluate=False), evaluate=False)
assert not self._check_args(expr, VecAdd(VecMul(2, v1), VecMul(2, v2)))
assert self._check_args(expr.doit(),
VecAdd(VecMul(2, v1), VecMul(2, v2)))
r = vn1 + vn2
assert VecAdd(vn1, vn2).doit() == r
expr = VecAdd(v1.mag, VecDot(vn1, vn2))
assert self._check_args(expr.doit(deep=False), expr)
assert self._check_args(expr.doit(), VecAdd(v1.mag, vn1 & vn2))
# test the rule default_sort_key used into VecAdd.doit
assert VecAdd(v1, one, v2).doit().args == (one, v1, v2)
assert VecAdd(v1, VecCross(v1, v2),
one).doit().args == (VecCross(v1, v2), one, v1)
# test the rule merge_explicit used into VecAdd.doit
assert VecAdd(vn1, vn2).doit() == vn1 + vn2
assert self._check_args(
VecAdd(vn1, vn2, one).doit(), VecAdd(one, vn1 + vn2))
def test_flatten(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
# test sympy.strategies.flatten applied to VecAdd
assert self._check_args(VecAdd(v1, VecAdd(v2, one)),
VecAdd(one, v1, v2))
assert self._check_args(VecAdd(v1, VecAdd(one), VecAdd(zero, v2)),
VecAdd(one, v1, v2))
assert self._check_args(VecAdd(v1, VecAdd(one, VecAdd(zero, v2))),
VecAdd(one, v1, v2))
assert self._check_args(
VecAdd(v1 ^ v2, VecAdd(one, VecAdd(2 * v1, (v2 & v1) * one))),
VecAdd(v1 ^ v2, one, 2 * v1, (v2 & v1) * one))
def func(args, **kwargs):
with self.assertRaises(TypeError) as context:
VecAdd(*args, **kwargs)
self.assertTrue(
'Mix of Vector and Scalar symbols' in str(context.exception))
def test_creation(self):
v1, v2, zero, one, nabla, C, vn1, vn2, x, y, z = self._get_vars()
# no args
assert VecAdd() == VectorZero()
# 1 arg
assert VecAdd(v1.mag) == v1.mag
assert VecAdd(v2) == v2
assert VecAdd(zero) == zero
assert VecAdd(0) == S.Zero
# one argument is 0 or VectorZero()
assert VecAdd(one, zero) == one
assert VecAdd(VectorOne(), VectorZero()) == one
assert VecAdd(0, x) == x
assert VecAdd(zero, v1) == v1
assert VecAdd(v1.mag, 0) == v1.mag
# no VectorExpr in arguments -> return Add
assert isinstance(VecAdd(2, x, y), Add)
# VectorExpr in args -> return VecAdd
assert isinstance(VecAdd(v1, v2), VecAdd)
assert isinstance(VecAdd(vn1, vn2), VecAdd)
assert isinstance(VecAdd(v1.mag, v2.mag), VecAdd)
assert isinstance(VecAdd(1, v2.mag), VecAdd)
assert isinstance(VecAdd(one, v2 ^ v1), VecAdd)
assert isinstance(VecAdd(1, v2 & v1), VecAdd)
assert isinstance(VecAdd(one, v1, v2), VecAdd)
assert isinstance(VecAdd(v1, v1), VecMul)
assert isinstance(VecAdd(v1, v1, evaluate=False), VecAdd)
def collect(expr, match):
""" Implement a custom collecting algorithm to collect dot and cross
products. There is a noticeable difference with expr.collect: given the
following expression:
expr = a * t + b * t + c * t**2 + d * t**3 + e * t**3
expr.collect(t)
t**3 * (d + e) + t**2 * c + t * (a + b)
Whereas:
collect(expr, t)
t * (a + b + t * c + t**2 * d + t**2 * e)
collect(expr, t**2)
a * t + b * t + t**2 * (c + t * d + t * e)
Parameters
----------
expr : the expression to process
match : the pattern to collect. If match is not an instance of VecDot,
VecCross or VecPow, standard expr.collect(match) will be used.
"""
# the method expr.collect doesn't know how to treat cross and dot products.
# Since dot-product is a scalar, it can also be raised to a power. We need
# to select which algorithm to run.
if not isinstance(match, (VecDot, VecCross, VecPow)):
return expr.collect(match)
# extract the pattern and exponent from the power term
pattern = match
n = 1
if isinstance(match, VecPow):
pattern = match.base
n = match.exp
if not expr.has(pattern):
return expr
if not isinstance(expr, VecAdd):
return expr
collect = []
not_collect = []
for arg in expr.args:
if not arg.has(pattern):
not_collect.append(arg)
else:
p = [
a for a in arg.args
if isinstance(a, VecPow) and isinstance(a.base, pattern.func)
]
if p:
# vecpow(pattern, exp)
if p[0].exp >= n:
collect.append(arg / match)
else:
not_collect.append(arg)
else:
if n == 1:
if arg != match:
term = arg.func(*[a for a in arg.args if a != match])
collect.append(term)
elif match.is_Vector:
collect.append(1)
else:
# TODO. need to test this!!!
collect.append(VectorOne())
else:
not_collect.append(arg)
return VecAdd(VecMul(match, VecAdd(*collect)), *not_collect)
def collect_const(expr, *vars, **kwargs):
""" This is the very same code of sympy.simplify.radsimp.py collect_const,
with modification: the original method used Mul._from_args
and Add._from_args, which do not call a post-processor, hence I obtained the
wrong result. Here, I use VecAdd, VecMul...
"""
if not expr.is_Add:
return expr
recurse = False
Numbers = kwargs.get('Numbers', True)
if not vars:
recurse = True
vars = set()
for a in expr.args:
for m in Mul.make_args(a):
if m.is_number:
vars.add(m)
else:
vars = sympify(vars)
if not Numbers:
vars = [v for v in vars if not v.is_Number]
vars = list(ordered(vars))
for v in vars:
terms = defaultdict(list)
Fv = Factors(v)
for m in Add.make_args(expr):
f = Factors(m)
q, r = f.div(Fv)
if r.is_one:
# only accept this as a true factor if
# it didn't change an exponent from an Integer
# to a non-Integer, e.g. 2/sqrt(2) -> sqrt(2)
# -- we aren't looking for this sort of change
fwas = f.factors.copy()
fnow = q.factors
if not any(k in fwas and fwas[k].is_Integer
and not fnow[k].is_Integer for k in fnow):
terms[v].append(q.as_expr())
continue
terms[S.One].append(m)
args = []
hit = False
uneval = False
for k in ordered(terms):
v = terms[k]
if k is S.One:
args.extend(v)
continue
if len(v) > 1:
v = Add(*v)
hit = True
if recurse and v != expr:
vars.append(v)
else:
v = v[0]
# be careful not to let uneval become True unless
# it must be because it's going to be more expensive
# to rebuild the expression as an unevaluated one
if Numbers and k.is_Number and v.is_Add:
# args.append(_keep_coeff(k, v, sign=True))
args.append(VecMul(*[k, v], evaluate=False))
uneval = True
else:
args.append(k * v)
if hit:
if uneval:
# expr = Add(*args)
expr = VecAdd(*args, evaluate=False)
else:
# expr = Add(*args)
expr = VecAdd(*args)
if not expr.is_Add:
break
return expr
def find_bac_cab(expr):
""" Given a vector expression, find the terms satisfying the pattern:
B * (A & C) - C * (A & B)
where & is the dot product. The list is ordered in such a way that
nested matches comes first (similarly to scanning the expression tree from
bottom to top).
"""
def _check_pairs(terms):
# c1 = Mul(*[a for a in terms[0].args if not hasattr(a, "is_Vector_Scalar")])
# c2 = Mul(*[a for a in terms[1].args if not hasattr(a, "is_Vector_Scalar")])
c1 = Mul(*[
a for a in terms[0].args if not isinstance(a, (Vector, VectorExpr))
])
c2 = Mul(*[
a for a in terms[1].args if not isinstance(a, (Vector, VectorExpr))
])
n1, _ = c1.as_coeff_mul()
n2, _ = c2.as_coeff_mul()
if n1 * n2 > 0:
# opposite sign
return False
if Abs(c1) != Abs(c2):
return False
# v1 = [a for a in terms[0].args if (hasattr(a, "is_Vector_Scalar") and a.is_Vector)][0]
# v2 = [a for a in terms[1].args if (hasattr(a, "is_Vector_Scalar") and a.is_Vector)][0]
v1 = [
a for a in terms[0].args
if (isinstance(a, (Vector, VectorExpr)) and a.is_Vector)
][0]
v2 = [
a for a in terms[1].args
if (isinstance(a, (Vector, VectorExpr)) and a.is_Vector)
][0]
if v1 == v2:
return False
dot1 = [a for a in terms[0].args if isinstance(a, VecDot)][0]
dot2 = [a for a in terms[1].args if isinstance(a, VecDot)][0]
if not ((v1 in dot2.args) and (v2 in dot1.args)):
return False
v = list(set(dot1.args).intersection(set(dot2.args)))[0]
if v == v1 or v == v2:
return False
return True
def _check(arg):
if (isinstance(arg, VecMul) and any([a.is_Vector for a in arg.args])
and any([isinstance(a, VecDot) for a in arg.args])):
return True
return False
found = set()
for arg in preorder_traversal(expr):
possible_args = list(filter(_check, arg.args))
# print("possible_args", possible_args)
# for p in possible_args:
# print("\t", p.func, p.is_Vector, p)
# # possible_args = list(arg.find(A * (B & C)))
# # possible_args = list(filter(lambda x: isinstance(x, VecMul), possible_args))
if len(possible_args) > 1:
combs = list(combinations(possible_args, 2))
# print("COMBINATIONS", combs)
for c in combs:
# print("\tC", c)
# print("\n\t\t".join(str(a.func) + ", " + str(a.is_Vector) + ", " + str(a) for a in c))
if _check_pairs(c):
# print("\t\t proceeding")
found.add(VecAdd(*c))
found = list(ordered(list(found)))
return found