o
    8Va                      @   s  d dl mZ d dlmZ d dlmZmZ d dlZd dl	m
Z
mZmZmZmZmZmZ d dlmZ d dlmZmZ d dlmZ d d	lmZmZ d d
lmZmZ d dlmZ d dl m!Z! d dl"m#Z# d dl$m%Z% d dl&m'Z' d dl(m)Z) d.ddZ*G dd deZ+e)e+edd Z,e)e+e+dd Z,dd Z-e-ege-egdej.e+< dd Z/G dd deZ0G dd  d e+Z1d!d" Z2G d#d$ d$Z3d%d& Z4d'd(l5m6Z6 d'd)l7m8Z8 d'd*l9m:Z: d'd+l;m<Z< d'd,l=m>Z> d'd-l?m@Z@mAZA dS )/    )Tuple)	FuzzyBool)wrapsreduceN)SSymbolIntegerBasicExprMulAdd)call_highest_priority)
SYMPY_INTSdefault_sort_key)Str)SympifyError_sympify)	conjugateadjoint)KroneckerDelta)NonSquareMatrixErrorsimplify)
MatrixKind)
filldedent)dispatchc                    s    fdd}|S )Nc                    s   t   fdd}|S )Nc                    s,   z
t |} | |W S  ty    Y S w N)r   r   ab)funcretval D/usr/lib/python3/dist-packages/sympy/matrices/expressions/matexpr.py__sympifyit_wrapper   s   z5_sympifyit.<locals>.deco.<locals>.__sympifyit_wrapper)r   )r    r$   r!   )r    r#   deco   s   z_sympifyit.<locals>.decor"   )argr!   r&   r"   r%   r#   
_sympifyit   s   r(   c                       s  e Zd ZdZdZdZdZdZdZdZ	dZ
dZdZdZdZdZdZdZe Zdd Zedeeef fd	d
Zedd Zedd Zdd Zdd Zedeeddd Z edeeddd Z!edeeddd Z"edeeddd Z#edeed d!d" Z$edeed d#d$ Z%edeed%d&d' Z&edeed%d(d) Z'edeed*d+d, Z(edeed-d.d/ Z)edeed0d1d2 Z*edeed3d4d5 Z+ed6d7 Z,ed8d9 Z-ed:d; Z.d<d= Z/dvd>d?Z0d@dA Z1dBdC Z2dDdE Z3dFdG Z4dHdI Z5dJdK Z6 fdLdMZ7e8dNdO Z9dPdQ Z:dRdS Z;dwdTdUZ<dVdW Z=dXdY Z>edZd[ Z?d\d] Z@d^d_ ZAed`da ZBdbdc ZCddde ZDdfdg ZEdhdi ZFdjdk ZGdldm ZHdndo ZIdpdq ZJeKdxdrdsZLdtdu ZM  ZNS )y
MatrixExpra  Superclass for Matrix Expressions

    MatrixExprs represent abstract matrices, linear transformations represented
    within a particular basis.

    Examples
    ========

    >>> from sympy import MatrixSymbol
    >>> A = MatrixSymbol('A', 3, 3)
    >>> y = MatrixSymbol('y', 3, 1)
    >>> x = (A.T*A).I * A * y

    See Also
    ========

    MatrixSymbol, MatAdd, MatMul, Transpose, Inverse
    Fg      &@TNc                 O   s"   t t|}tj| g|R i |S r   )mapr   r	   __new__)clsargskwargsr"   r"   r#   r+   Q   s   
zMatrixExpr.__new__returnc                 C      t r   NotImplementedErrorselfr"   r"   r#   shapeW      zMatrixExpr.shapec                 C      t S r   MatAddr3   r"   r"   r#   _add_handler[   r6   zMatrixExpr._add_handlerc                 C   r7   r   MatMulr3   r"   r"   r#   _mul_handler_   r6   zMatrixExpr._mul_handlerc                 C   s   t tj|  S r   )r<   r   NegativeOnedoitr3   r"   r"   r#   __neg__c      zMatrixExpr.__neg__c                 C   r0   r   r1   r3   r"   r"   r#   __abs__f      zMatrixExpr.__abs__other__radd__c                 C   s   t | |dd S NT)checkr9   r?   r4   rD   r"   r"   r#   __add__i      zMatrixExpr.__add__rJ   c                 C   s   t || dd S rF   rH   rI   r"   r"   r#   rE   n   rK   zMatrixExpr.__radd____rsub__c                 C   s   t | | dd S rF   rH   rI   r"   r"   r#   __sub__s      zMatrixExpr.__sub__rM   c                 C   s   t ||  dd S rF   rH   rI   r"   r"   r#   rL   x   rN   zMatrixExpr.__rsub____rmul__c                 C      t | | S r   r<   r?   rI   r"   r"   r#   __mul__}      zMatrixExpr.__mul__c                 C   rP   r   rQ   rI   r"   r"   r#   
__matmul__   rS   zMatrixExpr.__matmul__rR   c                 C      t ||  S r   rQ   rI   r"   r"   r#   rO      rS   zMatrixExpr.__rmul__c                 C   rU   r   rQ   rI   r"   r"   r#   __rmatmul__   rS   zMatrixExpr.__rmatmul____rpow__c                 C   rP   r   )MatPowr?   rI   r"   r"   r#   __pow__   rS   zMatrixExpr.__pow__rY   c                 C   s   t d)NzMatrix Power not definedr1   rI   r"   r"   r#   rW         zMatrixExpr.__rpow____rtruediv__c                 C   s   | |t j  S r   )r   r>   rI   r"   r"   r#   __truediv__   rS   zMatrixExpr.__truediv__r\   c                 C   s   t  r   r1   rI   r"   r"   r#   r[      s   zMatrixExpr.__rtruediv__c                 C   
   | j d S Nr   r5   r3   r"   r"   r#   rows      
zMatrixExpr.rowsc                 C   r]   N   r_   r3   r"   r"   r#   cols   ra   zMatrixExpr.colsc                 C   s   | j | jkS r   r`   rd   r3   r"   r"   r#   	is_square      zMatrixExpr.is_squarec                 C   s$   ddl m} ddlm} ||| S )Nr   Adjoint	Transpose)"sympy.matrices.expressions.adjointri   $sympy.matrices.expressions.transposerk   )r4   ri   rk   r"   r"   r#   _eval_conjugate   s   zMatrixExpr._eval_conjugatec                 K   s:   ddl m} tj| |    }| |   d|  }||fS )Nr   )I   )sympyro   r   ZHalfrn   )r4   deepZhintsro   realZimr"   r"   r#   as_real_imag   s   zMatrixExpr.as_real_imagc                 C      ddl m} || S )Nr   Inverse)Z"sympy.matrices.expressions.inverserw   )r4   rw   r"   r"   r#   _eval_inverse      zMatrixExpr._eval_inversec                 C      t | S r   rj   r3   r"   r"   r#   _eval_transpose      zMatrixExpr._eval_transposec                 C   s
   t | |S )z
        Override this in sub-classes to implement simplification of powers.  The cases where the exponent
        is -1, 0, 1 are already covered in MatPow.doit(), so implementations can exclude these cases.
        rX   )r4   Zexpr"   r"   r#   _eval_power   s   
