def expand(self,
deep=True,
modulus=None,
power_base=True,
power_exp=True,
mul=True,
log=True,
multinomial=True,
basic=True,
**hints):
if isinstance(self, VectorSymbol):
return self
_spaces.append(_token)
debug("VectorExpr expand", self)
# first expand dot and cross products
A, B = [WVS(t) for t in ["A", "B"]]
def func(expr, pattern, prev=None):
# debug("\n", hints)
old = expr
found = list(ordered(list(expr.find(pattern))))
# print("func", expr)
# debug("found", found)
# print("found", found)
for f in found:
fexp = f.expand(**hints)
# debug("\t fexp", fexp)
if f != fexp:
expr = expr.xreplace({f: fexp})
# debug("expr", expr)
if old != expr:
expr = func(expr, pattern)
return expr
expr = func(self, A ^ B)
# print("expr", expr)
expr = func(expr, A & B)
# print("expr", expr)
expr = func(expr, Grad)
expr = func(expr, Laplace)
# print("expr", expr)
del _spaces[-1]
return Expr.expand(expr,
deep=deep,
modulus=modulus,
power_base=power_base,
power_exp=power_exp,
mul=mul,
log=log,
multinomial=multinomial,
basic=basic,
**hints)
def remove_higher_order_terms(expr: sp.Expr,
symbols: Sequence[sp.Symbol],
order: int = 3) -> sp.Expr:
"""Removes all terms that contain more than 'order' factors of given 'symbols'
Example:
>>> x, y = sp.symbols("x y")
>>> term = x**2 * y + y**2 * x + y**3 + x + y ** 2
>>> remove_higher_order_terms(term, order=2, symbols=[x, y])
x + y**2
"""
from sympy.core.power import Pow
from sympy.core.add import Add, Mul
result = 0
expr = expr.expand()
def velocity_factors_in_product(product):
factor_count = 0
if type(product) is Mul:
for factor in product.args:
if type(factor) == Pow:
if factor.args[0] in symbols:
factor_count += factor.args[1]
if factor in symbols:
factor_count += 1
elif type(product) is Pow:
if product.args[0] in symbols:
factor_count += product.args[1]
return factor_count
if type(expr) == Mul or type(expr) == Pow:
if velocity_factors_in_product(expr) <= order:
return expr
else:
return sp.Rational(0, 1)
if type(expr) != Add:
return expr
for sum_term in expr.args:
if velocity_factors_in_product(sum_term) <= order:
result += sum_term
return result
def expand(self, **hints):
tp = TensorProduct(*[sympify(item).expand(**hints) for item in self.args])
return Expr.expand(tp, **hints)
def extract_coefficients(equation: sympy.Expr, local_map: dict,
global_coords: list) -> tuple:
"""
Args:
equation: The equation in local coordinates.
local_map: The mapping from local coordinates to the index of a
global coordinate.
global_coords: The list of global co-ordinates.
Returns:
The linear and nonlinear parts of the equation in the global
co-ordinate system.
Extracts the coordinates from the given equation and maps them into
the global coordinate space.
Equations are assumed to come in as sympy expressions of the form
:math:`\Phi(x) = 0`.
local_map is a dictionary mappings
.. math::
M: \\rightarrow i
where :math:`x` are the local co-ordinates and the keys of local_map, and
the values are the indices :math:`i` such that `global_coord[i]` is the
corresponding global coordinate. The result is :math:`L,N` such that:
.. math::
Ly + N(y) = 0
"""
coeff_dict = {}
nonlinear_terms = sympy.S(0)
subs = [(k, global_coords[v]) for k, v in local_map.items()]
subs.sort(key=lambda x: str(x[1])[-1], reverse=True)
logger.debug("Extracting coefficients from %s", repr(equation))
logger.debug("Using local-to-global substitutions %s", repr(subs))
terms = equation.expand().args
if not terms:
if equation in local_map:
coeff_dict[local_map[equation]] = sympy.S(1)
else:
nonlinear_terms = equation
else:
for term in terms:
factors = list(flatten(term.as_coeff_mul()))
coeff = sympy.S(1)
base = []
while factors:
factor = factors.pop()
if factor.is_number:
coeff *= factor
elif factor.is_symbol and factor not in local_map:
coeff *= factor
else:
base.append(factor)
if len(base) == 1 and base[0] in local_map:
coeff_dict[local_map[base[0]]] = coeff
else:
new_term = term
new_term = new_term.subs(subs)
nonlinear_terms = sympy.Add(new_term, nonlinear_terms)
logger.debug("Linear terms: %s", repr(coeff_dict))
logger.debug("Nonlinear terms: %s", repr(nonlinear_terms))
return coeff_dict, nonlinear_terms
def expand(self, **hints):
tp = TensorProduct(
*[sympify(item).expand(**hints) for item in self.args])
return Expr.expand(tp, **hints)
