o
    8Va                     @   s@  d dl mZ d dlZd dlmZ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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 m!Z! G dd deeZ"e#ee"fe" dd Z$dd Z%dd Z&dd Z'dd Z(edd eeee&e%e'e(ee fZ)eedd ee) Z*dS )    )reduceN)AddBasicsympify)addadjoint)
ShapeError)
MatrixBase	transpose)rm_idunpackflattensort	conditionexhaustdo_oneglom)
MatrixExpr)
ZeroMatrixGenericZeroMatrix)default_sort_keysiftc                   @   sh   e Zd ZdZdZe Zdddd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 )MatAdda3  A Sum of Matrix Expressions

    MatAdd inherits from and operates like SymPy Add

    Examples
    ========

    >>> from sympy import MatAdd, MatrixSymbol
    >>> A = MatrixSymbol('A', 5, 5)
    >>> B = MatrixSymbol('B', 5, 5)
    >>> C = MatrixSymbol('C', 5, 5)
    >>> MatAdd(A, B, C)
    A + B + C
    TF)evaluatecheck_sympifyc                   s   |s j S tt fdd|}|rttt|}tj g|R  }|r6tdd |D r2t	|S t
|  |rLtdd |D rHt|ddiS t|}|S )Nc                    s
    j | kS N)identity)icls C/usr/lib/python3/dist-packages/sympy/matrices/expressions/matadd.py<lambda>*      
 z MatAdd.__new__.<locals>.<lambda>c                 s       | ]	}t |t V  qd S r   
isinstancer   .0r    r#   r#   r$   	<genexpr>1       z!MatAdd.__new__.<locals>.<genexpr>c                 s   r'   r   r(   r*   r#   r#   r$   r,   6   r-   r   T)r   listfiltermapr   r   __new__allr   Zfromitervalidatecanonicalize)r"   r   r   r   argsobjr#   r!   r$   r1   $   s   
zMatAdd.__new__c                 C   s   | j d jS Nr   )r5   shapeselfr#   r#   r$   r8   <   s   zMatAdd.shapec                    s   t  fdd| jD  S )Nc                    s    g | ]}|j  fi qS r#   )_entryr+   argr    jkwargsr#   r$   
<listcomp>A   s     z!MatAdd._entry.<locals>.<listcomp>)r   r5   )r:   r    r?   r@   r#   r>   r$   r;   @   s   zMatAdd._entryc                 C      t dd | jD   S )Nc                 S      g | ]}t |qS r#   r   r<   r#   r#   r$   rA   D       z*MatAdd._eval_transpose.<locals>.<listcomp>r   r5   doitr9   r#   r#   r$   _eval_transposeC      zMatAdd._eval_transposec                 C   rB   )Nc                 S   rC   r#   r   r<   r#   r#   r$   rA   G   rD   z(MatAdd._eval_adjoint.<locals>.<listcomp>rE   r9   r#   r#   r$   _eval_adjointF   rH   zMatAdd._eval_adjointc                    s(   ddl m   t fdd| jD   S )N   tracec                    s   g | ]} |qS r#   r#   r<   rK   r#   r$   rA   K   rD   z&MatAdd._eval_trace.<locals>.<listcomp>)rL   r   r5   rF   r9   r#   rK   r$   _eval_traceI   s   zMatAdd._eval_tracec                    s8     dd}|r fdd| jD }n| j}tt| S )NdeepTc                    s   g | ]
}|j d i  qS )r#   )rF   r<   r@   r#   r$   rA   P   s    zMatAdd.doit.<locals>.<listcomp>)getr5   r4   r   )r:   r@   rN   r5   r#   rO   r$   rF   M   s
   zMatAdd.doitc                    s"    fdd| j D }dd |D S )Nc                    s   g | ]}|  qS r#   )_eval_derivative_matrix_linesr<   xr#   r$   rA   V   s    z8MatAdd._eval_derivative_matrix_lines.<locals>.<listcomp>c                 S   s   g | ]	}|D ]}|qqS r#   r#   )r+   r    r?   r#   r#   r$   rA   W   s    )r5   )r:   rS   Z	add_linesr#   rR   r$   rQ   U   s   z$MatAdd._eval_derivative_matrix_linesN)__name__
__module____qualname____doc__Z	is_MatAddr   r   r1   propertyr8   r;   rG   rI   rM   rF   rQ   r#   r#   r#   r$   r      s    
r   c                  G   sT   t dd | D std| d }| dd  D ]}|j|jkr'td||f qd S )Nc                 s   s    | ]}|j V  qd S r   )Z	is_Matrixr<   r#   r#   r$   r,   \   s    zvalidate.<locals>.<genexpr>z Mix of Matrix and Scalar symbolsr   rJ   z"Matrices %s and %s are not aligned)r2   	TypeErrorr8   r	   )r5   ABr#   r#   r$   r3   [   s   r3   c                 C   s   |   d S r7   )as_coeff_mmulr=   r#   r#   r$   r%   d   s    r%   c                 C   s   t |  d S NrJ   )r   r\   r]   r#   r#   r$   r%   e   s    c                 C   s   | dkr|S | | S r^   r#   )ZcntZmatr#   r#   r$   combinef   s   r_   c                 C   sB   t | jdd }t|d dkrt|d ttj|d g  S | S )a   Merge explicit MatrixBase arguments

    Examples
    ========

    >>> from sympy import MatrixSymbol, eye, Matrix, MatAdd, pprint
    >>> from sympy.matrices.expressions.matadd import merge_explicit
    >>> A = MatrixSymbol('A', 2, 2)
    >>> B = eye(2)
    >>> C = Matrix([[1, 2], [3, 4]])
    >>> X = MatAdd(A, B, C)
    >>> pprint(X)
        [1  0]   [1  2]
    A + [    ] + [    ]
        [0  1]   [3  4]
    >>> pprint(merge_explicit(X))
        [2  2]
    A + [    ]
        [3  5]
    c                 S   
   t | tS r   )r)   r
   r]   r#   r#   r$   r%      r&   z merge_explicit.<locals>.<lambda>TrJ   F)r   r5   lenr   r   operatorr   )Zmataddgroupsr#   r#   r$   merge_explicitm   s   rd   c                 C   s   | dkpt | tS r7   )r)   r   rR   r#   r#   r$   r%      s    c                 C   r`   r   )r)   r   rR   r#   r#   r$   r%      r&   )+	functoolsr   rb   Z
sympy.corer   r   r   Zsympy.core.addr   Zsympy.functionsr   Zsympy.matrices.commonr	   Zsympy.matrices.matricesr
   Z$sympy.matrices.expressions.transposer   Zsympy.strategiesr   r   r   r   r   r   r   r   Z"sympy.matrices.expressions.matexprr   Z"sympy.matrices.expressions.specialr   r   Zsympy.utilitiesr   r   r   Zregister_handlerclassr3   Z	factor_ofZ	matrix_ofr_   rd   Zrulesr4   r#   r#   r#   r$   <module>   s:    (H	