zMatrixExpr._eval_powerc                    s$   | j r| S | j fdd| jD  S )Nc                    s   g | ]
}t |fi  qS r"   r   ).0xr.   r"   r#   
<listcomp>       z-MatrixExpr._eval_simplify.<locals>.<listcomp>)Zis_Atomr    r-   r4   r.   r"   r   r#   _eval_simplify   s   zMatrixExpr._eval_simplifyc                 C   ru   )Nr   rh   )rl   ri   )r4   ri   r"   r"   r#   _eval_adjoint   ry   zMatrixExpr._eval_adjointc                 C   s   t | ||S r   )r	   _eval_derivative_n_times)r4   r   nr"   r"   r#   r      s   z#MatrixExpr._eval_derivative_n_timesc                    s    |  |rt |S t| j S r   )hassuper_eval_derivative
ZeroMatrixr5   r4   r   	__class__r"   r#   r      s   

zMatrixExpr._eval_derivativec                 C   s4   ddl m} ||ddd}|du rtd|dS )z2Helper function to check invalid matrix dimensionsr   )check_assumptionsT)ZintegerZnonnegativeFz?The dimension specification {} should be a nonnegative integer.N)Zsympy.core.assumptionsr   
ValueErrorformat)r,   Zdimr   okr"   r"   r#   
_check_dim   s   zMatrixExpr._check_dimc                 K   s   t d| jj )NzIndexing not implemented for %s)r2   r   __name__r4   ijr.   r"   r"   r#   _entry   s   
zMatrixExpr._entryc                 C   rz   r   )r   r3   r"   r"   r#   r      r|   zMatrixExpr.adjointc                 C   s
   t j| fS )z2Efficiently extract the coefficient of a product. )r   One)r4   Zrationalr"   r"   r#   as_coeff_Mul   ra   zMatrixExpr.as_coeff_Mulc                 C   rz   r   )r   r3   r"   r"   r#   r      r|   zMatrixExpr.conjugatec                 C   ru   )Nr   	transpose)rm   r   )r4   r   r"   r"   r#   r      ry   zMatrixExpr.transposec                 C      |   S )zMatrix transpositionr   r3   r"   r"   r#   T   rZ   zMatrixExpr.Tc                 C   s   | j std|  S )NzInverse of non-square matrix)rf   r   rx   r3   r"   r"   r#   inverse  s   zMatrixExpr.inversec                 C   r   r   r   r3   r"   r"   r#   inv  r|   zMatrixExpr.invc                 C   r   r   r   r3   r"   r"   r#   ro   
  s   zMatrixExpr.Ic                 C   sV   dd }||o*||o*| j d u pd|kdko|| j k dko*d|kdko*|| jk dkS )Nc                 S   s   t | ttttfS r   )
