<s> """<STR_LIT>""" <EOL> from __future__ import print_function , division <EOL> from sympy import ( Add , binomial , cos , exp , Expr , factorial , I , Integer , Mul , <EOL> pi , Rational , S , sin , simplify , sqrt , Sum , symbols , sympify , <EOL> Tuple , Dummy ) <EOL> from sympy . core . compatibility import u , unicode , range <EOL> from sympy . matrices import zeros <EOL> from sympy . printing . pretty . stringpict import prettyForm , stringPict <EOL> from sympy . printing . pretty . pretty_symbology import pretty_symbol <EOL> from sympy . physics . quantum . qexpr import QExpr <EOL> from sympy . physics . quantum . operator import ( HermitianOperator , Operator , <EOL> UnitaryOperator ) <EOL> from sympy . physics . quantum . state import Bra , Ket , State <EOL> from sympy . functions . special . tensor_functions import KroneckerDelta <EOL> from sympy . physics . quantum . constants import hbar <EOL> from sympy . physics . quantum . hilbert import ComplexSpace , DirectSumHilbertSpace <EOL> from sympy . physics . quantum . tensorproduct import TensorProduct <EOL> from sympy . physics . quantum . cg import CG <EOL> from sympy . physics . quantum . qapply import qapply <EOL> __all__ = [ <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' <EOL> ] <EOL> def m_values ( j ) : <EOL> j = sympify ( j ) <EOL> size = <NUM_LIT:2> * j + <NUM_LIT:1> <EOL> if not size . is_Integer or not size > <NUM_LIT:0> : <EOL> raise ValueError ( <EOL> '<STR_LIT>' % j <EOL> ) <EOL> return size , [ j - i for i in range ( int ( <NUM_LIT:2> * j + <NUM_LIT:1> ) ) ] <EOL> class SpinOpBase ( object ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def _eval_hilbert_space ( cls , label ) : <EOL> return ComplexSpace ( S . Infinity ) <EOL> @ property <EOL> def name ( self ) : <EOL> return self . args [ <NUM_LIT:0> ] <EOL> def _print_contents ( self , printer , * args ) : <EOL> return '<STR_LIT>' % ( unicode ( self . name ) , self . _coord ) <EOL> def _print_contents_pretty ( self , printer , * args ) : <EOL> a = stringPict ( unicode ( self . name ) ) <EOL> b = stringPict ( self . _coord ) <EOL> return self . _print_subscript_pretty ( a , b ) <EOL> def _print_contents_latex ( self , printer , * args ) : <EOL> return r'<STR_LIT>' % ( ( unicode ( self . name ) , self . _coord ) ) <EOL> def _represent_base ( self , basis , ** options ) : <EOL> j = options . get ( '<STR_LIT>' , Rational ( <NUM_LIT:1> , <NUM_LIT:2> ) ) <EOL> size , mvals = m_values ( j ) <EOL> result = zeros ( size , size ) <EOL> for p in range ( size ) : <EOL> for q in range ( size ) : <EOL> me = self . matrix_element ( j , mvals [ p ] , j , mvals [ q ] ) <EOL> result [ p , q ] = me <EOL> return result <EOL> def _apply_op ( self , ket , orig_basis , ** options ) : <EOL> state = ket . rewrite ( self . basis ) <EOL> if isinstance ( state , State ) : <EOL> ret = ( hbar * state . m ) * state <EOL> elif isinstance ( state , Sum ) : <EOL> ret = self . _apply_operator_Sum ( state , ** options ) <EOL> else : <EOL> ret = qapply ( self * state ) <EOL> if ret == self * state : <EOL> raise NotImplementedError <EOL> return ret . rewrite ( orig_basis ) <EOL> def _apply_operator_JxKet ( self , ket , ** options ) : <EOL> return self . _apply_op ( ket , '<STR_LIT>' , ** options ) <EOL> def _apply_operator_JxKetCoupled ( self , ket , ** options ) : <EOL> return self . _apply_op ( ket , '<STR_LIT>' , ** options ) <EOL> def _apply_operator_JyKet ( self , ket , ** options ) : <EOL> return self . _apply_op ( ket , '<STR_LIT>' , ** options ) <EOL> def _apply_operator_JyKetCoupled ( self , ket , ** options ) : <EOL> return self . _apply_op ( ket , '<STR_LIT>' , ** options ) <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> return self . _apply_op ( ket , '<STR_LIT>' , ** options ) <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> return self . _apply_op ( ket , '<STR_LIT>' , ** options ) <EOL> def _apply_operator_TensorProduct ( self , tp , ** options ) : <EOL> if not ( isinstance ( self , JxOp ) or isinstance ( self , JyOp ) or isinstance ( self , JzOp ) ) : <EOL> raise NotImplementedError <EOL> result = [ ] <EOL> for n in range ( len ( tp . args ) ) : <EOL> arg = [ ] <EOL> arg . extend ( tp . args [ : n ] ) <EOL> arg . append ( self . _apply_operator ( tp . args [ n ] ) ) <EOL> arg . extend ( tp . args [ n + <NUM_LIT:1> : ] ) <EOL> result . append ( tp . __class__ ( * arg ) ) <EOL> return Add ( * result ) . expand ( ) <EOL> def _apply_operator_Sum ( self , s , ** options ) : <EOL> new_func = qapply ( self * s . function ) <EOL> if new_func == self * s . function : <EOL> raise NotImplementedError <EOL> return Sum ( new_func , * s . limits ) <EOL> def _eval_trace ( self , ** options ) : <EOL> return self . _represent_default_basis ( ) . trace ( ) <EOL> class JplusOp ( SpinOpBase , Operator ) : <EOL> """<STR_LIT>""" <EOL> _coord = '<STR_LIT:+>' <EOL> basis = '<STR_LIT>' <EOL> def _eval_commutator_JminusOp ( self , other ) : <EOL> return <NUM_LIT:2> * hbar * JzOp ( self . name ) <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> j = ket . j <EOL> m = ket . m <EOL> if m . is_Number and j . is_Number : <EOL> if m >= j : <EOL> return S . Zero <EOL> return hbar * sqrt ( j * ( j + S . One ) - m * ( m + S . One ) ) * JzKet ( j , m + S . One ) <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> j = ket . j <EOL> m = ket . m <EOL> jn = ket . jn <EOL> coupling = ket . coupling <EOL> if m . is_Number and j . is_Number : <EOL> if m >= j : <EOL> return S . Zero <EOL> return hbar * sqrt ( j * ( j + S . One ) - m * ( m + S . One ) ) * JzKetCoupled ( j , m + S . One , jn , coupling ) <EOL> def matrix_element ( self , j , m , jp , mp ) : <EOL> result = hbar * sqrt ( j * ( j + S . One ) - mp * ( mp + S . One ) ) <EOL> result *= KroneckerDelta ( m , mp + <NUM_LIT:1> ) <EOL> result *= KroneckerDelta ( j , jp ) <EOL> return result <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( basis , ** options ) <EOL> def _eval_rewrite_as_xyz ( self , * args ) : <EOL> return JxOp ( args [ <NUM_LIT:0> ] ) + I * JyOp ( args [ <NUM_LIT:0> ] ) <EOL> class JminusOp ( SpinOpBase , Operator ) : <EOL> """<STR_LIT>""" <EOL> _coord = '<STR_LIT:->' <EOL> basis = '<STR_LIT>' <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> j = ket . j <EOL> m = ket . m <EOL> if m . is_Number and j . is_Number : <EOL> if m <= - j : <EOL> return S . Zero <EOL> return hbar * sqrt ( j * ( j + S . One ) - m * ( m - S . One ) ) * JzKet ( j , m - S . One ) <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> j = ket . j <EOL> m = ket . m <EOL> jn = ket . jn <EOL> coupling = ket . coupling <EOL> if m . is_Number and j . is_Number : <EOL> if m <= - j : <EOL> return S . Zero <EOL> return hbar * sqrt ( j * ( j + S . One ) - m * ( m - S . One ) ) * JzKetCoupled ( j , m - S . One , jn , coupling ) <EOL> def matrix_element ( self , j , m , jp , mp ) : <EOL> result = hbar * sqrt ( j * ( j + S . One ) - mp * ( mp - S . One ) ) <EOL> result *= KroneckerDelta ( m , mp - <NUM_LIT:1> ) <EOL> result *= KroneckerDelta ( j , jp ) <EOL> return result <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( basis , ** options ) <EOL> def _eval_rewrite_as_xyz ( self , * args ) : <EOL> return JxOp ( args [ <NUM_LIT:0> ] ) - I * JyOp ( args [ <NUM_LIT:0> ] ) <EOL> class JxOp ( SpinOpBase , HermitianOperator ) : <EOL> """<STR_LIT>""" <EOL> _coord = '<STR_LIT:x>' <EOL> basis = '<STR_LIT>' <EOL> def _eval_commutator_JyOp ( self , other ) : <EOL> return I * hbar * JzOp ( self . name ) <EOL> def _eval_commutator_JzOp ( self , other ) : <EOL> return - I * hbar * JyOp ( self . name ) <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> jp = JplusOp ( self . name ) . _apply_operator_JzKet ( ket , ** options ) <EOL> jm = JminusOp ( self . name ) . _apply_operator_JzKet ( ket , ** options ) <EOL> return ( jp + jm ) / Integer ( <NUM_LIT:2> ) <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> jp = JplusOp ( self . name ) . _apply_operator_JzKetCoupled ( ket , ** options ) <EOL> jm = JminusOp ( self . name ) . _apply_operator_JzKetCoupled ( ket , ** options ) <EOL> return ( jp + jm ) / Integer ( <NUM_LIT:2> ) <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> jp = JplusOp ( self . name ) . _represent_JzOp ( basis , ** options ) <EOL> jm = JminusOp ( self . name ) . _represent_JzOp ( basis , ** options ) <EOL> return ( jp + jm ) / Integer ( <NUM_LIT:2> ) <EOL> def _eval_rewrite_as_plusminus ( self , * args ) : <EOL> return ( JplusOp ( args [ <NUM_LIT:0> ] ) + JminusOp ( args [ <NUM_LIT:0> ] ) ) / <NUM_LIT:2> <EOL> class JyOp ( SpinOpBase , HermitianOperator ) : <EOL> """<STR_LIT>""" <EOL> _coord = '<STR_LIT:y>' <EOL> basis = '<STR_LIT>' <EOL> def _eval_commutator_JzOp ( self , other ) : <EOL> return I * hbar * JxOp ( self . name ) <EOL> def _eval_commutator_JxOp ( self , other ) : <EOL> return - I * hbar * J2Op ( self . name ) <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> jp = JplusOp ( self . name ) . _apply_operator_JzKet ( ket , ** options ) <EOL> jm = JminusOp ( self . name ) . _apply_operator_JzKet ( ket , ** options ) <EOL> return ( jp - jm ) / ( Integer ( <NUM_LIT:2> ) * I ) <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> jp = JplusOp ( self . name ) . _apply_operator_JzKetCoupled ( ket , ** options ) <EOL> jm = JminusOp ( self . name ) . _apply_operator_JzKetCoupled ( ket , ** options ) <EOL> return ( jp - jm ) / ( Integer ( <NUM_LIT:2> ) * I ) <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> jp = JplusOp ( self . name ) . _represent_JzOp ( basis , ** options ) <EOL> jm = JminusOp ( self . name ) . _represent_JzOp ( basis , ** options ) <EOL> return ( jp - jm ) / ( Integer ( <NUM_LIT:2> ) * I ) <EOL> def _eval_rewrite_as_plusminus ( self , * args ) : <EOL> return ( JplusOp ( args [ <NUM_LIT:0> ] ) - JminusOp ( args [ <NUM_LIT:0> ] ) ) / ( <NUM_LIT:2> * I ) <EOL> class JzOp ( SpinOpBase , HermitianOperator ) : <EOL> """<STR_LIT>""" <EOL> _coord = '<STR_LIT:z>' <EOL> basis = '<STR_LIT>' <EOL> def _eval_commutator_JxOp ( self , other ) : <EOL> return I * hbar * JyOp ( self . name ) <EOL> def _eval_commutator_JyOp ( self , other ) : <EOL> return - I * hbar * JxOp ( self . name ) <EOL> def _eval_commutator_JplusOp ( self , other ) : <EOL> return hbar * JplusOp ( self . name ) <EOL> def _eval_commutator_JminusOp ( self , other ) : <EOL> return - hbar * JminusOp ( self . name ) <EOL> def matrix_element ( self , j , m , jp , mp ) : <EOL> result = hbar * mp <EOL> result *= KroneckerDelta ( m , mp ) <EOL> result *= KroneckerDelta ( j , jp ) <EOL> return result <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( basis , ** options ) <EOL> class J2Op ( SpinOpBase , HermitianOperator ) : <EOL> """<STR_LIT>""" <EOL> _coord = '<STR_LIT:2>' <EOL> def _eval_commutator_JxOp ( self , other ) : <EOL> return S . Zero <EOL> def _eval_commutator_JyOp ( self , other ) : <EOL> return S . Zero <EOL> def _eval_commutator_JzOp ( self , other ) : <EOL> return S . Zero <EOL> def _eval_commutator_JplusOp ( self , other ) : <EOL> return S . Zero <EOL> def _eval_commutator_JminusOp ( self , other ) : <EOL> return S . Zero <EOL> def _apply_operator_JxKet ( self , ket , ** options ) : <EOL> j = ket . j <EOL> return hbar ** <NUM_LIT:2> * j * ( j + <NUM_LIT:1> ) * ket <EOL> def _apply_operator_JxKetCoupled ( self , ket , ** options ) : <EOL> j = ket . j <EOL> return hbar ** <NUM_LIT:2> * j * ( j + <NUM_LIT:1> ) * ket <EOL> def _apply_operator_JyKet ( self , ket , ** options ) : <EOL> j = ket . j <EOL> return hbar ** <NUM_LIT:2> * j * ( j + <NUM_LIT:1> ) * ket <EOL> def _apply_operator_JyKetCoupled ( self , ket , ** options ) : <EOL> j = ket . j <EOL> return hbar ** <NUM_LIT:2> * j * ( j + <NUM_LIT:1> ) * ket <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> j = ket . j <EOL> return hbar ** <NUM_LIT:2> * j * ( j + <NUM_LIT:1> ) * ket <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> j = ket . j <EOL> return hbar ** <NUM_LIT:2> * j * ( j + <NUM_LIT:1> ) * ket <EOL> def matrix_element ( self , j , m , jp , mp ) : <EOL> result = ( hbar ** <NUM_LIT:2> ) * j * ( j + <NUM_LIT:1> ) <EOL> result *= KroneckerDelta ( m , mp ) <EOL> result *= KroneckerDelta ( j , jp ) <EOL> return result <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( basis , ** options ) <EOL> def _print_contents_pretty ( self , printer , * args ) : <EOL> a = prettyForm ( unicode ( self . name ) ) <EOL> b = prettyForm ( u ( '<STR_LIT:2>' ) ) <EOL> return a ** b <EOL> def _print_contents_latex ( self , printer , * args ) : <EOL> return r'<STR_LIT>' % str ( self . name ) <EOL> def _eval_rewrite_as_xyz ( self , * args ) : <EOL> return JxOp ( args [ <NUM_LIT:0> ] ) ** <NUM_LIT:2> + JyOp ( args [ <NUM_LIT:0> ] ) ** <NUM_LIT:2> + JzOp ( args [ <NUM_LIT:0> ] ) ** <NUM_LIT:2> <EOL> def _eval_rewrite_as_plusminus ( self , * args ) : <EOL> a = args [ <NUM_LIT:0> ] <EOL> return JzOp ( a ) ** <NUM_LIT:2> + Rational ( <NUM_LIT:1> , <NUM_LIT:2> ) * ( JplusOp ( a ) * JminusOp ( a ) + JminusOp ( a ) * JplusOp ( a ) ) <EOL> class Rotation ( UnitaryOperator ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def _eval_args ( cls , args ) : <EOL> args = QExpr . _eval_args ( args ) <EOL> if len ( args ) != <NUM_LIT:3> : <EOL> raise ValueError ( '<STR_LIT>' % args ) <EOL> return args <EOL> @ classmethod <EOL> def _eval_hilbert_space ( cls , label ) : <EOL> return ComplexSpace ( S . Infinity ) <EOL> @ property <EOL> def alpha ( self ) : <EOL> return self . label [ <NUM_LIT:0> ] <EOL> @ property <EOL> def beta ( self ) : <EOL> return self . label [ <NUM_LIT:1> ] <EOL> @ property <EOL> def gamma ( self ) : <EOL> return self . label [ <NUM_LIT:2> ] <EOL> def _print_operator_name ( self , printer , * args ) : <EOL> return '<STR_LIT:R>' <EOL> def _print_operator_name_pretty ( self , printer , * args ) : <EOL> if printer . _use_unicode : <EOL> return prettyForm ( u ( '<STR_LIT>' ) + u ( '<STR_LIT:U+0020>' ) ) <EOL> else : <EOL> return prettyForm ( "<STR_LIT>" ) <EOL> def _print_operator_name_latex ( self , printer , * args ) : <EOL> return r'<STR_LIT>' <EOL> def _eval_inverse ( self ) : <EOL> return Rotation ( - self . gamma , - self . beta , - self . alpha ) <EOL> @ classmethod <EOL> def D ( cls , j , m , mp , alpha , beta , gamma ) : <EOL> """<STR_LIT>""" <EOL> return WignerD ( j , m , mp , alpha , beta , gamma ) <EOL> @ classmethod <EOL> def d ( cls , j , m , mp , beta ) : <EOL> """<STR_LIT>""" <EOL> return WignerD ( j , m , mp , <NUM_LIT:0> , beta , <NUM_LIT:0> ) <EOL> def matrix_element ( self , j , m , jp , mp ) : <EOL> result = self . __class__ . D ( <EOL> jp , m , mp , self . alpha , self . beta , self . gamma <EOL> ) <EOL> result *= KroneckerDelta ( j , jp ) <EOL> return result <EOL> def _represent_base ( self , basis , ** options ) : <EOL> j = sympify ( options . get ( '<STR_LIT>' , Rational ( <NUM_LIT:1> , <NUM_LIT:2> ) ) ) <EOL> evaluate = sympify ( options . get ( '<STR_LIT>' ) ) <EOL> size , mvals = m_values ( j ) <EOL> result = zeros ( size , size ) <EOL> for p in range ( size ) : <EOL> for q in range ( size ) : <EOL> me = self . matrix_element ( j , mvals [ p ] , j , mvals [ q ] ) <EOL> if evaluate : <EOL> result [ p , q ] = me . doit ( ) <EOL> else : <EOL> result [ p , q ] = me <EOL> return result <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( basis , ** options ) <EOL> def _apply_operator_uncoupled ( self , state , ket , ** options ) : <EOL> a = self . alpha <EOL> b = self . beta <EOL> g = self . gamma <EOL> j = ket . j <EOL> m = ket . m <EOL> if j . is_number : <EOL> s = [ ] <EOL> size = m_values ( j ) <EOL> sz = size [ <NUM_LIT:1> ] <EOL> for mp in sz : <EOL> r = Rotation . D ( j , m , mp , a , b , g ) <EOL> z = r . doit ( ) <EOL> s . append ( z * state ( j , mp ) ) <EOL> return Add ( * s ) <EOL> else : <EOL> if options . pop ( '<STR_LIT>' , True ) : <EOL> mp = Dummy ( '<STR_LIT>' ) <EOL> else : <EOL> mp = symbols ( '<STR_LIT>' ) <EOL> return Sum ( Rotation . D ( j , m , mp , a , b , g ) * state ( j , mp ) , ( mp , - j , j ) ) <EOL> def _apply_operator_JxKet ( self , ket , ** options ) : <EOL> return self . _apply_operator_uncoupled ( JxKet , ket , ** options ) <EOL> def _apply_operator_JyKet ( self , ket , ** options ) : <EOL> return self . _apply_operator_uncoupled ( JyKet , ket , ** options ) <EOL> def _apply_operator_JzKet ( self , ket , ** options ) : <EOL> return self . _apply_operator_uncoupled ( JzKet , ket , ** options ) <EOL> def _apply_operator_coupled ( self , state , ket , ** options ) : <EOL> a = self . alpha <EOL> b = self . beta <EOL> g = self . gamma <EOL> j = ket . j <EOL> m = ket . m <EOL> jn = ket . jn <EOL> coupling = ket . coupling <EOL> if j . is_number : <EOL> s = [ ] <EOL> size = m_values ( j ) <EOL> sz = size [ <NUM_LIT:1> ] <EOL> for mp in sz : <EOL> r = Rotation . D ( j , m , mp , a , b , g ) <EOL> z = r . doit ( ) <EOL> s . append ( z * state ( j , mp , jn , coupling ) ) <EOL> return Add ( * s ) <EOL> else : <EOL> if options . pop ( '<STR_LIT>' , True ) : <EOL> mp = Dummy ( '<STR_LIT>' ) <EOL> else : <EOL> mp = symbols ( '<STR_LIT>' ) <EOL> return Sum ( Rotation . D ( j , m , mp , a , b , g ) * state ( <EOL> j , mp , jn , coupling ) , ( mp , - j , j ) ) <EOL> def _apply_operator_JxKetCoupled ( self , ket , ** options ) : <EOL> return self . _apply_operator_coupled ( JxKetCoupled , ket , ** options ) <EOL> def _apply_operator_JyKetCoupled ( self , ket , ** options ) : <EOL> return self . _apply_operator_coupled ( JyKetCoupled , ket , ** options ) <EOL> def _apply_operator_JzKetCoupled ( self , ket , ** options ) : <EOL> return self . _apply_operator_coupled ( JzKetCoupled , ket , ** options ) <EOL> class WignerD ( Expr ) : <EOL> """<STR_LIT>""" <EOL> is_commutative = True <EOL> def __new__ ( cls , * args , ** hints ) : <EOL> if not len ( args ) == <NUM_LIT:6> : <EOL> raise ValueError ( '<STR_LIT>' % args ) <EOL> args = sympify ( args ) <EOL> evaluate = hints . get ( '<STR_LIT>' , False ) <EOL> if evaluate : <EOL> return Expr . __new__ ( cls , * args ) . _eval_wignerd ( ) <EOL> return Expr . __new__ ( cls , * args ) <EOL> @ property <EOL> def j ( self ) : <EOL> return self . args [ <NUM_LIT:0> ] <EOL> @ property <EOL> def m ( self ) : <EOL> return self . args [ <NUM_LIT:1> ] <EOL> @ property <EOL> def mp ( self ) : <EOL> return self . args [ <NUM_LIT:2> ] <EOL> @ property <EOL> def alpha ( self ) : <EOL> return self . args [ <NUM_LIT:3> ] <EOL> @ property <EOL> def beta ( self ) : <EOL> return self . args [ <NUM_LIT:4> ] <EOL> @ property <EOL> def gamma ( self ) : <EOL> return self . args [ <NUM_LIT:5> ] <EOL> def _latex ( self , printer , * args ) : <EOL> if self . alpha == <NUM_LIT:0> and self . gamma == <NUM_LIT:0> : <EOL> return r'<STR_LIT>' % ( <EOL> printer . _print ( self . j ) , printer . _print ( <EOL> self . m ) , printer . _print ( self . mp ) , <EOL> printer . _print ( self . beta ) ) <EOL> return r'<STR_LIT>' % ( <EOL> printer . _print ( <EOL> self . j ) , printer . _print ( self . m ) , printer . _print ( self . mp ) , <EOL> printer . _print ( self . alpha ) , printer . _print ( self . beta ) , printer . _print ( self . gamma ) ) <EOL> def _pretty ( self , printer , * args ) : <EOL> top = printer . _print ( self . j ) <EOL> bot = printer . _print ( self . m ) <EOL> bot = prettyForm ( * bot . right ( '<STR_LIT:U+002C>' ) ) <EOL> bot = prettyForm ( * bot . right ( printer . _print ( self . mp ) ) ) <EOL> pad = max ( top . width ( ) , bot . width ( ) ) <EOL> top = prettyForm ( * top . left ( '<STR_LIT:U+0020>' ) ) <EOL> bot = prettyForm ( * bot . left ( '<STR_LIT:U+0020>' ) ) <EOL> if pad > top . width ( ) : <EOL> top = prettyForm ( * top . right ( '<STR_LIT:U+0020>' * ( pad - top . width ( ) ) ) ) <EOL> if pad > bot . width ( ) : <EOL> bot = prettyForm ( * bot . right ( '<STR_LIT:U+0020>' * ( pad - bot . width ( ) ) ) ) <EOL> if self . alpha == <NUM_LIT:0> and self . gamma == <NUM_LIT:0> : <EOL> args = printer . _print ( self . beta ) <EOL> s = stringPict ( '<STR_LIT:d>' + '<STR_LIT:U+0020>' * pad ) <EOL> else : <EOL> args = printer . _print ( self . alpha ) <EOL> args = prettyForm ( * args . right ( '<STR_LIT:U+002C>' ) ) <EOL> args = prettyForm ( * args . right ( printer . _print ( self . beta ) ) ) <EOL> args = prettyForm ( * args . right ( '<STR_LIT:U+002C>' ) ) <EOL> args = prettyForm ( * args . right ( printer . _print ( self . gamma ) ) ) <EOL> s = stringPict ( '<STR_LIT:D>' + '<STR_LIT:U+0020>' * pad ) <EOL> args = prettyForm ( * args . parens ( ) ) <EOL> s = prettyForm ( * s . above ( top ) ) <EOL> s = prettyForm ( * s . below ( bot ) ) <EOL> s = prettyForm ( * s . right ( args ) ) <EOL> return s <EOL> def doit ( self , ** hints ) : <EOL> hints [ '<STR_LIT>' ] = True <EOL> return WignerD ( * self . args , ** hints ) <EOL> def _eval_wignerd ( self ) : <EOL> j = sympify ( self . j ) <EOL> m = sympify ( self . m ) <EOL> mp = sympify ( self . mp ) <EOL> alpha = sympify ( self . alpha ) <EOL> beta = sympify ( self . beta ) <EOL> gamma = sympify ( self . gamma ) <EOL> if not j . is_number : <EOL> raise ValueError ( <EOL> '<STR_LIT>' % j ) <EOL> r = <NUM_LIT:0> <EOL> if beta == pi / <NUM_LIT:2> : <EOL> for k in range ( <NUM_LIT:2> * j + <NUM_LIT:1> ) : <EOL> if k > j + mp or k > j - m or k < mp - m : <EOL> continue <EOL> r += ( - S ( <NUM_LIT:1> ) ) ** k * binomial ( j + mp , k ) * binomial ( j - mp , k + m - mp ) <EOL> r *= ( - S ( <NUM_LIT:1> ) ) ** ( m - mp ) / <NUM_LIT:2> ** j * sqrt ( factorial ( j + m ) * <EOL> factorial ( j - m ) / ( factorial ( j + mp ) * factorial ( j - mp ) ) ) <EOL> else : <EOL> size , mvals = m_values ( j ) <EOL> for mpp in mvals : <EOL> r += Rotation . d ( j , m , mpp , pi / <NUM_LIT:2> ) . doit ( ) * ( cos ( - mpp * beta ) + I * sin ( - mpp * beta ) ) * Rotation . d ( j , mpp , - mp , pi / <NUM_LIT:2> ) . doit ( ) <EOL> r = r * I ** ( <NUM_LIT:2> * j - m - mp ) * ( - <NUM_LIT:1> ) ** ( <NUM_LIT:2> * m ) <EOL> r = simplify ( r ) <EOL> r *= exp ( - I * m * alpha ) * exp ( - I * mp * gamma ) <EOL> return r <EOL> Jx = JxOp ( '<STR_LIT>' ) <EOL> Jy = JyOp ( '<STR_LIT>' ) <EOL> Jz = JzOp ( '<STR_LIT>' ) <EOL> J2 = J2Op ( '<STR_LIT>' ) <EOL> Jplus = JplusOp ( '<STR_LIT>' ) <EOL> Jminus = JminusOp ( '<STR_LIT>' ) <EOL> class SpinState ( State ) : <EOL> """<STR_LIT>""" <EOL> _label_separator = '<STR_LIT:U+002C>' <EOL> def __new__ ( cls , j , m ) : <EOL> j = sympify ( j ) <EOL> m = sympify ( m ) <EOL> if j . is_number : <EOL> if <NUM_LIT:2> * j != int ( <NUM_LIT:2> * j ) : <EOL> raise ValueError ( <EOL> '<STR_LIT>' % j ) <EOL> if j < <NUM_LIT:0> : <EOL> raise ValueError ( '<STR_LIT>' % j ) <EOL> if m . is_number : <EOL> if <NUM_LIT:2> * m != int ( <NUM_LIT:2> * m ) : <EOL> raise ValueError ( <EOL> '<STR_LIT>' % m ) <EOL> if j . is_number and m . is_number : <EOL> if abs ( m ) > j : <EOL> raise ValueError ( '<STR_LIT>' % ( j , m ) ) <EOL> if int ( j - m ) != j - m : <EOL> raise ValueError ( '<STR_LIT>' % ( j , m ) ) <EOL> return State . __new__ ( cls , j , m ) <EOL> @ property <EOL> def j ( self ) : <EOL> return self . label [ <NUM_LIT:0> ] <EOL> @ property <EOL> def m ( self ) : <EOL> return self . label [ <NUM_LIT:1> ] <EOL> @ classmethod <EOL> def _eval_hilbert_space ( cls , label ) : <EOL> return ComplexSpace ( <NUM_LIT:2> * label [ <NUM_LIT:0> ] + <NUM_LIT:1> ) <EOL> def _represent_base ( self , ** options ) : <EOL> j = self . j <EOL> m = self . m <EOL> alpha = sympify ( options . get ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> beta = sympify ( options . get ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> gamma = sympify ( options . get ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> size , mvals = m_values ( j ) <EOL> result = zeros ( size , <NUM_LIT:1> ) <EOL> for p , mval in enumerate ( mvals ) : <EOL> if m . is_number : <EOL> result [ p , <NUM_LIT:0> ] = Rotation . D ( <EOL> self . j , mval , self . m , alpha , beta , gamma ) . doit ( ) <EOL> else : <EOL> result [ p , <NUM_LIT:0> ] = Rotation . D ( self . j , mval , <EOL> self . m , alpha , beta , gamma ) <EOL> return result <EOL> def _eval_rewrite_as_Jx ( self , * args , ** options ) : <EOL> if isinstance ( self , Bra ) : <EOL> return self . _rewrite_basis ( Jx , JxBra , ** options ) <EOL> return self . _rewrite_basis ( Jx , JxKet , ** options ) <EOL> def _eval_rewrite_as_Jy ( self , * args , ** options ) : <EOL> if isinstance ( self , Bra ) : <EOL> return self . _rewrite_basis ( Jy , JyBra , ** options ) <EOL> return self . _rewrite_basis ( Jy , JyKet , ** options ) <EOL> def _eval_rewrite_as_Jz ( self , * args , ** options ) : <EOL> if isinstance ( self , Bra ) : <EOL> return self . _rewrite_basis ( Jz , JzBra , ** options ) <EOL> return self . _rewrite_basis ( Jz , JzKet , ** options ) <EOL> def _rewrite_basis ( self , basis , evect , ** options ) : <EOL> from sympy . physics . quantum . represent import represent <EOL> j = self . j <EOL> args = self . args [ <NUM_LIT:2> : ] <EOL> if j . is_number : <EOL> if isinstance ( self , CoupledSpinState ) : <EOL> if j == int ( j ) : <EOL> start = j ** <NUM_LIT:2> <EOL> else : <EOL> start = ( <NUM_LIT:2> * j - <NUM_LIT:1> ) * ( <NUM_LIT:2> * j + <NUM_LIT:1> ) / <NUM_LIT:4> <EOL> else : <EOL> start = <NUM_LIT:0> <EOL> vect = represent ( self , basis = basis , ** options ) <EOL> result = Add ( <EOL> * [ vect [ start + i ] * evect ( j , j - i , * args ) for i in range ( <NUM_LIT:2> * j + <NUM_LIT:1> ) ] ) <EOL> if isinstance ( self , CoupledSpinState ) and options . get ( '<STR_LIT>' ) is False : <EOL> return uncouple ( result ) <EOL> return result <EOL> else : <EOL> i = <NUM_LIT:0> <EOL> mi = symbols ( '<STR_LIT>' ) <EOL> while self . subs ( mi , <NUM_LIT:0> ) != self : <EOL> i += <NUM_LIT:1> <EOL> mi = symbols ( '<STR_LIT>' % i ) <EOL> break <EOL> if isinstance ( self , CoupledSpinState ) : <EOL> test_args = ( <NUM_LIT:0> , mi , ( <NUM_LIT:0> , <NUM_LIT:0> ) ) <EOL> else : <EOL> test_args = ( <NUM_LIT:0> , mi ) <EOL> if isinstance ( self , Ket ) : <EOL> angles = represent ( <EOL> self . __class__ ( * test_args ) , basis = basis ) [ <NUM_LIT:0> ] . args [ <NUM_LIT:3> : <NUM_LIT:6> ] <EOL> else : <EOL> angles = represent ( self . __class__ ( <EOL> * test_args ) , basis = basis ) [ <NUM_LIT:0> ] . args [ <NUM_LIT:0> ] . args [ <NUM_LIT:3> : <NUM_LIT:6> ] <EOL> if angles == ( <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ) : <EOL> return self <EOL> else : <EOL> state = evect ( j , mi , * args ) <EOL> lt = Rotation . D ( j , mi , self . m , * angles ) <EOL> return Sum ( lt * state , ( mi , - j , j ) ) <EOL> def _eval_innerproduct_JxBra ( self , bra , ** hints ) : <EOL> result = KroneckerDelta ( self . j , bra . j ) <EOL> if bra . dual_class ( ) is not self . __class__ : <EOL> result *= self . _represent_JxOp ( None ) [ bra . j - bra . m ] <EOL> else : <EOL> result *= KroneckerDelta ( <EOL> self . j , bra . j ) * KroneckerDelta ( self . m , bra . m ) <EOL> return result <EOL> def _eval_innerproduct_JyBra ( self , bra , ** hints ) : <EOL> result = KroneckerDelta ( self . j , bra . j ) <EOL> if bra . dual_class ( ) is not self . __class__ : <EOL> result *= self . _represent_JyOp ( None ) [ bra . j - bra . m ] <EOL> else : <EOL> result *= KroneckerDelta ( <EOL> self . j , bra . j ) * KroneckerDelta ( self . m , bra . m ) <EOL> return result <EOL> def _eval_innerproduct_JzBra ( self , bra , ** hints ) : <EOL> result = KroneckerDelta ( self . j , bra . j ) <EOL> if bra . dual_class ( ) is not self . __class__ : <EOL> result *= self . _represent_JzOp ( None ) [ bra . j - bra . m ] <EOL> else : <EOL> result *= KroneckerDelta ( <EOL> self . j , bra . j ) * KroneckerDelta ( self . m , bra . m ) <EOL> return result <EOL> def _eval_trace ( self , bra , ** hints ) : <EOL> return ( bra * self ) . doit ( ) <EOL> class JxKet ( SpinState , Ket ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JxBra <EOL> @ classmethod <EOL> def coupled_class ( self ) : <EOL> return JxKetCoupled <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JxOp ( None , ** options ) <EOL> def _represent_JxOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( ** options ) <EOL> def _represent_JyOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( alpha = <NUM_LIT:3> * pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( beta = pi / <NUM_LIT:2> , ** options ) <EOL> class JxBra ( SpinState , Bra ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JxKet <EOL> @ classmethod <EOL> def coupled_class ( self ) : <EOL> return JxBraCoupled <EOL> class JyKet ( SpinState , Ket ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JyBra <EOL> @ classmethod <EOL> def coupled_class ( self ) : <EOL> return JyKetCoupled <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JyOp ( None , ** options ) <EOL> def _represent_JxOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( gamma = pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JyOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( alpha = <NUM_LIT:3> * pi / <NUM_LIT:2> , beta = - pi / <NUM_LIT:2> , gamma = pi / <NUM_LIT:2> , ** options ) <EOL> class JyBra ( SpinState , Bra ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JyKet <EOL> @ classmethod <EOL> def coupled_class ( self ) : <EOL> return JyBraCoupled <EOL> class JzKet ( SpinState , Ket ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JzBra <EOL> @ classmethod <EOL> def coupled_class ( self ) : <EOL> return JzKetCoupled <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JxOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( beta = <NUM_LIT:3> * pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JyOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( alpha = <NUM_LIT:3> * pi / <NUM_LIT:2> , beta = pi / <NUM_LIT:2> , gamma = pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_base ( ** options ) <EOL> class JzBra ( SpinState , Bra ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JzKet <EOL> @ classmethod <EOL> def coupled_class ( self ) : <EOL> return JzBraCoupled <EOL> def _build_coupled ( jcoupling , length ) : <EOL> n_list = [ [ n + <NUM_LIT:1> ] for n in range ( length ) ] <EOL> coupled_jn = [ ] <EOL> coupled_n = [ ] <EOL> for n1 , n2 , j_new in jcoupling : <EOL> coupled_jn . append ( j_new ) <EOL> coupled_n . append ( ( n_list [ n1 - <NUM_LIT:1> ] , n_list [ n2 - <NUM_LIT:1> ] ) ) <EOL> n_sort = sorted ( n_list [ n1 - <NUM_LIT:1> ] + n_list [ n2 - <NUM_LIT:1> ] ) <EOL> n_list [ n_sort [ <NUM_LIT:0> ] - <NUM_LIT:1> ] = n_sort <EOL> return coupled_n , coupled_jn <EOL> class CoupledSpinState ( SpinState ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( cls , j , m , jn , * jcoupling ) : <EOL> SpinState ( j , m ) <EOL> if len ( jcoupling ) == <NUM_LIT:0> : <EOL> jcoupling = [ ] <EOL> for n in range ( <NUM_LIT:2> , len ( jn ) ) : <EOL> jcoupling . append ( ( <NUM_LIT:1> , n , Add ( * [ jn [ i ] for i in range ( n ) ] ) ) ) <EOL> jcoupling . append ( ( <NUM_LIT:1> , len ( jn ) , j ) ) <EOL> elif len ( jcoupling ) == <NUM_LIT:1> : <EOL> jcoupling = jcoupling [ <NUM_LIT:0> ] <EOL> else : <EOL> raise TypeError ( "<STR_LIT>" % ( len ( jcoupling ) + <NUM_LIT:3> ) ) <EOL> if not ( isinstance ( jn , list ) or isinstance ( jn , tuple ) or isinstance ( jn , Tuple ) ) : <EOL> raise TypeError ( '<STR_LIT>' % <EOL> jn . __class__ . __name__ ) <EOL> if not ( isinstance ( jcoupling , list ) or isinstance ( jcoupling , tuple ) or isinstance ( jcoupling , Tuple ) ) : <EOL> raise TypeError ( '<STR_LIT>' % <EOL> jcoupling . __class__ . __name__ ) <EOL> if not all ( isinstance ( term , list ) or isinstance ( term , tuple ) or isinstance ( term , Tuple ) for term in jcoupling ) : <EOL> raise TypeError ( <EOL> '<STR_LIT>' ) <EOL> if not len ( jn ) - <NUM_LIT:1> == len ( jcoupling ) : <EOL> raise ValueError ( '<STR_LIT>' % <EOL> ( len ( jn ) - <NUM_LIT:1> , len ( jcoupling ) ) ) <EOL> if not all ( len ( x ) == <NUM_LIT:3> for x in jcoupling ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> j = sympify ( j ) <EOL> m = sympify ( m ) <EOL> jn = Tuple ( * [ sympify ( ji ) for ji in jn ] ) <EOL> jcoupling = Tuple ( * [ Tuple ( sympify ( <EOL> n1 ) , sympify ( n2 ) , sympify ( ji ) ) for ( n1 , n2 , ji ) in jcoupling ] ) <EOL> if any ( <NUM_LIT:2> * ji != int ( <NUM_LIT:2> * ji ) for ji in jn if ji . is_number ) : <EOL> raise ValueError ( '<STR_LIT>' % jn ) <EOL> if any ( n1 != int ( n1 ) or n2 != int ( n2 ) for ( n1 , n2 , _ ) in jcoupling ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if any ( n1 < <NUM_LIT:1> or n2 < <NUM_LIT:1> or n1 > len ( jn ) or n2 > len ( jn ) for ( n1 , n2 , _ ) in jcoupling ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if any ( <NUM_LIT:2> * ji != int ( <NUM_LIT:2> * ji ) for ( _ , _ , ji ) in jcoupling if ji . is_number ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> coupled_n , coupled_jn = _build_coupled ( jcoupling , len ( jn ) ) <EOL> jvals = list ( jn ) <EOL> for n , ( n1 , n2 ) in enumerate ( coupled_n ) : <EOL> j1 = jvals [ min ( n1 ) - <NUM_LIT:1> ] <EOL> j2 = jvals [ min ( n2 ) - <NUM_LIT:1> ] <EOL> j3 = coupled_jn [ n ] <EOL> if sympify ( j1 ) . is_number and sympify ( j2 ) . is_number and sympify ( j3 ) . is_number : <EOL> if j1 + j2 < j3 : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' % ( n + <NUM_LIT:1> , j1 , j2 , j3 ) ) <EOL> if abs ( j1 - j2 ) > j3 : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( n + <NUM_LIT:1> , j1 , j2 , j3 ) ) <EOL> if int ( j1 + j2 ) == j1 + j2 : <EOL> pass <EOL> jvals [ min ( n1 + n2 ) - <NUM_LIT:1> ] = j3 <EOL> if len ( jcoupling ) > <NUM_LIT:0> and jcoupling [ - <NUM_LIT:1> ] [ <NUM_LIT:2> ] != j : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return State . __new__ ( cls , j , m , jn , jcoupling ) <EOL> def _print_label ( self , printer , * args ) : <EOL> label = [ printer . _print ( self . j ) , printer . _print ( self . m ) ] <EOL> for i , ji in enumerate ( self . jn , start = <NUM_LIT:1> ) : <EOL> label . append ( '<STR_LIT>' % ( <EOL> i , printer . _print ( ji ) <EOL> ) ) <EOL> for jn , ( n1 , n2 ) in zip ( self . coupled_jn [ : - <NUM_LIT:1> ] , self . coupled_n [ : - <NUM_LIT:1> ] ) : <EOL> label . append ( '<STR_LIT>' % ( <EOL> '<STR_LIT:U+002C>' . join ( str ( i ) for i in sorted ( n1 + n2 ) ) , printer . _print ( jn ) <EOL> ) ) <EOL> return '<STR_LIT:U+002C>' . join ( label ) <EOL> def _print_label_pretty ( self , printer , * args ) : <EOL> label = [ self . j , self . m ] <EOL> for i , ji in enumerate ( self . jn , start = <NUM_LIT:1> ) : <EOL> symb = '<STR_LIT>' % i <EOL> symb = pretty_symbol ( symb ) <EOL> symb = prettyForm ( symb + '<STR_LIT:=>' ) <EOL> item = prettyForm ( * symb . right ( printer . _print ( ji ) ) ) <EOL> label . append ( item ) <EOL> for jn , ( n1 , n2 ) in zip ( self . coupled_jn [ : - <NUM_LIT:1> ] , self . coupled_n [ : - <NUM_LIT:1> ] ) : <EOL> n = '<STR_LIT:U+002C>' . join ( pretty_symbol ( "<STR_LIT>" % i ) [ - <NUM_LIT:1> ] for i in sorted ( n1 + n2 ) ) <EOL> symb = prettyForm ( '<STR_LIT>' + n + '<STR_LIT:=>' ) <EOL> item = prettyForm ( * symb . right ( printer . _print ( jn ) ) ) <EOL> label . append ( item ) <EOL> return self . _print_sequence_pretty ( <EOL> label , self . _label_separator , printer , * args <EOL> ) <EOL> def _print_label_latex ( self , printer , * args ) : <EOL> label = [ self . j , self . m ] <EOL> for i , ji in enumerate ( self . jn , start = <NUM_LIT:1> ) : <EOL> label . append ( '<STR_LIT>' % ( i , printer . _print ( ji ) ) ) <EOL> for jn , ( n1 , n2 ) in zip ( self . coupled_jn [ : - <NUM_LIT:1> ] , self . coupled_n [ : - <NUM_LIT:1> ] ) : <EOL> n = '<STR_LIT:U+002C>' . join ( str ( i ) for i in sorted ( n1 + n2 ) ) <EOL> label . append ( '<STR_LIT>' % ( n , printer . _print ( jn ) ) ) <EOL> return self . _print_sequence ( <EOL> label , self . _label_separator , printer , * args <EOL> ) <EOL> @ property <EOL> def jn ( self ) : <EOL> return self . label [ <NUM_LIT:2> ] <EOL> @ property <EOL> def coupling ( self ) : <EOL> return self . label [ <NUM_LIT:3> ] <EOL> @ property <EOL> def coupled_jn ( self ) : <EOL> return _build_coupled ( self . label [ <NUM_LIT:3> ] , len ( self . label [ <NUM_LIT:2> ] ) ) [ <NUM_LIT:1> ] <EOL> @ property <EOL> def coupled_n ( self ) : <EOL> return _build_coupled ( self . label [ <NUM_LIT:3> ] , len ( self . label [ <NUM_LIT:2> ] ) ) [ <NUM_LIT:0> ] <EOL> @ classmethod <EOL> def _eval_hilbert_space ( cls , label ) : <EOL> j = Add ( * label [ <NUM_LIT:2> ] ) <EOL> if j . is_number : <EOL> return DirectSumHilbertSpace ( * [ ComplexSpace ( x ) for x in range ( int ( <NUM_LIT:2> * j + <NUM_LIT:1> ) , <NUM_LIT:0> , - <NUM_LIT:2> ) ] ) <EOL> else : <EOL> return ComplexSpace ( <NUM_LIT:2> * j + <NUM_LIT:1> ) <EOL> def _represent_coupled_base ( self , ** options ) : <EOL> evect = self . uncoupled_class ( ) <EOL> if not self . j . is_number : <EOL> raise ValueError ( <EOL> '<STR_LIT>' ) <EOL> if not self . hilbert_space . dimension . is_number : <EOL> raise ValueError ( <EOL> '<STR_LIT>' ) <EOL> result = zeros ( self . hilbert_space . dimension , <NUM_LIT:1> ) <EOL> if self . j == int ( self . j ) : <EOL> start = self . j ** <NUM_LIT:2> <EOL> else : <EOL> start = ( <NUM_LIT:2> * self . j - <NUM_LIT:1> ) * ( <NUM_LIT:1> + <NUM_LIT:2> * self . j ) / <NUM_LIT:4> <EOL> result [ start : start + <NUM_LIT:2> * self . j + <NUM_LIT:1> , <NUM_LIT:0> ] = evect ( <EOL> self . j , self . m ) . _represent_base ( ** options ) <EOL> return result <EOL> def _eval_rewrite_as_Jx ( self , * args , ** options ) : <EOL> if isinstance ( self , Bra ) : <EOL> return self . _rewrite_basis ( Jx , JxBraCoupled , ** options ) <EOL> return self . _rewrite_basis ( Jx , JxKetCoupled , ** options ) <EOL> def _eval_rewrite_as_Jy ( self , * args , ** options ) : <EOL> if isinstance ( self , Bra ) : <EOL> return self . _rewrite_basis ( Jy , JyBraCoupled , ** options ) <EOL> return self . _rewrite_basis ( Jy , JyKetCoupled , ** options ) <EOL> def _eval_rewrite_as_Jz ( self , * args , ** options ) : <EOL> if isinstance ( self , Bra ) : <EOL> return self . _rewrite_basis ( Jz , JzBraCoupled , ** options ) <EOL> return self . _rewrite_basis ( Jz , JzKetCoupled , ** options ) <EOL> class JxKetCoupled ( CoupledSpinState , Ket ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JxBraCoupled <EOL> @ classmethod <EOL> def uncoupled_class ( self ) : <EOL> return JxKet <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JxOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( ** options ) <EOL> def _represent_JyOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( alpha = <NUM_LIT:3> * pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( beta = pi / <NUM_LIT:2> , ** options ) <EOL> class JxBraCoupled ( CoupledSpinState , Bra ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JxKetCoupled <EOL> @ classmethod <EOL> def uncoupled_class ( self ) : <EOL> return JxBra <EOL> class JyKetCoupled ( CoupledSpinState , Ket ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JyBraCoupled <EOL> @ classmethod <EOL> def uncoupled_class ( self ) : <EOL> return JyKet <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JxOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( gamma = pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JyOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( alpha = <NUM_LIT:3> * pi / <NUM_LIT:2> , beta = - pi / <NUM_LIT:2> , gamma = pi / <NUM_LIT:2> , ** options ) <EOL> class JyBraCoupled ( CoupledSpinState , Bra ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JyKetCoupled <EOL> @ classmethod <EOL> def uncoupled_class ( self ) : <EOL> return JyBra <EOL> class JzKetCoupled ( CoupledSpinState , Ket ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JzBraCoupled <EOL> @ classmethod <EOL> def uncoupled_class ( self ) : <EOL> return JzKet <EOL> def _represent_default_basis ( self , ** options ) : <EOL> return self . _represent_JzOp ( None , ** options ) <EOL> def _represent_JxOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( beta = <NUM_LIT:3> * pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JyOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( alpha = <NUM_LIT:3> * pi / <NUM_LIT:2> , beta = pi / <NUM_LIT:2> , gamma = pi / <NUM_LIT:2> , ** options ) <EOL> def _represent_JzOp ( self , basis , ** options ) : <EOL> return self . _represent_coupled_base ( ** options ) <EOL> class JzBraCoupled ( CoupledSpinState , Bra ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def dual_class ( self ) : <EOL> return JzKetCoupled <EOL> @ classmethod <EOL> def uncoupled_class ( self ) : <EOL> return JzBra <EOL> def couple ( expr , jcoupling_list = None ) : <EOL> """<STR_LIT>""" <EOL> a = expr . atoms ( TensorProduct ) <EOL> for tp in a : <EOL> if not all ( [ isinstance ( state , SpinState ) for state in tp . args ] ) : <EOL> continue <EOL> if not all ( [ state . __class__ is tp . args [ <NUM_LIT:0> ] . __class__ for state in tp . args ] ) : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> expr = expr . subs ( tp , _couple ( tp , jcoupling_list ) ) <EOL> return expr <EOL> def _couple ( tp , jcoupling_list ) : <EOL> states = tp . args <EOL> coupled_evect = states [ <NUM_LIT:0> ] . coupled_class ( ) <EOL> if jcoupling_list is None : <EOL> jcoupling_list = [ ] <EOL> for n in range ( <NUM_LIT:1> , len ( states ) ) : <EOL> jcoupling_list . append ( ( <NUM_LIT:1> , n + <NUM_LIT:1> ) ) <EOL> if not len ( jcoupling_list ) == len ( states ) - <NUM_LIT:1> : <EOL> raise TypeError ( '<STR_LIT>' % <EOL> ( len ( states ) - <NUM_LIT:1> , len ( jcoupling_list ) ) ) <EOL> if not all ( len ( coupling ) == <NUM_LIT:2> for coupling in jcoupling_list ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if any ( [ n1 == n2 for n1 , n2 in jcoupling_list ] ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if all ( [ sympify ( n1 ) . is_number and sympify ( n2 ) . is_number for n1 , n2 in jcoupling_list ] ) : <EOL> j_test = [ <NUM_LIT:0> ] * len ( states ) <EOL> for n1 , n2 in jcoupling_list : <EOL> if j_test [ n1 - <NUM_LIT:1> ] == - <NUM_LIT:1> or j_test [ n2 - <NUM_LIT:1> ] == - <NUM_LIT:1> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> j_test [ max ( n1 , n2 ) - <NUM_LIT:1> ] = - <NUM_LIT:1> <EOL> jn = [ state . j for state in states ] <EOL> mn = [ state . m for state in states ] <EOL> coupling_list = [ ] <EOL> n_list = [ [ i + <NUM_LIT:1> ] for i in range ( len ( states ) ) ] <EOL> for j_coupling in jcoupling_list : <EOL> n1 , n2 = j_coupling <EOL> j1_n = list ( n_list [ n1 - <NUM_LIT:1> ] ) <EOL> j2_n = list ( n_list [ n2 - <NUM_LIT:1> ] ) <EOL> coupling_list . append ( ( j1_n , j2_n ) ) <EOL> n_list [ min ( n1 , n2 ) - <NUM_LIT:1> ] = sorted ( j1_n + j2_n ) <EOL> if all ( state . j . is_number and state . m . is_number for state in states ) : <EOL> diff_max = [ Add ( * [ jn [ n - <NUM_LIT:1> ] - mn [ n - <NUM_LIT:1> ] for n in coupling [ <NUM_LIT:0> ] + <EOL> coupling [ <NUM_LIT:1> ] ] ) for coupling in coupling_list ] <EOL> result = [ ] <EOL> for diff in range ( diff_max [ - <NUM_LIT:1> ] + <NUM_LIT:1> ) : <EOL> n = len ( coupling_list ) <EOL> tot = binomial ( diff + n - <NUM_LIT:1> , diff ) <EOL> for config_num in range ( tot ) : <EOL> diff_list = _confignum_to_difflist ( config_num , diff , n ) <EOL> if any ( [ d > m for d , m in zip ( diff_list , diff_max ) ] ) : <EOL> continue <EOL> cg_terms = [ ] <EOL> coupled_j = list ( jn ) <EOL> jcoupling = [ ] <EOL> for ( j1_n , j2_n ) , coupling_diff in zip ( coupling_list , diff_list ) : <EOL> j1 = coupled_j [ min ( j1_n ) - <NUM_LIT:1> ] <EOL> j2 = coupled_j [ min ( j2_n ) - <NUM_LIT:1> ] <EOL> j3 = j1 + j2 - coupling_diff <EOL> coupled_j [ min ( j1_n + j2_n ) - <NUM_LIT:1> ] = j3 <EOL> m1 = Add ( * [ mn [ x - <NUM_LIT:1> ] for x in j1_n ] ) <EOL> m2 = Add ( * [ mn [ x - <NUM_LIT:1> ] for x in j2_n ] ) <EOL> m3 = m1 + m2 <EOL> cg_terms . append ( ( j1 , m1 , j2 , m2 , j3 , m3 ) ) <EOL> jcoupling . append ( ( min ( j1_n ) , min ( j2_n ) , j3 ) ) <EOL> if any ( [ abs ( term [ <NUM_LIT:5> ] ) > term [ <NUM_LIT:4> ] for term in cg_terms ] ) : <EOL> continue <EOL> if any ( [ term [ <NUM_LIT:0> ] + term [ <NUM_LIT:2> ] < term [ <NUM_LIT:4> ] for term in cg_terms ] ) : <EOL> continue <EOL> if any ( [ abs ( term [ <NUM_LIT:0> ] - term [ <NUM_LIT:2> ] ) > term [ <NUM_LIT:4> ] for term in cg_terms ] ) : <EOL> continue <EOL> coeff = Mul ( * [ CG ( * term ) . doit ( ) for term in cg_terms ] ) <EOL> state = coupled_evect ( j3 , m3 , jn , jcoupling ) <EOL> result . append ( coeff * state ) <EOL> return Add ( * result ) <EOL> else : <EOL> cg_terms = [ ] <EOL> jcoupling = [ ] <EOL> sum_terms = [ ] <EOL> coupled_j = list ( jn ) <EOL> for j1_n , j2_n in coupling_list : <EOL> j1 = coupled_j [ min ( j1_n ) - <NUM_LIT:1> ] <EOL> j2 = coupled_j [ min ( j2_n ) - <NUM_LIT:1> ] <EOL> if len ( j1_n + j2_n ) == len ( states ) : <EOL> j3 = symbols ( '<STR_LIT>' ) <EOL> else : <EOL> j3_name = '<STR_LIT>' + '<STR_LIT>' . join ( [ "<STR_LIT:%s>" % n for n in j1_n + j2_n ] ) <EOL> j3 = symbols ( j3_name ) <EOL> coupled_j [ min ( j1_n + j2_n ) - <NUM_LIT:1> ] = j3 <EOL> m1 = Add ( * [ mn [ x - <NUM_LIT:1> ] for x in j1_n ] ) <EOL> m2 = Add ( * [ mn [ x - <NUM_LIT:1> ] for x in j2_n ] ) <EOL> m3 = m1 + m2 <EOL> cg_terms . append ( ( j1 , m1 , j2 , m2 , j3 , m3 ) ) <EOL> jcoupling . append ( ( min ( j1_n ) , min ( j2_n ) , j3 ) ) <EOL> sum_terms . append ( ( j3 , m3 , j1 + j2 ) ) <EOL> coeff = Mul ( * [ CG ( * term ) for term in cg_terms ] ) <EOL> state = coupled_evect ( j3 , m3 , jn , jcoupling ) <EOL> return Sum ( coeff * state , * sum_terms ) <EOL> def uncouple ( expr , jn = None , jcoupling_list = None ) : <EOL> """<STR_LIT>""" <EOL> a = expr . atoms ( SpinState ) <EOL> for state in a : <EOL> expr = expr . subs ( state , _uncouple ( state , jn , jcoupling_list ) ) <EOL> return expr <EOL> def _uncouple ( state , jn , jcoupling_list ) : <EOL> if isinstance ( state , CoupledSpinState ) : <EOL> jn = state . jn <EOL> coupled_n = state . coupled_n <EOL> coupled_jn = state . coupled_jn <EOL> evect = state . uncoupled_class ( ) <EOL> elif isinstance ( state , SpinState ) : <EOL> if jn is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not ( isinstance ( jn , list ) or isinstance ( jn , tuple ) ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> if jcoupling_list is None : <EOL> jcoupling_list = [ ] <EOL> for i in range ( <NUM_LIT:1> , len ( jn ) ) : <EOL> jcoupling_list . append ( <EOL> ( <NUM_LIT:1> , <NUM_LIT:1> + i , Add ( * [ jn [ j ] for j in range ( i + <NUM_LIT:1> ) ] ) ) ) <EOL> if not ( isinstance ( jcoupling_list , list ) or isinstance ( jcoupling_list , tuple ) ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> if not len ( jcoupling_list ) == len ( jn ) - <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> coupled_n , coupled_jn = _build_coupled ( jcoupling_list , len ( jn ) ) <EOL> evect = state . __class__ <EOL> else : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> j = state . j <EOL> m = state . m <EOL> coupling_list = [ ] <EOL> j_list = list ( jn ) <EOL> for j3 , ( n1 , n2 ) in zip ( coupled_jn , coupled_n ) : <EOL> j1 = j_list [ n1 [ <NUM_LIT:0> ] - <NUM_LIT:1> ] <EOL> j2 = j_list [ n2 [ <NUM_LIT:0> ] - <NUM_LIT:1> ] <EOL> coupling_list . append ( ( n1 , n2 , j1 , j2 , j3 ) ) <EOL> j_list [ min ( n1 + n2 ) - <NUM_LIT:1> ] = j3 <EOL> if j . is_number and m . is_number : <EOL> diff_max = [ <NUM_LIT:2> * x for x in jn ] <EOL> diff = Add ( * jn ) - m <EOL> n = len ( jn ) <EOL> tot = binomial ( diff + n - <NUM_LIT:1> , diff ) <EOL> result = [ ] <EOL> for config_num in range ( tot ) : <EOL> diff_list = _confignum_to_difflist ( config_num , diff , n ) <EOL> if any ( [ d > p for d , p in zip ( diff_list , diff_max ) ] ) : <EOL> continue <EOL> cg_terms = [ ] <EOL> for coupling in coupling_list : <EOL> j1_n , j2_n , j1 , j2 , j3 = coupling <EOL> m1 = Add ( * [ jn [ x - <NUM_LIT:1> ] - diff_list [ x - <NUM_LIT:1> ] for x in j1_n ] ) <EOL> m2 = Add ( * [ jn [ x - <NUM_LIT:1> ] - diff_list [ x - <NUM_LIT:1> ] for x in j2_n ] ) <EOL> m3 = m1 + m2 <EOL> cg_terms . append ( ( j1 , m1 , j2 , m2 , j3 , m3 ) ) <EOL> coeff = Mul ( * [ CG ( * term ) . doit ( ) for term in cg_terms ] ) <EOL> state = TensorProduct ( <EOL> * [ evect ( j , j - d ) for j , d in zip ( jn , diff_list ) ] ) <EOL> result . append ( coeff * state ) <EOL> return Add ( * result ) <EOL> else : <EOL> m_str = "<STR_LIT>" % ( len ( jn ) + <NUM_LIT:1> ) <EOL> mvals = symbols ( m_str ) <EOL> cg_terms = [ ( j1 , Add ( * [ mvals [ n - <NUM_LIT:1> ] for n in j1_n ] ) , <EOL> j2 , Add ( * [ mvals [ n - <NUM_LIT:1> ] for n in j2_n ] ) , <EOL> j3 , Add ( * [ mvals [ n - <NUM_LIT:1> ] for n in j1_n + j2_n ] ) ) for j1_n , j2_n , j1 , j2 , j3 in coupling_list [ : - <NUM_LIT:1> ] ] <EOL> cg_terms . append ( * [ ( j1 , Add ( * [ mvals [ n - <NUM_LIT:1> ] for n in j1_n ] ) , <EOL> j2 , Add ( * [ mvals [ n - <NUM_LIT:1> ] for n in j2_n ] ) , <EOL> j , m ) for j1_n , j2_n , j1 , j2 , j3 in [ coupling_list [ - <NUM_LIT:1> ] ] ] ) <EOL> cg_coeff = Mul ( * [ CG ( * cg_term ) for cg_term in cg_terms ] ) <EOL> sum_terms = [ ( m , - j , j ) for j , m in zip ( jn , mvals ) ] <EOL> state = TensorProduct ( * [ evect ( j , m ) for j , m in zip ( jn , mvals ) ] ) <EOL> return Sum ( cg_coeff * state , * sum_terms ) <EOL> def _confignum_to_difflist ( config_num , diff , list_len ) : <EOL> diff_list = [ ] <EOL> for n in range ( list_len ) : <EOL> prev_diff = diff <EOL> rem_spots = list_len - n - <NUM_LIT:1> <EOL> rem_configs = binomial ( diff + rem_spots - <NUM_LIT:1> , diff ) <EOL> while config_num >= rem_configs : <EOL> config_num -= rem_configs <EOL> diff -= <NUM_LIT:1> <EOL> rem_configs = binomial ( diff + rem_spots - <NUM_LIT:1> , diff ) <EOL> diff_list . append ( prev_diff - diff ) <EOL> return diff_list </s>
