Coverage for /usr/lib/python3/dist-packages/sympy/matrices/expressions/trace.py: 27%

81 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1from sympy.core.basic import Basic 

2from sympy.core.expr import Expr, ExprBuilder 

3from sympy.core.singleton import S 

4from sympy.core.sorting import default_sort_key 

5from sympy.core.symbol import uniquely_named_symbol 

6from sympy.core.sympify import sympify 

7from sympy.matrices.matrices import MatrixBase 

8from sympy.matrices.common import NonSquareMatrixError 

9 

10 

11class Trace(Expr): 

12 """Matrix Trace 

13 

14 Represents the trace of a matrix expression. 

15 

16 Examples 

17 ======== 

18 

19 >>> from sympy import MatrixSymbol, Trace, eye 

20 >>> A = MatrixSymbol('A', 3, 3) 

21 >>> Trace(A) 

22 Trace(A) 

23 >>> Trace(eye(3)) 

24 Trace(Matrix([ 

25 [1, 0, 0], 

26 [0, 1, 0], 

27 [0, 0, 1]])) 

28 >>> Trace(eye(3)).simplify() 

29 3 

30 """ 

31 is_Trace = True 

32 is_commutative = True 

33 

34 def __new__(cls, mat): 

35 mat = sympify(mat) 

36 

37 if not mat.is_Matrix: 

38 raise TypeError("input to Trace, %s, is not a matrix" % str(mat)) 

39 

40 if mat.is_square is False: 

41 raise NonSquareMatrixError("Trace of a non-square matrix") 

42 

43 return Basic.__new__(cls, mat) 

44 

45 def _eval_transpose(self): 

46 return self 

47 

48 def _eval_derivative(self, v): 

49 from sympy.concrete.summations import Sum 

50 from .matexpr import MatrixElement 

51 if isinstance(v, MatrixElement): 

52 return self.rewrite(Sum).diff(v) 

53 expr = self.doit() 

54 if isinstance(expr, Trace): 

55 # Avoid looping infinitely: 

56 raise NotImplementedError 

57 return expr._eval_derivative(v) 

58 

59 def _eval_derivative_matrix_lines(self, x): 

60 from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, ArrayContraction 

61 r = self.args[0]._eval_derivative_matrix_lines(x) 

62 for lr in r: 

63 if lr.higher == 1: 

64 lr.higher = ExprBuilder( 

65 ArrayContraction, 

66 [ 

67 ExprBuilder( 

68 ArrayTensorProduct, 

69 [ 

70 lr._lines[0], 

71 lr._lines[1], 

72 ] 

73 ), 

74 (1, 3), 

75 ], 

76 validator=ArrayContraction._validate 

77 ) 

78 else: 

79 # This is not a matrix line: 

80 lr.higher = ExprBuilder( 

81 ArrayContraction, 

82 [ 

83 ExprBuilder( 

84 ArrayTensorProduct, 

85 [ 

86 lr._lines[0], 

87 lr._lines[1], 

88 lr.higher, 

89 ] 

90 ), 

91 (1, 3), (0, 2) 

92 ] 

93 ) 

94 lr._lines = [S.One, S.One] 

95 lr._first_pointer_parent = lr._lines 

96 lr._second_pointer_parent = lr._lines 

97 lr._first_pointer_index = 0 

98 lr._second_pointer_index = 1 

99 return r 

100 

101 @property 

102 def arg(self): 

103 return self.args[0] 

104 

105 def doit(self, **hints): 

106 if hints.get('deep', True): 

107 arg = self.arg.doit(**hints) 

108 try: 

109 return arg._eval_trace() 

110 except (AttributeError, NotImplementedError): 

111 return Trace(arg) 

112 else: 

113 # _eval_trace would go too deep here 

114 if isinstance(self.arg, MatrixBase): 

115 return trace(self.arg) 

116 else: 

117 return Trace(self.arg) 

118 

119 def as_explicit(self): 

120 return Trace(self.arg.as_explicit()).doit() 

121 

122 def _normalize(self): 

123 # Normalization of trace of matrix products. Use transposition and 

124 # cyclic properties of traces to make sure the arguments of the matrix 

125 # product are sorted and the first argument is not a transposition. 

126 from sympy.matrices.expressions.matmul import MatMul 

127 from sympy.matrices.expressions.transpose import Transpose 

128 trace_arg = self.arg 

129 if isinstance(trace_arg, MatMul): 

130 

131 def get_arg_key(x): 

132 a = trace_arg.args[x] 

133 if isinstance(a, Transpose): 

134 a = a.arg 

135 return default_sort_key(a) 

136 

137 indmin = min(range(len(trace_arg.args)), key=get_arg_key) 

138 if isinstance(trace_arg.args[indmin], Transpose): 

139 trace_arg = Transpose(trace_arg).doit() 

140 indmin = min(range(len(trace_arg.args)), key=lambda x: default_sort_key(trace_arg.args[x])) 

141 trace_arg = MatMul.fromiter(trace_arg.args[indmin:] + trace_arg.args[:indmin]) 

142 return Trace(trace_arg) 

143 return self 

144 

145 def _eval_rewrite_as_Sum(self, expr, **kwargs): 

146 from sympy.concrete.summations import Sum 

147 i = uniquely_named_symbol('i', expr) 

148 s = Sum(self.arg[i, i], (i, 0, self.arg.rows - 1)) 

149 return s.doit() 

150 

151 

152def trace(expr): 

153 """Trace of a Matrix. Sum of the diagonal elements. 

154 

155 Examples 

156 ======== 

157 

158 >>> from sympy import trace, Symbol, MatrixSymbol, eye 

159 >>> n = Symbol('n') 

160 >>> X = MatrixSymbol('X', n, n) # A square matrix 

161 >>> trace(2*X) 

162 2*Trace(X) 

163 >>> trace(eye(3)) 

164 3 

165 """ 

166 return Trace(expr).doit()