isinstanceintr   r   r
   )idxr"   r"   r#   is_valid     z(MatrixExpr.valid_index.<locals>.is_validr   Fre   )r4   r   r   r   r"   r"   r#   valid_index  s   

zMatrixExpr.valid_indexc                 C   sH  t |tst |trddlm} || |dS t |trZt|dkrZ|\}}t |ts/t |tr;ddlm} || ||S t|t|}}| ||dkrR| ||S t	d||f t |t
tfr| j\}}t |tsqt	tdt|}|| }|| }| ||dkr| ||S t	d| t |ttfrt	td	t	d
|  )Nr   )MatrixSlice)r   Nrc   rp   FzInvalid indices (%s, %s)zo
                    Single indexing is only supported when the number
                    of columns is known.zInvalid index %szr
                    Only integers may be used when addressing the matrix
                    with a single index.zInvalid index, wanted %s[i,j])r   tuplesliceZ sympy.matrices.expressions.slicer   lenr   r   r   
IndexErrorr   r   r5   r   r   r
   )r4   keyr   r   r   r`   rd   r"   r"   r#   __getitem__  s2   

zMatrixExpr.__getitem__c                    sP   t  jttfrt  jttfstdddlm} | fddt jD S )a  
        Returns a dense Matrix with elements represented explicitly

        Returns an object of type ImmutableDenseMatrix.

        Examples
        ========

        >>> from sympy import Identity
        >>> I = Identity(3)
        >>> I
        I
        >>> I.as_explicit()
        Matrix([
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]])

        See Also
        ========
        as_mutable: returns mutable Matrix type

        z<Matrix with symbolic shape cannot be represented explicitly.r   ImmutableDenseMatrixc                    s&   g | ]  fd dt jD qS )c                    s   g | ]} |f qS r"   r"   r   r   )r   r4   r"   r#   r   W  s    z5MatrixExpr.as_explicit.<locals>.<listcomp>.<listcomp>)rangerd   )r   r3   )r   r#   r   W  s
    
z*MatrixExpr.as_explicit.<locals>.<listcomp>)	r   r`   r   r   rd   r   Zsympy.matrices.immutabler   r   )r4   r   r"   r3   r#   as_explicit9  s   zMatrixExpr.as_explicitc                 C   s   |    S )a  
        Returns a dense, mutable matrix with elements represented explicitly

        Examples
        ========

        >>> from sympy import Identity
        >>> I = Identity(3)
        >>> I
        I
        >>> I.shape
        (3, 3)
        >>> I.as_mutable()
        Matrix([
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]])

        See Also
        ========
        as_explicit: returns ImmutableDenseMatrix
        )r   
