def __new__ (cls, arg, **options):
# print "Inverse(", arg, ")"
# print INVERTIBLE
if isinstance(arg, Inverse):
return arg.args[0]
if arg.is_Number:
return 1 / arg
arg_rank = expr_rank(arg)
if arg_rank == 1:
raise NotInvertibleError
if is_one(arg):
return arg
if is_zero(arg):
raise NotInvertibleError
# FIXME: Funky case trying to catch lower triangular or diagonal
# muls like T(P_0)*A*P_0
if arg in INVERTIBLE:
pass
elif isinstance(arg, TensorExpr) and not arg.has_inverse:
raise NotInvertibleError
elif isinstance(arg, Mul):
if arg.args[0] == S(-1):
return - Inverse(reduce(operator.mul, arg.args[1:]))
if not expr_invertible(arg):
raise NotInvertibleError
options['commutative'] = arg.is_commutative
return Basic.__new__(cls, arg, **options)
def __new__(cls, arg, **options):
if isinstance(arg, Transpose) and isinstance(arg.args[0], Tensor):
return arg.args[0]
if isinstance(arg, matrix):
new_arg = arg.copy().transpose()
for i in xrange(new_arg.shape[0]):
for j in xrange(new_arg.shape[1]):
new_arg[i, j] = Transpose(new_arg[i, j])
return new_arg
if isinstance(arg, Tensor):
if arg.rank == 0:
return arg
if arg.name == '0' or arg.name == 'I':
return arg
if arg.is_Number:
return arg
if isinstance(arg, Mul):
rank_0_objs = filter(lambda arg: expr_rank(arg) == 0, arg.args)
if len(rank_0_objs) > 0:
other_objs = filter(lambda arg: arg not in rank_0_objs,
arg.args)
return Mul(*rank_0_objs) * Transpose(Mul(*other_objs))
if arg.is_Symbol and not isinstance(arg, TensorExpr):
return arg
return Basic.__new__(cls, arg, **options)
def expr_invertible(expr):
if expr.is_Number:
return True
if isinstance(expr, TensorExpr) and expr.has_inverse:
return True
elif isinstance(expr, Pow):
return expr_invertible(expr.args[0])
if isinstance(expr, Inverse):
return True
if expr in INVERTIBLE:
return True
er = expr_rank(expr)
if er == 1:
return False
if is_one(expr):
return True
if is_zero(expr):
return False
# FIXME: Funky case trying to catch lower triangular or diagonal
# muls like T(P_0)*A*P_0
if isinstance(expr, Mul):
return reduce(lambda acc, x: acc and expr_invertible(x), expr.args,
True)
if isinstance(expr, Pow):
return expr_invertible(expr.args[0])
def _visit (expr):
if expr is None:
return ""
elif isinstance(expr, (Number, float, int)):
return str(expr)
elif isinstance(expr, Mul):
if expr.args[0] is S(-1):
return str_dict["neg"] % _wrap_parens(-1 * expr)
if len(expr.args) == 2:
return str_dict["dot"] % tuple(map(_wrap_parens, expr.args))
else:
return str_dict["dot"] % (_wrap_parens(expr.args[0]),
_wrap_parens(reduce(operator.mul,
expr.args[1:])))
elif isinstance(expr, Add):
return str_dict["add"].join(map(_visit, expr.args))
elif isinstance(expr, Pow):
if expr.args[1] == S(-1):
return str_dict["div_under_one"] % _wrap_parens(expr.args[0])
if expr_rank(expr.args[1]) == 0 or expr_rank(expr) == 0:
return str_dict["pow"] % tuple(map(_visit, expr.args))
else:
raise NotImplementedError("Unable to print: %s" % expr)
elif isinstance(expr, Inner):
return str_dict["dot"] % (_visit(expr.args[0]),
_visit(expr.args[1]))
elif isinstance(expr, Transpose):
if expr_rank(expr) == 0:
return _visit(expr.args[0])
else:
return str_dict["transpose"] % _wrap_parens(expr.args[0])
elif isinstance(expr, Inverse):
er = expr_rank(expr)
if er == 0:
return str_dict["div_under_one"] % _wrap_parens(expr.args[0])
elif er == 2:
return str_dict["inverse"] % _wrap_parens(expr.args[0])
else:
raise NotImplementedError
elif isinstance(expr, Tensor):
return expr.__getattribute__(str_dict["name_attr"])
else:
raise NotImplementedError
def _eval_expand_basic(self, deep=True, **hints):
arg0 = expand(self.args[0])
arg1 = expand(self.args[1])
arg0_adds = [arg0]
arg1_adds = [arg1]
if isinstance(arg0, Add):
arg0_adds = arg0.args
if isinstance(arg1, Add):
arg1_adds = arg1.args
add_exprs = []
for arg0 in arg0_adds:
for arg1 in arg1_adds:
add_exprs.append(Inner(arg0, arg1))
ret_add_exprs = []
for expr in add_exprs:
coeffs = []
arg0 = expr.args[0]
arg1 = expr.args[1]
if isinstance(arg0, Mul):
arg0_coeffs = filter(lambda arg: expr_rank(arg) == 0,
arg0.args)
if len(arg0_coeffs) != 0:
coeffs.extend(arg0_coeffs)
arg0 = Mul(
*filter(lambda arg: expr_rank(arg) != 0, arg0.args))
if isinstance(arg1, Mul):
arg1_coeffs = filter(lambda arg: expr_rank(arg) == 0,
arg1.args)
if len(arg1_coeffs) != 0:
coeffs.extend(arg1_coeffs)
arg1 = Mul(
*filter(lambda arg: expr_rank(arg) != 0, arg1.args))
if len(coeffs) == 0:
ret_add_exprs.append(expr)
else:
ret_add_exprs.append(Mul(*(coeffs + [Inner(arg0, arg1)])))
return Add(*ret_add_exprs)
def solve_vec_eqn(eqn, var):
"""Returns the solution to a linear equation containing Tensors
Raises:
NonLinearEqnError if the variable is detected to be nonlinear
NotInvertibleError if an inverse is required that is not available
NotImplementedError if operation isn't supported by routine
"""
if DEBUG:
print "solve_vec_eqn: ", eqn, "for", var
if var.rank != expr_rank(eqn):
raise ValueError("Unmatched ranks of clauses")
if expr_nonlinear(eqn, var):
raise NonLinearEqnError()
def _only_solve_numerator(expr):
return isinstance(expr, Mul) and \
len(expr.args) == 2 and isinstance(expr.args[1], Inverse) and \
expr.args[1].rank == 0 and var in expr.args[0]
def _solve_recur(expr, rhs=S(0)):
if expr == var:
return rhs
expr = expand(expr)
if isinstance(expr, Mul):
lhs = S(1)
# Try by rank
l_coeff, var_expr, r_coeff = expr_coeff(expr, var)
lhs = var_expr
rhs = Inverse(l_coeff) * rhs * Inverse(r_coeff)
return _solve_recur(lhs, rhs)
if isinstance(expr, Add):
lhs = 0
for arg in expr.args:
if var in arg:
lhs += arg
else:
rhs -= arg
if isinstance(lhs, Add):
coeff = lhs.coeff(var)
if expand(coeff * var) == lhs:
rhs /= coeff
lhs = var
return _solve_recur(lhs, rhs)
if isinstance(expr, Transpose):
return _solve_recur(expr.args[0], Transpose(rhs))
raise NotImplementedError("Can't handle expr of type %s" % type(expr))
# Check if expr is of the form (a + b) / c with rank(c) == 0
# then solve just the numerator
if _only_solve_numerator(eqn):
return _solve_recur(eqn.args[0])
return _solve_recur(eqn)
def _eval_expand_basic(self, deep=True, **hints):
arg0 = expand(self.args[0])
arg1 = expand(self.args[1])
arg0_adds = [arg0]
arg1_adds = [arg1]
if isinstance(arg0, Add):
arg0_adds = arg0.args
if isinstance(arg1, Add):
arg1_adds = arg1.args
add_exprs = []
for arg0 in arg0_adds:
for arg1 in arg1_adds:
add_exprs.append(Inner(arg0, arg1))
ret_add_exprs = []
for expr in add_exprs:
coeffs = []
arg0 = expr.args[0]
arg1 = expr.args[1]
if isinstance(arg0, Mul):
arg0_coeffs = filter(lambda arg: expr_rank(arg) == 0, arg0.args)
if len(arg0_coeffs) != 0:
coeffs.extend(arg0_coeffs)
arg0 = Mul(*filter(lambda arg: expr_rank(arg) != 0, arg0.args))
if isinstance(arg1, Mul):
arg1_coeffs = filter(lambda arg: expr_rank(arg) == 0, arg1.args)
if len(arg1_coeffs) != 0:
coeffs.extend(arg1_coeffs)
arg1 = Mul(*filter(lambda arg: expr_rank(arg) != 0, arg1.args))
if len(coeffs) == 0:
ret_add_exprs.append(expr)
else:
ret_add_exprs.append(Mul(*(coeffs + [Inner(arg0, arg1)])))
return Add(*ret_add_exprs)
def __new__(cls, arg, **options):
if isinstance(arg, Transpose) and isinstance(arg.args[0], Tensor):
return arg.args[0]
if isinstance(arg, matrix):
new_arg = arg.copy().transpose()
for i in xrange(new_arg.shape[0]):
for j in xrange(new_arg.shape[1]):
new_arg[i, j] = Transpose(new_arg[i, j])
return new_arg
if isinstance(arg, Tensor):
if arg.rank == 0:
return arg
if arg.name == '0' or arg.name == 'I':
return arg
if arg.is_Number:
return arg
if isinstance(arg, Mul):
rank_0_objs = filter(lambda arg: expr_rank(arg) == 0, arg.args)
if len(rank_0_objs) > 0:
other_objs = filter(lambda arg: arg not in rank_0_objs, arg.args)
return Mul(*rank_0_objs) * Transpose(Mul(*other_objs))
if arg.is_Symbol and not isinstance(arg, TensorExpr):
return arg
return Basic.__new__(cls, arg, **options)
def expr_invertible(expr):
if expr.is_Number:
return True
if isinstance(expr, TensorExpr) and expr.has_inverse:
return True
elif isinstance(expr, Pow):
return expr_invertible(expr.args[0])
if isinstance(expr, Inverse):
return True
if expr in INVERTIBLE:
return True
er = expr_rank(expr)
if er == 1:
return False
if is_one(expr):
return True
if is_zero(expr):
return False
# FIXME: Funky case trying to catch lower triangular or diagonal
# muls like T(P_0)*A*P_0
if isinstance(expr, Mul):
return reduce(lambda acc, x: acc and expr_invertible(x), expr.args, True)
if isinstance(expr, Pow):
return expr_invertible(expr.args[0])
def rank(self):
return expr_rank(self.args[0])