<s> from __future__ import print_function , division <EOL> from sympy import Expr , Add , Mul , Pow , sympify , Matrix , Tuple <EOL> from sympy . core . compatibility import range <EOL> from sympy . utilities import default_sort_key <EOL> def _is_scalar ( e ) : <EOL> """<STR_LIT>""" <EOL> e = sympify ( e ) <EOL> if isinstance ( e , Expr ) : <EOL> if ( e . is_Integer or e . is_Float or <EOL> e . is_Rational or e . is_Number or <EOL> ( e . is_Symbol and e . is_commutative ) <EOL> ) : <EOL> return True <EOL> return False <EOL> def _cycle_permute ( l ) : <EOL> """<STR_LIT>""" <EOL> if len ( l ) == <NUM_LIT:1> : <EOL> return l <EOL> min_item = min ( l , key = default_sort_key ) <EOL> indices = [ i for i , x in enumerate ( l ) if x == min_item ] <EOL> le = list ( l ) <EOL> le . extend ( l ) <EOL> indices . append ( len ( l ) + indices [ <NUM_LIT:0> ] ) <EOL> sublist = [ [ le [ indices [ i ] : indices [ i + <NUM_LIT:1> ] ] ] for i in <EOL> range ( len ( indices ) - <NUM_LIT:1> ) ] <EOL> idx = sublist . index ( min ( sublist ) ) <EOL> ordered_l = le [ indices [ idx ] : indices [ idx ] + len ( l ) ] <EOL> return ordered_l <EOL> def _rearrange_args ( l ) : <EOL> """<STR_LIT>""" <EOL> if len ( l ) == <NUM_LIT:1> : <EOL> return l <EOL> x = list ( l [ - <NUM_LIT:1> : ] ) <EOL> x . extend ( l [ <NUM_LIT:0> : - <NUM_LIT:1> ] ) <EOL> return Mul ( * x ) . args <EOL> class Tr ( Expr ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( cls , * args ) : <EOL> """<STR_LIT>""" <EOL> if ( len ( args ) == <NUM_LIT:2> ) : <EOL> if not isinstance ( args [ <NUM_LIT:1> ] , ( list , Tuple , tuple ) ) : <EOL> indices = Tuple ( args [ <NUM_LIT:1> ] ) <EOL> else : <EOL> indices = Tuple ( * args [ <NUM_LIT:1> ] ) <EOL> expr = args [ <NUM_LIT:0> ] <EOL> elif ( len ( args ) == <NUM_LIT:1> ) : <EOL> indices = Tuple ( ) <EOL> expr = args [ <NUM_LIT:0> ] <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> if isinstance ( expr , Matrix ) : <EOL> return expr . trace ( ) <EOL> elif hasattr ( expr , '<STR_LIT>' ) and callable ( expr . trace ) : <EOL> return expr . trace ( ) <EOL> elif isinstance ( expr , Add ) : <EOL> return Add ( * [ Tr ( arg , indices ) for arg in expr . args ] ) <EOL> elif isinstance ( expr , Mul ) : <EOL> c_part , nc_part = expr . args_cnc ( ) <EOL> if len ( nc_part ) == <NUM_LIT:0> : <EOL> return Mul ( * c_part ) <EOL> else : <EOL> obj = Expr . __new__ ( cls , Mul ( * nc_part ) , indices ) <EOL> return Mul ( * c_part ) * obj if len ( c_part ) > <NUM_LIT:0> else obj <EOL> elif isinstance ( expr , Pow ) : <EOL> if ( _is_scalar ( expr . args [ <NUM_LIT:0> ] ) and <EOL> _is_scalar ( expr . args [ <NUM_LIT:1> ] ) ) : <EOL> return expr <EOL> else : <EOL> return Expr . __new__ ( cls , expr , indices ) <EOL> else : <EOL> if ( _is_scalar ( expr ) ) : <EOL> return expr <EOL> return Expr . __new__ ( cls , expr , indices ) <EOL> def doit ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if hasattr ( self . args [ <NUM_LIT:0> ] , '<STR_LIT>' ) : <EOL> return self . args [ <NUM_LIT:0> ] . _eval_trace ( indices = self . args [ <NUM_LIT:1> ] ) <EOL> return self <EOL> @ property <EOL> def is_number ( self ) : <EOL> return True <EOL> def permute ( self , pos ) : <EOL> """<STR_LIT>""" <EOL> if pos > <NUM_LIT:0> : <EOL> pos = pos % len ( self . args [ <NUM_LIT:0> ] . args ) <EOL> else : <EOL> pos = - ( abs ( pos ) % len ( self . args [ <NUM_LIT:0> ] . args ) ) <EOL> args = list ( self . args [ <NUM_LIT:0> ] . args [ - pos : ] + self . args [ <NUM_LIT:0> ] . args [ <NUM_LIT:0> : - pos ] ) <EOL> return Tr ( Mul ( * ( args ) ) ) <EOL> def _hashable_content ( self ) : <EOL> if isinstance ( self . args [ <NUM_LIT:0> ] , Mul ) : <EOL> args = _cycle_permute ( _rearrange_args ( self . args [ <NUM_LIT:0> ] . args ) ) <EOL> else : <EOL> args = [ self . args [ <NUM_LIT:0> ] ] <EOL> return tuple ( args ) + ( self . args [ <NUM_LIT:1> ] , ) </s>