as_mutabler3   r"   r"   r#   r   [  s   zMatrixExpr.as_mutablec                 C   sR   ddl m} || jtd}t| jD ]}t| jD ]}| ||f |||f< qq|S )Nr   )empty)Zdtype)Znumpyr   r5   objectr   r`   rd   )r4   r   r   r   r   r"   r"   r#   	__array__t  s   zMatrixExpr.__array__c                 C   s   |   |S )z
        Test elementwise equality between matrices, potentially of different
        types

        >>> from sympy import Identity, eye
        >>> Identity(3).equals(eye(3))
        True
        )r   equalsrI   r"   r"   r#   r   |  s   	zMatrixExpr.equalsc                 C      | S r   r"   r3   r"   r"   r#   canonicalize  rC   zMatrixExpr.canonicalizec                 C   s   dt | fS rb   r;   r3   r"   r"   r#   as_coeff_mmul     zMatrixExpr.as_coeff_mmulc           
   
      s   ddl mmm mm	m ddlm 	fddi f 	f	dd	| }t	| \}}
|}t|dksMtt|dgkrO|S |du rh|D ]
}|dur_|}	 nqU|g|	R  S |||S )	a  
        Parse expression of matrices with explicitly summed indices into a
        matrix expression without indices, if possible.

        This transformation expressed in mathematical notation:

        `\sum_{j=0}^{N-1} A_{i,j} B_{j,k} \Longrightarrow \mathbf{A}\cdot \mathbf{B}`

        Optional parameter ``first_index``: specify which free index to use as
        the index starting the expression.

        Examples
        ========

        >>> from sympy import MatrixSymbol, MatrixExpr, Sum
        >>> from sympy.abc import i, j, k, l, N
        >>> A = MatrixSymbol("A", N, N)
        >>> B = MatrixSymbol("B", N, N)
        >>> expr = Sum(A[i, j]*B[j, k], (j, 0, N-1))
        >>> MatrixExpr.from_index_summation(expr)
        A*B

        Transposition is detected:

        >>> expr = Sum(A[j, i]*B[j, k], (j, 0, N-1))
        >>> MatrixExpr.from_index_summation(expr)
        A.T*B

        Detect the trace:

        >>> expr = Sum(A[i, i], (i, 0, N-1))
        >>> MatrixExpr.from_index_summation(expr)
        Trace(A)

        More complicated expressions:

        >>> expr = Sum(A[i, j]*B[k, j]*A[l, k], (j, 0, N-1), (k, 0, N-1))
        >>> MatrixExpr.from_index_summation(expr)
        A*B.T*A.T
        r   )Sumr   r   r<   r   trace)	bottom_upc                    sR    fdd}|  |ddd } |  |dfdd} fdd}|| S )	Nc                    s    fdd}|S )Nc                    sV   t | tsdS | j  krdS | jd  dkr)| jd jd  dkr'dS dS dS )NF   r   rp   rc   T)r   MatrixElementr-   r5   r   )i1posr"   r#   r      s   
z\MatrixExpr.from_index_summation.<locals>.remove_matelement.<locals>.repl_match.<locals>.funcr"   )r   r    r   )r   r#   
repl_match  s   zNMatrixExpr.from_index_summation.<locals>.remove_matelement.<locals>.repl_matchrc   c                 S   r]   r^   r-   r   r"   r"   r#   <lambda>     
 zLMatrixExpr.from_index_summation.<locals>.remove_matelement.<locals>.<lambda>rp   c                    s    | j d S r^   r   r   r   r"   r#   r     s    c                    s"   t |  frtdd | jS | S )Nc                 S   s   | | S r   r"   r   r"   r"   r#   r     s    z^MatrixExpr.from_index_summation.<locals>.remove_matelement.<locals>.<lambda>.<locals>.<lambda>)r   r   r-   r   )r<   r   r"   r#   r        " )replace)exprr   i2r   Zrule)r<   r   r   r   r   r#   remove_matelement  s   


z:MatrixExpr.from_index_summation.<locals>.remove_matelementc           (         s.  | j r/g }g }g }i }g }d}g }| jD ]#}	|	|}
t|
ts#J t|
tr3|
D ]}|| q*q||
 q|D ]Y\}}|d u rI|| q;t|trS|jd }|| || |d gt|  t|D ]$\}}||v r|| }||| |< ||f||d  |d < ||f||< qk|d7 }q;d}i }|t|k rt|D ]\}}d |v r||d f} nq|}g }||d  |d  }	 |\}}|| dkr|d dkr|||  n|||  || d|  }|d7 }d||< |d u r
|| d|  }||||f< n|}q|t|k sfdd|	 D }
|d fgdd	 |	 D  S | jrfd
d	| jD }tt}|D ]/}d}|D ]'\}}|d u rW|}qJtt|td}|| ||g|R    d}qJqD fdd	|	 D S t| tr| j\} }!}"|"d u rg }"| jD ]!}#|#|vrq||# \}$}%|$dkrtd| |"|% qt|"dkrd }"n0t|"dkr|"d d |"d d f}"n|"d d |"d d f}"|"d |"d krtd| t|"d }&t|&| |!| |!ffgS t| try| j\}'} }!| |v r6||  \}$}%|$dks)|'jd |%d kr6td|$|%f|'jd |!|v r]||! \}$}%|$dksP|'jd |%d kr]td|$|%f|'jd | |!krn| |v rn|'d fgS t|'| |!| |!ffgS t| r| jd dd | jdd  D dS | d fgS )Nr   rc   Trp   )r   c                    s0   i | ]\}}|t |d kr |n|d qS )rc   r   )r   fromiterr   kvr;   r"   r#   
<dictcomp>  s   0 zIMatrixExpr.from_index_summation.<locals>.recurse_expr.<locals>.<dictcomp>c                 S   s(   g | ]\\}}}t |||||ffqS r"   r   )r   r   r   r   r"   r"   r#   r     s    zIMatrixExpr.from_index_summation.<locals>.recurse_expr.<locals>.<listcomp>c                       g | ]} |qS r"   r"   r   r   )recurse_exprr"   r#   r         )r   c                    s*   g | ]\}}t  |g|R  |fqS r"   )r   r   r   )r   r"   r#   r     s   * z%index ranges should start from zero: z$upper index ranges should be equal: z$index range mismatch: {} vs. (0, {})c                 S   s   i | ]}|d  |dd qS )r   rc   Nr"   r   r"   r"   r#   r   I      )index_ranges)Zis_Mulr-   r   listappendr   r   	enumerateindexitemsr   Zis_Addcollectionsdefaultdictr   sortedr   r   r   Identityr5   r   )(r   r   Z
nonmatargsZpos_argZpos_indZdlinksZlink_indZcounterZargs_indr'   retvalsr   Z
arg_symbolZarg_indicesZindZother_iZcounter2lineseZline_start_indexZcur_ind_posZcur_lineZindex1drZnext_ind_posZindex2resZ
res_addendZscalarelemindicesr   r   r5   Zkr_indr1r2identityZmatrix_symbol)	r   r<   r   r   
dimensionsr   r   r   r   r"   r#   r     s   



















z5MatrixExpr.from_index_summation.<locals>.recurse_exprN)rq   r   r   r   r<   r   r   Zsympy.strategies.traverser   zipr   r   r   set)
r   Zfirst_indexZ
last_indexr   r   Zfactorsr   Zretexprr   Zind0r"   )
r   r<   r   r   r   r   r   r   r   r   r#   from_index_summation  s"    * z
zMatrixExpr.from_index_summationc                 C   s   ddl m} ||| S )Nrc   )ElementwiseApplyFunction)	applyfuncr   )r4   r    r   r"   r"   r#   r   \  s   
zMatrixExpr.applyfunc)T)F)NNN)Or   
__module____qualname____doc__Z	_iterableZ_op_priority	is_Matrixis_MatrixExprZis_IdentityZ
is_InverseZis_Transposeis_ZeroMatrixZ	is_MatAddZ	is_MatMulis_commutativeZ	is_number	is_symbolZ	is_scalarr   kindr+   propertytTupler
   r5   r:   r=   r@   rB   r(   NotImplementedr   rJ   rE   rM   rL   rR   rT   rO   rV   rY   rW   r\   r[   r`   rd   rf   rn   rt   rx   r{   r~   r   r   r   r   classmethodr   r   r   r   r   r   r   r   r   ro   r   r   r   r   r   r   r   r   staticmethodr   r   __classcell__r"   r"   r   r#   r)   &   s    










#" Or)   c                 C   s   dS )NFr"   ZlhsZrhsr"   r"   r#   _eval_is_eq`  r6   r  c                 C   s"   | j |j krdS | | jrdS d S )NFT)r5   r   r  r"   r"   r#   r  d  s
   
c                    s    fdd}|S )Nc                    s   t ttti  }g }g }| jD ]}t|tr|| q|| q|s) |S |r\ t krNt	t
|D ]}|| jsL||  |||< g } nq5n ||| jddg S |tkrh|| jddS | |g|R  jddS )NF)rr   )r   r<   r   r9   r-   r   r)   r   Z
_from_argsr   r   r   rR   r?   )r   Z	mat_classZnonmatricesZmatricesZtermr   r,   r"   r#   _postprocessorl  s,   



z)get_postprocessor.<locals>._postprocessorr"   )r,   r	  r"   r  r#   get_postprocessork  s   #r
  )r   r   c                    s   ddl m} | |}dd |D }ddlm fdd|D }dd   fd	d
fdd|D }|d }dd |dkrLtfdd|D S || |S )Nr   )ArrayDerivativec                 S   s   g | ]}|  qS r"   )buildr   r"   r"   r#   r     r   z&_matrix_derivative.<locals>.<listcomp>convert_array_to_matrixc                    s   g | ]} fd d|D qS )c                    r   r"   r"   r   r  r"   r#   r     r   z1_matrix_derivative.<locals>.<listcomp>.<listcomp>r"   r   r  r"   r#   r     r   c                 S      t | tr| jS dS )Nrc   rc   r   r)   r5   r   r"   r"   r#   
_get_shape     
z&_matrix_derivative.<locals>._get_shapec                    s   t  fdd| D S )Nc                    s"   g | ]} |D ]}|d vqqS ))rc   Nr"   )r   r   r   r  r"   r#   r     r   z8_matrix_derivative.<locals>.get_rank.<locals>.<listcomp>)sum)partsr  r"   r#   get_rank  s   z$_matrix_derivative.<locals>.get_rankc                    r   r"   r"   r   )r  r"   r#   r     r   c                 S   s   t | dkr
| d S | d d \}}|jr|j}|tdkr!|}n|tdkr*|}n|| }t | dkr6|S |jr=td|t| dd   S )Nrc   r   rp    )r   r   r   r   r   r   r   )r  Zp1Zp2Zpbaser"   r"   r#   contract_one_dims  s   z-_matrix_derivative.<locals>.contract_one_dimsrp   c                    r   r"   r"   r   )r  r"   r#   r     r   )Z$sympy.tensor.array.array_derivativesr  _eval_derivative_matrix_linesZ3sympy.tensor.array.expressions.conv_array_to_matrixr  r   r   )r   r   r  r   r  Zranksrankr"   )r  r  r  r  r#   _matrix_derivative  s   

r  c                   @   s`   e Zd Zedd Zedd Zedd ZdZdZdZ	dd Z
dd	 Zed
d Zdd ZdS )r   c                 C   r]   r^   r   r3   r"   r"   r#   r     r   zMatrixElement.<lambda>c                 C   r]   rb   r   r3   r"   r"   r#   r     r   c                 C   r]   )Nrp   r   r3   r"   r"   r#   r     r   Tc                 C   s   t t||f\}}ddlm} t||fr!|jr!|jr!|||f S t|tr+t|}nt|}t|jt	s9t
dt| |||}|S )Nr   
MatrixBasez2First argument of MatrixElement should be a matrix)r*   r   rq   r  r   Z
is_Integerstrr   r   r   	TypeErrorr
   r+   )r,   namer   mr  objr"   r"   r#   r+     s   

zMatrixElement.__new__c                    sD     dd}|r fdd| jD }n| j}|d |d |d f S )Nrr   Tc                    s   g | ]
}|j d i  qS )r"   )r?   )r   r'   r   r"   r#   r     r   z&MatrixElement.doit.<locals>.<listcomp>r   rc   rp   )getr-   )r4   r.   rr   r-   r"   r   r#   r?     s
   zMatrixElement.doitc                 C   s   | j dd  S rb   r   r3   r"   r"   r#   r     s   zMatrixElement.indicesc                 C   sZ  ddl m}m}m} t|ts*ddl m} t| j|r'| j|| j	| j
f S tjS | jd }| jj\}}||jd kr\t| jd |jd d|d ft| jd |jd d|d f S t|tr| jdd  \}	}
|d|d\}}|jd }|j\}}|||	|f |||f | |||
f  |d|d f|d|d f S | |jd rd S tjS )Nr   )r   symbolsDummyr  rc   rp   zz1, z2r  )rq   r   r&  r'  r   r   r  parentZdiffr   r   r   Zeror-   r5   r   rw   r   )r4   r   r   r&  r'  r  Mr#  r   r   r   r   r   Yr   r   r"   r"   r#   r     s*   




HzMatrixElement._eval_derivativeN)r   r   r   r   r(  r   r   	_diff_wrtr   r   r+   r?   r   r   r"   r"   r"   r#   r     s    
r   c                   @   sh   e Zd ZdZdZdZdZdd Zedd Z	edd	 Z
d
d Zedd Zdd Zdd Zdd ZdS )MatrixSymbola  Symbolic representation of a Matrix object

    Creates a SymPy Symbol to represent a Matrix. This matrix has a shape and
    can be included in Matrix Expressions

    Examples
    ========

    >>> from sympy import MatrixSymbol, Identity
    >>> A = MatrixSymbol('A', 3, 4) # A 3 by 4 Matrix
    >>> B = MatrixSymbol('B', 4, 3) # A 4 by 3 Matrix
    >>> A.shape
    (3, 4)
    >>> 2*A*B + Identity(3)
    I + 2*A*B
    FTc                 C   sL   t |t |}}| | | | t|trt|}t| |||}|S r   )r   r   r   r   r   r	   r+   )r,   r"  r   r#  r$  r"   r"   r#   r+     s   


zMatrixSymbol.__new__c                 C   s   | j d | j d fS )Nrc   rp   r   r3   r"   r"   r#   r5   (  s   zMatrixSymbol.shapec                 C   s   | j d jS r^   )r-   r"  r3   r"   r"   r#   r"  ,  rg   zMatrixSymbol.namec                 K   s   t | ||S r   r   r   r"   r"   r#   r   0  r   zMatrixSymbol._entryc                 C   s   | hS r   r"   r3   r"   r"   r#   free_symbols3  s   zMatrixSymbol.free_symbolsc                 K   r   r   r"   r   r"   r"   r#   r   7  rC   zMatrixSymbol._eval_simplifyc                 C   s   t | jd | jd S Nr   rc   )r   r5   r   r"   r"   r#   r   :  s   zMatrixSymbol._eval_derivativec                 C   s   | |kr5| j d dkrt|j d | j d ntj}| j d dkr+t|j d | j d ntj}t||ggS | j d dkrCt| j d ntj}| j d dkrTt| j d ntj}t||ggS r/  )r5   r   r   r)  _LeftRightArgsr   r   )r4   r   firstsecondr"   r"   r#   r  >  s   **""z*MatrixSymbol._eval_derivative_matrix_linesN)r   r   r   r   r   r   r,  r+   r   r5   r"  r   r.  r   r   r  r"   r"   r"   r#   r-    s     


r-  c                 C   s   dd | j D S )Nc                 S   s   g | ]}|j r|qS r"   )r   )r   Zsymr"   r"   r#   r   N      z"matrix_symbols.<locals>.<listcomp>)r.  )r   r"   r"   r#   matrix_symbolsM  rA   r4  c                   @   s   e Zd ZdZejfddZedd Zej	dd Zedd Z
e
j	d	d Z
d
d Zdd Zedd Zdd Zdd Zdd Zdd Zdd Zdd ZdS )r0  a  
    Helper class to compute matrix derivatives.

    The logic: when an expression is derived by a matrix `X_{mn}`, two lines of
    matrix multiplications are created: the one contracted to `m` (first line),
    and the one contracted to `n` (second line).

    Transposition flips the side by which new matrices are connected to the
    lines.

    The trace connects the end of the two lines.
    c                 C   sB   dd |D | _ | j | _d| _d| _| j | _d| _d| _|| _d S )Nc                 S      g | ]}|qS r"   r"   r   r"   r"   r#   r   `      z+_LeftRightArgs.__init__.<locals>.<listcomp>r   rc   )_lines_first_pointer_parent_first_pointer_index_first_line_index_second_pointer_parent_second_pointer_index_second_line_indexhigher)r4   r   r>  r"   r"   r#   __init___  s   
z_LeftRightArgs.__init__c                 C      | j | j S r   r8  r9  r3   r"   r"   r#   first_pointeri  rg   z_LeftRightArgs.first_pointerc                 C      || j | j< d S r   rA  r4   valuer"   r"   r#   rB  m     c                 C   r@  r   r;  r<  r3   r"   r"   r#   second_pointerq  rg   z_LeftRightArgs.second_pointerc                 C   rC  r   rG  rD  r"   r"   r#   rH  u  rF  c                    s"    fdd j D }d| jf S )Nc                       g | ]}  |qS r"   _buildr   r3   r"   r#   r   z  r3  z+_LeftRightArgs.__repr__.<locals>.<listcomp>z#_LeftRightArgs(lines=%s, higher=%s))r7  r>  )r4   Zbuiltr"   r3   r#   __repr__y  s
   z_LeftRightArgs.__repr__c                 C   s:   | j | j| _| _ | j| j| _| _| j| j| _| _| S r   )r;  r8  r<  r9  r=  r:  r3   r"   r"   r#   r     s   z_LeftRightArgs.transposec                 C   sZ   ddl m} t| |r|  S t| tr+t| dkr| d S | d dd | d D  S | S )Nr   ExprBuilderrc   c                 S   s   g | ]}t |qS r"   )r0  rK  r   r"   r"   r#   r     r3  z)_LeftRightArgs._build.<locals>.<listcomp>)sympy.core.exprrN  r   r  r   r   )r   rN  r"   r"   r#   rK    s   

z_LeftRightArgs._buildc                    sB    fdd j D } jdkr|  jg7 }dd |D }|S )Nc                    rI  r"   rJ  r   r3   r"   r#   r     r3  z(_LeftRightArgs.build.<locals>.<listcomp>rc   c                 S   r5  r"   r"   r   r"   r"   r#   r     r6  )r7  r>  rK  )r4   datar"   r3   r#   r    s
   
z_LeftRightArgs.buildc                 C   s   | j dkr| jdkrtddd }|| j d || jd krC|| jdkr/| j | jd  S || j dkr?| j d | jj S td| j dkrO| j | jj S | jS )Nrc   z.higher dimensional array cannot be representedc                 S   r  )N)NNr  r  r"   r"   r#   r    r  z._LeftRightArgs.matrix_form.<locals>._get_shaper  )r   r   zincompatible shapes)r1  r>  r   r2  r   )r4   r  r"   r"   r#   matrix_form  s   
z_LeftRightArgs.matrix_formc                 C   sb   d}| j dkr|tdd | j jD 7 }| jdkr&|tdd | jjD 7 }| jdkr/|d7 }|S )zl
        Number of dimensions different from trivial (warning: not related to
        matrix rank).
        r   rc   c                 S      g | ]}|d kqS rc   r"   r   r"   r"   r#   r     r   z'_LeftRightArgs.rank.<locals>.<listcomp>c                 S   rR  rS  r"   r   r"   r"   r#   r     r   rp   )r1  r  r5   r2  r>  )r4   r  r"   r"   r#   r    s   


z_LeftRightArgs.rankc                 C   sF   ddl m} ddlm} ddlm} ||||||gdg|jd}|S )Nr   rM  r   )ArrayTensorProduct)ArrayContraction)rc   rp   )Z	validator)rO  rN  Z*tensor.array.expressions.array_expressionsrT  rU  Z	_validate)r4   ZpointerrD   rN  rT  rU  Zsubexprr"   r"   r#   _multiply_pointer  s    
z _LeftRightArgs._multiply_pointerc                 C      |  j |9  _ d S r   )rB  rI   r"   r"   r#   append_first  r   z_LeftRightArgs.append_firstc                 C   rW  r   )rH  rI   r"   r"   r#   append_second  r   z_LeftRightArgs.append_secondN)r   r   r   r   r   r   r?  r   rB  setterrH  rL  r   r  rK  r  rQ  r  rV  rX  rY  r"   r"   r"   r#   r0  Q  s*    





r0  c                 C   s&   ddl m} t| tr| S || ggS )Nr   r   )rq   r   r   r)   )r   r   r"   r"   r#   _make_matrix  s   
r[  rc   r;   r8   r}   rj   rv   )r   r   r   )Btypingr   r  Zsympy.core.logicr   	functoolsr   r   r   Z
sympy.corer   r   r   r	   r
   r   r   Zsympy.core.decoratorsr   Zsympy.core.compatibilityr   r   Zsympy.core.symbolr   Zsympy.core.sympifyr   r   Zsympy.functionsr   r   Z(sympy.functions.special.tensor_functionsr   Zsympy.matrices.commonr   Zsympy.simplifyr   Zsympy.matrices.matricesr   Zsympy.utilities.miscr   Zsympy.multipledispatchr   r(   r)   r  r
  Z"_constructor_postprocessor_mappingr  r   r-  r4  r0  r[  matmulr<   Zmataddr9   ZmatpowrX   r   rk   r   rw   Zspecialr   r   r"   r"   r"   r#   <module>   sV    $
    >

(/AE 